Coverage for src/cvx/simulator/utils/interpolation.py: 100%
56 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-31 13:14 +0000
1# Copyright 2023 Stanford University Convex Optimization Group
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Interpolation utilities for time series data.
16This module provides functions for interpolating missing values in time series
17and validating that time series don't have missing values in the middle.
18"""
20from typing import cast
22import pandas as pd
23import polars as pl
26def _check_series_type(ts: object) -> None:
27 """Raise ``TypeError`` unless ``ts`` is a pandas or polars Series."""
28 if not isinstance(ts, pd.Series | pl.Series):
29 msg = f"Expected pd.Series or pl.Series, got {type(ts)}"
30 raise TypeError(msg)
33def _pl_non_null_indices(ts: pl.Series) -> pl.Series:
34 """Return the positions of the non-null entries in ``ts``, in ascending order."""
35 return ts.is_not_null().arg_true()
38def interpolate(ts: pd.Series | pl.Series) -> pd.Series | pl.Series:
39 """Interpolate missing values in a time series between the first and last valid indices.
41 This function fills forward (ffill) missing values in a time series, but only
42 between the first and last valid indices. Values outside this range remain NaN/null.
44 Parameters
45 ----------
46 ts : pd.Series or pl.Series
47 The time series to interpolate
49 Returns:
50 -------
51 pd.Series or pl.Series
52 The interpolated time series
54 Examples:
55 --------
56 >>> import pandas as pd
57 >>> import numpy as np
58 >>> ts = pd.Series([1, np.nan, np.nan, 4, 5])
59 >>> interpolate(ts)
60 0 1.0
61 1 1.0
62 2 1.0
63 3 4.0
64 4 5.0
65 dtype: float64
67 """
68 _check_series_type(ts)
70 # If the input is a polars Series, use the polars-specific function
71 if isinstance(ts, pl.Series):
72 return interpolate_pl(ts)
73 first = ts.first_valid_index()
74 last = ts.last_valid_index()
76 if first is not None and last is not None:
77 ts_slice = ts.loc[first:last]
78 ts_slice = ts_slice.ffill()
79 result = ts.copy()
80 result.loc[first:last] = ts_slice
81 return result
82 return ts
85def valid(ts: pd.Series | pl.Series) -> bool:
86 """Check if a time series has no missing values between the first and last valid indices.
88 This function verifies that a time series doesn't have any NaN/null values in the middle.
89 It's acceptable to have NaNs/nulls at the beginning or end of the series.
91 Parameters
92 ----------
93 ts : pd.Series or pl.Series
94 The time series to check
96 Returns:
97 -------
98 bool
99 True if the time series has no missing values between the first and last valid indices,
100 False otherwise
102 Examples:
103 --------
104 >>> import pandas as pd
105 >>> import numpy as np
106 >>> ts1 = pd.Series([np.nan, 1, 2, 3, np.nan]) # NaNs only at beginning and end
107 >>> valid(ts1)
108 True
109 >>> ts2 = pd.Series([1, 2, np.nan, 4, 5]) # NaN in the middle
110 >>> valid(ts2)
111 False
113 """
114 _check_series_type(ts)
116 # If the input is a polars Series, use the polars-specific function
117 if isinstance(ts, pl.Series):
118 return valid_pl(ts)
119 # Check if the series with NaNs dropped has the same indices as the interpolated series with NaNs dropped
120 # If they're the same, there are no NaNs in the middle of the series
121 interpolated = cast(pd.Series, interpolate(ts))
122 return bool(ts.dropna().index.equals(interpolated.dropna().index))
125def interpolate_pl(ts: pl.Series) -> pl.Series:
126 """Interpolate missing values in a polars time series between the first and last valid indices.
128 This function fills forward (ffill) missing values in a time series, but only
129 between the first and last valid indices. Values outside this range remain null.
131 Parameters
132 ----------
133 ts : pl.Series
134 The time series to interpolate
136 Returns:
137 -------
138 pl.Series
139 The interpolated time series
141 Examples:
142 --------
143 >>> import polars as pl
144 >>> ts = pl.Series([1, None, None, 4, 5])
145 >>> interpolate_pl(ts)
146 shape: (5,)
147 Series: '' [i64]
148 [
149 1
150 1
151 1
152 4
153 5
154 ]
156 """
157 # Find first and last valid indices
158 non_null_indices = _pl_non_null_indices(ts)
160 if len(non_null_indices) == 0:
161 return ts
163 first = non_null_indices[0]
164 last = non_null_indices[-1]
166 # Create a new series with the same length as the original
167 values = ts.to_list()
169 # Fill forward within the slice between first and last valid indices
170 current_value = None
171 for i in range(first, last + 1):
172 if values[i] is not None:
173 current_value = values[i]
174 elif current_value is not None:
175 values[i] = current_value
177 # Create a new series with the filled values
178 return pl.Series(values, dtype=ts.dtype)
181def valid_pl(ts: pl.Series) -> bool:
182 """Check if a polars time series has no missing values between the first and last valid indices.
184 This function verifies that a time series doesn't have any null values in the middle.
185 It's acceptable to have nulls at the beginning or end of the series.
187 Parameters
188 ----------
189 ts : pl.Series
190 The time series to check
192 Returns:
193 -------
194 bool
195 True if the time series has no missing values between the first and last valid indices,
196 False otherwise
198 Examples:
199 --------
200 >>> import polars as pl
201 >>> ts1 = pl.Series([None, 1, 2, 3, None]) # Nulls only at beginning and end
202 >>> valid_pl(ts1)
203 True
204 >>> ts2 = pl.Series([1, 2, None, 4, 5]) # Null in the middle
205 >>> valid_pl(ts2)
206 False
208 """
209 # Get indices of non-null values
210 non_null_indices = _pl_non_null_indices(ts)
212 if len(non_null_indices) <= 1:
213 return True
215 # The series is valid when the non-null values form a contiguous run: their
216 # count equals the span between the first and last non-null index (inclusive).
217 first = non_null_indices[0]
218 last = non_null_indices[-1]
219 return bool(len(non_null_indices) == last - first + 1)
222def interpolate_df_pl(df: pl.DataFrame) -> pl.DataFrame:
223 """Interpolate missing values in a polars DataFrame between the first and last valid indices for each column.
225 This function applies interpolate_pl to each column of a DataFrame,
226 filling forward (ffill) missing values in each column, but only
227 between the first and last valid indices. Values outside this range remain null.
229 Parameters
230 ----------
231 df : pl.DataFrame
232 The DataFrame to interpolate
234 Returns:
235 -------
236 pl.DataFrame
237 The interpolated DataFrame
239 Examples:
240 --------
241 >>> import polars as pl
242 >>> df = pl.DataFrame({
243 ... 'A': [1.0, None, None, 4.0, 5.0],
244 ... 'B': [None, 2.0, None, 4.0, None]
245 ... })
246 >>> interpolate_df_pl(df)
247 shape: (5, 2)
248 ┌─────┬──────┐
249 │ A ┆ B │
250 │ --- ┆ --- │
251 │ f64 ┆ f64 │
252 ╞═════╪══════╡
253 │ 1.0 ┆ null │
254 │ 1.0 ┆ 2.0 │
255 │ 1.0 ┆ 2.0 │
256 │ 4.0 ┆ 4.0 │
257 │ 5.0 ┆ null │
258 └─────┴──────┘
260 """
261 # Apply interpolate_pl to each column
262 result = {}
263 for col in df.columns:
264 result[col] = interpolate_pl(df[col])
266 return pl.DataFrame(result)
269def valid_df_pl(df: pl.DataFrame) -> bool:
270 """Check if a polars DataFrame has no missing values between the first and last valid indices for each column.
272 This function verifies that each column in the DataFrame doesn't have any null values in the middle.
273 It's acceptable to have nulls at the beginning or end of each column.
275 Parameters
276 ----------
277 df : pl.DataFrame
278 The DataFrame to check
280 Returns:
281 -------
282 bool
283 True if all columns in the DataFrame have no missing values between their first and last valid indices,
284 False otherwise
286 Examples:
287 --------
288 >>> import polars as pl
289 >>> df1 = pl.DataFrame({
290 ... 'A': [None, 1, 2, 3, None], # Nulls only at beginning and end
291 ... 'B': [None, 2, 3, 4, None] # Nulls only at beginning and end
292 ... })
293 >>> valid_df_pl(df1)
294 True
295 >>> df2 = pl.DataFrame({
296 ... 'A': [1, 2, None, 4, 5], # Null in the middle
297 ... 'B': [1, 2, 3, 4, 5] # No nulls
298 ... })
299 >>> valid_df_pl(df2)
300 False
302 """
303 # Check each column
304 return all(valid_pl(df[col]) for col in df.columns)