From f677039ffc7146942e997b5d5b2cdeb3d29f4414 Mon Sep 17 00:00:00 2001
From: SeanMcOwen <seanmcowen@gmail.com>
Date: Thu, 18 Apr 2024 14:05:49 -0400
Subject: [PATCH] Created add_parameter_labels

---
 cadCAD/tools/__init__.py |  2 +-
 cadCAD/tools/utils.py    | 37 ++++++++++++++++++++++++++++++++-----
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/cadCAD/tools/__init__.py b/cadCAD/tools/__init__.py
index b5a75f8a..fec34456 100644
--- a/cadCAD/tools/__init__.py
+++ b/cadCAD/tools/__init__.py
@@ -1,3 +1,3 @@
 from cadCAD.tools.execution import easy_run
 from cadCAD.tools.profiling import profile_run
-from cadCAD.tools.utils import generic_suf
\ No newline at end of file
+from cadCAD.tools.utils import generic_suf, add_parameter_labels
diff --git a/cadCAD/tools/utils.py b/cadCAD/tools/utils.py
index ddbe0583..c50cf715 100644
--- a/cadCAD/tools/utils.py
+++ b/cadCAD/tools/utils.py
@@ -1,17 +1,44 @@
 from cadCAD.types import *
+import pandas as pd
 
-def generic_suf(variable: str,
-                signal: str='') -> StateUpdateFunction:
+
+def generic_suf(variable: str, signal: str = "") -> StateUpdateFunction:
     """
-    Generate a State Update Function that assigns the signal value to the 
+    Generate a State Update Function that assigns the signal value to the
     given variable. By default, the signal has the same identifier as the
     variable.
     """
-    if signal is '':
+    if signal is "":
         signal = variable
     else:
         pass
 
     def suf(_1, _2, _3, _4, signals: PolicyOutput) -> StateUpdateTuple:
         return (variable, signals[signal])
-    return suf
\ No newline at end of file
+
+    return suf
+
+
+def add_parameter_labels(configs: list, df: pd.DataFrame) -> pd.DataFrame:
+    """Utility function to add the parameters to a dataframe after processing
+
+    Args:
+        configs (list): The configurations of the simulations
+        df (pd.DataFrame): Simulation dataframe
+
+    Returns:
+        pd.DataFrame: Simulation dataframe with labels
+    """
+
+    # Find the relevant parameters
+    sim_params = pd.DataFrame([x.sim_config["M"] for x in configs])
+    sim_params[["subset", "simulation", "run"]] = [
+        [x.subset_id, x.simulation_id, x.run_id] for x in configs
+    ]
+    # Fix because run_id is 0 indexed, but cadCAD dataframe is 1 indexed for runs
+    sim_params["run"] += 1
+
+    # Join
+    sim_params = sim_params.set_index(["subset", "simulation", "run"])
+    df = df.join(sim_params, on=["subset", "simulation", "run"])
+    return df