core

Subpackages

Submodules

core.base_classes_algorithms module

class core.base_classes_algorithms.AlgorithmsBASE(*args, **kwargs)[source]

Bases: object

The Base-Class for all solvers.

spectrum_is_being_used

what the name says

Type:

bool

fft(signal, sk, rn, axis=-1)[source]
ifft(signal, sk, rn, axis=-1)[source]
generate_signal_t(descent_state, measurement_info, descent_info)[source]

Applies calculate_signal_t to a whole population via jax.vmap

calculate_S_prime_individual(signal_t, measured_trace, mu, measurement_info, descent_info, local_or_global)[source]
calculate_S_prime_population(signal_t, measured_trace, mu, measurement_info, descent_info, local_or_global, axes)[source]
do_checks_before_running()[source]

Called by algorithm.run() performs a few checks

run(init_vals, no_iterations=1, **kwargs)[source]

This function is invoked by most solvers to perform the iterative reconstruction. Excpetions to this are PtychographicIterativeEngine and COPRA. These invoke run_scan() twice in order to perform a local and global optimization phase.

use_measured_spectrum(frequency, spectrum, pulse_or_gate='pulse')[source]

Needs to be called if a pulse spectrum is meant to be used in the reconstruction.

In Classical Algorithms. The spectrum will be applied as a projection of the pulse onto an average of itself and the spectrum of the current guess. Setting algorithm.optimize_spectral_phase_directly=True will as the name suggests optimize the spectral phase directly without a projection.

In General Optimization Algorithms, providing the spectrum will cause an optimization of the spectral phase.

Parameters:

  • frequency – jnp.array, the frequency axis of spectrum

  • spectrum – jnp.array, the spectrum

  • pulse_or_gate (str) – whether the spectrum is from the pulse or the gate-pulse.

Returns:

the class instance

class core.base_classes_algorithms.ClassicAlgorithmsBASE(*args, **kwargs)[source]

Bases: AlgorithmsBASE

The Base-Class for all classical solvers (e.g. Generalized Projection, …). Inherits from AlgorithmsBASE.

local_gamma

step size for local iterations

Type:

float

global_gamma

step size for global iterations

Type:

float

linesearch

enables/disables a linesearch. Can be False, backtracking or zoom.

Type:

bool, str

max_steps_linesearch

maximum number of linesearch steos

Type:

int

c1

constant for the Armijo-condition

Type:

float

c2

constant for the strong Wolfe-condition

Type:

float

delta_gamma

a factor which by which gamma is increased each iteration

Type:

float

local_newton

enables/disables the use of a hessian in local iterations. Can be False, lbfgs or diagonal.

Type:

bool, str

global_newton

enables/disables the use of a hessian in global iterations. Can be False, lbfgs, diagonal or full.

Type:

bool, str

lambda_lm

a Levenberg-Marquardt style damping coefficient.

Type:

float

lbfgs_memory

the number of past iterations to use in LBFGS

Type:

int

linalg_solver

chooses a library/method for inverting the hessian. Can be scipy, lineax or a specific lineax solver.

Type:

str, lineax-solver

conjugate_gradients

enables/diables the use of the Nonlinear Conjugate Gradients method. Can be False, Fletcher-Reeves, Hestenes-Stiefel, Dai-Yuan, Polak-Ribiere or average.

Type:

bool, str

r_local_method

chooses the method on the calculation of S_prime in local iterations. Can be projection or iteration.

Type:

str

r_global_method

chooses the method on the calculation of S_prime in global iterations. Can be projection or iteration.

Type:

str

r_gradient

if r_method=iteration, chooses the type of residual to optimize. Can be amplitude or intensity.

Type:

str

r_newton

enables/diables the use of the diagonal hessian if r_method=iteration

Type:

bool

r_no_iterations

the number of iterations if r_method=iteration

Type:

int

xi

a damping coefficient for adaptive step-sizes, avoids division by zero

Type:

float

local_adaptive_scaling

enables/disables adaptive step sized in local iterations. Can be one of False, pade_10 (linear), pade_20 (nonlinear), pade_11, pade_01 or pade_02

Type:

bool, str

global_adaptive_scaling

enables/disables adaptive step sized in global iterations. Can be one of False, pade_10 (linear), pade_20 (nonlinear), pade_11, pade_01 or pade_02

Type:

bool, str

momentum_is_being_used

what the name says

Type:

bool

normalize_after_every_step

what the name says, is usually True, but e.g. in attosecond streaking it needs to be False

Type:

bool

set_nonlinear_optimization(local_method=False, global_method=False, damping=0.001, memory=10, solver='lineax')[source]

A helper to set attributes. Sets all attributes which are related to (Quasi)-Newton methods. Will overwrite all attributes and may thus cause some unwanted changes.

Parameters:

  • local_method (bool, str) – can be one of False, “diagonal”, “full” or “lbfgs”

  • global_method (bool, str) – can be one of False, “diagonal”, “full” or “lbfgs”

  • damping (float) – a Levenberg-Marquardt style damping factor

  • memory (int) – the memory used in LBFGS

  • solver (str, lineax-solver) – the solver used if method==”full”, can be one of “lineax”, “scipy” or an instance of a lineax-solver

Returns:

self

set_linesearch(method=False, max_steps=10, c1=0.0001, c2=0.9, delta_gamma=0.5)[source]

A helper to set attributes. Sets all attributes which are related to linesearches. Will overwrite all attributes and may thus cause some unwanted changes.

Parameters:

  • method (bool, str) – the linesearch method, can be one of False, “backtracking” or “zoom”

  • max_steps (int) – the max number of steps per linesearch

  • c1 (float) – the constant of the armijo-condition

  • c2 (float) – the constant of the wolfe-condition

  • delta_gamma (float) – a factor that scales the stepsize, for backtracking delta_gamma<1, for zoom delta_gamma>1

Returns:

self

set_adaptive_stepsize(local_method=False, global_method=False, local_factor=-1, global_factor=-1, damping=1e-12)[source]

A helper to set attributes. Sets all attributes which are related to the adaptive stepsize. Will overwrite all attributes and may thus cause some unwanted changes.

Parameters:

  • local_method (bool, str) – has to be one of False, pade_10, pade_20, pade_11, pade_01, pade_02

  • global_method (bool, str) – has to be one of False, pade_10, pade_20, pade_11, pade_01, pade_02

  • local_factor (float) – sets the target error reduction

  • global_factor (float) – sets the target error reduction

  • damping (float) – a damping factor

Returns:

self

set_S_prime_params(local_method='projection', global_method='projection', gradient='intensity', newton=False, no_iterations=1)[source]

A helper to set attributes. Sets all attributes which are related to the updating of the signal field (signal_t -> signal_t_new). Will overwrite all attributes and may thus cause some unwanted changes.

(e.g. in the case of COPRA this may unintentionally overwrite the default of global_method=”iteration”)

Parameters:

  • local_method (str) – has to be projection or iteration

  • global_method (str) – has to be projection or iteration

  • gradient (str) – has to be intensity or amplitude, defines which error function is used in iteration

  • newton (bool) – if enabled a newton-style optimization is performed if method==iteration

  • weights (float, jnp.array) – optional weights applied to the error function

  • no_iterations (int) – the number of iterations if method==iteration

Returns:

self

create_initial_population(population_size=1, guess_type='random')[source]

Creates an initial population of pulses, parametrized as complex values on a grid.

Parameters:

  • population_size (int) – the number of guesses to be optimized

  • guess_type (str) – Has to be one of random, random_phase, constant or constant_phase.

Returns:

tuple[jnp.array, jnp.array, jnp.array or None, jnp.array or None], initial populations for the pulse and possibly the gate-pulse in time domain or frequency domain for ChirpScans

shuffle_data_along_m(descent_state, measurement_info, descent_info)[source]

Some solvers randomize local iterations. This is done through this method. It returns shuffled but consistent(!) data.

apply_spectrum(pulse, spectral_amplitude, eta)[source]
do_step_and_apply_spectral_amplitude(descent_state, measurement_info, descent_info, do_step)[source]

If a spectrum is provided this wraps around the step-method of all solvers and projects the current guess onto the measured spectrum.

update_velocity_map(signal, velocity_map, update_velocity_map, eta)[source]

Updates the velocity map for such that it can be applied in self.apply_momentum()

apply_momentum(signal, velocity_map, eta)[source]

Applies momentum to a signal.

do_step_and_apply_momentum(descent_state, measurement_info, descent_info, do_step)[source]

If momentum is being used this wraps around the step-method of all solvers and updates the current guess accordingly.

momentum(eta)[source]

Needs to be called if momentum is meant to be used in the reconstruction.

Parameters:

eta (float) – parameter that controls the momentum strength

Returns:

the class instance

calculate_error_population(population, measurement_info, descent_info)[source]

What the names says. Returns the errors and the population.

post_process_create_trace(descent_state, measurement_info, descent_info, idx)[source]

Post processing to get the final trace

class core.base_classes_algorithms.GeneralOptimizationBASE(*args, **kwargs)[source]

Bases: AlgorithmsBASE

The Base-Class for all general solvers. Inherits from AlgorithmsBASE.

error_metric

an arbitrary loss-function, needs to expect (trace, measured_trace) as input

Type:

Callable

create_initial_population(population_size, amp_type='gaussian', phase_type='polynomial', no_funcs_amp=5, no_funcs_phase=6)[source]

Creates an initial guess either explicit or parametrized. For bsplines one has to provide the B-Spline order k via bsplines_k

Parameters:

  • population_size (int) – the number of individuals

  • amp_type (str) – the representation of the spectral amplitude, can be one of gaussian, lorentzian, bsplines, continuous or an input for Classical Algorithms

  • phase_type (str) – the representation of the spectral phase, can be one of polynomial, sinusoidal, sigmoidal, bsplines, continuous or an input for Classical Algorithms

  • no_funcs_amp (int) – the number of basis functions for the spectral amplitude (only used if not on a grid)

  • no_funcs_phase (int) – the number of basis functions for the spectral phase (only used if not on grid)

Returns:

Pytree, the initial guess population

split_population_in_amp_and_phase(population)[source]

Splits a population into an amplitude and phase population.

merge_population_from_amp_and_phase(population_amp, population_phase)[source]

Undoes split_population_in_amp_and_phase()

polynomial_term(coefficient, order, x0, x)[source]
polynomial_phase(coefficients, central_f, frequency, measurement_info)[source]
sinusoidal_term(a, b, c, x)[source]
sinusoidal_phase(coefficients, central_f, frequency, measurement_info)[source]
discrete_phase(coefficients, central_f, frequency, measurement_info)[source]
tanh_term(c, k, x)[source]
tanh_phase(coefficients, central_f, frequency, measurement_info)[source]
bspline_phase(coefficients, central_f, frequency, measurement_info, Nx)[source]
gaussian_term(a, b, c, frequency)[source]
gaussian_amplitude(coefficients, frequency, measurement_info)[source]
lorentzian_term(a, b, c, frequency)[source]
lorentzian_amplitude(coefficients, frequency, measurement_info)[source]
discrete_amplitude(coefficients, frequency, measurement_info)[source]
bspline_amplitude(coefficients, frequency, measurement_info, Nx)[source]
get_phase(coefficients, central_f, frequency, measurement_info, descent_info, pulse_or_gate)[source]

Evaluates the spectral phase onto the frequency axis.

get_amplitude(coefficients, frequency, measurement_info, descent_info, pulse_or_gate)[source]

Evaluates the spectral amplitude onto the frequency axis.

get_central_f_for_current_guess(amp_f, frequency)[source]
make_pulse_f_from_individual(individual, measurement_info, descent_info, pulse_or_gate='pulse')[source]

Evaluates a parametrized individual onto the frequency axis.

get_pulses_f_from_population(population, measurement_info, descent_info)[source]

Evaluates a parametrized population onto the frequency axis.

construct_trace(individual, measurement_info, descent_info)[source]

Generates a trace for a given individual. Calls the method specific function for calculating the nonlinear signal fields.

trace_error(trace, measured_trace)[source]

The least squares error.

calculate_error_individual(individual, measurement_info, descent_info)[source]

Calculates the error of an individual based on its trace. Allows modification of the error-function via error_metric() and loss_function_modification().

calculate_error_population(population, measurement_info, descent_info)[source]

Calls jax.vmap over calculate_error_individual() for an entire population.

do_checks_before_running()[source]

Called by algorithm.run() performs a few checks

initialize_general_optimizer(population)[source]

A common initialization step for all general solvers.

post_process_create_trace(descent_state, measurement_info, descent_info, idx)[source]

Post processing to get the final trace

core.base_classes_methods module

class core.base_classes_methods.RetrievePulses(nonlinear_method, *args, cross_correlation=False, interferometric=False, key=None, seed=None, central_frequency=(None, None), **kwargs)[source]

Bases: object

The Base-Class for all reconstruction methods. Defines general initialization, preprocessing and postprocessing.

nonlinear_method

SHG, THG, PG or SD

Type:

str

f0

rarely some solvers need the central frequency to be zero. This saves the original central frequency.

Type:

float

doubleblind

whether the reconstruction is supposed to yield the gate in addition to the pulse.

Type:

bool

spectrum_is_being_used
Type:

bool

momentum_is_being_used
Type:

bool

measurement_info

a container of variable (but static) structure. Holds measurement data and parameters.

Type:

Pytree

descent_info

a container of variable (but static) structure. Holds parameters of the reconstruction algorithm.

Type:

Pytree

descent_state

a container of variable (but static) structure. Contains the current state of the solver.

Type:

Pytree

prng_seed

seed for the key

Type:

int

key

a jax.random.PRNGKey

Type:

jnp.array

factor

for SHG/THG the a correction factor of 2/3 needs to applied occasionally.

Type:

int

theta

an alias for the shifts/delays, internally indexed via m

Type:

jnp.array

time

the time axis, internally indexed via k

Type:

jnp.array

frequency

the frequency axis, internally indexed via n

Type:

jnp.array

measured_trace

2D-array with the measured data. axis=0 corresponds to shift/delay (index m), axis=1 correpsonds to the frequencies (index n)

Type:

jnp.array

update_PRNG_key(seed)[source]
get_data(theta, frequency, measured_trace)[source]

Prepare/Convert data.

get_spectral_amplitude(measured_frequency, measured_spectrum, pulse_or_gate, frequency=None, **kwargs)[source]

Used to provide a measured pulse spectrum. A spectrum for the gate pulse can also be provided.

get_gate_pulse(frequency, gate_f)[source]

For crosscorrelation=True the actual gate pulse has to be provided.

get_calibration_curve(frequency, calibration_curve)[source]
plot_results(final_result, exact_pulse=None)[source]
get_individual_from_idx(idx, population)[source]
get_idx_best_individual(population)[source]

Calculates the error for a population. Returns the index of the fittest individual. And the population

get_idx_worst_individual(population)[source]

Calculates the error for a population. Returns the index of the worst individual. And the population

get_idx_average_individual(population)[source]

Calculates the error for a population. Returns the index of an average individual. And the population

post_process_get_pulse_and_gate(descent_state, measurement_info, descent_info, idx=None)[source]

post processing to get the final pulse/gate

post_process_center_pulse_and_gate(pulses)[source]

This essentially removes the linear phase. But only approximately since no fits are done.

post_process(descent_state, error_arr)[source]

Creates the final_result object from the final descent_state.

class core.base_classes_methods.RetrievePulsesFROG(delay, frequency, measured_trace, nonlinear_method, cross_correlation=False, interferometric=False, **kwargs)[source]

Bases: RetrievePulses

The reconstruction class for FROG. Inherits from RetrievePulses.

  1. Trebino, “Frequency-Resolved Optical Gating: The Measurement of Ultrashort Laser Pulses”, 10.1007/978-1-4615-1181-6 (2000)

tau_arr

the delays

Type:

jnp.array

gate

the gate-pulse (if its known).

Type:

jnp.array

transform_arr

an alias for tau_arr

Type:

jnp.array

idtheta

an array with indices for tau_arr

Type:

jnp.array

dt
Type:

float

df
Type:

float

sk

correction values for FFT->DFT

Type:

jnp.array

rn

correction values for FFT->DFT

Type:

jnp.array

cross_correlation
Type:

bool

interferometric
Type:

bool

create_initial_population_doublepulse(population_size, **kwargs)[source]

Calls initial_guess_doublepulse.make_population_doublepulse to create an initial guess. The guess is in the time domain. Assumes an autocorrelation FROG.

Parameters:

  • population_size (int)

  • **kwargs – passed to make_population_doublepulse()

Returns:

Pytree, the initial population

shift_signal_in_time(signal, tau, frequency, sk, rn)[source]

The Fourier-Shift theorem.

calculate_shifted_signal(signal, frequency, tau_arr, time, in_axes=(None, 0, None, None, None))[source]

The Fourier-Shift theorem applied to a list of signals.

calculate_signal_t(individual, tau_arr, measurement_info)[source]

Calculates the signal field of a FROG in the time domain.

Parameters:

  • individual (Pytree) – a population containing only one member. (jax.vmap over whole population)

  • tau_arr (jnp.array) – the delays

  • measurement_info (Pytree) – contains the measurement parameters (e.g. nonlinear method, interferometric, … )

Returns:

Pytree, contains the signal field in the time domain as well as the fields used to calculate it.

class core.base_classes_methods.RetrievePulsesTDP(delay, frequency, measured_trace, nonlinear_method, spectral_filter=None, **kwargs)[source]

Bases: RetrievePulsesFROG

The reconstruction class for Time-Domain-Ptychography.

  1. Spangenberg et al., Phys. Rev. A 91, 021803(R), 10.1103/PhysRevA.91.021803 (2015)

spectral_filter

the spectral filter in the gate arm.

Type:

jnp.array

apply_spectral_filter(signal, spectral_filter, sk, rn)[source]

Apply a spectral filter to a signal.

calculate_signal_t(individual, tau_arr, measurement_info)[source]

Calculates the signal field of TDP in the time domain.

Parameters:

  • individual (Pytree) – a population containing only one member. (jax.vmap over whole population)

  • tau_arr (jnp.array) – the delays

  • measurement_info (Pytree) – contains the measurement parameters (e.g. nonlinear method, interferometric, … )

Returns:

Pytree, contains the signal field in the time domain as well as the fields used to calculate it.

class core.base_classes_methods.RetrievePulsesCHIRPSCAN(theta, frequency, measured_trace, nonlinear_method, phase_type=None, chirp_parameters=None, **kwargs)[source]

Bases: RetrievePulses

The reconstruction class for Chirp-Scan methods.

V. V. Lozovoy et al., Optics Letters 29, 775-777 (2004) M. Miranda et al., Opt. Express 20, 18732-18743 (2012)

theta

the shifts

Type:

jnp.array

dt
Type:

float

df
Type:

float

sk

correction values for FFT->DFT

Type:

jnp.array

rn

correction values for FFT->DFT

Type:

jnp.array

phase_matrix

a 2D-array with the phase values applied to pulse

Type:

jnp.array

parameters

parameters for the chirp function

Type:

tuple

transform_arr

an alias for phase_matrix

Type:

jnp.array

idtheta

indices for theta

Type:

jnp.array

get_phase_matrix(parameters)[source]

Calls phase_matrix_func in order to calculate the phase matrix.

get_dispersed_pulse_t(pulse_f, phase_matrix, sk, rn)[source]

Applies phase-matrix to a signal.

calculate_signal_t(individual, phase_matrix, measurement_info)[source]

Calculates the signal field of a Chirp-Scan in the time domain.

Parameters:

  • individual (Pytree) – a population containing only one member. (jax.vmap over whole population)

  • phase_matrix (jnp.array) – the applied phases

  • measurement_info (Pytree) – contains the measurement parameters (e.g. nonlinear method, … )

Returns:

Pytree, contains the signal field in the time domain as well as the fields used to calculate it.

class core.base_classes_methods.RetrievePulses2DSI(delay, frequency, measured_trace, nonlinear_method, cross_correlation=False, spectral_filter1=None, spectral_filter2=None, tau_pulse_anc1=0, **kwargs)[source]

Bases: RetrievePulsesFROG

The reconstruction class for 2DSI.

[1] J. R. Birge et al., Opt. Lett. 31, 2063-2065 (2006)

spectral_filter1

1st filter in interferometer

Type:

jnp.array

spectral_filter2

2nd filter in interferometer

Type:

jnp.array

tau_pulse_anc1

delay between 1st interferometer arm and exterior pulse

Type:

float

anc_frequency1

central frequency of pulse in 1st arm (calculated from max of spectral_filter1)

Type:

float

anc_frequency2

central frequency of pulse in 2nd arm (calculated from max of spectral_filter2)

Type:

float

c0

the speed of light

Type:

float

refractive_index

returns the refractive index for a material given a wavelength in um

Type:

refractiveindex.RefractiveIndexMaterial, Callable

phase_matrix

a 2D-array with phase values that could potentially have been applied to a pulse

Type:

jnp.array

apply_spectral_filter(signal, spectral_filter1, spectral_filter2, sk, rn)[source]

Apply a spectral filter to a signal.

calculate_signal_t(individual, tau_arr, measurement_info)[source]

Calculates the signal field of 2DSI in the time domain.

Parameters:

  • individual (Pytree) – a population containing only one member. (jax.vmap over whole population)

  • tau_arr (jnp.array) – the delays

  • measurement_info (Pytree) – contains the measurement parameters (e.g. nonlinear method, … )

Returns:

Pytree, contains the signal field in the time domain as well as the fields used to calculate it.

class core.base_classes_methods.RetrievePulsesVAMPIRE(delay, frequency, measured_trace, nonlinear_method, cross_correlation=False, tau_interferometer=0, material_thickness=0, refractive_index=<refractiveindex.refractiveindex.RefractiveIndexMaterial object>, **kwargs)[source]

Bases: RetrievePulsesFROG

The reconstruction class for VAMPIRE.

[1] B. Seifert and H. Stolz, Meas. Sci. Technol. 20 (2009) 015303 (7pp), 10.1088/0957-0233/20/1/015303 (2008)

tau_interferometer

delay of the interferometer arms

Type:

float

c0

the speed of light

Type:

float

refractive_index

returns the refractive index for a material given a wavelength in um

Type:

refractiveindex.RefractiveIndexMaterial, Callable

phase_matrix

a 2D-array with phase values that could potentially have been applied to a pulse

Type:

jnp.array

get_phase_matrix(refractive_index, material_thickness, measurement_info)[source]

Calculates the phase matrix that is applied of a pulse passes through a material.

apply_phase(pulse_t, measurement_info, sk, rn)[source]

For an VAMPIRE reconstruction one may need to consider effects of material dispersion in the interferometer. This applies a dispersion based on phase_matrix in order to achieve this.

calculate_signal_t(individual, tau_arr, measurement_info)[source]

Calculates the signal field in the time domain.

Parameters:

  • individual (Pytree) – a population containing only one member. (jax.vmap over whole population)

  • tau_arr (jnp.array) – the delays

  • measurement_info (Pytree) – contains the measurement parameters (e.g. nonlinear method, … )

Returns:

Pytree, contains the signal field in the time domain as well as the fields used to calculate it.

class core.base_classes_methods.RetrievePulsesSTREAKING(delay_fs, energy_eV, measured_trace, Ip_eV=Array([0], dtype=int32), retrieve_dtme=False, cross_correlation='doubleblind', interferometric=False, f_range_nir_pulse=(None, None), f_range_euv_pulse=(None, None), eV_range_dtme=(None, None), **kwargs)[source]

Bases: RetrievePulsesFROG

The reconstruction class for Attosecond Streaking. Implements the single- and multi-channel Strong-Field-Approximation. All axis (time, frequency, … ) are converted to atomic units for convenience and because that standard with the Strong-Field-Approximation.

[1] R. Kienberger, E. Goulielmakis, M. Uiberacker, et al., Atomic transient recorder. Nature 427, 817-821 (2004), 10.1038/nature02277 [2] P. D. Keathley et al., (2016), New J. Phys. 18 073009, 10.1088/1367-2630/18/7/073009

energy_au

the energy axis in atomic units

Type:

jnp.array

momentum_au

the photoelectron momenta in atomic units

Type:

jnp.array

position_au

the conjugate axis to momentum_au

Type:

jnp.array

sk_position_momentum

same function as sk, but for the position-momentum fourier pair

Type:

jnp.array

rn_position_momentum

same function as rn, but for the position-momentum fourier pair

Type:

jnp.array

Ip_au

the ionization potential in atomic units, the number of channels is infered from the number of Ips

Type:

jnp.array

retrieve_dtme

if true, the dipole-transition-matrix-elements will be retrieved

Type:

bool

dtme_momentum

the optionally provided dipole-transition-matrix-elements

Type:

None, jnp.array

make_axis(f_range_nir_pulse, f_range_euv_pulse, df_au)[source]
make_axis_dtme(eV_range_dtme, dE_eV)[source]
convert_time_fs_au(time, unit_in, unit_out)[source]
convert_frequency_PHz_au(frequency, unit_in, unit_out)[source]
convert_energy_eV_au(energy, unit_in, unit_out)[source]
get_DTME(momentum_au, dtme_momentum)[source]
get_spectral_amplitude(measured_frequency_PHz, measured_spectrum, pulse_or_gate, cutoff_frequency_PHz=0.1, **kwargs)[source]

Used to provide a measured pulse spectrum. A spectrum for the gate pulse can also be provided. The cutoff_frequency_PHz is used if pulse_or_gate=”pulse”, for the vectorpotential the spectrum is divided by the frequency. In order to avoid divergences the spectrum is set to zero below this cutoff.

make_volkov_phase_0(Ip, measurement_info)[source]
make_volkov_phase_1(pulse_t_nir_vectorpotential, measurement_info)[source]
make_dressed_DTME(dtme_position, pulse_t_nir_vectorpotential, measurement_info)[source]
make_streaking_amplitude(dtme_position, pulse_t_nir_vectorpotential, pulse_t_euv, measurement_info)[source]
calculate_signal_t(individual, tau_arr, measurement_info)[source]

Calculates the signal field in the time domain.

Parameters:

  • individual (Pytree) – a population containing only one member. (jax.vmap over whole population)

  • tau_arr (jnp.array) – the delays

  • measurement_info (Pytree) – contains the measurement parameters (e.g. nonlinear method, … )

Returns:

Pytree, contains the signal field in the time domain as well as the fields used to calculate it.

post_process_get_pulse_and_gate(descent_state, measurement_info, descent_info, idx=None)[source]

post processing to get the final pulse/gate

post_process_center_pulse_and_gate(pulses)[source]

This essentially removes the linear phase. But only approximately since no fits are done.

post_process(descent_state, error_arr)[source]

Creates the final_result object from the final descent_state.

plot_results(final_result)[source]

core.base_classic_algorithms module

core.base_classic_algorithms.normalize_population(population, measurement_info, descent_info, pulse_or_gate)[source]
core.base_classic_algorithms.initialize_CG_state(population)[source]
core.base_classic_algorithms.initialize_pseudo_newton_state(population)[source]
core.base_classic_algorithms.initialize_lbfgs_state(population, descent_info)[source]
core.base_classic_algorithms.update_lbfgs_state(lbfgs_state, gamma, grad, descent_direction)[source]
core.base_classic_algorithms.initialize_linesearch_info(optimizer)[source]
core.base_classic_algorithms.initialize_S_prime_params(optimizer)[source]
core.base_classic_algorithms.initialize_newton_info(optimizer)[source]
core.base_classic_algorithms.initialize_descent_info(optimizer)[source]
core.base_classic_algorithms.initialize_momentum(optimizer)[source]
class core.base_classic_algorithms.LSGPABASE(delay, frequency, measured_trace, nonlinear_method, cross_correlation=False, **kwargs)[source]

Bases: ClassicAlgorithmsBASE

The Least-Squares Generalized Projection Algorithm. Only available for delay based non-interferometric methods.

  1. Gagnon et al., Appl. Phys. B 92, 25-32, 10.1007/s00340-008-3063-x (2008)

update_pulse(signal_t_new, gate_shifted, measurement_info)[source]

Generates an new (maybe improoved) guess for the pulse.

update_gate(gate, signal_t, signal_t_new, measurement_info, descent_info)[source]

Generates an new (maybe improoved) guess for the gate. Uses an iterative formula. Thus not linear-least-squares anymore.

step(descent_state, measurement_info, descent_info)[source]

Performs one iteration of the Vanilla Algorithm.

Parameters:

  • descent_state – Pytree,

  • measurement_info – Pytree,

  • descent_info – Pytree,

Returns:

tuple[Pytree, jnp.array], the updated descent state and the current errors

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population – Pytree, the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable], the initial descent state and the step-function of the algorithm.

class core.base_classic_algorithms.CPCGPABASE(delay, frequency, trace, nonlinear_method, cross_correlation=False, constraints=False, svd=False, antialias=False, shift_trace_to_zero=False, **kwargs)[source]

Bases: ClassicAlgorithmsBASE

The Constrained-PCGP-Algorithm. Only available for delay based non-interferometric methods. Transforms population to time domain during initialization.

    1. Kane and A. B. Vakhtin, Prog. Quantum Electron. 81 (100364), 10.1016/j.pquantelec.2021.100364 (2022)

constraints

if true the operator based constraints are used.

Type:

bool

svd

if true a full SVD is performed instead of a single iteration of the power method

Type:

bool

antialias

if true anti-aliasing is applied to the outer-product-matrix-form

Type:

bool

get_spectral_amplitude(measured_frequency, measured_spectrum, pulse_or_gate)[source]

Used to provide a measured pulse spectrum. A spectrum for the gate pulse can also be provided.

calculate_opf(pulse_t, gate, pulse_t_prime, gate_prime, iteration, nonlinear_method, measurement_info)[source]

Calculates the opf given a pulse and gate.

calculate_signal_t_using_opf(individual, iteration, measurement_info, descent_info)[source]

Calculates signal_t for and individual via the opf.

do_anti_alias(opf, half_N)[source]

Performs anti-aliasing to the opf by setting a lower and upper an triangle to zero.

shift_rows(row, idx)[source]

Vectorized version of shift_rows. Takes similar arguments as shift_rows but with additional array axes over which shift_rows is mapped.

convert_opf_to_signal_t(opf, idx_arr)[source]

Transforms opf to signal field, by shifting along the time axis. Switching and flipping the two halfs around.

convert_signal_t_to_opf(signal_t, idx_arr)[source]

Converts a signal field into an opf by reversing the operations from convert_opf_to_signal_t().

decompose_opf(opf, pulse_t, gate, measurement_info, descent_info)[source]

Decomposes the opf into its dominant components via an SVD or the Power-Method.

impose_constraints(pulse_t, gate, opf, measurement_info)[source]

Applies additional constraints according to the operator formalism of PCGP.

update_individual(opf, individual, measurement_info, descent_info)[source]

Updates and individual using an updated opf.

step(descent_state, measurement_info, descent_info)[source]

Performs one iteration of the C-PCGP Algorithm.

Parameters:

  • descent_state – Pytree,

  • measurement_info – Pytree,

  • descent_info – Pytree,

Returns:

tuple[Pytree, jnp.array], the updated descent state and the current errors

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population – Pytree, the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable], the initial descent state and the step-function of the algorithm.

post_process_create_trace(descent_state, measurement_info, descent_info, idx)[source]

For PCGP the trace is constructed using the opf.

do_step_and_apply_spectral_amplitude(descent_state, measurement_info, descent_info, do_step)[source]

If a spectrum is provided this wraps around the step-method of all solvers and projects the current guess onto the measured spectrum.

class core.base_classic_algorithms.GeneralizedProjectionBASE(*args, **kwargs)[source]

Bases: ClassicAlgorithmsBASE

Implements the Generalized Projection Algorithm.

    1. DeLong et al., Opt. Lett. 19, 2152-2154 (1994)

no_steps_descent

the numer of descent steps per iteration

Type:

int

optimize_spectral_phase_directly
Type:

bool

update_individual(individual, gamma, descent_direction, pulse_or_gate)[source]
update_population(population, gamma, descent_direction, pulse_or_gate)[source]

Applies the descent based update to the population.

get_Z_gradient(signal_t, signal_t_new, transform_arr, measurement_info, descent_info, pulse_or_gate)[source]

Calculates the Z-error gradient for the entire population.

calc_Z_error_for_linesearch(gamma, linesearch_info, measurement_info, pulse_or_gate)[source]

Calculates the Z-error such that it can be called in a linesearch.

calc_Z_grad_for_linesearch(gamma, linesearch_info, measurement_info, descent_info, pulse_or_gate)[source]

Calculates the Z-error gradient such that it can be called in a linesearch.

descent_Z_error_step(signal_t, signal_t_new, Z_error, descent_state, measurement_info, descent_info, pulse_or_gate)[source]

Performs a descent step in order to minimize the Z-error. Employs gradient descent, nonlinear conjugate gradients, LBFGS or damped Newtons method (diagonal or full). The step size is determined via a fixed/adaptive step size, a backtracking or a zoom linesearch.

do_descent_Z_error_step(descent_state, signal_t_new, measurement_info, descent_info)[source]

Does one Z-error descent step. Calls descent_Z_error_step for pulse and or gate.

do_descent_Z_error(descent_state, signal_t_new, measurement_info, descent_info)[source]

Performs a descent based optimization to find the pulse/gate that are able to produce S_prime.

step(descent_state, measurement_info, descent_info)[source]

Performs one iteration of the Generalized Projection Algorithm.

Parameters:

  • descent_state (Pytree)

  • measurement_info (Pytree)

  • descent_info (Pytree)

Returns:

tuple[Pytree, jnp.array], the updated descent state and the current trace errors of the population.

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population (Pytree) – the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable], the initial descent state and the step-function of the algorithm.

class core.base_classic_algorithms.PtychographicIterativeEngineBASE(*args, pie_method='rPIE', **kwargs)[source]

Bases: ClassicAlgorithmsBASE

Implements a version of the Ptychographic Iterative Engine (PIE).

A. Maiden et al., Optica 4, 736-745 (2017) T. Schweizer, “Time-Domain Ptychography and its Applications in Ultrafast Science”, PhD Thesis, Bern (2021)

alpha

a regularization parameter

Type:

float

pie_method

specifies the PIE variant. Can be one of None, PIE, ePIE, rPIE. Where None indicates that the pure gradient is used.

Type:

None, str

optimize_spectral_phase_directly
Type:

bool

calculate_PIE_mu(signal_f_abs, trace_root, measurement_info, descent_info, local_or_global)[source]
calculate_PIE_trace(signal_f, measured_trace, measurement_info, descent_info, local_or_global)[source]
calculate_PIE_error(mu, signal_f, measured_trace)[source]

Calculates the normalized least-squares error using the amplitude residuals.

get_PIE_weights(probe, alpha, pie_method)[source]

Calculates the weight-functions for the differen PIE-version.

update_individual(individual, gamma, descent_direction, pulse_or_gate)[source]
update_population(population, gamma, descent_direction, pulse_or_gate)[source]

Applies the PIE update to the population.

calculate_PIE_descent_direction(signal_t, signal_t_new, transform_arr, pie_method, measurement_info, descent_info, pulse_or_gate)[source]

Calculates the descent direction based on the PIE version.

calc_error_for_linesearch(gamma, linesearch_info, measurement_info, pulse_or_gate)[source]

Calculates the PIE-error such that it can be called in a linesearch.

calc_grad_for_linesearch(gamma, linesearch_info, measurement_info, descent_info, pulse_or_gate, local_or_global)[source]

Calculates the PIE direction such that it can be called in a linesearch.

do_iteration(signal_t, signal_t_new, transform_arr, measured_trace, population, local_or_global_state, measurement_info, descent_info, pulse_or_gate, local_or_global)[source]

Performs one local/global iteration of the PIE. On top of the different PIE-version nonlinear conjugate gradients, LBFGS or damped Newtons method (diagonal or full) may be used. The step size is determined via a fixed/adaptive step size, a backtracking or a zoom linesearch.

Newtons method with a full newton is not available for the reconstruction of the gate.

local_iteration(descent_state, transform_arr_m, trace_line, measurement_info, descent_info)[source]

Peforms one local iteration. Calls do_iteration() with the appropriate (randomized) signal fields.

local_step(descent_state, measurement_info, descent_info)[source]

Performs one local iteration of the PIE. This means the method loops over the randomized measurement data once and updates the population using each data point individually.

Parameters:

  • descent_state (Pytree)

  • measurement_info (Pytree)

  • descent_info (Pytree)

Returns:

tuple[Pytree, jnp.array], the updated descent state and the current trace errors of the population.

global_step(descent_state, measurement_info, descent_info)[source]

Performs one global iteration of the PIE. This means the method updates the population once using all measured data at once.

Parameters:

  • descent_state (Pytree)

  • measurement_info (Pytree)

  • descent_info (Pytree)

Returns:

tuple[Pytree, jnp.array], the updated descent state and the current trace errors of the population.

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population (Pytree) – the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable, Callable], the initial descent state, the local and global step-functions of the algorithm.

run(population, no_iterations_local, no_iterations_global, **kwargs)[source]

The Algorithm can use a local and a global sequentially.

Parameters:

  • population (Pytree) – the initial guess

  • no_iterations_local – int, the number of local iterations. Accepts zero as a value.

  • no_iterations_global – int, the number of globale iterations. Accepts zero as a value.

Returns:

Pytree, the final result

class core.base_classic_algorithms.COPRABASE(*args, **kwargs)[source]

Bases: ClassicAlgorithmsBASE

Implements a version of the Common Pulse Retrieval Algorithm (COPRA).

    1. Geib et al., Optica 6, 495-505 (2019)

update_individual(individual, gamma, descent_direction, pulse_or_gate)[source]

Updates an individual based on a descent direction and a step size.

update_population(population, gamma, descent_direction, pulse_or_gate)[source]

Applies the a descent based update to the population.

get_Z_gradient(signal_t, signal_t_new, transform_arr, measurement_info, descent_info, pulse_or_gate)[source]

Calculates the Z-error gradient for the current population.

calc_Z_error_for_linesearch(gamma, linesearch_info, measurement_info, descent_info, pulse_or_gate)[source]

Calculates the Z-error such that it can be called in a linesearch.

calc_Z_grad_for_linesearch(gamma, linesearch_info, measurement_info, descent_info, pulse_or_gate)[source]

Calculates the Z-error gradient such that it can be called in a linesearch.

do_iteration(signal_t, signal_t_new, transform_arr, population, local_or_global_state, measurement_info, descent_info, pulse_or_gate, local_or_global)[source]

Performs one local/global iteration of the Common Pulse Retrieval Algorithm. Uses gradient descent, nonlinear conjugate gradients, LBFGS or damped Newtons method (diagonal or full). The step size is determined via a fixed/adaptive step size, a backtracking or a zoom linesearch.

local_iteration(descent_state, transform_arr_m, trace_line, measurement_info, descent_info)[source]

Peforms one local iteration. Calls do_iteration() with the appropriate (randomized) signal fields.

local_step(descent_state, measurement_info, descent_info)[source]

Performs one local iteration of the Common Pulse Retrieval Algorithm. This means the method loops over the randomized measurement data once and updates the population using each data point individually.

Parameters:

  • descent_state (Pytree)

  • measurement_info (Pytree)

  • descent_info (Pytree)

Returns:

tuple[Pytree, jnp.array], the updated descent state and the current trace errors of the population.

global_step(descent_state, measurement_info, descent_info)[source]

Performs one global iteration of the Common Pulse Retrieval Algorithm. This means the method updates the population once using all measured data at once.

Parameters:

  • descent_state (Pytree)

  • measurement_info (Pytree)

  • descent_info (Pytree)

Returns:

tuple[Pytree, jnp.array], the updated descent state and the current trace errors of the population.

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population (Pytree) – the initial guess as created by create_initial_population()

Returns:

tuple[Pytree, Callable, Callable], the initial descent state, the local and global step-functions of the algorithm.

run(population, no_iterations_local, no_iterations_global, **kwargs)[source]

The Algorithm can use a local and a global sequentially.

Parameters:

  • population (Pytree) – the initial guess

  • no_iterations_local – int, the number of local iterations. Accepts zero as a value.

  • no_iterations_global – int, the number of globale iterations. Accepts zero as a value.

Returns:

Pytree, the final result

class core.base_classic_algorithms.LSFBASE(*args, **kwargs)[source]

Bases: ClassicAlgorithmsBASE

Implements a version of the Linesearch FROG Algorithm (LSF). Despite its name the algorithm is NOT restricted to FROG.

    1. Krook and V. Pasiskevicius, Opt. Express 33, 33258-33269 (2025)

no_sections

number of sections to split the search line into

Type:

int

number_of_disection_iterations

as the name says

Type:

int

direction_mode

can be random or continuous

Type:

str

ratio_points_for_continuous

smaller value means more randomness/eratic

Type:

int

only_allow_improvements

if true, only steps that decrease the error will be accepted

Type:

bool

apply_spectrum(pulse_f, spectral_amplitude, eta)[source]

Projects onto measured spectrum by dividing the current spectral amplitude out.

get_random_values(key, shape, minval, maxval, descent_info)[source]

LSF requires random directions. These are produced here.

get_search_direction_individual(keys, individual, measurement_info, descent_info)[source]

Creates a pytree with random search directions for one individual.

get_search_direction(key, population, measurement_info, descent_info)[source]

Creates a pytree with random search directions for an entire population.

get_scalars(direction, signal, descent_info)[source]

Calculates scalars to identify the min/max of a search direction.

select_new_brackets_N_section(E_arr, error_arr, descent_state, pulse_or_gate)[source]

Selects the two trial-individuals with the lowest errors in order to narrow the search bracket.

select_new_brackets_bisection(E_arr, error_arr, descent_state, pulse_or_gate)[source]

Narrows the search bracket by removing the trial-individual with the largest error or the one that is the furthest from the trial guess.

generate_E_arr(El, Er, n)[source]

Takes in a range El, Er and disects it into n equal sections.

bisection_step(El, Er, descent_state, measurement_info, descent_info, pulse_or_gate)[source]

Does one bisection step of the LSF algorithm.

Performs one bisection search to find the minimum along a given search direction. The number of iterations is set through self.number_of_bisection_iterations

search_along_direction(direction, descent_state, measurement_info, descent_info)[source]

Performs a bisection search along one direction for pulse and the for gate.

Combines the trial populations with the current population such that the errors may be computed.

calculate_error(E_arr, descent_state, measurement_info, descent_info, pulse_or_gate)[source]

Calculates the trace error. Since pulse and gate are optimized independently the population as provided to calculate_error_individual() needs to be constructed from the current population and the fields in the cureent optimization.

step(descent_state, measurement_info, descent_info)[source]

Performs one step of the LSF-algorithm.

Parameters:

  • descent_state – Pytree,

  • measurement_info – Pytree,

  • descent_info – Pytree,

Returns:

tuple[Pytree, jnp.array], the updated descent state, errors of the current population

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population – Pytree, the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable], the initial descent state, the step-function of the algorithm.

core.base_general_optimization module

core.base_general_optimization.loss_function_for_general(trace, measured_trace, exponent_trace=1, L_norm=1)[source]

An example of a modified loss function for general algorithms.

Inspired by ELFPIE: S. Zhang, T.T.J.M. Berendschot and J. Zhou, Signal Processing 210 109088, 10.1016/j.sigpro.2023.109088, (2023)

Parameters:

  • trace (jnp.array) – the current trace

  • measured_trace (jnp.array) – the measured trace

  • kwargs – exponent_trace (int, float): the exponent applied to trace before building the residual L_norm (int, float): the norm exponent to be used

core.base_general_optimization.make_key_tree(key, pytree)[source]

For a given pytree, each leaf is replaced by a prng-key.

core.base_general_optimization.shuffle_pytree(key, pytree, axis)[source]

Randomizes the leafs of a given pytree along a given axis.

core.base_general_optimization.merge_pytrees(pytree1, pytree2, bool_tree)[source]

bool_tree selects pytree1 if true and pytree2 if false.

class core.base_general_optimization.DifferentialEvolutionBASE(*args, strategy='best1_bin', selection_mechanism='greedy', mutation_rate=0.5, crossover_rate=0.7, **kwargs)[source]

Bases: GeneralOptimizationBASE

Implements a Differential-Evolution Algorithm. Based on Qiang, J., Mitchell, C., A Unified Differential Evolution Algorithm for Global Optimization, 2014, https://www.osti.gov/servlets/purl/1163659

strategy

the mutation and selection strategy, analogous to scipy’s differential evolution.

Type:

str

mutation_rate

the mutation rate

Type:

float

crossover_rate

the crossover rate

Type:

float

selection_mechanism

the selection mechanism, can be greedy or global, defined in select_population().

Type:

str

temperature

a temperature value for the global selection mechanism

Type:

float

best1(F, best_individual, population, key)[source]
best2(F, best_individual, population, key)[source]
rand1(F, best_individual, population, key)[source]
rand2(F, best_individual, population, key)[source]
randtobest1(F, best_individual, population, key)[source]
randtobest2(F, best_individual, population, key)[source]
currenttorand1(F, best_individual, population, key)[source]
currenttorand2(F, best_individual, population, key)[source]
currenttobest1(F, best_individual, population, key)[source]
currenttobest2(F, best_individual, population, key)[source]
make_bin_mask_tree(key, pytree, p)[source]

For a given pytree a random binary mask is generated on each leaf, via the probabilty p.

make_exp_mask_tree(key, pytree, p)[source]

For a given pytree a mask (e.g. of the form 111100000) is generated via the probability p on each row of each leaf.

bin_crossover(CR, parent_population, mutant_population, key)[source]

Peforms a binary crossover between two populations via the crossover rate CR.

exp_crossover(CR, parent_population, mutant_population, key)[source]

Peforms an exponential crossover between two populations via the crossover rate CR.

smooth_crossover(CR, parent_population, mutant_population, key)[source]

Similar to an exponential crossover but with a smooth nonbinary transition, which is fascilitated via a tanh function.

custom_mutation(F, best_individual, population, key)[source]

A placeholder to allow the introduction of custom mutation approaches.

custom_crossover(CR, parent_population, mutant_population, key)[source]

A placeholder to allow the introduction of custom crossover approaches.

do_mutation(mutation_strategy, F, best_individual, population, key)[source]

Creates and applies random mutations to a population.

do_crossover(crossover_strategy, CR, parent_population, mutant_population, key)[source]

Creates and applies random crossover events between two populations.

select_population(key, selection_mechanism, error_parent, error_trial, population_parent, population_trial, mu_parent, mu_trial, descent_info)[source]

Performs the evolutionary selection process for two populations. Currently selection_mechanism can be greedy or global.

greedy: implements a pairwise comparison between individuals of the two population, where always the individual with the higher fitness is selected. global: implements a comparison/ranking between all individuals of both populations. The next generation is selected via randomized sampling based on a Fermi-Distribution. A “temperature” allows tuning of this selection process.

step(descent_state, measurement_info, descent_info)[source]

Performs one step/generation of the Differential Evolution Algorithm.

Parameters:

  • descent_state – Pytree,

  • measurement_info – Pytree,

  • descent_info – Pytree,

Returns:

tuple[Pytree, jnp.array], the updated descent state, the mean, minimum and maximum error of the current population

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population – Pytree, the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable], the initial descent state, the step-function of the algorithm.

class core.base_general_optimization.EvosaxBASE(*args, solver=None, **kwargs)[source]

Bases: GeneralOptimizationBASE

Employs the evosax package to perform the optimization.

Robert Tjarko Lange, evosax: JAX-based Evolution Strategies, arXiv preprint arXiv:2212.04180 (2022)

solver

any evosax-solver should work

Type:

evosax-solver

solver_params

user defined parameters for the evosax-solver, if None the default params set in evosax are used

Type:

any

solver_kwargs

some evosax-solver require additional input arguments. These can be supplied via this aattribute.

Type:

dict

step_amp_or_phase(population_amp_or_phase, descent_state, measurement_info, descent_info, amp_or_phase)[source]

Evosax does not respect the structure of a pytree. Thus when optimizing parameters with vastly different properties simultaneously one should expose them seperately to evosax in order to avoid mixing. Thus the step function is applied once to update the amplitudes and once to update the phases.

step(descent_state, measurement_info, descent_info)[source]

Performs one optimization step through an evolutionary algorithm in evosax.

Parameters:

  • descent_state – Pytree,

  • measurement_info – Pytree,

  • descent_info – Pytree,

Returns:

tuple[Pytree, jnp.array], the updated descent state, the mean, minimum and maximum error of the current population

initialize_evosax_solver(key, solver, population, individual, amp_or_phase)[source]

Initializes the provided evosax solver. The solvers need to be exposed to the entire population, one individual or the initial fitnesses. Since the amplitude and phase are optimized separetly a solver needs to be initialized once for each. (This done by calling this method twice.)

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Initializes the evosax-solver appropriately. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population – Pytree, the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable], the initial descent state, the step-function of the algorithm.

class core.base_general_optimization.AutoDiffBASE(*args, solver=None, **kwargs)[source]

Bases: GeneralOptimizationBASE

Employs the optimistix package to perform the optimization via Automatic-Differentiation.

J. Rader, T. Lyons and P.Kidger, Optimistix: modular optimisation in JAX and Equinox, arXiv:2402.09983 (2024) DeepMind et al., The DeepMind JAX Ecosystem, http://github.com/google-deepmind (2020)

solver

solvers need to be initialized

Type:

optimistix-solver, optax-solver

alternating_optimization

if true, the optimizer alternates between amplitude and phase

Type:

bool

optimize_group_delay

if true, the group delay will be optimized instead of the spectral phase

Type:

bool

get_phase(coefficients, central_f, frequency, measurement_info, descent_info, pulse_or_gate)[source]

Wraps around GeneralOptimizationBASE.get_phase() in order to fascilliate the optimization of the group delay.

loss_function(individual, measurement_info, descent_info)[source]

Wraps around self.calculate_error_individual() to return the error of the current guess.

loss_function_amp(individual_amp, individual_phase, measurement_info, descent_info)[source]

Helper to fascilliate optimization with respect to the amplitude only.

loss_function_phase(individual_phase, individual_amp, measurement_info, descent_info)[source]

Helper to fascilliate optimization with respect to the phase only.

step(descent_state, measurement_info, descent_info)[source]

Performs one optimization step.

Parameters:

  • descent_state – Pytree,

  • measurement_info – Pytree,

  • descent_info – Pytree,

Returns:

tuple[Pytree, jnp.array], the updated descent state, errors of the current population

initialize_alternating_optimization(descent_state, population, solver, optimistix_args, measurement_info, descent_info)[source]

If the amplitude and phase are supposed to be optimized in an alternating fashion instead of simultaneously, different methods have to be initialized.

initialize_optimistix(descent_state, solver, measurement_info, descent_info)[source]

Initializes the optimistix-solver.

initialize_run(population)[source]

Prepares all provided data and parameters for the reconstruction. Here the final shape/structure of descent_state, measurement_info and descent_info are determined.

Parameters:

population – Pytree, the initial guess as created by self.create_initial_population()

Returns:

tuple[Pytree, Callable], the initial descent state, the step-function of the algorithm.

core.bsplines_1d module

core.bsplines_1d.get_C_in(i, n)[source]

Naive implementation of n choose i.

core.bsplines_1d.get_m_ij(i, j, k)[source]

Calculates matrix elements (i,j) for a uniform bspline weight matrix of order k.

core.bsplines_1d.get_M(k)[source]

Calculates the weight matrix for uniform bsplines of order k.

core.bsplines_1d.get_prefactor(k)[source]

Calculates the prefactor for a weight matrix for uniform bsplines of order k.

core.bsplines_1d.make_bsplines(cpoints, Nx, k, M, f)[source]

Evaluate arbitrary order bsplines in 1D.

core.construct_s_prime module

core.construct_s_prime.calculate_S_prime_projection(signal_t, measured_trace, mu, sk, rn)[source]

Calculates signal_t_new/S_prime via a projection onto the measured intensity.

Parameters:

  • signal_t (pytree) – contains the complex signal field of the current guess

  • measured_trace (jnp.array) – the measured intensity

  • mu (float) – the scaling factor between the measured intensity and the intensity of the current guess

  • measurement_info (Pytree) – contains measurement data and information

Returns:

jnp.array, the complex signal field in the time domain projected onto the measured intensity

core.construct_s_prime.calculate_r_newton_diagonal_intensity(trace, measured_trace)[source]
core.construct_s_prime.calculate_r_newton_diagonal_amplitude(trace, measured_trace)[source]
core.construct_s_prime.calculate_r_newton_diagonal(signal_f, measured_trace, sk, rn, descent_info)[source]
core.construct_s_prime.calculate_r_gradient_intensity(mu, signal_f, measured_trace, sk, rn)[source]
core.construct_s_prime.calculate_r_gradient_amplitude(mu, signal_f, measured_trace, sk, rn)[source]
core.construct_s_prime.calculate_r_gradient(mu, signal_f, measured_trace, sk, rn, descent_info)[source]
core.construct_s_prime.calculate_r_descent_direction(signal_f, mu, measured_trace, sk, rn, descent_info)[source]

Calculates descent direction of the iterative calculation of signal_t_new/S_prime. Uses either gradient descent or newtons method with the diagonal approximation. The error-functions can be based on intensity or amplitude based residuals.

core.construct_s_prime.calculate_r_error_intensity(mu, trace, measured_trace)[source]
core.construct_s_prime.calculate_r_error_amplitude(mu, trace, measured_trace)[source]
core.construct_s_prime.calculate_r_error(trace, measured_trace, mu, descent_info)[source]
core.construct_s_prime.calculate_S_prime_iterative_step(signal_t, measured_trace, mu, sk, rn, descent_info, local_or_global)[source]

One iteration of the iterative descent based calculation of signal_t_new/S_prime.

core.construct_s_prime.calculate_S_prime_iterative(signal_t, measured_trace, mu, sk, rn, descent_info, local_or_global)[source]

Calculates signal_t_new/S_prime via an iterative optimization of the least-squares error.

Parameters:

  • signal_t (jnp.array) – the complex signal field in the time domain of the current guess

  • signal_t – the complex signal field in the frequency domain of the current guess

  • measured_trace (jnp.array) – the measured intensity

  • mu (float) – the scaling factor between the measured intensity and the intensity of the current guess

  • measurement_info (Pytree) – contains measurement data and information

  • descent_info (Pytree) – contains information on the behaviour of the solver

  • local_or_global (str) – whether this is used in a local or global iteration

Returns:

jnp.array, the complex signal field in the time domain projected onto the measured intensity

core.construct_s_prime.calculate_S_prime(signal_t, measured_trace, mu, measurement_info, descent_info, local_or_global)[source]

Calculates signal_t_new/S_prime via projection or iterative optimization

Parameters:

  • signal_t (jnp.array) – the complex signal field in the time domain of the current guess

  • signal_f (jnp.array) – the complex signal field in the frequency domain of the current guess

  • measured_trace (jnp.array) – the measured intensity

  • mu (float) – the scaling factor between the measured intensity and the intensity of the current guess

  • measurement_info (Pytree) – contains measurement data and information

  • descent_info (Pytree) – contains information on the behaviour of the solver

  • local_or_global (str) – whether this is used in a local or global iteration

Returns:

jnp.array, the complex signal field in the time domain projected onto the measured intensity

core.create_population module

core.create_population.get_initial_amp(measurement_info, pulse_or_gate)[source]

Estimate spectral amplitude from the integrated measured intensity. For SHG/THG the amplitude is interpolated to lie in the correct frequency region.

core.create_population.random(key, shape)[source]
core.create_population.random_phase(key, shape, amp)[source]
core.create_population.constant(key, shape)[source]
core.create_population.constant_phase(key, shape, amp)[source]
core.create_population.create_population_classic(key, shape, guess_type, measurement_info, pulse_or_gate)[source]

Creates a stack of initial guesses with the shape (population_size, jnp.size(frequency)). The guesses are all in the frequency domain. The created populations can be optimized by both general and classical solvers.

Parameters:

  • key (jnp.array) – jax.random.PRNGKey

  • shape (tuple) – the number and size of individuals

  • guess_type (str) – the guess mode. Has to be one of random, random_phase, constant or constant_phase. doublepulse is moved to initial_guess_doublepulse.py

  • measurement_info (Pytree) – holds the measurement information, is filled during initialization of each solver

Returns:

jnp.array, stack of 1D-arrays

core.create_population.polynomial_guess(key, population, shape, measurement_info)[source]
core.create_population.sinusoidal_guess(key, population, shape, measurement_info)[source]
core.create_population.sigmoidal_guess(key, population, shape, measurement_info)[source]
core.create_population.bspline_guess_phase(key, population, shape, measurement_info)[source]
core.create_population.continuous_discrete_guess_phase(key, population, shape, measurement_info, pulse_or_gate)[source]
core.create_population.gaussian_or_lorentzian_guess(key, population, shape, measurement_info)[source]
core.create_population.continuous_discrete_guess_amp(key, population, shape, measurement_info, pulse_or_gate)[source]
core.create_population.bspline_guess_amp(key, population, shape, measurement_info)[source]
core.create_population.general_phase(key, population, shape, measurement_info, phase_type, pulse_or_gate)[source]
core.create_population.general_amp(key, population, shape, measurement_info, amp_type, pulse_or_gate)[source]
core.create_population.create_phase(key, phase_type, population, shape, measurement_info, pulse_or_gate)[source]
core.create_population.create_amp(key, amp_type, population, shape, measurement_info, pulse_or_gate)[source]
core.create_population.create_population_general(key, amp_type, phase_type, population, population_size, no_funcs_amp, no_funcs_phase, spectrum_provided, measurement_info, pulse_or_gate)[source]

Creates an initial guess population for general solvers. Since general solvers do not require a grid, different representations for amplitude and phase can be used. The population is represented in the freqeuncy domain.

Parameters:

  • key (jnp.array) – jax.random.PRNGKey

  • amp_type (str) – representation of the amplitude

  • phase_type (str) – representation of the phase

  • population (Pytree) – a MyNamespace object containing amp and phase

  • population_size (int) – the number of individuals to optimize

  • no_funcs_amp (int) – some representations can consist of multiple basis functions

  • no_funcs_phase (int) – some representations can consist of multiple basis functions

  • spectrum_provided (bool) – if a spectrum is provided then the guessed population will not include an amplitude

  • measurement_info (Pytree) – holds the measurement information, is filled during initialization of each solver

  • pulse_or_gate (str) – only needed in streaking, there the scale of the nir/ir-vectorpotential needs to be always optimizable

Returns:

tuple[jnp.array, Pytree]

core.initial_guess_doublepulse module

core.initial_guess_doublepulse.gaussian(x, a, b, c)[source]

Gaussian for fit.

core.initial_guess_doublepulse.get_double_pulse_amps(trace, sigma=10, init_std=0.001)[source]

Estimate temporal amplitude based on Trace. Takes average along frequency axis. Fits gaussians and constructs guess from fit result.

core.initial_guess_doublepulse.get_flat_spectral_phase(time, freq, trace, nonlinear_method, monochromatic=True, sigma=10)[source]

Construct linear spectral phase in time domain.

core.initial_guess_doublepulse.get_double_pulse_initial_guess(tau_arr, frequency, measured_trace, nonlinear_method, monochromatic_double_pulse=True, sigma=3, init_std=0.001, i_want_control=False)[source]

Finds an initial guess for an autocorrelation FROG trace for doublepulses. Works by finding the delays, amplitudes and central frequencies through scipy’s peak finder and gaussian fits. Does use or work with jax.

Parameters:

  • tau_arr (jnp.array) – the delay axis 1D-array

  • frequency (jnp.array) – the frequency axis 1D-array

  • measured_trace (jnp.array) – the measured trace 2D-array

  • nonlinear_method (str) – the nonlinear method

  • monochromatic_double_pulse (bool) – whether the doublepulses had different central frequencies

  • sigma (float) – a smoothing parameter

  • init_std (float) – initial value for width of gaussian fits

  • i_want_control (bool) – this method does not really work and is disabled for SHG/THG, can be overwritten by this

Returns:

tuple[jnp.array, jnp.array], guess in time and frequency domain

core.initial_guess_doublepulse.make_population_doublepulse(key, population_size, tau_arr, frequency, measured_trace, nonlinear_method, monochromatic_double_pulse=True, sigma=3, init_std=0.001, i_want_control=False)[source]

Create an initial guess population for a doublepulse using get_double_pulse_initial_guess, with its restrictions. The population is constructed from a deterministic doublepulse guess by adding noise.

Parameters:

  • key (jnp.array) – a jax.random.PRNGKey

  • population_size (int) – the number of initial guesses

  • tau_arr (jnp.array) – the delay axis 1D-array

  • frequency (jnp.array) – the frequency axis 1D-array

  • measured_trace (jnp.array) – the measured trace 2D-array

  • nonlinear_method (str) – the nonlinear method

  • monochromatic_double_pulse (bool) – whether the doublepulses had different central frequencies

  • sigma (float) – a smoothing parameter

  • init_std (float) – initial value for width of gaussian fits

  • i_want_control (bool) – this method does not really work and is disabled for SHG/THG, can be overwritten by this

Returns:

jnp.array

core.lbfgs module

core.lbfgs.backward_loop(q, i, rho, s, y)[source]

Constructing the LBFGS direction without materializing the inverse hessian is done through nested vector-multiplications.

core.lbfgs.forward_loop(r, i, alpha, rho, s, y)[source]

Constructing the LBFGS direction without materializing the inverse hessian is done through nested vector-multiplications.

core.lbfgs.calculate_quasi_newton_direction(grad_current, grad_prev, rho, s, y, newton_info)[source]

Does the actual LBFGS calculation.

core.lbfgs.do_lbfgs(grad_current, lbfgs_state, descent_info)[source]

Prepares and calls the LBFGS calculation.

core.lbfgs.get_quasi_newton_direction(grad, lbfgs_state, descent_info)[source]

Calculate the quasi-newton direciton using LBFGS.

Parameters:

  • grad (jnp.array) – the current gradient

  • lbfgs_stat (Pytree) – the current lbfgs state

  • descent_info (Pytree) – holds information on the solver (e.g. memory size for LBFGS)

Returns:

tuple[jnp.array, Pytree], the quasi-newton direction and the unchanged lbfgs_state

core.nonlinear_cg module

core.nonlinear_cg.get_nonlinear_CG_direction(descent_direction, cg, beta_parameter_version)[source]

Calculates the descent direction using Nonlinear Conjugate Gradients.

Parameters:

  • descent_direction (jnp.array) – the descent direction from a descent solver

  • cg (Pytree) – the current conjugate gradient state, holds the previous descent direction etc.

  • beta_parameter_version (str) – the concrete NCG method to use

Returns:

tuple[jnp.array, Pytree], the descent direction based on NCG, the updated NCG state

core.nonlinear_cg.get_beta(grad, grad_prev, descent_direction_prev, beta_parameter_version)[source]

Calls the NCG-parameter functions.

core.nonlinear_cg.beta_fletcher_reeves(grad, grad_prev, descent_direction_prev)[source]

Compute beta based on Fletcher-Reeves.

core.nonlinear_cg.beta_polak_ribiere(grad, grad_prev, descent_direction_prev)[source]

Compute beta based one Polak-Ribiere.

core.nonlinear_cg.beta_hestenes_stiefel(grad, grad_prev, descent_direction_prev)[source]

Compute beta based on Hestenes-Stiefel.

core.nonlinear_cg.beta_dai_yuan(grad, grad_prev, descent_direction_prev)[source]

Compute beta based on Dai-Yuan.

core.nonlinear_cg.beta_average(grad, grad_prev, descent_direction_prev)[source]

Compute beta as an average of the existing methods.

core.phase_matrix_funcs module

core.phase_matrix_funcs.calc_group_delay_phase(refractive_index, n_arr, k0_arr, wavelength_0, wavelength)[source]

Uses finite differences to get the group index at a specfific wavelength and returns the linearized group delay per mm at that wavelength.

core.phase_matrix_funcs.calculate_phase_matrix_material(measurement_info, parameters, central_f)[source]

Calculates a phase matrix via material dispersion. Not differentiable due to usage of refractiveindex.

Parameters:

  • measurement_info (Pytree) – holds measurement data and parameters, needs to contain the material thickness theta in mm.

  • parameters (refractiveindex.RefractiveIndexMaterial) – an object providing the refractive index, the speed of light in m/s

Returns:

jnp.array, the calculated phase matrix

core.phase_matrix_funcs.calc_sine_phase(omega, phase_shift, parameters)[source]
core.phase_matrix_funcs.calc_tanh_phase(omega, phase_shift, parameters)[source]
core.phase_matrix_funcs.calc_gaussian_phase(omega, phase_shift, parameters)[source]
core.phase_matrix_funcs.calc_MIIPS_phase(omega, phase_shift, parameters)[source]
core.phase_matrix_funcs.calc_G_MIIPS_phase(omega, phase_shift, parameters)[source]
core.phase_matrix_funcs.calculate_phase_matrix(measurement_info, parameters, phase_func=<function calc_MIIPS_phase>)[source]

Calculates a phase matrix using a specified phase function.

Parameters:

  • measurement_info (Pytree) – holds measurement data and parameters, needs to contain the shift-values in appropriate units

  • parameters (tuple) – the parameters which are expected by phase_func

  • phase_func (Callable) – defines how the phase is calculated

Returns:

jnp.array, the calculated phase matrix

core.phase_matrix_funcs.calc_GDD(omega, phase_shift, parameters, phase_func=<function calc_sine_phase>)[source]

Uses jax.grad to get the GDD of phase_func. Thus doesnt work with non-differentiable functions. Assumes phase_func is holomorphic.

core.spectral_filter_funcs module

core.spectral_filter_funcs.gaussian_filter(frequency, parameters, filter_dict)[source]
core.spectral_filter_funcs.lorentzian_filter(frequency, parameters, filter_dict)[source]
core.spectral_filter_funcs.rectangular_filter(frequency, parameters, filter_dict)[source]
core.spectral_filter_funcs.multi_filter(frequency, parameters, filter_dict)[source]
core.spectral_filter_funcs.get_filter(filter_type, frequency, parameters, custom_func=None, refractive_index=<refractiveindex.refractiveindex.RefractiveIndexMaterial object>, material_thickness=0)[source]

Generate a spectral filter.

Parameters:

  • filter_func (str, tuple[str]) – can be one of `gaussian, lorentzian, rectangular, multi or custom`. For multi, parameters needs to specify the repsective filetr function

  • frequency (jnp.array) – the frequency axis

  • parameters (tuple) – the parameters required by the filter function, have the input form (a, f0, fwhm, p)

  • custom_func (Callable, None) – in case of filter_func=”custom” the custom filter function needs to be provided here.

  • refractive_index (refractiveindex.RefractiveIndexMaterial) – provides the refractive index

  • material_thickness (float, int) – the material thickness in millimeters

Returns:

jnp.array, the spectral filter on the frequency axis

core.spectral_filter_funcs.get_phase_matrix(refractive_index, material_thickness, frequency, central_frequency)[source]

Calculates the phase matrix that is applied of a pulse passes through a material.

core.stepsize module

core.stepsize.backtracking_linesearch(linesearch_state, error_func, grad_func, linesearch_params, linesearch_info, measurement_info)[source]

Perform one iteration of an Armijo-style linesearch.

Parameters:

  • linesearch_state (Pytree) – contains the armijo condition, gamma and iteration

  • error_func (Callable) – calculates the error to be optimized, expects gamma, linesearch_info and measurement_info

  • grad_func (Callable) – unused

  • linesearch_params (Pytree) – contains parameters for the linesearch iteration

  • linesearch_info (Pytree) – contains variables related to the initial state of the linesearch

  • measurement_info (Pytree) – contains measurement data and parameters

Returns:

MyNamespace, the updated linesearch_state

core.stepsize.finding_phase(current_vals, bracket, linesearch_info, linesearch_params)[source]
core.stepsize.zoom_interpolation(low_vals, high_vals)[source]
core.stepsize.zoom_phase(current_vals, bracket, linesearch_info, linesearch_params)[source]
core.stepsize.zoom_linesearch(linesearch_state, error_func, grad_func, linesearch_params, linesearch_info, measurement_info)[source]

Perform one iteration of an zoom linesearch. Works in two phases. In the first phase an interval containing a minimum is located by increasing the step size and observation of the behavior of the error and error-gradient. If a suitable interval has been located the minimum is approximated/found through successive cubic interpolation. The search terminates if the Armijo- and Wolfe conditions are met or one runs out of iterations.

Parameters:

  • linesearch_state (Pytree) – contains information on the current state of the linesearch

  • error_func (Callable) – calculates the error to be optimized, expects gamma, linesearch_info and measurement_info

  • grad_func (Callable) – calculates the error-gradient to be optimized, expects gamma, linesearch_info and measurement_info

  • linesearch_params (Pytree) – contains parameters for the linesearch iteration

  • linesearch_info (Pytree) – contains variables related to the initial state of the linesearch

  • measurement_info (Pytree) – contains measurement data and parameters

Returns:

MyNamespace, the updated linesearch_state

core.stepsize.do_linesearch_step(linesearch_state, linesearch_info, measurement_info, linesearch_params, error_func, grad_func)[source]
core.stepsize.end_linesearch(linesearch_state, max_steps_linesearch)[source]
core.stepsize.do_linesearch(linesearch_info, measurement_info, descent_info, error_func, grad_func, local_or_global)[source]

Perform a linesearch to obtain an improved step size in a descent based optimization.

Parameters:

  • linesearch_info (Pytree) – holds information on the initial state at gamma=0

  • measurement_info (Pytree) – holds measurement data and parameters

  • descent_info (Pytree) – holds parameters of the descent algorithm

  • error_func (Callable) – calculates the error to be optimized, expects gamma, linesearch_info, measurement_info

  • grad_func (Callable) – calculates the gradient of error_func, expects gamma, linesearch_info, measurement_info

  • local_or_global (str) – whether this is used in a local or global iteration

Returns:

float, the approximated optimal step size

core.stepsize.get_scaling(gradient, descent_direction, xi, local_or_global_state, pulse_or_gate, local_or_global)[source]
core.stepsize.adaptive_step_size(error, gradient, descent_direction, xi, local_or_global_state, adaptive_scaling_info, pulse_or_gate, local_or_global)[source]

Calculate an improved step size based through a pade approximation of the error function at the current position.

Parameters:

  • error (float) – the current error

  • gradient (jnp.array) – the current gradient

  • descent_direction (jnp.array) – the current descent direction

  • local_or_global_state (Pytree) – holds information of the current descent_state

  • xi (float) – a damping factor to avoid division by zero

  • order (str) – the pade-approximation to be used, can be one of pade_10 (linear), pade_20 (nonlinear), pade_01, pade_11 or pade_02

  • pulse_or_gate (str) – whether this is applied to pulse or gate

  • local_or_global (str) – whether this happens inside a local or global iteration

Returns:

tuple[jnp.array, Pytree], the scaled descent direction and the local_or_global_state

Module contents