kdock 0.0.2__py3-none-any.whl

kdock/af3/json.py

@@ -0,0 +1,282 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/af3/00_json.ipynb.
+
+# %% auto 0
+__all__ = ['dump_json', 'get_protein_json', 'read_json', 'get_protein_smiles_json', 'get_protein_ccdcode_json',
+           'assign_atom_names_from_graph', 'mol_to_ccd_text', 'sdf2ccd', 'get_protein_ccd_json', 'split_nfolder']
+
+# %% ../../nbs/af3/00_json.ipynb 2
+import re, shutil, json, pandas as pd, numpy as np
+from pathlib import Path
+
+from rdkit import Chem as rd_chem
+from rdkit.Chem import AllChem,rdmolfiles
+from rdkit import Chem
+
+from Bio.PDB import PDBParser
+
+# %% ../../nbs/af3/00_json.ipynb 4
+def dump_json(data, save_path):
+    "Save json data into a file"
+    with open(save_path,'w') as f: 
+        json.dump(data,f,indent=4)
+
+# %% ../../nbs/af3/00_json.ipynb 5
+def get_protein_json(name, # job name
+                     seq, # aa sequence
+                     save_path=None, # .json
+                     seeds=[1]
+                     ):
+    "Generate json of single protein sequence for input of docker command"
+    
+    json_data = {
+        "name": name,
+        "modelSeeds": seeds,
+        "sequences": [
+            {
+                "protein": {
+                    "id": "A",
+                    "sequence": seq,
+                }
+            },
+        ],
+        "bondedAtomPairs": [],
+        "dialect": "alphafold3",
+        "version": 3
+    }
+    if save_path:
+        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
+        dump_json(json_data,save_path)
+    return json_data
+
+# %% ../../nbs/af3/00_json.ipynb 9
+def read_json(file_path):
+    with open(file_path,'r') as f: 
+        data = json.load(f)
+    return data
+
+# %% ../../nbs/af3/00_json.ipynb 11
+def get_protein_smiles_json(smi_id:str, 
+                            SMILES:str, 
+                            protein_json, # json type
+                            save_path=None, # .json
+                            seeds=[1]
+                            ):
+    
+    "Get json for protein-ligand docking task"
+    raw_smiles = r"{}".format(SMILES) # JSON escaping, \ to \\
+    protein_index = next(i for i, item in enumerate(protein_json["sequences"]) if "protein" in item)
+    json_data = {
+        "name": smi_id,
+        "modelSeeds": seeds,
+        "sequences": [
+            {
+                "ligand": {
+                    "id": "L",
+                    "smiles": raw_smiles,
+                }
+            }, 
+            {
+                "protein": protein_json["sequences"][protein_index]["protein"]
+            },
+        ],
+        "bondedAtomPairs": [],
+        "dialect": "alphafold3",
+        "version": 2
+    }
+    if save_path:
+        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
+        dump_json(json_data,save_path)
+    return json_data
+
+# %% ../../nbs/af3/00_json.ipynb 18
+def get_protein_ccdcode_json(protein_json,  # dict with protein sequence
+                              ccd_code,      # str or list of str
+                              job_id: str,   # job/task ID
+                              save_path=None,  # optional output path
+                              seeds=[1]):      # optional random seeds
+    "Create AlphaFold3 docking JSON with CCD code(s)."
+    
+    # Normalize ccd_code to a list
+    if isinstance(ccd_code, str):
+        ccd_code = [ccd_code]
+    elif not isinstance(ccd_code, list):
+        raise TypeError("ccd_code must be a string or a list of strings.")
+
+    protein_index = next(i for i, item in enumerate(protein_json["sequences"]) if "protein" in item)
+
+    json_data = {
+        "name": job_id,
+        "modelSeeds": seeds,
+        "sequences": [
+            {
+                "ligand": {
+                    "id": "L",
+                    "ccdCodes": ccd_code
+                }
+            },
+            {
+                "protein": protein_json["sequences"][protein_index]["protein"]
+            },
+        ],
+        "dialect": "alphafold3",
+        "version": 3
+    }
+
+    if save_path:
+        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
+        dump_json(json_data, save_path)
+
+    return json_data
+
+# %% ../../nbs/af3/00_json.ipynb 22
+# Mapping bond types to mmCIF-compatible values
+_RDKIT_BOND_TYPE_TO_MMCIF = {
+    rd_chem.BondType.SINGLE: 'SING',
+    rd_chem.BondType.DOUBLE: 'DOUB',
+    rd_chem.BondType.TRIPLE: 'TRIP',
+    rd_chem.BondType.AROMATIC: 'AROM'
+}
+
+def assign_atom_names_from_graph(mol):
+    for i, atom in enumerate(mol.GetAtoms()):
+        atom.SetProp('atom_name', f"{atom.GetSymbol()}{i+1}")
+    return mol
+
+def mol_to_ccd_text(mol, component_id, pdbx_smiles=None, include_hydrogens=False):
+    mol = rd_chem.Mol(mol)
+    if include_hydrogens:
+        mol = rd_chem.AddHs(mol)
+    rd_chem.Kekulize(mol, clearAromaticFlags=True)
+
+    if mol.GetNumConformers() == 0:
+        raise ValueError('The molecule has no conformers')
+    conf = mol.GetConformer()
+    coords = conf.GetPositions()
+
+    mol = assign_atom_names_from_graph(mol)
+    atom_map = {atom.GetIdx(): atom.GetProp('atom_name') for atom in mol.GetAtoms()}
+
+    lines = [
+        f"data_{component_id}",
+        "#",
+        f"_chem_comp.id {component_id}",
+        f"_chem_comp.name '{component_id}'",
+        "_chem_comp.type non-polymer",
+        "_chem_comp.formula '?'",
+        "_chem_comp.mon_nstd_parent_comp_id ?",
+        "_chem_comp.pdbx_synonyms ?",
+        "_chem_comp.formula_weight '?'",
+    ]
+    if pdbx_smiles:
+        lines.append(f"_chem_comp.pdbx_smiles {pdbx_smiles}")
+    lines += [
+        "#",
+        "loop_",
+        "_chem_comp_atom.comp_id",
+        "_chem_comp_atom.atom_id",
+        "_chem_comp_atom.type_symbol",
+        "_chem_comp_atom.charge",
+        "_chem_comp_atom.pdbx_leaving_atom_flag",
+        "_chem_comp_atom.pdbx_model_Cartn_x_ideal",
+        "_chem_comp_atom.pdbx_model_Cartn_y_ideal",
+        "_chem_comp_atom.pdbx_model_Cartn_z_ideal"
+    ]
+
+    for i, atom in enumerate(mol.GetAtoms()):
+        if not include_hydrogens and atom.GetSymbol() == 'H':
+            continue
+        x, y, z = coords[i]
+        lines.append(f"{component_id} {atom_map[atom.GetIdx()]} {atom.GetSymbol()} {atom.GetFormalCharge()} N {x:.3f} {y:.3f} {z:.3f}")
+
+    lines += [
+        "#",
+        "loop_",
+        "_chem_comp_bond.atom_id_1",
+        "_chem_comp_bond.atom_id_2",
+        "_chem_comp_bond.value_order",
+        "_chem_comp_bond.pdbx_aromatic_flag"
+    ]
+
+    for bond in mol.GetBonds():
+        a1, a2 = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
+        if not include_hydrogens and (mol.GetAtomWithIdx(a1).GetSymbol() == 'H' or mol.GetAtomWithIdx(a2).GetSymbol() == 'H'):
+            continue
+        bond_type = _RDKIT_BOND_TYPE_TO_MMCIF[bond.GetBondType()]
+        aromatic_flag = 'Y' if bond.GetIsAromatic() else 'N'
+        lines.append(f"{atom_map[a1]} {atom_map[a2]} {bond_type} {aromatic_flag}")
+    lines.append("#")
+
+    return "\n".join(lines)
+
+# %% ../../nbs/af3/00_json.ipynb 23
+def sdf2ccd(sdf_path,
+            CCD_name='lig-1', # do not use '_'; use as less letter as possible, 'lig-any' leads to extra ligands
+            ):
+
+    "Convert the compound to the AF3 required CCD format"
+    supplier = Chem.SDMolSupplier(sdf_path)
+    mol = supplier[0]  # Get the first molecule
+    return mol_to_ccd_text(mol,CCD_name)
+
+# %% ../../nbs/af3/00_json.ipynb 26
+def get_protein_ccd_json(protein_json, # dict with protein sequence
+                         rec_residue_num:int, # 1-indexed, for bondedAtomPairs, e.g., ["A", 145, "SG"] 
+                         rec_atom_id:str, # for bondedAtomPairs, e.g., ["A", 145, "SG"] 
+                         lig_sdf_path, # ccd text
+                         lig_atom_id:str, # 0-indexed, for bondedAtomPairs, ["L", 1, "C04"]
+                         job_id:str, # str, job/task ID
+                         save_path=None,# optional output path
+                         seeds=[1], # optional random seeds
+                         ):               
+    "Create AlphaFold3 docking JSON with customized CCD ligand and bondedAtomPairs."
+
+    # get userCCD
+    userCCD=sdf2ccd(lig_sdf_path)
+    ccd_id = re.search(r"_chem_comp.id\s+([^\s#]+)", userCCD).group(1)
+    
+    protein_index = next(i for i, item in enumerate(protein_json["sequences"]) if "protein" in item)
+
+    json_data = {
+        "name": job_id,
+        "modelSeeds": seeds,
+        "sequences": [
+            {
+                "ligand": {
+                    "id": "L",
+                    "ccdCodes": [ccd_id]
+                }
+            },
+            {
+                "protein": protein_json["sequences"][protein_index]["protein"]
+            },
+        ],
+        "bondedAtomPairs": [[["A", rec_residue_num, rec_atom_id],["L", 1, lig_atom_id]]],
+        "userCCD": userCCD,
+        "dialect": "alphafold3",
+        "version": 3
+    }
+
+    if save_path:
+        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
+        dump_json(json_data, save_path)
+
+    return json_data
+
+# %% ../../nbs/af3/00_json.ipynb 30
+def split_nfolder(folder_dir, 
+                  n=4):
+    "Move json files from a folder into subfolders (folder_0, folder_1, ..., folder_N)."
+    
+    folder_dir = Path(folder_dir)
+
+    files = sorted(folder_dir.glob("*.json"))
+    # print(len(files))
+    subfolders = [folder_dir / f"folder_{i}" for i in range(n)]
+    for folder in subfolders:
+        folder.mkdir(exist_ok=True)
+
+    for idx, file in enumerate(files):
+        target_folder = subfolders[idx % n]
+        shutil.move(str(file), target_folder / file.name)
+
+    print(f"Distributed {len(files)} files into {n} folders.")

kdock/af3/protein_pairs.py

@@ -0,0 +1,95 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/af3/02_protein_pairs.ipynb.
+
+# %% auto 0
+__all__ = ['get_colabfold_cmd', 'copy_a3m', 'a3m_to_seq', 'get_protein_subjson', 'dump_json_folder', 'get_multi_protein_json',
+           'generate_pair_df']
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 4
+import os, json, shutil, pandas as pd
+from tqdm import tqdm
+from itertools import combinations
+from pathlib import Path
+from .json import *
+from .docker import *
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 9
+def get_colabfold_cmd(csv_path,project_name):
+    print('Run below in terminal:')
+    print(f'\n colabfold_batch {csv_path} msa_{project_name} --msa-only')
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 13
+def copy_a3m(a3m_dir: str, # Path to the source directory containing .a3m files.
+             dest_dir: str, # Path to the destination directory where files will be copied
+             ):
+    "Copies all .a3m files from the source directory to the destination directory."
+    
+    a3m_dir,dest_dir = Path(a3m_dir),Path(dest_dir)
+    dest_dir.mkdir(parents=True, exist_ok=True)
+
+    files = list(a3m_dir.glob('*.a3m'))
+
+    for file in tqdm(files, desc="Copying files", unit="file"):
+        shutil.copy(file, dest_dir / file.name)
+
+    print(f"Copied {len(files)} a3m files from {a3m_dir} to {dest_dir}")
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 17
+def a3m_to_seq(file_path:Path):
+    "Get protein sequence from a3m file"
+    return file_path.read_text().splitlines()[2] # protein sequence is located on line 2
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 19
+def get_protein_subjson(gene_name, a3m_dir=".",idx = 'A',run_template=True):
+    "Get subjson (protein part) with colabfold unpairedMSA .a3m path"
+    file_path = Path(a3m_dir)/f"{gene_name}.a3m"
+    protein_sequence = a3m_to_seq(file_path)
+    
+    json_data = {
+        'id': idx,
+        'sequence': protein_sequence, 
+        'modifications': [],
+        'unpairedMsaPath': str("/root"/file_path), # for docker path, ECD under af_input
+        'pairedMsa': '',
+        'templates': None if run_template else []
+    }
+
+    return json_data
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 22
+def dump_json_folder(json_data, folder):
+    "Save json under a folder"
+    file_path = Path(folder)/f"{json_data['name']}.json"
+    with open(file_path,'w') as f: json.dump(json_data,f,indent=4)
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 23
+def get_multi_protein_json(gene_list,a3m_dir,run_template=True,save_folder=None):
+    'Get json of multiple proteins, with unpaired MSA path indicated (from colabfold MSA)'
+    sequences = []
+    alphabets = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
+    for index, gene in enumerate(gene_list):
+        sub_json=get_protein_subjson(gene,a3m_dir,idx=alphabets[index],run_template=run_template)
+        sequences.append({'protein':sub_json})
+    name = '_'.join(gene_list)
+    json_data = {
+            "name": name,
+            "modelSeeds": [1],
+            "sequences": sequences,
+            "bondedAtomPairs": [],
+            "dialect": "alphafold3",
+            "version": 2
+        }
+    if save_folder:
+        dump_json_folder(json_data,save_folder)
+    return json_data
+
+# %% ../../nbs/af3/02_protein_pairs.ipynb 27
+def generate_pair_df(gene_list,self_pair=True):
+    "Unique pair genes in a gene list"
+    pairs = list(combinations(gene_list, 2))
+    pair_df = pd.DataFrame(pairs,columns=["Gene1", "Gene2"])
+    
+    if self_pair:
+        self_pair_df = pd.DataFrame({'Gene1':gene_list, 'Gene2':gene_list})
+        pair_df = pd.concat([pair_df,self_pair_df])
+
+    return pair_df.reset_index(drop=True)

kdock/core/data.py

@@ -0,0 +1,64 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/core/00_data.ipynb.
+
+# %% auto 0
+__all__ = ['BASE_URL', 'fetch_csv', 'Collins', 'Kras']
+
+# %% ../../nbs/core/00_data.ipynb 3
+import pandas as pd
+import requests
+from functools import lru_cache
+
+# %% ../../nbs/core/00_data.ipynb 7
+BASE_URL = "https://github.com/sky1ove/kdock/raw/main/"
+
+# %% ../../nbs/core/00_data.ipynb 8
+@lru_cache()
+def fetch_csv(url):
+    return pd.read_csv(url)
+
+# %% ../../nbs/core/00_data.ipynb 9
+class Collins:
+    "A class of loading compound datasets from Collins lab."
+
+    @staticmethod
+    def get_antibiotics_2k():
+        """
+        Antibiotics dataset of 50 µM 2,560 compounds screening in E. coli K12 BW25113.
+        2,335 unique compounds after deduplicated.
+        Table S1B from 2020 Cell: A Deep Learning Approach to Antibiotic Discovery.
+        """
+        return fetch_csv(BASE_URL + "dataset/antibiotics_2k.csv")
+
+    @staticmethod
+    def get_antibiotics_39k():
+        """
+        Antibiotics dataset of 50 µM 39,128 compounds screening in E. coli K12 BW25113.
+        Supplementary dataset EV1 from 2022 Molecular Systems Biology: Benchmarking AlphaFold-enabled molecular docking predictions for antibiotic discovery.
+        """
+        return fetch_csv(BASE_URL + "dataset/antibiotics_39k.csv")
+
+    @staticmethod
+    def get_antibiotics_enzyme():
+        """
+        Antibiotics enzymatic inhibition dataset of 100 µM 218 compounds and 12 essential proteins in E. coli K12 BW25113.
+        Flattened benchmark dataset/Supplementary EV4 from 2022 Molecular Systems Biology: Benchmarking AlphaFold-enabled molecular docking predictions for antibiotic discovery.
+        """
+        return fetch_csv(BASE_URL + "dataset/antibiotics_enzyme.csv")
+
+# %% ../../nbs/core/00_data.ipynb 19
+class Kras:
+    "A class of fetching various KRAS datasets."
+    @staticmethod
+    def get_mirati_g12d():
+        "Deduplicated G12D dataset from the mirati paper and patents."
+        return fetch_csv(BASE_URL + "dataset/KRASi_g12d_dedup.csv")
+
+    @staticmethod
+    def get_mirati_g12d_raw():
+        "Raw G12D dataset from the paper and patents without deduplication."
+        return fetch_csv(BASE_URL + "dataset/KRASi_g12d.csv")
+
+    @staticmethod
+    def get_seq():
+        "Protein sequence of human KRAS and its mutants G12D and G12C."
+        return fetch_csv(BASE_URL + "dataset/kras_seq.csv")