models/tensor_layers.py

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