From 5c55c42cc715135edf1ced40a89779c5fce1f2a7 Mon Sep 17 00:00:00 2001
From: Wenrui Jiang <wjiang33@jhu.edu>
Date: Wed, 9 Oct 2024 17:11:39 -0400
Subject: [PATCH] test for the actual interpolation

---
 seaduck/lagrangian_budget.py    |  2 +-
 tests/test_lagrangian_budget.py | 22 ++++++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/seaduck/lagrangian_budget.py b/seaduck/lagrangian_budget.py
index 79970f7..fa6cc52 100644
--- a/seaduck/lagrangian_budget.py
+++ b/seaduck/lagrangian_budget.py
@@ -404,7 +404,7 @@ def prefetch_scalar(ds_slc, scalar_names):
 
 
 def read_wall_list(neo, tp, prefetch=None, scalar=True):
-    if "face" not in neo.data_vars and "fc" not in neo.data_vars: # pragma: no cover
+    if "face" not in neo.data_vars and "fc" not in neo.data_vars:  # pragma: no cover
         ind = (neo.iz - 1, neo.iy, neo.ix)
         deep_ind = (neo.iz, neo.iy, neo.ix)
         right_ind = tuple(
diff --git a/tests/test_lagrangian_budget.py b/tests/test_lagrangian_budget.py
index 8bed811..53f9ed7 100644
--- a/tests/test_lagrangian_budget.py
+++ b/tests/test_lagrangian_budget.py
@@ -7,10 +7,12 @@
 from seaduck.eulerian_budget import total_div
 from seaduck.lagrangian_budget import (
     check_particle_data_compat,
+    contr_p_relaxed,
     find_ind_frac_tres,
     first_last_neither,
     flatten,
     ind_tend_uv,
+    lhs_contribution,
     particle2xarray,
     prefetch_scalar,
     prefetch_vector,
@@ -222,3 +224,23 @@ def test_prefetch_vector(xrslc, same_size):
         xrslc, xname="sx", yname="sy", zname="sz", same_size=same_size
     )
     assert isinstance(larger, np.ndarray)
+
+
+def test_lhs_contribution():
+    lhs = np.random.random(10)
+    scalar_dic = {"lhs": lhs}
+    last = np.array([8, 10])
+    t = np.arange(10)
+    ans = lhs_contribution(t, scalar_dic, last)
+    assert np.allclose(ans[:3], lhs[:3])
+    assert ans[8] == 0
+
+
+def test_p_relax():
+    deltas = np.ones(3) * 3
+    tres = np.ones(3)
+    step_dic = {"term1": np.array([1, 2, 3, 4]), "term2": np.array([3, 2, 1, 0])}
+    termlist = ["term1", "term2"]
+    res = contr_p_relaxed(deltas, tres, step_dic, termlist)
+    assert np.allclose(0, res["error"])
+    assert res["term1"][1] == res["term2"][1]