attoworld.personal.marco

153def profile_analysis(
154    profile_file,
155    trace_file=None,
156    magnification=1.0,
157    pixelsize=5.0,
158    power=20.0,
159    reprate=4.0,
160    ROI_diam: int = 100,
161    forced_background=None,
162    trace_cutoff_radius=50.0,
163    cutoff_gaus_fit_profile=25.0,
164):
165    """if trace_file is None, this function only loads, fits and plots the beam profile; if trace_file is given, the function additionally calculates the peak intensity and field.
166
167    Args:
168        profile_file: name of the file containing the beam profile; the file (txt or csv) contains the 2D array of the beam profile
169        trace_file: None; name of the file containing the trace; the file (txt or csv) contains the time and field arrays as tab-separated columns
170        magnification: 1.; ratio of lens-camera/lens-z0 lengths; This is used to give the correct size of the beam profile at z0 in um
171        pixelsize: 5. um; pixel size of the camera
172        power: 20. mW; power of the laser beam; this is used to calculate the peak intensity
173        reprate: 4 kHz; repetition rate of the laser; this is used to calculate the peak intensity
174        ROI_diam: 100; diameter of the region of interest in pixels; the center of the ROI is the maximum of the camera data (which is assumed to be not saturated)
175        cutoff_gaus_fit_profile: 25. pixels; the 2d gaussian is fitted to the beam profile up to this distance from the peak
176        forced_background: None; if not None, the subtracted background is set to this value; otherwise it's computed from the outer ring of the ROI (camera data)
177        trace_cutoff_radius: 50. fs; for intensity calculation, the integral of the trace is computed only in the range t[peak]-trace_cutoff_radius < t < t[peak]+trace_cutoff_radius"""
178
179    pixelsize = pixelsize / magnification
180    data = pandas.read_csv(profile_file)
181
182    dataval = np.array(data.values)
183    # print('eliminating fixed broken pixels rayCi camera')
184    # dataval=eliminate_broken_pixels(dataval)
185    maxindex = np.unravel_index(np.argmax(dataval), dataval.shape)
186    print("diameter of Region Of Interest (ROI): ", ROI_diam, " pixels")
187    dataval = np.array(
188        data.values[
189            maxindex[0] - int(ROI_diam / 2) : maxindex[0] + int(ROI_diam / 2),
190            maxindex[1] - int(ROI_diam / 2) : maxindex[1] + int(ROI_diam / 2),
191        ]
192    )
193    maxindex = np.unravel_index(np.argmax(dataval), dataval.shape)
194
195    background = 0.0
196    count = 0
197    print(
198        "calculating background intensity outside a circle whose diameter is 3/4 of the ROI diameter"
199    )
200    for i, j in np.ndindex(dataval.shape):
201        if (i - maxindex[0]) ** 2 + (j - maxindex[1]) ** 2 > (ROI_diam * 3 / 8) ** 2:
202            background += dataval[i, j]
203            count += 1
204    background /= count
205
206    if forced_background is not None:
207        background = forced_background
208        print("subtracting background intensity: ", background)
209    else:
210        print("subtracting background intensity: ", background)
211
212    dataval = np.maximum(dataval - background, dataval * 0.0)
213    maxval = np.max(dataval)
214
215    # gaussian fit
216    param = moments_peak(dataval)
217    param = fitgaussian(cut_tail(dataval, cutoff_gaus_fit_profile))
218    X, Y = np.indices(dataval.shape)
219    fitted_funct = gaussian(*param)(X, Y)
220    print(
221        "beam w (um), resp. in x and y axis:",
222        2 * param[3] * pixelsize,
223        2 * param[4] * pixelsize,
224    )
225    print(
226        "beam 1/e^2 diameter [um]: ", 4 * param[3] * pixelsize, 4 * param[4] * pixelsize
227    )
228
229    if trace_file is not None:
230        trace = pandas.read_csv(trace_file, sep="\t")
231
232        # time in fs, field in a.u.
233        intensity = (
234            power
235            * 1e-3
236            / reprate
237            * 1e-3
238            / integral2d(fitted_funct, pixelsize=pixelsize)
239            * np.max(fitted_funct)
240            * 1e8
241            / trace_integral(
242                trace["delay (fs)"],
243                cut_trace(
244                    trace["delay (fs)"], trace["field (a.u.)"] ** 2, trace_cutoff_radius
245                ),
246            )
247            * np.max(trace["field (a.u.)"] ** 2)
248            * 1e15
249        )
250
251        intensityFromRawCameraData = (
252            power
253            * 1e-3
254            / reprate
255            * 1e-3
256            / integral2d(
257                dataval[
258                    maxindex[0] - int(ROI_diam / 2) : maxindex[0] + int(ROI_diam / 2),
259                    maxindex[1] - int(ROI_diam / 2) : maxindex[1] + int(ROI_diam / 2),
260                ],
261                pixelsize=pixelsize,
262            )
263            * np.max(dataval)
264            * 1e8
265            / trace_integral(
266                trace["delay (fs)"],
267                cut_trace(
268                    trace["delay (fs)"], trace["field (a.u.)"] ** 2, trace_cutoff_radius
269                ),
270            )
271            * np.max(trace["field (a.u.)"] ** 2)
272            * 1e15
273        )
274
275        temporal_fwhm = get_fwhm(
276            np.array(trace["delay (fs)"]), np.array(trace["field (a.u.)"])
277        )
278        intensityFromTemporalFWHM = (
279            power
280            * 1e-3
281            / reprate
282            * 1e-3
283            / integral2d(
284                dataval[
285                    maxindex[0] - int(ROI_diam / 2) : maxindex[0] + int(ROI_diam / 2),
286                    maxindex[1] - int(ROI_diam / 2) : maxindex[1] + int(ROI_diam / 2),
287                ],
288                pixelsize=pixelsize,
289            )
290            * np.max(dataval)
291            * 1e8
292            / (np.sqrt(2 * np.pi) * temporal_fwhm / np.sqrt(8 * np.log(2)))
293            * 1e15
294        )
295
296        print(
297            "\nusing the gaussian fit of the beam profile and the integral of the trace^2:\n "
298            "peak intensity (sub-cycle, W/cm2): {i:.3e}".format(**{"i": intensity})
299        )
300        peak_field = np.sqrt(
301            intensity * 1e4 / constants.speed_of_light / constants.epsilon_0
302        )
303        print("peak field (sub-cycle, V/m): {f:.3e}".format(**{"f": peak_field}))
304
305        print(
306            "\nusing directly camera data, and the integral of the trace^2:\n "
307            "peak intensity (sub-cycle, W/cm2): {i:.3e}".format(
308                **{"i": intensityFromRawCameraData}
309            )
310        )
311        peak_field = np.sqrt(
312            intensityFromRawCameraData
313            * 1e4
314            / constants.speed_of_light
315            / constants.epsilon_0
316        )
317        print("peak field (sub-cycle, V/m): {f:.3e}".format(**{"f": peak_field}))
318
319        print(
320            "\nusing directly camera data, and the temporal FWHM retrieved from the trace:\n "
321            "peak intensity (sub-cycle, W/cm2): {i:.3e}".format(
322                **{"i": intensityFromTemporalFWHM}
323            )
324        )
325        peak_field = np.sqrt(
326            intensityFromTemporalFWHM
327            * 1e4
328            / constants.speed_of_light
329            / constants.epsilon_0
330        )
331        print("peak field (sub-cycle, V/m): {f:.3e}".format(**{"f": peak_field}))
332
333        fig, ax = plt.subplots(1, 1)
334        time = trace["delay (fs)"]
335        field = trace["field (a.u.)"]
336        i_peak = np.argmax(field**2)
337        ax.plot(
338            time[np.abs(time - time[i_peak]) < trace_cutoff_radius],
339            field[np.abs(time - time[i_peak]) < trace_cutoff_radius] ** 2,
340        )
341        plt.show()
342
343    fig, ax = plt.subplots(1, 1)
344    ax.imshow(dataval)
345    plt.show()
346
347    plot_crosssect(dataval, fitfunct=gaussian(*param))

 129class TraceHandler:
 130    """Loads and stores the trace (field vs time) saved to file (two columns, tab separated) or given in the form of time and field [or wavelengths, spectrum and phase] arrays.
 131
 132    ALL THE TIMES ARE IN fs, FREQUENCIES IN PHz, WAVELENGTHS IN nm.
 133
 134    Most of the class methods modify the object data in place, and don't return anything.
 135
 136    Attributes:
 137        fieldTimeV: (WAVEFORM DATA in the time domain) time array (fs)
 138        fieldV: (WAVEFORM DATA in the time domain) field array
 139        fieldStdevV: (WAVEFORM DATA in the time domain) field standard deviation (or sigma) array
 140
 141        frequencyAxis: (WAVEFORM DATA in the frequency domain) frequency axis of the FFT in PHz (should be np.fft compatible, i.e., equally spaced, in the form [0, df, 2*df, ..., fmax, -fmax[-df], ..., -df])
 142        fftFieldV: (WAVEFORM DATA in the frequency domain) FFT field (complex)
 143        complexFieldTimeV: (WAVEFORM DATA in the frequency domain) time axis of the complex field (fs)
 144        complexFieldV: (WAVEFORM DATA in the frequency domain) complex field ( = IFFT{FFT(ω)θ(ω)}(t), where θ(ω) = 1 if ω >= 0, θ(ω) = 0 otherwise )
 145        wvlAxis: (WAVEFORM DATA in the frequency domain) wavelength axis for the FFT spectrum (nm)
 146        fftSpectrum: (WAVEFORM DATA in the frequency domain) wavelength-dependent positive FFT spectrum ( = |FFT{field}|^2 * df/dλ, λ > 0 )
 147        fftphase: (WAVEFORM DATA in the frequency domain) wavelength-dependent spectral phase
 148
 149        wvlSpectrometer: (SPECTROMETER DATA) wavelength (nm)
 150        ISpectrometer: (SPECTROMETER DATA) spectral intensity
 151
 152        fsZeroPadding: (= 150 by default) time span in fs for zero padding (appending zeroes both before and after the trace, in order to get a smoother fft)
 153        filename: filename to load the trace from
 154        filename_spectrum: filename to load the spectrum from
 155        normalization_trace: normalization factor for the trace
 156        zero_delay: time zero, corresponds to the maximum envelope
 157    """
 158
 159    def __init__(
 160        self,
 161        filename=None,
 162        filename_spectrum=None,
 163        time=None,
 164        field=None,
 165        stdev=None,
 166        wvl=None,
 167        spectrum=None,
 168        wvl_FFT_trace=None,
 169        spectrum_FFT_trace=None,
 170        phase_FFT_trace=None,
 171    ):
 172        """Constructor; loads the trace from file (or from time-field or wavelegth-spectrum-phase arrays) and stores it in the class.
 173
 174        One of the following must be provided (as a default all args are None):
 175        * filename: the name of the file containing the trace (two columns, tab separated)
 176        * time and field: the time vector (fs) and electric field vector (a.u.)
 177        * wvl_FFT_trace, spectrum_FFT_trace and phase_FFT_trace: the wavelength (nm) spectrum (a.u.) and phase (rad) vectors
 178
 179        Args:
 180            filename (str): the name of the file containing the trace (two columns - time and field, tab separated)
 181            filename_spectrum (str): the name of the file containing the spectrum (two columns - wavelength and spectral intensity, tab separated)
 182                this does not correspond to the fourier transform of the trace, but to some spectrometer data that you would like to compare to the loaded trace
 183            time (array): the time vector (fs)
 184            field (array): the electric field vector (a.u.)
 185            stdev (array): the standard deviation (or sigma) of the electric field vector (a.u.)
 186            wvl (array): the wavelength vector (nm) (spectrometer)
 187            spectrum (array): the spectral intensity vector (a.u.) (spectrometer)
 188            wvl_FFT_trace (array): the wavelength vector (nm) of the FFT trace (to reconstruct a trace from spectral data)
 189            spectrum_FFT_trace (array): the spectral intensity vector (a.u.) of the FFT trace (to reconstruct a trace from spectral data)
 190            phase_FFT_trace (array): the spectral phase vector (rad) of the FFT trace (to reconstruct a trace from spectral data)
 191        """
 192
 193        self.fsZeroPadding = 150
 194        self.filename = filename
 195        self.filename_spectrum = filename_spectrum
 196
 197        self.fieldTimeV = None
 198        self.fieldV = None
 199        self.fieldStdevV = None
 200
 201        self.frequencyAxis = None
 202        self.fftFieldV = None
 203        self.complexFieldTimeV = None
 204        self.complexFieldV = None
 205        self.wvlAxis = None
 206        self.fftSpectrum = None
 207        self.fftphase = None
 208
 209        self.wvlSpectrometer = None
 210        self.ISpectrometer = None
 211
 212        self.normalization_trace = None
 213        self.zero_delay = None
 214
 215        if filename is not None:
 216            self.load_trace()
 217        elif time is not None and field is not None:
 218            self.load_trace_from_arrays(time, field, stdev)
 219        elif (
 220            wvl_FFT_trace is not None
 221            and spectrum_FFT_trace is not None
 222            and phase_FFT_trace is not None
 223        ):
 224            self.load_trace_from_spectral_data(
 225                wvl_FFT_trace, spectrum_FFT_trace, phase_FFT_trace
 226            )
 227        else:
 228            print(
 229                "\n\nWARNING: in function TraceHandler.__init__() no data was loaded\n\n"
 230            )
 231
 232        if self.filename_spectrum is not None:
 233            self.load_spectrum()
 234        elif wvl is not None and spectrum is not None:
 235            self.load_spectrum_from_arrays(wvl, spectrum)
 236
 237    def load_spectrum_from_arrays(self, wvl, spectrum):
 238        """loads the spectrum from wavelength and spectrum arrays and stores it in the class.
 239
 240        The spectrum does not correspond to the fourier transform of the trace, but to some spectrometer data that you would like to compare with the loaded trace
 241
 242        Args:
 243            wvl: wavelength array (nm)
 244            spectrum: spectral intensity array
 245        """
 246        """wavelength is in nm"""
 247        check_equal_length(wvl, spectrum)
 248        self.wvlSpectrometer = np.array(wvl)
 249        self.ISpectrometer = np.array(spectrum)
 250
 251    def load_trace_from_arrays(self, fieldTimeV, fieldV, fieldStdevV=None):
 252        """loads the trace from time and field arrays and stores it in the class.
 253
 254        Args:
 255            fieldTimeV: time array (fs)
 256            fieldV: field array
 257            fieldStdevV: standard deviation - or sigma - to be associated to the trace (optional, default = None)
 258        """
 259        if fieldTimeV is None or fieldV is None:
 260            raise ValueError(
 261                "\nin function TraceHandler.load_from_arrays() too many arguments were None\n"
 262                "it might be that the constructor was erroneously used with no or too few arguments\n"
 263            )
 264        if np.max(np.abs(fieldTimeV)) < 1.0:
 265            print(
 266                "\n\nWARNING: do you remember that the time axis of TraceHandler is in fs?\n\n"
 267            )
 268        self.fieldTimeV = np.array(fieldTimeV)
 269        self.fieldV = np.array(fieldV)
 270        if fieldStdevV is not None:
 271            self.fieldStdevV = np.array(fieldStdevV)
 272        self.normalization_trace = np.max(np.abs(self.fieldV))
 273        self.update_fft()
 274
 275    def load_trace(self, fname=None):
 276        """loads from file"""
 277        if fname is not None:
 278            self.filename = fname
 279        data = pandas.read_csv(self.filename, sep="\t")
 280        self.fieldTimeV = data["delay (fs)"].to_numpy()
 281        self.fieldV = data["field (a.u.)"].to_numpy()
 282        self.fieldStdevV = data["stdev field"].to_numpy()
 283        self.normalization_trace = np.max(np.abs(self.fieldV))
 284        self.update_fft()
 285
 286    def load_trace_from_spectral_data(
 287        self, wvl_FFT_trace, spectrum_FFT_trace, phase_FFT_trace
 288    ):
 289        """loads the trace from wavelength, spectrum, and spectral phase arrays and stores it in the class.
 290
 291        This function has a different job than load_spectrum_from_arrays(): in fact, it retrieves the trace from the given spectral data via fourier transform.
 292
 293        Args:
 294            wvl_FFT_trace: wavelength array in nm
 295            spectrum_FFT_trace: (wavelength-)spectral intensity array. This is assumed to be |FFT{field}|^2 * df/dλ
 296            phase_FFT_trace: spectral phase in rad
 297        """
 298        check_equal_length(wvl_FFT_trace, spectrum_FFT_trace, phase_FFT_trace)
 299
 300        # check that the wvl_FFT_trace array is monotonically increasing
 301        if np.any(np.diff(wvl_FFT_trace) < 0):
 302            if np.all(np.diff(wvl_FFT_trace) < 0):
 303                wvl_FFT_trace = wvl_FFT_trace[::-1]
 304                spectrum_FFT_trace = spectrum_FFT_trace[::-1]
 305                phase_FFT_trace = phase_FFT_trace[::-1]
 306            else:
 307                raise ValueError(
 308                    "in function TraceHandler.load_from_spectral_data() wavelength array is not monotonous"
 309                )
 310
 311        # frequency spectrum (monotonically increasing)
 312        freq = constants.speed_of_light / wvl_FFT_trace[::-1] * 1e-6
 313        spectrum_freq = spectrum_FFT_trace[::-1] / (
 314            constants.speed_of_light / wvl_FFT_trace[::-1] ** 2 * 1e-6
 315        )
 316        phase_FFT_trace = phase_FFT_trace[::-1]
 317
 318        # create proper freq axis for the fourier transform
 319        df = (freq[1] - freq[0]) / 4
 320        nf = int(np.ceil(3 * freq[-1] / df))
 321        self.frequencyAxis = np.concatenate((np.arange(0, nf), np.arange(-nf, 0))) * df
 322        # corresponding time window
 323        twin = 1.0 / (df)
 324
 325        # complex fourier transform
 326        spectrum_freq = np.sqrt(spectrum_freq) * np.exp(1.0j * phase_FFT_trace)
 327
 328        # fill with zeros
 329        initial_zero_freq = np.linspace(
 330            freq[1] - freq[0], freq[0], int(np.ceil(4 * freq[0] / (freq[1] - freq[0])))
 331        )
 332        initial_zeros = np.zeros(len(initial_zero_freq))
 333        final_zero_freq = np.linspace(
 334            2 * freq[-1] - freq[-2],
 335            4 * freq[-1],
 336            int(np.ceil(15 * freq[-1] / (freq[-1] - freq[-2]))),
 337        )
 338        final_zeros = np.zeros(len(final_zero_freq))
 339        freq = np.concatenate((initial_zero_freq, freq, final_zero_freq))
 340        spectrum_freq = np.concatenate((initial_zeros, spectrum_freq, final_zeros))
 341        # add negative frequencies
 342        freq = np.concatenate((-freq[::-1], freq))
 343        spectrum_freq = np.concatenate(
 344            (np.conjugate(spectrum_freq[::-1]), spectrum_freq)
 345        )
 346
 347        # interpolate the spectrum to the frequency axis of the fft
 348        self.fftFieldV = np.interp(self.frequencyAxis, freq, spectrum_freq)
 349
 350        # add linear phase to center the pulse in the time window; this has to be done after interpolation because the spectral phase will be fast oscillating
 351        self.fftFieldV = self.fftFieldV * np.exp(
 352            1.0j * 2 * np.pi * self.frequencyAxis * twin / 2
 353        )
 354
 355        self.update_fft_spectrum()
 356        self.update_trace_from_fft()
 357
 358    def load_spectrum(self, fname=None):
 359        """loads spectrum from file.
 360
 361        The spectrum does not correspond to the fourier transform of the trace, but to the spectrometer data that you would like to compare with the loaded trace
 362
 363        Args:
 364            fname: filename
 365        """
 366        if fname is not None:
 367            self.filename_spectrum = fname
 368        data = pandas.read_csv(self.filename_spectrum, sep="\t")
 369        self.wvlSpectrometer = data["wavelength (nm)"].to_numpy()
 370        self.ISpectrometer = data["intensity (a.u.)"].to_numpy()
 371        self.normalize_spectrum()
 372
 373    def update_fft(self, zero_pad_field=True):
 374        """updates the fft of the trace from the time domain data.
 375
 376        Args:
 377            zero_pad_field: bool; if true (default) add zeros before and after the trace before computing fft (for a smoother spectrum)
 378                The length of appended zeros is defined by self.fsZeroPadding.
 379        """
 380        zp_fieldTimeV, zp_fieldV = deepcopy(self.fieldTimeV), deepcopy(self.fieldV)
 381        if zero_pad_field:
 382            zp_fieldTimeV, zp_fieldV = zero_padding(
 383                self.fieldTimeV, self.fieldV, fsExtension=self.fsZeroPadding
 384            )
 385        self.frequencyAxis, self.fftFieldV = fourier_transform(zp_fieldTimeV, zp_fieldV)
 386        self.update_fft_spectrum()
 387
 388    def update_fft_spectrum(self):
 389        """Updates the wavelength-dependent intensity spectrum. Results are stored in self.wvlAxis and self.fftSpectrum.
 390
 391        This is (derived from but) different from the fourier transform because:
 392         * it considers the squared modulus of the field
 393         * it multiplies the frequency-dependent spectral density by df/dλ = c/λ²
 394         * it considers only positive wavelengths
 395
 396        The function internally calls update_spectral_phase() as well.
 397        """
 398        self.wvlAxis = (
 399            constants.speed_of_light
 400            / self.frequencyAxis[self.frequencyAxis > 0]
 401            * 1.0e-6
 402        )
 403        self.fftSpectrum = (
 404            np.abs(self.fftFieldV[self.frequencyAxis > 0]) ** 2
 405            * constants.speed_of_light
 406            / self.wvlAxis**2
 407        )
 408        self.compute_complex_field()
 409        self.update_spectral_phase()
 410        self.normalize_fft_spectrum()
 411
 412    def update_spectral_phase(self):
 413        """computes the spectral phase of the trace. The result is stored in self.fftphase.
 414
 415        An attempt of automatic linear phase subtraction is made. Phase is unwrapped.
 416        """
 417        # determine the time distance (delay) of the waveform peak from the beginning of the trace (more precisely, from the beginning of the ifft)
 418        tmax = self.get_zero_delay()
 419        deltaT1 = tmax - self.complexFieldTimeV[0]
 420        deltaT2 = self.complexFieldTimeV[-1] - tmax
 421        ifft_twindow = 1 / (self.frequencyAxis[1] - self.frequencyAxis[0])
 422        total_delay = (
 423            ifft_twindow / 2 + (deltaT2 - deltaT1) / 2
 424        )  # total estimated delay of the peak from the beginning of the ifft trace
 425
 426        # compute the phase, use the computed 'delay' to correct for the non-interesting linear component, and unwrap it
 427        self.fftphase = np.angle(self.fftFieldV[self.frequencyAxis > 0])
 428        self.fftphase = (
 429            self.fftphase
 430            - total_delay * self.frequencyAxis[self.frequencyAxis > 0] * 2 * np.pi
 431        )
 432        self.fftphase = np.unwrap(self.fftphase)
 433
 434        # subtract constant phase offset (m * 2*π)
 435        n_pi = np.sum(self.fftSpectrum * self.fftphase) / np.sum(self.fftSpectrum)
 436        n_pi = int(n_pi / (2 * np.pi))
 437        self.fftphase -= n_pi * 2 * np.pi
 438
 439    def update_trace_from_fft(self):  # remember to use together with strip_from_trace
 440        # careful: the complex part of the IFFT is discarded
 441        """Updates the field trace from fft
 442
 443        In most cases, you might want to call strip_from_trace() afterward, to eliminate the zero-padding from the trace.
 444        """
 445        self.fieldTimeV, self.fieldV = inverse_fourier_transform(
 446            self.frequencyAxis, self.fftFieldV
 447        )
 448        self.fieldV = np.real(self.fieldV)
 449        self.normalization_trace = np.max(np.abs(self.fieldV))
 450        self.get_zero_delay()
 451
 452    def strip_from_trace(self, timeRange=None):
 453        """eliminates the zeros appended to the trace (zero-padding) when computing the FFT.
 454
 455        Useful when retrieving the trace via ifft (method update_trace_from_fft()).
 456
 457        Args:
 458            timeRange (float): the time range (fs) to strip from the beginning and end of the trace. If None (default), fsZeroPadding is used.
 459        """
 460        if timeRange is None:
 461            timeRange = self.fsZeroPadding
 462        # strip the first and last fsZeroPadding time points
 463        self.fieldV = self.fieldV[
 464            (self.fieldTimeV >= np.min(self.fieldTimeV) + timeRange)
 465            & (self.fieldTimeV <= np.max(self.fieldTimeV) - timeRange)
 466        ]
 467        self.fieldTimeV = self.fieldTimeV[
 468            (self.fieldTimeV >= np.min(self.fieldTimeV) + timeRange)
 469            & (self.fieldTimeV <= np.max(self.fieldTimeV) - timeRange)
 470        ]
 471        if self.fieldStdevV is None:
 472            return True
 473        if len(self.fieldStdevV) > len(self.fieldV):
 474            n_remove = len(self.fieldStdevV) - len(self.fieldV)
 475            print(
 476                "\n\nWarning: fieldStdevV is longer than fieldV, removing",
 477                n_remove,
 478                " elements\n\n",
 479            )
 480            while n_remove > 0:
 481                if n_remove % 2 == 0:
 482                    self.fieldStdevV = np.delete(self.fieldStdevV, 0)
 483                else:
 484                    self.fieldStdevV = np.delete(self.fieldStdevV, -1)
 485                n_remove -= 1
 486        elif len(self.fieldStdevV) < len(self.fieldV):
 487            n_add = len(self.fieldV) - len(self.fieldStdevV)
 488            print(
 489                "\n\nWarning: fieldStdevV is shorter than fieldV, adding",
 490                n_add,
 491                " elements\n\n",
 492            )
 493            while n_add > 0:
 494                if n_add % 2 == 0:
 495                    self.fieldStdevV = np.insert(self.fieldStdevV, 0, 0)
 496                else:
 497                    self.fieldStdevV = np.append(self.fieldStdevV, 0)
 498                n_add -= 1
 499        return True
 500
 501    def strip_from_complex_trace(self, timeRange=None):
 502        """from the complex trace eliminates the zeros appended (zero-padding) when computing the FFT (Analogous to strip_from_trace())."""
 503        if timeRange is None:
 504            timeRange = self.fsZeroPadding
 505        # strip the first and last fsZeroPadding time points
 506        self.complexFieldV = self.complexFieldV[
 507            (self.complexFieldTimeV >= np.min(self.complexFieldTimeV) + timeRange)
 508            & (self.complexFieldTimeV <= np.max(self.complexFieldTimeV) - timeRange)
 509        ]
 510        self.complexFieldTimeV = self.complexFieldTimeV[
 511            (self.complexFieldTimeV >= np.min(self.complexFieldTimeV) + timeRange)
 512            & (self.complexFieldTimeV <= np.max(self.complexFieldTimeV) - timeRange)
 513        ]
 514
 515    def tukey_time_window(self, lowEdge, upEdge, lowEdgeWidth, upEdgeWidth):
 516        """applies a tukey window to the trace in the time domain.
 517
 518        Args:
 519            lowEdge, upEdge: lower and upper edge of the window upEdge-lowEdge = FWHM of the window
 520            lowEdgeWidth, upEdgeWidth: width of the cosine-shaped edges (from 0 to 1 or viceversa)
 521        """
 522        window = asymmetric_tukey_window(
 523            np.abs(self.fieldTimeV), lowEdge, upEdge, lowEdgeWidth, upEdgeWidth
 524        )
 525        self.fieldV = self.fieldV * window
 526        self.update_fft()
 527
 528    def fourier_interpolation(self, ntimes_finer: int):
 529        """Shrinks the time step of the trace in time domain by a factor ntimes_finer.
 530
 531        This is done by 'zero-padding' the fourier transform of the trace, i.e., by extending the frequency axis by a factor ntimes_finer, and adding corresponding zeros to the FFT field.
 532
 533        Args:
 534            ntimes_finer (int): the factor by which to shrink the time step of the trace in time domain. Must be >= 1.
 535        """
 536        if ntimes_finer < 1:
 537            raise ValueError(
 538                "in function TraceHandler.fourier_interpolation() ntimes_finer must be >= 1"
 539            )
 540        if self.frequencyAxis is None or self.fftFieldV is None:
 541            raise ValueError(
 542                "in function TraceHandler.fourier_interpolation() frequency axis or fft field is not defined"
 543            )
 544        # extend the frequency axis by a factor ntimes_finer
 545        df = self.frequencyAxis[1] - self.frequencyAxis[0]
 546        i_last_pos = int(np.ceil(len(self.frequencyAxis) / 2) - 1)
 547        if self.frequencyAxis[i_last_pos] < 0 or self.frequencyAxis[i_last_pos + 1] > 0:
 548            print(self.frequencyAxis[0 : i_last_pos + 1])
 549            print(self.frequencyAxis[i_last_pos + 1 :])
 550            raise Exception(
 551                "in function fourier_interpolation() frequency axes was not extended correctly"
 552            )
 553
 554        appended_freqFFT = np.linspace(
 555            self.frequencyAxis[i_last_pos] + df,
 556            ntimes_finer * self.frequencyAxis[i_last_pos] + df,
 557            round((ntimes_finer - 1) * self.frequencyAxis[i_last_pos] / df),
 558        )
 559        prepended_freqFFT = np.linspace(
 560            -(ntimes_finer - 1) * self.frequencyAxis[i_last_pos]
 561            + self.frequencyAxis[i_last_pos + 1],
 562            +self.frequencyAxis[i_last_pos + 1],
 563            round((ntimes_finer - 1) * self.frequencyAxis[i_last_pos] / df),
 564        )
 565        self.frequencyAxis = np.concatenate(
 566            (
 567                self.frequencyAxis[: i_last_pos + 1],
 568                appended_freqFFT,
 569                prepended_freqFFT,
 570                self.frequencyAxis[i_last_pos + 1 :],
 571            )
 572        )
 573        self.fftFieldV = np.concatenate(
 574            (
 575                self.fftFieldV[: i_last_pos + 1],
 576                np.zeros(len(appended_freqFFT)),
 577                np.zeros(len(prepended_freqFFT)),
 578                self.fftFieldV[i_last_pos + 1 :],
 579            )
 580        )
 581
 582        self.update_trace_from_fft()
 583        self.strip_from_trace()
 584        self.update_fft_spectrum()
 585
 586    def differentiate_trace(self, spectrally=True):
 587        """computes the time derivative of the trace.
 588
 589        Args:
 590            spectrally: bool; if True (default) the derivative will be computed in the spectral domain (maybe more stable numerically)
 591        """
 592        if spectrally:
 593            self.fftFieldV = 1j * self.frequencyAxis * self.fftFieldV
 594            self.update_trace_from_fft()
 595            self.update_fft_spectrum()
 596        else:
 597            self.fieldV = np.gradient(self.fieldV, self.fieldTimeV)
 598            self.update_fft()
 599
 600    def integrate_trace(self, spectrally=True):
 601        """integrates the trace in time.
 602
 603        Args:
 604            spectrally (bool): if True (default) the integral will be computed in the spectral domain (maybe more stable numerically)
 605        """
 606        if spectrally:
 607            if self.frequencyAxis[0] != 0:
 608                print(
 609                    "WARNING: in function TraceHandler.integrate_trace() the FFT axis' first element is not zero"
 610                )
 611            self.fftFieldV[1:] = -1j / self.frequencyAxis[1:] * self.fftFieldV[1:]
 612            self.update_trace_from_fft()
 613            self.update_fft_spectrum()
 614        else:
 615            self.fieldV = np.cumsum(self.fieldV) * np.mean(np.diff(self.fieldTimeV))
 616            self.update_fft()
 617
 618    def save_trace_to_file(self, filename, low_lim=None, up_lim=None):
 619        """Saves the trace to file
 620
 621        Args:
 622            filename
 623            low_lim, up_lim: only save the trace between low_lim and up_lim (default: None, None)
 624        """
 625        if low_lim is not None and up_lim is not None:
 626            fieldTimeV_to_write = self.fieldTimeV[
 627                (self.fieldTimeV >= low_lim) & (self.fieldTimeV <= up_lim)
 628            ]
 629            fieldV_to_write = self.fieldV[
 630                (self.fieldTimeV >= low_lim) & (self.fieldTimeV <= up_lim)
 631            ]
 632        else:
 633            fieldTimeV_to_write = self.fieldTimeV
 634            fieldV_to_write = self.fieldV
 635        data = pandas.DataFrame(
 636            {"delay (fs)": fieldTimeV_to_write, "field (a.u.)": fieldV_to_write}
 637        )
 638        data.to_csv(filename, sep="\t", index=False)
 639
 640    def normalize_spectrum(self):
 641        """Normalizes the comparison spectrum to its integral."""
 642        spectrum_range = [60, 930]
 643        n_factor = (
 644            np.abs(
 645                np.diff(self.wvlSpectrometer)[
 646                    (self.wvlSpectrometer[:-1] > spectrum_range[0])
 647                    & (self.wvlSpectrometer[:-1] < spectrum_range[1])
 648                ]
 649            )
 650            * self.ISpectrometer[:-1][
 651                (self.wvlSpectrometer[:-1] > spectrum_range[0])
 652                & (self.wvlSpectrometer[:-1] < spectrum_range[1])
 653            ]
 654        ).sum()
 655        self.ISpectrometer /= n_factor
 656
 657    def normalize_fft_spectrum(self, spectrum_range=[60, 930]):
 658        """Normalizes the spectrum of the trace (|FFT{trace}|^2 * df/dλ) to its integral (area)
 659
 660        Args:
 661            spectrum_range (list): the range of wavelengths (nm) to consider for the normalization. Default is [60, 930] nm.
 662        """
 663        n_factor = (
 664            np.abs(
 665                np.diff(self.wvlAxis)[
 666                    (self.wvlAxis[:-1] > spectrum_range[0])
 667                    & (self.wvlAxis[:-1] < spectrum_range[1])
 668                ]
 669            )
 670            * self.fftSpectrum[:-1][
 671                (self.wvlAxis[:-1] > spectrum_range[0])
 672                & (self.wvlAxis[:-1] < spectrum_range[1])
 673            ]
 674        ).sum()
 675        self.fftSpectrum /= n_factor
 676
 677    def get_trace(self):
 678        """Returns the time (fs) and field arrays. Field is in a.u. unless set_fluence() was called or a calibrated field was given as an input."""
 679        return self.fieldTimeV, self.fieldV
 680
 681    def get_spectrum_trace(self):
 682        """Returns the wavelength and spectral intensity corresponding to the fourier transform of the trace."""
 683        return self.wvlAxis, self.fftSpectrum
 684
 685    def get_spectral_phase(self):
 686        """Returns the wavelength array and the spectral phase array corresponding to the FFT of the trace."""
 687        return self.wvlAxis, self.fftphase
 688
 689    def get_stdev(self):
 690        """Returns only the standard deviation (sigma) array"""
 691        return self.fieldStdevV
 692
 693    def get_positive_fft_field(self):
 694        """Returns the positive frequency axis f and the corresponding FFT field (complex).
 695
 696        Main task of this function is to 'smooth' the phase of the fft (= remove linear phase, corresponding to a time shift in the temporal domain), since the waveform is usually centered around t = 0"""
 697        if self.frequencyAxis is None or self.fftFieldV is None:
 698            raise ValueError(
 699                "in function TraceHandler.get_positive_fft_field() frequency axis or fft field is not defined"
 700            )
 701        positive_freq = self.frequencyAxis[self.frequencyAxis >= 0]
 702        positive_fft_field = self.fftFieldV[self.frequencyAxis >= 0]
 703        # remove linear phase
 704        tmax = self.get_zero_delay()
 705        deltaT1 = tmax - self.fieldTimeV[0]
 706        deltaT2 = self.fieldTimeV[-1] - tmax
 707        ifft_twindow = 1 / (self.frequencyAxis[1] - self.frequencyAxis[0])
 708        total_delay = ifft_twindow / 2 + (deltaT2 - deltaT1) / 2
 709        positive_fft_field = positive_fft_field * np.exp(
 710            -1.0j * 2 * np.pi * positive_freq * total_delay
 711        )
 712        return positive_freq, positive_fft_field
 713
 714    def get_fluence(self):
 715        """Calculates the fluence of the trace.
 716
 717        convention: field [V/Å], time [fs], fluence [J/cm²]
 718        F = c eps_0 integral(E^2)dt
 719
 720
 721        Returns:
 722            fluence: float
 723        """
 724        F = (
 725            constants.speed_of_light
 726            * constants.epsilon_0
 727            * np.trapz(self.fieldV**2, self.fieldTimeV)
 728            * 1e1
 729        )
 730        return F
 731
 732    def set_fluence(self, fluence):
 733        """Set the fluence of the trace.
 734
 735        Convention: field [V/Å], time [fs], fluence [J/cm²]
 736
 737        F = c eps_0 integral(E^2)dt"""
 738        F = (
 739            constants.speed_of_light
 740            * constants.epsilon_0
 741            * np.trapz(self.fieldV**2, self.fieldTimeV)
 742            * 1e1
 743        )
 744        self.fieldV *= np.sqrt(fluence / F)
 745        self.update_fft()
 746        self.update_fft_spectrum()
 747
 748    def compute_complex_field(self):
 749        """Computes the complex trace by inverting the positive-freq FFT and stores it in self.complexFieldV (for envelope or phase computation, for example)
 750
 751        Notice that the fourier transform of the trace is not computed, it is assumed to be already stored in the class. This should be true in most cases."""
 752        complex_spectrum = deepcopy(self.fftFieldV)
 753        complex_spectrum[self.frequencyAxis < 0] = 0
 754        self.complexFieldTimeV, self.complexFieldV = inverse_fourier_transform(
 755            self.frequencyAxis, complex_spectrum
 756        )
 757        self.complexFieldV = 2 * self.complexFieldV
 758        self.strip_from_complex_trace()
 759
 760    def get_envelope(self):
 761        """Get the envelope of the trace.
 762
 763        Returns:
 764            time array (fs)
 765            envelope array
 766        """
 767        if self.complexFieldV is None:
 768            self.compute_complex_field()
 769        return self.complexFieldTimeV, np.abs(self.complexFieldV)
 770
 771    def get_phase(self):
 772        """Get the instantaneous phase of the trace.
 773
 774        If complexFieldV is already computed and stored, no re-calculation occurs
 775
 776        Returns:
 777            time array (fs)
 778            instantaneous phase array (fs)
 779        """
 780        if self.complexFieldV is None:
 781            self.compute_complex_field()
 782        return self.complexFieldTimeV, np.angle(self.complexFieldV)
 783
 784    def get_zero_delay(self):
 785        """Get the time value corresponding to the envelope peak.
 786
 787        Returns:
 788
 789            zero_delay: float
 790        """
 791        # careful: the time grid should be fine enough to resolve the maximum of the envelope
 792        t, en = self.get_envelope()
 793        dt = t[1] - t[0]
 794        max_index, max_value = find_maximum_location(en)
 795        self.zero_delay = t[0] + dt * max_index
 796        return self.zero_delay
 797
 798    def get_FWHM(self):
 799        """get the FWHM of the trace.
 800
 801
 802        Returns:
 803            FWHM: float
 804        """
 805        t, en = self.get_envelope()
 806        dt = t[1] - t[0]
 807        return fwhm(en**2, dt)
 808
 809    def fft_tukey_bandpass(
 810        self, lowWavelengthEdge, upWavelengthEdge, lowEdgeWidth, highEdgeWidth
 811    ):
 812        """Applies a bandpass filter to the trace in the frequency domain using a tukey window.
 813
 814        The tukey window is 1 between  lowWavelengthEdge+lowEdgeWidth/2  and  upWavelengthEdge-highEdgeWidth/2
 815        and it is reaches zero at lowWavelengthEdge-lowEdgeWidth/2  and  upWavelengthEdge+highEdgeWidth/2.
 816        Notice that the edges are only cosine-shaped in the frequency domain. In the wavelength domain upWavelengthEdge - lowWavelengthEdge does not coincide with the FWHM of the tukey function
 817
 818        Args:
 819            lowWavelengthEdge: float (nm)
 820            upWavelengthEdge: float (nm)
 821            lowEdgeWidth: float (nm)
 822            highEdgeWidth: float (nm)
 823        """
 824        upFreqEdge = (
 825            constants.speed_of_light / (lowWavelengthEdge - lowEdgeWidth / 2) / 2 * 1e-6
 826            + constants.speed_of_light
 827            / (lowWavelengthEdge + lowEdgeWidth / 2)
 828            / 2
 829            * 1e-6
 830        )
 831        lowFreqEdge = (
 832            constants.speed_of_light / (upWavelengthEdge - highEdgeWidth / 2) / 2 * 1e-6
 833            + constants.speed_of_light
 834            / (upWavelengthEdge + highEdgeWidth / 2)
 835            / 2
 836            * 1e-6
 837        )
 838
 839        upFreqEdgeWidth = (
 840            constants.speed_of_light / (lowWavelengthEdge - lowEdgeWidth / 2) * 1e-6
 841            - constants.speed_of_light / (lowWavelengthEdge + lowEdgeWidth / 2) * 1e-6
 842        )
 843        lowFreqEdgeWidth = (
 844            constants.speed_of_light / (upWavelengthEdge - highEdgeWidth / 2) * 1e-6
 845            - constants.speed_of_light / (upWavelengthEdge + highEdgeWidth / 2) * 1e-6
 846        )
 847
 848        window = asymmetric_tukey_window(
 849            np.abs(self.frequencyAxis),
 850            lowFreqEdge,
 851            upFreqEdge,
 852            lowFreqEdgeWidth,
 853            upFreqEdgeWidth,
 854        )
 855        self.fftFieldV = self.fftFieldV * window
 856        self.update_trace_from_fft()
 857        self.update_fft_spectrum()
 858        self.strip_from_trace()
 859
 860    def apply_transmission(self, wavelengths, f):
 861        """Applies a spectral transmission function to the trace (e.g. spectral filter).
 862
 863        Args:
 864            wavelengths: ndarray = wavelength array (nm)
 865            f: ndarray = transmission function f(λ)
 866        """
 867        if np.any(np.diff(wavelengths)) < 0:
 868            if np.all(np.diff(wavelengths)) < 0:
 869                wavelengths = wavelengths[::-1]
 870                f = f[::-1]
 871            else:
 872                raise ValueError(
 873                    "in function TraceHandler.apply_transmission() wavelength array is not monotonous"
 874                )
 875        freq = constants.speed_of_light / wavelengths[::-1] * 1e-6
 876        spectrum_freq = f[::-1]
 877
 878        # fill with zeros
 879        initial_zero_freq = np.linspace(
 880            freq[1] - freq[0], freq[0], int(np.ceil(4 * freq[0] / (freq[1] - freq[0])))
 881        )
 882        initial_zeros = np.zeros(len(initial_zero_freq))
 883        final_zero_freq = np.linspace(
 884            2 * freq[-1] - freq[-2],
 885            4 * freq[-1],
 886            int(np.ceil(15 * freq[-1] / (freq[-1] - freq[-2]))),
 887        )
 888        final_zeros = np.zeros(len(final_zero_freq))
 889        freq = np.concatenate((initial_zero_freq, freq, final_zero_freq))
 890        spectrum_freq = np.concatenate((initial_zeros, spectrum_freq, final_zeros))
 891        # add negative frequencies
 892        freq = np.concatenate((-freq[::-1], freq))
 893        spectrum_freq = np.concatenate((spectrum_freq[::-1], spectrum_freq))
 894
 895        # interpolate the spectrum to the frequency axis of the fft
 896        spectrum_interp = np.interp(self.frequencyAxis, freq, spectrum_freq)
 897
 898        self.fftFieldV = self.fftFieldV * spectrum_interp
 899        self.update_trace_from_fft()
 900        self.update_fft_spectrum()
 901        self.strip_from_trace()
 902
 903    def apply_spectrum(
 904        self,
 905        wvl=None,
 906        spectrum=None,
 907        CEP_shift: float = 0.0,
 908        stripZeroPadding: bool = True,
 909    ):
 910        """Applies a spectrum to the phase of the trace. This means that a new trace is computed and stored in the TraceHandler object (replacing the existing one);
 911        the new trace has spectral intensity equal to the applied spectrum and spectral phase equal to the phase of the existing trace.
 912
 913        Args:
 914            wvl: wavelength array (nm). If None (default) the comparison spectrum stored in the class (self.wvlSpectrometer and self.ISpectrometer) is applied.
 915            spectrum: spectral intensity array. If None (default) the comparison spectrum stored in the class (self.wvlSpectrometer and self.ISpectrometer) is applied.
 916            CEP_shift: a possible artificial phase shift IN UNITS OF PI! (default = 0)
 917            stripZeroPadding (bool): whether to eliminate the zero padding (zeros appended to the trace). Default is True"""
 918
 919        if wvl is None or spectrum is None:
 920            wvl = self.wvlSpectrometer
 921            spectrum = self.ISpectrometer
 922        # check that the wvl array is monotonically increasing
 923        if np.any(np.diff(wvl) < 0):
 924            if np.all(np.diff(wvl) < 0):
 925                wvl = wvl[::-1]
 926                spectrum = spectrum[::-1]
 927            else:
 928                raise ValueError(
 929                    "in function TraceHandler.apply_spectrum() wavelength array is not monotonous"
 930                )
 931
 932        # frequency spectrum (monotonically increasing)
 933        freq = constants.speed_of_light / wvl[::-1] * 1e-6
 934        spectrum_freq = spectrum[::-1] / (
 935            constants.speed_of_light / wvl[::-1] ** 2 * 1e-6
 936        )
 937        spectrum_freq = np.maximum(spectrum_freq, np.zeros(len(spectrum_freq)))
 938
 939        # fill with zeros
 940        initial_zero_freq = np.linspace(
 941            freq[1] - freq[0], freq[0], int(np.ceil(4 * freq[0] / (freq[1] - freq[0])))
 942        )
 943        initial_zeros = np.zeros(len(initial_zero_freq))
 944        final_zero_freq = np.linspace(
 945            2 * freq[-1] - freq[-2],
 946            4 * freq[-1],
 947            int(np.ceil(15 * freq[-1] / (freq[-1] - freq[-2]))),
 948        )
 949        final_zeros = np.zeros(len(final_zero_freq))
 950        freq = np.concatenate((initial_zero_freq, freq, final_zero_freq))
 951        spectrum_freq = np.concatenate((initial_zeros, spectrum_freq, final_zeros))
 952        # add negative frequencies
 953        freq = np.concatenate((-freq[::-1], freq))
 954        spectrum_freq = np.concatenate((spectrum_freq[::-1], spectrum_freq))
 955
 956        # interpolate the spectrum to the frequency axis of the fft
 957        spectrum_interp = np.interp(self.frequencyAxis, freq, spectrum_freq)
 958        self.fftFieldV = np.sqrt(spectrum_interp) * np.exp(
 959            1j * np.angle(self.fftFieldV)
 960            + 1j * np.pi * CEP_shift * np.sign(self.frequencyAxis)
 961        )
 962
 963        self.update_trace_from_fft()
 964        if stripZeroPadding:
 965            self.strip_from_trace()
 966        else:
 967            self.fieldStdevV = None
 968        self.update_fft_spectrum()
 969
 970    def fresnel_reflection(
 971        self,
 972        material2,
 973        angle_in,
 974        material1=None,
 975        forward: bool = True,
 976        s_polarized: bool = True,
 977        path: str = "./RefractiveIndices/",
 978    ):
 979        """calculates the fresnel reflection of the pulse at the interface between two materials. The waveform is travelling from material 1 to material 2.
 980        The first medium (material1) should be non-absorptive. As usual the resulting waveform is stored in the TraceHandler object, replacing the previous one.
 981
 982        Refractive index files should contain 3 space-separated columns, respectively with headers: wvl n k, where wvl is the wavelength in um, n and k resp. the real and imaginary part of the refractive index.
 983        Args:
 984            material2: the filename (without '.txt') of the refractive index data for the material after the interface (e.g. Si Al MgF2); wavelength is in um in the refractive index file
 985            angle_in: the incidence angle in degrees
 986            material1: the filename (without '.txt') of the refractive index data for the material before the interface. If None (default), vacuum is assumed
 987            forward (bool):  if True (default) forward reflection is computed (the result waveform is the reflection of the previously stored waveform.
 988
 989                if False backward reflection is computed (the previous waveform is the reflection of the result waveform)
 990            s_polarized (bool): True. Reflection calculation only implemented for s-polarized light
 991            path (str): path for the refractive index files. Defaults to "./RefractiveIndices/"
 992        """
 993        if not s_polarized:
 994            raise ValueError(
 995                "in function TraceHandler.fresnel_reflection() p_polarized is not implemented yet\n"
 996            )
 997
 998        # read refractive index
 999        refIndexData = pandas.read_table(
1000            path + material2 + ".txt", sep=" ", keep_default_na=False
1001        )
1002        wvl2 = np.array(refIndexData["wvl"]) * 1e3
1003        n2 = np.array(refIndexData["n"]) + 1j * np.array(refIndexData["k"])
1004        if material1 is None:
1005            wvl1 = np.array(wvl2)
1006            n1 = n2 * 0 + 1
1007        else:
1008            refIndexData = pandas.read_table(
1009                path + material1 + ".txt", sep=" ", keep_default_na=False
1010            )
1011            wvl1 = np.array(refIndexData["wvl"]) * 1e3
1012            n1 = np.array(refIndexData["n"]) + 1j * np.array(refIndexData["k"])
1013
1014        # check that the wvl arrays are monotonically increasing
1015        if np.any(np.diff(wvl1) < 0):
1016            if np.all(np.diff(wvl1) < 0):
1017                wvl1 = wvl1[::-1]
1018                n1 = n1[::-1]
1019            else:
1020                raise ValueError(
1021                    "in function TraceHandler.fresnel_reflection() wavelength array is not monotonous"
1022                )
1023        if np.any(np.diff(wvl2) < 0):
1024            if np.all(np.diff(wvl2) < 0):
1025                wvl2 = wvl2[::-1]
1026                n2 = n2[::-1]
1027            else:
1028                raise ValueError(
1029                    "in function TraceHandler.fresnel_reflection() wavelength array is not monotonous"
1030                )
1031
1032        # frequency spectrum (monotonically increasing)
1033        freq1 = constants.speed_of_light / wvl1[::-1] * 1e-6
1034        freq2 = constants.speed_of_light / wvl2[::-1] * 1e-6
1035        n1 = n1[::-1]
1036        n2 = n2[::-1]
1037
1038        # fill with ones
1039        initial_ones_freq = np.linspace(
1040            freq1[1] - freq1[0],
1041            freq1[0],
1042            int(np.ceil(4 * freq1[0] / (freq1[1] - freq1[0]))),
1043        )
1044        initial_ones = np.ones(len(initial_ones_freq)) * n1[0]
1045        final_ones_freq = np.linspace(
1046            2 * freq1[-1] - freq1[-2],
1047            4 * freq1[-1],
1048            int(np.ceil(15 * freq1[-1] / (freq1[-1] - freq1[-2]))),
1049        )
1050        final_ones = np.ones(len(final_ones_freq)) * n1[-1]
1051        freq1 = np.concatenate((initial_ones_freq, freq1, final_ones_freq))
1052        n1 = np.concatenate((initial_ones, n1, final_ones))
1053        # same for n2
1054        initial_ones_freq = np.linspace(
1055            freq2[1] - freq2[0],
1056            freq2[0],
1057            int(np.ceil(4 * freq2[0] / (freq2[1] - freq2[0]))),
1058        )
1059        initial_ones = np.ones(len(initial_ones_freq)) * n2[0]
1060        final_ones_freq = np.linspace(
1061            2 * freq2[-1] - freq2[-2],
1062            4 * freq2[-1],
1063            int(np.ceil(15 * freq2[-1] / (freq2[-1] - freq2[-2]))),
1064        )
1065        final_ones = np.ones(len(final_ones_freq)) * n2[-1]
1066        freq2 = np.concatenate((initial_ones_freq, freq2, final_ones_freq))
1067        n2 = np.concatenate((initial_ones, n2, final_ones))
1068
1069        # add negative frequencies
1070        freq1 = np.concatenate((-freq1[::-1], freq1))
1071        n1 = np.concatenate((np.conjugate(n1[::-1]), n1))
1072        freq2 = np.concatenate((-freq2[::-1], freq2))
1073        n2 = np.concatenate((np.conjugate(n2[::-1]), n2))
1074
1075        # interpolate the n1 and n2 to the frequency axis of the fft
1076        n1Interp = np.interp(self.frequencyAxis, freq1, n1)
1077        n2Interp = np.interp(self.frequencyAxis, freq2, n2)
1078
1079        # calculate the fresnel reflection coefficients using only the angle of incidence (NOT SURE THIS WORKS WHEN THE FIRST MEDIUM IS LOSSY)
1080        r = (
1081            n1Interp * np.cos(angle_in * np.pi / 180)
1082            - n2Interp
1083            * np.sqrt(1 - (n1Interp / n2Interp * np.sin(angle_in * np.pi / 180)) ** 2)
1084        ) / (
1085            n1Interp * np.cos(angle_in * np.pi / 180)
1086            + n2Interp
1087            * np.sqrt(1 - (n1Interp / n2Interp * np.sin(angle_in * np.pi / 180)) ** 2)
1088        )
1089        if forward:
1090            self.fftFieldV = self.fftFieldV * r
1091        else:
1092            self.fftFieldV = self.fftFieldV / r
1093
1094        self.fieldStdevV = None
1095
1096        self.update_trace_from_fft()
1097        self.strip_from_trace()
1098        self.update_fft_spectrum()
1099
1100    def apply_zero_phase(self):
1101        """Applies zero-phase to the trace; this allows, for example, to retrieve the fourier limited pulse corresponding to the same FFT spectrum of the loaded trace"""
1102        tau = (
1103            self.fsZeroPadding + (np.max(self.fieldTimeV) - np.min(self.fieldTimeV)) / 2
1104        )
1105        self.fftFieldV = np.abs(self.fftFieldV) * np.exp(
1106            1.0j * (2 * np.pi * tau * self.frequencyAxis)
1107        )
1108        self.update_trace_from_fft()
1109        self.strip_from_trace()
1110        self.update_fft_spectrum()
1111
1112    def time_frequency_analysis(
1113        self, sigma_time, low_lim=None, up_lim=None, low_lim_freq=None, up_lim_freq=None
1114    ):
1115        """Performs time-frequency analysis by using scipy's short time fourier transform (fourier transform of the trace convoluted by a 'sigma_time' broad gaussian).
1116
1117        Args:
1118            sigma_time: sigma of the gaussian window
1119            low_lim, up_lim (float): xaxis limits for plotting. Default None
1120            low_lim_freq, up_lim_freq (float): xaxis limits for plotting. Default None
1121
1122        """
1123        dt = np.mean(np.diff(self.fieldTimeV))
1124        w = scipy.signal.windows.gaussian(
1125            int(sigma_time / dt * 6) + 1, sigma_time / dt, sym=True
1126        )
1127        TFA = scipy.signal.ShortTimeFFT(
1128            w, hop=1, fs=1.0 / dt, mfft=int(sigma_time / dt * 24), scale_to="magnitude"
1129        )
1130        TFData = TFA.stft(self.fieldV)
1131
1132        fig, ax = plt.subplots()
1133        t_lo, t_hi, f_lo, f_hi = TFA.extent(
1134            self.fieldV.size
1135        )  # time and freq range of plot
1136        if low_lim is None:
1137            low_lim = t_lo + self.fieldTimeV[0]
1138        if up_lim is None:
1139            up_lim = t_hi + self.fieldTimeV[0]
1140        if low_lim_freq is None:
1141            low_lim_freq = 0
1142        if up_lim_freq is None:
1143            up_lim_freq = 3.0
1144
1145        ax.set(xlim=(low_lim, up_lim), ylim=(low_lim_freq, up_lim_freq))
1146        ax.set_xlabel("Time (fs)")
1147        ax.set_ylabel("Frequency (PHz)")
1148        im1 = ax.imshow(
1149            abs(TFData) / np.max(abs(TFData)),
1150            origin="lower",
1151            aspect="auto",
1152            extent=(t_lo + self.fieldTimeV[0], t_hi + self.fieldTimeV[0], f_lo, f_hi),
1153            cmap="viridis",
1154        )
1155        cbar = fig.colorbar(im1)
1156
1157        cbar.ax.set_ylabel("Magnitude of the field (Arb. unit)")
1158        fig.tight_layout()
1159        return TFData
1160
1161    def plot_trace(self, low_lim=None, up_lim=None, normalize: bool = True):
1162        """Plots the field trace.
1163
1164        Args:
1165            low_lim, up_lim: float = xaxis limits for plotting. Default None
1166            normalize: bool = if True (default) normalize the peak of the trace to 1
1167        """
1168        fig, ax = plt.subplots()
1169        if normalize:
1170            norm_plot = self.normalization_trace
1171        else:
1172            norm_plot = 1
1173        main_line = ax.plot(
1174            self.fieldTimeV, self.fieldV / norm_plot, label="Field trace"
1175        )
1176        if self.fieldStdevV is not None:
1177            ax.fill_between(
1178                self.fieldTimeV,
1179                (self.fieldV - self.fieldStdevV) / norm_plot,
1180                (self.fieldV + self.fieldStdevV) / norm_plot,
1181                color=main_line[0].get_color(),
1182                alpha=0.3,
1183            )
1184        ax.set_xlabel("Time (fs)")
1185        ax.set_ylabel("Field (Arb. unit)")
1186
1187        if low_lim is not None and up_lim is not None:
1188            ax.set_xlim(low_lim, up_lim)
1189        return fig
1190
1191    def plot_spectrum(
1192        self, low_lim=40, up_lim=1000, no_phase: bool = False, phase_blanking_level=0.05
1193    ):
1194        """Plots the trace spectrum and phase together with the spectrometer measurement [if provided].
1195
1196        Args:
1197            low_lim, up_lim (float): xaxis limits for plotting. Default: 40, 1000
1198            no_phase: if True, don't plot the phase. Default: False
1199        """
1200        fig, ax = plt.subplots()
1201
1202        if not no_phase:
1203            ax2 = ax.twinx()
1204        lines = []
1205        min_intensity = phase_blanking_level * np.max(self.fftSpectrum)
1206        lines += ax.plot(
1207            self.wvlAxis[(self.wvlAxis > low_lim) & (self.wvlAxis < up_lim)],
1208            self.fftSpectrum[(self.wvlAxis > low_lim) & (self.wvlAxis < up_lim)],
1209            label="Fourier transform",
1210        )
1211        if not no_phase:
1212            ax2.plot([], [])
1213            if self.wvlSpectrometer is not None:
1214                ax2.plot([], [])
1215            lines += ax2.plot(
1216                self.wvlAxis[
1217                    (self.wvlAxis > low_lim)
1218                    & (self.wvlAxis < up_lim)
1219                    & (self.fftSpectrum > min_intensity)
1220                ],
1221                self.fftphase[
1222                    (self.wvlAxis > low_lim)
1223                    & (self.wvlAxis < up_lim)
1224                    & (self.fftSpectrum > min_intensity)
1225                ],
1226                "--",
1227                label="Phase",
1228            )
1229        if self.wvlSpectrometer is not None:
1230            lines += ax.plot(
1231                self.wvlSpectrometer[
1232                    (self.wvlSpectrometer > low_lim) & (self.wvlSpectrometer < up_lim)
1233                ],
1234                self.ISpectrometer[
1235                    (self.wvlSpectrometer > low_lim) & (self.wvlSpectrometer < up_lim)
1236                ],
1237                label="Spectrometer",
1238            )
1239        ax.set_xlabel("Wavelength (nm)")
1240        ax.set_ylabel("Intensity (Arb. unit)")
1241        ax.tick_params(axis="both")
1242        if not no_phase:
1243            ax2.set_ylabel("Phase (rad)")
1244            ax2.tick_params(axis="y")
1245        ax.legend(lines, [l.get_label() for l in lines])
1246        return fig