From f950c5f17a53d0548856bad7fa76e1ae4e66c5a5 Mon Sep 17 00:00:00 2001 From: Patrick Roddy Date: Tue, 19 Sep 2023 17:01:04 +0200 Subject: [PATCH] API: use new `numpy.random.Generator` interface (#36) As of v1.17.0, numpy uses a new random number `Generator` syntax. The old system will eventually be deprecated, so this ports over to the new system ahead of that. Closes: #35 Reviewed-by: Nicolas Tessore Co-authored-by: Nicolas Tessore --- heracles/maps.py | 32 +++++++++++++++++++++++++++----- tests/conftest.py | 6 ++++++ tests/test_catalog.py | 32 ++++++++++++++++---------------- tests/test_covariance.py | 12 ++++++------ tests/test_io.py | 36 ++++++++++++++++++------------------ tests/test_kmeans_radec.py | 6 +++--- tests/test_maps.py | 18 +++++++++--------- tests/test_twopoint.py | 26 +++++++++++++------------- 8 files changed, 98 insertions(+), 70 deletions(-) diff --git a/heracles/maps.py b/heracles/maps.py index f7e6732..81f7d7e 100644 --- a/heracles/maps.py +++ b/heracles/maps.py @@ -170,9 +170,19 @@ class RandomizableMap(Map): """ - def __init__(self, randomize: bool, **kwargs) -> None: + default_rng: np.random.Generator = np.random.default_rng() + """Default random number generator for randomizable maps.""" + + def __init__( + self, + randomize: bool, + *, + rng: t.Optional[np.random.Generator] = None, + **kwargs, + ) -> None: """Initialise map with the given randomize property.""" self._randomize = randomize + self._rng = rng super().__init__(**kwargs) @property @@ -184,6 +194,16 @@ def randomize(self, randomize: bool) -> None: """Set the randomize flag.""" self._randomize = randomize + @property + def rng(self) -> np.random.Generator: + """Random number generator of this map.""" + return self._rng or self.default_rng + + @rng.setter + def rng(self, rng: np.random.Generator) -> None: + """Set the random number generator of this map.""" + self._rng = rng + class NormalizableMap(Map): """Abstract base class for normalisable maps. @@ -226,9 +246,10 @@ def __init__( *, overdensity: bool = True, randomize: bool = False, + rng: t.Optional[np.random.Generator] = None, ) -> None: """Create a position map with the given properties.""" - super().__init__(columns=(lon, lat), nside=nside, randomize=randomize) + super().__init__(columns=(lon, lat), nside=nside, randomize=randomize, rng=rng) self._overdensity: bool = overdensity @property @@ -242,7 +263,6 @@ def overdensity(self, overdensity: bool) -> None: def __call__(self, catalog: "Catalog") -> MapGenerator: """Map the given catalogue.""" - # get catalogue column definition col = self.columns @@ -285,7 +305,7 @@ def mapper(page: "CatalogPage") -> None: p = np.full(npix, 1 / npix) else: p = vmap / np.sum(vmap) - pos[:] = np.random.multinomial(ngal, p) + pos[:] = self.rng.multinomial(ngal, p) # compute average number density nbar = ngal / npix @@ -448,6 +468,7 @@ def __init__( conjugate: bool = False, normalize: bool = True, randomize: bool = False, + rng: t.Optional[np.random.Generator] = None, ) -> None: """Create a new shear map.""" @@ -458,6 +479,7 @@ def __init__( nside=nside, normalize=normalize, randomize=randomize, + rng=rng, ) @property @@ -521,7 +543,7 @@ def mapper(page: "CatalogPage") -> None: im = -im if randomize: - a = np.random.uniform(0.0, 2 * np.pi, size=page.size) + a = self.rng.uniform(0.0, 2 * np.pi, size=page.size) r = np.hypot(re, im) re, im = r * np.cos(a), r * np.sin(a) del a, r diff --git a/tests/conftest.py b/tests/conftest.py index eccfcae..70a306a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import warnings from contextlib import contextmanager +import numpy as np import pytest from numba import config @@ -22,3 +23,8 @@ def warns(*types): yield finally: pass + + +@pytest.fixture(scope="session") +def rng(seed: int = 50) -> np.random.Generator: + return np.random.default_rng(seed) diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 912d04d..8ab0ba9 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -4,14 +4,14 @@ @pytest.fixture -def catalog(): +def catalog(rng): from heracles.catalog import CatalogBase, CatalogPage # fix a set of rows to be returned for testing size = 100 - x = np.random.rand(size) - y = np.random.rand(size) - z = np.random.rand(size) + x = rng.random(size) + y = rng.random(size) + z = rng.random(size) class TestCatalog(CatalogBase): SIZE = size @@ -299,19 +299,19 @@ def test_invalid_value_filter(catalog): npt.assert_array_equal(page.get(k), v[2:]) -def test_footprint_filter(catalog): +def test_footprint_filter(catalog, rng): from healpy import ang2pix from heracles.catalog import FootprintFilter # footprint for northern hemisphere nside = 8 - m = np.round(np.random.rand(12 * nside**2)) + m = np.round(rng.random(12 * nside**2)) # replace x and y in catalog with lon and lat - catalog.DATA["x"] = lon = np.random.uniform(-180, 180, size=catalog.SIZE) + catalog.DATA["x"] = lon = rng.uniform(-180, 180, size=catalog.SIZE) catalog.DATA["y"] = lat = np.degrees( - np.arcsin(np.random.uniform(-1, 1, size=catalog.SIZE)), + np.arcsin(rng.uniform(-1, 1, size=catalog.SIZE)), ) filt = FootprintFilter(m, "x", "y") @@ -329,12 +329,12 @@ def test_footprint_filter(catalog): np.testing.assert_array_equal(page[k], v[good]) -def test_array_catalog(): +def test_array_catalog(rng): from heracles.catalog import ArrayCatalog, Catalog arr = np.empty(100, [("lon", float), ("lat", float), ("x", float), ("y", float)]) for name in arr.dtype.names: - arr[name] = np.random.rand(len(arr)) + arr[name] = rng.random(len(arr)) catalog = ArrayCatalog(arr) @@ -372,15 +372,15 @@ def test_array_catalog(): assert copied.__dict__ == catalog.__dict__ -def test_fits_catalog(tmp_path): +def test_fits_catalog(rng, tmp_path): import fitsio from heracles.catalog import Catalog from heracles.catalog.fits import FitsCatalog size = 100 - ra = np.random.uniform(-180, 180, size=size) - dec = np.random.uniform(-90, 90, size=size) + ra = rng.uniform(-180, 180, size=size) + dec = rng.uniform(-90, 90, size=size) filename = str(tmp_path / "catalog.fits") @@ -436,7 +436,7 @@ def test_fits_catalog(tmp_path): assert copied._ext == catalog._ext -def test_fits_catalog_caching(tmp_path): +def test_fits_catalog_caching(rng, tmp_path): import gc import fitsio @@ -444,8 +444,8 @@ def test_fits_catalog_caching(tmp_path): from heracles.catalog.fits import FitsCatalog size = 100 - ra = np.random.uniform(-180, 180, size=size) - dec = np.random.uniform(-90, 90, size=size) + ra = rng.uniform(-180, 180, size=size) + dec = rng.uniform(-90, 90, size=size) filename = str(tmp_path / "cached.fits") diff --git a/tests/test_covariance.py b/tests/test_covariance.py index 15c240a..addf184 100644 --- a/tests/test_covariance.py +++ b/tests/test_covariance.py @@ -2,15 +2,15 @@ import pytest -def test_sample_covariance(): +def test_sample_covariance(rng): from heracles.covariance import SampleCovariance, add_sample n = 10 size = 3 size2 = 5 - samples = [np.random.randn(size) for _ in range(n)] - samples2 = [np.random.randn(size2) for _ in range(n)] + samples = [rng.standard_normal(size) for _ in range(n)] + samples2 = [rng.standard_normal(size2) for _ in range(n)] cov = SampleCovariance(size) @@ -53,7 +53,7 @@ def test_sample_covariance(): add_sample(cov, np.zeros(size + 1), np.zeros(size2 - 1)) -def test_update_covariance(): +def test_update_covariance(rng): from itertools import combinations_with_replacement from heracles.covariance import update_covariance @@ -62,7 +62,7 @@ def test_update_covariance(): cov = {} - sample = {i: np.random.randn(i + 1) for i in range(n)} + sample = {i: rng.standard_normal(i + 1) for i in range(n)} update_covariance(cov, sample) assert len(cov) == n * (n + 1) // 2 @@ -71,7 +71,7 @@ def test_update_covariance(): assert cov[k1, k2].shape == (sample[k1].size, sample[k2].size) assert np.all(cov[k1, k2] == 0) - sample2 = {i: np.random.randn(i + 1) for i in range(n)} + sample2 = {i: rng.standard_normal(i + 1) for i in range(n)} update_covariance(cov, sample2) assert len(cov) == n * (n + 1) // 2 diff --git a/tests/test_io.py b/tests/test_io.py index 7939455..4c64702 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -9,7 +9,7 @@ def zbins(): @pytest.fixture -def mock_alms(zbins): +def mock_alms(rng, zbins): import numpy as np lmax = 32 @@ -21,7 +21,7 @@ def mock_alms(zbins): alms = {} for n in names: for i in zbins: - a = np.random.randn(Nlm, 2) @ [1, 1j] + a = rng.standard_normal((Nlm, 2)) @ [1, 1j] a.dtype = np.dtype(a.dtype, metadata={"nside": 32}) alms[n, i] = a @@ -29,10 +29,10 @@ def mock_alms(zbins): @pytest.fixture -def mock_cls(): +def mock_cls(rng): import numpy as np - cl = np.random.rand(101) + cl = rng.random(101) cl.dtype = np.dtype(cl.dtype, metadata={"nside_1": 32, "nside_2": 64}) return { @@ -71,13 +71,13 @@ def datadir(tmp_path_factory): @pytest.fixture(scope="session") -def mock_mask_fields(nside): +def mock_mask_fields(nside, rng): import healpy as hp import numpy as np npix = hp.nside2npix(nside) - maps = np.random.rand(npix * NFIELDS_TEST).reshape((npix, NFIELDS_TEST)) - pixels = np.unique(np.random.randint(0, npix, npix // 3)) + maps = rng.random(npix * NFIELDS_TEST).reshape((npix, NFIELDS_TEST)) + pixels = np.unique(rng.integers(0, npix, size=npix // 3)) maskpix = np.delete(np.arange(0, npix), pixels) for i in range(NFIELDS_TEST): maps[:, i][maskpix] = 0 @@ -141,13 +141,13 @@ def mock_writemask_full(mock_mask_fields, nside, datadir): @pytest.fixture(scope="session") -def mock_mask_extra(nside): +def mock_mask_extra(nside, rng): import healpy as hp import numpy as np npix = hp.nside2npix(nside) - maps = np.random.rand(npix) - pixels = np.unique(np.random.randint(0, npix, npix // 3)) + maps = rng.random(npix) + pixels = np.unique(rng.integers(0, npix, size=npix // 3)) maskpix = np.delete(np.arange(0, npix), pixels) maps[maskpix] = 0 return [maps, pixels] @@ -178,7 +178,7 @@ def mock_writemask_extra(mock_mask_extra, nside, datadir): return filename -def test_write_read_maps(tmp_path): +def test_write_read_maps(rng, tmp_path): import healpy as hp import numpy as np @@ -187,9 +187,9 @@ def test_write_read_maps(tmp_path): nside = 4 npix = 12 * nside**2 - p = np.random.rand(npix) - v = np.random.rand(npix) - g = np.random.rand(2, npix) + p = rng.random(npix) + v = rng.random(npix) + g = rng.random((2, npix)) p.dtype = np.dtype(p.dtype, metadata={"spin": 0}) v.dtype = np.dtype(v.dtype, metadata={"spin": 0}) @@ -254,7 +254,7 @@ def test_write_read_cls(mock_cls, tmp_path): assert cl.dtype.metadata == mock_cls[key].dtype.metadata -def test_write_read_mms(tmp_path): +def test_write_read_mms(rng, tmp_path): import numpy as np from heracles.io import read_mms, write_mms @@ -263,9 +263,9 @@ def test_write_read_mms(tmp_path): workdir = str(tmp_path) mms = { - ("00", 0, 1): np.random.randn(10, 10), - ("0+", 1, 2): np.random.randn(20, 5), - ("++", 2, 3): np.random.randn(10, 5, 2), + ("00", 0, 1): rng.standard_normal((10, 10)), + ("0+", 1, 2): rng.standard_normal((20, 5)), + ("++", 2, 3): rng.standard_normal((10, 5, 2)), } write_mms(filename, mms, workdir=workdir) diff --git a/tests/test_kmeans_radec.py b/tests/test_kmeans_radec.py index 8d089f1..e825eb1 100644 --- a/tests/test_kmeans_radec.py +++ b/tests/test_kmeans_radec.py @@ -2,7 +2,7 @@ @pytest.mark.flaky(reruns=2) -def test_kmeans_sample(): +def test_kmeans_sample(rng): import numpy as np from heracles._kmeans_radec import kmeans_sample @@ -11,8 +11,8 @@ def test_kmeans_sample(): ncen = 20 pts = np.empty((npts, 2)) - pts[:, 0] = np.random.uniform(-180, 180, size=npts) - pts[:, 1] = np.degrees(np.arcsin(np.random.uniform(-1, 1, size=npts))) + pts[:, 0] = rng.uniform(-180, 180, size=npts) + pts[:, 1] = np.degrees(np.arcsin(rng.uniform(-1, 1, size=npts))) km = kmeans_sample(pts, ncen, verbose=2) diff --git a/tests/test_maps.py b/tests/test_maps.py index 4f796c6..5b3b148 100644 --- a/tests/test_maps.py +++ b/tests/test_maps.py @@ -30,12 +30,12 @@ def sigma_e(): @pytest.fixture -def vmap(nside): - return np.round(np.random.rand(12 * nside**2)) +def vmap(nside, rng): + return np.round(rng.random(12 * nside**2)) @pytest.fixture -def page(nside): +def page(nside, rng): from unittest.mock import Mock ipix = np.ravel( @@ -47,9 +47,9 @@ def page(nside): size = ra.size - w = np.random.rand(size // 4, 4) - g1 = np.random.randn(size // 4, 4) - g2 = np.random.randn(size // 4, 4) + w = rng.random((size // 4, 4)) + g1 = rng.standard_normal((size // 4, 4)) + g2 = rng.standard_normal((size // 4, 4)) g1 -= np.sum(w * g1, axis=-1, keepdims=True) / np.sum(w, axis=-1, keepdims=True) g2 -= np.sum(w * g2, axis=-1, keepdims=True) / np.sum(w, axis=-1, keepdims=True) w, g1, g2 = w.reshape(-1), g1.reshape(-1), g2.reshape(-1) @@ -301,15 +301,15 @@ def test_weight_map(nside, catalog): np.testing.assert_array_almost_equal(m, w) -def test_transform_maps(): +def test_transform_maps(rng): from heracles.maps import transform_maps, update_metadata nside = 32 npix = 12 * nside**2 - t = np.random.randn(npix) + t = rng.standard_normal(npix) update_metadata(t, spin=0, a=1) - p = np.random.randn(2, npix) + p = rng.standard_normal((2, npix)) update_metadata(p, spin=2, b=2) # single scalar map diff --git a/tests/test_twopoint.py b/tests/test_twopoint.py index c4ab9a0..1b2501e 100644 --- a/tests/test_twopoint.py +++ b/tests/test_twopoint.py @@ -15,7 +15,7 @@ def zbins(): @pytest.fixture -def mock_alms(zbins): +def mock_alms(rng, zbins): import numpy as np lmax = 32 @@ -27,7 +27,7 @@ def mock_alms(zbins): alms = {} for n in names: for i in zbins: - a = np.random.randn(Nlm, 2) @ [1, 1j] + a = rng.standard_normal((Nlm, 2)) @ [1, 1j] a.dtype = np.dtype(a.dtype, metadata={"nside": 32}) alms[n, i] = a @@ -111,14 +111,14 @@ def test_debias_cls(): assert np.all(cls["PP", 0, 0] == -1.23) -def test_mixing_matrices(): +def test_mixing_matrices(rng): from heracles.twopoint import mixing_matrices # this only tests the function logic # the mixing matrix computation itself is tested elsewhere lmax = 20 - cl = np.random.randn(lmax + 1) + cl = rng.standard_normal(lmax + 1) # compute pos-pos cls = {("V", "V", 0, 1): cl} @@ -179,15 +179,15 @@ def test_pixelate_mms_healpix(): @pytest.mark.parametrize("weights", [None, "l(l+1)", "2l+1", ""]) -def test_binned_cls(weights): +def test_binned_cls(rng, weights): from heracles.twopoint import binned_cls - cls = {"key": np.random.randn(21)} + cls = {"key": rng.standard_normal(21)} bins = [2, 5, 10, 15, 20] if weights == "": - weights_ = np.random.rand(40) + weights_ = rng.random(40) else: weights_ = weights @@ -223,15 +223,15 @@ def test_binned_cls(weights): @pytest.mark.parametrize("weights", [None, "l(l+1)", "2l+1", ""]) @pytest.mark.parametrize("ndim", [1, 2, 3]) -def test_bin2pt(ndim, weights): +def test_bin2pt(ndim, rng, weights): from heracles.twopoint import bin2pt - data = np.random.randn(*[21, 31, 41][:ndim]) + data = rng.standard_normal((21, 31, 41)[:ndim]) bins = [2, 5, 10, 15, 20] if weights == "": - weights_ = np.random.rand(51) + weights_ = rng.random(51) else: weights_ = weights @@ -266,7 +266,7 @@ def test_bin2pt(ndim, weights): @pytest.mark.parametrize("full", [False, True]) -def test_random_bias(full): +def test_random_bias(full, rng): from heracles.twopoint import random_bias nside = 64 @@ -275,8 +275,8 @@ def test_random_bias(full): catalog = unittest.mock.Mock() catalog.visibility = None - map_a = unittest.mock.Mock(side_effect=lambda _: np.random.rand(npix)) - map_b = unittest.mock.Mock(side_effect=lambda _: np.random.rand(npix)) + map_a = unittest.mock.Mock(side_effect=lambda _: rng.random(npix)) + map_b = unittest.mock.Mock(side_effect=lambda _: rng.random(npix)) initial_randomize = [map_a.randomize, map_b.randomize]