Big potential speedups for coupling matrices

[mreineck]

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)}")

[mreineck]

spin0and2 and apin0and2_pure cases now also work. The former has a speedup of around 6 with nspec=1, and the latter around 3, but for the pure case there is still some remaining optimization potential.

# 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)}")

[thibautlouis]

this is very impressive !

[mreineck]

The spin0and2_pure version is now vectorized and also exhibits a speedup factor of roughly 6.

During testing I noticed that the 5th array computed by calc_mcm_spin0and2_pure always seems to be practically zero, no matter which power spectra I'm providing. My own code behaves in the same way. But if this array is in fact identically zero, why do we spend time computing it?

[thibautlouis]

So Martin, this will be available in a future version of ducc ? we are about to finish the ACT dr6 analysis, so the code is kind of frozen at the moment, but really happy to make the change after that !

[mreineck]

I'm absolutely planning to include this in a future ducc release! The reason for this early "unofficial" demo is that I'd like to get as many things right in the official version as possible, and that I want to include feedback from potential users. A few questions that are still open for me:

• the spin0 and spin_0and2 cases return matrices that are "almost" symmetric (up to a factor 2*l1 + 1), and this factor can be applied easily at a later point. Should I make use of this symmetry and only return a triangular matrix, saving a factor 2 in memory?
• pspy does not compute values for el1<2 or el2<2; should I compute those nevertheless, in case that other codes might want them?
• what is the matter with the 5th component in spin_0and2_pure? Why is it always zero?

I realize now that I chose a fairly unconventional (and intrusive) way of asking for this kind of feedback, by opening an issue in a perfectly well-working code, and I'm sorry about that! If you prefer, I can certainly close the issue here and we can continue discussion on https://github.com/mreineck/ducc.

[mreineck]

The code has now moved to the ducc0 branch (although it is still experimental), and a demo has been added to the ducc0 repository at python/demos/mcm_demo.py. I'll continue to gather feedback before freezing the interface at some point.