1use pyo3::exceptions::PyRuntimeError;
6use pyo3::prelude::*;
7use pyo3::types::{PyAny, PyDict};
8use scirs2_core::Array2;
9use scirs2_numpy::{PyArray1, PyArray2, PyArrayMethods};
10use scirs2_stats::contingency::{fisher_exact, odds_ratio, relative_risk};
11use scirs2_stats::{
12 chi2_independence, chi2_yates, friedman, ks_2samp, linregress, pearsonr, polyfit, spearmanr,
13 tukey_hsd,
14};
15
16#[pyfunction]
29#[pyo3(signature = (x, y, alternative = "two-sided"))]
30pub fn ks_2samp_py(
31 py: Python,
32 x: &Bound<'_, PyArray1<f64>>,
33 y: &Bound<'_, PyArray1<f64>>,
34 alternative: &str,
35) -> PyResult<Py<PyAny>> {
36 let x_data = x.readonly();
37 let x_arr = x_data.as_array();
38 let y_data = y.readonly();
39 let y_arr = y_data.as_array();
40 let (statistic, pvalue) = ks_2samp(&x_arr.view(), &y_arr.view(), alternative)
41 .map_err(|e| PyRuntimeError::new_err(format!("Two-sample KS test failed: {}", e)))?;
42 let dict = PyDict::new(py);
43 dict.set_item("statistic", statistic)?;
44 dict.set_item("pvalue", pvalue)?;
45 Ok(dict.into())
46}
47#[pyfunction]
58pub fn friedman_py(py: Python, data: &Bound<'_, PyArray2<f64>>) -> PyResult<Py<PyAny>> {
59 let data_readonly = data.readonly();
60 let data_view = data_readonly.as_array();
61 let data_arr = scirs2_core::ndarray::Array2::from_shape_vec(
62 data_view.dim(),
63 data_view.iter().copied().collect(),
64 )
65 .map_err(|e| PyRuntimeError::new_err(format!("Array conversion failed: {}", e)))?;
66 let (statistic, pvalue) = friedman(&data_arr.view())
67 .map_err(|e| PyRuntimeError::new_err(format!("Friedman test failed: {}", e)))?;
68 let dict = PyDict::new(py);
69 dict.set_item("statistic", statistic)?;
70 dict.set_item("pvalue", pvalue)?;
71 Ok(dict.into())
72}
73#[pyfunction]
84pub fn chi2_independence_py(
85 py: Python,
86 observed: &Bound<'_, PyArray2<i64>>,
87) -> PyResult<Py<PyAny>> {
88 let data_readonly = observed.readonly();
89 let data_view = data_readonly.as_array();
90 let data_arr = scirs2_core::ndarray::Array2::from_shape_vec(
91 data_view.dim(),
92 data_view.iter().copied().collect(),
93 )
94 .map_err(|e| PyRuntimeError::new_err(format!("Array conversion failed: {}", e)))?;
95 let result = chi2_independence::<f64, i64>(&data_arr.view()).map_err(|e| {
96 PyRuntimeError::new_err(format!("Chi-square independence test failed: {}", e))
97 })?;
98 let dict = PyDict::new(py);
99 dict.set_item("statistic", result.statistic)?;
100 dict.set_item("pvalue", result.p_value)?;
101 dict.set_item("df", result.df)?;
102 let shape = result.expected.dim();
103 let expected_vec: Vec<Vec<f64>> = (0..shape.0)
104 .map(|i| (0..shape.1).map(|j| result.expected[(i, j)]).collect())
105 .collect();
106 let expected_py = PyArray2::from_vec2(py, &expected_vec)
107 .map_err(|e| PyRuntimeError::new_err(format!("Failed to create expected array: {}", e)))?;
108 dict.set_item("expected", expected_py)?;
109 Ok(dict.into())
110}
111#[pyfunction]
122pub fn chi2_yates_py(py: Python, observed: &Bound<'_, PyArray2<i64>>) -> PyResult<Py<PyAny>> {
123 let data_readonly = observed.readonly();
124 let data_view = data_readonly.as_array();
125 let shape = data_view.dim();
126 if shape.0 != 2 || shape.1 != 2 {
127 return Err(PyRuntimeError::new_err(
128 "Yates' correction requires a 2x2 contingency table",
129 ));
130 }
131 let data_arr = scirs2_core::ndarray::Array2::from_shape_vec(
132 data_view.dim(),
133 data_view.iter().copied().collect(),
134 )
135 .map_err(|e| PyRuntimeError::new_err(format!("Array conversion failed: {}", e)))?;
136 let result = chi2_yates::<f64, i64>(&data_arr.view())
137 .map_err(|e| PyRuntimeError::new_err(format!("Chi-square Yates' test failed: {}", e)))?;
138 let dict = PyDict::new(py);
139 dict.set_item("statistic", result.statistic)?;
140 dict.set_item("pvalue", result.p_value)?;
141 dict.set_item("df", result.df)?;
142 let expected_vec: Vec<f64> = result.expected.iter().copied().collect();
143 let expected_py = PyArray2::from_vec2(
144 py,
145 &[expected_vec[0..2].to_vec(), expected_vec[2..4].to_vec()],
146 )
147 .map_err(|e| PyRuntimeError::new_err(format!("Failed to create expected array: {}", e)))?;
148 dict.set_item("expected", expected_py)?;
149 Ok(dict.into())
150}
151#[pyfunction]
169#[pyo3(signature = (table, alternative = "two-sided"))]
170pub fn fisher_exact_py(
171 py: Python,
172 table: &Bound<'_, PyArray2<f64>>,
173 alternative: &str,
174) -> PyResult<Py<PyAny>> {
175 let table_readonly = table.readonly();
176 let table_arr = Array2::from_shape_vec(
177 table_readonly.as_array().dim(),
178 table_readonly.as_array().iter().copied().collect(),
179 )
180 .map_err(|e| PyRuntimeError::new_err(format!("Array conversion failed: {}", e)))?;
181 let (odds_ratio, pvalue) = fisher_exact(&table_arr.view(), alternative)
182 .map_err(|e| PyRuntimeError::new_err(format!("Fisher's exact test failed: {}", e)))?;
183 let dict = PyDict::new(py);
184 dict.set_item("odds_ratio", odds_ratio)?;
185 dict.set_item("pvalue", pvalue)?;
186 Ok(dict.into())
187}
188#[pyfunction]
207pub fn odds_ratio_py(table: &Bound<'_, PyArray2<f64>>) -> PyResult<f64> {
208 let table_readonly = table.readonly();
209 let table_arr = Array2::from_shape_vec(
210 table_readonly.as_array().dim(),
211 table_readonly.as_array().iter().copied().collect(),
212 )
213 .map_err(|e| PyRuntimeError::new_err(format!("Array conversion failed: {}", e)))?;
214 let or = odds_ratio(&table_arr.view())
215 .map_err(|e| PyRuntimeError::new_err(format!("Odds ratio calculation failed: {}", e)))?;
216 Ok(or)
217}
218#[pyfunction]
237pub fn relative_risk_py(table: &Bound<'_, PyArray2<f64>>) -> PyResult<f64> {
238 let table_readonly = table.readonly();
239 let table_arr = Array2::from_shape_vec(
240 table_readonly.as_array().dim(),
241 table_readonly.as_array().iter().copied().collect(),
242 )
243 .map_err(|e| PyRuntimeError::new_err(format!("Array conversion failed: {}", e)))?;
244 let rr = relative_risk(&table_arr.view())
245 .map_err(|e| PyRuntimeError::new_err(format!("Relative risk calculation failed: {}", e)))?;
246 Ok(rr)
247}
248#[pyfunction]
267pub fn linregress_py(
268 py: Python,
269 x: &Bound<'_, PyArray1<f64>>,
270 y: &Bound<'_, PyArray1<f64>>,
271) -> PyResult<Py<PyAny>> {
272 let x_readonly = x.readonly();
273 let x_arr = x_readonly.as_array();
274 let y_readonly = y.readonly();
275 let y_arr = y_readonly.as_array();
276 let (slope, intercept, rvalue, pvalue, stderr) = linregress(&x_arr.view(), &y_arr.view())
277 .map_err(|e| PyRuntimeError::new_err(format!("Linear regression failed: {}", e)))?;
278 let dict = PyDict::new(py);
279 dict.set_item("slope", slope)?;
280 dict.set_item("intercept", intercept)?;
281 dict.set_item("rvalue", rvalue)?;
282 dict.set_item("pvalue", pvalue)?;
283 dict.set_item("stderr", stderr)?;
284 Ok(dict.into())
285}
286#[pyfunction]
312pub fn polyfit_py(
313 py: Python,
314 x: &Bound<'_, PyArray1<f64>>,
315 y: &Bound<'_, PyArray1<f64>>,
316 deg: usize,
317) -> PyResult<Py<PyAny>> {
318 let x_readonly = x.readonly();
319 let x_arr = x_readonly.as_array();
320 let y_readonly = y.readonly();
321 let y_arr = y_readonly.as_array();
322 let result = polyfit::<f64>(&x_arr.view(), &y_arr.view(), deg)
323 .map_err(|e| PyRuntimeError::new_err(format!("Polynomial fit failed: {}", e)))?;
324 let dict = PyDict::new(py);
325 let coef_vec: Vec<f64> = result.coefficients.to_vec();
326 dict.set_item("coefficients", coef_vec)?;
327 dict.set_item("r_squared", result.r_squared)?;
328 dict.set_item("adj_r_squared", result.adj_r_squared)?;
329 let residuals_vec: Vec<f64> = result.residuals.to_vec();
330 dict.set_item("residuals", residuals_vec)?;
331 let fitted_vec: Vec<f64> = result.fitted_values.to_vec();
332 dict.set_item("fitted_values", fitted_vec)?;
333 Ok(dict.into())
334}
335#[pyfunction]
352#[pyo3(signature = (*args, alpha = 0.05))]
353pub fn tukey_hsd_py(
354 py: Python,
355 args: &Bound<'_, pyo3::types::PyTuple>,
356 alpha: f64,
357) -> PyResult<Py<PyAny>> {
358 if args.len() < 2 {
359 return Err(PyRuntimeError::new_err(
360 "Need at least 2 groups for Tukey's HSD",
361 ));
362 }
363 let mut arrays = Vec::new();
364 for item in args.iter() {
365 let arr: &Bound<'_, PyArray1<f64>> = item.cast()?;
366 let readonly = arr.readonly();
367 let owned = readonly.as_array().to_owned();
368 arrays.push(owned);
369 }
370 let views: Vec<_> = arrays.iter().map(|a| a.view()).collect();
371 let view_refs: Vec<&_> = views.iter().collect();
372 let results = tukey_hsd(&view_refs, alpha)
373 .map_err(|e| PyRuntimeError::new_err(format!("Tukey's HSD failed: {}", e)))?;
374 let result_list = pyo3::types::PyList::empty(py);
375 for (group1, group2, mean_diff, pvalue, significant) in results {
376 let dict = PyDict::new(py);
377 dict.set_item("group1", group1)?;
378 dict.set_item("group2", group2)?;
379 dict.set_item("mean_diff", mean_diff)?;
380 dict.set_item("pvalue", pvalue)?;
381 dict.set_item("significant", significant)?;
382 result_list.append(dict)?;
383 }
384 Ok(result_list.into())
385}
386#[pyfunction]
400#[pyo3(signature = (x, y, alternative = "two-sided"))]
401pub fn pearsonr_py(
402 py: Python,
403 x: &Bound<'_, PyArray1<f64>>,
404 y: &Bound<'_, PyArray1<f64>>,
405 alternative: &str,
406) -> PyResult<Py<PyAny>> {
407 let x_readonly = x.readonly();
408 let x_arr = x_readonly.as_array();
409 let y_readonly = y.readonly();
410 let y_arr = y_readonly.as_array();
411 let (r, pvalue) = pearsonr(&x_arr.view(), &y_arr.view(), alternative)
412 .map_err(|e| PyRuntimeError::new_err(format!("Pearson correlation test failed: {}", e)))?;
413 let dict = PyDict::new(py);
414 dict.set_item("correlation", r)?;
415 dict.set_item("pvalue", pvalue)?;
416 Ok(dict.into())
417}
418#[pyfunction]
432#[pyo3(signature = (x, y, alternative = "two-sided"))]
433pub fn spearmanr_py(
434 py: Python,
435 x: &Bound<'_, PyArray1<f64>>,
436 y: &Bound<'_, PyArray1<f64>>,
437 alternative: &str,
438) -> PyResult<Py<PyAny>> {
439 let x_readonly = x.readonly();
440 let x_arr = x_readonly.as_array();
441 let y_readonly = y.readonly();
442 let y_arr = y_readonly.as_array();
443 let (rho, pvalue) = spearmanr(&x_arr.view(), &y_arr.view(), alternative)
444 .map_err(|e| PyRuntimeError::new_err(format!("Spearman correlation test failed: {}", e)))?;
445 let dict = PyDict::new(py);
446 dict.set_item("correlation", rho)?;
447 dict.set_item("pvalue", pvalue)?;
448 Ok(dict.into())
449}