Pocket-conditioned generation

app.py

@@ -1,4 +1,5 @@
 import argparse
+import shutil
 
 import gradio as gr
 import numpy as np
@@ -9,10 +10,12 @@ import output
 
 from rdkit import Chem
 from src import const
-from src.datasets import get_dataloader, collate_with_fragment_edges, parse_molecule
+from src.datasets import get_dataloader, collate_with_fragment_edges, parse_molecule, MOADDataset
 from src.lightning import DDPM
 from src.linker_size_lightning import SizeClassifier
-from src.generation import N_SAMPLES, generate_linkers, try_to_convert_to_sdf
+from src.generation import N_SAMPLES, generate_linkers, try_to_convert_to_sdf, get_pocket
+from zipfile import ZipFile
+
 
 MODELS_METADATA = {
     'geom_difflinker': {
@@ -85,65 +88,167 @@ def read_molecule(path):
     raise Exception('Unknown file extension')
 
 
-def show_input(input_file):
-    if input_file is None:
-        return ['', gr.Radio.update(visible=False, value='Sample 1'), None]
-    if isinstance(input_file, str):
-        path = input_file
+def read_molecule_file(in_file, allowed_extentions):
+    if isinstance(in_file, str):
+        path = in_file
     else:
-        path = input_file.name
+        path = in_file.name
     extension = path.split('.')[-1]
-    if extension not in ['sdf', 'pdb', 'mol', 'mol2']:
+
+    if extension not in allowed_extentions:
         msg = output.INVALID_FORMAT_MSG.format(extension=extension)
-        return [
-            output.IFRAME_TEMPLATE.format(html=msg),
-            gr.Radio.update(visible=False),
-            None,
-        ]
+        return None, None, msg
 
     try:
-        molecule = read_molecule_content(path)
+        mol = read_molecule(path)
     except Exception as e:
-        return [
-            f'Could not read the molecule: {e}',
-            gr.Radio.update(visible=False),
-            None,
-        ]
+        e = str(e).replace('\'', '')
+        msg = output.ERROR_FORMAT_MSG.format(message=e)
+        return None, None, msg
+
+    if extension == 'pdb':
+        content = Chem.MolToPDBBlock(mol)
+    elif extension in ['mol', 'mol2', 'sdf']:
+        content = Chem.MolToMolBlock(mol, kekulize=False)
+        extension = 'mol'
+    else:
+        raise NotImplementedError
 
-    html = output.INITIAL_RENDERING_TEMPLATE.format(molecule=molecule, fmt=extension)
-    return [
-        output.IFRAME_TEMPLATE.format(html=html),
-        gr.Radio.update(visible=False),
-        None,
-    ]
+    return content, extension, None
+
+
+def show_input(in_fragments, in_protein):
+    vis = ''
+    if in_fragments is not None and in_protein is None:
+        vis = show_fragments(in_fragments)
+    elif in_fragments is None and in_protein is not None:
+        vis = show_target(in_protein)
+    elif in_fragments is not None and in_protein is not None:
+        vis = show_fragments_and_target(in_fragments, in_protein)
+    return [vis, gr.Radio.update(visible=False), None]
+
+
+def show_fragments(in_fragments):
+    molecule, extension, html = read_molecule_file(in_fragments, allowed_extentions=['sdf', 'pdb', 'mol', 'mol2'])
+    if molecule is not None:
+        html = output.FRAGMENTS_RENDERING_TEMPLATE.format(molecule=molecule, fmt=extension)
+
+    return output.IFRAME_TEMPLATE.format(html=html)
+
+
+def show_target(in_protein):
+    molecule, extension, html = read_molecule_file(in_protein, allowed_extentions=['pdb'])
+    if molecule is not None:
+        html = output.TARGET_RENDERING_TEMPLATE.format(molecule=molecule, fmt=extension)
+
+    return output.IFRAME_TEMPLATE.format(html=html)
+
+
+def show_fragments_and_target(in_fragments, in_protein):
+    fragments_molecule, fragments_extension, msg = read_molecule_file(in_fragments, ['sdf', 'pdb', 'mol', 'mol2'])
+    if fragments_molecule is None:
+        return output.IFRAME_TEMPLATE.format(html=msg)
+
+    target_molecule, target_extension, msg = read_molecule_file(in_protein, allowed_extentions=['pdb'])
+    if fragments_molecule is None:
+        return output.IFRAME_TEMPLATE.format(html=msg)
+
+    html = output.FRAGMENTS_AND_TARGET_RENDERING_TEMPLATE.format(
+        molecule=fragments_molecule,
+        fmt=fragments_extension,
+        target=target_molecule,
+        target_fmt=target_extension,
+    )
+
+    return output.IFRAME_TEMPLATE.format(html=html)
+
+
+def clear_fragments_input(in_protein):
+    vis = ''
+    if in_protein is not None:
+        vis = show_target(in_protein)
+    return [None, vis, gr.Radio.update(visible=False), None]
+
+
+def clear_protein_input(in_fragments):
+    vis = ''
+    if in_fragments is not None:
+        vis = show_fragments(in_fragments)
+    return [None, vis, gr.Radio.update(visible=False), None]
+
+
+def click_on_example(example):
+    print('Clicked:', example)
+    fragment_fname, target_fname = example
+    fragment_path = f'examples/{fragment_fname}' if fragment_fname != '' else None
+    target_path = f'examples/{target_fname}' if target_fname != '' else None
+    return [fragment_path, target_path, 50, 0] + show_input(fragment_path, target_path)
 
 
 def draw_sample(idx, out_files):
+    with_protein = (len(out_files) == N_SAMPLES + 3)
+
     if isinstance(idx, str):
         idx = int(idx.strip().split(' ')[-1]) - 1
 
-    in_file = out_files[0]
+    in_file = out_files[1]
     in_sdf = in_file if isinstance(in_file, str) else in_file.name
+    input_fragments_content = read_molecule_content(in_sdf)
+    fragments_fmt = in_sdf.split('.')[-1]
 
-    out_file = out_files[idx + 1]
-    out_sdf = out_file if isinstance(out_file, str) else out_file.name
+    offset = 2
+    input_target_content = None
+    target_fmt = None
+    if with_protein:
+        offset += 1
+        in_pdb = out_files[2] if isinstance(out_files[2], str) else out_files[2].name
+        input_target_content = read_molecule_content(in_pdb)
+        target_fmt = in_pdb.split('.')[-1]
 
-    input_fragments_content = read_molecule_content(in_sdf)
+    out_file = out_files[idx + offset]
+    out_sdf = out_file if isinstance(out_file, str) else out_file.name
     generated_molecule_content = read_molecule_content(out_sdf)
-
-    fragments_fmt = in_sdf.split('.')[-1]
     molecule_fmt = out_sdf.split('.')[-1]
 
-    html = output.SAMPLES_RENDERING_TEMPLATE.format(
-        fragments=input_fragments_content,
-        fragments_fmt=fragments_fmt,
-        molecule=generated_molecule_content,
-        molecule_fmt=molecule_fmt,
-    )
+    if with_protein:
+        html = output.SAMPLES_WITH_TARGET_RENDERING_TEMPLATE.format(
+            fragments=input_fragments_content,
+            fragments_fmt=fragments_fmt,
+            molecule=generated_molecule_content,
+            molecule_fmt=molecule_fmt,
+            target=input_target_content,
+            target_fmt=target_fmt,
+        )
+    else:
+        html = output.SAMPLES_RENDERING_TEMPLATE.format(
+            fragments=input_fragments_content,
+            fragments_fmt=fragments_fmt,
+            molecule=generated_molecule_content,
+            molecule_fmt=molecule_fmt,
+        )
     return output.IFRAME_TEMPLATE.format(html=html)
 
 
-def generate(input_file, n_steps, n_atoms, radio_samples, selected_atoms):
+def compress(output_fnames, name):
+    archive_path = f'results/all_files_{name}.zip'
+    with ZipFile(archive_path, 'w') as archive:
+        for fname in output_fnames:
+            archive.write(fname)
+
+    return archive_path
+
+
+def generate(in_fragments, in_protein, n_steps, n_atoms, radio_samples, selected_atoms):
+    if in_fragments is None:
+        return [None, None, None, None]
+
+    if in_protein is None:
+        return generate_without_pocket(in_fragments, n_steps, n_atoms, radio_samples, selected_atoms)
+    else:
+        return generate_with_pocket(in_fragments, in_protein, n_steps, n_atoms, radio_samples, selected_atoms)
+
+
+def generate_without_pocket(input_file, n_steps, n_atoms, radio_samples, selected_atoms):
     # Parsing selected atoms (javascript output)
     selected_atoms = selected_atoms.strip()
     if selected_atoms == '':
@@ -157,9 +262,6 @@ def generate(input_file, n_steps, n_atoms, radio_samples, selected_atoms):
     else:
         selected_model_name = 'geom_difflinker_given_anchors'
 
-    if input_file is None:
-        return [None, None, None, None]
-
     print(f'Start generating with model {selected_model_name}, selected_atoms:', selected_atoms)
     ddpm = diffusion_models[selected_model_name]
     path = input_file.name
@@ -170,20 +272,25 @@ def generate(input_file, n_steps, n_atoms, radio_samples, selected_atoms):
 
     try:
         molecule = read_molecule(path)
-        molecule = Chem.RemoveAllHs(molecule)
+        try:
+            molecule = Chem.RemoveAllHs(molecule)
+        except:
+            pass
         name = '.'.join(path.split('/')[-1].split('.')[:-1])
         inp_sdf = f'results/input_{name}.sdf'
     except Exception as e:
+        e = str(e).replace('\'', '')
         error = f'Could not read the molecule: {e}'
         msg = output.ERROR_FORMAT_MSG.format(message=error)
         return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
 
-    if molecule.GetNumAtoms() > 50:
-        error = f'Too large molecule: upper limit is 50 heavy atoms'
+    if molecule.GetNumAtoms() > 100:
+        error = f'Too large molecule: upper limit is 100 heavy atoms'
         msg = output.ERROR_FORMAT_MSG.format(message=error)
         return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
 
     with Chem.SDWriter(inp_sdf) as w:
+        w.SetKekulize(False)
         w.write(molecule)
 
     positions, one_hot, charges = parse_molecule(molecule, is_geom=True)
@@ -227,14 +334,152 @@ def generate(input_file, n_steps, n_atoms, radio_samples, selected_atoms):
 
     for data in dataloader:
         try:
-            generate_linkers(ddpm=ddpm, data=data, sample_fn=sample_fn, name=name)
+            generate_linkers(ddpm=ddpm, data=data, sample_fn=sample_fn, name=name, with_pocket=False)
         except Exception as e:
+            e = str(e).replace('\'', '')
             error = f'Caught exception while generating linkers: {e}'
             msg = output.ERROR_FORMAT_MSG.format(message=error)
             return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
 
     out_files = try_to_convert_to_sdf(name)
     out_files = [inp_sdf] + out_files
+    out_files = [compress(out_files, name=name)] + out_files
+
+    return [
+        draw_sample(radio_samples, out_files),
+        out_files,
+        gr.Radio.update(visible=True),
+        None
+    ]
+
+
+def generate_with_pocket(in_fragments, in_protein, n_steps, n_atoms, radio_samples, selected_atoms):
+    # Parsing selected atoms (javascript output)
+    selected_atoms = selected_atoms.strip()
+    if selected_atoms == '':
+        selected_atoms = []
+    else:
+        selected_atoms = list(map(int, selected_atoms.split(',')))
+
+    # Selecting model
+    if len(selected_atoms) == 0:
+        selected_model_name = 'pockets_difflinker'
+    else:
+        selected_model_name = 'pockets_difflinker_given_anchors'
+
+    print(f'Start generating with model {selected_model_name}, selected_atoms:', selected_atoms)
+    ddpm = diffusion_models[selected_model_name]
+
+    fragments_path = in_fragments.name
+    fragments_extension = fragments_path.split('.')[-1]
+    if fragments_extension not in ['sdf', 'pdb', 'mol', 'mol2']:
+        msg = output.INVALID_FORMAT_MSG.format(extension=fragments_extension)
+        return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
+
+    protein_path = in_protein.name
+    protein_extension = protein_path.split('.')[-1]
+    if protein_extension not in ['pdb']:
+        msg = output.INVALID_FORMAT_MSG.format(extension=protein_extension)
+        return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
+
+    try:
+        fragments_mol = read_molecule(fragments_path)
+        name = '.'.join(fragments_path.split('/')[-1].split('.')[:-1])
+    except Exception as e:
+        e = str(e).replace('\'', '')
+        error = f'Could not read the molecule: {e}'
+        msg = output.ERROR_FORMAT_MSG.format(message=error)
+        return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
+
+    if fragments_mol.GetNumAtoms() > 100:
+        error = f'Too large molecule: upper limit is 100 heavy atoms'
+        msg = output.ERROR_FORMAT_MSG.format(message=error)
+        return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
+
+    inp_sdf = f'results/input_{name}.sdf'
+    with Chem.SDWriter(inp_sdf) as w:
+        w.SetKekulize(False)
+        w.write(fragments_mol)
+
+    inp_pdb = f'results/target_{name}.pdb'
+    shutil.copy(protein_path, inp_pdb)
+
+    frag_pos, frag_one_hot, frag_charges = parse_molecule(fragments_mol, is_geom=True)
+    pocket_pos, pocket_one_hot, pocket_charges = get_pocket(fragments_mol, protein_path)
+    print(f'Detected pocket with {len(pocket_pos)} atoms')
+
+    positions = np.concatenate([frag_pos, pocket_pos], axis=0)
+    one_hot = np.concatenate([frag_one_hot, pocket_one_hot], axis=0)
+    charges = np.concatenate([frag_charges, pocket_charges], axis=0)
+    anchors = np.zeros_like(charges)
+    anchors[selected_atoms] = 1
+
+    fragment_only_mask = np.concatenate([
+        np.ones_like(frag_charges),
+        np.zeros_like(pocket_charges),
+    ])
+    pocket_mask = np.concatenate([
+        np.zeros_like(frag_charges),
+        np.ones_like(pocket_charges),
+    ])
+    linker_mask = np.concatenate([
+        np.zeros_like(frag_charges),
+        np.zeros_like(pocket_charges),
+    ])
+    fragment_mask = np.concatenate([
+        np.ones_like(frag_charges),
+        np.ones_like(pocket_charges),
+    ])
+    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.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),
+    }] * N_SAMPLES
+    dataset = MOADDataset(data=dataset)
+    ddpm.val_dataset = dataset
+
+    dataloader = get_dataloader(dataset, batch_size=N_SAMPLES, collate_fn=collate_with_fragment_edges)
+    print('Created dataloader')
+
+    ddpm.edm.T = n_steps
+
+    if n_atoms == 0:
+        def sample_fn(_data):
+            out, _ = size_nn.forward(_data, return_loss=False)
+            probabilities = torch.softmax(out, 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
+    else:
+        def sample_fn(_data):
+            return torch.ones(_data['positions'].shape[0], device=device, dtype=torch.long) * n_atoms
+
+    for data in dataloader:
+        try:
+            generate_linkers(ddpm=ddpm, data=data, sample_fn=sample_fn, name=name, with_pocket=True)
+        except Exception as e:
+            e = str(e).replace('\'', '')
+            error = f'Caught exception while generating linkers: {e}'
+            msg = output.ERROR_FORMAT_MSG.format(message=error)
+            return [output.IFRAME_TEMPLATE.format(html=msg), None, None, None]
+
+    out_files = try_to_convert_to_sdf(name)
+    out_files = [inp_sdf, inp_pdb] + out_files
+    out_files = [compress(out_files, name=name)] + out_files
 
     return [
         draw_sample(radio_samples, out_files),
@@ -260,19 +505,34 @@ with demo:
     with gr.Box():
         with gr.Row():
             with gr.Column():
-                gr.Markdown('## Input Fragments')
+                gr.Markdown('## Input')
                 gr.Markdown('Upload the file with 3D-coordinates of the input fragments in .pdb, .mol2 or .sdf format:')
-                input_file = gr.File(file_count='single', label='Input Fragments')
-                n_steps = gr.Slider(minimum=10, maximum=500, label="Number of Denoising Steps", step=10)
+                with gr.Column():
+                    input_fragments_file = gr.File(
+                        file_count='single',
+                        label='Input Fragments',
+                        file_types=['.sdf', '.pdb', '.mol', '.mol2']
+                    )
+                # gr.Markdown('(Optionally) upload the file of the target protein in .pdb format:')
+                with gr.Column():
+                    input_protein_file = gr.File(file_count='single', label='Target Protein', file_types=['.pdb'])
+
+                n_steps = gr.Slider(minimum=50, maximum=500, label="Number of Denoising Steps", step=10)
                 n_atoms = gr.Slider(
                     minimum=0, maximum=20,
                     label="Linker Size: DiffLinker will predict it if set to 0",
                     step=1
                 )
                 examples = gr.Dataset(
-                    components=[gr.File(visible=False)],
-                    samples=[['examples/example_1.sdf'], ['examples/example_2.sdf']],
-                    type='index',
+                    components=[gr.File(visible=False), gr.File(visible=False)],
+                    samples=[
+                        ['examples/example_1.sdf', None],
+                        ['examples/example_2.sdf', None],
+                        ['examples/3hz1_fragments.sdf', 'examples/3hz1_protein.pdb'],
+                        ['examples/5ou2_fragments.sdf', 'examples/5ou2_protein.pdb'],
+                    ],
+                    headers=['Fragments', 'Target Protein'],
+                    type='values',
                 ) 
 
                 button = gr.Button('Generate Linker!')
@@ -294,24 +554,34 @@ with demo:
                 )
                 visualization = gr.HTML()
 
-    input_file.change(
+    input_fragments_file.change(
         fn=show_input,
-        inputs=[input_file],
+        inputs=[input_fragments_file, input_protein_file],
         outputs=[visualization, samples, hidden],
     )
-    input_file.clear(
-        fn=lambda: [None, '', gr.Radio.update(visible=False), None],
-        inputs=[],
-        outputs=[input_file, visualization, samples, hidden],
+    input_protein_file.change(
+        fn=show_input,
+        inputs=[input_fragments_file, input_protein_file],
+        outputs=[visualization, samples, hidden],
+    )
+    input_fragments_file.clear(
+        fn=clear_fragments_input,
+        inputs=[input_protein_file],
+        outputs=[input_fragments_file, visualization, samples, hidden],
+    )
+    input_protein_file.clear(
+        fn=clear_protein_input,
+        inputs=[input_fragments_file],
+        outputs=[input_protein_file, visualization, samples, hidden],
     )
     examples.click(
-        fn=lambda idx: [f'examples/example_{idx+1}.sdf', 10, 0] + show_input(f'examples/example_{idx+1}.sdf'),
+        fn=click_on_example,
         inputs=[examples],
-        outputs=[input_file, n_steps, n_atoms, visualization, samples, hidden]
+        outputs=[input_fragments_file, input_protein_file, n_steps, n_atoms, visualization, samples, hidden]
     )
     button.click(
         fn=generate,
-        inputs=[input_file, n_steps, n_atoms, samples, hidden],
+        inputs=[input_fragments_file, input_protein_file, n_steps, n_atoms, samples, hidden],
         outputs=[visualization, output_files, samples, hidden],
         _js=output.RETURN_SELECTION_JS,
     )

examples/3hz1_fragments.sdf

@@ -0,0 +1,54 @@
+fragments
+  PyMOL2.5          3D                             0
+
+ 23 25  0  0  0  0  0  0  0  0999 V2000
+    0.7050   10.1160   25.5000 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.4250   10.6930   24.7810 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.6420   10.9060   25.5370 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.7510   10.5210   26.8370 N   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.6900    9.9510   27.4380 C   0  0  0  0  0  0  0  0  0  0  0  0
+    0.4770    9.7630   26.7990 N   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.6830   11.1870   23.5600 N   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.9660   11.6240   23.5390 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -2.5810   11.4250   24.7070 N   0  0  0  0  0  0  0  0  0  0  0  0
+    1.9520    9.8170   24.8700 N   0  0  0  0  0  0  0  0  0  0  0  0
+    3.1230    9.3980   25.6290 C   0  0  0  0  0  0  0  0  0  0  0  0
+    2.1100    9.7530   23.4320 C   0  0  0  0  0  0  0  0  0  0  0  0
+    7.8600   10.1360   22.6040 C   0  0  0  0  0  0  0  0  0  0  0  0
+    6.5530    9.6800   22.8080 C   0  0  0  0  0  0  0  0  0  0  0  0
+    5.8720   10.7150   23.6130 O   0  0  0  0  0  0  0  0  0  0  0  0
+    6.8390   11.6780   23.7840 C   0  0  0  0  0  0  0  0  0  0  0  0
+    8.0580   11.3690   23.2280 C   0  0  0  0  0  0  0  0  0  0  0  0
+    6.6560   12.9400   24.5720 C   0  0  0  0  0  0  0  0  0  0  0  0
+    7.6630   13.4980   25.2340 N   0  0  0  0  0  0  0  0  0  0  0  0
+    7.1190   14.6210   25.8930 N   0  0  0  0  0  0  0  0  0  0  0  0
+    5.8050   14.8140   25.6500 C   0  0  0  0  0  0  0  0  0  0  0  0
+    5.4220   13.6990   24.7720 C   0  0  0  0  0  0  0  0  0  0  0  0
+    4.9170   15.9400   26.1920 C   0  0  0  0  0  0  0  0  0  0  0  0
+  1  2  4  0  0  0  0
+  1  6  4  0  0  0  0
+  1 10  1  0  0  0  0
+  2  3  4  0  0  0  0
+  2  7  4  0  0  0  0
+  3  4  4  0  0  0  0
+  3  9  4  0  0  0  0
+  4  5  4  0  0  0  0
+  5  6  4  0  0  0  0
+  7  8  4  0  0  0  0
+  8  9  4  0  0  0  0
+ 10 11  1  0  0  0  0
+ 10 12  1  0  0  0  0
+ 13 14  4  0  0  0  0
+ 13 17  4  0  0  0  0
+ 14 15  4  0  0  0  0
+ 15 16  4  0  0  0  0
+ 16 17  4  0  0  0  0
+ 16 18  1  0  0  0  0
+ 18 19  4  0  0  0  0
+ 18 22  4  0  0  0  0
+ 19 20  4  0  0  0  0
+ 20 21  4  0  0  0  0
+ 21 22  4  0  0  0  0
+ 21 23  1  0  0  0  0
+M  END
+$$$$

examples/5ou2_fragments.sdf

@@ -0,0 +1,56 @@
+5ou2_fragments
+  PyMOL2.5          3D                             0
+
+ 24 26  0  0  0  0  0  0  0  0999 V2000
+  135.6651  -15.3583    0.1325 N   0  0  0  0  0  0  0  0  0  0  0  0
+  134.8356  -14.4706   -0.4078 C   0  0  0  0  0  0  0  0  0  0  0  0
+  134.5969  -13.5549    0.5236 N   0  0  0  0  0  0  0  0  0  0  0  0
+  135.2672  -13.8787    1.6104 C   0  0  0  0  0  0  0  0  0  0  0  0
+  135.9361  -15.0095    1.3626 C   0  0  0  0  0  0  0  0  0  0  0  0
+  135.2407  -13.1072    2.8878 C   0  0  0  0  0  0  0  0  0  0  0  0
+  135.5339  -13.7328    4.0539 C   0  0  0  0  0  0  0  0  0  0  0  0
+  135.5239  -13.0695    5.2284 C   0  0  0  0  0  0  0  0  0  0  0  0
+  135.1995  -11.7489    5.2810 C   0  0  0  0  0  0  0  0  0  0  0  0
+  134.9023  -11.1173    4.1089 C   0  0  0  0  0  0  0  0  0  0  0  0
+  134.9113  -11.7774    2.9035 C   0  0  0  0  0  0  0  0  0  0  0  0
+  135.1362  -10.8138    6.9517 Br  0  0  0  0  0  0  0  0  0  0  0  0
+  126.8521  -19.0355    0.2522 N   0  0  0  0  0  0  0  0  0  0  0  0
+  126.0921  -18.0299   -0.2360 C   0  0  0  0  0  0  0  0  0  0  0  0
+  126.8721  -17.2548   -1.0322 N   0  0  0  0  0  0  0  0  0  0  0  0
+  128.1098  -17.7707   -1.0325 C   0  0  0  0  0  0  0  0  0  0  0  0
+  128.0889  -18.8815   -0.2256 C   0  0  0  0  0  0  0  0  0  0  0  0
+  129.3145  -17.2106   -1.7791 C   0  0  0  0  0  0  0  0  0  0  0  0
+  130.5850  -17.7185   -1.5264 C   0  0  0  0  0  0  0  0  0  0  0  0
+  131.6879  -17.2095   -2.1865 C   0  0  0  0  0  0  0  0  0  0  0  0
+  131.5211  -16.1844   -3.1052 C   0  0  0  0  0  0  0  0  0  0  0  0
+  130.2586  -15.6644   -3.3699 C   0  0  0  0  0  0  0  0  0  0  0  0
+  129.1548  -16.1795   -2.7058 C   0  0  0  0  0  0  0  0  0  0  0  0
+  133.0656  -15.5029   -4.0086 Br  0  0  0  0  0  0  0  0  0  0  0  0
+  1  2  4  0  0  0  0
+  2  3  4  0  0  0  0
+  3  4  4  0  0  0  0
+  4  6  1  0  0  0  0
+  1  5  4  0  0  0  0
+  4  5  4  0  0  0  0
+  6  7  4  0  0  0  0
+  6 11  4  0  0  0  0
+  7  8  4  0  0  0  0
+  8  9  4  0  0  0  0
+  9 10  4  0  0  0  0
+  9 12  1  0  0  0  0
+ 10 11  4  0  0  0  0
+ 13 14  4  0  0  0  0
+ 14 15  4  0  0  0  0
+ 15 16  4  0  0  0  0
+ 16 18  1  0  0  0  0
+ 13 17  4  0  0  0  0
+ 16 17  4  0  0  0  0
+ 18 19  4  0  0  0  0
+ 18 23  4  0  0  0  0
+ 19 20  4  0  0  0  0
+ 20 21  4  0  0  0  0
+ 21 22  4  0  0  0  0
+ 21 24  1  0  0  0  0
+ 22 23  4  0  0  0  0
+M  END
+$$$$

output.py

@@ -1,4 +1,4 @@
-INITIAL_RENDERING_TEMPLATE = """<!DOCTYPE html>
+FRAGMENTS_RENDERING_TEMPLATE = """<!DOCTYPE html>
 <html>
     <head>
         <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
@@ -26,7 +26,6 @@ INITIAL_RENDERING_TEMPLATE = """<!DOCTYPE html>
             let defaultStyle = {{ stick: {{ colorscheme: "greenCarbon" }} }};
             viewer.addModel(`{molecule}`, "{fmt}");
             viewer.getModel(0).setStyle(defaultStyle);
-            // document.cookie = document.cookie + "|selected_atoms:";
             
             viewer.getModel(0).setClickable(
                 {{}}, 
@@ -38,20 +37,16 @@ INITIAL_RENDERING_TEMPLATE = """<!DOCTYPE html>
                             {{"serial": _atom.serial, "model": 0}}, 
                             {{"sphere": {{"color": "magenta", "radius": 0.4}} }}
                         );
-                        // document.cookie = document.cookie + "atom_" + String(_atom.serial) + "-";
                         window.parent.postMessage({{
                             name: "atom_selection",
                             data: {{"atom": _atom.serial, "add": true}} 
-                            // data: JSON.stringify({{"add": _atom.serial}})
                         }}, "*");
                     }} else {{
                         delete _atom.isClicked;
                         _viewer.setStyle({{"serial": _atom.serial, "model": 0}}, defaultStyle);
-                        // document.cookie = document.cookie.replace("atom_" + String(_atom.serial) + "-", "");
                         window.parent.postMessage({{
                             name: "atom_selection",
                             data: {{"atom": _atom.serial, "add": false}} 
-                            // data: JSON.stringify({{"remove": _atom.serial}})
                         }}, "*");
                     }}
                     _viewer.render();
@@ -67,6 +62,112 @@ INITIAL_RENDERING_TEMPLATE = """<!DOCTYPE html>
 </html>
 """
 
+TARGET_RENDERING_TEMPLATE = """<!DOCTYPE html>
+<html>
+    <head>
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+        <script src="https://3Dmol.org/build/3Dmol.js"></script>
+        <style>
+            .mol-container {{
+                width: 600px;
+                height: 600px;
+                position: relative;
+            }}
+            .mol-container select{{
+                background-image:None;
+            }}
+        </style>
+    </head>
+
+    <body>
+    <div id="container" class="mol-container"></div>
+    <script>
+        $(document).ready(function() {{
+            let element = $("#container");
+            let config = {{ backgroundColor: "white" }};
+            let viewer = $3Dmol.createViewer(element, config);
+            let proteinStyle = {{ cartoon: {{ colorscheme: "ssPyMOL" }} }};
+            viewer.addModel(`{molecule}`, "{fmt}");
+            viewer.getModel(0).setStyle(proteinStyle);
+
+            viewer.zoomTo();
+            viewer.zoom(0.7);
+            viewer.render();
+        }});
+    </script>
+    </body>
+</html>
+"""
+
+FRAGMENTS_AND_TARGET_RENDERING_TEMPLATE = """<!DOCTYPE html>
+<html>
+    <head>
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+        <script src="https://3Dmol.org/build/3Dmol.js"></script>
+        <style>
+            .mol-container {{
+                width: 600px;
+                height: 600px;
+                position: relative;
+            }}
+            .mol-container select{{
+                background-image:None;
+            }}
+        </style>
+    </head>
+
+    <body>
+    <div id="container" class="mol-container"></div>
+    <script>
+        $(document).ready(function() {{
+            let element = $("#container");
+            let config = {{ backgroundColor: "white" }};
+            let viewer = $3Dmol.createViewer(element, config);
+            let defaultStyle = {{ stick: {{ colorscheme: "greenCarbon" }} }};
+            let proteinStyle = {{ cartoon: {{ colorscheme: "ssPyMOL" }} }};
+            
+            viewer.addModel(`{molecule}`, "{fmt}");
+            viewer.getModel(0).setStyle(defaultStyle);
+            viewer.getModel(0).setClickable(
+                {{}}, 
+                true,
+                function (_atom, _viewer, _event, _container) {{
+                    if (!_atom.isClicked) {{
+                        _atom.isClicked = true;
+                        _viewer.addStyle(
+                            {{"serial": _atom.serial, "model": 0}}, 
+                            {{"sphere": {{"color": "magenta", "radius": 0.4}} }}
+                        );
+                        window.parent.postMessage({{
+                            name: "atom_selection",
+                            data: {{"atom": _atom.serial, "add": true}} 
+                        }}, "*");
+                    }} else {{
+                        delete _atom.isClicked;
+                        _viewer.setStyle({{"serial": _atom.serial, "model": 0}}, defaultStyle);
+                        window.parent.postMessage({{
+                            name: "atom_selection",
+                            data: {{"atom": _atom.serial, "add": false}} 
+                        }}, "*");
+                    }}
+                    _viewer.render();
+                }}
+            );
+            
+            viewer.addModel(`{target}`, "{target_fmt}");
+            viewer.getModel(1).setStyle(proteinStyle);
+
+            viewer.zoomTo();
+            viewer.zoom(0.7);
+            viewer.render();
+        }});
+    </script>
+    </body>
+</html>
+"""
+
 
 SAMPLES_RENDERING_TEMPLATE = """<!DOCTYPE html>
 <html>
@@ -88,6 +189,7 @@ SAMPLES_RENDERING_TEMPLATE = """<!DOCTYPE html>
 
     <body>
     <div id="container" class="mol-container"></div>
+    <br>
     <button id="fragments">Input Fragments</button>
     <button id="molecule">Output Molecule</button> 
     <script>
@@ -120,6 +222,74 @@ SAMPLES_RENDERING_TEMPLATE = """<!DOCTYPE html>
 </html>
 """
 
+SAMPLES_WITH_TARGET_RENDERING_TEMPLATE = """<!DOCTYPE html>
+<html>
+    <head>
+        <meta http-equiv="content-type" content="text/html; charset=UTF-8" />
+        <script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/3.6.3/jquery.min.js" integrity="sha512-STof4xm1wgkfm7heWqFJVn58Hm3EtS31XFaagaa8VMReCXAkQnJZ+jEy8PCC/iT18dFy95WcExNHFTqLyp72eQ==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
+        <script src="https://3Dmol.org/build/3Dmol.js"></script>
+        <style>
+            .mol-container {{
+                width: 600px;
+                height: 600px;
+                position: relative;
+            }}
+            .mol-container select{{
+                background-image:None;
+            }}
+        </style>
+    </head>
+
+    <body>
+    <div id="container" class="mol-container"></div>
+    <br>
+    <button id="fragments">Input Fragments</button>
+    <button id="molecule">Output Molecule</button>
+    <button id="show-target">Show Target</button>
+    <button id="hide-target">Hide Target</button> 
+    <script>
+        let element = $("#container");
+        let config = {{ backgroundColor: "white" }};
+        let viewer = $3Dmol.createViewer( element, config );
+
+        $(document).ready(function() {{
+            viewer.addModel(`{fragments}`, "{fragments_fmt}")
+            viewer.getModel(0).setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
+            viewer.getModel(0).hide();
+            
+            viewer.addModel(`{molecule}`, "{molecule_fmt}")
+            viewer.getModel(1).setStyle({{ stick: {{ colorscheme:"greenCarbon" }} }})
+            
+            viewer.addModel(`{target}`, "{target_fmt}")
+            viewer.getModel(2).setStyle({{ cartoon: {{ colorscheme: "ssPyMOL" }} }})
+            
+            viewer.zoomTo();
+            viewer.zoom(0.7);
+            viewer.render();
+        }});
+        $("#fragments").click(function() {{
+            viewer.getModel(0).show();
+            viewer.getModel(1).hide();
+            viewer.render();
+        }});
+        $("#molecule").click(function() {{
+            viewer.getModel(1).show();
+            viewer.getModel(0).hide();
+            viewer.render();
+        }});
+        $("#show-target").click(function() {{
+            viewer.getModel(2).show();
+            viewer.render();
+        }});
+        $("#hide-target").click(function() {{
+            viewer.getModel(2).hide();
+            viewer.render();
+        }});
+    </script>
+    </body>
+</html>
+"""
+
 
 INVALID_FORMAT_MSG = """
 <!DOCTYPE html>
@@ -135,13 +305,18 @@ INVALID_FORMAT_MSG = """
 
     <body>
         <h3>Invalid file format: {extension}</h3>
-        Please upload the file in one of the following formats:
+        Allowed formats for the fragments file:
         <ul>
           <li>.pdb</li>
           <li>.sdf</li>
           <li>.mol</li>
           <li>.mol2</li>
         </ul>
+        
+        Allowed formats for the optional protein file:
+        <ul>
+          <li>.pdb</li>
+        </ul>
     </body>
 </html>
 """
@@ -190,7 +365,7 @@ STARTUP_JS = """
 """
 
 RETURN_SELECTION_JS = """
-(input_file, n_steps, n_atoms, samples, hidden) => {
+(input_file, input_protein_file, n_steps, n_atoms, samples, hidden) => {
     let selected = []
     for (const [atom, add] of Object.entries(window.selected_elements)) {
       if (add) {
@@ -203,6 +378,6 @@ RETURN_SELECTION_JS = """
       }
     }
     console.log("Finished parsing");
-    return [input_file, n_steps, n_atoms, samples, selected.join(",")];
+    return [input_file, input_protein_file, n_steps, n_atoms, samples, selected.join(",")];
 }
 """

src/datasets.py

@@ -101,15 +101,25 @@ class ZincDataset(Dataset):
 
 
 class MOADDataset(Dataset):
-    def __init__(self, data_path, prefix, device):
-        prefix, pocket_mode = prefix.split('.')
+    def __init__(self, data=None, data_path=None, prefix=None, device=None):
+        assert (data is not None) or all(x is not None for x in (data_path, prefix, device))
+        if data is not None:
+            self.data = data
+            return
+
+        if '.' in prefix:
+            prefix, pocket_mode = prefix.split('.')
+        else:
+            parts = prefix.split('_')
+            prefix = '_'.join(parts[:-1])
+            pocket_mode = parts[-1]
 
         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)
+            self.data = self.preprocess(data_path, prefix, pocket_mode, device)
             torch.save(self.data, dataset_path)
 
     def __len__(self):
@@ -264,7 +274,7 @@ def collate_with_fragment_edges(batch):
     out = {}
 
     # Filter out big molecules
-    batch = [data for data in batch if data['num_atoms'] <= 50]
+    # batch = [data for data in batch if data['num_atoms'] <= 50]
 
     for i, data in enumerate(batch):
         for key, value in data.items():

src/generation.py

@@ -1,24 +1,44 @@
+import numpy as np
 import os.path
 import subprocess
 import torch
 
+from Bio.PDB import PDBParser
+from src import const
 from src.visualizer import save_xyz_file
+from src.utils import FoundNaNException
+from src.datasets import get_one_hot
 
 N_SAMPLES = 5
 
 
-def generate_linkers(ddpm, data, sample_fn, name):
-    chain, node_mask = ddpm.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
+def generate_linkers(ddpm, data, sample_fn, name, with_pocket=False):
+    chain = node_mask = None
+    for i in range(5):
+        try:
+            chain, node_mask = ddpm.sample_chain(data, sample_fn=sample_fn, keep_frames=1)
+            break
+        except FoundNaNException:
+            continue
+
     print('Generated linker')
     x = chain[0][:, :, :ddpm.n_dims]
     h = chain[0][:, :, ddpm.n_dims:]
 
     # Put the molecule back to the initial orientation
-    pos_masked = data['positions'] * data['fragment_mask']
-    N = data['fragment_mask'].sum(1, keepdims=True)
+    if with_pocket:
+        com_mask = data['fragment_only_mask'] if ddpm.center_of_mass == 'fragments' else data['anchors']
+    else:
+        com_mask = data['fragment_mask'] if ddpm.center_of_mass == 'fragments' else data['anchors']
+
+    pos_masked = data['positions'] * com_mask
+    N = com_mask.sum(1, keepdims=True)
     mean = torch.sum(pos_masked, dim=1, keepdim=True) / N
     x = x + mean * node_mask
 
+    if with_pocket:
+        node_mask[torch.where(data['pocket_mask'])] = 0
+
     names = [f'output_{i + 1}_{name}' for i in range(N_SAMPLES)]
     save_xyz_file('results', h, x, node_mask, names=names, is_geom=True, suffix='')
     print('Saved XYZ files')
@@ -36,3 +56,62 @@ def try_to_convert_to_sdf(name):
             out_files.append(out_xyz)
 
     return out_files
+
+
+def get_pocket(mol, pdb_path):
+    struct = PDBParser().get_structure('', pdb_path)
+    residue_ids = []
+    atom_coords = []
+
+    for residue in struct.get_residues():
+        resid = residue.get_id()[1]
+        for atom in residue.get_atoms():
+            atom_coords.append(atom.get_coord())
+            residue_ids.append(resid)
+
+    residue_ids = np.array(residue_ids)
+    atom_coords = np.array(atom_coords)
+    mol_atom_coords = mol.GetConformer().GetPositions()
+
+    distances = np.linalg.norm(atom_coords[:, None, :] - mol_atom_coords[None, :, :], axis=-1)
+    contact_residues = np.unique(residue_ids[np.where(distances.min(1) <= 6)[0]])
+
+    pocket_coords_full = []
+    pocket_types_full = []
+
+    pocket_coords_bb = []
+    pocket_types_bb = []
+
+    for residue in struct.get_residues():
+        resid = residue.get_id()[1]
+        if resid not in contact_residues:
+            continue
+
+        for atom in residue.get_atoms():
+            atom_name = atom.get_name()
+            atom_type = atom.element.upper()
+            atom_coord = atom.get_coord()
+
+            pocket_coords_full.append(atom_coord.tolist())
+            pocket_types_full.append(atom_type)
+
+            if atom_name in {'N', 'CA', 'C', 'O'}:
+                pocket_coords_bb.append(atom_coord.tolist())
+                pocket_types_bb.append(atom_type)
+
+    pocket_pos = []
+    pocket_one_hot = []
+    pocket_charges = []
+    for coord, atom_type in zip(pocket_coords_full, pocket_types_full):
+        if atom_type not in const.GEOM_ATOM2IDX.keys():
+            continue
+
+        pocket_pos.append(coord)
+        pocket_one_hot.append(get_one_hot(atom_type, const.GEOM_ATOM2IDX))
+        pocket_charges.append(const.GEOM_CHARGES[atom_type])
+
+    pocket_pos = np.array(pocket_pos)
+    pocket_one_hot = np.array(pocket_one_hot)
+    pocket_charges = np.array(pocket_charges)
+
+    return pocket_pos, pocket_one_hot, pocket_charges

src/lightning.py

@@ -21,7 +21,6 @@ from pdb import set_trace
 
 
 def get_activation(activation):
-    print(activation)
     if activation == 'silu':
         return torch.nn.SiLU()
     else:
@@ -158,7 +157,7 @@ class DDPM(pl.LightningModule):
             context = fragment_mask
 
         # Add information about pocket to the context
-        if '.' in self.train_data_prefix:
+        if isinstance(self.train_dataset, MOADDataset):
             fragment_pocket_mask = fragment_mask
             fragment_only_mask = data['fragment_only_mask']
             pocket_only_mask = fragment_pocket_mask - fragment_only_mask
@@ -170,6 +169,8 @@ class DDPM(pl.LightningModule):
         # Removing COM of fragment from the atom coordinates
         if self.inpainting:
             center_of_mass_mask = node_mask
+        elif isinstance(self.train_dataset, MOADDataset) and self.center_of_mass == 'fragments':
+            center_of_mass_mask = data['fragment_only_mask']
         elif self.center_of_mass == 'fragments':
             center_of_mass_mask = fragment_mask
         elif self.center_of_mass == 'anchors':
@@ -423,9 +424,9 @@ class DDPM(pl.LightningModule):
             context = fragment_mask
 
         # Add information about pocket to the context
-        if '.' in self.train_data_prefix:
+        if isinstance(self.val_dataset, MOADDataset):
             fragment_pocket_mask = fragment_mask
-            fragment_only_mask = data['fragment_only_mask']
+            fragment_only_mask = template_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)
@@ -435,6 +436,8 @@ class DDPM(pl.LightningModule):
         # Removing COM of fragment from the atom coordinates
         if self.inpainting:
             center_of_mass_mask = node_mask
+        elif isinstance(self.val_dataset, MOADDataset) and self.center_of_mass == 'fragments':
+            center_of_mass_mask = template_data['fragment_only_mask']
         elif self.center_of_mass == 'fragments':
             center_of_mass_mask = fragment_mask
         elif self.center_of_mass == 'anchors':