From 7d69ff6de18476ef7194abc6264c967f4d22f389 Mon Sep 17 00:00:00 2001 From: Ethan Peterson Date: Wed, 26 Feb 2020 12:33:51 -0500 Subject: [PATCH] applied suggestions from code review --- openmc/surface.py | 48 ++++++++++++++++++++++---------- tests/unit_tests/test_surface.py | 8 +++--- 2 files changed, 38 insertions(+), 18 deletions(-) diff --git a/openmc/surface.py b/openmc/surface.py index 03be4dbb00..562b4c859d 100644 --- a/openmc/surface.py +++ b/openmc/surface.py @@ -51,15 +51,17 @@ def get_rotation_matrix(rotation, order='xyz'): check_type('surface rotation', rotation, Iterable, Real) check_length('surface rotation', rotation, 3) - phi, theta, psi = np.array(rotation)*(np.pi/180.) - cx, sx = np.cos(phi), np.sin(phi) - cy, sy = np.cos(theta), np.sin(theta) - cz, sz = np.cos(psi), np.sin(psi) - R = {} - R['x'] = np.array([[1., 0., 0.], [0., cx, -sx], [0., sx, cx]]) - R['y'] = np.array([[cy, 0., sy], [0., 1., 0.], [-sy, 0., cy]]) - R['z'] = np.array([[cz, -sz, 0.], [sz, cz, 0.], [0., 0., 1.]]) - R1, R2, R3 = tuple(R[xi] for xi in order) + phi, theta, psi = np.array(rotation)*(math.pi/180.) + cx, sx = math.cos(phi), math.sin(phi) + cy, sy = math.cos(theta), math.sin(theta) + cz, sz = math.cos(psi), math.sin(psi) + R = { + 'x': np.array([[1., 0., 0.], [0., cx, -sx], [0., sx, cx]]), + 'y': np.array([[cy, 0., sy], [0., 1., 0.], [-sy, 0., cy]]), + 'z': np.array([[cz, -sz, 0.], [sz, cz, 0.], [0., 0., 1.]]), + } + + R1, R2, R3 = (R[xi] for xi in order) return R3 @ R2 @ R1 @@ -522,7 +524,7 @@ class PlaneMixin(metaclass=ABCMeta): return self a, b, c, d = self._get_base_coeffs() - d = d + np.array([a, b, c]) @ np.array(vector) + d = d + np.dot([a, b, c], vector) surf = self if inplace else self.clone() @@ -532,12 +534,22 @@ class PlaneMixin(metaclass=ABCMeta): def rotate(self, rotation, pivot=(0., 0., 0.), order='xyz', inplace=False): pivot = np.asarray(pivot) + rotation = np.asarray(rotation) + + # Allow rotaiton matrix to be passed in directly, otherwise build it + if np.rank(rotation) == 2: + check_type('surface rotation', rotation, Iterable, Real) + check_length('surface rotation', rotation.ravel(), 9) + Rmat = rotation + else: + Rmat = get_rotation_matrix(rotation, order=order) + + # Translate surface to pivot surf = self.translate(-pivot, inplace=inplace) - Rmat = get_rotation_matrix(rotation, order=order) a, b, c, d = surf._get_base_coeffs() # Compute new rotated coefficients a, b, c - a, b, c = Rmat @ np.array([a, b, c]) + a, b, c = Rmat @ [a, b, c] kwargs = {'boundary_type': surf.boundary_type, 'name': surf.name} if inplace: @@ -1089,14 +1101,22 @@ class QuadricMixin(metaclass=ABCMeta): def rotate(self, rotation, pivot=(0., 0., 0.), order='xyz', inplace=False): # Get pivot and rotation matrix pivot = np.asarray(pivot) - Rmat = get_rotation_matrix(rotation, order=order) + rotation = np.asarray(rotation) + + # Allow rotaiton matrix to be passed in directly, otherwise build it + if np.rank(rotation) == 2: + check_type('surface rotation', rotation, Iterable, Real) + check_length('surface rotation', rotation.ravel(), 9) + Rmat = rotation + else: + Rmat = get_rotation_matrix(rotation, order=order) # Translate surface to the pivot point tsurf = self.translate(-pivot, inplace=inplace) # If the surface is already generalized just clone it if type(tsurf) is tsurf._virtual_base: - surf = tsurf if inplace else tsurf.clone() + surf = tsurf if inplace else tsurf.clone() else: base_cls = type(tsurf)._virtual_base # Copy necessary surface attributes to new kwargs dictionary diff --git a/tests/unit_tests/test_surface.py b/tests/unit_tests/test_surface.py index 4b0da6f4b0..da764e3bbc 100644 --- a/tests/unit_tests/test_surface.py +++ b/tests/unit_tests/test_surface.py @@ -38,13 +38,13 @@ def test_plane(): assert (st.a, st.b, st.c, st.d) == (s.a, s.b, s.c, 4) # rotate method - yp = openmc.YPlane(np.abs(s.d)/np.sqrt(s.a**2 + s.b**2 + s.c**2)) - psi = np.rad2deg(np.arctan2(1, 2)) - phi = np.rad2deg(np.arctan2(1, np.sqrt(5))) + yp = openmc.YPlane(abs(s.d)/math.sqrt(s.a**2 + s.b**2 + s.c**2)) + psi = math.degrees(np.arctan2(1, 2)) + phi = math.degrees(np.arctan2(1, np.sqrt(5))) sr = s.rotate((phi, 0., psi), order='zyx') assert yp.normalize() == pytest.approx(sr.normalize()) # test rotation ordering - phi = np.rad2deg(np.arctan2(1, np.sqrt(2))) + phi = math.degrees(np.arctan2(1, np.sqrt(2))) sr = s.rotate((0., -45., phi), order='xyz') assert yp.normalize() == pytest.approx(sr.normalize())