File size: 7,662 Bytes
3e02ab8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 | # This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
"""Interaction Block"""
import torch
from e3nn.o3._irreps import Irreps
from e3nn.o3._linear import Linear
from e3nn.o3._tensor_product._sub import FullyConnectedTensorProduct
from onescience.datapipes.materials.nequip import AtomicDataDict
from ._graph_mixin import GraphModuleMixin
from .mlp import ScalarMLPFunction
from ._ghost_exchange_base import NoOpGhostExchangeModule
from ._tp_scatter_base import TensorProductScatter
from .norm import AvgNumNeighborsNorm
from typing import Optional, Sequence, Union, Dict
class InteractionBlock(GraphModuleMixin, torch.nn.Module):
use_sc: bool
def __init__(
self,
irreps_in,
irreps_out,
radial_mlp_depth: int = 1,
radial_mlp_width: int = 8,
use_sc: bool = True,
is_first_layer: bool = False,
type_names: Optional[Sequence[str]] = None,
avg_num_neighbors: Optional[Union[float, Dict[str, float]]] = None,
) -> None:
"""InteractionBlock.
Args:
irreps_in: input irreps
irreps_out: output irreps
radial_mlp_depth (int): number of radial layers
radial_mlp_width (int): number of hidden neurons in radial function
use_sc (bool): use self-connection or not
is_first_layer (bool): whether to use first layer (default ``False``)
avg_num_neighbors (float/Dict[str, float]): global (float) or per-type (dict) average number of neighbors
type_names (List[str]): list of type names
"""
super().__init__()
self._init_irreps(
irreps_in=irreps_in,
required_irreps_in=[
AtomicDataDict.EDGE_EMBEDDING_KEY,
AtomicDataDict.EDGE_ATTRS_KEY,
AtomicDataDict.NODE_FEATURES_KEY,
AtomicDataDict.NODE_ATTRS_KEY,
],
my_irreps_in={
AtomicDataDict.EDGE_EMBEDDING_KEY: Irreps(
[
(
irreps_in[AtomicDataDict.EDGE_EMBEDDING_KEY].num_irreps,
(0, 1),
)
] # (0, 1) is even (invariant) scalars. We are forcing the EDGE_EMBEDDING to be invariant scalars so we can use a dense network
)
},
irreps_out={AtomicDataDict.NODE_FEATURES_KEY: irreps_out},
)
# === normalization module ===
self.avg_num_neighbors_norm = AvgNumNeighborsNorm(
avg_num_neighbors=avg_num_neighbors, type_names=type_names
)
self.use_sc = use_sc
feature_irreps_in = self.irreps_in[AtomicDataDict.NODE_FEATURES_KEY]
feature_irreps_out = self.irreps_out[AtomicDataDict.NODE_FEATURES_KEY]
irreps_edge_attr = self.irreps_in[AtomicDataDict.EDGE_ATTRS_KEY]
# - Build modules -
self.linear_1 = Linear(
irreps_in=feature_irreps_in,
irreps_out=feature_irreps_in,
internal_weights=True,
shared_weights=True,
)
irreps_mid = []
instructions = []
for i, (mul, ir_in) in enumerate(feature_irreps_in):
for j, (_, ir_edge) in enumerate(irreps_edge_attr):
for ir_out in ir_in * ir_edge:
if ir_out in feature_irreps_out:
k = len(irreps_mid)
irreps_mid.append((mul, ir_out))
instructions.append((i, j, k, "uvu", True))
# We sort the output irreps of the tensor product so that we can simplify them
# when they are provided to the second o3.Linear
irreps_mid = Irreps(irreps_mid)
irreps_mid, p, _ = irreps_mid.sort()
# Permute the output indexes of the instructions to match the sorted irreps:
instructions = [
(i_in1, i_in2, p[i_out], mode, train)
for i_in1, i_in2, i_out, mode, train in instructions
]
self.tp_scatter = TensorProductScatter(
feature_irreps_in,
irreps_edge_attr,
irreps_mid,
instructions,
)
# init_irreps already confirmed that the edge embeddding is all invariant scalars
self.edge_mlp = ScalarMLPFunction(
input_dim=self.irreps_in[AtomicDataDict.EDGE_EMBEDDING_KEY].num_irreps,
output_dim=self.tp_scatter.tp.weight_numel,
hidden_layers_depth=radial_mlp_depth,
hidden_layers_width=radial_mlp_width,
nonlinearity="silu", # hardcode SiLU
bias=False,
forward_weight_init=True,
)
self.linear_2 = Linear(
# irreps_mid has uncoallesed irreps because of the uvu instructions,
# but there's no reason to treat them seperately for the Linear
# Note that normalization of o3.Linear changes if irreps are coallesed
# (likely for the better)
irreps_in=irreps_mid.simplify(),
irreps_out=feature_irreps_out,
internal_weights=True,
shared_weights=True,
)
self.sc = None
if self.use_sc:
self.sc = FullyConnectedTensorProduct(
feature_irreps_in,
self.irreps_in[AtomicDataDict.NODE_ATTRS_KEY],
feature_irreps_out,
)
self.ghost_exchange = NoOpGhostExchangeModule(
field=AtomicDataDict.NODE_FEATURES_KEY, irreps_in=self.irreps_in
)
self.is_first_layer = is_first_layer
@torch.jit.unused
def _get_mliap_num_local(self, data: AtomicDataDict.Type) -> int:
return data[AtomicDataDict.LMP_MLIAP_DATA_KEY].nlocal
def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
if AtomicDataDict.LMP_MLIAP_DATA_KEY in data:
num_local_nodes = self._get_mliap_num_local(data)
else:
num_local_nodes = AtomicDataDict.num_nodes(data)
x = data[AtomicDataDict.NODE_FEATURES_KEY]
# truncate if not first layer
if not self.is_first_layer:
x = x[:num_local_nodes]
if self.sc is not None:
node_attrs = data[AtomicDataDict.NODE_ATTRS_KEY]
# truncate if not first layer
if not self.is_first_layer:
node_attrs = node_attrs[:num_local_nodes]
sc = self.sc(x, node_attrs)
x = self.linear_1(x)
# normalize before TP-scatter
data[AtomicDataDict.NODE_FEATURES_KEY] = x
data = self.avg_num_neighbors_norm(data)
x = data[AtomicDataDict.NODE_FEATURES_KEY]
# === comms for ghost-exchange ===
# only done if not first layer
# because initial embedding include ghosts since atom types come with ghosts
if not self.is_first_layer:
data[AtomicDataDict.NODE_FEATURES_KEY] = x
data = self.ghost_exchange(data, ghost_included=False)
x = data[AtomicDataDict.NODE_FEATURES_KEY]
# === TP and scatter ===
x = self.tp_scatter(
x=x,
edge_attr=data[AtomicDataDict.EDGE_ATTRS_KEY],
edge_weight=self.edge_mlp(data[AtomicDataDict.EDGE_EMBEDDING_KEY]),
edge_dst=data[AtomicDataDict.EDGE_INDEX_KEY][0],
edge_src=data[AtomicDataDict.EDGE_INDEX_KEY][1],
)[:num_local_nodes]
x = self.linear_2(x)
if self.sc is not None:
x = x + sc
data[AtomicDataDict.NODE_FEATURES_KEY] = x
return data
|