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,19 +1,121 @@
 """Functions to support I-ALiRT SWAPI processing."""
 
 import logging
+from typing import Optional
 
 import numpy as np
+import pandas as pd
 import xarray as xr
+from scipy.optimize import curve_fit
+from scipy.special import erf
 from xarray import DataArray
 
+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.swapi.l1.swapi_l1 import process_sweep_data
+from imap_processing.swapi.l2.swapi_l2 import TIME_PER_BIN
 
 logger = logging.getLogger(__name__)
 
 
+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, energy_passbands: Optional[np.ndarray] = None
+) -> (dict)[str, list[float]]:
+    """
+    Find the pseudo speed (u), density (n) and temperature (T) of solar wind.
+
+    Fit a curve to calculated count rate values as a function of energy passbands.
+
+    Parameters
+    ----------
+    count_rates : np.ndarray
+        Particle coincidence count rates.
+    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.any():
+        # Read in energy passbands
+        energy_data = pd.read_csv(
+            f"{imap_module_directory}/tests/ialirt/test_data/ialirt_test_data.csv"
+        )
+        energy_passbands = energy_data["Energy [eV/q]"].to_numpy()
+
+    # 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, :]
+        sol = curve_fit(
+            count_rate,
+            energy_passbands,
+            current_sweep_count_rates,
+            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) -> dict[str, DataArray]:
     """
     Extract I-ALiRT variables and calculate coincidence count rate.
@@ -58,12 +160,16 @@ def process_swapi_ialirt(unpacked_data: xr.Dataset) -> dict[str, DataArray]:
         :, 0
     ]
 
-    # raw_coin_count = process_sweep_data(grouped_dataset, "coin_cnt")
-    # raw_coin_rate = raw_coin_count / TIME_PER_BIN
+    raw_coin_count = process_sweep_data(grouped_dataset, "swapi_coin_cnt")
+    raw_coin_rate = raw_coin_count / TIME_PER_BIN
+
+    solution = optimize_pseudo_parameters(raw_coin_rate)
 
     swapi_data = {
-        "met": met_values
-        # more variables to go here
+        "met": met_values,
+        "pseudo_speed": solution["pseudo_speed"],
+        "pseudo_density": solution["pseudo_density"],
+        "pseudo_temperature": solution["pseudo_temperature"],
     }
 
     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

imap_processing/tests/ialirt/unit/test_process_swapi.py

@@ -3,7 +3,11 @@
 import pytest
 
 from imap_processing import imap_module_directory
-from imap_processing.ialirt.l0.process_swapi import process_swapi_ialirt
+from imap_processing.ialirt.l0.process_swapi import (
+    count_rate,
+    optimize_pseudo_parameters,
+    process_swapi_ialirt,
+)
 from imap_processing.utils import packet_file_to_datasets
 
 
@@ -87,7 +91,36 @@ def test_decom_packets(xarray_data, swapi_test_data):
 
 
 def test_process_swapi_ialirt(xarray_data):
-    """Placeholder test for the process_swapi_ialirt function."""
+    """Test that the process_swapi_ialirt function returns expected keys."""
 
     swapi_result = process_swapi_ialirt(xarray_data)
     assert swapi_result["met"] is not None
+    assert len(swapi_result["met"]) == len(swapi_result["pseudo_temperature"])
+    assert len(swapi_result["pseudo_density"]) == len(swapi_result["pseudo_speed"])
+
+
+def test_count_rate():
+    """Use random realistic values to test for expected output of count_rate."""
+    actual_result = count_rate(1370, *[550, 5.27, 1e5])
+    expected_result = 621.0028766348703
+    assert actual_result == expected_result
+
+
+def test_optimize_parameters(xarray_data):
+    """Test the optimize_pseudo_parameters function."""
+
+    energy_data = pd.read_csv(
+        f"{imap_module_directory}/tests/ialirt/test_data/ialirt_test_data.csv"
+    )
+    count_rates = energy_data["Count Rates [Hz]"].to_numpy()
+    count_rates = np.tile(count_rates, (2, 1))
+    result = optimize_pseudo_parameters(count_rates)
+
+    # Test values corresponding to this exact set and values of the test input.
+    expected_speed = [542.9522302014949, 542.9522302014949]
+    expected_density = [4.504282147321004, 4.504282147321004]
+    expected_temperature = [143238.45841298936, 143238.45841298936]
+
+    assert result["pseudo_speed"] == expected_speed
+    assert result["pseudo_density"] == expected_density
+    assert result["pseudo_temperature"] == expected_temperature