core.models.escaip.configs

Classes

GlobalConfigs
MolecularGraphConfigs
GraphNeuralNetworksConfigs
RegularizationConfigs
EScAIPConfigs

Functions

resolve_type_hint(cls, field)	Resolves forward reference type hints from string to actual class objects.
init_configs(cls, kwargs)	Initialize a dataclass with the given kwargs.

Module Contents

class core.models.escaip.configs.GlobalConfigs

regress_forces: bool

direct_forces: bool

hidden_size: int

num_layers: int

activation: Literal['squared_relu', 'gelu', 'leaky_relu', 'relu', 'smelu', 'star_relu'] = 'gelu'

regress_stress: bool = False

use_compile: bool = True

use_padding: bool = True

use_fp16_backbone: bool = False

dataset_list: list

class core.models.escaip.configs.MolecularGraphConfigs

use_pbc: bool

max_num_elements: int

max_atoms: int

max_batch_size: int

max_radius: float

knn_k: int

knn_soft: bool

knn_sigmoid_scale: float

knn_lse_scale: float

knn_use_low_mem: bool

knn_pad_size: int

distance_function: Literal['gaussian', 'sigmoid', 'linearsigmoid', 'silu'] = 'gaussian'

use_envelope: bool = True

class core.models.escaip.configs.GraphNeuralNetworksConfigs

atten_name: Literal['math', 'memory_efficient', 'flash']

atten_num_heads: int

atom_embedding_size: int = 128

node_direction_embedding_size: int = 64

node_direction_expansion_size: int = 10

edge_distance_expansion_size: int = 600

edge_distance_embedding_size: int = 512

readout_hidden_layer_multiplier: int = 2

output_hidden_layer_multiplier: int = 2

ffn_hidden_layer_multiplier: int = 2

use_angle_embedding: Literal['scalar', 'bias', 'none'] = 'none'

angle_expansion_size: int = 10

angle_embedding_size: int = 8

use_graph_attention: bool = False

use_message_gate: bool = False

use_global_readout: bool = False

use_frequency_embedding: bool = True

freequency_list: list

energy_reduce: Literal['sum', 'mean'] = 'sum'

class core.models.escaip.configs.RegularizationConfigs

normalization: Literal['layernorm', 'rmsnorm', 'skip'] = 'rmsnorm'

mlp_dropout: float = 0.0

atten_dropout: float = 0.0

stochastic_depth_prob: float = 0.0

node_ffn_dropout: float = 0.0

edge_ffn_dropout: float = 0.0

scalar_output_dropout: float = 0.0

vector_output_dropout: float = 0.0

class core.models.escaip.configs.EScAIPConfigs

global_cfg: GlobalConfigs

molecular_graph_cfg: MolecularGraphConfigs

gnn_cfg: GraphNeuralNetworksConfigs

reg_cfg: RegularizationConfigs

core.models.escaip.configs.resolve_type_hint(cls, field)

Resolves forward reference type hints from string to actual class objects.

core.models.escaip.configs.init_configs(cls, kwargs)

Initialize a dataclass with the given kwargs.