imap-processing 0.11.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

imap_processing/tests/hit/test_hit_l1a.py

@@ -27,9 +27,7 @@ from imap_processing.tests.hit.helpers.l1_validation import (
 @pytest.fixture(scope="module")
 def hk_packet_filepath():
     """Set path to test data file"""
-    return (
-        imap_module_directory / "tests/hit/test_data/imap_hit_l0_raw_20100105_v001.pkts"
-    )
+    return imap_module_directory / "tests/hit/test_data/hskp_sample.ccsds"
 
 
 @pytest.fixture(scope="module")
@@ -54,7 +52,7 @@ def validation_data():
 def test_subcom_sectorates(sci_packet_filepath):
     """Test the subcom_sectorates function.
 
-    This function organizes the sector rates data
+    This function organizes the sectored rates data
     by species and adds the data as new variables
     to the dataset.
     """
@@ -69,26 +67,26 @@ def test_subcom_sectorates(sci_packet_filepath):
     # Number of science frames in the dataset
     frames = sci_dataset["epoch"].shape[0]
 
-    # Check if the dataset has the expected new variables
-    for species in ["h", "he4", "cno", "nemgsi", "fe"]:
-        assert f"{species}_counts_sectored" in sci_dataset
-        assert f"{species}_energy_min" in sci_dataset
-        assert f"{species}_energy_max" in sci_dataset
+    # Shape of the new data variables
+    expected_shapes = {
+        "h": (3, 15, 8),
+        "he4": (2, 15, 8),
+        "cno": (2, 15, 8),
+        "nemgsi": (2, 15, 8),
+        "fe": (1, 15, 8),
+    }
 
+    for species, shape in expected_shapes.items():
+        # Check if the dataset has the new  data variables
+        assert f"{species}_sectored_counts" in sci_dataset
+        assert f"{species}_energy_mean" in sci_dataset
+        assert f"{species}_energy_delta_minus" in sci_dataset
+        assert f"{species}_energy_delta_plus" in sci_dataset
         # Check the shape of the new data variables
-        if species == "h":
-            assert sci_dataset[f"{species}_counts_sectored"].shape == (frames, 3, 8, 15)
-            assert sci_dataset[f"{species}_energy_min"].shape == (3,)
-        elif species in ("4he", "cno", "nemgsi"):
-            assert sci_dataset[f"{species}_counts_sectored"].shape == (frames, 2, 8, 15)
-            assert sci_dataset[f"{species}_energy_min"].shape == (2,)
-        elif species == "fe":
-            assert sci_dataset[f"{species}_counts_sectored"].shape == (frames, 1, 8, 15)
-            assert sci_dataset[f"{species}_energy_min"].shape == (1,)
-        assert (
-            sci_dataset[f"{species}_energy_max"].shape
-            == sci_dataset[f"{species}_energy_min"].shape
-        )
+        assert sci_dataset[f"{species}_sectored_counts"].shape == (frames, *shape)
+        assert sci_dataset[f"{species}_energy_mean"].shape == (shape[0],)
+        assert sci_dataset[f"{species}_energy_delta_minus"].shape == (shape[0],)
+        assert sci_dataset[f"{species}_energy_delta_plus"].shape == (shape[0],)
 
 
 def test_calculate_uncertainties():
@@ -151,6 +149,7 @@ def test_validate_l1a_housekeeping_data(hk_packet_filepath):
     )
     validation_data = pd.read_csv(validation_file)
     validation_data.columns = validation_data.columns.str.lower()
+    validation_data.columns = validation_data.columns.str.strip()
 
     # Get a list of leak columns in ascending order
     # (LEAK_I_00, LEAK_I_01, ..., LEAK_I_63)
@@ -166,7 +165,6 @@ def test_validate_l1a_housekeeping_data(hk_packet_filepath):
     # Define the keys that should have dropped from the housekeeping dataset
     dropped_fields = {
         "pkt_apid",
-        "sc_tick",
         "version",
         "type",
         "sec_hdr_flg",
@@ -190,7 +188,6 @@ def test_validate_l1a_housekeeping_data(hk_packet_filepath):
         "ccsds_grp_flag",
         "ccsds_seq_cnt",
         "ccsds_length",
-        "shcoarse",
     }
 
     # Check that dropped variables are not in the dataset
@@ -243,7 +240,7 @@ def test_validate_l1a_counts_data(sci_packet_filepath, validation_data):
         "seq_flgs",
         "src_seq_ctr",
         "pkt_len",
-        "energy_idx",
+        "energy_bin",
     ]
 
     # Compare processed data to validation data
@@ -272,10 +269,9 @@ def test_hit_l1a(hk_packet_filepath, sci_packet_filepath):
         else:
             assert len(processed_datasets) == 2
             assert (
-                processed_datasets[0].attrs["Logical_source"]
-                == "imap_hit_l1a_count-rates"
+                processed_datasets[0].attrs["Logical_source"] == "imap_hit_l1a_counts"
             )
             assert (
                 processed_datasets[1].attrs["Logical_source"]
-                == "imap_hit_l1a_pulse-height-events"
+                == "imap_hit_l1a_direct-events"
             )

imap_processing/tests/hit/test_hit_l1b.py

@@ -5,7 +5,17 @@ import xarray as xr
 
 from imap_processing import imap_module_directory
 from imap_processing.hit.l1a import hit_l1a
-from imap_processing.hit.l1b import hit_l1b
+from imap_processing.hit.l1b.hit_l1b import (
+    PARTICLE_ENERGY_RANGE_MAPPING,
+    SummedCounts,
+    add_energy_variables,
+    add_rates_to_dataset,
+    calculate_summed_counts,
+    create_particle_data_arrays,
+    hit_l1b,
+    process_standard_rates_data,
+    process_summed_rates_data,
+)
 from imap_processing.tests.hit.helpers.l1_validation import (
     prepare_standard_rates_validation_data,
 )
@@ -48,7 +58,7 @@ def dependencies(packet_filepath, sci_packet_filepath):
 @pytest.fixture()
 def l1b_hk_dataset(dependencies):
     """Get the housekeeping dataset"""
-    datasets = hit_l1b.hit_l1b(dependencies, "001")
+    datasets = hit_l1b(dependencies, "001")
     for dataset in datasets:
         if dataset.attrs["Logical_source"] == "imap_hit_l1b_hk":
             return dataset
@@ -58,7 +68,7 @@ def l1b_hk_dataset(dependencies):
 def l1b_standard_rates_dataset(dependencies):
     """Get the standard rates dataset"""
     # TODO: use this fixture in future unit test to validate the standard rates dataset
-    datasets = hit_l1b.hit_l1b(dependencies, "001")
+    datasets = hit_l1b(dependencies, "001")
     for dataset in datasets:
         if dataset.attrs["Logical_source"] == "imap_hit_l1b_standard-rates":
             return dataset
@@ -69,15 +79,232 @@ def l1a_counts_dataset(sci_packet_filepath):
     """Get L1A counts dataset to test l1b processing functions"""
     l1a_datasets = hit_l1a.hit_l1a(sci_packet_filepath, "001")
     for dataset in l1a_datasets:
-        if dataset.attrs["Logical_source"] == "imap_hit_l1a_count-rates":
+        if dataset.attrs["Logical_source"] == "imap_hit_l1a_counts":
             return dataset
 
 
-def test_process_standard_rates_data(l1a_counts_dataset):
-    """Test function for processing standard rates data"""
-    l1b_standard_rates_dataset = hit_l1b.process_standard_rates_data(l1a_counts_dataset)
+@pytest.fixture()
+def livetime(l1a_counts_dataset):
+    """Calculate livetime for L1A counts dataset"""
+    return l1a_counts_dataset["livetime_counter"] / 270
+
+
+def test_calculate_summed_counts():
+    # Create a mock raw_counts_dataset
+    data = {
+        "l2fgrates": (
+            ("epoch", "index"),
+            np.array([[1, 2, 3, 4, 5]] * 5, dtype=np.int64),
+        ),
+        "l3fgrates": (
+            ("epoch", "index"),
+            np.array([[6, 7, 8, 9, 10]] * 5, dtype=np.int64),
+        ),
+        "penfgrates": (
+            ("epoch", "index"),
+            np.array([[11, 12, 13, 14, 15]] * 5, dtype=np.int64),
+        ),
+        "l2fgrates_delta_minus": (
+            ("epoch", "index"),
+            np.zeros((5, 5), dtype=np.float32),
+        ),
+        "l3fgrates_delta_minus": (
+            ("epoch", "index"),
+            np.full((5, 5), 0.01, dtype=np.float32),
+        ),
+        "penfgrates_delta_minus": (
+            ("epoch", "index"),
+            np.full((5, 5), 0.001, dtype=np.float32),
+        ),
+        "l2fgrates_delta_plus": (
+            ("epoch", "index"),
+            np.full((5, 5), 0.02, dtype=np.float32),
+        ),
+        "l3fgrates_delta_plus": (
+            ("epoch", "index"),
+            np.full((5, 5), 0.002, dtype=np.float32),
+        ),
+        "penfgrates_delta_plus": (
+            ("epoch", "index"),
+            np.full((5, 5), 0.003, dtype=np.float32),
+        ),
+    }
+    coords = {"epoch": np.arange(5), "index": np.arange(5)}
+    raw_counts_dataset = xr.Dataset(data, coords=coords)
+
+    # Define count_indices
+    count_indices = {
+        "R2": [0, 1],
+        "R3": [2, 3],
+        "R4": [4],
+    }
+
+    # Call the function
+    summed_counts, summed_counts_delta_minus, summed_counts_delta_plus = (
+        calculate_summed_counts(raw_counts_dataset, count_indices)
+    )
+
+    # Expected values based on `count_indices`
+    expected_summed_counts = np.array([35, 35, 35, 35, 35])
+    expected_summed_counts_delta_minus = np.array([0.021, 0.021, 0.021, 0.021, 0.021])
+    expected_summed_counts_delta_plus = np.array([0.047, 0.047, 0.047, 0.047, 0.047])
+
+    # Assertions
+    assert summed_counts.shape == (5,)
+    assert summed_counts_delta_minus.shape == (5,)
+    assert summed_counts_delta_plus.shape == (5,)
+
+    np.testing.assert_array_almost_equal(summed_counts.values, expected_summed_counts)
+    np.testing.assert_array_almost_equal(
+        summed_counts_delta_minus.values, expected_summed_counts_delta_minus
+    )
+    np.testing.assert_array_almost_equal(
+        summed_counts_delta_plus.values, expected_summed_counts_delta_plus
+    )
+
+    # Check dtype consistency
+    assert summed_counts.dtype == np.int64, f"Unexpected dtype: {summed_counts.dtype}"
+    assert (
+        summed_counts_delta_minus.dtype == np.float32
+    ), f"Unexpected dtype: {summed_counts_delta_minus.dtype}"
+    assert (
+        summed_counts_delta_plus.dtype == np.float32
+    ), f"Unexpected dtype: {summed_counts_delta_plus.dtype}"
+
+
+def test_add_rates_to_dataset():
+    # Create a sample dataset
+    dataset = xr.Dataset(
+        {
+            "epoch": ("epoch", np.arange(10)),
+            "livetime": ("epoch", np.random.rand(10) + 1),  # Avoid division by zero
+        }
+    )
+
+    # Add empty data arrays for a sample particle
+    particle = "test_particle"
+    dataset[particle] = xr.DataArray(
+        data=np.zeros((10, 5), dtype=np.float32),
+        dims=["epoch", f"{particle}_energy_index"],
+    )
+    dataset[f"{particle}_delta_minus"] = xr.DataArray(
+        data=np.zeros((10, 5), dtype=np.float32),
+        dims=["epoch", f"{particle}_energy_index"],
+    )
+    dataset[f"{particle}_delta_plus"] = xr.DataArray(
+        data=np.zeros((10, 5), dtype=np.float32),
+        dims=["epoch", f"{particle}_energy_index"],
+    )
+
+    # Set the random seed for reproducibility
+    np.random.seed(42)
+
+    # Define the summed counts with random values in a namedtuple
+    summed_counts = SummedCounts(
+        xr.DataArray(np.random.rand(10), dims=["epoch"]),
+        xr.DataArray(np.random.rand(10), dims=["epoch"]),
+        xr.DataArray(np.random.rand(10), dims=["epoch"]),
+    )
+
+    # Call the function
+    updated_dataset = add_rates_to_dataset(
+        dataset, particle, 0, summed_counts, dataset["livetime"]
+    )
+
+    # Check the results
+    np.testing.assert_array_almost_equal(
+        updated_dataset[particle][:, 0].values,
+        summed_counts.summed_counts / dataset["livetime"].values,
+    )
+    np.testing.assert_array_almost_equal(
+        updated_dataset[f"{particle}_delta_minus"][:, 0].values,
+        summed_counts.summed_counts_delta_minus / dataset["livetime"].values,
+    )
+    np.testing.assert_array_almost_equal(
+        updated_dataset[f"{particle}_delta_plus"][:, 0].values,
+        summed_counts.summed_counts_delta_plus / dataset["livetime"].values,
+    )
+
+
+def test_add_energy_variables():
+    dataset = xr.Dataset()
+    particle = "test_particle"
+    energy_min = np.array([1.8, 4.0, 6.0], dtype=np.float32)
+    energy_max = np.array([2.2, 6.0, 10.0], dtype=np.float32)
+    result = add_energy_variables(dataset, particle, energy_min, energy_max)
+    assert f"{particle}_energy_min" in result.data_vars
+    assert f"{particle}_energy_max" in result.data_vars
+    assert np.all(result[f"{particle}_energy_min"].values == energy_min)
+    assert np.all(result[f"{particle}_energy_max"].values == energy_max)
+
+
+def test_create_particle_data_arrays():
+    dataset = xr.Dataset()
+    particle = "test_particle"
+    result = create_particle_data_arrays(
+        dataset, particle, num_energy_ranges=3, epoch_size=10
+    )
+
+    assert f"{particle}" in result.data_vars
+    assert f"{particle}_delta_minus" in result.data_vars
+    assert f"{particle}_delta_plus" in result.data_vars
+    assert f"{particle}_energy_index" in result.coords
+
+    for var in result.data_vars:
+        assert result[var].shape == (10, 3)
+
+    assert result[f"{particle}_energy_index"].shape == (3,)
+
+
+def test_process_summed_rates_data(l1a_counts_dataset, livetime):
+    """Test the variables in the summed rates dataset"""
+
+    l1b_summed_rates_dataset = process_summed_rates_data(l1a_counts_dataset, livetime)
+
+    # Check that a xarray dataset is returned
+    assert isinstance(l1b_summed_rates_dataset, xr.Dataset)
+
+    valid_coords = {
+        "epoch",
+        "h_energy_index",
+        "he3_energy_index",
+        "he4_energy_index",
+        "he_energy_index",
+        "c_energy_index",
+        "o_energy_index",
+        "fe_energy_index",
+        "n_energy_index",
+        "si_energy_index",
+        "mg_energy_index",
+        "s_energy_index",
+        "ar_energy_index",
+        "ca_energy_index",
+        "na_energy_index",
+        "al_energy_index",
+        "ne_energy_index",
+        "ni_energy_index",
+    }
+
+    # Check that the dataset has the correct coords and variables
+    assert valid_coords == set(l1b_summed_rates_dataset.coords), "Coordinates mismatch"
+
+    assert "dynamic_threshold_state" in l1b_summed_rates_dataset.data_vars
+
+    for particle in PARTICLE_ENERGY_RANGE_MAPPING.keys():
+        assert f"{particle}" in l1b_summed_rates_dataset.data_vars
+        assert f"{particle}_delta_minus" in l1b_summed_rates_dataset.data_vars
+        assert f"{particle}_delta_plus" in l1b_summed_rates_dataset.data_vars
+        assert f"{particle}_energy_min" in l1b_summed_rates_dataset.data_vars
+        assert f"{particle}_energy_max" in l1b_summed_rates_dataset.data_vars
+
+
+def test_process_standard_rates_data(l1a_counts_dataset, livetime):
+    """Test the variables in the standard rates dataset"""
+    l1b_standard_rates_dataset = process_standard_rates_data(
+        l1a_counts_dataset, livetime
+    )
 
-    # Check that a xarray dataset with the correct logical source is returned
+    # Check that a xarray dataset is returned
     assert isinstance(l1b_standard_rates_dataset, xr.Dataset)
 
     # Define the data variables that should be present in the dataset
@@ -152,7 +379,6 @@ def test_hit_l1b_hk_dataset_variables(l1b_hk_dataset):
     # Define the keys that should have dropped from the housekeeping dataset
     dropped_keys = {
         "pkt_apid",
-        "sc_tick",
         "version",
         "type",
         "sec_hdr_flg",
@@ -167,6 +393,7 @@ def test_hit_l1b_hk_dataset_variables(l1b_hk_dataset):
     }
     # Define the keys that should be present in the housekeeping dataset
     valid_keys = {
+        "sc_tick",
         "heater_on",
         "fsw_version_b",
         "ebox_m12va",
@@ -239,16 +466,13 @@ def test_validate_l1b_hk_data(l1b_hk_dataset):
 
     Parameters
     ----------
-    hk_dataset : xr.Dataset
-        Housekeeping dataset created by the L1B processing.
+    l1b_hk_dataset : xr.Dataset
+        Housekeeping dataset created by L1B processing.
     """
-    # TODO: finish test. HIT will provide an updated validation file to fix issues:
-    #  - some fields have strings as values but in the processed data they're integers
-    #  - Some columns have blank cells where there should be data
 
     # Load the validation data
     validation_file = (
-        imap_module_directory / "tests/hit/validation_data/hskp_sample_eu.csv"
+        imap_module_directory / "tests/hit/validation_data/hskp_sample_eu_3_6_2025.csv"
     )
     validation_data = pd.read_csv(validation_file)
     validation_data.columns = validation_data.columns.str.lower().str.strip()
@@ -260,14 +484,13 @@ def test_validate_l1b_hk_data(l1b_hk_dataset):
         col for col in validation_data.columns if col.startswith("leak_i_")
     ][::-1]
     validation_data["leak_i"] = validation_data[leak_columns].apply(
-        lambda row: row.values, axis=1
+        lambda row: row.values.astype(np.float64), axis=1
     )
     validation_data.drop(columns=leak_columns, inplace=True)
 
     # Define the keys that should have dropped from the housekeeping dataset
     dropped_fields = {
         "pkt_apid",
-        "sc_tick",
         "version",
         "type",
         "sec_hdr_flg",
@@ -284,29 +507,47 @@ def test_validate_l1b_hk_data(l1b_hk_dataset):
     # Check that dropped variables are not in the dataset
     assert set(dropped_fields).isdisjoint(set(l1b_hk_dataset.data_vars.keys()))
 
-    # TODO: uncomment block after new validation data is provided
     # Define the keys that should be ignored in the validation
     # like ccsds headers
-    # ignore_validation_fields = {
-    #     "ccsds_version",
-    #     "ccsds_type",
-    #     "ccsds_sec_hdr_flag",
-    #     "ccsds_appid",
-    #     "ccsds_grp_flag",
-    #     "ccsds_seq_cnt",
-    #     "ccsds_length",
-    #     "sc_tick",
-    # }
-
-    # # Compare the housekeeping dataset with the expected validation data
-    # for field in validation_data.columns:
-    #     if field not in ignore_validation_fields:
-    #         print(field)
-    #         assert field in hk_dataset.data_vars.keys()
-    #         for pkt in range(validation_data.shape[0]):
-    #             assert np.array_equal(
-    #                 hk_dataset[field][pkt].data, validation_data[field][pkt]
-    #             )
+    ignore_validation_fields = {
+        "ccsds_version",
+        "ccsds_type",
+        "ccsds_sec_hdr_flag",
+        "ccsds_appid",
+        "ccsds_grp_flag",
+        "ccsds_seq_cnt",
+        "ccsds_length",
+        "sc_tick",
+    }
+
+    for field in validation_data.columns:
+        if field not in ignore_validation_fields:
+            assert field in l1b_hk_dataset.data_vars.keys()
+            if field == "leak_i":
+                # Compare leak_i arrays
+                # Reshape validation_data to match the shape of l1b_hk_dataset
+                reshaped_validation_data = np.vstack(validation_data[field].values)
+                # Compare leak_i arrays
+                np.testing.assert_allclose(
+                    l1b_hk_dataset[field].values.astype(np.float64),
+                    reshaped_validation_data,
+                    atol=1e-2,
+                    err_msg=f"Mismatch in {field}",
+                )
+            elif l1b_hk_dataset[field].dtype.kind == "U":
+                np.testing.assert_array_equal(
+                    l1b_hk_dataset[field].values,
+                    validation_data[field].str.strip().values,
+                    err_msg=f"Mismatch in {field}",
+                )
+            else:
+                # Compare float values
+                np.testing.assert_allclose(
+                    l1b_hk_dataset[field].values.astype(np.float64),
+                    validation_data[field].values,
+                    atol=1e-2,
+                    err_msg=f"Mismatch in {field}",
+                )
 
 
 def test_validate_l1b_standard_rates_data(l1b_standard_rates_dataset):
@@ -333,6 +574,28 @@ def test_validate_l1b_standard_rates_data(l1b_standard_rates_dataset):
             )
 
 
+def test_hit_l1b_missing_apid(sci_packet_filepath):
+    """Test missing housekeeping apid from packet file
+
+    Check that no L1B datasets are created when the housekeeping
+    apid is missing from the L0 file path dependency since APIDs
+    are currently in separate packet files.
+
+    In the future, all HIT APIDs will be included in the same packet file.
+
+    Parameters
+    ----------
+    sci_packet_filepath : str
+        Science CCSDS packet file path. Only contains science APID and is
+        missing the housekeeping APID.
+    """
+    # Create a dependency dictionary with a science CCSDS packet file
+    # excluding the housekeeping apid
+    dependency = {"imap_hit_l0_raw": sci_packet_filepath}
+    datasets = hit_l1b(dependency, "001")
+    assert len(datasets) == 0
+
+
 def test_hit_l1b(dependencies):
     """Test creating L1B CDF files
 
@@ -344,10 +607,11 @@ def test_hit_l1b(dependencies):
         Dictionary of L1A datasets and CCSDS packet file path
     """
     # TODO: update assertions after science data processing is completed
-    datasets = hit_l1b.hit_l1b(dependencies, "001")
+    datasets = hit_l1b(dependencies, "001")
 
-    assert len(datasets) == 2
+    assert len(datasets) == 3
     for dataset in datasets:
         assert isinstance(dataset, xr.Dataset)
     assert datasets[0].attrs["Logical_source"] == "imap_hit_l1b_hk"
     assert datasets[1].attrs["Logical_source"] == "imap_hit_l1b_standard-rates"
+    assert datasets[2].attrs["Logical_source"] == "imap_hit_l1b_summed-rates"