"""
MDSuite: A Zincwarecode package.
License
-------
This program and the accompanying materials are made available under the terms
of the Eclipse Public License v2.0 which accompanies this distribution, and is
available at https://www.eclipse.org/legal/epl-v20.html
SPDX-License-Identifier: EPL-2.0
Copyright Contributors to the Zincwarecode Project.
Contact Information
-------------------
email: zincwarecode@gmail.com
github: https://github.com/zincware
web: https://zincwarecode.com/
Citation
--------
If you use this module please cite us with:
Summary
-------
"""
import logging
import numpy as np
import tensorflow as tf
from tqdm import tqdm
log = logging.getLogger(__name__)
[docs]def get_triu_indicies(n_atoms):
"""
Version of np.triu_indices with k=1
Returns
---------
Returns a vector of size (2, None) instead of a tuple of two values like
np.triu_indices
"""
bool_mat = tf.ones((n_atoms, n_atoms), dtype=tf.bool)
# Just construct a boolean true matrix the size of one time_step
indices = tf.where(~tf.linalg.band_part(bool_mat, -1, 0))
# Get the indices of the lower triangle (without diagonals)
indices = tf.cast(indices, dtype=tf.int32) # Get the correct dtype
return tf.transpose(indices) # Return the transpose for convenience later
[docs]def get_neighbour_list(positions: tf.Tensor, cell=None, batch_size=None) -> tf.Tensor:
"""
Generate the neighbour list
Parameters
----------
positions: tf.Tensor
Tensor with shape (n_confs, n_atoms, 3) representing the coordinates
cell: list
If periodic boundary conditions are used, please supply the cell dimensions,
e.g. [13.97, 13.97, 13.97]. If the cell is provided minimum image convention
will be applied!
batch_size: int
Has to be evenly divisible by the the number of configurations.
Returns
-------
generator object which results all distances for the current batch of time steps
To get the real r_ij matrix for one time_step you can use the following:
r_ij_mat = np.zeros((n_atoms, n_atoms, 3))
r_ij_mat[np.triu_indices(n_atoms, k = 1)] = get_neighbour_list(``*args``)
r_ij_mat -= r_ij_mat.transpose(1, 0, 2)
"""
def get_rij_mat(positions, triu_mask, cell):
"""
Use the upper triangle of the virtual r_ij matrix constructed of
n_atoms * n_atoms matrix and subtract the transpose to get all distances once!
If PBC are used, apply the minimum image convention.
"""
r_ij_mat = tf.gather(positions, triu_mask[0], axis=1) - tf.gather(
positions, triu_mask[1], axis=1
)
if cell:
r_ij_mat -= tf.math.rint(r_ij_mat / cell) * cell
return r_ij_mat
n_atoms = positions.shape[1]
triu_mask = get_triu_indicies(n_atoms)
if batch_size is not None:
if not positions.shape[0] % batch_size == 0:
msg = (
"positions must be evenly divisible by batch_size, but are"
f" {positions.shape[0]} and {batch_size}"
)
log.error(msg)
raise RuntimeError(msg)
for positions_batch in tf.split(positions, batch_size):
yield get_rij_mat(positions_batch, triu_mask, cell)
else:
yield get_rij_mat(positions, triu_mask, cell)
# @tf.function
[docs]def get_triplets(
full_r_ij: tf.Tensor, r_cut: float, n_atoms: int, n_batches=200, disable_tqdm=True
) -> tf.Tensor:
"""Compute the triple indices within a cutoff
Mostly vectorized method to compute the triples inside the cutoff sphere. Therefore
a matrix of all possible distances *r_ijk* over all timesteps
(n_timesteps, n_atoms, n_atoms, n_atoms) is constructed. The first layer is
identical to the r_ij matrix, all consecutive layers are shifted along the "j" axis
which leads to the full r_ijk matrix.
To check for a triple, the depth can be converted into the third particle, yielding
an index over all *ijk* indices
Parameters
----------
full_r_ij: tf.Tensor
The full distance matrix (r_ij and r_ji) with the shape
(n_timesteps, n_atoms, n_atoms)
r_cut: float
The cutoff for the maximal triple distance
n_atoms: int
Number of atoms
n_batches: int
Number of batches to split the computation of the triples in.
A low number of batches can result in memory issues.
Returns
-------
tf.Tensor
Warnings
---------
Using the @tf.function decorator can speed things up! But it can also lead to
memory issues.
"""
if n_batches >= n_atoms:
n_batches = n_atoms - 1
r_ij = tf.norm(full_r_ij, axis=-1)
r_ij = tf.cast(r_ij, dtype=tf.float16) # Using float16 for maximal memory safety.
r_ij = tf.where(r_ij == 0, tf.ones_like(r_ij) * r_cut, r_ij)
batches = np.array_split(np.arange(1, n_atoms), n_batches)
triples = []
# batches would be [(1, 100), (101, 200), (201, 300), ...]
for batch in tqdm(batches, ncols=70, disable=disable_tqdm):
r_ijk = tf.TensorArray(tf.float16, size=len(batch))
for idx, atom in enumerate(batch):
tmp = tf.roll(r_ij, shift=-atom, axis=2)
r_ijk = r_ijk.write(idx, tmp)
r_ijk = r_ijk.stack()
r_ijk = tf.transpose(r_ijk, perm=(1, 0, 2, 3)) # shape is (t, n, i, j)
intermediate_indices = tf.where(
tf.math.logical_and(r_ij[:, None] < r_cut, r_ijk < r_cut)
)
n_atoms = tf.cast(n_atoms, dtype=tf.int64)
t, n, i, j = tf.unstack(intermediate_indices, axis=1)
k = j + n + batch[0] # this comes from tf.roll.
k = tf.where(k >= n_atoms, k - n_atoms, k)
triples.append(tf.stack([t, i, j, k], axis=1))
return tf.concat(triples, axis=0)