# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Storage
#
# This notebook illustrates how simulations and results can be saved to file.

# +
import pypesto
import pypesto.optimize as optimize
import pypesto.visualize as visualize
from pypesto.store import (save_to_hdf5, read_from_hdf5)

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import tempfile

# %matplotlib inline
# -

# ## Define the objective and problem

# +
objective = pypesto.Objective(fun=sp.optimize.rosen, 
                              grad=sp.optimize.rosen_der, 
                              hess=sp.optimize.rosen_hess)

dim_full = 10
lb = -3 * np.ones((dim_full, 1))
ub = 3 * np.ones((dim_full, 1))

problem = pypesto.Problem(objective=objective, lb=lb, ub=ub)

# create optimizers
optimizer = optimize.ScipyOptimizer(method='l-bfgs-b')

# set number of starts
n_starts = 20
# -

# ## Objective function traces
#
# During optimization, it is possible to regularly write the objective function trace to file. This is useful e.g. when runs fail, or for various diagnostics. Currently, pyPESTO can save histories to 3 backends: in-memory, as CSV files, or to HDF5 files.

# ### Memory History

# To record the history in-memory, just set `trace_record=True` in the `pypesto.HistoryOptions`. Then, the optimization result contains those histories:

# +
# record the history
history_options = pypesto.HistoryOptions(trace_record=True)

# Run optimizaitons
result = optimize.minimize(
    problem=problem, optimizer=optimizer,
    n_starts=n_starts, history_options=history_options)
# -

# Now, in addition to queries on the result, we can also access the

# +
print("History type: ", type(result.optimize_result.list[0].history))
# print("Function value trace of best run: ", result.optimize_result.list[0].history.get_fval_trace())

fig, ax = plt.subplots(1, 2)
visualize.waterfall(result, ax=ax[0])
visualize.optimizer_history(result, ax=ax[1])
fig.set_size_inches((15, 5))
# -

# ### CSV History

# The in-memory storage is however not stored anywhere. To do that, it is possible to store either to CSV or HDF5. This is specified via the `storage_file` option. If it ends in `.csv`, a `pypesto.objective.history.CsvHistory` will be employed; if it ends in `.hdf5` a `pypesto.objective.history.Hdf5History`. Occurrences of the substring `{id}` in the filename are replaced by the multistart id, allowing to maintain a separate file per run (this is necessary for CSV as otherwise runs are overwritten).

# +
# record the history and store to CSV
history_options = pypesto.HistoryOptions(trace_record=True, storage_file='history_{id}.csv')

# Run optimizaitons
result = optimize.minimize(
    problem=problem, optimizer=optimizer,
    n_starts=n_starts, history_options=history_options)
# -

# Note that for this simple cost function, saving to CSV takes a considerable amount of time. This overhead decreases for more costly simulators, e.g. using ODE simulations via AMICI.

# +
print("History type: ", type(result.optimize_result.list[0].history))
# print("Function value trace of best run: ", result.optimize_result.list[0].history.get_fval_trace())

fig, ax = plt.subplots(1, 2)
visualize.waterfall(result, ax=ax[0])
visualize.optimizer_history(result, ax=ax[1])
fig.set_size_inches((15, 5))
# -

# ### HDF5 History

# + [markdown] pycharm={"name": "#%% md\n"}
# Just as in CSV, writing the history to HDF5 takes a considerable amount of time.
# If a user specifies a HDF5 output file named `my_results.hdf5` and uses a parallelization engine, then:
# * a folder is created to contain partial results, named `my_results/` (the stem of the output filename)
# * files are created to store the results of each start, named `my_results/my_results_{START_INDEX}.hdf5`
# * a file is created to store the combined result from all starts, named `my_results.hdf5`.
# Note that this file depends on the files in the `my_results/` directory, so **cease to function** if
# `my_results/` is deleted.

# + pycharm={"name": "#%%\n"}
# record the history and store to CSV
history_options = pypesto.HistoryOptions(trace_record=True, storage_file='history.hdf5')

# Run optimizaitons
result = optimize.minimize(
    problem=problem, optimizer=optimizer,
    n_starts=n_starts, history_options=history_options)

# + pycharm={"name": "#%%\n"}
print("History type: ", type(result.optimize_result.list[0].history))
# print("Function value trace of best run: ", result.optimize_result.list[0].history.get_fval_trace())

fig, ax = plt.subplots(1, 2)
visualize.waterfall(result, ax=ax[0])
visualize.optimizer_history(result, ax=ax[1])
fig.set_size_inches((15, 5))

# + [markdown] pycharm={"name": "#%% md\n"}
# ## Result storage
#
# Result objects can be stored as HDF5 files. When applicable, this is preferable to just pickling results, which is
# not guaranteed to be reproducible in the future.

# + pycharm={"name": "#%%\n"}
# Run optimizaitons
result = optimize.minimize(
    problem=problem, optimizer=optimizer,
    n_starts=n_starts)

# + pycharm={"name": "#%%\n"}
result.optimize_result.list[0:2]

# + [markdown] pycharm={"name": "#%% md\n"}
# As usual, having obtained our result, we can directly perform some plots:

# + pycharm={"name": "#%%\n"}
# plot waterfalls
visualize.waterfall(result, size=(15,6))
# -

# ### Save optimization result as HDF5 file
#
# The optimization result can be saved with `pypesto.store.write_result()`. This will write the problem and the
# optimization result, and the profiling and sampling results if available, to HDF5.
# All of them can be disabled with boolean flags
# (see [the documentation](https://pypesto.readthedocs.io/en/latest/api_store.html#pypesto.store.write_result))

# + pycharm={"name": "#%%\n"}
fn = tempfile.mktemp(".hdf5")

# Write result
save_to_hdf5.write_result(result, fn)

# + [markdown] pycharm={"name": "#%% md\n"}
# ### Read optimization result from HDF5 file
#
# When reading in the stored result again, we recover the original optimization result:

# + pycharm={"name": "#%%\n"}
# Read result and problem
result = read_from_hdf5.read_result(fn)
# -

result.optimize_result.list[0:2]

# plot waterfalls
pypesto.visualize.waterfall(result, size=(15,6))

StarcoderdataJupyterScriptsDedupFiltered