pytc.vmc.optimization.optimize

pytc.vmc.optimization.optimize(ansatz, cost_fn=None, n_walkers=100, n_steps=1000, step_size=1.0, burn_in_steps=1000, use_importance_sampling=False, initial_walkers=None, key=None, n_opt_steps=100, max_vmap_batch_size=0, learning_rate=0.01, optimizer_type='adam', opt_kwargs=None, params=None, frozen_params=None, move_type='one', use_custom_jvp=True, adaptive_step_size=True, step_size_adjust_interval=10, save_frequency=100, save_path=None)

Perform wavefunction optimization using MCMC sampling.

Parameters:

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

• cost_fn – Cost function (defaults to average local energy if None)

• n_walkers (int) – Number of parallel walkers

• n_steps (int) – Number of MCMC steps for each walker in each opt iteration

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

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

• use_importance_sampling (bool) – Whether to use importance sampling with drift

• thinning – Keep only every thinning steps to reduce autocorrelation

• n_samples – If provided, collect this many uncorrelated samples

• 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.)

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

• jastrow_params – Initial Jastrow parameters

• linear_coeffs – Initial linear coefficients

• move_type (str)

• use_custom_jvp (bool)

• adaptive_step_size (bool)

• step_size_adjust_interval (int)

• save_frequency (int)

• save_path (str | None)

Returns:

Dictionary with optimization results and statistics

Return type:

Dict[str, Any]