Formula models, JSON candidates, training loop and control charts

README.md

@@ -23,10 +23,15 @@ project.
 
 - `PymcWrapper`: a thin wrapper around a PyMC model that handles prior/posterior
   sampling, scoring, LOO cross-validation, and saving/loading fitted models to disk.
-- A library of ready-to-use `bayesbuilding.models` functions covering common building
-  energy patterns: seasonal change-point energy signatures, heating/cooling with
-  occupation change points, solar-radiation-augmented models, artificial lighting, and
-  PV panel production (constant efficiency and NOCT models).
+- Forward models written as formulas (`bayesbuilding.formula.FormulaModel`), e.g.
+  `"g[occ]*max(tau - text, 0) - fs*rad + base[occ]"`: no Python function to write for
+  a new model variant. (`bayesbuilding.models` holds the older hand-written functions,
+  frozen and kept only to reload traces saved with them.)
+- JSON candidate configs (`bayesbuilding.candidates`) and a training loop that fits,
+  scores (LOO, R2, NMBE, CV(RMSE), MAE at daily/weekly/monthly resolution) and ranks
+  candidates (`bayesbuilding.training.train_candidates`).
+- Control charts on residuals (`bayesbuilding.control_charts`): X-bar, EWMA (with alarm
+  confirmation/reset) and CUSUM, with limits following the model's own sigma.
 - Plotting helpers (`bayesbuilding.plotting`) for prior/posterior comparison, HDI time
   series plots, and change-point diagnostic plots, with both `matplotlib` and `plotly`
   backends.
@@ -47,7 +52,7 @@ import numpy as np
 import pandas as pd
 import pymc as pm
 
-from bayesbuilding.models import season_cp_heating_es
+from bayesbuilding.formula import FormulaModel
 from bayesbuilding.wrapper import PymcWrapper
 
 # Monthly external temperature and heating consumption
@@ -57,10 +62,10 @@ data = pd.DataFrame(
 )
 data["heating"] = 50 * np.maximum(14 - data["Text"], 0) + 50 + np.random.randn(12) * 5
 
-# Define the model and priors for a seasonal change-point energy signature:
-# heating = g * max(tau - Text, 0) + base
+# Define the model and priors for a seasonal change-point energy signature.
+# Symbols that are not priors are inputs, i.e. columns of x.
 model = PymcWrapper(
-    model_function=season_cp_heating_es,
+    model_function=FormulaModel(mu="g*max(tau - Text, 0) + base", sigma="sigma"),
     priors_dict={
         "g": (pm.Normal, dict(name="g", mu=40, sigma=5)),
         "tau": (pm.Normal, dict(name="tau", mu=12, sigma=1)),
@@ -82,7 +87,28 @@ reloaded = PymcWrapper()
 reloaded.load_model("my_model")
 ```
 
-See `bayesbuilding/models.py` for the full list of built-in model functions and
-`tests/test_wrapper.py` for a complete end-to-end example, including scoring on held-out
+## Formula grammar
+
+- Operators `+ - * / **`, numeric constants, and the functions `max(a, b)`,
+  `min(a, b)`, `switch(cond, a, b)`, `sigmoid`, `sqrt`, `exp`, `log`, `abs`.
+- `g[occ]` / `g[occ, 1]` index a vector/matrix prior by a categorical input.
+- `sigma` is the likelihood scale and may be heteroscedastic, e.g.
+  `"sqrt(s0[occ]**2 + (s1*sigmoid((tau - text)/1.5))**2)"`; `lower` (default `0.0`)
+  is returned for a `TruncatedNormal` likelihood, `None` disables it.
+- Input order (the columns of `x`): continuous drivers first, then categorical inputs,
+  each in order of first appearance.
+
+Expressions are parsed against a whitelist, never `eval`-ed. In a candidate config:
+
+```json
+{
+  "name": "dt_occ",
+  "model": {"mu": "g[occ]*dt - fs[occ]*rad", "sigma": "s0[occ]", "lower": 0.0},
+  "inputs": {"dt": "dt__C", "rad": "GHI__W/m2", "occ": "occupation"},
+  "priors": {"g": {"dist": "HalfNormal", "kwargs": {"sigma": 500, "shape": 2}}, "...": {}}
+}
+```
+
+See `tests/test_wrapper.py` for a complete end-to-end example, including scoring on held-out
 data and plotting predictions with `bayesbuilding.plotting.time_series_hdi` and
 `changepoint_graph`.

bayesbuilding/candidates.py

@@ -0,0 +1,178 @@
+"""JSON configuration of candidate models to fit and compare.
+
+A config file looks like::
+
+    {
+      "config_id": 69370,
+      "target": "heating_energy",
+      "feature_pipe": {...},
+      "candidates": [
+        {
+          "name": "dt_occ",
+          "model": {"mu": "g[occ]*dt - fs[occ]*rad", "sigma": "s0[occ]", "lower": 0.0},
+          "inputs": {"dt": "<column>", "rad": "<column>", "occ": "<column>"},
+          "priors": {
+              "g": {"dist": "HalfNormal", "kwargs": {"sigma": 10, "shape": 2}},
+              ...
+          },
+          "draws": 4000, "tune": 1000
+        }
+      ],
+      "unused_candidates": []
+    }
+
+``model`` is a :class:`~bayesbuilding.formula.FormulaModel` spec. ``inputs`` maps
+each input symbol of the formula to a real column name, since column names vary
+from one site to the next. ``feature_pipe`` is an opaque dict left to the caller
+(e.g. a Tide pipe computing the columns once for every candidate).
+"""
+
+import json
+import warnings
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import pymc as pm
+
+from bayesbuilding.formula import FormulaModel
+from bayesbuilding.legacy_formulas import legacy_to_formula
+
+
+def build_priors_dict(priors: dict) -> dict:
+    """Turn a JSON-serializable priors spec into PymcWrapper's priors_dict.
+
+    Input  : {"g": {"dist": "Normal", "kwargs": {"mu": 100.0, "sigma": 10.0}}, ...}
+    Output : {"g": (pm.Normal, {"name": "g", "mu": 100.0, "sigma": 10.0}), ...}
+    """
+    return {
+        name: (getattr(pm, spec["dist"]), {"name": name, **spec["kwargs"]})
+        for name, spec in priors.items()
+    }
+
+
+@dataclass
+class CandidateConfig:
+    """One forward model to fit and compare, within a BayesConfig.
+
+    ``model`` is given as a formula spec dict (or a FormulaModel) and is stored
+    as a FormulaModel bound to the prior names, so ``model.inputs`` is the
+    ordered list of input symbols. A legacy model name (a string) is still
+    accepted and translated with :func:`~bayesbuilding.legacy_formulas.
+    legacy_to_formula`, with a DeprecationWarning.
+    """
+
+    name: str
+    model: dict | FormulaModel
+    inputs: dict[str, str]
+    priors: dict
+    prior_draws: int = 4000
+    draws: int = 4000
+    tune: int = 1000
+    chains: int | None = None
+    random_seed: int | None = None
+
+    def __post_init__(self):
+        if isinstance(self.model, str):
+            warnings.warn(
+                f"Candidate {self.name!r}: legacy model name {self.model!r}, "
+                "replace it by its formula (see legacy_to_formula)",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+            self.model = legacy_to_formula(self.model)
+        if isinstance(self.model, dict):
+            self.model = FormulaModel.from_dict(self.model)
+        try:
+            self.model.bind(self.priors)
+        except ValueError as e:
+            raise ValueError(f"Candidate {self.name!r}: {e}") from e
+
+    def to_dict(self) -> dict:
+        spec = self.model.to_dict()
+        spec.pop("inputs", None)  # derived from the priors on load
+        return {
+            "name": self.name,
+            "model": spec,
+            "inputs": dict(self.inputs),
+            "priors": self.priors,
+            "prior_draws": self.prior_draws,
+            "draws": self.draws,
+            "tune": self.tune,
+            "chains": self.chains,
+            "random_seed": self.random_seed,
+        }
+
+
+@dataclass
+class BayesConfig:
+    """A set of candidate models sharing the same target and feature pipe."""
+
+    target: str
+    candidates: list[CandidateConfig]
+    config_id: int | str | None = None
+    feature_pipe: dict = field(default_factory=dict)
+    unused_candidates: list[CandidateConfig] = field(default_factory=list)
+
+    @classmethod
+    def from_dict(cls, raw: dict) -> "BayesConfig":
+        raw = dict(raw)
+        candidates = [CandidateConfig(**c) for c in raw.pop("candidates")]
+        unused = [CandidateConfig(**c) for c in raw.pop("unused_candidates", [])]
+        return cls(candidates=candidates, unused_candidates=unused, **raw)
+
+    @classmethod
+    def from_json(cls, path: Path | str) -> "BayesConfig":
+        return cls.from_dict(json.loads(Path(path).read_text(encoding="utf-8")))
+
+    def to_dict(self) -> dict:
+        return {
+            "config_id": self.config_id,
+            "target": self.target,
+            "feature_pipe": self.feature_pipe,
+            "candidates": [c.to_dict() for c in self.candidates],
+            "unused_candidates": [c.to_dict() for c in self.unused_candidates],
+        }
+
+    def to_json(self, path: Path | str):
+        Path(path).write_text(
+            json.dumps(self.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8"
+        )
+
+    def get_candidate(self, name: str) -> CandidateConfig:
+        for candidate in self.candidates + self.unused_candidates:
+            if candidate.name == name:
+                return candidate
+        raise KeyError(f"No candidate named {name!r}")
+
+
+def posterior_as_priors(candidate: CandidateConfig, wrapper) -> dict:
+    """Turn a fitted wrapper's parameter posterior into a new priors spec.
+
+    For each of ``candidate.priors``' variables, keeps the original prior's
+    distribution family (and any kwarg other than ``mu``/``sigma``, e.g.
+    ``TruncatedNormal``'s ``lower``) but replaces whichever of ``mu``/``sigma``
+    the original spec already declared with the posterior sample's empirical
+    mean/std -- i.e. a moment-matched parametric refit of "today's posterior"
+    usable as "tomorrow's prior" (see :func:`~bayesbuilding.training.
+    fit_candidate`'s ``priors`` argument). This is a standard approximation for
+    sequential Bayesian updating: PyMC priors must be given as a parametric
+    family, so the posterior's exact (possibly non-Gaussian) shape isn't
+    preserved, only its first two moments. Only kwargs already present in the
+    original spec are touched -- a one-parameter family like ``HalfNormal``
+    (kwarg ``sigma`` only, no ``mu``) doesn't get a spurious ``mu`` injected,
+    which would break the next ``fit_candidate(..., priors=updated)`` call.
+    """
+    posterior = wrapper.traces["sampling"].posterior
+    updated = {}
+    for name, spec in candidate.priors.items():
+        samples = posterior[name].to_numpy().reshape(-1)
+        moments = {}
+        if "mu" in spec["kwargs"]:
+            moments["mu"] = float(samples.mean())
+        if "sigma" in spec["kwargs"]:
+            moments["sigma"] = float(samples.std())
+        updated[name] = {
+            "dist": spec["dist"],
+            "kwargs": {**spec["kwargs"], **moments},
+        }
+    return updated