API Reference

compile_formula(model, data_example::NamedTuple) -> UnifiedCompiled

Primary API for compiling statistical models into high-performance evaluators.

Position Mapping System

This function implements a position mapping system that converts statistical formulas into zero-allocation execution plans. The system works in three phases:

Phase 1: Formula Decomposition

• Extracts the schema-applied formula from the fitted model
• Converts StatsModels terms into typed operations (LoadOp, ConstantOp, etc.)
• Assigns unique scratch positions to intermediate values and output positions to final results

Phase 2: Position Allocation

• Uses CompilationContext.position_map to track term → position mappings
• Allocates consecutive scratch positions starting from 1
• Maps each model matrix column to a specific output position

Phase 3: Type Specialization

• Embeds all positions as compile-time type parameters
• Creates operations like LoadOp{:x, 3}() (load column :x into scratch position 3)
• Enables zero-allocation execution through complete type specialization

Position Mapping Examples

# Simple formula: y ~ 1 + x
# Position mapping:
# scratch[1] = 1.0          (intercept, ConstantOp{1.0, 1})
# scratch[2] = data.x[row]  (variable x, LoadOp{:x, 2})  
# output[1] = scratch[1]    (CopyOp{1, 1})
# output[2] = scratch[2]    (CopyOp{2, 2})

# Interaction: y ~ x * z  
# Position mapping:
# scratch[1] = data.x[row]     (LoadOp{:x, 1})
# scratch[2] = data.z[row]     (LoadOp{:z, 2}) 
# scratch[3] = scratch[1] * scratch[2]  (BinaryOp{:*, 1, 2, 3})
# output[1] = scratch[1], output[2] = scratch[2], output[3] = scratch[3]

# Function: y ~ log(x)
# Position mapping:
# scratch[1] = data.x[row]     (LoadOp{:x, 1})
# scratch[2] = log(scratch[1]) (UnaryOp{:log, 1, 2})
# output[1] = scratch[2]       (CopyOp{2, 1})

Performance Characteristics

• Scratch space: Fixed size allocated once, reused for all rows
• Type stability: All positions known at compile time → zero allocations
• Execution: Pure array indexing with no dynamic dispatch
• Memory: O(maxscratchpositions) + O(output_size) per formula

Arguments

• model: Fitted statistical model (GLM, LMM, etc.) with schema-applied formula
• data_example: NamedTuple with sample data for type inference and schema validation

Returns

UnifiedCompiled{T, OpsTuple, ScratchSize, OutputSize} containing:

• Type-specialized operation tuple
• Pre-allocated scratch buffer
• Position mappings embedded in operation types

compile_formula(formula::StatsModels.FormulaTerm, data_example::NamedTuple) -> UnifiedCompiled

Convenience overload to compile directly from a StatsModels.FormulaTerm and column-table data. This mirrors the model-based entry point but skips get_fixed_effects_formula.

modelrow(model, data, row_idx) -> Vector{Float64}

Evaluate a single row and return a new vector (allocating version). Uses compiled formulas for optimal performance.

Example

row_values = modelrow(model, data, 1)  # Returns Vector{Float64}

modelrow(model, data, row_indices) -> Matrix{Float64}

Evaluate multiple rows and return a new matrix (allocating version). Uses compiled formulas for optimal performance.

Example

matrix = modelrow(model, data, [1, 5, 10])  # Returns Matrix{Float64}

modelrow(compiled_formula, data, row_idx) -> Vector{Float64}

Evaluate a single row with pre-compiled compiled formula.

Example

compiled = compile_formula(model, data)
row_values = modelrow(compiled, data, 1)  # Returns Vector{Float64}

modelrow(compiled_formula, data, row_indices) -> Matrix{Float64}

Evaluate multiple rows with pre-compiled compiled formula.

Example

compiled = compile_formula(model, data)
matrix = modelrow(compiled, data, [1, 5, 10])  # Returns Matrix{Float64}

modelrow(model, scenario::DataScenario, row_idx) -> Vector{Float64}

Evaluate model row using a data scenario (allocating version).

modelrow(compiled::UnifiedCompiled, scenario::DataScenario, row_idx) -> Vector{Float64}

Evaluate model row using a data scenario with UnifiedCompiled (allocating version).

modelrow!(row_vec, compiled_formula, data, row_idx)

Evaluate a single row of the model matrix in-place (zero-allocation).

Arguments

• row_vec::AbstractVector{Float64}: Pre-allocated output vector (modified in-place)
• compiled_formula: Compiled formula from compile_formula
• data: Data in Tables.jl format (preferably from Tables.columntable)
• row_idx::Int: Row index to evaluate

Returns

• row_vec: The same vector passed in, now containing the evaluated row

Example

compiled = compile_formula(model, data)
row_vec = Vector{Float64}(undef, length(compiled))
modelrow!(row_vec, compiled, data, 1)  # Zero allocations

modelrow!(row_vec, model, data, row_idx; cache=true)

Evaluate a single row of the model matrix in-place with automatic compilation.

Arguments

• row_vec::AbstractVector{Float64}: Pre-allocated output vector (modified in-place)
• model: Statistical model (GLM, MixedModel, etc.)
• data: Data in Tables.jl format
• row_idx::Int: Row index to evaluate
• cache::Bool: Whether to cache compiled formula (default: true)

Returns

• row_vec: The same vector passed in, now containing the evaluated row

Example

model = lm(@formula(y ~ x + group), df)
data = Tables.columntable(df)
row_vec = Vector{Float64}(undef, size(modelmatrix(model), 2))
modelrow!(row_vec, model, data, 1)

ModelRowEvaluator{D, O}

Pre-compiled evaluator using compiled formulas only.

OverrideVector{T} <: AbstractVector{T}

A lazy vector that returns the same override value for all indices. This avoids allocating full arrays when setting all observations to a representative value.

Example

# Instead of: fill(2.5, 1_000_000)  # Allocates 8MB
# Use: OverrideVector(2.5, 1_000_000)  # Allocates ~32 bytes

DataScenario

Represents a data scenario with specific variable overrides. Contains the modified data that can be used directly with compiled formulas.

Fields

• name::String: Descriptive name for the scenario
• overrides::Dict{Symbol,Any}: Variable overrides (mutable for iterative development)
• data::NamedTuple: Modified column-table data with OverrideVectors applied
• original_data::NamedTuple: Original unmodified data for reference

create_scenario(name, original_data; overrides...)
create_scenario(name, original_data, overrides::Dict)

Create a data scenario with specified variable overrides for counterfactual analysis.

Arguments

• name::String: Descriptive name for this scenario
• original_data::NamedTuple: Original data in column-table format (from Tables.columntable)
• overrides...: Keyword arguments for variable overrides (or Dict in second method)

Returns

• DataScenario: Object containing original data, overrides, and modified data with OverrideVectors

Example

data = Tables.columntable(df)

# Override single variable to mean
scenario1 = create_scenario("x_at_mean", data; x = mean(data.x))

# Override multiple variables for policy analysis
scenario2 = create_scenario("policy", data; x = 2.5, group = "A", treatment = true)

# Use dictionary for dynamic overrides
overrides = Dict(:dose => 100.0, :region => "North")
scenario3 = create_scenario("high_dose_north", data, overrides)

# Evaluate with compiled formula (zero-allocation)
compiled = compile_formula(model)
row_vec = Vector{Float64}(undef, length(compiled))
compiled(row_vec, scenario1.data, row_idx)

build_derivative_evaluator(compiled, data; vars, chunk=:auto) -> DerivativeEvaluator

Build a ForwardDiff-based derivative evaluator for a fixed set of variables.

Arguments:

• compiled::UnifiedCompiled: Result of compile_formula(model, data).
• data::NamedTuple: Column-table data (e.g., Tables.columntable(df)).
• vars::Vector{Symbol}: Variables to differentiate with respect to (typically continuous predictors).
• chunk: ForwardDiff.Chunk{N}() or :auto (uses Chunk{length(vars)}).

Returns:

• DerivativeEvaluator: Prebuilt evaluator object reusable across rows.

Notes:

• Compile once per model + variable set; reuse across calls.
• Zero allocations in steady state after warmup (typed closure + config; no per-call merges).
• Keep vars fixed for best specialization.

derivative_modelrow!(J, deval, row) -> AbstractMatrix{Float64}

Fill J with the Jacobian of one model row with respect to deval.vars.

Arguments:

• J::AbstractMatrix{Float64}: Preallocated buffer of size (n_terms, n_vars).
• deval::DerivativeEvaluator: Built by build_derivative_evaluator.
• row::Int: Row index (1-based).

Returns:

• The same J buffer, with J[i, j] = ∂X[i]/∂vars[j] for the given row.

Notes:

• Orientation is (n_terms, n_vars); n_terms == length(compiled).
• Small allocations (~368 bytes) due to ForwardDiff internals. For strict zero-allocation requirements, use derivative_modelrow_fd! instead.

derivative_modelrow(deval, row) -> Matrix{Float64}

Allocating convenience wrapper that returns the Jacobian for one row.

derivative_modelrow_fd!(J, compiled, data, row; vars, step=:auto)

Finite-difference Jacobian for a single row using central differences (standalone).

Arguments:

• J::AbstractMatrix{Float64}: Preallocated (n_terms, n_vars) buffer.
• compiled::UnifiedCompiled: Result of compile_formula.
• data::NamedTuple: Column-table data.
• row::Int: Row index.
• vars::Vector{Symbol}: Variables to differentiate with respect to.
• step: Numeric step size or :auto (eps()^(1/3) * max(1, |x|)).

Notes:

• Two evaluations per variable; useful as a robust fallback and for cross-checks.
• This standalone path allocates per call (builds per-call overrides and small temporaries). For zero allocations after warmup, prefer the evaluator FD path (derivative_modelrow_fd_pos!).

contrast_modelrow!(Δ, compiled, data, row; var, from, to)

Compute a discrete contrast at one row for a single variable: Δ = X(to) − X(from).

Arguments:

• Δ::AbstractVector{Float64}: Preallocated buffer of length n_terms.
• compiled::UnifiedCompiled: Result of compile_formula.
• data::NamedTuple: Column-table data.
• row::Int: Row index.
• var::Symbol: Variable to change (e.g., :group3).
• from, to: Values to contrast (level names or CategoricalValue for categorical; numbers for discrete).

Notes:

• Uses a row-local override; for categorical columns, values are normalized to the column's levels.

continuous_variables(compiled, data) -> Vector{Symbol}

Return a list of continuous variable symbols present in the compiled ops, excluding categoricals detected via ContrastOps. Filters by eltype(data[sym]) <: Real.

marginal_effects_eta!(g, de, beta, row; backend=:ad)

Fill g with marginal effects of η = Xβ w.r.t. de.vars at row. Implements: g = J' * β, where J = ∂X/∂vars.

Arguments:

• backend::Symbol: :ad (ForwardDiff) or :fd (finite differences)

Backends and allocations:

• :ad: Uses ForwardDiff automatic differentiation. Small allocations (~368 bytes) due to AD internals, but faster and more accurate.
• :fd: Uses zero-allocation finite differences. Strict 0 bytes after warmup, but slightly slower due to multiple function evaluations.
• Allocating convenience (marginal_effects_eta) allocates the result vector by design.

Recommendations:

• Use :fd backend for strict zero-allocation requirements
• Use :ad backend for speed and numerical accuracy (default)

marginal_effects_mu!(g, de, beta, row; link, backend=:ad)

Compute marginal effects of μ = g⁻¹(η) at row via chain rule: dμ/dx = (dμ/dη) * (dη/dx).

Arguments:

• link: Link function (e.g., IdentityLink(), LogLink(), LogitLink())
• backend::Symbol: :ad (ForwardDiff) or :fd (finite differences)

Backends and allocations:

• :ad: Uses ForwardDiff via η path. Small allocations (~368 bytes) due to AD internals, but faster and more accurate.
• :fd: Uses zero-allocation finite differences. Strict 0 bytes after warmup, but slightly slower due to multiple function evaluations.
• Allocating convenience (marginal_effects_mu) allocates the result vector by design.

Recommendations:

• Use :fd backend for strict zero-allocation requirements
• Use :ad backend for speed and numerical accuracy (default)