main.py

import glob, pathlib, warnings, dataclasses, enum, pprint, os
import itertools as it
from typing import Iterator, Self

from Bio.PDB import MMCIFParser, PDBParser, NeighborSearch, Selection, Structure
from Bio.PDB.mmcifio import MMCIFIO
from Bio.PDB.Structure import Structure
from Bio.PDB.Residue import Residue
from Bio.PDB.Atom import Atom

import vg
import numpy as np
import pandas as pd

import tqdm

import torch
import torch_geometric as pyg

import hydra
from omegaconf import DictConfig, OmegaConf

from src.load import get_paths, collate_paths, load_structures, load_contact_tables, load_result_tables
from src.residue_graph import ResidueGraph, build_residue_graph
from src.mutation import mutate_residue_graph, get_mutated_res_idx
from src.stability import make_stability_table, calculate_torsion

def load_emcast_tsuboyama_dataset(
    cfg: DictConfig,
) -> tuple[list[Structure],list[pd.DataFrame],list[pd.DataFrame]]:
    """Load structures, voronota contacts, and result tables describing EmCAST predictions on a subset of the Tsuboyama dataset.
    """
    structure_paths, contact_paths, result_paths = get_paths(
        pathlib.Path(cfg.data.dir),
        cfg.data.structures_suffix,
        cfg.data.contacts_suffix,
        cfg.data.results_suffix,
    )
    print("num paths pre-collation:", len(structure_paths),len(contact_paths),len(result_paths))

    structure_paths, contact_paths, result_paths = collate_paths(
        structure_paths,
        contact_paths,
        result_paths,
        cfg.data.exclude,
    )
    print("num paths post-collation:", len(structure_paths),len(contact_paths),len(result_paths))

    print("loading structures.")
    structures = load_structures(
        tqdm.tqdm(structure_paths,leave=False),
        cfg.data.structures_suffix,
    )
    print("loading contact tables.")
    contact_tables = load_contact_tables(
        tqdm.tqdm(contact_paths,leave=False),
        cfg.data.contact_sep,
        cfg.data.contact_column_names,
        cfg.data.contact_column_query,
    )
    print("loading result tables.")
    result_tables = load_result_tables(
        tqdm.tqdm(result_paths,leave=False),
        cfg.data.results_sep,
    )
    return (
        structures,
        contact_tables,
        result_tables,
    )

def generate_graph_data(
    cfg: DictConfig,
    structures: list[Structure],
    contact_tables: list[pd.DataFrame],
    result_tables: list[pd.DataFrame],
    n: int,
) -> tuple[list[ResidueGraph],list[list[ResidueGraph]]]:
    """Construct residue graphs from structures and voronota contact tables, then induce mutations according to the EmCAST / Tsuboyama result tables.
    """
    resnames = cfg.residues.names.three_letter_code
    reschars = cfg.residues.names.one_letter_code
    resname_index = {
        x: i
        for (i,x) in enumerate(sorted(resnames))
    }
    reschar_to_resname = {
        res_char: res_name
        for (res_char, res_name) in zip(reschars, resnames)
    }
    planar_angle_atoms = cfg.residues.angles.planar
    chi1_angle_atoms = cfg.residues.angles.chi1
    print("building residue graphs.")
    residue_graphs = [
        build_residue_graph(
            structures[i],
            contact_tables[i],
            resname_index,
            planar_angle_atoms,
            chi1_angle_atoms,
        )
        for i in tqdm.trange(n,leave=False)
    ]

    # TODO parallelize
    print("generating residue graph mutations.")
    residue_graph_mutations = [[] for i in range(n)]
    for i in tqdm.trange(n,leave=False):
        for (chain,mut_key) in result_tables[i][['chain','pos']].to_numpy():
            residue_graph_mutations[i].append(mutate_residue_graph(
                residue_graphs[i],
                chain,
                mut_key,
                resname_index,
                reschar_to_resname,
            ))
    num_mutations = [len(x) for x in residue_graph_mutations]
    print("num mutations:",sum(num_mutations))

    return (
        residue_graphs,
        residue_graph_mutations,
    )

def generate_labels(
    cfg: DictConfig,
    residue_graph_mutations: list[list[Structure]],
    result_tables: list[pd.DataFrame],
    n: int,
) -> tuple[list[pd.DataFrame],list[list[float]]]:
    """Generate residue graph labels. Stability is calculated from the EmCAST / Tsuboyama results. Torsion is calculated at mutated positions using the dihedral angle feature of Bio.PDB.
    """
    print("calculating stability.")
    stability_tables = [
        make_stability_table(
            result_tables[i],
            *cfg.data.results_emcast_query,
            *cfg.data.results_tsuboyama_query,
        )
        for i in tqdm.trange(n,leave=False)
    ]
    num_stability_observations = sum([len(x) for x in stability_tables])
    print("num observations:",num_stability_observations)

    print("calculating torsion.")
    torsion = [
        [
            calculate_torsion(
                x,
                get_mutated_res_idx(x),
            )
            for x in residue_graph_mutations[i]
        ]
        for i in tqdm.trange(n,leave=False)
    ]
    num_torsion = sum([len(x) for x in torsion])
    num_null_torsion = len([x for batch in torsion for x in batch if x == (0.,0.)])
    print(f"num torsion: {num_torsion} ({num_null_torsion} null, {num_torsion - num_null_torsion} usable.)")

    return (
        stability_tables,
        torsion,
    )

@hydra.main(version_base=None, config_path="config", config_name="config")
def main(cfg: DictConfig) -> None:
    print("configuration:")
    pprint.pprint(
        OmegaConf.to_container(cfg),
        width=os.get_terminal_size().columns,
    )

    structures, contact_tables, result_tables = load_emcast_tsuboyama_dataset(cfg)
    assert len(structures) == len(contact_tables) == len(result_tables)
    n = len(structures)

    residue_graphs, residue_graph_mutations = generate_graph_data(
        cfg,
        structures,
        contact_tables,
        result_tables,
        n,
    )

    stability_tables, torsion = generate_labels(
        cfg,
        residue_graph_mutations,
        result_tables,
        n,
    )


if __name__ == "__main__":
    main()