From 013ad8d61392ad10ec1bf6f413f954adbd12e4c9 Mon Sep 17 00:00:00 2001
From: Egor Kraev <egor.kraev@gmail.com>
Date: Thu, 19 Dec 2024 17:59:04 +0100
Subject: [PATCH 1/5] A pretty decent cut of parallelized Ray calc for score
 testing

---
 notebooks/RunExperiments/cluster_config.yaml  |  10 +-
 .../runners/experiment_runner.py              | 398 +++++++++---------
 notebooks/RunExperiments/runners/iv.py        |   8 +-
 notebooks/RunExperiments/runners/kc.py        |   6 +-
 notebooks/RunExperiments/runners/kckp.py      |   6 +-
 notebooks/RunExperiments/runners/rct.py       |   6 +-
 6 files changed, 217 insertions(+), 217 deletions(-)

diff --git a/notebooks/RunExperiments/cluster_config.yaml b/notebooks/RunExperiments/cluster_config.yaml
index 9e4dbbf..d3c9775 100644
--- a/notebooks/RunExperiments/cluster_config.yaml
+++ b/notebooks/RunExperiments/cluster_config.yaml
@@ -6,7 +6,7 @@ cluster_name: default
 
 # The maximum number of workers nodes to launch in addition to the head
 # node.
-max_workers: 2
+max_workers: 20
 
 # The autoscaler will scale up the cluster faster with higher upscaling speed.
 # E.g., if the task requires adding more nodes then autoscaler will gradually
@@ -41,7 +41,7 @@ idle_timeout_minutes: 5
 provider:
     type: gcp
     region: us-west1
-    availability_zone: us-west1-a
+    availability_zone: us-west1-b
     project_id: motleys # Globally unique project id
 
 # How Ray will authenticate with newly launched nodes.
@@ -60,7 +60,7 @@ auth:
 available_node_types:
     ray_head_default:
         # The resources provided by this node type.
-        resources: {"CPU": 2}
+        resources: {"CPU": 0}
         # Provider-specific config for the head node, e.g. instance type. By default
         # Ray will auto-configure unspecified fields such as subnets and ssh-keys.
         # For more documentation on available fields, see:
@@ -93,7 +93,7 @@ available_node_types:
         min_workers: 1
         # The maximum number of worker nodes of this type to launch.
         # This takes precedence over min_workers.
-        max_workers: 5
+        max_workers: 20
         # The resources provided by this node type.
         resources: {"CPU": 2}
         # Provider-specific config for the head node, e.g. instance type. By default
@@ -101,7 +101,7 @@ available_node_types:
         # For more documentation on available fields, see:
         # https://cloud.google.com/compute/docs/reference/rest/v1/instances/insert
         node_config:
-            machineType: n1-standard-2
+            machineType: n2-highmem-2
             disks:
               - boot: true
                 autoDelete: true
diff --git a/notebooks/RunExperiments/runners/experiment_runner.py b/notebooks/RunExperiments/runners/experiment_runner.py
index 3d774e0..282003c 100644
--- a/notebooks/RunExperiments/runners/experiment_runner.py
+++ b/notebooks/RunExperiments/runners/experiment_runner.py
@@ -5,23 +5,27 @@
 import pickle
 import sys
 import warnings
-from datetime import datetime
 from typing import List, Union
+import cloudpickle
 
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+import ray
 from sklearn.model_selection import train_test_split
 
+
+from causaltune import CausalTune
+from causaltune.data_utils import CausalityDataset
+from causaltune.models.passthrough import passthrough_model
+
 # Ensure CausalTune is in the Python path
 root_path = os.path.realpath("../../../../..")  # noqa: E402
 sys.path.append(os.path.join(root_path, "causaltune"))  # noqa: E402
 
 # Import CausalTune and other custom modules after setting up the path
-from causaltune import CausalTune  # noqa: E402
 from causaltune.datasets import load_dataset  # noqa: E402
-from causaltune.models.passthrough import passthrough_model  # noqa: E402
 from causaltune.search.params import SimpleParamService  # noqa: E402
 from causaltune.score.scoring import (
     metrics_to_minimize,  # noqa: E402
@@ -47,9 +51,7 @@ def parse_arguments():
         help="Datasets to use (format: Size Name, e.g., Small Linear_RCT)",
     )
     parser.add_argument("--n_runs", type=int, default=1, help="Number of runs")
-    parser.add_argument(
-        "--num_samples", type=int, default=-1, help="Maximum number of iterations"
-    )
+    parser.add_argument("--num_samples", type=int, default=-1, help="Maximum number of iterations")
 
     parser.add_argument("--outcome_model", type=str, default="nested", help="Outcome model type")
     parser.add_argument(
@@ -90,7 +92,7 @@ def get_estimator_list(dataset_name):
     return [est for est in estimator_list if "Dummy" not in est]
 
 
-def run_experiment(args, dataset_path: str, use_ray: bool = False):
+def run_experiment(args, dataset_path: str, use_ray: bool):
     # Process datasets
     data_sets = {}
     for dataset in args.datasets:
@@ -104,11 +106,7 @@ def run_experiment(args, dataset_path: str, use_ray: bool = False):
         file_path = f"{dataset_path}/{size}/{name}.pkl"
         data_sets[f"{size} {name}"] = load_dataset(file_path)
 
-    if args.timestamp_in_dirname:
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-        out_dir = f"EXPERIMENT_RESULTS_{timestamp}_{args.identifier}"
-    else:
-        out_dir = f"EXPERIMENT_RESULTS_{args.identifier}"
+    out_dir = f"../EXPERIMENT_RESULTS_{args.identifier}"
 
     os.makedirs(out_dir, exist_ok=True)
     out_dir = os.path.realpath(os.path.join(out_dir, size))
@@ -116,18 +114,12 @@ def run_experiment(args, dataset_path: str, use_ray: bool = False):
 
     print(f"Loaded datasets: {list(data_sets.keys())}")
 
-    # Set time budgets properly
-    if args.time_budget is not None and args.components_time_budget is not None:
-        raise ValueError("Please specify either time_budget or components_time_budget, not both.")
-    elif args.time_budget is None and args.components_time_budget is None:
-        args.components_time_budget = 30  # Set default components budget
-
-    # If only time_budget is specified, derive components_time_budget from it
-    if args.time_budget is not None:
-        args.components_time_budget = max(30, args.time_budget / 4)  # Ensure minimum budget
-        args.time_budget = None  # Use only components_time_budget
+    tasks = []
+    out = []
+    i_run = 1
 
     for dataset_name, cd in data_sets.items():
+        estimators = get_estimator_list(dataset_name)
         # Extract case while preserving original string checking logic
         if "KCKP" in dataset_name:
             case = "KCKP"
@@ -139,122 +131,48 @@ def run_experiment(args, dataset_path: str, use_ray: bool = False):
             case = "RCT"
 
         os.makedirs(f"{out_dir}/{case}", exist_ok=True)
-
-        for i_run in range(1, args.n_runs + 1):
-            cd_i = copy.deepcopy(cd)
-            train_df, test_df = train_test_split(cd_i.data, test_size=args.test_size)
-            test_df = test_df.reset_index(drop=True)
-            cd_i.data = train_df
-
-            for metric in args.metrics:
-                if metric == "ate":  # this is not something to optimize
-                    continue
-
-                print(f"Optimizing {metric} for {dataset_name} (run {i_run})")
-                try:
-                    fn = make_filename(metric, dataset_name, i_run)
-                    out_fn = os.path.join(out_dir, case, fn)
-                    if os.path.isfile(out_fn):
-                        print(f"File {out_fn} exists, skipping...")
-                        continue
-
-                    # Set propensity model using string checking like original version
-                    if "KCKP" in dataset_name:
-                        print(f"Using passthrough propensity model for {dataset_name}")
-                        propensity_model = passthrough_model(
-                            cd_i.propensity_modifiers, include_control=False
-                        )
-                    elif "KC" in dataset_name:
-                        print(f"Using auto propensity model for {dataset_name}")
-                        propensity_model = "auto"
-                    else:
-                        print(f"Using dummy propensity model for {dataset_name}")
-                        propensity_model = "dummy"
-
-                    ct = CausalTune(
-                        metric=metric,
-                        estimator_list=get_estimator_list(dataset_name),
-                        num_samples=args.num_samples,
-                        components_time_budget=args.components_time_budget,  # Use this instead
-                        verbose=1,
-                        components_verbose=1,
-                        store_all_estimators=True,
-                        propensity_model=propensity_model,
-                        outcome_model=args.outcome_model,
-                        use_ray=use_ray,
-                    )
-
-                    ct.fit(
-                        data=cd_i,
-                        treatment="treatment",
-                        outcome="outcome",
+        for metric in args.metrics:
+            fn = make_filename(metric, dataset_name, i_run)
+            out_fn = os.path.join(out_dir, case, fn)
+            if os.path.isfile(out_fn):
+                print(f"File {out_fn} exists, skipping...")
+                continue
+            if use_ray:
+                tasks.append(
+                    remote_single_run.remote(
+                        dataset_name,
+                        cd,
+                        metric,
+                        args.test_size,
+                        args.num_samples,
+                        args.components_time_budget,
+                        out_dir,
+                        out_fn,
+                        estimators,
                     )
+                )
+            else:
+                results = single_run(
+                    dataset_name,
+                    cd,
+                    metric,
+                    args.test_size,
+                    args.num_samples,
+                    args.components_time_budget,
+                    out_dir,
+                    out_fn,
+                )
+                out.append(results)
+    if use_ray:
+        out = ray.get(tasks)
 
-                    # Compute scores and save results
-                    results = compute_scores(ct, metric, test_df)
-
-                    with open(out_fn, "wb") as f:
-                        pickle.dump(results, f)
-                except Exception as e:
-                    print(f"Error processing {dataset_name}_{metric}_{i_run}: {e}")
+    for out_fn, results in out:
+        with open(out_fn, "wb") as f:
+            pickle.dump(results, f)
 
     return out_dir
 
 
-def compute_scores(ct, metric, test_df):
-    datasets = {"train": ct.train_df, "validation": ct.test_df, "test": test_df}
-    estimator_scores = {est: [] for est in ct.scores.keys() if "NewDummy" not in est}
-
-    all_scores = []
-    for trial in ct.results.trials:
-        try:
-            estimator_name = trial.last_result["estimator_name"]
-            if "estimator" in trial.last_result and trial.last_result["estimator"]:
-                estimator = trial.last_result["estimator"]
-                scores = {}
-                for ds_name, df in datasets.items():
-                    scores[ds_name] = {}
-                    est_scores = ct.scorer.make_scores(
-                        estimator,
-                        df,
-                        metrics_to_report=ct.metrics_to_report,
-                    )
-                    est_scores["estimator_name"] = estimator_name
-
-                    scores[ds_name]["CATE_estimate"] = np.squeeze(estimator.estimator.effect(df))
-                    scores[ds_name]["CATE_groundtruth"] = np.squeeze(df["true_effect"])
-                    est_scores["MSE"] = np.mean(
-                        (scores[ds_name]["CATE_estimate"] - scores[ds_name]["CATE_groundtruth"])
-                        ** 2
-                    )
-                    scores[ds_name]["scores"] = est_scores
-                scores["optimization_score"] = trial.last_result.get("optimization_score")
-                estimator_scores[estimator_name].append(copy.deepcopy(scores))
-            # Will use this in the nex
-            all_scores.append(scores)
-        except Exception as e:
-            print(f"Error processing trial: {e}")
-
-    for k in estimator_scores.keys():
-        estimator_scores[k] = sorted(
-            estimator_scores[k],
-            key=lambda x: x["validation"]["scores"][metric],
-            reverse=metric not in metrics_to_minimize(),
-        )
-
-    # Debugging: Log final result structure
-    print(f"Returning scores for metric {metric}: Best estimator: {ct.best_estimator}")
-
-    return {
-        "best_estimator": ct.best_estimator,
-        "best_config": ct.best_config,
-        "best_score": ct.best_score,
-        "optimised_metric": metric,
-        "scores_per_estimator": estimator_scores,
-        "all_scores": all_scores,
-    }
-
-
 def extract_metrics_datasets(out_dir: str):
     metrics = set()
     datasets = set()
@@ -333,11 +251,8 @@ def generate_plots(
 
     def plot_grid(title):
         # Use determined problem type instead of hardcoding "backdoor"
-        all_metrics = [
-            m
-            for m in supported_metrics(problem, False, False)
-            if m.lower() != "ate" and m.lower() != "norm_erupt"
-        ]
+        files = os.listdir(out_dir)
+        all_metrics = sorted(list(set([f.split("-")[0] for f in files])))
 
         fig, axs = plt.subplots(
             len(all_metrics), len(datasets), figsize=(20, 5 * len(all_metrics)), dpi=300
@@ -351,7 +266,7 @@ def plot_grid(title):
         # For multiple metrics in args.metrics, use the first one that has a results file
         results_files = {}
         for dataset in datasets:
-            for metric in args.metrics:
+            for metric in all_metrics:
                 filename = make_filename(metric, dataset, 1)
                 filepath = os.path.join(out_dir, filename)
                 if os.path.exists(filepath):
@@ -375,11 +290,7 @@ def plot_grid(title):
                 try:
                     # Find best estimator for this metric
                     best_estimator = None
-                    best_score = (
-                        float("inf")
-                        if metric in metrics_to_minimize()
-                        else float("-inf")
-                    )
+                    best_score = float("inf") if metric in metrics_to_minimize() else float("-inf")
                     estimator_name = None
 
                     for score in results["all_scores"]:
@@ -389,24 +300,16 @@ def plot_grid(title):
                                 if current_score < best_score:
                                     best_score = current_score
                                     best_estimator = score
-                                    estimator_name = score["test"]["scores"][
-                                        "estimator_name"
-                                    ]
+                                    estimator_name = score["test"]["scores"]["estimator_name"]
                             else:
                                 if current_score > best_score:
                                     best_score = current_score
                                     best_estimator = score
-                                    estimator_name = score["test"]["scores"][
-                                        "estimator_name"
-                                    ]
+                                    estimator_name = score["test"]["scores"]["estimator_name"]
 
                     if best_estimator:
-                        CATE_gt = np.array(
-                            best_estimator["test"]["CATE_groundtruth"]
-                        ).flatten()
-                        CATE_est = np.array(
-                            best_estimator["test"]["CATE_estimate"]
-                        ).flatten()
+                        CATE_gt = np.array(best_estimator["test"]["CATE_groundtruth"]).flatten()
+                        CATE_est = np.array(best_estimator["test"]["CATE_estimate"]).flatten()
 
                         # Plotting
                         ax.scatter(CATE_gt, CATE_est, s=40, alpha=0.5)
@@ -445,9 +348,7 @@ def plot_grid(title):
                             )
 
                 except Exception as e:
-                    print(
-                        f"Error processing metric {metric} for dataset {dataset}: {e}"
-                    )
+                    print(f"Error processing metric {metric} for dataset {dataset}: {e}")
                     ax.text(
                         0.5,
                         0.5,
@@ -466,9 +367,7 @@ def plot_grid(title):
                         labelpad=5,  # Reduce padding between label and plot
                     )
                 if i == 0:
-                    ax.set_title(
-                        dataset, fontsize=font_size + 14, fontweight="bold", pad=15
-                    )
+                    ax.set_title(dataset, fontsize=font_size + 14, fontweight="bold", pad=15)
                 ax.set_xticks([])
                 ax.set_yticks([])
 
@@ -488,7 +387,10 @@ def plot_mse_grid(title):
         est_names = sorted(df["estimator_name"].unique())
 
         # Problem type already determined at top level
-        all_metrics = [c for c in df.columns if c in supported_metrics(problem, False, False)and c.lower() != "ate"
+        all_metrics = [
+            c
+            for c in df.columns
+            if c in supported_metrics(problem, False, False) and c.lower() != "ate"
         ]
 
         fig, axs = plt.subplots(
@@ -570,9 +472,7 @@ def plot_mse_grid(title):
                         labelpad=5,
                     )
                 if i == 0:
-                    ax.set_title(
-                        dataset, fontsize=font_size + 14, fontweight="bold", pad=15
-                    )
+                    ax.set_title(dataset, fontsize=font_size + 14, fontweight="bold", pad=15)
 
         plt.suptitle(
             f"MSE vs. Scores: {title}",
@@ -582,6 +482,11 @@ def plot_mse_grid(title):
 
         # Match spacing style with plot_grid
         plt.tight_layout(rect=[0.1, 0, 1, 0.96], h_pad=1.0, w_pad=0.5)
+
+        fig_legend, ax_legend = plt.subplots(figsize=(6, 6))
+        ax_legend.legend(handles=legend_elements, loc="center", fontsize=10)
+        ax_legend.axis("off")
+
         plt.savefig(os.path.join(out_dir, "MSE_grid.pdf"), format="pdf", bbox_inches="tight")
         plt.savefig(os.path.join(out_dir, "MSE_grid.png"), format="png", bbox_inches="tight")
         plt.close()
@@ -600,10 +505,7 @@ def plot_mse_grid(title):
 
 
 def run_batch(
-    identifier: str,
-    kind: str,
-    metrics: List[str],
-    dataset_path: str,
+    identifier: str, kind: str, metrics: List[str], dataset_path: str, use_ray: bool = False
 ):
     args = parse_arguments()
     args.identifier = identifier
@@ -613,54 +515,140 @@ def run_batch(
     args.num_samples = 100
     args.timestamp_in_dirname = False
     args.outcome_model = "auto"  # or use "nested" for the old-style nested model
+    args.components_time_budget = 120
 
-    # os.environ["RAY_ADDRESS"] = "ray://127.0.0.1:8265"
-
-    use_ray = True
     if use_ray:
         import ray
 
         # Assuming we port-mapped already by running ray dashboard
         ray.init(
-            "ray://localhost:10001", runtime_env={"pip": ["causaltune", "catboost"]}
-        )  # "34.82.184.148:6379"
+            "ray://localhost:10001",
+            runtime_env={"working_dir": ".", "pip": ["causaltune", "catboost"]},
+        )
+
     out_dir = run_experiment(args, dataset_path=dataset_path, use_ray=use_ray)
     return out_dir
 
 
-if __name__ == "__main__":
+@ray.remote
+def remote_single_run(*args):
+    return single_run(*args)
 
-    args = parse_arguments()
-    args.identifier = "Egor_test"
-    args.metrics = supported_metrics("backdoor", False, False)
-    # run_experiment assumes we don't mix large and small datasets in the same call
-    args.datasets = ["Large Linear_RCT", "Large NonLinear_RCT"]
-    args.num_samples = 100
-    args.timestamp_in_dirname = False
-    args.outcome_model = "auto"  # or use "nested" for the old-style nested model
 
-    use_ray = True
-    if use_ray:
-        import ray
+def single_run(
+    dataset_name: str,
+    cd: CausalityDataset,
+    metric: str,
+    test_size: float,
+    num_samples: int,
+    components_time_budget: int,
+    out_dir: str,
+    out_fn: str,
+    estimators: List[str],
+    outcome_model: str = "auto",
+    i_run: int = 1,
+):
+
+    cd_i = copy.deepcopy(cd)
+    train_df, test_df = train_test_split(cd_i.data, test_size=test_size)
+    test_df = test_df.reset_index(drop=True)
+    cd_i.data = train_df
+    print(f"Optimizing {metric} for {dataset_name} (run {i_run})")
+    try:
+
+        # Set propensity model using string checking like original version
+        if "KCKP" in dataset_name:
+            print(f"Using passthrough propensity model for {dataset_name}")
+            propensity_model = passthrough_model(cd_i.propensity_modifiers, include_control=False)
+        elif "KC" in dataset_name:
+            print(f"Using auto propensity model for {dataset_name}")
+            propensity_model = "auto"
+        else:
+            print(f"Using dummy propensity model for {dataset_name}")
+            propensity_model = "dummy"
+
+        ct = CausalTune(
+            metric=metric,
+            estimator_list=estimators,
+            num_samples=num_samples,
+            components_time_budget=components_time_budget,  # Use this instead
+            verbose=1,
+            components_verbose=1,
+            store_all_estimators=True,
+            propensity_model=propensity_model,
+            outcome_model=outcome_model,
+            use_ray=False,
+        )
 
-        ray.init()
-    out_dir = run_experiment(args, dataset_path="../RunDatasets", use_ray=use_ray)
+        ct.fit(
+            data=cd_i,
+            treatment="treatment",
+            outcome="outcome",
+        )
 
-    # plot results
-    upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
-    lower_bounds = {"erupt": 0.06, "bite": 0.75}
+        # Embedding this so it ships well to Ray remotes
 
-    # Determine case from datasets
-    if any("IV" in dataset for dataset in args.datasets):
-        case = "IV"
-    elif any("KC" in dataset for dataset in args.datasets):
-        case = "KC"
-    elif any("KCKP" in dataset for dataset in args.datasets):
-        case = "KCKP"
-    else:
-        case = "RCT"
-    # upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
-    # lower_bounds = {"erupt": 0.06, "bite": 0.75}
-    generate_plots(
-        os.path.join(out_dir, case), font_size=8
-    )  # , upper_bounds=upper_bounds, lower_bounds=lower_bounds)
+        def compute_scores(ct, metric, test_df):
+            datasets = {"train": ct.train_df, "validation": ct.test_df, "test": test_df}
+            estimator_scores = {est: [] for est in ct.scores.keys() if "NewDummy" not in est}
+
+            all_scores = []
+            for trial in ct.results.trials:
+                try:
+                    estimator_name = trial.last_result["estimator_name"]
+                    if "estimator" in trial.last_result and trial.last_result["estimator"]:
+                        estimator = trial.last_result["estimator"]
+                        scores = {}
+                        for ds_name, df in datasets.items():
+                            scores[ds_name] = {}
+                            est_scores = ct.scorer.make_scores(
+                                estimator,
+                                df,
+                                metrics_to_report=ct.metrics_to_report,
+                            )
+                            est_scores["estimator_name"] = estimator_name
+
+                            scores[ds_name]["CATE_estimate"] = np.squeeze(
+                                estimator.estimator.effect(df)
+                            )
+                            scores[ds_name]["CATE_groundtruth"] = np.squeeze(df["true_effect"])
+                            est_scores["MSE"] = np.mean(
+                                (
+                                    scores[ds_name]["CATE_estimate"]
+                                    - scores[ds_name]["CATE_groundtruth"]
+                                )
+                                ** 2
+                            )
+                            scores[ds_name]["scores"] = est_scores
+                        scores["optimization_score"] = trial.last_result.get("optimization_score")
+                        estimator_scores[estimator_name].append(copy.deepcopy(scores))
+                    # Will use this in the nex
+                    all_scores.append(scores)
+                except Exception as e:
+                    print(f"Error processing trial: {e}")
+
+            for k in estimator_scores.keys():
+                estimator_scores[k] = sorted(
+                    estimator_scores[k],
+                    key=lambda x: x["validation"]["scores"][metric],
+                    reverse=metric not in metrics_to_minimize(),
+                )
+
+            # Debugging: Log final result structure
+            print(f"Returning scores for metric {metric}: Best estimator: {ct.best_estimator}")
+
+            return {
+                "best_estimator": ct.best_estimator,
+                "best_config": ct.best_config,
+                "best_score": ct.best_score,
+                "optimised_metric": metric,
+                "scores_per_estimator": estimator_scores,
+                "all_scores": all_scores,
+            }
+
+        # Compute scores and save results
+        results = compute_scores(ct, metric, test_df)
+
+        return out_fn, results
+    except Exception as e:
+        print(f"Error processing {dataset_name}_{metric}_{i_run}: {e}")
diff --git a/notebooks/RunExperiments/runners/iv.py b/notebooks/RunExperiments/runners/iv.py
index 8e975e8..fde3434 100644
--- a/notebooks/RunExperiments/runners/iv.py
+++ b/notebooks/RunExperiments/runners/iv.py
@@ -1,12 +1,18 @@
 import os
+import ray
 
 from experiment_runner import run_batch, generate_plots
 
 identifier = "Egor_test"
 kind = "IV"
 metrics = ["energy_distance", "frobenius_norm", "codec"]
+use_ray = False
+remote_function = ray.remote(run_batch)
+calls = []
 
-out_dir = run_batch(identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"))
+out_dir = run_batch(
+    identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"), use_ray=use_ray
+)
 # plot results
 # upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
 # lower_bounds = {"erupt": 0.06, "bite": 0.75}
diff --git a/notebooks/RunExperiments/runners/kc.py b/notebooks/RunExperiments/runners/kc.py
index 3168719..570c575 100644
--- a/notebooks/RunExperiments/runners/kc.py
+++ b/notebooks/RunExperiments/runners/kc.py
@@ -15,8 +15,10 @@
     "codec",  # NEW
     "bite",  # NEW
 ]
-
-out_dir = run_batch(identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"))
+use_ray = True
+out_dir = run_batch(
+    identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"), use_ray=use_ray
+)
 # plot results
 # upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
 # lower_bounds = {"erupt": 0.06, "bite": 0.75}
diff --git a/notebooks/RunExperiments/runners/kckp.py b/notebooks/RunExperiments/runners/kckp.py
index 79955ca..e6ac376 100644
--- a/notebooks/RunExperiments/runners/kckp.py
+++ b/notebooks/RunExperiments/runners/kckp.py
@@ -15,8 +15,10 @@
     "codec",  # NEW
     "bite",  # NEW
 ]
-
-out_dir = run_batch(identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"))
+use_ray = True
+out_dir = run_batch(
+    identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"), use_ray=use_ray
+)
 # plot results
 # upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
 # lower_bounds = {"erupt": 0.06, "bite": 0.75}
diff --git a/notebooks/RunExperiments/runners/rct.py b/notebooks/RunExperiments/runners/rct.py
index 797bc69..efc6afb 100644
--- a/notebooks/RunExperiments/runners/rct.py
+++ b/notebooks/RunExperiments/runners/rct.py
@@ -15,8 +15,10 @@
     "codec",  # NEW
     "bite",  # NEW
 ]
-
-out_dir = run_batch(identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"))
+use_ray = True
+out_dir = run_batch(
+    identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"), use_ray=use_ray
+)
 # plot results
 # upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
 # lower_bounds = {"erupt": 0.06, "bite": 0.75}

From 145e6cd56832f67440ab0b5001837f9c29f6d343 Mon Sep 17 00:00:00 2001
From: whimo <whimox@gmail.com>
Date: Wed, 25 Dec 2024 23:46:28 +0800
Subject: [PATCH 2/5] Simple retry mechanism for ray.get()

---
 .../runners/experiment_runner.py              | 45 ++++++++++++++++++-
 1 file changed, 43 insertions(+), 2 deletions(-)

diff --git a/notebooks/RunExperiments/runners/experiment_runner.py b/notebooks/RunExperiments/runners/experiment_runner.py
index 282003c..7cdd751 100644
--- a/notebooks/RunExperiments/runners/experiment_runner.py
+++ b/notebooks/RunExperiments/runners/experiment_runner.py
@@ -92,7 +92,7 @@ def get_estimator_list(dataset_name):
     return [est for est in estimator_list if "Dummy" not in est]
 
 
-def run_experiment(args, dataset_path: str, use_ray: bool):
+def run_experiment(args, dataset_path: str, use_ray: bool, ray_get_retries: int = 3):
     # Process datasets
     data_sets = {}
     for dataset in args.datasets:
@@ -164,7 +164,27 @@ def run_experiment(args, dataset_path: str, use_ray: bool):
                 )
                 out.append(results)
     if use_ray:
-        out = ray.get(tasks)
+        remaining_tasks = tasks
+        n_fetch_errors = 0
+        out = []
+
+        while remaining_tasks:
+            print(f"Ray: {len(remaining_tasks)} tasks still running...")
+            ready_tasks, remaining_tasks = ray.wait(remaining_tasks, num_returns=1, timeout=5)
+            for ready_task in ready_tasks:
+                print(f"Ray: task ready: {ready_task}")
+                for retry in range(ray_get_retries + 1):
+                    try:
+                        result = ray.get(ready_task)
+                        results.append(result)
+                        break
+                    except Exception as e:
+                        print(
+                            f"Ray: error fetching task {ready_task} result (retry {retry} of {ray_get_retries}): {e}"
+                        )
+                        if retry == ray_get_retries:
+                            print("Ray: error: task result could not be fetched")
+        print(f"Ray: tasks completed with {n_fetch_errors} fetch errors")
 
     for out_fn, results in out:
         with open(out_fn, "wb") as f:
@@ -535,6 +555,27 @@ def remote_single_run(*args):
     return single_run(*args)
 
 
+@ray.remote
+class TaskRunner:
+    def __init__(self):
+        self.futures = {}
+
+    def remote_single_run(self, *args):
+        ref = remote_single_run(*args)
+        self.futures[ref.hex()] = ref
+        return ref.hex()
+
+    def get_results(self):
+        return ray.get(list(self.futures.values()))
+
+    def get_single_result(self, ref_hex):
+        return ray.get(self.futures[ref_hex])
+
+    def is_ready(self, ref_hex):
+        ready, _ = ray.wait([self.futures[ref_hex]], timeout=0)
+        return bool(ready)
+
+
 def single_run(
     dataset_name: str,
     cd: CausalityDataset,

From 4b543492b2779dd2f737220bf6fa2656f078dcd5 Mon Sep 17 00:00:00 2001
From: whimo <whimox@gmail.com>
Date: Thu, 26 Dec 2024 23:09:54 +0800
Subject: [PATCH 3/5] Persistent runner

---
 .../runners/experiment_runner.py              | 160 +++++++++++-------
 1 file changed, 97 insertions(+), 63 deletions(-)

diff --git a/notebooks/RunExperiments/runners/experiment_runner.py b/notebooks/RunExperiments/runners/experiment_runner.py
index 7cdd751..a77af58 100644
--- a/notebooks/RunExperiments/runners/experiment_runner.py
+++ b/notebooks/RunExperiments/runners/experiment_runner.py
@@ -4,6 +4,7 @@
 import os
 import pickle
 import sys
+import time
 import warnings
 from typing import List, Union
 import cloudpickle
@@ -35,6 +36,8 @@
 # Configure warnings
 warnings.filterwarnings("ignore")
 
+RAY_NAMESPACE = "causaltune_experiments"
+
 
 def parse_arguments():
     parser = argparse.ArgumentParser(description="Run CausalTune experiments")
@@ -92,7 +95,7 @@ def get_estimator_list(dataset_name):
     return [est for est in estimator_list if "Dummy" not in est]
 
 
-def run_experiment(args, dataset_path: str, use_ray: bool, ray_get_retries: int = 3):
+def run_experiment(args, dataset_path: str, use_ray: bool):
     # Process datasets
     data_sets = {}
     for dataset in args.datasets:
@@ -114,32 +117,68 @@ def run_experiment(args, dataset_path: str, use_ray: bool, ray_get_retries: int
 
     print(f"Loaded datasets: {list(data_sets.keys())}")
 
-    tasks = []
-    out = []
-    i_run = 1
+    already_running = False
+    if use_ray:
+        try:
+            runner = ray.get_actor("TaskRunner")
+            print("\n" * 4)
+            print(
+                "!!! Found an existing detached TaskRunner. Will assume the tasks have already been submitted."
+            )
+            print(
+                "!!! If you want to re-run the experiments from scratch, "
+                'run ray.kill(ray.get_actor("TaskRunner", namespace="{}")) or recreate the cluster.'.format(
+                    RAY_NAMESPACE
+                )
+            )
+            print("\n" * 4)
+            already_running = True
+        except ValueError:
+            print("Ray: no detached TaskRunner found, creating...")
+            # This thing will be alive even if the host program exits
+            # Must be killed explicitly: ray.kill(ray.get_actor("TaskRunner"))
+            runner = TaskRunner.options(name="TaskRunner", lifetime="detached").remote()
 
-    for dataset_name, cd in data_sets.items():
-        estimators = get_estimator_list(dataset_name)
-        # Extract case while preserving original string checking logic
-        if "KCKP" in dataset_name:
-            case = "KCKP"
-        elif "KC" in dataset_name:
-            case = "KC"
-        elif "IV" in dataset_name:
-            case = "IV"
-        else:
-            case = "RCT"
-
-        os.makedirs(f"{out_dir}/{case}", exist_ok=True)
-        for metric in args.metrics:
-            fn = make_filename(metric, dataset_name, i_run)
-            out_fn = os.path.join(out_dir, case, fn)
-            if os.path.isfile(out_fn):
-                print(f"File {out_fn} exists, skipping...")
-                continue
-            if use_ray:
-                tasks.append(
-                    remote_single_run.remote(
+    out = []
+    if not already_running:
+        tasks = []
+        i_run = 1
+
+        for dataset_name, cd in data_sets.items():
+            estimators = get_estimator_list(dataset_name)
+            # Extract case while preserving original string checking logic
+            if "KCKP" in dataset_name:
+                case = "KCKP"
+            elif "KC" in dataset_name:
+                case = "KC"
+            elif "IV" in dataset_name:
+                case = "IV"
+            else:
+                case = "RCT"
+
+            os.makedirs(f"{out_dir}/{case}", exist_ok=True)
+            for metric in args.metrics:
+                fn = make_filename(metric, dataset_name, i_run)
+                out_fn = os.path.join(out_dir, case, fn)
+                if os.path.isfile(out_fn):
+                    print(f"File {out_fn} exists, skipping...")
+                    continue
+                if use_ray:
+                    tasks.append(
+                        runner.remote_single_run.remote(
+                            dataset_name,
+                            cd,
+                            metric,
+                            args.test_size,
+                            args.num_samples,
+                            args.components_time_budget,
+                            out_dir,
+                            out_fn,
+                            estimators,
+                        )
+                    )
+                else:
+                    results = single_run(
                         dataset_name,
                         cd,
                         metric,
@@ -148,48 +187,32 @@ def run_experiment(args, dataset_path: str, use_ray: bool, ray_get_retries: int
                         args.components_time_budget,
                         out_dir,
                         out_fn,
-                        estimators,
                     )
-                )
-            else:
-                results = single_run(
-                    dataset_name,
-                    cd,
-                    metric,
-                    args.test_size,
-                    args.num_samples,
-                    args.components_time_budget,
-                    out_dir,
-                    out_fn,
-                )
-                out.append(results)
+                    out.append(results)
+
     if use_ray:
-        remaining_tasks = tasks
-        n_fetch_errors = 0
-        out = []
-
-        while remaining_tasks:
-            print(f"Ray: {len(remaining_tasks)} tasks still running...")
-            ready_tasks, remaining_tasks = ray.wait(remaining_tasks, num_returns=1, timeout=5)
-            for ready_task in ready_tasks:
-                print(f"Ray: task ready: {ready_task}")
-                for retry in range(ray_get_retries + 1):
-                    try:
-                        result = ray.get(ready_task)
-                        results.append(result)
-                        break
-                    except Exception as e:
-                        print(
-                            f"Ray: error fetching task {ready_task} result (retry {retry} of {ray_get_retries}): {e}"
-                        )
-                        if retry == ray_get_retries:
-                            print("Ray: error: task result could not be fetched")
-        print(f"Ray: tasks completed with {n_fetch_errors} fetch errors")
+        while True:
+            completed, in_progress = ray.get(runner.get_progress.remote())
+            print(f"Ray: {completed}/{completed + in_progress} tasks completed")
+            if not in_progress:
+                print("Ray: all tasks completed!")
+                break
+            time.sleep(10)
+
+        print("Ray: fetching results...")
+        out = ray.get(runner.get_results.remote())
 
     for out_fn, results in out:
         with open(out_fn, "wb") as f:
             pickle.dump(results, f)
 
+    if use_ray:
+        destroy = input("Ray: seems like the results fetched OK. Destroy TaskRunner? ")
+        if destroy.lower().startswith("y"):
+            print("Destroying TaskRunner... ", end="")
+            ray.kill(runner)
+            print("success!")
+
     return out_dir
 
 
@@ -544,6 +567,7 @@ def run_batch(
         ray.init(
             "ray://localhost:10001",
             runtime_env={"working_dir": ".", "pip": ["causaltune", "catboost"]},
+            namespace=RAY_NAMESPACE,
         )
 
     out_dir = run_experiment(args, dataset_path=dataset_path, use_ray=use_ray)
@@ -561,7 +585,7 @@ def __init__(self):
         self.futures = {}
 
     def remote_single_run(self, *args):
-        ref = remote_single_run(*args)
+        ref = remote_single_run.remote(*args)
         self.futures[ref.hex()] = ref
         return ref.hex()
 
@@ -572,9 +596,19 @@ def get_single_result(self, ref_hex):
         return ray.get(self.futures[ref_hex])
 
     def is_ready(self, ref_hex):
-        ready, _ = ray.wait([self.futures[ref_hex]], timeout=0)
+        ready, _ = ray.wait([self.futures[ref_hex]], timeout=0, fetch_local=False)
         return bool(ready)
 
+    def all_tasks_ready(self):
+        _, in_progress = ray.wait(list(self.futures.values()), timeout=0, fetch_local=False)
+        return not bool(in_progress)
+
+    def get_progress(self):
+        completed, in_progress = ray.wait(
+            list(self.futures.values()), num_returns=len(self.futures), timeout=0, fetch_local=False
+        )
+        return len(completed), len(in_progress)
+
 
 def single_run(
     dataset_name: str,

From 06b9028a55ba8cc5afcfde3e1c6f6d67d14777be Mon Sep 17 00:00:00 2001
From: Egor Kraev <egor.kraev@gmail.com>
Date: Mon, 30 Dec 2024 02:04:58 +0100
Subject: [PATCH 4/5] Iterations on scoring. Fixed BITE.

---
 causaltune/remote.py                          |   2 +-
 causaltune/score/bite.py                      | 138 +++++++
 causaltune/score/frobenius.py                 |   0
 causaltune/score/scoring.py                   | 124 ++----
 notebooks/RunExperiments/cluster_config.yaml  |   6 +-
 .../runners/experiment_plots.py               | 349 ++++++++++++++++
 .../runners/experiment_runner.py              | 388 ++----------------
 notebooks/RunExperiments/runners/iv.py        |   3 +-
 notebooks/RunExperiments/runners/kc.py        |   3 +-
 notebooks/RunExperiments/runners/kckp.py      |  11 +-
 .../RunExperiments/runners/kckp_no_meta.py    |  35 ++
 notebooks/RunExperiments/runners/rct.py       |  14 +-
 .../RunExperiments/runners/rct_no_meta.py     |  35 ++
 13 files changed, 649 insertions(+), 459 deletions(-)
 create mode 100644 causaltune/score/bite.py
 create mode 100644 causaltune/score/frobenius.py
 create mode 100644 notebooks/RunExperiments/runners/experiment_plots.py
 create mode 100644 notebooks/RunExperiments/runners/kckp_no_meta.py
 create mode 100644 notebooks/RunExperiments/runners/rct_no_meta.py

diff --git a/causaltune/remote.py b/causaltune/remote.py
index e1c1c1b..ddd79b8 100644
--- a/causaltune/remote.py
+++ b/causaltune/remote.py
@@ -7,6 +7,6 @@ def remote_exec(function, args, use_ray=False):
     else:
         from joblib import Parallel, delayed
 
-        return Parallel(n_jobs=2, backend="threading")(delayed(function)(*args) for i in range(1))[
+        return Parallel(n_jobs=1, backend="threading")(delayed(function)(*args) for i in range(1))[
             0
         ]
diff --git a/causaltune/score/bite.py b/causaltune/score/bite.py
new file mode 100644
index 0000000..2e54158
--- /dev/null
+++ b/causaltune/score/bite.py
@@ -0,0 +1,138 @@
+from typing import List, Optional
+
+import numpy as np
+import pandas as pd
+from scipy.stats import kendalltau
+
+
+def bite(
+    working_df: pd.DataFrame,
+    treatment_name: str,
+    outcome_name: str,
+    N_values: Optional[List[int]] = None,
+) -> float:
+    if N_values is None:
+        N_values = exponential_spacing(10, 100, 20)
+    # Calculate weights with clipping to avoid extremes
+    working_df["weights"] = np.where(
+        working_df[treatment_name] == 1,
+        1 / np.clip(working_df["propensity"], 0.05, 0.95),
+        1 / np.clip(1 - working_df["propensity"], 0.05, 0.95),
+    )
+
+    kendall_tau_values = []
+
+    for N in N_values:
+        iter_df = working_df.copy()
+
+        try:
+            # Ensure enough unique values for binning
+            unique_ites = np.unique(iter_df["estimated_ITE"])
+            if len(unique_ites) < N:
+                continue
+
+            # Create bins
+            iter_df["ITE_bin"] = pd.qcut(
+                iter_df["estimated_ITE"], q=N, labels=False, duplicates="drop"
+            )
+
+            # Compute bin statistics
+            bin_stats = []
+            for bin_idx in iter_df["ITE_bin"].unique():
+                bin_data = iter_df[iter_df["ITE_bin"] == bin_idx]
+
+                # Skip if bin is too small
+                if len(bin_data) < 2:
+                    continue
+
+                naive_est = compute_naive_estimate(bin_data, treatment_name, outcome_name)
+
+                # Only compute average ITE if weights are valid
+                bin_weights = bin_data["weights"].values
+                if bin_weights.sum() > 0 and not np.isnan(naive_est):
+                    try:
+                        avg_est_ite = np.average(bin_data["estimated_ITE"], weights=bin_weights)
+                        bin_stats.append(
+                            {
+                                "ITE_bin": bin_idx,
+                                "naive_estimate": naive_est,
+                                "average_estimated_ITE": avg_est_ite,
+                            }
+                        )
+                    except ZeroDivisionError:
+                        continue
+
+            # Calculate Kendall's Tau if we have enough valid bins
+            bin_stats_df = pd.DataFrame(bin_stats)
+            if len(bin_stats_df) >= 2:
+                tau, _ = kendalltau(
+                    bin_stats_df["naive_estimate"],
+                    bin_stats_df["average_estimated_ITE"],
+                )
+                if not np.isnan(tau):
+                    kendall_tau_values.append(tau)
+
+        except (ValueError, ZeroDivisionError):
+            continue
+
+    # Return final score
+    if len(kendall_tau_values) == 0:
+        return -np.inf  # Return -inf for failed computations
+
+    # top_3_taus = sorted(kendall_tau_values, reverse=True)[:3]
+    return np.mean(kendall_tau_values)
+
+
+def compute_naive_estimate(
+    group_data: pd.DataFrame, treatment_name: str, outcome_name: str
+) -> float:
+    """Compute naive estimate for a group with safeguards against edge cases."""
+    treated = group_data[group_data[treatment_name] == 1]
+    control = group_data[group_data[treatment_name] == 0]
+
+    if len(treated) == 0 or len(control) == 0:
+        return np.nan
+
+    treated_weights = treated["weights"].values
+    control_weights = control["weights"].values
+
+    # Check if weights sum to 0 or if all weights are 0
+    if (
+        treated_weights.sum() == 0
+        or control_weights.sum() == 0
+        or not (treated_weights > 0).any()
+        or not (control_weights > 0).any()
+    ):
+        return np.nan
+
+    # Weighted averages with explicit handling of edge cases
+    try:
+        y1 = np.average(treated[outcome_name], weights=treated_weights)
+        y0 = np.average(control[outcome_name], weights=control_weights)
+        return y1 - y0
+    except ZeroDivisionError:
+        return np.nan
+
+
+def exponential_spacing(start, end, num_points):
+    """
+    Generate approximately exponentially spaced integers between start and end.
+
+    Parameters:
+        start (int): The starting value.
+        end (int): The ending value.
+        num_points (int): Number of integers to generate.
+
+    Returns:
+        list: A list of approximately exponentially spaced integers.
+    """
+    # Use a logarithmic scale for exponential spacing
+    log_start = np.log(start)
+    log_end = np.log(end)
+    log_space = np.linspace(log_start, log_end, num_points)
+
+    # Exponentiate back and round to nearest integers
+    spaced_integers = np.round(np.exp(log_space)).astype(int)
+
+    # Ensure unique integers
+    return list(np.unique(spaced_integers))
diff --git a/causaltune/score/frobenius.py b/causaltune/score/frobenius.py
new file mode 100644
index 0000000..e69de29
diff --git a/causaltune/score/scoring.py b/causaltune/score/scoring.py
index 9aa5cd1..283959f 100644
--- a/causaltune/score/scoring.py
+++ b/causaltune/score/scoring.py
@@ -14,6 +14,7 @@
 from causaltune.score.thompson import thompson_policy, extract_means_stds
 from causaltune.thirdparty.causalml import metrics
 from causaltune.score.erupt import ERUPT
+from causaltune.score.bite import bite
 from causaltune.utils import treatment_values, psw_joint_weights
 
 import dcor
@@ -21,7 +22,6 @@
 from scipy.spatial import distance
 from sklearn.neighbors import NearestNeighbors
 
-from scipy.stats import kendalltau
 
 from sklearn.preprocessing import StandardScaler
 
@@ -142,6 +142,26 @@ def __init__(
                 + self.psw_estimator._observed_common_causes_names,
             )
 
+    def inverse_propensity_score(self, df: pd.DataFrame, clip: float = 0.05) -> np.ndarray:
+        """
+        Calculate the inverse propensity score weights for the given dataframe.
+
+        Args:
+            df (pandas.DataFrame): input dataframe
+            clip (float): clipping value for propensity scores
+        """
+
+        propensity_model = self.psw_estimator.estimator.propensity_model
+        p = propensity_model.predict_proba(
+            df[self.causal_model.get_effect_modifiers() + self.causal_model.get_common_causes()]
+        )
+        treatment = df[self.psw_estimator._treatment_name].values
+        ex_ante_p = p[np.arange(p.shape[0]), treatment]
+
+        psw = 1.0 / np.clip(ex_ante_p, clip, 1 - clip)
+
+        return psw
+
     def ate(self, df: pd.DataFrame) -> tuple:
         """
         Calculate the Average Treatment Effect. Provide naive std estimates in
@@ -308,6 +328,7 @@ def frobenius_norm_score(
 
         # Get data splits and check validity
         Y0X, treatment_name, split_test_by = self._Y0_X_potential_outcomes(estimate, df)
+
         Y0X_1 = Y0X[Y0X[split_test_by] == 1]
         Y0X_0 = Y0X[Y0X[split_test_by] == 0]
 
@@ -1041,8 +1062,6 @@ def bite_score(
         Returns:
             float: The BITE score. Higher values indicate better model performance.
         """
-        if N_values is None:
-            N_values = list(range(10, 21)) + list(range(25, 51, 5)) + list(range(60, 101, 10))
 
         est = estimate.estimator
         treatment_name = est._treatment_name
@@ -1068,102 +1087,9 @@ def bite_score(
         else:
             raise ValueError("Propensity model is not available.")
 
-        # Calculate weights with clipping to avoid extremes
-        working_df["weights"] = np.where(
-            working_df[treatment_name] == 1,
-            1 / np.clip(working_df["propensity"], 0.05, 0.95),
-            1 / np.clip(1 - working_df["propensity"], 0.05, 0.95),
-        )
-
-        kendall_tau_values = []
-
-        def compute_naive_estimate(group_data):
-            """Compute naive estimate for a group with safeguards against edge cases."""
-            treated = group_data[group_data[treatment_name] == 1]
-            control = group_data[group_data[treatment_name] == 0]
-
-            if len(treated) == 0 or len(control) == 0:
-                return np.nan
-
-            treated_weights = treated["weights"].values
-            control_weights = control["weights"].values
-
-            # Check if weights sum to 0 or if all weights are 0
-            if (
-                treated_weights.sum() == 0
-                or control_weights.sum() == 0
-                or not (treated_weights > 0).any()
-                or not (control_weights > 0).any()
-            ):
-                return np.nan
-
-            # Weighted averages with explicit handling of edge cases
-            try:
-                y1 = np.average(treated[outcome_name], weights=treated_weights)
-                y0 = np.average(control[outcome_name], weights=control_weights)
-                return y1 - y0
-            except ZeroDivisionError:
-                return np.nan
-
-        for N in N_values:
-            iter_df = working_df.copy()
-
-            try:
-                # Ensure enough unique values for binning
-                unique_ites = np.unique(iter_df["estimated_ITE"])
-                if len(unique_ites) < N:
-                    continue
-
-                # Create bins
-                iter_df["ITE_bin"] = pd.qcut(
-                    iter_df["estimated_ITE"], q=N, labels=False, duplicates="drop"
-                )
-
-                # Compute bin statistics
-                bin_stats = []
-                for bin_idx in iter_df["ITE_bin"].unique():
-                    bin_data = iter_df[iter_df["ITE_bin"] == bin_idx]
-
-                    # Skip if bin is too small
-                    if len(bin_data) < 2:
-                        continue
-
-                    naive_est = compute_naive_estimate(bin_data)
-
-                    # Only compute average ITE if weights are valid
-                    bin_weights = bin_data["weights"].values
-                    if bin_weights.sum() > 0 and not np.isnan(naive_est):
-                        try:
-                            avg_est_ite = np.average(bin_data["estimated_ITE"], weights=bin_weights)
-                            bin_stats.append(
-                                {
-                                    "ITE_bin": bin_idx,
-                                    "naive_estimate": naive_est,
-                                    "average_estimated_ITE": avg_est_ite,
-                                }
-                            )
-                        except ZeroDivisionError:
-                            continue
-
-                # Calculate Kendall's Tau if we have enough valid bins
-                bin_stats_df = pd.DataFrame(bin_stats)
-                if len(bin_stats_df) >= 2:
-                    tau, _ = kendalltau(
-                        bin_stats_df["naive_estimate"],
-                        bin_stats_df["average_estimated_ITE"],
-                    )
-                    if not np.isnan(tau):
-                        kendall_tau_values.append(tau)
-
-            except (ValueError, ZeroDivisionError):
-                continue
-
-        # Return final score
-        if len(kendall_tau_values) == 0:
-            return -np.inf  # Return -inf for failed computations
-
-        top_3_taus = sorted(kendall_tau_values, reverse=True)[:3]
-        return np.mean(top_3_taus)
+        # Calculate the BITE score
+        bite_score = bite(working_df, treatment_name, outcome_name)
+        return bite_score
 
     def make_scores(
         self,
diff --git a/notebooks/RunExperiments/cluster_config.yaml b/notebooks/RunExperiments/cluster_config.yaml
index d3c9775..e316f6e 100644
--- a/notebooks/RunExperiments/cluster_config.yaml
+++ b/notebooks/RunExperiments/cluster_config.yaml
@@ -6,7 +6,7 @@ cluster_name: default
 
 # The maximum number of workers nodes to launch in addition to the head
 # node.
-max_workers: 20
+max_workers: 9
 
 # The autoscaler will scale up the cluster faster with higher upscaling speed.
 # E.g., if the task requires adding more nodes then autoscaler will gradually
@@ -93,7 +93,7 @@ available_node_types:
         min_workers: 1
         # The maximum number of worker nodes of this type to launch.
         # This takes precedence over min_workers.
-        max_workers: 20
+        max_workers: 9
         # The resources provided by this node type.
         resources: {"CPU": 2}
         # Provider-specific config for the head node, e.g. instance type. By default
@@ -161,7 +161,7 @@ initialization_commands: []
 
 # List of shell commands to run to set up nodes.
 setup_commands:
-    - pip install causaltune catboost
+    - pip install causaltune catboost ray[tune]
 
     # Note: if you're developing Ray, you probably want to create a Docker image that
     # has your Ray repo pre-cloned. Then, you can replace the pip installs
diff --git a/notebooks/RunExperiments/runners/experiment_plots.py b/notebooks/RunExperiments/runners/experiment_plots.py
new file mode 100644
index 0000000..7bc1c6b
--- /dev/null
+++ b/notebooks/RunExperiments/runners/experiment_plots.py
@@ -0,0 +1,349 @@
+import glob
+import os
+import pickle
+from typing import Union, List
+
+import matplotlib
+import numpy as np
+
+import pandas as pd
+from matplotlib import pyplot as plt
+
+from causaltune.score.scoring import metrics_to_minimize, supported_metrics
+
+
+def extract_metrics_datasets(out_dir: str):
+    metrics = set()
+    datasets = set()
+
+    for file in glob.glob(f"{out_dir}/*.pkl"):
+        parts = os.path.basename(file).split("-")
+        metrics.add(parts[0])
+        datasets.add(parts[-1].replace(".pkl", "").replace("_", " "))
+
+    return sorted(list(metrics)), sorted(list(datasets))
+
+
+def make_filename(metric, dataset, i_run):
+    return f"{metric}-run-{i_run}-{dataset.replace(' ', '_')}.pkl"
+
+
+def get_all_test_scores(out_dir, dataset_name):
+    size, ds_type, case = dataset_name.split(" ")
+    all_scores = []
+    for file in glob.glob(f"{out_dir}/*_{ds_type}_{case}.pkl"):
+        with open(file, "rb") as f:
+            results = pickle.load(f)
+            for x in results["all_scores"]:
+                all_scores.append(
+                    {k: v for k, v in x["test"]["scores"].items() if k not in ["values"]}
+                )
+    out = pd.DataFrame(all_scores)
+    return out
+
+
+def generate_plots(
+    out_dir: str,
+    log_scale: Union[List[str], None] = None,
+    upper_bounds: Union[dict, None] = None,
+    lower_bounds: Union[dict, None] = None,
+    font_size=0,
+):
+    if log_scale is None:
+        log_scale = ["energy_distance", "psw_energy_distance", "frobenius_norm"]
+    if upper_bounds is None:
+        upper_bounds = {}  # Use an empty dictionary if None
+    if lower_bounds is None:
+        lower_bounds = {}  # Use an empty dictionary if None
+
+    metrics, datasets = extract_metrics_datasets(out_dir)
+    # Remove 'ate' from metrics
+    metrics = [m for m in metrics if m.lower() not in ["ate", "norm_erupt"]]
+
+    metric_names = {
+        "psw_frobenius_norm": "PSW\nFrobenius\nNorm",
+        "frobenius_norm": "Frobenius\nNorm",
+        "erupt": "ERUPT",
+        "codec": "CODEC",
+        "auc": "AUC",
+        "qini": "Qini",
+        "bite": "BITE",
+        "policy_risk": "Policy\nRisk",
+        "energy_distance": "Energy\nDistance",
+        "psw_energy_distance": "PSW\nEnergy\nDistance",
+        "norm_erupt": "Normalized\nERUPT",
+    }
+
+    colors = (
+        [matplotlib.colors.CSS4_COLORS["black"]]
+        + list(matplotlib.colors.TABLEAU_COLORS)
+        + [
+            matplotlib.colors.CSS4_COLORS["lime"],
+            matplotlib.colors.CSS4_COLORS["yellow"],
+            matplotlib.colors.CSS4_COLORS["pink"],
+        ]
+    )
+    markers = ["o", "s", "D", "^", "v", "<", ">", "P", "*", "h", "X", "|", "_", "8"]
+
+    # Determine the problem type from the dataset name
+    problem = "iv" if any("IV" in dataset for dataset in datasets) else "backdoor"
+
+    def plot_grid(title):
+        # Use determined problem type instead of hardcoding "backdoor"
+        # files = os.listdir(out_dir)
+        all_metrics = metrics  # sorted(list(set([f.split("-")[0] for f in files])))
+        if "psw_energy_distance" in all_metrics and "energy_distance" in all_metrics:
+            all_metrics.remove("energy_distance")
+
+        fig, axs = plt.subplots(
+            len(all_metrics), len(datasets), figsize=(20, 5 * len(all_metrics)), dpi=300
+        )
+
+        if len(all_metrics) == 1 and len(datasets) == 1:
+            axs = np.array([[axs]])
+        elif len(all_metrics) == 1 or len(datasets) == 1:
+            axs = axs.reshape(-1, 1) if len(datasets) == 1 else axs.reshape(1, -1)
+
+        # For multiple metrics in args.metrics, use the first one that has a results file
+        results_files = {}
+        for dataset in datasets:
+            for metric in all_metrics:
+                filename = make_filename(metric, dataset, 1)
+                filepath = os.path.join(out_dir, filename)
+                if os.path.exists(filepath):
+                    results_files[dataset] = filepath
+                    break
+            if dataset not in results_files:
+                print(f"No results file found for dataset {dataset}")
+
+        for j, dataset in enumerate(datasets):
+            if dataset not in results_files:
+                continue
+
+            with open(results_files[dataset], "rb") as f:
+                results = pickle.load(f)
+
+            print(f"Loading results for Dataset: {dataset}")
+
+            for i, metric in enumerate(all_metrics):
+                ax = axs[i, j]
+
+                try:
+                    # Find best estimator for this metric
+                    best_estimator = None
+                    best_score = float("inf") if metric in metrics_to_minimize() else float("-inf")
+                    estimator_name = None
+
+                    for score in results["all_scores"]:
+                        if "test" in score and metric in score["test"]["scores"]:
+                            current_score = score["test"]["scores"][metric]
+                            if metric in metrics_to_minimize():
+                                if current_score < best_score:
+                                    best_score = current_score
+                                    best_estimator = score
+                                    estimator_name = score["test"]["scores"]["estimator_name"]
+                            else:
+                                if current_score > best_score:
+                                    best_score = current_score
+                                    best_estimator = score
+                                    estimator_name = score["test"]["scores"]["estimator_name"]
+
+                    if best_estimator:
+                        CATE_gt = np.array(best_estimator["test"]["CATE_groundtruth"]).flatten()
+                        CATE_est = np.array(best_estimator["test"]["CATE_estimate"]).flatten()
+
+                        # Plotting
+                        ax.scatter(CATE_gt, CATE_est, s=40, alpha=0.5)
+                        ax.plot(
+                            [min(CATE_gt), max(CATE_gt)],
+                            [min(CATE_gt), max(CATE_gt)],
+                            "k-",
+                            linewidth=1.0,
+                        )
+
+                        # Calculate correlation coefficient
+                        corr = np.corrcoef(CATE_gt, CATE_est)[0, 1]
+
+                        # Add correlation
+                        ax.text(
+                            0.05,
+                            0.95,
+                            f"Corr: {corr:.2f}",
+                            transform=ax.transAxes,
+                            verticalalignment="top",
+                            fontsize=font_size + 12,
+                            fontweight="bold",
+                        )
+
+                        # Add estimator name at bottom center
+                        if estimator_name:
+                            estimator_base = estimator_name.split(".")[-1]
+                            ax.text(
+                                0.5,
+                                0.02,
+                                estimator_base,
+                                transform=ax.transAxes,
+                                horizontalalignment="center",
+                                color="blue",
+                                fontsize=font_size + 10,
+                            )
+
+                except Exception as e:
+                    print(f"Error processing metric {metric} for dataset {dataset}: {e}")
+                    ax.text(
+                        0.5,
+                        0.5,
+                        "Error processing data",
+                        ha="center",
+                        va="center",
+                        fontsize=font_size + 12,
+                    )
+
+                if j == 0:
+                    # Create tight layout for ylabel
+                    ax.set_ylabel(
+                        metric_names.get(metric, metric),
+                        fontsize=font_size + 12,
+                        fontweight="bold",
+                        labelpad=5,  # Reduce padding between label and plot
+                    )
+                if i == 0:
+                    ax.set_title(dataset, fontsize=font_size + 14, fontweight="bold", pad=15)
+                ax.set_xticks([])
+                ax.set_yticks([])
+
+        plt.suptitle(
+            f"Estimated CATEs vs. True CATEs: {title}",
+            fontsize=font_size + 18,
+            fontweight="bold",
+        )
+        # Adjust spacing between subplots
+        plt.tight_layout(rect=[0.1, 0, 1, 0.96], h_pad=1.0, w_pad=0.5)
+        plt.savefig(os.path.join(out_dir, "CATE_grid.pdf"), format="pdf", bbox_inches="tight")
+        plt.savefig(os.path.join(out_dir, "CATE_grid.png"), format="png", bbox_inches="tight")
+        plt.close()
+
+    def plot_mse_grid(title):
+        df = get_all_test_scores(out_dir, datasets[0])
+        est_names = sorted(df["estimator_name"].unique())
+
+        # Problem type already determined at top level
+        all_metrics = [
+            c
+            for c in df.columns
+            if c in supported_metrics(problem, False, False)
+            and c.lower() not in ["ate", "norm_erupt"]
+        ]
+
+        if "psw_energy_distance" in all_metrics:
+            all_metrics.remove("energy_distance")
+
+        fig, axs = plt.subplots(
+            len(all_metrics), len(datasets), figsize=(20, 5 * len(all_metrics)), dpi=300
+        )
+
+        # Handle single plot cases
+        if len(all_metrics) == 1 and len(datasets) == 1:
+            axs = np.array([[axs]])
+        elif len(all_metrics) == 1 or len(datasets) == 1:
+            axs = axs.reshape(-1, 1) if len(datasets) == 1 else axs.reshape(1, -1)
+
+        legend_elements = []
+        for j, dataset in enumerate(datasets):
+            df = get_all_test_scores(out_dir, dataset)
+            # Apply bounds filtering
+            for m, value in upper_bounds.items():
+                if m in df.columns:
+                    df = df[df[m] < value].copy()
+            for m, value in lower_bounds.items():
+                if m in df.columns:
+                    df = df[df[m] > value].copy()
+
+            for i, metric in enumerate(all_metrics):
+                ax = axs[i, j]
+                this_df = df[["estimator_name", metric, "MSE"]].dropna()
+                this_df = this_df[~np.isinf(this_df[metric].values)]
+
+                if len(this_df):
+                    for idx, est_name in enumerate(est_names):
+                        df_slice = this_df[this_df["estimator_name"] == est_name]
+                        if "Dummy" not in est_name and len(df_slice):
+                            marker = markers[idx % len(markers)]
+                            ax.scatter(
+                                df_slice["MSE"],
+                                df_slice[metric],
+                                color=colors[idx],
+                                s=50,
+                                marker=marker,
+                                linewidths=0.5,
+                            )
+                            if metric not in metrics_to_minimize():
+                                ax.invert_yaxis()
+
+                            trimmed_est_name = est_name.split(".")[-1]
+                            if i == 0 and j == 0:
+                                legend_elements.append(
+                                    plt.Line2D(
+                                        [0],
+                                        [0],
+                                        color=colors[idx],
+                                        marker=marker,
+                                        label=trimmed_est_name,
+                                        linestyle="None",
+                                        markersize=6,
+                                    )
+                                )
+
+                    ax.set_xscale("log")
+                    if metric in log_scale:
+                        ax.set_yscale("log")
+                    ax.grid(True)
+                else:
+                    ax.text(
+                        0.5,
+                        0.5,
+                        "No data",
+                        ha="center",
+                        va="center",
+                        fontsize=font_size + 12,
+                    )
+
+                if j == 0:
+                    # Match ylabel style with plot_grid
+                    ax.set_ylabel(
+                        metric_names.get(metric, metric),
+                        fontsize=font_size + 12,
+                        fontweight="bold",
+                        labelpad=5,
+                    )
+                if i == 0:
+                    ax.set_title(dataset, fontsize=font_size + 14, fontweight="bold", pad=15)
+
+        plt.suptitle(
+            f"MSE vs. Scores: {title}",
+            fontsize=font_size + 18,
+            fontweight="bold",
+        )
+
+        # # Match spacing style with plot_grid
+        # plt.tight_layout(rect=[0.1, 0, 1, 0.96], h_pad=1.0, w_pad=0.5)
+        #
+        # fig_legend, ax_legend = plt.subplots(figsize=(6, 6))
+        # ax_legend.legend(handles=legend_elements, loc="center", fontsize=10)
+        # ax_legend.axis("off")
+
+        plt.savefig(os.path.join(out_dir, "MSE_grid.pdf"), format="pdf", bbox_inches="tight")
+        plt.savefig(os.path.join(out_dir, "MSE_grid.png"), format="png", bbox_inches="tight")
+        plt.close()
+
+        # # Create separate legend
+        # fig_legend, ax_legend = plt.subplots(figsize=(6, 6))
+        # ax_legend.legend(handles=legend_elements, loc="center", fontsize=10)
+        # ax_legend.axis("off")
+        # plt.savefig(os.path.join(out_dir, "MSE_legend.pdf"), format="pdf", bbox_inches="tight")
+        # plt.savefig(os.path.join(out_dir, "MSE_legend.png"), format="png", bbox_inches="tight")
+        # plt.close()
+
+    # Generate plots
+    plot_grid("Experiment Results")
+    plot_mse_grid("Experiment Results")
diff --git a/notebooks/RunExperiments/runners/experiment_runner.py b/notebooks/RunExperiments/runners/experiment_runner.py
index a77af58..04abc29 100644
--- a/notebooks/RunExperiments/runners/experiment_runner.py
+++ b/notebooks/RunExperiments/runners/experiment_runner.py
@@ -1,18 +1,13 @@
 import argparse
 import copy
-import glob
 import os
 import pickle
 import sys
 import time
 import warnings
-from typing import List, Union
-import cloudpickle
+from typing import List, Optional
 
-import matplotlib
-import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 import ray
 from sklearn.model_selection import train_test_split
 
@@ -20,6 +15,7 @@
 from causaltune import CausalTune
 from causaltune.data_utils import CausalityDataset
 from causaltune.models.passthrough import passthrough_model
+from experiment_plots import make_filename
 
 # Ensure CausalTune is in the Python path
 root_path = os.path.realpath("../../../../..")  # noqa: E402
@@ -30,7 +26,7 @@
 from causaltune.search.params import SimpleParamService  # noqa: E402
 from causaltune.score.scoring import (
     metrics_to_minimize,  # noqa: E402
-    supported_metrics,  # noqa: E402
+    # noqa: E402
 )
 
 # Configure warnings
@@ -80,7 +76,14 @@ def parse_arguments():
     return parser.parse_args()
 
 
-def get_estimator_list(dataset_name):
+def get_estimator_list(
+    dataset_name,
+    include_patterns: Optional[List[str]] = None,
+    exclude_patterns: Optional[List[str]] = None,
+):
+    assert (
+        include_patterns is None or exclude_patterns is None
+    ), "Cannot specify both include and exclude patterns"
     if "IV" in dataset_name:
         problem = "iv"
     else:
@@ -92,10 +95,24 @@ def get_estimator_list(dataset_name):
         multivalue=False,
     )
     estimator_list = cfg.estimator_names_from_patterns(problem, "all", 1001)
-    return [est for est in estimator_list if "Dummy" not in est]
 
+    out = [est for est in estimator_list if "Dummy" not in est]
+
+    if include_patterns is not None:
+        out = [est for est in out if any(pat in est for pat in include_patterns)]
+
+    if exclude_patterns is not None:
+        out = [est for est in out if not any(pat in est for pat in exclude_patterns)]
+
+    return out
 
-def run_experiment(args, dataset_path: str, use_ray: bool):
+
+def run_experiment(
+    args,
+    estimators: List[str],
+    dataset_path: str,
+    use_ray: bool,
+):
     # Process datasets
     data_sets = {}
     for dataset in args.datasets:
@@ -110,7 +127,6 @@ def run_experiment(args, dataset_path: str, use_ray: bool):
         data_sets[f"{size} {name}"] = load_dataset(file_path)
 
     out_dir = f"../EXPERIMENT_RESULTS_{args.identifier}"
-
     os.makedirs(out_dir, exist_ok=True)
     out_dir = os.path.realpath(os.path.join(out_dir, size))
     os.makedirs(out_dir, exist_ok=True)
@@ -145,7 +161,7 @@ def run_experiment(args, dataset_path: str, use_ray: bool):
         i_run = 1
 
         for dataset_name, cd in data_sets.items():
-            estimators = get_estimator_list(dataset_name)
+
             # Extract case while preserving original string checking logic
             if "KCKP" in dataset_name:
                 case = "KCKP"
@@ -172,7 +188,6 @@ def run_experiment(args, dataset_path: str, use_ray: bool):
                             args.test_size,
                             args.num_samples,
                             args.components_time_budget,
-                            out_dir,
                             out_fn,
                             estimators,
                         )
@@ -185,8 +200,8 @@ def run_experiment(args, dataset_path: str, use_ray: bool):
                         args.test_size,
                         args.num_samples,
                         args.components_time_budget,
-                        out_dir,
                         out_fn,
+                        estimators,
                     )
                     out.append(results)
 
@@ -216,339 +231,13 @@ def run_experiment(args, dataset_path: str, use_ray: bool):
     return out_dir
 
 
-def extract_metrics_datasets(out_dir: str):
-    metrics = set()
-    datasets = set()
-
-    for file in glob.glob(f"{out_dir}/*.pkl"):
-        parts = os.path.basename(file).split("-")
-        metrics.add(parts[0])
-        datasets.add(parts[-1].replace(".pkl", "").replace("_", " "))
-
-    return sorted(list(metrics)), sorted(list(datasets))
-
-
-def make_filename(metric, dataset, i_run):
-    return f"{metric}-run-{i_run}-{dataset.replace(' ', '_')}.pkl"
-
-
-def get_all_test_scores(out_dir, dataset_name):
-    size, ds_type, case = dataset_name.split(" ")
-    all_scores = []
-    for file in glob.glob(f"{out_dir}/*_{ds_type}_{case}.pkl"):
-        with open(file, "rb") as f:
-            results = pickle.load(f)
-            for x in results["all_scores"]:
-                all_scores.append(
-                    {k: v for k, v in x["test"]["scores"].items() if k not in ["values"]}
-                )
-    out = pd.DataFrame(all_scores)
-    return out
-
-
-def generate_plots(
-    out_dir: str,
-    log_scale: Union[List[str], None] = None,
-    upper_bounds: Union[dict, None] = None,
-    lower_bounds: Union[dict, None] = None,
-    font_size=0,
-):
-    if log_scale is None:
-        log_scale = ["energy_distance", "psw_energy_distance", "frobenius_norm"]
-    if upper_bounds is None:
-        upper_bounds = {}  # Use an empty dictionary if None
-    if lower_bounds is None:
-        lower_bounds = {}  # Use an empty dictionary if None
-
-    metrics, datasets = extract_metrics_datasets(out_dir)
-    # Remove 'ate' from metrics
-    metrics = [m for m in metrics if m.lower() != "ate"]
-
-    metric_names = {
-        "psw_frobenius_norm": "PSW\nFrobenius\nNorm",
-        "frobenius_norm": "Frobenius\nNorm",
-        "erupt": "ERUPT",
-        "codec": "CODEC",
-        "auc": "AUC",
-        "qini": "Qini",
-        "bite": "BITE",
-        "policy_risk": "Policy\nRisk",
-        "energy_distance": "Energy\nDistance",
-        "psw_energy_distance": "PSW\nEnergy\nDistance",
-        "norm_erupt": "Normalized\nERUPT",
-    }
-
-    colors = (
-        [matplotlib.colors.CSS4_COLORS["black"]]
-        + list(matplotlib.colors.TABLEAU_COLORS)
-        + [
-            matplotlib.colors.CSS4_COLORS["lime"],
-            matplotlib.colors.CSS4_COLORS["yellow"],
-            matplotlib.colors.CSS4_COLORS["pink"],
-        ]
-    )
-    markers = ["o", "s", "D", "^", "v", "<", ">", "P", "*", "h", "X", "|", "_", "8"]
-
-    # Determine the problem type from the dataset name
-    problem = "iv" if any("IV" in dataset for dataset in datasets) else "backdoor"
-
-    def plot_grid(title):
-        # Use determined problem type instead of hardcoding "backdoor"
-        files = os.listdir(out_dir)
-        all_metrics = sorted(list(set([f.split("-")[0] for f in files])))
-
-        fig, axs = plt.subplots(
-            len(all_metrics), len(datasets), figsize=(20, 5 * len(all_metrics)), dpi=300
-        )
-
-        if len(all_metrics) == 1 and len(datasets) == 1:
-            axs = np.array([[axs]])
-        elif len(all_metrics) == 1 or len(datasets) == 1:
-            axs = axs.reshape(-1, 1) if len(datasets) == 1 else axs.reshape(1, -1)
-
-        # For multiple metrics in args.metrics, use the first one that has a results file
-        results_files = {}
-        for dataset in datasets:
-            for metric in all_metrics:
-                filename = make_filename(metric, dataset, 1)
-                filepath = os.path.join(out_dir, filename)
-                if os.path.exists(filepath):
-                    results_files[dataset] = filepath
-                    break
-            if dataset not in results_files:
-                print(f"No results file found for dataset {dataset}")
-
-        for j, dataset in enumerate(datasets):
-            if dataset not in results_files:
-                continue
-
-            with open(results_files[dataset], "rb") as f:
-                results = pickle.load(f)
-
-            print(f"Loading results for Dataset: {dataset}")
-
-            for i, metric in enumerate(all_metrics):
-                ax = axs[i, j]
-
-                try:
-                    # Find best estimator for this metric
-                    best_estimator = None
-                    best_score = float("inf") if metric in metrics_to_minimize() else float("-inf")
-                    estimator_name = None
-
-                    for score in results["all_scores"]:
-                        if "test" in score and metric in score["test"]["scores"]:
-                            current_score = score["test"]["scores"][metric]
-                            if metric in metrics_to_minimize():
-                                if current_score < best_score:
-                                    best_score = current_score
-                                    best_estimator = score
-                                    estimator_name = score["test"]["scores"]["estimator_name"]
-                            else:
-                                if current_score > best_score:
-                                    best_score = current_score
-                                    best_estimator = score
-                                    estimator_name = score["test"]["scores"]["estimator_name"]
-
-                    if best_estimator:
-                        CATE_gt = np.array(best_estimator["test"]["CATE_groundtruth"]).flatten()
-                        CATE_est = np.array(best_estimator["test"]["CATE_estimate"]).flatten()
-
-                        # Plotting
-                        ax.scatter(CATE_gt, CATE_est, s=40, alpha=0.5)
-                        ax.plot(
-                            [min(CATE_gt), max(CATE_gt)],
-                            [min(CATE_gt), max(CATE_gt)],
-                            "k-",
-                            linewidth=1.0,
-                        )
-
-                        # Calculate correlation coefficient
-                        corr = np.corrcoef(CATE_gt, CATE_est)[0, 1]
-
-                        # Add correlation
-                        ax.text(
-                            0.05,
-                            0.95,
-                            f"Corr: {corr:.2f}",
-                            transform=ax.transAxes,
-                            verticalalignment="top",
-                            fontsize=font_size + 12,
-                            fontweight="bold",
-                        )
-
-                        # Add estimator name at bottom center
-                        if estimator_name:
-                            estimator_base = estimator_name.split(".")[-1]
-                            ax.text(
-                                0.5,
-                                0.02,
-                                estimator_base,
-                                transform=ax.transAxes,
-                                horizontalalignment="center",
-                                color="blue",
-                                fontsize=font_size + 10,
-                            )
-
-                except Exception as e:
-                    print(f"Error processing metric {metric} for dataset {dataset}: {e}")
-                    ax.text(
-                        0.5,
-                        0.5,
-                        "Error processing data",
-                        ha="center",
-                        va="center",
-                        fontsize=font_size + 12,
-                    )
-
-                if j == 0:
-                    # Create tight layout for ylabel
-                    ax.set_ylabel(
-                        metric_names.get(metric, metric),
-                        fontsize=font_size + 12,
-                        fontweight="bold",
-                        labelpad=5,  # Reduce padding between label and plot
-                    )
-                if i == 0:
-                    ax.set_title(dataset, fontsize=font_size + 14, fontweight="bold", pad=15)
-                ax.set_xticks([])
-                ax.set_yticks([])
-
-        plt.suptitle(
-            f"Estimated CATEs vs. True CATEs: {title}",
-            fontsize=font_size + 18,
-            fontweight="bold",
-        )
-        # Adjust spacing between subplots
-        plt.tight_layout(rect=[0.1, 0, 1, 0.96], h_pad=1.0, w_pad=0.5)
-        plt.savefig(os.path.join(out_dir, "CATE_grid.pdf"), format="pdf", bbox_inches="tight")
-        plt.savefig(os.path.join(out_dir, "CATE_grid.png"), format="png", bbox_inches="tight")
-        plt.close()
-
-    def plot_mse_grid(title):
-        df = get_all_test_scores(out_dir, datasets[0])
-        est_names = sorted(df["estimator_name"].unique())
-
-        # Problem type already determined at top level
-        all_metrics = [
-            c
-            for c in df.columns
-            if c in supported_metrics(problem, False, False) and c.lower() != "ate"
-        ]
-
-        fig, axs = plt.subplots(
-            len(all_metrics), len(datasets), figsize=(20, 5 * len(all_metrics)), dpi=300
-        )
-
-        # Handle single plot cases
-        if len(all_metrics) == 1 and len(datasets) == 1:
-            axs = np.array([[axs]])
-        elif len(all_metrics) == 1 or len(datasets) == 1:
-            axs = axs.reshape(-1, 1) if len(datasets) == 1 else axs.reshape(1, -1)
-
-        legend_elements = []
-        for j, dataset in enumerate(datasets):
-            df = get_all_test_scores(out_dir, dataset)
-            # Apply bounds filtering
-            for m, value in upper_bounds.items():
-                if m in df.columns:
-                    df = df[df[m] < value].copy()
-            for m, value in lower_bounds.items():
-                if m in df.columns:
-                    df = df[df[m] > value].copy()
-
-            for i, metric in enumerate(all_metrics):
-                ax = axs[i, j]
-                this_df = df[["estimator_name", metric, "MSE"]].dropna()
-                this_df = this_df[~np.isinf(this_df[metric].values)]
-
-                if len(this_df):
-                    for idx, est_name in enumerate(est_names):
-                        df_slice = this_df[this_df["estimator_name"] == est_name]
-                        if "Dummy" not in est_name and len(df_slice):
-                            marker = markers[idx % len(markers)]
-                            ax.scatter(
-                                df_slice["MSE"],
-                                df_slice[metric],
-                                color=colors[idx],
-                                s=50,
-                                marker=marker,
-                                linewidths=0.5,
-                            )
-                            if metric not in metrics_to_minimize():
-                                ax.invert_yaxis()
-
-                            trimmed_est_name = est_name.split(".")[-1]
-                            if i == 0 and j == 0:
-                                legend_elements.append(
-                                    plt.Line2D(
-                                        [0],
-                                        [0],
-                                        color=colors[idx],
-                                        marker=marker,
-                                        label=trimmed_est_name,
-                                        linestyle="None",
-                                        markersize=6,
-                                    )
-                                )
-
-                    ax.set_xscale("log")
-                    if metric in log_scale:
-                        ax.set_yscale("log")
-                    ax.grid(True)
-                else:
-                    ax.text(
-                        0.5,
-                        0.5,
-                        "No data",
-                        ha="center",
-                        va="center",
-                        fontsize=font_size + 12,
-                    )
-
-                if j == 0:
-                    # Match ylabel style with plot_grid
-                    ax.set_ylabel(
-                        metric_names.get(metric, metric),
-                        fontsize=font_size + 12,
-                        fontweight="bold",
-                        labelpad=5,
-                    )
-                if i == 0:
-                    ax.set_title(dataset, fontsize=font_size + 14, fontweight="bold", pad=15)
-
-        plt.suptitle(
-            f"MSE vs. Scores: {title}",
-            fontsize=font_size + 18,
-            fontweight="bold",
-        )
-
-        # Match spacing style with plot_grid
-        plt.tight_layout(rect=[0.1, 0, 1, 0.96], h_pad=1.0, w_pad=0.5)
-
-        fig_legend, ax_legend = plt.subplots(figsize=(6, 6))
-        ax_legend.legend(handles=legend_elements, loc="center", fontsize=10)
-        ax_legend.axis("off")
-
-        plt.savefig(os.path.join(out_dir, "MSE_grid.pdf"), format="pdf", bbox_inches="tight")
-        plt.savefig(os.path.join(out_dir, "MSE_grid.png"), format="png", bbox_inches="tight")
-        plt.close()
-
-        # # Create separate legend
-        # fig_legend, ax_legend = plt.subplots(figsize=(6, 6))
-        # ax_legend.legend(handles=legend_elements, loc="center", fontsize=10)
-        # ax_legend.axis("off")
-        # plt.savefig(os.path.join(out_dir, "MSE_legend.pdf"), format="pdf", bbox_inches="tight")
-        # plt.savefig(os.path.join(out_dir, "MSE_legend.png"), format="png", bbox_inches="tight")
-        # plt.close()
-
-    # Generate plots
-    plot_grid("Experiment Results")
-    plot_mse_grid("Experiment Results")
-
-
 def run_batch(
-    identifier: str, kind: str, metrics: List[str], dataset_path: str, use_ray: bool = False
+    identifier: str,
+    kind: str,
+    metrics: List[str],
+    estimators: List[str],
+    dataset_path: str,
+    use_ray: bool = False,
 ):
     args = parse_arguments()
     args.identifier = identifier
@@ -566,11 +255,13 @@ def run_batch(
         # Assuming we port-mapped already by running ray dashboard
         ray.init(
             "ray://localhost:10001",
-            runtime_env={"working_dir": ".", "pip": ["causaltune", "catboost"]},
+            runtime_env={"working_dir": ".", "pip": ["causaltune", "catboost", "ray[tune]"]},
             namespace=RAY_NAMESPACE,
         )
 
-    out_dir = run_experiment(args, dataset_path=dataset_path, use_ray=use_ray)
+    out_dir = run_experiment(
+        args, estimators=estimators, dataset_path=dataset_path, use_ray=use_ray
+    )
     return out_dir
 
 
@@ -617,7 +308,6 @@ def single_run(
     test_size: float,
     num_samples: int,
     components_time_budget: int,
-    out_dir: str,
     out_fn: str,
     estimators: List[str],
     outcome_model: str = "auto",
diff --git a/notebooks/RunExperiments/runners/iv.py b/notebooks/RunExperiments/runners/iv.py
index fde3434..4889f45 100644
--- a/notebooks/RunExperiments/runners/iv.py
+++ b/notebooks/RunExperiments/runners/iv.py
@@ -1,7 +1,8 @@
 import os
 import ray
 
-from experiment_runner import run_batch, generate_plots
+from experiment_runner import run_batch
+from experiment_plots import generate_plots
 
 identifier = "Egor_test"
 kind = "IV"
diff --git a/notebooks/RunExperiments/runners/kc.py b/notebooks/RunExperiments/runners/kc.py
index 570c575..065333e 100644
--- a/notebooks/RunExperiments/runners/kc.py
+++ b/notebooks/RunExperiments/runners/kc.py
@@ -1,6 +1,7 @@
 import os
 
-from experiment_runner import run_batch, generate_plots
+from experiment_runner import run_batch
+from experiment_plots import generate_plots
 
 identifier = "Egor_test"
 kind = "KC"
diff --git a/notebooks/RunExperiments/runners/kckp.py b/notebooks/RunExperiments/runners/kckp.py
index e6ac376..32f5b1d 100644
--- a/notebooks/RunExperiments/runners/kckp.py
+++ b/notebooks/RunExperiments/runners/kckp.py
@@ -1,6 +1,7 @@
 import os
 
-from experiment_runner import run_batch, generate_plots
+from experiment_runner import run_batch, get_estimator_list
+from experiment_plots import generate_plots
 
 identifier = "Egor_test"
 kind = "KCKP"
@@ -16,8 +17,14 @@
     "bite",  # NEW
 ]
 use_ray = True
+estimators = get_estimator_list(kind)
 out_dir = run_batch(
-    identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"), use_ray=use_ray
+    identifier,
+    kind,
+    metrics,
+    estimators=estimators,
+    dataset_path=os.path.realpath("../RunDatasets"),
+    use_ray=use_ray,
 )
 # plot results
 # upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
diff --git a/notebooks/RunExperiments/runners/kckp_no_meta.py b/notebooks/RunExperiments/runners/kckp_no_meta.py
new file mode 100644
index 0000000..13f2bd7
--- /dev/null
+++ b/notebooks/RunExperiments/runners/kckp_no_meta.py
@@ -0,0 +1,35 @@
+import os
+
+from experiment_runner import run_batch, get_estimator_list
+from experiment_plots import generate_plots
+
+identifier = "Egor_test_no_meta"
+kind = "KCKP"
+metrics = [
+    "erupt",
+    # "greedy_erupt",  # regular erupt was made probabilistic,
+    "policy_risk",  # NEW
+    "qini",
+    "auc",
+    "psw_energy_distance",
+    "frobenius_norm",  # NEW
+    "codec",  # NEW
+    "bite",  # NEW
+]
+
+estimators = get_estimator_list(kind, exclude_patterns=["SLearner", "TLearner", "XLearner"])
+
+use_ray = True
+out_dir = run_batch(
+    identifier,
+    kind,
+    metrics,
+    estimators=estimators,
+    dataset_path=os.path.realpath("../RunDatasets"),
+    use_ray=use_ray,
+)
+# plot results
+# upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
+# lower_bounds = {"erupt": 0.06, "bite": 0.75}
+generate_plots(os.path.join(out_dir, kind))  # , upper_bounds, lower_bounds)
+print("yay!")
diff --git a/notebooks/RunExperiments/runners/rct.py b/notebooks/RunExperiments/runners/rct.py
index efc6afb..8aebd69 100644
--- a/notebooks/RunExperiments/runners/rct.py
+++ b/notebooks/RunExperiments/runners/rct.py
@@ -1,9 +1,11 @@
 import os
 
-from experiment_runner import run_batch, generate_plots
+from experiment_runner import run_batch, get_estimator_list
+from experiment_plots import generate_plots
 
 identifier = "Egor_test"
 kind = "RCT"
+
 metrics = [
     "erupt",
     # "greedy_erupt",  # regular erupt was made probabilistic,
@@ -15,9 +17,15 @@
     "codec",  # NEW
     "bite",  # NEW
 ]
-use_ray = True
+estimators = get_estimator_list(kind)
+use_ray = False
 out_dir = run_batch(
-    identifier, kind, metrics, dataset_path=os.path.realpath("../RunDatasets"), use_ray=use_ray
+    identifier,
+    kind,
+    metrics,
+    estimators=estimators,
+    dataset_path=os.path.realpath("../RunDatasets"),
+    use_ray=use_ray,
 )
 # plot results
 # upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
diff --git a/notebooks/RunExperiments/runners/rct_no_meta.py b/notebooks/RunExperiments/runners/rct_no_meta.py
new file mode 100644
index 0000000..a9f431c
--- /dev/null
+++ b/notebooks/RunExperiments/runners/rct_no_meta.py
@@ -0,0 +1,35 @@
+import os
+
+from experiment_runner import run_batch, get_estimator_list
+from experiment_plots import generate_plots
+
+identifier = "Egor_test_no_meta"
+kind = "RCT"
+metrics = [
+    "erupt",
+    # "greedy_erupt",  # regular erupt was made probabilistic,
+    "policy_risk",  # NEW
+    "qini",
+    "auc",
+    "psw_energy_distance",
+    "frobenius_norm",  # NEW
+    "codec",  # NEW
+    "bite",  # NEW
+]
+
+estimators = get_estimator_list(kind, exclude_patterns=["SLearner", "TLearner", "XLearner"])
+
+use_ray = False
+out_dir = run_batch(
+    identifier,
+    kind,
+    metrics,
+    estimators=estimators,
+    dataset_path=os.path.realpath("../RunDatasets"),
+    use_ray=use_ray,
+)
+# plot results
+# upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
+# lower_bounds = {"erupt": 0.06, "bite": 0.75}
+generate_plots(os.path.join(out_dir, kind))  # , upper_bounds, lower_bounds)
+print("yay!")

From fc33b98d35cfce4860df000da79977659c1e06ba Mon Sep 17 00:00:00 2001
From: Egor Kraev <egor.kraev@gmail.com>
Date: Sun, 19 Jan 2025 11:29:22 +0100
Subject: [PATCH 5/5]  Fixes to BITE, Frobenius, CODEC

---
 causaltune/optimiser.py                       | 10 +++--
 causaltune/score/bite.py                      | 11 ++++--
 causaltune/score/scoring.py                   | 16 ++++----
 notebooks/RunExperiments/cluster_config.yaml  |  6 +--
 .../runners/experiment_plots.py               |  2 +-
 .../runners/experiment_runner.py              | 37 ++++++++++++------
 .../RunExperiments/runners/kc_no_meta.py      | 39 +++++++++++++++++++
 notebooks/RunExperiments/runners/rct.py       |  2 +-
 .../RunExperiments/runners/rct_no_meta.py     |  2 +-
 9 files changed, 94 insertions(+), 31 deletions(-)
 create mode 100644 notebooks/RunExperiments/runners/kc_no_meta.py

diff --git a/causaltune/optimiser.py b/causaltune/optimiser.py
index ff5f037..72cd8d3 100644
--- a/causaltune/optimiser.py
+++ b/causaltune/optimiser.py
@@ -94,6 +94,7 @@ def __init__(
         test_size=None,
         num_samples=-1,
         propensity_model="dummy",
+        propensity_automl_estimators: Optional[List[str]] = None,
         outcome_model="nested",
         components_task="regression",
         components_verbose=0,
@@ -185,6 +186,7 @@ def __init__(
         self._settings["component_models"]["n_jobs"] = components_njobs
         self._settings["component_models"]["time_budget"] = components_time_budget
         self._settings["component_models"]["eval_method"] = "holdout"
+        self._settings["propensity_automl_estimators"] = propensity_automl_estimators
 
         if 0 < train_size < 1:
             component_test_size = 1 - train_size
@@ -224,9 +226,11 @@ def init_propensity_model(self, propensity_model: str):
         if propensity_model == "dummy":
             self.propensity_model = DummyClassifier(strategy="prior")
         elif propensity_model == "auto":
-            self.propensity_model = AutoML(
-                **{**self._settings["component_models"], "task": "classification"}
-            )
+            automl_args = {**self._settings["component_models"], "task": "classification"}
+            if self._settings["propensity_automl_estimators"]:
+                automl_args["estimator_list"] = self._settings["propensity_automl_estimators"]
+
+            self.propensity_model = AutoML(**automl_args)
         elif hasattr(propensity_model, "fit") and hasattr(propensity_model, "predict_proba"):
             self.propensity_model = propensity_model
         else:
diff --git a/causaltune/score/bite.py b/causaltune/score/bite.py
index 2e54158..b501ec5 100644
--- a/causaltune/score/bite.py
+++ b/causaltune/score/bite.py
@@ -9,15 +9,20 @@ def bite(
     working_df: pd.DataFrame,
     treatment_name: str,
     outcome_name: str,
+    min_N: int = 10,
+    max_N: int = 1000,
+    num_N: int = 20,
     N_values: Optional[List[int]] = None,
+    clip_propensity: float = 0.05,
 ) -> float:
+    max_N = int(min(max_N, len(working_df) / 10))
     if N_values is None:
-        N_values = exponential_spacing(10, 100, 20)
+        N_values = exponential_spacing(min_N, max_N, num_N)
     # Calculate weights with clipping to avoid extremes
     working_df["weights"] = np.where(
         working_df[treatment_name] == 1,
-        1 / np.clip(working_df["propensity"], 0.05, 0.95),
-        1 / np.clip(1 - working_df["propensity"], 0.05, 0.95),
+        1 / np.clip(working_df["propensity"], clip_propensity, 1 - clip_propensity),
+        1 / np.clip(1 - working_df["propensity"], clip_propensity, 1 - clip_propensity),
     )
 
     kendall_tau_values = []
diff --git a/causaltune/score/scoring.py b/causaltune/score/scoring.py
index 283959f..67009c1 100644
--- a/causaltune/score/scoring.py
+++ b/causaltune/score/scoring.py
@@ -14,17 +14,17 @@
 from causaltune.score.thompson import thompson_policy, extract_means_stds
 from causaltune.thirdparty.causalml import metrics
 from causaltune.score.erupt import ERUPT
-from causaltune.score.bite import bite
+from .bite import bite
 from causaltune.utils import treatment_values, psw_joint_weights
 
 import dcor
 
 from scipy.spatial import distance
 from sklearn.neighbors import NearestNeighbors
-
-
 from sklearn.preprocessing import StandardScaler
 
+logger = logging.getLogger(__name__)
+
 
 class DummyEstimator:
     def __init__(self, cate_estimate: np.ndarray, effect_intervals: Optional[np.ndarray] = None):
@@ -93,7 +93,7 @@ def __init__(
         Access methods and attributes via `CausalTune.scorer`.
 
         """
-
+        logger.info("Initializing Scorer")
         self.problem = problem
         self.multivalue = multivalue
         self.causal_model = copy.deepcopy(causal_model)
@@ -341,8 +341,8 @@ def frobenius_norm_score(
         # Normalize features
         select_cols = estimate.estimator._effect_modifier_names + ["yhat"]
         scaler = StandardScaler()
-        Y0X_1_normalized = scaler.fit_transform(Y0X_1[select_cols])
-        Y0X_0_normalized = scaler.transform(Y0X_0[select_cols])
+        Y0X_0_normalized = scaler.fit_transform(Y0X_0[select_cols])
+        Y0X_1_normalized = scaler.transform(Y0X_1[select_cols])
 
         # Calculate pairwise differences
         differences_xy = Y0X_1_normalized[:, np.newaxis, :] - Y0X_0_normalized[np.newaxis, :, :]
@@ -927,7 +927,7 @@ def codec_score(estimate: CausalEstimate, df: pd.DataFrame) -> float:
         if standard_deviations < 0.01:
             return np.inf
 
-        return Scorer.codec(Y, Z, X)
+        return abs(Scorer.codec(Y, Z, X))
 
     @staticmethod
     def auc_make_score(
@@ -945,7 +945,7 @@ def auc_make_score(
             float: area under the uplift curve
 
         """
-
+        print("running auuc_score")
         est = estimate.estimator
         new_df = pd.DataFrame()
         new_df["y"] = df[est._outcome_name]
diff --git a/notebooks/RunExperiments/cluster_config.yaml b/notebooks/RunExperiments/cluster_config.yaml
index e316f6e..80d602d 100644
--- a/notebooks/RunExperiments/cluster_config.yaml
+++ b/notebooks/RunExperiments/cluster_config.yaml
@@ -6,7 +6,7 @@ cluster_name: default
 
 # The maximum number of workers nodes to launch in addition to the head
 # node.
-max_workers: 9
+max_workers: 8
 
 # The autoscaler will scale up the cluster faster with higher upscaling speed.
 # E.g., if the task requires adding more nodes then autoscaler will gradually
@@ -93,7 +93,7 @@ available_node_types:
         min_workers: 1
         # The maximum number of worker nodes of this type to launch.
         # This takes precedence over min_workers.
-        max_workers: 9
+        max_workers: 8
         # The resources provided by this node type.
         resources: {"CPU": 2}
         # Provider-specific config for the head node, e.g. instance type. By default
@@ -161,7 +161,7 @@ initialization_commands: []
 
 # List of shell commands to run to set up nodes.
 setup_commands:
-    - pip install causaltune catboost ray[tune]
+    - pip install causaltune catboost ray[tune] flaml[blendsearch]
 
     # Note: if you're developing Ray, you probably want to create a Docker image that
     # has your Ray repo pre-cloned. Then, you can replace the pip installs
diff --git a/notebooks/RunExperiments/runners/experiment_plots.py b/notebooks/RunExperiments/runners/experiment_plots.py
index 7bc1c6b..98617c5 100644
--- a/notebooks/RunExperiments/runners/experiment_plots.py
+++ b/notebooks/RunExperiments/runners/experiment_plots.py
@@ -70,7 +70,7 @@ def generate_plots(
         "bite": "BITE",
         "policy_risk": "Policy\nRisk",
         "energy_distance": "Energy\nDistance",
-        "psw_energy_distance": "PSW\nEnergy\nDistance",
+        "psw_energy_distance": "Energy\nDistance",
         "norm_erupt": "Normalized\nERUPT",
     }
 
diff --git a/notebooks/RunExperiments/runners/experiment_runner.py b/notebooks/RunExperiments/runners/experiment_runner.py
index 04abc29..4ebb081 100644
--- a/notebooks/RunExperiments/runners/experiment_runner.py
+++ b/notebooks/RunExperiments/runners/experiment_runner.py
@@ -12,6 +12,9 @@
 from sklearn.model_selection import train_test_split
 
 
+sys.path.insert(0, os.getcwd())
+import causaltune  # noqa: E402
+
 from causaltune import CausalTune
 from causaltune.data_utils import CausalityDataset
 from causaltune.models.passthrough import passthrough_model
@@ -112,6 +115,7 @@ def run_experiment(
     estimators: List[str],
     dataset_path: str,
     use_ray: bool,
+    propensity_automl_estimators: Optional[List[str]] = None,
 ):
     # Process datasets
     data_sets = {}
@@ -125,6 +129,7 @@ def run_experiment(
         name = " ".join(parts[1:])
         file_path = f"{dataset_path}/{size}/{name}.pkl"
         data_sets[f"{size} {name}"] = load_dataset(file_path)
+    run_kind = dataset.split("_")[1]
 
     out_dir = f"../EXPERIMENT_RESULTS_{args.identifier}"
     os.makedirs(out_dir, exist_ok=True)
@@ -136,24 +141,22 @@ def run_experiment(
     already_running = False
     if use_ray:
         try:
-            runner = ray.get_actor("TaskRunner")
+            runner = ray.get_actor(f"TaskRunner {run_kind}")
             print("\n" * 4)
             print(
                 "!!! Found an existing detached TaskRunner. Will assume the tasks have already been submitted."
             )
             print(
-                "!!! If you want to re-run the experiments from scratch, "
-                'run ray.kill(ray.get_actor("TaskRunner", namespace="{}")) or recreate the cluster.'.format(
-                    RAY_NAMESPACE
-                )
+                f"!!! If you want to re-run the experiments from scratch, "
+                'run ray.kill(ray.get_actor("TaskRunner {run_kind}", namespace="{RAY_NAMESPACE}")) or recreate the cluster.'
             )
             print("\n" * 4)
             already_running = True
         except ValueError:
             print("Ray: no detached TaskRunner found, creating...")
             # This thing will be alive even if the host program exits
-            # Must be killed explicitly: ray.kill(ray.get_actor("TaskRunner"))
-            runner = TaskRunner.options(name="TaskRunner", lifetime="detached").remote()
+            # Must be killed explicitly: ray.kill(ray.get_actor(f"TaskRunner {run_kind}"))
+            runner = TaskRunner.options(name=f"TaskRunner {run_kind}", lifetime="detached").remote()
 
     out = []
     if not already_running:
@@ -190,6 +193,7 @@ def run_experiment(
                             args.components_time_budget,
                             out_fn,
                             estimators,
+                            propensity_automl_estimators,
                         )
                     )
                 else:
@@ -202,6 +206,7 @@ def run_experiment(
                         args.components_time_budget,
                         out_fn,
                         estimators,
+                        propensity_automl_estimators,
                     )
                     out.append(results)
 
@@ -238,6 +243,7 @@ def run_batch(
     estimators: List[str],
     dataset_path: str,
     use_ray: bool = False,
+    propensity_automl_estimators: Optional[List[str]] = None,
 ):
     args = parse_arguments()
     args.identifier = identifier
@@ -255,12 +261,19 @@ def run_batch(
         # Assuming we port-mapped already by running ray dashboard
         ray.init(
             "ray://localhost:10001",
-            runtime_env={"working_dir": ".", "pip": ["causaltune", "catboost", "ray[tune]"]},
+            runtime_env={
+                "working_dir": ".",
+                "pip": ["causaltune", "catboost", "ray[tune]", "flaml[blendsearch]"],
+            },
             namespace=RAY_NAMESPACE,
         )
 
     out_dir = run_experiment(
-        args, estimators=estimators, dataset_path=dataset_path, use_ray=use_ray
+        args,
+        estimators=estimators,
+        dataset_path=dataset_path,
+        use_ray=use_ray,
+        propensity_automl_estimators=propensity_automl_estimators,
     )
     return out_dir
 
@@ -275,8 +288,8 @@ class TaskRunner:
     def __init__(self):
         self.futures = {}
 
-    def remote_single_run(self, *args):
-        ref = remote_single_run.remote(*args)
+    def remote_single_run(self, *args, **kwargs):
+        ref = remote_single_run.remote(*args, **kwargs)
         self.futures[ref.hex()] = ref
         return ref.hex()
 
@@ -310,6 +323,7 @@ def single_run(
     components_time_budget: int,
     out_fn: str,
     estimators: List[str],
+    propensity_automl_estimators: Optional[List[str]] = None,
     outcome_model: str = "auto",
     i_run: int = 1,
 ):
@@ -342,6 +356,7 @@ def single_run(
             store_all_estimators=True,
             propensity_model=propensity_model,
             outcome_model=outcome_model,
+            propensity_automl_estimators=propensity_automl_estimators,
             use_ray=False,
         )
 
diff --git a/notebooks/RunExperiments/runners/kc_no_meta.py b/notebooks/RunExperiments/runners/kc_no_meta.py
new file mode 100644
index 0000000..7be9466
--- /dev/null
+++ b/notebooks/RunExperiments/runners/kc_no_meta.py
@@ -0,0 +1,39 @@
+import os
+
+from experiment_runner import run_batch, get_estimator_list
+from experiment_plots import generate_plots
+
+identifier = "Egor_test"
+kind = "KC"
+metrics = [
+    "erupt",
+    # "greedy_erupt",  # regular erupt was made probabilistic,
+    "policy_risk",  # NEW
+    "qini",
+    "auc",
+    "psw_energy_distance",
+    "frobenius_norm",  # NEW
+    "codec",  # NEW
+    "bite",  # NEW
+]
+estimators = get_estimator_list(kind, exclude_patterns=["SLearner", "TLearner", "XLearner"])
+ptt_estimators = [
+    "lgbm",
+    "lrl2",
+]
+
+use_ray = True
+out_dir = run_batch(
+    identifier,
+    kind,
+    metrics,
+    estimators=estimators,
+    propensity_automl_estimators=ptt_estimators,
+    dataset_path=os.path.realpath("../RunDatasets"),
+    use_ray=use_ray,
+)
+# plot results
+# upper_bounds = {"MSE": 1e2, "policy_risk": 0.2}
+# lower_bounds = {"erupt": 0.06, "bite": 0.75}
+generate_plots(os.path.join(out_dir, kind))  # , upper_bounds, lower_bounds)
+print("yay!")
diff --git a/notebooks/RunExperiments/runners/rct.py b/notebooks/RunExperiments/runners/rct.py
index 8aebd69..41e1c9c 100644
--- a/notebooks/RunExperiments/runners/rct.py
+++ b/notebooks/RunExperiments/runners/rct.py
@@ -18,7 +18,7 @@
     "bite",  # NEW
 ]
 estimators = get_estimator_list(kind)
-use_ray = False
+use_ray = True
 out_dir = run_batch(
     identifier,
     kind,
diff --git a/notebooks/RunExperiments/runners/rct_no_meta.py b/notebooks/RunExperiments/runners/rct_no_meta.py
index a9f431c..fae51e9 100644
--- a/notebooks/RunExperiments/runners/rct_no_meta.py
+++ b/notebooks/RunExperiments/runners/rct_no_meta.py
@@ -19,7 +19,7 @@
 
 estimators = get_estimator_list(kind, exclude_patterns=["SLearner", "TLearner", "XLearner"])
 
-use_ray = False
+use_ray = True
 out_dir = run_batch(
     identifier,
     kind,