Repository navigation
STalign 2.0.0: much faster alignment, more accurate defaults, bug fixes - #57
Merged
Merged
Conversation
rasterize and rasterize_with_signal looped over cells in Python. The Gaussian kernel and its per-window normalization are separable, so each cell's contribution is an outer product of two 1D kernels; _raster.py accumulates these as tiled matrix products in numpy (no new dependencies). Images match the previous code to 1e-12 relative; 86k cells take ~35 ms instead of ~2 s, 2M cells ~0.85 s instead of ~47 s, and peak memory is bounded. Also fix utility bugs found while auditing the module: - rasterizePCA called the undefined name tools (NameError). - rasterize_with_signal took np.abs of signed signals. - normalize returned NaN for constant arrays. - saveRasters ignored a folder path without a trailing slash. - download_aba_* saved HTTP error pages as data. - calculate_tre rejected torch tensors. - plot_brain_regions cycled 10 colors over ~150 regions. - point_annotator and curve_annotator raised NameError on output names.
- Source-to-target maps (transform_points_source_to_target, build_transform(direction='f'), transform_image_target_to_source, build_transform3D(direction='f')) now use the exact inverse of the backward map that LDDMM fits, by fixed-point iteration, instead of a forward Euler map that is not its inverse. exact_inverse=False keeps the old map. The simulation tutorial's ground-truth error drops from 9.68 to 6.68 um. - transform_points_target_to_source integrated the velocity in the wrong time order; it now equals build_transform(direction='b'). - build_transform3D(direction='f') always raised an error. - analyze3Dalign indexed a 10-unit grid with dx-based indices, so atlas coordinates were wrong whenever dx != 10 (by 4.3 mm in the STARmap tutorial); it now evaluates the map at each cell. - v_to_phii accepts the velocity layout LDDMM returns; transform helpers accept numpy, torch and integer inputs. The implementations move to _transforms.py; STalign.py keeps the public functions and docstrings.
The per-iteration cost was dominated by grid_sample's CPU backward, which runs single-threaded at batch size 1, and by redrawing figures every 10 iterations. The new torch engine chains the gradient by hand through torch's grid_sampler kernels with the points split into chunks, uses the analytic smoothed regularizer gradient, and draws once at the end by default (draw=True restores the old behavior). _lddmm_common.py holds what both engines share: input checks, kernels, drawing, threads (num_threads, STALIGN_NUM_THREADS) and backend selection (backend='auto'|'cpp'|'torch', backend_info()). Fixes: - LDDMM_3D_to_slice's regularizer transformed v over 2 of its 3 axes; it now computes the documented 0.5/sigmaR^2 sum_t |L v_t|^2. The 3D default sigmaR is unchanged for now. - The returned A, v (and 3D Xs) come from the same, last evaluated iteration; v used to be one update ahead of A. - muA/muB are matched to the target's channels (a 3-entry mean against a grayscale target used to count the distance three times). - 3D landmarks, which have no term, are still ignored, but now with a warning. tests/reference.py is a plain-autograd oracle of the fixed iteration; it reproduces the previous code bit for bit where no fix applies, and the engine matches it to ~1e-12.
A fused C++17 engine (STalign/_cpp, pybind11 bindings, std::thread pool, vendored BSD-3 pocketfft) runs the whole fitting loop with the GIL released: one pass over target pixels for the forward map and contrast terms, one for the gradient with thread-local splatting, and a fused spectral velocity update. It computes the same iteration as the torch engine (gradients within 1e-13 of tests/reference.py) and gives bitwise identical results for any thread count. It is fork-safe and shared between package and sys.path imports. backend='auto' uses it for CPU float64 problems and falls back to torch otherwise. setup.py builds it as an optional extension: without a compiler or pybind11, pip install still succeeds with a warning. pyproject.toml declares the build requirements and MANIFEST.in ships the sources in the sdist. Per iteration on an M1 Ultra (16 threads) it is 81-91x faster than the previous code in 2D and 39-45x in 3D: the MERFISH tutorial's 10,000 iterations take about 11 s instead of ~14 min, and the 3D atlas tutorial's 2,000 about 35 s instead of ~22 min.
requirements.txt (read by setup.py) pinned ~50 packages, mostly for the docs, and old numpy/torch versions. It now lists only what the package imports, unpinned: matplotlib, numpy, pandas, pynrrd, requests and torch. STalign.py drops unused imports (plotly, tornado, glob, copy, ...) and imports pandas, requests and pynrrd only inside the atlas helpers that use them. tests/test_api.py checks that every public function and parameter of the previous version remains, both import styles of the tutorials, the fallback from cpp to torch, and the draw modes.
README: how the optional engine is installed, backend_info(), backend and thread options, the in-place build for clones, and a summary of changes. CHANGES.md lists the speedups, new options, changes that affect results and bug fixes.
With the corrected 3D regularizer (Fourier transform over all three axes), the old default sigmaR=1e8 regularizes LDDMM_3D_to_slice much less than before: the rms velocity of the 3D tutorial fits became up to twice as large. DEFAULT_SIGMAR_3D is now 2e6, which gives deformations of about the previous size. With the default epT the affine translation never converged: it jumped back and forth by tens of um per iteration until the last one, so changing the initial T in its last digit moved the cells by 24-84 um. The new keyword-only argument affine_step_control (default True in LDDMM_3D_to_slice, False in LDDMM) scales the step of each entry of L and T by its own multiplier in [1e-3, 1], which halves when that parameter's gradient changes sign and grows by 5% per iteration otherwise (Rprop's sign rule), so steps never exceed epL and epT. The constants are module defaults in _lddmm_common; the state lives in the torch and C++ engines and persists across run() segments. It is off in 2D, where the plain steps already converge (MERFISH) or the control made the noisy H&E tutorial worse. Off, results are bitwise unchanged. On the two 3D tutorials the fits now reproduce the research runs behind these defaults (dev/research/3d) to rounding, a 1-ulp change of T moves the cells by ~1e-12 um, and the C++ and torch engines agree to ~1e-14. Matching energy, cells in brain, Dice and correlation with the atlas stay close to the previous version (CHANGES.md). tests/reference.py implements the control, and both engines are tested against it in 2D and 3D with oscillating affine steps, diffeo_start, landmarks and segmented runs, plus defaults, validation, cpp/torch agreement and thread-count invariance. The C++ binding API version is now 2. Also: test_raster.py and test_transforms.py find the frozen original through STALIGN_ORIGINAL like the other tests, and the torch engine's docstring points to CHANGES.md for the corrections.
Built and tested on Linux (Pop!_OS 22.04, gcc 11.4, torch 2.14.1+cpu): 200 tests pass after these two changes. - test_fft_pool_after_fork hung on Linux. The cause is torch, not STalign: Linux torch wheels use GNU OpenMP, which deadlocks the first multi-threaded torch op in a forked child once the parent has used threads (a plain elementwise op hangs the same way). The child now runs torch on one thread, as torch's DataLoader workers do; STalign's FFT pool is still called with 2 workers, so the test still catches a pool that is not re-created after fork. - Zero-initialize SplatAcc::acc to silence 8 false-positive gcc 11 -Wuninitialized warnings (acc is zeroed before every read); results are unchanged. README: the engine is now tested on macOS and Linux.
The kernel width of rasterize is given in pixels (standard deviation 2*blur pixels), so the same blur gives a wider kernel at a coarser pixel size. The new keyword-only blur_um gives the standard deviation in the units of the coordinates instead (blur = blur_um/(2*dx)), so a raster can be redone at another pixel size with the same physical kernel. This is what the re-rasterizing alignment workflow needs. blur now defaults to None, which means 1.0 unless blur_um is given; passing both raises ValueError. Calls without blur_um are unchanged bit for bit (checked against bec34e4 on synthetic and eval-suite cell sets), and blur_um gives exactly the result of the equivalent blur.
On the synthetic problems of the round 2 evaluation suite, the most accurate 2D workflow re-rasterizes both cell sets at 15 um with a 30 um kernel and fits LDDMM on a two level pyramid of mean-pooled copies of those images (4x for 500 iterations, then 2x for 500), each level starting from the previous level's affine and velocity field. Users rasterize their cells themselves today, usually at a coarser pixel size with a kernel tied to it. align_cells(xI, yI, xJ, yJ, dx, blur_um, levels, **lddmm_kwargs) does this with the public pieces: rasterize with blur_um, and one ordinary LDDMM call per level with LDDMM's own defaults, so anything LDDMM derives from the images is derived per level. The first level gets the user's initial state, prefit and diffeo_start; later levels get the previous A, v and xv (an initial velocity field, so LDDMM does not pre-fit again) and diffeo_start=0; with landmarks, sigmaP is scaled by the pooling factor; only the last level draws. It returns the last level's LDDMM result (so the transform functions work on it unchanged) plus the rasters and the pooled grids of each level. The matching weights start afresh at each level because LDDMM takes no initial weights. The research version (syn.py's W path) carries them over; with its shipped settings that changed the errors on eight eval problems by 0.1% (geometric mean). With carry-over disabled, the same settings and the same velocity grid (LDDMM builds it from the first level's source; syn.py from the full raster), align_cells reproduces syn.py's pyramid bit for bit.
CHANGES lists the two additions with the measured effect of the workflow: on the 27 held-out synthetic 2D problems of the round 2 evaluation suite (Modal runs cells_synW_holdout_modal vs ship_core_holdout_modal), re-rasterizing the cells at 15 um and fitting the two level pyramid had a geometric mean error of 5.2 um instead of 11.9 um for LDDMM with the round 2 defaults on the suite's rasters, better on 24 problems, in half the fitting time. It notes the two limits seen there: problems already rasterized at 15 um got slightly worse, and partial tissue remains hard. The README points users with two sets of cells to the new function.
bec34e4 zero-initialized SplatAcc::acc to silence a false-positive -Wmaybe-uninitialized from GCC (acc is always zeroed before it is read). With clang the initializer changes the code generated for pass 2 and costs about 4% per 2D iteration: 2.48 vs 2.38 ms/iter on the MERFISH bench input (342x326, 4 threads, Apple M1 Ultra, timing lock). Keep the initializer for GCC only. GCC 16 compiles the engine without warnings and gets the same code as before; clang gets the code of 3a5f88b back. Results are unchanged (bitwise, both backends).
LDDMM's landmark term moved pointsI by forward Euler steps,
p <- p + v_t(p)/nt, which invert the fitted backward map only up to a
time discretization error. Since transform_points_source_to_target uses
the exact inverse, the landmark energy measured a different map from
the one applied to the points after the fit.
Each step now solves p_{t+1} = p_t + v_t(p_{t+1})/nt: fixed-point
iteration from the Euler step (until no landmark moves by more than
1e-10 of a velocity grid cell, at most 100 sweeps), then one Newton
step, so the result does not depend on the tolerance. The gradient
uses the implicit function theorem with the same point and Jacobian.
Both engines and the oracle (tests/reference.py) implement it. It is
the new default, Problem.landmark_map = 'exact'; 'euler' keeps the old
map. Solves that do not converge are counted in the engine snapshot,
and the fit warns.
Validation: gradients match the oracle to ~1e-12 and the engines agree
to 1e-9 over 150 iterations. On the four real landmark cases (Modal)
both backends reproduce the research run syn_lm_torch_exact_modal:
same landmark errors, maps within 0.003 um, e.g. merfish_visium_points
in-fit 67.8 um (Euler 73.8) and held-out 265.1 um (Euler 270.7).
xenium_heimage differs by 6-8 um, within its 1-ulp noise floor (8.7
um). With landmark_map='euler' results are bitwise equal to bec34e4 in
2D and 3D, on both backends.
The oracle-vs-original test uses landmark_map='euler'. The affine
step-control fixture uses epL=5e-5: with 1e-4 and the exact map, L
oscillates with growing amplitude for ~20 iterations, which amplifies
rounding differences ~1e4-fold. The C++ binding API version is now 3.
New keyword-only argument mixture=None of LDDMM. 'em' runs the mixture-model lens's working EM, natively in both engines and in the oracle. WM starts at 1. From iteration 50, every 5 iterations, a log-space E-step weighs each target pixel between the matching class and a background class. The matching variance is estimated from the residual. The background mean is the median of the target's border (or the user's muB) and is kept fixed. The background variance is (sigmaB/sigmaM)^2 times the matching variance, the matching prior is kept >= 0.5, and there is no artifact class. A dict overrides the options start, every, classes, pi_min, pi_init, fix_muA and fix_muB. mixture=None leaves the original weight updates unchanged (bitwise). LDDMM_3D_to_slice takes the argument but raises NotImplementedError, since the EM was not studied in 3D. The research prototype ran the E-step as a Python hook between engine segments, which cost +35-40% per iteration. Native, it costs nothing measurable (MERFISH bench input, 4 threads, timing lock): C++ 2.39 -> 2.37 ms/iter (the EM skips the original's class-mean sums), torch +1.6%. Validation: both engines match the oracle over 200 iterations (~1e-9) and each other, with results independent of the thread count and of how a run is split into segments. Parity with the research implementation (mixture/mixcand.py with its mixcpp engine and syn.py's MIX_V3 options), on 5 eval problems including 3 partial-tissue ones, 5000 iterations: the maps are identical (0.0 um) on 4 of them. On xen_m150_dx30_partial the energies agree through iteration 80 and the maps end 51 um apart. That problem is chaotic: a 1-ulp change of T moves its map by 43 um. The C++ binding API version is now 4.
A new private module with the data-driven settings of round 2, ported
from the research reference (synthesis/syn_search.py and syn_scale.py)
with the same math:
- affine_prefit: a coarse multi-start affine pre-alignment. 2D: the
initial affine rotated about the target tissue centroid in 12 steps,
each start fitted by batched Adam (100 iterations) on images pooled
to about 80 pixels across, the best one refined (60 iterations) at
twice that resolution; with landmarks a single start with STalign's
landmark term. 3D: the initial plane tilted by {-6, 0, 6} degrees
about both in-plane axes and shifted by {-300, 0, 300} um (27 starts,
200 + 100 iterations). The match is LDDMM's contrast-fitted data term
with background weight 0.1 and a linearly varying intensity bias.
Ties within 1% go to the start closest to the initial affine. There
is no randomness.
- physical_sigmaR: sigmaR for a prior displacement scale in um,
sigma_disp sqrt(C_d pi) a^(d/2) sqrt(nt) (dx/30) sigmaM / s_J.
- curvature_steps: steps scaled by sigmaM^2 * 1.1 / G, G a curvature
proxy of the 2D image term.
On 12 eval-suite problems the pre-fitted affine and sigmaR are bitwise
those of syn.py (with float64 tilt angles in 3D, where syn.py computes
them in float32), and the steps agree to 4e-16 (bilinear sampling by
grid_sample instead of scipy's map_coordinates: no scipy dependency).
New defaults of LDDMM and LDDMM_3D_to_slice (round 2, A1-A4), all in the engine-agnostic front-end (resolve_defaults in _lddmm_common, run between prepare_problem and the engine): - prefit='auto' (keyword-only): the coarse affine pre-fit of _prefit runs in 3D, and in 2D unless landmarks are given; never when an initial velocity field is given (it would move the affine under a fitted deformation, e.g. a warm start). True/False force it. It starts from the user's L/T/A; if it fails it warns and the fit starts from the given affine. - sigmaR=None (2D and 3D): the physical rule with the new keyword-only sigma_disp=94.03 um. An explicit sigmaR is used as is. - epL, epT, epV = None in LDDMM: the old defaults scaled to the curvature of the image term. Explicit values are used as is. 3D steps are unchanged. - niter=None in LDDMM: 2000. LDDMM_3D_to_slice keeps 5000. legacy_defaults=True (keyword-only, A6) restores the earlier fit bit for bit: no pre-fit for 'auto', niter 5000, the old 2D steps, sigmaR 5e5 (2D) / 2e6 (3D), and the forward Euler landmark map when the problem has a landmark_map field. LEGACY_DEFAULTS replaces DEFAULT_SIGMAR_3D. The result dict gains 'prefit' (what the pre-fit did) and 'settings' (the niter, sigmaR, epL, epT, epV used); passing them back with A=prefit['A'] and prefit=False repeats the fit exactly. run_lddmm refuses a problem whose settings are still None. Tests: test_round2_core.py covers the pre-fit (large 2D rotation, tilted 3D plane, ties, determinism, the 'auto' rules, failure), the sigmaR and step rules, the settings round trip, cpp/torch agreement on the default path, and legacy_defaults (against the explicit old values, and against commit bec34e4's package when STALIGN_LEGACY_SRC is set). The oracle comparisons in test_lddmm_torch.py run with prefit=False, test_api's 3D problem too, and the changed defaults are listed.
A "Round 2" section in CHANGES.md: the affine pre-fit, the sigmaR rule with sigma_disp, the curvature-scaled 2D steps, niter 2000 in 2D, the legacy_defaults switch and the new 'prefit' and 'settings' entries, with the measured effect on the evaluation suite's held-out problems (paired Modal runs of the shipped configuration against 3a5f88b: 2D error ratio 0.43, 3D 0.21) and the partial-tissue caveat (5 of 8 partial-section problems got worse). The 3D sigmaR=2e6 default of the earlier entry is marked as superseded; README gets a short note.
curvature_steps divided by max(G, 1e-300), so a constant target (G = 0) got steps of order 1e300 and the fit blew up. Without a finite positive curvature the steps are now returned unchanged.
'auto' skipped the pre-fit whenever v was passed. A zero v (passed with xv to choose the velocity grid) carries no fitted deformation, so the pre-fit is harmless there; only a nonzero v (a warm start) now turns it off. resolve_defaults reads this from the problem instead of an initial_velocity argument.
The overview of _lddmm_common now names the step between prepare_problem and run_lddmm that runs the pre-fit and fills in the data-driven settings.
The test module's docstring gives the commands that export and build commit bec34e4 for STALIGN_LEGACY_SRC.
The core branch set problem.landmark_map = 'euler' only when the Problem had that field, because the exact landmark map came from a parallel branch. Both are merged now, so legacy_defaults=True sets it directly. The bec34e4 comparison that covers the landmark case only runs when STALIGN_LEGACY_SRC is set, so test_legacy_defaults_are_the_old_values now also checks, on every run, that a fit with landmarks under legacy_defaults=True equals one with the old values and the Euler map on both backends, and that the default (exact) map gives a different fit.
align_cells relied on LDDMM's prefit='auto' skipping the pre-fit when it is given an initial velocity field. The merged rule skips it only for a nonzero v, so after a level of affine-only iterations (diffeo_start at least that level's niter) the next level ran the pre-fit again from the warm start: with levels ((4, 20), (2, 20)) and diffeo_start=100 the second level reported prefit ran=True. Later levels now pass prefit=False. With the default levels the previous level's v is nonzero, so 'auto' already skipped the pre-fit and the results are unchanged.
The three branches documented their parts separately: the core branch wrote the Round 2 section, the cells branch added align_cells and blur_um under "New options", and the engines branch left its notes to the merge. CHANGES.md now has one Round 2 section. It adds the exact landmark map and the mixture='em' option, with their measured effects; the align_cells and blur_um entries move into it, the partial-tissue caveat points to the mixture option, and the legacy_defaults entry lists the Euler landmark map. The README's change list and the LDDMM docstring (sigmaP, mixture, 'settings') say the same. The landmark errors come from the merged package on the full evaluation suite (results/r2_pkg_full_modal) against the core branch's run with the same defaults and the Euler map (r2core_pkg_full_modal): only the four landmark cases change. The mixture numbers are the verification's marginal of the EM in the research configuration (partial sections 0.25, 8 of 8 better; density problems 1.09; the rest 1.01).
The oracle's docstring and SIGMAR_3D comment still called 2e6 the package's default. Since Round 2 the package chooses sigmaR from the data, and 2e6 is what legacy_defaults=True gives; the oracle has no pre-fit and no data-driven settings, so comparisons with it pass prefit=False and explicit settings.
The exact landmark map finishes each implicit step with a Newton step that solves M = I - Dv_t(z)/nt. When M is exactly singular, the torch engine raised from torch.linalg.solve and the C++ engine divided by a zero pivot, so the fit stopped or went non-finite. Both engines (and the oracle) now test det M = M00 M11 - M01 M10 with the same formula and, where it is 0 or NaN, take the forward Euler step p_t + v_t(p_t)/nt for that landmark and time step instead, with the Euler step's gradient. Such steps are counted with the solves that did not converge, so the existing RuntimeWarning reports them; its text now mentions a singular Jacobian. Non-singular steps run the same operations as before, so fits are unchanged. The new test builds a velocity field with Dv/nt = diag(1, 0) exactly at a landmark and checks both engines against the oracle (with the image term switched off), the count and the warning.
The engines' gradients of the exact landmark map were checked only against the oracle, which differentiates the same implicit-function construction. The new test checks dEP/dL, dEP/dT and the derivative along a smooth direction of v against central finite differences of the landmark energy on both engines (the image term is switched off with WM = 0). The errors are about 1e-10 (L, T) and 3e-9 (v) of the gradient's size, against a tolerance of 1e-6; the Euler map's gradient differs from the finite difference by 3%, so the test tells the two maps apart.
curvature_steps keeps the given step sizes when the target has no curvature, but curvature_proxy took a channel's variance from np.var, which is a tiny positive number for most constants (np.var of an image of 0.3 is about 1e-34, from the rounded mean). On such a target the 2D steps were multiplied by about 6e33. A target channel now adds variance only if it is not constant (np.ptp > 0), so a constant target gives G = 0 and the steps are kept. Likewise a source channel that is constant in I is skipped, as the docstring says, even where its bilinear samples differ in the last bit. Images without constant channels give the same steps as before.
On a constant target the pre-fit's tissue weights are all zero, so the tissue centroid was 0/0: numpy emitted four RuntimeWarnings, the search failed with "attempt to get argmin of an empty sequence" and the generic failure warning followed. _prefit.has_tissue checks the target, pooled as the coarse search pools it, before the pre-fit runs. Without tissue the pre-fit is skipped (also with prefit=True) with one warning that says why, 'prefit' reports the reason 'no target tissue', and the fit starts from the given affine. Targets with tissue take the same path as before.
* The 2D affine step control: the Xenium-to-H&E landmark errors 34.2 and 30.3 are in pixels of the H&E image, not um. The harness compares against pointsI, which are H&E pixel coordinates (dev/bench/cases.py, _xenium_heimage; dev/research/3d/tables/ twod_xenium_heimage.md; dev/integration/integrate_output.json says "TRE is 34.2 vs 30.3 px"). In um they are about 106 um, as the Round 2 bullets report. * The 3D table's "this version" rows came from the fits with sigmaR=2e6 and no pre-fit (dev/research/3d/tables/final.md), which are what legacy_defaults=True now gives: on all 41 problems of the evaluation suite, including both 3D tutorials, its fit files are bitwise identical to those of 3a5f88b (dev/round2_impl/validate/identical_r2val_legacy_vs_review_3a5f88b_full.md). The rows are relabelled, and the Round 2 defaults' own values are added from the verified run of b87713c (dev/research/round2/eval/results/vfy_full_modal/summary.md): MERFISH 97.3% in brain, Dice 0.885, correlation 0.632; STARmap 98.8%, 0.806, 0.666. The 3a5f88b column of that comparison reproduces the old rows exactly, so the metrics are the same. * The velocity-curvature cap's cost was given per million target pixels; it grows with the velocity grid instead. The verification of b87713c measured 1.9 s at a=500 and 14.6 s at a=200 on a 679 x 841 target with 25 um pixels, 2 threads (dev/round2_impl/steprule_journal.jsonl, last record).
The changelog was an "Unreleased" section assembled from separate branches. It is now one document for 2.0.0, dated 2026-10-14, in the order speed, Round 2 defaults, new functions and options, corrections, bug fixes, tutorials, packaging and testing. * Faster comes first and says that the tutorials pass their own niter, so the timings do not include the 2000-iteration default; the Round 2 overhead (up to 6 s / 12 s) is the pre-fit and step estimation (dev/round2_impl/bench_final/README.md, summary.json). * The pre-fit times are "about 1-4 s" in 2D: the Mac's range is 0.8-3.8 s at 4 or 16 threads (bench_final/summary.json, prefit_s). * LDDMM_3D_to_slice's epV default is now None (1e3, capped by the regularizer bound), in its own bullet and in a list of all changed defaults, which matches CHANGED_DEFAULTS in tests/test_api.py. legacy_defaults=True in 3D gives epV=1e3 without the cap (STalign/_lddmm_common.py, LEGACY_DEFAULTS and resolve_defaults). * align_cells' losses where the suite's rasters had 15 um pixels are 10-14% on the held-out set and 5-14% over all three sets (dx=15 ratios 1.119/1.136/1.100, 1.137/1.051/1.107 and 1.112/1.126/1.068; dev/round2_impl/steprule_journal.jsonl). * blur_um moves to the options; the corrections section says its effects were measured with the pre-Round 2 defaults, and the 3D sigmaR history is told from version 1.0 (1e8) instead of from the intermediate 2e6 default that was never released. * LDDMM_3D_to_slice takes mixture but raises NotImplementedError for anything but None; the text said it "does not have it". * Tutorials: from the notebooks branch. The fit cells take 0.6-35 s (579 ms to 34.5 s in the executed notebooks' %%time outputs; the curve annotator 3 min 43 s); the medians 47.5, 14.1, 98.2 and 69.3 um and the simulation's RMSE 6.383 um are printed in the notebooks; the start-plane offsets 252.8 and 47.2 um follow from slices 177 and 140 of the 50 um CCFv3 with scale_z 0.9. The old notebooks' recorded times, from other machines, are not compared; the Faster table is the controlled comparison. The simulation's 6.68 um (exact inverse alone, dev/integration/E2E.md) and 6.38 um (this version) are both stated. * Packaging and testing: the version, the pocketfft license in wheels, the GitHub Actions workflow, the documentation workflow's actions, docs/requirements.txt, the Pipfile and the setuptools note.
The README said align_cells had 10-14% more error than LDDMM where the suite's rasters had 15 um pixels, speaking of "our tests" in general, but 10-14% holds for the held-out set only. Over the main, held-out and further held-out sets the dx=15 ratios are 1.137/1.051/1.107, 1.119/1.136/1.100 and 1.112/1.126/1.068, so 5-14% (verification of b87713c, dev/round2_impl/steprule_journal.jsonl). CHANGES gives both.
* Installation: one section with the pip line, the tutorials' extra packages and the import; the C++ engine is a subsection that says the compiler is optional, shows backend_info() and the threads setting in one paragraph, and keeps the in-place build for clones. The engine was also built and tested with gcc 13 (Ubuntu 24.04, by the CI simulation of the pr-ci branch), so it says gcc 11 and 13. * The pipenv instructions are condensed to their two commands and the note about Pipfile.lock. * "What changed in version 2.0.0" keeps the five changes most likely to matter for existing scripts, ordered by impact, and points to CHANGES.md for the rest. The 2D time halving is left to CHANGES: next to the Performance table, whose tutorials pass their own niter, it would read as a contradiction. * The Performance section is unchanged.
README and CHANGES said the Round 2 defaults cut the alignment error to 0.43 (2D) and 0.21 (3D) times "the previous one", and that align_cells had 0.17 times the error of "the previous version". Every one of these comparisons was made against this version with legacy_defaults=True (the pre-Round 2 package, which already has the exact inverse, the corrected 3D regularizer, 3D sigmaR=2e6 and the 3D affine step control), not against version 1.0, which was not run on the test problems. Readers would take "the previous version" to mean 1.0, and CHANGES also used the phrase for 1.0 itself in the speed and 3D tables. Say legacy_defaults=True wherever the baseline is the old defaults, state once what that baseline is, and say "version 1.0" wherever the baseline is the original code (the speed tables, the 3D atlas table and the rounding paragraph). No number changes.
- align_cells was not less accurate on every problem whose suite rasters had 15 um pixels: it was 5-14% worse on 9 of these 11, but had 0.42-0.44 times LDDMM's error on the 2 with a wider kernel (mer_m150_dx15_serial_b2 and its held-out twin). - Of the two partial heart sections that the step caps "improved", hrt_m50_dx60_partial started at 42.5 um and fails with or without the caps (105.8 and 84.1 um); say that the error fell, and that. - The 142 of 351 stress fits that were unstable had the uncapped Round 2 steps (scaled by G alone), not version 1.0's steps; say so, and call the steps the caps lower "uncapped" instead of "old".
The wheel ships STalign/_cpp/third_party/pocketfft/LICENSE.md as the license of the pocketfft code compiled into the C++ engine, but that file named only "Copyright (C) 2010-2018 Max-Planck-Society". The vendored pocketfft_hdronly.h, which is what the engine compiles in, names more holders (Max-Planck-Society 2010-2024, Peter Bell, Matteo Frigo and MIT, Tan Ping Liang, Cris Luengo), and BSD-3-Clause asks binary distributions to reproduce that notice. LICENSE.md is now the header's notice verbatim; the conditions and disclaimer are unchanged.
The tutorial had a cell under the heading "TO DO: remove section below before uploading" that ran importlib.reload(STalign). Its output was the module's repr, which showed the local path of the checkout that ran the notebook (the author's machine on main, a development worktree here) and would also be published on the documentation site. The reload has no effect on the tutorial, so the heading and the cell are removed. The notebook was re-run with this version; apart from the two cells, only object addresses and timings differ in its outputs, and every figure is identical.
The heart tutorial saved tpointsJ, the cells of CN73_E2 (the target) mapped into the coordinates of CN73_E1 (the source), as CN73_E1_to_CN73_E2.csv.gz, a name that reads as the opposite direction. It now writes CN73_E2_STalign_to_CN73_E1.csv.gz, like the other tutorials' "<moved>_STalign_to_<reference>" files, and a comment in the save cell says that the rows are (y, x) positions. The file is new in this version (the save cell never ran before), so nothing reads the old name. Its content is unchanged (same 11,770 rows); the notebook was re-run, and only object addresses and timings differ in its outputs.
The point-annotator tutorial thresholds the matching weights (WM_values > 0.95) into results_filtered, plots it, and then saved results, all 85,958 cells, as ..._with_point_annotator_filtered.csv.gz. It now saves results_filtered: the same rows and columns as the cells with WM_values > 0.95 in the previous file (38,657 cells). This bug is older than this version, but this version regenerates the file. The same re-run includes the numpy 2 cleanup of this notebook: the affine is converted with A.cpu().numpy() instead of np.array(A.cpu()), which warned that torch's __array__ takes no copy argument (the warning also printed the local temporary directory). Apart from the warning, only object addresses and timings differ in the outputs, and every figure is identical.
Five tutorials printed warnings, each with the local temporary directory of the kernel in its file name: - np.array(A.cpu()) and np.array(tpointsI[:,1]) (affine only with points, simulation): numpy 2 warns that torch's __array__ takes no copy argument. They now call .numpy(). - torch.sum(WA*J, ...) with a numpy J (MERFISH to Visium, with and without the curve annotator): the product went through numpy and torch's __array_wrap__, which numpy 2 deprecates. J and I are now converted with torch.as_tensor; muA and muB are bit for bit the same (float64, as before). - ax.scatter(..., vmin=0, vmax=1) without c (Xenium to Xenium): matplotlib warns that it ignores vmin and vmax. They are removed. The five notebooks were re-run. They have no stderr output now; apart from the warnings, only object addresses and timings differ, and every figure is identical. The CSV files they write are left as committed: their content is the same, except the affine-only-with-points file, which differs by under 1e-10 um, as in the previous regeneration.
- Name the heart tutorial's result file and say what it holds. - Mention the point annotator's filtered file, the tutorials that now run without warnings and the removed reload cell of the L/T tutorial. - Say that the uncompressed CSVs in docs/xenium_data and docs/starmap_data come from earlier versions of the tutorials and do not match this version's results (they are left in place). - Update the range of the tutorials' alignment times to the re-run notebooks: 0.5-36 s, and 3.8 min for the curve annotator.
On main and so far on this branch, STalign/__init__.py exported
nothing, so `import STalign; STalign.LDDMM(...)` worked only when
STalign.py itself was imported as a top-level module. The tutorials
that offer only "OPTION B" (sys.path.append("../../STalign"); import
STalign) do that from a clone, but with STalign installed the package
comes first on sys.path and the tutorials fail at their first
STalign.<function> call. Five tutorials, including
merfish-xenium-alignment since this version, use only that import.
The package now star-imports STalign.STalign, so both imports give the
same functions; `from STalign import STalign` is unchanged. A new test
imports the package the way an installed copy is found first and
checks that it has every function of the module.
The docstrings of legacy_defaults (LDDMM, LDDMM_3D_to_slice, resolve_defaults) cited "commit bec34e4", which upstream will not have if the pull request is squashed, and "Round 2", and align_cells cited "the round 2 evaluation suite". They now say what the options do in terms of version 1.0 and point to CHANGES.md for the measurements. The 3D legacy_defaults docstring also names epV=1e3, which LEGACY_DEFAULTS sets. Docstrings only.
The tests that compare the rasterization, the transforms, the plain-autograd oracle and the function signatures with version 1.0's STalign.py skipped in CI, because they looked for it only in the untracked dev/ directory or in STALIGN_ORIGINAL; CHANGES says the suite makes these comparisons. Both jobs of tests.yaml now fetch commit b2068ed (version 1.0, JEFworks-Lab/STalign main) and point STALIGN_ORIGINAL at its STalign/STalign.py. Locally, with that file, the full suite passes (274 passed; the one skip needs the Allen data), and the comparisons add about 4 s. The test docstrings now say where the frozen original comes from (STALIGN_ORIGINAL, e.g. git show b2068ed:STalign/STalign.py, or dev/ref/stalign_orig.py), and the workflow's header comment says which tests still skip (the development data and STALIGN_LEGACY_SRC).
The documentation build reported docutils errors in three docstrings:
the regularizer formulas of LDDMM and LDDMM_3D_to_slice contain
|L v_t| and |Lv|, which reStructuredText reads as substitution
references, and align_cells' Returns section nested a definition list
without blank lines ("Unexpected indentation"). The formulas are now
literals, and the two extra entries of align_cells' result form a
proper definition list. Docstrings only.
The Sphinx build succeeded with 501 warnings and 5 docutils errors: - The numpy-style docstrings were rendered without napoleon, so their Parameters sections were parsed as plain reStructuredText (errors in LDDMM, LDDMM_3D_to_slice, align_cells and analyze3Dalign, whose X_ and Y_ read as link targets). conf.py now loads sphinx.ext.napoleon. - 442 warnings said that the notebooks' 'ipython3' lexer is unknown: ipython was not installed. docs/requirements.txt now lists it. - The README, included as the overview page, links to CHANGES.md and to its own headings, which myst could not resolve. docs/CHANGES.md includes CHANGES.md as a page of the site (in the index toctree), and myst_heading_anchors = 3 gives the headings their GitHub anchors. The documentation workflow installs docs/requirements.txt instead of its own list, which also installed "collection", an unrelated PyPI package (the standard library's collections was meant), and the CUDA build of torch; it now installs the CPU build of torch first. It also uses actions/checkout@v7, like tests.yaml. A build of the docs with these changes (Python 3.10, sphinx 8.1.3, nbsphinx 0.9.8, myst-parser 4.0.1) succeeds with 32 warnings and no errors. As on main, the ones left come from tutorials.rst (section titles inside its toctree, and 9 notebooks that it does not list), from the README, whose headings start at level 2, and from the missing _static folder.
Version 1.0's STalign.py imports plotly and tornado at the top, which this version's requirements no longer include, so in a CI job the seven tests that load it failed with ModuleNotFoundError (found by running the job's steps in a fresh Python 3.14 venv on macOS). The step that fetches it now installs both. With them, the job's full test suite passes there: 272 passed, 3 skipped (the Allen data and STALIGN_LEGACY_SRC).
The test matrix had Python 3.10, whose support ends this month, and 3.12, so no job tested a current release. It now also runs 3.14 on Linux and macOS. Its steps were run on macOS in a fresh Python 3.14.7 venv (torch 2.14.1): pip install -v . compiles the C++ engine (backend 'cpp'), build_ext --inplace works, and the full suite passes with version 1.0's STalign.py (272 passed, 3 skipped).
- import STalign gives the functions of STalign.STalign. - The test suite compares with version 1.0's STalign.py, which the CI workflow takes from git, and the workflow tests Python 3.10, 3.12 and 3.14. - The documentation workflow's checkout and CPU torch, napoleon, the CHANGES page and ipython in docs/requirements.txt.
The heart, Xenium-to-Xenium, Xenium-to-H&E and Xenium-to-STARmap tutorials wrote bare numeric arrays with np.savetxt. They now write the target's own columns plus aligned_x and aligned_y, like the MERFISH tutorials, and were re-executed. Also clear the stale output that an empty cell of the MERFISH-to-Xenium tutorial kept from version 1.0.
CHANGES.md says "the new defaults" instead of the development label "Round 2", and describes the tutorials' new result tables. The test file for the defaults is renamed test_defaults.py; the reference to the last commit before the new defaults names its subject as well as its hash.
Current torch and numpy need Python 3.10 or later, and CI tests 3.10, 3.12 and 3.14. The README's pip command installed upstream's version 1.0; it now installs this fork, which the README describes.
On GitHub's hosted macOS runners MPS is available, but the virtual GPU's
Metal compiler rejects one of torch's matmul kernels ("Failed to create
pipeline state object"), so the test failed on all three macOS jobs
with everything else passing. It now skips on that compiler error only;
any other error, or a wrong result, still fails it, and it still runs
on real Apple GPUs.
The fork's README pointed pip at slowkow/STalign; upstream installs from its own repository.
Contributor
|
I have confirmed that:
Additional questions/requests:
|
The align_cells sections of the MERFISH-to-MERFISH, Xenium-to-Xenium, Xenium-to-STARmap and heart tutorials were headed "Recommended". They are now "Alternative", and say that align_cells chooses its settings for you while the walkthrough shows what each one does, so students are pointed back to the parameters. The comparisons no longer say that accuracy cannot be judged without a ground truth: gene expression can be used to evaluate an alignment. Markdown only; no outputs changed.
Contributor
Author
|
Thank you for testing this so thoroughly, and for confirming the installation, the agreement with previous results and the behavior on new data! Data files in In the files that changed by a few KB (the MERFISH A few output files changed more:
If you would prefer to keep the files in
Comparison statement. Done: I removed "Without a ground truth we cannot tell which is more accurate here." from both notebooks. The tutorial datasets in the repository contain cell positions and metadata but no gene expression matrices, so I didn't add an expression-based comparison. I'd be glad to add one if you'd like. These changes are markdown only, so no outputs changed. Thanks again! |
Contributor
|
Thanks for the clarification and contribution. Merging pull request. TODO: check if gh-pages auto-recompiles and deploys appropriately post merge |
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
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.
Hi! Thank you for STalign. While trying to make it faster, we ended up with a larger set of changes than we planned. We hope they are useful upstream. Everything is explained in
CHANGES.md. This PR is large, and we are happy to split it into smaller ones (for example speed, bug fixes, new defaults, tutorials) if that is easier to review.Faster
LDDMMandLDDMM_3D_to_sliceget an optional C++ engine (pybind11, no OpenMP, vendored BSD-3 pocketfft).pipbuilds it when a C++17 compiler is available. Without one, installation still succeeds, and a new pure-torchimplementation computes the same iterations. Both engines are tested against a plain-autograd reference implementation of the original iteration: gradients agree to within 1e-9, and the reference reproduces version 1.0 bit for bit where no fix applies.The 10 tutorial alignments, with their own parameters and iteration counts:
torchengineFor example, the MERFISH-to-MERFISH tutorial goes from 9.2 min to 10 s, and MERFISH to the Allen atlas from 21.6 min to 34 s.
rasterizeis exact and about 50× faster: 86,000 cells take 31 ms instead of 1.8 s.Bug fixes
These are listed in
CHANGES.mdwith their measured effects. The main ones:LDDMM_3D_to_slicetook the FFT of v over only 2 of its 3 axes, so the penalty had no smoothness along one axis and was position-dependent.transform_points_source_to_targetused a forward Euler map that is not the inverse of the fitted map. It now inverts it exactly;exact_inverse=Falsekeeps the old behavior. The simulation tutorial's ground-truth error drops from 9.68 to 6.68 µm.transform_points_target_to_sourceintegrated time steps in the opposite order tobuild_transform.analyze3Daligncoordinates were wrong wheneverdx≠ 10: 4.3 mm off in the STARmap tutorial.rasterizePCA(toolsundefined);build_transform3D(direction='f');normalizeon a constant array.LDDMMreturnedAandvfrom different iterations.More accurate defaults (opt out with
legacy_defaults=True)We built a ground-truth benchmark: real MERFISH, Xenium, heart and Allen-atlas layouts, warped by known deformations at several pixel sizes. We kept held-out problems that were scored only after the configuration was frozen. Against
legacy_defaults=Trueon the held-out problems:What changed, and why:
prefit='auto'). Unknown rotations and 3D slice tilts were the largest source of error.sigmaRderived from a displacement scale in µm (sigma_disp). The bestsigmaRmoved with pixel size (∝ 1/dx), because the matching term is a pixel sum while the regularizer is an integral. An explicitsigmaRis still used as given.epL/epT/epVare still honored.legacy_defaults=Truereproduces the previous defaults bit for bit. Every 1.0 function and argument still exists, and new arguments are keyword-only.New
align_cells(xI, yI, xJ, yJ): rasterizes two sets of cells at 15 µm and fits coarse to fine. In our tests its error is 0.39× that ofLDDMM.mixture='em': a working background/tissue model for targets with missing tissue. With the default sigmas, the old one had near-zero background weights.rasterize(blur_um=...),backend_info(), and thedraw,backendandnum_threadsoptions.drawalso covers Opt out of plots and error progresion in output #49.Tutorials, packaging and CI
align_cellssection.requirements.txtlists only the 6 runtime dependencies; it previously pinned about 50 packages, mostly for the docs.Note that re-executing the notebooks and regenerating the CSVs adds about 395 MB of new file versions; a squash merge keeps
main's history small. We'd be glad to answer questions or adjust anything.