MolecularDiffusion.modules.models.oareactdiff.dynamics

The denoiser wrapper: three per-object encoder/decoder pairs + LEFTNet.

Ported from oa_reactdiff/dynamics/_base.py and oa_reactdiff/dynamics/egnn_dynamics.py (commit 543aaa8, MIT).

What this layer actually does, since the name does not say it: the reactant, transition state and product each get their own encoder MLP and their own decoder MLP (fragment_names = ["R", "TS", "P"] fixes the order and hence which weights belong to which object). Their encoded node features are then concatenated into one flat tensor, LEFTNet message-passes over the whole thing with subgraph_mask telling it which edges are intra-object, and the predicted displacement is split back apart and centred per object – which is what keeps the three SE(3) frames independent.

Deviations from upstream, both non-behavioural:

  • Upstream’s EGNN backbone is not ported (OA-ReactDiff’s released model is LEFTNet-only), so LEFTNet is the default model=. The factory passes it explicitly anyway.

  • Four unreachable methods are dropped: enpose_pbc (its only call site is commented out at egnn_dynamics.py:170) and the edge-attribute rebuild trio adjust_edge_attr_on_new_eij / create_new_edge_attr / init_edge_attr. The latter need edge features, and this model has none – edge_nf is 0 in the checkpoint and every self.dynamics(...) call in en_diffusion.py passes edge_attr=None literally. The edge_encoder / edge_decoder build path is kept as-is.

Classes

BaseDynamics

Base dynamics class set up for denoising process.

EGNNDynamics

Base dynamics class set up for denoising process.

Module Contents

class MolecularDiffusion.modules.models.oareactdiff.dynamics.BaseDynamics(model_config: Dict, fragment_names: List[str], node_nfs: List[int], edge_nf: int, condition_nf: int = 0, pos_dim: int = 3, update_pocket_coords: bool = True, condition_time: bool = True, edge_cutoff: float | None = None, model: torch.nn.Module | None = LEFTNet, device: torch.device = torch.device('cuda'), enforce_same_encoding: List | None = None, source: Dict | None = None)

Bases: torch.nn.Module

Base dynamics class set up for denoising process.

Parameters:

  • model_config (Dict) – config for the equivariant model.

  • fragment_names (List[str]) – list of names for fragments

  • node_nfs (List[int]) – list of number of input node attributues.

  • edge_nf (int) – number of input edge attributes.

  • condition_nf (int) – number of attributes for conditional generation.

  • 0. (Defaults to)

  • pos_dim (int) – dimension for position vector. Defaults to 3.

  • update_pocket_coords (bool) – whether to update positions of everything. Defaults to True.

  • condition_time (bool) – whether to condition on time. Defaults to True.

  • edge_cutoff (Optional[float]) – cutoff for building intra-fragment edges. Defaults to None.

  • model (Optional[nn.Module]) – Module for equivariant model. Defaults to None.

build_encoders_decoders(enfoce_name_encoding: List | None = None, source: Dict | None = None)

Build encoders and decoders for nodes and edges.

abstractmethod forward()
condition_nf = 0
condition_time = True
device
dist_dim = 0
edge_cutoff = None
edge_embed_dim
edge_nf
embed_dim
fragment_names
model
model_config
node_nfs
pos_dim = 3
update_pocket_coords = True
class MolecularDiffusion.modules.models.oareactdiff.dynamics.EGNNDynamics(model_config: Dict, fragment_names: List[str], node_nfs: List[int], edge_nf: int, condition_nf: int = 0, pos_dim: int = 3, update_pocket_coords: bool = True, condition_time: bool = True, edge_cutoff: float | None = None, model: torch.nn.Module | None = LEFTNet, device: torch.device = torch.device('cuda'), enforce_same_encoding: List | None = None, source: Dict | None = None)

Bases: BaseDynamics

Base dynamics class set up for denoising process.

Parameters:

  • model_config (Dict) – config for the equivariant model.

  • fragment_names (List[str]) – list of names for fragments

  • node_nfs (List[int]) – list of number of input node attributues.

  • edge_nf (int) – number of input edge attributes.

  • condition_nf (int) – number of attributes for conditional generation.

  • 0. (Defaults to)

  • pos_dim (int) – dimension for position vector. Defaults to 3.

  • update_pocket_coords (bool) – whether to update positions of everything. Defaults to True.

  • condition_time (bool) – whether to condition on time. Defaults to True.

  • edge_cutoff (Optional[float]) – cutoff for building intra-fragment edges. Defaults to None.

  • model (Optional[nn.Module]) – Module for equivariant model. Defaults to None.

static compute_frag_index(n_frag_switch: torch.Tensor) → numpy.ndarray
forward(xh: List[torch.Tensor], edge_index: torch.Tensor, t: torch.Tensor, conditions: torch.Tensor, n_frag_switch: torch.Tensor, combined_mask: torch.Tensor, edge_attr: torch.Tensor | None = None) → Tuple[List[torch.Tensor], torch.Tensor]

predict noise /mu.

Parameters:

  • xh (List[Tensor]) – list of concatenated tensors for pos and h

  • edge_index (Tensor) – [n_edge, 2]

  • t (Tensor) – time tensor. If dim is 1, same for all samples; otherwise different t for different samples

  • conditions (Tensor) – condition tensors

  • n_frag_switch (Tensor) – [n_nodes], fragment index for each nodes

  • combined_mask (Tensor) – [n_nodes], sample index for each node

  • edge_attr (Optional[Tensor]) – [n_edge, dim_edge_attribute]. Defaults to None.

Raises:

NotImplementedError – The fragement-position-fixed mode is not implement.

Returns:

updated pos-h and edge attributes

Return type:

Tuple[List[Tensor], Tensor]

static remove_mean_batch(x, indices)