Support arbitrary symmetry axis for CylindricalIndependent class (#3474)

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

docs/source/io_formats/settings.rst

@@ -814,6 +814,7 @@ attributes/sub-elements:
 
       For a "cylindrical" distribution, no parameters are specified. Instead,
       the ``r``, ``phi``, ``z``, and ``origin`` elements must be specified.
+      Optionally, the ``r_dir`` and ``z_dir`` elements could be specified.
 
       For a "spherical" distribution, no parameters are specified. Instead,
       the ``r``, ``theta``, ``phi``, and ``origin`` elements must be specified.
@@ -845,6 +846,10 @@ attributes/sub-elements:
       of a univariate probability distribution (see the description in
       :ref:`univariate`).
 
+    :r_dir:
+      For "cylindrical" distributions, this element specifies the direction
+      of the cylinder r-axis at phi=0. Defaults to (1.0, 0.0, 0.0).
+
     :theta:
       For a "spherical" distribution, this element specifies the distribution
       of theta-coordinates. The necessary sub-elements/attributes are those of a
@@ -857,6 +862,10 @@ attributes/sub-elements:
       sub-elements/attributes are those of a univariate probability
       distribution (see the description in :ref:`univariate`).
 
+    :z_dir:
+      For "cylindrical" distributions, this element specifies the direction
+      of the cylinder z-axis. Defaults to (0.0, 0.0, 1.0).
+
     :origin:
       For "cylindrical and "spherical" distributions, this element specifies
       the coordinates for the origin of the coordinate system.

docs/source/pythonapi/stats.rst

@@ -67,3 +67,4 @@ Spatial Distributions
    :template: myfunction.rst
 
    openmc.stats.spherical_uniform
+   openmc.stats.cylindrical_uniform

include/openmc/distribution_spatial.h

@@ -67,12 +67,18 @@ public:
   Distribution* phi() const { return phi_.get(); }
   Distribution* z() const { return z_.get(); }
   Position origin() const { return origin_; }
+  Direction r_dir() const { return r_dir_; }
+  Direction phi_dir() const { return phi_dir_; }
+  Direction z_dir() const { return z_dir_; }
 
 private:
-  UPtrDist r_;      //!< Distribution of r coordinates
-  UPtrDist phi_;    //!< Distribution of phi coordinates
-  UPtrDist z_;      //!< Distribution of z coordinates
-  Position origin_; //!< Cartesian coordinates of the cylinder center
+  UPtrDist r_;        //!< Distribution of r coordinates
+  UPtrDist phi_;      //!< Distribution of phi coordinates
+  UPtrDist z_;        //!< Distribution of z coordinates
+  Position origin_;   //!< Cartesian coordinates of the cylinder center
+  Direction r_dir_;   //!< Direction of r-axis at phi=0
+  Direction phi_dir_; //!< Direction of phi-axis at phi=0
+  Direction z_dir_;   //!< Direction of z-axis
 };
 
 //==============================================================================

openmc/stats/multivariate.py

@@ -12,7 +12,7 @@ import openmc
 import openmc.checkvalue as cv
 from .._xml import get_elem_list, get_text
 from ..mesh import MeshBase
-from .univariate import PowerLaw, Uniform, Univariate
+from .univariate import PowerLaw, Uniform, Univariate, delta_function
 
 
 class UnitSphere(ABC):
@@ -610,6 +610,10 @@ class CylindricalIndependent(Spatial):
     origin: Iterable of float, optional
         coordinates (x0, y0, z0) of the center of the cylindrical reference
         frame. Defaults to (0.0, 0.0, 0.0)
+    r_dir : Iterable of float, optional
+        Unit vector of the cylinder r axis at phi=0.
+    z_dir : Iterable of float, optional
+        Unit vector of the cylinder z axis direction.
 
     Attributes
     ----------
@@ -623,14 +627,21 @@ class CylindricalIndependent(Spatial):
     origin: Iterable of float, optional
         coordinates (x0, y0, z0) of the center of the cylindrical reference
         frame. Defaults to (0.0, 0.0, 0.0)
+    r_dir : Iterable of float, optional
+        Unit vector of the cylinder r axis at phi=0.
+    z_dir : Iterable of float, optional
+        Unit vector of the cylinder z axis direction.
 
     """
 
-    def __init__(self, r, phi, z, origin=(0.0, 0.0, 0.0)):
+    def __init__(self, r, phi, z, origin=(0.0, 0.0, 0.0), r_dir=(1.0, 0.0, 0.0),
+                 z_dir=(0.0, 0.0, 1.0)):
         self.r = r
         self.phi = phi
         self.z = z
         self.origin = origin
+        self.z_dir = z_dir
+        self.r_dir = r_dir
 
     @property
     def r(self):
@@ -669,6 +680,33 @@ class CylindricalIndependent(Spatial):
         origin = np.asarray(origin)
         self._origin = origin
 
+    @property
+    def z_dir(self):
+        return self._z_dir
+
+    @z_dir.setter
+    def z_dir(self, z_dir):
+        cv.check_type('z-axis direction', z_dir, Iterable, Real)
+        z_dir = np.array(z_dir)
+        norm = np.linalg.norm(z_dir)
+        cv.check_greater_than('z-axis direction magnitude', norm, 0.0)
+        z_dir /= norm
+        self._z_dir = z_dir
+
+    @property
+    def r_dir(self):
+        return self._r_dir
+
+    @r_dir.setter
+    def r_dir(self, r_dir):
+        cv.check_type('r-axis direction', r_dir, Iterable, Real)
+        r_dir = np.array(r_dir)
+        r_dir -= np.dot(r_dir, self.z_dir) * self.z_dir
+        norm = np.linalg.norm(r_dir)
+        cv.check_greater_than('r-axis direction magnitude', norm, 0.0)
+        r_dir /= norm
+        self._r_dir = r_dir
+
     def to_xml_element(self):
         """Return XML representation of the spatial distribution
 
@@ -683,7 +721,12 @@ class CylindricalIndependent(Spatial):
         element.append(self.r.to_xml_element('r'))
         element.append(self.phi.to_xml_element('phi'))
         element.append(self.z.to_xml_element('z'))
-        element.set("origin", ' '.join(map(str, self.origin)))
+        if not np.allclose(self.origin, [0., 0., 0.]):
+            element.set("origin", ' '.join(map(str, self.origin)))
+        if not np.allclose(self.r_dir, [1., 0., 0.]):
+            element.set("r_dir", ' '.join(map(str, self.r_dir)))
+        if not np.allclose(self.z_dir, [0., 0., 1.]):
+            element.set("z_dir", ' '.join(map(str, self.z_dir)))
         return element
 
     @classmethod
@@ -704,8 +747,10 @@ class CylindricalIndependent(Spatial):
         r = Univariate.from_xml_element(elem.find('r'))
         phi = Univariate.from_xml_element(elem.find('phi'))
         z = Univariate.from_xml_element(elem.find('z'))
-        origin = get_elem_list(elem, "origin", float)
-        return cls(r, phi, z, origin=origin)
+        origin = get_elem_list(elem, "origin", float) or [0.0, 0.0, 0.0]
+        r_dir = get_elem_list(elem, "r_dir", float) or [1.0, 0.0, 0.0]
+        z_dir = get_elem_list(elem, "z_dir", float) or [0.0, 0.0, 1.0]
+        return cls(r, phi, z, origin=origin, r_dir=r_dir, z_dir=z_dir)
 
 
 class MeshSpatial(Spatial):
@@ -1219,3 +1264,49 @@ def spherical_uniform(
     phis_dist = Uniform(phis[0], phis[1])
 
     return SphericalIndependent(r_dist, cos_thetas_dist, phis_dist, origin)
+
+
+def cylindrical_uniform(
+    r_outer: float,
+    height: float,
+    r_inner: float = 0.0,
+    phis: Sequence[float] = (0., 2*pi),
+    **kwargs,
+):
+    """Return a uniform spatial distribution over a cylindrical shell.
+
+    This function provides a uniform spatial distribution over a cylindrical
+    shell between `r_inner` and `r_outer`. When `height` is zero, a delta
+    function is used for the z-distribution, giving a uniform distribution over
+    a flat ring (annulus) at z=0 in the local coordinate frame. Optionally, the
+    range of angles can be restricted by the `phis` argument.
+
+    .. versionadded:: 0.15.4
+
+    Parameters
+    ----------
+    r_outer : float
+        Outer radius of the cylindrical shell in [cm]
+    height : float
+        Height of the cylindrical shell in [cm]. When 0, the distribution is a
+        flat ring at z=0 in the local frame.
+    r_inner : float
+        Inner radius of the cylindrical shell in [cm]
+    phis : iterable of float
+        Starting and ending phi coordinates (azimuthal angle) in radians in a
+        reference frame centered at `origin`.
+    **kwargs
+        Keyword arguments passed directly to
+        :class:`~openmc.stats.CylindricalIndependent` (e.g., ``origin``,
+        ``r_dir``, ``z_dir``).
+
+    Returns
+    -------
+    openmc.stats.CylindricalIndependent
+        Uniform distribution over the cylindrical shell
+    """
+
+    r_dist = PowerLaw(r_inner, r_outer, 1)
+    phis_dist = Uniform(phis[0], phis[1])
+    z_dist = delta_function(0.0) if height == 0.0 else Uniform(-height/2, height/2)
+    return CylindricalIndependent(r_dist, phis_dist, z_dist, **kwargs)

src/distribution_spatial.cpp

@@ -141,6 +141,41 @@ CylindricalIndependent::CylindricalIndependent(pugi::xml_node node)
     // If no coordinates were specified, default to (0, 0, 0)
     origin_ = {0.0, 0.0, 0.0};
   }
+
+  // Read cylinder z_dir
+  if (check_for_node(node, "z_dir")) {
+    auto z_dir = get_node_array<double>(node, "z_dir");
+    if (z_dir.size() == 3) {
+      z_dir_ = z_dir;
+      z_dir_ /= z_dir_.norm();
+    } else {
+      fatal_error("z_dir for cylindrical source distribution must be length 3");
+    }
+  } else {
+    // If no z_dir was specified, default to (0, 0, 1)
+    z_dir_ = {0.0, 0.0, 1.0};
+  }
+
+  // Read cylinder r_dir
+  if (check_for_node(node, "r_dir")) {
+    auto r_dir = get_node_array<double>(node, "r_dir");
+    if (r_dir.size() == 3) {
+      r_dir_ = r_dir;
+      r_dir_ /= r_dir_.norm();
+    } else {
+      fatal_error("r_dir for cylindrical source distribution must be length 3");
+    }
+  } else {
+    // If no r_dir was specified, default to (1, 0, 0)
+    r_dir_ = {1.0, 0.0, 0.0};
+  }
+
+  if (r_dir_.dot(z_dir_) > 1e-12)
+    fatal_error("r_dir must be perpendicular to z_dir");
+
+  auto phi_dir = z_dir_.cross(r_dir_);
+  phi_dir /= phi_dir.norm();
+  phi_dir_ = phi_dir;
 }
 
 std::pair<Position, double> CylindricalIndependent::sample(uint64_t* seed) const
@@ -148,10 +183,8 @@ std::pair<Position, double> CylindricalIndependent::sample(uint64_t* seed) const
   auto [r, r_wgt] = r_->sample(seed);
   auto [phi, phi_wgt] = phi_->sample(seed);
   auto [z, z_wgt] = z_->sample(seed);
-  double x = r * cos(phi) + origin_.x;
-  double y = r * sin(phi) + origin_.y;
-  z += origin_.z;
-  Position xi {x, y, z};
+  Position xi =
+    r * (cos(phi) * r_dir_ + sin(phi) * phi_dir_) + z * z_dir_ + origin_;
   return {xi, r_wgt * phi_wgt * z_wgt};
 }
 

tests/unit_tests/test_source.py

@@ -35,21 +35,6 @@ def test_source():
     assert src.strength == 1.0
 
 
-def test_spherical_uniform():
-    r_outer = 2.0
-    r_inner = 1.0
-    thetas = (0.0, pi/2)
-    phis = (0.0, pi)
-    origin = (0.0, 1.0, 2.0)
-
-    sph_indep_function = openmc.stats.spherical_uniform(r_outer,
-                                                        r_inner,
-                                                        thetas,
-                                                        phis,
-                                                        origin)
-
-    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]

tests/unit_tests/test_stats.py

@@ -560,6 +560,113 @@ def test_point():
     assert d.xyz == pytest.approx(p)
 
 
+def test_spherical_uniform():
+    r_outer = 2.0
+    r_inner = 1.0
+    thetas = (0.0, pi/2)
+    phis = (0.0, pi)
+    origin = (0.0, 1.0, 2.0)
+
+    sph_indep_function = openmc.stats.spherical_uniform(r_outer,
+                                                        r_inner,
+                                                        thetas,
+                                                        phis,
+                                                        origin)
+
+    assert isinstance(sph_indep_function, openmc.stats.SphericalIndependent)
+
+
+def test_cylindrical_uniform():
+    r_outer = 2.0
+    r_inner = 1.0
+    height = 1.0
+    phis = (0.0, pi)
+    origin = (0.0, 1.0, 2.0)
+
+    dist = openmc.stats.cylindrical_uniform(r_outer, height, r_inner, phis,
+                                            origin=origin)
+
+    assert isinstance(dist, openmc.stats.CylindricalIndependent)
+
+    # Check r distribution (PowerLaw with exponent 1 for uniform area sampling)
+    assert isinstance(dist.r, openmc.stats.PowerLaw)
+    assert dist.r.a == pytest.approx(r_inner)
+    assert dist.r.b == pytest.approx(r_outer)
+    assert dist.r.n == pytest.approx(1.0)
+
+    # Check phi distribution
+    assert isinstance(dist.phi, openmc.stats.Uniform)
+    assert dist.phi.a == pytest.approx(phis[0])
+    assert dist.phi.b == pytest.approx(phis[1])
+
+    # Check z distribution (centered on origin along z_dir)
+    assert isinstance(dist.z, openmc.stats.Uniform)
+    assert dist.z.a == pytest.approx(-height / 2)
+    assert dist.z.b == pytest.approx(height / 2)
+
+    # Check origin and default directions
+    np.testing.assert_allclose(dist.origin, origin)
+    np.testing.assert_allclose(dist.r_dir, [1., 0., 0.])
+    np.testing.assert_allclose(dist.z_dir, [0., 0., 1.])
+
+    # XML round-trip preserves all parameters
+    elem = dist.to_xml_element()
+    dist2 = openmc.stats.CylindricalIndependent.from_xml_element(elem)
+    np.testing.assert_allclose(dist2.origin, origin)
+    np.testing.assert_allclose(dist2.r_dir, dist.r_dir)
+    np.testing.assert_allclose(dist2.z_dir, dist.z_dir)
+
+
+def test_cylindrical_uniform_tilted():
+    # Test with non-default axis orientation (y-axis as cylinder axis)
+    dist = openmc.stats.cylindrical_uniform(
+        r_outer=3.0, height=2.0, r_dir=(1., 0., 0.), z_dir=(0., 1., 0.)
+    )
+    np.testing.assert_allclose(dist.z_dir, [0., 1., 0.])
+    np.testing.assert_allclose(dist.r_dir, [1., 0., 0.])
+
+    # XML round-trip preserves tilted directions
+    elem = dist.to_xml_element()
+    dist2 = openmc.stats.CylindricalIndependent.from_xml_element(elem)
+    np.testing.assert_allclose(dist2.z_dir, dist.z_dir)
+    np.testing.assert_allclose(dist2.r_dir, dist.r_dir)
+
+
+def test_cylindrical_uniform_ring():
+    # height=0 should produce a flat ring (delta function at z=0)
+    r_outer = 2.0
+    r_inner = 1.0
+    phis = (0.0, pi)
+    origin = (0.0, 1.0, 2.0)
+
+    dist = openmc.stats.cylindrical_uniform(r_outer, 0.0, r_inner, phis,
+                                            origin=origin)
+
+    assert isinstance(dist, openmc.stats.CylindricalIndependent)
+
+    # Check r distribution
+    assert isinstance(dist.r, openmc.stats.PowerLaw)
+    assert dist.r.a == pytest.approx(r_inner)
+    assert dist.r.b == pytest.approx(r_outer)
+    assert dist.r.n == pytest.approx(1.0)
+
+    # Check phi distribution
+    assert isinstance(dist.phi, openmc.stats.Uniform)
+    assert dist.phi.a == pytest.approx(phis[0])
+    assert dist.phi.b == pytest.approx(phis[1])
+
+    # z distribution must be a delta function at 0.0 (local frame)
+    assert isinstance(dist.z, openmc.stats.Discrete)
+    assert dist.z.x[0] == pytest.approx(0.0)
+
+    # XML round-trip
+    elem = dist.to_xml_element()
+    dist2 = openmc.stats.CylindricalIndependent.from_xml_element(elem)
+    np.testing.assert_allclose(dist2.origin, origin)
+    np.testing.assert_allclose(dist2.r_dir, dist.r_dir)
+    np.testing.assert_allclose(dist2.z_dir, dist.z_dir)
+
+
 @pytest.mark.flaky(reruns=1)
 def test_normal():
     mean = 10.0