import logging import os from typing import Tuple, Union, List import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from e3nn import o3 from e3nn.nn import BatchNorm from e3nn.o3 import TensorProduct, Linear from torch_scatter import scatter, scatter_mean from models.layers import FCBlock def get_irrep_seq(ns, nv, use_second_order_repr, reduce_pseudoscalars): if use_second_order_repr: irrep_seq = [ f'{ns}x0e', f'{ns}x0e + {nv}x1o + {nv}x2e', f'{ns}x0e + {nv}x1o + {nv}x2e + {nv}x1e + {nv}x2o', f'{ns}x0e + {nv}x1o + {nv}x2e + {nv}x1e + {nv}x2o + {nv if reduce_pseudoscalars else ns}x0o' ] else: irrep_seq = [ f'{ns}x0e', f'{ns}x0e + {nv}x1o', f'{ns}x0e + {nv}x1o + {nv}x1e', f'{ns}x0e + {nv}x1o + {nv}x1e + {nv if reduce_pseudoscalars else ns}x0o' ] return irrep_seq def irrep_to_size(irrep): irreps = irrep.split(' + ') size = 0 for ir in irreps: m, (l, p) = ir.split('x') size += int(m) * (2 * int(l) + 1) return size class FasterTensorProduct(torch.nn.Module): # Implemented by Bowen Jing def __init__(self, in_irreps, sh_irreps, out_irreps, **kwargs): super().__init__() #for ir in in_irreps: # m, (l, p) = ir # assert l in [0, 1], "Higher order in irreps are not supported" #for ir in out_irreps: # m, (l, p) = ir # assert l in [0, 1], "Higher order out irreps are not supported" assert o3.Irreps(sh_irreps) == o3.Irreps('1x0e+1x1o'), "sh_irreps don't look like 1st order spherical harmonics" self.in_irreps = o3.Irreps(in_irreps) self.out_irreps = o3.Irreps(out_irreps) in_muls = {'0e': 0, '1o': 0, '1e': 0, '0o': 0} out_muls = {'0e': 0, '1o': 0, '1e': 0, '0o': 0} for (m, ir) in self.in_irreps: in_muls[str(ir)] = m for (m, ir) in self.out_irreps: out_muls[str(ir)] = m self.weight_shapes = { '0e': (in_muls['0e'] + in_muls['1o'], out_muls['0e']), '1o': (in_muls['0e'] + in_muls['1o'] + in_muls['1e'], out_muls['1o']), '1e': (in_muls['1o'] + in_muls['1e'] + in_muls['0o'], out_muls['1e']), '0o': (in_muls['1e'] + in_muls['0o'], out_muls['0o']) } self.weight_numel = sum(a * b for (a, b) in self.weight_shapes.values()) def forward(self, in_, sh, weight): in_dict, out_dict = {}, {'0e': [], '1o': [], '1e': [], '0o': []} for (m, ir), sl in zip(self.in_irreps, self.in_irreps.slices()): in_dict[str(ir)] = in_[..., sl] if ir[0] == 1: in_dict[str(ir)] = in_dict[str(ir)].reshape(list(in_dict[str(ir)].shape)[:-1] + [-1, 3]) sh_0e, sh_1o = sh[..., 0], sh[..., 1:] if '0e' in in_dict: out_dict['0e'].append(in_dict['0e'] * sh_0e.unsqueeze(-1)) out_dict['1o'].append(in_dict['0e'].unsqueeze(-1) * sh_1o.unsqueeze(-2)) if '1o' in in_dict: out_dict['0e'].append((in_dict['1o'] * sh_1o.unsqueeze(-2)).sum(-1) / np.sqrt(3)) out_dict['1o'].append(in_dict['1o'] * sh_0e.unsqueeze(-1).unsqueeze(-1)) out_dict['1e'].append(torch.linalg.cross(in_dict['1o'], sh_1o.unsqueeze(-2), dim=-1) / np.sqrt(2)) if '1e' in in_dict: out_dict['1o'].append(torch.linalg.cross(in_dict['1e'], sh_1o.unsqueeze(-2), dim=-1) / np.sqrt(2)) out_dict['1e'].append(in_dict['1e'] * sh_0e.unsqueeze(-1).unsqueeze(-1)) out_dict['0o'].append((in_dict['1e'] * sh_1o.unsqueeze(-2)).sum(-1) / np.sqrt(3)) if '0o' in in_dict: out_dict['1e'].append(in_dict['0o'].unsqueeze(-1) * sh_1o.unsqueeze(-2)) out_dict['0o'].append(in_dict['0o'] * sh_0e.unsqueeze(-1)) weight_dict = {} start = 0 for key in self.weight_shapes: in_, out = self.weight_shapes[key] weight_dict[key] = weight[..., start:start + in_ * out].reshape( list(weight.shape)[:-1] + [in_, out]) / np.sqrt(in_) start += in_ * out if out_dict['0e']: out_dict['0e'] = torch.cat(out_dict['0e'], dim=-1) out_dict['0e'] = torch.matmul(out_dict['0e'].unsqueeze(-2), weight_dict['0e']).squeeze(-2) if out_dict['1o']: out_dict['1o'] = torch.cat(out_dict['1o'], dim=-2) out_dict['1o'] = (out_dict['1o'].unsqueeze(-2) * weight_dict['1o'].unsqueeze(-1)).sum(-3) out_dict['1o'] = out_dict['1o'].reshape(list(out_dict['1o'].shape)[:-2] + [-1]) if out_dict['1e']: out_dict['1e'] = torch.cat(out_dict['1e'], dim=-2) out_dict['1e'] = (out_dict['1e'].unsqueeze(-2) * weight_dict['1e'].unsqueeze(-1)).sum(-3) out_dict['1e'] = out_dict['1e'].reshape(list(out_dict['1e'].shape)[:-2] + [-1]) if out_dict['0o']: out_dict['0o'] = torch.cat(out_dict['0o'], dim=-1) # out_dict['0o'] = (out_dict['0o'].unsqueeze(-1) * weight_dict['0o']).sum(-2) out_dict['0o'] = torch.matmul(out_dict['0o'].unsqueeze(-2), weight_dict['0o']).squeeze(-2) out = [] for _, ir in self.out_irreps: out.append(out_dict[str(ir)]) return torch.cat(out, dim=-1) def tp_scatter_simple(tp, fc_layer, node_attr, edge_index, edge_attr, edge_sh, out_nodes=None, reduce='mean', edge_weight=1.0): """ Perform TensorProduct + scatter operation, aka graph convolution. This function is only for edge_groups == 1. For multiple edge groups, and for larger graphs, use tp_scatter_multigroup instead. """ assert isinstance(edge_attr, torch.Tensor), \ "This function is only for a single edge group, so edge_attr must be a tensor and not a list." _device = node_attr.device _dtype = node_attr.dtype edge_src, edge_dst = edge_index out_irreps = fc_layer(edge_attr).to(_device).to(_dtype) out_irreps.mul_(edge_weight) tp = tp(node_attr[edge_dst], edge_sh, out_irreps) out_nodes = out_nodes or node_attr.shape[0] out = scatter(tp, edge_src, dim=0, dim_size=out_nodes, reduce=reduce) return out def tp_scatter_multigroup(tp: o3.TensorProduct, fc_layer: Union[nn.Module, nn.ModuleList], node_attr: torch.Tensor, edge_index: torch.Tensor, edge_attr_groups: List[torch.Tensor], edge_sh: torch.Tensor, out_nodes=None, reduce='mean', edge_weight=1.0): """ Perform TensorProduct + scatter operation, aka graph convolution. To keep the peak memory usage reasonably low, this function does not concatenate the edge_attr_groups. Rather, we sum the output of the tensor product for each edge group, and then divide by the number of edges Parameters ---------- tp: o3.TensorProduct fc_layer: nn.Module, or nn.ModuleList If a list, must be the same length as edge_attr_groups node_attr: torch.Tensor edge_index: torch.Tensor of shape (2, num_edges) Indicates the source and destination nodes of each edge edge_attr_groups: List[torch.Tensor] List of tensors, with shape (X_i, num_edge_attributes). Each tensor is a different group of edge attributes X may be different for each tensor, although sum(X_i) must be equal to edge_index.shape[1] edge_sh: torch.Tensor Spherical harmonics for the edges (see o3.spherical_harmonics) out_nodes: Number of output nodes reduce: str 'mean' or 'sum'. Reduce function for scatter. edge_weight : float or torch.Tensor Edge weights. If a tensor, must be the same shape as `edge_index` Returns ------- torch.Tensor Result of the graph convolution """ assert isinstance(edge_attr_groups, list), "This function is only for a list of edge groups" assert reduce in {"mean", "sum"}, "Only 'mean' and 'sum' are supported for reduce" # It would be possible to support mul/min/max but that would require more work and more code, # so only going to do it if it's needed. _device = node_attr.device _dtype = node_attr.dtype edge_src, edge_dst = edge_index edge_attr_lengths = [_edge_attr.shape[0] for _edge_attr in edge_attr_groups] total_rows = sum(edge_attr_lengths) assert total_rows == edge_index.shape[1], "Sum of edge_attr_groups must be equal to edge_index.shape[1]" num_edge_groups = len(edge_attr_groups) edge_weight_is_indexable = hasattr(edge_weight, '__getitem__') out_nodes = out_nodes or node_attr.shape[0] total_output_dim = sum([x.dim for x in tp.irreps_out]) final_out = torch.zeros((out_nodes, total_output_dim), device=_device, dtype=_dtype) div_factors = torch.zeros(out_nodes, device=_device, dtype=_dtype) cur_start = 0 for ii in range(num_edge_groups): cur_length = edge_attr_lengths[ii] cur_end = cur_start + cur_length cur_edge_range = slice(cur_start, cur_end) cur_edge_src, cur_edge_dst = edge_src[cur_edge_range], edge_dst[cur_edge_range] cur_fc = fc_layer[ii] if isinstance(fc_layer, nn.ModuleList) else fc_layer cur_out_irreps = cur_fc(edge_attr_groups[ii]) if edge_weight_is_indexable: cur_out_irreps.mul_(edge_weight[cur_edge_range]) else: cur_out_irreps.mul_(edge_weight) summand = tp(node_attr[cur_edge_dst, :], edge_sh[cur_edge_range, :], cur_out_irreps) # We take a simple sum, and then add up the count of edges which contribute, # so that we can take the mean later. final_out += scatter(summand, cur_edge_src, dim=0, dim_size=out_nodes, reduce="sum") div_factors += torch.bincount(cur_edge_src, minlength=out_nodes) cur_start = cur_end del cur_out_irreps, summand if reduce == 'mean': div_factors = torch.clamp(div_factors, torch.finfo(_dtype).eps) final_out = final_out / div_factors[:, None] return final_out class TensorProductConvLayer(torch.nn.Module): def __init__(self, in_irreps, sh_irreps, out_irreps, n_edge_features, residual=True, batch_norm=True, dropout=0.0, hidden_features=None, faster=False, edge_groups=1, tp_weights_layers=2, activation='relu', depthwise=False): super(TensorProductConvLayer, self).__init__() self.in_irreps = in_irreps self.out_irreps = out_irreps self.sh_irreps = sh_irreps self.residual = residual self.edge_groups = edge_groups self.out_size = irrep_to_size(out_irreps) self.depthwise = depthwise if hidden_features is None: hidden_features = n_edge_features if depthwise: in_irreps = o3.Irreps(in_irreps) sh_irreps = o3.Irreps(sh_irreps) out_irreps = o3.Irreps(out_irreps) irreps_mid = [] instructions = [] for i, (mul, ir_in) in enumerate(in_irreps): for j, (_, ir_edge) in enumerate(sh_irreps): for ir_out in ir_in * ir_edge: if ir_out in out_irreps: 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 = o3.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 = TensorProduct( in_irreps, sh_irreps, irreps_mid, instructions, shared_weights=False, internal_weights=False, ) 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=out_irreps, internal_weights=True, shared_weights=True, ) else: if faster: print("Faster Tensor Product") self.tp = FasterTensorProduct(in_irreps, sh_irreps, out_irreps) else: self.tp = o3.FullyConnectedTensorProduct(in_irreps, sh_irreps, out_irreps, shared_weights=False) if edge_groups == 1: self.fc = FCBlock(n_edge_features, hidden_features, self.tp.weight_numel, tp_weights_layers, dropout, activation) else: self.fc = [FCBlock(n_edge_features, hidden_features, self.tp.weight_numel, tp_weights_layers, dropout, activation) for _ in range(edge_groups)] self.fc = nn.ModuleList(self.fc) self.batch_norm = BatchNorm(out_irreps) if batch_norm else None def forward(self, node_attr, edge_index, edge_attr, edge_sh, out_nodes=None, reduce='mean', edge_weight=1.0): if edge_index.shape[1] == 0 and node_attr.shape[0] == 0: raise ValueError("No edges and no nodes") _dtype = node_attr.dtype if edge_index.shape[1] == 0: out = torch.zeros((node_attr.shape[0], self.out_size), dtype=_dtype, device=node_attr.device) else: if self.edge_groups == 1: out = tp_scatter_simple(self.tp, self.fc, node_attr, edge_index, edge_attr, edge_sh, out_nodes, reduce, edge_weight) else: out = tp_scatter_multigroup(self.tp, self.fc, node_attr, edge_index, edge_attr, edge_sh, out_nodes, reduce, edge_weight) if self.depthwise: out = self.linear_2(out) if self.batch_norm: out = self.batch_norm(out) if self.residual: padded = F.pad(node_attr, (0, out.shape[-1] - node_attr.shape[-1])) out = out + padded out = out.to(_dtype) return out class OldTensorProductConvLayer(torch.nn.Module): def __init__(self, in_irreps, sh_irreps, out_irreps, n_edge_features, residual=True, batch_norm=True, dropout=0.0, hidden_features=None): super(OldTensorProductConvLayer, self).__init__() self.in_irreps = in_irreps self.out_irreps = out_irreps self.sh_irreps = sh_irreps self.residual = residual if hidden_features is None: hidden_features = n_edge_features self.tp = tp = o3.FullyConnectedTensorProduct(in_irreps, sh_irreps, out_irreps, shared_weights=False) self.fc = nn.Sequential( nn.Linear(n_edge_features, hidden_features), nn.ReLU(), nn.Dropout(dropout), nn.Linear(hidden_features, tp.weight_numel) ) self.batch_norm = BatchNorm(out_irreps) if batch_norm else None def forward(self, node_attr, edge_index, edge_attr, edge_sh, out_nodes=None, reduce='mean', edge_weight=1.0): # Break up the edge_attr into chunks to limit the maximum memory usage edge_chunk_size = 100_000 num_edges = edge_attr.shape[0] num_chunks = (num_edges // edge_chunk_size) if num_edges % edge_chunk_size == 0 \ else (num_edges // edge_chunk_size) + 1 edge_ranges = np.array_split(np.arange(num_edges), num_chunks) edge_attr_groups = [edge_attr[cur_range] for cur_range in edge_ranges] out = tp_scatter_multigroup(self.tp, self.fc, node_attr, edge_index, edge_attr_groups, edge_sh, out_nodes, reduce, edge_weight) if self.residual: padded = F.pad(node_attr, (0, out.shape[-1] - node_attr.shape[-1])) out = out + padded if self.batch_norm: out = self.batch_norm(out) out = out.to(node_attr.dtype) return out