pytc.vmc.loss.make_variance_loss

pytc.vmc.loss.make_variance_loss(ansatz, optimizer_type='adam', use_custom_jvp=True, max_vmap_batch_size=0, clip_multiplier=5.0, mesh=None)[source]

Factory to create variance-based loss function for reference variance optimization.

This minimizes the variance of local energies with respect to a reference determinant, which can improve the quality of the Jastrow factor.

Uses vmap (or batched_vmap for memory efficiency) following the same pattern as make_energy_loss.

Parameters:

  • ansatz – Wavefunction object with local_energy method

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

  • use_custom_jvp (bool) – Whether to use custom JVP for memory-efficient gradients

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

  • clip_multiplier (float) – Multiplier for energy clipping range (clips to mean ± multiplier * mean absolute deviation (MAD) of the local energy). Set to 0 to disable clipping. Default 5.0, matching make_energy_loss.

  • mesh (jax.sharding.Mesh | None) – Optional device mesh for sharding.

Returns:

Loss function with signature (params, batch_data) -> (variance, (mean_energy, energy_mad))