daftar 0.1.0__py3-none-any.whl

daftar/__init__.py

@@ -0,0 +1,36 @@
+"""daftar -- record what produced each computational result.
+
+    import daftar
+
+    with daftar.track("celegans", params={"dt": 0.025}, seed=42) as run:
+        run.add_input("data/connectome.csv")
+        v = simulate(dt=0.025)
+        run.log_result("mean_rate_hz", float(v.mean()))
+
+Then, from a shell::
+
+    daftar list
+    daftar diff r-4f21ab r-88c07e
+"""
+
+from .__version__ import __version__
+from .diff import Diff, compare_many, diff_manifests, render_diff, render_manifest
+from .manifest import Manifest
+from .run import Run, track, tracked
+from .store import RunStore
+from .sweep import (
+    ReplayPlan, SweepResult, export_bundle, grid, load_bundle, plan_replay, sweep,
+)
+
+__all__ = [
+    "__version__",
+    "track", "tracked", "Run",
+    "Manifest", "RunStore",
+    "diff_manifests", "render_diff", "render_manifest", "Diff", "compare_many",
+    "sweep", "grid", "SweepResult",
+    "plan_replay", "ReplayPlan",
+    "export_bundle", "load_bundle",
+    "adapters",
+]
+
+from . import adapters  # noqa: E402  (needs the names above)

daftar/__version__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

daftar/adapters/__init__.py

@@ -0,0 +1,37 @@
+"""Framework adapters.
+
+Import is lazy on purpose: someone with Jaxley but not MeltingPot installed must
+still be able to ``import daftar``. Nothing here imports a target framework
+at module load.
+"""
+
+from __future__ import annotations
+
+from . import cpm_adapter, jaxley_adapter, meltingpot_adapter
+from .base import Adapter, AdapterRegistry, record_optional, safe
+
+registry = AdapterRegistry()
+registry.register("jaxley", jaxley_adapter)
+registry.register("cpm", cpm_adapter)
+registry.register("meltingpot", meltingpot_adapter)
+
+jaxley = jaxley_adapter
+cpm = cpm_adapter
+meltingpot = meltingpot_adapter
+
+
+def available() -> list[str]:
+    """Adapters whose target framework is importable right now."""
+    return registry.available()
+
+
+def get(name: str):
+    return registry.get(name)
+
+
+__all__ = [
+    "registry", "available", "get",
+    "jaxley", "cpm", "meltingpot",
+    "jaxley_adapter", "cpm_adapter", "meltingpot_adapter",
+    "Adapter", "AdapterRegistry", "safe", "record_optional",
+]

daftar/adapters/base.py

@@ -0,0 +1,91 @@
+"""What an adapter is, and what makes one worth writing.
+
+The core package can already track any Python function. An adapter earns its
+existence only by knowing something domain-specific that a generic tracker
+cannot infer:
+
+* which arguments are *scientifically* meaningful parameters, as opposed to
+  file paths and verbosity flags;
+* where the framework hides state that silently changes results -- a solver
+  tolerance, an estimator choice, a substrate revision;
+* what counts as a comparable scalar result in that field.
+
+An adapter that just calls ``log_params(kwargs)`` is not worth the import. The
+test is whether it records something the researcher would have forgotten.
+
+Adapters must never import their framework at module import time. Someone with
+Jaxley installed but not MeltingPot has to be able to ``import daftar``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+
+from ..run import Run
+
+
+@runtime_checkable
+class Adapter(Protocol):
+    """Structural protocol. Adapters are modules, not classes."""
+
+    name: str
+
+    def is_available(self) -> bool:
+        """True if the target framework is importable."""
+
+    def describe(self, obj: Any, run: Run, prefix: str = "") -> None:
+        """Record everything provenance-relevant about ``obj`` into ``run``."""
+
+
+class AdapterRegistry:
+    def __init__(self) -> None:
+        self._adapters: dict[str, Any] = {}
+
+    def register(self, name: str, module: Any) -> None:
+        self._adapters[name] = module
+
+    def get(self, name: str) -> Any:
+        if name not in self._adapters:
+            raise KeyError(
+                f"no adapter named {name!r}; available: "
+                f"{', '.join(sorted(self._adapters)) or 'none'}"
+            )
+        return self._adapters[name]
+
+    def available(self) -> list[str]:
+        out = []
+        for name, mod in self._adapters.items():
+            try:
+                if mod.is_available():
+                    out.append(name)
+            except Exception:
+                continue
+        return sorted(out)
+
+    def all(self) -> list[str]:
+        return sorted(self._adapters)
+
+
+def safe(fn, default=None):
+    """Run a probe that may fail against an unfamiliar framework version.
+
+    Adapters read private-ish attributes of fast-moving research code. A
+    provenance tool that crashes a four-hour simulation because a framework
+    renamed an attribute has done far more harm than the missing field is worth.
+    Every probe is best-effort; a missing field is recorded as missing.
+    """
+    try:
+        return fn()
+    except Exception:
+        return default
+
+
+def record_optional(run: Run, key: str, fn, *, kind: str = "param") -> None:
+    """Record ``fn()`` under ``key``, or note that it could not be read."""
+    value = safe(fn, default="<unavailable>")
+    if kind == "param":
+        run.log_param(key, value)
+    elif kind == "result":
+        run.log_result(key, value)
+    else:
+        run.manifest.set(f"{kind}.{key}", value)

daftar/adapters/cpm_adapter.py

@@ -0,0 +1,244 @@
+"""cpm (Computational Psychiatry Modelling) adapter.
+
+The cpm paper names the problem this package exists for: models are implemented
+differently across labs with unstated assumptions, and undetected bugs propagate
+as new researchers build on existing implementations. cpm solved model
+*specification*. It did not solve knowing which fit produced which number.
+
+What a generic tracker misses here:
+
+* **Bounds and priors are the model.** Two fits with identical code and data but
+  a learning-rate bound of ``(0, 1)`` versus ``(0, 2)`` are different
+  experiments. cpm keeps these in ``Parameters``, not in the call arguments.
+* **The estimator is a result-changing choice.** Fmin, FminBound, genetic,
+  BADS, and the hierarchical variants converge differently. Which one ran, with
+  which scipy method and tolerances, has to be in the record.
+* **Fits have convergence status per participant.** A group-level parameter
+  computed from 60 fits of which 7 failed to converge is not the same number as
+  one where all 60 converged, and nothing in the output distinguishes them.
+* **Participant count and identifiers.** Dropping two subjects changes every
+  group statistic. The count and a hash of the identifier list are recorded.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from typing import Any
+
+from ..run import Run
+from .base import safe
+
+name = "cpm"
+
+
+def is_available() -> bool:
+    try:
+        import cpm  # noqa: F401
+        return True
+    except Exception:
+        return False
+
+
+# --------------------------------------------------------------------------
+# parameters
+# --------------------------------------------------------------------------
+
+def describe_parameters(parameters: Any, run: Run, prefix: str = "model") -> None:
+    """Record every free parameter's value, bounds, and prior."""
+    names = safe(lambda: list(parameters.keys()), [])
+    run.log_param(f"{prefix}.parameter_names", sorted(names))
+    run.log_param(f"{prefix}.n_parameters", len(names))
+
+    free = safe(lambda: list(parameters.free()), None)
+    if free is not None:
+        run.log_param(f"{prefix}.n_free_parameters", len(free))
+
+    # bounds() returns [[lowers], [uppers]] in cpm.
+    bounds = safe(lambda: parameters.bounds())
+    if bounds is not None:
+        try:
+            lowers, uppers = bounds[0], bounds[1]
+            for i, nm in enumerate(safe(lambda: list(parameters.free()), names)):
+                run.log_param(f"{prefix}.bounds.{nm}", f"[{lowers[i]}, {uppers[i]}]")
+        except Exception:
+            run.log_param(f"{prefix}.bounds", str(bounds))
+
+    for nm in names:
+        value = safe(lambda n=nm: getattr(parameters, n))
+        if value is None:
+            continue
+        run.log_param(f"{prefix}.value.{nm}", safe(lambda v=value: float(v), str(value)))
+        prior = safe(lambda v=value: getattr(v, "prior", None))
+        if prior is not None:
+            run.log_param(
+                f"{prefix}.prior.{nm}",
+                safe(lambda p=prior: getattr(p, "dist", type(p).__name__), str(prior)),
+            )
+            args = safe(lambda v=value: getattr(v, "args", None))
+            if args:
+                run.log_param(f"{prefix}.prior_args.{nm}", args)
+
+
+# --------------------------------------------------------------------------
+# data
+# --------------------------------------------------------------------------
+
+def describe_data(data: Any, run: Run, prefix: str = "data") -> None:
+    """Record shape and a stable hash of participant identifiers.
+
+    Hashing the identifier list rather than storing it keeps subject IDs out of
+    the manifest -- these are behavioural studies and manifests get committed to
+    public repositories -- while still detecting a changed or reordered cohort.
+    """
+    run.log_param(f"{prefix}.type", type(data).__name__)
+
+    n = safe(lambda: len(data))
+    if n is not None:
+        run.log_param(f"{prefix}.n_records", n)
+
+    columns = safe(lambda: sorted(map(str, data.columns)))
+    if columns:
+        run.log_param(f"{prefix}.columns", columns)
+
+    for id_col in ("ppt", "participant", "subject", "id"):
+        ids = safe(lambda c=id_col: sorted(map(str, data[c].unique())))
+        if ids:
+            run.log_param(f"{prefix}.n_participants", len(ids))
+            digest = hashlib.sha256("\n".join(ids).encode()).hexdigest()[:12]
+            run.log_param(f"{prefix}.participant_id_sha256", digest)
+            break
+
+    if safe(lambda: "observed" in data) or safe(lambda: "observed" in data.columns):
+        run.log_param(f"{prefix}.has_observed", True)
+
+
+# --------------------------------------------------------------------------
+# fitting
+# --------------------------------------------------------------------------
+
+def describe_optimiser(optimiser: Any, run: Run, prefix: str = "fit") -> None:
+    """Record which estimator ran and how it was configured."""
+    run.log_param(f"{prefix}.estimator", type(optimiser).__name__)
+    run.log_param(
+        f"{prefix}.loss",
+        safe(lambda: getattr(optimiser.loss, "__name__", str(optimiser.loss)), "unknown"),
+    )
+    run.log_param(f"{prefix}.uses_prior", bool(safe(lambda: optimiser.prior, False)))
+
+    kwargs = safe(lambda: dict(optimiser.kwargs or {}), {})
+    for k, v in sorted(kwargs.items()):
+        run.log_param(f"{prefix}.kwargs.{k}", v)
+
+    # The scipy method and tolerance hide inside kwargs and change the answer.
+    for key in ("method", "tol", "maxiter", "options"):
+        if key in kwargs:
+            run.log_param(f"{prefix}.{key}", kwargs[key])
+
+    model = safe(lambda: optimiser.model)
+    if model is not None:
+        params = safe(lambda: model.parameters)
+        if params is not None:
+            describe_parameters(params, run, prefix="model")
+        data = safe(lambda: model.data)
+        if data is not None:
+            describe_data(data, run, prefix="data")
+
+
+def describe_fit_results(optimiser: Any, run: Run, prefix: str = "fit") -> None:
+    """Record convergence and group-level outcomes after ``optimise()``.
+
+    Convergence counts matter more than they look. A group mean over 60
+    participants of whom 7 hit the iteration limit is a different number from
+    one where all converged, and the output alone does not say which you have.
+    """
+    fits = safe(lambda: list(optimiser.fit), [])
+    run.log_result(f"{prefix}.n_fits", len(fits))
+    if not fits:
+        return
+
+    converged = 0
+    unknown = 0
+    for f in fits:
+        status = None
+        for key in ("success", "converged", "status"):
+            if isinstance(f, dict) and key in f:
+                status = f[key]
+                break
+        if status is None:
+            unknown += 1
+        elif bool(status) is True or status == 0:
+            converged += 1
+    run.log_result(f"{prefix}.n_converged", converged)
+    if unknown:
+        run.log_result(f"{prefix}.n_convergence_unknown", unknown)
+
+    # Group-level central tendency of each fitted parameter.
+    def _summaries():
+        import numpy as np
+        params = safe(lambda: list(optimiser.parameters), [])
+        if not params:
+            return {}
+        keys = set()
+        for p in params:
+            if isinstance(p, dict):
+                keys |= set(p)
+        out = {}
+        for k in sorted(keys):
+            vals = [p[k] for p in params if isinstance(p, dict) and k in p]
+            numeric = [float(v) for v in vals if isinstance(v, (int, float))]
+            if numeric:
+                out[f"group_mean.{k}"] = round(float(np.mean(numeric)), 6)
+                out[f"group_sd.{k}"] = round(float(np.std(numeric)), 6)
+        return out
+
+    for k, v in safe(_summaries, {}).items():
+        run.log_result(f"{prefix}.{k}", v)
+
+    def _loss():
+        import numpy as np
+        vals = []
+        for f in fits:
+            if isinstance(f, dict):
+                for key in ("fun", "loss", "nll", "value"):
+                    if key in f and isinstance(f[key], (int, float)):
+                        vals.append(float(f[key]))
+                        break
+        return round(float(np.mean(vals)), 6) if vals else None
+
+    mean_loss = safe(_loss)
+    if mean_loss is not None:
+        run.log_result(f"{prefix}.mean_loss", mean_loss)
+
+
+def optimise(optimiser: Any, run: Run, **kwargs: Any):
+    """Run ``optimiser.optimise()`` with configuration and outcome recorded.
+
+    ::
+
+        with daftar.track("bandit-fit", seed=7) as run:
+            cpm_adapter.optimise(fmin, run)
+    """
+    import cpm
+
+    run.log_param("fit.cpm_version", safe(lambda: cpm.__version__, "unknown"))
+    describe_optimiser(optimiser, run)
+    result = optimiser.optimise(**kwargs)
+    describe_fit_results(optimiser, run)
+    return result
+
+
+def describe(obj: Any, run: Run, prefix: str = "model") -> None:
+    """Dispatch on whatever cpm object is handed in."""
+    cls = type(obj).__name__
+    if cls in ("Parameters",):
+        describe_parameters(obj, run, prefix)
+    elif cls in ("Wrapper", "Simulator"):
+        params = safe(lambda: obj.parameters)
+        if params is not None:
+            describe_parameters(params, run, prefix)
+        data = safe(lambda: obj.data)
+        if data is not None:
+            describe_data(data, run)
+        run.log_param(f"{prefix}.generator", cls)
+    else:
+        describe_optimiser(obj, run)

daftar/adapters/jaxley_adapter.py

@@ -0,0 +1,176 @@
+"""Jaxley adapter.
+
+What a generic tracker misses about a Jaxley run, and what this records instead:
+
+* ``jx.integrate`` takes ``solver`` and ``voltage_solver`` with *defaults*
+  (``bwd_euler``, ``jaxley.dhs``). Someone who never passes them has no record
+  of them, and Jaxley changing a default between versions silently moves every
+  result. Defaults are recorded explicitly.
+* ``delta_t`` defaults to 0.025 ms and is the single most common cause of a
+  result moving. Recorded whether or not it was passed.
+* The morphology is provenance. A ``Cell`` built from an SWC file, its branch
+  and compartment counts, and which channels are inserted where, all determine
+  the answer and none of it appears in the call arguments.
+* Jaxley is JAX, so x64 mode matters. A run made with
+  ``jax_enable_x64=True`` and one without are different experiments that look
+  identical in any generic log.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+from ..run import Run
+from .base import safe
+
+name = "jaxley"
+
+#: jx.integrate defaults, as of Jaxley 0.6.x. Recorded when not passed
+#: explicitly, so that a future change of default is visible as a diff.
+INTEGRATE_DEFAULTS = {
+    "delta_t": 0.025,
+    "solver": "bwd_euler",
+    "voltage_solver": "jaxley.dhs",
+    "t_max": None,
+    "checkpoint_lengths": None,
+    "return_states": False,
+}
+
+
+def is_available() -> bool:
+    try:
+        import jaxley  # noqa: F401
+        return True
+    except Exception:
+        return False
+
+
+def _jax_config(run: Run) -> None:
+    jax = sys.modules.get("jax")
+    if jax is None:
+        return
+    cfg = getattr(jax, "config", None)
+    if cfg is None:
+        return
+    # x64 silently changes numerical results; platform changes them too.
+    run.manifest.set("env.jax_enable_x64", safe(lambda: cfg.jax_enable_x64, "unknown"))
+    run.manifest.set(
+        "env.jax_platform",
+        safe(lambda: jax.default_backend(), "unknown"),
+    )
+    run.manifest.set(
+        "env.jax_devices",
+        safe(lambda: ",".join(sorted({d.platform for d in jax.devices()})), "unknown"),
+    )
+
+
+def describe_module(module: Any, run: Run, prefix: str = "morphology") -> None:
+    """Record the structure of a ``jx.Module`` (Cell, Network, Branch, Compartment)."""
+    run.log_param(f"{prefix}.type", type(module).__name__)
+
+    # Compartment/branch/cell counts live in the nodes DataFrame.
+    nodes = safe(lambda: module.nodes)
+    if nodes is not None:
+        run.log_param(f"{prefix}.n_compartments", safe(lambda: len(nodes), "unknown"))
+        for col, label in (
+            ("global_branch_index", "n_branches"),
+            ("global_cell_index", "n_cells"),
+        ):
+            run.log_param(
+                f"{prefix}.{label}",
+                safe(lambda c=col: int(nodes[c].nunique()), "unknown"),
+            )
+
+        # Which channels are inserted, and how many compartments carry each.
+        channels = safe(lambda: [c._name for c in module.channels], [])
+        if channels:
+            run.log_param(f"{prefix}.channels", sorted(channels))
+            for ch in sorted(set(channels)):
+                run.log_param(
+                    f"{prefix}.channel_compartments.{ch}",
+                    safe(lambda c=ch: int(nodes[c].sum()) if c in nodes else 0, "unknown"),
+                )
+
+    edges = safe(lambda: module.edges)
+    if edges is not None and len(edges):
+        run.log_param(f"{prefix}.n_synapses", len(edges))
+        run.log_param(
+            f"{prefix}.synapse_types",
+            safe(lambda: sorted(set(edges["type"].tolist())), "unknown"),
+        )
+
+    # Trainable parameters change what a gradient step does.
+    run.log_param(
+        f"{prefix}.n_trainable_params",
+        safe(lambda: int(module.num_trainable_params), 0),
+    )
+    # Externals: stimuli and clamps are inputs, not incidental.
+    externals = safe(lambda: module.externals, {})
+    if externals:
+        run.log_param(f"{prefix}.externals", sorted(externals.keys()))
+        for k, v in sorted(externals.items()):
+            run.log_param(
+                f"{prefix}.external.{k}.shape",
+                safe(lambda vv=v: str(tuple(vv.shape)), "unknown"),
+            )
+
+
+def describe_integration(run: Run, **kwargs: Any) -> dict[str, Any]:
+    """Record integrator settings, filling in Jaxley's defaults explicitly.
+
+    Returns the fully-resolved kwargs so the caller can pass them straight to
+    ``jx.integrate`` -- there is then no way for what was recorded and what was
+    run to drift apart.
+    """
+    resolved = dict(INTEGRATE_DEFAULTS)
+    resolved.update({k: v for k, v in kwargs.items() if v is not None or k in resolved})
+
+    for key, value in resolved.items():
+        run.log_param(f"integrate.{key}", value)
+        if key in INTEGRATE_DEFAULTS and key not in kwargs:
+            run.log_param(f"integrate.{key}.was_default", True)
+
+    _jax_config(run)
+    return resolved
+
+
+def integrate(module: Any, run: Run, **kwargs: Any):
+    """``jx.integrate`` with the module and settings recorded first.
+
+    ::
+
+        with daftar.track("celegans", seed=0) as run:
+            v = jaxley_adapter.integrate(cell, run, t_max=10.0, delta_t=0.025)
+    """
+    import jaxley as jx
+
+    describe_module(module, run)
+    resolved = describe_integration(run, **kwargs)
+    run.log_param("integrate.jaxley_version", safe(lambda: jx.__version__, "unknown"))
+
+    call_kwargs = {k: v for k, v in resolved.items() if v is not None}
+    result = jx.integrate(module, **call_kwargs)
+
+    # Summary statistics, not the trace. A recorded array would make every diff
+    # unreadable; these four numbers are what you actually compare.
+    def _stats():
+        import numpy as np
+        arr = np.asarray(result)
+        return {
+            "shape": str(arr.shape),
+            "v_mean": float(arr.mean()),
+            "v_min": float(arr.min()),
+            "v_max": float(arr.max()),
+            "v_final_mean": float(arr[..., -1].mean()),
+            "n_nonfinite": int((~np.isfinite(arr)).sum()),
+        }
+
+    stats = safe(_stats, {})
+    for k, v in stats.items():
+        run.log_result(f"voltage.{k}", v)
+    return result
+
+
+def describe(obj: Any, run: Run, prefix: str = "morphology") -> None:
+    describe_module(obj, run, prefix)