Harrison 2382 - SWYN/PSCH - Rotate temperature and velocity in moments calculation

Author: Menlo Innovations - CAVA Project

imap_processing/swe/l3/science/moment_calculations.py

@@ -1,7 +1,9 @@
 import math
 from dataclasses import dataclass
+from datetime import datetime
 
 import numpy as np
+from spiceypy import spiceypy
 
 from imap_processing.constants import ELECTRON_MASS_KG, \
     BOLTZMANN_CONSTANT_JOULES_PER_KELVIN
@@ -109,7 +111,6 @@ class Moments:
     ao: float
 
 
-# TODO figure out what the moments are and return a dataclass with them
 def calculate_fit_temperature_density_velocity(parameters: np.ndarray[float]):
     moments = Moments(alpha=0,
                       beta=0,
@@ -226,3 +227,23 @@ def filter_and_flatten_regress_parameters(corrected_energy_bins: np.ndarray,
     yreg[filtered_phase_space_density <= 1e-35] = -80.6
 
     return velocity_vectors[valid_mask], weights[valid_mask], yreg
+
+
+def rotate_dps_vector_to_rtn(epoch: datetime, vector: np.ndarray[float]) -> np.ndarray[float]:
+    et_time = spiceypy.datetime2et(epoch)
+    rotation_matrix = spiceypy.pxform("IMAP_DPS", "IMAP_RTN", et_time)
+    return rotation_matrix @ vector
+
+
+def rotate_temperature(epoch: datetime, alpha: float, beta: float):
+    sin_dec = np.sin(beta)
+    x = sin_dec * np.cos(alpha)
+    y = sin_dec * np.sin(alpha)
+    z = np.cos(beta)
+
+    rtn_temperature = rotate_dps_vector_to_rtn(epoch, np.array([x, y, z]))
+
+    phi = np.asin(rtn_temperature[2])
+    theta = np.atan2(rtn_temperature[1], rtn_temperature[0])
+
+    return theta, phi

imap_processing/swe/swe_processor.py

@@ -3,12 +3,14 @@
 import numpy as np
 from imap_data_access import upload
 
+from imap_processing import spice_wrapper
 from imap_processing.data_utils import find_closest_neighbor
 from imap_processing.processor import Processor
 from imap_processing.swapi.l3a.science.calculate_pickup_ion import calculate_solar_wind_velocity_vector
 from imap_processing.swe.l3.models import SweL3Data
 from imap_processing.swe.l3.science.moment_calculations import compute_maxwellian_weight_factors, \
-    filter_and_flatten_regress_parameters, regress, calculate_fit_temperature_density_velocity
+    filter_and_flatten_regress_parameters, regress, calculate_fit_temperature_density_velocity, rotate_temperature, \
+    rotate_dps_vector_to_rtn
 from imap_processing.swe.l3.science.pitch_calculations import average_over_look_directions, find_breakpoints, \
     correct_and_rebin, \
     integrate_distribution_to_get_1d_spectrum, integrate_distribution_to_get_inbound_and_outbound_1d_spectrum, \
@@ -26,22 +28,21 @@ def process(self):
         upload(output_cdf)
 
     def calculate_moment_products(self, dependencies: SweL3Dependencies):
+        spice_wrapper.furnish()
         swe_l2_data = dependencies.swe_l2_data
         swe_epoch = swe_l2_data.epoch
         config = dependencies.configuration
 
         spacecraft_potential_history = [config["spacecraft_potential_initial_guess"] for _ in
-                                        range(4)]
-        halo_core_history = [config["core_halo_breakpoint_initial_guess"] for _ in range(4)]
+                                        range(3)]
+        halo_core_history = [config["core_halo_breakpoint_initial_guess"] for _ in range(3)]
 
         for i in range(len(swe_epoch)):
-            averaged_phase_space_density = average_over_look_directions(swe_l2_data.phase_space_density[i],
-                                                                        np.array(config["geometric_fractions"]))
-
-            spacecraft_potential, halo_core = find_breakpoints(swe_l2_data.energy, averaged_phase_space_density,
-                                                               np.average(spacecraft_potential_history[:3]),
-                                                               np.average(halo_core_history[:3]),
-                                                               spacecraft_potential_history[-1], halo_core_history[-1],
+            averaged_psd = average_over_look_directions(swe_l2_data.phase_space_density[i],
+                                                        np.array(config["geometric_fractions"]))
+            spacecraft_potential, halo_core = find_breakpoints(swe_l2_data.energy, averaged_psd,
+                                                               spacecraft_potential_history,
+                                                               halo_core_history,
                                                                config)
 
             spacecraft_potential_history = [*spacecraft_potential_history[1:], spacecraft_potential]
@@ -53,7 +54,7 @@ def calculate_moment_products(self, dependencies: SweL3Dependencies):
                                                                                 swe_l2_data.inst_az_spin_sector[i])
 
             # TODO: read these from L1 dataset
-            ccounts = np.reshape(np.arange(24 * 30 * 7), (24, 30, 7)) * 1000
+            ccounts = np.reshape(np.arange(20 * 30 * 7), (20, 30, 7)) * 1000
 
             weights: np.ndarray[float] = compute_maxwellian_weight_factors(ccounts)
 
@@ -80,6 +81,11 @@ def calculate_moment_products(self, dependencies: SweL3Dependencies):
                 else:
                     halo_core_breakpoint_index -= 1
 
+            rtn_velocity = rotate_dps_vector_to_rtn(swe_epoch[i],
+                                                    np.array(
+                                                        [moments.velocity_x, moments.velocity_y, moments.velocity_z]))
+            rotate_temperature(swe_epoch[i], moments.alpha, moments.beta)
+
     def calculate_pitch_angle_products(self, dependencies: SweL3Dependencies) -> SweL3Data:
         swe_l2_data = dependencies.swe_l2_data
         swe_epoch = swe_l2_data.epoch

run_local.py

@@ -298,9 +298,9 @@ def create_hit_direct_event_cdf():
 
     if "swe_moments" in sys.argv:
         dependencies = SweL3Dependencies.from_file_paths(
-            get_test_data_path("swe/imap_swe_l2_sci-with-ace-data_20250101_v002.cdf"),
-            get_test_data_path("mag/imap_mag_l1d_mago-normal_20250101_v001.cdf"),
-            get_test_data_path("swe/imap_swapi_l3a_proton-sw_20250101_v001.cdf"),
+            get_test_data_path("swe/imap_swe_l2_sci-with-ace-data_19990609_v002.cdf"),
+            get_test_data_path("swe/imap_mag_l1d_mago-normal_19990609_v001.cdf"),
+            get_test_data_path("swe/imap_swapi_l3a_proton-sw_19990609_v001.cdf"),
             get_test_data_path("swe/example_swe_config.json"),
         )
         print(create_swe_moments_cdf(dependencies))

tests/hit/l3/test_hit_processor.py

@@ -1,7 +1,6 @@
 from copy import deepcopy
-from dataclasses import fields
 from datetime import datetime
-from typing import Type, TypeVar
+from typing import TypeVar
 from unittest import TestCase
 from unittest.mock import sentinel, patch, call, Mock
 
@@ -18,7 +17,7 @@
 from imap_processing.hit.l3.sectored_products.models import HitPitchAngleDataProduct
 from imap_processing.models import MagL1dData, InputMetadata
 from imap_processing.processor import Processor
-from tests.test_helpers import NumpyArrayMatcher
+from tests.test_helpers import NumpyArrayMatcher, create_dataclass_mock
 
 
 class TestHitProcessor(TestCase):
@@ -57,11 +56,11 @@ def test_process_pitch_angle_product(self, mock_rebin_by_pitch_angle_and_gyropha
         averaged_mag_vectors = [sentinel.mag_vector1, sentinel.mag_vector2]
 
         mock_dependencies = Mock(spec=HITL3SectoredDependencies)
-        mock_mag_data = self.create_dataclass_mock(MagL1dData)
+        mock_mag_data = create_dataclass_mock(MagL1dData)
         mock_mag_data.rebin_to = Mock()
         mock_mag_data.rebin_to.return_value = averaged_mag_vectors
         mock_dependencies.mag_l1d_data = mock_mag_data
-        mock_hit_data = self.create_dataclass_mock(HitL2Data)
+        mock_hit_data = create_dataclass_mock(HitL2Data)
         mock_hit_data.epoch = epochs
         mock_hit_data.epoch_delta = epoch_deltas
 
@@ -588,6 +587,3 @@ def test_raise_error_if_descriptor_doesnt_match(self):
                          f"Don't know how to generate 'spectral-index' /n Known HIT l3 data products: 'pitch-angle', 'direct-event'.")
 
     T = TypeVar("T")
-
-    def create_dataclass_mock(self, obj: Type[T]) -> T:
-        return Mock(spec=[field.name for field in fields(obj)])

tests/swe/l3/science/test_moments_calculation.py

@@ -1,9 +1,13 @@
+import math
 import unittest
+from datetime import datetime
+from unittest.mock import patch
 
 import numpy as np
 
 from imap_processing.swe.l3.science.moment_calculations import compute_maxwellian_weight_factors, \
-    filter_and_flatten_regress_parameters, regress, calculate_fit_temperature_density_velocity
+    filter_and_flatten_regress_parameters, regress, calculate_fit_temperature_density_velocity, rotate_temperature, \
+    rotate_dps_vector_to_rtn
 from tests.test_helpers import get_test_data_path
 
 
@@ -139,3 +143,44 @@ def test_filter_and_flatten_regress_parameters(self):
         np.testing.assert_array_equal(vectors, [[3, 0, 0], [4, 0, 0], [10, 0, 0], [10, 0, 0], [0, 0, 0]])
         np.testing.assert_array_equal(actual_weights, [3, 1e-36, 10, 11, 12])
         np.testing.assert_array_equal(yreg, [np.log(3), -80.6, np.log(10), np.log(11), np.log(12)])
+
+    @patch('spiceypy.spiceypy.pxform')
+    @patch('spiceypy.spiceypy.datetime2et')
+    def test_rotate_dps_vector_to_rtn(self, mock_datetime2et, mock_pxform):
+        epoch = datetime(year=2020, month=3, day=10)
+        dsp_vector = np.array([0, 1, 0])
+        rotation_matrix = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]])
+        mock_pxform.return_value = rotation_matrix
+
+        rtn_vector = rotate_dps_vector_to_rtn(epoch, dsp_vector)
+        mock_datetime2et.assert_called_once_with(epoch)
+
+        mock_pxform.assert_called_once_with("IMAP_DPS", "IMAP_RTN", mock_datetime2et.return_value)
+
+        np.testing.assert_array_equal(rtn_vector, rotation_matrix @ dsp_vector)
+
+    @patch('spiceypy.spiceypy.pxform')
+    @patch('spiceypy.spiceypy.datetime2et')
+    def test_rotate_temperature(self, mock_datetime2et, mock_pxform):
+        epoch = datetime(year=2020, month=3, day=11)
+        temperature_alpha = math.pi / 4
+        temperature_beta = math.pi / 8
+
+        rotation_matrix = np.array([[0, 0, 1], [1, 0, 0], [0, 1, 0]])
+        mock_pxform.return_value = rotation_matrix
+
+        theta, phi = rotate_temperature(epoch, temperature_alpha, temperature_beta)
+        mock_datetime2et.assert_called_once_with(epoch)
+
+        mock_pxform.assert_called_once_with("IMAP_DPS", "IMAP_RTN", mock_datetime2et.return_value)
+
+        sin_dec = np.sin(temperature_beta)
+        x = sin_dec * np.cos(temperature_alpha)
+        y = sin_dec * np.sin(temperature_alpha)
+        z = np.cos(temperature_beta)
+
+        expected_rtn_temperature = rotation_matrix @ np.array([x, y, z])
+        expected_rtn_temperature /= np.linalg.norm(expected_rtn_temperature)
+
+        self.assertEqual(np.asin(expected_rtn_temperature[2]), phi)
+        self.assertEqual(np.atan2(expected_rtn_temperature[1], expected_rtn_temperature[0]), theta)