add dataset generation script

.gitignore

@@ -0,0 +1,6 @@
+# Ignore __pycache__ folders
+__pycache__/
+.idea/
+
+# Ignore .DS_Store files
+.DS_Store

README.md

@@ -1,11 +1,11 @@
 ---
-title: Pinder Inference Template
-emoji: 👁
+title: DockFormerPP
+emoji: ⚡
 colorFrom: indigo
 colorTo: indigo
 sdk: docker
 pinned: false
-license: mit
+license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

env_consts.py

@@ -3,7 +3,7 @@ import os
 TEST_INPUT_DIR = None
 TEST_OUTPUT_DIR = None
 THIS_FILE_DIR = os.path.dirname(os.path.abspath(__file__))
-CKPT_PATH = os.path.join(THIS_FILE_DIR, "resources", "77-182500_only_weights.ckpt")
+CKPT_PATH = os.path.join(THIS_FILE_DIR, "resources", "only_weights_102-240750.ckpt")
 RUN_CONFIG_PATH = os.path.join(THIS_FILE_DIR, "resources", "run_config.json")
 
 OUTPUT_PATH = os.path.join(THIS_FILE_DIR, "predicted_out.pdb")

prepare_pinder_dataset.py

@@ -0,0 +1,622 @@
+import json
+import os
+import shutil
+import random
+import sys
+from typing import List, Tuple, Optional
+
+import Bio.PDB
+import Bio.SeqUtils
+import pandas as pd
+import numpy as np
+
+
+
+OUTPUT_FOLDER = "/tmp/output"
+PINDER_ANNOTATIONS = "/tmp/index.parquet"
+GSUTIL_PATH = "/tmp/google-cloud-sdk/bin/gsutil"
+
+
+MAX_SYSTEMS_FOR_CLUSTER = 2
+MAX_LENGTH = 350
+MAX_TRIES_OF_METHOD = 5
+
+
+def do_robust_chain_object_renumber(chain: Bio.PDB.Chain.Chain, new_chain_id: str) -> Optional[Bio.PDB.Chain.Chain]:
+    all_residues = [res for res in chain.get_residues()
+                    if "CA" in res and Bio.SeqUtils.seq1(res.get_resname()) not in ("X", "", " ")]
+    if not all_residues:
+        return None
+
+    res_and_res_id = [(res, res.get_id()[1]) for res in all_residues]
+
+    min_res_id = min([i[1] for i in res_and_res_id])
+    if min_res_id < 1:
+        print("Negative res id", chain, min_res_id)
+        factor = -1 * min_res_id + 1
+        res_and_res_id = [(res, res_id + factor) for res, res_id in res_and_res_id]
+
+    res_and_res_id_no_collisions = []
+    for res, res_id in res_and_res_id[::-1]:
+        if res_and_res_id_no_collisions and res_and_res_id_no_collisions[-1][1] == res_id:
+            # there is a collision, usually an insertion residue
+            res_and_res_id_no_collisions = [(i, j + 1) for i, j in res_and_res_id_no_collisions]
+        res_and_res_id_no_collisions.append((res, res_id))
+
+    first_res_id = min([i[1] for i in res_and_res_id_no_collisions])
+    factor = 1 - first_res_id  # start from 1
+    new_chain = Bio.PDB.Chain.Chain(new_chain_id)
+
+    res_and_res_id_no_collisions.sort(key=lambda x: x[1])
+
+    for res, res_id in res_and_res_id_no_collisions:
+        chain.detach_child(res.id)
+        res.id = (" ", res_id + factor, " ")
+        new_chain.add(res)
+
+    return new_chain
+
+
+def robust_renumber_protein(pdb_path: str, output_path: str):
+    if pdb_path.endswith(".pdb"):
+        pdb_parser = Bio.PDB.PDBParser(QUIET=True)
+        pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
+    elif pdb_path.endswith(".cif"):
+        pdb_struct = Bio.PDB.MMCIFParser().get_structure("original_pdb", pdb_path)
+    else:
+        raise ValueError("Unknown file type", pdb_path)
+    assert len(list(pdb_struct)) == 1, "can't extract if more than one model"
+    model = next(iter(pdb_struct))
+    chains = list(model.get_chains())
+    new_model = Bio.PDB.Model.Model(0)
+    chain_ids = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
+    for chain, chain_id in zip(chains, chain_ids):
+        new_chain = do_robust_chain_object_renumber(chain, chain_id)
+        if new_chain is None:
+            continue
+        new_model.add(new_chain)
+    new_struct = Bio.PDB.Structure.Structure("renumbered_pdb")
+    new_struct.add(new_model)
+    io = Bio.PDB.PDBIO()
+    io.set_structure(new_struct)
+    io.save(output_path)
+
+
+def get_chain_object_to_seq(chain: Bio.PDB.Chain.Chain) -> str:
+    res_id_to_res = {res.get_id()[1]: res for res in chain.get_residues() if "CA" in res}
+
+    if len(res_id_to_res) == 0:
+        print("skipping empty chain", chain.get_id())
+        return ""
+    seq = ""
+    for i in range(1, max(res_id_to_res) + 1):
+        if i in res_id_to_res:
+            seq += Bio.SeqUtils.seq1(res_id_to_res[i].get_resname())
+        else:
+            seq += "X"
+    return seq
+
+
+def get_sequence_from_pdb(pdb_path: str) -> Tuple[str, List[int]]:
+    pdb_parser = Bio.PDB.PDBParser(QUIET=True)
+    pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
+    # chain_to_seq = {chain.id: get_chain_object_to_seq(chain) for chain in pdb_struct.get_chains()}
+    all_chain_seqs = [get_chain_object_to_seq(chain) for chain in pdb_struct.get_chains()]
+    chain_lengths = [len(seq) for seq in all_chain_seqs]
+    return ("X" * 20).join(all_chain_seqs), chain_lengths
+
+
+from Bio import PDB
+from Bio import pairwise2
+
+
+def extract_sequence(chain):
+    seq = ''
+    residues = []
+    for res in chain.get_residues():
+        seq_res = Bio.SeqUtils.seq1(res.get_resname())
+        if seq_res in ('X', "", " "):
+            continue
+        seq += seq_res
+        residues.append(res)
+    return seq, residues
+
+
+def map_residues(alignment, residues_gt, residues_pred):
+    idx_gt = 0
+    idx_pred = 0
+    mapping = []
+    for i in range(len(alignment.seqA)):
+        aa_gt = alignment.seqA[i]
+        aa_pred = alignment.seqB[i]
+        res_gt = None
+        res_pred = None
+        if aa_gt != '-':
+            res_gt = residues_gt[idx_gt]
+            idx_gt += 1
+        if aa_pred != '-':
+            res_pred = residues_pred[idx_pred]
+            idx_pred += 1
+        if res_gt and res_pred:
+            mapping.append((res_gt, res_pred))
+    return mapping
+
+
+class ResidueSelect(PDB.Select):
+    def __init__(self, residues_to_select):
+        self.residues_to_select = set(residues_to_select)
+
+    def accept_residue(self, residue):
+        return residue in self.residues_to_select
+
+
+def count_gapped_single_aa(alignment):
+    count_non_gap = 0
+    count_fully_gapped = 0
+    for i in range(1, len(alignment.seqA) - 1):
+        if alignment.seqA[i] != '-':
+            count_non_gap += 1
+            if alignment.seqA[i - 1] == '-' and alignment.seqA[i + 1] == '-':
+                count_fully_gapped += 1
+    top_ratio = count_fully_gapped / count_non_gap
+
+    count_non_gap = 0
+    count_fully_gapped = 0
+    for i in range(1, len(alignment.seqB) - 1):
+        if alignment.seqA[i] != '-':
+            count_non_gap += 1
+            if alignment.seqA[i - 1] == '-' and alignment.seqA[i + 1] == '-':
+                count_fully_gapped += 1
+
+    if count_fully_gapped / count_non_gap > top_ratio:
+        top_ratio = count_fully_gapped / count_non_gap
+
+    return top_ratio
+
+
+def copy_residue_numbering(gt_pdb_path, input_pdb_path):
+    parser = PDB.PDBParser(QUIET=True)
+    gt_structure = parser.get_structure('gt', gt_pdb_path)
+    input_structure = parser.get_structure('input', input_pdb_path)
+
+    for res in list(input_structure.get_residues()):
+        res.id = (' ', res.get_id()[1] + 10000, ' ')
+
+    for gt_res, input_res in zip(gt_structure.get_residues(), input_structure.get_residues()):
+        input_res.id = gt_res.id
+
+    io = PDB.PDBIO()
+    io.set_structure(input_structure)
+    io.save(input_pdb_path)
+
+
+def align_gt_and_input(gt_pdb_path, input_pdb_path, output_gt_path, output_input_path):
+    # print("aligning", gt_pdb_path, input_pdb_path, output_gt_path, output_input_path)
+    parser = PDB.PDBParser(QUIET=True)
+    gt_structure = parser.get_structure('gt', gt_pdb_path)
+    pred_structure = parser.get_structure('pred', input_pdb_path)
+    matched_residues_gt = []
+    matched_residues_pred = []
+
+    total_gt_size = len([res for res in gt_structure.get_residues() if "CA" in res])
+
+    used_chain_pred = []
+    total_mapping_size = 0
+    for chain_gt in gt_structure.get_chains():
+        seq_gt, residues_gt = extract_sequence(chain_gt)
+        best_alignment = None
+        best_chain_pred = None
+        best_score = -1
+        best_residues_pred = None
+        # Find the best matching chain in pred
+        for chain_pred in pred_structure.get_chains():
+            # print("checking", chain_pred.get_id(), chain_gt.get_id())
+            if chain_pred in used_chain_pred:
+                continue
+            seq_pred, residues_pred = extract_sequence(chain_pred)
+            # print(seq_gt)
+            # print(seq_pred)
+            # alignments = pairwise2.align.globalxx(seq_gt, seq_pred, one_alignment_only=True)
+            alignments = pairwise2.align.globalms(seq_gt, seq_pred, 2, -10000, -1, 0, one_alignment_only=True)
+            if not alignments:
+                continue
+            # print("checking2", chain_pred.get_id(), chain_gt.get_id())
+
+            alignment = alignments[0]
+            score = alignment.score
+            if score > best_score:
+                best_score = score
+                best_alignment = alignment
+                best_chain_pred = chain_pred
+                best_residues_pred = residues_pred
+        if best_alignment and count_gapped_single_aa(best_alignment) < 0.2:
+            mapping = map_residues(best_alignment, residues_gt, best_residues_pred)
+            total_mapping_size += len(mapping)
+            used_chain_pred.append(best_chain_pred)
+            for res_gt, res_pred in mapping:
+                matched_residues_gt.append(res_gt)
+                matched_residues_pred.append(res_pred)
+        else:
+            print(f"No matching chain found for chain {chain_gt.get_id()}")
+    assert total_mapping_size / total_gt_size > 0.8, \
+        f"Mapping size too low ({total_mapping_size}/{total_gt_size}), skipping"
+    print(f"Total mapping size: {total_mapping_size}")
+
+    # Write new PDB files with only matched residues
+    io = PDB.PDBIO()
+    io.set_structure(gt_structure)
+    io.save(output_gt_path, ResidueSelect(matched_residues_gt))
+    io = PDB.PDBIO()
+    io.set_structure(pred_structure)
+    io.save(output_input_path, ResidueSelect(matched_residues_pred))
+
+    copy_residue_numbering(output_gt_path, output_input_path)
+
+
+def validate_matching_input_gt(gt_pdb_path, input_pdb_path):
+    gt_residues = [res for res in PDB.PDBParser().get_structure('gt', gt_pdb_path).get_residues()]
+    input_residues = [res for res in PDB.PDBParser().get_structure('input', input_pdb_path).get_residues()]
+
+    if len(gt_residues) != len(input_residues):
+        print(f"Residue count mismatch: {len(gt_residues)} vs {len(input_residues)}")
+        return -1
+
+    for res_gt, res_input in zip(gt_residues, input_residues):
+        if res_gt.get_resname() != res_input.get_resname():
+            print(f"Residue name mismatch: {res_gt.get_resname()} vs {res_input.get_resname()}")
+            return -1
+    return len(input_residues)
+
+
+def download_pdb(pdb_name, output_folder):
+    output_path = os.path.join(output_folder, pdb_name)
+    if os.path.exists(output_path):
+        return output_path
+    print("downloading", pdb_name)
+    os.system(f'{GSUTIL_PATH} -m -q cp "gs://pinder/2024-02/pdbs/{pdb_name}" {output_path}')
+    return output_path
+
+
+INTERFACE_MIN_ATOM_DIST = 5
+
+
+def get_filtered_res(gt_r_res, gt_l_res, max_length: int):
+    gt_r_ca = np.array([res["CA"].coord for res in gt_r_res])
+    gt_l_ca = np.array([res["CA"].coord for res in gt_l_res])
+
+    if len(gt_r_res) + len(gt_l_res) < max_length:
+        # continue without cropping
+        print("no cropping needed", len(gt_r_res), len(gt_l_res))
+        return gt_r_res, gt_l_res
+
+    # close_residues = np.argwhere(scipy.spatial.distance.cdist(gt_r_ca, gt_l_ca) < INTERFACE_MIN_ATOM_DIST)
+    # gt_r_interface, gt_l_interface = set(), set()
+    # for i, j in close_residues:
+    #     gt_r_interface.add(gt_r_res[i].id[1])
+    #     gt_l_interface.add(gt_l_res[j].id[1])
+
+    inter_dists = gt_r_ca[:, np.newaxis, :] - gt_l_ca[np.newaxis, :, :]
+    inter_dists = np.sqrt((inter_dists ** 2).sum(-1))
+    min_inter_dist_per_gt_l_res = inter_dists.min(axis=0)
+    min_inter_dist_per_gt_r_res = inter_dists.min(axis=1)
+
+    assert min_inter_dist_per_gt_l_res.shape[0] == len(gt_l_res)
+    assert min_inter_dist_per_gt_r_res.shape[0] == len(gt_r_res)
+
+    min_r_res, max_r_res = min(min_inter_dist_per_gt_r_res), max(min_inter_dist_per_gt_r_res)
+    min_l_res, max_l_res = min(min_inter_dist_per_gt_l_res), max(min_inter_dist_per_gt_l_res)
+
+    r_pocket = [res for res in gt_r_res if min_r_res <= res.id[1] <= max_r_res]
+    l_pocket = [res for res in gt_l_res if min_l_res <= res.id[1] <= max_l_res]
+
+    if len(r_pocket) + len(l_pocket) < max_length:
+        # add extra residues to both chains to get a total of max_length
+        res_r_before = [res for res in gt_r_res if res.id[1] < min_r_res]
+        res_r_after = [res for res in gt_r_res if res.id[1] > max_r_res]
+        res_l_before = [res for res in gt_l_res if res.id[1] < min_l_res]
+        res_l_after = [res for res in gt_l_res if res.id[1] > max_l_res]
+
+        extra_to_add = max_length - len(r_pocket) - len(l_pocket)
+
+        actions = []
+        if len(res_r_before) > 0:
+            actions.append("add_r_before")
+        if len(res_r_after) > 0:
+            actions.append("add_r_after")
+        if len(res_l_before) > 0:
+            actions.append("add_l_before")
+        if len(res_l_after) > 0:
+            actions.append("add_l_after")
+        while extra_to_add > 0 and actions:
+            action = random.choice(actions)
+
+            if action == "add_r_before":
+                r_pocket.insert(0, res_r_before.pop())
+                if not len(res_r_before):
+                    actions.remove("add_r_before")
+            elif action == "add_r_after":
+                r_pocket.append(res_r_after.pop())
+                if not len(res_r_after):
+                    actions.remove("add_r_after")
+            elif action == "add_l_before":
+                l_pocket.insert(0, res_l_before.pop())
+                if not len(res_l_before):
+                    actions.remove("add_l_before")
+            elif action == "add_l_after":
+                l_pocket.append(res_l_after.pop())
+                if not len(res_l_after):
+                    actions.remove("add_l_after")
+            extra_to_add -= 1
+        print("Extended pocket sizes", len(r_pocket), len(l_pocket), "extra_to_add", extra_to_add)
+        return r_pocket, l_pocket
+
+    print("cropping simply")
+    # remove residues that are farthest from the interface
+    res_and_dist_r = [(res, min_inter_dist_per_gt_r_res[res_idx]) for res_idx, res in enumerate(gt_r_res)]
+    res_and_dist_l = [(res, min_inter_dist_per_gt_l_res[res_idx]) for res_idx, res in enumerate(gt_l_res)]
+
+    res_and_dist_r = [(res, dist) for res, dist in res_and_dist_r if res in r_pocket]
+    res_and_dist_l = [(res, dist) for res, dist in res_and_dist_l if res in l_pocket]
+
+    res_and_dist_r = sorted(res_and_dist_r, key=lambda x: x[1], reverse=True)
+    res_and_dist_l = sorted(res_and_dist_l, key=lambda x: x[1], reverse=True)
+
+    while len(res_and_dist_r) + len(res_and_dist_l) > max_length:
+        if res_and_dist_r[0][1] > res_and_dist_l[0][1]:
+            res_and_dist_r.pop(0)
+        else:
+            res_and_dist_l.pop(0)
+
+    return [res for res, _ in res_and_dist_r], [res for res, _ in res_and_dist_l]
+
+
+def prepare_holo(row, tmp_dir_path, max_length: int):
+    tmp_gt_r_pdb = os.path.join(tmp_dir_path, f"tmp_{row.id}_gt_r.pdb")
+    tmp_gt_l_pdb = os.path.join(tmp_dir_path, f"tmp_{row.id}_gt_l.pdb")
+
+    if os.path.exists(tmp_gt_r_pdb) and os.path.exists(tmp_gt_l_pdb):
+        return tmp_gt_r_pdb, tmp_gt_l_pdb
+
+    holo_r_pdb = download_pdb(row.holo_R_pdb, tmp_dir_path)
+    holo_l_pdb = download_pdb(row.holo_L_pdb, tmp_dir_path)
+
+    # make gt and apo that matches
+    robust_renumber_protein(holo_r_pdb, tmp_gt_r_pdb)
+    robust_renumber_protein(holo_l_pdb, tmp_gt_l_pdb)
+
+    parser = PDB.PDBParser(QUIET=True)
+    gt_r_prot = parser.get_structure('r', tmp_gt_r_pdb)
+    gt_l_prot = parser.get_structure('l', tmp_gt_l_pdb)
+
+    assert len(list(gt_r_prot.get_chains())) == 1, "can't extract if more than one chain"
+    assert len(list(gt_l_prot.get_chains())) == 1, "can't extract if more than one chain"
+
+    gt_r_res = [res for res in gt_r_prot.get_residues() if "CA" in res]
+    gt_l_res = [res for res in gt_l_prot.get_residues() if "CA" in res]
+
+    to_keep_r, to_keep_l = get_filtered_res(gt_r_res, gt_l_res, max_length)
+
+    io = PDB.PDBIO()
+    io.set_structure(gt_r_prot)
+    io.save(tmp_gt_r_pdb, ResidueSelect(to_keep_r))
+    io = PDB.PDBIO()
+    io.set_structure(gt_l_prot)
+    io.save(tmp_gt_l_pdb, ResidueSelect(to_keep_l))
+
+    return tmp_gt_r_pdb, tmp_gt_l_pdb
+
+
+def generate_input_pdbs(tmp_input_r_pdb, tmp_input_l_pdb, tmp_gt_r_pdb, tmp_gt_l_pdb,
+                        input_r_output_pdb, input_l_output_pdb, gt_r_output_pdb, gt_l_output_pdb):
+    # print("preparing input pdbs", gt_r_output_pdb)
+    if not os.path.exists(tmp_input_r_pdb) or not os.path.exists(tmp_input_l_pdb):
+        raise False
+
+    try:
+        align_gt_and_input(tmp_gt_r_pdb, tmp_input_r_pdb, gt_r_output_pdb, input_r_output_pdb)
+        protein_size_r = validate_matching_input_gt(gt_r_output_pdb, input_r_output_pdb)
+        assert protein_size_r > -1, "Failed to validate matching input and gt"
+
+        align_gt_and_input(tmp_gt_l_pdb, tmp_input_l_pdb, gt_l_output_pdb, input_l_output_pdb)
+        protein_size_l = validate_matching_input_gt(gt_l_output_pdb, input_l_output_pdb)
+        assert protein_size_l > -1, "Failed to validate matching input and gt"
+    except Exception as e:
+        print("Failed to align", e)
+        if os.path.exists(gt_r_output_pdb):
+            os.remove(gt_r_output_pdb)
+        if os.path.exists(gt_l_output_pdb):
+            os.remove(gt_l_output_pdb)
+        if os.path.exists(input_r_output_pdb):
+            os.remove(input_r_output_pdb)
+        if os.path.exists(input_l_output_pdb):
+            os.remove(input_l_output_pdb)
+        return False
+
+    return True
+
+
+def _get_rel_path(abs_path):
+    return os.path.join(os.path.basename(os.path.dirname(abs_path)), os.path.basename(abs_path))
+
+
+def main(start_ind: Optional[int] = None, end_ind: Optional[int] = None):
+    print("running with", start_ind, end_ind)
+
+    os.makedirs(OUTPUT_FOLDER, exist_ok=True)
+    output_models_folder = os.path.join(OUTPUT_FOLDER, "pinder_models")
+    output_train_jsons_folder = os.path.join(OUTPUT_FOLDER, "pinder_jsons_train")
+    output_val_jsons_folder = os.path.join(OUTPUT_FOLDER, "pinder_jsons_val")
+    output_test_jsons_folder = os.path.join(OUTPUT_FOLDER, "pinder_jsons_test")
+    output_info = os.path.join(OUTPUT_FOLDER, "pinder_generation_info.csv")
+
+    os.makedirs(output_models_folder, exist_ok=True)
+    os.makedirs(output_train_jsons_folder, exist_ok=True)
+    os.makedirs(output_val_jsons_folder, exist_ok=True)
+    os.makedirs(output_test_jsons_folder, exist_ok=True)
+
+    split_to_folder = {
+        "train": output_train_jsons_folder,
+        "val": output_val_jsons_folder,
+        "test": output_test_jsons_folder
+    }
+
+    # output_info_file = open(output_info, "a+")
+    systems = pd.read_parquet(PINDER_ANNOTATIONS)
+    systems = systems[systems.split.isin(['train', 'val', 'test'])]
+
+    cluster_ids = systems["cluster_id"].value_counts()
+    cluster_ids = cluster_ids[cluster_ids >= 1]
+    print("There are", len(cluster_ids), "clusters")
+
+    # clusters_with_data = 0
+    # for cluster_id in cluster_ids.index:
+    #     cluster_systems = systems[systems["cluster_id"] == cluster_id]
+    #     with_apo = cluster_systems[cluster_systems.apo_R & cluster_systems.apo_L]
+    #     if len(with_apo) > 0:
+    #         print("Cluster", cluster_id, "has", len(with_apo), "systems with apo")
+    #         clusters_with_data += 1
+    #         continue
+    #     with_pred = cluster_systems[cluster_systems.predicted_R & cluster_systems.predicted_L]
+    #     if len(with_pred) > 0:
+    #         print("Cluster", cluster_id, "has", len(with_pred), "systems with pred")
+    #         clusters_with_data += 1
+    #         continue
+    # print("There are", clusters_with_data, "clusters with data out of", len(cluster_ids))
+
+    for cluster_ind, cluster_id in enumerate(sorted(cluster_ids.index)):
+        if (start_ind is not None and cluster_ind < start_ind) or (end_ind is not None and cluster_ind >= end_ind):
+            continue
+        # if cluster_id != "cluster_10004_p":
+        #     continue
+        tmp_dir_path = os.path.join(OUTPUT_FOLDER, "tmp_" + cluster_id)
+        os.makedirs(tmp_dir_path, exist_ok=True)
+        system_id_to_method = {}
+
+        cluster_systems = systems[systems["cluster_id"] == cluster_id]
+        print("--- Starting cluster", cluster_ind, cluster_id, "size", cluster_systems.shape)
+
+        with_apo = cluster_systems[cluster_systems.apo_R & cluster_systems.apo_L]
+        print("*** APO *** Cluster", cluster_id, "has", len(with_apo), "systems with apo")
+        for try_counter, row in enumerate(with_apo.itertuples()):
+            if row.split not in ("test", "val") \
+                    and (try_counter >= MAX_TRIES_OF_METHOD or len(system_id_to_method) >= MAX_SYSTEMS_FOR_CLUSTER):
+                continue
+            print("-- Trying to prepare apo", row.id, row.split)
+            try:
+                tmp_gt_r_pdb, tmp_gt_l_pdb = prepare_holo(row, tmp_dir_path, MAX_LENGTH)
+
+                gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
+                gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")
+
+                input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
+                input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")
+
+                input_r_pdb_path = download_pdb(row.apo_R_pdb, tmp_dir_path)
+                input_l_pdb_path = download_pdb(row.apo_L_pdb, tmp_dir_path)
+
+                if generate_input_pdbs(input_r_pdb_path, input_l_pdb_path, tmp_gt_r_pdb, tmp_gt_l_pdb,
+                                       input_r_output_path, input_l_output_path, gt_r_output_path, gt_l_output_path):
+                    system_id_to_method[row.id] = "apo"
+
+            except Exception as e:
+                print("Failed to prepare apo", row.id, e)
+                continue
+
+        with_pred = cluster_systems[cluster_systems.predicted_R & cluster_systems.predicted_L]
+        print("*** Pred *** Cluster", cluster_id, "has", len(with_pred), "systems with pred")
+        for try_counter, row in enumerate(with_pred.itertuples()):
+            if row.id in system_id_to_method:
+                continue
+            if row.split not in ("test", "val") \
+                    and (try_counter >= MAX_TRIES_OF_METHOD or len(system_id_to_method) >= MAX_SYSTEMS_FOR_CLUSTER):
+                continue
+            print("-- Trying to prepare pred", row.id, row.split)
+            try:
+                tmp_gt_r_pdb, tmp_gt_l_pdb = prepare_holo(row, tmp_dir_path, MAX_LENGTH)
+
+                gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
+                gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")
+
+                input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
+                input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")
+
+                input_r_pdb_path = download_pdb(row.predicted_R_pdb, tmp_dir_path)
+                input_l_pdb_path = download_pdb(row.predicted_L_pdb, tmp_dir_path)
+
+                if generate_input_pdbs(input_r_pdb_path, input_l_pdb_path, tmp_gt_r_pdb, tmp_gt_l_pdb,
+                                       input_r_output_path, input_l_output_path, gt_r_output_path, gt_l_output_path):
+                    system_id_to_method[row.id] = "pred"
+
+            except Exception as e:
+                print("Failed to prepare pred", row.id, e)
+
+        # default - use gt
+        print("*** GT *** ")
+        for row in cluster_systems.itertuples():
+            if row.id in system_id_to_method:
+                continue
+            if row.split not in ("test", "val") and len(system_id_to_method) >= MAX_SYSTEMS_FOR_CLUSTER:
+                continue
+            try:
+                tmp_gt_r_pdb, tmp_gt_l_pdb = prepare_holo(row, tmp_dir_path, MAX_LENGTH)
+
+                gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
+                gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")
+
+                input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
+                input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")
+
+                shutil.copyfile(tmp_gt_r_pdb, gt_r_output_path)
+                shutil.copyfile(tmp_gt_r_pdb, input_r_output_path)
+
+                shutil.copyfile(tmp_gt_l_pdb, gt_l_output_path)
+                shutil.copyfile(tmp_gt_l_pdb, input_l_output_path)
+
+                system_id_to_method[row.id] = "gt"
+
+            except Exception as e:
+                print("Failed to prepare gt", row.id, e)
+
+        # save jsons
+        for row in cluster_systems.itertuples():
+            if row.id not in system_id_to_method:
+                continue
+
+            output_json_path = os.path.join(split_to_folder[row.split], f"{row.id}.json")
+
+            gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
+            gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")
+
+            input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
+            input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")
+
+            protein_r_seq_len = validate_matching_input_gt(gt_r_output_path, input_r_output_path)
+            protein_l_seq_len = validate_matching_input_gt(gt_l_output_path, input_l_output_path)
+
+            json_data = {
+                "input_r_structure": _get_rel_path(input_r_output_path),
+                "input_l_structure": _get_rel_path(input_l_output_path),
+                "gt_r_structure": _get_rel_path(gt_r_output_path),
+                "gt_l_structure": _get_rel_path(gt_l_output_path),
+                "resolution": 1.0,
+                "protein_r_seq_len": protein_r_seq_len,
+                "protein_l_seq_len": protein_l_seq_len,
+                "uniprot_r": row.uniprot_R,
+                "uniprot_l": row.uniprot_L,
+                "cluster": row.cluster_id,
+                "input_protein_source": system_id_to_method[row.id],
+                "pdb_id": row.id,
+            }
+            open(output_json_path, "w").write(json.dumps(json_data, indent=4))
+
+            print("******* saved", row.id, system_id_to_method[row.id], flush=True)
+        shutil.rmtree(tmp_dir_path)
+
+    print("total systems", len(systems))
+
+
+if __name__ == '__main__':
+    if len(sys.argv) == 3:
+        main(int(sys.argv[1]), int(sys.argv[2]))
+    else:
+        main()

resources/{77-182500_only_weights.ckpt → only_weights_102-240750.ckpt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3412e505a8fce15011c589bb008f0f9c92d0ce9703bc75d70a085db733c2b16d
-size 52579155
+oid sha256:3e32cd27b63f684ed813e65db4369f5af28b1fceb3c81df66bdd4952f6b78853
+size 52579856