Optimize DecohesionTransform using F.conv1d (#201)

[JGiraldo29]

Performance fix for the DecohesionTransform introduced in #201. It uses the same math and has the same output while being faster on both CPU and GPU.

What changed

In batchgeneratorsv2/transforms/spatial/decohesion.py the Python for loop in _causal_smear has been replaced with a single F.conv1d call after permuting the smear axis to the innermost position.

How it works

The original _causal_smear uses a Python loop to accumulate shifted slices. This is mathematically a causal 1-D convolution, but written as a loop it incurs one Python iteration and one kernel dispatch per tap.
This PR replaces the loop with a single F.conv1d call:

• The target smear axis is swapped to the innermost (last) position before the convolution.
• The axis is restored to its original position immediately after.

This permutation step is required due to how cuDNN handles 3D convolutions. The backend is highly sensitive to kernel orientation:

• Applying a kernel along an outer spatial dimension (such as 7x1x1) forces cuDNN into a slower, non-coalesced execution path.
• Applying the kernel along the innermost dimension ensures the operation hits a fully optimized, vectorized path.

Why it's useful

The changes implemented here aim to speed up training when decohesion is in the augmentation pipeline. The reduction is large on GPU (where the original loop incurs K-1 kernel launches per call) and significant on modern CPUs.

Evidence

Equivalence

Bit-equivalent to the original within float32 summation roundoff. Tested across 22 synthetic cases (11 CPU + 11 CUDA) covering 2D/3D shapes, every smear axis, kernel widths from 1 to 16, and channel counts from 1 to 4. Max absolute error across all cases: 5e-7 (float32 epsilon).

Verified on real scroll data: a (160, 160, 160) crop from the PHerc172 OME-Zarr (53keV_7.91um.zarr, scale 0) produces max|old - new| = 1.19e-07.

The existing __main__ self-test in decohesion.py still passes all five structural assertions (identity at p=0, segmentation untouched, one-sided trailing-vs-leading smear, density modulation, finite output) with identical numerical output to the original.

Speed (CUDA — Tesla T4)

shape	K	smear axis	old (ms)	new (ms)	speedup
(1, 160, 160, 160)	7	1	2.95	1.32	2.23×
(1, 160, 160, 160)	7	3	3.10	0.43	7.23×
(1, 160, 160, 160)	13	1	5.71	1.29	4.42×
(4, 128, 128, 128)	7	1	5.93	2.35	2.52×
(1, 512, 512)	13	—	0.83	0.18	4.62×

Speed (CPU — Intel Core Ultra 7 255H)

shape	K	old (ms)	new (ms)	speedup
(1, 160, 160, 160)	7	94.0	49.4	1.90×
(1, 160, 160, 160)	13	182.2	56.4	3.23×
(1, 512, 512)	13	147.1	25.3	5.81×

Decohesion old vs new

Decohesion old minus new

Details

• Same taxis contract as the original (params['axis'] + 1, indexing into a (C, *spatial) tensor).
• Works for 2D (C, X, Y) and 3D (C, X, Y, Z) tensors.
• K=1 short-circuit preserved.
• Numerical drift is summation-order roundoff only (the conv accumulator and the loop accumulator sum the same terms in different orders).

How to reproduce

Two scripts are included:

python test_causal_smear_rewrite.py     # 22-case correctness + benchmark table
python validate_on_real_data.py         # PHerc172 crop, before/after/diff figures

The second script streams a small crop from dl.ash2txt.org and generates the figures attached to this PR.

test_causal_smear_rewrite.py
validate_on_real_data.py

Scope

This PR only touches _causal_smear. The Python-loop accumulators in _density_weight and the avg-pool path are already vectorized; no change needed there. Subsequent PRs can address other per-transform optimizations (e.g. the device-sync .item() calls in SqueezeTransform.get_parameters).

Addresses the optimization side of the Decohesion item of #201.

---

The use of LLMs was limited for REVIEWING and TESTING, as well as assisting in some parts of the changes. No agentic coding was used for this.

[vercel]

@JGiraldo29 is attempting to deploy a commit to the scroll Team on Vercel.

A member of the Team first needs to authorize it.