Make COLines compatible with the other emissions

* moved inside models
* renamed arguments
* fixed unit conversion, input maps are in uK_CMB

pysm3/__init__.py

@@ -1,6 +1,5 @@
 # flake8: noqa
 from ._astropy_init import *  # noqa
-
 import warnings
 from astropy.utils.exceptions import AstropyDeprecationWarning
 
@@ -22,4 +21,3 @@
 from .distribution import MapDistribution
 from .mpi import mpi_smoothing
 from .utils import normalize_weights, bandpass_unit_conversion, check_freq_input
-from .co_lines import COLines

pysm3/models/__init__.py

@@ -8,3 +8,4 @@
 from .interpolating import InterpolatingComponent
 from .cmb import CMBMap, CMBLensed
 from .dust_layers import ModifiedBlackBodyLayers
+from .co_lines import COLines

pysm3/models/co_lines.py

@@ -1,15 +1,9 @@
 import numpy as np
-
 import healpy as hp
 
-try:  # PySM >= 3.2.1
-    import pysm3.units as u
-    import pysm3 as pysm
-except ImportError:
-    import pysm.units as u
-    import pysm
-
-from . import utils
+from .. import units as u
+from .. import utils
+from .template import Model
 
 
 def build_lines_dict(lines, maps):
@@ -22,11 +16,10 @@ def build_lines_dict(lines, maps):
     return dict(zip(lines, np.atleast_2d(maps)))
 
 
-class COLines(pysm.Model):
+class COLines(Model):
     def __init__(
         self,
-        target_nside,
-        output_units,
+        nside,
         has_polarization=True,
         lines=["10", "21", "32"],
         include_high_galactic_latitude_clouds=False,
@@ -45,10 +38,8 @@ def __init__(
 
         Parameters
         ----------
-        target_nside : int
+        nside : int
             HEALPix NSIDE of the output maps
-        output_units : str
-            unit string as defined by `pysm.convert_units`, e.g. uK_RJ, K_CMB
         has_polarization : bool
             whether or not to simulate also polarization maps
         lines : list of strings
@@ -79,16 +70,16 @@ def __init__(
             "21": 230.538 * u.GHz,
             "32": 345.796 * u.GHz,
         }
-        self.target_nside = target_nside
+        self.nside = nside
 
-        self.template_nside = 512 if self.target_nside <= 512 else 2048
+        self.template_nside = 512 if self.nside <= 512 else 2048
 
-        super().__init__(nside=target_nside, map_dist=map_dist)
+        super().__init__(nside=nside, map_dist=map_dist)
 
         self.remote_data = utils.RemoteData()
 
-        self.planck_templatemap_filename = "co/HFI_CompMap_CO-Type1_{}_R2.00_ring.fits".format(
-            self.template_nside
+        self.planck_templatemap_filename = (
+            "co/HFI_CompMap_CO-Type1_{}_R2.00_ring.fits".format(self.template_nside)
         )
         self.planck_templatemap = build_lines_dict(
             self.lines,
@@ -98,7 +89,7 @@ def __init__(
                     field=[self.line_index[line] for line in self.lines],
                     unit=u.K_CMB,
                 ),
-                nside_out=self.target_nside,
+                nside_out=self.nside,
             )
             << u.K_CMB,
         )
@@ -136,20 +127,19 @@ def __init__(
                 ),
             )
 
-        self.output_units = u.Unit(output_units)
         self.verbose = verbose
 
     @u.quantity_input
     def get_emission(self, freqs: u.GHz, weights=None) -> u.uK_RJ:
         freqs = utils.check_freq_input(freqs)
         weights = utils.normalize_weights(freqs, weights)
-        out = np.zeros((3, hp.nside2npix(self.target_nside)), dtype=np.double)
+        out = np.zeros((3, hp.nside2npix(self.nside)), dtype=np.double)
         for line in self.lines:
             line_freq = self.line_frequency[line].to_value(u.GHz)
             if line_freq >= freqs[0] and line_freq <= freqs[-1]:
                 weight = np.interp(line_freq, freqs, weights)
                 convert_to_uK_RJ = (1 * u.K_CMB).to_value(
-                    self.output_units,
+                    u.K_RJ,
                     equivalencies=u.cmb_equivalencies(line_freq * u.GHz),
                 )
                 I_map = self.planck_templatemap[line]
@@ -203,7 +193,7 @@ def simulate_high_galactic_latitude_CO(self, line):
             R_em = 6.6
             model = "LogSpiral"
 
-            nside = self.target_nside
+            nside = self.nside
             Itot_o, _ = cl.integrate_intensity_map(
                 self.planck_templatemap[line],
                 hp.get_nside(self.planck_templatemap[line]),

pysm3/tests/test_co.py

@@ -1,5 +1,6 @@
 from ..utils import RemoteData
 from .. import units as u
+from .. import Sky
 
 import numpy as np
 import healpy as hp
@@ -13,11 +14,12 @@
 @pytest.mark.parametrize("include_high_galactic_latitude_clouds", [False, True])
 def test_co(include_high_galactic_latitude_clouds):
 
+    line = "10"
+
     co = COLines(
-        target_nside=16,
-        output_units="K_CMB",
+        nside=16,
         has_polarization=True,
-        lines=["10"],
+        lines=[line],
         include_high_galactic_latitude_clouds=include_high_galactic_latitude_clouds,
         polarization_fraction=0.001,
         theta_high_galactic_latitude_deg=20.0,
@@ -26,16 +28,17 @@ def test_co(include_high_galactic_latitude_clouds):
         run_mcmole3d=False,
     )
 
-    co_map = co.get_emission(115.271 * u.GHz)
+    line_freq = co.line_frequency[line]
+    co_map = co.get_emission(line_freq)
 
     tag = "wHGL" if include_high_galactic_latitude_clouds else "noHGL"
     remote_data = RemoteData()
     expected_map_filename = remote_data.get(
         "co/testing/CO10_TQUmaps_{}_nside16_ring.fits.zip".format(tag)
     )
     expected_co_map = (
-        hp.read_map(expected_map_filename, field=(0, 1, 2), dtype=np.float64) * u.uK_RJ
-    )
+        hp.read_map(expected_map_filename, field=(0, 1, 2), dtype=np.float64) * u.uK_CMB
+    ).to(u.uK_RJ, equivalencies=u.cmb_equivalencies(line_freq))
 
     assert_quantity_allclose(co_map, expected_co_map, rtol=1e-5)
 
@@ -57,3 +60,23 @@ def test_co(include_high_galactic_latitude_clouds):
         expected_co_map[nonzero_values] * 0.05475254098360655,
         rtol=1e-4,
     )
+
+
+@pytest.mark.parametrize("model_tag", ["co2", "co3"])
+def test_co_model(model_tag):
+    include_high_galactic_latitude_clouds = model_tag == "co3"
+
+    model = Sky(preset_strings=[model_tag], nside=16, output_unit="uK_CMB")
+
+    co_map = model.get_emission(115.271 * u.GHz)
+
+    tag = "wHGL" if include_high_galactic_latitude_clouds else "noHGL"
+    remote_data = RemoteData()
+    expected_map_filename = remote_data.get(
+        "co/testing/CO10_TQUmaps_{}_nside16_ring.fits.zip".format(tag)
+    )
+    expected_co_map = (
+        hp.read_map(expected_map_filename, field=(0, 1, 2), dtype=np.float64) * u.uK_CMB
+    )
+
+    assert_quantity_allclose(co_map, expected_co_map, rtol=1e-5)