mihenk 0.1.0__py3-none-any.whl

mihenk/__init__.py

@@ -0,0 +1,31 @@
+"""mihenk: a touchstone for discovered equations.
+
+Before you interpret the law that PySINDy (or any equation-discovery tool)
+found, mihenk checks whether it could be an artifact of noise, of how
+derivatives were estimated, of which states the data happened to visit, or of
+analysis choices.
+
+    import mihenk
+    report = mihenk.scan(bundle, model)            # tier 0: data + equation
+    report = mihenk.scan_pysindy(sindy, x, t)      # tiers 0-1 with a fitted PySINDy model
+"""
+
+__version__ = "0.1.0"
+
+from .bundle import Bundle  # noqa: E402
+from .model import Model  # noqa: E402
+from .report import Check, Finding, Report  # noqa: E402
+from .scan import Thresholds, scan  # noqa: E402
+
+
+def scan_pysindy(sindy_model, x, t, *, feature_names=None, name="pysindy", val=None,
+                 thresholds=None, variants=True, seed=0, meta=None):
+    """Scan a fitted PySINDy model on the data it was fitted to (tiers 0 and 1)."""
+    from .adapters import pysindy as _ps
+
+    return _ps.scan_pysindy(sindy_model, x, t, feature_names=feature_names, name=name, val=val,
+                            thresholds=thresholds, variants=variants, seed=seed, meta=meta)
+
+
+__all__ = ["Bundle", "Model", "Report", "Finding", "Check", "Thresholds", "scan", "scan_pysindy",
+           "__version__"]

mihenk/adapters/__init__.py

@@ -0,0 +1 @@
+"""Adapters for discovery tools. Each is optional and imports its tool lazily."""

mihenk/adapters/pysindy.py

@@ -0,0 +1,146 @@
+"""PySINDy adapter (``pip install "mihenk[pysindy]"``), written against PySINDy 2.x.
+
+Turns a fitted ``pysindy.SINDy`` into a :class:`~mihenk.model.Model`, and builds
+the tier-1 callables: ``refit`` (the same configuration on resampled data) and
+``variants`` (the same data under other defensible derivative methods and
+sparsity thresholds).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from ..bundle import Bundle
+from ..model import Model, parse_term
+
+
+def _require():
+    try:
+        import pysindy  # noqa: F401
+    except ImportError as exc:  # pragma: no cover
+        raise ImportError('this needs PySINDy: pip install "mihenk[pysindy]"') from exc
+
+
+def from_pysindy(sindy, variables=None, source="pysindy") -> Model:
+    """Read variables, candidate library and coefficients from a fitted PySINDy model."""
+    _require()
+    if getattr(sindy, "n_control_features_", 0):
+        raise NotImplementedError("models with control inputs (u) are not supported in mihenk 0.1")
+    lib_name = type(sindy.feature_library).__name__
+    if "Weak" in lib_name or "PDE" in lib_name:
+        raise NotImplementedError(f"{lib_name} is not supported in mihenk 0.1 (ODE libraries only)")
+    variables = list(variables or sindy.feature_names)
+    names = sindy.get_feature_names()
+    try:
+        terms = [parse_term(n, variables) for n in names]
+    except Exception as exc:
+        raise ValueError(f"cannot read PySINDy feature names as expressions ({names[:5]}...): {exc}") from exc
+    coefs = np.asarray(sindy.coefficients(), dtype=float)
+    m = Model(variables, terms, coefs, library=list(terms), source=source)
+    m.meta = {"library": lib_name, "optimizer": type(sindy.optimizer).__name__,
+              "differentiation": type(sindy.differentiation_method).__name__}
+    return m
+
+
+def _fit_clone(sindy, pairs, variables, **overrides):
+    from sklearn.base import clone
+
+    m = clone(sindy)
+    if overrides:
+        m.set_params(**overrides)
+    xs = [x for x, _ in pairs]
+    ts = [t for _, t in pairs]
+    import warnings
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        if len(xs) == 1:
+            m.fit(xs[0], t=ts[0], feature_names=list(variables))
+        else:
+            m.fit(xs, t=ts, feature_names=list(variables))
+    model = from_pysindy(m, variables)
+    # STLSQ can stop at max_iter with coefficients far below its own threshold still non-zero;
+    # in a refit those are optimizer leftovers, not structure, so they are not counted as terms.
+    thr = m.optimizer.get_params().get("threshold")
+    if thr:
+        small = (np.abs(model.coefficients) < thr) & (model.coefficients != 0)
+        model.meta["pruned_below_threshold"] = int(small.sum())
+        model.coefficients[small] = 0.0
+    return model
+
+
+def make_refit(sindy, variables=None):
+    variables = list(variables or sindy.feature_names)
+    return lambda pairs: _fit_clone(sindy, pairs, variables)
+
+
+def make_variants(sindy, variables=None, threshold_factors=(0.5, 1.0, 1.5)) -> dict:
+    """Other defensible analyses of the same data: derivative method x sparsity threshold."""
+    import pysindy as ps
+
+    variables = list(variables or sindy.feature_names)
+    diffs = {"FiniteDifference": ps.FiniteDifference, "SmoothedFiniteDifference": ps.SmoothedFiniteDifference}
+    current_diff = type(sindy.differentiation_method).__name__
+    params = sindy.optimizer.get_params()
+    th0 = params.get("threshold")
+    out = {}
+    for dname, dcls in diffs.items():
+        for f in threshold_factors if th0 is not None else (1.0,):
+            if dname == current_diff and f == 1.0:
+                continue  # that is the original analysis
+            over = {"differentiation_method": dcls()}
+            label = dname
+            if th0 is not None:
+                over["optimizer__threshold"] = th0 * f
+                label += f", threshold {th0 * f:.3g}"
+            out[label] = (lambda pairs, over=over: _fit_clone(sindy, pairs, variables, **over))
+    return out
+
+
+def scan_pysindy(sindy, x, t, *, feature_names=None, name="pysindy", val=None, thresholds=None,
+                 variants=True, seed=0, meta=None):
+    from ..scan import scan
+
+    variables = list(feature_names or sindy.feature_names)
+    xs = x if isinstance(x, (list, tuple)) else [x]
+    if isinstance(t, (int, float)):
+        ts = [np.arange(len(xi)) * float(t) for xi in xs]
+    else:
+        ts = t if isinstance(t, (list, tuple)) else [t]
+    b = Bundle(list(xs), list(ts), variables, val=val, meta={"name": name, **(meta or {})})
+    model = from_pysindy(sindy, variables)
+    return scan(b, model, refit=make_refit(sindy, variables),
+                variants=make_variants(sindy, variables) if variants else None,
+                thresholds=thresholds, seed=seed)
+
+
+# --------------------------------------------------------------- CLI helpers
+DIFFS = {"fd": "FiniteDifference", "smoothed": "SmoothedFiniteDifference"}
+
+
+def build_sindy(degree: int = 2, threshold: float = 0.1, diff: str = "fd"):
+    """A PySINDy model from a small, serializable configuration (used by ``mihenk fit``)."""
+    _require()
+    import pysindy as ps
+
+    if diff not in DIFFS:
+        raise ValueError(f"diff must be one of {', '.join(DIFFS)}")
+    return ps.SINDy(optimizer=ps.STLSQ(threshold=threshold),
+                    feature_library=ps.PolynomialLibrary(degree=degree),
+                    differentiation_method=getattr(ps, DIFFS[diff])())
+
+
+def fit_bundle(bundle: Bundle, degree: int = 2, threshold: float = 0.1, diff: str = "fd"):
+    """Fit PySINDy to a bundle's trajectories; returns ``(sindy, model)``."""
+    import warnings
+
+    sindy = build_sindy(degree, threshold, diff)
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        if bundle.n_traj == 1:
+            sindy.fit(bundle.xs[0], t=bundle.ts[0], feature_names=bundle.variables)
+        else:
+            sindy.fit(bundle.xs, t=bundle.ts, feature_names=bundle.variables)
+    model = from_pysindy(sindy, bundle.variables, source="pysindy (mihenk fit)")
+    model.meta["pysindy_config"] = {"degree": degree, "threshold": threshold, "diff": diff}
+    return sindy, model

mihenk/bundle.py

@@ -0,0 +1,168 @@
+"""The mihenk bundle: the observations an equation was discovered from, as plain arrays.
+
+A bundle is one ``.npz`` file holding one or more trajectories. Shapes use T
+time samples and n state variables; trajectories may differ in length.
+
+Keys
+    x             (T, n)   one trajectory, or ``x_0, x_1, ...`` for several
+    t             (T,)     sample times (or ``t_0, t_1, ...``); a scalar ``dt`` also works
+    x_dot         (T, n)   optional measured derivatives (``x_dot_0, ...``)
+    val           (K,)     optional bool: trajectories held out from discovery.
+                           With one trajectory, the last 20% is held out instead.
+    meta          JSON string: variables, name, units, true_equations (synthetic only), ...
+
+``true_equations`` in meta (``{"x": "-10*x + 10*y", ...}``) unlocks the tier-2
+ground-truth checks; only synthetic data has it.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+
+@dataclass
+class Bundle:
+    xs: list                       # list of (T_k, n) arrays
+    ts: list                       # list of (T_k,) arrays
+    variables: list
+    x_dots: list | None = None
+    val: np.ndarray | None = None  # (K,) bool, held-out trajectories
+    meta: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        if isinstance(self.xs, np.ndarray):
+            self.xs = [self.xs]
+        if isinstance(self.ts, np.ndarray):
+            self.ts = [self.ts]
+        self.xs = [np.atleast_2d(np.asarray(x, float)) for x in self.xs]
+        self.ts = [np.asarray(t, float).ravel() for t in self.ts]
+        if len(self.xs) != len(self.ts):
+            raise ValueError(f"{len(self.xs)} trajectories but {len(self.ts)} time vectors")
+        n = self.xs[0].shape[1]
+        for k, (x, t) in enumerate(zip(self.xs, self.ts)):
+            if x.shape[1] != n:
+                raise ValueError(f"trajectory {k} has {x.shape[1]} variables, expected {n}")
+            if x.shape[0] != t.shape[0]:
+                raise ValueError(f"trajectory {k}: {x.shape[0]} samples but {t.shape[0]} times")
+            if x.shape[0] < 20:
+                raise ValueError(f"trajectory {k} has only {x.shape[0]} samples; mihenk needs at least 20")
+            if np.any(np.diff(t) <= 0):
+                raise ValueError(f"trajectory {k}: times must be strictly increasing")
+            if not np.all(np.isfinite(x)):
+                raise ValueError(f"trajectory {k} contains NaN or inf")
+        if len(self.variables) != n:
+            raise ValueError(f"{len(self.variables)} variable names for {n} columns")
+        self.variables = [str(v) for v in self.variables]
+        if self.x_dots is not None:
+            self.x_dots = [np.asarray(d, float) for d in self.x_dots]
+        if self.val is not None:
+            self.val = np.asarray(self.val, bool).ravel()
+            if self.val.shape[0] != len(self.xs):
+                raise ValueError("val must have one entry per trajectory")
+            if self.val.all():
+                raise ValueError("every trajectory is marked held-out; nothing is left for fitting")
+
+    # ------------------------------------------------------------ split
+    @property
+    def n_traj(self) -> int:
+        return len(self.xs)
+
+    @property
+    def n_vars(self) -> int:
+        return self.xs[0].shape[1]
+
+    @property
+    def name(self) -> str:
+        return self.meta.get("name", "bundle")
+
+    @property
+    def true_equations(self) -> dict | None:
+        return self.meta.get("true_equations")
+
+    def split(self, holdout_fraction: float = 0.2):
+        """(train, test) lists of (x, t) pairs.
+
+        Held-out trajectories when ``val`` is given (or, with several
+        trajectories, the last 20% of them); otherwise the tail of the single
+        trajectory.
+        """
+        pairs = list(zip(self.xs, self.ts))
+        if self.val is not None:
+            train = [p for p, v in zip(pairs, self.val) if not v]
+            test = [p for p, v in zip(pairs, self.val) if v]
+            return train, test, "held-out trajectories (val)"
+        if self.n_traj >= 3:
+            k = max(1, int(round(holdout_fraction * self.n_traj)))
+            return pairs[:-k], pairs[-k:], f"last {k} of {self.n_traj} trajectories"
+        train, test = [], []
+        for x, t in pairs:
+            cut = int(round((1 - holdout_fraction) * len(t)))
+            train.append((x[:cut], t[:cut]))
+            test.append((x[cut:], t[cut:]))
+        return train, test, f"last {int(holdout_fraction * 100)}% of each trajectory"
+
+    # ------------------------------------------------------------ io
+    def save(self, path) -> Path:
+        path = Path(path)
+        arrays = {}
+        for k, (x, t) in enumerate(zip(self.xs, self.ts)):
+            arrays[f"x_{k}"] = x
+            arrays[f"t_{k}"] = t
+            if self.x_dots is not None:
+                arrays[f"x_dot_{k}"] = self.x_dots[k]
+        if self.val is not None:
+            arrays["val"] = self.val
+        meta = dict(self.meta)
+        meta["variables"] = self.variables
+        arrays["meta"] = np.array(json.dumps(meta))
+        np.savez_compressed(path, **arrays)
+        return path
+
+    @classmethod
+    def load(cls, path) -> "Bundle":
+        path = Path(path)
+        if path.suffix == ".csv":
+            return cls.from_csv(path)
+        with np.load(path, allow_pickle=False) as z:
+            keys = set(z.files)
+            meta = json.loads(str(z["meta"])) if "meta" in keys else {}
+            if "x" in keys:
+                xs = [z["x"]]
+                ts = [_times(z, "t", xs[0].shape[0])]
+                dots = [z["x_dot"]] if "x_dot" in keys else None
+            else:
+                idx = sorted(int(k[2:]) for k in keys if k.startswith("x_") and k[2:].isdigit())
+                if not idx:
+                    raise ValueError(f"{path}: no 'x' or 'x_0' array")
+                xs = [z[f"x_{k}"] for k in idx]
+                ts = [_times(z, f"t_{k}", xs[i].shape[0]) for i, k in enumerate(idx)]
+                dots = [z[f"x_dot_{k}"] for k in idx] if all(f"x_dot_{k}" in keys for k in idx) else None
+            val = z["val"] if "val" in keys else None
+        variables = meta.get("variables") or [f"x{i}" for i in range(xs[0].shape[1])]
+        meta.setdefault("name", path.stem)
+        return cls(xs, ts, variables, x_dots=dots, val=val, meta=meta)
+
+    @classmethod
+    def from_csv(cls, path, time_column: str = "t") -> "Bundle":
+        """One trajectory per file; a header row names the columns, one of them the time."""
+        path = Path(path)
+        with open(path) as fh:
+            header = [h.strip() for h in fh.readline().split(",")]
+        data = np.loadtxt(path, delimiter=",", skiprows=1, ndmin=2)
+        if time_column not in header:
+            raise ValueError(f"{path}: no '{time_column}' column; columns are {header}")
+        ti = header.index(time_column)
+        cols = [i for i in range(len(header)) if i != ti]
+        return cls([data[:, cols]], [data[:, ti]], [header[i] for i in cols], meta={"name": path.stem})
+
+
+def _times(z, key, n):
+    if key in z.files:
+        return z[key]
+    if "dt" in z.files:
+        return np.arange(n) * float(z["dt"])
+    raise ValueError(f"no '{key}' times and no 'dt'")

mihenk/cli.py

@@ -0,0 +1,185 @@
+"""mihenk command line.
+
+    mihenk scan BUNDLE --model MODEL [--out DIR] [--json] [--set name=value ...] [--daftar]
+    mihenk fit BUNDLE -o MODEL.json [--degree D] [--threshold T] [--diff fd|smoothed]
+    mihenk synth SYSTEM -o BUNDLE.npz [--noise F] [--n-traj K] [--seed S]
+    mihenk validate [--quick]
+    mihenk info
+
+``scan`` exits with status 1 when it finds errors, so it can gate a pipeline.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from contextlib import nullcontext
+from pathlib import Path
+
+from . import __version__, provenance
+from .bundle import Bundle
+from .model import Model
+
+
+def _parse_sets(items):
+    out = {}
+    for item in items or []:
+        if "=" not in item:
+            raise SystemExit(f"--set expects name=value, got {item!r}")
+        k, v = item.split("=", 1)
+        out[k.strip()] = v.strip()
+    return out
+
+
+def cmd_scan(args) -> int:
+    from .scan import Thresholds, scan
+
+    bundle = Bundle.load(args.bundle)
+    model = Model.load(args.model)
+    th = Thresholds().update(_parse_sets(args.set))
+    refit = variants = None
+    cfg = model.meta.get("pysindy_config")
+    if cfg and not args.no_refit:
+        try:
+            from .adapters import pysindy as A
+
+            sindy, _ = A.fit_bundle(bundle, **cfg)
+            refit = A.make_refit(sindy, bundle.variables)
+            variants = A.make_variants(sindy, bundle.variables)
+        except ImportError:
+            print("note: PySINDy is not installed, so tier-1 checks are skipped "
+                  '(pip install "mihenk[pysindy]")', file=sys.stderr)
+    if args.daftar and not provenance.available():
+        print('daftar is not installed; run `pip install "mihenk[provenance]"`', file=sys.stderr)
+        return 2
+    params = {"thresholds": _parse_sets(args.set), "seed": args.seed, "bundle": str(args.bundle),
+              "model": str(args.model)}
+    ctx = provenance.tracked("mihenk scan", params, inputs=[args.bundle, args.model], seed=args.seed) \
+        if args.daftar else nullcontext(None)
+    with ctx as run:
+        report = scan(bundle, model, refit=refit, variants=variants, thresholds=th, seed=args.seed)
+        out = None
+        if args.out:
+            out = report.save(args.out)
+        provenance.log_report(run, report, out)
+    if args.json:
+        print(json.dumps(report.to_dict(), indent=2, default=str))
+    else:
+        print(report.render_text(), end="")
+        if out:
+            print(f"\nreport written to {out}/report.html (also .json, .md, .txt)")
+    return 1 if report.has_errors else 0
+
+
+def cmd_fit(args) -> int:
+    from .adapters import pysindy as A
+
+    bundle = Bundle.load(args.bundle)
+    _, model = A.fit_bundle(bundle, degree=args.degree, threshold=args.threshold, diff=args.diff)
+    model.save(args.output)
+    print(model.pretty())
+    print(f"\nmodel written to {args.output}; next: mihenk scan {args.bundle} --model {args.output}")
+    return 0
+
+
+def cmd_synth(args) -> int:
+    from . import synth
+
+    if args.list or not args.system:
+        for name, s in synth.SYSTEMS.items():
+            print(f"{name:<24} {', '.join(s.variables):<10} {s.note}")
+        return 0
+    b = synth.simulate(args.system, noise=args.noise, n_traj=args.n_traj, dt=args.dt,
+                       duration=args.duration, seed=args.seed)
+    out = args.output or f"{b.name}.npz"
+    b.save(out)
+    print(f"wrote {out}: {b.n_traj} trajectory(ies), {len(b.ts[0])} samples each, variables "
+          f"{', '.join(b.variables)}")
+    print("true law: " + "; ".join(f"{v}' = {e}" for v, e in b.true_equations.items()))
+    return 0
+
+
+def cmd_validate(args) -> int:
+    from . import validate
+
+    if args.seeds > 1:
+        text, ok = validate.sweep(range(args.seed, args.seed + args.seeds), quick=args.quick)
+        print(text)
+        return 0 if ok else 1
+    results = validate.run_all(quick=args.quick, seed=args.seed)
+    print(validate.render(results))
+    return 0 if all(r.passed for r in results) else 1
+
+
+def cmd_info(_args) -> int:
+    print(f"mihenk {__version__}")
+    for mod, why in (("pysindy", 'extra [pysindy]: fit, tier-1 refits'),
+                     ("daftar", "extra [provenance]: scan --daftar")):
+        try:
+            m = __import__(mod)
+            status = getattr(m, "__version__", "installed")
+        except ImportError:
+            status = "not installed"
+        print(f"  {mod:<10} {status:<16} {why}")
+    return 0
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(prog="mihenk", description="A touchstone for discovered equations.")
+    p.add_argument("--version", action="version", version=f"mihenk {__version__}")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    s = sub.add_parser("scan", help="check a discovered equation against its data")
+    s.add_argument("bundle", help="data: .npz bundle or .csv with a 't' column")
+    s.add_argument("--model", required=True, help="equation: .json from `mihenk fit`/Model.save, or a text "
+                                                   "file with lines like  x' = -10 x + 10 y")
+    s.add_argument("--out", help="directory for report.html/.json/.md/.txt")
+    s.add_argument("--json", action="store_true", help="print the report as JSON")
+    s.add_argument("--set", action="append", metavar="NAME=VALUE", help="override a threshold")
+    s.add_argument("--seed", type=int, default=0)
+    s.add_argument("--no-refit", action="store_true", help="skip tier-1 refits even if possible")
+    s.add_argument("--daftar", action="store_true", help="record the scan with daftar")
+    s.set_defaults(fn=cmd_scan)
+
+    f = sub.add_parser("fit", help="fit PySINDy (polynomial library, STLSQ) and save the equation")
+    f.add_argument("bundle")
+    f.add_argument("-o", "--output", default="model.json")
+    f.add_argument("--degree", type=int, default=2)
+    f.add_argument("--threshold", type=float, default=0.1)
+    f.add_argument("--diff", choices=["fd", "smoothed"], default="fd")
+    f.set_defaults(fn=cmd_fit)
+
+    y = sub.add_parser("synth", help="simulate a system with a known law")
+    y.add_argument("system", nargs="?")
+    y.add_argument("-o", "--output")
+    y.add_argument("--noise", type=float, default=0.0, help="noise SD as a fraction of each variable's SD")
+    y.add_argument("--n-traj", type=int, default=1)
+    y.add_argument("--dt", type=float)
+    y.add_argument("--duration", type=float)
+    y.add_argument("--seed", type=int, default=0)
+    y.add_argument("--list", action="store_true")
+    y.set_defaults(fn=cmd_synth)
+
+    v = sub.add_parser("validate", help="run mihenk's self-validation cases")
+    v.add_argument("--quick", action="store_true")
+    v.add_argument("--seed", type=int, default=0)
+    v.add_argument("--seeds", type=int, default=1, help="run over this many seeds and summarize rates")
+    v.set_defaults(fn=cmd_validate)
+
+    i = sub.add_parser("info", help="version and optional dependencies")
+    i.set_defaults(fn=cmd_info)
+    return p
+
+
+def main(argv=None) -> int:
+    args = build_parser().parse_args(argv)
+    try:
+        return args.fn(args)
+    except (ValueError, KeyError, FileNotFoundError, NotImplementedError) as exc:
+        print(f"mihenk: error: {exc}", file=sys.stderr)
+        return 2
+
+
+if __name__ == "__main__":  # pragma: no cover
+    sys.exit(main())