update util to take a reference vector, and change ww to use vector

Author: nglaser3

mcdc/numba_types.py

@@ -593,9 +593,7 @@
     ('Nx', int64),
     ('Ny', int64),
     ('Nz', int64),
-    ('px', float64),
-    ('py', float64),
-    ('pz', float64),
+    ('polar_reference', float64, (3,)),
     ('lower_weights_offset', int64),
     ('lower_weights_length', int64),
     ('target_weights_offset', int64),

mcdc/object_/technique.py

@@ -100,9 +100,7 @@ class WeightWindows(ObjectSingleton):
     Nz: int
 
     # reference vector for calculating mu and azimuthal angle
-    px: float
-    py: float
-    pz: float
+    polar_reference: Annotated[NDArray[np.float64], (3,)]
 
     # arrays of ww params
     lower_weights: Annotated[
@@ -128,9 +126,7 @@ def __init__(self):
         self.mesh_ID = -1  # skirt around having to create a MeshBase instance
         self.Nx, self.Ny, self.Nz = 1, 1, 1
         self.Nt = 1
-        self.px = 0.0
-        self.py = 0.0
-        self.pz = 1.0
+        self.polar_reference = np.array([0.0, 0.0, 1.0])
         shape = (self.Nt, self.Ne, self.Nmu, self.Na, self.Nx, self.Ny, self.Nz)
         self.lower_weights = np.array([1.0]).reshape(*shape)
         self.target_weights = np.array([1.0]).reshape(*shape)
@@ -217,10 +213,8 @@ def __call__(
             )
 
     def set_polar_reference(self, px, py, pz):
-        norm = np.linalg.norm([px, py, pz])
-        self.px = px / norm
-        self.py = py / norm
-        self.pz = pz / norm
+        polar_reference = np.array([px, py, pz])
+        self.polar_reference = polar_reference / np.linalg.norm(polar_reference)
 
     @staticmethod
     def __check_array(array: NDArray[np.float64], name: str):

mcdc/transport/tally/filter.py

@@ -12,7 +12,11 @@
     COINCIDENCE_TOLERANCE_ENERGY,
     COINCIDENCE_TOLERANCE_TIME,
 )
-from mcdc.transport.util import find_bin_with_tolerance, find_bin_with_rules, calculate_angles
+from mcdc.transport.util import (
+    find_bin_with_tolerance,
+    find_bin_with_rules,
+    calculate_angles,
+)
 
 
 @njit
@@ -34,11 +38,7 @@ def get_filter_indices(particle_container, tally, data, MG_mode):
 @njit
 def get_direction_index(particle_container, tally, data):
     # Polar reference
-    nx = tally["polar_reference"][0]
-    ny = tally["polar_reference"][1]
-    nz = tally["polar_reference"][2]
-
-    mu, azi = calculate_angles(particle_container, nx, ny, nz)
+    mu, azi = calculate_angles(particle_container, tally["polar_reference"])
 
     tolerance = COINCIDENCE_TOLERANCE_DIRECTION
 

mcdc/transport/technique.py

@@ -161,9 +161,7 @@ def get_ww_indices(particle_container, ww_obj, simulation, data):
     ie = util.find_bin(energy, energy_bounds)
 
     # get angular indices
-    mu, azimuthal = util.calculate_angles(
-        particle_container, ww_obj["px"], ww_obj["py"], ww_obj["pz"]
-    )
+    mu, azimuthal = util.calculate_angles(particle_container, ww_obj["polar_reference"])
 
     # mu
     mu_bounds = ww_get.mu_bounds_all(ww_obj, data)

mcdc/transport/util.py

@@ -138,20 +138,16 @@ def log_interpolation(x, x1, x2, y1, y2):
 
 
 @njit
-def calculate_angles(particle_container, px, py, pz):
+def calculate_angles(particle_container, polar_reference):
     """
     Calculate particle mu and azimuthal angle from given reference vector
 
     Parameters
     ----------
     particle_container : ndarray
       Container holding the particle.
-    px : float
-      X-component of reference vector
-    py : float
-      Y-component of reference vector
-    pz : float
-      Z-component of reference vector
+    polar_reference : ndarray
+      3D polar reference vector
 
     Returns
     -------
@@ -165,6 +161,10 @@ def calculate_angles(particle_container, px, py, pz):
     uy = particle["uy"]
     uz = particle["uz"]
 
+    px = polar_reference[0]
+    py = polar_reference[1]
+    pz = polar_reference[2]
+
     mu = ux * px + uy * py + uz * pz
 
     azimuthal = _calculate_azimuthal(ux, uy, uz, px, py, pz)