From f5c9d1e627fb41c2623781ba695ffecfd1fd1a68 Mon Sep 17 00:00:00 2001
From: Mathias Hauser <mathause@users.noreply.github.com>
Date: Thu, 21 Sep 2023 23:17:10 +0200
Subject: [PATCH] add calc_gaspari_cohn_correlation_matrices (#300)

* add calc_gaspari_cohn_correlation_matrices

* CHANGELOG

* add see also

* linting
---
 CHANGELOG.rst                  |  3 +++
 mesmer/core/computation.py     | 31 +++++++++++++++++++++++++++++++
 tests/unit/test_computation.py | 29 ++++++++++++++++++++++++++++-
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 9c850aca..138fd547 100644
--- a/CHANGELOG.rst
+++ b/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
diff --git a/mesmer/core/computation.py b/mesmer/core/computation.py
index 186fb36a..ac03e4c1 100644
--- a/mesmer/core/computation.py
+++ b/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
 
diff --git a/tests/unit/test_computation.py b/tests/unit/test_computation.py
index 020d0dfb..3f80fb8b 100644
--- a/tests/unit/test_computation.py
+++ b/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)