Adding fix and tests for spherical mesh as spatial distribution

openmc/mesh.py

@@ -5,6 +5,7 @@
 from functools import wraps
 from math import pi, sqrt, atan2
 from numbers import Integral, Real
+from typing import Protocol
 
 import h5py
 import lxml.etree as ET
@@ -842,6 +843,11 @@ def _check_vtk_dataset(self, label: str, dataset: np.ndarray):
             )
 
 
+class HasBoundingBox(Protocol):
+    """Object that has a ``bounding_box`` attribute."""
+    bounding_box: openmc.BoundingBox
+
+
 class RegularMesh(StructuredMesh):
     """A regular Cartesian mesh in one, two, or three dimensions
 
@@ -1088,17 +1094,16 @@ def from_rect_lattice(
     @classmethod
     def from_domain(
         cls,
-        domain: 'openmc.Cell' | 'openmc.Region' | 'openmc.Universe' | 'openmc.Geometry',
+        domain: HasBoundingBox,
         dimension: Sequence[int] = (10, 10, 10),
         mesh_id: int | None = None,
         name: str = ''
     ):
-        """Create mesh from an existing openmc cell, region, universe or
-        geometry by making use of the objects bounding box property.
+        """Create RegularMesh from a domain using its bounding box.
 
         Parameters
         ----------
-        domain : {openmc.Cell, openmc.Region, openmc.Universe, openmc.Geometry}
+        domain : HasBoundingBox
             The object passed in will be used as a template for this mesh. The
             bounding box of the property of the object passed will be used to
             set the lower_left and upper_right and of the mesh instance
@@ -1115,11 +1120,8 @@ def from_domain(
             RegularMesh instance
 
         """
-        cv.check_type(
-            "domain",
-            domain,
-            (openmc.Cell, openmc.Region, openmc.Universe, openmc.Geometry),
-        )
+        if not hasattr(domain, 'bounding_box'):
+            raise TypeError("Domain must have a bounding_box property")
 
         mesh = cls(mesh_id=mesh_id, name=name)
         mesh.lower_left = domain.bounding_box[0]
@@ -1787,17 +1789,18 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
     @classmethod
     def from_domain(
         cls,
-        domain: 'openmc.Cell' | 'openmc.Region' | 'openmc.Universe' | 'openmc.Geometry',
+        domain: HasBoundingBox,
         dimension: Sequence[int] = (10, 10, 10),
         mesh_id: int | None = None,
         phi_grid_bounds: Sequence[float] = (0.0, 2*pi),
-        name: str = ''
+        name: str = '',
+        enclose_domain: bool = False
     ):
-        """Creates a regular CylindricalMesh from an existing openmc domain.
+        """Create CylindricalMesh from a domain using its bounding box.
 
         Parameters
         ----------
-        domain : openmc.Cell or openmc.Region or openmc.Universe or openmc.Geometry
+        domain : HasBoundingBox
             The object passed in will be used as a template for this mesh. The
             bounding box of the property of the object passed will be used to
             set the r_grid, z_grid ranges.
@@ -1811,32 +1814,30 @@ def from_domain(
             is (0, 2π), i.e., the full phi range.
         name : str
             Name of the mesh
+        enclose_domain : bool
+            If True, the mesh will encompass the bounding box of the domain. If
+            False, the mesh will be inscribed within the domain's bounding box.
 
         Returns
         -------
         openmc.CylindricalMesh
             CylindricalMesh instance
 
         """
-        cv.check_type(
-            "domain",
-            domain,
-            (openmc.Cell, openmc.Region, openmc.Universe, openmc.Geometry),
-        )
+        if not hasattr(domain, 'bounding_box'):
+            raise TypeError("Domain must have a bounding_box property")
 
         # loaded once to avoid recalculating bounding box
         cached_bb = domain.bounding_box
-        max_bounding_box_radius = max(
-            [
-                cached_bb[0][0],
-                cached_bb[0][1],
-                cached_bb[1][0],
-                cached_bb[1][1],
-            ]
-        )
+
+        if enclose_domain:
+            outer_radius = 0.5 * np.linalg.norm(cached_bb.width[:2])
+        else:
+            outer_radius = 0.5 * min(cached_bb.width[:2])
+
         r_grid = np.linspace(
             0,
-            max_bounding_box_radius,
+            outer_radius,
             num=dimension[0]+1
         )
         phi_grid = np.linspace(
@@ -2167,6 +2168,77 @@ def from_hdf5(cls, group: h5py.Group, mesh_id: int, name: str):
 
         return mesh
 
+    @classmethod
+    def from_domain(
+        cls,
+        domain: HasBoundingBox,
+        dimension: Sequence[int] = (10, 10, 10),
+        mesh_id: int | None = None,
+        phi_grid_bounds: Sequence[float] = (0.0, 2*pi),
+        theta_grid_bounds: Sequence[float] = (0.0, pi),
+        name: str = '',
+        enclose_domain: bool = False
+    ):
+        """Create SphericalMesh from a domain using its bounding box.
+
+        Parameters
+        ----------
+        domain : HasBoundingBox
+            The object passed in will be used as a template for this mesh. The
+            bounding box of the property of the object passed will be used to
+            set the r_grid, phi_grid, and theta_grid ranges.
+        dimension : Iterable of int
+            The number of equally spaced mesh cells in each direction (r_grid,
+            phi_grid, theta_grid). Spacing is in angular space (radians) for
+            phi and theta, and in absolute space for r.
+        mesh_id : int
+            Unique identifier for the mesh
+        phi_grid_bounds : numpy.ndarray
+            Mesh bounds points along the phi-axis in radians. The default value
+            is (0, 2π), i.e., the full phi range.
+        theta_grid_bounds : numpy.ndarray
+            Mesh bounds points along the theta-axis in radians. The default value
+            is (0, π), i.e., the full theta range.
+        name : str
+            Name of the mesh
+        enclose_domain : bool
+            If True, the mesh will encompass the bounding box of the domain. If
+            False, the mesh will be inscribed within the domain's bounding box.
+
+        Returns
+        -------
+        openmc.SphericalMesh
+            SphericalMesh instance
+
+        """
+        if not hasattr(domain, 'bounding_box'):
+            raise TypeError("Domain must have a bounding_box property")
+
+        # loaded once to avoid recalculating bounding box
+        cached_bb = domain.bounding_box
+
+        if enclose_domain:
+            outer_radius = 0.5 * np.linalg.norm(cached_bb.width)
+        else:
+            outer_radius = 0.5 * min(cached_bb.width)
+
+        r_grid = np.linspace(0, outer_radius, num=dimension[0] + 1)
+        theta_grid = np.linspace(
+            theta_grid_bounds[0],
+            theta_grid_bounds[1],
+            num=dimension[1]+1
+        )
+        phi_grid = np.linspace(
+            phi_grid_bounds[0],
+            phi_grid_bounds[1],
+            num=dimension[2]+1
+        )
+        origin = np.array([
+            cached_bb.center[0], cached_bb.center[1], cached_bb.center[2]])
+
+        return cls(r_grid=r_grid, phi_grid=phi_grid, theta_grid=theta_grid,
+                   origin=origin, mesh_id=mesh_id, name=name)
+
     def to_xml_element(self):
         """Return XML representation of the mesh
 

src/mesh.cpp

@@ -1794,7 +1794,8 @@ Position SphericalMesh::sample_element(
   double phi_min = this->phi(ijk[2] - 1);
   double phi_max = this->phi(ijk[2]);
 
-  double cos_theta = uniform_distribution(theta_min, theta_max, seed);
+  double cos_theta =
+    uniform_distribution(std::cos(theta_min), std::cos(theta_max), seed);
   double sin_theta = std::sin(std::acos(cos_theta));
   double phi = uniform_distribution(phi_min, phi_max, seed);
   double r_min_cub = std::pow(r_min, 3);

tests/unit_tests/test_mesh_from_domain.py

@@ -40,7 +40,7 @@ def test_cylindrical_mesh_from_cell():
 
     assert isinstance(mesh, openmc.CylindricalMesh)
     assert np.array_equal(mesh.dimension, (1, 1, 1))
-    assert np.array_equal(mesh.r_grid, [0., 150.])
+    assert np.array_equal(mesh.r_grid, [0., 50.])
     assert np.array_equal(mesh.origin, [100., 0., 10.])
 
     # Cell is not centralized on Z, X or Y axis
@@ -49,7 +49,7 @@ def test_cylindrical_mesh_from_cell():
     mesh = openmc.CylindricalMesh.from_domain(domain=cell, dimension=[1, 1, 1])
 
     assert isinstance(mesh, openmc.CylindricalMesh)
-    assert np.array_equal(mesh.r_grid, [0., 220.])
+    assert np.array_equal(mesh.r_grid, [0., 50.])
     assert np.array_equal(mesh.origin, [100., 170., 10.])
 
 
@@ -87,6 +87,34 @@ def test_cylindrical_mesh_from_region():
     assert np.array_equal(mesh.origin, (0.0, 0.0, -30.))
 
 
+def test_spherical_mesh_from_domain():
+    """Tests a SphericalMesh can be made from a Region and the specified
+    dimensions are propagated through. Cell is not centralized"""
+    sphere = openmc.Sphere(r=5, x0=2, y0=3, z0=4)
+    region = -sphere
+
+    geometry = openmc.Geometry(openmc.Universe(cells=[openmc.Cell(region=region)]))
+
+    region_mesh = openmc.SphericalMesh.from_domain(
+        domain=region, dimension=(4, 3, 4))
+    universe_mesh = openmc.SphericalMesh.from_domain(
+        domain=geometry.root_universe, dimension=(4, 3, 4))
+    geometry_mesh = openmc.SphericalMesh.from_domain(
+        domain=geometry, dimension=(4, 3, 4))
+
+
+    for mesh in (region_mesh, universe_mesh, geometry_mesh):
+        assert isinstance(mesh, openmc.SphericalMesh)
+        assert np.array_equal(mesh.dimension, (4, 3, 4))
+        assert np.array_equal(mesh.r_grid, [0., 1.25, 2.5, 3.75, 5.0])
+        assert np.array_equal(mesh.theta_grid, [0., np.pi/3., 2*np.pi/3., np.pi])
+        assert np.array_equal(mesh.phi_grid, [0., np.pi/2., np.pi, 3*np.pi/2., 2*np.pi])
+        assert np.array_equal(mesh.origin, (2.0, 3.0, 4.0))
+
+        for p in mesh.centroids.reshape(-1, 3):
+            assert p in mesh.bounding_box
+
+
 def test_reg_mesh_from_universe():
     """Tests a RegularMesh can be made from a Universe and the default
     dimensions are propagated through. Universe is centralized"""

tests/unit_tests/test_source_mesh.py

@@ -396,3 +396,48 @@ def test_mesh_source_file(run_in_tmpdir):
         assert site.u == source_particle.u
         assert site.time == source_particle.time
         assert site.r in bbox
+
+
+@pytest.mark.parametrize("mesh_type", ('rectangular', 'cylindrical', 'spherical'))
+def test_mesh_spatial(run_in_tmpdir, mesh_type):
+    """Test that a spherical mesh source works as expected."""
+    model = openmc.Model()
+
+    # Set up geometry, a box that is shifted in x, y, and z
+    box = openmc.model.RectangularParallelepiped(5.0, 25.0, -20.0, 20.0, -30.0, 30.0, boundary_type='vacuum')
+    mat = openmc.Material()
+    mat.add_nuclide('H1', 1.0)
+    model.geometry = openmc.Geometry([openmc.Cell(fill=mat, region=-box)])
+
+    # Create a mesh of each type in turn
+    if mesh_type == 'rectangular':
+        mesh = openmc.RegularMesh.from_domain(model.geometry, (10, 2, 2))
+    elif mesh_type == 'cylindrical':
+        mesh = openmc.CylindricalMesh.from_domain(model.geometry, (10, 2, 2))
+        assert max(mesh.r_grid) == 10.0, "Cylindrical mesh radius exceeds geometry bounds"
+        assert mesh.origin[0] == 15.0, "Cylindrical mesh origin x-coordinate is incorrect"
+    elif mesh_type == 'spherical':
+        mesh = openmc.SphericalMesh.from_domain(model.geometry, (10, 2, 2))
+        assert max(mesh.r_grid) == 10.0, "Spherical mesh radius exceeds geometry bounds"
+        assert mesh.origin[0] == 15.0, "Spherical mesh origin x-coordinate is incorrect"
+
+    # Create a mesh source with a single particle
+    ind_source = openmc.IndependentSource(space=openmc.stats.MeshSpatial(mesh, np.prod(mesh.dimension)*[1.0]))
+    model.settings.source = ind_source
+
+    model.settings.particles = 100
+    model.settings.batches = 10
+    model.settings.run_mode = 'fixed source'
+
+    model.export_to_model_xml()
+
+    openmc.lib.init()
+    openmc.lib.simulation_init()
+    sites = openmc.lib.sample_external_source(10)
+    openmc.lib.simulation_finalize()
+    openmc.lib.finalize()
+
+    # Check that the sites are within the spherical mesh bounds
+    bbox = mesh.bounding_box
+    for site in sites:
+        assert site.r in bbox