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/cli.py b/imap_processing/cli.py index 324a46378..81ffd5fb6 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 @@ -1212,52 +1213,109 @@ def do_processing( class Lo(ProcessInstrument): """Process IMAP-Lo.""" + @staticmethod + def _pointings_at_pivot_angle( + dependencies: ProcessingInputCollection, map_pivot_angle: int + ) -> set[int]: + """ + Find the pointings that were taken at the pivot angle of a map. + + 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 + ---------- + dependencies : ProcessingInputCollection + Object containing dependencies to process. + map_pivot_angle : int + The pivot angle [degrees] of the map being made. + + Returns + ------- + set[int] + The repointings whose pivot angle is that of the map. + """ + 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 = 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 at_pivot_angle + 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. + Extends the base pre-processing by dropping, for map products, the Lo + 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 ------- dependencies : ProcessingInputCollection Object containing dependencies to process. """ - datasets = super().pre_processing() - new_datasets = ProcessingInputCollection() + 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 - for processing_input in datasets.get_processing_inputs(): + if not isinstance(map_pivot_angle, int): + # 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 ( - processing_input.source == "lo" - and processing_input.descriptor == "pset" + processing_input.input_type != ProcessingInputType.SCIENCE_FILE + or processing_input.source != "lo" ): - 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) + filtered_dependencies.add(processing_input) + continue - return new_datasets + kept_filenames = [ + str(imap_file_path.filename) + for imap_file_path in processing_input.imap_file_paths + if imap_file_path.repointing in at_pivot_angle + ] + if kept_filenames: + filtered_dependencies.add(type(processing_input)(*kept_filenames)) + + return filtered_dependencies def do_processing( self, dependencies: ProcessingInputCollection @@ -1323,18 +1381,26 @@ def do_processing( datasets = lo_l1c.lo_l1c(data_dict, anc_dependencies) elif self.data_level == "l2": - data_dict = {} - science_files = dependencies.get_file_paths(source="lo", descriptor="pset") 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 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 + ) - 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) return datasets diff --git a/imap_processing/lo/constants.py b/imap_processing/lo/constants.py index aa0e6e4a6..40de8b1f9 100644 --- a/imap_processing/lo/constants.py +++ b/imap_processing/lo/constants.py @@ -1,18 +1,52 @@ """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) 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 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. @@ -45,6 +79,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 @@ -53,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..ad10569de 100644 --- a/imap_processing/lo/l1b/lo_l1b.py +++ b/imap_processing/lo/l1b/lo_l1b.py @@ -2521,10 +2521,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 50021a0bb..4821350a3 100644 --- a/imap_processing/lo/l2/lo_l2.py +++ b/imap_processing/lo/l2/lo_l2.py @@ -1,1455 +1,622 @@ """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 ( + 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 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_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 # ============================================================================= 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 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 + ``cli.Lo.pre_processing``) so that the map records only the files it was + made from as its parents. 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. + 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 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 ------ - ValueError - If no pointing set data found in science dependencies. 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") - 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}") - # 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) - - 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 - ) - - 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, - ) - - 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 - ) + # 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("IMAP-Lo L2 processing pipeline completed successfully") - return [dataset] + sky_map = map_descriptor.to_empty_map() + if not isinstance(sky_map, RectangularSkyMap): + raise NotImplementedError("HEALPix map output not supported for Lo") + pointings = _complete_pointings(sci_dependencies) + logger.info(f"Building {descriptor} from {len(pointings)} pointings") -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. + _initialize_accumulators(sky_map) + esa_mode = 0 - 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. + 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) - 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. + variables = _calculate_rates_and_intensities(sky_map, esa_mode) + dataset = _build_map_dataset(sky_map, variables, esa_mode) - 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, - ) + logger.info("IMAP-Lo L2 processing pipeline completed successfully") + 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. - """ - efficiency_files = [ - anc_file - for anc_file in anc_dependencies - if "efficiency-factor" in str(anc_file) - ] - - 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: +def _complete_pointings( + sci_dependencies: dict[int, dict[str, xr.Dataset]], +) -> dict[int, tuple]: """ - Load sputter correction factors from an ancillary file. + Reduce the grouped inputs to the pointings that can be mapped. 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). + sci_dependencies : dict[int, dict[str, xr.Dataset]] + The input datasets of each pointing, keyed by repointing and then by + product descriptor. 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*")) + dict[int, tuple] + The (goodtimes, bgrates, histrates) datasets of each mappable pointing, + keyed by repointing. - if not sputter_files: - raise ValueError("No sputter correction files found") + Raises + ------ + KeyError + If any of the three required products is missing entirely. + """ + 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 sci_dependencies.items(): + missing = set(REQUIRED_PRODUCTS) - set(products) + if missing: + logger.warning( + f"Dropping repoint{repointing:05d}, it has no {sorted(missing)}" + ) + continue + pointings[repointing] = tuple(products[p] for p in REQUIRED_PRODUCTS) - 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 + return pointings -def load_bootstrap_correction_data() -> pd.DataFrame: +def _esa_energy() -> np.ndarray: """ - Load bootstrap correction factors from an ancillary file. + Get the energy of each ESA level the map is binned in. 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. + np.ndarray + The energy [keV] of each ESA level. """ - anc_path = Path(__file__).parent.parent / "ancillary_data" - bootstrap_files = sorted(anc_path.glob("*bootstrap-correction-factors*")) - - if not bootstrap_files: - raise ValueError("No bootstrap correction factor files found") - - return pd.concat( - [lo_ancillary.read_ancillary_file(f) for f in bootstrap_files], - ignore_index=True, - ) + return np.array(c.ESA_ENERGY[: c.N_ESA_LEVELS]) -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], - map_descriptor: MapDescriptor, - efficiency_data: pd.DataFrame, - cg_correct: bool, -) -> AbstractSkyMap: - """ - Create a sky map by processing all pointing sets. - - Parameters - ---------- - psets : list[xr.Dataset] - List of pointing set datasets to process. - 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. - - 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). +class LoSpinAnglePointingSet(PointingSet): """ - # 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) + The spin-angle bins of one pointing, as an in-memory pointing set. - 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. + 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 ---------- - 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. - """ - # 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 + 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, ) - # Prepare energy_sc for exposure time weighted projection - pset_processed["energy_sc_exposure_factor"] = ( - pset_processed["energy_sc"] * pset_processed["exposure_factor"] + 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], ) - # Always calculate ram-mask to identify ram/anti-ram bins - pset_processed = calculate_ram_mask(pset_processed) - - 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. - """ - # 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 - + @property + def midpoint_j2000_et(self) -> float: + """ + The time the pointing is projected from. -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. + The base class derives this from an ``epoch_delta``; a pointing built + here is handed the single epoch it is projected from directly. - Parameters - ---------- - pset : xr.Dataset - Pointing set dataset to add efficiency factors to. - efficiency_data : pd.DataFrame - Efficiency factor data containing date-indexed efficiency values. + Returns + ------- + float + The epoch of the pointing set [J2000 ET]. + """ + return float(ttj2000ns_to_et(self.epoch)) - 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. +def _initialize_accumulators(sky_map: RectangularSkyMap) -> None: """ - 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 + Seed the map with the empty accumulators each pointing is added into. - -def calculate_efficiency_corrected_quantities( - pset: xr.Dataset, -) -> xr.Dataset: - """ - Calculate efficiency-corrected quantities for each particle type. + ``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 ---------- - 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"] - ) - - return pset + 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 project_pset_to_map( - pset: xr.Dataset, - output_map: AbstractSkyMap, - directional_mask: xr.DataArray, - cg_correct: bool = False, +def _accumulate_pointing( + goodtimes: xr.Dataset, + bgrates: xr.Dataset, + histrates: xr.Dataset, + sky_map: RectangularSkyMap, + map_descriptor: MapDescriptor, ) -> None: """ - Project pointing set data to the output map. + Add one pointing's counts and exposure to the 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. + 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, modified in place. + map_descriptor : MapDescriptor + The parsed descriptor of the map being made. + """ + 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 + + 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] + + # 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) + pointing_set = LoSpinAnglePointingSet( + epoch, + pivot_angle, + 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 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()))) + 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 ------- - None - Function modifies output_map in place. + np.ndarray + The IMAP_DPS azimuth [degrees] of each of the histogram spin bins. """ - # 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, - ) - logger.debug(f"Projected {len(value_keys)} quantities to sky map") + 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) -# ============================================================================= -# GEOMETRIC FACTORS -# ============================================================================= - - -def add_geometric_factors(dataset: xr.Dataset, species: str) -> xr.Dataset: - """ - Add geometric factors to the sky map after projection. - - 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) + Get the spin-angle bins belonging on a map of the given spin phase. - # Populate geometric factors for each energy step - dataset = populate_geometric_factors(dataset, species) - - return dataset - - -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}'" - ) - - 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 + f"Invalid spin phase: {spin_phase}. Must be 'ram', 'anti' or 'full'." ) - # 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: - """ - Calculate particle intensities and uncertainties for the specified species. - - 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: +def _geometric_factors(esa_mode: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ - Calculate background rates and intensities 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 background rates and intensities for the - specified species. - """ - # 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"] - ) - - dataset["bg_intensity_sys_err_plus"] = ( - dataset["bg_intensity"] * plus_multiplier - ) - dataset["bg_intensity"] - - 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 - - -def calculate_sputtering_corrections( - dataset: xr.Dataset, o_dataset: xr.Dataset -) -> 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. + 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 sputtering-corrected intensities and their - uncertainties. + tuple[np.ndarray, np.ndarray, np.ndarray] + The geometric factor of each ESA level, and its upper and lower error + bounds. """ - 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"], - ) + 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 - # 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, - ) + 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) - # 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"], - ) + return geometric_factor, error_upper, error_lower - # 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: +def _calculate_rates_and_intensities( + sky_map: RectangularSkyMap, esa_mode: int +) -> dict[str, np.ndarray]: """ - 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. + 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 ------- - 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 (epoch, esa level, pixel). + """ + 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] + 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 (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. 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"], - ) - ) - - return dataset - - -def cleanup_intermediate_variables(dataset: xr.Dataset) -> xr.Dataset: + The map variables on the (epoch, energy, longitude, latitude) grid, + with the energy coordinate and its widths. """ - Remove intermediate variables that were only needed for calculations. - - Parameters - ---------- - dataset : xr.Dataset - Dataset containing intermediate calculation variables. + dims = sky_map.data_1d["ena_count"].dims + for name, values in variables.items(): + 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") - 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_l1b.py b/imap_processing/tests/lo/test_lo_l1b.py index 4291dbdd8..c4405e323 100644 --- a/imap_processing/tests/lo/test_lo_l1b.py +++ b/imap_processing/tests/lo/test_lo_l1b.py @@ -3023,7 +3023,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 c053ff8db..d8ed50d7f 100644 --- a/imap_processing/tests/lo/test_lo_l2.py +++ b/imap_processing/tests/lo/test_lo_l2.py @@ -1,2915 +1,473 @@ -"""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.ena_maps import match_coords_to_indices 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, + LoSpinAnglePointingSet, + _complete_pointings, + _dps_spin_angles, + _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, ) +from imap_processing.spice.time import met_to_ttj2000ns -# ============================================================================= -# 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) +# 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" - # 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 +N_ESA = LoConstants.N_ESA_LEVELS +N_SPIN_BINS = LoConstants.N_SPIN_ANGLE_BINS +PIVOT = 90.0 - # 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 +# 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] - dataset = 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), - }, - coords={ - "epoch": [8.1794907049e17], - "esa_energy_step": ESA_ENERGY_STEPS, - "spin_angle": SPIN_ANGLE_BIN_CENTERS, - "off_angle": OFF_ANGLE_BIN_CENTERS, - }, - ) - return dataset +def product_attrs(repointing, product): + """The global attributes an L1B input of a pointing is written with.""" -@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), + return { + "Repointing": f"repoint{repointing:05d}", + "Logical_source": f"imap_lo_l1b_{product}", } - # 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, - }, - ) - return dataset +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 + 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) -@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 + 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 - dataset = xr.Dataset( + 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), + "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": [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" + coords={"epoch": met_to_ttj2000ns(mets)}, + attrs=product_attrs(repointing, "histrates"), ) - 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), - } - - # 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, + goodtimes = xr.Dataset( + { + "pivot": ("epoch", [pivot]), + "gt_start_met": ("epoch", [GT_START]), + "gt_end_met": ("epoch", [GT_END]), }, + coords={"epoch": [0]}, + attrs=product_attrs(repointing, "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=product_attrs(repointing, "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 { + "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, } - # 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 +def as_dependencies(*pointings, products=("goodtimes", "bgrates", "histrates")): + """Turn pointings into the sci_dependencies lo_l2 takes. - -@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), + The CLI keys the inputs by repointing and then by product descriptor, see + ``cli.Lo.do_processing``. + """ + return { + pointing["repointing"]: {product: pointing[product] for product in products} + for pointing in pointings } - # 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"]) - ) +def identity_pointing(et, az_el, *args, **kwargs): + """Stand in for the SPICE DPS transform: spin angle -> lon, off -> lat. - # 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, - ) - - return dataset - - -@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" + Like the real ``frame_transform_az_el``, the singleton off-angle dimension + is squeezed out. + """ + return np.asarray(az_el)[:, 0, :] -@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), +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") } - - # 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, - ): - """Test functionality of reduce_geometric_factor_dataset with real gf data.""" - - result = reduce_geometric_factor_dataset(species, esa_mode) - - # Verify it returns an xarray Dataset - assert isinstance(result, xr.Dataset) - - # Verify it has 7 energy steps - assert len(result["Observed_E-Step"]) == 7 - - # Verify the index is 1-7 - expected_indices = list(range(1, 8)) - np.testing.assert_array_equal( - result["Observed_E-Step"].values, expected_indices - ) - - # Verify that incident_E-Step == Observed_E-Step - np.testing.assert_array_equal( - result["incident_E-Step"].values, result["Observed_E-Step"].values - ) - - @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 - ) - - # 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 - - # Verify no duplicate indices - assert len(np.unique(result["Observed_E-Step"].values)) == 7 - - @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) - - @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) - - # Check that dimensions were renamed - assert "energy" in result.dims - assert "esa_energy_step" not in result.dims - - # 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, - }, - ) - - # 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, - }, - ) - - result = normalize_pset_coordinates(pset, species) - - # 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 - - -class TestAddEfficiencyFactorsToPset: - """Tests for the add_efficiency_factors_to_pset function.""" - - 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, - ): - # 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)) - - def test_add_efficiency_factors_missing_date( - self, minimal_pset, sample_efficiency_data + with patch( + "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el", + side_effect=identity_pointing, ): - """Test error when efficiency factor not found for date.""" - pset = minimal_pset.copy() - - 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, - ): - # 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))}, - ) - - result = calculate_efficiency_corrected_quantities(pset) - - # 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 - - # 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) - - # 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 - ) - - # 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 - ) - - expected_bg_uncert_exposure = ( - pset["bg_rate_stat_uncert"] ** 2 * pset["exposure_factor"] ** 2 + return LoSpinAnglePointingSet( + met_to_ttj2000ns(GT_START), + pivot, + spin_angles, + values, + sky_map.spice_reference_frame, ) - xr.testing.assert_allclose( - result["bg_rate_stat_uncert_exposure_factor2"], expected_bg_uncert_exposure - ) - - 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 - ) - - -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 - ) - - 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 - } - ) - - # 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_calculate_intensities_basic(self, sample_dataset_with_geometric_factors): - """Test intensity calculation with current implementation.""" - result = calculate_intensities(sample_dataset_with_geometric_factors) - - # 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 - - # 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 - ) - - # 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 - } - ) - - # 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 - - # 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) - - # 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) - - # 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))}, - ) - - result = calculate_backgrounds(dataset) - - # 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)) - - -class TestCalculateFluxCorrections: - """Tests for the calculate_flux_corrections function.""" - - 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) - - # Run flux correction - result = calculate_flux_corrections(original_dataset, lo_flux_factors_file) - # Verify that the function returns a dataset - assert isinstance(result, xr.Dataset) - # Verify that intensity variables are present - assert "ena_intensity" in result.data_vars - assert "ena_intensity_stat_uncert" in result.data_vars - - # Verify that data shape is preserved - original_shape = sample_dataset_with_intensities["ena_intensity"].shape - assert result["ena_intensity"].shape == original_shape - - # 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}" - ) +@pytest.fixture +def one_pointing(): + """A single synthetic pointing.""" + return make_pointing() - # 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_calculate_flux_corrections_preserves_other_vars( - self, sample_dataset_with_intensities, lo_flux_factors_file +@pytest.fixture +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 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)) - - result = calculate_flux_corrections( - sample_dataset_with_intensities, lo_flux_factors_file - ) + (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", + ) - # 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_map_dims_and_source(self, full_map): + """The map is a 6 degree rectangular map of the descriptor.""" + dataset, _ = full_map - 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), + assert dataset.attrs["Logical_source"] == f"imap_lo_l2_{FULL_DESCRIPTOR}" + assert dict(dataset.sizes) == { + "epoch": 1, + "energy": N_ESA, + "longitude": 60, + "latitude": 30, } - # 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 - - 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 - ) - - # 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}" - ) - - -class TestCalculateSputteringCorrections: - """Tests for the calculate_sputtering_corrections function.""" - - 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_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", ) - # 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" + def test_energy_coordinate(self, full_map): + """The energy coordinate and its widths come from the ESA constants.""" + dataset, _ = full_map - 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", - ) - - # 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 + dataset["energy"].values, LoConstants.ESA_ENERGY[:N_ESA] ) - - 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["energy_delta_plus"].values, LoConstants.ESA_ENERGY_DELTA[0] ) - 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_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" - 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], - } + cdf_path = write_cdf(dataset) - dataset = xr.Dataset(coords=coords) + 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) - # 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) +class TestAccumulation: + """What the map accumulates from its inputs.""" - dataset["ena_intensity"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values, - ) + def test_counts_are_conserved(self, full_map): + """Every in-window count lands somewhere on the map.""" + dataset, pointing = full_map - # 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, - ) + per_energy = dataset["ena_count"].values.sum(axis=(0, 2, 3)) + np.testing.assert_allclose(per_energy, pointing["expected_counts"]) - result = calculate_bootstrap_corrections(dataset) + def test_exposure_is_conserved(self, full_map): + """Every in-window second of exposure lands somewhere on the map.""" + dataset, pointing = full_map - # Lower energy channels should show larger corrections - # because they receive spillover from higher energy channels - original_ratios = [] - corrected_ratios = [] + per_energy = dataset["exposure_factor"].values.sum(axis=(0, 2, 3)) + np.testing.assert_allclose(per_energy, pointing["expected_exposure"], rtol=1e-6) - 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) + 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 - # Bootstrap should affect the intensities - # Check that the bootstrap corrections are actually being applied + 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"]) - # 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 + def test_pointings_accumulate(self, one_pointing): + """Two pointings contribute twice the counts of one.""" + other = make_pointing(repointing=101, seed=7) - # 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" - ) + with patch( + "imap_processing.lo.l1c.lo_l1c.frame_transform_az_el", + side_effect=identity_pointing, + ): + (one,) = lo_l2(as_dependencies(one_pointing), [], FULL_DESCRIPTOR) + (both,) = lo_l2(as_dependencies(one_pointing, other), [], FULL_DESCRIPTOR) - # 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" + 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"], ) - 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)) + 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, + ): + (full,) = lo_l2(as_dependencies(one_pointing), [], FULL_DESCRIPTOR) + (ram,) = lo_l2(as_dependencies(one_pointing), [], RAM_DESCRIPTOR) - # 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 - ) + full_counts = full["ena_count"].values.sum() + ram_counts = ram["ena_count"].values.sum() + assert 0 < ram_counts < full_counts - # 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, - ) +class TestRatesAndIntensities: + """The maths turning accumulated counts into intensities.""" - result = calculate_bootstrap_corrections(dataset) + def test_rate_and_intensity(self, full_map): + """Where exposed, rate = counts/exposure and intensity = rate/(G*E).""" + dataset, _ = full_map - # 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" - ) + counts = dataset["ena_count"].values + exposure = dataset["exposure_factor"].values + exposed = exposure > 0 + assert exposed.any() - # 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" + 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) - 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, + geometric_factor = ( + np.array(LoConstants.GEO_FACTOR[:N_ESA]) * LoConstants.GEO_FACTOR_SCALE ) - dataset["ena_intensity_sys_err"] = ( - ("epoch", "energy", "longitude", "latitude"), - intensity_values * 0.03, + energy = np.array(LoConstants.ESA_ENERGY[:N_ESA]) + expected = dataset["ena_count_rate"] / xr.DataArray( + geometric_factor * energy, dims=["energy"] ) - - # 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" + np.testing.assert_allclose( + dataset["ena_intensity"].values[exposed], + expected.values[exposed], + rtol=1e-5, ) - 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) + def test_statistical_uncertainty_is_poisson(self, full_map): + """The rate uncertainty is the Poisson count error over the exposure.""" + dataset, _ = full_map - 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)) + counts = dataset["ena_count"].values + exposure = dataset["exposure_factor"].values + exposed = exposure > 0 - # 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, + np.testing.assert_allclose( + dataset["ena_count_rate_stat_uncert"].values[exposed], + np.sqrt(counts[exposed]) / exposure[exposed], + rtol=1e-5, ) - 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 + def test_background_rate(self, full_map): + """The background rate is the input rate wherever the map was exposed.""" + dataset, pointing = full_map - # 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]}" + 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) - # 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) + def test_systematic_error_bounds(self, full_map): + """The systematic error is bracketed by the G-factor excursions.""" + dataset, _ = full_map - # 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" - ) + 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 - 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, - ), - } + 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]) - 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 +class TestGeometry: + """The spin-angle to sky-pixel geometry.""" - # 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"]) + def test_dps_spin_angles_carry_the_offset(self): + """The hardware spin bins are rotated onto the instrument frame.""" + angles = _dps_spin_angles() - # 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] + 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) - 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"]) + 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() + spin_angles = _dps_spin_angles() + pointing_set = make_pointing_set(sky_map, spin_angles) - 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, - ) + pixels = match_coords_to_indices(pointing_set, sky_map).values - # Both corrections should be applied - assert "ena_intensity" in result.data_vars - mock_interp.assert_called_once() + # 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_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()) -@pytest.mark.external_kernel -class TestIntegrationWithMocks: - """Integration tests using mocked external dependencies.""" + pixels = match_coords_to_indices(pointing_set, sky_map).values - 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 + assert len(np.unique(pixels)) == N_SPIN_BINS - # Run the function - should not crash - result = lo_l2(sci_dependencies, anc_dependencies, descriptor) + @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() - # Basic validation - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], xr.Dataset) + mask = _spin_phase_mask(angles, PIVOT, spin_phase) - # Mock the rates calculation to return the dataset unchanged - mock_calc_rates.side_effect = lambda x, **kwargs: x + assert mask.sum() == expected - # Run the function - should not crash - result = lo_l2(sci_dependencies, anc_dependencies, descriptor) + 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") - # Basic validation - assert isinstance(result, list) - assert len(result) == 1 - assert isinstance(result[0], xr.Dataset) +class TestPointingSelection: + """Reducing the grouped inputs to the pointings that can be mapped.""" -@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" + 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 = _complete_pointings(as_dependencies(one_pointing, other)) -@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(ValueError, match="No pointing set data found"): - 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 - } + assert set(pointings) == {100, 101} + assert pointings[100] == ( + one_pointing["goodtimes"], + one_pointing["bgrates"], + one_pointing["histrates"], ) - 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)), - } + def test_incomplete_pointings_are_dropped(self, one_pointing, caplog): + """A pointing missing one of the three products cannot be mapped.""" + incomplete = make_pointing(repointing=101) + dependencies = as_dependencies(one_pointing) | as_dependencies( + incomplete, products=("goodtimes", "histrates") ) - result = calculate_rates(dataset) + pointings = _complete_pointings(dependencies) - # 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 + assert set(pointings) == {100} + assert "repoint00101" in caplog.text + assert "bgrates" in caplog.text - # 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]) + def test_missing_product_raises(self, one_pointing): + """A map cannot be made without all three products.""" + dependencies = as_dependencies(one_pointing, products=("goodtimes", "bgrates")) + with pytest.raises(KeyError, match="histrates"): + _complete_pointings(dependencies) -# ============================================================================= -# TESTS FOR CG CORRECTION FUNCTIONALITY -# ============================================================================= +class TestUnsupported: + """Map flavours the Lo pipeline does not make.""" -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 + 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", ) - # 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 + 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", ) - - ( - 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 92acd8cbe..15aedc50b 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,28 @@ 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 = [ - "imap_lo_l1c_pset_20250415_v001.cdf", - "imap_lo_l1c_pset_20250416_v001.cdf", + 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_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 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] + # Loaded in descriptor order: goodtimes, then bgrates, then histrates + mocks["mock_load_cdf"].side_effect = [ + mock_goodtimes, + mock_bgrates, + mock_histrates_1, + mock_histrates_2, + ] mock_lo_pre_processing.return_value = processing_input output_l2_dataset = xr.Dataset() @@ -499,8 +510,16 @@ 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( - {"some_pset_logical_source": [mock_loaded_pset_1, sentinel.loaded_pset_2]}, + { + 217: { + "goodtimes": mock_goodtimes, + "bgrates": mock_bgrates, + "histrates": mock_histrates_1, + }, + 218: {"histrates": mock_histrates_2}, + }, [], descriptor, ) @@ -510,37 +529,81 @@ 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): - 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" + """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(non_pset), + 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_angle": xr.DataArray(30.0)}), + xr.Dataset({"pivot": ("epoch", [90.1])}), + # A neighbouring pivot angle, which belongs on its own map + xr.Dataset({"pivot": ("epoch", [75.0])}), ] instrument = Lo( "l2", - "some-descriptor", + "l090-ena-h-sf-nsp-ram-hae-6deg-3mo", base_collection.serialize(), "20250415", - "20250416", + "20250715", "v001", False, ) result = instrument.pre_processing() - result_inputs = list(result.get_processing_inputs()) - assert len(result_inputs) == 2 + # 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() + ] == [[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() - 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] + assert list(result.get_processing_inputs()) == [] @mock.patch("imap_processing.cli.quaternions.process_quaternions", autospec=True)