npls1-python 0.1.0

//! Python bindings for the `npls1` crate (N-PLS1 regression).
//! Built for pyo3 0.29 / numpy 0.29.

use ndarray::{Array1, Array3};
use numpy::{IntoPyArray, PyArray1, PyArray2, PyReadonlyArray1, PyReadonlyArray3};
use pyo3::exceptions::{PyRuntimeError, PyValueError};
use pyo3::prelude::*;
use pyo3::types::PyDict;

use npls1::{Npls as RustNpls, NplsError};

/// Convert a Rust `NplsError` into an appropriate Python exception.
fn map_err(e: NplsError) -> PyErr {
    match e {
        NplsError::LinalgError(msg) => PyRuntimeError::new_err(format!("LinalgError: {msg}")),
        NplsError::ValueError(msg) => PyValueError::new_err(msg),
        NplsError::NotFitted => PyRuntimeError::new_err(
            "Model has not been fitted yet. Call fit() before predict().",
        ),
    }
}

/// N-PLS1 (N-way Partial Least Squares) regressor — Rust-backed.
#[pyclass(name = "Npls")]
pub struct PyNpls {
    inner: RustNpls,
}

#[pymethods]
impl PyNpls {
    #[new]
    #[pyo3(signature = (
        n_components,
        a,
        derivative_rang = None,
        norm_func = None,
        crash_norm_name = None,
        crash_norm_value = None,
        excitation_wavelenth = None,
        emission_wavelenth = None,
    ))]
    #[allow(clippy::too_many_arguments)]
    fn new(
        n_components: usize,
        a: f64,
        derivative_rang: Option<Vec<usize>>,
        norm_func: Option<Vec<String>>,
        crash_norm_name: Option<String>,
        crash_norm_value: Option<f64>,
        excitation_wavelenth: Option<PyReadonlyArray1<f64>>,
        emission_wavelenth: Option<PyReadonlyArray1<f64>>,
    ) -> PyResult<Self> {
        let deriv = derivative_rang.unwrap_or_default();
        let norms = norm_func.unwrap_or_default();

        let use_snr = !deriv.is_empty()
            || crash_norm_name.is_some()
            || excitation_wavelenth.is_some()
            || emission_wavelenth.is_some();

        let inner = if use_snr {
            let exc = match excitation_wavelenth {
                Some(arr) => arr.as_array().to_owned(),
                None => Array1::zeros(1),
            };
            let emi = match emission_wavelenth {
                Some(arr) => arr.as_array().to_owned(),
                None => Array1::zeros(1),
            };
            RustNpls::with_snr(
                n_components,
                a,
                deriv,
                norms,
                crash_norm_name,
                crash_norm_value,
                exc,
                emi,
            )
        } else {
            RustNpls::new(n_components, a)
        };

        Ok(PyNpls { inner })
    }

    /// Fit the model. `xtrain`: (n_samples, n_emission, n_excitation); `ytrain`: (n_samples,).
    fn fit<'py>(
        mut slf: PyRefMut<'py, Self>,
        xtrain: PyReadonlyArray3<'py, f64>,
        ytrain: PyReadonlyArray1<'py, f64>,
    ) -> PyResult<PyRefMut<'py, Self>> {
        let x: Array3<f64> = xtrain.as_array().to_owned();
        let y: Array1<f64> = ytrain.as_array().to_owned();
        slf.inner.fit(&x, &y).map_err(map_err)?;
        Ok(slf)
    }

    /// Predict target values. `xtest`: (n_samples, n_emission, n_excitation).
    fn predict<'py>(
        &self,
        py: Python<'py>,
        xtest: PyReadonlyArray3<'py, f64>,
    ) -> PyResult<Bound<'py, PyArray1<f64>>> {
        let x: Array3<f64> = xtest.as_array().to_owned();
        let preds = self.inner.predict(&x).map_err(map_err)?;
        Ok(preds.into_pyarray(py))
    }

    // ── read-only accessors ─────────────────────────────────────────

    #[getter]
    fn n_components(&self) -> usize {
        self.inner.n_components
    }

    #[getter]
    fn a(&self) -> f64 {
        self.inner.a
    }

    #[getter]
    fn train_error(&self) -> Option<f64> {
        self.inner.train_error
    }

    #[getter]
    fn w_k<'py>(&self, py: Python<'py>) -> Option<Vec<Bound<'py, PyArray2<f64>>>> {
        self.inner
            .w_k
            .as_ref()
            .map(|vec| vec.iter().map(|m| m.clone().into_pyarray(py)).collect())
    }

    #[getter]
    fn w_i<'py>(&self, py: Python<'py>) -> Option<Vec<Bound<'py, PyArray2<f64>>>> {
        self.inner
            .w_i
            .as_ref()
            .map(|vec| vec.iter().map(|m| m.clone().into_pyarray(py)).collect())
    }

    #[getter]
    fn bf_array<'py>(&self, py: Python<'py>) -> Option<Vec<Bound<'py, PyArray2<f64>>>> {
        self.inner
            .bf_array
            .as_ref()
            .map(|vec| vec.iter().map(|m| m.clone().into_pyarray(py)).collect())
    }

    #[getter]
    fn snr_emission<'py>(&self, py: Python<'py>) -> PyResult<Option<Vec<Bound<'py, PyDict>>>> {
        snr_to_py(py, &self.inner.snr_emission)
    }

    #[getter]
    fn snr_excitation<'py>(&self, py: Python<'py>) -> PyResult<Option<Vec<Bound<'py, PyDict>>>> {
        snr_to_py(py, &self.inner.snr_excitation)
    }

    fn __repr__(&self) -> String {
        format!(
            "Npls(n_components={}, a={}, fitted={})",
            self.inner.n_components,
            self.inner.a,
            self.inner.bf_array.is_some()
        )
    }
}

/// Convert SNR results into a list of Python dicts.
fn snr_to_py<'py>(
    py: Python<'py>,
    snr: &Option<Vec<npls1::SnrResponse>>,
) -> PyResult<Option<Vec<Bound<'py, PyDict>>>> {
    match snr {
        None => Ok(None),
        Some(list) => {
            let mut out = Vec::with_capacity(list.len());
            for comp in list {
                let d = PyDict::new(py);
                for (k, v) in comp.iter() {
                    d.set_item(k, v.clone())?;
                }
                out.push(d);
            }
            Ok(Some(out))
        }
    }
}

/// The Python module `npls1_python`.
#[pymodule]
fn npls1_python(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyNpls>()?;
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    Ok(())
}