SimModel¶

class simmodel.SimModel(nneighbor_cutoff, output_forces=True, virial=False, check_nlist=False, dtype=<sphinx.ext.autodoc.importer._MockObject object>, name='htf-model', **kwargs)¶

Bases: sphinx.ext.autodoc.importer._MockObject

SimModel is the main way that HOOMD-TF interacts with a simulation.

__init__(nneighbor_cutoff, output_forces=True, virial=False, check_nlist=False, dtype=<sphinx.ext.autodoc.importer._MockObject object>, name='htf-model', **kwargs)¶

SimModel is the main way that HOOMD-TF interacts with a simulation. Any kwargs are passed to setup().

Parameters:

nneighbor_cutoff (int) – The maximum number of neighbors to consider (can be 0)

output_forces (bool) – True if your graph will compute forces to be used in the Hoomd simulation

check_nlist (bool) – True will raise error if neighbor list overflows (nneighbor_cutoff too low)

dtype (dtype) – The floating point specification for model (e.g., tf.float32)

compute(nlist, positions, box, training=True)¶

The main method were computation occurs occurs. This method must be implemented by subclass. You may take less args, e.g. (nlist, positions). It should return one or more values as a tuple. The first element is interpreted as forces (if output_forces=True, default). Second element is interpreted as virial (if virial=True, not default). Subsequent elements of tuple are only accessible if tfcompute.attach() is passed save_output_period, after which you can obtain from the tfcompute object as the outputs attribute. Use compute_nlist_forces() or compute_positions_forces() to compute forces from an energy.

Parameters:

nlist (tensor) – an N x NN x 4 tensor containing the nearest neighbors. An entry of all zeros indicates that less than NN nearest neighbors where present for a particular particle. The last axis 4 dimensions are x,y,z and w, which is the particle type. Particle type is an integer starting at 0. Note that the x,y,z values are a vector originating at the particle and ending at its neighbor.
positions (tensor) – an N x 4 tensor of particle positions (x,y,z) and type.
box (tensor) – a 3x3 tensor containing the low box coordinate (row 0), high box coordinate (row 1), and then tilt factors (row 2). Call box_size() to convert to size
training (bool) – a boolean indicating if doing training or inference.

Returns:

Tuple of tensors

mapped_nlist(nlist)¶

If tfcompute.enable_mapped_nlist() is called, this method will split the nlist into the all atom and mapped.

Parameters:	nlist (tensor) – The `nlist` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom nlist. Second is the mapped nlist.

mapped_positions(positions)¶

If tfcompute.enable_mapped_nlist() is called, this method will split the positions into the all atom and mapped.

Parameters:	positions (tensor) – The `positions` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom positions. Second is the mapped positions.

retrace_compute()¶

Force a retrace of the compute function. This is necessary if your compute function depends variables inside self. For example:

def compute(self, nlist):
    if self.flag:
        nlist *= 2

If self.flag is changed after executing your model, you must call this function to force TF retrace your function.

setup(**kwargs)¶: This method can be implemented by a subclass to perform tasks after object creation. Any kwargs passed to SimModel.__init__() will be args here. This method will be called automatically.

MolSimModel¶

class simmodel.MolSimModel(MN, mol_indices, nneighbor_cutoff, output_forces=True, virial=False, check_nlist=False, dtype=<sphinx.ext.autodoc.importer._MockObject object>, name='htf-mol-model', **kwargs)¶

Bases: simmodel.SimModel

A molecular batched SimModel

Warning

Hoomd re-orders positions to improve performance. MolSimModel will disable this to keep a specific ordering of positions.

__init__(MN, mol_indices, nneighbor_cutoff, output_forces=True, virial=False, check_nlist=False, dtype=<sphinx.ext.autodoc.importer._MockObject object>, name='htf-mol-model', **kwargs)¶

The change from SimModel.__init__() are the first two parameters.

Parameters:

MN (int) – The number of atoms in a molecule. MN must be chosen to be large enough to encompass all molecules. If your molecule is 6 atoms and you chose MN=18, then the extra entries will be zeros.
mol_indices (list of lists) – mol_indices describes the molecules in your system as a list of atom indices. This can be created directly from a hoomd system via find_molecules(). The mol_indices are a, possibly ragged, 2D python list where each element in the list is a list of atom indices for a molecule. For example, [[0,1], [1]] means that there are two molecules with the first containing atoms 0 and 1 and the second containing atom 1. Note that the molecules can be different size and atoms can exist in multiple molecules.

mapped_nlist(nlist)¶

If tfcompute.enable_mapped_nlist() is called, this method will split the nlist into the all atom and mapped.

Parameters:	nlist (tensor) – The `nlist` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom nlist. Second is the mapped nlist.

mapped_positions(positions)¶

If tfcompute.enable_mapped_nlist() is called, this method will split the positions into the all atom and mapped.

Parameters:	positions (tensor) – The `positions` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom positions. Second is the mapped positions.

mol_compute(nlist, positions, mol_nlist, mol_positions, box, training)¶

See SimModel.compute() for details. Make sure that your forces still use nlist when computing, instead of mol_nlist. You may take less args in your implementation, like mol_compute(self, nlist, positions, mol_nlist).

Parameters:

nlist (tensor) – an N x NN x 4 tensor containing the nearest neighbors. An entry of all zeros indicates that less than NN nearest neighbors where present for a particular particle. The last axis 4 dimensions are x,y,z and w, which is the particle type. Particle type is an integer starting at 0. Note that the x,y,z values are a vector originating at the particle and ending at its neighbor.
positions (tensor) – an N x 4 tensor of particle positions (x,y,z) and type.
mol_nlist – a mol_number x MN x NN x 4 tensor containing the nearest neighbors broken out by molecule. An entry of all zeros indicates that less than NN nearest neighbors where present for a particular particle. The last axis 4 dimensions are x,y,z and w, which is the particle type. Particle type is an integer starting at 0. Note that the x,y,z values are a vector originating at the particle and ending at its neighbor.
mol_positions – a mol_number x MN x N x 4 tensor of particle positions (x,y,z) and type.
box (tensor) – a 3x3 tensor containing the low box coordinate (row 0), high box coordinate (row 1), and then tilt factors (row 2). Call box_size() to convert to size
training – a boolean indicating if doing training or inference.

Returns:

Tuple of tensors

retrace_compute()¶

Force a retrace of the compute function. This is necessary if your compute function depends variables inside self. For example:

def compute(self, nlist):
    if self.flag:
        nlist *= 2

If self.flag is changed after executing your model, you must call this function to force TF retrace your function.

setup(**kwargs)¶: This method can be implemented by a subclass to perform tasks after object creation. Any kwargs passed to SimModel.__init__() will be args here. This method will be called automatically.

Functions to use with SimModel¶

These are functions for working with nlist, position, box etc. and computing forces.

class simmodel.MolSimModel(MN, mol_indices, nneighbor_cutoff, output_forces=True, virial=False, check_nlist=False, dtype=<sphinx.ext.autodoc.importer._MockObject object>, name='htf-mol-model', **kwargs)

A molecular batched SimModel

Warning

Hoomd re-orders positions to improve performance. MolSimModel will disable this to keep a specific ordering of positions.

compute(nlist, positions, box, training)¶

The main method were computation occurs occurs. This method must be implemented by subclass. You may take less args, e.g. (nlist, positions). It should return one or more values as a tuple. The first element is interpreted as forces (if output_forces=True, default). Second element is interpreted as virial (if virial=True, not default). Subsequent elements of tuple are only accessible if tfcompute.attach() is passed save_output_period, after which you can obtain from the tfcompute object as the outputs attribute. Use compute_nlist_forces() or compute_positions_forces() to compute forces from an energy.

Parameters:

nlist (tensor) – an N x NN x 4 tensor containing the nearest neighbors. An entry of all zeros indicates that less than NN nearest neighbors where present for a particular particle. The last axis 4 dimensions are x,y,z and w, which is the particle type. Particle type is an integer starting at 0. Note that the x,y,z values are a vector originating at the particle and ending at its neighbor.
positions (tensor) – an N x 4 tensor of particle positions (x,y,z) and type.
box (tensor) – a 3x3 tensor containing the low box coordinate (row 0), high box coordinate (row 1), and then tilt factors (row 2). Call box_size() to convert to size
training (bool) – a boolean indicating if doing training or inference.

Returns:

Tuple of tensors

mapped_nlist(nlist)

If tfcompute.enable_mapped_nlist() is called, this method will split the nlist into the all atom and mapped.

Parameters:	nlist (tensor) – The `nlist` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom nlist. Second is the mapped nlist.

mapped_positions(positions)

If tfcompute.enable_mapped_nlist() is called, this method will split the positions into the all atom and mapped.

Parameters:	positions (tensor) – The `positions` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom positions. Second is the mapped positions.

mol_compute(nlist, positions, mol_nlist, mol_positions, box, training)

See SimModel.compute() for details. Make sure that your forces still use nlist when computing, instead of mol_nlist. You may take less args in your implementation, like mol_compute(self, nlist, positions, mol_nlist).

Parameters:

nlist (tensor) – an N x NN x 4 tensor containing the nearest neighbors. An entry of all zeros indicates that less than NN nearest neighbors where present for a particular particle. The last axis 4 dimensions are x,y,z and w, which is the particle type. Particle type is an integer starting at 0. Note that the x,y,z values are a vector originating at the particle and ending at its neighbor.
positions (tensor) – an N x 4 tensor of particle positions (x,y,z) and type.
mol_nlist – a mol_number x MN x NN x 4 tensor containing the nearest neighbors broken out by molecule. An entry of all zeros indicates that less than NN nearest neighbors where present for a particular particle. The last axis 4 dimensions are x,y,z and w, which is the particle type. Particle type is an integer starting at 0. Note that the x,y,z values are a vector originating at the particle and ending at its neighbor.
mol_positions – a mol_number x MN x N x 4 tensor of particle positions (x,y,z) and type.
box (tensor) – a 3x3 tensor containing the low box coordinate (row 0), high box coordinate (row 1), and then tilt factors (row 2). Call box_size() to convert to size
training – a boolean indicating if doing training or inference.

Returns:

Tuple of tensors

retrace_compute()

Force a retrace of the compute function. This is necessary if your compute function depends variables inside self. For example:

def compute(self, nlist):
    if self.flag:
        nlist *= 2

If self.flag is changed after executing your model, you must call this function to force TF retrace your function.

setup(**kwargs): This method can be implemented by a subclass to perform tasks after object creation. Any kwargs passed to SimModel.__init__() will be args here. This method will be called automatically.

class simmodel.SimModel(nneighbor_cutoff, output_forces=True, virial=False, check_nlist=False, dtype=<sphinx.ext.autodoc.importer._MockObject object>, name='htf-model', **kwargs)

SimModel is the main way that HOOMD-TF interacts with a simulation.

compute(nlist, positions, box, training=True)

The main method were computation occurs occurs. This method must be implemented by subclass. You may take less args, e.g. (nlist, positions). It should return one or more values as a tuple. The first element is interpreted as forces (if output_forces=True, default). Second element is interpreted as virial (if virial=True, not default). Subsequent elements of tuple are only accessible if tfcompute.attach() is passed save_output_period, after which you can obtain from the tfcompute object as the outputs attribute. Use compute_nlist_forces() or compute_positions_forces() to compute forces from an energy.

Parameters:

nlist (tensor) – an N x NN x 4 tensor containing the nearest neighbors. An entry of all zeros indicates that less than NN nearest neighbors where present for a particular particle. The last axis 4 dimensions are x,y,z and w, which is the particle type. Particle type is an integer starting at 0. Note that the x,y,z values are a vector originating at the particle and ending at its neighbor.
positions (tensor) – an N x 4 tensor of particle positions (x,y,z) and type.
box (tensor) – a 3x3 tensor containing the low box coordinate (row 0), high box coordinate (row 1), and then tilt factors (row 2). Call box_size() to convert to size
training (bool) – a boolean indicating if doing training or inference.

Returns:

Tuple of tensors

mapped_nlist(nlist)

If tfcompute.enable_mapped_nlist() is called, this method will split the nlist into the all atom and mapped.

Parameters:	nlist (tensor) – The `nlist` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom nlist. Second is the mapped nlist.

mapped_positions(positions)

If tfcompute.enable_mapped_nlist() is called, this method will split the positions into the all atom and mapped.

Parameters:	positions (tensor) – The `positions` from py:meth:.compute
Returns:	Tuple of 2 tensors. First is all atom positions. Second is the mapped positions.

retrace_compute()

Force a retrace of the compute function. This is necessary if your compute function depends variables inside self. For example:

def compute(self, nlist):
    if self.flag:
        nlist *= 2

If self.flag is changed after executing your model, you must call this function to force TF retrace your function.

setup(**kwargs): This method can be implemented by a subclass to perform tasks after object creation. Any kwargs passed to SimModel.__init__() will be args here. This method will be called automatically.

simmodel.compute_nlist_forces(nlist, energy, virial=False)¶

Computes pairwise forces given a potential energy function that computes per-particle or overall energy. Returns forces as a N x 4 tensor, where the last dimension of axis 1 is per-particle energy if available, otherwise 0.

Parameters:	nlist (tensor) – `N x NN x 4` or `N` x `NN` x 3 neighbor list energy (tensor) – The potential energy. Can be size `1` `N` or `N x L` virial (bool) – True if virial contribution will be computed.
Returns:	Either the forces or a tuple with forces, virial.

simmodel.compute_positions_forces(positions, energy)¶

Computes position dependent forces given a potential energy function. Returns forces as a N x 4 tensor, where the last dimension of axis 1 is per-particle energy if available, otherwise 0.

Parameters:	positions – `N x 4` or `N x 3` positions tensor energy (tensor) – The potential energy. Can be size `1`, `N`, or `N x ?`
Returns:	Forces as tensor

simmodel.compute_rdf(nlist, r_range, type_tensor=None, nbins=100, type_i=None, type_j=None)¶

Computes the pairwise radial distribution function

Parameters:

nlist (tensor) – Neighbor list to use for RDF calculation.
r_range (2 element list) – A list containing two elements, begin and end, for r range.
type_tensor (tensor) – N x 1 tensor containing types. Can use positions[:, 3]
bins (int) – The bins to use for the RDF
type_i (int) – Use this to select the first particle type.
type_j (int) – Use this to select the second particle type.

Returns:

length nbins tensor of the RDF (not normalized).

simmodel.masked_nlist(nlist, type_tensor, type_i=None, type_j=None)¶

Returns a neighbor list masked by the given particle type(s).

Parameters:	nlist (tensor) – Neighbor list to use for RDF calculation. type_tensor (tensor) – `N x 1` tensor containing types. Can use `positions[:, 3]` type_i (int) – Use this to select the first particle type. type_j (int) – Use this to select the second particle type.
Returns:	The masked neighbor list tensor.

simmodel.nlist_rinv(nlist)¶: Returns an N x NN tensor of 1 / r for each neighbor while correctly treating zeros. Empty neighbors are still zero and it is differentiable.

simmodel.safe_norm(tensor, delta=1e-07, **kwargs)¶

There are some numerical instabilities that can occur during learning when gradients are propagated. The delta is problem specific. Note you should not take a safe norm and then pass to tf.math.divide_no_nan See this TensorFlow issue.

Parameters:	tensor – the tensor over which to take the norm delta – small value to add so near-zero is treated without too much accuracy loss.
Returns:	The safe norm op (TensorFlow operation)

simmodel.wrap_vector(r, box)¶

Computes the minimum image version of the given vector.

Parameters:	r (tensor) – The vector to wrap around the Hoomd box.
Returns:	The wrapped vector as a TF tensor