diff --git a/src/scgenerator/physics/pulse.py b/src/scgenerator/physics/pulse.py index b33114d..8328606 100644 --- a/src/scgenerator/physics/pulse.py +++ b/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, - old_w: np.ndarray, - w_ind: np.ndarray, - new_w: np.ndarray, + signal: 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 - 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 + delays : np.ndarray, shape (nd,) + delays overs which to compute the spectrogram + signal : np.ndarray, shape (nt,) + signal array that matches the time array + old_w : np.ndarray, shape (nt,) + 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)) - spec[i] = math.linear_interp_1d(old_w, math.abs2(fft(masked)[w_ind]), new_w) + masked = signal * np.exp(-(((time - delay) / t_gate) ** 2)) + spec[i] = math.abs2(fft(masked)) - return spec, delays + return spec def g12(values): diff --git a/src/scgenerator/plotting.py b/src/scgenerator/plotting.py index 4b2dc1f..811a2cd 100644 --- a/src/scgenerator/plotting.py +++ b/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) - spec, new_t = pulse.spectrogram_interp( - t_axis, delays, values, old_w, w_ind, new_w, gate_width=gate_width - ) + new_t = np.linspace(t_range.unit(t_range.left), t_range.unit(t_range.right), t_res) + spec, new_t = pulse.spectrogram(t_axis, new_t, values, 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)