Use of experimental pyxmolpp2 pipes

pipes-2.0.py

from pyxmolpp2 import *
import pyxmolpp2.pipe as pp
from tqdm import tqdm
from typing import List, Union


def has_atoms(*atom_names):
    def predicate(residue: Residue):
        try:
            for atom_name in atom_names:
                residue[atom_name]
            return True
        except IndexError:
            return False

    return predicate


def atom_pairs_from_consequent_residues(atom_names_1: Union[str, List[str]], atom_names_2: Union[str, List[str]]):
    if isinstance(atom_names_1, str):
        atom_names_1 = [atom_names_1]
    if isinstance(atom_names_2, str):
        atom_names_2 = [atom_names_2]

    def to_atoms(residue):
        if residue is not None:
            return residue.atoms
        return AtomSelection()

    def selector(frame: Frame):
        residues = frame.residues.filter(has_atoms(*atom_names_1))
        return [
            (first_atom, second_atom)
            for r in residues
            for first_atom in r.atoms.filter(aName.is_in(*atom_names_1))
            for second_atom in to_atoms(r.next).filter(aName.is_in(*atom_names_2))
        ]

    return selector


def atom_pairs_from_one_residue(atom_names_1: Union[str, List[str]], atom_names_2: Union[str, List[str]]):
    if isinstance(atom_names_1, str):
        atom_names_1 = [atom_names_1]
    if isinstance(atom_names_2, str):
        atom_names_2 = [atom_names_2]

    def selector(frame: Frame):
        residues = frame.residues.filter(has_atoms(*atom_names_1))
        return [
            (first_atom, second_atom)
            for r in residues
            for first_atom in r.atoms.filter(aName.is_in(*atom_names_1))
            for second_atom in r.atoms.filter(aName.is_in(*atom_names_2))
        ]

    return selector


def output_filename_formatter(atom1: Atom, atom2: Atom):
    if atom1.residue == atom2.residue:
        return f"./vectors/{atom1.residue.id.serial:02d}_{atom1.name}-{atom2.name}.csv"
    return f"./vectors/" \
           f"{atom1.residue.id.serial:02d}_{atom1.name}-" \
           f"{atom2.residue.id.serial:02d}_{atom2.name}" \
           f".csv"


path_to_traj = "./archive/example_xtc_trj_ubq_NPT_bussi_box_8/"
ref: Frame = PdbFile(f"{path_to_traj}/run00001.pdb").frames()[0]
traj = Trajectory(ref)

for i in range(1, 10):
    traj.extend(GromacsXtcFile(f"{path_to_traj}/run{i:05d}.xtc", 1000))

processing = (
        pp.WriteVectorsToCsv(atom_pairs_from_one_residue("N", "H"), output_filename_formatter) |
        pp.WriteVectorsToCsv(atom_pairs_from_one_residue("CA", ("HA", "HA2", "HA3")), output_filename_formatter) |
        pp.WriteVectorsToCsv(atom_pairs_from_one_residue("H", ("HA", "HA2", "HA3")), output_filename_formatter) |
        pp.WriteVectorsToCsv(atom_pairs_from_one_residue("C", "O"), output_filename_formatter) |
        pp.WriteVectorsToCsv(atom_pairs_from_one_residue("CA", "C"), output_filename_formatter)
)


tqdm(traj) | processing | pp.Run()