add covariance_xr method and a test

aurora/time_series/xarray_helpers.py

@@ -89,6 +89,39 @@ 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],

tests/time_series/test_xarray_helpers.py

@@ -1,8 +1,14 @@
 # -*- coding: utf-8 -*-
 """
+This module contains unittests for the xarray_helpers module.
 """
 
+import numpy as np
 import unittest
+
+import xarray as xr
+
+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
 
@@ -15,25 +21,50 @@ class TestXarrayHelpers(unittest.TestCase):
 
     @classmethod
     def setUpClass(self):
-        pass
+        self.standard_channel_names = ["ex", "ey", "hx", "hy", "hz"]
 
     def setUp(self):
         pass
 
     def test_initialize_xrda_1d(self):
-        channels = ["ex", "ey", "hx", "hy", "hz"]
         dtype = float
         value = -1
-        tmp = initialize_xrda_1d(channels, dtype=dtype, value=value)
+        tmp = initialize_xrda_1d(self.standard_channel_names, dtype=dtype, value=value)
         self.assertTrue((tmp.data == value).all())
 
     def test_initialize_xrda_2d(self):
-        channels = ["ex", "ey", "hx", "hy", "hz"]
         dtype = float
         value = -1
-        tmp = initialize_xrda_2d(channels, dtype=dtype, value=value)
+        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