attoworld.wave

 7def align_waves(
 8    waves, dt: float, frequency_roi_start: float, frequency_roi_stop: float
 9):
10    """
11    Align a set of waveforms, inside of a 2D numpy array.
12    Args:
13        waves: set of waveforms
14        dt (float): time step (assuming constant spacing), (s)
15        frequency_roi_start (float): start frequency of region of interest for alignment (Hz)
16        frequency_roi_stop (float): stop frequency of region of interest for alignment (Hz)
17
18    Returns:
19        np.ndarray: set of aligned waves
20    """
21
22    waves_f = np.array(waves)
23
24    # windowing and offset removal
25    for i in range(waves_f.shape[0]):
26        waves_f[i, :] = sig.windows.tukey(waves_f.shape[1]) * (
27            waves_f[i, :] - np.mean(waves_f[i, :])
28        )
29
30    # find frequency roi
31    f = np.fft.fftfreq(waves.shape[1], dt)
32    f0 = np.argmin(np.abs(f - frequency_roi_start))
33    f1 = np.argmin(np.abs(f - frequency_roi_stop))
34    w_roi = 2 * np.pi * f[f0:f1]
35
36    # get active spectral region
37    waves_f = np.fft.fft(waves, axis=1)
38    waves_roi = np.array(waves_f[:, f0:f1])
39    waves_roi /= np.max(np.max(np.abs(waves_roi)))
40
41    # apply tau phase shifts
42    def apply_taus(spectrum, taus, w):
43        spectrum_shifted = np.zeros(spectrum.shape, dtype=complex)
44        for i in range(spectrum.shape[0]):
45            spectrum_shifted[i, :] = np.exp(-1j * 1e-18 * taus[i] * w) * spectrum[i, :]
46        return spectrum_shifted
47
48    # return fitting weights
49    def get_residual(taus):
50        shifted = apply_taus(waves_roi, taus, w_roi)
51        mean_amplitudes = np.mean(shifted, axis=0)
52        return 5.0 - np.abs(mean_amplitudes)
53
54    # apply fitting
55    res = opt.least_squares(
56        get_residual, np.zeros(waves.shape[0]), ftol=1e-12, max_nfev=16384
57    )
58
59    # remove mean shift
60    res.x -= np.mean(res.x)
61
62    print(f"Rms shift in attoseconds: {np.std(res.x)}")
63    return np.real(np.fft.ifft(apply_taus(waves_f, res.x, 2 * np.pi * f)))

171def reconstruct_shg_frog(
172    measurement: Spectrogram,
173    test_iterations: int = 100,
174    polish_iterations=5000,
175    repeats: int = 256,
176):
177    """
178    Run the core FROG loop several times and pick the best result
179
180    Args:
181        measurement (np.ndarray): measured spectrogram, sqrt + fftshift(axes=0)
182        test_iterations (int): number of iterations for the multiple tests
183        polish_iteration (int): number of extra iterations to apply to the winner
184        repeats (int): number of different initial guesses to try
185
186    Returns:
187    FrogData: the completed reconstruction
188    """
189    sqrt_sg = np.fft.fftshift(
190        np.sqrt(measurement.data - np.min(measurement.data[:])), axes=0
191    )
192    sqrt_sg /= np.max(sqrt_sg)
193    results = np.zeros((sqrt_sg.shape[0], repeats), dtype=np.complex128)
194    errors = np.zeros(repeats, dtype=float)
195    for _i in range(repeats):
196        results[:, _i] = reconstruct_shg_frog_core(
197            sqrt_sg, max_iterations=test_iterations
198        )
199        errors[_i] = calculate_g_error(sqrt_sg, results[:, _i])
200    min_error_index = np.argmin(errors)
201    result = reconstruct_shg_frog_core(
202        sqrt_sg, guess=results[:, min_error_index], max_iterations=polish_iterations
203    )
204    return bundle_frog_reconstruction(
205        t=measurement.time,
206        result=result,
207        measurement=sqrt_sg,
208        f0=float(np.mean(measurement.freq) / 2.0),
209    )

65def generate_shg_spectrogram(Et, Gt):
66    """
67    Generate a SHG spectrogram, same pattern as in FROG book
68
69    Args:
70        Et: field
71        Gt: gate (same as field unless XFROG)
72    Returns:
73        np.ndarray: the complex spectrogram
74    """
75    spectrogram_timetime = np.outer(Et, Gt)
76    for _i in range(Et.shape[0]):
77        spectrogram_timetime[_i, :] = blank_roll(
78            spectrogram_timetime[_i, :], -_i + int(Et.shape[0] / 2)
79        )
80
81    return np.fft.fft(spectrogram_timetime, axis=0)

20def bundle_frog_reconstruction(
21    t, result, measurement, f0: float = 375e12, interpolation_factor: int = 100
22):
23    """
24    Turn the results of a FROG reconstruction into a FrogData struct
25
26    Args:
27        t: the time vector (s)
28        result: the reconstructed complex envelope
29        measurement: the measured spectrogram, sqrt-ed and fftshift-ed
30        f0 (float): the central frequency (Hz)
31        interpolation_factor (int): factor by which to interpolate to get a valid waveform (Nyquist)
32
33    Returns:
34        FrogData: the bundled data
35    """
36    f = np.fft.fftfreq(len(t), d=(t[1] - t[0]))
37    sg_freq = np.fft.fftshift(f) + 2 * f0
38    result_sg = Spectrogram(
39        data=np.fft.fftshift(
40            np.abs(generate_shg_spectrogram(result, result)) ** 2, axes=0
41        ),
42        time=t,
43        freq=sg_freq,
44    )
45    measurement_sg = Spectrogram(
46        data=np.fft.fftshift(np.abs(measurement) ** 2, axes=0), time=t, freq=sg_freq
47    )
48    result_ce = ComplexEnvelope(
49        time=t, dt=(t[1] - t[0]), carrier_frequency=f0, envelope=result
50    ).to_waveform(interpolation_factor=interpolation_factor)
51    result_cs = result_ce.to_complex_spectrum()
52
53    return FrogData(
54        measured_spectrogram=measurement_sg,
55        pulse=result_ce,
56        reconstructed_spectrogram=result_sg,
57        spectrum=result_cs,
58        raw_reconstruction=result,
59        f0=f0,
60        dt=(t[1] - t[0]),
61    )