fixing input lists to intervention templates

ciemss · Aug 8, 2024 · 02ed976 · 02ed976
1 parent b8e9b18
commit 02ed976
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 20 deletions.
diff --git a/docs/source/optimize_interface.ipynb b/docs/source/optimize_interface.ipynb
@@ -1104,11 +1104,11 @@
     "intervened_params = [\"beta_c\", \"gamma\", \"gamma_c\"]\n",
     "static_parameter_interventions1 = param_value_objective(\n",
     "    param_name=[intervened_params[0]],\n",
-    "    start_time=torch.tensor(10.0),\n",
+    "    start_time=[torch.tensor(10.0)],\n",
     ")\n",
     "static_parameter_interventions2 = start_time_objective(\n",
     "    param_name=[intervened_params[1]],\n",
-    "    param_value=torch.tensor([0.4]),\n",
+    "    param_value=[torch.tensor([0.4])],\n",
     ")\n",
     "\n",
     "# Combine different intervention templates\n",
@@ -1270,11 +1270,11 @@
     "intervened_params = [\"beta_c\", \"gamma\"]\n",
     "static_parameter_interventions1 = param_value_objective(\n",
     "    param_name=[intervened_params[0]],\n",
-    "    start_time=torch.tensor(10.0),\n",
+    "    start_time=[torch.tensor(10.0)],\n",
     ")\n",
     "static_parameter_interventions2 = start_time_objective(\n",
     "    param_name=[intervened_params[1]],\n",
-    "    param_value=torch.tensor([0.45]),\n",
+    "    param_value=[torch.tensor([0.45])],\n",
     ")\n",
     "# Combine different intervention templates into a list of Callables\n",
     "static_parameter_interventions = intervention_func_combinator(\n",

diff --git a/pyciemss/integration_utils/intervention_builder.py b/pyciemss/integration_utils/intervention_builder.py
@@ -30,15 +30,13 @@ def intervention_generator(
         for count in range(param_size):
             if start_time[count].item() in static_parameter_interventions:
                 static_parameter_interventions[start_time[count].item()].update(
-                    {param_name[count]: param_value[count](x[count].item()).item()}
+                    {param_name[count]: param_value[count](x[count].item())}
                 )
             else:
                 static_parameter_interventions.update(
                     {
                         start_time[count].item(): {
-                            param_name[count]: param_value[count](
-                                x[count].item()
-                            ).item()
+                            param_name[count]: param_value[count](x[count].item())
                         }
                     }
                 )
@@ -52,7 +50,7 @@ def start_time_objective(
     param_value: List[Intervention],
 ) -> Callable[[torch.Tensor], Dict[float, Dict[str, Intervention]]]:
     param_size = len(param_name)
-    # Note that code below will only work for tensors and not callable functions
+    # Note: code below will only work for tensors and not callable functions
     param_value = [torch.atleast_1d(y) for y in param_value]
 
     def intervention_generator(
@@ -67,11 +65,11 @@ def intervention_generator(
         for count in range(param_size):
             if x[count].item() in static_parameter_interventions:
                 static_parameter_interventions[x[count].item()].update(
-                    {param_name[count]: param_value[count].item()}
+                    {param_name[count]: param_value[count]}
                 )
             else:
                 static_parameter_interventions.update(
-                    {x[count].item(): {param_name[count]: param_value[count].item()}}
+                    {x[count].item(): {param_name[count]: param_value[count]}}
                 )
         return static_parameter_interventions
 
@@ -102,19 +100,15 @@ def intervention_generator(
         for count in range(param_size):
             if x[count * 2].item() in static_parameter_interventions:
                 static_parameter_interventions[x[count * 2].item()].update(
-                    {
-                        param_name[count]: param_value[count](
-                            x[count * 2 + 1].item()
-                        ).item()
-                    }
+                    {param_name[count]: param_value[count](x[count * 2 + 1].item())}
                 )
             else:
                 static_parameter_interventions.update(
                     {
                         x[count * 2].item(): {
                             param_name[count]: param_value[count](
                                 x[count * 2 + 1].item()
-                            ).item()
+                            )
                         }
                     }
                 )
@@ -129,7 +123,10 @@ def intervention_func_combinator(
     ],
     intervention_func_lengths: List[int],
 ) -> Callable[[torch.Tensor], Dict[float, Dict[str, Intervention]]]:
-    assert len(intervention_funcs) == len(intervention_func_lengths)
+    assert len(intervention_funcs) == len(intervention_func_lengths), (
+        f"Size mismatch between number of intervention functions ('{len(intervention_funcs)}')"
+        f"and number of intervention function lengths ('{len(intervention_func_lengths)}') "
+    )
 
     total_length = sum(intervention_func_lengths)
 

diff --git a/tests/fixtures.py b/tests/fixtures.py
@@ -134,11 +134,11 @@ def __init__(
 intervened_params = ["beta_c", "gamma"]
 static_parameter_interventions1 = param_value_objective(
     param_name=[intervened_params[0]],
-    start_time=torch.tensor(10.0),
+    start_time=[torch.tensor(10.0)],
 )
 static_parameter_interventions2 = start_time_objective(
     param_name=[intervened_params[1]],
-    param_value=torch.tensor([0.45]),
+    param_value=[torch.tensor([0.45])],
 )
 optkwargs_SEIRHD_paramtimeComb_maxQoI = {
     "qoi": [lambda x: obs_max_qoi(x, ["I_state"])],