diff --git a/metapeg_v2.py b/metapeg_v2.py
index e0c23b0..a6de29c 100644
--- a/metapeg_v2.py
+++ b/metapeg_v2.py
@@ -745,15 +745,17 @@ if __name__ == "__main__":
 
     plt.figure()
     plt.subplot(1,2,1)
-    plt.plot(objfunc_history[:,0], label="f1 - 1.50 um")
-    plt.plot(objfunc_history[:,1], label="f2 - 1.55 um")
-    plt.plot(objfunc_history[:,2], label="f3 - 1.60 um")
+    line1,line2,line3, = plt.plot(objfunc_history)
+    plt.xlabel("Iteration")
+    plt.ylabel("Value")
+    plt.legend([line1, line2, line3], ["Obj. func. - 1.50 um", "Obj. func. - 1.55 um", "Obj. func. - 1.60 um"])
     plt.subplot(1,2,2)
     plt.plot(epivar_history, label="Epigraph")
     plt.xlabel("Iteration")
     plt.ylabel("Value")
     plt.legend()
-    plt.savefig("results/optimization_history.pdf")
+    plt.savefig(f"results/optimization_history_beta{sigmoid_bias}.pdf")
+    plt.close()
 
     saveResults = True
     if saveResults:
diff --git a/metapeg_v2_unc.py b/metapeg_v2_unc.py
index 037cacc..bc36cc0 100644
--- a/metapeg_v2_unc.py
+++ b/metapeg_v2_unc.py
@@ -749,15 +749,16 @@ if __name__ == "__main__":
 
             plt.figure()
             plt.subplot(1,2,1)
-            plt.plot(objfunc_history[:,0], label="f1 - 1.50 um")
-            plt.plot(objfunc_history[:,1], label="f2 - 1.55 um")
-            plt.plot(objfunc_history[:,2], label="f3 - 1.60 um")
+            line1,line2,line3, = plt.plot(objfunc_history)
+            plt.xlabel("Iteration")
+            plt.ylabel("Value")
+            plt.legend([line1, line2, line3], ["Obj. func. - 1.50 um", "Obj. func. - 1.55 um", "Obj. func. - 1.60 um"])
             plt.subplot(1,2,2)
             plt.plot(epivar_history, label="Epigraph")
             plt.xlabel("Iteration")
             plt.ylabel("Value")
             plt.legend()
-            plt.savefig("results/optimization_history.pdf")
+            plt.savefig(f"results/optimization_history_beta{sigmoid_bias}.pdf")
             plt.close()
 
     saveResults = True