add calc_gaspari_cohn_correlation_matrices (MESMER-group#300)

* add calc_gaspari_cohn_correlation_matrices

* CHANGELOG

* add see also

* linting

CHANGELOG.rst

@@ -78,6 +78,9 @@ New Features
   By `Mathias Hauser`_.
 - Allow passing `xr.DataArray` to ``calc_geodist_exact`` (`#299 <https://github.com/MESMER-group/mesmer/pull/299>`__).
   By `Zeb Nicholls`_ and `Mathias Hauser`_.
+- Add ``calc_gaspari_cohn_correlation_matrices`` a function to calculate Gaspari-Cohn correlation
+  matrices for a range of localisation radii (`#300 <https://github.com/MESMER-group/mesmer/pull/300>`__).
+  By `Zeb Nicholls`_ and `Mathias Hauser`_.
 
 
 Breaking changes

mesmer/core/computation.py

@@ -10,6 +10,37 @@
 from .utils import create_equal_dim_names
 
 
+def calc_gaspari_cohn_correlation_matrices(geodist, localisation_radii):
+    """Gaspari-Cohn correlation matrices for a range of localisation radii
+
+    Parameters
+    ----------
+    geodist : xr.DataArray, np.ndarray
+        2D array of great circle distances. Calculated from e.g. ``calc_geodist_exact``.
+    localisation_radii : iterable of float
+        Localisation radii to test (in km)
+
+    Returns
+    -------
+    gaspari_cohn_correlation_matrices: dict[float : :obj:`xr.DataArray`, :obj:`np.ndarray`]
+        Gaspari-Cohn correlation matrix (values) for each localisation radius (keys)
+
+    Notes
+    -----
+    Values in ``localisation_radii`` should not exceed 10'000 km by much because
+    it can lead to correlation matrices which are not positive semidefinite.
+
+    See Also
+    --------
+    gaspari_cohn, calc_geodist_exact
+
+    """
+
+    out = {lr: gaspari_cohn(geodist / lr) for lr in localisation_radii}
+
+    return out
+
+
 def gaspari_cohn(r):
     """smooth, exponentially decaying Gaspari-Cohn correlation function
 

tests/unit/test_computation.py

@@ -2,7 +2,12 @@
 import pytest
 import xarray as xr
 
-from mesmer.core.computation import calc_geodist_exact, gaspari_cohn
+from mesmer.core.computation import (
+    calc_gaspari_cohn_correlation_matrices,
+    calc_geodist_exact,
+    gaspari_cohn,
+)
+from mesmer.testing import assert_dict_allclose
 
 
 def test_gaspari_cohn_error():
@@ -147,3 +152,25 @@ def test_calc_geodist_exact(as_dataarray):
     else:
 
         np.testing.assert_allclose(result, expected)
+
+
+@pytest.mark.parametrize("localisation_radii", [[1000, 2000], range(5000, 5001)])
+@pytest.mark.parametrize("as_dataarray", [True, False])
+def test_gaspari_cohn_correlation_matrices(localisation_radii, as_dataarray):
+
+    lon = np.arange(0, 31, 5)
+    lat = np.arange(-45, 46, 30)
+    lat, lon = np.meshgrid(lat, lon)
+    lon, lat = lon.flatten(), lat.flatten()
+
+    if as_dataarray:
+        lon = xr.DataArray(lon, dims="gridpoint")
+        lat = xr.DataArray(lat, dims="gridpoint")
+
+    geodist = calc_geodist_exact(lon, lat)
+
+    result = calc_gaspari_cohn_correlation_matrices(geodist, localisation_radii)
+
+    expected = {lr: gaspari_cohn(geodist / lr) for lr in localisation_radii}
+
+    assert_dict_allclose(expected, result)