diff --git a/INSTALL.md b/INSTALL.md index c0f6c5dfe..38fb0e609 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -69,6 +69,13 @@ Now, to install cluster-toolkit, cluster-toolkit has a gsl dependency, you'll al Note, you may choose to install some or all of the ccl, numcosmo, and/or cluster_toolkit packages. You need at least one. If you install cluster_toolkit and others, then you need to install cluster_toolkit *last*. If you have already installed cluster_toolkit before the other packages, simply run, `pip uninstall cluster_toolkit` then re-install cluster_toolkit. +### Optional dependencies +To compute k-means jackknife regions (`clmm.utils.compute_kmeans_centers`, or `ClusterEnsemble.compute_jackknife_covariance` with `n_regions`), install kmeans_radec: + +```bash + pip install git+https://github.com/esheldon/kmeans_radec.git +``` + ## Developer installation Here we provide a quick guide for a developer instalation, using `environment.yml` to create a suitable conda environment and install all necessary packages. diff --git a/README.md b/README.md index 24f9ca48e..58d5d9801 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,10 @@ CLMM has the following dependencies: - [healpy](https://healpy.readthedocs.io/en/latest/index.html) - [qp-prob](https://github.com/LSSTDESC/qp) +Optional: + +- [kmeans_radec](https://github.com/esheldon/kmeans_radec) (to compute k-means jackknife regions) + ### Back-ends For the theoretical predictions of the signal, CLMM relies on existing libraries and **at least one of the following must be installed as well**: diff --git a/clmm/clusterensemble.py b/clmm/clusterensemble.py index d37591545..0cc5d23af 100644 --- a/clmm/clusterensemble.py +++ b/clmm/clusterensemble.py @@ -10,7 +10,7 @@ from .dataops import make_stacked_radial_profile from .gcdata import GCData -from .utils import DiffArray +from .utils import DiffArray, assign_kmeans_regions, compute_kmeans_centers class ClusterEnsemble: @@ -328,11 +328,22 @@ def compute_bootstrap_covariance( self.cov["tan_bs"] = coeff * np.cov(np.array(gt_boot), bias=False, ddof=0) self.cov["cross_bs"] = coeff * np.cov(np.array(gx_boot), bias=False) - def compute_jackknife_covariance(self, tan_component="gt", cross_component="gx", n_side=16): + def compute_jackknife_covariance( + self, + tan_component="gt", + cross_component="gx", + n_side=16, + n_regions=None, + centers=None, + seed=None, + ): """Compute the jackknife covariance matrix, add boostrap covariance matrix for tangential and cross stacked profiles and updates .cov dict (`tan_jk`, `cross_jk`). - Uses healpix sky area sub-division : https://healpix.sourceforge.io + The sky is divided into regions, and each region containing clusters is removed in + turn. By default, the regions are healpix pixels (https://healpix.sourceforge.io). + If `n_regions` or `centers` is given, they are k-means regions instead: each cluster + belongs to the region of its nearest center. Parameters ---------- @@ -343,17 +354,37 @@ def compute_jackknife_covariance(self, tan_component="gt", cross_component="gx", Name of the cross component column in `data`. Default: 'gx' n_side : int - healpix sky area division parameter (number of sky area : 12*n_side^2) + healpix sky area division parameter (number of sky area : 12*n_side^2). + Only used for healpix regions. + n_regions : int, None, optional + Number of k-means regions computed from the cluster positions with + `clmm.utils.compute_kmeans_centers` (requires `kmeans_radec`). Default: None + centers : array_like, None, optional + Centers (ra, dec) in degrees of predefined k-means regions, shape (n, 2), + e.g. to use the same regions for several ensembles. Default: None + seed : int, None, optional + Seed for the k-means initialisation when `n_regions` is given. Default: None """ # may induce artificial noise if there are some healpix pixels # not covering entirely the 2D map of clusters self._check_empty_data() - pixels = healpy.ang2pix(n_side, self.data["ra"], self.data["dec"], nest=True, lonlat=True) - pixels_list_unique = np.unique(pixels) + if n_regions is not None and centers is not None: + raise ValueError("Only one of n_regions and centers can be given.") + if n_regions is not None: + centers = compute_kmeans_centers( + self.data["ra"], self.data["dec"], n_regions, seed=seed + ) + if centers is None: + regions = healpy.ang2pix( + n_side, self.data["ra"], self.data["dec"], nest=True, lonlat=True + ) + else: + regions = assign_kmeans_regions(self.data["ra"], self.data["dec"], centers) + regions_list_unique = np.unique(regions) gt_jack, gx_jack = [], [] - for hp_list_delete in pixels_list_unique: - mask = ~np.isin(pixels, hp_list_delete) + for region_delete in regions_list_unique: + mask = ~np.isin(regions, region_delete) gt, gx = make_stacked_radial_profile( self["radius"][mask], self["W_l"][mask], @@ -361,7 +392,7 @@ def compute_jackknife_covariance(self, tan_component="gt", cross_component="gx", )[1] gt_jack.append(gt) gx_jack.append(gx) - n_jack = pixels_list_unique.size + n_jack = regions_list_unique.size coeff = (n_jack - 1) ** 2 / (n_jack) self.cov["tan_jk"] = coeff * np.cov(np.transpose(gt_jack), bias=False, ddof=0) self.cov["cross_jk"] = coeff * np.cov(np.transpose(gx_jack), bias=False, ddof=0) diff --git a/clmm/utils/__init__.py b/clmm/utils/__init__.py index 3ca66453b..2df5dc8fd 100644 --- a/clmm/utils/__init__.py +++ b/clmm/utils/__init__.py @@ -22,6 +22,7 @@ compute_lensed_ellipticity, convert_shapes_to_epsilon, ) +from .jackknife import assign_kmeans_regions, compute_kmeans_centers from .redshift_tools import _integ_pzfuncs, compute_for_good_redshifts from .statistic import ( _draw_random_points_from_distribution, diff --git a/clmm/utils/jackknife.py b/clmm/utils/jackknife.py new file mode 100644 index 000000000..0d7ee30c3 --- /dev/null +++ b/clmm/utils/jackknife.py @@ -0,0 +1,92 @@ +"""Functions to define sky regions for jackknife resampling""" + +import numpy as np + + +def _radec_to_xyz(ra, dec): + """Convert sky coordinates to cartesian coordinates on the unit sphere. + + Parameters + ---------- + ra, dec : array_like + Right ascension and declination in degrees. + + Returns + ------- + numpy.ndarray + Cartesian coordinates, shape (n, 3). + """ + ra_rad, dec_rad = np.radians(ra), np.radians(dec) + return np.column_stack( + [np.cos(dec_rad) * np.cos(ra_rad), np.cos(dec_rad) * np.sin(ra_rad), np.sin(dec_rad)] + ) + + +def compute_kmeans_centers(ra, dec, n_regions, seed=None, maxiter=100): + """Define `n_regions` jackknife regions on the sky with k-means clustering of the points + (`ra`, `dec`) on the sphere. Requires the `kmeans_radec` package + (https://github.com/esheldon/kmeans_radec). + + Parameters + ---------- + ra, dec : array_like + Right ascension and declination of the points, in degrees. + n_regions : int + Number of regions, at most the number of points. + seed : int, None, optional + Seed for the random first-pass sample and initial centers. A local random generator + is used, so global random states are not changed. Default: None + maxiter : int, optional + Maximum number of k-means iterations. Default: 100 + + Returns + ------- + numpy.ndarray + Region centers (ra, dec) in degrees, shape (`n_regions`, 2). + """ + try: + from kmeans_radec import KMeans # pylint: disable=import-outside-toplevel + except ImportError as err: + raise ImportError( + "compute_kmeans_centers requires kmeans_radec: " + "pip install git+https://github.com/esheldon/kmeans_radec.git" + ) from err + coords = np.column_stack([ra, dec]) + if not 0 < n_regions <= len(coords): + raise ValueError( + f"n_regions ({n_regions}) must be between 1 and the number of points ({len(coords)})" + ) + # Two passes as in kmeans_radec.kmeans_sample, which draws its samples from the global + # `random` state: a first pass on a random subsample, then a pass on all the points. + # kmeans_radec's first-pass sample size, max(2 sqrt(N), 10 n_regions), + # can exceed the number of points, so cap it + rng = np.random.default_rng(seed) + nsample = int(min(len(coords), max(2 * np.sqrt(len(coords)), 10 * n_regions))) + sample = coords[rng.choice(len(coords), nsample, replace=False)] + start = coords[rng.choice(len(coords), n_regions, replace=False)] + kmeans = KMeans(start, verbose=0) + kmeans.run(sample, maxiter=maxiter) + kmeans = KMeans(kmeans.get_centers(), verbose=0) + kmeans.run(coords, maxiter=maxiter) + return kmeans.get_centers() + + +def assign_kmeans_regions(ra, dec, centers): + """Assign each point (`ra`, `dec`) to the region of its nearest center on the sphere. + + Parameters + ---------- + ra, dec : array_like + Right ascension and declination of the points, in degrees. + centers : array_like + Region centers (ra, dec) in degrees, shape (n_regions, 2), e.g. from + `compute_kmeans_centers`. + + Returns + ------- + numpy.ndarray + Index of the region of each point. + """ + centers = np.atleast_2d(centers) + # the nearest center has the largest scalar product on the unit sphere + return np.argmax(_radec_to_xyz(ra, dec) @ _radec_to_xyz(*centers.T).T, axis=1) diff --git a/docs/doc-config.ini b/docs/doc-config.ini index 2b4f14883..362a5c611 100644 --- a/docs/doc-config.ini +++ b/docs/doc-config.ini @@ -31,5 +31,6 @@ OTHER ../examples/demo_mass_conversion.ipynb ../examples/demo_mock_ensemble_realistic.ipynb ../examples/demo_coordinate_system.ipynb +../examples/demo_kmeans_jackknife.ipynb #../examples/other_compare_NFW_critical_massdef.ipynb #../examples/other_flat_sky_limitations.ipynb diff --git a/environment.yml b/environment.yml index a2b42be74..4f6c3988f 100644 --- a/environment.yml +++ b/environment.yml @@ -27,4 +27,5 @@ dependencies: - pip: - black-configparser - qp-prob @ git+https://github.com/LSSTDESC/qp.git - - cluster-toolkit @ git+https://github.com/tmcclintock/cluster_toolkit.git \ No newline at end of file + - cluster-toolkit @ git+https://github.com/tmcclintock/cluster_toolkit.git + - kmeans_radec @ git+https://github.com/esheldon/kmeans_radec.git \ No newline at end of file diff --git a/examples/demo_kmeans_jackknife.ipynb b/examples/demo_kmeans_jackknife.ipynb new file mode 100644 index 000000000..4e342b2b5 --- /dev/null +++ b/examples/demo_kmeans_jackknife.ipynb @@ -0,0 +1,302 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a118fddc", + "metadata": {}, + "source": [ + "_Authored by Kabelo Tsiane, 2026-10-19_\n", + "\n", + "# Jackknife covariance with HEALPix and k-means regions\n", + "\n", + "`ClusterEnsemble.compute_jackknife_covariance` estimates the covariance of a stacked profile by dividing the sky into regions and removing each region that contains clusters in turn. The regions can be\n", + "\n", + "- **HEALPix pixels** (default), set by `n_side`: a fixed grid, the same for every ensemble, but the pixels do not follow the footprint of the clusters;\n", + "- **k-means regions**, set by `n_regions`: regions of similar numbers of clusters that follow the footprint, computed from the cluster positions with `clmm.utils.compute_kmeans_centers`;\n", + "- **predefined k-means regions**, set by `centers`: the same regions for several ensembles, e.g. for different redshift or richness bins.\n", + "\n", + "Computing k-means regions requires the optional [kmeans_radec](https://github.com/esheldon/kmeans_radec) package:\n", + "```\n", + "pip install git+https://github.com/esheldon/kmeans_radec.git\n", + "```\n", + "Using predefined `centers` does not." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dcbbb059", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "import healpy\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import clmm\n", + "from clmm import ClusterEnsemble, Cosmology, GalaxyCluster\n", + "from clmm.support import mock_data as mock\n", + "\n", + "clmm.__version__" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1ff3ff2f", + "metadata": {}, + "outputs": [], + "source": [ + "np.random.seed(11)\n", + "warnings.simplefilter(\"ignore\") # silence warnings from the mock catalogs" + ] + }, + { + "cell_type": "markdown", + "id": "4c4d4068", + "metadata": {}, + "source": [ + "## Mock cluster ensemble\n", + "\n", + "We draw `n_clusters` clusters uniformly in log-mass and redshift, spread over a 40 x 40 deg patch of sky, and generate a background galaxy catalog for each of them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "016ba2bb", + "metadata": {}, + "outputs": [], + "source": [ + "cosmo = Cosmology(H0=71.0, Omega_dm0=0.265 - 0.0448, Omega_b0=0.0448, Omega_k0=0.0)\n", + "\n", + "n_clusters = 200\n", + "cluster_m = 10 ** np.random.uniform(14.0, 14.5, n_clusters)\n", + "cluster_z = np.random.uniform(0.2, 0.4, n_clusters)\n", + "concentration = np.full(n_clusters, 4.0)\n", + "\n", + "# uniform on the sphere in 0 < ra < 40 deg, -20 < dec < 20 deg\n", + "cluster_ra = np.random.uniform(0.0, 40.0, n_clusters)\n", + "sin_dec = np.sin(np.radians(20.0))\n", + "cluster_dec = np.degrees(np.arcsin(np.random.uniform(-sin_dec, sin_dec, n_clusters)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "89b944e7", + "metadata": {}, + "outputs": [], + "source": [ + "gclist = []\n", + "for m, z, c, ra, dec in zip(cluster_m, cluster_z, concentration, cluster_ra, cluster_dec):\n", + " galcat = mock.generate_galaxy_catalog(\n", + " m,\n", + " z,\n", + " c,\n", + " cosmo,\n", + " \"chang13\",\n", + " delta_so=200,\n", + " massdef=\"critical\",\n", + " halo_profile_model=\"nfw\",\n", + " zsrc_min=z + 0.1,\n", + " zsrc_max=3.0,\n", + " field_size=10.0,\n", + " shapenoise=0.25,\n", + " ngal_density=5,\n", + " cluster_ra=ra,\n", + " cluster_dec=dec,\n", + " )\n", + " cl = GalaxyCluster(f\"mock_{len(gclist)}\", ra, dec, z, galcat)\n", + " cl.compute_tangential_and_cross_components(add=True)\n", + " cl.compute_galaxy_weights(cosmo=cosmo, add=True)\n", + " gclist.append(cl)" + ] + }, + { + "cell_type": "markdown", + "id": "22105822", + "metadata": {}, + "source": [ + "Individual and stacked tangential shear profiles:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "53ca2ec7", + "metadata": {}, + "outputs": [], + "source": [ + "bins = clmm.dataops.make_bins(0.3, 5.0, nbins=6, method=\"evenlog10width\")\n", + "profile_args = dict(bins=bins, bin_units=\"Mpc\", cosmo=cosmo)\n", + "\n", + "ensemble = ClusterEnsemble(\"all\", gclist, **profile_args)\n", + "ensemble.make_stacked_radial_profile()\n", + "ensemble.data[:3]" + ] + }, + { + "cell_type": "markdown", + "id": "7d4fb11e", + "metadata": {}, + "source": [ + "## HEALPix regions\n", + "\n", + "With `n_side=8` the pixels are about 7 deg across. Only the pixels that contain clusters are used." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ae0120a5", + "metadata": {}, + "outputs": [], + "source": [ + "n_side = 8\n", + "ensemble.compute_jackknife_covariance(n_side=n_side)\n", + "cov_healpix = ensemble.cov[\"tan_jk\"]\n", + "\n", + "pixels = healpy.ang2pix(n_side, ensemble[\"ra\"], ensemble[\"dec\"], nest=True, lonlat=True)\n", + "n_healpix = np.unique(pixels).size\n", + "print(f\"{n_healpix} occupied HEALPix pixels\")" + ] + }, + { + "cell_type": "markdown", + "id": "47d09b55", + "metadata": {}, + "source": [ + "## k-means regions\n", + "\n", + "For a fair comparison we use as many k-means regions as occupied HEALPix pixels. `seed` makes the regions reproducible." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bbcc4ee5", + "metadata": {}, + "outputs": [], + "source": [ + "ensemble.compute_jackknife_covariance(n_regions=n_healpix, seed=11)\n", + "cov_kmeans = ensemble.cov[\"tan_jk\"]\n", + "\n", + "# the same regions, to plot them\n", + "centers = clmm.utils.compute_kmeans_centers(ensemble[\"ra\"], ensemble[\"dec\"], n_healpix, seed=11)\n", + "labels = clmm.utils.assign_kmeans_regions(ensemble[\"ra\"], ensemble[\"dec\"], centers)" + ] + }, + { + "cell_type": "markdown", + "id": "9f0829c3", + "metadata": {}, + "source": [ + "Number of clusters in each region. The HEALPix grid is fixed, so its pixels can cut through groups of clusters or be only partly filled at the edges of the footprint, while the k-means regions are built from the cluster positions. Here the clusters are spread uniformly, so both are similar; the difference grows for an irregular footprint:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d740317a", + "metadata": {}, + "outputs": [], + "source": [ + "fig, axes = plt.subplots(1, 3, figsize=(15, 4.5))\n", + "for ax, regions, title in zip(axes[:2], (pixels, labels), (\"HEALPix pixels\", \"k-means regions\")):\n", + " _, regions = np.unique(regions, return_inverse=True)\n", + " ax.scatter(ensemble[\"ra\"], ensemble[\"dec\"], c=regions, cmap=\"tab20\", s=15)\n", + " ax.set_xlabel(\"ra [deg]\")\n", + " ax.set_ylabel(\"dec [deg]\")\n", + " ax.set_title(title)\n", + "axes[1].scatter(*centers.T, marker=\"x\", c=\"k\", label=\"centers\")\n", + "axes[1].legend()\n", + "\n", + "axes[2].hist(\n", + " [np.unique(pixels, return_counts=True)[1], np.bincount(labels)],\n", + " bins=np.arange(0, 30, 2),\n", + " label=[\"HEALPix\", \"k-means\"],\n", + ")\n", + "axes[2].set_xlabel(\"clusters per region\")\n", + "axes[2].legend()\n", + "fig.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "b2018c1f", + "metadata": {}, + "source": [ + "Errors on the stacked tangential shear from both sets of regions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "da82c1b2", + "metadata": {}, + "outputs": [], + "source": [ + "radius = ensemble.stacked_data[\"radius\"]\n", + "fig, ax = plt.subplots()\n", + "ax.plot(radius, np.sqrt(np.diag(cov_healpix)), \"o-\", label=\"HEALPix\")\n", + "ax.plot(radius, np.sqrt(np.diag(cov_kmeans)), \"s-\", label=\"k-means\")\n", + "ax.set_xscale(\"log\")\n", + "ax.set_xlabel(\"R [Mpc]\")\n", + "ax.set_ylabel(r\"$\\sigma(g_t)$\")\n", + "ax.legend()" + ] + }, + { + "cell_type": "markdown", + "id": "77910a12", + "metadata": {}, + "source": [ + "## Shared regions for several ensembles\n", + "\n", + "When an ensemble is split into bins (here in redshift), the k-means regions can be computed once from all the clusters and passed to every bin with `centers`. Every bin then uses the same regions, as with HEALPix pixels, and the jackknife samples of different bins remove the same part of the sky. Regions with no clusters of a bin are skipped." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "512ed1d0", + "metadata": {}, + "outputs": [], + "source": [ + "for z_min, z_max in [(0.2, 0.3), (0.3, 0.4)]:\n", + " in_bin = [gc for gc in gclist if z_min <= gc.z < z_max]\n", + " ensemble_bin = ClusterEnsemble(f\"z{z_min}-{z_max}\", in_bin, **profile_args)\n", + " ensemble_bin.make_stacked_radial_profile()\n", + " ensemble_bin.compute_jackknife_covariance(centers=centers)\n", + " n_occupied = np.unique(\n", + " clmm.utils.assign_kmeans_regions(ensemble_bin[\"ra\"], ensemble_bin[\"dec\"], centers)\n", + " ).size\n", + " print(f\"z {z_min}-{z_max}: {len(in_bin)} clusters in {n_occupied} of {len(centers)} regions\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "NERSC Python", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.15" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index b3e036a01..27485a513 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -88,6 +88,7 @@ ignored-modules = [ "gi.repository.NumCosmo", "gi.repository.NumCosmoMath", "healpy", + "kmeans_radec", "matplotlib", "matplotlib.pyplot", "numpy", diff --git a/tests/test_clusterensemble.py b/tests/test_clusterensemble.py index 63c3d3958..63ebb6732 100644 --- a/tests/test_clusterensemble.py +++ b/tests/test_clusterensemble.py @@ -221,3 +221,43 @@ def test_covariance(): assert_equal(ce.cov["cross_sc"].shape, ce2.cov["cross_sc"].shape) assert_equal(ce.cov["cross_bs"].shape, ce2.cov["cross_bs"].shape) assert_equal(ce.cov["cross_jk"].shape, ce2.cov["cross_jk"].shape) + + +def test_jackknife_regions(): + """test the healpix and k-means regions of the jackknife covariance""" + rng = np.random.default_rng(0) + n_clusters, n_bins = 40, 5 + ce = ClusterEnsemble("jk") + ce.data = GCData( + { + "id": np.arange(n_clusters), + "ra": rng.uniform(0.0, 30.0, n_clusters), + "dec": rng.uniform(-10.0, 10.0, n_clusters), + "z": np.full(n_clusters, 0.3), + "radius": np.tile(np.arange(1.0, n_bins + 1.0), (n_clusters, 1)), + "gt": rng.normal(size=(n_clusters, n_bins)), + "gx": rng.normal(size=(n_clusters, n_bins)), + "W_l": np.ones((n_clusters, n_bins)), + } + ) + + # one region per cluster: k-means and fine healpix pixels both give leave-one-out + ce.compute_jackknife_covariance(n_side=1024) + tan_jk_healpix = ce.cov["tan_jk"] + ce.compute_jackknife_covariance(centers=np.column_stack([ce.data["ra"], ce.data["dec"]])) + assert_allclose(ce.cov["tan_jk"], tan_jk_healpix, **TOLERANCE) + + # k-means regions computed from the clusters, reproducible with a seed + ce.compute_jackknife_covariance(n_regions=8, seed=11) + tan_jk_kmeans = ce.cov["tan_jk"] + assert_equal(tan_jk_kmeans.shape, (n_bins, n_bins)) + assert_equal(ce.cov["cross_jk"].shape, (n_bins, n_bins)) + ce.compute_jackknife_covariance(n_regions=8, seed=11) + assert_equal(ce.cov["tan_jk"], tan_jk_kmeans) + + # shared regions with no clusters in them are skipped + centers = clmm.utils.compute_kmeans_centers(ce.data["ra"], ce.data["dec"], 8, seed=11) + ce.compute_jackknife_covariance(centers=np.vstack([centers, [[200.0, 50.0]]])) + assert_allclose(ce.cov["tan_jk"], tan_jk_kmeans, **TOLERANCE) + + assert_raises(ValueError, ce.compute_jackknife_covariance, n_regions=8, centers=centers) diff --git a/tests/test_utils.py b/tests/test_utils.py index 2bd7aeb4d..6a7059f6b 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,5 +1,7 @@ # pylint: disable=no-member, protected-access """ Tests for utils.py """ +import sys + import numpy as np from numpy.testing import assert_raises, assert_allclose, assert_equal from scipy.integrate import quad @@ -634,3 +636,39 @@ def integrand2(z_i, z_inf=z_inf, z_cl=z_cl, cosmo=cosmo): ), **TOLERANCE ) + + +def test_assign_kmeans_regions(): + """Test the assignment of points to the nearest region center""" + centers = [[1.0, 0.0], [180.0, 0.0], [0.0, 89.0]] + # points near ra=0 on both sides of the wrap, the opposite side, and the pole + ra = [359.5, 0.5, 170.0, 100.0] + dec = [0.0, 0.0, 10.0, 85.0] + assert_equal(utils.assign_kmeans_regions(ra, dec, centers), [0, 0, 1, 2]) + # a single center gets every point + assert_equal(utils.assign_kmeans_regions(ra, dec, [10.0, 0.0]), [0, 0, 0, 0]) + + +def test_compute_kmeans_centers(monkeypatch): + """Test the k-means region centers""" + rng = np.random.default_rng(1) + ra = rng.uniform(0.0, 360.0, 500) + dec = np.degrees(np.arcsin(rng.uniform(-1.0, 1.0, 500))) + + centers = utils.compute_kmeans_centers(ra, dec, 10, seed=11) + assert_equal(centers.shape, (10, 2)) + # every region gets points + assert_equal(np.unique(utils.assign_kmeans_regions(ra, dec, centers)).size, 10) + # reproducible with a seed, without changing the global random state + state = np.random.get_state()[1].copy() + assert_equal(utils.compute_kmeans_centers(ra, dec, 10, seed=11), centers) + assert_equal(np.random.get_state()[1], state) + # as many regions as points + assert_equal(utils.compute_kmeans_centers(ra[:5], dec[:5], 5, seed=11).shape, (5, 2)) + + assert_raises(ValueError, utils.compute_kmeans_centers, ra, dec, 0) + assert_raises(ValueError, utils.compute_kmeans_centers, ra[:5], dec[:5], 6) + + # kmeans_radec is optional + monkeypatch.setitem(sys.modules, "kmeans_radec", None) + assert_raises(ImportError, utils.compute_kmeans_centers, ra, dec, 10)