diff --git a/causalts/algorithms/__init__.py b/causalts/algorithms/__init__.py index 7cf8f92..ee92dd3 100644 --- a/causalts/algorithms/__init__.py +++ b/causalts/algorithms/__init__.py @@ -3,23 +3,55 @@ """Algorithm plugin registry for causal-ts. -Third-party algorithms can register themselves with :func:`register_algorithm` -so they are automatically available in the CLI without editing the package:: +Third-party algorithms register a discovery function under a short name and then +work anywhere a built-in does -- :func:`run_algorithm`, ``list_algorithms()``, and +``causal-ts discover --algorithm ``. + +The function must accept the calling convention :func:`run_algorithm` uses -- +``fn(df=..., ci_test=..., max_lag=..., **kwargs)`` -- and return a +:class:`~causalts.result.CausalResult` subclass:: from causalts import register_algorithm, CausalResult @register_algorithm("my_algo") def run_my_algo(df, ci_test, max_lag, **kwargs): ... - return MyResult(cg_tig=..., var_names=..., _df=df) - -After the import, ``causal-ts discover --algorithm my_algo data.csv`` works. + return MyResult(graph, df, list(df.columns)) + +``@register_algorithm`` only registers within the running process, so it covers +the Python API and any script that imports your module. To reach the installed +``causal-ts`` command -- a separate process that never imports your code -- also +advertise it in your own distribution's ``pyproject.toml``:: + + [project.entry-points."causalts.algorithms"] + my_algo = "my_package.plugin:run_my_algo" + +Entry points are discovered on first use (see +:func:`_load_entry_point_algorithms`), so an installed plugin needs no import on +the caller's side. + +**CLI limitations.** ``causal-ts discover --algorithm `` forwards only +``df``/``ci_test``/``max_lag`` -- CDNOTS/CEDAR/GRACE-specific flags like +``--alpha``, ``--include-c``, or ``--max-degree`` are parsed but not passed +through, since they have no defined meaning for an arbitrary plugin. A plugin +that needs more configuration should read it from its own environment +variable, config file, or a separate CLI, and document that. ``--validate`` +is also not supported for plugins -- the stability bootstrap only knows how +to re-run the built-in algorithms, so it is rejected outright rather than +silently reporting another algorithm's persistence values. """ from __future__ import annotations +import warnings + _ALGO_REGISTRY: dict[str, callable] = {} +#: Entry-point group third-party distributions advertise algorithms under. +ENTRY_POINT_GROUP = "causalts.algorithms" + +_entry_points_loaded = False + # Algorithms whose module is expensive to import (GRACE pulls in # pytorch-lightning) are registered lazily: name -> (module path, function). # ``causalts.grace`` used to self-register at import time; registering here @@ -41,6 +73,44 @@ def _resolve_lazy(name: str): return fn +def _load_entry_point_algorithms() -> None: + """Register algorithms advertised under the :data:`ENTRY_POINT_GROUP`. + + Lets an installed third-party distribution reach the ``causal-ts`` command, + which runs in its own process and never imports the caller's modules. Runs + once, on first use, so plain ``import causalts`` pays nothing until an + algorithm is actually looked up. + + A plugin that fails to import is warned about and skipped -- one broken + third-party package must not make the CLI unusable. Plugins never shadow a + built-in name. + """ + global _entry_points_loaded + if _entry_points_loaded: + return + _entry_points_loaded = True # set first: a raising plugin must not retry forever + + from importlib.metadata import entry_points + + try: + eps = entry_points(group=ENTRY_POINT_GROUP) + except Exception as exc: # pragma: no cover - defensive + warnings.warn(f"could not scan {ENTRY_POINT_GROUP!r} entry points: {exc}") + return + + for ep in eps: + if ep.name in _ALGO_REGISTRY or ep.name in _LAZY_ALGO_REGISTRY: + warnings.warn( + f"algorithm plugin {ep.name!r} clashes with a built-in name and " + f"was ignored" + ) + continue + try: + _ALGO_REGISTRY[ep.name] = ep.load() + except Exception as exc: + warnings.warn(f"could not load algorithm plugin {ep.name!r}: {exc}") + + def register_algorithm(name: str): """Decorator to register a discovery function under *name*. @@ -64,7 +134,12 @@ def decorator(fn): def list_algorithms() -> list[str]: - """Return sorted list of all registered algorithm names.""" + """Return sorted list of all registered algorithm names. + + Includes installed entry-point plugins. ``causalts.cli`` calls this at import + time to build ``--algorithm``'s choices, so plugins appear there too. + """ + _load_entry_point_algorithms() return sorted(set(_ALGO_REGISTRY) | set(_LAZY_ALGO_REGISTRY)) @@ -92,7 +167,7 @@ def run_algorithm(name: str, df, ci_test, max_lag, **kwargs): if name in _LAZY_ALGO_REGISTRY: _resolve_lazy(name) else: - raise ValueError( - f"Unknown algorithm {name!r}. Available: {list_algorithms()}" - ) + _load_entry_point_algorithms() + if name not in _ALGO_REGISTRY: + raise ValueError(f"Unknown algorithm {name!r}. Available: {list_algorithms()}") return _ALGO_REGISTRY[name](df=df, ci_test=ci_test, max_lag=max_lag, **kwargs) diff --git a/causalts/cli.py b/causalts/cli.py index 8d58034..074236e 100644 --- a/causalts/cli.py +++ b/causalts/cli.py @@ -73,6 +73,8 @@ from .algorithms import list_algorithms as _list_algorithms # noqa: E402 ALGORITHM_CHOICES = _list_algorithms() +# Everything else in ALGORITHM_CHOICES is a third-party plugin (see .algorithms). +_BUILTIN_ALGORITHMS = ("cdnots", "cdnots+", "cedar", "grace", "grace-ss") PLOT_FORMAT_CHOICES = ["png", "pdf", "svg"] DATASET_CHOICES = ["ex1", "ex2", "ex3", "henon"] LAG_SEL_CHOICES = ["partial_dcor", "dcor", "dcor_biased", "pearson", "lasso"] @@ -480,6 +482,20 @@ def discover( "check with --algorithm cdnots or cedar.", param_hint="--validate", ) + if do_validate and algorithm not in _BUILTIN_ALGORITHMS: + # The bootstrap re-discovery closure below only special-cases + # cdnots/cdnots+ and otherwise re-runs CEDAR -- correct for the + # built-in fallback (nothing else reaches it), but wrong for a + # plugin: it would silently annotate the plugin's edges with CEDAR's + # persistence values instead of the plugin's own. + raise click.BadParameter( + f"--validate is not supported for third-party algorithm {algorithm!r} " + "(the stability bootstrap only knows how to re-run the built-in " + "algorithms). Re-run the stability check with --algorithm cdnots or " + "cedar, or bootstrap the plugin directly with " + "causalts.bootstrap.temporal_bootstrap.", + param_hint="--validate", + ) if want_pvalues and algorithm in ("grace", "grace-ss"): raise click.BadParameter( "--pvalues is not supported for GRACE (it emits continuous gate " @@ -806,6 +822,28 @@ def discover( f"GRACE-SS complete. Graph shape: {grace_ss_res.cg_tig.shape}, {n_edges} edges, {elapsed:.1f}s", ) + else: + # Third-party algorithm from the plugin registry. The built-ins get + # bespoke branches above because each saves extra artifacts (gate values, + # stability scores, ...); a plugin only has to return a CausalResult. + from .algorithms import run_algorithm + + _log(ctx, f"Running {algorithm} (registered plugin)...") + plugin_res = run_algorithm(algorithm, df=df, ci_test=ci, max_lag=max_lag) + plugin_graph = np.asarray(plugin_res.cg_tig) + np.save(os.path.join(outdir, "estimated_graph.npy"), plugin_graph) + summary["output_files"]["graph"] = "estimated_graph.npy" + + n_edges = int(plugin_graph.astype(bool).sum()) + elapsed = _time.time() - t0 + summary["elapsed_seconds"] = round(elapsed, 1) + summary["n_edges"] = n_edges + _log( + ctx, + f"{algorithm} complete. Graph shape: {plugin_graph.shape}, " + f"{n_edges} edges, {elapsed:.1f}s", + ) + # Named edge list for interpretation (agent-legible). Load the graph back # from disk so this works regardless of which algorithm branch ran. _graph_path = os.path.join(outdir, "estimated_graph.npy") diff --git a/examples/custom_algorithm.ipynb b/examples/custom_algorithm.ipynb index a3b9d1b..d9f40c3 100644 --- a/examples/custom_algorithm.ipynb +++ b/examples/custom_algorithm.ipynb @@ -11,7 +11,8 @@ "\n", "1. **Subclass `CausalResult`** — set `cg_tig`, `var_names`, `_df`, `_scm_cache`.\n", "2. **Decorate with `@register_algorithm(\"name\")`** — makes the algorithm available to\n", - " `list_algorithms()` and the `causal-ts discover --algorithm name` CLI.\n", + " `list_algorithms()` and `run_algorithm(...)` in-process. Add an entry point (Step 6)\n", + " to reach the `causal-ts` command as well.\n", "3. **Return your result object** — you automatically inherit `.plot()`, `.estimate_effect()`,\n", " `.fit_scm()`, `.counterfactual()`, `.attribute_anomaly()`, and all other DoWhy methods.\n", "\n", @@ -326,30 +327,40 @@ "source": [ "## Step 6 — CLI integration\n", "\n", - "Once your module is imported, the CLI picks up your algorithm automatically.\n", - "No edits to causal-ts source code are required.\n", + "`@register_algorithm` registers within the **running process**, so it covers the Python\n", + "API and any script that imports your module:\n", "\n", "```python\n", - "# my_granger_plugin.py\n", - "from causalts import register_algorithm, CausalResult\n", - "...\n", - "@register_algorithm(\"granger\")\n", - "def run_granger(...): ...\n", + "from causalts.algorithms import run_algorithm\n", + "import my_granger_plugin # triggers registration\n", + "\n", + "res = run_algorithm(\"granger\", df=df, ci_test=None, max_lag=3)\n", + "```\n", + "\n", + "The installed `causal-ts` command is a **separate process** that never imports your code,\n", + "so the decorator alone cannot reach it. To make your algorithm available there, package\n", + "it and advertise the `causalts.algorithms` entry point in your own `pyproject.toml`:\n", + "\n", + "```toml\n", + "# my_granger_plugin/pyproject.toml\n", + "[project.entry-points.\"causalts.algorithms\"]\n", + "granger = \"my_granger_plugin.plugin:run_granger\"\n", "```\n", "\n", "```bash\n", - "# In your shell session after importing your module:\n", + "pip install -e . # your plugin package\n", "causal-ts discover data.csv --algorithm granger --max-lag 3\n", "```\n", "\n", - "Or from Python:\n", + "causal-ts discovers installed entry points on first use, so no import is needed on the\n", + "caller's side — `granger` simply appears in `--algorithm`, and `list_algorithms()`\n", + "reports it. A plugin that fails to import is warned about and skipped rather than\n", + "breaking the CLI, and it can never shadow a built-in name.\n", "\n", - "```python\n", - "from causalts.algorithms import run_algorithm\n", - "import my_granger_plugin # triggers registration\n", - "\n", - "res = run_algorithm(\"granger\", df=df, ci_test=None, max_lag=3)\n", - "```" + "> **Signature contract.** The CLI calls your function as\n", + "> `fn(df=..., ci_test=..., max_lag=...)`, so it must accept those keyword names (extra\n", + "> options via `**kwargs`). Ignore `ci_test` if your method doesn't use one, as Granger\n", + "> does above.\n" ] }, { @@ -385,7 +396,7 @@ "- `result.estimate_effect(treatment, outcome, ...)` — DoWhy ATE\n", "- `result.fit_scm()` / `result.counterfactual(...)` — structural causal model\n", "- `result.attribute_anomaly(...)` — Shapley root cause attribution\n", - "- `causal-ts discover --algorithm my_algo data.csv` — CLI\n", + "- `causal-ts discover --algorithm my_algo data.csv` — CLI (once the entry point in Step 6 is declared)\n", "\n", "**Optional overrides:**\n", "- `__repr__` — custom summary string\n", diff --git a/tests/test_algorithm_plugins.py b/tests/test_algorithm_plugins.py new file mode 100644 index 0000000..ddcde6b --- /dev/null +++ b/tests/test_algorithm_plugins.py @@ -0,0 +1,305 @@ +# Copyright 2025 Bloomberg Finance L.P. +# SPDX-License-Identifier: GPL-3.0-or-later + +"""Third-party algorithm plugins: registry, entry points, and CLI dispatch. + +The plugin path exists so someone can add an algorithm without editing the +package. Two halves have to work: registering it (in-process decorator, or an +entry point for the installed ``causal-ts`` command), and actually *running* it +from ``discover`` -- which previously fell through the built-in ``if/elif`` +chain and silently wrote no graph. +""" + +import warnings + +import numpy as np +import pandas as pd +import pytest +from click.testing import CliRunner + +import causalts.algorithms as algorithms +from causalts import CausalResult +from causalts.algorithms import list_algorithms, register_algorithm, run_algorithm +from causalts.cli import main + +BUILTINS = ("cdnots", "cdnots+", "cedar", "grace", "grace-ss") + + +class _StubResult(CausalResult): + def __init__(self, graph, df, var_names): + self.cg_tig = graph + self.var_names = list(var_names) + self._df = df + self._scm_cache = {} + + +def _stub_algorithm(df, ci_test=None, max_lag=2, **kwargs): + """Minimal plugin: one edge at lag 1, so the graph is unmistakably ours.""" + d = df.shape[1] + graph = np.zeros((d, d, max_lag + 1), dtype=np.int8) + graph[0, 1, 1] = 1 + return _StubResult(graph, df, list(df.columns)) + + +class _FakeEntryPoint: + def __init__(self, name, loader): + self.name = name + self.value = "stub:stub" + self._loader = loader + + def load(self): + return self._loader() + + +@pytest.fixture +def data(): + rng = np.random.default_rng(0) + arr = rng.standard_normal((120, 4)) + for t in range(1, 120): + arr[t, 1] += 0.6 * arr[t - 1, 0] + return pd.DataFrame(arr, columns=[f"X{i}" for i in range(4)]) + + +@pytest.fixture +def csv(tmp_path, data): + path = tmp_path / "d.csv" + data.to_csv(path, index=False) + return str(path) + + +@pytest.fixture +def registered(): + """Register a plugin via the decorator, then remove it again.""" + register_algorithm("stubalgo")(_stub_algorithm) + yield "stubalgo" + algorithms._ALGO_REGISTRY.pop("stubalgo", None) + + +@pytest.fixture +def entry_points(monkeypatch): + """Install fake entry points and reset the one-shot load flag.""" + + def _install(eps): + monkeypatch.setattr( + "importlib.metadata.entry_points", lambda group=None: list(eps) + ) + monkeypatch.setattr(algorithms, "_entry_points_loaded", False) + + yield _install + algorithms._entry_points_loaded = False + for name in ("stubalgo", "ep_algo", "boom"): + algorithms._ALGO_REGISTRY.pop(name, None) + + +# ── decorator registration ─────────────────────────────────────────────────── +def test_decorator_registers_and_runs(registered, data): + assert registered in list_algorithms() + res = run_algorithm(registered, df=data, ci_test=None, max_lag=2) + assert isinstance(res, CausalResult) + assert res.cg_tig[0, 1, 1] == 1 + + +def test_unknown_algorithm_raises(): + with pytest.raises(ValueError, match="Unknown algorithm"): + run_algorithm("does-not-exist", df=None, ci_test=None, max_lag=1) + + +# ── entry-point discovery ──────────────────────────────────────────────────── +def test_entry_point_plugin_is_discovered(entry_points, data): + """An installed distribution reaches the registry with no import by the caller.""" + entry_points([_FakeEntryPoint("ep_algo", lambda: _stub_algorithm)]) + assert "ep_algo" in list_algorithms() + assert ( + run_algorithm("ep_algo", df=data, ci_test=None, max_lag=2).cg_tig[0, 1, 1] == 1 + ) + + +def test_broken_plugin_warns_and_is_skipped(entry_points): + """One bad third-party package must not make causal-ts unusable.""" + + def explode(): + raise ImportError("simulated broken plugin") + + entry_points( + [ + _FakeEntryPoint("boom", explode), + _FakeEntryPoint("ep_algo", lambda: _stub_algorithm), + ] + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + names = list_algorithms() + assert any("could not load" in str(w.message) for w in caught) + assert "boom" not in names + assert "ep_algo" in names # the healthy plugin still loaded + assert all(b in names for b in BUILTINS) + + +def test_plugin_cannot_shadow_a_builtin(entry_points): + from causalts.cedar.discovery import run_cedar + + entry_points([_FakeEntryPoint("cedar", lambda: _stub_algorithm)]) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + list_algorithms() + assert any("clashes with a built-in" in str(w.message) for w in caught) + assert algorithms._ALGO_REGISTRY.get("cedar") is run_cedar + + +def test_entry_points_scanned_once(entry_points, monkeypatch): + """Scanning is cached: plain lookups must not re-walk installed metadata.""" + calls = [] + + def counting(group=None): + calls.append(group) + return [] + + monkeypatch.setattr("importlib.metadata.entry_points", counting) + monkeypatch.setattr(algorithms, "_entry_points_loaded", False) + list_algorithms() + list_algorithms() + list_algorithms() + assert len(calls) == 1 + + +# ── CLI dispatch (the regression this file exists for) ─────────────────────── +@pytest.fixture +def cli_with_plugin(entry_points, monkeypatch): + """A CLI whose ``--algorithm`` choices include an entry-point plugin. + + ``causalts.cli`` snapshots ``ALGORITHM_CHOICES = list_algorithms()`` at import + time, so an installed plugin is picked up only because entry points are read + during that first call. Reloading the module here reproduces a fresh + interpreter -- which is what the ``causal-ts`` command actually is -- rather + than pretending a late in-process registration would be visible. + """ + import importlib + + import causalts.cli as cli_module + + entry_points([_FakeEntryPoint("ep_algo", lambda: _stub_algorithm)]) + reloaded = importlib.reload(cli_module) + yield reloaded.main, "ep_algo" + algorithms._entry_points_loaded = False + algorithms._ALGO_REGISTRY.pop("ep_algo", None) + importlib.reload(cli_module) # restore the un-patched choice list + + +def test_cli_offers_installed_plugin(cli_with_plugin): + cli_main, name = cli_with_plugin + result = CliRunner().invoke(cli_main, ["discover", "--help"]) + assert name in result.output + + +def test_cli_runs_plugin_and_writes_a_graph(cli_with_plugin, csv, tmp_path): + """Regression: discover used to fall through the built-in chain, report + success, and write summary.json with no graph at all.""" + cli_main, name = cli_with_plugin + out = tmp_path / "out" + result = CliRunner().invoke( + cli_main, + ["-o", str(out), "-q", "discover", csv, "--algorithm", name, "--max-lag", "2"], + ) + assert result.exit_code == 0, result.output + graphs = list(out.glob("*/estimated_graph.npy")) + assert graphs, "plugin produced no graph (the silent no-op regression)" + g = np.load(graphs[0]) + assert g[0, 1, 1] == 1 + assert int(g.astype(bool).sum()) == 1 + + +def test_cli_plugin_json_has_edges_and_diagnostics(cli_with_plugin, csv, tmp_path): + """The generic post-dispatch block must treat a plugin like any built-in.""" + import json + + cli_main, name = cli_with_plugin + result = CliRunner().invoke( + cli_main, + [ + "-o", + str(tmp_path / "o"), + "-q", + "discover", + csv, + "--algorithm", + name, + "--max-lag", + "2", + "--json", + ], + ) + assert result.exit_code == 0, result.output + payload = json.loads(result.output[result.output.index("{") :]) + assert payload["n_edges"] == 1 + assert payload["output_files"]["graph"] == "estimated_graph.npy" + assert payload["edges"] == [ + {"source": "X0", "target": "X1", "lag": 1, "pvalue": None} + ] + assert payload["diagnostics"]["n_edges"] == 1 + + +def test_cli_still_rejects_unknown_algorithm(csv, tmp_path): + result = CliRunner().invoke( + main, + ["-o", str(tmp_path), "-q", "discover", csv, "--algorithm", "not-an-algo"], + ) + assert result.exit_code != 0 + assert "not-an-algo" in result.output + + +def test_cli_validate_rejects_plugin(cli_with_plugin, csv, tmp_path): + """--validate must not silently annotate a plugin with CEDAR's persistence. + + Regression for a review finding: the bootstrap re-discovery closure only + special-cases cdnots/cdnots+ and otherwise falls through to CEDAR -- fine + for built-ins (nothing else used to reach it), but wrong for a plugin now + that the else-branch makes it reachable there too. + """ + cli_main, name = cli_with_plugin + result = CliRunner().invoke( + cli_main, + [ + "-o", + str(tmp_path), + "-q", + "discover", + csv, + "--algorithm", + name, + "--max-lag", + "2", + "--validate", + ], + ) + assert result.exit_code != 0 + assert "--validate" in result.output + assert name in result.output + assert not list(tmp_path.glob("*/estimated_graph.npy")) + + +@pytest.mark.parametrize("algorithm", ["cdnots", "cedar"]) +def test_cli_builtins_unaffected(algorithm, csv, tmp_path): + """The new else-branch must not capture any built-in.""" + out = tmp_path / algorithm.replace("+", "p") + result = CliRunner().invoke( + main, + [ + "-o", + str(out), + "-q", + "discover", + csv, + "--algorithm", + algorithm, + "--max-lag", + "2", + ], + ) + assert result.exit_code == 0, result.output + assert list(out.glob("*/estimated_graph.npy")) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))