1use crate::error::SciRS2Error;
28#[cfg(feature = "series")]
29use crate::series::PyTimeSeries;
30use pyo3::prelude::*;
31use pyo3::types::PyDict;
32#[cfg(feature = "series")]
33use pyo3::types::PyList;
34use scirs2_numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods};
35
36#[cfg(feature = "series")]
47#[pyfunction]
48pub fn pandas_to_timeseries(py: Python, series: Py<PyAny>) -> PyResult<PyTimeSeries> {
49 let values_obj = series.getattr(py, "values")?;
51 let values_array = values_obj.cast_bound::<PyArray1<f64>>(py)?;
52
53 let index = series.getattr(py, "index")?;
55
56 let timestamps = if let Ok(index_values) = index.getattr(py, "values") {
58 if index.getattr(py, "to_numpy").is_ok() {
60 let timestamp_method = index.getattr(py, "astype")?;
62 let timestamps_ns = timestamp_method.call1(py, ("int64",))?;
63 let ts_values = timestamps_ns.getattr(py, "values")?;
64 let timestamps_array = ts_values.cast_bound::<PyArray1<i64>>(py)?;
65
66 let binding = timestamps_array.readonly();
68 let ts_arr = binding.as_array();
69 let ts_vec: Vec<f64> = ts_arr.iter().map(|&ns| ns as f64 / 1e9).collect();
70
71 Some(scirs2_core::Array1::from_vec(ts_vec))
72 } else if let Ok(index_array) = index_values.cast_bound::<PyArray1<f64>>(py) {
73 let binding = index_array.readonly();
75 let idx_arr = binding.as_array();
76 Some(idx_arr.to_owned())
77 } else {
78 None
79 }
80 } else {
81 None
82 };
83
84 let binding = values_array.readonly();
86 let values_arr = binding.as_array();
87 let values_owned = values_arr.to_owned();
88
89 Ok(PyTimeSeries::from_arrays(values_owned, timestamps))
90}
91
92#[cfg(feature = "series")]
100#[pyfunction]
101pub fn timeseries_to_pandas(py: Python, ts: &PyTimeSeries) -> PyResult<Py<PyAny>> {
102 let pandas = py.import("pandas")?;
104
105 let values = ts.values_owned().into_pyarray(py).unbind();
107
108 let series = if let Some(timestamps) = ts.timestamps_owned() {
110 let timestamps_ns: Vec<i64> = timestamps.iter().map(|&s| (s * 1e9) as i64).collect();
112
113 let datetime_index = pandas
114 .getattr("DatetimeIndex")?
115 .call1((PyList::new(py, ×tamps_ns)?,))?;
116
117 let kwargs = PyDict::new(py);
119 kwargs.set_item("index", datetime_index)?;
120 pandas.getattr("Series")?.call((values,), Some(&kwargs))
121 } else {
122 pandas.getattr("Series")?.call1((values,))
124 }?;
125
126 Ok(series.into())
127}
128
129#[pyfunction]
140pub fn dataframe_to_array(py: Python, df: Py<PyAny>) -> PyResult<Py<PyArray2<f64>>> {
141 let values = df.getattr(py, "values")?;
143 let array = values.cast_bound::<PyArray2<f64>>(py)?;
144
145 Ok(array.to_owned().unbind())
146}
147
148#[pyfunction]
158#[pyo3(signature = (array, columns=None, index=None))]
159pub fn array_to_dataframe(
160 py: Python,
161 array: &Bound<'_, PyArray2<f64>>,
162 columns: Option<Vec<String>>,
163 index: Option<Py<PyAny>>,
164) -> PyResult<Py<PyAny>> {
165 let pandas = py.import("pandas")?;
167
168 let kwargs = PyDict::new(py);
170 if let Some(cols) = columns {
171 kwargs.set_item("columns", cols)?;
172 }
173 if let Some(idx) = index {
174 kwargs.set_item("index", idx)?;
175 }
176
177 let df = pandas.getattr("DataFrame")?.call((array,), Some(&kwargs))?;
178
179 Ok(df.into())
180}
181
182#[pyfunction]
199pub fn apply_to_dataframe(py: Python, df: Py<PyAny>, func: Py<PyAny>) -> PyResult<Py<PyAny>> {
200 let pandas = py.import("pandas")?;
202
203 let columns = df.getattr(py, "columns")?;
205 let col_list: Vec<String> = columns.extract(py)?;
206
207 let results = PyDict::new(py);
209 for col_name in col_list {
210 let column = df.call_method1(py, "__getitem__", (&col_name,))?;
211 let values = column.getattr(py, "values")?;
212 let result = func.call1(py, (values,))?;
213 results.set_item(&col_name, result)?;
214 }
215
216 let series = pandas.getattr("Series")?.call1((results,))?;
218 Ok(series.into())
219}
220
221#[pyfunction]
231#[pyo3(signature = (df, func, axis=0))]
232pub fn apply_along_axis(
233 py: Python,
234 df: Py<PyAny>,
235 func: Py<PyAny>,
236 axis: usize,
237) -> PyResult<Py<PyAny>> {
238 let pandas = py.import("pandas")?;
239
240 if axis == 0 {
241 apply_to_dataframe(py, df, func)
243 } else {
244 let values = df.getattr(py, "values")?;
246 let array = values.cast_bound::<PyArray2<f64>>(py)?;
247 let binding = array.readonly();
248 let arr = binding.as_array();
249
250 let results: Vec<f64> = arr
251 .rows()
252 .into_iter()
253 .map(|row| {
254 let row_array = row.to_owned().into_pyarray(py);
255 func.call1(py, (row_array,))
256 .and_then(|r| r.extract::<f64>(py))
257 })
258 .collect::<Result<Vec<_>, _>>()?;
259
260 let series = pandas.getattr("Series")?.call1((results,))?;
261 Ok(series.into())
262 }
263}
264
265#[pyfunction]
278pub fn rolling_apply(
279 py: Python,
280 series: Py<PyAny>,
281 window: usize,
282 func: Py<PyAny>,
283) -> PyResult<Py<PyAny>> {
284 let pandas = py.import("pandas")?;
285
286 let values = series.getattr(py, "values")?;
288 let array = values.cast_bound::<PyArray1<f64>>(py)?;
289 let binding = array.readonly();
290 let arr = binding.as_array();
291
292 if arr.len() < window {
293 return Err(SciRS2Error::ValueError(format!(
294 "Window size {} is larger than array length {}",
295 window,
296 arr.len()
297 ))
298 .into());
299 }
300
301 let mut results = Vec::with_capacity(arr.len() - window + 1);
303 for i in 0..=(arr.len() - window) {
304 let window_slice = arr.slice(ndarray::s![i..i + window]);
305 let window_array = window_slice.to_owned().into_pyarray(py);
306 let result: f64 = func.call1(py, (window_array,))?.extract(py)?;
307 results.push(result);
308 }
309
310 let mut padded = vec![f64::NAN; window - 1];
312 padded.extend(results);
313
314 let series_result = pandas.getattr("Series")?.call1((padded,))?;
316 Ok(series_result.into())
317}
318
319pub fn register_pandas_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
321 #[cfg(feature = "series")]
322 m.add_function(wrap_pyfunction!(pandas_to_timeseries, m)?)?;
323 #[cfg(feature = "series")]
324 m.add_function(wrap_pyfunction!(timeseries_to_pandas, m)?)?;
325 m.add_function(wrap_pyfunction!(dataframe_to_array, m)?)?;
326 m.add_function(wrap_pyfunction!(array_to_dataframe, m)?)?;
327 m.add_function(wrap_pyfunction!(apply_to_dataframe, m)?)?;
328 m.add_function(wrap_pyfunction!(apply_along_axis, m)?)?;
329 m.add_function(wrap_pyfunction!(rolling_apply, m)?)?;
330 Ok(())
331}