Currently, all matrix elements in calc_mcm_spin0and2_pure are calculated separately. But the 3j symbols for m1==m2==0 and m1==-m2==2 that are used in those calculations have symmetries with respect to exchanging l1 and l2. So if calc_coupling_elem_spin0and2_pure was changed to not only return the result for element (l1,l2) but also for element (l2,l1), roughly 30% of run time could be saved.

pspy requires ducc here, but it does not require its installation

I think the following line is probably not intentional, but I want to make sure.

This line causes taskrange to generate a list of floating point numbers for the task ids, because in python 3 the / always generates a float. It should probably be //.

The pasted code below demonstrates that mode-coupling matrices can be computed much more efficiently than it is currently done in pspy. The demo pasted below runs roughly ten times faster with the ducc-based implementation when performed for a single spectrum. Increasing the number of simultaneous computations widens the gap even more.

Please note that this is currently fairly experimental code: I'm collecting feedback on the user interface to make it callable easily from other codes such as pspy, if this should be desired in the future.

The current demo is for spin 0 only, but I hope to have equivalents for spin0and2 and spin0and2_pure ready soon. Their respective speedups likely won't be as dramatic, but should still be around a factor of 5 or so.

@thibautlouis @xzackli

# NOTE: this will not work with the official ducc0 release, but only with
# an experimental version from the "vector_wigners" branch.
# To install, first uninstall possibly existing ducc0 packages, then do
# pip install --no-binary ducc0 --user git+https://gitlab.mpcdf.mpg.de/mtr/ducc.git@vector_wigners

import numpy as np
from time import time

# We have to set nthreads before importing pspy, otherwise the
# environment variable change will have no effect.
nthreads=4
import os
os.environ["OMP_NUM_THREADS"]=str(nthreads)

# This must happen after setting OMP_NUM_THREADS!
from pspy.mcm_fortran.mcm_fortran import mcm_compute as mcm_fortran
import ducc0


def get_random_spectra(nspec, lmax):
    return np.random.normal(size=(nspec, lmax))


# This routine is more complicated than mcm00_ducc, since a few multiplication
# steps are carried out in Python in pspy, and since the array indices are
# a bit different. Overall this should not have noticeable impact on performance
# at higher lmax. 
def mcm00_pspy(spec, lmax):
    nspec = spec.shape[0]
    lrange_spec = np.arange(spec.shape[1])
    res=np.zeros((nspec, lmax+1, lmax+1))
    mcmtmp = np.empty((lmax+1, lmax+1))
    for i in range(nspec):
        wcl = spec[i]*(2*lrange_spec+1)
        mcm_fortran.calc_coupling_spin0(wcl, lmax+1, lmax+1, lmax+1, mcmtmp.T)
        mcm_fortran.fill_upper(mcmtmp.T)
        mcmtmp *= (np.arange(2, lmax+3)*2+1.)/(4*np.pi)
        res[i, 2:, 2:] = mcmtmp[:-2,:-2]
    return res


def mcm00_ducc(spec, lmax):
    return ducc0.misc.coupling_matrix_00(spec, lmax, nthreads=nthreads, algo=1)


# lmax up to which the MCM will be computed
lmax=2000
# number of spectra to process simultaneously
nspec=1

# we generate the spectra up to 2*lmax+1 to use all Wigner 3j symbols
# but this could also be lower.
spec = get_random_spectra(nspec, 2*lmax+1)

t0=time()
mcm_pspy = mcm00_pspy(spec, lmax)
print(f"pspy time: {time()-t0}s")

t0=time()
mcm_ducc = mcm00_ducc(spec, lmax)
print(f"ducc time: {time()-t0}s")

# pspy does not provide the first two rows and columns, so we zero them
# in the ducc result for comparison
mcm_ducc[:,0:2,:]=0
mcm_ducc[:,:,0:2]=0

# compare the results
print(f"L2 error between solutions: {ducc0.misc.l2error(mcm_pspy,mcm_ducc)}")

When playing with the routine calc_mcm_spin0and2_pure I noticed something interesting: no matter which power spectra are fed into the function, the fifth output matrix always seems to contain extremely small values. This can be verified with the following script:

import numpy as np

from pspy.mcm_fortran.mcm_fortran import mcm_compute as mcm_fortran

lmax=100
mcm = np.zeros((5, lmax+1, lmax+1))
spec = np.random.normal(size=(4, 2*lmax+1))
mcm_fortran.calc_mcm_spin0and2_pure(spec[0], spec[1], spec[2], spec[3], mcm.T)
for i in range(5):
    print (np.max(np.abs(mcm[i,:,:])))

Is this due to a bug in the function, or is there actually some kind of symmetry in the Wigner 3j symbols which makes this component vanish by construction? If the latter, it might be good to simply skip its calculation, saving some CPU time and memory.

While running benchmarks on the code I noticed a block in hte Fortran sources that looks dangerous to me: https://github.com/simonsobs/pspy/blob/1d5d8734f489c6a15fe6ee53d8d61b25478fd9b4/pspy/mcm_fortran/mcm_fortran.f90#L77C1-L86C1

        if (i2 < 0) then
            fac_b = 0d0
        end if

        if (i3 < 0) then
            fac_c = 0d0
        end if

        combin = thrcofa(i1) + fac_b * thrcofb(i2) + fac_c * thrcofc(i3)

Here, i2 and i3 can obviously take values below 1, but they are used to access the arrays thrcofb and thrcofc whose first index is 1. This is undefined behaviour and can cause segmentation faults.
Even if this currently doesn't cause any issues, I'd still suggest to avoid these accesses, since they will make debugging of potential other problems much harder.

I think we should flush print message by adding flush=True in the print function (it is mainly relevant for so_mpi file) otherwise message will appear at the end of the mpi jobs

We should preserve wcs when manipulating so_map object. This line