1 files changed, 28 insertions, 15 deletions
diff --git a/experiments/analyze_coupled_ladder_scaling.py b/experiments/analyze_coupled_ladder_scaling.py
index 9f36863..7f32c50 100644
--- a/experiments/analyze_coupled_ladder_scaling.py
+++ b/experiments/analyze_coupled_ladder_scaling.py
@@ -39,11 +39,13 @@ STYLE = {
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--core", type=Path,
-        default=Path("results/coupled_ladder/p1_imperfection_pilot.json"))
+        "--core", type=Path, nargs="+",
+        default=[Path(
+            "results/coupled_ladder/p1_imperfection_pilot.json")])
     parser.add_argument(
-        "--baselines", type=Path,
-        default=Path("results/coupled_ladder/p1_bias_baseline_pilot.json"))
+        "--baselines", type=Path, nargs="*",
+        default=[Path(
+            "results/coupled_ladder/p1_bias_baseline_pilot.json")])
     parser.add_argument(
         "--output-analysis", type=Path,
         default=Path("results/coupled_ladder/p1_scaling_analysis.json"))
@@ -58,9 +60,9 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def merge_reports(core: dict, baselines: dict) -> list[dict]:
+def merge_reports(reports: list[dict]) -> list[dict]:
     records = {}
-    for report in (core, baselines):
+    for report in reports:
         for record in report["records"]:
             key = (
                 record["side"],
@@ -204,17 +206,28 @@ def slope_bootstrap(
 def trace_summary(
     records: list[dict], side: int, method: str
 ) -> list[dict]:
-    method_records = [
-        record["methods"][method]
+    source_records = [
+        record
         for record in records
         if record["side"] == side and method in record["methods"]
     ]
-    epochs = [record["epoch"] for record in method_records[0]["trace"]]
+    task_indices = sorted({
+        record["task_index"] for record in source_records})
+    epochs = [
+        record["epoch"]
+        for record in source_records[0]["methods"][method]["trace"]
+    ]
     output = []
     for index, epoch in enumerate(epochs):
         values = np.asarray([
-            record["trace"][index]["classification_error"]
-            for record in method_records
+            np.mean([
+                record["methods"][method]["trace"][index][
+                    "classification_error"
+                ]
+                for record in source_records
+                if record["task_index"] == task_index
+            ])
+            for task_index in task_indices
         ])
         output.append({
             "epoch": epoch,
@@ -459,15 +472,15 @@ def plot_figure(path: Path, analysis: dict) -> None:
 
 def main() -> None:
     args = parse_args()
-    core = json.loads(args.core.read_text())
-    baselines = json.loads(args.baselines.read_text())
-    records = merge_reports(core, baselines)
+    source_paths = list(args.core) + list(args.baselines)
+    reports = [json.loads(path.read_text()) for path in source_paths]
+    records = merge_reports(reports)
     analysis = build_analysis(
         records,
         replicates=args.bootstrap_replicates,
         seed=args.bootstrap_seed,
     )
-    analysis["sources"] = [str(args.core), str(args.baselines)]
+    analysis["sources"] = [str(path) for path in source_paths]
     args.output_analysis.parent.mkdir(parents=True, exist_ok=True)
     args.output_analysis.write_text(json.dumps(analysis, indent=2) + "\n")
     write_csv(args.output_csv, analysis)