Port nd2 reader to the 1.x reader API and expose it in the GUI

docs/parameters.md

@@ -8,7 +8,7 @@ Suite2p can be run with different configurations using the db and settings dicti
 |---|---|---|---|---|
 | `data_path` | Data path | `<class 'list'>` | `[]` | List of folders with tiffs or other files to process. |
 | `look_one_level_down` | Look one level down | `<class 'bool'>` | `False` | Whether to look in all subfolders of all data_path folders when searching for tiffs. |
-| `input_format` | Input format | `<class 'str'>` | `tif` | Can be ['tif', 'h5', 'nwb', 'bruker', 'movie', 'dcimg']. |
+| `input_format` | Input format | `<class 'str'>` | `tif` | Can be ['tif', 'h5', 'nwb', 'nd2', 'bruker', 'movie', 'dcimg']. |
 | `keep_movie_raw` | Keep movie raw | `<class 'bool'>` | `False` | Whether to keep binary file of non-registered frames. |
 | `nplanes` | Number of planes | `<class 'int'>` | `1` | Each tiff / file has this many planes in sequence. |
 | `nrois` | Number of ScanImage ROIs | `<class 'int'>` | `1` | Each tiff / file has this many different ROIs. |

suite2p/io/dcam.py

@@ -32,6 +32,9 @@ def dcimg_to_binary(dbs, settings, reg_file, reg_file_chan2):
             "nframes", "meanImg", "meanImg_chan2"
     """
 
+    if not DCIMG:
+        raise ImportError("dcimg is required for this file type, please 'pip install dcimg'")
+
     t0 = time.time()
     # # copy settings to list where each element is settings for each plane
     # settings1 = utils.init_settings(settings)

suite2p/io/nd2.py

@@ -1,126 +1,131 @@
 """
 Copyright © 2023 Howard Hughes Medical Institute, Authored by Carsen Stringer and Marius Pachitariu.
 """
-import os
 import gc
 import math
 import time
+import logging
 import numpy as np
-from . import utils
+
+logger = logging.getLogger(__name__)
 
 try:
-    import nd2 
-    ND2 = True
+    import nd2
+    HAS_ND2 = True
 except ImportError:
-    ND2 = False
+    HAS_ND2 = False
+
 
-def nd2_to_binary(settings):
+def nd2_to_binary(dbs, settings, reg_file, reg_file_chan2):
     """finds nd2 files and writes them to binaries
 
     Parameters
     ----------
-    settings: dictionary
-        "nplanes", "data_path", "save_path", "save_folder", "fast_disk",
-        "nchannels", "keep_movie_raw", "look_one_level_down"
+    dbs : list of dict
+        Per-plane database dictionaries. Must contain "file_list", "first_files",
+        "nplanes", "nchannels", "batch_size", "functional_chan". Updated in-place
+        with "Ly", "Lx", "nframes", "frames_per_file", "frames_per_folder",
+        "meanImg", and "meanImg_chan2".
+    settings : dict
+        Suite2p settings dictionary.
+    reg_file : list of file objects
+        Opened binary files for writing each plane's functional channel data.
+    reg_file_chan2 : list of file objects
+        Opened binary files for writing each plane's second channel data
+        (used only when nchannels > 1).
 
     Returns
     -------
-        settings : dictionary of first plane
-            settings["reg_file"] or settings["raw_file"] is created binary
-            assigns keys "Ly", "Lx", "tiffreader", "first_tiffs",
-            "nframes", "meanImg", "meanImg_chan2"
+    dbs : list of dict
+        Updated database dictionaries.
     """
+    if not HAS_ND2:
+        raise ImportError("nd2 is required for this file type, please 'pip install nd2'")
+
+    fs = dbs[0]["file_list"]
+    first_files = dbs[0]["first_files"]
+    nplanes = dbs[0]["nplanes"]
+    nchannels = dbs[0]["nchannels"]
+    if nchannels > 2:
+        raise ValueError(f"nd2 input supports at most 2 channels, got nchannels={nchannels}.")
+    # batch_size counts interleaved frames like the tiff/h5 readers, so one chunk
+    # of timepoints holds batch_size / (nplanes * nchannels) volumes
+    tbatch = max(1, math.ceil(dbs[0]["batch_size"] / (nplanes * nchannels)))
+    nfunc = dbs[0]["functional_chan"] - 1 if nchannels > 1 else 0
 
     t0 = time.time()
-    # copy settings to list where each element is settings for each plane
-    settings1 = utils.init_settings(settings)
-
-    # open all binary files for writing
-    # look for nd2s in all requested folders
-    settings1, fs, reg_file, reg_file_chan2 = utils.find_files_open_binaries(settings1, False)
-    settings = settings1[0]
-
-    # loop over all nd2 files
     iall = 0
-    ik = 0
-    for file_name in fs:
-        # open nd2
-        nd2_file = nd2.ND2File(file_name)
-        nd2_dims = {k: i for i, k in enumerate(nd2_file.sizes)}
-
-        valid_dimensions = "TZCYX"
-        assert set(nd2_dims) <= set(
-            valid_dimensions
-        ), f"Unknown dimensions {set(nd2_dims)-set(valid_dimensions)} in file {file_name}."
-
-        # Sort the dimensions in the order of TZCYX, skipping the missing ones.
-        im = nd2_file.asarray().transpose(
-            [nd2_dims[x] for x in valid_dimensions if x in nd2_dims])
-
-        # Expand array to include the missing dimensions.
-        for i, dim in enumerate("TZC"):
-            if dim not in nd2_dims:
-                im = np.expand_dims(im, i)
-
-        nplanes = nd2_file.sizes["Z"] if "Z" in nd2_file.sizes else 1
-        nchannels = nd2_file.sizes["C"] if "C" in nd2_file.sizes else 1
-        nframes = nd2_file.sizes["T"] if "T" in nd2_file.sizes else 1
-
-        iblocks = np.arange(0, nframes, settings1[0]["batch_size"])
-        if iblocks[-1] < nframes:
-            iblocks = np.append(iblocks, nframes)
-
-        if nchannels > 1:
-            nfunc = settings1[0]["functional_chan"] - 1
-        else:
-            nfunc = 0
-
-        assert im.max() < 32768 and im.min() >= -32768, "image data is out of range"
-        im = im.astype(np.int16)
-
-        # loop over all frames
-        for ichunk, onset in enumerate(iblocks[:-1]):
-            offset = iblocks[ichunk + 1]
-            im_p = np.array(im[onset:offset, :, :, :, :])
-            im2mean = im_p.mean(axis=0).astype(np.float32) / len(iblocks)
-            for ichan in range(nchannels):
-                nframes = im_p.shape[0]
-                im2write = im_p[:, :, ichan, :, :]
-                for j in range(0, nplanes):
+    which_folder = -1
+    for ifile, file_name in enumerate(fs):
+        if first_files[ifile]:
+            which_folder += 1
+        with nd2.ND2File(file_name) as nd2_file:
+            sizes = nd2_file.sizes
+            valid_dimensions = "TZCYX"
+            unknown = set(sizes) - set(valid_dimensions)
+            if unknown:
+                raise ValueError(f"Unknown dimensions {unknown} in file {file_name}.")
+            nplanes_file = sizes.get("Z", 1)
+            nchannels_file = sizes.get("C", 1)
+            nframes = sizes.get("T", 1)
+            if nplanes_file != nplanes or nchannels_file != nchannels:
+                raise ValueError(
+                    f"{file_name} has {nplanes_file} planes and {nchannels_file} channels, "
+                    f"but db has nplanes={nplanes} and nchannels={nchannels}.")
+
+            # lazy view reordered to T x Z x C x Y x X, inserting missing axes
+            data = nd2_file.to_dask()
+            data = data.transpose([list(sizes).index(d) for d in valid_dimensions if d in sizes])
+            data = data[tuple(slice(None) if d in sizes else None for d in "TZC") + (Ellipsis,)]
+
+            iblocks = np.arange(0, nframes, tbatch)
+            if iblocks[-1] < nframes:
+                iblocks = np.append(iblocks, nframes)
+            for ichunk, onset in enumerate(iblocks[:-1]):
+                offset = iblocks[ichunk + 1]
+                im = np.asarray(data[onset:offset])
+                if im.dtype.type == np.uint16:
+                    im = (im // 2).astype(np.int16)
+                elif im.dtype.type == np.int32:
+                    if im.min() < -65536 or im.max() > 65535:
+                        raise ValueError(
+                            f"int32 values in {file_name} do not fit int16 after halving.")
+                    im = (im // 2).astype(np.int16)
+                elif im.dtype.type in (np.uint8, np.int16):
+                    im = im.astype(np.int16)
+                else:
+                    raise ValueError(f"unsupported nd2 pixel dtype {im.dtype} in {file_name}.")
+                nframes_chunk, _, _, Ly, Lx = im.shape
+                for j in range(nplanes):
                     if iall == 0:
-                        settings1[j]["meanImg"] = np.zeros((im_p.shape[3], im_p.shape[4]),
-                                                      np.float32)
+                        dbs[j]["Ly"] = Ly
+                        dbs[j]["Lx"] = Lx
+                        dbs[j]["meanImg"] = np.zeros((Ly, Lx), np.float32)
                         if nchannels > 1:
-                            settings1[j]["meanImg_chan2"] = np.zeros(
-                                (im_p.shape[3], im_p.shape[4]), np.float32)
-                        settings1[j]["nframes"] = 0
-                    if ichan == nfunc:
-                        settings1[j]["meanImg"] += np.squeeze(im2mean[j, ichan, :, :])
-                        reg_file[j].write(
-                            bytearray(im2write[:, j, :, :].astype("int16")))
-                    else:
-                        settings1[j]["meanImg_chan2"] += np.squeeze(im2mean[j, ichan, :, :])
-                        reg_file_chan2[j].write(
-                            bytearray(im2write[:, j, :, :].astype("int16")))
-
-                    settings1[j]["nframes"] += im2write.shape[0]
-            ik += nframes
-            iall += nframes
-
-        nd2_file.close()
-
-    # write settings files
-    do_registration = settings1[0]["do_registration"]
-    for settings in settings1:
-        settings["Ly"] = im.shape[3]
-        settings["Lx"] = im.shape[4]
-        if not do_registration:
-            settings["yrange"] = np.array([0, settings["Ly"]])
-            settings["xrange"] = np.array([0, settings["Lx"]])
-        np.save(settings["settings_path"], settings)
-    # close all binary files and write settings files
-    for j in range(0, nplanes):
-        reg_file[j].close()
+                            dbs[j]["meanImg_chan2"] = np.zeros((Ly, Lx), np.float32)
+                        dbs[j]["nframes"] = 0
+                        dbs[j]["frames_per_file"] = np.zeros(len(fs), np.int32)
+                        dbs[j]["frames_per_folder"] = np.zeros(first_files.sum(), np.int32)
+                    im2write = np.ascontiguousarray(im[:, j, nfunc])
+                    reg_file[j].write(bytearray(im2write))
+                    dbs[j]["meanImg"] += im2write.astype(np.float32).sum(axis=0)
+                    if nchannels > 1:
+                        im2write = np.ascontiguousarray(im[:, j, 1 - nfunc])
+                        reg_file_chan2[j].write(bytearray(im2write))
+                        dbs[j]["meanImg_chan2"] += im2write.astype(np.float32).sum(axis=0)
+                    dbs[j]["nframes"] += nframes_chunk
+                    dbs[j]["frames_per_file"][ifile] += nframes_chunk
+                    dbs[j]["frames_per_folder"][which_folder] += nframes_chunk
+                iall += nframes_chunk
+                if (ichunk + 1) % 4 == 0 or offset == nframes:
+                    logger.info("%d frames of binary, time %0.2f sec." % (iall, time.time() - t0))
+        gc.collect()
+
+    for db in dbs:
+        db["meanImg"] /= db["nframes"]
         if nchannels > 1:
-            reg_file_chan2[j].close()
-    return settings1[0]
+            db["meanImg_chan2"] /= db["nframes"]
+        np.save(db["db_path"], db)
+        np.save(db["settings_path"], settings)
+
+    return dbs

suite2p/parameters.py

@@ -44,7 +44,7 @@
             "min": None,
             "max": None,
             "default": "tif",
-            "description": "Can be ['tif', 'h5', 'nwb', 'bruker', 'movie', 'dcimg'].",
+            "description": "Can be ['tif', 'h5', 'nwb', 'nd2', 'bruker', 'movie', 'dcimg'].",
         },
         "keep_movie_raw": {
             "gui_name": "Keep movie raw",

tests/test_io.py

@@ -7,8 +7,9 @@
 import pytest
 from natsort import natsorted
 from pynwb import NWBHDF5IO
-from suite2p import io
+from suite2p import default_db, default_settings, io
 from suite2p.io.nwb import read_nwb, save_nwb
+from suite2p.run_s2p import files_to_binary
 from suite2p.io.utils import get_suite2p_path
 from suite2p.detection.detect import bin_movie
 
@@ -258,3 +259,101 @@ def test_get_suite2p_path(input_path, expected_path, success):
     else:
         with pytest.raises(FileNotFoundError):
             get_suite2p_path(Path(input_path))
+
+
+def _convert_to_binary(db, settings):
+    """Runs the file -> binary conversion block of run_s2p for db["input_format"]."""
+    import contextlib
+
+    fs, first_files = io.get_file_list(db)
+    db["file_list"] = fs
+    db["first_files"] = first_files
+    dbs = io.init_dbs(db)
+    Path(db["save_path0"]).joinpath(db["save_folder"]).mkdir(parents=True, exist_ok=True)
+    with contextlib.ExitStack() as stack:
+        files = [stack.enter_context(open(d["reg_file"], "wb")) for d in dbs]
+        files_chan2 = None
+        if db["nchannels"] > 1:
+            files_chan2 = [stack.enter_context(open(d["reg_file_chan2"], "wb")) for d in dbs]
+        return files_to_binary[db["input_format"]](dbs, settings, files, files_chan2)
+
+
+def _install_fake_nd2(monkeypatch, sizes, arr):
+    """Makes suite2p.io.nd2 read `arr` (axes in `sizes` order) without the nd2 package."""
+    da = pytest.importorskip("dask.array")
+    from suite2p.io import nd2 as nd2_module
+
+    class FakeND2File:
+        def __init__(self, path):
+            self.sizes = dict(sizes)
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *args):
+            return False
+
+        def to_dask(self):
+            return da.from_array(arr, chunks=(1,) * (arr.ndim - 2) + arr.shape[-2:])
+
+    monkeypatch.setattr(nd2_module, "nd2", type("nd2", (), {"ND2File": FakeND2File}),
+                        raising=False)
+    monkeypatch.setattr(nd2_module, "HAS_ND2", True)
+
+
+def _nd2_db(tmp_path, **overrides):
+    data_dir = tmp_path / "data"
+    data_dir.mkdir()
+    (data_dir / "rec.nd2").touch()
+    return {**default_db(), "data_path": [str(data_dir)], "input_format": "nd2",
+            "save_path0": str(tmp_path / "out"), "batch_size": 3, **overrides}
+
+
+def test_nd2_to_binary_reorders_axes_and_splits_planes_and_channels(tmp_path, monkeypatch):
+    T, Z, C, Y, X = 7, 2, 2, 5, 6
+    arr = np.random.default_rng(0).integers(0, 65535, size=(Z, T, C, Y, X), dtype=np.uint16)
+    _install_fake_nd2(monkeypatch, {"Z": Z, "T": T, "C": C, "Y": Y, "X": X}, arr)
+
+    db = _nd2_db(tmp_path, nplanes=Z, nchannels=C, functional_chan=2)
+    dbs = _convert_to_binary(db, default_settings())
+
+    for j in range(Z):
+        func = np.fromfile(dbs[j]["reg_file"], np.int16).reshape(T, Y, X)
+        chan2 = np.fromfile(dbs[j]["reg_file_chan2"], np.int16).reshape(T, Y, X)
+        np.testing.assert_array_equal(func, (arr[j, :, 1] // 2).astype(np.int16))
+        np.testing.assert_array_equal(chan2, (arr[j, :, 0] // 2).astype(np.int16))
+        np.testing.assert_allclose(dbs[j]["meanImg"], func.mean(axis=0), atol=1e-3)
+        assert dbs[j]["nframes"] == T
+        assert list(dbs[j]["frames_per_file"]) == [T]
+        assert Path(dbs[j]["db_path"]).exists()
+
+
+def test_nd2_to_binary_rejects_plane_count_mismatch(tmp_path, monkeypatch):
+    arr = np.zeros((4, 2, 5, 6), dtype=np.uint16)
+    _install_fake_nd2(monkeypatch, {"T": 4, "Z": 2, "Y": 5, "X": 6}, arr)
+
+    with pytest.raises(ValueError, match="nplanes"):
+        _convert_to_binary(_nd2_db(tmp_path, nplanes=1, nchannels=1), default_settings())
+
+
+def test_nd2_to_binary_rejects_float_data(tmp_path, monkeypatch):
+    arr = np.zeros((4, 5, 6), dtype=np.float32)
+    _install_fake_nd2(monkeypatch, {"T": 4, "Y": 5, "X": 6}, arr)
+
+    with pytest.raises(ValueError, match="dtype"):
+        _convert_to_binary(_nd2_db(tmp_path), default_settings())
+
+
+@pytest.mark.parametrize("value, fits", [(65535, True), (65536, False),
+                                         (-65536, True), (-65537, False)])
+def test_nd2_to_binary_checks_int32_range(tmp_path, monkeypatch, value, fits):
+    arr = np.full((2, 5, 6), value, dtype=np.int32)
+    _install_fake_nd2(monkeypatch, {"T": 2, "Y": 5, "X": 6}, arr)
+
+    if fits:
+        dbs = _convert_to_binary(_nd2_db(tmp_path), default_settings())
+        got = np.fromfile(dbs[0]["reg_file"], np.int16)
+        assert got.min() == got.max() == value // 2
+    else:
+        with pytest.raises(ValueError, match="int32"):
+            _convert_to_binary(_nd2_db(tmp_path), default_settings())