Skip to content

Commit

Permalink
Merge pull request #29 from simonsobs/cint
Browse files Browse the repository at this point in the history
Add first example of unit test for `ps_utils` functions & include a montecarlo covariance stage in the pipeline.
  • Loading branch information
adrien-laposta authored Feb 13, 2024
2 parents 397be8d + fcdcd61 commit 066b8bd
Show file tree
Hide file tree
Showing 7 changed files with 329 additions and 126 deletions.
40 changes: 40 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,47 @@ jobs:
- name: Checkout SOOPERCOOL repository
uses: actions/checkout@v3

- name: Install python
uses: actions/setup-python@v2
with:
python-version: 3.8

- uses: awvwgk/setup-fortran@main
id: setup-fortran
with:
compiler: gcc
version: 11

- run: ${{ env.FC }} --version
env:
FC: ${{ steps.setup-fortran.outputs.fc }}

- name: Install dependencies
run: |
sudo -H apt-get install libgsl-dev libfftw3-dev libcfitsio-dev
- name: Install soopercool & dependencies
run: |
python -m pip install -U pip
pip install -U wheel
pip install -U setuptools
pip install -U numpy
pip install -U scipy
pip install -U healpy
pip install -U sacc
pip install -U camb
pip install -U pymaster
pip install -U flake8
pip install -U pytest
pip install -U pytest-cov
pip install -U coveralls
python setup.py install
- name: Lint
uses: py-actions/flake8@v2
with:
args: "--config .flake8"

- name: Tests
run: pytest -vv --cov=soopercool

1 change: 1 addition & 0 deletions paramfiles/paramfile_SAT.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ output_dirs:
cell_data_directory: "cells_data"
cell_sims_directory: "cells_sims"
coupling_directory: "couplings"
covmat_directory: "covariances"

##################################
# Metadata related to maps #
Expand Down
121 changes: 121 additions & 0 deletions pipeline/covfefe.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
import argparse
from soopercool import BBmeta
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable


def covFeFe(args):
"""
This function handles the covariance matrix computation.
We can use a set of simulations to build a numerical estimate
of the covariance matrices.
TODO: implement an analytical estimate of the covariances
"""
meta = BBmeta(args.globals)
cov_dir = meta.covmat_directory
n_bins = meta.get_n_bandpowers()

Nsims = meta.num_sims
cl_dir = meta.cell_sims_directory

field_pairs = [f"{m1}{m2}" for m1 in "TEB" for m2 in "TEB"]

cross_ps_names = meta.get_ps_names_list(type="all", coadd=True)

# Build the covariance matrix elements to compute
cov_names = []
for i, (ms1, ms2) in enumerate(cross_ps_names):
for j, (ms3, ms4) in enumerate(cross_ps_names):
if i > j:
continue
cov_names.append((ms1, ms2, ms3, ms4))

# Load the simulations
cl_dict = {}
for ms1, ms2 in cross_ps_names:
cl_dict[ms1, ms2] = []
for iii in range(Nsims):
cells_dict = np.load(
f"{cl_dir}/decoupled_cross_pcls_nobeam_{ms1}_{ms2}_{iii:04d}.npz", # noqa
)
cl_vec = np.concatenate(
[
cells_dict[field_pair] for field_pair in field_pairs
]
)
cl_dict[ms1, ms2].append(cl_vec)

cl_dict[ms1, ms2] = np.array(cl_dict[ms1, ms2])

full_cov_dict = {}
for ms1, ms2, ms3, ms4 in cov_names:

cl12 = cl_dict[ms1, ms2]
cl34 = cl_dict[ms3, ms4]

cl12_mean = np.mean(cl12, axis=0)
cl34_mean = np.mean(cl34, axis=0)
cov = np.mean(
np.einsum("ij,ik->ijk", cl12-cl12_mean, cl34-cl34_mean),
axis=0
)
full_cov_dict[ms1, ms2, ms3, ms4] = cov

cov_dict = {}
for i, field_pair_1 in enumerate(field_pairs):
for j, field_pair_2 in enumerate(field_pairs):

cov_block = cov[i*n_bins:(i+1)*n_bins, j*n_bins:(j+1)*n_bins]
cov_dict[field_pair_1 + field_pair_2] = cov_block

np.savez(
f"{cov_dir}/mc_cov_{ms1}_{ms2}_{ms3}_{ms4}.npz",
**cov_dict
)

if args.plots:
plot_dir = meta.plot_dir_from_output_dir(
meta.covmat_directory_rel
)
n_fields = len(field_pairs)
n_spec = len(cross_ps_names)

full_size = n_spec*n_fields*n_bins
full_cov = np.zeros((full_size, full_size))

for i, (ms1, ms2) in enumerate(cross_ps_names):
for j, (ms3, ms4) in enumerate(cross_ps_names):
if i > j:
continue

full_cov[
i*n_fields*n_bins:(i+1)*n_fields*n_bins,
j*n_fields*n_bins:(j+1)*n_fields*n_bins
] = full_cov_dict[ms1, ms2, ms3, ms4]

# Symmetrize
full_cov = np.triu(full_cov)
full_cov += full_cov.T - np.diag(full_cov.diagonal())
covdiag = full_cov.diagonal()
full_corr = full_cov / np.outer(np.sqrt(covdiag), np.sqrt(covdiag))

plt.figure(figsize=(8, 8))
im = plt.imshow(full_corr, vmin=-1, vmax=1, cmap="RdBu_r")
divider = make_axes_locatable(plt.gca())
cax = divider.append_axes("right", size=0.3, pad=0.1)
plt.colorbar(im, cax=cax)
plt.savefig(f"{plot_dir}/full_corr.png", dpi=300,
bbox_inches="tight")


if __name__ == "__main__":

parser = argparse.ArgumentParser(
description="Covariance matrix calculator"
)
parser.add_argument("--globals", type=str, help="Path to the yaml file")
parser.add_argument("--plots", action="store_true", help="Generate plots")

args = parser.parse_args()
covFeFe(args)
140 changes: 14 additions & 126 deletions pipeline/pcler.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,129 +4,12 @@
import os
from soopercool import BBmeta
from soopercool.utils import get_noise_cls, theory_cls
from soopercool import ps_utils
import pymaster as nmt
from itertools import product
import matplotlib.pyplot as plt
import warnings


def get_coupled_pseudo_cls(fields1, fields2, nmt_binning):
"""
Compute the binned coupled pseudo-C_ell estimates from two
(spin-0 or spin-2) NaMaster fields and a multipole binning scheme.
Parameters
----------
fields1, fields2 : NmtField
Spin-0 or spin-2 fields to correlate.
nmt_binning : NmtBin
Multipole binning scheme.
"""
spins = list(fields1.keys())

pcls = {}
for spin1 in spins:
for spin2 in spins:

f1 = fields1[spin1]
f2 = fields2[spin2]

coupled_cell = nmt.compute_coupled_cell(f1, f2)
coupled_cell = coupled_cell[:, :nmt_binning.lmax+1]

pcls[f"{spin1}x{spin2}"] = nmt_binning.bin_cell(coupled_cell)
return pcls


def decouple_pseudo_cls(coupled_pseudo_cells, coupling_inv):
"""
Decouples the coupled pseudo-C_ell estimators computed between two fields
of spin 0 or 2. Returns decoupled binned power spectra labeled by field
pairs (e.g. 'TT', 'TE', 'EE', 'EB', 'BB' etc.).
Parameters
----------
coupled_pseudo_cells : dict with keys f"spin{s1}xspin{s2}",
items array-like. Coupled pseudo-C_ell estimators.
coupling_inv : array-like
Inverse binned bandpower coupling matrix.
"""
decoupled_pcls = {}
for spin_comb, coupled_pcl in coupled_pseudo_cells.items():
n_bins = coupled_pcl.shape[-1]
decoupled_pcl = coupling_inv[spin_comb] @ coupled_pcl.flatten()
if spin_comb == "spin0xspin0":
size = 1
elif spin_comb in ["spin0xspin2", "spin2xspin0"]:
size = 2
elif spin_comb == "spin2xspin2":
size = 4
decoupled_pcl = decoupled_pcl.reshape((size, n_bins))

decoupled_pcls[spin_comb] = decoupled_pcl

decoupled_pcls = field_pairs_from_spins(decoupled_pcls)

return decoupled_pcls


def field_pairs_from_spins(cls_in_dict):
"""
Reorders power spectrum dictionary with a given input spin
pair into pairs of output (pseudo-)scalar fields on the sky
(T, E, or B).
Parameters
----------
cls_in_dict: dictionary
"""
cls_out_dict = {}

field_spin_mapping = {
"spin0xspin0": ["TT"],
"spin0xspin2": ["TE", "TB"],
"spin2xspin0": ["ET", "BT"],
"spin2xspin2": ["EE", "EB", "BE", "BB"]
}

for spin_pair in cls_in_dict:
for index, field_pair in enumerate(field_spin_mapping[spin_pair]):

cls_out_dict[field_pair] = cls_in_dict[spin_pair][index]

return cls_out_dict


def get_pcls_mat_transfer(fields, nmt_binning):
"""
Compute coupled binned pseudo-C_ell estimates from
pure-E and pure-B transfer function estimation simulations,
and cast them into matrix shape.
Parameters
----------
fields: dictionary of NmtField objects (keys "pureE", "pureB")
nmt_binning: NmtBin object
"""
n_bins = nmt_binning.get_n_bands()
pcls_mat_00 = np.zeros((1, 1, n_bins))
pcls_mat_02 = np.zeros((2, 2, n_bins))
pcls_mat_22 = np.zeros((4, 4, n_bins))

index = 0
cases = ["pureE", "pureB"]
for index, (pure_type1, pure_type2) in enumerate(product(cases, cases)):
pcls = get_coupled_pseudo_cls(fields[pure_type1],
fields[pure_type2],
nmt_binning)
pcls_mat_22[index] = pcls["spin2xspin2"]
pcls_mat_02[cases.index(pure_type2)] = pcls["spin0xspin2"]

pcls_mat_00[0] = pcls["spin0xspin0"]

return {"spin0xspin0": pcls_mat_00,
"spin0xspin2": pcls_mat_02,
"spin2xspin2": pcls_mat_22}


def pcler(args):
"""
Compute all decoupled binned pseudo-C_ell estimates needed for
Expand Down Expand Up @@ -228,13 +111,13 @@ def pcler(args):
coadd=False):
map_set1, id_split1 = map_name1.split("__")
map_set2, id_split2 = map_name2.split("__")
pcls = get_coupled_pseudo_cls(
pcls = ps_utils.get_coupled_pseudo_cls(
fields[map_set1, id_split1],
fields[map_set2, id_split2],
nmt_binning
)

decoupled_pcls = decouple_pseudo_cls(
decoupled_pcls = ps_utils.decouple_pseudo_cls(
pcls, inv_couplings[map_set1, map_set2]
)

Expand Down Expand Up @@ -357,10 +240,14 @@ def pcler(args):
fields["unfiltered"][pure_type] = field
fields["filtered"][pure_type] = field_filtered

pcls_mat_filtered = get_pcls_mat_transfer(fields["filtered"],
nmt_binning)
pcls_mat_unfiltered = get_pcls_mat_transfer(fields["unfiltered"],
nmt_binning)
pcls_mat_filtered = ps_utils.get_pcls_mat_transfer(
fields["filtered"],
nmt_binning
)
pcls_mat_unfiltered = ps_utils.get_pcls_mat_transfer(
fields["unfiltered"],
nmt_binning
)

np.savez(f"{cl_dir}/pcls_mat_tf_est_filtered_{id_sim:04d}.npz",
**pcls_mat_filtered)
Expand Down Expand Up @@ -398,9 +285,10 @@ def pcler(args):
purify_b=meta.tf_est_pure_B)
}

pcls = get_coupled_pseudo_cls(field, field, nmt_binning)
pcls = ps_utils.get_coupled_pseudo_cls(field, field,
nmt_binning)

decoupled_pcls = decouple_pseudo_cls(
decoupled_pcls = ps_utils.decouple_pseudo_cls(
pcls, inv_couplings[filter_flag])

np.savez(f"{cl_dir}/pcls_{cl_type}_{id_sim:04d}_{filter_flag}.npz", # noqa
Expand Down
Loading

0 comments on commit 066b8bd

Please sign in to comment.