Unified SZ implementation that can handle both websky and Agora

.github/workflows/ci_tests.yml

@@ -48,7 +48,7 @@ jobs:
         cache-environment: true
         post-cleanup: 'all'
     - name: Install package
-      run: pip install .[test]
+      run: pip install -e .[test]
       shell: bash -el {0}
     - name: Run Pytest
       run: |

.github/workflows/import_base_requirements.yaml

@@ -37,5 +37,5 @@ jobs:
       shell: bash -el {0}
     - name: Check import
       run: |
-        python -c "import pysm3"
+        PYTHONPATH=src python -c "import pysm3"
       shell: bash -el {0}

.gitignore

@@ -79,4 +79,5 @@ uv.lock
 build/
 dist/
 pysm3.egg-info/
-src/pysm3/_version.py
+src/pysm3/_version.py
+data

docs/models.rst

@@ -136,10 +136,22 @@ Cosmic Infrared Background
 Sunyaev–Zeldovich emission
 ==========================
 
+* `Comparison of WebSky and Agora SZ templates <preprocess-templates/verify_templates/compare_websky_agora_sz.ipynb>`_
+
 - **tsz1**: Thermal SZ emission from WebSky 0.4. Available at $N_{side}=8192$. For more details see :ref:`websky`.
 
+- **tsz2**: Thermal SZ emission from the `Agora simulations <https://doi.org/10.1093/mnras/stae1031>`_ (based on BAHAMAS hydrodynamical simulations). Available at $N_{side}=8192$.
+
+- **tsz3**: Lensed Thermal SZ emission from the `Agora simulations <https://doi.org/10.1093/mnras/stae1031>`_. Available at $N_{side}=8192$.
+
+- **tsz4**: Thermal SZ emission from the `HalfDome simulations <https://doi.org/10.48550/arXiv.2407.17462>`_. Available at $N_{side}=8192$. Multiple realizations are available by overriding the `template_name` with seeds: 100 (default), 102, 104, 106, 108, 110, 112, 114, 116, 118, 120. For example: `halfdome/0.1/tsz/y_b16_halo_res1_s102.fits`. Multiple realizations are available by overriding the `template_name` with seeds: 100 (default), 102, 104, 106, 108, 110, 112, 114, 116, 118, 120. For example: `halfdome/0.1/tsz/y_b16_halo_res1_s102.fits`.
+
 - **ksz1**: Kinetic SZ emission from WebSky 0.4. Available at $N_{side}=4096$. For more details see :ref:`websky`.
 
+- **ksz2**: Kinetic SZ emission from the `Agora simulations <https://doi.org/10.1093/mnras/stae1031>`_ (based on BAHAMAS hydrodynamical simulations). Available at $N_{side}=8192$.
+
+- **ksz3**: Lensed Kinetic SZ emission from the `Agora simulations <https://doi.org/10.1093/mnras/stae1031>`_. Available at $N_{side}=8192$.
+
 Radio galaxies
 ==============
 

pyproject.toml

@@ -90,6 +90,21 @@ version-file = "src/pysm3/_version.py"
 include = [
     "/src/pysm3",
 ]
+packages = [
+    "src/pysm3",
+]
+
+[tool.hatch.build.targets.wheel]
+include = [
+    "/src/pysm3/**",
+    "src/pysm3/data/*",
+]
+packages = [
+    "src/pysm3",
+]
+
+[tool.hatch.build.targets.wheel.force-include]
+"src/pysm3/data/presets.cfg" = "pysm3/data/presets.cfg"
 
 [tool.pytest.ini_options]
 filterwarnings = [

src/pysm3/data/presets.cfg

@@ -385,15 +385,40 @@ pre_applied_beam = {2048=5.1, 4096=2.6, 8192=0.9}
 pre_applied_beam_units = "arcmin"
 max_nside = 8192
 [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"
-max_nside = 4096
+template_name = "websky/0.4/tsz_8192_hp.fits"
+max_nside = 8192
+[ksz2]
+class = "SimpleSZ"
+sz_type = "kinetic"
+template_name = "agora/agora_ukszNG_bahamas80_bnd_unb_1.0e+12_1.0e+18.fits"
+max_nside = 8192
+[tsz2]
+class = "SimpleSZ"
+sz_type = "thermal"
+template_name = "agora/agora_utszNG_bahamas80_bnd_unb_1.0e+12_1.0e+18.fits"
+max_nside = 8192
+[ksz3]
+class = "SimpleSZ"
+sz_type = "kinetic"
+template_name = "agora/agora_lkszNG_bahamas80_bnd_unb_1.0e+12_1.0e+18_lensed.fits"
+max_nside = 8192
+[tsz3]
+class = "SimpleSZ"
+sz_type = "thermal"
+template_name = "agora/agora_ltszNG_bahamas80_bnd_unb_1.0e+12_1.0e+18_lensed.fits"
+max_nside = 8192
+[tsz4]
+class = "SimpleSZ"
+sz_type = "thermal"
+template_name = "halfdome/0.1/tsz/y_b16_halo_res1_s100.fits"
+max_nside = 8192
 [dip1]
 # From NPIPE https://doi.org/10.1051/0004-6361/202038073
 #

src/pysm3/models/__init__.py

@@ -15,7 +15,7 @@
     SPT_CIB_map_scaling,
     WebSkyCMB,
     WebSkyCIB,
-    WebSkySZ,
+    SimpleSZ,
     WebSkyRadioGalaxies,
 )
 from .catalog import PointSourceCatalog

src/pysm3/models/websky.py

@@ -1,11 +1,11 @@
 import os.path
-from pathlib import Path
 
 import numpy as np
 from numba import njit
 
-import pysm3 as pysm
-import pysm3.units as u
+from ..utils import check_freq_input, normalize_weights
+from ..utils import trapz_step_inplace
+from .. import units as u
 
 from .. import utils
 from .cmb import CMBMap
@@ -218,47 +218,52 @@ def get_filenames(self, path):
         return filenames
 
 
-class WebSkySZ(Model):
+class SimpleSZ(Model):
+    """Simple, frequency-independent SZ model using a single template map.
+
+    This component uses a precomputed SZ template map that is independent of
+    observing frequency. The same sky template is used at all frequencies and
+    is retrieved via ``template_name``.
+
+    Parameters
+    ----------
+    nside : int
+        HEALPix NSIDE of the output maps.
+    template_name : str
+        Name or key identifying the SZ template map to download/load via
+        :class:`pysm3.models.utils.RemoteData`. The template is expected to be
+        a single-frequency SZ map in units of uK_CMB.
+    sz_type : str
+        Type of SZ effect to model, either ``"kinetic"`` or ``"thermal"``.
+    max_nside : int
+        Maximum HEALPix NSIDE at which the input template map is available.
+    map_dist : object, optional
+        HEALPix map distribution helper or MPI communicator-like object used
+        by the base :class:`Model` class to handle distributed maps.
+    """
+
     def __init__(
         self,
         nside,
-        version="0.4",
-        sz_type="kinetic",
-        max_nside=None,
+        template_name,
+        sz_type,
+        max_nside,
         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.Quantity[u.GHz], weights=None
     ) -> u.Quantity[u.uK_RJ]:
 
-        freqs = pysm.check_freq_input(freqs)
-        weights = pysm.normalize_weights(freqs, weights)
+        freqs = check_freq_input(freqs)
+        weights = 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
@@ -280,7 +285,7 @@ def get_sz_emission_numba(freqs, weights, m, is_thermal):
     output = np.zeros((3, len(m)), dtype=np.float64)
     for i in range(len(freqs)):
         signal = m * y2uK_CMB(freqs[i]) if is_thermal else m.astype(np.float64)
-        pysm.utils.trapz_step_inplace(freqs, weights, i, signal, output[0])
+        trapz_step_inplace(freqs, weights, i, signal, output[0])
     return output
 
 

src/pysm3/sky.py

@@ -7,8 +7,6 @@
 """
 
 import toml
-
-from . import data
 from . import units as u
 from .models import *
 from .models import Model
@@ -22,6 +20,15 @@ def remove_class_from_dict(d):
 
 def create_components_from_config(config, nside, map_dist=None):
     output_components = []
+    if "class" in config:
+        class_name = config["class"]
+        component_class = globals()[class_name]
+        output_component = component_class(
+            **remove_class_from_dict(config), nside=nside, map_dist=map_dist
+        )
+        output_components.append(output_component)
+        return output_components
+
     for model_name, model_config in config.items():
         try:
             class_name = model_config["class"]
@@ -43,7 +50,9 @@ def create_components_from_config(config, nside, map_dist=None):
         else:
             component_class = globals()[class_name]
             output_component = component_class(
-                **remove_class_from_dict(model_config), nside=nside, map_dist=map_dist
+                **remove_class_from_dict(model_config),
+                nside=nside,
+                map_dist=map_dist,
             )
         output_components.append(output_component)
     return output_components
@@ -56,7 +65,9 @@ def create_components_from_config(config, nside, map_dist=None):
     import importlib_resources as pkg_resources
 
 
-PRESET_MODELS = toml.loads(pkg_resources.read_text(data, "presets.cfg"))
+PRESET_MODELS = toml.loads(
+    pkg_resources.files("pysm3").joinpath("data", "presets.cfg").read_text()
+)
 
 
 def get_pysm_emission(preset_string, nside):

src/pysm3/utils/__init__.py

@@ -30,12 +30,14 @@
 def set_verbosity(level=logging.INFO):
     logger = logging.getLogger("pysm3")
     logger.setLevel(level)
-    if not logger.hasHandlers():
-        handler = logging.StreamHandler()
-        handler.setFormatter(
-            logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
-        )
-        logger.addHandler(handler)
+    if logger.hasHandlers():
+        logger.handlers.clear()
+    handler = logging.StreamHandler()
+    handler.setFormatter(
+        logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    )
+    logger.addHandler(handler)
+    logger.propagate = False
 
 
 def get_relevant_frequencies(freqs, low, high):

tests/test_websky.py

@@ -12,8 +12,8 @@
 from pysm3 import (
     SPT_CIB_map_scaling,
     WebSkyCIB,
+    SimpleSZ,
     WebSkyRadioGalaxies,
-    WebSkySZ,
     utils,
 )  # , WebSkyCMBTensor
 
@@ -127,7 +127,9 @@ def test_sz(tmp_path, monkeypatch, sz_type):
         test_map *= 1e-6
     hp.write_map(path / filename, test_map)
 
-    tsz = WebSkySZ(nside, "0.4", sz_type=sz_type)
+    tsz = SimpleSZ(
+        nside, template_name=str(path / filename), sz_type=sz_type, max_nside=8192
+    )
 
     freq = 100 * u.GHz
     tsz_map = tsz.get_emission(freq).to(