Converts a Fluent mesh in ASCII format into a regular mesh in VTK format

fluent_ascii_to_regular_mesh_vtk.py

# script provides a function to convert ASCII column data to a regular mesh and writes it to a VTK file.
# example usage at the bottom of the script

import os
import typing
import warnings
from collections.abc import Iterable
from numbers import Integral, Real

import numpy as np
from scipy import spatial


"""
This python script contains code snippets from OpenMC, below is the OpenMC license
functions check_length, check_type and class RegularMesh are all adapted from OpenMC
The convert_ascii_to_regular_mesh is not from OpenMC

Copyright (c) 2011-2024 Massachusetts Institute of Technology, UChicago Argonne
LLC, and OpenMC contributors

Permission is hereby granted, free of charge, to any person obtaining a copy of
this software and associated documentation files (the "Software"), to deal in
the Software without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Software, and to permit persons to whom the Software is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
"""


def check_length(name, value, length_min, length_max=None):
    """Ensure that a sized object has length within a given range.

    Parameters
    ----------
    name : str
        Description of value being checked
    value : collections.Sized
        Object to check length of
    length_min : int
        Minimum length of object
    length_max : int or None, optional
        Maximum length of object. If None, it is assumed object must be of
        length length_min.

    """

    if length_max is None:
        if len(value) < length_min:
            msg = (
                f'Unable to set "{name}" to "{value}" since it must be at '
                f'least of length "{length_min}"'
            )
            raise ValueError(msg)
    elif not length_min <= len(value) <= length_max:
        if length_min == length_max:
            msg = (
                f'Unable to set "{name}" to "{value}" since it must be of '
                f'length "{length_min}"'
            )
        else:
            msg = (
                f'Unable to set "{name}" to "{value}" since it must have '
                f'length between "{length_min}" and "{length_max}"'
            )
        raise ValueError(msg)


def check_type(name, value, expected_type, expected_iter_type=None, *, none_ok=False):
    """Ensure that an object is of an expected type. Optionally, if the object is
    iterable, check that each element is of a particular type.

    Parameters
    ----------
    name : str
        Description of value being checked
    value : object
        Object to check type of
    expected_type : type or Iterable of type
        type to check object against
    expected_iter_type : type or Iterable of type or None, optional
        Expected type of each element in value, assuming it is iterable. If
        None, no check will be performed.
    none_ok : bool, optional
        Whether None is allowed as a value

    """
    if none_ok and value is None:
        return

    if not isinstance(value, expected_type):
        if isinstance(expected_type, Iterable):
            msg = (
                'Unable to set "{}" to "{}" which is not one of the '
                'following types: "{}"'.format(
                    name, value, ", ".join([t.__name__ for t in expected_type])
                )
            )
        else:
            msg = (
                f'Unable to set "{name}" to "{value}" which is not of type "'
                f'{expected_type.__name__}"'
            )
        raise TypeError(msg)

    if expected_iter_type:
        if isinstance(value, np.ndarray):
            if not issubclass(value.dtype.type, expected_iter_type):
                msg = (
                    f'Unable to set "{name}" to "{value}" since each item '
                    f'must be of type "{expected_iter_type.__name__}"'
                )
                raise TypeError(msg)
            else:
                return

        for item in value:
            if not isinstance(item, expected_iter_type):
                if isinstance(expected_iter_type, Iterable):
                    msg = (
                        'Unable to set "{}" to "{}" since each item must be '
                        'one of the following types: "{}"'.format(
                            name,
                            value,
                            ", ".join([t.__name__ for t in expected_iter_type]),
                        )
                    )
                else:
                    msg = (
                        f'Unable to set "{name}" to "{value}" since each '
                        f'item must be of type "{expected_iter_type.__name__}"'
                    )
                raise TypeError(msg)


class RegularMesh:
    """A regular Cartesian mesh in one, two, or three dimensions

    Parameters
    ----------
    mesh_id : int
        Unique identifier for the mesh
    name : str
        Name of the mesh

    Attributes
    ----------
    dimension : Iterable of int
        The number of mesh cells in each direction (x, y, z).
    n_dimension : int
        Number of mesh dimensions.
    lower_left : Iterable of float
        The lower-left corner of the structured mesh. If only two coordinate
        are given, it is assumed that the mesh is an x-y mesh.
    upper_right : Iterable of float
        The upper-right corner of the structured mesh. If only two coordinate
        are given, it is assumed that the mesh is an x-y mesh.
    bounding_box : openmc.BoundingBox
        Axis-aligned bounding box of the mesh as defined by the upper-right and
        lower-left coordinates.
    width : Iterable of float
        The width of mesh cells in each direction.
    indices : Iterable of tuple
        An iterable of mesh indices for each mesh element, e.g. [(1, 1, 1),
        (2, 1, 1), ...]

    """

    def __init__(self):

        self._dimension = None
        self._lower_left = None
        self._upper_right = None
        self._width = None

    @property
    def vertices(self):
        """Return coordinates of mesh vertices in Cartesian coordinates. Also
           see :meth:`CylindricalMesh.vertices_cylindrical` and
           :meth:`SphericalMesh.vertices_spherical` for coordinates in other coordinate
           systems.

        Returns
        -------
        vertices : numpy.ndarray
            Returns a numpy.ndarray representing the coordinates of the mesh
            vertices with a shape equal to (dim1 + 1, ..., dimn + 1, ndim). X, Y, Z values
            can be unpacked with xx, yy, zz = np.rollaxis(mesh.vertices, -1).

        """
        return self._generate_vertices(*self._grids)

    @property
    def num_mesh_cells(self):
        return np.prod(self.dimension)

    @staticmethod
    def _generate_vertices(i_grid, j_grid, k_grid):
        """Returns an array with shape (i_grid.size, j_grid.size, k_grid.size, 3)
        containing the corner vertices of mesh elements.
        """
        return np.stack(np.meshgrid(i_grid, j_grid, k_grid, indexing="ij"), axis=-1)

    @property
    def centroids(self):
        """Return coordinates of mesh element centroids.

        Returns
        -------
        centroids : numpy.ndarray
            Returns a numpy.ndarray representing the mesh element centroid
            coordinates with a shape equal to (dim1, ..., dimn, ndim). X,
            Y, Z values can be unpacked with xx, yy, zz =
            np.rollaxis(mesh.centroids, -1).
        """
        ndim = self.n_dimension
        # this line ensures that the vertices aren't adjusted by the origin or
        # converted to the Cartesian system for cylindrical and spherical meshes
        vertices = self.vertices
        s0 = (slice(0, -1),) * ndim + (slice(None),)
        s1 = (slice(1, None),) * ndim + (slice(None),)
        return (vertices[s0] + vertices[s1]) / 2

    @property
    def dimension(self):
        return tuple(self._dimension)

    @dimension.setter
    def dimension(self, dimension: typing.Iterable[int]):
        check_type("mesh dimension", dimension, Iterable, Integral)
        check_length("mesh dimension", dimension, 1, 3)
        self._dimension = dimension

    @property
    def n_dimension(self):
        if self._dimension is not None:
            return len(self._dimension)
        else:
            return None

    @property
    def lower_left(self):
        return self._lower_left

    @lower_left.setter
    def lower_left(self, lower_left: typing.Iterable[Real]):
        check_type("mesh lower_left", lower_left, Iterable, Real)
        check_length("mesh lower_left", lower_left, 1, 3)
        self._lower_left = lower_left

        if self.upper_right is not None and any(
            np.isclose(self.upper_right, lower_left)
        ):
            raise ValueError("Mesh cannot have zero thickness in any dimension")

    @property
    def upper_right(self):
        if self._upper_right is not None:
            return self._upper_right
        elif self._width is not None:
            if self._lower_left is not None and self._dimension is not None:
                ls = self._lower_left
                ws = self._width
                dims = self._dimension
                return [l + w * d for l, w, d in zip(ls, ws, dims)]

    @upper_right.setter
    def upper_right(self, upper_right: typing.Iterable[Real]):
        check_type("mesh upper_right", upper_right, Iterable, Real)
        check_length("mesh upper_right", upper_right, 1, 3)
        self._upper_right = upper_right

        if self._width is not None:
            self._width = None
            warnings.warn("Unsetting width attribute.")

        if self.lower_left is not None and any(
            np.isclose(self.lower_left, upper_right)
        ):
            raise ValueError("Mesh cannot have zero thickness in any dimension")

    @property
    def width(self):
        if self._width is not None:
            return self._width
        elif self._upper_right is not None:
            if self._lower_left is not None and self._dimension is not None:
                us = self._upper_right
                ls = self._lower_left
                dims = self._dimension
                return [(u - l) / d for u, l, d in zip(us, ls, dims)]

    @width.setter
    def width(self, width: typing.Iterable[Real]):
        check_type("mesh width", width, Iterable, Real)
        check_length("mesh width", width, 1, 3)
        self._width = width

        if self._upper_right is not None:
            self._upper_right = None
            warnings.warn("Unsetting upper_right attribute.")

    @property
    def volumes(self):
        """Return Volumes for every mesh cell

        Returns
        -------
        volumes : numpy.ndarray
            Volumes

        """
        self._volume_dim_check()
        return np.full(self.dimension, np.prod(self.width))

    @property
    def total_volume(self):
        return np.prod(self.dimension) * np.prod(self.width)

    @property
    def indices(self):
        ndim = len(self._dimension)
        if ndim == 3:
            nx, ny, nz = self.dimension
            return (
                (x, y, z)
                for z in range(1, nz + 1)
                for y in range(1, ny + 1)
                for x in range(1, nx + 1)
            )
        elif ndim == 2:
            nx, ny = self.dimension
            return ((x, y) for y in range(1, ny + 1) for x in range(1, nx + 1))
        else:
            (nx,) = self.dimension
            return ((x,) for x in range(1, nx + 1))

    @property
    def _grids(self):
        ndim = len(self._dimension)
        if ndim == 3:
            x0, y0, z0 = self.lower_left
            x1, y1, z1 = self.upper_right
            nx, ny, nz = self.dimension
            xarr = np.linspace(x0, x1, nx + 1)
            yarr = np.linspace(y0, y1, ny + 1)
            zarr = np.linspace(z0, z1, nz + 1)
            return (xarr, yarr, zarr)
        elif ndim == 2:
            x0, y0 = self.lower_left
            x1, y1 = self.upper_right
            nx, ny = self.dimension
            xarr = np.linspace(x0, x1, nx + 1)
            yarr = np.linspace(y0, y1, ny + 1)
            return (xarr, yarr)
        else:
            (nx,) = self.dimension
            (x0,) = self.lower_left
            (x1,) = self.upper_right
            return (np.linspace(x0, x1, nx + 1),)

    def write_data_to_vtk(
        self,
        filename: typing.Union[str, os.PathLike],
        datasets: typing.Optional[dict] = None,
        volume_normalization: bool = True,
    ):
        """Creates a VTK object of the mesh

        Parameters
        ----------
        filename : str
            Name of the VTK file to write.
        datasets : dict
            Dictionary whose keys are the data labels
            and values are the data sets.
        volume_normalization : bool, optional
            Whether or not to normalize the data by
            the volume of the mesh elements.

        Raises
        ------
        ValueError
            When the size of a dataset doesn't match the number of mesh cells

        Returns
        -------
        vtk.StructuredGrid or vtk.UnstructuredGrid
            a VTK grid object representing the mesh
        """
        import vtk
        from vtk.util import numpy_support as nps

        # check that the data sets are appropriately sized
        if datasets is not None:
            self._check_vtk_datasets(datasets)

        # write linear elements using a structured grid
        vtk_grid = self._create_vtk_structured_grid()
        writer = vtk.vtkStructuredGridWriter()

        if datasets is not None:
            # maintain a list of the datasets as added
            # to the VTK arrays to ensure they persist
            # in memory until the file is written
            datasets_out = []
            for label, dataset in datasets.items():
                dataset = np.asarray(dataset).flatten()
                datasets_out.append(dataset)

                if volume_normalization:
                    dataset /= self.volumes.T.flatten()

                dataset_array = vtk.vtkDoubleArray()
                dataset_array.SetName(label)
                dataset_array.SetArray(nps.numpy_to_vtk(dataset), dataset.size, True)
                vtk_grid.GetCellData().AddArray(dataset_array)

        writer.SetFileName(str(filename))
        writer.SetInputData(vtk_grid)
        writer.Write()

        return vtk_grid

    def _check_vtk_datasets(self, datasets: dict):
        """Perform some basic checks that the datasets are valid for this mesh

        Parameters
        ----------
        datasets : dict
            Dictionary whose keys are the data labels
            and values are the data sets.

        """
        for label, dataset in datasets.items():
            errmsg = (
                f"The size of the dataset '{label}' ({dataset.size}) should be"
                f" equal to the number of mesh cells ({self.num_mesh_cells})"
            )
            if isinstance(dataset, np.ndarray):
                if not dataset.size == self.num_mesh_cells:
                    raise ValueError(errmsg)
            else:
                if len(dataset) == self.num_mesh_cells:
                    raise ValueError(errmsg)
            check_type("data label", label, str)

    def _create_vtk_structured_grid(self):
        """Create a structured grid

        Returns
        -------
        vtk.vtkStructuredGrid
            a VTK structured grid object representing the mesh
        """
        import vtk
        from vtk.util import numpy_support as nps

        vtkPts = vtk.vtkPoints()
        vtkPts.SetData(
            nps.numpy_to_vtk(np.swapaxes(self.vertices, 0, 2).reshape(-1, 3), deep=True)
        )
        vtk_grid = vtk.vtkStructuredGrid()
        vtk_grid.SetPoints(vtkPts)
        vtk_grid.SetDimensions(*[dim + 1 for dim in self.dimension])

        return vtk_grid

    def _volume_dim_check(self):
        if self.n_dimension != 3 or any([d == 0 for d in self.dimension]):
            raise RuntimeError(
                f"Mesh {self.id} is not 3D. " "Volumes cannot be provided."
            )


def convert_ascii_to_regular_mesh(
    filepath: str,
    requested_scalars: typing.Iterable[str],
    regular_mesh_dimensions: typing.Tuple[int, int, int],
    output_filename: str = "rmesh.vtk",
):
    """
    Converts ASCII column data to a regular mesh and writes it to a VTK file.

    Args:
        filepath (str): The path to the ASCII csv file, separated with commas.
        requested_scalars (Iterable[str]): The names of the scalars you want to extract, should match the headers of the columns from the ASCII csv file.
        regular_mesh_dimensions (Tuple[int, int, int]): The number of regular mesh elements in each direction (x,y,z).
        output_filename (str, optional): The name of the output VTK file. Defaults to 'rmesh.vtk'.
    """

    # Load the file to get the column names (headers)
    with open(filepath, "r") as f:
        column_names = f.readline().strip().split(",")
    column_names = [name.strip() for name in column_names]

    # Load the content of the file
    content = np.genfromtxt(
        fname=filepath,
        delimiter=",",
        skip_header=1,
    )

    # Extract the file contents into separate arrays
    x_coords = content[:, column_names.index("x-coordinate")]
    y_coords = content[:, column_names.index("y-coordinate")]
    z_coords = content[:, column_names.index("z-coordinate")]
    cell_volumes = content[:, column_names.index("cell-volume")]

    total_atoms = {}
    scalars = {}
    for requested_scalar in requested_scalars:
        if requested_scalar not in column_names:
            raise ValueError(
                f"Requested scalar '{requested_scalar}' not found in the file. Available scalars are: {column_names}"
            )
        scalar = content[:, column_names.index(requested_scalar)]
        scalars[requested_scalar] = scalar
        atoms = cell_volumes * scalar
        total_atoms[requested_scalar] = np.sum(atoms)

    rmesh = RegularMesh()
    rmesh.lower_left = (min(x_coords), min(y_coords), min(z_coords))
    rmesh.upper_right = (max(x_coords), max(y_coords), max(z_coords))
    rmesh.dimension = regular_mesh_dimensions

    coords = np.column_stack((x_coords, y_coords, z_coords))
    flat_centroids = rmesh.centroids.reshape(rmesh.num_mesh_cells, 3, order="F")

    kdtree = spatial.KDTree(coords)

    datasets = {}
    for key, value in scalars.items():
        rmesh_values = []
        for point in flat_centroids:

            distance, index = kdtree.query(point)

            if distance > 0.5 * max(rmesh.width):
                # TODO consider user defined distance / tolerance
                # Point is not in the unstructured mesh and does not even have partial coverage with a regular mesh cell
                rmesh_values.append(0.0)
            else:
                # Assign the regular mesh cell the value of the closest unstructured mesh cell (nearest neighbor method)
                rmesh_values.append(value[index])

        array_rmesh_values = np.array(rmesh_values)
        normalisation_coeff = total_atoms[key] / np.sum(array_rmesh_values)
        datasets[key] = array_rmesh_values * normalisation_coeff

    rmesh.write_data_to_vtk(filename=output_filename, datasets=datasets)


# example usage, reads in a csv file produced by Fluent with this structure:
# cellnumber,    x-coordinate,    y-coordinate,    z-coordinate,    uds-1-scalar,    uds-2-scalar,    uds-3-scalar, cell-volume
#          1, 9.471632417E-01, 2.780554034E-03,-6.899168244E-02, 3.497009281E+11, 6.194563914E+07, 3.422243003E+10, 1.345343532E-11
#          2, 9.471632956E-01, 2.100672898E-03,-6.868461691E-02, 3.498690659E+11, 6.196891909E+07, 3.424317962E+10, 1.342342423E-11
#          3, 9.471630839E-01, 3.418335265E-03,-6.938289824E-02, 3.495451807E+11, 6.192396353E+07, 3.420349805E+10, 1.546456645E-11
#          4, 9.471634845E-01, 2.751234181E-03,-6.904688086E-02, 3.497131805E+11, 6.194895704E+07, 3.422291764E+10, 1.564564565E-11
#          5, 9.466798659E-01, 3.119070606E-03,-6.918231709E-02, 3.496175822E+11, 6.193411426E+07, 3.421219831E+10, 2.456454664E-11
#          ...
# converts the data from uds-1-scalar uds-2-scalar to a regular mesh with 100x10x10 elements and writes it to rmesh.vtk
convert_ascii_to_regular_mesh(
    filepath="InnerHead_domain_1kgs_2Mcells_CellVolume_centroid_N16_N17_O19",
    requested_scalars=["uds-1-scalar", "uds-2-scalar"],
    regular_mesh_dimensions=(100, 10, 10),
    output_filename="rmesh.vtk",
)