update name of initialize_xrda_2d to initialize_xrda_2d_cov

aurora/time_series/xarray_helpers.py

@@ -7,6 +7,9 @@
 from loguru import logger
 from typing import Optional, Union
 
+from mth5.timeseries.xarray_helpers import covariance_xr
+from mth5.timeseries.xarray_helpers import initialize_xrda_1d
+
 
 def handle_nan(X, Y, RR, drop_dim=""):
     """
@@ -89,81 +92,7 @@ def handle_nan(X, Y, RR, drop_dim=""):
     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(
+def initialize_xrda_2d_cov(
     channels, dtype=complex, value: Optional[Union[complex, float, bool]] = 0, dims=None
 ):
 

tests/time_series/test_xarray_helpers.py

@@ -10,7 +10,7 @@
 
 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
+from aurora.time_series.xarray_helpers import initialize_xrda_2d_cov
 
 
 class TestXarrayHelpers(unittest.TestCase):
@@ -32,10 +32,12 @@ def test_initialize_xrda_1d(self):
         tmp = initialize_xrda_1d(self.standard_channel_names, dtype=dtype, value=value)
         self.assertTrue((tmp.data == value).all())
 
-    def test_initialize_xrda_2d(self):
+    def test_initialize_xrda_2d_cov(self):
         dtype = float
         value = -1
-        tmp = initialize_xrda_2d(self.standard_channel_names, dtype=dtype, value=value)
+        tmp = initialize_xrda_2d_cov(
+            self.standard_channel_names, dtype=dtype, value=value
+        )
         self.assertTrue((tmp.data == value).all())
 
     def test_covariance_xr(self):

tests/transfer_function/test_cross_power.py

@@ -1,4 +1,4 @@
-from aurora.time_series.xarray_helpers import initialize_xrda_2d
+from aurora.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,
         )