feat: implement SimpleSZ that can handle both WebSky and Agora

pysm3/data/presets.cfg

@@ -366,12 +366,22 @@ interpolation_kind = "linear"
 apply_SPT_correction = false
 max_nside = 4096
 [ksz1]
-class = "WebSkySZ"
-version = "0.4"
+class = "SimpleSZ"
 sz_type = "kinetic"
+template_name= "websky/0.4/ksz.fits"
 max_nside = 4096
 [tsz1]
-class = "WebSkySZ"
-version = "0.4"
+class = "SimpleSZ"
 sz_type = "thermal"
+template_name = "websky/0.4/tsz_8192_hp.fits"
 max_nside = 4096
+[ksz2]
+class = "SimpleSZ"
+sz_type = "kinetic"
+template_name = "agora/agora_lkszNGbahamas80_nside8192_uk.fits"
+max_nside = 8192
+[tsz2]
+class = "SimpleSZ"
+sz_type = "thermal"
+template_name = "agora/agora_ltszNGbahamas80_nside8192_uk.fits"
+max_nside = 8192

pysm3/models/websky.py

@@ -217,40 +217,39 @@ def get_filenames(self, path):
         return filenames
 
 
-class WebSkySZ(Model):
+class SimpleSZ(Model):
+    """Simple SZ model with a single frequency-independent map
+
+    Parameters
+    ----------
+
+    nside : int
+        HEALPix NSIDE of the output maps
+    template_path : str
+        path to the FITS file containing the template
+    sz_type : str
+        "kinetic" or "thermal"
+    max_nside : int
+        maximum NSIDE at which the input maps are available
+    """
+
     def __init__(
         self,
         nside,
-        version="0.4",
-        sz_type="kinetic",
-        max_nside=None,
+        template_name,
+        sz_type,
+        max_nside,
+        version=None,
         map_dist=None,
     ):
 
-        if max_nside is None:
-            if sz_type == "kinetic":
-                max_nside = 4096
-            if sz_type == "thermal":
-                max_nside = 8192
         super().__init__(nside=nside, max_nside=max_nside, map_dist=map_dist)
         self.version = str(version)
         self.sz_type = sz_type
         self.remote_data = utils.RemoteData()
-        filename = self.remote_data.get(self.get_filename())
+        filename = self.remote_data.get(template_name)
         self.m = self.read_map(filename, field=0, unit=u.uK_CMB)
 
-    def get_filename(self):
-        """Get SZ filenames for a websky version"""
-
-        path = Path("websky") / self.version
-
-        if self.sz_type == "kinetic":
-            path = path / "ksz.fits"
-        elif self.sz_type == "thermal":
-            path = path / "tsz_8192_hp.fits"
-
-        return str(path)
-
     @u.quantity_input
     def get_emission(self, freqs: u.GHz, weights=None) -> u.uK_RJ:
 
@@ -314,56 +313,3 @@ def __init__(
             max_nside=max_nside,
             map_dist=map_dist,
         )
-
-
-class AgoraSZ(Model):
-    def __init__(
-        self,
-        nside,
-        sz_type="kinetic",
-        max_nside=None,
-        map_dist=None,
-    ):
-
-        if max_nside is None:
-            if sz_type == "kinetic":
-                max_nside = 8192
-            if sz_type == "thermal":
-                max_nside = 8192
-        super().__init__(nside=nside, max_nside=max_nside, map_dist=map_dist)
-        self.sz_type = sz_type
-        # self.remote_data = utils.RemoteData()
-        # filename = self.remote_data.get(self.get_filename())
-        filename = self.get_filename()
-        self.m = self.read_map(filename, field=0, unit=u.uK_CMB)
-
-    def get_filename(self):
-        """Get SZ filenames for an agora version"""
-
-        if self.sz_type == "kinetic":
-            path = "/Users/kristen/Documents/GitHub/ACT-Simulations/agora/agora_lkszNGbahamas80_nside8192_uk.fits"
-        elif self.sz_type == "thermal":
-            path = "/Users/kristen/Documents/GitHub/ACT-Simulations/agora/agora_ltszNGbahamas80_nside8192_uk.fits"
-
-        return str(path)
-
-    @u.quantity_input
-    def get_emission(self, freqs: u.GHz, weights=None) -> u.uK_RJ:
-
-        freqs = pysm.check_freq_input(freqs)
-        weights = pysm.normalize_weights(freqs, weights)
-
-        # input map is in uK_CMB, we multiply the weights which are
-        # in uK_RJ by the conversion factor of uK_CMB->uK_RJ
-        # this is the equivalent of
-        weights = (weights * u.uK_CMB).to_value(
-            u.uK_RJ, equivalencies=u.cmb_equivalencies(freqs * u.GHz)
-        )
-
-        is_thermal = self.sz_type == "thermal"
-        output = (
-            get_sz_emission_numba(freqs, weights, self.m.value, is_thermal) << u.uK_RJ
-        )
-
-        # the output of out is always 2D, (IQU, npix)
-        return output