cov() and corr() - finalization

doc/api.rst

@@ -30,6 +30,8 @@ Top-level functions
    zeros_like
    ones_like
    dot
+   cov
+   corr
    map_blocks
 
 Dataset

doc/whats-new.rst

@@ -25,6 +25,8 @@ Breaking changes
 
 New Features
 ~~~~~~~~~~~~
+- Added - :py:func:`xarray.cov` and :py:func:`xarray.corr` (:pull:`3550`).
+  By `Robin Beer <https://github.com/r-beer>`_.
 
 
 Bug fixes

xarray/__init__.py

@@ -10,7 +10,7 @@
 from .core.common import full_like, zeros_like, ones_like
 from .core.concat import concat
 from .core.combine import combine_by_coords, combine_nested, auto_combine
-from .core.computation import apply_ufunc, dot, where
+from .core.computation import apply_ufunc, dot, where, cov, corr
 from .core.extensions import register_dataarray_accessor, register_dataset_accessor
 from .core.variable import as_variable, Variable, IndexVariable, Coordinate
 from .core.dataset import Dataset

xarray/core/computation.py

@@ -22,9 +22,9 @@
 )
 
 import numpy as np
-
+import xarray as xr
 from . import duck_array_ops, utils
-from .alignment import deep_align
+from .alignment import broadcast, deep_align
 from .merge import merge_coordinates_without_align
 from .pycompat import dask_array_type
 from .utils import is_dict_like
@@ -1047,6 +1047,175 @@ def earth_mover_distance(first_samples,
         return apply_array_ufunc(func, *args, dask=dask)
 
 
+def cov(da_a, da_b, dim=None, ddof=1):
+    """Compute covariance between two DataArray objects along a shared dimension.
+
+    Parameters
+    ----------
+    da_a: DataArray (or Variable) object
+        Array to compute.
+    da_b: DataArray (or Variable) object
+        Array to compute.
+    dim : str, optional
+        The dimension along which the covariance will be computed
+
+    Returns
+    -------
+    covariance: DataArray
+
+    See also
+    --------
+    pandas.Series.cov: corresponding pandas function
+    xr.corr: respective function to calculate correlation
+
+    Examples
+    --------
+
+    >>> da_a = DataArray(np.random.random((3, 5)),
+    ...                  dims=("space", "time"),
+    ...                  coords=[('space', ['IA', 'IL', 'IN']),
+    ...                          ('time', pd.date_range("2000-01-01", freq="1D", periods=5))])
+    >>> da_a
+    <xarray.DataArray (space: 3, time: 5)>
+    array([[0.04356841, 0.11479286, 0.70359101, 0.59072561, 0.16601438],
+            [0.81552383, 0.72304926, 0.77644406, 0.05788198, 0.74065536],
+            [0.96252519, 0.36877741, 0.22248412, 0.55185954, 0.23547536]])
+    Coordinates:
+    * space    (space) <U2 'IA' 'IL' 'IN'
+    * time     (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-05
+
+    >>> da_b = DataArray(np.random.random((3, 5)),
+    ...                  dims=("space", "time"),
+    ...                  coords=[('space', ['IA', 'IL', 'IN']),
+    ...                          ('time', pd.date_range("2000-01-01", freq="1D", periods=5))])
+    >>> da_b
+    <xarray.DataArray (space: 3, time: 5)>
+    array([[0.41505599, 0.43002193, 0.45250454, 0.57701084, 0.5327754 ],
+            [0.0998048 , 0.67225522, 0.4234324 , 0.13514615, 0.4399088 ],
+            [0.24675048, 0.58555283, 0.1942955 , 0.86128908, 0.05068975]])
+    Coordinates:
+    * space    (space) <U2 'IA' 'IL' 'IN'
+    * time     (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-05
+    >>> xr.cov(da_a, da_b)
+    <xarray.DataArray ()>
+    array(0.03823054)
+    >>> xr.cov(da_a, da_b, dim='time')
+    <xarray.DataArray (space: 3)>
+    array([0.00207952, 0.01024296, 0.08214707])
+    Coordinates:
+    * space    (space) <U2 'IA' 'IL' 'IN'
+    """
+
+    # 1. Broadcast the two arrays
+    da_a, da_b = broadcast(da_a, da_b)
+
+    # 2. Ignore the nans
+    valid_values = da_a.notnull() & da_b.notnull()
+    # TODO: avoid drop
+    da_a = da_a.where(valid_values, drop=True)
+    da_b = da_b.where(valid_values, drop=True)
+    valid_count = valid_values.sum(dim) - ddof
+
+    # if dim is not None:
+    #     valid_count = da_a[dim].size - ddof
+    # else:
+    #     valid_count = da_a.size
+
+    # 3. Compute mean and standard deviation along the given dim
+    demeaned_da_a = da_a - da_a.mean(dim=dim)
+    demeaned_da_b = da_b - da_b.mean(dim=dim)
+
+    # 4. Compute covariance along the given dim
+    cov = (demeaned_da_a * demeaned_da_b).sum(dim=dim) / (valid_count)
+
+    return xr.DataArray(cov)
+
+
+def corr(da_a, da_b, dim=None, ddof=0):
+    """Compute the Pearson correlation coefficient between two DataArray objects along a shared dimension.
+
+    Parameters
+    ----------
+    da_a: DataArray (or Variable) object
+        Array to compute.
+    da_b: DataArray (or Variable) object
+        Array to compute.
+    dim: str, optional
+        The dimension along which the correlation will be computed
+
+    Returns
+    -------
+    correlation: DataArray
+
+    See also
+    --------
+    pandas.Series.corr: corresponding pandas function
+    xr.cov: underlying covariance function
+
+    Examples
+    --------
+
+    >>> da_a = DataArray(np.random.random((3, 5)),
+    ...                  dims=("space", "time"),
+    ...                  coords=[('space', ['IA', 'IL', 'IN']),
+    ...                          ('time', pd.date_range("2000-01-01", freq="1D", periods=5))])
+    >>> da_a
+    <xarray.DataArray (space: 3, time: 5)>
+    array([[0.04356841, 0.11479286, 0.70359101, 0.59072561, 0.16601438],
+            [0.81552383, 0.72304926, 0.77644406, 0.05788198, 0.74065536],
+            [0.96252519, 0.36877741, 0.22248412, 0.55185954, 0.23547536]])
+    Coordinates:
+    * space    (space) <U2 'IA' 'IL' 'IN'
+    * time     (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-05
+
+    >>> da_b = DataArray(np.random.random((3, 5)),
+    ...                  dims=("space", "time"),
+    ...                  coords=[('space', ['IA', 'IL', 'IN']),
+    ...                          ('time', pd.date_range("2000-01-01", freq="1D", periods=5))])
+    >>> da_b
+    <xarray.DataArray (space: 3, time: 5)>
+    array([[0.41505599, 0.43002193, 0.45250454, 0.57701084, 0.5327754 ],
+            [0.0998048 , 0.67225522, 0.4234324 , 0.13514615, 0.4399088 ],
+            [0.24675048, 0.58555283, 0.1942955 , 0.86128908, 0.05068975]])
+    Coordinates:
+    * space    (space) <U2 'IA' 'IL' 'IN'
+    * time     (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-05
+    >>> xr.corr(da_a, da_b)
+    <xarray.DataArray ()>
+    array(0.67407116)
+    >>> xr.corr(da_a, da_b, dim='time')
+    <xarray.DataArray (space: 3)>
+    array([0.23150267, 0.24900968, 0.9061562 ])
+    Coordinates:
+    * space    (space) <U2 'IA' 'IL' 'IN'
+    """
+    from .dataarray import DataArray
+
+    if any(not isinstance(arr, (Variable, DataArray)) for arr in [da_a, da_b]):
+        raise TypeError(
+            "Only xr.DataArray and xr.Variable are supported."
+            "Given {}.".format([type(arr) for arr in [da_a, da_b]])
+        )
+
+    # 1. Broadcast the two arrays
+    da_a, da_b = broadcast(da_a, da_b)
+
+    # 2. Ignore the nans
+    valid_values = da_a.notnull() & da_b.notnull()
+    da_a = da_a.where(valid_values, drop=True)
+    da_b = da_b.where(
+        valid_values, drop=True
+    )  # TODO: avoid drop as explained in https://github.com/pydata/xarray/pull/2652#discussion_r245492002
+
+    # 3. Compute correlation based on standard deviations and cov()
+    da_a_std = da_a.std(dim=dim)
+    da_b_std = da_b.std(dim=dim)
+
+    corr = cov(da_a, da_b, dim=dim, ddof=ddof) / (da_a_std * da_b_std)
+
+    return xr.DataArray(corr)
+
+
 def dot(*arrays, dims=None, **kwargs):
     """Generalized dot product for xarray objects. Like np.einsum, but
     provides a simpler interface based on array dimensions.

xarray/core/dataarray.py

@@ -2834,7 +2834,6 @@ def dot(
         >>> da = DataArray(da_vals, dims=['x', 'y', 'z'])
         >>> dm_vals = np.arange(4)
         >>> dm = DataArray(dm_vals, dims=['z'])
-
         >>> dm.dims
         ('z')
         >>> da.dims

xarray/tests/test_backends.py

@@ -2439,6 +2439,7 @@ def skip_if_not_engine(engine):
         pytest.importorskip(engine)
 
 
+@requires_scipy_or_netCDF4
 @requires_dask
 @pytest.mark.filterwarnings("ignore:use make_scale(name) instead")
 def test_open_mfdataset_manyfiles(

xarray/tests/test_computation.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 import pytest
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_array_equal, assert_allclose
 
 import xarray as xr
 from xarray.core.computation import (
@@ -789,6 +789,68 @@ def func(x):
     assert_identical(expected, actual)
 
 
+da = xr.DataArray(np.random.random((3, 21, 4)),
+                  coords={"time": pd.date_range("2000-01-01", freq="1D", periods=21)},
+                  dims=("a", "time", "x"),)
+
+arrays = [
+    da.isel(time=range(0, 18)),
+    da.isel(time=range(2, 20)).rolling(time=3, center=True).mean(dim="time"),
+    xr.DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time"),
+    xr.DataArray([1, 1, np.nan, 2, np.nan, 3, 5, 4, 6, np.nan, 7], dims="time"),
+    xr.DataArray([[1, 2], [1, np.nan]], dims=["x", "time"]),
+    xr.DataArray([[1, 2], [np.nan, np.nan]], dims=["x", "time"]),
+]
+
+array_tuples = [
+    (arrays[0], arrays[0]),
+    (arrays[0], arrays[1]),
+    (arrays[1], arrays[1]),
+    (arrays[2], arrays[2]),
+    (arrays[2], arrays[3]),
+    (arrays[3], arrays[3]),
+    (arrays[4], arrays[4]),
+    (arrays[4], arrays[5]),
+    (arrays[5], arrays[5]),
+]
+
+@pytest.mark.parametrize("da_a, da_b", array_tuples)
+@pytest.mark.parametrize("dim", [None, "time", "x"])
+def test_cov(da_a, da_b, dim):
+    def pandas_cov(ts1, ts2):
+        """Ensure the ts are aligned and missing values ignored"""
+        ts1, ts2 = xr.align(ts1, ts2)
+        valid_values = ts1.notnull() & ts2.notnull()
+
+        ts1 = ts1.where(valid_values, drop=True)
+        ts2 = ts2.where(valid_values, drop=True)
+
+        return ts1.to_series().cov(ts2.to_series())
+
+    expected = pandas_cov(da_a, da_b)
+    actual = xr.cov(da_a, da_b, dim)
+
+    assert_allclose(actual, expected)
+
+
+@pytest.mark.parametrize("da_a, da_b", array_tuples)
+@pytest.mark.parametrize("dim", [None, "time", "x"])
+def test_corr(da_a, da_b, dim):
+    def pandas_corr(ts1, ts2):
+        """Ensure the ts are aligned and missing values ignored"""
+        ts1, ts2 = xr.align(ts1, ts2)
+        valid_values = ts1.notnull() & ts2.notnull()
+
+        ts1 = ts1.where(valid_values, drop=True)
+        ts2 = ts2.where(valid_values, drop=True)
+
+        return ts1.to_series().corr(ts2.to_series())
+
+    expected = pandas_corr(da_a, da_b)
+    actual = xr.corr(da_a, da_b, dim)
+    assert_allclose(actual, expected)
+
+
 def pandas_median(x):
     return pd.Series(x).median()