RustyStats 🦀📊

High-performance Generalized Linear Models with a Rust backend and Python API

Codebase Documentation: pricingfrontier.github.io/rustystats/

Features

• Dict-First API - Programmatic model building for automated workflows and agents
• Fast - Parallel Rust backend for high-throughput fitting
• Memory Efficient - Low memory footprint at scale
• Stable - Step-halving IRLS and warm starts to aid convergence
• Splines - B-splines and natural splines with auto-tuned smoothing and monotonicity
• Target Encoding - Ordered target encoding for high-cardinality categoricals
• Regularization - Ridge, Lasso, and Elastic Net via coordinate descent
• Lasso Credibility - Shrink toward a prior model instead of zero (CAS Monograph 13)
• Out-of-Core - Fit from batches within a memory budget, bitwise identical to the in-memory fit
• Validation - Design matrix checks with fix suggestions before fitting
• Complete - 8 families, robust SEs, full diagnostics, VIF, partial dependence
• Minimal - Only numpy and polars 2 required

Installation

uv add rustystats

Quick Start

import rustystats as rs
import polars as pl

# Load data
data = pl.read_parquet("insurance.parquet")

# Fit a Poisson GLM for claim frequency
result = rs.glm_dict(
    response="ClaimCount",
    terms={
        "VehAge": {"type": "linear"},
        "VehPower": {"type": "linear"},
        "Area": {"type": "categorical"},
        "Region": {"type": "categorical"},
    },
    data=data,
    family="poisson",
    exposure="Exposure",
).fit()

# View results
print(result.summary())

Families & Links

Dict-First API

result = rs.glm_dict(
    response="ClaimCount",
    terms={
        "VehAge": {"type": "bs", "monotonicity": "increasing"},  # Monotonic (auto-tuned)
        "DrivAge": {"type": "bs"},                               # Penalized smooth (default)
        "Income": {"type": "bs", "df": 5},                       # Fixed 5 df
        "BonusMalus": {"type": "linear", "monotonicity": "increasing"},  # Constrained coefficient
        "Region": {"type": "categorical"},
        "Brand": {"type": "target_encoding"},
        "Age2": {"type": "expression", "expr": "DrivAge**2"},
    },
    interactions=[
        {
            "VehAge": {"type": "linear"}, 
            "Region": {"type": "categorical"}, 
            "include_main": True
        },
    ],
    data=data,
    family="poisson",
    exposure="Exposure",
    seed=42,
).fit(regularization="elastic_net")

Multinomial Choice

Use multinomial_dict for mutually exclusive class outcomes such as insurance product-tier conversion.

result = rs.multinomial_dict(
    response="PurchasedTier",
    terms={
        "DriverAge": {"type": "bs"},
        "VehicleValue": {"type": "linear"},
        "Channel": {"type": "categorical"},
    },
    alternative_terms={
        "price": {
            "columns": {
                "basic": "price_basic",
                "standard": "price_standard",
                "premium": "price_premium",
            },
            "coefficient": "generic",
            "transform": "log",
        }
    },
    data=quotes,
    classes=["none", "basic", "standard", "premium"],
    reference="none",
).fit()

probs = result.predict_proba(new_quotes)
mix = result.tier_mix(new_quotes)
scenario = result.scenario(new_quotes, changes={"price_premium": 1.03})

See examples/tier_conversion_multinomial.py for a full train/holdout workflow.

Term Types

Interactions

Each interaction is a dict with variable specs. Use include_main to also add main effects.

interactions=[
    # Standard interaction: product terms (main effects + interaction)
    {
        "DrivAge": {"type": "bs", "df": 5}, 
        "Brand": {"type": "target_encoding"},
        "include_main": True
    },
    # Categorical × continuous (interaction only)
    {
        "VehAge": {"type": "linear"}, 
        "Region": {"type": "categorical"}, 
        "include_main": False
    },
    # TE interaction: combined target encoding TE(Brand:Region)
    {
        "Brand": {"type": "categorical"},
        "Region": {"type": "categorical"},
        "target_encoding": True,
        "prior_weight": 1.0,  # optional
    },
    # FE interaction: combined frequency encoding FE(Brand:Region)
    {
        "Brand": {"type": "categorical"},
        "Region": {"type": "categorical"},
        "frequency_encoding": True,
    },
]

Splines

# Default: penalized smooth with automatic tuning via GCV
result = rs.glm_dict(
    response="ClaimNb",
    terms={
        "Age": {"type": "bs"},           # B-spline (auto-tuned)
        "VehPower": {"type": "ns"},      # Natural spline (auto-tuned)
        "Region": {"type": "categorical"},
    },
    data=data, family="poisson", exposure="Exposure",
).fit()

# Fixed degrees of freedom (no penalty)
result = rs.glm_dict(
    response="ClaimNb",
    terms={
        "Age": {"type": "bs", "df": 5},       # Fixed 5 df
        "VehPower": {"type": "ns", "df": 4},  # Fixed 4 df
        "Region": {"type": "categorical"},
    },
    data=data, family="poisson", exposure="Exposure",
).fit()

Spline parameters:

• No parameters → penalized smooth with automatic tuning (k=10)
• df=5 → fixed 5 degrees of freedom
• k=15 → penalized smooth with 15 basis functions
• knots=[2.0, 5.0, 8.0] → explicit interior knot positions (mutually exclusive with df/k)
• boundary_knots=(0.0, 10.0) → custom boundary knots (optional, defaults to data range)
• monotonicity="increasing" or "decreasing" → constrained effect (bs only)

When to use each type:

• B-splines (bs): Standard choice, more flexible at boundaries, supports monotonicity
• Natural splines (ns): Better extrapolation, linear beyond boundaries

Monotonic Splines

Constrain the fitted curve to be monotonically increasing or decreasing. Useful when business rules require a monotonic relationship.

# Monotonically increasing effect (e.g., age → risk)
result = rs.glm_dict(
    response="ClaimNb",
    terms={
        "Age": {"type": "bs", "monotonicity": "increasing"},
        "Region": {"type": "categorical"},
    },
    data=data, family="poisson", exposure="Exposure",
).fit()

# Monotonically decreasing effect (e.g., vehicle value with age)
result = rs.glm_dict(
    response="ClaimAmt",
    terms={"VehAge": {"type": "bs", "df": 4, "monotonicity": "decreasing"}},
    data=data, family="gamma",
).fit()

Coefficient Constraints

Constrain coefficient signs using monotonicity on linear and expression terms.

result = rs.glm_dict(
    response="y",
    terms={
        "age": {"type": "linear", "monotonicity": "increasing"},  # β ≥ 0
        "age2": {"type": "expression", "expr": "age ** 2", "monotonicity": "decreasing"},  # β ≤ 0
        "income": {"type": "linear"},
    },
    data=data, family="poisson",
).fit()

Target Encoding

Ordered target encoding for high-cardinality categoricals.

# Dict API
result = rs.glm_dict(
    response="ClaimNb",
    terms={
        "Brand": {"type": "target_encoding"},
        "Model": {"type": "target_encoding", "prior_weight": 2.0},
        "Age": {"type": "linear"},
        "Region": {"type": "categorical"},
    },
    data=data, family="poisson", exposure="Exposure",
).fit()

# Sklearn-style API
encoder = rs.TargetEncoder(prior_weight=1.0, n_permutations=4)
train_encoded = encoder.fit_transform(train_categories, train_target)
test_encoded = encoder.transform(test_categories)

Key benefits:

• No target leakage: Ordered target statistics
• Regularization: Prior weight controls shrinkage toward global mean
• High-cardinality: Single column instead of thousands of dummies
• Exposure-aware: For frequency models, pass exposure="Exposure" so the encoder weights by claim rate (ClaimCount/Exposure) instead of raw counts. Exposure weighting comes only from exposure=; an offset= is a link-scale adjustment and never feeds the encoder.
• Interactions: Use target_encoding: True in interactions to encode variable combinations

Expression Terms

result = rs.glm_dict(
    response="y",
    terms={
        "age": {"type": "linear"},
        "age2": {"type": "expression", "expr": "age ** 2"},
        "age3": {"type": "expression", "expr": "age ** 3"},
        "income_k": {"type": "expression", "expr": "income / 1000"},
        "power_to_weight": {"type": "expression", "expr": "power / weight"},
    },
    data=data, family="gaussian",
).fit()

Supported operations: +, -, *, /, ** (power). An expression is one operation on columns or numbers. For anything else, such as a log, add the column to your DataFrame first.

Regularization

CV-Based Regularization

# Just specify regularization type - cv=5 is automatic
result = rs.glm_dict(
    response="y",
    terms={"x1": {"type": "linear"}, "x2": {"type": "linear"}, "cat": {"type": "categorical"}},
    data=data,
    family="poisson",
).fit(regularization="ridge")  # "ridge", "lasso", or "elastic_net"

print(f"Selected alpha: {result.alpha}")
print(f"CV deviance: {result.cv_deviance}")

Options:

• regularization: "ridge" (L2), "lasso" (L1), or "elastic_net" (mix)
• selection: "min" (best fit) or "1se" (more conservative, default: "min")
• cv: Number of folds (default: 5)
• standardize: Internally standardize penalized columns before the penalty acts, reporting original-scale coefficients (default: True; set False to penalize on the raw coefficient scale)

Explicit Alpha

# Skip CV, use specific alpha
result = rs.glm_dict(response="y", terms={"x1": {"type": "linear"}, "x2": {"type": "linear"}}, data=data).fit(alpha=0.1, l1_ratio=0.0)  # Ridge
result = rs.glm_dict(response="y", terms={"x1": {"type": "linear"}, "x2": {"type": "linear"}}, data=data).fit(alpha=0.1, l1_ratio=1.0)  # Lasso
result = rs.glm_dict(response="y", terms={"x1": {"type": "linear"}, "x2": {"type": "linear"}}, data=data).fit(alpha=0.1, l1_ratio=0.5)  # Elastic Net

Lasso Credibility

Shrink model coefficients toward a prior model (complement of credibility) instead of toward zero. Based on the methodology in CAS Monograph 13 (Holmes & Casotto, 2025).

When lasso zeroes a coefficient, the prediction for that term falls back to the complement rather than vanishing, so regularized models can be used directly as rating plans.

# 1. Fit a countrywide (prior) model
cw_result = rs.glm_dict(
    response="ClaimCount",
    terms={"VehAge": {"type": "bs"}, "DrivAge": {"type": "bs"}},
    data=countrywide_data,
    family="poisson",
    exposure="Exposure",
).fit()

# 2. Fit a state model with lasso, shrinking toward countrywide rates
state_result = rs.glm_dict(
    response="ClaimCount",
    terms={
        "VehAge": {"type": "bs"},
        "DrivAge": {"type": "bs"},
        "Region": {"type": "categorical"},
    },
    data=state_data,
    family="poisson",
    exposure="Exposure",
    complement="countrywide_rate",  # Column with prior rates (response scale)
).fit(regularization="lasso")

# 3. Inspect which terms the data supports vs. trusts the complement
print(state_result.summary())            # Shows "Lasso Credibility Results"
print(state_result.credibility_summary())  # Deviation from complement per term

Complement sources:

• str: column name in the DataFrame (rates for log-link, probabilities for logit)
• np.ndarray: array of prior values on the response scale
• GLMModel: fitted model; predictions are computed automatically

Works with all families/links, splines, categoricals, interactions, and target encoding.

Out-of-Core Fitting

Fit data too large for memory from a source that can be re-read in batches. Memory is bounded by a budget you set, and the result is bitwise identical to an in-memory fit.

import polars as pl
import rustystats as rs

path = "training.parquet"

def open_batches(columns):
    # Every call must yield the same rows in the same order.
    yield from pl.scan_parquet(path).select(columns).collect_batches()

result = rs.glm_dict(
    response="ClaimCount",
    terms={"Age": {"type": "bs", "df": 5}, "Region": {"type": "categorical"}},
    data=rs.BatchSource(open_batches, pl.read_parquet_schema(path)),
    family="poisson",
    exposure="Exposure",
).fit(out_of_core=rs.OutOfCoreOptions(max_bytes=2 * 1024**3))

See Out-of-Core Fitting for the options and what a result serves.

Design Matrix Validation

# Check for issues before fitting
model = rs.glm_dict(
    response="y",
    terms={"x": {"type": "ns", "df": 4}, "cat": {"type": "categorical"}},
    data=data, family="poisson",
)
results = model.validate()  # Prints diagnostics

if not results['valid']:
    print("Issues:", results['suggestions'])

# Validation runs automatically on fit failure with suggested fixes

Checks performed:

• Rank deficiency (linearly dependent columns)
• High multicollinearity (condition number)
• Zero variance columns
• NaN/Inf values
• Highly correlated column pairs (>0.999)

Model Diagnostics

# Compute all diagnostics at once
diagnostics = result.diagnostics(
    train_data=data,
    categorical_factors=["Region", "VehBrand", "Area"],  # Including non-fitted
    continuous_factors=["Age", "Income", "VehPower"],    # Including non-fitted
)

# Export as compact JSON (optimized for LLM consumption)
json_str = diagnostics.to_json()

# Pre-fit data exploration (no model needed)
exploration = rs.explore_data(
    data=data,
    response="ClaimNb",
    categorical_factors=["Region", "VehBrand", "Area"],
    continuous_factors=["Age", "VehPower", "Income"],
    exposure="Exposure",
    family="poisson",
    detect_interactions=True,
)

Diagnostic Features:

• Calibration: Overall A/E ratio, calibration by decile with CIs, Hosmer-Lemeshow test
• Discrimination: Gini coefficient, AUC, KS statistic, lift metrics
• Factor Diagnostics: A/E by level/bin for ALL factors (fitted and non-fitted)
• VIF/Multicollinearity: Variance inflation factors for design matrix columns
• Partial Dependence: Effect plots with shape detection and recommendations
• Overfitting Detection: Compare train vs test metrics when test data provided
• Interaction Detection: Greedy residual-based detection of potential interactions
• Warnings: Auto-generated alerts for high dispersion, poor calibration, missing factors
• Base Model Comparison: Compare new model against existing/benchmark predictions

Diagnostics JSON Shape

The diagnostics.to_json() output includes:

{
  "model_summary": {
    "formula": "...", "family": "poisson", "link": "log",
    "n_obs": 2000, "n_params": 6, "df_resid": 1994,
    "converged": true, "iterations": 5,
    "scale": 1.0,
    "scale_pearson": 1.0148,
    "null_deviance": 1408.77,
    "robust_se_type": "HC1"
  },
  "train_test": {
    "train": {
      "n_obs": 2000, "deviance": 2118.42,
      "log_likelihood": -1059.21,
      "aic": 2130.42, "bic": 2162.70,
      "gini": 0.3241, "auc": 0.6621, "ae_ratio": 1.0
    }
  },
  "coefficient_summary": [
    {
      "feature": "Age", "estimate": 0.00996,
      "std_error": 0.00339, "z_value": 2.941,
      "p_value": 0.0033, "significant": true,
      "conf_int": [0.003318, 0.01659],
      "relativity": 1.01, "relativity_ci": [1.0033, 1.0167],
      "robust_std_error": 0.003378, "robust_z_value": 2.947,
      "robust_p_value": 0.0032, "robust_significant": true
    }
  ]
}

• BIC (train_test.train.bic): Bayesian Information Criterion alongside AIC
• Scale (model_summary.scale): Deviance-based dispersion parameter
• Scale Pearson (model_summary.scale_pearson): Pearson-based dispersion estimate
• Null Deviance (model_summary.null_deviance): Intercept-only model deviance
• Confidence Intervals (coefficient_summary[].conf_int): 95% CI [lower, upper]
• Robust SEs (coefficient_summary[].robust_*): HC1 sandwich estimators for each coefficient
• Robust SE Type (model_summary.robust_se_type): Present only when robust SEs were computed

Robust SE fields are null when store_design_matrix=False (lean mode) or for deserialized models.

Comparing Against a Base Model

Compare your new model against predictions from an existing model (e.g., current production model):

# Add base model predictions to your data
data = data.with_columns(pl.lit(old_model_predictions).alias("base_pred"))

# Run diagnostics with base_predictions
diagnostics = result.diagnostics(
    train_data=data,
    categorical_factors=["Region", "VehBrand"],
    continuous_factors=["Age", "VehPower"],
    base_predictions="base_pred",  # Column name with base model predictions
)

# Access comparison results
bc = diagnostics.base_predictions_comparison

# Side-by-side metrics
print(f"Model loss: {bc.model_metrics.loss}, Base loss: {bc.base_metrics.loss}")
print(f"Model Gini: {bc.model_metrics.gini}, Base Gini: {bc.base_metrics.gini}")

# Improvement metrics (positive = new model is better)
print(f"Loss improvement: {bc.loss_improvement_pct}%")
print(f"Gini improvement: {bc.gini_improvement}")
print(f"AUC improvement: {bc.auc_improvement}")

# Decile analysis sorted by model/base prediction ratio
for d in bc.model_vs_base_deciles:
    print(f"Decile {d.decile}: actual={d.actual:.4f}, "
          f"model={d.model_predicted:.4f}, base={d.base_predicted:.4f}")

The comparison includes:

• Side-by-side metrics: Loss (mean deviance), Gini, AUC, A/E ratio for both models
• Improvement metrics: loss_improvement_pct, gini_improvement, auc_improvement
• Decile analysis: Data sorted by model/base ratio, showing where the new model diverges
• Calibration comparison: Count of deciles where each model has better A/E

Calibration Primitives

Explicit calibration tools for assessing and adjusting model balance, kept separate from the GLM coefficients so the underlying fit stays untouched.

# Standalone summary on arrays (overall A/E, per-bin, optional per-factor)
summary = rs.calibration_summary(
    y, mu,
    exposure=exposure,
    weights=weights,         # optional; weighted Σwy/Σwμ
    by={"Region": region},   # optional per-factor breakdown
    n_bins=10,
    ranking="auto",          # rate-rank when exposure is present
    min_exposure=10.0,       # flag low-exposure cells as suppressed
)

# From a fitted GLM (response/exposure resolved automatically)
result.calibration_summary(data, by="Region")

# Multiplicative or monotone calibration objects (opt-in, serialized separately)
cal = result.fit_calibration(holdout, method="global")     # GlobalCalibration
iso = result.fit_calibration(holdout, method="isotonic")   # IsotonicCalibration
calibrated_pred = cal.predict(result.predict(new_data))

# Log-link intercept relevel. Same factor c = Σ(w·y)/Σ(w·μ); updates only the
# intercept. Every other coefficient is bit-identical, relativities preserved.
releveled = result.relevel(holdout)
assert all(releveled.params[1:] == result.params[1:])

Calibration is never applied silently to result.predict(). Calibration objects are separate, serializable (to_dict/from_dict), and not folded into GLM coefficients. Fitting calibration on the same rows used to fit the model overstates calibration quality, so prefer a held-out fold.

Per-Prediction Contributions

Decompose each row's prediction into per-term contributions for trace explainability:

result = rs.glm_dict(
    response="sale_flag",
    terms={"diff_to_market": {"type": "ns", "df": 10}},
    data=train,
    family="binomial",
).fit()

rows = result.predict_contributions(new_data)
print(rows[0])
# {
#   "family": "binomial", "link": "logit",
#   "output_space": "linear_predictor", "prediction_space": "response",
#   "base_value": 0.0417,
#   "sum_contributions": -1.4280,
#   "prediction_from_contributions": -1.3863,   # eta
#   "prediction_value": 0.2000,                  # mu = inverse_link(eta)
#   "contributions": [
#       {"term": "diff_to_market", "term_type": "ns",
#        "feature_value": -25.0, "contribution": -1.4280, "rank": 1}
#   ]
# }

Key properties:

• base_value + sum(contributions) == linear predictor (validated to 1e-9 by default)
• inverse_link(linear predictor) == predict() (also validated)
• Spline bases, categorical dummies, target/frequency encoding columns, and interaction tensor products are grouped back to their source term; the ladder shows factor-level rows, not basis-level rows
• Offset is an explicit row (term_type="offset", contribution = log(Exposure) for log-link)
• For complement-of-credibility models, base_value is per-row = link(complement[row]), and the intercept appears as a contribution row representing the deviation

Options:

• group_terms=False: expand multi-column terms into one row per design column
• include_design_columns=True: keep grouped rows but attach a per-column breakdown
• return_format="dataframe": long-format pl.DataFrame (faster for batch scoring)
• validate=False, atol, rtol: control the additivity check

Model Serialization

Save and load fitted models for later use:

# Fit and save
model_bytes = result.to_bytes()

with open("model.bin", "wb") as f:
    f.write(model_bytes)

# Load later
with open("model.bin", "rb") as f:
    loaded = rs.GLMModel.from_bytes(f.read())

# Predict with loaded model
predictions = loaded.predict(new_data)

What's preserved:

• Coefficients and feature names
• Categorical encoding levels
• Spline knot positions
• Target encoding statistics
• Formula, family, link function
• Complement of credibility specification

Compact storage: Only prediction-essential state is stored (~KB, not MB).

Model Export (PMML & ONNX)

Export fitted models to standard formats for deployment, with no extra dependencies. PMML uses stdlib XML; ONNX protobuf serialization is implemented from scratch in Rust.

PMML

# Export to PMML 4.4 XML
pmml_xml = result.to_pmml()
result.to_pmml(path="model.pmml")

# Load & predict (consumer side)
# uv add pypmml
from pypmml import Model
pmml_model = Model.fromFile("model.pmml")

new_data = pl.DataFrame({"VehAge": [3, 5, 1], "Area": ["C", "A", "B"]})
preds = pmml_model.predict(new_data.to_dict(as_series=False))

ONNX

# Export: "scoring" requires a pre-built design matrix, "full" embeds preprocessing
result.to_onnx(path="model.onnx", mode="scoring")
result.to_onnx(path="model_full.onnx", mode="full")

# Predict (consumer side)
# uv add onnxruntime
import onnxruntime as ort
session = ort.InferenceSession("model_full.onnx")
preds = session.run(None, {"input": raw_features})[0]

For MultinomialModel, PMML/ONNX export is currently Level-1 scoring only: the consumer supplies the pre-built shared design matrix without the intercept column. Multinomial mode="full", target encoding, alternative terms, availability masks, and class-specific offsets fail closed with validation errors.

Dependencies

Rust

• ndarray, nalgebra - Linear algebra
• rayon - Parallel iterators (multi-threading)
• statrs - Statistical distributions
• pyo3 - Python bindings

Python

• numpy - Array operations (required)
• polars 2.x - DataFrame support (required)

License

AGPL-3.0