PulsedJAX: Pulse-Retrieval with JAX
Uses JAX to implement various algorithms for pulse-retrieval in ultrafast optics.
Here, the point of JAX is to take advantage of jax.jit and other JAX-transformations. A consequence of this is that the implementations of the algorithms are differentiable as part of a pure function.
Installation and Documentation
To install (create a new environment and) run:
pip install git+https://github.com/matillda123/pulse-retrieval-with-jax.git
The documentation is available here.
Overview
Available methods are FROG, Chirp-Scans (material dispersion or pulse-shaping), Time-Domain-Ptychography (TDP), Two-Dimensional Spectral-Shearing Interferometry (2D-SI), VAMPIRE and Attosecond Streaking (via multichannel-SFA).
For each method SHG, THG, PG/TG and SD as well as cross-correlation and doubleblind retrievals are implemented (with some exceptions). For FROG and TDP interferometric retrievals are implemented as well.
In addition most algorithms support the retrieval of signals from Difference-Frequency-Generation and that of traces containing multiple nonlinear signals at once as well as the retrieval of the calibration curve of the measured trace. For Attosecond Streaking the retrieval of the momentum dependent dipole-transition-matrix-elements is available as well.
On top of the naive implementation of each algorithm, some allow the usage of standard nonlinear optimization approaches like nonlinear conjugate gradients or LBFGS. In other cases the pulses may be parametrized via analytic functions instead of a discrete grid.
The available algorithms are:
Algorithm |
Citation |
|---|---|
Generalized Projection |
|
Ptychographic Iterative Engine |
A. Maiden et al., Optica 4, 736-745 (2017) and P. Sidorenko et al., Optica 3, 1320-1330 (2016) |
Common Pulse Retrieval Algorithm |
|
Linesearch-Frog-Algorithm (for all methods) |
C. O. Krook and V. Pasiskevicius, Opt. Express 33, 33258-33269 (2025) |
Differential Evolution |
J. Qiang and C. Mitchell (2014) and I. Grigorenko and M.E. Garcia, Physica A 284 131–139 (2000) |
AD-Solvers (optax and optimistix) |
DeepMind et al., The DeepMind JAX Ecosystem (2020) and |
Evolutionary solvers (evosax) |
For some methods, specific additional algorithms are implemented:
Method |
Algorithms |
Citation |
|---|---|---|
FROG |
Vanilla |
|
FROG, TDP and VAMPIRE |
LSGPA |
J. Gagnon et al., Appl. Phys. B 92, 25-32, 10.1007/s00340-008-3063-x (2008) |
C-PCGPA |
D. J. Kane and A. B. Vakhtin, Prog. Quantum Electron. 81 (100364), (2022) |
|
2D-SI |
Direct-Reconstruction |
J. R. Birge et al., Opt. Lett. 31, 2063-2065, 10.1364/OL.31.002063 (2006) |
Adding new methods or algorithms should be relatively easy. The general approach for doing this is described here.
If you end up using this code for a publication, please use the citation below as well as the appropriate citations for the algorithm used.
@software{pulseretrievalwithjax,
title = {Pulse {R}etrieval with {JAX}},
author = {T. J. Stehling},
url = {https://github.com/matillda123/Pulse-Retrieval-with-Jax/},
year = {2025}
}