Skip to content

feat(cli): opt-in float16 transfer functions - #589

Open
srivarra wants to merge 2 commits into
mainfrom
feat/float16-recon
Open

srivarra wants to merge 2 commits into
mainfrom
feat/float16-recon

Conversation

@srivarra

@srivarra srivarra commented Oct 23, 2026 •

Copy link
Copy Markdown
Contributor

Adds support for float16 OTFs. compute-tf writes them when a config sets transfer_function_dtype: float16, and apply-inv-tf and waveorder view converts them back to complex64 when they load them. If you don't pass in the dtype, then it continues to default to float32.

Each complex array keeps complex64's real-then-imaginary layout at half width. numpy and zarr don't support complex32 out of the box. We could manage a custom dtype, but that would need integration in iohub, and TensorStore has no way to register custom dtypes. So the store declares each complex array as float16 with X doubled and lists it in float16_complex_arrays.

Size

Getting some nice tf size reductions when testing it out with waveorder.api.*.simulate.

Case TF on disk
phase 3D, 64×512×512 218 MB → 30.5 MB
phase 2D, 16×512×512 63.6 MB → 27.6 MB
fluorescence 3D, 64×512×512 109 MB → 15.8 MB
birefringence + phase 3D, 32×256×256 314 MB → 140 MB

One thing that's nice is the stores shrink by more than half because blosc compresses the rounded float16 better.

Another option: zarr's complex_float16

This one's more work, but could be useful in the longer term...

zarr-extensions registers a complex_float16 data type for zarr v3, and ml_dtypes 0.6.0 adds a matching complex32 to numpy. The layout is the same as this pr's, so a store could declare complex_float16 at the original shape and drop the float16_complex_arrays attribute. It takes more work than this PR does:

  • zarr-python 3.1.6+ doesn't know this type natively, so we have to register a new Zarr type built on top of ml_dtypes.complex32. With that registered, iohub's zarrs-python and zarr-python backends can handle that dtype. The registration could live in waveorder, or in iohub so biahub and the GUI get it without importing waveorder.
  • The extension is zarr v3 only, so OME-Zarr v0.4 stores would still need the encoding in this PR.
  • It adds ml_dtypes>=0.6.0 as a dependency.
  • iohub's tensorstore backend can't read it /:

PyTorch can't handle a complex32 array either, so loading still upcasts it to complex64.

Add ReconstructionSettings.transfer_function_dtype (float32 default). With
float16, compute-tf stores complex arrays as interleaved real/imag float16
pairs along X and lists them under the float16_complex_arrays position
attribute; apply-inv-tf and waveorder view upcast them back to complex64.
compute-tf refuses values beyond float16's maximum of 65504.
# Conflicts:
#	tests/cli_tests/test_compute_tf.py
@srivarra
srivarra marked this pull request as ready for review October 23, 2026 23:51
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant