plotting improvements

examples/dudley2006.py

@@ -1,3 +1,4 @@
+import colorcet as cc
 import matplotlib.pyplot as plt
 import numpy as np
 from plotapp import PlotApp
@@ -104,16 +105,6 @@ def compute_manual():
         prop.append(spec)
         z.append(stat["z"])
 
-    # with PlotApp(i=range(len(z))) as app:
-    #     o = np.argsort(p.l)
-    #     sax = app["spectrum"]
-    #     tax = app["field"]
-
-    #     @app.update
-    #     def draw(i):
-    #         sax.set_line_data("spectrum", p.l[o], sc.abs2(prop[i][o]))
-    #         tax.set_line_data("field", p.t, sc.abs2(np.fft.ifft(prop[i])))
-
 
 def compute_auto():
     sc.compute(params, True, "examples/dudley2006")
@@ -122,11 +113,10 @@ def compute_auto():
 def plot():
     prop = sc.propagation("examples/dudley_manual.zip")
     newprop = sc.propagation("examples/dudley2006.zip")
-    wl_r = sc.PlotRange(450, 2500, "nm")
     fig, (top, bot) = plt.subplots(2, 2)
     z = np.linspace(0, prop.parameters.length, len(prop))
-    sc.plotting.summary_plot(prop[:], z, axes=top, wl_range=wl_r)
-    sc.plotting.summary_plot(newprop[:], newprop.parameters.z_targets, axes=bot, wl_range=wl_r)
+    sc.plotting.summary_plot(prop[:], z, axes=top)
+    sc.plotting.summary_plot(newprop[:], newprop.parameters.z_targets, axes=bot)
     plt.show()
 
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "scgenerator"
-version = "0.3.13"
+version = "0.3.14"
 description = "Simulate nonlinear pulse propagation in optical fibers"
 readme = "README.md"
 authors = [{ name = "Benoit Sierro", email = "benoit.sierro@iap.unibe.ch" }]

src/scgenerator/physics/fiber.py

@@ -8,7 +8,7 @@ from numpy.polynomial.chebyshev import Chebyshev, cheb2poly
 
 from scgenerator import io
 from scgenerator.io import DataFile
-from scgenerator.math import argclosest, u_nm
+from scgenerator.math import all_zeros, argclosest, u_nm
 from scgenerator.physics import materials as mat
 from scgenerator.physics import units
 from scgenerator.physics.units import c, pi
@@ -56,6 +56,10 @@ def D_to_beta2(D, l):
     return -(l**2) / (pipi * c) * D
 
 
+def find_zdw(l: np.ndarray, beta2_arr: np.ndarray) -> list[float]:
+    return [el for el in all_zeros(l, beta2_arr) if el > 0]
+
+
 def plasma_dispersion(l, number_density, simple=False):
     """
     computes dispersion (beta2) for constant plasma

src/scgenerator/plotting.py

@@ -2,6 +2,7 @@ import os
 from pathlib import Path
 from typing import Any, Callable, Literal, Optional, Sequence, Union
 
+import matplotlib.colors
 import matplotlib.gridspec as gs
 import matplotlib.pyplot as plt
 import numpy as np
@@ -1004,11 +1005,13 @@ def apply_log(values: np.ndarray, log: Union[str, bool, float, int]) -> np.ndarr
         elif log == "2D":
             values = math.to_dB(values, ref=values.max())
         elif log == "1D" or log is True:
-            values = math.to_dB(values)
+            values = math.to_dB(values, axis=1)
         elif log == "smooth 1D":
             ref = np.max(values, axis=1)
             ind = np.argmax((ref[:-1] - ref[1:]) < 0)
             values = math.to_dB(values, ref=np.max(ref[ind:]))
+        elif log == "linear 1D":
+            values = (values.T / values.max(axis=1)).T
         else:
             raise ValueError(f"Log argument {log} not recognized")
     return values
@@ -1158,9 +1161,11 @@ def summary_plot(
     wl_range: PlotRange | None = None,
     time_range: PlotRange | None = None,
     db_min: float = -50.0,
+    lin_min: float = 1e-3,
     axes: tuple[Axes, Axes] | None = None,
     wl_db="1D",
     time_db=False,
+    cmap: str | matplotlib.colors.LinearSegmentedColormap = "viridis",
 ):
     wl_int = specs.wl_int
     time_int = specs.time_int
@@ -1171,12 +1176,18 @@ def summary_plot(
         z = z[: specs.shape[0]]
     z = np.asarray(z)
 
+    calc_limit, wl_disp_limit = (
+        (10 ** (0.1 * db_min - 1), db_min) if wl_db and wl_db != "linear 1D" else (lin_min, 0)
+    )
     if wl_range is None:
-        imin, imax = math.span_above(wl_int, wl_int.max() * 1e-6)
+        imin, imax = math.span_above(wl_int, wl_int.max() * calc_limit)
         wl_range = PlotRange(specs.wl_disp[imin] * 1e9, specs.wl_disp[imax] * 1e9, "nm")
 
+    calc_limit, time_disp_limit = (
+        (10 ** (0.1 * db_min - 1), db_min) if time_db and time_db != "linear 1D" else (lin_min, 0)
+    )
     if time_range is None:
-        imin, imax = math.span_above(time_int, time_int.max() * 1e-6)
+        imin, imax = math.span_above(time_int, time_int.max() * calc_limit)
         time_range = PlotRange(specs.t[imin] * 1e15, specs.t[imax] * 1e15, "fs")
 
     if axes is None:
@@ -1185,7 +1196,21 @@ def summary_plot(
         left, right = axes
 
     x, y, values = transform_2D_propagation(wl_int, wl_range, specs.wl_disp, z, log=wl_db)
-    left.imshow(values, extent=get_extent(x, y), origin="lower", aspect="auto", vmin=db_min)
+    left.imshow(
+        values,
+        extent=get_extent(x, y),
+        origin="lower",
+        aspect="auto",
+        vmin=wl_disp_limit,
+        cmap=cmap,
+    )
 
     x, y, values = transform_2D_propagation(time_int, time_range, specs.t, z, log=time_db)
-    right.imshow(values, extent=get_extent(x, y), origin="lower", aspect="auto", vmin=db_min)
+    right.imshow(
+        values,
+        extent=get_extent(x, y),
+        origin="lower",
+        aspect="auto",
+        vmin=time_disp_limit,
+        cmap=cmap,
+    )