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

imap_processing/mag/l2/mag_l2.py

@@ -0,0 +1,140 @@
+"""Module to run MAG L2 processing."""
+
+import numpy as np
+import xarray as xr
+
+from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
+from imap_processing.mag import imap_mag_sdc_configuration_v001 as configuration
+from imap_processing.mag.constants import DataMode
+from imap_processing.mag.l1b.mag_l1b import calibrate_vector
+from imap_processing.mag.l2.mag_l2_data import MagL2
+
+
+def mag_l2(
+    calibration_datasets: list[xr.Dataset],
+    offsets_dataset: xr.Dataset,
+    input_data: xr.Dataset,
+) -> list[xr.Dataset]:
+    """
+    Complete MAG L2 processing.
+
+    Processing uses 4 data input sources:
+    1. Calibration dataset
+        These calibration files are ancillary files and may require multiple files to
+        cover the entire timespan. They are not expected to change often. They are used
+        to provide rotation matrices to correct the frame of the vectors. The same
+        file(s) are used for both burst and norm calculations.
+    2. Offsets dataset
+        This is one, hand-created file which must correspond exactly to an L1B
+        (for burst) or L1C (for norm) data file. For each vector, this file includes
+        offsets, timedelta, and quality flags. The offsets are added to the vectors,
+        the timedelta is used to correct the epoch time, and the quality flags are
+        directly passed into the output file.
+    3. Input data
+        This is the L1B or L1C data file. It is used to provide the vectors and epoch
+        time. It should always be MAGo in the nominal case, but it is possible that we
+        will switch permanently to using MAGi (in the case of sensor failure, for
+        example.) The offsets dataset and the input
+        data are tightly related, so the input data filename is actually retrieved from
+        the offset dataset to ensure they always match.
+    4. sdc-configuration
+        This is a local configuration file for changes we never expect to make in
+        flight. This is in the IMAP local repo because changes to these settings will
+        require other code updates to validate the changes. In L2, the only setting used
+        is "always_output_mago", which indicates whether we should always output MAGo.
+        Note that if this ever is set to False, we will need to update the dependency
+        system to set MAGi files as an upstream dependency.
+
+    Input data can be burst or normal mode, but MUST match the file in offset_dataset.
+    TODO: retrieve the file from offset_dataset in cli.py.
+    Calibration dataset is the same for all runs.
+
+    MAGi data is not used unless we indicate it.
+
+    Parameters
+    ----------
+    calibration_datasets : list[xr.Dataset]
+        Calibration ancillary file inputs.
+    offsets_dataset : xr.Dataset
+        Offset ancillary file input.
+    input_data : xr.Dataset
+        Input data from MAG L1C or L1B.
+
+    Returns
+    -------
+    list[xr.Dataset]
+        List of xarray datasets ready to write to CDF file. Expected to be four outputs
+        for different frames.
+    """
+    # TODO we may need to combine multiple calibration datasets into one timeline.
+
+    always_output_mago = configuration.ALWAYS_OUTPUT_MAGO
+
+    # TODO Check that the input file matches the offsets file
+    if not np.array_equal(input_data["epoch"].data, offsets_dataset["epoch"].data):
+        raise ValueError("Input file and offsets file must have the same timestamps.")
+
+    calibration_matrix = retrieve_matrix_from_l2_calibration(
+        calibration_datasets, always_output_mago
+    )
+
+    vectors = np.apply_along_axis(
+        func1d=calibrate_vector,
+        axis=1,
+        arr=input_data["vectors"].data,
+        calibration_matrix=calibration_matrix,
+    )
+
+    basic_test_data = MagL2(
+        vectors[:, :3],  # level 2 vectors don't include range
+        input_data["epoch"].data,
+        input_data["vectors"].data[:, 3],
+        {},
+        np.zeros(len(input_data["epoch"].data)),
+        np.zeros(len(input_data["epoch"].data)),
+        DataMode.NORM,
+        offsets=offsets_dataset["offsets"].data,
+        timedelta=offsets_dataset["timedeltas"].data,
+    )
+    attributes = ImapCdfAttributes()
+    attributes.add_instrument_global_attrs("mag")
+    # temporarily point to l1c
+    attributes.add_instrument_variable_attrs("mag", "l1c")
+    return [basic_test_data.generate_dataset(attributes)]
+
+
+def retrieve_matrix_from_l2_calibration(
+    calibration_datasets: list[xr.Dataset], use_mago: bool = True
+) -> xr.DataArray:
+    """
+    Get the calibration matrix for the file.
+
+    Parameters
+    ----------
+    calibration_datasets : list[xr.Dataset]
+        Ancillary file inputs for calibration.
+    use_mago : bool
+        Use the MAGo calibration matrix. Default is True.
+
+    Returns
+    -------
+    np.ndarray
+        Calibration matrix in the shape (3, 3, 4) to rotate vectors.
+    """
+    # TODO: allow for multiple inputs
+    if isinstance(calibration_datasets, list):
+        calibration_dataset = calibration_datasets[0]
+        if len(calibration_datasets) > 1:
+            raise NotImplementedError
+    else:
+        calibration_dataset = calibration_datasets
+
+    if use_mago:
+        calibration_data = calibration_dataset["URFTOORFO"]
+    else:
+        calibration_data = calibration_dataset["URFTOORFI"]
+
+    # TODO will need to combine multiple files here
+    # TODO: Check validity of the calibration file?
+
+    return calibration_data

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/quality_flags.py

@@ -52,6 +52,7 @@ class ImapAttitudeUltraFlags(FlagNameMixin):
 
     NONE = CommonFlags.NONE
     SPINRATE = 2**0  # bit 0
+    AUXMISMATCH = 2**1  # bit 1 # aux packet does not match Universal Spin Table
 
 
 class ImapRatesUltraFlags(FlagNameMixin):