Skip to content

Commit

Permalink
modify tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@laspsandoval
laspsandoval committed Sep 12, 2024
1 parent 1f43556 commit b89426c
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 43 deletions.
38 changes: 27 additions & 11 deletions imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,22 +12,32 @@

def test_build_energy_bins():
"""Tests build_energy_bins function."""
energy_bin_start, energy_bin_end = build_energy_bins()
energy_bin_start, energy_bin_end, energy_bin_mean = build_energy_bins()

assert energy_bin_start[0] == 3.5
assert len(energy_bin_start) == 90
assert len(energy_bin_end) == 90
assert len(energy_bin_mean) == 90

# Comparison to expected values.
np.testing.assert_allclose(energy_bin_end[0], 3.6795, atol=1e-4)
np.testing.assert_allclose(energy_bin_start[-1], 299.9724, atol=1e-4)
np.testing.assert_allclose(energy_bin_end[-1], 315.3556, atol=1e-4)

# Calculate expected geometric mean for the first bin.
expected_mean_first = np.sqrt(3.5 * 3.6795)
np.testing.assert_allclose(energy_bin_mean[0], expected_mean_first, atol=1e-4)

# Calculate expected geometric mean for the last bin.
expected_mean_last = np.sqrt(299.9724 * 315.3556)
np.testing.assert_allclose(energy_bin_mean[-1], expected_mean_last, atol=1e-4)


def test_build_spatial_bins():
"""Tests build_spatial_bins function."""
az_bin_edges, el_bin_edges, \
az_bin_midpoints, el_bin_midpoints = build_spatial_bins()
az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
build_spatial_bins()
)

assert az_bin_edges[0] == 0
assert az_bin_edges[-1] == 360
Expand All @@ -37,6 +47,14 @@ def test_build_spatial_bins():
assert el_bin_edges[-1] == 90
assert len(el_bin_edges) == 361

assert len(az_bin_midpoints) == 720
np.testing.assert_allclose(az_bin_midpoints[0], 0.25, atol=1e-4)
np.testing.assert_allclose(az_bin_midpoints[-1], 359.75, atol=1e-4)

assert len(el_bin_midpoints) == 360
np.testing.assert_allclose(el_bin_midpoints[0], -89.75, atol=1e-4)
np.testing.assert_allclose(el_bin_midpoints[-1], 89.75, atol=1e-4)


def test_cartesian_to_spherical():
"""Tests cartesian_to_spherical function."""
Expand All @@ -45,23 +63,21 @@ def test_cartesian_to_spherical():
vy_sc = np.array([-707.5707, -516.0282])
vz_sc = np.array([618.0569, 892.6931])

az, el = cartesian_to_spherical(vx_sc, vy_sc, vz_sc)
az_sc, el_sc = cartesian_to_spherical(vx_sc, vy_sc, vz_sc)

# MATLAB code outputs:
np.testing.assert_allclose(az_sc, np.array([1.31300, 2.34891]), atol=1e-05, rtol=0)
np.testing.assert_allclose(
az, np.array([1.31300, 2.34891]), atol=1e-05, rtol=0
)
np.testing.assert_allclose(
el, np.array([-0.70136, -0.88901]), atol=1e-05, rtol=0
el_sc, np.array([-0.70136, -0.88901]), atol=1e-05, rtol=0
)


def test_bin_space():
"""Tests bin_space function."""
# Example particle velocity in the pointing frame wrt s/c.
vx_sc = np.array([-186.5575, 508.5697])
vy_sc = np.array([-707.5707, -516.0282])
vz_sc = np.array([618.0569, 892.6931])
vx_sc = np.array([-186.5575, 508.5697, 0])
vy_sc = np.array([-707.5707, -516.0282, 0])
vz_sc = np.array([618.0569, 892.6931, -1])

az_midpoint, el_midpoint = bin_space(vx_sc, vy_sc, vz_sc)
az, el = cartesian_to_spherical(vx_sc, vy_sc, vz_sc)
Expand Down
102 changes: 70 additions & 32 deletions imap_processing/ultra/l1c/ultra_l1c_pset_bins.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,15 @@
import numpy as np


def build_energy_bins() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
def build_energy_bins() -> tuple[np.ndarray, np.ndarray]:
"""
Build energy bin boundaries.
Returns
-------
energy_bin_start : np.ndarray
Array of energy bin start values.
energy_bin_end : np.ndarray
Array of energy bin end values.
energy_bin_midpoints : np.ndarray
energy_bin_edges : np.ndarray
Array of energy bin edges.
energy_bin_mean : np.ndarray
Array of energy bin midpoint values.
"""
alpha = 0.05 # deltaE/E
Expand All @@ -24,13 +22,12 @@ def build_energy_bins() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
energy_step = (1 + alpha / 2) / (1 - alpha / 2)

# Create energy bins.
bin_edges = energy_start * energy_step ** np.arange(n_bins + 1)
energy_bin_start = bin_edges[:-1]
energy_bin_end = bin_edges[1:]
energy_bin_edges = energy_start * energy_step ** np.arange(n_bins + 1)

energy_bin_midpoints = np.sqrt(energy_bin_start * energy_bin_end)
# Calculate the geometric mean.
energy_bin_mean = np.sqrt(energy_bin_edges[:-1] * energy_bin_edges[1:])

return energy_bin_start, energy_bin_end, energy_bin_midpoints
return energy_bin_edges, energy_bin_mean


def build_spatial_bins(
Expand Down Expand Up @@ -67,18 +64,18 @@ def build_spatial_bins(


def cartesian_to_spherical(
vx_dps_sc: np.ndarray, vy_dps_sc: np.ndarray, vz_dps_sc: np.ndarray
vx: np.ndarray, vy: np.ndarray, vz: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
"""
Convert cartesian coordinates to spherical coordinates.
Parameters
----------
vx_dps_sc : np.ndarray
vx : np.ndarray
The x-components of the velocity vector.
vy_dps_sc : np.ndarray
vy : np.ndarray
The y-components of the velocity vector.
vz_dps_sc : np.ndarray
vz : np.ndarray
The z-components of the velocity vector.
Returns
Expand All @@ -89,11 +86,11 @@ def cartesian_to_spherical(
The elevation angles in degrees.
"""
# Magnitude of the velocity vector
magnitude_v = np.sqrt(vx_dps_sc**2 + vy_dps_sc**2 + vz_dps_sc**2)
magnitude_v = np.sqrt(vx**2 + vy**2 + vz**2)

vhat_x = -vx_dps_sc / magnitude_v
vhat_y = -vy_dps_sc / magnitude_v
vhat_z = -vz_dps_sc / magnitude_v
vhat_x = -vx / magnitude_v
vhat_y = -vy / magnitude_v
vhat_z = -vz / magnitude_v

# Convert from cartesian to spherical coordinates (azimuth, elevation)
# Radius (magnitude)
Expand All @@ -111,38 +108,79 @@ def cartesian_to_spherical(
return az, el


def bin_space(vx_dps_sc, vy_dps_sc, vz_dps_sc) -> tuple[np.ndarray, np.ndarray]:
def bin_space(
vx: np.ndarray, vy: np.ndarray, vz: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
"""
Bin particle.
Bin the particle.
Parameters
----------
vx_dps_sc : float, optional
The bin spacing in degrees (default is 0.5 degrees).
vy_dps_sc : float, optional
The bin spacing in degrees (default is 0.5 degrees).
vz_dps_sc : float, optional
The bin spacing in degrees (default is 0.5 degrees).
vx : np.ndarray
The x-components of the velocity vector.
vy : np.ndarray
The y-components of the velocity vector.
vz : np.ndarray
The z-components of the velocity vector.
Returns
-------
az_midpoint : np.ndarray
Array of azimuth bin boundary values.
Array of azimuth midpoint values.
el_midpoint : np.ndarray
Array of elevation bin boundary values.
Array of elevation midpoint values.
"""
az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
build_spatial_bins()
)

az, el = cartesian_to_spherical(vx_dps_sc, vy_dps_sc, vz_dps_sc)
az, el = cartesian_to_spherical(vx, vy, vz)

az_degrees = np.degrees(az)
el_degrees = np.degrees(el)

# If azimuth is exactly 360 degrees it is placed in last bin.
az_degrees[az_degrees >= az_bin_edges[-1]] = az_bin_midpoints[-1]
# If elevation is exactly 90 degrees it is placed in last bin.
el_degrees[el_degrees >= el_bin_edges[-1]] = el_bin_midpoints[-1]

# Find the appropriate bin index.
az_bin_idx = np.searchsorted(az_bin_edges, np.degrees(az), side="right") - 1
el_bin_idx = np.searchsorted(el_bin_edges, np.degrees(el), side="right") - 1
az_bin_idx = np.searchsorted(az_bin_edges, az_degrees, side="right") - 1
el_bin_idx = np.searchsorted(el_bin_edges, el_degrees, side="right") - 1

# Assign the corresponding midpoints.
az_midpoint = az_bin_midpoints[az_bin_idx]
el_midpoint = el_bin_midpoints[el_bin_idx]

return az_midpoint, el_midpoint


def bin_energy(energy: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""
Bin the particle.
Parameters
----------
energy : np.ndarray
Particle energy.
Returns
-------
energy_mean : np.ndarray
Mean energy value.
"""
# TODO: Use quality flags to filter out energies beyond threshold.

energy_bin_edges, energy_bin_mean = build_energy_bins()

# If energy is exactly equal to the last bin edge it is placed in last bin.
energy[energy >= energy_bin_edges[-1]] = energy_bin_mean[-1]

# Find the appropriate bin index.
energy_bin_idx = np.searchsorted(energy_bin_edges, energy, side="right") - 1

# Assign the corresponding means.
az_mean = energy_bin_mean[energy_bin_idx]
el_mean = energy_bin_mean[energy_bin_idx]

return az_mean, el_mean

0 comments on commit b89426c

Please sign in to comment.