diff --git a/docs/source/io_formats/settings.rst b/docs/source/io_formats/settings.rst index ed7c6273a..db72984a6 100644 --- a/docs/source/io_formats/settings.rst +++ b/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. diff --git a/docs/source/pythonapi/stats.rst b/docs/source/pythonapi/stats.rst index c8318ba86..e0ae74e39 100644 --- a/docs/source/pythonapi/stats.rst +++ b/docs/source/pythonapi/stats.rst @@ -67,3 +67,4 @@ Spatial Distributions :template: myfunction.rst openmc.stats.spherical_uniform + openmc.stats.cylindrical_uniform diff --git a/include/openmc/distribution_spatial.h b/include/openmc/distribution_spatial.h index 3e31b17a9..663cf5568 100644 --- a/include/openmc/distribution_spatial.h +++ b/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 }; //============================================================================== diff --git a/openmc/stats/multivariate.py b/openmc/stats/multivariate.py index 2057f4a49..d48697037 100644 --- a/openmc/stats/multivariate.py +++ b/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) diff --git a/src/distribution_spatial.cpp b/src/distribution_spatial.cpp index e25e08d74..ba8658f10 100644 --- a/src/distribution_spatial.cpp +++ b/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(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(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 CylindricalIndependent::sample(uint64_t* seed) const @@ -148,10 +183,8 @@ std::pair 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}; } diff --git a/tests/unit_tests/test_source.py b/tests/unit_tests/test_source.py index bb8a1b785..6eb03f388 100644 --- a/tests/unit_tests/test_source.py +++ b/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] diff --git a/tests/unit_tests/test_stats.py b/tests/unit_tests/test_stats.py index b9d9c4e1f..71679024f 100644 --- a/tests/unit_tests/test_stats.py +++ b/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