onnx-export-rs 0.1.1

Export canonical Rust machine-learning models to ONNX
Documentation
//! Adapters for `smartcore` fitted models using its standard dense matrix.

use ndarray::{Array1, Array2};
use smartcore::decomposition::pca::PCA;
use smartcore::decomposition::svd::SVD;
use smartcore::linalg::basic::arrays::Array;
use smartcore::linalg::basic::matrix::DenseMatrix;
use smartcore::linear::elastic_net::ElasticNet;
use smartcore::linear::lasso::Lasso;
use smartcore::linear::linear_regression::LinearRegression;
use smartcore::linear::logistic_regression::LogisticRegression;
use smartcore::linear::ridge_regression::RidgeRegression;

use crate::canonical::{AffineModel, LinearModelWeights, LogisticModelWeights};
use crate::Result;

/// Standard `smartcore` fitted linear-regression type supported by this adapter.
pub type DenseLinearRegression = LinearRegression<f64, f64, DenseMatrix<f64>, Vec<f64>>;
/// Standard `smartcore` fitted logistic-regression type supported by this adapter.
pub type DenseLogisticRegression = LogisticRegression<f64, i32, DenseMatrix<f64>, Vec<i32>>;
/// Standard `smartcore` fitted ridge-regression type supported by this adapter.
pub type DenseRidgeRegression = RidgeRegression<f64, f64, DenseMatrix<f64>, Vec<f64>>;
/// Standard `smartcore` fitted Lasso type supported by this adapter.
pub type DenseLasso = Lasso<f64, f64, DenseMatrix<f64>, Vec<f64>>;
/// Standard `smartcore` fitted Elastic Net type supported by this adapter.
pub type DenseElasticNet = ElasticNet<f64, f64, DenseMatrix<f64>, Vec<f64>>;
/// Standard `smartcore` fitted PCA transform supported by this adapter.
pub type DensePca = PCA<f64, DenseMatrix<f64>>;
/// Standard `smartcore` fitted truncated-SVD transform supported by this adapter.
pub type DenseSvd = SVD<f64, DenseMatrix<f64>>;

fn ndarray_matrix(matrix: &DenseMatrix<f64>) -> Array2<f64> {
    let (rows, columns) = matrix.shape();
    Array2::from_shape_fn((rows, columns), |(row, column)| *matrix.get((row, column)))
}

fn column_weights(coefficients: &DenseMatrix<f64>, intercept: f64) -> LinearModelWeights {
    let (rows, columns) = coefficients.shape();
    debug_assert_eq!(columns, 1);
    LinearModelWeights::new(
        Array1::from_iter((0..rows).map(|row| *coefficients.get((row, 0)))),
        intercept,
    )
}

/// Extracts canonical weights from a fitted `smartcore` linear model.
#[must_use]
pub fn linear_weights(model: &DenseLinearRegression) -> LinearModelWeights {
    column_weights(model.coefficients(), *model.intercept())
}

/// Extracts canonical weights from a fitted `smartcore` logistic model.
///
/// The returned rows follow the order reported by [`logistic_classes`].
///
/// # Errors
///
/// Returns an error if `smartcore` exposes an inconsistent fitted shape.
pub fn logistic_weights(model: &DenseLogisticRegression) -> Result<LogisticModelWeights> {
    let coefficients = model.coefficients();
    let intercept = model.intercept();
    let (rows, columns) = coefficients.shape();
    let coefficient_values =
        (0..rows).flat_map(|row| (0..columns).map(move |column| *coefficients.get((row, column))));
    LogisticModelWeights::new(
        Array2::from_shape_vec((rows, columns), coefficient_values.collect())
            .expect("matrix shape and iteration count agree"),
        Array1::from_iter((0..rows).map(|row| *intercept.get((row, 0)))),
        model.classes().len(),
    )
}

/// Returns class labels in the score-column order used by [`logistic_weights`].
#[must_use]
pub fn logistic_classes(model: &DenseLogisticRegression) -> &[i32] {
    model.classes()
}

/// Extracts canonical weights from fitted ridge regression.
#[must_use]
pub fn ridge_weights(model: &DenseRidgeRegression) -> LinearModelWeights {
    column_weights(model.coefficients(), *model.intercept())
}

/// Extracts canonical weights from a fitted Lasso model.
#[must_use]
pub fn lasso_weights(model: &DenseLasso) -> LinearModelWeights {
    column_weights(model.coefficients(), *model.intercept())
}

/// Extracts canonical weights from a fitted Elastic Net model.
#[must_use]
pub fn elastic_net_weights(model: &DenseElasticNet) -> LinearModelWeights {
    column_weights(model.coefficients(), *model.intercept())
}

/// Extracts the affine projection performed by fitted PCA.
///
/// # Errors
///
/// Returns an error if SmartCore rejects a zero-row probe or exposes an
/// inconsistent projection.
pub fn pca_transform(model: &DensePca) -> Result<AffineModel> {
    let matrix = ndarray_matrix(model.components());
    let zero = DenseMatrix::new(1, matrix.nrows(), vec![0.0; matrix.nrows()], false)
        .map_err(|error| crate::Error::InvalidModel(error.to_string()))?;
    let bias_matrix = model
        .transform(&zero)
        .map_err(|error| crate::Error::InvalidModel(error.to_string()))?;
    let bias = Array1::from_iter((0..matrix.ncols()).map(|column| *bias_matrix.get((0, column))));
    AffineModel::new(matrix, bias)
}

/// Extracts the matrix multiplication performed by fitted truncated SVD.
///
/// # Errors
///
/// Returns an error if the exposed component matrix is empty or non-finite.
pub fn svd_transform(model: &DenseSvd) -> Result<AffineModel> {
    let matrix = ndarray_matrix(model.components());
    let bias = Array1::zeros(matrix.ncols());
    AffineModel::new(matrix, bias)
}