fix spectrogram function

src/scgenerator/physics/pulse.py

@@ -607,56 +607,9 @@ def ideal_compressed_pulse(spectra):
 
 
 def spectrogram(
-    time, values, t_res=256, t_win=24e-12, gate_width=200e-15, shift=False, w_ind: np.ndarray = None
-):
-    """
-    returns the spectorgram of the field given in values
-
-    Parameters
-    ----------
-        time : 1D array-like
-            time in the co-moving frame of reference
-        values : 1D array-like
-            field array that matches the time array
-        t_res : int, optional
-            how many "bins" the time array is subdivided into. Default : 256
-        t_win : float, optional
-            total time window (=length of time) over which the spectrogram is computed. Default : 24e-12
-        gate_width : float, optional
-            width of the gaussian gate function (=sqrt(2 log(2)) * FWHM). Default : 200e-15
-
-    Returns
-    ----------
-        spec : 2D array
-            real 2D spectrogram
-        delays : 1D array of size t_res
-            new time axis
-    """
-    if isinstance(t_win, tuple):
-        left, right = t_win
-    else:
-        t_lim = t_win / 2
-        left, right = -t_lim, t_lim
-
-    delays = np.linspace(left, right, t_res)
-
-    spec = np.zeros((t_res, len(time) if w_ind is None else len(w_ind)))
-    for i, delay in enumerate(delays):
-        masked = values * np.exp(-(((time - delay) / gate_width) ** 2))
-        spec[i] = math.abs2(fft(masked))
-
-    if shift:
-        spec = fftshift(spec, axes=1)
-    return spec, delays
-
-
-def spectrogram_interp(
     time: np.ndarray,
     delays: np.ndarray,
-    values: np.ndarray,
+    signal: np.ndarray,
-    old_w: np.ndarray,
-    w_ind: np.ndarray,
-    new_w: np.ndarray,
     gate_width=200e-15,
 ):
     """
@@ -664,31 +617,32 @@ def spectrogram_interp(
 
     Parameters
     ----------
-        time : 1D array-like
+        time : np.ndarray, shape (nt,)
             time in the co-moving frame of reference
-        values : 1D array-like
+        delays : np.ndarray, shape (nd,)
-            field array that matches the time array
+            delays overs which to compute the spectrogram
-        t_res : int, optional
+        signal : np.ndarray, shape (nt,)
-            how many "bins" the time array is subdivided into. Default : 256
+            signal array that matches the time array
-        t_win : float, optional
+        old_w : np.ndarray, shape (nt,)
-            total time window (=length of time) over which the spectrogram is computed. Default : 24e-12
+            angular frequency array corresponding to the fourier space of the `time` array
+        new_w : np.ndarray, shape (nw,)
+            angural frequency grid of the spectrogram
         gate_width : float, optional
-            width of the gaussian gate function (=sqrt(2 log(2)) * FWHM). Default : 200e-15
+            full width at half maximum of the gaussian gate pulse
 
     Returns
     ----------
-        spec : 2D array
+        spec : np.ndarray, shape (nd, nt)
             real 2D spectrogram
-        delays : 1D array of size t_res
-            new time axis
     """
 
-    spec = np.zeros((len(delays), len(new_w)))
+    spec = np.zeros((len(delays), len(time)))
+    t_gate = gate_width / (2 * np.sqrt(np.log(2)))
     for i, delay in enumerate(delays):
-        masked = values * np.exp(-(((time - delay) / gate_width) ** 2))
+        masked = signal * np.exp(-(((time - delay) / t_gate) ** 2))
-        spec[i] = math.linear_interp_1d(old_w, math.abs2(fft(masked)[w_ind]), new_w)
+        spec[i] = math.abs2(fft(masked))
 
-    return spec, delays
+    return spec
 
 
 def g12(values):

src/scgenerator/plotting.py

@@ -905,10 +905,9 @@ def plot_spectrogram(
     new_w = np.linspace(w_range.left, w_range.right, w_res)
 
     # Actually compute the spectrogram
-    delays = np.linspace(t_range.unit(t_range.left), t_range.unit(t_range.right), t_res)
+    new_t = np.linspace(t_range.unit(t_range.left), t_range.unit(t_range.right), t_res)
-    spec, new_t = pulse.spectrogram_interp(
+    spec, new_t = pulse.spectrogram(t_axis, new_t, values, gate_width=gate_width)
-        t_axis, delays, values, old_w, w_ind, new_w, gate_width=gate_width
+    spec = interp1d(old_w, spec[:, w_ind], fill_value=0, bounds_error=False, axis=1)(new_w)
-    )
 
     new_t = t_range.unit.inv(new_t)