toolsandogh 0.1.0__py3-none-any.whl

toolsandogh/__init__.py

@@ -0,0 +1,23 @@
+"""
+A collection of tools for working with large-scale microscopy data.
+
+Provided to you by the Sandoghdar Division of the Max Planck Institute for the Physics of Light.
+"""
+
+from ._canonicalize_video import canonicalize_video
+from ._generate_video import generate_video
+from ._load_video import load_video
+from ._rolling import differential_rolling_average, rolling_average, rolling_sum
+from ._rvt import radial_variance_transform
+from ._store_video import store_video
+
+__all__: list[str] = [
+    "canonicalize_video",
+    "rolling_sum",
+    "rolling_average",
+    "differential_rolling_average",
+    "generate_video",
+    "load_video",
+    "store_video",
+    "radial_variance_transform",
+]

toolsandogh/_canonicalize_video.py

@@ -0,0 +1,242 @@
+import sys
+from datetime import datetime, timezone
+
+import dask.array as da
+import numpy as np
+import numpy.typing as npt
+import xarray as xr
+from ome_types.model import (
+    OME,
+    Image,
+    Pixels,
+    Pixels_DimensionOrder,
+    PixelType,
+    UnitsLength,
+    UnitsTime,
+)
+
+from ._validate_video import validate_video
+
+
+def canonicalize_video(
+    video: npt.ArrayLike,
+    # optional metadata
+    acquisition_date: datetime | None = None,
+    creator: str | None = None,
+    dt: float | None = None,
+    dz: float | None = None,
+    dy: float | None = None,
+    dx: float | None = None,
+) -> xr.DataArray:
+    """
+    Turn the supplied data into its canonical TCZYX video representation.
+
+    Parameters
+    ----------
+    video : xr.DataArray
+        An xarray.
+    acquisition_date : datetime.datetime
+        A timestamp of when the video was created.  Defaults to datetime.now().
+    creator : str
+        A string describing who created the video.
+    dt : float
+        The time interval in milliseconds between one video frame and the next.
+        Defaults to 1/60 of a second.
+    dz : float
+        The spatial distance in micrometer between one Z slice and the next.
+        Defaults to 1 micrometer.
+    dy : float
+        The spatial distance in micrometer between one row of pixels and the next.
+        Defaults to 1 micrometer.
+    dx : float
+        The spatial distance in micrometer between one column of pixels and the next.
+        Defaults to 1 micrometer.
+
+    Returns
+    -------
+    xarray.DataArray
+        A TCZYX video with the supplied parameters.
+    """
+    # Turn video into an xarray.
+    if not isinstance(video, xr.DataArray):
+        video = video_from_array(video)
+
+    # Ensure the video's data is a Dask array.
+    if not isinstance(video.data, da.Array):
+        video = video.copy(data=da.from_array(video.data), deep=False)
+
+    # Ensure the TZYX axes exist and are continuous.
+    for dim in ("T", "Z", "Y", "X"):
+        if dim not in video.dims:
+            video = video.expand_dims({dim: np.arange(1.0)})
+        elif video[dim].dtype != np.float64:
+            values = np.array(video[dim])
+            video = video.assign_coords({dim: values.astype(np.float64)})
+
+    # Ensure the C axis exists.
+    if "C" not in video.dims:
+        video = video.expand_dims({"C": 1})
+
+    # If there is a S axis, merge its entries.
+    if "S" in video.dims:
+        video = video.mean("S", dtype=np.float32).astype(video.dtype)
+
+    # Ensure the correct ordering of axes.
+    video = video.transpose("T", "C", "Z", "Y", "X")
+
+    # Ensure there is metadata attached.
+    if not hasattr(video, "processed"):
+        (size_t, size_c, size_z, size_y, size_x) = video.shape
+        pixels = Pixels(
+            type=dtype_pixel_type(video.dtype),
+            big_endian=dtype_is_big_endian(video.dtype),
+            dimension_order=Pixels_DimensionOrder.XYZCT,
+            size_t=size_t,
+            size_c=size_c,
+            size_z=size_z,
+            size_y=size_y,
+            size_x=size_x,
+            time_increment=dt or (1000 / 60),
+            physical_size_z=dz or 1.0,
+            physical_size_y=dy or 1.0,
+            physical_size_x=dx or 1.0,
+            time_increment_unit=UnitsTime.MILLISECOND,
+            physical_size_z_unit=UnitsLength.MICROMETER,
+            physical_size_y_unit=UnitsLength.MICROMETER,
+            physical_size_x_unit=UnitsLength.MICROMETER,
+        )
+        image = Image(
+            acquisition_date=(acquisition_date or datetime.now(timezone.utc)),
+            description="Video with auto-generated metadata.",
+            pixels=pixels,
+        )
+        ome = OME(
+            images=[image],
+            creator=(creator or "MPL Erlangen, Sandoghdar Division, toolsandogh"),
+        )
+        video = video.assign_attrs({"processed": ome})
+    # TODO: Update metadata with any supplied parameters.
+
+    # Raise an exception if the video is still not in canonical form.
+    validate_video(video)
+
+    # Done.
+    return video
+
+
+def video_from_array(array: npt.ArrayLike) -> xr.DataArray:
+    """
+    Turn a supplied array into an xarray.
+
+    Parameters
+    ----------
+    array : npt.ArrayLike
+        An object designating an array.
+
+    Returns
+    -------
+    xr.DataArray
+        An xarray with the same content and dtype as the supplied array.
+    """
+    # Determine the Dask array holding the video's data.
+    if isinstance(array, da.Array):
+        data = array
+    else:
+        data = da.from_array(array)
+
+    # Determine the appropriate dims
+    rank = len(data.shape)
+    match rank:
+        case 0:
+            dims = ()
+        case 1:
+            dims = ("X",)
+        case 2:
+            dims = ("Y", "X")
+        case 3:
+            dims = ("T", "Y", "X")
+        case 4:
+            dims = ("T", "Z", "Y", "X")
+        case 5:
+            dims = ("T", "C", "Z", "Y", "X")
+        case 6:
+            dims = ("T", "C", "Z", "Y", "X", "S")
+        case _:
+            raise RuntimeError(f"Cannot interpret {rank}-dimensional data as a video.")
+
+    # Create the xarray.
+    return xr.DataArray(data=data, dims=dims)
+
+
+def dtype_pixel_type(dtype: npt.DTypeLike) -> PixelType:
+    """
+    Return the OME Pixel type corresponding to the supplied dtype.
+
+    Parameters
+    ----------
+    dtype : npt.DTypeLike
+        The Numpy dtype to be used for representing pixel data.
+
+    Returns
+    -------
+    ome_types.model.PixelType
+        A suitable OME pixel type.
+    """
+    match np.dtype(dtype):
+        case np.int8:
+            return PixelType.INT8
+        case np.int16:
+            return PixelType.INT16
+        case np.int32:
+            return PixelType.INT32
+        case np.uint8:
+            return PixelType.UINT8
+        case np.uint16:
+            return PixelType.UINT16
+        case np.uint32:
+            return PixelType.UINT32
+        case np.float32:
+            return PixelType.FLOAT
+        case np.float64:
+            return PixelType.DOUBLE
+        case np.complex64:
+            return PixelType.COMPLEXFLOAT
+        case np.complex128:
+            return PixelType.COMPLEXDOUBLE
+        # The dtypes np.int64 and np.uint64 have no OME equivalent.
+        case np.int64:
+            return PixelType.BIT
+        case np.uint64:
+            return PixelType.BIT
+        case _:
+            raise RuntimeError(f"Cannot interpret {dtype} as a OME pixel type.")
+
+
+def dtype_is_big_endian(dtype: npt.DTypeLike) -> bool:
+    """
+    Return whether the video's data is stored in big endian byte order.
+
+    Parameters
+    ----------
+    dtype : npt.DTypeLike
+        The Numpy dtype to be used for representing pixel data.
+
+    Returns
+    -------
+    bool
+        True when the dtype is big endian, False otherwise.
+    """
+    dtype = np.dtype(dtype)
+
+    if dtype.itemsize == 1:
+        return False
+
+    match np.dtype(dtype).byteorder:
+        case ">":
+            return True
+        case "<":
+            return False
+        case "=":
+            return sys.byteorder == "big"
+        case _:
+            raise RuntimeError(f"Cannot determine endianness of {dtype}.")

toolsandogh/_generate_video.py

@@ -0,0 +1,72 @@
+import dask.array as da
+import numpy as np
+import numpy.typing as npt
+import xarray as xr
+
+from ._canonicalize_video import canonicalize_video
+
+
+def generate_video(
+    T: int = 1,
+    C: int = 1,
+    Z: int = 1,
+    Y: int = 1,
+    X: int = 1,
+    dt: float = 1.0,
+    dz: float = 1.0,
+    dy: float = 1.0,
+    dx: float = 1.0,
+    dtype: npt.DTypeLike = np.float32,
+) -> xr.DataArray:
+    """
+    Generate a video filled with random noise.
+
+    Parameters
+    ----------
+    T : int
+        The T (time) extent of the video.
+    C : int
+        The C (channel) extent of the video.
+    Z : int
+        The Z (height) extent of the video.
+    Y : int
+        The Y (row) extent of the video.
+    X : int
+        The X (column) extent of the video.
+    dt : float
+        The T (time) step size of the video.
+    dz : float
+        The Z (height) step size of the video.
+    dy : float
+        The Y (row) step size of the video.
+    dx : float
+        The X (column) step size of the video.
+    dtype : npt.DtypeLike
+        The dtype of the resulting video.
+
+    Returns
+    -------
+    xarray.DataArray
+        A TCZYX video with the supplied parameters.
+    """
+    rng = da.random.default_rng()
+    video = canonicalize_video(
+        rng.random(
+            size=(T, C, Z, Y, X),
+            dtype=dtype,  # type: ignore
+        ),
+    )
+
+    # Check that the video matches the supplied parameters.
+    for dim, size, step in zip("TCZYX", (T, C, Z, Y, X), (dt, None, dz, dy, dx)):
+        coord = video[dim]
+        assert len(coord) == size
+        if len(coord) > 1 and step is not None:
+            array = coord.values
+            delta = array[1] - array[0]
+            error = abs(delta - step)
+            assert (error / step) <= 1e-3
+    assert video.dtype == dtype
+
+    # Done.
+    return video

toolsandogh/_load_video.py

@@ -0,0 +1,263 @@
+"""Define a function for loading videos from a wide range of sources."""
+
+import os
+import os.path
+import pathlib
+import urllib.parse
+
+import bioio
+import bioio_czi
+import bioio_nd2
+import dask.array as da
+import fsspec
+import numpy as np
+import numpy.typing as npt
+import xarray as xr
+from fsspec.utils import math
+
+from ._canonicalize_video import canonicalize_video
+
+
+def load_video(
+    path: str | os.PathLike,
+    T: int | None = None,
+    C: int | None = None,
+    Z: int | None = None,
+    Y: int | None = None,
+    X: int | None = None,
+    dt: float | None = None,
+    dz: float | None = None,
+    dy: float | None = None,
+    dx: float | None = None,
+    dtype: npt.DTypeLike | None = None,
+) -> xr.DataArray:
+    """
+    Load a video from the given path.
+
+    The resulting video is a 5D xarray with dimensions TCZYX.  If any keyword
+    arguments are supplied, this function asserts that the resulting video
+    matches these arguments, e.g., Z=1 may be used assert that the result is a
+    2D video.
+
+    The suffix of the supplied path determines the method that is used for
+    reading the file.  For .raw and .bin files, the X, Y arguments are
+    mandatory, the number of frames T is inferred automatically, and the
+    remaining arguments default to Z=1, C=1, dz=1.0, dy=dx=1.0, dt=1.0, and
+    dtype=uint8.
+
+    Parameters
+    ----------
+    path : os.PathLike
+        The name of a file, a URI, or a path.
+    T : int
+        The expected T (time) extent of the video.
+    C : int
+        The expected C (channel) extent of the video.
+    Z : int
+        The expected Z (height) extent of the video.
+    Y : int
+        The expected Y (row) extent of the video.
+    X : int
+        The expected X (column) extent of the video.
+    dt : float
+        The expected T (time) step size of the video.
+    dz : float
+        The expected Z (height) step size of the video.
+    dy : float
+        The expected Y (row) step size of the video.
+    dx : float
+        The expected X (column) step size of the video.
+    dtype : npt.DtypeLike
+        The expected dtype of the video.
+
+    Returns
+    -------
+    xarray.DataArray
+        A canonical TCZYX array that matches the supplied parameters.
+    """
+    # Parse the supplied path.
+    pathstr = str(path)
+    url = urllib.parse.urlparse(pathstr)
+    path = pathlib.Path(url.path)
+
+    # Gather all keyword arguments for those load functions that require them.
+    kwargs = {
+        "T": T,
+        "C": C,
+        "Z": Z,
+        "Y": Y,
+        "X": X,
+        "dt": dt,
+        "dz": dz,
+        "dy": dy,
+        "dx": dx,
+        "dtype": dtype,
+    }
+
+    # Load the video.
+    protocols = fsspec.available_protocols()
+    match (url.scheme or "file", path.suffix):
+        case (scheme, ".bin" | ".raw") if scheme in protocols:
+            video = load_raw_video(path, **kwargs)
+        case (_, ".czi"):
+            img = bioio.BioImage(pathstr, reader=bioio_czi.Reader)
+            video = img.xarray_dask_data
+        case (_, ".nd2"):
+            img = bioio.BioImage(pathstr, reader=bioio_nd2.Reader)
+            video = img.xarray_dask_data
+        case (_, _):
+            # Let bioio figure out the rest or raise an exception
+            img = bioio.BioImage(pathstr)
+            video = img.xarray_dask_data
+
+    # Ensure the video is canonical and return.
+    return canonicalize_video(video)
+
+
+def load_raw_video(
+    path: os.PathLike,
+    T: int | None = None,
+    C: int | None = None,
+    Z: int | None = None,
+    Y: int | None = None,
+    X: int | None = None,
+    dt: float | None = None,
+    dz: float | None = None,
+    dy: float | None = None,
+    dx: float | None = None,
+    dtype: npt.DTypeLike | None = None,
+) -> xr.DataArray:
+    """
+    Load a video from the specified raw file.
+
+    Parameters
+    ----------
+    path : os.PathLike
+        The name of a file, a URI, or a path.
+    T : int
+        The expected T (time) extent of the video.
+    C : int
+        The expected C (channel) extent of the video.
+    Z : int
+        The expected Z (height) extent of the video.
+    Y : int
+        The expected Y (row) extent of the video.
+    X : int
+        The expected X (column) extent of the video.
+    dt : float
+        The expected T (time) step size of the video.
+    dz : float
+        The expected Z (height) step size of the video.
+    dy : float
+        The expected Y (row) step size of the video.
+    dx : float
+        The expected X (column) step size of the video.
+    dtype : npt.DtypeLike
+        The expected dtype of the video.
+
+    Returns
+    -------
+    xarray.DataArray
+        A canonical TCZYX array that matches the supplied parameters.
+    """
+    # The parameters X and Y are mandatory for loading raw files.
+    if Y is None:
+        raise RuntimeError("The parameter Y is required for loading raw files.")
+    if X is None:
+        raise RuntimeError("The parameter X is required for loading raw files.")
+    # Z and C default to one.
+    if Z is None:
+        Z = 1
+    if C is None:
+        C = 1
+    # The default dtype is np.uint8
+    if dtype is None:
+        dtype = np.uint8
+    # The number of frames T can be inferred from the file's size.
+    bits_per_item = 12 if (dtype == "uint12") else (np.dtype(dtype).itemsize * 8)
+    nbits = os.path.getsize(path) * 8
+    nitems = nbits // bits_per_item
+    items_per_T = C * Z * Y * X
+    if T is None:
+        T = nitems // items_per_T
+    nexpected = T * items_per_T
+    if nexpected > nitems:
+        raise RuntimeError(
+            f"The file {path} has only {nitems} items, but {nexpected} were expected."
+        )
+    # All step sizes default to 1.0.
+    if dt is None:
+        dt = 1.0
+    if dz is None:
+        dz = 1.0
+    if dy is None:
+        dy = dx or 1.0
+    if dx is None:
+        dx = dy or 1.0
+    # Load the file as a dask array.  Treat the uint12 case specially.
+    shape = (T, C, Z, Y, X)
+    chunks = ("auto", "auto", None, None, None)
+    if dtype == "uint12":
+        nbytes = (nexpected * 12) // 8
+        mmap_array = np.memmap(path, dtype=np.uint8, mode="r", shape=(nbytes,))
+        dask_array = da.from_array(mmap_array)
+        a = dask_array[0::3].astype(np.uint16)
+        b = dask_array[1::3].astype(np.uint16)
+        c = dask_array[2::3].astype(np.uint16)
+        assert len(a) == len(b) == len(c)
+        evens = ((b & 0x0F) << 8) | a
+        odds = ((b & 0xF0) >> 4) | (c << 4)
+        flat = da.stack([evens, odds], axis=1).ravel()
+        dask_array = flat.reshape((T, C, Z, Y, X), chunks=chunks)
+    else:
+        dtype = np.dtype(dtype)
+        mmap_array = np.memmap(path, dtype=dtype, mode="r", shape=shape)
+        dask_array = da.from_array(mmap_array, chunks=chunks)  # type: ignore
+    # Wrap the dask array as a xarray.DataArray and return it.
+    return xr.DataArray(
+        dask_array,
+        coords={
+            "T": dt * np.arange(T),
+            "C": range(C),
+            "Z": dz * np.arange(Z),
+            "Y": dy * np.arange(Y),
+            "X": dx * np.arange(X),
+        },
+    )
+
+
+def load_raw_chunk(
+    path: os.PathLike,
+    shape: tuple[int, ...],
+    dtype: npt.DTypeLike,
+    offset: int,
+) -> np.ndarray:
+    """
+    Load a portion of the supplied raw file as a Numpy array.
+
+    Parameters
+    ----------
+    path : os.PathLike
+        Path (local or remote) to the raw file.
+    shape : tuple[int, ...]
+        Desired shape of the returned array (e.g. ``(T, C, Z, Y, X)``).  The
+        product of the dimensions must match the number of items that will be
+        read.
+    dtype : npt.DTypeLike
+        Numpy data type of the stored items.
+    offset : int
+        Index of the first element to read **in items**, *not* in bytes.
+
+    Returns
+    -------
+    np.ndarray
+        A dense NumPy array with the requested shape and dtype.
+    """
+    count = math.prod(shape)
+    with fsspec.open(path, newline="") as f:
+        return np.fromfile(
+            f,  # type: ignore
+            dtype=dtype,
+            offset=offset,
+            count=count,
+        ).reshape(shape)