Add train utilities

[SimonBoothroyd]

Description

This PR adds a set of convenience tools / classes to mimic the kind of convenient API that forcebalance exposes via the new Trainable class.

The Trainable class takes configuration for parameters and attributes:

parameter_config = {
    "vdW": ParameterConfig(
        # define which 'columns' of parameters should be flagged as trainable
        cols=["epsilon", "sigma"],
        # define how each type of parameter should be scaled (similar to FB priors)
        scales={"epsilon": 10.0, "sigma": 1.0},
        # define limits for each type of parameter (e.g. epsilon should not be less than zero).
        limits={"epsilon": (0.0, None), "sigma": (0.0, None)},
        # optionally define which parameters should be trained.
        include=[mock_ff.potentials_by_type["vdW"].parameter_keys[0]],
        # OR to include all parameters (the default)
        # include=None   
    ),
    "Bonds": ParameterConfig(
        cols=["length"],
        scales={"length": 1.0},
        limits={"length": (0.1, 0.7)},
        exclude=[mock_ff.potentials_by_type["Bonds"].parameter_keys[0]],
    ),
}
attribute_config = {
    "vdW": AttributeConfig(
        cols=["scale_14"],
        scales={"scale_14": 0.1},
        limits={"scale_14": (0.0, None)},
    )
}

Once created,

trainable = Trainable(ff, parameters=parameter_configs, attributes=attribute_configs)

a flat tensor of only the parameters / attributes that are un-frozen (i.e. allow to be trained) and scaled based on the config can be retrieved:

values = trainable.to_values()

These values are cloned from the original force field, so any changes made to values won't be reflected on the force field itself. To create a force field with the updated parameters / attributes:

ff = trainable.to_force_field(values)

Status

• Ready to go

[SimonBoothroyd]

@jthorton does this look reasonable to you?

[jthorton]

@SimonBoothroyd this is working great in testing!