ULTRA l1c counts nside update

imap_processing/cdf/config/imap_ultra_l1c_variable_attrs.yaml

@@ -136,8 +136,11 @@ scatter_threshold:
 counts:
   <<: *default
   CATDESC: Counts for a spin.
+  VAR_NOTES: Counts healpix maps are sampled at finer resolution to help maintain the pointing accuracy for each event.
+    Since count maps are a binned integral quantity, they necessarily require a non-spun approach per Pointing, unlike
+    exposure time and sensitivities.
   DEPEND_1: energy_bin_geometric_mean
-  DEPEND_2: pixel_index
+  DEPEND_2: counts_pixel_index
   FIELDNAM: counts
   LABLAXIS: counts
   # TODO: come back to format
@@ -244,6 +247,18 @@ pixel_index:
   VALIDMAX: *max_int
   VAR_TYPE: support_data
 
+counts_pixel_index:
+  FILLVAL: *min_int
+  CATDESC: Counts variable HEALPix pixel index.
+  DISPLAY_TYPE: no_plot
+  FIELDNAM: counts_pixel_index
+  FORMAT: I12
+  LABLAXIS: Counts Pixel Index
+  UNITS: " "
+  VALIDMIN: *min_int
+  VALIDMAX: *max_int
+  VAR_TYPE: support_data
+
 spin_phase_step:
   FILLVAL: *min_int
   CATDESC: Spin phase step index (1ms resolution).

imap_processing/ena_maps/ena_maps.py

@@ -558,6 +558,8 @@ def __init__(
             np.stack((azimuth_pixel_center, elevation_pixel_center), axis=-1),
             dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value],
         )
+        # downsample counts variable to match the nside of the pointing set
+        self.downsample_counts()
 
     @property
     def num_points(self) -> int:
@@ -600,6 +602,73 @@ def __repr__(self) -> str:
             f"num_points={self.num_points})"
         )
 
+    def downsample_counts(self) -> None:
+        """
+        Downsample the counts variable to match the pset nside.
+
+        Counts at l1c are sampled at a finer resolution to help maintain the
+        pointing accuracy for each event. Since count maps are a binned integral
+        quantity, they necessarily require a non-spun approach per Pointing, unlike
+        exposure time and sensitivities. We need to downsample the counts from the
+        nside of the input pset counts variable (e.g. 128) to the nside of the pset.
+        """
+        pset_data = self.data
+        counts_n_pix = pset_data.sizes["counts_pixel_index"]
+        pset_n_pix = hp.nside2npix(self.nside)
+        if counts_n_pix != pset_n_pix:
+            # Raise an error if the nside the counts were sampled at is lower than the
+            # nside of the output map. We never want counts to be upsampled.
+            if counts_n_pix < pset_n_pix:
+                raise ValueError(
+                    f"Counts in the input PSET are sampled at nside "
+                    f"{hp.npix2nside(counts_n_pix)}, and the pset is {self.nside}. "
+                    f"This would require upsampling the counts, which we do not want."
+                )
+            counts_nside = hp.npix2nside(counts_n_pix)
+            n_energy_bins = pset_data.sizes["energy_bin_geometric_mean"]
+            order_diff = int(np.log2(counts_nside // self.nside))
+            counts = pset_data["counts"].values[
+                0
+            ]  # shape: (n_energy_bins, counts_n_pix)
+            # Get counts in nested ordering. In nested ordering, the
+            # pixels that need to be binned together to go from the counts nside to
+            # the pset nside are contiguous in the array.
+            if not self.nested:
+                counts_n = counts[
+                    :, hp.ring2nest(counts_nside, np.arange(counts_n_pix))
+                ]
+            else:
+                counts_n = counts
+
+            # reshape the counts by the amount pixels to bin together which is
+            # 4**order_diff because each step in order multiplies the pixel count
+            # by 4
+            # Shape: (n_energy_bins, pset_n_pix, 4**order_diff) ->
+            # (n_energy_bins, pset_n_pix)
+            binned_counts_n = counts_n.reshape(
+                (n_energy_bins, pset_n_pix, 4**order_diff)
+            ).sum(axis=-1)
+
+            if not self.nested:
+                # convert back to ring ordering if necessary and store in the
+                # downsampled counts array
+                binned_counts_n = binned_counts_n[
+                    :, hp.nest2ring(self.nside, np.arange(pset_n_pix))
+                ]
+
+            self.data["counts"] = xr.DataArray(
+                binned_counts_n[np.newaxis, :, :],
+                dims=(
+                    *self.data["counts"].dims[:-1],
+                    CoordNames.HEALPIX_INDEX.value,
+                ),
+            )
+        else:
+            # Update the counts variable with the correct dims
+            self.data["counts"] = self.data["counts"].rename(
+                {CoordNames.COUNTS_HEALPIX_INDEX.value: CoordNames.HEALPIX_INDEX.value}
+            )
+
 
 class LoHiBasePointingSet(PointingSet):
     """

imap_processing/ena_maps/utils/coordinates.py

@@ -12,6 +12,7 @@ class CoordNames(Enum):
     ENERGY_ULTRA_L1C = "energy_bin_geometric_mean"
     ENERGY_L2 = "energy"
     HEALPIX_INDEX = "pixel_index"
+    COUNTS_HEALPIX_INDEX = "counts_pixel_index"
 
     # The names of the az/el angular coordinates may differ between L1C and L2 data
     AZIMUTH_L1C = "longitude"

imap_processing/tests/ena_maps/test_ena_maps.py

@@ -8,6 +8,7 @@
 from copy import deepcopy
 from pathlib import Path
 from unittest import mock
+from unittest.mock import patch
 
 import astropy_healpix.healpy as hp
 import numpy as np
@@ -101,10 +102,14 @@ def test_init_cdf(
 
         cdf_filepath = write_cdf(ultra_pset, istp=False)
 
-        ultra_pset_from_dataset = ena_maps.UltraPointingSet(ultra_pset)
-
-        ultra_pset_from_str = ena_maps.UltraPointingSet(cdf_filepath)
-        ultra_pset_from_path = ena_maps.UltraPointingSet(Path(cdf_filepath))
+        # Mock the downsample_counts method to avoid dimension bugs
+        # since this is a dummy cdf, the dimensions are not present.
+        with patch.object(
+            ena_maps.UltraPointingSet, "downsample_counts", lambda self: None
+        ):
+            ultra_pset_from_str = ena_maps.UltraPointingSet(cdf_filepath)
+            ultra_pset_from_path = ena_maps.UltraPointingSet(Path(cdf_filepath))
+            ultra_pset_from_dataset = ena_maps.UltraPointingSet(ultra_pset)
 
         np.testing.assert_allclose(
             ultra_pset_from_dataset.data["counts"].values,
@@ -118,7 +123,6 @@ def test_init_cdf(
             rtol=1e-6,
         )
 
-    @pytest.mark.usefixtures("_setup_ultra_l1c_pset_products")
     @pytest.mark.usefixtures("_setup_ultra_l1c_pset_products")
     def test_different_spacing_raises_error(self):
         """Test that different spaced az/el from the L1C dataset raises ValueError"""
@@ -135,6 +139,23 @@ def test_different_spacing_raises_error(self):
                 spice_reference_frame=geometry.SpiceFrame.IMAP_DPS,
             )
 
+    def test_downsample_counts(self):
+        ultra_pset = self.l1c_pset_products[0]
+
+        # First check that counts are at a finer resolution than exposure factor
+        counts_npix_before = ultra_pset["counts"].shape[-1]
+        assert counts_npix_before != ultra_pset["exposure_factor"].shape[-1]
+        pset = ena_maps.UltraPointingSet(ultra_pset)
+
+        # Verify counts are now at the same resolution as pset
+        counts_nside_after = hp.npix2nside(pset.data["counts"].shape[-1])
+        assert counts_nside_after == pset.nside
+
+        # Verify counts are the same after downsampling.
+        np.testing.assert_allclose(
+            pset.data["counts"].values.sum(), ultra_pset["counts"].values.sum()
+        )
+
 
 @pytest.fixture(scope="module")
 def hi_pset_cdf_path(imap_tests_path):

imap_processing/tests/ultra/mock_data.py

@@ -193,7 +193,8 @@ def get_binomial_counts(distance_scaling, lat_bin, central_lat_bin):
 # TODO: Add ability to mock with/without energy dim to exposure_factor
 # The Helio frame L1C will have the energy dimension, but the spacecraft frame will not.
 def mock_l1c_pset_product_healpix(
-    nside: int = DEFAULT_HEALPIX_NSIDE_L1C,
+    nside: int = 32,
+    counts_nside: int = 128,
     stripe_center_lat: int = 0,
     width_scale: float = 10.0,
     counts_scaling_params: tuple[int, float] = (100, 0.01),
@@ -264,19 +265,19 @@ def mock_l1c_pset_product_healpix(
     energy_bin_delta = np.diff(energy_intervals, axis=1).squeeze()
     num_energy_bins = len(energy_bin_midpoints)
     npix = hp.nside2npix(nside)
-    counts = np.zeros(npix)
-    exposure_time = np.zeros(npix)
-
+    counts_npix = hp.nside2npix(counts_nside)
+    # counts are binned at a higher resolution than the other variables. See L1c
+    # code for more details.
     # Get latitude for each healpix pixel
-    pix_indices = np.arange(npix)
-    lon_pix, lat_pix = hp.pix2ang(nside, pix_indices, lonlat=True)
-
-    counts = np.zeros(shape=(num_energy_bins, npix))
+    counts_pix_indices = np.arange(counts_npix)
+    counts_lon_pix, counts_lat_pix = hp.pix2ang(
+        counts_nside, counts_pix_indices, lonlat=True
+    )
 
     # Calculate probability based on distance from target latitude
-    lat_diff = np.abs(lat_pix - stripe_center_lat)
+    counts_lat_diff = np.abs(counts_lat_pix - stripe_center_lat)
     prob_scaling_factor = counts_scaling_params[1] * np.exp(
-        -(lat_diff**2) / (2 * width_scale**2)
+        -(counts_lat_diff**2) / (2 * width_scale**2)
     )
     # Generate counts using binomial distribution
     rng = np.random.default_rng(seed=42)
@@ -286,6 +287,12 @@ def mock_l1c_pset_product_healpix(
             for _ in range(num_energy_bins)
         ]
     )
+    # Get latitude for each healpix pixel
+    pix_indices = np.arange(npix)
+    lon_pix, lat_pix = hp.pix2ang(nside, pix_indices, lonlat=True)
+
+    # Calculate probability based on distance from target latitude
+    lat_diff = np.abs(lat_pix - stripe_center_lat)
 
     # Generate exposure times using gaussian distribution, but wider
     prob_scaling_factor_exptime = counts_scaling_params[1] * np.exp(
@@ -317,7 +324,7 @@ def mock_l1c_pset_product_healpix(
     counts = counts.astype(int)
     # add an epoch dimension
     counts = np.expand_dims(counts, axis=0)
-    ones_ds = np.ones_like(counts)[0]  # pointing independent
+    ones_ds = np.ones((num_energy_bins, npix))  # pointing independent
     sensitivity = ones_ds
     geometric_function = ones_ds
     efficiency = ones_ds
@@ -338,7 +345,7 @@ def mock_l1c_pset_product_healpix(
                 [
                     CoordNames.TIME.value,
                     CoordNames.ENERGY_ULTRA_L1C.value,
-                    CoordNames.HEALPIX_INDEX.value,
+                    CoordNames.COUNTS_HEALPIX_INDEX.value,
                 ],
                 counts,
             ),
@@ -348,7 +355,7 @@ def mock_l1c_pset_product_healpix(
                     CoordNames.ENERGY_ULTRA_L1C.value,
                     CoordNames.HEALPIX_INDEX.value,
                 ],
-                np.full_like(counts, 0.05, dtype=float),
+                np.full((1, num_energy_bins, npix), 0.05, dtype=float),
             ),
             "exposure_factor": (
                 exposure_dims,  # special case: optionally energy dependent exposure

imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py

@@ -118,13 +118,9 @@ def test_get_spacecraft_histogram(test_data):
     energy_bin_edges, _, _ = build_energy_bins()
     subset_energy_bin_edges = energy_bin_edges[:3]
 
-    hist, latitude, longitude, n_pix = get_spacecraft_histogram(
-        v, energy, subset_energy_bin_edges, nside=1
-    )
+    hist, n_pix = get_spacecraft_histogram(v, energy, subset_energy_bin_edges, nside=1)
     assert hist.shape == (len(subset_energy_bin_edges), hp.nside2npix(1))
     assert n_pix == hp.nside2npix(1)
-    assert latitude.shape == (n_pix,)
-    assert longitude.shape == (n_pix,)
 
     # Spot check that 2 counts are in the second energy bin
     assert np.sum(hist[2, :]) == 2
@@ -135,14 +131,10 @@ def test_get_spacecraft_histogram(test_data):
         (2.5, 4.137),
         (3.385, 5.057),
     ]
-    hist, latitude, longitude, n_pix = get_spacecraft_histogram(
-        v, energy, overlapping_bins, nside=1
-    )
+    hist, n_pix = get_spacecraft_histogram(v, energy, overlapping_bins, nside=1)
     # Spot check that 3 counts are in the third energy bin
     assert np.sum(hist[2, :]) == 3
     assert n_pix == hp.nside2npix(1)
-    assert latitude.shape == (n_pix,)
-    assert longitude.shape == (n_pix,)
 
 
 def mock_imap_state(time, ref_frame):

imap_processing/ultra/constants.py

@@ -92,6 +92,8 @@ class UltraConstants:
         300.0,
         1e5,
     ]
+    # Counts at l1c are sampled at a finer resolution.
+    L1C_COUNTS_NSIDE = 128
 
     PSET_ENERGY_BIN_EDGES: ClassVar[list] = [
         3.0,

imap_processing/ultra/l1c/helio_pset.py

@@ -2,6 +2,7 @@
 
 import logging
 
+import astropy_healpix.healpy as hp
 import numpy as np
 import xarray as xr
 
@@ -156,17 +157,24 @@ def calculate_helio_pset(
         )
     )
 
-    counts, latitude, longitude, n_pix = get_spacecraft_histogram(
+    counts, counts_n_pix = get_spacecraft_histogram(
         vhat_dps_helio,
         species_dataset["energy_heliosphere"].values,
         intervals,
-        nside=nside,
+        nside=UltraConstants.L1C_COUNTS_NSIDE,
     )
+    n_pix = hp.nside2npix(nside)
     helio_pset_quality_flags = np.full(
         n_pix, ImapPSETUltraFlags.NONE.value, dtype=np.uint16
     )
+    counts_healpix = np.arange(counts_n_pix)
+    # Determine nside for non "counts" variables from the lookup table
     healpix = np.arange(n_pix)
 
+    # Calculate the corresponding longitude (az) latitude (el)
+    # center coordinates
+    longitude, latitude = hp.pix2ang(nside, healpix, lonlat=True)
+
     logger.info("Calculating spacecraft exposure times with deadtime correction.")
     exposure_time, deadtime_ratios = get_spacecraft_exposure_times(
         rates_dataset,
@@ -236,6 +244,7 @@ def calculate_helio_pset(
     pset_dict["background_rates"] = background_rates[np.newaxis, ...]
     pset_dict["exposure_factor"] = exposure_time[np.newaxis, ...]
     pset_dict["pixel_index"] = healpix
+    pset_dict["counts_pixel_index"] = counts_healpix
     pset_dict["energy_bin_delta"] = np.diff(intervals, axis=1).squeeze()[
         np.newaxis, ...
     ]

imap_processing/ultra/l1c/spacecraft_pset.py

@@ -139,15 +139,22 @@ def calculate_spacecraft_pset(
             reject_scattering,
         )
     )
-    # Determine nside from the lookup table
-    nside = hp.npix2nside(for_indices_by_spin_phase.sizes["pixel"])
-    counts, latitude, longitude, n_pix = get_spacecraft_histogram(
+    counts, counts_n_pix = get_spacecraft_histogram(
         vhat_dps_spacecraft,
         species_dataset["energy_spacecraft"].values,
         intervals,
-        nside=nside,
+        nside=UltraConstants.L1C_COUNTS_NSIDE,
     )
+    counts_healpix = np.arange(counts_n_pix)
+    # Determine nside for non "counts" variables from the lookup table
+    n_pix = for_indices_by_spin_phase.sizes["pixel"]
+    nside = hp.npix2nside(n_pix)
     healpix = np.arange(n_pix)
+
+    # Calculate the corresponding longitude (az) latitude (el)
+    # center coordinates
+    longitude, latitude = hp.pix2ang(nside, healpix, lonlat=True)
+
     # Get the start and stop times of the pointing period
     repoint_id = species_dataset.attrs.get("Repointing", None)
     if repoint_id is None:
@@ -225,6 +232,7 @@ def calculate_spacecraft_pset(
     pset_dict["background_rates"] = background_rates[np.newaxis, ...]
     pset_dict["exposure_factor"] = exposure_pointing[np.newaxis, ...]
     pset_dict["pixel_index"] = healpix
+    pset_dict["counts_pixel_index"] = counts_healpix
     pset_dict["energy_bin_delta"] = np.diff(intervals, axis=1).squeeze()[
         np.newaxis, ...
     ]