progress hook progress

pyciemss/interfaces.py

@@ -791,6 +791,7 @@ def optimize(
     maxfeval: int = 25,
     verbose: bool = False,
     roundup_decimal: int = 4,
+    progress_hook: Callable = lambda i, feval: None,
 ) -> Dict[str, Any]:
     r"""
     Load a model from a file, compile it into a probabilistic program, and optimize under uncertainty with risk-based
@@ -932,6 +933,7 @@ def objfun_penalty(x):
             maxiter=maxiter,
             maxfeval=maxfeval,
             u_bounds=bounds_np,
+            progress_hook=progress_hook,
         ).solve()
 
         # Rounding up to given number of decimal places

pyciemss/ouu/ouu.py

@@ -186,6 +186,7 @@ def __init__(
         maxfeval: int = 100,
         maxiter: int = 100,
         u_bounds: np.ndarray = np.atleast_2d([[0], [1]]),
+        progress_hook: Callable = lambda i: None,
     ):
         self.x0 = np.squeeze(np.array([x0]))
         self.objfun = objfun
@@ -197,6 +198,7 @@ def __init__(
         self.maxiter = maxiter
         self.maxfeval = maxfeval
         self.u_bounds = u_bounds
+        self.progress_hook = progress_hook
         # self.kwargs = kwargs
 
     def solve(self):
@@ -211,7 +213,7 @@ def update_progress(xk):
             constraints=self.constraints,
             method="COBYLA",
             tol=1e-5,
-            callback=update_progress,
+            callback=self.progress_hook,
             options={
                 "rhobeg": 0.1
                 * np.linalg.norm(self.u_bounds[1, :] - self.u_bounds[0, :]),

tests/test_interfaces.py

@@ -562,6 +562,19 @@ def test_output_format(
 def test_optimize(model_fixture, start_time, end_time, num_samples):
     logging_step_size = 1.0
     model_url = model_fixture.url
+
+    class TestProgressHook:
+        def __init__(self):
+            self.iterations = []
+            # self.function_evals = []
+
+        def __call__(self, iteration):
+            # Log the iteration number
+            self.iterations.append(iteration)
+            # self.function_evals.append(feval)
+
+    progress_hook = TestProgressHook()
+
     optimize_kwargs = {
         **model_fixture.optimize_kwargs,
         "solver_method": "euler",
@@ -570,6 +583,7 @@ def test_optimize(model_fixture, start_time, end_time, num_samples):
         "n_samples_ouu": int(2),
         "maxiter": 1,
         "maxfeval": 2,
+        "progress_hook": progress_hook,
     }
     bounds_interventions = optimize_kwargs["bounds_interventions"]
     opt_result = optimize(
@@ -617,6 +631,9 @@ def test_optimize(model_fixture, start_time, end_time, num_samples):
         intervened_result_subset, start_time, end_time, logging_step_size, num_samples
     )
 
+    assert len(progress_hook.iterations) == (optimize_kwargs["maxfeval"] + 1) * (optimize_kwargs["maxiter"] + 1)
+    # assert len(progress_hook.function_evals) == optimize_kwargs["maxfeval"]
+
 
 @pytest.mark.parametrize("model_fixture", MODELS)
 @pytest.mark.parametrize("bad_num_iterations", NON_POS_INTS)