cosmock 0.0.1__tar.gz

cosmock-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Ivan Espinoza
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cosmock-0.0.1/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: cosmock
+Version: 0.0.1
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: numpy>=2.0
+Requires-Dist: scipy>=1.15
+Requires-Dist: joblib>=1.2
+Requires-Dist: healpy>=1.16
+Dynamic: license-file

cosmock-0.0.1/README.md

@@ -0,0 +1,12 @@
+# cosmock
+
+Mock-field generation tools for cosmological kappa maps, implementing the
+formalism and algorithm presented in 2411.04759.
+
+## Installation
+
+Install the base package for fitting, transforms, and spectrum conversion:
+
+```bash
+python -m pip install -e .
+```

cosmock-0.0.1/cosmock/Cls.py

@@ -0,0 +1,169 @@
+import numpy as np 
+from numpy.polynomial.hermite import hermgauss
+from scipy.special import eval_legendre
+from scipy.interpolate import interp1d
+from joblib import Parallel, delayed
+from .Gn import Gn
+
+def get_gh_nodes_weights(n_nodes):
+    """
+    Computes the quadrature points and weights to 
+    solve integrals for the form I = E[f(x)] where x
+    is standard normal. So, to solve integral
+    simply do np.sum(y_weights * f(y_nodes)).
+    
+    Parameters
+    ----------
+    x : int
+            Number of quadrature points
+    
+    Returns
+    -------
+    points_weights : tuple
+            Returns the quadrature points, weights.
+    """
+    t, w = hermgauss(n_nodes)
+    y_nodes = np.sqrt(2.0) * t
+    y_weights = w / np.sqrt(np.pi)
+    return y_nodes, y_weights
+
+
+def F_gauss_hermite_single(N, params_i, params_j, xi_g,
+                           n_nodes=40,
+                           precomputed=None):
+    if precomputed is None:
+        y_nodes, y_weights = get_gh_nodes_weights(n_nodes)
+    else:
+        y_nodes, y_weights = precomputed
+
+    cov = np.array([[1.0, xi_g],
+                    [xi_g, 1.0]])
+    L = np.linalg.cholesky(cov)  
+
+    # We build tensor product nodes and weights 
+    YI, YJ = np.meshgrid(y_nodes, y_nodes, indexing='ij')
+    WI, WJ = np.meshgrid(y_weights, y_weights, indexing='ij')
+
+    # We map to X coordinates
+    Ystack = np.stack([YI.ravel(), YJ.ravel()], axis=1)   
+    Xstack = (L @ Ystack.T).T                             
+
+    xi_vals = Xstack[:, 0]
+    xj_vals = Xstack[:, 1]
+
+    Gn_i = Gn(xi_vals, N, params_i)
+    Gn_j = Gn(xj_vals, N, params_j)
+
+    integrand = Gn_i * Gn_j * (WI * WJ).ravel()
+    result = integrand.sum()
+    return result
+
+def build_lookup_table(N, params_i,params_j, xi_g_values, pre,nnodes=20):
+    results = []
+    for xi_g in xi_g_values:
+        result = F_gauss_hermite_single(N, params_i,params_j, xi_g, nnodes,pre)
+        results.append(result)
+    return np.array(results)
+
+
+def C_NG_to_C_G(cl_NG, fitted_params, N_bins, N, 
+                          xig_grid_size=75, quad_order=3, Nnodes=16, 
+                          n_jobs=4, v=0):
+    """
+    Converts non-Gaussian power spectrum to a gaussianized version
+    given the parameters
+    
+    Parameters
+    ----------
+    cl_NG : array
+        Non-Gaussian power spectrum
+    fitted_params : array
+        Parameters for Gn function for each z-bin
+    N_bins : int
+        Number of redshift bins
+    N : int
+        Parameter for Gn function
+    xig_grid_size : int, optional
+        Number of points in xi_g interpolation grid (default: 75)
+    quad_order : int, optional
+        Gauss-Legendre quadrature order multiplier (default: 3)
+        Total points = quad_order * lmax
+    Nnodes : int, optional
+        Number of Gauss-Hermite nodes for F integral (default: 16)
+    n_jobs : int, optional
+        Number of parallel jobs. -1 uses all cores (default: 4)
+    v : int, optional
+        Verbosity level for joblib (default: 0)
+    
+    Returns
+    -------
+    cl_G : array
+        Gaussian power spectrum
+    """    
+    cl_G = np.zeros_like(cl_NG)
+    
+    # Determine number of quadrature points needed
+    lmax_cl = cl_NG.shape[-1] - 1
+    n_quad  = quad_order * lmax_cl
+    
+    # We compute Gauss-Legendre quadrature points and weights
+    mu, w = np.polynomial.legendre.leggauss(n_quad)
+    
+    # Pre-compute Legendre polynomials, instead of computing them
+    # each time 
+    ell_array = np.arange(lmax_cl + 1)
+    P_ell     = np.array([eval_legendre(ell, mu) for ell in ell_array])
+
+    # Setup for lookup table
+    xi_g_grid = np.linspace(-0.99999, 0.99999, xig_grid_size)
+    pre       = get_gh_nodes_weights(Nnodes)
+    
+    def process_pair_optimized(i, j):
+        params_i = fitted_params[i]
+        params_j = fitted_params[j]
+                
+        # Cl_NG -> xi_NG 
+        ell_col = ell_array[:, np.newaxis]
+        arg = (2*ell_col + 1) * P_ell * cl_NG[i, j, :, np.newaxis]
+        xi_NG = np.sum(arg, axis=0) / (4*np.pi)
+        
+        # xi_NG -> xi_G
+        # Only in the case N = 2 we have an analitycal relation.
+        # In the other cases, we have to build lookup table.
+        if N == 2:
+            alpha_i, beta_i = params_i
+            alpha_j, beta_j = params_j
+            xi_G = np.log(1 + xi_NG / (beta_i * beta_j)) / (alpha_i * alpha_j)
+        else:
+            F_values = build_lookup_table(N, params_i, params_j, xi_g_grid, pre, Nnodes)
+            F_to_xi_g = interp1d(F_values, xi_g_grid, kind='linear',
+                                fill_value='extrapolate')
+            xi_G = F_to_xi_g(xi_NG)
+        
+        # xi_G -> Cl_G 
+        integrand = P_ell * xi_G[np.newaxis, :]
+        clG_ij = 2 * np.pi * np.sum(w[np.newaxis, :] * integrand, axis=1)
+        clG_ij[:2] = 1e-20
+        
+        return i, j, clG_ij
+    
+    # Generate list of pairs to process
+    pairs = [(i, j) for i in range(N_bins) for j in range(i + 1)]
+    
+    # Parallel computation
+    results = Parallel(n_jobs=n_jobs, verbose=v)(
+        delayed(process_pair_optimized)(i, j) for i, j in pairs
+    )
+    
+    # We fill the cl_G array with our results.
+    for i, j, clG_ij in results:
+        cl_G[i, j] = clG_ij
+        cl_G[j, i] = clG_ij
+    
+    # ell = 0,1 are zero. This gives problems
+    # Instead, we make a positive definite matrix
+    # with very small eigenvalues. 
+    for l in [0, 1]:
+        cl_G[:, :, l] = 1e-20 * np.eye(N_bins)
+    
+    return cl_G

cosmock-0.0.1/cosmock/Gn.py

@@ -0,0 +1,31 @@
+import numpy as np
+
+def Gn(x, N, lbda):
+    """
+    Evaluate Gn transformation at given x values
+    
+    Parameters
+    ----------
+    x : array of x values (standard normal inputs)
+    n : str, which G function to use ('2', '3', '4', '5')
+    params : array of parameters for the transformation
+    N_nodes : number of Gauss-Hermite quad points to compute
+             integral for normalization.
+    
+    Returns
+    -------
+    y : transformed values
+    """
+    # Evaluate transformation
+    if N == 2:
+        alpha, beta = lbda
+        return beta * np.exp(alpha * x - 0.5 * alpha**2) - beta
+    
+    elif N == 3:
+        a, b, c = lbda
+        arg = np.exp(a * x - 0.5 * a**2) + b*x + c
+        norm = 1/(1+c)
+        return norm * arg - 1
+    
+    else:
+        raise ValueError(f"Unknown model type: {N}")

cosmock-0.0.1/cosmock/__init__.py

@@ -0,0 +1 @@
+from .wrappers import fit_parameters,generate_mocks

cosmock-0.0.1/cosmock/fitter.py

@@ -0,0 +1,181 @@
+import numpy as np
+from .Gn import Gn
+from scipy.optimize import minimize
+from scipy.stats import norm
+
+def variance_from_Cl(Cl):
+    """
+    Computes the variance of the field
+    as predicted by the cls.
+    
+    Parameters
+    ----------
+    Cl : array
+             Array of cls
+    ell_min : float 
+             First ell to start calculation. 
+             Defaults to ell=0. 
+    
+    Returns
+    -------
+    variance : float
+             Variance of the field as predicted by cls.
+    """
+    Cl = np.asarray(Cl)
+    ell = np.arange(len(Cl))
+    return np.sum((2*ell + 1) * Cl) / (4*np.pi)
+
+def get_binned_data(x_data, y_data, x_range, n_bins):
+    """
+    Bin x values and compute mean y within each bin
+    
+    Parameters
+    ----------
+    x_data : array of Gaussianized x values
+    y_data : array of corresponding y values
+    n_bins : number of bins
+    x_range : tuple (min, max) to keep 
+    
+    Returns
+    -------
+    x_bin_centers : x value at center of each bin
+    y_bin_means : mean y value in each bin
+    """
+
+    mask = (x_data >= x_range[0]) & (x_data <= x_range[1])
+    x_filtered = x_data[mask]
+    y_filtered = y_data[mask]
+    
+    # Create bins
+    bin_edges = np.linspace(x_filtered.min(), x_filtered.max(), n_bins + 1)
+    
+    # Find which bin each x value belongs to
+    bin_indices = np.digitize(x_filtered, bin_edges) - 1
+    
+    # Clip to valid range
+    bin_indices = np.clip(bin_indices, 0, n_bins - 1)
+    
+    x_bin_centers = []
+    y_bin_means   = []
+    
+    for i in range(n_bins):
+        bin_mask = (bin_indices == i)
+        # We only want to include bins with data
+        if np.sum(bin_mask) > 0:  
+            x_bin_centers.append(bin_edges[i:i+2].mean())  
+            y_bin_means.append(y_filtered[bin_mask].mean())
+    
+    return np.array(x_bin_centers), np.array(y_bin_means)
+
+def empirical_cdf(map):
+    """
+    Computes the CDF of a map
+   
+    Parameters
+    ----------
+    map : array of that represents map
+   
+    Returns
+    -------
+    sorted_map : x-coordinates of the CDF
+    cdf_values : y-coordinates of the CDF
+    """
+    sorted_map = np.sort(map)
+    cdf_values = np.arange(1, len(sorted_map) + 1) / len(sorted_map)
+    cdf_values = np.clip(cdf_values, 1e-10, 1 - 1e-10)
+   
+    return sorted_map, cdf_values
+
+def histogramer2d(map,Nbins,x_range=(-4.5,4.5)):
+    """
+    Given a NL field, computes the black triangles in
+    FIG 1 in 2411.04759
+    
+    Parameters
+    ----------
+    map : field that we want to model
+    
+    Returns
+    -------
+    x_avg : x value of triangle (standard normal)
+    y_avg : y value of triangle (NL field)
+    """
+    y_data, cdf  = empirical_cdf(map)
+    x_gaussianized = norm.ppf(cdf)
+    x_avg,y_avg = get_binned_data(x_gaussianized,y_data,x_range,Nbins)
+    return x_avg, y_avg
+
+
+def fit_gn_with_constraint(x_data, y_data, N, cls, initial_lbda = None):
+    """
+    Fit a Gn transformation to (x, y) data points 
+    in a self-consistent manner by also including
+    the variance as predicted by the power spectrum.
+    
+    Parameters
+    ----------
+    x_data : array
+             array of x-values (standard normal)
+    y_data : array 
+             array of y-values (NL field)
+    N : str
+             Which G function to use ('2' and '3' are the 
+             only supported)
+    cls : array
+             The power spectrum of the field.
+    initial_lbda : array
+             Initialization. If None, defaults to ones.
+    
+    Returns
+    -------
+    fitted_lbda : array
+             Best fit parameters.
+    """
+    var = variance_from_Cl(cls)
+
+    def calc_constrained_lbda(unconstrained_lbda, var, N):
+        if N == 2:
+            beta = unconstrained_lbda[0]
+            alpha = np.sqrt(np.log(1 + var / beta**2))
+            return np.array([alpha, beta])
+        elif N == 3:
+            a, b = unconstrained_lbda
+            c = np.sqrt((np.exp(a**2) - 1 + 2*a*b + b**2) / var) - 1
+            return np.array([a, b, c])
+        else:
+            raise ValueError(f"Cannot do a constrained fit for G{N}.")
+        
+    def cost_function(unconstrained_lbda):
+        """Least squares cost"""
+        lbda = calc_constrained_lbda(unconstrained_lbda, var, N)
+
+        try:
+            y_pred = Gn(x_data, N, lbda)
+            return np.sum((y_pred - y_data)**2)
+        except:
+            return np.inf
+    if initial_lbda is None:
+        if N == 2:
+            beta_grid = np.geomspace(1e-6, 1, 50)
+            cost_grid = np.array([cost_function([b]) for b in beta_grid])
+            beta_init = beta_grid[np.argmin(cost_grid)]
+            initial_lbda = np.array([beta_init, np.nan])
+        elif N == 3:
+            a_grid = np.linspace(0.01, 1.7, 25)
+            b_grid = np.linspace(0.01, 2.7, 25)
+            cost_grid = np.array([
+                [cost_function([aa, bb]) for aa in a_grid]
+                for bb in b_grid
+            ])
+            j_min, i_min = np.unravel_index(np.argmin(cost_grid), cost_grid.shape)
+            initial_lbda = np.array([a_grid[i_min], b_grid[j_min], np.nan])
+        else:
+            initial_lbda = np.ones(int(N))
+    initial_unconstrained_lbda = initial_lbda[:int(N)-1]
+
+    result = minimize(
+        fun=cost_function,
+        x0=initial_unconstrained_lbda,
+        method='BFGS'
+    )
+    return calc_constrained_lbda(result.x, var, N)

cosmock-0.0.1/cosmock/mocker.py

@@ -0,0 +1,82 @@
+import numpy as np 
+import healpy as hp 
+from .Gn import Gn
+
+def get_xlm(xlm_real, xlm_imag,gen_lmax,nbins):
+    ell, emm = hp.Alm.getlm(gen_lmax)
+    #==============================
+    _xlm_real = np.zeros((nbins, len(ell)))
+    _xlm_imag = np.zeros_like(_xlm_real)
+    _xlm_real[:,ell > 1] = xlm_real
+    _xlm_imag[:,(ell > 1) & (emm > 0)] = xlm_imag
+    xlm = _xlm_real + 1j * _xlm_imag
+    #==============================
+    return xlm
+
+def generate_xlm(nbins,gen_lmax):
+    ell, emm = hp.Alm.getlm(gen_lmax)
+    xlm_real = np.random.normal(size=(nbins, (ell > 1).sum()))
+    xlm_imag = np.random.normal(size=(nbins, ((ell > 1) & (emm > 0)).sum()))
+    xlm = get_xlm(xlm_real, xlm_imag, gen_lmax, nbins)
+    return xlm, [xlm_real,xlm_imag]
+
+def apply_cl_G(xlm, Cl_G, gen_lmax):
+    Cl_G_T = np.moveaxis(Cl_G, 2, 0)
+    L_T = np.linalg.cholesky(Cl_G_T)
+    L_G = np.moveaxis(L_T, 0, 2)
+    gen_ell, gen_emm = hp.Alm.getlm(gen_lmax)
+    L_expanded = L_G[:, :, gen_ell]
+    ylm_real = np.einsum('ijm,jm->im', L_expanded, xlm.real) / np.sqrt(2)
+    ylm_imag = np.einsum('ijm,jm->im', L_expanded, xlm.imag) / np.sqrt(2)
+    ylm_real = np.where(gen_emm == 0, ylm_real * np.sqrt(2), ylm_real)
+    ylm_imag = np.where(gen_emm == 0, 0.0, ylm_imag)
+    return ylm_real + 1j * ylm_imag 
+
+def get_y_maps(cl,nside,nbins,gen_lmax,xlms=None):
+    if xlms is not None:
+        xlm = xlms
+        _xlm = None
+    else:
+        xlm, _xlm = generate_xlm(nbins,gen_lmax)
+    y_lm, xlm = apply_cl_G(xlm, cl,gen_lmax), _xlm
+    y_maps = []
+    for i in range(nbins):
+        y_map = hp.alm2map(np.ascontiguousarray(y_lm[i]), nside, lmax=gen_lmax, pol=False)
+        y_maps.append(y_map)    
+    return np.array(y_maps)    
+
+def get_kappa(y_maps,nbins,N,fitted_params):
+    k_list = []
+    for i in range(nbins):
+        k_nf = Gn(y_maps[i], N, fitted_params[i])
+        k = k_nf
+        k_list.append(k)  
+    k_arr  = np.array(k_list)
+    return k_arr  
+
+def get_kappa_pixwin(y_maps,nbins,N,fitted_params,nside,pixwinatell):
+    k_list = []
+    lmax = 2*nside
+    for i in range(nbins):
+        k_nf = Gn(y_maps[i], N, fitted_params[i])
+        k = k_nf
+        klm = hp.map2alm(k,lmax=lmax)
+        klm = klm * pixwinatell 
+        k = hp.alm2map(klm,nside)
+        k_list.append(k)  
+    k_arr  = np.array(k_list)
+
+    return k_arr  
+
+def get_kappa_lm_pixwin(y_maps,nbins,N,fitted_params,nside,pixwinatell):
+    k_lm_list = []
+    lmax = 2*nside
+    for i in range(nbins):
+        k_nf = Gn(y_maps[i], N, fitted_params[i])
+        k = k_nf
+        klm = hp.map2alm(k,lmax=lmax)
+        klm = klm * pixwinatell 
+        k_lm_list.append(klm)  
+    k_lm_arr  = np.array(k_lm_list)
+
+    return k_lm_arr  

cosmock-0.0.1/cosmock/structs.py

@@ -0,0 +1,8 @@
+from typing import NamedTuple
+import numpy as np
+
+class NonlinParameters(NamedTuple):
+    lbda: np.ndarray
+    N: int
+    Cl_Gauss: np.ndarray
+    mapshape: tuple

cosmock-0.0.1/cosmock/wrappers.py

@@ -0,0 +1,52 @@
+import numpy as np
+import healpy as hp
+from .structs import NonlinParameters
+from .fitter import histogramer2d, fit_gn_with_constraint
+from .Cls import C_NG_to_C_G
+from .mocker import get_y_maps, get_kappa_pixwin
+
+
+def fit_parameters(maps, Cl_delta, N):
+
+    Nbins = maps.shape[0]
+    Nside = hp.npix2nside(maps.shape[1])
+
+    lbda = np.zeros((Nbins, N))
+    for i in range(Nbins):
+        x, y = histogramer2d(maps[i], 1000)
+        lbda[i] = fit_gn_with_constraint(x, y, N, Cl_delta[i,i,:3*Nside])
+    
+    print(f'Done fitting the G{N} parameters')
+
+    Cl_G = C_NG_to_C_G(Cl_delta, lbda, Nbins, N)
+    Cl_G = Cl_G[:,:,:3*Nside]
+
+    print('Done finding the power spectrum of the underlying gaussian random field')
+
+    params = NonlinParameters(lbda = lbda, N = N, Cl_Gauss = Cl_G, mapshape = maps.shape)
+
+    return params
+
+
+def generate_mocks(params, Nmocks, pixwin=None):
+
+    lbda = params.lbda
+    N = params.N
+    Cl_G = params.Cl_Gauss
+    mapshape = params.mapshape
+    Nbins    = mapshape[0]
+    Npix     = mapshape[1]
+    Nside    = hp.npix2nside(Npix)
+    gen_lmax = 3*Nside-1
+    lmax     = 2*Nside
+    ell, emm = hp.Alm.getlm(lmax)
+
+    if pixwin is None:
+        pixwin = hp.pixwin(Nside, False, gen_lmax)
+
+    mocks = np.zeros((Nmocks, *mapshape))
+    for i in range(Nmocks):
+        y_map = get_y_maps(Cl_G, Nside, Nbins, gen_lmax)
+        mocks[i] = get_kappa_pixwin(y_map, Nbins, N, lbda, Nside, pixwin[ell])
+
+    return mocks

cosmock-0.0.1/cosmock.egg-info/PKG-INFO

@@ -0,0 +1,10 @@
+Metadata-Version: 2.4
+Name: cosmock
+Version: 0.0.1
+Requires-Python: >=3.10
+License-File: LICENSE
+Requires-Dist: numpy>=2.0
+Requires-Dist: scipy>=1.15
+Requires-Dist: joblib>=1.2
+Requires-Dist: healpy>=1.16
+Dynamic: license-file

cosmock-0.0.1/cosmock.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+cosmock/Cls.py
+cosmock/Gn.py
+cosmock/__init__.py
+cosmock/fitter.py
+cosmock/mocker.py
+cosmock/structs.py
+cosmock/wrappers.py
+cosmock.egg-info/PKG-INFO
+cosmock.egg-info/SOURCES.txt
+cosmock.egg-info/dependency_links.txt
+cosmock.egg-info/requires.txt
+cosmock.egg-info/top_level.txt
+tests/test_CG.py

cosmock-0.0.1/cosmock.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

cosmock-0.0.1/cosmock.egg-info/requires.txt

@@ -0,0 +1,4 @@
+numpy>=2.0
+scipy>=1.15
+joblib>=1.2
+healpy>=1.16

cosmock-0.0.1/cosmock.egg-info/top_level.txt

@@ -0,0 +1 @@
+cosmock

cosmock-0.0.1/pyproject.toml

@@ -0,0 +1,15 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "cosmock"
+version = "0.0.1"
+requires-python = ">=3.10"
+
+dependencies = [
+    "numpy>=2.0",
+    "scipy>=1.15",
+    "joblib>=1.2",
+    "healpy>=1.16",
+]

cosmock-0.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

cosmock-0.0.1/tests/test_CG.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+import tomllib
+
+import cosmock
+
+
+def test_public_api_imports():
+    assert "Gn" in cosmock.__all__
+    assert "C_NG_to_C_G" in cosmock.__all__
+
+
+def test_version_matches_pyproject():
+    pyproject = tomllib.loads(Path("pyproject.toml").read_text())
+    assert cosmock.__version__ == pyproject["project"]["version"]