Additional plot functions

examples/test_wl_del_freq.py

@@ -0,0 +1,98 @@
+"""
+This example presents potential of new plot functions.
+Especially:
+    -different scales(linear and logarithmic)
+    -different x arguments(delay, frequency, wavelength)
+    -color map as argument of plot function
+    -slice plot for chosen z(propagation) distances
+    or z=0 and z=end if no specific z were chosen.
+Data used in this example are taken from test_Dudley.py file from examples.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+import gnlse
+
+if __name__ == '__main__':
+    setup = gnlse.GNLSESetup()
+
+    # Numerical parameters
+    setup.resolution = 2**14
+    setup.time_window = 12.5  # ps
+    setup.z_saves = 200
+
+    # Physical parameters
+    setup.wavelength = 835  # nm
+    setup.fiber_length = 0.15  # m
+    setup.nonlinearity = 0.11  # 1/W/m
+    setup.raman_model = gnlse.raman_blowwood
+    setup.self_steepening = True
+
+    # The dispersion model is built from a Taylor expansion with coefficients
+    # given below.
+    loss = 0
+    betas = np.array([
+        -11.830e-3, 8.1038e-5, -9.5205e-8, 2.0737e-10, -5.3943e-13, 1.3486e-15,
+        -2.5495e-18, 3.0524e-21, -1.7140e-24
+    ])
+    setup.dispersion_model = gnlse.DispersionFiberFromTaylor(loss, betas)
+
+    # Input pulse parameters
+    power = 10000
+    # pulse duration [ps]
+    tfwhm = 0.05
+    # hyperbolic secant
+    setup.pulse_model = gnlse.SechEnvelope(power, tfwhm)
+    solver = gnlse.GNLSE(setup)
+    solution = solver.run()
+
+    plt.figure(figsize=(14, 8), facecolor='w', edgecolor='k')
+
+    plt.subplot(4, 3, 1)
+    gnlse.plot_delay_vs_distance(solution, time_range=[-.5, 5], cmap="jet")
+
+    plt.subplot(4, 3, 2)
+    gnlse.plot_frequency_vs_distance(solution, frequency_range=[-300, 200],
+                                     cmap="plasma")
+
+    plt.subplot(4, 3, 3)
+    gnlse.plot_wavelength_vs_distance(solution, WL_range=[400, 1400])
+
+    plt.subplot(4, 3, 4)
+    gnlse.plot_delay_vs_distance_logarithmic(solution, time_range=[-.5, 5],
+                                             cmap="jet")
+
+    plt.subplot(4, 3, 5)
+    gnlse.plot_frequency_vs_distance_logarithmic(solution,
+                                                 frequency_range=[-300, 200],
+                                                 cmap="plasma")
+
+    plt.subplot(4, 3, 6)
+    gnlse.plot_wavelength_vs_distance_logarithmic(solution,
+                                                  WL_range=[400, 1400])
+
+    plt.subplot(4, 3, 7)
+    gnlse.plot_delay_for_distance_slice(solution, time_range=[-.5, 5])
+
+    plt.subplot(4, 3, 8)
+    gnlse.plot_frequency_for_distance_slice(solution,
+                                            frequency_range=[-300, 200])
+
+    plt.subplot(4, 3, 9)
+    gnlse.plot_wavelength_for_distance_slice(solution, WL_range=[400, 1400])
+
+    plt.subplot(4, 3, 10)
+    gnlse.plot_delay_for_distance_slice_logarithmic(
+        solution, time_range=[-.5, 5])
+
+    plt.subplot(4, 3, 11)
+    gnlse.plot_frequency_for_distance_slice_logarithmic(
+        solution, frequency_range=[-300, 200])
+
+    plt.subplot(4, 3, 12)
+    gnlse.plot_wavelength_for_distance_slice_logarithmic(solution,
+                                                         WL_range=[400, 1400])
+
+    plt.tight_layout()
+    plt.show()

gnlse/__init__.py

@@ -7,14 +7,37 @@
 from gnlse.nonlinearity import NonlinearityFromEffectiveArea
 from gnlse.raman_response import (raman_blowwood, raman_holltrell,
                                   raman_linagrawal)
-from gnlse.visualization import (plot_delay_vs_distance,
-                                 plot_wavelength_vs_distance, quick_plot)
+from gnlse.visualization import (
+    plot_delay_vs_distance,
+    plot_delay_vs_distance_logarithmic,
+    plot_delay_for_distance_slice,
+    plot_delay_for_distance_slice_logarithmic,
+    plot_frequency_vs_distance,
+    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,
+    plot_wavelength_for_distance_slice,
+    plot_wavelength_for_distance_slice_logarithmic,
+    quick_plot)
 
 __all__ = [
     'DispersionFiberFromTaylor', 'DispersionFiberFromInterpolation',
     'SechEnvelope', 'GaussianEnvelope', 'LorentzianEnvelope', 'GNLSESetup',
     'GNLSE', 'Solution', 'read_mat', 'write_mat', 'raman_blowwood',
-    'raman_holltrell', 'raman_linagrawal', 'plot_delay_vs_distance',
-    'plot_wavelength_vs_distance', 'quick_plot',
-    'NonlinearityFromEffectiveArea', 'CWEnvelope'
+    'raman_holltrell', 'raman_linagrawal',
+    'plot_delay_vs_distance',
+    'plot_delay_vs_distance_logarithmic',
+    'plot_delay_for_distance_slice',
+    'plot_delay_for_distance_slice_logarithmic',
+    'plot_frequency_vs_distance',
+    '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',
+    'plot_wavelength_for_distance_slice',
+    'plot_wavelength_for_distance_slice_logarithmic',
+    'quick_plot', 'NonlinearityFromEffectiveArea', 'CWEnvelope'
 ]

gnlse/gnlse.py

@@ -79,10 +79,11 @@ class Solution:
         Intermediate steps in the frequency domain.
     """
 
-    def __init__(self, t=None, W=None, Z=None, At=None, AW=None,
+    def __init__(self, t=None, W=None, w_0=None, Z=None, At=None, AW=None,
                  Aty=None, AWy=None):
         self.t = t
         self.W = W
+        self.w_0 = w_0
         self.Z = Z
         self.At = At
         self.AW = AW
@@ -163,26 +164,26 @@ def __init__(self, setup):
                                        self.N / 2
                                        ) / (self.N * (self.t[1] - self.t[0]))
         # Central angular frequency [10^12 rad]
-        w_0 = (2.0 * np.pi * c) / setup.wavelength
-        self.Omega = self.V + w_0
+        self.w_0 = (2.0 * np.pi * c) / setup.wavelength
+        self.Omega = self.V + self.w_0
 
         # Absolute angular frequency grid
-        if setup.self_steepening and np.abs(w_0) > np.finfo(float).eps:
-            W = self.V + w_0
+        if setup.self_steepening and np.abs(self.w_0) > np.finfo(float).eps:
+            W = self.V + self.w_0
         else:
-            W = np.full(self.V.shape, w_0)
+            W = np.full(self.V.shape, self.w_0)
         self.W = np.fft.fftshift(W)
 
         # Nonlinearity
         if hasattr(setup.nonlinearity, 'gamma'):
             # in case in of frequency dependent nonlinearity
             gamma, self.scale = setup.nonlinearity.gamma(self.V)
-            self.gamma = gamma / w_0
+            self.gamma = gamma / self.w_0
             self.gamma = np.fft.fftshift(self.gamma)
             self.scale = np.fft.fftshift(self.scale)
         else:
             # in case in of direct introduced value
-            self.gamma = setup.nonlinearity / w_0
+            self.gamma = setup.nonlinearity / self.w_0
             self.scale = 1
 
         # Raman scattering
@@ -276,4 +277,4 @@ def rhs(z, AW):
             At[i, :] = np.fft.fft(AW[i, :])
             AW[i, :] = np.fft.fftshift(AW[i, :]) * self.N * dt
 
-        return Solution(self.t, self.Omega, Z, At, AW)
+        return Solution(self.t, self.Omega, self.w_0, Z, At, AW)