[ieivanov]

Merge order and dependencies

1. This PR has no dependencies. It works with the iohub waveorder already requires, and can merge first.
2. refactor(cli)!: run apply-inv-tf through iohub's process_single_position, with --resume and per-reconstruction provenance keys #583 is stacked on this branch. It runs apply-inv-tf through iohub's process_single_position and adds --resume. It keeps this PR's gains: the TF is passed as torch tensors, so it crosses iohub's pool as a shared-memory handle, and each volume is read once. Merge it after this PR, then retarget it to main.
3. refactor(cli)!: run apply-inv-tf through iohub's process_single_position, with --resume and per-reconstruction provenance keys #583 depends on fix(ngff): sync a unit's shards before recording it as done czbiohub-sf/iohub#476, pinned until iohub releases it. This PR doesn't.

Problem

apply_inverse_transfer_function_single_position binds the loaded transfer function into the per-time-point partial. ProcessPoolExecutor pickles each task's arguments. The TF was an xr.Dataset of numpy arrays, so the full TF was serialised by the parent and sent through a pipe to a worker once per time point. _to_tensor then copied it again on every call.

Observed in biahub/Nextflow runs (one position per job, 16 processes):

run	T	TF size	peak RSS	rchar / wchar per position
2026_08_14_dynatrack	768	2 × (181, 768, 794) c64 ≈ 1.77 GB	175.6 GB	1.6 TB / 1.5 TB
2026_08_11_A549_SEC61B_DENV (v6)	67	2 × (96, 1664, 1193) c64 ≈ 3.05 GB	~270 GB	~252 GB / ~228 GB

The on-disk data is far smaller. For dynatrack, ~210 GB is read and ~261 GB written per position. The excess is T × TF size of pipe traffic, not disk I/O: 768 × 1.77 GB ≈ 1.36 TB for dynatrack, and 67 × 3.05 GB ≈ 204 GB for v6. rchar/wchar (what the Nextflow trace reports) count every byte passed through read()/write() system calls, including pipes. The TF is never written to storage. The cost is CPU to pickle and unpickle ~3 GB per time point, temporary memory copies in both processes, and dispatch serialised through the executor's single feeder thread.

Separately, each time point's input volume was read from the store twice (see Changes).

Changes

• cli/apply_inverse_transfer_function.py: _load_transfer_function_dataset → _load_transfer_function, which returns a dict of torch tensors instead of an xr.Dataset of numpy arrays. It's still bound into each time point's task as before; the executor code is unchanged from main.
  • This alone fixes the pickling. Importing torch.multiprocessing registers pickling reductions that move a pickled CPU tensor into shared memory and send a small handle to it, not its data. That applies to any ForkingPickler route, including the executor's per-task call queue. So every worker maps the same single copy: one storage read, one copy in RAM, and nothing large through the pipes. Numpy arrays, or an xr.Dataset of them, pickle by value, which is why main didn't get this.
• api/_utils.py: _to_tensor / _to_singular_system also accept a mapping of tensors, returned as-is. xr.Dataset variables are now wrapped with torch.from_numpy without the extra .copy(). None of the apply-inverse models, or the helpers they pass the TF to, write to it in place. A test checks that apply leaves an xr.Dataset TF unchanged.
• api/{phase,fluorescence,birefringence,birefringence_and_phase}.py: the apply_inverse_transfer_function type hints and docstrings say they accept xr.Dataset | Mapping[str, torch.Tensor].
• cli/utils.py: apply_inverse_to_zyx_and_save loads the dask-backed input slice once (.load()). Before, the NaN/zero check and the model each called .values, so every time point was read and decompressed twice (0.48 GB extra per time point on v6).

Not changed: the executor type (ProcessPoolExecutor, spawn context), thread settings, and the public signature of apply_inverse_transfer_function_single_position.

Commits, which trace how we got here (squash on merge is fine):

1. 6521aaf loaded the TF in each worker via a pool initializer. That removed the pickling but read and held the TF once per worker.
2. 9f496df shared one copy with an explicit share_memory_().
3. d080361 dropped share_memory_(), which torch's pickling hook already does.
4. 89637f6 dropped the initializer too, since per-task tensors already cross as handles. The loaders were merged and the API docs updated.
5. 097ff97 made each time point read its input once and removed the _to_tensor copy.

Benchmark

wo apply-inv-tf -j 16 on one v6 position (C/2/000002). The TF was recomputed with the correct pixel sizes (0.1133 / 0.17; same shape). Each table is one SLURM job on one node (16 CPUs, 336 GB), with the versions run alternately.

How the numbers were measured:

• I/O and CPU are the monitor's own /proc/self/io and RUSAGE_CHILDREN deltas after reaping the command, which counts the whole process tree once.
• Pipe traffic = wchar − write_bytes.
• Storage reads ≈ rchar − pipe traffic. On Lustre, write_bytes is exact, but read_bytes undercounts because readahead runs in kernel threads.
• PSS (sampled every 2 s) splits shared pages between the processes that map them. Summed RSS, which Nextflow's peak_rss reports, would count the shared TF once per worker.

Final, 16 time points (node cpu-h-15): main vs the previous commit vs this PR's head

	main (60c3af6)	89637f6 (before the single input read)	this PR (097ff97)
wall time	100 s	39.9 / 40.8 s	39.4 / 39.9 s (2.5× faster than main)
CPU (process tree, user+sys)	659 s	471 / 466 s	443 / 450 s (−32% vs main)
pipe traffic	48.8 GB (= 16 × 3.05 GB TF pickles)	0	0
storage reads	18.8 GB	18.8 GB	11.2 GB (−41%: input read once, TF read once)
storage writes	9.6 GB	9.6 GB	9.6 GB (the output)
peak PSS, process tree	44 GB	112–134 GB	124–140 GB
peak RSS, parent	9.7 GB	4.4–4.8 GB	4.8–5.3 GB (holds the shared TF)
output		bit-identical	bit-identical (16/16 time points)

Why main has the lowest memory: it's limited by its pipe. The executor's single feeder thread pushes the TF pickles into worker pipes at ~0.4–0.5 GB/s, so workers mostly wait for their payload and few are busy at once. This PR runs all 16 at once. Each worker's private working set is ~11 GB (FFT buffers), plus the one shared 3 GB TF.

Earlier runs, per commit (all bit-identical to main):

T, node	main	6521aaf per-worker load	9f496df explicit shm	d080361 implicit shm	89637f6 per-task handles
32, cpu-d-1: wall / CPU	270 s / 1861 s	106–109 s / 1525–1578 s	103–106 s / 1360–1391 s
32, cpu-d-1: storage reads / peak PSS	34.5 GB / 55 GB	73.1 GB / 153–181 GB	34.1 GB / 110–117 GB
16, cpu-h-15: wall / CPU	98–101 s / 627–652 s		42–43 s / 481–511 s	40–43 s / 473–483 s	40–41 s / 463–465 s

Projected for T = 100 (node cpu-d-1, from the per-time-point costs measured in the 32-time-point job; CPU and wall time are before 097ff97, which cut a further ~4% of CPU at T = 16):

T = 100	main	this PR
storage reads	~98 GB	~51 GB (TF once + each input volume once)
storage writes	~60 GB	~60 GB
pipe traffic	~305 GB	0
CPU	~5,800 s	~4,250 s (≈ −27%)
wall time	~850 s (pipe-bound)	~300 s (compute-bound)

Testing

• New test_apply_inv_tf_process_pool_matches_serial: a 3D phase reconstruction with 2 spawned workers gives output identical to the serial path.
• New test_apply_inverse_to_zyx_and_save_reads_input_once: the dask-backed input is computed once per time point. It fails with 2 == 1 on the old code.
• New test_apply_inverse_leaves_dataset_transfer_function_unchanged[phase|fluorescence]: apply doesn't modify an xr.Dataset TF, and the tensor shares its memory (no copy).
• New API tests: phase 3D gives identical results from a tensor dict and an xr.Dataset, and _to_tensor returns tensors as-is.
• tests/api_tests + tests/cli_tests: 185 passed, 1 skipped, 2 xfailed.
• Benchmarks above: output bit-identical to main on real data for every commit.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 90.47619% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 56.05%. Comparing base (60c3af6) to head (d080361).

Files with missing lines	Patch %	Lines
waveorder/cli/apply_inverse_transfer_function.py	86.66%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #582      +/-   ##
==========================================
+ Coverage   55.95%   56.05%   +0.10%     
==========================================
  Files          75       75              
  Lines       12618    12631      +13     
==========================================
+ Hits         7060     7080      +20     
+ Misses       5558     5551       -7     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ieivanov]

@srivarra this PR is ready for review too 🙏 . The key change is that the internal _load_transfer_function function changes from returning x-array dataset to returning a torch Tensor which is a lot more efficient. The public API is expanded to allow passing either x-array or Torch tensor as the transfer function.

[srivarra]

Looks great!

[srivarra]

this is a cool find. I always thought Dask itself would be able to know where/when it needs things kept in memory.