mlcolvar.core.loss.CommittorLoss

class mlcolvar.core.loss.CommittorLoss(alpha: float, atomic_masses: Tensor | None, cell: float = None, gamma: float = 10000.0, delta_f: float = 0.0, separate_boundary_dataset: bool = True, descriptors_derivatives: SmartDerivatives | Tensor = None, log_var: bool = False, use_gradients_wrt_positions: bool = True, z_regularization: float = 0.0, z_threshold: float = None, n_dim: int = None)[source]

Bases: Module

Compute a loss function based on Kolmogorov’s variational principle (or its approximation) for the determination of the committor function

__init__(alpha: float, atomic_masses: Tensor | None, cell: float = None, gamma: float = 10000.0, delta_f: float = 0.0, separate_boundary_dataset: bool = True, descriptors_derivatives: SmartDerivatives | Tensor = None, log_var: bool = False, use_gradients_wrt_positions: bool = True, z_regularization: float = 0.0, z_threshold: float = None, n_dim: int = None)[source]

Compute Kolmogorov’s variational principle loss and impose boundary conditions on the metastable states

Parameters:

  • alpha (float) – Hyperparamer that scales the boundary conditions contribution to loss, i.e. alpha*(loss_bound_A + loss_bound_B)

  • atomic_masses (torch.Tensor, optional) – List of masses of all the atoms we are using, for each atom we need to repeat three times for x,y,z, by default None. The mlcolvar.cvs.committor.utils.initialize_committor_masses can be used to simplify this. If the position-less loss is used, this must be set to None.

  • cell (float, optional) – CUBIC cell size length, used to scale the positions from reduce coordinates to real coordinates, by default None

  • gamma (float, optional) – Hyperparamer that scales the whole loss to avoid too small numbers, i.e. gamma*(loss_var + loss_bound), by default 10000

  • delta_f (float, optional) – Delta free energy between A (label 0) and B (label 1), units is kBT, by default 0. State B is supposed to be higher in energy.

  • separate_boundary_dataset (bool, optional) – Switch to exculde boundary condition labeled data from the variational loss, by default True

  • descriptors_derivatives (Union[SmartDerivatives, torch.Tensor], optional) – Derivatives of descriptors wrt atomic positions (if used) to speed up calculation of gradients, by default None. Can be either:

    • A SmartDerivatives object to save both memory and time, see also mlcolvar.core.loss.committor_loss.SmartDerivatives

    • A torch.Tensor with the derivatives to save time, memory-wise could be less efficient

  • ref_idx (torch.Tensor, optional) – Reference indeces for the unshuffled dataset for properly handling batching/splitting/shuffling when descriptors derivatives are provided, by default None. Ref_idx can be generated automatically using SmartDerivatives or by setting create_ref_idx=True when initializing a DictDataset. See also mlcolvar.core.loss.utils.smart_derivatives.SmartDerivatives

  • log_var (bool, optional) – Switch to minimize the log of the variational functional, by default False.

  • use_gradients_wrt_positions (bool, optional) – Whether to use gradients with respect to positions as prescribed in the original Kolmogorov variational functional, by default True.

  • z_regularization (float, optional) – Scales a regularization on the learned z space preventing it from exceeding the threshold given with ‘z_threshold’. The magnitude of the regularization is scaled by the given number, by default 0.0

  • z_threshold (float, optional) – Sets a maximum threshold for the z value during the training, by default None. The magnitude of the regularization term is scaled via the z_regularization key.

  • n_dim (int) – Number of dimensions, by default None. If None, it defaults to 3 for the position-based loss and to 1 for the position-less loss.

Methods

__init__(alpha, atomic_masses[, cell, ...])

Compute Kolmogorov's variational principle loss and impose boundary conditions on the metastable states

forward(x, z, q, labels, w[, ref_idx, ...])

Committor loss forward pass
forward(x: Tensor, z: Tensor, q: Tensor, labels: Tensor, w: Tensor, ref_idx: Tensor = None, create_graph: bool = True) → Tuple[Tensor, Tensor, Tensor, Tensor][source]

Committor loss forward pass

Parameters:

  • x (torch.Tensor) – Model input, i.e., either positions or descriptors if using descriptors_derivatives

  • z (torch.Tensor) – Model unactivated output, i.e., z value

  • q (torch.Tensor) – Model final output, i.e., committor value

  • labels (torch.Tensor) – Input labels

  • w (torch.Tensor) – Input weights

  • ref_idx (torch.Tensor, optional) – Reference indeces for the unshuffled dataset for properly handling batching/splitting/shuffling when descriptors derivatives are provided, by default None. Ref_idx can be generated automatically using SmartDerivatives or by setting create_ref_idx=True when initializing a DictDataset. See also mlcolvar.core.loss.utils.smart_derivatives.SmartDerivatives

  • create_graph (bool, optional) – Whether to create the graph during the computation for backpropagation, by default True

Returns:

Total loss and its components, i.e., variational, boundary A, and boundary B

Return type:

Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Attributes

T_destination

call_super_init

dump_patches

training