fastsar 0.1.0__py3-none-any.whl

fastsar/__init__.py

@@ -0,0 +1,16 @@
+"""Fast spotlight SAR image formation: factorized backprojection with kernels for Cloud TPUs (Pallas), Nvidia GPUs
+(CUDA) and x86 CPUs (C++/OpenMP), and polar format with its geometric resampling.  See fastsar.api.form_image."""
+import os as _os
+
+# JAX takes 75% of a GPU's memory when it first runs there unless told otherwise; FastSAR's CUDA kernels (CuPy) and
+# its JAX programs share the GPU, so JAX allocates as it goes (a setting the user made before importing JAX stands);
+# JAX keeps what it has allocated, and the CuPy formers shrink their working sets to what is left
+_os.environ.setdefault('XLA_PYTHON_CLIENT_PREALLOCATE', 'false')
+from .api import form_image, available_backends, ImageFormer  # noqa: F401
+from . import io, autofocus, stripmap, burst, patches, quality, products  # noqa: F401
+from .bp import backproject, plane_points  # noqa: F401
+from .exact import ExactFormer  # noqa: F401
+from .cphd import form_cphd  # noqa: F401
+from .memory import MemoryWarning  # noqa: F401
+
+__version__ = '0.1.0'

fastsar/_build.py

@@ -0,0 +1,59 @@
+"""Compiling the C++ kernels on first use: one shared object per source, compiler, flags and host, cached in
+~/.cache/fastsar. The source is compiled with the given flags and linked in a separate step without them, so that
+-ffast-math, when a kernel is compiled with it, does not link crtfastmath.o (which would switch the whole process to
+flush-to-zero on load). Concurrent first builds each write their own temporary files and the last os.replace wins,
+with identical contents."""
+import hashlib
+import os
+import platform
+import subprocess
+
+
+def _host_tag():
+    """The machine and its CPU's instruction set extensions (a -march=native build must not be loaded on another
+    CPU sharing the same home directory)."""
+    tag = platform.machine()
+    try:
+        with open('/proc/cpuinfo') as fh:
+            for line in fh:
+                if line.startswith(('flags', 'Features')):
+                    return tag + line
+    except OSError:
+        pass
+    return tag + platform.processor()
+
+
+def shared_object(name, src, flags):
+    """Path of the shared object built from the C++ source text src with the compiler flags (a list), building it
+    if it is not cached. The compiler is $CXX (default g++); OpenMP and -fPIC are added."""
+    cxx = os.environ.get('CXX', 'g++')
+    key = '\0'.join([src, cxx, ' '.join(flags), _host_tag()])
+    tag = hashlib.sha1(key.encode()).hexdigest()[:12]
+    d = os.path.join(os.path.expanduser('~'), '.cache', 'fastsar')
+    os.makedirs(d, exist_ok=True)
+    so = os.path.join(d, f'lib{name}_{tag}.so')
+    if os.path.exists(so):
+        return so
+    stem = os.path.join(d, f'{name}_{tag}.{os.getpid()}')
+    cpp, obj, tmp = stem + '.cpp', stem + '.o', stem + '.so.tmp'
+    with open(cpp, 'w') as fh:
+        fh.write(src)
+    try:
+        for cmd in ([cxx] + list(flags) + ['-fopenmp', '-fPIC', '-c', cpp, '-o', obj],
+                    [cxx, '-shared', '-fopenmp', obj, '-o', tmp]):
+            try:
+                subprocess.run(cmd, check=True, capture_output=True, text=True)
+            except FileNotFoundError:
+                raise RuntimeError(f"fastsar's cpu backend compiles its C++ kernels on first use and needs a C++ compiler "
+                                   f'with OpenMP: {cxx!r} was not found. Install g++ (Debian/Ubuntu: apt install g++; '
+                                   'RHEL/Fedora: dnf install gcc-c++) or point CXX at one') from None
+            except subprocess.CalledProcessError as e:
+                raise RuntimeError(f'compiling the {name} kernel failed ({" ".join(cmd)}):\n{e.stderr.strip()}') from None
+        os.replace(tmp, so)
+    finally:
+        for f in (cpp, obj, tmp):
+            try:
+                os.remove(f)
+            except OSError:
+                pass
+    return so

fastsar/api.py

@@ -0,0 +1,487 @@
+"""One call for spotlight image formation on any of the supported devices.
+
+    import fastsar
+    img = fastsar.form_image(S, ant, fmin, df, nx, ny, spx, spy, e1, e2)            # factorized backprojection
+    img = fastsar.form_image(..., algorithm='pfa')                                  # polar format
+    img = fastsar.form_image(..., algorithm='bp')                                   # exact backprojection
+
+S is the phase history [pulses, samples] (complex, frequency domain, motion compensated to the scene reference
+point), ant the antenna phase centers [pulses, 3] in a frame whose origin is the scene reference point, fmin and df
+the first frequency and the sample spacing (Hz), and nx, ny, spx, spy, e1, e2 the output grid: pixel counts and
+spacings (m) along the unit vectors e1 (azimuth) and e2 (range) of the image plane. The result is a complex64 image
+[nx, ny] on that grid.
+
+Backends for factorized backprojection (the same plan, filters and float64 geometry on every device):
+
+    'tpu'   Pallas kernels (JAX on a Cloud TPU)
+    'cuda'  CUDA kernels through CuPy (Nvidia GPU)
+    'cpu'   C++ kernels with OpenMP (x86-64; compiled with g++ on first use)
+    'jax'   the plain JAX program, on whatever device JAX has
+    'auto'  tpu if JAX sees a TPU, else cuda if CuPy sees a GPU, else cpu
+
+ref [P]: the range (one way) each pulse's samples are referenced to, when it is not |ant| (a bistatic collection's
+half path, a vendor's reference point); both algorithms honor it (polar format re-references the samples to |ant|).
+
+pfa_guard: margin (m) around the scene that polar format keeps free of wrap-around; 300 m suits orbital scenes of a
+few kilometers and must be smaller for small simulated scenes.
+
+Precision: 'float32' (default), 'float16' (CUDA: float16 storage throughout with float32 accumulation; JAX: float16
+throughout), 'single-pass' (TPU: one bfloat16 pass per product, what the matrix unit does by itself), 'three-pass'
+(TPU: float32-class accuracy). On the TPU 'float32' means three-pass.
+
+Tile size: T='auto' (default) picks the largest final tile (32 or 16 pixels) whose predicted error against exact
+backprojection meets target_db (-40 dB by default). The error of the final stage's plane-wave model grows as the
+square of the tile size and falls with range, so orbital collections keep T=32 and short-range (airborne) ones
+drop to 16. The prediction is kept as ImageFormer.predicted_error_db.
+"""
+import os
+import threading
+
+import numpy as np
+
+from . import memory as _mem
+from .sim import Collect
+
+
+def _window(P, K, sll=35.0, nbar=4):
+    from scipy.signal.windows import taylor
+    return taylor(P, nbar=nbar, sll=sll, norm=False).astype(np.float32), taylor(K, nbar=nbar, sll=sll, norm=False).astype(np.float32)
+
+
+def available_backends():
+    """The backends this machine can run, in the order 'auto' tries them."""
+    out = []
+    try:
+        import jax
+        if any(d.platform == 'tpu' for d in jax.devices()):
+            out.append('tpu')
+    except Exception:
+        pass
+    try:
+        import cupy
+        if cupy.cuda.runtime.getDeviceCount() > 0:
+            out.append('cuda')
+    except Exception:
+        pass
+    out.append('cpu')
+    return out
+
+
+_JAX_PROGRAMS = {}         # compiled JAX/TPU programs by plan signature (ImageFormer)
+_JAX_LOCK = threading.Lock()   # mosaic workers build formers concurrently: one build per signature, no torn eviction
+
+BACKENDS = ('auto', 'tpu', 'cuda', 'cpu', 'jax')
+
+
+def _backend(backend):
+    """backend checked against BACKENDS and this machine, 'auto' resolved."""
+    if backend not in BACKENDS:
+        raise ValueError(f'unknown backend {backend!r}: one of {", ".join(map(repr, BACKENDS))}')
+    if backend in ('cpu', 'jax'):            # no device probe (a mosaic builds a former per patch)
+        return backend
+    have = available_backends()
+    if backend == 'auto':
+        return have[0]
+    if backend == 'cuda' and 'cuda' not in have:
+        raise ValueError("backend 'cuda' needs CuPy and an Nvidia GPU; this machine has " + ', '.join(have + ['jax']))
+    if backend == 'tpu' and 'tpu' not in have and not os.environ.get('FFBP_FORCE_TPU_KERNELS'):
+        raise ValueError("backend 'tpu' needs JAX on a Cloud TPU; this machine has " + ', '.join(have + ['jax']))
+    return backend
+
+
+def _check_history(S, P=None, name='phase history', finite=True):
+    """S [P, K] complex (numpy or cupy) with at least 2 pulses and 2 samples, all finite (finite=False: the scan
+    for NaN and inf is left to the caller); -> S unchanged. Real or integer samples are refused (an I/Q pair
+    interleaved along the last axis is S[..., 0] + 1j * S[..., 1])."""
+    if not hasattr(S, 'ndim'):
+        S = np.asarray(S)
+    if S.ndim != 2:
+        raise ValueError(f'{name} must be 2-D [pulses, samples], got shape {tuple(S.shape)}')
+    if S.dtype.kind != 'c':
+        raise TypeError(f'{name} must be complex (complex64 or complex128), got {S.dtype}')
+    if P is not None and S.shape[0] != P:
+        raise ValueError(f'{name} has {S.shape[0]} pulses but there are {P} antenna positions')
+    if S.shape[0] < 2 or S.shape[1] < 2:
+        raise ValueError(f'{name} needs at least 2 pulses and 2 samples, got shape {tuple(S.shape)}')
+    # a sum per block of rows: one pass, no full-size temporary
+    for i in range(0, S.shape[0] if finite else 0, 4096):
+        if not np.isfinite(complex(S[i:i + 4096].sum())):
+            first = i + int(np.argmax((~np.isfinite(S[i:i + 4096])).any(1)))
+            raise ValueError(f'{name} has non-finite samples (NaN or inf), the first in pulse {first}; '
+                             'zero them (S[~np.isfinite(S)] = 0) or drop the pulses')
+    return S
+
+
+def _check_positions(a, name='antenna positions', P=None):
+    """a [P, 3] finite float64, at least 2 rows (P of them when given)."""
+    a = np.asarray(a, np.float64)
+    if a.ndim != 2 or a.shape[1] != 3:
+        raise ValueError(f'{name} must be [pulses, 3], got shape {a.shape}')
+    if P is not None and len(a) != P:
+        raise ValueError(f'{name}: {len(a)} rows for {P} pulses')
+    if len(a) < 2:
+        raise ValueError(f'{name}: at least 2 pulses are needed, got {len(a)}')
+    if not np.isfinite(a).all():
+        raise ValueError(f'{name}: non-finite values (NaN or inf)')
+    return a
+
+
+def _check_grid(nx, ny, spx, spy, e1, e2):
+    """Pixel counts (positive integers), spacings (positive) and the axes (orthonormal 3-vectors) -> nx, ny (int),
+    e1, e2 (float64)."""
+    for n, v in (('nx', nx), ('ny', ny)):
+        if not np.isscalar(v) or not np.isfinite(v) or v < 1 or int(v) != v:
+            raise ValueError(f'{n} must be a positive integer, got {v!r}')
+    for n, v in (('spx', spx), ('spy', spy)):
+        if not np.isfinite(v) or v <= 0:
+            raise ValueError(f'{n} must be a positive spacing in metres, got {v!r}')
+    e1, e2 = np.asarray(e1, np.float64), np.asarray(e2, np.float64)
+    if e1.shape != (3,) or e2.shape != (3,):
+        raise ValueError(f'e1 and e2 must be 3-vectors, got shapes {e1.shape} and {e2.shape}')
+    if abs(e1 @ e1 - 1) > 1e-6 or abs(e2 @ e2 - 1) > 1e-6 or abs(e1 @ e2) > 1e-6:
+        raise ValueError(f'e1 and e2 must be orthonormal (unit length, perpendicular), got {e1.tolist()} and {e2.tolist()}')
+    return int(nx), int(ny), e1, e2
+
+
+def _jax_group(plan, K):
+    """First-level children per group on a TPU: memory.TPU_GROUP (groups of 4 are as fast as 8; docs/performance.md),
+    halved until memory.full_speed's model (the padded history planes and XLA's temporaries per child) fits in 95% of
+    the device's memory limit. The limit, not the memory free at this moment, so that patches of a mosaic with equal
+    plans get equal groups (and share a program) however many histories the workers have staged. FASTSAR_TPU_GROUP
+    overrides it; ImageFormer halves it again if the device still runs out of memory."""
+    forced = _mem._env_number('FASTSAR_TPU_GROUP', 1)
+    if forced:
+        return forced
+    hbm = _mem.tpu_capacity()
+    hist, per = _mem.tpu_history_bytes(plan, K) + _mem.TPU_FIXED, _mem.child_bytes(plan, 'tpu')
+    ng = _mem.TPU_GROUP
+    while ng > 1 and hist + ng * per > 0.95 * hbm:
+        ng //= 2
+    return ng
+
+
+def final_weights(plan, weight, P, grad=False, points=None):
+    """Weight of each final subaperture at each final tile [ntiles, Pf] (float32), for ImageFormer's aperture_weight:
+    weight(points [m, 3], pulses [n] int) -> [n, m], the per-pulse weights at the tile centers, averaged over the
+    pulses that each final subaperture spans (its center on the input pulse axis composed through the levels'
+    decimations, width the product of the decimation factors) by 3-point Gauss-Legendre quadrature. grad=True
+    also returns the gradient along e1 and e2 (per metre) at the tile centers: [3, ntiles, Pf]."""
+    cen = np.asarray(plan['final']['cen'], np.float64)
+    if grad:
+        h1, h2 = 0.5 * plan['T'] * plan['spx'], 0.5 * plan['T'] * plan['spy']
+        e1, e2 = np.asarray(plan['e1'], np.float64), np.asarray(plan['e2'], np.float64)
+        w = [final_weights(plan, weight, P, points=q) for q in (cen, cen + h1 * e1, cen - h1 * e1, cen + h2 * e2, cen - h2 * e2)]
+        return np.stack([w[0], (w[1] - w[2]) / (2 * h1), (w[3] - w[4]) / (2 * h2)]).astype(np.float32)
+    q = cen if points is None else points
+    idx, D = np.arange(plan['final']['P'], dtype=np.float64), 1.0
+    for lv in reversed(plan['levels']):
+        idx = lv['Dp'] * idx + float(lv['pidx'][0])
+        D *= lv['Dp']
+    # pulse p covers [p - 1/2, p + 1/2); a subaperture spans [idx - D/2, idx + D/2] within the collection
+    a = np.clip(idx - D / 2, -0.5, P - 0.5); b = np.clip(idx + D / 2, -0.5, P - 0.5)
+    t, gw = np.array([-np.sqrt(0.6), 0.0, np.sqrt(0.6)]), np.array([5.0, 8.0, 5.0]) / 18.0
+    nodes = np.clip(np.rint(0.5 * (a + b)[:, None] + 0.5 * (b - a)[:, None] * t[None, :]), 0, P - 1).astype(np.int64)
+    U, inv = np.unique(nodes.ravel(), return_inverse=True)
+    W = np.asarray(weight(q, U), np.float64)[inv].reshape(len(idx), 3, -1)    # [Pf, 3, ntiles]
+    # a subaperture centred beyond the collection holds the filter tails of the edge pulses (the decimators run m
+    # outputs past each end): it takes the weight of the nearest pulse (its nodes clip to it), not zero
+    m = np.tensordot(gw, W, axes=(0, 1))                                          # [Pf, ntiles]
+    return np.ascontiguousarray(m.T, np.float32)
+
+
+def _jax_runtime_error():
+    """The exception type JAX raises for a failed computation (JaxRuntimeError; XlaRuntimeError in older releases)."""
+    import jax
+    err = getattr(jax.errors, 'JaxRuntimeError', None)
+    if err is None:
+        from jaxlib.xla_extension import XlaRuntimeError as err
+    return err
+
+
+class _Staged:
+    """A phase history staged for a JAX or TPU former (ImageFormer.stage): the host array, its device planes and
+    scale, and the program they were padded for."""
+    __slots__ = ('S', 'hre', 'him', 'scale', 'fn')
+
+    def __init__(self, S, hre, him, scale, fn):
+        self.S, self.hre, self.him, self.scale, self.fn = S, hre, him, scale, fn
+
+
+class ImageFormer:
+    """Factorized backprojection set up once for a collection geometry and output grid, then called on phase
+    histories:  former = ImageFormer(ant, fmin, df, K, nx, ny, spx, spy, e1, e2); img = former(S).
+
+    Building plans the tiles and filters, computes the float64 geometry and compiles the kernels; each call then
+    pays only the image formation. Reuse one former for repeated images of the same geometry (or for timing);
+    a different antenna path needs a new former. Arguments as for form_image (those after e2 by keyword only), and
+    aperture_weight(points [m, 3], pulses [n]) -> W [n, m]: a per-pixel weight of pulses (int indices) (a stripmap aperture window), applied
+    in the final stage as the mean weight of each final subaperture's pulses at each final tile's center, with its
+    first-order variation across the tile (final_weights), and
+    ref [P]: the range (one way) each pulse is referenced to, when it is not |ant| (the distance to the origin):
+    a bistatic collection's half path |tx| / 2 + |rcv| / 2, or a reference point other than the origin."""
+
+    def __init__(self, ant, fmin, df, K, nx, ny, spx, spy, e1=(1.0, 0.0, 0.0), e2=(0.0, 1.0, 0.0), *, backend='auto',
+                 precision='float32', window=True, T='auto', levels=3, pmax=0.4, target_db=-40.0, aperture_weight=None,
+                 ref=None):
+        from . import ffbp2
+        import warnings
+        self.ant = _check_positions(ant)
+        self.P, self.K = self.ant.shape[0], int(K)
+        if self.K < 2:
+            raise ValueError(f'K must be at least 2 frequency samples, got {K}')
+        if not (np.isfinite(fmin) and np.isfinite(df) and fmin > 0 and df > 0):
+            raise ValueError(f'fmin and df must be positive frequencies in Hz, got {fmin!r} and {df!r}')
+        self.window = window
+        col = Collect(fmin=float(fmin), df=float(df), K=self.K, ant=self.ant, res=0.5)
+        nx, ny, e1, e2 = _check_grid(nx, ny, spx, spy, e1, e2)
+        self.backend = _backend(backend)
+        self.precision = precision
+        self.predicted_error_db = None
+        if T == 'auto':
+            fmax = float(fmin) + self.K * float(df)
+            if self.backend == 'cuda' and precision == 'float16':
+                T, err = 32, ffbp2.final_phase_error(self.ant, fmax, nx, ny, spx, spy, e1, e2, 32)
+                if err > target_db:
+                    warnings.warn(f'cuda float16 needs T=32, predicted error {err:.1f} dB misses target {target_db:.1f} dB; '
+                                  "use precision='float32' for this geometry")
+            else:
+                T, err = ffbp2.choose_T(self.ant, fmax, nx, ny, spx, spy, e1, e2, target_db)
+                if err > target_db:
+                    warnings.warn(f'predicted error {err:.1f} dB at T={T} misses target {target_db:.1f} dB '
+                                  '(short range for this pixel size)')
+            self.predicted_error_db = float(err)
+        self.T = T
+        plan = ffbp2.make_plan(col, nx, ny, spx, spy, T=T, nlev=levels, pmax=pmax, e1=e1, e2=e2)
+        self._plan = plan
+        if ref is not None and np.shape(ref) != (self.P,):
+            raise ValueError(f'ref must hold one range per pulse ({self.P}), got shape {np.shape(ref)}')
+        coll = ffbp2.collection_arrays(plan, self.ant, ref)
+        wf = None if aperture_weight is None else final_weights(plan, aperture_weight, self.P, grad=os.environ.get('FASTSAR_WEIGHT_GRAD', '1') == '1')
+        if window:
+            self.wp, self.wk = _window(self.P, self.K)
+        if self.backend == 'cuda':
+            from . import ffbp_cuda
+            if precision not in ('float32', 'float16'):
+                raise ValueError("cuda precision: 'float32' or 'float16'")
+            if precision == 'float16' and T != 32:
+                raise ValueError('cuda float16 uses the tensor-core final stage, which is built for T=32')
+            self._form = ffbp_cuda.make_ffbp_cuda(plan, coll, final_mode='f16tc' if precision == 'float16' else 'fp32',
+                                                  store='f16' if precision == 'float16' else 'fp32', wf=wf)
+        elif self.backend == 'cpu':
+            from . import ffbp_cpu
+            if T not in (16, 32):
+                raise ValueError(f'the cpu backend supports T=16 or T=32, not {T}')
+            if precision != 'float32':
+                raise ValueError("cpu precision: 'float32'")
+            self._form = ffbp_cpu.make_ffbp_cpu(plan, coll, wf=wf)
+        elif self.backend in ('tpu', 'jax'):
+            pol = {'float32': 'fp32_high' if self.backend == 'tpu' else 'fp32', 'three-pass': 'fp32_high',
+                   'single-pass': 'fp32_fast', 'float16': 'f16'}.get(precision)
+            if pol is None:
+                raise ValueError(f'unknown precision {precision!r}')
+            filt = 'pallas2' if self.backend == 'tpu' else 'conv'
+            self._pol = pol
+            # one compiled program per plan signature: patches of a mosaic with equal shapes and filters share it
+            # (with the device arrays that depend on the plan alone: filters, tile geometry)
+            self._filt = filt
+            ng = _jax_group(plan, self.K) if filt == 'pallas2' else 1
+            if filt == 'pallas2' and ng < min(_mem.TPU_GROUP, plan['levels'][0]['C']) and not os.environ.get('FASTSAR_TPU_GROUP'):
+                _mem._warn(f'tpu: first-level groups of {_mem._nchildren(ng)} instead of {_mem.TPU_GROUP} for lack of device '
+                           f'memory, which is slower; full speed needs about {_mem._gb(_mem.full_speed(plan, self.K, "tpu")[0])} '
+                           f'of TPU memory, the device has {_mem._gb(_mem.tpu_capacity())}')
+            self._fn, static = self._program(ng)
+            self._arrs = ffbp2.device_arrays(pol, plan, coll, static)
+            if wf is not None:
+                import jax.numpy as jnp
+                self._arrs['final']['w'] = jnp.asarray(wf if wf.ndim == 3 else np.stack([wf, 0 * wf, 0 * wf]))
+        else:
+            raise ValueError(f'unknown backend {backend!r}')
+
+    def stage(self, S):
+        """The work of a call that precedes formation, done ahead of it: on the JAX and TPU backends the checks of S,
+        its scaling to float32 planes and their upload to the device; former(former.stage(S)) equals former(S).
+        A mosaic stages the next patches on its worker threads while one forms. Other backends return S."""
+        if self.backend not in ('jax', 'tpu'):
+            return S
+        from . import ffbp2
+        S = self._host(S)
+        hre, him, scale = ffbp2.prepare(self._pol, S)
+        hre, him = self._fn.pad(hre, him)
+        return _Staged(S, hre, him, scale, self._fn)
+
+    def _host(self, S):
+        if not (self.backend == 'cuda' and type(S).__module__.startswith('cupy')):
+            S = np.asarray(S)
+        if S.shape != (self.P, self.K):
+            raise ValueError(f'phase history must be {(self.P, self.K)}, got {S.shape}')
+        _check_history(S)
+        if self.window:
+            wp, wk = self.wp, self.wk
+            if not isinstance(S, np.ndarray):
+                import cupy as cp
+                wp, wk = cp.asarray(wp), cp.asarray(wk)
+            S = (S * wp[:, None] * wk[None, :]).astype(np.complex64)
+        else:
+            S = S.astype(np.complex64, copy=False)
+        if not S.flags.c_contiguous:              # a transposed or strided view (the kernels read rows in place)
+            S = S.copy()
+        return S
+
+    def __call__(self, S):
+        """S [P, K] complex phase history (numpy, or cupy on the cuda backend), or former.stage(S) -> complex64
+        image [nx, ny]."""
+        staged = S if isinstance(S, _Staged) else None
+        if self.backend == 'cuda' and staged is None and not type(S).__module__.startswith('cupy'):
+            # a host history is windowed and checked on the device as its row blocks are uploaded (two passes over
+            # 1.75 GB on the host cost a 4-vCPU instance more than the whole formation); one that streams is prepared
+            # on the host
+            S = np.asarray(S)
+            if S.shape != (self.P, self.K):
+                raise ValueError(f'phase history must be {(self.P, self.K)}, got {S.shape}')
+            _check_history(S, finite=False)
+            import cupy as cp
+            return cp.asnumpy(self._form(S, window=(self.wp, self.wk) if self.window else None, check=True)).astype(np.complex64, copy=False)
+        S = staged.S if staged is not None else self._host(S)
+        if self.backend == 'cuda':
+            import cupy as cp
+            return cp.asnumpy(self._form(S)).astype(np.complex64, copy=False)       # a host S may stream (ffbp_cuda)
+        if self.backend == 'cpu':
+            return self._form(S).astype(np.complex64, copy=False)
+        import jax
+        from . import ffbp2
+        runtime_error = _jax_runtime_error()
+        while True:
+            if staged is not None and staged.fn is self._fn:
+                hre, him, scale = staged.hre, staged.him, staged.scale
+                staged = None
+            else:
+                staged = None
+                hre, him, scale = ffbp2.prepare(self._pol, S)
+                hre, him = self._fn.pad(hre, him)
+            try:
+                # dispatch is asynchronous: a device that runs out of memory may say so only when the result is read
+                re, im = jax.block_until_ready(self._fn(hre, him, self._arrs))
+                break
+            except runtime_error as e:      # device memory: fewer first-level children per group, down to one
+                # a kernel's on-chip scratch (VMEM) is sized at compile time and does not depend on the group size
+                if 'RESOURCE_EXHAUSTED' not in str(e) or 'vmem' in str(e).lower():
+                    raise
+                ng = self._fn.stages['ng']          # the group the program runs (a divisor of the children)
+                if ng <= 1:
+                    need = _mem.tpu_history_bytes(self._plan, self.K) + _mem.child_bytes(self._plan, 'tpu')
+                    raise MemoryError(f'{self.backend}: the phase history ({_mem._gb(8.0 * self.P * self.K)} as float32 '
+                                      f'planes) and one first-level child do not fit in device memory; this needs about '
+                                      f'{_mem._gb(need)}. Use the cuda or cpu backend (both stream or read the history in '
+                                      'place), a device with more memory, or fewer pulses.') from e
+                del hre, him
+                self._fn = self._program(ng // 2)[0]
+                _mem._warn(f'{self.backend}: out of device memory; retrying with first-level groups of '
+                           f'{_mem._nchildren(self._fn.stages["ng"])} (slower); full speed needs about '
+                           f'{_mem._gb(_mem.full_speed(self._plan, self.K, "tpu")[0])}')
+        return ((np.asarray(re) + 1j * np.asarray(im)) * scale).astype(np.complex64)
+
+    def memory(self):
+        """What full speed needs on this former's device and what is free: dict(backend, needed, available,
+        full_speed, parts) in bytes (fastsar.memory)."""
+        isz = 4 if self.backend == 'cuda' and self.precision == 'float16' else 8
+        need, parts = _mem.full_speed(self._plan, self.K, self.backend, isz)
+        avail = _mem.device_available(self.backend)
+        return dict(backend=self.backend, needed=need, available=avail, full_speed=need <= avail, parts=parts)
+
+    def _program(self, ng):
+        """The compiled JAX/TPU program for ng first-level children per group, with the plan's device arrays: one per
+        plan signature, so that patches of a mosaic with equal shapes and filters share it."""
+        from . import ffbp2
+        key = (self._pol, self._filt, ng, ffbp2.plan_signature(self._plan))
+        with _JAX_LOCK:
+            hit = _JAX_PROGRAMS.get(key)
+            if hit is None:
+                hit = (ffbp2.make_ffbp(self._pol, self._plan, self._filt, 1 << 26, 'direct', pallas_pb=256, pallas_nc=8,
+                                       pallas_ng=ng, pallas_final=2, pallas_gen=3),
+                       ffbp2.static_arrays(self._pol, self._plan, self._filt))
+                if len(_JAX_PROGRAMS) >= 32:
+                    _JAX_PROGRAMS.pop(next(iter(_JAX_PROGRAMS)))
+                _JAX_PROGRAMS[key] = hit
+        return hit
+
+
+def form_image(S, ant, fmin, df, nx, ny, spx, spy, e1=(1.0, 0.0, 0.0), e2=(0.0, 1.0, 0.0), algorithm='ffbp',
+               backend='auto', precision='float32', window=True, T='auto', levels=3, pmax=0.4, pfa_guard=300.0, target_db=-40.0,
+               ref=None, interp=None, upsample=None, center=None):
+    """Form the complex image [nx, ny] (complex64). See the module docstring for the arguments. precision, T, levels,
+    pmax and target_db apply to factorized backprojection, pfa_guard to polar format, and interp ('cubic' or
+    'linear'), upsample and center to exact backprojection (algorithm='bp'; see ExactFormer). For more than one
+    image of the same geometry, build an ImageFormer (or ExactFormer) once and call it; this function sets one up
+    on every call."""
+    if algorithm not in ('ffbp', 'pfa', 'bp'):
+        raise ValueError("algorithm must be 'ffbp', 'pfa' or 'bp'")
+    if algorithm != 'bp':
+        for n, v in (('interp', interp), ('upsample', upsample), ('center', center)):
+            if v is not None:
+                raise ValueError(f"{n} applies to exact backprojection (algorithm='bp'), not to algorithm={algorithm!r}")
+    # the formers scan S for NaN and inf themselves; polar format is checked here
+    S = _check_history(np.asarray(S), len(_check_positions(ant)), finite=algorithm == 'pfa')
+    if algorithm == 'bp':
+        from .exact import ExactFormer
+        return ExactFormer(ant, fmin, df, S.shape[1], nx, ny, spx, spy, e1, e2, backend=backend, window=window,
+                           interp='cubic' if interp is None else interp, upsample=upsample, ref=ref, center=center)(S)
+    if algorithm == 'pfa':
+        nx, ny, e1, e2 = _check_grid(nx, ny, spx, spy, e1, e2)
+        P, K = S.shape
+        col = Collect(fmin=float(fmin), df=float(df), K=K, ant=np.asarray(ant, np.float64), res=0.5)
+        if ref is not None:               # polar format needs the samples referenced to |ant|: move them there
+            from .io import rereference
+            if np.shape(ref) != (P,):
+                raise ValueError(f'ref must hold one range per pulse ({P}), got shape {np.shape(ref)}')
+            S = rereference(S, float(fmin), float(df), np.asarray(ref, np.float64) - np.linalg.norm(col.ant, axis=1))
+        return _pfa(S, col, nx, ny, spx, spy, np.asarray(e1, np.float64), np.asarray(e2, np.float64), pfa_guard, window)
+    return ImageFormer(ant, fmin, df, S.shape[1], nx, ny, spx, spy, e1, e2, backend=backend, precision=precision,
+                       window=window, T=T, levels=levels, pmax=pmax, target_db=target_db, ref=ref)(S)
+
+
+_PFA_PROGRAMS = {}
+_PFA_LOCK = threading.Lock()
+
+
+def _pfa(S, col, nx, ny, spx, spy, e1, e2, guard=300.0, window=True):
+    """Polar format with its final resampling (removes the planar-wavefront displacement), pulse resampling in
+    gather form. The frequency samples are weighted by f_c / f_k so that polar format applies the same spectral
+    weighting as backprojection (the polar Jacobian). The compiled program and its geometry arrays are kept per
+    collection geometry (the last four), so repeated images of one geometry compile once; the window, the weighting
+    and the scaling run on the device."""
+    import hashlib, sys
+    from . import pfa2
+    import jax
+    import jax.numpy as jnp
+    P, K = S.shape
+    key = (P, K, float(col.fmin), float(col.df), nx, ny, float(spx), float(spy), tuple(e1), tuple(e2), float(guard), bool(window),
+           hashlib.sha1(np.ascontiguousarray(col.ant).tobytes()).hexdigest())
+    with _PFA_LOCK:
+        prog = _PFA_PROGRAMS.get(key)
+    if prog is None:
+        geo = pfa2.geometry(col, nx, ny, spx, spy, e1=e1, e2=e2, guard=guard)
+        dist = pfa2.distortion(col, nx, ny, spx, spy, e1, e2)
+        fn = pfa2.make_pfa(geo, nx, ny, spx, spy, 'taps', None, jax.lax.Precision.HIGHEST, dist=dist)
+        arrs = pfa2.arrays(geo, P, 'taps')
+        f_k = col.fmin + np.arange(K) * col.df
+        wp, wk = _window(P, K) if window else (np.ones(P), np.ones(K))
+        wkf = jnp.asarray((wk * (col.fmin + (K // 2) * col.df) / f_k).astype(np.float32))
+        wpd = jnp.asarray(np.asarray(wp, np.float32))
+
+        @jax.jit
+        def prep(Sd):
+            x = Sd * wpd[:, None] * wkf[None, :]
+            scale = jnp.max(jnp.abs(x))
+            scale = jnp.where(scale > 0, scale, 1.0)          # an all-zero history gives a zero image, not 0/0
+            return jnp.real(x) / scale, jnp.imag(x) / scale, scale
+        prog = (fn, arrs, prep)
+        with _PFA_LOCK:
+            if len(_PFA_PROGRAMS) >= 4:
+                _PFA_PROGRAMS.pop(next(iter(_PFA_PROGRAMS)))
+            _PFA_PROGRAMS[key] = prog
+    fn, (W, alpha, eps_r, shift, eps_a), prep = prog
+    if 'cupy' in sys.modules and jax.default_backend() == 'gpu':
+        sys.modules['cupy'].get_default_memory_pool().free_all_blocks()     # blocks CuPy cached are not JAX's to use
+    re, im, scale = prep(jax.device_put(np.asarray(S, np.complex64)))
+    out = fn(re, im, W, alpha, eps_r, shift, eps_a)
+    return (np.asarray(out) * float(scale)).astype(np.complex64)