diff --git a/docs/parameters.md b/docs/parameters.md index d63fb037..c51f51eb 100644 --- a/docs/parameters.md +++ b/docs/parameters.md @@ -8,7 +8,7 @@ Suite2p can be run with different configurations using the db and settings dicti |---|---|---|---|---| | `data_path` | Data path | `` | `[]` | List of folders with tiffs or other files to process. | | `look_one_level_down` | Look one level down | `` | `False` | Whether to look in all subfolders of all data_path folders when searching for tiffs. | -| `input_format` | Input format | `` | `tif` | Can be ['tif', 'h5', 'nwb', 'bruker', 'movie', 'dcimg']. | +| `input_format` | Input format | `` | `tif` | Can be ['tif', 'h5', 'nwb', 'nd2', 'bruker', 'movie', 'dcimg']. | | `keep_movie_raw` | Keep movie raw | `` | `False` | Whether to keep binary file of non-registered frames. | | `nplanes` | Number of planes | `` | `1` | Each tiff / file has this many planes in sequence. | | `nrois` | Number of ScanImage ROIs | `` | `1` | Each tiff / file has this many different ROIs. | diff --git a/suite2p/io/dcam.py b/suite2p/io/dcam.py index b0355588..cbdfbb9a 100644 --- a/suite2p/io/dcam.py +++ b/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) diff --git a/suite2p/io/nd2.py b/suite2p/io/nd2.py index 627041e6..f880fd7a 100644 --- a/suite2p/io/nd2.py +++ b/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 diff --git a/suite2p/parameters.py b/suite2p/parameters.py index 60ae4040..f06fc958 100644 --- a/suite2p/parameters.py +++ b/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", diff --git a/tests/test_io.py b/tests/test_io.py index bdff4abe..fc8c7f66 100644 --- a/tests/test_io.py +++ b/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())