Merge pull request #18 from DifferentiableUniverseInitiative/kde_nz

Kde nz

Author: EiffL

jax_cosmo/angular_cl.py

@@ -129,8 +129,8 @@ def gaussian_cl_covariance(cosmo, ell, probes, f_sky=0.25, return_cls=True):
 
   def get_cov_block(inds):
       a, b, c, d = inds
-      cov = (cl_obs[a]*cl_obs[b] + cl_obs[c]*cl_obs[d])*np.eye(n_ell) / norm
-      return cov
+      cov = (cl_obs[a]*cl_obs[b] + cl_obs[c]*cl_obs[d]) / norm
+      return cov*np.eye(n_ell)
 
   cov_mat = lax.map(get_cov_block, cov_blocks)
 

jax_cosmo/kernels.py

@@ -86,3 +86,37 @@ class smail_nz(redshift_distribution):
   def pz_fn(self, z):
     a, b, z0 = self.params
     return z**a * np.exp(-(z / z0)**b)
+
+@register_pytree_node_class
+class kde_nz(redshift_distribution):
+  """
+  A redshift distribution based on a KDE estimate of the nz of a given catalog
+  currently uses a Gaussian kernel, TODO: add more if necessary
+
+  Parameters:
+  -----------
+  zcat: redshift catalog
+  weights: weight for each galaxy between 0 and 1
+
+  Configuration:
+  --------------
+  bw: Bandwidth for the KDE
+
+  Example:
+  nz = kde_nz(redshift_catalog, w, bw=0.1)
+  """
+
+  def _kernel(self, bw, X, x):
+    """
+    Gaussian kernel for KDE
+    """
+    return (1. / np.sqrt(2 * np.pi )/bw) * np.exp(-(X - x)**2 / (bw**2 * 2.))
+
+  def pz_fn(self, z):
+    # Extract parameters
+    zcat, weight = self.params[:2]
+    w = np.atleast_1d(weight)
+    q = np.sum(w)
+    X = np.expand_dims(zcat, axis=-1)
+    k = self._kernel(self.config['bw'], X , z)
+    return np.dot(k.T, w)/(q)