add unit tests and replace resolve_poly2

space/models/planetary.py

@@ -386,12 +386,20 @@ def mp_liu2015(theta, phi, **kwargs):
 
 
 
-_models = {"mp_shue": mp_shue1998,
+_models = {"mp_shue1998": mp_shue1998,
            "mp_formisano1979": mp_formisano1979,
+           "mp_liu2015" : mp_liu2015,
            "bs_formisano1979": bs_formisano1979,
-           "bs_jerab": bs_Jerab2005}
+           "bs_jerab2005": bs_Jerab2005}
 
 
+def available(model):
+    if model == "magnetopause":
+        return tuple([m for m in _models.keys() if m.startswith("mp_")])
+    elif model == "bow_shock":
+        return tuple([m for m in _models.keys() if m.startswith("bs_")])
+    raise ValueError("invalid model type")
+
 
 def _interest_points(model, **kwargs):
     dup = kwargs.copy()
@@ -408,7 +416,7 @@ def _parabolic_approx(theta, phi, x, xf, **kwargs):
     a = np.sin(theta) ** 2
     b = 4 * K * np.cos(theta)
     c = -4 * K * x
-    r = resolve_poly2(a, b, c)[0]
+    r = resolve_poly2(a, b, c, 0)
     return coords.BaseChoice(kwargs.get("base", "cartesian"), r, theta, phi)
 
 
@@ -425,9 +433,18 @@ def wrapper(self, theta, phi, **kwargs):
 
 class Magnetosheath:
     def __init__(self, **kwargs):
-        self._magnetopause = _models[kwargs.get("magnetopause", "shue")]
-        self._bow_shock = _models[kwargs.get("bow_shock", "jerab")]
 
+        kwargs["magnetopause"] = kwargs.get("magnetopause", "mp_shue1998")
+        kwargs["bow_shock"] = kwargs.get("bow_shock", "bs_jerab2005")
+
+        if not kwargs["magnetopause"].startswith("mp_") \
+            or not kwargs["bow_shock"].startswith("bs_"):
+            raise ValueError("invalid model name")
+
+        self._magnetopause = _models[kwargs["magnetopause"]]
+        self._bow_shock = _models[kwargs["bow_shock"]]
+        self.model_magnetopause = kwargs["magnetopause"]
+        self.model_bow_shock = kwargs["bow_shock"]
 
     @check_parabconfoc
     def magnetopause(self, theta, phi, **kwargs):

space/smath.py

@@ -4,28 +4,11 @@
 def norm(u, v, w):
     return np.sqrt(u**2 + v**2 + w**2)
 
-
-def resolve_poly2(a, b, c):
-    if isinstance(a, (np.ndarray | pd.Series)):
-        r1 = np.zeros_like(a)
-        r2 = np.zeros_like(a)
-        a_null = np.where(np.abs(a) < 1e-6)[0]
-
-        delta = b ** 2 - 4 * a * c
-        np.testing.assert_array_less(-delta, 0)
-
-        r1, r2 = (-b + np.sqrt(delta)) / (2 * a), (-b + np.sqrt(delta)) / (2 * a)
-        if isinstance(c, (np.ndarray | pd.Series)):
-            c = c[a_null]
-        if isinstance(b, (np.ndarray | pd.Series)):
-            b = b[a_null]
-        r1[a_null] = r2[a_null] = -c / b
-    else:
-        delta = b ** 2 - 4 * a * c
-        np.testing.assert_array_less(-delta, 0)
-        if np.abs(a) < 1e-6:
-            r1 = r2 = -c / b
-        else :
-            r1, r2 = (-b + np.sqrt(delta)) / (2 * a), (-b - np.sqrt(delta)) / (2 * a)
-
-    return r1, r2
+def resolve_poly2(a, b, c, roots=None):
+    def fun(a, b, c):
+        if roots is None:
+            return np.roots([a, b, c])
+        return np.roots([a,b,c])[roots]
+
+    vfun = np.vectorize(fun, otypes=(np.ndarray,))
+    return vfun(a, b, c)

tests/test_space.py

@@ -4,18 +4,71 @@
 
 
 import unittest
-
+from ddt import data, unpack, ddt
+from space.models.planetary import Magnetosheath
+import numpy as np
 from space import space
 
-
-class TestSpace(unittest.TestCase):
+@ddt
+class TestMagnetosheath(unittest.TestCase):
     """Tests for `space` package."""
 
     def setUp(self):
         """Set up test fixtures, if any."""
+        self.typical_msh = Magnetosheath(magnetopause="mp_formisano1979", bow_shock="bs_formisano1979")
+        theta = np.arange(0, np.pi/2, 0.1*np.pi)
+        self.typical_mp_sph, self.typical_bs_sph  = self.typical_msh.boundaries(theta, 0, base = "spherical")
+        self.typical_mp_cart, self.typical_bs_cart = self.typical_msh.boundaries(theta, 0, base = "cartesian")
 
     def tearDown(self):
         """Tear down test fixtures, if any."""
 
-    def test_000_something(self):
+    def test_magnetosheath_can_be_constructed(self):
+        """Test something."""
+        self.assertIsNotNone(self.typical_msh)
+
+
+
+    @data(("bs_shue1998", "mp_formisano1979"), #both wrong
+          ("mp_shue1998", "mp_formisano1979")) # bow shock wrong
+    @unpack
+    def test_msh_invalid_boundaries_raise(self, mp, bs):
         """Test something."""
+        self.assertRaises(ValueError, Magnetosheath, magnetopause=mp,
+                                        bow_shock=bs)
+
+
+    @data("mp_formisano1979", "mp_shue1998")
+    def test_mp_nose_are_in_expected_dayside_region(self, mp):
+        msh = Magnetosheath(magnetopause=mp, bow_shock="bs_formisano1979")
+        x,y,z = msh.magnetopause(0,0)
+        self.assertGreater(x, 5)
+        self.assertLess(x, 15)
+
+
+    @data("bs_formisano1979", "bs_jerab2005")
+    def test_mp_nose_are_in_expected_dayside_region(self, bs):
+        msh = Magnetosheath(magnetopause="mp_formisano1979", bow_shock=bs)
+        x,y,z = msh.bow_shock(0,0)
+        self.assertGreater(x, 7)
+        self.assertLess(x, 30)
+
+
+    def test_spherical_and_cartesian_are_consistent(self):
+        for model_sph, model_cart in zip((self.typical_mp_sph, self.typical_bs_sph),
+                                         (self.typical_mp_cart, self.typical_bs_cart)):
+            rfromcart = np.sqrt(model_cart[0]**2 + model_cart[1]**2+ model_cart[2]**2)
+            np.testing.assert_allclose(rfromcart, model_sph[0], atol=1e-12)
+
+
+    def test_msh_return_mp_and_bs(self):
+        rmp = self.typical_mp_sph[0]
+        rbs = self.typical_bs_sph[0]
+        np.testing.assert_array_less(rmp, rbs)
+
+
+    def test_non_regression(self):
+        pass
+        # read expected values in pickle file
+        # run code for same params
+        # check np.testing.allclose(actual, expected)