Commit 3b2aa9aa by Goran Jelic-Cizmek

```For now, it only has the `covariance_matrix` function, but I'm sure more will
parent a92f45d7
coffe_utils.py 0 → 100644
 """ Various helper utilities that are outside of the scope of the core functionality of COFFE. """ from __future__ import annotations from typing import Callable, List, Optional, Union import numpy as np from scipy.linalg import block_diag import pandas as pd from coffe import Covariance def covariance_matrix( cov : List[Covariance], l : Optional[List[int]] = None, z_mean : Optional[List[float]] = None, deltaz : Optional[List[float]] = None, rmin : Optional[Union[Callable, float, int]] = None, rmax : Optional[Union[Callable, float, int]] = None, rstep : Optional[Union[float, int]] = None, ) -> np.array: """ Converts an array of covariances into a numpy matrix for easy matrix multiplication. Parameters ---------- l : Optional[List[int]], default = None the list of values of l z_mean : Optional[List[float]], default = None the list of values of mean redshifts deltaz : Optional[List[float]], default = None the list of redshift bin widths. Only useful if one of `rmin`, `rmax` is a function. rmin : Optional[Union[Callable, float, int]], default = None If an int or a float, the smallest separation to take from the covariance matrix. If a function, must take exactly two mandatory arguments, which should be the mean redshift and the size of the bin. An example can be `lambda z_mean, deltaz: 1 / np.sqrt(1 + z)` rmax : Optional[Union[Callable, float, int]], default = None If an int or a float, the largest separation to take from the covariance matrix. If a function, must take exactly two mandatory arguments, which should be the mean redshift and the size of the bin. An example is the function `maximum_separation` from the `Coffe` class. rstep : Optional[float] = None if set, only takes separations from the covariance matrix that are multiples of `rstep`. Examples -------- >>> covariance_matrix(coffe.Coffe(has_density=True).compute_corrfunc_bulk()) """ def convert_array_to_matrix( arr ): size = round(np.sqrt(len(arr))) return np.reshape(arr, (size, size)) if not all(isinstance(_, Covariance) for _ in cov): raise ValueError( f'{cov} is not an array of covariances' ) df = pd.DataFrame([_.to_dict() for _ in cov]) # N.B. this may require fixing due to floating point math if z_mean is not None: df = df.loc[df.z.isin(z_mean)] else: z_mean = df.z.unique() if l is not None: df = df.loc[(df.l1.isin(l)) & (df.l2.isin(l))] if rmin is not None: if isinstance(rmin, (int, float)): df = df.loc[(df.r1 >= rmin) & (df.r2 >= rmin)] else: func_rmin = rmin df_rmin = pd.DataFrame() for z, dz in zip(z_mean, deltaz): df_rmin = df_rmin.append( df.loc[ (df.r1 >= func_rmin(z, dz)) & (df.r2 >= func_rmin(z, dz)) & (df.z == z) ] ) df = df_rmin if rmax is not None: if isinstance(rmax, (int, float)): df = df.loc[(df.r1 <= rmax) & (df.r2 <= rmax)] else: func_rmax = rmax df_rmax = pd.DataFrame() for z, dz in zip(z_mean, deltaz): df_rmax = df_rmax.append( df.loc[ (df.r1 <= func_rmax(z, dz)) & (df.r2 <= func_rmax(z, dz)) & (df.z == z) ] ) df = df_rmax if rstep is not None: df = df[np.isclose(df.r1 % rstep, 0)] return block_diag( *[ convert_array_to_matrix( df.loc[df.z == z].value.to_numpy(dtype=float) ) for z in z_mean ] )
 ... ... @@ -4,6 +4,7 @@ import pandas as pd import pytest import coffe from coffe_utils import covariance_matrix DATA_DIR = 'tests/benchmarks/' ... ... @@ -227,11 +228,20 @@ class TestCoffe: assert np.allclose(df.loc[(df.l1 == mp1) & (df.l2 == mp2)].r2.values, y) assert np.allclose(df.loc[(df.l1 == mp1) & (df.l2 == mp2)].value.values, z, rtol=5e-4) assert covariance_matrix(result).ndim == 2 assert np.shape(covariance_matrix(result)) == (18, 18) assert np.allclose( 0, 0 covariance_matrix(result), np.transpose(covariance_matrix(result)) ) # we don't need to test the trace or the determinant since that follows # automatically if the matrix has all positive eigenvalues assert np.all(np.linalg.eigvalsh(covariance_matrix(result)) > 0) assert covariance_matrix(result, rstep=111).size == 0 class TestRepresentation: ... ...
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!