Add Suna synthetic dataset generator and wire CLI support.

dataset_examples/suna_synthetic_metadata.json

@@ -0,0 +1,56 @@
+{
+  "dataset_name": "suna_synthetic",
+  "description": "Synthetic dataset generated for Suna confounder discovery demos. It contains known confounders that influence both treatment assignment and outcomes, an instrumental variable, and nuisance features used to test bootstrap selection heuristics.",
+  "columns": [
+    "treatment",
+    "outcome",
+    "confounder_1",
+    "confounder_2",
+    "instrument",
+    "proxy",
+    "seasonal_component",
+    "linear_trend",
+    "noise"
+  ],
+  "row_count": 2000,
+  "generator": {
+    "n_samples": 2000,
+    "seed": 7,
+    "treatment_intercept": -0.3,
+    "treatment_beta_conf1": 1.2,
+    "treatment_beta_conf2": 0.9,
+    "treatment_beta_instrument": 0.6,
+    "outcome_base": 0.5,
+    "ate": 1.5,
+    "outcome_beta_conf1": 2.0,
+    "outcome_beta_conf2": 1.3,
+    "outcome_beta_instrument": 0.2,
+    "noise_std": 1.0
+  },
+  "checksum": {
+    "sha256_first100": "cb3029805c08577248ff8147146469f67e10dd413e2b2ef20a8acee95a7dc2bd",
+    "note": "Hash computed from the first 100 rows to provide a lightweight fingerprint."
+  },
+  "tags": [
+    "synthetic",
+    "causal",
+    "suna"
+  ],
+  "source": "Programmatically generated by seeker.src.suna.synthetic",
+  "suna": {
+    "treatment_column": "treatment",
+    "outcome_column": "outcome",
+    "candidate_covariates": [
+      "confounder_1",
+      "confounder_2",
+      "instrument",
+      "proxy",
+      "seasonal_component",
+      "linear_trend"
+    ],
+    "tau": 0.05,
+    "n_boot": 100,
+    "default_method": "bcd",
+    "random_state": 7
+  }
+}

seeker/src/query_processor/query_processor.py

@@ -1,13 +1,17 @@
 
 from seeker.src.seeker_service_modules import DataSeeker
+from seeker.src.suna.service import execute_from_cli, SunaDiscoveryResult
 class QueryProcessorClass:
     def __init__(self,dataset_models, query, search_in_metadata):
         self.service_modules = {
             "semantic": self.semantic_search,
             "vector": self.vector_search,
             "cause_and_consequences": self.cause_and_consequences_search,
-            "query_by_example": self.query_by_example_search
+            "query_by_example": self.query_by_example_search,
         }
+        self.service_modules.update({
+            "suna_discover": self.suna_discover_search
+        })
         self.dataset_models = dataset_models
         self.search_query = query
         self.search_in_metadata = search_in_metadata
@@ -42,3 +46,38 @@ def query_by_example_search(self, *args):
         search_query = args[0] if args else self.search_query
         seeker = DataSeeker(search_query)
         print(seeker.query_by_example_search(self.dataset_models, search_query))
+
+    def suna_discover_search(self, *args):
+        try:
+            result: SunaDiscoveryResult = execute_from_cli(
+                self.dataset_models,
+                args,
+                search_query="suna_discover",
+            )
+        except ValueError as exc:
+            print(f"Suna discovery error: {exc}")
+            return
+
+        findings_df = result.to_dataframe()
+        if findings_df.empty:
+            print(
+                f"Suna discovery completed for dataset '{result.dataset_name}' "
+                "with no confounders selected."
+            )
+        else:
+            display_cols = [
+                "confounder",
+                "mi_drop",
+                "bootstrap_quantile",
+                "bootstrap_mean",
+            ]
+            print("Suna discovery findings:")
+            print(findings_df[display_cols].to_string(index=False))
+        if result.ate is not None:
+            print(
+                f"Estimated ATE of treatment '{result.treatment}' on outcome "
+                f"'{result.outcome}': {result.ate:.4f} "
+                f"(std err: {result.ate_std_err:.4f})"
+                if result.ate_std_err is not None
+                else f"{result.ate:.4f}"
+            )

seeker/src/scripts/run_suna_discovery.py

@@ -0,0 +1,151 @@
+import argparse
+import json
+from pathlib import Path
+
+from seeker.src.metadata_dataset_separation.data_import import DataLoader
+from seeker.src.suna.service import SunaConfounderDiscoveryService
+from seeker.src.suna.synthetic import (
+    SYNTHETIC_DATASET_NAME,
+    SyntheticDatasetSpec,
+    ensure_synthetic_dataset,
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run Suna confounder discovery on a dataset loaded by SEEKER.",
+    )
+    parser.add_argument(
+        "--data-dir",
+        default="dataset_examples",
+        help="Directory containing CSV datasets and metadata JSON files.",
+    )
+    parser.add_argument(
+        "--dataset",
+        required=True,
+        help="Dataset name (without .csv) to run discovery on.",
+    )
+    parser.add_argument(
+        "--tau",
+        type=float,
+        default=None,
+        help="Override quantile threshold tau (default from metadata).",
+    )
+    parser.add_argument(
+        "--n-samples",
+        type=int,
+        default=None,
+        help="Override number of bootstrap samples.",
+    )
+    parser.add_argument(
+        "--disable-cache",
+        action="store_true",
+        help="Disable local sketch caching for this run.",
+    )
+    parser.add_argument(
+        "--generate-synthetic",
+        action="store_true",
+        help=(
+            "Regenerate the synthetic Suna demo dataset before running discovery. "
+            "Only applies when --dataset is set to 'suna_synthetic'."
+        ),
+    )
+    parser.add_argument(
+        "--synthetic-samples",
+        type=int,
+        default=None,
+        help=(
+            "Override the number of rows used when generating the synthetic "
+            "dataset (requires --dataset suna_synthetic)."
+        ),
+    )
+    parser.add_argument(
+        "--synthetic-seed",
+        type=int,
+        default=None,
+        help=(
+            "Override the RNG seed used for the synthetic dataset generation "
+            "(requires --dataset suna_synthetic)."
+        ),
+    )
+    return parser.parse_args()
+
+
+def build_overrides(args: argparse.Namespace):
+    overrides = {}
+    if args.tau is not None:
+        overrides["bootstrap.tau"] = args.tau
+    if args.n_samples is not None:
+        overrides["bootstrap.n_samples"] = args.n_samples
+    if args.disable_cache:
+        overrides["sketch.enable_cache"] = False
+    return overrides
+
+
+def main():
+    args = parse_args()
+    data_loader = DataLoader()
+
+    if args.dataset == SYNTHETIC_DATASET_NAME:
+        spec_kwargs = {}
+        if args.synthetic_samples is not None:
+            spec_kwargs["n_samples"] = args.synthetic_samples
+        if args.synthetic_seed is not None:
+            spec_kwargs["seed"] = args.synthetic_seed
+        spec = SyntheticDatasetSpec(**spec_kwargs) if spec_kwargs else None
+        ensure_synthetic_dataset(
+            args.data_dir,
+            spec=spec,
+            force=args.generate_synthetic,
+        )
+
+    dataset_models = data_loader.upload_multiple(args.data_dir, include_metadata=True)
+    if args.dataset not in dataset_models:
+        raise ValueError(f"Dataset '{args.dataset}' not found in {args.data_dir}.")
+
+    overrides = build_overrides(args)
+    dataset_model = dataset_models[args.dataset]
+    service = SunaConfounderDiscoveryService.from_dataset(
+        dataset_model,
+        overrides=overrides,
+    )
+    result = service.discover()
+
+    df = result.to_dataframe()
+    if df.empty:
+        print("No confounders selected for the supplied configuration.")
+    else:
+        print(df.to_string(index=False))
+    if result.ate is not None:
+        print(
+            f"\nEstimated ATE ({result.treatment} -> {result.outcome}): "
+            f"{result.ate:.4f}"
+            + (
+                f" ± {result.ate_std_err:.4f}"
+                if result.ate_std_err is not None
+                else ""
+            )
+        )
+
+    report_path = Path("dist") / "suna" / f"{args.dataset}_latest.json"
+    report_path.parent.mkdir(parents=True, exist_ok=True)
+    report_path.write_text(
+        json.dumps(
+            {
+                "dataset": result.dataset_name,
+                "treatment": result.treatment,
+                "outcome": result.outcome,
+                "method": result.method,
+                "findings": df.to_dict(orient="records"),
+                "ate": result.ate,
+                "ate_std_err": result.ate_std_err,
+                "metadata": result.metadata,
+            },
+            indent=2,
+        )
+    )
+    print(f"\nDiscovery report saved to {report_path}")
+
+
+if __name__ == "__main__":
+    main()

seeker/src/suna/README_SUNA.md

@@ -0,0 +1,85 @@
+# Suna Integration Overview
+
+This package brings the Suna confounder discovery algorithm into SEEKER’s service architecture. It mirrors the layout of other features (e.g., `qgram_index`) and exposes a service API for the interpreter, scripts, and notebooks.
+
+## Module Layout
+- `__init__.py` — Re-exports configs, services, and result dataclasses for convenient imports.
+- `config.py` — Dataclasses (`SunaConfig`, `BootstrapConfig`, `SketchConfig`) plus CLI parsing helpers. Supports metadata-driven defaults with override hooks.
+- `encoders.py` — Lightweight categorical encoder utilities and helpers for selecting candidate covariates.
+- `sketches.py` — Sketch preparation and caching helpers. Currently materializes single-table sketches but keeps cache management aligned with the paper’s design.
+- `scoring.py` — Bivariate causal discovery scoring (bootstrap MI deltas) and stubs for optional subspace methods.
+- `discovery.py` — Iterative Algorithm 1 loop that selects confounders based on bootstrap quantiles and MI-drop heuristics; produces typed results.
+- `pipelines.py` — High-level orchestration that wires preprocessing, discovery, and caching for a `DatasetModel`.
+- `results.py` — Dataclasses describing findings, bootstrap summaries, and overall discovery results.
+- `service.py` — SEEKER-facing service + `SunaDataSeeker` adapter used by the interpreter and scripts.
+
+## Usage
+
+### Interpreter Command
+Add the operation to a query plan:
+```
+suna_discover:dataset=<name>:treatment=<col>:outcome=<col>:candidate_covariates=Z1,Z2
+```
+- `dataset` (required): key from the loaded dataset models.
+- `treatment`, `outcome`: required if metadata lacks `suna.treatment_column` / `suna.outcome_column`.
+- Optional overrides include `method`, `candidate_covariates`, `tau`, `bootstrap.n_samples`, `sketch.enable_cache=false`, etc.
+Results print in the terminal and, when available, render via `SearchResultsVisualizer`.
+
+### Script
+Execute the helper script to run discovery offline:
+```bash
+python -m seeker.src.scripts.run_suna_discovery \
+  --dataset my_dataset \
+  --tau 0.1 \
+  --n-samples 200
+```
+The script reads metadata from `dataset_examples/<dataset>.json`, stores discovery reports under `dist/suna/`, and echoes findings to stdout.
+
+## Metadata Schema
+Annotate datasets with an optional `suna` block:
+```json
+{
+  "suna": {
+    "treatment_column": "T",
+    "outcome_column": "O",
+    "candidate_covariates": ["Z1", "Z2"],
+    "join_keys": [["user_id"]],
+    "tau": 0.05,
+    "n_boot": 100,
+    "default_method": "bcd"
+  }
+}
+```
+Overrides supplied via CLI/interpreter take precedence.
+
+### Synthetic Demo Dataset
+
+The Suna prototype relied on synthetic data with known confounders.  The SEEKER
+integration now ships an equivalent generator in `seeker.src.suna.synthetic`.
+Run discovery against the synthetic asset the same way you would target any
+other dataset:
+
+```bash
+python -m seeker.src.scripts.run_suna_discovery \
+  --dataset suna_synthetic \
+  --generate-synthetic \
+  --tau 0.05
+```
+
+- The first invocation materializes `dataset_examples/suna_synthetic.csv` and
+  `dataset_examples/suna_synthetic_metadata.json`.  Subsequent runs reuse the
+  files unless `--generate-synthetic` is specified.
+- Use `--synthetic-samples <n>` and `--synthetic-seed <seed>` to customise the
+  number of rows or the RNG seed while still producing metadata that matches the
+  SEEKER ingestion contract.
+- Columns cover treatment, outcome, two true confounders, an instrumental
+  variable, a proxy, seasonal and trend terms, plus noise—mirroring the
+  structure of the prototype experiments.
+- Metadata for the synthetic dataset already declares the treatment, outcome,
+  candidate covariates, and bootstrap defaults under the `suna` block so the CLI
+  requires no extra overrides.
+
+## Follow-Up Ideas
+- Expand `sketches.py` to build semi-ring sketches across relational joins once SEEKER exposes join graphs.
+- Replace the MI-drop heuristic with a proper conditional MI estimator from the Suna prototype.
+- Wire optional extensions (`subgroup`, `mprp`, `shap`) once corresponding extras are packaged.

seeker/src/suna/__init__.py

@@ -0,0 +1,35 @@
+"""
+Suna confounder discovery package.
+
+This package refactors the standalone Suna prototype into SEEKER’s
+service-oriented architecture. It exposes configuration factories,
+discovery pipelines, and service bindings used by the interpreter.
+"""
+
+from .config import SunaConfig, BootstrapConfig, SketchConfig
+from .service import (
+    SunaConfounderDiscoveryService,
+    SunaDataSeeker,
+    OptionalDependencyUnavailable,
+)
+from .results import SunaDiscoveryResult, ConfounderFinding, BootstrapSummary
+from .synthetic import (
+    SYNTHETIC_DATASET_NAME,
+    SyntheticDatasetSpec,
+    ensure_synthetic_dataset,
+)
+
+__all__ = [
+    "SunaConfig",
+    "BootstrapConfig",
+    "SketchConfig",
+    "SunaConfounderDiscoveryService",
+    "SunaDataSeeker",
+    "OptionalDependencyUnavailable",
+    "SunaDiscoveryResult",
+    "ConfounderFinding",
+    "BootstrapSummary",
+    "SYNTHETIC_DATASET_NAME",
+    "SyntheticDatasetSpec",
+    "ensure_synthetic_dataset",
+]