Issue #95 fixed and improved plot_loss (#106)

* plotter.py plot_loss fixed Issue #95
* Improved the plot_loss to show title and labels.

Author: Bovhasselt

pina/plotter.py

@@ -186,24 +186,37 @@ def plot(self, pinn, components=None, fixed_variables={}, method='contourf',
         else:
             plt.show()
 
-    def plot_loss(self, pinn, label=None, log_scale=True):
+    def plot_loss(self, pinn, label=None, log_scale=True, filename=None):
         """
         Plot the loss function values during traininig.
 
         :param PINN pinn: the PINN object.
         :param str label: the label to use in the legend, defaults to None.
         :param bool log_scale: If True, the y axis is in log scale. Default is
             True.
+        :param str filename: the file name to save the plot. If None, the plot
+            is not saved. Default is None.
         """
 
         if not label:
             label = str(pinn)
 
         epochs = list(pinn.history_loss.keys())
         loss = np.array(list(pinn.history_loss.values()))
+
+        # if multiple outputs, sum the loss
         if loss.ndim != 1:
-            loss = loss[:, 0]
+            loss = np.sum(loss, axis=1)
 
+        # plot loss
         plt.plot(epochs, loss, label=label)
+        plt.legend()
         if log_scale:
             plt.yscale('log')
+        plt.title('Loss function')
+        plt.xlabel('Epochs')
+        plt.ylabel('Loss')
+
+        # save plot
+        if filename:
+            plt.savefig(filename)