Added c-api methods for rotation matrix set/get

include/openmc/capi.h

@@ -15,6 +15,7 @@ extern "C" {
   int openmc_cell_get_id(int32_t index, int32_t* id);
   int openmc_cell_get_temperature(int32_t index, const int32_t* instance, double* T);
   int openmc_cell_get_translation(int32_t index, double xyz[]);
+  int openmc_cell_get_rotation(int32_t index, double rot[], size_t* n);
   int openmc_cell_get_name(int32_t index, const char** name);
   int openmc_cell_get_num_instances(int32_t index, int32_t* num_instances);
   int openmc_cell_set_name(int32_t index, const char* name);
@@ -22,6 +23,7 @@ extern "C" {
   int openmc_cell_set_id(int32_t index, int32_t id);
   int openmc_cell_set_temperature(int32_t index, double T, const int32_t* instance, bool set_contained = false);
   int openmc_cell_set_translation(int32_t index, const double xyz[]);
+  int openmc_cell_set_rotation(int32_t index, const double rot[], const size_t rot_len);
   int openmc_energy_filter_get_bins(int32_t index, const double** energies, size_t* n);
   int openmc_energy_filter_set_bins(int32_t index, size_t n, const double* energies);
   int openmc_energyfunc_filter_get_energy(int32_t index, size_t* n, const double** energy);

include/openmc/cell.h

@@ -159,6 +159,10 @@ public:
   //!   temperatures.
   void set_temperature(double T, int32_t instance = -1, bool set_contained = false);
 
+  //! Set the rotation matrix of a cell instance
+  //! \param[in] rot The rotation matrix of length 3 or 9
+  void set_rotation(const vector<double>& rot);
+
   //! Get the name of a cell
   //! \return Cell name
   const std::string& name() const { return name_; };

openmc/lib/cell.py

@@ -1,7 +1,7 @@
 import sys
 
 from collections.abc import Mapping, Iterable
-from ctypes import c_int, c_int32, c_double, c_char_p, POINTER, c_bool
+from ctypes import c_int, c_int32, c_double, c_char_p, POINTER, c_bool, c_size_t
 from weakref import WeakValueDictionary
 
 import numpy as np
@@ -38,6 +38,10 @@ _dll.openmc_cell_get_name.errcheck = _error_handler
 _dll.openmc_cell_get_translation.argtypes = [c_int32, POINTER(c_double)]
 _dll.openmc_cell_get_translation.restype = c_int
 _dll.openmc_cell_get_translation.errcheck = _error_handler
+_dll.openmc_cell_get_rotation.argtypes = [c_int32, POINTER(c_double),
+    POINTER(c_size_t)]
+_dll.openmc_cell_get_rotation.restype = c_int
+_dll.openmc_cell_get_rotation.errcheck = _error_handler
 _dll.openmc_cell_set_name.argtypes = [c_int32, c_char_p]
 _dll.openmc_cell_set_name.restype = c_int
 _dll.openmc_cell_set_name.errcheck = _error_handler
@@ -55,6 +59,10 @@ _dll.openmc_cell_set_temperature.errcheck = _error_handler
 _dll.openmc_cell_set_translation.argtypes = [c_int32, POINTER(c_double)]
 _dll.openmc_cell_set_translation.restype = c_int
 _dll.openmc_cell_set_translation.errcheck = _error_handler
+_dll.openmc_cell_set_rotation.argtypes = [
+    c_int32, POINTER(c_double), c_size_t]
+_dll.openmc_cell_set_rotation.restype = c_int
+_dll.openmc_cell_set_rotation.errcheck = _error_handler
 _dll.openmc_get_cell_index.argtypes = [c_int32, POINTER(c_int32)]
 _dll.openmc_get_cell_index.restype = c_int
 _dll.openmc_get_cell_index.errcheck = _error_handler
@@ -242,6 +250,44 @@ class Cell(_FortranObjectWithID):
         _dll.openmc_cell_set_translation(
             self._index, translation.ctypes.data_as(POINTER(c_double)))
 
+    def get_rotation(self):
+        """Get the rotation matrix of a cell
+
+        """
+
+        rotation = np.zeros(12)
+        rot_size = c_size_t()
+
+        _dll.openmc_cell_get_rotation(
+            self._index, rotation.ctypes.data_as(POINTER(c_double)), rot_size)
+        rot_size = rot_size.value
+
+        if rot_size == 0:
+            return None
+        elif rot_size == 9:
+            return rotation[:n]
+        elif rot_size == 12:
+            return rotation[9:]
+        else:
+            raise ValueError(
+                'Invalid size of rotation matrix: {}'.format(rot_size))
+
+    def set_rotation(self, rotation):
+        """Set the rotation matrix of a cell
+
+        Parameters
+        ----------
+        rotation : numpy.ndarray
+            The rotation angles (if length is 3) or the 3x3 rotation matrix
+
+        """
+
+        flat_rotation = rotation.flatten()
+
+        _dll.openmc_cell_set_rotation(
+            self._index, flat_rotation.ctypes.data_as(POINTER(c_double)),
+            c_size_t(len(flat_rotation)))
+
     @property
     def bounding_box(self):
         inf = sys.float_info.max

src/cell.cpp

@@ -248,6 +248,48 @@ BoundingBox Universe::bounding_box() const {
 // Cell implementation
 //==============================================================================
 
+void
+Cell::set_rotation(const vector<double>& rot) {
+  if (fill_ == C_NONE) {
+    fatal_error(fmt::format("Cannot apply a rotation to cell {}"
+                            " because it is not filled with another universe",
+      id_));
+  }
+
+  if (rot.size() != 3 && rot.size() != 9) {
+    fatal_error(fmt::format("Non-3D rotation vector applied to cell {}", id_));
+  }
+
+  // Compute and store the rotation matrix.
+  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_));
+  }
+}
+
 double
 Cell::temperature(int32_t instance) const
 {
@@ -572,44 +614,8 @@ CSGCell::CSGCell(pugi::xml_node cell_node)
 
   // Read the rotation transform.
   if (check_for_node(cell_node, "rotation")) {
-    if (fill_ == C_NONE) {
-      fatal_error(fmt::format("Cannot apply a rotation to cell {}"
-        " because it is not filled with another universe", id_));
-    }
-
     auto rot {get_node_array<double>(cell_node, "rotation")};
-    if (rot.size() != 3 && rot.size() != 9) {
-      fatal_error(fmt::format(
-        "Non-3D rotation vector applied to cell {}", id_));
-    }
-
-    // Compute and store the rotation matrix.
-    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_));
-    }
+    set_rotation(rot);
   }
 }
 
@@ -1332,6 +1338,40 @@ extern "C" int openmc_cell_set_translation(int32_t index, const double xyz[])
   }
 }
 
+//! Return the rotation matrix of a cell
+extern "C" int openmc_cell_get_rotation(int32_t index, double rot[], size_t* n)
+{
+  if (index >= 0 && index < model::cells.size()) {
+    auto& cell = openmc::model::cells[index];
+    *n = cell->rotation_.size();
+    std::memcpy(rot, cell->rotation_.data(), *n * sizeof(cell->rotation_[0]));
+    return 0;
+  } else {
+    set_errmsg("Index in cells array is out of bounds.");
+    return OPENMC_E_OUT_OF_BOUNDS;
+  }
+}
+
+//! Set the flattened rotation matrix of a cell
+extern "C" int openmc_cell_set_rotation(int32_t index, const double rot[],
+    const size_t rot_len)
+{
+  if (index >= 0 && index < model::cells.size()) {
+    if (model::cells[index]->fill_ == C_NONE) {
+      set_errmsg(fmt::format("Cannot apply a rotation to cell {}"
+                             " because it is not filled with another universe",
+        index));
+      return OPENMC_E_GEOMETRY;
+    }
+    std::vector<double> vec_rot(rot, rot + rot_len);
+    model::cells[index]->set_rotation(vec_rot);
+    return 0;
+  } else {
+    set_errmsg("Index in cells array is out of bounds.");
+    return OPENMC_E_OUT_OF_BOUNDS;
+  }
+}
+
 //! Get the number of instances of the requested cell
 extern "C" int
 openmc_cell_get_num_instances(int32_t index, int32_t* num_instances)

tests/unit_tests/test_lib.py

@@ -727,3 +727,18 @@ def test_cell_translation(pincell_model_w_univ, mpi_intracomm):
     # This time we *can* set it
     cell.set_translation(np.array([1., 0., -1.]))
     assert cell.get_translation() == pytest.approx([1., 0., -1.])
+
+
+def test_cell_rotation(pincell_model_w_univ):
+    # Cell 1 is filled with a material so we cannot rotate it, but we can get
+    # its rotation matrix (which will be the identity matrix)
+    cell = openmc.lib.cells[1]
+    assert cell.get_rotation() is None
+    with pytest.raises(exc.GeometryError, match='not filled with'):
+        cell.set_rotation(np.array([180., 0., 0.]))
+
+    # Now repeat with Cell 2 and we will be allowed to do it
+    cell = openmc.lib.cells[2]
+    assert cell.get_rotation() is None
+    cell.set_rotation(np.array([180., 0., 0.]))
+    assert cell.get_rotation() == pytest.approx([180., 0., 0.])