fix: normalization of weights in test and handle non-healpix

galsci · Sep 18, 2024 · 21d0ade · 21d0ade
1 parent c0bed53
commit 21d0ade
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 12 deletions.
diff --git a/src/pysm3/models/catalog.py b/src/pysm3/models/catalog.py
@@ -125,7 +125,6 @@ def __init__(
     ):
         self.catalog_filename = catalog_filename
         self.nside = nside
-        self.shape = (3, hp.nside2npix(nside))
         self.wcs = target_wcs
 
         with h5py.File(self.catalog_filename) as f:
@@ -138,9 +137,10 @@ def __init__(
 
     def get_fluxes(self, freqs: u.GHz, coeff="logpolycoefflux", weights=None):
         """Get catalog fluxes in Jy integrated over a bandpass"""
-        weights /= trapezoid(weights, x=freqs.to_value(u.GHz))
+        freqs = utils.check_freq_input(freqs)
+        weights = utils.normalize_weights(freqs, weights)
         with h5py.File(self.catalog_filename) as f:
-            flux = evaluate_model(freqs.to_value(u.GHz), weights, np.array(f[coeff]))
+            flux = evaluate_model(freqs, weights, np.array(f[coeff]))
         return flux * u.Jy
 
     @u.quantity_input
@@ -178,8 +178,6 @@ def get_emission(
         -------
         output_map: np.array
             Output HEALPix or CAR map"""
-        with h5py.File(self.catalog_filename) as f:
-            pix = hp.ang2pix(self.nside, f["theta"], f["phi"])
         scaling_factor = utils.bandpass_unit_conversion(
             freqs, weights, output_unit=output_units, input_unit=u.Jy / u.sr
         )
@@ -195,7 +193,12 @@ def get_emission(
         fluxes_I = self.get_fluxes(freqs, weights=weights, coeff="logpolycoefflux")
 
         if fwhm is None:
-            output_map = np.zeros(self.shape, dtype=np.float32) * output_units
+            with h5py.File(self.catalog_filename) as f:
+                pix = hp.ang2pix(self.nside, f["theta"], f["phi"])
+            output_map = (
+                np.zeros((3, hp.nside2npix(self.nside)), dtype=np.float32)
+                * output_units
+            )
             # sum, what if we have 2 sources on the same pixel?
             output_map[0, pix] += fluxes_I / pix_size * scaling_factor
         else:
@@ -258,8 +261,8 @@ def get_emission(
                     ),
                     ((r, p)),
                 )
-        if return_car:
-            return output_map
-        from pixell import reproject
+            if not return_car:
+                from pixell import reproject
 
-        return reproject.map2healpix(output_map, self.nside)
+                return reproject.map2healpix(output_map, self.nside)
+        return output_map
diff --git a/tests/test_catalog.py b/tests/test_catalog.py
@@ -97,11 +97,13 @@ def test_catalog_class_fluxes(test_catalog):
     catalog = PointSourceCatalog(test_catalog, nside=nside)
     freqs = np.exp(np.array([3, 4])) * u.GHz  # ~ 20 and ~ 55 GHz
     weights = np.array([1, 1], dtype=np.float64)
-    weights /= trapezoid(weights, x=freqs.to_value(u.GHz))
+    normalized_weights = utils.normalize_weights(utils.check_freq_input(freqs), weights)
     flux = catalog.get_fluxes(freqs, weights=weights)
     assert_allclose(flux[0], 3.7 * u.Jy)
     assert (
-        flux[1] == trapezoid(weights * np.array([6, 8]), x=freqs.to_value(u.GHz)) * u.Jy
+        flux[1]
+        == trapezoid(normalized_weights * np.array([6, 8]), x=freqs.to_value(u.GHz))
+        * u.Jy
     )