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 _boundary_indices(
        n: int,
        lower: int,
        upper: int,
        lower_bc: str,
        upper_bc: str,
    ) -> np.ndarray:
        """Map an extended axis to source indices for MICRESS boundary conditions."""
        if n <= 0:
            raise ValueError("Axis length must be positive.")
        if lower < 0 or upper < 0:
            raise ValueError("Boundary extension must be non-negative.")

        valid = {"periodic", "insulating", "symmetric"}
        if lower_bc not in valid or upper_bc not in valid:
            raise ValueError(
                "Boundary condition must be 'periodic', 'insulating', or 'symmetric'."
            )
        if (lower_bc == "periodic") != (upper_bc == "periodic"):
            raise ValueError(
                "Periodic boundary conditions must be used on both sides of an axis."
            )

        positions = np.arange(-lower, n + upper, dtype=np.int64)
        if lower_bc == upper_bc == "periodic":
            return np.mod(positions, n)

        indices = positions.copy()
        while True:
            low = indices < 0
            high = indices >= n
            if not (np.any(low) or np.any(high)):
                return indices

            if np.any(low):
                if lower_bc == "insulating":
                    indices[low] = -indices[low] - 1
                else:  # symmetric
                    indices[low] = -indices[low]

            if np.any(high):
                if upper_bc == "insulating":
                    indices[high] = 2 * n - 1 - indices[high]
                else:  # symmetric
                    indices[high] = 2 * n - 2 - indices[high]

    @classmethod
    def _extend_array(cls, array, extension_zyx, boundary_zyx):
        """Extend the first three axes of an array using MICRESS boundary conditions."""
        values = np.asarray(array)
        if values.ndim < 3:
            raise ValueError("array must have at least three dimensions (z, y, x)")

        indices = []
        for axis, n, pad, bc in zip(
            "ZYX", values.shape[:3], extension_zyx, boundary_zyx
        ):
            lower, upper = (int(v) for v in pad)
            if n == 1 and (lower or upper):
                raise ValueError(
                    f"Cannot extend degenerate {axis} axis (dimension is 1). "
                    "Set its lower/upper extension to 0."
                )
            lower_bc, upper_bc = bc
            indices.append(cls._boundary_indices(n, lower, upper, lower_bc, upper_bc))

        result = values[np.ix_(*indices)]
        # np.ix_ covers the spatial axes only; retain any trailing component axes.
        if values.ndim > 3:
            result = values[np.ix_(*indices, *[np.arange(n) for n in values.shape[3:]])]
        return result

    def extend(
        self,
        x=(0, 0),
        y=(0, 0),
        z=(0, 0),
        boundary_x=("periodic", "periodic"),
        boundary_y=("periodic", "periodic"),
        boundary_z=("periodic", "periodic"),
    ) -> "Field":
        """Return a field extended using MICRESS-style boundary conditions.

        Args:
            x (tuple of int, optional): Lower and upper extension in the x
                direction. Defaults to ``(0, 0)``.
            y (tuple of int, optional): Lower and upper extension in the y
                direction. Defaults to ``(0, 0)``.
            z (tuple of int, optional): Lower and upper extension in the z
                direction. Defaults to ``(0, 0)``.
            boundary_x (tuple of str, optional): Lower and upper boundary
                conditions for x. Defaults to ``("periodic", "periodic")``.
            boundary_y (tuple of str, optional): Lower and upper boundary
                conditions for y. Defaults to ``("periodic", "periodic")``.
            boundary_z (tuple of str, optional): Lower and upper boundary
                conditions for z. Defaults to ``("periodic", "periodic")``.

                ``"insulating"`` mirrors while duplicating the outermost
                sample, ``"symmetric"`` mirrors without duplicating it, and
                ``"periodic"`` wraps. Extensions and boundary pairs are
                specified in physical x/y/z order.

        Returns:
            Field: Extended field with the original time and spacing.

        Raises:
            ValueError: If the field is not three-dimensional, an extension is
                negative, a boundary condition is invalid, or periodic
                conditions are not used on both sides of an axis.
        """
        if self.ndim != 3:
            raise ValueError("field must be 3D (nz, ny, nx)")

        data = self._extend_array(
            np.asarray(self),
            extension_zyx=(z, y, x),
            boundary_zyx=(boundary_z, boundary_y, boundary_x),
        )
        return Field(data, time=self.time, spacing=self.spacing)

    def resample(
        self,
        shape: Tuple[int, int, int],
        method: str = "nearest",
    ) -> "Field":
        """Resample the field to a new cell geometry.

        The physical extent of the field is preserved. Consequently, changing
        the number of cells changes the spacing accordingly.

        Args:
            shape (Tuple[int, int, int]): Target shape in NumPy axis order,
                i.e. ``(z, y, x)``.
            method (str, optional): Resampling method. Defaults to ``"nearest"``.
                ``"nearest"`` samples the nearest source cell, preserving the
                input dtype and suiting categorical data such as grain or phase
                IDs. ``"linear"`` interpolates between source cell centres and
                returns floating-point data for continuous fields. ``"mean"``
                computes an overlap-weighted mean of source cells covered by
                each target cell, returning floating-point data for conservative
                coarsening of continuous cell data.

        Returns:
            Field: Resampled field with updated spacing and preserved time.

        Raises:
            ValueError: If the target shape or method is invalid, or if
                ``"mean"`` would refine an axis.
        """
        if self.ndim != 3:
            raise ValueError("resample() requires a three-dimensional Field")

        try:
            shape = tuple(int(n) for n in shape)
        except (TypeError, ValueError) as exc:
            raise ValueError("shape must contain three positive integers") from exc

        if len(shape) != 3 or any(n <= 0 for n in shape):
            raise ValueError("shape must contain three positive integers")

        method = method.lower()
        if method not in {"nearest", "linear", "mean"}:
            raise ValueError("method must be 'nearest', 'linear', or 'mean'")

        old_shape = tuple(self.shape)

        if method == "mean" and any(new > old for old, new in zip(old_shape, shape)):
            raise ValueError(
                "'mean' can only be used for coarsening; "
                "use 'nearest' or 'linear' when refining a field"
            )

        data = np.asarray(self)

        if shape == old_shape:
            result = data.copy()

        elif method == "nearest":
            result = data

            for axis, (old_n, new_n) in enumerate(zip(old_shape, shape)):
                if old_n == new_n:
                    continue

                # Map target cell centres into source-cell coordinates.
                coords = (
                    np.arange(new_n, dtype=np.float64) + 0.5
                ) * old_n / new_n - 0.5

                indices = np.floor(coords + 0.5).astype(np.intp)
                indices = np.clip(indices, 0, old_n - 1)

                result = np.take(result, indices, axis=axis)

        elif method == "linear":
            result = data.astype(np.result_type(data.dtype, np.float64), copy=False)

            for axis, (old_n, new_n) in enumerate(zip(old_shape, shape)):
                if old_n == new_n:
                    continue

                # A singleton axis contains no information from which to
                # interpolate. Repeating the sole value is the natural result.
                if old_n == 1:
                    result = np.repeat(result, new_n, axis=axis)
                    continue

                # Coordinates of target cell centres expressed in source
                # cell-centre index coordinates.
                coords = (
                    np.arange(new_n, dtype=np.float64) + 0.5
                ) * old_n / new_n - 0.5

                # Constant extrapolation at the physical boundaries.
                coords = np.clip(coords, 0.0, old_n - 1.0)

                lower = np.floor(coords).astype(np.intp)
                upper = np.minimum(lower + 1, old_n - 1)
                weight = coords - lower

                lower_values = np.take(result, lower, axis=axis)
                upper_values = np.take(result, upper, axis=axis)

                weight_shape = [1] * result.ndim
                weight_shape[axis] = new_n
                weight = weight.reshape(weight_shape)

                result = lower_values * (1.0 - weight) + upper_values * weight

        else:  # method == "mean"
            result = data.astype(np.result_type(data.dtype, np.float64), copy=False)

            for axis, (old_n, new_n) in enumerate(zip(old_shape, shape)):
                if old_n == new_n:
                    continue

                # Source cells occupy intervals [i, i + 1]. Target cells cover
                # the same physical interval [0, old_n], divided into new_n
                # equal cells.
                target_edges = np.linspace(0.0, float(old_n), new_n + 1)

                source_left = np.arange(old_n, dtype=np.float64)
                source_right = source_left + 1.0

                target_left = target_edges[:-1, None]
                target_right = target_edges[1:, None]

                overlap = np.maximum(
                    0.0,
                    np.minimum(target_right, source_right)
                    - np.maximum(target_left, source_left),
                )

                # Divide by target-cell width so every row sums to one.
                weights = overlap / (float(old_n) / new_n)

                # Move the current spatial axis to the front, apply the
                # one-dimensional overlap operator, then restore the axis.
                moved = np.moveaxis(result, axis, 0)
                moved_shape = moved.shape

                result = weights @ moved.reshape(old_n, -1)
                result = result.reshape((new_n,) + moved_shape[1:])
                result = np.moveaxis(result, 0, axis)

        spacing = None
        if self.spacing is not None:
            spacing = tuple(
                old_spacing * old_n / new_n
                for old_spacing, old_n, new_n in zip(self.spacing, old_shape, shape)
            )

        return Field(result, time=self.time, spacing=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",
        data_mode: DataMode = "cell",
        length_unit: str = "cm",
        *,
        point_data: bool | None = None,
    ) -> "vtkImageData":
        """Convert the field to a VTK ImageData object.

        Args:
            name:
                Name of the VTK data array.
            data_mode:
                How MICRESS cell values are represented:
                - ``"cell"``: Store the original values as VTK CellData.
                - ``"point_average"``: Convert the cell values to PointData by averaging adjacent cells.
                - ``"point_direct"``: Store each original field value directly as a VTK point value, without averaging.
            length_unit:
                The unit of length for the VTK image spacing. Supported units
                are "um" (micrometers), "mm" (millimeters), "cm" (centimeters),
                and "m" (meters).
            point_data:
                Deprecated. Use ``data_mode`` instead.

        Returns:
            A ``vtkImageData`` object containing the field values.
        """

        if point_data is not None:
            warnings.warn(
                "'point_data' is deprecated and will be removed in MicPy 0.5.0. "
                "Use data_mode='cell' or data_mode='point_average' instead.",
                DeprecationWarning,
                stacklevel=2,
            )

            data_mode = "point_average" if point_data else "cell"

        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)")

        if self.spacing is None:
            raise ValueError("spacing is required for VTK conversion")

        if data_mode not in {"cell", "point_average", "point_direct"}:
            raise ValueError(
                "data_mode must be 'cell', 'point_average', or 'point_direct'"
            )

        nz, ny, nx = cells.shape
        dz, dy, dx = self.spacing
        length_scale = _length_unit_scale_from_cm(length_unit)
        dz, dy, dx = (dz * length_scale, dy * length_scale, dx * length_scale)

        def cell_to_point_average(c: np.ndarray) -> np.ndarray:
            nz_, ny_, nx_ = c.shape

            dtype = np.result_type(c.dtype, np.float32)

            acc = np.zeros(
                (nz_ + 1, ny_ + 1, nx_ + 1),
                dtype=dtype,
            )
            weights = np.zeros(
                (nz_ + 1, ny_ + 1, nx_ + 1),
                dtype=np.float32,
            )

            for z_offset in (0, 1):
                for y_offset in (0, 1):
                    for x_offset in (0, 1):
                        z_slice = slice(z_offset, z_offset + nz_)
                        y_slice = slice(y_offset, y_offset + ny_)
                        x_slice = slice(x_offset, x_offset + nx_)

                        acc[z_slice, y_slice, x_slice] += c
                        weights[z_slice, y_slice, x_slice] += 1.0

            return acc / weights

        if data_mode == "cell":
            data = cells
            dimensions = (nx + 1, ny + 1, nz + 1)
            vtk_attributes = "cell"

        elif data_mode == "point_average":
            data = cell_to_point_average(cells)
            dimensions = (nx + 1, ny + 1, nz + 1)
            vtk_attributes = "point"

        else:
            data = cells
            dimensions = (nx, ny, nz)
            vtk_attributes = "point"

        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(*dimensions)

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

        if vtk_attributes == "point":
            image.GetPointData().SetScalars(vtk_arr)
        else:
            image.GetCellData().SetScalars(vtk_arr)

        # Keep the NumPy allocation alive because numpy_to_vtk used deep=False.
        image._numpy_pin = data_c

        return image

    def save_vti(
        self,
        filename: str,
        name: str = "values",
        data_mode: DataMode = "cell",
        file_format: VTKFileFormat = "xml",
        *,
        point_data: bool | None = None,
    ) -> str:
        """Save the field as VTK ImageData.

        Args:
            filename:
                Filename of the VTK file.
            name:
                Name of the VTK data array.
            data_mode:
                How MICRESS cell values are represented:
                - ``"cell"``: Store the original values as VTK CellData.
                - ``"point_average"``: Convert the cell values to PointData by
                averaging adjacent cells.
                - ``"point_direct"``: Store each original field value directly
                as a VTK point value, without averaging.
            file_format:
                VTK file format:
                - ``"xml"``: VTK XML ImageData format (``.vti``).
                - ``"legacy"``: Legacy VTK file format (typically ``.vtk``).
            point_data:
                Deprecated. Use ``data_mode`` instead.

        Returns:
            The filename of the saved VTK file.
        """

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

        filename = str(Path(filename))

        if file_format == "xml":
            writer = vtkXMLImageDataWriter()
            writer.SetDataModeToAppended()
            writer.EncodeAppendedDataOff()

        elif file_format == "legacy":
            writer = vtkDataSetWriter()
            writer.SetFileTypeToBinary()

        else:
            raise ValueError("file_format must be 'xml' or 'legacy'")

        writer.SetFileName(filename)
        writer.SetInputData(image)

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

        return filename