update ksw/cosmo.py/compute_transfer_tensor

Makefile

@@ -42,8 +42,9 @@ TEST_OBJECTS = $(TDIR)/obj/seatest.o \
                $(TDIR)/obj/run_tests.o
 
 LINK_COMMON = -lm
-LINK_FFTW = -L$(FFTWROOT) -lfftw3 -lfftw3f
-LINK_MKL = -L${MKLROOT}/lib/intel64 -lmkl_rt -Wl,--no-as-needed -lpthread -lm -ldl
+LINK_FFTW = -L$(FFTW_HOME)/lib -lfftw3 -lfftw3f -Wl,-rpath,${FFTW_HOME}/lib
+CFLAGS_FFTW = -I"${FFTW_HOME}/include"
+LINK_MKL = -L${MKLROOT}/lib -lmkl_rt -Wl,--no-as-needed,-rpath,${MKLROOT}/lib -lpthread -lm -ldl
 CFLAGS_MKL = -m64  -I"${MKLROOT}/include"
 
 all: $(LDIR)/libradial_functional.so ${LDIR}/libksw_estimator.so ${LDIR}/libksw_fisher.so
@@ -64,25 +65,25 @@ $(ODIR)/hyperspherical.o: $(HS_SDIR)/hyperspherical.c $(HS_IDIR)/*.h
 	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(HS_IDIR)
 
 $(LDIR)/libksw_estimator.so: $(EST_OBJECTS) $(LP_OBJECTS) $(IDIR)/ksw_estimator.h
-	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(OMPFLAG) $(OPTFLAG) -shared -o $@ $< $(LP_OBJECTS) $(LINK_COMMON) $(LINK_FFTW) $(LINK_MKL) -lgomp
+	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(CFLAGS_FFTW) $(OMPFLAG) $(OPTFLAG) -shared -o $@ $< $(LP_OBJECTS) $(LINK_COMMON) $(LINK_FFTW) $(LINK_MKL) -lgomp
 
 $(LDIR)/libksw_fisher.so: $(FIS_OBJECTS) $(LP_OBJECTS) $(IDIR)/ksw_fisher.h
-	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(OMPFLAG) $(OPTFLAG) -shared -o $@ $< $(LP_OBJECTS) $(LINK_COMMON) $(LINK_FFTW) $(LINK_MKL) -lgomp
+	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(CFLAGS_FFTW) $(OMPFLAG) $(OPTFLAG) -shared -o $@ $< $(LP_OBJECTS) $(LINK_COMMON) $(LINK_FFTW) $(LINK_MKL) -lgomp
 
 $(ODIR)/estimator.o: $(SDIR)/estimator.c $(IDIR)/*.h
-	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) $(LINK_COMMON) $(LINK_FFTW) $(LINK_MKL) -lgomp
+	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(CFLAGS_FFTW) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR)
 
 $(ODIR)/fisher.o: $(SDIR)/fisher.c $(IDIR)/*.h
-	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) $(LINK_COMMON) $(LINK_MKL)
+	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) $(LINK_COMMON)
 
 $(ODIR)/ylmgen_c.o: $(LP_SDIR)/ylmgen_c.c $(LP_IDIR)/*.h
-	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) $(LINK_COMMON) 
+	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR)
 
 $(ODIR)/c_utils.o: $(LP_SDIR)/c_utils.c $(LP_IDIR)/*.h
-	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) $(LINK_COMMON) 
+	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR)
 
 $(ODIR)/walltime_c.o: $(LP_SDIR)/walltime_c.c $(LP_IDIR)/*.h
-	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) $(LINK_COMMON) 
+	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR)
 
 check: check_c check_python
 
@@ -93,7 +94,7 @@ check_python:
 	cd $(TDIR); python -m pytest python/
 
 $(TDIR)/bin/run_tests: $(TEST_OBJECTS) $(LDIR)/libradial_functional.so $(LDIR)/libksw_estimator.so $(LDIR)/libksw_fisher.so
-	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -o $@ $(TEST_OBJECTS) -L$(LDIR) -lradial_functional -lksw_estimator -lksw_fisher $(LINK_COMMON) $(LINK_FFTW) $(LINK_MKL) -lgomp -Wl,-rpath,$(LDIR)
+	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(CFLAGS_FFTW) $(OMPFLAG) $(OPTFLAG) -o $@ $(TEST_OBJECTS) -L$(LDIR) -lradial_functional -lksw_estimator -lksw_fisher $(LINK_COMMON) $(LINK_FFTW) $(LINK_MKL) -lgomp -Wl,-rpath,$(LDIR)
 
 $(TDIR)/obj/run_tests.o: $(TDIR)/src/run_tests.c $(TDIR)/include/seatest.h
 	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(TDIR)/include -I$(IDIR) -I$(LP_IDIR)
@@ -105,10 +106,10 @@ $(TDIR)/obj/test_radial_functional.o: $(TDIR)/src/test_radial_functional.c $(IDI
 	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(HS_IDIR) -I$(TDIR)/include
 
 $(TDIR)/obj/test_estimator.o: $(TDIR)/src/test_estimator.c $(IDIR)/*.h
-	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) -I$(TDIR)/include
+	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(CFLAGS_FFTW) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) -I$(TDIR)/include
 
 $(TDIR)/obj/test_fisher.o: $(TDIR)/src/test_fisher.c $(IDIR)/*.h
-	$(CC) $(CFLAGS) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) -I$(TDIR)/include
+	$(CC) $(CFLAGS) $(CFLAGS_MKL) $(CFLAGS_FFTW) $(OMPFLAG) $(OPTFLAG) -c -o $@ $< -I$(IDIR) -I$(LP_IDIR) -I$(TDIR)/include
 
 clean:
 	rm -rf $(ODIR)

README.md

@@ -23,9 +23,10 @@ An implementation of the Komatsu-Spergel-Wandelt bispectrum estimator for modern
 
 Start by making sure the MKL library is loaded in your environment. On most clusters this can be achieved by loading a predefined module. On the Princeton `della` and `tiger` clusters you can use `module load intel-mkl` (see [here](https://researchcomputing.princeton.edu/faq/how-to-build-using-intel-mkl) for more information). On `NERSC` you can use `load intel` (see [here](https://docs-dev.nersc.gov/cgpu/software/math/)). Once you have loaded the module, check if the `MKLROOT` environment variable has been set (`echo $MKLROOT`).
 
+Then do the same for the FFTW library. Right now, the Makefile assumes that the `FFTW_HOME` environment variable is set. So that make sure that is the case and that `FFTW_HOME/lib` and `FFTW_HOME/include` point to the correct directories.
 
 
-Once the MKL environment has been set, `git clone` this repository, go into the directory and run:
+Once the MKL and FFTW environments have been set, `git clone` this repository, go into the directory and run:
 
 
 ```

cython/radial_functional.pxd

@@ -10,3 +10,4 @@ cdef extern from "radial_functional.h":
 			 int nr,
 			 int npol,
 			 int ncomp);
+

cython/radial_functional.pyx

@@ -70,8 +70,9 @@ def radial_func(f_k, tr_ell_k, k, radii, ells):
                          npol,
                          ncomp)
     return f_ell_r
+
 
-def check_and_return_shape(arr, exp_shape):
+def check_and_return_shape(arr, exp_shape):    
     '''
     Return shape of input array if it matches expections.
 
@@ -109,3 +110,91 @@ def check_and_return_shape(arr, exp_shape):
                     idx, n_exp, n))            
     return arr.shape
 
+
+def radial_func_dL(f_k, tr_ell_k, k, radii, ells):
+    '''
+    Compute f_ell^X(r) = (2/pi) int k^2 dk f(k) transfer^X_ell(k) j_ell(k r),
+    where f(k) is an arbitrary function of wavenumber k.
+
+    Parameters
+    ---------
+    f_k : (nk, ncomp) array
+        Input functions.
+    tr_ell_k : (nell, nk, npol) array
+        Transfer functions.
+    k : (nk) array
+        Wavenumbers in 1/Mpc.
+    radii : (nr) array
+        Multipoles
+
+    Returns
+    -------
+    f_ell_r : (nr, nell, npol, ncomp, delta_L_list) array
+        Evaluated integral for all radii, multipoles, polarizations 
+        and input function components.
+
+    Raises
+    ------
+    ValueError
+        If input array shapes are incorrect.
+        If the input arrays contain nans or infs.
+    '''
+
+    # Check input for nans and infs.
+    f_k = np.asarray_chkfinite(f_k, dtype=float, order='C')
+    tr_ell_k = np.asarray_chkfinite(tr_ell_k, dtype=float, order='C')
+    k = np.asarray_chkfinite(k, dtype=float, order='C')
+    radii = np.asarray_chkfinite(radii, dtype=float, order='C')
+    ells = np.asarray_chkfinite(ells, dtype=np.int32, order='C')
+    delta_L_list = np.asarray_chkfinite([-2, -1, 0, 1, 2], dtype=np.int32, order='C')
+
+    # Check for input shapes.
+    nk, = check_and_return_shape(k, (None,))
+    nr, = check_and_return_shape(radii, (None,))
+    nell, = check_and_return_shape(ells, (None,))
+    nk, ncomp = check_and_return_shape(f_k, (nk, None))
+    nell, nk, npol = check_and_return_shape(tr_ell_k, (nell, nk, None))
+    ndeltaL, = check_and_return_shape(delta_L_list, (None,))
+
+    # Create output array.
+    f_ell_r = np.empty((ndeltaL, nr, nell, npol, ncomp), dtype=float)
+
+    cdef double [::1] f_k_ = f_k.reshape(-1)
+    cdef double [::1] tr_ell_k_ = tr_ell_k.reshape(-1)
+    cdef double [::1] k_ = k.reshape(-1)
+    cdef double [::1] radii_ = radii.reshape(-1)
+    cdef double [::1] f_ell_r_ = f_ell_r.reshape(-1)
+    cdef int [::1] ells_= ells
+
+    # Define local variables
+    cdef Py_ssize_t i, index        # loop indices (Py_ssize_t = signed size type)
+    cdef int L, delta_L             # temporary variables for ell and its shifted version
+
+    # Preallocate a temporary buffer to hold valid shifted ell values   
+    cdef np.ndarray[np.int32_t, ndim=1, mode='c'] L_arr = np.empty(nell, dtype=np.int32)
+    cdef int[::1] L_buf = L_arr    # contiguous int memoryview (to pass to the C function)
+
+    # Loop over each delta_L value
+    for index in range(ndeltaL):
+        delta_L = delta_L_list[index]
+
+        # build the shifted ell list (L = ell+delta_L), excluding negative results
+        for i in range(nell):
+            L = ells_[i] + delta_L
+            if L >= 0:
+                L_buf[i] = L
+            else:
+                L_buf[i] = 0      
+
+        radial_func = compute_radial_func(&f_k_[0],
+                                         &tr_ell_k_[0],
+                                         &k_[0],
+                                         &radii_[0],
+                                         &f_ell_r_[index * (nr * nell * npol * ncomp)],
+                                         &L_buf[0],
+                                         nk,
+                                         nell,
+                                         nr,
+                                         npol,
+                                         ncomp)
+    return f_ell_r

include/radial_functional.h

@@ -81,3 +81,37 @@ void _trapezoidal_weights(double const *k,
  */
 
 void * _malloc_checked(size_t size);
+
+
+/* 
+* Compute f_ell^X(r) = (2/pi) int k^2 dk f(k) transfer^X_ell(k) j_ell(k r)
+* where f(k) is an arbitrary function of wavenumber k.
+* 
+* Arguments
+* ---------
+* f_k      : (nk * ncomp) input functions.
+* tr_ell_k : (nell * nk * npol) transfer function.
+* k        : (nk) array of monotonically increasing wavenumbers.
+* radii    : (nr) array of radii
+* f_r_ell  : (nr * nell * npol * ncomp) output array.
+* ells     : (nell) multipoles.
+* nk       : Number of wavenumbers.
+* nell     : Number of multipoles.
+* nr       : Number of radii.
+* npol     : Number of polarization components (1=T, 2= E, 3=B).
+* ncomp    : Number of input functions.
+* delta_l  : shift in ell for transfer functions
+*/
+
+void compute_radial_func_dL2(double const *f_k,
+                            double const *tr_ell_k,
+                            double *k,
+                            double const *radii,
+                            double *f_r_ell,
+                            int *ells,
+                            int nk,
+                            int nell,
+                            int nr,
+                            int npol,
+                            int ncomp,
+                            int delta_l);      // New: delta ell