torchbp.data module

Lazy out-of-core data sources for backprojection functions.

A LazyData stands in for the range compressed data tensor when the raw data does not fit in memory. It mimics the small part of the Tensor interface that the library uses (shape, dtype, device, dim(), len() and contiguous slicing along the sweep axis) so that a Tensor and a LazyData are interchangeable at every Python call site; the data is materialized with LazyData.load() only where a kernel consumes it.

Slicing (lazy[i0:i1]) is free: it returns a view restricted to the sweep range without loading anything. torchbp.ops.ffbp() slices subaperture views down its recursion tree and loads only leaf-sized chunks, so its peak memory stays at image scale; torchbp.autofocus.gpga() and torchbp.autofocus.gpga_tde() additionally read the data in bounded chunks in their estimation stage. Functions that need the full tensor at once (direct backprojection, afbp, cfbp) accept a LazyData but load the whole extent at entry, so they are not memory efficient with it.

Gradients do not flow to the data through a LazyData (requires_grad is always False). Gradients with respect to positions are unaffected.

class torchbp.data.CachedData(source, path=None, storage='disk', storage_device=None)[source]

Bases: LazyData

Cache of another LazyData’s transformed output.

The first load of a sweep runs the source’s pipeline and writes the result to the cache; later loads of the same sweeps read the cache instead. Use it when an expensive per-chunk transform (range compression, RVP removal) feeds an iterative consumer like torchbp.autofocus.gpga(), which re-reads the data every iteration: the pipeline then runs once per sweep instead of once per iteration.

Only 2D [nsweeps, samples] sources are supported. The cache is keyed per sweep, so overlapping or out-of-order loads fill it incrementally. Call fill() to populate it eagerly and release the source, so the raw data backing the pipeline can be freed before image formation.

Parameters:

  • source (LazyData) – The source (typically with a transform pipeline) to cache.

  • path (str or None) – Cache file path for storage="disk". Default None creates a temporary file that is deleted when this object is garbage collected. A given path is left on disk (but always created anew, never trusted stale). The default avoids a RAM-backed system temp dir (tmpfs /tmp) by falling back to /var/tmp.

  • storage (str) –

    • “disk” (default): memory-mapped scratch file. Bounded RAM use; for data larger than memory.

    • ”ram”: a tensor on storage_device. For data that fits, a filled RAM cache on the compute device is equivalent to an eagerly transformed tensor, loads are zero-copy views, so the lazy plumbing adds no per-load overhead, while fill() still lets the raw source be freed.

  • storage_device (str, torch.device or None) – Device of the “ram” cache tensor. Default None uses the source’s (compute) device. Pass “cpu” with a CUDA compute device when the data fits in system RAM but not in VRAM: the cache then stages in pinned CPU memory and loads copy to the GPU chunk by chunk. Not used with storage="disk".

fill(block_bytes=67108864)[source]

Eagerly populate the whole cache in bounded blocks and release the source.

After this the cache is the only backing store: the source (and whatever raw data its pipeline reads) is dropped and can be freed by the caller, restoring the memory timeline of an eager range compression, compress once, free the raw data, with the compressed data in the cache storage instead of pinned by the pipeline. Returns self for chaining.

Parameters:

block_bytes (int)

Return type:

CachedData

class torchbp.data.CallbackData(load_fn, shape, dtype=torch.complex64, device='cpu', transform=None)[source]

Bases: LazyData

Lazy data source backed by a user callback.

Parameters:

  • load_fn (callable) – load_fn(start, stop) -> Tensor returning the range compressed sweeps [start, stop), shape [stop - start, samples]. May return anything torch.as_tensor accepts; the result is converted to dtype and moved to device.

  • shape (tuple) – Full data shape [nsweeps, samples].

  • dtype (torch.dtype) – Dtype of the data. Default is torch.complex64.

  • device (str or torch.device) – Device the backprojection runs on. Default is "cpu".

  • transform (callable or None) – Optional per-chunk hook, see LazyData.

class torchbp.data.LazyData(shape, dtype, device, transform=None)[source]

Bases: object

Base class for lazy radar data sources.

Subclasses implement _load_range() returning the sweeps [start, stop) as a Tensor (or anything torch.as_tensor accepts) and call super().__init__ with the full extent shape, the dtype/device the loaded chunks should have, and an optional transform.

Parameters:

  • shape (tuple) – Shape of the full data, [nsweeps, samples].

  • dtype (torch.dtype) – Dtype of the loaded data (after conversion), typically torch.complex64.

  • device (str or torch.device) – Device the loaded chunks are moved to. This is the device the backprojection runs on; lazy.device reports it so existing device = data.device call sites keep working.

  • transform (callable or None) – Optional per-chunk pipeline transform(chunk, start, stop) -> Tensor applied after the chunk is moved to device but before dtype conversion. Use it for preprocessing that would otherwise force materializing the full tensor up front: windowing, range compression of raw sweeps, RVP removal, modulation. start/stop are the absolute sweep indices of the chunk, so per-sweep factors (an azimuth window over the full aperture) can be indexed. The transform may change the dtype and the number of samples per sweep — shape and dtype describe its output — but must not change the number of sweeps. For an iterative consumer (torchbp.autofocus.gpga()) that re-reads the data every iteration, wrap the source in CachedData so an expensive pipeline runs only once per sweep.

property device: device
dim()[source]
Return type:

int

property dtype: dtype
load()[source]

Materialize the full extent of this source (or view) as a Tensor.

Return type:

Tensor

property ndim: int
property requires_grad: bool
resolve_conj()[source]
Return type:

LazyData

property shape: Size
class torchbp.data.MemmapData(array, device='cpu', dtype=None, transform=None, shape=None)[source]

Bases: LazyData

Lazy data source backed by a sliceable array on disk.

Wraps any array supporting array[start:stop], .shape and .dtype — numpy.memmap, an h5py dataset, a zarr array — and loads slices on demand.

Parameters:

  • array (array-like) – Sliceable array of shape [nsweeps, samples].

  • device (str or torch.device) – Device the backprojection runs on. Default is "cpu".

  • dtype (torch.dtype or None) – Dtype of the loaded data after the transform. Default None keeps the array’s own dtype; pass it explicitly when the transform changes the dtype.

  • transform (callable or None) – Optional per-chunk pipeline, see LazyData.

  • shape (tuple or None) – Output shape override, [nsweeps, samples]. Default None uses the array’s shape; pass it when the transform changes the number of samples per sweep.

torchbp.data.available_ram()[source]

Available system RAM in bytes, counting reclaimable caches as available.

Use it to size a CachedData RAM/disk decision, e.g. storage="ram" if nbytes < available_ram() // 2 else "disk".

  • Linux: MemAvailable from /proc/meminfo, the kernel’s estimate of memory usable without swapping; it counts reclaimable page cache and slab, unlike MemFree.

  • Windows: GlobalMemoryStatusEx ullAvailPhys, which includes the standby (file cache) list — the same semantics.

  • Elsewhere (or on failure): half of total physical RAM as a conservative estimate, or 2 GiB if even that is unavailable.

Return type:

int

torchbp.data.materialize(data)[source]

Return data as a Tensor, loading it if it is a LazyData.

Return type:

Tensor