gemmi-protools 1.1.0__py3-none-any.whl

gemmi_protools/tools/__init__.py

@@ -0,0 +1,3 @@
+"""
+@Author: Luo Jiejian
+"""

gemmi_protools/tools/align.py

@@ -0,0 +1,336 @@
+"""
+@Author: Luo Jiejian
+"""
+
+import os
+import re
+import shutil
+import string
+import subprocess
+import tempfile
+from typing import Literal, Optional, Dict, Any, List
+
+import numpy as np
+from Bio.Align import PairwiseAligner, substitution_matrices
+from Bio.PDB import Superimposer
+from gemmi_protools import StructureParser
+from gemmi_protools.io.convert import gemmi2bio, bio2gemmi
+
+
+def check_sequence(seq: str):
+    """
+    Remove space, star at the end, and \n, upper the letters
+    Check sequence is valid or not
+
+    :param seq:str
+    :return:
+
+    """
+    seq_clean = re.sub(pattern=r" |\*|-", repl='', string=seq.upper().strip())
+    if len(seq_clean) == 0:
+        raise ValueError("Sequence is empty")
+
+    s = re.sub(pattern=r"[A-Z]", repl="", string=seq_clean)
+    if len(s) > 0:
+        raise ValueError("Sequence has Non-alphabetic characters: %s" % str(set(s)))
+
+    return seq_clean
+
+
+def align_sequences(seq1: str,
+                    seq2: str,
+                    seq_type: Literal["dna", "rna", "protein"] = "protein",
+                    mode: Literal["global", "local"] = "local",
+                    substitution_matrix: Optional[str] = None,
+                    open_gap_score: Optional[float] = None,
+                    extend_gap_score: Optional[float] = None,
+                    ):
+    """
+    [To Do]: when one insertion greater than 52, raise mapping error
+    """
+    default_params = {
+        "dna": {
+            "matrix": "NUC.4.4",
+            "open_gap_score": -10.0,
+            "extend_gap_score": -0.5,
+            "mode": "global"
+        },
+        "rna": {
+            "matrix": "NUC.4.4",
+            "open_gap_score": -10.0,
+            "extend_gap_score": -0.5,
+            "mode": "global"
+        },
+        "protein": {
+            "matrix": "BLOSUM62",
+            "open_gap_score": -11.0,
+            "extend_gap_score": -1.0,
+            "mode": "global"
+        }
+
+    }
+
+    available_matrices = {
+        "dna": ["NUC.4.4"],
+        "rna": ["NUC.4.4"],
+        "protein": ["BLOSUM45", "BLOSUM50", "BLOSUM62",
+                    "BLOSUM80", "BLOSUM90",
+                    "PAM30", "PAM70", "PAM250"]
+    }
+
+    seq1 = check_sequence(seq1)
+    seq2 = check_sequence(seq2)
+
+    params = default_params[seq_type].copy()
+    a_mats = available_matrices[seq_type]
+
+    if substitution_matrix is not None:
+        if substitution_matrix not in a_mats:
+            raise ValueError("substitution matrix `%s` not support for %s" % (substitution_matrix, seq_type))
+        else:
+            params["matrix"] = substitution_matrix
+
+    if open_gap_score is not None:
+        params["open_gap_score"] = open_gap_score
+    if extend_gap_score is not None:
+        params["extend_gap_score"] = extend_gap_score
+
+    params["mode"] = mode
+    # Finish parameters checking and setting
+    aligner = PairwiseAligner()
+    aligner.mode = params["mode"]
+    aligner.substitution_matrix = substitution_matrices.load(params["matrix"])
+    aligner.open_gap_score = params["open_gap_score"]
+    aligner.extend_gap_score = params["extend_gap_score"]
+
+    best_alignment = aligner.align(seq1, seq2)[0]
+
+    aligned_seq1, aligned_seq2 = best_alignment
+
+    # start from 1
+    aa_mapper = dict()
+    i = 0
+    j = 0
+
+    ins_letters = string.ascii_uppercase + string.ascii_lowercase
+    k = 0
+
+    for aa1, aa2 in zip(aligned_seq1, aligned_seq2):
+        if aa1 != "-":
+            i += 1
+        if aa2 != "-":
+            j += 1
+            # reset k
+            if k > 0:
+                k = 0
+
+        if aa1 != "-" and aa2 != "-":
+            aa_mapper[i] = (j, "")
+
+        # for insertion of seq1
+        if aa1 != "-" and aa2 == "-":
+            aa_mapper[i] = (j, ins_letters[k])
+            k += 1
+
+    # from align idx to sequence idx
+    # tmap_1 = re.search(aligned_seq1.replace("-", ""), seq1)
+    # tmap_2 = re.search(aligned_seq2.replace("-", ""), seq2)
+    # shift_1 = tmap_1.span()[0]
+    # shift_2 = tmap_2.span()[0]
+
+    start_1, start_2 = best_alignment.coordinates[:, 0]
+    _mapper = {k + start_1: "%d%s" % (v[0] + start_2, v[1]) for k, v in aa_mapper.items()}
+
+    out_mapper = dict()
+    # check head and tail of seq1 with E prefix
+    for i in range(1, len(seq1) + 1):
+        if i not in _mapper:
+            out_mapper[i] = "E%d" % i
+        else:
+            out_mapper[i] = _mapper[i]
+
+    ident = best_alignment.counts().identities / best_alignment.length
+    n_aligned = best_alignment.length - best_alignment.counts().gaps
+
+    coverage_1 = n_aligned / len(seq1)
+    coverage_2 = n_aligned / len(seq2)
+
+    return dict(seq1=seq1,
+                seq2=seq2,
+                aligned_seq1=aligned_seq1,
+                aligned_seq2=aligned_seq2,
+                alignment_length=best_alignment.length,
+                aligned_aa_mapper=out_mapper,
+                identity=round(ident, 3),
+                coverage_1=round(coverage_1, 3),
+                coverage_2=round(coverage_2, 3),
+                )
+
+
+class StructureAligner(object):
+    def __init__(self, query_path: str, ref_path: str):
+        self._query_st = StructureParser()
+        self._query_st.load_from_file(query_path)
+
+        self._ref_st = StructureParser()
+        self._ref_st.load_from_file(ref_path)
+
+        self.values = dict()
+        self.rot_mat = None
+        self.is_aligned = False
+        self.by_query = None
+        self.by_ref = None
+        self.query_path = query_path
+        self.ref_path = ref_path
+
+    @property
+    def __mmalign_path(self):
+        _path = shutil.which("MMAlign") or shutil.which("MMalign")
+        if _path is None:
+            raise RuntimeError("Executable program MMAlign is not found. "
+                               "Download from https://zhanggroup.org/MM-align/ ."
+                               "Build it and add MMAlign to environment PATH")
+        else:
+            return _path
+
+    @staticmethod
+    def __parser_rotation_matrix(matrix_file: str):
+        rotation_matrix = []
+        translation_vector = []
+
+        with open(matrix_file, 'r') as file:
+            lines = file.readlines()
+            values = lines[2:5]
+            for cur_line in values:
+                tmp = re.split(pattern=r"\s+", string=cur_line.strip())
+                assert len(tmp) == 5
+                rotation_matrix.append(tmp[2:])
+                translation_vector.append(tmp[1])
+        return dict(R=np.array(rotation_matrix).astype(np.float32),
+                    T=np.array(translation_vector).astype(np.float32))
+
+    @staticmethod
+    def __parse_terminal_outputs(output_string: str) -> Dict[str, Any]:
+        lines = re.split(pattern=r"\n", string=output_string)
+        # chain mapping
+        patterns = dict(query_chain_ids=r"Structure_1.+\.pdb:([\w:]+)",
+                        ref_chain_ids=r"Structure_2.+\.pdb:([\w:]+)",
+                        query_total_length=r"Length of Structure_1.*?(\d+).*residues",
+                        ref_total_length=r"Length of Structure_2.*?(\d+).*residues",
+                        aligned_length=r"Aligned length=.*?(\d+)",
+                        rmsd=r"RMSD=.*?([\d.]+)",
+                        tmscore_by_query=r"TM-score=.*?([\d.]+).+Structure_1",
+                        tmscore_by_ref=r"TM-score=.*?([\d.]+).+Structure_2",
+                        aligned_seq_start=r"denotes other aligned residues",
+                        )
+
+        values = dict()
+        for idx, line in enumerate(lines):
+            current_keys = list(patterns.keys())
+            for key in current_keys:
+                tmp = re.search(patterns[key], line)
+                if tmp:
+                    if key in ['query_chain_ids', 'ref_chain_ids']:
+                        values[key] = re.split(pattern=":", string=tmp.groups()[0])
+                        del patterns[key]
+                    elif key in ['query_total_length', 'ref_total_length', 'aligned_length']:
+                        values[key] = int(tmp.groups()[0])
+                        del patterns[key]
+                    elif key in ['rmsd', 'tmscore_by_query', 'tmscore_by_ref']:
+                        values[key] = float(tmp.groups()[0])
+                        del patterns[key]
+                    elif key == "aligned_seq_start":
+                        # idx + 1 and idx + 3 for aligned sequences 1 and 2
+                        seq_1 = lines[idx + 1]
+                        seq_2 = lines[idx + 3]
+
+                        sp1 = re.split(pattern=r"\*", string=seq_1)
+                        sp2 = re.split(pattern=r"\*", string=seq_2)
+                        values["query_sequences"] = sp1[:-1] if "*" in seq_1 else sp1
+                        values["ref_sequences"] = sp2[:-1] if "*" in seq_2 else sp2
+                        del patterns[key]
+        return values
+
+    def make_alignment(self, query_chains: Optional[List[str]] = None,
+                       ref_chains: Optional[List[str]] = None, timeout=300.0):
+        """
+
+        :param
+        query_chains: list, None
+        for all chains
+            :param
+        ref_chains: list, None
+        for all chains
+            :param
+        timeout: default
+        300
+        :return:
+        """
+
+        program_path = self.__mmalign_path
+
+        # clone
+        if isinstance(query_chains, list):
+            q_st = self._query_st.pick_chains(query_chains)
+        else:
+            q_st = self._query_st
+
+        if isinstance(ref_chains, list):
+            r_st = self._ref_st.pick_chains(ref_chains)
+        else:
+            r_st = self._ref_st
+
+        q_ch_mapper = q_st.make_one_letter_chain()
+        r_ch_mapper = r_st.make_one_letter_chain()
+
+        q_ch_mapper_r = {v: k for k, v in q_ch_mapper.items()}
+        r_ch_mapper_r = {v: k for k, v in r_ch_mapper.items()}
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            _tmp_a = os.path.join(tmp_dir, "a.pdb")
+            q_st.to_pdb(_tmp_a)
+
+            _tmp_b = os.path.join(tmp_dir, "b.pdb")
+            r_st.to_pdb(_tmp_b)
+
+            matrix_file = os.path.join(tmp_dir, "m.txt")
+            _command = "%s %s %s -m %s" % (program_path, _tmp_a, _tmp_b, matrix_file)
+
+            try:
+                result = subprocess.run(_command, shell=True, check=True,
+                                        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                                        timeout=timeout)
+            except Exception as e:
+                print("%s: between files %s and %s; between chains: %s and %s" % (
+                    str(e), self.query_path, self.ref_path,
+                    str(q_st.chain_ids), str(r_st.chain_ids))
+                      )
+            else:
+                self.values = self.__parse_terminal_outputs(result.stdout.decode())
+                self.rot_mat = self.__parser_rotation_matrix(matrix_file)
+                self.is_aligned = True
+                self.by_query = q_st.chain_ids if query_chains is None else query_chains
+                self.by_ref = r_st.chain_ids if ref_chains is None else ref_chains
+                self.values["query_chain_ids"] = [q_ch_mapper_r.get(ch, ch) for ch in self.values["query_chain_ids"]]
+                self.values["ref_chain_ids"] = [r_ch_mapper_r.get(ch, ch) for ch in self.values["ref_chain_ids"]]
+
+    def save_aligned_query(self, out_file: str):
+        """
+
+        :param
+        out_file:.cif
+        file
+        :return:
+        """
+        if not self.is_aligned:
+            raise RuntimeError("structure not aligned, run make_alignment first")
+
+        super_imposer = Superimposer()
+        super_imposer.rotran = (self.rot_mat["R"].T, self.rot_mat["T"])
+
+        bio_s = gemmi2bio(self._query_st.STRUCT)
+        super_imposer.apply(bio_s)
+        query_st_aligned = bio2gemmi(bio_s)
+
+        block = query_st_aligned.make_mmcif_block()
+        block.write_file(out_file)

gemmi_protools/tools/dockq.py

@@ -0,0 +1,128 @@
+"""
+@Author: Luo Jiejian
+"""
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from copy import deepcopy
+from typing import List, Tuple
+
+import gemmi
+import pandas as pd
+
+from gemmi_protools.io.reader import StructureParser
+
+
+def dockq_score_interface(query_model: str,
+                          native_model: str,
+                          partner_1_mapping: List[Tuple[str, str]],
+                          partner_2_mapping: List[Tuple[str, str]],
+                          ):
+    """
+    Calculate Dockq Score for an interface (partner 1 vs partner 2)
+
+    :param query_model: str
+        path of query model, support .pdb, .pdb.gz, .cif, .cif.gz
+    :param native_model:
+    :param partner_1_mapping: a list of chain ID mapping between query and native for partner1 of the interface
+        e.g. [(q chain1, n chain1), (q chain2, n chain2)]
+    :param partner_2_mapping:
+    :return:
+    """
+    dockq_program = shutil.which("DockQ")
+    if dockq_program is None:
+        raise RuntimeError("DockQ is need")
+
+    assert len(partner_1_mapping) > 0, "partner_1_mapping must be a list of chain ID tuples, can't be empty"
+    assert len(partner_2_mapping) > 0, "partner_2_mapping must be a list of chain ID tuples, can't be empty"
+
+    def load_struct(path: str, partner_1: List[str], partner_2: List[str]):
+        st = StructureParser()
+        st.load_from_file(path)
+        st.clean_structure()
+
+        for ch in partner_1 + partner_2:
+            if ch not in st.chain_ids:
+                raise ValueError("Chain %s not found for %s (only [%s])" % (ch, path, " ".join(st.chain_ids)))
+
+        # merge chains in each each partner into on chain
+        # partner_1 with chain ID A
+        # partner_2 with chain ID B
+
+        chain_a = gemmi.Chain("A")
+        idx_a = 1
+        for ch in partner_1:
+            for res in st.get_chain(ch):
+                nr = deepcopy(res)
+                nr.seqid.icode = " "
+                nr.seqid.num = idx_a
+                chain_a.add_residue(nr)
+                idx_a += 1
+
+        chain_b = gemmi.Chain("B")
+        idx_b = 1
+        for ch in partner_2:
+            for res in st.get_chain(ch):
+                nr = deepcopy(res)
+                nr.seqid.icode = " "
+                nr.seqid.num = idx_b
+                chain_b.add_residue(nr)
+                idx_b += 1
+
+        model = gemmi.Model(1)
+        model.add_chain(chain_a)
+        model.add_chain(chain_b)
+
+        struct = gemmi.Structure()
+        struct.add_model(model)
+
+        output = StructureParser(struct)
+        return output
+
+    partner_1_query, partner_1_native = list(zip(*partner_1_mapping))
+    partner_2_query, partner_2_native = list(zip(*partner_2_mapping))
+
+    q_st = load_struct(query_model, list(partner_1_query), list(partner_2_query))
+    n_st = load_struct(native_model, list(partner_1_native), list(partner_2_native))
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        result_file = os.path.join(tmp_dir, "result.json")
+        q_file = os.path.join(tmp_dir, "q.pdb")
+        n_file = os.path.join(tmp_dir, "n.pdb")
+        q_st.to_pdb(q_file, write_minimal_pdb=True)
+        n_st.to_pdb(n_file, write_minimal_pdb=True)
+
+        mapping = "AB:AB"
+
+        _command = "%s --mapping %s --json %s %s %s" % (dockq_program, mapping, result_file, q_file, n_file)
+        metrics = ['DockQ', 'F1', 'chain1', 'chain2']
+
+        try:
+            _ = subprocess.run(_command, shell=True, check=True,
+                               stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                               timeout=300.0)
+        except subprocess.CalledProcessError as e:
+            # Handle errors in the called executable
+            msg = e.stderr.decode()
+            outputs = pd.DataFrame(columns=metrics)
+        except Exception as e:
+            # Handle other exceptions such as file not found or permissions issues
+            msg = str(e)
+            outputs = pd.DataFrame(columns=metrics)
+        else:
+            with open(result_file, "r") as fin:
+                vals = json.load(fin)
+            msg = "Finished"
+            result = []
+            for v in vals["best_result"].values():
+                result.append(v)
+            outputs = pd.DataFrame(result)[metrics]
+
+        if len(outputs) > 0:
+            score = "%.4f" % outputs.iloc[0]["DockQ"]
+        else:
+            score = ""
+
+        return dict(score=score, status=msg)

gemmi_protools/tools/mesh.py

@@ -0,0 +1,197 @@
+"""
+@Author: Luo Jiejian
+"""
+import os
+import subprocess
+import tempfile
+from collections import defaultdict
+from typing import List, Optional, Union
+
+import freesasa
+import numpy as np
+import trimesh
+from Bio.PDB import Selection
+from Bio.PDB.ResidueDepth import _get_atom_radius, _read_vertex_array
+
+from gemmi_protools import StructureParser
+from gemmi_protools import gemmi2bio
+
+
+def _read_face_array(filename: str):
+    with open(filename) as fp:
+        face_list = []
+        for line in fp:
+            sl = line.split()
+            if len(sl) != 5:
+                # skip header
+                continue
+            vl = [int(x) for x in sl[0:3]]
+            face_list.append(vl)
+    return np.array(face_list)
+
+
+def get_mesh(struct_file: str, chains: Optional[List[str]] = None, MSMS: str = "msms"):
+    """
+
+    :param struct_file: str
+        .pdb, .cif, .pdb.gz, .cif.gz
+    :param chains: a list of chain names
+        default None to include all chains
+    :param MSMS: str
+        path of msms executable
+    :return:
+        https://ccsb.scripps.edu/msms/downloads/
+    """
+    xyz_tmp = tempfile.NamedTemporaryFile(delete=False).name
+    surface_tmp = tempfile.NamedTemporaryFile(delete=False).name
+    msms_tmp = tempfile.NamedTemporaryFile(delete=False).name
+    face_file = surface_tmp + ".face"
+    surface_file = surface_tmp + ".vert"
+
+    try:
+        st = StructureParser()
+        st.load_from_file(struct_file)
+        st.clean_structure(remove_ligand=True)
+
+        if chains is None:
+            st_p = st
+        else:
+            for ch in chains:
+                if ch not in st.chain_ids:
+                    raise ValueError("Chain %s not found (only [%s])" % (ch, " ".join(st.chain_ids)))
+            st_p = st.pick_chains(chains)
+
+        bio_st = gemmi2bio(st_p.STRUCT)
+        model = bio_st[0]
+
+        # Replace pdb_to_xyzr
+        # Make x,y,z,radius file
+        atom_list = Selection.unfold_entities(model, "A")
+
+        with open(xyz_tmp, "w") as pdb_to_xyzr:
+            for atom in atom_list:
+                x, y, z = atom.coord
+                radius = _get_atom_radius(atom, rtype="united")
+                pdb_to_xyzr.write(f"{x:6.3f}\t{y:6.3f}\t{z:6.3f}\t{radius:1.2f}\n")
+
+        # Make surface
+        MSMS = MSMS + " -no_header -probe_radius 1.5 -if %s -of %s > " + msms_tmp
+        make_surface = MSMS % (xyz_tmp, surface_tmp)
+        subprocess.call(make_surface, shell=True)
+        if not os.path.isfile(surface_file):
+            raise RuntimeError(
+                f"Failed to generate surface file using command:\n{make_surface}"
+            )
+
+    except Exception as e:
+        print(str(e))
+        mesh = None
+    else:
+        # Read surface vertices from vertex file
+        vertices = _read_vertex_array(surface_file)
+        faces = _read_face_array(face_file)
+        mesh = trimesh.Trimesh(vertices=vertices, faces=faces - 1)
+        mesh.merge_vertices()
+        mesh.update_faces(mesh.unique_faces())
+        mesh.update_faces(mesh.nondegenerate_faces())
+        mesh.remove_unreferenced_vertices()
+
+    # Remove temporary files
+    for fn in [xyz_tmp, surface_tmp, msms_tmp, face_file, surface_file]:
+        try:
+            os.remove(fn)
+        except OSError:
+            pass
+
+    return mesh
+
+
+def get_surface_residues(struct_file: str,
+                         chains: Optional[List[str]] = None,
+                         relative_sasa_cutoff: Union[int, float] = 0.15):
+    ####################
+    # check and pick
+    ####################
+    st = StructureParser()
+    st.load_from_file(struct_file)
+    st.clean_structure()
+
+    if chains is None:
+        chains = st.chain_ids
+
+    if isinstance(chains, list):
+        if len(chains) == 0:
+            raise ValueError("chains is not set")
+        else:
+            # check if chains valid
+            for ch in chains:
+                if ch not in st.chain_ids:
+                    raise ValueError("Chain %s not found" % ch)
+
+    st_p = st.pick_chains(chains)
+    # sequences = {k: s.replace("-", "").upper() for k, s in st_p.polymer_sequences().items()}
+
+    # start from 1
+    seq_num_mapper = dict()
+    for chain in st_p.MODEL:
+        for i, res in enumerate(chain):
+            key = (chain.name, str(res.seqid.num) + res.seqid.icode.strip(), res.name)
+            seq_num_mapper[key] = i + 1
+
+    # make one upper letter chain ID
+    mapper = st_p.make_one_letter_chain(only_uppercase=True)
+    mapper_r = {v: k for k, v in mapper.items()}
+
+    ####################
+    # save to pdb
+    ####################
+    with tempfile.NamedTemporaryFile(delete=True, suffix=".pdb", mode='w') as tmp_file:
+        st_p.to_pdb(tmp_file.name)
+        structure = freesasa.Structure(tmp_file.name)
+
+    result = freesasa.calc(structure)
+
+    residue_areas = result.residueAreas()
+
+    surface_residues_relative_sasa = dict()
+    surface_atoms = defaultdict(list)
+    for atom_index in range(structure.nAtoms()):
+        ch = structure.chainLabel(atom_index)
+        ch = mapper_r.get(ch, ch)
+
+        res_num = structure.residueNumber(atom_index).strip()
+        res_name = structure.residueName(atom_index)
+        atom_sasa = result.atomArea(atom_index)
+
+        res_id = (ch, res_num, res_name)
+        res_relative_total = residue_areas[ch][res_num].relativeTotal
+        if res_relative_total > relative_sasa_cutoff:
+            if res_id not in surface_residues_relative_sasa:
+                surface_residues_relative_sasa[res_id] = res_relative_total
+            if atom_sasa > 0:
+                atom_name = structure.atomName(atom_index).strip()
+                pos = structure.coord(atom_index)
+                surface_atoms[res_id].append((atom_sasa, atom_name, pos))
+
+    results = []
+    for res_id, query_atoms in surface_atoms.items():
+        seq_loc = seq_num_mapper[res_id]
+
+        query_atoms.sort(reverse=True)
+        centroid = tuple(np.array([a[2] for a in query_atoms[0:3]]).mean(axis=0).tolist())
+        results.append((res_id[0],
+                        res_id[1],
+                        res_id[2],
+                        seq_loc,
+                        centroid,
+                        surface_residues_relative_sasa[res_id]
+                        )
+                       )
+    dtype = [("chain_name", "U5"),
+             ("residue_numi", "U8"),
+             ("residue_name", "U5"),
+             ("sequential_residue_num", "i4"),
+             ("centroid", ("f4", (3,))),
+             ("relative_sasa", "f4"),
+             ]
+    return np.array(results, dtype=dtype)