"""This module is adapted from https://github.com/Open-Catalyst-Project/ocp/tree/master/ocpmodels/models """ import torch import torch.nn as nn from torch_scatter import scatter # from torch_geometric.nn.acts import swish from torch.nn import SiLU as swish from torch_geometric.nn.inits import glorot_orthogonal from torch_geometric.nn.models.dimenet import ( BesselBasisLayer, EmbeddingBlock, ResidualLayer, SphericalBasisLayer, ) from torch_sparse import SparseTensor from diffcsp.common.data_utils import ( get_pbc_distances, frac_to_cart_coords, radius_graph_pbc_wrapper, ) try: import sympy as sym except ImportError: sym = None class InteractionPPBlock(torch.nn.Module): def __init__( self, hidden_channels, int_emb_size, basis_emb_size, num_spherical, num_radial, num_before_skip, num_after_skip, act=swish, ): super(InteractionPPBlock, self).__init__() self.act = act # Transformations of Bessel and spherical basis representations. self.lin_rbf1 = nn.Linear(num_radial, basis_emb_size, bias=False) self.lin_rbf2 = nn.Linear(basis_emb_size, hidden_channels, bias=False) self.lin_sbf1 = nn.Linear( num_spherical * num_radial, basis_emb_size, bias=False ) self.lin_sbf2 = nn.Linear(basis_emb_size, int_emb_size, bias=False) # Dense transformations of input messages. self.lin_kj = nn.Linear(hidden_channels, hidden_channels) self.lin_ji = nn.Linear(hidden_channels, hidden_channels) # Embedding projections for interaction triplets. self.lin_down = nn.Linear(hidden_channels, int_emb_size, bias=False) self.lin_up = nn.Linear(int_emb_size, hidden_channels, bias=False) # Residual layers before and after skip connection. self.layers_before_skip = torch.nn.ModuleList( [ ResidualLayer(hidden_channels, act) for _ in range(num_before_skip) ] ) self.lin = nn.Linear(hidden_channels, hidden_channels) self.layers_after_skip = torch.nn.ModuleList( [ ResidualLayer(hidden_channels, act) for _ in range(num_after_skip) ] ) self.reset_parameters() def reset_parameters(self): glorot_orthogonal(self.lin_rbf1.weight, scale=2.0) glorot_orthogonal(self.lin_rbf2.weight, scale=2.0) glorot_orthogonal(self.lin_sbf1.weight, scale=2.0) glorot_orthogonal(self.lin_sbf2.weight, scale=2.0) glorot_orthogonal(self.lin_kj.weight, scale=2.0) self.lin_kj.bias.data.fill_(0) glorot_orthogonal(self.lin_ji.weight, scale=2.0) self.lin_ji.bias.data.fill_(0) glorot_orthogonal(self.lin_down.weight, scale=2.0) glorot_orthogonal(self.lin_up.weight, scale=2.0) for res_layer in self.layers_before_skip: res_layer.reset_parameters() glorot_orthogonal(self.lin.weight, scale=2.0) self.lin.bias.data.fill_(0) for res_layer in self.layers_after_skip: res_layer.reset_parameters() def forward(self, x, rbf, sbf, idx_kj, idx_ji): # Initial transformations. x_ji = self.act(self.lin_ji(x)) x_kj = self.act(self.lin_kj(x)) # Transformation via Bessel basis. rbf = self.lin_rbf1(rbf) rbf = self.lin_rbf2(rbf) x_kj = x_kj * rbf # Down-project embeddings and generate interaction triplet embeddings. x_kj = self.act(self.lin_down(x_kj)) # Transform via 2D spherical basis. sbf = self.lin_sbf1(sbf) sbf = self.lin_sbf2(sbf) x_kj = x_kj[idx_kj] * sbf # Aggregate interactions and up-project embeddings. x_kj = scatter(x_kj, idx_ji, dim=0, dim_size=x.size(0)) x_kj = self.act(self.lin_up(x_kj)) h = x_ji + x_kj for layer in self.layers_before_skip: h = layer(h) h = self.act(self.lin(h)) + x for layer in self.layers_after_skip: h = layer(h) return h class OutputPPBlock(torch.nn.Module): def __init__( self, num_radial, hidden_channels, out_emb_channels, out_channels, num_layers, act=swish, ): super(OutputPPBlock, self).__init__() self.act = act self.lin_rbf = nn.Linear(num_radial, hidden_channels, bias=False) self.lin_up = nn.Linear(hidden_channels, out_emb_channels, bias=True) self.lins = torch.nn.ModuleList() for _ in range(num_layers): self.lins.append(nn.Linear(out_emb_channels, out_emb_channels)) self.lin = nn.Linear(out_emb_channels, out_channels, bias=False) self.reset_parameters() def reset_parameters(self): glorot_orthogonal(self.lin_rbf.weight, scale=2.0) glorot_orthogonal(self.lin_up.weight, scale=2.0) for lin in self.lins: glorot_orthogonal(lin.weight, scale=2.0) lin.bias.data.fill_(0) self.lin.weight.data.fill_(0) def forward(self, x, rbf, i, num_nodes=None): x = self.lin_rbf(rbf) * x x = scatter(x, i, dim=0, dim_size=num_nodes) x = self.lin_up(x) for lin in self.lins: x = self.act(lin(x)) return self.lin(x) class DimeNetPlusPlus(torch.nn.Module): r"""DimeNet++ implementation based on https://github.com/klicperajo/dimenet. Args: hidden_channels (int): Hidden embedding size. out_channels (int): Size of each output sample. num_blocks (int): Number of building blocks. int_emb_size (int): Embedding size used for interaction triplets basis_emb_size (int): Embedding size used in the basis transformation out_emb_channels(int): Embedding size used for atoms in the output block num_spherical (int): Number of spherical harmonics. num_radial (int): Number of radial basis functions. cutoff: (float, optional): Cutoff distance for interatomic interactions. (default: :obj:`5.0`) envelope_exponent (int, optional): Shape of the smooth cutoff. (default: :obj:`5`) num_before_skip: (int, optional): Number of residual layers in the interaction blocks before the skip connection. (default: :obj:`1`) num_after_skip: (int, optional): Number of residual layers in the interaction blocks after the skip connection. (default: :obj:`2`) num_output_layers: (int, optional): Number of linear layers for the output blocks. (default: :obj:`3`) act: (function, optional): The activation funtion. (default: :obj:`swish`) """ url = "https://github.com/klicperajo/dimenet/raw/master/pretrained" def __init__( self, hidden_channels, out_channels, num_blocks, int_emb_size, basis_emb_size, out_emb_channels, num_spherical, num_radial, cutoff=5.0, envelope_exponent=5, num_before_skip=1, num_after_skip=2, num_output_layers=3, act=swish, ): super(DimeNetPlusPlus, self).__init__() self.cutoff = cutoff if sym is None: raise ImportError("Package `sympy` could not be found.") self.num_blocks = num_blocks self.rbf = BesselBasisLayer(num_radial, cutoff, envelope_exponent) self.sbf = SphericalBasisLayer( num_spherical, num_radial, cutoff, envelope_exponent ) self.emb = EmbeddingBlock(num_radial, hidden_channels, act) self.output_blocks = torch.nn.ModuleList( [ OutputPPBlock( num_radial, hidden_channels, out_emb_channels, out_channels, num_output_layers, act, ) for _ in range(num_blocks + 1) ] ) self.interaction_blocks = torch.nn.ModuleList( [ InteractionPPBlock( hidden_channels, int_emb_size, basis_emb_size, num_spherical, num_radial, num_before_skip, num_after_skip, act, ) for _ in range(num_blocks) ] ) self.reset_parameters() def reset_parameters(self): self.rbf.reset_parameters() self.emb.reset_parameters() for out in self.output_blocks: out.reset_parameters() for interaction in self.interaction_blocks: interaction.reset_parameters() def triplets(self, edge_index, num_nodes): row, col = edge_index # j->i # row, col = col, row # Swap because my definition of edge_index is i->j value = torch.arange(row.size(0), device=row.device) adj_t = SparseTensor( row=col, col=row, value=value, sparse_sizes=(num_nodes, num_nodes) ) adj_t_row = adj_t[row] num_triplets = adj_t_row.set_value(None).sum(dim=1).to(torch.long) # Node indices (k->j->i) for triplets. idx_i = col.repeat_interleave(num_triplets) idx_j = row.repeat_interleave(num_triplets) idx_k = adj_t_row.storage.col() mask = idx_i != idx_k # Remove i == k triplets. idx_i, idx_j, idx_k = idx_i[mask], idx_j[mask], idx_k[mask] # Edge indices (k-j, j->i) for triplets. idx_kj = adj_t_row.storage.value()[mask] idx_ji = adj_t_row.storage.row()[mask] return col, row, idx_i, idx_j, idx_k, idx_kj, idx_ji def forward(self, z, pos, batch=None): """""" raise NotImplementedError class DimeNetPlusPlusWrap(DimeNetPlusPlus): def __init__( self, num_targets, hidden_channels=128, num_blocks=4, int_emb_size=64, basis_emb_size=8, out_emb_channels=256, num_spherical=7, num_radial=6, otf_graph=False, cutoff=10.0, max_num_neighbors=20, envelope_exponent=5, num_before_skip=1, num_after_skip=2, num_output_layers=3, readout='mean', ): self.num_targets = num_targets self.cutoff = cutoff self.max_num_neighbors = max_num_neighbors self.otf_graph = otf_graph self.readout = readout super(DimeNetPlusPlusWrap, self).__init__( hidden_channels=hidden_channels, out_channels=num_targets, num_blocks=num_blocks, int_emb_size=int_emb_size, basis_emb_size=basis_emb_size, out_emb_channels=out_emb_channels, num_spherical=num_spherical, num_radial=num_radial, cutoff=cutoff, envelope_exponent=envelope_exponent, num_before_skip=num_before_skip, num_after_skip=num_after_skip, num_output_layers=num_output_layers, ) def forward(self, data): batch = data.batch if self.otf_graph: edge_index, cell_offsets, neighbors = radius_graph_pbc_wrapper( data, self.cutoff, self.max_num_neighbors, data.num_atoms.device ) data.edge_index = edge_index data.to_jimages = cell_offsets data.num_bonds = neighbors pos = frac_to_cart_coords( data.frac_coords, data.lengths, data.angles, data.num_atoms) out = get_pbc_distances( data.frac_coords, data.edge_index, data.lengths, data.angles, data.to_jimages, data.num_atoms, data.num_bonds, return_offsets=True ) edge_index = out["edge_index"] dist = out["distances"] offsets = out["offsets"] j, i = edge_index _, _, idx_i, idx_j, idx_k, idx_kj, idx_ji = self.triplets( edge_index, num_nodes=data.atom_types.size(0) ) # Calculate angles. pos_i = pos[idx_i].detach() pos_j = pos[idx_j].detach() pos_ji, pos_kj = ( pos[idx_j].detach() - pos_i + offsets[idx_ji], pos[idx_k].detach() - pos_j + offsets[idx_kj], ) a = (pos_ji * pos_kj).sum(dim=-1) b = torch.cross(pos_ji, pos_kj).norm(dim=-1) angle = torch.atan2(b, a) rbf = self.rbf(dist) sbf = self.sbf(dist, angle, idx_kj) # Embedding block. x = self.emb(data.atom_types.long(), rbf, i, j) P = self.output_blocks[0](x, rbf, i, num_nodes=pos.size(0)) # Interaction blocks. for interaction_block, output_block in zip( self.interaction_blocks, self.output_blocks[1:] ): x = interaction_block(x, rbf, sbf, idx_kj, idx_ji) P += output_block(x, rbf, i, num_nodes=pos.size(0)) # Use mean if batch is None: if self.readout == 'mean': energy = P.mean(dim=0) elif self.readout == 'sum': energy = P.sum(dim=0) elif self.readout == 'cat': import pdb pdb.set_trace() energy = torch.cat([P.sum(dim=0), P.mean(dim=0)]) else: raise NotImplementedError else: # TODO: if want to use cat, need two lines here energy = scatter(P, batch, dim=0, reduce=self.readout) return energy @property def num_params(self): return sum(p.numel() for p in self.parameters())