MolecularDiffusion.runmodes.train.tasks_esen

Attributes

logger

Classes

ModelTaskFactory

Factory to construct eSEN models and task handlers for regression, guidance,

Module Contents

class MolecularDiffusion.runmodes.train.tasks_esen.ModelTaskFactory(task_type: str, atom_vocab: list, condition_names: list = [], hidden_size: int = 128, num_layers: int = 2, lmax: int = 2, mmax: int = 2, grid_resolution: int | None = None, cutoff: float = 5.0, edge_channels: int = 128, hidden_channels: int = 128, distance_function: str = 'gaussian', num_distance_basis: int = 512, norm_type: str = 'rms_norm_sh', act_type: str = 's2', mlp_type: str = 'grid', use_envelope: bool = False, activation_checkpointing: bool = False, global_attributes: bool = False, sphere_embedding_type: str = 'mixed', aggregation_method: str = 'sum', chkpt_path: str = None, **kwargs)

Factory to construct eSEN models and task handlers for regression, guidance, and SSL3D tasks.

Constructor Parameters:

task_type (str): “regression”, “guidance”, or “ssl3d”. atom_vocab (list): List of atom vocabulary used for encoding. condition_names (list): List of conditional labels. hidden_size (int): Hidden dimension size. num_layers (int): Number of layers.

# eSEN specific lmax (int): Maximum degree of spherical harmonics. mmax (int): Maximum order of spherical harmonics. grid_resolution (int): Resolution of the grid. cutoff (float): Cutoff radius. edge_channels (int): Number of edge channels. distance_function (str): Distance function type. num_distance_basis (int): Number of distance basis functions. norm_type (str): Normalization type. act_type (str): Activation type. mlp_type (str): MLP type. use_envelope (bool): Whether to use envelope function. activation_checkpointing (bool): Whether to use activation checkpointing. global_attributes (bool): Whether to use global attributes. sphere_embedding_type (str): Type of sphere embedding.

Property-prediction kwargs:

task_learn (List) criterion (str) metric (List) num_mlp_layer (int) mlp_dropout (float) normalization (bool) chkpt_path (str)

build()

Build and return (model, task) based on task_type.

Parameters:

task_type (str) – “regression”, “guidance”, or “ssl3d”.

act_type = 's2'

activation_checkpointing = False

aggregation_method = 'sum'

atom_vocab

chkpt_path = None

condition_names = []

context_node_nf

cutoff = 5.0

distance_function = 'gaussian'

edge_channels = 128

global_attributes = False

grid_resolution = None

hidden_channels = 128

hidden_size = 128

in_node_nf

kwargs

lmax = 2

mlp_type = 'grid'

mmax = 2

norm_type = 'rms_norm_sh'

num_distance_basis = 512

num_layers = 2

sphere_embedding_type = 'mixed'

task_type

use_envelope = False

MolecularDiffusion.runmodes.train.tasks_esen.logger