Merge pull request #1380 from paulromano/rotation-matrix

Support for setting rotation matrix directly

include/openmc/cell.h

@@ -179,10 +179,10 @@ public:
 
   //! \brief Rotational tranfsormation of the filled universe.
   //
-  //! The vector is empty if there is no rotation.  Otherwise, the first three
-  //! values are the rotation angles respectively about the x-, y-, and z-, axes
-  //! in degrees.  The next 9 values give the rotation matrix in row-major
-  //! order.
+  //! The vector is empty if there is no rotation. Otherwise, the first 9 values
+  //! give the rotation matrix in row-major order. When the user specifies
+  //! rotation angles about the x-, y- and z- axes in degrees, these values are
+  //! also present at the end of the vector, making it of length 12.
   std::vector<double> rotation_;
 
   std::vector<int32_t> offset_;  //!< Distribcell offset table

openmc/cell.py

@@ -63,6 +63,10 @@ class Cell(IDManagerMixin):
            \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.
     rotation_matrix : numpy.ndarray
         The rotation matrix defined by the angles specified in the
         :attr:`Cell.rotation` property.
@@ -227,21 +231,24 @@ class Cell(IDManagerMixin):
 
     @rotation.setter
     def rotation(self, rotation):
-        cv.check_type('cell rotation', rotation, Iterable, Real)
         cv.check_length('cell rotation', rotation, 3)
         self._rotation = np.asarray(rotation)
 
         # Save rotation matrix -- the reason we do this instead of having it be
         # automatically calculated when the rotation_matrix property is accessed
         # is so that plotting on a rotated geometry can be done faster.
-        phi, theta, psi = self.rotation*(-pi/180.)
-        c3, s3 = cos(phi), sin(phi)
-        c2, s2 = cos(theta), sin(theta)
-        c1, s1 = cos(psi), sin(psi)
-        self._rotation_matrix = np.array([
-            [c1*c2, c1*s2*s3 - c3*s1, s1*s3 + c1*c3*s2],
-            [c2*s1, c1*c3 + s1*s2*s3, c3*s1*s2 - c1*s3],
-            [-s2, c2*s3, c2*c3]])
+        if self._rotation.ndim == 2:
+            # User specified rotation matrix directly
+            self._rotation_matrix = self._rotation
+        else:
+            phi, theta, psi = self.rotation*(-pi/180.)
+            c3, s3 = cos(phi), sin(phi)
+            c2, s2 = cos(theta), sin(theta)
+            c1, s1 = cos(psi), sin(psi)
+            self._rotation_matrix = np.array([
+                [c1*c2, c1*s2*s3 - c3*s1, s1*s3 + c1*c3*s2],
+                [c2*s1, c1*c3 + s1*s2*s3, c3*s1*s2 - c1*s3],
+                [-s2, c2*s3, c2*c3]])
 
     @translation.setter
     def translation(self, translation):
@@ -516,7 +523,7 @@ class Cell(IDManagerMixin):
             element.set("translation", ' '.join(map(str, self.translation)))
 
         if self.rotation is not None:
-            element.set("rotation", ' '.join(map(str, self.rotation)))
+            element.set("rotation", ' '.join(map(str, self.rotation.ravel())))
 
         return element
 

openmc/summary.py

@@ -154,8 +154,9 @@ class Summary(object):
 
                 if 'rotation' in group:
                     rotation = group['rotation'][()]
-                    rotation = np.asarray(rotation, dtype=np.int)
-                    cell._rotation = rotation
+                    if rotation.size == 9:
+                        rotation.shape = (3, 3)
+                    cell.rotation = rotation
 
             elif fill_type == 'material':
                 cell.temperature = group['temperature'][()]

src/cell.cpp

@@ -1,8 +1,10 @@
 
 #include "openmc/cell.h"
 
+#include <algorithm>
 #include <cctype>
 #include <cmath>
+#include <iterator>
 #include <sstream>
 #include <set>
 #include <string>
@@ -438,35 +440,39 @@ CSGCell::CSGCell(pugi::xml_node cell_node)
     }
 
     auto rot {get_node_array<double>(cell_node, "rotation")};
-    if (rot.size() != 3) {
+    if (rot.size() != 3 && rot.size() != 9) {
       std::stringstream err_msg;
       err_msg << "Non-3D rotation vector applied to cell " << id_;
       fatal_error(err_msg);
     }
 
-    // Store the rotation angles.
-    rotation_.reserve(12);
-    rotation_.push_back(rot[0]);
-    rotation_.push_back(rot[1]);
-    rotation_.push_back(rot[2]);
-
     // Compute and store the rotation matrix.
-    auto phi = -rot[0] * PI / 180.0;
-    auto theta = -rot[1] * PI / 180.0;
-    auto 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));
+    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_));
+    }
   }
 }
 
@@ -578,8 +584,12 @@ CSGCell::to_hdf5(hid_t cell_group) const
       write_dataset(group, "translation", translation_);
     }
     if (!rotation_.empty()) {
-      std::array<double, 3> rot {rotation_[0], rotation_[1], rotation_[2]};
-      write_dataset(group, "rotation", rot);
+      if (rotation_.size() == 12) {
+        std::array<double, 3> rot {rotation_[9], rotation_[10], rotation_[11]};
+        write_dataset(group, "rotation", rot);
+      } else {
+        write_dataset(group, "rotation", rotation_);
+      }
     }
 
   } else if (type_ == FILL_LATTICE) {

src/position.cpp

@@ -88,9 +88,9 @@ Position
 Position::rotate(const std::vector<double>& rotation) const
 {
   return {
+    x*rotation[0] + y*rotation[1] + z*rotation[2],
     x*rotation[3] + y*rotation[4] + z*rotation[5],
-    x*rotation[6] + y*rotation[7] + z*rotation[8],
-    x*rotation[9] + y*rotation[10] + z*rotation[11]
+    x*rotation[6] + y*rotation[7] + z*rotation[8]
   };
 }
 

tests/unit_tests/test_geometry.py

@@ -280,3 +280,24 @@ def test_from_xml(run_in_tmpdir, mixed_lattice_model):
     ll, ur = geom.bounding_box
     assert ll == pytest.approx((-6.0, -6.0, -np.inf))
     assert ur == pytest.approx((6.0, 6.0, np.inf))
+
+
+def test_rotation_matrix():
+    """Test ability to set a rotation matrix directly"""
+    y = openmc.YPlane()
+    cyl1 = openmc.ZCylinder(r=1.0)
+    cyl2 = openmc.ZCylinder(r=2.0, boundary_type='vacuum')
+
+    # Create a universe and then reflect in the y-direction
+    c1 = openmc.Cell(region=-cyl1 & +y)
+    c2 = openmc.Cell(region=+cyl1 & +y)
+    c3 = openmc.Cell(region=-y)
+    univ = openmc.Universe(cells=[c1, c2, c3])
+    c = openmc.Cell(fill=univ, region=-cyl2)
+    c.rotation = [[1, 0, 0], [0, -1, 0], [0, 0, 1]]
+    assert np.all(c.rotation_matrix == c.rotation)
+    geom = openmc.Geometry([c])
+
+    assert geom.find((0.0, 0.5, 0.0))[-1] == c3
+    assert geom.find((0.0, -0.5, 0.0))[-1] == c1
+    assert geom.find((0.0, -1.5, 0.0))[-1] == c2