Support for a custom beam in 3.4.0

pysm3/models/template.py

@@ -7,6 +7,7 @@
 Objects:
     Model
 """
+
 import logging
 import numpy as np
 import healpy as hp
@@ -124,6 +125,7 @@ def get_emission(self, freqs: u.GHz, weights=None) -> u.uK_RJ:
 def apply_smoothing_and_coord_transform(
     input_map,
     fwhm=None,
+    beam_window=None,
     rot=None,
     lmax=None,
     output_nside=None,
@@ -150,6 +152,8 @@ def apply_smoothing_and_coord_transform(
     fwhm : astropy.units.Quantity
         Full width at half-maximum, defining the
         Gaussian kernels to be applied.
+    beam_window: array, optional
+        Custom beam window function (:math:`B_\ell`)
     rot: hp.Rotator
         Apply a coordinate rotation give a healpy `Rotator`, e.g. if the
         inputs are in Galactic, `hp.Rotator(coord=("G", "C"))` rotates
@@ -240,8 +244,21 @@ def apply_smoothing_and_coord_transform(
                         error,
                     )
         if fwhm is not None:
+            assert beam_window is None, "Either FWHM or beam_window"
             log.info("Smoothing with fwhm of %s", str(fwhm))
             hp.smoothalm(alm, fwhm=fwhm.to_value(u.rad), inplace=True, pol=True)
+        if beam_window is not None:
+            assert fwhm is None, "Either FWHM or beam_window"
+            log.info("Smoothing with a custom isotropic beam")
+            # smoothalm does not support polarized beam
+            for i in range(3):
+                try:
+                    beam_window_i = beam_window[:, i]
+                    log.info("Using polarized beam")
+                except IndexError:
+                    beam_window_i = beam_window
+                    log.info("Using the same beam for all components")
+                hp.smoothalm(alm[i], beam_window=beam_window_i, inplace=True)
         if rot is not None:
             log.info("Rotate Alm")
             rot.rotate_alm(alm, inplace=True)

pysm3/tests/test_smoothing.py

@@ -15,6 +15,7 @@
 from astropy.tests.helper import assert_quantity_allclose
 import pytest
 
+INITIAL_FWHM = (1 * u.deg).to_value(u.radian)
 FWHM = (5 * u.deg).to_value(u.radian)
 NSIDE = 128
 CAR_RESOL = 12 * u.arcmin
@@ -25,16 +26,31 @@
     scope="module"
 )  # scope makes the fixture just run once per execution of module
 def input_map():
-    beam_window = hp.gauss_beam(fwhm=FWHM, lmax=LMAX) ** 2
+    beam_window = hp.gauss_beam(fwhm=INITIAL_FWHM, lmax=LMAX) ** 2
     cl = np.zeros((6, len(beam_window)))
     cl[0:3] = beam_window
     np.random.seed(7)
     m = hp.synfast(cl, NSIDE, lmax=LMAX, new=True) * u.uK_RJ
     return m
 
 
-def test_smoothing_healpix(input_map):
+def test_smoothing_healpix_beamwindow(input_map):
+    beam_window = hp.gauss_beam(fwhm=FWHM, lmax=LMAX, pol=True)
+
+    smoothed_map = apply_smoothing_and_coord_transform(
+        input_map, lmax=LMAX, beam_window=beam_window
+    )
+    assert input_map.shape[0] == 3
+    assert smoothed_map.shape == input_map.shape
+    assert_quantity_allclose(
+        actual=smoothed_map,
+        desired=hp.smoothing(input_map, fwhm=FWHM, lmax=LMAX, use_pixel_weights=True)
+        * input_map.unit,
+        rtol=1e-7,
+    )
 
+
+def test_smoothing_healpix(input_map):
     smoothed_map = apply_smoothing_and_coord_transform(
         input_map, lmax=LMAX, fwhm=FWHM * u.radian
     )
@@ -48,7 +64,6 @@ def test_smoothing_healpix(input_map):
 
 
 def test_car_nosmoothing(input_map):
-
     # `enmap_from_healpix` has no iteration or weights
     # so for test purpose we reproduce it here
     alm = (
@@ -78,7 +93,6 @@ def test_car_nosmoothing(input_map):
 
 
 def test_healpix_output_nside(input_map):
-
     output_nside = 64
     output_map = apply_smoothing_and_coord_transform(
         input_map, fwhm=None, output_nside=output_nside, lmax=LMAX

pysm3/utils/__init__.py

@@ -151,7 +151,9 @@ def bandpass_unit_conversion(
         if weights.min() < cut:
             good = np.logical_not(weights < cut)
             log.info(
-                "Removing {}/{} points below {}".format(good.sum(), len(good), cut)
+                "Removing {}/{} points below {}".format(
+                    (good == 0).sum(), len(good), cut
+                )
             )
             weights = weights[good]
             freqs = freqs[good]

sharp.py

@@ -0,0 +1,10 @@
+import libsharp
+import numpy as np
+from mpi4py import MPI
+
+lmax = 10
+mpi_comm = MPI.COMM_WORLD
+local_m_indices = np.arange(mpi_comm.rank, lmax + 1, mpi_comm.size, dtype=np.int32)
+order = libsharp.packed_real_order(lmax, ms=local_m_indices)
+print(MPI.COMM_WORLD.rank, order.mval(), order.mvstart(), order.local_size())
+