pytc.vmc.optimization

Optimization algorithms for VMC parameter optimization.

This module contains functions for optimizing wavefunction parameters using various optimization strategies and cost functions.

Design Patterns

The optimization module supports two complementary training patterns:

1. Energy Minimization (optimize function): - Pattern: Multiple MCMC steps → Single optimization step - Parameter: n_mcmc_per_opt (default: n_steps in optimize) - Rationale: Decorrelate walkers before computing gradients - Use case: Minimizing ground state energy - Example: n_mcmc_per_opt=20 means 20 MCMC steps, then 1 parameter update

2. Variance Minimization (optimize_ref_var function): - Pattern: Multiple optimization steps → Single MCMC step - Parameter: n_opt_per_mcmc (default: n_steps in optimize_ref_var) - Rationale: Multiple gradient steps on same walker configuration - Use case: Reducing variance for fixed reference determinant - Example: n_opt_per_mcmc=20 means 20 parameter updates, then 1 MCMC step

The factory function make_training_step() supports both patterns.

Functions

make_opt_update_step	Factory to create a JIT-compilable optimizer step for Optax optimizers.
make_second_order_training_step	Factory to create unified training step for optimizers with a .step() method.
make_training_step	Factory to create unified training step combining optimization and MCMC.
optimize	Perform wavefunction optimization using MCMC sampling.
optimize_ref_var	Perform variational Monte Carlo optimization using MCMC sampling.