"""
layers.py
~~~~~~~~~
E(3)-equivariant message-passing blocks:
* uses `torch_scatter.scatter` for edge->node aggregation
* dropout is `e3nn.nn.Dropout` (acts on *all* irrep coeffs)
* normalisation + activation selected via `make_nonlinearity`
"""
from __future__ import annotations
import torch
from torch import nn
from torch_scatter import scatter
from collections import OrderedDict
from e3nn.o3 import Irreps
from e3nn.o3 import FullyConnectedTensorProduct
from e3nn.nn import Dropout
from e3nn.nn import Gate
from net.common import Config, E3MLP, RadialMLP, SeparateWeightTensorProduct
from net.activations import make_nonlinearity
from net.layer_norm import E3LayerNorm
def projected_attention_reduce(
node,
edge_messages,
edge_length_emb,
dst,
query_proj,
key_proj,
value_projs,
num_heads,
head_dim,
):
"""Shared invariant projected multihead attention reduction."""
query = query_proj(node).reshape(node.shape[0], num_heads, head_dim)
key_input = torch.cat([edge_length_emb, edge_messages], dim=-1)
key = key_proj(key_input).reshape(edge_messages.shape[0], num_heads, head_dim)
values = torch.stack([proj(edge_messages) for proj in value_projs], dim=1)
scores = (query[dst] * key).sum(dim=-1) / float(head_dim) ** 0.5
max_per_dst = scatter(scores, dst, dim=0, dim_size=node.size(0), reduce="max")
scores = scores - max_per_dst.index_select(0, dst)
exp_scores = torch.exp(scores)
denom = scatter(exp_scores, dst, dim=0, dim_size=node.size(0), reduce="sum")
weights = exp_scores / denom.index_select(0, dst).clamp_min(1e-12)
weighted_values = values * weights.unsqueeze(-1)
head_outputs = scatter(
weighted_values,
dst,
dim=0,
dim_size=node.size(0),
reduce="sum",
)
return head_outputs.mean(dim=1)
def _assert_irreps_equal(a: Irreps, b: Irreps, *, what: str) -> None:
if Irreps(a).simplify() != Irreps(b).simplify():
raise AssertionError(f"{what} requires matching irreps, got {a} vs {b}")
def _magnitude_splits(
features: torch.Tensor,
irreps: Irreps,
) -> dict[str, torch.Tensor]:
"""
Split features by irrep and compute magnitude per irreducible component.
Returns a mapping from angular momentum l to a 1D tensor of magnitudes.
"""
mags: dict[str, torch.Tensor] = OrderedDict()
# features: (M, D)
start = 0
for mul, ir in irreps:
dim = ir.dim
size = mul * dim
# slice for this irrep
chunk = features[:, start : start + size]
# reshape to (M * mul, dim)
if mul > 0 and dim > 0:
reshaped = chunk.reshape(-1, dim)
# magnitude across dim
mag = torch.linalg.norm(reshaped, dim=1)
mags[f"{mul}x{ir.l}{'e' if ir.p == 1 else 'o'}"] = mag
start += size
return mags
def _tp_path_exists(irreps_in1: Irreps, irreps_in2: Irreps, ir_out) -> bool:
"""Check if a tensor product path exists for given irreps."""
from e3nn.o3 import Irrep
irreps_in1 = Irreps(irreps_in1)
irreps_in2 = Irreps(irreps_in2)
# Handle both Irrep and Irreps
if isinstance(ir_out, Irrep):
target_irreps = [ir_out]
elif isinstance(ir_out, str):
# Parse string like "1o" as single Irrep
try:
target_irreps = [Irrep(ir_out)]
except (ValueError, TypeError):
# If it fails, treat as Irreps
target_irreps = [ir for _, ir in Irreps(ir_out)]
else:
# Assume it's Irreps
target_irreps = [ir for _, ir in Irreps(ir_out)]
# Check if ANY of the target irreps can be produced
for target_ir in target_irreps:
for _, ir1 in irreps_in1:
for _, ir2 in irreps_in2:
if target_ir in ir1 * ir2:
return True
return False
def _tensor_product_expected_dims(tp: nn.Module) -> tuple[int, int]:
"""
Return the input feature dimensions expected by a tensor-product module.
This works for both e3nn's native tensor products and the
SeparateWeightTensorProduct wrapper used in this codebase.
"""
inner_tp = getattr(tp, "tp", tp)
in1_dim = getattr(inner_tp, "_in1_dim", None)
in2_dim = getattr(inner_tp, "_in2_dim", None)
if in1_dim is None or in2_dim is None:
return -1, -1
return int(in1_dim), int(in2_dim)
# ════════════════════════════════════════════════════════════════════════
# EquiConv - E(3)-equivariant convolution with radial weighting
# ════════════════════════════════════════════════════════════════════════
[docs]
class EquiConv(nn.Module):
"""
E(3)-equivariant convolution with radial MLP weighting.
Similar to DeepH-E3's EquiConv. Combines:
- Tensor product: irreps_in1 (x) irreps_in2 -> irreps_out
- Radial MLP: generates weights from edge length embeddings
- Optional Gate nonlinearity
- Element-wise multiplication with learned weights
Args:
n_radial: Input dimension for radial MLP (edge length embedding size)
irreps_in1: Input irreps for geometric features (node/edge concatenation)
irreps_in2: Input irreps for spherical harmonics
irreps_out: Output irreps
cfg: Configuration object
nonlin: Whether to use Gate nonlinearity after TP
"""
def __init__(
self,
n_radial: int,
irreps_in1: Irreps,
irreps_in2: Irreps,
irreps_out: Irreps,
cfg: Config,
nonlin: bool = True,
info: dict | None = None,
n_edge_types: int | None = None,
):
super().__init__()
self.cfg = cfg
self.info = info or {}
self.pair_conditioned_radial_mlp = bool(
getattr(self.cfg, "pair_conditioned_radial_mlp", False)
)
self.pair_emb = None
radial_in_dim = n_radial
if self.pair_conditioned_radial_mlp:
if n_edge_types is None:
raise ValueError("pair_conditioned_radial_mlp requires n_edge_types.")
self.pair_emb = nn.Embedding(
n_edge_types,
self.cfg.edge_type_emb_dim,
dtype=self.cfg.dtype,
)
nn.init.normal_(self.pair_emb.weight, std=0.2)
radial_in_dim += self.cfg.edge_type_emb_dim
irreps_in1 = Irreps(irreps_in1)
irreps_in2 = Irreps(irreps_in2)
irreps_out = Irreps(irreps_out)
# Determine if we need a gate nonlinearity
self.nonlin = None
if nonlin:
# Separate scalars and gated irreps
irreps_scalars = Irreps(
[
(mul, ir)
for mul, ir in irreps_out
if ir.l == 0 and _tp_path_exists(irreps_in1, irreps_in2, ir)
]
).simplify()
irreps_gated = Irreps(
[
(mul, ir)
for mul, ir in irreps_out
if ir.l > 0 and _tp_path_exists(irreps_in1, irreps_in2, ir)
]
).simplify()
if irreps_gated.dim > 0:
# Need gates (scalars) for gated tensors
if _tp_path_exists(irreps_in1, irreps_in2, "0e"):
ir_gate = "0e"
elif _tp_path_exists(irreps_in1, irreps_in2, "0o"):
ir_gate = "0o"
else:
raise ValueError(
f"Cannot produce gates for irreps_gated={irreps_gated}"
)
irreps_gates = Irreps(
[(mul, ir_gate) for mul, _ in irreps_gated]
).simplify()
else:
irreps_gates = Irreps([])
# Create gate nonlinearity
act_scalars = []
for _, ir in irreps_scalars:
if ir.p == 1: # even parity
act_scalars.append(torch.nn.functional.silu)
else: # odd parity
act_scalars.append(torch.tanh)
act_gates = []
for _, ir in irreps_gates:
if ir.p == 1: # even parity
act_gates.append(torch.sigmoid)
else: # odd parity
act_gates.append(torch.tanh)
self.nonlin = Gate(
irreps_scalars, act_scalars, irreps_gates, act_gates, irreps_gated
)
irreps_tp_out = self.nonlin.irreps_in
else:
irreps_tp_out = Irreps(
[
(mul, ir)
for mul, ir in irreps_out
if _tp_path_exists(irreps_in1, irreps_in2, ir)
]
)
# Create tensor product
if cfg.tp_type == "separate_weight":
self.tp = SeparateWeightTensorProduct(irreps_in1, irreps_in2, irreps_tp_out)
elif cfg.tp_type == "fully_connected":
self.tp = FullyConnectedTensorProduct(
irreps_in1,
irreps_in2,
irreps_tp_out,
internal_weights=True,
shared_weights=True,
)
else:
raise ValueError(f"Unknown tp_type: {cfg.tp_type}")
# Determine output irreps
if nonlin:
self.irreps_out = self.nonlin.irreps_out
else:
self.irreps_out = irreps_tp_out
# Keep the historical output width in both modes so existing checkpoint
# shapes and initialization are unchanged. ``term`` shares a gate across
# every multiplicity in a grouped term; ``copy`` uses all emitted gates.
self.radial_weight_granularity = str(cfg.radial_weight_granularity).lower()
if self.radial_weight_granularity not in {"term", "copy"}:
raise ValueError("radial_weight_granularity must be 'term' or 'copy'.")
self.radial_mlp = RadialMLP(
in_dim=radial_in_dim,
out_dim=self.irreps_out.num_irreps,
layers=self.cfg.radial_layers,
act="silu",
dtype=self.cfg.dtype,
)
radial_weight_index = []
copy_offset = 0
for idx, (mul, ir) in enumerate(self.irreps_out):
if self.radial_weight_granularity == "term":
radial_weight_index.extend([idx] * (mul * ir.dim))
else:
for copy in range(mul):
radial_weight_index.extend([copy_offset + copy] * ir.dim)
copy_offset += mul
self.register_buffer(
"_radial_weight_index",
torch.tensor(radial_weight_index, dtype=torch.long),
persistent=False,
)
def forward(
self,
fea_in1: torch.Tensor,
fea_in2: torch.Tensor,
edge_length_emb: torch.Tensor,
edge_type_idx: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Args:
fea_in1: Geometric features (batch, irreps_in1.dim)
fea_in2: Spherical harmonics (batch, irreps_in2.dim)
edge_length_emb: Radial embeddings (batch, n_radial)
Returns:
Tensor of shape (batch, irreps_out.dim)
"""
expected_in1_dim, expected_in2_dim = _tensor_product_expected_dims(self.tp)
if expected_in1_dim >= 0 and fea_in1.shape[-1] != expected_in1_dim:
print(
"[TP SHAPE MISMATCH] "
f"{self.info=} "
f"fea_in1.shape={tuple(fea_in1.shape)} "
f"fea_in2.shape={tuple(fea_in2.shape)} "
f"edge_length_emb.shape={tuple(edge_length_emb.shape)} "
f"expected_in1_dim={expected_in1_dim} "
f"expected_in2_dim={expected_in2_dim} "
f"tp={type(self.tp).__name__} "
f"irreps_in1={self.tp.tp.irreps_in1 if hasattr(self.tp, 'tp') else self.tp.irreps_in1} "
f"irreps_in2={self.tp.tp.irreps_in2 if hasattr(self.tp, 'tp') else self.tp.irreps_in2} "
f"irreps_out={self.tp.tp.irreps_out if hasattr(self.tp, 'tp') else self.tp.irreps_out}",
flush=True,
)
elif expected_in2_dim >= 0 and fea_in2.shape[-1] != expected_in2_dim:
print(
"[TP SHAPE MISMATCH] "
f"{self.info=} "
f"fea_in1.shape={tuple(fea_in1.shape)} "
f"fea_in2.shape={tuple(fea_in2.shape)} "
f"edge_length_emb.shape={tuple(edge_length_emb.shape)} "
f"expected_in1_dim={expected_in1_dim} "
f"expected_in2_dim={expected_in2_dim} "
f"tp={type(self.tp).__name__} "
f"irreps_in1={self.tp.tp.irreps_in1 if hasattr(self.tp, 'tp') else self.tp.irreps_in1} "
f"irreps_in2={self.tp.tp.irreps_in2 if hasattr(self.tp, 'tp') else self.tp.irreps_in2} "
f"irreps_out={self.tp.tp.irreps_out if hasattr(self.tp, 'tp') else self.tp.irreps_out}",
flush=True,
)
# Tensor product
z = self.tp(fea_in1, fea_in2)
# Handle empty output
if z.shape[-1] == 0:
return z
# Apply gate nonlinearity if present
if self.nonlin is not None:
z = self.nonlin(z)
# Radial weighting (element-wise multiplication)
if self.pair_conditioned_radial_mlp:
if edge_type_idx is None:
raise ValueError(
"pair_conditioned_radial_mlp requires edge_type_idx in forward()."
)
pair_emb = self.pair_emb(edge_type_idx)
edge_length_emb = torch.cat([edge_length_emb, pair_emb], dim=-1)
weights = self.radial_mlp(edge_length_emb)
if self._radial_weight_index.numel() == 0:
return torch.zeros(z.shape[0], 0, dtype=z.dtype, device=z.device)
component_weights = weights.index_select(1, self._radial_weight_index)
return z * component_weights
# ════════════════════════════════════════════════════════════════════════
# Edge update
# ════════════════════════════════════════════════════════════════════════
[docs]
class EdgeUpdateBlock(nn.Module):
"""
Edge update block following DeepH-E3 architecture.
For each edge (i -> j):
• Concatenate node[i], node[j], and edge features
• Apply EquiConv with spherical harmonics and radial weighting
• Apply linear transformation and nonlinearity
• Optional self-connection and residual
"""
def __init__(
self,
node_irreps: Irreps,
edge_irreps: Irreps,
num_species: int,
cfg: Config,
info: dict = None,
edge_irreps_out: Irreps | None = None,
n_edge_types: int | None = None,
):
super().__init__()
self.cfg = cfg
self.info = info
self.node_irreps = node_irreps
self.edge_irreps = edge_irreps
self.node_combine_mode = str(cfg.edge_update_node_combine).lower()
if edge_irreps_out is None:
edge_irreps_out = edge_irreps
self.edge_irreps_out = edge_irreps_out
if self.node_combine_mode not in {"concat", "sum", "tensor_product"}:
raise ValueError(
"edge_update_node_combine must be one of 'concat', 'sum', or "
f"'tensor_product', got {cfg.edge_update_node_combine!r}."
)
# Resolve sh_irreps and n_radial from config
sh_irreps = Irreps.spherical_harmonics(cfg.l_max)
n_radial = cfg.n_radial
# Pre-linear transformation
self.lin_pre = E3MLP(
edge_irreps,
edge_irreps,
edge_irreps,
1, # Single layer
cfg,
activate_last=False,
)
self.node_pair_tp = None
if self.node_combine_mode == "concat":
self.node_context_irreps = (node_irreps + node_irreps).simplify()
elif self.node_combine_mode == "sum":
_assert_irreps_equal(
node_irreps,
node_irreps,
what="edge_update_node_combine='sum'",
)
self.node_context_irreps = Irreps(node_irreps).simplify()
else:
self.node_pair_tp = FullyConnectedTensorProduct(
node_irreps,
node_irreps,
node_irreps,
internal_weights=True,
shared_weights=True,
)
self.node_context_irreps = Irreps(node_irreps).simplify()
_assert_irreps_equal(
self.node_context_irreps,
node_irreps,
what="edge_update_node_combine='tensor_product'",
)
# Edge convolution always sees edge features plus a node-pair context.
irreps_in1 = (self.node_context_irreps + edge_irreps).simplify()
irreps_in2 = sh_irreps
# EquiConv: main tensor product with radial weighting
self.conv = EquiConv(
n_radial=n_radial,
irreps_in1=irreps_in1,
irreps_in2=irreps_in2,
irreps_out=edge_irreps_out,
cfg=cfg,
nonlin=True, # Use gate nonlinearity
info=info,
n_edge_types=n_edge_types,
)
# Post-linear transformation
self.lin_post = E3MLP(
self.conv.irreps_out,
self.conv.irreps_out,
self.conv.irreps_out,
1,
cfg,
activate_last=False,
)
# Self-connection (edge-type specific)
self.sc = None
if cfg.use_self_connection:
# num_species^2 edge types (all pairs of elements)
n_edge_types = num_species * num_species
self.sc = FullyConnectedTensorProduct(
edge_irreps,
Irreps(f"{n_edge_types}x0e"),
self.conv.irreps_out,
internal_weights=True,
shared_weights=True,
)
# Normalization and activation
self.norm_act = make_nonlinearity(self.conv.irreps_out, cfg)
self.layer_norm = E3LayerNorm(self.conv.irreps_out) if cfg.e3layernorm else None
# Dropout
if cfg.dropout > 0.0:
self.dropout = Dropout(self.conv.irreps_out, p=cfg.dropout)
else:
self.dropout = None
internal_variant = cfg.internal_e3mlp_variant or cfg.e3mlp_variant
if cfg.internal_e3mlp_layers > 0:
self.refine = E3MLP(
self.conv.irreps_out,
edge_irreps_out,
edge_irreps_out,
cfg.internal_e3mlp_layers,
cfg,
activate_last=False,
variant=internal_variant,
)
else:
self.refine = nn.Identity()
# Downstream blocks consume the post-refine representation, not the raw
# convolution output. Using the pre-refine irreps here causes later
# tensor products to be built for the wrong input width.
self.irreps_out = edge_irreps_out
def forward(
self,
node: torch.Tensor,
edge: torch.Tensor,
edge_index: torch.Tensor,
edge_sh: torch.Tensor,
edge_length_emb: torch.Tensor,
edge_one_hot: torch.Tensor = None,
edge_type_idx: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Args:
node: Node features (N, node_irreps.dim)
edge: Edge features (E, edge_irreps.dim)
edge_index: Edge indices (2, E)
edge_sh: Spherical harmonics (E, sh_irreps.dim)
edge_length_emb: Radial embeddings (E, n_radial)
edge_one_hot: Edge type one-hot (E, num_species^2) for self-connection
Returns:
Updated edge features (E, irreps_out.dim)
"""
edge_old = edge
# Self-connection
if self.sc is not None and edge_one_hot is not None:
edge_self_connection = self.sc(edge, edge_one_hot)
# Pre-linear
edge = self.lin_pre(edge)
# Concatenate node features
src, dst = edge_index
if self.node_combine_mode == "concat":
node_context = torch.cat([node[src], node[dst]], dim=-1)
elif self.node_combine_mode == "sum":
node_context = node[src] + node[dst]
else:
node_context = self.node_pair_tp(node[src], node[dst])
fea_in = torch.cat([node_context, edge], dim=-1)
# EquiConv
edge = self.conv(fea_in, edge_sh, edge_length_emb, edge_type_idx=edge_type_idx)
# Post-linear
edge = self.lin_post(edge)
# Add self-connection
if self.sc is not None and edge_one_hot is not None:
edge = edge + edge_self_connection
# Normalization and activation
edge = self.norm_act(edge)
if self.layer_norm is not None:
edge = self.layer_norm(edge)
# Dropout
if self.dropout:
edge = self.dropout(edge)
edge = self.refine(edge)
# Residual connection
if self.cfg.edge_update_residual and edge.shape == edge_old.shape:
edge = edge + edge_old
return edge
# ════════════════════════════════════════════════════════════════════════
# Node update
# ════════════════════════════════════════════════════════════════════════
[docs]
class NodeUpdateBlock(nn.Module):
"""
Node update block following DeepH-E3 architecture.
Aggregates edge messages to nodes:
• Concatenate node[i], node[j], and edge features
• Apply EquiConv to create edge messages
• Scatter-aggregate messages to destination nodes
• Apply transformations and optional self-connection
"""
def __init__(
self,
node_irreps: Irreps,
edge_irreps: Irreps,
num_species: int,
cfg: Config,
info: dict = None,
node_irreps_out: Irreps = None, # Output irreps for nodes (defaults to edge_irreps)
n_edge_types: int | None = None,
):
super().__init__()
self.cfg = cfg
self.info = info
self.node_irreps = node_irreps
self.edge_irreps = edge_irreps
self.message_agg_mode = str(cfg.node_update_message_agg).lower()
if self.message_agg_mode not in {"sum", "average", "attention"}:
raise ValueError(
"node_update_message_agg must be one of 'sum', 'average', or "
f"'attention', got {cfg.node_update_message_agg!r}."
)
self.attn_scalar_dim = int(cfg.node_update_attention_scalar_dim)
if int(cfg.node_update_attention_heads) <= 0:
raise ValueError("node_update_attention_heads must be > 0.")
if self.attn_scalar_dim <= 0:
raise ValueError("node_update_attention_scalar_dim must be > 0.")
if self.attn_scalar_dim % int(cfg.node_update_attention_heads) != 0:
raise ValueError(
"node_update_attention_scalar_dim must be divisible by "
"node_update_attention_heads."
)
# Default: output nodes with same irreps as edges
if node_irreps_out is None:
node_irreps_out = edge_irreps
# Resolve sh_irreps and n_radial from config
sh_irreps = Irreps.spherical_harmonics(cfg.l_max)
n_radial = cfg.n_radial
# Pre-linear transformation
self.lin_pre = E3MLP(
node_irreps,
node_irreps,
node_irreps,
1,
cfg,
activate_last=False,
)
# Concatenated input: node[i] + node[j] + edge
irreps_in1 = (node_irreps + node_irreps + edge_irreps).simplify()
irreps_in2 = sh_irreps
# EquiConv: main tensor product with radial weighting
self.conv = EquiConv(
n_radial=n_radial,
irreps_in1=irreps_in1,
irreps_in2=irreps_in2,
irreps_out=node_irreps_out,
cfg=cfg,
nonlin=True,
info=info,
n_edge_types=n_edge_types,
)
# Post-linear transformation
self.lin_post = E3MLP(
self.conv.irreps_out,
self.conv.irreps_out,
self.conv.irreps_out,
1,
cfg,
activate_last=False,
)
# Self-connection (node-type specific)
self.sc = None
if cfg.use_self_connection:
self.sc = FullyConnectedTensorProduct(
node_irreps,
Irreps(f"{num_species}x0e"),
self.conv.irreps_out,
internal_weights=True,
shared_weights=True,
) # Normalization and activation
self.norm_act = make_nonlinearity(self.conv.irreps_out, cfg)
self.layer_norm = E3LayerNorm(self.conv.irreps_out) if cfg.e3layernorm else None
# Dropout
if cfg.dropout > 0.0:
self.dropout = Dropout(self.conv.irreps_out, p=cfg.dropout)
else:
self.dropout = None
self.attn_num_heads = int(cfg.node_update_attention_heads)
self.attn_head_dim = self.attn_scalar_dim // self.attn_num_heads
self.query_proj = None
self.key_proj = None
self.value_projs = None
if self.message_agg_mode == "attention":
attn_irreps = Irreps(f"{self.attn_scalar_dim}x0e")
attn_geom_irreps = Irreps(f"{n_radial}x0e")
attn_key_input_irreps = (attn_geom_irreps + self.conv.irreps_out).simplify()
self.query_proj = E3MLP(
node_irreps,
Irreps("64x0e"),
attn_irreps,
1,
cfg,
activate_last=False,
)
self.key_proj = E3MLP(
attn_key_input_irreps,
Irreps("64x0e"),
attn_irreps,
1,
cfg,
activate_last=False,
)
self.value_projs = nn.ModuleList(
[
E3MLP(
self.conv.irreps_out,
self.conv.irreps_out,
self.conv.irreps_out,
1,
cfg,
activate_last=False,
)
for _ in range(self.attn_num_heads)
]
)
internal_variant = cfg.internal_e3mlp_variant or cfg.e3mlp_variant
if cfg.internal_e3mlp_layers > 0:
self.refine = E3MLP(
self.conv.irreps_out,
node_irreps_out,
node_irreps_out,
cfg.internal_e3mlp_layers,
cfg,
activate_last=False,
variant=internal_variant,
)
else:
self.refine = nn.Identity()
# Downstream blocks consume the post-refine representation, not the raw
# convolution output. Using the pre-refine irreps here causes later
# tensor products to be built for the wrong input width.
self.irreps_out = node_irreps_out
def forward(
self,
node: torch.Tensor,
edge: torch.Tensor,
edge_index: torch.Tensor,
edge_sh: torch.Tensor,
edge_length_emb: torch.Tensor,
node_one_hot: torch.Tensor = None,
edge_type_idx: torch.Tensor | None = None,
) -> torch.Tensor:
"""
Args:
node: Node features (N, node_irreps.dim)
edge: Edge features (E, edge_irreps.dim)
edge_index: Edge indices (2, E)
edge_sh: Spherical harmonics (E, sh_irreps.dim)
edge_length_emb: Radial embeddings (E, n_radial)
node_one_hot: Node type one-hot (N, num_species) for self-connection
Returns:
Updated node features (N, irreps_out.dim)
"""
node_old = node
# Self-connection
if self.sc is not None and node_one_hot is not None:
node_self_connection = self.sc(node, node_one_hot)
# Pre-linear
node = self.lin_pre(node)
# Create edge messages
src, dst = edge_index
fea_in = torch.cat([node[src], node[dst], edge], dim=-1)
# EquiConv to create messages
edge_messages = self.conv(
fea_in, edge_sh, edge_length_emb, edge_type_idx=edge_type_idx
)
# Aggregate messages to nodes
if self.message_agg_mode == "sum":
node = scatter(
edge_messages, dst, dim=0, dim_size=node_old.size(0), reduce="sum"
)
elif self.message_agg_mode == "average":
node = scatter(
edge_messages, dst, dim=0, dim_size=node_old.size(0), reduce="mean"
)
else:
node = projected_attention_reduce(
node,
edge_messages,
edge_length_emb,
dst,
self.query_proj,
self.key_proj,
self.value_projs,
self.attn_num_heads,
self.attn_head_dim,
)
# Post-linear
node = self.lin_post(node)
# Add self-connection
if self.sc is not None and node_one_hot is not None:
node = node + node_self_connection
# Normalization and activation
node = self.norm_act(node)
if self.layer_norm is not None:
node = self.layer_norm(node)
# Dropout
if self.dropout:
node = self.dropout(node)
node = self.refine(node)
# Residual connection (only if dimensions match)
if self.cfg.node_update_residual and node.shape == node_old.shape:
node = node + node_old
return node
# ════════════════════════════════════════════════════════════════════════
# Message block
# ════════════════════════════════════════════════════════════════════════
class MessageBlock(nn.Module):
"""
One message-passing step: node update followed by edge update.
Combines EdgeUpdateBlock and NodeUpdateBlock following DeepH-E3.
Args:
node: Tensor of shape (N, hidden_dim)
edge: Tensor of shape (E, hidden_dim)
edge_index: LongTensor of shape (2, E) with source/dest indices
edge_sh: Spherical harmonics (E, sh_irreps.dim)
edge_length_emb: Radial embeddings (E, n_radial)
node_one_hot: Node type one-hot (N, num_species)
edge_one_hot: Edge type one-hot (E, num_species^2)
Returns:
Tuple (node_updated, edge_updated) of same shapes.
"""
def __init__(
self,
node_irreps: Irreps,
edge_irreps: Irreps,
num_species: int,
cfg: Config,
info: dict = None,
node_irreps_out: Irreps | None = None,
edge_irreps_out: Irreps | None = None,
n_edge_types: int | None = None,
):
super().__init__()
self.cfg = cfg
self.info = info or {}
self.node_irreps = node_irreps
self.edge_irreps = edge_irreps
self.node_upd = NodeUpdateBlock(
node_irreps=node_irreps,
edge_irreps=edge_irreps,
node_irreps_out=node_irreps_out,
num_species=num_species,
cfg=cfg,
info=info,
n_edge_types=n_edge_types,
)
self.edge_upd = EdgeUpdateBlock(
node_irreps=self.node_upd.irreps_out,
edge_irreps=edge_irreps,
edge_irreps_out=edge_irreps_out,
num_species=num_species,
cfg=cfg,
info=info,
n_edge_types=n_edge_types,
)
self.node_irreps_out = self.node_upd.irreps_out
self.edge_irreps_out = self.edge_upd.irreps_out
def forward(
self,
node: torch.Tensor,
edge: torch.Tensor,
edge_index: torch.Tensor,
edge_sh: torch.Tensor,
edge_length_emb: torch.Tensor,
node_one_hot: torch.Tensor = None,
edge_one_hot: torch.Tensor = None,
edge_type_idx: torch.Tensor | None = None,
activation_mags: dict = None,
):
"""
Forward pass through edge and node updates.
Args:
node: Node features (N, node_irreps.dim)
edge: Edge features (E, edge_irreps.dim)
edge_index: Edge indices (2, E)
edge_sh: Spherical harmonics (E, sh_irreps.dim)
edge_length_emb: Radial embeddings (E, n_radial)
node_one_hot: Node type one-hot (N, num_species)
edge_one_hot: Edge type one-hot (E, num_species^2)
activation_mags: Optional dict to track activation magnitudes
Returns:
Tuple of (updated_node, updated_edge)
"""
node = self.node_upd(
node,
edge,
edge_index,
edge_sh,
edge_length_emb,
node_one_hot,
edge_type_idx=edge_type_idx,
)
if activation_mags is not None and self.cfg.log_activation_mag and self.info:
prefix = f"mag_node_{self.info.get('graph', 'graph')}_layer_{self.info.get('layer', 0)}"
splits = _magnitude_splits(node, self.node_upd.irreps_out)
for ir_str, mag in splits.items():
tag = f"{prefix}_{ir_str}"
activation_mags[tag] = mag
edge = self.edge_upd(
node,
edge,
edge_index,
edge_sh,
edge_length_emb,
edge_one_hot,
edge_type_idx=edge_type_idx,
)
if activation_mags is not None and self.cfg.log_activation_mag and self.info:
prefix = f"mag_edge_{self.info.get('graph', 'graph')}_layer_{self.info.get('layer', 0)}"
splits = _magnitude_splits(edge, self.edge_upd.irreps_out)
for ir_str, mag in splits.items():
tag = f"{prefix}_{ir_str}"
activation_mags[tag] = mag
return node, edge