smc.py: knn-based covariance estimation

Author: gpmp-dev

gpmp/misc/knn_cov.py

@@ -0,0 +1,90 @@
+import gpmp.num as gnp
+
+try:
+    import hnswlib
+    HNSWLIB_AVAILABLE = True
+except ImportError:
+    HNSWLIB_AVAILABLE = False
+    print(
+        "hnswlib is not installed. estimate_cov_matrix_knn will fall back to classical covariance estimation."
+    )
+
+def estimate_cov_matrix(x):
+    """
+    Default classical covariance estimation using the sample covariance.
+    x: shape = (N, d)
+    Returns a (d, d) covariance matrix
+    """
+    # gpmp.num.cov expects shape (d, N), so we transpose x
+    return gnp.cov(x.T)
+
+def estimate_cov_matrix_knn(
+    x,
+    n_random=50,
+    n_neighbors=50,
+    ef=50,
+    max_ef_construction=200,
+    M=16,
+):
+    """
+    Estimate the covariance matrix of x by:
+      1) Randomly sampling 'n_random' points from x.
+      2) Building an HNSW index (if hnswlib is available).
+      3) For each sampled point, finding its n_neighbors nearest neighbors.
+      4) Computing each local k-NN covariance.
+      5) Averaging the results into a single covariance matrix.
+
+    Fallback: if HNSWLIB_AVAILABLE is False, returns classical covariance.
+
+    Parameters
+    ----------
+    x : ndarray, shape (N, d)
+        The dataset.
+    n_random : int, optional
+        Number of random points from x to sample for local covariance.
+    n_neighbors : int, optional
+        Number of neighbors for each point to compute local covariances.
+    ef : int, optional
+        'ef' parameter for HNSW search (larger means more accurate at some cost).
+    max_ef_construction : int, optional
+        'ef_construction' parameter for HNSW index building.
+    M : int, optional
+        HNSW parameter controlling the connectivity of the graph.
+
+    Returns
+    -------
+    C_avg : ndarray, shape (d, d)
+        The averaged covariance matrix (or classical covariance if fallback).
+    """
+    if not HNSWLIB_AVAILABLE:
+        # If import hnswlib failed at module load, fallback to classical
+        return estimate_cov_matrix(x)
+
+    N, d = x.shape
+    if n_random > N:
+        n_random = N
+
+    # 1) Build the HNSW index
+    p = hnswlib.Index(space='l2', dim=d)
+    p.init_index(max_elements=N, ef_construction=max_ef_construction, M=M)
+    p.add_items(x)
+    p.set_ef(ef)
+
+    # 2) Randomly select points and compute local covariances
+    random_indices = gnp.choice(N, size=n_random, replace=False)
+    local_covs = []
+    
+    for idx in random_indices:
+        query_point = x[idx]
+        # knn_query returns two arrays: (labels, distances) each of shape (1, k)
+        labels, distances = p.knn_query(query_point, k=n_neighbors)
+        neighbors_idx = labels[0]        # shape (n_neighbors,)
+        neighbors_x = x[neighbors_idx]   # shape (n_neighbors, d)
+
+        # local covariance (gpmp.num.cov => shape (d, n_neighbors))
+        C_local = gnp.cov(neighbors_x.T)
+        local_covs.append(C_local)
+
+    # 3) Average all local covariance matrices
+    C_avg = gnp.mean(local_covs, axis=0)  # shape (d, d)
+    return gnp.asarray(C_avg)

gpmp/misc/smc.py

@@ -21,6 +21,7 @@
 from scipy.stats import qmc
 from scipy.optimize import brentq
 import gpmp.num as gnp
+import gpmp.misc.knn_cov
 import matplotlib.pyplot as plt
 import matplotlib.gridspec as gridspec
 
@@ -34,6 +35,9 @@ class ParticlesSetConfig:
     param_s_lower_bound: float = 1e-3
     jitter_initial_value: float = 1e-16
     jitter_max_iterations: int = 10
+    covariance_method: str = "knn"
+    covariance_knn_n_random: int = 50
+    covariance_knn_n_neighbors: int = 100
 
 
 @dataclass
@@ -364,7 +368,15 @@ def perturb(self):
             raise ParticlesSetError(self.param_s, param_s_lower, param_s_upper)
 
         # Covariance matrix of the pertubation noise
-        C = self.param_s * gnp.cov(self.x.reshape(self.n, -1).T)
+        if self.config.covariance_method == "knn":
+            base_cov = gpmp.misc.knn_cov.estimate_cov_matrix_knn(
+                self.x,
+                n_random=self.config.covariance_knn_n_random,
+                n_neighbors=self.config.covariance_knn_n_neighbors,
+            )  # shape (dim, dim)
+        elif self.config.covariance_method == "normal":
+            base_cov = gpmp.misc.knn_cov.estimate_cov_matrix(self.x)
+        C = self.param_s * base_cov
 
         # Call ParticlesSet.multivariate_normal_rvs(C, self.n, self.rng)
         # with control on the possible degeneracy of C
@@ -424,7 +436,7 @@ def move(self):
         self.x = gnp.set_row_2d(self.x, accept_mask, y[accept_mask, :])
         # self.logpx[accept_mask] = logpy[accept_mask]
         self.logpx = gnp.set_elem_1d(self.logpx, accept_mask, logpy[accept_mask])
-        
+
         # Compute the acceptance rate
         acceptance_rate = gnp.sum(accept_mask) / self.n
 
@@ -1287,7 +1299,7 @@ def test_run_smc_sampling_gaussian_mixture():
     from scipy import stats
 
     # Gaussian mixture parameters
-    m1, s1, w1 = 0.0, 0.05, 0.3
+    m1, s1, w1 = 0.0, 0.04, 0.3
     m2, s2, w2 = 1.0, 0.1, 0.7
 
     # Tempered log pdf: log p_T(x) = beta * log(p(x))
@@ -1301,7 +1313,7 @@ def logpdf_mixture(x, beta):
 
     # Domain: 1D in [-1, 2]
     init_box = [[-1], [2]]
-    initial_logpdf_param = 0.1  # initial beta
+    initial_logpdf_param = 0.01  # initial beta
     target_logpdf_param = 1.0  # target beta
 
     # SMC settings
@@ -1329,13 +1341,14 @@ def logpdf_mixture(x, beta):
     print("Sample variance:", gnp.var(particles))
 
     # Plot target density and histogram of particles
-    x_vals = gnp.linspace(-1, 2, 300)
-    target_density = w1 * stats.norm.pdf(
+    x_vals = gnp.linspace(-0.5, 1.5, 600)
+    target_density = lambda x_vals: w1 * stats.norm.pdf(
         x_vals, loc=m1, scale=s1
     ) + w2 * stats.norm.pdf(x_vals, loc=m2, scale=s2)
     plt.figure(figsize=(8, 3))
     plt.hist(particles, bins=100, density=True, histtype="step", label="SMC particles")
-    plt.plot(x_vals, target_density, "r--", label="Target density")
+    plt.plot(x_vals, target_density(x_vals), "r--", label="Target density")
+    #  plt.plot(particles, target_density(particles), 'b.')
     plt.xlabel("x")
     plt.ylabel("Density")
     plt.legend()