micress-micpy 0.3.2b2__py3-none-any.whl → 0.4.0a1__py3-none-any.whl

micpy/bin copy.py

@@ -0,0 +1,865 @@
+"""The `micpy.bin` module provides methods to read and write binary files."""
+
+from collections.abc import Callable as ABCCallable
+from dataclasses import dataclass
+from pathlib import Path
+from time import time
+from typing import Callable, Generator, IO, List, Optional, Tuple, Union, overload
+import os
+import sys
+
+import gzip
+import rapidgzip
+
+import numpy as np
+
+from micpy import geo
+from micpy import utils
+from micpy.matplotlib import matplotlib, pyplot
+
+
+__all__ = ["File", "Field", "Series", "plot", "PlotArgs"]
+
+
+class Footer:
+    """A field footer."""
+
+    TYPE = [("length", np.int32)]
+    SIZE = np.dtype(TYPE).itemsize
+
+    def __init__(self, length: int = 0):
+        data = np.array((length,), dtype=self.TYPE)
+        self.body_length = data["length"]
+
+    def to_bytes(self):
+        """Convert the footer to bytes."""
+        return np.array((self.body_length,), dtype=self.TYPE).tobytes()
+
+
+class Header:
+    """A field header."""
+
+    TYPE = [("size", np.int32), ("time", np.float32), ("length", np.int32)]
+    SIZE = np.dtype(TYPE).itemsize
+
+    def __init__(self, size: int, time: float, length: int):
+        self.size = size
+        self.time = round(float(time), 7)
+        self.body_length = length
+
+        self.field_size = Header.SIZE + 4 * self.body_length + Footer.SIZE
+
+        if not self.size == self.field_size - 8:
+            raise ValueError("Invalid header")
+
+    @staticmethod
+    def from_bytes(data: bytes):
+        """Create a new header from bytes."""
+        kwargs = np.frombuffer(data[: Header.SIZE], dtype=Header.TYPE)
+        return Header(*kwargs[0].item())
+
+    def to_bytes(self):
+        """Convert the header to bytes."""
+        return np.array(
+            (self.size, self.time, self.body_length), dtype=Header.TYPE
+        ).tobytes()
+
+    @staticmethod
+    def read(file: IO[bytes]) -> "Header":
+        """Read the header of a binary file."""
+        file.seek(0)
+        data = file.read(Header.SIZE)
+        return Header.from_bytes(data)
+
+    def get_field_count(self, file_size: int) -> int:
+        """Get the number of fields in the file."""
+        return file_size // self.field_size
+
+
+class Field(np.ndarray):
+    """A field."""
+
+    def __new__(cls, data, time: float, spacing: Tuple[float, float, float]):
+        obj = np.asarray(data).view(cls)
+        obj.time = time
+        obj.spacing = spacing
+        return obj
+
+    def __array_finalize__(self, obj):
+        if obj is None:
+            return
+
+        # pylint: disable=attribute-defined-outside-init
+        self.time = getattr(obj, "time", None)
+        self.spacing = getattr(obj, "spacing", None)
+
+    @staticmethod
+    def from_bytes(data: bytes, shape=None, spacing=None):
+        """Create a new time step from bytes."""
+
+        header = Header.from_bytes(data)
+
+        start, end, count = (
+            header.SIZE,
+            header.field_size,
+            header.body_length,
+        )
+
+        data = data[start:end]
+        data = np.frombuffer(data, count=count, dtype="float32")
+        if np.all(np.isclose(data, data.astype("int32"))):
+            data = data.astype("int32")
+
+        if shape is not None:
+            data = data.reshape(shape)
+
+        return Field(data, time=header.time, spacing=spacing)
+
+    @staticmethod
+    def from_bytes2(data: bytes, shape=None, spacing=None):
+        """Create a new time step from bytes."""
+
+        header = Header.from_bytes(data)
+
+        start, end, count = (
+            header.SIZE,
+            header.field_size,
+            header.body_length,
+        )
+
+        data = data[start:end]
+        data = np.frombuffer(data, count=count, dtype="float32")
+        # if np.all(np.isclose(data, data.astype("int32"))):
+        #     data = data.astype("int32")
+
+        if shape is not None:
+            data = data.reshape(shape)
+
+        return Field(data, time=header.time, spacing=spacing)
+
+    def to_bytes(self):
+        """Convert the field to bytes."""
+        header = Header(
+            size=self.size * self.itemsize + 8,
+            time=self.time,
+            length=self.size,
+        )
+        footer = Footer(length=self.size)
+
+        return header.to_bytes() + self.tobytes() + footer.to_bytes()
+
+    @staticmethod
+    def dimensions(shape: Tuple[int, int, int]) -> int:
+        """Get the number of dimensions of a shape."""
+        x, y, _ = shape
+
+        if y == 1:
+            if x == 1:
+                return 1
+            return 2
+        return 3
+
+    @staticmethod
+    def read(
+        file: IO[bytes],
+        field_id: int,
+        field_size: int,
+        field_count: int,
+        shape=None,
+        spacing=None,
+    ) -> "Field":
+        """Read a field from a binary file."""
+
+        if field_id < 0:
+            field_id += field_count
+
+        if not 0 <= field_id < field_count:
+            raise IndexError("field_id out of range")
+
+        offset = field_size * field_id
+        file.seek(offset)
+
+        field_data = file.read(field_size)
+        if len(field_data) < field_size:
+            raise EOFError("Unexpected end of file")
+
+        return Field.from_bytes(field_data, shape=shape, spacing=spacing)
+
+    @staticmethod
+    def read2(
+        file: IO[bytes],
+        field_id: int,
+        field_size: int,
+        shape=None,
+        spacing=None,
+    ) -> "Field":
+        """Read a field from a binary file."""
+
+        if field_id < 0:
+            start = time()
+            info("Indexing entire file...")
+            file.seek(0, 2)
+            file_size = file.tell()
+            field_count = file_size // field_size
+            field_id += field_count
+            info(f"Indexing completed in {time() - start:.2f} seconds.")
+
+        start = time()
+        info(f"Seeking to field ID {field_id}...")
+        offset = field_size * field_id
+        file.seek(offset)
+        info(f"Seek completed in {time() - start:.2f} seconds.")
+
+        start = time()
+        info(f"Reading data for field ID {field_id}...")
+        field_data = file.read(field_size)
+        if len(field_data) < field_size:
+            raise EOFError("Unexpected end of file")
+        info(f"Read completed in {time() - start:.2f} seconds.")
+
+        start = time()
+        info(f"Creating Field object for field ID {field_id}...")
+        field = Field.from_bytes2(field_data, shape=shape, spacing=spacing)
+        info(f"Field object created in {time() - start:.2f} seconds.")
+        return field
+
+    def to_file(self, file: IO[bytes], geometry: bool = True):
+        """Write the field to a binary file.
+
+        Args:
+            file (IO[bytes]): Binary file.
+            geometry (bool, optional): `True` if geometry should be written, `False`
+                otherwise. Defaults to `True`.
+        """
+
+        file.write(self.to_bytes())
+
+        if geometry:
+            geo_filename, geo_data = geo.build(file.name, self.shape, self.spacing)
+            geo.write(geo_filename, geo_data, geo.Type.BASIC)
+
+    def write(
+        self,
+        filename: str,
+        compressed: bool = True,
+        geometry: bool = True,
+    ):
+        """Write the field to a binary file.
+
+        Args:
+            filename (str): Filename of the binary file.
+            compressed (bool, optional): `True` if file should be compressed, `False`
+                otherwise. Defaults to `True`.
+            geometry (bool, optional): `True` if geometry should be written, `False`
+                otherwise. Defaults to `True`.
+        """
+
+        file_open = gzip.open if compressed else open
+        with file_open(filename, "wb") as file:
+            self.to_file(file, geometry)
+
+
+class Series(np.ndarray):
+    def __new__(cls, fields: List[Field]):
+        obj = np.asarray(fields).view(cls)
+        obj.times = [field.time for field in fields]
+        obj.spacings = [field.spacing for field in fields]
+        return obj
+
+    def __array_finalize__(self, obj):
+        if obj is None:
+            return
+
+        # pylint: disable=attribute-defined-outside-init
+        self.times = getattr(obj, "times", None)
+        self.spacings = getattr(obj, "spacings", None)
+
+    def iterate_fields(self):
+        """Iterate over fields in the series.
+
+        Yields:
+            Field.
+        """
+        for item, time, spacing in zip(self, self.times, self.spacings):
+            yield Field(item, time, spacing)
+
+    def get_field(self, index: int) -> Field:
+        """Get a field from the series.
+
+        Args:
+            index (int): Index of the field.
+
+        Returns:
+            Field.
+        """
+        return Field(self[index], self.times[index], self.spacings[index])
+
+    def get_series(self, key: Union[int, slice, list]) -> "Series":
+        """Get a series of fields.
+
+        Args:
+            key (Union[int, slice, list]): Key to list of field IDs, a slice object, or a
+                list of field IDs.
+
+        Returns:
+            Series of fields.
+        """
+        if isinstance(key, int):
+            return Series([self.get_field(key)])
+        if isinstance(key, slice):
+            return Series([self.get_field(i) for i in range(*key.indices(len(self)))])
+        if isinstance(key, list):
+            return Series([self.get_field(i) for i in key])
+        raise TypeError("Invalid argument type")
+
+    def write(self, filename: str, compressed: bool = True, geometry: bool = True):
+        """Write the series to a binary file.
+
+        Args:
+            filename (str): Filename of the binary file.
+            compressed (bool, optional): `True` if file should be compressed, `False`
+                otherwise. Defaults to `True`.
+            geometry (bool, optional): `True` if geometry should be written, `False`
+                otherwise. Defaults to `True`.
+        """
+
+        file_open = gzip.open if compressed else open
+        with file_open(filename, "wb") as file:
+            for field in self.iterate_fields():
+                field.to_file(file, geometry)
+                geometry = False
+
+
+verbose = False
+
+
+def info(*args):
+    if verbose:
+        print(*args)
+
+
+def warn(*args):
+    if verbose:
+        print(*args)  # , file=sys.stderr)
+
+
+def index_filename(filename: str) -> Path:
+    return Path(f"{filename}.idx")
+
+
+def import_index(file, index: str):
+    path = Path(index)
+    if not path.is_file():
+        warn("Index file not found.")
+        return
+    if not os.access(path, os.R_OK):
+        warn("Index file is not readable.")
+        return
+    try:
+        file.import_index(str(path))
+        info("Index file imported.")
+    except ValueError:
+        warn("Index file is invalid.")
+        return
+    except (PermissionError, FileNotFoundError, IsADirectoryError, OSError):
+        warn("Failed to import index file.")
+        return
+
+
+def export_index(file, index: str):
+    path = Path(index)
+
+    if path.exists() and not path.is_file():
+        warn("Index path is not a file.")
+        return
+    if not os.access(path.parent, os.W_OK):
+        warn("Index file is not writable.")
+        return
+    try:
+        file.export_index(str(path))
+        info("Index file exported.")
+    except (PermissionError, IsADirectoryError, FileNotFoundError, OSError):
+        warn("Failed to export index file.")
+        return
+
+
+def read_field(
+    filename: str,
+    field_id: int,
+    shape: tuple = None,
+    spacing: tuple = None,
+    threads: int = None,
+) -> Field:
+
+    if not os.path.isfile(filename):
+        raise FileNotFoundError(f"File not found: {filename}")
+
+    threads = threads if threads is not None else min(4, os.cpu_count())
+
+    if shape is None or spacing is None:
+        try:
+            geometry = geo.read(geo.find(filename), type=geo.Type.BASIC)
+            shape = shape or geometry["shape"][::-1]
+            spacing = spacing or geometry["spacing"][::-1]
+        except (geo.GeometryFileNotFoundError, geo.MultipleGeometryFilesError):
+            warn("Caution: A geometry file was not found.")
+
+    compressed = utils.is_compressed(filename)
+
+    if compressed:
+        file = rapidgzip.open(filename, threads)
+    else:
+        file = open(filename, "rb")
+
+    if compressed:
+        import_index(file, index_filename(filename))
+        info(f"Imported file size: {file.size()} bytes")
+
+    field = Field.read2(
+        file=file,
+        field_id=field_id,
+        field_size=Header.read(file).field_size,
+        shape=shape,
+        spacing=spacing,
+    )
+
+    if compressed:
+        export_index(file, index_filename(filename))
+        info(f"Exported file size: {file.size()} bytes")
+
+    return field
+
+
+class File:
+    """A binary file."""
+
+    def __init__(
+        self,
+        filename: str,
+        parallelization: int = None,
+        verbose: bool = True,
+    ):
+        """Initialize a binary file.
+
+        Args:
+            filename (str): File name.
+            parallelization (int, optional): Number of threads for reading compressed
+                files. Defaults to `None` (auto).
+            verbose (bool, optional): Verbose output. Defaults to `True`.
+
+        Raises:
+            `FileNotFoundError`: If file is not found.
+        """
+        if not os.path.isfile(filename):
+            raise FileNotFoundError(f"File not found: {filename}")
+
+        self._filename: str = filename
+        self._parallelization: int = (
+            parallelization if parallelization is not None else min(4, os.cpu_count())
+        )
+        self._verbose: bool = verbose
+
+        self._file: IO[bytes] = None
+        self._compressed: bool = None
+        self._header: Header = None
+        self._file_size: int = None
+
+        self.shape: np.ndarray[(3,), np.int32] = None
+        self.spacing: np.ndarray[(3,), np.float32] = None
+
+        try:
+            self.find_geometry()
+        except (geo.GeometryFileNotFoundError, geo.MultipleGeometryFilesError):
+            self._warn("Caution: A geometry file was not found.")
+
+    def __getitem__(self, key: Union[int, slice, list, Callable[[Field], bool]]):
+        return self.read(key)
+
+    def __enter__(self):
+        return self.open()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
+    def __iter__(self):
+        if not self._file:
+            self._open_file()
+        return self.iterate()
+
+    def _info(self, *args):
+        if self._verbose:
+            print(*args)
+
+    def _warn(self, *args):
+        if self._verbose:
+            print(*args, file=sys.stderr)
+
+    def _open_file(self):
+        self._compressed = utils.is_compressed(self._filename)
+
+        if self._compressed:
+            self._file = rapidgzip.open(self._filename, self._parallelization)
+        else:
+            self._file = open(self._filename, "rb")
+
+        self.import_index()
+        self.export_index()
+
+        self._header = Header.read(self._file)
+
+    def _close_file(self):
+        if self._file:
+            self._file.close()
+        self._reset()
+
+    def _reset(self):
+        self._file = None
+        self._compressed = None
+        self._header = None
+        self._file_size = None
+
+    def open(self):
+        """Open the file."""
+        if not self._file:
+            self._open_file()
+        return self
+
+    def close(self):
+        """Close the file."""
+        self._close_file()
+
+    def times(self) -> List[float]:
+        """Get the times of the fields in the file.
+
+        Returns:
+            List of times.
+        """
+        times = []
+        self._file.seek(0)
+        while True:
+            data = self._file.read(Header.SIZE)
+            if len(data) < Header.SIZE:
+                break
+            header = Header.from_bytes(data)
+            times.append(header.time)
+            self._file.seek(header.field_size - Header.SIZE, 1)
+        return times
+
+    def _index_path(self) -> Path:
+        return Path(f"{self._filename}.idx")
+
+    def import_index(self):
+        """Import index from a file."""
+        if not self._compressed:
+            return
+
+        path = self._index_path()
+
+        if not path.is_file():
+            self._warn("Index file not found.")
+            return
+        if not os.access(path, os.R_OK):
+            self._warn("Index file is not readable.")
+            return
+        try:
+            self._file.import_index(str(path))
+            self._info("Index file imported.")
+        except ValueError:
+            self._warn("Index file is invalid.")
+            return
+        except (PermissionError, FileNotFoundError, IsADirectoryError, OSError):
+            self._warn("Failed to import index file.")
+            return
+        finally:
+            self._info("Creating index...")
+            start = time()
+            self._file.seek(0, 2)
+            self._file_size = self._file.tell()
+            self._info(f"Index created in {time() - start:.2f} seconds.")
+
+    def export_index(self):
+        """Export index to a file."""
+        if not self._compressed:
+            return
+
+        path = self._index_path()
+
+        if path.exists() and not path.is_file():
+            self._warn("Index path is not a file.")
+            return
+        if not os.access(path.parent, os.W_OK):
+            self._warn("Index file is not writable.")
+            return
+        try:
+            self._file.export_index(str(path))
+            self._info("Index file exported.")
+        except (PermissionError, IsADirectoryError, FileNotFoundError, OSError):
+            self._warn("Failed to export index file.")
+            return
+
+    def set_geometry(
+        self, shape: Tuple[int, int, int], spacing: Tuple[float, float, float]
+    ):
+        """Set the geometry.
+
+        Args:
+            shape (Tuple[int, int, int]): Shape of the geometry (z, y, x).
+            spacing (Tuple[float, float, float]): Spacing of the geometry (dz, dy, dx) in μm.
+        """
+
+        self.shape = np.array(shape)
+        self.spacing = np.array(spacing)
+
+        self.print_geometry()
+
+    def read_geometry(self, filename: str, compressed: Optional[bool] = None):
+        """Read geometry from a file.
+
+        Args:
+            filename (str): Filename of a geometry file.
+            compressed (bool, optional): `True` if file is compressed, `False`
+                otherwise. Defaults to `None` (auto).
+        """
+        if compressed is None:
+            compressed = utils.is_compressed(filename)
+
+        geometry = geo.read(filename, type=geo.Type.BASIC, compressed=compressed)
+
+        shape = geometry["shape"][::-1]
+        spacing = geometry["spacing"][::-1]
+
+        self.set_geometry(shape, spacing)
+
+    def find_geometry(self, compressed: Optional[bool] = None):
+        """Find geometry file and read it.
+
+        Args:
+            compressed (bool, optional): True if file is compressed, False otherwise.
+                Defaults to `None` (auto).
+
+        Raises:
+            `GeometryFileNotFoundError`: If no geometry file is found.
+            `MultipleGeometryFilesError`: If multiple geometry files are found.
+        """
+        filename = geo.find(self._filename)
+
+        self.read_geometry(filename, compressed=compressed)
+
+    def print_geometry(self):
+        """Get a string representation of the geometry."""
+
+        if self.shape is None or self.spacing is None:
+            self._info("Geometry: None")
+            return
+
+        dimensions = Field.dimensions(self.shape)
+        cells = self.shape
+        spacing = 1e4 * np.round(self.spacing.astype(float), 7)
+        size = cells * spacing
+
+        self._info(f"Geometry: {dimensions}-Dimensional Grid")
+        self._info(f"Grid Size [μm]: {tuple(size)}")
+        self._info(f"Grid Shape (Cell Count): {tuple(cells)}")
+        self._info(f"Grid Spacing (Cell Size) [μm]: {tuple(spacing)}")
+
+    def read_field(self, field_id: int) -> Field:
+        """Read a field from the file.
+
+        Args:
+            field_id (int): Field ID.
+
+        Returns:
+            Field.
+        """
+
+        return Field.read(
+            file=self._file,
+            field_id=field_id,
+            field_size=self._header.field_size,
+            field_count=self._header.get_field_count(self._file_size),
+            shape=self.shape,
+            spacing=self.spacing,
+        )
+
+    def iterate(self) -> Generator[Field, None, None]:
+        """Iterate over fields in the file.
+
+        Returns:
+            A generator of fields.
+        """
+
+        field_id = 0
+        while True:
+            try:
+                yield self.read_field(field_id)
+                field_id += 1
+            except (IndexError, EOFError):
+                break
+
+    @overload
+    def read(self, key: None = None) -> Series: ...
+    @overload
+    def read(self, key: int) -> Series: ...
+    @overload
+    def read(self, key: slice) -> Series: ...
+    @overload
+    def read(self, key: list[int]) -> Series: ...
+    @overload
+    def read(self, key: Callable[[Field], bool]) -> Series: ...
+
+    def read(
+        self, key: Optional[Union[int, slice, list, Callable[[Field], bool]]] = None
+    ) -> Series:
+        """Read a series of fields from the file.
+
+        Args:
+            key (Union[int, slice, list, Callable[[Field], bool]]): Key to list of
+            field IDs, a slice object, a list of field IDs, or a function that filters
+            fields. Defaults to `None`.
+
+        Returns:
+            Series of fields.
+        """
+
+        def iterable(iterable):
+            if self._verbose:
+                return utils.progress_indicator(
+                    iterable, description="Reading", unit="Field"
+                )
+            return iterable
+
+        if key is None:
+            fields = list(field for field in iterable(self.iterate()))
+        elif isinstance(key, int):
+            fields = [self.read_field(i) for i in iterable([key])]
+        elif isinstance(key, slice):
+            indices = range(*key.indices(self._header.get_field_count(self._file_size)))
+            fields = [self.read_field(i) for i in iterable(indices)]
+        elif isinstance(key, list):
+            fields = [self.read_field(i) for i in iterable(key)]
+        elif isinstance(key, ABCCallable):
+            fields = [field for field in iterable(self.iterate()) if key(field)]
+        else:
+            raise TypeError("Invalid argument type")
+        return Series(fields)
+
+
+@dataclass
+class PlotArgs:
+    """Arguments for plotting a field.
+
+    Args:
+        title (str, optional): Title of the plot. Defaults to `None`.
+        xlabel (str, optional): Label of the x-axis. Defaults to `None`.
+        ylabel (str, optional): Label of the y-axis. Defaults to `None`.
+        figsize (Tuple[float, float], optional): Figure size. Defaults to `None`.
+        dpi (int, optional): Figure DPI. Defaults to `None`.
+        aspect (str, optional): Aspect ratio. Defaults to `equal`.
+        ax (matplotlib.Axes, optional): Axes of the plot. Defaults to `None`.
+        cax (matplotlib.Axes, optional): Axes of the color bar. Defaults to `None`.
+        vmin (float, optional): Minimum value of the color bar. Defaults to `None`.
+        vmax (float, optional): Maximum value of the color bar. Defaults to `None`.
+        cmap (str, optional): Colormap. Defaults to `micpy`.
+        alpha (float, optional): Transparency of the plot. Defaults to `1.0`.
+        interpolation (str, optional): Interpolation method. Defaults to `none`.
+        extent (Tuple[float, float, float, float], optional): Extent of the plot.
+    """
+
+    title: Optional[str] = None
+    xlabel: Optional[str] = None
+    ylabel: Optional[str] = None
+    figsize: Optional[Tuple[float, float]] = None
+    dpi: Optional[int] = None
+    aspect: str = "equal"
+    ax: "matplotlib.Axes" = None
+    cax: "matplotlib.Axes" = None
+    vmin: Optional[float] = None
+    vmax: Optional[float] = None
+    cmap: str = "micpy"
+    alpha: float = 1.0
+    interpolation: str = "none"
+    extent: Optional[Tuple[float, float, float, float]] = None
+
+
+def plot(
+    field: Field,
+    axis: str = "y",
+    index: int = 0,
+    args: Optional[PlotArgs] = None,
+) -> Tuple["matplotlib.Figure", "matplotlib.Axes"]:
+    """Plot a slice of the field.
+
+    Args:
+        field (Field): Field to plot.
+        axis (str, optional): Axis to plot. Possible values are `x`, `y`, and `z`.
+            Defaults to `y`.
+        index (int, optional): Index of the slice. Defaults to `0`.
+        args (PlotArgs, optional): Arguments for plotting. Defaults to `None`.
+
+    Returns:
+        Matplotlib figure, axes, and color bar.
+    """
+
+    if matplotlib is None:
+        raise ImportError("matplotlib is not installed")
+
+    if field.ndim != 3:
+        raise ValueError("Invalid field shape")
+
+    if axis == "z":
+        x, y = "x", "y"
+        slice_2d = field[index, :, :]
+    elif axis == "y":
+        x, y = "x", "z"
+        slice_2d = field[:, index, :]
+    elif axis == "x":
+        x, y = "y", "z"
+        slice_2d = field[:, :, index]
+    else:
+        raise ValueError("Invalid axis")
+
+    if args is None:
+        args = PlotArgs()
+
+    fig, ax = (
+        pyplot.subplots(figsize=args.figsize, dpi=args.dpi)
+        if args.ax is None
+        else (args.ax.get_figure(), args.ax)
+    )
+
+    if args.title is not None:
+        ax.set_title(args.title)
+    else:
+        if isinstance(field, Field):
+            ax.set_title(f"t={np.round(field.time, 7)}s")
+    if args.xlabel is not None:
+        ax.set_xlabel(args.xlabel)
+    else:
+        ax.set_xlabel(x)
+    if args.ylabel is not None:
+        ax.set_ylabel(args.ylabel)
+    else:
+        ax.set_ylabel(y)
+    if args.aspect is not None:
+        ax.set_aspect(args.aspect)
+    ax.set_frame_on(False)
+
+    image = ax.imshow(
+        slice_2d,
+        cmap=args.cmap,
+        vmin=args.vmin,
+        vmax=args.vmax,
+        interpolation=args.interpolation,
+        alpha=args.alpha,
+        extent=args.extent,
+        origin="lower",
+    )
+
+    bar = pyplot.colorbar(image, ax=ax, cax=args.cax)
+    bar.locator = matplotlib.ticker.MaxNLocator(
+        integer=np.issubdtype(slice_2d.dtype, np.integer)
+    )
+    bar.outline.set_visible(False)
+    bar.update_ticks()
+
+    return fig, ax, bar