From 1bdc6bef421bf403f71e44444b4460e4e0d95577 Mon Sep 17 00:00:00 2001
From: Emmanuel Lujan <lujan.emmanuel@gmail.com>
Date: Thu, 27 Jun 2024 17:09:09 -0400
Subject: [PATCH] Improvements in loss function and plots in hyperparameter
 optimization.

---
 examples/Opt-ACE-aHfO2/fit-opt-ace-ahfo2.jl |  2 +-
 examples/Opt-ACE-aHfO2/utils.jl             | 37 +-----------------
 examples/utils/plots.jl                     | 43 +++++++++++++++++++++
 3 files changed, 46 insertions(+), 36 deletions(-)

diff --git a/examples/Opt-ACE-aHfO2/fit-opt-ace-ahfo2.jl b/examples/Opt-ACE-aHfO2/fit-opt-ace-ahfo2.jl
index 34aa10e6..fca907df 100644
--- a/examples/Opt-ACE-aHfO2/fit-opt-ace-ahfo2.jl
+++ b/examples/Opt-ACE-aHfO2/fit-opt-ace-ahfo2.jl
@@ -60,7 +60,7 @@ pars = OrderedDict( :body_order        => [2, 3, 4],
 # Use random sampling to find the optimal hyper-parameters.
 iap, res = hyperlearn!(model, pars, conf_train; 
                        n_samples = 10, sampler = RandomSampler(),
-                       ws = [1.0, 1.0], int = true)
+                       loss = loss, ws = [1.0, 1.0], int = true)
 
 # Save and show results.
 @save_var res_path iap.β
diff --git a/examples/Opt-ACE-aHfO2/utils.jl b/examples/Opt-ACE-aHfO2/utils.jl
index 112c52d7..4672b8e4 100644
--- a/examples/Opt-ACE-aHfO2/utils.jl
+++ b/examples/Opt-ACE-aHfO2/utils.jl
@@ -26,39 +26,6 @@ function get_results(ho)
     return sort!(results)
 end
 
-# Plot fitting error vs force time (Pareto front)
-function plot_err_time(res)
-    e_mae = res[!, :e_mae]
-    f_mae = res[!, :f_mae]
-    times = res[!, :time_us]
-    plot(times,
-         e_mae,
-         seriestype = :scatter,
-         alpha = 0.55,
-         thickness_scaling = 1.35,
-         markersize = 3,
-         markerstrokewidth = 1,
-         markerstrokecolor = :black,
-         markershape = :circle,
-         markercolor = :gray,
-         label = "MAE(E_Pred, E_DFT)")
-    plot!(times,
-          f_mae,
-          seriestype = :scatter,
-          alpha = 0.55,
-          thickness_scaling = 1.35,
-          markersize = 3,
-          markerstrokewidth = 1,
-          markerstrokecolor = :red,
-          markershape = :utriangle,
-          markercolor = :red2,
-          label = "MAE(F_Pred, F_DFT)")
-    plot!(dpi = 300,
-          label = "",
-          xlabel = "Time per force per atom | µs",
-          ylabel = "MAE")
-end
-
 
 function get_species(confs)
     return unique(vcat(unique.(atomic_symbol.(get_system.(confs)))...))
@@ -68,8 +35,8 @@ create_ho(x) = Hyperoptimizer(1)
 
 # hyperlearn!
 function hyperlearn!(model, pars, conf_train;
-                     n_samples = 5, sampler = RandomSampler(), loss = loss,
-                     ws = [1.0, 1.0], int = true)
+                     n_samples = 5, sampler = RandomSampler(),
+                     loss = loss, ws = [1.0, 1.0], int = true)
 
     s = "create_ho(sampler) = Hyperoptimizer($n_samples, sampler, " *
          join("$k = $v, " for (k, v) in pars) * ")"
diff --git a/examples/utils/plots.jl b/examples/utils/plots.jl
index 356c3156..e29efe31 100644
--- a/examples/utils/plots.jl
+++ b/examples/utils/plots.jl
@@ -238,3 +238,46 @@ function plot_cos(
     return p
 end
 
+
+"""
+    plot_err_time(
+        res
+    )
+
+`res`: Dataframe with fitting error and force time information.
+
+Returns a plot with fitting errors vs force times. Required in hyper-parameter optimization.
+
+"""
+function plot_err_time(res)
+    e_mae = res[!, :e_mae]
+    f_mae = res[!, :f_mae]
+    times = res[!, :time_us]
+    plot(times,
+         e_mae,
+         seriestype = :scatter,
+         alpha = 0.55,
+         thickness_scaling = 1.35,
+         markersize = 3,
+         markerstrokewidth = 1,
+         markerstrokecolor = :black,
+         markershape = :circle,
+         markercolor = :gray,
+         label = "MAE(E_Pred, E_DFT) | eV/atom")
+    plot!(times,
+          f_mae,
+          seriestype = :scatter,
+          alpha = 0.55,
+          thickness_scaling = 1.35,
+          markersize = 3,
+          markerstrokewidth = 1,
+          markerstrokecolor = :red,
+          markershape = :utriangle,
+          markercolor = :red2,
+          label = "MAE(F_Pred, F_DFT) | eV/Å")
+    plot!(dpi = 300,
+          label = "",
+          xlabel = "Time per force per atom | µs",
+          ylabel = "MAE")
+end
+