model.prior

This module provides a simple classical force-field with pairwise non-bonded interactions and bonded, angular, and dihedral interactions.

Creating a force field consists of three simple steps. First, the force field parameters are read from a TOML file:

force_field = prior.ForceField.load_ff(ff_path)

Then, the force field is combined with a topology. This topoly can be infered from a graph, e.g., provided by the package mdtraj:

top = mdtraj.load_topology(top_path)
topology = prior.Topology.from_mdtraj(top, mapping=force_field.mapping(by_name=True))

Finally, topology and force field parameters are combined to create a potential:

prior_energy_fn_template = prior.init_prior_potential(
    displacement_fn, nonbonded_type="lennard_jones"
)
energy_fn = energy_fn_template(force_field, topology)

For a more detailed example, see the Prior Simulation example.

Topological Information

class Topology(num_particles, species=None, bond_idx=None, angle_idx=None, dihedral_idx=None, improper_dihedral_idx=None)

Class documenting the topology of a system.

This pytree stores the topological information of a system and provides utilities, to create this information e.g. from an mdtraj topology.

For a simple way to load a topology from a molecular graph, see Topology.from_mdtraj().

Parameters:

• num_particles (int) – The number of particles in the system

• species (Union[Array, ndarray, bool, number, bool, int, float, complex]) – Integer or array of species number.

• bond_idx (Union[Array, ndarray, bool, number, bool, int, float, complex]) – Array of shape [B, 2] with indices of the atoms that span the \(B\) bonds.

• angle_idx (Union[Array, ndarray, bool, number, bool, int, float, complex]) – Array of shape [A, 3] with indices of the atoms that span the \(A\) angles.

• dihedral_idx (Union[Array, ndarray, bool, number, bool, int, float, complex]) – Array of shape [D, 4] with indices of the atoms that span the \(D\) dihedral angles.

• improper_dihedral_idx (Union[Array, ndarray, bool, number, bool, int, float, complex]) – Not used.

classmethod concatenate(*topologies)

Concatenate the topologies.

classmethod from_mdtraj(topology, mapping=None)

Creates a topology instance from a mdtraj topology instance.

This function uses mdtraj to create a graph-representation of the system. It then discovers bonds, angles and dihedral angles by searching for all simple paths between tuples with corresponding length.

Parameters:

• topology – mdtraj topology object

• mapping – Function to map the atom data, e.g., atom symbol, to a species number.

Returns:

Returns a JAX topology based on the mdtraj topology.

get_angles()

Returns the indices, species and mask for all angles.

get_atom_species(idx=None)

Returns the species numbers for the atoms of the system.

get_bonds()

Returns the indices, species and mask for all angles.

get_dihedrals()

Returns the indices, species and mask for all angles.

get_padded_topology(num_particles, max_species=1, max_bonds=0, max_angles=0, max_dihedrals=0, max_improper_dihedrals=0)

Returns a new topology with padded graph data.

When dealing with systems that have a different number of bonds, etc. padding the topologies enables to still batch-process the systems.

prune_topology(force_field)

Remove all bonds, angles, etc. not parametrized by the force field.

Parameters:

force_field (ForceField) – Force field to identify, which bonds, etc. are parametrized.

Returns:

Returns a topology with potentially reduced size.

Force Field

class ForceField(data=None, lookup=None, mapping=None)

Parameter set for classical molecular dynamics potentials.

This class simplifies access to classical molecular dynamics potential parameters. This class is registered as a pytree with parameters as leaves, thus enable computing gradients with respect to the force field parameters.

To construct a force field from a TOML file, see ForceField.load_ff().

get_angle_params(s1, s2, s3)

Get the parameters of an angle between three species.

get_bond_params(s1, s2)

Get the parameters of a bond between two species.

get_data(keys=None)

Returns the trainable parameters.

It is possible to select only a subset of the trainable parameters. For example, by selecting only the bond parameters of the bonded interactions:

>>> selection = (
...     ("bonded", ("bonds",)),
... )
>>> ff.get_data(selection)

Parameters:

keys (Union[str, Tuple[Union[str, Tuple[str]]]]) – Key, or tuple structure of keys corresponding to the partial data that should be returned.

Returns:

Returns the learnable interaction parameters.

get_dihedral_params(s1, s2, s3, s4, multiplicity=None)

Get the parameters of an angle between four species.

get_nonbonded_params(s)

Get the nonbonded parameters for a species.

classmethod load_ff(fname)

Loads a force field from a toml file.

This is the recommended way to create a force field.

Parameters:

fname (str) – Force field parameter file

Returns:

Returns a force field instance with parameters read from the file.

mapping(by_name=False, renaming_pattern=None)

Returns a function that maps atom data into a species number.

Parameters:

• by_name – If true, then the atom name (e. g. "CA") is used as an identifier. Otherwise, the element symbol (e. g. "C") is used to look up the species number.

• renaming_pattern (List[Tuple[str, str]]) – Translate the atom names between different naming conventions, e.g. from PDB to AMBER.

Returns:

Returns a mapping function that maps atom data into a species number.

property max_species

Maximum number of species.

set_data(data)

Update the learnable interaction parameters.

In combination with get_data, set_data enables to learn only parts of the force field.

Parameters:

data – Dictionary, corresponding to the internal data structure.

Returns:

Returns a new instance of the force field.

write_ff(fname)

Writes the parameters of the force field back into a toml file.

Parameters:

fname (str) – Path, where the force field should be stored.

Classical Potential

init_prior_potential(displacement_fn, mask_bonded=True, nonbonded_type='repulsion', r_onset=0.45, r_cutoff=0.5)

Initializes the prior template.

Parameters:

• displacement_fn – jax_md displacement function

• mask_bonded (bool) – Excluded non-bonded interactions between bonded atoms via masking them in the neighbor list. This might be ineffective if the number of bonds is very large.

• nonbonded_type (str) – Type of the nonbonded interactions. Possible values are "repulsion", "lennard_jones", "truncated_lennard_jones".

• r_onset (float) – Distance to start smoothing the nonbonded interactions to 0.

• r_cutoff (float) – Cutoff distance of the neighbor list.

Helper Functions

init_bond_potential(displacement_fn, topology, force_field)

Initializes all bond-potentials for the given topology.

init_angle_potential(displacement_fn, topology, force_field)

Initializes all angle-potentials for the given topology.

init_dihedral_potential(displacement_fn, topology, force_field)

Initializes all dihedral-potentials for the given topology.

init_nonbonded_potential(displacement_fn, topology, force_field, nonbonded_type='repulsion', r_onset=0.45, r_cutoff=0.5)

Initializes the non-bonded interactions for the given topology.

constrain_ff_params(unconstrained_data)

Maps unconstrained force field data onto a constrained space.

The constrained space enforces the following properties:

• Force constants are positive \(k_B, k_\theta, k_D > 0\)

• Bond lengths are positive \(b_0> 0\)

• Angles are subject to \(\theta, \theta_D \in [0, 180]\)

Parameters:

unconstrained_data – Data obtained by force_field.get_data() to be constrained.

Returns:

Returns values mapped onto the constrained space.

unconstrain_ff_params(constrained_data)

Recovers the unconstrained data from the constrained space.

This function maps the parameters back onto the unconstrained space. Additionally, the mapping stops the backpropagation of gradients for the dihedral multiplicity.

Parameters:

constrained_data – Force field parameters in the constrained space.