fix spectrogram function

This commit is contained in:
Benoît Sierro
2023-07-27 11:48:31 +02:00
parent b95446df0d
commit c09dbf4841
2 changed files with 20 additions and 67 deletions

View File

@@ -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):

View File

@@ -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)