Feat: training experiments

[carshadi]

Summary

This PR changes the project from a live, per-patch training pipeline into a two-stage system with reproducible cached inputs and cache ordering:

1. An offline precompute stage reads cloud data, samples patches, runs BM4D, builds annotation-derived foreground masks, and writes count-space patch caches.
2. Training consumes every cached patch once per epoch, applies one frozen intensity mapping, and selects checkpoints using both signal fidelity and compressibility.

The normalization contract is now explicit and shared by dataset generation, training, validation, and inference. Bright intensities are no longer flattened by a per-patch percentile clip, and raw-volume background offsets are handled without changing the mapping learned during training.

Area	main	feat-training-experiments
Dataset generation	Cloud reads and BM4D inside Dataset.__getitem__	Required offline train/validation caches containing raw counts, BM4D teacher, and foreground mask
Foreground definition	Used mainly to bias patch sampling	Segmentation plus traced skeleton defines both the optional preserved target and count-space validation regions
Normalization	Per-patch percentiles and a hard upper clip; separate inference behavior	One frozen, serialized transform used across train/validation/inference, with invertible HDR options
Background pedestal	No per-acquisition alignment	Per-brain offsets are subtracted before caching; inference composes the runtime offset around the frozen transform
Training target/loss	BM4D everywhere with uniform L1	Optional raw-on-foreground target plus foreground-weighted Charbonnier loss; both are disabled in the current experiment configuration
Model	Direct-output U-Net with BatchNorm	Residual-output U-Net with GroupNorm and serialized constructor configuration
Epoch	A configured number of newly sampled examples	A deterministic shuffled pass over the complete training cache, including the final partial batch
Validation	Global L1 and median compression ratio; best L1 checkpoint only	Count-space signal/background/tail metrics plus compression-aware scoring; evaluated checkpoints after the first epoch are retained
Inference/output	Independent percentile normalization, hard clipping, high overlap-add memory use	Checkpoint-supplied transform, production offset support, lower-memory overlap-add, and metadata-preserving OME-Zarr output

End-to-end data flow

PRECOMPUTE
raw OME-Zarr + FFN segmentation + SWC traces + per-brain offsets
    -> foreground-biased candidate sampling
    -> reject spatially incoherent labeled artifacts
    -> subtract the brain pedestal
    -> run BM4D and build the annotation mask
    -> raw.npy + teacher.npy + fg.npy + transform.json + config.json

TRAIN / VALIDATE
memory-mapped caches
    -> shared frozen intensity transform
    -> target = BM4D, or raw-on-foreground/BM4D-on-background ablation
    -> residual 3D U-Net + Charbonnier objective
    -> invert predictions to count space
    -> fidelity metrics + compression-aware checkpoint score

INFERENCE
raw volume + precomputed full-tile pedestal
    -> subtract pedestal around the checkpoint's frozen transform
    -> tiled U-Net prediction and overlap-add
    -> inverse transform and restore pedestal
    -> uint16 volume, MIPs, compression results, and optional OME-Zarr

Dataset preparation and target construction

Patch generation is now an explicit offline stage

scripts/precompute.py replaces live data preparation with a single --split {train,val} workflow. Both splits use the same source, sampling, offset, BM4D, mask, and transform configuration; only their sampling stream and downstream dataset interface differ.

Each cache contains:

File	Type	Meaning
raw.npy	float32, (N, D, H, W)	Raw counts after subtracting the per-brain background pedestal
teacher.npy	float32, (N, D, H, W)	BM4D output clipped to the physical count range
fg.npy	uint8, (N, D, H, W)	Foreground mask
transform.json	JSON	Exact frozen transform to use later
config.json	JSON	Generation configuration: inputs, sampling thresholds, mask settings, offsets, BM4D sigma, seed, worker count, and cache dtype

The count arrays were deliberately changed from float16 to float32. float16 increasingly quantizes values across the uint16 sensor range and can lose multiple counts in the bright tail before either the transform or the model sees the sample. float32 exactly represents source uint16 counts and also preserves negative offset-corrected values and fractional BM4D estimates. The tradeoff is a larger cache: the current 30,000-patch training configuration is documented as roughly 71 GB.

Training now requires both caches and validates that all four required data/transform files exist and that train and validation use identical transform configurations. There is no live cloud/BM4D fallback in scripts/train_bm4dnet.py. This makes a training run self-contained and prevents a subtle train/validation mapping mismatch.

Foreground masks now come from annotations

The branch's first foreground-mask implementation combined an intensity mask with segmentation labels. The final mask is instead built from annotations:

• all positive FFN segmentation labels, with label dilation opt-in rather than automatic; and
• rasterized SWC nodes, dilated by a configurable neurite radius (currently two voxels).

The union protects traced neurites that the segmentation misses while avoiding the interim behavior of treating every unusually bright structure as signal. That intensity-based mask preserved bright noise, autofluorescence, and off-target label as though they were neurites. A robust median/MAD intensity mask remains only as a fallback for brains with neither segmentation nor SWC annotations.

This does not make the final mask immune to annotation errors. It trusts all positive FFN labels unless the patch-level coherence gate detects the specific spatially incoherent artifact described below; smooth mislabeled structures can still enter the foreground mask.

Validation uses the same completed annotation mask as training, rather than recomputing an intensity-only mask. This is important because the training target, foreground loss term, and validation metrics must agree on what "signal" means.

SWCs are assumed to be densely sampled in voxel space. Loading now warns when a parent-child link spans more than one voxel in Chebyshev distance, while correctly accepting one-step 3D diagonals. The code does not interpolate long edges silently; it rasterizes the supplied nodes.

Foreground-preserving targets are an ablation, not the current default run

The shared target builder supports:

target = np.where(fg_mask, raw, bm4d_teacher)

when preserve_foreground=True. The rationale is that BM4D can erase thin neurites, so the teacher should supervise denoising on the background without asking the model to reproduce that foreground attenuation.

This mechanism remains available independently from foreground loss weighting, but the current branch-tip experiment sets:

preserve_foreground = False
fg_weight = 0

Therefore, the configured run currently trains against BM4D everywhere with a uniform Charbonnier objective. This is an experiment-driven operating-point choice, not removal of the new infrastructure. Earlier runs found that a raw foreground target combined with strong foreground weighting rewarded an almost identity mapping: the output changed the input by fewer than a few counts and compressed approximately like raw data rather than like BM4D. The branch first reduced the weight and ultimately disabled both knobs in the current entrypoint. Because raw, teacher, and mask remain separate in the caches, either target or loss-weighting ablation can be changed without rebuilding the patch pool.

Sampling is more representative of thin fibers

Foreground-biased sampling remains a mixture of skeleton, segmentation, and image-driven candidates, but several biases were corrected:

• Skeleton jitter is now symmetric around the traced node instead of shifting only in the positive octant.
• The segmentation acceptance threshold is configurable.
• The image fallback uses a per-patch robust median/MAD mask rather than a fixed >100 count threshold.

Spatially incoherent artifacts are rejected before BM4D

Some bright salt-and-pepper processing artifacts are labeled as objects by the FFN segmenter. They cannot be removed merely by trusting segmentation or by an intensity threshold: they are often brighter than real signal.

The precompute configuration enables reject_incoherent_patches. Each labeled segment with enough voxels is tested for both:

• low lag-2 spatial autocorrelation; and
• high high-frequency energy relative to a Gaussian-smoothed patch.

A segment is considered an artifact only when it fails both tests, protecting thin but spatially smooth neurites from false rejection. The whole patch is resampled because the raw input itself is corrupted; simply deleting the mask would still train on a bad input. The check runs before BM4D, so rejected examples avoid the most expensive preprocessing step. A bounded retry count prevents fixed-size cache generation from stalling indefinitely; after the configured 50 failed attempts, the last patch is accepted even if contaminated.

Normalization and background transforms

One frozen transform replaces per-patch percentile normalization

On main, each training patch computed its own low/high percentiles, normalized against them, and clipped the normalized bright tail. Inference had separate normalization and a hard-coded clip. This caused two core problems:

1. The same physical count could map to a different network value depending on patch contents.
2. Bright neurites above the clip collapsed onto a non-invertible plateau.

machine_learning/transforms.py introduces fixed, stateless, configuration-driven transforms that are shared unchanged by cache datasets, validation, checkpoint metrics, and inference.

Transform	Role	Bright-tail behavior
AsinhTransform	Current configured transform; approximately linear near the noise floor and logarithmic in the tail	Monotonic and invertible up to the physical sensor clamp; no brightness plateau
AnscombeTransform	Generalized variance-stabilizing alternative for Poisson-Gaussian noise	Square-root-like compression; exact algebraic or asymptotically unbiased inverse
LinearClipTransform	Frozen baseline for A/B comparisons with the old approach	Retains an explicit hard clip and is intentionally not the recommended HDR mapping

The transform is normalized against the physical sensor maximum (65535 for uint16) rather than a patch statistic. Small negative normalized values for below-background counts are intentional, preserving noise-floor symmetry. The final inverse alone clips to the physical range and rounds to uint16.

Per-brain pedestal correction aligns acquisitions

Training data can mix background-subtracted and raw acquisitions whose background pedestals differ. Under a fixed nonlinear transform, a raw background around 100 counts does not map to the same network value as a subtracted background near zero.

The PR adds scripts/estimate_background_offsets.py to estimate a low, nonzero percentile from a coarse multiscale level for every brain. Zeros are ignored so padded regions do not pull the estimate to zero.

During cache generation, the selected per-brain value is subtracted from raw counts before BM4D and before the shared base transform. The base transform therefore stays fixed at offset zero for all cached examples.

At inference, OffsetTransform composes the runtime pedestal around the frozen base transform:

forward(raw) = trained_transform.forward(raw - runtime_offset)
inverse(y)   = trained_transform.inverse_float(y) + runtime_offset

Production inference should supply an offset precomputed from the full image tile at a lower-resolution level.

Transform identity is persisted and enforced

• Both caches stamp the resolved transform.json.
• Training refuses train/validation caches with different transforms.
• Checkpoints store the transform configuration next to model weights.
• Inference reconstructs the transform from the checkpoint.
• Loading pretrained weights rejects a checkpoint transform that does not match either attached dataset.

This turns normalization from ephemeral per-example metadata into part of the model's versioned contract.

Training objective, architecture, and loop

Loss function

Uniform nn.L1Loss is replaced by SignalPreservingLoss, a smooth Charbonnier penalty with voxel weight 1 + fg_weight * fg_mask. It is less sensitive to outliers than L2 while retaining an L1-like objective, and it can upweight sparse neurites that would otherwise be numerically dominated by background voxels.

U-Net behavior

The overall 3D U-Net topology remains, with two material changes:

• The default output is residual: prediction = input + learned_correction. Denoising is naturally a correction problem, and the desired correction is near zero where signal should be preserved.
• BatchNorm3d is replaced by GroupNorm. GroupNorm is independent of batch statistics and is less sensitive to small batches and rare bright patches. The group count uses gcd(8, channels) so custom integer width multipliers always divide the channel count.

width_multiplier is now validated as a positive integer instead of silently truncating fractional channel counts. The model exposes a constructor config containing width, upsampling mode, and residual mode so inference can recreate the exact architecture from a checkpoint.

These changes intentionally break compatibility with pre-overhaul weights: GroupNorm changes state-dict keys, residual output changes output semantics, and the learned coordinate system changed with normalization.

Full-cache epoch semantics and reproducibility

CachedPatchDataset is index-addressable and its length is the actual cache size. A training epoch now visits every cached item exactly once. The loader uses a deterministic permutation from (seed, epoch), so:

• rerunning with the same cache and seed gives the same order;
• each epoch gets a different order;
• validation stays ordered; and
• the last partial batch is retained.

This replaces the former n_examples_per_epoch sequence of fresh random draws. Epoch counts and learning-rate schedules are therefore not directly comparable between the branches.

Optimization and numerical stability

• AdamW remains the optimizer.
• CosineAnnealingLR.T_max now equals max_epochs, producing one decay across the run. The former fixed T_max=25 returned the learning rate to its peak every 50 epochs and correlated with repeated, growing loss spikes.
• Float16 autocast now has a GradScaler, preventing small gradients from underflowing and adapting to overflow.
• Tensor movement respects Trainer.device instead of forcing CUDA in helper paths. AMP autocast and its scaler remain CUDA-specific, so a non-CUDA run must set use_amp=False.
• val_every decouples CPU-heavy validation from epoch cost, while always validating the final epoch. The current full-cache experiment chooses every epoch, but the general trainer still supports a longer interval.

Run and checkpoint provenance

Each session records a config.json with cache locations, transform, target policy, model constructor configuration, batch size, learning rate, epoch count, loader settings, validation cadence, seed, foreground weight, and checkpoint weights.

Checkpoints now contain:

model          model state_dict
model_config   width/trilinear/residual constructor arguments
transform      frozen intensity-transform configuration

Every evaluated checkpoint after epoch index 0 is saved, with its selection score in the filename, so operating points can be compared offline. Index 0 has already completed one training pass, but the still-early model was empirically susceptible to nearly constant, trivially compressible output and a spuriously good compression-aware score. The is_best flag is now informational rather than deciding whether a checkpoint exists.

resume_path restores compatible model weights, but it is not a full training resume: optimizer state, scheduler state, epoch number, scaler state, and the previous best score are not serialized.

Training previews changed from full 3D TIFFs to percentile-stretched PNG maximum-intensity projections, which are much faster to inspect.

Validation, evaluation, and checkpoint selection

Validation is performed in count space

The model still runs and trains in transform space, but validation inverses its predictions before computing the new metrics. This makes comparisons meaningful across different transforms and expresses errors in sensor counts.

Per-example metrics include:

Metric	Reference and intent
Foreground MAE	Prediction versus raw counts inside the annotation mask; measures signal preservation
Background MAE	Prediction versus the training target outside the mask; measures teacher-like cleanup
Top-0.1% error/preservation	Compares the bright tail to raw input; catches systematic neurite attenuation
Global-maximum error (mip_max_error)	Absolute difference between the volumes' scalar maxima; does not compare MIPs spatially and can miss a lost structure when another peak remains
False-bright rate	Fraction of background predicted above a brightness threshold
Compression ratio	Zstd/Blosc compressibility of the inverse-transformed uint16 prediction

Metrics are averaged over examples, while compression ratio is summarized by the median because its per-patch distribution is heavy-tailed.

Checkpoint selection changes from lowest global validation L1 to a lower-is-better joint score:

score = w_fg  * foreground_mae
      + w_bg  * background_mae
      + w_top * top_0.1_percent_error
      - w_cr  * compression_ratio

Offline evaluators were repaired and share transform/offset handling

evaluate.py was updated to the current inference API and recognizes both current dictionary checkpoints and the legacy bare-state-dict format. For this module, the caller still has to construct the correct model; it does not recreate model_config. Bare loading is therefore only a format fallback and is not semantically faithful for main checkpoints trained with BatchNorm, direct output, and percentile normalization.

• SupervisedEvaluator applies an estimated offset per standalone raw block and reports compression ratio and 3D SSIM against the raw/noisy block. Despite the historical class name, that SSIM is not against clean ground truth.
• UnsupervisedEvaluator applies the training offset once per brain rather than estimating it per patch, then reports raw/BM4D/model compression, SSIM against raw and BM4D, MAE, and maximum absolute error.
• Both have a raw_input switch for data that has already been background subtracted.

These evaluator results are not numerically interchangeable with training validation: evaluate.py uses Blosc/Zstd compression level 6 rather than 5, and its unsupervised BM4D reference uses sigma 10 rather than the current cache sigma 24.

The custom 3D SSIM implementation now converts uint16 inputs to floating point before powers and products. Previously, integer overflow could corrupt local variance and covariance. MAE and maximum-error helpers likewise perform safe floating-point subtraction.

A standalone crop evaluation script was added

scripts/evaluate_bm4dnet.py provides an end-to-end crop evaluation path:

• choose an explicit checkpoint or parse the lowest score from checkpoint filenames, including negative scores;
• load both model architecture and transform from the checkpoint;
• read a lazy local/S3/GCS Zarr array and crop without loading the whole tile;
• convert Neuroglancer (x, y, z) coordinates to source-level (z, y, x) voxels and validate crop bounds;
• apply a precomputed full-tile background offset, with crop-local estimation only as a fallback;
• report raw and denoised compression ratios and their gain;
• save orthogonal raw/denoised MIPs; and
• optionally write the denoised crop as OME-Zarr while preserving the source scale, units, and crop-adjusted translation.

Performance and scalability

CPU and I/O work is moved out of the GPU loop

Cloud reads and per-patch BM4D were the dominant training cost on main and left the GPU idle. Precomputing both train and validation sets pays that cost once. Cache-backed training then performs memory-mapped reads, the inexpensive transform/target assembly, and GPU work. The cache entrypoint uses num_workers=0 intentionally because process overhead is not worthwhile for this cheap path.

The general custom loader was also rewritten for live/expensive datasets. A bounded background thread prefetches batches; optional worker processes are created once per iterator/epoch rather than once per batch, and the dataset is installed once in each worker instead of being pickled with every task.

TensorStore caching is disabled for random patch access

TensorStore handles previously received separate 1 GB local and remote cache pools per worker and per brain, allowing aggregate memory to grow far beyond host RAM. Both limits are now zero. Random 64-cubed sampling has low cache reuse, so bounded memory was judged more valuable than those caches.

Full-volume inference uses lower-precision, in-place accumulation

Overlap-add accumulators are now float32 rather than NumPy's default float64, and averaging happens in place after the transformed input is freed. This roughly halves peak host memory and removes extra full-volume temporaries, addressing out-of-memory failures for 1024-cubed volumes on a 30 GB host.

Image I/O, dependencies, and developer tooling

Zarr 3 migration

The supported environment moves to Python 3.11+, zarr>=3.0.8, and ome-zarr>=0.12.0. Image utilities now use Zarr 3 path-based local/cloud stores and BloscCodec:

• local, S3, and GCS Zarr reads use zarr.open;
• OME-Zarr writing uses Zarr 3 codec/chunk options;
• write_zarr writes a single 5D local or cloud array;
• write_ome_zarr supports explicit scale, translation, and spatial unit;
• N5 uses TensorStore because Zarr 3 removed its built-in N5 store.

The N5 migration is structurally covered but is explicitly noted in code as not exercised against current live N5 data.

.gitignore also excludes CodeOcean/workspace artifacts, local data and scratch directories, generated PNG previews, and related environment metadata.

Added operational scripts

Script	Purpose
estimate_background_offsets.py	Build and inspect per-brain pedestal JSON from coarse, nonzero data
precompute.py	Reproducibly build either train or validation cache from one shared configuration
train_bm4dnet.py	Validate the cache contract, persist the run config, and train fully offline
evaluate_bm4dnet.py	Evaluate a checkpoint on a real crop, measure compression, and optionally write metadata-correct OME-Zarr
visualize_patches.py	Compare cached raw, BM4D teacher, optional preserved target, mask, and overlay
visualize_swc_masks.py	Rebuild and inspect the SWC-only contribution to foreground masks

The visualization tools are important QA for a design in which the mask is baked into caches: they make it possible to inspect whether annotations cover real neurites without annexing noise before spending a training run. visualize_patches.py always renders the raw-on-foreground composite in its "target" panel. With the current preserve_foreground=False experiment, the actual training target is the separate BM4D teacher panel.

Reproducibility and test coverage

Cache generation seeds every task from a SeedSequence keyed by base seed, split-specific stream, and task index. Train and validation have different streams, so a fixed software/data environment is designed to reproduce sample choices and ordering independently of worker scheduling/count. Bitwise BM4D identity across platforms or library versions is not established. The intended cache-only loader path (num_workers=0) also has reproducible epoch ordering. This does not make the model run bit-reproducible: model initialization and PyTorch/CUDA kernels are not seeded or forced into deterministic modes. Full precompute configuration, including the count dtype, is written for provenance.

What remains conceptually unchanged

• BM4D is still the teacher/reference denoiser; it is now moved offline and can be selectively overridden on annotated foreground.
• The core model remains a patch-based 3D U-Net operating on 64-cubed patches.
• Compression is still measured with shuffled Zstd/Blosc and is a first-class project outcome.
• Final predictions are returned in physical uint16 count space.

The main architectural shift is therefore not a replacement of BM4D or the U-Net. It is the introduction of a reproducible cached-data, count-aware contract around them: curated cached examples, one stable intensity coordinate system, annotation-aware validation, and checkpoint selection that acknowledges the fidelity/compression tradeoff.

[carshadi]

Py310 tests fail due to the new >=3.11 requirement for zarr>=3. The actions could be fixed in a separate PR.

[anna-grim]

The code looks great and is a huge improvement to the repo. I really don't see any obvious issues or bugs, so i'll just go ahead and merge this branch