pytc.vmc.loss.make_energy_loss

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

Factory to create energy-based loss function for VMC optimization.

This creates a loss function that computes the average local energy, with optional energy clipping and custom JVP for memory efficiency.

Parameters:

  • ansatz – Wavefunction object with local_energy method

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

  • cost_fn (Callable | None) – Optional cost function to apply to energies (defaults to mean)

  • clip_multiplier (float) – Multiplier for energy clipping range (clips to mean ± multiplier * std)

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

  • max_vmap_batch_size (int) – If 0, use standard vmap everywhere. If >0, use folx.batched_vmap (or sharded_batched_vmap if multi-device).

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

Returns:

Loss function with signature (params, batch_data) -> (loss, AuxData) where AuxData is a namedtuple with (mean_energy, energy_std, clipped_energies, diff)