3999396905

Build Type

push

github-actions

Committed by GitHub

Commit Message

[merge before other PRs] ruff updates (#580)

Run Details

960 of 1311 branches covered (73.23%)

Branch coverage included in aggregate %.

15 of 15 new or added lines in 11 files covered. (100.0%)

4034 of 4930 relevant lines covered (81.83%)

3.19 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

44.64

/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!

1	"""
2	MDSuite: A Zincwarecode package.
3
4	License
5	-------
6	This program and the accompanying materials are made available under the terms
7	of the Eclipse Public License v2.0 which accompanies this distribution, and is
8	available at https://www.eclipse.org/legal/epl-v20.html
9
10	SPDX-License-Identifier: EPL-2.0
11
12	Copyright Contributors to the Zincwarecode Project.
13
14	Contact Information
15	-------------------
16	email: zincwarecode@gmail.com
17	github: https://github.com/zincware
18	web: https://zincwarecode.com/
19
20	Citation
21	--------
22	If you use this module please cite us with:
23
24	Summary
25	-------
26	"""
27
28	import tensorflow as tf	4✔
29
30
31	def unit_vector(vector):	4✔
32	"""Returns the unit vector of the vector."""
33	return vector / tf.expand_dims(tf.norm(vector, axis=-1), -1)	4✔
34
35
36	def angle_between(v1, v2, acos=True):	4✔
37	"""Returns the angle in radians between vectors 'v1' and 'v2'."""
38	v1_u = unit_vector(v1)	4✔
39	v2_u = unit_vector(v2)	4✔
40	# return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
41	# return tf.math.acos(tf.clip_by_value(tf.einsum("ijk, ijk -> ij", v1_u, v2_u),
42	# -1.0, 1.0))
43	if acos:	4!
44	return tf.math.acos(	4✔
45	tf.clip_by_value(tf.einsum("ij, ij -> i", v1_u, v2_u), -1.0, 1.0)
46	)
47	else:
48	return tf.einsum("ij, ij -> i", v1_u, v2_u)	×
49
50
51	def get_angles(r_ij_mat, indices, acos=True):	4✔
52	"""
53	Compute the cosine angle for the given triples.
54
55	Using :math theta = acos(r_ij * r_ik / (\|r_ij\| * \|r_ik\|))
56
57	Parameters
58	----------
59	acos: bool
60	Apply tf.math.acos to the output if true. default true.
61	r_ij_mat: tf.Tensor
62	r_ij matrix. Shape is (n_timesteps, n_atoms, n_atoms, 3)
63	indices: tf.Tensor
64	Indices of the triples. Shape is (triples, 4) where a single element is composed
65	of (timestep, idx, jdx, kdx)
66
67	Returns
68	-------
69	tf.Tensor: Tensor with the shape (triples)
70	"""
71	r_ij = tf.gather_nd(	4✔
72	r_ij_mat, tf.stack([indices[:, 0], indices[:, 1], indices[:, 2]], axis=1)
73	)
74	r_ik = tf.gather_nd(	4✔
75	r_ij_mat, tf.stack([indices[:, 0], indices[:, 1], indices[:, 3]], axis=1)
76	)
77
78	return (	4✔
79	angle_between(r_ij, r_ik, acos),
80	tf.linalg.norm(r_ij, axis=-1) * tf.linalg.norm(r_ik, axis=-1),
81	)
82
83
84	@tf.function(experimental_relax_shapes=True)	4✔
85	def apply_minimum_image(r_ij, box_array):	3✔
86	"""
87
88	Parameters
89	----------
90	r_ij: tf.Tensor
91	an r_ij matrix of size (batches, atoms, 3)
92	box_array: tf.Tensor
93	a box array (3,)
94
95	Returns
96	-------
97
98	"""
99	return r_ij - tf.math.rint(r_ij / box_array) * box_array	×
100
101
102	def get_partial_triu_indices(n_atoms: int, m_atoms: int, idx: int) -> tf.Tensor:	4✔
103	"""Calculate the indices of a slice of the triu values.
104
105	Parameters
106	----------
107	n_atoms: total number of atoms in the system
108	m_atoms: size of the slice (horizontal)
109	idx: start index of slize
110
111	Returns
112	-------
113	tf.Tensor
114
115	"""
116	bool_mat = tf.ones((m_atoms, n_atoms), dtype=tf.bool)	4✔
117	bool_vector = ~tf.linalg.band_part(bool_mat, -1, idx) # rename!	4✔
118
119	indices = tf.where(bool_vector)	4✔
120	indices = tf.cast(indices, dtype=tf.int32) # is this large enough?!	4✔
121	indices = tf.transpose(indices)	4✔
122	return indices	4✔
123
124
125	def apply_system_cutoff(tensor: tf.Tensor, cutoff: float) -> tf.Tensor:	4✔
126	"""
127	Enforce a cutoff on a tensor.
128
129	Parameters
130	----------
131	tensor : tf.Tensor
132	cutoff : flaot
133	"""
134	cutoff_mask = tf.cast(tf.less(tensor, cutoff), dtype=tf.bool) # Construct the mask	×
135
136	return tf.boolean_mask(tensor, cutoff_mask)	×
137
138
139	def cartesian_to_spherical_coordinates(	4✔
140	point_cartesian, name="cartesian_to_spherical_coordinates"
141	):
142	"""
143	References
144	----------
145	https://www.tensorflow.org/graphics/api_docs/python/tfg/math/math_helpers/cartesian_to_spherical_coordinates
146
147	Function to transform Cartesian coordinates to spherical coordinates.
148	This function assumes a right handed coordinate system with `z` pointing up.
149	When `x` and `y` are both `0`, the function outputs `0` for `phi`. Note that
150	the function is not smooth when `x = y = 0`.
151	Note:
152	In the following, A1 to An are optional batch dimensions.
153	Args:
154	point_cartesian: A tensor of shape `[A1, ..., An, 3]`. In the last
155	dimension, the data follows the `x`, `y`, `z` order.
156	eps: A small `float`, to be added to the denominator. If left as `None`, its
157	value is automatically selected using `point_cartesian.dtype`.
158	name: A name for this op. Defaults to "cartesian_to_spherical_coordinates".
159
160	Returns
161	-------
162	A tensor of shape `[A1, ..., An, 3]`. The last dimensions contains
163	(`r`,`theta`,`phi`), where `r` is the sphere radius, `theta` is the polar
164	angle and `phi` is the azimuthal angle. Returns `NaN` gradient if x = y = 0.
165	"""
166	with tf.name_scope(name):	×
167	point_cartesian = tf.convert_to_tensor(value=point_cartesian)	×
168
169	x, y, z = tf.unstack(point_cartesian, axis=-1)	×
170	radius = tf.norm(tensor=point_cartesian, axis=-1)	×
171	theta = tf.acos(tf.clip_by_value(tf.math.divide_no_nan(z, radius), -1.0, 1.0))	×
172	phi = tf.atan2(y, x)	×
173	return tf.stack((radius, theta, phi), axis=-1)	×
174
175
176	def spherical_to_cartesian_coordinates(	4✔
177	point_spherical, name="spherical_to_cartesian_coordinates"
178	):
179	"""
180	References
181	----------
182	https://www.tensorflow.org/graphics/api_docs/python/tfg/math/math_helpers/spherical_to_cartesian_coordinates
183
184	Function to transform Cartesian coordinates to spherical coordinates.
185	Note:
186	In the following, A1 to An are optional batch dimensions.
187	Args:
188	point_spherical: A tensor of shape `[A1, ..., An, 3]`. The last dimension
189	contains r, theta, and phi that respectively correspond to the radius,
190	polar angle and azimuthal angle; r must be non-negative.
191	name: A name for this op. Defaults to "spherical_to_cartesian_coordinates".
192
193	Raises
194	------
195	tf.errors.InvalidArgumentError: If r, theta or phi contains out of range
196	data.
197
198	Returns
199	-------
200	A tensor of shape `[A1, ..., An, 3]`, where the last dimension contains the
201	cartesian coordinates in x,y,z order.
202	"""
203	with tf.name_scope(name):	×
204	point_spherical = tf.convert_to_tensor(value=point_spherical)	×
205
206	r, theta, phi = tf.unstack(point_spherical, axis=-1)	×
207	tmp = r * tf.sin(theta)	×
208	x = tmp * tf.cos(phi)	×
209	y = tmp * tf.sin(phi)	×
210	z = r * tf.cos(theta)	×
211	return tf.stack((x, y, z), axis=-1)	×
212
213
214	def get2dHistogram(x, y, value_range, nbins=100, dtype=tf.dtypes.int32):	4✔
215	"""
216	Bins x, y coordinates of points onto simple square 2d histogram.
217
218	Given the tensor x and y:
219	x: x coordinates of points
220	y: y coordinates of points
221	this operation returns a rank 2 `Tensor`
222	representing the indices of a histogram into which each element
223	of `values` would be binned. The bins are equal width and
224	determined by the arguments `value_range` and `nbins`.
225
226
227	Parameters
228	----------
229	x: Numeric `Tensor`.
230	y: Numeric `Tensor`.
231	value_range[0] lims for x
232	value_range[1] lims for y
233
234	nbins: Scalar `int32 Tensor`. Number of histogram bins.
235	dtype: dtype for returned histogram.
236
237	References
238	----------
239	https://gist.github.com/isentropic/a86effab2c007e86912a50f995cac52b
240
241	"""
242	x_range = value_range[0]	×
243	y_range = value_range[1]	×
244
245	histy_bins = tf.histogram_fixed_width_bins(y, y_range, nbins=nbins, dtype=dtype)	×
246
247	H = tf.map_fn(	×
248	lambda i: tf.histogram_fixed_width(x[histy_bins == i], x_range, nbins=nbins),
249	tf.range(nbins),
250	)
251	return H # Matrix!	×

zincware / MDSuite / 3999396905

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous