BUG - Lo Hist Rate L1B - Swap ESA Step and Az dims, intialize exposure factor to 1 instead of 0 (#2466)

* swapped esa step and az dims, inititalize exposure factor to 1

* minor xarray usage fixes

Author: sdhoyt

imap_processing/lo/l1b/lo_l1b.py

@@ -1177,12 +1177,12 @@ def initialize_l1b_histrates(
     """
     l1b_histrates = xr.Dataset(
         coords={
-            "epoch": l1a_hist["epoch"],
+            "epoch": xr.DataArray(l1a_hist["epoch"].values, dims=["epoch"]),
+            "esa_step": l1a_hist["esa_step"],
             "spin_bin_6": xr.DataArray(
                 l1a_hist["azimuth_6"].values,
                 dims=["spin_bin_6"],
             ),
-            "esa_step": l1a_hist["esa_step"],
         },
         attrs=attr_mgr_l1b.get_global_attributes(logical_source),
     )
@@ -1195,13 +1195,13 @@ def initialize_l1b_histrates(
     # Copy over fields from L1A DE that will not change in L1B processing
     l1b_histrates["h_counts"] = xr.DataArray(
         l1a_hist["hydrogen"].values,
-        dims=["epoch", "spin_bin_6", "esa_step"],
+        dims=["epoch", "esa_step", "spin_bin_6"],
         # TODO: Add hydrogen to YAML file
         # attrs=attr_mgr.get_variable_attributes("hydrogen"),
     )
     l1b_histrates["o_counts"] = xr.DataArray(
         l1a_hist["oxygen"].values,
-        dims=["epoch", "spin_bin_6", "esa_step"],
+        dims=["epoch", "esa_step", "spin_bin_6"],
         # TODO: Add oxygen to YAML file
         # attrs=attr_mgr.get_variable_attributes("oxygen"),
     )
@@ -1254,8 +1254,12 @@ def resweep_histogram_data(
     o_counts_reswept = np.zeros_like(l1b_histrates["o_counts"].values)
 
     # Get the number of azimuth bins from the l1b_histrates dataset
-    num_azimuth = l1b_histrates["h_counts"].shape[1]
-    exposure_factor = np.zeros((len(epochs), num_azimuth, 7), dtype=int)
+    num_azimuth = l1b_histrates.sizes["spin_bin_6"]
+    # initialize exposure factor to 1 as this will be used to scale (multiply)
+    # the exposure time later
+    exposure_factor = np.full(
+        (len(epochs), l1b_histrates.sizes["esa_step"], num_azimuth), 1, dtype=int
+    )
 
     for epoch_idx, epoch in enumerate(epoch_utc):
         # Get only the date portion of the epoch string for comparison with the
@@ -1289,7 +1293,6 @@ def resweep_histogram_data(
             logger.warning(f"No LUT entries found for table index {lut_table_idx}")
             h_counts_reswept[epoch_idx] = l1b_histrates["h_counts"].values[epoch_idx]
             o_counts_reswept[epoch_idx] = l1b_histrates["o_counts"].values[epoch_idx]
-            exposure_factor[epoch_idx] = 1
             continue
 
         # Sort the LUT entries by E-Step_Idx to ensure correct mapping order
@@ -1309,8 +1312,8 @@ def resweep_histogram_data(
         # TODO: Change all instances of azimuth to spin bin
         # Resweep the counts for each spin bin using the energy step mapping
         for az_idx in range(num_azimuth):
-            h_original = l1b_histrates["h_counts"].values[epoch_idx, az_idx, :]
-            o_original = l1b_histrates["o_counts"].values[epoch_idx, az_idx, :]
+            h_original = l1b_histrates["h_counts"].values[epoch_idx, :, az_idx]
+            o_original = l1b_histrates["o_counts"].values[epoch_idx, :, az_idx]
 
             # Loop through the original ESA step indices and map to the true ESA steps
             for orig_idx, true_esa_step in energy_step_mapping.items():
@@ -1319,16 +1322,17 @@ def resweep_histogram_data(
                     # Resweep the counts into the true ESA step
                     # (convert to 0-based index)
                     reswept_idx = true_esa_step - 1
-                    h_counts_reswept[epoch_idx, az_idx, reswept_idx] += h_original[
+                    h_counts_reswept[epoch_idx, reswept_idx, az_idx] += h_original[
                         orig_idx
                     ]
-                    o_counts_reswept[epoch_idx, az_idx, reswept_idx] += o_original[
+                    o_counts_reswept[epoch_idx, reswept_idx, az_idx] += o_original[
                         orig_idx
                     ]
                     # If a reswept was needed for this index, increment the exposure
                     # factor to so the exposure time can be scaled accordingly
                     if orig_idx != reswept_idx:
-                        exposure_factor[epoch_idx, az_idx, reswept_idx] += 1
+                        exposure_factor[epoch_idx, reswept_idx, az_idx] += 1
+
                 else:
                     logger.warning(
                         f"Original ESA index {orig_idx} or "
@@ -1386,8 +1390,8 @@ def calculate_histogram_rates(
 
     h_rates = np.zeros_like(h_counts, dtype=float)
     o_rates = np.zeros_like(o_counts, dtype=float)
-    num_azimuth = h_counts.shape[1]
-    exposure_times = np.zeros((len(epochs), num_azimuth, 7), dtype=float)
+    num_azimuth = h_counts.shape[2]
+    exposure_times = np.zeros((len(epochs), 7, num_azimuth), dtype=float)
 
     # Calculate rates for each epoch
     for epoch_idx, epoch in enumerate(epochs):
@@ -1409,7 +1413,6 @@ def calculate_histogram_rates(
         base_exposure_time = (
             4 * avg_spin_durations_per_cycle.values[spin_cycle_idx] / 60
         )
-
         # Scale the exposure time by the exposure factor from resweeping
         scaled_exposure = base_exposure_time * exposure_factor[epoch_idx, ...]
         # Avoid division by zero by setting zero exposure times to NaN
@@ -1420,7 +1423,7 @@ def calculate_histogram_rates(
 
     l1b_histrates["exposure_time"] = xr.DataArray(
         exposure_times,
-        dims=["epoch", "spin_bin_6", "esa_step"],
+        dims=["epoch", "esa_step", "spin_bin_6"],
     )
     l1b_histrates["h_rates"] = xr.DataArray(
         h_rates,

imap_processing/tests/lo/test_lo_l1b.py

@@ -96,13 +96,13 @@ def l1b_histrates():
     )
     l1b_histrates = xr.Dataset(
         {
-            "h_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
-            "o_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
+            "h_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
+            "o_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
         },
         coords={
             "epoch": epoch_date,
-            "azimuth_6": np.arange(60),
             "esa_step": np.arange(1, 8),
+            "azimuth_6": np.arange(60),
         },
     )
 
@@ -114,13 +114,13 @@ def l1a_hist():
     epoch_date = et_to_ttj2000ns(str_to_et(["2025-04-15T02:00:00"]))
     l1a_hist = xr.Dataset(
         {
-            "hydrogen": (("epoch", "azimuth_6", "esa_step"), np.zeros((1, 60, 7))),
-            "oxygen": (("epoch", "azimuth_6", "esa_step"), np.zeros((1, 60, 7))),
+            "hydrogen": (("epoch", "esa_step", "azimuth_6"), np.zeros((1, 7, 60))),
+            "oxygen": (("epoch", "esa_step", "azimuth_6"), np.zeros((1, 7, 60))),
         },
         coords={
             "epoch": epoch_date,
-            "azimuth_6": np.arange(60),
             "esa_step": np.arange(1, 8),
+            "azimuth_6": np.arange(60),
         },
     )
     return l1a_hist
@@ -787,37 +787,39 @@ def test_resweep_histogram_success(anc_dependencies):
     epoch_date = et_to_ttj2000ns(
         str_to_et(["2025-04-15T02:00:00", "2025-04-15T03:00:00"])
     )
-    l1b_de = xr.Dataset(
+    l1b_histrate = xr.Dataset(
         {
-            "h_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
-            "o_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
+            "h_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
+            "o_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
         },
         coords={
             "epoch": epoch_date,
-            "azimuth_6": np.arange(60),
             "esa_step": np.arange(1, 8),
+            "spin_bin_6": np.arange(60),
         },
     )
-    exposure_factor_expected = np.zeros((2, 60, 7))
-    exposure_factor_expected[:, :, 0] = 1
+    exposure_factor_expected = np.full((2, 7, 60), 1)
+    exposure_factor_expected[:, 0, :] = 2
 
-    l1b_de.h_counts[0, 0, 0] = 5
-    l1b_de.h_counts[0, 0, 1] = 10
-    l1b_de.h_counts[0, 0, 2] = 2
+    l1b_histrate.h_counts[0, 0, 0] = 5
+    l1b_histrate.h_counts[0, 1, 0] = 10
+    l1b_histrate.h_counts[0, 2, 0] = 2
 
-    l1b_de.o_counts[1, 0, 0] = 2
-    l1b_de.o_counts[1, 0, 1] = 3
-    l1b_de.o_counts[1, 0, 2] = 4
+    l1b_histrate.o_counts[1, 0, 0] = 2
+    l1b_histrate.o_counts[1, 1, 0] = 3
+    l1b_histrate.o_counts[1, 2, 0] = 4
 
-    l1b_histrates, exposure_factor = resweep_histogram_data(l1b_de, anc_dependencies)
+    l1b_histrates, exposure_factor = resweep_histogram_data(
+        l1b_histrate, anc_dependencies
+    )
 
     assert l1b_histrates.h_counts[0, 0, 0] == 15
-    assert l1b_histrates.h_counts[0, 0, 1] == 0
-    assert l1b_histrates.h_counts[0, 0, 2] == 2
+    assert l1b_histrates.h_counts[0, 1, 0] == 0
+    assert l1b_histrates.h_counts[0, 2, 0] == 2
 
     assert l1b_histrates.o_counts[1, 0, 0] == 5
-    assert l1b_histrates.o_counts[1, 0, 1] == 0
-    assert l1b_histrates.o_counts[1, 0, 2] == 4
+    assert l1b_histrates.o_counts[1, 1, 0] == 0
+    assert l1b_histrates.o_counts[1, 2, 0] == 4
 
     assert np.array_equal(exposure_factor, exposure_factor_expected)
 
@@ -827,43 +829,43 @@ def test_resweep_histogram_no_date(anc_dependencies):
     epoch_date = et_to_ttj2000ns(
         str_to_et(["2025-04-25T02:00:00", "2025-04-25T03:00:00"])
     )
-    l1b_de = xr.Dataset(
+    l1b_histrate = xr.Dataset(
         {
-            "h_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
-            "o_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
+            "h_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
+            "o_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
         },
         coords={
             "epoch": epoch_date,
-            "azimuth_6": np.arange(60),
             "esa_step": np.arange(1, 8),
+            "spin_bin_6": np.arange(60),
         },
     )
 
-    l1b_de.h_counts[0, 0, 0] = 5
-    l1b_de.h_counts[0, 0, 1] = 10
-    l1b_de.h_counts[0, 0, 2] = 2
+    l1b_histrate.h_counts[0, 0, 0] = 5
+    l1b_histrate.h_counts[0, 1, 0] = 10
+    l1b_histrate.h_counts[0, 2, 0] = 2
 
     with pytest.raises(
         ValueError,
         match="No sweep table entry found for date "
         "2025-04-25T02:00:00.000 at epoch idx 0",
     ):
-        resweep_histogram_data(l1b_de, anc_dependencies)
+        resweep_histogram_data(l1b_histrate, anc_dependencies)
 
 
 def test_resweep_histogram_multiple_lut(anc_dependencies):
     epoch_date = et_to_ttj2000ns(
         str_to_et(["2025-04-16T02:00:00", "2025-04-16T03:00:00"])
     )
-    l1b_de = xr.Dataset(
+    l1b_histrate = xr.Dataset(
         {
-            "h_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
-            "o_counts": (("epoch", "azimuth_6", "esa_step"), np.zeros((2, 60, 7))),
+            "h_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
+            "o_counts": (("epoch", "esa_step", "azimuth_6"), np.zeros((2, 7, 60))),
         },
         coords={
             "epoch": epoch_date,
-            "azimuth_6": np.arange(60),
             "esa_step": np.arange(1, 8),
+            "spin_bin_6": np.arange(60),
         },
     )
 
@@ -872,7 +874,7 @@ def test_resweep_histogram_multiple_lut(anc_dependencies):
         match=f"Expected exactly 1 unique LUT_table "
         f"value for date 2025-04-16, but found 2:{[1, 2]}",
     ):
-        resweep_histogram_data(l1b_de, anc_dependencies)
+        resweep_histogram_data(l1b_histrate, anc_dependencies)
 
 
 def test_calculate_histogram_rates(l1b_histrates):
@@ -889,32 +891,32 @@ def test_calculate_histogram_rates(l1b_histrates):
         ]
     )
     avg_spin_durations_per_cycle = xr.DataArray([30, 15])
-    exposure_factor = np.zeros((2, 60, 7))
+    exposure_factor = np.zeros((2, 7, 60))
     exposure_factor[0, 0, 0] = 1
     l1b_histrates.h_counts[0, 0, 0] = 30
-    l1b_histrates.h_counts[0, 0, 1] = 10
-    l1b_histrates.h_counts[0, 0, 2] = 2
+    l1b_histrates.h_counts[0, 1, 0] = 10
+    l1b_histrates.h_counts[0, 2, 0] = 2
     l1b_histrates.h_counts[1, 0, 0] = 15
-    l1b_histrates.h_counts[1, 0, 1] = 30
-    l1b_histrates.h_counts[1, 0, 2] = 45
+    l1b_histrates.h_counts[1, 1, 0] = 30
+    l1b_histrates.h_counts[1, 2, 0] = 45
 
     l1b_histrates.o_counts[0, 0, 0] = 100
-    l1b_histrates.o_counts[0, 0, 1] = 50
-    l1b_histrates.o_counts[0, 0, 2] = 25
+    l1b_histrates.o_counts[0, 1, 0] = 50
+    l1b_histrates.o_counts[0, 2, 0] = 25
     l1b_histrates.o_counts[1, 0, 0] = 2
-    l1b_histrates.o_counts[1, 0, 1] = 3
-    l1b_histrates.o_counts[1, 0, 2] = 4
+    l1b_histrates.o_counts[1, 1, 0] = 3
+    l1b_histrates.o_counts[1, 2, 0] = 4
 
     l1b_histrate = calculate_histogram_rates(
         l1b_histrates, acq_start, acq_end, avg_spin_durations_per_cycle, exposure_factor
     )
 
     hist_rates_h_epoch_0 = l1b_histrate["h_rates"]
     hist_rates_h_epoch_0[0, :, :] = hist_rates_h_epoch_0[0, :, :] / 2
-    hist_rates_h_epoch_0[0, 0, :] = hist_rates_h_epoch_0[0, 0, :] / 2
+    hist_rates_h_epoch_0[0, :, 0] = hist_rates_h_epoch_0[0, :, 0] / 2
     hist_rates_o_epoch_0 = l1b_histrate["o_rates"]
     hist_rates_o_epoch_0[0, :, :] = hist_rates_o_epoch_0[0, :, :] / 2
-    hist_rates_o_epoch_0[0, 0, :] = hist_rates_o_epoch_0[0, 0, :] / 2
+    hist_rates_o_epoch_0[0, :, 0] = hist_rates_o_epoch_0[0, :, 0] / 2
 
     np.testing.assert_array_equal(
         l1b_histrate["h_rates"][0, :, :], hist_rates_h_epoch_0[0, :, :]
@@ -944,13 +946,13 @@ def test_calculate_histogram_rates_no_interval_found(l1b_histrates):
         ]
     )
     avg_spin_durations_per_cycle = xr.DataArray([30, 15])
-    exposure_factor = np.zeros((2, 60, 7))
+    exposure_factor = np.zeros((2, 7, 60))
     l1b_histrate = calculate_histogram_rates(
         l1b_histrates, acq_start, acq_end, avg_spin_durations_per_cycle, exposure_factor
     )
 
-    np.testing.assert_array_equal(l1b_histrate["h_rates"], np.full((2, 60, 7), np.nan))
-    np.testing.assert_array_equal(l1b_histrate["o_rates"], np.full((2, 60, 7), np.nan))
+    np.testing.assert_array_equal(l1b_histrate["h_rates"], np.full((2, 7, 60), np.nan))
+    np.testing.assert_array_equal(l1b_histrate["o_rates"], np.full((2, 7, 60), np.nan))
 
 
 def test_calculate_histogram_rates_zero_exposure_time(l1b_histrates):
@@ -967,10 +969,10 @@ def test_calculate_histogram_rates_zero_exposure_time(l1b_histrates):
         ]
     )
     avg_spin_durations_per_cycle = xr.DataArray([0, 15])
-    exposure_factor = np.zeros((2, 60, 7))
+    exposure_factor = np.zeros((2, 7, 60))
     l1b_histrate = calculate_histogram_rates(
         l1b_histrates, acq_start, acq_end, avg_spin_durations_per_cycle, exposure_factor
     )
 
-    np.testing.assert_array_equal(l1b_histrate["h_rates"], np.full((2, 60, 7), np.nan))
-    np.testing.assert_array_equal(l1b_histrate["o_rates"], np.full((2, 60, 7), np.nan))
+    np.testing.assert_array_equal(l1b_histrate["h_rates"], np.full((2, 7, 60), np.nan))
+    np.testing.assert_array_equal(l1b_histrate["o_rates"], np.full((2, 7, 60), np.nan))