Create DAGMC geometry from CAD
"""Making the stellarmesh logo."""
import logging

import build123d as bd
import stellarmesh as sm

logging.getLogger("stellarmesh").setLevel(logging.INFO)

b1 = bd.Solid.make_box(10, 10, 10)
b2 = b1.transformed(offset=(0, 10, 0))
b3 = b1.transformed(offset=(0, 5, 10))

cmp_initial = bd.Compound.make_compound([b1, b2, b3])
solids = cmp_initial.solids()
geom = sm.Geometry(solids, material_names=[""] * len(solids))
geom_imprinted = geom.imprint()

examples/imprinted-boxes.py
Thea-Energy/stellarmesh
"""Test of a simple layered torus."""
import logging

import build123d as bd
import stellarmesh as sm

logging.getLogger("stellarmesh").setLevel(logging.INFO)

solids = [bd.Solid.make_torus(1000, 100)]
for _ in range(3):
    solids.append(solids[-1].faces()[0].thicken(100))
solids = solids[1:]

geometry = sm.Geometry(solids[::-1], material_names=["a", "a", "c"])
mesh = sm.Mesh.mesh_geometry(geometry, min_mesh_size=50, max_mesh_size=50)
mesh.write("test.msh")
mesh.render("doc/torus-mesh-reversed.png", rotation_xyz=(90, 0, -90), normals=15)

h5m = sm.MOABModel.make_from_mesh(mesh)
h5m.write("dagmc.h5m")
h5m.write("dagmc.vtk")

examples/layered-torus.py
Thea-Energy/stellarmesh
"""Test of a simple layered torus. Outputs refined mesh to refined.msh."""
import logging

import build123d as bd
import numpy as np
import stellarmesh as sm

logging.getLogger("stellarmesh").setLevel(logging.INFO)

with bd.BuildPart() as my_part:
    with bd.BuildLine() as my_line:
        bd.Bezier(
            *[
                (i, 0, 2 * x)
                for i, x in enumerate(np.sin(np.arange(2 * np.pi, step=0.5)))
            ]
        )
    with bd.BuildSketch() as my_sketch:
        bd.Rectangle(1, 5)
    bd.sweep()

geom = sm.Geometry(my_part.solids(), material_names=["hi"])
mesh = sm.Mesh.mesh_geometry(geom, min_mesh_size=0.5, max_mesh_size=0.5)
refined_mesh = mesh.refine(const_mesh_size=1, hausdorff_value=100)
refined_mesh.write("refined.msh")

examples/mesh-refinement.py
Thea-Energy/stellarmesh
"""Making the stellarmesh logo."""
import logging

import build123d as bd
import stellarmesh as sm

logging.getLogger("stellarmesh").setLevel(logging.INFO)

cmp = bd.Compound.make_text("Stellarmesh", 14, font="Arial Black")
solids = [f.thicken(10) for f in cmp.faces()]

geometry = sm.Geometry(solids, [""] * len(solids))
mesh = sm.Mesh.mesh_geometry(geometry, min_mesh_size=1, max_mesh_size=2)
mesh.render("doc/logo.png", rotation_xyz=(0, -2, 0), clipping=False)

examples/stellarmesh-logo.py
Thea-Energy/stellarmesh
"""GMSH wrapper and DAGMC geometry creator."""
from .geometry import Geometry  # noqa: F401
from .mesh import Mesh  # noqa: F401
from .moab import MOABModel  # noqa: F401

src/stellarmesh/__init__.py
Thea-Energy/stellarmesh
"""Stellarmesh geometry.

name: geometry.py
author: Alex Koen

desc: Geometry class represents a CAD geometry to be meshed.
"""
import logging
from typing import Sequence, Union

from OCP.BOPAlgo import BOPAlgo_MakeConnected
from OCP.BRep import BRep_Builder
from OCP.BRepTools import BRepTools
from OCP.IFSelect import IFSelect_RetDone
from OCP.STEPControl import STEPControl_Reader
from OCP.TopAbs import TopAbs_ShapeEnum
from OCP.TopExp import TopExp_Explorer
from OCP.TopoDS import TopoDS, TopoDS_Shape, TopoDS_Solid

logger = logging.getLogger(__name__)


class Geometry:
    """Geometry, representing an ordered list of solids, to be meshed."""

    solids: Sequence[TopoDS_Solid]
    material_names: Sequence[str]

    def __init__(
        self,
        solids: Sequence[Union["bd.Solid", "cq.Solid", TopoDS_Solid]],  # noqa: F821
        material_names: Sequence[str],
    ):
        """Construct geometry from solids.

        Args:
            solids: List of solids, where each solid is a build123d Solid, cadquery
            Solid, or OCP TopoDS_Solid.
            material_names: List of materials. Must match length of solids.
        """
        logger.info(f"Importing {len(solids)} solids to geometry")
        if len(material_names) != len(solids):
            raise ValueError(
                f"Number of material names ({len(material_names)}) must match length of"
                + " solids ({len(solids)})."
            )

        self.solids = []
        for i, s in enumerate(solids):
            if isinstance(s, TopoDS_Solid):
                self.solids.append(s)
            elif hasattr(s, "wrapped"):
                self.solids.append(s.wrapped)
            else:
                raise ValueError(
                    f"Solid {i} is of type {type(s).__name__}, not a cadquery Solid, "
                    + "build123d Solid, or TopoDS_Solid"
                )

        self.material_names = material_names

    @staticmethod
    def _solids_from_shape(shape: TopoDS_Shape) -> list[TopoDS_Solid]:
        """Return all the solids in this shape."""
        solids = []
        if shape.ShapeType() == TopAbs_ShapeEnum.TopAbs_SOLID:
            solids.append(TopoDS.Solid_s(shape))
        if shape.ShapeType() == TopAbs_ShapeEnum.TopAbs_COMPOUND:
            explorer = TopExp_Explorer(shape, TopAbs_ShapeEnum.TopAbs_SOLID)
            while explorer.More():
                assert explorer.Current().ShapeType() == TopAbs_ShapeEnum.TopAbs_SOLID
                solids.append(TopoDS.Solid_s(explorer.Current()))
                explorer.Next()
        return solids

    # TODO(akoen): import_step and import_brep are not DRY
    # https://github.com/Thea-Energy/stellarmesh/issues/2
    @classmethod
    def import_step(
        cls,
        filename: str,
        material_names: Sequence[str],
    ) -> "Geometry":
        """Import model from a step file.

        Args:
            filename: File path to import.
            material_names: Ordered list of material names matching solids in file.

        Returns:
            Model.
        """
        logger.info(f"Importing {filename}")

        reader = STEPControl_Reader()
        read_status = reader.ReadFile(filename)
        if read_status != IFSelect_RetDone:
            raise ValueError(f"STEP File {filename} could not be loaded")
        for i in range(reader.NbRootsForTransfer()):
            reader.TransferRoot(i + 1)

        solids = []
        for i in range(reader.NbShapes()):
            shape = reader.Shape(i + 1)
            solids.extend(cls._solids_from_shape(shape))

        return cls(solids, material_names)

    @classmethod
    def import_brep(
        cls,
        filename: str,
        material_names: Sequence[str],
    ) -> "Geometry":
        """Import model from a brep (cadquery, build123d native) file.

        Args:
            filename: File path to import.
            material_names: Ordered list of material names matching solids in file.

        Returns:
            Model.
        """
        logger.info(f"Importing {filename}")

        shape = TopoDS_Shape()
        builder = BRep_Builder()
        BRepTools.Read_s(shape, filename, builder)

        if shape.IsNull():
            raise ValueError(f"Could not import {filename}")
        solids = cls._solids_from_shape(shape)

        logger.info(f"Importing {len(solids)} from {filename}")
        return cls(solids, material_names)

    def imprint(self) -> "Geometry":
        """Imprint faces of current geometry.

        Returns:
            A new geometry with the imprinted and merged geometry.
        """
        bldr = BOPAlgo_MakeConnected()
        bldr.SetRunParallel(theFlag=True)
        bldr.SetUseOBB(theUseOBB=True)

        for solid in self.solids:
            bldr.AddArgument(solid)

        bldr.Perform()
        res = bldr.Shape()
        res_solids = self._solids_from_shape(res)

        if (l0 := len(res_solids)) != (l1 := len(self.solids)):
            raise RuntimeError(
                f"Length of imprinted solids {l0} != length of original solids {l1}"
            )

        return type(self)(res_solids, self.material_names)

src/stellarmesh/geometry.py
Thea-Energy/stellarmesh
"""Stellarmesh mesh.

name: mesh.py
author: Alex Koen

desc: Mesh class wraps Gmsh functionality for geometry meshing.
"""
import logging
import subprocess
import tempfile
from contextlib import contextmanager
from pathlib import Path
from typing import Optional

import gmsh

from .geometry import Geometry

logger = logging.getLogger(__name__)


class Mesh:
    """A Gmsh mesh.

    As gmsh allows for only a single process, this class provides a context manager to
    set the gmsh api to operate on this mesh.
    """

    _mesh_filename: str

    def __init__(self, mesh_filename: Optional[str] = None):
        """Initialize a mesh from a .msh file.

        Args:
            mesh_filename: Optional .msh filename. If not provided defaults to a
            temporary file. Defaults to None.
        """
        if not mesh_filename:
            with tempfile.NamedTemporaryFile(suffix=".msh", delete=False) as mesh_file:
                mesh_filename = mesh_file.name
        self._mesh_filename = mesh_filename

    def __enter__(self):
        """Enter mesh context, setting gmsh commands to operate on this mesh."""
        if not gmsh.is_initialized():
            gmsh.initialize()

        gmsh.option.set_number(
            "General.Terminal",
            1 if logger.getEffectiveLevel() <= logging.INFO else 0,
        )
        gmsh.open(self._mesh_filename)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Cleanup (finalize) gmsh."""
        gmsh.finalize()

    def _save_changes(self, *, save_all: bool = True):
        gmsh.option.set_number("Mesh.SaveAll", 1 if save_all else 0)
        gmsh.write(self._mesh_filename)

    def write(self, filename: str, *, save_all: bool = True):
        """Write mesh to a .msh file.

        Args:
            filename: Path to write file.
            save_all: Whether to save all entities (or just physical groups). See
            Gmsh documentation for Mesh.SaveAll. Defaults to True.
        """
        with self:
            gmsh.option.set_number("Mesh.SaveAll", 1 if save_all else 0)
            gmsh.write(filename)

    @classmethod
    def mesh_geometry(
        cls,
        geometry: Geometry,
        min_mesh_size: float = 50,
        max_mesh_size: float = 50,
    ):
        """Mesh solids with Gmsh.

        See Gmsh documentation on mesh sizes:
        https://gmsh.info/doc/texinfo/gmsh.html#Specifying-mesh-element-sizes

        Args:
            geometry: Geometry to be meshed.
            mesh_filename: Optional filename to store .msh file. Defaults to None.
            min_mesh_size: Min mesh element size. Defaults to 50.
            max_mesh_size: Max mesh element size. Defaults to 50.
        """
        logger.info(f"Meshing solids with mesh size {min_mesh_size}, {max_mesh_size}")

        with cls() as mesh:
            gmsh.model.add("stellarmesh_model")

            material_solid_map = {}
            for s, m in zip(geometry.solids, geometry.material_names):
                dim_tags = gmsh.model.occ.import_shapes_native_pointer(s._address())
                if dim_tags[0][0] != 3:
                    raise TypeError("Importing non-solid geometry.")

                solid_tag = dim_tags[0][1]
                if m not in material_solid_map:
                    material_solid_map[m] = [solid_tag]
                else:
                    material_solid_map[m].append(solid_tag)

            gmsh.model.occ.synchronize()

            for material, solid_tags in material_solid_map.items():
                gmsh.model.add_physical_group(3, solid_tags, name=f"mat:{material}")

            gmsh.option.set_number("Mesh.MeshSizeMin", min_mesh_size)
            gmsh.option.set_number("Mesh.MeshSizeMax", max_mesh_size)
            gmsh.model.mesh.generate(2)

            mesh._save_changes(save_all=True)
            return mesh

    def render(
        self,
        output_filename: Optional[str] = None,
        rotation_xyz: tuple[float, float, float] = (0, 0, 0),
        normals: int = 0,
        *,
        clipping: bool = True,
    ) -> str:
        """Render mesh as an image.

        Args:
            output_filename: Optional output filename. Defaults to None.
            rotation_xyz: Rotation in Euler angles. Defaults to (0, 0, 0).
            normals: Normal render size. Defaults to 0.
            clipping: Whether to enable mesh clipping. Defaults to True.

        Returns:
            Path to image file, either passed output_filename or a temporary file.
        """
        with self:
            gmsh.option.set_number("Mesh.SurfaceFaces", 1)
            gmsh.option.set_number("Mesh.Clip", 1 if clipping else 0)
            gmsh.option.set_number("Mesh.Normals", normals)
            gmsh.option.set_number("General.Trackball", 0)
            gmsh.option.set_number("General.RotationX", rotation_xyz[0])
            gmsh.option.set_number("General.RotationY", rotation_xyz[1])
            gmsh.option.set_number("General.RotationZ", rotation_xyz[2])
            if not output_filename:
                with tempfile.NamedTemporaryFile(
                    delete=False, mode="w", suffix=".png"
                ) as temp_file:
                    output_filename = temp_file.name

            try:
                gmsh.fltk.initialize()
                gmsh.write(output_filename)
            finally:
                gmsh.fltk.finalize()
            return output_filename

    @staticmethod
    def _check_is_initialized():
        if not gmsh.is_initialized():
            raise RuntimeError("Gmsh not initialized.")

    @contextmanager
    def _stash_physical_groups(self):
        self._check_is_initialized()
        physical_groups: dict[tuple[int, int], tuple[list[int], str]] = {}
        dim_tags = gmsh.model.get_physical_groups()
        for dim_tag in dim_tags:
            tags = gmsh.model.get_entities_for_physical_group(*dim_tag)
            name = gmsh.model.get_physical_name(*dim_tag)
            physical_groups[dim_tag] = (tags, name)
        gmsh.model.remove_physical_groups(dim_tags)

        try:
            yield
        except Exception as e:
            raise RuntimeError("Cannot unstash physical groups due to error.") from e
        else:
            if len(gmsh.model.get_physical_groups()) > 0:
                raise RuntimeError(
                    "Not overwriting existing physical groups on stash restore."
                )
            for physical_group in physical_groups.items():
                dim, tag = physical_group[0]
                tags, name = physical_group[1]
                gmsh.model.add_physical_group(dim, tags, tag, name)

    def refine(  # noqa: PLR0913
        self,
        *,
        min_mesh_size: Optional[float] = None,
        max_mesh_size: Optional[float] = None,
        const_mesh_size: Optional[float] = None,
        hausdorff_value: float = 0.01,
        gradation_value: float = 1.3,
        optim: bool = False,
    ) -> "Mesh":
        """Refine mesh using mmgs.

        See mmgs documentation:
        https://www.mmgtools.org/mmg-remesher-try-mmg/mmg-remesher-tutorials/mmg-remesher-mmg2d/mesh-adaptation-to-a-solution
        for more info.

        Pay particular attention to the hausdorff value, which overrides most of the
        other options and is typically set too low. Set to a large value, on the order
        of the size of your bounding box, to disable completely.

        Args:
            min_mesh_size: -hmin: Min size of output mesh elements. Defaults to None.
            max_mesh_size: -hmax: Max size of output mesh elements. Defaults to None.
            const_mesh_size: -hsize: Constant size map
            hausdorff_value: -hausd: Hausdorff value. Defaults to 0.01, which is
            suitable for a circle of radius 1. Set to a large value to disable effect.
            gradation_value: -hgrad Gradation value. Defaults to 1.3.
            optim: -optim Do not change elements sizes. Defaults to False.

        Raises:
            RuntimeError: If refinement fails.

        Returns:
            New refined mesh with filename <original-filename>.refined.msh.
        """
        with (
            self
        ), self._stash_physical_groups(), tempfile.TemporaryDirectory() as tmpdir:
            filename = f"{tmpdir}/model.mesh"
            gmsh.write(filename)

            refined_filename = str(Path(filename).with_suffix(".o.mesh").resolve())
            command = ["mmgs"]

            params = {
                "-hmin": min_mesh_size,
                "-hmax": max_mesh_size,
                "-hsiz": const_mesh_size,
                "-hausd": hausdorff_value,
                "-hgrad": gradation_value,
            }

            for param in params.items():
                if param[1]:
                    command.extend([param[0], str(param[1])])
            if optim:
                command.append("-optim")

            command.extend(
                [
                    "-in",
                    filename,
                    "-out",
                    refined_filename,
                ]
            )

            # TODO(akoen): log subprocess realtime
            # https://github.com/Thea-Energy/stellarmesh/issues/13
            try:
                logger.info(
                    f"Refining mesh {filename} with mmgs, output to {refined_filename}."
                )
                output = subprocess.run(
                    command,
                    text=True,
                    check=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                )

                if output.stdout:
                    logger.info(output.stdout)

            except subprocess.CalledProcessError as e:
                logger.exception(
                    "Command failed with error code %d\nSTDERR:%s",
                    e.returncode,
                    e.stdout,
                )
                raise RuntimeError("Command failed to run. See output above.") from e

            gmsh.model.mesh.clear()
            gmsh.merge(refined_filename)

            new_filename = str(
                Path(self._mesh_filename).with_suffix(".refined.msh").resolve()
            )
            gmsh.option.set_number("Mesh.SaveAll", 1)
            gmsh.write(new_filename)
            return type(self)(new_filename)

src/stellarmesh/mesh.py
Thea-Energy/stellarmesh
"""Stellarmesh MOAB model.

name: moab.py
author: Alex Koen

desc: MOABModel class represents a MOAB model.
"""
import subprocess
import tempfile
from dataclasses import dataclass, field

import gmsh
import numpy as np
import pymoab.core
import pymoab.types

from .mesh import Mesh


class _MOABEntity:
    _core: pymoab.core.Core
    handle: np.uint64

    def __init__(self, core: pymoab.core.Core, handle: np.uint64):
        self._core = core
        self.handle = handle


class MOABSurface(_MOABEntity):
    """MOAB surface entity."""

    @property
    def adjacent_volumes(self) -> list["MOABVolume"]:
        """Get adjacent volumes.

        Returns:
            Adjacent volumes.
        """
        parent_entities = self._core.get_parent_meshsets(self.handle)
        return [MOABVolume(self._core, e) for e in parent_entities]


class MOABVolume(_MOABEntity):
    """MOAB volume entity."""

    @property
    def adjacent_surfaces(self) -> list["MOABSurface"]:
        """Get adjacent surfaces.

        Returns:
            Adjacent surfaces.
        """
        child_entities = self._core.get_child_meshsets(self.handle)
        return [MOABSurface(self._core, e) for e in child_entities]


@dataclass
class _Surface:
    """Internal class for surface sense handling."""

    handle: np.uint64
    forward_volume: np.uint64 = field(default=np.uint64(0))
    reverse_volume: np.uint64 = field(default=np.uint64(0))

    def sense_data(self) -> list[np.uint64]:
        """Get MOAB tag sense data.

        Returns:
            Sense data.
        """
        return [self.forward_volume, self.reverse_volume]


class MOABModel:
    """MOAB Model."""

    # h5m_filename: str
    _core: pymoab.core.Core

    def __init__(self, core: pymoab.core.Core):
        """Initialize model from a pymoab core object.

        Args:
            core: Pymoab core.
        """
        self._core = core

    @classmethod
    def read_file(cls, h5m_file: str) -> "MOABModel":
        """Initialize model from .h5m file.

        Args:
            h5m_file: File to load.

        Returns:
            Initialized model.
        """
        core = pymoab.core.Core()
        core.load_file(h5m_file)
        return cls(core)

    def write(self, filename: str):
        """Write MOAB model to .h5m, .vtk, or other file.

        Args:
            filename: Filename with format-appropriate extension.
        """
        self._core.write_file(filename)

    @staticmethod
    def make_watertight(
        input_filename: str,
        output_filename: str,
        binary_path: str = "make_watertight",
    ):
        """Make mesh watertight.

        Args:
            input_filename: Input .h5m filename.
            output_filename: Output watertight .h5m filename.
            binary_path: Path to make_watertight or default to find in path. Defaults to
            "make_watertight".
        """
        subprocess.run(
            [binary_path, input_filename, "-o", output_filename],
            check=True,
        )

    @staticmethod
    def _get_moab_tag_handles(core: pymoab.core.Core) -> dict[str, np.uint64]:
        tag_handles = {}

        sense_tag_name = "GEOM_SENSE_2"
        sense_tag_size = 2
        tag_handles["surf_sense"] = core.tag_get_handle(
            sense_tag_name,
            sense_tag_size,
            pymoab.types.MB_TYPE_HANDLE,
            pymoab.types.MB_TAG_SPARSE,
            create_if_missing=True,
        )

        tag_handles["category"] = core.tag_get_handle(
            pymoab.types.CATEGORY_TAG_NAME,
            pymoab.types.CATEGORY_TAG_SIZE,
            pymoab.types.MB_TYPE_OPAQUE,
            pymoab.types.MB_TAG_SPARSE,
            create_if_missing=True,
        )

        tag_handles["name"] = core.tag_get_handle(
            pymoab.types.NAME_TAG_NAME,
            pymoab.types.NAME_TAG_SIZE,
            pymoab.types.MB_TYPE_OPAQUE,
            pymoab.types.MB_TAG_SPARSE,
            create_if_missing=True,
        )

        geom_dimension_tag_size = 1
        tag_handles["geom_dimension"] = core.tag_get_handle(
            pymoab.types.GEOM_DIMENSION_TAG_NAME,
            geom_dimension_tag_size,
            pymoab.types.MB_TYPE_INTEGER,
            pymoab.types.MB_TAG_SPARSE,
            create_if_missing=True,
        )

        faceting_tol_tag_name = "FACETING_TOL"
        faceting_tol_tag_size = 1
        tag_handles["faceting_tol"] = core.tag_get_handle(
            faceting_tol_tag_name,
            faceting_tol_tag_size,
            pymoab.types.MB_TYPE_DOUBLE,
            pymoab.types.MB_TAG_SPARSE,
            create_if_missing=True,
        )

        # Default tag, does not need to be created
        tag_handles["global_id"] = core.tag_get_handle(pymoab.types.GLOBAL_ID_TAG_NAME)

        return tag_handles

    @classmethod
    def make_from_mesh(  # noqa: PLR0915
        cls,
        mesh: Mesh,
    ):
        """Compose DAGMC MOAB .h5m file from mesh.

        Args:
            mesh: Mesh from which to build DAGMC geometry.
            filename: Filename of the output .h5m file.
        """
        core = pymoab.core.Core()

        tag_handles = cls._get_moab_tag_handles(core)

        known_surfaces: dict[int, _Surface] = {}
        known_groups: dict[int, np.uint64] = {}

        with mesh:
            volume_dimtags = gmsh.model.get_entities(3)
            volume_tags = [v[1] for v in volume_dimtags]
            for i, volume_tag in enumerate(volume_tags):
                # Add volume set
                volume_set_handle = core.create_meshset()
                global_id = volume_set_handle
                core.tag_set_data(tag_handles["global_id"], global_id, i)
                core.tag_set_data(tag_handles["geom_dimension"], volume_set_handle, 3)
                core.tag_set_data(tag_handles["category"], volume_set_handle, "Volume")

                # Add volume to its physical group, which stores metadata incl. material
                # TODO(akoen): should this be a parent-child relationship?
                # https://github.com/Thea-Energy/neutronics-cad/issues/2
                vol_groups = gmsh.model.get_physical_groups_for_entity(3, volume_tag)
                if (num_groups := len(vol_groups)) != 1:
                    raise ValueError(
                        f"Volume with tag {volume_tag} and global_id {global_id} "
                        f"belongs to {num_groups} physical groups, should be 1"
                    )

                if (vol_group := vol_groups[0]) not in known_groups:
                    mat_name = gmsh.model.get_physical_name(3, vol_group)
                    group_set = core.create_meshset()
                    core.tag_set_data(tag_handles["category"], group_set, "Group")
                    core.tag_set_data(tag_handles["name"], group_set, f"{mat_name}")
                    core.tag_set_data(tag_handles["global_id"], group_set, vol_group)
                    known_groups[vol_group] = group_set
                else:
                    group_set = known_groups[vol_group]

                core.add_entity(group_set, volume_set_handle)

                # Add surfaces to MOAB core, respecting surface sense.
                # Logic: Gmsh meshes volumes in order. When it gets to the first volume,
                # it points all adjacent surfaces normals outward. For each subsequent
                # volume, it points the surface normals outwards iff the surface hasn't
                # yet been encountered. Thus, the first time a surface is encountered,
                # the current volume has a forward sense and the second time a reverse
                # sense.
                adjacencies = gmsh.model.get_adjacencies(3, volume_tag)
                surface_tags = adjacencies[1]
                for surface_tag in surface_tags:
                    if surface_tag not in known_surfaces:
                        surface_set_handle = core.create_meshset()
                        surface = _Surface(handle=surface_set_handle)
                        surface.forward_volume = volume_set_handle
                        known_surfaces[surface_tag] = surface

                        core.tag_set_data(
                            tag_handles["global_id"], surface.handle, surface_tag
                        )
                        core.tag_set_data(
                            tag_handles["geom_dimension"], surface.handle, 2
                        )
                        core.tag_set_data(
                            tag_handles["category"], surface.handle, "Surface"
                        )
                        core.tag_set_data(
                            tag_handles["surf_sense"],
                            surface.handle,
                            surface.sense_data(),
                        )

                        # Write surface to MOAB. STL export/import is very efficient.
                        with tempfile.NamedTemporaryFile(
                            suffix=".stl", delete=True
                        ) as stl_file:
                            group_tag = gmsh.model.add_physical_group(2, [surface_tag])
                            gmsh.write(stl_file.name)
                            gmsh.model.remove_physical_groups([(2, group_tag)])
                            core.load_file(stl_file.name, surface_set_handle)

                    else:
                        # Surface already has a forward volume, so this must be the
                        # reverse volume.
                        surface = known_surfaces[surface_tag]
                        surface.reverse_volume = volume_set_handle
                        core.tag_set_data(
                            tag_handles["surf_sense"],
                            surface.handle,
                            surface.sense_data(),
                        )

                    core.add_parent_child(volume_set_handle, surface.handle)

            all_entities = core.get_entities_by_handle(0)
            file_set = core.create_meshset()
            # TODO(akoen): faceting tol set to a random value
            # https://github.com/Thea-Energy/neutronics-cad/issues/5
            # faceting_tol required to be set for make_watertight, although its
            # significance is not clear
            core.tag_set_data(tag_handles["faceting_tol"], file_set, 1e-3)
            core.add_entities(file_set, all_entities)

            return cls(core)

    def _get_entities_of_geom_dimension(self, dim: int) -> list[np.uint64]:
        dim_tag = self._core.tag_get_handle(pymoab.types.GEOM_DIMENSION_TAG_NAME)
        return self._core.get_entities_by_type_and_tag(
            0, pymoab.types.MBENTITYSET, dim_tag, [dim]
        )

    @property
    def surfaces(self):
        """Get surfaces in this model.

        Returns:
            Surfaces.
        """
        surface_handles = self._get_entities_of_geom_dimension(2)
        return [MOABSurface(self._core, h) for h in surface_handles]

    @property
    def volumes(self):
        """Get volumes in this model.

        Returns:
            Volumes.
        """
        volume_handles = self._get_entities_of_geom_dimension(3)
        return [MOABVolume(self._core, h) for h in volume_handles]

src/stellarmesh/moab.py
Thea-Energy/stellarmesh
import build123d as bd
import cadquery as cq
import pytest
import stellarmesh as sm


@pytest.fixture
def geom_bd_layered_torus():
    solids = [bd.Solid.make_torus(1000, 100)]
    for _ in range(3):
        solids.append(solids[-1].faces()[0].thicken(100))
    solids = solids[1:]
    return solids


@pytest.fixture
def geom_cq_layered_torus():
    solids = [cq.Solid.makeTorus(1000, 100)]
    for _ in range(3):
        solids.append(cq.Workplane(solids[-1]).faces().vals()[0].thicken(100))
    solids = solids[1:]
    return solids


@pytest.fixture
def geom_ocp_layered_torus(geom_bd_layered_torus):
    return [s.wrapped for s in geom_bd_layered_torus]


@pytest.mark.parametrize(
    "fixture",
    [("geom_bd_layered_torus"), ("geom_cq_layered_torus"), ("geom_ocp_layered_torus")],
)
def test_geometry_init(fixture, request):
    solids = request.getfixturevalue(fixture)
    material_names = ["material"] * len(solids)
    geom = sm.Geometry(solids, material_names)
    if hasattr(solids[0], "wrapped"):
        assert geom.solids == [s.wrapped for s in solids]
    else:
        assert geom.solids == solids
    assert geom.material_names == material_names


def test_geometry_init_wrong_materials(geom_bd_layered_torus):
    solids = geom_bd_layered_torus
    material_names = ["material"] * (len(solids) - 1)
    with pytest.raises(ValueError):
        sm.Geometry(solids, material_names)


def test_step_import_compound(geom_bd_layered_torus):
    bd.Compound.make_compound(geom_bd_layered_torus).export_step("model.step")
    sm.Geometry.import_step("model.step", material_names=[""] * 3)
    with pytest.raises(ValueError):
        sm.Geometry.import_step("model.step", material_names=[""] * 2)


def test_step_import_solid(geom_bd_layered_torus):
    geom_bd_layered_torus[0].export_step("layer.step")
    sm.Geometry.import_step("layer.step", material_names=[""])


def test_brep_import_compound(geom_bd_layered_torus):
    bd.Compound.make_compound(geom_bd_layered_torus).export_brep("model.brep")
    sm.Geometry.import_brep("model.brep", material_names=[""] * 3)
    with pytest.raises(ValueError):
        sm.Geometry.import_brep("model.brep", material_names=[""] * 2)


def test_brep_import_solid(geom_bd_layered_torus):
    geom_bd_layered_torus[0].export_brep("layer.brep")
    sm.Geometry.import_brep("layer.brep", material_names=[""])


def test_geometry_imprint(geom_bd_layered_torus):
    solids = geom_bd_layered_torus
    material_names = ["material"] * len(solids)
    geom = sm.Geometry(solids, material_names)
    imprinted_geom = geom.imprint()

tests/test_geometry.py
Thea-Energy/stellarmesh