pytc.vmc.optimization.optimize_ref_var

pytc.vmc.optimization.optimize_ref_var(ansatz, cost_fn=None, n_walkers=100, n_steps=20, step_size=1.0, burn_in_steps=1000, initial_walkers=None, key=None, n_opt_steps=100, max_vmap_batch_size=0, learning_rate=0.01, optimizer_type='adam', move_type='one', opt_kwargs=None, params=None, adaptive_step_size=True, step_size_adjust_interval=10, jacobian_sample_size=None, save_frequency=100, save_path=None, n_mcmc_per_opt=None, n_opt_per_mcmc=None)[source]

Perform variational Monte Carlo optimization using MCMC sampling.

Parameters:

  • ansatz – Wavefunction object with __call__ method that returns ψ(R)

  • cost_fn – Cost function (defaults to reference variance if None). Should accept (params, walkers) and return (cost, aux_data).

  • n_walkers (int) – Number of parallel walkers

  • n_steps (int) – Legacy training cadence parameter. If neither n_mcmc_per_opt nor n_opt_per_mcmc is provided, optimize_ref_var preserves its historical behavior and uses n_opt_per_mcmc=n_steps.

  • step_size (float) – Standard deviation of Gaussian proposal for MCMC

  • burn_in_steps (int) – Number of initial MCMC steps to discard (equilibration)

  • initial_walkers – Optional initial positions, otherwise initialized near nuclei

  • key – PRNG key

  • n_opt_steps (int) – Number of optimization steps

  • max_vmap_batch_size (int) – If 0, use standard vmap. If >0, use folx.batched_vmap with the given batch size for memory efficiency. Recommended: 10-50

  • learning_rate (float) – Learning rate for optimizer

  • optimizer_type (str) – Type of optimizer (“adam”, “sgd”, etc.)

  • move_type (str) – “one” or “all” for MCMC electron moves

  • opt_kwargs (Dict[str, Any] | None) – Additional optimizer parameters

  • params – Initial combined parameters [jastrow_params, linear_coeffs].

  • jacobian_sample_size (int | None) – Optional[int]. If provided and using Newton optimizer, subsample this many walkers for Jacobian computation (curvature matrix approximation). Speeds up Newton steps when n_walkers is large. Typical: 500-2000 for 100k walkers.

  • n_mcmc_per_opt (int | None) – Optional explicit number of MCMC steps before each optimization update.

  • n_opt_per_mcmc (int | None) – Optional explicit number of optimization steps before each MCMC refresh.

  • adaptive_step_size (bool)

  • step_size_adjust_interval (int)

  • save_frequency (int)

  • save_path (str | None)

Returns:

Dictionary with optimization results and statistics