[frozen] Gauss Legendre pixelization implementation of d11

pysm3/data/presets.cfg

@@ -171,6 +171,20 @@ galplane_fix = "dust_gnilc/raw/gnilc_dust_galplane.fits.gz"
 # Configuration for reproducing d10
 # seeds = [8192,777,888]
 # synalm_lmax = 16384
+[d11gl]
+class = "ModifiedBlackBodyRealizationGL"
+largescale_alm = "dust_gnilc/raw/gnilc_dust_largescale_template_logpoltens_alm_nside2048_lmax3072_complex64.fits"
+amplitude_modulation_temp_alm = "dust_gnilc/raw/gnilc_dust_temperature_modulation_alms_lmax3072.fits"
+amplitude_modulation_pol_alm = "dust_gnilc/raw/gnilc_dust_polarization_modulation_alms_lmax3072.fits"
+small_scale_cl = "dust_gnilc/raw/gnilc_dust_small_scales_logpoltens_cl_lmax16384.fits"
+largescale_alm_mbb_index = "dust_gnilc/raw/gnilc_dust_largescale_template_beta_alm_nside2048_lmax1024_complex64.fits"
+small_scale_cl_mbb_index = "dust_gnilc/raw/gnilc_dust_small_scales_beta_cl_lmax16384.fits"
+largescale_alm_mbb_temperature = "dust_gnilc/raw/gnilc_dust_largescale_template_Td_alm_nside2048_lmax1024_complex64.fits"
+small_scale_cl_mbb_temperature = "dust_gnilc/raw/gnilc_dust_small_scales_Td_cl_lmax16384.fits"
+freq_ref = "353 GHz"
+# Configuration for reproducing d10
+# seeds = [8192,777,888]
+# synalm_lmax = 16384
 [d12]
 class = "ModifiedBlackBodyLayers"
 map_layers = "mkd_dust/2048/thermaldust_ampl{layer}.fits"

pysm3/models/__init__.py

@@ -10,4 +10,5 @@
 from .dust_layers import ModifiedBlackBodyLayers
 from .co_lines import COLines
 from .dust_realization import ModifiedBlackBodyRealization
+from .dust_realization_gl import ModifiedBlackBodyRealizationGL
 from .power_law_realization import PowerLawRealization

pysm3/models/dust.py

@@ -10,7 +10,7 @@
 
 
 class ModifiedBlackBody(Model):
-    """ This is a model for modified black body emission.
+    """This is a model for modified black body emission.
 
     Attributes
     ----------
@@ -36,7 +36,7 @@ def __init__(
         unit_mbb_temperature=None,
         map_dist=None,
     ):
-        """ This function initializes the modified black body model.
+        """This function initializes the modified black body model.
 
         The initialization of this model consists of reading in emission
         templates from file, reading in spectral parameter maps from
@@ -131,9 +131,10 @@ def get_emission_numba(
     mbb_index,
     mbb_temperature,
 ):
-    output = np.zeros((3, len(I_ref)), dtype=I_ref.dtype)
+    map_shape = (3,) + I_ref.shape  # support both 1D HEALPix and 2D GL
+    output = np.zeros(map_shape, dtype=I_ref.dtype)
     if len(freqs) > 1:
-        temp = np.zeros((3, len(I_ref)), dtype=I_ref.dtype)
+        temp = np.zeros(map_shape, dtype=I_ref.dtype)
     else:
         temp = output
 
@@ -174,7 +175,7 @@ def __init__(
         unit_mbb_temperature=None,
         correlation_length=None,
     ):
-        """ See parent class for other documentation.
+        """See parent class for other documentation.
 
         Parameters
         ----------
@@ -203,7 +204,7 @@ def __init__(
 
     @u.quantity_input
     def get_emission(self, freqs: u.GHz, weights=None) -> u.uK_RJ:
-        """ Function to calculate the emission of a decorrelated modified black
+        """Function to calculate the emission of a decorrelated modified black
         body model.
         """
         freqs = utils.check_freq_input(freqs)
@@ -237,7 +238,7 @@ def get_emission(self, freqs: u.GHz, weights=None) -> u.uK_RJ:
 
 @u.quantity_input
 def frequency_decorr_model(freqs: u.GHz, correlation_length: u.dimensionless_unscaled):
-    """ Function to calculate the frequency decorrelation method of
+    """Function to calculate the frequency decorrelation method of
     Vansyngel+17.
     """
     log_dep = np.log(freqs[:, None] / freqs[None, :])
@@ -250,7 +251,7 @@ def get_decorrelation_matrix(
     freqs_unconstrained: u.GHz,
     correlation_length: u.dimensionless_unscaled,
 ):
-    """ Function to calculate the correlation matrix between observed
+    """Function to calculate the correlation matrix between observed
     frequencies. This model is based on the proposed model for decorrelation
     of Vansyngel+17. The proposed frequency covariance matrix in this paper
     is implemented, and a constrained Gaussian realization for the unobserved
@@ -323,7 +324,7 @@ def invert_safe(matrix):
 
 @njit
 def blackbody_ratio(freq_to, freq_from, temp):
-    """ Function to calculate the flux ratio between two frequencies for a
+    """Function to calculate the flux ratio between two frequencies for a
     blackbody at a given temperature.
 
     Parameters
@@ -391,6 +392,6 @@ def blackbody_nu(freq, temp):
     boltzm1 = np.expm1(log_boltz)
 
     # Calculate blackbody flux
-    bb_nu = 2.0 * h * (freq * 1e9) ** 3 / (c ** 2 * boltzm1)
+    bb_nu = 2.0 * h * (freq * 1e9) ** 3 / (c**2 * boltzm1)
 
     return bb_nu

pysm3/models/dust_realization.py

@@ -19,6 +19,7 @@ def __init__(
         largescale_alm_mbb_temperature,
         small_scale_cl_mbb_temperature,
         nside,
+        fwhm=None,
         galplane_fix=None,
         seeds=None,
         synalm_lmax=None,
@@ -62,6 +63,9 @@ def __init__(
             excess power in the full sky spectra due to the generated small scales being
             too strong on the galactic plane.
             By default in d9,d10,d11 we use the input GNILC map with a resolution of 21.8'
+        fwhm: gaussian beam width
+            Produce small scales already smoothed with a gaussian beam, this makes the spherical
+            harmonics transforms more accurate compared to smoothing the outputs
         seeds: list of ints
             List of seeds used for generating the small scales, first is used for the template,
             the second for the spectral index, the third for the black body temperature.
@@ -100,8 +104,7 @@ def __init__(
         else:
             self.galplane_fix_map = None
         self.largescale_alm_mbb_index = self.read_alm(
-            largescale_alm_mbb_index,
-            has_polarization=False,
+            largescale_alm_mbb_index, has_polarization=False
         ).to(u.dimensionless_unscaled)
         self.small_scale_cl_mbb_index = self.read_cl(small_scale_cl_mbb_index).to(
             u.dimensionless_unscaled
@@ -112,6 +115,22 @@ def __init__(
         self.small_scale_cl_mbb_temperature = self.read_cl(
             small_scale_cl_mbb_temperature
         ).to(u.K ** 2)
+        self.fwhm = fwhm
+        if self.fwhm is not None:
+            self.fwhm = self.fwhm.to(u.rad)
+            B_ell_squared = (
+                hp.gauss_beam(
+                    fwhm=self.fwhm.value, lmax=self.small_scale_cl.shape[-1] - 1
+                )
+                ** 2
+            )
+            self.small_scale_cl *= B_ell_squared
+            self.small_scale_cl_mbb_temperature *= B_ell_squared[
+                : len(self.small_scale_cl_mbb_temperature)
+            ]
+            self.small_scale_cl_mbb_index *= B_ell_squared[
+                : len(self.small_scale_cl_mbb_index)
+            ]
         self.nside = int(nside)
         (
             self.I_ref,
@@ -151,8 +170,7 @@ def draw_realization(self, synalm_lmax=None, seeds=None):
         map_small_scale[1:] *= hp.alm2map(self.modulate_alm[1].value, self.nside)
 
         map_small_scale += hp.alm2map(
-            self.template_largescale_alm.value,
-            nside=self.nside,
+            self.template_largescale_alm.value, nside=self.nside
         )
 
         if self.galplane_fix_map is not None:
@@ -177,11 +195,7 @@ def draw_realization(self, synalm_lmax=None, seeds=None):
             np.random.seed(seed)
             input_cl = getattr(self, f"small_scale_cl_{key}")
             output_unit = np.sqrt(1 * input_cl.unit).unit
-            alm_small_scale = hp.synalm(
-                input_cl.value,
-                lmax=synalm_lmax,
-                new=True,
-            )
+            alm_small_scale = hp.synalm(input_cl.value, lmax=synalm_lmax, new=True)
 
             alm_small_scale = hp.almxfl(
                 alm_small_scale, np.ones(min(3 * self.nside - 1, synalm_lmax + 1))
@@ -190,8 +204,7 @@ def draw_realization(self, synalm_lmax=None, seeds=None):
             output[key] *= modulate_map_I
             output[key] += (
                 hp.alm2map(
-                    getattr(self, f"largescale_alm_{key}").value,
-                    nside=self.nside,
+                    getattr(self, f"largescale_alm_{key}").value, nside=self.nside
                 )
                 * output_unit
             )

pysm3/models/dust_realization_gl.py

@@ -0,0 +1,125 @@
+import numpy as np
+import healpy as hp
+
+from .. import utils
+from .dust_realization import ModifiedBlackBodyRealization
+
+
+class ModifiedBlackBodyRealizationGL(ModifiedBlackBodyRealization):
+    def __init__(
+        self,
+        largescale_alm,
+        freq_ref,
+        amplitude_modulation_temp_alm,
+        amplitude_modulation_pol_alm,
+        small_scale_cl,
+        largescale_alm_mbb_index,
+        small_scale_cl_mbb_index,
+        largescale_alm_mbb_temperature,
+        small_scale_cl_mbb_temperature,
+        fwhm,
+        nside,
+        galplane_fix=None,
+        seeds=None,
+        synalm_lmax=None,
+        has_polarization=True,
+        map_dist=None,
+    ):
+        super().__init__(
+            largescale_alm=largescale_alm,
+            freq_ref=freq_ref,
+            amplitude_modulation_temp_alm=amplitude_modulation_temp_alm,
+            amplitude_modulation_pol_alm=amplitude_modulation_pol_alm,
+            small_scale_cl=small_scale_cl,
+            largescale_alm_mbb_index=largescale_alm_mbb_index,
+            small_scale_cl_mbb_index=small_scale_cl_mbb_index,
+            largescale_alm_mbb_temperature=largescale_alm_mbb_temperature,
+            small_scale_cl_mbb_temperature=small_scale_cl_mbb_temperature,
+            fwhm=fwhm,
+            nside=nside,
+            galplane_fix=galplane_fix,
+            seeds=seeds,
+            synalm_lmax=synalm_lmax,
+            has_polarization=has_polarization,
+            map_dist=map_dist,
+        )
+
+    def draw_realization(self, synalm_lmax=None, seeds=None):
+
+        if seeds is None:
+            seeds = (None, None, None)
+
+        if synalm_lmax is None:
+            synalm_lmax = min(16384, 3 * self.nside - 1)
+
+        np.random.seed(seeds[0])
+
+        alm_small_scale = hp.synalm(
+            list(self.small_scale_cl.value)
+            + [np.zeros_like(self.small_scale_cl[0])] * 3,
+            lmax=synalm_lmax,
+            new=True,
+        )
+
+        lmax = min(synalm_lmax, 3 * self.nside - 1)
+        alm_small_scale = np.array(
+            [hp.almxfl(each, np.ones(lmax)) for each in alm_small_scale]
+        )
+        nlon, nlat = utils.lmax2nlon(lmax), utils.lmax2nlat(lmax)
+        map_small_scale = utils.gl_alm2map(alm_small_scale, lmax, nlon=nlon, nlat=nlat)
+
+        # need later for beta, Td
+        modulate_map_I = utils.gl_alm2map(
+            self.modulate_alm[0].value, lmax, nlon=nlon, nlat=nlat
+        )[0]
+
+        map_small_scale[0] *= modulate_map_I
+        map_small_scale[1:] *= utils.gl_alm2map(
+            self.modulate_alm[1].value, lmax, nlon=nlon, nlat=nlat
+        )[0]
+
+        map_small_scale += utils.gl_alm2map(
+            self.template_largescale_alm.value, lmax, nlon=nlon, nlat=nlat
+        )
+
+        output_IQU = (
+            utils.log_pol_tens_to_map(map_small_scale)
+            * self.template_largescale_alm.unit
+        ) * 0.911  # includes color correction
+        # See https://github.com/galsci/pysm/issues/99
+
+        # Fixed values for comparison with d9
+        # mbb_index = 1.48 * u.dimensionless_unscaled
+        # mbb_temperature = 19.6 * u.K
+
+        output = {}
+
+        for seed, key in zip(seeds[1:], ["mbb_index", "mbb_temperature"]):
+            np.random.seed(seed)
+            input_cl = getattr(self, f"small_scale_cl_{key}")
+            output_unit = np.sqrt(1 * input_cl.unit).unit
+            alm_small_scale = hp.synalm(input_cl.value, lmax=synalm_lmax, new=True)
+
+            alm_small_scale = hp.almxfl(alm_small_scale, np.ones(lmax))
+            output[key] = (
+                utils.gl_alm2map(alm_small_scale, lmax, nlon=nlon, nlat=nlat)[0]
+                * output_unit
+            )
+            output[key] *= modulate_map_I
+            output[key] += (
+                utils.gl_alm2map(
+                    getattr(self, f"largescale_alm_{key}").value,
+                    lmax,
+                    nlon=nlon,
+                    nlat=nlat,
+                )[0]
+                * output_unit
+            )
+
+        return (
+            output_IQU[0],
+            output_IQU[1],
+            output_IQU[2],
+            output["mbb_index"],
+            output["mbb_temperature"],
+        )

pysm3/tests/test_dust_gl.py

@@ -0,0 +1,63 @@
+import healpy as hp
+from astropy.tests.helper import assert_quantity_allclose
+
+import pysm3
+from pysm3 import utils
+from pysm3 import units as u
+
+import pytest
+
+
+@pytest.mark.parametrize("freq", [353 * u.GHz, 545 * u.GHz, 857 * u.GHz])
+def test_d11gl_lownside(freq):
+    nside = 512
+
+    beamwidth = 14 * u.arcmin  # 2 pixels per beam
+
+    d11_configuration = pysm3.sky.PRESET_MODELS["d11"].copy()
+    d11_configuration["fwhm"] = beamwidth
+    del d11_configuration["class"]
+    del d11_configuration["galplane_fix"]
+    output_healpix = pysm3.models.ModifiedBlackBodyRealization(
+        nside=nside, seeds=[8192, 777, 888], **d11_configuration
+    ).get_emission(freq)
+    output_gl = pysm3.models.ModifiedBlackBodyRealizationGL(
+        nside=nside, seeds=[8192, 777, 888], **d11_configuration
+    ).get_emission(freq)
+    lmax = 3 * nside - 1
+    alm_dx11gl = utils.gl_map2alm(output_gl.value, lmax)
+    output_gl_to_healpix = hp.alm2map(alm_dx11gl, nside=nside) * output_gl.unit
+
+    rtol = 1e-4
+    atol = {353: 20 * u.uK_RJ, 545: 60 * u.uK_RJ, 857: 150 * u.uK_RJ}
+    assert_quantity_allclose(
+        output_healpix[0], output_gl_to_healpix[0], rtol=rtol, atol=atol[freq.value]
+    )
+    assert_quantity_allclose(
+        output_healpix[1:],
+        output_gl_to_healpix[1:],
+        rtol=rtol,
+        atol=atol[freq.value] / 10,
+    )
+
+
+# @pytest.mark.skipif(
+#     psutil.virtual_memory().total * u.byte < 20 * u.GB,
+#     reason="Running d11 at high lmax requires 20 GB of RAM",
+# )
+# def test_d10_vs_d11():
+#     nside = 2048
+#
+#     freq = 857 * u.GHz
+#
+#     output_d10 = pysm3.Sky(preset_strings=["d10"], nside=nside).get_emission(freq)
+#     d11_configuration = pysm3.sky.PRESET_MODELS["d11"].copy()
+#     del d11_configuration["class"]
+#     d11 = pysm3.models.ModifiedBlackBodyRealization(
+#         nside=nside, seeds=[8192, 777, 888], synalm_lmax=16384, **d11_configuration
+#     )
+#     output_d11 = d11.get_emission(freq)
+#
+#     rtol = 1e-5
+#
+#     assert_quantity_allclose(output_d10, output_d11, rtol=rtol, atol=0.05 * u.uK_RJ)