pytc.vmc.optimizer.create_gradient_mask

pytc.vmc.optimizer.create_gradient_mask(ansatz, params, frozen_params)[source]

Create a gradient mask PyTree for the combined params structure.

Assumes params = [jastrow_params, linear_coeffs]. The mask is applied only to the jastrow_params part based on frozen_params identifiers. The linear_coeffs part of the mask is always True (not frozen).

Parameters:

  • ansatz – The wavefunction ansatz object.

  • params – The combined parameters PyTree [jastrow_params, linear_coeffs].

  • frozen_params – A list of identifiers (int index or str name/type) for Jastrow factors whose parameters should be frozen.

Returns:

A PyTree with the same structure as params, where frozen parameters are wrapped with jax.lax.stop_gradient.