mxtaltools.models.modules.components

class mxtaltools.models.modules.components.Activation(activation_func: str, filters: int)

Bases: Module

Wrapper module for several activation options

Parameters:

• activation_func (str) – type of activation function

• filters (int) – feature depth of objects to be normalized

forward(x: Tensor) → Tensor

class mxtaltools.models.modules.components.EmbeddingBlock(init_node_embedding_dim: int, num_input_classes: int, num_scalar_input_features: int, atom_type_embedding_dim: int)

Bases: Module

forward(x: Tensor) → Tensor

class mxtaltools.models.modules.components.KernelActivation(n_basis: int, span: float, filters: int)

Bases: Module

Function for learning an activation function for every node in a given layer, as a linear combination of basis functions over a given span. Uses nn.Conv1d groups option for efficient evaluation.

Parameters:

• n_basis (int) – number of basis functions

• span (float) – span over which to define localized basis functions

• filters (int) – feature depth of inputs to be activated

forward(x: Tensor) → Tensor

kernel(x: Tensor) → Tensor

class mxtaltools.models.modules.components.Normalization(norm: str, filters: int)

Bases: Module

Wrapper module for several normalization options

Parameters:

• norm (str) – type of normalization function

• filters (int) – feature depth of objects to be normalized

forward(x: Tensor, batch: LongTensor | None = None)

class mxtaltools.models.modules.components.ResidualMLP(input_dim, hidden_dim, output_dim, num_layers, norm, dropout_p)

Bases: Module

a simpler general-purpose MLP

forward(x)

class mxtaltools.models.modules.components.Scalarizer(hidden_dim: int, embedding_dim: int, norm_mode: str, act_func: str, dropout: float = 0, output_dim: int = None)

Bases: Module

__init__(hidden_dim: int, embedding_dim: int, norm_mode: str, act_func: str, dropout: float = 0, output_dim: int = None)

Generate a learned invariant representation of dimension \((k)\) from a list of vectors of dimension \(v=(k x 3)\).

Generate m vectors as a linear combination of the k vectors of v, take their normalized dot products with the components of v, concatenate to the norms of v, and linearly combine to the so-called scalarized representation of v.

Parameters:

• hidden_dim (int) – feature depth of input and output

• embedding_dim (int) – number of vectors to use for dot-product projection

• norm_mode (str) – type of normalization to use

• act_func (str) – type of activation to use

• dropout (float) – dropout probability

forward(v: Tensor) → Tensor

class mxtaltools.models.modules.components.VectorActivation(hidden_dim: int, act_func: str)

Bases: Module

Modified implementation of the vector activation function from https://github.com/FlyingGiraffe/vnn/blob/master/models/vn_layers.py

Generates an axis as a learned linear combination of input v, then the normalized overlaps of all the components of v.

Applies an activation function on the vector overlaps, such that, e.g., for ReLU activation, vectors with negative overlap are rotated to be perpendicular to the learned axis (zero overlap) and vectors with positive overlap are untouched.

Parameters:

• hidden_dim (int) – feature depth of input/output vectors, \((k\times 3)\)

• act_func (str) – activation function to apply to the normalized vector overlaps with the learned axis

forward(v: Tensor) → Tensor

class mxtaltools.models.modules.components.scalarMLP(layers: int, filters: int, input_dim: int, output_dim: int, activation: str = 'gelu', seed: int = 0, dropout: float = 0, norm: str = None, norm_after_linear: bool = True, bias: bool = True, ramp_depth: bool = False)

Bases: Module

Flexible multi-layer perceptron module, with several options.

Parameters:

• layers (int) – number of fully-connected layers

• filters (int) – feature depth with FC layers

• input_dim (int) – feature depth of inputs

• output_dim (int) – feature depth of outputs

• activation (str) – activation function

• seed (int) – random seed

• dropout (float) – dropout probability

• ramp_depth (bool) – whether to ramp the feature depth exponentially from input_dim to output_dim through the network

forward(x: Tensor, return_latent: bool = False, batch: LongTensor | None = None) → Tuple[Tensor, Tensor | None]

initialize to correct feature dimension

init_scalar_filters(filters)

returns a list with layers + 1 integer elements

init_scalar_transforms()

scalar MLP layers

class mxtaltools.models.modules.components.vectorMLP(layers: int, filters: int, input_dim: int, output_dim: int, vector_input_dim: int, vector_output_dim: int, activation: str = 'gelu', seed: int = 0, dropout: float = 0, norm: str = None, norm_after_linear: bool = True, bias: bool = True, vector_norm: str = None, ramp_depth: bool = False, v_to_s_combination: str = 'sum')

Bases: scalarMLP

scalarMLP model with l=1 vector track added with o3 equivariance

forward(x: Tensor, v: Tensor, return_latent: bool = False, batch: LongTensor | None = None) → Tuple[Tensor, Tensor, Tensor | None]

initialize to correct feature dimension

get_residues(i: int, x: Tensor, v: Tensor) → Tuple[Tensor, Tensor]

init_vector_filters(filters)

init_vector_transforms()

vector MLP layers