pytc.vmc.optimizer.create_gradient_mask¶

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

Create a trainable-leaf mask for the combined parameter PyTree.

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 boolean PyTree with the same structure as params. False leaves are frozen and True leaves remain trainable. Returns None when no parameters are frozen.