clean up discmodel basics

pyproject.toml

@@ -30,8 +30,13 @@ classifiers = [
 dynamic = ["version"]
 
 [project.urls]
-Homepage = "https://github.com/michael-petersen/discmodel"
+Homepage = "https://github.com/ObservationalExpansions/discmodel"
 Documentation = "https://discmodel.readthedocs.io"
 
-[tool.setuptools_scm]
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
 
+[tool.setuptools_scm]

pytest.ini

@@ -1,2 +1,3 @@
 [pytest]
 addopts = --ignore=flex
+pythonpath = src

setup.py

@@ -1,9 +1,4 @@
 from setuptools import setup
-import numpy as np
 
 
-setup(
-    name="discmodel",
-    packages=["discmodel"],
-    package_dir={"": "src"}
-)
+setup()

src/discmodel/__init__.py

@@ -6,13 +6,17 @@
 See README.md for details
 """
 
+from importlib.metadata import PackageNotFoundError, version
+
 from .optional_imports import _check_lintsampler, _check_flex
 HAS_LINTSAMPLER = _check_lintsampler()
 HAS_FLEX = _check_flex()
 
 from .discmodel import DiscGalaxy
-from importlib.metadata import version
 
-__version__ = version("discmodel")
-__all__ = ["DiscGalaxy"]
+try:
+    __version__ = version("discmodel")
+except PackageNotFoundError:
+    __version__ = "0+unknown"
 
+__all__ = ["DiscGalaxy"]

src/discmodel/discmodel.py

@@ -1,9 +1,5 @@
-
 import numpy as np
 
-# for the Laguerre polynomials
-from scipy.special import eval_genlaguerre
-
 # for interpolation
 from scipy import interpolate
 
@@ -22,11 +18,12 @@
 class DiscGalaxy(object):
 
 
-    def __init__(self,N=None,phasespace=None,a=3.,M=1.,vcirc=200.,rmax=30.):
+    def __init__(self,N=None,phasespace=None,a=3.,M=1.,vcirc=200.,rmax=30.,seed=42):
 
         self.a = a
         self.M = M
         self.vcirc = vcirc
+        self.seed = seed
         self.rmax = rmax*a    # rmax is now a multiple of the scale length; this should only need to be changed in the case of very large N
 
         if N is not None:
@@ -53,12 +50,12 @@ def menclosed(r,a=self.a):
         f = interpolate.interp1d(menclosed(x),x)
 
         # pull a bunch of points: pick a random radius in the disc
-        np.random.seed(42)  # for reproducibility: this might need to go somewhere else?
-        m = np.random.rand(self.N)
+        rng = np.random.default_rng(self.seed)
+        m = rng.random(self.N)
         r = f(m)
 
         # pick a random azimuthal angle
-        p = 2.*np.pi*np.random.rand(self.N)
+        p = 2.*np.pi*rng.random(self.N)
 
         x = r*np.cos(p)
         y = r*np.sin(p)
@@ -135,30 +132,17 @@ def rotate_disc(self,xrotation=0.,yrotation=0.,zrotation=0.,euler=False):
 
         Rmatrix = self.make_rotation_matrix(xrotation,yrotation,zrotation,euler=euler)
 
-        #
         # do the transformation in position
         tmp = np.dot(np.array([x,y,z]).T,Rmatrix)
-
-        #try:
         xout = tmp[:,0]
         yout = tmp[:,1]
         zout = tmp[:,2]
-        #except:
-        #    xout = tmp[0]
-        #    yout = tmp[1]
-        #    zout = tmp[2]
 
         # and in velocity
         tmpv = np.dot(np.array([u,v,w]).T,Rmatrix)
-
-        #try:
         uout = tmpv[:,0]
         vout = tmpv[:,1]
         wout = tmpv[:,2]
-        #except:
-        #    uout = tmpv[0]
-        #    vout = tmpv[1]
-        #    wout = tmpv[2]        
 
         self.x = xout
         self.y = yout
@@ -184,7 +168,6 @@ def generate_image(self,rmax,nbins,noiselevel=-1.0):
         self.x_centers = (self.x_edges[:-1] + self.x_edges[1:]) / 2
         self.y_centers = (self.y_edges[:-1] + self.y_edges[1:]) / 2
 
-        dx = self.x_edges[1]-self.x_edges[0]
         xpix,ypix = np.meshgrid(self.x_centers,self.y_centers,indexing='ij')
         rr,pp = np.sqrt(xpix**2+ypix**2),np.arctan2(ypix,xpix)
 
@@ -198,19 +181,18 @@ def make_expansion(self,mmax,nmax,rscl,xmax=10000.,noisy=False):
         if not HAS_FLEX:
             raise ImportError("flex is not available. Please install flex to use this method.")
 
-        try:
-            snapshot = self.img
-        except:
-            print('No image data to expand... run generate_image first.')
-            return
+        if not hasattr(self, "img"):
+            raise RuntimeError("No image data to expand. Run generate_image first.")
 
         if noisy:
+            if not hasattr(self, "noisyimage"):
+                raise RuntimeError("No noisy image data to expand. Run generate_image with noiselevel first.")
             snapshot = self.noisyimage
+        else:
+            snapshot = self.img
 
         # recreate these temporarily, so we can use them
-        dx = self.x_edges[1]-self.x_edges[0]
         xpix,ypix = np.meshgrid(self.x_centers,self.y_centers,indexing='ij')
-        rr,pp = np.sqrt(xpix**2+ypix**2),np.arctan2(ypix,xpix)
 
         rval = np.sqrt(xpix**2+ypix**2).reshape(-1,)
         phi  = np.arctan2(ypix,xpix).reshape(-1,)
@@ -219,8 +201,6 @@ def make_expansion(self,mmax,nmax,rscl,xmax=10000.,noisy=False):
         # create a mask for pixels outside the maximum radius
         gvals = np.where(rval>xmax)
 
-        #rval[gvals]         = np.nan
-        #phi[gvals]          = np.nan
         snapshotflat[gvals] = np.nan
 
         laguerre = flex.FLEX(rscl,mmax,nmax,rval,phi,mass=snapshotflat)
@@ -247,10 +227,11 @@ def resample_expansion(self,E):
             raise ImportError("lintsampler is not available. Please install lintsampler to use this method.")
 
         def rndmpdf(X): return np.random.uniform()
-        g = lintsampler.DensityGrid((self.x_centers,self.x_centers), rndmpdf)
+        g = lintsampler.DensityGrid((self.x_centers,self.y_centers), rndmpdf)
 
-        E.laguerre_reconstruction(self.r, self.p)
-        g.vertex_densities = E.reconstruction.T/(2.*np.pi)
+        E.laguerre_reconstruction(self.r.reshape(-1), self.p.reshape(-1))
+        reconstruction = E.reconstruction.reshape(self.r.shape)
+        g.vertex_densities = reconstruction.T/(2.*np.pi)
 
         g.masses = g._calculate_faverages() * g._calculate_volumes()
         g._total_mass = np.sum(g.masses)
@@ -264,4 +245,4 @@ def compute_a1(self,E):
 
 
 
-
+

tests/test_discmodel.py

@@ -1,6 +1,8 @@
 import numpy as np
+import pytest
 
 import discmodel
+import discmodel.discmodel as discmodel_module
 
 def test_discmodel_initialization():
     """Test initialization of DiscGalaxy class."""
@@ -84,6 +86,31 @@ def test_discmodel_phasespace_input():
     disc = discmodel.DiscGalaxy(N=N, a=a, M=M, vcirc=vcirc, rmax=rmax)
 
     newdisc = discmodel.DiscGalaxy(phasespace=(disc.x,disc.y,disc.z,disc.u,disc.v,disc.w))
+    assert newdisc.N == N
+
+
+def test_discmodel_initialization_does_not_reset_global_rng():
+    """Test DiscGalaxy uses its own RNG for reproducible initialization."""
+    np.random.seed(1)
+    expected = np.random.rand()
+
+    np.random.seed(1)
+    discmodel.DiscGalaxy(N=10)
+    actual = np.random.rand()
+
+    assert actual == expected
+
+
+def test_discmodel_seed_reproducibility():
+    """Test seeded DiscGalaxy instances generate reproducible phase space."""
+    disc1 = discmodel.DiscGalaxy(N=20, seed=123)
+    disc2 = discmodel.DiscGalaxy(N=20, seed=123)
+    disc3 = discmodel.DiscGalaxy(N=20, seed=124)
+
+    for attr in ("x", "y", "z", "u", "v", "w"):
+        assert np.allclose(getattr(disc1, attr), getattr(disc2, attr))
+
+    assert not np.allclose(disc1.x, disc3.x)
 
 
 def test_discmodel_version():
@@ -112,7 +139,82 @@ def test_discmodel_image():
     assert hasattr(disc,'p')
 
 
+def test_discmodel_expansion_requires_flex(monkeypatch):
+    monkeypatch.setattr(discmodel_module, "HAS_FLEX", False)
+    disc = discmodel.DiscGalaxy(N=10)
+
+    with pytest.raises(ImportError, match="flex is not available"):
+        disc.make_expansion(mmax=4, nmax=4, rscl=1.0)
+
+    with pytest.raises(ImportError, match="flex is not available"):
+        disc.make_particle_expansion(mmax=4, nmax=4, rscl=1.0)
+
+
+def test_discmodel_expansion_requires_image_data(monkeypatch):
+    monkeypatch.setattr(discmodel_module, "HAS_FLEX", True)
+    disc = discmodel.DiscGalaxy(N=10)
+
+    with pytest.raises(RuntimeError, match="Run generate_image first"):
+        disc.make_expansion(mmax=4, nmax=4, rscl=1.0)
+
+
+def test_discmodel_noisy_expansion_requires_noisy_image(monkeypatch):
+    monkeypatch.setattr(discmodel_module, "HAS_FLEX", True)
+    disc = discmodel.DiscGalaxy(N=10)
+    disc.generate_image(rmax=30.0, nbins=5)
+
+    with pytest.raises(RuntimeError, match="Run generate_image with noiselevel first"):
+        disc.make_expansion(mmax=4, nmax=4, rscl=1.0, noisy=True)
+
+
+def test_discmodel_particle_expansion_uses_flex(monkeypatch):
+    class DummyFlex:
+        def __init__(self, rscl, mmax, nmax, rval, phi, mass):
+            self.rscl = rscl
+            self.mmax = mmax
+            self.nmax = nmax
+            self.rval = rval
+            self.phi = phi
+            self.mass = mass
+
+    class DummyFlexModule:
+        FLEX = DummyFlex
+
+    monkeypatch.setattr(discmodel_module, "HAS_FLEX", True)
+    monkeypatch.setattr(discmodel_module, "flex", DummyFlexModule)
+
+    disc = discmodel.DiscGalaxy(N=10, M=2.0)
+    expansion = disc.make_particle_expansion(mmax=3, nmax=2, rscl=1.5)
+
+    assert expansion.rscl == 1.5
+    assert expansion.mmax == 3
+    assert expansion.nmax == 2
+    assert expansion.rval.shape == (10,)
+    assert expansion.phi.shape == (10,)
+    assert np.all(expansion.mass == 0.2)
+
+
+def test_discmodel_resampling_requires_lintsampler(monkeypatch):
+    monkeypatch.setattr(discmodel_module, "HAS_LINTSAMPLER", False)
+    disc = discmodel.DiscGalaxy(N=10)
+
+    with pytest.raises(ImportError, match="lintsampler is not available"):
+        disc.resample_expansion(object())
+
+
+def test_discmodel_compute_a1_with_dummy_expansion():
+    class DummyExpansion:
+        coscoefs = np.array([[1.0, 0.0], [2.0, 0.0]])
+        sincoefs = np.array([[0.0, 0.0], [0.0, 0.0]])
+
+    disc = discmodel.DiscGalaxy(N=10)
+
+    assert disc.compute_a1(DummyExpansion()) == 2.0
+
+
 def test_discmodel_expansion():
+    pytest.importorskip("flex")
+
     N = 1000
     a = 3.0
     M = 1.0
@@ -149,6 +251,9 @@ def test_discmodel_expansion():
 
 
 def test_discmodel_resampling():
+    pytest.importorskip("flex")
+    pytest.importorskip("lintsampler")
+
     N = 1000
     a = 3.0
     M = 1.0
@@ -167,3 +272,46 @@ def test_discmodel_resampling():
 
     # check newdisc is Nx2 (sampled from 2d image only)
     assert newdisc.shape == (N, 2)
+
+
+def test_discmodel_resampling_uses_flattened_reconstruction_grid(monkeypatch):
+    class DummyExpansion:
+        def laguerre_reconstruction(self, rr, pp):
+            rr = np.asarray(rr)
+            pp = np.asarray(pp)
+            assert rr.ndim == 1
+            assert pp.ndim == 1
+            assert rr.shape == pp.shape
+            self.reconstruction = np.ones_like(rr)
+
+    class DummyDensityGrid:
+        def __init__(self, axes, rndmpdf):
+            self.axes = axes
+            self.rndmpdf = rndmpdf
+
+        def _calculate_faverages(self):
+            return np.ones_like(self.vertex_densities)
+
+        def _calculate_volumes(self):
+            return np.ones_like(self.vertex_densities)
+
+    class DummyLintSampler:
+        def __init__(self, grid):
+            self.grid = grid
+
+        def sample(self, n):
+            return np.zeros((n, 2))
+
+    class DummyLintSamplerModule:
+        DensityGrid = DummyDensityGrid
+        LintSampler = DummyLintSampler
+
+    monkeypatch.setattr(discmodel_module, "HAS_LINTSAMPLER", True)
+    monkeypatch.setattr(discmodel_module, "lintsampler", DummyLintSamplerModule)
+
+    disc = discmodel.DiscGalaxy(N=10)
+    disc.generate_image(rmax=30.0, nbins=5)
+
+    pos = disc.resample_expansion(DummyExpansion())
+
+    assert pos.shape == (10, 2)