RustyStats 🦀📊

Features

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",
 offset="Exposure",
).fit()
# View results
print(result.summary())

Families & Links

Dict-Based 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",
 offset="Exposure",
 seed=42,
).fit(regularization="elastic_net")

Term Types

Interactions

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", offset="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", offset="Exposure",
).fit()

Monotonic Splines

# 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", offset="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

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

# 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", offset="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)

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"},
 "bmi": {"type": "expression", "expr": "weight / (height ** 2)"},
 },
 data=data, family="gaussian",
).fit()

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}")

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

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 helpful suggestions

Results

# Coefficients & Inference
result.params # Coefficients
result.fittedvalues # Predicted means
result.deviance # Model deviance
result.bse() # Standard errors
result.tvalues() # z-statistics
result.pvalues() # P-values
result.conf_int(alpha) # Confidence intervals
# Robust Standard Errors (sandwich estimators)
result.bse_robust("HC1") # Robust SE (HC0, HC1, HC2, HC3)
result.tvalues_robust() # z-stats with robust SE
result.pvalues_robust() # P-values with robust SE
result.conf_int_robust() # Confidence intervals with robust SE
result.cov_robust() # Full robust covariance matrix
# Diagnostics (statsmodels-compatible)
result.resid_response() # Raw residuals (y - μ)
result.resid_pearson() # Pearson residuals
result.resid_deviance() # Deviance residuals
result.resid_working() # Working residuals
result.llf() # Log-likelihood
result.aic() # Akaike Information Criterion
result.bic() # Bayesian Information Criterion
result.null_deviance() # Null model deviance
result.pearson_chi2() # Pearson chi-squared
result.scale() # Dispersion (deviance-based)
result.scale_pearson() # Dispersion (Pearson-based)
result.family # Family name

Model Diagnostics

# Compute all diagnostics at once
diagnostics = result.diagnostics(
 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,
)

Comparing Against a Base 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}")

Model Serialization

# 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)

Model Export (PMML & ONNX)

PMML

# Export to PMML 4.4 XML
pmml_xml = result.to_pmml()
result.to_pmml(path="model.pmml")
# Load & predict (consumer side)
# pip install 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

# Scoring mode (default) — input is pre-built design matrix
onnx_bytes = result.to_onnx(mode="scoring")
result.to_onnx(path="model.onnx", mode="scoring")
# Full mode — input is raw feature values, preprocessing embedded
result.to_onnx(path="model_full.onnx", mode="full")

# Load & predict on a DataFrame with onnxruntime (consumer side)
# pip install onnxruntime
import onnxruntime as ort
import numpy as np
new_data = pl.DataFrame({
 "VehAge": [3, 5, 1],
 "Area": ["C", "A", "B"],
})
# ── Scoring mode: build design matrix from DataFrame ──
session = ort.InferenceSession("model.onnx")
# Columns match model.feature_names (excluding Intercept): [VehAge, Area_B, Area_C]
X = np.column_stack([
 new_data["VehAge"].to_numpy().astype(np.float64),
 (new_data["Area"] == "B").cast(pl.Float64).to_numpy(),
 (new_data["Area"] == "C").cast(pl.Float64).to_numpy(),
])
preds = session.run(None, {"X": X})[0] # shape (3, 1)
# ── Full mode: pass raw values, categoricals as integer codes ──
session = ort.InferenceSession("model_full.onnx")
# Map categorical levels to 0-based codes: A=0, B=1, C=2
level_map = {"A": 0, "B": 1, "C": 2}
raw = np.column_stack([
 new_data["VehAge"].to_numpy().astype(np.float64),
 new_data["Area"].map_elements(lambda v: level_map[v], return_dtype=pl.Int64).to_numpy().astype(np.float64),
])
preds = session.run(None, {"input": raw})[0] # shape (3, 1)

Performance Benchmarks

Memory Usage

Dependencies

Rust

Python

License