From e7fc372d13e3baae87d4eeafad0ca9082329c06f Mon Sep 17 00:00:00 2001
From: ValentinaHutter <valentinahutter@gmail.com>
Date: Wed, 18 Sep 2024 16:14:47 +0200
Subject: [PATCH] add array_interpolate

---
 .../process_implementations/arrays.py             | 12 ++++++++++++
 tests/test_arrays.py                              | 15 +++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/openeo_processes_dask/process_implementations/arrays.py b/openeo_processes_dask/process_implementations/arrays.py
index 14decd5..3d27f8c 100644
--- a/openeo_processes_dask/process_implementations/arrays.py
+++ b/openeo_processes_dask/process_implementations/arrays.py
@@ -33,6 +33,7 @@
     "array_find",
     "array_labels",
     "array_apply",
+    "array_interpolate_linear",
     "first",
     "last",
     "order",
@@ -234,6 +235,17 @@ def array_apply(
         )
 
 
+def array_interpolate_linear(data: ArrayLike):
+    x = np.arange(len(data))
+    valid = np.isfinite(data)
+    if len(x[valid]) < 2:
+        return data
+    data[~valid] = np.interp(
+        x[~valid], x[valid], data[valid], left=np.nan, right=np.nan
+    )
+    return data
+
+
 def first(
     data: ArrayLike,
     ignore_nodata: Optional[bool] = True,
diff --git a/tests/test_arrays.py b/tests/test_arrays.py
index 6f66c99..9af5f90 100644
--- a/tests/test_arrays.py
+++ b/tests/test_arrays.py
@@ -264,6 +264,21 @@ def test_array_apply(process_registry):
     assert (output_cube == [2, 3, 4, 5, 6, 7]).all()
 
 
+@pytest.mark.parametrize(
+    "data, expected",
+    [
+        ([np.nan, 1, np.nan, 6, np.nan, -8], [np.nan, 1, 3.5, 6, -1, -8]),
+        ([np.nan, 1, np.nan, np.nan], [np.nan, 1, np.nan, np.nan]),
+    ],
+)
+def test_array_interpolate_linear(data, expected):
+    assert np.array_equal(
+        array_interpolate_linear(data),
+        expected,
+        equal_nan=True,
+    )
+
+
 def test_first():
     assert first(np.array([1, 0, 3, 2])) == 1
     assert pd.isnull(first(np.array([np.nan, 2, 3]), ignore_nodata=False))