Enable specifying reference direction for azimuthal angle in PolarAzimuthal distribution (#3582)

Co-authored-by: shimwell <[email protected]>

Author: GuySten

include/openmc/distribution_multi.h

@@ -51,6 +51,8 @@ class PolarAzimuthal : public UnitSphereDistribution {
   Distribution* phi() const { return phi_.get(); }
 
 private:
+  Direction v_ref_ {1.0, 0.0, 0.0}; //!< reference direction
+  Direction w_ref_;
   UPtrDist mu_;  //!< Distribution of polar angle
   UPtrDist phi_; //!< Distribution of azimuthal angle
 };

openmc/stats/multivariate.py

@@ -79,6 +79,9 @@ class PolarAzimuthal(UnitSphere):
     reference_uvw : Iterable of float
         Direction from which polar angle is measured. Defaults to the positive
         z-direction.
+    reference_vwu : Iterable of float
+        Direction from which azimuthal angle is measured. Defaults to the positive
+        x-direction.    
 
     Attributes
     ----------
@@ -89,8 +92,9 @@ class PolarAzimuthal(UnitSphere):
 
     """
 
-    def __init__(self, mu=None, phi=None, reference_uvw=(0., 0., 1.)):
+    def __init__(self, mu=None, phi=None, reference_uvw=(0., 0., 1.), reference_vwu=(1., 0., 0.)):
         super().__init__(reference_uvw)
+        self.reference_vwu = reference_vwu
         if mu is not None:
             self.mu = mu
         else:
@@ -100,6 +104,20 @@ def __init__(self, mu=None, phi=None, reference_uvw=(0., 0., 1.)):
             self.phi = phi
         else:
             self.phi = Uniform(0., 2*pi)
+            
+    @property
+    def reference_vwu(self):
+        return self._reference_vwu
+
+    @reference_vwu.setter
+    def reference_vwu(self, vwu):
+        cv.check_type('reference v direction', vwu, Iterable, Real)
+        vwu = np.asarray(vwu)
+        uvw = self.reference_uvw
+        cv.check_greater_than('reference v direction must not be parallel to reference u direction', np.linalg.norm(np.cross(vwu,uvw)), 1e-6*np.linalg.norm(vwu))        
+        vwu -= vwu.dot(uvw)*uvw        
+        cv.check_less_than('reference v direction must be orthogonal to reference u direction', np.abs(vwu.dot(uvw)), 1e-6)
+        self._reference_vwu = vwu/np.linalg.norm(vwu)
 
     @property
     def mu(self):
@@ -132,6 +150,8 @@ def to_xml_element(self):
         element.set("type", "mu-phi")
         if self.reference_uvw is not None:
             element.set("reference_uvw", ' '.join(map(str, self.reference_uvw)))
+        if self.reference_vwu is not None:
+            element.set("reference_vwu", ' '.join(map(str, self.reference_vwu)))    
         element.append(self.mu.to_xml_element('mu'))
         element.append(self.phi.to_xml_element('phi'))
         return element
@@ -155,6 +175,9 @@ def from_xml_element(cls, elem):
         uvw = get_elem_list(elem, "reference_uvw", float)
         if uvw is not None:
             mu_phi.reference_uvw = uvw
+        vwu = get_elem_list(elem, "reference_vwu", float)
+        if vwu is not None:
+            mu_phi.reference_vwu = vwu            
         mu_phi.mu = Univariate.from_xml_element(elem.find('mu'))
         mu_phi.phi = Univariate.from_xml_element(elem.find('phi'))
         return mu_phi

src/distribution_multi.cpp

@@ -58,6 +58,15 @@ PolarAzimuthal::PolarAzimuthal(Direction u, UPtrDist mu, UPtrDist phi)
 PolarAzimuthal::PolarAzimuthal(pugi::xml_node node)
   : UnitSphereDistribution {node}
 {
+  // Read reference directional unit vector
+  if (check_for_node(node, "reference_vwu")) {
+    auto v_ref = get_node_array<double>(node, "reference_vwu");
+    if (v_ref.size() != 3)
+      fatal_error("Angular distribution reference v direction must have "
+                  "three parameters specified.");
+    v_ref_ = Direction(v_ref.data());
+  }
+  w_ref_ = u_ref_.cross(v_ref_);
   if (check_for_node(node, "mu")) {
     pugi::xml_node node_dist = node.child("mu");
     mu_ = distribution_from_xml(node_dist);
@@ -79,11 +88,15 @@ Direction PolarAzimuthal::sample(uint64_t* seed) const
   double mu = mu_->sample(seed);
   if (mu == 1.0)
     return u_ref_;
+  if (mu == -1.0)
+    return -u_ref_;
 
   // Sample azimuthal angle
   double phi = phi_->sample(seed);
 
-  return rotate_angle(u_ref_, mu, &phi, seed);
+  double f = std::sqrt(1 - mu * mu);
+
+  return mu * u_ref_ + f * std::cos(phi) * v_ref_ + f * std::sin(phi) * w_ref_;
 }
 
 //==============================================================================

tests/regression_tests/source/inputs_true.dat

@@ -25,7 +25,7 @@
           <parameters>-2.0 0.0 2.0 0.2 0.3 0.2</parameters>
         </z>
       </space>
-      <angle type="mu-phi" reference_uvw="0.0 0.0 1.0">
+      <angle type="mu-phi" reference_uvw="0.0 0.0 1.0" reference_vwu="1.0 0.0 0.0">
         <mu type="discrete">
           <parameters>-1.0 0.0 1.0 0.5 0.25 0.25</parameters>
         </mu>

tests/regression_tests/source/results_true.dat

@@ -1,2 +1,2 @@
 k-combined:
-3.034717E-01 2.799386E-03
+3.080655E-01 4.837707E-03

tests/unit_tests/test_stats.py

@@ -516,3 +516,33 @@ def test_combine_distributions():
     # uncertainty of the expected value
     samples = combined.sample(10_000)
     assert_sample_mean(samples, 0.25)
+
+def test_reference_vwu_projection():
+    """When a non-orthogonal vector is provided, the setter should project out
+    any component along reference_uvw so the stored vector is orthogonal.
+    """
+    pa = openmc.stats.PolarAzimuthal()  # default reference_uvw == (0, 0, 1)
+
+    # Provide a vector that is not orthogonal to (0,0,1)
+    pa.reference_vwu = (2.0, 0.5, 0.3)
+
+    reference_v = np.asarray(pa.reference_vwu)
+    reference_u = np.asarray(pa.reference_uvw)
+
+    # reference_v should be orthogonal to reference_u
+    assert abs(np.dot(reference_v, reference_u)) < 1e-6
+
+
+def test_reference_vwu_normalization():
+    """When a non-normalized vector is provided, the setter should normalize
+    the projected vector to unit length.
+    """
+    pa = openmc.stats.PolarAzimuthal()  # default reference_uvw == (0, 0, 1)
+
+    # Provide a vector that is neither orthogonal to (0,0,1) nor unit-length
+    pa.reference_vwu = (2.0, 0.5, 0.3)
+
+    reference_v = np.asarray(pa.reference_vwu)
+
+    # reference_v should be unit length
+    assert np.isclose(np.linalg.norm(reference_v), 1.0, atol=1e-12)