pytc.vmc.optimization.make_training_step

pytc.vmc.optimization.make_training_step(mcmc_step, opt_update_step, n_mcmc_per_opt=1, n_opt_per_mcmc=1)[source]

Factory to create unified training step combining optimization and MCMC.

This creates a single JIT-compiled function that can perform either: 1. Multiple MCMC steps → single optimization (for energy minimization) 2. Multiple optimization steps → single MCMC step (for variance minimization) 3. Any combination of the above

Parameters:

  • mcmc_step – JIT-compiled MCMC step function from make_mcmc_step() Signature: (ansatz, walkers, key, params) -> (walkers, pmove)

  • opt_update_step – JIT-compiled optimizer step from make_opt_update_step() Signature: (ansatz, params, walkers, opt_state, key) -> (params, opt_state, loss, aux_data)

  • n_mcmc_per_opt – Number of MCMC steps before each optimization update (default: 1)

  • n_opt_per_mcmc – Number of optimization steps per MCMC update (default: 1) Note: If n_mcmc_per_opt > 1, this should typically be 1.

Returns:

training_step(ansatz, walkers, params, opt_state, key) ->

(walkers, params, opt_state, loss, aux_data, pmove)

Return type:

A JIT-compiled function with signature

Design patterns:

  • Energy minimization: n_mcmc_per_opt=10-100, n_opt_per_mcmc=1 (decorrelate walkers before each parameter update)

  • Variance minimization: n_mcmc_per_opt=1, n_opt_per_mcmc=5-20 (multiple gradient steps on same walker configuration)