From 9f084ce61b49270e88982868fc90900bc5447140 Mon Sep 17 00:00:00 2001
From: Tim Plummer <Timothy.Plummer@lasp.colorado.edu>
Date: Tue, 12 Nov 2024 17:08:27 -0700
Subject: [PATCH] Add placeholder function for converting from esa_step to
 esa_energy_step

---
 imap_processing/hi/l1b/hi_l1b.py        | 43 +++++++++++++++++++++----
 imap_processing/tests/hi/test_hi_l1b.py | 18 +++++++++++
 2 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/imap_processing/hi/l1b/hi_l1b.py b/imap_processing/hi/l1b/hi_l1b.py
index 97bc4e18e..486096063 100644
--- a/imap_processing/hi/l1b/hi_l1b.py
+++ b/imap_processing/hi/l1b/hi_l1b.py
@@ -123,14 +123,14 @@ def annotate_direct_events(l1a_dataset: xr.Dataset) -> xr.Dataset:
     l1b_dataset.update(compute_coincidence_type_and_time_deltas(l1b_dataset))
     l1b_dataset.update(de_nominal_bin_and_spin_phase(l1b_dataset))
     l1b_dataset.update(compute_hae_coordinates(l1b_dataset))
-    l1b_de_var_names = [
-        "esa_energy_step",
-        "quality_flag",
-    ]
-    new_data_vars = create_dataset_variables(
-        l1b_de_var_names, l1b_dataset["epoch"].size, att_manager_lookup_str="hi_de_{0}"
+    l1b_dataset.update(de_esa_energy_step(l1b_dataset))
+    l1b_dataset.update(
+        create_dataset_variables(
+            ["quality_flag"],
+            l1b_dataset["epoch"].size,
+            att_manager_lookup_str="hi_de_{0}",
+        )
     )
-    l1b_dataset.update(new_data_vars)
     l1b_dataset = l1b_dataset.drop_vars(
         ["tof_1", "tof_2", "tof_3", "de_tag", "ccsds_met", "meta_event_met"]
     )
@@ -341,3 +341,32 @@ def compute_hae_coordinates(dataset: xr.Dataset) -> dict[str, xr.DataArray]:
     new_vars["hae_longitude"].values = pointing_coordinates[:, 1]
 
     return new_vars
+
+
+def de_esa_energy_step(dataset: xr.Dataset) -> dict[str, xr.DataArray]:
+    """
+    Compute esa_energy_step for each direct event.
+
+    TODO: For now this function just returns the esa_step from the input dataset.
+        Eventually, it will take L1B housekeeping data and determine the esa
+        energy steps from that data.
+
+    Parameters
+    ----------
+    dataset : xarray.Dataset
+        The partial L1B dataset.
+
+    Returns
+    -------
+    new_vars : dict[str, xarray.DataArray]
+        Keys are variable names and values are `xarray.DataArray`.
+    """
+    new_vars = create_dataset_variables(
+        ["esa_energy_step"],
+        len(dataset.epoch),
+        att_manager_lookup_str="hi_de_{0}",
+    )
+    # TODO: Implement this algorithm
+    new_vars["esa_energy_step"].values = dataset.esa_step.values
+
+    return new_vars
diff --git a/imap_processing/tests/hi/test_hi_l1b.py b/imap_processing/tests/hi/test_hi_l1b.py
index 9855fead2..e232ec764 100644
--- a/imap_processing/tests/hi/test_hi_l1b.py
+++ b/imap_processing/tests/hi/test_hi_l1b.py
@@ -12,6 +12,7 @@
     CoincidenceBitmap,
     compute_coincidence_type_and_time_deltas,
     compute_hae_coordinates,
+    de_esa_energy_step,
     de_nominal_bin_and_spin_phase,
     hi_l1b,
 )
@@ -235,3 +236,20 @@ def side_effect_func(et, inst_frame: SpiceFrame, to_frame):
     assert "hae_longitude" in new_vars
     assert new_vars["hae_longitude"].shape == fake_dataset.epoch.shape
     np.testing.assert_allclose(new_vars["hae_longitude"].values, sensor_number)
+
+
+def test_de_esa_energy_step():
+    """Test coverage for de_esa_energy_step function."""
+    n_epoch = 20
+    fake_dataset = xr.Dataset(
+        coords={
+            "epoch": xr.DataArray(np.arange(n_epoch), name="epoch", dims=["epoch"])
+        },
+        data_vars={"esa_step": xr.DataArray(np.arange(n_epoch) % 9, dims=["epoch"])},
+    )
+    esa_energy_step_var = de_esa_energy_step(fake_dataset)
+    # TODO: The below check is for the temporary implementation and should be
+    #    removed when the function is update.
+    np.testing.assert_array_equal(
+        esa_energy_step_var["esa_energy_step"].values, fake_dataset.esa_step.values
+    )