core.models.gemnet.utils

core.models.gemnet.utils#

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Functions#

`read_json`(path)
`update_json`(→ None)
`write_json`(→ None)
`read_value_json`(path, key)
`ragged_range`(→ torch.Tensor)	Multiple concatenated ranges.
`repeat_blocks`(→ torch.Tensor)	Repeat blocks of indices.
`calculate_interatomic_vectors`(→ tuple[torch.Tensor, ...)	Calculate the vectors connecting the given atom pairs,
`inner_product_normalized`(→ torch.Tensor)	Calculate the inner product between the given normalized vectors,
`mask_neighbors`(→ torch.Tensor)

Module Contents#

core.models.gemnet.utils.read_json(path: str)#

core.models.gemnet.utils.update_json(path: str, data) → None#

core.models.gemnet.utils.write_json(path: str, data) → None#

core.models.gemnet.utils.read_value_json(path: str, key: str)#

core.models.gemnet.utils.ragged_range(sizes: torch.Tensor) → torch.Tensor#

Multiple concatenated ranges.

Examples

sizes = [1 4 2 3] Return: [0 0 1 2 3 0 1 0 1 2]

core.models.gemnet.utils.repeat_blocks(sizes: torch.Tensor, repeats: int | torch.Tensor, continuous_indexing: bool = True, start_idx: int = 0, block_inc: int = 0, repeat_inc: int = 0) → torch.Tensor#

Repeat blocks of indices. Adapted from https://stackoverflow.com/questions/51154989/numpy-vectorized-function-to-repeat-blocks-of-consecutive-elements

continuous_indexing: Whether to keep increasing the index after each block start_idx: Starting index block_inc: Number to increment by after each block,

either global or per block. Shape: len(sizes) - 1

repeat_inc: Number to increment by after each repetition,: either global or per block

Examples

sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = False Return: [0 0 0 0 1 2 0 1 2 0 1 0 1 0 1] sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True Return: [0 0 0 1 2 3 1 2 3 4 5 4 5 4 5] sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; repeat_inc = 4 Return: [0 4 8 1 2 3 5 6 7 4 5 8 9 12 13] sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; start_idx = 5 Return: [5 5 5 6 7 8 6 7 8 9 10 9 10 9 10] sizes = [1,3,2] ; repeats = [3,2,3] ; continuous_indexing = True ; block_inc = 1 Return: [0 0 0 2 3 4 2 3 4 6 7 6 7 6 7] sizes = [0,3,2] ; repeats = [3,2,3] ; continuous_indexing = True Return: [0 1 2 0 1 2 3 4 3 4 3 4] sizes = [2,3,2] ; repeats = [2,0,2] ; continuous_indexing = True Return: [0 1 0 1 5 6 5 6]

core.models.gemnet.utils.calculate_interatomic_vectors(R: torch.Tensor, id_s: torch.Tensor, id_t: torch.Tensor, offsets_st: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]#

Calculate the vectors connecting the given atom pairs, considering offsets from periodic boundary conditions (PBC).

Parameters:

R (Tensor, shape = (nAtoms, 3)) – Atom positions.
id_s (Tensor, shape = (nEdges,)) – Indices of the source atom of the edges.
id_t (Tensor, shape = (nEdges,)) – Indices of the target atom of the edges.
offsets_st (Tensor, shape = (nEdges,)) – PBC offsets of the edges. Subtract this from the correct direction.

Returns:

(D_st, V_st) –

D_st: Tensor, shape = (nEdges,): Distance from atom t to s.
V_st: Tensor, shape = (nEdges,): Unit direction from atom t to s.

Return type:

tuple

core.models.gemnet.utils.inner_product_normalized(x: torch.Tensor, y: torch.Tensor) → torch.Tensor#: Calculate the inner product between the given normalized vectors, giving a result between -1 and 1.

core.models.gemnet.utils.mask_neighbors(neighbors: torch.Tensor, edge_mask: torch.Tensor) → torch.Tensor#