onnx-export-rs 0.1.1

//! Adapters for `linfa` fitted models.

use ndarray::{Array1, Array2, ArrayView2, Axis};

use crate::canonical::{
    flatten_tree, AffineModel, CentroidModel, GaussianMixture, LinearModelWeights,
    LogisticModelWeights, RecursiveNode, TreeStructure,
};
use crate::{Error, Result};

/// Extracts canonical weights from a fitted `linfa-linear` model.
#[must_use]
pub fn linear_weights(model: &linfa_linear::FittedLinearRegression<f64>) -> LinearModelWeights {
    LinearModelWeights::new(model.params().clone(), model.intercept())
}

/// Extracts canonical weights from a fitted binary `linfa-logistic` model.
///
/// Its single output is the probability of the model's positive label.
///
/// # Errors
///
/// Returns an error if the fitted parameter shape is inconsistent.
pub fn binary_logistic_weights<C: PartialOrd + Clone>(
    model: &linfa_logistic::FittedLogisticRegression<f64, C>,
) -> Result<LogisticModelWeights> {
    LogisticModelWeights::new(
        model.params().clone().insert_axis(Axis(0)),
        Array1::from_vec(vec![model.intercept()]),
        2,
    )
}

/// Extracts canonical weights from a fitted multinomial `linfa-logistic` model.
///
/// # Errors
///
/// Returns an error if the fitted parameter shape is inconsistent.
pub fn multiclass_logistic_weights<C: PartialOrd + Clone>(
    model: &linfa_logistic::MultiFittedLogisticRegression<f64, C>,
) -> Result<LogisticModelWeights> {
    LogisticModelWeights::new(
        model.params().t().to_owned(),
        model.intercept().clone(),
        model.classes().len(),
    )
}

/// Returns multinomial class labels in exported probability-column order.
#[must_use]
pub fn multiclass_logistic_classes<C: PartialOrd + Clone>(
    model: &linfa_logistic::MultiFittedLogisticRegression<f64, C>,
) -> &[C] {
    model.classes()
}

/// Converts a fitted Linfa classification tree into canonical class scores.
///
/// Each leaf emits a one-hot score in the supplied class order. The class list
/// must contain every label fitted by the tree and should use the order desired
/// for exported output columns.
///
/// # Errors
///
/// Returns an error for missing children, a leaf label absent from `classes`,
/// or an empty class list.
pub fn classification_tree<L>(
    model: &linfa_trees::DecisionTree<f64, L>,
    classes: &[L],
) -> Result<TreeStructure>
where
    L: linfa::Label + std::fmt::Debug,
{
    if classes.is_empty() {
        return Err(Error::InvalidModel(
            "classification tree needs classes".into(),
        ));
    }

    fn convert<L>(node: &linfa_trees::TreeNode<f64, L>, classes: &[L]) -> Result<RecursiveNode>
    where
        L: linfa::Label + std::fmt::Debug,
    {
        if let Some(prediction) = node.prediction() {
            let class_index = classes
                .iter()
                .position(|class| *class == prediction)
                .ok_or_else(|| Error::InvalidModel("tree leaf has an unknown class".into()))?;
            let mut scores = vec![0.0; classes.len()];
            scores[class_index] = 1.0;
            return Ok(RecursiveNode::Leaf(scores));
        }
        let children = node.children();
        let left = children
            .first()
            .and_then(|child| child.as_deref())
            .ok_or_else(|| Error::InvalidModel("tree branch is missing its left child".into()))?;
        let right = children
            .get(1)
            .and_then(|child| child.as_deref())
            .ok_or_else(|| Error::InvalidModel("tree branch is missing its right child".into()))?;
        let (feature, threshold, _) = node.split();
        Ok(RecursiveNode::BranchLessThan {
            feature,
            threshold,
            left: Box::new(convert(left, classes)?),
            right: Box::new(convert(right, classes)?),
        })
    }

    flatten_tree(&convert(model.root_node(), classes)?)
}

/// Extracts canonical weights from a fitted Linfa Elastic Net, Lasso, or
/// ridge model.
#[must_use]
pub fn elastic_net_weights(model: &linfa_elasticnet::ElasticNet<f64>) -> LinearModelWeights {
    LinearModelWeights::new(model.hyperplane().clone(), model.intercept())
}

/// Extracts a multi-output affine transform from fitted multi-task Elastic
/// Net, multi-task Lasso, or multi-task ridge.
///
/// # Errors
///
/// Returns an error if the fitted arrays are inconsistent.
pub fn multi_task_elastic_net(
    model: &linfa_elasticnet::MultiTaskElasticNet<f64>,
) -> Result<AffineModel> {
    AffineModel::new(model.hyperplane().clone(), model.intercept().clone())
}

/// Converts fitted Linfa PCA into an affine projection.
///
/// # Errors
///
/// Returns an error if the fitted arrays are inconsistent.
pub fn pca(model: &linfa_reduction::Pca<f64>) -> Result<AffineModel> {
    let matrix = model.components().t().to_owned();
    let bias = -model.mean().dot(&matrix);
    AffineModel::new(matrix, bias)
}

/// Converts fitted Linfa L2 k-means centroids.
///
/// # Errors
///
/// Returns an error for invalid fitted centroids.
pub fn kmeans(
    model: &linfa_clustering::KMeans<f64, linfa_nn::distance::L2Dist>,
) -> Result<CentroidModel> {
    CentroidModel::new(model.centroids().clone())
}

/// Converts a fitted full-covariance Linfa Gaussian mixture.
///
/// # Errors
///
/// Returns an error if a precision matrix is singular or fitted shapes are
/// inconsistent.
pub fn gaussian_mixture(
    model: &linfa_clustering::GaussianMixtureModel<f64>,
) -> Result<GaussianMixture> {
    let features = model.means().ncols();
    let offsets = Array1::from_iter(
        model
            .precisions()
            .outer_iter()
            .zip(model.weights())
            .map(|(precision, &weight)| {
                let determinant = determinant(precision)?;
                if determinant <= 0.0 || weight <= 0.0 {
                    return Err(Error::InvalidModel(
                        "Gaussian mixture precision/weight is not positive".into(),
                    ));
                }
                Ok(weight.ln() + 0.5 * determinant.ln()
                    - 0.5 * features as f64 * std::f64::consts::TAU.ln())
            })
            .collect::<Result<Vec<_>>>()?,
    );
    GaussianMixture::new(model.means().clone(), model.precisions().clone(), offsets)
}

/// Extracts binary logistic probabilities from fitted Linfa FTRL state.
///
/// FTRL prediction is a sigmoid over the learned proximal weights with no
/// intercept, so it reuses the canonical binary-logistic exporter.
pub fn ftrl(model: &linfa_ftrl::Ftrl<f64>) -> Result<LogisticModelWeights> {
    LogisticModelWeights::new(
        model.get_weights().insert_axis(Axis(0)),
        Array1::zeros(1),
        2,
    )
}

/// Extracts the decision function of a fitted *linear-kernel* Linfa SVM.
///
/// The feature count is explicit because Linfa does not expose the separating
/// hyperplane. This function probes the public decision function on basis
/// vectors. Passing a nonlinear-kernel model would not produce an equivalent
/// affine function and is rejected when an interaction check fails.
pub fn linear_svm_score<T>(
    model: &linfa_svm::Svm<f64, T>,
    n_features: usize,
) -> Result<LinearModelWeights> {
    if n_features == 0 {
        return Err(Error::InvalidModel(
            "linear SVM needs input features".into(),
        ));
    }
    let zero = Array1::zeros(n_features);
    let origin = model.weighted_sum(&zero);
    let coefficients = Array1::from_iter((0..n_features).map(|feature| {
        let mut basis = Array1::zeros(n_features);
        basis[feature] = 1.0;
        model.weighted_sum(&basis) - origin
    }));
    // A linear function must be additive. This catches standard polynomial
    // and RBF kernels without requiring access to Linfa's private kernel enum.
    let ones = Array1::ones(n_features);
    let expected = origin + coefficients.sum();
    if (model.weighted_sum(&ones) - expected).abs() > 1e-8 * (1.0 + expected.abs()) {
        return Err(Error::InvalidModel(
            "Linfa SVM does not use a linear kernel".into(),
        ));
    }
    Ok(LinearModelWeights::new(coefficients, origin - model.rho))
}

fn determinant(matrix: ArrayView2<'_, f64>) -> Result<f64> {
    if matrix.nrows() != matrix.ncols() {
        return Err(Error::InvalidModel("precision matrix is not square".into()));
    }
    let mut owned: Array2<f64> = matrix.to_owned();
    let size = owned.nrows();
    let mut result = 1.0;
    for column in 0..size {
        let pivot = (column..size)
            .max_by(|&left, &right| {
                owned[(left, column)]
                    .abs()
                    .total_cmp(&owned[(right, column)].abs())
            })
            .unwrap();
        if owned[(pivot, column)] == 0.0 {
            return Ok(0.0);
        }
        if pivot != column {
            for index in 0..size {
                owned.swap((pivot, index), (column, index));
            }
            result = -result;
        }
        let diagonal = owned[(column, column)];
        result *= diagonal;
        for row in column + 1..size {
            let factor = owned[(row, column)] / diagonal;
            for index in column + 1..size {
                owned[(row, index)] -= factor * owned[(column, index)];
            }
        }
    }
    Ok(result)
}