imap-processing 0.11.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

imap_processing/swe/l2/swe_l2.py

@@ -11,57 +11,10 @@ import xarray as xr
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 from imap_processing.spice.geometry import SpiceFrame
 from imap_processing.spice.spin import get_instrument_spin_phase, get_spin_angle
-from imap_processing.swe.utils.swe_utils import (
-    ESA_VOLTAGE_ROW_INDEX_DICT,
-    read_lookup_table,
-)
-
-# TODO: add these to instrument status summary
-ENERGY_CONVERSION_FACTOR = 4.75
-# 7 CEMs geometric factors in cm^2 sr eV/eV units.
-GEOMETRIC_FACTORS = np.array(
-    [
-        435e-6,
-        599e-6,
-        808e-6,
-        781e-6,
-        876e-6,
-        548e-6,
-        432e-6,
-    ]
-)
-ELECTRON_MASS = 9.10938356e-31  # kg
-
-# See doc string of calculate_phase_space_density() for more details.
-VELOCITY_CONVERSION_FACTOR = 1.237e31
-# See doc string of calculate_flux() for more details.
-FLUX_CONVERSION_FACTOR = 6.187e30
-
-CEM_DETECTORS_ANGLE = np.array([-63, -42, -21, 0, 21, 42, 63])
-
-
-def get_particle_energy() -> npt.NDArray:
-    """
-    Get particle energy.
-
-    Calculate particle energy and add to the lookup table.
-    To convert Volts to Energy, multiply ESA voltage in Volts by
-    energy conversion factor to get electron energy in eV.
-
-    Returns
-    -------
-    lookup_table : pandas.DataFrame
-        Lookup table with energy column added.
-    """
-    # The lookup table gives voltage applied to analyzers.
-    lookup_table = read_lookup_table()
-
-    # Convert voltage to electron energy in eV by apply conversion factor.
-    lookup_table["energy"] = lookup_table["esa_v"].values * ENERGY_CONVERSION_FACTOR
-    return lookup_table
+from imap_processing.swe.utils import swe_constants
 
 
-def calculate_phase_space_density(l1b_dataset: xr.Dataset) -> xr.Dataset:
+def calculate_phase_space_density(l1b_dataset: xr.Dataset) -> npt.NDArray:
     """
     Convert counts to phase space density.
 
@@ -73,7 +26,7 @@ def calculate_phase_space_density(l1b_dataset: xr.Dataset) -> xr.Dataset:
     Where:
         C / tau = corrected count rate which in the input L1B science data.
         G = geometric factor, in (cm^2 * ster). 7 CEMs geometric factor value.
-        eV = eV in electron-volts, calculated by get_particle_energy().
+        eV = eV in electron-volts.
         E = Energy in Joules. eV * 1.60219e-19(J/eV).
         m = mass of electron (9.10938356e-31 kg).
         s = second.
@@ -101,53 +54,29 @@ def calculate_phase_space_density(l1b_dataset: xr.Dataset) -> xr.Dataset:
 
     Returns
     -------
-    phase_space_density_dataset : xarray.Dataset
+    phase_space_density : np.ndarray
         Phase space density. We need to call this phase space density because
         there will be density in L3 processing.
     """
-    # Get esa_table_num for each full sweep.
-    esa_table_nums = l1b_dataset["esa_table_num"].values[:, 0]
-    # Get energy values from lookup table.
-    particle_energy = get_particle_energy()
-    # Get 720 (24 energy steps x 30 angle) particle energy for each full
-    # sweep data.
-    particle_energy_data = np.array(
-        [
-            particle_energy[particle_energy["table_index"] == val]["energy"].tolist()
-            for val in esa_table_nums
-        ]
-    )
-    particle_energy_data = particle_energy_data.reshape(-1, 24, 30)
+    # Get energy values.
+    particle_energy_data = l1b_dataset["esa_energy"].values
 
     # Calculate phase space density using formula:
     #   2 * (C/tau) / (G * 1.237e31 * eV^2)
     # See doc string for more details.
     density = (2 * l1b_dataset["science_data"]) / (
-        GEOMETRIC_FACTORS[np.newaxis, np.newaxis, np.newaxis, :]
-        * VELOCITY_CONVERSION_FACTOR
+        swe_constants.GEOMETRIC_FACTORS[np.newaxis, np.newaxis, np.newaxis, :]
+        * swe_constants.VELOCITY_CONVERSION_FACTOR
         * particle_energy_data[:, :, :, np.newaxis] ** 2
     )
+    phase_space_density = density.data
 
-    # Return density as xr.dataset with phase space density and
-    # energy in eV value that flux calculation can use.
-    phase_space_density_dataset = xr.Dataset(
-        {
-            "phase_space_density": (
-                ["epoch", "esa_step", "spin_sector", "cem_id"],
-                density.data,
-            ),
-            "energy_in_eV": (
-                ["epoch", "esa_step", "spin_sector"],
-                particle_energy_data,
-            ),
-        },
-        coords=l1b_dataset.coords,
-    )
+    return phase_space_density
 
-    return phase_space_density_dataset
 
-
-def calculate_flux(l1b_dataset: xr.Dataset) -> npt.NDArray:
+def calculate_flux(
+    phase_space_density: np.ndarray, esa_energy: np.ndarray
+) -> npt.NDArray:
     """
     Calculate flux.
 
@@ -159,7 +88,7 @@ def calculate_flux(l1b_dataset: xr.Dataset) -> npt.NDArray:
     Where:
         fv = the phase space density of solar wind electrons
             given by calculate_phase_space_density() result.
-        eV = Energy in electron-volts, calculated by get_particle_energy().
+        eV = Energy in electron-volts.
         E  = Energy in Joules. eV * 1.60219e-19(J/eV).
         v  = sqrt( (3.20438 * 10e-15 / 9.10938e-31) * eV ) cm/s. See
             calculate_phase_space_density() for this calculation.
@@ -184,19 +113,20 @@ def calculate_flux(l1b_dataset: xr.Dataset) -> npt.NDArray:
 
     Parameters
     ----------
-    l1b_dataset : xarray.Dataset
-        The L1B dataset to process.
+    phase_space_density : numpy.ndarray
+        The phase space density.
+    esa_energy : numpy.ndarray
+        The energy values in eV.
 
     Returns
     -------
     flux : numpy.ndarray
         Flux values.
     """
-    phase_space_density_ds = calculate_phase_space_density(l1b_dataset)
     flux = (
-        FLUX_CONVERSION_FACTOR
-        * phase_space_density_ds["energy_in_eV"].data[:, :, :, np.newaxis]
-        * phase_space_density_ds["phase_space_density"].data
+        swe_constants.FLUX_CONVERSION_FACTOR
+        * esa_energy[:, :, :, np.newaxis]
+        * phase_space_density
     )
     return flux
 
@@ -220,22 +150,32 @@ def put_data_into_angle_bins(
     Parameters
     ----------
     data : numpy.ndarray
-        Data to put in bins. Shape: (full_cycle_data, energy_step, angle_bin, 7).
+        Data to put in bins. Shape:
+        (full_cycle_data, N_ESA_STEPS, N_ANGLE_BINS, N_CEMS).
     angle_bin_indices : numpy.ndarray
         Indices of angle bins to put data in. Shape:
-        (full_cycle_data, energy_step, angle_bin).
+        (full_cycle_data, N_ESA_STEPS, N_ANGLE_BINS).
 
     Returns
     -------
     numpy.ndarray
-        Data in bins. Shape: (full_cycle_data, 24, 30, 7).
+        Data in bins. Shape:
+        (full_cycle_data, N_ESA_STEPS, N_ANGLE_BINS, N_CEMS).
     """
     # Initialize with zeros instead of NaN because np.add.at() does not
     # work with nan values. It results in nan + value = nan
-    binned_data = np.zeros((data.shape[0], 24, 30, 7), dtype=np.float64)
+    binned_data = np.zeros(
+        (
+            data.shape[0],
+            swe_constants.N_ESA_STEPS,
+            swe_constants.N_ANGLE_BINS,
+            swe_constants.N_CEMS,
+        ),
+        dtype=np.float64,
+    )
 
     time_indices = np.arange(data.shape[0])[:, None, None]
-    energy_indices = np.arange(24)[None, :, None]
+    energy_indices = np.arange(swe_constants.N_ESA_STEPS)[None, :, None]
 
     # Use np.add.at() to accumulate values into bins
     np.add.at(binned_data, (time_indices, energy_indices, angle_bin_indices), data)
@@ -312,7 +252,7 @@ def find_angle_bin_indices(
     return spin_angle_bins_indices
 
 
-def swe_l2(l1b_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
+def swe_l2(l1b_dataset: xr.Dataset) -> xr.Dataset:
     """
     Will process data to L2.
 
@@ -320,8 +260,6 @@ def swe_l2(l1b_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
     ----------
     l1b_dataset : xarray.Dataset
         The L1B dataset to process.
-    data_version : str
-        Version of the data product being created.
 
     Returns
     -------
@@ -331,49 +269,55 @@ def swe_l2(l1b_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
     cdf_attributes = ImapCdfAttributes()
     cdf_attributes.add_instrument_global_attrs("swe")
     cdf_attributes.add_instrument_variable_attrs("swe", "l2")
-    cdf_attributes.add_global_attribute("Data_version", data_version)
 
     # Energy values in eV.
     energy_xr = xr.DataArray(
-        np.array(list(ESA_VOLTAGE_ROW_INDEX_DICT.keys())) * ENERGY_CONVERSION_FACTOR,
+        np.array(list(swe_constants.ESA_VOLTAGE_ROW_INDEX_DICT.keys()))
+        * swe_constants.ENERGY_CONVERSION_FACTOR,
         name="energy",
         dims=["energy"],
-        attrs=cdf_attributes.get_variable_attributes("energy"),
+        attrs=cdf_attributes.get_variable_attributes("energy", check_schema=False),
     )
 
     energy_label = xr.DataArray(
-        np.array(list(ESA_VOLTAGE_ROW_INDEX_DICT.keys())).astype(str),
+        np.array(list(swe_constants.ESA_VOLTAGE_ROW_INDEX_DICT.keys())).astype(str),
         name="energy_label",
         dims=["energy"],
-        attrs=cdf_attributes.get_variable_attributes("energy_label"),
+        attrs=cdf_attributes.get_variable_attributes(
+            "energy_label", check_schema=False
+        ),
     )
 
     # Angle of each CEM detectors.
     inst_el_xr = xr.DataArray(
-        CEM_DETECTORS_ANGLE,
+        swe_constants.CEM_DETECTORS_ANGLE,
         name="inst_el",
         dims=["inst_el"],
-        attrs=cdf_attributes.get_variable_attributes("inst_el"),
+        attrs=cdf_attributes.get_variable_attributes("inst_el", check_schema=False),
     )
     inst_el_label = xr.DataArray(
-        CEM_DETECTORS_ANGLE.astype(str),
+        swe_constants.CEM_DETECTORS_ANGLE.astype(str),
         name="inst_el_label",
         dims=["inst_el"],
-        attrs=cdf_attributes.get_variable_attributes("inst_el_label"),
+        attrs=cdf_attributes.get_variable_attributes(
+            "inst_el_label", check_schema=False
+        ),
     )
 
     # Spin Angle bins storing bin center values.
     inst_az_xr = xr.DataArray(
-        np.arange(0, 360, 12) + 6,
+        np.arange(0, 360, 12, dtype=np.float32) + 6,
         name="inst_az",
         dims=["inst_az"],
-        attrs=cdf_attributes.get_variable_attributes("inst_az"),
+        attrs=cdf_attributes.get_variable_attributes("inst_az", check_schema=False),
     )
     inst_az_label = xr.DataArray(
         inst_az_xr.values.astype(str),
         name="inst_az_label",
         dims=["inst_az"],
-        attrs=cdf_attributes.get_variable_attributes("inst_az_label"),
+        attrs=cdf_attributes.get_variable_attributes(
+            "inst_az_label", check_schema=False
+        ),
     )
 
     dataset = xr.Dataset(
@@ -399,9 +343,7 @@ def swe_l2(l1b_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
     # Calculate phase space density and flux. Store data in shape
     # (epoch, esa_step, spin_sector, cem_id). This is for L3 purposes.
     ############################################################
-    phase_space_density = calculate_phase_space_density(l1b_dataset)[
-        "phase_space_density"
-    ]
+    phase_space_density = calculate_phase_space_density(l1b_dataset)
     dataset["phase_space_density_spin_sector"] = xr.DataArray(
         phase_space_density,
         name="phase_space_density_spin_sector",
@@ -409,7 +351,7 @@ def swe_l2(l1b_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
         attrs=cdf_attributes.get_variable_attributes("phase_space_density_spin_sector"),
     )
 
-    flux = calculate_flux(l1b_dataset)
+    flux = calculate_flux(phase_space_density, l1b_dataset["esa_energy"].data)
     dataset["flux_spin_sector"] = xr.DataArray(
         flux,
         name="flux_spin_sector",
@@ -419,36 +361,33 @@ def swe_l2(l1b_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
 
     # Carry over acquisition times for L3 purposes.
     dataset["acquisition_time"] = l1b_dataset["acquisition_time"]
+    # Update the acquisition_time variable attributes.
+    dataset["acquisition_time"].attrs = cdf_attributes.get_variable_attributes(
+        "acquisition_time"
+    )
+    # Carry over acq_duration for L3 purposes.
+    dataset["acq_duration"] = l1b_dataset["acq_duration"]
+    # Update the acq_duration variable attributes.
+    dataset["acq_duration"].attrs = cdf_attributes.get_variable_attributes(
+        "acq_duration"
+    )
 
     # Calculate spin phase using SWE acquisition_time from the
     # L1B dataset. The L1B dataset stores acquisition_time with
-    # dimensions (epoch, esa_step, spin_sector). Use center time
-    # to calculate spin phase. This center time calculation is
-    # necessary to accurately determine the center angle of the data.
-    #
-    # To determine the center acquisition time, we adjust the
-    # recorded acquisition_time as follows:
-    #   acquisition_time + (acq_duration / 1000000) / 2
-    #
-    # Here, acq_duration is given in microseconds and is stored
-    # in the L1B dataset with dimensions (epoch, cycle). Since acq_duration
-    # remains the same for all quarter cycles within a full sweep,
-    # we use the first acq_duration value for each full sweep to perform
-    # this adjustment.
-
-    acq_duration = l1b_dataset["acq_duration"].data[:, 0] / 2000000
-    data_acq_time = (
-        l1b_dataset["acquisition_time"].data + acq_duration[:, np.newaxis, np.newaxis]
-    )
+    # dimensions (epoch, esa_step, spin_sector). acquisition_time is
+    # center time of acquisition time of each science measurement which
+    # is necessary to accurately determine the center angle of the data.
 
     # Calculate spin phase
     inst_spin_phase = get_instrument_spin_phase(
-        query_met_times=data_acq_time.ravel(),
+        query_met_times=l1b_dataset["acquisition_time"].data.ravel(),
         instrument=SpiceFrame.IMAP_SWE,
     )
 
     # Convert spin phase to spin angle in degrees.
-    inst_spin_angle = get_spin_angle(inst_spin_phase, degrees=True).reshape(-1, 24, 30)
+    inst_spin_angle = get_spin_angle(inst_spin_phase, degrees=True).reshape(
+        -1, swe_constants.N_ESA_STEPS, swe_constants.N_ANGLE_SECTORS
+    )
 
     # Save spin angle in dataset per SWE request.
     dataset["inst_az_spin_sector"] = xr.DataArray(
@@ -471,7 +410,7 @@ def swe_l2(l1b_dataset: xr.Dataset, data_version: str) -> xr.Dataset:
 
     # Put phase space density data in its spin angle bins using the indices.
     phase_space_density_binned_data = put_data_into_angle_bins(
-        phase_space_density.data, spin_angle_bins_indices
+        phase_space_density, spin_angle_bins_indices
     )
     dataset["phase_space_density"] = xr.DataArray(
         phase_space_density_binned_data,