Better plotting for the example examples/DeepONet/ROBER/eval_rober.py

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>
IBM · Mar 23, 2023 · 106a763 · 106a763
1 parent 5f8ead2
commit 106a763
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/examples/DeepONet/ROBER/eval_rober.py b/examples/DeepONet/ROBER/eval_rober.py
@@ -58,9 +58,13 @@
 plt.plot(time, evaluation, label="Approximated")
 plt.xlabel("t (s)")
 plt.savefig("rober_approximation.png")
+plt.close()
 
 for i in range(n_outputs):
-    plt.plot(time, evaluation[:, i])
+    plt.plot(time, evaluation[:, i], label=f"s_{i+1}")
     plt.xlabel("t (s)")
-    plt.savefig(f"rober_approximation_{i}.png")
-    plt.close()
+
+plt.yticks(np.linspace(0,1,5))
+plt.legend()
+plt.grid(True)
+plt.savefig(f"rober_approximation_custom.png")