1use crate::error::SciRS2Error;
28use crate::series::PyTimeSeries;
29use pyo3::prelude::*;
30use pyo3::types::{PyDict, PyList};
31use scirs2_numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayMethods};
32
33#[pyfunction]
44pub fn pandas_to_timeseries(py: Python, series: Py<PyAny>) -> PyResult<PyTimeSeries> {
45 let values_obj = series.getattr(py, "values")?;
47 let values_array = values_obj.cast_bound::<PyArray1<f64>>(py)?;
48
49 let index = series.getattr(py, "index")?;
51
52 let timestamps = if let Ok(index_values) = index.getattr(py, "values") {
54 if index.getattr(py, "to_numpy").is_ok() {
56 let timestamp_method = index.getattr(py, "astype")?;
58 let timestamps_ns = timestamp_method.call1(py, ("int64",))?;
59 let ts_values = timestamps_ns.getattr(py, "values")?;
60 let timestamps_array = ts_values.cast_bound::<PyArray1<i64>>(py)?;
61
62 let binding = timestamps_array.readonly();
64 let ts_arr = binding.as_array();
65 let ts_vec: Vec<f64> = ts_arr.iter().map(|&ns| ns as f64 / 1e9).collect();
66
67 Some(scirs2_core::Array1::from_vec(ts_vec))
68 } else if let Ok(index_array) = index_values.cast_bound::<PyArray1<f64>>(py) {
69 let binding = index_array.readonly();
71 let idx_arr = binding.as_array();
72 Some(idx_arr.to_owned())
73 } else {
74 None
75 }
76 } else {
77 None
78 };
79
80 let binding = values_array.readonly();
82 let values_arr = binding.as_array();
83 let values_owned = values_arr.to_owned();
84
85 Ok(PyTimeSeries::from_arrays(values_owned, timestamps))
86}
87
88#[pyfunction]
96pub fn timeseries_to_pandas(py: Python, ts: &PyTimeSeries) -> PyResult<Py<PyAny>> {
97 let pandas = py.import("pandas")?;
99
100 let values = ts.values_owned().into_pyarray(py).unbind();
102
103 let series = if let Some(timestamps) = ts.timestamps_owned() {
105 let timestamps_ns: Vec<i64> = timestamps.iter().map(|&s| (s * 1e9) as i64).collect();
107
108 let datetime_index = pandas
109 .getattr("DatetimeIndex")?
110 .call1((PyList::new(py, ×tamps_ns)?,))?;
111
112 let kwargs = PyDict::new(py);
114 kwargs.set_item("index", datetime_index)?;
115 pandas.getattr("Series")?.call((values,), Some(&kwargs))
116 } else {
117 pandas.getattr("Series")?.call1((values,))
119 }?;
120
121 Ok(series.into())
122}
123
124#[pyfunction]
135pub fn dataframe_to_array(py: Python, df: Py<PyAny>) -> PyResult<Py<PyArray2<f64>>> {
136 let values = df.getattr(py, "values")?;
138 let array = values.cast_bound::<PyArray2<f64>>(py)?;
139
140 Ok(array.to_owned().unbind())
141}
142
143#[pyfunction]
153#[pyo3(signature = (array, columns=None, index=None))]
154pub fn array_to_dataframe(
155 py: Python,
156 array: &Bound<'_, PyArray2<f64>>,
157 columns: Option<Vec<String>>,
158 index: Option<Py<PyAny>>,
159) -> PyResult<Py<PyAny>> {
160 let pandas = py.import("pandas")?;
162
163 let kwargs = PyDict::new(py);
165 if let Some(cols) = columns {
166 kwargs.set_item("columns", cols)?;
167 }
168 if let Some(idx) = index {
169 kwargs.set_item("index", idx)?;
170 }
171
172 let df = pandas.getattr("DataFrame")?.call((array,), Some(&kwargs))?;
173
174 Ok(df.into())
175}
176
177#[pyfunction]
194pub fn apply_to_dataframe(py: Python, df: Py<PyAny>, func: Py<PyAny>) -> PyResult<Py<PyAny>> {
195 let pandas = py.import("pandas")?;
197
198 let columns = df.getattr(py, "columns")?;
200 let col_list: Vec<String> = columns.extract(py)?;
201
202 let results = PyDict::new(py);
204 for col_name in col_list {
205 let column = df.call_method1(py, "__getitem__", (&col_name,))?;
206 let values = column.getattr(py, "values")?;
207 let result = func.call1(py, (values,))?;
208 results.set_item(&col_name, result)?;
209 }
210
211 let series = pandas.getattr("Series")?.call1((results,))?;
213 Ok(series.into())
214}
215
216#[pyfunction]
226#[pyo3(signature = (df, func, axis=0))]
227pub fn apply_along_axis(
228 py: Python,
229 df: Py<PyAny>,
230 func: Py<PyAny>,
231 axis: usize,
232) -> PyResult<Py<PyAny>> {
233 let pandas = py.import("pandas")?;
234
235 if axis == 0 {
236 apply_to_dataframe(py, df, func)
238 } else {
239 let values = df.getattr(py, "values")?;
241 let array = values.cast_bound::<PyArray2<f64>>(py)?;
242 let binding = array.readonly();
243 let arr = binding.as_array();
244
245 let results: Vec<f64> = arr
246 .rows()
247 .into_iter()
248 .map(|row| {
249 let row_array = row.to_owned().into_pyarray(py);
250 func.call1(py, (row_array,))
251 .and_then(|r| r.extract::<f64>(py))
252 })
253 .collect::<Result<Vec<_>, _>>()?;
254
255 let series = pandas.getattr("Series")?.call1((results,))?;
256 Ok(series.into())
257 }
258}
259
260#[pyfunction]
273pub fn rolling_apply(
274 py: Python,
275 series: Py<PyAny>,
276 window: usize,
277 func: Py<PyAny>,
278) -> PyResult<Py<PyAny>> {
279 let pandas = py.import("pandas")?;
280
281 let values = series.getattr(py, "values")?;
283 let array = values.cast_bound::<PyArray1<f64>>(py)?;
284 let binding = array.readonly();
285 let arr = binding.as_array();
286
287 if arr.len() < window {
288 return Err(SciRS2Error::ValueError(format!(
289 "Window size {} is larger than array length {}",
290 window,
291 arr.len()
292 ))
293 .into());
294 }
295
296 let mut results = Vec::with_capacity(arr.len() - window + 1);
298 for i in 0..=(arr.len() - window) {
299 let window_slice = arr.slice(ndarray::s![i..i + window]);
300 let window_array = window_slice.to_owned().into_pyarray(py);
301 let result: f64 = func.call1(py, (window_array,))?.extract(py)?;
302 results.push(result);
303 }
304
305 let mut padded = vec![f64::NAN; window - 1];
307 padded.extend(results);
308
309 let series_result = pandas.getattr("Series")?.call1((padded,))?;
311 Ok(series_result.into())
312}
313
314pub fn register_pandas_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
316 m.add_function(wrap_pyfunction!(pandas_to_timeseries, m)?)?;
317 m.add_function(wrap_pyfunction!(timeseries_to_pandas, m)?)?;
318 m.add_function(wrap_pyfunction!(dataframe_to_array, m)?)?;
319 m.add_function(wrap_pyfunction!(array_to_dataframe, m)?)?;
320 m.add_function(wrap_pyfunction!(apply_to_dataframe, m)?)?;
321 m.add_function(wrap_pyfunction!(apply_along_axis, m)?)?;
322 m.add_function(wrap_pyfunction!(rolling_apply, m)?)?;
323 Ok(())
324}