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

imap_processing/mag/l1b/mag_l1b.py

@@ -1,5 +1,6 @@
 """MAG L1B Processing."""
 
+import logging
 from pathlib import Path
 
 import numpy as np
@@ -9,8 +10,12 @@ from xarray import Dataset
 from imap_processing.cdf.imap_cdf_manager import ImapCdfAttributes
 from imap_processing.cdf.utils import load_cdf
 
+logger = logging.getLogger(__name__)
 
-def mag_l1b(input_dataset: xr.Dataset, version: str) -> Dataset:
+
+def mag_l1b(
+    input_dataset: xr.Dataset, version: str, calibration_dataset: xr.Dataset = None
+) -> Dataset:
     """
     Will process MAG L1B data from L1A data.
 
@@ -20,6 +25,10 @@ def mag_l1b(input_dataset: xr.Dataset, version: str) -> Dataset:
         The input dataset to process.
     version : str
         The version of the output data.
+    calibration_dataset : xr.Dataset
+        The calibration dataset containing calibration matrices and timeshift values for
+        mago and magi.
+        When None, this defaults to the test calibration file.
 
     Returns
     -------
@@ -28,62 +37,85 @@ def mag_l1b(input_dataset: xr.Dataset, version: str) -> Dataset:
     """
     # TODO:
     # Read in calibration file
-    # multiply all vectors by calibration file
-    if "raw" in input_dataset.attrs["Logical_source"]:
+
+    # TODO: This should definitely be loaded from AWS
+    if calibration_dataset is None:
+        calibration_dataset = load_cdf(
+            Path(__file__).parent / "imap_calibration_mag_20240229_v01.cdf"
+        )
+        logger.info("Using default test calibration file.")
+
+    source = input_dataset.attrs["Logical_source"]
+    if isinstance(source, list):
+        source = source[0]
+
+    if "raw" in source:
         # Raw files should not be processed in L1B.
         raise ValueError("Raw L1A file passed into L1B. Unable to process.")
 
-    output_dataset = mag_l1b_processing(input_dataset)
-    attribute_manager = ImapCdfAttributes()
-    attribute_manager.add_instrument_global_attrs("mag")
-    attribute_manager.add_global_attribute("Data_version", version)
-
-    # Variable attributes can remain the same as L1A
-    input_logical_source = input_dataset.attrs["Logical_source"]
-    if isinstance(input_dataset.attrs["Logical_source"], list):
-        input_logical_source = input_dataset.attrs["Logical_source"][0]
+    mag_attributes = ImapCdfAttributes()
+    mag_attributes.add_instrument_global_attrs("mag")
+    mag_attributes.add_instrument_variable_attrs("mag", "l1b")
+    mag_attributes.add_global_attribute("Data_version", version)
+    source = source.replace("l1a", "l1b")
 
-    logical_source = input_logical_source.replace("l1a", "l1b")
-    output_dataset.attrs = attribute_manager.get_global_attributes(logical_source)
+    output_dataset = mag_l1b_processing(
+        input_dataset, calibration_dataset, mag_attributes, source
+    )
 
     return output_dataset
 
 
-def mag_l1b_processing(input_dataset: xr.Dataset) -> xr.Dataset:
+def mag_l1b_processing(
+    input_dataset: xr.Dataset,
+    calibration_dataset: xr.Dataset,
+    mag_attributes: ImapCdfAttributes,
+    logical_source: str,
+) -> xr.Dataset:
     """
     Will process MAG L1B data from L1A data.
 
     MAG L1B is almost identical to L1A, with only the vectors and attributes getting
     updated. All non-vector variables are the same.
 
+    This step rescales the vector data according to the compression width, and then
+    multiplies the vector according to the calibration matrix for a given range. It
+    also shifts the timestamps by the values defined in calibration_dataset.
+
     Parameters
     ----------
     input_dataset : xr.Dataset
         The input dataset to process.
+    calibration_dataset : xr.Dataset
+        The calibration dataset containing calibration matrices and timeshift values for
+        mago and magi.
+    mag_attributes : ImapCdfAttributes
+        Attribute class for output CDF containing MAG L1B attributes.
+    logical_source : str
+        The expected logical source of the output file. Should look something like:
+        imap_mag_l1b_norm-magi.
 
     Returns
     -------
     output_dataset : xr.Dataset
         L1b dataset.
     """
-    # TODO: There is a time alignment step that will add a lot of complexity.
-    # This needs to be done once we have some SPICE time data.
-
-    mag_attributes = ImapCdfAttributes()
-    mag_attributes.add_instrument_variable_attrs("mag", "l1")
-
     dims = [["direction"], ["compression"]]
     new_dims = [["direction"], ["compression"]]
-    # TODO: This should definitely be loaded from AWS
-    calibration_dataset = load_cdf(
-        Path(__file__).parent / "imap_calibration_mag_20240229_v01.cdf"
-    )
-    # TODO: add time shift
-    # TODO: Check validity of time range for calibration
-    if "mago" in input_dataset.attrs["Logical_source"][0]:
+
+    if "mago" in logical_source:
         calibration_matrix = calibration_dataset["MFOTOURFO"]
-    else:
+        time_shift = calibration_dataset["OTS"]
+    elif "magi" in logical_source:
         calibration_matrix = calibration_dataset["MFITOURFI"]
+        time_shift = calibration_dataset["ITS"]
+    else:
+        raise ValueError(
+            f"Calibration matrix not found, invalid logical source "
+            f"{input_dataset.attrs['Logical_source']}"
+        )
+
+    # TODO: Check validity of time range for calibration
 
     l1b_fields = xr.apply_ufunc(
         update_vector,
@@ -96,23 +128,85 @@ def mag_l1b_processing(input_dataset: xr.Dataset) -> xr.Dataset:
         kwargs={"calibration_matrix": calibration_matrix},
     )
 
-    output_dataset = input_dataset.copy()
-    output_dataset["vectors"].data = l1b_fields[0].data
+    epoch_time = shift_time(input_dataset["epoch"], time_shift)
+    epoch_time.attrs = mag_attributes.get_variable_attributes("epoch")
 
-    output_dataset["epoch"].attrs = mag_attributes.get_variable_attributes("epoch")
-    output_dataset["direction"].attrs = mag_attributes.get_variable_attributes(
-        "direction_attrs"
+    direction = xr.DataArray(
+        np.arange(4),
+        name="direction",
+        dims=["direction"],
+        attrs=mag_attributes.get_variable_attributes(
+            "direction_attrs", check_schema=False
+        ),
     )
-    output_dataset["compression"].attrs = mag_attributes.get_variable_attributes(
-        "compression_attrs"
+
+    compression = xr.DataArray(
+        np.arange(2),
+        name="compression",
+        dims=["compression"],
+        attrs=mag_attributes.get_variable_attributes(
+            "compression_attrs", check_schema=False
+        ),
     )
-    output_dataset["direction_label"].attrs = mag_attributes.get_variable_attributes(
-        "direction_label", check_schema=False
+
+    direction_label = xr.DataArray(
+        direction.values.astype(str),
+        name="direction_label",
+        dims=["direction_label"],
+        attrs=mag_attributes.get_variable_attributes(
+            "direction_label", check_schema=False
+        ),
     )
-    output_dataset["compression_label"].attrs = mag_attributes.get_variable_attributes(
-        "compression_label", check_schema=False
+
+    compression_label = xr.DataArray(
+        compression.values.astype(str),
+        name="compression_label",
+        dims=["compression_label"],
+        attrs=mag_attributes.get_variable_attributes(
+            "compression_label", check_schema=False
+        ),
     )
 
+    global_attributes = mag_attributes.get_global_attributes(logical_source)
+    try:
+        global_attributes["is_mago"] = input_dataset.attrs["is_mago"]
+        global_attributes["is_active"] = input_dataset.attrs["is_active"]
+        global_attributes["vectors_per_second"] = input_dataset.attrs[
+            "vectors_per_second"
+        ]
+        global_attributes["missing_sequences"] = input_dataset.attrs[
+            "missing_sequences"
+        ]
+    except KeyError as e:
+        logger.info(
+            f"Key error when assigning global attributes, attribute not found in "
+            f"L1A file: {e}"
+        )
+
+    output_dataset = xr.Dataset(
+        coords={
+            "epoch": epoch_time,
+            "direction": direction,
+            "compression": compression,
+            "direction_label": direction_label,
+            "compression_label": compression_label,
+        },
+        attrs=global_attributes,
+    )
+    # Fill the output with data
+    output_dataset["vectors"] = xr.DataArray(
+        l1b_fields[0].data,
+        name="vectors",
+        dims=["epoch", "direction"],
+        attrs=mag_attributes.get_variable_attributes("vector_attrs"),
+    )
+
+    output_dataset["compression_flags"] = xr.DataArray(
+        input_dataset["compression_flags"].data,
+        name="compression_flags",
+        dims=["epoch", "compression"],
+        attrs=mag_attributes.get_variable_attributes("compression_flags_attrs"),
+    )
     return output_dataset
 
 
@@ -177,11 +271,13 @@ def rescale_vector(
     output_vector : numpy.ndarray
         Updated vector.
     """
-    if not compression_flags[0]:
-        return input_vector
-    else:
+    output_vector: np.ndarray = input_vector.astype(np.float64)
+
+    if compression_flags[0]:
         factor = np.float_power(2, (16 - compression_flags[1]))
-        return input_vector * factor  # type: ignore
+        output_vector[:3] = input_vector.astype(np.float64)[:3] * factor
+
+    return output_vector
 
 
 def calibrate_vector(
@@ -206,8 +302,54 @@ def calibrate_vector(
     updated_vector : numpy.ndarray
         Calibrated vector.
     """
-    updated_vector = input_vector.copy()
-    updated_vector[:3] = np.dot(
-        calibration_matrix.values[:, :, int(input_vector[3])], input_vector[:3]
-    )
+    updated_vector: np.ndarray = input_vector.copy()
+    if input_vector[3] % 1 != 0:
+        raise ValueError("Range must be an integer.")
+
+    range = int(input_vector[3])
+    x_y_z = input_vector[:3]
+    updated_vector[:3] = np.dot(calibration_matrix.values[:, :, range], x_y_z)
     return updated_vector
+
+
+def shift_time(epoch_times: xr.DataArray, time_shift: xr.DataArray) -> xr.DataArray:
+    """
+    Shift epoch times by the provided time_shift calibration value.
+
+    Sometimes the time values calculated from the sensor vary slightly from the "actual"
+    time the data was captured. To correct for this, the MAG team provides time shift
+    values in the calibration file. This function applies the time shift to the epoch
+    times.
+
+    The time shift is provided in seconds. A positive shift is adding time, while a
+    negative shift subtracts it (so the values move backwards.)
+
+    This may mean vectors shift out of the specific day that is being processed. To
+    manage this, all MAG L0, L1A, L1B, and L1C science data files contain an extra 30
+    minute buffer on either side (so the data ranges from
+    midnight - 30 minutes to midnight + 24 hours + 30 minutes.)
+    The extra buffer is removed at L1D and L2 so those science files are exactly 24
+    hours long.
+
+    For more information please refer to the algorithm document.
+
+    Parameters
+    ----------
+    epoch_times : xr.DataArray
+        The input epoch times, in J2000 ns.
+    time_shift : xr.DataArray
+        The time shift to apply for the given sensor. This should be one value and is
+        in seconds.
+
+    Returns
+    -------
+    shifted_times : xr.DataArray
+        The shifted epoch times, equal to epoch_times with time_shift added to each
+        value.
+    """
+    if time_shift.size != 1:
+        raise ValueError("Time shift must be a single value.")
+    # Time shift is in seconds
+    time_shift_ns = time_shift.data * 1e9
+
+    return epoch_times + time_shift_ns

imap_processing/mag/l1c/interpolation_methods.py

@@ -0,0 +1,211 @@
+# mypy: ignore-errors
+"""Module containing interpolation methods for MAG L1C."""
+
+from enum import Enum
+
+import numpy as np
+from scipy.interpolate import make_interp_spline
+
+
+def linear(
+    input_vectors: np.ndarray,
+    input_timestamps: np.ndarray,
+    output_timestamps: np.ndarray,
+) -> np.ndarray:
+    """
+    Linear interpolation of input vectors to output timestamps.
+
+    Parameters
+    ----------
+    input_vectors : numpy.ndarray
+        Input vectors of shape (n, 3) where n is equal to the number of input
+        timestamps. Contains x, y, z components of the vector.
+    input_timestamps : numpy.ndarray
+        Input timestamps of shape (n,) which correspond to the timestamps of the input
+        vectors.
+    output_timestamps : numpy.ndarray
+        Output timestamps of shape (m,) to generate interpolated vectors for.
+
+    Returns
+    -------
+    numpy.ndarray
+        Interpolated vectors of shape (m, 3) where m is equal to the number of output
+        timestamps. Contains x, y, z components of the vector.
+    """
+    spline = make_interp_spline(input_timestamps, input_vectors, k=1)
+    return spline(output_timestamps)
+
+
+def quadratic(
+    input_vectors: np.ndarray,
+    input_timestamps: np.ndarray,
+    output_timestamps: np.ndarray,
+) -> np.ndarray:
+    """
+    Quadratic interpolation of input vectors to output timestamps.
+
+    Parameters
+    ----------
+    input_vectors : numpy.ndarray
+        Input vectors of shape (n, 3) where n is equal to the number of input
+        timestamps. Contains x, y, z components of the vector.
+    input_timestamps : numpy.ndarray
+        Input timestamps of shape (n,) which correspond to the timestamps of the input
+        vectors.
+    output_timestamps : numpy.ndarray
+        Output timestamps of shape (m,) to generate interpolated vectors for.
+
+    Returns
+    -------
+    numpy.ndarray
+        Interpolated vectors of shape (m, 3) where m is equal to the number of output
+        timestamps. Contains x, y, z components of the vector.
+    """
+    spline = make_interp_spline(input_timestamps, input_vectors, k=2)
+    return spline(output_timestamps)
+
+
+def cubic(
+    input_vectors: np.ndarray,
+    input_timestamps: np.ndarray,
+    output_timestamps: np.ndarray,
+) -> np.ndarray:
+    """
+    Cubic interpolation of input vectors to output timestamps.
+
+    Parameters
+    ----------
+    input_vectors : numpy.ndarray
+        Input vectors of shape (n, 3) where n is equal to the number of input
+        timestamps. Contains x, y, z components of the vector.
+    input_timestamps : numpy.ndarray
+        Input timestamps of shape (n,) which correspond to the timestamps of the input
+        vectors.
+    output_timestamps : numpy.ndarray
+        Output timestamps of shape (m,) to generate interpolated vectors for.
+
+    Returns
+    -------
+    numpy.ndarray
+        Interpolated vectors of shape (m, 3) where m is equal to the number of output
+        timestamps. Contains x, y, z components of the vector.
+    """
+    spline = make_interp_spline(input_timestamps, input_vectors, k=3)
+    return spline(output_timestamps)
+
+
+def linear_filtered(
+    input_vectors: np.ndarray,
+    input_timestamps: np.ndarray,
+    output_timestamps: np.ndarray,
+) -> np.ndarray:
+    """
+    Linear filtered interpolation of input vectors to output timestamps.
+
+    Parameters
+    ----------
+    input_vectors : numpy.ndarray
+        Input vectors of shape (n, 3) where n is equal to the number of input
+        timestamps. Contains x, y, z components of the vector.
+    input_timestamps : numpy.ndarray
+        Input timestamps of shape (n,) which correspond to the timestamps of the input
+        vectors.
+    output_timestamps : numpy.ndarray
+        Output timestamps of shape (m,) to generate interpolated vectors for.
+
+    Returns
+    -------
+    numpy.ndarray
+        Interpolated vectors of shape (m, 3) where m is equal to the number of output
+        timestamps. Contains x, y, z components of the vector.
+    """
+    pass
+
+
+def quadratic_filtered(
+    input_vectors: np.ndarray,
+    input_timestamps: np.ndarray,
+    output_timestamps: np.ndarray,
+) -> np.ndarray:
+    """
+    Quadratic filtered interpolation of input vectors to output timestamps.
+
+    Parameters
+    ----------
+    input_vectors : numpy.ndarray
+        Input vectors of shape (n, 3) where n is equal to the number of input
+        timestamps. Contains x, y, z components of the vector.
+    input_timestamps : numpy.ndarray
+        Input timestamps of shape (n,) which correspond to the timestamps of the input
+        vectors.
+    output_timestamps : numpy.ndarray
+        Output timestamps of shape (m,) to generate interpolated vectors for.
+
+    Returns
+    -------
+    numpy.ndarray
+        Interpolated vectors of shape (m, 3) where m is equal to the number of output
+        timestamps. Contains x, y, z components of the vector.
+    """
+    pass
+
+
+def cubic_filtered(
+    input_vectors: np.ndarray,
+    input_timestamps: np.ndarray,
+    output_timestamps: np.ndarray,
+) -> np.ndarray:
+    """
+    Cubic filtered interpolation of input vectors to output timestamps.
+
+    Parameters
+    ----------
+    input_vectors : numpy.ndarray
+        Input vectors of shape (n, 3) where n is equal to the number of input
+        timestamps. Contains x, y, z components of the vector.
+    input_timestamps : numpy.ndarray
+        Input timestamps of shape (n,) which correspond to the timestamps of the input
+        vectors.
+    output_timestamps : numpy.ndarray
+        Output timestamps of shape (m,) to generate interpolated vectors for.
+
+    Returns
+    -------
+    numpy.ndarray
+        Interpolated vectors of shape (m, 3) where m is equal to the number of output
+        timestamps. Contains x, y, z components of the vector.
+    """
+    pass
+
+
+class InterpolationFunction(Enum):
+    """Enum which describes the options for interpolation functions on MAG L1C."""
+
+    linear = (linear,)
+    quadratic = (quadratic,)
+    cubic = (cubic,)
+    linear_filtered = (linear_filtered,)
+    quadratic_filtered = (quadratic_filtered,)
+    cubic_filtered = (cubic_filtered,)
+
+    def __call__(self, *args, **kwargs):
+        """
+        Overwritten call which allows you to call the interpolation function directly.
+
+        example: InterpolationFunction.linear(input_vectors, input_timestamps,
+        output_timestamps)
+
+        Parameters
+        ----------
+        *args : list
+            List of arguments to pass to the interpolation function. Must match
+            (input_vectors, input_timestamps, output_timestamps).
+        **kwargs : dict
+            Keyword arguments to pass to the interpolation function.
+
+        Returns
+        -------
+        Any
+            The return value of the interpolation function.
+        """
+        return self.value[0](*args, **kwargs)