Python interface for calling the C++ code

.gitignore

@@ -29,6 +29,7 @@
 
 # Executables
 cov
+bin
 triple
 *.exe
 *.out
@@ -47,6 +48,11 @@ _build
 _static
 _templates
 
+# Python package build files
+build
+dist
+*.egg-info
+
 # IPython/Jupyter notebook checkpoints
 .ipynb_checkpoints
 

Makefile

@@ -1,8 +1,10 @@
 ## MAKEFILE FOR RascalC. This compiles the grid_covariance.cpp file into the ./cov exececutable.
 
 CC = gcc
-CFLAGS = -g -O3 -Wall -MMD
-CXXFLAGS = -O3 -Wall -MMD -DOPENMP -DLEGENDRE_MIX -DJACKKNIFE -DPRINTPERCENTS
+CFLAGS = -O3 -Wall -MMD
+CXX = g++
+CXXFLAGS	= -O3 -Wall -MMD -std=c++11 -ffast-math $(shell pkg-config --cflags gsl)
+CXXFLAGS	+= -DOPENMP -DLEGENDRE_MIX -DJACKKNIFE -DPRINTPERCENTS
 #-DOPENMP  # use this to run multi-threaded with OPENMP
 #-DPERIODIC # use this to enable periodic behavior
 #-DLEGENDRE # use this to compute 2PCF covariances in Legendre bins (original mode, corresponding to direct accumulation into multipoles from pair counts)
@@ -12,25 +14,26 @@ CXXFLAGS = -O3 -Wall -MMD -DOPENMP -DLEGENDRE_MIX -DJACKKNIFE -DPRINTPERCENTS
 #-DTHREE_PCF # use this to compute 3PCF autocovariances
 #-DPRINTPERCENTS # use this to print percentage of progress in each loop. This can be a lot of output
 
+LFLAGS	= $(shell pkg-config --libs gsl) # common part
+
 # Known OS-specific choices
 ifeq ($(shell uname -s),Darwin)
-# Here we use LLVM compiler to load the Mac OpenMP. Tested after installation commands:
-# brew install llvm
+# Here we load the Mac OpenMP. Tested after installation commands:
 # brew install libomp
 # This may need to be modified with a different installation
 ifndef HOMEBREW_PREFIX
 HOMEBREW_PREFIX = /usr/local
 endif
-CXX = ${HOMEBREW_PREFIX}/opt/llvm/bin/clang++ -std=c++0x -fopenmp -ffast-math $(shell pkg-config --cflags gsl)
-LD	= ${HOMEBREW_PREFIX}/opt/llvm/bin/clang++
-LFLAGS	= $(shell pkg-config --libs gsl) -fopenmp -lomp
+CXXFLAGS += -I$(HOMEBREW_PREFIX)/opt/libomp/include
+LFLAGS	+= -L$(HOMEBREW_PREFIX)/opt/libomp/lib -lomp
 else
 # default (Linux) case
-CXX = g++ -fopenmp -lgomp -std=c++0x -ffast-math $(shell pkg-config --cflags gsl)
-LD	= g++
-LFLAGS	= -L/usr/local/lib -L/usr/lib/x86_64-linux-gnu $(shell pkg-config --libs gsl) -lgomp
+CXXFLAGS += -fopenmp
+LFLAGS	+= -lgomp
 endif
 
+LD = $(CXX)
+
 AUNTIE	= cov
 AOBJS	= grid_covariance.o ./cubature/hcubature.o ./ransampl/ransampl.o
 ADEPS   = ${AOBJS:.o=.d}

README.md

@@ -9,6 +9,61 @@ Any usage of this code should cite [Philcox et al. 2019](https://arxiv.org/abs/1
 
 ***New for version 2***: Legendre moment covariances and the 3PCF
 
+## Python interface and [Jupyter notebook tutorial](https://github.com/misharash/RascalC/blob/master/tutorial.ipynb) (alpha-testing before version 3.0)
+
+### Installation for DESI members at NERSC
+
+Recommended to use with `cosmodesi` environment.
+In particular, load it before installing:
+```
+source /global/common/software/desi/users/adematti/cosmodesi_environment.sh main
+pip install -e /global/common/software/desi/users/mrash/RascalC
+```
+This installs the library from my software folder in the development mode, so that after I update it e.g. with some fix, you will have the new version without the need to re-install or any other action.
+
+### Generic installation (if the above is unavailable)
+
+```
+git clone https://github.com/misharash/RascalC
+cd RascalC
+python3 -m pip install .
+```
+Make sure to reinstall after pulling updates.
+At this early stage, the fixes might be needed quite often.
+
+Or you can install with just one command:
+```
+python3 -m pip install https://github.com/misharash/RascalC.git
+```
+
+The code requires [`pycorr`](https://github.com/cosmodesi/pycorr) to deal with pair counts and data correlation function estimators.
+To compute pair counts of catalogs, you need a [custom version of `Corrfunc`](https://github.com/adematti/Corrfunc) (see also [`pycorr` installation instructions](https://py2pcf.readthedocs.io/en/latest/user/building.html)).
+Both can be installed quickly via
+```
+python3 -m pip install 'git+https://github.com/cosmodesi/pycorr#egg=pycorr[corrfunc]'
+```
+
+### Usage guide
+
+```
+import RascalC
+result = RascalC.run_cov(...)
+```
+
+`run_cov` is the main function for the covariance matrix computation.
+Use `help(RascalC.run_cov)` to learn more about the inputs and outputs; many of them are similar to [pycorr](https://github.com/cosmodesi/pycorr) `TwoPointCorrelationFunction` and some others are `pycorr.TwoPointEstimator`s.
+
+An example pipeline is showcased in a [tutorial notebook](https://github.com/misharash/RascalC/blob/master/tutorial.ipynb) (will be smoothed out soon).
+
+It is strongly recommended NOT to use multi-threaded operations in the `python` process before launching `RascalC.run_cov` – this may cause the code to run effectively single-threaded.
+E.g. at NERSC this would mean not setting `OMP_*` and other `*_THREADS` environment variables; the code will set them by itself according to the number of threads you passed.
+This caveat does not seem to be unique for RascalC – different multi-threading backends can interfere.
+
+Some specific examples are also available in the new separate script gallery: <https://github.com/misharash/RascalC-scripts>.
+
+More documentation is coming, in the meantime please contact Michael 'Misha' Rashkovetskyi <[email protected]> with any questions.
+Please also feel free to open [GitHub issues](https://github.com/misharash/RascalC/issues) both for problems and clarification requests.
+
 ## Authors
 
 - Oliver Philcox (Columbia / Simons Foundation)
@@ -17,4 +72,5 @@ Any usage of this code should cite [Philcox et al. 2019](https://arxiv.org/abs/1
 - Alexander Wiegand (Garching)
 - Misha Rashkovetskyi (Harvard)
 
-We thank Yuting Wang and Ryuichiro Hada for pointing out and fixing a number of issues with the code and its documentation. We are particularly grateful to Uendert Andrade for finding a wide variety of improvements and bugs!
+We thank Yuting Wang and Ryuichiro Hada for pointing out and fixing a number of issues with the code and its documentation.
+We are particularly grateful to Uendert Andrade for finding a wide variety of improvements and bugs, and to Jiaxi Yu for feedback on the Python interface and Jupyter notebook tutorial!

RascalC/Makefile

@@ -0,0 +1,102 @@
+## MAKEFILE FOR RascalC. This compiles the grid_covariance.cpp file into the ./cov exececutable.
+
+CC = gcc
+CFLAGS = -O3 -Wall -MMD
+CXX = g++
+CXXFLAGS	= -O3 -Wall -MMD -std=c++11 -ffast-math $(shell pkg-config --cflags gsl)
+CXXFLAGS	+= -DOPENMP
+#-DOPENMP  # use this to run multi-threaded with OPENMP
+#-DPERIODIC # use this to enable periodic behavior
+#-DLEGENDRE # use this to compute 2PCF covariances in Legendre bins (original mode, corresponding to direct accumulation into multipoles from pair counts)
+#-DLEGENDRE_MIX # also compute 2PCF covariances in Legendre bins, but in other, "mixed" mode, corresponding to projection of s,µ bins into multipoles
+# without either of the two Legendre flags above, the covariance is computed in s,µ bins
+#-DJACKKNIFE # use this to compute (r,mu)-space 2PCF covariances and jackknife covariances. Incompatible with -DLEGENDRE but works with -DLEGENDRE_MIX
+#-DTHREE_PCF # use this to compute 3PCF autocovariances
+#-DPRINTPERCENTS # use this to print percentage of progress in each loop. This can be a lot of output
+
+LFLAGS	= $(shell pkg-config --libs gsl) # common part
+
+# Known OS-specific choices
+ifeq ($(shell uname -s),Darwin)
+# Here we load the Mac OpenMP. Tested after installation commands:
+# brew install libomp
+# This may need to be modified with a different installation
+ifndef HOMEBREW_PREFIX
+HOMEBREW_PREFIX = /usr/local
+endif
+CXXFLAGS += -I$(HOMEBREW_PREFIX)/opt/libomp/include
+LFLAGS	+= -L$(HOMEBREW_PREFIX)/opt/libomp/lib -lomp
+else
+# default (Linux) case
+CXXFLAGS += -fopenmp
+LFLAGS	+= -lgomp
+endif
+
+LD = $(CXX)
+
+# The code below compiles all the valid variants for the 2PCF covariance
+
+BASE_VARIANTS	= s_mu s_mu_jackknife legendre_accumulated legendre_projected legendre_projected_jackknife
+
+DEFINES_FOR_s_mu=
+DEFINES_FOR_s_mu_jackknife=-DJACKKNIFE
+DEFINES_FOR_legendre_accumulated=-DLEGENDRE
+DEFINES_FOR_legendre_projected=-DLEGENDRE_MIX
+DEFINES_FOR_legendre_projected_jackknife=-DLEGENDRE_MIX -DJACKKNIFE
+
+define add_periodic_variant
+DEFINES_FOR_$(1)_periodic	= $$(DEFINES_FOR_$(1)) -DPERIODIC
+VARIANTS_INT	+= $(1) $(1)_periodic
+endef
+
+$(foreach variant,$(BASE_VARIANTS),$(eval $(call add_periodic_variant,$(variant))))
+
+define add_verbose_variant
+DEFINES_FOR_$(1)_verbose	= $$(DEFINES_FOR_$(1)) -DPRINTPERCENTS
+VARIANTS	+= $(1) $(1)_verbose
+endef
+
+$(foreach variant,$(VARIANTS_INT),$(eval $(call add_verbose_variant,$(variant))))
+
+OBJDIR	= obj
+EXECDIR	= bin
+
+VAR_SRC_BASE	= grid_covariance
+VAR_SRC	= ../$(VAR_SRC_BASE).cpp
+VAR_OBJ_BASE	= $(OBJDIR)/$(VAR_SRC_BASE)
+
+EXEC_BASE	= $(EXECDIR)/cov
+COMMON_OBJS	= ../cubature/hcubature.o ../ransampl/ransampl.o
+
+ALL_OBJS	= $(COMMON_OBJS)
+
+define make_targets
+OBJ_$(1)	= $$(VAR_OBJ_BASE).$(1).o
+ALL_OBJS	+= $$(OBJ_$(1))
+EXEC_$(1)	= $$(EXEC_BASE).$(1)
+ALL_EXECS	+= $$(EXEC_$(1))
+
+$$(OBJ_$(1)): $$(VAR_SRC)
+	mkdir -p $$(OBJDIR)
+	$$(CXX) $$(CXXFLAGS) $$(DEFINES_FOR_$(1)) -c $$(VAR_SRC) -o $$@
+
+$$(EXEC_$(1)): $$(OBJ_$(1)) $$(COMMON_OBJS)
+	mkdir -p $$(EXECDIR)
+	$$(LD) $$(OBJ_$(1)) $$(COMMON_OBJS) $$(LFLAGS) -o $$@
+
+endef
+
+$(foreach variant,$(VARIANTS),$(eval $(call make_targets,$(variant))))
+
+.PHONY: all clean
+
+all: $(ALL_EXECS)
+.DEFAULT_GOAL	:= all
+
+clean:
+	rm -f $(ALL_EXECS) $(ALL_OBJS) ${ALL_DEPS}
+	rm -rf $(OBJDIR) $(EXECDIR)
+
+$(ALL_OBJS): Makefile
+ALL_DEPS	= ${ALL_OBJS:.o=.d}
+-include ${ALL_DEPS}

RascalC/__init__.py

@@ -0,0 +1,3 @@
+__version__ = "3.0.alpha1"
+
+from .interface import run_cov

RascalC/comb/combine_covs.py

@@ -0,0 +1,28 @@
+"This reads two sets of RascalC results and two cosmodesi/pycorr .npy files to combine two covs following NS/GCcomb procedure. Covariance of N and S 2PCF is neglected."
+
+from pycorr import TwoPointCorrelationFunction
+import numpy as np
+from ..pycorr_utils.utils import reshape_pycorr
+from ..cov_utils import get_cov_header, load_cov
+from ..pycorr_utils.counts import get_counts_from_pycorr
+
+
+def combine_covs(rascalc_results1: str, rascalc_results2: str, pycorr_file1: str, pycorr_file2: str, output_cov_file: str, n_mu_bins: int | None = None, r_step: float = 1, skip_r_bins: int | tuple[int, int] = 0, output_cov_file1: str | None = None, output_cov_file2: str | None = None, print_function = print):
+    # Read RascalC results
+    header1 = get_cov_header(rascalc_results1)
+    cov1 = load_cov(rascalc_results1, print_function)
+    header2 = get_cov_header(rascalc_results2)
+    cov2 = load_cov(rascalc_results2, print_function)
+    # Save to their files if any
+    if output_cov_file1: np.savetxt(output_cov_file1, cov1, header = header1)
+    if output_cov_file2: np.savetxt(output_cov_file2, cov2, header = header2)
+    header = f"combined from {rascalc_results1} with {header1} and {rascalc_results2} with {header2}" # form the final header to include both
+
+    # Read pycorr files to figure out weights
+    weight1 = get_counts_from_pycorr(reshape_pycorr(TwoPointCorrelationFunction.load(pycorr_file1), n_mu_bins, r_step, skip_r_bins = skip_r_bins).normalize(), counts_factor = 1).ravel()
+    weight2 = get_counts_from_pycorr(reshape_pycorr(TwoPointCorrelationFunction.load(pycorr_file2), n_mu_bins, r_step, skip_r_bins = skip_r_bins).normalize(), counts_factor = 1).ravel()
+
+    # Produce and save combined cov
+    # following xi = (xi1 * weight1 + xi2 * weight2) / (weight1 + weight2)
+    cov = (cov1 * weight1[None, :] * weight1[:, None] + cov2 * weight2[None, :] * weight2[:, None]) / (weight1 + weight2)[None, :] / (weight1 + weight2)[:, None]
+    np.savetxt(output_cov_file, cov, header = header) # includes source parts and their shot-noise rescaling values in the header

RascalC/comb/combine_covs_legendre.py

@@ -0,0 +1,47 @@
+"This reads two sets of RascalC results and two cosmodesi/pycorr .npy files to combine two covs following NS/GCcomb procedure in Legendre mode. Covariance of N and S 2PCF is neglected."
+
+from pycorr import TwoPointCorrelationFunction
+import numpy as np
+from ..pycorr_utils.utils import reshape_pycorr
+from ..cov_utils import get_cov_header, load_cov_legendre
+from ..pycorr_utils.counts import get_counts_from_pycorr
+from ..mu_bin_legendre_factors import compute_mu_bin_legendre_factors
+
+
+def combine_covs_legendre(rascalc_results1: str, rascalc_results2: str, pycorr_file1: str, pycorr_file2: str, output_cov_file: str, max_l: int, r_step: float = 1, skip_r_bins: int | tuple[int, int] = 0, output_cov_file1: str | None = None, output_cov_file2: str | None = None, print_function = print):
+    # Read RascalC results
+    header1 = get_cov_header(rascalc_results1)
+    cov1 = load_cov_legendre(rascalc_results1, max_l, print_function)
+    n_bins = len(cov1)
+    header2 = get_cov_header(rascalc_results2)
+    cov2 = load_cov_legendre(rascalc_results2, max_l, print_function)
+    # Save to their files if any
+    if output_cov_file1: np.savetxt(output_cov_file1, cov1, header = header1)
+    if output_cov_file2: np.savetxt(output_cov_file2, cov2, header = header2)
+    header = f"combined from {rascalc_results1} with {header1} and {rascalc_results2} with {header2}" # form the final header to include both
+
+    # Read pycorr files to figure out weights of s, mu binned 2PCF
+    xi_estimator1 = reshape_pycorr(TwoPointCorrelationFunction.load(pycorr_file1), n_mu = None, r_step = r_step, skip_r_bins = skip_r_bins).normalize()
+    n_r_bins = xi_estimator1.shape[0]
+    mu_edges = xi_estimator1.edges[1]
+    weight1 = get_counts_from_pycorr(xi_estimator1, counts_factor = 1)
+    weight2 = get_counts_from_pycorr(reshape_pycorr(TwoPointCorrelationFunction.load(pycorr_file2).normalize(), n_mu = None, r_step = r_step, skip_r_bins = skip_r_bins), counts_factor = 1)
+
+    # Normalize weights
+    sum_weight = weight1 + weight2
+    weight1 /= sum_weight
+    weight2 /= sum_weight
+
+    mu_leg_factors, leg_mu_factors = compute_mu_bin_legendre_factors(mu_edges, max_l, do_inverse = True)
+
+    # Derivatives of angularly binned 2PCF wrt Legendre are leg_mu_factors[ell//2, mu_bin]
+    # Angularly binned 2PCF are added with weights (normalized) weight1/2[r_bin, mu_bin]
+    # Derivatives of Legendre wrt binned 2PCF are mu_leg_factors[mu_bin, ell//2]
+    # So we need to sum such product over mu bins, while radial bins stay independent, and the partial derivative of combined 2PCF wrt the 2PCFs 1/2 will be
+    pd1 = np.einsum('il,kl,lj,km->ikjm', leg_mu_factors, weight1, mu_leg_factors, np.eye(n_r_bins)).reshape(n_bins, n_bins)
+    pd2 = np.einsum('il,kl,lj,km->ikjm', leg_mu_factors, weight2, mu_leg_factors, np.eye(n_r_bins)).reshape(n_bins, n_bins)
+    # We have correct [l_in, r_in, l_out, r_out] ordering and want to make these matrices in the end thus the reshape
+
+    # Produce and save combined cov
+    cov = pd1.T.dot(cov1).dot(pd1) + pd2.T.dot(cov2).dot(pd2)
+    np.savetxt(output_cov_file, cov, header=header) # includes source parts and their shot-noise rescaling values in the header

RascalC/comb/combine_covs_legendre_multi.py

@@ -0,0 +1,55 @@
+"This reads two sets of RascalC results and two triplets of cosmodesi/pycorr .npy files to combine two full 2-tracer covs following NS/GCcomb procedure for 2 tracers in Legendre mode. Covariance of N(GC) and S(GC) 2PCF is neglected."
+
+from pycorr import TwoPointCorrelationFunction
+import numpy as np
+from ..pycorr_utils.utils import reshape_pycorr
+from ..cov_utils import get_cov_header, load_cov_legendre_multi
+from ..pycorr_utils.counts import get_counts_from_pycorr
+from ..mu_bin_legendre_factors import compute_mu_bin_legendre_factors
+
+
+def combine_covs_legendre_multi(rascalc_results1: str, rascalc_results2: str, pycorr_files1: list[str], pycorr_files2: list[str], output_cov_file: str, max_l: int, r_step: float = 1, skip_r_bins: int | tuple[int, int] = 0, output_cov_file1: str | None = None, output_cov_file2: str | None = None, print_function = print):
+    # Read RascalC results
+    header1 = get_cov_header(rascalc_results1)
+    cov1 = load_cov_legendre_multi(rascalc_results1, max_l, print_function)
+    n_bins = len(cov1)
+    header2 = get_cov_header(rascalc_results2)
+    cov2 = load_cov_legendre_multi(rascalc_results2, max_l, print_function)
+    # Save to their files if any
+    if output_cov_file1: np.savetxt(output_cov_file1, cov1, header = header1)
+    if output_cov_file2: np.savetxt(output_cov_file2, cov2, header = header2)
+    header = f"combined from {rascalc_results1} with {header1} and {rascalc_results2} with {header2}" # form the final header to include both
+
+    # Read pycorr files to figure out weights
+    weight1 = []
+    for pycorr_file1 in pycorr_files1:
+        xi_estimator1 = reshape_pycorr(TwoPointCorrelationFunction.load(pycorr_file1), n_mu = None, r_step = r_step, skip_r_bins = skip_r_bins).normalize()
+        weight1.append(get_counts_from_pycorr(xi_estimator1, counts_factor = 1))
+    weight1 = np.array(weight1)
+
+    n_r_bins = xi_estimator1.shape[0]
+    mu_edges = xi_estimator1.edges[1]
+
+    weight2 = []
+    for pycorr_file2 in pycorr_files2:
+        weight2.append(get_counts_from_pycorr(reshape_pycorr(TwoPointCorrelationFunction.load(pycorr_file2), n_mu = None, r_step = r_step, skip_r_bins = skip_r_bins).normalize(), counts_factor = 1))
+    weight2 = np.array(weight2)
+
+    # Normalize weights
+    sum_weight = weight1 + weight2
+    weight1 /= sum_weight
+    weight2 /= sum_weight
+
+    mu_leg_factors, leg_mu_factors = compute_mu_bin_legendre_factors(mu_edges, max_l, do_inverse = True)
+
+    # Derivatives of angularly binned 2PCF wrt Legendre are leg_mu_factors[ell//2, mu_bin]
+    # Angularly binned 2PCF are added with weights (normalized) weight1/2[tracer, r_bin, mu_bin]
+    # Derivatives of Legendre wrt binned 2PCF are mu_leg_factors[mu_bin, ell//2]
+    # So we need to sum such product over mu bins, while tracers and radial bins stay independent, and the partial derivative of combined 2PCF wrt the 2PCFs 1/2 will be
+    pd1 = np.einsum('il,tkl,lj,km,tr->tikrjm', leg_mu_factors, weight1, mu_leg_factors, np.eye(n_r_bins), np.eye(3)).reshape(n_bins, n_bins)
+    pd2 = np.einsum('il,tkl,lj,km,tr->tikrjm', leg_mu_factors, weight2, mu_leg_factors, np.eye(n_r_bins), np.eye(3)).reshape(n_bins, n_bins)
+    # We have correct [t_in, l_in, r_in, t_out, l_out, r_out] ordering and want to make these matrices in the end thus the reshape
+
+    # Produce and save combined cov
+    cov = pd1.T.dot(cov1).dot(pd1) + pd2.T.dot(cov2).dot(pd2)
+    np.savetxt(output_cov_file, cov, header=header) # includes source parts and their shot-noise rescaling values in the header