ULTRA L1C HELIO index maps (#2492)

* helio map func

Author: lacoak21

codecov.yml

@@ -15,3 +15,5 @@ coverage:
 ignore:
   - "**/conftest.py"
   - "**/test_spacecraft_pset.py"
+  - "**/test_helio_pset.py"
+  - "**/*make_helio_index_maps.py"

imap_processing/tests/ultra/unit/conftest.py

@@ -599,23 +599,23 @@ def random_spin_data():
 @pytest.fixture
 def mock_spacecraft_pointing_lookups():
     """Test lookup tables fixture."""
+    np.random.seed(42)
     pix = hp.nside2npix(128)  # reduced for testing
     steps = 2  # Reduced for testing
-    for_indices_by_spin_phase = np.random.choice(
-        [True, False], size=(pix, steps), p=[0.1, 0.9]
+    for_indices_by_spin_phase = xr.DataArray(
+        np.random.choice([True, False], size=(steps, pix), p=[0.1, 0.9]),
+        dims=("spin_phase_step", "pixel"),
     )
-    theta_vals = np.random.uniform(-60, 60, size=(pix, steps))
-    phi_vals = np.random.uniform(-60, 60, size=(pix, steps))
+    theta_vals = np.random.uniform(-60, 60, size=(steps, pix))
+    phi_vals = np.random.uniform(-60, 60, size=(steps, pix))
     # Ra and Dec pixel shape needs to be the default healpix pixel count
-    ra_and_dec = np.random.uniform(-80, 80, size=(pix, 2))
-    boundary_scale_factors = np.ones((pix, steps))
+    ra_and_dec = np.random.uniform(-80, 80, size=(steps, pix))
+    boundary_scale_factors = np.ones((steps, pix))
+
     with (
         mock.patch(
             "imap_processing.ultra.l1c.spacecraft_pset.get_spacecraft_pointing_lookup_tables"
         ) as mock_lookup,
-        mock.patch(
-            "imap_processing.ultra.l1c.helio_pset.get_spacecraft_pointing_lookup_tables"
-        ) as mock_lookup_helio,
     ):
         mock_lookup.return_value = (
             for_indices_by_spin_phase,
@@ -624,11 +624,60 @@ def mock_spacecraft_pointing_lookups():
             ra_and_dec,
             boundary_scale_factors,
         )
-        mock_lookup_helio.return_value = (
-            for_indices_by_spin_phase,
-            theta_vals,
-            phi_vals,
-            ra_and_dec,
-            boundary_scale_factors,
+        yield mock_lookup
+
+
+@pytest.fixture
+def mock_helio_pointing_lookups():
+    """Test lookup tables fixture returning an xarray Dataset."""
+    np.random.seed(42)
+    pix = hp.nside2npix(32)  # reduced for testing
+    steps = 2  # Reduced for testing
+    energy = 46
+
+    # Ra and Dec pixel shape needs to be the default healpix pixel count
+    ra_and_dec = np.random.uniform(-80, 80, size=(steps, pix))
+
+    index_map = np.random.choice([True, False], size=(steps, energy, pix), p=[0.1, 0.9])
+    index_map = index_map.astype(bool)
+    theta_map = np.random.uniform(-60, 60, size=(steps, energy, pix))
+    phi_map = np.random.uniform(-60, 60, size=(steps, energy, pix))
+    bsf_map = np.ones((steps, energy, pix))
+    with (
+        mock.patch(
+            "imap_processing.ultra.l1c.helio_pset.make_helio_index_maps_with_nominal_kernels"
+        ) as mock_lookup,
+    ):
+        ds = xr.Dataset(
+            data_vars={
+                "index": (
+                    ["spin_phase_step", "energy", "pixel"],
+                    index_map,
+                    {"long_name": "Pixel in FOV flag"},
+                ),
+                "theta": (
+                    ["spin_phase_step", "energy", "pixel"],
+                    theta_map,
+                    {"long_name": "Instrument theta angle", "units": "degrees"},
+                ),
+                "phi": (
+                    ["spin_phase_step", "energy", "pixel"],
+                    phi_map,
+                    {"long_name": "Instrument phi angle", "units": "degrees"},
+                ),
+                "bsf": (
+                    ["spin_phase_step", "energy", "pixel"],
+                    bsf_map,
+                    {"long_name": "Boundary scale factor", "units": "fractional"},
+                ),
+                "ra_and_dec": (["spin_phase_step", "pixel"], ra_and_dec),
+            },
+            coords={
+                "spin_phase_step": np.arange(steps),
+                "energy": np.arange(energy),
+                "pixel": np.arange(pix),
+            },
         )
+        mock_lookup.return_value = ds
+
         yield mock_lookup

imap_processing/tests/ultra/unit/test_helio_pset.py

@@ -0,0 +1,122 @@
+from unittest import mock
+
+import numpy as np
+import pandas as pd
+import pytest
+import xarray as xr
+
+from imap_processing import imap_module_directory
+from imap_processing.cdf.utils import load_cdf
+from imap_processing.tests.conftest import _download_external_data
+from imap_processing.ultra.constants import UltraConstants
+from imap_processing.ultra.l1c.helio_pset import calculate_helio_pset
+
+TEST_PATH = imap_module_directory / "tests" / "ultra" / "data" / "l1"
+
+
+@pytest.mark.skip(reason="Long running test for validation purposes.")
+def test_validate_exposure_time_and_sensitivities(
+    ancillary_files, deadtime_datasets, imap_ena_sim_metakernel
+):
+    """Validates exposure time and sensitivities for ebin 0."""
+    sens_filename = "SENS-IMAP_ULTRA_90-IMAP_DPS-HELIO-nside32-ebin0.csv"
+    exposure_filename = "Exposures-IMAP_ULTRA_90-IMAP_DPS-HELIO-nside32-ebin0.csv"
+    de_filename = "imap_ultra_l1b_90sensor-de_20000101-repoint00000_v000.cdf"
+    test_data = [
+        (sens_filename, "ultra/data/l1/"),
+        (exposure_filename, "ultra/data/l1/"),
+        (de_filename, "ultra/data/l1/"),
+    ]
+    _download_external_data(test_data)
+    l1b_de = TEST_PATH / de_filename
+    l1b_de = load_cdf(l1b_de)
+    sensitivities_ebin_0 = pd.read_csv(TEST_PATH / sens_filename)
+    exposure_factor_ebin_0 = pd.read_csv(TEST_PATH / exposure_filename)
+
+    test_deadtimes = (
+        pd.read_csv(TEST_PATH / "test_p0_ebin0_deadtimes.csv", header=None)
+        .to_numpy()
+        .squeeze()
+    )
+    npix = 12288  # nside 32
+    # Create a minimal dataset to pass to the function
+    dataset = xr.Dataset(
+        {
+            "spin_number": (["epoch"], np.array([1, 2, 3])),
+        }
+    )
+    dataset.attrs["Repointing"] = "repoint00000"
+
+    pointing_range_met = (472374890.0, 582378000.0)
+    # Create mock spin data that has 5525 nominal spins
+    # Create DataFrame
+    nspins = 5522
+    nominal_spin_seconds = 15.0
+    spin_data = pd.DataFrame(
+        {
+            "spin_start_met": np.linspace(
+                pointing_range_met[0], pointing_range_met[1], nspins
+            ),
+            "spin_period_sec": np.full(nspins, nominal_spin_seconds),
+            "spin_phase_valid": np.ones(nspins),
+            "spin_period_valid": np.ones(nspins),
+        }
+    )
+    with (
+        # Mock the pointing times
+        mock.patch(
+            "imap_processing.ultra.l1c.helio_pset.get_pointing_times_from_id",
+            return_value=pointing_range_met,
+        ),
+        mock.patch(
+            "imap_processing.ultra.l1c.ultra_l1c_pset_bins.ttj2000ns_to_met",
+            side_effect=lambda x: x,
+        ),
+        # Mock deadtimes to be all ones
+        mock.patch(
+            "imap_processing.ultra.l1c.ultra_l1c_pset_bins."
+            "get_deadtime_ratios_by_spin_phase",
+            return_value=xr.DataArray(test_deadtimes, dims="spin_phase_step"),
+        ),
+        # Mock spin data to match nominal spins in a pointing period
+        mock.patch(
+            "imap_processing.ultra.l1c.ultra_l1c_pset_bins.get_spin_data",
+            return_value=spin_data,
+        ),
+        # Mock background rates to be constant 0.1
+        mock.patch(
+            "imap_processing.ultra.l1c.helio_pset.get_spacecraft_background_rates",
+            return_value=np.ones((46, npix)),
+        ),
+        # Mock culling mask (no culling)
+        mock.patch("imap_processing.ultra.l1c.helio_pset.compute_culling_mask"),
+    ):
+        pset = calculate_helio_pset(
+            l1b_de,
+            dataset,
+            deadtime_datasets["rates"],
+            deadtime_datasets["params"],
+            "imap_ultra_l1c_90sensor-heliopset",
+            ancillary_files,
+            90,
+            UltraConstants.TOFXPH_SPECIES_GROUPS["proton"],
+        )
+
+    # Validate exposure times for ebin 0
+    exposure_times = pset["exposure_factor"][0, 0, :].values
+    expected_exposure_times = exposure_factor_ebin_0["P0"].to_numpy()
+    np.testing.assert_allclose(
+        exposure_times,
+        expected_exposure_times,
+        atol=95,  # TODO This is due to the helio index map differences
+        err_msg="Exposure times do not match expected values for ebin 0.",
+    )
+    # Validate sensitivities for ebin 0
+    sensitivity = pset["sensitivity"][0, :].values
+    expected_sensitivity = sensitivities_ebin_0["Sensitivity (cm2)"].to_numpy()
+    np.testing.assert_allclose(
+        sensitivity,
+        expected_sensitivity,
+        atol=0.0006,  # TODO This is due to the helio index map differences
+        err_msg="Sensitivities times do not match expected values for ebin 0.",
+    )

imap_processing/tests/ultra/unit/test_l1c_lookup_utils.py

@@ -1,8 +1,10 @@
 import astropy_healpix.healpy as hp
 import numpy as np
 import pytest
+import xarray as xr
 
 from imap_processing.ultra.l1c.l1c_lookup_utils import (
+    calculate_fwhm_spun_scattering,
     get_scattering_thresholds_for_energy,
     get_spacecraft_pointing_lookup_tables,
     get_static_deadtime_ratios,
@@ -25,10 +27,10 @@ def test_get_spacecraft_pointing_lookup_tables(ancillary_files):
     # Test shapes
     # There should be 498 spin phase steps. In the real test files there will be 15000
     cols = 498
-    assert for_indices_by_spin_phase.shape == (npix, cols)
-    assert theta_vals.shape == (npix, cols)
-    assert phi_vals.shape == (npix, cols)
-    assert ra_and_dec.shape == (npix, 2)
+    assert for_indices_by_spin_phase.shape == (cols, npix)
+    assert theta_vals.shape == (cols, npix)
+    assert phi_vals.shape == (cols, npix)
+    assert ra_and_dec.shape == (2, npix)
 
     # Value tests
     assert for_indices_by_spin_phase.dtype == bool
@@ -111,3 +113,47 @@ def test_get_static_deadtime_ratios(ancillary_files):
     np.testing.assert_array_equal(dt_ratio.shape, (721,))
     # Test the values
     assert np.all((dt_ratio >= 0.0) & (dt_ratio <= 1.0))
+
+
+def test_calculate_fwhm_spun_scattering(ancillary_files):
+    """Test calculate_fwhm_spun_scattering function."""
+    # Make array with ones (we are only testing the shape here)
+    for_pixels = np.ones((50, 10))
+    theta_vals = np.ones((50, 10)) * 20  # All theta values are 20
+    phi_vals = np.ones((50, 5)) * 15  # All phi
+    with pytest.raises(ValueError, match="Shape mismatch"):
+        calculate_fwhm_spun_scattering(
+            for_pixels, theta_vals, phi_vals, ancillary_files, 45
+        )
+
+
+@pytest.mark.external_test_data
+def test_calculate_fwhm_spun_scattering_reject(ancillary_files):
+    """Test calculate_fwhm_spun_scattering function."""
+    nside = 8
+    pix = hp.nside2npix(nside)
+    steps = 5  # Reduced for testing
+    energy_dim = 46
+    np.random.seed(42)
+    mock_theta = np.random.uniform(-60, 60, (steps, energy_dim, pix))
+    mock_phi = np.random.uniform(-60, 60, (steps, energy_dim, pix))
+    for_pixels = xr.DataArray(
+        np.zeros((steps, energy_dim, pix)).astype(bool),
+        dims=("spin_phase_step", "energy", "pixel"),
+    )
+    # Simulate first 100 pixels are in the FOR for all spin phases
+    inside_inds = 100
+    for_pixels[:, :, :inside_inds] = True
+    valid_spun_pixels, fwhm_theta, fwhm_phi, thresholds = (
+        calculate_fwhm_spun_scattering(
+            for_pixels,
+            mock_theta,
+            mock_phi,
+            ancillary_files,
+            45,
+            reject_scattering=True,
+        )
+    )
+    assert valid_spun_pixels.shape == (steps, energy_dim, pix)
+    # Check that some pixels are rejected
+    assert not np.array_equal(valid_spun_pixels, for_pixels)