Validate cost_model argument in _get_time function

optimagic-dev · Dec 3, 2024 · c86421c · c86421c
1 parent fe27a3b
commit c86421c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/optimagic/optimization/history.py b/src/optimagic/optimization/history.py
@@ -192,7 +192,8 @@ def flat_param_names(self) -> list[str]:
     def _get_time(
         self, cost_model: CostModel | Literal["wall_time"]
     ) -> NDArray[np.float64]:
-        # TODO: validate that cost_model is either a CostModel or "wall_time"
+        if not isinstance(cost_model, CostModel) and cost_model != "wall_time":
+            raise ValueError("cost_model must be a CostModel or 'wall_time'.")
 
         if cost_model == "wall_time":
             return np.array(self.stop_time, dtype=np.float64) - self.start_time[0]