From 16954357470a1c59a85cd6ef83affb142a259244 Mon Sep 17 00:00:00 2001 From: michael-petersen Date: Wed, 1 Jul 2026 11:40:10 +0100 Subject: [PATCH 1/4] clean up discmodel basics --- pyproject.toml | 2 +- pytest.ini | 1 + setup.py | 7 +----- src/discmodel/__init__.py | 10 +++++--- src/discmodel/discmodel.py | 44 ++++++++++------------------------ tests/test_discmodel.py | 19 +++++++++++++++ tests/test_optional_imports.py | 31 +++++++++++++++++++----- 7 files changed, 66 insertions(+), 48 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 14a7fb4..84138f4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,7 +30,7 @@ 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] diff --git a/pytest.ini b/pytest.ini index 08ba66f..620ce09 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,3 @@ [pytest] addopts = --ignore=flex +pythonpath = src diff --git a/setup.py b/setup.py index 562953b..b024da8 100644 --- a/setup.py +++ b/setup.py @@ -1,9 +1,4 @@ from setuptools import setup -import numpy as np -setup( - name="discmodel", - packages=["discmodel"], - package_dir={"": "src"} -) +setup() diff --git a/src/discmodel/__init__.py b/src/discmodel/__init__.py index 893f28e..fa313ac 100644 --- a/src/discmodel/__init__.py +++ b/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"] diff --git a/src/discmodel/discmodel.py b/src/discmodel/discmodel.py index 7179097..c468ee8 100644 --- a/src/discmodel/discmodel.py +++ b/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.RandomState(self.seed) + m = rng.rand(self.N) r = f(m) # pick a random azimuthal angle - p = 2.*np.pi*np.random.rand(self.N) + p = 2.*np.pi*rng.rand(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) @@ -264,4 +244,4 @@ def compute_a1(self,E): - \ No newline at end of file + diff --git a/tests/test_discmodel.py b/tests/test_discmodel.py index e5899e3..c8fb16d 100644 --- a/tests/test_discmodel.py +++ b/tests/test_discmodel.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import discmodel @@ -84,6 +85,19 @@ 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_version(): @@ -113,6 +127,8 @@ def test_discmodel_image(): def test_discmodel_expansion(): + pytest.importorskip("flex") + N = 1000 a = 3.0 M = 1.0 @@ -149,6 +165,9 @@ def test_discmodel_expansion(): def test_discmodel_resampling(): + pytest.importorskip("flex") + pytest.importorskip("lintsampler") + N = 1000 a = 3.0 M = 1.0 diff --git a/tests/test_optional_imports.py b/tests/test_optional_imports.py index 91efd65..b0f04ad 100644 --- a/tests/test_optional_imports.py +++ b/tests/test_optional_imports.py @@ -1,7 +1,7 @@ -# tests/test_optional_imports.py import importlib -import pytest -import discmodel.discmodel as dm +import types + +from discmodel.optional_imports import _check_flex, _check_lintsampler def test_check_lintsampler_missing(monkeypatch): # make importlib.import_module raise for lintsampler @@ -10,14 +10,33 @@ def raise_for(name): raise ImportError return importlib.import_module(name) monkeypatch.setattr(importlib, "import_module", raise_for) - assert dm._check_lintsampler() is False + assert _check_lintsampler() is False def test_check_lintsampler_present(monkeypatch): # make importlib.import_module return a dummy module for lintsampler def stub(name): if name == "lintsampler": - import types return types.ModuleType("lintsampler") return importlib.import_module(name) monkeypatch.setattr(importlib, "import_module", stub) - assert dm._check_lintsampler() is True + assert _check_lintsampler() is True + + +def test_check_flex_missing(monkeypatch): + # make importlib.import_module raise for flex + def raise_for(name): + if name == "flex": + raise ImportError + return importlib.import_module(name) + monkeypatch.setattr(importlib, "import_module", raise_for) + assert _check_flex() is False + + +def test_check_flex_present(monkeypatch): + # make importlib.import_module return a dummy module for flex + def stub(name): + if name == "flex": + return types.ModuleType("flex") + return importlib.import_module(name) + monkeypatch.setattr(importlib, "import_module", stub) + assert _check_flex() is True From 91d79af42da7f85e8b56ee30042a99539b0e3220 Mon Sep 17 00:00:00 2001 From: michael-petersen Date: Wed, 1 Jul 2026 12:35:54 +0100 Subject: [PATCH 2/4] oops broken test --- src/discmodel/discmodel.py | 7 +++--- tests/test_discmodel.py | 44 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/src/discmodel/discmodel.py b/src/discmodel/discmodel.py index c468ee8..926525e 100644 --- a/src/discmodel/discmodel.py +++ b/src/discmodel/discmodel.py @@ -227,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) diff --git a/tests/test_discmodel.py b/tests/test_discmodel.py index c8fb16d..2fa2163 100644 --- a/tests/test_discmodel.py +++ b/tests/test_discmodel.py @@ -2,6 +2,7 @@ import pytest import discmodel +import discmodel.discmodel as discmodel_module def test_discmodel_initialization(): """Test initialization of DiscGalaxy class.""" @@ -186,3 +187,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) From 578a813bf990e16b3c2876cf611f427986df22c7 Mon Sep 17 00:00:00 2001 From: michael-petersen Date: Wed, 1 Jul 2026 13:17:46 +0100 Subject: [PATCH 3/4] fix rng; up coverage --- src/discmodel/discmodel.py | 6 +-- tests/test_discmodel.py | 85 ++++++++++++++++++++++++++++++++++ tests/test_optional_imports.py | 16 +++++-- 3 files changed, 100 insertions(+), 7 deletions(-) diff --git a/src/discmodel/discmodel.py b/src/discmodel/discmodel.py index 926525e..d37e98c 100644 --- a/src/discmodel/discmodel.py +++ b/src/discmodel/discmodel.py @@ -50,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 - rng = np.random.RandomState(self.seed) - m = rng.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*rng.rand(self.N) + p = 2.*np.pi*rng.random(self.N) x = r*np.cos(p) y = r*np.sin(p) diff --git a/tests/test_discmodel.py b/tests/test_discmodel.py index 2fa2163..134d109 100644 --- a/tests/test_discmodel.py +++ b/tests/test_discmodel.py @@ -101,6 +101,18 @@ def test_discmodel_initialization_does_not_reset_global_rng(): 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(): """Test that the version string is correctly set.""" import discmodel @@ -127,6 +139,79 @@ 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") diff --git a/tests/test_optional_imports.py b/tests/test_optional_imports.py index b0f04ad..f2621e6 100644 --- a/tests/test_optional_imports.py +++ b/tests/test_optional_imports.py @@ -5,38 +5,46 @@ def test_check_lintsampler_missing(monkeypatch): # make importlib.import_module raise for lintsampler + original_import_module = importlib.import_module + def raise_for(name): if name == "lintsampler": raise ImportError - return importlib.import_module(name) + return original_import_module(name) monkeypatch.setattr(importlib, "import_module", raise_for) assert _check_lintsampler() is False def test_check_lintsampler_present(monkeypatch): # make importlib.import_module return a dummy module for lintsampler + original_import_module = importlib.import_module + def stub(name): if name == "lintsampler": return types.ModuleType("lintsampler") - return importlib.import_module(name) + return original_import_module(name) monkeypatch.setattr(importlib, "import_module", stub) assert _check_lintsampler() is True def test_check_flex_missing(monkeypatch): # make importlib.import_module raise for flex + original_import_module = importlib.import_module + def raise_for(name): if name == "flex": raise ImportError - return importlib.import_module(name) + return original_import_module(name) monkeypatch.setattr(importlib, "import_module", raise_for) assert _check_flex() is False def test_check_flex_present(monkeypatch): # make importlib.import_module return a dummy module for flex + original_import_module = importlib.import_module + def stub(name): if name == "flex": return types.ModuleType("flex") - return importlib.import_module(name) + return original_import_module(name) monkeypatch.setattr(importlib, "import_module", stub) assert _check_flex() is True From 0393f7f49894ab995801ca1d215c34d1d90978bf Mon Sep 17 00:00:00 2001 From: michael-petersen Date: Wed, 1 Jul 2026 13:33:17 +0100 Subject: [PATCH 4/4] solidify build --- pyproject.toml | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 84138f4..d81121d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,5 +33,10 @@ dynamic = ["version"] 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]