MolecularDiffusion.modules.models.esen¶

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes¶

`GaussianFeatureEmbedding`
`Linear_Force_Head`
`MLP_EFS_Head`
`MLP_Energy_Head`
`eSEN_Backbone`	Mixin Model class implementing some general convenience properties and methods.
`eSEN_Encoder`	eSEN Encoder for VAE tasks.
`eSEN_VelocityFFN`	Native eSEN equivariant FFN for coordinate updates.
`eSEN_dynamics`	Mixin Model class implementing some general convenience properties and methods.

Module Contents¶

class MolecularDiffusion.modules.models.esen.GaussianFeatureEmbedding(num_channels, num_basis=20, start=0.0, stop=10.0, trainable=False)¶

Bases: torch.nn.Module

forward(x)¶

num_basis = 20¶

num_channels¶

trainable = False¶

class MolecularDiffusion.modules.models.esen.Linear_Force_Head(backbone)¶

Bases: torch.nn.Module, MolecularDiffusion.modules.layers.graph_module.HeadInterface

forward(data_dict, emb: dict[str, torch.Tensor])¶

Head forward.

Parameters:

data (AtomicData) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

linear¶

class MolecularDiffusion.modules.models.esen.MLP_EFS_Head(backbone)¶

Bases: torch.nn.Module, MolecularDiffusion.modules.layers.graph_module.HeadInterface

forward(data, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]¶

Head forward.

Parameters:

data (AtomicData) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

energy_block¶

hidden_channels¶

regress_forces¶

regress_stress¶

sphere_channels¶

class MolecularDiffusion.modules.models.esen.MLP_Energy_Head(backbone, reduce: str = 'sum')¶

Bases: torch.nn.Module, MolecularDiffusion.modules.layers.graph_module.HeadInterface

forward(data_dict, emb: dict[str, torch.Tensor])¶

Head forward.

Parameters:

data (AtomicData) – Atomic systems as input
emb (dict[str->torch.Tensor]) – Embeddings of the input as generated by the backbone

Returns:

outputs – Return one or more targets generated by this head

Return type:

dict[str->torch.Tensor]

energy_block¶

hidden_channels¶

reduce = 'sum'¶

sphere_channels¶

class MolecularDiffusion.modules.models.esen.eSEN_Backbone(max_num_elements: int = 100, sphere_channels: int = 128, lmax: int = 2, mmax: int = 2, grid_resolution: int | None = None, otf_graph: bool = False, max_neighbors: int = 300, use_pbc: bool = True, use_pbc_single: bool = False, cutoff: float = 5.0, edge_channels: int = 128, distance_function: str = 'gaussian', num_distance_basis: int = 512, direct_forces: bool = True, regress_forces: bool = True, regress_stress: bool = False, num_layers: int = 2, hidden_channels: int = 128, 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 = 'embedding', in_node_channels: int | None = None, update_position: bool = False, tanh: bool = False, coords_range: float = 10, normalization_factor: float = 1, aggregation_method: str = 'sum')¶

Bases: torch.nn.Module, MolecularDiffusion.modules.layers.graph_module.GraphModelMixin

Mixin Model class implementing some general convenience properties and methods.

csd_embedding(charge, spin, dataset)¶

forward(data) → dict[str, torch.Tensor]¶

generate_graph(*args, **kwargs)¶

get_rotmat_and_wigner(edge_distance_vecs)¶

no_weight_decay() → set¶: Returns a list of parameters with no weight decay.

SO3_grid¶

act_type = 's2'¶

activation_checkpointing = False¶

blocks¶

cutoff = 5.0¶

d_model = 1152¶

direct_forces = True¶

distance_function = 'gaussian'¶

edge_channels = 128¶

edge_channels_list¶

edge_degree_embedding¶

enforce_max_neighbors_strictly = False¶

global_attributes = False¶

grid_resolution = None¶

hidden_channels = 128¶

hidden_nf = 1152¶

in_node_channels = None¶

lmax = 2¶

mappingReduced¶

max_neighbors = 300¶

max_num_elements = 100¶

mlp_type = 'grid'¶

mmax = 2¶

norm¶

norm_type = 'rms_norm_sh'¶

num_distance_basis = 512¶

num_layers = 2¶

property num_params¶

otf_graph = False¶

regress_forces = True¶

regress_stress = False¶

sph_feature_size = 9¶

sphere_channels = 128¶

sphere_embedding_type = 'embedding'¶

update_position = False¶

use_envelope = False¶

use_pbc = True¶

use_pbc_single = False¶

class MolecularDiffusion.modules.models.esen.eSEN_Encoder(max_num_elements: int = 100, sphere_channels: int = 128, lmax: int = 2, mmax: int = 2, grid_resolution: int | None = None, max_neighbors: int = 300, cutoff: float = 5.0, edge_channels: int = 128, distance_function: str = 'gaussian', num_distance_basis: int = 512, num_layers: int = 2, hidden_channels: int = 128, norm_type: str = 'rms_norm_sh', act_type: str = 's2', mlp_type: str = 'grid', use_envelope: bool = False, activation_checkpointing: bool = False, sphere_embedding_type: str = 'embedding', **kwargs)¶

Bases: eSEN_Backbone

eSEN Encoder for VAE tasks.

Wraps eSEN_Backbone with fixed configuration for VAE encoding: - No force/stress regression - No position updates - Standard embedding - No global attributes - No OTF graph (assumes pre-computed or standard radius graph)

class MolecularDiffusion.modules.models.esen.eSEN_VelocityFFN(sphere_channels: int, hidden_channels: int, lmax: int)¶

Bases: torch.nn.Module

Native eSEN equivariant FFN for coordinate updates. Mirrors the equiformer FeedForwardNetwork pattern but operates on eSEN’s native [N, (lmax+1)^2, sphere_channels] tensors.

Forward: x_message -> SO3_Linear -> GateActivation -> SO3_Linear(out=1): -> l=1 slice [:, 1:4, 0] -> [N, 3] position delta

forward(x: torch.Tensor) → torch.Tensor¶

Parameters:: x – [N, (lmax+1)^2, sphere_channels]
Returns:: [N, 3]
Return type:: pos_delta

gate_act¶

gating_linear¶

lmax¶

so3_linear_1¶

so3_linear_2¶

class MolecularDiffusion.modules.models.esen.eSEN_dynamics(condition_time=True, context_node_nf=0, use_adapter_module: bool = False, max_num_elements: int = 100, sphere_channels: int = 128, lmax: int = 2, mmax: int = 2, grid_resolution: int | None = None, max_neighbors: int = 300, cutoff: float = 5.0, edge_channels: int = 128, distance_function: str = 'gaussian', num_distance_basis: int = 512, num_layers: int = 2, hidden_channels: int = 128, norm_type: str = 'rms_norm_sh', act_type: str = 's2', mlp_type: str = 'grid', use_envelope: bool = False, activation_checkpointing: bool = False, sphere_embedding_type: str = 'embedding', in_node_channels: int | None = None, tanh: bool = False, coords_range: float = 10, normalization_factor: float = 1, n_mlp_layers_out: int = 3, aggregation_method: str = 'sum', otf_graph: bool = False)¶

Bases: torch.nn.Module, MolecularDiffusion.modules.layers.graph_module.GraphModelMixin

Mixin Model class implementing some general convenience properties and methods.

condition_time = True¶

context_node_nf = 0¶

egnn¶

embedding_out¶

n_dims = 3¶

sph_feature_size = 9¶

use_adapter_module = False¶