phlower.nn.EnEquivariantMLP

class phlower.nn.EnEquivariantMLP(nodes, activations=None, dropouts=None, bias=True, create_linear_weight=False, norm_function_name='identity', cross_interaction=False, normalize=True)

Bases: IPhlowerCoreModule, Module

EnEquivariantMLP is a neural network module that performs an E(n)-equivariant operation on the input tensor.

Parameters:

• nodes (list[int]) – List of feature dimension sizes (The last value of tensor shape).

• activations (list[str] | None (optional)) – List of activation functions to apply to the output. Defaults to None.

• dropouts (list[float] | None (optional)) – List of dropout rates to apply to the output. Defaults to None.

• bias (bool) – Whether to use bias in the output.

• create_linear_weight (bool) – Whether to create a linear weight. Default is False.

• norm_function_name (str) – Name of the normalization function to apply to the output.

• cross_interaction (bool) – If True, use interaction layer proposed by https://arxiv.org/abs/2203.06442.

• normalize (bool) – If True, use eq (12) in https://arxiv.org/abs/2203.06442 when cross_interaction is True.

Examples

>>> en_equivariant_mlp = EnEquivariantMLP(
...     nodes=[10, 20, 30],
...     activations=["relu", "relu", "relu"],
...     dropouts=[0.1, 0.1, 0.1],
...     bias=True,
...     create_linear_weight=True,
...     norm_function_name="identity",
...     cross_interaction=False,
...     normalize=True,
... )
>>> en_equivariant_mlp(data)

Methods

forward(data, *[, field_data])	forward function which overloads torch.nn.Module
from_setting(setting)	Generate model from setting object
get_nn_name()	Return neural network name

Attributes

T_destination
call_super_init
dump_patches
training

forward(data, *, field_data=None, **kwards)

forward function which overloads torch.nn.Module

Parameters:

• data (IPhlowerTensorCollections) – IPhlowerTensorCollections data which receives from predecessors

• field_data (ISimulationField | None) – ISimulationField | None Constant information through training or prediction

Returns:

Tensor object

Return type:

PhlowerTensor

classmethod from_setting(setting)

Generate model from setting object

Parameters:

setting (EnEquivariantMLPSetting) – setting object

Returns:

EnEquivariantMLP object

Return type:

Self

classmethod get_nn_name()

Return neural network name

Returns:

name

Return type:

str