m3gnet.utils package

Module contents

Utilities for m3gnet

class m3gnet.utils.Gaussian(centers: Union[tensorflow.python.framework.ops.Tensor, numpy.ndarray], width: float, **kwargs)

Bases: object

Gaussian expansion function

class m3gnet.utils.SphericalBesselFunction(max_l: int, max_n: int = 5, cutoff: float = 5.0, smooth: bool = False)

Bases: object

Calculate the spherical Bessel function based on the sympy + tensorflow implementations

static rbf_j0(r: tensorflow.python.framework.ops.Tensor, cutoff: float = 5.0, max_n: int = 3) → tensorflow.python.framework.ops.Tensor

Spherical Bessel function of order 0, ensuring the function value vanishes at cutoff

Args:

r: tf.Tensor tensorflow tensors cutoff: float, the cutoff radius max_n: int max number of basis

Returns: basis function expansion using first spherical Bessel function

class m3gnet.utils.SphericalHarmonicsFunction(max_l: int, use_phi: bool = True)

Bases: object

Spherical Harmonics function

m3gnet.utils.broadcast_states_to_atoms(graph: List)

Broadcast state attributes of shape [Ns, Nstate] to bond attributes shape [Na, Nstate] Args:

graph: List

Returns: broadcasted state attributes

m3gnet.utils.broadcast_states_to_bonds(graph: List)

Broadcast state attributes of shape [Ns, Nstate] to bond attributes shape [Nb, Nstate] Args:

graph: List

Returns: broadcasted state attributes

m3gnet.utils.check_array_equal(array1: Optional[numpy.ndarray], array2: Optional[numpy.ndarray], rtol: float = 0.001) → bool

Check the equality of two arrays Args:

array1 (np.ndarray): first array array2 (np.narray): second array rtol (float): relative tolerance

Returns: Bool

m3gnet.utils.check_shape_consistency(array: Optional[numpy.ndarray], shape: Sequence) → bool

Check if array complies with shape. Shape is a sequence of integer that may end with None. If None is at the end of shape, then any shapes in array after that dimension will match with shape.

Example: array with shape [10, 20, 30, 40] matches with [10, 20, None], but

does not match with shape [10, 20, 30, 20]

Args:

array (np.ndarray or None): array to be checked shape (Sequence): integer array shape, it may ends with None

Returns: bool

m3gnet.utils.combine_sbf_shf(sbf: tensorflow.python.framework.ops.Tensor, shf: tensorflow.python.framework.ops.Tensor, max_n: int, max_l: int, use_phi: bool)

Combine the spherical Bessel function and the spherical Harmonics function For the spherical Bessel function, the column is ordered by

[n=[0, …, max_n-1], n=[0, …, max_n-1], …], max_l blocks,

For the spherical Harmonics function, the column is ordered by

[m=[0], m=[-1, 0, 1], m=[-2, -1, 0, 1, 2], …] max_l blocks, and each block has 2*l + 1 if use_phi is False, then the columns become [m=[0], m=[0], …] max_l columns

Args:

sbf: tf.Tensor spherical bessel function results shf: tf.Tensor spherical harmonics function results max_n: int, max number of n max_l: int, max number of l use_phi: whether to use phi

Returns:

m3gnet.utils.get_length(t)

equivalent to tf.linalg.norm(t, axis=1), but faster Args:

t: (tf.Tensor), two dimensional tensor

Returns: tf.Tensor

m3gnet.utils.get_range_indices_from_n(ns: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Give ns = [2, 3], return [0, 1, 0, 1, 2] Args:

ns: tf.Tensor, the number of atoms/bonds array

Returns: range indices

m3gnet.utils.get_segment_indices_from_n(ns: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Get segment indices from number array. For example if ns = [2, 3], then the function will return [0, 0, 1, 1, 1]

Args:

ns: tf.Tensor, the number of atoms/bonds array

Returns:

object:

Returns: segment indices tensor

m3gnet.utils.get_spherical_bessel_roots() → numpy.ndarray

Get precomputed Spherical Bessel function roots. The results is a 2D matrix with dimension [128, 128]. The n-th (0-based index） root of order l Spherical Bessel function is the [l, n] entry

Returns: 2D matrix of pre-computed roots

m3gnet.utils.register(cls)

Register m3gnet classes Args:

cls: class name

Returns:

m3gnet.utils.register_plain(cls)

register a simple class_name: instance pair to the custom object

Args:

cls: class name

m3gnet.utils.repeat_with_n(tensor: tensorflow.python.framework.ops.Tensor, n: tensorflow.python.framework.ops.Tensor) → tensorflow.python.framework.ops.Tensor

Repeat the first dimension according to n array. The Args:

tensor: Tensor, tensor to augment n: Tensor, array to specify the repeat times for

each item in the first dimension of tensor

Returns: repeated tensor

m3gnet.utils.reshape_array(array: numpy.ndarray, shape: Sequence) → numpy.ndarray

Take an array and reshape it according to shape. Here shape may contain None field at the end.

if array shape is [3, 4] and shape is [3, 4, None], then array is shaped to [3, 4, 1]. If the two shapes do not match then report an error

Args:

array (np.ndarray): array to be reshaped shape (Sequence): shape dimensions

Returns: np.ndarray, reshaped array

m3gnet.utils.spherical_bessel_roots(max_l: int, max_n: int)

Calculate the spherical Bessel roots. The n-th root of the l-th spherical bessel function is the [l, n] entry of the return matrix. The calculation is based on the fact that the n-root for l-th spherical Bessel function j_l, i.e., z_{j, n} is in the range [z_{j-1,n}, z_{j-1, n+1}]. On the other hand we know precisely the roots for j0, i.e., sinc(x)

Args:

max_l: max order of spherical bessel function max_n: max number of roots

Returns: root matrix of size [max_l, max_n]

m3gnet.utils.spherical_bessel_smooth(r, cutoff: float = 5.0, max_n: int = 10)

This is an orthogonal basis with first and second derivative at the cutoff equals to zero. The function was derived from the order 0 spherical Bessel function, and was expanded by the different zero roots

Ref:

https://arxiv.org/pdf/1907.02374.pdf

Args:

r: tf.Tensor distance tensor cutoff: float, cutoff radius max_n: int, max number of basis, expanded by the zero roots

Returns: expanded spherical harmonics with derivatives smooth at boundary

m3gnet.utils.unsorted_segment_fraction(data, segment_ids, num_segments)

Segment softmax Args:

data (tf.Tensor): original data segment_ids (tf.Tensor): segment ids num_segments (tf.Tensor): number of segments

Returns:

m3gnet.utils.unsorted_segment_softmax(data, segment_ids, num_segments, weights=None)

Unsorted segment softmax with optional weights Args:

data (tf.Tensor): original data segment_ids (tf.Tensor): tensor segment ids num_segments (int): number of segments

Returns: tf.Tensor