Big potential speedups for coupling matrices about pspy HOT 6 CLOSED

# 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=1
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


# 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 mcm02_pspy(spec, lmax):
    nspec = spec.shape[0]
    lrange_spec = np.arange(spec.shape[2])
    res=np.zeros((nspec, 5, lmax+1, lmax+1))
    mcmtmp = np.empty((5, lmax+1, lmax+1))
    for i in range(nspec):
        wcl = spec[i]*((2*lrange_spec+1).reshape((1,-1)))
        mcm_fortran.calc_coupling_spin0and2(wcl[0], wcl[1], wcl[2], wcl[3], lmax+1, lmax+1, lmax+1, mcmtmp.T)
        for j in range(5):
            mcm_fortran.fill_upper(mcmtmp[j].T)
        mcmtmp *= (np.arange(2, lmax+3)*2+1.)/(4*np.pi)
        res[i, :, 2:, 2:] = mcmtmp[:,:-2,:-2]
    return res

def mcm02_pure_pspy(spec, lmax):
    nspec = spec.shape[0]
    lrange_spec = np.arange(spec.shape[2])
    res=np.zeros((nspec, 5, lmax+1, lmax+1))
    mcmtmp = np.empty((5, lmax+1, lmax+1))
    for i in range(nspec):
        wcl = spec[i]*((2*lrange_spec+1).reshape((1,-1)))
        mcm_fortran.calc_mcm_spin0and2_pure(wcl[0], wcl[1], wcl[2], wcl[3], 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)

def mcm02_ducc(spec, lmax):
    return ducc0.misc.coupling_matrix_spin0and2(spec, lmax, nthreads=nthreads, algo=1)

def mcm02_pure_ducc(spec, lmax):
    return ducc0.misc.coupling_matrix_spin0and2_pure(spec, lmax, nthreads=nthreads, algo=1)

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

print()
print("Mode coupling matrix computation comparison")
print(f"nspec={nspec}, lmax={lmax}, nthreads={nthreads}")

print()
print("Spin 0 case:")
# we generate the spectra up to 2*lmax+1 to use all Wigner 3j symbols
# but this could also be lower.
spec = np.random.normal(size=(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)}")

print()
print("Spin 0and2 case:")
# we generate the spectra up to 2*lmax+1 to use all Wigner 3j symbols
# but this could also be lower.
spec = np.random.normal(size=(nspec, 4, 2*lmax+1))

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

t0=time()
mcm_ducc = mcm02_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)}")

print()
print("Spin 0and2_pure case:")
# we generate the spectra up to 2*lmax+1 to use all Wigner 3j symbols
# but this could also be lower.
spec = np.random.normal(size=(nspec, 4, 2*lmax+1))

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

t0=time()
mcm_ducc = mcm02_pure_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)}")