imap-processing 0.12.0__py3-none-any.whl → 0.13.0__py3-none-any.whl

imap_processing/mag/l2/mag_l2_data.py

@@ -0,0 +1,288 @@
+"""Data structures for MAG L2 and L1D processing."""
+
+from dataclasses import InitVar, dataclass, field
+from enum import Enum
+
+import numpy as np
+import xarray as xr
+
+from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
+from imap_processing.mag.constants import DataMode
+
+
+class ValidFrames(Enum):
+    """SPICE reference frames for output."""
+
+    dsrf = "dsrf"
+    srf = "srf"
+    rtn = "rtn"
+    gse = "gse"
+
+
+@dataclass
+class MagL2:
+    """
+    Dataclass for MAG L2 data.
+
+    Since L2 and L1D should have the same structure, this can be used for either level.
+
+    Some of the methods are also static, so they can be used in i-ALiRT processing.
+
+    Attributes
+    ----------
+    vectors: np.ndarray
+        Magnetic field vectors of size (n, 3) where n is the number of vectors.
+        Describes (x, y, z) components of the magnetic field.
+    epoch: np.ndarray
+        Time of each vector in J2000 seconds. Should be of length n.
+    range: np.ndarray
+        Range of each vector. Should be of length n.
+    global_attributes: dict
+        Any global attributes we want to carry forward into the output CDF file.
+    quality_flags: np.ndarray
+        Quality flags for each vector. Should be of length n.
+    quality_bitmask: np.ndarray
+        Quality bitmask for each vector. Should be of length n. Copied from offset
+        file in L2, marked as good always in L1D.
+    magnitude: np.ndarray
+        Magnitude of each vector. Should be of length n. Calculated from L2 vectors.
+    is_l1d: bool
+        Flag to indicate if the data is L1D. Defaults to False.
+    """
+
+    vectors: np.ndarray
+    epoch: np.ndarray
+    range: np.ndarray
+    global_attributes: dict
+    quality_flags: np.ndarray
+    quality_bitmask: np.ndarray
+    data_mode: DataMode
+    magnitude: np.ndarray = field(init=False)
+    is_l1d: bool = False
+    offsets: InitVar[np.ndarray] = None
+    timedelta: InitVar[np.ndarray] = None
+
+    def __post_init__(self, offsets: np.ndarray, timedelta: np.ndarray) -> None:
+        """
+        Calculate the magnitude of the vectors after initialization.
+
+        Parameters
+        ----------
+        offsets : np.ndarray
+            Offsets to apply to the vectors. Should be of shape (n, 3) where n is the
+            number of vectors.
+        timedelta : np.ndarray
+            Time deltas to shift the timestamps by. Should be of length n.
+            Given in seconds.
+        """
+        if offsets is not None:
+            self.vectors = self.apply_offsets(self.vectors, offsets)
+        if timedelta is not None:
+            self.epoch = self.shift_timestamps(self.epoch, timedelta)
+
+        self.magnitude = self.calculate_magnitude(self.vectors)
+
+    @staticmethod
+    def calculate_magnitude(
+        vectors: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Given a list of vectors (x, y, z), calculate the magnitude of each vector.
+
+        For an input list of vectors of size (n, 3) returns a list of magnitudes of
+        size (n,).
+
+        Parameters
+        ----------
+        vectors : np.ndarray
+            Array of vectors to calculate the magnitude of.
+
+        Returns
+        -------
+        np.ndarray
+            Array of magnitudes of the input vectors.
+        """
+        return np.zeros(vectors.shape[0])  # type: ignore
+
+    @staticmethod
+    def apply_offsets(vectors: np.ndarray, offsets: np.ndarray) -> np.ndarray:
+        """
+        Apply the offsets to the vectors by adding them together.
+
+        These offsets are used to shift the vectors in the x, y, and z directions.
+        They can either be provided through a custom offsets datafile, or calculated
+        using a gradiometry algorithm.
+
+        Parameters
+        ----------
+        vectors : np.ndarray
+            Array of vectors to apply the offsets to. Should be of shape (n, 3) where n
+            is the number of vectors.
+        offsets : np.ndarray
+            Array of offsets to apply to the vectors. Should be of shape (n, 3) where n
+            is the number of vectors.
+
+        Returns
+        -------
+        np.ndarray
+            Array of vectors with offsets applied. Should be of shape (n, 3).
+        """
+        if vectors.shape[0] != offsets.shape[0]:
+            raise ValueError("Vectors and offsets must have the same length.")
+
+        offset_vectors: np.ndarray = vectors[:, :3] + offsets
+
+        # TODO: CDF files don't have NaNs. Emailed MAG to ask what this will look like.
+        # Any values where offsets is nan must also be nan
+        offset_vectors[np.isnan(offsets).any(axis=1)] = np.nan
+
+        return offset_vectors
+
+    @staticmethod
+    def shift_timestamps(epoch: np.ndarray, timedelta: np.ndarray) -> np.ndarray:
+        """
+        Shift the timestamps by the given timedelta.
+
+        If timedelta is positive, the epochs are shifted forward in time.
+
+        Parameters
+        ----------
+        epoch : np.ndarray
+            Array of timestamps to shift. Should be of length n.
+        timedelta : np.ndarray
+            Array of time deltas to shift the timestamps by. Should be the same length
+            as epoch. Given in seconds.
+
+        Returns
+        -------
+        np.ndarray
+            Shifted timestamps.
+        """
+        if epoch.shape[0] != timedelta.shape[0]:
+            raise ValueError(
+                "Input Epoch and offsets timedeltas must be the same length."
+            )
+
+        timedelta_ns = timedelta * 1e9
+        shifted_timestamps = epoch + timedelta_ns
+        return shifted_timestamps
+
+    def truncate_to_24h(self, timestamp: str) -> None:
+        """
+        Truncate all data to a 24 hour period.
+
+        24 hours is given by timestamp in the format YYYYmmdd.
+
+        Parameters
+        ----------
+        timestamp : str
+            Timestamp in the format YYYYMMDD.
+        """
+        pass
+
+    def generate_dataset(
+        self,
+        attribute_manager: ImapCdfAttributes,
+        frame: ValidFrames = ValidFrames.dsrf,
+    ) -> xr.Dataset:
+        """
+        Generate an xarray dataset from the dataclass.
+
+        This method can be used for L2 and L1D, since they have extremely similar
+        output.
+
+        Parameters
+        ----------
+        attribute_manager : ImapCdfAttributes
+            CDF attributes object for the correct level.
+        frame : ValidFrames
+            SPICE reference frame to rotate the data into.
+
+        Returns
+        -------
+        xr.Dataset
+            Complete dataset ready to write to CDF file.
+        """
+        logical_source_id = f"imap_mag_l2_{self.data_mode.value.lower()}-{frame.name}"
+        direction = xr.DataArray(
+            np.arange(3),
+            name="direction",
+            dims=["direction"],
+            attrs=attribute_manager.get_variable_attributes(
+                "direction_attrs", check_schema=False
+            ),
+        )
+
+        direction_label = xr.DataArray(
+            direction.values.astype(str),
+            name="direction_label",
+            dims=["direction_label"],
+            attrs=attribute_manager.get_variable_attributes(
+                "direction_label", check_schema=False
+            ),
+        )
+
+        epoch_time = xr.DataArray(
+            self.epoch,
+            name="epoch",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("epoch"),
+        )
+
+        vectors = xr.DataArray(
+            self.vectors,
+            name="vectors",
+            dims=["epoch", "direction"],
+            attrs=attribute_manager.get_variable_attributes("vector_attrs"),
+        )
+
+        quality_flags = xr.DataArray(
+            self.quality_flags,
+            name="quality_flags",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("compression"),
+        )
+
+        quality_bitmask = xr.DataArray(
+            self.quality_flags,
+            name="quality_flags",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("compression"),
+        )
+
+        rng = xr.DataArray(
+            self.range,
+            name="range",
+            dims=["epoch"],
+            # TODO temp attrs
+            attrs=attribute_manager.get_variable_attributes("compression_width"),
+        )
+
+        magnitude = xr.DataArray(
+            self.magnitude,
+            name="magnitude",
+            dims=["epoch"],
+            attrs=attribute_manager.get_variable_attributes("compression_width"),
+        )
+
+        global_attributes = (
+            attribute_manager.get_global_attributes(logical_source_id)
+            | self.global_attributes
+        )
+
+        output = xr.Dataset(
+            coords={
+                "epoch": epoch_time,
+                "direction": direction,
+                "direction_label": direction_label,
+            },
+            attrs=global_attributes,
+        )
+
+        output["vectors"] = vectors
+        output["quality_flags"] = quality_flags
+        output["quality_bitmask"] = quality_bitmask
+        output["range"] = rng
+        output["magnitude"] = magnitude
+
+        return output

imap_processing/spacecraft/quaternions.py

@@ -66,7 +66,7 @@ def assemble_quaternions(ds: xr.Dataset) -> xr.Dataset:
     base_name = "FSW_ACS_QUAT_10_HZ_BUFFERED".lower()
     for quat_i, label in enumerate(["x", "y", "z", "s"]):
         # 0, 1, 2, .. 9 // 10, 11, 12, .. 19 // 20, 21, 22, .. 29 // 30, 31, 32, .. 39
-        names = [f"{base_name}_{i + quat_i*10}" for i in range(10)]
+        names = [f"{base_name}_{i + quat_i * 10}" for i in range(10)]
         quat = np.stack([ds[name] for name in names], axis=1).ravel()
         output_ds[f"quat_{label}"] = ("epoch", quat)
     return output_ds
@@ -100,8 +100,6 @@ def process_quaternions(packet_file: Path | str) -> tuple[xr.Dataset, xr.Dataset
     # Update dataset global attributes
     attr_mgr = ImapCdfAttributes()
     attr_mgr.add_instrument_global_attrs("spacecraft")
-    # TODO: Allow version to be passed in
-    attr_mgr.add_global_attribute("Data_version", 1)
     attr_mgr.add_instrument_variable_attrs(instrument="spacecraft", level=None)
 
     l1a_ds.attrs.update(

imap_processing/spice/geometry.py

@@ -63,7 +63,7 @@ class SpiceFrame(IntEnum):
 
 
 BORESIGHT_LOOKUP = {
-    SpiceFrame.IMAP_LO: np.array([0, -1, 0]),
+    SpiceFrame.IMAP_LO_BASE: np.array([0, -1, 0]),
     SpiceFrame.IMAP_HI_45: np.array([0, 1, 0]),
     SpiceFrame.IMAP_HI_90: np.array([0, 1, 0]),
     SpiceFrame.IMAP_ULTRA_45: np.array([0, 0, 1]),
@@ -136,7 +136,7 @@ def get_spacecraft_to_instrument_spin_phase_offset(instrument: SpiceFrame) -> fl
     """
     # TODO: Implement retrieval from SPICE?
     offset_lookup = {
-        SpiceFrame.IMAP_LO: 330 / 360,
+        SpiceFrame.IMAP_LO_BASE: 330 / 360,
         SpiceFrame.IMAP_HI_45: 255 / 360,
         SpiceFrame.IMAP_HI_90: 285 / 360,
         SpiceFrame.IMAP_ULTRA_45: 33 / 360,
@@ -325,7 +325,7 @@ def instrument_pointing(
     """
     Compute the instrument pointing at the specified times.
 
-    By default, the coordinates returned are Latitude/Longitude coordinates in
+    By default, the coordinates returned are (Longitude, Latitude) coordinates in
     the reference frame `to_frame`. Cartesian coordinates can be returned if
     desired by setting `cartesian=True`.
 

imap_processing/spice/kernels.py

@@ -3,14 +3,9 @@
 import functools
 import logging
 import os
-from collections.abc import Generator
-from contextlib import contextmanager
-from pathlib import Path
 from typing import Any, Callable, Optional, Union, overload
 
-import numpy as np
 import spiceypy
-from numpy.typing import NDArray
 from spiceypy.utils.exceptions import SpiceyError
 
 from imap_processing import imap_module_directory
@@ -181,277 +176,6 @@ def ensure_spice(
         return _decorator
 
 
-@contextmanager
-def open_spice_ck_file(pointing_frame_path: Path) -> Generator[int, None, None]:
-    """
-    Context manager for handling SPICE CK files.
-
-    Parameters
-    ----------
-    pointing_frame_path : str
-        Path to the CK file.
-
-    Yields
-    ------
-    handle : int
-        Handle to the opened CK file.
-    """
-    # TODO: We will need to figure out if ck kernel changes
-    # and how that will affect appending to the pointing
-    # frame kernel.
-    if pointing_frame_path.exists():
-        handle = spiceypy.dafopw(str(pointing_frame_path))
-    else:
-        handle = spiceypy.ckopn(str(pointing_frame_path), "CK", 0)
-    try:
-        yield handle
-    finally:
-        spiceypy.ckcls(handle)
-
-
-@ensure_spice
-def create_pointing_frame(pointing_frame_path: Path, ck_path: Path) -> None:
-    """
-    Create the pointing frame.
-
-    Parameters
-    ----------
-    pointing_frame_path : pathlib.Path
-        Location of pointing frame kernel.
-    ck_path : pathlib.Path
-        Location of the CK kernel.
-
-    Notes
-    -----
-    Kernels required to be furnished:
-    "imap_science_0001.tf",
-    "imap_sclk_0000.tsc",
-    "imap_sim_ck_2hr_2secsampling_with_nutation.bc" or
-    "sim_1yr_imap_attitude.bc",
-    "imap_wkcp.tf",
-    "naif0012.tls"
-
-    Assumptions:
-    - The MOC has removed timeframe in which nutation/procession are present.
-    TODO: We may come back and have a check for this.
-    - We will continue to append to the pointing frame kernel.
-    TODO: Figure out how we want to handle the file size becoming too large.
-    - For now we can only furnish a single ck kernel.
-    TODO: This will not be the case once we add the ability to query the .csv.
-
-    References
-    ----------
-    https://numpydoc.readthedocs.io/en/latest/format.html#references
-    """
-    # Get IDs.
-    # https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy.gipool
-    id_imap_dps = spiceypy.gipool("FRAME_IMAP_DPS", 0, 1)
-    id_imap_sclk = spiceypy.gipool("CK_-43000_SCLK", 0, 1)
-
-    # Verify that only ck_path kernel is loaded.
-    count = spiceypy.ktotal("ck")
-    loaded_ck_kernel, _, _, _ = spiceypy.kdata(count - 1, "ck")
-
-    if count != 1 or str(ck_path) != loaded_ck_kernel:
-        raise ValueError(f"Error: Expected CK kernel {ck_path}")
-
-    # If the pointing frame kernel already exists, find the last time.
-    if pointing_frame_path.exists():
-        # Get the last time in the pointing frame kernel.
-        pointing_cover = spiceypy.ckcov(
-            str(pointing_frame_path), int(id_imap_dps), True, "SEGMENT", 0, "TDB"
-        )
-        num_segments = spiceypy.wncard(pointing_cover)
-        _, et_end_pointing_frame = spiceypy.wnfetd(pointing_cover, num_segments - 1)
-    else:
-        et_end_pointing_frame = None
-
-    # TODO: Query for .csv file to get the pointing start and end times.
-    # TODO: Remove next four lines once query is added.
-    id_imap_spacecraft = spiceypy.gipool("FRAME_IMAP_SPACECRAFT", 0, 1)
-    ck_cover = spiceypy.ckcov(
-        str(ck_path), int(id_imap_spacecraft), True, "INTERVAL", 0, "TDB"
-    )
-    num_intervals = spiceypy.wncard(ck_cover)
-
-    with open_spice_ck_file(pointing_frame_path) as handle:
-        # TODO: this will change to the number of pointings.
-        for i in range(num_intervals):
-            # Get the coverage window
-            # TODO: this will change to pointing start and end time.
-            et_start, et_end = spiceypy.wnfetd(ck_cover, i)
-            et_times = _get_et_times(et_start, et_end)
-
-            # TODO: remove after query is added.
-            if (
-                et_end_pointing_frame is not None
-                and et_times[0] < et_end_pointing_frame
-            ):
-                break
-
-            # Create a rotation matrix
-            rotation_matrix = _create_rotation_matrix(et_times)
-
-            # Convert the rotation matrix to a quaternion.
-            # https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy.m2q
-            q_avg = spiceypy.m2q(rotation_matrix)
-
-            # https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy.sce2c
-            # Convert start and end times to SCLK.
-            sclk_begtim = spiceypy.sce2c(int(id_imap_sclk), et_times[0])
-            sclk_endtim = spiceypy.sce2c(int(id_imap_sclk), et_times[-1])
-
-            # Create the pointing frame kernel.
-            # https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy.ckw02
-            spiceypy.ckw02(
-                # Handle of an open CK file.
-                handle,
-                # Start time of the segment.
-                sclk_begtim,
-                # End time of the segment.
-                sclk_endtim,
-                # Pointing frame ID.
-                int(id_imap_dps),
-                # Reference frame.
-                "ECLIPJ2000",  # Reference frame
-                # Identifier.
-                "IMAP_DPS",
-                # Number of pointing intervals.
-                1,
-                # Start times of individual pointing records within segment.
-                # Since there is only a single record this is equal to sclk_begtim.
-                np.array([sclk_begtim]),
-                # End times of individual pointing records within segment.
-                # Since there is only a single record this is equal to sclk_endtim.
-                np.array([sclk_endtim]),  # Single stop time
-                # Average quaternion.
-                q_avg,
-                # 0.0 Angular rotation terms.
-                np.array([0.0, 0.0, 0.0]),
-                # Rates (seconds per tick) at which the quaternion and
-                # angular velocity change.
-                np.array([1.0]),
-            )
-
-
-def _get_et_times(et_start: float, et_end: float) -> NDArray[np.float64]:
-    """
-    Get times for pointing start and stop.
-
-    Parameters
-    ----------
-    et_start : float
-        Pointing start time.
-    et_end : float
-        Pointing end time.
-
-    Returns
-    -------
-    et_times : numpy.ndarray
-        Array of times between et_start and et_end.
-    """
-    # TODO: Queried pointing start and stop times here.
-    # TODO removing the @ensure_spice decorator when using the repointing table.
-
-    # 1 spin/15 seconds; 10 quaternions / spin.
-    num_samples = (et_end - et_start) / 15 * 10
-    # There were rounding errors when using spiceypy.pxform so np.ceil and np.floor
-    # were used to ensure the start and end times were included in the array.
-    et_times = np.linspace(
-        np.ceil(et_start * 1e6) / 1e6, np.floor(et_end * 1e6) / 1e6, int(num_samples)
-    )
-
-    return et_times
-
-
-@ensure_spice
-def _average_quaternions(et_times: np.ndarray) -> NDArray:
-    """
-    Average the quaternions.
-
-    Parameters
-    ----------
-    et_times : numpy.ndarray
-        Array of times between et_start and et_end.
-
-    Returns
-    -------
-    q_avg : np.ndarray
-        Average quaternion.
-    """
-    aggregate = np.zeros((4, 4))
-    for tdb in et_times:
-        # we use a quick and dirty method here for grabbing the quaternions
-        # from the attitude kernel.  Depending on how well the kernel input
-        # data is built and sampled, there may or may not be aliasing with this
-        # approach.  If it turns out that we need to pull the quaternions
-        # directly from the CK there are several routines that exist to do this
-        # but it's not straight forward.  We'll revisit this if needed.
-
-        # Rotation matrix from IMAP spacecraft frame to ECLIPJ2000.
-        # https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy.pxform
-        body_rots = spiceypy.pxform("IMAP_SPACECRAFT", "ECLIPJ2000", tdb)
-        # Convert rotation matrix to quaternion.
-        # https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy.m2q
-        body_quat = spiceypy.m2q(body_rots)
-
-        # Standardize the quaternion so that they may be compared.
-        body_quat = body_quat * np.sign(body_quat[0])
-        # Aggregate quaternions into a single matrix.
-        aggregate += np.outer(body_quat, body_quat)
-
-    # Reference: "On Averaging Rotations".
-    # Link: https://link.springer.com/content/pdf/10.1023/A:1011129215388.pdf
-    aggregate /= len(et_times)
-
-    # Compute eigen values and vectors of the matrix A
-    # Eigenvalues tell you how much "influence" each
-    # direction (eigenvector) has.
-    # The largest eigenvalue corresponds to the direction
-    # that has the most influence.
-    # The eigenvector corresponding to the largest
-    # eigenvalue points in the direction that has the most
-    # combined rotation influence.
-    eigvals, eigvecs = np.linalg.eig(aggregate)
-    # q0: The scalar part of the quaternion.
-    # q1, q2, q3: The vector part of the quaternion.
-    q_avg = eigvecs[:, np.argmax(eigvals)]
-
-    return q_avg
-
-
-def _create_rotation_matrix(et_times: np.ndarray) -> NDArray:
-    """
-    Create a rotation matrix.
-
-    Parameters
-    ----------
-    et_times : numpy.ndarray
-        Array of times between et_start and et_end.
-
-    Returns
-    -------
-    rotation_matrix : np.ndarray
-        Rotation matrix.
-    """
-    # Averaged quaternions.
-    q_avg = _average_quaternions(et_times)
-
-    # Converts the averaged quaternion (q_avg) into a rotation matrix
-    # and get inertial z axis.
-    # https://spiceypy.readthedocs.io/en/main/documentation.html#spiceypy.spiceypy.q2m
-    z_avg = spiceypy.q2m(list(q_avg))[:, 2]
-    # y_avg is perpendicular to both z_avg and the standard Z-axis.
-    y_avg = np.cross(z_avg, [0, 0, 1])
-    # x_avg is perpendicular to y_avg and z_avg.
-    x_avg = np.cross(y_avg, z_avg)
-
-    # Construct the rotation matrix from x_avg, y_avg, z_avg
-    rotation_matrix = np.asarray([x_avg, y_avg, z_avg])
-
-    return rotation_matrix
-
-
 def furnish_time_kernel() -> None:
     """Furnish the time kernels."""
     spice_test_data_path = imap_module_directory / "tests/spice/test_data"