Optimize Interventions type -> same as simulation interventions

service/models/base.py

@@ -37,6 +37,11 @@ class InterventionObject(BaseModel):
  value: float
 
 
+class OptimizeInterventionObject(BaseModel):
+ timestep: float
+ name: str
+
+
 class InterventionSelection(BaseModel):
  timestep: float
  name: str

service/models/converters.py

@@ -11,6 +11,13 @@ def convert_to_static_interventions(interventions):
  return static_interventions
 
 
+def convert_optimize_to_static_interventions(interventions):
+ static_interventions = defaultdict(dict)
+ for i in interventions:
+ static_interventions[torch.tensor(i.timestep)] = i.name
+ return static_interventions
+
+
 def convert_to_solution_mapping(config):
  individual_to_ensemble = {
  individual_state: ensemble_state

service/models/operations/optimize.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
 
 # from enum import Enum
-from typing import ClassVar, Dict, List, Optional, Tuple
+from typing import ClassVar, Dict, List, Optional
 
 import numpy as np
 import torch
 from pydantic import BaseModel, Field, Extra
-
-
-from models.base import OperationRequest, Timespan
+from models.base import OperationRequest, Timespan, OptimizeInterventionObject
+from models.converters import convert_optimize_to_static_interventions
 from utils.tds import fetch_model, fetch_inferred_parameters
 
 
@@ -58,8 +57,8 @@ class Optimize(OperationRequest):
  pyciemss_lib_function: ClassVar[str] = "optimize"
  model_config_id: str = Field(..., example="ba8da8d4-047d-11ee-be56")
  timespan: Timespan = Timespan(start=0, end=90)
- interventions: List[Tuple[float, str]] = Field(
- default_factory=list, example=[(1.0, "beta")]
+ interventions: List[OptimizeInterventionObject] = Field(
+ default_factory=list, example=[{"timestep": 1, "name": "beta"}]
  )
  step_size: float = 1.0
  qoi: List[str] # QOIMethod
@@ -75,7 +74,7 @@ def gen_pyciemss_args(self, job_id):
  # Get model from TDS
  amr_path = fetch_model(self.model_config_id, job_id)
 
- interventions = {torch.tensor(k): v for k, v in self.interventions}
+ interventions = convert_optimize_to_static_interventions(self.interventions)
 
  extra_options = self.extra.dict()
  inferred_parameters = fetch_inferred_parameters(

tests/examples/optimize/input/request.json

@@ -3,7 +3,10 @@
  "user_id": "not_provided",
  "model_config_id": "sidarthe",
  "interventions": [
- [1.0, "beta"]
+ {
+ "timestep": 1.0,
+ "name": "beta"
+ }
  ],
  "timespan": {
  "start": 0,