update energy shovel

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)