Issue/711/jk cov kmeans

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 <a name="dev_install"></a>
 
 Here we provide a quick guide for a developer instalation, using `environment.yml` to create a suitable conda environment and install all necessary packages.

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 <a name="backends"></a>
 
 For the theoretical predictions of the signal, CLMM relies on existing libraries and **at least one of the following must be installed as well**:

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,25 +354,45 @@ 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],
                 [self[tan_component][mask], self[cross_component][mask]],
             )[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)

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,

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)

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

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
+    - cluster-toolkit @ git+https://github.com/tmcclintock/cluster_toolkit.git
+    - kmeans_radec @ git+https://github.com/esheldon/kmeans_radec.git