Ultra l1c bin data

imap_processing/tests/ultra/unit/test_ultra_l1c_pset_bins.py

@@ -3,28 +3,35 @@
 import numpy as np
 
 from imap_processing.ultra.l1c.ultra_l1c_pset_bins import (
+    bin_energy,
+    bin_space,
     build_energy_bins,
     build_spatial_bins,
+    cartesian_to_spherical,
 )
 
 
 def test_build_energy_bins():
     """Tests build_energy_bins function."""
-    energy_bin_start, energy_bin_end = build_energy_bins()
+    energy_bin_edges = build_energy_bins()
+    energy_bin_start = energy_bin_edges[:-1]
+    energy_bin_end = energy_bin_edges[1:]
 
-    assert energy_bin_start[0] == 3.5
-    assert len(energy_bin_start) == 90
-    assert len(energy_bin_end) == 90
+    assert energy_bin_start[0] == 0
+    assert energy_bin_start[1] == 3.385
+    assert len(energy_bin_edges) == 25
 
     # 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)
+    np.testing.assert_allclose(energy_bin_end[1], 4.137, atol=1e-4)
+    np.testing.assert_allclose(energy_bin_start[-1], 279.810, atol=1e-4)
+    np.testing.assert_allclose(energy_bin_end[-1], 341.989, atol=1e-4)
 
 
 def test_build_spatial_bins():
     """Tests build_spatial_bins function."""
-    az_bin_edges, el_bin_edges = 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
@@ -33,3 +40,69 @@ def test_build_spatial_bins():
     assert el_bin_edges[0] == -90
     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."""
+    # 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])
+    v = np.column_stack((vx_sc, vy_sc, vz_sc))
+
+    az_sc, el_sc, r = cartesian_to_spherical(v)
+
+    # MATLAB code outputs:
+    np.testing.assert_allclose(az_sc, np.array([1.31300, 2.34891]), atol=1e-05, rtol=0)
+    np.testing.assert_allclose(
+        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, 508.5697, 0])
+    vy_sc = np.array([-707.5707, -516.0282, -516.0282, 0])
+    vz_sc = np.array([618.0569, 892.6931, 892.6931, -1])
+
+    v = np.column_stack((vx_sc, vy_sc, vz_sc))
+    bin_id, counts = bin_space(v)
+
+    az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
+        build_spatial_bins()
+    )
+    actual_az, actual_el, _ = cartesian_to_spherical(v)
+
+    # Transform flat index bin_id back to 2D grid coordinates.
+    az_indices, el_indices = divmod(bin_id, len(el_bin_midpoints))
+
+    az_bin = az_bin_midpoints[az_indices]
+    el_bin = el_bin_midpoints[el_indices]
+
+    az_within_tolerance = (
+        np.abs(np.sort(az_bin) - np.degrees(np.unique(actual_az))) <= 0.25
+    )
+    el_within_tolerance = (
+        np.abs(np.sort(el_bin) - np.degrees(np.unique(actual_el))) <= 0.25
+    )
+
+    assert np.all(az_within_tolerance)
+    assert np.all(el_within_tolerance)
+    assert np.array_equal(counts, np.array([1, 2, 1]))
+
+
+def test_bin_energy():
+    """Tests bin_energy function."""
+    energy = np.array([3.384, 3.385, 341.989, 342])
+    bin = bin_energy(energy)
+
+    np.testing.assert_equal(bin, (0, 3.385, 341.989, 341.989))

imap_processing/ultra/l1c/ultra_l1c_pset_bins.py

@@ -1,35 +1,37 @@
-"""Module to create bins for pointing sets."""
+"""Module to create energy bins for pointing sets."""
 
 import numpy as np
+from numpy.typing import NDArray
 
 
-def build_energy_bins() -> tuple[np.ndarray, np.ndarray]:
+def build_energy_bins() -> NDArray[np.float64]:
     """
     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_edges : np.ndarray
+        Array of energy bin edges.
     """
-    alpha = 0.05  # deltaE/E
-    energy_start = 3.5  # energy start for the Ultra grids
-    n_bins = 90  # number of energy bins
+    # TODO: these value will almost certainly change.
+    alpha = 0.2  # deltaE/E
+    energy_start = 3.385  # energy start for the Ultra grids
+    n_bins = 23  # number of energy bins
 
     # Calculate energy step
     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)
+    # Add a zero to the left side for outliers and round to nearest 3 decimal places.
+    energy_bin_edges = np.around(np.insert(energy_bin_edges, 0, 0), 3)
 
-    return energy_bin_start, energy_bin_end
+    return energy_bin_edges
 
 
-def build_spatial_bins(spacing: float = 0.5) -> tuple[np.ndarray, np.ndarray]:
+def build_spatial_bins(
+    spacing: float = 0.5,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """
     Build spatial bin boundaries for azimuth and elevation.
 
@@ -44,11 +46,129 @@ def build_spatial_bins(spacing: float = 0.5) -> tuple[np.ndarray, np.ndarray]:
         Array of azimuth bin boundary values.
     el_bin_edges : np.ndarray
         Array of elevation bin boundary values.
+    az_bin_midpoints : np.ndarray
+        Array of azimuth bin midpoint values.
+    el_bin_midpoints : np.ndarray
+        Array of elevation bin midpoint values.
     """
     # Azimuth bins from 0 to 360 degrees.
     az_bin_edges = np.arange(0, 360 + spacing, spacing)
+    az_bin_midpoints = az_bin_edges[:-1] + spacing / 2  # Midpoints between edges
 
     # Elevation bins from -90 to 90 degrees.
     el_bin_edges = np.arange(-90, 90 + spacing, spacing)
+    el_bin_midpoints = el_bin_edges[:-1] + spacing / 2  # Midpoints between edges
 
-    return az_bin_edges, el_bin_edges
+    return az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints
+
+
+def cartesian_to_spherical(
+    v: tuple[np.ndarray, np.ndarray, np.ndarray],
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Convert cartesian coordinates to spherical coordinates.
+
+    Parameters
+    ----------
+    v : tuple[np.ndarray, np.ndarray, np.ndarray]
+        The x,y,z-components of the velocity vector.
+
+    Returns
+    -------
+    az : np.ndarray
+        The azimuth angles in degrees.
+    el : np.ndarray
+        The elevation angles in degrees.
+    r : np.ndarray
+        The radii, or magnitudes, of the vectors.
+    """
+    vx, vy, vz = np.hsplit(v, 3)
+    vx, vy, vz = vx.flatten(), vy.flatten(), vz.flatten()  # Flatten the arrays
+
+    # Magnitude of the velocity vector
+    magnitude_v = np.sqrt(vx**2 + vy**2 + vz**2)
+
+    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)
+    r = np.sqrt(vhat_x**2 + vhat_y**2 + vhat_z**2)
+
+    # Elevation angle (angle from the z-axis, range: [-pi/2, pi/2])
+    el = np.arcsin(vhat_z / r)
+
+    # Azimuth angle (angle in the xy-plane, range: [0, 2*pi])
+    az = np.arctan2(vhat_y, vhat_x)
+
+    # Ensure azimuth is from 0 to 2PI
+    az = az % (2 * np.pi)
+
+    return az, el, r
+
+
+def bin_space(
+    v: tuple[np.ndarray, np.ndarray, np.ndarray],
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Bin the particle.
+
+    Parameters
+    ----------
+    v : tuple[np.ndarray, np.ndarray, np.ndarray]
+        The x,y,z-components of the velocity vector.
+
+    Returns
+    -------
+    unique_bin_ids : np.ndarray
+        Bin ids.
+    counts : np.ndarray
+        Event counts.
+    """
+    az_bin_edges, el_bin_edges, az_bin_midpoints, el_bin_midpoints = (
+        build_spatial_bins()
+    )
+
+    az, el, _ = cartesian_to_spherical(v)
+
+    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, az_degrees, side="right") - 1
+    el_bin_idx = np.searchsorted(el_bin_edges, el_degrees, side="right") - 1
+
+    # Create flattened version of the 2D data array.
+    bin_id = az_bin_idx * len(el_bin_midpoints) + el_bin_idx
+    # Counts for each unique bin_id.
+    unique_bin_ids, counts = np.unique(bin_id, return_counts=True)
+
+    return unique_bin_ids, counts
+
+
+def bin_energy(energy: np.ndarray) -> NDArray[np.float64]:
+    """
+    Bin the particle.
+
+    Parameters
+    ----------
+    energy : np.ndarray
+        Particle energy.
+
+    Returns
+    -------
+    energy_start : np.ndarray
+        Start of energy bin.
+    """
+    energy_bin_edges = build_energy_bins()
+
+    # Find the appropriate bin index.
+    energy_bin_idx = np.searchsorted(energy_bin_edges, energy, side="right") - 1
+
+    return energy_bin_edges[energy_bin_idx]