Add TorusMixin.rotate method that can handle simple rotations

openmc/surface.py

@@ -2142,7 +2142,48 @@ class TorusMixin:
         return surf
 
     def rotate(self, rotation, pivot=(0., 0., 0.), order='xyz', inplace=False):
-        raise NotImplementedError('Torus surfaces cannot be rotated.')
+        pivot = np.asarray(pivot)
+        rotation = np.asarray(rotation, dtype=float)
+
+        # Allow rotation matrix to be passed in directly, otherwise build it
+        if rotation.ndim == 2:
+            check_length('surface rotation', rotation.ravel(), 9)
+            Rmat = rotation
+        else:
+            Rmat = get_rotation_matrix(rotation, order=order)
+
+        # Only can handle trivial rotation matrices
+        close = np.isclose
+        if not np.all(close(Rmat, -1.0) | close(Rmat, 0.0) | close(Rmat, 1.0)):
+            raise NotImplementedError('Torus surfaces cannot handle generic rotations')
+
+        # Translate surface to pivot
+        surf = self.translate(-pivot, inplace=inplace)
+
+        # Determine "center" of torus and a point above it (along main axis)
+        center = [surf.x0, surf.y0, surf.z0]
+        above_center = center.copy()
+        index = ['x-torus', 'y-torus', 'z-torus'].index(surf._type)
+        above_center[index] += 1
+
+        # Compute new rotated torus center
+        center = Rmat @ center
+
+        # Figure out which axis should be used after rotation
+        above_center = Rmat @ above_center
+        new_index = np.where(np.isclose(np.abs(above_center - center), 1.0))[0][0]
+        cls = [XTorus, YTorus, ZTorus][new_index]
+
+        # Create rotated torus
+        kwargs = {
+            'boundary_type': surf.boundary_type, 'name': surf.name,
+            'a': surf.a, 'b': surf.b, 'c': surf.c
+        }
+        if inplace:
+            kwargs['surface_id'] = surf.id
+        surf = cls(x0=center[0], y0=center[1], z0=center[2], **kwargs)
+
+        return surf.translate(pivot, inplace=inplace)
 
     def _get_base_coeffs(self):
         raise NotImplementedError

tests/unit_tests/test_surface.py

@@ -630,9 +630,15 @@ def torus_common(center, R, r1, r2, cls):
     assert st.b == s.b
     assert st.c == s.c
 
-    # can't rotate at present
+    # trivial rotations
+    for rotation in [(0., 0., 0.), (180., 0., 0.), (0., 180., 0.), (0., 0., 180.)]:
+        sr = s.rotate(rotation)
+        assert type(sr) == type(s)
+        assert (sr.a, sr.b, sr.c) == (s.a, s.b, s.c)
+
+    # can't do generic rotate at present
     with pytest.raises(NotImplementedError):
-        s.rotate((0., 1., 0.))
+        s.rotate((0., 45., 0.))
 
     # Check bounding box
     ll, ur = (+s).bounding_box
@@ -676,6 +682,12 @@ def test_xtorus():
     assert s.evaluate((x + r1 + 0.01, y, z + R)) > 0.0
     assert s.evaluate((x + r1 + 0.01, y + R, z)) > 0.0
 
+    # rotation
+    sr = s.rotate((0., 0., 90.))
+    assert isinstance(sr, openmc.YTorus)
+    sr = s.rotate((0., 90., 0.))
+    assert isinstance(sr, openmc.ZTorus)
+
 
 def test_ytorus():
     x, y, z = 2, -4, 5
@@ -703,6 +715,12 @@ def test_ytorus():
     assert s.evaluate((x, y + r1 + 0.01, z + R)) > 0.0
     assert s.evaluate((x + R, y + r1 + 0.01, z)) > 0.0
 
+    # rotation
+    sr = s.rotate((90., 0., 0.))
+    assert isinstance(sr, openmc.ZTorus)
+    sr = s.rotate((0., 0., 90.))
+    assert isinstance(sr, openmc.XTorus)
+
 
 def test_ztorus():
     x, y, z = 2, -4, 5
@@ -729,3 +747,9 @@ def test_ztorus():
     assert s.evaluate((x + R + r2 + 0.01, y, z)) > 0.0
     assert s.evaluate((x + R, y, z + r1 + 0.01)) > 0.0
     assert s.evaluate((x, y + R, z + r1 + 0.01)) > 0.0
+
+    # rotation
+    sr = s.rotate((90., 0., 0.))
+    assert isinstance(sr, openmc.YTorus)
+    sr = s.rotate((0., 90., 0.))
+    assert isinstance(sr, openmc.XTorus)