fix spectrogram function
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user