From 8e1ca4bd60a4a2edbe33e256bfedf5d29a50a78c Mon Sep 17 00:00:00 2001
From: linarphy <linarphy@linarphy.net>
Date: Tue, 23 May 2023 16:59:05 +0200
Subject: [PATCH] Update code to add function (still WIP)

---
 wavelength_calibration.py | 137 ++++++++++++++++----------------------
 1 file changed, 57 insertions(+), 80 deletions(-)

diff --git a/wavelength_calibration.py b/wavelength_calibration.py
index de81e62..d09dcfc 100644
--- a/wavelength_calibration.py
+++ b/wavelength_calibration.py
@@ -4,8 +4,6 @@ import matplotlib.backends.backend_qtagg as qt
 import matplotlib.figure as fig
 import utils
 
-rect = [0.125,0.11,0.775,0.77]
-
 calib_top  = np.load( 'asset/calib_top.npy' )
 calib_down = np.load( 'asset/calib_down.npy' )
 
@@ -15,68 +13,76 @@ peaks_reference         = np.sort( peaks_reference_sorted )
 mean_calib_top  = np.mean( calib_top  , axis = 0 )
 mean_calib_down = np.mean( calib_down , axis = 0 )
 
-peaks_calib_top = np.array(
+peaks_calib_top  = np.array(
     utils.retrieve_peaks( mean_calib_top )
 ).astype( int )
+peaks_calib_down = np.array(
+    utils.retrieve_peaks( mean_calib_down )
+).astype( int )
 
 """
 Signal without peaks
 """
 
-peakless_calib_top = mean_calib_top.copy()
-for peak in peaks_calib_top:
-    first = peak
-    peak , old = first - 2 , first - 1
-    while peakless_calib_top[ peak ] <= peakless_calib_top[ old ]:
-        old   = peak
-        peak -= 1
-    peakless_calib_top[ peak : first ] = peakless_calib_top[ peak ] * np.ones( first - peak )
+def remove_peaks( signal , peaks ):
+    peakless_signal = signal.copy()
+    for peak in peaks:
+        first = peak
+        peak , old = first - 2 , first - 1
+        while peakless_signal[ peak ] <= peakless_signal[ old ]:
+            old   = peak
+            peak -= 1
+        peakless_signal[ peak : first ] = peakless_signal[ peak ] * np.ones( first - peak )
 
-    peak , old = first + 2 , first + 1
-    while peakless_calib_top[ peak ] <= peakless_calib_top[ old ]:
-        old   = peak
-        peak += 1
-    peakless_calib_top[ first : peak ] = peakless_calib_top[ peak ] * np.ones( peak - first )
+        peak , old = first + 2 , first + 1
+        while peakless_signal[ peak ] <= peakless_signal[ old ]:
+            old   = peak
+            peak += 1
+        peakless_signal[ first : peak ] = peakless_signal[ peak ] * np.ones( peak - first )
 
-peakless_calib_top = sig.medfilt( peakless_calib_top , 111 )
+    return sig.medfilt( peakless_signal , 111 )
+peakless_calib_top = remove_peaks( mean_calib_top , peaks_calib_top )
 
-"""
-It's possible to have an idea of the part the spectrum begins with the
-small large peak at the start of it, the peak at 3810 A should be inside
-it. The same goes for the area at the end.
-ONLY TRUE FOR THE CURRENT CALIBRATION (Hg)
-"""
-argmin = np.argmin(
-    peakless_calib_top[ : len( peakless_calib_top ) // 2 ]
-)
-argmax = np.argmax( peakless_calib_top[ : argmin ] )
-peaks_inside = np.where(
-    np.logical_and(
-        argmax < peaks_calib_top,
-        argmin > peaks_calib_top,
+def get_extremities( signal ):
+    """
+    It's possible to have an idea of the part the spectrum begins with the
+    small large peak at the start of it, the peak at 3810 A should be inside
+    it. The same goes for the area at the end.
+    ONLY TRUE FOR THE CURRENT CALIBRATION (Hg)
+    """
+    argmin = np.argmin(
+        signal[ : len( signal ) // 2 ]
     )
-)[0]
-if len( peaks_inside ) < 1:
-    raise Exception( 'unknown plage, cannot autocalibrate' )
-first_peak_cal = peaks_inside[0]
+    argmax = np.argmax( peakless_calib_top[ : argmin ] )
+    peaks_inside = np.where(
+        np.logical_and(
+            argmax < peaks_calib_top,
+            argmin > peaks_calib_top,
+        )
+    )[0]
+    if len( peaks_inside ) < 1:
+        raise Exception( 'unknown plage, cannot autocalibrate' )
+    first_peak = peaks_inside[0]
+
+    """
+    The next peak after the minimum at the end of the spectrum should be
+    5079 A.
+    ONLY TRUE FOR THE CURRENT CALIBRATION (Hg)
+    """
+    argmin_1 = np.argmin(
+        peakless_calib_top[ len( peakless_calib_top ) // 2 : ]
+    ) + len( peakless_calib_top ) // 2
+
+    peaks_inside = np.where(
+        argmin_1 < peaks_calib_top
+    )[0]
+    if len( peaks_inside ) < 1:
+        raise Exception( 'unknown plage, cannot autocalibrate' )
+    return ( first_peak , peaks_inside[0] )
+
 first_peak_ref = 3 # hard-coded
-
-"""
-The next peak after the minimum at the end of the spectrum should be
-5079 A.
-ONLY TRUE FOR THE CURRENT CALIBRATION (Hg)
-"""
-argmin_1 = np.argmin(
-    peakless_calib_top[ len( peakless_calib_top ) // 2 : ]
-) + len( peakless_calib_top ) // 2
-
-peaks_inside = np.where(
-    argmin_1 < peaks_calib_top
-)[0]
-if len( peaks_inside ) < 1:
-    raise Exception( 'unknown plage, cannot autocalibrate' )
-last_peak_cal = peaks_inside[0]
 last_peak_ref = 20 # hard-coded
+first_peak_cal , last_peak_cal = get_extremities( peakless_calib_top )
 
 """
 delete if too much peak in calib
@@ -131,35 +137,6 @@ wavelength = np.polyval(
     )
 )
 
-manager = qt.FigureManager(
-    qt.FigureCanvas(
-        fig.Figure(
-            figsize = ( 10 , 5 )
-        ),
-    ),
-    0,
-)
-manager.canvas.figure.gca().plot(
-    wavelength    ,
-    mean_calib_top,
-)
-manager.canvas.figure.gca().plot(
-    wavelength  ,
-    peakless_calib_top,
-)
-manager.canvas.figure.gca().vlines(
-    [
-        wavelength[ argmin ]  ,
-        wavelength[ argmax ]  ,
-        wavelength[ argmin_1 ],
-    ]                           ,
-    np.min( peakless_calib_top ),
-    np.max( peakless_calib_top ),
-    color = 'red'               ,
-)
-manager.show()
-manager.start_main_loop()
-
 manager = qt.FigureManager(
     qt.FigureCanvas(
         fig.Figure(