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#[allow(dead_code)]
314#[pyfunction]
315pub fn polyfit_py(
316 py: Python,
317 x: &Bound<'_, PyArray1<f64>>,
318 y: &Bound<'_, PyArray1<f64>>,
319 deg: usize,
320) -> PyResult<Py<PyAny>> {
321 let x_readonly = x.readonly();
322 let x_arr = x_readonly.as_array();
323 let y_readonly = y.readonly();
324 let y_arr = y_readonly.as_array();
325 let result = polyfit::<f64>(&x_arr.view(), &y_arr.view(), deg)
326 .map_err(|e| PyRuntimeError::new_err(format!("Polynomial fit failed: {}", e)))?;
327 let dict = PyDict::new(py);
328 let coef_vec: Vec<f64> = result.coefficients.to_vec();
329 dict.set_item("coefficients", coef_vec)?;
330 dict.set_item("r_squared", result.r_squared)?;
331 dict.set_item("adj_r_squared", result.adj_r_squared)?;
332 let residuals_vec: Vec<f64> = result.residuals.to_vec();
333 dict.set_item("residuals", residuals_vec)?;
334 let fitted_vec: Vec<f64> = result.fitted_values.to_vec();
335 dict.set_item("fitted_values", fitted_vec)?;
336 Ok(dict.into())
337}
338#[pyfunction]
355#[pyo3(signature = (*args, alpha = 0.05))]
356pub fn tukey_hsd_py(
357 py: Python,
358 args: &Bound<'_, pyo3::types::PyTuple>,
359 alpha: f64,
360) -> PyResult<Py<PyAny>> {
361 if args.len() < 2 {
362 return Err(PyRuntimeError::new_err(
363 "Need at least 2 groups for Tukey's HSD",
364 ));
365 }
366 let mut arrays = Vec::new();
367 for item in args.iter() {
368 let arr: &Bound<'_, PyArray1<f64>> = item.cast()?;
369 let readonly = arr.readonly();
370 let owned = readonly.as_array().to_owned();
371 arrays.push(owned);
372 }
373 let views: Vec<_> = arrays.iter().map(|a| a.view()).collect();
374 let view_refs: Vec<&_> = views.iter().collect();
375 let results = tukey_hsd(&view_refs, alpha)
376 .map_err(|e| PyRuntimeError::new_err(format!("Tukey's HSD failed: {}", e)))?;
377 let result_list = pyo3::types::PyList::empty(py);
378 for (group1, group2, mean_diff, pvalue, significant) in results {
379 let dict = PyDict::new(py);
380 dict.set_item("group1", group1)?;
381 dict.set_item("group2", group2)?;
382 dict.set_item("mean_diff", mean_diff)?;
383 dict.set_item("pvalue", pvalue)?;
384 dict.set_item("significant", significant)?;
385 result_list.append(dict)?;
386 }
387 Ok(result_list.into())
388}
389#[pyfunction]
403#[pyo3(signature = (x, y, alternative = "two-sided"))]
404pub fn pearsonr_py(
405 py: Python,
406 x: &Bound<'_, PyArray1<f64>>,
407 y: &Bound<'_, PyArray1<f64>>,
408 alternative: &str,
409) -> PyResult<Py<PyAny>> {
410 let x_readonly = x.readonly();
411 let x_arr = x_readonly.as_array();
412 let y_readonly = y.readonly();
413 let y_arr = y_readonly.as_array();
414 let (r, pvalue) = pearsonr(&x_arr.view(), &y_arr.view(), alternative)
415 .map_err(|e| PyRuntimeError::new_err(format!("Pearson correlation test failed: {}", e)))?;
416 let dict = PyDict::new(py);
417 dict.set_item("correlation", r)?;
418 dict.set_item("pvalue", pvalue)?;
419 Ok(dict.into())
420}
421#[pyfunction]
435#[pyo3(signature = (x, y, alternative = "two-sided"))]
436pub fn spearmanr_py(
437 py: Python,
438 x: &Bound<'_, PyArray1<f64>>,
439 y: &Bound<'_, PyArray1<f64>>,
440 alternative: &str,
441) -> PyResult<Py<PyAny>> {
442 let x_readonly = x.readonly();
443 let x_arr = x_readonly.as_array();
444 let y_readonly = y.readonly();
445 let y_arr = y_readonly.as_array();
446 let (rho, pvalue) = spearmanr(&x_arr.view(), &y_arr.view(), alternative)
447 .map_err(|e| PyRuntimeError::new_err(format!("Spearman correlation test failed: {}", e)))?;
448 let dict = PyDict::new(py);
449 dict.set_item("correlation", rho)?;
450 dict.set_item("pvalue", pvalue)?;
451 Ok(dict.into())
452}