Support rotation in MeshFilter (#3176)

Co-authored-by: Jonathan Shimwell <[email protected]>
Co-authored-by: Patrick Shriwise <[email protected]>

Author: 3 people

include/openmc/tallies/filter_mesh.h

@@ -51,13 +51,21 @@ class MeshFilter : public Filter {
 
   virtual bool translated() const { return translated_; }
 
+  virtual void set_rotation(const vector<double>& rotation);
+
+  virtual const vector<double>& rotation() const { return rotation_; }
+
+  virtual bool rotated() const { return rotated_; }
+
 protected:
   //----------------------------------------------------------------------------
   // Data members
 
   int32_t mesh_;            //!< Index of the mesh
   bool translated_ {false}; //!< Whether or not the filter is translated
   Position translation_ {0.0, 0.0, 0.0}; //!< Filter translation
+  bool rotated_ {false};    //!< Whether or not the filter is rotated
+  vector<double> rotation_; //!< Filter rotation
 };
 
 } // namespace openmc

openmc/filter.py

@@ -833,6 +833,25 @@ class MeshFilter(Filter):
     translation : Iterable of float
         This array specifies a vector that is used to translate (shift) the mesh
         for this filter
+    rotation : Iterable of float
+        This array specifies the angles in degrees about the x, y, and z axes
+        that the mesh should be rotated. The rotation applied is an intrinsic
+        rotation with specified Tait-Bryan angles. That is to say, if the angles
+        are :math:`(\phi, \theta, \psi)`, then the rotation matrix applied is
+        :math:`R_z(\psi) R_y(\theta) R_x(\phi)` or
+
+        .. math::
+
+           \left [ \begin{array}{ccc} \cos\theta \cos\psi & -\cos\phi \sin\psi
+           + \sin\phi \sin\theta \cos\psi & \sin\phi \sin\psi + \cos\phi
+           \sin\theta \cos\psi \\ \cos\theta \sin\psi & \cos\phi \cos\psi +
+           \sin\phi \sin\theta \sin\psi & -\sin\phi \cos\psi + \cos\phi
+           \sin\theta \sin\psi \\ -\sin\theta & \sin\phi \cos\theta & \cos\phi
+           \cos\theta \end{array} \right ]
+
+        A rotation matrix can also be specified directly by setting this
+        attribute to a nested list (or 2D numpy array) that specifies each
+        element of the matrix.
     bins : list of tuple
         A list of mesh indices for each filter bin, e.g. [(1, 1, 1), (2, 1, 1),
         ...]
@@ -845,6 +864,7 @@ def __init__(self, mesh, filter_id=None):
         self.mesh = mesh
         self.id = filter_id
         self._translation = None
+        self._rotation = None
 
     def __hash__(self):
         string = type(self).__name__ + '\n'
@@ -856,6 +876,7 @@ def __repr__(self):
         string += '{: <16}=\t{}\n'.format('\tMesh ID', self.mesh.id)
         string += '{: <16}=\t{}\n'.format('\tID', self.id)
         string += '{: <16}=\t{}\n'.format('\tTranslation', self.translation)
+        string += '{: <16}=\t{}\n'.format('\tRotation', self.rotation)
         return string
 
     @classmethod
@@ -879,6 +900,10 @@ def from_hdf5(cls, group, **kwargs):
         if translation:
             out.translation = translation[()]
 
+        rotation = group.get('rotation')
+        if rotation:
+            out.rotation = rotation[()]
+
         return out
 
     @property
@@ -911,6 +936,15 @@ def translation(self, t):
         cv.check_length('mesh filter translation', t, 3)
         self._translation = np.asarray(t)
 
+    @property
+    def rotation(self):
+        return self._rotation
+
+    @rotation.setter
+    def rotation(self, rotation):
+        cv.check_length('mesh filter rotation', rotation, 3)
+        self._rotation = np.asarray(rotation)
+
     def can_merge(self, other):
         # Mesh filters cannot have more than one bin
         return False
@@ -996,6 +1030,8 @@ def to_xml_element(self):
         subelement.text = str(self.mesh.id)
         if self.translation is not None:
             element.set('translation', ' '.join(map(str, self.translation)))
+        if self.rotation is not None:
+            element.set('rotation', ' '.join(map(str, self.rotation.ravel())))
         return element
 
     @classmethod
@@ -1008,6 +1044,13 @@ def from_xml_element(cls, elem: ET.Element, **kwargs) -> MeshFilter:
         translation = get_elem_list(elem, "translation", float) or []
         if translation:
             out.translation = translation
+
+        rotation = get_elem_list(elem, 'rotation', float) or []
+        if rotation:
+            if len(rotation) == 3:
+                out.rotation = rotation
+            elif len(rotation) == 9:
+                out.rotation = np.array(rotation).reshape(3, 3)
         return out
 
 

openmc/lib/filter.py

@@ -97,6 +97,14 @@
 _dll.openmc_mesh_filter_set_translation.argtypes = [c_int32, POINTER(c_double*3)]
 _dll.openmc_mesh_filter_set_translation.restype = c_int
 _dll.openmc_mesh_filter_set_translation.errcheck = _error_handler
+_dll.openmc_mesh_filter_get_rotation.argtypes = [c_int32, POINTER(c_double),
+    POINTER(c_size_t)]
+_dll.openmc_mesh_filter_get_rotation.restype = c_int
+_dll.openmc_mesh_filter_get_rotation.errcheck = _error_handler
+_dll.openmc_mesh_filter_set_rotation.argtypes = [
+    c_int32, POINTER(c_double), c_size_t]
+_dll.openmc_mesh_filter_set_rotation.restype = c_int
+_dll.openmc_mesh_filter_set_rotation.errcheck = _error_handler
 _dll.openmc_meshborn_filter_get_mesh.argtypes = [c_int32, POINTER(c_int32)]
 _dll.openmc_meshborn_filter_get_mesh.restype = c_int
 _dll.openmc_meshborn_filter_get_mesh.errcheck = _error_handler
@@ -393,6 +401,10 @@ class MeshFilter(Filter):
         Mesh used for the filter
     translation : Iterable of float
         3-D coordinates of the translation vector
+    rotation : Iterable of float
+        The rotation matrix or angles of the filter mesh. This can either be 
+        a fully specified 3 x 3 rotation matrix or an Iterable of length 3 
+        with the angles in degrees about the x, y, and z axes, respectively.
 
     """
     filter_type = 'mesh'
@@ -422,6 +434,34 @@ def translation(self):
     def translation(self, translation):
         _dll.openmc_mesh_filter_set_translation(self._index, (c_double*3)(*translation))
 
+    @property
+    def rotation(self):
+        rotation_data = np.zeros(12)
+        rot_size = c_size_t()
+
+        _dll.openmc_mesh_filter_get_rotation(
+            self._index, rotation_data.ctypes.data_as(POINTER(c_double)),
+            rot_size)
+        rot_size = rot_size.value
+
+        if rot_size == 9:
+            return rotation_data[:rot_size].shape(3, 3)
+        elif rot_size in (0, 12):
+            # If size is 0, rotation_data[9:] will be zeros. This indicates no
+            # rotation and is the most straightforward way to always return
+            # an iterable of floats
+            return rotation_data[9:]
+        else:
+            raise ValueError(
+                f'Invalid size of rotation matrix: {rot_size}')
+        
+    @rotation.setter
+    def rotation(self, rotation_data):
+        flat_rotation = np.asarray(rotation_data, dtype=float).flatten()
+
+        _dll.openmc_mesh_filter_set_rotation(
+            self._index, flat_rotation.ctypes.data_as(POINTER(c_double)),
+            c_size_t(len(flat_rotation)))
 
 class MeshBornFilter(Filter):
     """MeshBorn filter stored internally.

src/cell.cpp

@@ -57,7 +57,7 @@ void Cell::set_rotation(const vector<double>& rot)
     fatal_error(fmt::format("Non-3D rotation vector applied to cell {}", id_));
   }
 
-  // Compute and store the rotation matrix.
+  // Compute and store the inverse rotation matrix for the angles given.
   rotation_.clear();
   rotation_.reserve(rot.size() == 9 ? 9 : 12);
   if (rot.size() == 3) {

src/tallies/filter_mesh.cpp

@@ -6,6 +6,7 @@
 #include "openmc/constants.h"
 #include "openmc/error.h"
 #include "openmc/mesh.h"
+#include "openmc/position.h"
 #include "openmc/xml_interface.h"
 
 namespace openmc {
@@ -30,6 +31,10 @@ void MeshFilter::from_xml(pugi::xml_node node)
   if (check_for_node(node, "translation")) {
     set_translation(get_node_array<double>(node, "translation"));
   }
+  // Read the rotation transform.
+  if (check_for_node(node, "rotation")) {
+    set_rotation(get_node_array<double>(node, "rotation"));
+  }
 }
 
 void MeshFilter::get_all_bins(
@@ -45,6 +50,12 @@ void MeshFilter::get_all_bins(
     last_r -= translation();
     r -= translation();
   }
+  // apply rotation if present
+  if (!rotation_.empty()) {
+    last_r = last_r.rotate(rotation_);
+    r = r.rotate(rotation_);
+    u = u.rotate(rotation_);
+  }
 
   if (estimator != TallyEstimator::TRACKLENGTH) {
     auto bin = model::meshes[mesh_]->get_bin(r);
@@ -65,6 +76,9 @@ void MeshFilter::to_statepoint(hid_t filter_group) const
   if (translated_) {
     write_dataset(filter_group, "translation", translation_);
   }
+  if (rotated_) {
+    write_dataset(filter_group, "rotation", rotation_);
+  }
 }
 
 std::string MeshFilter::text_label(int bin) const
@@ -93,6 +107,40 @@ void MeshFilter::set_translation(const double translation[3])
   this->set_translation({translation[0], translation[1], translation[2]});
 }
 
+void MeshFilter::set_rotation(const vector<double>& rot)
+{
+  rotated_ = true;
+
+  // Compute and store the inverse rotation matrix for the angles given.
+  rotation_.clear();
+  rotation_.reserve(rot.size() == 9 ? 9 : 12);
+  if (rot.size() == 3) {
+    double phi = -rot[0] * PI / 180.0;
+    double theta = -rot[1] * PI / 180.0;
+    double psi = -rot[2] * PI / 180.0;
+    rotation_.push_back(std::cos(theta) * std::cos(psi));
+    rotation_.push_back(-std::cos(phi) * std::sin(psi) +
+                        std::sin(phi) * std::sin(theta) * std::cos(psi));
+    rotation_.push_back(std::sin(phi) * std::sin(psi) +
+                        std::cos(phi) * std::sin(theta) * std::cos(psi));
+    rotation_.push_back(std::cos(theta) * std::sin(psi));
+    rotation_.push_back(std::cos(phi) * std::cos(psi) +
+                        std::sin(phi) * std::sin(theta) * std::sin(psi));
+    rotation_.push_back(-std::sin(phi) * std::cos(psi) +
+                        std::cos(phi) * std::sin(theta) * std::sin(psi));
+    rotation_.push_back(-std::sin(theta));
+    rotation_.push_back(std::sin(phi) * std::cos(theta));
+    rotation_.push_back(std::cos(phi) * std::cos(theta));
+
+    // When user specifies angles, write them at end of vector
+    rotation_.push_back(rot[0]);
+    rotation_.push_back(rot[1]);
+    rotation_.push_back(rot[2]);
+  } else {
+    std::copy(rot.begin(), rot.end(), std::back_inserter(rotation_));
+  }
+}
+
 //==============================================================================
 // C-API functions
 //==============================================================================
@@ -201,4 +249,48 @@ extern "C" int openmc_mesh_filter_set_translation(
   return 0;
 }
 
+//! Return the rotation matrix of a mesh filter
+extern "C" int openmc_mesh_filter_get_rotation(
+  int32_t index, double rot[], size_t* n)
+{
+  // Make sure this is a valid index to an allocated filter
+  if (int err = verify_filter(index))
+    return err;
+
+  // Check the filter type
+  const auto& filter = model::tally_filters[index];
+  if (filter->type() != FilterType::MESH) {
+    set_errmsg("Tried to get a rotation from a non-mesh filter.");
+    return OPENMC_E_INVALID_TYPE;
+  }
+  // Get rotation from the mesh filter and set value
+  auto mesh_filter = dynamic_cast<MeshFilter*>(filter.get());
+  *n = mesh_filter->rotation().size();
+  std::memcpy(rot, mesh_filter->rotation().data(),
+    *n * sizeof(mesh_filter->rotation()[0]));
+  return 0;
+}
+
+//! Set the flattened rotation matrix of a mesh filter
+extern "C" int openmc_mesh_filter_set_rotation(
+  int32_t index, const double rot[], size_t rot_len)
+{
+  // Make sure this is a valid index to an allocated filter
+  if (int err = verify_filter(index))
+    return err;
+
+  const auto& filter = model::tally_filters[index];
+  // Check the filter type
+  if (filter->type() != FilterType::MESH) {
+    set_errmsg("Tried to set a rotation from a non-mesh filter.");
+    return OPENMC_E_INVALID_TYPE;
+  }
+
+  // Get a pointer to the filter and downcast
+  auto mesh_filter = dynamic_cast<MeshFilter*>(filter.get());
+  std::vector<double> vec_rot(rot, rot + rot_len);
+  mesh_filter->set_rotation(vec_rot);
+  return 0;
+}
+
 } // namespace openmc

src/weight_windows.cpp

@@ -933,7 +933,8 @@ void WeightWindowsGenerator::create_tally()
   for (const auto& f : model::tally_filters) {
     if (f->type() == FilterType::MESH) {
       const auto* mesh_filter = dynamic_cast<MeshFilter*>(f.get());
-      if (mesh_filter->mesh() == mesh_idx && !mesh_filter->translated()) {
+      if (mesh_filter->mesh() == mesh_idx && !mesh_filter->translated() &&
+          !mesh_filter->rotated()) {
         ww_tally->add_filter(f.get());
         found_mesh_filter = true;
         break;

tests/regression_tests/filter_rotations/__init__.py

@@ -0,0 +1,59 @@
+<?xml version='1.0' encoding='utf-8'?>
+<model>
+  <materials>
+    <material id="1" depletable="true">
+      <density value="10.0" units="g/cm3"/>
+      <nuclide name="U235" ao="1.0"/>
+    </material>
+    <material id="2">
+      <density value="1.0" units="g/cm3"/>
+      <nuclide name="Zr90" ao="1.0"/>
+    </material>
+  </materials>
+  <geometry>
+    <cell id="1" material="1" region="1 -2 3 -4 10 -9" universe="1"/>
+    <cell id="2" material="2" region="(-1 | 2 | -3 | 4) (5 -6 7 -8) 10 -9" universe="1"/>
+    <surface id="1" name="minimum x" type="x-plane" coeffs="-5.0"/>
+    <surface id="2" name="maximum x" type="x-plane" coeffs="5.0"/>
+    <surface id="3" name="minimum y" type="y-plane" coeffs="-5.0"/>
+    <surface id="4" name="maximum y" type="y-plane" coeffs="5.0"/>
+    <surface id="5" name="minimum x" type="x-plane" boundary="reflective" coeffs="-10.0"/>
+    <surface id="6" name="maximum x" type="x-plane" boundary="reflective" coeffs="10.0"/>
+    <surface id="7" name="minimum y" type="y-plane" boundary="reflective" coeffs="-10.0"/>
+    <surface id="8" name="maximum y" type="y-plane" boundary="reflective" coeffs="10.0"/>
+    <surface id="9" type="z-plane" boundary="vacuum" coeffs="10.0"/>
+    <surface id="10" type="z-plane" boundary="vacuum" coeffs="-10.0"/>
+  </geometry>
+  <settings>
+    <run_mode>eigenvalue</run_mode>
+    <particles>1000</particles>
+    <batches>5</batches>
+    <inactive>0</inactive>
+  </settings>
+  <tallies>
+    <mesh id="1">
+      <dimension>3 4 5</dimension>
+      <lower_left>-9 -9 -9</lower_left>
+      <upper_right>9 9 9</upper_right>
+    </mesh>
+    <mesh id="2">
+      <dimension>3 4 5</dimension>
+      <lower_left>-9 -9 -9</lower_left>
+      <upper_right>9 9 9</upper_right>
+    </mesh>
+    <filter id="1" type="mesh">
+      <bins>1</bins>
+    </filter>
+    <filter id="2" type="mesh" rotation="0 0 10">
+      <bins>2</bins>
+    </filter>
+    <tally id="1">
+      <filters>1</filters>
+      <scores>total</scores>
+    </tally>
+    <tally id="2">
+      <filters>2</filters>
+      <scores>total</scores>
+    </tally>
+  </tallies>
+</model>