diff --git a/imap_processing/ialirt/l0/process_swe.py b/imap_processing/ialirt/l0/process_swe.py index 6b9d2944f..3b56b9c51 100644 --- a/imap_processing/ialirt/l0/process_swe.py +++ b/imap_processing/ialirt/l0/process_swe.py @@ -475,7 +475,9 @@ def process_swe(accumulated_data: xr.Dataset, in_flight_cal_files: list) -> list # Get total full cycle data available for processing. # There are 60 packets in a set so (0, 59) is the range. - grouped_data = find_groups(accumulated_data, (0, 59), "swe_seq", "time_seconds") + grouped_data = find_groups( + accumulated_data, (0, 59), "swe_seq", "met", "swe_nom_flag" + ) unique_groups = np.unique(grouped_data["group"]) swe_data: list[dict] = [] incomplete_groups = [] diff --git a/imap_processing/ialirt/utils/grouping.py b/imap_processing/ialirt/utils/grouping.py index 27bf3fcbc..cb28193b6 100644 --- a/imap_processing/ialirt/utils/grouping.py +++ b/imap_processing/ialirt/utils/grouping.py @@ -8,7 +8,9 @@ logger = logging.getLogger(__name__) -def filter_valid_groups(grouped_data: xr.Dataset) -> xr.Dataset: +def filter_valid_groups( + grouped_data: xr.Dataset, flag: str | None = None +) -> xr.Dataset: """ Filter out groups where `src_seq_ctr` diff are not 1. @@ -16,6 +18,8 @@ def filter_valid_groups(grouped_data: xr.Dataset) -> xr.Dataset: ---------- grouped_data : xr.Dataset Dataset with a "group" coordinate. + flag : str | None + Optional name of flag data variable. Returns ------- @@ -42,6 +46,12 @@ def filter_valid_groups(grouped_data: xr.Dataset) -> xr.Dataset: drop=True, ) + if flag: + filtered_data = filtered_data.where( + filtered_data[flag] != 0, + drop=True, + ) + return filtered_data @@ -50,6 +60,7 @@ def find_groups( sequence_range: tuple, sequence_name: str, time_name: str, + flag: str | None = None, ) -> xr.Dataset: """ Group data based on time and sequence number values. @@ -64,6 +75,8 @@ def find_groups( Name of the sequence variable. time_name : str Name of the time variable. + flag : str | None + Optional name of flag data variable. Returns ------- @@ -82,9 +95,14 @@ def find_groups( # Use sequence_range == 0 to define the beginning of the group. # Find time at this index and use it as the beginning time for the group. - start_times = sorted_data[time_name][ - (sorted_data[sequence_name] == sequence_range[0]) - ] + if flag: + start_times = sorted_data[time_name][ + (sorted_data[sequence_name] == sequence_range[0]) & (sorted_data[flag] != 0) + ] + else: + start_times = sorted_data[time_name][ + (sorted_data[sequence_name] == sequence_range[0]) + ] # Use max sequence_range to define the end of the group. end_times = sorted_data[time_name][ ([sorted_data[sequence_name] == sequence_range[-1]][-1]) @@ -115,6 +133,6 @@ def find_groups( grouped_data = grouped_data.assign_coords(group=("epoch", group_labels)) # Filter out groups with non-sequential src_seq_ctr values. - filtered_data = filter_valid_groups(grouped_data) + filtered_data = filter_valid_groups(grouped_data, flag) return filtered_data diff --git a/imap_processing/tests/ialirt/unit/test_process_swe.py b/imap_processing/tests/ialirt/unit/test_process_swe.py index a9946deff..1df92c3fb 100644 --- a/imap_processing/tests/ialirt/unit/test_process_swe.py +++ b/imap_processing/tests/ialirt/unit/test_process_swe.py @@ -168,6 +168,7 @@ def test_process_spacecraft_packet( np.arange(462466219, 462466219 + n, dtype=np.uint32), ) sc_xarray_data["swe_seq"] = ("epoch", np.arange(n) % 60) + sc_xarray_data["swe_nom_flag"] = xr.ones_like(sc_xarray_data["swe_nom_flag"]) in_flight_cal_file = ( imap_module_directory