Clip times in interpolate_states

sot · Jan 9, 2025 · d3e76ec · d3e76ec
1 parent 510ea8c
commit d3e76ec
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/kadi/commands/states.py b/kadi/commands/states.py
@@ -2532,6 +2532,9 @@ def get_continuity(
 def interpolate_states(states, times):
     """Interpolate ``states`` table at given times.
 
+    Any ``times`` that are outside the range of the states table are clipped to the
+    first or last state.
+
     Parameters
     ----------
     states
@@ -2554,7 +2557,7 @@ def interpolate_states(states, times):
         tstops = date2secs(states["datestop"])
 
     indexes = np.searchsorted(tstops, times)
-    out = states[indexes]
+    out = states[indexes.clip(0, len(states) - 1)]
     out.add_column(Column(secs2date(times), name="date"), index=0)
 
     return out

diff --git a/kadi/commands/tests/test_states.py b/kadi/commands/tests/test_states.py
@@ -1897,3 +1897,12 @@ def test_get_continuity_spm_eclipse():
             "eclipse_enable_spm": "2017:087:07:49:55.838",
         },
     }
+
+
+def test_interpolate_states_extrapolate():
+    """Test that the states are correctly interpolated and extrapolated."""
+    # fmt: off
+    sts =  states.get_states("2024:001:00:01:00", "2024:001:00:05:00", state_keys=["obsid"])
+    times = ["2024:001:00:00:00", "2024:001:00:02:30", "2024:001:00:05:00", "2024:001:00:06:00"]
+    sts_interp = states.interpolate_states(sts, times)
+    assert np.all(sts_interp["obsid"] == [43839] * 4)