From 774a22a759f31bad670092c910ef4e96f484a9e9 Mon Sep 17 00:00:00 2001
From: Kabelo Tsiane <ktsiane00@gmail.com>
Date: Thu, 30 Jul 2026 12:45:30 -0500
Subject: [PATCH 1/4] first commit, code runs

---
 clmm/gt_CLMM_versus_treecorr.py | 543 ++++++++++++++++++++++++++++++++
 clmm/hp_vs_kmeans.py            |  62 ++++
 clmm/jk_kmeans.py               |  40 +++
 3 files changed, 645 insertions(+)
 create mode 100644 clmm/gt_CLMM_versus_treecorr.py
 create mode 100644 clmm/hp_vs_kmeans.py
 create mode 100644 clmm/jk_kmeans.py

diff --git a/clmm/gt_CLMM_versus_treecorr.py b/clmm/gt_CLMM_versus_treecorr.py
new file mode 100644
index 000000000..7a06a2853
--- /dev/null
+++ b/clmm/gt_CLMM_versus_treecorr.py
@@ -0,0 +1,543 @@
+"""Compare TreeCorr to CLMM for stacked reduced tangential shear measurement.
+
+This script compares, in a very ideal setup, the computation of the stacked
+reduced shear profile either using CLMM or using TreeCorr to cross-correlate
+cluster positions with background source shears. The mock data is generated
+using CLMM.
+"""
+import os
+import warnings
+
+import healpy as hp
+import numpy as np
+import pyccl as ccl
+import matplotlib.pyplot as plt
+import scipy
+import treecorr
+from astropy.table import Table, vstack
+
+import clmm
+from clmm import ClusterEnsemble, Cosmology
+from clmm.support import mock_data as mock
+from clmm.utils.redshift_distributions import desc_srd
+
+from jk_kmeans import jk_cov
+
+PLOTS_DIR = "./plots/"
+os.makedirs(PLOTS_DIR, exist_ok=True)
+
+# For reproducibility:
+np.random.seed(11)
+
+# -----------------------------------------------------------------------------
+# Draw pairs of mass and redshift according to a mass function
+# -----------------------------------------------------------------------------
+
+# Mass function - CCL implementation
+cosmo = ccl.Cosmology(
+    Omega_c=0.265,
+    Omega_b=0.0448,
+    h=0.71,
+    sigma8=0.8,
+    n_s=0.96,
+    Neff=3.04,
+    m_nu=1.0e-05,
+    mass_split="single",
+)
+hmd_200c = ccl.halos.MassDef(200, "critical")
+
+
+# For a different multiplicty function, the user must change this function below
+def tinker08_ccl(logm, z):
+    mass = 10 ** (logm)
+    hmf_200c = ccl.halos.MassFuncTinker08(mass_def=hmd_200c)
+    nm = hmf_200c(cosmo, mass, 1.0 / (1 + z))
+    return nm  # dn/dlog10M
+
+
+# Computing the volume element
+def dV_over_dOmega_dz(z):
+    a = 1.0 / (1.0 + z)
+    da = ccl.background.angular_diameter_distance(cosmo, a)
+    E = ccl.background.h_over_h0(cosmo, a)
+    return ((1.0 + z) ** 2) * (da**2) * ccl.physical_constants.CLIGHT_HMPC / cosmo["h"] / E
+
+
+def pdf_tinker08_ccl(logm, z):
+    return tinker08_ccl(logm, z) * dV_over_dOmega_dz(z)
+
+
+# Acceptance-rejection method to sample (M,z) from the mass function
+def bivariate_draw(
+    pdf, N=1000, logm_min=14.0, logm_max=15.0, zmin=0.01, zmax=1, Ngrid_m=30, Ngrid_z=30
+):
+    """
+    Uses the rejection method for generating random numbers derived from an arbitrary
+    probability distribution.
+
+    Parameters
+    ----------
+    pdf : func
+      2d distribution function to sample
+    N : int
+      number of points to generate
+    log_min,logm_max : float
+      log10 mass range
+    zmin,zmax : float
+      redshift range
+
+    Returns:
+    --------
+    ran_logm : list
+        accepted logm values
+    ran_z : list
+        accepted redshift values
+    acceptance : float
+        acceptance ratio of the method
+    """
+
+    # maximum value of the pdf over the mass and redshift space.
+    # the pdf is not monotonous in mass and redshift, so we need
+    # to find the maximum numerically. Here we scan the space
+    # with a regular grid and use the maximum value.
+    # Accuracy of the results depends on Ngrid_m and Ngrid_z
+    # This should probably be improved
+
+    logM_arr = np.linspace(logm_min, logm_max, Ngrid_m)
+    z_arr = np.logspace(np.log10(zmin), np.log10(zmax), Ngrid_z)
+    p = []
+    for logM in logM_arr:
+        for z in z_arr:
+            p.append(pdf(logM, z))
+    pmax = np.max(p)
+    pmin = np.min(p)
+    # Counters
+    naccept = 0
+    ntrial = 0
+
+    # Keeps drawing until N points are accepted
+    ran_logm = []  # output list of random numbers
+    ran_z = []  # output list of random numbers
+    while naccept < N:
+        # draw (logm,z) from uniform distribution
+        # draw p from uniform distribution
+        logm = np.random.uniform(logm_min, logm_max)  # x'
+        z = np.random.uniform(zmin, zmax)  # x'
+        p = np.random.uniform(0.0, pmax)  # y'
+
+        if p < pdf(logm, z):
+            # keep the point
+            ran_logm.append(logm)
+            ran_z.append(z)
+            naccept = naccept + 1
+        ntrial = ntrial + 1
+
+    ran_logm = np.asarray(ran_logm)
+    ran_z = np.asarray(ran_z)
+
+    acceptance = float(N / ntrial)
+    print(f"acceptance = {acceptance}")
+    return ran_logm, ran_z, acceptance
+
+
+# Draw N=30 (M,z) pairs
+
+# Define the numbers of pairs to draw and the mass and redshift ranges
+N = 100
+logm_min = 14.5
+logm_max = 14.8
+zmin = 0.2
+zmax = 0.3
+
+pdf = pdf_tinker08_ccl
+
+# Normalisation factor for the pdf to be used later
+norm = (
+    1.0
+    / scipy.integrate.dblquad(
+        pdf, zmin, zmax, lambda x: logm_min, lambda x: logm_max, epsrel=1.0e-4
+    )[0]
+)
+
+# Random draw
+ran_logm, ran_z, acceptance = bivariate_draw(
+    pdf, N=N, logm_min=logm_min, logm_max=logm_max, zmin=zmin, zmax=zmax
+)
+
+# -----------------------------------------------------------------------------
+# Generate a cluster catalog and associated source catalogs
+# -----------------------------------------------------------------------------
+
+# Cluster ensemble masses, redshifts and positions
+cluster_m = 10 ** ran_logm
+cluster_z = ran_z
+
+# Concentration CCL object to compute the theoretical concentration
+conc_obj = ccl.halos.ConcentrationDuffy08(mass_def=hmd_200c)
+conc_list = []
+for number in range(0, len(cluster_m)):
+    a = 1.0 / (1.0 + (cluster_z[number]))
+    # mean value of the concentration for that cluster
+    lnc_mean = np.log(conc_obj(cosmo, M=(cluster_m[number]), a=a))
+    # random draw of actual concentration from normal distribution around lnc_mean, with a 0.14 scatter
+    lnc = np.random.normal(lnc_mean, 0.14)
+    conc_list.append(np.exp(lnc))
+
+conc_list = np.array(conc_list)
+
+# randomly draw cluster positions over the full sky
+ra = np.random.random(N) * 360  # from 0 to 360 deg
+sindec = np.random.random(N) * 2 - 1
+dec = np.arcsin(sindec) * 180 / np.pi  # from -90 to 90 deg
+
+print(len(cluster_m), N)
+
+# Background galaxy catalog generation
+#
+# For each cluster of the ensemble, we use `mock_data` to generate a background
+# galaxy catalog and store the results in a `GalaxyCluster` object. Note that:
+# - The cluster density profiles follow the NFW parametrisation
+# - The source redshifts follow the Chang et al. distribution and have associated pdfs
+# - The shapes include shape noise and shape measurement errors
+# - Background galaxy catalogs are independent, even if the clusters are close
+#   (i.e., no common galaxy between two catalogs).
+# - For each cluster we then compute
+#     - the tangential and cross Delta_Sigma for each background galaxy
+#     - the weights `w_ls` to be used to compute the corresponding radial profiles
+#
+# The cluster objects are then stored in `gclist`.
+
+warnings.simplefilter(
+    "ignore"
+)  # just to prevent warning print out when looping over the cluster ensemble below
+
+gclist = []
+tables = []
+# number of galaxies in each cluster field (alternatively, can use the galaxy density instead)
+ngal_density = 10
+cosmo_clmm = Cosmology(H0=71.0, Omega_dm0=0.265 - 0.0448, Omega_b0=0.0448, Omega_k0=0.0)
+cosmo_clmm.set_be_cosmo(cosmo)
+gal_args = dict(
+    cosmo=cosmo_clmm,
+    zsrc="chang13",
+    delta_so=200,
+    massdef="critical",
+    halo_profile_model="nfw",
+    zsrc_max=2.0,
+    field_size=16.0,
+    shapenoise=0.04,
+    photoz_sigma_unscaled=0.02,
+    ngal_density=ngal_density,
+    mean_e_err=0.1,
+)
+
+for i in range(N):
+    # generate background galaxy catalog for cluster i
+    cl = clmm.GalaxyCluster(
+        f"mock_cluster_{i:04}",
+        ra[i],
+        dec[i],
+        cluster_z[i],
+        galcat=mock.generate_galaxy_catalog(
+            cluster_m[i],
+            cluster_z[i],
+            conc_list[i],
+            cluster_ra=ra[i],
+            cluster_dec=dec[i],
+            zsrc_min=0.35,  # cluster zmax = 0.3 in this example
+            **gal_args,
+        ),
+    )
+
+    # compute DeltaSigma for each background galaxy
+    cl.compute_tangential_and_cross_components(
+        shape_component1="e1",
+        shape_component2="e2",
+        tan_component="g_t",
+        cross_component="g_x",
+        cosmo=cosmo_clmm,
+        is_deltasigma=False,
+        use_pdz=True,
+    )
+
+    # compute the weights to be used to bluid the DeltaSigma radial profiles
+    cl.compute_galaxy_weights(
+        use_pdz=True,
+        use_shape_noise=True,
+        shape_component1="e1",
+        shape_component2="e2",
+        use_shape_error=True,
+        shape_component1_err="e_err",
+        shape_component2_err="e_err",
+        weight_name="w_ls",
+        cosmo=cosmo_clmm,
+        is_deltasigma=False,
+        add=True,
+    )
+
+    # append the cluster in the list
+    gclist.append(cl)
+    tables.append(cl.galcat)
+
+# Stack as plain astropy Tables: GCData's metadata guard for "coordinate_system"
+# rejects vstack's metadata merge even when all tables share the same value.
+galcat_tot = vstack([Table(t) for t in tables])
+
+# -----------------------------------------------------------------------------
+# Create ClusterEnsemble object and estimation of individual excess surface
+# density profiles
+# -----------------------------------------------------------------------------
+
+bins = np.logspace(np.log10(0.5), np.log10(60), 12)
+ensemble_id = 1
+clusterensemble = ClusterEnsemble(ensemble_id)
+for cluster in gclist:
+    clusterensemble.make_individual_radial_profile(
+        galaxycluster=cluster,
+        tan_component_in="g_t",
+        cross_component_in="g_x",
+        tan_component_out="g_t",
+        cross_component_out="g_x",
+        weights_in="w_ls",
+        weights_out="W_l",
+        bins=bins,
+        bin_units="arcmin",
+        cosmo=cosmo_clmm,
+    )
+
+# Stacked profile of the cluster ensemble
+clusterensemble.make_stacked_radial_profile(tan_component="g_t", cross_component="g_x", weights="W_l")
+
+# Jackknife covariance of the stack between radial bins
+jk_n_side = 16
+clusterensemble.compute_jackknife_covariance(
+    n_side=jk_n_side, tan_component="g_t", cross_component="g_x"
+)
+njk = 10
+kmeans_gt_cov, kmeans_gx_cov, km_centers = jk_cov(clusterensemble, njk=njk)
+# -----------------------------------------------------------------------------
+# Visualizing the stacked profiles
+#
+# In the figure below, we plot:
+# - the individual g_t profiles of the clusters (light blue)
+# - the stacked signal (red symbols)
+# -----------------------------------------------------------------------------
+
+moo = clmm.Modeling(massdef="critical", delta_mdef=200, halo_profile_model="nfw")
+moo.set_cosmo(cosmo_clmm)
+# Average values of mass and concentration of the ensemble to be used below
+# to overplot the model on the stacked profile
+moo.set_concentration(conc_list.mean())
+moo.set_mass(cluster_m.mean())
+
+r_stack, gt_stack, gx_stack = (clusterensemble.stacked_data[c] for c in ("radius", "g_t", "g_x"))
+plt.rcParams["axes.linewidth"] = 2
+fig, axs = plt.subplots(1, 2, figsize=(17, 6))
+
+err_gt = clusterensemble.cov["tan_jk"].diagonal() ** 0.5
+err_gx = clusterensemble.cov["cross_jk"].diagonal() ** 0.5
+km_err_gt = np.sqrt(np.diag(kmeans_gt_cov))
+km_err_gx = np.sqrt(np.diag(kmeans_gx_cov))
+axs[0].errorbar(
+    r_stack,
+    gt_stack,
+    err_gt,
+    markersize=5,
+    c="r",
+    fmt="o",
+    capsize=10,
+    elinewidth=1,
+    zorder=1000,
+    alpha=1,
+    label="Stack",
+)
+axs[0].errorbar(
+    r_stack,
+    gt_stack,
+    km_err_gt,
+    markersize=5,
+    c="k",
+    fmt="o",
+    capsize=10,
+    elinewidth=1,
+    zorder=1000,
+    alpha=1,
+    label="Stack",
+)
+axs[1].errorbar(
+    r_stack,
+    gx_stack,
+    err_gx,
+    markersize=5,
+    c="r",
+    fmt="o",
+    capsize=10,
+    elinewidth=1,
+    zorder=1000,
+    alpha=1,
+    label="Stack",
+)
+axs[1].errorbar(
+    r_stack,
+    gx_stack,
+    km_err_gx,
+    markersize=5,
+    c="k",
+    fmt="o",
+    capsize=10,
+    elinewidth=1,
+    zorder=1000,
+    alpha=1,
+    label="Stack",
+)
+axs[0].set_xscale("log")
+axs[0].set_yscale("log")
+axs[1].set_xscale("log")
+
+
+for i in range(N):
+    axs[0].plot(
+        clusterensemble.data["radius"][i],
+        clusterensemble.data["g_t"][i],
+        color="cyan",
+        label="Individual",
+        alpha=1,
+        linewidth=1,
+    )
+    axs[1].plot(
+        clusterensemble.data["radius"][i],
+        clusterensemble.data["g_x"][i],
+        color="cyan",
+        label="Individual",
+        alpha=1,
+        linewidth=1,
+    )
+    if i == 0:
+        axs[0].legend(frameon=False, fontsize=15)
+        axs[1].legend(frameon=False, fontsize=15)
+
+axs[0].set_xlabel("separation [arcmin]", fontsize=20)
+axs[1].set_xlabel("separation [arcmin]", fontsize=20)
+axs[0].tick_params(axis="both", which="major", labelsize=18)
+axs[1].tick_params(axis="both", which="major", labelsize=18)
+axs[0].set_ylabel(r"$g_+$", fontsize=20)
+axs[1].set_ylabel(r"$g_\times$", fontsize=20)
+axs[0].set_title(r"Tangential", fontsize=20)
+axs[1].set_title(r"Cross", fontsize=20)
+
+for ax in axs:
+    ax.minorticks_on()
+    ax.grid(lw=0.5)
+    ax.grid(which="minor", lw=0.1)
+fig.tight_layout()
+fig.savefig(os.path.join(PLOTS_DIR, "individual_and_stacked_profiles.png"))
+plt.close(fig)
+
+# -----------------------------------------------------------------------------
+# Saving/Loading ClusterEnsemble
+#
+# The `ClusterEnsemble` object also has an option for saving/loading:
+# -----------------------------------------------------------------------------
+
+clusterensemble.save("ce.pkl")
+
+# -----------------------------------------------------------------------------
+# Compute reduced shear profile with TreeCorr and compare to CLMM stacked
+# profile
+# -----------------------------------------------------------------------------
+# Cluster catalog
+cat_cluster = treecorr.Catalog(
+    ra=clusterensemble["ra"],
+    dec=clusterensemble["dec"],
+    ra_units="deg",
+    dec_units="deg",
+    patch_centers=km_centers
+)
+# Source catalog
+cat_source = treecorr.Catalog(
+    ra=galcat_tot["ra"],
+    dec=galcat_tot["dec"],
+    g1=galcat_tot["e1"],
+    g2=galcat_tot["e2"],
+    ra_units="deg",
+    dec_units="deg",
+    w=galcat_tot["w_ls"],
+    patch_centers=cat_cluster.patch_centers
+)
+
+
+
+ng = treecorr.NGCorrelation(nbins=11, min_sep=0.5, max_sep=60, sep_units="arcmin", bin_type="Log")
+ng.process(cat_cluster, cat_source)
+theta = ng.meanr  # average angular separation in each bin
+cov_jk = ng.estimate_cov('jackknife')
+var_jk = np.diag(cov_jk)  
+print(cosmo_clmm.be_cosmo)
+
+da = cosmo.angular_diameter_distance(1.0 / (1.0 + np.mean(clusterensemble["z"])))
+da_clmm = cosmo_clmm.eval_da(np.mean(clusterensemble["z"]))
+print(da, da_clmm)
+
+arcmin_to_Mpc = np.pi / (60.0 * 180) * da_clmm
+
+fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 5))
+ax.errorbar(
+    theta,
+    ng.xi,
+    np.sqrt(var_jk),
+    color="blue",
+    marker="x",
+    markersize=15,
+    linestyle="",
+    label="Treecorr - cluster positions x source shears",
+)
+
+ax.errorbar(
+    r_stack,
+    gt_stack,
+    err_gt,
+    markersize=5,
+    c="r",
+    fmt=".",
+    capsize=10,
+    elinewidth=1,
+    zorder=1000,
+    alpha=1,
+    label="CLMM - stacked reduced rangential shear",
+)
+
+ax.plot(
+    clusterensemble.data["radius"][0],
+    moo.eval_reduced_tangential_shear(
+        clusterensemble.data["radius"][0] * arcmin_to_Mpc,
+        cluster_z.mean(),
+        z_src=desc_srd,
+        z_src_info="distribution",
+    ),
+    "--k",
+    linewidth=1,
+    label="CLMM prediction for stack 'mean' cluster",
+    zorder=100,
+)
+
+ax.set_xscale("log")
+ax.set_yscale("log")
+ax.set_ylim([-0.02, 0.15])
+ax.axhline(0, color="k", ls=":")
+ax.set_xlabel("separation [arcmin]")
+ax.set_ylabel("$\\widehat{g_t}$")
+ax.legend()
+fig.savefig(os.path.join(PLOTS_DIR, "treecorr_vs_clmm.png"))
+plt.close(fig)
+
+fig, ax = plt.subplots()
+ax.scatter(err_gt, np.sqrt(var_jk), label='healpy')
+ax.scatter(km_err_gt, np.sqrt(var_jk), label='kmeans')
+ax.plot([0, 0.007], [0, 0.007], "k:")
+ax.set_xlabel("Error from CLMM")
+ax.set_ylabel("Error from TreeCorr")
+ax.set_xscale('log')
+ax.set_yscale('log')
+plt.legend()
+fig.savefig(os.path.join(PLOTS_DIR, "error_comparison.png"))
+plt.close(fig)
diff --git a/clmm/hp_vs_kmeans.py b/clmm/hp_vs_kmeans.py
new file mode 100644
index 000000000..74e91e8aa
--- /dev/null
+++ b/clmm/hp_vs_kmeans.py
@@ -0,0 +1,62 @@
+import healpy as hp
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+from kmeans_radec.test import plot_centers
+from kmeans_radec import KMeans, kmeans_sample, find_nearest
+
+PLOTS_DIR = "./plots/"
+os.makedirs(PLOTS_DIR, exist_ok=True)
+
+# randomly draw cluster positions over the full sky
+N = 30
+ra = np.random.random(N) * 360  # from 0 to 360 deg
+sindec = np.random.random(N) * 2 - 1
+dec = np.arcsin(sindec) * 180 / np.pi  # from -90 to 90 deg
+jk_n_side = 16
+# CLMM's jackknife regions are HEALPix pixels (not kmeans patches like TreeCorr),
+# and are not stored on the object, so recompute them here to visualize the full
+# extent of each region using healpy's own map plotting (cartview).
+jk_pixels = hp.ang2pix(
+    jk_n_side, ra, dec, nest=True, lonlat=True
+)
+jk_unique_pixels = np.unique(jk_pixels)
+jk_center_ra, jk_center_dec = hp.pix2ang(jk_n_side, jk_unique_pixels, nest=True, lonlat=True)
+jk_map = np.full(hp.nside2npix(jk_n_side), hp.UNSEEN)
+jk_map[jk_unique_pixels] = np.arange(jk_unique_pixels.size)
+
+fig = plt.figure(figsize=(12, 5))
+
+# Left panel: healpy manages its own projection axes internally, so instead of
+# handing it a matplotlib Axes, place it into a subplot slot of our figure via
+# fig= and sub= (same (nrows, ncols, index) convention as plt.subplot).
+hp.cartview(
+    jk_map,
+    fig=fig.number,
+    sub=(1, 2, 1),
+    nest=True,
+    cmap="tab20",
+    cbar=False,
+    title="Jackknife patches (HEALPix pixels)",
+)
+hp.graticule()
+hp.projscatter(ra, dec, lonlat=True, c="k", s=10)
+
+njk = 3 
+X = np.vstack((ra, dec)).T
+km = kmeans_sample(X, ncen=njk, method='slow')
+
+ax2= fig.add_subplot(1, 2, 2)
+plot_centers(km.centers, ax2, y_min=-90, y_max=90)
+ax2.scatter(ra, dec, c="r", s=30)
+ax2.set_title("kmeans_radec Voronoi cells")
+ax2.set_axis_off()
+
+fig.savefig(os.path.join(PLOTS_DIR, "hp_vs_kmeans.png"))
+plt.show()
+
+print('----- Distances ------')
+print(km.distances)
+labels = find_nearest(X, km.centers)
+print('--- labels ---')
+print(labels)
diff --git a/clmm/jk_kmeans.py b/clmm/jk_kmeans.py
new file mode 100644
index 000000000..95365f376
--- /dev/null
+++ b/clmm/jk_kmeans.py
@@ -0,0 +1,40 @@
+from kmeans_radec import kmeans_sample, find_nearest
+from clmm.dataops import make_stacked_radial_profile
+import numpy as np
+
+def radec_to_xyz(coords):
+    ra_rad = np.radians(coords[:, 0])
+    dec_rad = np.radians(coords[:, 1])
+    x = np.cos(dec_rad) * np.cos(ra_rad)
+    y = np.cos(dec_rad) * np.sin(ra_rad)
+    z = np.sin(dec_rad)
+    return np.column_stack([x, y, z])
+
+
+def jk_regions(X, njk):
+    km = kmeans_sample(X, ncen=njk)
+    labels = find_nearest(X, km.centers)
+    return labels, km.centers
+
+
+def jk_cov(clusterensemble, njk):
+    X = np.vstack((clusterensemble['ra'], clusterensemble['dec'])).T
+    labels, centers = jk_regions(X, njk=njk)
+    unique_labels = np.unique(labels)
+    gt_jack, gx_jack = [], []
+    for drop_region in unique_labels:
+        mask = np.isin(labels, drop_region, invert=True)
+        gt, gx = make_stacked_radial_profile(
+            clusterensemble['radius'][mask],
+            clusterensemble['W_l'][mask],
+            [clusterensemble['g_t'][mask], clusterensemble['g_x'][mask]])[1]
+        gt_jack.append(gt)
+        gx_jack.append(gx)
+    n_jack = unique_labels.size
+    coeff = (n_jack - 1)**2 / n_jack
+    tan_jk = coeff * np.cov(np.transpose(gt_jack), bias=False, ddof=0)
+    cross_jk = coeff * np.cov(np.transpose(gx_jack), bias=False, ddof=0)
+    cart_centers = radec_to_xyz(centers) 
+    return tan_jk, cross_jk, cart_centers
+
+

From d324666be6df008b230bff1c261b098acc11308b Mon Sep 17 00:00:00 2001
From: Kabelo Tsiane <ktsiane00@gmail.com>
Date: Thu, 13 Aug 2026 11:06:04 -0700
Subject: [PATCH 2/4] made mocks bigger, made comparison more direct

---
 clmm/gt_CLMM_versus_treecorr.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/clmm/gt_CLMM_versus_treecorr.py b/clmm/gt_CLMM_versus_treecorr.py
index 7a06a2853..3eee28828 100644
--- a/clmm/gt_CLMM_versus_treecorr.py
+++ b/clmm/gt_CLMM_versus_treecorr.py
@@ -21,7 +21,7 @@
 from clmm.support import mock_data as mock
 from clmm.utils.redshift_distributions import desc_srd
 
-from jk_kmeans import jk_cov
+from clmm.jk_kmeans import jk_cov
 
 PLOTS_DIR = "./plots/"
 os.makedirs(PLOTS_DIR, exist_ok=True)
@@ -143,9 +143,9 @@ def bivariate_draw(
 # Draw N=30 (M,z) pairs
 
 # Define the numbers of pairs to draw and the mass and redshift ranges
-N = 100
-logm_min = 14.5
-logm_max = 14.8
+N = 10000
+logm_min = 14.0
+logm_max = 15.0
 zmin = 0.2
 zmax = 0.3
 
@@ -541,3 +541,14 @@ def bivariate_draw(
 plt.legend()
 fig.savefig(os.path.join(PLOTS_DIR, "error_comparison.png"))
 plt.close(fig)
+
+fig, ax = plt.subplots()
+ax.scatter(km_err_gt, err_gt)
+ax.plot([0, 0.007], [0, 0.007], "k:")
+ax.set_xlabel("Error from kmeans")
+ax.set_ylabel("Error from healpy")
+ax.set_xscale('log')
+ax.set_yscale('log')
+plt.legend()
+fig.savefig(os.path.join(PLOTS_DIR, "kmeans_healpy_error_comparison.png"))
+plt.close(fig)

From 1c8cf948164999900dcbb8e456fb492790073d84 Mon Sep 17 00:00:00 2001
From: Kabelo Tsiane <ktsiane00@gmail.com>
Date: Mon, 19 Oct 2026 20:35:59 -0700
Subject: [PATCH 3/4] Added kmeans jackknife utils. converts sky coords to xyz,
 computes patch centers and assigns each point to a region of nearest center

---
 clmm/utils/__init__.py  |  1 +
 clmm/utils/jackknife.py | 92 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)
 create mode 100644 clmm/utils/jackknife.py

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)

From 3e672f83b802a3b63e323e191bc476b2d0647d9a Mon Sep 17 00:00:00 2001
From: Kabelo Tsiane <ktsiane00@gmail.com>
Date: Mon, 19 Oct 2026 22:48:13 -0700
Subject: [PATCH 4/4] adding tests for kmeans regions. removed prototype code,
 now a demo notebook. added kmeans_radec as optional dependency

---
 INSTALL.md                           |   7 +
 README.md                            |   4 +
 clmm/clusterensemble.py              |  49 ++-
 clmm/gt_CLMM_versus_treecorr.py      | 554 ---------------------------
 clmm/hp_vs_kmeans.py                 |  62 ---
 clmm/jk_kmeans.py                    |  40 --
 docs/doc-config.ini                  |   1 +
 environment.yml                      |   3 +-
 examples/demo_kmeans_jackknife.ipynb | 302 +++++++++++++++
 pyproject.toml                       |   1 +
 tests/test_clusterensemble.py        |  40 ++
 tests/test_utils.py                  |  38 ++
 12 files changed, 435 insertions(+), 666 deletions(-)
 delete mode 100644 clmm/gt_CLMM_versus_treecorr.py
 delete mode 100644 clmm/hp_vs_kmeans.py
 delete mode 100644 clmm/jk_kmeans.py
 create mode 100644 examples/demo_kmeans_jackknife.ipynb

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 <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.
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 <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**:
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/gt_CLMM_versus_treecorr.py b/clmm/gt_CLMM_versus_treecorr.py
deleted file mode 100644
index 3eee28828..000000000
--- a/clmm/gt_CLMM_versus_treecorr.py
+++ /dev/null
@@ -1,554 +0,0 @@
-"""Compare TreeCorr to CLMM for stacked reduced tangential shear measurement.
-
-This script compares, in a very ideal setup, the computation of the stacked
-reduced shear profile either using CLMM or using TreeCorr to cross-correlate
-cluster positions with background source shears. The mock data is generated
-using CLMM.
-"""
-import os
-import warnings
-
-import healpy as hp
-import numpy as np
-import pyccl as ccl
-import matplotlib.pyplot as plt
-import scipy
-import treecorr
-from astropy.table import Table, vstack
-
-import clmm
-from clmm import ClusterEnsemble, Cosmology
-from clmm.support import mock_data as mock
-from clmm.utils.redshift_distributions import desc_srd
-
-from clmm.jk_kmeans import jk_cov
-
-PLOTS_DIR = "./plots/"
-os.makedirs(PLOTS_DIR, exist_ok=True)
-
-# For reproducibility:
-np.random.seed(11)
-
-# -----------------------------------------------------------------------------
-# Draw pairs of mass and redshift according to a mass function
-# -----------------------------------------------------------------------------
-
-# Mass function - CCL implementation
-cosmo = ccl.Cosmology(
-    Omega_c=0.265,
-    Omega_b=0.0448,
-    h=0.71,
-    sigma8=0.8,
-    n_s=0.96,
-    Neff=3.04,
-    m_nu=1.0e-05,
-    mass_split="single",
-)
-hmd_200c = ccl.halos.MassDef(200, "critical")
-
-
-# For a different multiplicty function, the user must change this function below
-def tinker08_ccl(logm, z):
-    mass = 10 ** (logm)
-    hmf_200c = ccl.halos.MassFuncTinker08(mass_def=hmd_200c)
-    nm = hmf_200c(cosmo, mass, 1.0 / (1 + z))
-    return nm  # dn/dlog10M
-
-
-# Computing the volume element
-def dV_over_dOmega_dz(z):
-    a = 1.0 / (1.0 + z)
-    da = ccl.background.angular_diameter_distance(cosmo, a)
-    E = ccl.background.h_over_h0(cosmo, a)
-    return ((1.0 + z) ** 2) * (da**2) * ccl.physical_constants.CLIGHT_HMPC / cosmo["h"] / E
-
-
-def pdf_tinker08_ccl(logm, z):
-    return tinker08_ccl(logm, z) * dV_over_dOmega_dz(z)
-
-
-# Acceptance-rejection method to sample (M,z) from the mass function
-def bivariate_draw(
-    pdf, N=1000, logm_min=14.0, logm_max=15.0, zmin=0.01, zmax=1, Ngrid_m=30, Ngrid_z=30
-):
-    """
-    Uses the rejection method for generating random numbers derived from an arbitrary
-    probability distribution.
-
-    Parameters
-    ----------
-    pdf : func
-      2d distribution function to sample
-    N : int
-      number of points to generate
-    log_min,logm_max : float
-      log10 mass range
-    zmin,zmax : float
-      redshift range
-
-    Returns:
-    --------
-    ran_logm : list
-        accepted logm values
-    ran_z : list
-        accepted redshift values
-    acceptance : float
-        acceptance ratio of the method
-    """
-
-    # maximum value of the pdf over the mass and redshift space.
-    # the pdf is not monotonous in mass and redshift, so we need
-    # to find the maximum numerically. Here we scan the space
-    # with a regular grid and use the maximum value.
-    # Accuracy of the results depends on Ngrid_m and Ngrid_z
-    # This should probably be improved
-
-    logM_arr = np.linspace(logm_min, logm_max, Ngrid_m)
-    z_arr = np.logspace(np.log10(zmin), np.log10(zmax), Ngrid_z)
-    p = []
-    for logM in logM_arr:
-        for z in z_arr:
-            p.append(pdf(logM, z))
-    pmax = np.max(p)
-    pmin = np.min(p)
-    # Counters
-    naccept = 0
-    ntrial = 0
-
-    # Keeps drawing until N points are accepted
-    ran_logm = []  # output list of random numbers
-    ran_z = []  # output list of random numbers
-    while naccept < N:
-        # draw (logm,z) from uniform distribution
-        # draw p from uniform distribution
-        logm = np.random.uniform(logm_min, logm_max)  # x'
-        z = np.random.uniform(zmin, zmax)  # x'
-        p = np.random.uniform(0.0, pmax)  # y'
-
-        if p < pdf(logm, z):
-            # keep the point
-            ran_logm.append(logm)
-            ran_z.append(z)
-            naccept = naccept + 1
-        ntrial = ntrial + 1
-
-    ran_logm = np.asarray(ran_logm)
-    ran_z = np.asarray(ran_z)
-
-    acceptance = float(N / ntrial)
-    print(f"acceptance = {acceptance}")
-    return ran_logm, ran_z, acceptance
-
-
-# Draw N=30 (M,z) pairs
-
-# Define the numbers of pairs to draw and the mass and redshift ranges
-N = 10000
-logm_min = 14.0
-logm_max = 15.0
-zmin = 0.2
-zmax = 0.3
-
-pdf = pdf_tinker08_ccl
-
-# Normalisation factor for the pdf to be used later
-norm = (
-    1.0
-    / scipy.integrate.dblquad(
-        pdf, zmin, zmax, lambda x: logm_min, lambda x: logm_max, epsrel=1.0e-4
-    )[0]
-)
-
-# Random draw
-ran_logm, ran_z, acceptance = bivariate_draw(
-    pdf, N=N, logm_min=logm_min, logm_max=logm_max, zmin=zmin, zmax=zmax
-)
-
-# -----------------------------------------------------------------------------
-# Generate a cluster catalog and associated source catalogs
-# -----------------------------------------------------------------------------
-
-# Cluster ensemble masses, redshifts and positions
-cluster_m = 10 ** ran_logm
-cluster_z = ran_z
-
-# Concentration CCL object to compute the theoretical concentration
-conc_obj = ccl.halos.ConcentrationDuffy08(mass_def=hmd_200c)
-conc_list = []
-for number in range(0, len(cluster_m)):
-    a = 1.0 / (1.0 + (cluster_z[number]))
-    # mean value of the concentration for that cluster
-    lnc_mean = np.log(conc_obj(cosmo, M=(cluster_m[number]), a=a))
-    # random draw of actual concentration from normal distribution around lnc_mean, with a 0.14 scatter
-    lnc = np.random.normal(lnc_mean, 0.14)
-    conc_list.append(np.exp(lnc))
-
-conc_list = np.array(conc_list)
-
-# randomly draw cluster positions over the full sky
-ra = np.random.random(N) * 360  # from 0 to 360 deg
-sindec = np.random.random(N) * 2 - 1
-dec = np.arcsin(sindec) * 180 / np.pi  # from -90 to 90 deg
-
-print(len(cluster_m), N)
-
-# Background galaxy catalog generation
-#
-# For each cluster of the ensemble, we use `mock_data` to generate a background
-# galaxy catalog and store the results in a `GalaxyCluster` object. Note that:
-# - The cluster density profiles follow the NFW parametrisation
-# - The source redshifts follow the Chang et al. distribution and have associated pdfs
-# - The shapes include shape noise and shape measurement errors
-# - Background galaxy catalogs are independent, even if the clusters are close
-#   (i.e., no common galaxy between two catalogs).
-# - For each cluster we then compute
-#     - the tangential and cross Delta_Sigma for each background galaxy
-#     - the weights `w_ls` to be used to compute the corresponding radial profiles
-#
-# The cluster objects are then stored in `gclist`.
-
-warnings.simplefilter(
-    "ignore"
-)  # just to prevent warning print out when looping over the cluster ensemble below
-
-gclist = []
-tables = []
-# number of galaxies in each cluster field (alternatively, can use the galaxy density instead)
-ngal_density = 10
-cosmo_clmm = Cosmology(H0=71.0, Omega_dm0=0.265 - 0.0448, Omega_b0=0.0448, Omega_k0=0.0)
-cosmo_clmm.set_be_cosmo(cosmo)
-gal_args = dict(
-    cosmo=cosmo_clmm,
-    zsrc="chang13",
-    delta_so=200,
-    massdef="critical",
-    halo_profile_model="nfw",
-    zsrc_max=2.0,
-    field_size=16.0,
-    shapenoise=0.04,
-    photoz_sigma_unscaled=0.02,
-    ngal_density=ngal_density,
-    mean_e_err=0.1,
-)
-
-for i in range(N):
-    # generate background galaxy catalog for cluster i
-    cl = clmm.GalaxyCluster(
-        f"mock_cluster_{i:04}",
-        ra[i],
-        dec[i],
-        cluster_z[i],
-        galcat=mock.generate_galaxy_catalog(
-            cluster_m[i],
-            cluster_z[i],
-            conc_list[i],
-            cluster_ra=ra[i],
-            cluster_dec=dec[i],
-            zsrc_min=0.35,  # cluster zmax = 0.3 in this example
-            **gal_args,
-        ),
-    )
-
-    # compute DeltaSigma for each background galaxy
-    cl.compute_tangential_and_cross_components(
-        shape_component1="e1",
-        shape_component2="e2",
-        tan_component="g_t",
-        cross_component="g_x",
-        cosmo=cosmo_clmm,
-        is_deltasigma=False,
-        use_pdz=True,
-    )
-
-    # compute the weights to be used to bluid the DeltaSigma radial profiles
-    cl.compute_galaxy_weights(
-        use_pdz=True,
-        use_shape_noise=True,
-        shape_component1="e1",
-        shape_component2="e2",
-        use_shape_error=True,
-        shape_component1_err="e_err",
-        shape_component2_err="e_err",
-        weight_name="w_ls",
-        cosmo=cosmo_clmm,
-        is_deltasigma=False,
-        add=True,
-    )
-
-    # append the cluster in the list
-    gclist.append(cl)
-    tables.append(cl.galcat)
-
-# Stack as plain astropy Tables: GCData's metadata guard for "coordinate_system"
-# rejects vstack's metadata merge even when all tables share the same value.
-galcat_tot = vstack([Table(t) for t in tables])
-
-# -----------------------------------------------------------------------------
-# Create ClusterEnsemble object and estimation of individual excess surface
-# density profiles
-# -----------------------------------------------------------------------------
-
-bins = np.logspace(np.log10(0.5), np.log10(60), 12)
-ensemble_id = 1
-clusterensemble = ClusterEnsemble(ensemble_id)
-for cluster in gclist:
-    clusterensemble.make_individual_radial_profile(
-        galaxycluster=cluster,
-        tan_component_in="g_t",
-        cross_component_in="g_x",
-        tan_component_out="g_t",
-        cross_component_out="g_x",
-        weights_in="w_ls",
-        weights_out="W_l",
-        bins=bins,
-        bin_units="arcmin",
-        cosmo=cosmo_clmm,
-    )
-
-# Stacked profile of the cluster ensemble
-clusterensemble.make_stacked_radial_profile(tan_component="g_t", cross_component="g_x", weights="W_l")
-
-# Jackknife covariance of the stack between radial bins
-jk_n_side = 16
-clusterensemble.compute_jackknife_covariance(
-    n_side=jk_n_side, tan_component="g_t", cross_component="g_x"
-)
-njk = 10
-kmeans_gt_cov, kmeans_gx_cov, km_centers = jk_cov(clusterensemble, njk=njk)
-# -----------------------------------------------------------------------------
-# Visualizing the stacked profiles
-#
-# In the figure below, we plot:
-# - the individual g_t profiles of the clusters (light blue)
-# - the stacked signal (red symbols)
-# -----------------------------------------------------------------------------
-
-moo = clmm.Modeling(massdef="critical", delta_mdef=200, halo_profile_model="nfw")
-moo.set_cosmo(cosmo_clmm)
-# Average values of mass and concentration of the ensemble to be used below
-# to overplot the model on the stacked profile
-moo.set_concentration(conc_list.mean())
-moo.set_mass(cluster_m.mean())
-
-r_stack, gt_stack, gx_stack = (clusterensemble.stacked_data[c] for c in ("radius", "g_t", "g_x"))
-plt.rcParams["axes.linewidth"] = 2
-fig, axs = plt.subplots(1, 2, figsize=(17, 6))
-
-err_gt = clusterensemble.cov["tan_jk"].diagonal() ** 0.5
-err_gx = clusterensemble.cov["cross_jk"].diagonal() ** 0.5
-km_err_gt = np.sqrt(np.diag(kmeans_gt_cov))
-km_err_gx = np.sqrt(np.diag(kmeans_gx_cov))
-axs[0].errorbar(
-    r_stack,
-    gt_stack,
-    err_gt,
-    markersize=5,
-    c="r",
-    fmt="o",
-    capsize=10,
-    elinewidth=1,
-    zorder=1000,
-    alpha=1,
-    label="Stack",
-)
-axs[0].errorbar(
-    r_stack,
-    gt_stack,
-    km_err_gt,
-    markersize=5,
-    c="k",
-    fmt="o",
-    capsize=10,
-    elinewidth=1,
-    zorder=1000,
-    alpha=1,
-    label="Stack",
-)
-axs[1].errorbar(
-    r_stack,
-    gx_stack,
-    err_gx,
-    markersize=5,
-    c="r",
-    fmt="o",
-    capsize=10,
-    elinewidth=1,
-    zorder=1000,
-    alpha=1,
-    label="Stack",
-)
-axs[1].errorbar(
-    r_stack,
-    gx_stack,
-    km_err_gx,
-    markersize=5,
-    c="k",
-    fmt="o",
-    capsize=10,
-    elinewidth=1,
-    zorder=1000,
-    alpha=1,
-    label="Stack",
-)
-axs[0].set_xscale("log")
-axs[0].set_yscale("log")
-axs[1].set_xscale("log")
-
-
-for i in range(N):
-    axs[0].plot(
-        clusterensemble.data["radius"][i],
-        clusterensemble.data["g_t"][i],
-        color="cyan",
-        label="Individual",
-        alpha=1,
-        linewidth=1,
-    )
-    axs[1].plot(
-        clusterensemble.data["radius"][i],
-        clusterensemble.data["g_x"][i],
-        color="cyan",
-        label="Individual",
-        alpha=1,
-        linewidth=1,
-    )
-    if i == 0:
-        axs[0].legend(frameon=False, fontsize=15)
-        axs[1].legend(frameon=False, fontsize=15)
-
-axs[0].set_xlabel("separation [arcmin]", fontsize=20)
-axs[1].set_xlabel("separation [arcmin]", fontsize=20)
-axs[0].tick_params(axis="both", which="major", labelsize=18)
-axs[1].tick_params(axis="both", which="major", labelsize=18)
-axs[0].set_ylabel(r"$g_+$", fontsize=20)
-axs[1].set_ylabel(r"$g_\times$", fontsize=20)
-axs[0].set_title(r"Tangential", fontsize=20)
-axs[1].set_title(r"Cross", fontsize=20)
-
-for ax in axs:
-    ax.minorticks_on()
-    ax.grid(lw=0.5)
-    ax.grid(which="minor", lw=0.1)
-fig.tight_layout()
-fig.savefig(os.path.join(PLOTS_DIR, "individual_and_stacked_profiles.png"))
-plt.close(fig)
-
-# -----------------------------------------------------------------------------
-# Saving/Loading ClusterEnsemble
-#
-# The `ClusterEnsemble` object also has an option for saving/loading:
-# -----------------------------------------------------------------------------
-
-clusterensemble.save("ce.pkl")
-
-# -----------------------------------------------------------------------------
-# Compute reduced shear profile with TreeCorr and compare to CLMM stacked
-# profile
-# -----------------------------------------------------------------------------
-# Cluster catalog
-cat_cluster = treecorr.Catalog(
-    ra=clusterensemble["ra"],
-    dec=clusterensemble["dec"],
-    ra_units="deg",
-    dec_units="deg",
-    patch_centers=km_centers
-)
-# Source catalog
-cat_source = treecorr.Catalog(
-    ra=galcat_tot["ra"],
-    dec=galcat_tot["dec"],
-    g1=galcat_tot["e1"],
-    g2=galcat_tot["e2"],
-    ra_units="deg",
-    dec_units="deg",
-    w=galcat_tot["w_ls"],
-    patch_centers=cat_cluster.patch_centers
-)
-
-
-
-ng = treecorr.NGCorrelation(nbins=11, min_sep=0.5, max_sep=60, sep_units="arcmin", bin_type="Log")
-ng.process(cat_cluster, cat_source)
-theta = ng.meanr  # average angular separation in each bin
-cov_jk = ng.estimate_cov('jackknife')
-var_jk = np.diag(cov_jk)  
-print(cosmo_clmm.be_cosmo)
-
-da = cosmo.angular_diameter_distance(1.0 / (1.0 + np.mean(clusterensemble["z"])))
-da_clmm = cosmo_clmm.eval_da(np.mean(clusterensemble["z"]))
-print(da, da_clmm)
-
-arcmin_to_Mpc = np.pi / (60.0 * 180) * da_clmm
-
-fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 5))
-ax.errorbar(
-    theta,
-    ng.xi,
-    np.sqrt(var_jk),
-    color="blue",
-    marker="x",
-    markersize=15,
-    linestyle="",
-    label="Treecorr - cluster positions x source shears",
-)
-
-ax.errorbar(
-    r_stack,
-    gt_stack,
-    err_gt,
-    markersize=5,
-    c="r",
-    fmt=".",
-    capsize=10,
-    elinewidth=1,
-    zorder=1000,
-    alpha=1,
-    label="CLMM - stacked reduced rangential shear",
-)
-
-ax.plot(
-    clusterensemble.data["radius"][0],
-    moo.eval_reduced_tangential_shear(
-        clusterensemble.data["radius"][0] * arcmin_to_Mpc,
-        cluster_z.mean(),
-        z_src=desc_srd,
-        z_src_info="distribution",
-    ),
-    "--k",
-    linewidth=1,
-    label="CLMM prediction for stack 'mean' cluster",
-    zorder=100,
-)
-
-ax.set_xscale("log")
-ax.set_yscale("log")
-ax.set_ylim([-0.02, 0.15])
-ax.axhline(0, color="k", ls=":")
-ax.set_xlabel("separation [arcmin]")
-ax.set_ylabel("$\\widehat{g_t}$")
-ax.legend()
-fig.savefig(os.path.join(PLOTS_DIR, "treecorr_vs_clmm.png"))
-plt.close(fig)
-
-fig, ax = plt.subplots()
-ax.scatter(err_gt, np.sqrt(var_jk), label='healpy')
-ax.scatter(km_err_gt, np.sqrt(var_jk), label='kmeans')
-ax.plot([0, 0.007], [0, 0.007], "k:")
-ax.set_xlabel("Error from CLMM")
-ax.set_ylabel("Error from TreeCorr")
-ax.set_xscale('log')
-ax.set_yscale('log')
-plt.legend()
-fig.savefig(os.path.join(PLOTS_DIR, "error_comparison.png"))
-plt.close(fig)
-
-fig, ax = plt.subplots()
-ax.scatter(km_err_gt, err_gt)
-ax.plot([0, 0.007], [0, 0.007], "k:")
-ax.set_xlabel("Error from kmeans")
-ax.set_ylabel("Error from healpy")
-ax.set_xscale('log')
-ax.set_yscale('log')
-plt.legend()
-fig.savefig(os.path.join(PLOTS_DIR, "kmeans_healpy_error_comparison.png"))
-plt.close(fig)
diff --git a/clmm/hp_vs_kmeans.py b/clmm/hp_vs_kmeans.py
deleted file mode 100644
index 74e91e8aa..000000000
--- a/clmm/hp_vs_kmeans.py
+++ /dev/null
@@ -1,62 +0,0 @@
-import healpy as hp
-import matplotlib.pyplot as plt
-import numpy as np
-import os
-from kmeans_radec.test import plot_centers
-from kmeans_radec import KMeans, kmeans_sample, find_nearest
-
-PLOTS_DIR = "./plots/"
-os.makedirs(PLOTS_DIR, exist_ok=True)
-
-# randomly draw cluster positions over the full sky
-N = 30
-ra = np.random.random(N) * 360  # from 0 to 360 deg
-sindec = np.random.random(N) * 2 - 1
-dec = np.arcsin(sindec) * 180 / np.pi  # from -90 to 90 deg
-jk_n_side = 16
-# CLMM's jackknife regions are HEALPix pixels (not kmeans patches like TreeCorr),
-# and are not stored on the object, so recompute them here to visualize the full
-# extent of each region using healpy's own map plotting (cartview).
-jk_pixels = hp.ang2pix(
-    jk_n_side, ra, dec, nest=True, lonlat=True
-)
-jk_unique_pixels = np.unique(jk_pixels)
-jk_center_ra, jk_center_dec = hp.pix2ang(jk_n_side, jk_unique_pixels, nest=True, lonlat=True)
-jk_map = np.full(hp.nside2npix(jk_n_side), hp.UNSEEN)
-jk_map[jk_unique_pixels] = np.arange(jk_unique_pixels.size)
-
-fig = plt.figure(figsize=(12, 5))
-
-# Left panel: healpy manages its own projection axes internally, so instead of
-# handing it a matplotlib Axes, place it into a subplot slot of our figure via
-# fig= and sub= (same (nrows, ncols, index) convention as plt.subplot).
-hp.cartview(
-    jk_map,
-    fig=fig.number,
-    sub=(1, 2, 1),
-    nest=True,
-    cmap="tab20",
-    cbar=False,
-    title="Jackknife patches (HEALPix pixels)",
-)
-hp.graticule()
-hp.projscatter(ra, dec, lonlat=True, c="k", s=10)
-
-njk = 3 
-X = np.vstack((ra, dec)).T
-km = kmeans_sample(X, ncen=njk, method='slow')
-
-ax2= fig.add_subplot(1, 2, 2)
-plot_centers(km.centers, ax2, y_min=-90, y_max=90)
-ax2.scatter(ra, dec, c="r", s=30)
-ax2.set_title("kmeans_radec Voronoi cells")
-ax2.set_axis_off()
-
-fig.savefig(os.path.join(PLOTS_DIR, "hp_vs_kmeans.png"))
-plt.show()
-
-print('----- Distances ------')
-print(km.distances)
-labels = find_nearest(X, km.centers)
-print('--- labels ---')
-print(labels)
diff --git a/clmm/jk_kmeans.py b/clmm/jk_kmeans.py
deleted file mode 100644
index 95365f376..000000000
--- a/clmm/jk_kmeans.py
+++ /dev/null
@@ -1,40 +0,0 @@
-from kmeans_radec import kmeans_sample, find_nearest
-from clmm.dataops import make_stacked_radial_profile
-import numpy as np
-
-def radec_to_xyz(coords):
-    ra_rad = np.radians(coords[:, 0])
-    dec_rad = np.radians(coords[:, 1])
-    x = np.cos(dec_rad) * np.cos(ra_rad)
-    y = np.cos(dec_rad) * np.sin(ra_rad)
-    z = np.sin(dec_rad)
-    return np.column_stack([x, y, z])
-
-
-def jk_regions(X, njk):
-    km = kmeans_sample(X, ncen=njk)
-    labels = find_nearest(X, km.centers)
-    return labels, km.centers
-
-
-def jk_cov(clusterensemble, njk):
-    X = np.vstack((clusterensemble['ra'], clusterensemble['dec'])).T
-    labels, centers = jk_regions(X, njk=njk)
-    unique_labels = np.unique(labels)
-    gt_jack, gx_jack = [], []
-    for drop_region in unique_labels:
-        mask = np.isin(labels, drop_region, invert=True)
-        gt, gx = make_stacked_radial_profile(
-            clusterensemble['radius'][mask],
-            clusterensemble['W_l'][mask],
-            [clusterensemble['g_t'][mask], clusterensemble['g_x'][mask]])[1]
-        gt_jack.append(gt)
-        gx_jack.append(gx)
-    n_jack = unique_labels.size
-    coeff = (n_jack - 1)**2 / n_jack
-    tan_jk = coeff * np.cov(np.transpose(gt_jack), bias=False, ddof=0)
-    cross_jk = coeff * np.cov(np.transpose(gx_jack), bias=False, ddof=0)
-    cart_centers = radec_to_xyz(centers) 
-    return tan_jk, cross_jk, cart_centers
-
-
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)