more tests

Signed-off-by: Nathaniel Starkman (@nstarman) <[email protected]>

Author: nstarman

sample_scf/base.py

@@ -19,6 +19,7 @@
 import numpy as np
 import numpy.typing as npt
 from astropy.coordinates import PhysicsSphericalRepresentation
+from galpy.potential import SCFPotential
 from scipy._lib._util import check_random_state
 from scipy.stats import rv_continuous
 
@@ -102,6 +103,14 @@ class SCFSamplerBase:
     pot : `galpy.potential.SCFPotential`
     """
 
+    def __init__(
+        self,
+        pot: SCFPotential,
+    ):
+        self._pot = pot
+
+    # /def
+
     _rsampler: rv_continuous_modrvs
     _thetasampler: rv_continuous_modrvs
     _phisampler: rv_continuous_modrvs
@@ -147,7 +156,7 @@ def cdf(
         (N, 3) ndarray
         """
         R: NDArray64 = self.rsampler.cdf(r)
-        Theta: NDArray64 = self.thetasampler.cdf(theta=theta, r=r)
+        Theta: NDArray64 = self.thetasampler.cdf(theta, r=r)
         Phi: NDArray64 = self.phisampler.cdf(phi, r=r, theta=theta)
 
         RTP: NDArray64 = np.c_[R, Theta, Phi]

sample_scf/conftest.py

@@ -75,7 +75,7 @@ def pytest_configure(config):
 # Fixtures
 
 
-@pytest.fixture(scope="session")
+@pytest.fixture(autouse=True, scope="session")
 def hernquist_scf_potential():
     """Make a SCF of a Hernquist potential."""
     Acos = np.zeros((5, 6, 6))
@@ -90,10 +90,27 @@ def hernquist_scf_potential():
 # /def
 
 
-@pytest.fixture(scope="session")
-def nfw_scf_potential():
-    """Make a SCF of a triaxial NFW potential."""
-    raise NotImplementedError("TODO")
-
-
-# /def
+# @pytest.fixture(autouse=True, scope="session")
+# def nfw_scf_potential():
+#     """Make a SCF of a triaxial NFW potential."""
+#     raise NotImplementedError("TODO")
+#
+#
+# # /def
+
+
+@pytest.fixture(
+    # autouse=True,
+    scope="session",
+    params=[
+        "hernquist_scf_potential",  # TODO! use hernquist_scf_potential
+    ],
+)
+def potentials(request):
+    if request.param == "hernquist_scf_potential":
+        Acos = np.zeros((5, 6, 6))
+        Acos_hern = Acos.copy()
+        Acos_hern[0, 0, 0] = 1
+        potential = SCFPotential(Acos=Acos_hern)
+
+    yield potential

sample_scf/core.py

@@ -19,11 +19,23 @@
 from galpy.potential import SCFPotential
 
 # LOCAL
-from .base import SCFSamplerBase
+from .base import SCFSamplerBase, rv_continuous_modrvs
+from .sample_exact import SCFSampler as SCFSamplerExact
+from .sample_intrp import SCFSampler as SCFSamplerIntrp
 
 __all__: T.List[str] = ["SCFSampler"]
 
 
+##############################################################################
+# Parameters
+
+
+class MethodsMapping(T.TypedDict):
+    r: rv_continuous_modrvs
+    theta: rv_continuous_modrvs
+    phi: rv_continuous_modrvs
+
+
 ##############################################################################
 # CODE
 ##############################################################################
@@ -85,15 +97,32 @@ class SCFSampler(SCFSamplerBase):  # metaclass=SCFSamplerSwitch
     def __init__(
         self,
         pot: SCFPotential,
-        method: T.Union[T.Literal["interp", "exact"], T.Mapping],
+        method: T.Union[T.Literal["interp", "exact"], MethodsMapping],
         **kwargs: T.Any
     ) -> None:
-        if not isinstance(method, Mapping):
-            raise NotImplementedError
-
-        self._rsampler = method["r"](pot, **kwargs)
-        self._thetasampler = method["theta"](pot, **kwargs)
-        self._phisampler = method["phi"](pot, **kwargs)
+        super().__init__(pot)
+
+        if isinstance(method, Mapping):
+            sampler = None
+            rsampler = method["r"](pot, **kwargs)
+            thetasampler = method["theta"](pot, **kwargs)
+            phisampler = method["phi"](pot, **kwargs)
+        else:
+            sampler_cls: rv_continuous_modrvs
+            if method == "interp":
+                sampler_cls = SCFSamplerIntrp
+            elif method == "exact":
+                sampler_cls = SCFSamplerExact
+
+            sampler = sampler_cls(pot, **kwargs)
+            rsampler = self._sampler._rsampler
+            thetasampler = self._sampler._thetasampler
+            phisampler = self._sampler._phisampler
+
+        self._sampler: T.Optional[SCFSamplerBase] = sampler
+        self._rsampler = rsampler
+        self._thetasampler = thetasampler
+        self._phisampler = phisampler
 
     # /def
 

sample_scf/sample_exact.py

@@ -19,7 +19,6 @@
 import numpy as np
 import numpy.typing as npt
 from galpy.potential import SCFPotential
-from scipy.stats import rv_continuous
 
 # LOCAL
 from ._typing import NDArray64, RandomLike
@@ -72,7 +71,7 @@ def __init__(self, pot: SCFPotential, **kw: T.Any) -> None:
 # radial sampler
 
 
-class SCFRSampler(rv_continuous):
+class SCFRSampler(rv_continuous_modrvs):
     """Sample radial coordinate from an SCF potential.
 
     Parameters

sample_scf/sample_intrp.py

@@ -270,7 +270,7 @@ def __init__(
             raise ValueError(f"Qls must be shape ({len(rgrid)}, {lmax})")
 
         # l = 0 : spherical symmetry
-        term0 = T.cast(npt.NDArray, 0.5 * (xs + 1))  # (T,)
+        term0 = T.cast(NDArray64, 0.5 * (xs + 1))  # (T,)
         # l = 1+ : non-symmetry
         factor = 1.0 / (2.0 * Qls[:, 0])  # (R,)
         term1p = np.sum(

sample_scf/tests/test_base.py

@@ -7,28 +7,61 @@
 # IMPORTS
 
 # THIRD PARTY
-# import astropy.units as u
-# import numpy as np
+import astropy.coordinates as coord
+import astropy.units as u
+import numpy as np
 import pytest
+from numpy.testing import assert_allclose
 
 # LOCAL
 from sample_scf import base
 
-# from galpy.potential import SCFPotential
-
-
 ##############################################################################
 # TESTS
 ##############################################################################
 
 
-class Test_rv_continuous_modrvs:
+class testrvsampler(base.rv_continuous_modrvs):
+    def _cdf(self, x, *args, **kwargs):
+        return x
+
+    # /def
+
+    cdf = _cdf
+
+    def _rvs(self, *args, size=None, random_state=None):
+        if random_state is None:
+            random_state = np.random
+
+        return np.atleast_1d(random_state.uniform(size=size))
+
+    # /def
+
+
+# /class
+
+
+class Test_RVContinuousModRVS:
     """Test `sample_scf.base.rv_continuous_modrvs`."""
 
-    @pytest.mark.skip("TODO!")
-    def test_rvs(self):
+    def setup_class(self):
+        self.sampler = testrvsampler()
+
+    # /def
+
+    # ===============================================================
+
+    @pytest.mark.parametrize(
+        "size, random, expected",
+        [
+            (None, 0, 0.5488135039273248),
+            (1, 2, 0.43599490214200376),
+            ((3, 1), 4, (0.9670298390136767, 0.5472322491757223, 0.9726843599648843)),
+        ],
+    )
+    def test_rvs(self, size, random, expected):
         """Test :meth:`sample_scf.base.rv_continuous_modrvs.rvs`."""
-        assert False
+        assert_allclose(self.sampler.rvs(size=size, random_state=random), expected, atol=1e-16)
 
     # /def
 
@@ -42,46 +75,117 @@ def test_rvs(self):
 class Test_SCFSamplerBase:
     """Test :class:`sample_scf.base.SCFSamplerBase`."""
 
-    _cls = base.SCFSamplerBase
+    def setup_class(self):
+        self.cls = base.SCFSamplerBase
+        self.cls_args = ()
 
-    @pytest.mark.skip("TODO!")
-    def test_rsampler(self):
+        self.expected_rvs = {
+            0: dict(r=0.548813503927, theta=1.021982822867 * u.rad, phi=0.548813503927 * u.rad),
+            1: dict(r=0.548813503927, theta=1.021982822867 * u.rad, phi=0.548813503927 * u.rad),
+            2: dict(
+                r=[0.9670298390136, 0.5472322491757, 0.9726843599648, 0.7148159936743],
+                theta=[0.603766487781, 1.023564077619, 0.598111966830, 0.855980333120] * u.rad,
+                phi=[0.9670298390136, 0.547232249175, 0.9726843599648, 0.7148159936743] * u.rad,
+            ),
+        }
+
+    # /def
+
+    @pytest.fixture(autouse=True, scope="class")
+    def sampler(self, potentials):
+        """Set up r, theta, phi sampler."""
+        sampler = self.cls(potentials, *self.cls_args)
+        sampler._rsampler = testrvsampler()
+        sampler._thetasampler = testrvsampler()
+        sampler._phisampler = testrvsampler()
+
+        return sampler
+
+    # /def
+
+    # ===============================================================
+
+    def test_rsampler(self, sampler):
         """Test :meth:`sample_scf.base.SCFSamplerBase.rsampler`."""
-        assert False
+        assert isinstance(sampler.rsampler, base.rv_continuous_modrvs)
 
     # /def
 
-    @pytest.mark.skip("TODO!")
-    def test_thetasampler(self):
+    def test_thetasampler(self, sampler):
         """Test :meth:`sample_scf.base.SCFSamplerBase.thetasampler`."""
-        assert False
+        assert isinstance(sampler.thetasampler, base.rv_continuous_modrvs)
 
     # /def
 
-    @pytest.mark.skip("TODO!")
-    def test_phisampler(self):
+    def test_phisampler(self, sampler):
         """Test :meth:`sample_scf.base.SCFSamplerBase.phisampler`."""
-        assert False
+        assert isinstance(sampler.phisampler, base.rv_continuous_modrvs)
 
     # /def
 
-    @pytest.mark.skip("TODO!")
-    def test_cdf(self):
+    @pytest.mark.parametrize(
+        "r, theta, phi, expected",
+        [
+            (0, 0, 0, [0, 0, 0]),
+            (1, 0, 0, [1, 0, 0]),
+            ([0, 1], [0, 0], [0, 0], [[0, 0, 0], [1, 0, 0]]),
+        ],
+    )
+    def test_cdf(self, sampler, r, theta, phi, expected):
         """Test :meth:`sample_scf.base.SCFSamplerBase.cdf`."""
-        assert False
+        assert np.allclose(sampler.cdf(r, theta, phi), expected, atol=1e-16)
 
     # /def
 
-    @pytest.mark.skip("TODO!")
-    def test_rvs(self):
+    @pytest.mark.parametrize(
+        "id, size, random",
+        [
+            (0, None, 0),
+            (1, 1, 0),
+            (2, 4, 4),
+        ],
+    )
+    def test_rvs(self, sampler, id, size, random):
         """Test :meth:`sample_scf.base.SCFSamplerBase.rvs`."""
-        assert False
+        samples = sampler.rvs(size=size, random_state=random)
+        sce = coord.PhysicsSphericalRepresentation(**self.expected_rvs[id])
+
+        assert_allclose(samples.r, sce.r, atol=1e-16)
+        assert_allclose(samples.theta.value, sce.theta.value, atol=1e-16)
+        assert_allclose(samples.phi.value, sce.phi.value, atol=1e-16)
 
     # /def
 
 
 # /class
 
 
+class SCFSamplerTestBase(Test_SCFSamplerBase):
+    def setup_class(self):
+
+        self.expected_rvs = {
+            0: dict(r=0.548813503927, theta=1.021982822867 * u.rad, phi=0.548813503927 * u.rad),
+            1: dict(r=0.548813503927, theta=1.021982822867 * u.rad, phi=0.548813503927 * u.rad),
+            2: dict(
+                r=[0.9670298390136, 0.5472322491757, 0.9726843599648, 0.7148159936743],
+                theta=[0.603766487781, 1.023564077619, 0.598111966830, 0.855980333120] * u.rad,
+                phi=[0.9670298390136, 0.547232249175, 0.9726843599648, 0.7148159936743] * u.rad,
+            ),
+        }
+
+    # /def
+
+    @pytest.fixture(autouse=True, scope="class")
+    def sampler(self, potentials):
+        """Set up r, theta, phi sampler."""
+        sampler = self.cls(potentials, *self.cls_args)
+
+        return sampler
+
+    # /def
+
+
+# /class
+
 ##############################################################################
 # END