Direct ptychography: real-space montage, empirical probes, non-electron units - #281
Open
gvarnavi wants to merge 31 commits into
Open
Direct ptychography: real-space montage, empirical probes, non-electron units#281gvarnavi wants to merge 31 commits into
gvarnavi wants to merge 31 commits into
Conversation
`_serialize_container` recorded `type(value).__name__` as `_container_type`, so a `torch.Size` was written as "Size" and `_deserialize_container` — which only knows "list", "tuple", "dict" and "set" — rejected it on load. This made every `DirectPtychography` unloadable: its `gpts` is `bf_mask.shape[:2]`. Namedtuples were affected the same way. Normalize subclasses to their base container on write instead, so they come back as plain tuples/lists. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…e bugs Splits `DirectPtychography` so a second, real-space reconstruction class can share its machinery, and fixes the bugs that surfaced while doing so. Refactor (behaviour-preserving): - new `DirectPtychographyBase` holds what does not depend on how the deconvolution is performed: the geometry/sampling properties, `HyperparameterState`, bright-field mask indexing, `visualize`, and the optuna / grid-search drivers. Subclasses supply `reconstruct`, `variance_loss` and `corrected_bf`; the contract is documented on the class and declared as annotations. - `build_vbf_stack_from_dataset4d` and friends move to `direct_ptycho_utils`, so any class consuming a virtual bright-field stack runs the same origin-correction, masking and normalization pipeline. - verified by running both trees (`git worktree` at the parent commit) over 5 kernels x 2 upsampling factors plus variance losses, half-sets, lateral shifts, grid search and the least-squares fit: all 35 outputs bit-identical. Fixes: - `visualize`'s scalebar ignored `upsampling_factor`, so an upsampled object reported double its true field of view. `reconstruct` now records `_last_upsampling_factor` and `_obj_sampling` divides by it. - `from_virtual_bfs` defaulted `semiangle_cutoff` to None, which always raised a `TypeError` in the setter. It is now required, with a real error message. - `gpts`/`scan_gpts` are plain tuples rather than `torch.Size`, so they round-trip through `AutoSerialize` as the type they started as. Because save/load never worked, no stored file can name the old module path, so the backwards-compatibility re-exports were dropped. Also adds `scatter_add_splat`: a batched, device-agnostic bilinear/nearest scatter-add, the torch counterpart of `imaging_utils.bilinear_kde`'s accumulation stage. It carries a leading batch axis, offers a drop (`"pad"`) as well as a wrap boundary via clamped indices with zeroed weights (no host sync), and accumulates the sum of squares needed for a variance estimate. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…th classes
Real-space ("shadow montage") direct ptychography: each virtual bright-field image is
translated by its own aberration-dependent lateral shift and accumulated onto a shared
canvas. This is the real-space dual of the existing parallax kernel — the Fourier
multiplier `exp(-1j * grad_chi . q)` is exactly a translation by `grad_chi / (2*pi)`
Angstrom, and Fourier-space tiling by U is exactly zero-insertion at every U-th pixel.
Working in real space buys two things:
- no scan-space FFT, so the scan need not be gridded. `from_dataset3d` takes (N, Qx, Qy)
patterns plus (N, 2) probe positions in Angstrom.
- the phase-flip and Butterworth filters do not depend on the bright-field index, so they
collapse into a single post-hoc filter on the finished image.
Notes:
- `boundary="wrap"` reproduces `DirectPtychography` exactly; `"pad"` grows the canvas to
cover the shifted positions and normalizes by accumulated weight. Defaults are "wrap"
for the gridded constructors, "pad" for `from_dataset3d`.
- `variance_loss` comes from `sum(w)`, `sum(wv)` and `sum(wv^2)` accumulated in the same
splat pass, so optuna and grid search work without materializing an (N_bf, Ny, Nx)
stack. The four ways it differs from the Fourier one are documented on the method.
- only the parallax kernel applies; SSB/OBF/matched-filter have bright-field-dependent
multipliers that are not translations.
- the post-hoc filter is built at the k-grid's native float32 while the accumulators stay
float64: chi(q) reaches tens of radians, so sign(sin(chi)) is ill-conditioned at its
zero crossings and a float64 evaluation flips a few pixels relative to
`DirectPtychography` — a large perturbation per flipped mode.
Tests (105 total, previously none for either class). Shared synthetic data lives in
`conftest.py` and is imported via `from .conftest import ...`, as the tomography suite
does. The headline check is equivalence with `deconvolution_kernel="prlx"`, which is
exact because a contrived C10 puts every bright-field shift on an integer canvas pixel:
for pure defocus the shift is `wavelength * C10 * k` and `k = m * dk` exactly, so
`C10 = p * scan_sampling / (U * wavelength * dk)` gives a shift of exactly `p * m`.
The two agree to 2e-7 relative at realistic scale.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`_last_upsampling_factor` was mutable state that `reconstruct` had to write and `_preprocess` had to reset, in two places per class, purely so `visualize` could report the right scalebar. Anything that forgot to update it reported a silently wrong scale. Store the physical anchor instead: `fov` is fixed by the acquisition (scan_gpts x scan_sampling), and `_obj_sampling` divides it by the reconstructed object's own shape. Upsampling then needs no bookkeeping at all -- the object's shape already encodes it -- and the scalebar cannot fall out of sync with the image it labels. The scan field of view is not the whole story for the montage: with `boundary="pad"` the canvas grows past the scan to cover the shifted positions, so `fov / shape` would under-report the pixel size by the padding fraction (21-26% on the test data). Hence the `_obj_fov` hook, which the montage overrides with the canvas extent. `_return_canvas` now returns that extent alongside the canvas shape, so the two are produced by one expression and cannot disagree about the upsampling factor. `DirectPtychography` reconstruction outputs remain bit-identical to c1192af across all 35 checks; this only affects reported sampling. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
On a raster scan `positions_px * U` is an exact integer, so `round(n + s) == n + round(s)`
and `interpolation="nearest"` is exactly a roll of each virtual bright-field image by
`round(shift)` upsampled pixels -- the quantization error is `1/(2U)` scan pixels and
shrinks as you upsample.
Measured against the exact Fourier shift on the 4 mrad apoferritin dataset, as in-band
power retained (0 < q <= scan Nyquist) at upsampling factor 1 / 2 / 4 / 8:
nearest 0.81 0.95 0.99 1.00
bilinear 0.67 0.90 0.97 0.99
Both converge, but nearest is closer at every factor and avoids the visible softening at
low upsampling, so it becomes the default. The docstring points at bilinear for a blocky
result, and for ungridded scans, where snapping leaves parts of the canvas unvisited
(17-31% empty against 10-15% for bilinear on a jittered scan, widening with upsampling).
Also retargets `weight_normalize`, which corrects uneven sampling density, at the property
that actually drives it: whether the scan is gridded, rather than the boundary mode. On a
raster scan the density is already uniform, so normalizing only rescales the low-weight
edges of a padded canvas and amplifies their noise; leaving it off lets those edges fade
instead. `from_dataset3d` still defaults it on, where density genuinely varies.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Adds `defocus_gradient=(g_row, g_col)` in A/A to ShadowMontagePtychography, so
each scan position is shifted by its own local defocus:
C10(r) = C10_global + g . (r - r_centroid)
This is montage-only by construction. The Fourier parallax operator
`exp(-1j * grad_chi . q)` is a global multiplier over the scan, and a
position-dependent shift is not a Fourier multiplier at all.
It is also nearly free. chi is linear in every aberration magnitude, so
`d shift / d C10 = wavelength * k` exactly and independently of all other
aberrations (verified to float32 eps), making the correction one broadcast add
inside the existing splat loop rather than a re-fit.
Measuring the offset from the position centroid keeps mean(delta C10) = 0, so
the gradient is orthogonal to the global C10 and an optuna or grid search over
C10 stays well posed. Because `defocus_gradient` rides through
**reconstruct_kwargs, both existing optimizers forward it with no plumbing.
Fitting, following the patch-then-plane strategy:
defocus_map() per-patch best-fit C10 over a trial grid, with the
full loss curves returned for inspection
fit_defocus_gradient() least-squares plane through those, storing the gradient
and (by default) the refitted mean defocus
Two details make the patch loss comparable across trial defocus values: the
canvas is frozen rather than sized from the shifts (a growing canvas adds
low-weight edge pixels whose small variance would make the loss fall
monotonically with defocus instead of having a minimum), and only pixels at 90%
of the peak weight are scored.
Testing needed a fixture the 32x32 4D simulation cannot provide -- at that size
patches are ~16 positions and the estimator has a ~430 A bias, larger than the
signal, as the zero-gradient control shows. `make_model_vbf_stack` builds the
parallax relation `v_m(r) = obj[r + shift_m(C10(r))]` directly at any scan size;
the `+` sign was measured against the 4D pipeline (0.30 vs 0.09 correlation).
At 96x96 the fit recovers seeded gradients to within 6% and lifts correlation
with ground truth from 0.89-0.94 to 0.992-0.993.
Also fixes a pre-existing crash in both classes: save() then load() before
reconstruct() raised TypeError, because AutoSerialize evaluates every property
during load and `obj` called to_numpy(None). It now returns None, mirroring
corrected_bf.
All 106 outputs of the parallax/SSB/OBF/MF/iCoM battery are bit-identical to
bab7926 with no gradient set.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`obj = sum_m ifft2(G_m * K_m) = sum_m (v_m conv kappa_m)` is an identity, so the deconvolution kernels can be applied as real-space convolutions. Truncating kappa to a box stencil is then the only approximation, which lets the montage run SSB, OBF and matched-filter on scans the Fourier class cannot take. New `scatter_add_convolve` deposits `value * kappa[tap]` at `coords + offset[tap]` for every tap, looping over taps so peak memory stays O(B * N_pos) whatever the stencil size. It is kept separate from `scatter_add_splat` rather than folded in: the accumulator is complex, and the weight and sum-of-squares buffers behind `variance_loss` have no meaning when the taps are kernel weights rather than a partition of unity. The shared index arithmetic is factored into `_resolve_indices` and `_deposition_corners`. The integer parallax shift is divided out of the kernel by a phase ramp and added back to the deposit coordinates, so the stencil sizes only the residual chirp and aperture ringing rather than the shift itself, which is tens of pixels at realistic defocus. Verified against DirectPtychography on identical data: relative error falls 0.34 -> 0.075 -> 0.0079 for SSB at radius 4 -> 8 -> 15, and similarly for OBF and MF, confirming truncation is the sole approximation. These kernels converge slowly, and the docstrings and warning say so rather than implying otherwise. Dividing by |gamma| leaves a unit-magnitude phase on a hard-edged support, whose transform has r**-1.5 tails, so the error falls only like 1/radius. Being in focus does not help -- it is the slowest case, since there is no chirp to concentrate the kernel. Measured relative operator error at 20 mrad in focus is 0.29 at radius 8 and still 0.15 at radius 32; reaching 10% needs radius 20-50, i.e. 1.7k-10k taps per bright-field pixel. So the warning is quantitative, reporting the estimated error for the chosen radius alongside the worst bright-field pixel, and points at DirectPtychography. On this fixture real-space SSB costs 353 ms against 63 ms for the exact FFT. Two measured choices are pinned by tests. The stencil is a hard box with no taper, because tapering is consistently worse at equal radius (0.40 vs 0.29 at radius 8) -- it discards mid-radius content that matters more than the ringing it suppresses. And the auto-radius keys off the mean rather than worst-case bright-field pixel, since the worst are at the disk centre, where the two shifted apertures nearly coincide and gamma collapses to a thin annulus. iCoM stays rejected: its `k . q / |q|**2` kernel is unbounded as q -> 0. `variance_loss` raises for the convolution kernels instead of returning something meaningless, following the precedent that fit_hyperparameters_cross_correlation forces the parallax kernel. All 106 outputs of the parallax battery remain bit-identical to bab7926. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…scans Every kernel in DirectPtychography is a multiplier on the scan-space Fourier transform, so it needs a regular grid. `from_dataset3d` resamples the bright-field stack onto one first -- splatting each image at its probe positions with no aberration shift and dividing by the accumulated weight -- after which SSB, OBF, matched-filter and iCoM all run unchanged. That makes it exact where it can be: with positions on a lattice the splat is an identity map, and all five kernels reproduce `from_dataset4d` bit-for-bit. The ungridded preamble both classes need -- position validation, "auto" scan sampling, origin fitting, bright-field masking, normalization, anchoring to the bounding box -- is hoisted into `build_vbf_stack_from_dataset3d`, mirroring the existing `build_vbf_stack_from_dataset4d`, so the two entry points cannot drift. `ShadowMontagePtychography.from_dataset3d` now calls it, and its private `_validate_positions` / `_infer_scan_sampling` are gone in favour of the shared `validate_probe_positions` / `infer_scan_sampling`. Regridding is the only approximation, and its failure mode is holes: a grid pixel no probe reached stays zero and the FFT reads that as signal. `regrid_vbf_stack` reports the fraction and warns above 1%, pointing at a coarser `scan_gpts` or at ShadowMontagePtychography, which needs no grid at all. An explicit `scan_gpts` rescales the positions so the field of view is preserved rather than cropped. `regrid_vbf_stack` splats one bright-field image at a time on purpose: `scatter_add_splat` accumulates its whole batch into a single canvas, which is what the montage wants but would sum the stack away here. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
… zero
`from_dataset3d` was badly wrong on masked or irregular scans, which is the case
it exists to serve.
The cause is an interaction with `_preprocess`, which zeroes the DC bin -- that
subtracts the mean over the *whole* grid, holes included. Zero-filled holes
therefore end up sitting at `-mean` rather than at the signal's own level, a
hard-edged step that the deconvolution then smears across the reconstruction.
Filling holes with each bright-field image's mean over the visited pixels puts
them exactly at the zero level instead, and the step disappears. Measured against
ground truth on a disk-masked scan with 20% holes:
zero fill (before) 0.25
mean fill (now) 0.69
montage, same data 0.69
Contiguous holes turn out to be far more damaging than scattered ones at equal
fraction -- 0.25 versus 0.51 at 20% -- because an excluded region is a
low-frequency mask the deconvolution spreads everywhere, while scattered holes
behave like noise. That makes a masked, non-rectangular scan exactly the worst
case, and exactly what users hit.
Nearest-neighbour filling was measured too and is a wash (0.66 vs 0.66), so it is
not offered. `hole_fill="zero"` remains available.
The hole warning was also recalibrated: at 1% it cried wolf, since 5% holes with
mean fill reconstructs as well as a full raster. It now fires at 10% and says
what the fill did, rather than implying holes are read as zero signal.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Regridding is where an ungridded reconstruction goes wrong, and the failure is visible in the occupancy map long before it is visible in the reconstruction. Stashes hole fraction, occupancy, positions and resolved grid on `_regrid_info` so they can be inspected without re-running the pipeline by hand. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ampled
`from_dataset3d` combined with `upsampling_factor > 1` returns a replica of the
contrast-transfer function rather than an extension of it, on any scan that is
not already close to a lattice.
This is intrinsic, not a defect in the regridding. `upsampling_factor` recovers
detail above the scan Nyquist from the aliased content of the bright-field
images, which depends on where each probe actually sat. Binning those
measurements onto a grid discards exactly that information, so there is nothing
left for the upsampling to unfold.
Measured on a white-noise object against the analytical parallax CTF, correlation
over the band above the scan Nyquist, with *no holes at all*:
sub-pixel scatter 0.00 0.25 0.50 1.00 grid px
from_dataset3d 0.997 0.965 0.822 0.711
ShadowMontage 0.995 0.994 0.990 0.983
The montage is unaffected because it places each measurement at its own position
and never regrids. On an exact lattice the two agree (0.997 vs 0.995), which is
what the existing bit-exactness test pins.
So `from_dataset3d` now measures how far the positions sit from the grid they
were binned onto, stores it as `_regrid_info["lattice_rms_px"]`, and warns at
reconstruct time when upsampling is requested above 0.1 px of scatter -- pointing
at the montage with a finer `scan_sampling`, which is the operation that actually
does what was wanted. Gridded reconstructions carry no regrid info and are
unaffected.
Found by testing against the white-noise/analytical-CTF harness in
ipynb-playground/white-noise-ptycho.ipynb, where |FFT(reconstruction)| is the CTF
directly. Earlier tests missed it because their scan pitch equalled the object
sampling, leaving nothing above the scan Nyquist to unfold.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Regridding an ungridded scan and then calling `reconstruct(upsampling_factor=U)`
cannot unfold, because binning to the coarse grid discards the sub-pixel probe
positions the unfolding depends on. But `upsampling_factor` is itself only
Fourier tiling, which is exactly zero-insertion -- so the fix is to build that
same sparse comb *directly on the fine grid*, with the teeth on the real probe
positions instead of a lattice.
`from_dataset3d(..., upsample_positions=U)` refines the grid by U, scales the
positions onto it, and splats. Positions are then quantized only to 1/U of a scan
pixel instead of being binned away.
The DC handling is what makes it work, and it is not optional. `_preprocess`
zeroes the DC bin, subtracting the mean over the whole fine grid; if the comb is
built from raw values its gaps end up at -mean rather than 0, and the
reconstruction inverts outright (in-band correlation -0.07). Removing each
image's mean over its own positions *before* splatting makes the comb zero-mean,
so the DC zeroing is a no-op and the gaps stay at exactly zero. That is
`hole_fill="comb"`, which `upsample_positions` selects automatically.
Measured against the analytical parallax CTF, correlation over the band above the
scan Nyquist, on a scan scattered by a full pixel:
bin then upsampling_factor=2 0.720
upsample_positions=2 0.961
ShadowMontagePtychography 0.982
and it holds across the range -- 0.997 / 0.963 / 0.963 / 0.962 at 0 / 0.25 / 0.5
/ 1.0 px scatter, against 0.997 / 0.966 / 0.825 / 0.720 for the old route.
Tests move onto the white-noise / analytical-CTF harness for this, since it is
the only fixture that can tell unfolding from replication: the object has
constant Fourier amplitude, so |FFT(reconstruction)| is the CTF directly, and the
scan pitch is deliberately half the aperture's transfer limit so there is a real
factor of two to recover. conftest gains that harness, including Fourier-shift
simulation at fractional positions and jittered positions with the outer ring
pinned so `boundary="wrap"` stays comparable. The previous fixtures scanned at
the object sampling, leaving nothing above the scan Nyquist, which is why this
went unnoticed through several rounds of testing.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Drops `upsample_positions` and `hole_fill="comb"`, both added last commit, and
defaults `interpolation` to "nearest". All three were redundant or wrong.
`upsample_positions=U` did nothing that `scan_sampling / U` does not already do:
it refined the grid, scaled the positions onto it, and splatted. Measured on the
white-noise CTF harness with a scan scattered by a full pixel, radial CTF
correlation is 0.9985 for `upsample_positions=2` against 0.9967 for halving
`scan_sampling` directly -- the gap is a one-pixel difference in how the grid is
rounded, not a difference in method. Specifying the sampling is also the same way
ShadowMontagePtychography says it.
`hole_fill="comb"` was algebraically identical to "mean". Comb subtracted each
image's mean before splatting so the gaps sat at zero; mean-fill puts the gaps
*at* that mean, and `_preprocess` then subtracts it as the DC bin, landing them
at zero too. Reconstructions agree to 1.3e-3 relative. Mean-fill is in fact the
tidier of the two, since it centers on the mean of what is actually on the grid.
`interpolation="nearest"` is now the default, and wins everywhere -- not only on
a finer grid, but even at the native spacing where it leaves holes and bilinear
leaves none:
jitter grid nearest bilinear
0.5 native 0.0% holes 0.9839 0.0% holes 0.9616
1.0 native 31.8% holes 0.9865 0.0% holes 0.9478
1.0 fine 76.8% holes 0.9961 34.5% holes 0.9860
Blurring each measurement across four pixels costs more than the holes do, once
mean-fill centers them. That also retires the hole warning as it stood: at a 10%
threshold it fired on the *better* configuration and pushed callers toward
bilinear. It now fires only when there were enough positions to cover the grid
and more than half of it is still empty -- positions clustered, rather than
merely sparse -- which is the case the caller can actually act on. The verbose
line likewise reports comb gaps as gaps rather than as unvisited pixels.
`scan_gpts` no longer rescales the positions to fill the requested shape, which
silently contradicted `scan_sampling` and cost real fidelity (0.853 against 0.962
on the same data). It pads the canvas instead, and raises if asked for a grid too
small to hold the positions.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Three separate faults, all reached with device='gpu' on Apple silicon. **float64 on a device that has none.** ShadowMontagePtychography hardcoded torch.float64 for every positional tensor -- coordinates, shifts, sampling, canvas origins -- so every one of its entry points raised on MPS. They now take their dtype from the device via `preferred_float_dtype`, which already existed for the splat accumulators under the name `preferred_accumulator_dtype` (kept as an alias). The float32 fallback resolves canvas positions to about 1e-3 pixels on a 10k-pixel canvas, far below the sub-pixel detail any of this preserves; measured CPU-vs-MPS agreement is ~4e-6 relative across the parallax paths. **A CPU tensor meeting an MPS one.** `CenterOfMassOriginModel.fit_origin_background` validated explicitly-passed `probe_positions` without moving them to its device. Probe positions arrive as numpy from the ungridded constructors, so they became a CPU tensor and met `origin_measured` on MPS: "Passed CPU tensor to MPS op". The `probe_positions is None` branch already built its grid with `device=self.device`; the explicit branch now matches, as its own `origin_fitted` setter already did. This was invisible to every earlier check because they all passed `force_fitted_origin`, which skips the centre-of-mass fit entirely -- the same blind spot that hid a coverage gap earlier in this branch. The new tests fit the origin for real. **A canvas that changed shape with the device.** `_return_canvas` floors and ceils bounds that are integers in exact arithmetic but arrive as 4.0000001 from a float32 k-grid, costing a whole pixel -- and a different number of them in float32 than in float64, so the same data gave a 42x42 canvas on CPU and 41x41 on MPS. `_snap_to_integer` rounds bounds that are integral to within 1e-4 before the floor/ceil. That also retires a long-standing wart: `boundary="pad"` canvases were one pixel larger than needed on each side. 42 of the 106 battery outputs change size accordingly, all of them `pad`; `wrap` and every DirectPtychography output remain bit-identical, so the Fourier-equivalence results are untouched. Also falls `torch.linalg.lstsq` back to CPU in `fit_hyperparameters_least_squares`, which is unimplemented on MPS. That one is pre-existing, and follows the "Fall back to CPU (to support MPS)" precedent already used for `eigh` in origin_models. Tests now parametrize over whatever accelerators are present and assert both classes agree with CPU, that the padded canvas shape is device-independent, and that the ungridded constructors fit the origin on device. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`dft_upsample` was wrong on two counts, and both only bit when
`cross_correlation_shift` was called with `upsample_factor > 1`:
- the offset was added to the input *frequency* index instead of
re-centering the output sampling positions, which factors out as a
per-row phase twist rather than a translation of the fine grid;
- the exponent carried the forward transform's sign, so it evaluated
the reflection of the correlation rather than the correlation.
`cross_correlation_shift` then converted the local peak index back using
`upsample_factor` where the grid is actually centered on its middle tap,
`ceil(1.5 * upsample_factor)` -- a further half-pixel bias in both axes,
visible as a (0.5, 0.5) result for a pair of identical images.
Measured against Fourier-shifted ground truth over shifts from 0 to 11.5
px, worst-case error falls from 2.0 px to 0.0022 px, and `dft_upsample`
now agrees with `ifft2` to 5e-11 where the two grids coincide. The
numpy path also now matches the (independently correct) torch port.
This affects `DriftCorrection.align_translation` / `align_affine`, which
correlate at `upsample_factor=8`; `tomography_utils` calls it at the
default factor of 1 and is unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
A multi-frame acquisition -- several interleaved passes of the same
region, so that drift shows up between frames rather than smeared within
one -- needs its frames reconstructed onto comparable canvases before
they can be aligned. Two things stood in the way:
- `from_dataset3d` anchors the scan grid at the probe-position bounding
box, and that corner differs from frame to frame. On the hexagonal
apoferritin dataset the anchors span 70 A, so nothing recorded where
a reconstruction sat in the specimen's own coordinates.
- `boundary="pad"` sizes the canvas from the positions, so frames came
out 156x155, 151x155, 156x154 and 151x154 -- not stackable, let alone
correlatable.
`scan_origin` now records the anchor, and `obj_origin` maps object pixels
back onto the positions that were passed in. `reconstruct` gains
`obj_origin` and `obj_fov`, which name a fixed window in those same
coordinates; frames reconstructed with both land pixel-identical. The
canvas corner is deliberately *not* rounded to a whole scan pixel, since
the per-frame lattices are themselves offset from one another and
snapping to them would reintroduce a sub-pixel error of the same order as
the drift being measured.
`estimate_frame_drift` then cross-correlates the frames -- leave-one-out
against the mean of the others, iterated, mean-referenced -- and returns
a per-frame drift in Angstrom to subtract from the positions. It refuses
frames whose canvases disagree in shape, origin or sampling, since a
correlation between those measures the canvas offset instead and would
otherwise succeed silently.
Measured on the 9-frame hexagonal apoferritin dataset (6247 positions,
+/-1000 units of scan): drift of 21 A in row and 14 A in column across
the acquisition, and correcting for it raises the combined
reconstruction's RMS gradient by 12.8% (0.00388 -> 0.00438) while the
opposite sign lowers it by 12.2% (-> 0.00341).
Behaviour-preserving for every existing call path: an 86-output battery
over both classes -- gridded and ungridded, all kernels, both boundaries,
both interpolations, upsampled, padded and tilted -- is bit-identical
against 3294590.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`_patch_variance_loss` averaged the per-pixel variance over pixels at 90% of the *peak* accumulated weight, to keep partially-filled canvas edges out of the average. On a raster scan the interior weight is flat and equals the peak, so that cut means exactly "interior only" -- which is where it was validated. On an ungridded scan the weight is uneven everywhere, not just at the edges. The cut then keeps whichever spots the probe positions happened to crowd into -- 13.7% of the canvas on the hexagonal apoferritin dataset -- and which spots those are moves with the trial defocus, so the loss tracked that selection rather than the image. Measured over C10 = 9 to 17 kA on that dataset: peak-weight cut minimum at 9.5 kA weight-averaged minimum at 13.0 kA global variance_loss 13.0 kA peak image contrast 13.0 kA (28% swing, unambiguous) Weighting each pixel's variance by its accumulated weight needs no threshold, is the inverse-variance combination of estimates whose uncertainty goes as 1/n, and is the same functional as `variance_loss` -- so the patch fit and a hyperparameter search can no longer disagree about what a good defocus is. On the aligned hexagonal montage this takes `fit_defocus_gradient` from 4 of 9 patches bracketing, a C10 offset 1.6 kA away from the optuna fit and a gradient that changed with `patch_grid`, to 9 of 9, an offset within 1.7% of it, and a gradient stable to 10% between 2x2 and 3x3. It costs a little on the gridded synthetic the cut was tuned against: correlation with a seeded defocus plane falls 0.995 -> 0.980, since discarding the edges there is a small genuine win. Being unbiased on both is worth more than 0.015 of correlation on one. Behaviour-preserving elsewhere: the 86-output battery is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`optimize_hyperparameters` and `grid_search_hyperparameters` hardcoded
`variance_loss`. Both now take `loss=`, accepting `"variance"` (the
default, unchanged), `"rms_gradient"`, or any callable of the
reconstruction.
`rms_gradient_loss` is the classic autofocus metric -- the negated RMS
image gradient, per Angstrom so it is comparable across samplings and
upsampling factors. Two reasons to have it:
- it is far better conditioned. Over a defocus series on the hexagonal
apoferritin dataset it swings 28% where the variance loss swings
0.08%, and the two agree on the optimum to within one step of a
250 A grid. The gap is structural: the variance loss compares
bright-field images with each other and saturates once they agree,
while this measures the image you actually want. Measured across
`boundary="wrap"`, an automatic `"pad"` canvas, a frozen `pad_px` and
a pinned `obj_fov`, the optimum moved by at most one grid step.
- it is defined for every deconvolution kernel, where
`ShadowMontagePtychography.variance_loss` is defined only for the
parallax one. Without it there is no way at all to tune aberrations
for the real-space ssb/obf/mf kernels.
On the real dataset at 200 optuna trials the two losses land in the same
place -- the resulting sharpness differs by less than the spread over
random seeds -- though C10 across three seeds is 2.4x tighter with the
gradient (126 A against 297 A). It does reward amplitude as well as
sharpness, which is safe for aberrations and rotation but is documented.
Behaviour-preserving: the 86-output battery is unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`fit_defocus_gradient` already stored the per-patch measurement it fits a plane to, but only privately, so plotting the loss curves meant calling `defocus_map` a second time and paying for the whole sweep twice -- 30 s each on the hexagonal apoferritin dataset. Those curves are worth looking at: a patch whose minimum sits on an endpoint of `c10_values` is dropped from the plane, and nothing but the reported count says so. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Both classes derived the wavelength from `energy` through
`electron_wavelength_angstrom`, which is the relativistic electron de
Broglie formula. That is correct for electrons -- 0.019687 A at 300 kV,
the textbook value -- but it is the wrong physics for anything else: a
7.9 keV photon gets 0.1375 A where `hc/E` is 1.5694.
Every constructor now takes `energy` or `wavelength` (in Angstrom), and
raises if given both or neither. `wavelength` is a validated property on
the base, so AutoSerialize already persists it; there is a round-trip
test rather than new serialization code.
The mrad machinery needs nothing: at the cSAXS geometry that motivated
this, `angular_sampling` comes out at 10.4 urad and the probe edge at
2.55 mrad, which the existing code handles unchanged.
Two knock-on signature changes, since Python will not take a defaulted
parameter before an undefaulted one and `wavelength` has to sit
alongside `energy`:
- `semiangle_cutoff` and `rotation_angle` gain `None` defaults. Both
were already validated at runtime -- `semiangle_cutoff` by its
setter, `rotation_angle` by a new check that reproduces the error it
used to get from being positional.
- `TestSemiangleCutoff.test_from_virtual_bfs_requires_it` asserted the
signature had no default, as a proxy for "never optional". It now
asserts the error instead, which is the actual guarantee, and covers
`from_virtual_bfs` -- which the behavioural test did not.
Behaviour-preserving: the 86-output battery is unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Both classes described the probe as a circular aperture times
`exp(-i chi)`, which no zone plate, central stop or speckled X-ray
illumination is. `FourierProbe` now sits behind `gamma_factor` with two
constructors -- `from_aberrations`, the analytic probe unchanged, and
`from_array`, a measured or independently reconstructed `psi(k)`. Pass
one as `fourier_probe=` and `semiangle_cutoff` becomes optional.
The check that makes this trustworthy is that there is no ground truth
for a measured probe, so the test feeds an analytic probe back in *as an
array*: SSB, OBF and matched-filter all reproduce the analytic
reconstruction to 1e-7 in both classes.
Two things fell out of writing it:
- **Zero-extension, not wraparound.** An array can only be read on its
own grid, and `gamma_factor` asks for `k -/+ q`, which leaves the
detector constantly. Beyond the detector's Nyquist nothing was
measured, so `psi` is zero there; indexing modulo the grid instead
folds the opposite edge of the aperture back in. On the electron
fixtures, whose bright-field mask is cropped tight to the disk, that
alone moved the reconstruction by 13%.
- **Commensurate grids, or say so.** `k -/+ q` lands on the probe grid
only when the probe's field of view is an integer multiple of the
canvas's. Off-grid sampling raises by default and names the fix
(choose a commensurate canvas, or zero-pad the probe in real space)
rather than quietly interpolating -- a speckled probe varies over a
few detector pixels, so that is a real approximation.
`interpolation="bilinear"` opts into it.
Everything that reads the aberration surface -- the parallax and iCoM
kernels, the `sign(sin(chi))` phase flip, the defocus gradient -- raises
for an empirical probe rather than silently treating the absent
coefficients as zero, which for parallax would have summed the
bright-field stack into a plain incoherent image.
Behaviour-preserving: the 86-output battery is unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The real-space SSB/OBF/MF stencil deposits `(2R+1)**2` taps per scan
point through `index_add_`. On the X-ray dataset that motivated this --
167k bright-field pixels, 1174 positions, a 180x140 canvas -- a
full-canvas stencil extrapolates to about 90 minutes on MPS. These
kernels are never compact (`r**-1.5` tails), so a full canvas is what
they need.
`reconstruct` gains `convolution_mode`:
- `"fft"` splats each bright-field image onto the canvas and multiplies
by the kernel in `q`. Measured at 54 s for the case above, and it is
**exact**: on a gridded scan it reproduces `DirectPtychography` to
2e-7, where a radius-5 stencil is 20-34% off and a radius-12 one
still 2-5% off. With `boundary="pad"` it doubles the canvas and crops
back, since a Fourier convolution is otherwise circular.
- `"stencil"` keeps the truncated box but evaluates it as a grouped
`conv2d` rather than a loop over taps: 5.4 s against 78 s at radius 8
for the same case.
- `"auto"` reads `stencil_radius` -- naming one is a request to
truncate, leaving it `"auto"` takes the FFT.
Unlike `DirectPtychography`, this streams over detector pixels into a
single canvas rather than materializing an `(N_bf, Ry, Rx)` stack, which
at 167k bright-field pixels would be 34 GB.
`splat_and_convolve` is a reorganization of `scatter_add_convolve`, not a
change of operator, and there is a test on adversarial inputs to pin
that. The subtlety it has to reproduce is that the scatter tests the
boundary at the *deposit* position, so a point outside the canvas still
contributes inward through the kernel -- hence splatting onto a canvas
grown by the radius.
Battery: 16 of 86 outputs move, all montage ssb/obf, by 1.5-3.4e-6
relative. That is float32 reassociation from summing taps inside a
grouped convolution instead of one `index_add_` per tap; every parallax
output is bit-identical.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Two things had to be settled before an empirical probe could be used on a real canvas, and both were wrong in the plan. **The commensurability condition is on the fields of view, not the pixel sizes.** `dq_det = pixel / (wavelength * distance)` is fixed by the experiment, so the probe's field of view is too, and cropping the detector does not change it. `psi(k -/+ q)` is an exact lookup only when the canvas field of view is a whole multiple of the probe's. `resampled_to` now refines the probe onto the canvas grid by zero-padding it in real space, which is the sinc interpolation the probe's band limit licenses -- exact where the real-space probe is confined to its field of view, and the honest best estimate otherwise, since nothing between the detector samples was ever measured. A canvas equal to the probe's field of view needs no resampling at all and is the assumption-free case. **Rotation silently broke it.** `_return_k_grid` passively rotates the k-grid into the scan frame, which is what the analytic aperture wants since it is *evaluated* there. An array can only be *read* on its own lattice, and a rotation takes `k -/+ q` off it. Measured against the analytic reconstruction: 0 and 90 degrees agree to 1e-7, 15 degrees was 7.6% wrong with no indication. Rotations that map the lattice onto itself are allowed, the rest raise. Also caught while testing: for the montage, `upsampling_factor=U` at reconstruct time is *bit-identical* to constructing with `scan_sampling / U`, for both parallax and SSB on an ungridded scan -- the positions are never gridded, so neither route loses sub-pixel information. That is not true of `DirectPtychography.from_dataset3d`, where regridding discards exactly that. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The previous commit refused a `rotation_angle` that was not a multiple of
90 degrees when a `fourier_probe` was set. That was wrong: rotation is a
real experimental parameter, not something to opt out of.
The underlying constraint is real but narrower than the guard implied.
`v_m` pairs with the specimen-frame `k = R(theta) k_det`, so Gamma needs
`psi(k_det -/+ R(-theta) q)`: `k` itself stays on the detector lattice
and only the *offset* arrives rotated. No rotation of the array fixes
both at once, so a general angle has to be interpolated -- but
interpolated well, not refused.
`probe_oversample` refines the probe before sampling it, by the exact
real-space zero-padding already used to match the canvas. Bilinear error
then falls as the square of the refinement. Measured on the speckled
X-ray probe against an exact Fourier-shift evaluation, at a generic
sub-pixel offset:
unrefined 20.9% 4x 1.7% 16x 0.064%
2x 7.5% 8x 0.52%
A rotation that maps the lattice onto itself still takes the exact path
untouched, and an analytic probe ignores the setting entirely since it is
evaluated rather than sampled.
One caveat the tests now state rather than hide: refinement converges on
the *band-limited* interpolant, which is the best inference from the
samples but is not the true probe where the real-space probe is not
confined to its field of view. On the electron fixture, whose aperture is
cropped to an 11x11 grid, that floor is about 6% and no refinement
removes it. On the X-ray probe 1.5% of the real-space power sits at the
field edge, so expect that order.
Battery unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…anvas Two gaps the X-ray data walked straight into. `from_dataset3d` derived the bright-field mask by thresholding the mean diffraction pattern, which assumes a bright disk on a dark field. The cSAXS probe is an annulus with a direct beam far brighter than the rest, so the default threshold selected **6 detector pixels** out of 410k. The mask can now be given -- `abs(psi) > threshold` from an empirical probe is the natural choice -- and it has to be given at construction, since the vBF stack is built from it and overriding `bf_mask` afterwards leaves a stack with the wrong number of columns. `boundary="wrap"` refused `obj_origin` / `obj_fov`. That guard confused two independent things: the boundary rule says what happens to deposits outside the canvas, and the window says where the canvas is. Wrapping into a *given* window is perfectly meaningful, and it is what an empirical probe needs -- a canvas matching the probe's own field of view samples psi exactly, where padding the canvas for a linear convolution would put the reciprocal grids out of step. The default is unchanged: no window given, `"wrap"` is still the scan grid. Battery unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The montage refused iCoM outright: `k . q / |q|**2` is unbounded as q -> 0, so no box stencil captures it. That is an objection to *truncating*, not to the kernel -- and `convolution_mode="fft"` does not truncate. It now runs there, matching `DirectPtychography` to 1.7e-7, and is refused only for the stencil. `DirectPtychography` also required an analytic probe for iCoM, which it never needed: the kernel reads only the detector coordinates and the canvas frequencies, never psi. With a matched probe normalization the two are bit-identical. That matters for a probe no aperture describes, where iCoM and the first-order limit of SSB are the natural comparison -- and on a scan whose accessible q is far below the aperture's transfer limit, they should agree. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
`boundary="wrap"` on a canvas that spans the whole scan is harmless -- it is what makes the montage reproduce `DirectPtychography`, which is periodic in the scan by construction. On a canvas that does *not*, the positions outside come back on the opposite side and lay a second, offset copy of the specimen over the reconstruction. That is what happened on the X-ray data: the canvas was pinned to the probe's 15.07 um field of view while the scan spans 20 um in one axis, so 4.9 um of it wrapped over the top and printed a ghost pillar above the real one. It reads as a real reconstruction with an artifact in it, which is the kind of thing that gets explained away. The warning names the overhang in Angstrom and the two fixes: widen `obj_fov` -- to the next whole multiple of the probe's field of view, if the probe is empirical and the grids have to stay commensurate -- or use `boundary="pad"`, which drops those positions instead of folding them. Silent whenever the canvas covers the scan, including the default where `"wrap"` sizes itself to the scan grid. Battery unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Yu et al. (Microsc Microanal 28, 1526) cross-correlate the centre-of-mass shift map with `(r_p - r_xy) / |r_p - r_xy|**2` truncated to an n x n box, and call it riCOM. That kernel is the iCoM operator's real-space form: for `G = ln|r| / 2pi`, `grad G = r / (2 pi |r|**2)`, whose transform is `-i q / |q|**2` -- the operator already implemented here. And because each bright-field pixel's kernel is linear in `k_m`, summing over the detector collapses the montage's per-pixel convolutions into a single convolution of the COM shift, which is riCOM exactly. So it was already supported, except that I had explicitly forbidden it two commits ago: iCoM plus `convolution_mode="stencil"` raised, on the grounds that `k . q / |q|**2` is unbounded as q -> 0 and so has no useful truncation. That reasoning is wrong. Truncating it is not a compromise, it is the method -- the radius is a high-pass cutoff, and suppressing the long-range drift that blurs an iCoM image is the entire point. Measured on the 96x96 model fixture: a radius spanning the canvas reproduces untruncated iCoM to a correlation of 0.99, and shrinking the kernel raises the high-to-low band ratio monotonically, by 1.4x from radius 40 down to 5. The kernel identity itself is checked directly against `r / (2 pi |r|**2)`. The truncation warning is now suppressed for iCoM, where reporting the deliberate cutoff as "error" would be backwards. It still fires for ssb/obf/mf, whose truncation genuinely is an approximation. Battery unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The iCoM kernel `K_m(q) = k_m . (-i q / |q|**2)` is linear in `k_m`, so
sum_m FFT(V_m) K_m(q) = A(q) FFT(splat(com_x)) + B(q) FFT(splat(com_y))
and the sum over the detector can be done *first*. That turns the whole
reconstruction into two convolutions of the centre-of-mass shift instead
of one per bright-field pixel -- which is what riCOM does, and the
difference between "supported" and "real time".
Measured on the cSAXS data, 167313 bright-field pixels, against the
per-pixel SSB path on the same canvas:
U=1 64x32 0.24 s vs 3.3 s 14x
U=2 128x64 0.14 s vs 11.3 s 81x
U=4 256x128 0.10 s vs 41.1 s 397x
The time *falls* with canvas size because the cost is now the one-off
collapse -- a single matmul over the detector axis -- and the two
transforms are noise beside it. 1174 positions in 0.14 s is about 8 kHz
from Python, the same order as the paper's C++ implementation.
The stencil route collapses the same way, giving riCOM proper: the box is
taken about the origin, since unlike SSB the iCoM kernel carries no
parallax shift to divide out and add back.
The collapse is skipped when a defocus gradient gives each position its
own shift, which it cannot represent; that falls back to the per-pixel
path, and a test with a negligible gradient checks the two agree.
Correctness is pinned against `DirectPtychography` (2.3e-7) and the
collapse itself against an explicit detector sum.
Battery unchanged.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ychographyMontage The class was named for its parallax kernel, but it now runs five: SSB, OBF and matched-filter as exact FFT convolutions, iCoM, and iCoM truncated to riCOM. Only `prlx` builds a shadow montage, so the name promised one method and delivered one of five. `DirectPtychographyMontage` follows the repo's noun-first, qualifier-suffix convention -- ObjectPixelated, ProbeParametric, PtychographyDatasetRaster -- so it sorts and autocompletes next to DirectPtychography, which is the class a reader is choosing between. Introduced on this branch in f3482a1 and never released, so nothing published depends on the old name. "Shadow montage" survives as what it actually names: the parallax construction. Pure rename: an AST fingerprint with docstrings and comments stripped is byte identical across all seven touched files. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…that did The docstrings were written a commit at a time, each against the state of the one before, and several outlived the code: - The montage class docstring framed parallax as the method and SSB/OBF/MF as an expensive approximation to avoid. 529055d made them exact by FFT and the default; iCoM and riCOM went undocumented entirely. - DirectPtychography.from_dataset3d told the reader the montage cost them "the parallax kernel being the only exact one there." Untrue since 529055d. - Which class to use was scattered over five places, phrased differently in each, two of them stale. Now stated once, in the montage class docstring, with the rest pointing at it -- including the memory ceiling that is the actual reason an X-ray dataset can only be reconstructed there (34 GB against 800 MB at 167k BF pixels). - DirectPtychography's class docstring was empty, so every one of those pointers landed on nothing. Adds a References block: the shadow-montage and riCOM papers, plus convolution WDD and segmented-detector OBF marked as related rather than as what this implements. Four code changes, each fixing something that misdirects a caller: - _require_analytic_probe listed 'ssb', 'obf', 'mf' as the escape hatch and omitted 'icom', which since 28cf4a8 runs with an empirical probe and never reads psi. - variance_loss's guard sent callers to the parallax kernel; since 1cfe888 the answer is loss='rms_gradient', defined for every kernel. - preferred_accumulator_dtype was kept as "the older name" for a contract no caller outside the module ever had. - optimize_hyperparameters documented three parameters absent from its signature. Also repairs a comment spliced from two fragments mid-sentence, a :func: reference to a function that does not exist, and trims ten comment blocks that had grown into commit-message-length narratives. Verified inert: an AST fingerprint with docstrings and comments stripped shows only the three code edits above and nothing else across seven files. 795 tests pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Collaborator
|
It overall looks good to me! Only thing I would ask before merging into dev is that you update the tutorials (or at least make sure these changes don't break anything) Also maybe clean up docstrings and comments to remove claudisms. e.g. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Adds
DirectPtychographyMontage, a second direct-ptychography class that accumulates the scan onto a real-space canvas instead of Fourier transforming over it, and extends both classes to handle probes and wavelengths that the electron-only code path could not express.DirectPtychographygains an ungridded constructor, a test suite (it had none), and three bug fixes — one of which meant no reconstruction could be saved and reloaded.The idea
Every kernel in
DirectPtychographyis a multiplier onG(k, q), the virtual bright-field stack Fourier transformed over the scan. That requires the scan to lie on a regular grid, and it requires the whole transformed stack in memory at once.DirectPtychographyMontagedrops the scan-space transform: each bright-field image is deposited onto one shared canvas at its own probe position. The detector axis is handled identically. This buys three things that are not available in the Fourier formulation:(N_bf, Ry, Rx)complex64 — 34 GB at 167k bright-field pixels on a 128×170 canvas, where the montage needs ~800 MB.Which class to use
DirectPtychography— exact, and cheapest at one scan FFTDirectPtychography.from_dataset3dregrids onto a lattice first, the montage never gridshole_fillcannot serve filled holes and deliberate gaps at onceBoth classes take the same kernels (
prlx,ssb,obf,mf,icom) and agree numerically where they overlap. On the montage the convolution kernels are exact by default, evaluated by FFT on the canvas;convolution_mode="stencil"truncates them to a box instead. Truncatingicomis the exception that is not an approximation — it gives riCOM, where the radius is a deliberate high-pass cutoff.What's new
DirectPtychographyMontagewithfrom_dataset4d/from_virtual_bfs/from_dataset3d, all five kernels, wrap and pad boundaries, and canvas upsampling that is bit-identical to constructing at a finerscan_sampling.FourierProbe— one interface over an analytic aperture-plus-aberrations probe and an empirical complex array. Makes a zone plate, a central stop, or a speckled X-ray optic usable as ψ(k). Sampled exactly by real-space zero-padding where the canvas and detector grids are commensurate, with bilinear interpolation as an opt-in fallback.wavelength=alongsideenergy=, since the electron de Broglie formula returns 0.137 Å for a 7.9 keV photon against the correct 1.569 Å.DirectPtychography.from_dataset3d— ungridded scans by regridding, with diagnostics and warnings for the two ways that goes wrong.defocus_gradient,fit_defocus_gradient,defocus_map.obj_origin/obj_fov) andestimate_frame_drift, so the frames of a multi-frame acquisition reconstruct onto pixel-identical canvases and can be aligned.rms_gradient_lossand aloss=parameter on both hyperparameter searches. Far better conditioned than the variance loss (28% dynamic range against 0.08% over a defocus series) and defined for every kernel, wherevariance_lossis parallax-only.conv2dfor stencils (14× over the tap loop), FFT convolution for kernels that span the canvas, and a detector collapse for riCOM that turnsnum_bfconvolutions into two (14×/81×/397× at upsampling 1/2/4 on real data).Fixes to existing behavior
These change results for code already on
dev:serializerecordedtype(value).__name__as the container type, so atorch.Sizewas written as"Size"and rejected on load.DirectPtychography.gptsisbf_mask.shape[:2], so every reconstruction was unloadable.dft_upsample/cross_correlation_shiftapplied the sub-pixel offset to the input frequency index rather than re-centering the output grid, carried the forward transform's sign, and converted the peak index against the wrong center. Worst-case error over shifts of 0–11.5 px falls from 2.0 px to 0.0022 px. This also affectsDriftCorrection.align_translation/align_affine.visualize's scalebar ignoredupsampling_factor, so an upsampled object reported double its true field of view.from_virtual_bfsdefaultedsemiangle_cutoff=None, which always raised aTypeErrorin the setter.Testing
DirectPtychographyhad no tests before this branch. Adds ~4,300 lines acrosstest_direct_ptychography.py,test_direct_ptychography_montage.py,test_imaging_utils.pyand a sharedconftest.py.370 passed / 26 skipped in
tests/diffractive_imaging; 795 / 34 for the full suite.Equivalence between the two classes is pinned throughout — the montage reproduces
DirectPtychographyto float precision on a gridded scan for every kernel, and the refactor inb0dff8fcwas verified bit-identical across 35 outputs by running both trees in parallelgit worktrees.References
ozaf126 (2026). https://doi.org/10.1093/mam/ozaf126
1526 (2022).
Related, though neither is the kernel implemented here — the first convolves a WDD kernel where this convolves SSB/OBF/MF kernels, and the second uses a segmented detector rather than a pixelated one:
Notes for review
quantem-tutorials.ShadowMontagePtychographywas renamed toDirectPtychographyMontagebefore merge. It was introduced on this branch and never released, so nothing published depends on the old name.