From c59ee7a973ce0917a849f0ea6f38c1d14f2b80ba Mon Sep 17 00:00:00 2001
From: Andrea Zonca <code@andreazonca.com>
Date: Mon, 22 Jul 2024 23:02:50 -0500
Subject: [PATCH] feat: implement SimpleSZ that can handle both WebSky and
 Agora

---
 pysm3/data/presets.cfg | 18 ++++++--
 pysm3/models/websky.py | 96 +++++++++---------------------------------
 2 files changed, 35 insertions(+), 79 deletions(-)

diff --git a/pysm3/data/presets.cfg b/pysm3/data/presets.cfg
index 4433d39b..9e1f9f48 100644
--- a/pysm3/data/presets.cfg
+++ b/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
\ No newline at end of file
diff --git a/pysm3/models/websky.py b/pysm3/models/websky.py
index 5174874d..6fa543c9 100644
--- a/pysm3/models/websky.py
+++ b/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