Pixelized fsl

pipelines/toast_planck_reduce.py

@@ -151,6 +151,12 @@ def add_sim_params(parser):
         help="Path to beam alm files. Tag DETECTOR will be "
         "replaced with detector name.",
     )
+    # FSL beam mask
+    parser.add_argument(
+        "--fslbeam_mask",
+        required=False,
+        help="Path to pixelized FSL beam mask",
+    )
     # Bandpass for foreground
     parser.add_argument(
         "--freq_sigma",
@@ -1860,6 +1866,7 @@ def run_reproc(
         fg_deriv,
         cmb,
         fslnames,
+        fslbeam_mask_path,
 ):
     """ Reprocess preprocessed or simulated signal.
 
@@ -1924,6 +1931,7 @@ def run_reproc(
         forcepol=args.reproc_forcepol,
         forcefsl=args.reproc_forcefsl,
         fslnames=fslnames,
+        fslbeam_mask_path=fslbeam_mask_path,
         asymmetric_fsl=args.reproc_asymmetric_fsl,
         bpcorrect=args.reproc_bpcorrect,
         pscorrect=args.reproc_pscorrect,
@@ -2332,6 +2340,13 @@ def main():
             fslnames.append(fslname)
     else:
         fslnames = None
+
+    if args.fslbeam_mask:
+        fslbeam_mask_path = {}
+        for det in data.obs[0]["tod"].detectors:
+            fslbeam_mask_path[det] = args.fslbeam_mask.replace("DETECTOR", det)
+    else:
+        fslbeam_mask_path = None
 
     for mc in range(args.MC_start, args.MC_start + args.MC_count):
         mpiworld.Barrier()
@@ -2367,6 +2382,7 @@ def main():
             fg_deriv,
             cmb,
             fslnames,
+            fslbeam_mask_path,
         )
         del cmb
         purge_caches(data, mcmode, mpiworld)

toast_planck/reproc_ring.py

@@ -275,6 +275,7 @@ def __init__(
         forcepol=False,
         forcefsl=False,
         fslnames=None,
+        fslbeam_mask_path=None,
         asymmetric_fsl=False,
         pscorrect=False,
         psradius=30,
@@ -403,6 +404,10 @@ def __init__(
         self.forcepol = forcepol
         self.forcefsl = forcefsl
         self.fslnames = fslnames
+        self.fslbeam_mask_path = fslbeam_mask_path
+        self.fslbeam_mask = None
+        self.nside_fsl = None
+        self.fsl_pixels = None
         self.asymmetric_fsl = asymmetric_fsl
         self.pscorrect = pscorrect
         self.psradius = psradius
@@ -434,6 +439,7 @@ def __init__(
         self.cache_dipole = cache_dipole
 
         self.mapsampler_freq = None
+        self.downgraded_mapsampler_freq = None
         self.mapsamplers = {}
         self.freq_symmetric = ["pol0", "pol1", "pol2", "dipo_foreground"]
         if not self.independent_zodi:
@@ -1564,6 +1570,51 @@ def _add_fsl_template(
                 del fsl, ring_fsl
         return
 
+    @function_timer
+    def _add_pixelized_fsl_templates(
+        self,
+        iring,
+        templates,
+        ring_theta,
+        ring_phi,
+        ring_psi,
+        det,
+        namplitude,
+    ):
+        if self.fslbeam_mask_path is not None:
+            ## Compute detector quaternions
+            nbin = ring_theta.shape[0]
+            psi = ring_psi - np.radians((self.rimo[det].psi_pol + self.rimo[det].psi_uv))
+            det_quat = np.zeros((nbin, 4))
+            # ZYZ conversion from
+            # http://ntrs.nasa.gov/archive/nasa/casi.ntrs.nasa.gov/19770024290.pdf
+            # Note: The above document has the scalar part of the quaternion at
+            # first position but quaternionarray module has it at the end, we
+            # use the quaternionarray convention
+            # scalar part:
+            det_quat[:,3] = np.cos(0.5 * ring_theta) * np.cos(0.5 * (ring_phi + psi))
+            # vector part
+            det_quat[:,0] = -np.sin(0.5 * ring_theta) * np.sin(0.5 * (ring_phi - psi))
+            det_quat[:,1] = np.sin(0.5 * ring_theta) * np.cos(0.5 * (ring_phi - psi))
+            det_quat[:,2] = np.cos(0.5 * ring_theta) * np.sin(0.5 * (ring_phi + psi))
+
+            ## Build FSL templates pixel by pixel
+            for fsl_pixel in self.fsl_pixels[det]:
+                fslname = f'pixelized_fsl-{fsl_pixel}'
+                dtheta_fsl, dphi_fsl = hp.pix2ang(self.nside_fsl[det], fsl_pixel)
+                # Sidelobe pixel rotation
+                sidelobe_quat = qa.mult(qa.rotation(ZAXIS, dphi_fsl), qa.rotation(YAXIS, dtheta_fsl))
+                full_quat = qa.mult(det_quat, sidelobe_quat)
+                vec_sl = qa.rotate(full_quat, ZAXIS).T
+                pix_sl = hp.vec2pix(self.nside_fsl[det], *vec_sl, nest=self.mapsampler_freq.nest)
+                ring_fsl = self.downgraded_mapsampler_freq[det][pix_sl] #sampler.at(theta_fsl, phi_fsl) #sky[pix_sl] - np.mean(sky[pix_sl])
+                ring_fsl -= np.mean(ring_fsl)
+                templates[iring][det][fslname] = RingTemplate(ring_fsl, 0)
+                if fslname not in namplitude:
+                    namplitude[fslname] = 1
+            del det_quat, full_quat, vec_sl, pix_sl, ring_fsl
+        return
+
     @function_timer
     def _add_bandpass_templates(
         self,
@@ -1955,6 +2006,20 @@ def build_templates(self, rings):
             self.mapsampler_freq = mapsampler_fg
             self.mapsampler_freq_has_dipole = False
 
+        # Load FSL beam mask and corresponding low resolution sky model (one for each detector)
+        if self.fslbeam_mask_path is not None:
+            if self.fslbeam_mask == None:
+                self.fslbeam_mask = {}
+                self.nside_fsl = {}
+                self.fsl_pixels = {}
+                if self.downgraded_mapsampler_freq is None:
+                    self.downgraded_mapsampler_freq = {}
+                for det in self.dets:
+                    self.fslbeam_mask[det] = hp.read_map(self.fslbeam_mask_path[det])
+                    self.nside_fsl[det] = hp.get_nside(self.fslbeam_mask[det])
+                    self.fsl_pixels[det] = np.arange(12*self.nside_fsl[det]**2)[self.fslbeam_mask[det]]
+                    self.downgraded_mapsampler_freq[det] = hp.ud_grade(self.mapsampler_freq.Map[:], self.nside_fsl[det], order_in=self.mapsampler_freq.order)
+
         self.local_dipo_amp = np.zeros(self.nring)
         self.local_calib_rms = np.zeros(self.nring)
 
@@ -1989,9 +2054,9 @@ def build_templates(self, rings):
                 ring = rings[iring][det]
                 nbin = ring.pixels.size
 
-                ring_theta, ring_phi = qa.to_position(ring.quat)
+                # ring_theta, ring_phi = qa.to_position(ring.quat)
                 # DEBUG begin
-                # ring_theta, ring_phi, ring_psi = qa.to_angles(ring.quat)
+                ring_theta, ring_phi, ring_psi = qa.to_angles(ring.quat)
                 # fname_pickle = f"ring_pointing_{iring}_{det}.pck"
                 # with open(fname_pickle, "wb") as f:
                 #     pickle.dump(
@@ -2090,6 +2155,16 @@ def build_templates(self, rings):
                     glob2loc,
                 )
 
+                self._add_pixelized_fsl_templates(
+                    iring,
+                    templates,
+                    ring_theta,
+                    ring_phi,
+                    ring_psi,
+                    det,
+                    namplitude,
+                )
+
                 self._add_zodi_templates(
                     ring, iring, position, nbin, det, namplitude, templates, iquweights
                 )
@@ -3334,6 +3409,43 @@ def _subtract_fsl(
                     best_fits[fslname] += amp * orbital_gain
         return
 
+    @function_timer
+    def _subtract_pixelized_fsl(
+        self,
+        det,
+        iring,
+        old_fits,
+        templates,
+        baselines,
+        gain,
+        signal,
+        best_fits,
+        orbital_gain,
+        phase,
+        ring
+    ):
+        if self.fslbeam_mask_path is not None:
+            for fsl_pixel in self.fsl_pixels[det]:
+                fslname = f'pixelized_fsl-{fsl_pixel}'
+                if fslname in old_fits:
+                    old_amp = old_fits[fslname]
+                else:
+                    old_amp = 0
+                offset = templates[iring][det][fslname].offset
+                amp = self.get_amp(det, offset, baselines, fslname)
+                fsltemplate = self._interpolate_ring_template(
+                    ring, templates[iring][det][fslname].template.astype(np.float64), phase
+                )
+                corr = (1 - gain) * old_amp + amp
+                signal -= corr * fsltemplate
+                if fslname not in best_fits:
+                    if fslname in old_fits:
+                        best_fits[fslname] = old_fits[fslname]
+                    else:
+                        best_fits[fslname] = 0.0
+                    best_fits[fslname] += amp * orbital_gain
+        return
+
     @function_timer
     def _interpolate_ring_template(self, ring, template, phase):
         """ Expand a phase-binned ring template
@@ -3813,6 +3925,20 @@ def clean_tod(self, rings, templates, gains, baselines, orbital_gain):
                     rings[iring][det],
                 )
 
+                self._subtract_pixelized_fsl(
+                    det,
+                    iring,
+                    old_fits,
+                    templates,
+                    baselines,
+                    gain,
+                    signal,
+                    best_fits,
+                    orbital_gain,
+                    phase,
+                    rings[iring][det],
+                )
+
                 self._subtract_dipole(
                     templates,
                     iring,
@@ -4553,6 +4679,11 @@ def _update_frequency_map(self, pars):
         self.mapsampler_freq_has_dipole = True
         del full_map
 
+        # Update low resolution sky model to sample FSL signal
+        if self.fslbeam_mask_path is not None:
+            for det in self.dets:
+                self.downgraded_mapsampler_freq[det] = hp.ud_grade(self.mapsampler_freq.Map[:], self.nside_fsl[det], order_in=self.mapsampler_freq.order)
+
         """
         Smoothing the polarization template may compromise single detector maps
         if self.pol_fwhm: