imap-processing 0.18.0__py3-none-any.whl → 0.19.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of imap-processing might be problematic. Click here for more details.

Files changed (122) hide show
  1. imap_processing/_version.py +2 -2
  2. imap_processing/ancillary/ancillary_dataset_combiner.py +161 -1
  3. imap_processing/cdf/config/imap_codice_global_cdf_attrs.yaml +6 -0
  4. imap_processing/cdf/config/imap_codice_l1a_variable_attrs.yaml +221 -1057
  5. imap_processing/cdf/config/imap_codice_l1b_variable_attrs.yaml +307 -283
  6. imap_processing/cdf/config/imap_codice_l2_variable_attrs.yaml +1044 -203
  7. imap_processing/cdf/config/imap_constant_attrs.yaml +4 -2
  8. imap_processing/cdf/config/imap_enamaps_l2-common_variable_attrs.yaml +11 -0
  9. imap_processing/cdf/config/imap_glows_l1b_variable_attrs.yaml +15 -1
  10. imap_processing/cdf/config/imap_hi_global_cdf_attrs.yaml +5 -0
  11. imap_processing/cdf/config/imap_hit_global_cdf_attrs.yaml +10 -4
  12. imap_processing/cdf/config/imap_idex_l2a_variable_attrs.yaml +33 -4
  13. imap_processing/cdf/config/imap_idex_l2b_variable_attrs.yaml +8 -91
  14. imap_processing/cdf/config/imap_idex_l2c_variable_attrs.yaml +106 -16
  15. imap_processing/cdf/config/imap_lo_global_cdf_attrs.yaml +5 -4
  16. imap_processing/cdf/config/imap_lo_l1a_variable_attrs.yaml +4 -15
  17. imap_processing/cdf/config/imap_lo_l1c_variable_attrs.yaml +189 -98
  18. imap_processing/cdf/config/imap_mag_global_cdf_attrs.yaml +85 -2
  19. imap_processing/cdf/config/imap_mag_l1c_variable_attrs.yaml +24 -1
  20. imap_processing/cdf/config/imap_ultra_global_cdf_attrs.yaml +20 -8
  21. imap_processing/cdf/config/imap_ultra_l1b_variable_attrs.yaml +45 -35
  22. imap_processing/cdf/config/imap_ultra_l1c_variable_attrs.yaml +110 -7
  23. imap_processing/cli.py +138 -93
  24. imap_processing/codice/codice_l0.py +2 -1
  25. imap_processing/codice/codice_l1a.py +167 -69
  26. imap_processing/codice/codice_l1b.py +42 -32
  27. imap_processing/codice/codice_l2.py +215 -9
  28. imap_processing/codice/constants.py +790 -603
  29. imap_processing/codice/data/lo_stepping_values.csv +1 -1
  30. imap_processing/decom.py +1 -4
  31. imap_processing/ena_maps/ena_maps.py +71 -43
  32. imap_processing/ena_maps/utils/corrections.py +291 -0
  33. imap_processing/ena_maps/utils/map_utils.py +20 -4
  34. imap_processing/ena_maps/utils/naming.py +8 -2
  35. imap_processing/glows/ancillary/imap_glows_exclusions-by-instr-team_20250923_v002.dat +10 -0
  36. imap_processing/glows/ancillary/imap_glows_map-of-excluded-regions_20250923_v002.dat +393 -0
  37. imap_processing/glows/ancillary/imap_glows_map-of-uv-sources_20250923_v002.dat +593 -0
  38. imap_processing/glows/ancillary/imap_glows_pipeline-settings_20250923_v002.json +54 -0
  39. imap_processing/glows/ancillary/imap_glows_suspected-transients_20250923_v002.dat +10 -0
  40. imap_processing/glows/l1b/glows_l1b.py +123 -18
  41. imap_processing/glows/l1b/glows_l1b_data.py +358 -47
  42. imap_processing/glows/l2/glows_l2.py +11 -0
  43. imap_processing/hi/hi_l1a.py +124 -3
  44. imap_processing/hi/hi_l1b.py +154 -71
  45. imap_processing/hi/hi_l1c.py +4 -109
  46. imap_processing/hi/hi_l2.py +104 -60
  47. imap_processing/hi/utils.py +262 -8
  48. imap_processing/hit/l0/constants.py +3 -0
  49. imap_processing/hit/l0/decom_hit.py +3 -6
  50. imap_processing/hit/l1a/hit_l1a.py +311 -21
  51. imap_processing/hit/l1b/hit_l1b.py +54 -126
  52. imap_processing/hit/l2/hit_l2.py +6 -6
  53. imap_processing/ialirt/calculate_ingest.py +219 -0
  54. imap_processing/ialirt/constants.py +12 -2
  55. imap_processing/ialirt/generate_coverage.py +15 -2
  56. imap_processing/ialirt/l0/ialirt_spice.py +6 -2
  57. imap_processing/ialirt/l0/parse_mag.py +293 -42
  58. imap_processing/ialirt/l0/process_hit.py +5 -3
  59. imap_processing/ialirt/l0/process_swapi.py +41 -25
  60. imap_processing/ialirt/process_ephemeris.py +70 -14
  61. imap_processing/ialirt/utils/create_xarray.py +1 -1
  62. imap_processing/idex/idex_l0.py +2 -2
  63. imap_processing/idex/idex_l1a.py +2 -3
  64. imap_processing/idex/idex_l1b.py +2 -3
  65. imap_processing/idex/idex_l2a.py +130 -4
  66. imap_processing/idex/idex_l2b.py +158 -143
  67. imap_processing/idex/idex_utils.py +1 -3
  68. imap_processing/lo/ancillary_data/imap_lo_hydrogen-geometric-factor_v001.csv +75 -0
  69. imap_processing/lo/ancillary_data/imap_lo_oxygen-geometric-factor_v001.csv +75 -0
  70. imap_processing/lo/l0/lo_science.py +25 -24
  71. imap_processing/lo/l1b/lo_l1b.py +93 -19
  72. imap_processing/lo/l1c/lo_l1c.py +273 -93
  73. imap_processing/lo/l2/lo_l2.py +949 -135
  74. imap_processing/lo/lo_ancillary.py +55 -0
  75. imap_processing/mag/l1a/mag_l1a.py +1 -0
  76. imap_processing/mag/l1a/mag_l1a_data.py +26 -0
  77. imap_processing/mag/l1b/mag_l1b.py +3 -2
  78. imap_processing/mag/l1c/interpolation_methods.py +14 -15
  79. imap_processing/mag/l1c/mag_l1c.py +23 -6
  80. imap_processing/mag/l1d/mag_l1d.py +57 -14
  81. imap_processing/mag/l1d/mag_l1d_data.py +202 -32
  82. imap_processing/mag/l2/mag_l2.py +2 -0
  83. imap_processing/mag/l2/mag_l2_data.py +14 -5
  84. imap_processing/quality_flags.py +23 -1
  85. imap_processing/spice/geometry.py +89 -39
  86. imap_processing/spice/pointing_frame.py +4 -8
  87. imap_processing/spice/repoint.py +78 -2
  88. imap_processing/spice/spin.py +28 -8
  89. imap_processing/spice/time.py +12 -22
  90. imap_processing/swapi/l1/swapi_l1.py +10 -4
  91. imap_processing/swapi/l2/swapi_l2.py +15 -17
  92. imap_processing/swe/l1b/swe_l1b.py +1 -2
  93. imap_processing/ultra/constants.py +30 -24
  94. imap_processing/ultra/l0/ultra_utils.py +9 -11
  95. imap_processing/ultra/l1a/ultra_l1a.py +1 -2
  96. imap_processing/ultra/l1b/badtimes.py +35 -11
  97. imap_processing/ultra/l1b/de.py +95 -31
  98. imap_processing/ultra/l1b/extendedspin.py +31 -16
  99. imap_processing/ultra/l1b/goodtimes.py +112 -0
  100. imap_processing/ultra/l1b/lookup_utils.py +281 -28
  101. imap_processing/ultra/l1b/quality_flag_filters.py +10 -1
  102. imap_processing/ultra/l1b/ultra_l1b.py +7 -7
  103. imap_processing/ultra/l1b/ultra_l1b_culling.py +169 -7
  104. imap_processing/ultra/l1b/ultra_l1b_extended.py +311 -69
  105. imap_processing/ultra/l1c/helio_pset.py +139 -37
  106. imap_processing/ultra/l1c/l1c_lookup_utils.py +289 -0
  107. imap_processing/ultra/l1c/spacecraft_pset.py +140 -29
  108. imap_processing/ultra/l1c/ultra_l1c.py +33 -24
  109. imap_processing/ultra/l1c/ultra_l1c_culling.py +92 -0
  110. imap_processing/ultra/l1c/ultra_l1c_pset_bins.py +400 -292
  111. imap_processing/ultra/l2/ultra_l2.py +54 -11
  112. imap_processing/ultra/utils/ultra_l1_utils.py +37 -7
  113. imap_processing/utils.py +3 -4
  114. {imap_processing-0.18.0.dist-info → imap_processing-0.19.2.dist-info}/METADATA +2 -2
  115. {imap_processing-0.18.0.dist-info → imap_processing-0.19.2.dist-info}/RECORD +118 -109
  116. imap_processing/idex/idex_l2c.py +0 -84
  117. imap_processing/spice/kernels.py +0 -187
  118. imap_processing/ultra/l1b/cullingmask.py +0 -87
  119. imap_processing/ultra/l1c/histogram.py +0 -36
  120. {imap_processing-0.18.0.dist-info → imap_processing-0.19.2.dist-info}/LICENSE +0 -0
  121. {imap_processing-0.18.0.dist-info → imap_processing-0.19.2.dist-info}/WHEEL +0 -0
  122. {imap_processing-0.18.0.dist-info → imap_processing-0.19.2.dist-info}/entry_points.txt +0 -0
@@ -1,213 +1,1027 @@
1
1
  """IMAP-Lo L2 data processing."""
2
2
 
3
+ import logging
4
+ from pathlib import Path
5
+
3
6
  import numpy as np
7
+ import pandas as pd
4
8
  import xarray as xr
5
9
 
6
10
  from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
7
11
  from imap_processing.ena_maps import ena_maps
8
- from imap_processing.ena_maps.ena_maps import RectangularSkyMap
9
- from imap_processing.spice import geometry
10
- from imap_processing.spice.geometry import SpiceFrame
12
+ from imap_processing.ena_maps.ena_maps import AbstractSkyMap, RectangularSkyMap
13
+ from imap_processing.ena_maps.utils.naming import MapDescriptor
14
+ from imap_processing.lo import lo_ancillary
15
+ from imap_processing.spice.time import et_to_datetime64, ttj2000ns_to_et
16
+
17
+ logger = logging.getLogger(__name__)
18
+
19
+ # =============================================================================
20
+ # MAIN ENTRY POINT
21
+ # =============================================================================
11
22
 
12
23
 
13
- def lo_l2(sci_dependencies: dict, anc_dependencies: list) -> list[xr.Dataset]:
24
+ def lo_l2(
25
+ sci_dependencies: dict, anc_dependencies: list, descriptor: str
26
+ ) -> list[xr.Dataset]:
14
27
  """
15
- Will process IMAP-Lo L1C data into Le CDF data products.
28
+ Process IMAP-Lo L1C data into L2 CDF data products.
29
+
30
+ This is the main entry point for L2 processing. It orchestrates the entire
31
+ processing pipeline from L1C pointing sets to L2 sky maps with intensities.
16
32
 
17
33
  Parameters
18
34
  ----------
19
35
  sci_dependencies : dict
20
36
  Dictionary of datasets needed for L2 data product creation in xarray Datasets.
37
+ Must contain "imap_lo_l1c_pset" key with list of pointing set datasets.
21
38
  anc_dependencies : list
22
- Ancillary files needed for L2 data product creation.
39
+ List of ancillary file paths needed for L2 data product creation.
40
+ Should include efficiency factor files.
41
+ descriptor : str
42
+ The map descriptor to be produced
43
+ (e.g., "ilo90-ena-h-sf-nsp-full-hae-6deg-3mo").
23
44
 
24
45
  Returns
25
46
  -------
26
- created_file_paths : list[Path]
27
- Location of created CDF files.
47
+ list[xr.Dataset]
48
+ List containing the processed L2 dataset with rates, intensities,
49
+ and uncertainties.
50
+
51
+ Raises
52
+ ------
53
+ ValueError
54
+ If no pointing set data found in science dependencies.
55
+ NotImplementedError
56
+ If HEALPix map output is requested (only rectangular maps supported).
57
+ """
58
+ logger.info("Starting IMAP-Lo L2 processing pipeline")
59
+ if "imap_lo_l1c_pset" not in sci_dependencies:
60
+ raise ValueError("No pointing set data found in science dependencies")
61
+ psets = sci_dependencies["imap_lo_l1c_pset"]
62
+
63
+ # Parse the map descriptor to get species and other attributes
64
+ map_descriptor = MapDescriptor.from_string(descriptor)
65
+ logger.info(f"Processing map for species: {map_descriptor.species}")
66
+
67
+ logger.info("Step 1: Loading ancillary data")
68
+ efficiency_data = load_efficiency_data(anc_dependencies)
69
+
70
+ logger.info(f"Step 2: Creating sky map from {len(psets)} pointing sets")
71
+ sky_map = create_sky_map_from_psets(psets, map_descriptor, efficiency_data)
72
+
73
+ logger.info("Step 3: Converting to dataset and adding geometric factors")
74
+ dataset = sky_map.to_dataset()
75
+ dataset = add_geometric_factors(dataset, map_descriptor.species)
76
+
77
+ logger.info("Step 4: Calculating rates and intensities")
78
+ dataset = calculate_all_rates_and_intensities(dataset)
79
+
80
+ logger.info("Step 5: Finalizing dataset with attributes")
81
+ dataset = finalize_dataset(dataset, descriptor)
82
+
83
+ logger.info("IMAP-Lo L2 processing pipeline completed successfully")
84
+ return [dataset]
85
+
86
+
87
+ # =============================================================================
88
+ # SETUP AND INITIALIZATION HELPERS
89
+ # =============================================================================
90
+
91
+
92
+ def load_efficiency_data(anc_dependencies: list) -> pd.DataFrame:
93
+ """
94
+ Load efficiency factor data from ancillary files.
95
+
96
+ Parameters
97
+ ----------
98
+ anc_dependencies : list
99
+ List of ancillary file paths to search for efficiency factor files.
100
+
101
+ Returns
102
+ -------
103
+ pd.DataFrame
104
+ Concatenated efficiency factor data from all matching files.
105
+ Returns empty DataFrame if no efficiency files found.
106
+ """
107
+ efficiency_files = [
108
+ anc_file
109
+ for anc_file in anc_dependencies
110
+ if "efficiency-factor" in str(anc_file)
111
+ ]
112
+
113
+ if not efficiency_files:
114
+ logger.warning("No efficiency factor files found in ancillary dependencies")
115
+ return pd.DataFrame()
116
+
117
+ logger.debug(f"Loading {len(efficiency_files)} efficiency factor files")
118
+ return pd.concat(
119
+ [lo_ancillary.read_ancillary_file(anc_file) for anc_file in efficiency_files],
120
+ ignore_index=True,
121
+ )
122
+
123
+
124
+ def finalize_dataset(dataset: xr.Dataset, descriptor: str) -> xr.Dataset:
125
+ """
126
+ Add attributes and perform final dataset preparation.
127
+
128
+ Parameters
129
+ ----------
130
+ dataset : xr.Dataset
131
+ The dataset to finalize with attributes.
132
+ descriptor : str
133
+ The descriptor for this map dataset.
134
+
135
+ Returns
136
+ -------
137
+ xr.Dataset
138
+ The finalized dataset with all attributes added.
28
139
  """
29
- # create the attribute manager for this data level
140
+ # Initialize the attribute manager
30
141
  attr_mgr = ImapCdfAttributes()
31
142
  attr_mgr.add_instrument_global_attrs(instrument="lo")
32
143
  attr_mgr.add_instrument_variable_attrs(instrument="enamaps", level="l2-common")
33
144
  attr_mgr.add_instrument_variable_attrs(instrument="enamaps", level="l2-rectangular")
34
145
 
35
- # if the dependencies are used to create Annotated Direct Events
36
- if "imap_lo_l1c_pset" in sci_dependencies:
37
- logical_source = "imap_lo_l2_l090-ena-h-sf-nsp-ram-hae-6deg-3mo"
38
- psets = sci_dependencies["imap_lo_l1c_pset"]
146
+ # Add global and variable attributes
147
+ dataset.attrs.update(attr_mgr.get_global_attributes("imap_lo_l2_enamap"))
39
148
 
40
- # Create the rectangular sky map from the pointing set.
41
- lo_rect_map = project_pset_to_rect_map(
42
- psets, spacing_deg=6, spice_frame=geometry.SpiceFrame.ECLIPJ2000
149
+ # Our global attributes have placeholders for descriptor
150
+ # so iterate through here and fill that in with the map-specific descriptor
151
+ for key in ["Data_type", "Logical_source", "Logical_source_description"]:
152
+ dataset.attrs[key] = dataset.attrs[key].format(descriptor=descriptor)
153
+ for var in dataset.data_vars:
154
+ try:
155
+ dataset[var].attrs = attr_mgr.get_variable_attributes(var)
156
+ except KeyError:
157
+ # If no attributes found, try without schema validation
158
+ try:
159
+ dataset[var].attrs = attr_mgr.get_variable_attributes(
160
+ var, check_schema=False
161
+ )
162
+ except KeyError:
163
+ logger.warning(f"No attributes found for variable {var}")
164
+
165
+ return dataset
166
+
167
+
168
+ # =============================================================================
169
+ # SKY MAP CREATION PIPELINE
170
+ # =============================================================================
171
+
172
+
173
+ def create_sky_map_from_psets(
174
+ psets: list[xr.Dataset],
175
+ map_descriptor: MapDescriptor,
176
+ efficiency_data: pd.DataFrame,
177
+ ) -> AbstractSkyMap:
178
+ """
179
+ Create a sky map by processing all pointing sets.
180
+
181
+ Parameters
182
+ ----------
183
+ psets : list[xr.Dataset]
184
+ List of pointing set datasets to process.
185
+ map_descriptor : MapDescriptor
186
+ Map descriptor object defining the projection and binning.
187
+ efficiency_data : pd.DataFrame
188
+ Efficiency factor data for correcting counts.
189
+
190
+ Returns
191
+ -------
192
+ AbstractSkyMap
193
+ The populated sky map with projected data from all pointing sets.
194
+
195
+ Raises
196
+ ------
197
+ NotImplementedError
198
+ If HEALPix map output is requested (only rectangular maps supported).
199
+ """
200
+ # Initialize the output map
201
+ output_map = map_descriptor.to_empty_map()
202
+
203
+ if not isinstance(output_map, RectangularSkyMap):
204
+ raise NotImplementedError("HEALPix map output not supported for Lo")
205
+
206
+ logger.debug(f"Processing {len(psets)} pointing sets")
207
+ # Process each pointing set
208
+ for i, pset in enumerate(psets):
209
+ logger.debug(f"Processing pointing set {i + 1}/{len(psets)}")
210
+ processed_pset = process_single_pset(
211
+ pset, efficiency_data, map_descriptor.species
43
212
  )
44
- # Add the hydrogen rates to the rectangular map dataset.
45
- lo_rect_map.data_1d["h_rate"] = calculate_rates(
46
- lo_rect_map.data_1d["h_counts"], lo_rect_map.data_1d["exposure_time"]
213
+ project_pset_to_map(processed_pset, output_map)
214
+
215
+ return output_map
216
+
217
+
218
+ def process_single_pset(
219
+ pset: xr.Dataset,
220
+ efficiency_data: pd.DataFrame,
221
+ species: str,
222
+ ) -> xr.Dataset:
223
+ """
224
+ Process a single pointing set for projection to the sky map.
225
+
226
+ Parameters
227
+ ----------
228
+ pset : xr.Dataset
229
+ Single pointing set dataset to process.
230
+ efficiency_data : pd.DataFrame
231
+ Efficiency factor data for correcting counts.
232
+ species : str
233
+ The species to process (e.g., "h", "o").
234
+
235
+ Returns
236
+ -------
237
+ xr.Dataset
238
+ Processed pointing set ready for projection with efficiency corrections applied.
239
+ """
240
+ # Step 1: Normalize coordinate system
241
+ pset_processed = normalize_pset_coordinates(pset, species)
242
+
243
+ # Step 2: Add efficiency factors
244
+ pset_processed = add_efficiency_factors_to_pset(pset_processed, efficiency_data)
245
+
246
+ # Step 3: Calculate efficiency-corrected quantities
247
+ pset_processed = calculate_efficiency_corrected_quantities(pset_processed)
248
+
249
+ return pset_processed
250
+
251
+
252
+ def normalize_pset_coordinates(pset: xr.Dataset, species: str) -> xr.Dataset:
253
+ """
254
+ Normalize pointing set coordinates to match the output map.
255
+
256
+ Parameters
257
+ ----------
258
+ pset : xr.Dataset
259
+ Input pointing set dataset with potentially mismatched coordinates.
260
+ species : str
261
+ The species to process (e.g., "h", "o").
262
+
263
+ Returns
264
+ -------
265
+ xr.Dataset
266
+ Pointing set with normalized energy coordinates and dimension names.
267
+ """
268
+ # Ensure consistent energy coordinates (maps want energy not esa_energy_step)
269
+ pset_renamed = pset.rename_dims({"esa_energy_step": "energy"})
270
+
271
+ # Drop the esa_energy_step coordinate first to avoid conflicts
272
+ pset_renamed = pset_renamed.drop_vars("esa_energy_step")
273
+
274
+ # Ensure the pset energy coordinates match the output map
275
+ # TODO: Do we even need this if we are assigning the true
276
+ # energy levels later?
277
+ pset_renamed = pset_renamed.assign_coords(energy=range(7))
278
+
279
+ # Rename the variables in the pset for projection to the map
280
+ # L2 wants different variable names than l1c
281
+ rename_map = {
282
+ "exposure_time": "exposure_factor",
283
+ f"{species}_counts": "counts",
284
+ f"{species}_background_rates": "bg_rates",
285
+ f"{species}_background_rates_stat_uncert": "bg_rates_stat_uncert",
286
+ }
287
+ pset_renamed = pset_renamed.rename_vars(rename_map)
288
+
289
+ return pset_renamed
290
+
291
+
292
+ def add_efficiency_factors_to_pset(
293
+ pset: xr.Dataset, efficiency_data: pd.DataFrame
294
+ ) -> xr.Dataset:
295
+ """
296
+ Add efficiency factors to the pointing set based on observation date.
297
+
298
+ Parameters
299
+ ----------
300
+ pset : xr.Dataset
301
+ Pointing set dataset to add efficiency factors to.
302
+ efficiency_data : pd.DataFrame
303
+ Efficiency factor data containing date-indexed efficiency values.
304
+
305
+ Returns
306
+ -------
307
+ xr.Dataset
308
+ Pointing set with efficiency factors added as new data variable.
309
+
310
+ Raises
311
+ ------
312
+ ValueError
313
+ If no efficiency factor found for the pointing set observation date.
314
+ """
315
+ if efficiency_data.empty:
316
+ # If no efficiency data, create unity efficiency
317
+ logger.warning("No efficiency data available, using unity efficiency")
318
+ pset["efficiency"] = xr.DataArray(np.ones(7), dims=["energy"])
319
+ return pset
320
+
321
+ # Convert the epoch to datetime64
322
+ date = et_to_datetime64(ttj2000ns_to_et(pset["epoch"].values[0]))
323
+ # The efficiency file only has date as YYYYDDD, so drop the time for this
324
+ date = date.astype("M8[D]") # Convert to date only (no time)
325
+
326
+ ef_df = efficiency_data[efficiency_data["Date"] == date]
327
+ if ef_df.empty:
328
+ raise ValueError(f"No efficiency factor found for pset date {date}")
329
+
330
+ efficiency_values = ef_df[
331
+ [
332
+ "E-Step1_eff",
333
+ "E-Step2_eff",
334
+ "E-Step3_eff",
335
+ "E-Step4_eff",
336
+ "E-Step5_eff",
337
+ "E-Step6_eff",
338
+ "E-Step7_eff",
339
+ ]
340
+ ].values[0]
341
+
342
+ pset["efficiency"] = xr.DataArray(
343
+ efficiency_values,
344
+ dims=["energy"],
345
+ )
346
+ logger.debug(f"Applied efficiency factors for date {date}")
347
+ return pset
348
+
349
+
350
+ def calculate_efficiency_corrected_quantities(
351
+ pset: xr.Dataset,
352
+ ) -> xr.Dataset:
353
+ """
354
+ Calculate efficiency-corrected quantities for each particle type.
355
+
356
+ Parameters
357
+ ----------
358
+ pset : xr.Dataset
359
+ Pointing set with efficiency factors applied.
360
+
361
+ Returns
362
+ -------
363
+ xr.Dataset
364
+ Pointing set with efficiency-corrected count variables added.
365
+ """
366
+ # counts / efficiency
367
+ pset["counts_over_eff"] = pset["counts"] / pset["efficiency"]
368
+ # counts / efficiency**2 (for variance propagation)
369
+ pset["counts_over_eff_squared"] = pset["counts"] / (pset["efficiency"] ** 2)
370
+
371
+ # background * exposure_factor for weighted average
372
+ pset["bg_rates_exposure_factor"] = pset["bg_rates"] * pset["exposure_factor"]
373
+ # background_uncertainty ** 2 * exposure_factor ** 2
374
+ pset["bg_rates_stat_uncert_exposure_factor2"] = (
375
+ pset["bg_rates_stat_uncert"] ** 2 * pset["exposure_factor"] ** 2
376
+ )
377
+
378
+ return pset
379
+
380
+
381
+ def project_pset_to_map(
382
+ pset: xr.Dataset,
383
+ output_map: AbstractSkyMap,
384
+ ) -> None:
385
+ """
386
+ Project pointing set data to the output map.
387
+
388
+ Parameters
389
+ ----------
390
+ pset : xr.Dataset
391
+ Processed pointing set ready for projection.
392
+ output_map : AbstractSkyMap
393
+ Target sky map to receive the projected data.
394
+
395
+ Returns
396
+ -------
397
+ None
398
+ Function modifies output_map in place.
399
+ """
400
+ # Define base quantities to project
401
+ value_keys = [
402
+ "exposure_factor",
403
+ "counts",
404
+ "counts_over_eff",
405
+ "counts_over_eff_squared",
406
+ "bg_rates",
407
+ "bg_rates_stat_uncert",
408
+ "bg_rates_exposure_factor",
409
+ "bg_rates_stat_uncert_exposure_factor2",
410
+ ]
411
+
412
+ # Create LoPointingSet and project to map
413
+ lo_pset = ena_maps.LoPointingSet(pset)
414
+ output_map.project_pset_values_to_map(
415
+ pointing_set=lo_pset,
416
+ value_keys=value_keys,
417
+ index_match_method=ena_maps.IndexMatchMethod.PUSH,
418
+ )
419
+ logger.debug(f"Projected {len(value_keys)} quantities to sky map")
420
+
421
+
422
+ # =============================================================================
423
+ # GEOMETRIC FACTORS
424
+ # =============================================================================
425
+
426
+
427
+ def add_geometric_factors(dataset: xr.Dataset, species: str) -> xr.Dataset:
428
+ """
429
+ Add geometric factors to the sky map after projection.
430
+
431
+ Parameters
432
+ ----------
433
+ dataset : xr.Dataset
434
+ Sky map dataset to add geometric factors to.
435
+ species : str
436
+ The species to process (only "h" and "o" have geometric factors).
437
+
438
+ Returns
439
+ -------
440
+ xr.Dataset
441
+ Dataset with geometric factor variables added for the specified species.
442
+ """
443
+ # Only add geometric factors for hydrogen and oxygen
444
+ if species not in ["h", "o"]:
445
+ logger.warning(f"No geometric factors to add for species: {species}")
446
+ return dataset
447
+
448
+ logger.info(f"Loading and applying geometric factors for species: {species}")
449
+
450
+ # Load geometric factor data for the specific species
451
+ gf_data = load_geometric_factor_data(species)
452
+
453
+ # Initialize geometric factor variables
454
+ dataset = initialize_geometric_factor_variables(dataset)
455
+
456
+ # Populate geometric factors for each energy step
457
+ dataset = populate_geometric_factors(dataset, gf_data, species)
458
+
459
+ return dataset
460
+
461
+
462
+ def load_geometric_factor_data(species: str) -> pd.DataFrame:
463
+ """
464
+ Load geometric factor data for the specified species.
465
+
466
+ Parameters
467
+ ----------
468
+ species : str
469
+ The species to load geometric factors for ("h" or "o").
470
+
471
+ Returns
472
+ -------
473
+ pd.DataFrame
474
+ Geometric factor dataframe for the specified species.
475
+
476
+ Raises
477
+ ------
478
+ ValueError
479
+ If species is not "h" or "o".
480
+ """
481
+ if species not in ["h", "o"]:
482
+ raise ValueError(
483
+ f"Geometric factors only available for 'h' and 'o', got '{species}'"
47
484
  )
48
- # Add the hydrogen flux to the rectangular map dataset.
49
- lo_rect_map.data_1d["h_flux"] = calculate_fluxes(lo_rect_map.data_1d["h_rate"])
50
- # Create the dataset from the rectangular map.
51
- lo_rect_map_ds = lo_rect_map.to_dataset()
52
- # Add the attributes to the dataset.
53
- lo_rect_map_ds = add_attributes(
54
- lo_rect_map_ds, attr_mgr, logical_source=logical_source
485
+
486
+ anc_path = Path(__file__).parent.parent / "ancillary_data"
487
+
488
+ if species == "h":
489
+ gf_file = anc_path / "imap_lo_hydrogen-geometric-factor_v001.csv"
490
+ else: # species == "o"
491
+ gf_file = anc_path / "imap_lo_oxygen-geometric-factor_v001.csv"
492
+
493
+ return lo_ancillary.read_ancillary_file(gf_file)
494
+
495
+
496
+ def initialize_geometric_factor_variables(
497
+ dataset: xr.Dataset,
498
+ ) -> xr.Dataset:
499
+ """
500
+ Initialize geometric factor variables for the specified species.
501
+
502
+ Parameters
503
+ ----------
504
+ dataset : xr.Dataset
505
+ Input dataset to add geometric factor variables to.
506
+
507
+ Returns
508
+ -------
509
+ xr.Dataset
510
+ Dataset with initialized geometric factor variables for the specified species.
511
+ """
512
+ gf_vars = [
513
+ "energy",
514
+ "energy_stat_uncert",
515
+ "geometric_factor",
516
+ "geometric_factor_stat_uncert",
517
+ ]
518
+
519
+ # Initialize variables with proper dimensions (energy only)
520
+ for var in gf_vars:
521
+ dataset[var] = xr.DataArray(
522
+ np.zeros(7),
523
+ dims=["energy"],
55
524
  )
56
525
 
57
- return [lo_rect_map_ds]
526
+ return dataset
527
+
528
+
529
+ def populate_geometric_factors(
530
+ dataset: xr.Dataset,
531
+ gf_data: pd.DataFrame,
532
+ species: str,
533
+ ) -> xr.Dataset:
534
+ """
535
+ Populate geometric factor values for each energy step.
536
+
537
+ Parameters
538
+ ----------
539
+ dataset : xr.Dataset
540
+ Dataset with initialized geometric factor variables.
541
+ gf_data : pd.DataFrame
542
+ Geometric factor data for the specified species.
543
+ species : str
544
+ The species to process (only "h" and "o" have geometric factors).
545
+
546
+ Returns
547
+ -------
548
+ xr.Dataset
549
+ Dataset with populated geometric factor values for the specified species.
550
+ """
551
+ # Only populate if the species has geometric factors
552
+ if species not in ["h", "o"]:
553
+ logger.debug(f"No geometric factors to populate for species: {species}")
554
+ return dataset
555
+
556
+ # Mapping of dataset variables to dataframe columns for this species
557
+ if species == "h":
558
+ gf_vars = {
559
+ "energy": "Cntr_E",
560
+ "energy_stat_uncert": "Cntr_E_unc",
561
+ "geometric_factor": "GF_Trpl_H",
562
+ "geometric_factor_stat_uncert": "GF_Trpl_H_unc",
563
+ }
564
+ else: # species == "o"
565
+ gf_vars = {
566
+ "energy": "Cntr_E",
567
+ "energy_stat_uncert": "Cntr_E_unc",
568
+ "geometric_factor": "GF_Trpl_O",
569
+ "geometric_factor_stat_uncert": "GF_Trpl_O_unc",
570
+ }
571
+
572
+ # Get ESA mode from the map (assuming it's constant or we take the first)
573
+ # TODO: Figure out how to handle esa_mode properly
574
+ if "esa_mode" in dataset:
575
+ esa_mode = dataset["esa_mode"].values[0]
576
+ else:
577
+ # Default to mode 0 if not available (HiRes mode)
578
+ esa_mode = 0
58
579
 
580
+ # Populate the geometric factors for each energy step
581
+ for i in range(7):
582
+ # Get geometric factor data for this energy step and ESA mode
583
+ gf_row = gf_data[
584
+ (gf_data["esa_mode"] == esa_mode) & (gf_data["Observed_E-Step"] == i + 1)
585
+ ].iloc[0]
59
586
 
60
- def project_pset_to_rect_map(
61
- psets: list[xr.Dataset], spacing_deg: int, spice_frame: SpiceFrame
62
- ) -> RectangularSkyMap:
587
+ # Fill energy step with the geometric factor values
588
+ for var, col in gf_vars.items():
589
+ dataset[var].values[i] = gf_row[col]
590
+
591
+ return dataset
592
+
593
+
594
+ # =============================================================================
595
+ # RATES AND INTENSITIES CALCULATIONS
596
+ # =============================================================================
597
+
598
+
599
+ def calculate_all_rates_and_intensities(
600
+ dataset: xr.Dataset,
601
+ sputtering_correction: bool = False,
602
+ bootstrap_correction: bool = False,
603
+ ) -> xr.Dataset:
604
+ """
605
+ Calculate rates and intensities with proper error propagation.
606
+
607
+ Parameters
608
+ ----------
609
+ dataset : xr.Dataset
610
+ Sky map dataset with count data and geometric factors.
611
+ sputtering_correction : bool, optional
612
+ Whether to apply sputtering corrections to oxygen intensities.
613
+ Default is False.
614
+ bootstrap_correction : bool, optional
615
+ Whether to apply bootstrap corrections to intensities.
616
+ Default is False.
617
+
618
+ Returns
619
+ -------
620
+ xr.Dataset
621
+ Dataset with calculated rates, intensities, and uncertainties for the
622
+ specified species.
63
623
  """
64
- Project the pointing set to a rectangular sky map.
624
+ # Step 1: Calculate rates for the specified species
625
+ dataset = calculate_rates(dataset)
626
+
627
+ # Step 2: Calculate intensities
628
+ dataset = calculate_intensities(dataset)
629
+
630
+ # Step 3: Calculate background rates and intensities
631
+ dataset = calculate_backgrounds(dataset)
632
+
633
+ # Optional Step 4: Calculate sputtering corrections
634
+ if sputtering_correction:
635
+ # TODO: The second dataset is for Oxygen specifically,
636
+ # if we get an H dataset in, we may need to calculate
637
+ # the O dataset separately before calling here.
638
+ dataset = calculate_sputtering_corrections(dataset, dataset)
639
+
640
+ # Optional Step 5: Clean up intermediate variables
641
+ if bootstrap_correction:
642
+ dataset = calculate_bootstrap_corrections(dataset)
643
+
644
+ # Step 6: Clean up intermediate variables
645
+ dataset = cleanup_intermediate_variables(dataset)
646
+
647
+ return dataset
65
648
 
66
- This function is used to create a rectangular sky map from the pointing set
67
- data in the L1C dataset.
649
+
650
+ def calculate_rates(dataset: xr.Dataset) -> xr.Dataset:
651
+ """
652
+ Calculate count rates and their statistical uncertainties.
68
653
 
69
654
  Parameters
70
655
  ----------
71
- psets : list[xr.Dataset]
72
- List of pointing sets in xarray Dataset format.
73
- spacing_deg : int
74
- The spacing in degrees for the rectangular sky map.
75
- spice_frame : SpiceFrame
76
- The SPICE frame to use for the rectangular sky map projection.
656
+ dataset : xr.Dataset
657
+ Dataset with count data and exposure times.
77
658
 
78
659
  Returns
79
660
  -------
80
- RectangularSkyMap
81
- The rectangular sky map created from the pointing set data.
661
+ xr.Dataset
662
+ Dataset with calculated count rates and statistical uncertainties
663
+ for the specified species.
82
664
  """
83
- lo_rect_map = ena_maps.RectangularSkyMap(
84
- spacing_deg=spacing_deg,
85
- spice_frame=spice_frame,
665
+ # Rate = counts / exposure_factor
666
+ # TODO: Account for ena / isn naming differences
667
+ dataset["ena_count_rate"] = dataset["counts"] / dataset["exposure_factor"]
668
+
669
+ # Poisson uncertainty on the counts propagated to the rate
670
+ # TODO: Is there uncertainty in the exposure time too?
671
+ dataset["ena_count_rate_stat_uncert"] = (
672
+ np.sqrt(dataset["counts"]) / dataset["exposure_factor"]
86
673
  )
87
- for pset in psets:
88
- lo_pset = ena_maps.LoPointingSet(pset)
89
- lo_rect_map.project_pset_values_to_map(
90
- pointing_set=lo_pset,
91
- value_keys=["h_counts", "exposure_time"],
92
- index_match_method=ena_maps.IndexMatchMethod.PUSH,
93
- )
94
- return lo_rect_map
674
+
675
+ return dataset
95
676
 
96
677
 
97
- def calculate_rates(counts: xr.DataArray, exposure_time: xr.DataArray) -> xr.DataArray:
678
+ def calculate_intensities(dataset: xr.Dataset) -> xr.Dataset:
98
679
  """
99
- Calculate the hydrogen rates from the counts and exposure time.
680
+ Calculate particle intensities and uncertainties for the specified species.
100
681
 
101
682
  Parameters
102
683
  ----------
103
- counts : xr.DataArray
104
- The counts of hydrogen or oxygen ENAs.
105
- exposure_time : xr.DataArray
106
- The exposure time for the counts.
684
+ dataset : xr.Dataset
685
+ Dataset with count rates, geometric factors, and center energies.
107
686
 
108
687
  Returns
109
688
  -------
110
- xr.DataArray
111
- The calculated hydrogen rates.
689
+ xr.Dataset
690
+ Dataset with calculated particle intensities and their statistical
691
+ and systematic uncertainties for the specified species.
112
692
  """
113
- # Calculate the rates based on the h_counts and exposure_time
114
- rate = counts / exposure_time
115
- return rate
693
+ # Equation 3 from mapping document (average intensity)
694
+ dataset["ena_intensity"] = dataset["counts_over_eff"] / (
695
+ dataset["geometric_factor"] * dataset["energy"] * dataset["exposure_factor"]
696
+ )
116
697
 
698
+ # Equation 4 from mapping document (statistical uncertainty)
699
+ # Note that we need to take the square root to get the uncertainty as
700
+ # the equation is for the variance
701
+ dataset["ena_intensity_stat_uncert"] = np.sqrt(
702
+ dataset["counts_over_eff_squared"]
703
+ / (dataset["geometric_factor"] * dataset["energy"] * dataset["exposure_factor"])
704
+ )
117
705
 
118
- def calculate_fluxes(rates: xr.DataArray) -> xr.DataArray:
706
+ # Equation 5 from mapping document (systematic uncertainty)
707
+ dataset["ena_intensity_sys_err"] = (
708
+ dataset["ena_intensity"]
709
+ * dataset["geometric_factor_stat_uncert"]
710
+ / dataset["geometric_factor"]
711
+ )
712
+
713
+ return dataset
714
+
715
+
716
+ def calculate_backgrounds(dataset: xr.Dataset) -> xr.Dataset:
119
717
  """
120
- Calculate the flux from the hydrogen rate.
718
+ Calculate background rates and intensities for the specified species.
121
719
 
122
720
  Parameters
123
721
  ----------
124
- rates : xr.Dataset
125
- The hydrogen or oxygen rates.
722
+ dataset : xr.Dataset
723
+ Dataset with count rates, geometric factors, and center energies.
126
724
 
127
725
  Returns
128
726
  -------
129
- xr.DataArray
130
- The calculated flux.
727
+ xr.Dataset
728
+ Dataset with calculated background rates and intensities for the
729
+ specified species.
131
730
  """
132
- # Temporary values. These will all come from ancillary data when
133
- # the data is available and integrated.
134
- geometric_factor = 1.0
135
- efficiency_factor = 1.0
136
- energy_dict = {1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7}
137
- energies = np.array([energy_dict[i] for i in range(1, 8)])
138
- energies = energies.reshape(1, 7, 1)
731
+ # Equation 6 from mapping document (background rate)
732
+ # exposure time weighted average of the background rates
733
+ dataset["bg_rates"] = (
734
+ dataset["bg_rates_exposure_factor"] / dataset["exposure_factor"]
735
+ )
736
+ # Equation 7 from mapping document (background statistical uncertainty)
737
+ dataset["bg_rates_stat_uncert"] = np.sqrt(
738
+ dataset["bg_rates_stat_uncert_exposure_factor2"]
739
+ / dataset["exposure_factor"] ** 2
740
+ )
741
+ # Equation 8 from mapping document (background systematic uncertainty)
742
+ dataset["bg_rates_sys_err"] = (
743
+ dataset["bg_rates"]
744
+ * dataset["geometric_factor_stat_uncert"]
745
+ / dataset["geometric_factor"]
746
+ )
139
747
 
140
- flux = rates / (geometric_factor * energies * efficiency_factor)
141
- return flux
748
+ return dataset
142
749
 
143
750
 
144
- def add_attributes(
145
- lo_map: xr.Dataset, attr_mgr: ImapCdfAttributes, logical_source: str
751
+ def calculate_sputtering_corrections(
752
+ dataset: xr.Dataset, o_dataset: xr.Dataset
146
753
  ) -> xr.Dataset:
147
754
  """
148
- Add attributes to the map dataset.
755
+ Calculate sputtering corrections from oxygen intensities.
756
+
757
+ Only for Oxygen sputtering and correction only at ESA levels 5 and 6
758
+ for 90 degree maps. If off-angle maps are made, we may have to extend
759
+ this to levels 3 and 4 as well.
760
+
761
+ Follows equations 9-13 from the mapping document.
149
762
 
150
763
  Parameters
151
764
  ----------
152
- lo_map : xr.Dataset
153
- The dataset to add attributes to.
154
- attr_mgr : ImapCdfAttributes
155
- The attribute manager to use for adding attributes.
156
- logical_source : str
157
- The logical source for the dataset.
765
+ dataset : xr.Dataset
766
+ Dataset with count rates, geometric factors, and center energies.
767
+ This could be either an H or O dataset.
768
+ o_dataset : xr.Dataset
769
+ Dataset specifically for oxygen, needed to access oxygen intensities
770
+ and uncertainties.
158
771
 
159
772
  Returns
160
773
  -------
161
774
  xr.Dataset
162
- The dataset with added attributes.
163
- """
164
- # Add the global attributes to the dataset.
165
- lo_map.attrs.update(attr_mgr.get_global_attributes(logical_source))
166
-
167
- # TODO: Lo is using different field names than what's in the attributes.
168
- # check if the Lo should use exposure factor instead of exposure time.
169
- # check if hydrogen and oxygen specific ena intensities should be added
170
- # to the attributes or if general ena intensities can be used or updated
171
- # in the code. This dictionary is temporary solution for SIT-4
172
- map_fields = {
173
- "epoch": "epoch",
174
- "h_flux": "ena_intensity",
175
- "h_rate": "ena_rate",
176
- "h_counts": "ena_count",
177
- "exposure_time": "exposure_factor",
178
- "energy": "energy",
179
- "solid_angle": "solid_angle",
180
- "longitude": "longitude",
181
- "latitude": "latitude",
182
- }
775
+ Dataset with calculated sputtering-corrected intensities and their
776
+ uncertainties for hydrogen and oxygen.
777
+ """
778
+ logger.info("Applying sputtering corrections to oxygen intensities")
779
+ # Only apply sputtering correction to esa levels 5 and 6 (indices 4 and 5)
780
+ energy_indices = [4, 5]
781
+ small_dataset = dataset.isel(epoch=0, energy=energy_indices)
782
+ o_small_dataset = o_dataset.isel(epoch=0, energy=energy_indices)
183
783
 
184
- # TODO: The mapping utility is supposed to handle at least some of these
185
- # attributes but is not working. Need to investigate this after SIT-4
186
- # Add the attributes to the dataset variables.
187
- for field, attr_name in map_fields.items():
188
- if field in lo_map.data_vars or field in lo_map.coords:
189
- lo_map[field].attrs.update(
190
- attr_mgr.get_variable_attributes(attr_name, check_schema=False)
191
- )
784
+ # NOTE: We only have background rates, so turn them into intensities
785
+ o_small_dataset["bg_intensity"] = o_small_dataset["bg_rates"] / (
786
+ o_small_dataset["geometric_factor"] * o_small_dataset["energy"]
787
+ )
788
+ o_small_dataset["bg_intensity_stat_uncert"] = o_small_dataset[
789
+ "bg_rates_stat_uncert"
790
+ ] / (o_small_dataset["geometric_factor"] * o_small_dataset["energy"])
192
791
 
193
- labels = {
194
- "energy": np.arange(1, 8).astype(str),
195
- "longitude": lo_map["longitude"].values.astype(str),
196
- "latitude": lo_map["latitude"].values.astype(str),
197
- }
198
- # add the coordinate labels to the dataset
199
- for dim, values in labels.items():
200
- lo_map = lo_map.assign_coords(
201
- {
202
- f"{dim}_label": xr.DataArray(
203
- values,
204
- name=f"{dim}_label",
205
- dims=[dim],
206
- attrs=attr_mgr.get_variable_attributes(
207
- f"{dim}_label", check_schema=False
208
- ),
209
- )
210
- }
792
+ # Equation 9
793
+ j_o_prime = o_small_dataset["ena_intensity"] - o_small_dataset["bg_intensity"]
794
+ j_o_prime.values[j_o_prime.values < 0] = 0 # No negative intensities
795
+
796
+ # Equation 10
797
+ j_o_prime_var = (
798
+ o_small_dataset["ena_intensity_stat_uncert"] ** 2
799
+ + o_small_dataset["bg_intensity_stat_uncert"] ** 2
800
+ )
801
+
802
+ # NOTE: From table 2 of the mapping document, for energy level 5 and 6
803
+ sputter_correction_factor = xr.DataArray(
804
+ [0.15, 0.01], dims=["energy"], coords={"energy": energy_indices}
805
+ )
806
+ # Equation 11
807
+ # Remove the sputtered oxygen intensity to correct the original O intensity
808
+ sputter_corrected_intensity = (
809
+ small_dataset["ena_intensity"] - sputter_correction_factor * j_o_prime
810
+ )
811
+
812
+ # Equation 12
813
+ sputter_corrected_intensity_var = (
814
+ small_dataset["ena_intensity_stat_uncert"] ** 2
815
+ + (sputter_correction_factor**2) * j_o_prime_var
816
+ )
817
+
818
+ # Equation 13
819
+ sputter_corrected_intensity_sys_err = (
820
+ sputter_corrected_intensity
821
+ / small_dataset["ena_intensity"]
822
+ * small_dataset["ena_intensity_sys_err"]
823
+ )
824
+
825
+ # Now put the corrected values into the original dataset
826
+ dataset["ena_intensity"][0, energy_indices, ...] = sputter_corrected_intensity
827
+ dataset["ena_intensity_stat_uncert"][0, energy_indices, ...] = np.sqrt(
828
+ sputter_corrected_intensity_var
829
+ )
830
+ dataset["ena_intensity_sys_err"][0, energy_indices, ...] = (
831
+ sputter_corrected_intensity_sys_err
832
+ )
833
+
834
+ return dataset
835
+
836
+
837
+ def calculate_bootstrap_corrections(dataset: xr.Dataset) -> xr.Dataset:
838
+ """
839
+ Calculate bootstrap corrections for hydrogen and oxygen intensities.
840
+
841
+ Follows equations 14-35 from the mapping document.
842
+
843
+ Parameters
844
+ ----------
845
+ dataset : xr.Dataset
846
+ Dataset with count rates, geometric factors, and center energies.
847
+
848
+ Returns
849
+ -------
850
+ xr.Dataset
851
+ Dataset with calculated bootstrap-corrected intensities and their
852
+ uncertainties for hydrogen.
853
+ """
854
+ logger.info("Applying bootstrap corrections")
855
+
856
+ # Table 3 bootstrap terms h_i,k
857
+ bootstrap_factor = np.array(
858
+ [
859
+ [0, 0.03, 0.01, 0, 0, 0, 0, 0],
860
+ [0, 0, 0.05, 0.02, 0.01, 0, 0, 0],
861
+ [0, 0, 0, 0.09, 0.03, 0.016, 0.01, 0],
862
+ [0, 0, 0, 0, 0.16, 0.068, 0.016, 0.01],
863
+ [0, 0, 0, 0, 0, 0.29, 0.068, 0.016],
864
+ [0, 0, 0, 0, 0, 0, 0.52, 0.061],
865
+ [0, 0, 0, 0, 0, 0, 0, 0.75],
866
+ ]
867
+ )
868
+
869
+ # Equation 14
870
+ bg_intensity = dataset["bg_rates"] / (
871
+ dataset["geometric_factor"] * dataset["energy"]
872
+ )
873
+ j_c_prime = dataset["ena_intensity"] - bg_intensity
874
+ j_c_prime.values[j_c_prime.values < 0] = 0
875
+
876
+ # Equation 15
877
+ j_c_prime_var = dataset["ena_intensity_stat_uncert"] ** 2
878
+
879
+ # Equation 16 - systematic error propagation
880
+ # Handle division by zero: only compute where ena_intensity > 0
881
+ j_c_prime_err = xr.where(
882
+ dataset["ena_intensity"] > 0,
883
+ j_c_prime / dataset["ena_intensity"] * dataset["ena_intensity_sys_err"],
884
+ 0,
885
+ )
886
+
887
+ # NOTE: E8 virtual channel calculation is from the text. This is to
888
+ # start the calculations off from the higher energies and avoid
889
+ # reliance on IMAP Hi energy channels.
890
+ # E8 is a virtual energy channel at 2.1 * E7
891
+ e8 = 2.1 * dataset["energy"].values[-1]
892
+
893
+ j_c_6 = j_c_prime.isel(energy=5)
894
+ j_c_7 = j_c_prime.isel(energy=6)
895
+ e_6 = dataset["energy"].isel(energy=5)
896
+ e_7 = dataset["energy"].isel(energy=6)
897
+
898
+ # Calculate gamma, ignoring any invalid values
899
+ # Fill in the invalid values with zeros after the fact
900
+ with np.errstate(divide="ignore", invalid="ignore"):
901
+ gamma = np.log(j_c_6 / j_c_7) / np.log(e_6 / e_7)
902
+ j_8_b = j_c_7 * (e8 / e_7) ** gamma
903
+
904
+ # Set j_8_b to zero where the calculation was invalid
905
+ j_8_b = j_8_b.where(np.isfinite(j_8_b) & (j_8_b > 0), 0)
906
+
907
+ # Initialize bootstrap intensity and uncertainty arrays
908
+ dataset["bootstrap_intensity"] = xr.zeros_like(dataset["ena_intensity"])
909
+ dataset["bootstrap_intensity_var"] = xr.zeros_like(dataset["ena_intensity"])
910
+ dataset["bootstrap_intensity_sys_err"] = xr.zeros_like(dataset["ena_intensity"])
911
+
912
+ for i in range(6, -1, -1):
913
+ # Initialize the variable with the non-summation term and virtual
914
+ # channel energy subtraction first, then iterate through the other
915
+ # channels which can be looked up via indexing
916
+ # i.e. the summation is always k=i+1 to 7, because we've already
917
+ # included the k=8 term here.
918
+ # NOTE: The paper uses 1-based indexing and we use 0-based indexing
919
+ # so there is an off-by-one difference in the indices.
920
+ dataset["bootstrap_intensity"][0, i, ...] = (
921
+ j_c_prime[0, i, ...] - bootstrap_factor[i, 7] * j_8_b[0, ...]
211
922
  )
923
+ # NOTE: We will square root at the end to get the uncertainty, but
924
+ # all equations are with variances
925
+ dataset["bootstrap_intensity_var"][0, i, ...] = j_c_prime_var[0, i, ...]
926
+
927
+ for k in range(i + 1, 7):
928
+ logger.debug(
929
+ f"Subtracting bootstrap factor h_{i},{k} * J_{k}_b from J_{i}_b"
930
+ )
931
+ # Subtraction terms from equations 18-23
932
+ dataset["bootstrap_intensity"][0, i, ...] -= (
933
+ bootstrap_factor[i, k] * dataset["bootstrap_intensity"][0, k, ...]
934
+ )
935
+
936
+ # Summation terms from equations 25-30
937
+ dataset["bootstrap_intensity_var"][0, i, ...] += (
938
+ bootstrap_factor[i, k] ** 2
939
+ ) * dataset["bootstrap_intensity_var"][0, k, ...]
940
+
941
+ # Again zero any bootstrap fluxes that are negative
942
+ dataset["bootstrap_intensity"][0, i, ...].values[
943
+ dataset["bootstrap_intensity"][0, i, ...] < 0
944
+ ] = 0.0
945
+
946
+ # Equation 31 - systematic error propagation for bootstrap intensity
947
+ # Handle division by zero: only compute where j_c_prime > 0
948
+ dataset["bootstrap_intensity_sys_err"] = xr.where(
949
+ j_c_prime > 0, dataset["bootstrap_intensity"] / j_c_prime * j_c_prime_err, 0
950
+ )
951
+
952
+ # Update the original intensity values
953
+ # Equation 32 / 33
954
+ # ena_intensity = ena_intensity (J_c) - (j_c_prime - J_b)
955
+ dataset["ena_intensity"] -= j_c_prime - dataset["bootstrap_intensity"]
956
+
957
+ # Ensure corrected intensities are non-negative
958
+ dataset["ena_intensity"] = dataset["ena_intensity"].where(
959
+ dataset["ena_intensity"] >= 0, 0
960
+ )
961
+
962
+ # Equation 34 - statistical uncertainty
963
+ # Take the square root, since we were in variances up to this point
964
+ dataset["ena_intensity_stat_uncert"] = np.sqrt(dataset["bootstrap_intensity_var"])
965
+
966
+ # Equation 35 - systematic error for corrected intensity
967
+ # Handle division by zero and ensure reasonable values
968
+ dataset["ena_intensity_sys_err"] = xr.zeros_like(dataset["ena_intensity"])
969
+ valid_bootstrap = (dataset["bootstrap_intensity"] > 0) & np.isfinite(
970
+ dataset["bootstrap_intensity"]
971
+ )
972
+
973
+ # Only compute where bootstrap intensity is valid
974
+ dataset["ena_intensity_sys_err"] = xr.where(
975
+ valid_bootstrap,
976
+ (
977
+ dataset["ena_intensity"]
978
+ / dataset["bootstrap_intensity"]
979
+ * dataset["bootstrap_intensity_sys_err"]
980
+ ),
981
+ 0,
982
+ )
983
+
984
+ # Drop the intermediate bootstrap variables
985
+ dataset = dataset.drop_vars(
986
+ [
987
+ "bootstrap_intensity",
988
+ "bootstrap_intensity_var",
989
+ "bootstrap_intensity_sys_err",
990
+ ]
991
+ )
992
+
993
+ return dataset
994
+
995
+
996
+ def cleanup_intermediate_variables(dataset: xr.Dataset) -> xr.Dataset:
997
+ """
998
+ Remove intermediate variables that were only needed for calculations.
999
+
1000
+ Parameters
1001
+ ----------
1002
+ dataset : xr.Dataset
1003
+ Dataset containing intermediate calculation variables.
1004
+
1005
+ Returns
1006
+ -------
1007
+ xr.Dataset
1008
+ Cleaned dataset with intermediate variables removed.
1009
+ """
1010
+ # Remove the intermediate variables from the map
1011
+ # i.e. the ones that were projected from the pset only for the purposes
1012
+ # of math and not desired in the output.
1013
+ vars_to_remove = []
1014
+
1015
+ # Only remove variables that exist in the dataset for the specific species
1016
+ potential_vars = [
1017
+ "counts_over_eff",
1018
+ "counts_over_eff_squared",
1019
+ "bg_rates_exposure_factor",
1020
+ "bg_rates_stat_uncert_exposure_factor2",
1021
+ ]
1022
+
1023
+ for potential_var in potential_vars:
1024
+ if potential_var in dataset.data_vars:
1025
+ vars_to_remove.append(potential_var)
212
1026
 
213
- return lo_map
1027
+ return dataset.drop_vars(vars_to_remove)