Binding Crash Course: OLS

This page is the shortest binding walkthrough in the library. OLS uses the crate’s native QR least-squares helpers behind a thin PyO3 boundary, so it is the cleanest starting point for understanding how the Rust extension exposes estimators to Python.

1 Module Export

OLS is exported from the extension module in src/lib.rs:

use crate::estimators::{
    ElasticNet, FixedEffectsOLS, Logit, MEstimator, MultinomialLogit, Poisson,
    SyntheticControl, TwoSLS, KernelBasis, OLS, PcaTransformer,
};

#[pymodule]
fn crabbymetrics(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<OLS>()?;
    Ok(())
}

That one m.add_class::<OLS>()? line is what makes from crabbymetrics import OLS valid in Python.

2 The Stored Rust State

The Rust-side class is tiny:

#[pyclass]
pub struct OLS {
    fit_intercept: bool,
    intercept: f64,
    coef: Option<Array1<f64>>,
    x: Option<Array2<f64>>,
    y: Option<Array1<f64>>,
    sample_weight: Option<Array1<f64>>,
}

This is the whole design in miniature:

  • fit_intercept controls whether a constant column is added
  • intercept and coef store the fitted parameters directly
  • x and y cache the training data for summary() and bootstrap()
  • sample_weight records weights for weighted covariance calculations

The constructor is minimal too:

#[pymethods]
impl OLS {
    #[new]
    fn new() -> Self {
        Self {
            fit_intercept: true,
            intercept: 0.0,
            coef: None,
            x: None,
            y: None,
            sample_weight: None,
        }
    }
}

So from Python, the public constructor is just:

OLS()

3 The Key Imports

The interesting imports for OLS are the ones that define the wrapper pattern:

use crate::utils::{add_intercept, solve_least_squares_vec, to_array1, to_array2};
use ndarray::{Array1, Array2};
use numpy::{PyArray1, PyReadonlyArray1, PyReadonlyArray2};
use pyo3::prelude::*;

Their roles are:

  • PyReadonlyArray*: accept NumPy arrays from Python safely
  • to_array1 / to_array2: copy NumPy inputs into owned ndarray arrays
  • add_intercept: construct the regression design matrix
  • solve_least_squares_vec: solve the dense least-squares problem with QR

There is no iterative optimizer or callback bridge. OLS converts inputs, builds the design, solves least squares, and stores the parameters directly.

4 What fit() Actually Does

The entire fitting path is the simplest possible Rust-to-linfa wrapper:

fn fit(&mut self, x: PyReadonlyArray2<f64>, y: PyReadonlyArray1<f64>) -> PyResult<()> {
    let x = to_array2(&x);
    let y = to_array1(&y);
    if x.nrows() != y.len() {
        return Err(PyValueError::new_err("x rows must match y length"));
    }
    let params =
        fit_linear_params(&x, &y, self.fit_intercept, None).map_err(PyValueError::new_err)?;
    let (intercept, coef) = split_params(&params, self.fit_intercept);
    self.intercept = intercept;
    self.coef = Some(coef);
    self.x = Some(x);
    self.y = Some(y);
    self.sample_weight = None;
    Ok(())
}

The sequence is:

  1. convert Python/NumPy arrays into owned ndarray arrays with to_array1 and to_array2
  2. validate that x and y have matching row counts
  3. add an intercept column when requested
  4. solve the dense least-squares problem through the shared QR helper
  5. split and store the intercept, slopes, and training data on self

This is the most direct example of the wrapper pattern in the repo. The numerical work is native Rust and shared with weighted linear-estimator paths.

5 Prediction Is Just Another Thin Wrapper

predict() follows the same pattern:

fn predict<'py>(
    &self,
    py: Python<'py>,
    x: PyReadonlyArray2<f64>,
) -> PyResult<Bound<'py, PyArray1<f64>>> {
    let coef = self
        .coef
        .as_ref()
        .ok_or_else(|| PyValueError::new_err("OLS model is not fitted"))?;
    let x = to_array2(&x);
    let pred = x.dot(coef) + self.intercept;
    Ok(pyarray1_from_f64(py, &pred))
}

So the boundary mechanics are:

  • pull the fitted coefficient vector out of self.coef
  • convert the new NumPy feature matrix into ndarray
  • evaluate \(x\hat\beta+\hat\alpha\)
  • convert the result back into a Python numpy.ndarray

6 Summary And Variance Estimators

summary() combines the stored native point estimates with covariance calculations. Python can choose vanilla, hc0–hc3, newey_west, or cluster covariance:

let design = if self.fit_intercept {
    add_intercept(x)
} else {
    x.clone()
};
let fitted = design.dot(&params);
let residuals = y - &fitted;
let cov = linear_covariance(
    &design,
    &residuals,
    vcov,
    lags,
    cluster_ids.as_ref(),
    None,
)?;
let se_all = diag_sqrt(&cov)?;

Then it combines:

  • the native point estimates stored in intercept and coef
  • the locally computed covariance from the chosen OLS variance formula

and returns a plain Python dict:

let dict = pyo3::types::PyDict::new(py);
dict.set_item("intercept", intercept)?;
dict.set_item("coef", pyarray1_from_f64(py, &coef))?;
dict.set_item("intercept_se", intercept_se)?;
dict.set_item("coef_se", pyarray1_from_f64(py, &coef_se))?;

This is a common pattern in the library:

  • point estimates and inference are computed in native Rust
  • inference helpers often live in utils.rs
  • Python sees a simple dict/array interface either way

7 Bootstrap Reuses The Same Native Solver

The bootstrap path resamples the stored training rows and calls the same QR least-squares helper for every draw. That is possible because fit() cached x and y on the Rust side during the original fit.

8 Python-Side Mental Model

From Python, all of that machinery collapses down to the usual API:

import numpy as np
from crabbymetrics import OLS

x = np.random.randn(200, 3)
y = 0.3 + x @ np.array([1.0, -2.0, 0.5]) + np.random.randn(200) * 0.1

model = OLS()
model.fit(x, y)

summary = model.summary()
summary_vanilla = model.summary(vcov="vanilla")
summary_hac = model.summary(vcov="newey_west", lags=4)
clusters = np.repeat(np.arange(20, dtype=np.int64), len(y) // 20)
summary_cluster = model.summary(vcov="cluster", clusters=clusters)
pred = model.predict(x[:5])
boot = model.bootstrap(50, seed=1)

The implementation detail to keep in mind is that OLS owns its fitted parameters and delegates only reusable linear algebra and covariance calculations to crate-local helpers. Ridge, FixedEffectsOLS, and TwoSLS use the same covariance vocabulary even though their point estimators differ.

9 Why Start Here

OLS is the first mechanics page because it shows the base pattern with almost no distractions:

  • export a Rust struct with #[pyclass]
  • accept NumPy arrays via PyReadonlyArray*
  • convert them to ndarray
  • call a native numerical helper
  • store the fitted parameters
  • convert predictions and summaries back to Python values

From there, Poisson adds a native optimization problem in Rust, and MEstimator adds the full callback bridge where Rust owns the solver loop but Python supplies the objective and score functions.