diff --git a/imap_processing/cli.py b/imap_processing/cli.py index 324a46378..3e41d5aff 100644 --- a/imap_processing/cli.py +++ b/imap_processing/cli.py @@ -64,7 +64,6 @@ from imap_processing.idex.idex_l1b import idex_l1b from imap_processing.idex.idex_l2a import idex_l2a from imap_processing.idex.idex_l2b import idex_l2b -from imap_processing.lo.constants import LoConstants from imap_processing.lo.l1a import lo_l1a from imap_processing.lo.l1b import lo_l1b from imap_processing.lo.l1c import lo_l1c @@ -1212,53 +1211,6 @@ def do_processing( class Lo(ProcessInstrument): """Process IMAP-Lo.""" - def pre_processing(self) -> ProcessingInputCollection: - """ - Complete pre-processing. - - Extends the base pre-processing by filtering Lo PSET science inputs to - only those whose ``pivot_angle`` is within - ``LoConstants.PSET_PIVOT_ANGLE_TOLERANCE`` of - ``LoConstants.PSET_PIVOT_ANGLE``. PSET files that fall outside this - range are dropped before processing begins. - - Returns - ------- - dependencies : ProcessingInputCollection - Object containing dependencies to process. - """ - datasets = super().pre_processing() - new_datasets = ProcessingInputCollection() - - for processing_input in datasets.get_processing_inputs(): - if ( - processing_input.source == "lo" - and processing_input.descriptor == "pset" - ): - valid_filenames = [] - for imap_file_path in processing_input.imap_file_paths: - pset = load_cdf(imap_file_path.construct_path()) - if "pivot_angle" in pset: - if ( - abs( - pset["pivot_angle"].item() - - LoConstants.PSET_PIVOT_ANGLE - ) - < LoConstants.PSET_PIVOT_ANGLE_TOLERANCE - ): - valid_filenames.append(str(imap_file_path.filename)) - else: - logger.info( - f"Dropping pset {imap_file_path.filename} because " - f"pivot angle is not in range." - ) - if valid_filenames: - new_datasets.add(type(processing_input)(*valid_filenames)) - else: - new_datasets.add(processing_input) - - return new_datasets - def do_processing( self, dependencies: ProcessingInputCollection ) -> list[xr.Dataset]: @@ -1323,18 +1275,21 @@ def do_processing( datasets = lo_l1c.lo_l1c(data_dict, anc_dependencies) elif self.data_level == "l2": - data_dict = {} + sci_dependencies: dict[str, list[xr.Dataset]] = {} science_files = dependencies.get_file_paths(source="lo", descriptor="pset") + science_files += dependencies.get_file_paths(source="lo", data_type="l1b") anc_dependencies = dependencies.get_file_paths(data_type="ancillary") - # Load all pset files into datasets - if not science_files: - logger.info("No valid psets found for L2 processing.") - return datasets + # Load every input, at every pivot angle, grouped by product. + for file in science_files: + dataset = load_cdf(file) + sci_dependencies.setdefault(dataset.attrs["Logical_source"], []).append( + dataset + ) - psets = [load_cdf(file) for file in science_files] - data_dict[psets[0].attrs["Logical_source"]] = psets - datasets = lo_l2.lo_l2(data_dict, anc_dependencies, self.descriptor) + datasets = lo_l2.lo_l2( + sci_dependencies, anc_dependencies, self.descriptor, self.start_date + ) return datasets diff --git a/imap_processing/lo/constants.py b/imap_processing/lo/constants.py index aa0e6e4a6..70cfba5c9 100644 --- a/imap_processing/lo/constants.py +++ b/imap_processing/lo/constants.py @@ -8,11 +8,9 @@ class LoConstants: """Constants for Lo which can be used across different levels.""" - # Expected pivot angle [degrees] for pointing sets for generating map products. - PSET_PIVOT_ANGLE: float = 90.0 - # Absolute tolerance [degrees] for accepting a pset's pivot angle - # as sufficiently close to the required value. - PSET_PIVOT_ANGLE_TOLERANCE: float = 45.0 + # Absolute tolerance [degrees] for accepting a pset's pivot angle as + # sufficiently close to the pivot angle of the map being made. + PSET_PIVOT_ANGLE_TOLERANCE: float = 5.0 # Ion species tracked. "H" is mandatory (and should be the first element); # any others for which we have histrates may be added here. diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index 50021a0bb..b9b3a57d3 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -18,9 +18,14 @@ get_pset_directional_mask, interpolate_map_flux_to_helio_frame, ) -from imap_processing.ena_maps.utils.naming import MapDescriptor +from imap_processing.ena_maps.utils.naming import DAYS_IN_MONTH, MapDescriptor from imap_processing.lo import lo_ancillary -from imap_processing.spice.time import et_to_datetime64, ttj2000ns_to_et +from imap_processing.lo.constants import LoConstants +from imap_processing.spice.time import ( + et_to_datetime64, + str_yyyymmdd_to_ttj2000ns, + ttj2000ns_to_et, +) logger = logging.getLogger(__name__) @@ -30,7 +35,10 @@ def lo_l2( - sci_dependencies: dict, anc_dependencies: list, descriptor: str + sci_dependencies: dict, + anc_dependencies: list, + descriptor: str, + start_date: str | None = None, ) -> list[xr.Dataset]: """ Process IMAP-Lo L1C data into L2 CDF data products. @@ -38,40 +46,56 @@ def lo_l2( This is the main entry point for L2 processing. It orchestrates the entire processing pipeline from L1C pointing sets to L2 sky maps with intensities. + Only the pointing sets taken at the pivot angle of the map being made are + projected onto it, see ``_inputs_at_map_pivot_angle``. When none of the + psets match (or no psets were passed-in to begin with), a quickmap is + produced instead. + Parameters ---------- sci_dependencies : dict Dictionary of datasets needed for L2 data product creation in xarray Datasets. - Must contain "imap_lo_l1c_pset" key with list of pointing set datasets. + Should contain "imap_lo_l1c_pset" key with list of pointing set datasets, + at any pivot angle. anc_dependencies : list List of ancillary file paths needed for L2 data product creation. Should include efficiency factor files. descriptor : str The map descriptor to be produced (e.g., "ilo90-ena-h-sf-nsp-full-hae-6deg-3mo"). + start_date : str, optional + The start of the map window in YYYYMMDD format. Only required for + quickmaps, which have no pointing sets to get their time coverage from. Returns ------- list[xr.Dataset] List containing the processed L2 dataset with rates, intensities, - and uncertainties. + and uncertainties, or the quickmap if there were no pointing sets at + the pivot angle of the map. Raises ------ ValueError - If no pointing set data found in science dependencies. + If a quickmap is required but no start_date was given. NotImplementedError If HEALPix map output is requested (only rectangular maps supported). """ logger.info("Starting IMAP-Lo L2 processing pipeline") - if "imap_lo_l1c_pset" not in sci_dependencies: - raise ValueError("No pointing set data found in science dependencies") - psets = sci_dependencies["imap_lo_l1c_pset"] # Parse the map descriptor to get species and other attributes map_descriptor = MapDescriptor.from_string(descriptor) logger.info(f"Processing map for species: {map_descriptor.species}") + psets = _inputs_at_map_pivot_angle( + sci_dependencies.get("imap_lo_l1c_pset", []), map_descriptor + ) + if not psets: + logger.info("No psets - trying to create a quickmap.") + if start_date is None: + raise ValueError(f"A start_date is required to create the map {descriptor}") + return _lo_l2_quickmap(sci_dependencies, descriptor, start_date) + # Determine if corrections are needed and prepare oxygen data if required ( sputtering_correction, @@ -116,6 +140,161 @@ def lo_l2( return [dataset] +def _inputs_at_map_pivot_angle( + datasets: list[xr.Dataset], map_descriptor: MapDescriptor +) -> list[xr.Dataset]: + """ + Keep only the inputs taken at the pivot angle the map is for. + + Every input we have for the map window is passed in, whatever pivot + angle it was taken at. A map is made from one pivot angle only, which is + the sensor field of its descriptor ("l090" is the 90 degree pivot angle). + Inputs within ``LoConstants.PSET_PIVOT_ANGLE_TOLERANCE`` of that angle are + kept, and inputs with no pivot angle at all are dropped. + + Parameters + ---------- + datasets : list[xr.Dataset] + The input products available for the map window. Every Lo product + records the pivot angle it was taken at. + map_descriptor : MapDescriptor + The parsed descriptor of the map being made. + + Returns + ------- + list[xr.Dataset] + The inputs that belong on this map. + """ + if not isinstance(map_descriptor.sensor, int): + # No pivot angle in the descriptor to select inputs with + return datasets + + kept = [] + for dataset in datasets: + if "pivot_angle" not in dataset: + logger.info("Dropping input with no pivot angle.") + continue + pivot_angle = dataset["pivot_angle"].item() + if ( + abs(pivot_angle - map_descriptor.sensor) + < LoConstants.PSET_PIVOT_ANGLE_TOLERANCE + ): + kept.append(dataset) + else: + logger.info(f"Dropping input with pivot angle {pivot_angle}") + + return kept + + +def _lo_l2_quickmap( + sci_dependencies: dict, descriptor: str, start_date: str +) -> list[xr.Dataset]: + """ + Create a correctly shaped, zero-filled L2 quickmap. + + Parameters + ---------- + sci_dependencies : dict + Dictionary of the input datasets, keyed by logical source, at any + pivot angle. + descriptor : str + The map descriptor to be produced + (e.g., "l090-enansnbs-h-sf-nsp-ram-hae-6deg-6mo"). + start_date : str + The start of the map window in YYYYMMDD format. + + Returns + ------- + list[xr.Dataset] + List containing the single quickmap dataset. + + Raises + ------ + NotImplementedError + If a HEALPix map is requested (only rectangular maps supported for Lo), + or if the map is of a quantity whose variables are not defined yet. + """ + logger.info(f"Creating IMAP-Lo L2 quickmap for descriptor: {descriptor}") + map_descriptor = MapDescriptor.from_string(descriptor) + + sky_map = map_descriptor.to_empty_map() + if not isinstance(sky_map, RectangularSkyMap): + raise NotImplementedError("HEALPix map output not supported for Lo") + + # TODO: No CDF variable attributes are defined for the variables of the + # quantities in non-ENA maps, e.g. isn_rate for an ISN map. + if not map_descriptor.principal_data.startswith("ena"): + raise NotImplementedError( + f"Cannot make a quickmap of {map_descriptor.principal_data} for " + f"{descriptor}. No CDF variable attributes are defined for " + f"{map_descriptor.principal_data_var} in the imap_enamaps_l2 " + f"variable attribute files." + ) + + for logical_source, datasets in sci_dependencies.items(): + selected = _inputs_at_map_pivot_angle(datasets, map_descriptor) + logger.info( + f"{len(selected)} of {len(datasets)} {logical_source} inputs are at " + f"the pivot angle of {descriptor}" + ) + + # No pointing sets available, so the epoch bounds are set using `start_date`. + duration = cast(str, map_descriptor.duration) + if duration.endswith("yr"): + duration_months = int(duration.removesuffix("yr")) * 12 + else: + duration_months = int(duration.removesuffix("mo")) + sky_map.min_epoch = int(str_yyyymmdd_to_ttj2000ns(start_date)) + sky_map.max_epoch = sky_map.min_epoch + int( + np.timedelta64(duration_months * DAYS_IN_MONTH, "D") / np.timedelta64(1, "ns") + ) + + # TODO: Figure out how to handle esa_mode properly + energies = reduce_geometric_factor_dataset(map_descriptor.species, esa_mode=0)[ + "Cntr_E" + ].values + + # The variables an ENA map has, all with dims (epoch, energy, pixel). + float_variables = ( + "ena_intensity", + "ena_intensity_stat_uncert", + "ena_intensity_sys_err", + "ena_count_rate", + "ena_count_rate_stat_uncert", + "ena_count", + "bg_rate", + "bg_rate_stat_uncert", + "bg_rate_sys_err", + "bg_intensity", + "bg_intensity_stat_uncert", + "bg_intensity_sys_err", + "exposure_factor", + ) + int_variables = ("obs_date", "obs_date_range") + variable_dtypes = dict.fromkeys(float_variables, np.float32) | dict.fromkeys( + int_variables, np.int64 + ) + + for variable, dtype in variable_dtypes.items(): + sky_map.data_1d[variable] = xr.DataArray( + np.zeros((1, energies.size, sky_map.num_points), dtype=dtype), + dims=["epoch", "energy", "pixel"], + coords={"energy": energies}, + ) + + dataset = add_geometric_factors(sky_map.to_dataset(), map_descriptor.species) + dataset = cleanup_intermediate_variables(dataset) + + return [ + sky_map.build_cdf_dataset( + instrument="lo", + level="l2", + descriptor=descriptor, + external_map_dataset=dataset, + ) + ] + + def _prepare_corrections( map_descriptor: MapDescriptor, descriptor: str, diff --git a/imap_processing/tests/lo/test_lo_l2.py b/imap_processing/tests/lo/test_lo_l2.py index c053ff8db..9c08f972a 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -1,5 +1,6 @@ """Comprehensive test suite for IMAP-Lo L2 data processing.""" +import logging from pathlib import Path from unittest.mock import Mock, patch @@ -8,7 +9,7 @@ import pytest import xarray as xr -from imap_processing.cdf.utils import load_cdf +from imap_processing.cdf.utils import load_cdf, write_cdf from imap_processing.ena_maps.ena_maps import RectangularSkyMap from imap_processing.ena_maps.utils.corrections import ( add_spacecraft_position_and_velocity_to_pset, @@ -24,6 +25,8 @@ SPIN_ANGLE_BIN_CENTERS, ) from imap_processing.lo.l2.lo_l2 import ( + _inputs_at_map_pivot_angle, + _lo_l2_quickmap, _prepare_corrections, add_efficiency_factors_to_pset, calculate_all_rates_and_intensities, @@ -90,6 +93,7 @@ def sample_pset(): ), "hae_longitude": (("epoch", "spin_angle", "off_angle"), hae_longitude), "hae_latitude": (("epoch", "spin_angle", "off_angle"), hae_latitude), + "pivot_angle": ("epoch", [90.0]), }, coords={ "epoch": [8.1794907049e17], @@ -132,6 +136,7 @@ def sample_pset_for_species(species_name): "exposure_factor": (PSET_DIMS, exposure_factor), "hae_longitude": (("epoch", "spin_angle", "off_angle"), hae_longitude), "hae_latitude": (("epoch", "spin_angle", "off_angle"), hae_latitude), + "pivot_angle": ("epoch", [90.0]), } # Add background rates only for h and o @@ -186,6 +191,7 @@ def minimal_pset(): ), "hae_longitude": (("epoch", "spin_angle", "off_angle"), hae_longitude), "hae_latitude": (("epoch", "spin_angle", "off_angle"), hae_latitude), + "pivot_angle": ("epoch", [90.0]), }, coords={ "epoch": [8.1794907049e17], @@ -227,6 +233,7 @@ def minimal_pset_for_species(species_name): "exposure_factor": (PSET_DIMS, exposure_factor), "hae_longitude": (("epoch", "spin_angle", "off_angle"), hae_longitude), "hae_latitude": (("epoch", "spin_angle", "off_angle"), hae_latitude), + "pivot_angle": ("epoch", [90.0]), } # Add background rates for all species @@ -2582,14 +2589,19 @@ def test_lo_l2_integration_full( class TestErrorHandling: """Tests for error handling in various functions.""" - def test_lo_l2_no_pset_data(self): - """Test error when no pointing set data is provided.""" + def test_lo_l2_no_pset_data(self, caplog): + """Test that a quickmap is made when no pointing set data is provided.""" sci_dependencies = {} # Missing imap_lo_l1c_pset anc_dependencies = [] descriptor = "l090-ena-h-sf-nsp-ram-hae-6deg-3mo" - with pytest.raises(ValueError, match="No pointing set data found"): - lo_l2(sci_dependencies, anc_dependencies, descriptor) + with caplog.at_level(logging.INFO): + (dataset,) = lo_l2( + sci_dependencies, anc_dependencies, descriptor, "20260101" + ) + + assert "No psets" in caplog.text + assert dataset.attrs["Logical_source"] == f"imap_lo_l2_{descriptor}" def test_create_sky_map_healpix_not_supported(self, minimal_pset_for_species): """Test error when HEALPix map is requested.""" @@ -2913,3 +2925,103 @@ def test_project_pset_to_map_value_keys(self, minimal_pset_for_species): for key in expected_keys: assert key in value_keys, f"Expected key '{key}' not in value_keys" + + +class TestPsetsAtMapPivotAngle: + """Tests for selecting the psets that belong on a map.""" + + def test_psets_are_selected_by_map_pivot_angle(self): + """Test that only psets near the descriptor pivot angle are kept.""" + map_descriptor = MapDescriptor.from_string("l090-ena-h-sf-nsp-ram-hae-6deg-1yr") + in_range = xr.Dataset({"pivot_angle": xr.DataArray(90.1)}) + psets = [ + in_range, + # The neighbouring pivot angles belong on their own maps + xr.Dataset({"pivot_angle": xr.DataArray(75.0)}), + xr.Dataset({"pivot_angle": xr.DataArray(105.0)}), + xr.Dataset({"pivot_angle": xr.DataArray(30.0)}), + xr.Dataset(), # No pivot angle at all + ] + + assert _inputs_at_map_pivot_angle(psets, map_descriptor) == [in_range] + + +class TestLoL2Quickmap: + """Tests quickmap when there are no pointing sets.""" + + descriptor = "l090-enansnbs-h-sf-nsp-ram-hae-6deg-6mo" + expected_variables = ( + "ena_intensity", + "ena_intensity_stat_uncert", + "ena_intensity_sys_err", + "ena_count_rate", + "ena_count_rate_stat_uncert", + "ena_count", + "bg_rate", + "bg_rate_stat_uncert", + "bg_rate_sys_err", + "bg_intensity", + "bg_intensity_stat_uncert", + "bg_intensity_sys_err", + "exposure_factor", + "obs_date", + "obs_date_range", + ) + + def test_lo_l2_makes_quickmap_without_psets(self): + """Test that lo_l2 makes a quickmap when there are no pointing sets.""" + (dataset,) = lo_l2({}, [], self.descriptor, start_date="20260101") + + assert dataset.attrs["Logical_source"] == f"imap_lo_l2_{self.descriptor}" + + def test_lo_l2_quickmap_requires_start_date(self): + """Test error when a quickmap is required but no start date is given.""" + with pytest.raises(ValueError, match="start_date is required"): + lo_l2({}, [], self.descriptor) + + def test_quickmap_selects_inputs_by_pivot_angle(self, caplog): + """Test that the quickmap reports the inputs belonging on the map.""" + sci_dependencies = { + "imap_lo_l1b_de": [ + xr.Dataset({"pivot_angle": xr.DataArray(90.0)}), + xr.Dataset({"pivot_angle": xr.DataArray(75.0)}), + ] + } + + with caplog.at_level(logging.INFO): + _lo_l2_quickmap(sci_dependencies, self.descriptor, "20260101") + + assert "1 of 2 imap_lo_l1b_de inputs" in caplog.text + + def test_quickmap_shape(self): + """Test that the quickmap has the shape of a real 6deg map.""" + (dataset,) = _lo_l2_quickmap({}, self.descriptor, "20260101") + + assert dict(dataset.sizes) == { + "epoch": 1, + "energy": 7, + "longitude": 60, + "latitude": 30, + } + for variable in self.expected_variables: + assert dataset[variable].dims == ( + "epoch", + "energy", + "longitude", + "latitude", + ) + + # Intermediate variables should not make it into the output + assert "geometric_factor" not in dataset.data_vars + + def test_quickmap_writes_to_cdf(self): + """Test that the quickmap can be written out as a valid CDF.""" + (dataset,) = _lo_l2_quickmap({}, self.descriptor, "20260101") + dataset.attrs["Data_version"] = "001.0001" + dataset.attrs["Start_date"] = "20260101" + + cdf_path = write_cdf(dataset) + + assert cdf_path.exists() + assert cdf_path.name == (f"imap_lo_l2_{self.descriptor}_20260101_v001.0001.cdf") + assert dict(load_cdf(cdf_path).sizes) == dict(dataset.sizes) diff --git a/imap_processing/tests/test_cli.py b/imap_processing/tests/test_cli.py index 92acd8cbe..9d79119ef 100644 --- a/imap_processing/tests/test_cli.py +++ b/imap_processing/tests/test_cli.py @@ -6,7 +6,7 @@ import sys from pathlib import Path from unittest import mock -from unittest.mock import Mock, sentinel +from unittest.mock import Mock import imap_data_access.io import numpy as np @@ -472,17 +472,24 @@ def test_lo_l2(mock_lo_pre_processing, mock_lo_l2, mock_instrument_dependencies) descriptor = "some-ena-map-descriptor" - mock_loaded_pset_1 = Mock(attrs={"Logical_source": "some_pset_logical_source"}) - pset_file_paths = [ + mock_loaded_pset_1 = Mock(attrs={"Logical_source": "imap_lo_l1c_pset"}) + mock_loaded_pset_2 = Mock(attrs={"Logical_source": "imap_lo_l1c_pset"}) + mock_loaded_de = Mock(attrs={"Logical_source": "imap_lo_l1b_de"}) + science_file_paths = [ "imap_lo_l1c_pset_20250415_v001.cdf", "imap_lo_l1c_pset_20250416_v001.cdf", + "imap_lo_l1b_de_20250415_v001.cdf", ] processing_input = ProcessingInputCollection( - *[ScienceInput(file_path) for file_path in pset_file_paths], + *[ScienceInput(file_path) for file_path in science_file_paths], ) - mocks["mock_load_cdf"].side_effect = [mock_loaded_pset_1, sentinel.loaded_pset_2] + mocks["mock_load_cdf"].side_effect = [ + mock_loaded_pset_1, + mock_loaded_pset_2, + mock_loaded_de, + ] mock_lo_pre_processing.return_value = processing_input output_l2_dataset = xr.Dataset() @@ -500,47 +507,15 @@ def test_lo_l2(mock_lo_pre_processing, mock_lo_l2, mock_instrument_dependencies) instrument.process() mock_lo_l2.assert_called_once_with( - {"some_pset_logical_source": [mock_loaded_pset_1, sentinel.loaded_pset_2]}, + { + "imap_lo_l1c_pset": [mock_loaded_pset_1, mock_loaded_pset_2], + "imap_lo_l1b_de": [mock_loaded_de], + }, [], descriptor, - ) - mocks["mock_write_cdf"].assert_called_once_with(output_l2_dataset) - - -@mock.patch("imap_processing.cli.load_cdf") -@mock.patch("imap_processing.cli.ProcessInstrument.pre_processing") -def test_lo_pre_processing_pivot_angle_filter(mock_super_pre_processing, mock_load_cdf): - valid_pset = "imap_lo_l1c_pset_20250415_v001.cdf" - invalid_pset = "imap_lo_l1c_pset_20250416_v001.cdf" - non_pset = "imap_lo_l1a_de_20260415-repoint00217_v001.cdf" - - base_collection = ProcessingInputCollection( - ScienceInput(valid_pset, invalid_pset), - ScienceInput(non_pset), - ) - mock_super_pre_processing.return_value = base_collection - mock_load_cdf.side_effect = [ - xr.Dataset({"pivot_angle": xr.DataArray(90.1)}), - xr.Dataset({"pivot_angle": xr.DataArray(30.0)}), - ] - - instrument = Lo( - "l2", - "some-descriptor", - base_collection.serialize(), "20250415", - "20250416", - "v001", - False, ) - result = instrument.pre_processing() - - result_inputs = list(result.get_processing_inputs()) - assert len(result_inputs) == 2 - - pset_input, non_pset_input = result_inputs - assert [str(fp.filename) for fp in pset_input.imap_file_paths] == [valid_pset] - assert [str(fp.filename) for fp in non_pset_input.imap_file_paths] == [non_pset] + mocks["mock_write_cdf"].assert_called_once_with(output_l2_dataset) @mock.patch("imap_processing.cli.quaternions.process_quaternions", autospec=True)