Some minor changes after first round of own review

Author: timmens

src/optimagic/optimization/history.py

@@ -93,7 +93,7 @@ def _get_next_batch_id(self) -> int:
         return batch
 
     # ==================================================================================
-    # Properties to access the history
+    # Properties and methods to access the history
     # ==================================================================================
 
     # Function data, function value, and monotone function value
@@ -243,10 +243,11 @@ def _get_time_per_task(
 
         """
         dummy_task = np.array([1 if t == task else 0 for t in self.task])
+        factor: float | NDArray[np.float64]
         if cost_factor is None:
-            factor: float | NDArray[np.float64] = np.array(
-                self.stop_time, dtype=np.float64
-            ) - np.array(self.start_time, dtype=np.float64)
+            factor = np.array(self.stop_time, dtype=np.float64) - np.array(
+                self.start_time, dtype=np.float64
+            )
         else:
             factor = cost_factor
 
@@ -342,16 +343,16 @@ def _calculate_monotone_sequence(
     sequence: list[float | None], direction: Direction
 ) -> NDArray[np.float64]:
     sequence_arr = np.array(sequence, dtype=np.float64)  # converts None to nan
-    none_mask = np.isnan(sequence_arr)
+    nan_mask = np.isnan(sequence_arr)
 
     if direction == Direction.MINIMIZE:
-        sequence_arr[none_mask] = np.inf
+        sequence_arr[nan_mask] = np.inf
         out = np.minimum.accumulate(sequence_arr)
     elif direction == Direction.MAXIMIZE:
-        sequence_arr[none_mask] = -np.inf
+        sequence_arr[nan_mask] = -np.inf
         out = np.maximum.accumulate(sequence_arr)
 
-    out[none_mask] = np.nan
+    out[nan_mask] = np.nan
     return out
 
 
@@ -404,7 +405,7 @@ def _apply_to_batch(
 
     """
     batch_starts = _get_batch_start(batch_ids)
-    batch_stops = [*batch_starts, len(data)][1:]
+    batch_stops = [*batch_starts[1:], len(data)]
 
     batch_results = []
     for batch, (start, stop) in zip(

src/optimagic/optimization/process_results.py

@@ -112,8 +112,8 @@ def process_multistart_result(
             direction=extra_fields.direction,
             fun=[opt.fun for opt in info.local_optima],
             params=[opt.params for opt in info.local_optima],
-            start_time=[np.nan for _ in info.local_optima],
-            stop_time=[np.nan for _ in info.local_optima],
+            start_time=len(info.local_optima) * [np.nan],
+            stop_time=len(info.local_optima) * [np.nan],
             batches=list(range(len(info.local_optima))),
             task=len(info.local_optima) * [EvalTask.FUN],
         )

src/optimagic/visualization/history_plots.py

@@ -407,10 +407,11 @@ def _extract_plotting_data_from_database(res, stack_multistart, show_exploration
         fun=_history["fun"],
         params=_history["params"],
         start_time=_history["time"],
-        # TODO: This needs to be updated
+        # TODO (@janosg): Retrieve that information from `hist` once it is available.
+        # https://github.com/optimagic-dev/optimagic/pull/553
         stop_time=len(_history["fun"]) * [None],
-        batches=len(_history["fun"]) * [None],
         task=len(_history["fun"]) * [None],
+        batches=list(range(len(_history["fun"]))),
     )
 
     data = {
@@ -449,8 +450,10 @@ def _get_stacked_local_histories(local_histories, direction, history=None):
         fun=stacked["criterion"],
         params=stacked["params"],
         start_time=stacked["runtime"],
-        # TODO: This needs to be fixed
+        # TODO (@janosg): Retrieve that information from `hist` once it is available
+        # for the IterationHistory.
+        # https://github.com/optimagic-dev/optimagic/pull/553
         stop_time=len(stacked["criterion"]) * [None],
         task=len(stacked["criterion"]) * [None],
-        batches=len(stacked["criterion"]) * [None],
+        batches=list(range(len(stacked["criterion"]))),
     )

tests/optimagic/optimization/test_history.py

@@ -169,7 +169,7 @@ def history(history_data):
 
 
 @pytest.fixture
-def history_with_batch_data(history_data):
+def history_parallel(history_data):
     data = history_data.copy()
     data["batches"] = [0, 0, 1, 1, 2, 2]
     return History(direction=Direction.MINIMIZE, **data)
@@ -223,8 +223,8 @@ def test_history_fun_data_with_fun_evaluations_cost_model_and_monotone(history):
     assert_frame_equal(got, exp, check_dtype=False, check_categorical=False)
 
 
-def test_history_fun_data_with_fun_batches_cost_model(history_with_batch_data):
-    got = history_with_batch_data.fun_data(
+def test_history_fun_data_with_fun_batches_cost_model(history_parallel):
+    got = history_parallel.fun_data(
         cost_model=om.timing.fun_batches,
         monotone=False,
     )
@@ -385,8 +385,8 @@ def test_get_time_fun_batches(history):
     assert_array_equal(got, exp)
 
 
-def test_get_time_fun_batches_with_batch_data(history_with_batch_data):
-    got = history_with_batch_data._get_time(cost_model=om.timing.fun_batches)
+def test_get_time_fun_batches_parallel(history_parallel):
+    got = history_parallel._get_time(cost_model=om.timing.fun_batches)
     exp = np.array([1, 1, 2, 2, 3, 3])
     assert_array_equal(got, exp)