added option to have cylindrical sources with general axis

Author: GuySten

docs/source/io_formats/settings.rst

@@ -695,6 +695,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.
@@ -726,6 +727,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, it 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
@@ -737,6 +742,10 @@ attributes/sub-elements:
       the distribution of phi-coordinates. The necessary
       sub-elements/attributes are those of a univariate probability
       distribution (see the description in :ref:`univariate`).
+      
+    :phi_dir:
+      For "cylindrical" distributions, this element specifies the direction 
+      of the cylinder phi-axis at phi=0, it defaults to (0.0, 1.0, 0.0).
 
     :origin:
       For "cylindrical and "spherical" distributions, this element specifies

docs/source/pythonapi/stats.rst

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

include/openmc/distribution_spatial.h

@@ -65,12 +65,18 @@ class CylindricalIndependent : public SpatialDistribution {
   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.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):
@@ -532,6 +532,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. Defaults to (1.0, 0.0, 0.0).        
+    z_dir : Iterable of float, optional
+        Unit vector of the cylinder z axis direction. Defaults to (0.0, 0.0, 1.0).    
 
     Attributes
     ----------
@@ -545,14 +549,20 @@ 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. Defaults to (1.0, 0.0, 0.0).        
+    z_dir : Iterable of float, optional
+        Unit vector of the cylinder z axis direction. Defaults to (0.0, 0.0, 1.0).    
 
     """
 
-    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):
@@ -590,6 +600,33 @@ def origin(self, origin):
         cv.check_type('origin coordinates', origin, Iterable, Real)
         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.asarray(z_dir)
+        norm2 = np.dot(z_dir,z_dir)
+        cv.check_greater_than('z-axis direction magnitude', norm2, 0.0)
+        z_dir /= np.sqrt(norm2)
+        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.asarray(r_dir)
+        r_dir -= np.dot(r_dir,self.z_dir)*self.z_dir
+        norm2 = np.dot(r_dir,r_dir)
+        cv.check_greater_than('r-axis direction magnitude', norm2,0.0)
+        r_dir /= np.sqrt(norm2)
+        self._r_dir = r_dir        
 
     def to_xml_element(self):
         """Return XML representation of the spatial distribution
@@ -606,6 +643,8 @@ def to_xml_element(self):
         element.append(self.phi.to_xml_element('phi'))
         element.append(self.z.to_xml_element('z'))
         element.set("origin", ' '.join(map(str, self.origin)))
+        element.set("r_dir", ' '.join(map(str, self.r_dir)))
+        element.set("z_dir", ' '.join(map(str, self.z_dir)))
         return element
 
     @classmethod
@@ -627,7 +666,9 @@ def from_xml_element(cls, elem: ET.Element):
         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)
+        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):
@@ -1086,3 +1127,95 @@ 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),
+        origin: Sequence[float] = (0., 0., 0.),
+        r_dir: Sequence[float] = (1., 0., 0.),
+        z_dir: Sequence[float] = (0., 0., 1.),
+    ):
+    """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`. Optionally, the range of angles
+    can be restricted by the `phis` arguments.
+
+
+    Parameters
+    ----------
+    r_outer : float
+        Outer radius of the cylindrical shell in [cm]
+    height : float
+        Height of the cylindrical shell in [cm]        
+    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`        
+    origin: iterable of float
+        Coordinates (x0, y0, z0) of the center of the cylindrical
+        reference frame for the distribution. Defaults to (0.0, 0.0, 0.0)
+    r_dir : iterable of float
+        Direction of the r-axis at phi=0. Defaults to (1.0, 0.0, 0.0)
+    z_dir : 
+        Direction of the z-axis. Defaults to (0.0, 0.0, 1.0)
+    
+    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 = Uniform(-height/2, height/2)
+    
+    return CylindricalIndependent(r_dist, phis_dist, z_dist, origin, r_dir, z_dir) 
+    
+def ring_uniform(
+        r_outer: float,
+        r_inner: float = 0.0,
+        phis: Sequence[float] = (0., 2*pi),
+        origin: Sequence[float] = (0., 0., 0.),
+        r_dir: Sequence[float] = (1., 0., 0.),
+        z_dir: Sequence[float] = (0., 0., 1.),
+    ):
+    """Return a uniform spatial distribution over a ring.
+
+    This function provides a uniform spatial distribution over a ring
+    shell between `r_inner` and `r_outer`. Optionally, the range of angles
+    can be restricted by the `phis` arguments.
+
+
+    Parameters
+    ----------
+    r_outer : float
+        Outer radius of the ring in [cm]       
+    r_inner : float
+        Inner radius of the ring in [cm]
+    phis : iterable of float
+        Starting and ending phi coordinates (azimuthal angle) in
+        radians in a reference frame centered at `origin`        
+    origin: iterable of float
+        Coordinates (x0, y0, z0) of the center of the cylindrical
+        reference frame for the distribution. Defaults to (0.0, 0.0, 0.0)
+    r_dir : iterable of float
+        Direction of the r-axis at phi=0. Defaults to (1.0, 0.0, 0.0)
+    z_dir : 
+        Direction of the z-axis. Defaults to (0.0, 0.0, 1.0)
+    
+    Returns
+    -------
+    openmc.stats.CylindricalIndependent
+        Uniform distribution over the ring
+    """
+
+    r_dist = PowerLaw(r_inner, r_outer, 1)
+    phis_dist = Uniform(phis[0], phis[1])
+    z_dist = delta_function(np.dot(origin,z_dir))
+    
+    return CylindricalIndependent(r_dist, phis_dist, z_dist, origin, r_dir, z_dir)          

src/distribution_spatial.cpp

@@ -137,16 +137,49 @@ 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;
 }
 
 Position CylindricalIndependent::sample(uint64_t* seed) const
 {
   double r = r_->sample(seed);
   double phi = phi_->sample(seed);
-  double x = r * cos(phi) + origin_.x;
-  double y = r * sin(phi) + origin_.y;
-  double z = z_->sample(seed) + origin_.z;
-  return {x, y, z};
+  double z = z_->sample(seed);
+  return r * (cos(phi) * r_dir_ + sin(phi) * phi_dir_) + z * z_dir_ + origin_;
 }
 
 //==============================================================================

tests/regression_tests/source/inputs_true.dat

@@ -63,7 +63,7 @@
       </energy>
     </source>
     <source type="independent" strength="0.1" particle="neutron">
-      <space type="cylindrical" origin="1.0 1.0 0.0">
+      <space type="cylindrical" origin="1.0 1.0 0.0" r_dir="1.0 0.0 0.0" z_dir="0.0 0.0 1.0">
         <r type="uniform" parameters="2.0 3.0"/>
         <phi type="uniform" parameters="0.0 6.283185307179586"/>
         <z type="tabular" interpolation="linear-linear">
@@ -76,7 +76,7 @@
       </energy>
     </source>
     <source type="independent" strength="0.1" particle="neutron">
-      <space type="cylindrical" origin="1.0 1.0 0.0">
+      <space type="cylindrical" origin="1.0 1.0 0.0" r_dir="1.0 0.0 0.0" z_dir="0.0 0.0 1.0">
         <r type="uniform" parameters="2.0 3.0"/>
         <phi type="uniform" parameters="0.0 6.283185307179586"/>
         <z type="tabular" interpolation="linear-linear">
@@ -123,7 +123,7 @@
       <time type="uniform" parameters="2 5"/>
     </source>
     <source type="independent" strength="0.1" particle="neutron">
-      <space type="cylindrical" origin="1.0 1.0 0.0">
+      <space type="cylindrical" origin="1.0 1.0 0.0" r_dir="1.0 0.0 0.0" z_dir="0.0 0.0 1.0">
         <r type="powerlaw" parameters="2.0 3.0 1.0"/>
         <phi type="uniform" parameters="0.0 6.283185307179586"/>
         <z type="tabular" interpolation="linear-linear">

tests/unit_tests/test_source.py

@@ -49,6 +49,34 @@ def test_spherical_uniform():
                                                         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)
+
+    cyl_indep_function = openmc.stats.cylindrical_uniform(r_outer,
+                                                          height,
+                                                          r_inner,
+                                                          phis,
+                                                          origin)
+
+    assert isinstance(cyl_indep_function, openmc.stats.CylindricalIndependent)
+    
+def test_ring_uniform():
+    r_outer = 2.0
+    r_inner = 1.0
+    phis = (0.0, pi)
+    origin = (0.0, 1.0, 2.0)
+
+    ring_indep_function = openmc.stats.ring_uniform(r_outer,
+                                                    r_inner,
+                                                    phis,
+                                                    origin)
+
+    assert isinstance(ring_indep_function, openmc.stats.CylindricalIndependent)        
 
 def test_point_cloud():
     positions = [(1, 0, 2), (0, 1, 0), (0, 0, 3), (4, 9, 2)]