diff --git a/Dipolar-Gas-Simulator/+Plotter/plotLive.m b/Dipolar-Gas-Simulator/+Plotter/plotLive.m
index 8d98648..4b293c4 100644
--- a/Dipolar-Gas-Simulator/+Plotter/plotLive.m
+++ b/Dipolar-Gas-Simulator/+Plotter/plotLive.m
@@ -1,9 +1,16 @@
-function plotLive(psi,Params,Transf,Observ,SimulationMode)
+function plotLive(psi, Params, Transf, Observ, SimulationMode, varargin)
     set(0,'defaulttextInterpreter','latex')
     set(groot, 'defaultAxesTickLabelInterpreter','latex'); set(groot, 'defaultLegendInterpreter','latex');
     
     format long
 
+    % Check if AdditionalParams is provided
+    if nargin >= 6
+        AdditionalParams = varargin{1}; % Get AdditionalParams from varargin
+    else
+        AdditionalParams = struct(); % Default to an empty struct if not provided
+    end
+
     % Axes scaling and coordinates in micrometers
     x = Transf.x * Params.l0 * 1e6; 
     y = Transf.y * Params.l0 * 1e6;
@@ -122,20 +129,29 @@ function plotLive(psi,Params,Transf,Observ,SimulationMode)
             title('$|\Psi(x,y)|^2$', 'Interpreter', 'latex', 'FontSize', 14) 
         
             nexttile;
-            plot(Observ.theta,'-b')
-            ylabel('$\theta_{opt}$', 'FontSize', 14); xlabel('Time steps', 'FontSize', 14);
-            title('Optimal angle', 'FontSize', 14);
+            switch AdditionalParams.GradientDescentMethod
+                case 'HeavyBall'
+                    plot(-log10(Observ.residual),'-b')
+                    ylabel('$-\mathrm{log}_{10}(r)$', 'FontSize', 14); xlabel('Time steps', 'FontSize', 14);
+                    title('Residual', 'FontSize', 14);
+                case 'NonLinearCGD'
+                    plot(Observ.theta,'-b')
+                    ylabel('$\theta_{opt}$', 'FontSize', 14); xlabel('Steps', 'FontSize', 14);
+                    title('Optimal angle', 'FontSize', 14);
+                otherwise
+                    warning('Energy minimization method not specified - cannot choose between plotting residuals or optimal angles.')
+            end
             grid on
             
             nexttile;
             plot(Observ.EVec,'-b')
-            ylabel('$E_{tot}$', 'FontSize', 14); xlabel('Time steps', 'FontSize', 14);
+            ylabel('$E_{tot}$', 'FontSize', 14); xlabel('Steps', 'FontSize', 14);
             title('Total Energy', 'FontSize', 14);
             grid on
         
             nexttile;
             plot(Observ.mucVec,'-b')
-            ylabel('$\mu$', 'FontSize', 14); xlabel('Time steps', 'FontSize', 14);
+            ylabel('$\mu$', 'FontSize', 14); xlabel('Steps', 'FontSize', 14);
             title('Chemical Potential', 'FontSize', 14);
             grid on
         
diff --git a/Dipolar-Gas-Simulator/+Simulator/@DipolarGas/runGradientDescent.m b/Dipolar-Gas-Simulator/+Simulator/@DipolarGas/runGradientDescent.m
index e3d3d76..e2eb775 100644
--- a/Dipolar-Gas-Simulator/+Simulator/@DipolarGas/runGradientDescent.m
+++ b/Dipolar-Gas-Simulator/+Simulator/@DipolarGas/runGradientDescent.m
@@ -1,6 +1,9 @@
 function [psi]  = runGradientDescent(this,psi,Params,Transf,VDk,V,Observ)
     
     format long;
+
+    AdditionalParams.GradientDescentMethod = this.GradientDescentMethod;
+
     switch this.GradientDescentMethod
         case 'HeavyBall'
             % Convergence Criteria:
@@ -14,7 +17,7 @@ function [psi]  = runGradientDescent(this,psi,Params,Transf,VDk,V,Observ)
             
             % Live Plotter
             if this.PlotLive
-                Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode)
+                Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode, AdditionalParams)
                 drawnow
             end
         
@@ -60,7 +63,7 @@ function [psi]  = runGradientDescent(this,psi,Params,Transf,VDk,V,Observ)
                         Observ.res_idx  = Observ.res_idx + 1;
                         
                         if this.PlotLive
-                            Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode)
+                            Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode, AdditionalParams)
                             drawnow
                         end
             
@@ -94,14 +97,14 @@ function [psi]  = runGradientDescent(this,psi,Params,Transf,VDk,V,Observ)
             i                 = 1;
             Observ.residual   = 1; 
             Observ.res        = 1; 
-            Observ.theta      = 0;
+            Observ.theta      = NaN;
         
             % Initialize the PrematureExitFlag to false
             PrematureExitFlag = false;
             
             % Live plotter
             if this.PlotLive
-                Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode)
+                Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode, AdditionalParams)
                 drawnow
             end
             
@@ -179,7 +182,7 @@ function [psi]  = runGradientDescent(this,psi,Params,Transf,VDk,V,Observ)
                     
                     % Live plotter
                     if this.PlotLive
-                        Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode)
+                        Plotter.plotLive(psi,Params,Transf,Observ,this.SimulationMode, AdditionalParams)
                         drawnow
                     end
                     save(sprintf(strcat(this.SaveDirectory, '/Run_%03i/psi_gs.mat'),Params.njob),'psi','muchem','Observ','Transf','Params','VDk','V');