pytc.vmc.metropolis.make_mcmc_step

pytc.vmc.metropolis.make_mcmc_step(ansatz, step_size, move_type='one', max_vmap_batch_size=0, mesh=None)[source]

Factory to create a JIT-compilable MCMC step function.

This function validates move_type at creation time (not JIT time) and returns a JIT-compilable step function that performs Metropolis-Hastings sampling.

The ansatz argument is captured at factory time and used for all subsequent MCMC steps. This is critical for optimize_ref_var, where the factory receives ref_det (a SlaterDet) so that walkers are sampled from |Det|² regardless of what ansatz object is passed at call time by the training loop.

Parameters:

  • ansatz – Wavefunction object — captured and used for MCMC proposals.

  • step_size – Standard deviation of Gaussian proposal for MCMC moves

  • move_type – “all” to move all electrons at once, “one” to move one electron at a time

  • max_vmap_batch_size – If > 0, use folx.batched_vmap with this batch size

Returns:

mcmc_step(ansatz, walkers, key, params) -> (new_walkers, acceptance_rate)

Return type:

A JIT-compiled function with signature

Raises:

ValueError – If move_type is not “all” or “one”