diff --git a/python/energy_flow_proper/08_dynamic_one_bath/energy_shovel.py b/python/energy_flow_proper/08_dynamic_one_bath/energy_shovel.py index f01a9e5..280b52a 100644 --- a/python/energy_flow_proper/08_dynamic_one_bath/energy_shovel.py +++ b/python/energy_flow_proper/08_dynamic_one_bath/energy_shovel.py @@ -16,6 +16,8 @@ import scipy import utilities as ut import matplotlib import itertools +import pathlib +import shutil from hops.util.dynamic_matrix import ( SmoothStep, @@ -89,7 +91,7 @@ def models_table(models, **kwargs): def make_row(title, accessor): row = [title] for model, data in aux.model_data_iterator(models): - row += [f"${accessor(model, data)}$"] + row += [f"${accessor(model, data):.2f}$"] return [row] @@ -98,6 +100,7 @@ def models_table(models, **kwargs): table += make_row(r"$T$", lambda m, d: m.T) table += make_row(r"$N$", lambda m, d: d.samples) table += make_row(r"$k_{\mathrm{max}}$", lambda m, d: m.k_max) + table += make_row(r"BCF Terms", lambda m, d: m.bcf_terms) return tabulate(table, **(dict(tablefmt="latex_raw") | kwargs)) @@ -146,3 +149,41 @@ def energy_friction_plot(models, strobe_frequency, strobe_indices): ) return fig, (ax, ax2) + + +def optimize_for_interaction_energy( + goal, periods_for_optimization, samples_for_optimization, **model_args +): + opt_args = model_args | dict(periods=periods_for_optimization) + proto, _, _ = energy_shovel(**opt_args) + + def cost(δ): + proto.δ = δ + aux.integrate(proto, samples_for_optimization) + with aux.get_data(proto) as data: + inter = proto.interaction_energy(data) + + return (abs(inter.value).max() - goal) ** 2 + + old_cache_path = pathlib.Path(f"./.cache/{proto.hexhash}_delta_{goal}.npy") + cache_path = pathlib.Path( + f"./.cache/{proto.hexhash}_delta_{goal}_{periods_for_optimization}_{samples_for_optimization}.npy" + ) + + if old_cache_path.exists(): + shutil.copy(old_cache_path, cache_path) + + if cache_path.exists(): + with cache_path.open("rb") as cache: + δ = np.load(cache, allow_pickle=True) + else: + δ = scipy.optimize.minimize_scalar( + cost, bounds=(0.1, 3), method="bounded", options=dict(xatol=1e-2, disp=3) + ).x + cache_path.parent.mkdir(parents=True, exist_ok=True) + with cache_path.open("wb") as cache: + np.save(cache, δ) + + model_args["δ"] = δ + + return energy_shovel(**model_args)