probe_profile

#!/usr/bin/env python

import argparse
from pathlib import Path
from typing import List

import gridData
from Bio.PDB.Structure import Structure
from joblib import Parallel, delayed

from script import maxpmap
from script.alignresenv import align_res_env
from script.profile import create_residue_interaction_profile
from script.resenv import resenv
from script.setting import parse_yaml
from script.utilities import util
from script.utilities.Bio import PDB as uPDB
from script.utilities.logger import logger

VERSION = "0.1.0"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Calculate residue interaction profile around probe")
    parser.add_argument("setting_yaml", type=Path)
    parser.add_argument("-v,--verbose", dest="verbose", action="store_true")
    parser.add_argument("--debug", action="store_true")
    parser.add_argument("--version", action="version", version=VERSION)
    args = parser.parse_args()

    if args.debug:
        logger.setLevel("debug")
    elif args.verbose:
        logger.setLevel("info")
    # else: logger level is "warn"

    logger.info(f"read yaml: {args.setting_yaml}")
    setting = parse_yaml(args.setting_yaml)

    indices = util.expand_index(setting["general"]["iter_index"])
    WORKING_DIR = setting["general"]["workdir"]
    probe_resn = setting["input"]["probe"]["cid"]
    JOB_NAME = setting["general"]["name"]
    n_jobs = setting["general"]["multiprocessing"]
    mapname = setting["probe_profile"]["resenv"]["map"]
    threshold = setting["probe_profile"]["threshold"]
    env_distance = setting["probe_profile"]["resenv"]["env_dist"]
    profile_types = setting["probe_profile"]["profile"]["types"]

    # generate max_pmap
    pmap_pathes = [f"{WORKING_DIR}/system{idx}/PMAP_{JOB_NAME}_{mapname}.dx" for idx in indices]
    pmaps = [gridData.Grid(path) for path in pmap_pathes]
    max_pmap = maxpmap.grid_max(pmaps)

    # extract environments around probes
    trajectory_files = [f"{WORKING_DIR}/system{idx}/{JOB_NAME}_woWAT_10ps.pdb" for idx in indices]
    trajectories = [uPDB.MultiModelPDBReader(path) for path in trajectory_files]
    probe_environment_structs: List[Structure] = Parallel(n_jobs=n_jobs)(  # type: ignore
        delayed(resenv)(
            grid=max_pmap,
            trajectory=trajectory,
            resn=probe_resn,
            res_atomnames=[" CB "],
            threshold=threshold,
            env_distance=env_distance,
            verbose=args.verbose,
        )
        for trajectory in trajectories
    )
    probe_environment_struct = uPDB.concatenate_structures(probe_environment_structs)

    # remove unnecessary atoms
    target_residue_atoms = set()  # convert from list to tuple
    for profile_type in profile_types:
        target_residue_atoms.update([(*lst,) for lst in profile_type["atoms"]])
    sele = uPDB.Selector(
        lambda a: (uPDB.get_atom_attr(a, "resname"), uPDB.get_atom_attr(a, "fullname")) in target_residue_atoms
        or uPDB.get_atom_attr(a, "resname") == probe_resn
    )
    probe_environment_struct = uPDB.extract_substructure(probe_environment_struct, sele)

    # align structures in accordance with the probe structures
    ref_struct = probe_environment_struct[0].copy()  # all structures are superimposed to this
    aligned_environment = align_res_env(probe_environment_struct, ref_struct, probe_resn)

    # create residue interaction profile for each residue type
    for profile_type in profile_types:
        residue_type = profile_type["name"]
        target_residue_atoms = [(*lst,) for lst in profile_type["atoms"]]  # convert from list to tuple

        try:
            g = create_residue_interaction_profile(aligned_environment, target_residue_atoms)
            g.export(f"{WORKING_DIR}/{JOB_NAME}_{probe_resn}_mesh_{residue_type}.dx", type="short")
        except Exception as e:
            # glysine must be in here because it does not have CB atom
            logger.error(f"Error: {e} - skip this residue_type / target_residue_atoms pair")
            logger.error(f"residue_type: {residue_type}")
            logger.error(f"target_residue_atoms: {target_residue_atoms}")