From e3af035095bc7783b127d97df956f06edef21f70 Mon Sep 17 00:00:00 2001
From: Ethan Peterson <ethan.peterson@mit.edu>
Date: Mon, 24 Feb 2020 15:30:28 -0500
Subject: [PATCH] finished unit tests and fixed wrong math

---
 openmc/surface.py                |  44 ++++++-----
 tests/unit_tests/test_surface.py | 121 +++++++++++++++++++++++++++++++
 2 files changed, 146 insertions(+), 19 deletions(-)

diff --git a/openmc/surface.py b/openmc/surface.py
index cc92a810b3..e4c11692a1 100644
--- a/openmc/surface.py
+++ b/openmc/surface.py
@@ -887,6 +887,15 @@ class ZPlane(PlaneMixin, Surface):
 class QuadricMixin(metaclass=ABCMeta):
     """A Mixin class implementing common functionality for quadric surfaces"""
 
+    @property
+    def _origin(self):
+        return np.array((self.x0, self.y0, self.z0))
+
+    @property
+    def _axis(self):
+        axis = np.array((self.dx, self.dy, self.dz))
+        return axis / np.linalg.norm(axis)
+
     def get_Abc(self, coeffs=None):
         """Compute matrix, vector, and scalar coefficients for this surface or
         for a specified set of coefficients.
@@ -1135,8 +1144,8 @@ class Cylinder(QuadricMixin, Surface):
 
     def _get_base_coeffs(self):
         # Get x, y, z coordinates of two points
-        x1, y1, z1 = self.x0, self.y0, self.z0
-        x2, y2, z2 = x1 + self.dx, y1 + self.dy, z1 + self.dz
+        x1, y1, z1 = self._origin
+        x2, y2, z2 = self._origin + self._axis
         r = self.r 
 
         # Define intermediate terms
@@ -1868,24 +1877,21 @@ class Cone(QuadricMixin, Surface):
         # The argument r2 for cones is actually tan^2(theta) so that
         # cos^2(theta) = 1 / (1 + r2)
 
-        x0, y0, z0, r2 = self.x0, self.y0, self.z0, self.r2
-        dx, dy, dz = self.dx, self.dy, self.dz
-        dnorm = dx*dx + dy*dy + dz*dz
-        dx /= dnorm
-        dy /= dnorm
-        dz /= dnorm
-        cos2 = 1 / (1 + r2)
+        x0, y0, z0 = self._origin
+        dx, dy, dz = self._axis
+        cos2 = 1 / (1 + self.r2)
 
-        a = dx*dx - cos2
-        b = dy*dy - cos2
-        c = dz*dz - cos2
-        d = 2*dx*dy
-        e = 2*dy*dz
-        f = 2*dx*dz
-        g = -2*(dx*dx*x0 + dx*dy*y0 + dx*dz*z0 - cos2)
-        h = -2*(dy*dy*y0 + dx*dy*x0 + dy*dz*z0 - cos2)
-        j = -2*(dz*dz*z0 + dx*dz*x0 + dy*dz*y0 - cos2)
-        k = (dx*x0 + dy*y0 + dz*z0)**2 - cos2*(x0*x0 + y0*y0 + z0*z0)
+        a = cos2 - dx*dx
+        b = cos2 - dy*dy
+        c = cos2 - dz*dz
+        d = -2*dx*dy
+        e = -2*dy*dz
+        f = -2*dx*dz
+        g = 2*(dx*(dy*y0 + dz*z0) - a*x0)
+        h = 2*(dy*(dx*x0 + dz*z0) - b*y0)
+        j = 2*(dz*(dx*x0 + dy*y0) - c*z0)
+        k = a*x0*x0 + b*y0*y0 + c*z0*z0 - 2*(dx*dy*x0*y0 + dy*dz*y0*z0 +
+                                             dx*dz*x0*z0)
 
         return (a, b, c, d, e, f, g, h, j, k)
 
diff --git a/tests/unit_tests/test_surface.py b/tests/unit_tests/test_surface.py
index 86e01b6af3..b6c7398778 100644
--- a/tests/unit_tests/test_surface.py
+++ b/tests/unit_tests/test_surface.py
@@ -2,6 +2,7 @@ from functools import partial
 from random import uniform, seed
 
 import numpy as np
+import math
 import openmc
 import pytest
 
@@ -139,6 +140,63 @@ def test_zplane():
     repr(s)
 
 
+def test_cylinder():
+    x0, y0, z0, r = 2, 3, 4, 2
+    dx, dy, dz = 1, -1, 1
+    s = openmc.Cylinder(x0=x0, y0=y0, z0=z0, dx=dx, dy=dy, dz=dz, r=r)
+    assert s.x0 == 2
+    assert s.y0 == 3
+    assert s.z0 == 4
+    assert s.dx == 1
+    assert s.dy == -1
+    assert s.dz == 1
+    assert s.r == 2
+
+    # Check bounding box
+    assert_infinite_bb(s)
+
+    # evaluate method
+    # |(p - p1) ⨯ (p - p2)|^2 / |p2 - p1|^2 - r^2
+    # point inside
+    p = s._origin + 5*s._axis
+    p1 = s._origin
+    p2 = p1 + s._axis
+    c1 = np.linalg.norm(np.cross(p - p1, p - p2)) / np.linalg.norm(p2 - p1)
+    val = c1*c1 - s.r*s.r
+    assert s.evaluate(p) < 0
+    assert s.evaluate(p) == pytest.approx(val)
+    # point outside
+    p = np.array((4., 0., 2.5))
+    p1 = s._origin
+    p2 = p1 + s._axis
+    c1 = np.linalg.norm(np.cross(p - p1, p - p2)) / np.linalg.norm(p2 - p1)
+    val = c1*c1 - s.r*s.r
+    assert s.evaluate(p) > 0
+    assert s.evaluate(p) == pytest.approx(val)
+    # point on cylinder
+    perp = np.array((1, -2, 1))*(1 / s._axis)
+    p = s._origin + s.r*perp / np.linalg.norm(perp)
+    p1 = s._origin
+    p2 = p1 + s._axis
+    c1 = np.linalg.norm(np.cross(p - p1, p - p2)) / np.linalg.norm(p2 - p1)
+    val = c1*c1 - s.r*s.r
+    assert 0 == pytest.approx(s.evaluate(p))
+    assert val == pytest.approx(s.evaluate(p))
+
+    # translate method
+    st = s.translate((1.0, 1.0, 1.0))
+    assert st.x0 == s.x0 + 1
+    assert st.y0 == s.y0 + 1
+    assert st.z0 == s.z0 + 1
+    assert st.dx == s.dx
+    assert st.dy == s.dy
+    assert st.dz == s.dz
+    assert st.r == s.r
+
+    # Make sure repr works
+    repr(s)
+
+
 def test_xcylinder():
     y, z, r = 3, 5, 2
     s = openmc.XCylinder(y0=y, z0=z, r=r)
@@ -284,6 +342,69 @@ def cone_common(apex, r2, cls):
     repr(s)
 
 
+def test_cone():
+    x0, y0, z0, r2 = 2, 3, 4, 4
+    dx, dy, dz = 1, -1, 1
+    s = openmc.Cone(x0=x0, y0=y0, z0=z0, dx=dx, dy=dy, dz=dz, r2=r2)
+    assert s.x0 == 2
+    assert s.y0 == 3
+    assert s.z0 == 4
+    assert s.dx == 1
+    assert s.dy == -1
+    assert s.dz == 1
+    assert s.r2 == 4
+
+    # Check bounding box
+    assert_infinite_bb(s)
+
+    # evaluate method
+    # cos^2(theta) * ((p - p1))**2 - (d @ (p - p1))^2 
+    # The argument r2 for cones is actually tan^2(theta) so that
+    # cos^2(theta) = 1 / (1 + r2)
+    #
+    # This makes the evaluation equation shown below where p is the evaluation
+    # point (x, y, z) p1 is the apex (origin) of the cone and r2 is related to
+    # the aperature of the cone as described above
+    # (p - p1) @ (p - p1) / (1 + r2) - (d @ (p - p1))^2
+    # point inside
+    p1 = s._origin
+    d = s._axis
+    p = p1 + 5*d
+    val = np.sum((p - p1)**2) / (1 + s.r2) - np.sum((d @ (p - p1))**2)
+    assert s.evaluate(p) < 0
+    assert s.evaluate(p) == pytest.approx(val)
+    # point outside
+    p1 = s._origin
+    d = s._axis
+    perp = np.array((1, -2, 1))*(1 / d)
+    p = p1 + 3.2*perp
+    val = np.sum((p - p1)**2) / (1 + s.r2) - np.sum((d @ (p - p1))**2)
+    assert s.evaluate(p) > 0
+    assert s.evaluate(p) == pytest.approx(val)
+    # point on cone
+    p1 = s._origin
+    d = s._axis
+    perp = np.array((1, -2, 1))*(1 / d)
+    perp /= np.linalg.norm(perp)
+    p = p1 + 3.2*d + 3.2*math.sqrt(s.r2)*perp
+    val = np.sum((p - p1)**2) / (1 + s.r2) - np.sum((d @ (p - p1))**2)
+    assert 0 == pytest.approx(s.evaluate(p))
+    assert val == pytest.approx(s.evaluate(p))
+
+    # translate method
+    st = s.translate((1.0, 1.0, 1.0))
+    assert st.x0 == s.x0 + 1
+    assert st.y0 == s.y0 + 1
+    assert st.z0 == s.z0 + 1
+    assert st.dx == s.dx
+    assert st.dy == s.dy
+    assert st.dz == s.dz
+    assert st.r2 == s.r2
+
+    # Make sure repr works
+    repr(s)
+
+
 def test_xcone():
     apex = (10, 0, 0)
     r2 = 4