febolt 0.1.56__tar.gz → 0.1.57__tar.gz

{febolt-0.1.56 → febolt-0.1.57}/Cargo.lock

@@ -362,7 +362,7 @@ checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be"
 
 [[package]]
 name = "febolt"
-version = "0.1.56"
+version = "0.1.57"
 dependencies = [
  "blas",
  "intel-mkl-src",
@@ -1046,9 +1046,9 @@ dependencies = [
 
 [[package]]
 name = "openssl"
-version = "0.10.70"
+version = "0.10.71"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "61cfb4e166a8bb8c9b55c500bc2308550148ece889be90f609377e58140f42c6"
+checksum = "5e14130c6a98cd258fdcb0fb6d744152343ff729cbfcb28c656a9d12b999fbcd"
 dependencies = [
  "bitflags",
  "cfg-if",
@@ -1087,9 +1087,9 @@ dependencies = [
 
 [[package]]
 name = "openssl-sys"
-version = "0.9.105"
+version = "0.9.106"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "8b22d5b84be05a8d6947c7cb71f7c849aa0f112acd4bf51c2a7c1c988ac0a9dc"
+checksum = "8bb61ea9811cc39e3c2069f40b8b8e2e70d8569b361f879786cc7ed48b777cdd"
 dependencies = [
  "cc",
  "libc",

{febolt-0.1.56 → febolt-0.1.57}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "febolt"
-version = "0.1.56"
+version = "0.1.57"
 edition = "2021"
 description = "Statistics library for Python powered by Rust"
 license = "MIT"

{febolt-0.1.56 → febolt-0.1.57}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: febolt
-Version: 0.1.56
+Version: 0.1.57
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{febolt-0.1.56 → febolt-0.1.57}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "febolt"
-version = "0.1.56"
+version = "0.1.57"
 requires-python = ">=3.8"
 description = "A Rust-based Statistics and ML package, callable from Python."
 keywords = ["rust", "python", "Machine Learning", "Statistics", "pyo3"]

febolt-0.1.57/src/lib.rs

@@ -0,0 +1,296 @@
+// src/lib.rs
+
+use pyo3::prelude::*;
+use pyo3::types::{IntoPyDict, PyAny, PyDict, PyList, PyType};
+use pyo3::wrap_pyfunction;
+
+use numpy::{PyArray1, PyArray2, IntoPyArray};
+use ndarray::{Array1, Array2, ArrayView1, ArrayView2, s, Axis};
+use statrs::distribution::{Continuous, Normal, ContinuousCDF};
+use std::collections::HashMap;
+use std::fmt::Write as _;  // for building cluster-keys
+
+/// Add significance stars.
+fn add_significance_stars(p: f64) -> &'static str {
+    if p < 0.01 {
+        "***"
+    } else if p < 0.05 {
+        "**"
+    } else if p < 0.1 {
+        "*"
+    } else {
+        ""
+    }
+}
+
+/// A tiny helper to check if we have a Logit or Probit from statsmodels
+fn detect_model_type(model: &PyAny) -> Result<bool, PyErr> {
+    // Return is_logit: true if logit, false if probit
+    let model_obj = model.getattr("model").unwrap_or(model);
+    let cls: String = model_obj
+        .getattr("__class__")?
+        .getattr("__name__")?
+        .extract()?;
+    let lc = cls.to_lowercase();
+    if lc == "logit" {
+        Ok(true) // is_logit = true
+    } else if lc == "probit" {
+        Ok(false) // is_logit = false
+    } else {
+        Err(pyo3::exceptions::PyValueError::new_err(
+            format!("ame: only Logit or Probit supported, got {cls}"),
+        ))
+    }
+}
+
+/// Downcast a Python object to NumPy PyArray2<f64> => ndarray::ArrayView2<f64>.
+///
+/// Marked unsafe because .as_array() in pyo3-numpy is unsafe, trusting Python memory is valid.
+fn as_array2_f64<'py>(obj: &'py PyAny) -> PyResult<ArrayView2<'py, f64>> {
+    let pyarray = obj.downcast::<PyArray2<f64>>()?;
+    let view = unsafe { pyarray.as_array() };
+    Ok(view)
+}
+
+/// Similarly, for 1D arrays.
+fn as_array1_f64<'py>(obj: &'py PyAny) -> PyResult<ArrayView1<'py, f64>> {
+    let pyarray = obj.downcast::<PyArray1<f64>>()?;
+    let view = unsafe { pyarray.as_array() };
+    Ok(view)
+}
+
+/// Evaluate logistic or normal cdf
+fn cdf_logit_probit(is_logit: bool, z: f64) -> f64 {
+    if is_logit {
+        // logistic cdf
+        1.0 / (1.0 + (-z).exp())
+    } else {
+        // normal cdf
+        let dist = Normal::new(0.0, 1.0).unwrap();
+        dist.cdf(z)
+    }
+}
+/// Evaluate logistic or normal pdf
+fn pdf_logit_probit(is_logit: bool, z: f64) -> f64 {
+    if is_logit {
+        // logistic pdf
+        let e = z.exp();
+        e / (1.0 + e).powi(2)
+    } else {
+        // normal pdf
+        (-0.5 * z * z).exp() / (2.0 * std::f64::consts::PI).sqrt()
+    }
+}
+
+/// The main function that calculates Logit/Probit AME in the same style as the original Probit code
+#[pyfunction]
+fn ame<'py>(
+    py: Python<'py>,
+    model: &'py PyAny,   // Could be Logit or Probit
+    chunk_size: Option<usize>, // optional chunk
+) -> PyResult<&'py PyAny> {
+    // 1) detect Logit vs Probit
+    let is_logit = detect_model_type(model)?;
+
+    // 2) read params
+    let params_pyarray: &PyArray1<f64> = model.getattr("params")?.downcast()?;
+    let beta = unsafe { params_pyarray.as_array() };  // shape(k)
+
+    // 3) read cov
+    let cov_pyarray: &PyArray2<f64> = model.call_method0("cov_params")?.downcast()?;
+    let cov_beta = unsafe { cov_pyarray.as_array() }; // shape(k,k)
+
+    // 4) Get model object and handle exog (X) and exog_names
+    let model_obj = model.getattr("model").unwrap_or(model);
+
+    // Handle pandas DataFrame input
+    let exog_py = model_obj.getattr("exog")?;
+    let (x_pyarray, exog_names) = if let Ok(values) = exog_py.getattr("values") {
+        // Pandas DataFrame path
+        (
+            values.downcast::<PyArray2<f64>>()?,
+            exog_py.getattr("columns")?.extract::<Vec<String>>()?
+        )
+    } else {
+        // Numpy array path
+        (
+            exog_py.downcast::<PyArray2<f64>>()?,
+            model_obj.getattr("exog_names")?.extract::<Vec<String>>()?
+        )
+    };
+
+    let X = unsafe { x_pyarray.as_array() };
+    let (n, k) = (X.nrows(), X.ncols());
+    let chunk = chunk_size.unwrap_or(n);
+
+    // 5) Identify intercept columns
+    let intercept_indices: Vec<usize> = exog_names
+        .iter()
+        .enumerate()
+        .filter_map(|(i, nm)| {
+            let ln = nm.to_lowercase();
+            if ln == "const" || ln == "intercept" {
+                Some(i)
+            } else {
+                None
+            }
+        })
+        .collect();
+
+    // 6) Identify discrete columns => strictly 0/1
+    let is_discrete: Vec<usize> = exog_names
+        .iter()
+        .enumerate()
+        .filter_map(|(j, _)| {
+            if intercept_indices.contains(&j) {
+                None
+            } else {
+                let col_j = X.column(j);
+                if col_j.iter().all(|&v| v == 0.0 || v == 1.0) {
+                    Some(j)
+                } else {
+                    None
+                }
+            }
+        })
+        .collect();
+
+    // 7) Prepare accumulators
+    let mut sum_ame = vec![0.0; k];          // sum partial effects
+    let mut partial_jl_sums = vec![0.0; k * k];
+    let normal = Normal::new(0.0, 1.0).unwrap();
+
+    // 8) single pass with chunk
+    let mut idx_start = 0;
+    while idx_start < n {
+        let idx_end = (idx_start + chunk).min(n);
+        let x_chunk = X.slice(s![idx_start..idx_end, ..]);
+        let z_chunk = x_chunk.dot(&beta);  // shape(n_chunk)
+
+        // pdf => we might do partial for continuous
+        let pdf_chunk = z_chunk.mapv(|z| pdf_logit_probit(is_logit, z));
+
+        // handle discrete
+        for &j in &is_discrete {
+            let xj_col = x_chunk.column(j);
+            let b_j = beta[j];
+            // z_j1 => z + (1-xj)*b_j
+            // z_j0 => z - xj*b_j
+            let delta_j1 = (1.0 - &xj_col).mapv(|x| x * b_j);
+            let delta_j0 = xj_col.mapv(|x| -x * b_j);
+            let z_j1 = &z_chunk + &delta_j1;
+            let z_j0 = &z_chunk + &delta_j0;
+
+            let cdf_j1 = z_j1.mapv(|z| cdf_logit_probit(is_logit, z));
+            let cdf_j0 = z_j0.mapv(|z| cdf_logit_probit(is_logit, z));
+            // sum
+            let effect_sum = cdf_j1.sum() - cdf_j0.sum();
+            sum_ame[j] += effect_sum;
+
+            // partial_jl_sums => row j, col l
+            let pdf_j1 = z_j1.mapv(|z| pdf_logit_probit(is_logit, z));
+            let pdf_j0 = z_j0.mapv(|z| pdf_logit_probit(is_logit, z));
+            for l in 0..k {
+                let grad = if l == j {
+                    // special case
+                    pdf_j1.sum()
+                } else {
+                    let x_l = x_chunk.column(l);
+                    let diff_pdf = &pdf_j1 - &pdf_j0;
+                    diff_pdf.dot(&x_l)
+                };
+                partial_jl_sums[j * k + l] += grad;
+            }
+        }
+
+        // handle continuous
+        for j in 0..k {
+            if intercept_indices.contains(&j) || is_discrete.contains(&j) {
+                continue;
+            }
+            let b_j = beta[j];
+            // sum_ame
+            sum_ame[j] += b_j * pdf_chunk.sum();
+            // partial_jl_sums => row j, col l
+            for l in 0..k {
+                let grad = if j == l {
+                    pdf_chunk.sum()
+                } else {
+                    // - b_j * sum(z_chunk * x_col(l) * pdf_chunk)
+                    let x_l = x_chunk.column(l);
+                    // careful about sign from the original code
+                    -b_j * (&z_chunk * &x_l).dot(&pdf_chunk)
+                };
+                partial_jl_sums[j * k + l] += grad;
+            }
+        }
+
+        idx_start = idx_end;
+    }
+
+    // 9) average sums
+    let ame: Vec<f64> = sum_ame.iter().map(|v| v / (n as f64)).collect();
+
+    // gradient matrix => shape(k,k)
+    let mut grad_ame = Array2::<f64>::zeros((k,k));
+    for j in 0..k {
+        for l in 0..k {
+            grad_ame[[j,l]] = partial_jl_sums[j * k + l] / (n as f64);
+        }
+    }
+
+    // cov => grad_ame * cov_beta * grad_ame^T
+    let cov_ame = grad_ame.dot(&cov_beta).dot(&grad_ame.t());
+    let var_ame: Vec<f64> = cov_ame.diag().iter().map(|&v| v.max(0.0)).collect();
+    let se_ame: Vec<f64> = var_ame.iter().map(|&v| v.sqrt()).collect();
+
+    // 10) Build final results
+    let (mut dy_dx, mut se_err, mut z_vals, mut p_vals, mut sig) =
+        (Vec::new(), Vec::new(), Vec::new(), Vec::new(), Vec::new());
+    let mut names_out = Vec::new();
+
+    for j in 0..k {
+        if intercept_indices.contains(&j) {
+            continue;
+        }
+        let dy = ame[j];
+        let s = se_ame[j];
+        dy_dx.push(dy);
+        se_err.push(s);
+        if s > 1e-15 {
+            let z = dy / s;
+            let p = 2.0*(1.0 - normal.cdf(z.abs()));
+            z_vals.push(z);
+            p_vals.push(p);
+            sig.push(add_significance_stars(p));
+        } else {
+            z_vals.push(0.0);
+            p_vals.push(1.0);
+            sig.push("");
+        }
+        names_out.push(exog_names[j].clone());
+    }
+
+    // 11) Create DataFrame
+    let pd = py.import("pandas")?;
+    let data = PyDict::new(py);
+    data.set_item("dy/dx", &dy_dx)?;
+    data.set_item("Std. Err", &se_err)?;
+    data.set_item("z", &z_vals)?;
+    data.set_item("Pr(>|z|)", &p_vals)?;
+    data.set_item("Significance", &sig)?;
+
+    let kwargs = PyDict::new(py);
+    kwargs.set_item("data", data)?;
+    kwargs.set_item("index", &names_out)?;
+
+    let df = pd.call_method("DataFrame", (), Some(kwargs))?;
+    Ok(df)
+}
+
+
+#[pymodule]
+fn febolt(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(ame, m)?)?;
+    Ok(())
+}

febolt-0.1.56/src/lib.rs

@@ -1,281 +0,0 @@
-use pyo3::prelude::*;
-use pyo3::types::{PyAny, PyDict};
-use pyo3::wrap_pyfunction;
-
-use numpy::{PyArray1, PyArray2};
-use ndarray::{Array1, Array2, Axis, s};
-use statrs::distribution::{Normal, ContinuousCDF};
-use std::f64;
-
-/// Distinguish whether statsmodels model is Logit or Probit
-fn is_logit_model(model: &PyAny) -> PyResult<bool> {
-    // Return true if Logit, false if Probit
-    let model_obj = model.getattr("model").unwrap_or(model);
-    let cls_name: String = model_obj.getattr("__class__")?.getattr("__name__")?.extract()?;
-    let lower = cls_name.to_lowercase();
-    if lower == "logit" {
-        Ok(true)
-    } else if lower == "probit" {
-        Ok(false)
-    } else {
-        Err(pyo3::exceptions::PyValueError::new_err(
-            format!("ame() only supports Logit or Probit. Got: {cls_name}"),
-        ))
-    }
-}
-
-/// Evaluate logistic or normal cdf
-fn cdf_logit_probit(is_logit: bool, z: f64) -> f64 {
-    if is_logit {
-        // logistic cdf => 1/(1+exp(-z))
-        1.0 / (1.0 + (-z).exp())
-    } else {
-        // normal cdf
-        let dist = Normal::new(0.0, 1.0).unwrap();
-        dist.cdf(z)
-    }
-}
-
-/// Evaluate logistic or normal pdf
-fn pdf_logit_probit(is_logit: bool, z: f64) -> f64 {
-    if is_logit {
-        let e = z.exp();
-        e / (1.0 + e).powi(2)
-    } else {
-        (-0.5 * z * z).exp() / (2.0 * f64::consts::PI).sqrt()
-    }
-}
-
-/// Significance stars
-fn add_significance_stars(p: f64) -> &'static str {
-    if p < 0.001 {
-        "***"
-    } else if p < 0.01 {
-        "**"
-    } else if p < 0.05 {
-        "*"
-    } else if p < 0.1 {
-        "."
-    } else {
-        ""
-    }
-}
-
-/// Main AME function
-#[pyfunction]
-fn ame<'py>(
-    py: Python<'py>,
-    model: &'py PyAny, // statsmodels fitted results
-    chunk_size: Option<usize>,
-) -> PyResult<&'py PyAny> {
-    // 1) Detect Logit vs Probit
-    let is_logit = is_logit_model(model)?;
-
-    // 2) Extract 1D params => shape(k,)
-    let params_obj = model.getattr("params")?;
-    // Usually `params` is already a 1D numpy array, so:
-    let params_nd = params_obj.downcast::<PyArray1<f64>>()?;
-    let beta_view = unsafe { params_nd.as_array() };
-    let beta = beta_view.to_owned(); // shape(k)
-
-    let k = beta.len();
-
-    // 3) Extract covariance => shape(k,k)
-    let cov_py = model.call_method0("cov_params")?.downcast::<PyArray2<f64>>()?;
-    let cov_view = unsafe { cov_py.as_array() };
-    if cov_view.nrows() != k || cov_view.ncols() != k {
-        return Err(pyo3::exceptions::PyValueError::new_err(
-            "cov_params dimension mismatch with `params`"
-        ));
-    }
-    let cov_beta = cov_view.to_owned();
-
-    // 4) Extract exog_names from statsmodels
-    let model_obj = model.getattr("model")?;
-    let exog_names: Vec<String> = model_obj.getattr("exog_names")?.extract()?;
-
-    // 5) Extract exog
-    // In many cases, `model_obj.getattr("exog")` is already a NumPy array.
-    // But if it's a DataFrame, we forcibly do .to_numpy().
-    let exog_obj = model_obj.getattr("exog")?;
-    // We'll try direct downcast first:
-    let exog_pyarray = match exog_obj.downcast::<PyArray2<f64>>() {
-        Ok(arr) => arr,
-        Err(_e) => {
-            // fallback => call .to_numpy()
-            let np = exog_obj.call_method0("to_numpy")?;
-            np.downcast::<PyArray2<f64>>()?
-        }
-    };
-    let X_view = unsafe { exog_pyarray.as_array() };
-    let (n, k_exog) = X_view.dim();
-    if k_exog != k {
-        return Err(pyo3::exceptions::PyValueError::new_err(
-            format!("exog dimension mismatch: exog has {} columns, params has length={}", k_exog, k)
-        ));
-    }
-
-    let chunk = chunk_size.unwrap_or(n);
-
-    // 6) Identify intercept columns
-    let intercept_idx: Vec<usize> = exog_names
-        .iter()
-        .enumerate()
-        .filter_map(|(i, nm)| {
-            let ln = nm.to_lowercase();
-            if ln=="intercept" || ln=="const" {
-                Some(i)
-            } else {
-                None
-            }
-        })
-        .collect();
-
-    // 7) Identify discrete columns => strictly 0/1
-    let mut disc_idx = Vec::new();
-    for j in 0..k {
-        if intercept_idx.contains(&j) {
-            continue;
-        }
-        let col_j = X_view.column(j);
-        if col_j.iter().all(|&v| v==0.0 || v==1.0) {
-            disc_idx.push(j);
-        }
-    }
-
-    // 8) Single pass, chunked
-    let mut sum_ame = vec![0.0; k];
-    let mut partial_jl_sums = vec![0.0; k*k];
-
-    let mut idx_start = 0;
-    while idx_start < n {
-        let idx_end = (idx_start + chunk).min(n);
-        let x_chunk = X_view.slice(s![idx_start..idx_end, ..]);
-        let z_chunk = x_chunk.dot(&beta); // shape(n_chunk)
-
-        // pdf for continuous partials
-        let pdf_chunk = z_chunk.mapv(|z| pdf_logit_probit(is_logit, z));
-
-        // handle discrete
-        for &jj in &disc_idx {
-            let b_j = beta[jj];
-            let col_j = x_chunk.column(jj);
-            // z_j1 => z_chunk + (1-col_j)*b_j
-            // z_j0 => z_chunk - col_j*b_j
-            let delta_j1 = (1.0 - &col_j).mapv(|x| x*b_j);
-            let delta_j0 = col_j.mapv(|x| -x*b_j);
-            let z_j1 = &z_chunk + &delta_j1;
-            let z_j0 = &z_chunk + &delta_j0;
-
-            let cdf_j1 = z_j1.mapv(|z| cdf_logit_probit(is_logit, z));
-            let cdf_j0 = z_j0.mapv(|z| cdf_logit_probit(is_logit, z));
-            let eff = cdf_j1.sum() - cdf_j0.sum();
-            sum_ame[jj]+= eff;
-
-            let pdf_j1 = z_j1.mapv(|z| pdf_logit_probit(is_logit, z));
-            let pdf_j0 = z_j0.mapv(|z| pdf_logit_probit(is_logit, z));
-            for l in 0..k {
-                let grad = if l==jj {
-                    pdf_j1.sum()
-                } else {
-                    let col_l = x_chunk.column(l);
-                    (&pdf_j1 - &pdf_j0).dot(&col_l)
-                };
-                partial_jl_sums[jj*k + l]+= grad;
-            }
-        }
-
-        // handle continuous
-        for j in 0..k {
-            if intercept_idx.contains(&j) || disc_idx.contains(&j) {
-                continue;
-            }
-            let b_j = beta[j];
-            sum_ame[j]+= b_j* pdf_chunk.sum();
-
-            for l in 0..k {
-                let grad = if j==l {
-                    pdf_chunk.sum()
-                } else {
-                    let col_l = x_chunk.column(l);
-                    // - b_j * sum( z_chunk * col_l * pdf_chunk)
-                    -b_j * (&z_chunk * &col_l).dot(&pdf_chunk)
-                };
-                partial_jl_sums[j*k + l]+= grad;
-            }
-        }
-
-        idx_start= idx_end;
-    }
-
-    // 9) Average
-    let nf = n as f64;
-    let ame_vals: Vec<f64> = sum_ame.iter().map(|&v| v/nf).collect();
-
-    let mut grad_ame = ndarray::Array2::<f64>::zeros((k,k));
-    for j in 0..k {
-        for l in 0..k {
-            grad_ame[[j,l]] = partial_jl_sums[j*k + l]/ nf;
-        }
-    }
-
-    // cov(ame)
-    let j_c = grad_ame.dot(&cov_beta);
-    let cov_ame = j_c.dot(&grad_ame.t());
-    let normal = Normal::new(0.0, 1.0).unwrap();
-    let var_ame: Vec<f64> = cov_ame.diag().iter().map(|&x| x.max(0.0)).collect();
-    let se_vals: Vec<f64> = var_ame.iter().map(|&x| x.sqrt()).collect();
-
-    // 10) Build final output => skip intercept
-    let mut dy_dx = Vec::new();
-    let mut se_err= Vec::new();
-    let mut z_vals= Vec::new();
-    let mut p_vals= Vec::new();
-    let mut sigs= Vec::new();
-    let mut names_out= Vec::new();
-
-    for j in 0..k {
-        if intercept_idx.contains(&j) {
-            continue;
-        }
-        let dy = ame_vals[j];
-        let s = se_vals[j];
-        dy_dx.push(dy);
-        se_err.push(s);
-        if s>1e-14 {
-            let z = dy/s;
-            let p = 2.0*(1.0 - normal.cdf(z.abs()));
-            z_vals.push(z);
-            p_vals.push(p);
-            sigs.push(add_significance_stars(p));
-        } else {
-            z_vals.push(0.0);
-            p_vals.push(1.0);
-            sigs.push("");
-        }
-        names_out.push(exog_names[j].clone());
-    }
-
-    // 11) Build a Pandas DataFrame
-    let pd = py.import("pandas")?;
-    let data = PyDict::new(py);
-    data.set_item("dy/dx", &dy_dx)?;
-    data.set_item("Std. Err", &se_err)?;
-    data.set_item("z", &z_vals)?;
-    data.set_item("Pr(>|z|)", &p_vals)?;
-    data.set_item("Significance", &sigs)?;
-
-    let kwargs = PyDict::new(py);
-    kwargs.set_item("data", data)?;
-    kwargs.set_item("index", names_out)?;
-
-    let df = pd.call_method("DataFrame", (), Some(kwargs))?;
-    Ok(df)
-}
-
-/// Python module
-#[pymodule]
-fn febolt(_py: Python, m: &PyModule) -> PyResult<()> {
-    m.add_function(wrap_pyfunction!(ame, m)?)?;
-    Ok(())
-}