MolecularDiffusion.modules.models.oareactdiff.leftnet¶

LEFTNet – the object-aware equivariant backbone behind OA-ReactDiff.

Ported from oa_reactdiff/model/leftnet.py (commit 543aaa8, MIT, (c) 2023 Chenru Duan). The blocks below (RBFEmb, NeighborEmb, CFConvS2V, GCLMessage, EquiMessage, EquiUpdate, vector, EquiOutput) are used by nothing else in the tree, so they stay here rather than moving to modules/layers/.

object_aware=True is the paper’s contribution: subgraph_mask gates every message so intra-object and inter-object edges are treated differently, which is what lets one network denoise a reactant, a transition state and a product simultaneously while keeping each object’s own SE(3) frame.

Three deliberate deviations from the upstream file, none of them behavioural:

MLP / ACTIVATION_MAPPING / unsorted_segment_sum are inlined from oa_reactdiff/model/core.py and util_funcs.py instead of being imported, so the port is two files instead of four. dynamics.py imports MLP from here.
EquiMessage(...).jittable() lost its .jittable(): PyG >= 2.5 makes it a deprecated no-op that only emits a warning per layer.
torch.cross(a, b) became torch.cross(a, b, dim=1). Without dim torch picks the first axis of size 3, which is axis 1 for the (N, 3) tensors used here – except when N == 3, where it silently picks axis 0 and produces a non-equivariant result. Naming the axis is what upstream means everywhere and removes the deprecation warning.

Attributes¶

`ACTIVATION_MAPPING`
`EPS`

Classes¶

`CFConvS2V`	Scalar to vector.
`EquiMessage`
`EquiOutput`
`EquiUpdate`
`GCLMessage`
`GatedEquivariantBlock`	Gated Equivariant Block as defined in Schütt et al. (2021):
`LEFTNet`	LEFTNet
`MLP`	Multi-layer perceptron. (`oa_reactdiff/model/core.py`)
`NeighborEmb`	Initialize node features based on neighboring nodes.
`OneLayerActivation`	One layer NN with activation. (`oa_reactdiff/model/core.py`)
`RBFEmb`	radial basis function to embed distances
`vector`

Functions¶

`assert_rot_equiv`(func, dist, edge_index, pos)	Added by Chenru: test a func for constructing y1 is equivariant.
`com`(x)
`nn_vector`(dist, edge_index, pos)	Added by Chenru: Getting the nearest neighbor position to construct nodeframe.
`remove_mean_batch`(x, indices)
`swish`(x)
`unsorted_segment_sum`(data, segment_ids, num_segments, ...)	TensorFlow's `unsorted_segment_sum`. (`model/util_funcs.py`)

Module Contents¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.CFConvS2V(hid_dim: int)¶

Bases: torch_geometric.nn.conv.MessagePassing

Scalar to vector.

forward(s, v, edge_index, emb)¶

_summary_

Parameters:

s (_type_) – _description_, [n_atom, n_z, n_embed]
v (_type_) – _description_, [n_edge, n_pos, n_embed]
edge_index (_type_) – _description_, [2, n_edge]
emb (_type_) – _description_, [n_edge, n_embed]

Returns:

_description_

Return type:

_type_

message(x_j, norm)¶

hid_dim¶

lin1¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.EquiMessage(hidden_channels, num_radial, reflect_equiv)¶

Bases: torch_geometric.nn.conv.MessagePassing

aggregate(features: Tuple[torch.Tensor, torch.Tensor], index: torch.Tensor, ptr: torch.Tensor | None, dim_size: int | None) → Tuple[torch.Tensor, torch.Tensor]¶

forward(x, vec, edge_index, edge_rbf, weight, edge_vector, edge_cross)¶

message(xh_j, xh_i, vec_j, rbfh_ij, r_ij, edge_cross)¶

reset_parameters()¶

update(inputs: Tuple[torch.Tensor, torch.Tensor]) → Tuple[torch.Tensor, torch.Tensor]¶

dir_proj¶

hidden_channels¶

inv_sqrt_3¶

inv_sqrt_h¶

num_radial¶

rbf_proj¶

reflect_equiv¶

x_layernorm¶

x_proj¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.EquiOutput(hidden_channels, out_channels=1, single_layer_output=True)¶

Bases: torch.nn.Module

forward(x, vec)¶

reset_parameters()¶

hidden_channels¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.EquiUpdate(hidden_channels, reflect_equiv: bool = True)¶

Bases: torch.nn.Module

forward(x, vec, nodeframe)¶

reset_parameters()¶

hidden_channels¶

inv_sqrt_2¶

inv_sqrt_h¶

lin3¶

reflect_equiv = True¶

vec_proj¶

xvec_proj¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.GCLMessage(hidden_channels, num_radial, act_fn: str = 'swish', legacy: bool = False)¶

Bases: torch.nn.Module

edge_message(xh_i, xh_j, edgeh)¶

forward(x, edge_index, weight)¶

node_message(xh, edge_index, m_ij)¶

att_mlp¶

edge_mlp¶

edge_out_trans¶

node_mlp¶

x_layernorm¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.GatedEquivariantBlock(hidden_channels, out_channels)¶

Bases: torch.nn.Module

Gated Equivariant Block as defined in Schütt et al. (2021): Equivariant message passing for the prediction of tensorial properties and molecular spectra.

Borrowed from TorchMD-Net

forward(x, v)¶

reset_parameters()¶

act¶

out_channels¶

update_net¶

vec1_proj¶

vec2_proj¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.LEFTNet(pos_require_grad=False, cutoff=10.0, num_layers=4, hidden_channels=128, num_radial=96, in_hidden_channels: int = 8, reflect_equiv: bool = True, legacy: bool = True, update: bool = True, pos_grad: bool = False, single_layer_output: bool = True, for_conf: bool = False, ff: bool = False, object_aware: bool = True, **kwargs)¶

Bases: torch.nn.Module

LEFTNet

Parameters:

pos_require_grad (bool, optional) – If set to True, will require to take derivative of model output with respect to the atomic positions. (default: False)
cutoff (float, optional) – Cutoff distance for interatomic interactions. (default: 5.0)
num_layers (int, optional) – Number of building blocks. (default: 4)
hidden_channels (int, optional) – Hidden embedding size. (default: 128)
num_radial (int, optional) – Number of radial basis functions. (default: 96)
y_mean (float, optional) – Mean value of the labels of training data. (default: 0)
y_std (float, optional) – Standard deviation of the labels of training data. (default: 1)

static assemble_nodemask(edge_index: torch.Tensor, pos: torch.Tensor)¶

forward(h: torch.Tensor, pos: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor | None = None, node_mask: torch.Tensor | None = None, edge_mask: torch.Tensor | None = None, update_coords_mask: torch.Tensor | None = None, subgraph_mask: torch.Tensor | None = None)¶

reset_parameters()¶

scalarization(pos, edge_index)¶

cutoff = 10.0¶

embedding¶

embedding_out¶

ff = False¶

for_conf = False¶

gcl_layers¶

hidden_channels = 128¶

inv_sqrt_2¶

last_layer¶

legacy = True¶

lin3¶

message_layers¶

neighbor_emb¶

num_layers = 4¶

object_aware = True¶

out_pos¶

pos_expansion¶

pos_grad = False¶

pos_require_grad = False¶

radial_emb¶

radial_lin¶

reflect_equiv = True¶

s2v¶

update = True¶

update_layers¶

vec¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.MLP(in_dim: int, out_dims: list, bias: bool = True, activation: List[str | None] | str | None = 'swish', last_layer_no_activation: bool = False)¶

Bases: torch.nn.Module

Multi-layer perceptron. (oa_reactdiff/model/core.py)

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

mlp¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.NeighborEmb(hid_dim, in_hidden_channels=5)¶

Bases: torch_geometric.nn.conv.MessagePassing

Initialize node features based on neighboring nodes.

forward(z, s, edge_index, embs)¶

message(x_j, norm)¶

embedding¶

hid_dim¶

ln_emb¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.OneLayerActivation(in_dim: int, out_dim: int, bias: int = True, activation=Union[str, None])¶

Bases: torch.nn.Module

One layer NN with activation. (oa_reactdiff/model/core.py)

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

activation¶

linear¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.RBFEmb(num_rbf, rbound_upper, rbf_trainable=False)¶

Bases: torch.nn.Module

radial basis function to embed distances modified: delete cutoff with r

forward(dist)¶

reset_parameters()¶

num_rbf¶

rbf_trainable = False¶

rbound_lower = 0¶

rbound_upper¶

class MolecularDiffusion.modules.models.oareactdiff.leftnet.vector¶

Bases: torch_geometric.nn.conv.MessagePassing

forward(vector, edge_index)¶

MolecularDiffusion.modules.models.oareactdiff.leftnet.assert_rot_equiv(func: Callable, dist: torch.Tensor, edge_index: torch.Tensor, pos: torch.Tensor)¶

Added by Chenru: test a func for constructing y1 is equivariant.

Parameters:

func (Callable) – _description_
dist (Tensor) – _description_
edge_index (Tensor) – _description_
pos (Tensor) – _description_

MolecularDiffusion.modules.models.oareactdiff.leftnet.com(x)¶

MolecularDiffusion.modules.models.oareactdiff.leftnet.nn_vector(dist: torch.Tensor, edge_index: torch.Tensor, pos: torch.Tensor)¶

Added by Chenru: Getting the nearest neighbor position to construct nodeframe.

Parameters:

dist (Tensor) – (n_edge)
edge_index (Tensor) – (2, n_edge)
pos (Tensor) – (n_atom, 3)

Returns:

(n_atom, 3)

Return type:

Tensor

MolecularDiffusion.modules.models.oareactdiff.leftnet.remove_mean_batch(x, indices)¶

MolecularDiffusion.modules.models.oareactdiff.leftnet.swish(x)¶

MolecularDiffusion.modules.models.oareactdiff.leftnet.unsorted_segment_sum(data, segment_ids, num_segments, normalization_factor, aggregation_method: str)¶: TensorFlow’s unsorted_segment_sum. (model/util_funcs.py)

MolecularDiffusion.modules.models.oareactdiff.leftnet.ACTIVATION_MAPPING¶

MolecularDiffusion.modules.models.oareactdiff.leftnet.EPS = 1e-06¶