agedi.models.schnetpack.regressor_heads¶

Classes¶

`Forces`	Predict the atomic force on the atoms in the structure.
`Energy`	Predict the potential energy of the structure.

Functions¶

build_gated_equivariant_mlp(→ torch.nn.Sequential)

Build neural network analog to MLP with `GatedEquivariantBlock`s instead of dense layers.

Module Contents¶

agedi.models.schnetpack.regressor_heads.build_gated_equivariant_mlp(s_in: int, v_in: int, n_out: int, n_layers: int = 2, activation: Callable = F.silu, sactivation: Callable = F.silu) → torch.nn.Sequential¶

Build neural network analog to MLP with `GatedEquivariantBlock`s instead of dense layers.

Parameters:

n_in (int) – Number of input nodes.
n_out (int) – Number of output nodes.
n_layers (int) – Number of layers.
activation (Callable) – Activation function.
sactivation (Callable) – Activation function for the skip connection.
n_hidden (int) – Number of hidden nodes.

Return type:

nn.Module

class agedi.models.schnetpack.regressor_heads.Forces(input_dim_scalar: int = 64, input_dim_vector: int = 64, gated_blocks: int = 3, **kwargs)¶

Bases: agedi.models.head.Head

Predict the atomic force on the atoms in the structure.

Parameters:

input_dim_scalar (int) – The dimension of the scalar input.
input_dim_vector (int) – The dimension of the vector input.
gated_blocks (int) – The number of gated blocks in the network.

Return type:

Head

_key = 'forces'¶

input_dim_scalar = 64¶

input_dim_vector = 64¶

gated_blocks = 3¶

net¶

get_hparams() → Dict¶

Return hyperparameters sufficient to reconstruct this head.

Returns:: Hyperparameter dictionary.
Return type:: dict

_score(translated_batch: dict) → torch.Tensor¶

Predict the force on the atoms in the structure.

Parameters:: translated_batch (dict) – The translated input batch.
Returns:: The predicted forces tensor.
Return type:: torch.Tensor

predict(translated_batch: dict) → torch.Tensor¶

Predict forces – alias for Forces._score() kept for backwards compatibility.

Parameters:: translated_batch (dict) – The translated input batch.
Returns:: The predicted forces tensor.
Return type:: torch.Tensor

class agedi.models.schnetpack.regressor_heads.Energy(input_dim_scalar: int = 64, **kwargs)¶