From 579a4e05f771a947237b46f789cbbbd22be0fe6e Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Fri, 24 Jul 2026 17:25:38 -0400 Subject: [PATCH 01/10] a placeholder quickmap product with all 0s, just to start producing something --- imap_processing/cli.py | 67 ++------- imap_processing/lo/constants.py | 8 +- imap_processing/lo/l2/lo_l2.py | 197 +++++++++++++++++++++++-- imap_processing/tests/lo/test_lo_l2.py | 122 ++++++++++++++- imap_processing/tests/test_cli.py | 59 +++----- 5 files changed, 336 insertions(+), 117 deletions(-) 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) From d60aaafdcf0f2ebf9d82319dd3faaa7eaee89ee5 Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Sun, 26 Jul 2026 11:49:51 -0400 Subject: [PATCH 02/10] filtering by pivot angle happens at pre-processing stage, for map products --- imap_processing/cli.py | 89 ++++++++++++++++++++++++++ imap_processing/lo/l2/lo_l2.py | 74 ++++----------------- imap_processing/tests/lo/test_lo_l2.py | 26 +------- imap_processing/tests/test_cli.py | 45 +++++++++++++ 4 files changed, 149 insertions(+), 85 deletions(-) diff --git a/imap_processing/cli.py b/imap_processing/cli.py index 3e41d5aff..8185509f5 100644 --- a/imap_processing/cli.py +++ b/imap_processing/cli.py @@ -53,6 +53,7 @@ # In code: # call cdf.utils.write_cdf from imap_processing.codice import codice_l1a, codice_l1b, codice_l2 +from imap_processing.ena_maps.utils.naming import MapDescriptor from imap_processing.glows.l1a.glows_l1a import glows_l1a from imap_processing.glows.l1b.glows_l1b import glows_l1b, glows_l1b_de from imap_processing.glows.l2.glows_l2 import glows_l2 @@ -64,6 +65,7 @@ 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 @@ -1211,6 +1213,93 @@ def do_processing( class Lo(ProcessInstrument): """Process IMAP-Lo.""" + @staticmethod + def _input_at_pivot_angle( + imap_file_path: imap_data_access.ImapFilePath, map_pivot_angle: int + ) -> bool: + """ + Check whether a Lo input file was taken at the pivot angle of a map. + + Every Lo science product records the pivot angle it was taken at. Files + without a pivot angle can't be placed on a map and so are rejected. + + Parameters + ---------- + imap_file_path : imap_data_access.ImapFilePath + The Lo input file to check. + map_pivot_angle : int + The pivot angle [degrees] of the map being made. + + Returns + ------- + bool + Whether the file belongs on the map. + """ + dataset = load_cdf(imap_file_path.construct_path()) + if "pivot_angle" not in dataset: + logger.info(f"Dropping {imap_file_path.filename} - no pivot angle.") + return False + + pivot_angle = dataset["pivot_angle"].item() + if abs(pivot_angle - map_pivot_angle) >= LoConstants.PSET_PIVOT_ANGLE_TOLERANCE: + logger.info( + f"Dropping {imap_file_path.filename}, its pivot angle {pivot_angle} " + f"is not the {map_pivot_angle} degree pivot angle of the map." + ) + return False + + return True + + def pre_processing(self) -> ProcessingInputCollection: + """ + Complete pre-processing. + + Extends the base pre-processing by dropping, for map products, the Lo + science inputs that were not taken at the pivot angle of the map being + made. + + Returns + ------- + dependencies : ProcessingInputCollection + Object containing dependencies to process. + """ + dependencies = super().pre_processing() + if self.data_level != "l2": + return dependencies + + try: + map_pivot_angle = MapDescriptor.from_string(self.descriptor).sensor + except ValueError: + # Not a map product, so there is no pivot angle to select inputs with + logger.info( + f"Not filtering inputs by pivot angle, {self.descriptor} is not a " + f"map descriptor." + ) + return dependencies + + if not isinstance(map_pivot_angle, int): + # A map of no particular pivot angle, e.g. "ilo-ena-h-sf-nsp-ram-..." + return dependencies + + filtered_dependencies = ProcessingInputCollection() + for processing_input in dependencies.get_processing_inputs(): + if ( + processing_input.input_type != ProcessingInputType.SCIENCE_FILE + or processing_input.source != "lo" + ): + filtered_dependencies.add(processing_input) + continue + + kept_filenames = [ + str(imap_file_path.filename) + for imap_file_path in processing_input.imap_file_paths + if self._input_at_pivot_angle(imap_file_path, map_pivot_angle) + ] + if kept_filenames: + filtered_dependencies.add(type(processing_input)(*kept_filenames)) + + return filtered_dependencies + def do_processing( self, dependencies: ProcessingInputCollection ) -> list[xr.Dataset]: diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index b9b3a57d3..414cab934 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -20,7 +20,6 @@ ) from imap_processing.ena_maps.utils.naming import DAYS_IN_MONTH, MapDescriptor from imap_processing.lo import lo_ancillary -from imap_processing.lo.constants import LoConstants from imap_processing.spice.time import ( et_to_datetime64, str_yyyymmdd_to_ttj2000ns, @@ -46,17 +45,18 @@ 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 + The inputs are expected to have already been filtered down to the pivot + angle of the map being made, which is done in pre-processing (see + ``cli.Lo.pre_processing``) so that the map records only the files it was + made from as its parents. When no pointing sets are left, a quickmap is produced instead. Parameters ---------- sci_dependencies : dict Dictionary of datasets needed for L2 data product creation in xarray Datasets. - Should contain "imap_lo_l1c_pset" key with list of pointing set datasets, - at any pivot angle. + Should contain "imap_lo_l1c_pset" key with list of pointing set datasets + taken at the pivot angle of the map. anc_dependencies : list List of ancillary file paths needed for L2 data product creation. Should include efficiency factor files. @@ -71,8 +71,7 @@ def lo_l2( ------- list[xr.Dataset] List containing the processed L2 dataset with rates, intensities, - and uncertainties, or the quickmap if there were no pointing sets at - the pivot angle of the map. + and uncertainties, or the quickmap if there were no pointing sets. Raises ------ @@ -87,9 +86,7 @@ def lo_l2( 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 - ) + psets = sci_dependencies.get("imap_lo_l1c_pset", []) if not psets: logger.info("No psets - trying to create a quickmap.") if start_date is None: @@ -140,52 +137,6 @@ 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]: @@ -195,8 +146,8 @@ def _lo_l2_quickmap( Parameters ---------- sci_dependencies : dict - Dictionary of the input datasets, keyed by logical source, at any - pivot angle. + Dictionary of the input datasets, keyed by logical source, taken at the + pivot angle of the map. descriptor : str The map descriptor to be produced (e.g., "l090-enansnbs-h-sf-nsp-ram-hae-6deg-6mo"). @@ -232,10 +183,9 @@ def _lo_l2_quickmap( ) 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}" + f"{len(datasets)} {logical_source} inputs are available at the pivot " + f"angle of {descriptor}" ) # No pointing sets available, so the epoch bounds are set using `start_date`. diff --git a/imap_processing/tests/lo/test_lo_l2.py b/imap_processing/tests/lo/test_lo_l2.py index 9c08f972a..b7b8c186b 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -25,7 +25,6 @@ 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, @@ -2927,25 +2926,6 @@ def test_project_pset_to_map_value_keys(self, minimal_pset_for_species): 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.""" @@ -2979,19 +2959,19 @@ def test_lo_l2_quickmap_requires_start_date(self): with pytest.raises(ValueError, match="start_date is required"): lo_l2({}, [], self.descriptor) - def test_quickmap_selects_inputs_by_pivot_angle(self, caplog): + def test_quickmap_reports_its_inputs(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)}), + xr.Dataset({"pivot_angle": xr.DataArray(90.1)}), ] } 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 + assert "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.""" diff --git a/imap_processing/tests/test_cli.py b/imap_processing/tests/test_cli.py index 9d79119ef..d4fa61df9 100644 --- a/imap_processing/tests/test_cli.py +++ b/imap_processing/tests/test_cli.py @@ -518,6 +518,51 @@ def test_lo_l2(mock_lo_pre_processing, mock_lo_l2, mock_instrument_dependencies) 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): + """Test that only inputs at the pivot angle of the map are kept.""" + valid_pset = "imap_lo_l1c_pset_20250415_v001.cdf" + invalid_pset = "imap_lo_l1c_pset_20250416_v001.cdf" + valid_de = "imap_lo_l1b_de_20250415-repoint00217_v001.cdf" + no_pivot_angle = "imap_lo_l1b_goodtimes_20250415-repoint00217_v001.cdf" + ancillary = "imap_lo_efficiency-factors_20250415_v001.csv" + + base_collection = ProcessingInputCollection( + ScienceInput(valid_pset, invalid_pset), + ScienceInput(valid_de), + ScienceInput(no_pivot_angle), + AncillaryInput(ancillary), + ) + mock_super_pre_processing.return_value = base_collection + mock_load_cdf.side_effect = [ + xr.Dataset({"pivot_angle": xr.DataArray(90.1)}), + # A neighbouring pivot angle, which belongs on its own map + xr.Dataset({"pivot_angle": xr.DataArray(75.0)}), + xr.Dataset({"pivot_angle": xr.DataArray(90.0)}), + xr.Dataset(), # No pivot angle at all + ] + + instrument = Lo( + "l2", + "l090-ena-h-sf-nsp-ram-hae-6deg-3mo", + base_collection.serialize(), + "20250415", + "20250715", + "v001", + False, + ) + result = instrument.pre_processing() + + # The goodtimes input had no pivot angle, so it drops out entirely + assert [ + [str(file_path.filename) for file_path in processing_input.imap_file_paths] + for processing_input in result.get_processing_inputs() + ] == [[valid_pset], [valid_de], [ancillary]] + # Ancillary files are not loaded to be checked for a pivot angle + assert mock_load_cdf.call_count == 4 + + @mock.patch("imap_processing.cli.quaternions.process_quaternions", autospec=True) def test_spacecraft(mock_spacecraft_l1a, mock_instrument_dependencies): """Test coverage for cli.Spacecraft class""" From 3c2f2572bea521a8719387230fc8ef6f7ce324db Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Sun, 26 Jul 2026 12:08:13 -0400 Subject: [PATCH 03/10] quickmap path gone (a quickmap is a regular ena map) --- imap_processing/cli.py | 4 +- imap_processing/lo/l2/lo_l2.py | 140 ++----------------------- imap_processing/tests/lo/test_lo_l2.py | 98 +---------------- imap_processing/tests/test_cli.py | 1 - 4 files changed, 13 insertions(+), 230 deletions(-) diff --git a/imap_processing/cli.py b/imap_processing/cli.py index 8185509f5..c37aa9e1c 100644 --- a/imap_processing/cli.py +++ b/imap_processing/cli.py @@ -1376,9 +1376,7 @@ def do_processing( dataset ) - datasets = lo_l2.lo_l2( - sci_dependencies, anc_dependencies, self.descriptor, self.start_date - ) + datasets = lo_l2.lo_l2(sci_dependencies, anc_dependencies, self.descriptor) return datasets diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index 414cab934..1a53c6a94 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -18,13 +18,9 @@ get_pset_directional_mask, interpolate_map_flux_to_helio_frame, ) -from imap_processing.ena_maps.utils.naming import DAYS_IN_MONTH, MapDescriptor +from imap_processing.ena_maps.utils.naming import MapDescriptor from imap_processing.lo import lo_ancillary -from imap_processing.spice.time import ( - et_to_datetime64, - str_yyyymmdd_to_ttj2000ns, - ttj2000ns_to_et, -) +from imap_processing.spice.time import et_to_datetime64, ttj2000ns_to_et logger = logging.getLogger(__name__) @@ -34,10 +30,7 @@ def lo_l2( - sci_dependencies: dict, - anc_dependencies: list, - descriptor: str, - start_date: str | None = None, + sci_dependencies: dict, anc_dependencies: list, descriptor: str ) -> list[xr.Dataset]: """ Process IMAP-Lo L1C data into L2 CDF data products. @@ -48,14 +41,13 @@ def lo_l2( The inputs are expected to have already been filtered down to the pivot angle of the map being made, which is done in pre-processing (see ``cli.Lo.pre_processing``) so that the map records only the files it was - made from as its parents. When no pointing sets are left, a quickmap is - produced instead. + made from as its parents. Parameters ---------- sci_dependencies : dict Dictionary of datasets needed for L2 data product creation in xarray Datasets. - Should contain "imap_lo_l1c_pset" key with list of pointing set datasets + Must contain "imap_lo_l1c_pset" key with list of pointing set datasets taken at the pivot angle of the map. anc_dependencies : list List of ancillary file paths needed for L2 data product creation. @@ -63,20 +55,15 @@ def lo_l2( 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, or the quickmap if there were no pointing sets. + and uncertainties. Raises ------ - ValueError - If a quickmap is required but no start_date was given. NotImplementedError If HEALPix map output is requested (only rectangular maps supported). """ @@ -86,12 +73,7 @@ def lo_l2( map_descriptor = MapDescriptor.from_string(descriptor) logger.info(f"Processing map for species: {map_descriptor.species}") - psets = sci_dependencies.get("imap_lo_l1c_pset", []) - 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) + psets = sci_dependencies["imap_lo_l1c_pset"] # Determine if corrections are needed and prepare oxygen data if required ( @@ -137,114 +119,6 @@ def lo_l2( return [dataset] -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, taken at the - pivot angle of the map. - 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(): - logger.info( - f"{len(datasets)} {logical_source} inputs are available at the pivot " - f"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 b7b8c186b..fef600bd0 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -1,6 +1,5 @@ """Comprehensive test suite for IMAP-Lo L2 data processing.""" -import logging from pathlib import Path from unittest.mock import Mock, patch @@ -9,7 +8,7 @@ import pytest import xarray as xr -from imap_processing.cdf.utils import load_cdf, write_cdf +from imap_processing.cdf.utils import load_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, @@ -25,7 +24,6 @@ SPIN_ANGLE_BIN_CENTERS, ) from imap_processing.lo.l2.lo_l2 import ( - _lo_l2_quickmap, _prepare_corrections, add_efficiency_factors_to_pset, calculate_all_rates_and_intensities, @@ -2588,19 +2586,14 @@ def test_lo_l2_integration_full( class TestErrorHandling: """Tests for error handling in various functions.""" - def test_lo_l2_no_pset_data(self, caplog): - """Test that a quickmap is made when no pointing set data is provided.""" + def test_lo_l2_no_pset_data(self): + """Test error 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 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}" + with pytest.raises(KeyError, match="imap_lo_l1c_pset"): + lo_l2(sci_dependencies, anc_dependencies, descriptor) def test_create_sky_map_healpix_not_supported(self, minimal_pset_for_species): """Test error when HEALPix map is requested.""" @@ -2924,84 +2917,3 @@ 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 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_reports_its_inputs(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(90.1)}), - ] - } - - with caplog.at_level(logging.INFO): - _lo_l2_quickmap(sci_dependencies, self.descriptor, "20260101") - - assert "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 d4fa61df9..3088d68f5 100644 --- a/imap_processing/tests/test_cli.py +++ b/imap_processing/tests/test_cli.py @@ -513,7 +513,6 @@ def test_lo_l2(mock_lo_pre_processing, mock_lo_l2, mock_instrument_dependencies) }, [], descriptor, - "20250415", ) mocks["mock_write_cdf"].assert_called_once_with(output_l2_dataset) From 2d8f660b0fd75c9685df2dae6ed75c7fa9c55ec7 Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Sun, 26 Jul 2026 21:14:42 -0400 Subject: [PATCH 04/10] quickmap implementation take 1, without any corrections --- imap_processing/cli.py | 86 +- imap_processing/lo/constants.py | 55 +- imap_processing/lo/l2/lo_l2.py | 1677 +++---------- imap_processing/tests/lo/test_lo_l2.py | 3135 +++--------------------- imap_processing/tests/test_cli.py | 94 +- 5 files changed, 891 insertions(+), 4156 deletions(-) diff --git a/imap_processing/cli.py b/imap_processing/cli.py index c37aa9e1c..e020c266e 100644 --- a/imap_processing/cli.py +++ b/imap_processing/cli.py @@ -1214,49 +1214,64 @@ class Lo(ProcessInstrument): """Process IMAP-Lo.""" @staticmethod - def _input_at_pivot_angle( - imap_file_path: imap_data_access.ImapFilePath, map_pivot_angle: int - ) -> bool: + def _pointings_at_pivot_angle( + dependencies: ProcessingInputCollection, map_pivot_angle: int + ) -> set[int]: """ - Check whether a Lo input file was taken at the pivot angle of a map. + Find the pointings that were taken at the pivot angle of a map. - Every Lo science product records the pivot angle it was taken at. Files - without a pivot angle can't be placed on a map and so are rejected. + A pointing's pivot angle is recorded in its goodtimes product, which is + also the product the other inputs of that pointing are selected by. Parameters ---------- - imap_file_path : imap_data_access.ImapFilePath - The Lo input file to check. + dependencies : ProcessingInputCollection + Object containing dependencies to process. map_pivot_angle : int The pivot angle [degrees] of the map being made. Returns ------- - bool - Whether the file belongs on the map. + set[int] + The repointings whose pivot angle is that of the map. """ - dataset = load_cdf(imap_file_path.construct_path()) - if "pivot_angle" not in dataset: - logger.info(f"Dropping {imap_file_path.filename} - no pivot angle.") - return False + at_pivot_angle = set() + for goodtimes_path in dependencies.get_file_paths( + source="lo", descriptor="goodtimes" + ): + goodtimes = load_cdf(goodtimes_path) + repointing = imap_data_access.ScienceFilePath( + goodtimes_path.name + ).repointing + if "pivot" not in goodtimes: + logger.info(f"Dropping {goodtimes_path.name} - no pivot angle.") + continue - pivot_angle = dataset["pivot_angle"].item() - if abs(pivot_angle - map_pivot_angle) >= LoConstants.PSET_PIVOT_ANGLE_TOLERANCE: - logger.info( - f"Dropping {imap_file_path.filename}, its pivot angle {pivot_angle} " - f"is not the {map_pivot_angle} degree pivot angle of the map." - ) - return False + pivot_angle = np.atleast_1d(goodtimes["pivot"].values)[0] + if ( + abs(pivot_angle - map_pivot_angle) + < LoConstants.PSET_PIVOT_ANGLE_TOLERANCE + ): + at_pivot_angle.add(repointing) + else: + logger.info( + f"Dropping repoint{repointing}, its pivot angle {pivot_angle} " + f"is not the {map_pivot_angle} degree pivot angle of the map." + ) - return True + return at_pivot_angle def pre_processing(self) -> ProcessingInputCollection: """ Complete pre-processing. Extends the base pre-processing by dropping, for map products, the Lo - science inputs that were not taken at the pivot angle of the map being - made. + science inputs of the pointings that were not taken at the pivot angle + of the map being made. A pointing is dropped whole: its goodtimes give + the pivot angle, and its other inputs go with them. + + Filtering here, rather than during processing, keeps the `Parents` + attribute of the produced map limited to the files it was made from. Returns ------- @@ -1281,6 +1296,8 @@ def pre_processing(self) -> ProcessingInputCollection: # A map of no particular pivot angle, e.g. "ilo-ena-h-sf-nsp-ram-..." return dependencies + at_pivot_angle = self._pointings_at_pivot_angle(dependencies, map_pivot_angle) + filtered_dependencies = ProcessingInputCollection() for processing_input in dependencies.get_processing_inputs(): if ( @@ -1293,7 +1310,7 @@ def pre_processing(self) -> ProcessingInputCollection: kept_filenames = [ str(imap_file_path.filename) for imap_file_path in processing_input.imap_file_paths - if self._input_at_pivot_angle(imap_file_path, map_pivot_angle) + if imap_file_path.repointing in at_pivot_angle ] if kept_filenames: filtered_dependencies.add(type(processing_input)(*kept_filenames)) @@ -1364,17 +1381,22 @@ def do_processing( datasets = lo_l1c.lo_l1c(data_dict, anc_dependencies) elif self.data_level == "l2": - 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") + sci_dependencies: dict[str, list[xr.Dataset]] = { + lo_l2.GOODTIMES: [], + lo_l2.BGRATES: [], + lo_l2.HISTRATES: [], + } + science_files = [] + for descriptor in ("goodtimes", "bgrates", "histrates"): + science_files += dependencies.get_file_paths( + source="lo", data_type="l1b", descriptor=descriptor + ) anc_dependencies = dependencies.get_file_paths(data_type="ancillary") - # Load every input, at every pivot angle, grouped by product. + # Load every pointing of the map window, grouped by product. for file in science_files: dataset = load_cdf(file) - sci_dependencies.setdefault(dataset.attrs["Logical_source"], []).append( - dataset - ) + sci_dependencies[dataset.attrs["Logical_source"]].append(dataset) datasets = lo_l2.lo_l2(sci_dependencies, anc_dependencies, self.descriptor) return datasets diff --git a/imap_processing/lo/constants.py b/imap_processing/lo/constants.py index 70cfba5c9..270417474 100644 --- a/imap_processing/lo/constants.py +++ b/imap_processing/lo/constants.py @@ -8,10 +8,14 @@ class LoConstants: """Constants for Lo which can be used across different levels.""" - # Absolute tolerance [degrees] for accepting a pset's pivot angle as + # Absolute tolerance [degrees] for accepting an input's pivot angle as # sufficiently close to the pivot angle of the map being made. PSET_PIVOT_ANGLE_TOLERANCE: float = 5.0 + # Empirical offset [degrees] added to the measured pivot angle when + # projecting a look direction onto the RAM direction. + PIVOT_RAM_OFFSET: float = 4.0 + # Ion species tracked. "H" is mandatory (and should be the first element); # any others for which we have histrates may be added here. ELEMS = ("H", "O") @@ -43,6 +47,55 @@ class LoConstants: RAM_HISTOGRAM_BINS: tuple[slice, ...] = (slice(0, 20), slice(50, 60)) ANTI_RAM_HISTOGRAM_BINS: tuple[slice, ...] = (slice(20, 50),) + # The following are indexed by ESA level (0-indexed, ESA level = index + 1). + # The 8th entry is the virtual E8 channel, unused by the map. + ESA_ENERGY: ClassVar[list[float]] = [ + 0.016, + 0.030, + 0.056, + 0.106, + 0.200, + 0.405, + 0.787, + 1.821, + ] + GEO_FACTOR: ClassVar[list[float]] = [ + 7.0e-5, + 7.9e-5, + 9.7e-5, + 11.2e-5, + 14.0e-5, + 17.7e-5, + 22.5e-5, + 6.721e-5, + ] + GEO_FACTOR_ERR: ClassVar[list[float]] = [ + 4.9e-5, + 5.5e-5, + 6.8e-5, + 3.0e-5, + 4.5e-5, + 2.0e-5, + 1.4e-5, + 6.721e-5, + ] + + # GEO_FACTOR/GEO_FACTOR_ERR above are the raw, pre-recalibration values; the + # map multiplies them by GEO_FACTOR_SCALE, and derives the asymmetric + # upper/lower G-factor bounds using the two scale factors below. + GEO_FACTOR_SCALE: float = 0.63529412 + GEO_FACTOR_SCALE_UPPER: float = 1.57407407 + GEO_FACTOR_SCALE_LOWER: float = 0.36728395 + + # Half-widths [keV] of the ESA energy passbands, by ESA level, for the two + # ESA modes. NOTE: From an e-mail from Nathan on 2025-09-11 (converted to keV). + ESA_ENERGY_DELTA: ClassVar[dict[int, list[float]]] = { + # esa_mode 0, HiRes + 0: [0.00543, 0.01002, 0.01861, 0.03331, 0.06498, 0.13164, 0.26235], + # esa_mode 1, HiThr + 1: [0.00881, 0.01604, 0.02850, 0.05313, 0.10560, 0.21967, 0.41360], + } + # Nominal background rates [counts/s] for each species BG_RATES: ClassVar[dict[str, float]] = {"H": 0.0014925, "O": 0.000136635} # When no exposure is available, scale the nominal rate down as a conservative diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index 1a53c6a94..064581178 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -1,29 +1,28 @@ """IMAP-Lo L2 data processing.""" import logging -from pathlib import Path from typing import cast import numpy as np -import pandas as pd import xarray as xr -from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes -from imap_processing.ena_maps import ena_maps -from imap_processing.ena_maps.ena_maps import AbstractSkyMap, RectangularSkyMap -from imap_processing.ena_maps.utils.corrections import ( - PowerLawFluxCorrector, - apply_compton_getting_correction, - calculate_ram_mask, - get_pset_directional_mask, - interpolate_map_flux_to_helio_frame, -) +from imap_processing.ena_maps.ena_maps import RectangularSkyMap from imap_processing.ena_maps.utils.naming import 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 as c # noqa: N813 +from imap_processing.lo.l1c.lo_l1c import compute_pointing_directions +from imap_processing.spice.geometry import ( + SpiceFrame, + get_spacecraft_to_instrument_spin_phase_offset, +) +from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_met logger = logging.getLogger(__name__) +# The L1B products a map is built from, one set per pointing. +GOODTIMES = "imap_lo_l1b_goodtimes" +BGRATES = "imap_lo_l1b_bgrates" +HISTRATES = "imap_lo_l1b_histrates" + # ============================================================================= # MAIN ENTRY POINT # ============================================================================= @@ -33,10 +32,12 @@ def lo_l2( sci_dependencies: dict, anc_dependencies: list, descriptor: str ) -> list[xr.Dataset]: """ - Process IMAP-Lo L1C data into L2 CDF data products. + Process IMAP-Lo L1B data into an L2 sky map. - 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. + A map accumulates the histogram counts and exposure of every pointing in + its window, binned by the sky direction each spin-angle bin was looking in, + and converts the accumulated counts into an intensity with the instrument's + geometric factors. The inputs are expected to have already been filtered down to the pivot angle of the map being made, which is done in pre-processing (see @@ -46,1413 +47,521 @@ def lo_l2( 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 - taken at the pivot angle of the map. + Dictionary of the input datasets, keyed by logical source, each a list + of datasets covering the pointings of the map window. Must contain + ``"imap_lo_l1b_goodtimes"``, ``"imap_lo_l1b_bgrates"`` and + ``"imap_lo_l1b_histrates"``. anc_dependencies : list - List of ancillary file paths needed for L2 data product creation. - Should include efficiency factor files. + List of ancillary file paths. Unused, the calibration constants of the + map live in ``LoConstants``. descriptor : str The map descriptor to be produced - (e.g., "ilo90-ena-h-sf-nsp-full-hae-6deg-3mo"). + (e.g., "l090-ena-h-sf-nsp-ram-hae-6deg-3mo"). Returns ------- list[xr.Dataset] - List containing the processed L2 dataset with rates, intensities, - and uncertainties. + List containing the processed L2 map. Raises ------ NotImplementedError - If HEALPix map output is requested (only rectangular maps supported). + If a HEALPix map is requested (only rectangular maps supported for Lo), + or if the map is of a species other than hydrogen. """ logger.info("Starting IMAP-Lo L2 processing pipeline") - # 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 = sci_dependencies["imap_lo_l1c_pset"] - - # Determine if corrections are needed and prepare oxygen data if required - ( - sputtering_correction, - bootstrap_correction, - flux_correction, - o_map_dataset, - flux_factors, - cg_correction, - ) = _prepare_corrections( - map_descriptor, descriptor, sci_dependencies, anc_dependencies - ) - - logger.info("Step 1: Loading ancillary data") - efficiency_data = load_efficiency_data(anc_dependencies) + # The geometric factors in LoConstants are hydrogen only. + if map_descriptor.species != "h": + raise NotImplementedError( + f"Cannot make a map of species {map_descriptor.species} for " + f"{descriptor}. Only hydrogen geometric factors are defined." + ) - logger.info(f"Step 2: Creating sky map from {len(psets)} pointing sets") - sky_map = create_sky_map_from_psets( - psets, map_descriptor, efficiency_data, cg_correction - ) + sky_map = map_descriptor.to_empty_map() + if not isinstance(sky_map, RectangularSkyMap): + raise NotImplementedError("HEALPix map output not supported for Lo") - logger.info("Step 3: Converting to dataset and adding geometric factors") - dataset = sky_map.to_dataset() - dataset = add_geometric_factors(dataset, map_descriptor.species) - - logger.info("Step 4: Calculating rates and intensities") - dataset = calculate_all_rates_and_intensities( - dataset, - sputtering_correction=sputtering_correction, - bootstrap_correction=bootstrap_correction, - flux_correction=flux_correction, - o_map_dataset=o_map_dataset, - flux_factors=flux_factors, - cg_correction=cg_correction, - ) + pointings = _group_inputs_by_pointing(sci_dependencies) + logger.info(f"Building {descriptor} from {len(pointings)} pointings") + + shape = (c.N_ESA_LEVELS, sky_map.num_points) + counts = np.zeros(shape) + exposure = np.zeros(shape) + # Background is a rate per ESA level per pointing, so it is accumulated + # weighted by exposure and divided by the total exposure at the end. + bg_rate_exposure = np.zeros(shape) + esa_mode = 0 + + for repointing, (goodtimes, bgrates, histrates) in sorted(pointings.items()): + logger.debug(f"Accumulating {repointing}") + esa_mode = _get_esa_mode(histrates) + _accumulate_pointing( + goodtimes, + bgrates, + histrates, + sky_map, + map_descriptor, + counts, + exposure, + bg_rate_exposure, + ) - logger.info("Step 5: Finalizing dataset with attributes") - dataset = cast(RectangularSkyMap, sky_map).build_cdf_dataset( - instrument="lo", level="l2", descriptor=descriptor, external_map_dataset=dataset + variables = _calculate_rates_and_intensities( + counts, exposure, bg_rate_exposure, esa_mode ) + dataset = _build_map_dataset(sky_map, variables, esa_mode) logger.info("IMAP-Lo L2 processing pipeline completed successfully") - return [dataset] - - -def _prepare_corrections( - map_descriptor: MapDescriptor, - descriptor: str, - sci_dependencies: dict, - anc_dependencies: list, -) -> tuple[bool, bool, bool, xr.Dataset | None, Path | None, bool]: - """ - Determine what corrections are needed and prepare oxygen dataset if required. - - This helper function encapsulates the logic for determining when sputtering - and bootstrap corrections should be applied, and handles the creation of - the oxygen dataset needed for sputtering corrections. - - Parameters - ---------- - map_descriptor : MapDescriptor - The parsed map descriptor containing species and data type information. - descriptor : str - The original descriptor string for creating the oxygen variant. - sci_dependencies : dict - Dictionary of datasets needed for L2 data product creation. - anc_dependencies : list - List of ancillary file paths. - - Returns - ------- - tuple[bool, bool, bool, xr.Dataset | None, Path | None, bool] - A tuple containing: - - sputtering_correction: Whether to apply sputtering corrections - - bootstrap_correction: Whether to apply bootstrap corrections - - flux_correction: Whether to apply flux corrections - - o_map_dataset: Oxygen dataset if needed, None otherwise - - flux_factors: Path to flux factors ancillary file if needed, - None otherwise - - cg_correction: Whether to apply CG correction to the dataset. - """ - # Default values - no corrections needed - sputtering_correction = False - bootstrap_correction = False - flux_correction = False - o_map_dataset = None - flux_factors: None | Path = None - - # Sputtering and bootstrap corrections are only applied to hydrogen ENA data - # Guard against recursion: don't process oxygen for oxygen maps - if ( - map_descriptor.species == "h" - and map_descriptor.principal_data == "ena" - and "-o-" not in descriptor - ): # Safety check to prevent infinite recursion - logger.info("Creating map for oxygen for sputtering corrections") - o_descriptor = descriptor.replace("-h-", "-o-") - o_map_dataset = lo_l2(sci_dependencies, anc_dependencies, o_descriptor)[0] - sputtering_correction = True - bootstrap_correction = True - - if "raw" not in map_descriptor.principal_data: - flux_correction = True - try: - flux_factors = next( - x for x in anc_dependencies if "esa-eta-fit-factors" in str(x) - ) - except StopIteration: - raise ValueError( - "No flux correction factor file found in ancillary dependencies" - ) from None - - cg_correction = True if map_descriptor.frame_descriptor == "hf" else False - - return ( - sputtering_correction, - bootstrap_correction, - flux_correction, - o_map_dataset, - flux_factors, - cg_correction, - ) + return [ + sky_map.build_cdf_dataset( + instrument="lo", + level="l2", + descriptor=descriptor, + external_map_dataset=dataset, + ) + ] # ============================================================================= -# SETUP AND INITIALIZATION HELPERS +# INPUT HANDLING # ============================================================================= -def load_efficiency_data(anc_dependencies: list) -> pd.DataFrame: - """ - Load efficiency factor data from ancillary files. - - Parameters - ---------- - anc_dependencies : list - List of ancillary file paths to search for efficiency factor files. - - Returns - ------- - pd.DataFrame - Concatenated efficiency factor data from all matching files. - Returns empty DataFrame if no efficiency files found. +def _group_inputs_by_pointing(sci_dependencies: dict) -> dict[str, tuple]: """ - efficiency_files = [ - anc_file - for anc_file in anc_dependencies - if "efficiency-factor" in str(anc_file) - ] + Group the L1B inputs into the (goodtimes, bgrates, histrates) of a pointing. - if not efficiency_files: - logger.warning("No efficiency factor files found in ancillary dependencies") - return pd.DataFrame() - - logger.debug(f"Loading {len(efficiency_files)} efficiency factor files") - return pd.concat( - [lo_ancillary.read_ancillary_file(anc_file) for anc_file in efficiency_files], - ignore_index=True, - ) - - -def load_sputter_correction_data( - source_species: str, target_species: str -) -> pd.DataFrame: - """ - Load sputter correction factors from an ancillary file. + Each input product records the pointing it covers in its ``Repointing`` + global attribute. Pointings missing any of the three products cannot be + mapped and are dropped. Parameters ---------- - source_species : str - The species doing the sputtering (e.g. "o" for oxygen). - target_species : str - The species being corrected (e.g. "h" for hydrogen). - - Returns - ------- - pd.DataFrame - Rows matching the given species pair, sorted ascending by esa_step, - with columns: source_species, target_species, esa_step, - sputter_factor, sputter_factor_uncertainty. - """ - anc_path = Path(__file__).parent.parent / "ancillary_data" - sputter_files = sorted(anc_path.glob("*sputter-correction-factors*")) - - if not sputter_files: - raise ValueError("No sputter correction files found") - - df = pd.concat( - [lo_ancillary.read_ancillary_file(f) for f in sputter_files], - ignore_index=True, - ) - mask = (df["source_species"] == source_species) & ( - df["target_species"] == target_species - ) - result = df[mask].sort_values("esa_step").reset_index(drop=True) - return result - - -def load_bootstrap_correction_data() -> pd.DataFrame: - """ - Load bootstrap correction factors from an ancillary file. + sci_dependencies : dict + Dictionary of the input datasets, keyed by logical source. Returns ------- - pd.DataFrame - Bootstrap correction factors with columns: esa_step_i, esa_step_k, - bootstrap_factor. Indices are 1-based ESA step numbers where esa_step_k=8 - refers to the virtual E8 channel. - """ - anc_path = Path(__file__).parent.parent / "ancillary_data" - bootstrap_files = sorted(anc_path.glob("*bootstrap-correction-factors*")) + dict[str, tuple] + The (goodtimes, bgrates, histrates) datasets of each pointing, keyed by + repointing. - if not bootstrap_files: - raise ValueError("No bootstrap correction factor files found") + Raises + ------ + KeyError + If any of the three required products is missing entirely. + """ + by_pointing: dict[str, dict[str, xr.Dataset]] = {} + for logical_source in (GOODTIMES, BGRATES, HISTRATES): + for dataset in sci_dependencies[logical_source]: + repointing = dataset.attrs.get("Repointing", "") + by_pointing.setdefault(repointing, {})[logical_source] = dataset + + pointings = {} + for repointing, products in by_pointing.items(): + missing = {GOODTIMES, BGRATES, HISTRATES} - set(products) + if missing: + logger.warning(f"Dropping {repointing}, it has no {sorted(missing)}") + continue + pointings[repointing] = ( + products[GOODTIMES], + products[BGRATES], + products[HISTRATES], + ) - return pd.concat( - [lo_ancillary.read_ancillary_file(f) for f in bootstrap_files], - ignore_index=True, - ) + return pointings -def finalize_dataset(dataset: xr.Dataset, descriptor: str) -> xr.Dataset: +def _get_esa_mode(histrates: xr.Dataset) -> int: """ - Add attributes and perform final dataset preparation. + Read the ESA mode of a pointing, defaulting to HiRes. Parameters ---------- - dataset : xr.Dataset - The dataset to finalize with attributes. - descriptor : str - The descriptor for this map dataset. + histrates : xr.Dataset + The L1B histogram rates of the pointing. Returns ------- - xr.Dataset - The finalized dataset with all attributes added. + int + The ESA mode, 0 for HiRes and 1 for HiThr. """ - # Initialize the attribute manager - attr_mgr = ImapCdfAttributes() - attr_mgr.add_instrument_global_attrs(instrument="lo") - attr_mgr.add_instrument_variable_attrs(instrument="enamaps", level="l2-common") - attr_mgr.add_instrument_variable_attrs(instrument="enamaps", level="l2-rectangular") - - # Add global and variable attributes - dataset.attrs.update(attr_mgr.get_global_attributes("imap_lo_l2_enamap")) - - # Our global attributes have placeholders for descriptor - # so iterate through here and fill that in with the map-specific descriptor - for key in ["Data_type", "Logical_source", "Logical_source_description"]: - dataset.attrs[key] = dataset.attrs[key].format(descriptor=descriptor) - for var in dataset.data_vars: - try: - dataset[var].attrs = attr_mgr.get_variable_attributes(var) - except KeyError: - # If no attributes found, try without schema validation - try: - dataset[var].attrs = attr_mgr.get_variable_attributes( - var, check_schema=False - ) - except KeyError: - logger.warning(f"No attributes found for variable {var}") - - return dataset + if "esa_mode" not in histrates: + return 0 + return int(np.atleast_1d(histrates["esa_mode"].values)[0]) # ============================================================================= -# SKY MAP CREATION PIPELINE +# SKY MAP ACCUMULATION # ============================================================================= -def create_sky_map_from_psets( - psets: list[xr.Dataset], +def _accumulate_pointing( + goodtimes: xr.Dataset, + bgrates: xr.Dataset, + histrates: xr.Dataset, + sky_map: RectangularSkyMap, map_descriptor: MapDescriptor, - efficiency_data: pd.DataFrame, - cg_correct: bool, -) -> AbstractSkyMap: + counts: np.ndarray, + exposure: np.ndarray, + bg_rate_exposure: np.ndarray, +) -> None: """ - Create a sky map by processing all pointing sets. + Add one pointing's counts and exposure to the map accumulators. Parameters ---------- - psets : list[xr.Dataset] - List of pointing set datasets to process. + goodtimes : xr.Dataset + The L1B goodtimes of the pointing, giving its pivot angle and the + good-time windows its histograms are accepted within. + bgrates : xr.Dataset + The L1B background rates of the pointing, one rate per ESA level. + histrates : xr.Dataset + The L1B histogram rates of the pointing, giving the counts and exposure + of each spin-angle bin. + sky_map : RectangularSkyMap + The map being built, used for its pixel grid. map_descriptor : MapDescriptor - Map descriptor object defining the projection and binning. - efficiency_data : pd.DataFrame - Efficiency factor data for correcting counts. - cg_correct : bool - Whether to apply the CG correction to each PSET. + The parsed descriptor of the map being made. + counts : np.ndarray + Accumulator of shape (esa level, pixel), modified in place. + exposure : np.ndarray + Accumulator of shape (esa level, pixel), modified in place. + bg_rate_exposure : np.ndarray + Accumulator of shape (esa level, pixel), modified in place. + """ + species = map_descriptor.species + pivot_angle = float(np.atleast_1d(goodtimes["pivot"].values)[0]) + gt_start = np.atleast_1d(goodtimes["gt_start_met"].values) + gt_end = np.atleast_1d(goodtimes["gt_end_met"].values) + + histogram_met = ttj2000ns_to_met(histrates["epoch"].values) + in_goodtime = np.any( + (histogram_met[:, np.newaxis] >= gt_start) + & (histogram_met[:, np.newaxis] <= gt_end), + axis=1, + ) + if not in_goodtime.any(): + logger.warning("No histogram epochs fall within the good-time windows.") + return + + pointing_counts = histrates[f"{species}_counts"].values[in_goodtime].sum(axis=0) + pointing_exposure = histrates["exposure_time_6deg"].values[in_goodtime].sum(axis=0) + background_rates = np.atleast_2d(bgrates[f"{species}_background_rates"].values)[0] + + spin_angles = _dps_spin_angles() + # The whole pointing is projected from the middle of its good times, which + # is where the despun frame is sampled. + epoch = met_to_ttj2000ns((gt_start.min() + gt_end.max()) / 2.0) + az_el = compute_pointing_directions( + epoch, + pivot_angle, + spin_angles=spin_angles, + off_angles=np.array([0.0]), + to_frame=cast(SpiceFrame, map_descriptor.map_spice_coord_frame), + ) + # The single boresight off-angle is squeezed out of the frame transform, so + # the directions come back as (spin angle, lon/lat). + az_el = np.asarray(az_el).reshape(spin_angles.size, 2) + pixels = _pixel_indices(sky_map, az_el[:, 0], az_el[:, 1]) + + keep = _spin_phase_mask(spin_angles, pivot_angle, map_descriptor.spin_phase) + if not keep.any(): + return + + # np.add.at accumulates repeated pixels, which is what happens whenever + # several spin-angle bins land in the same map pixel. + np.add.at(counts, (slice(None), pixels[keep]), pointing_counts[:, keep]) + np.add.at(exposure, (slice(None), pixels[keep]), pointing_exposure[:, keep]) + np.add.at( + bg_rate_exposure, + (slice(None), pixels[keep]), + background_rates[:, np.newaxis] * pointing_exposure[:, keep], + ) + + sky_map.min_epoch = min(sky_map.min_epoch, int(met_to_ttj2000ns(gt_start.min()))) + sky_map.max_epoch = max(sky_map.max_epoch, int(met_to_ttj2000ns(gt_end.max()))) + + +def _dps_spin_angles() -> np.ndarray: + """ + Get the despun-frame azimuth of each histogram spin-angle bin center. + + The L1B histogram spin bins are hardware spin-phase bins referenced to the + spacecraft spin pulse, NOT the instrument (DPS) spin angle. A bin center is + converted to the IMAP_DPS azimuth by adding the spacecraft to instrument + spin-phase offset, exactly as the L1B star-sensor product does. Returns ------- - AbstractSkyMap - The populated sky map with projected data from all pointing sets. - - Raises - ------ - NotImplementedError - If HEALPix map output is requested (only rectangular maps supported). - """ - # Initialize the output map - output_map = map_descriptor.to_empty_map() - - if not isinstance(output_map, RectangularSkyMap): - raise NotImplementedError("HEALPix map output not supported for Lo") - - logger.debug(f"Processing {len(psets)} pointing sets") - # Process each pointing set - for i, pset in enumerate(psets): - logger.debug(f"Processing pointing set {i + 1}/{len(psets)}") - processed_pset = process_single_pset( - pset, - efficiency_data, - map_descriptor.species, - cg_correct, - ) - directional_mask = get_pset_directional_mask( - processed_pset, map_descriptor.spin_phase - ) - project_pset_to_map(processed_pset, output_map, directional_mask, cg_correct) - - return output_map - - -def process_single_pset( - pset: xr.Dataset, - efficiency_data: pd.DataFrame, - species: str, - cg_correct: bool = False, -) -> xr.Dataset: - """ - Process a single pointing set for projection to the sky map. - - Parameters - ---------- - pset : xr.Dataset - Single pointing set dataset to process. - efficiency_data : pd.DataFrame - Efficiency factor data for correcting counts. - species : str - The species to process (e.g., "h", "o"). - cg_correct : bool - Whether to apply the CG correction to each PSET. A value of True will - cause the pre-projection Compton Getting Correction to be applied to - the PSET data. - - Returns - ------- - xr.Dataset - Processed pointing set ready for projection with efficiency corrections applied. + np.ndarray + The IMAP_DPS azimuth [degrees] of each of the histogram spin bins. """ - # Step 1: Normalize coordinate system - pset_processed = normalize_pset_coordinates(pset, species) - - # Step 2: Add efficiency factors - pset_processed = add_efficiency_factors_to_pset(pset_processed, efficiency_data) - - # Step 3: Calculate efficiency-corrected quantities - pset_processed = calculate_efficiency_corrected_quantities(pset_processed) - - # Step 4: Optionally apply CG correction and calculate ram-mask - if cg_correct: - # NOTE: Heliospheric frame energy selection for CG correction - # The heliospheric (HF) energies passed to the CG correction algorithm - # could in principle be completely different from the ESA central energies. - # However, for Lo, the instrument team has chosen to use the same HF - # energies as the ESA central energies (from the geometric factor files). - # This decision aligns the energy grid between the spacecraft frame and - # heliospheric frame representations. - - # Convert energy coordinate from keV to eV for CG correction - # (energy coordinate was set in normalize_pset_coordinates in keV) - energy_values_ev: xr.DataArray = pset_processed["energy"] * 1000.0 - pset_processed = apply_compton_getting_correction( - pset_processed, energy_values_ev - ) - # Prepare energy_sc for exposure time weighted projection - pset_processed["energy_sc_exposure_factor"] = ( - pset_processed["energy_sc"] * pset_processed["exposure_factor"] - ) - - # Always calculate ram-mask to identify ram/anti-ram bins - pset_processed = calculate_ram_mask(pset_processed) + bin_width = 360.0 / c.N_SPIN_ANGLE_BINS + bin_centers = (np.arange(c.N_SPIN_ANGLE_BINS) + 0.5) * bin_width + offset = get_spacecraft_to_instrument_spin_phase_offset(SpiceFrame.IMAP_LO) * 360.0 + return np.mod(bin_centers + offset, 360.0) - return pset_processed - - -def normalize_pset_coordinates(pset: xr.Dataset, species: str) -> xr.Dataset: - """ - Normalize pointing set coordinates to match the output map. - Parameters - ---------- - pset : xr.Dataset - Input pointing set dataset with potentially mismatched coordinates. - species : str - The species to process (e.g., "h", "o"). - - Returns - ------- - xr.Dataset - Pointing set with normalized energy coordinates and dimension names. +def _pixel_indices( + sky_map: RectangularSkyMap, longitude: np.ndarray, latitude: np.ndarray +) -> np.ndarray: """ - # Load true energy values for this species (in keV, matching map convention) - # TODO: Figure out how to handle esa_mode properly - if "esa_mode" in pset: - esa_mode = pset["esa_mode"].values[0] - else: - # Default to mode 0 if not available (HiRes mode) - esa_mode = 0 - gf_dataset = reduce_geometric_factor_dataset(species, esa_mode=esa_mode) - - # Ensure consistent energy coordinates (maps want energy not esa_energy_step) - pset_renamed = pset.rename_dims({"esa_energy_step": "energy"}) - - # Drop the esa_energy_step coordinate first to avoid conflicts - pset_renamed = pset_renamed.drop_vars("esa_energy_step") - - # Assign TRUE energy values as coordinates (in keV, matching map convention) - pset_renamed = pset_renamed.assign_coords(energy=gf_dataset["Cntr_E"].values) - - # Rename the variables in the pset for projection to the map - # L2 wants different variable names than l1c - rename_map = { - "exposure_time": "exposure_factor", - f"{species}_counts": "counts", - f"{species}_background_rates": "bg_rate", - f"{species}_background_rates_stat_uncert": "bg_rate_stat_uncert", - f"{species}_background_rates_sys_err": "bg_rate_sys_err", - } - pset_renamed = pset_renamed.rename_vars(rename_map) - - return pset_renamed + Get the map pixel each sky direction falls in. - -def add_efficiency_factors_to_pset( - pset: xr.Dataset, efficiency_data: pd.DataFrame -) -> xr.Dataset: - """ - Add efficiency factors to the pointing set based on observation date. + A rectangular map stores its pixels as a 1D array raveled from the + (azimuth, elevation) grid, elevation varying fastest. Parameters ---------- - pset : xr.Dataset - Pointing set dataset to add efficiency factors to. - efficiency_data : pd.DataFrame - Efficiency factor data containing date-indexed efficiency values. + sky_map : RectangularSkyMap + The map being built. + longitude : np.ndarray + Longitudes [degrees] in the map's frame. + latitude : np.ndarray + Latitudes [degrees] in the map's frame. Returns ------- - xr.Dataset - Pointing set with efficiency factors added as new data variable. - - Raises - ------ - ValueError - If no efficiency factor found for the pointing set observation date. + np.ndarray + The pixel index of each direction. """ - if efficiency_data.empty: - # If no efficiency data, create unity efficiency - logger.warning("No efficiency data available, using unity efficiency") - pset["efficiency"] = xr.DataArray(np.ones(7), dims=["energy"]) - return pset - - # Convert the epoch to datetime64 - date = et_to_datetime64(ttj2000ns_to_et(pset["epoch"].values[0])) - # The efficiency file only has date as YYYYDDD, so drop the time for this - date = date.astype("M8[D]") # Convert to date only (no time) - - ef_df = efficiency_data[efficiency_data["Date"] == date] - if ef_df.empty: - raise ValueError(f"No efficiency factor found for pset date {date}") - - efficiency_values = ef_df[ - [ - "E-Step1_eff", - "E-Step2_eff", - "E-Step3_eff", - "E-Step4_eff", - "E-Step5_eff", - "E-Step6_eff", - "E-Step7_eff", - ] - ].values[0] - - pset["efficiency"] = xr.DataArray( - efficiency_values, - dims=["energy"], - ) - logger.debug(f"Applied efficiency factors for date {date}") - return pset - - -def calculate_efficiency_corrected_quantities( - pset: xr.Dataset, -) -> xr.Dataset: - """ - Calculate efficiency-corrected quantities for each particle type. + spacing = sky_map.spacing_deg + num_azimuth, num_elevation = sky_map.binning_grid_shape - Parameters - ---------- - pset : xr.Dataset - Pointing set with efficiency factors applied. - - Returns - ------- - xr.Dataset - Pointing set with efficiency-corrected count variables added. - """ - # counts / efficiency - pset["counts_over_eff"] = pset["counts"] / pset["efficiency"] - # counts / efficiency**2 (for variance propagation) - pset["counts_over_eff_squared"] = pset["counts"] / (pset["efficiency"] ** 2) - - # background * exposure_factor for weighted average - pset["bg_rate_exposure_factor"] = pset["bg_rate"] * pset["exposure_factor"] - # background_uncertainty ** 2 * exposure_factor ** 2 - pset["bg_rate_stat_uncert_exposure_factor2"] = ( - pset["bg_rate_stat_uncert"] ** 2 * pset["exposure_factor"] ** 2 - ) - # background systematic * exposure_factor for weighted average. - pset["bg_rate_sys_err_exposure_factor"] = ( - pset["bg_rate_sys_err"] * pset["exposure_factor"] + azimuth_index = np.clip( + (np.mod(longitude, 360.0) // spacing).astype(int), 0, num_azimuth - 1 ) - - return pset - - -def project_pset_to_map( - pset: xr.Dataset, - output_map: AbstractSkyMap, - directional_mask: xr.DataArray, - cg_correct: bool = False, -) -> None: - """ - Project pointing set data to the output map. - - Parameters - ---------- - pset : xr.Dataset - Processed pointing set ready for projection. - output_map : AbstractSkyMap - Target sky map to receive the projected data. - directional_mask : xr.DataArray - Boolean mask indicating which PSET bins to use for projection. This is - how ram/anti-ram bins are removed depending on the descriptor spin phase. - cg_correct : bool - Whether the CG correction is being applied. If set to True, "energy_sc" - is added to the list of variables to be projected. - - Returns - ------- - None - Function modifies output_map in place. - """ - # Define base quantities to project - value_keys = [ - "exposure_factor", - "counts", - "counts_over_eff", - "counts_over_eff_squared", - "bg_rate", - "bg_rate_stat_uncert", - "bg_rate_sys_err", - "bg_rate_exposure_factor", - "bg_rate_stat_uncert_exposure_factor2", - "bg_rate_sys_err_exposure_factor", - ] - if cg_correct: - value_keys.append("energy_sc_exposure_factor") - - # Create LoPointingSet and project to map - lo_pset = ena_maps.LoPointingSet(pset) - output_map.project_pset_values_to_map( - pointing_set=lo_pset, - value_keys=value_keys, - index_match_method=ena_maps.IndexMatchMethod.PUSH, - pset_valid_mask=directional_mask, + elevation_index = np.clip( + ((latitude + 90.0) // spacing).astype(int), 0, num_elevation - 1 ) - logger.debug(f"Projected {len(value_keys)} quantities to sky map") - - -# ============================================================================= -# GEOMETRIC FACTORS -# ============================================================================= - - -def add_geometric_factors(dataset: xr.Dataset, species: str) -> xr.Dataset: - """ - Add geometric factors to the sky map after projection. + return azimuth_index * num_elevation + elevation_index - Parameters - ---------- - dataset : xr.Dataset - Sky map dataset to add geometric factors to. - species : str - The species to process (only "h" and "o" have geometric factors). - Returns - ------- - xr.Dataset - Dataset with geometric factor variables added for the specified species. +def _spin_phase_mask( + spin_angles: np.ndarray, pivot_angle: float, spin_phase: str +) -> np.ndarray: """ - # Only add geometric factors for hydrogen and oxygen - if species not in ["h", "o"]: - logger.warning(f"No geometric factors to add for species: {species}") - return dataset - - logger.info(f"Loading and applying geometric factors for species: {species}") - - # Initialize geometric factor variables - dataset = initialize_geometric_factor_variables(dataset) - - # Populate geometric factors for each energy step - dataset = populate_geometric_factors(dataset, species) - - return dataset - + Get the spin-angle bins belonging on a map of the given spin phase. -def load_geometric_factor_data(species: str) -> pd.DataFrame: - """ - Load geometric factor data for the specified species. + A bin's RAM projection factor is ``sin(pivot) * sin(spin angle)``, positive + looking into the RAM direction and negative looking away from it. Parameters ---------- - species : str - The species to load geometric factors for ("h" or "o"). + spin_angles : np.ndarray + The IMAP_DPS azimuth [degrees] of each spin-angle bin. + pivot_angle : float + The pivot angle [degrees] of the pointing. + spin_phase : str + The spin phase of the map, "ram", "anti" or "full". Returns ------- - pd.DataFrame - Geometric factor dataframe for the specified species. + np.ndarray + Boolean mask of the bins to keep. Raises ------ ValueError - If species is not "h" or "o". + If the spin phase is not one of "ram", "anti" or "full". """ - if species not in ["h", "o"]: + if spin_phase == "full": + return np.ones(spin_angles.size, dtype=bool) + if spin_phase not in ("ram", "anti"): raise ValueError( - f"Geometric factors only available for 'h' and 'o', got '{species}'" + f"Invalid spin phase: {spin_phase}. Must be 'ram', 'anti' or 'full'." ) - anc_path = Path(__file__).parent.parent / "ancillary_data" - - if species == "h": - gf_file = sorted(anc_path.glob("*hydrogen-geometric-factor*"))[-1] - else: # species == "o" - gf_file = sorted(anc_path.glob("*oxygen-geometric-factor*"))[-1] - - return lo_ancillary.read_ancillary_file(gf_file) - - -def reduce_geometric_factor_dataset(species: str, esa_mode: int) -> xr.Dataset: - """ - Get geometric factor data as xarray Dataset for a specific species and ESA mode. - - This helper function loads geometric factor data, filters by ESA mode, converts - to xarray, and selects all 7 energy steps for vectorized operations. - - Parameters - ---------- - species : str - The species to load geometric factors for ("h" or "o"). - esa_mode : int - ESA mode (0 for HiRes, 1 for HiThr). - - Returns - ------- - xarray.Dataset - Geometric factor data indexed by Observed_E-Step (1-7), containing all - columns from the geometric factor CSV file. - """ - # Load geometric factor data for this species - gf_data = load_geometric_factor_data(species) - - # Filter for the specific ESA mode - if "esa_mode" in gf_data.columns: - gf_data = gf_data[gf_data["esa_mode"] == esa_mode].copy() - - # Convert to xarray Dataset indexed by energy step for vectorized selection - gf_ds = gf_data.set_index("Observed_E-Step").to_xarray() - - # Lo Instrument team: Use only geometric factors where - # incident_E-Step == Observed_E-Step - gf_ds = gf_ds.where(gf_ds["incident_E-Step"] == gf_ds["Observed_E-Step"], drop=True) - - # Select energy steps 1-7 and return - return gf_ds.sel({"Observed_E-Step": range(1, 8)}) - - -def initialize_geometric_factor_variables( - dataset: xr.Dataset, -) -> xr.Dataset: - """ - Initialize geometric factor variables for the specified species. - - Parameters - ---------- - dataset : xr.Dataset - Input dataset to add geometric factor variables to. - - Returns - ------- - xr.Dataset - Dataset with initialized geometric factor variables for the specified species. - """ - gf_vars = [ - "energy", - "energy_delta_minus", - "energy_delta_plus", - "geometric_factor", - "geometric_factor_stat_uncert_minus", - "geometric_factor_stat_uncert_plus", - ] - - # Initialize variables with proper dimensions (energy only) - for var in gf_vars: - dataset[var] = xr.DataArray( - np.zeros(7), - dims=["energy"], - ) - - return dataset - - -def populate_geometric_factors( - dataset: xr.Dataset, - species: str, -) -> xr.Dataset: - """ - Populate geometric factor values for each energy step. - - Parameters - ---------- - dataset : xr.Dataset - Dataset with initialized geometric factor variables. - species : str - The species to process (only "h" and "o" have geometric factors). - - Returns - ------- - xr.Dataset - Dataset with populated geometric factor values for the specified species. - """ - # Only populate if the species has geometric factors - if species not in ["h", "o"]: - logger.debug(f"No geometric factors to populate for species: {species}") - return dataset - - # Mapping of dataset variables to dataframe columns for this species - gf_coords = {"energy": "Cntr_E"} - gf_vars = { - "geometric_factor": f"GF_Trpl_{species.upper()}", - "geometric_factor_stat_uncert_minus": f"GF_Trpl_{species.upper()}_unc_minus", - "geometric_factor_stat_uncert_plus": f"GF_Trpl_{species.upper()}_unc_plus", - } - if species == "h": - # NOTE: From an e-mail from Nathan on 2025-09-11 (values converted to keV) - energy_delta_hires_values = ( - np.array([5.43, 10.02, 18.61, 33.31, 64.98, 131.64, 262.35]) * 1e-3 - ) - energy_delta_hithr_values = ( - np.array([8.81, 16.04, 28.50, 53.13, 105.60, 219.67, 413.60]) * 1e-3 - ) - else: # species == "o" - energy_delta_hires_values = ( - np.array([5.82, 11.10, 21.78, 41.47, 85.61, 180.67, 361.93]) * 1e-3 - ) - energy_delta_hithr_values = ( - np.array([9.45, 17.84, 33.51, 66.61, 139.95, 302.24, 569.48]) * 1e-3 - ) - - # Get ESA mode from the map (assuming it's constant or we take the first) - # TODO: Figure out how to handle esa_mode properly - if "esa_mode" in dataset: - esa_mode = dataset["esa_mode"].values[0] - else: - # Default to mode 0 if not available (HiRes mode) - esa_mode = 0 - - # Filter for the specific ESA mode - gf_dataset = reduce_geometric_factor_dataset(species, esa_mode) - - # Populate geometric factors in dataset - dataset = dataset.assign_coords(energy=gf_dataset[gf_coords["energy"]].values) - for var, col in gf_vars.items(): - dataset[var].values = gf_dataset[col].values - - # Update delta_minus and delta_plus based on ESA mode - # converting eV to keV - if esa_mode == 0: # HiRes - dataset["energy_delta_minus"].values = energy_delta_hires_values - dataset["energy_delta_plus"].values = energy_delta_hires_values - else: # HiThr - dataset["energy_delta_minus"].values = energy_delta_hithr_values - dataset["energy_delta_plus"].values = energy_delta_hithr_values - - return dataset + ram_projection = np.sin(np.radians(pivot_angle + c.PIVOT_RAM_OFFSET)) * np.sin( + np.radians(spin_angles) + ) + return ram_projection > 0 if spin_phase == "ram" else ram_projection < 0 # ============================================================================= -# RATES AND INTENSITIES CALCULATIONS +# RATES AND INTENSITIES # ============================================================================= -def calculate_all_rates_and_intensities( - dataset: xr.Dataset, - sputtering_correction: bool = False, - bootstrap_correction: bool = False, - flux_correction: bool = False, - o_map_dataset: xr.Dataset | None = None, - flux_factors: Path | None = None, - cg_correction: bool = False, -) -> xr.Dataset: - """ - Calculate rates and intensities with proper error propagation. - - Parameters - ---------- - dataset : xr.Dataset - Sky map dataset with count data and geometric factors. - sputtering_correction : bool, optional - Whether to apply sputtering corrections to oxygen intensities. - Default is False. - bootstrap_correction : bool, optional - Whether to apply bootstrap corrections to intensities. - Default is False. - flux_correction : bool, optional - Whether to apply flux corrections to intensities. - Default is False. - o_map_dataset : xr.Dataset, optional - Dataset specifically for oxygen, needed for sputtering corrections. - flux_factors : Path, optional - Path to flux factor file for flux corrections. - cg_correction : bool, optional - Whether to apply CG correction to intensities. - - Returns - ------- - xr.Dataset - Dataset with calculated rates, intensities, and uncertainties for the - specified species. - """ - # Step 1: Calculate rates for the specified species - dataset = calculate_rates(dataset) - - # Step 2: Calculate intensities - dataset = calculate_intensities(dataset) - - # Step 3: Calculate background rates and intensities - dataset = calculate_backgrounds(dataset) - - # Optional Step 4: Calculate sputtering corrections - if sputtering_correction: - logger.info("Calculating sputtering corrections") - dataset = calculate_sputtering_corrections(dataset, o_map_dataset) - - # Optional Step 5: Calculate bootstrap corrections - if bootstrap_correction: - logger.info("Calculating bootstrap corrections") - dataset = calculate_bootstrap_corrections(dataset) - - # Optional Step 6: Calculate flux corrections - if flux_correction: - if flux_factors is None: - raise ValueError("Flux factors file must be provided for flux corrections") - dataset = calculate_flux_corrections(dataset, flux_factors) - - # Optional Step 7: Finish CG correction - if cg_correction: - logger.info("Interpolating map intensities to helio-frame energies") - # Finish calculation of the exposure factor weighted projection of energy_sc - # and convert to units of keV - dataset["energy_sc"] = ( - dataset["energy_sc_exposure_factor"] / dataset["exposure_factor"] / 1e3 - ) - dataset = interpolate_map_flux_to_helio_frame( - dataset, - dataset["energy"], - dataset["energy"], - ["ena_intensity", "bg_intensity"], - ) - - # Step 7: Clean up intermediate variables - dataset = cleanup_intermediate_variables(dataset) - - return dataset - - -def calculate_rates(dataset: xr.Dataset) -> xr.Dataset: - """ - Calculate count rates and their statistical uncertainties. - - Parameters - ---------- - dataset : xr.Dataset - Dataset with count data and exposure times. - - Returns - ------- - xr.Dataset - Dataset with calculated count rates and statistical uncertainties - for the specified species. - """ - # Rate = counts / exposure_factor - # TODO: Account for ena / isn naming differences - dataset["ena_count_rate"] = dataset["counts"] / dataset["exposure_factor"] - - # Poisson uncertainty on the counts propagated to the rate - # TODO: Is there uncertainty in the exposure time too? - dataset["ena_count_rate_stat_uncert"] = ( - np.sqrt(dataset["counts"]) / dataset["exposure_factor"] - ) - - return dataset - - -def calculate_intensities(dataset: xr.Dataset) -> xr.Dataset: +def _geometric_factors(esa_mode: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ - Calculate particle intensities and uncertainties for the specified species. + Get the recalibrated geometric factors and their asymmetric bounds. Parameters ---------- - dataset : xr.Dataset - Dataset with count rates, geometric factors, and center energies. - - Returns - ------- - xr.Dataset - Dataset with calculated particle intensities and their statistical - and systematic uncertainties for the specified species. - """ - # Equation 3 from mapping document (average intensity) - dataset["ena_intensity"] = dataset["counts_over_eff"] / ( - dataset["geometric_factor"] * dataset["energy"] * dataset["exposure_factor"] - ) - - # Equation 4 from mapping document (statistical uncertainty) - # Note that we need to take the square root to get the uncertainty as - # the equation is for the variance - dataset["ena_intensity_stat_uncert"] = np.sqrt( - dataset["counts_over_eff_squared"] - ) / (dataset["geometric_factor"] * dataset["energy"] * dataset["exposure_factor"]) - - plus_multiplier = dataset["geometric_factor"] / ( - dataset["geometric_factor"] - dataset["geometric_factor_stat_uncert_minus"] - ) - minus_multiplier = dataset["geometric_factor"] / ( - dataset["geometric_factor"] + dataset["geometric_factor_stat_uncert_plus"] - ) - - dataset["ena_intensity_sys_err_plus"] = ( - dataset["ena_intensity"] * plus_multiplier - ) - dataset["ena_intensity"] - - dataset["ena_intensity_sys_err_minus"] = dataset["ena_intensity"] - ( - dataset["ena_intensity"] * minus_multiplier - ) - - # Symmetric systematic error - dataset["ena_intensity_sys_err"] = np.sqrt( - dataset["ena_intensity_sys_err_minus"] * dataset["ena_intensity_sys_err_plus"] - ) - - return dataset - - -def calculate_backgrounds(dataset: xr.Dataset) -> xr.Dataset: - """ - Calculate background rates and intensities for the specified species. - - Parameters - ---------- - dataset : xr.Dataset - Dataset with count rates, geometric factors, and center energies. + esa_mode : int + The ESA mode, 0 for HiRes and 1 for HiThr. Unused for now, the + geometric factors are not yet split by ESA mode. Returns ------- - xr.Dataset - Dataset with calculated background rates and intensities for the - specified species. + tuple[np.ndarray, np.ndarray, np.ndarray] + The geometric factor of each ESA level, and its upper and lower error + bounds. """ - # Equation 62 from mapping document (background rate) - # exposure time weighted average of the background rates - dataset["bg_rate"] = dataset["bg_rate_exposure_factor"] / dataset["exposure_factor"] - # Equation 63 from mapping document (background statistical uncertainty) - dataset["bg_rate_stat_uncert"] = np.sqrt( - dataset["bg_rate_stat_uncert_exposure_factor2"] - / dataset["exposure_factor"] ** 2 - ) - # Equation 64 from mapping document (background systematic error). - dataset["bg_rate_sys_err"] = ( - dataset["bg_rate_sys_err_exposure_factor"] / dataset["exposure_factor"] - ) - # The background rate systematic is a single symmetric value. - dataset["bg_rate_sys_err_plus"] = dataset["bg_rate_sys_err"].copy() - dataset["bg_rate_sys_err_minus"] = dataset["bg_rate_sys_err"].copy() - - plus_multiplier = dataset["geometric_factor"] / ( - dataset["geometric_factor"] - dataset["geometric_factor_stat_uncert_minus"] - ) - minus_multiplier = dataset["geometric_factor"] / ( - dataset["geometric_factor"] + dataset["geometric_factor_stat_uncert_plus"] - ) - - # Background intensity - dataset["bg_intensity"] = dataset["bg_rate"] / ( - dataset["geometric_factor"] * dataset["energy"] - ) - dataset["bg_intensity_stat_uncert"] = dataset["bg_rate_stat_uncert"] / ( - dataset["geometric_factor"] * dataset["energy"] - ) + levels = slice(0, c.N_ESA_LEVELS) + geometric_factor = np.array(c.GEO_FACTOR[levels]) * c.GEO_FACTOR_SCALE + error = np.array(c.GEO_FACTOR_ERR[levels]) * c.GEO_FACTOR_SCALE - dataset["bg_intensity_sys_err_plus"] = ( - dataset["bg_intensity"] * plus_multiplier - ) - dataset["bg_intensity"] + error_upper = np.hypot(geometric_factor * (c.GEO_FACTOR_SCALE_UPPER - 1.0), error) + error_lower = np.hypot(geometric_factor * (1.0 - c.GEO_FACTOR_SCALE_LOWER), error) - dataset["bg_intensity_sys_err_minus"] = dataset["bg_intensity"] - ( - dataset["bg_intensity"] * minus_multiplier - ) - - # Symmetric systematic error - dataset["bg_intensity_sys_err"] = np.sqrt( - dataset["bg_intensity_sys_err_minus"] * dataset["bg_intensity_sys_err_plus"] - ) - - return dataset + return geometric_factor, error_upper, error_lower -def calculate_sputtering_corrections( - dataset: xr.Dataset, o_dataset: xr.Dataset +def _calculate_rates_and_intensities( + counts: np.ndarray, + exposure: np.ndarray, + bg_rate_exposure: np.ndarray, + esa_mode: int, ) -> xr.Dataset: """ - Calculate sputtering corrections from oxygen intensities. - - Correction factors are read from imap_lo_sputter-correction-factors_v001.csv. - Only for Oxygen sputtering and correction only at ESA levels 5 and 6 - for 90 degree maps. If off-angle maps are made, we may have to extend - this to levels 3 and 4 as well. - - Follows equations 9-13 from the mapping document. - - Parameters - ---------- - dataset : xr.Dataset - Dataset with count rates, geometric factors, and center energies. - This is an H dataset that we are applying the corrections to. - o_dataset : xr.Dataset - Dataset specifically for oxygen, needed to access oxygen intensities - and uncertainties. - - Returns - ------- - xr.Dataset - Dataset with calculated sputtering-corrected intensities and their - uncertainties. - """ - logger.info("Applying sputtering corrections to hydrogen intensities") - sputter_df = load_sputter_correction_data("o", "h") - energy_indices = (sputter_df["esa_step"].values - 1).tolist() - - small_dataset = dataset.isel(epoch=0, energy=energy_indices) - o_small_dataset = o_dataset.isel(epoch=0, energy=energy_indices) - - # We need to align the energy dimensions from the oxygen dataset to the - # Hydrogen dataset so the calculations below get aligned by xarray correctly. - o_small_dataset["energy"] = small_dataset["energy"] - - # Equation 9 - j_o_prime = o_small_dataset["ena_intensity"] - o_small_dataset["bg_intensity"] - j_o_prime.values[j_o_prime.values < 0] = 0 # No negative intensities - j_o_prime_valid = np.isfinite(j_o_prime) & (j_o_prime > 0) - - # Equation 10 - j_o_prime_var = ( - o_small_dataset["ena_intensity_stat_uncert"] ** 2 - + o_small_dataset["bg_intensity_stat_uncert"] ** 2 - ) - - sputter_correction_factor = xr.DataArray( - sputter_df["sputter_factor"].values, - dims=["energy"], - coords={"energy": small_dataset["energy"]}, - ) - # Equation 11 - # Remove the sputtered oxygen intensity to correct the original H intensity - sputter_corrected_intensity = xr.where( - j_o_prime_valid, - small_dataset["ena_intensity"] - sputter_correction_factor * j_o_prime, - small_dataset["ena_intensity"], - ) - - # Equation 12 - sputter_corrected_intensity_var = xr.where( - j_o_prime_valid, - small_dataset["ena_intensity_stat_uncert"] ** 2 - + (sputter_correction_factor**2) * j_o_prime_var, - small_dataset["ena_intensity_stat_uncert"] ** 2, - ) - - # Equation 13 - sputter_corrected_intensity_sys_err = xr.where( - j_o_prime_valid, - sputter_corrected_intensity - / small_dataset["ena_intensity"] - * small_dataset["ena_intensity_sys_err"], - small_dataset["ena_intensity_sys_err"], - ) - - # Now put the corrected values into the original dataset - dataset["ena_intensity"].values[0, energy_indices, ...] = ( - sputter_corrected_intensity.values - ) - dataset["ena_intensity_stat_uncert"].values[0, energy_indices, ...] = np.sqrt( - sputter_corrected_intensity_var.values - ) - dataset["ena_intensity_sys_err"].values[0, energy_indices, ...] = ( - sputter_corrected_intensity_sys_err.values - ) - - return dataset - - -def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset: - """ - Calculate bootstrap corrections for hydrogen and oxygen intensities. + Turn the accumulated counts and exposure into rates and intensities. - Follows equations 14-35 from the mapping document. + Every quantity is zero in the pixels that were never exposed. Parameters ---------- - dataset : xr.Dataset - Dataset with count rates, geometric factors, and center energies. + counts : np.ndarray + Accumulated counts of shape (esa level, pixel). + exposure : np.ndarray + Accumulated exposure time [s] of shape (esa level, pixel). + bg_rate_exposure : np.ndarray + Accumulated exposure-weighted background rate, same shape. + esa_mode : int + The ESA mode, 0 for HiRes and 1 for HiThr. Returns ------- - xr.Dataset - Dataset with calculated bootstrap-corrected intensities and their - uncertainties for hydrogen. - """ - logger.info("Applying bootstrap corrections") - - # Table 3 bootstrap terms h_i,k - load from an ancillary file - bootstrap_df = load_bootstrap_correction_data() - - # Create xarray DataArray with named dimensions for proper broadcasting - bootstrap_factor = ( - bootstrap_df.set_index(["esa_step_i", "esa_step_k"])["bootstrap_factor"] - .to_xarray() - .fillna(0) - .reindex(esa_step_i=range(1, 8), esa_step_k=range(1, 9), fill_value=0) - .rename({"esa_step_i": "energy_i", "esa_step_k": "energy_k"}) - .assign_coords( - energy_i=dataset["energy"].values, - # Add an extra coordinate for the virtual E8 channel, unused - # in the broadcasting calculations - energy_k=np.concatenate([dataset["energy"].values, [np.nan]]), + dict[str, np.ndarray] + The map variables, each of shape (esa level, pixel). + """ + energy = np.array(c.ESA_ENERGY[: c.N_ESA_LEVELS])[:, np.newaxis] + geometric_factor, error_upper, error_lower = _geometric_factors(esa_mode) + geometric_factor = geometric_factor[:, np.newaxis] + error_upper = error_upper[:, np.newaxis] + error_lower = error_lower[:, np.newaxis] + + exposed = exposure > 0 + + def _divide(numerator: np.ndarray, denominator: np.ndarray) -> np.ndarray: + """ + Divide only where the map was exposed, zero elsewhere. + + Parameters + ---------- + numerator : np.ndarray + The array being divided. + denominator : np.ndarray + The array to divide it by. + + Returns + ------- + np.ndarray + The quotient, zero in the pixels that were never exposed. + """ + return np.divide( + numerator, + denominator, + out=np.zeros_like(exposure), + where=exposed, ) - ) - # Equation 14 - j_c_prime = dataset["ena_intensity"] - dataset["bg_intensity"] - j_c_prime.values[j_c_prime.values < 0] = 0 - - # Equation 15 - j_c_prime_var = dataset["ena_intensity_stat_uncert"] ** 2 - - # Equation 16 - systematic error propagation - # Handle division by zero: only compute where ena_intensity > 0 - j_c_prime_err = xr.where( - dataset["ena_intensity"] > 0, - j_c_prime / dataset["ena_intensity"] * dataset["ena_intensity_sys_err"], - 0, - ) - - # NOTE: E8 virtual channel calculation is from the text. This is to - # start the calculations off from the higher energies and avoid - # reliance on IMAP Hi energy channels. - # E8 is a virtual energy channel at 2.1 * E7 - e8 = 2.1 * dataset["energy"].values[-1] - - j_c_6 = j_c_prime.isel(energy=5) - j_c_7 = j_c_prime.isel(energy=6) - e_6 = dataset["energy"].isel(energy=5) - e_7 = dataset["energy"].isel(energy=6) - - # Calculate gamma, ignoring any invalid values - # Fill in the invalid values with zeros after the fact - with np.errstate(divide="ignore", invalid="ignore"): - gamma = np.log(j_c_6 / j_c_7) / np.log(e_6 / e_7) - j_8_b = j_c_7 * (e8 / e_7) ** gamma - - # Set j_8_b to zero where the calculation was invalid - j_8_b = j_8_b.where(np.isfinite(j_8_b) & (j_8_b > 0), 0) - - # Initialize bootstrap intensity and uncertainty arrays - dataset["bootstrap_intensity"] = xr.zeros_like(dataset["ena_intensity"]) - dataset["bootstrap_intensity_var"] = xr.zeros_like(dataset["ena_intensity"]) - dataset["bootstrap_intensity_sys_err"] = xr.zeros_like(dataset["ena_intensity"]) - - for i in range(6, -1, -1): - # Create views for the current energy channel to avoid repeated indexing - bootstrap_intensity_i = dataset["bootstrap_intensity"][0, i, ...] - bootstrap_intensity_var_i = dataset["bootstrap_intensity_var"][0, i, ...] - j_c_prime_i = j_c_prime[0, i, ...] - j_c_prime_var_i = j_c_prime_var[0, i, ...] - - # Initialize the variable with the non-summation term and virtual - # channel energy subtraction first, then iterate through the other - # channels which can be looked up via indexing - # i.e. the summation is always k=i+1 to 7, because we've already - # included the k=8 term here. - # NOTE: The paper uses 1-based indexing and we use 0-based indexing - # so there is an off-by-one difference in the indices. - bootstrap_intensity_i[:] = ( - j_c_prime_i - bootstrap_factor.isel(energy_i=i, energy_k=7) * j_8_b[0, ...] - ) - # NOTE: We will square root at the end to get the uncertainty, but - # all equations are with variances - bootstrap_intensity_var_i[:] = j_c_prime_var_i - - # Vectorized summation using xarray's built-in broadcasting - # Select the relevant k indices for summation (k = i+1 to 6) - k_indices = list(range(i + 1, 7)) - - # Get bootstrap factors for this i and the relevant k values - # Rename energy_k dimension to energy for alignment with intensity - bootstrap_factors_k = bootstrap_factor.isel( - energy_i=i, energy_k=k_indices - ).rename({"energy_k": "energy"}) - - # Get intensity slices - these will have an 'energy' dimension still - intensity_k = dataset["bootstrap_intensity"][0, k_indices, ...] - intensity_var_k = dataset["bootstrap_intensity_var"][0, k_indices, ...] - - # Subtraction terms from equations 18-23 (xarray vectorized) - bootstrap_intensity_i -= (bootstrap_factors_k * intensity_k).sum(dim="energy") - - # Summation terms from equations 25-30 (xarray vectorized) - bootstrap_intensity_var_i += (bootstrap_factors_k**2 * intensity_var_k).sum( - dim="energy" + count_rate = _divide(counts, exposure) + # Poisson uncertainty on the counts, propagated to the rate + count_rate_stat_uncert = _divide(np.sqrt(counts), exposure) + + intensity = _divide(count_rate, geometric_factor * energy) + intensity_stat_uncert = _divide(count_rate_stat_uncert, geometric_factor * energy) + + # The systematic error is the flux excursion from the recalibrated G-factor + # bounds: the upper/lower excursions come from the lower/upper G-factor + # bounds respectively, and the symmetric error is their geometric mean. It + # is undefined where the lower bound would drive the G-factor non-positive. + valid = geometric_factor > error_lower + if not valid.all(): + logger.warning( + "The geometric factor of ESA levels " + f"{(np.flatnonzero(~valid[:, 0]) + 1).tolist()} is below its lower " + f"error bound; their systematic errors are left at zero." ) - - # Again zero any bootstrap fluxes that are negative - bootstrap_intensity_i.values[bootstrap_intensity_i < 0] = 0.0 - - # Equation 31 - systematic error propagation for bootstrap intensity - # Handle division by zero: only compute where j_c_prime > 0 - dataset["bootstrap_intensity_sys_err"] = xr.where( - j_c_prime > 0, dataset["bootstrap_intensity"] / j_c_prime * j_c_prime_err, 0 - ) - - valid_bootstrap = (dataset["bootstrap_intensity"] > 0) & np.isfinite( - dataset["bootstrap_intensity"] - ) - # Update the original intensity values - # Equation 32 / 33 - # ena_intensity = ena_intensity (J_c) - (j_c_prime - J_b) - dataset["ena_intensity"] = xr.where( - valid_bootstrap, - dataset["ena_intensity"] - j_c_prime + dataset["bootstrap_intensity"], - dataset["ena_intensity"], - ) - - # Ensure corrected intensities are non-negative - dataset["ena_intensity"] = xr.where( - dataset["ena_intensity"] < 0, 0, dataset["ena_intensity"] - ) - - # Equation 34 - statistical uncertainty - # Take the square root, since we were in variances up to this point - dataset["ena_intensity_stat_uncert"] = xr.where( - valid_bootstrap, - np.sqrt(dataset["bootstrap_intensity_var"]), - dataset["ena_intensity_stat_uncert"], - ) - - # Equation 35 - systematic error for corrected intensity - # Handle division by zero and ensure reasonable values - dataset["ena_intensity_sys_err"] = xr.zeros_like(dataset["ena_intensity"]) - - # Only compute where bootstrap intensity is valid - dataset["ena_intensity_sys_err"] = xr.where( - valid_bootstrap, - ( - dataset["ena_intensity"] - / dataset["bootstrap_intensity"] - * dataset["bootstrap_intensity_sys_err"] + intensity_sys_err_plus = np.where( + valid, + intensity * geometric_factor / (geometric_factor - error_lower) - intensity, + 0.0, + ) + intensity_sys_err_minus = np.where( + valid, + intensity - intensity * geometric_factor / (geometric_factor + error_upper), + 0.0, + ) + + bg_rate = _divide(bg_rate_exposure, exposure) + bg_rate_stat_uncert = np.sqrt(_divide(bg_rate, exposure)) + bg_intensity = _divide(bg_rate, geometric_factor * energy) + bg_intensity_stat_uncert = _divide(bg_rate_stat_uncert, geometric_factor * energy) + + return { + "ena_count": counts, + "exposure_factor": exposure, + "ena_count_rate": count_rate, + "ena_count_rate_stat_uncert": count_rate_stat_uncert, + "ena_intensity": intensity, + "ena_intensity_stat_uncert": intensity_stat_uncert, + "ena_intensity_sys_err": np.sqrt( + intensity_sys_err_plus * intensity_sys_err_minus ), - 0, - ) - - # Drop the intermediate bootstrap variables - dataset = dataset.drop_vars( - [ - "bootstrap_intensity", - "bootstrap_intensity_var", - "bootstrap_intensity_sys_err", - ] - ) - - return dataset + "ena_intensity_sys_err_plus": intensity_sys_err_plus, + "ena_intensity_sys_err_minus": intensity_sys_err_minus, + "bg_rate": bg_rate, + "bg_rate_stat_uncert": bg_rate_stat_uncert, + "bg_intensity": bg_intensity, + "bg_intensity_stat_uncert": bg_intensity_stat_uncert, + } -def calculate_flux_corrections(dataset: xr.Dataset, flux_factors: Path) -> xr.Dataset: +def _build_map_dataset( + sky_map: RectangularSkyMap, variables: dict[str, np.ndarray], esa_mode: int +) -> xr.Dataset: """ - Calculate flux corrections for intensities. + Lay the map variables out on the map's sky grid. - Uses the shared ena maps ``PowerLawFluxCorrector`` class to do the - correction calculations. + The variables are handed to the map as 1D pixel arrays, which the map + rewraps onto its longitude/latitude grid and adds its solid angles to. Parameters ---------- - dataset : xr.Dataset - Dataset with count rates, geometric factors, and center energies. - flux_factors : Path - Path to the eta flux factor file to use for corrections. Read in as - an ancillary file in the preprocessing step. + sky_map : RectangularSkyMap + The map being built. + variables : dict[str, np.ndarray] + The map variables, each of shape (esa level, pixel). + esa_mode : int + The ESA mode, 0 for HiRes and 1 for HiThr, which sets the widths of the + ESA energy passbands. Returns ------- xr.Dataset - Dataset with calculated flux-corrected intensities and their - uncertainties for the specified species. - """ - logger.info("Applying flux corrections") - - # Flux correction - corrector = PowerLawFluxCorrector(flux_factors) - - # NOTE: We need to apply this to both total flux and background flux - for var in ["ena", "bg"]: - # Apply flux correction with xarray inputs - dataset[f"{var}_intensity"], dataset[f"{var}_intensity_stat_uncert"] = ( - corrector.apply_flux_correction( - dataset[f"{var}_intensity"], - dataset[f"{var}_intensity_stat_uncert"], - dataset["energy"], - ) + The map variables on the (epoch, energy, longitude, latitude) grid, + with the energy coordinate and its widths. + """ + energy = np.array(c.ESA_ENERGY[: c.N_ESA_LEVELS]) + for name, values in variables.items(): + sky_map.data_1d[name] = xr.DataArray( + values[np.newaxis, ...].astype(np.float32), + dims=["epoch", "energy", "pixel"], + coords={"energy": energy}, ) - return dataset - - -def cleanup_intermediate_variables(dataset: xr.Dataset) -> xr.Dataset: - """ - Remove intermediate variables that were only needed for calculations. - - Parameters - ---------- - dataset : xr.Dataset - Dataset containing intermediate calculation variables. - - Returns - ------- - xr.Dataset - Cleaned dataset with intermediate variables removed. - """ - # Remove the intermediate variables from the map - # i.e. the ones that were projected from the pset only for the purposes - # of math and not desired in the output. - vars_to_remove = [] - - # Only remove variables that exist in the dataset for the specific species - potential_vars = [ - "geometric_factor", - "geometric_factor_stat_uncert", - "geometric_factor_stat_uncert_minus", - "geometric_factor_stat_uncert_plus", - "counts_over_eff", - "counts_over_eff_squared", - "bg_rate_exposure_factor", - "bg_rate_stat_uncert_exposure_factor2", - "bg_rate_sys_err_exposure_factor", - ] + dataset = sky_map.to_dataset() - for potential_var in potential_vars: - if potential_var in dataset.data_vars: - vars_to_remove.append(potential_var) + energy_delta = np.array(c.ESA_ENERGY_DELTA[esa_mode]) + dataset["energy_delta_minus"] = xr.DataArray(energy_delta, dims=["energy"]) + dataset["energy_delta_plus"] = xr.DataArray(energy_delta, dims=["energy"]) - return dataset.drop_vars(vars_to_remove) + return dataset diff --git a/imap_processing/tests/lo/test_lo_l2.py b/imap_processing/tests/lo/test_lo_l2.py index fef600bd0..8903684e8 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -1,2919 +1,430 @@ -"""Comprehensive test suite for IMAP-Lo L2 data processing.""" +"""Test suite for IMAP-Lo L2 map processing.""" -from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import patch import numpy as np -import pandas as pd import pytest import xarray as xr -from imap_processing.cdf.utils import load_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, -) +from imap_processing.cdf.utils import load_cdf, write_cdf from imap_processing.ena_maps.utils.naming import MapDescriptor -from imap_processing.lo.l1c.lo_l1c import ( - ESA_ENERGY_STEPS, - N_OFF_ANGLE_BINS, - N_SPIN_ANGLE_BINS, - OFF_ANGLE_BIN_CENTERS, - PSET_DIMS, - PSET_SHAPE, - SPIN_ANGLE_BIN_CENTERS, -) +from imap_processing.lo.constants import LoConstants from imap_processing.lo.l2.lo_l2 import ( - _prepare_corrections, - add_efficiency_factors_to_pset, - calculate_all_rates_and_intensities, - calculate_backgrounds, - calculate_bootstrap_corrections, - calculate_efficiency_corrected_quantities, - calculate_flux_corrections, - calculate_intensities, - calculate_rates, - calculate_sputtering_corrections, - cleanup_intermediate_variables, - create_sky_map_from_psets, - initialize_geometric_factor_variables, + BGRATES, + GOODTIMES, + HISTRATES, + _dps_spin_angles, + _group_inputs_by_pointing, + _pixel_indices, + _spin_phase_mask, lo_l2, - load_efficiency_data, - normalize_pset_coordinates, - populate_geometric_factors, - process_single_pset, - project_pset_to_map, - reduce_geometric_factor_dataset, ) - -# ============================================================================= -# FIXTURES FOR MOCK DATA -# ============================================================================= - - -@pytest.fixture(params=["h", "o", "doubles", "triples"]) -def species_name(request): - """Parametrized fixture for different species names.""" - return request.param - - -@pytest.fixture -def sample_pset(): - """Create a sample pointing set with typical data variables.""" - # Create counts data with some non-zero values - counts = np.zeros(PSET_SHAPE) - counts[:, 2:4, 10:20, 5:15] = 5 # Add some counts for testing - - exposure_factor = np.full(PSET_SHAPE, 0.5) - - # Create background rates data - background_rates = np.full(PSET_SHAPE, 0.1) # 0.1 counts/s background - background_rates_stat_uncert = np.full(PSET_SHAPE, 0.01) # 10% uncertainty - - # Create coordinate arrays - lons, lats = np.meshgrid( - SPIN_ANGLE_BIN_CENTERS, OFF_ANGLE_BIN_CENTERS, indexing="ij" - ) - hae_longitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_latitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_longitude[0, :, :] = lons - hae_latitude[0, :, :] = lats - - dataset = xr.Dataset( +from imap_processing.spice.time import met_to_ttj2000ns + +# A full-spin map, so that every spin-angle bin lands on it. +FULL_DESCRIPTOR = "l090-ena-h-sf-nsp-full-hae-6deg-3mo" +RAM_DESCRIPTOR = "l090-ena-h-sf-nsp-ram-hae-6deg-3mo" + +N_ESA = LoConstants.N_ESA_LEVELS +N_SPIN_BINS = LoConstants.N_SPIN_ANGLE_BINS +PIVOT = 90.0 + +# Good-time window [MET seconds] that the "in-window" histogram epochs fall in. +GT_START = 511_000_000.0 +GT_END = 511_000_600.0 +IN_METS = [511_000_150.0, 511_000_200.0, 511_000_250.0] +OUT_METS = [510_990_000.0, 511_010_000.0] + + +def make_pointing(repointing="repoint00100", pivot=PIVOT, seed=42): + """Build the three synthetic L1B inputs of one pointing. + + The in-window epochs carry modest counts and exposure; the out-of-window + epochs carry large values that good-time filtering must exclude. + """ + mets = np.array(IN_METS + OUT_METS) + in_idx = np.arange(len(IN_METS)) + out_idx = np.arange(len(IN_METS), mets.size) + + rng = np.random.default_rng(seed) + counts = np.zeros((mets.size, N_ESA, N_SPIN_BINS)) + exposure = np.zeros_like(counts) + for i in in_idx: + counts[i] = rng.integers(0, 4, size=(N_ESA, N_SPIN_BINS)).astype(float) + exposure[i] = 2.0 * (np.arange(N_ESA)[:, None] + 1) + for i in out_idx: + counts[i] = 999.0 + exposure[i] = 999.0 + + histrates = xr.Dataset( { - "counts": (PSET_DIMS, counts), - "exposure_factor": (PSET_DIMS, exposure_factor), - "background_rates": (PSET_DIMS, background_rates), - "background_rates_stat_uncert": ( - PSET_DIMS, - background_rates_stat_uncert, - ), - "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], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, - }, - ) - return dataset - - -@pytest.fixture -def sample_pset_for_species(species_name): - """Create a sample pointing set for a specific species.""" - # Create counts data with some non-zero values - counts = np.zeros(PSET_SHAPE) - if species_name == "h": - counts[:, 2:4, 10:20, 5:15] = 5 - elif species_name == "o": - counts[:, 1:3, 15:25, 8:18] = 3 - elif species_name == "doubles": - counts[:, 0:2, 5:15, 10:20] = 2 - elif species_name == "triples": - counts[:, 3:5, 20:30, 15:25] = 1 - - exposure_factor = np.full(PSET_SHAPE, 0.5) - - # Create coordinate arrays - lons, lats = np.meshgrid( - SPIN_ANGLE_BIN_CENTERS, OFF_ANGLE_BIN_CENTERS, indexing="ij" - ) - hae_longitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_latitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_longitude[0, :, :] = lons - hae_latitude[0, :, :] = lats - - # Base dataset with coords and exposure time - dataset_dict = { - "counts": (PSET_DIMS, counts), - "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 - if species_name in ["h", "o"]: - bg_rate = np.full(PSET_SHAPE, 0.1 if species_name == "h" else 0.05) - bg_uncert = np.full(PSET_SHAPE, 0.01 if species_name == "h" else 0.005) - dataset_dict["background_rates"] = (PSET_DIMS, bg_rate) - dataset_dict["background_rates_stat_uncert"] = ( - PSET_DIMS, - bg_uncert, - ) - - dataset = xr.Dataset( - dataset_dict, - coords={ - "epoch": [8.1794907049e17], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, + "h_counts": (["epoch", "esa_step", "spin_bin_6"], counts), + "exposure_time_6deg": (["epoch", "esa_step", "spin_bin_6"], exposure), + "esa_mode": ("epoch", np.zeros(mets.size, dtype=int)), }, + coords={"epoch": met_to_ttj2000ns(mets)}, + attrs={"Repointing": repointing, "Logical_source": HISTRATES}, ) - return dataset - - -@pytest.fixture -def minimal_pset(): - """Create a minimal pointing set with typical data for testing.""" - counts = np.ones(PSET_SHAPE) # All ones for easy testing - exposure_factor = np.full(PSET_SHAPE, 1.0) # 1 second exposure for easy math - - # Create simple background rates for testing - background_rates = np.full(PSET_SHAPE, 0.2) # 0.2 counts/s - background_rates_stat_uncert = np.full(PSET_SHAPE, 0.02) # 10% uncertainty - - # Simple coordinate arrays - lons, lats = np.meshgrid( - SPIN_ANGLE_BIN_CENTERS, OFF_ANGLE_BIN_CENTERS, indexing="ij" - ) - hae_longitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_latitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_longitude[0, :, :] = lons - hae_latitude[0, :, :] = lats - - dataset = xr.Dataset( + goodtimes = xr.Dataset( { - "counts": (PSET_DIMS, counts), - "exposure_factor": (PSET_DIMS, exposure_factor), - "background_rates": (PSET_DIMS, background_rates), - "background_rates_stat_uncert": ( - PSET_DIMS, - background_rates_stat_uncert, - ), - "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], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, - }, - ) - return dataset - - -@pytest.fixture -def minimal_pset_for_species(species_name): - """Create a minimal pointing set for a specific species.""" - # Create simple counts data - if species_name == "h": - counts = np.ones(PSET_SHAPE) - elif species_name == "o": - counts = np.ones(PSET_SHAPE) * 0.5 - elif species_name == "doubles": - counts = np.ones(PSET_SHAPE) * 0.2 - elif species_name == "triples": - counts = np.ones(PSET_SHAPE) * 0.1 - - exposure_factor = np.full(PSET_SHAPE, 1.0) # 1 second exposure for easy math - - # Simple coordinate arrays - lons, lats = np.meshgrid( - SPIN_ANGLE_BIN_CENTERS, OFF_ANGLE_BIN_CENTERS, indexing="ij" - ) - hae_longitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_latitude = np.empty((1, N_SPIN_ANGLE_BINS, N_OFF_ANGLE_BINS)) - hae_longitude[0, :, :] = lons - hae_latitude[0, :, :] = lats - - # Base dataset with coords and exposure time - dataset_dict = { - "counts": (PSET_DIMS, counts), - "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 - bg_rate = np.full(PSET_SHAPE, 0.2 if species_name == "h" else 0.1) - bg_uncert = np.full(PSET_SHAPE, 0.02 if species_name == "h" else 0.01) - dataset_dict["background_rates"] = (PSET_DIMS, bg_rate) - dataset_dict["background_rates_stat_uncert"] = (PSET_DIMS, bg_uncert) - - dataset = xr.Dataset( - dataset_dict, - coords={ - "epoch": [8.1794907049e17], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, + "pivot": ("epoch", [pivot]), + "gt_start_met": ("epoch", [GT_START]), + "gt_end_met": ("epoch", [GT_END]), }, + coords={"epoch": [0]}, + attrs={"Repointing": repointing, "Logical_source": GOODTIMES}, ) - return dataset - - -@pytest.fixture -def sample_efficiency_data(): - """Create sample efficiency factor data for testing.""" - data = { - "Date": [np.datetime64("2025-01-01"), np.datetime64("2025-01-02")], - "E-Step1_eff": [0.8, 0.85], - "E-Step2_eff": [0.82, 0.87], - "E-Step3_eff": [0.84, 0.89], - "E-Step4_eff": [0.86, 0.91], - "E-Step5_eff": [0.88, 0.93], - "E-Step6_eff": [0.90, 0.95], - "E-Step7_eff": [0.92, 0.97], - } - return pd.DataFrame(data) - - -@pytest.fixture -def sample_geometric_factor_data(): - """Create sample geometric factor data for testing.""" - h_gf_data = [] - o_gf_data = [] - - for i in range(7): # 7 energy steps - h_gf_data.append( - { - "esa_mode": 0, - "Observed_E-Step": i + 1, - "incident_E-Step": i + 1, - "Cntr_E": 0.01 * (i + 1), # Simple energy values - "Cntr_E_unc": 0.001 * (i + 1), - "GF_Trpl_H": 1e-4 * (i + 1), - "GF_Trpl_H_unc_minus": 1e-5 * (i + 1), - "GF_Trpl_H_unc_plus": 2e-5 * (i + 1), - "GF_Dbl_all": 2e-4 * (i + 1), - "GF_Dbl_all_unc": 2e-5 * (i + 1), - "GF_Trpl_all": 3e-4 * (i + 1), - "GF_Trpl_all_unc": 3e-5 * (i + 1), - } - ) - - o_gf_data.append( - { - "esa_mode": 0, - "Observed_E-Step": i + 1, - "incident_E-Step": i + 1, - "Cntr_E": 0.015 * (i + 1), # Slightly different for oxygen - "Cntr_E_unc": 0.0015 * (i + 1), - "GF_Trpl_O": 1.5e-4 * (i + 1), - "GF_Trpl_O_unc_minus": 1.5e-5 * (i + 1), - "GF_Trpl_O_unc_plus": 3e-5 * (i + 1), - } - ) - - return pd.DataFrame(h_gf_data), pd.DataFrame(o_gf_data) - - -@pytest.fixture -def sample_sky_map_dataset(): - """Create a sample sky map dataset for testing calculations.""" - # Create a simple rectangular map - n_lon, n_lat = 60, 30 - n_energy = 7 - - dataset = xr.Dataset( - coords={ - "epoch": [8.1794907049e17], - "energy": list(range(n_energy)), - "longitude": np.linspace(0, 360, n_lon, endpoint=False), - "latitude": np.linspace(-90, 90, n_lat), - } - ) - - # Current lo_l2.py uses generic variable names, not species-specific - counts = np.ones((1, n_energy, n_lon, n_lat)) * 10 # 10 counts for easy math - dataset["counts"] = (("epoch", "energy", "longitude", "latitude"), counts) - - # Add efficiency-corrected quantities for intensity calculations - eff_corr = counts / 0.9 # Assuming 90% efficiency - dataset["counts_over_eff"] = ( - ("epoch", "energy", "longitude", "latitude"), - eff_corr, - ) - dataset["counts_over_eff_squared"] = ( - ("epoch", "energy", "longitude", "latitude"), - eff_corr, - ) - - # Add exposure time using the current naming convention - exposure = np.ones((1, n_energy, n_lon, n_lat)) * 1.0 # 1 second - dataset["exposure_factor"] = ( - ("epoch", "energy", "longitude", "latitude"), - exposure, - ) - - return dataset - - -@pytest.fixture -def sample_dataset_with_geometric_factors(): - """Create a dataset with geometric factors for testing calculations.""" - dataset = xr.Dataset( - coords={ - "epoch": [8.1794907049e17], - "energy": list(range(7)), - } - ) - - # Add current generic variable names used by lo_l2.py - dataset["counts_over_eff"] = (("epoch", "energy"), np.ones((1, 7)) * 100) - dataset["counts_over_eff_squared"] = (("epoch", "energy"), np.ones((1, 7)) * 100) - dataset["exposure_factor"] = (("epoch", "energy"), np.ones((1, 7)) * 1.0) - dataset["geometric_factor"] = (("energy",), np.ones(7) * 1e-4) - dataset["energy"] = (("energy",), np.ones(7) * 0.1) # Energy values - dataset["geometric_factor_stat_uncert_minus"] = (("energy",), np.ones(7) * 1e-5) - dataset["geometric_factor_stat_uncert_plus"] = (("energy",), np.ones(7) * 3e-5) - - return dataset - - -@pytest.fixture -def sample_dataset_with_background_intermediates(): - """Create a dataset with background intermediate variables for testing.""" - # Create a simple rectangular map with background data - n_energy = 7 - - dataset = xr.Dataset( - coords={ - "epoch": [8.1794907049e17], - "energy": list(range(n_energy)), - } - ) - - # Add the intermediate background variables using current naming convention - bg_rate_exposure_factor = np.ones((1, n_energy)) * 0.2 # 0.2 counts - dataset["bg_rate_exposure_factor"] = (("epoch", "energy"), bg_rate_exposure_factor) - - # Background uncertainty squared times exposure time squared - bg_rate_stat_uncert_exposure_factor2 = np.ones((1, n_energy)) * 0.004 # 0.02^2 - dataset["bg_rate_stat_uncert_exposure_factor2"] = ( - ("epoch", "energy"), - bg_rate_stat_uncert_exposure_factor2, - ) - - # Background systematic error times exposure time - bg_rate_sys_err_exposure_factor = np.ones((1, n_energy)) * 0.05 - dataset["bg_rate_sys_err_exposure_factor"] = ( - ("epoch", "energy"), - bg_rate_sys_err_exposure_factor, - ) - - # Add exposure time (using current naming convention) - exposure = np.ones((1, n_energy)) * 1.0 # 1 second - dataset["exposure_factor"] = (("epoch", "energy"), exposure) - - # Add geometric factors for systematic uncertainty calculation - dataset["geometric_factor"] = (("energy",), np.ones(n_energy) * 1e-4) - dataset["geometric_factor_stat_uncert_minus"] = ( - ("energy",), - np.ones(n_energy) * 1e-5, - ) - dataset["geometric_factor_stat_uncert_plus"] = ( - ("energy",), - np.ones(n_energy) * 3e-5, + background = np.array([0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07]) + bgrates = xr.Dataset( + {"h_background_rates": (["epoch", "esa_step"], background[np.newaxis, :])}, + coords={"epoch": [0], "esa_step": np.arange(1, N_ESA + 1)}, + attrs={"Repointing": repointing, "Logical_source": BGRATES}, ) - return dataset - - -@pytest.fixture -def sample_dataset_with_sputtering_data(): - """Create datasets with ENA intensities for sputtering correction testing.""" - # Create a simple map dataset with the required variables for sputtering correction - n_energy = 7 - n_lon, n_lat = 10, 5 # Smaller for testing - - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(n_energy)), - "longitude": np.linspace(0, 360, n_lon, endpoint=False), - "latitude": np.linspace(-90, 90, n_lat), + return { + GOODTIMES: goodtimes, + BGRATES: bgrates, + HISTRATES: histrates, + "expected_counts": counts[in_idx].sum(axis=(0, 2)), + "expected_exposure": exposure[in_idx].sum(axis=(0, 2)), + "background": background, } - # Create hydrogen dataset - h_intensity_values = np.ones((1, n_energy, n_lon, n_lat)) * 1e6 # Base intensity - h_intensity_values[0, 4, :, :] *= 3 # Higher at energy level 4 (ESA level 5) - h_intensity_values[0, 5, :, :] *= 2 # Higher at energy level 5 (ESA level 6) - - h_dataset = xr.Dataset(coords=coords) - h_dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - h_intensity_values, - ) - h_dataset["bg_intensity"] = h_dataset["ena_intensity"] * 0.1 # 10% background - - # Add statistical uncertainties - h_dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.sqrt(h_intensity_values) * 0.1, - ) - h_dataset["bg_intensity_stat_uncert"] = np.sqrt(h_dataset["bg_intensity"]) * 0.1 - - # Add systematic error - h_dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - h_intensity_values * 0.05, - ) - h_dataset["bg_intensity_sys_err"] = h_dataset["bg_intensity"] * 0.05 - - # Create oxygen dataset - o_intensity_values = np.ones((1, n_energy, n_lon, n_lat)) * 1e6 # Base intensity - o_intensity_values[0, 4, :, :] *= 5 # Higher at energy level 4 (ESA level 5) - o_intensity_values[0, 5, :, :] *= 3 # Higher at energy level 5 (ESA level 6) - - o_dataset = xr.Dataset(coords=coords) - o_dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - o_intensity_values, - ) - o_dataset["bg_intensity"] = o_dataset["ena_intensity"] * 0.1 # 10% background - o_dataset["bg_intensity_stat_uncert"] = np.sqrt(o_dataset["bg_intensity"]) * 0.1 - - # Add statistical uncertainties - o_dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.sqrt(o_intensity_values) * 0.1, - ) - o_dataset["bg_intensity_stat_uncert"] = ( - np.sqrt(o_dataset["bg_intensity_stat_uncert"]) * 0.1 - ) - - # Add systematic error - o_dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - o_intensity_values * 0.05, - ) - o_dataset["bg_intensity_sys_err"] = o_dataset["bg_intensity"] * 0.05 - - # Add geometric factors for intensity calculations - o_dataset["geometric_factor"] = (("energy",), np.ones(n_energy)) - - return h_dataset, o_dataset - -@pytest.fixture -def sample_dataset_with_bootstrap_data(): - """Create a dataset with ENA intensities for bootstrap correction testing.""" - # Create a simple map dataset with the required variables for bootstrap correction - n_energy = 7 - n_lon, n_lat = 10, 5 # Smaller for testing - - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(n_energy)), - "longitude": np.linspace(0, 360, n_lon, endpoint=False), - "latitude": np.linspace(-90, 90, n_lat), +def as_dependencies(*pointings): + """Turn pointings into the sci_dependencies lo_l2 takes.""" + return { + source: [pointing[source] for pointing in pointings] + for source in (GOODTIMES, BGRATES, HISTRATES) } - # Create realistic energy values for hydrogen - energy_values = np.array([0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]) # keV - - # Create intensity values that follow a power law distribution - # Higher intensities at lower energies - base_intensity = 1e6 # particles/(cm^2 sr s keV) - intensity_values = np.ones((1, n_energy, n_lon, n_lat)) - for i in range(n_energy): - # Power law: I = I0 * (E/E0)^(-2.5) - intensity_values[0, i, :, :] = base_intensity * (energy_values[i] / 1.0) ** ( - -2.5 - ) - - dataset = xr.Dataset(coords=coords) - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - dataset["energy"] = (("energy",), energy_values) - dataset["geometric_factor"] = (("energy",), np.ones(n_energy)) - - # Add background rates (much lower values) - dataset["bg_intensity"] = ( - dataset["ena_intensity"] - * 0.1 - / (dataset["geometric_factor"] * dataset["energy"]) - ) - - # Add statistical uncertainties (Poisson-like) - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.sqrt(intensity_values) * 0.1, - ) - # Add systematic error (5% of intensity) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.05, - ) +def identity_pointing(et, az_el, *args, **kwargs): + """Stand in for the SPICE DPS transform: spin angle -> lon, off -> lat. - return dataset + Like the real ``frame_transform_az_el``, the singleton off-angle dimension + is squeezed out. + """ + return np.asarray(az_el)[:, 0, :] @pytest.fixture -def lo_flux_factors_file(): - """Path to the LO flux factors test file.""" - # Use the actual test data file from the ena_maps test data - test_data_path = Path(__file__).parent.parent / "ena_maps" / "data" - return test_data_path / "imap_lo_esa-eta-fit-factors_20240101_v001.csv" +def one_pointing(): + """A single synthetic pointing.""" + return make_pointing() @pytest.fixture -def sample_dataset_with_intensities(): - """Create a dataset with intensities for flux correction testing.""" - n_energy = 7 - n_lon, n_lat = 6, 4 # Small for testing - - # Create realistic energy values matching the flux factors file - energy_values = np.array([16.35, 30.56, 56.4, 105, 199.8, 407.5, 795.3]) - - coords = { - "epoch": [8.1794907049e17], - "energy": energy_values, - "longitude": np.linspace(0, 360, n_lon, endpoint=False), - "latitude": np.linspace(-90, 90, n_lat), - } - - # Create intensity values with some spatial and energy structure - intensity_values = np.ones((1, n_energy, n_lon, n_lat)) - for i in range(n_energy): - # Power law: I = I0 * (E/E0)^(-2.0) - intensity_values[0, i, :, :] = 1e6 * (energy_values[i] / 100.0) ** (-2.0) - - # Add some spatial structure - for j in range(n_lon): - for k in range(n_lat): - intensity_values[0, :, j, k] *= 1.0 + 0.1 * np.sin(j) * np.cos(k) - - dataset = xr.Dataset(coords=coords) - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - - # Add statistical uncertainties (10% of intensity) - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.1, - ) - - dataset["bg_intensity"] = dataset["ena_intensity"] * 0.5 # 50% background - - # Add statistical uncertainties (10% of intensity) - dataset["bg_intensity_stat_uncert"] = dataset["bg_intensity"] * 0.1 - - return dataset - - -# ============================================================================= -# UNIT TESTS FOR INDIVIDUAL FUNCTIONS -# ============================================================================= - - -class TestLoadEfficiencyData: - """Tests for the load_efficiency_data function.""" - - def test_load_efficiency_data_with_files(self, tmp_path): - """Test loading efficiency data when files are present.""" - # Create temporary efficiency files - eff_file1 = tmp_path / "efficiency-factor_v001.csv" - eff_file2 = tmp_path / "efficiency-factor_v002.csv" - - # Create sample data - data1 = pd.DataFrame( - { - "Date": [np.datetime64("2025-01-01")], - "E-Step1_eff": [0.8], - "E-Step2_eff": [0.82], - "E-Step3_eff": [0.84], - "E-Step4_eff": [0.86], - "E-Step5_eff": [0.88], - "E-Step6_eff": [0.90], - "E-Step7_eff": [0.92], - } - ) - - data2 = pd.DataFrame( - { - "Date": [np.datetime64("2025-01-02")], - "E-Step1_eff": [0.85], - "E-Step2_eff": [0.87], - "E-Step3_eff": [0.89], - "E-Step4_eff": [0.91], - "E-Step5_eff": [0.93], - "E-Step6_eff": [0.95], - "E-Step7_eff": [0.97], - } - ) - - # Save to CSV - data1.to_csv(eff_file1, index=False) - data2.to_csv(eff_file2, index=False) - - # Mock the ancillary file reader - with patch( - "imap_processing.lo.l2.lo_l2.lo_ancillary.read_ancillary_file" - ) as mock_read: - mock_read.side_effect = [data1, data2] - - # Test the function - result = load_efficiency_data([str(eff_file1), str(eff_file2)]) - - # Verify results - assert len(result) == 2 - assert "Date" in result.columns - assert "E-Step1_eff" in result.columns - assert mock_read.call_count == 2 - - def test_load_efficiency_data_no_files(self): - """Test loading efficiency data when no files are present.""" - result = load_efficiency_data([]) - - assert isinstance(result, pd.DataFrame) - assert result.empty - - def test_load_efficiency_data_non_efficiency_files(self): - """Test that non-efficiency files are ignored.""" - files = ["some_other_file.csv", "another_file.txt"] - result = load_efficiency_data(files) - - assert isinstance(result, pd.DataFrame) - assert result.empty - - -class TestReduceGeometricFactor: - """Tests for the reduce_geometric_factor_dataset function.""" - - @pytest.mark.parametrize("species", ["h", "o"]) - @pytest.mark.parametrize("esa_mode", [0, 1]) - def test_reduce_geometric_factor_dataset( - self, - species, - esa_mode, +def full_map(one_pointing): + """The full-spin map of one pointing, with the sky pointing mocked.""" + with patch( + "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el", + side_effect=identity_pointing, ): - """Test functionality of reduce_geometric_factor_dataset with real gf data.""" + (dataset,) = lo_l2(as_dependencies(one_pointing), [], FULL_DESCRIPTOR) + return dataset, one_pointing + + +class TestMapStructure: + """The shape and contents of the produced map.""" + + expected_variables = ( + "ena_count", + "exposure_factor", + "ena_count_rate", + "ena_count_rate_stat_uncert", + "ena_intensity", + "ena_intensity_stat_uncert", + "ena_intensity_sys_err", + "ena_intensity_sys_err_plus", + "ena_intensity_sys_err_minus", + "bg_rate", + "bg_rate_stat_uncert", + "bg_intensity", + "bg_intensity_stat_uncert", + ) - result = reduce_geometric_factor_dataset(species, esa_mode) + def test_map_dims_and_source(self, full_map): + """The map is a 6 degree rectangular map of the descriptor.""" + dataset, _ = full_map - # Verify it returns an xarray Dataset - assert isinstance(result, xr.Dataset) + assert dataset.attrs["Logical_source"] == f"imap_lo_l2_{FULL_DESCRIPTOR}" + assert dict(dataset.sizes) == { + "epoch": 1, + "energy": N_ESA, + "longitude": 60, + "latitude": 30, + } - # Verify it has 7 energy steps - assert len(result["Observed_E-Step"]) == 7 + def test_map_variables(self, full_map): + """Every map variable is on the (epoch, energy, sky) grid.""" + dataset, _ = full_map + + for variable in self.expected_variables: + assert variable in dataset.data_vars, f"missing {variable}" + assert dataset[variable].dims == ( + "epoch", + "energy", + "longitude", + "latitude", + ) - # Verify the index is 1-7 - expected_indices = list(range(1, 8)) - np.testing.assert_array_equal( - result["Observed_E-Step"].values, expected_indices - ) + def test_energy_coordinate(self, full_map): + """The energy coordinate and its widths come from the ESA constants.""" + dataset, _ = full_map - # Verify that incident_E-Step == Observed_E-Step - np.testing.assert_array_equal( - result["incident_E-Step"].values, result["Observed_E-Step"].values + np.testing.assert_allclose( + dataset["energy"].values, LoConstants.ESA_ENERGY[:N_ESA] ) - - @patch("imap_processing.lo.l2.lo_l2.load_geometric_factor_data") - def test_reduce_geometric_factor_dataset_hydrogen( - self, mock_load_gf, sample_geometric_factor_data - ): - """Test that hydrogen geometric factors are correctly loaded and processed.""" - h_gf_data, _ = sample_geometric_factor_data - mock_load_gf.return_value = h_gf_data - - result = reduce_geometric_factor_dataset("h", esa_mode=0) - - # Check that hydrogen-specific columns are present - assert "Cntr_E" in result.data_vars - assert "GF_Trpl_H" in result.data_vars - assert "GF_Trpl_H_unc_minus" in result.data_vars - assert "GF_Trpl_H_unc_plus" in result.data_vars - - # Verify energy values match expected for hydrogen - expected_energies = [0.01 * (i + 1) for i in range(7)] - np.testing.assert_array_almost_equal(result["Cntr_E"].values, expected_energies) - - @patch("imap_processing.lo.l2.lo_l2.load_geometric_factor_data") - def test_reduce_geometric_factor_dataset_oxygen( - self, mock_load_gf, sample_geometric_factor_data - ): - """Test that oxygen geometric factors are correctly loaded and processed.""" - _, o_gf_data = sample_geometric_factor_data - mock_load_gf.return_value = o_gf_data - - result = reduce_geometric_factor_dataset("o", esa_mode=0) - - # Check that oxygen-specific columns are present - assert "Cntr_E" in result.data_vars - assert "GF_Trpl_O" in result.data_vars - assert "GF_Trpl_O_unc_minus" in result.data_vars - assert "GF_Trpl_O_unc_plus" in result.data_vars - - # Verify energy values match expected for oxygen - expected_energies = [0.015 * (i + 1) for i in range(7)] - np.testing.assert_array_almost_equal(result["Cntr_E"].values, expected_energies) - - @patch("imap_processing.lo.l2.lo_l2.load_geometric_factor_data") - def test_reduce_geometric_factor_dataset_esa_mode_filtering( - self, mock_load_gf, sample_geometric_factor_data - ): - """Test that ESA mode filtering works correctly.""" - h_gf_data, _ = sample_geometric_factor_data - - # Create data with two ESA modes - h_gf_mode_0 = h_gf_data.copy() - h_gf_mode_1 = h_gf_data.copy() - - # Modify mode 1 data to have different GF values - for col in ["GF_Trpl_H", "GF_Dbl_all", "GF_Trpl_all"]: - h_gf_mode_1[col] = h_gf_mode_1[col] * 1.5 - h_gf_mode_1["esa_mode"] = 1 - - combined_data = pd.concat([h_gf_mode_0, h_gf_mode_1], ignore_index=True) - mock_load_gf.return_value = combined_data - - # Get data for mode 0 - result_mode_0 = reduce_geometric_factor_dataset("h", esa_mode=0) - - # Get data for mode 1 - result_mode_1 = reduce_geometric_factor_dataset("h", esa_mode=1) - - # Verify that mode 1 has different (larger) GF values - assert np.all( - result_mode_1["GF_Trpl_H"].values > result_mode_0["GF_Trpl_H"].values + np.testing.assert_allclose( + dataset["energy_delta_plus"].values, LoConstants.ESA_ENERGY_DELTA[0] ) - # Verify the ratio is approximately 1.5 - ratio = result_mode_1["GF_Trpl_H"].values / result_mode_0["GF_Trpl_H"].values - np.testing.assert_array_almost_equal(ratio, np.ones(7) * 1.5) - - @patch("imap_processing.lo.l2.lo_l2.load_geometric_factor_data") - def test_reduce_geometric_factor_dataset_duplicate_removal( - self, mock_load_gf, sample_geometric_factor_data - ): - """Test that duplicate Observed_E-Step values are handled correctly.""" - h_gf_data, _ = sample_geometric_factor_data - - # Add duplicate rows with same Observed_E-Step but different incident_E-Step - duplicates = h_gf_data.copy() - duplicates["incident_E-Step"] = [ - i + 8 for i in range(7) - ] # Different incident steps - - combined_data = pd.concat([h_gf_data, duplicates], ignore_index=True) - mock_load_gf.return_value = combined_data - - result = reduce_geometric_factor_dataset("h", esa_mode=0) - - # Should still have exactly 7 energy steps (duplicates removed) - assert len(result["Observed_E-Step"]) == 7 + def test_map_writes_to_cdf(self, full_map): + """The map can be written out as a valid CDF.""" + dataset, _ = full_map + dataset.attrs["Data_version"] = "001.0001" + dataset.attrs["Start_date"] = "20260101" - # Verify no duplicate indices - assert len(np.unique(result["Observed_E-Step"].values)) == 7 + cdf_path = write_cdf(dataset) - @patch("imap_processing.lo.l2.lo_l2.load_geometric_factor_data") - def test_reduce_geometric_factor_dataset_invalid_species(self, mock_load_gf): - """Test that invalid species raises appropriate error.""" - # Mock should raise ValueError for invalid species - mock_load_gf.side_effect = ValueError( - "Geometric factors only available for 'h' and 'o', got 'invalid'" - ) - - with pytest.raises(ValueError, match="Geometric factors only available"): - reduce_geometric_factor_dataset("invalid", esa_mode=0) + assert cdf_path.exists() + assert cdf_path.name == f"imap_lo_l2_{FULL_DESCRIPTOR}_20260101_v001.0001.cdf" + assert dict(load_cdf(cdf_path).sizes) == dict(dataset.sizes) - @patch("imap_processing.lo.l2.lo_l2.load_geometric_factor_data") - def test_reduce_geometric_factor_dataset_no_esa_mode_column( - self, mock_load_gf, sample_geometric_factor_data - ): - """Test handling when esa_mode column is missing from data.""" - h_gf_data, _ = sample_geometric_factor_data - - # Remove esa_mode column if it exists - if "esa_mode" in h_gf_data.columns: - h_gf_data = h_gf_data.drop(columns=["esa_mode"]) - - mock_load_gf.return_value = h_gf_data - - # Should still work, just won't filter by esa_mode - result = reduce_geometric_factor_dataset("h", esa_mode=0) - - # Should return 7 energy steps - assert len(result["Observed_E-Step"]) == 7 - assert isinstance(result, xr.Dataset) - - -class TestNormalizePsetCoordinates: - """Tests for the normalize_pset_coordinates function.""" - - @pytest.mark.parametrize("species", ["h", "o"]) - def test_normalize_coordinates_basic(self, species): - """Test basic coordinate normalization for a specific species.""" - # Create a pset with the specified species - pset = xr.Dataset( - { - f"{species}_counts": (PSET_DIMS, np.ones(PSET_SHAPE)), - "exposure_time": (PSET_DIMS, np.ones(PSET_SHAPE)), - f"{species}_background_rates": (PSET_DIMS, np.ones(PSET_SHAPE) * 0.1), - f"{species}_background_rates_stat_uncert": ( - PSET_DIMS, - np.ones(PSET_SHAPE) * 0.01, - ), - f"{species}_background_rates_sys_err": ( - PSET_DIMS, - np.ones(PSET_SHAPE) * 0.02, - ), - }, - coords={ - "epoch": [8.1794907049e17], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, - }, - ) - result = normalize_pset_coordinates(pset, species) +class TestAccumulation: + """What the map accumulates from its inputs.""" - # Check that dimensions were renamed - assert "energy" in result.dims - assert "esa_energy_step" not in result.dims + def test_counts_are_conserved(self, full_map): + """Every in-window count lands somewhere on the map.""" + dataset, pointing = full_map - # Check that energy coordinate is present - assert "energy" in result.coords - expected_energies = ( - np.array([0.01633, 0.03047, 0.05576, 0.10626, 0.20004, 0.40496, 0.78729]) - if species == "h" - else np.array([0.01919, 0.03675, 0.07121, 0.14141, 0.274, 0.58503, 1.13506]) - ) - np.testing.assert_array_equal(result.coords["energy"], expected_energies) - - # Check that old coordinate variable was dropped - assert "esa_energy_step" not in result.variables - - # Check that variables were renamed - assert "counts" in result.data_vars - assert "exposure_factor" in result.data_vars - assert "bg_rate" in result.data_vars - assert "bg_rate_stat_uncert" in result.data_vars - assert "bg_rate_sys_err" in result.data_vars - - # Check that old variable names are gone - assert f"{species}_counts" not in result.data_vars - assert "exposure_time" not in result.data_vars - - @pytest.mark.parametrize("species", ["doubles", "triples"]) - def test_normalize_coordinates_no_background(self, species): - """Test normalization for species without background rates.""" - # Create a pset with only counts and exposure time - pset = xr.Dataset( - { - f"{species}_counts": (PSET_DIMS, np.ones(PSET_SHAPE)), - "exposure_time": (PSET_DIMS, np.ones(PSET_SHAPE)), - }, - coords={ - "epoch": [8.1794907049e17], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, - }, - ) + per_energy = dataset["ena_count"].values.sum(axis=(0, 2, 3)) + np.testing.assert_allclose(per_energy, pointing["expected_counts"]) - # For species without background rates, the function should fail - # because geometric factors can only be retrieved for "h" and "o" - with pytest.raises(ValueError, match="Geometric factors only available"): - normalize_pset_coordinates(pset, species) - - def test_normalize_coordinates_removes_old_coordinate(self): - """Test that old esa_energy_step coordinate is removed.""" - species = "h" - pset = xr.Dataset( - { - f"{species}_counts": (PSET_DIMS, np.ones(PSET_SHAPE)), - "exposure_time": (PSET_DIMS, np.ones(PSET_SHAPE)), - f"{species}_background_rates": (PSET_DIMS, np.ones(PSET_SHAPE) * 0.1), - f"{species}_background_rates_stat_uncert": ( - PSET_DIMS, - np.ones(PSET_SHAPE) * 0.01, - ), - f"{species}_background_rates_sys_err": ( - PSET_DIMS, - np.ones(PSET_SHAPE) * 0.02, - ), - "esa_energy_step_var": xr.DataArray([1, 2, 3, 4, 5, 6, 7]), # Variable - }, - coords={ - "epoch": [8.1794907049e17], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, - }, - ) + def test_exposure_is_conserved(self, full_map): + """Every in-window second of exposure lands somewhere on the map.""" + dataset, pointing = full_map - result = normalize_pset_coordinates(pset, species) + per_energy = dataset["exposure_factor"].values.sum(axis=(0, 2, 3)) + np.testing.assert_allclose(per_energy, pointing["expected_exposure"], rtol=1e-6) - # Should remove the esa_energy_step coordinate and variable - assert "esa_energy_step" not in result.variables - assert "esa_energy_step_var" in result.variables # Data variable should remain + def test_out_of_goodtime_epochs_are_excluded(self, full_map): + """The 999-per-bin epochs outside the good times do not reach the map.""" + dataset, pointing = full_map + per_energy = dataset["ena_count"].values.sum(axis=(0, 2, 3)) + assert per_energy.max() < 999.0 * N_SPIN_BINS + np.testing.assert_allclose(per_energy, pointing["expected_counts"]) -class TestAddEfficiencyFactorsToPset: - """Tests for the add_efficiency_factors_to_pset function.""" + def test_pointings_accumulate(self, one_pointing): + """Two pointings contribute twice the counts of one.""" + other = make_pointing(repointing="repoint00101", seed=7) - def test_add_efficiency_factors_with_data( - self, minimal_pset, sample_efficiency_data - ): - """Test adding efficiency factors when data is available.""" - # Set the epoch to match our sample data - pset = minimal_pset.copy() - # Convert date to TT2000 nanoseconds (approximate) - epoch_ns = 8.1794907049e17 # This should correspond to 2025-01-01 - pset = pset.assign_coords(epoch=[epoch_ns]) - - with ( - patch("imap_processing.lo.l2.lo_l2.ttj2000ns_to_et") as mock_ttj2000_to_et, - patch("imap_processing.lo.l2.lo_l2.et_to_datetime64") as mock_et_to_dt64, + with patch( + "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el", + side_effect=identity_pointing, ): - # Mock the time conversion - mock_ttj2000_to_et.return_value = 1234567890.0 - mock_et_to_dt64.return_value = np.datetime64("2025-01-01") - - result = add_efficiency_factors_to_pset(pset, sample_efficiency_data) - - # Check that efficiency was added - assert "efficiency" in result.data_vars - assert result["efficiency"].dims == ("energy",) - assert len(result["efficiency"]) == 7 - - # Check efficiency values match expected (first row of sample data) - expected_eff = [0.8, 0.82, 0.84, 0.86, 0.88, 0.90, 0.92] - np.testing.assert_array_almost_equal( - result["efficiency"].values, expected_eff - ) - - def test_add_efficiency_factors_no_data(self, minimal_pset): - """Test adding efficiency factors when no data is available.""" - empty_df = pd.DataFrame() - - result = add_efficiency_factors_to_pset(minimal_pset, empty_df) - - # Should create unity efficiency - assert "efficiency" in result.data_vars - np.testing.assert_array_equal(result["efficiency"].values, np.ones(7)) + (one,) = lo_l2(as_dependencies(one_pointing), [], FULL_DESCRIPTOR) + (both,) = lo_l2(as_dependencies(one_pointing, other), [], FULL_DESCRIPTOR) - def test_add_efficiency_factors_missing_date( - self, minimal_pset, sample_efficiency_data - ): - """Test error when efficiency factor not found for date.""" - pset = minimal_pset.copy() + np.testing.assert_allclose( + both["ena_count"].values.sum(axis=(0, 2, 3)), + one["ena_count"].values.sum(axis=(0, 2, 3)) + other["expected_counts"], + ) - with ( - patch("imap_processing.lo.l2.lo_l2.ttj2000ns_to_et") as mock_ttj2000_to_et, - patch("imap_processing.lo.l2.lo_l2.et_to_datetime64") as mock_et_to_dt64, + def test_ram_map_keeps_half_the_spin(self, one_pointing): + """A ram map takes fewer counts than the full spin it is cut from.""" + with patch( + "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el", + side_effect=identity_pointing, ): - # Mock conversion to a date not in sample data - mock_ttj2000_to_et.return_value = 1234567890.0 - mock_et_to_dt64.return_value = np.datetime64("2025-12-31") - - with pytest.raises(ValueError, match="No efficiency factor found"): - add_efficiency_factors_to_pset(pset, sample_efficiency_data) - - -class TestCalculateEfficiencyCorrectedQuantities: - """Tests for the calculate_efficiency_corrected_quantities function.""" - - def test_calculate_efficiency_corrected_quantities(self): - """Test calculation of efficiency-corrected quantities.""" - # Create a dataset with the current generic variable names - pset = xr.Dataset( - { - "counts": (("energy",), np.ones(7) * 10), # 10 counts - "exposure_factor": (("energy",), np.ones(7) * 1.0), # 1 second - "bg_rate": (("energy",), np.ones(7) * 0.1), # 0.1 counts/s - "bg_rate_stat_uncert": (("energy",), np.ones(7) * 0.01), # uncertainty - "bg_rate_sys_err": (("energy",), np.ones(7) * 0.02), # systematic - "efficiency": ( - ("energy",), - np.array([0.8, 0.85, 0.9, 0.95, 0.88, 0.92, 0.87]), - ), - }, - coords={"energy": list(range(7))}, - ) + (full,) = lo_l2(as_dependencies(one_pointing), [], FULL_DESCRIPTOR) + (ram,) = lo_l2(as_dependencies(one_pointing), [], RAM_DESCRIPTOR) - result = calculate_efficiency_corrected_quantities(pset) + full_counts = full["ena_count"].values.sum() + ram_counts = ram["ena_count"].values.sum() + assert 0 < ram_counts < full_counts - # Check that corrected quantities were added - assert "counts_over_eff" in result.data_vars - assert "counts_over_eff_squared" in result.data_vars - assert "bg_rate_exposure_factor" in result.data_vars - assert "bg_rate_stat_uncert_exposure_factor2" in result.data_vars - assert "bg_rate_sys_err_exposure_factor" in result.data_vars - # Check dimensions - assert result["counts_over_eff"].dims == pset["counts"].dims +class TestRatesAndIntensities: + """The maths turning accumulated counts into intensities.""" - # Check that division by efficiency happened correctly - expected_over_eff = pset["counts"] / pset["efficiency"] - xr.testing.assert_allclose(result["counts_over_eff"], expected_over_eff) + def test_rate_and_intensity(self, full_map): + """Where exposed, rate = counts/exposure and intensity = rate/(G*E).""" + dataset, _ = full_map - # Check that division by efficiency squared happened correctly - expected_over_eff_sq = pset["counts"] / (pset["efficiency"] ** 2) - xr.testing.assert_allclose( - result["counts_over_eff_squared"], expected_over_eff_sq - ) + counts = dataset["ena_count"].values + exposure = dataset["exposure_factor"].values + exposed = exposure > 0 + assert exposed.any() - # Check background rate calculations - expected_bg_exposure = pset["bg_rate"] * pset["exposure_factor"] - xr.testing.assert_allclose( - result["bg_rate_exposure_factor"], expected_bg_exposure + np.testing.assert_allclose( + dataset["ena_count_rate"].values[exposed], + counts[exposed] / exposure[exposed], + rtol=1e-5, ) + assert np.all(dataset["ena_count_rate"].values[~exposed] == 0) - expected_bg_uncert_exposure = ( - pset["bg_rate_stat_uncert"] ** 2 * pset["exposure_factor"] ** 2 + geometric_factor = ( + np.array(LoConstants.GEO_FACTOR[:N_ESA]) * LoConstants.GEO_FACTOR_SCALE ) - xr.testing.assert_allclose( - result["bg_rate_stat_uncert_exposure_factor2"], expected_bg_uncert_exposure + energy = np.array(LoConstants.ESA_ENERGY[:N_ESA]) + expected = dataset["ena_count_rate"] / xr.DataArray( + geometric_factor * energy, dims=["energy"] ) - - expected_bg_sys_err_exposure = pset["bg_rate_sys_err"] * pset["exposure_factor"] - xr.testing.assert_allclose( - result["bg_rate_sys_err_exposure_factor"], expected_bg_sys_err_exposure + np.testing.assert_allclose( + dataset["ena_intensity"].values[exposed], + expected.values[exposed], + rtol=1e-5, ) + def test_statistical_uncertainty_is_poisson(self, full_map): + """The rate uncertainty is the Poisson count error over the exposure.""" + dataset, _ = full_map -class TestCalculateRates: - """Tests for the calculate_rates function.""" - - def test_calculate_rates_basic(self, sample_sky_map_dataset): - """Test rate calculation with current implementation.""" - result = calculate_rates(sample_sky_map_dataset) - - # Check that the expected output variables were created - assert "ena_count_rate" in result.data_vars - assert "ena_count_rate_stat_uncert" in result.data_vars - - # Check dimensions match input - assert result["ena_count_rate"].dims == sample_sky_map_dataset["counts"].dims - - # Check rate calculation (counts / exposure_factor) - # With counts=10 and exposure_factor=1, rate should be 10 - expected_rate = ( - sample_sky_map_dataset["counts"] / sample_sky_map_dataset["exposure_factor"] - ) - xr.testing.assert_allclose(result["ena_count_rate"], expected_rate) - - # Check uncertainty calculation (sqrt(counts) / exposure_factor) - expected_uncert = ( - np.sqrt(sample_sky_map_dataset["counts"]) - / sample_sky_map_dataset["exposure_factor"] - ) - xr.testing.assert_allclose( - result["ena_count_rate_stat_uncert"], expected_uncert - ) + counts = dataset["ena_count"].values + exposure = dataset["exposure_factor"].values + exposed = exposure > 0 - def test_calculate_rates_missing_variables(self): - """Rate calculation when required variables are missing.""" - # Create dataset missing required variables - dataset = xr.Dataset( - { - "counts": (("epoch", "energy"), np.ones((1, 7)) * 5), - # Missing exposure_factor - } + np.testing.assert_allclose( + dataset["ena_count_rate_stat_uncert"].values[exposed], + np.sqrt(counts[exposed]) / exposure[exposed], + rtol=1e-5, ) - # Should raise KeyError for missing exposure_factor - with pytest.raises(KeyError, match="exposure_factor"): - calculate_rates(dataset) - - -class TestCalculateIntensities: - """Tests for the calculate_intensities function.""" + def test_background_rate(self, full_map): + """The background rate is the input rate wherever the map was exposed.""" + dataset, pointing = full_map - def test_calculate_intensities_basic(self, sample_dataset_with_geometric_factors): - """Test intensity calculation with current implementation.""" - result = calculate_intensities(sample_dataset_with_geometric_factors) + exposure = dataset["exposure_factor"].values + for energy_index in range(N_ESA): + exposed = exposure[0, energy_index] > 0 + bg_rate = dataset["bg_rate"].values[0, energy_index] + np.testing.assert_allclose( + bg_rate[exposed], pointing["background"][energy_index], rtol=1e-5 + ) + assert np.all(bg_rate[~exposed] == 0) - # Check that the expected output variables were created - assert "ena_intensity" in result.data_vars - assert "ena_intensity_stat_uncert" in result.data_vars - assert "ena_intensity_sys_err" in result.data_vars + def test_systematic_error_bounds(self, full_map): + """The systematic error is bracketed by the G-factor excursions.""" + dataset, _ = full_map - # Check intensity calculation: - # counts_over_eff / (geometric_factor * energy * exposure_factor) - # 100 / (1e-4 * 0.1 * 1.0) = 100 / 1e-5 = 1e7 - expected_intensity = sample_dataset_with_geometric_factors[ - "counts_over_eff" - ] / ( - sample_dataset_with_geometric_factors["geometric_factor"] - * sample_dataset_with_geometric_factors["energy"] - * sample_dataset_with_geometric_factors["exposure_factor"] - ) - xr.testing.assert_allclose(result["ena_intensity"], expected_intensity) - - # Check statistical uncertainty calculation - expected_stat_uncert = np.sqrt( - sample_dataset_with_geometric_factors["counts_over_eff_squared"] - ) / ( - sample_dataset_with_geometric_factors["geometric_factor"] - * sample_dataset_with_geometric_factors["energy"] - * sample_dataset_with_geometric_factors["exposure_factor"] - ) - xr.testing.assert_allclose( - result["ena_intensity_stat_uncert"], expected_stat_uncert - ) + intensity = dataset["ena_intensity"].values + plus = dataset["ena_intensity_sys_err_plus"].values + minus = dataset["ena_intensity_sys_err_minus"].values + symmetric = dataset["ena_intensity_sys_err"].values + lit = intensity > 0 - # Check systematic uncertainty calculation - gf = sample_dataset_with_geometric_factors["geometric_factor"] - dg_minus = sample_dataset_with_geometric_factors[ - "geometric_factor_stat_uncert_minus" - ] - dg_plus = sample_dataset_with_geometric_factors[ - "geometric_factor_stat_uncert_plus" - ] - expected_sys_err_plus = result["ena_intensity"] * dg_minus / (gf - dg_minus) - expected_sys_err_minus = result["ena_intensity"] * dg_plus / (gf + dg_plus) - expected_sys_err = np.sqrt(expected_sys_err_minus * expected_sys_err_plus) - xr.testing.assert_allclose(result["ena_intensity_sys_err"], expected_sys_err) - - def test_calculate_intensities_missing_variables(self): - """Test intensity calculation when required variables are missing.""" - # Create dataset missing geometric_factor - dataset = xr.Dataset( - { - "counts_over_eff": (("energy",), np.ones(7) * 100), - "counts_over_eff_squared": (("energy",), np.ones(7) * 100), - "exposure_factor": (("energy",), np.ones(7) * 1.0), - "energy": (("energy",), np.ones(7) * 0.1), - # Missing geometric_factor - } + assert np.all(plus[lit] > 0) + assert np.all(minus[lit] > 0) + # The symmetric error is the geometric mean of the two excursions + np.testing.assert_allclose( + symmetric[lit], np.sqrt(plus[lit] * minus[lit]), rtol=1e-4 ) + # The lower G-factor bound gives the bigger flux excursion + assert np.all(plus[lit] >= minus[lit]) - # Should raise KeyError for missing geometric_factor - with pytest.raises(KeyError, match="geometric_factor"): - calculate_intensities(dataset) - - -class TestCalculateBackgrounds: - """Tests for the calculate_backgrounds function.""" - - def test_calculate_backgrounds_basic( - self, sample_dataset_with_background_intermediates - ): - """Test basic background calculations with standard data.""" - dataset = sample_dataset_with_background_intermediates - - result = calculate_backgrounds(dataset) - # Check that background variables were calculated - assert "bg_rate" in result.data_vars - assert "bg_rate_stat_uncert" in result.data_vars - assert "bg_rate_sys_err" in result.data_vars +class TestGeometry: + """The spin-angle to sky-pixel geometry.""" - # Check background rate calculation - # bg_rate_exposure_factor / exposure_factor = 0.2 / 1.0 = 0.2 - expected_bg_rate = ( - dataset["bg_rate_exposure_factor"] / dataset["exposure_factor"] - ) - xr.testing.assert_allclose(result["bg_rate"], expected_bg_rate) + def test_dps_spin_angles_carry_the_offset(self): + """The hardware spin bins are rotated onto the instrument frame.""" + angles = _dps_spin_angles() - # Check statistical uncertainty calculation - # sqrt(bg_rate_stat_uncert_exposure_factor2) / exposure_factor^2 - expected_stat_uncert = np.sqrt( - dataset["bg_rate_stat_uncert_exposure_factor2"] - / dataset["exposure_factor"] ** 2 - ) - xr.testing.assert_allclose(result["bg_rate_stat_uncert"], expected_stat_uncert) + assert angles.size == N_SPIN_BINS + # IMAP-Lo sits 60 degrees from the spacecraft spin pulse, so bin 0's + # center (3 degrees) becomes 63 degrees in the despun frame. + np.testing.assert_allclose(angles[0], 63.0) + np.testing.assert_allclose(np.diff(np.sort(angles)), 6.0) - # Check systematic error calculation - expected_sys_err = ( - dataset["bg_rate_sys_err_exposure_factor"] / dataset["exposure_factor"] - ) - xr.testing.assert_allclose(result["bg_rate_sys_err"], expected_sys_err) - - def test_calculate_backgrounds_zero_exposure(self): - """Test background calculations with zero exposure time.""" - dataset = xr.Dataset( - { - "bg_rate_exposure_factor": ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.2, - ), - "bg_rate_stat_uncert_exposure_factor2": ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.004, - ), - "bg_rate_sys_err_exposure_factor": ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.05, - ), - "exposure_factor": (("epoch", "energy"), np.zeros((1, 7))), - "geometric_factor": (("energy",), np.ones(7) * 1e-4), - "geometric_factor_stat_uncert_minus": (("energy",), np.ones(7) * 1e-5), - "geometric_factor_stat_uncert_plus": (("energy",), np.ones(7) * 3e-5), - }, - coords={"epoch": [8.1794907049e17], "energy": list(range(7))}, - ) + def test_pixel_indices_match_the_map_grid(self): + """Directions are placed in the pixel whose center they are nearest.""" + sky_map = MapDescriptor.from_string(FULL_DESCRIPTOR).to_empty_map() + centers = sky_map.az_el_points.values - result = calculate_backgrounds(dataset) + pixels = _pixel_indices(sky_map, centers[:, 0], centers[:, 1]) - # Should handle division by zero gracefully - assert "bg_rate" in result.data_vars - assert "bg_rate_stat_uncert" in result.data_vars - assert "bg_rate_sys_err" in result.data_vars - # Results should be infinite where exposure time is zero - assert np.all(np.isinf(result["bg_rate"].values)) - assert np.all(np.isinf(result["bg_rate_stat_uncert"].values)) + np.testing.assert_array_equal(pixels, np.arange(sky_map.num_points)) + def test_pixel_indices_wrap_longitude(self): + """Longitudes outside 0-360 wrap onto the map.""" + sky_map = MapDescriptor.from_string(FULL_DESCRIPTOR).to_empty_map() -class TestCalculateFluxCorrections: - """Tests for the calculate_flux_corrections function.""" + wrapped = _pixel_indices(sky_map, np.array([-357.0]), np.array([0.0])) + direct = _pixel_indices(sky_map, np.array([3.0]), np.array([0.0])) - def test_calculate_flux_corrections_basic( - self, sample_dataset_with_intensities, lo_flux_factors_file - ): - """Test basic flux correction calculation.""" - # Make a copy to avoid modifying the original fixture - original_dataset = sample_dataset_with_intensities.copy(deep=True) + np.testing.assert_array_equal(wrapped, direct) - # Run flux correction - result = calculate_flux_corrections(original_dataset, lo_flux_factors_file) + @pytest.mark.parametrize( + "spin_phase, expected", + [("full", 60), ("ram", 30), ("anti", 30)], + ) + def test_spin_phase_mask(self, spin_phase, expected): + """Ram and anti-ram split the spin; a full map keeps all of it.""" + angles = _dps_spin_angles() - # Verify that the function returns a dataset - assert isinstance(result, xr.Dataset) + mask = _spin_phase_mask(angles, PIVOT, spin_phase) - # Verify that intensity variables are present - assert "ena_intensity" in result.data_vars - assert "ena_intensity_stat_uncert" in result.data_vars + assert mask.sum() == expected - # Verify that data shape is preserved - original_shape = sample_dataset_with_intensities["ena_intensity"].shape - assert result["ena_intensity"].shape == original_shape + def test_spin_phase_mask_rejects_unknown(self): + """An unknown spin phase is an error, not a silently empty map.""" + with pytest.raises(ValueError, match="Invalid spin phase"): + _spin_phase_mask(_dps_spin_angles(), PIVOT, "sideways") - # Check that corrections were applied by comparing to the original fixture - # (not the potentially modified copy) - original_intensity = sample_dataset_with_intensities["ena_intensity"].values - corrected_intensity = result["ena_intensity"].values - # Check for meaningful differences - relative_diff = np.abs( - (corrected_intensity - original_intensity) / original_intensity - ) - max_relative_diff = np.max(relative_diff) - # Should have at least 10% change somewhere - assert max_relative_diff > 0.1, ( - f"Max relative difference was only {max_relative_diff}" - ) +class TestInputGrouping: + """Sorting the input products into pointings.""" - # Verify that uncertainties were also corrected - original_uncert = sample_dataset_with_intensities[ - "ena_intensity_stat_uncert" - ].values - corrected_uncert = result["ena_intensity_stat_uncert"].values - uncert_relative_diff = np.abs( - (corrected_uncert - original_uncert) / original_uncert - ) - max_uncert_diff = np.max(uncert_relative_diff) - # Should have at least 10% change in uncertainties too - assert max_uncert_diff > 0.1, ( - f"Max uncertainty relative difference was only {max_uncert_diff}" - ) + def test_inputs_are_grouped_by_repointing(self, one_pointing): + """Each pointing's three products are grouped together.""" + other = make_pointing(repointing="repoint00101") - def test_calculate_flux_corrections_preserves_other_vars( - self, sample_dataset_with_intensities, lo_flux_factors_file - ): - """Test that flux correction preserves other variables in the dataset.""" - # Add an extra variable to the dataset - sample_dataset_with_intensities["extra_var"] = (("energy",), np.ones(7)) + pointings = _group_inputs_by_pointing(as_dependencies(one_pointing, other)) - result = calculate_flux_corrections( - sample_dataset_with_intensities, lo_flux_factors_file - ) + assert set(pointings) == {"repoint00100", "repoint00101"} + assert pointings["repoint00100"][0] is one_pointing[GOODTIMES] + assert pointings["repoint00100"][2] is one_pointing[HISTRATES] - # Verify that other variables are preserved - assert "extra_var" in result.data_vars - np.testing.assert_array_equal( - result["extra_var"].values, - sample_dataset_with_intensities["extra_var"].values, - ) + def test_incomplete_pointings_are_dropped(self, one_pointing, caplog): + """A pointing missing one of the three products cannot be mapped.""" + dependencies = as_dependencies(one_pointing) + dependencies[BGRATES] = [] - def test_calculate_flux_corrections_energy_dimension_handling( - self, lo_flux_factors_file - ): - """Test that flux correction properly handles energy dimension reshaping.""" - # Create a dataset with different spatial dimensions - n_energy = 7 - n_x, n_y = 12, 8 # Different spatial dimensions - - energy_values = np.array([16.35, 30.56, 56.4, 105, 199.8, 407.5, 795.3]) - - coords = { - "epoch": [8.1794907049e17], - "energy": energy_values, - "x": np.arange(n_x), - "y": np.arange(n_y), - } + pointings = _group_inputs_by_pointing(dependencies) - # Create intensity values with energy-dependent structure (power law) - intensity_values = np.ones((1, n_energy, n_x, n_y)) - for i in range(n_energy): - intensity_values[0, i, :, :] = 1e6 * (energy_values[i] / 100.0) ** (-2.0) - uncert_values = intensity_values * 0.1 + assert pointings == {} + assert "imap_lo_l1b_bgrates" in caplog.text - original_dataset = xr.Dataset(coords=coords) - original_dataset["ena_intensity"] = ( - ("epoch", "energy", "x", "y"), - intensity_values.copy(), - ) - original_dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "x", "y"), - uncert_values.copy(), - ) - original_dataset["bg_intensity"] = original_dataset["ena_intensity"] * 0.5 - original_dataset["bg_intensity_stat_uncert"] = ( - original_dataset["bg_intensity"] * 0.1 - ) + def test_missing_product_raises(self, one_pointing): + """A map cannot be made without all three products.""" + dependencies = as_dependencies(one_pointing) + del dependencies[HISTRATES] - # Run flux correction on a copy - dataset_copy = original_dataset.copy(deep=True) - result = calculate_flux_corrections(dataset_copy, lo_flux_factors_file) - - # Verify shape is preserved - assert result["ena_intensity"].shape == (1, n_energy, n_x, n_y) - assert result["ena_intensity_stat_uncert"].shape == (1, n_energy, n_x, n_y) - - # Verify corrections were applied by checking for meaningful differences - original_values = original_dataset["ena_intensity"].values - corrected_values = result["ena_intensity"].values - relative_diff = np.abs((corrected_values - original_values) / original_values) - max_relative_diff = np.max(relative_diff) - # Should have at least 10% change somewhere (flux corrections are significant) - assert max_relative_diff > 0.1, ( - f"Max relative difference was only {max_relative_diff}" - ) + with pytest.raises(KeyError, match=HISTRATES): + _group_inputs_by_pointing(dependencies) -class TestCalculateSputteringCorrections: - """Tests for the calculate_sputtering_corrections function.""" +class TestUnsupported: + """Map flavours the Lo pipeline does not make.""" - def test_calculate_sputtering_corrections_basic( - self, sample_dataset_with_sputtering_data - ): - """Test basic sputtering corrections for hydrogen and oxygen intensities.""" - h_dataset, o_dataset = sample_dataset_with_sputtering_data - - # Test with hydrogen dataset first - original_h_intensity = h_dataset["ena_intensity"].copy() - original_h_stat_uncert = h_dataset["ena_intensity_stat_uncert"].copy() - original_h_sys_err = h_dataset["ena_intensity_sys_err"].copy() - - result_h = calculate_sputtering_corrections(h_dataset, o_dataset) - - # Check that only energy levels 4 and 5 (ESA levels 5 and 6) were modified - # for hydrogen - for energy_idx in [0, 1, 2, 3, 6]: - np.testing.assert_array_equal( - result_h["ena_intensity"][0, energy_idx, :, :].values, - original_h_intensity[0, energy_idx, :, :].values, - err_msg=f"Hydrogen energy level {energy_idx} should not be modified", + def test_oxygen_not_supported(self, one_pointing): + """Only hydrogen geometric factors are defined.""" + with pytest.raises(NotImplementedError, match="species o"): + lo_l2( + as_dependencies(one_pointing), + [], + "l090-ena-o-sf-nsp-full-hae-6deg-3mo", ) - # Check that energy levels 4 and 5 were modified (should be lower) - assert np.all( - result_h["ena_intensity"][0, 4, :, :].values - < original_h_intensity[0, 4, :, :].values - ), ( - "Hydrogen energy level 4 intensity should be reduced by sputtering " - "correction" - ) - - assert np.all( - result_h["ena_intensity"][0, 5, :, :].values - < original_h_intensity[0, 5, :, :].values - ), ( - "Hydrogen energy level 5 intensity should be reduced by sputtering " - "correction" - ) - - # Check that uncertainties were also updated for levels 4 and 5 - assert not np.array_equal( - result_h["ena_intensity_stat_uncert"][0, 4, :, :].values, - original_h_stat_uncert[0, 4, :, :].values, - ), "Statistical uncertainty should be updated for hydrogen energy level 4" - - assert not np.array_equal( - result_h["ena_intensity_sys_err"][0, 4, :, :].values, - original_h_sys_err[0, 4, :, :].values, - ), "Systematic error should be updated for hydrogen energy level 4" - - # Test with oxygen dataset - original_o_intensity = o_dataset["ena_intensity"].copy() - - result_o = calculate_sputtering_corrections(o_dataset, o_dataset) - - # Check that only energy levels 4 and 5 were modified for oxygen - for energy_idx in [0, 1, 2, 3, 6]: - np.testing.assert_array_equal( - result_o["ena_intensity"][0, energy_idx, :, :].values, - original_o_intensity[0, energy_idx, :, :].values, - err_msg=f"Oxygen energy level {energy_idx} should not be modified", + def test_healpix_not_supported(self, one_pointing): + """Lo makes rectangular maps only.""" + with pytest.raises(NotImplementedError, match="HEALPix"): + lo_l2( + as_dependencies(one_pointing), + [], + "l090-ena-h-sf-nsp-full-hae-nside8-3mo", ) - - # Check that energy levels 4 and 5 were modified - assert np.all( - result_o["ena_intensity"][0, 4, :, :].values - < original_o_intensity[0, 4, :, :].values - ), "Oxygen energy level 4 intensity should be reduced by sputtering correction" - - assert np.all( - result_o["ena_intensity"][0, 5, :, :].values - < original_o_intensity[0, 5, :, :].values - ), "Oxygen energy level 5 intensity should be reduced by sputtering correction" - - def test_calculate_sputtering_corrections_equations( - self, sample_dataset_with_sputtering_data - ): - """Test that sputtering corrections follow the correct equations.""" - h_dataset, o_dataset = sample_dataset_with_sputtering_data - - # Get the subset that will be processed (energy levels 4 and 5) - o_small_dataset = o_dataset.isel(epoch=0, energy=[4, 5]) - - # Calculate expected j_o_prime (Equation 9) - expected_j_o_prime = ( - o_small_dataset["ena_intensity"] - o_small_dataset["bg_intensity"] - ) - expected_j_o_prime = expected_j_o_prime.where(expected_j_o_prime >= 0, 0) - - # Expected correction factors from the mapping document table 2 - sputter_correction_factor = np.array([0.15, 0.01]) - - # Calculate expected corrected intensity (Equation 11) for hydrogen - h_small_dataset = h_dataset.isel(epoch=0, energy=[4, 5]) - expected_corrected_intensity = ( - h_small_dataset["ena_intensity"] - - sputter_correction_factor[:, np.newaxis, np.newaxis] * expected_j_o_prime - ) - - # Run the function - result = calculate_sputtering_corrections(h_dataset, o_dataset) - - # Check that the corrected intensities match expected values - np.testing.assert_allclose( - result["ena_intensity"][0, [4, 5], :, :].values, - expected_corrected_intensity.values, - rtol=1e-10, - err_msg="Sputtering-corrected intensities don't match expected calculation", - ) - - def test_calculate_sputtering_corrections_negative_j_o_prime(self): - """Test handling when j_o_prime becomes negative.""" - # Create dataset where background > intensity (would give negative j_o_prime) - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": np.linspace(0, 360, 5, endpoint=False), - "latitude": np.linspace(-90, 90, 3), - } - - # Create hydrogen dataset - h_dataset = xr.Dataset(coords=coords) - h_dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 1e6, - ) - - # Add required uncertainty variables for hydrogen - h_dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 0.1e6, - ) - h_dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 0.05e6, - ) - - # Create oxygen dataset where background > intensity - o_dataset = xr.Dataset(coords=coords) - o_dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 1e6, - ) - o_dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 2e6, # Higher than signal - ) - - # Add required uncertainty variables for oxygen - o_dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 0.1e6, - ) - o_dataset["bg_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 0.1e6, - ) - o_dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 5, 3)) * 0.05e6, - ) - o_dataset["geometric_factor"] = (("energy",), np.ones(7)) - - result = calculate_sputtering_corrections(h_dataset, o_dataset) - - # Function should handle negative j_o_prime by setting it to zero - # This means no sputtering correction should be applied - # The corrected intensity should equal the original intensity - np.testing.assert_allclose( - result["ena_intensity"][0, [4, 5], :, :].values, - h_dataset["ena_intensity"][0, [4, 5], :, :].values, - rtol=1e-10, - err_msg=( - "When background > signal, no sputtering correction should be applied" - ), - ) - - def test_calculate_sputtering_corrections_only_oxygen( - self, sample_dataset_with_sputtering_data - ): - """Test that sputtering corrections work for different species datasets.""" - h_dataset, o_dataset = sample_dataset_with_sputtering_data - - # Store original hydrogen values - original_h_intensity = h_dataset["ena_intensity"].copy() - original_h_stat_uncert = h_dataset["ena_intensity_stat_uncert"].copy() - original_h_sys_err = h_dataset["ena_intensity_sys_err"].copy() - - # Store original oxygen values - original_o_intensity = o_dataset["ena_intensity"].copy() - - # Test hydrogen dataset with oxygen reference - result_h = calculate_sputtering_corrections(h_dataset, o_dataset) - - # Check that hydrogen values changed for energy levels 4 and 5 - assert not np.array_equal( - result_h["ena_intensity"][0, 4, :, :].values, - original_h_intensity[0, 4, :, :].values, - ), "Hydrogen intensity should be affected by sputtering corrections" - - assert not np.array_equal( - result_h["ena_intensity_stat_uncert"][0, 4, :, :].values, - original_h_stat_uncert[0, 4, :, :].values, - ), "Hydrogen stat uncertainty should be updated" - - assert not np.array_equal( - result_h["ena_intensity_sys_err"][0, 4, :, :].values, - original_h_sys_err[0, 4, :, :].values, - ), "Hydrogen sys error should be updated" - - # Test oxygen dataset with itself as reference - result_o = calculate_sputtering_corrections(o_dataset, o_dataset) - - # Check that oxygen values also changed - assert not np.array_equal( - result_o["ena_intensity"][0, 4, :, :].values, - original_o_intensity[0, 4, :, :].values, - ), "Oxygen intensity should also be affected by sputtering corrections" - - def test_calculate_sputtering_corrections_uncertainty_propagation( - self, sample_dataset_with_sputtering_data - ): - """Test that uncertainties are properly propagated in sputtering corrections.""" - h_dataset, o_dataset = sample_dataset_with_sputtering_data - - # Get subset for manual calculation - o_small_dataset = o_dataset.isel(epoch=0, energy=[4, 5]) - h_small_dataset = h_dataset.isel(epoch=0, energy=[4, 5]) - - # Manual calculation following equations 10, 12 - j_o_prime = o_small_dataset["ena_intensity"] - o_small_dataset["bg_intensity"] - j_o_prime = j_o_prime.where(j_o_prime >= 0, 0) - - j_o_prime_var = ( - o_small_dataset["ena_intensity_stat_uncert"] ** 2 - + (o_small_dataset["bg_intensity_stat_uncert"]) ** 2 - ) - - sputter_correction_factor = np.array([0.15, 0.01])[:, np.newaxis, np.newaxis] - - expected_corrected_var = ( - h_small_dataset["ena_intensity_stat_uncert"] ** 2 - + (sputter_correction_factor**2) * j_o_prime_var - ) - - result = calculate_sputtering_corrections(h_dataset, o_dataset) - - # Check that statistical uncertainties match expected propagation - np.testing.assert_allclose( - result["ena_intensity_stat_uncert"][0, [4, 5], :, :].values ** 2, - expected_corrected_var.values, - rtol=1e-10, - err_msg="Statistical uncertainty propagation is incorrect", - ) - - def test_calculate_sputtering_corrections_energy_levels(self): - """Test that sputtering corrections are applied to correct energy levels.""" - # Create minimal dataset for testing specific energy level targeting - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": [0, 90, 180, 270], - "latitude": [-45, 0, 45], - } - - # Create hydrogen dataset - h_dataset = xr.Dataset(coords=coords) - h_intensity_data = np.zeros((1, 7, 4, 3)) - h_bg_rate_data = np.zeros((1, 7, 4, 3)) - - # Set specific values for energy levels 4 and 5 only - h_intensity_data[0, 4, :, :] = 200_000_000 # 200M for energy index 4 - h_intensity_data[0, 5, :, :] = 250_000_000 # 250M for energy index 5 - h_bg_rate_data[0, 4, :, :] = 20_000_000 # 20M background - h_bg_rate_data[0, 5, :, :] = 25_000_000 # 25M background - - h_dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - h_intensity_data, - ) - h_dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - h_bg_rate_data, - ) - h_dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 4, 3)) * 100_000, - ) - h_dataset["bg_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 4, 3)) * 10_000, - ) - h_dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 4, 3)) * 50_000, - ) - - # Create oxygen dataset with higher values - o_dataset = xr.Dataset(coords=coords) - o_intensity_data = np.zeros((1, 7, 4, 3)) - o_bg_rate_data = np.zeros((1, 7, 4, 3)) - - # Set specific values for energy levels 4 and 5 only - o_intensity_data[0, 4, :, :] = 250_000_000 # 250M for energy index 4 - o_intensity_data[0, 5, :, :] = 300_000_000 # 300M for energy index 5 - o_bg_rate_data[0, 4, :, :] = 25_000_000 # 25M background - o_bg_rate_data[0, 5, :, :] = 30_000_000 # 30M background - - o_dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - o_intensity_data, - ) - o_dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - o_bg_rate_data, - ) - o_dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 4, 3)) * 100_000, - ) - o_dataset["bg_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 4, 3)) * 10_000, - ) - o_dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 4, 3)) * 50_000, - ) - o_dataset["geometric_factor"] = (("energy",), np.ones(7)) - - # Make a copy to preserve original values for comparison - original_h_dataset = h_dataset.copy(deep=True) - result = calculate_sputtering_corrections(h_dataset, o_dataset) - - # Only energy indices 4 and 5 should be modified - modified_indices = [4, 5] - unchanged_indices = [0, 1, 2, 3, 6] - - for idx in unchanged_indices: - np.testing.assert_array_equal( - result["ena_intensity"][0, idx, :, :].values, - original_h_dataset["ena_intensity"][0, idx, :, :].values, - err_msg=f"Energy index {idx} should not be modified", - ) - - for idx in modified_indices: - # Check that values changed with some tolerance for numerical precision - original_values = original_h_dataset["ena_intensity"][0, idx, :, :].values - corrected_values = result["ena_intensity"][0, idx, :, :].values - - assert not np.allclose(corrected_values, original_values, rtol=1e-10), ( - f"Energy index {idx} should be modified" - ) - - def test_calculate_sputtering_corrections_no_csv_files( - self, sample_dataset_with_sputtering_data, caplog - ): - """Test that missing sputter CSV files raise a ValueError.""" - h_dataset, o_dataset = sample_dataset_with_sputtering_data - - with patch("imap_processing.lo.l2.lo_l2.Path.glob", return_value=iter([])): - with pytest.raises(ValueError, match="No sputter correction files found"): - calculate_sputtering_corrections(h_dataset, o_dataset) - - -class TestInitializeGeometricFactorVariables: - """Tests for the initialize_geometric_factor_variables function.""" - - def test_initialize_geometric_factor_variables(self): - """Test initialization of geometric factor variables.""" - # Create a simple dataset - dataset = xr.Dataset( - { - "test_var": (("energy",), np.ones(7)), - }, - coords={"energy": range(7)}, - ) - - result = initialize_geometric_factor_variables(dataset) - - # Check that all geometric factor variables were initialized - expected_vars = [ - "energy_delta_minus", - "energy_delta_plus", - "geometric_factor", - "geometric_factor_stat_uncert_minus", - "geometric_factor_stat_uncert_plus", - ] - - for var in expected_vars: - assert var in result.data_vars - assert result[var].dims == ("energy",) - assert result[var].shape == (7,) - assert np.all(result[var].values == 0) # Should be initialized to zeros - - # The energy coordinate should also be updated - assert "energy" in result.coords - assert result.coords["energy"].shape == (7,) - assert np.all(result.coords["energy"].values == 0) # Should be zeros - - -class TestPopulateGeometricFactors: - """Tests for the populate_geometric_factors function.""" - - @pytest.mark.parametrize("species", ["h", "o"]) - @patch("imap_processing.lo.l2.lo_l2.reduce_geometric_factor_dataset") - def test_populate_geometric_factors( - self, mock_get_geometric_factor_dataset, species, sample_geometric_factor_data - ): - """Test population of geometric factor values for a specific species.""" - h_gf_data, o_gf_data = sample_geometric_factor_data - gf_data = h_gf_data if species == "h" else o_gf_data - mock_get_geometric_factor_dataset.return_value = gf_data.to_xarray() - - # Create initialized dataset - dataset = xr.Dataset(coords={"energy": range(7)}) - dataset = initialize_geometric_factor_variables(dataset) - - result = populate_geometric_factors(dataset, species) - - # Check that values were populated correctly - for i in range(7): - if species == "h": - # Check hydrogen values - assert result["energy"].values[i] == 0.01 * (i + 1) - assert result["geometric_factor"].values[i] == 1e-4 * (i + 1) - assert result["geometric_factor_stat_uncert_minus"].values[i] == ( - 1e-5 * (i + 1) - ) - assert result["geometric_factor_stat_uncert_plus"].values[i] == ( - 2e-5 * (i + 1) - ) - else: # oxygen - assert result["energy"].values[i] == 0.015 * (i + 1) - assert result["geometric_factor"].values[i] == 1.5e-4 * (i + 1) - assert result["geometric_factor_stat_uncert_minus"].values[i] == ( - 1.5e-5 * (i + 1) - ) - assert result["geometric_factor_stat_uncert_plus"].values[i] == ( - 3e-5 * (i + 1) - ) - # Ensure that energy_deltas are in units of keV - assert np.all(result["energy_delta_plus"].values < 1) - assert np.all(result["energy_delta_minus"].values < 1) - - def test_populate_geometric_factors_no_gf_species(self): - """Test population for species without geometric factors.""" - # Create initialized dataset - dataset = xr.Dataset(coords={"energy": range(7)}) - dataset = initialize_geometric_factor_variables(dataset) - - # Test with doubles (no geometric factors) - result = populate_geometric_factors(dataset, "doubles") - - # Should return dataset unchanged (all zeros) - assert np.all(result["geometric_factor"].values == 0) - - -class TestCleanupIntermediateVariables: - """Tests for the cleanup_intermediate_variables function.""" - - def test_cleanup_intermediate_variables(self): - """Test removal of intermediate variables.""" - # Create dataset with intermediate variables using current naming - dataset = xr.Dataset( - { - "counts": (("energy",), np.ones(7)), - "counts_over_eff": (("energy",), np.ones(7)), - "counts_over_eff_squared": (("energy",), np.ones(7)), - "bg_rate_exposure_factor": (("energy",), np.ones(7)), - "bg_rate_stat_uncert_exposure_factor2": (("energy",), np.ones(7)), - "bg_rate_sys_err_exposure_factor": (("energy",), np.ones(7)), - "ena_intensity": (("energy",), np.ones(7)), # Should be kept - "exposure_factor": (("energy",), np.ones(7)), # Should be kept - } - ) - - result = cleanup_intermediate_variables(dataset) - - # Should keep these variables - assert "counts" in result.data_vars - assert "ena_intensity" in result.data_vars - assert "exposure_factor" in result.data_vars - - # Should remove these intermediate variables - assert "counts_over_eff" not in result.data_vars - assert "counts_over_eff_squared" not in result.data_vars - assert "bg_rate_exposure_factor" not in result.data_vars - assert "bg_rate_stat_uncert_exposure_factor2" not in result.data_vars - assert "bg_rate_sys_err_exposure_factor" not in result.data_vars - - def test_cleanup_partial_variables(self): - """Test cleanup when only some intermediate variables exist.""" - # Create dataset with only some intermediate variables - dataset = xr.Dataset( - { - "counts": (("energy",), np.ones(7)), - "counts_over_eff": (("energy",), np.ones(7)), - "exposure_factor": (("energy",), np.ones(7)), - } - ) - - result = cleanup_intermediate_variables(dataset) - - # Should keep these - assert "counts" in result.data_vars - assert "exposure_factor" in result.data_vars - - # Should remove only the existing intermediate variable - assert "counts_over_eff" not in result.data_vars - - -class TestCalculateBootstrapCorrections: - """Tests for the calculate_bootstrap_corrections function.""" - - def test_calculate_bootstrap_corrections_basic( - self, sample_dataset_with_bootstrap_data - ): - """Test basic bootstrap correction functionality.""" - dataset = sample_dataset_with_bootstrap_data.copy() - - # Store original values for comparison - original_intensity = dataset["ena_intensity"].copy() - - result = calculate_bootstrap_corrections(dataset) - - # Check that bootstrap corrections were applied - assert "ena_intensity" in result.data_vars - assert "ena_intensity_stat_uncert" in result.data_vars - assert "ena_intensity_sys_err" in result.data_vars - - # Check that intermediate bootstrap variables were removed - assert "bootstrap_intensity" not in result.data_vars - assert "bootstrap_intensity_stat_uncert" not in result.data_vars - assert "bootstrap_intensity_sys_err" not in result.data_vars - - # Check that corrected intensities are different from originals - # (bootstrap should reduce intensities due to spillover correction) - corrected_intensity = result["ena_intensity"] - assert not np.allclose( - corrected_intensity.values, original_intensity.values, rtol=1e-10 - ), "Bootstrap corrections should modify intensities" - - # Check that corrected intensities are generally lower - # (bootstrap removes spillover from higher to lower energies) - for energy_idx in range(5): # Lower energy channels should be reduced - assert np.all( - corrected_intensity[0, energy_idx, :, :].values - <= original_intensity[0, energy_idx, :, :].values - ), f"Bootstrap should reduce intensity at energy index {energy_idx}" - - # This is a spot check value to ensure we are getting actual - # corrected intensities. Previously we were missing the application - # of the lower energy channels in the summation. - np.testing.assert_allclose(corrected_intensity[0, 5, 0, 0], [895.96302438]) - - def test_calculate_bootstrap_corrections_equations( - self, sample_dataset_with_bootstrap_data - ): - """Test that bootstrap equations are correctly implemented.""" - dataset = sample_dataset_with_bootstrap_data.copy() - - # Calculate expected j_c_prime (equation 14) - j_c_prime_expected = dataset["ena_intensity"] - dataset["bg_intensity"] - j_c_prime_expected = j_c_prime_expected.where(j_c_prime_expected >= 0, 0) - - # Apply bootstrap corrections and check the calculation was done correctly - result = calculate_bootstrap_corrections(dataset) - - # Verify the final result makes sense given the bootstrap factors - # Higher energy channels should have less correction - assert np.all(result["ena_intensity"][0, 6, :, :] >= 0), ( - "Bootstrap intensities should be non-negative" - ) - - def test_calculate_bootstrap_corrections_negative_handling(self): - """Test proper handling of negative values during bootstrap calculation.""" - # Create dataset with some negative j_c_prime values - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": [0, 90], - "latitude": [0, 45], - } - - dataset = xr.Dataset(coords=coords) - - # Create energy values - energy_values = np.array([0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]) - dataset["energy"] = (("energy",), energy_values) - dataset["geometric_factor"] = (("energy",), np.ones(7)) - - # Create intensities where some background > intensity (negative j_c_prime) - intensity_values = np.ones((1, 7, 2, 2)) * 1e6 - bg_intensity_values = np.ones((1, 7, 2, 2)) * 1.5e6 # Higher than intensity - - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - bg_intensity_values, - ) - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 2, 2)) * 1e5, - ) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 2, 2)) * 5e4, - ) - - result = calculate_bootstrap_corrections(dataset) - - # All corrected intensities should be non-negative - assert np.all(result["ena_intensity"].values >= 0), ( - "Bootstrap corrections should not produce negative intensities" - ) - - def test_calculate_bootstrap_corrections_energy_dependence(self): - """Test that bootstrap corrections show proper energy dependence.""" - # Create dataset with realistic energy spectrum - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": [0], - "latitude": [0], - } - - dataset = xr.Dataset(coords=coords) - - # Create energy values - energy_values = np.array([0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]) - dataset["energy"] = (("energy",), energy_values) - dataset["geometric_factor"] = (("energy",), np.ones(7)) - - # Create a steep power law spectrum (typical for ENAs) - base_intensity = 1e8 - intensity_values = np.ones((1, 7, 1, 1)) - for i in range(7): - intensity_values[0, i, 0, 0] = base_intensity * (energy_values[i]) ** (-3) - - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - - # Low background - dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.01, # 1% background - ) - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.sqrt(intensity_values) * 0.1, - ) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.05, - ) - - result = calculate_bootstrap_corrections(dataset) - - # Lower energy channels should show larger corrections - # because they receive spillover from higher energy channels - original_ratios = [] - corrected_ratios = [] - - for i in range(6): # Compare adjacent energy channels - original_ratio = ( - dataset["ena_intensity"][0, i, 0, 0].values - / dataset["ena_intensity"][0, i + 1, 0, 0].values - ) - corrected_ratio = ( - result["ena_intensity"][0, i, 0, 0].values - / result["ena_intensity"][0, i + 1, 0, 0].values - ) - original_ratios.append(original_ratio) - corrected_ratios.append(corrected_ratio) - - # Bootstrap should affect the intensities - # Check that the bootstrap corrections are actually being applied - - # For power law spectra with small backgrounds, corrections may be very small - # Let's check that the algorithm at least completes without error - # and produces reasonable output - - # Check that all intensities are finite and positive - assert np.all(np.isfinite(result["ena_intensity"].values)), ( - "All corrected intensities should be finite" - ) - assert np.all(result["ena_intensity"].values >= 0), ( - "All corrected intensities should be non-negative" - ) - - # Check that uncertainties are reasonable - assert np.all(np.isfinite(result["ena_intensity_stat_uncert"].values)), ( - "All statistical uncertainties should be finite" - ) - assert np.all(np.isfinite(result["ena_intensity_sys_err"].values)), ( - "All systematic errors should be finite" - ) - - def test_calculate_bootstrap_corrections_uncertainty_propagation(self): - """Test proper uncertainty propagation in bootstrap corrections.""" - # Create simple dataset for uncertainty testing - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": [0], - "latitude": [0], - } - - dataset = xr.Dataset(coords=coords) - - # Create energy values - energy_values = np.array([0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]) - dataset["energy"] = (("energy",), energy_values) - dataset["geometric_factor"] = (("energy",), np.ones(7)) - - # Simple flat spectrum for easier uncertainty analysis - intensity_values = np.ones((1, 7, 1, 1)) * 1e6 - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.1, # 10% background - ) - - # Known uncertainties - stat_uncert = np.ones((1, 7, 1, 1)) * 1e5 - sys_err = np.ones((1, 7, 1, 1)) * 5e4 - - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - stat_uncert, - ) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - sys_err, - ) - - result = calculate_bootstrap_corrections(dataset) - - # Check that uncertainties are properly propagated - assert np.all(result["ena_intensity_stat_uncert"].values >= 0), ( - "Statistical uncertainties should be non-negative" - ) - assert np.all(result["ena_intensity_sys_err"].values >= 0), ( - "Systematic errors should be non-negative" - ) - - # Uncertainties should be reasonable relative to the intensities - relative_stat_uncert = ( - result["ena_intensity_stat_uncert"] / result["ena_intensity"] - ) - assert np.all(relative_stat_uncert.values < 1.0), ( - "Relative statistical uncertainty should be reasonable" - ) - - def test_calculate_bootstrap_corrections_virtual_channel(self): - """Test the virtual channel E8 calculation and its impact.""" - # Create dataset focused on energy channels 5 and 6 for E8 calculation - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": [0], - "latitude": [0], - } - - dataset = xr.Dataset(coords=coords) - - # Create energy values where E6/E7 ratio is well-defined - energy_values = np.array([0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]) - dataset["energy"] = (("energy",), energy_values) - dataset["geometric_factor"] = (("energy",), np.ones(7)) - - # Create intensities with clear power law for gamma calculation - intensity_values = np.ones((1, 7, 1, 1)) - for i in range(7): - intensity_values[0, i, 0, 0] = 1e6 * (energy_values[i] / 1.0) ** (-2) - - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.05, # 5% background - ) - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.sqrt(intensity_values) * 0.1, - ) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.03, - ) - - # Calculate expected E8 and gamma manually - j_c_prime = dataset["ena_intensity"] - dataset["bg_intensity"] - j_c_prime = j_c_prime.where(j_c_prime >= 0, 0) - - result = calculate_bootstrap_corrections(dataset) - - # E8 should follow the power law relationship - # The virtual channel should have meaningful impact on energy channel 6 - original_e6 = j_c_prime[0, 6, 0, 0].values - corrected_e6 = result["ena_intensity"][0, 6, 0, 0].values - - # Energy channel 6 should be reduced due to E8 spillover subtraction - assert corrected_e6 < original_e6, ( - "Energy channel 6 should be reduced by E8 virtual channel correction" - ) - - def test_calculate_bootstrap_corrections_bootstrap_factors(self): - """Test that the bootstrap factor matrix is applied correctly.""" - # Create a simple dataset to verify bootstrap factor application - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": [0], - "latitude": [0], - } - - dataset = xr.Dataset(coords=coords) - - energy_values = np.array([0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]) - dataset["energy"] = (("energy",), energy_values) - dataset["geometric_factor"] = (("energy",), np.ones(7)) - - # Use unit intensities to make bootstrap factor effects clear - intensity_values = np.ones((1, 7, 1, 1)) * 1.0 - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.zeros((1, 7, 1, 1)), # No background - ) - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 1, 1)) * 0.1, - ) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.ones((1, 7, 1, 1)) * 0.05, - ) - - result = calculate_bootstrap_corrections(dataset) - - # With unit intensities and no background, the corrections should - # directly reflect the bootstrap factors - # The bootstrap algorithm is complex due to interdependencies - # Let's just verify that corrections are applied and reasonable - corrected_intensities = result["ena_intensity"][0, :, 0, 0].values - - # All channels should be reduced from their original value of 1.0 - for i in range(7): - assert corrected_intensities[i] < 1.0, ( - f"Energy {i} should be corrected (reduced from 1.0), " - f"got {corrected_intensities[i]}" - ) - - # Energy 6 should have the largest correction due to 0.75 factor - assert corrected_intensities[6] < 0.5, ( - f"Energy 6 should have large correction, got {corrected_intensities[6]}" - ) - - def test_calculate_bootstrap_corrections_edge_cases(self): - """Test edge cases in bootstrap correction calculation.""" - # Test with zero intensities - coords = { - "epoch": [8.1794907049e17], - "energy": list(range(7)), - "longitude": [0], - "latitude": [0], - } - - dataset = xr.Dataset(coords=coords) - - energy_values = np.array([0.5, 1.0, 2.0, 4.0, 8.0, 16.0, 32.0]) - dataset["energy"] = (("energy",), energy_values) - dataset["geometric_factor"] = (("energy",), np.ones(7)) - - # Zero intensities - intensity_values = np.zeros((1, 7, 1, 1)) - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) - dataset["bg_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.zeros((1, 7, 1, 1)), - ) - dataset["ena_intensity_stat_uncert"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.zeros((1, 7, 1, 1)), - ) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - np.zeros((1, 7, 1, 1)), - ) - - result = calculate_bootstrap_corrections(dataset) - - # Should handle zero intensities gracefully - assert np.all(result["ena_intensity"].values >= 0), ( - "Zero intensities should remain non-negative" - ) - assert np.all(np.isfinite(result["ena_intensity"].values)), ( - "All intensities should be finite" - ) - - def test_calculate_bootstrap_corrections_no_csv_files( - self, sample_dataset_with_bootstrap_data, caplog - ): - """Test that missing bootstrap CSV files raise a ValueError.""" - dataset = sample_dataset_with_bootstrap_data.copy() - - with patch("imap_processing.lo.l2.lo_l2.Path.glob", return_value=iter([])): - with pytest.raises( - ValueError, match="No bootstrap correction factor files found" - ): - calculate_bootstrap_corrections(dataset) - - -# ============================================================================= -# INTEGRATION TESTS -# ============================================================================= - - -class TestCalculateAllRatesAndIntensities: - """Integration tests for the calculate_all_rates_and_intensities function.""" - - def test_calculate_all_rates_and_intensities_complete(self): - """Test the complete rates and intensities calculation pipeline.""" - # Create a comprehensive dataset with current naming convention - dataset = xr.Dataset( - { - # Count data (current generic naming) - "counts": (("energy",), np.ones(7) * 10), - # Efficiency corrected data - "counts_over_eff": (("energy",), np.ones(7) * 12), # 10/0.83 ≈ 12 - "counts_over_eff_squared": (("energy",), np.ones(7) * 12), - # Other required data - "exposure_factor": (("energy",), np.ones(7) * 1.0), - "geometric_factor": (("energy",), np.ones(7) * 1e-4), - "energy": (("energy",), np.ones(7) * 0.1), - "geometric_factor_stat_uncert_minus": (("energy",), np.ones(7) * 1e-5), - "geometric_factor_stat_uncert_plus": (("energy",), np.ones(7) * 3e-5), - # Background intermediate data - "bg_rate_exposure_factor": (("energy",), np.ones(7) * 0.3), - "bg_rate_stat_uncert_exposure_factor2": ( - ("energy",), - np.ones(7) * 0.009, - ), - "bg_rate_sys_err_exposure_factor": ( - ("energy",), - np.ones(7) * 0.06, - ), - } - ) - - result = calculate_all_rates_and_intensities(dataset) - - # Check that rates were calculated - assert "ena_count_rate" in result.data_vars - assert "ena_count_rate_stat_uncert" in result.data_vars - - # Check that intensities were calculated - assert "ena_intensity" in result.data_vars - assert "ena_intensity_stat_uncert" in result.data_vars - assert "ena_intensity_sys_err" in result.data_vars - - # Check that background rates were calculated - assert "bg_rate" in result.data_vars - assert "bg_rate_stat_uncert" in result.data_vars - assert "bg_rate_sys_err" in result.data_vars - - # Check that intermediate variables were cleaned up - assert "counts_over_eff" not in result.data_vars - assert "counts_over_eff_squared" not in result.data_vars - assert "bg_rate_exposure_factor" not in result.data_vars - assert "bg_rate_stat_uncert_exposure_factor2" not in result.data_vars - - def test_calculate_all_rates_with_cg_correction( - self, sample_dataset_with_intensities - ): - """Test that CG correction is applied when cg_correction=True.""" - # Add necessary variables for the calculation - dataset = sample_dataset_with_intensities.copy(deep=True) - dataset["counts"] = (("epoch", "energy"), np.ones((1, 7)) * 10) - dataset["counts_over_eff"] = (("epoch", "energy"), np.ones((1, 7)) * 12) - dataset["counts_over_eff_squared"] = (("epoch", "energy"), np.ones((1, 7)) * 12) - dataset["exposure_factor"] = (("epoch", "energy"), np.ones((1, 7)) * 1.0) - dataset["geometric_factor"] = (("energy",), np.ones(7) * 1e-4) - dataset["geometric_factor_stat_uncert_minus"] = (("energy",), np.ones(7) * 1e-5) - dataset["geometric_factor_stat_uncert_plus"] = (("energy",), np.ones(7) * 3e-5) - dataset["bg_rate_exposure_factor"] = ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.3, - ) - dataset["bg_rate_stat_uncert_exposure_factor2"] = ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.009, - ) - dataset["bg_rate_sys_err_exposure_factor"] = ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.06, - ) - dataset["energy_sc_exposure_factor"] = xr.ones_like(dataset["ena_intensity"]) - - # Mock the interpolation function - with patch( - "imap_processing.lo.l2.lo_l2.interpolate_map_flux_to_helio_frame" - ) as mock_interp: - # Make the mock return the input dataset - mock_interp.side_effect = lambda ds, *args, **kwargs: ds - - # Call with CG correction enabled - _ = calculate_all_rates_and_intensities(dataset, cg_correction=True) - - # Verify that interpolation was called - mock_interp.assert_called_once() - - # Check the call arguments - call_args = mock_interp.call_args - assert call_args[0][1].equals( - dataset["energy"] - ) # spacecraft frame energies - assert call_args[0][2].equals(dataset["energy"]) # helio frame energies - # Check that s/c energies get computed in keV units - expected_sc_energies = ( - dataset["energy_sc_exposure_factor"] / dataset["exposure_factor"] / 1e3 - ) - xr.testing.assert_allclose( - call_args[0][0]["energy_sc"], expected_sc_energies - ) - assert "ena_intensity" in call_args[0][3] # variables to interpolate - assert "bg_intensity" in call_args[0][3] - - def test_calculate_all_rates_cg_with_other_corrections( - self, sample_dataset_with_intensities, lo_flux_factors_file - ): - """Test CG correction works alongside other corrections.""" - # Add necessary variables - dataset = sample_dataset_with_intensities.copy(deep=True) - dataset["counts"] = (("epoch", "energy"), np.ones((1, 7)) * 10) - dataset["counts_over_eff"] = (("epoch", "energy"), np.ones((1, 7)) * 12) - dataset["counts_over_eff_squared"] = (("epoch", "energy"), np.ones((1, 7)) * 12) - dataset["exposure_factor"] = (("epoch", "energy"), np.ones((1, 7)) * 1.0) - dataset["geometric_factor"] = (("energy",), np.ones(7) * 1e-4) - dataset["geometric_factor_stat_uncert_minus"] = (("energy",), np.ones(7) * 1e-5) - dataset["geometric_factor_stat_uncert_plus"] = (("energy",), np.ones(7) * 3e-5) - dataset["bg_rate_exposure_factor"] = ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.3, - ) - dataset["bg_rate_stat_uncert_exposure_factor2"] = ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.009, - ) - dataset["bg_rate_sys_err_exposure_factor"] = ( - ("epoch", "energy"), - np.ones((1, 7)) * 0.06, - ) - dataset["energy_sc_exposure_factor"] = xr.ones_like(dataset["ena_intensity"]) - - with patch( - "imap_processing.lo.l2.lo_l2.interpolate_map_flux_to_helio_frame" - ) as mock_interp: - mock_interp.side_effect = lambda ds, *args, **kwargs: ds - - # Call with both flux correction and CG correction - result = calculate_all_rates_and_intensities( - dataset, - flux_correction=True, - flux_factors=lo_flux_factors_file, - cg_correction=True, - ) - - # Both corrections should be applied - assert "ena_intensity" in result.data_vars - mock_interp.assert_called_once() - - -@pytest.mark.external_kernel -class TestIntegrationWithMocks: - """Integration tests using mocked external dependencies.""" - - def test_lo_l2_integration_minimal( - self, minimal_pset_for_species, lo_flux_factors_file - ): - """Test the main lo_l2 function with minimal mocking.""" - # Test with hydrogen data - sci_dependencies = {"imap_lo_l1c_pset": [minimal_pset_for_species]} - anc_dependencies = [lo_flux_factors_file] # Include flux factors file - descriptor = "l090-ena-h-sf-nsp-ram-hae-6deg-3mo" - - # Mock the complex external dependencies to return simple results - with ( - patch( - "imap_processing.lo.l2.lo_l2.create_sky_map_from_psets" - ) as mock_create_map, - patch("imap_processing.lo.l2.lo_l2.add_geometric_factors") as mock_add_gf, - patch( - "imap_processing.lo.l2.lo_l2.calculate_all_rates_and_intensities" - ) as mock_calc_rates, - patch("imap_processing.lo.l2.lo_l2.finalize_dataset") as mock_finalize, - ): - # Setup mock returns - mock_sky_map = Mock() - mock_dataset = xr.Dataset({"test_var": (("energy",), np.ones(7))}) - mock_sky_map.to_dataset.return_value = mock_dataset - mock_sky_map.build_cdf_dataset.return_value = mock_dataset - mock_create_map.return_value = mock_sky_map - mock_add_gf.return_value = mock_dataset - mock_calc_rates.return_value = mock_dataset - mock_finalize.return_value = mock_dataset - - # Run the function - should not crash - result = lo_l2(sci_dependencies, anc_dependencies, descriptor) - - # Basic validation - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], xr.Dataset) - - # Mock the rates calculation to return the dataset unchanged - mock_calc_rates.side_effect = lambda x, **kwargs: x - - # Run the function - should not crash - result = lo_l2(sci_dependencies, anc_dependencies, descriptor) - - # Basic validation - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], xr.Dataset) - - -@pytest.fixture -def ibex_pset_file(): - """Path to the LO/IBEX pset test file.""" - # Use the actual test data file from the ena_maps test data - test_data_path = Path(__file__).parent / "test_cdfs" - return test_data_path / "imap_lo_l1c_pset_20260101-repoint01261_v001.cdf" - - -@pytest.mark.external_test_data -@pytest.mark.external_kernel -class TestIntegration: - """Integration tests using IBEX data and simulated kernels.""" - - def test_lo_l2_integration_full( - self, ibex_pset_file, imap_ena_sim_metakernel, lo_flux_factors_file - ): - """Test the main lo_l2 function with no mocking.""" - # Test with oxygen data to reduce test run-time - psets = load_cdf(ibex_pset_file) - psets.attrs["Logical_source"] = "imap_lo_l1c_pset" - psets = add_spacecraft_position_and_velocity_to_pset(psets) - - sci_dependencies = {"imap_lo_l1c_pset": [psets]} - anc_dependencies = [lo_flux_factors_file] # Include flux factors file - descriptor = "l090-ena-o-hf-nsp-ram-hae-6deg-3mo" - - # Run the function - should not crash - result = lo_l2(sci_dependencies, anc_dependencies, descriptor) - - # Basic validation - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], xr.Dataset) - - # Make sure that bg_rate variables are present - for var in ["bg_rate", "bg_rate_stat_uncert", "bg_rate_sys_err"]: - assert var in result[0].data_vars - - -# ============================================================================= -# ERROR HANDLING TESTS -# ============================================================================= - - -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.""" - sci_dependencies = {} # Missing imap_lo_l1c_pset - anc_dependencies = [] - descriptor = "l090-ena-h-sf-nsp-ram-hae-6deg-3mo" - - with pytest.raises(KeyError, match="imap_lo_l1c_pset"): - lo_l2(sci_dependencies, anc_dependencies, descriptor) - - def test_create_sky_map_healpix_not_supported(self, minimal_pset_for_species): - """Test error when HEALPix map is requested.""" - with patch.object(MapDescriptor, "from_string") as mock_from_string: - mock_map_desc = Mock() - mock_healpix_map = Mock() # Not a RectangularSkyMap - mock_map_desc.to_empty_map.return_value = mock_healpix_map - mock_map_desc.species = "h" - mock_from_string.return_value = mock_map_desc - - with pytest.raises( - NotImplementedError, match="HEALPix map output not supported" - ): - create_sky_map_from_psets( - [minimal_pset_for_species], mock_map_desc, pd.DataFrame(), False - ) - - -# ============================================================================= -# PROPERTY-BASED AND EDGE CASE TESTS -# ============================================================================= - - -class TestEdgeCases: - """Tests for edge cases and boundary conditions.""" - - def test_empty_efficiency_data_handling(self, minimal_pset): - """Test handling of empty efficiency data.""" - empty_df = pd.DataFrame() - result = add_efficiency_factors_to_pset(minimal_pset, empty_df) - - # Should create unity efficiency - assert "efficiency" in result.data_vars - np.testing.assert_array_equal(result["efficiency"].values, np.ones(7)) - - def test_zero_exposure_time_handling(self): - """Test handling of zero exposure times.""" - dataset = xr.Dataset( - { - "counts": (("energy",), np.ones(7) * 10), - "exposure_factor": (("energy",), np.zeros(7)), # Zero exposure - } - ) - - result = calculate_rates(dataset) - - # Should handle division by zero gracefully - assert "ena_count_rate" in result.data_vars - # Rates should be infinite where exposure time is zero - assert np.all(np.isinf(result["ena_count_rate"].values)) - - def test_negative_counts_handling(self): - """Test handling of negative count values.""" - dataset = xr.Dataset( - { - "counts": (("energy",), np.array([-1, 0, 1, 2, 3, 4, 5])), - "exposure_factor": (("energy",), np.ones(7)), - } - ) - - result = calculate_rates(dataset) - - # Should calculate rates even with negative counts - assert "ena_count_rate" in result.data_vars - assert "ena_count_rate_stat_uncert" in result.data_vars - - # Uncertainty calculation should handle negative counts - # (sqrt of negative gives NaN, which is expected behavior) - assert np.isnan(result["ena_count_rate_stat_uncert"].values[0]) - - -# ============================================================================= -# TESTS FOR CG CORRECTION FUNCTIONALITY -# ============================================================================= - - -class TestPrepareCorrections: - """Tests for the _prepare_corrections function.""" - - def test_prepare_corrections_hf_frame(self, lo_flux_factors_file): - """Test that CG correction is enabled for heliocentric frame (hf).""" - # Create a map descriptor with heliocentric frame - descriptor = "ilo90-ena-h-hf-nsp-full-hae-6deg-3mo" - map_descriptor = MapDescriptor.from_string(descriptor) - - sci_dependencies = {"imap_lo_l1c_pset": []} - anc_dependencies = [lo_flux_factors_file] - - with patch("imap_processing.lo.l2.lo_l2.lo_l2") as mock_lo_l2: - mock_o_dataset = xr.Dataset() - mock_lo_l2.return_value = [mock_o_dataset] - result = _prepare_corrections( - map_descriptor, descriptor, sci_dependencies, anc_dependencies - ) - - # Unpack the tuple - ( - sputtering_correction, - bootstrap_correction, - flux_correction, - o_map_dataset, - flux_factors, - cg_correction, - ) = result - - # Check that CG correction is enabled for hf frame - assert cg_correction is True, "CG correction should be enabled for 'hf' frame" - - # Check other correction flags - assert flux_correction is True # ENA data should have flux correction - assert sputtering_correction is True # h-ena product, apply sputtering - assert bootstrap_correction is True # h-ena product, apply bootstrap - assert o_map_dataset is not None # oxygen dataset produced - assert flux_factors is not None # Flux factors should be found - - def test_prepare_corrections_sc_frame(self, lo_flux_factors_file): - """Test that CG correction is disabled for spacecraft frame (sc).""" - # Create a map descriptor with spacecraft frame - descriptor = "ilo90-ena-o-sf-nsp-full-hae-6deg-3mo" - map_descriptor = MapDescriptor.from_string(descriptor) - - sci_dependencies = {"imap_lo_l1c_pset": []} - anc_dependencies = [lo_flux_factors_file] - - result = _prepare_corrections( - map_descriptor, descriptor, sci_dependencies, anc_dependencies - ) - - # Unpack the tuple - ( - sputtering_correction, - bootstrap_correction, - flux_correction, - o_map_dataset, - flux_factors, - cg_correction, - ) = result - - # Check that CG correction is disabled for sc frame - assert cg_correction is False, ( - "CG correction should be disabled for non-'hf' frame" - ) - - def test_prepare_corrections_with_hydrogen_ena(self, lo_flux_factors_file): - """Test corrections for hydrogen ENA data.""" - descriptor = "ilo90-ena-h-sf-nsp-full-hae-6deg-3mo" - map_descriptor = MapDescriptor.from_string(descriptor) - - # Mock the recursive call to lo_l2 for oxygen - with patch("imap_processing.lo.l2.lo_l2.lo_l2") as mock_lo_l2: - mock_o_dataset = xr.Dataset({"test": (("energy",), np.ones(7))}) - mock_lo_l2.return_value = [mock_o_dataset] - - sci_dependencies = {"imap_lo_l1c_pset": []} - anc_dependencies = [lo_flux_factors_file] - - result = _prepare_corrections( - map_descriptor, descriptor, sci_dependencies, anc_dependencies - ) - - ( - sputtering_correction, - bootstrap_correction, - flux_correction, - o_map_dataset, - flux_factors, - cg_correction, - ) = result - - # Check that all corrections are enabled for hydrogen ENA - assert sputtering_correction is True - assert bootstrap_correction is True - assert flux_correction is True - assert o_map_dataset is not None - assert cg_correction is False # sf frame, not hf - - -class TestProcessSinglePset: - """Tests for the process_single_pset function with CG correction.""" - - def test_process_single_pset_hf_frame(self, minimal_pset, sample_efficiency_data): - """Test that CG correction is applied for heliocentric frame.""" - pset = minimal_pset.copy() - pset = pset.rename({"esa_energy_step": "energy"}) - # apply_compton_getting_correction gets mocked out so we need to add the - # energy_sc variable to the pset - pset["energy_sc"] = xr.ones_like(pset["counts"]) - - with ( - patch( - "imap_processing.lo.l2.lo_l2.normalize_pset_coordinates" - ) as mock_norm, - patch( - "imap_processing.lo.l2.lo_l2.add_efficiency_factors_to_pset" - ) as mock_add_ef, - patch( - "imap_processing.lo.l2.lo_l2.calculate_efficiency_corrected_quantities" - ) as mock_calc_ef, - patch( - "imap_processing.lo.l2.lo_l2.apply_compton_getting_correction" - ) as mock_cg, - patch("imap_processing.lo.l2.lo_l2.calculate_ram_mask") as mock_ram_mask, - ): - mock_norm.return_value = pset - mock_add_ef.return_value = pset - mock_calc_ef.return_value = pset - mock_cg.return_value = pset - mock_ram_mask.return_value = pset - - # Mock the spacecraft velocity - pset["sc_velocity"] = xr.DataArray( - data=[400, 0, 0], # 400 km/s in x direction - dims="component", - coords={"component": ["vx", "vy", "vz"]}, - ) - - # Process with hf frame - _ = process_single_pset(pset, sample_efficiency_data, "h", cg_correct=True) - - # Check that CG correction was called - mock_cg.assert_called_once() - assert mock_cg.call_args[0][0] is pset - # Energy should be passed as second argument - assert "energy" in mock_cg.call_args[0][1].dims - - # Check that calculate_ram_mask was called - mock_ram_mask.assert_called_once() - - def test_process_single_pset_sc_frame(self, minimal_pset, sample_efficiency_data): - """Test that CG correction is not applied for spacecraft frame.""" - pset = minimal_pset.copy() - pset = pset.rename({"esa_energy_step": "energy"}) - - with ( - patch( - "imap_processing.lo.l2.lo_l2.normalize_pset_coordinates" - ) as mock_norm, - patch( - "imap_processing.lo.l2.lo_l2.add_efficiency_factors_to_pset" - ) as mock_add_ef, - patch( - "imap_processing.lo.l2.lo_l2.calculate_efficiency_corrected_quantities" - ) as mock_calc_ef, - patch( - "imap_processing.lo.l2.lo_l2.apply_compton_getting_correction" - ) as mock_cg, - patch("imap_processing.lo.l2.lo_l2.calculate_ram_mask") as mock_ram_mask, - ): - mock_norm.return_value = pset - mock_add_ef.return_value = pset - mock_calc_ef.return_value = pset - mock_cg.return_value = pset - - # Mock the spacecraft velocity - pset["sc_velocity"] = xr.DataArray( - data=[400, 0, 0], # 400 km/s in x direction - dims="component", - coords={"component": ["vx", "vy", "vz"]}, - ) - - # Process with sc frame - _ = process_single_pset(pset, sample_efficiency_data, "h", cg_correct=False) - - # Check that CG correction was NOT called - mock_cg.assert_not_called() - - # Check that the ram mask was called instead - mock_ram_mask.assert_called_once() - - -class TestProjectPsetToMap: - """Tests for the project_pset_to_map function with directional mask.""" - - def test_project_pset_to_map_with_mask(self, minimal_pset_for_species): - """Test that directional mask is passed to projection.""" - # Create a mock sky map - mock_map = Mock(spec=RectangularSkyMap) - - # Create a directional mask - directional_mask = xr.DataArray( - np.ones(7, dtype=bool), - dims=["energy"], - coords={"energy": list(range(7))}, - ) - - # Call project_pset_to_map - project_pset_to_map(minimal_pset_for_species, mock_map, directional_mask) - - # Verify that project_pset_values_to_map was called with the mask - mock_map.project_pset_values_to_map.assert_called_once() - call_kwargs = mock_map.project_pset_values_to_map.call_args[1] - assert "pset_valid_mask" in call_kwargs - assert call_kwargs["pset_valid_mask"] is directional_mask - - def test_project_pset_to_map_value_keys(self, minimal_pset_for_species): - """Test that correct value keys are projected.""" - mock_map = Mock(spec=RectangularSkyMap) - directional_mask = xr.DataArray( - np.ones(7, dtype=bool), - dims=["energy"], - ) - - project_pset_to_map(minimal_pset_for_species, mock_map, directional_mask) - - # Check that the expected value keys are in the call - call_kwargs = mock_map.project_pset_values_to_map.call_args[1] - value_keys = call_kwargs["value_keys"] - - expected_keys = [ - "exposure_factor", - "counts", - "counts_over_eff", - "counts_over_eff_squared", - "bg_rate", - "bg_rate_stat_uncert", - "bg_rate_sys_err", - "bg_rate_exposure_factor", - "bg_rate_stat_uncert_exposure_factor2", - "bg_rate_sys_err_exposure_factor", - ] - - for key in expected_keys: - assert key in value_keys, f"Expected key '{key}' not in value_keys" diff --git a/imap_processing/tests/test_cli.py b/imap_processing/tests/test_cli.py index 3088d68f5..6a5acb70a 100644 --- a/imap_processing/tests/test_cli.py +++ b/imap_processing/tests/test_cli.py @@ -472,23 +472,27 @@ 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": "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"}) + mock_goodtimes = Mock(attrs={"Logical_source": "imap_lo_l1b_goodtimes"}) + mock_bgrates = Mock(attrs={"Logical_source": "imap_lo_l1b_bgrates"}) + mock_histrates_1 = Mock(attrs={"Logical_source": "imap_lo_l1b_histrates"}) + mock_histrates_2 = Mock(attrs={"Logical_source": "imap_lo_l1b_histrates"}) science_file_paths = [ - "imap_lo_l1c_pset_20250415_v001.cdf", - "imap_lo_l1c_pset_20250416_v001.cdf", - "imap_lo_l1b_de_20250415_v001.cdf", + "imap_lo_l1b_goodtimes_20250415-repoint00217_v001.cdf", + "imap_lo_l1b_bgrates_20250415-repoint00217_v001.cdf", + "imap_lo_l1b_histrates_20250415-repoint00217_v001.cdf", + "imap_lo_l1b_histrates_20250416-repoint00218_v001.cdf", ] processing_input = ProcessingInputCollection( *[ScienceInput(file_path) for file_path in science_file_paths], ) + # Loaded in descriptor order: goodtimes, then bgrates, then histrates mocks["mock_load_cdf"].side_effect = [ - mock_loaded_pset_1, - mock_loaded_pset_2, - mock_loaded_de, + mock_goodtimes, + mock_bgrates, + mock_histrates_1, + mock_histrates_2, ] mock_lo_pre_processing.return_value = processing_input @@ -508,8 +512,9 @@ def test_lo_l2(mock_lo_pre_processing, mock_lo_l2, mock_instrument_dependencies) mock_lo_l2.assert_called_once_with( { - "imap_lo_l1c_pset": [mock_loaded_pset_1, mock_loaded_pset_2], - "imap_lo_l1b_de": [mock_loaded_de], + "imap_lo_l1b_goodtimes": [mock_goodtimes], + "imap_lo_l1b_bgrates": [mock_bgrates], + "imap_lo_l1b_histrates": [mock_histrates_1, mock_histrates_2], }, [], descriptor, @@ -520,26 +525,31 @@ def test_lo_l2(mock_lo_pre_processing, mock_lo_l2, mock_instrument_dependencies) @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): - """Test that only inputs at the pivot angle of the map are kept.""" - valid_pset = "imap_lo_l1c_pset_20250415_v001.cdf" - invalid_pset = "imap_lo_l1c_pset_20250416_v001.cdf" - valid_de = "imap_lo_l1b_de_20250415-repoint00217_v001.cdf" - no_pivot_angle = "imap_lo_l1b_goodtimes_20250415-repoint00217_v001.cdf" + """Test that only the pointings at the pivot angle of the map are kept.""" + kept = "-repoint00217_v001.cdf" + dropped = "-repoint00218_v001.cdf" + goodtimes = [ + f"imap_lo_l1b_goodtimes_20250415{kept}", + f"imap_lo_l1b_goodtimes_20250416{dropped}", + ] + histrates = [ + f"imap_lo_l1b_histrates_20250415{kept}", + f"imap_lo_l1b_histrates_20250416{dropped}", + ] + bgrates = [f"imap_lo_l1b_bgrates_20250415{kept}"] ancillary = "imap_lo_efficiency-factors_20250415_v001.csv" base_collection = ProcessingInputCollection( - ScienceInput(valid_pset, invalid_pset), - ScienceInput(valid_de), - ScienceInput(no_pivot_angle), + ScienceInput(*goodtimes), + ScienceInput(*histrates), + ScienceInput(*bgrates), AncillaryInput(ancillary), ) mock_super_pre_processing.return_value = base_collection mock_load_cdf.side_effect = [ - xr.Dataset({"pivot_angle": xr.DataArray(90.1)}), + xr.Dataset({"pivot": ("epoch", [90.1])}), # A neighbouring pivot angle, which belongs on its own map - xr.Dataset({"pivot_angle": xr.DataArray(75.0)}), - xr.Dataset({"pivot_angle": xr.DataArray(90.0)}), - xr.Dataset(), # No pivot angle at all + xr.Dataset({"pivot": ("epoch", [75.0])}), ] instrument = Lo( @@ -553,13 +563,43 @@ def test_lo_pre_processing_pivot_angle_filter(mock_super_pre_processing, mock_lo ) result = instrument.pre_processing() - # The goodtimes input had no pivot angle, so it drops out entirely + # Only repoint00217 is at the map's pivot angle, so repoint00218 drops out + # of every product it appears in. assert [ [str(file_path.filename) for file_path in processing_input.imap_file_paths] for processing_input in result.get_processing_inputs() - ] == [[valid_pset], [valid_de], [ancillary]] - # Ancillary files are not loaded to be checked for a pivot angle - assert mock_load_cdf.call_count == 4 + ] == [[goodtimes[0]], [histrates[0]], [bgrates[0]], [ancillary]] + # Only the goodtimes files are loaded, to read their pivot angle + assert mock_load_cdf.call_count == 2 + + +@mock.patch("imap_processing.cli.load_cdf") +@mock.patch("imap_processing.cli.ProcessInstrument.pre_processing") +def test_lo_pre_processing_drops_goodtimes_without_pivot( + mock_super_pre_processing, mock_load_cdf +): + """Test that a pointing whose goodtimes has no pivot angle is dropped.""" + goodtimes = "imap_lo_l1b_goodtimes_20250415-repoint00217_v001.cdf" + histrates = "imap_lo_l1b_histrates_20250415-repoint00217_v001.cdf" + + base_collection = ProcessingInputCollection( + ScienceInput(goodtimes), ScienceInput(histrates) + ) + mock_super_pre_processing.return_value = base_collection + mock_load_cdf.side_effect = [xr.Dataset()] # No pivot angle at all + + instrument = Lo( + "l2", + "l090-ena-h-sf-nsp-ram-hae-6deg-3mo", + base_collection.serialize(), + "20250415", + "20250715", + "v001", + False, + ) + result = instrument.pre_processing() + + assert list(result.get_processing_inputs()) == [] @mock.patch("imap_processing.cli.quaternions.process_quaternions", autospec=True) From b2bd558615d34ac207f414d73daf6389c7295c89 Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Tue, 28 Jul 2026 11:08:12 -0400 Subject: [PATCH 05/10] PivotAngleSpec class; get_instrument_spin_angle_bins function --- imap_processing/lo/constants.py | 62 +++++++++++++++++++----- imap_processing/lo/l1b/lo_l1b.py | 28 +++++------ imap_processing/lo/l2/lo_l2.py | 33 +++---------- imap_processing/spice/spin.py | 38 +++++++++++++++ imap_processing/tests/lo/test_lo_l1b.py | 9 ++-- imap_processing/tests/lo/test_lo_l2.py | 22 ++++----- imap_processing/tests/spice/test_spin.py | 26 ++++++++++ 7 files changed, 151 insertions(+), 67 deletions(-) diff --git a/imap_processing/lo/constants.py b/imap_processing/lo/constants.py index 270417474..40de8b1f9 100644 --- a/imap_processing/lo/constants.py +++ b/imap_processing/lo/constants.py @@ -1,7 +1,39 @@ """Constants for IMAP-Lo.""" from dataclasses import dataclass -from typing import ClassVar +from typing import ClassVar, NamedTuple + + +class PivotAngleSpec(NamedTuple): + """ + Pivot angle [degrees] and associated settings for a nominal pivot index. + + Attributes + ---------- + pointing_index : int + The unique pointing "index" (matching technical documentation for Imap-Lo) + nominal : float + Nominal pivot angle. + min : float + Lower bound of the acceptable pivot angle range. + max : float + Upper bound of the acceptable pivot angle range. + bg_rate_ram : float, optional + RAM background-rate threshold [counts/s] for this pivot angle. ``None`` + if no pivot-specific value is known, in which case + ``LoConstants.THRESHOLD_BG_RATE_RAM_DEFAULT`` applies. + bg_rate_anti_ram : float, optional + Anti-RAM background-rate threshold [counts/s] for this pivot angle. + ``None`` if no pivot-specific value is known, in which case + ``LoConstants.THRESHOLD_BG_RATE_ANTI_RAM_DEFAULT`` applies. + """ + + pointing_index: int + nominal: float + min: float + max: float + bg_rate_ram: float | None = None + bg_rate_anti_ram: float | None = None @dataclass(frozen=True) @@ -104,19 +136,25 @@ class LoConstants: # Minimum non-zero background rate floor = nominal / divisor BG_RATE_FLOOR_DIVISOR: ClassVar[dict[str, float]] = {"H": 50.0, "O": 150.0} - # Background-rate thresholds [counts/s] by pivot-angle range (low, high) [deg]. - # Each value is (ram_threshold, anti_ram_threshold). - # The first matching open interval (low < pivot < high) is used; if none matches, - # THRESHOLD_BG_RATE_RAM_DEFAULT / THRESHOLD_BG_RATE_ANTI_RAM_DEFAULT apply. - PIVOT_ANGLE_THRESHOLDS: ClassVar[dict[tuple[float, float], tuple[float, float]]] = { - (88.0, 92.0): (0.028, 0.014), - (73.0, 77.0): (0.035, 0.0175), - (103.0, 107.0): (0.0224, 0.0112), + # Pivot angle specs keyed by nominal pivot angle. A measured pivot angle is + # assigned to the first spec whose [min, max] range contains it; the ranges are + # disjoint, so the ordering here does not matter. + PIVOT_ANGLES: ClassVar[dict[int, PivotAngleSpec]] = { + 60: PivotAngleSpec(1, 60.0, 55.0, 65.0, None, None), + 75: PivotAngleSpec(2, 75.0, 70.0, 80.0, 0.035, 0.0175), + 90: PivotAngleSpec(3, 90.0, 85.0, 95.0, 0.028, 0.014), + 105: PivotAngleSpec(4, 105.0, 100.0, 110.0, 0.0224, 0.0112), + 120: PivotAngleSpec(5, 120.0, 115.0, 125.0, None, None), + 135: PivotAngleSpec(6, 135.0, 130.0, 140.0, None, None), + 148: PivotAngleSpec(7, 148.0, 143.0, 153.0, None, None), + 160: PivotAngleSpec(8, 160.0, 155.0, 165.0, None, None), } - # Default background-rate thresholds [counts/s] when no pivot range matches. - THRESHOLD_BG_RATE_RAM_DEFAULT: float = 0.0175 - THRESHOLD_BG_RATE_ANTI_RAM_DEFAULT: float = 0.00875 + # Default background-rate thresholds [counts/s] when the pivot angle matches no + # spec in PIVOT_ANGLES, or the matching spec has no pivot-specific value. + # Currently set to nominal values for the 90-deg pivot angle. + THRESHOLD_BG_RATE_RAM_DEFAULT: float = 0.028 + THRESHOLD_BG_RATE_ANTI_RAM_DEFAULT: float = 0.014 # Maximum time gap [s] between consecutive histogram epochs before treating them as # separate intervals. diff --git a/imap_processing/lo/l1b/lo_l1b.py b/imap_processing/lo/l1b/lo_l1b.py index f375388ad..7923fda2e 100644 --- a/imap_processing/lo/l1b/lo_l1b.py +++ b/imap_processing/lo/l1b/lo_l1b.py @@ -22,7 +22,6 @@ SpiceFrame, cartesian_to_latitudinal, frame_transform, - get_spacecraft_to_instrument_spin_phase_offset, lo_instrument_pointing, ) from imap_processing.spice.repoint import ( @@ -32,6 +31,7 @@ interpolate_repoint_data, ) from imap_processing.spice.spin import ( + get_instrument_spin_angle_bins, get_spin_data, get_spin_number, interpolate_spin_data, @@ -2151,7 +2151,6 @@ def calculate_star_sensor_profiles_by_group( sampling_cadence: float, spin_period: float, group_size: int = 64, - start_angle_offset: float = 62.0, end_bins_to_exclude: int = c.STAR_END_BINS_TO_EXCLUDE, min_count_threshold: int = c.STAR_MIN_COUNT_THRESHOLD, bin_offset: float = 0.5, @@ -2172,8 +2171,6 @@ def calculate_star_sensor_profiles_by_group( Spin period in seconds. group_size : int Number of records per group (default: 64). - start_angle_offset : float - Starting angle offset in degrees (default: 62.0 = 90° - 28°). end_bins_to_exclude : int Number of ending bins to exclude from each average (default: 2). min_count_threshold : int @@ -2202,11 +2199,12 @@ def calculate_star_sensor_profiles_by_group( valid_indices = np.where(valid_mask)[0] n_valid = len(valid_indices) - # Calculate spin angles (same for all groups) + # Calculate spin angles (same for all groups). The samples are taken at a + # fixed cadence, so they do not evenly divide a spin. deg_per_bin = 360.0 * (sampling_cadence / 1000.0) / spin_period - bin_indices = np.arange(720) - sample_centers = (bin_indices + bin_offset) * deg_per_bin - spin_angle = (start_angle_offset + sample_centers) % 360.0 + spin_angle = get_instrument_spin_angle_bins( + SpiceFrame.IMAP_LO, 720, deg_per_bin=deg_per_bin, bin_offset=bin_offset + ) if n_valid == 0: logger.warning( @@ -2335,9 +2333,6 @@ def l1b_star( logger.info(f"Using spin duration from spin data: {spin_duration:.6f} s") # TODO: Read from ancillary config file when available - sc_to_inst_angle_offset = 360 * get_spacecraft_to_instrument_spin_phase_offset( - SpiceFrame.IMAP_LO - ) end_bins_to_exclude = c.STAR_END_BINS_TO_EXCLUDE min_count_threshold = c.STAR_MIN_COUNT_THRESHOLD @@ -2361,7 +2356,6 @@ def l1b_star( sampling_cadence, spin_duration, group_size=group_size, - start_angle_offset=sc_to_inst_angle_offset, end_bins_to_exclude=end_bins_to_exclude, min_count_threshold=min_count_threshold, bin_offset=bin_offset, @@ -2521,10 +2515,12 @@ def l1b_bgrates_and_goodtimes( # noqa: PLR0912 # Choose background rate thresholds based on pivot orientation. bg_rate_ram_nominal = c.THRESHOLD_BG_RATE_RAM_DEFAULT bg_rate_anti_ram_nominal = c.THRESHOLD_BG_RATE_ANTI_RAM_DEFAULT - for (low, high), (ram_thresh, anti_ram_thresh) in c.PIVOT_ANGLE_THRESHOLDS.items(): - if low < pivot < high: - bg_rate_ram_nominal = ram_thresh - bg_rate_anti_ram_nominal = anti_ram_thresh + for pivot_spec in c.PIVOT_ANGLES.values(): + if pivot_spec.min <= pivot <= pivot_spec.max: + if pivot_spec.bg_rate_ram is not None: + bg_rate_ram_nominal = pivot_spec.bg_rate_ram + if pivot_spec.bg_rate_anti_ram is not None: + bg_rate_anti_ram_nominal = pivot_spec.bg_rate_anti_ram break # Manual overrides of the anti-RAM threshold for anomalous days. diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index 064581178..e6de3aaf7 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -10,10 +10,8 @@ from imap_processing.ena_maps.utils.naming import MapDescriptor from imap_processing.lo.constants import LoConstants as c # noqa: N813 from imap_processing.lo.l1c.lo_l1c import compute_pointing_directions -from imap_processing.spice.geometry import ( - SpiceFrame, - get_spacecraft_to_instrument_spin_phase_offset, -) +from imap_processing.spice.geometry import SpiceFrame +from imap_processing.spice.spin import get_instrument_spin_angle_bins from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_met logger = logging.getLogger(__name__) @@ -253,7 +251,12 @@ def _accumulate_pointing( pointing_exposure = histrates["exposure_time_6deg"].values[in_goodtime].sum(axis=0) background_rates = np.atleast_2d(bgrates[f"{species}_background_rates"].values)[0] - spin_angles = _dps_spin_angles() + # The L1B histogram spin bins are hardware spin-phase bins referenced to + # the spacecraft spin pulse, so they are rotated onto the instrument (DPS) + # spin angle, exactly as the L1B star-sensor product does. + spin_angles = get_instrument_spin_angle_bins( + SpiceFrame.IMAP_LO, c.N_SPIN_ANGLE_BINS + ) # The whole pointing is projected from the middle of its good times, which # is where the despun frame is sampled. epoch = met_to_ttj2000ns((gt_start.min() + gt_end.max()) / 2.0) @@ -287,26 +290,6 @@ def _accumulate_pointing( sky_map.max_epoch = max(sky_map.max_epoch, int(met_to_ttj2000ns(gt_end.max()))) -def _dps_spin_angles() -> np.ndarray: - """ - Get the despun-frame azimuth of each histogram spin-angle bin center. - - The L1B histogram spin bins are hardware spin-phase bins referenced to the - spacecraft spin pulse, NOT the instrument (DPS) spin angle. A bin center is - converted to the IMAP_DPS azimuth by adding the spacecraft to instrument - spin-phase offset, exactly as the L1B star-sensor product does. - - Returns - ------- - np.ndarray - The IMAP_DPS azimuth [degrees] of each of the histogram spin bins. - """ - bin_width = 360.0 / c.N_SPIN_ANGLE_BINS - bin_centers = (np.arange(c.N_SPIN_ANGLE_BINS) + 0.5) * bin_width - offset = get_spacecraft_to_instrument_spin_phase_offset(SpiceFrame.IMAP_LO) * 360.0 - return np.mod(bin_centers + offset, 360.0) - - def _pixel_indices( sky_map: RectangularSkyMap, longitude: np.ndarray, latitude: np.ndarray ) -> np.ndarray: diff --git a/imap_processing/spice/spin.py b/imap_processing/spice/spin.py index 20264e473..968a4841f 100644 --- a/imap_processing/spice/spin.py +++ b/imap_processing/spice/spin.py @@ -348,3 +348,41 @@ def get_instrument_spin_phase( instrument ) return (spacecraft_spin_phase + instrument_spin_phase_offset) % 1 + + +def get_instrument_spin_angle_bins( + instrument: SpiceFrame, + n_bins: int, + deg_per_bin: float | None = None, + bin_offset: float = 0.5, +) -> npt.NDArray: + """ + Get the instrument spin angle of each hardware spin-angle bin. + + Instruments accumulate data into bins referenced to the spacecraft spin + pulse. A bin is converted to the instrument spin angle by adding the + spacecraft to instrument spin phase offset. + + Parameters + ---------- + instrument : SpiceFrame + Instrument frame the bins are wanted in. + n_bins : int + Number of spin-angle bins. + deg_per_bin : float, optional + Width [degrees] of a bin. Defaults to ``360 / n_bins``, i.e. bins that + evenly divide a spin. Bins sampled at a fixed cadence do not evenly + divide a spin and must pass their own width. + bin_offset : float + Fractional position within a bin the angle is reported at. Default is + 0.5 for the bin center; use 0.0 for the leading edge. + + Returns + ------- + spin_angle : np.ndarray + The instrument spin angle [degrees, 0-360) of each bin. + """ + if deg_per_bin is None: + deg_per_bin = 360.0 / n_bins + start_angle = 360.0 * get_spacecraft_to_instrument_spin_phase_offset(instrument) + return np.mod(start_angle + (np.arange(n_bins) + bin_offset) * deg_per_bin, 360.0) diff --git a/imap_processing/tests/lo/test_lo_l1b.py b/imap_processing/tests/lo/test_lo_l1b.py index 4291dbdd8..937b83376 100644 --- a/imap_processing/tests/lo/test_lo_l1b.py +++ b/imap_processing/tests/lo/test_lo_l1b.py @@ -1802,8 +1802,12 @@ def test_profiles_by_group_handles_no_valid_records(self, mock_repoint): assert len(group_epochs) == 0 # No valid records assert avg_amplitudes.shape == (0, 720) + @patch( + "imap_processing.spice.spin.get_spacecraft_to_instrument_spin_phase_offset", + return_value=350.0 / 360.0, + ) @patch("imap_processing.lo.l1b.lo_l1b.interpolate_repoint_data") - def test_profiles_by_group_angle_wrapping(self, mock_repoint): + def test_profiles_by_group_angle_wrapping(self, mock_repoint, mock_offset): """Test that spin angles wrap correctly to [0, 360) range.""" # Arrange mock_repoint.return_value = pd.DataFrame({"repoint_in_progress": [False]}) @@ -1824,7 +1828,6 @@ def test_profiles_by_group_angle_wrapping(self, mock_repoint): l1a_star, sampling_cadence=21.0, spin_period=15.0, - start_angle_offset=350.0, # Large offset to test wrapping ) # Assert @@ -3023,7 +3026,7 @@ def test_l1b_bgrates_sigma_when_anti_ram_nominal_is_zero( "imap_processing.lo.l1b.lo_l1b.get_pointing_times_from_id", return_value=(met_start, met_start + 1), ), - patch.object(LoConstants, "PIVOT_ANGLE_THRESHOLDS", {}), + patch.object(LoConstants, "PIVOT_ANGLES", {}), patch.object(LoConstants, "THRESHOLD_BG_RATE_ANTI_RAM_DEFAULT", 0.0), ): bgrates_ds, _ = l1b_bgrates_and_goodtimes( diff --git a/imap_processing/tests/lo/test_lo_l2.py b/imap_processing/tests/lo/test_lo_l2.py index 8903684e8..2559e2579 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -13,12 +13,13 @@ BGRATES, GOODTIMES, HISTRATES, - _dps_spin_angles, _group_inputs_by_pointing, _pixel_indices, _spin_phase_mask, lo_l2, ) +from imap_processing.spice.geometry import SpiceFrame +from imap_processing.spice.spin import get_instrument_spin_angle_bins from imap_processing.spice.time import met_to_ttj2000ns # A full-spin map, so that every spin-angle bin lands on it. @@ -29,6 +30,9 @@ N_SPIN_BINS = LoConstants.N_SPIN_ANGLE_BINS PIVOT = 90.0 +# The instrument spin angle of each histogram spin bin, as L2 computes it. +SPIN_ANGLES = get_instrument_spin_angle_bins(SpiceFrame.IMAP_LO, N_SPIN_BINS) + # Good-time window [MET seconds] that the "in-window" histogram epochs fall in. GT_START = 511_000_000.0 GT_END = 511_000_600.0 @@ -330,15 +334,13 @@ def test_systematic_error_bounds(self, full_map): class TestGeometry: """The spin-angle to sky-pixel geometry.""" - def test_dps_spin_angles_carry_the_offset(self): + def test_spin_angles_carry_the_offset(self): """The hardware spin bins are rotated onto the instrument frame.""" - angles = _dps_spin_angles() - - assert angles.size == N_SPIN_BINS + assert SPIN_ANGLES.size == N_SPIN_BINS # IMAP-Lo sits 60 degrees from the spacecraft spin pulse, so bin 0's # center (3 degrees) becomes 63 degrees in the despun frame. - np.testing.assert_allclose(angles[0], 63.0) - np.testing.assert_allclose(np.diff(np.sort(angles)), 6.0) + np.testing.assert_allclose(SPIN_ANGLES[0], 63.0) + np.testing.assert_allclose(np.diff(np.sort(SPIN_ANGLES)), 6.0) def test_pixel_indices_match_the_map_grid(self): """Directions are placed in the pixel whose center they are nearest.""" @@ -364,16 +366,14 @@ def test_pixel_indices_wrap_longitude(self): ) def test_spin_phase_mask(self, spin_phase, expected): """Ram and anti-ram split the spin; a full map keeps all of it.""" - angles = _dps_spin_angles() - - mask = _spin_phase_mask(angles, PIVOT, spin_phase) + mask = _spin_phase_mask(SPIN_ANGLES, PIVOT, spin_phase) assert mask.sum() == expected def test_spin_phase_mask_rejects_unknown(self): """An unknown spin phase is an error, not a silently empty map.""" with pytest.raises(ValueError, match="Invalid spin phase"): - _spin_phase_mask(_dps_spin_angles(), PIVOT, "sideways") + _spin_phase_mask(SPIN_ANGLES, PIVOT, "sideways") class TestInputGrouping: diff --git a/imap_processing/tests/spice/test_spin.py b/imap_processing/tests/spice/test_spin.py index 808e093dc..9aa17d6d7 100644 --- a/imap_processing/tests/spice/test_spin.py +++ b/imap_processing/tests/spice/test_spin.py @@ -312,3 +312,29 @@ def test_get_spin_start_met(fake_spin_data): # when we have a linear ramp of data points between spins start_mets = spin.interpolate_spin_data(np.linspace(32, 37, 3))["spin_start_met"] np.testing.assert_array_equal(start_mets, [30, 30, 30]) + + +def test_get_instrument_spin_angle_bins(): + """Bins that evenly divide a spin are reported at their centers.""" + spin_angles = spin.get_instrument_spin_angle_bins(SpiceFrame.IMAP_LO, 60) + + assert spin_angles.shape == (60,) + # IMAP-Lo sits 60 degrees from the spacecraft spin pulse, so the center of + # bin 0 (3 degrees) is at 63 degrees in the instrument frame. + np.testing.assert_allclose(spin_angles[0], 63.0) + np.testing.assert_allclose(np.diff(np.sort(spin_angles)), 6.0) + + +def test_get_instrument_spin_angle_bins_uneven_bins(): + """Bins sampled at a fixed cadence wrap around the spin.""" + # 720 samples of 21 ms each cover more than one 15 second spin. + deg_per_bin = 360.0 * 0.021 / 15.0 + spin_angles = spin.get_instrument_spin_angle_bins( + SpiceFrame.IMAP_LO, 720, deg_per_bin=deg_per_bin, bin_offset=0.0 + ) + + assert np.all((spin_angles >= 0) & (spin_angles < 360)) + # The first sample sits at the instrument offset, and the bins wrap past it. + np.testing.assert_allclose(spin_angles[0], 60.0) + np.testing.assert_allclose(spin_angles[1], 60.0 + deg_per_bin) + assert np.any(spin_angles < 60.0) From 3b012c1f56f7284ec315a8d5aeef5562d640d48f Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Tue, 28 Jul 2026 14:23:39 -0400 Subject: [PATCH 06/10] code simplification (with key groupings etc) --- imap_processing/cli.py | 30 ++++----- imap_processing/lo/l2/lo_l2.py | 66 +++++++++----------- imap_processing/tests/lo/test_lo_l2.py | 84 +++++++++++++++----------- imap_processing/tests/test_cli.py | 10 ++- 4 files changed, 103 insertions(+), 87 deletions(-) diff --git a/imap_processing/cli.py b/imap_processing/cli.py index e020c266e..81ffd5fb6 100644 --- a/imap_processing/cli.py +++ b/imap_processing/cli.py @@ -1381,22 +1381,24 @@ def do_processing( datasets = lo_l1c.lo_l1c(data_dict, anc_dependencies) elif self.data_level == "l2": - sci_dependencies: dict[str, list[xr.Dataset]] = { - lo_l2.GOODTIMES: [], - lo_l2.BGRATES: [], - lo_l2.HISTRATES: [], - } - science_files = [] - for descriptor in ("goodtimes", "bgrates", "histrates"): - science_files += dependencies.get_file_paths( - source="lo", data_type="l1b", descriptor=descriptor - ) anc_dependencies = dependencies.get_file_paths(data_type="ancillary") - # Load every pointing of the map window, grouped by product. - for file in science_files: - dataset = load_cdf(file) - sci_dependencies[dataset.attrs["Logical_source"]].append(dataset) + # Load every pointing of the map window, grouped into the products + # of each pointing. + sci_dependencies: dict[int, dict[str, xr.Dataset]] = {} + for descriptor in lo_l2.REQUIRED_PRODUCTS: + for file in dependencies.get_file_paths( + source="lo", data_type="l1b", descriptor=descriptor + ): + repointing = imap_data_access.ScienceFilePath(file.name).repointing + if repointing is None: + logger.warning( + f"Dropping {file.name}, it covers no single pointing." + ) + continue + sci_dependencies.setdefault(repointing, {})[descriptor] = load_cdf( + file + ) datasets = lo_l2.lo_l2(sci_dependencies, anc_dependencies, self.descriptor) return datasets diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index e6de3aaf7..7ce3637db 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -16,10 +16,8 @@ logger = logging.getLogger(__name__) -# The L1B products a map is built from, one set per pointing. -GOODTIMES = "imap_lo_l1b_goodtimes" -BGRATES = "imap_lo_l1b_bgrates" -HISTRATES = "imap_lo_l1b_histrates" +# The descriptors of the L1B products a map is built from, one set per pointing. +REQUIRED_PRODUCTS = ("goodtimes", "bgrates", "histrates") # ============================================================================= # MAIN ENTRY POINT @@ -27,7 +25,9 @@ def lo_l2( - sci_dependencies: dict, anc_dependencies: list, descriptor: str + sci_dependencies: dict[int, dict[str, xr.Dataset]], + anc_dependencies: list, + descriptor: str, ) -> list[xr.Dataset]: """ Process IMAP-Lo L1B data into an L2 sky map. @@ -44,11 +44,9 @@ def lo_l2( Parameters ---------- - sci_dependencies : dict - Dictionary of the input datasets, keyed by logical source, each a list - of datasets covering the pointings of the map window. Must contain - ``"imap_lo_l1b_goodtimes"``, ``"imap_lo_l1b_bgrates"`` and - ``"imap_lo_l1b_histrates"``. + sci_dependencies : dict[int, dict[str, xr.Dataset]] + The input datasets covering the pointings of the map window, keyed by + repointing and then by product descriptor. anc_dependencies : list List of ancillary file paths. Unused, the calibration constants of the map live in ``LoConstants``. @@ -83,7 +81,7 @@ def lo_l2( if not isinstance(sky_map, RectangularSkyMap): raise NotImplementedError("HEALPix map output not supported for Lo") - pointings = _group_inputs_by_pointing(sci_dependencies) + pointings = _complete_pointings(sci_dependencies) logger.info(f"Building {descriptor} from {len(pointings)} pointings") shape = (c.N_ESA_LEVELS, sky_map.num_points) @@ -95,7 +93,7 @@ def lo_l2( esa_mode = 0 for repointing, (goodtimes, bgrates, histrates) in sorted(pointings.items()): - logger.debug(f"Accumulating {repointing}") + logger.debug(f"Accumulating repoint{repointing:05d}") esa_mode = _get_esa_mode(histrates) _accumulate_pointing( goodtimes, @@ -129,47 +127,43 @@ def lo_l2( # ============================================================================= -def _group_inputs_by_pointing(sci_dependencies: dict) -> dict[str, tuple]: +def _complete_pointings( + sci_dependencies: dict[int, dict[str, xr.Dataset]], +) -> dict[int, tuple]: """ - Group the L1B inputs into the (goodtimes, bgrates, histrates) of a pointing. - - Each input product records the pointing it covers in its ``Repointing`` - global attribute. Pointings missing any of the three products cannot be - mapped and are dropped. + Reduce the grouped inputs to the pointings that can be mapped. Parameters ---------- - sci_dependencies : dict - Dictionary of the input datasets, keyed by logical source. + sci_dependencies : dict[int, dict[str, xr.Dataset]] + The input datasets of each pointing, keyed by repointing and then by + product descriptor. Returns ------- - dict[str, tuple] - The (goodtimes, bgrates, histrates) datasets of each pointing, keyed by - repointing. + dict[int, tuple] + The (goodtimes, bgrates, histrates) datasets of each mappable pointing, + keyed by repointing. Raises ------ KeyError If any of the three required products is missing entirely. """ - by_pointing: dict[str, dict[str, xr.Dataset]] = {} - for logical_source in (GOODTIMES, BGRATES, HISTRATES): - for dataset in sci_dependencies[logical_source]: - repointing = dataset.attrs.get("Repointing", "") - by_pointing.setdefault(repointing, {})[logical_source] = dataset + found = {product for products in sci_dependencies.values() for product in products} + missing_products = set(REQUIRED_PRODUCTS) - found + if missing_products: + raise KeyError(f"No input files for {sorted(missing_products)}") pointings = {} - for repointing, products in by_pointing.items(): - missing = {GOODTIMES, BGRATES, HISTRATES} - set(products) + for repointing, products in sci_dependencies.items(): + missing = set(REQUIRED_PRODUCTS) - set(products) if missing: - logger.warning(f"Dropping {repointing}, it has no {sorted(missing)}") + logger.warning( + f"Dropping repoint{repointing:05d}, it has no {sorted(missing)}" + ) continue - pointings[repointing] = ( - products[GOODTIMES], - products[BGRATES], - products[HISTRATES], - ) + pointings[repointing] = tuple(products[p] for p in REQUIRED_PRODUCTS) return pointings diff --git a/imap_processing/tests/lo/test_lo_l2.py b/imap_processing/tests/lo/test_lo_l2.py index 2559e2579..29be956a7 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -10,10 +10,7 @@ from imap_processing.ena_maps.utils.naming import MapDescriptor from imap_processing.lo.constants import LoConstants from imap_processing.lo.l2.lo_l2 import ( - BGRATES, - GOODTIMES, - HISTRATES, - _group_inputs_by_pointing, + _complete_pointings, _pixel_indices, _spin_phase_mask, lo_l2, @@ -40,7 +37,16 @@ OUT_METS = [510_990_000.0, 511_010_000.0] -def make_pointing(repointing="repoint00100", pivot=PIVOT, seed=42): +def product_attrs(repointing, product): + """The global attributes an L1B input of a pointing is written with.""" + + return { + "Repointing": f"repoint{repointing:05d}", + "Logical_source": f"imap_lo_l1b_{product}", + } + + +def make_pointing(repointing=100, pivot=PIVOT, seed=42): """Build the three synthetic L1B inputs of one pointing. The in-window epochs carry modest counts and exposure; the out-of-window @@ -67,7 +73,7 @@ def make_pointing(repointing="repoint00100", pivot=PIVOT, seed=42): "esa_mode": ("epoch", np.zeros(mets.size, dtype=int)), }, coords={"epoch": met_to_ttj2000ns(mets)}, - attrs={"Repointing": repointing, "Logical_source": HISTRATES}, + attrs=product_attrs(repointing, "histrates"), ) goodtimes = xr.Dataset( { @@ -76,30 +82,35 @@ def make_pointing(repointing="repoint00100", pivot=PIVOT, seed=42): "gt_end_met": ("epoch", [GT_END]), }, coords={"epoch": [0]}, - attrs={"Repointing": repointing, "Logical_source": GOODTIMES}, + attrs=product_attrs(repointing, "goodtimes"), ) background = np.array([0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07]) bgrates = xr.Dataset( {"h_background_rates": (["epoch", "esa_step"], background[np.newaxis, :])}, coords={"epoch": [0], "esa_step": np.arange(1, N_ESA + 1)}, - attrs={"Repointing": repointing, "Logical_source": BGRATES}, + attrs=product_attrs(repointing, "bgrates"), ) return { - GOODTIMES: goodtimes, - BGRATES: bgrates, - HISTRATES: histrates, + "repointing": repointing, + "goodtimes": goodtimes, + "bgrates": bgrates, + "histrates": histrates, "expected_counts": counts[in_idx].sum(axis=(0, 2)), "expected_exposure": exposure[in_idx].sum(axis=(0, 2)), "background": background, } -def as_dependencies(*pointings): - """Turn pointings into the sci_dependencies lo_l2 takes.""" +def as_dependencies(*pointings, products=("goodtimes", "bgrates", "histrates")): + """Turn pointings into the sci_dependencies lo_l2 takes. + + The CLI keys the inputs by repointing and then by product descriptor, see + ``cli.Lo.do_processing``. + """ return { - source: [pointing[source] for pointing in pointings] - for source in (GOODTIMES, BGRATES, HISTRATES) + pointing["repointing"]: {product: pointing[product] for product in products} + for pointing in pointings } @@ -224,7 +235,7 @@ def test_out_of_goodtime_epochs_are_excluded(self, full_map): def test_pointings_accumulate(self, one_pointing): """Two pointings contribute twice the counts of one.""" - other = make_pointing(repointing="repoint00101", seed=7) + other = make_pointing(repointing=101, seed=7) with patch( "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el", @@ -376,36 +387,41 @@ def test_spin_phase_mask_rejects_unknown(self): _spin_phase_mask(SPIN_ANGLES, PIVOT, "sideways") -class TestInputGrouping: - """Sorting the input products into pointings.""" +class TestPointingSelection: + """Reducing the grouped inputs to the pointings that can be mapped.""" - def test_inputs_are_grouped_by_repointing(self, one_pointing): - """Each pointing's three products are grouped together.""" - other = make_pointing(repointing="repoint00101") + def test_complete_pointings_are_kept_in_product_order(self, one_pointing): + """Each pointing's products are ordered goodtimes, bgrates, histrates.""" + other = make_pointing(repointing=101) - pointings = _group_inputs_by_pointing(as_dependencies(one_pointing, other)) + pointings = _complete_pointings(as_dependencies(one_pointing, other)) - assert set(pointings) == {"repoint00100", "repoint00101"} - assert pointings["repoint00100"][0] is one_pointing[GOODTIMES] - assert pointings["repoint00100"][2] is one_pointing[HISTRATES] + assert set(pointings) == {100, 101} + assert pointings[100] == ( + one_pointing["goodtimes"], + one_pointing["bgrates"], + one_pointing["histrates"], + ) def test_incomplete_pointings_are_dropped(self, one_pointing, caplog): """A pointing missing one of the three products cannot be mapped.""" - dependencies = as_dependencies(one_pointing) - dependencies[BGRATES] = [] + incomplete = make_pointing(repointing=101) + dependencies = as_dependencies(one_pointing) | as_dependencies( + incomplete, products=("goodtimes", "histrates") + ) - pointings = _group_inputs_by_pointing(dependencies) + pointings = _complete_pointings(dependencies) - assert pointings == {} - assert "imap_lo_l1b_bgrates" in caplog.text + assert set(pointings) == {100} + assert "repoint00101" in caplog.text + assert "bgrates" in caplog.text def test_missing_product_raises(self, one_pointing): """A map cannot be made without all three products.""" - dependencies = as_dependencies(one_pointing) - del dependencies[HISTRATES] + dependencies = as_dependencies(one_pointing, products=("goodtimes", "bgrates")) - with pytest.raises(KeyError, match=HISTRATES): - _group_inputs_by_pointing(dependencies) + with pytest.raises(KeyError, match="histrates"): + _complete_pointings(dependencies) class TestUnsupported: diff --git a/imap_processing/tests/test_cli.py b/imap_processing/tests/test_cli.py index 6a5acb70a..15aedc50b 100644 --- a/imap_processing/tests/test_cli.py +++ b/imap_processing/tests/test_cli.py @@ -510,11 +510,15 @@ def test_lo_l2(mock_lo_pre_processing, mock_lo_l2, mock_instrument_dependencies) ) instrument.process() + # Grouped by the repointing in the filename and the descriptor queried by. mock_lo_l2.assert_called_once_with( { - "imap_lo_l1b_goodtimes": [mock_goodtimes], - "imap_lo_l1b_bgrates": [mock_bgrates], - "imap_lo_l1b_histrates": [mock_histrates_1, mock_histrates_2], + 217: { + "goodtimes": mock_goodtimes, + "bgrates": mock_bgrates, + "histrates": mock_histrates_1, + }, + 218: {"histrates": mock_histrates_2}, }, [], descriptor, From 7a1f49e27241dab9abaed2c308241b6051db627f Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Tue, 28 Jul 2026 14:48:12 -0400 Subject: [PATCH 07/10] removed get_instrument_spin_angle_bins function (avoid touching common code) --- imap_processing/lo/l1b/lo_l1b.py | 18 +++++++---- imap_processing/lo/l2/lo_l2.py | 33 +++++++++++++++----- imap_processing/spice/spin.py | 38 ------------------------ imap_processing/tests/lo/test_lo_l1b.py | 7 ++--- imap_processing/tests/lo/test_lo_l2.py | 22 +++++++------- imap_processing/tests/spice/test_spin.py | 26 ---------------- 6 files changed, 50 insertions(+), 94 deletions(-) diff --git a/imap_processing/lo/l1b/lo_l1b.py b/imap_processing/lo/l1b/lo_l1b.py index 7923fda2e..ad10569de 100644 --- a/imap_processing/lo/l1b/lo_l1b.py +++ b/imap_processing/lo/l1b/lo_l1b.py @@ -22,6 +22,7 @@ SpiceFrame, cartesian_to_latitudinal, frame_transform, + get_spacecraft_to_instrument_spin_phase_offset, lo_instrument_pointing, ) from imap_processing.spice.repoint import ( @@ -31,7 +32,6 @@ interpolate_repoint_data, ) from imap_processing.spice.spin import ( - get_instrument_spin_angle_bins, get_spin_data, get_spin_number, interpolate_spin_data, @@ -2151,6 +2151,7 @@ def calculate_star_sensor_profiles_by_group( sampling_cadence: float, spin_period: float, group_size: int = 64, + start_angle_offset: float = 62.0, end_bins_to_exclude: int = c.STAR_END_BINS_TO_EXCLUDE, min_count_threshold: int = c.STAR_MIN_COUNT_THRESHOLD, bin_offset: float = 0.5, @@ -2171,6 +2172,8 @@ def calculate_star_sensor_profiles_by_group( Spin period in seconds. group_size : int Number of records per group (default: 64). + start_angle_offset : float + Starting angle offset in degrees (default: 62.0 = 90° - 28°). end_bins_to_exclude : int Number of ending bins to exclude from each average (default: 2). min_count_threshold : int @@ -2199,12 +2202,11 @@ def calculate_star_sensor_profiles_by_group( valid_indices = np.where(valid_mask)[0] n_valid = len(valid_indices) - # Calculate spin angles (same for all groups). The samples are taken at a - # fixed cadence, so they do not evenly divide a spin. + # Calculate spin angles (same for all groups) deg_per_bin = 360.0 * (sampling_cadence / 1000.0) / spin_period - spin_angle = get_instrument_spin_angle_bins( - SpiceFrame.IMAP_LO, 720, deg_per_bin=deg_per_bin, bin_offset=bin_offset - ) + bin_indices = np.arange(720) + sample_centers = (bin_indices + bin_offset) * deg_per_bin + spin_angle = (start_angle_offset + sample_centers) % 360.0 if n_valid == 0: logger.warning( @@ -2333,6 +2335,9 @@ def l1b_star( logger.info(f"Using spin duration from spin data: {spin_duration:.6f} s") # TODO: Read from ancillary config file when available + sc_to_inst_angle_offset = 360 * get_spacecraft_to_instrument_spin_phase_offset( + SpiceFrame.IMAP_LO + ) end_bins_to_exclude = c.STAR_END_BINS_TO_EXCLUDE min_count_threshold = c.STAR_MIN_COUNT_THRESHOLD @@ -2356,6 +2361,7 @@ def l1b_star( sampling_cadence, spin_duration, group_size=group_size, + start_angle_offset=sc_to_inst_angle_offset, end_bins_to_exclude=end_bins_to_exclude, min_count_threshold=min_count_threshold, bin_offset=bin_offset, diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index 7ce3637db..de3ba04c9 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -10,8 +10,10 @@ from imap_processing.ena_maps.utils.naming import MapDescriptor from imap_processing.lo.constants import LoConstants as c # noqa: N813 from imap_processing.lo.l1c.lo_l1c import compute_pointing_directions -from imap_processing.spice.geometry import SpiceFrame -from imap_processing.spice.spin import get_instrument_spin_angle_bins +from imap_processing.spice.geometry import ( + SpiceFrame, + get_spacecraft_to_instrument_spin_phase_offset, +) from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_met logger = logging.getLogger(__name__) @@ -245,12 +247,7 @@ def _accumulate_pointing( pointing_exposure = histrates["exposure_time_6deg"].values[in_goodtime].sum(axis=0) background_rates = np.atleast_2d(bgrates[f"{species}_background_rates"].values)[0] - # The L1B histogram spin bins are hardware spin-phase bins referenced to - # the spacecraft spin pulse, so they are rotated onto the instrument (DPS) - # spin angle, exactly as the L1B star-sensor product does. - spin_angles = get_instrument_spin_angle_bins( - SpiceFrame.IMAP_LO, c.N_SPIN_ANGLE_BINS - ) + spin_angles = _dps_spin_angles() # The whole pointing is projected from the middle of its good times, which # is where the despun frame is sampled. epoch = met_to_ttj2000ns((gt_start.min() + gt_end.max()) / 2.0) @@ -284,6 +281,26 @@ def _accumulate_pointing( sky_map.max_epoch = max(sky_map.max_epoch, int(met_to_ttj2000ns(gt_end.max()))) +def _dps_spin_angles() -> np.ndarray: + """ + Get the despun-frame azimuth of each histogram spin-angle bin center. + + The L1B histogram spin bins are hardware spin-phase bins referenced to the + spacecraft spin pulse, NOT the instrument (DPS) spin angle. A bin center is + converted to the IMAP_DPS azimuth by adding the spacecraft to instrument + spin-phase offset, exactly as the L1B star-sensor product does. + + Returns + ------- + np.ndarray + The IMAP_DPS azimuth [degrees] of each of the histogram spin bins. + """ + bin_width = 360.0 / c.N_SPIN_ANGLE_BINS + bin_centers = (np.arange(c.N_SPIN_ANGLE_BINS) + 0.5) * bin_width + offset = get_spacecraft_to_instrument_spin_phase_offset(SpiceFrame.IMAP_LO) * 360.0 + return np.mod(bin_centers + offset, 360.0) + + def _pixel_indices( sky_map: RectangularSkyMap, longitude: np.ndarray, latitude: np.ndarray ) -> np.ndarray: diff --git a/imap_processing/spice/spin.py b/imap_processing/spice/spin.py index 968a4841f..20264e473 100644 --- a/imap_processing/spice/spin.py +++ b/imap_processing/spice/spin.py @@ -348,41 +348,3 @@ def get_instrument_spin_phase( instrument ) return (spacecraft_spin_phase + instrument_spin_phase_offset) % 1 - - -def get_instrument_spin_angle_bins( - instrument: SpiceFrame, - n_bins: int, - deg_per_bin: float | None = None, - bin_offset: float = 0.5, -) -> npt.NDArray: - """ - Get the instrument spin angle of each hardware spin-angle bin. - - Instruments accumulate data into bins referenced to the spacecraft spin - pulse. A bin is converted to the instrument spin angle by adding the - spacecraft to instrument spin phase offset. - - Parameters - ---------- - instrument : SpiceFrame - Instrument frame the bins are wanted in. - n_bins : int - Number of spin-angle bins. - deg_per_bin : float, optional - Width [degrees] of a bin. Defaults to ``360 / n_bins``, i.e. bins that - evenly divide a spin. Bins sampled at a fixed cadence do not evenly - divide a spin and must pass their own width. - bin_offset : float - Fractional position within a bin the angle is reported at. Default is - 0.5 for the bin center; use 0.0 for the leading edge. - - Returns - ------- - spin_angle : np.ndarray - The instrument spin angle [degrees, 0-360) of each bin. - """ - if deg_per_bin is None: - deg_per_bin = 360.0 / n_bins - start_angle = 360.0 * get_spacecraft_to_instrument_spin_phase_offset(instrument) - return np.mod(start_angle + (np.arange(n_bins) + bin_offset) * deg_per_bin, 360.0) diff --git a/imap_processing/tests/lo/test_lo_l1b.py b/imap_processing/tests/lo/test_lo_l1b.py index 937b83376..c4405e323 100644 --- a/imap_processing/tests/lo/test_lo_l1b.py +++ b/imap_processing/tests/lo/test_lo_l1b.py @@ -1802,12 +1802,8 @@ def test_profiles_by_group_handles_no_valid_records(self, mock_repoint): assert len(group_epochs) == 0 # No valid records assert avg_amplitudes.shape == (0, 720) - @patch( - "imap_processing.spice.spin.get_spacecraft_to_instrument_spin_phase_offset", - return_value=350.0 / 360.0, - ) @patch("imap_processing.lo.l1b.lo_l1b.interpolate_repoint_data") - def test_profiles_by_group_angle_wrapping(self, mock_repoint, mock_offset): + def test_profiles_by_group_angle_wrapping(self, mock_repoint): """Test that spin angles wrap correctly to [0, 360) range.""" # Arrange mock_repoint.return_value = pd.DataFrame({"repoint_in_progress": [False]}) @@ -1828,6 +1824,7 @@ def test_profiles_by_group_angle_wrapping(self, mock_repoint, mock_offset): l1a_star, sampling_cadence=21.0, spin_period=15.0, + start_angle_offset=350.0, # Large offset to test wrapping ) # Assert diff --git a/imap_processing/tests/lo/test_lo_l2.py b/imap_processing/tests/lo/test_lo_l2.py index 29be956a7..f3e220df5 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -11,12 +11,11 @@ from imap_processing.lo.constants import LoConstants from imap_processing.lo.l2.lo_l2 import ( _complete_pointings, + _dps_spin_angles, _pixel_indices, _spin_phase_mask, lo_l2, ) -from imap_processing.spice.geometry import SpiceFrame -from imap_processing.spice.spin import get_instrument_spin_angle_bins from imap_processing.spice.time import met_to_ttj2000ns # A full-spin map, so that every spin-angle bin lands on it. @@ -27,9 +26,6 @@ N_SPIN_BINS = LoConstants.N_SPIN_ANGLE_BINS PIVOT = 90.0 -# The instrument spin angle of each histogram spin bin, as L2 computes it. -SPIN_ANGLES = get_instrument_spin_angle_bins(SpiceFrame.IMAP_LO, N_SPIN_BINS) - # Good-time window [MET seconds] that the "in-window" histogram epochs fall in. GT_START = 511_000_000.0 GT_END = 511_000_600.0 @@ -345,13 +341,15 @@ def test_systematic_error_bounds(self, full_map): class TestGeometry: """The spin-angle to sky-pixel geometry.""" - def test_spin_angles_carry_the_offset(self): + def test_dps_spin_angles_carry_the_offset(self): """The hardware spin bins are rotated onto the instrument frame.""" - assert SPIN_ANGLES.size == N_SPIN_BINS + angles = _dps_spin_angles() + + assert angles.size == N_SPIN_BINS # IMAP-Lo sits 60 degrees from the spacecraft spin pulse, so bin 0's # center (3 degrees) becomes 63 degrees in the despun frame. - np.testing.assert_allclose(SPIN_ANGLES[0], 63.0) - np.testing.assert_allclose(np.diff(np.sort(SPIN_ANGLES)), 6.0) + np.testing.assert_allclose(angles[0], 63.0) + np.testing.assert_allclose(np.diff(np.sort(angles)), 6.0) def test_pixel_indices_match_the_map_grid(self): """Directions are placed in the pixel whose center they are nearest.""" @@ -377,14 +375,16 @@ def test_pixel_indices_wrap_longitude(self): ) def test_spin_phase_mask(self, spin_phase, expected): """Ram and anti-ram split the spin; a full map keeps all of it.""" - mask = _spin_phase_mask(SPIN_ANGLES, PIVOT, spin_phase) + angles = _dps_spin_angles() + + mask = _spin_phase_mask(angles, PIVOT, spin_phase) assert mask.sum() == expected def test_spin_phase_mask_rejects_unknown(self): """An unknown spin phase is an error, not a silently empty map.""" with pytest.raises(ValueError, match="Invalid spin phase"): - _spin_phase_mask(SPIN_ANGLES, PIVOT, "sideways") + _spin_phase_mask(_dps_spin_angles(), PIVOT, "sideways") class TestPointingSelection: diff --git a/imap_processing/tests/spice/test_spin.py b/imap_processing/tests/spice/test_spin.py index 9aa17d6d7..808e093dc 100644 --- a/imap_processing/tests/spice/test_spin.py +++ b/imap_processing/tests/spice/test_spin.py @@ -312,29 +312,3 @@ def test_get_spin_start_met(fake_spin_data): # when we have a linear ramp of data points between spins start_mets = spin.interpolate_spin_data(np.linspace(32, 37, 3))["spin_start_met"] np.testing.assert_array_equal(start_mets, [30, 30, 30]) - - -def test_get_instrument_spin_angle_bins(): - """Bins that evenly divide a spin are reported at their centers.""" - spin_angles = spin.get_instrument_spin_angle_bins(SpiceFrame.IMAP_LO, 60) - - assert spin_angles.shape == (60,) - # IMAP-Lo sits 60 degrees from the spacecraft spin pulse, so the center of - # bin 0 (3 degrees) is at 63 degrees in the instrument frame. - np.testing.assert_allclose(spin_angles[0], 63.0) - np.testing.assert_allclose(np.diff(np.sort(spin_angles)), 6.0) - - -def test_get_instrument_spin_angle_bins_uneven_bins(): - """Bins sampled at a fixed cadence wrap around the spin.""" - # 720 samples of 21 ms each cover more than one 15 second spin. - deg_per_bin = 360.0 * 0.021 / 15.0 - spin_angles = spin.get_instrument_spin_angle_bins( - SpiceFrame.IMAP_LO, 720, deg_per_bin=deg_per_bin, bin_offset=0.0 - ) - - assert np.all((spin_angles >= 0) & (spin_angles < 360)) - # The first sample sits at the instrument offset, and the bins wrap past it. - np.testing.assert_allclose(spin_angles[0], 60.0) - np.testing.assert_allclose(spin_angles[1], 60.0 + deg_per_bin) - assert np.any(spin_angles < 60.0) From d357ce8bcb0d1b5382f2038ec17c2c66a0973ecf Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Tue, 28 Jul 2026 15:32:45 -0400 Subject: [PATCH 08/10] pre-commit fixes --- imap_processing/lo/l2/lo_l2.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index de3ba04c9..1abdc30d3 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -256,7 +256,7 @@ def _accumulate_pointing( pivot_angle, spin_angles=spin_angles, off_angles=np.array([0.0]), - to_frame=cast(SpiceFrame, map_descriptor.map_spice_coord_frame), + to_frame=map_descriptor.map_spice_coord_frame, ) # The single boresight off-angle is squeezed out of the frame transform, so # the directions come back as (spin angle, lon/lat). @@ -325,7 +325,9 @@ def _pixel_indices( The pixel index of each direction. """ spacing = sky_map.spacing_deg - num_azimuth, num_elevation = sky_map.binning_grid_shape + # binning_grid_shape is annotated as a 1-tuple in the base class, but a + # rectangular map always returns (num azimuth bins, num elevation bins). + num_azimuth, num_elevation = cast(tuple[int, int], sky_map.binning_grid_shape) azimuth_index = np.clip( (np.mod(longitude, 360.0) // spacing).astype(int), 0, num_azimuth - 1 From ed24bc01a15616dc616aed8ab62cc912539f76ab Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Tue, 28 Jul 2026 17:03:51 -0400 Subject: [PATCH 09/10] using an in-memory pset to reuse existing projection code --- imap_processing/lo/l2/lo_l2.py | 283 +++++++++++++++---------- imap_processing/tests/lo/test_lo_l2.py | 49 ++++- 2 files changed, 209 insertions(+), 123 deletions(-) diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index 1abdc30d3..c801a3c6a 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -1,12 +1,16 @@ """IMAP-Lo L2 data processing.""" import logging -from typing import cast import numpy as np import xarray as xr -from imap_processing.ena_maps.ena_maps import RectangularSkyMap +from imap_processing.ena_maps.ena_maps import ( + PointingSet, + RectangularSkyMap, + SkyTilingType, +) +from imap_processing.ena_maps.utils.coordinates import CoordNames from imap_processing.ena_maps.utils.naming import MapDescriptor from imap_processing.lo.constants import LoConstants as c # noqa: N813 from imap_processing.lo.l1c.lo_l1c import compute_pointing_directions @@ -14,13 +18,21 @@ SpiceFrame, get_spacecraft_to_instrument_spin_phase_offset, ) -from imap_processing.spice.time import met_to_ttj2000ns, ttj2000ns_to_met +from imap_processing.spice.time import ( + met_to_ttj2000ns, + ttj2000ns_to_et, + ttj2000ns_to_met, +) logger = logging.getLogger(__name__) # The descriptors of the L1B products a map is built from, one set per pointing. REQUIRED_PRODUCTS = ("goodtimes", "bgrates", "histrates") +# The map variables accumulated directly from the pointings, before any rates +# or intensities are derived from them. +ACCUMULATED_VARIABLES = ("ena_count", "exposure_factor", "bg_rate_exposure") + # ============================================================================= # MAIN ENTRY POINT # ============================================================================= @@ -86,31 +98,15 @@ def lo_l2( pointings = _complete_pointings(sci_dependencies) logger.info(f"Building {descriptor} from {len(pointings)} pointings") - shape = (c.N_ESA_LEVELS, sky_map.num_points) - counts = np.zeros(shape) - exposure = np.zeros(shape) - # Background is a rate per ESA level per pointing, so it is accumulated - # weighted by exposure and divided by the total exposure at the end. - bg_rate_exposure = np.zeros(shape) + _initialize_accumulators(sky_map) esa_mode = 0 for repointing, (goodtimes, bgrates, histrates) in sorted(pointings.items()): logger.debug(f"Accumulating repoint{repointing:05d}") esa_mode = _get_esa_mode(histrates) - _accumulate_pointing( - goodtimes, - bgrates, - histrates, - sky_map, - map_descriptor, - counts, - exposure, - bg_rate_exposure, - ) + _accumulate_pointing(goodtimes, bgrates, histrates, sky_map, map_descriptor) - variables = _calculate_rates_and_intensities( - counts, exposure, bg_rate_exposure, esa_mode - ) + variables = _calculate_rates_and_intensities(sky_map, esa_mode) dataset = _build_map_dataset(sky_map, variables, esa_mode) logger.info("IMAP-Lo L2 processing pipeline completed successfully") @@ -170,6 +166,18 @@ def _complete_pointings( return pointings +def _esa_energy() -> np.ndarray: + """ + Get the energy of each ESA level the map is binned in. + + Returns + ------- + np.ndarray + The energy [keV] of each ESA level. + """ + return np.array(c.ESA_ENERGY[: c.N_ESA_LEVELS]) + + def _get_esa_mode(histrates: xr.Dataset) -> int: """ Read the ESA mode of a pointing, defaulting to HiRes. @@ -194,18 +202,117 @@ def _get_esa_mode(histrates: xr.Dataset) -> int: # ============================================================================= +class LoSpinAnglePointingSet(PointingSet): # type: ignore[misc] + """ + The spin-angle bins of one pointing, as an in-memory pointing set. + + Lo builds its maps straight from the L1B products of a pointing rather than + from a written L1C pointing set, so the sky direction of each spin-angle + bin and the values looking in it are assembled here. + + Parameters + ---------- + epoch : int + The time [TTJ2000 ns] the pointing is projected from. + pivot_angle : float + The pivot angle [degrees] of the pointing. + spin_angles : np.ndarray + The IMAP_DPS azimuth [degrees] of each spin-angle bin. + values : dict[str, np.ndarray] + The values of the pointing, each of shape (esa level, spin angle). + frame : SpiceFrame + The frame to compute the sky directions in, i.e. the map's frame. + """ + + tiling_type: SkyTilingType = SkyTilingType.RECTANGULAR + + def __init__( + self, + epoch: int, + pivot_angle: float, + spin_angles: np.ndarray, + values: dict[str, np.ndarray], + frame: SpiceFrame, + ): + dims = [CoordNames.TIME.value, CoordNames.ENERGY_L2.value, "spin_angle"] + super().__init__( + xr.Dataset( + { + name: (dims, value[np.newaxis, ...]) # add epoch axis + for name, value in values.items() + }, + coords={ + CoordNames.TIME.value: [epoch], + CoordNames.ENERGY_L2.value: _esa_energy(), + }, + ), + spice_reference_frame=frame, + ) + self.spatial_coords = ("spin_angle",) + + az_el = compute_pointing_directions( + epoch, + pivot_angle, + spin_angles=spin_angles, + off_angles=np.array([0.0]), + to_frame=frame, + ) + self.az_el_points = xr.DataArray( + np.asarray(az_el), + dims=[CoordNames.GENERIC_PIXEL.value, CoordNames.AZ_EL_VECTOR.value], + ) + + @property + def midpoint_j2000_et(self) -> float: + """ + The time the pointing is projected from. + + The base class derives this from an ``epoch_delta``; a pointing built + here is handed the single epoch it is projected from directly. + + Returns + ------- + float + The epoch of the pointing set [J2000 ET]. + """ + return float(ttj2000ns_to_et(self.epoch)) + + +def _initialize_accumulators(sky_map: RectangularSkyMap) -> None: + """ + Seed the map with the empty accumulators each pointing is added into. + + ``project_pset_values_to_map`` creates a map variable the first time it + projects one, so seeding them is what lets the rest of the pipeline read + the accumulators unconditionally, however many pointings turn out to be + usable. + + Parameters + ---------- + sky_map : RectangularSkyMap + The map being built, modified in place. + """ + for name in ACCUMULATED_VARIABLES: + sky_map.data_1d[name] = xr.DataArray( + np.zeros((1, c.N_ESA_LEVELS, sky_map.num_points)), + dims=[ + CoordNames.TIME.value, + CoordNames.ENERGY_L2.value, + CoordNames.GENERIC_PIXEL.value, + ], + coords={CoordNames.ENERGY_L2.value: _esa_energy()}, + ) + + def _accumulate_pointing( goodtimes: xr.Dataset, bgrates: xr.Dataset, histrates: xr.Dataset, sky_map: RectangularSkyMap, map_descriptor: MapDescriptor, - counts: np.ndarray, - exposure: np.ndarray, - bg_rate_exposure: np.ndarray, ) -> None: """ - Add one pointing's counts and exposure to the map accumulators. + Add one pointing's counts and exposure to the map. Parameters ---------- @@ -218,15 +325,9 @@ def _accumulate_pointing( The L1B histogram rates of the pointing, giving the counts and exposure of each spin-angle bin. sky_map : RectangularSkyMap - The map being built, used for its pixel grid. + The map being built, modified in place. map_descriptor : MapDescriptor The parsed descriptor of the map being made. - counts : np.ndarray - Accumulator of shape (esa level, pixel), modified in place. - exposure : np.ndarray - Accumulator of shape (esa level, pixel), modified in place. - bg_rate_exposure : np.ndarray - Accumulator of shape (esa level, pixel), modified in place. """ species = map_descriptor.species pivot_angle = float(np.atleast_1d(goodtimes["pivot"].values)[0]) @@ -243,38 +344,38 @@ def _accumulate_pointing( logger.warning("No histogram epochs fall within the good-time windows.") return + spin_angles = _dps_spin_angles() + keep = _spin_phase_mask(spin_angles, pivot_angle, map_descriptor.spin_phase) + if not keep.any(): + return + pointing_counts = histrates[f"{species}_counts"].values[in_goodtime].sum(axis=0) pointing_exposure = histrates["exposure_time_6deg"].values[in_goodtime].sum(axis=0) background_rates = np.atleast_2d(bgrates[f"{species}_background_rates"].values)[0] - spin_angles = _dps_spin_angles() # The whole pointing is projected from the middle of its good times, which # is where the despun frame is sampled. epoch = met_to_ttj2000ns((gt_start.min() + gt_end.max()) / 2.0) - az_el = compute_pointing_directions( + pointing_set = LoSpinAnglePointingSet( epoch, pivot_angle, - spin_angles=spin_angles, - off_angles=np.array([0.0]), - to_frame=map_descriptor.map_spice_coord_frame, + spin_angles, + { + "ena_count": pointing_counts, + "exposure_factor": pointing_exposure, + # Background is a rate per ESA level per pointing, so it is + # accumulated weighted by exposure and divided by the total + # exposure at the end. + "bg_rate_exposure": background_rates[:, np.newaxis] * pointing_exposure, + }, + sky_map.spice_reference_frame, ) - # The single boresight off-angle is squeezed out of the frame transform, so - # the directions come back as (spin angle, lon/lat). - az_el = np.asarray(az_el).reshape(spin_angles.size, 2) - pixels = _pixel_indices(sky_map, az_el[:, 0], az_el[:, 1]) - - keep = _spin_phase_mask(spin_angles, pivot_angle, map_descriptor.spin_phase) - if not keep.any(): - return - - # np.add.at accumulates repeated pixels, which is what happens whenever - # several spin-angle bins land in the same map pixel. - np.add.at(counts, (slice(None), pixels[keep]), pointing_counts[:, keep]) - np.add.at(exposure, (slice(None), pixels[keep]), pointing_exposure[:, keep]) - np.add.at( - bg_rate_exposure, - (slice(None), pixels[keep]), - background_rates[:, np.newaxis] * pointing_exposure[:, keep], + # The projection sums the spin-angle bins that land in the same map pixel, + # and adds this pointing on top of what the earlier pointings left there. + sky_map.project_pset_values_to_map( + pointing_set, + value_keys=list(ACCUMULATED_VARIABLES), + pset_valid_mask=keep, ) sky_map.min_epoch = min(sky_map.min_epoch, int(met_to_ttj2000ns(gt_start.min()))) @@ -301,43 +402,6 @@ def _dps_spin_angles() -> np.ndarray: return np.mod(bin_centers + offset, 360.0) -def _pixel_indices( - sky_map: RectangularSkyMap, longitude: np.ndarray, latitude: np.ndarray -) -> np.ndarray: - """ - Get the map pixel each sky direction falls in. - - A rectangular map stores its pixels as a 1D array raveled from the - (azimuth, elevation) grid, elevation varying fastest. - - Parameters - ---------- - sky_map : RectangularSkyMap - The map being built. - longitude : np.ndarray - Longitudes [degrees] in the map's frame. - latitude : np.ndarray - Latitudes [degrees] in the map's frame. - - Returns - ------- - np.ndarray - The pixel index of each direction. - """ - spacing = sky_map.spacing_deg - # binning_grid_shape is annotated as a 1-tuple in the base class, but a - # rectangular map always returns (num azimuth bins, num elevation bins). - num_azimuth, num_elevation = cast(tuple[int, int], sky_map.binning_grid_shape) - - azimuth_index = np.clip( - (np.mod(longitude, 360.0) // spacing).astype(int), 0, num_azimuth - 1 - ) - elevation_index = np.clip( - ((latitude + 90.0) // spacing).astype(int), 0, num_elevation - 1 - ) - return azimuth_index * num_elevation + elevation_index - - def _spin_phase_mask( spin_angles: np.ndarray, pivot_angle: float, spin_phase: str ) -> np.ndarray: @@ -411,11 +475,8 @@ def _geometric_factors(esa_mode: int) -> tuple[np.ndarray, np.ndarray, np.ndarra def _calculate_rates_and_intensities( - counts: np.ndarray, - exposure: np.ndarray, - bg_rate_exposure: np.ndarray, - esa_mode: int, -) -> xr.Dataset: + sky_map: RectangularSkyMap, esa_mode: int +) -> dict[str, np.ndarray]: """ Turn the accumulated counts and exposure into rates and intensities. @@ -423,21 +484,21 @@ def _calculate_rates_and_intensities( Parameters ---------- - counts : np.ndarray - Accumulated counts of shape (esa level, pixel). - exposure : np.ndarray - Accumulated exposure time [s] of shape (esa level, pixel). - bg_rate_exposure : np.ndarray - Accumulated exposure-weighted background rate, same shape. + sky_map : RectangularSkyMap + The map the pointings were projected onto, read for its accumulators. esa_mode : int The ESA mode, 0 for HiRes and 1 for HiThr. Returns ------- dict[str, np.ndarray] - The map variables, each of shape (esa level, pixel). + The map variables, each of shape (epoch, esa level, pixel). """ - energy = np.array(c.ESA_ENERGY[: c.N_ESA_LEVELS])[:, np.newaxis] + counts = sky_map.data_1d["ena_count"].values + exposure = sky_map.data_1d["exposure_factor"].values + bg_rate_exposure = sky_map.data_1d["bg_rate_exposure"].values + + energy = _esa_energy()[:, np.newaxis] geometric_factor, error_upper, error_lower = _geometric_factors(esa_mode) geometric_factor = geometric_factor[:, np.newaxis] error_upper = error_upper[:, np.newaxis] @@ -535,7 +596,7 @@ def _build_map_dataset( sky_map : RectangularSkyMap The map being built. variables : dict[str, np.ndarray] - The map variables, each of shape (esa level, pixel). + The map variables, each of shape (epoch, esa level, pixel). esa_mode : int The ESA mode, 0 for HiRes and 1 for HiThr, which sets the widths of the ESA energy passbands. @@ -546,13 +607,11 @@ def _build_map_dataset( The map variables on the (epoch, energy, longitude, latitude) grid, with the energy coordinate and its widths. """ - energy = np.array(c.ESA_ENERGY[: c.N_ESA_LEVELS]) + dims = sky_map.data_1d["ena_count"].dims for name, values in variables.items(): - sky_map.data_1d[name] = xr.DataArray( - values[np.newaxis, ...].astype(np.float32), - dims=["epoch", "energy", "pixel"], - coords={"energy": energy}, - ) + sky_map.data_1d[name] = xr.DataArray(values.astype(np.float32), dims=dims) + # `bg_rate_exposure` is an accumulator, not a map variable. + sky_map.data_1d = sky_map.data_1d.drop_vars("bg_rate_exposure") dataset = sky_map.to_dataset() diff --git a/imap_processing/tests/lo/test_lo_l2.py b/imap_processing/tests/lo/test_lo_l2.py index f3e220df5..d8ed50d7f 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -7,12 +7,13 @@ import xarray as xr from imap_processing.cdf.utils import load_cdf, write_cdf +from imap_processing.ena_maps.ena_maps import match_coords_to_indices from imap_processing.ena_maps.utils.naming import MapDescriptor from imap_processing.lo.constants import LoConstants from imap_processing.lo.l2.lo_l2 import ( + LoSpinAnglePointingSet, _complete_pointings, _dps_spin_angles, - _pixel_indices, _spin_phase_mask, lo_l2, ) @@ -119,6 +120,25 @@ def identity_pointing(et, az_el, *args, **kwargs): return np.asarray(az_el)[:, 0, :] +def make_pointing_set(sky_map, spin_angles, pivot=PIVOT): + """Build the in-memory pointing set of one pointing, sky pointing mocked.""" + values = { + name: np.ones((N_ESA, spin_angles.size)) + for name in ("ena_count", "exposure_factor", "bg_rate_exposure") + } + with patch( + "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el", + side_effect=identity_pointing, + ): + return LoSpinAnglePointingSet( + met_to_ttj2000ns(GT_START), + pivot, + spin_angles, + values, + sky_map.spice_reference_frame, + ) + + @pytest.fixture def one_pointing(): """A single synthetic pointing.""" @@ -351,23 +371,30 @@ def test_dps_spin_angles_carry_the_offset(self): np.testing.assert_allclose(angles[0], 63.0) np.testing.assert_allclose(np.diff(np.sort(angles)), 6.0) - def test_pixel_indices_match_the_map_grid(self): - """Directions are placed in the pixel whose center they are nearest.""" + def test_bins_land_in_the_pixel_they_are_nearest(self): + """Each spin-angle bin is projected into the pixel it points into.""" sky_map = MapDescriptor.from_string(FULL_DESCRIPTOR).to_empty_map() - centers = sky_map.az_el_points.values + spin_angles = _dps_spin_angles() + pointing_set = make_pointing_set(sky_map, spin_angles) - pixels = _pixel_indices(sky_map, centers[:, 0], centers[:, 1]) + pixels = match_coords_to_indices(pointing_set, sky_map).values - np.testing.assert_array_equal(pixels, np.arange(sky_map.num_points)) + # A direction never lands further than half a pixel from its pixel's + # center, in either axis. + centers = sky_map.az_el_points.values[pixels] + directions = pointing_set.az_el_points.values + half_pixel = sky_map.spacing_deg / 2 + assert np.all(np.abs(centers[:, 0] - directions[:, 0]) <= half_pixel) + assert np.all(np.abs(centers[:, 1] - directions[:, 1]) <= half_pixel) - def test_pixel_indices_wrap_longitude(self): - """Longitudes outside 0-360 wrap onto the map.""" + def test_bins_of_a_spin_land_in_distinct_pixels(self): + """The 60 six-degree spin bins fill a row of the six-degree map.""" sky_map = MapDescriptor.from_string(FULL_DESCRIPTOR).to_empty_map() + pointing_set = make_pointing_set(sky_map, _dps_spin_angles()) - wrapped = _pixel_indices(sky_map, np.array([-357.0]), np.array([0.0])) - direct = _pixel_indices(sky_map, np.array([3.0]), np.array([0.0])) + pixels = match_coords_to_indices(pointing_set, sky_map).values - np.testing.assert_array_equal(wrapped, direct) + assert len(np.unique(pixels)) == N_SPIN_BINS @pytest.mark.parametrize( "spin_phase, expected", From 5bd190e283fdf9bc4223915448de1e67251c8026 Mon Sep 17 00:00:00 2001 From: Vineet Bansal Date: Tue, 28 Jul 2026 17:27:02 -0400 Subject: [PATCH 10/10] pre-commit fix --- imap_processing/_version.py | 4 ++-- imap_processing/lo/l2/lo_l2.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/imap_processing/_version.py b/imap_processing/_version.py index e3e53154f..31a9c8e8e 100644 --- a/imap_processing/_version.py +++ b/imap_processing/_version.py @@ -1,3 +1,3 @@ # These version placeholders will be replaced later during substitution. -__version__ = "0.0.0" -__version_tuple__ = (0, 0, 0) +__version__ = "1.0.35.post11.dev0+ed24bc01" +__version_tuple__ = (1, 0, 35, "post11", "dev0", "ed24bc01") diff --git a/imap_processing/lo/l2/lo_l2.py b/imap_processing/lo/l2/lo_l2.py index c801a3c6a..4821350a3 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -202,7 +202,7 @@ def _get_esa_mode(histrates: xr.Dataset) -> int: # ============================================================================= -class LoSpinAnglePointingSet(PointingSet): # type: ignore[misc] +class LoSpinAnglePointingSet(PointingSet): """ The spin-angle bins of one pointing, as an in-memory pointing set.