Draft: use ducc0.totalconvolve for convolution

beamconv/instrument.py

@@ -11,6 +11,7 @@
 import healpy as hp
 
 from . import tools
+from . import totalconvolve
 from .detector import Beam
 
 class MPIBase(object):
@@ -3244,9 +3245,19 @@ def scan(self, beam, add_to_tod=False, interp=False,
         else:
             hwp_ang = self.hwp_ang
 
+        # obtain theta, phi, psi pointings
+        if interp:
+            theta, phi = pix[1], pix[0]
+        else:
+            nside = hp.npix2nside(len(spinmaps['s0a0']['maps'][0]))
+            theta, phi = hp.pix2ang(nside,pix,nest=False)
+        psi = pa
+
         # Find out if old or new HWP behaviour is desired.
         if set(spinmaps.keys()) == set(['s0a0', 's2a4']):
             # Old behaviour.
+# For some reason, this causes a single test to fail...
+#            tod_c += spinmaps['s2a4']['interpolators'].interpol(theta, phi, psi)
             self._scan_modulate_pa(tod_c, pix, pa,
                                    spinmaps['s2a4']['maps'],
                                    spinmaps['s2a4']['s_vals'],
@@ -3259,39 +3270,29 @@ def scan(self, beam, add_to_tod=False, interp=False,
             tod = np.real(tod_c) # Note, shares memory with tod_c.
 
             # Add unpolarized tod.
-            self._scan_modulate_pa(tod, pix, pa,
-                                   spinmaps['s0a0']['maps'],
-                                   spinmaps['s0a0']['s_vals'],
-                                   reality=True, interp=interp)
+            tod += spinmaps['s0a0']['interpolators'].interpol(theta, phi, psi)
+            # self._scan_modulate_pa(tod, pix, pa,
+                                   # spinmaps['s0a0']['maps'],
+                                   # spinmaps['s0a0']['s_vals'],
+                                   # reality=True, interp=interp)
 
         else:
             # New behaviour.
-            self._scan_modulate_pa(tod_c, pix, pa,
-                                   spinmaps['s2a4']['maps'],
-                                   spinmaps['s2a4']['s_vals'],
-                                   reality=False, interp=interp)
+            tod_c = spinmaps['s2a4']['interpolators'].interpol(theta, phi, psi)
 
             # Modulate by HWP angle and polarization angle.
             expm2 = np.exp(1j * (4 * hwp_ang + 2 * np.radians(polang)))
             tod_c *= expm2
 
             tod = np.real(tod_c).copy()
 
-            tod_c *= 0.
-            self._scan_modulate_pa(tod_c, pix, pa,
-                                   spinmaps['s2a2']['maps'],
-                                   spinmaps['s2a2']['s_vals'],
-                                   reality=False, interp=interp)
+            tod_c = spinmaps['s2a2']['interpolators'].interpol(theta, phi, psi)
 
             expm2 = np.exp(1j * (2 * hwp_ang))
             tod_c *= expm2
             tod += np.real(tod_c)
 
-            tod_c *= 0
-            self._scan_modulate_pa(tod_c, pix, pa,
-                                   spinmaps['s2a0']['maps'],
-                                   spinmaps['s2a0']['s_vals'],
-                                   reality=False, interp=interp)
+            tod_c = spinmaps['s2a0']['interpolators'].interpol(theta, phi, psi)
 
             tod += np.real(tod_c)
             del tod_c
@@ -3346,11 +3347,7 @@ def scan(self, beam, add_to_tod=False, interp=False,
             tod += tod_tmp
             del tod_tmp
 
-            # Normal unpolarized part.
-            self._scan_modulate_pa(tod, pix, pa,
-                                   spinmaps['s0a0']['maps'],
-                                   spinmaps['s0a0']['s_vals'],
-                                   reality=True, interp=interp)
+            tod += spinmaps['s0a0']['interpolators'].interpol(theta, phi, psi)
 
         if add_to_tod and hasattr(self, 'tod'):
             self.tod += tod
@@ -3666,9 +3663,14 @@ def _init_spinmaps(alm, blm, max_spin, nside,
             if input_v:
                 blm_s0a0_v = ScanStrategy.blmxhwp(blm, hwp_spin, 's0a0_v',
                                                     beam_v=beam_v)
+                spinmap_dict['s0a0']['interpolators'] = totalconvolve.Interpolator_real(
+                    [alm[0], alm[3]], [blm_s0a0, blm_s0a0_v], lmax, spin_values_unpol)
                 spinmap_dict['s0a0']['maps'] += ScanStrategy._spinmaps_real(
                     alm[3], blm_s0a0_v, spin_values_unpol, nside)
                 del blm_s0a0_v
+            else:
+                spinmap_dict['s0a0']['interpolators'] = totalconvolve.Interpolator_real(
+                    [alm[0]], [blm_s0a0], lmax, spin_values_unpol)
 
             spinmap_dict['s0a0']['s_vals'] = spin_values_unpol
             del blm_s0a0
@@ -3679,6 +3681,8 @@ def _init_spinmaps(alm, blm, max_spin, nside,
             blmE, blmB = tools.spin2eb(blmm2, blmp2)
             spinmap_dict['s2a4']['maps'] = ScanStrategy._spinmaps_complex(
                 almE, almB, blmE, blmB, spin_values_pol, nside)
+            spinmap_dict['s2a4']['interpolators'] = totalconvolve.Interpolator_complex(
+                almE, almB, blmE, blmB, lmax, spin_values_pol)
             spinmap_dict['s2a4']['s_vals'] = spin_values_pol
 
             # s0a2.
@@ -3688,12 +3692,16 @@ def _init_spinmaps(alm, blm, max_spin, nside,
             # Minus sign is because eb2spin applied to alm_I results in (-1) alm_I.
             spinmap_dict['s0a2']['maps'] = ScanStrategy._spinmaps_complex(
                 -alm[0], alm[0] * 0, blmE, blmB, spin_values_s0a2, nside)
+            spinmap_dict['s0a2']['interpolators0'] = totalconvolve.Interpolator_complex(
+                -alm[0], alm[0]*0, blmE, blmB, lmax, spin_values_s0a2)
 
             if input_v:
                 blmm2, blmp2 = ScanStrategy.blmxhwp(blm, hwp_spin, 's0a2_v')
                 blmE, blmB = tools.spin2eb(blmp2, blmm2)
                 spinmap_dict['s0a2']['maps'] += ScanStrategy._spinmaps_complex(
                     -alm[3], alm[3] * 0, blmE, blmB, spin_values_s0a2, nside)
+                spinmap_dict['s0a2']['interpolators3'] = totalconvolve.Interpolator_complex(
+                    -alm[3], alm[3]*0, blmE, blmB, lmax, spin_values_s0a2)
             spinmap_dict['s0a2']['s_vals'] = spin_values_s0a2
 
             # s2a2.
@@ -3703,6 +3711,8 @@ def _init_spinmaps(alm, blm, max_spin, nside,
             blmE, blmB = tools.spin2eb(blmm2, blmp2)
             spinmap_dict['s2a2']['maps'] = ScanStrategy._spinmaps_complex(
                 almE, almB, blmE, blmB, spin_values_pol, nside)
+            spinmap_dict['s2a2']['interpolators'] = totalconvolve.Interpolator_complex(
+                almE, almB, blmE, blmB, lmax, spin_values_pol)
             spinmap_dict['s2a2']['s_vals'] = spin_values_pol
 
             # s2a0.
@@ -3711,6 +3721,8 @@ def _init_spinmaps(alm, blm, max_spin, nside,
             blmE, blmB = tools.spin2eb(blmm2, blmp2)
             spinmap_dict['s2a0']['maps'] = ScanStrategy._spinmaps_complex(
                 almE, almB, blmE, blmB, spin_values_pol, nside)
+            spinmap_dict['s2a0']['interpolators'] = totalconvolve.Interpolator_complex(
+                almE, almB, blmE, blmB, lmax, spin_values_pol)
             spinmap_dict['s2a0']['s_vals'] = spin_values_pol
 
         else:
@@ -3723,8 +3735,13 @@ def _init_spinmaps(alm, blm, max_spin, nside,
                 alm[0], blm[0], spin_values_unpol, nside)
 
             if beam_v and input_v:
+                spinmap_dict['s0a0']['interpolators'] = totalconvolve.Interpolator_real(
+                    [alm[0], alm[3]], [blm[0], blm[3]], lmax, spin_values_unpol)
                 spinmap_dict['s0a0']['maps'] += ScanStrategy._spinmaps_real(
                     alm[3], blm[3], spin_values_unpol, nside)
+            else:
+                spinmap_dict['s0a0']['interpolators'] = totalconvolve.Interpolator_real(
+                    [alm[0]], [blm[0]], lmax, spin_values_unpol)
 
             spinmap_dict['s0a0']['s_vals'] = spin_values_unpol
 
@@ -3733,6 +3750,8 @@ def _init_spinmaps(alm, blm, max_spin, nside,
             spinmap_dict['s2a4']['maps'] = ScanStrategy._spinmaps_complex(
                 almE, almB, blmE, blmB, spin_values_pol, nside)
             spinmap_dict['s2a4']['s_vals'] = spin_values_pol
+            spinmap_dict['s2a4']['interpolators'] = totalconvolve.Interpolator_complex(
+                almE, almB, blmE, blmB, lmax, spin_values_pol)
 
         return spinmap_dict
 
@@ -3883,7 +3902,7 @@ def _spinmaps_real(alm, blm, spin_values, nside):
             if s == 0: # Scalar transform.
 
                 flms = hp.almxfl(alm, bell, inplace=False)
-                func[sidx,:] = hp.alm2map(flms, nside, verbose=False)
+                func[sidx,:] = hp.alm2map(flms, nside)
 
             else: # Spin transforms.
 
@@ -3966,8 +3985,8 @@ def _spinmaps_complex(almE, almB, blmE, blmB, spin_values, nside):
 
             if s == 0:
                 # The (-1) factor for spin 0 is explained in HEALPix doc.
-                spinmaps = [hp.alm2map(-ps_flm_p, nside, verbose=False),
-                            hp.alm2map(ms_flm_m, nside, verbose=False)]
+                spinmaps = [hp.alm2map(-ps_flm_p, nside),
+                            hp.alm2map(ms_flm_m, nside)]
                 func_c[sidx,:] = spinmaps[0] + 1j * spinmaps[1]
 
             if s > 0:

beamconv/test.py

@@ -122,8 +122,8 @@ def scan_bicep(lmax=700, mmax=5, fwhm=43, ra0=-10, dec0=-57.5,
 
         # plot smoothed input maps
         nside = hp.get_nside(maps[0])
-        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.), verbose=False)
-        maps_raw = hp.alm2map(alm, nside, verbose=False)
+        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.))
+        maps_raw = hp.alm2map(alm, nside)
 
         plot_iqu(maps_raw, img_out_path, 'raw_bicep',
                  sym_limits=[250, 5, 5],
@@ -265,8 +265,8 @@ def scan_atacama(lmax=700, mmax=5, fwhm=40,
 
         # plot smoothed input maps
         nside = hp.get_nside(maps[0])
-        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.), verbose=False)
-        maps_raw = hp.alm2map(alm, nside, verbose=False)
+        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.))
+        maps_raw = hp.alm2map(alm, nside)
 
         plot_iqu(maps_raw, img_out_path, 'raw_atacama',
                  sym_limits=[250, 5, 5],
@@ -925,8 +925,8 @@ def test_ghosts(lmax=700, mmax=5, fwhm=43, ra0=-10, dec0=-57.5,
 
         # plot smoothed input maps
         nside = hp.get_nside(maps[0])
-        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.), verbose=False)
-        maps_raw = hp.alm2map(alm, nside, verbose=False)
+        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.))
+        maps_raw = hp.alm2map(alm, nside)
 
         plot_iqu(maps_raw, '../scratch/img/', 'raw_ghost',
                  sym_limits=[250, 5, 5],
@@ -1162,8 +1162,8 @@ def idea_jon():
 
     # plot smoothed input maps
     nside = hp.get_nside(maps_g[0])
-    hp.smoothalm(alm, fwhm=np.radians(fwhm/60.), verbose=False)
-    maps_raw = hp.alm2map(alm, nside, verbose=False)
+    hp.smoothalm(alm, fwhm=np.radians(fwhm/60.))
+    maps_raw = hp.alm2map(alm, nside)
 
     plot_iqu(maps_raw, '../scratch/img/', 'raw_delta',
              sym_limits=[1, 1, 1],
@@ -1334,8 +1334,8 @@ def test_satellite_scan(lmax=700, mmax=2, fwhm=43,
 
         # plot smoothed input maps
         nside = hp.get_nside(maps[0])
-        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.), verbose=False)
-        maps_raw = hp.alm2map(alm, nside, verbose=False)
+        hp.smoothalm(alm, fwhm=np.radians(fwhm/60.))
+        maps_raw = hp.alm2map(alm, nside)
 
         plot_iqu(maps_raw, '../scratch/img/', 'raw_satellite',
                  sym_limits=[250, 5, 5],

beamconv/totalconvolve.py

@@ -0,0 +1,62 @@
+import ducc0.totalconvolve
+import numpy as np
+
+
+def _convert_blm(blm_in, lmax, spin_values):
+    kmax = max(spin_values)
+    nblm = ((kmax+1)*(kmax+2))//2 + (kmax+1)*(lmax-kmax)
+    blm = np.array([x[:nblm].copy() for x in blm_in])
+    lfac = np.sqrt((1.+2*np.arange(lmax+1.))/(4*np.pi))
+    ofs=0
+    for m in range(kmax+1):
+       blm[:, ofs:ofs+lmax+1-m] *= lfac[m:].reshape((1,-1))
+       if (not m in spin_values):
+           blm[:, ofs:ofs+lmax+1-m] = 0
+       ofs += lmax+1-m
+    return blm
+def _convert_blm2(blm_in, lmax, kmax):
+    nblm = ((kmax+1)*(kmax+2))//2 + (kmax+1)*(lmax-kmax)
+    blm = blm_in.copy()
+    lfac = np.sqrt((1.+2*np.arange(lmax+1.))/(4*np.pi))
+    ofs=0
+    for m in range(kmax+1):
+       blm[:, ofs:ofs+lmax+1-m] = (-1)**m * np.conj(blm[:, ofs:ofs+lmax+1-m])
+       ofs += lmax+1-m
+    blm= blm[-1::-1]
+    return blm
+
+
+class Interpolator_real:
+    def __init__(self, slm, blm, lmax, spin_values, epsilon=1e-11, nthreads=1):
+        _slm = np.array([x.copy() for x in slm])
+        _blm = _convert_blm(blm, lmax, spin_values)
+        self._inter = ducc0.totalconvolve.Interpolator(
+            _slm, _blm, False, lmax, max(spin_values), epsilon, 2., nthreads)
+
+    def interpol(self, theta, phi, psi):
+        ptg = np.zeros((theta.shape[0], 3))
+        ptg[:, 0] = theta
+        ptg[:, 1] = phi
+        ptg[:, 2] = psi
+        return self._inter.interpol(ptg)[0]
+
+
+class Interpolator_complex:
+    def __init__(self, slmE, slmB, blmE, blmB, lmax, spin_values, epsilon=1e-11, nthreads=1):
+        _slm = np.array([slmE, slmB])
+        _blm = _convert_blm([blmE, blmB], lmax, spin_values)
+        kmax = max(spin_values)
+        self._inter = ducc0.totalconvolve.Interpolator(
+            _slm, _blm, False, lmax, kmax, epsilon, 2., nthreads)
+        _blm2 = _convert_blm2(_blm, lmax, kmax)
+        self._inter2 = ducc0.totalconvolve.Interpolator(
+            _slm, _blm2, False, lmax, kmax, epsilon, 2., nthreads)
+
+    def interpol(self, theta, phi, psi):
+        ptg = np.zeros((theta.shape[0], 3))
+        ptg[:, 0] = theta
+        ptg[:, 1] = phi
+        ptg[:, 2] = psi
+        res = self._inter.interpol(ptg)
+        res2 = self._inter2.interpol(ptg)
+        return res[0]+1j*res2[0]

tests/test_scan_strategy.py

@@ -156,7 +156,7 @@ def test_init_detpair(self):
         # Since we have a infinitely narrow Gaussian the convolved
         # maps should just match the input (up to healpix quadrature
         # wonkyness).
-        input_map = hp.alm2map(self.alm, nside, verbose=False) # I, Q, U
+        input_map = hp.alm2map(self.alm, nside) # I, Q, U
         zero_map = np.zeros_like(input_map[0])
         np.testing.assert_array_almost_equal(input_map[0],
                                              func[0], decimal=6)
@@ -264,7 +264,7 @@ def test_scan_spole(self):
 
         # Construct TOD manually.
         polang = beam.polang
-        maps_sm = np.asarray(hp.alm2map(self.alm, nside, verbose=False,
+        maps_sm = np.asarray(hp.alm2map(self.alm, nside,
                                         fwhm=np.radians(beam.fwhm / 60.)))
 
         np.testing.assert_almost_equal(maps_sm[0],
@@ -346,7 +346,7 @@ def test_scan_spole_pol(self):
 
         # Construct TOD manually.
         polang = beam.polang
-        maps_sm = np.asarray(hp.alm2map(alm, nside, verbose=False,
+        maps_sm = np.asarray(hp.alm2map(alm, nside,
                                         fwhm=np.radians(beam.fwhm / 60.)))
 
         np.testing.assert_almost_equal(maps_sm[0],
@@ -407,9 +407,8 @@ def test_scan_spole_bin(self):
         # Solve for the maps.
         maps, cond = scs.solve_for_map(fill=np.nan)
 
-        alm = hp.smoothalm(self.alm, fwhm=np.radians(fwhm/60.),
-                     verbose=False)
-        maps_raw = np.asarray(hp.alm2map(self.alm, nside, verbose=False))
+        alm = hp.smoothalm(self.alm, fwhm=np.radians(fwhm/60.))
+        maps_raw = np.asarray(hp.alm2map(self.alm, nside))
 
         cond[~np.isfinite(cond)] = 10
 
@@ -1256,7 +1255,7 @@ def test_scan_spole_hwp_mueller(self):
 
         # Construct TOD manually.
         polang = beam.polang
-        maps_sm = np.asarray(hp.alm2map(self.alm, nside, verbose=False,
+        maps_sm = np.asarray(hp.alm2map(self.alm, nside,
                                         fwhm=np.radians(beam.fwhm / 60.)))
 
         np.testing.assert_almost_equal(maps_sm[0],
@@ -1326,9 +1325,8 @@ def test_scan_spole_bin_hwp_mueller(self):
         # Solve for the maps.
         maps, cond = scs.solve_for_map(fill=np.nan)
 
-        alm = hp.smoothalm(self.alm, fwhm=np.radians(fwhm/60.),
-                     verbose=False)
-        maps_raw = np.asarray(hp.alm2map(self.alm, nside, verbose=False))
+        alm = hp.smoothalm(self.alm, fwhm=np.radians(fwhm/60.))
+        maps_raw = np.asarray(hp.alm2map(self.alm, nside))
 
         cond[~np.isfinite(cond)] = 10