Add PointCloud spatial distribution (#3161)

Co-authored-by: Patrick Shriwise <pshriwise@gmail.com>
Co-authored-by: Paul Romano <paul.k.romano@gmail.com>

docs/source/io_formats/settings.rst

@@ -579,24 +579,38 @@ attributes/sub-elements:
 
     :type:
       The type of spatial distribution. Valid options are "box", "fission",
-      "point", "cartesian", "cylindrical", and "spherical". A "box" spatial
-      distribution has coordinates sampled uniformly in a parallelepiped. A
-      "fission" spatial distribution samples locations from a "box"
+      "point", "cartesian", "cylindrical", "spherical", "mesh", and "cloud".
+
+      A "box" spatial distribution has coordinates sampled uniformly in a
+      parallelepiped.
+
+      A "fission" spatial distribution samples locations from a "box"
       distribution but only locations in fissionable materials are accepted.
+
       A "point" spatial distribution has coordinates specified by a triplet.
+
       A "cartesian" spatial distribution specifies independent distributions of
-      x-, y-, and z-coordinates. A "cylindrical" spatial distribution specifies
-      independent distributions of r-, phi-, and z-coordinates where phi is the
-      azimuthal angle and the origin for the cylindrical coordinate system is
-      specified by origin. A "spherical" spatial distribution specifies
-      independent distributions of r-, cos_theta-, and phi-coordinates where
-      cos_theta is the cosine of the angle with respect to the z-axis, phi is
-      the azimuthal angle, and the sphere is centered on the coordinate
-      (x0,y0,z0). A "mesh" spatial distribution samples source sites from a mesh element
+      x-, y-, and z-coordinates.
+
+      A "cylindrical" spatial distribution specifies independent distributions
+      of r-, phi-, and z-coordinates where phi is the azimuthal angle and the
+      origin for the cylindrical coordinate system is specified by origin.
+
+      A "spherical" spatial distribution specifies independent distributions of
+      r-, cos_theta-, and phi-coordinates where cos_theta is the cosine of the
+      angle with respect to the z-axis, phi is the azimuthal angle, and the
+      sphere is centered on the coordinate (x0,y0,z0).
+
+      A "mesh" spatial distribution samples source sites from a mesh element
       based on the relative strengths provided in the node. Source locations
       within an element are sampled isotropically. If no strengths are provided,
       the space within the mesh is uniformly sampled.
 
+      A "cloud" spatial distribution samples source sites from a list of spatial
+      positions provided in the node, based on the relative strengths provided
+      in the node. If no strengths are provided, the positions are uniformly
+      sampled.
+
       *Default*: None
 
     :parameters:
@@ -662,6 +676,26 @@ attributes/sub-elements:
       For "cylindrical and "spherical" distributions, this element specifies
       the coordinates for the origin of the coordinate system.
 
+    :mesh_id:
+      For "mesh" spatial distributions, this element specifies which mesh ID to
+      use for the geometric description of the mesh.
+
+    :coords:
+      For "cloud" distributions, this element specifies a list of coordinates
+      for each of the points in the cloud.
+
+    :strengths:
+      For "mesh" and "cloud" spatial distributions, this element specifies the
+      relative source strength of each mesh element or each point in the cloud.
+
+    :volume_normalized:
+      For "mesh" spatial distrubtions, this optional boolean element specifies
+      whether the vector of relative strengths should be multiplied by the mesh
+      element volume. This is most common if the strengths represent a source
+      per unit volume.
+
+      *Default*: false
+
   :angle:
     An element specifying the angular distribution of source sites. This element
     has the following attributes:

docs/source/pythonapi/stats.rst

@@ -59,6 +59,7 @@ Spatial Distributions
    openmc.stats.Box
    openmc.stats.Point
    openmc.stats.MeshSpatial
+   openmc.stats.PointCloud
 
 .. autosummary::
    :toctree: generated

docs/source/usersguide/settings.rst

@@ -192,7 +192,9 @@ distributions using spherical or cylindrical coordinates, you can use
 :class:`openmc.stats.SphericalIndependent` or
 :class:`openmc.stats.CylindricalIndependent`, respectively. Meshes can also be
 used to represent spatial distributions with :class:`openmc.stats.MeshSpatial`
-by specifying a mesh and source strengths for each mesh element.
+by specifying a mesh and source strengths for each mesh element. It is also
+possible to define a "cloud" of source points, each with a different relative
+probability, using :class:`openmc.stats.PointCloud`.
 
 The angular distribution can be set equal to a sub-class of
 :class:`openmc.stats.UnitSphere` such as :class:`openmc.stats.Isotropic`,

include/openmc/distribution_spatial.h

@@ -136,6 +136,26 @@ private:
                                 //!< mesh element indices
 };
 
+//==============================================================================
+//! Distribution of points
+//==============================================================================
+
+class PointCloud : public SpatialDistribution {
+public:
+  explicit PointCloud(pugi::xml_node node);
+  explicit PointCloud(
+    std::vector<Position> point_cloud, gsl::span<const double> strengths);
+
+  //! Sample a position from the distribution
+  //! \param seed Pseudorandom number seed pointer
+  //! \return Sampled position
+  Position sample(uint64_t* seed) const override;
+
+private:
+  std::vector<Position> point_cloud_;
+  DiscreteIndex point_idx_dist_; //!< Distribution of Position indices
+};
+
 //==============================================================================
 //! Uniform distribution of points over a box
 //==============================================================================

include/openmc/xml_interface.h

@@ -50,6 +50,9 @@ xt::xarray<T> get_node_xarray(
   return xt::adapt(v, shape);
 }
 
+std::vector<Position> get_node_position_array(
+  pugi::xml_node node, const char* name, bool lowercase = false);
+
 Position get_node_position(
   pugi::xml_node node, const char* name, bool lowercase = false);
 

openmc/stats/multivariate.py

@@ -279,6 +279,8 @@ class Spatial(ABC):
             return Point.from_xml_element(elem)
         elif distribution == 'mesh':
             return MeshSpatial.from_xml_element(elem, meshes)
+        elif distribution == 'cloud':
+            return PointCloud.from_xml_element(elem)
 
 
 class CartesianIndependent(Spatial):
@@ -756,6 +758,118 @@ class MeshSpatial(Spatial):
         return cls(meshes[mesh_id], strengths, volume_normalized)
 
 
+class PointCloud(Spatial):
+    """Spatial distribution from a point cloud.
+
+    This distribution specifies a discrete list of points, with corresponding
+    relative probabilities.
+
+    .. versionadded:: 0.15.1
+
+    Parameters
+    ----------
+    positions : iterable of 3-tuples
+        The points in space to be sampled
+    strengths : iterable of float, optional
+        An iterable of values that represents the relative probabilty of each
+        point.
+
+    Attributes
+    ----------
+    positions : numpy.ndarray
+        The points in space to be sampled with shape (N, 3)
+    strengths : numpy.ndarray or None
+        An array of relative probabilities for each mesh point
+    """
+
+    def __init__(
+        self,
+        positions: Sequence[Sequence[float]],
+        strengths: Sequence[float] | None = None
+    ):
+        self.positions = positions
+        self.strengths = strengths
+
+    @property
+    def positions(self) -> np.ndarray:
+        return self._positions
+
+    @positions.setter
+    def positions(self, positions):
+        positions = np.array(positions, dtype=float)
+        if positions.ndim != 2:
+            raise ValueError('positions must be a 2D array')
+        elif positions.shape[1] != 3:
+            raise ValueError('Each position must have 3 values')
+        self._positions = positions
+
+    @property
+    def strengths(self) -> np.ndarray:
+        return self._strengths
+
+    @strengths.setter
+    def strengths(self, strengths):
+        if strengths is not None:
+            strengths = np.array(strengths, dtype=float)
+            if strengths.ndim != 1:
+                raise ValueError('strengths must be a 1D array')
+            elif strengths.size != self.positions.shape[0]:
+                raise ValueError('strengths must have the same length as positions')
+        self._strengths = strengths
+
+    @property
+    def num_strength_bins(self) -> int:
+        if self.strengths is None:
+            raise ValueError('Strengths are not set')
+        return self.strengths.size
+
+    def to_xml_element(self) -> ET.Element:
+        """Return XML representation of the spatial distribution
+
+        Returns
+        -------
+        element : lxml.etree._Element
+            XML element containing spatial distribution data
+
+        """
+        element = ET.Element('space')
+        element.set('type', 'cloud')
+
+        subelement = ET.SubElement(element, 'coords')
+        subelement.text = ' '.join(str(e) for e in self.positions.flatten())
+
+        if self.strengths is not None:
+            subelement = ET.SubElement(element, 'strengths')
+            subelement.text = ' '.join(str(e) for e in self.strengths)
+
+        return element
+
+    @classmethod
+    def from_xml_element(cls, elem: ET.Element) -> PointCloud:
+        """Generate spatial distribution from an XML element
+
+        Parameters
+        ----------
+        elem : lxml.etree._Element
+            XML element
+
+        Returns
+        -------
+        openmc.stats.PointCloud
+            Spatial distribution generated from XML element
+
+
+        """
+        coord_data = get_text(elem, 'coords')
+        positions = np.array([float(b) for b in coord_data.split()]).reshape((-1, 3))
+
+        strengths = get_text(elem, 'strengths')
+        if strengths is not None:
+            strengths = [float(b) for b in strengths.split()]
+
+        return cls(positions, strengths)
+
+
 class Box(Spatial):
     """Uniform distribution of coordinates in a rectangular cuboid.
 

src/distribution_spatial.cpp

@@ -26,6 +26,8 @@ unique_ptr<SpatialDistribution> SpatialDistribution::create(pugi::xml_node node)
     return UPtrSpace {new SphericalIndependent(node)};
   } else if (type == "mesh") {
     return UPtrSpace {new MeshSpatial(node)};
+  } else if (type == "cloud") {
+    return UPtrSpace {new PointCloud(node)};
   } else if (type == "box") {
     return UPtrSpace {new SpatialBox(node)};
   } else if (type == "fission") {
@@ -298,6 +300,49 @@ Position MeshSpatial::sample(uint64_t* seed) const
   return this->sample_mesh(seed).second;
 }
 
+//==============================================================================
+// PointCloud implementation
+//==============================================================================
+
+PointCloud::PointCloud(pugi::xml_node node)
+{
+  if (check_for_node(node, "coords")) {
+    point_cloud_ = get_node_position_array(node, "coords");
+  } else {
+    fatal_error("No coordinates were provided for the PointCloud "
+                "spatial distribution");
+  }
+
+  std::vector<double> strengths;
+
+  if (check_for_node(node, "strengths"))
+    strengths = get_node_array<double>(node, "strengths");
+  else
+    strengths.resize(point_cloud_.size(), 1.0);
+
+  if (strengths.size() != point_cloud_.size()) {
+    fatal_error(
+      fmt::format("Number of entries for the strengths array {} does "
+                  "not match the number of spatial points provided {}.",
+        strengths.size(), point_cloud_.size()));
+  }
+
+  point_idx_dist_.assign(strengths);
+}
+
+PointCloud::PointCloud(
+  std::vector<Position> point_cloud, gsl::span<const double> strengths)
+{
+  point_cloud_.assign(point_cloud.begin(), point_cloud.end());
+  point_idx_dist_.assign(strengths);
+}
+
+Position PointCloud::sample(uint64_t* seed) const
+{
+  int32_t index = point_idx_dist_.sample(seed);
+  return point_cloud_[index];
+}
+
 //==============================================================================
 // SpatialBox implementation
 //==============================================================================

src/xml_interface.cpp

@@ -4,6 +4,7 @@
 
 #include "openmc/error.h"
 #include "openmc/string_utils.h"
+#include "openmc/vector.h"
 
 namespace openmc {
 
@@ -48,6 +49,24 @@ bool get_node_value_bool(pugi::xml_node node, const char* name)
   return false;
 }
 
+vector<Position> get_node_position_array(
+  pugi::xml_node node, const char* name, bool lowercase)
+{
+  vector<double> coords = get_node_array<double>(node, name, lowercase);
+  if (coords.size() % 3 != 0) {
+    fatal_error(fmt::format(
+      "Incorect number of coordinates in Position array ({}) for \"{}\"",
+      coords.size(), name));
+  }
+  vector<Position> positions;
+  positions.reserve(coords.size() / 3);
+  auto it = coords.begin();
+  for (size_t i = 0; i < coords.size(); i += 3) {
+    positions.push_back({coords[i], coords[i + 1], coords[i + 2]});
+  }
+  return positions;
+}
+
 Position get_node_position(
   pugi::xml_node node, const char* name, bool lowercase)
 {

tests/unit_tests/test_source.py

@@ -1,3 +1,4 @@
+from collections import Counter
 from math import pi
 
 import openmc
@@ -49,6 +50,61 @@ def test_spherical_uniform():
 
     assert isinstance(sph_indep_function, openmc.stats.SphericalIndependent)
 
+def test_point_cloud():
+    positions = [(1, 0, 2), (0, 1, 0), (0, 0, 3), (4, 9, 2)]
+    strengths = [1, 2, 3, 4]
+
+    space = openmc.stats.PointCloud(positions, strengths)
+    np.testing.assert_equal(space.positions, positions)
+    np.testing.assert_equal(space.strengths, strengths)
+
+    src = openmc.IndependentSource(space=space)
+    assert src.space == space
+    np.testing.assert_equal(src.space.positions, positions)
+    np.testing.assert_equal(src.space.strengths, strengths)
+
+    elem = src.to_xml_element()
+    src = openmc.IndependentSource.from_xml_element(elem)
+    np.testing.assert_equal(src.space.positions, positions)
+    np.testing.assert_equal(src.space.strengths, strengths)
+
+
+def test_point_cloud_invalid():
+    with pytest.raises(ValueError, match='2D'):
+        openmc.stats.PointCloud([1, 0, 2, 0, 1, 0])
+
+    with pytest.raises(ValueError, match='3 values'):
+        openmc.stats.PointCloud([(1, 0, 2, 3), (4, 5, 2, 3)])
+
+    with pytest.raises(ValueError, match='1D'):
+        openmc.stats.PointCloud([(1, 0, 2), (4, 5, 2)], [(1, 2), (3, 4)])
+
+    with pytest.raises(ValueError, match='same length'):
+        openmc.stats.PointCloud([(1, 0, 2), (4, 5, 2)], [1, 2, 4])
+
+
+def test_point_cloud_strengths(run_in_tmpdir, sphere_box_model):
+    positions = [(1., 0., 2.), (0., 1., 0.), (0., 0., 3.), (-1., -1., 2.)]
+    strengths = [1, 2, 3, 4]
+    space = openmc.stats.PointCloud(positions, strengths)
+
+    model = sphere_box_model[0]
+    model.settings.run_mode = 'fixed source'
+    model.settings.source = openmc.IndependentSource(space=space)
+
+    try:
+        model.init_lib()
+        n_samples = 50_000
+        sites = openmc.lib.sample_external_source(n_samples)
+    finally:
+        model.finalize_lib()
+
+    count = Counter(s.r for s in sites)
+    for i, (strength, position) in enumerate(zip(strengths, positions)):
+        sampled_strength = count[position] / n_samples
+        expected_strength = pytest.approx(strength/sum(strengths), abs=0.02)
+        assert sampled_strength == expected_strength, f'Strength incorrect for {positions[i]}'
+
 
 def test_source_file():
     filename = 'source.h5'