fast-acbf 0.7.0__py3-none-any.whl

fast_acbf/__init__.py

@@ -0,0 +1,24 @@
+from .solver import BFSolver
+from .core.aberrations import AberrationState
+from .optimization.metrics import QualityMetrics
+from .data.dataset4d import Dataset4D
+from .data.bf_extractor import BFExtractor
+from .data.geometry import DetectorGeometry, ScanGeometry, CoordinateTransform
+from .data.imagefft import ImageFFT
+from .recon.pipeline import PipelineManager, PipelineResolution
+
+__version__ = "0.7.0" # 2026.07.30
+
+__all__ = [
+    "BFSolver",
+    "AberrationState",
+    "QualityMetrics",
+    "Dataset4D",
+    "BFExtractor",
+    "DetectorGeometry",
+    "ScanGeometry",
+    "CoordinateTransform",
+    "ImageFFT",
+    "PipelineManager",
+    "PipelineResolution",
+]

fast_acbf/core/aberrations.py

@@ -0,0 +1,165 @@
+"""AberrationState — nn.Module holding detector-frame aberration coefficients."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class AberrationState(torch.nn.Module):
+    def __init__(self, ab_dict: dict, max_order: int = None, device='cpu', tolerance_factors: dict = None):
+        super().__init__()
+        self.orig_ab_dict = ab_dict.copy()
+        self.device = torch.device(device)
+        self.order_keys = []
+        self.tolerance_factors = tolerance_factors
+        # tolerance_factors: dict {n: Tn_in_ang} — Kirkland tolerance for each order.
+        # https://doi.org/10.1016/j.ultramic.2017.12.002
+        # When set, parameters are stored in normalized (dimensionless) form internally;
+        # all external accessors (get_cartesian_dict, get_flat_coeffs, get_physical) return physical Å values.
+
+        if max_order is None:
+            max_order = max([n for (n, m) in ab_dict.keys()]) if ab_dict else 2
+
+        self.max_order = max_order
+
+        self.coeffs = nn.ParameterDict()
+
+        for n in range(1, max_order + 1):
+            start_m = (n + 1) % 2
+            Tn = tolerance_factors[n] if tolerance_factors is not None else 1.0
+
+            for m in range(start_m, n + 2, 2):
+                self.order_keys.append((n, m))
+
+                if m == 0:
+                    val = ab_dict.get((n, m), 0.0)
+                    val_clean = float(val) if isinstance(val, (int, float, np.number)) else 0.0
+
+                    key = f"C_{n}_{m}"
+                    self.coeffs[key] = nn.Parameter(torch.tensor(val_clean / Tn, device=self.device))
+
+                else:
+                    val = ab_dict.get((n, m), {'a': 0.0, 'b': 0.0})
+                    if isinstance(val, (int, float, np.number)):
+                        val_a, val_b = float(val), 0.0
+                    else:
+                        val_a = float(val.get('a', 0.0))
+                        val_b = float(val.get('b', 0.0))
+
+                    key_a = f"C_{n}_{m}_a"
+                    key_b = f"C_{n}_{m}_b"
+                    self.coeffs[key_a] = nn.Parameter(torch.tensor(val_a / Tn, device=self.device))
+                    self.coeffs[key_b] = nn.Parameter(torch.tensor(val_b / Tn, device=self.device))
+
+    def get_cartesian_dict(self):
+        """Reconstructs the nested dictionary format in physical Å units."""
+        out = {}
+        for (n, m) in self.order_keys:
+            Tn = self.tolerance_factors[n] if self.tolerance_factors is not None else 1.0
+            if m == 0:
+                out[(n, m)] = self.coeffs[f"C_{n}_{m}"].item() * Tn
+            else:
+                out[(n, m)] = {
+                    'a': self.coeffs[f"C_{n}_{m}_a"].item() * Tn,
+                    'b': self.coeffs[f"C_{n}_{m}_b"].item() * Tn
+                }
+        return out
+
+    def get_flat_coeffs(self):
+        """Returns all coefficients as a 1D tensor in physical Å units."""
+        coeffs_list = []
+        tn_list = []
+        for (n, m) in self.order_keys:
+            Tn = self.tolerance_factors[n] if self.tolerance_factors is not None else 1.0
+            if m == 0:
+                coeffs_list.append(self.coeffs[f"C_{n}_{m}"])
+                tn_list.append(Tn)
+            else:
+                coeffs_list.append(self.coeffs[f"C_{n}_{m}_a"])
+                tn_list.append(Tn)
+                coeffs_list.append(self.coeffs[f"C_{n}_{m}_b"])
+                tn_list.append(Tn)
+        stacked = torch.stack(coeffs_list)
+        if self.tolerance_factors is not None:
+            tn_tensor = stacked.new_tensor(tn_list)
+            stacked = stacked * tn_tensor
+        return stacked
+
+    def to_scan_frame(self, rotation_deg: float) -> torch.Tensor:
+        """
+        Return flat physical coefficients rotated from detector frame to scan frame.
+
+        Differentiable with respect to the stored detector-frame parameters.
+        Symmetric terms (m == 0) are rotation-invariant. Each asymmetric Cartesian
+        pair (C_a, C_b) is rotated by m * rotation_deg.
+        """
+        if rotation_deg == 0.0:
+            return self.get_flat_coeffs()
+
+        flat = self.get_flat_coeffs()
+        out = flat.clone()
+        idx = 0
+        for _, m in self.order_keys:
+            if m == 0:
+                idx += 1
+                continue
+
+            theta = flat.new_tensor(m * rotation_deg * np.pi / 180.0)
+            c = torch.cos(theta)
+            s = torch.sin(theta)
+            ca = flat[idx]
+            cb = flat[idx + 1]
+            out[idx] = ca * c - cb * s
+            out[idx + 1] = ca * s + cb * c
+            idx += 2
+        return out
+
+    def get_physical(self, key: str) -> float:
+        """Return the physical Å value for a single coefficient key."""
+        n = int(key.split('_')[1])
+        raw = self.coeffs[key].detach().item()
+        return raw * self.tolerance_factors[n] if self.tolerance_factors is not None else raw
+
+    def set_physical(self, key: str, phys_value):
+        """Write a physical Å value to a coefficient, normalizing internally."""
+        n = int(key.split('_')[1])
+        val = phys_value.item() if isinstance(phys_value, torch.Tensor) else float(phys_value)
+        norm_val = val / self.tolerance_factors[n] if self.tolerance_factors is not None else val
+        with torch.no_grad():
+            self.coeffs[key].copy_(torch.tensor(norm_val, dtype=torch.float32, device=self.device))
+
+    def flat_to_cartesian_dict(self, flat: torch.Tensor) -> dict:
+        """Convert a flat physical-unit coefficient tensor to the cartesian dict format."""
+        out = {}
+        idx = 0
+        for (n, m) in self.order_keys:
+            if m == 0:
+                out[(n, m)] = flat[idx].item()
+                idx += 1
+            else:
+                out[(n, m)] = {'a': flat[idx].item(), 'b': flat[idx + 1].item()}
+                idx += 2
+        return out
+
+    def rebuild_normalization(self, new_tolerance_factors: dict):
+        """Re-normalize all parameters under a new set of tolerance factors.
+
+        Reads current physical values, updates self.tolerance_factors, then stores
+        parameters normalized by the new Tₙ values.  Use this after overriding
+        individual entries in tolerance_factors to keep physical ↔ internal mapping
+        consistent.
+        """
+        phys = self.get_cartesian_dict()
+        self.tolerance_factors = new_tolerance_factors
+        with torch.no_grad():
+            for (n, m), val in phys.items():
+                Tn = new_tolerance_factors[n]
+                if m == 0:
+                    key = f"C_{n}_{m}"
+                    self.coeffs[key].copy_(torch.tensor(val / Tn, dtype=torch.float32, device=self.device))
+                else:
+                    for suffix, v in [('_a', val['a']), ('_b', val['b'])]:
+                        key = f"C_{n}_{m}{suffix}"
+                        self.coeffs[key].copy_(torch.tensor(v / Tn, dtype=torch.float32, device=self.device))

fast_acbf/core/acbf.py

@@ -0,0 +1,226 @@
+"""acBF reconstruction — pure math, no state."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from fast_acbf.core.functional import generate_aberration_basis, make_soft_aperture_torch
+from fast_acbf.recon.cache import ACBFGeometryCache, ACBFOpticsCache
+
+if TYPE_CHECKING:
+    from fast_acbf.data.imagefft import ImageFFT
+
+
+def compute_transfer(
+    geom_chunk: dict,
+    optics_chunk: dict | None,
+    coeffs: torch.Tensor,
+    qx_grid: torch.Tensor,
+    qy_grid: torch.Tensor,
+    geometry: ACBFGeometryCache,
+    device: str,
+) -> torch.Tensor:
+    """
+    Compute the detector-wise complex transfer for one chunk.
+
+    Two paths driven by the presence of an optics chunk:
+        Fast path: bases and apertures pre-cached on optics_chunk.
+        Lazy path: optics_chunk is None — recompute from geom_chunk + scan grids + physics.
+
+    Args:
+        geom_chunk:   One ACBFGeometryCache chunk dict ({kxt, kyt, start, end}).
+        optics_chunk: Matching ACBFOpticsCache chunk dict, or None for lazy path.
+        coeffs:       Flat scan-frame aberration coefficients, shape (num_coeffs,).
+        qx_grid:      Scan-frame frequency grid, shape (1, 1, Rx).
+        qy_grid:      Scan-frame frequency grid, shape (1, Ry, 1).
+        geometry:     ACBFGeometryCache — physics params used by the lazy path.
+        device:       Target device string.
+
+    Returns:
+        Complex transfer T, shape (chunk_size, Ny, Nx).
+    """
+    j1 = torch.tensor(1.0j, dtype=torch.complex64, device=device)
+
+    if optics_chunk is not None:
+        b_tr  = optics_chunk['b_tr']
+        b_t   = optics_chunk['b_t']
+        b_mt  = optics_chunk['b_mt']
+        ap_t  = optics_chunk['ap_t']
+        ap_mt = optics_chunk['ap_mt']
+    else:
+        kxt = geom_chunk['kxt'].to(device)
+        kyt = geom_chunk['kyt'].to(device)
+        kx_t,  ky_t  = qx_grid + kxt, qy_grid + kyt
+        kx_mt, ky_mt = qx_grid - kxt, qy_grid - kyt
+        ap_t  = make_soft_aperture_torch(
+            torch.sqrt(kx_t**2  + ky_t**2)  * geometry.wavelength,
+            geometry.max_alpha, geometry.rolloff)
+        ap_mt = make_soft_aperture_torch(
+            torch.sqrt(kx_mt**2 + ky_mt**2) * geometry.wavelength,
+            geometry.max_alpha, geometry.rolloff)
+        b_tr = generate_aberration_basis(
+            geometry.max_order, geometry.order_keys, kxt,    kyt,    geometry.wavelength)
+        b_t  = generate_aberration_basis(
+            geometry.max_order, geometry.order_keys, kx_t,   ky_t,   geometry.wavelength)
+        b_mt = generate_aberration_basis(
+            geometry.max_order, geometry.order_keys, -kx_mt, -ky_mt, geometry.wavelength)
+
+    chi_tr_az = torch.einsum('k, kbxy -> bxy', coeffs, b_tr)
+    chi_t     = torch.einsum('k, kbxy -> bxy', coeffs, b_t)
+    chi_mt    = torch.einsum('k, kbxy -> bxy', coeffs, b_mt)
+
+    term_mt = ap_mt * torch.exp(-j1 * (chi_tr_az - chi_mt))
+    term_t  = ap_t  * torch.exp( j1 * (chi_tr_az - chi_t))
+    D = term_mt - term_t
+
+    # T = -i * D. Keeps the transfer aligned with the acBF phasor convention.
+    return (-j1) * D
+
+
+def _iter_chunks(geometry: ACBFGeometryCache, optics: ACBFOpticsCache | None):
+    """Yield (geom_chunk, optics_chunk_or_None) pairs."""
+    if optics is None:
+        for g in geometry.chunks:
+            yield g, None
+        return
+    if len(geometry.chunks) != len(optics.chunks):
+        raise ValueError(
+            f"geometry/optics chunk count mismatch: "
+            f"{len(geometry.chunks)} vs {len(optics.chunks)}."
+        )
+    for g, o in zip(geometry.chunks, optics.chunks):
+        if g['start'] != o['start'] or g['end'] != o['end']:
+            raise ValueError(
+                f"geometry/optics chunk bounds mismatch: "
+                f"({g['start']},{g['end']}) vs ({o['start']},{o['end']})."
+            )
+        yield g, o
+
+
+def reconstruct_acbf(
+    provider: ImageFFT,
+    qx_grid: torch.Tensor,
+    qy_grid: torch.Tensor,
+    geometry: ACBFGeometryCache,
+    optics: ACBFOpticsCache | None,
+    coeffs: torch.Tensor,
+    eps: float,
+    device: str,
+) -> torch.Tensor:
+    """
+    Phase-only acBF reconstruction.
+
+    Aligns detector contributions by their phase before summation.
+
+    Args:
+        provider:  ImageFFT serving (chunk_size, Ry, Rx) complex64 chunks.
+        qx_grid:   Scan-frame frequency grid, shape (1, 1, Rx).
+        qy_grid:   Scan-frame frequency grid, shape (1, Ry, 1).
+        geometry:  ACBFGeometryCache built by pipeline.build_acbf_geometry_cache.
+        optics:    ACBFOpticsCache (full mode) or None (lazy mode).
+        coeffs:    Flat scan-frame aberration coefficients, shape (num_coeffs,).
+        eps:       Small constant for phase normalization stability.
+        device:    Target device string.
+
+    Returns:
+        Reconstructed acBF image, shape (Ry, Rx), float32.
+    """
+    Ry_out = qy_grid.shape[-2]
+    Rx_out = qx_grid.shape[-1]
+    acBF_total = torch.zeros((Ry_out, Rx_out), dtype=torch.float32, device=device)
+    for geom_chunk, optics_chunk in _iter_chunks(geometry, optics):
+        img_fft_chunk = provider.get_chunk(geom_chunk['start'], geom_chunk['end'])
+        transfer = compute_transfer(geom_chunk, optics_chunk, coeffs, qx_grid, qy_grid, geometry, device)
+        phasor = transfer / (transfer.abs() + eps)
+        acBF_total += torch.sum(torch.fft.ifft2(img_fft_chunk * phasor, dim=(-2, -1)).real, dim=0)
+    return acBF_total
+
+
+def reconstruct_acbf_complex_inversion(
+    provider: ImageFFT,
+    qx_grid: torch.Tensor,
+    qy_grid: torch.Tensor,
+    geometry: ACBFGeometryCache,
+    optics: ACBFOpticsCache | None,
+    coeffs: torch.Tensor,
+    device: str,
+    regularization: float = 1e-3,
+    support_threshold: float = 1e-6,
+    return_diagnostics: bool = False,
+):
+    """
+    Complex-inversion acBF reconstruction via regularized transfer inversion.
+
+    Solves a regularized matched-filter inversion of the detector-wise complex transfer:
+
+        V_hat(q) = M(q) / (S(q) + lambda * S_ref)
+
+    where:
+        M(q) = sum_b T_b(q) * I_b(q)
+        S(q) = sum_b |T_b(q)|^2
+
+    Args:
+        provider:           ImageFFT serving (chunk_size, Ry, Rx) complex64 chunks.
+        qx_grid:            Scan-frame frequency grid, shape (1, 1, Rx).
+        qy_grid:            Scan-frame frequency grid, shape (1, Ry, 1).
+        geometry:           ACBFGeometryCache.
+        optics:             ACBFOpticsCache (full mode) or None (lazy mode).
+        coeffs:             Flat scan-frame aberration coefficients.
+        device:             Target device string.
+        regularization:     Non-negative regularization weight lambda.
+        support_threshold:  Fraction of median transfer power below which Fourier
+                            components are zeroed.
+        return_diagnostics: If True, return full diagnostic dict instead of just image.
+
+    Returns:
+        Reconstructed image (Ry, Rx) float32, or dict if return_diagnostics=True.
+    """
+    if regularization < 0:
+        raise ValueError(f"regularization must be non-negative, got {regularization}.")
+    if support_threshold < 0:
+        raise ValueError(f"support_threshold must be non-negative, got {support_threshold}.")
+
+    Ry_out = qy_grid.shape[-2]
+    Rx_out = qx_grid.shape[-1]
+    out_shape = (Ry_out, Rx_out)
+    numerator = torch.zeros(out_shape, dtype=torch.complex64, device=device)
+    transfer_power = torch.zeros(out_shape, dtype=torch.float32, device=device)
+
+    for geom_chunk, optics_chunk in _iter_chunks(geometry, optics):
+        img_fft_chunk = provider.get_chunk(geom_chunk['start'], geom_chunk['end'])
+        transfer = compute_transfer(geom_chunk, optics_chunk, coeffs, qx_grid, qy_grid, geometry, device)
+        numerator.add_(torch.sum(transfer * img_fft_chunk, dim=0))
+        transfer_power.add_(torch.sum(transfer.abs().square(), dim=0))
+
+    positive_power = transfer_power[transfer_power > 0]
+    if positive_power.numel() == 0:
+        transfer_reference = torch.tensor(1.0, dtype=torch.float32, device=device)
+    else:
+        transfer_reference = positive_power.median()
+
+    support = transfer_power > (support_threshold * transfer_reference)
+    denom = transfer_power + (regularization * transfer_reference)
+    fourier_estimate = torch.where(
+        support,
+        numerator / denom.to(torch.complex64),
+        torch.zeros_like(numerator),
+    )
+
+    complex_image = torch.fft.ifft2(fourier_estimate, dim=(-2, -1))
+    reconstructed = complex_image.real
+
+    if not return_diagnostics:
+        return reconstructed
+
+    return {
+        'image': reconstructed,
+        'complex_image': complex_image,
+        'real_channel': complex_image.real,
+        'imag_channel': complex_image.imag,
+        'fourier_estimate': fourier_estimate,
+        'transfer_power': transfer_power,
+        'support_mask': support,
+        'support_reference': transfer_reference,
+    }

fast_acbf/core/functional.py

@@ -0,0 +1,190 @@
+"""Differentiable physics primitives — no state, no side effects. Portable to PtyRAD."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from ptyrad.core.functional import fftshift2, ifftshift2, torch_phasor
+
+
+def imshift_with_batch(imgs, shifts, grid, batch_size=None):
+    """
+    Generates a batch of shifted images from an input stack of images.
+
+    This function shifts complex/real-valued input images by applying phase shifts
+    in the Fourier domain, achieving subpixel shifts in both x and y directions.
+    Optionally applies a CTF correction in the same operation.
+
+    Args:
+        imgs (torch.Tensor): The input images to be shifted.
+                             Shape must be (Nb, ..., Ny, Nx).
+        shifts (torch.Tensor): The shifts to be applied. Shape is (Nb, 2),
+                               where each slice is (shift_y, shift_x) in pixels.
+        grid (torch.Tensor): The k-space grid used for computing the shifts.
+                             Shape=(2, Ny, Nx), spanning [-0.5, 0.5).
+        batch_size (int, optional): The number of images to process at once to avoid OOM.
+                                    If None, processes the entire batch simultaneously.
+
+    Returns:
+        torch.Tensor: The shifted (and optionally CTF-corrected) images.
+                      Shape remains (Nb, ..., Ny, Nx).
+    """
+    assert imgs.shape[-2:] == grid.shape[-2:], \
+        f"Incompatible dimensions: imgs {imgs.shape[-2:]} vs grid {grid.shape[-2:]}"
+    assert imgs.shape[0] == shifts.shape[0], \
+        f"Batch size mismatch: imgs has {imgs.shape[0]}, shifts has {shifts.shape[0]}"
+
+    Nb = imgs.shape[0]
+    ndim = imgs.ndim
+
+    batch_size = batch_size or Nb
+
+    shift_expand_shape = [-1] + [1] * (ndim - 1)
+    grid_expand_shape = [1] * (ndim - 2) + list(grid.shape[-2:])
+
+    ky = grid[0].view(*grid_expand_shape)
+    kx = grid[1].view(*grid_expand_shape)
+
+    shifted_imgs_list = []
+
+    for i in range(0, Nb, batch_size):
+        end = min(i + batch_size, Nb)
+
+        imgs_b = imgs[i:end]
+        shifts_b = shifts[i:end]
+
+        shift_y = shifts_b[:, 0].view(*shift_expand_shape)
+        shift_x = shifts_b[:, 1].view(*shift_expand_shape)
+
+        phase = -2 * torch.pi * (shift_x * kx + shift_y * ky)
+        w = torch_phasor(phase)
+
+        img_fft = torch.fft.fft2(imgs_b)
+        shifted_img_b = torch.fft.ifft2(img_fft * w)
+
+        shifted_imgs_list.append(shifted_img_b)
+
+    return torch.cat(shifted_imgs_list, dim=0)
+
+
+def make_soft_aperture_torch(alpha: torch.Tensor, max_alpha_mrad: float, rolloff_mrad: float) -> torch.Tensor:
+    """
+    Creates a differentiable circular aperture mask with a cosine edge taper.
+
+    Args:
+        alpha: Tensor of scattering angles in radians.
+        max_alpha_mrad: The semiangle cutoff in milliradians.
+        rolloff_mrad: The width of the soft edge transition in milliradians.
+    """
+    cutoff = max_alpha_mrad / 1000.0
+
+    if rolloff_mrad <= 0.0:
+        return (alpha <= cutoff).to(torch.float32)
+
+    rolloff = rolloff_mrad / 1000.0
+
+    transition = 0.5 * (1.0 + torch.cos(torch.pi * (alpha - cutoff + rolloff) / rolloff))
+    mask = torch.where(alpha > cutoff, torch.zeros_like(alpha), transition)
+    mask = torch.where(alpha < (cutoff - rolloff), torch.ones_like(alpha), mask)
+
+    return mask
+
+
+def make_probe_from_chi(chi, mask):
+    assert chi.shape[-2:] == mask.shape[-2:]
+
+    psi = torch_phasor(-1 * chi)
+    probe = mask * psi
+    probe = fftshift2(torch.fft.ifft2(ifftshift2(probe)))
+    probe = probe / torch.sqrt(torch.sum((torch.abs(probe)) ** 2))
+    return probe
+
+
+def generate_shift_basis(order_keys: list, kX: torch.Tensor, kY: torch.Tensor, wavelength: float):
+    """Generates the unweighted (C=1) analytic shift basis vectors for dx and dy."""
+    alphaX = kX * wavelength
+    alphaY = kY * wavelength
+    alpha_sq = alphaX**2 + alphaY**2
+
+    alpha_sq_safe = alpha_sq + 1e-12
+
+    max_m = max([m for (n, m) in order_keys]) if order_keys else 0
+
+    X, Y = {}, {}
+    X[0] = torch.ones_like(alphaX)
+    Y[0] = torch.zeros_like(alphaX)
+
+    for m in range(max_m):
+        X[m+1] = X[m] * alphaX - Y[m] * alphaY
+        Y[m+1] = X[m] * alphaY + Y[m] * alphaX
+
+    basis_dx_list = []
+    basis_dy_list = []
+
+    for (n, m) in order_keys:
+        p = (n + 1 - m) / 2.0
+
+        if n + 1 - m > 0:
+            rad_deriv_x = (n + 1 - m) * alphaX * (alpha_sq_safe ** (p - 1.0))
+            rad_deriv_y = (n + 1 - m) * alphaY * (alpha_sq_safe ** (p - 1.0))
+        else:
+            rad_deriv_x = torch.zeros_like(alphaX)
+            rad_deriv_y = torch.zeros_like(alphaX)
+
+        rad_base = alpha_sq ** p
+
+        if m == 0:
+            dx = (rad_deriv_x * X[m]) / (n + 1)
+            dy = (rad_deriv_y * X[m]) / (n + 1)
+            basis_dx_list.append(dx)
+            basis_dy_list.append(dy)
+
+        else:
+            part1_x_a = rad_deriv_x * X[m]
+            part1_y_a = rad_deriv_y * X[m]
+            part2_x_a = rad_base * (m * X[m-1])
+            part2_y_a = rad_base * (m * -Y[m-1])
+
+            basis_dx_list.append((part1_x_a + part2_x_a) / (n + 1))
+            basis_dy_list.append((part1_y_a + part2_y_a) / (n + 1))
+
+            part1_x_b = rad_deriv_x * Y[m]
+            part1_y_b = rad_deriv_y * Y[m]
+            part2_x_b = rad_base * (m * Y[m-1])
+            part2_y_b = rad_base * (m * X[m-1])
+
+            basis_dx_list.append((part1_x_b + part2_x_b) / (n + 1))
+            basis_dy_list.append((part1_y_b + part2_y_b) / (n + 1))
+
+    return torch.stack(basis_dx_list, dim=0), torch.stack(basis_dy_list, dim=0)
+
+
+def generate_aberration_basis(max_order: int, order_keys: list, kX: torch.Tensor, kY: torch.Tensor, wavelength: float):
+    """Generates the unweighted (C=1) Cartesian polynomial basis tensors."""
+    alphaX = kX * wavelength
+    alphaY = kY * wavelength
+    alpha_sq = alphaX**2 + alphaY**2
+
+    X, Y = {}, {}
+    X[0] = torch.ones_like(alpha_sq)
+    Y[0] = torch.zeros_like(alpha_sq)
+
+    for m in range(max_order + 1):
+        X[m+1] = X[m] * alphaX - Y[m] * alphaY
+        Y[m+1] = X[m] * alphaY + Y[m] * alphaX
+
+    multiplier = (2 * torch.pi / wavelength)
+    basis_list = []
+
+    for (n, m) in order_keys:
+        power_rad = (n + 1 - m) / 2.0
+        term_radial = alpha_sq ** power_rad
+
+        if m == 0:
+            basis_list.append(term_radial * X[m] / (n + 1) * multiplier)
+        else:
+            basis_list.append(term_radial * X[m] / (n + 1) * multiplier)  # a
+            basis_list.append(term_radial * Y[m] / (n + 1) * multiplier)  # b
+
+    return torch.stack(basis_list, dim=0)

fast_acbf/core/tcbf.py

@@ -0,0 +1,53 @@
+"""tcBF reconstruction — pure math, no state."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from fast_acbf.recon.cache import TCBFCache
+
+if TYPE_CHECKING:
+    from fast_acbf.data.imagefft import ImageFFT
+
+
+def reconstruct_tcbf(
+    provider: ImageFFT,
+    qx_grid: torch.Tensor,
+    qy_grid: torch.Tensor,
+    cache: TCBFCache,
+    coeffs: torch.Tensor,
+    device: str,
+) -> torch.Tensor:
+    """
+    Ultra-lean AD forward pass for tcBF.
+
+    Calculates exact analytical shifts via Einstein summation over the pre-built
+    shift basis stored in the cache.
+
+    Args:
+        provider:  ImageFFT serving (chunk_size, Ry, Rx) complex64 chunks.
+        qx_grid:   Scan-frame frequency grid, shape (1, 1, Rx).
+        qy_grid:   Scan-frame frequency grid, shape (1, Ry, 1).
+        cache:     TCBFCache built by pipeline.build_tcbf_cache (basis only).
+        coeffs:    Flat scan-frame aberration coefficients, shape (num_coeffs,).
+        device:    Target device string.
+
+    Returns:
+        Reconstructed tcBF image, shape (Ry, Rx), float32.
+    """
+    neg_two_pi_j = torch.tensor(-2.0j * torch.pi, dtype=torch.complex64, device=device)
+    Ry_out = qy_grid.shape[-2]
+    Rx_out = qx_grid.shape[-1]
+
+    tcBF_total = torch.zeros((Ry_out, Rx_out), dtype=torch.float32, device=device)
+    for chunk in cache.chunks:
+        img_fft_chunk = provider.get_chunk(chunk['start'], chunk['end'])
+        shift_dx = torch.einsum('k, kb -> b', coeffs, chunk['b_dx']).view(-1, 1, 1)
+        shift_dy = torch.einsum('k, kb -> b', coeffs, chunk['b_dy']).view(-1, 1, 1)
+        ramp = shift_dx * qx_grid + shift_dy * qy_grid
+        shift_op = torch.exp(neg_two_pi_j * ramp)
+        tcBF_total += torch.sum(torch.fft.ifft2(img_fft_chunk * shift_op, dim=(-2, -1)).real, dim=0)
+
+    return tcBF_total

fast_acbf/data/__init__.py

@@ -0,0 +1,5 @@
+from fast_acbf.data.bf_extractor import BFExtractor
+from fast_acbf.data.bf_preparer import BFPreparer
+from fast_acbf.data.imagefft import ImageFFT
+
+__all__ = ["BFExtractor", "BFPreparer", "ImageFFT"]