From 801e440490ca770a363596b50a5ef5cf12813901 Mon Sep 17 00:00:00 2001
From: Rafal-Cichowski <samael97@tlen.pl>
Date: Wed, 30 Nov 2022 15:14:35 +0100
Subject: [PATCH] Plot functions and example of their use.

---
 examples/test_wl_del_freq.py | 31 +++++++++++++------------------
 gnlse/__init__.py            | 27 ++++++++++++++++-----------
 gnlse/visualization.py       |  4 ++--
 3 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/examples/test_wl_del_freq.py b/examples/test_wl_del_freq.py
index af2534e..0839a04 100644
--- a/examples/test_wl_del_freq.py
+++ b/examples/test_wl_del_freq.py
@@ -1,15 +1,12 @@
 """
-This example is here to test new plot functions. This test is based on file
-test_dispersion:
-Example:  Example of supercontinuum generation in anomalous dispersion regime
-at a central wavelength of 835 nm in a 15 centimeter long fiber.
-Comparision of results obtained with two dispersion input:
-1. dispersion calculated from Taylor expansion
-2. dispersion calculated from effective refractive indicies
-
-The python code based on MATLAB code published in
-'Supercontinuum Generation in Optical Fibers'
-by J. M. Dudley and J. R. Taylor, available at http://scgbook.info/.
+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
@@ -17,8 +14,6 @@
 
 import gnlse
 
-import os
-
 if __name__ == '__main__':
     setup = gnlse.GNLSESetup()
 
@@ -48,7 +43,7 @@
     # pulse duration [ps]
     tfwhm = 0.05
     # hyperbolic secant
-    setup.pulse_model = gnlse.SechEnvelope(power, tfwhm) 
+    setup.pulse_model = gnlse.SechEnvelope(power, tfwhm)
     solver = gnlse.GNLSE(setup)
     solution = solver.run()
 
@@ -88,12 +83,12 @@
     gnlse.plot_wavelength_for_distance_slice(solution, WL_range=[400, 1400])
 
     plt.subplot(4, 3, 10)
-    ax = gnlse.plot_delay_for_distance_slice_logarithmic(solution,
-                                                    time_range=[-.5, 5])
+    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])
+    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,
diff --git a/gnlse/__init__.py b/gnlse/__init__.py
index 3f482b7..720c85b 100644
--- a/gnlse/__init__.py
+++ b/gnlse/__init__.py
@@ -7,17 +7,20 @@
 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_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,
-                                 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',
@@ -34,5 +37,7 @@
     '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'
 ]
diff --git a/gnlse/visualization.py b/gnlse/visualization.py
index f368698..56f3784 100644
--- a/gnlse/visualization.py
+++ b/gnlse/visualization.py
@@ -48,7 +48,7 @@ def plot_frequency_vs_distance_logarithmic(solver, ax=None, norm=None,
                              frequency <= frequency_range[1])
         # indices of interest
 
-        frequency= frequency[iis]
+        frequency = frequency[iis]
         lIW = lIW[:, iis]
 
     ax.imshow(lIW, origin='lower', aspect='auto', cmap=cmap,
@@ -96,7 +96,7 @@ def plot_frequency_vs_distance(solver, frequency_range=None,
                              frequency <= frequency_range[1])
         # indices of interest
 
-        frequency= frequency[iis]
+        frequency = frequency[iis]
         IW = IW[:, iis]
 
     ax.imshow(IW, origin='lower', aspect='auto', cmap=cmap,