Merge pull request #194 from galsci/integrate_mapsims

Make sure PySM integrates properly with mapsims

src/pysm3/data/presets.cfg

@@ -147,6 +147,7 @@ unit_mbb_temperature = "K"
 freq_ref_I = "353 GHz"
 freq_ref_P = "353 GHz"
 max_nside = 8192
+available_nside = [2048, 4096, 8192]
 [d10]
 class = "ModifiedBlackBody"
 map_I = "dust_gnilc/gnilc_dust_template_nside{nside}_2023.02.10.fits"
@@ -155,6 +156,7 @@ map_mbb_temperature = "dust_gnilc/gnilc_dust_Td_nside{nside}_2023.06.06.fits"
 freq_ref_I = "353 GHz"
 freq_ref_P = "353 GHz"
 max_nside = 8192
+available_nside = [2048, 4096, 8192]
 [d11]
 class = "ModifiedBlackBodyRealization"
 amplitude_modulation_temp_alm = "dust_gnilc/raw/gnilc_dust_temperature_modulation_alms_lmax768.fits.gz"
@@ -237,13 +239,15 @@ map_pl_index = -3.1
 freq_ref_I = "23 GHz"
 freq_ref_P = "23 GHz"
 max_nside = 8192
+available_nside = [2048, 4096, 8192]
 [s5]
 class = "PowerLaw"
 map_I = "synch/synch_template_nside{nside}_2023.02.25.fits"
 map_pl_index = "synch/synch_beta_nside{nside}_2023.02.16.fits"
 freq_ref_I = "23 GHz"
 freq_ref_P = "23 GHz"
 max_nside = 8192
+available_nside = [2048, 4096, 8192]
 [s6]
 class = "PowerLawRealization"
 largescale_alm = "synch/raw/synch_largescale_template_logpoltens_alm_lmax128_2023.02.24.fits.gz"
@@ -269,6 +273,7 @@ freq_ref_P = "23 GHz"
 spectral_curvature = "synch/synch_curvature_nside{nside}_2023.02.17.fits"
 freq_curve = "23 GHz"
 max_nside = 8192
+available_nside = [2048, 4096, 8192]
 [f1]
 class = "PowerLaw"
 map_I = "pysm_2/ff_t_new.fits"
@@ -373,7 +378,7 @@ max_nside = 8192
 class = "InterpolatingComponent"
 input_units = "uK_RJ"
 freqs = [5.0, 18.7, 24.5, 44.0, 70.0, 100.0, 143.0, 217.0, 353.0, 545.0, 643.0, 729.0, 857.0, 906.0]
-path = "websky/0.4/radio_catalog/{nside}"
+path = "websky/0.4/radio_catalog/background/{nside}"
 interpolation_kind = "linear"
 available_nside = [2048, 4096, 8192]
 pre_applied_beam = {2048=5.1, 4096=2.6, 8192=0.9}

src/pysm3/models/catalog.py

@@ -144,6 +144,8 @@ class PointSourceCatalog(Model):
         Path to the catalog HDF5 file
     """
 
+    includes_smoothing = True
+
     def __init__(
         self,
         catalog_filename,
@@ -184,10 +186,10 @@ def get_emission(
         freqs: u.Quantity[u.GHz],
         fwhm: Optional[u.Quantity[u.arcmin]] = None,
         weights=None,
-        output_units=u.uK_RJ,
         coord=None,
         car_map_resolution: Optional[u.Quantity[u.arcmin]] = None,
         return_car=False,
+        return_healpix=True,
     ):
         """Generate a HEALPix or CAR map of the catalog emission integrated on the bandpass
         and convolved with the beam
@@ -202,22 +204,24 @@ def get_emission(
         weights: np.array
             Array of relative bandpass weights already normalized
             Same length of freqs, if None, uniform weights are assumed
-        output_units: astropy.units
-            Output units of the map
         car_map_resolution: float
             Resolution of the CAR map used by pixell to generate the map, if None,
             it is set to half of the resolution of the HEALPix map given by `self.nside`
         coord: tuple of str
             coordinate rotation, it uses the healpy convention, "Q" for Equatorial,
             "G" for Galactic.
+        return_healpix: bool
+            If True return a HEALPix map
         return_car: bool
-            If True return a CAR map, if False return a HEALPix map
+            If True return a CAR map, in case they are both true, return a tuple,
+            (healpix_map, car_map)
 
         Returns
         -------
         output_map: np.array
-            Output HEALPix or CAR map"""
+            Output HEALPix and/or CAR map, in case both (healpix_map, car_map)"""
 
+        output_units = u.uK_RJ
         convolve_beam = fwhm is not None
         scaling_factor = utils.bandpass_unit_conversion(
             freqs, weights, output_unit=output_units, input_unit=u.Jy / u.sr
@@ -317,22 +321,27 @@ def get_emission(
                 assert (
                     coord is None
                 ), "Coord rotation for CAR not implemented yet, open issue if you need it"
-            else:
+            if return_healpix:
                 from pixell import reproject
 
                 frames_dict = {"Q": "equ", "C": "equ", "G": "gal"}
                 if coord is not None:
                     coord = [frames_dict[frame] for frame in coord]
 
                 log.info("Reprojecting to HEALPix")
-                output_map = (
+                output_map_healpix = (
                     reproject.map2healpix(
                         output_map, self.nside, rot=coord, method="spline"
                     )
                     * output_units
                 )
+                if return_car:
+                    output_map = (output_map_healpix, output_map)
+                else:
+                    output_map = output_map_healpix
 
         else:
+            assert not return_car and return_healpix, "Only HEALPix maps supported"
             aggregate(
                 pix,
                 output_map[1],

src/pysm3/models/dust.py

@@ -28,6 +28,7 @@ def __init__(
         map_mbb_temperature,
         nside,
         max_nside=None,
+        available_nside=None,
         map_Q=None,
         map_U=None,
         has_polarization=True,
@@ -65,8 +66,17 @@ def __init__(
             integer
         nside: int
             Resolution parameter at which this model is to be calculated.
+        max_nside : int
+            Maximum resolution parameter at which this model is defined.
+        available_nside : list of int
+            List of resolution parameters at which the input maps are defined.
         """
-        super().__init__(nside=nside, max_nside=max_nside, map_dist=map_dist)
+        super().__init__(
+            nside=nside,
+            max_nside=max_nside,
+            map_dist=map_dist,
+            available_nside=available_nside,
+        )
         # do model setup
         self.is_IQU = has_polarization and map_Q is None
         self.I_ref = self.read_map(
@@ -168,6 +178,7 @@ def __init__(
         map_mbb_temperature=None,
         nside=None,
         max_nside=None,
+        available_nside=None,
         mpi_comm=None,
         map_dist=None,
         unit_I=None,
@@ -196,6 +207,7 @@ def __init__(
             map_mbb_temperature=map_mbb_temperature,
             nside=nside,
             max_nside=max_nside,
+            available_nside=available_nside,
             unit_I=unit_I,
             unit_Q=unit_Q,
             unit_U=unit_U,
@@ -287,7 +299,7 @@ def get_decorrelation_matrix(
     freqs_all = np.insert(freqs_unconstrained, 0, freq_constrained)
     indref = np.where(freqs_all == freq_constrained)
     corrmatrix = frequency_decorr_model(freqs_all, correlation_length)
-    rho_inv = invert_safe(corrmatrix)
+    rho_inv = invert_safe(corrmatrix.value)
     rho_uu = np.delete(np.delete(rho_inv, indref, axis=0), indref, axis=1)
     rho_uu = invert_safe(rho_uu)
     rho_inv_cu = rho_inv[:, indref]

src/pysm3/models/interpolating.py

@@ -12,6 +12,7 @@
 from .. import utils
 from ..utils import trapz_step_inplace, map2alm
 from .template import Model
+import pixell
 
 log = logging.getLogger("pysm3")
 
@@ -88,6 +89,10 @@ def __init__(
         self.interpolation_kind = interpolation_kind
         self.verbose = verbose
 
+    @property
+    def includes_smoothing(self):
+        return self.pre_applied_beam is not None
+
     def get_filenames(self, path):
         # Override this to implement name convention
         filenames = {}
@@ -104,8 +109,41 @@ def get_emission(
         weights=None,
         fwhm: Optional[u.Quantity[u.arcmin]] = None,
         output_nside: Optional[int] = None,
+        output_car_resol=None,
+        coord=None,
         lmax: Optional[int] = None,
+        return_healpix=True,
+        return_car=False,
     ) -> u.Quantity[u.uK_RJ]:
+        """Return the emission at a frequency or integrated over a bandpass
+
+        Parameters
+        ----------
+            freqs : u.Quantity[u.GHz]
+                Frequency or frequencies at which to compute the emission
+            weights : array-like
+                Weights for the bandpass integration
+            fwhm : u.Quantity[u.arcmin]
+                Beam FWHM
+            output_nside : int
+                NSIDE of the output map
+            coord: tuple of str
+                coordinate rotation, it uses the healpy convention, "Q" for Equatorial,
+                "G" for Galactic.
+            output_car_resol : astropy.Quantity
+                CAR output map resolution, generally in arcmin
+            lmax : int
+                Maximum l of the alm transform
+            return_healpix : bool
+                If True, return the map in healpix format
+            return_car : bool
+                If True, return the map in CAR format
+
+        Returns
+        -------
+            m : u.Quantity[u.uK_RJ]
+                Emission at the requested frequency or integrated over the bandpass
+        """
         nu = utils.check_freq_input(freqs)
         weights = utils.normalize_weights(nu, weights)
 
@@ -157,23 +195,33 @@ def get_emission(
                 zeros = np.zeros_like(out)
                 out = np.array([out, zeros, zeros])
 
-        if (self.pre_applied_beam is not None) and (fwhm is not None):
-            assert lmax is not None, "lmax must be provided when applying a beam"
-            if output_nside is None:
-                output_nside = self.nside
-            pre_beam = (
+        pre_beam = (
+            None
+            if self.pre_applied_beam is None
+            else (
                 self.pre_applied_beam.get(
-                    self.nside, self.pre_applied_beam[self.available_nside[0]]
+                    self.nside, self.pre_applied_beam[str(self.templates_nside)]
                 )
                 * self.pre_applied_beam_units
             )
+        )
+        if (
+            ((pre_beam is None) or (pre_beam == fwhm))
+            and (coord is None)
+            and (return_car is False)
+        ):
+            # no need to go to alm if no beam and no rotation
+            output_maps = [out << u.uK_RJ]
+        else:
+
+            assert lmax is not None, "lmax must be provided when applying a beam"
+            alm = map2alm(out, self.nside, lmax)
             if pre_beam != fwhm:
                 log.info(
                     "Applying the differential beam between: %s %s",
                     str(pre_beam),
                     str(fwhm),
                 )
-                alm = map2alm(out, self.nside, lmax)
 
                 beam = hp.gauss_beam(
                     fwhm.to_value(u.radian), lmax=lmax, pol=True
@@ -184,10 +232,30 @@ def get_emission(
                 )
                 for each_alm, each_beam in zip(alm, beam.T):
                     hp.almxfl(each_alm, each_beam, mmax=lmax, inplace=True)
-                out = hp.alm2map(alm, nside=output_nside, pixwin=False)
+            if coord is not None:
+                rot = hp.Rotator(coord=coord)
+                alm = rot.rotate_alm(alm)
+            output_maps = []
+            if return_healpix:
+                log.info("Alm to map HEALPix")
+                output_maps.append(
+                    hp.alm2map(alm, nside=output_nside if output_nside is not None else self.nside, pixwin=False) << u.uK_RJ
+                )
+            if return_car:
+                log.info("Alm to map CAR")
+                shape, wcs = pixell.enmap.fullsky_geometry(
+                    output_car_resol.to_value(u.radian),
+                    dims=(3,),
+                    variant="fejer1",
+                )
+                ainfo = pixell.curvedsky.alm_info(lmax=lmax)
+                output_maps.append(
+                    pixell.curvedsky.alm2map(
+                        alm, pixell.enmap.empty(shape, wcs), ainfo=ainfo
+                    )
+                )
 
-        # the output of out is always 2D, (IQU, npix)
-        return out << u.uK_RJ
+        return output_maps[0] if len(output_maps) == 1 else tuple(output_maps)
 
     def read_map_by_frequency(self, freq):
         filename = self.maps[freq]

src/pysm3/models/power_law.py

@@ -16,6 +16,7 @@ def __init__(
         map_pl_index,
         nside,
         max_nside=None,
+        available_nside=None,
         has_polarization=True,
         map_Q=None,
         map_U=None,
@@ -48,8 +49,21 @@ def __init__(
             Path to the map to be used as the power law index.
         nside: int
             Resolution parameter at which this model is to be calculated.
+        max_nside: int
+            Maximum resolution parameter at which this model is to be calculated.
+        has_polarization: bool
+            If True, the model will include polarization.
+        available_nside: list of int
+            List of available nside for the input maps.
+        map_dist: pysm.MapDistribution
+            Distribution object used for parallel computing with MPI
         """
-        super().__init__(nside, max_nside=max_nside, map_dist=map_dist)
+        super().__init__(
+            nside,
+            max_nside=max_nside,
+            available_nside=available_nside,
+            map_dist=map_dist,
+        )
         # do model setup
         self.is_IQU = has_polarization and map_Q is None
         self.I_ref = self.read_map(
@@ -141,6 +155,7 @@ def __init__(
         spectral_curvature,
         freq_curve,
         max_nside=None,
+        available_nside=None,
         has_polarization=True,
         map_Q=None,
         map_U=None,
@@ -156,6 +171,7 @@ def __init__(
             map_pl_index=map_pl_index,
             nside=nside,
             max_nside=max_nside,
+            available_nside=available_nside,
             has_polarization=has_polarization,
             map_Q=map_Q,
             map_U=map_U,

src/pysm3/models/template.py

@@ -47,7 +47,7 @@ class Model:
     If libsharp is not available, pixels are distributed uniformly across
     processes, see :py:func:`pysm.mpi.distribute_pixels_uniformly`"""
 
-    def __init__(self, nside, available_nside=None, max_nside=None, map_dist=None):
+    def __init__(self, nside, max_nside=None, available_nside=None, map_dist=None):
         """
         Parameters
         ----------
@@ -66,6 +66,11 @@ def __init__(self, nside, available_nside=None, max_nside=None, map_dist=None):
         self.nside = nside
         self.available_nside = available_nside
         assert nside is not None
+        self.templates_nside = (
+            min(available_nside, key=lambda x: abs(x - self.nside))
+            if available_nside
+            else None
+        )
         self.max_nside = 512 if max_nside is None else max_nside
         self.map_dist = map_dist
 
@@ -79,7 +84,7 @@ def read_map(self, path, unit=None, field=0, nside=None):
         """
         nside = nside if nside is not None else self.nside
         if "{nside}" in path:
-            path = path.format(nside=max(2048, nside))
+            path = path.format(nside=self.templates_nside)
         return read_map(
             path, nside=nside, unit=unit, field=field, map_dist=self.map_dist
         )