I-ALiRT: Create SWAPI count rate optimization function

imap_processing/ialirt/constants.py

@@ -0,0 +1,38 @@
+"""Module for constants and useful shared classes used in I-ALiRT processing."""
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class IalirtSwapiConstants:
+    """
+    Constants for I-ALiRT SWAPI which can be used across different levels or classes.
+
+    Attributes
+    ----------
+    BOLTZ: float
+        Boltzmann constant [J/K]
+    AT_MASS: float
+        Atomic mass [kg]
+    PROT_MASS: float
+        Mass of proton [kg]
+    EFF_AREA: float
+        Instrument effective area [m^2]
+    AZ_FOV: float
+        Azimuthal width of the field of view for solar wind [radians]
+    FWHM_WIDTH: float
+        Full Width at Half Maximum of energy width [unitless]
+    SPEED_EW: float
+        Speed width of energy passband [unitless]
+    """
+
+    # Scientific constants used in optimization model
+    boltz = 1.380649e-23  # Boltzmann constant, J/K
+    at_mass = 1.6605390666e-27  # atomic mass, kg
+    prot_mass = 1.007276466621 * at_mass  # mass of proton, kg
+    eff_area = 3.3e-5 * 1e-4  # effective area, meters squared
+    az_fov = np.deg2rad(30)  # azimuthal width of the field of view, radians
+    fwhm_width = 0.085  # FWHM of energy width
+    speed_ew = 0.5 * fwhm_width  # speed width of energy passband

imap_processing/ialirt/l0/process_swapi.py

@@ -1,20 +1,144 @@
 """Functions to support I-ALiRT SWAPI processing."""
 
 import logging
+from decimal import Decimal
+from typing import Optional
 
 import numpy as np
+import pandas as pd
 import xarray as xr
-from xarray import DataArray
+from scipy.optimize import curve_fit
+from scipy.special import erf
 
+from imap_processing import imap_module_directory
+from imap_processing.ialirt.constants import IalirtSwapiConstants as Consts
 from imap_processing.ialirt.utils.grouping import find_groups
-
-# from imap_processing.swapi.l1.swapi_l1 import process_sweep_data
-# from imap_processing.swapi.l2.swapi_l2 import TIME_PER_BIN
+from imap_processing.ialirt.utils.time import calculate_time
+from imap_processing.spice.time import met_to_ttj2000ns, met_to_utc
+from imap_processing.swapi.l1.swapi_l1 import process_sweep_data
+from imap_processing.swapi.l2.swapi_l2 import TIME_PER_BIN
 
 logger = logging.getLogger(__name__)
 
 
-def process_swapi_ialirt(unpacked_data: xr.Dataset) -> dict[str, DataArray]:
+def count_rate(
+    energy_pass: float, speed: float, density: float, temp: float
+) -> float | np.ndarray:
+    """
+    Compute SWAPI count rate for provided energy passband, speed, density and temp.
+
+    This model for coincidence count rate was developed by the SWAPI instrument
+    science team, detailed on page 52 of the IMAP SWAPI Instrument Algorithms Document.
+
+    Parameters
+    ----------
+    energy_pass : float
+        Energy passband [eV].
+    speed : float
+        Bulk solar wind speed [km/s].
+    density : float
+        Proton density [cm^-3].
+    temp : float
+        Temperature [K].
+
+    Returns
+    -------
+    count_rate : float | np.ndarray
+        Particle coincidence count rate.
+    """
+    # thermal velocity of solar wind ions
+    thermal_velocity = np.sqrt(2 * Consts.boltz * temp / Consts.prot_mass)
+    beta = 1 / (thermal_velocity**2)
+    # convert energy to Joules
+    center_speed = np.sqrt(2 * energy_pass * 1.60218e-19 / Consts.prot_mass)
+    speed = speed * 1000  # convert km/s to m/s
+    density = density * 1e6  # convert 1/cm**3 to 1/m**3
+
+    return (
+        (density * Consts.eff_area * (beta / np.pi) ** (3 / 2))
+        * (np.exp(-beta * (center_speed**2 + speed**2 - 2 * center_speed * speed)))
+        * np.sqrt(np.pi / (beta * speed * center_speed))
+        * erf(np.sqrt(beta * speed * center_speed) * (Consts.az_fov / 2))
+        * (
+            center_speed**4
+            * Consts.speed_ew
+            * np.arcsin(thermal_velocity / center_speed)
+        )
+    )
+
+
+def optimize_pseudo_parameters(
+    count_rates: np.ndarray,
+    count_rate_error: np.ndarray,
+    energy_passbands: Optional[np.ndarray] = None,
+) -> (dict)[str, list[float]]:
+    """
+    Find the pseudo speed (u), density (n) and temperature (T) of solar wind particles.
+
+    Fit a curve to calculated count rate values as a function of energy passband.
+
+    Parameters
+    ----------
+    count_rates : np.ndarray
+        Particle coincidence count rates.
+    count_rate_error : np.ndarray
+        Standard deviation of the coincidence count rates parameter.
+    energy_passbands : np.ndarray, default None
+        Energy passbands, passed in only for testing purposes.
+
+    Returns
+    -------
+    solution_dict : dict
+        Dictionary containing the optimized speed, density, and temperature values for
+        each sweep included in the input count_rates array.
+    """
+    if not energy_passbands:
+        # Read in energy passbands
+        energy_data = pd.read_csv(
+            f"{imap_module_directory}/tests/swapi/lut/imap_swapi_esa-unit"
+            f"-conversion_20250211_v000.csv"
+        )
+        energy_passbands = (
+            energy_data["Energy"][0:63]
+            .replace(",", "", regex=True)
+            .to_numpy()
+            .astype(float)
+        )
+
+    # Initial guess pulled from page 52 of the IMAP SWAPI Instrument Algorithms Document
+    initial_param_guess = np.array([550, 5.27, 1e5])
+    solution_dict = {  # type: ignore
+        "pseudo_speed": [],
+        "pseudo_density": [],
+        "pseudo_temperature": [],
+    }
+
+    for sweep in np.arange(count_rates.shape[0]):
+        current_sweep_count_rates = count_rates[sweep, :]
+        current_sweep_count_rate_errors = count_rate_error[sweep, :]
+        # Find the max count rate, and use the 6 points surrounding it (inclusive)
+        max_index = np.argmax(current_sweep_count_rates)
+        sol = curve_fit(
+            f=count_rate,
+            xdata=energy_passbands.take(
+                range(max_index - 3, max_index + 3), mode="wrap"
+            ),
+            ydata=current_sweep_count_rates.take(
+                range(max_index - 3, max_index + 3), mode="wrap"
+            ),
+            sigma=current_sweep_count_rate_errors.take(
+                range(max_index - 3, max_index + 3), mode="wrap"
+            ),
+            p0=initial_param_guess,
+        )
+        solution_dict["pseudo_speed"].append(sol[0][0])
+        solution_dict["pseudo_density"].append(sol[0][1])
+        solution_dict["pseudo_temperature"].append(sol[0][2])
+
+    return solution_dict
+
+
+def process_swapi_ialirt(unpacked_data: xr.Dataset) -> list[dict]:
     """
     Extract I-ALiRT variables and calculate coincidence count rate.
 
@@ -32,14 +156,32 @@ def process_swapi_ialirt(unpacked_data: xr.Dataset) -> dict[str, DataArray]:
 
     sci_dataset = unpacked_data.sortby("epoch", ascending=True)
 
-    grouped_dataset = find_groups(sci_dataset, (0, 11), "swapi_seq_number", "swapi_acq")
+    met = calculate_time(
+        sci_dataset["sc_sclk_sec"], sci_dataset["sc_sclk_sub_sec"], 256
+    )
+
+    # Add required parameters.
+    sci_dataset["met"] = met
+    met_values = []
+
+    grouped_dataset = find_groups(sci_dataset, (0, 11), "swapi_seq_number", "met")
+
+    if grouped_dataset.group.size == 0:
+        logger.warning(
+            "There was an issue with the SWAPI grouping process, returning empty data."
+        )
+        return [{}]
 
     for group in np.unique(grouped_dataset["group"]):
         # Sequence values for the group should be 0-11 with no duplicates.
         seq_values = grouped_dataset["swapi_seq_number"][
             (grouped_dataset["group"] == group)
         ]
 
+        met_values.append(
+            int(grouped_dataset["met"][(grouped_dataset["group"] == group).values][0])
+        )
+
         # Ensure no duplicates and all values from 0 to 11 are present
         if not np.array_equal(seq_values.astype(int), np.arange(12)):
             logger.info(
@@ -48,22 +190,29 @@ def process_swapi_ialirt(unpacked_data: xr.Dataset) -> dict[str, DataArray]:
             )
             continue
 
-    total_packets = len(grouped_dataset["swapi_seq_number"].data)
-
-    # It takes 12 sequence data to make one full SWAPI sweep
-    total_sequence = 12
-    total_full_sweeps = total_packets // total_sequence
-
-    met_values = grouped_dataset["swapi_shcoarse"].data.reshape(total_full_sweeps, 12)[
-        :, 0
-    ]
-
-    # raw_coin_count = process_sweep_data(grouped_dataset, "coin_cnt")
-    # raw_coin_rate = raw_coin_count / TIME_PER_BIN
-
-    swapi_data = {
-        "met": met_values
-        # more variables to go here
-    }
+    raw_coin_count = process_sweep_data(grouped_dataset, "swapi_coin_cnt")
+    raw_coin_rate = raw_coin_count / TIME_PER_BIN
+    count_rate_error = np.sqrt(raw_coin_count) / TIME_PER_BIN
+
+    solution = optimize_pseudo_parameters(raw_coin_rate, count_rate_error)
+
+    swapi_data = []
+
+    for entry in np.arange(0, len(solution["pseudo_speed"])):
+        swapi_data.append(
+            {
+                "apid": 478,
+                "met": met_values[entry],
+                "utc": met_to_utc(met_values[entry]).split(".")[0],
+                "ttj2000ns": int(met_to_ttj2000ns(met_values[entry])),
+                "swapi_pseudo_proton_speed": Decimal(solution["pseudo_speed"][entry]),
+                "swapi_pseudo_proton_density": Decimal(
+                    solution["pseudo_density"][entry]
+                ),
+                "swapi_pseudo_proton_temperature": Decimal(
+                    solution["pseudo_temperature"][entry]
+                ),
+            }
+        )
 
     return swapi_data

imap_processing/tests/ialirt/test_data/ialirt_test_data.csv

@@ -0,0 +1,73 @@
+Energy [eV/q],Count Rates [Hz],Count Rates Error [Hz]
+19098.3579,0.00E+00,0
+19098.3579,0.00E+00,0
+17541.17689,0.00E+00,0
+16113.17733,0.00E+00,0
+14798.37998,0.00E+00,0
+13591.36578,0.00E+00,0
+12485.77704,0.00E+00,0
+11467.61829,0.00E+00,0
+10532.60822,0.00E+00,0
+9675.514168,0.00E+00,0
+8885.04638,0.00E+00,0
+8165.393845,0.00E+00,0
+7501.760233,0.00E+00,0
+6888.477149,0.00E+00,0
+6327.926581,0.00E+00,0
+5811.486083,0.00E+00,0
+5338.867546,0.00E+00,0
+4901.30318,0.00E+00,0
+4504.29887,6.00E+00,5.988024
+4138.38252,0.00E+00,0
+3800.760624,0.00E+00,0
+3490.866227,6.00E+00,5.988024
+3205.462334,2.40E+01,11.976048
+2944.699516,3.00E+01,13.389629
+2705.519228,1.20E+01,8.468345
+2485.023495,6.00E+00,5.988024
+2281.728846,6.00E+00,5.988024
+2094.335628,1.80E+01,10.371562
+1921.410538,2.58E+02,39.266099
+1764.61444,8.70E+02,72.105357
+1621.075258,1.37E+03,90.615245
+1489.379616,1.25E+03,86.567858
+1369.25523,6.72E+02,63.371289
+1257.562068,2.46E+02,38.342061
+1155.04315,3.00E+01,13.389629
+1061.325411,6.00E+00,5.988024
+974.875126,0.00E+00,0
+895.314188,0.00E+00,0
+822.018852,0.00E+00,0
+754.982368,0.00E+00,0
+693.547324,0.00E+00,0
+636.793361,0.00E+00,0
+584.81978,0.00E+00,0
+537.016673,0.00E+00,0
+493.208286,0.00E+00,0
+453.103315,0.00E+00,0
+416.133867,0.00E+00,0
+382.037059,0.00E+00,0
+350.921008,0.00E+00,0
+322.396008,0.00E+00,0
+296.176976,0.00E+00,0
+271.952917,0.00E+00,0
+249.936708,0.00E+00,0
+229.494886,0.00E+00,0
+210.757138,0.00E+00,0
+193.581931,0.00E+00,0
+177.766309,0.00E+00,0
+163.295895,0.00E+00,0
+150.015166,0.00E+00,0
+137.803904,0.00E+00,0
+126.579577,0.00E+00,0
+116.253172,0.00E+00,0
+106.797953,0.00E+00,0
+1764.61444,8.70E+02,72.105357
+1727.604298,9.18E+02,73.913169
+1691.370389,9.66E+02,75.720982
+1655.896431,1.01E+03,77.528795
+1621.166486,1.06E+03,79.336607
+1587.16495,1.11E+03,81.14442
+1553.876545,1.16E+03,82.952233
+1521.286315,1.21E+03,84.760045
+1489.379616,1.25E+03,86.567858