onnx-export-rs 0.1.1

use ndarray::Array1;

/// Inverse-link activation applied to a generalized linear model's linear
/// predictor `X · coefficients + intercept`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LinkFunction {
    /// `g⁻¹(η) = η`. Exported as a bare `Gemm`, matching plain linear
    /// regression.
    Identity,
    /// `g⁻¹(η) = exp(η)`. The canonical link for Poisson/Gamma/Tweedie
    /// regression; exported as `Gemm` followed by `Exp`.
    Log,
    /// `g⁻¹(η) = 1 / (1 + exp(-η))`. Exported as `Gemm` followed by `Sigmoid`,
    /// matching binary logistic regression.
    Logit,
}

/// A generalized linear model: a linear predictor followed by an inverse-link
/// activation. Captures Poisson/Gamma/Tweedie regressors whose only difference
/// from linear regression at inference time is the output activation.
///
/// ```
/// use ndarray::array;
/// use onnx_export_rs::canonical::{GeneralizedLinearModel, LinkFunction};
/// let model = GeneralizedLinearModel::new(array![0.3, -0.1], 0.5, LinkFunction::Log);
/// assert_eq!(model.n_features(), 2);
/// ```
#[derive(Clone, Debug, PartialEq)]
pub struct GeneralizedLinearModel {
    /// One coefficient per input feature.
    pub coefficients: Array1<f64>,
    /// Scalar bias added to the linear predictor.
    pub intercept: f64,
    /// Inverse-link activation applied to the linear predictor.
    pub link: LinkFunction,
}

impl GeneralizedLinearModel {
    /// Constructs a generalized linear model.
    #[must_use]
    pub const fn new(coefficients: Array1<f64>, intercept: f64, link: LinkFunction) -> Self {
        Self {
            coefficients,
            intercept,
            link,
        }
    }

    /// Returns the required number of input features.
    #[must_use]
    pub fn n_features(&self) -> usize {
        self.coefficients.len()
    }
}