Add small box optimize endpoint (#58)

* Add broken optimize tests

* Bump PyCIEMSS version

* Update PyCIEMSS to use patch

* Use updated PyCIEMSS and flesh out optimize

* Pass conversion test

* Fix serialization

* Fix execution in API

* Fix postproc + tests

* Make optimize test run faster

.gitignore

@@ -75,5 +75,3 @@ create.sql
 .eslintcache
 
 .version
-venv
-flake.lock

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pyciemss-service"
-version = "2.0.0"
+version = "2.1.0"
 description = "PyCIEMSS simulation service to run CIEMSS simulations"
 authors = ["Powell Fendley", "Five Grant"]
 readme = "README.md"
@@ -21,6 +21,7 @@ filelock = "^3.12.2"
 poethepoet = "^0.21.1"
 # juliacall = { version="^0.9.14", optional = true }
 dill = "^0.3.7"
+numpy = "^1.26.4"
 
 
 [tool.poetry.scripts]
@@ -39,7 +40,7 @@ httpx = "^0.24.1"
 
 
 [tool.poe.tasks]
-install-pyciemss = "pip install --no-cache-dir git+https://github.com/fivegrant/pyciemss.git@087bc64d935f2ab5090330f1f7d6bde930404115 --use-pep517"
+install-pyciemss = "pip install --no-cache-dir git+https://github.com/ciemss/pyciemss.git@1fabf279590c613a8ed38d88c6b6faf0c52ba867 --use-pep517"
 
 
 [tool.pytest.ini_options]
@@ -56,4 +57,4 @@ build-backend = "poetry.core.masonry.api"
 ignore = ["E501"]
 
 [tool.ruff.per-file-ignores]
-"__init__.py" = ["F401", "F403"]
+"__init__.py" = ["F401", "F403"]

service/api.py

@@ -11,6 +11,7 @@
  Calibrate,
  Simulate,
  EnsembleSimulate,
+ Optimize,
  StatusSimulationIdGetResponse,
 )
 
@@ -20,6 +21,7 @@
  "simulate": Simulate,
  "calibrate": Calibrate,
  "ensemble-simulate": EnsembleSimulate,
+ "optimize": Optimize,
 }
 
 logging.basicConfig()
@@ -114,12 +116,3 @@ def ensemble_calibrate_not_yet_implemented():
  This will be reimplemented in the future.
  """
  raise HTTPException(status=501, detail="Not yet reimplemented")
-
-
-@app.get("/optimize", response_model=StatusSimulationIdGetResponse) # NOT YET IN SPEC
-def optimize_not_yet_implemented(): # NOT YET IN SPEC
- """
- DO NOT USE. Placeholder for `optimize` endpoint.
- This will be implemented in the future.
- """
- raise HTTPException(status=501, detail="Not yet implemented")

service/execute.py

@@ -7,7 +7,12 @@
  attach_files,
 )
 
-from pyciemss.interfaces import sample, calibrate, ensemble_sample # noqa: F401
+from pyciemss.interfaces import ( # noqa: F401
+ sample,
+ calibrate,
+ ensemble_sample,
+ optimize,
+)
 
 # jl = newmodule("SciMLIntegration")
 # jl.seval("using SciMLIntegration, PythonCall")

service/models/base.py

@@ -42,12 +42,6 @@ class InterventionSelection(BaseModel):
  name: str
 
 
-class QuantityOfInterest(BaseModel):
- function: str
- state: str
- arg: int # TODO: Make this a list of args?
-
-
 class OperationRequest(BaseModel):
  pyciemss_lib_function: ClassVar[str] = ""
  engine: str = Field("ciemss", example="ciemss")

service/models/operations/__init__.py

@@ -1,3 +1,4 @@
 from models.operations.simulate import Simulate
 from models.operations.calibrate import Calibrate
 from models.operations.ensemble_simulate import EnsembleSimulate
+from models.operations.optimize import Optimize

service/models/operations/optimize.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+# from enum import Enum
+from typing import ClassVar, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+from pydantic import BaseModel, Field, Extra
+
+
+from models.base import OperationRequest, Timespan
+from utils.tds import fetch_model, fetch_inferred_parameters
+
+
+# TODO: Add more methods later if needed
+# class QOIMethod(Enum):
+# obs_nday_average = "obs_nday_average"
+
+
+def obs_nday_average_qoi(
+ samples: Dict[str, torch.Tensor], contexts: List, ndays: int = 7
+) -> np.ndarray:
+ """
+ Return estimate of last n-day average of each sample.
+ samples is is the output from a Pyro Predictive object.
+ samples[VARIABLE] is expected to have dimension (nreplicates, ntimepoints)
+ Note: last ndays timepoints is assumed to represent last n-days of simulation.
+
+ Taken from:
+ https://github.com/ciemss/pyciemss/blob/main/docs/source/interfaces.ipynb
+ """
+ dataQoI = samples[contexts[0] + "_state"].detach().numpy()
+
+ return np.mean(dataQoI[:, -ndays:], axis=1)
+
+
+# qoi_implementations = {QOIMethod.obs_nday_average.value: obs_nday_average_qoi}
+
+
+class OptimizeExtra(BaseModel):
+ num_samples: int = Field(
+ 100,
+ description="""
+ The number of samples to draw from the model to estimate risk for each optimization iteration.
+ """,
+ example=100,
+ )
+ inferred_parameters: Optional[str] = Field(
+ None,
+ description="id from a previous calibration",
+ example=None,
+ )
+ maxiter: int = 5
+ maxfeval: int = 5
+
+
+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")]
+ )
+ step_size: float = 1.0
+ qoi: List[str] # QOIMethod
+ risk_bound: float
+ initial_guess_interventions: List[float]
+ bounds_interventions: List[List[float]]
+ extra: OptimizeExtra = Field(
+ None,
+ description="optional extra system specific arguments for advanced use cases",
+ )
+
+ 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}
+
+ extra_options = self.extra.dict()
+ inferred_parameters = fetch_inferred_parameters(
+ extra_options.pop("inferred_parameters"), job_id
+ )
+ n_samples_ouu = extra_options.pop("num_samples")
+
+ return {
+ "model_path_or_json": amr_path,
+ "logging_step_size": self.step_size,
+ "start_time": self.timespan.start,
+ "end_time": self.timespan.end,
+ "objfun": lambda x: np.sum(np.abs(x)),
+ "qoi": lambda samples: obs_nday_average_qoi(samples, self.qoi, 1),
+ "risk_bound": self.risk_bound,
+ "initial_guess_interventions": self.initial_guess_interventions,
+ "bounds_interventions": self.bounds_interventions,
+ "static_parameter_interventions": interventions,
+ "inferred_parameters": inferred_parameters,
+ "n_samples_ouu": n_samples_ouu,
+ **extra_options,
+ }
+
+ class Config:
+ extra = Extra.forbid
+ # use_enum_values = True

service/utils/tds.py

@@ -198,17 +198,24 @@ def attach_files(output: dict, job_id, status="complete"):
 
  params_filename = os.path.join(job_dir, "./parameters.dill")
  params_result = output.get("inferred_parameters", None)
- if params_result:
+ if params_result is not None:
  with open(params_filename, "wb") as file:
  dill.dump(params_result, file)
  files[params_filename] = "parameters.dill"
 
- policy_filename = os.path.join(job_dir, "./policy.dill")
+ policy_filename = os.path.join(job_dir, "./policy.json")
  policy = output.get("policy", None)
  if policy is not None:
- with open(policy_filename, "wb") as file:
- dill.dump(params_result, file)
- files[policy_filename] = "policy.dill"
+ with open(policy_filename, "w") as file:
+ json.dump(policy.tolist(), file)
+ files[policy_filename] = "policy.json"
+
+ results_filename = os.path.join(job_dir, "./optimize_results.dill")
+ results = output.get("OptResults", None)
+ if results is not None:
+ with open(results_filename, "wb") as file:
+ dill.dump(results, file)
+ files[results_filename] = "optimize_results.dill"
 
  visualization_filename = os.path.join(job_dir, "./visualization.json")
  viz_result = output.get("visual", None)