train

Helpers for training parameters.

Classes:

• AttributeConfig –

  Configuration for how a potential's attributes should be trained.

• ParameterConfig –

  Configuration for how a potential's parameters should be trained.

• Trainable –

  A convenient wrapper around a tensor force field that gives greater control

AttributeConfig pydantic-model

Bases: BaseModel

Configuration for how a potential's attributes should be trained.

Fields:

• cols (list[str])
• scales (dict[str, float])
• limits (dict[str, tuple[float | None, float | None]])
• regularize (dict[str, float])

Validators:

• _validate_keys

cols pydantic-field

cols: list[str]

The parameters to train, e.g. 'k', 'length', 'epsilon'.

scales pydantic-field

scales: dict[str, float] = {}

The scales to apply to each parameter, e.g. 'k': 1.0, 'length': 1.0, 'epsilon': 1.0.

limits pydantic-field

limits: dict[str, tuple[float | None, float | None]] = {}

The min and max values to clamp each parameter within, e.g. 'k': (0.0, None), 'angle': (0.0, pi), 'epsilon': (0.0, None), where none indicates no constraint.

regularize pydantic-field

regularize: dict[str, float] = {}

The regularization strength to apply to each parameter, e.g. 'k': 0.01, 'epsilon': 0.001. Parameters not listed are not regularized.

ParameterConfig pydantic-model

Bases: AttributeConfig

Configuration for how a potential's parameters should be trained.

Fields:

• cols (list[str])
• scales (dict[str, float])
• limits (dict[str, tuple[float | None, float | None]])
• regularize (dict[str, float])
• include (PotentialKeyList | None)
• exclude (PotentialKeyList | None)

Validators:

• _validate_keys
• _validate_include_exclude

include pydantic-field

include: PotentialKeyList | None = None

The keys (see smee.TensorPotential.parameter_keys for details) corresponding to specific parameters to be trained. If None, all parameters will be trained.

exclude pydantic-field

exclude: PotentialKeyList | None = None

The keys (see smee.TensorPotential.parameter_keys for details) corresponding to specific parameters to be excluded from training. If None, no parameters will be excluded.

cols pydantic-field

cols: list[str]

The parameters to train, e.g. 'k', 'length', 'epsilon'.

scales pydantic-field

scales: dict[str, float] = {}

The scales to apply to each parameter, e.g. 'k': 1.0, 'length': 1.0, 'epsilon': 1.0.

limits pydantic-field

limits: dict[str, tuple[float | None, float | None]] = {}

The min and max values to clamp each parameter within, e.g. 'k': (0.0, None), 'angle': (0.0, pi), 'epsilon': (0.0, None), where none indicates no constraint.

regularize pydantic-field

regularize: dict[str, float] = {}

The regularization strength to apply to each parameter, e.g. 'k': 0.01, 'epsilon': 0.001. Parameters not listed are not regularized.

Trainable

Trainable(
    force_field: TensorForceField,
    parameters: dict[str, ParameterConfig],
    attributes: dict[str, AttributeConfig],
    vsites: ParameterConfig | None = None,
)

A convenient wrapper around a tensor force field that gives greater control over how parameters should be trained.

This includes imposing limits on the values of parameters, scaling the values so parameters passed to the optimizer have similar magnitudes, and freezing parameters so they are not updated during training.

parameters: Configure which parameters to train.
attributes: Configure which attributes to train.
vsites: Configure which vsite parameters to train.

Methods:

• to_values –

  Returns unfrozen parameter and attribute values as a flat tensor.

• to_force_field –

  Returns a force field with the parameters and attributes set to the given

• clamp –

  Clamps the given values to the configured min and max values.

Attributes:

• regularized_idxs (Tensor) –

  The indices (within the tensor returned by to_values)

• regularization_weights (Tensor) –

  The regularization weights for parameters/attributes to regularize.

Source code in descent/train.py

def __init__(
    self,
    force_field: smee.TensorForceField,
    parameters: dict[str, ParameterConfig],
    attributes: dict[str, AttributeConfig],
    vsites: ParameterConfig | None = None,
):
    """

    Args:
        force_field: The force field to wrap.
        parameters: Configure which parameters to train.
        attributes: Configure which attributes to train.
        vsites: Configure which vsite parameters to train.
    """
    self._force_field = force_field
    self._fit_vsites = False

    param_block = self._prepare(force_field, parameters, "parameters")
    attr_block = self._prepare(force_field, attributes, "attributes")

    self._param_types = sorted(parameters)
    self._param_shapes = param_block.shapes
    self._attr_types = sorted(attributes)
    self._attr_shapes = attr_block.shapes

    blocks: list[_PreparedBlock] = [param_block, attr_block]

    if vsites is not None:
        blocks.append(self._prepare_vsites(force_field, vsites))
        self._fit_vsites = True

    # Each block's indices are zero-based within that block. Apply the
    # running value-tensor offset at the point of combination so that the
    # offset arithmetic is in one place and adding further blocks requires
    # no changes beyond appending to `blocks`.
    all_values = []
    all_unfrozen_idxs = []
    all_scales = []
    all_clamp_lower = []
    all_clamp_upper = []
    all_regularized_idxs = []
    all_regularization_weights = []

    offset = 0
    for block in blocks:
        all_values.append(block.values)
        all_unfrozen_idxs.append(block.unfrozen_idxs + offset)
        all_scales.append(block.scales)
        all_clamp_lower.append(block.clamp_lower)
        all_clamp_upper.append(block.clamp_upper)
        all_regularized_idxs.append(block.regularized_idxs + offset)
        all_regularization_weights.append(block.regularization_weights)
        offset += len(block.values)

    self._values = torch.cat(all_values)
    self._unfrozen_idxs = torch.cat(all_unfrozen_idxs).long()
    self._scales = torch.cat(all_scales)
    self._clamp_lower = torch.cat(all_clamp_lower)
    self._clamp_upper = torch.cat(all_clamp_upper)

    # Map global flat indices -> positions within the unfrozen vector.
    # regularized_idxs are guaranteed to be a subset of unfrozen_idxs,
    # so no missing-key guard is needed.
    combined_regularized_idxs = torch.cat(all_regularized_idxs).long()
    combined_regularization_weights = torch.cat(all_regularization_weights)

    idx_mapping = {idx.item(): i for i, idx in enumerate(self._unfrozen_idxs)}
    self._regularized_idxs = torch.tensor(
        [idx_mapping[idx.item()] for idx in combined_regularized_idxs]
    ).long()
    self._regularization_weights = (
        combined_regularization_weights
        if len(combined_regularization_weights) > 0
        else torch.tensor([])
    )

regularized_idxs property

regularized_idxs: Tensor

The indices (within the tensor returned by to_values) of parameters/attributes to regularize.

regularization_weights property

regularization_weights: Tensor

The regularization weights for parameters/attributes to regularize.

to_values

to_values() -> Tensor

Returns unfrozen parameter and attribute values as a flat tensor.

Source code in descent/train.py

@torch.no_grad()
def to_values(self) -> torch.Tensor:
    """Returns unfrozen parameter and attribute values as a flat tensor."""
    values_flat = self.clamp(self._values[self._unfrozen_idxs] * self._scales)
    return values_flat.detach().clone().requires_grad_()

to_force_field

to_force_field(values_flat: Tensor) -> TensorForceField

Returns a force field with the parameters and attributes set to the given values.

Parameters:

• values_flat (Tensor) –

  A flat tensor of parameter and attribute values. See to_values for the expected shape and ordering.

Source code in descent/train.py

def to_force_field(self, values_flat: torch.Tensor) -> smee.TensorForceField:
    """Returns a force field with the parameters and attributes set to the given
    values.

    Args:
        values_flat: A flat tensor of parameter and attribute values. See
            ``to_values`` for the expected shape and ordering.
    """
    potentials = self._force_field.potentials_by_type

    values = self._values.detach().clone()
    values[self._unfrozen_idxs] = (values_flat / self._scales).clamp(
        min=self._clamp_lower, max=self._clamp_upper
    )
    shapes = self._param_shapes + self._attr_shapes

    if self._fit_vsites:
        shapes.append(self._force_field.v_sites.parameters.shape)

    values = _unflatten_tensors(values, shapes)

    params = values[: len(self._param_shapes)]

    for potential_type, param in zip(self._param_types, params, strict=True):
        potentials[potential_type].parameters = param

    attrs = values[
        len(self._param_shapes) : len(self._param_shapes) + len(self._attr_shapes)
    ]

    for potential_type, attr in zip(self._attr_types, attrs, strict=True):
        potentials[potential_type].attributes = attr

    if self._fit_vsites:
        vsite_params = values[len(self._param_shapes) + len(self._attr_shapes) :]
        self._force_field.v_sites.parameters = vsite_params[0]

    return self._force_field

clamp

clamp(values_flat: Tensor) -> Tensor

Clamps the given values to the configured min and max values.

Source code in descent/train.py

@torch.no_grad()
def clamp(self, values_flat: torch.Tensor) -> torch.Tensor:
    """Clamps the given values to the configured min and max values."""
    return (values_flat / self._scales).clamp(
        min=self._clamp_lower, max=self._clamp_upper
    ) * self._scales