diff --git a/src/optimagic/optimization/history.py b/src/optimagic/optimization/history.py index 4016d4152..2b0215722 100644 --- a/src/optimagic/optimization/history.py +++ b/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( diff --git a/src/optimagic/optimization/process_results.py b/src/optimagic/optimization/process_results.py index 8da70059c..64d764174 100644 --- a/src/optimagic/optimization/process_results.py +++ b/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], ) diff --git a/src/optimagic/visualization/history_plots.py b/src/optimagic/visualization/history_plots.py index 73069c522..22d98bf71 100644 --- a/src/optimagic/visualization/history_plots.py +++ b/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"]))), ) diff --git a/tests/optimagic/optimization/test_history.py b/tests/optimagic/optimization/test_history.py index bd7137922..72cdbb4a2 100644 --- a/tests/optimagic/optimization/test_history.py +++ b/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)