mdsuite/utils/linalg.py
"""
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 tensorflow as tf
def unit_vector(vector):
"""Returns the unit vector of the vector."""
return vector / tf.expand_dims(tf.norm(vector, axis=-1), -1)
def angle_between(v1, v2, acos=True):
"""Returns the angle in radians between vectors 'v1' and 'v2'."""
v1_u = unit_vector(v1)
v2_u = unit_vector(v2)
# return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
# return tf.math.acos(tf.clip_by_value(tf.einsum("ijk, ijk -> ij", v1_u, v2_u),
# -1.0, 1.0))
if acos:
return tf.math.acos(
tf.clip_by_value(tf.einsum("ij, ij -> i", v1_u, v2_u), -1.0, 1.0)
)
else:
return tf.einsum("ij, ij -> i", v1_u, v2_u)
def get_angles(r_ij_mat, indices, acos=True):
"""
Compute the cosine angle for the given triples.
Using :math theta = acos(r_ij * r_ik / (|r_ij| * |r_ik|))
Parameters
----------
acos: bool
Apply tf.math.acos to the output if true. default true.
r_ij_mat: tf.Tensor
r_ij matrix. Shape is (n_timesteps, n_atoms, n_atoms, 3)
indices: tf.Tensor
Indices of the triples. Shape is (triples, 4) where a single element is composed
of (timestep, idx, jdx, kdx)
Returns
-------
tf.Tensor: Tensor with the shape (triples)
"""
r_ij = tf.gather_nd(
r_ij_mat, tf.stack([indices[:, 0], indices[:, 1], indices[:, 2]], axis=1)
)
r_ik = tf.gather_nd(
r_ij_mat, tf.stack([indices[:, 0], indices[:, 1], indices[:, 3]], axis=1)
)
return (
angle_between(r_ij, r_ik, acos),
tf.linalg.norm(r_ij, axis=-1) * tf.linalg.norm(r_ik, axis=-1),
)
@tf.function(experimental_relax_shapes=True)
def apply_minimum_image(r_ij, box_array):
"""
Parameters
----------
r_ij: tf.Tensor
an r_ij matrix of size (batches, atoms, 3)
box_array: tf.Tensor
a box array (3,)
Returns
-------
"""
return r_ij - tf.math.rint(r_ij / box_array) * box_array
def get_partial_triu_indices(n_atoms: int, m_atoms: int, idx: int) -> tf.Tensor:
"""Calculate the indices of a slice of the triu values.
Parameters
----------
n_atoms: total number of atoms in the system
m_atoms: size of the slice (horizontal)
idx: start index of slize
Returns
-------
tf.Tensor
"""
bool_mat = tf.ones((m_atoms, n_atoms), dtype=tf.bool)
bool_vector = ~tf.linalg.band_part(bool_mat, -1, idx) # rename!
indices = tf.where(bool_vector)
indices = tf.cast(indices, dtype=tf.int32) # is this large enough?!
indices = tf.transpose(indices)
return indices
def apply_system_cutoff(tensor: tf.Tensor, cutoff: float) -> tf.Tensor:
"""
Enforce a cutoff on a tensor.
Parameters
----------
tensor : tf.Tensor
cutoff : flaot
"""
cutoff_mask = tf.cast(tf.less(tensor, cutoff), dtype=tf.bool) # Construct the mask
return tf.boolean_mask(tensor, cutoff_mask)
def cartesian_to_spherical_coordinates(
point_cartesian, name="cartesian_to_spherical_coordinates"
):
"""
References:
----------
https://www.tensorflow.org/graphics/api_docs/python/tfg/math/math_helpers/cartesian_to_spherical_coordinates.
Function to transform Cartesian coordinates to spherical coordinates.
This function assumes a right handed coordinate system with `z` pointing up.
When `x` and `y` are both `0`, the function outputs `0` for `phi`. Note that
the function is not smooth when `x = y = 0`.
Note:
----
In the following, A1 to An are optional batch dimensions.
Args:
----
point_cartesian: A tensor of shape `[A1, ..., An, 3]`. In the last
dimension, the data follows the `x`, `y`, `z` order.
eps: A small `float`, to be added to the denominator. If left as `None`, its
value is automatically selected using `point_cartesian.dtype`.
name: A name for this op. Defaults to "cartesian_to_spherical_coordinates".
Returns:
-------
A tensor of shape `[A1, ..., An, 3]`. The last dimensions contains
(`r`,`theta`,`phi`), where `r` is the sphere radius, `theta` is the polar
angle and `phi` is the azimuthal angle. Returns `NaN` gradient if x = y = 0.
"""
with tf.name_scope(name):
point_cartesian = tf.convert_to_tensor(value=point_cartesian)
x, y, z = tf.unstack(point_cartesian, axis=-1)
radius = tf.norm(tensor=point_cartesian, axis=-1)
theta = tf.acos(tf.clip_by_value(tf.math.divide_no_nan(z, radius), -1.0, 1.0))
phi = tf.atan2(y, x)
return tf.stack((radius, theta, phi), axis=-1)
def spherical_to_cartesian_coordinates(
point_spherical, name="spherical_to_cartesian_coordinates"
):
"""
References:
----------
https://www.tensorflow.org/graphics/api_docs/python/tfg/math/math_helpers/spherical_to_cartesian_coordinates.
Function to transform Cartesian coordinates to spherical coordinates.
Note:
----
In the following, A1 to An are optional batch dimensions.
Args:
----
point_spherical: A tensor of shape `[A1, ..., An, 3]`. The last dimension
contains r, theta, and phi that respectively correspond to the radius,
polar angle and azimuthal angle; r must be non-negative.
name: A name for this op. Defaults to "spherical_to_cartesian_coordinates".
Raises:
------
tf.errors.InvalidArgumentError: If r, theta or phi contains out of range
data.
Returns:
-------
A tensor of shape `[A1, ..., An, 3]`, where the last dimension contains the
cartesian coordinates in x,y,z order.
"""
with tf.name_scope(name):
point_spherical = tf.convert_to_tensor(value=point_spherical)
r, theta, phi = tf.unstack(point_spherical, axis=-1)
tmp = r * tf.sin(theta)
x = tmp * tf.cos(phi)
y = tmp * tf.sin(phi)
z = r * tf.cos(theta)
return tf.stack((x, y, z), axis=-1)
def get2dHistogram(x, y, value_range, nbins=100, dtype=tf.dtypes.int32):
"""
Bins x, y coordinates of points onto simple square 2d histogram.
Given the tensor x and y:
x: x coordinates of points
y: y coordinates of points
this operation returns a rank 2 `Tensor`
representing the indices of a histogram into which each element
of `values` would be binned. The bins are equal width and
determined by the arguments `value_range` and `nbins`.
Parameters
----------
x: Numeric `Tensor`.
y: Numeric `Tensor`.
value_range[0] lims for x
value_range[1] lims for y
nbins: Scalar `int32 Tensor`. Number of histogram bins.
dtype: dtype for returned histogram.
References
----------
https://gist.github.com/isentropic/a86effab2c007e86912a50f995cac52b
"""
x_range = value_range[0]
y_range = value_range[1]
histy_bins = tf.histogram_fixed_width_bins(y, y_range, nbins=nbins, dtype=dtype)
H = tf.map_fn(
lambda i: tf.histogram_fixed_width(x[histy_bins == i], x_range, nbins=nbins),
tf.range(nbins),
)
return H # Matrix!