Utilities¶
There are a few convenience functions for plotting potential energies of pairwise potentials and constructing CG mappings.
RDF¶
To compute an RDF, use compute_rdf()
:
class LJRDF(htf.SimModel):
def setup(self):
self.avg_rdf = tf.keras.metrics.MeanTensor()
def compute(self, nlist, positions, box):
# get r
r = tf.norm(tensor=nlist[:, :, :3], axis=2)
# compute 1 / r while safely treating r = 0.
# pairwise energy. Double count -> divide by 2
inv_r6 = tf.math.divide_no_nan(1., r**6)
p_energy = 4.0 / 2.0 * (inv_r6 * inv_r6 - inv_r6)
# rdf from r = 3 to r = 5
rdf, rs = htf.compute_rdf(nlist, [3, 5])
# compute running mean
self.avg_rdf.update_state(rdf)
forces = htf.compute_nlist_forces(nlist, p_energy)
return forces
Then after your simulation, access the RDF as a numpy array with:
rdf = model.avg_rdf.result().numpy()
Pairwise Potential and Forces¶
To take your model and compute pairwise outputs,
use compute_pairwise()
, which can
be convenient for computing pairwise energy or forces.
model = build_examples.LJModel(4)
r = np.linspace(0.5, 1.5, 5)
output = htf.compute_pairwise(model, r)
Trajectory Parsing¶
To process information from a trajectory, use
iter_from_trajectory()
. This generator will process information from a trajectory and
yield a tuple of [nlist, positions, box]
(see SimModel.compute()
for details)
and MDAnalysis.TimeStep
object.
The first list can be directly called with a SimModel
(e.g., model(inputs)
).
The MDAnalysis.TimeStep
object can be used to compute other properties with MDAnalysis.
Here’s an example:
model = MyModel(16)
for inputs, ts in htf.iter_from_trajectory(16, universe):
result = model(inputs)
positions = inputs[1]
# compute something with position...
and here’s an example of you can do training, assuming forces exist
in your MDAnalysisUniverse
:
model = MyModel(16)
losses = []
for inputs, ts in htf.iter_from_trajectory(16, universe):
forces = ts.forces
l = model.train_on_batch(inputs, forces)
losses.append(l)
Coarse-Graining¶
Find Molecules¶
To go from atom index to particle index, use the
find_molecules()
:
# The method takes in a hoomd system as an argument.
...
molecule_mapping_index = htf.find_molecules(system)
...
Sparse Mapping¶
The sparse_mapping()
method creates the necessary indices and
values for defining a sparse tensor in tensorflow that is a
mass-weighted \(M \times N\) mapping operator where \(M\) is the number of
coarse-grained particles and \(N\) is the number of atoms in the system. In
the following example,mapping_per_molecule
is a list of \(k \times n\) matrices where
\(k\) is the number of coarse-grained sites for each molecule and \(n\) is the
number of atoms in the corresponding molecule. There should be one
matrix per molecule. Since the example is for a 1 bead mapping per
molecule the shape is \(1 \times n\). The ordering of the atoms should follow the
output from the find_molecules method. The variable
molecule_mapping_index
is the output from
find_molecules()
.
#The example is shown for 1 coarse-grained site per molecule.
...
molecule_mapping_matrix = numpy.ones([1, len(molecule_mapping_index[0])], dtype=np.int)
mapping_per_molecule = [molecule_mapping_matrix for _ in molecule_mapping_index]
cg_mapping = htf.sparse_mapping(mapping_per_molecule, \
molecule_mapping_index, system = system)
...
Center of Mass¶
center_of_mass()
maps the given positions according to
the specified mapping operator to coarse-grain site positions, while
considering periodic boundary conditions. The coarse grain site position
is placed at the center of mass of its constituent atoms.
...
mapped_position = htf.center_of_mass(graph.positions[:,:3], cg_mapping, system)
#cg_mapping is the output from the sparse_matrix(...) method and indicates how each molecule is mapped.
...
Compute Mapped Neighbor List¶
compute_nlist()
returns the neighbor list for a set of
mapped coarse-grained particles. In the following example, mapped_positions
is
the mapped particle positions obeying the periodic boundary condition, as
returned by center_of_mass()
, rcut
is the cutoff
radius and NN
is the number of nearest neighbors to be considered
for the coarse-grained system.
...
mapped_nlist= htf.compute_nlist(mapped_positions, rcut, NN, system)
...
Tensorboard¶
You can visualize your models with Tensorboard to observe metrics and other quantities you choose in a web browser. Find out more about Tensorboard.