From ebb6736ba65f16afc0063427b0e4e8644db33180 Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Wed, 29 Dec 2021 10:37:54 -0500
Subject: [PATCH] Add TorusMixin.rotate method that can handle simple rotations

---
 openmc/surface.py                | 43 +++++++++++++++++++++++++++++++-
 tests/unit_tests/test_surface.py | 28 +++++++++++++++++++--
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/openmc/surface.py b/openmc/surface.py
index e00dc8962e..2f9618224d 100644
--- a/openmc/surface.py
+++ b/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
diff --git a/tests/unit_tests/test_surface.py b/tests/unit_tests/test_surface.py
index c24d6e3521..12cd8c9d20 100644
--- a/tests/unit_tests/test_surface.py
+++ b/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)