From b6c76b838d3a58680348951a656d465d4d27517a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Sierro?= <benoitsierro@bluewin.ch>
Date: Mon, 25 Oct 2021 13:19:32 +0200
Subject: [PATCH] misc

---
 .gitignore                       |  1 +
 src/scgenerator/__init__.py      |  2 +-
 src/scgenerator/math.py          |  5 +++++
 src/scgenerator/physics/pulse.py |  4 ++++
 src/scgenerator/plotting.py      | 13 ++++++++++---
 5 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index 3fc7b58..1e04fc0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
 **/*.npy
 
 plots*
+/make_*.py
 Archive
 *.mp4
 *.png
diff --git a/src/scgenerator/__init__.py b/src/scgenerator/__init__.py
index c5f3af3..ce3619a 100644
--- a/src/scgenerator/__init__.py
+++ b/src/scgenerator/__init__.py
@@ -1,5 +1,5 @@
 from . import math
-from .math import abs2, argclosest, span
+from .math import abs2, argclosest, span, normalized
 from .physics import fiber, materials, pulse, simulate, units
 from .physics.simulate import RK4IP, parallel_RK4IP, run_simulation
 from .plotting import (
diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py
index 4b6f069..9963f5c 100644
--- a/src/scgenerator/math.py
+++ b/src/scgenerator/math.py
@@ -76,6 +76,11 @@ def abs2(z: np.ndarray) -> np.ndarray:
     return z.real ** 2 + z.imag ** 2
 
 
+def normalized(z: np.ndarray) -> np.ndarray:
+    ab = abs2(z)
+    return ab / ab.max()
+
+
 def sigmoid(x):
     return 1 / (np.exp(-x) + 1)
 
diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py
index 53cd159..0139968 100644
--- a/src/scgenerator/physics/pulse.py
+++ b/src/scgenerator/physics/pulse.py
@@ -1146,3 +1146,7 @@ def remove_2nd_order_dispersion2(
     opti = minimize_scalar(score, bounds=(-max_gdd * 1e30, max_gdd * 1e30))
     opti["x"] *= 1e-30
     return propagate(opti.x * 1e30), opti
+
+
+def gdd(w: np.ndarray, gdd: float) -> np.ndarray:
+    return np.exp(0.5j * w ** 2 * gdd)
diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py
index 30afa42..4dff2c1 100644
--- a/src/scgenerator/plotting.py
+++ b/src/scgenerator/plotting.py
@@ -1042,12 +1042,20 @@ def measure_and_annotate_fwhm(
         field = abs2(field)
     _, (left, right), *_ = pulse.find_lobe_limits(unit.inv(t), field)
     arrow_label = f"{right - left:.1f} {unit.name}"
+
+    annotate_fwhm(ax, left, right, arrow_label, field.max(), side, arrow_length_pts, arrow_props)
+    return right - left
+
+
+def annotate_fwhm(
+    ax, left, right, arrow_label, v_max=1, side="right", arrow_length_pts=20.0, arrow_props=None
+):
     arrow_dict = dict(arrowstyle="->")
     if arrow_props is not None:
         arrow_dict |= arrow_props
     ax.annotate(
         "" if side == "right" else arrow_label,
-        (left, field.max() / 2),
+        (left, v_max / 2),
         xytext=(-arrow_length_pts, 0),
         ha="right",
         va="center",
@@ -1056,13 +1064,12 @@ def measure_and_annotate_fwhm(
     )
     ax.annotate(
         "" if side == "left" else arrow_label,
-        (right, field.max() / 2),
+        (right, v_max / 2),
         xytext=(arrow_length_pts, 0),
         textcoords="offset points",
         arrowprops=arrow_dict,
         va="center",
     )
-    return right - left
 
 
 def partial_plot(root: os.PathLike):