Additional plot functions

WUST-FOG · Jan 13, 2023 · cab0dcd · cab0dcd
1 parent af2d3c4
commit cab0dcd
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 90 deletions.
diff --git a/gnlse/__init__.py b/gnlse/__init__.py
@@ -23,7 +23,7 @@
     'DispersionFiberFromTaylor', 'DispersionFiberFromInterpolation',
     'SechEnvelope', 'GaussianEnvelope', 'LorentzianEnvelope', 'GNLSESetup',
     'GNLSE', 'Solution', 'read_mat', 'write_mat', 'raman_blowwood',
-    'raman_holltrell', 'raman_linagrawal', 
+    'raman_holltrell', 'raman_linagrawal',
     'plot_delay_vs_distance',
     'plot_delay_vs_distance_logarithmic',
     'plot_delay_for_distance_slice',
@@ -32,7 +32,7 @@
     'plot_frequency_vs_distance_logarithmic',
     'plot_frequency_for_distance_slice',
     'plot_frequency_for_distance_slice_logarithmic',
-    'plot_wavelength_vs_distance', 
-    'plot_wavelength_vs_distance_logarithmic', 
-    'quick_plot','NonlinearityFromEffectiveArea', 'CWEnvelope'
+    'plot_wavelength_vs_distance',
+    'plot_wavelength_vs_distance_logarithmic',
+    'quick_plot', 'NonlinearityFromEffectiveArea', 'CWEnvelope'
 ]
diff --git a/gnlse/visualization.py b/gnlse/visualization.py
@@ -43,12 +43,13 @@ def plot_frequency_vs_distance_logarithmic(solver, ax=None, norm=None):
     frequency = (solver.W - solver.w_0) / 2 / np.pi  # frequency grid
 
     ax.imshow(lIW, origin='lower', aspect='auto', cmap="magma",
-              extent=[np.min(frequency), np.max(frequency), 0, np.max(solver.Z)],
-              vmin=-40)
+              extent=[np.min(frequency), np.max(frequency),
+                      0, np.max(solver.Z)], vmin=-40)
     ax.set_xlabel("Frequency [THz]")
     ax.set_ylabel("Distance [m]")
     return ax
 
+
 def plot_frequency_vs_distance(solver, ax=None, norm=None):
     """Plotting results in frequency domain. Linear scale.
 
@@ -75,22 +76,23 @@ def plot_frequency_vs_distance(solver, ax=None, norm=None):
 
     if norm is None:
         norm = np.max(np.abs(solver.AW)**2)
- 
+
     IW = np.fliplr(
         np.abs(solver.AW)**2 / norm)
     frequency = (solver.W - solver.w_0) / 2 / np.pi  # frequency grid
-    
+
     ax.imshow(IW, origin='lower', aspect='auto', cmap="magma",
-              extent=[np.min(frequency), np.max(frequency), 0, np.max(solver.Z)],
-              vmin=0)
+              extent=[np.min(frequency), np.max(frequency),
+                      0, np.max(solver.Z)], vmin=0)
     ax.set_xlabel("Frequency [THz]")
     ax.set_ylabel("Distance [m]")
     return ax
 
+
 def plot_delay_for_distance_slice(solver, time_range=None, ax=None,
-                                z_slice = None, norm=None):
+                                  z_slice=None, norm=None):
     """Plotting intensity in linear scale in time domain.
-    
+
     Parameters
     ----------
     solver : Solution
@@ -102,7 +104,7 @@ def plot_delay_for_distance_slice(solver, time_range=None, ax=None,
     norm : float
         Normalization factor for output spectrum. As default maximum of
         square absolute of ``solver.At`` variable is taken.
-    
+
     Returns
     -------
     ax : :class:`~matplotlib.axes.Axes`
@@ -114,36 +116,37 @@ def plot_delay_for_distance_slice(solver, time_range=None, ax=None,
 
     if time_range is None:
         time_range = [np.min(solver.t), np.max(solver.t)]
-    
+
     if norm is None:
         norm = np.max(np.abs(solver.At)**2)
-
-    It = np.fliplr(
-        np.abs(solver.At)**2 / norm) 
-
 
-    #indices of interest if no z_slice positions were given      
-    if z_slice is None: 
-        iis = [0,-1] 
-    #indices of interest nearest to given z_slice positions    
-    else:    
-        iis = [ np.nonzero(np.min(np.abs(solver.Z-z)) == np.abs(solver.Z-z))[0][0] for z in z_slice ] 
+    It = np.fliplr(
+        np.abs(solver.At)**2 / norm)
 
-    for i in iis: 
-        label_i = "z = " + str(solver.Z[i]) +"m"
-        ax.plot(solver.t, It[i][:], label = label_i)  
+    # indices of interest if no z_slice positions were given
+    if z_slice is None:
+        iis = [0, -1]
+    # indices of interest nearest to given z_slice positions
+    else:
+        iis = [np.nonzero(
+            np.min(np.abs(solver.Z - z)) == np.abs(solver.Z - z)
+        )[0][0] for z in z_slice]
 
+    for i in iis:
+        label_i = "z = " + str(solver.Z[i]) + "m"
+        ax.plot(solver.t, It[i][:], label=label_i)
 
     ax.set_xlim(time_range)
     ax.set_xlabel("Delay [ps]")
     ax.set_ylabel("Normalized Power")
     ax.legend()
     return ax
 
+
 def plot_delay_for_distance_slice_logarithmic(solver, time_range=None, ax=None,
-                                z_slice = None, norm=None):
+                                              z_slice=None, norm=None):
     """Plotting intensity in logarithmic scale in time domain.
-    
+
     Parameters
     ----------
     solver : Solution
@@ -155,7 +158,7 @@ def plot_delay_for_distance_slice_logarithmic(solver, time_range=None, ax=None,
     norm : float
         Normalization factor for output spectrum. As default maximum of
         square absolute of ``solver.At`` variable is taken.
-    
+
     Returns
     -------
     ax : :class:`~matplotlib.axes.Axes`
@@ -167,24 +170,25 @@ def plot_delay_for_distance_slice_logarithmic(solver, time_range=None, ax=None,
 
     if time_range is None:
         time_range = [np.min(solver.t), np.max(solver.t)]
-    
+
     if norm is None:
         norm = np.max(np.abs(solver.At)**2)
-  
+
     lIt = 10 * np.log10(np.abs(solver.At)**2 / norm,
                         where=(np.abs(solver.At)**2 > 0))
 
-    #indices of interest if no z_slice positions were given      
-    if z_slice is None: 
-        iis = [0,-1] 
-    #indices of interest nearest to given z_slice positions    
-    else:    
-        iis = [ np.nonzero(np.min(np.abs(solver.Z-z)) == np.abs(solver.Z-z))[0][0] for z in z_slice ] 
-
-    for i in iis: 
-        label_i = "z = " + str(solver.Z[i]) +"m"
-        ax.plot(solver.t, lIt[i][:], label = label_i)  
+    # indices of interest if no z_slice positions were given
+    if z_slice is None:
+        iis = [0, -1]
+    # indices of interest nearest to given z_slice positions
+    else:
+        iis = [np.nonzero(
+            np.min(np.abs(solver.Z - z)) == np.abs(solver.Z - z)
+        )[0][0] for z in z_slice]
 
+    for i in iis:
+        label_i = "z = " + str(solver.Z[i]) + "m"
+        ax.plot(solver.t, lIt[i][:], label=label_i)
 
     ax.set_xlim(time_range)
     ax.set_ylim(-40)
@@ -193,9 +197,10 @@ def plot_delay_for_distance_slice_logarithmic(solver, time_range=None, ax=None,
     ax.legend()
     return ax
 
+
 def plot_frequency_for_distance_slice(solver, frequency_range=None, ax=None,
-                                                 z_slice = None, norm=None):
-    """Plotting chosen slices of intensity in linear scale in frequency domain. 
+                                      z_slice=None, norm=None):
+    """Plotting chosen slices of intensity in linear scale in frequency domain.
 
     Parameters
     ----------
@@ -219,36 +224,40 @@ def plot_frequency_for_distance_slice(solver, frequency_range=None, ax=None,
         ax = plt.gca()
 
     if frequency_range is None:
-        frequency_range = [np.min((solver.W-solver.w_0)/2/np.pi),
-                           np.max((solver.W-solver.w_0)/2/np.pi)]
- 
+        frequency_range = [np.min((solver.W - solver.w_0) / 2 / np.pi),
+                           np.max((solver.W - solver.w_0) / 2 / np.pi)]
+
     if norm is None:
         norm = np.max(np.abs(solver.AW)**2)
- 
+
     IW = np.fliplr(
         np.abs(solver.AW)**2 / norm)
 
-    #indices of interest if no z_slice positions were given  
+    # indices of interest if no z_slice positions were given
     if z_slice is None:
-        iis = [0,-1] #beginning, end
-    #indices of interest nearest to given z_slice positions 
-    else:  
-        iis = [ np.nonzero(np.min(np.abs(solver.Z-z)) == np.abs(solver.Z-z))[0][0] for z in z_slice ] 
-
-    for i in iis: 
-        label_i = "z = " + str(solver.Z[i]) +"m"
-        ax.plot((solver.W-solver.w_0)/2/np.pi, IW[i][:], label = label_i)  
+        iis = [0, -1]  # beginning, end
+    # indices of interest nearest to given z_slice positions
+    else:
+        iis = [np.nonzero(
+            np.min(np.abs(solver.Z - z)) == np.abs(solver.Z - z)
+        )[0][0] for z in z_slice]
 
+    for i in iis:
+        label_i = "z = " + str(solver.Z[i]) + "m"
+        ax.plot((solver.W - solver.w_0) / 2 / np.pi, IW[i][:], label=label_i)
 
     ax.set_xlim(frequency_range)
     ax.set_xlabel("Frequency [Thz]")
     ax.set_ylabel("Normalized Power")
     ax.legend()
-    return ax   
+    return ax
+
 
-def plot_frequency_for_distance_slice_logarithmic(solver, frequency_range=None, 
-                                           ax=None, z_slice = None, norm=None):
-    """Plotting chosen slices of intensity in logarithmic scale in frequency domain. 
+def plot_frequency_for_distance_slice_logarithmic(solver, frequency_range=None,
+                                                  ax=None, z_slice=None,
+                                                  norm=None):
+    """Plotting chosen slices of intensity
+    in logarithmic scale in frequency domain.
 
     Parameters
     ----------
@@ -272,38 +281,39 @@ def plot_frequency_for_distance_slice_logarithmic(solver, frequency_range=None,
         ax = plt.gca()
 
     if frequency_range is None:
-        frequency_range = [np.min((solver.W-solver.w_0)/2/np.pi),
-                           np.max((solver.W-solver.w_0)/2/np.pi)]
- 
+        frequency_range = [np.min((solver.W - solver.w_0) / 2 / np.pi),
+                           np.max((solver.W - solver.w_0) / 2 / np.pi)]
+
     if norm is None:
         norm = np.max(np.abs(solver.AW)**2)
- 
+
     lIW = np.fliplr(
         10 * np.log10(np.abs(solver.AW)**2 / norm,
                       where=(np.abs(solver.AW)**2 > 0)))
-    
-    #indices of interest if no z_slice positions were given
+
+    # indices of interest if no z_slice positions were given
     if z_slice is None:
-        iis = [0,-1] #beginning, end
-    
-    #indices of interest nearest to given z_slice positions 
-    else:  
-        iis = [ np.nonzero(np.min(np.abs(solver.Z-z)) == np.abs(solver.Z-z))[0][0] for z in z_slice ] 
-
-
-    for i in iis: 
-        label_i = "z = " + str(solver.Z[i]) +"m"
-        ax.plot((solver.W-solver.w_0)/2/np.pi, lIW[i][:], label = label_i)  
-
-
-    ax.set_xlim(frequency_range) 
+        iis = [0, -1]  # beginning, end
+
+    # indices of interest nearest to given z_slice positions
+    else:
+        iis = [np.nonzero(
+            np.min(np.abs(solver.Z - z)) == np.abs(solver.Z - z)
+        )[0][0] for z in z_slice]
+
+    for i in iis:
+        label_i = "z = " + str(solver.Z[i]) + "m"
+        ax.plot((solver.W - solver.w_0) / 2 / np.pi, lIW[i][:], label=label_i)
+
+    ax.set_xlim(frequency_range)
     ax.set_xlabel("Frequency [Thz]")
     ax.set_ylabel("Normalized Power")
     ax.legend()
-    return ax   
-
-def plot_delay_vs_distance_logarithmic(solver, time_range=None, ax=None, 
-                                                             norm=None):
+    return ax
+
+
+def plot_delay_vs_distance_logarithmic(solver, time_range=None, ax=None,
+                                       norm=None):
     """Plotting intensity in logarithmic scale in time domain.
 
     Parameters
@@ -342,6 +352,7 @@ def plot_delay_vs_distance_logarithmic(solver, time_range=None, ax=None,
     ax.set_ylabel("Distance [m]")
     return ax
 
+
 def plot_delay_vs_distance(solver, time_range=None, ax=None, norm=None):
     """Plotting normalized intensity in linear scale in time domain.
 
@@ -373,13 +384,14 @@ def plot_delay_vs_distance(solver, time_range=None, ax=None, norm=None):
 
     lIT = np.abs(solver.At)**2 / norm
 
-    ax.pcolormesh(solver.t, solver.Z, lIT, shading="auto", vmin=0,  
+    ax.pcolormesh(solver.t, solver.Z, lIT, shading="auto", vmin=0,
                   cmap="jet")
     ax.set_xlim(time_range)
     ax.set_xlabel("Delay [ps]")
     ax.set_ylabel("Distance [m]")
     return ax
 
+
 def plot_wavelength_vs_distance(solver, WL_range=None, ax=None,
                                 norm=None):
     """Plotting results in linear scale in wavelength domain.
@@ -406,8 +418,8 @@ def plot_wavelength_vs_distance(solver, WL_range=None, ax=None,
         ax = plt.gca()
 
     if WL_range is None:
-        WL_range = [np.min(c/(solver.W/2/np.pi)), 
-                    np.max(c/(solver.W/2/np.pi))]
+        WL_range = [np.min(c / (solver.W / 2 / np.pi)),
+                    np.max(c / (solver.W / 2 / np.pi))]
 
     if norm is None:
         norm = np.max(np.abs(solver.AW)**2)
@@ -433,8 +445,9 @@ def plot_wavelength_vs_distance(solver, WL_range=None, ax=None,
     ax.set_ylabel("Distance [m]")
     return ax
 
+
 def plot_wavelength_vs_distance_logarithmic(solver, WL_range=None,
-                                              ax=None, norm=None):
+                                            ax=None, norm=None):
     """Plotting results in logarithmic scale in wavelength domain.
 
     Parameters
@@ -459,8 +472,8 @@ def plot_wavelength_vs_distance_logarithmic(solver, WL_range=None,
         ax = plt.gca()
 
     if WL_range is None:
-        WL_range = [np.min(c/(solver.W/2/np.pi)), 
-                    np.max(c/(solver.W/2/np.pi))]
+        WL_range = [np.min(c / (solver.W / 2 / np.pi)),
+                    np.max(c / (solver.W / 2 / np.pi))]
 
     if norm is None:
         norm = np.max(np.abs(solver.AW)**2)
@@ -487,6 +500,7 @@ def plot_wavelength_vs_distance_logarithmic(solver, WL_range=None,
     ax.set_ylabel("Distance [m]")
     return ax
 
+
 def quick_plot(solution):
     """Plotting results in time and frequency domain for default value
     of parameters.