1import ducc0
2import numpy as np
3from time import time
4
5rng = np.random.default_rng(48)
6
7
8def nalm(lmax, mmax):
9    return ((mmax+1)*(mmax+2))//2 + (mmax+1)*(lmax-mmax)
10
11
12def random_alm(lmax, mmax, spin, ncomp):
13    res = rng.uniform(-1., 1., (ncomp, nalm(lmax, mmax))) \
14     + 1j*rng.uniform(-1., 1., (ncomp, nalm(lmax, mmax)))
15    # make a_lm with m==0 real-valued
16    res[:, 0:lmax+1].imag = 0.
17    ofs=0
18    for s in range(spin):
19        res[:, ofs:ofs+spin-s] = 0.
20        ofs += lmax+1-s
21    return res
22
23
24def test(lmax, geometry, spin, nthreads=1):
25    print("testing lmax={}, spin={}, nthreads={}, geometry={}".format(lmax,spin,nthreads,geometry))
26    ncomp = 1 if spin == 0 else 2
27
28    nrings = lmax+1
29    if geometry=="CC":
30        nrings = lmax+2
31    elif geometry=="DH":
32        nrings = 2*lmax+2
33    elif geometry=="F2":
34        nrings = 2*lmax+1
35    nphi=2*lmax+2
36
37    alm = random_alm(lmax, lmax, spin, ncomp)
38    map = np.empty((alm.shape[0], nrings,nphi))
39    alm2 = np.empty((map.shape[0],nalm(lmax,lmax)),dtype=np.complex128)
40    t0=time()
41    map = ducc0.sht.experimental.synthesis_2d(alm=alm, lmax=lmax, spin=spin, map=map, nthreads=nthreads, geometry=geometry)
42    alm2 = ducc0.sht.experimental.analysis_2d(alm=alm2, map=map, lmax=lmax, spin=spin, geometry=geometry, nthreads=nthreads)
43    print(time()-t0)
44    print("L2 error after full round-trip:", ducc0.misc.l2error(alm2,alm))
45    print("L_inf error after full round-trip:", np.max(np.abs(alm2-alm)))
46
47nthr=8
48for l0 in [4095]:
49    for geometry in ["CC", "F1", "MW", "MWflip", "GL", "DH", "F2"]:
50        for spin in [0,1,2]:
51            test(lmax=l0, spin=spin, nthreads=nthr, geometry=geometry)
52