Module: micpy.bin¶

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)

    @property
    def spacing_cm(self) -> Tuple[float, float, float] | None:
        """Get the spacing in centimeters."""
        return self.spacing

    @property
    def spacing_mm(self) -> Tuple[float, float, float] | None:
        """Get the spacing in millimeters."""
        if self.spacing is None:
            return None
        return tuple(x * 1e-1 for x in self.spacing)

    @property
    def spacing_m(self) -> Tuple[float, float, float] | None:
        """Get the spacing in meters."""
        if self.spacing is None:
            return None
        return tuple(x * 1e-2 for x in self.spacing)

    @property
    def spacing_um(self) -> Tuple[float, float, float] | None:
        """Get the spacing in micrometers."""
        if self.spacing is None:
            return None
        return tuple(x * 1e4 for x in self.spacing)

    @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 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_index: int,
        field_size: int,
        shape=None,
        spacing=None,
    ) -> "Field":
        """Read a field from a binary file."""

        if field_index < 0:
            start = t()
            _debug("Indexing file...")
            file.seek(0, 2)
            file_size = file.tell()
            field_count = file_size // field_size
            field_index += field_count
            _debug(f"Index completed in {t() - start:.2f} seconds.")

        start = t()
        _debug(f"Seeking to field index {field_index}...")
        offset = field_size * field_index
        file.seek(offset)
        _debug(f"Seek completed in {t() - start:.2f} seconds.")

        start = t()
        _debug(f"Reading data for field index {field_index}...")
        field_data = file.read(field_size)
        if len(field_data) < field_size:
            raise EOFError("Unexpected end of file")
        _debug(f"Read completed in {t() - start:.2f} seconds.")

        start = t()
        _debug(f"Creating Field object for field index {field_index}...")
        field = Field.from_bytes(field_data, shape=shape, spacing=spacing)
        _debug(f"Field object created in {t() - 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)

    def plot(
        self,
        axis: str = "y",
        index: int = 0,
        length_unit: str = "um",
        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: "Axes" = None,
        cax: "Axes" = None,
        vmin: Optional[float] = None,
        vmax: Optional[float] = None,
        cmap: str = "micpy",
        alpha: float = 1.0,
        interpolation: str = "none",
        scalebar: bool = True,
        scalebar_fraction: float = 0.2,
        scalebar_loc: str = "lower right",
        scalebar_pad: float = 0.3,
        scalebar_borderpad: float = 0.0,
        scalebar_sep: float = 2,
        scalebar_frameon: bool = True,
        scalebar_alpha: float = 0.5,
    ) -> Tuple["Figure", "Axes", "Colorbar"]:
        """Plot a slice of the field using Matplotlib.

        Args:
            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`.
            length_unit (str, optional): Physical length unit for axes and scalebar.
                One of `um`, `mm`, `cm`, or `m`. Defaults to `um`.
            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 (Axes, optional): Axes of the plot. Defaults to `None`.
            cax (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`.
            scalebar (bool, optional): `True` if a scale bar should be added,
                `False` otherwise. Defaults to `True`.
            scalebar_fraction (float, optional): Fraction of the visible width used
                for the automatically chosen scalebar length. Defaults to `0.2`.
            scalebar_loc (str, optional): Location of the scale bar.
                Defaults to `lower right`.
            scalebar_pad (float, optional): Padding between the scale bar and the plot.
                Defaults to `0.3`.
            scalebar_borderpad (float, optional): Padding between the frame and content.
                Defaults to `0.0`.
            scalebar_sep (float, optional): Separation between scale bar and label.
                Defaults to `2`.
            scalebar_frameon (bool, optional): `True` if the scale bar should have a
                background frame, `False` otherwise. Defaults to `True`.
            scalebar_alpha (float, optional): Transparency of the scale bar
                background. Defaults to `0.5`.

        Returns:
            Matplotlib figure, axes, and color bar.
        """
        if not HAS_MATPLOTLIB:
            raise ImportError("Matplotlib is not installed.")

        if self.ndim != 3:
            raise ValueError("'field' must be a 3D array")

        if self.spacing is None:
            raise ValueError("'spacing' is required for plotting.")

        if axis == "z":
            x, y = "x", "y"
            slice_2d = self[index, :, :]
        elif axis == "y":
            x, y = "x", "z"
            slice_2d = self[:, index, :]
        elif axis == "x":
            x, y = "y", "z"
            slice_2d = self[:, :, index]
        else:
            raise ValueError("'axis' must be 'x', 'y', or 'z'")

        fig, ax = (
            pyplot.subplots(figsize=figsize, dpi=dpi)
            if ax is None
            else (ax.get_figure(), ax)
        )

        unit_scale = _length_unit_scale_from_cm(length_unit)
        unit_label = _length_unit_label(length_unit)

        spacing_native_cm = self.spacing[0]
        spacing_display = spacing_native_cm * unit_scale

        ny, nx = slice_2d.shape
        width_display = nx * spacing_display
        height_display = ny * spacing_display
        extent = (0.0, width_display, 0.0, height_display)

        if title is not None:
            ax.set_title(title)
        else:
            ax.set_title(f"t={np.round(self.time, 7)}s")

        ax.set_xlabel(xlabel if xlabel is not None else f"{x} [{unit_label}]")
        ax.set_ylabel(ylabel if ylabel is not None else f"{y} [{unit_label}]")

        if aspect is not None:
            ax.set_aspect(aspect)

        ax.set_frame_on(False)

        image = ax.imshow(
            slice_2d,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            interpolation=interpolation,
            alpha=alpha,
            extent=extent,
            origin="lower",
        )

        if scalebar:
            scalebar_length = _nice_length(width_display * scalebar_fraction)

            sbar = AnchoredSizeBar(
                transform=ax.transData,
                size=scalebar_length,
                label=_format_length_in_unit(scalebar_length, length_unit),
                loc=scalebar_loc,
                pad=scalebar_pad,
                borderpad=scalebar_borderpad,
                sep=scalebar_sep,
                frameon=scalebar_frameon,
            )

            ax.add_artist(sbar)
            ax.scalebar = sbar

            if scalebar_frameon:
                patch = sbar.patch
                patch.set_facecolor(ax.get_facecolor())
                patch.set_alpha(scalebar_alpha)
                patch.set_edgecolor("none")

        cbar = pyplot.colorbar(image, ax=ax, cax=cax)
        cbar.locator = matplotlib.ticker.MaxNLocator(
            integer=np.issubdtype(slice_2d.dtype, np.integer)
        )
        cbar.outline.set_visible(False)
        cbar.update_ticks()

        return fig, ax, cbar

    def as_vti(
        self,
        name: str = "values",
        point_data: bool = False,
    ) -> "vtkImageData":
        """Convert the field to a VTK ImageData object.

        Args:
            name (str, optional): Name of the data array. Defaults to `values`.
            point_data (bool, optional): `True` if data should be stored as PointData,
                `False` if data should be stored as CellData. Defaults to `False`.
        Returns:
            VTK ImageData object.
        """

        if not HAS_VTK:
            raise ImportError("VTK is not installed.")

        cells = np.asarray(self)
        if cells.ndim != 3:
            raise ValueError("field must be 3D (nz, ny, nx)")

        nz, ny, nx = cells.shape
        dz, dy, dx = self.spacing

        def cell_to_point_average(c: np.ndarray) -> np.ndarray:
            nz_, ny_, nx_ = c.shape
            acc = np.zeros(
                (nz_ + 1, ny_ + 1, nx_ + 1), dtype=np.result_type(c.dtype, np.float32)
            )
            w = np.zeros((nz_ + 1, ny_ + 1, nx_ + 1), dtype=np.float32)

            acc[0:nz_, 0:ny_, 0:nx_] += c
            acc[0:nz_, 0:ny_, 1 : nx_ + 1] += c
            acc[0:nz_, 1 : ny_ + 1, 0:nx_] += c
            acc[0:nz_, 1 : ny_ + 1, 1 : nx_ + 1] += c
            acc[1 : nz_ + 1, 0:ny_, 0:nx_] += c
            acc[1 : nz_ + 1, 0:ny_, 1 : nx_ + 1] += c
            acc[1 : nz_ + 1, 1 : ny_ + 1, 0:nx_] += c
            acc[1 : nz_ + 1, 1 : ny_ + 1, 1 : nx_ + 1] += c

            w[0:nz_, 0:ny_, 0:nx_] += 1
            w[0:nz_, 0:ny_, 1 : nx_ + 1] += 1
            w[0:nz_, 1 : ny_ + 1, 0:nx_] += 1
            w[0:nz_, 1 : ny_ + 1, 1 : nx_ + 1] += 1
            w[1 : nz_ + 1, 0:ny_, 0:nx_] += 1
            w[1 : nz_ + 1, 0:ny_, 1 : nx_ + 1] += 1
            w[1 : nz_ + 1, 1 : ny_ + 1, 0:nx_] += 1
            w[1 : nz_ + 1, 1 : ny_ + 1, 1 : nx_ + 1] += 1

            return acc / w

        if point_data:
            data = cell_to_point_average(cells)
        else:
            data = cells

        data_c = np.ascontiguousarray(data)
        flat = data_c.ravel(order="C")

        image = vtkImageData()
        image.SetOrigin(0.0, 0.0, 0.0)
        image.SetSpacing(float(dx), float(dy), float(dz))
        image.SetDimensions(nx + 1, ny + 1, nz + 1)

        vtk_arr = numpy_to_vtk(flat, deep=False)
        vtk_arr.SetName(name)

        if point_data:
            image.GetPointData().SetScalars(vtk_arr)
        else:
            image.GetCellData().SetScalars(vtk_arr)

        image._numpy_pin = data_c

        return image

    def save_vti(
        self,
        filename: str,
        name: str = "values",
        point_data: bool = False,
    ) -> str:
        """Save the field as a VTK ImageData file.

        Args:
            filename (str): Filename of the VTK ImageData file.
            name (str, optional): Name of the data array. Defaults to `values`.
            point_data (bool, optional): `True` if data should be stored as PointData
                `False` if data should be stored as CellData. Defaults to `False`.
        Returns:
            Filename of the VTK ImageData file.
        """

        image = self.as_vti(name=name, point_data=point_data)

        filename = str(Path(filename).with_suffix(".vti"))

        writer = vtkXMLImageDataWriter()
        writer.SetFileName(filename)

        writer.SetInputData(image)
        writer.SetDataModeToAppended()
        writer.EncodeAppendedDataOff()

        ok = writer.Write()
        if ok != 1:
            raise OSError(f"Failed to write {filename}")

        return filename