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],
)

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.

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],
):
    """

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

    (
        self._param_types,
        param_values,
        self._param_shapes,
        param_unfrozen_idxs,
        param_scales,
        param_clamp_lower,
        param_clamp_upper,
        param_regularized_idxs,
        param_regularization_weights,
    ) = self._prepare(force_field, parameters, "parameters")
    (
        self._attr_types,
        attr_values,
        self._attr_shapes,
        attr_unfrozen_idxs,
        attr_scales,
        attr_clamp_lower,
        attr_clamp_upper,
        attr_regularized_idxs,
        attr_regularization_weights,
    ) = self._prepare(force_field, attributes, "attributes")

    self._values = torch.cat([param_values, attr_values])

    self._unfrozen_idxs = torch.cat(
        [param_unfrozen_idxs, attr_unfrozen_idxs + len(param_scales)]
    ).long()

    self._scales = torch.cat([param_scales, attr_scales])[self._unfrozen_idxs]

    self._clamp_lower = torch.cat([param_clamp_lower, attr_clamp_lower])[
        self._unfrozen_idxs
    ]
    self._clamp_upper = torch.cat([param_clamp_upper, attr_clamp_upper])[
        self._unfrozen_idxs
    ]

    # Store regularization information
    all_regularized_idxs = torch.cat(
        [param_regularized_idxs, attr_regularized_idxs + len(param_scales)]
    ).long()
    all_regularization_weights = torch.cat(
        [param_regularization_weights, attr_regularization_weights]
    )

    # Map global indices to unfrozen indices
    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 all_regularized_idxs
            if idx.item() in idx_mapping
        ]
    ).long()
    regularization_weights = [
        all_regularization_weights[i]
        for i, idx in enumerate(all_regularized_idxs)
        if idx.item() in idx_mapping
    ]
    self._regularization_weights = (
        torch.stack(regularization_weights)
        if regularization_weights
        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
    )
    values = _unflatten_tensors(values, self._param_shapes + self._attr_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) :]

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

    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