<repo_name>wlgfour/eqnets<file_sep>/e3nn/install.sh
pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
pip install --no-index torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
pip install --no-index torch-sparse -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
pip install --no-index torch-cluster -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
pip install --no-index torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.7.0+cu101.html
pip install torch-geometric
pip install e3nn
<file_sep>/e3nn/net1.py
import os
from typing import List
import matplotlib as mpl
import numpy as np

from torch.utils import data
from e3nn import non_linearities
from e3nn.networks import GatedConvNetwork
from e3nn.point.data_helpers import DataNeighbors
from e3nn.point.message_passing import Convolution as MessageConv
from utils import get_data_loader, mask_by_len, N_FEATURES
import torch
from setproctitle import setproctitle
import argparse
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


class MaskedAverageNet(torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.network = GatedConvNetwork(*args, **kwargs)

    def forward(self, lengths, *args, **kwargs):
        output = self.network(*args, **kwargs)
        output = mask_by_len(output, lengths)
        return output.sum(1)[:, 0] / lengths

class GatedConvNeighbors(torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()
        kwargs['convolution'] = MessageConv
        self.network = GatedConvNetwork(*args, **kwargs)

    def forward(self, nbrs):
        """ nbrs is a neighbors list generated by e3nn.point.data_helpers.DataNeighbors
        """
        output = self.network(nbrs.x, nbrs.edge_index, nbrs.edge_attr)
        return output

def main(
    name: str = 'GatedConvDefault',
    data_path: str='',
    save_dir: str='',
    save_every: int=10000,
    tblog: str = '',
    tbinterval: int = 100,
    gpu: int = -1,
    lmax: int = 3,
    rmax: float = 1000,
    layers: int = 3,
    lr: float = 1e-2,
    batch_size: int = 1,
    visualization_structs: List[str] = [],
    ):
    """ Try to do the same as main1, but with the conbolution calculated only over a fixed-radius
        neighborhood around a certain point.
    """

    if tbinterval != 0:
        ldir = os.path.join(tblog, name)
        print(f'Logging in {ldir}...')
        tbwriter = SummaryWriter(log_dir=ldir, comment=f'pid_{os.getpid()}')
    else:
        tbwriter = None

    torch.set_default_dtype(torch.float64)
    if gpu == -1:
        device = 'cpu'
    else:
        device = f'cuda:{gpu}'
    device = torch.device(device)

    dataset, data_loader = get_data_loader(data_path, batch_size=batch_size, max_length=-1)

    Rs_in = [(N_FEATURES, 0)]
    Rs_hidden = [(16, 0), (8, 1), (4, 2)]
    Rs_out = [(1, 0)]

    net = GatedConvNeighbors(Rs_in, Rs_hidden, Rs_out, lmax=lmax, max_radius=rmax, layers=layers)
    net = net.to(device)

    optimizer = torch.optim.Adam(net.parameters(), lr=lr)

    epoch = 0
    global_step = 0
    metrics = {
        'Loss': [],
        'Avg_Neighbors': [],
        'Out/range': [],
        'Out/var': []
    }
    while True:
        epoch += 1
        for batch in data_loader:
            for j in range(batch_size):
                global_step += 1
                coords = batch['coords'][j]
                feature = batch['features'][j]
                # lengths = batch['length'][j]
                nbrs = DataNeighbors(feature, coords, rmax)
                # for tens in [nbrs]:  # coords, feature, labels, lengths, 
                    # tens.to(device)
                nbrs = nbrs.to(device)
                out = net(nbrs)
                # if global_step == 1 and tbwriter != None:
                    # tbwriter.add_graph(net, [nbrs.x, nbrs.edge_index, nbrs.edge_attr])

                # https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/20
                labels = batch['drmsd'][j].to(device)
                loss = (out.mean() - labels) ** 2
                loss.backward()
                if global_step % batch_size == 0:
                    optimizer.step()
                    optimizer.zero_grad()


                # update metrics
                metrics['Loss'].append(loss.detach().cpu())
                metrics['Avg_Neighbors'].append(nbrs.edge_attr.shape[0] / nbrs.x.shape[0])
                metrics['Out/range'].append((out.max() - out.min()).cpu())
                metrics['Out/var'].append(out.var().cpu())
                if global_step % save_every == 0 and save_dir != '' and save_every != 0:
                    torch.save(net.state_dict(), os.path.join(save_dir, sdir))
                    if tbwriter is not None:
                        pass
                        # tbwriter.add_hparams({'lr': lr, 'lmax': lmax, 'rmax': rmax, 'layers': layers, 'epochs': epoch, 'steps': i},
                        #          {'Loss': 0., 'Avg_Neighbors': 0., 'Out/var': 0.},
                        #         )
                if global_step % tbinterval == 0:
                    if tbwriter is None:
                        print(f"epoch:step={epoch}:{global_step} loss={loss:.2f}")
                    for m in metrics:
                        val = sum(metrics[m]) / max(1, len(metrics[m]))
                        if tbwriter is not None:
                            tbwriter.add_scalar(m, val, global_step)
                        else:
                            print(f'\t{m}: {val:.2f}')
                    metrics = {m: [] for m in metrics}

                    # plot backbone colored by predictions
                    # net.eval()
                    # with torch.no_grad():
                    #     for s in visualization_structs:
                    #         if s not in data_loader:
                    #             coords = dataset[s]
                    #             nbrs = DataNeighbors(
                    #                 coords['features'],
                    #                 coords['coords'],
                    #                 rmax
                    #             )
                    #             nbrs.to(device)
                    #             out = net(nbrs).detach().cpu().numpy()
                    #             coords = coords['coords'].numpy()

                    #             fig = plt.figure(figsize=(6, 6))
                    #             ax = fig.add_subplot(1, 1, 1, projection='3d')
                    #             _cmap = mpl.cm.get_cmap('cool')
                    #             sm = mpl.cm.ScalarMappable(norm=mpl.colors.Normalize(vmin=0, vmax=750), cmap=_cmap)
                    #             cmap = lambda c: _cmap(min(c, 750) / 750)
                    #             # cmap = sm.cmap
                    #             colors = [np.array(cmap(o)) for o in out.flatten()]
                    #             ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2], c=colors)
                    #             for i in range(len(coords) - 1):
                    #                 seg = coords[i:i+2]
                    #                 ax.plot(seg[:, 0], seg[:, 1], seg[:, 2], c=(colors[i] + colors[i+1]) / 2)
                    #             fig.colorbar(sm)

                    #             tbwriter.add_figure(s, fig, global_step=global_step)
                    # net.train()




if __name__ == '__main__':
    # command line interface
    # Parse inputs
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--vs', dest='visualization_structs', type=str, nargs='+', default=[''],
                        help='List fo structures from the training set to include in the tensorboard plots.')
    parser.add_argument('--data', type=str,
                        help=f'Direcotry to find the representations dataset outputted by protein_geometry/src/scripts/rgn_data.py')
    parser.add_argument('--name', type=str, default='',
                        help='Override the name of the run. Affects the save file name and the name of the tensorboard run. '
                             'Will default to depend on the model specs.')
    parser.add_argument('--load', type=None, default=None,
                        help='Not Implemented')
    parser.add_argument('--save-dir', dest='save_dir', type=str, default=os.path.join('.', 'save'),
                        help='Where to save the model.')
    parser.add_argument('--save-interval', dest='save_interval', type=int, default=10000,
                        help='How often to save the model. Setting to 0 will disable saving.')
    parser.add_argument('--tb', type=str, default=os.path.join('~', 'eqnetworks', 'runs'),
                        help='Where to log tensorboard output. Will create a directory inside of the tb directory for this run.')
    parser.add_argument('--tb-interval', dest='tb_interval', type=int, default=100,
                        help='How often to log to tensorboard. Setting to 0 will disable logging.')
    parser.add_argument('--gpu', type=int, default=-1,
                        help='Which gpu to use. -1 means cpu.')
    parser.add_argument('--lmax', type=int, default=3,
                        help='Maximum l to use for spherical hermonics.')
    parser.add_argument('--rmax', type=float, default=1000,
                        help='Radius (Picometer) for convolution scope.')
    parser.add_argument('--layers', type=int, default=3,
                        help='Number of gated equivariant convolutional layers.')
    parser.add_argument('--reps', type=None, default=None,
                        help='Not Implemented')
    parser.add_argument('--features', type=None, default=None,
                        help='Not Implemented')
    parser.add_argument('--bs', type=int, default=1,
                        help='Batch size.')
    parser.add_argument('--lr', type=float, default=1e-2,
                        help='Batch size.')
    opts = parser.parse_args()

    data_path = opts.data  # '~/protein_geometry/data/representations/rgn'
    save_dir = opts.save_dir
    if opts.name == '':
        name = f'lmax:{opts.lmax}-rmax:{opts.rmax}-layers:{opts.layers}-bs:{opts.bs}-lr:{opts.lr}'
    else:
        name = opts.name

    # Considering parameters to a file that can be loaded.
    kwargs = dict(
        data_path=data_path,
        save_every=opts.save_interval,
        save_dir=save_dir,
        tblog=opts.tb,
        gpu=opts.gpu,
        lmax=opts.lmax,
        rmax=opts.rmax,
        name=name,
        layers=opts.layers,
        tbinterval=opts.tb_interval,
        visualization_structs=opts.visualization_structs,
        batch_size=opts.bs,
        lr=opts.lr,
    )

    setproctitle(name)
    
    # save dir
    if kwargs['save_every'] != 0:
        i = 0
        sdir = os.path.join(kwargs['save_dir'], f'{name}{i}')
        while os.path.exists(sdir):
            i += 1
            sdir = os.path.join(kwargs['save_dir'], f'{name}{i}')
        os.mkdir(sdir)
        kwargs['save_dir'] = sdir


    print(f'Starting {name} with pid {os.getpid()}')

    main(**kwargs)
<file_sep>/e3nn/utils.py
from collections import defaultdict
from torch import Tensor
from torch.utils.data import Dataset, DataLoader
import torch
import numpy as np
import pandas as pd
import os
from typing import Any, Dict, List, Tuple, Union

DSSP_CODES = ['H', 'G', 'I', 'E', 'O', 'S', 'B', 'T', '-', 'L']
AAs = ['A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P', 'S', 'T', 'W', 'Y', 'V', 'B', 'Z']
N_FEATURES = len(AAs) + len(DSSP_CODES)

class ProteinLoader(Dataset):
    """
    ref: https://pytorch.org/tutorials/beginner/data_loading_tutorial.html
    Loads a list of protein representations and their corresponding ground-truth drmsds.

    Inheritance:
        Dataset:

    Args:
        directory (str):
        representation (str='Euclidean'):

    >>> 
    """
    
    def __init__(self, dir: str, representation: str='Euclidean', max_length: int=-1):
        """
        Args:
            self (undefined):
            directory (str):
            representation (str='Euclidean'):
        """
        self.dir = dir
        self.labels = pd.read_csv(os.path.join(dir, 'labels.csv'))
        self.rep = representation
        self.max_length = max_length

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, i: Union[int, str]):
        if torch.is_tensor(i):
            i = i.tolist()
        
        if isinstance(i, int):
            name = self.labels.iloc[i, 0]
        elif isinstance(i, str):
            name = i
            i = self.labels[self.labels == i].index[0]
        else:
            raise TypeError(f'Type ({type(i)}) not supported by ProteinLoader.')

        dpath = os.path.join(self.dir, self.rep, f'{name}.npy')
        fpath = os.path.join(self.dir, 'Features', f'{int(name[1:name.index("_")])}.csv')
        coords = np.load(dpath)
        if os.path.isfile(fpath):
            with open(fpath, 'r') as f:
                lines = f.readlines()
            features = read_features(lines)  # concatenated 1-hots for each position
        else:
            features = torch.from_numpy(np.ones_like(coords))[:, 0:1]  # a single '1' at each position
        if (self.max_length != -1) and (len(coords) > self.max_length):
            coords = coords[:self.max_length]
        drmsd = self.labels.iloc[i, 1]
        return {'name': name, 'coords': torch.from_numpy(coords), 'drmsd': torch.tensor(drmsd), 'features': features}

    def __contains__(self, key: Any):
        if isinstance(key, str):
            return (self.labels['protein.iteration'] == key).any()
        else:
            raise TypeError(f'Type ({type(i)}) not supported by ProteinLoader.')
        
def _collate(batch: Dict[str, Union[str, Tensor]]) -> List[Union[Tensor, List[str]]]:
    ret = defaultdict(list)
    for sample in batch:
        ret['coords'].append(sample['coords'])
        ret['features'].append(sample['features'])
        ret['length'].append(torch.tensor(sample['coords'].shape[0]))
        ret['drmsd'].append(sample['drmsd'])
        ret['name'].append(sample['name'])
    for k in ['coords', 'features']:
        ret[k] = torch.nn.utils.rnn.pad_sequence(ret[k], batch_first=True)
    for k in ['length', 'drmsd']:
        ret[k] = torch.stack(ret[k])
    return ret

def read_features(lines: List[str]) -> torch.Tensor:
    """ Takes a list of lines read from a file containing features with the columns 'DSSP' and 'Sequence' and outputs
        a concatenated list of one-hot vectors that is consistent no matter what parameters are passed.
    """
    cols = {k: i for (i, k) in enumerate(lines[0].strip().split(','))}
    features = np.zeros((len(lines) - 1, len(AAs) + len(DSSP_CODES)))
    for i, line in enumerate(lines[1:]):
        l = line.strip().split(',')
        # index in AAs of the sequence of this residue
        features[i, AAs.index(l[cols['Sequence']])] = 1.
        features[i, len(AAs) + DSSP_CODES.index(l[cols['DSSP']])] = 1.
    return torch.from_numpy(features)
    

def mask_by_len(inputs: Tensor, lens: Tensor) -> Tensor:
    ml = torch.max(lens)
    masks = torch.arange(ml)[None, :] < lens[:, None]
    return inputs * masks[:, :, None]

def get_data_loader(dir: str, representation: str='Euclidean', batch_size: int=1, max_length: int=100) -> Tuple[Dataset, DataLoader]:
    """
    Description of get_data_loader

    Args:
        dir (str):
        representation (str='Euclidean'):

    Returns:
        DataLoader

    """
    loader = ProteinLoader(dir=dir, representation=representation, max_length=max_length)
    return loader, DataLoader(loader, batch_size=batch_size, shuffle=True, collate_fn=_collate, pin_memory=True)

if __name__ == '__main__':
    d = '/home/wlg/development/HMS/protein_geometry/data/representations/150_1500'
    dataset, data_loader = get_data_loader(d)
    for batch in data_loader:
        pass<file_sep>/e3nn/uninstall.sh
pip uninstall -y e3nn torch-grometric torch-scatter torch-sparse torch-cluster torch-spline-conv