Spaces:

igashov
/

DiffLinker

Sleeping

App Files Files Community

igashov commited on Oct 15, 2022

Commit

95ba5bc

•

1 Parent(s): e361c7a

DiffLinker code

Browse files

Files changed (19) hide show

app.py +83 -7
src/__init__.py +0 -0
src/const.py +218 -0
src/datasets.py +350 -0
src/delinker.py +278 -0
src/delinker_utils/__init__.py +0 -0
src/delinker_utils/calc_SC_RDKit.py +40 -0
src/delinker_utils/frag_utils.py +413 -0
src/delinker_utils/sascorer.py +173 -0
src/edm.py +730 -0
src/egnn.py +541 -0
src/lightning.py +473 -0
src/linker_size.py +95 -0
src/linker_size_lightning.py +460 -0
src/metrics.py +167 -0
src/molecule_builder.py +102 -0
src/noise.py +169 -0
src/utils.py +348 -0
src/visualizer.py +227 -0

app.py CHANGED Viewed

@@ -1,6 +1,14 @@
 import gradio as gr
 import os
-import sys
 HTML_TEMPLATE = """<!DOCTYPE html>
@@ -43,20 +51,88 @@ allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
 allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
-def read_molecule(path):
     with open(path, "r") as f:
         return "".join(f.readlines())
 def generate(input_file):
     try:
         path = input_file.name
         molecule = read_molecule(path)
-        fmt = path.split('.')[-1]
-    except:
-        return 'Error: could not open the provided file'
-    html = HTML_TEMPLATE.format(molecule=molecule, fmt=fmt)
     return IFRAME_TEMPLATE.format(html=html)

 import gradio as gr
 import os
+import torch
+import subprocess
+from rdkit import Chem
+from src import const
+from src.visualizer import save_xyz_file
+from src.datasets import get_dataloader, collate_with_fragment_edges, parse_molecule
+from src.lightning import DDPM
+from src.linker_size_lightning import SizeClassifier
 HTML_TEMPLATE = """<!DOCTYPE html>
 allowpaymentrequest="" frameborder="0" srcdoc='{html}'></iframe>"""
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+os.makedirs("results", exist_ok=True)
+print('Created results directory')
+size_nn = SizeClassifier.load_from_checkpoint('models/geom_size_gnn.ckpt', map_location=device).eval().to(device)
+print('Loaded SizeGNN model')
+ddpm = DDPM.load_from_checkpoint('models/geom_difflinker.ckpt', map_location=device).eval().to(device)
+print('Loaded diffusion model')
+def sample_fn(_data):
+    output, _ = size_nn.forward(_data)
+    probabilities = torch.softmax(output, dim=1)
+    distribution = torch.distributions.Categorical(probs=probabilities)
+    samples = distribution.sample()
+    sizes = []
+    for label in samples.detach().cpu().numpy():
+        sizes.append(size_nn.linker_id2size[label])
+    sizes = torch.tensor(sizes, device=samples.device, dtype=torch.long)
+    return sizes
+def read_molecule_content(path):
     with open(path, "r") as f:
         return "".join(f.readlines())
+def read_molecule(path):
+    if path.endswith('.pdb'):
+        return Chem.MolFromPDBFile(path, sanitize=False, removeHs=True)
+    elif path.endswith('.mol'):
+        return Chem.MolFromMolFile(path, sanitize=False, removeHs=True)
+    elif path.endswith('.mol2'):
+        return Chem.MolFromMol2File(path, sanitize=False, removeHs=True)
+    elif path.endswith('.sdf'):
+        return Chem.SDMolSupplier(path, sanitize=False, removeHs=True)[0]
+    raise Exception('Unknown file extension')
 def generate(input_file):
     try:
         path = input_file.name
         molecule = read_molecule(path)
+        name = '.'.join(molecule.split('/')[-1].split('.')[:-1])
+        out_sdf = f'results/{name}_generated.sdf'
+        print(f'Input path={path}, name={name}')
+    except Exception as e:
+        return f'Could not read the molecule: {e}'
+    positions, one_hot, charges = parse_molecule(molecule, is_geom=True)
+    positions = torch.tensor(positions, dtype=const.TORCH_FLOAT, device=device)
+    one_hot = torch.tensor(one_hot, dtype=const.TORCH_FLOAT, device=device)
+    print('Read and parsed molecule')
+    dataset = [{
+        'uuid': '0',
+        'name': '0',
+        'positions': torch.tensor(positions, dtype=const.TORCH_FLOAT, device=device),
+        'one_hot': torch.tensor(one_hot, dtype=const.TORCH_FLOAT, device=device),
+        'charges': torch.tensor(charges, dtype=const.TORCH_FLOAT, device=device),
+        'anchors': torch.zeros_like(charges, dtype=const.TORCH_FLOAT, device=device),
+        'fragment_mask': torch.ones_like(charges, dtype=const.TORCH_FLOAT, device=device),
+        'linker_mask': torch.zeros_like(charges, dtype=const.TORCH_FLOAT, device=device),
+        'num_atoms': len(positions),
+    }]
+    dataloader = get_dataloader(dataset, batch_size=1, collate_fn=collate_with_fragment_edges)
+    print('Created dataloader')
+    for data in dataloader:
+        chain, node_mask = ddpm.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
+        print('Generated linker')
+        x = chain[0][:, :, :ddpm.n_dims]
+        h = chain[0][:, :, ddpm.n_dims:]
+        save_xyz_file('results', h, x, node_mask, names=[name], is_geom=True, suffix='generated')
+        print('Saved XYZ file')
+        subprocess.run(f'obabel results/{name}_generated.xyz -O {out_sdf}', shell=True)
+        print('Converted to SDF')
+        break
+    generated_molecule = read_molecule_content(out_sdf)
+    html = HTML_TEMPLATE.format(molecule=generated_molecule, fmt='sdf')
     return IFRAME_TEMPLATE.format(html=html)

src/__init__.py ADDED Viewed

File without changes

src/const.py ADDED Viewed

	@@ -0,0 +1,218 @@

+import torch
+from rdkit import Chem
+TORCH_FLOAT = torch.float32
+TORCH_INT = torch.int8
+# #################################################################################### #
+# ####################################### ZINC ####################################### #
+# #################################################################################### #
+# Atom idx for one-hot encoding
+ATOM2IDX = {'C': 0, 'O': 1, 'N': 2, 'F': 3, 'S': 4, 'Cl': 5, 'Br': 6, 'I': 7}
+IDX2ATOM = {0: 'C', 1: 'O', 2: 'N', 3: 'F', 4: 'S', 5: 'Cl', 6: 'Br', 7: 'I'}
+# Atomic numbers (Z)
+CHARGES = {'C': 6, 'O': 8, 'N': 7, 'F': 9, 'S': 16, 'Cl': 17, 'Br': 35, 'I': 53}
+# One-hot atom types
+NUMBER_OF_ATOM_TYPES = len(ATOM2IDX)
+# #################################################################################### #
+# ####################################### GEOM ####################################### #
+# #################################################################################### #
+# Atom idx for one-hot encoding
+GEOM_ATOM2IDX = {'C': 0, 'O': 1, 'N': 2, 'F': 3, 'S': 4, 'Cl': 5, 'Br': 6, 'I': 7, 'P': 8}
+GEOM_IDX2ATOM = {0: 'C', 1: 'O', 2: 'N', 3: 'F', 4: 'S', 5: 'Cl', 6: 'Br', 7: 'I', 8: 'P'}
+# Atomic numbers (Z)
+GEOM_CHARGES = {'C': 6, 'O': 8, 'N': 7, 'F': 9, 'S': 16, 'Cl': 17, 'Br': 35, 'I': 53, 'P': 15}
+# One-hot atom types
+GEOM_NUMBER_OF_ATOM_TYPES = len(GEOM_ATOM2IDX)
+# Dataset keys
+DATA_LIST_ATTRS = {
+    'uuid', 'name', 'fragments_smi', 'linker_smi', 'num_atoms'
+}
+DATA_ATTRS_TO_PAD = {
+    'positions', 'one_hot', 'charges', 'anchors', 'fragment_mask', 'linker_mask', 'pocket_mask', 'fragment_only_mask'
+}
+DATA_ATTRS_TO_ADD_LAST_DIM = {
+    'charges', 'anchors', 'fragment_mask', 'linker_mask', 'pocket_mask', 'fragment_only_mask'
+}
+# Distribution of linker size in train data
+LINKER_SIZE_DIST = {
+    4: 85540,
+    3: 113928,
+    6: 70946,
+    7: 30408,
+    5: 77671,
+    9: 5177,
+    10: 1214,
+    8: 12712,
+    11: 158,
+    12: 7,
+}
+# Bond lengths from:
+# http://www.wiredchemist.com/chemistry/data/bond_energies_lengths.html
+# And:
+# http://chemistry-reference.com/tables/Bond%20Lengths%20and%20Enthalpies.pdf
+BONDS_1 = {
+    'H': {
+        'H': 74, 'C': 109, 'N': 101, 'O': 96, 'F': 92,
+        'B': 119, 'Si': 148, 'P': 144, 'As': 152, 'S': 134,
+        'Cl': 127, 'Br': 141, 'I': 161
+    },
+    'C': {
+        'H': 109, 'C': 154, 'N': 147, 'O': 143, 'F': 135,
+        'Si': 185, 'P': 184, 'S': 182, 'Cl': 177, 'Br': 194,
+        'I': 214
+    },
+    'N': {
+        'H': 101, 'C': 147, 'N': 145, 'O': 140, 'F': 136,
+        'Cl': 175, 'Br': 214, 'S': 168, 'I': 222, 'P': 177
+    },
+    'O': {
+        'H': 96, 'C': 143, 'N': 140, 'O': 148, 'F': 142,
+        'Br': 172, 'S': 151, 'P': 163, 'Si': 163, 'Cl': 164,
+        'I': 194
+    },
+    'F': {
+        'H': 92, 'C': 135, 'N': 136, 'O': 142, 'F': 142,
+        'S': 158, 'Si': 160, 'Cl': 166, 'Br': 178, 'P': 156,
+        'I': 187
+    },
+    'B': {
+        'H':  119, 'Cl': 175
+    },
+    'Si': {
+        'Si': 233, 'H': 148, 'C': 185, 'O': 163, 'S': 200,
+        'F': 160, 'Cl': 202, 'Br': 215, 'I': 243,
+    },
+    'Cl': {
+        'Cl': 199, 'H': 127, 'C': 177, 'N': 175, 'O': 164,
+        'P': 203, 'S': 207, 'B': 175, 'Si': 202, 'F': 166,
+        'Br': 214
+    },
+    'S': {
+        'H': 134, 'C': 182, 'N': 168, 'O': 151, 'S': 204,
+        'F': 158, 'Cl': 207, 'Br': 225, 'Si': 200, 'P': 210,
+        'I': 234
+    },
+    'Br': {
+        'Br': 228, 'H': 141, 'C': 194, 'O': 172, 'N': 214,
+        'Si': 215, 'S': 225, 'F': 178, 'Cl': 214, 'P': 222
+    },
+    'P': {
+        'P': 221, 'H': 144, 'C': 184, 'O': 163, 'Cl': 203,
+        'S': 210, 'F': 156, 'N': 177, 'Br': 222
+    },
+    'I': {
+        'H': 161, 'C': 214, 'Si': 243, 'N': 222, 'O': 194,
+        'S': 234, 'F': 187, 'I': 266
+    },
+    'As': {
+        'H': 152
+    }
+}
+BONDS_2 = {
+    'C': {'C': 134, 'N': 129, 'O': 120, 'S': 160},
+    'N': {'C': 129, 'N': 125, 'O': 121},
+    'O': {'C': 120, 'N': 121, 'O': 121, 'P': 150},
+    'P': {'O': 150, 'S': 186},
+    'S': {'P': 186}
+}
+BONDS_3 = {
+    'C': {'C': 120, 'N': 116, 'O': 113},
+    'N': {'C': 116, 'N': 110},
+    'O': {'C': 113}
+}
+BOND_DICT = [
+    None,
+    Chem.rdchem.BondType.SINGLE,
+    Chem.rdchem.BondType.DOUBLE,
+    Chem.rdchem.BondType.TRIPLE,
+    Chem.rdchem.BondType.AROMATIC,
+]
+BOND2IDX = {
+    Chem.rdchem.BondType.SINGLE: 1,
+    Chem.rdchem.BondType.DOUBLE: 2,
+    Chem.rdchem.BondType.TRIPLE: 3,
+    Chem.rdchem.BondType.AROMATIC: 4,
+}
+ALLOWED_BONDS = {
+    'H': 1,
+    'C': 4,
+    'N': 3,
+    'O': 2,
+    'F': 1,
+    'B': 3,
+    'Al': 3,
+    'Si': 4,
+    'P': [3, 5],
+    'S': 4,
+    'Cl': 1,
+    'As': 3,
+    'Br': 1,
+    'I': 1,
+    'Hg': [1, 2],
+    'Bi': [3, 5]
+}
+MARGINS_EDM = [10, 5, 2]
+COLORS = ['C0', 'C1', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8']
+# RADII = [0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3, 0.3]
+RADII = [0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77, 0.77]
+ZINC_TRAIN_LINKER_ID2SIZE = [3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
+ZINC_TRAIN_LINKER_SIZE2ID = {
+    size: idx
+    for idx, size in enumerate(ZINC_TRAIN_LINKER_ID2SIZE)
+}
+ZINC_TRAIN_LINKER_SIZE_WEIGHTS = [
+    3.47347831e-01,
+    4.63079100e-01,
+    5.12370917e-01,
+    5.62392614e-01,
+    1.30294388e+00,
+    3.24247801e+00,
+    8.12391184e+00,
+    3.45634358e+01,
+    2.72428571e+02,
+    6.26585714e+03
+]
+GEOM_TRAIN_LINKER_ID2SIZE = [
+    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, 36, 38, 41
+]
+GEOM_TRAIN_LINKER_SIZE2ID = {
+    size: idx
+    for idx, size in enumerate(GEOM_TRAIN_LINKER_ID2SIZE)
+}
+GEOM_TRAIN_LINKER_SIZE_WEIGHTS = [
+    1.07790681e+00, 4.54693604e-01, 3.62575713e-01, 3.75199484e-01,
+    3.67812588e-01, 3.92388528e-01, 3.83421054e-01, 4.26924670e-01,
+    4.92768040e-01, 4.99761944e-01, 4.92342726e-01, 5.71456905e-01,
+    7.30631393e-01, 8.45412928e-01, 9.97252243e-01, 1.25423985e+00,
+    1.57316129e+00, 2.19902962e+00, 3.22640431e+00, 4.25481066e+00,
+    6.34749573e+00, 9.00676236e+00, 1.43084017e+01, 2.25763173e+01,
+    3.36867096e+01, 9.50713805e+01, 2.08693274e+02, 2.51659537e+02,
+    7.77856749e+02, 8.55642424e+03, 8.55642424e+03, 4.27821212e+03,
+    4.27821212e+03
+]

src/datasets.py ADDED Viewed

	@@ -0,0 +1,350 @@

+import os
+import numpy as np
+import pandas as pd
+import pickle
+import torch
+from rdkit import Chem
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+from src import const
+from pdb import set_trace
+def read_sdf(sdf_path):
+    with Chem.SDMolSupplier(sdf_path, sanitize=False) as supplier:
+        for molecule in supplier:
+            yield molecule
+def get_one_hot(atom, atoms_dict):
+    one_hot = np.zeros(len(atoms_dict))
+    one_hot[atoms_dict[atom]] = 1
+    return one_hot
+def parse_molecule(mol, is_geom):
+    one_hot = []
+    charges = []
+    atom2idx = const.GEOM_ATOM2IDX if is_geom else const.ATOM2IDX
+    charges_dict = const.GEOM_CHARGES if is_geom else const.CHARGES
+    for atom in mol.GetAtoms():
+        one_hot.append(get_one_hot(atom.GetSymbol(), atom2idx))
+        charges.append(charges_dict[atom.GetSymbol()])
+    positions = mol.GetConformer().GetPositions()
+    return positions, np.array(one_hot), np.array(charges)
+class ZincDataset(Dataset):
+    def __init__(self, data_path, prefix, device):
+        dataset_path = os.path.join(data_path, f'{prefix}.pt')
+        if os.path.exists(dataset_path):
+            self.data = torch.load(dataset_path, map_location=device)
+        else:
+            print(f'Preprocessing dataset with prefix {prefix}')
+            self.data = ZincDataset.preprocess(data_path, prefix, device)
+            torch.save(self.data, dataset_path)
+    def __len__(self):
+        return len(self.data)
+    def __getitem__(self, item):
+        return self.data[item]
+    @staticmethod
+    def preprocess(data_path, prefix, device):
+        data = []
+        table_path = os.path.join(data_path, f'{prefix}_table.csv')
+        fragments_path = os.path.join(data_path, f'{prefix}_frag.sdf')
+        linkers_path = os.path.join(data_path, f'{prefix}_link.sdf')
+        is_geom = ('geom' in prefix) or ('MOAD' in prefix)
+        is_multifrag = 'multifrag' in prefix
+        table = pd.read_csv(table_path)
+        generator = tqdm(zip(table.iterrows(), read_sdf(fragments_path), read_sdf(linkers_path)), total=len(table))
+        for (_, row), fragments, linker in generator:
+            uuid = row['uuid']
+            name = row['molecule']
+            frag_pos, frag_one_hot, frag_charges = parse_molecule(fragments, is_geom=is_geom)
+            link_pos, link_one_hot, link_charges = parse_molecule(linker, is_geom=is_geom)
+            positions = np.concatenate([frag_pos, link_pos], axis=0)
+            one_hot = np.concatenate([frag_one_hot, link_one_hot], axis=0)
+            charges = np.concatenate([frag_charges, link_charges], axis=0)
+            anchors = np.zeros_like(charges)
+            if is_multifrag:
+                for anchor_idx in map(int, row['anchors'].split('-')):
+                    anchors[anchor_idx] = 1
+            else:
+                anchors[row['anchor_1']] = 1
+                anchors[row['anchor_2']] = 1
+            fragment_mask = np.concatenate([np.ones_like(frag_charges), np.zeros_like(link_charges)])
+            linker_mask = np.concatenate([np.zeros_like(frag_charges), np.ones_like(link_charges)])
+            data.append({
+                'uuid': uuid,
+                'name': name,
+                'positions': torch.tensor(positions, dtype=const.TORCH_FLOAT, device=device),
+                'one_hot': torch.tensor(one_hot, dtype=const.TORCH_FLOAT, device=device),
+                'charges': torch.tensor(charges, dtype=const.TORCH_FLOAT, device=device),
+                'anchors': torch.tensor(anchors, dtype=const.TORCH_FLOAT, device=device),
+                'fragment_mask': torch.tensor(fragment_mask, dtype=const.TORCH_FLOAT, device=device),
+                'linker_mask': torch.tensor(linker_mask, dtype=const.TORCH_FLOAT, device=device),
+                'num_atoms': len(positions),
+            })
+        return data
+class MOADDataset(Dataset):
+    def __init__(self, data_path, prefix, device):
+        prefix, pocket_mode = prefix.split('.')
+        dataset_path = os.path.join(data_path, f'{prefix}_{pocket_mode}.pt')
+        if os.path.exists(dataset_path):
+            self.data = torch.load(dataset_path, map_location=device)
+        else:
+            print(f'Preprocessing dataset with prefix {prefix}')
+            self.data = MOADDataset.preprocess(data_path, prefix, pocket_mode, device)
+            torch.save(self.data, dataset_path)
+    def __len__(self):
+        return len(self.data)
+    def __getitem__(self, item):
+        return self.data[item]
+    @staticmethod
+    def preprocess(data_path, prefix, pocket_mode, device):
+        data = []
+        table_path = os.path.join(data_path, f'{prefix}_table.csv')
+        fragments_path = os.path.join(data_path, f'{prefix}_frag.sdf')
+        linkers_path = os.path.join(data_path, f'{prefix}_link.sdf')
+        pockets_path = os.path.join(data_path, f'{prefix}_pockets.pkl')
+        is_geom = True
+        is_multifrag = 'multifrag' in prefix
+        with open(pockets_path, 'rb') as f:
+            pockets = pickle.load(f)
+        table = pd.read_csv(table_path)
+        generator = tqdm(
+            zip(table.iterrows(), read_sdf(fragments_path), read_sdf(linkers_path), pockets),
+            total=len(table)
+        )
+        for (_, row), fragments, linker, pocket_data in generator:
+            uuid = row['uuid']
+            name = row['molecule']
+            frag_pos, frag_one_hot, frag_charges = parse_molecule(fragments, is_geom=is_geom)
+            link_pos, link_one_hot, link_charges = parse_molecule(linker, is_geom=is_geom)
+            # Parsing pocket data
+            pocket_pos = pocket_data[f'{pocket_mode}_coord']
+            pocket_one_hot = []
+            pocket_charges = []
+            for atom_type in pocket_data[f'{pocket_mode}_types']:
+                pocket_one_hot.append(get_one_hot(atom_type, const.GEOM_ATOM2IDX))
+                pocket_charges.append(const.GEOM_CHARGES[atom_type])
+            pocket_one_hot = np.array(pocket_one_hot)
+            pocket_charges = np.array(pocket_charges)
+            positions = np.concatenate([frag_pos, pocket_pos, link_pos], axis=0)
+            one_hot = np.concatenate([frag_one_hot, pocket_one_hot, link_one_hot], axis=0)
+            charges = np.concatenate([frag_charges, pocket_charges, link_charges], axis=0)
+            anchors = np.zeros_like(charges)
+            if is_multifrag:
+                for anchor_idx in map(int, row['anchors'].split('-')):
+                    anchors[anchor_idx] = 1
+            else:
+                anchors[row['anchor_1']] = 1
+                anchors[row['anchor_2']] = 1
+            fragment_only_mask = np.concatenate([
+                np.ones_like(frag_charges),
+                np.zeros_like(pocket_charges),
+                np.zeros_like(link_charges)
+            ])
+            pocket_mask = np.concatenate([
+                np.zeros_like(frag_charges),
+                np.ones_like(pocket_charges),
+                np.zeros_like(link_charges)
+            ])
+            linker_mask = np.concatenate([
+                np.zeros_like(frag_charges),
+                np.zeros_like(pocket_charges),
+                np.ones_like(link_charges)
+            ])
+            fragment_mask = np.concatenate([
+                np.ones_like(frag_charges),
+                np.ones_like(pocket_charges),
+                np.zeros_like(link_charges)
+            ])
+            data.append({
+                'uuid': uuid,
+                'name': name,
+                'positions': torch.tensor(positions, dtype=const.TORCH_FLOAT, device=device),
+                'one_hot': torch.tensor(one_hot, dtype=const.TORCH_FLOAT, device=device),
+                'charges': torch.tensor(charges, dtype=const.TORCH_FLOAT, device=device),
+                'anchors': torch.tensor(anchors, dtype=const.TORCH_FLOAT, device=device),
+                'fragment_only_mask': torch.tensor(fragment_only_mask, dtype=const.TORCH_FLOAT, device=device),
+                'pocket_mask': torch.tensor(pocket_mask, dtype=const.TORCH_FLOAT, device=device),
+                'fragment_mask': torch.tensor(fragment_mask, dtype=const.TORCH_FLOAT, device=device),
+                'linker_mask': torch.tensor(linker_mask, dtype=const.TORCH_FLOAT, device=device),
+                'num_atoms': len(positions),
+            })
+        return data
+    @staticmethod
+    def create_edges(positions, fragment_mask_only, linker_mask_only):
+        ligand_mask = fragment_mask_only.astype(bool) | linker_mask_only.astype(bool)
+        ligand_adj = ligand_mask[:, None] & ligand_mask[None, :]
+        proximity_adj = np.linalg.norm(positions[:, None, :] - positions[None, :, :], axis=-1) <= 6
+        full_adj = ligand_adj | proximity_adj
+        full_adj &= ~np.eye(len(positions)).astype(bool)
+        curr_rows, curr_cols = np.where(full_adj)
+        return [curr_rows, curr_cols]
+def collate(batch):
+    out = {}
+    # Filter out big molecules
+    if 'pocket_mask' not in batch[0].keys():
+        batch = [data for data in batch if data['num_atoms'] <= 50]
+    else:
+        batch = [data for data in batch if data['num_atoms'] <= 1000]
+    for i, data in enumerate(batch):
+        for key, value in data.items():
+            out.setdefault(key, []).append(value)
+    for key, value in out.items():
+        if key in const.DATA_LIST_ATTRS:
+            continue
+        if key in const.DATA_ATTRS_TO_PAD:
+            out[key] = torch.nn.utils.rnn.pad_sequence(value, batch_first=True, padding_value=0)
+            continue
+        raise Exception(f'Unknown batch key: {key}')
+    atom_mask = (out['fragment_mask'].bool() | out['linker_mask'].bool()).to(const.TORCH_INT)
+    out['atom_mask'] = atom_mask[:, :, None]
+    batch_size, n_nodes = atom_mask.size()
+    # In case of MOAD edge_mask is batch_idx
+    if 'pocket_mask' in batch[0].keys():
+        batch_mask = torch.cat([
+            torch.ones(n_nodes, dtype=const.TORCH_INT) * i
+            for i in range(batch_size)
+        ]).to(atom_mask.device)
+        out['edge_mask'] = batch_mask
+    else:
+        edge_mask = atom_mask[:, None, :] * atom_mask[:, :, None]
+        diag_mask = ~torch.eye(edge_mask.size(1), dtype=const.TORCH_INT, device=atom_mask.device).unsqueeze(0)
+        edge_mask *= diag_mask
+        out['edge_mask'] = edge_mask.view(batch_size * n_nodes * n_nodes, 1)
+    for key in const.DATA_ATTRS_TO_ADD_LAST_DIM:
+        if key in out.keys():
+            out[key] = out[key][:, :, None]
+    return out
+def collate_with_fragment_edges(batch):
+    out = {}
+    # Filter out big molecules
+    batch = [data for data in batch if data['num_atoms'] <= 50]
+    for i, data in enumerate(batch):
+        for key, value in data.items():
+            out.setdefault(key, []).append(value)
+    for key, value in out.items():
+        if key in const.DATA_LIST_ATTRS:
+            continue
+        if key in const.DATA_ATTRS_TO_PAD:
+            out[key] = torch.nn.utils.rnn.pad_sequence(value, batch_first=True, padding_value=0)
+            continue
+        raise Exception(f'Unknown batch key: {key}')
+    frag_mask = out['fragment_mask']
+    edge_mask = frag_mask[:, None, :] * frag_mask[:, :, None]
+    diag_mask = ~torch.eye(edge_mask.size(1), dtype=const.TORCH_INT, device=frag_mask.device).unsqueeze(0)
+    edge_mask *= diag_mask
+    batch_size, n_nodes = frag_mask.size()
+    out['edge_mask'] = edge_mask.view(batch_size * n_nodes * n_nodes, 1)
+    # Building edges and covalent bond values
+    rows, cols, bonds = [], [], []
+    for batch_idx in range(batch_size):
+        for i in range(n_nodes):
+            for j in range(n_nodes):
+                rows.append(i + batch_idx * n_nodes)
+                cols.append(j + batch_idx * n_nodes)
+    edges = [torch.LongTensor(rows).to(frag_mask.device), torch.LongTensor(cols).to(frag_mask.device)]
+    out['edges'] = edges
+    atom_mask = (out['fragment_mask'].bool() | out['linker_mask'].bool()).to(const.TORCH_INT)
+    out['atom_mask'] = atom_mask[:, :, None]
+    for key in const.DATA_ATTRS_TO_ADD_LAST_DIM:
+        if key in out.keys():
+            out[key] = out[key][:, :, None]
+    return out
+def get_dataloader(dataset, batch_size, collate_fn=collate, shuffle=False):
+    return DataLoader(dataset, batch_size, collate_fn=collate_fn, shuffle=shuffle)
+def create_template(tensor, fragment_size, linker_size, fill=0):
+    values_to_keep = tensor[:fragment_size]
+    values_to_add = torch.ones(linker_size, tensor.shape[1], dtype=values_to_keep.dtype, device=values_to_keep.device)
+    values_to_add = values_to_add * fill
+    return torch.cat([values_to_keep, values_to_add], dim=0)
+def create_templates_for_linker_generation(data, linker_sizes):
+    """
+    Takes data batch and new linker size and returns data batch where fragment-related data is the same
+    but linker-related data is replaced with zero templates with new linker sizes
+    """
+    decoupled_data = []
+    for i, linker_size in enumerate(linker_sizes):
+        data_dict = {}
+        fragment_mask = data['fragment_mask'][i].squeeze()
+        fragment_size = fragment_mask.sum().int()
+        for k, v in data.items():
+            if k == 'num_atoms':
+                # Computing new number of atoms (fragment_size + linker_size)
+                data_dict[k] = fragment_size + linker_size
+                continue
+            if k in const.DATA_LIST_ATTRS:
+                # These attributes are written without modification
+                data_dict[k] = v[i]
+                continue
+            if k in const.DATA_ATTRS_TO_PAD:
+                # Should write fragment-related data + (zeros x linker_size)
+                fill_value = 1 if k == 'linker_mask' else 0
+                template = create_template(v[i], fragment_size, linker_size, fill=fill_value)
+                if k in const.DATA_ATTRS_TO_ADD_LAST_DIM:
+                    template = template.squeeze(-1)
+                data_dict[k] = template
+        decoupled_data.append(data_dict)
+    return collate(decoupled_data)

src/delinker.py ADDED Viewed

	@@ -0,0 +1,278 @@

+import csv
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem import MolStandardize
+from src import metrics
+from src.delinker_utils import sascorer, calc_SC_RDKit
+from tqdm import tqdm
+from pdb import set_trace
+def get_valid_as_in_delinker(data, progress=False):
+    valid = []
+    generator = tqdm(enumerate(data), total=len(data)) if progress else enumerate(data)
+    for i, m in generator:
+        pred_mol = Chem.MolFromSmiles(m['pred_mol_smi'], sanitize=False)
+        true_mol = Chem.MolFromSmiles(m['true_mol_smi'], sanitize=False)
+        frag = Chem.MolFromSmiles(m['frag_smi'], sanitize=False)
+        pred_mol_frags = Chem.GetMolFrags(pred_mol, asMols=True, sanitizeFrags=False)
+        pred_mol_filtered = max(pred_mol_frags, default=pred_mol, key=lambda mol: mol.GetNumAtoms())
+        try:
+            Chem.SanitizeMol(pred_mol_filtered)
+            Chem.SanitizeMol(true_mol)
+            Chem.SanitizeMol(frag)
+        except:
+            continue
+        if len(pred_mol_filtered.GetSubstructMatch(frag)) > 0:
+            valid.append({
+                'pred_mol': m['pred_mol'],
+                'true_mol': m['true_mol'],
+                'pred_mol_smi': Chem.MolToSmiles(pred_mol_filtered),
+                'true_mol_smi': Chem.MolToSmiles(true_mol),
+                'frag_smi': Chem.MolToSmiles(frag)
+            })
+    return valid
+def extract_linker_smiles(molecule, fragments):
+    match = molecule.GetSubstructMatch(fragments)
+    elinker = Chem.EditableMol(molecule)
+    for atom_id in sorted(match, reverse=True):
+        elinker.RemoveAtom(atom_id)
+    linker = elinker.GetMol()
+    Chem.RemoveStereochemistry(linker)
+    try:
+        linker = MolStandardize.canonicalize_tautomer_smiles(Chem.MolToSmiles(linker))
+    except:
+        linker = Chem.MolToSmiles(linker)
+    return linker
+def compute_and_add_linker_smiles(data, progress=False):
+    data_with_linkers = []
+    generator = tqdm(data) if progress else data
+    for m in generator:
+        pred_mol = Chem.MolFromSmiles(m['pred_mol_smi'], sanitize=True)
+        true_mol = Chem.MolFromSmiles(m['true_mol_smi'], sanitize=True)
+        frag = Chem.MolFromSmiles(m['frag_smi'], sanitize=True)
+        pred_linker = extract_linker_smiles(pred_mol, frag)
+        true_linker = extract_linker_smiles(true_mol, frag)
+        data_with_linkers.append({
+            **m,
+            'pred_linker': pred_linker,
+            'true_linker': true_linker,
+        })
+    return data_with_linkers
+def compute_uniqueness(data, progress=False):
+    mol_dictionary = {}
+    generator = tqdm(data) if progress else data
+    for m in generator:
+        frag = m['frag_smi']
+        pred_mol = m['pred_mol_smi']
+        true_mol = m['true_mol_smi']
+        key = f'{true_mol}.{frag}'
+        mol_dictionary.setdefault(key, []).append(pred_mol)
+    total_mol = 0
+    unique_mol = 0
+    for molecules in mol_dictionary.values():
+        total_mol += len(molecules)
+        unique_mol += len(set(molecules))
+    return unique_mol / total_mol
+def compute_novelty(data, progress=False):
+    novel = 0
+    true_linkers = set([m['true_linker'] for m in data])
+    generator = tqdm(data) if progress else data
+    for m in generator:
+        pred_linker = m['pred_linker']
+        if pred_linker in true_linkers:
+            continue
+        else:
+            novel += 1
+    return novel / len(data)
+def compute_recovery_rate(data, progress=False):
+    total = set()
+    recovered = set()
+    generator = tqdm(data) if progress else data
+    for m in generator:
+        pred_mol = Chem.MolFromSmiles(m['pred_mol_smi'], sanitize=True)
+        Chem.RemoveStereochemistry(pred_mol)
+        pred_mol = Chem.MolToSmiles(Chem.RemoveHs(pred_mol))
+        true_mol = Chem.MolFromSmiles(m['true_mol_smi'], sanitize=True)
+        Chem.RemoveStereochemistry(true_mol)
+        true_mol = Chem.MolToSmiles(Chem.RemoveHs(true_mol))
+        true_link = m['true_linker']
+        total.add(f'{true_mol}.{true_link}')
+        if pred_mol == true_mol:
+            recovered.add(f'{true_mol}.{true_link}')
+    return len(recovered) / len(total)
+def calc_sa_score_mol(mol):
+    if mol is None:
+        return None
+    return sascorer.calculateScore(mol)
+def check_ring_filter(linker):
+    check = True
+    # Get linker rings
+    ssr = Chem.GetSymmSSSR(linker)
+    # Check rings
+    for ring in ssr:
+        for atom_idx in ring:
+            for bond in linker.GetAtomWithIdx(atom_idx).GetBonds():
+                if bond.GetBondType() == 2 and bond.GetBeginAtomIdx() in ring and bond.GetEndAtomIdx() in ring:
+                    check = False
+    return check
+def check_pains(mol, pains_smarts):
+    for pain in pains_smarts:
+        if mol.HasSubstructMatch(pain):
+            return False
+    return True
+def calc_2d_filters(toks, pains_smarts):
+    pred_mol = Chem.MolFromSmiles(toks['pred_mol_smi'])
+    frag = Chem.MolFromSmiles(toks['frag_smi'])
+    linker = Chem.MolFromSmiles(toks['pred_linker'])
+    result = [False, False, False]
+    if len(pred_mol.GetSubstructMatch(frag)) > 0:
+        sa_score = False
+        ra_score = False
+        pains_score = False
+        try:
+            sa_score = calc_sa_score_mol(pred_mol) < calc_sa_score_mol(frag)
+        except Exception as e:
+            print(f'Could not compute SA score: {e}')
+        try:
+            ra_score = check_ring_filter(linker)
+        except Exception as e:
+            print(f'Could not compute RA score: {e}')
+        try:
+            pains_score = check_pains(pred_mol, pains_smarts)
+        except Exception as e:
+            print(f'Could not compute PAINS score: {e}')
+        result = [sa_score, ra_score, pains_score]
+    return result
+def calc_filters_2d_dataset(data):
+    with open('models/wehi_pains.csv', 'r') as f:
+        pains_smarts = [Chem.MolFromSmarts(line[0], mergeHs=True) for line in csv.reader(f)]
+    pass_all = pass_SA = pass_RA = pass_PAINS = 0
+    for m in data:
+        filters_2d = calc_2d_filters(m, pains_smarts)
+        pass_all += filters_2d[0] & filters_2d[1] & filters_2d[2]
+        pass_SA += filters_2d[0]
+        pass_RA += filters_2d[1]
+        pass_PAINS += filters_2d[2]
+    return pass_all / len(data), pass_SA / len(data), pass_RA / len(data), pass_PAINS / len(data)
+def calc_sc_rdkit_full_mol(gen_mol, ref_mol):
+    try:
+        score = calc_SC_RDKit.calc_SC_RDKit_score(gen_mol, ref_mol)
+        return score
+    except:
+        return -0.5
+def sc_rdkit_score(data):
+    scores = []
+    for m in data:
+        score = calc_sc_rdkit_full_mol(m['pred_mol'], m['true_mol'])
+        scores.append(score)
+    return np.mean(scores)
+def get_delinker_metrics(pred_molecules, true_molecules, true_fragments):
+    default_values = {
+        'DeLinker/validity': 0,
+        'DeLinker/uniqueness': 0,
+        'DeLinker/novelty': 0,
+        'DeLinker/recovery': 0,
+        'DeLinker/2D_filters': 0,
+        'DeLinker/2D_filters_SA': 0,
+        'DeLinker/2D_filters_RA': 0,
+        'DeLinker/2D_filters_PAINS': 0,
+        'DeLinker/SC_RDKit': 0,
+    }
+    if len(pred_molecules) == 0:
+        return default_values
+    data = []
+    for pred_mol, true_mol, true_frag in zip(pred_molecules, true_molecules, true_fragments):
+        data.append({
+            'pred_mol': pred_mol,
+            'true_mol': true_mol,
+            'pred_mol_smi': Chem.MolToSmiles(pred_mol),
+            'true_mol_smi': Chem.MolToSmiles(true_mol),
+            'frag_smi': Chem.MolToSmiles(true_frag)
+        })
+    # Validity according to DeLinker paper:
+    # Passing rdkit.Chem.Sanitize and the biggest fragment contains both fragments
+    valid_data = get_valid_as_in_delinker(data)
+    validity_as_in_delinker = len(valid_data) / len(data)
+    if len(valid_data) == 0:
+        return default_values
+    # Compute linkers and add to results
+    valid_data = compute_and_add_linker_smiles(valid_data)
+    # Compute uniqueness
+    uniqueness = compute_uniqueness(valid_data)
+    # Compute novelty
+    novelty = compute_novelty(valid_data)
+    # Compute recovered molecules
+    recovery_rate = compute_recovery_rate(valid_data)
+    # 2D filters
+    pass_all, pass_SA, pass_RA, pass_PAINS = calc_filters_2d_dataset(valid_data)
+    # 3D Filters
+    sc_rdkit = sc_rdkit_score(valid_data)
+    return {
+        'DeLinker/validity': validity_as_in_delinker,
+        'DeLinker/uniqueness': uniqueness,
+        'DeLinker/novelty': novelty,
+        'DeLinker/recovery': recovery_rate,
+        'DeLinker/2D_filters': pass_all,
+        'DeLinker/2D_filters_SA': pass_SA,
+        'DeLinker/2D_filters_RA': pass_RA,
+        'DeLinker/2D_filters_PAINS': pass_PAINS,
+        'DeLinker/SC_RDKit': sc_rdkit,
+    }

src/delinker_utils/__init__.py ADDED Viewed

File without changes

src/delinker_utils/calc_SC_RDKit.py ADDED Viewed

	@@ -0,0 +1,40 @@

+import os
+from rdkit import Chem
+from rdkit.Chem import AllChem, rdShapeHelpers
+from rdkit.Chem.FeatMaps import FeatMaps
+from rdkit import RDConfig
+# Set up features to use in FeatureMap
+fdefName = os.path.join(RDConfig.RDDataDir, 'BaseFeatures.fdef')
+fdef = AllChem.BuildFeatureFactory(fdefName)
+fmParams = {}
+for k in fdef.GetFeatureFamilies():
+    fparams = FeatMaps.FeatMapParams()
+    fmParams[k] = fparams
+keep = ('Donor', 'Acceptor', 'NegIonizable', 'PosIonizable',
+        'ZnBinder', 'Aromatic', 'Hydrophobe', 'LumpedHydrophobe')
+def get_FeatureMapScore(query_mol, ref_mol):
+    featLists = []
+    for m in [query_mol, ref_mol]:
+        rawFeats = fdef.GetFeaturesForMol(m)
+        # filter that list down to only include the ones we're intereted in
+        featLists.append([f for f in rawFeats if f.GetFamily() in keep])
+    fms = [FeatMaps.FeatMap(feats=x, weights=[1] * len(x), params=fmParams) for x in featLists]
+    fms[0].scoreMode = FeatMaps.FeatMapScoreMode.Best
+    fm_score = fms[0].ScoreFeats(featLists[1]) / min(fms[0].GetNumFeatures(), len(featLists[1]))
+    return fm_score
+def calc_SC_RDKit_score(query_mol, ref_mol):
+    fm_score = get_FeatureMapScore(query_mol, ref_mol)
+    protrude_dist = rdShapeHelpers.ShapeProtrudeDist(query_mol, ref_mol,
+                                                     allowReordering=False)
+    SC_RDKit_score = 0.5 * fm_score + 0.5 * (1 - protrude_dist)
+    return SC_RDKit_score

src/delinker_utils/frag_utils.py ADDED Viewed

	@@ -0,0 +1,413 @@

+import csv
+import networkx as nx
+from joblib import Parallel, delayed
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from src.delinker_utils import sascorer
+def read_triples_file(filename):
+    '''Reads .smi file '''
+    '''Returns array containing smiles strings of molecules'''
+    smiles, names = [], []
+    with open(filename, 'r') as f:
+        for line in f:
+            if line:
+                smiles.append(line.strip().split(' ')[0:3])
+    return smiles
+def remove_dummys(smi_string):
+    return Chem.MolToSmiles(Chem.RemoveHs(AllChem.ReplaceSubstructs(Chem.MolFromSmiles(smi_string),Chem.MolFromSmiles('*'),Chem.MolFromSmiles('[H]'),True)[0]))
+def sa_filter(results, verbose=True):
+    count = 0
+    total = 0
+    for processed, res in enumerate(results):
+        total += len(res)
+        for m in res:
+            # Check SA score has improved
+            if calc_mol_props(m[1])[1] < calc_mol_props(m[0])[1]:
+                count += 1
+        # Progress
+        if verbose:
+            if processed % 10 == 0:
+                print("\rProcessed %d" % processed, end="")
+    print("\r",end="")
+    return count/total
+def ring_check_res(res, clean_frag):
+    check = True
+    gen_mol = Chem.MolFromSmiles(res[1])
+    linker = Chem.DeleteSubstructs(gen_mol, clean_frag)
+    # Get linker rings
+    ssr = Chem.GetSymmSSSR(linker)
+    # Check rings
+    for ring in ssr:
+        for atom_idx in ring:
+            for bond in linker.GetAtomWithIdx(atom_idx).GetBonds():
+                if bond.GetBondType() == 2 and bond.GetBeginAtomIdx() in ring and bond.GetEndAtomIdx() in ring:
+                    check = False
+    return check
+def ring_filter(results, verbose=True):
+    count = 0
+    total = 0
+    du = Chem.MolFromSmiles('*')
+    for processed, res in enumerate(results):
+        total += len(res)
+        for m in res:
+            # Clean frags
+            clean_frag = Chem.RemoveHs(AllChem.ReplaceSubstructs(Chem.MolFromSmiles(m[0]),du,Chem.MolFromSmiles('[H]'),True)[0])
+            if ring_check_res(m, clean_frag):
+                count += 1
+        # Progress
+        if verbose:
+            if processed % 10 == 0:
+                print("\rProcessed %d" % processed, end="")
+    print("\r",end="")
+    return count/total
+def check_ring_filter(linker):
+    check = True
+    # Get linker rings
+    ssr = Chem.GetSymmSSSR(linker)
+    # Check rings
+    for ring in ssr:
+        for atom_idx in ring:
+            for bond in linker.GetAtomWithIdx(atom_idx).GetBonds():
+                if bond.GetBondType() == 2 and bond.GetBeginAtomIdx() in ring and bond.GetEndAtomIdx() in ring:
+                    check = False
+    return check
+def check_pains(mol, pains_smarts):
+    for pain in pains_smarts:
+        if mol.HasSubstructMatch(pain):
+            return False
+    return True
+def calc_2d_filters(toks, pains_smarts):
+    try:
+        # Input format: (Full Molecule (SMILES), Linker (SMILES), Unlinked Fragments (SMILES))
+        frags = Chem.MolFromSmiles(toks[2])
+        linker = Chem.MolFromSmiles(toks[1])
+        full_mol = Chem.MolFromSmiles(toks[0])
+        # Remove dummy atoms from unlinked fragments
+        du = Chem.MolFromSmiles('*')
+        clean_frag = Chem.RemoveHs(AllChem.ReplaceSubstructs(frags, du, Chem.MolFromSmiles('[H]'), True)[0])
+        res = []
+        # Check: Unlinked fragments in full molecule
+        if len(full_mol.GetSubstructMatch(clean_frag)) > 0:
+            # Check: SA score improved from unlinked fragments to full molecule
+            if calc_sa_score_mol(full_mol) < calc_sa_score_mol(frags):
+                res.append(True)
+            else:
+                res.append(False)
+            # Check: No non-aromatic rings with double bonds
+            if check_ring_filter(linker):
+                res.append(True)
+            else:
+                res.append(False)
+            # Check: Pass pains filters
+            if check_pains(full_mol, pains_smarts):
+                res.append(True)
+            else:
+                res.append(False)
+        return res
+    except:
+        return [False, False, False]
+def calc_filters_2d_dataset(results, pains_smarts_loc, n_cores=1):
+    # Load pains filters
+    with open(pains_smarts_loc, 'r') as f:
+        pains_smarts = [Chem.MolFromSmarts(line[0], mergeHs=True) for line in csv.reader(f)]
+    # calc_2d_filters([results[0][2], results[0][4], results[0][1]], pains_smarts)
+    with Parallel(n_jobs=n_cores, backend='multiprocessing') as parallel:
+        filters_2d = parallel(delayed(calc_2d_filters)([toks[2], toks[4], toks[1]], pains_smarts) for toks in results)
+    return filters_2d
+def calc_mol_props(smiles):
+    # Create RDKit mol
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        print("Error passing: %s" % smiles)
+        return None
+    # QED
+    qed = Chem.QED.qed(mol)
+    # Synthetic accessibility score - number of cycles (rings with > 6 atoms)
+    sas = sascorer.calculateScore(mol)
+    # Cyles with >6 atoms
+    ri = mol.GetRingInfo()
+    nMacrocycles = 0
+    for x in ri.AtomRings():
+        if len(x) > 6:
+            nMacrocycles += 1
+    prop_array = [qed, sas]
+    return prop_array
+def calc_sa_score_mol(mol, verbose=False):
+    if mol is None:
+        if verbose:
+            print("Error passing: %s" % mol)
+        return None
+    # Synthetic accessibility score
+    return sascorer.calculateScore(mol)
+def get_linker(full_mol, clean_frag, starting_point):
+    # INPUT FORMAT: molecule (RDKit mol object), clean fragments (RDKit mol object), starting fragments (SMILES)
+    # Get matches of fragments
+    matches = list(full_mol.GetSubstructMatches(clean_frag))
+    # If no matches, terminate
+    if len(matches) == 0:
+        print("No matches")
+        return ""
+    # Get number of atoms in linker
+    linker_len = full_mol.GetNumHeavyAtoms() - clean_frag.GetNumHeavyAtoms()
+    if linker_len == 0:
+        return ""
+    # Setup
+    mol_to_break = Chem.Mol(full_mol)
+    Chem.Kekulize(full_mol, clearAromaticFlags=True)
+    poss_linker = []
+    if len(matches) > 0:
+        # Loop over matches
+        for match in matches:
+            mol_rw = Chem.RWMol(full_mol)
+            # Get linker atoms
+            linker_atoms = list(set(list(range(full_mol.GetNumHeavyAtoms()))).difference(match))
+            linker_bonds = []
+            atoms_joined_to_linker = []
+            # Loop over starting fragments atoms
+            # Get (i) bonds between starting fragments and linker, (ii) atoms joined to linker
+            for idx_to_delete in sorted(match, reverse=True):
+                nei = [x.GetIdx() for x in mol_rw.GetAtomWithIdx(idx_to_delete).GetNeighbors()]
+                intersect = set(nei).intersection(set(linker_atoms))
+                if len(intersect) == 1:
+                    linker_bonds.append(mol_rw.GetBondBetweenAtoms(idx_to_delete, list(intersect)[0]).GetIdx())
+                    atoms_joined_to_linker.append(idx_to_delete)
+                elif len(intersect) > 1:
+                    for idx_nei in list(intersect):
+                        linker_bonds.append(mol_rw.GetBondBetweenAtoms(idx_to_delete, idx_nei).GetIdx())
+                        atoms_joined_to_linker.append(idx_to_delete)
+            # Check number of atoms joined to linker
+            # If not == 2, check next match
+            if len(set(atoms_joined_to_linker)) != 2:
+                continue
+            # Delete starting fragments atoms
+            for idx_to_delete in sorted(match, reverse=True):
+                mol_rw.RemoveAtom(idx_to_delete)
+            linker = Chem.Mol(mol_rw)
+            # Check linker required num atoms
+            if linker.GetNumHeavyAtoms() == linker_len:
+                mol_rw = Chem.RWMol(full_mol)
+                # Delete linker atoms
+                for idx_to_delete in sorted(linker_atoms, reverse=True):
+                    mol_rw.RemoveAtom(idx_to_delete)
+                frags = Chem.Mol(mol_rw)
+                # Check there are two disconnected fragments
+                if len(Chem.rdmolops.GetMolFrags(frags)) == 2:
+                    # Fragment molecule into starting fragments and linker
+                    fragmented_mol = Chem.FragmentOnBonds(mol_to_break, linker_bonds)
+                    # Remove starting fragments from fragmentation
+                    linker_to_return = Chem.Mol(fragmented_mol)
+                    qp = Chem.AdjustQueryParameters()
+                    qp.makeDummiesQueries = True
+                    for f in starting_point.split('.'):
+                        qfrag = Chem.AdjustQueryProperties(Chem.MolFromSmiles(f), qp)
+                        linker_to_return = AllChem.DeleteSubstructs(linker_to_return, qfrag, onlyFrags=True)
+                    # Check linker is connected and two bonds to outside molecule
+                    if len(Chem.rdmolops.GetMolFrags(linker)) == 1 and len(linker_bonds) == 2:
+                        Chem.Kekulize(linker_to_return, clearAromaticFlags=True)
+                        # If for some reason a starting fragment isn't removed (and it's larger than the linker), remove (happens v. occassionally)
+                        if len(Chem.rdmolops.GetMolFrags(linker_to_return)) > 1:
+                            for frag in Chem.MolToSmiles(linker_to_return).split('.'):
+                                if Chem.MolFromSmiles(frag).GetNumHeavyAtoms() == linker_len:
+                                    return frag
+                        return Chem.MolToSmiles(Chem.MolFromSmiles(Chem.MolToSmiles(linker_to_return)))
+                    # If not, add to possible linkers (above doesn't capture some complex cases)
+                    else:
+                        fragmented_mol = Chem.MolFromSmiles(Chem.MolToSmiles(fragmented_mol), sanitize=False)
+                        linker_to_return = AllChem.DeleteSubstructs(fragmented_mol, Chem.MolFromSmiles(starting_point))
+                        poss_linker.append(Chem.MolToSmiles(linker_to_return))
+    # If only one possibility, return linker
+    if len(poss_linker) == 1:
+        return poss_linker[0]
+    # If no possibilities, process failed
+    elif len(poss_linker) == 0:
+        print("FAIL:", Chem.MolToSmiles(full_mol), Chem.MolToSmiles(clean_frag), starting_point)
+        return ""
+    # If multiple possibilities, process probably failed
+    else:
+        print("More than one poss linker. ", poss_linker)
+        return poss_linker[0]
+def get_linker_v2(full_mol, clean_frag):
+    # INPUT FORMAT: molecule (RDKit mol object), clean fragments (RDKit mol object), starting fragments (SMILES)
+    # Get matches of fragments
+    matches = list(full_mol.GetSubstructMatches(clean_frag))
+    # If no matches, terminate
+    if len(matches) == 0:
+        print("No matches")
+        return ""
+    # Get number of atoms in linker
+    linker_len = full_mol.GetNumHeavyAtoms() - clean_frag.GetNumHeavyAtoms()
+    if linker_len == 0:
+        return ""
+    # Setup
+    mol_to_break = Chem.Mol(full_mol)
+    Chem.Kekulize(full_mol, clearAromaticFlags=True)
+    poss_linker = []
+    if len(matches) > 0:
+        # Loop over matches
+        for match in matches:
+            mol_rw = Chem.RWMol(full_mol)
+            # Get linker atoms
+            linker_atoms = list(set(list(range(full_mol.GetNumHeavyAtoms()))).difference(match))
+            linker_bonds = []
+            atoms_joined_to_linker = []
+            # Loop over starting fragments atoms
+            # Get (i) bonds between starting fragments and linker, (ii) atoms joined to linker
+            for idx_to_delete in sorted(match, reverse=True):
+                nei = [x.GetIdx() for x in mol_rw.GetAtomWithIdx(idx_to_delete).GetNeighbors()]
+                intersect = set(nei).intersection(set(linker_atoms))
+                if len(intersect) == 1:
+                    linker_bonds.append(mol_rw.GetBondBetweenAtoms(idx_to_delete, list(intersect)[0]).GetIdx())
+                    atoms_joined_to_linker.append(idx_to_delete)
+                elif len(intersect) > 1:
+                    for idx_nei in list(intersect):
+                        linker_bonds.append(mol_rw.GetBondBetweenAtoms(idx_to_delete, idx_nei).GetIdx())
+                        atoms_joined_to_linker.append(idx_to_delete)
+            # Check number of atoms joined to linker
+            # If not == 2, check next match
+            if len(set(atoms_joined_to_linker)) != 2:
+                continue
+            # Delete starting fragments atoms
+            for idx_to_delete in sorted(match, reverse=True):
+                mol_rw.RemoveAtom(idx_to_delete)
+            linker = Chem.Mol(mol_rw)
+            # Check linker required num atoms
+            if linker.GetNumHeavyAtoms() == linker_len:
+                mol_rw = Chem.RWMol(full_mol)
+                # Delete linker atoms
+                for idx_to_delete in sorted(linker_atoms, reverse=True):
+                    mol_rw.RemoveAtom(idx_to_delete)
+                frags = Chem.Mol(mol_rw)
+                # Check linker is connected and two bonds to outside molecule
+                if len(Chem.rdmolops.GetMolFrags(linker)) == 1 and len(linker_bonds) == 2:
+                    Chem.Kekulize(linker, clearAromaticFlags=True)
+                    # If for some reason a starting fragment isn't removed (and it's larger than the linker), remove (happens v. occassionally)
+                    if len(Chem.rdmolops.GetMolFrags(linker)) > 1:
+                        for frag in Chem.MolToSmiles(linker).split('.'):
+                            if Chem.MolFromSmiles(frag).GetNumHeavyAtoms() == linker_len:
+                                return frag
+                    return Chem.MolToSmiles(Chem.MolFromSmiles(Chem.MolToSmiles(linker)))
+                # If not, add to possible linkers (above doesn't capture some complex cases)
+                else:
+                    poss_linker.append(Chem.MolToSmiles(linker))
+    # If only one possibility, return linker
+    if len(poss_linker) == 1:
+        return poss_linker[0]
+    # If no possibilities, process failed
+    elif len(poss_linker) == 0:
+        print("FAIL:", Chem.MolToSmiles(full_mol), Chem.MolToSmiles(clean_frag))
+        return ""
+    # If multiple possibilities, process probably failed
+    else:
+        print("More than one poss linker. ", poss_linker)
+        return poss_linker[0]
+def unique(results):
+    total_dupes = 0
+    total = 0
+    for res in results:
+        original_num = len(res)
+        test_data = set(res)
+        new_num = len(test_data)
+        total_dupes += original_num - new_num
+        total += original_num
+    return 1 - total_dupes/float(total)
+def check_recovered_original_mol_with_idx(results):
+    outcomes = []
+    rec_idx = []
+    for res in results:
+        success = False
+        # Load original mol and canonicalise
+        orig_mol = Chem.MolFromSmiles(res[0][0][0])
+        Chem.RemoveStereochemistry(orig_mol)
+        orig_mol = Chem.MolToSmiles(Chem.RemoveHs(orig_mol))
+        #orig_mol = MolStandardize.canonicalize_tautomer_smiles(orig_mol)
+        # Check generated mols
+        for m in res:
+            # print(1)
+            gen_mol = Chem.MolFromSmiles(m[0][2])
+            Chem.RemoveStereochemistry(gen_mol)
+            gen_mol = Chem.MolToSmiles(Chem.RemoveHs(gen_mol))
+            #gen_mol = MolStandardize.canonicalize_tautomer_smiles(gen_mol)
+            if gen_mol == orig_mol:
+                # outcomes.append(True)
+                success = True
+                rec_idx.append(m[1])
+                # break
+        if not success:
+            outcomes.append(False)
+        else:
+            outcomes.append(True)
+    return outcomes, rec_idx
+def topology_from_rdkit(rdkit_molecule):
+    topology = nx.Graph()
+    for atom in rdkit_molecule.GetAtoms():
+        # Add the atoms as nodes
+        topology.add_node(atom.GetIdx(), atom_type=atom.GetAtomicNum())
+        # Add the bonds as edges
+    for bond in rdkit_molecule.GetBonds():
+        topology.add_edge(bond.GetBeginAtomIdx(), bond.GetEndAtomIdx(), bond_type=bond.GetBondType())
+    return topology

src/delinker_utils/sascorer.py ADDED Viewed

	@@ -0,0 +1,173 @@

+#
+# calculation of synthetic accessibility score as described in:
+#
+# Estimation of Synthetic Accessibility Score of Drug-like Molecules based on Molecular Complexity and Fragment Contributions
+# Peter Ertl and Ansgar Schuffenhauer
+# Journal of Cheminformatics 1:8 (2009)
+# http://www.jcheminf.com/content/1/1/8
+#
+# several small modifications to the original paper are included
+# particularly slightly different formula for marocyclic penalty
+# and taking into account also molecule symmetry (fingerprint density)
+#
+# for a set of 10k diverse molecules the agreement between the original method
+# as implemented in PipelinePilot and this implementation is r2 = 0.97
+#
+# peter ertl & greg landrum, september 2013
+#
+from __future__ import print_function
+from rdkit import Chem
+from rdkit.Chem import rdMolDescriptors
+from rdkit.six.moves import cPickle
+from rdkit.six import iteritems
+import math
+from collections import defaultdict
+import os.path as op
+_fscores = None
+def readFragmentScores(name='models/fpscores'):
+  import gzip
+  global _fscores
+  # generate the full path filename:
+  if name == "fpscores":
+    name = op.join(op.dirname(__file__), name)
+  _fscores = cPickle.load(gzip.open('%s.pkl.gz' % name))
+  outDict = {}
+  for i in _fscores:
+    for j in range(1, len(i)):
+      outDict[i[j]] = float(i[0])
+  _fscores = outDict
+def numBridgeheadsAndSpiro(mol, ri=None):
+  nSpiro = rdMolDescriptors.CalcNumSpiroAtoms(mol)
+  nBridgehead = rdMolDescriptors.CalcNumBridgeheadAtoms(mol)
+  return nBridgehead, nSpiro
+def calculateScore(m):
+  if _fscores is None:
+    readFragmentScores()
+  # fragment score
+  fp = rdMolDescriptors.GetMorganFingerprint(m,
+                                             2)  #<- 2 is the *radius* of the circular fingerprint
+  fps = fp.GetNonzeroElements()
+  score1 = 0.
+  nf = 0
+  for bitId, v in iteritems(fps):
+    nf += v
+    sfp = bitId
+    score1 += _fscores.get(sfp, -4) * v
+  score1 /= nf
+  # features score
+  nAtoms = m.GetNumAtoms()
+  nChiralCenters = len(Chem.FindMolChiralCenters(m, includeUnassigned=True))
+  ri = m.GetRingInfo()
+  nBridgeheads, nSpiro = numBridgeheadsAndSpiro(m, ri)
+  nMacrocycles = 0
+  for x in ri.AtomRings():
+    if len(x) > 8:
+      nMacrocycles += 1
+  sizePenalty = nAtoms**1.005 - nAtoms
+  stereoPenalty = math.log10(nChiralCenters + 1)
+  spiroPenalty = math.log10(nSpiro + 1)
+  bridgePenalty = math.log10(nBridgeheads + 1)
+  macrocyclePenalty = 0.
+  # ---------------------------------------
+  # This differs from the paper, which defines:
+  #  macrocyclePenalty = math.log10(nMacrocycles+1)
+  # This form generates better results when 2 or more macrocycles are present
+  if nMacrocycles > 0:
+    macrocyclePenalty = math.log10(2)
+  score2 = 0. - sizePenalty - stereoPenalty - spiroPenalty - bridgePenalty - macrocyclePenalty
+  # correction for the fingerprint density
+  # not in the original publication, added in version 1.1
+  # to make highly symmetrical molecules easier to synthetise
+  score3 = 0.
+  if nAtoms > len(fps):
+    score3 = math.log(float(nAtoms) / len(fps)) * .5
+  sascore = score1 + score2 + score3
+  # need to transform "raw" value into scale between 1 and 10
+  min = -4.0
+  max = 2.5
+  sascore = 11. - (sascore - min + 1) / (max - min) * 9.
+  # smooth the 10-end
+  if sascore > 8.:
+    sascore = 8. + math.log(sascore + 1. - 9.)
+  if sascore > 10.:
+    sascore = 10.0
+  elif sascore < 1.:
+    sascore = 1.0
+  return sascore
+def processMols(mols):
+  print('smiles\tName\tsa_score')
+  for i, m in enumerate(mols):
+    if m is None:
+      continue
+    s = calculateScore(m)
+    smiles = Chem.MolToSmiles(m)
+    print(smiles + "\t" + m.GetProp('_Name') + "\t%3f" % s)
+if __name__ == '__main__':
+  import sys, time
+  t1 = time.time()
+  readFragmentScores("fpscores")
+  t2 = time.time()
+  suppl = Chem.SmilesMolSupplier(sys.argv[1])
+  t3 = time.time()
+  processMols(suppl)
+  t4 = time.time()
+  print('Reading took %.2f seconds. Calculating took %.2f seconds' % ((t2 - t1), (t4 - t3)),
+        file=sys.stderr)
+#
+#  Copyright (c) 2013, Novartis Institutes for BioMedical Research Inc.
+#  All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met:
+#
+#     * Redistributions of source code must retain the above copyright
+#       notice, this list of conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above
+#       copyright notice, this list of conditions and the following
+#       disclaimer in the documentation and/or other materials provided
+#       with the distribution.
+#     * Neither the name of Novartis Institutes for BioMedical Research Inc.
+#       nor the names of its contributors may be used to endorse or promote
+#       products derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#

src/edm.py ADDED Viewed

	@@ -0,0 +1,730 @@

+import torch
+import torch.nn.functional as F
+import numpy as np
+import math
+from src import utils
+from src.egnn import Dynamics
+from src.noise import GammaNetwork, PredefinedNoiseSchedule
+from typing import Union
+from pdb import set_trace
+class EDM(torch.nn.Module):
+    def __init__(
+            self,
+            dynamics: Union[Dynamics],
+            in_node_nf: int,
+            n_dims: int,
+            timesteps: int = 1000,
+            noise_schedule='learned',
+            noise_precision=1e-4,
+            loss_type='vlb',
+            norm_values=(1., 1., 1.),
+            norm_biases=(None, 0., 0.),
+    ):
+        super().__init__()
+        if noise_schedule == 'learned':
+            assert loss_type == 'vlb', 'A noise schedule can only be learned with a vlb objective'
+            self.gamma = GammaNetwork()
+        else:
+            self.gamma = PredefinedNoiseSchedule(noise_schedule, timesteps=timesteps, precision=noise_precision)
+        self.dynamics = dynamics
+        self.in_node_nf = in_node_nf
+        self.n_dims = n_dims
+        self.T = timesteps
+        self.norm_values = norm_values
+        self.norm_biases = norm_biases
+    def forward(self, x, h, node_mask, fragment_mask, linker_mask, edge_mask, context=None):
+        # Normalization and concatenation
+        x, h = self.normalize(x, h)
+        xh = torch.cat([x, h], dim=2)
+        # Volume change loss term
+        delta_log_px = self.delta_log_px(linker_mask).mean()
+        # Sample t
+        t_int = torch.randint(0, self.T + 1, size=(x.size(0), 1), device=x.device).float()
+        s_int = t_int - 1
+        t = t_int / self.T
+        s = s_int / self.T
+        # Masks for t=0 and t>0
+        t_is_zero = (t_int == 0).squeeze().float()
+        t_is_not_zero = 1 - t_is_zero
+        # Compute gamma_t and gamma_s according to the noise schedule
+        gamma_t = self.inflate_batch_array(self.gamma(t), x)
+        gamma_s = self.inflate_batch_array(self.gamma(s), x)
+        # Compute alpha_t and sigma_t from gamma
+        alpha_t = self.alpha(gamma_t, x)
+        sigma_t = self.sigma(gamma_t, x)
+        # Sample noise
+        # Note: only for linker
+        eps_t = self.sample_combined_position_feature_noise(n_samples=x.size(0), n_nodes=x.size(1), mask=linker_mask)
+        # Sample z_t given x, h for timestep t, from q(z_t | x, h)
+        # Note: keep fragments unchanged
+        z_t = alpha_t * xh + sigma_t * eps_t
+        z_t = xh * fragment_mask + z_t * linker_mask
+        # Neural net prediction
+        eps_t_hat = self.dynamics.forward(
+            xh=z_t,
+            t=t,
+            node_mask=node_mask,
+            linker_mask=linker_mask,
+            context=context,
+            edge_mask=edge_mask,
+        )
+        eps_t_hat = eps_t_hat * linker_mask
+        # Computing basic error (further used for computing NLL and L2-loss)
+        error_t = self.sum_except_batch((eps_t - eps_t_hat) ** 2)
+        # Computing L2-loss for t>0
+        normalization = (self.n_dims + self.in_node_nf) * self.numbers_of_nodes(linker_mask)
+        l2_loss = error_t / normalization
+        l2_loss = l2_loss.mean()
+        # The KL between q(z_T | x) and p(z_T) = Normal(0, 1) (should be close to zero)
+        kl_prior = self.kl_prior(xh, linker_mask).mean()
+        # Computing NLL middle term
+        SNR_weight = (self.SNR(gamma_s - gamma_t) - 1).squeeze(1).squeeze(1)
+        loss_term_t = self.T * 0.5 * SNR_weight * error_t
+        loss_term_t = (loss_term_t * t_is_not_zero).sum() / t_is_not_zero.sum()
+        # Computing noise returned by dynamics
+        noise = torch.norm(eps_t_hat, dim=[1, 2])
+        noise_t = (noise * t_is_not_zero).sum() / t_is_not_zero.sum()
+        if t_is_zero.sum() > 0:
+            # The _constants_ depending on sigma_0 from the
+            # cross entropy term E_q(z0 | x) [log p(x | z0)]
+            neg_log_constants = -self.log_constant_of_p_x_given_z0(x, linker_mask)
+            # Computes the L_0 term (even if gamma_t is not actually gamma_0)
+            # and selected only relevant via masking
+            loss_term_0 = -self.log_p_xh_given_z0_without_constants(h, z_t, gamma_t, eps_t, eps_t_hat, linker_mask)
+            loss_term_0 = loss_term_0 + neg_log_constants
+            loss_term_0 = (loss_term_0 * t_is_zero).sum() / t_is_zero.sum()
+            # Computing noise returned by dynamics
+            noise_0 = (noise * t_is_zero).sum() / t_is_zero.sum()
+        else:
+            loss_term_0 = 0.
+            noise_0 = 0.
+        return delta_log_px, kl_prior, loss_term_t, loss_term_0, l2_loss, noise_t, noise_0
+    @torch.no_grad()
+    def sample_chain(self, x, h, node_mask, fragment_mask, linker_mask, edge_mask, context, keep_frames=None):
+        n_samples = x.size(0)
+        n_nodes = x.size(1)
+        # Normalization and concatenation
+        x, h, = self.normalize(x, h)
+        xh = torch.cat([x, h], dim=2)
+        # Initial linker sampling from N(0, I)
+        z = self.sample_combined_position_feature_noise(n_samples, n_nodes, mask=linker_mask)
+        z = xh * fragment_mask + z * linker_mask
+        if keep_frames is None:
+            keep_frames = self.T
+        else:
+            assert keep_frames <= self.T
+        chain = torch.zeros((keep_frames,) + z.size(), device=z.device)
+        # Sample p(z_s | z_t)
+        for s in reversed(range(0, self.T)):
+            s_array = torch.full((n_samples, 1), fill_value=s, device=z.device)
+            t_array = s_array + 1
+            s_array = s_array / self.T
+            t_array = t_array / self.T
+            z = self.sample_p_zs_given_zt_only_linker(
+                s=s_array,
+                t=t_array,
+                z_t=z,
+                node_mask=node_mask,
+                fragment_mask=fragment_mask,
+                linker_mask=linker_mask,
+                edge_mask=edge_mask,
+                context=context,
+            )
+            write_index = (s * keep_frames) // self.T
+            chain[write_index] = self.unnormalize_z(z)
+        # Finally sample p(x, h | z_0)
+        x, h = self.sample_p_xh_given_z0_only_linker(
+            z_0=z,
+            node_mask=node_mask,
+            fragment_mask=fragment_mask,
+            linker_mask=linker_mask,
+            edge_mask=edge_mask,
+            context=context,
+        )
+        chain[0] = torch.cat([x, h], dim=2)
+        return chain
+    def sample_p_zs_given_zt_only_linker(self, s, t, z_t, node_mask, fragment_mask, linker_mask, edge_mask, context):
+        """Samples from zs ~ p(zs | zt). Only used during sampling. Samples only linker features and coords"""
+        gamma_s = self.gamma(s)
+        gamma_t = self.gamma(t)
+        sigma2_t_given_s, sigma_t_given_s, alpha_t_given_s = self.sigma_and_alpha_t_given_s(gamma_t, gamma_s, z_t)
+        sigma_s = self.sigma(gamma_s, target_tensor=z_t)
+        sigma_t = self.sigma(gamma_t, target_tensor=z_t)
+        # Neural net prediction.
+        eps_hat = self.dynamics.forward(
+            xh=z_t,
+            t=t,
+            node_mask=node_mask,
+            linker_mask=linker_mask,
+            context=context,
+            edge_mask=edge_mask,
+        )
+        eps_hat = eps_hat * linker_mask
+        # Compute mu for p(z_s | z_t)
+        mu = z_t / alpha_t_given_s - (sigma2_t_given_s / alpha_t_given_s / sigma_t) * eps_hat
+        # Compute sigma for p(z_s | z_t)
+        sigma = sigma_t_given_s * sigma_s / sigma_t
+        # Sample z_s given the parameters derived from zt
+        z_s = self.sample_normal(mu, sigma, linker_mask)
+        z_s = z_t * fragment_mask + z_s * linker_mask
+        return z_s
+    def sample_p_xh_given_z0_only_linker(self, z_0, node_mask, fragment_mask, linker_mask, edge_mask, context):
+        """Samples x ~ p(x|z0). Samples only linker features and coords"""
+        zeros = torch.zeros(size=(z_0.size(0), 1), device=z_0.device)
+        gamma_0 = self.gamma(zeros)
+        # Computes sqrt(sigma_0^2 / alpha_0^2)
+        sigma_x = self.SNR(-0.5 * gamma_0).unsqueeze(1)
+        eps_hat = self.dynamics.forward(
+            t=zeros,
+            xh=z_0,
+            node_mask=node_mask,
+            linker_mask=linker_mask,
+            edge_mask=edge_mask,
+            context=context
+        )
+        eps_hat = eps_hat * linker_mask
+        mu_x = self.compute_x_pred(eps_t=eps_hat, z_t=z_0, gamma_t=gamma_0)
+        xh = self.sample_normal(mu=mu_x, sigma=sigma_x, node_mask=linker_mask)
+        xh = z_0 * fragment_mask + xh * linker_mask
+        x, h = xh[:, :, :self.n_dims], xh[:, :, self.n_dims:]
+        x, h = self.unnormalize(x, h)
+        h = F.one_hot(torch.argmax(h, dim=2), self.in_node_nf) * node_mask
+        return x, h
+    def compute_x_pred(self, eps_t, z_t, gamma_t):
+        """Computes x_pred, i.e. the most likely prediction of x."""
+        sigma_t = self.sigma(gamma_t, target_tensor=eps_t)
+        alpha_t = self.alpha(gamma_t, target_tensor=eps_t)
+        x_pred = 1. / alpha_t * (z_t - sigma_t * eps_t)
+        return x_pred
+    def kl_prior(self, xh, mask):
+        """
+        Computes the KL between q(z1 | x) and the prior p(z1) = Normal(0, 1).
+        This is essentially a lot of work for something that is in practice negligible in the loss.
+        However, you compute it so that you see it when you've made a mistake in your noise schedule.
+        """
+        # Compute the last alpha value, alpha_T
+        ones = torch.ones((xh.size(0), 1), device=xh.device)
+        gamma_T = self.gamma(ones)
+        alpha_T = self.alpha(gamma_T, xh)
+        # Compute means
+        mu_T = alpha_T * xh
+        mu_T_x, mu_T_h = mu_T[:, :, :self.n_dims], mu_T[:, :, self.n_dims:]
+        # Compute standard deviations (only batch axis for x-part, inflated for h-part)
+        sigma_T_x = self.sigma(gamma_T, mu_T_x).view(-1)  # Remove inflate, only keep batch dimension for x-part
+        sigma_T_h = self.sigma(gamma_T, mu_T_h)
+        # Compute KL for h-part
+        zeros, ones = torch.zeros_like(mu_T_h), torch.ones_like(sigma_T_h)
+        kl_distance_h = self.gaussian_kl(mu_T_h, sigma_T_h, zeros, ones)
+        # Compute KL for x-part
+        zeros, ones = torch.zeros_like(mu_T_x), torch.ones_like(sigma_T_x)
+        d = self.dimensionality(mask)
+        kl_distance_x = self.gaussian_kl_for_dimension(mu_T_x, sigma_T_x, zeros, ones, d=d)
+        return kl_distance_x + kl_distance_h
+    def log_constant_of_p_x_given_z0(self, x, mask):
+        batch_size = x.size(0)
+        degrees_of_freedom_x = self.dimensionality(mask)
+        zeros = torch.zeros((batch_size, 1), device=x.device)
+        gamma_0 = self.gamma(zeros)
+        # Recall that sigma_x = sqrt(sigma_0^2 / alpha_0^2) = SNR(-0.5 gamma_0)
+        log_sigma_x = 0.5 * gamma_0.view(batch_size)
+        return degrees_of_freedom_x * (- log_sigma_x - 0.5 * np.log(2 * np.pi))
+    def log_p_xh_given_z0_without_constants(self, h, z_0, gamma_0, eps, eps_hat, mask, epsilon=1e-10):
+        # Discrete properties are predicted directly from z_0
+        z_h = z_0[:, :, self.n_dims:]
+        # Take only part over x
+        eps_x = eps[:, :, :self.n_dims]
+        eps_hat_x = eps_hat[:, :, :self.n_dims]
+        # Compute sigma_0 and rescale to the integer scale of the data
+        sigma_0 = self.sigma(gamma_0, target_tensor=z_0) * self.norm_values[1]
+        # Computes the error for the distribution N(x | 1 / alpha_0 z_0 + sigma_0/alpha_0 eps_0, sigma_0 / alpha_0),
+        # the weighting in the epsilon parametrization is exactly '1'
+        log_p_x_given_z_without_constants = -0.5 * self.sum_except_batch((eps_x - eps_hat_x) ** 2)
+        # Categorical features
+        # Compute delta indicator masks
+        h = h * self.norm_values[1] + self.norm_biases[1]
+        estimated_h = z_h * self.norm_values[1] + self.norm_biases[1]
+        # Centered h_cat around 1, since onehot encoded
+        centered_h = estimated_h - 1
+        # Compute integrals from 0.5 to 1.5 of the normal distribution
+        # N(mean=centered_h_cat, stdev=sigma_0_cat)
+        log_p_h_proportional = torch.log(
+            self.cdf_standard_gaussian((centered_h + 0.5) / sigma_0) -
+            self.cdf_standard_gaussian((centered_h - 0.5) / sigma_0) +
+            epsilon
+        )
+        # Normalize the distribution over the categories
+        log_Z = torch.logsumexp(log_p_h_proportional, dim=2, keepdim=True)
+        log_probabilities = log_p_h_proportional - log_Z
+        # Select the log_prob of the current category using the onehot representation
+        log_p_h_given_z = self.sum_except_batch(log_probabilities * h * mask)
+        # Combine log probabilities for x and h
+        log_p_xh_given_z = log_p_x_given_z_without_constants + log_p_h_given_z
+        return log_p_xh_given_z
+    def sample_combined_position_feature_noise(self, n_samples, n_nodes, mask):
+        z_x = utils.sample_gaussian_with_mask(
+            size=(n_samples, n_nodes, self.n_dims),
+            device=mask.device,
+            node_mask=mask
+        )
+        z_h = utils.sample_gaussian_with_mask(
+            size=(n_samples, n_nodes, self.in_node_nf),
+            device=mask.device,
+            node_mask=mask
+        )
+        z = torch.cat([z_x, z_h], dim=2)
+        return z
+    def sample_normal(self, mu, sigma, node_mask):
+        """Samples from a Normal distribution."""
+        eps = self.sample_combined_position_feature_noise(mu.size(0), mu.size(1), node_mask)
+        return mu + sigma * eps
+    def normalize(self, x, h):
+        new_x = x / self.norm_values[0]
+        new_h = (h.float() - self.norm_biases[1]) / self.norm_values[1]
+        return new_x, new_h
+    def unnormalize(self, x, h):
+        new_x = x * self.norm_values[0]
+        new_h = h * self.norm_values[1] + self.norm_biases[1]
+        return new_x, new_h
+    def unnormalize_z(self, z):
+        assert z.size(2) == self.n_dims + self.in_node_nf
+        x, h = z[:, :, :self.n_dims], z[:, :, self.n_dims:]
+        x, h = self.unnormalize(x, h)
+        return torch.cat([x, h], dim=2)
+    def delta_log_px(self, mask):
+        return -self.dimensionality(mask) * np.log(self.norm_values[0])
+    def dimensionality(self, mask):
+        return self.numbers_of_nodes(mask) * self.n_dims
+    def sigma(self, gamma, target_tensor):
+        """Computes sigma given gamma."""
+        return self.inflate_batch_array(torch.sqrt(torch.sigmoid(gamma)), target_tensor)
+    def alpha(self, gamma, target_tensor):
+        """Computes alpha given gamma."""
+        return self.inflate_batch_array(torch.sqrt(torch.sigmoid(-gamma)), target_tensor)
+    def SNR(self, gamma):
+        """Computes signal to noise ratio (alpha^2/sigma^2) given gamma."""
+        return torch.exp(-gamma)
+    def sigma_and_alpha_t_given_s(self, gamma_t: torch.Tensor, gamma_s: torch.Tensor, target_tensor: torch.Tensor):
+        """
+        Computes sigma t given s, using gamma_t and gamma_s. Used during sampling.
+        These are defined as:
+            alpha t given s = alpha t / alpha s,
+            sigma t given s = sqrt(1 - (alpha t given s) ^2 ).
+        """
+        sigma2_t_given_s = self.inflate_batch_array(
+            -self.expm1(self.softplus(gamma_s) - self.softplus(gamma_t)),
+            target_tensor
+        )
+        # alpha_t_given_s = alpha_t / alpha_s
+        log_alpha2_t = F.logsigmoid(-gamma_t)
+        log_alpha2_s = F.logsigmoid(-gamma_s)
+        log_alpha2_t_given_s = log_alpha2_t - log_alpha2_s
+        alpha_t_given_s = torch.exp(0.5 * log_alpha2_t_given_s)
+        alpha_t_given_s = self.inflate_batch_array(alpha_t_given_s, target_tensor)
+        sigma_t_given_s = torch.sqrt(sigma2_t_given_s)
+        return sigma2_t_given_s, sigma_t_given_s, alpha_t_given_s
+    @staticmethod
+    def numbers_of_nodes(mask):
+        return torch.sum(mask.squeeze(2), dim=1)
+    @staticmethod
+    def inflate_batch_array(array, target):
+        """
+        Inflates the batch array (array) with only a single axis (i.e. shape = (batch_size,),
+        or possibly more empty axes (i.e. shape (batch_size, 1, ..., 1)) to match the target shape.
+        """
+        target_shape = (array.size(0),) + (1,) * (len(target.size()) - 1)
+        return array.view(target_shape)
+    @staticmethod
+    def sum_except_batch(x):
+        return x.view(x.size(0), -1).sum(-1)
+    @staticmethod
+    def expm1(x: torch.Tensor) -> torch.Tensor:
+        return torch.expm1(x)
+    @staticmethod
+    def softplus(x: torch.Tensor) -> torch.Tensor:
+        return F.softplus(x)
+    @staticmethod
+    def cdf_standard_gaussian(x):
+        return 0.5 * (1. + torch.erf(x / math.sqrt(2)))
+    @staticmethod
+    def gaussian_kl(q_mu, q_sigma, p_mu, p_sigma):
+        """
+        Computes the KL distance between two normal distributions.
+        Args:
+            q_mu: Mean of distribution q.
+            q_sigma: Standard deviation of distribution q.
+            p_mu: Mean of distribution p.
+            p_sigma: Standard deviation of distribution p.
+        Returns:
+            The KL distance, summed over all dimensions except the batch dim.
+        """
+        kl = torch.log(p_sigma / q_sigma) + 0.5 * (q_sigma ** 2 + (q_mu - p_mu) ** 2) / (p_sigma ** 2) - 0.5
+        return EDM.sum_except_batch(kl)
+    @staticmethod
+    def gaussian_kl_for_dimension(q_mu, q_sigma, p_mu, p_sigma, d):
+        """
+        Computes the KL distance between two normal distributions taking the dimension into account.
+        Args:
+            q_mu: Mean of distribution q.
+            q_sigma: Standard deviation of distribution q.
+            p_mu: Mean of distribution p.
+            p_sigma: Standard deviation of distribution p.
+            d: dimension
+        Returns:
+            The KL distance, summed over all dimensions except the batch dim.
+        """
+        mu_norm_2 = EDM.sum_except_batch((q_mu - p_mu) ** 2)
+        return d * torch.log(p_sigma / q_sigma) + 0.5 * (d * q_sigma ** 2 + mu_norm_2) / (p_sigma ** 2) - 0.5 * d
+class InpaintingEDM(EDM):
+    def forward(self, x, h, node_mask, fragment_mask, linker_mask, edge_mask, context=None):
+        # Normalization and concatenation
+        x, h = self.normalize(x, h)
+        xh = torch.cat([x, h], dim=2)
+        # Volume change loss term
+        delta_log_px = self.delta_log_px(node_mask).mean()
+        # Sample t
+        t_int = torch.randint(0, self.T + 1, size=(x.size(0), 1), device=x.device).float()
+        s_int = t_int - 1
+        t = t_int / self.T
+        s = s_int / self.T
+        # Masks for t=0 and t>0
+        t_is_zero = (t_int == 0).squeeze().float()
+        t_is_not_zero = 1 - t_is_zero
+        # Compute gamma_t and gamma_s according to the noise schedule
+        gamma_t = self.inflate_batch_array(self.gamma(t), x)
+        gamma_s = self.inflate_batch_array(self.gamma(s), x)
+        # Compute alpha_t and sigma_t from gamma
+        alpha_t = self.alpha(gamma_t, x)
+        sigma_t = self.sigma(gamma_t, x)
+        # Sample noise
+        eps_t = self.sample_combined_position_feature_noise(n_samples=x.size(0), n_nodes=x.size(1), mask=node_mask)
+        # Sample z_t given x, h for timestep t, from q(z_t | x, h)
+        # Note: keep fragments unchanged
+        z_t = alpha_t * xh + sigma_t * eps_t
+        # Neural net prediction
+        eps_t_hat = self.dynamics.forward(
+            xh=z_t,
+            t=t,
+            node_mask=node_mask,
+            linker_mask=None,
+            context=context,
+            edge_mask=edge_mask,
+        )
+        # Computing basic error (further used for computing NLL and L2-loss)
+        error_t = self.sum_except_batch((eps_t - eps_t_hat) ** 2)
+        # Computing L2-loss for t>0
+        normalization = (self.n_dims + self.in_node_nf) * self.numbers_of_nodes(node_mask)
+        l2_loss = error_t / normalization
+        l2_loss = l2_loss.mean()
+        # The KL between q(z_T | x) and p(z_T) = Normal(0, 1) (should be close to zero)
+        kl_prior = self.kl_prior(xh, node_mask).mean()
+        # Computing NLL middle term
+        SNR_weight = (self.SNR(gamma_s - gamma_t) - 1).squeeze(1).squeeze(1)
+        loss_term_t = self.T * 0.5 * SNR_weight * error_t
+        loss_term_t = (loss_term_t * t_is_not_zero).sum() / t_is_not_zero.sum()
+        # Computing noise returned by dynamics
+        noise = torch.norm(eps_t_hat, dim=[1, 2])
+        noise_t = (noise * t_is_not_zero).sum() / t_is_not_zero.sum()
+        if t_is_zero.sum() > 0:
+            # The _constants_ depending on sigma_0 from the
+            # cross entropy term E_q(z0 | x) [log p(x | z0)]
+            neg_log_constants = -self.log_constant_of_p_x_given_z0(x, node_mask)
+            # Computes the L_0 term (even if gamma_t is not actually gamma_0)
+            # and selected only relevant via masking
+            loss_term_0 = -self.log_p_xh_given_z0_without_constants(h, z_t, gamma_t, eps_t, eps_t_hat, node_mask)
+            loss_term_0 = loss_term_0 + neg_log_constants
+            loss_term_0 = (loss_term_0 * t_is_zero).sum() / t_is_zero.sum()
+            # Computing noise returned by dynamics
+            noise_0 = (noise * t_is_zero).sum() / t_is_zero.sum()
+        else:
+            loss_term_0 = 0.
+            noise_0 = 0.
+        return delta_log_px, kl_prior, loss_term_t, loss_term_0, l2_loss, noise_t, noise_0
+    @torch.no_grad()
+    def sample_chain(self, x, h, node_mask, edge_mask, fragment_mask, linker_mask, context, keep_frames=None):
+        n_samples = x.size(0)
+        n_nodes = x.size(1)
+        # Normalization and concatenation
+        x, h, = self.normalize(x, h)
+        xh = torch.cat([x, h], dim=2)
+        # Sampling initial noise
+        z = self.sample_combined_position_feature_noise(n_samples, n_nodes, node_mask)
+        if keep_frames is None:
+            keep_frames = self.T
+        else:
+            assert keep_frames <= self.T
+        chain = torch.zeros((keep_frames,) + z.size(), device=z.device)
+        # Sample p(z_s | z_t)
+        for s in reversed(range(0, self.T)):
+            s_array = torch.full((n_samples, 1), fill_value=s, device=z.device)
+            t_array = s_array + 1
+            s_array = s_array / self.T
+            t_array = t_array / self.T
+            z_linker_only_sampled = self.sample_p_zs_given_zt(
+                s=s_array,
+                t=t_array,
+                z_t=z,
+                node_mask=node_mask,
+                edge_mask=edge_mask,
+                context=context,
+            )
+            z_fragments_only_sampled = self.sample_q_zs_given_zt_and_x(
+                s=s_array,
+                t=t_array,
+                z_t=z,
+                x=xh * fragment_mask,
+                node_mask=fragment_mask,
+            )
+            z = z_linker_only_sampled * linker_mask + z_fragments_only_sampled * fragment_mask
+            # Project down to avoid numerical runaway of the center of gravity
+            z_x = utils.remove_mean_with_mask(z[:, :, :self.n_dims], node_mask)
+            z_h = z[:, :, self.n_dims:]
+            z = torch.cat([z_x, z_h], dim=2)
+            # Saving step to the chain
+            write_index = (s * keep_frames) // self.T
+            chain[write_index] = self.unnormalize_z(z)
+        # Finally sample p(x, h | z_0)
+        x_out_linker, h_out_linker = self.sample_p_xh_given_z0(
+            z_0=z,
+            node_mask=node_mask,
+            edge_mask=edge_mask,
+            context=context,
+        )
+        x_out_fragments, h_out_fragments = self.sample_q_xh_given_z0_and_x(z_0=z, node_mask=node_mask)
+        xh_out_linker = torch.cat([x_out_linker, h_out_linker], dim=2)
+        xh_out_fragments = torch.cat([x_out_fragments, h_out_fragments], dim=2)
+        xh_out = xh_out_linker * linker_mask + xh_out_fragments * fragment_mask
+        # Overwrite last frame with the resulting x and h
+        chain[0] = xh_out
+        return chain
+    def sample_p_zs_given_zt(self, s, t, z_t, node_mask, edge_mask, context):
+        """Samples from zs ~ p(zs | zt). Only used during sampling"""
+        gamma_s = self.gamma(s)
+        gamma_t = self.gamma(t)
+        sigma2_t_given_s, sigma_t_given_s, alpha_t_given_s = self.sigma_and_alpha_t_given_s(gamma_t, gamma_s, z_t)
+        sigma_s = self.sigma(gamma_s, target_tensor=z_t)
+        sigma_t = self.sigma(gamma_t, target_tensor=z_t)
+        # Neural net prediction.
+        eps_hat = self.dynamics.forward(
+            xh=z_t,
+            t=t,
+            node_mask=node_mask,
+            linker_mask=None,
+            edge_mask=edge_mask,
+            context=context
+        )
+        # Checking that epsilon is centered around linker COM
+        utils.assert_mean_zero_with_mask(eps_hat[:, :, :self.n_dims], node_mask)
+        # Compute mu for p(z_s | z_t)
+        mu = z_t / alpha_t_given_s - (sigma2_t_given_s / alpha_t_given_s / sigma_t) * eps_hat
+        # Compute sigma for p(z_s | z_t)
+        sigma = sigma_t_given_s * sigma_s / sigma_t
+        # Sample z_s given the parameters derived from z_t
+        z_s = self.sample_normal(mu, sigma, node_mask)
+        return z_s
+    def sample_q_zs_given_zt_and_x(self, s, t, z_t, x, node_mask):
+        """Samples from zs ~ q(zs | zt, x). Only used during sampling. Samples only linker features and coords"""
+        gamma_s = self.gamma(s)
+        gamma_t = self.gamma(t)
+        sigma2_t_given_s, sigma_t_given_s, alpha_t_given_s = self.sigma_and_alpha_t_given_s(gamma_t, gamma_s, z_t)
+        sigma_s = self.sigma(gamma_s, target_tensor=z_t)
+        sigma_t = self.sigma(gamma_t, target_tensor=z_t)
+        alpha_s = self.alpha(gamma_s, x)
+        mu = (
+            alpha_t_given_s * (sigma_s ** 2) / (sigma_t ** 2) * z_t +
+            alpha_s * sigma2_t_given_s / (sigma_t ** 2) * x
+        )
+        # Compute sigma for p(zs | zt)
+        sigma = sigma_t_given_s * sigma_s / sigma_t
+        # Sample zs given the parameters derived from zt
+        z_s = self.sample_normal(mu, sigma, node_mask)
+        return z_s
+    def sample_p_xh_given_z0(self, z_0, node_mask, edge_mask, context):
+        """Samples x ~ p(x|z0). Samples only linker features and coords"""
+        zeros = torch.zeros(size=(z_0.size(0), 1), device=z_0.device)
+        gamma_0 = self.gamma(zeros)
+        # Computes sqrt(sigma_0^2 / alpha_0^2)
+        sigma_x = self.SNR(-0.5 * gamma_0).unsqueeze(1)
+        eps_hat = self.dynamics.forward(
+            xh=z_0,
+            t=zeros,
+            node_mask=node_mask,
+            linker_mask=None,
+            edge_mask=edge_mask,
+            context=context
+        )
+        utils.assert_mean_zero_with_mask(eps_hat[:, :, :self.n_dims], node_mask)
+        mu_x = self.compute_x_pred(eps_hat, z_0, gamma_0)
+        xh = self.sample_normal(mu=mu_x, sigma=sigma_x, node_mask=node_mask)
+        x, h = xh[:, :, :self.n_dims], xh[:, :, self.n_dims:]
+        x, h = self.unnormalize(x, h)
+        h = F.one_hot(torch.argmax(h, dim=2), self.in_node_nf) * node_mask
+        return x, h
+    def sample_q_xh_given_z0_and_x(self, z_0, node_mask):
+        """Samples x ~ q(x|z0). Samples only linker features and coords"""
+        zeros = torch.zeros(size=(z_0.size(0), 1), device=z_0.device)
+        gamma_0 = self.gamma(zeros)
+        alpha_0 = self.alpha(gamma_0, z_0)
+        sigma_0 = self.sigma(gamma_0, z_0)
+        eps = self.sample_combined_position_feature_noise(z_0.size(0), z_0.size(1), node_mask)
+        xh = (1 / alpha_0) * z_0 - (sigma_0 / alpha_0) * eps
+        x, h = xh[:, :, :self.n_dims], xh[:, :, self.n_dims:]
+        x, h = self.unnormalize(x, h)
+        h = F.one_hot(torch.argmax(h, dim=2), self.in_node_nf) * node_mask
+        return x, h
+    def sample_combined_position_feature_noise(self, n_samples, n_nodes, mask):
+        z_x = utils.sample_center_gravity_zero_gaussian_with_mask(
+            size=(n_samples, n_nodes, self.n_dims),
+            device=mask.device,
+            node_mask=mask
+        )
+        z_h = utils.sample_gaussian_with_mask(
+            size=(n_samples, n_nodes, self.in_node_nf),
+            device=mask.device,
+            node_mask=mask
+        )
+        z = torch.cat([z_x, z_h], dim=2)
+        return z
+    def dimensionality(self, mask):
+        return (self.numbers_of_nodes(mask) - 1) * self.n_dims

src/egnn.py ADDED Viewed

	@@ -0,0 +1,541 @@

+import math
+import numpy as np
+import torch
+import torch.nn as nn
+from src import utils
+from pdb import set_trace
+class GCL(nn.Module):
+    def __init__(self, input_nf, output_nf, hidden_nf, normalization_factor, aggregation_method, activation,
+                 edges_in_d=0, nodes_att_dim=0, attention=False, normalization=None):
+        super(GCL, self).__init__()
+        input_edge = input_nf * 2
+        self.normalization_factor = normalization_factor
+        self.aggregation_method = aggregation_method
+        self.attention = attention
+        self.edge_mlp = nn.Sequential(
+            nn.Linear(input_edge + edges_in_d, hidden_nf),
+            activation,
+            nn.Linear(hidden_nf, hidden_nf),
+            activation)
+        if normalization is None:
+            self.node_mlp = nn.Sequential(
+                nn.Linear(hidden_nf + input_nf + nodes_att_dim, hidden_nf),
+                activation,
+                nn.Linear(hidden_nf, output_nf)
+            )
+        elif normalization == 'batch_norm':
+            self.node_mlp = nn.Sequential(
+                nn.Linear(hidden_nf + input_nf + nodes_att_dim, hidden_nf),
+                nn.BatchNorm1d(hidden_nf),
+                activation,
+                nn.Linear(hidden_nf, output_nf),
+                nn.BatchNorm1d(output_nf),
+            )
+        else:
+            raise NotImplementedError
+        if self.attention:
+            self.att_mlp = nn.Sequential(nn.Linear(hidden_nf, 1), nn.Sigmoid())
+    def edge_model(self, source, target, edge_attr, edge_mask):
+        if edge_attr is None:  # Unused.
+            out = torch.cat([source, target], dim=1)
+        else:
+            out = torch.cat([source, target, edge_attr], dim=1)
+        mij = self.edge_mlp(out)
+        if self.attention:
+            att_val = self.att_mlp(mij)
+            out = mij * att_val
+        else:
+            out = mij
+        if edge_mask is not None:
+            out = out * edge_mask
+        return out, mij
+    def node_model(self, x, edge_index, edge_attr, node_attr):
+        row, col = edge_index
+        agg = unsorted_segment_sum(edge_attr, row, num_segments=x.size(0),
+                                   normalization_factor=self.normalization_factor,
+                                   aggregation_method=self.aggregation_method)
+        if node_attr is not None:
+            agg = torch.cat([x, agg, node_attr], dim=1)
+        else:
+            agg = torch.cat([x, agg], dim=1)
+        out = x + self.node_mlp(agg)
+        return out, agg
+    def forward(self, h, edge_index, edge_attr=None, node_attr=None, node_mask=None, edge_mask=None):
+        row, col = edge_index
+        edge_feat, mij = self.edge_model(h[row], h[col], edge_attr, edge_mask)
+        h, agg = self.node_model(h, edge_index, edge_feat, node_attr)
+        if node_mask is not None:
+            h = h * node_mask
+        return h, mij
+class EquivariantUpdate(nn.Module):
+    def __init__(self, hidden_nf, normalization_factor, aggregation_method,
+                 edges_in_d=1, activation=nn.SiLU(), tanh=False, coords_range=10.0):
+        super(EquivariantUpdate, self).__init__()
+        self.tanh = tanh
+        self.coords_range = coords_range
+        input_edge = hidden_nf * 2 + edges_in_d
+        layer = nn.Linear(hidden_nf, 1, bias=False)
+        torch.nn.init.xavier_uniform_(layer.weight, gain=0.001)
+        self.coord_mlp = nn.Sequential(
+            nn.Linear(input_edge, hidden_nf),
+            activation,
+            nn.Linear(hidden_nf, hidden_nf),
+            activation,
+            layer)
+        self.normalization_factor = normalization_factor
+        self.aggregation_method = aggregation_method
+    def coord_model(self, h, coord, edge_index, coord_diff, edge_attr, edge_mask, linker_mask):
+        row, col = edge_index
+        input_tensor = torch.cat([h[row], h[col], edge_attr], dim=1)
+        if self.tanh:
+            trans = coord_diff * torch.tanh(self.coord_mlp(input_tensor)) * self.coords_range
+        else:
+            trans = coord_diff * self.coord_mlp(input_tensor)
+        if edge_mask is not None:
+            trans = trans * edge_mask
+        agg = unsorted_segment_sum(trans, row, num_segments=coord.size(0),
+                                   normalization_factor=self.normalization_factor,
+                                   aggregation_method=self.aggregation_method)
+        if linker_mask is not None:
+            agg = agg * linker_mask
+        coord = coord + agg
+        return coord
+    def forward(
+            self, h, coord, edge_index, coord_diff, edge_attr=None, linker_mask=None, node_mask=None, edge_mask=None
+    ):
+        coord = self.coord_model(h, coord, edge_index, coord_diff, edge_attr, edge_mask, linker_mask)
+        if node_mask is not None:
+            coord = coord * node_mask
+        return coord
+class EquivariantBlock(nn.Module):
+    def __init__(self, hidden_nf, edge_feat_nf=2, device='cpu', activation=nn.SiLU(), n_layers=2, attention=True,
+                 norm_diff=True, tanh=False, coords_range=15, norm_constant=1, sin_embedding=None,
+                 normalization_factor=100, aggregation_method='sum'):
+        super(EquivariantBlock, self).__init__()
+        self.hidden_nf = hidden_nf
+        self.device = device
+        self.n_layers = n_layers
+        self.coords_range_layer = float(coords_range)
+        self.norm_diff = norm_diff
+        self.norm_constant = norm_constant
+        self.sin_embedding = sin_embedding
+        self.normalization_factor = normalization_factor
+        self.aggregation_method = aggregation_method
+        for i in range(0, n_layers):
+            self.add_module("gcl_%d" % i, GCL(self.hidden_nf, self.hidden_nf, self.hidden_nf, edges_in_d=edge_feat_nf,
+                                              activation=activation, attention=attention,
+                                              normalization_factor=self.normalization_factor,
+                                              aggregation_method=self.aggregation_method))
+        self.add_module("gcl_equiv", EquivariantUpdate(hidden_nf, edges_in_d=edge_feat_nf, activation=activation, tanh=tanh,
+                                                       coords_range=self.coords_range_layer,
+                                                       normalization_factor=self.normalization_factor,
+                                                       aggregation_method=self.aggregation_method))
+        self.to(self.device)
+    def forward(self, h, x, edge_index, node_mask=None, linker_mask=None, edge_mask=None, edge_attr=None):
+        # Edit Emiel: Remove velocity as input
+        distances, coord_diff = coord2diff(x, edge_index, self.norm_constant)
+        if self.sin_embedding is not None:
+            distances = self.sin_embedding(distances)
+        edge_attr = torch.cat([distances, edge_attr], dim=1)
+        for i in range(0, self.n_layers):
+            h, _ = self._modules["gcl_%d" % i](h, edge_index, edge_attr=edge_attr, node_mask=node_mask, edge_mask=edge_mask)
+        x = self._modules["gcl_equiv"](
+            h, x,
+            edge_index=edge_index,
+            coord_diff=coord_diff,
+            edge_attr=edge_attr,
+            linker_mask=linker_mask,
+            node_mask=node_mask,
+            edge_mask=edge_mask,
+        )
+        # Important, the bias of the last linear might be non-zero
+        if node_mask is not None:
+            h = h * node_mask
+        return h, x
+class EGNN(nn.Module):
+    def __init__(self, in_node_nf, in_edge_nf, hidden_nf, device='cpu', activation=nn.SiLU(), n_layers=3, attention=False,
+                 norm_diff=True, out_node_nf=None, tanh=False, coords_range=15, norm_constant=1, inv_sublayers=2,
+                 sin_embedding=False, normalization_factor=100, aggregation_method='sum'):
+        super(EGNN, self).__init__()
+        if out_node_nf is None:
+            out_node_nf = in_node_nf
+        self.hidden_nf = hidden_nf
+        self.device = device
+        self.n_layers = n_layers
+        self.coords_range_layer = float(coords_range/n_layers)
+        self.norm_diff = norm_diff
+        self.normalization_factor = normalization_factor
+        self.aggregation_method = aggregation_method
+        if sin_embedding:
+            self.sin_embedding = SinusoidsEmbeddingNew()
+            edge_feat_nf = self.sin_embedding.dim * 2
+        else:
+            self.sin_embedding = None
+            edge_feat_nf = 2
+        self.embedding = nn.Linear(in_node_nf, self.hidden_nf)
+        self.embedding_out = nn.Linear(self.hidden_nf, out_node_nf)
+        for i in range(0, n_layers):
+            self.add_module("e_block_%d" % i, EquivariantBlock(hidden_nf, edge_feat_nf=edge_feat_nf, device=device,
+                                                               activation=activation, n_layers=inv_sublayers,
+                                                               attention=attention, norm_diff=norm_diff, tanh=tanh,
+                                                               coords_range=coords_range, norm_constant=norm_constant,
+                                                               sin_embedding=self.sin_embedding,
+                                                               normalization_factor=self.normalization_factor,
+                                                               aggregation_method=self.aggregation_method))
+        self.to(self.device)
+    def forward(self, h, x, edge_index, node_mask=None, linker_mask=None, edge_mask=None):
+        # Edit Emiel: Remove velocity as input
+        distances, _ = coord2diff(x, edge_index)
+        if self.sin_embedding is not None:
+            distances = self.sin_embedding(distances)
+        h = self.embedding(h)
+        for i in range(0, self.n_layers):
+            h, x = self._modules["e_block_%d" % i](
+                h, x, edge_index,
+                node_mask=node_mask,
+                linker_mask=linker_mask,
+                edge_mask=edge_mask,
+                edge_attr=distances
+            )
+        # Important, the bias of the last linear might be non-zero
+        h = self.embedding_out(h)
+        if node_mask is not None:
+            h = h * node_mask
+        return h, x
+class GNN(nn.Module):
+    def __init__(self, in_node_nf, in_edge_nf, hidden_nf, aggregation_method='sum', device='cpu',
+                 activation=nn.SiLU(), n_layers=4, attention=False, normalization_factor=1,
+                 out_node_nf=None, normalization=None):
+        super(GNN, self).__init__()
+        if out_node_nf is None:
+            out_node_nf = in_node_nf
+        self.hidden_nf = hidden_nf
+        self.device = device
+        self.n_layers = n_layers
+        # Encoder
+        self.embedding = nn.Linear(in_node_nf, self.hidden_nf)
+        self.embedding_out = nn.Linear(self.hidden_nf, out_node_nf)
+        for i in range(0, n_layers):
+            self.add_module("gcl_%d" % i, GCL(
+                self.hidden_nf, self.hidden_nf, self.hidden_nf,
+                normalization_factor=normalization_factor,
+                aggregation_method=aggregation_method,
+                edges_in_d=in_edge_nf, activation=activation,
+                attention=attention, normalization=normalization))
+        self.to(self.device)
+    def forward(self, h, edges, edge_attr=None, node_mask=None, edge_mask=None):
+        # Edit Emiel: Remove velocity as input
+        h = self.embedding(h)
+        for i in range(0, self.n_layers):
+            h, _ = self._modules["gcl_%d" % i](h, edges, edge_attr=edge_attr, node_mask=node_mask, edge_mask=edge_mask)
+        h = self.embedding_out(h)
+        # Important, the bias of the last linear might be non-zero
+        if node_mask is not None:
+            h = h * node_mask
+        return h
+class SinusoidsEmbeddingNew(nn.Module):
+    def __init__(self, max_res=15., min_res=15. / 2000., div_factor=4):
+        super().__init__()
+        self.n_frequencies = int(math.log(max_res / min_res, div_factor)) + 1
+        self.frequencies = 2 * math.pi * div_factor ** torch.arange(self.n_frequencies)/max_res
+        self.dim = len(self.frequencies) * 2
+    def forward(self, x):
+        x = torch.sqrt(x + 1e-8)
+        emb = x * self.frequencies[None, :].to(x.device)
+        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
+        return emb.detach()
+def coord2diff(x, edge_index, norm_constant=1):
+    row, col = edge_index
+    coord_diff = x[row] - x[col]
+    radial = torch.sum((coord_diff) ** 2, 1).unsqueeze(1)
+    norm = torch.sqrt(radial + 1e-8)
+    coord_diff = coord_diff/(norm + norm_constant)
+    return radial, coord_diff
+def unsorted_segment_sum(data, segment_ids, num_segments, normalization_factor, aggregation_method: str):
+    """Custom PyTorch op to replicate TensorFlow's `unsorted_segment_sum`.
+        Normalization: 'sum' or 'mean'.
+    """
+    result_shape = (num_segments, data.size(1))
+    result = data.new_full(result_shape, 0)  # Init empty result tensor.
+    segment_ids = segment_ids.unsqueeze(-1).expand(-1, data.size(1))
+    result.scatter_add_(0, segment_ids, data)
+    if aggregation_method == 'sum':
+        result = result / normalization_factor
+    if aggregation_method == 'mean':
+        norm = data.new_zeros(result.shape)
+        norm.scatter_add_(0, segment_ids, data.new_ones(data.shape))
+        norm[norm == 0] = 1
+        result = result / norm
+    return result
+class Dynamics(nn.Module):
+    def __init__(
+            self, n_dims, in_node_nf, context_node_nf, hidden_nf=64, device='cpu', activation=nn.SiLU(),
+            n_layers=4, attention=False, condition_time=True, tanh=False, norm_constant=0, inv_sublayers=2,
+            sin_embedding=False, normalization_factor=100, aggregation_method='sum', model='egnn_dynamics',
+            normalization=None, centering=False,
+    ):
+        super().__init__()
+        self.device = device
+        self.n_dims = n_dims
+        self.context_node_nf = context_node_nf
+        self.condition_time = condition_time
+        self.model = model
+        self.centering = centering
+        in_node_nf = in_node_nf + context_node_nf + condition_time
+        if self.model == 'egnn_dynamics':
+            self.dynamics = EGNN(
+                in_node_nf=in_node_nf,
+                in_edge_nf=1,
+                hidden_nf=hidden_nf, device=device,
+                activation=activation,
+                n_layers=n_layers,
+                attention=attention,
+                tanh=tanh,
+                norm_constant=norm_constant,
+                inv_sublayers=inv_sublayers,
+                sin_embedding=sin_embedding,
+                normalization_factor=normalization_factor,
+                aggregation_method=aggregation_method,
+            )
+        elif self.model == 'gnn_dynamics':
+            self.dynamics = GNN(
+                in_node_nf=in_node_nf+3,
+                in_edge_nf=0,
+                hidden_nf=hidden_nf,
+                out_node_nf=in_node_nf+3,
+                device=device,
+                activation=activation,
+                n_layers=n_layers,
+                attention=attention,
+                normalization_factor=normalization_factor,
+                aggregation_method=aggregation_method,
+                normalization=normalization,
+            )
+        else:
+            raise NotImplementedError
+        self.edge_cache = {}
+    def forward(self, t, xh, node_mask, linker_mask, edge_mask, context):
+        """
+        - t: (B)
+        - xh: (B, N, D), where D = 3 + nf
+        - node_mask: (B, N, 1)
+        - edge_mask: (B*N*N, 1)
+        - context: (B, N, C)
+        """
+        bs, n_nodes = xh.shape[0], xh.shape[1]
+        edges = self.get_edges(n_nodes, bs)  # (2, B*N)
+        node_mask = node_mask.view(bs * n_nodes, 1)  # (B*N, 1)
+        if linker_mask is not None:
+            linker_mask = linker_mask.view(bs * n_nodes, 1)  # (B*N, 1)
+        # Reshaping node features & adding time feature
+        xh = xh.view(bs * n_nodes, -1).clone() * node_mask  # (B*N, D)
+        x = xh[:, :self.n_dims].clone()  # (B*N, 3)
+        h = xh[:, self.n_dims:].clone()  # (B*N, nf)
+        if self.condition_time:
+            if np.prod(t.size()) == 1:
+                # t is the same for all elements in batch.
+                h_time = torch.empty_like(h[:, 0:1]).fill_(t.item())
+            else:
+                # t is different over the batch dimension.
+                h_time = t.view(bs, 1).repeat(1, n_nodes)
+                h_time = h_time.view(bs * n_nodes, 1)
+            h = torch.cat([h, h_time], dim=1)  # (B*N, nf+1)
+        if context is not None:
+            context = context.view(bs*n_nodes, self.context_node_nf)
+            h = torch.cat([h, context], dim=1)
+        # Forward EGNN
+        # Output: h_final (B*N, nf), x_final (B*N, 3), vel (B*N, 3)
+        if self.model == 'egnn_dynamics':
+            h_final, x_final = self.dynamics(
+                h,
+                x,
+                edges,
+                node_mask=node_mask,
+                linker_mask=linker_mask,
+                edge_mask=edge_mask
+            )
+            vel = (x_final - x) * node_mask  # This masking operation is redundant but just in case
+        elif self.model == 'gnn_dynamics':
+            xh = torch.cat([x, h], dim=1)
+            output = self.dynamics(xh, edges, node_mask=node_mask)
+            vel = output[:, 0:3] * node_mask
+            h_final = output[:, 3:]
+        else:
+            raise NotImplementedError
+        # Slice off context size
+        if context is not None:
+            h_final = h_final[:, :-self.context_node_nf]
+        # Slice off last dimension which represented time.
+        if self.condition_time:
+            h_final = h_final[:, :-1]
+        vel = vel.view(bs, n_nodes, -1)  # (B, N, 3)
+        h_final = h_final.view(bs, n_nodes, -1)  # (B, N, D)
+        node_mask = node_mask.view(bs, n_nodes, 1)  # (B, N, 1)
+        if torch.any(torch.isnan(vel)) or torch.any(torch.isnan(h_final)):
+            raise utils.FoundNaNException(vel, h_final)
+        if self.centering:
+            vel = utils.remove_mean_with_mask(vel, node_mask)
+        return torch.cat([vel, h_final], dim=2)
+    def get_edges(self, n_nodes, batch_size):
+        if n_nodes in self.edge_cache:
+            edges_dic_b = self.edge_cache[n_nodes]
+            if batch_size in edges_dic_b:
+                return edges_dic_b[batch_size]
+            else:
+                # get edges for a single sample
+                rows, cols = [], []
+                for batch_idx in range(batch_size):
+                    for i in range(n_nodes):
+                        for j in range(n_nodes):
+                            rows.append(i + batch_idx * n_nodes)
+                            cols.append(j + batch_idx * n_nodes)
+                edges = [torch.LongTensor(rows).to(self.device), torch.LongTensor(cols).to(self.device)]
+                edges_dic_b[batch_size] = edges
+                return edges
+        else:
+            self.edge_cache[n_nodes] = {}
+            return self.get_edges(n_nodes, batch_size)
+class DynamicsWithPockets(Dynamics):
+    def forward(self, t, xh, node_mask, linker_mask, edge_mask, context):
+        """
+        - t: (B)
+        - xh: (B, N, D), where D = 3 + nf
+        - node_mask: (B, N, 1)
+        - edge_mask: (B*N*N, 1)
+        - context: (B, N, C)
+        """
+        bs, n_nodes = xh.shape[0], xh.shape[1]
+        node_mask = node_mask.view(bs * n_nodes, 1)  # (B*N, 1)
+        if linker_mask is not None:
+            linker_mask = linker_mask.view(bs * n_nodes, 1)  # (B*N, 1)
+        # Reshaping node features & adding time feature
+        xh = xh.view(bs * n_nodes, -1).clone() * node_mask  # (B*N, D)
+        x = xh[:, :self.n_dims].clone()  # (B*N, 3)
+        h = xh[:, self.n_dims:].clone()  # (B*N, nf)
+        edges = self.get_dist_edges(x, node_mask, edge_mask)
+        if self.condition_time:
+            if np.prod(t.size()) == 1:
+                # t is the same for all elements in batch.
+                h_time = torch.empty_like(h[:, 0:1]).fill_(t.item())
+            else:
+                # t is different over the batch dimension.
+                h_time = t.view(bs, 1).repeat(1, n_nodes)
+                h_time = h_time.view(bs * n_nodes, 1)
+            h = torch.cat([h, h_time], dim=1)  # (B*N, nf+1)
+        if context is not None:
+            context = context.view(bs*n_nodes, self.context_node_nf)
+            h = torch.cat([h, context], dim=1)
+        # Forward EGNN
+        # Output: h_final (B*N, nf), x_final (B*N, 3), vel (B*N, 3)
+        if self.model == 'egnn_dynamics':
+            h_final, x_final = self.dynamics(
+                h,
+                x,
+                edges,
+                node_mask=node_mask,
+                linker_mask=linker_mask,
+                edge_mask=None
+            )
+            vel = (x_final - x) * node_mask  # This masking operation is redundant but just in case
+        elif self.model == 'gnn_dynamics':
+            xh = torch.cat([x, h], dim=1)
+            output = self.dynamics(xh, edges, node_mask=node_mask)
+            vel = output[:, 0:3] * node_mask
+            h_final = output[:, 3:]
+        else:
+            raise NotImplementedError
+        # Slice off context size
+        if context is not None:
+            h_final = h_final[:, :-self.context_node_nf]
+        # Slice off last dimension which represented time.
+        if self.condition_time:
+            h_final = h_final[:, :-1]
+        vel = vel.view(bs, n_nodes, -1)  # (B, N, 3)
+        h_final = h_final.view(bs, n_nodes, -1)  # (B, N, D)
+        node_mask = node_mask.view(bs, n_nodes, 1)  # (B, N, 1)
+        if torch.any(torch.isnan(vel)) or torch.any(torch.isnan(h_final)):
+            raise utils.FoundNaNException(vel, h_final)
+        if self.centering:
+            vel = utils.remove_mean_with_mask(vel, node_mask)
+        return torch.cat([vel, h_final], dim=2)
+    @staticmethod
+    def get_dist_edges(x, node_mask, batch_mask):
+        node_mask = node_mask.squeeze().bool()
+        batch_adj = (batch_mask[:, None] == batch_mask[None, :])
+        nodes_adj = (node_mask[:, None] & node_mask[None, :])
+        dists_adj = (torch.cdist(x, x) <= 4)
+        rm_self_loops = ~torch.eye(x.size(0), dtype=torch.bool, device=x.device)
+        adj = batch_adj & nodes_adj & dists_adj & rm_self_loops
+        edges = torch.stack(torch.where(adj))
+        return edges

src/lightning.py ADDED Viewed

	@@ -0,0 +1,473 @@

+import numpy as np
+import os
+import pytorch_lightning as pl
+import torch
+import wandb
+from src import metrics, utils, delinker
+from src.const import LINKER_SIZE_DIST
+from src.egnn import Dynamics, DynamicsWithPockets
+from src.edm import EDM, InpaintingEDM
+from src.datasets import (
+    ZincDataset, MOADDataset, create_templates_for_linker_generation, get_dataloader, collate
+)
+from src.linker_size import DistributionNodes
+from src.molecule_builder import build_molecules
+from src.visualizer import save_xyz_file, visualize_chain
+from typing import Dict, List, Optional
+from tqdm import tqdm
+from pdb import set_trace
+def get_activation(activation):
+    print(activation)
+    if activation == 'silu':
+        return torch.nn.SiLU()
+    else:
+        raise Exception("activation fn not supported yet. Add it here.")
+class DDPM(pl.LightningModule):
+    train_dataset = None
+    val_dataset = None
+    test_dataset = None
+    starting_epoch = None
+    metrics: Dict[str, List[float]] = {}
+    FRAMES = 100
+    def __init__(
+        self,
+        in_node_nf, n_dims, context_node_nf, hidden_nf, activation, tanh, n_layers, attention, norm_constant,
+        inv_sublayers, sin_embedding, normalization_factor, aggregation_method,
+        diffusion_steps, diffusion_noise_schedule, diffusion_noise_precision, diffusion_loss_type,
+        normalize_factors, include_charges, model,
+        data_path, train_data_prefix, val_data_prefix, batch_size, lr, torch_device, test_epochs, n_stability_samples,
+        normalization=None, log_iterations=None, samples_dir=None, data_augmentation=False,
+        center_of_mass='fragments', inpainting=False, anchors_context=True,
+    ):
+        super(DDPM, self).__init__()
+        self.save_hyperparameters()
+        self.data_path = data_path
+        self.train_data_prefix = train_data_prefix
+        self.val_data_prefix = val_data_prefix
+        self.batch_size = batch_size
+        self.lr = lr
+        self.torch_device = torch_device
+        self.include_charges = include_charges
+        self.test_epochs = test_epochs
+        self.n_stability_samples = n_stability_samples
+        self.log_iterations = log_iterations
+        self.samples_dir = samples_dir
+        self.data_augmentation = data_augmentation
+        self.center_of_mass = center_of_mass
+        self.inpainting = inpainting
+        self.loss_type = diffusion_loss_type
+        self.n_dims = n_dims
+        self.num_classes = in_node_nf - include_charges
+        self.include_charges = include_charges
+        self.anchors_context = anchors_context
+        self.is_geom = ('geom' in self.train_data_prefix) or ('MOAD' in self.train_data_prefix)
+        if type(activation) is str:
+            activation = get_activation(activation)
+        dynamics_class = DynamicsWithPockets if '.' in train_data_prefix else Dynamics
+        dynamics = dynamics_class(
+            in_node_nf=in_node_nf,
+            n_dims=n_dims,
+            context_node_nf=context_node_nf,
+            device=torch_device,
+            hidden_nf=hidden_nf,
+            activation=activation,
+            n_layers=n_layers,
+            attention=attention,
+            tanh=tanh,
+            norm_constant=norm_constant,
+            inv_sublayers=inv_sublayers,
+            sin_embedding=sin_embedding,
+            normalization_factor=normalization_factor,
+            aggregation_method=aggregation_method,
+            model=model,
+            normalization=normalization,
+            centering=inpainting,
+        )
+        edm_class = InpaintingEDM if inpainting else EDM
+        self.edm = edm_class(
+            dynamics=dynamics,
+            in_node_nf=in_node_nf,
+            n_dims=n_dims,
+            timesteps=diffusion_steps,
+            noise_schedule=diffusion_noise_schedule,
+            noise_precision=diffusion_noise_precision,
+            loss_type=diffusion_loss_type,
+            norm_values=normalize_factors,
+        )
+        self.linker_size_sampler = DistributionNodes(LINKER_SIZE_DIST)
+    def setup(self, stage: Optional[str] = None):
+        dataset_type = MOADDataset if '.' in self.train_data_prefix else ZincDataset
+        if stage == 'fit':
+            self.is_geom = ('geom' in self.train_data_prefix) or ('MOAD' in self.train_data_prefix)
+            self.train_dataset = dataset_type(
+                data_path=self.data_path,
+                prefix=self.train_data_prefix,
+                device=self.torch_device
+            )
+            self.val_dataset = dataset_type(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        elif stage == 'val':
+            self.is_geom = ('geom' in self.val_data_prefix) or ('MOAD' in self.val_data_prefix)
+            self.val_dataset = dataset_type(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        else:
+            raise NotImplementedError
+    def train_dataloader(self, collate_fn=collate):
+        return get_dataloader(self.train_dataset, self.batch_size, collate_fn=collate_fn, shuffle=True)
+    def val_dataloader(self, collate_fn=collate):
+        return get_dataloader(self.val_dataset, self.batch_size, collate_fn=collate_fn)
+    def test_dataloader(self, collate_fn=collate):
+        return get_dataloader(self.test_dataset, self.batch_size, collate_fn=collate_fn)
+    def forward(self, data, training):
+        x = data['positions']
+        h = data['one_hot']
+        node_mask = data['atom_mask']
+        edge_mask = data['edge_mask']
+        anchors = data['anchors']
+        fragment_mask = data['fragment_mask']
+        linker_mask = data['linker_mask']
+        # Anchors and fragments labels are used as context
+        if self.anchors_context:
+            context = torch.cat([anchors, fragment_mask], dim=-1)
+        else:
+            context = fragment_mask
+        # Add information about pocket to the context
+        if '.' in self.train_data_prefix:
+            fragment_pocket_mask = fragment_mask
+            fragment_only_mask = data['fragment_only_mask']
+            pocket_only_mask = fragment_pocket_mask - fragment_only_mask
+            if self.anchors_context:
+                context = torch.cat([anchors, fragment_only_mask, pocket_only_mask], dim=-1)
+            else:
+                context = torch.cat([fragment_only_mask, pocket_only_mask], dim=-1)
+        # Removing COM of fragment from the atom coordinates
+        if self.inpainting:
+            center_of_mass_mask = node_mask
+        elif self.center_of_mass == 'fragments':
+            center_of_mass_mask = fragment_mask
+        elif self.center_of_mass == 'anchors':
+            center_of_mass_mask = anchors
+        else:
+            raise NotImplementedError(self.center_of_mass)
+        x = utils.remove_partial_mean_with_mask(x, node_mask, center_of_mass_mask)
+        utils.assert_partial_mean_zero_with_mask(x, node_mask, center_of_mass_mask)
+        # Applying random rotation
+        if training and self.data_augmentation:
+            x = utils.random_rotation(x)
+        return self.edm.forward(
+            x=x,
+            h=h,
+            node_mask=node_mask,
+            fragment_mask=fragment_mask,
+            linker_mask=linker_mask,
+            edge_mask=edge_mask,
+            context=context
+        )
+    def training_step(self, data, *args):
+        delta_log_px, kl_prior, loss_term_t, loss_term_0, l2_loss, noise_t, noise_0 = self.forward(data, training=True)
+        vlb_loss = kl_prior + loss_term_t + loss_term_0 - delta_log_px
+        if self.loss_type == 'l2':
+            loss = l2_loss
+        elif self.loss_type == 'vlb':
+            loss = vlb_loss
+        else:
+            raise NotImplementedError(self.loss_type)
+        training_metrics = {
+            'loss': loss,
+            'delta_log_px': delta_log_px,
+            'kl_prior': kl_prior,
+            'loss_term_t': loss_term_t,
+            'loss_term_0': loss_term_0,
+            'l2_loss': l2_loss,
+            'vlb_loss': vlb_loss,
+            'noise_t': noise_t,
+            'noise_0': noise_0
+        }
+        if self.log_iterations is not None and self.global_step % self.log_iterations == 0:
+            for metric_name, metric in training_metrics.items():
+                self.metrics.setdefault(f'{metric_name}/train', []).append(metric)
+                self.log(f'{metric_name}/train', metric, prog_bar=True)
+        return training_metrics
+    def validation_step(self, data, *args):
+        delta_log_px, kl_prior, loss_term_t, loss_term_0, l2_loss, noise_t, noise_0 = self.forward(data, training=False)
+        vlb_loss = kl_prior + loss_term_t + loss_term_0 - delta_log_px
+        if self.loss_type == 'l2':
+            loss = l2_loss
+        elif self.loss_type == 'vlb':
+            loss = vlb_loss
+        else:
+            raise NotImplementedError(self.loss_type)
+        return {
+            'loss': loss,
+            'delta_log_px': delta_log_px,
+            'kl_prior': kl_prior,
+            'loss_term_t': loss_term_t,
+            'loss_term_0': loss_term_0,
+            'l2_loss': l2_loss,
+            'vlb_loss': vlb_loss,
+            'noise_t': noise_t,
+            'noise_0': noise_0
+        }
+    def test_step(self, data, *args):
+        delta_log_px, kl_prior, loss_term_t, loss_term_0, l2_loss, noise_t, noise_0 = self.forward(data, training=False)
+        vlb_loss = kl_prior + loss_term_t + loss_term_0 - delta_log_px
+        if self.loss_type == 'l2':
+            loss = l2_loss
+        elif self.loss_type == 'vlb':
+            loss = vlb_loss
+        else:
+            raise NotImplementedError(self.loss_type)
+        return {
+            'loss': loss,
+            'delta_log_px': delta_log_px,
+            'kl_prior': kl_prior,
+            'loss_term_t': loss_term_t,
+            'loss_term_0': loss_term_0,
+            'l2_loss': l2_loss,
+            'vlb_loss': vlb_loss,
+            'noise_t': noise_t,
+            'noise_0': noise_0
+        }
+    def training_epoch_end(self, training_step_outputs):
+        for metric in training_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(training_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/train', []).append(avg_metric)
+            self.log(f'{metric}/train', avg_metric, prog_bar=True)
+    def validation_epoch_end(self, validation_step_outputs):
+        for metric in validation_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(validation_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/val', []).append(avg_metric)
+            self.log(f'{metric}/val', avg_metric, prog_bar=True)
+        if (self.current_epoch + 1) % self.test_epochs == 0:
+            sampling_results = self.sample_and_analyze(self.val_dataloader())
+            for metric_name, metric_value in sampling_results.items():
+                self.log(f'{metric_name}/val', metric_value, prog_bar=True)
+                self.metrics.setdefault(f'{metric_name}/val', []).append(metric_value)
+            # Logging the results corresponding to the best validation_and_connectivity
+            best_metrics, best_epoch = self.compute_best_validation_metrics()
+            self.log('best_epoch', int(best_epoch), prog_bar=True, batch_size=self.batch_size)
+            for metric, value in best_metrics.items():
+                self.log(f'best_{metric}', value, prog_bar=True, batch_size=self.batch_size)
+    def test_epoch_end(self, test_step_outputs):
+        for metric in test_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(test_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/test', []).append(avg_metric)
+            self.log(f'{metric}/test', avg_metric, prog_bar=True)
+        if (self.current_epoch + 1) % self.test_epochs == 0:
+            sampling_results = self.sample_and_analyze(self.test_dataloader())
+            for metric_name, metric_value in sampling_results.items():
+                self.log(f'{metric_name}/test', metric_value, prog_bar=True)
+                self.metrics.setdefault(f'{metric_name}/test', []).append(metric_value)
+    def generate_animation(self, chain_batch, node_mask, batch_i):
+        batch_indices, mol_indices = utils.get_batch_idx_for_animation(self.batch_size, batch_i)
+        for bi, mi in zip(batch_indices, mol_indices):
+            chain = chain_batch[:, bi, :, :]
+            name = f'mol_{mi}'
+            chain_output = os.path.join(self.samples_dir, f'epoch_{self.current_epoch}', name)
+            os.makedirs(chain_output, exist_ok=True)
+            one_hot = chain[:, :, 3:-1] if self.include_charges else chain[:, :, 3:]
+            positions = chain[:, :, :3]
+            chain_node_mask = torch.cat([node_mask[bi].unsqueeze(0) for _ in range(self.FRAMES)], dim=0)
+            names = [f'{name}_{j}' for j in range(self.FRAMES)]
+            save_xyz_file(chain_output, one_hot, positions, chain_node_mask, names=names, is_geom=self.is_geom)
+            visualize_chain(chain_output, wandb=wandb, mode=name, is_geom=self.is_geom)
+    def sample_and_analyze(self, dataloader):
+        pred_molecules = []
+        true_molecules = []
+        true_fragments = []
+        for b, data in tqdm(enumerate(dataloader), total=len(dataloader), desc='Sampling'):
+            atom_mask = data['atom_mask']
+            fragment_mask = data['fragment_mask']
+            # Save molecules without pockets
+            if '.' in self.train_data_prefix:
+                atom_mask = data['atom_mask'] - data['pocket_mask']
+                fragment_mask = data['fragment_only_mask']
+            true_molecules_batch = build_molecules(
+                data['one_hot'],
+                data['positions'],
+                atom_mask,
+                is_geom=self.is_geom,
+            )
+            true_fragments_batch = build_molecules(
+                data['one_hot'],
+                data['positions'],
+                fragment_mask,
+                is_geom=self.is_geom,
+            )
+            for sample_idx in tqdm(range(self.n_stability_samples)):
+                try:
+                    chain_batch, node_mask = self.sample_chain(data, keep_frames=self.FRAMES)
+                except utils.FoundNaNException as e:
+                    for idx in e.x_h_nan_idx:
+                        smiles = data['name'][idx]
+                        print(f'FoundNaNException: [xh], e={self.current_epoch}, b={b}, i={idx}: {smiles}')
+                    for idx in e.only_x_nan_idx:
+                        smiles = data['name'][idx]
+                        print(f'FoundNaNException: [x ], e={self.current_epoch}, b={b}, i={idx}: {smiles}')
+                    for idx in e.only_h_nan_idx:
+                        smiles = data['name'][idx]
+                        print(f'FoundNaNException: [ h], e={self.current_epoch}, b={b}, i={idx}: {smiles}')
+                    continue
+                # Get final molecules from chains – for computing metrics
+                x, h = utils.split_features(
+                    z=chain_batch[0],
+                    n_dims=self.n_dims,
+                    num_classes=self.num_classes,
+                    include_charges=self.include_charges,
+                )
+                # Save molecules without pockets
+                if '.' in self.train_data_prefix:
+                    node_mask = node_mask - data['pocket_mask']
+                one_hot = h['categorical']
+                pred_molecules_batch = build_molecules(one_hot, x, node_mask, is_geom=self.is_geom)
+                # Adding only results for valid ground truth molecules
+                for pred_mol, true_mol, frag in zip(pred_molecules_batch, true_molecules_batch, true_fragments_batch):
+                    if metrics.is_valid(true_mol):
+                        pred_molecules.append(pred_mol)
+                        true_molecules.append(true_mol)
+                        true_fragments.append(frag)
+                # Generate animation – will always do it for molecules with idx 0, 110 and 360
+                if self.samples_dir is not None and sample_idx == 0:
+                    self.generate_animation(chain_batch=chain_batch, node_mask=node_mask, batch_i=b)
+        # Our own & DeLinker metrics
+        our_metrics = metrics.compute_metrics(
+            pred_molecules=pred_molecules,
+            true_molecules=true_molecules
+        )
+        delinker_metrics = delinker.get_delinker_metrics(
+            pred_molecules=pred_molecules,
+            true_molecules=true_molecules,
+            true_fragments=true_fragments
+        )
+        return {
+            **our_metrics,
+            **delinker_metrics
+        }
+    def sample_chain(self, data, sample_fn=None, keep_frames=None):
+        if sample_fn is None:
+            linker_sizes = data['linker_mask'].sum(1).view(-1).int()
+        else:
+            linker_sizes = sample_fn(data)
+        if self.inpainting:
+            template_data = data
+        else:
+            template_data = create_templates_for_linker_generation(data, linker_sizes)
+        x = template_data['positions']
+        node_mask = template_data['atom_mask']
+        edge_mask = template_data['edge_mask']
+        h = template_data['one_hot']
+        anchors = template_data['anchors']
+        fragment_mask = template_data['fragment_mask']
+        linker_mask = template_data['linker_mask']
+        # Anchors and fragments labels are used as context
+        if self.anchors_context:
+            context = torch.cat([anchors, fragment_mask], dim=-1)
+        else:
+            context = fragment_mask
+        # Add information about pocket to the context
+        if '.' in self.train_data_prefix:
+            fragment_pocket_mask = fragment_mask
+            fragment_only_mask = data['fragment_only_mask']
+            pocket_only_mask = fragment_pocket_mask - fragment_only_mask
+            if self.anchors_context:
+                context = torch.cat([anchors, fragment_only_mask, pocket_only_mask], dim=-1)
+            else:
+                context = torch.cat([fragment_only_mask, pocket_only_mask], dim=-1)
+        # Removing COM of fragment from the atom coordinates
+        if self.inpainting:
+            center_of_mass_mask = node_mask
+        elif self.center_of_mass == 'fragments':
+            center_of_mass_mask = fragment_mask
+        elif self.center_of_mass == 'anchors':
+            center_of_mass_mask = anchors
+        else:
+            raise NotImplementedError(self.center_of_mass)
+        x = utils.remove_partial_mean_with_mask(x, node_mask, center_of_mass_mask)
+        chain = self.edm.sample_chain(
+            x=x,
+            h=h,
+            node_mask=node_mask,
+            edge_mask=edge_mask,
+            fragment_mask=fragment_mask,
+            linker_mask=linker_mask,
+            context=context,
+            keep_frames=keep_frames,
+        )
+        return chain, node_mask
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.edm.parameters(), lr=self.lr, amsgrad=True, weight_decay=1e-12)
+    def compute_best_validation_metrics(self):
+        loss = self.metrics[f'validity_and_connectivity/val']
+        best_epoch = np.argmax(loss)
+        best_metrics = {
+            metric_name: metric_values[best_epoch]
+            for metric_name, metric_values in self.metrics.items()
+            if metric_name.endswith('/val')
+        }
+        return best_metrics, best_epoch
+    @staticmethod
+    def aggregate_metric(step_outputs, metric):
+        return torch.tensor([out[metric] for out in step_outputs]).mean()

src/linker_size.py ADDED Viewed

	@@ -0,0 +1,95 @@

+import numpy as np
+import torch
+import torch.nn as nn
+from torch.distributions.categorical import Categorical
+from src.egnn import GCL
+class DistributionNodes:
+    def __init__(self, histogram):
+        self.n_nodes = []
+        prob = []
+        self.keys = {}
+        for i, nodes in enumerate(histogram):
+            self.n_nodes.append(nodes)
+            self.keys[nodes] = i
+            prob.append(histogram[nodes])
+        self.n_nodes = torch.tensor(self.n_nodes)
+        prob = np.array(prob)
+        prob = prob/np.sum(prob)
+        self.prob = torch.from_numpy(prob).float()
+        entropy = torch.sum(self.prob * torch.log(self.prob + 1e-30))
+        print("Entropy of n_nodes: H[N]", entropy.item())
+        self.m = Categorical(torch.tensor(prob))
+    def sample(self, n_samples=1):
+        idx = self.m.sample((n_samples,))
+        return self.n_nodes[idx]
+    def log_prob(self, batch_n_nodes):
+        assert len(batch_n_nodes.size()) == 1
+        idcs = [self.keys[i.item()] for i in batch_n_nodes]
+        idcs = torch.tensor(idcs).to(batch_n_nodes.device)
+        log_p = torch.log(self.prob + 1e-30)
+        log_p = log_p.to(batch_n_nodes.device)
+        log_probs = log_p[idcs]
+        return log_probs
+class SizeGNN(nn.Module):
+    def __init__(self, in_node_nf, hidden_nf, out_node_nf, n_layers, normalization, device='cpu'):
+        super(SizeGNN, self).__init__()
+        self.hidden_nf = hidden_nf
+        self.out_node_nf = out_node_nf
+        self.device = device
+        self.embedding_in = nn.Linear(in_node_nf, self.hidden_nf)
+        self.gcl1 = GCL(
+            input_nf=self.hidden_nf,
+            output_nf=self.hidden_nf,
+            hidden_nf=self.hidden_nf,
+            normalization_factor=1,
+            aggregation_method='sum',
+            edges_in_d=1,
+            activation=nn.ReLU(),
+            attention=False,
+            normalization=normalization
+        )
+        layers = []
+        for i in range(n_layers - 1):
+            layer = GCL(
+                input_nf=self.hidden_nf,
+                output_nf=self.hidden_nf,
+                hidden_nf=self.hidden_nf,
+                normalization_factor=1,
+                aggregation_method='sum',
+                edges_in_d=1,
+                activation=nn.ReLU(),
+                attention=False,
+                normalization=normalization
+            )
+            layers.append(layer)
+        self.gcl_layers = nn.ModuleList(layers)
+        self.embedding_out = nn.Linear(self.hidden_nf, self.out_node_nf)
+        self.to(self.device)
+    def forward(self, h, edges, distances, node_mask, edge_mask):
+        h = self.embedding_in(h)
+        h, _ = self.gcl1(h, edges, edge_attr=distances, node_mask=node_mask, edge_mask=edge_mask)
+        for gcl in self.gcl_layers:
+            h, _ = gcl(h, edges, edge_attr=distances, node_mask=node_mask, edge_mask=edge_mask)
+        h = self.embedding_out(h)
+        return h

src/linker_size_lightning.py ADDED Viewed

	@@ -0,0 +1,460 @@

+import pytorch_lightning as pl
+import torch
+from src.const import ZINC_TRAIN_LINKER_ID2SIZE, ZINC_TRAIN_LINKER_SIZE2ID
+from src.linker_size import SizeGNN
+from src.egnn import coord2diff
+from src.datasets import ZincDataset, get_dataloader, collate_with_fragment_edges
+from typing import Dict, List, Optional
+from torch.nn.functional import cross_entropy, mse_loss, sigmoid
+from pdb import set_trace
+class SizeClassifier(pl.LightningModule):
+    train_dataset = None
+    val_dataset = None
+    test_dataset = None
+    metrics: Dict[str, List[float]] = {}
+    def __init__(
+            self, data_path, train_data_prefix, val_data_prefix,
+            in_node_nf, hidden_nf, out_node_nf, n_layers, batch_size, lr, torch_device,
+            normalization=None,
+            loss_weights=None,
+            min_linker_size=None,
+            linker_size2id=ZINC_TRAIN_LINKER_SIZE2ID,
+            linker_id2size=ZINC_TRAIN_LINKER_ID2SIZE,
+            task='classification',
+    ):
+        super(SizeClassifier, self).__init__()
+        self.save_hyperparameters()
+        self.data_path = data_path
+        self.train_data_prefix = train_data_prefix
+        self.val_data_prefix = val_data_prefix
+        self.min_linker_size = min_linker_size
+        self.linker_size2id = linker_size2id
+        self.linker_id2size = linker_id2size
+        self.batch_size = batch_size
+        self.lr = lr
+        self.torch_device = torch_device
+        self.loss_weights = None if loss_weights is None else torch.tensor(loss_weights, device=torch_device)
+        self.gnn = SizeGNN(
+            in_node_nf=in_node_nf,
+            hidden_nf=hidden_nf,
+            out_node_nf=out_node_nf,
+            n_layers=n_layers,
+            device=torch_device,
+            normalization=normalization,
+        )
+    def setup(self, stage: Optional[str] = None):
+        if stage == 'fit':
+            self.train_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.train_data_prefix,
+                device=self.torch_device
+            )
+            self.val_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        elif stage == 'val':
+            self.val_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        else:
+            raise NotImplementedError
+    def train_dataloader(self):
+        return get_dataloader(self.train_dataset, self.batch_size, collate_fn=collate_with_fragment_edges, shuffle=True)
+    def val_dataloader(self):
+        return get_dataloader(self.val_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
+    def test_dataloader(self):
+        return get_dataloader(self.test_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
+    def forward(self, data):
+        h = data['one_hot']
+        x = data['positions']
+        fragment_mask = data['fragment_mask']
+        linker_mask = data['linker_mask']
+        edge_mask = data['edge_mask']
+        edges = data['edges']
+        # Considering only fragments
+        x = x * fragment_mask
+        h = h * fragment_mask
+        # Reshaping
+        bs, n_nodes = x.shape[0], x.shape[1]
+        fragment_mask = fragment_mask.view(bs * n_nodes, 1)
+        x = x.view(bs * n_nodes, -1)
+        h = h.view(bs * n_nodes, -1)
+        # Prediction
+        distances, _ = coord2diff(x, edges)
+        distance_edge_mask = (edge_mask.bool() & (distances < 6)).long()
+        output = self.gnn.forward(h, edges, distances, fragment_mask, distance_edge_mask)
+        output = output.view(bs, n_nodes, -1).mean(1)
+        true = self.get_true_labels(linker_mask)
+        loss = cross_entropy(output, true, weight=self.loss_weights)
+        return output, loss
+    def get_true_labels(self, linker_mask):
+        labels = []
+        sizes = linker_mask.squeeze().sum(-1).long().detach().cpu().numpy()
+        for size in sizes:
+            label = self.linker_size2id.get(size)
+            if label is None:
+                label = self.linker_size2id[max(self.linker_id2size)]
+            labels.append(label)
+        labels = torch.tensor(labels, device=linker_mask.device, dtype=torch.long)
+        return labels
+    def training_step(self, data, *args):
+        _, loss = self.forward(data)
+        return {'loss': loss}
+    def validation_step(self, data, *args):
+        _, loss = self.forward(data)
+        return {'loss': loss}
+    def test_step(self, data, *args):
+        loss = self.forward(data)
+        return {'loss': loss}
+    def training_epoch_end(self, training_step_outputs):
+        for metric in training_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(training_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/train', []).append(avg_metric)
+            self.log(f'{metric}/train', avg_metric, prog_bar=True)
+    def validation_epoch_end(self, validation_step_outputs):
+        for metric in validation_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(validation_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/val', []).append(avg_metric)
+            self.log(f'{metric}/val', avg_metric, prog_bar=True)
+        correct = 0
+        total = 0
+        for data in self.val_dataloader():
+            output, _ = self.forward(data)
+            pred = output.argmax(dim=-1)
+            true = self.get_true_labels(data['linker_mask'])
+            correct += (pred == true).sum()
+            total += len(pred)
+        accuracy = correct / total
+        self.metrics.setdefault(f'accuracy/val', []).append(accuracy)
+        self.log(f'accuracy/val', accuracy, prog_bar=True)
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.gnn.parameters(), lr=self.lr, amsgrad=True, weight_decay=1e-12)
+    @staticmethod
+    def aggregate_metric(step_outputs, metric):
+        return torch.tensor([out[metric] for out in step_outputs]).mean()
+class SizeOrdinalClassifier(pl.LightningModule):
+    train_dataset = None
+    val_dataset = None
+    test_dataset = None
+    metrics: Dict[str, List[float]] = {}
+    def __init__(
+            self, data_path, train_data_prefix, val_data_prefix,
+            in_node_nf, hidden_nf, out_node_nf, n_layers, batch_size, lr, torch_device,
+            normalization=None,
+            min_linker_size=None,
+            linker_size2id=ZINC_TRAIN_LINKER_SIZE2ID,
+            linker_id2size=ZINC_TRAIN_LINKER_ID2SIZE,
+            task='ordinal',
+    ):
+        super(SizeOrdinalClassifier, self).__init__()
+        self.save_hyperparameters()
+        self.data_path = data_path
+        self.train_data_prefix = train_data_prefix
+        self.val_data_prefix = val_data_prefix
+        self.min_linker_size = min_linker_size
+        self.batch_size = batch_size
+        self.lr = lr
+        self.torch_device = torch_device
+        self.linker_size2id = linker_size2id
+        self.linker_id2size = linker_id2size
+        self.gnn = SizeGNN(
+            in_node_nf=in_node_nf,
+            hidden_nf=hidden_nf,
+            out_node_nf=out_node_nf,
+            n_layers=n_layers,
+            device=torch_device,
+            normalization=normalization,
+        )
+    def setup(self, stage: Optional[str] = None):
+        if stage == 'fit':
+            self.train_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.train_data_prefix,
+                device=self.torch_device
+            )
+            self.val_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        elif stage == 'val':
+            self.val_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        else:
+            raise NotImplementedError
+    def train_dataloader(self):
+        return get_dataloader(self.train_dataset, self.batch_size, collate_fn=collate_with_fragment_edges, shuffle=True)
+    def val_dataloader(self):
+        return get_dataloader(self.val_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
+    def test_dataloader(self):
+        return get_dataloader(self.test_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
+    def forward(self, data):
+        h = data['one_hot']
+        x = data['positions']
+        fragment_mask = data['fragment_mask']
+        linker_mask = data['linker_mask']
+        edge_mask = data['edge_mask']
+        edges = data['edges']
+        # Considering only fragments
+        x = x * fragment_mask
+        h = h * fragment_mask
+        # Reshaping
+        bs, n_nodes = x.shape[0], x.shape[1]
+        fragment_mask = fragment_mask.view(bs * n_nodes, 1)
+        x = x.view(bs * n_nodes, -1)
+        h = h.view(bs * n_nodes, -1)
+        # Prediction
+        distances, _ = coord2diff(x, edges)
+        distance_edge_mask = (edge_mask.bool() & (distances < 6)).long()
+        output = self.gnn.forward(h, edges, distances, fragment_mask, distance_edge_mask)
+        output = output.view(bs, n_nodes, -1).mean(1)
+        output = sigmoid(output)
+        true = self.get_true_labels(linker_mask)
+        loss = self.ordinal_loss(output, true)
+        return output, loss
+    def ordinal_loss(self, pred, true):
+        target = torch.zeros_like(pred, device=self.torch_device)
+        for i, label in enumerate(true):
+            target[i, 0:label + 1] = 1
+        return mse_loss(pred, target, reduction='none').sum(1).mean()
+    def get_true_labels(self, linker_mask):
+        labels = []
+        sizes = linker_mask.squeeze().sum(-1).long().detach().cpu().numpy()
+        for size in sizes:
+            label = self.linker_size2id.get(size)
+            if label is None:
+                label = self.linker_size2id[max(self.linker_id2size)]
+            labels.append(label)
+        labels = torch.tensor(labels, device=linker_mask.device, dtype=torch.long)
+        return labels
+    @staticmethod
+    def prediction2label(pred):
+        return torch.cumprod(pred > 0.5, dim=1).sum(dim=1) - 1
+    def training_step(self, data, *args):
+        _, loss = self.forward(data)
+        return {'loss': loss}
+    def validation_step(self, data, *args):
+        _, loss = self.forward(data)
+        return {'loss': loss}
+    def test_step(self, data, *args):
+        loss = self.forward(data)
+        return {'loss': loss}
+    def training_epoch_end(self, training_step_outputs):
+        for metric in training_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(training_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/train', []).append(avg_metric)
+            self.log(f'{metric}/train', avg_metric, prog_bar=True)
+    def validation_epoch_end(self, validation_step_outputs):
+        for metric in validation_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(validation_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/val', []).append(avg_metric)
+            self.log(f'{metric}/val', avg_metric, prog_bar=True)
+        correct = 0
+        total = 0
+        for data in self.val_dataloader():
+            output, _ = self.forward(data)
+            pred = self.prediction2label(output)
+            true = self.get_true_labels(data['linker_mask'])
+            correct += (pred == true).sum()
+            total += len(pred)
+        accuracy = correct / total
+        self.metrics.setdefault(f'accuracy/val', []).append(accuracy)
+        self.log(f'accuracy/val', accuracy, prog_bar=True)
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.gnn.parameters(), lr=self.lr, amsgrad=True, weight_decay=1e-12)
+    @staticmethod
+    def aggregate_metric(step_outputs, metric):
+        return torch.tensor([out[metric] for out in step_outputs]).mean()
+class SizeRegressor(pl.LightningModule):
+    train_dataset = None
+    val_dataset = None
+    test_dataset = None
+    metrics: Dict[str, List[float]] = {}
+    def __init__(
+            self, data_path, train_data_prefix, val_data_prefix,
+            in_node_nf, hidden_nf, n_layers, batch_size, lr, torch_device,
+            normalization=None, task='regression',
+    ):
+        super(SizeRegressor, self).__init__()
+        self.save_hyperparameters()
+        self.data_path = data_path
+        self.train_data_prefix = train_data_prefix
+        self.val_data_prefix = val_data_prefix
+        self.batch_size = batch_size
+        self.lr = lr
+        self.torch_device = torch_device
+        self.gnn = SizeGNN(
+            in_node_nf=in_node_nf,
+            hidden_nf=hidden_nf,
+            out_node_nf=1,
+            n_layers=n_layers,
+            device=torch_device,
+            normalization=normalization,
+        )
+    def setup(self, stage: Optional[str] = None):
+        if stage == 'fit':
+            self.train_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.train_data_prefix,
+                device=self.torch_device
+            )
+            self.val_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        elif stage == 'val':
+            self.val_dataset = ZincDataset(
+                data_path=self.data_path,
+                prefix=self.val_data_prefix,
+                device=self.torch_device
+            )
+        else:
+            raise NotImplementedError
+    def train_dataloader(self):
+        return get_dataloader(self.train_dataset, self.batch_size, collate_fn=collate_with_fragment_edges, shuffle=True)
+    def val_dataloader(self):
+        return get_dataloader(self.val_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
+    def test_dataloader(self):
+        return get_dataloader(self.test_dataset, self.batch_size, collate_fn=collate_with_fragment_edges)
+    def forward(self, data):
+        h = data['one_hot']
+        x = data['positions']
+        fragment_mask = data['fragment_mask']
+        linker_mask = data['linker_mask']
+        edge_mask = data['edge_mask']
+        edges = data['edges']
+        # Considering only fragments
+        x = x * fragment_mask
+        h = h * fragment_mask
+        # Reshaping
+        bs, n_nodes = x.shape[0], x.shape[1]
+        fragment_mask = fragment_mask.view(bs * n_nodes, 1)
+        x = x.view(bs * n_nodes, -1)
+        h = h.view(bs * n_nodes, -1)
+        # Prediction
+        distances, _ = coord2diff(x, edges)
+        distance_edge_mask = (edge_mask.bool() & (distances < 6)).long()
+        output = self.gnn.forward(h, edges, distances, fragment_mask, distance_edge_mask)
+        output = output.view(bs, n_nodes, -1).mean(1).squeeze()
+        true = linker_mask.squeeze().sum(-1).float()
+        loss = mse_loss(output, true)
+        return output, loss
+    def training_step(self, data, *args):
+        _, loss = self.forward(data)
+        return {'loss': loss}
+    def validation_step(self, data, *args):
+        _, loss = self.forward(data)
+        return {'loss': loss}
+    def test_step(self, data, *args):
+        loss = self.forward(data)
+        return {'loss': loss}
+    def training_epoch_end(self, training_step_outputs):
+        for metric in training_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(training_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/train', []).append(avg_metric)
+            self.log(f'{metric}/train', avg_metric, prog_bar=True)
+    def validation_epoch_end(self, validation_step_outputs):
+        for metric in validation_step_outputs[0].keys():
+            avg_metric = self.aggregate_metric(validation_step_outputs, metric)
+            self.metrics.setdefault(f'{metric}/val', []).append(avg_metric)
+            self.log(f'{metric}/val', avg_metric, prog_bar=True)
+        correct = 0
+        total = 0
+        for data in self.val_dataloader():
+            output, _ = self.forward(data)
+            pred = torch.round(output).long()
+            true = data['linker_mask'].squeeze().sum(-1).long()
+            correct += (pred == true).sum()
+            total += len(pred)
+        accuracy = correct / total
+        self.metrics.setdefault(f'accuracy/val', []).append(accuracy)
+        self.log(f'accuracy/val', accuracy, prog_bar=True)
+    def configure_optimizers(self):
+        return torch.optim.AdamW(self.gnn.parameters(), lr=self.lr, amsgrad=True, weight_decay=1e-12)
+    @staticmethod
+    def aggregate_metric(step_outputs, metric):
+        return torch.tensor([out[metric] for out in step_outputs]).mean()

src/metrics.py ADDED Viewed

	@@ -0,0 +1,167 @@

+import numpy as np
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from src import const
+from src.molecule_builder import get_bond_order
+from scipy.stats import wasserstein_distance
+from pdb import set_trace
+def is_valid(mol):
+    try:
+        Chem.SanitizeMol(mol)
+    except ValueError:
+        return False
+    return True
+def is_connected(mol):
+    try:
+        mol_frags = Chem.GetMolFrags(mol, asMols=True)
+    except Chem.rdchem.AtomValenceException:
+        return False
+    if len(mol_frags) != 1:
+        return False
+    return True
+def get_valid_molecules(molecules):
+    valid = []
+    for mol in molecules:
+        if is_valid(mol):
+            valid.append(mol)
+    return valid
+def get_connected_molecules(molecules):
+    connected = []
+    for mol in molecules:
+        if is_connected(mol):
+            connected.append(mol)
+    return connected
+def get_unique_smiles(valid_molecules):
+    unique = set()
+    for mol in valid_molecules:
+        unique.add(Chem.MolToSmiles(mol))
+    return list(unique)
+def get_novel_smiles(unique_true_smiles, unique_pred_smiles):
+    return list(set(unique_pred_smiles).difference(set(unique_true_smiles)))
+def compute_energy(mol):
+    mp = AllChem.MMFFGetMoleculeProperties(mol)
+    energy = AllChem.MMFFGetMoleculeForceField(mol, mp, confId=0).CalcEnergy()
+    return energy
+def wasserstein_distance_between_energies(true_molecules, pred_molecules):
+    true_energy_dist = []
+    for mol in true_molecules:
+        try:
+            energy = compute_energy(mol)
+            true_energy_dist.append(energy)
+        except:
+            continue
+    pred_energy_dist = []
+    for mol in pred_molecules:
+        try:
+            energy = compute_energy(mol)
+            pred_energy_dist.append(energy)
+        except:
+            continue
+    if len(true_energy_dist) > 0 and len(pred_energy_dist) > 0:
+        return wasserstein_distance(true_energy_dist, pred_energy_dist)
+    else:
+        return 0
+def compute_metrics(pred_molecules, true_molecules):
+    if len(pred_molecules) == 0:
+        return {
+            'validity': 0,
+            'validity_and_connectivity': 0,
+            'validity_as_in_delinker': 0,
+            'uniqueness': 0,
+            'novelty': 0,
+            'energies': 0,
+        }
+    # Passing rdkit.Chem.Sanitize filter
+    true_valid = get_valid_molecules(true_molecules)
+    pred_valid = get_valid_molecules(pred_molecules)
+    validity = len(pred_valid) / len(pred_molecules)
+    # Checking if molecule consists of a single connected part
+    true_valid_and_connected = get_connected_molecules(true_valid)
+    pred_valid_and_connected = get_connected_molecules(pred_valid)
+    validity_and_connectivity = len(pred_valid_and_connected) / len(pred_molecules)
+    # Unique molecules
+    true_unique = get_unique_smiles(true_valid_and_connected)
+    pred_unique = get_unique_smiles(pred_valid_and_connected)
+    uniqueness = len(pred_unique) / len(pred_valid_and_connected) if len(pred_valid_and_connected) > 0 else 0
+    # Novel molecules
+    pred_novel = get_novel_smiles(true_unique, pred_unique)
+    novelty = len(pred_novel) / len(pred_unique) if len(pred_unique) > 0 else 0
+    # Difference between Energy distributions
+    energies = wasserstein_distance_between_energies(true_valid_and_connected, pred_valid_and_connected)
+    return {
+        'validity': validity,
+        'validity_and_connectivity': validity_and_connectivity,
+        'uniqueness': uniqueness,
+        'novelty': novelty,
+        'energies': energies,
+    }
+# def check_stability(positions, atom_types):
+#     assert len(positions.shape) == 2
+#     assert positions.shape[1] == 3
+#     x = positions[:, 0]
+#     y = positions[:, 1]
+#     z = positions[:, 2]
+#
+#     nr_bonds = np.zeros(len(x), dtype='int')
+#     for i in range(len(x)):
+#         for j in range(i + 1, len(x)):
+#             p1 = np.array([x[i], y[i], z[i]])
+#             p2 = np.array([x[j], y[j], z[j]])
+#             dist = np.sqrt(np.sum((p1 - p2) ** 2))
+#             atom1, atom2 = const.IDX2ATOM[atom_types[i].item()], const.IDX2ATOM[atom_types[j].item()]
+#             order = get_bond_order(atom1, atom2, dist)
+#             nr_bonds[i] += order
+#             nr_bonds[j] += order
+#     nr_stable_bonds = 0
+#     for atom_type_i, nr_bonds_i in zip(atom_types, nr_bonds):
+#         possible_bonds = const.ALLOWED_BONDS[const.IDX2ATOM[atom_type_i.item()]]
+#         if type(possible_bonds) == int:
+#             is_stable = possible_bonds == nr_bonds_i
+#         else:
+#             is_stable = nr_bonds_i in possible_bonds
+#         nr_stable_bonds += int(is_stable)
+#
+#     molecule_stable = nr_stable_bonds == len(x)
+#     return molecule_stable, nr_stable_bonds, len(x)
+#
+#
+# def count_stable_molecules(one_hot, x, node_mask):
+#     stable_molecules = 0
+#     for i in range(len(one_hot)):
+#         mol_size = node_mask[i].sum()
+#         atom_types = one_hot[i][:mol_size, :].argmax(dim=1).detach().cpu()
+#         positions = x[i][:mol_size, :].detach().cpu()
+#         stable, _, _ = check_stability(positions, atom_types)
+#         stable_molecules += int(stable)
+#
+#     return stable_molecules

src/molecule_builder.py ADDED Viewed

	@@ -0,0 +1,102 @@

+import torch
+import numpy as np
+from rdkit import Chem, Geometry
+from src import const
+def create_conformer(coords):
+    conformer = Chem.Conformer()
+    for i, (x, y, z) in enumerate(coords):
+        conformer.SetAtomPosition(i, Geometry.Point3D(x, y, z))
+    return conformer
+def build_molecules(one_hot, x, node_mask, is_geom, margins=const.MARGINS_EDM):
+    molecules = []
+    for i in range(len(one_hot)):
+        mask = node_mask[i].squeeze() == 1
+        atom_types = one_hot[i][mask].argmax(dim=1).detach().cpu()
+        positions = x[i][mask].detach().cpu()
+        mol = build_molecule(positions, atom_types, is_geom, margins=margins)
+        molecules.append(mol)
+    return molecules
+def build_molecule(positions, atom_types, is_geom, margins=const.MARGINS_EDM):
+    idx2atom = const.GEOM_IDX2ATOM if is_geom else const.IDX2ATOM
+    X, A, E = build_xae_molecule(positions, atom_types, is_geom=is_geom, margins=margins)
+    mol = Chem.RWMol()
+    for atom in X:
+        a = Chem.Atom(idx2atom[atom.item()])
+        mol.AddAtom(a)
+    all_bonds = torch.nonzero(A)
+    for bond in all_bonds:
+        mol.AddBond(bond[0].item(), bond[1].item(), const.BOND_DICT[E[bond[0], bond[1]].item()])
+    mol.AddConformer(create_conformer(positions.detach().cpu().numpy().astype(np.float64)))
+    return mol
+def build_xae_molecule(positions, atom_types, is_geom, margins=const.MARGINS_EDM):
+    """ Returns a triplet (X, A, E): atom_types, adjacency matrix, edge_types
+        args:
+        positions: N x 3  (already masked to keep final number nodes)
+        atom_types: N
+        returns:
+        X: N         (int)
+        A: N x N     (bool) (binary adjacency matrix)
+        E: N x N     (int)  (bond type, 0 if no bond) such that A = E.bool()
+    """
+    n = positions.shape[0]
+    X = atom_types
+    A = torch.zeros((n, n), dtype=torch.bool)
+    E = torch.zeros((n, n), dtype=torch.int)
+    idx2atom = const.GEOM_IDX2ATOM if is_geom else const.IDX2ATOM
+    pos = positions.unsqueeze(0)
+    dists = torch.cdist(pos, pos, p=2).squeeze(0)
+    for i in range(n):
+        for j in range(i):
+            pair = sorted([atom_types[i], atom_types[j]])
+            order = get_bond_order(idx2atom[pair[0].item()], idx2atom[pair[1].item()], dists[i, j], margins=margins)
+            # TODO: a batched version of get_bond_order to avoid the for loop
+            if order > 0:
+                # Warning: the graph should be DIRECTED
+                A[i, j] = 1
+                E[i, j] = order
+    return X, A, E
+def get_bond_order(atom1, atom2, distance, check_exists=True, margins=const.MARGINS_EDM):
+    distance = 100 * distance  # We change the metric
+    # Check exists for large molecules where some atom pairs do not have a
+    # typical bond length.
+    if check_exists:
+        if atom1 not in const.BONDS_1:
+            return 0
+        if atom2 not in const.BONDS_1[atom1]:
+            return 0
+    # margin1, margin2 and margin3 have been tuned to maximize the stability of the QM9 true samples
+    if distance < const.BONDS_1[atom1][atom2] + margins[0]:
+        # Check if atoms in bonds2 dictionary.
+        if atom1 in const.BONDS_2 and atom2 in const.BONDS_2[atom1]:
+            thr_bond2 = const.BONDS_2[atom1][atom2] + margins[1]
+            if distance < thr_bond2:
+                if atom1 in const.BONDS_3 and atom2 in const.BONDS_3[atom1]:
+                    thr_bond3 = const.BONDS_3[atom1][atom2] + margins[2]
+                    if distance < thr_bond3:
+                        return 3  # Triple
+                return 2  # Double
+        return 1  # Single
+    return 0  # No bond

src/noise.py ADDED Viewed

	@@ -0,0 +1,169 @@

+import torch
+import torch.nn.functional as F
+import math
+import numpy as np
+def clip_noise_schedule(alphas2, clip_value=0.001):
+    """
+    For a noise schedule given by alpha^2, this clips alpha_t / alpha_t-1. This may help improve stability during
+    sampling.
+    """
+    alphas2 = np.concatenate([np.ones(1), alphas2], axis=0)
+    alphas_step = (alphas2[1:] / alphas2[:-1])
+    alphas_step = np.clip(alphas_step, a_min=clip_value, a_max=1.)
+    alphas2 = np.cumprod(alphas_step, axis=0)
+    return alphas2
+def polynomial_schedule(timesteps: int, s=1e-4, power=3.):
+    """
+    A noise schedule based on a simple polynomial equation: 1 - x^power.
+    """
+    steps = timesteps + 1
+    x = np.linspace(0, steps, steps)
+    alphas2 = (1 - np.power(x / steps, power)) ** 2
+    alphas2 = clip_noise_schedule(alphas2, clip_value=0.001)
+    precision = 1 - 2 * s
+    alphas2 = precision * alphas2 + s
+    return alphas2
+def cosine_beta_schedule(timesteps, s=0.008, raise_to_power: float = 1):
+    """
+    cosine schedule
+    as proposed in https://openreview.net/forum?id=-NEXDKk8gZ
+    """
+    steps = timesteps + 2
+    x = np.linspace(0, steps, steps)
+    alphas_cumprod = np.cos(((x / steps) + s) / (1 + s) * np.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    betas = np.clip(betas, a_min=0, a_max=0.999)
+    alphas = 1. - betas
+    alphas_cumprod = np.cumprod(alphas, axis=0)
+    if raise_to_power != 1:
+        alphas_cumprod = np.power(alphas_cumprod, raise_to_power)
+    return alphas_cumprod
+class PositiveLinear(torch.nn.Module):
+    """Linear layer with weights forced to be positive."""
+    def __init__(self, in_features: int, out_features: int, bias: bool = True,
+                 weight_init_offset: int = -2):
+        super(PositiveLinear, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = torch.nn.Parameter(
+            torch.empty((out_features, in_features)))
+        if bias:
+            self.bias = torch.nn.Parameter(torch.empty(out_features))
+        else:
+            self.register_parameter('bias', None)
+        self.weight_init_offset = weight_init_offset
+        self.reset_parameters()
+    def reset_parameters(self) -> None:
+        torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        with torch.no_grad():
+            self.weight.add_(self.weight_init_offset)
+        if self.bias is not None:
+            fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(self.weight)
+            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+            torch.nn.init.uniform_(self.bias, -bound, bound)
+    def forward(self, x):
+        positive_weight = F.softplus(self.weight)
+        return F.linear(x, positive_weight, self.bias)
+class PredefinedNoiseSchedule(torch.nn.Module):
+    """
+    Predefined noise schedule. Essentially creates a lookup array for predefined (non-learned) noise schedules.
+    """
+    def __init__(self, noise_schedule, timesteps, precision):
+        super(PredefinedNoiseSchedule, self).__init__()
+        self.timesteps = timesteps
+        if noise_schedule == 'cosine':
+            alphas2 = cosine_beta_schedule(timesteps)
+        elif 'polynomial' in noise_schedule:
+            splits = noise_schedule.split('_')
+            assert len(splits) == 2
+            power = float(splits[1])
+            alphas2 = polynomial_schedule(timesteps, s=precision, power=power)
+        else:
+            raise ValueError(noise_schedule)
+        # print('alphas2', alphas2)
+        sigmas2 = 1 - alphas2
+        log_alphas2 = np.log(alphas2)
+        log_sigmas2 = np.log(sigmas2)
+        log_alphas2_to_sigmas2 = log_alphas2 - log_sigmas2
+        # print('gamma', -log_alphas2_to_sigmas2)
+        self.gamma = torch.nn.Parameter(
+            torch.from_numpy(-log_alphas2_to_sigmas2).float(),
+            requires_grad=False)
+    def forward(self, t):
+        t_int = torch.round(t * self.timesteps).long()
+        return self.gamma[t_int]
+class GammaNetwork(torch.nn.Module):
+    """The gamma network models a monotonic increasing function. Construction as in the VDM paper."""
+    def __init__(self):
+        super().__init__()
+        self.l1 = PositiveLinear(1, 1)
+        self.l2 = PositiveLinear(1, 1024)
+        self.l3 = PositiveLinear(1024, 1)
+        self.gamma_0 = torch.nn.Parameter(torch.tensor([-5.]))
+        self.gamma_1 = torch.nn.Parameter(torch.tensor([10.]))
+        self.show_schedule()
+    def show_schedule(self, num_steps=50):
+        t = torch.linspace(0, 1, num_steps).view(num_steps, 1)
+        gamma = self.forward(t)
+        print('Gamma schedule:')
+        print(gamma.detach().cpu().numpy().reshape(num_steps))
+    def gamma_tilde(self, t):
+        l1_t = self.l1(t)
+        return l1_t + self.l3(torch.sigmoid(self.l2(l1_t)))
+    def forward(self, t):
+        zeros, ones = torch.zeros_like(t), torch.ones_like(t)
+        # Not super efficient.
+        gamma_tilde_0 = self.gamma_tilde(zeros)
+        gamma_tilde_1 = self.gamma_tilde(ones)
+        gamma_tilde_t = self.gamma_tilde(t)
+        # Normalize to [0, 1]
+        normalized_gamma = (gamma_tilde_t - gamma_tilde_0) / (
+                gamma_tilde_1 - gamma_tilde_0)
+        # Rescale to [gamma_0, gamma_1]
+        gamma = self.gamma_0 + (self.gamma_1 - self.gamma_0) * normalized_gamma
+        return gamma

src/utils.py ADDED Viewed

	@@ -0,0 +1,348 @@

+import sys
+from datetime import datetime
+import torch
+import numpy as np
+class Logger(object):
+    def __init__(self, logpath, syspart=sys.stdout):
+        self.terminal = syspart
+        self.log = open(logpath, "a")
+    def write(self, message):
+        self.terminal.write(message)
+        self.log.write(message)
+        self.log.flush()
+    def flush(self):
+        # this flush method is needed for python 3 compatibility.
+        # this handles the flush command by doing nothing.
+        # you might want to specify some extra behavior here.
+        pass
+def log(*args):
+    print(f'[{datetime.now()}]', *args)
+class EMA:
+    def __init__(self, beta):
+        super().__init__()
+        self.beta = beta
+    def update_model_average(self, ma_model, current_model):
+        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
+            old_weight, up_weight = ma_params.data, current_params.data
+            ma_params.data = self.update_average(old_weight, up_weight)
+    def update_average(self, old, new):
+        if old is None:
+            return new
+        return old * self.beta + (1 - self.beta) * new
+def sum_except_batch(x):
+    return x.reshape(x.size(0), -1).sum(dim=-1)
+def remove_mean(x):
+    mean = torch.mean(x, dim=1, keepdim=True)
+    x = x - mean
+    return x
+def remove_mean_with_mask(x, node_mask):
+    masked_max_abs_value = (x * (1 - node_mask)).abs().sum().item()
+    assert masked_max_abs_value < 1e-5, f'Error {masked_max_abs_value} too high'
+    N = node_mask.sum(1, keepdims=True)
+    mean = torch.sum(x, dim=1, keepdim=True) / N
+    x = x - mean * node_mask
+    return x
+def remove_partial_mean_with_mask(x, node_mask, center_of_mass_mask):
+    """
+    Subtract center of mass of fragments from coordinates of all atoms
+    """
+    x_masked = x * center_of_mass_mask
+    N = center_of_mass_mask.sum(1, keepdims=True)
+    mean = torch.sum(x_masked, dim=1, keepdim=True) / N
+    x = x - mean * node_mask
+    return x
+def assert_mean_zero(x):
+    mean = torch.mean(x, dim=1, keepdim=True)
+    assert mean.abs().max().item() < 1e-4
+def assert_mean_zero_with_mask(x, node_mask, eps=1e-10):
+    assert_correctly_masked(x, node_mask)
+    largest_value = x.abs().max().item()
+    error = torch.sum(x, dim=1, keepdim=True).abs().max().item()
+    rel_error = error / (largest_value + eps)
+    assert rel_error < 1e-2, f'Mean is not zero, relative_error {rel_error}'
+def assert_partial_mean_zero_with_mask(x, node_mask, center_of_mass_mask, eps=1e-10):
+    assert_correctly_masked(x, node_mask)
+    x_masked = x * center_of_mass_mask
+    largest_value = x_masked.abs().max().item()
+    error = torch.sum(x_masked, dim=1, keepdim=True).abs().max().item()
+    rel_error = error / (largest_value + eps)
+    assert rel_error < 1e-2, f'Partial mean is not zero, relative_error {rel_error}'
+def assert_correctly_masked(variable, node_mask):
+    assert (variable * (1 - node_mask)).abs().max().item() < 1e-4, \
+        'Variables not masked properly.'
+def check_mask_correct(variables, node_mask):
+    for i, variable in enumerate(variables):
+        if len(variable) > 0:
+            assert_correctly_masked(variable, node_mask)
+def center_gravity_zero_gaussian_log_likelihood(x):
+    assert len(x.size()) == 3
+    B, N, D = x.size()
+    assert_mean_zero(x)
+    # r is invariant to a basis change in the relevant hyperplane.
+    r2 = sum_except_batch(x.pow(2))
+    # The relevant hyperplane is (N-1) * D dimensional.
+    degrees_of_freedom = (N-1) * D
+    # Normalizing constant and logpx are computed:
+    log_normalizing_constant = -0.5 * degrees_of_freedom * np.log(2*np.pi)
+    log_px = -0.5 * r2 + log_normalizing_constant
+    return log_px
+def sample_center_gravity_zero_gaussian(size, device):
+    assert len(size) == 3
+    x = torch.randn(size, device=device)
+    # This projection only works because Gaussian is rotation invariant around
+    # zero and samples are independent!
+    x_projected = remove_mean(x)
+    return x_projected
+def center_gravity_zero_gaussian_log_likelihood_with_mask(x, node_mask):
+    assert len(x.size()) == 3
+    B, N_embedded, D = x.size()
+    assert_mean_zero_with_mask(x, node_mask)
+    # r is invariant to a basis change in the relevant hyperplane, the masked
+    # out values will have zero contribution.
+    r2 = sum_except_batch(x.pow(2))
+    # The relevant hyperplane is (N-1) * D dimensional.
+    N = node_mask.squeeze(2).sum(1)  # N has shape [B]
+    degrees_of_freedom = (N-1) * D
+    # Normalizing constant and logpx are computed:
+    log_normalizing_constant = -0.5 * degrees_of_freedom * np.log(2*np.pi)
+    log_px = -0.5 * r2 + log_normalizing_constant
+    return log_px
+def sample_center_gravity_zero_gaussian_with_mask(size, device, node_mask):
+    assert len(size) == 3
+    x = torch.randn(size, device=device)
+    x_masked = x * node_mask
+    # This projection only works because Gaussian is rotation invariant around
+    # zero and samples are independent!
+    # TODO: check it
+    x_projected = remove_mean_with_mask(x_masked, node_mask)
+    return x_projected
+def standard_gaussian_log_likelihood(x):
+    # Normalizing constant and logpx are computed:
+    log_px = sum_except_batch(-0.5 * x * x - 0.5 * np.log(2*np.pi))
+    return log_px
+def sample_gaussian(size, device):
+    x = torch.randn(size, device=device)
+    return x
+def standard_gaussian_log_likelihood_with_mask(x, node_mask):
+    # Normalizing constant and logpx are computed:
+    log_px_elementwise = -0.5 * x * x - 0.5 * np.log(2*np.pi)
+    log_px = sum_except_batch(log_px_elementwise * node_mask)
+    return log_px
+def sample_gaussian_with_mask(size, device, node_mask):
+    x = torch.randn(size, device=device)
+    x_masked = x * node_mask
+    return x_masked
+def concatenate_features(x, h):
+    xh = torch.cat([x, h['categorical']], dim=2)
+    if 'integer' in h:
+        xh = torch.cat([xh, h['integer']], dim=2)
+    return xh
+def split_features(z, n_dims, num_classes, include_charges):
+    assert z.size(2) == n_dims + num_classes + include_charges
+    x = z[:, :, 0:n_dims]
+    h = {'categorical': z[:, :, n_dims:n_dims+num_classes]}
+    if include_charges:
+        h['integer'] = z[:, :, n_dims+num_classes:n_dims+num_classes+1]
+    return x, h
+# For gradient clipping
+class Queue:
+    def __init__(self, max_len=50):
+        self.items = []
+        self.max_len = max_len
+    def __len__(self):
+        return len(self.items)
+    def add(self, item):
+        self.items.insert(0, item)
+        if len(self) > self.max_len:
+            self.items.pop()
+    def mean(self):
+        return np.mean(self.items)
+    def std(self):
+        return np.std(self.items)
+def gradient_clipping(flow, gradnorm_queue):
+    # Allow gradient norm to be 150% + 2 * stdev of the recent history.
+    max_grad_norm = 1.5 * gradnorm_queue.mean() + 2 * gradnorm_queue.std()
+    # Clips gradient and returns the norm
+    grad_norm = torch.nn.utils.clip_grad_norm_(
+        flow.parameters(), max_norm=max_grad_norm, norm_type=2.0)
+    if float(grad_norm) > max_grad_norm:
+        gradnorm_queue.add(float(max_grad_norm))
+    else:
+        gradnorm_queue.add(float(grad_norm))
+    if float(grad_norm) > max_grad_norm:
+        print(f'Clipped gradient with value {grad_norm:.1f} while allowed {max_grad_norm:.1f}')
+    return grad_norm
+def disable_rdkit_logging():
+    """
+    Disables RDKit whiny logging.
+    """
+    import rdkit.rdBase as rkrb
+    import rdkit.RDLogger as rkl
+    logger = rkl.logger()
+    logger.setLevel(rkl.ERROR)
+    rkrb.DisableLog('rdApp.error')
+class FoundNaNException(Exception):
+    def __init__(self, x, h):
+        x_nan_idx = self.find_nan_idx(x)
+        h_nan_idx = self.find_nan_idx(h)
+        self.x_h_nan_idx = x_nan_idx & h_nan_idx
+        self.only_x_nan_idx = x_nan_idx.difference(h_nan_idx)
+        self.only_h_nan_idx = h_nan_idx.difference(x_nan_idx)
+    @staticmethod
+    def find_nan_idx(z):
+        idx = set()
+        for i in range(z.shape[0]):
+            if torch.any(torch.isnan(z[i])):
+                idx.add(i)
+        return idx
+def get_batch_idx_for_animation(batch_size, batch_idx):
+    batch_indices = []
+    mol_indices = []
+    for idx in [0, 110, 360]:
+        if idx // batch_size == batch_idx:
+            batch_indices.append(idx % batch_size)
+            mol_indices.append(idx)
+    return batch_indices, mol_indices
+# Rotation data augmntation
+def random_rotation(x):
+    bs, n_nodes, n_dims = x.size()
+    device = x.device
+    angle_range = np.pi * 2
+    if n_dims == 2:
+        theta = torch.rand(bs, 1, 1).to(device) * angle_range - np.pi
+        cos_theta = torch.cos(theta)
+        sin_theta = torch.sin(theta)
+        R_row0 = torch.cat([cos_theta, -sin_theta], dim=2)
+        R_row1 = torch.cat([sin_theta, cos_theta], dim=2)
+        R = torch.cat([R_row0, R_row1], dim=1)
+        x = x.transpose(1, 2)
+        x = torch.matmul(R, x)
+        x = x.transpose(1, 2)
+    elif n_dims == 3:
+        # Build Rx
+        Rx = torch.eye(3).unsqueeze(0).repeat(bs, 1, 1).to(device)
+        theta = torch.rand(bs, 1, 1).to(device) * angle_range - np.pi
+        cos = torch.cos(theta)
+        sin = torch.sin(theta)
+        Rx[:, 1:2, 1:2] = cos
+        Rx[:, 1:2, 2:3] = sin
+        Rx[:, 2:3, 1:2] = - sin
+        Rx[:, 2:3, 2:3] = cos
+        # Build Ry
+        Ry = torch.eye(3).unsqueeze(0).repeat(bs, 1, 1).to(device)
+        theta = torch.rand(bs, 1, 1).to(device) * angle_range - np.pi
+        cos = torch.cos(theta)
+        sin = torch.sin(theta)
+        Ry[:, 0:1, 0:1] = cos
+        Ry[:, 0:1, 2:3] = -sin
+        Ry[:, 2:3, 0:1] = sin
+        Ry[:, 2:3, 2:3] = cos
+        # Build Rz
+        Rz = torch.eye(3).unsqueeze(0).repeat(bs, 1, 1).to(device)
+        theta = torch.rand(bs, 1, 1).to(device) * angle_range - np.pi
+        cos = torch.cos(theta)
+        sin = torch.sin(theta)
+        Rz[:, 0:1, 0:1] = cos
+        Rz[:, 0:1, 1:2] = sin
+        Rz[:, 1:2, 0:1] = -sin
+        Rz[:, 1:2, 1:2] = cos
+        x = x.transpose(1, 2)
+        x = torch.matmul(Rx, x)
+        #x = torch.matmul(Rx.transpose(1, 2), x)
+        x = torch.matmul(Ry, x)
+        #x = torch.matmul(Ry.transpose(1, 2), x)
+        x = torch.matmul(Rz, x)
+        #x = torch.matmul(Rz.transpose(1, 2), x)
+        x = x.transpose(1, 2)
+    else:
+        raise Exception("Not implemented Error")
+    return x.contiguous()

src/visualizer.py ADDED Viewed

	@@ -0,0 +1,227 @@

+import torch
+import os
+import imageio
+import matplotlib.pyplot as plt
+import numpy as np
+import glob
+import random
+from sklearn.decomposition import PCA
+from src import const
+from src.molecule_builder import get_bond_order
+def save_xyz_file(path, one_hot, positions, node_mask, names, is_geom, suffix=''):
+    idx2atom = const.GEOM_IDX2ATOM if is_geom else const.IDX2ATOM
+    for batch_i in range(one_hot.size(0)):
+        mask = node_mask[batch_i].squeeze()
+        n_atoms = mask.sum()
+        atom_idx = torch.where(mask)[0]
+        f = open(os.path.join(path, f'{names[batch_i]}_{suffix}.xyz'), "w")
+        f.write("%d\n\n" % n_atoms)
+        atoms = torch.argmax(one_hot[batch_i], dim=1)
+        for atom_i in atom_idx:
+            atom = atoms[atom_i].item()
+            atom = idx2atom[atom]
+            f.write("%s %.9f %.9f %.9f\n" % (
+                atom, positions[batch_i, atom_i, 0], positions[batch_i, atom_i, 1], positions[batch_i, atom_i, 2]
+            ))
+        f.close()
+def load_xyz_files(path, suffix=''):
+    files = []
+    for fname in os.listdir(path):
+        if fname.endswith(f'_{suffix}.xyz'):
+            files.append(fname)
+    files = sorted(files, key=lambda f: -int(f.replace(f'_{suffix}.xyz', '').split('_')[-1]))
+    return [os.path.join(path, fname) for fname in files]
+def load_molecule_xyz(file, is_geom):
+    atom2idx = const.GEOM_ATOM2IDX if is_geom else const.ATOM2IDX
+    idx2atom = const.GEOM_IDX2ATOM if is_geom else const.IDX2ATOM
+    with open(file, encoding='utf8') as f:
+        n_atoms = int(f.readline())
+        one_hot = torch.zeros(n_atoms, len(idx2atom))
+        charges = torch.zeros(n_atoms, 1)
+        positions = torch.zeros(n_atoms, 3)
+        f.readline()
+        atoms = f.readlines()
+        for i in range(n_atoms):
+            atom = atoms[i].split(' ')
+            atom_type = atom[0]
+            one_hot[i, atom2idx[atom_type]] = 1
+            position = torch.Tensor([float(e) for e in atom[1:]])
+            positions[i, :] = position
+        return positions, one_hot, charges
+def draw_sphere(ax, x, y, z, size, color, alpha):
+    u = np.linspace(0, 2 * np.pi, 100)
+    v = np.linspace(0, np.pi, 100)
+    xs = size * np.outer(np.cos(u), np.sin(v))
+    ys = size * np.outer(np.sin(u), np.sin(v)) #* 0.8
+    zs = size * np.outer(np.ones(np.size(u)), np.cos(v))
+    ax.plot_surface(x + xs, y + ys, z + zs, rstride=2, cstride=2, color=color, alpha=alpha)
+def plot_molecule(ax, positions, atom_type, alpha, spheres_3d, hex_bg_color, is_geom, fragment_mask=None):
+    x = positions[:, 0]
+    y = positions[:, 1]
+    z = positions[:, 2]
+    # Hydrogen, Carbon, Nitrogen, Oxygen, Flourine
+    idx2atom = const.GEOM_IDX2ATOM if is_geom else const.IDX2ATOM
+    colors_dic = np.array(const.COLORS)
+    radius_dic = np.array(const.RADII)
+    area_dic = 1500 * radius_dic ** 2
+    areas = area_dic[atom_type]
+    radii = radius_dic[atom_type]
+    colors = colors_dic[atom_type]
+    if fragment_mask is None:
+        fragment_mask = torch.ones(len(x))
+    for i in range(len(x)):
+        for j in range(i + 1, len(x)):
+            p1 = np.array([x[i], y[i], z[i]])
+            p2 = np.array([x[j], y[j], z[j]])
+            dist = np.sqrt(np.sum((p1 - p2) ** 2))
+            atom1, atom2 = idx2atom[atom_type[i]], idx2atom[atom_type[j]]
+            draw_edge_int = get_bond_order(atom1, atom2, dist)
+            line_width = (3 - 2) * 2 * 2
+            draw_edge = draw_edge_int > 0
+            if draw_edge:
+                if draw_edge_int == 4:
+                    linewidth_factor = 1.5
+                else:
+                    linewidth_factor = 1
+                linewidth_factor *= 0.5
+                ax.plot(
+                    [x[i], x[j]], [y[i], y[j]], [z[i], z[j]],
+                    linewidth=line_width * linewidth_factor * 2,
+                    c=hex_bg_color,
+                    alpha=alpha
+                )
+    # from pdb import set_trace
+    # set_trace()
+    if spheres_3d:
+        # idx = torch.where(fragment_mask[:len(x)] == 0)[0]
+        # ax.scatter(
+        #     x[idx],
+        #     y[idx],
+        #     z[idx],
+        #     alpha=0.9 * alpha,
+        #     edgecolors='#FCBA03',
+        #     facecolors='none',
+        #     linewidths=2,
+        #     s=900
+        # )
+        for i, j, k, s, c, f in zip(x, y, z, radii, colors, fragment_mask):
+            if f == 1:
+                alpha = 1.0
+            draw_sphere(ax, i.item(), j.item(), k.item(), 0.5 * s, c, alpha)
+    else:
+        ax.scatter(x, y, z, s=areas, alpha=0.9 * alpha, c=colors)
+def plot_data3d(positions, atom_type, is_geom, camera_elev=0, camera_azim=0, save_path=None, spheres_3d=False,
+                bg='black', alpha=1., fragment_mask=None):
+    black = (0, 0, 0)
+    white = (1, 1, 1)
+    hex_bg_color = '#FFFFFF' if bg == 'black' else '#000000' #'#666666'
+    fig = plt.figure(figsize=(10, 10))
+    ax = fig.add_subplot(projection='3d')
+    ax.set_aspect('auto')
+    ax.view_init(elev=camera_elev, azim=camera_azim)
+    if bg == 'black':
+        ax.set_facecolor(black)
+    else:
+        ax.set_facecolor(white)
+    ax.xaxis.pane.set_alpha(0)
+    ax.yaxis.pane.set_alpha(0)
+    ax.zaxis.pane.set_alpha(0)
+    ax._axis3don = False
+    if bg == 'black':
+        ax.w_xaxis.line.set_color("black")
+    else:
+        ax.w_xaxis.line.set_color("white")
+    plot_molecule(
+        ax, positions, atom_type, alpha, spheres_3d, hex_bg_color, is_geom=is_geom, fragment_mask=fragment_mask
+    )
+    max_value = positions.abs().max().item()
+    axis_lim = min(40, max(max_value / 1.5 + 0.3, 3.2))
+    ax.set_xlim(-axis_lim, axis_lim)
+    ax.set_ylim(-axis_lim, axis_lim)
+    ax.set_zlim(-axis_lim, axis_lim)
+    dpi = 120 if spheres_3d else 50
+    if save_path is not None:
+        plt.savefig(save_path, bbox_inches='tight', pad_inches=0.0, dpi=dpi)
+        # plt.savefig(save_path, bbox_inches='tight', pad_inches=0.0, dpi=dpi, transparent=True)
+        if spheres_3d:
+            img = imageio.imread(save_path)
+            img_brighter = np.clip(img * 1.4, 0, 255).astype('uint8')
+            imageio.imsave(save_path, img_brighter)
+    else:
+        plt.show()
+    plt.close()
+def visualize_chain(
+        path, spheres_3d=False, bg="black", alpha=1.0, wandb=None, mode="chain", is_geom=False, fragment_mask=None
+):
+    files = load_xyz_files(path)
+    save_paths = []
+    # Fit PCA to the final molecule – to obtain the best orientation for visualization
+    positions, one_hot, charges = load_molecule_xyz(files[-1], is_geom=is_geom)
+    pca = PCA(n_components=3)
+    pca.fit(positions)
+    for i in range(len(files)):
+        file = files[i]
+        positions, one_hot, charges = load_molecule_xyz(file, is_geom=is_geom)
+        atom_type = torch.argmax(one_hot, dim=1).numpy()
+        # Transform positions of each frame according to the best orientation of the last frame
+        positions = pca.transform(positions)
+        positions = torch.tensor(positions)
+        fn = file[:-4] + '.png'
+        plot_data3d(
+            positions, atom_type,
+            save_path=fn,
+            spheres_3d=spheres_3d,
+            alpha=alpha,
+            bg=bg,
+            camera_elev=90,
+            camera_azim=90,
+            is_geom=is_geom,
+            fragment_mask=fragment_mask,
+        )
+        save_paths.append(fn)
+    imgs = [imageio.imread(fn) for fn in save_paths]
+    dirname = os.path.dirname(save_paths[0])
+    gif_path = dirname + '/output.gif'
+    imageio.mimsave(gif_path, imgs, subrectangles=True)
+    if wandb is not None:
+        wandb.log({mode: [wandb.Video(gif_path, caption=gif_path)]})