Muellerconvolver comparison

beamconv/instrument.py

@@ -1811,11 +1811,15 @@ def scan_instrument_mpi(self, alm, ground_alm=None, verbose=1, binning=True,
         # nside_spin from init_spinmaps kwargs (if not None).
         max_spin = kwargs.pop('max_spin', None)
         nside_spin = kwargs.pop('nside_spin', None)
+        mu_con_hwp = kwargs.pop('mu_con_hwp', True)
+        mu_con_spin = kwargs.pop('mu_con_spin', True)
         spinmaps_opts = {}
         if max_spin:
             spinmaps_opts.update({'max_spin' : max_spin})
         if nside_spin:
             spinmaps_opts.update({'nside_spin' : nside_spin})
+        spinmaps_opts.update({'mu_con_hwp' : mu_con_hwp})
+        spinmaps_opts.update({'mu_con_spin' : mu_con_spin})
 
         if verbose and self.mpi_rank == 0:
             print('Scanning with {:d} detectors'.format(
@@ -3799,7 +3803,7 @@ def _scan_modulate_pa(tod, pix, pa, maps, s_vals,
 
 
     def init_spinmaps(self, alm, beam, ground_alm=None, max_spin=5, nside_spin=256,
-                      symmetric=False, input_v=False, beam_v=False):
+                      symmetric=False, input_v=False, beam_v=False, mu_con_hwp=True, mu_con_spin=True):
         '''
         Compute appropriate spinmaps for beam and
         all its ghosts.
@@ -3827,6 +3831,12 @@ def init_spinmaps(self, alm, beam, ground_alm=None, max_spin=5, nside_spin=256,
             include the 4th alm component
         beam_v : bool
             include the 4th blm component
+        mu_con_hwp : bool
+            Switch between the Mueller convolver matching 
+            HWP tweak and the old (pre-July 2023) behavior
+        mu_con_spin : bool
+            Switch between the Mueller convolver matching 
+            spin +-2 tweak and the old (pre-July 2023) behavior
 
 
         Notes
@@ -3852,7 +3862,7 @@ def init_spinmaps(self, alm, beam, ground_alm=None, max_spin=5, nside_spin=256,
         spinmap_dict = self._init_spinmaps(alm,
                     blm, max_s, nside_spin, symmetric=beam.symmetric,
                     hwp_mueller=beam.hwp_mueller, input_v=input_v,
-                    beam_v=beam_v)
+                    beam_v=beam_v, mu_con_hwp=mu_con_hwp, mu_con_spin=mu_con_spin)
 
         # Names: s0a0, s0a2, s2a0, s2a2, s2a4.
         # s refers to spin value under psi, a to spin value under HWP rot.
@@ -3886,20 +3896,20 @@ def init_spinmaps(self, alm, beam, ground_alm=None, max_spin=5, nside_spin=256,
                 spinmap_dict = self._init_spinmaps(alm,
                             blm, max_s, nside_spin, symmetric=ghost.symmetric,
                             hwp_mueller=ghost.hwp_mueller,
-                            input_v=input_v, beam_v=beam_v)
+                            input_v=input_v, beam_v=beam_v, mu_con_hwp=mu_con_hwp, mu_con_spin=mu_con_spin)
                 self.spinmaps['ghosts'][u] = spinmap_dict
 
         if ground_alm is not None:
             self.spinmaps['ground'] = self._init_spinmaps(ground_alm,
                                     blm, max_s, nside_spin, symmetric=beam.symmetric,
                                     hwp_mueller=beam.hwp_mueller, input_v=input_v,
-                                    beam_v=beam_v)
+                                    beam_v=beam_v, mu_con_hwp=mu_con_hwp, mu_con_spin=mu_con_spin)
 
 
     @staticmethod
     def _init_spinmaps(alm, blm, max_spin, nside,
                        symmetric=False, hwp_mueller=None,
-                       input_v=False, beam_v=False):
+                       input_v=False, beam_v=False, mu_con_hwp=True, mu_con_spin=True):
         '''
         Compute convolution of map with different spin modes
         of the beam.
@@ -3928,6 +3938,13 @@ def _init_spinmaps(alm, blm, max_spin, nside,
         beam_v : bool
             include the 4th blm component if
             it exists
+        mu_con_hwp : bool
+            Switch between the Mueller convolver matching 
+            HWP tweak and the old (pre-July 2023) behavior
+        mu_con_spin : bool
+            Switch between the Mueller convolver matching 
+            spin +-2 tweak and the old (pre-July 2023) behavior
+
 
         returns
         -------
@@ -4036,15 +4053,21 @@ def _init_spinmaps(alm, blm, max_spin, nside,
 
             # s0a2.
             spinmap_dict['s0a2'] = {}
-            blmm2, blmp2 = ScanStrategy.blmxhwp(blm, hwp_spin, 's0a2')
-            blmE, blmB = tools.spin2eb(blmp2, blmm2)
+            blmm2, blmp2 = ScanStrategy.blmxhwp(blm, hwp_spin, 's0a2', mu_con_hwp=mu_con_hwp)
+            if mu_con_spin==False:
+                blmE, blmB = tools.spin2eb(blmp2, blmm2)
+            else:
+                blmE, blmB = tools.spin2eb(blmm2, blmp2)
             # 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)
 
             if input_v:
-                blmm2, blmp2 = ScanStrategy.blmxhwp(blm, hwp_spin, 's0a2_v')
-                blmE, blmB = tools.spin2eb(blmp2, blmm2)
+                blmm2, blmp2 = ScanStrategy.blmxhwp(blm, hwp_spin, 's0a2_v', mu_con_hwp=mu_con_hwp)
+                if mu_con_spin==False:
+                    blmE, blmB = tools.spin2eb(blmp2, blmm2)
+                else:
+                    blmE, blmB = tools.spin2eb(blmm2, blmp2)
                 spinmap_dict['s0a2']['maps'] += ScanStrategy._spinmaps_complex(
                     -alm[3], alm[3] * 0, blmE, blmB, spin_values_s0a2, nside)
             spinmap_dict['s0a2']['s_vals'] = spin_values_s0a2
@@ -4090,7 +4113,7 @@ def _init_spinmaps(alm, blm, max_spin, nside,
         return spinmap_dict
 
     @staticmethod
-    def blmxhwp(blm, hwp_spin, mode, beam_v=False):
+    def blmxhwp(blm, hwp_spin, mode, beam_v=False, mu_con_hwp=True):
         '''
         Arguments
         ---------
@@ -4103,6 +4126,9 @@ def blmxhwp(blm, hwp_spin, mode, beam_v=False):
             Pick between 's0a0', 's0a0_v', s2a4', 's0a2', 's0a2_v', 's2a2' or 's2a0'.
         beam_v : bool
             include the 4th blm component when constructing the spinmaps
+        mu_con_hwp : bool
+            Switch between the Mueller convolver matching 
+            HWP tweak and the old (pre-July 2023) behavior
 
         Returns
         -------
@@ -4145,13 +4171,19 @@ def blmxhwp(blm, hwp_spin, mode, beam_v=False):
 
         elif mode == 's0a2':
             blmm2, blmp2 = tools.shift_blm(blm[1], blm[2], 2, eb=False)            
-            blmm2 *= hwp_spin[0,2] * np.sqrt(2)
+            if mu_con_hwp == False:
+                blmm2 *= hwp_spin[0,2] * np.sqrt(2)
+            else:
+                blmm2 *= hwp_spin[1,0] * np.sqrt(2)
             blmp2 *= hwp_spin[2,0] * np.sqrt(2)
             return blmm2, blmp2
 
         elif mode == 's0a2_v':
             blmm2_v, blmp2_v = tools.shift_blm(blm[1], blm[2], 2, eb=False)
-            blmm2_v *= hwp_spin[3,2] * np.sqrt(2)
+            if mu_con_hwp == False:
+                blmm2_v *= hwp_spin[3,2] * np.sqrt(2)
+            else:
+                blmm2_v *= hwp_spin[1,3] * np.sqrt(2)
             blmp2_v *= hwp_spin[2,3] * np.sqrt(2)
             return blmm2_v, blmp2_v
 
@@ -4235,7 +4267,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.
 
@@ -4318,8 +4350,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:

comparison_script.py

@@ -0,0 +1,180 @@
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import healpy as hp
+import mueller_convolver
+import ducc0
+
+def nalm(lmax, mmax):
+    return ((mmax+1)*(mmax+2))//2 + (mmax+1)*(lmax-mmax)
+
+def make_full_random_alm(lmax, mmax, rng):
+    res = rng.uniform(-1., 1., (4, nalm(lmax, mmax))) \
+     + 1j*rng.uniform(-1., 1., (4, nalm(lmax, mmax)))
+    # make a_lm with m==0 real-valued
+    res[:, 0:lmax+1].imag = 0.
+    ofs=0
+    # components 1 and 2 are spin-2, fix them accordingly
+    spin=2
+    for s in range(spin):
+        res[1:3, ofs:ofs+spin-s] = 0.
+        ofs += lmax+1-s
+    return res
+
+def get_gauss_beam_from_beamconv(fwhm, lmax):
+    import beamconv
+    blmT, blmm2 = beamconv.tools.gauss_blm(fwhm*180*60/np.pi, lmax, pol=True)
+    res = np.zeros((4,blmT.shape[0]), dtype=np.complex128)
+    blmE, blmB = beamconv.tools.spin2eb(blmm2, blmm2*0, spin=2)
+    res[0] = blmT
+    res[1] = blmE
+    res[2] = blmB
+    res[3] = blmT  # correct?
+    return res
+
+def blm_gauss_new(fwhm, lmax, pol=False):
+    fwhm = float(fwhm)
+    lmax = int(lmax)
+    pol = bool(pol)
+    mmax = 2 if pol else 0
+    ncomp = 3 if pol else 1
+    nval = hp.Alm.getsize(lmax, mmax)
+
+    if mmax > lmax:
+        raise ValueError("lmax value too small")
+
+    blm = np.zeros((ncomp, nval), dtype=np.complex128)
+    sigmasq = fwhm * fwhm / (8 * np.log(2.0))
+
+    for l in range(lmax+1):
+        blm[0, hp.Alm.getidx(lmax, l, 0)] = np.exp(-0.5*sigmasq*l*(l+1))
+
+    if pol:
+        for l in range(2, lmax+1):
+            blm[1, hp.Alm.getidx(lmax, l, 2)] = np.exp(-0.5 * sigmasq * (l*(l+1)-4))
+        blm[2] = 1j * blm[1]
+
+    return blm
+
+# blm_gauss_new times sqrt((2*l+1)/(4pi))
+def Blm_gauss_new(fwhm, lmax, pol=False):
+    blm = blm_gauss_new(fwhm, lmax, pol)
+    for l in range(lmax+1):
+        blm[0, hp.Alm.getidx(lmax, l, 0)] *= np.sqrt((2*l+1) / (4*np.pi))
+
+    if pol:
+        for l in range(2, lmax+1):
+            blm[1:3, hp.Alm.getidx(lmax, l, 2)] *= np.sqrt((2*l+1) / (4*np.pi))
+
+    return blm
+
+# code by Marta to get beamconv results for user-specified angles
+def get_beamconv_values(lmax, kmax, slm, blm, ptg, hwp_angles, mueller,
+                        mu_con_hwp, mu_con_spin):
+    import beamconv
+    import qpoint as qp
+
+    # prepare PO beam file
+    blm2 = np.zeros((blm.shape[0], hp.Alm.getsize(lmax=lmax)), dtype=np.complex128)
+    blm2[:,:blm.shape[1]] = blm
+    blmm, blmp = beamconv.tools.eb2spin(blm2[1],blm2[2])
+    blm2[1] = blmm
+    blm2[2] = blmp
+    np.save("temp_beam.npy", blm2)
+
+    # set up beam and HWP mueller matrix (identity, i.e. no HWP)
+    beam = beamconv.Beam(btype='PO', lmax=lmax, mmax=lmax, deconv_q=True, normalize=False, po_file="temp_beam.npy", hwp_mueller=mueller)
+
+    nsamp = ptg.shape[0]
+
+    # from (theta,phi) to (ra,dec) convention
+    # also, all angles are converted to degrees
+    ra = np.degrees(ptg[:,1])
+    dec = 90. - np.degrees(ptg[:,0])
+    # Adjustment for difference in convention between qpoint and MuellerConvolver?
+    psi = 180. - np.degrees(ptg[:,2])
+
+    # calculate the quaternion
+    q_bore_array = qp.QPoint().radecpa2quat(ra, dec, psi)
+
+    def ctime_test(**kwargs):
+        return np.zeros(kwargs.pop('end')-kwargs.pop('start'))
+
+    def q_bore_test(**kwargs):
+        return q_bore_array[kwargs.pop('start'):kwargs.pop('end')]
+
+    S = beamconv.ScanStrategy(duration=nsamp, sample_rate=1, external_pointing=True, use_l2_scan=False)
+    S.add_to_focal_plane(beam, combine=False)
+    S.set_hwp_mod(mode='stepped', freq=1, angles=hwp_angles*180/np.pi)
+
+    # determine nside_spin necessary for good accuracy
+    nside_spin = 1
+    while nside_spin < 4*lmax:
+        nside_spin *= 2
+
+    S.scan_instrument_mpi(slm.copy(), save_tod=True, ctime_func=ctime_test, q_bore_func=q_bore_test,
+                      ctime_kwargs={'useless':0}, q_bore_kwargs={'useless':0},nside_spin=nside_spin, interp=True, input_v=True, beam_v=True, max_spin=kmax+4, binning=False, verbose=0, mu_con_hwp=mu_con_hwp, mu_con_spin=mu_con_spin)
+
+    return S.data(S.chunks[0], beam=beam, data_type='tod').copy()
+
+
+np.random.seed(10)
+rng = np.random.default_rng(np.random.SeedSequence(42))
+lmax = 30
+kmax = 18
+
+# completely random sky
+slm =make_full_random_alm(lmax, lmax, rng)
+
+# completely random Mueller matrix
+mueller = np.random.uniform(-1,1,size=(4,4))
+#mueller[1:3,0]=mueller[1:3,-1] = 0
+#mueller[0,2]=mueller[2,0] = 0
+
+
+# completely random beam
+blm = make_full_random_alm(lmax, kmax, rng)
+# ... or use a Gauss beam
+#blmtmp = blm_gauss_new(np.radians(10.), lmax, True)
+#blm *= 0
+#blm [0:3, 0:blmtmp.shape[1]] = blmtmp
+
+nptg=100
+# completely random pointings
+ptg = np.empty((nptg,3))
+ptg[:,0]=np.random.uniform(0,np.pi,size=(nptg,))    # theta
+ptg[:,1]=np.random.uniform(0,2*np.pi,size=(nptg,))  # phi
+ptg[:,2]=np.random.uniform(0,2*np.pi,size=(nptg,))  # psi
+hwp_angles = np.random.uniform(0,2*np.pi,size=(nptg,))  # alpha
+
+# get the signal from beamconv
+signal_beamconv_FF = get_beamconv_values(lmax=lmax, kmax=kmax, slm=slm, blm=blm, ptg=ptg, hwp_angles=hwp_angles, mueller=mueller, mu_con_hwp=False, mu_con_spin=False)
+signal_beamconv_FT = get_beamconv_values(lmax=lmax, kmax=kmax, slm=slm, blm=blm, ptg=ptg, hwp_angles=hwp_angles, mueller=mueller, mu_con_hwp=False, mu_con_spin=True)
+signal_beamconv_TF = get_beamconv_values(lmax=lmax, kmax=kmax, slm=slm, blm=blm, ptg=ptg, hwp_angles=hwp_angles, mueller=mueller, mu_con_hwp=True, mu_con_spin=False)
+signal_beamconv_TT = get_beamconv_values(lmax=lmax, kmax=kmax, slm=slm, blm=blm, ptg=ptg, hwp_angles=hwp_angles, mueller=mueller, mu_con_hwp=True, mu_con_spin=True)
+
+# Now do the same thing with MuellerConvolver
+fullconv = mueller_convolver.MuellerConvolver(
+    lmax=lmax,
+    kmax=kmax,
+    slm=slm,
+    blm=blm,
+    mueller=mueller,
+    single_precision=False,
+    epsilon=1e-7,
+    nthreads=1,
+)
+signal_muellerconvolver = fullconv.signal(ptg=ptg, alpha=hwp_angles)
+
+# L2 error
+print("L2 error to original beamconv:", ducc0.misc.l2error(signal_beamconv_FF, signal_muellerconvolver))
+print("L2 error to beamconv + HWP:", ducc0.misc.l2error(signal_beamconv_TF, signal_muellerconvolver))
+print("L2 error to beamconv + spin:", ducc0.misc.l2error(signal_beamconv_FT, signal_muellerconvolver))
+print("L2 error to beamconv + HWP + spin:", ducc0.misc.l2error(signal_beamconv_TT, signal_muellerconvolver))
+plt.plot(signal_beamconv_FF, label="beamconv, original")
+plt.plot(signal_beamconv_TF, label="beamconv + HWP")
+plt.plot(signal_beamconv_FT, label="beamconv + spin")
+plt.plot(signal_beamconv_TT, label="beamconv + HWP + spin")
+plt.plot(signal_muellerconvolver, label="MuellerConvolver")
+plt.legend()
+plt.show()