Skip to content

Commit d29fedb

Browse files
pleasant-menlopleasant-menlo
authored
Merge pull request #75 from mecuesta/swe_l3_potential
SWE L3 Breakpoint Finder Replacement Reviewed by PMCL/EYEA
2 parents 5c1426f + 1b10e92 commit d29fedb

4 files changed

Lines changed: 211 additions & 12 deletions

File tree

‎imap_l3_processing/swe/l3/science/pitch_calculations.py‎

Lines changed: 147 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,12 @@
33
from typing import TypeVar
44

55
import numpy as np
6+
import scipy
67

78
from imap_l3_processing.constants import ELECTRON_MASS_KG, PROTON_CHARGE_COULOMBS, METERS_PER_KILOMETER
89
from imap_l3_processing.pitch_angles import calculate_pitch_angle, calculate_unit_vector, calculate_gyrophase
910
from imap_l3_processing.swe.l3.models import SweConfiguration
10-
11+
from imap_l3_processing.swe.quality_flags import SweL3Flags
1112

1213
def piece_wise_model(x: np.ndarray, b0: float, b1: float,
1314
b2: float, b3: float, b4: float, b5: float) -> np.ndarray:
@@ -18,6 +19,151 @@ def piece_wise_model(x: np.ndarray, b0: float, b1: float,
1819
lambda x: b0 * np.exp(b2 * (b3 - b1)) * np.exp(b4 * (b5 - b3)) * np.exp(-b5 * x),
1920
]))
2021

22+
def mec_breakpoint_finder(energies: np.ndarray, averaged_psd: np.ndarray) -> tuple[float, float, SweL3Flags]:
23+
"""
24+
Input:
25+
energies - energy bins
26+
averaged_psd - phase space density, either averaged over all CEMs or individual CEMs
27+
Output:
28+
sc_pot_output, ch_break_output, total_flag - tuple for spacecraft potential, core-halo break point, and all flags thrown
29+
"""
30+
log_energy = np.log(energies)
31+
log_psd = np.log(averaged_psd)
32+
33+
# Check to see if first 4 points are nearly linear
34+
# If True, then potential is likely less than 2.7 V (lower than first energy bin)
35+
def line_model(params, x):
36+
return params[0] + params[1] * x
37+
from scipy import odr
38+
odr_model = odr.Model(line_model)
39+
x = log_energy[:4]
40+
y = log_psd[:4]
41+
mydata = odr.RealData(x=x,y=y)
42+
myodr = odr.ODR(mydata, odr_model, beta0=[y.max(),-5])
43+
myodr.set_job(fit_type=2)
44+
myoutput= myodr.run()
45+
if myoutput.res_var <= 0.01:
46+
FALLBACK_POTENTIAL_ESTIMATE = SweL3Flags.FALLBACK_POTENTIAL_ESTIMATE
47+
return_value = 2.5
48+
else:
49+
FALLBACK_POTENTIAL_ESTIMATE = SweL3Flags.NONE
50+
51+
# Use a smoothed spline on log_psd for spectral break finding routine as a fall back only!
52+
# Mirror real point as fake point to left of first energy bin to improve spline concavity
53+
ewidth = np.nanmean(log_energy[1:] - log_energy[:-1])
54+
from scipy.interpolate import UnivariateSpline as uspline
55+
spline = uspline(np.concatenate([[log_energy[0]-ewidth],log_energy]),
56+
np.concatenate([[log_psd[2]],log_psd]), s=.25)
57+
spline_energies = np.geomspace(energies.min()*np.exp(-ewidth), energies.max(), 100)
58+
spline_derivative = spline.derivative(2)
59+
curvature = spline_derivative(np.log(spline_energies))
60+
try:
61+
peaks = scipy.signal.find_peaks(curvature)[0]
62+
sc_pot = spline_energies[peaks[0]]
63+
ch_break = spline_energies[peaks[1]]
64+
BACKUP_SPLINE_UNRESOLVED = SweL3Flags.NONE
65+
except:
66+
# Spline peak finder did not work
67+
BACKUP_SPLINE_UNRESOLVED = SweL3Flags.BACKUP_SPLINE_UNRESOLVED
68+
sc_pot = np.nan
69+
ch_break = np.nan
70+
71+
def piece_wise_model_mec(x, b0, b1, b2, b3):
72+
"""
73+
Modified Piecewise to fit Potential and Core-Halo Break separately
74+
The breakpoint is b2
75+
"""
76+
return np.piecewise(x, [x<=b2, x>b2],
77+
[lambda x: b0 - b1*x, lambda x: b0 + b2*(b3-b1) - b3*x])
78+
79+
def refine_breakpoint_value(energy, psd, breakpoint_value, num_points):
80+
"""
81+
Function to use lines from num_points to left and right to find intersection
82+
energy at which the lines intersect is the refined_breakpoint
83+
"""
84+
# Find Nearest energy bin to breakpoint_value
85+
nearest_energy_idx = np.argmin(np.abs(energy-breakpoint_value))
86+
if np.abs(breakpoint_value - energy[nearest_energy_idx])/energy[nearest_energy_idx] <= .075:
87+
# Breakpoint_value was within FWHM of nearest energy bin
88+
# return that energy bin as refined_breakpoint
89+
return energy[nearest_energy_idx]
90+
# Get left and right spectrum of breakpoint_value
91+
e_left = energy[energy < breakpoint_value]
92+
e_right = energy[energy > breakpoint_value]
93+
psd_left = psd[energy < breakpoint_value]
94+
psd_right = psd[energy > breakpoint_value]
95+
if len(e_left) < num_points:
96+
# Not enough points to the left (really only possible for s/c potential)
97+
return breakpoint_value
98+
# Fit a line to the num_points left and right of breakpoint_value
99+
z_left = np.polyfit(e_left[-num_points:], psd_left[-num_points:], 1)
100+
z_right = np.polyfit(e_right[:num_points], psd_right[:num_points], 1)
101+
# Determine energy of their intersection
102+
refined_breakpoint = (z_left[1]-z_right[1]) / (z_right[0]-z_left[0])
103+
if (refined_breakpoint < z_left[-1]) | (refined_breakpoint > z_right[0]):
104+
# Refined breakpoint lies outside of expected range
105+
return breakpoint_value
106+
return refined_breakpoint
107+
108+
# Prepare masking for the two separate fits
109+
mask_sc = log_energy <= np.log(30)
110+
mask_ch = (log_energy > np.log(30)) & (log_energy < np.log(400))
111+
112+
log_energy_sc = log_energy[mask_sc]; log_psd_sc = log_psd[mask_sc]
113+
log_energy_ch = log_energy[mask_ch]; log_psd_ch = log_psd[mask_ch]
114+
115+
fitting_model = piece_wise_model_mec
116+
# Try Spacecraft Potential Fit
117+
POTENTIAL_FIT_UNCONVERGED = SweL3Flags.NONE
118+
if FALLBACK_POTENTIAL_ESTIMATE == SweL3Flags.NONE:
119+
try:
120+
initial_guess = [log_psd[0],1,7,1]
121+
z, cov = scipy.optimize.curve_fit(fitting_model, np.exp(log_energy_sc), log_psd_sc, p0=initial_guess)
122+
# Make sure the fit converged
123+
if ((z[0] == initial_guess[0]) | (z[1] == initial_guess[1])
124+
| (z[2] == initial_guess[2]) | (z[3] == initial_guess[3])):
125+
# Fall back on Spline method
126+
# Fit did not converge
127+
POTENTIAL_FIT_UNCONVERGED = SweL3Flags.POTENTIAL_FIT_UNCONVERGED
128+
sc_pot_output = sc_pot
129+
else:
130+
# Fit worked
131+
# Check whether breakpoint has two points to left and right
132+
# If so, then find intersection of linear fits to each side
133+
sc_pot_output = refine_breakpoint_value(np.exp(log_energy_sc), log_psd_sc, z[2], 2)
134+
# spline not used
135+
BACKUP_SPLINE_UNRESOLVED = SweL3Flags.NONE
136+
except:
137+
# Fall back on Spline method
138+
# Fit did not converge
139+
BACKUP_SPLINE_UNRESOLVED = SweL3Flags.BACKUP_SPLINE_UNRESOLVED
140+
sc_pot_output = sc_pot
141+
else:
142+
sc_pot_output = return_value
143+
# Try Core-Halo Breakpoint Fit
144+
BREAKPOINT_FIT_UNCONVERGED = SweL3Flags.NONE
145+
try:
146+
initial_guess = [log_psd_ch[0],1,65,1]
147+
z, cov = scipy.optimize.curve_fit(fitting_model, np.exp(log_energy_ch), log_psd_ch, p0=initial_guess)
148+
# Make sure the fit converged
149+
if ((z[0] == initial_guess[0]) & (z[1] == initial_guess[1])
150+
& (z[2] == initial_guess[2]) & (z[3] == initial_guess[3])):
151+
# Fall back on Spline method
152+
# Fit did not converge
153+
BREAKPOINT_FIT_UNCONVERGED = SweL3Flags.BREAKPOINT_FIT_UNCONVERGED
154+
ch_break_output = ch_break
155+
else:
156+
# Fit worked
157+
# Check whether breakpoint has three points to left and right
158+
# If so, then find intersection of linear fits to each side
159+
ch_break_output = refine_breakpoint_value(np.exp(log_energy_ch), log_psd_ch, z[2], 3)
160+
except:
161+
# Fall back on Spline method
162+
# Fit did not converge
163+
BREAKPOINT_FIT_UNCONVERGED = SweL3Flags.BREAKPOINT_FIT_UNCONVERGED
164+
ch_break_output = ch_break
165+
return sc_pot_output, ch_break_output, FALLBACK_POTENTIAL_ESTIMATE | BACKUP_SPLINE_UNRESOLVED | POTENTIAL_FIT_UNCONVERGED | BREAKPOINT_FIT_UNCONVERGED
166+
21167

22168
def find_breakpoints(energies: np.ndarray, averaged_psd: np.ndarray, latest_spacecraft_potentials: list[float],
23169
latest_core_halo_break_points: list[float],

‎imap_l3_processing/swe/quality_flags.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
from imap_processing.quality_flags import FlagNameMixin, CommonFlags
2+
3+
4+
class SweL3Flags(FlagNameMixin):
5+
NONE = CommonFlags.NONE
6+
FALLBACK_POTENTIAL_ESTIMATE = 2 ** 2
7+
BACKUP_SPLINE_UNRESOLVED = 2 ** 3
8+
POTENTIAL_FIT_UNCONVERGED = 2 ** 4
9+
BREAKPOINT_FIT_UNCONVERGED = 2 ** 5
10+
ULTRA_HV_OFF = 2 ** 6
11+
TEMPERATURE_OUTLIER = 2 ** 7
12+
PRELIMINARY_MAG = 2 ** 8
13+
FALLBACK_SWAPI_SPEED = 2 ** 9
14+
NEGATIVE_MOMENT = 2 ** 10
15+
PREDICTIVE_EPHEMERIS = 2 ** 15
16+

‎imap_l3_processing/swe/swe_processor.py‎

Lines changed: 27 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,14 @@
1616
integrate, scale_core_density, scale_halo_density, rotate_vector_to_rtn_spherical_coordinates, \
1717
calculate_primary_eigenvector, \
1818
ScaleDensityOutput, rotate_temperature_tensor_to_mag
19-
from imap_l3_processing.swe.l3.science.pitch_calculations import average_over_look_directions, find_breakpoints, \
19+
from imap_l3_processing.swe.l3.science.pitch_calculations import average_over_look_directions, find_breakpoints, mec_breakpoint_finder, \
2020
correct_and_rebin, \
2121
integrate_distribution_to_get_1d_spectrum, integrate_distribution_to_get_inbound_and_outbound_1d_spectrum, \
2222
calculate_velocity_in_dsp_frame_km_s, swe_rebin_intensity_by_pitch_angle_and_gyrophase
2323
from imap_l3_processing.swe.l3.swe_l3_dependencies import SweL3Dependencies
2424
from imap_l3_processing.swe.l3.utils import compute_epoch_delta_in_ns
2525
from imap_l3_processing.utils import save_data
26+
from imap_l3_processing.swe.quality_flags import SweL3Flags
2627

2728
logger = logging.getLogger(__name__)
2829

@@ -50,6 +51,7 @@ def calculate_products(self, dependencies: SweL3Dependencies) -> SweL3Data:
5051
halo_core_history = [config["core_halo_breakpoint_initial_guess"] for _ in range(3)]
5152

5253
average_psd = []
54+
swe_quality_flags = []
5355
spacecraft_potential: np.ndarray[np.float64] = np.empty_like(swe_epoch, dtype=np.float64)
5456
halo_core: np.ndarray[np.float64] = np.empty_like(swe_epoch, dtype=np.float64)
5557
corrected_energy_bins = []
@@ -59,19 +61,32 @@ def calculate_products(self, dependencies: SweL3Dependencies) -> SweL3Data:
5961

6062
geometric_fractions = np.array(config["geometric_fractions"])
6163
for i in range(len(swe_epoch)):
62-
average_psd.append(average_over_look_directions(swe_l2_data.phase_space_density[i],
64+
# average PSDs in time smoothed by 7 samples centered to help with break point finder
65+
# edges will have less than 7 samples
66+
npts = 7//2
67+
left_idx = 0 if i-npts < 0 else i-npts
68+
right_idx = i + (i-left_idx+1)
69+
time_avg_psd = np.nanmean(swe_l2_data.phase_space_density[left_idx:right_idx], axis=0)
70+
71+
average_psd.append(average_over_look_directions(time_avg_psd,
6372
geometric_fractions,
6473
config["minimum_phase_space_density_value"]))
6574

66-
spacecraft_potential[i], halo_core[i] = find_breakpoints(swe_l2_data.energy, average_psd[i],
67-
spacecraft_potential_history,
68-
halo_core_history,
69-
config)
75+
#spacecraft_potential[i], halo_core[i] = find_breakpoints(swe_l2_data.energy, average_psd[i],
76+
# spacecraft_potential_history,
77+
# halo_core_history,
78+
# config)
79+
# FLAG_UPDATE is a placeholder to be replaced when formal swe_flags are generated
80+
(
81+
spacecraft_potential[i], halo_core[i], breakpoint_quality_flag
82+
) = mec_breakpoint_finder(swe_l2_data.energy, average_psd[i])
7083

7184
spacecraft_potential_history = [*spacecraft_potential_history[1:], spacecraft_potential[i]]
7285
halo_core_history = [*halo_core_history[1:], halo_core[i]]
7386
corrected_energy_bins.append(swe_l2_data.energy - spacecraft_potential[i])
87+
swe_quality_flags.append(breakpoint_quality_flag)
7488

89+
7590
corrected_energy_bins = np.array(corrected_energy_bins)
7691

7792
swe_l3_moments_data = self.calculate_moment_products(swe_l2_data, dependencies.swe_l1b_data,
@@ -191,8 +206,7 @@ def calculate_moment_products(self, swe_l2_data: SweL2Data, swe_l1b_data: SweL1b
191206
swe_l2_data.inst_az_spin_sector[i])
192207

193208
weights: np.ndarray[float] = compute_maxwellian_weight_factors(swe_l1b_data.count_rates[i],
194-
swe_l2_data.acquisition_duration[
195-
i] / 1e6)
209+
swe_l2_data.acquisition_duration[i] / 1e6)
196210
ifit = next(
197211
index for index, energy in enumerate(swe_l2_data.energy) if energy >= spacecraft_potential[i])
198212
jbreak = next(index for index, energy in enumerate(swe_l2_data.energy) if energy >= halo_core[i])
@@ -286,7 +300,8 @@ def calculate_moment_products(self, swe_l2_data: SweL2Data, swe_l1b_data: SweL1b
286300
core_t_perpendicular_to_mag_ratio]
287301

288302
total_integration_output = integrate(ifit + 1,
289-
len(corrected_energy_bins[i]) - 1,
303+
# MEC: Changed 1 to 6 in next line
304+
len(corrected_energy_bins[i]) - 6,
290305
corrected_energy_bins[i],
291306
sin_theta,
292307
cos_theta, config['aperture_field_of_view_radians'],
@@ -365,7 +380,9 @@ def calculate_moment_products(self, swe_l2_data: SweL2Data, swe_l1b_data: SweL1b
365380
halo_temp_phi_rtns[i] = halo_temp_phi_rtn
366381
halo_temp_avg = (2 * halo_moment.t_perpendicular + halo_moment.t_parallel) / 3
367382
if 1e4 < halo_temp_avg < 1e8:
368-
halo_integrate_result = integrate(jbreak, len(corrected_energy_bins[i]) - 1,
383+
halo_integrate_result = integrate(jbreak,
384+
# MEC: Changed 1 to 6 in next line
385+
len(corrected_energy_bins[i]) - 6,
369386
corrected_energy_bins[i],
370387
sin_theta, cos_theta,
371388
config['aperture_field_of_view_radians'],

‎tests/swe/l3/science/test_pitch_calculations.py‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,14 @@
33

44
import numpy as np
55

6-
from imap_l3_processing.swe.l3.science.pitch_calculations import find_breakpoints, \
6+
from imap_l3_processing.swe.l3.science.pitch_calculations import find_breakpoints, mec_breakpoint_finder, \
77
average_over_look_directions, calculate_velocity_in_dsp_frame_km_s, calculate_look_directions, rebin_by_pitch_angle, \
88
correct_and_rebin, calculate_energy_in_ev_from_velocity_in_km_per_second, integrate_distribution_to_get_1d_spectrum, \
99
integrate_distribution_to_get_inbound_and_outbound_1d_spectrum, try_curve_fit_until_valid, \
1010
rebin_by_pitch_angle_and_gyrophase, swe_rebin_intensity_by_pitch_angle_and_gyrophase, ls_fit
1111
from tests.test_helpers import build_swe_configuration, NumpyArrayMatcher
1212

13+
from imap_l3_processing.swe.quality_flags import SweL3Flags
1314

1415
class TestPitchCalculations(unittest.TestCase):
1516
def test_average_flux(self):
@@ -126,6 +127,25 @@ def test_compute_velocity_and_confirm_energy_calculation(self):
126127
np.testing.assert_almost_equal(calculated_energy[1], energy[1])
127128
self.assertEqual((24, 30, 7, 3), velocity.shape)
128129

130+
def test_mec_breakpoint_finder(self):
131+
average_psd = np.asarray([3.86611617e-05, 1.75927552e-05, 6.06703495e-06, 1.47614412e-06,
132+
2.46270751e-07, 4.97864654e-08, 3.90323281e-08, 4.21037028e-08,
133+
4.07355835e-08, 3.45166807e-08, 2.52075419e-08, 1.54987172e-08,
134+
7.84932795e-09, 3.32903783e-09, 1.40908421e-09, 8.53206391e-10,
135+
7.21564216e-10, 6.31903627e-10, 5.07238502e-10, 3.58827457e-10])
136+
energies = np.asarray([ 2.66 , 3.37852249, 4.29113316, 5.45025936,
137+
6.92249015, 8.79240175, 11.16741618, 14.18397245,
138+
18.01536462, 22.88169719, 29.06252952, 36.91293593,
139+
46.88390382, 59.54824189, 75.63348661, 96.06369752,
140+
122.01254227, 154.97072104, 196.83160382, 250. ])
141+
142+
sc_pot_to_test, ch_break_to_test, quality_flag_to_test = mec_breakpoint_finder(energies, average_psd)
143+
expected_values = [np.float64(8.79240175),np.float64(75.63348660999998)]
144+
self.assertAlmostEqual(expected_values[0], sc_pot_to_test)
145+
self.assertAlmostEqual(expected_values[1], ch_break_to_test)
146+
self.assertEqual(SweL3Flags.NONE, quality_flag_to_test)
147+
148+
129149
@patch('imap_l3_processing.swe.l3.science.pitch_calculations.try_curve_fit_until_valid')
130150
def test_find_breakpoints_determines_b_deltas_correctly(self, mock_try_curve_fit_until_valid):
131151
config = build_swe_configuration(refit_core_halo_breakpoint_index=4, slope_ratio_cutoff_for_potential_calc=0)

0 commit comments

Comments
 (0)