From 7a37063fffd967a5ad9bc23443042b11eb6b02e5 Mon Sep 17 00:00:00 2001
From: corentincarton <corentin.carton@ecmwf.int>
Date: Thu, 19 Dec 2024 18:45:02 +0100
Subject: [PATCH] Add implementation of lmoment since scipy v0.15 is not yet
 released

The plan is to exclusively rely on scipy for the functionality in the future,
but for now to provide a substitute that doesn't rely on a version of scipy
that is still under development (as of Dec 2024).

Co-authored-by: Christopher Polster <christopher.polster@ecmwf.int>
---
 .../meteo/stats/array/distributions.py        | 58 ++++++++++++++++++-
 1 file changed, 57 insertions(+), 1 deletion(-)

diff --git a/src/earthkit/meteo/stats/array/distributions.py b/src/earthkit/meteo/stats/array/distributions.py
index 433aae6..fe82792 100644
--- a/src/earthkit/meteo/stats/array/distributions.py
+++ b/src/earthkit/meteo/stats/array/distributions.py
@@ -34,6 +34,59 @@ def ppf(self, x):
         """Evaluate the inverse CDF (percent point function)"""
 
 
+# Temporary drop-in replacement for scipy.stats.lmoment from scipy v0.15
+def _lmoment(sample, order=(1, 2), axis=0):
+    """Compute first 2 L-moments of dataset along first axis"""
+    if len(order) != 2 or order[0] != 1 or order[1] != 2:
+        raise NotImplementedError
+    if axis != 0:
+        raise NotImplementedError
+
+    nmoments = 3
+
+    # At least four values needed to make a sample L-moments estimation
+    nvalues, *rest_shape = sample.shape
+    if nvalues < 4:
+        raise ValueError("Insufficient number of values to perform sample L-moments estimation")
+
+    sample = np.sort(sample, axis=0)  # ascending order
+
+    sums = np.zeros_like(sample, shape=(nmoments, *rest_shape))
+
+    for i in range(1, nvalues + 1):
+        z = i
+        term = sample[i - 1]
+        sums[0] = sums[0] + term
+        for j in range(1, nmoments):
+            z -= 1
+            term = term * z
+            sums[j] = sums[j] + term
+
+    y = float(nvalues)
+    z = float(nvalues)
+    sums[0] = sums[0] / z
+    for j in range(1, nmoments):
+        y = y - 1.0
+        z = z * y
+        sums[j] = sums[j] / z
+
+    k = nmoments
+    p0 = -1.0
+    for _ in range(2):
+        ak = float(k)
+        p0 = -p0
+        p = p0
+        temp = p * sums[0]
+        for i in range(1, k):
+            ai = i
+            p = -p * (ak + ai - 1.0) * (ak - ai) / (ai * ai)
+            temp = temp + (p * sums[i])
+        sums[k - 1] = temp
+        k = k - 1
+
+    return np.stack([sums[0], sums[1]])
+
+
 class MaxGumbel(ContinuousDistribution):
     """Gumbel distribution for extreme values
 
@@ -61,7 +114,10 @@ def fit(cls, sample, axis=0):
         axis: int
             The axis along which to compute the parameters.
         """
-        from scipy.stats import lmoment
+        try:
+            from scipy.stats import lmoment
+        except ImportError:
+            lmoment = _lmoment
 
         lmom = lmoment(sample, axis=axis, order=[1, 2])
         sigma = lmom[1] / np.log(2)