Features

aurora/pipelines/process_mth5.py

@@ -6,7 +6,6 @@
 can be repurposed for other TF estimation schemes.  The "legacy" version
 corresponds to aurora default processing.
 
-
 Notes on process_mth5_legacy:
 Note 1: process_mth5 assumes application of cascading decimation, and that the
 decimated data will be accessed from the previous decimation level.  This should be
@@ -22,7 +21,7 @@
 
 Note 3: This point in the loop marks the interface between _generation_ of the FCs and
  their _usage_. In future the code above this comment would be pushed into
- create_fourier_coefficients() and the code below this would access those FCs and
+ the creation of the spectrograms and the code below this would access those FCs and
  execute compute_transfer_function().
  This would also be an appropriate place to place a feature extraction layer, and
  compute weights for the FCs.
@@ -93,6 +92,49 @@ def make_stft_objects(processing_config, i_dec_level, run_obj, run_xrds, units="
     -------
     stft_obj: xarray.core.dataset.Dataset
         Time series of calibrated Fourier coefficients per each channel in the run
+
+    Development Notes:
+    Here are the parameters that are defined via the mt_metadata fourier coefficients structures:
+
+    "bands",
+    "decimation.anti_alias_filter": "default",
+    "decimation.factor": 4.0,
+    "decimation.level": 2,
+    "decimation.method": "default",
+    "decimation.sample_rate": 0.0625,
+    "stft.per_window_detrend_type": "linear",
+    "stft.prewhitening_type": "first difference",
+    "stft.window.clock_zero_type": "ignore",
+    "stft.window.num_samples": 128,
+    "stft.window.overlap": 32,
+    "stft.window.type": "boxcar"
+
+    Creating the decimations config requires a decision about decimation factors and the number of levels.
+    We have been getting this from the EMTF band setup file by default.  It is desirable to continue supporting this,
+    however, note that the EMTF band setup is really about a time series operation, and not about making STFTs.
+
+    For the record, here is the legacy decimation config from EMTF, a.k.a. decset.cfg:
+    ```
+    4     0      # of decimation level, & decimation offset
+    128  32.   1   0   0   7   4   32   1
+    1.0
+    128  32.   4   0   0   7   4   32   4
+    .2154  .1911   .1307   .0705
+    128  32.   4   0   0   7   4   32   4
+    .2154  .1911   .1307   .0705
+    128  32.   4   0   0   7   4   32   4
+    .2154  .1911   .1307   .0705
+    ```
+
+    This essentially corresponds to a "Decimations Group" which is a list of decimations.
+    Related to the generation of FCs is the ARMA prewhitening (Issue #60) which was controlled in
+    EMTF with pwset.cfg
+    4    5             # of decimation level, # of channels
+    3 3 3 3 3
+    3 3 3 3 3
+    3 3 3 3 3
+    3 3 3 3 3
+
     """
     stft_config = processing_config.get_decimation_level(i_dec_level)
     stft_obj = run_ts_to_stft(stft_config, run_xrds)
@@ -278,7 +320,7 @@ def load_stft_obj_from_mth5(
     """
     Load stft_obj from mth5 (instead of compute)
 
-    Note #1: See note #1 in time_series.frequency_band_helpers.extract_band
+    Note #1: See note #1 in mth5.timeseries.spectre.spectrogram.py in extract_band function.
 
     Parameters
     ----------

aurora/pipelines/time_series_helpers.py

@@ -19,180 +19,11 @@
     Decimation as FCDecimation,
 )
 from mth5.groups import RunGroup
+from mth5.timeseries.spectre.prewhitening import apply_prewhitening
+from mth5.timeseries.spectre.prewhitening import apply_recoloring
 from typing import Literal, Optional, Union
 
 
-def apply_prewhitening(
-    decimation_obj: Union[AuroraDecimationLevel, FCDecimation],
-    run_xrds_input: xr.Dataset,
-) -> xr.Dataset:
-    """
-    Applies pre-whitening to time series to avoid spectral leakage when FFT is applied.
-
-    TODO: If "first difference", consider clipping first and last sample from the
-     differentiated time series.
-
-    Parameters
-    ----------
-    decimation_obj : Union[AuroraDecimationLevel, FCDecimation]
-        Information about how the decimation level is to be processed
-
-    run_xrds_input : xarray.core.dataset.Dataset
-        Time series to be pre-whitened
-
-    Returns
-    -------
-    run_xrds : xarray.core.dataset.Dataset
-        pre-whitened time series
-
-    """
-    # TODO: remove this try/except once mt_metadata issue 238 PR is merged
-    try:
-        pw_type = decimation_obj.prewhitening_type
-    except AttributeError:
-        pw_type = decimation_obj.stft.prewhitening_type
-
-    if not pw_type:
-        msg = "No prewhitening specified - skipping this step"
-        logger.info(msg)
-        return run_xrds_input
-
-    if pw_type == "first difference":
-        run_xrds = run_xrds_input.differentiate("time")
-    else:
-        msg = f"{pw_type} pre-whitening not implemented"
-        logger.exception(msg)
-        raise NotImplementedError(msg)
-    return run_xrds
-
-
-def apply_recoloring(
-    decimation_obj: Union[AuroraDecimationLevel, FCDecimation],
-    stft_obj: xr.Dataset,
-) -> xr.Dataset:
-    """
-    Inverts the pre-whitening operation in frequency domain.
-
-    Parameters
-    ----------
-    decimation_obj : mt_metadata.transfer_functions.processing.fourier_coefficients.decimation.Decimation
-        Information about how the decimation level is to be processed
-    stft_obj : xarray.core.dataset.Dataset
-        Time series of Fourier coefficients to be recoloured
-
-
-    Returns
-    -------
-    stft_obj : xarray.core.dataset.Dataset
-        Recolored time series of Fourier coefficients
-    """
-    # TODO: remove this try/except once mt_metadata issue 238 PR is merged
-    try:
-        pw_type = decimation_obj.prewhitening_type
-    except AttributeError:
-        pw_type = decimation_obj.stft.prewhitening_type
-
-    # No recoloring needed if prewhitening not appiled, or recoloring set to False
-    if not pw_type:
-        return stft_obj
-    # TODO Delete after tests (20241220) -- this check has been moved above the call to this function
-    # if not decimation_obj.recoloring:
-    #     return stft_obj
-
-    if pw_type == "first difference":
-        # first difference prewhitening correction is to divide by jw
-        freqs = stft_obj.frequency.data  # was freqs = decimation_obj.fft_frequencies
-        jw = 1.0j * 2 * np.pi * freqs
-        stft_obj /= jw
-
-        # suppress nan and inf to mute later warnings
-        if jw[0] == 0.0:
-            cond = stft_obj.frequency != 0.0
-            stft_obj = stft_obj.where(cond, complex(0.0))
-    # elif decimation_obj.prewhitening_type == "ARMA":
-    #     from statsmodels.tsa.arima.model import ARIMA
-    #     AR = 3 # add this to processing config
-    #     MA = 4 # add this to processing config
-
-    else:
-        msg = f"{pw_type} recoloring not yet implemented"
-        logger.error(msg)
-        raise NotImplementedError(msg)
-
-    return stft_obj
-
-
-def run_ts_to_stft_scipy(
-    decimation_obj: Union[AuroraDecimationLevel, FCDecimation],
-    run_xrds_orig: xr.Dataset,
-) -> xr.Dataset:
-    """
-    TODO: Replace with mth5 run_ts_to_stft_scipy method
-    Converts a runts object into a time series of Fourier coefficients.
-    This method uses scipy.signal.spectrogram.
-
-
-    Parameters
-    ----------
-    decimation_obj : mt_metadata.transfer_functions.processing.aurora.DecimationLevel
-        Information about how the decimation level is to be processed
-        Note: This works with FCdecimation and AuroraDecimationLevel becuase test_fourier_coefficients
-         and test_stft_methods_agree both use them)
-        Note: Both of these objects are basically spectrogram metadata with provenance for decimation levels.
-    run_xrds_orig : : xarray.core.dataset.Dataset
-        Time series to be processed
-
-    Returns
-    -------
-    stft_obj : xarray.core.dataset.Dataset
-        Time series of Fourier coefficients
-    """
-    run_xrds = apply_prewhitening(decimation_obj, run_xrds_orig)
-    windowing_scheme = window_scheme_from_decimation(
-        decimation_obj
-    )  # TODO: deprecate in favor of stft.window.taper
-
-    stft_obj = xr.Dataset()
-    for channel_id in run_xrds.data_vars:
-        ff, tt, specgm = ssig.spectrogram(
-            run_xrds[channel_id].data,
-            fs=decimation_obj.decimation.sample_rate,
-            window=windowing_scheme.taper,
-            nperseg=decimation_obj.stft.window.num_samples,
-            noverlap=decimation_obj.stft.window.overlap,
-            detrend="linear",
-            scaling="density",
-            mode="complex",
-        )
-
-        # drop Nyquist>
-        ff = ff[:-1]
-        specgm = specgm[:-1, :]
-        specgm *= np.sqrt(2)  # compensate energy for keeping only half the spectrum
-
-        # make time_axis
-        tt = tt - tt[0]
-        tt *= decimation_obj.decimation.sample_rate
-        time_axis = run_xrds.time.data[tt.astype(int)]
-
-        xrd = xr.DataArray(
-            specgm.T,
-            dims=["time", "frequency"],
-            coords={"frequency": ff, "time": time_axis},
-        )
-        stft_obj.update({channel_id: xrd})
-
-    # TODO : remove try/except after mt_metadata issue 238 addressed
-    try:
-        to_recolor_or_not_to_recolor = decimation_obj.recoloring
-    except AttributeError:
-        to_recolor_or_not_to_recolor = decimation_obj.stft.recoloring
-    if to_recolor_or_not_to_recolor:
-        stft_obj = apply_recoloring(decimation_obj, stft_obj)
-
-    return stft_obj
-
-
 def truncate_to_clock_zero(
     decimation_obj: Union[AuroraDecimationLevel, FCDecimation],
     run_xrds: RunGroup,
@@ -204,7 +35,7 @@ def truncate_to_clock_zero(
 
     Parameters
     ----------
-    decimation_obj: mt_metadata.transfer_functions.processing.aurora.DecimationLevel
+    decimation_obj: Union[AuroraDecimationLevel, FCDecimation]
         Information about how the decimation level is to be processed
     run_xrds : xarray.core.dataset.Dataset
         normally extracted from mth5.RunTS
@@ -277,7 +108,7 @@ def run_ts_to_stft(
 
     Parameters
     ----------
-    decimation_obj : mt_metadata.transfer_functions.processing.aurora.DecimationLevel
+    decimation_obj : AuroraDecimationLevel
         Information about how the decimation level is to be processed
     run_ts : xarray.core.dataset.Dataset
         normally extracted from mth5.RunTS
@@ -292,7 +123,7 @@ def run_ts_to_stft(
     # need to remove any nans before windowing, or else if there is a single
     # nan then the whole channel becomes nan.
     run_xrds = nan_to_mean(run_xrds_orig)
-    run_xrds = apply_prewhitening(decimation_obj, run_xrds)
+    run_xrds = apply_prewhitening(decimation_obj.stft.prewhitening_type, run_xrds)
     run_xrds = truncate_to_clock_zero(decimation_obj, run_xrds)
     windowing_scheme = window_scheme_from_decimation(decimation_obj)
     windowed_obj = windowing_scheme.apply_sliding_window(
@@ -313,7 +144,7 @@ def run_ts_to_stft(
     )
 
     if decimation_obj.stft.recoloring:
-        stft_obj = apply_recoloring(decimation_obj, stft_obj)
+        stft_obj = apply_recoloring(decimation_obj.stft.prewhitening_type, stft_obj)
 
     return stft_obj
 

aurora/time_series/frequency_band_helpers.py

@@ -6,17 +6,24 @@
 from mt_metadata.transfer_functions.processing.aurora import (
     DecimationLevel as AuroraDecimationLevel,
 )
+from mt_metadata.transfer_functions.processing.aurora import Band
+from mth5.timeseries.spectre.spectrogram import extract_band
+from typing import Optional, Tuple
+import xarray as xr
 
 
 def get_band_for_tf_estimate(
-    band, dec_level_config: AuroraDecimationLevel, local_stft_obj, remote_stft_obj
-):
+    band: Band,
+    dec_level_config: AuroraDecimationLevel,
+    local_stft_obj: xr.Dataset,
+    remote_stft_obj: Optional[xr.Dataset],
+) -> Tuple[xr.Dataset, xr.Dataset, Optional[xr.Dataset]]:
     """
     Returns spectrograms X, Y, RR for harmonics within the given band
 
     Parameters
     ----------
-    band : mt_metadata.transfer_functions.processing.aurora.FrequencyBands
+    band : mt_metadata.transfer_functions.processing.aurora.Band
         object with lower_bound and upper_bound to tell stft object which
         subarray to return
     config : AuroraDecimationLevel
@@ -53,49 +60,6 @@ def get_band_for_tf_estimate(
     return X, Y, RR
 
 
-def extract_band(frequency_band, fft_obj, channels=[], epsilon=1e-7):
-    """
-    Extracts a frequency band from xr.DataArray representing a spectrogram.
-
-    Stand alone version of the method that is used by WIP Spectrogram class.
-
-    Development Notes:
-    #1: 20230902
-    TODO: Decide if base dataset object should be a xr.DataArray (not xr.Dataset)
-    - drop=True does not play nice with h5py and Dataset, results in a type error.
-    File "stringsource", line 2, in h5py.h5r.Reference.__reduce_cython__
-    TypeError: no default __reduce__ due to non-trivial __cinit__
-    However, it works OK with DataArray, so maybe use data array in general
-
-    Parameters
-    ----------
-    frequency_band: mt_metadata.transfer_functions.processing.aurora.band.Band
-        Specifies interval corresponding to a frequency band
-    fft_obj: xarray.core.dataset.Dataset
-        To be replaced with an fft_obj() class in future
-    epsilon: float
-        Use this when you are worried about missing a frequency due to
-        round off error.  This is in general not needed if we use a df/2 pad
-        around true harmonics.
-
-    Returns
-    -------
-    band: xr.DataArray
-        The frequencies within the band passed into this function
-    """
-    cond1 = fft_obj.frequency >= frequency_band.lower_bound - epsilon
-    cond2 = fft_obj.frequency <= frequency_band.upper_bound + epsilon
-    try:
-        band = fft_obj.where(cond1 & cond2, drop=True)
-    except TypeError:  # see Note #1
-        tmp = fft_obj.to_array()
-        band = tmp.where(cond1 & cond2, drop=True)
-        band = band.to_dataset("variable")
-    if channels:
-        band = band[channels]
-    return band
-
-
 def check_time_axes_synched(X, Y):
     """
     Utility function for checking that time axes agree.

aurora/time_series/xarray_helpers.py

@@ -2,13 +2,17 @@
 Placeholder module for methods manipulating xarray time series
 """
 
-import numpy as np
 import xarray as xr
 from loguru import logger
-from typing import Optional, Union
+from typing import Optional
 
 
-def handle_nan(X, Y, RR, drop_dim=""):
+def handle_nan(
+    X: xr.Dataset,
+    Y: Optional[xr.Dataset],
+    RR: Optional[xr.Dataset],
+    drop_dim: Optional[str] = "",
+) -> tuple:
     """
     Drops Nan from multiple channel series'.
 
@@ -87,118 +91,3 @@ def handle_nan(X, Y, RR, drop_dim=""):
     RR = RR.rename(data_var_rm_label_mapper)
 
     return X, Y, RR
-
-
-def covariance_xr(
-    X: xr.DataArray, aweights: Optional[Union[np.ndarray, None]] = None
-) -> xr.DataArray:
-    """
-    Compute the covariance matrix with numpy.cov.
-
-    Parameters
-    ----------
-    X: xarray.core.dataarray.DataArray
-        Multivariate time series as an xarray
-    aweights: array_like, optional
-        Doc taken from numpy cov follows:
-        1-D array of observation vector weights. These relative weights are
-        typically large for observations considered "important" and smaller for
-        observations considered less "important". If ``ddof=0`` the array of
-        weights can be used to assign probabilities to observation vectors.
-
-    Returns
-    -------
-    S: xarray.DataArray
-        The covariance matrix of the data in xarray form.
-    """
-
-    channels = list(X.coords["variable"].values)
-
-    S = xr.DataArray(
-        np.cov(X, aweights=aweights),
-        dims=["channel_1", "channel_2"],
-        coords={"channel_1": channels, "channel_2": channels},
-    )
-    return S
-
-
-def initialize_xrda_1d(
-    channels: list,
-    dtype=Optional[type],
-    value: Optional[Union[complex, float, bool]] = 0,
-) -> xr.DataArray:
-    """
-
-    Returns a 1D xr.DataArray with variable "channel", having values channels named by the input list.
-
-    Parameters
-    ----------
-    channels: list
-        The channels in the multivariate array
-    dtype: type
-        The datatype to initialize the array.
-        Common cases are complex, float, and bool
-    value: Union[complex, float, bool]
-        The default value to assign the array
-
-    Returns
-    -------
-    xrda: xarray.core.dataarray.DataArray
-        An xarray container for the channels, initialized to zeros.
-    """
-    k = len(channels)
-    logger.debug(f"Initializing xarray with values {value}")
-    xrda = xr.DataArray(
-        np.zeros(k, dtype=dtype),
-        dims=[
-            "variable",
-        ],
-        coords={
-            "variable": channels,
-        },
-    )
-    if value != 0:
-        data = value * np.ones(k, dtype=dtype)
-        xrda.data = data
-    return xrda
-
-
-def initialize_xrda_2d(
-    channels, dtype=complex, value: Optional[Union[complex, float, bool]] = 0, dims=None
-):
-
-    """
-     TODO: consider merging with initialize_xrda_1d
-     TODO: consider changing nomenclature from dims=["channel_1", "channel_2"],
-     to dims=["variable_1", "variable_2"], to be consistent with initialize_xrda_1d
-
-    Parameters
-     ----------
-     channels: list
-         The channels in the multivariate array
-     dtype: type
-         The datatype to initialize the array.
-         Common cases are complex, float, and bool
-     value: Union[complex, float, bool]
-         The default value to assign the array
-
-    Returns
-     -------
-     xrda: xarray.core.dataarray.DataArray
-         An xarray container for the channel variances etc., initialized to zeros.
-    """
-    if dims is None:
-        dims = [channels, channels]
-
-    K = len(channels)
-    logger.debug(f"Initializing 2D xarray to {value}")
-    xrda = xr.DataArray(
-        np.zeros((K, K), dtype=dtype),
-        dims=["channel_1", "channel_2"],
-        coords={"channel_1": dims[0], "channel_2": dims[1]},
-    )
-    if value != 0:
-        data = value * np.ones(xrda.shape, dtype=dtype)
-        xrda.data = data
-
-    return xrda

aurora/transfer_function/base.py

@@ -50,19 +50,15 @@ def __init__(
         """
         Constructor.
 
-        Development Notes:
-        change 2021-07-23 to require a frequency_bands object.  We may want
-        to just pass the band_edges.
-
         Parameters
         ----------
         _emtf_header : legacy header information used by Egbert's matlab class.  Header contains
         local site header, remote site header if appropriate, and information about estimation approach
         decimation_level_id: int
             Identifies the relevant decimation level.  Used for accessing the
             appropriate info in self.processing config.
-        frequency_bands: aurora.time_series.frequency_band.FrequencyBands
-            frequency bands object
+        frequency_bands: FrequencyBands
+            frequency bands object defining the tf estimation bands.
         """
         self._emtf_tf_header = None
         self.decimation_level_id = decimation_level_id

docs/tutorials/synthetic_data_processing.ipynb

@@ -1908,9 +1908,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from aurora.pipelines.fourier_coefficients import add_fcs_to_mth5\n",
-    "from aurora.pipelines.fourier_coefficients import fc_decimations_creator\n",
-    "from aurora.pipelines.fourier_coefficients import read_back_fcs"
+    "from mth5.timeseries.spectre.helpers import add_fcs_to_mth5\n",
+    "from mth5.timeseries.spectre.helpers import fc_decimations_creator\n",
+    "from mth5.timeseries.spectre.helpers import read_back_fcs"
    ]
   },
   {

tests/config/test_config_creator.py

@@ -1,4 +1,5 @@
 # import logging
+import pandas as pd
 import unittest
 
 from aurora.config.config_creator import ConfigCreator
@@ -68,8 +69,13 @@ def test_frequency_bands(self):
         delta_f = dec_level_0.frequency_sample_interval
         lower_edges = (dec_level_0.lower_bounds * delta_f) - delta_f / 2.0
         upper_edges = (dec_level_0.upper_bounds * delta_f) + delta_f / 2.0
-        band_edges_b = np.vstack((lower_edges, upper_edges)).T
-        assert (band_edges_b - band_edges_a == 0).all()
+        band_edges_b = pd.DataFrame(
+            data={
+                "lower_bound": lower_edges,
+                "upper_bound": upper_edges,
+            }
+        )
+        assert (band_edges_b - band_edges_a == 0).all().all()
 
 
 def main():

tests/synthetic/test_fourier_coefficients.py

@@ -1,26 +1,21 @@
 import unittest
 
 from aurora.config.config_creator import ConfigCreator
-from aurora.pipelines.fourier_coefficients import add_fcs_to_mth5
-from aurora.pipelines.fourier_coefficients import fc_decimations_creator
-
-# from aurora.pipelines.fourier_coefficients import read_back_fcs
 from aurora.pipelines.process_mth5 import process_mth5
 from aurora.test_utils.synthetic.make_processing_configs import (
     create_test_run_config,
 )
 from aurora.test_utils.synthetic.paths import SyntheticTestPaths
+from loguru import logger
 from mth5.data.make_mth5_from_asc import create_test1_h5
 from mth5.data.make_mth5_from_asc import create_test2_h5
 from mth5.data.make_mth5_from_asc import create_test3_h5
 from mth5.data.make_mth5_from_asc import create_test12rr_h5
-from mth5.timeseries.spectre.helpers import read_back_fcs
-
-# from mtpy-v2
-from mtpy.processing import RunSummary, KernelDataset
-
-from loguru import logger
 from mth5.helpers import close_open_files
+from mth5.timeseries.spectre.helpers import add_fcs_to_mth5
+from mth5.timeseries.spectre.helpers import fc_decimations_creator
+from mth5.timeseries.spectre.helpers import read_back_fcs
+from mtpy.processing import RunSummary, KernelDataset  # from mtpy-v2
 
 synthetic_test_paths = SyntheticTestPaths()
 synthetic_test_paths.mkdirs()

tests/synthetic/test_stft_methods_agree.py

@@ -7,18 +7,15 @@
 
 from aurora.pipelines.time_series_helpers import prototype_decimate
 from aurora.pipelines.time_series_helpers import run_ts_to_stft
-from aurora.pipelines.time_series_helpers import run_ts_to_stft_scipy
 from aurora.test_utils.synthetic.make_processing_configs import (
     create_test_run_config,
 )
-
-# from mtpy-v2
-from mtpy.processing import RunSummary, KernelDataset
-
 from loguru import logger
 from mth5.data.make_mth5_from_asc import create_test1_h5
 from mth5.mth5 import MTH5
 from mth5.helpers import close_open_files
+from mth5.timeseries.spectre.stft import run_ts_to_stft_scipy
+from mtpy.processing import RunSummary, KernelDataset  # from mtpy-v2
 
 
 def test_stft_methods_agree():

tests/time_series/test_xarray_helpers.py

@@ -4,74 +4,83 @@
 """
 
 import numpy as np
-import unittest
-
 import xarray as xr
+import pytest
+
+from aurora.time_series.xarray_helpers import handle_nan
+
+
+def test_handle_nan_basic():
+    """Test basic functionality of handle_nan with NaN values."""
+    # Create sample data with NaN values
+    times = np.array([0, 1, 2, 3, 4])
+    data_x = np.array([1.0, np.nan, 3.0, 4.0, 5.0])
+    data_y = np.array([1.0, 2.0, np.nan, 4.0, 5.0])
+
+    X = xr.Dataset({"hx": ("time", data_x)}, coords={"time": times})
+    Y = xr.Dataset({"ex": ("time", data_y)}, coords={"time": times})
+
+    # Test with X and Y only
+    X_clean, Y_clean, _ = handle_nan(X, Y, None, drop_dim="time")
+
+    # Check that NaN values were dropped
+    assert len(X_clean.time) == 3
+    assert len(Y_clean.time) == 3
+    assert not np.any(np.isnan(X_clean.hx.values))
+    assert not np.any(np.isnan(Y_clean.ex.values))
+
+
+def test_handle_nan_with_remote_reference():
+    """Test handle_nan with remote reference data."""
+    # Create sample data
+    times = np.array([0, 1, 2, 3])
+    data_x = np.array([1.0, np.nan, 3.0, 4.0])
+    data_y = np.array([1.0, 2.0, 3.0, 4.0])
+    data_rr = np.array([1.0, 2.0, np.nan, 4.0])
+
+    X = xr.Dataset({"hx": ("time", data_x)}, coords={"time": times})
+    Y = xr.Dataset({"ex": ("time", data_y)}, coords={"time": times})
+    RR = xr.Dataset({"hx": ("time", data_rr)}, coords={"time": times})
+
+    # Test with all datasets
+    X_clean, Y_clean, RR_clean = handle_nan(X, Y, RR, drop_dim="time")
+
+    # Check that NaN values were dropped
+    assert len(X_clean.time) == 2
+    assert len(Y_clean.time) == 2
+    assert len(RR_clean.time) == 2
+    assert not np.any(np.isnan(X_clean.hx.values))
+    assert not np.any(np.isnan(Y_clean.ex.values))
+    assert not np.any(np.isnan(RR_clean.hx.values))
+
+    # Check that the values are correct
+    expected_times = np.array([0, 3])
+    assert np.allclose(X_clean.time.values, expected_times)
+    assert np.allclose(Y_clean.time.values, expected_times)
+    assert np.allclose(RR_clean.time.values, expected_times)
+    assert np.allclose(X_clean.hx.values, np.array([1.0, 4.0]))
+    assert np.allclose(Y_clean.ex.values, np.array([1.0, 4.0]))
+    assert np.allclose(RR_clean.hx.values, np.array([1.0, 4.0]))
+
+
+def test_handle_nan_time_mismatch():
+    """Test handle_nan with time coordinate mismatches."""
+    # Create sample data with slightly different timestamps
+    times_x = np.array([0, 1, 2, 3])
+    times_rr = times_x + 0.1  # Small offset
+    data_x = np.array([1.0, 2.0, 3.0, 4.0])
+    data_rr = np.array([1.0, 2.0, 3.0, 4.0])
+
+    X = xr.Dataset({"hx": ("time", data_x)}, coords={"time": times_x})
+    RR = xr.Dataset({"hx": ("time", data_rr)}, coords={"time": times_rr})
+
+    # Test handling of time mismatch
+    X_clean, _, RR_clean = handle_nan(X, None, RR, drop_dim="time")
+
+    # Check that data was preserved despite time mismatch
+    assert len(X_clean.time) == 4
+    assert "hx" in RR_clean.data_vars
+    assert np.allclose(RR_clean.hx.values, data_rr)
 
-from aurora.time_series.xarray_helpers import covariance_xr
-from aurora.time_series.xarray_helpers import initialize_xrda_1d
-from aurora.time_series.xarray_helpers import initialize_xrda_2d
-
-
-class TestXarrayHelpers(unittest.TestCase):
-    """
-    Test methods in xarray helpers
-    - may get broken into separate tests if this module grows
-    """
-
-    @classmethod
-    def setUpClass(self):
-        self.standard_channel_names = ["ex", "ey", "hx", "hy", "hz"]
-
-    def setUp(self):
-        pass
-
-    def test_initialize_xrda_1d(self):
-        dtype = float
-        value = -1
-        tmp = initialize_xrda_1d(self.standard_channel_names, dtype=dtype, value=value)
-        self.assertTrue((tmp.data == value).all())
-
-    def test_initialize_xrda_2d(self):
-        dtype = float
-        value = -1
-        tmp = initialize_xrda_2d(self.standard_channel_names, dtype=dtype, value=value)
-        self.assertTrue((tmp.data == value).all())
-
-    def test_covariance_xr(self):
-        np.random.seed(0)
-        n_observations = 100
-        xrds = xr.Dataset(
-            {
-                "hx": (
-                    [
-                        "time",
-                    ],
-                    np.abs(np.random.randn(n_observations)),
-                ),
-                "hy": (
-                    [
-                        "time",
-                    ],
-                    np.abs(np.random.randn(n_observations)),
-                ),
-            },
-            coords={
-                "time": np.arange(n_observations),
-            },
-        )
-
-        X = xrds.to_array()
-        cov = covariance_xr(X)
-        self.assertTrue((cov.data == cov.data.transpose().conj()).all())
-
-    def test_sometehing_else(self):
-        """
-        Place holder
-
-        """
-        pass
-
-
-if __name__ == "__main__":
-    unittest.main()
+    # Check that the time values match X's time values
+    assert np.allclose(RR_clean.time.values, X_clean.time.values)

tests/transfer_function/test_cross_power.py

@@ -1,4 +1,4 @@
-from aurora.time_series.xarray_helpers import initialize_xrda_2d
+from mth5.timeseries.xarray_helpers import initialize_xrda_2d_cov
 from aurora.transfer_function.cross_power import tf_from_cross_powers
 from aurora.transfer_function.cross_power import _channel_names
 from aurora.transfer_function.cross_power import (
@@ -32,7 +32,7 @@ def setUpClass(self):
         station_1_channels = [f"{self.station_ids[0]}_{x}" for x in components]
         station_2_channels = [f"{self.station_ids[1]}_{x}" for x in components]
         channels = station_1_channels + station_2_channels
-        sdm = initialize_xrda_2d(
+        sdm = initialize_xrda_2d_cov(
             channels=channels,
             dtype=complex,
         )