quivers.formulas fit entry point

The user-facing fit entry point, the formula_to_qvr source emitter, and the BayesianFit result wrapper.

_fit

User-facing fit entry point and the BayesianFit result wrapper.

The compilation path is fully AST-driven: the formula lens emits a quivers.dsl.ast_nodes.Module, the existing quivers.dsl.compiler.Compiler consumes it directly (no source-string round-trip), and inference runs on the resulting quivers.continuous.programs.MonadicProgram. Source text is generated only when the user requests it via formula_to_qvr or BayesianFit.dump_qvr, in which case quivers.dsl.emit.module_to_source walks the same AST to produce canonical .qvr source.

BayesianFit didactic-model

BayesianFit(*, formula: Formula = field(opaque=True), family: Family, program: MonadicProgram = field(opaque=True), posterior: MCMCResult | Guide = field(opaque=True), observations: Mapping[str, Tensor] = field(default_factory=dict, opaque=True), reparameterize: Literal['centered', 'noncentered'] = 'noncentered', binomial_trials: int | str = 1, thresholds_by: str | None = None, mixture_components: int = 2, fixed_prior: str = 'Normal(0.0, 5.0)', random_scale_prior: str = 'HalfNormal(1.0)', priors: Mapping[str, str] = field(default_factory=dict, opaque=True), predictor: Module | None = field(default=None, opaque=True), predictor_data: Tensor | None = field(default=None, opaque=True), predictor_name: str | None = None)

Bases: Model

A fitted Bayesian regression: the compiled program, the parsed formula, the family, the user-supplied data, and the posterior samples.

ATTRIBUTE	DESCRIPTION
formula	Parsed formula IR. TYPE: Formula
family	Response family used at compile time. TYPE: Family
program	The compiled program. TYPE: MonadicProgram
posterior	quivers.inference.MCMCResult from NUTS / HMC, or a quivers.inference.guides.base.Guide from SVI. TYPE: MCMCResult or Guide
observations	Inference-time observations dict (response + per-column covariates + per-group plate indices). TYPE: Mapping[str, Tensor]

qvr_source property

qvr_source: str

Lazily emit the AST-equivalent .qvr source for display.

dump_qvr

dump_qvr(path: str | Path) -> Path

Write the AST-equivalent .qvr source to path and return the resulting Path.

Source code in src/quivers/formulas/_fit.py

def dump_qvr(self, path: str | Path) -> Path:
    """Write the AST-equivalent ``.qvr`` source to ``path`` and
    return the resulting `Path`.
    """
    out = Path(path)
    out.write_text(self.qvr_source)
    return out

fit

fit(formula: str, *, data: IntoDataFrame, family: str | Family = 'gaussian', method: Literal['nuts', 'hmc', 'svi'] = 'nuts', num_warmup: int = 500, num_samples: int = 1000, num_chains: int = 4, fixed_prior: str = 'Normal(0.0, 5.0)', random_scale_prior: str = 'HalfNormal(1.0)', priors: Mapping[str, str] | None = None, guide: type | None = None, reparameterize: Literal['centered', 'noncentered'] = 'noncentered', binomial_trials: int | str = 1, thresholds_by: str | None = None, mixture_components: int = 2, predictor: Module | None = None, predictor_data: Tensor | None = None, seed: int = 0) -> BayesianFit

Compile a brms-style formula, fit it, and return the result.

See quivers.formulas for surface details. This entry point composes formula_from_data, FormulaToQVRModule, Compiler, and the inference layer in one call.

family="categorical" infers contiguous labels 0..K-1 and uses label zero as the reference. An attached categorical predictor may return (N, K-1) reference logits or (N, K) full logits. For family="mixture", mixture_components selects any finite component count of at least two; the likelihood integrates assignments exactly.

PARAMETER	DESCRIPTION
formula	brms/lme4-style response formula. TYPE: str
data	Pandas, Polars, or another Narwhals-compatible dataframe. TYPE: IntoDataFrame
family	Registered response family or a family value. TYPE: str or Family DEFAULT: 'gaussian'
method	Inference algorithm. TYPE: ('nuts', 'hmc', 'svi') DEFAULT: "nuts"
mixture_components	Number of Gaussian components for the mixture family. TYPE: int DEFAULT: 2
predictor	Optional differentiable contribution on the linear-predictor scale. TYPE: Module or None DEFAULT: None
predictor_data	Input passed to predictor on each SVI step. TYPE: Tensor or None DEFAULT: None

RETURNS	DESCRIPTION
BayesianFit	Parsed formula, compiled program, fitted posterior, and observations.

Source code in src/quivers/formulas/_fit.py

def fit(
    formula: str,
    *,
    data: IntoDataFrame,
    family: str | Family = "gaussian",
    method: Literal["nuts", "hmc", "svi"] = "nuts",
    num_warmup: int = 500,
    num_samples: int = 1000,
    num_chains: int = 4,
    fixed_prior: str = "Normal(0.0, 5.0)",
    random_scale_prior: str = "HalfNormal(1.0)",
    priors: Mapping[str, str] | None = None,
    guide: type | None = None,
    reparameterize: Literal["centered", "noncentered"] = "noncentered",
    binomial_trials: int | str = 1,
    thresholds_by: str | None = None,
    mixture_components: int = 2,
    predictor: nn.Module | None = None,
    predictor_data: torch.Tensor | None = None,
    seed: int = 0,
) -> BayesianFit:
    """Compile a brms-style formula, fit it, and return the result.

    See [`quivers.formulas`][quivers.formulas] for surface details.  This entry
    point composes `formula_from_data`, `FormulaToQVRModule`,
    `Compiler`, and the inference layer in one call.

    ``family="categorical"`` infers contiguous labels ``0..K-1`` and uses
    label zero as the reference. An attached categorical predictor may return
    ``(N, K-1)`` reference logits or ``(N, K)`` full logits. For
    ``family="mixture"``, ``mixture_components`` selects any finite component
    count of at least two; the likelihood integrates assignments exactly.

    Parameters
    ----------
    formula : str
        brms/lme4-style response formula.
    data : IntoDataFrame
        Pandas, Polars, or another Narwhals-compatible dataframe.
    family : str or Family
        Registered response family or a family value.
    method : {"nuts", "hmc", "svi"}
        Inference algorithm.
    mixture_components : int
        Number of Gaussian components for the mixture family.
    predictor : nn.Module or None
        Optional differentiable contribution on the linear-predictor scale.
    predictor_data : torch.Tensor or None
        Input passed to ``predictor`` on each SVI step.

    Returns
    -------
    BayesianFit
        Parsed formula, compiled program, fitted posterior, and observations.
    """
    if isinstance(family, str):
        if family not in families:
            raise ValueError(
                f"fit: unknown family {family!r}; choices are {sorted(families)}"
            )
        family_obj = families[family]
    else:
        family_obj = family

    parsed = formula_from_data(formula, data)
    predictor_name = "neural_eta" if predictor is not None else None
    if predictor is not None and predictor_data is None:
        raise ValueError("fit: predictor_data is required when predictor is supplied")
    if predictor is None and predictor_data is not None:
        raise ValueError("fit: predictor_data requires predictor")
    if (
        predictor is not None
        and method != "svi"
        and any(parameter.requires_grad for parameter in predictor.parameters())
    ):
        raise ValueError(
            "fit: a trainable external predictor requires method='svi'; "
            "freeze its parameters before using NUTS or HMC"
        )
    lens = FormulaToQVRModule(
        family_obj,
        fixed_prior=fixed_prior,
        random_scale_prior=random_scale_prior,
        user_priors=priors,
        reparameterize=reparameterize,
        binomial_trials=binomial_trials,
        thresholds_by=thresholds_by,
        mixture_components=mixture_components,
        predictor_name=predictor_name,
    )
    module, _ = lens.forward(parsed)
    compiler = Compiler(module)
    program_runtime = compiler.compile()
    morphism = program_runtime.morphism
    if not isinstance(morphism, MonadicProgram):
        raise TypeError(
            f"fit: compiled morphism has type "
            f"{type(morphism).__name__}, expected MonadicProgram"
        )
    program = morphism

    observations: dict[str, torch.Tensor] = {}
    observations.update(lens.fixed_column_observations(parsed))
    response_name = parsed.response_name
    observations[response_name] = torch.as_tensor(
        parsed.response_values.copy(), dtype=torch.float32
    ).reshape(-1)
    for group, codes in parsed.group_indices.items():
        observations[f"{_qvr_name(group)}_idx"] = torch.as_tensor(
            list(codes), dtype=torch.long
        )
    if family_obj.name == "binomial":
        trials = _binomial_trial_tensor(
            data, binomial_trials, n_obs=parsed.response_values.shape[0]
        )
        response = observations[response_name]
        _validate_binomial_response(response, trials, caller="fit")
        if isinstance(binomial_trials, str):
            observations[_qvr_name(binomial_trials)] = trials
    if predictor is not None:
        assert predictor_data is not None
        assert predictor_name is not None
        predictor_categories = (
            FormulaToQVRModule._category_count(parsed, family="categorical")
            if family_obj.name == "categorical"
            else None
        )
        observations[predictor_name] = _predictor_output(
            predictor,
            predictor_data,
            n_obs=parsed.response_values.shape[0],
            n_categories=predictor_categories,
        ).detach()
    else:
        predictor_categories = None

    torch.manual_seed(seed)
    if method == "svi":
        posterior = _fit_svi(
            program,
            observations,
            num_samples,
            guide_cls=guide,
            predictor=predictor,
            predictor_data=predictor_data,
            predictor_name=predictor_name,
            predictor_categories=predictor_categories,
        )
    else:
        posterior = _fit_mcmc(
            program,
            observations,
            sampler=method,
            num_warmup=num_warmup,
            num_samples=num_samples,
            num_chains=num_chains,
        )

    return BayesianFit(
        formula=parsed,
        family=family_obj,
        program=program,
        posterior=posterior,
        observations=observations,
        reparameterize=reparameterize,
        binomial_trials=binomial_trials,
        thresholds_by=thresholds_by,
        mixture_components=mixture_components,
        fixed_prior=fixed_prior,
        random_scale_prior=random_scale_prior,
        priors=dict(priors or {}),
        predictor=predictor,
        predictor_data=predictor_data,
        predictor_name=predictor_name,
    )

formula_to_qvr

formula_to_qvr(formula: str, *, data: IntoDataFrame, family: str | Family = 'gaussian', fixed_prior: str = 'Normal(0.0, 5.0)', random_scale_prior: str = 'HalfNormal(1.0)', priors: Mapping[str, str] | None = None, reparameterize: Literal['centered', 'noncentered'] = 'noncentered', binomial_trials: int | str = 1, thresholds_by: str | None = None, mixture_components: int = 2, predictor_name: str | None = None, path: str | Path | None = None) -> str

Emit .qvr source for a brms-style formula without fitting.

Builds the formula AST → QVR Module via FormulaToQVRModule, then serialises the module via quivers.dsl.emit.module_to_source. Optionally writes the result to path. mixture_components selects any integer component count of at least two for family="mixture".

Source code in src/quivers/formulas/_fit.py

def formula_to_qvr(
    formula: str,
    *,
    data: IntoDataFrame,
    family: str | Family = "gaussian",
    fixed_prior: str = "Normal(0.0, 5.0)",
    random_scale_prior: str = "HalfNormal(1.0)",
    priors: Mapping[str, str] | None = None,
    reparameterize: Literal["centered", "noncentered"] = "noncentered",
    binomial_trials: int | str = 1,
    thresholds_by: str | None = None,
    mixture_components: int = 2,
    predictor_name: str | None = None,
    path: str | Path | None = None,
) -> str:
    """Emit ``.qvr`` source for a brms-style formula without fitting.

    Builds the formula AST → QVR Module via `FormulaToQVRModule`,
    then serialises the module via [`quivers.dsl.emit.module_to_source`][quivers.dsl.emit.module_to_source].
    Optionally writes the result to ``path``. ``mixture_components`` selects
    any integer component count of at least two for ``family="mixture"``.
    """
    if isinstance(family, str):
        if family not in families:
            raise ValueError(
                f"formula_to_qvr: unknown family {family!r}; choices are "
                f"{sorted(families)}"
            )
        family_obj = families[family]
    else:
        family_obj = family
    parsed = formula_from_data(formula, data)
    if family_obj.name == "binomial":
        trials = _binomial_trial_tensor(
            data, binomial_trials, n_obs=parsed.response_values.shape[0]
        )
        response = torch.as_tensor(parsed.response_values.copy())
        _validate_binomial_response(response, trials, caller="formula_to_qvr")
    lens = FormulaToQVRModule(
        family_obj,
        fixed_prior=fixed_prior,
        random_scale_prior=random_scale_prior,
        user_priors=priors,
        reparameterize=reparameterize,
        binomial_trials=binomial_trials,
        thresholds_by=thresholds_by,
        mixture_components=mixture_components,
        predictor_name=predictor_name,
    )
    module, _ = lens.forward(parsed)
    source = module_to_source(module)
    if path is not None:
        Path(path).write_text(source)
    return source