<jupyter_start><jupyter_text>Summary----  Imports<jupyter_code>import concurrent.futures
import itertools
import os
from pathlib import Path

import numpy as np
import pandas as pd

import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import proteinsolver
import psutil
import pyarrow as pa
import pyarrow.parquet as pq
import torch
from kmbio import PDB
from scipy import stats
from torch.utils.data import DataLoader, Dataset
DEBUG = "CI" not in os.environ
if DEBUG:
    %load_ext autoreload
    %autoreload 2<jupyter_output><empty_output><jupyter_text>Parameters<jupyter_code>UNIQUE_ID = "191f05de"  # No attention
# UNIQUE_ID = "0007604c"  # 5-layer graph-conv with attention, batch_size=1
# UNIQUE_ID = "91fc9ab9"  # 4-layer graph-conv with attention, batch_size=4
BEST_STATE_FILES = {
    #
    "191f05de": "protein_train/191f05de/e53-s1952148-d93703104.state"
}
NOTEBOOK_NAME = "06_global_analysis_of_protein_folding"
NOTEBOOK_PATH = Path(NOTEBOOK_NAME)
NOTEBOOK_PATH.mkdir(exist_ok=True)
NOTEBOOK_PATH
INPUT_PATH = Path(os.getenv("DATAPKG_INPUT_DIR"))
INPUT_PATH
DATAPKG_DATA_DIR = Path(f"~/datapkg_data_dir").expanduser().resolve()
DATAPKG_DATA_DIR
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
device
proteinsolver.settings.data_url = DATAPKG_DATA_DIR.as_posix()
proteinsolver.settings.data_url<jupyter_output><empty_output><jupyter_text>Load data<jupyter_code>!ls {INPUT_PATH}/global_analysis_of_protein_folding
!ls {INPUT_PATH}/global_analysis_of_protein_folding/aan0693_designed-PDB-files
!ls {INPUT_PATH}/global_analysis_of_protein_folding/aan0693_SI_datasets<jupyter_output>counts_and_ec50s.tar.gz		  protein_and_dna_sequences.tar.gz
design_scripts.tgz		  stability_scores
design_structural_metrics	  stability_scores.tar.gz
design_structural_metrics.tar.gz  unfolded_state_model_params
fig1_thermodynamic_data.csv<jupyter_text>aan0693_SI_datasets<jupyter_code>!ls {INPUT_PATH}/global_analysis_of_protein_folding/aan0693_SI_datasets/stability_scores
def remove_controls(df):
    df = df[
        (~df["name"].str.endswith("_hp")) & (~df["name"].str.endswith("_random")) & (~df["name"].str.endswith("_buryD"))
    ]
    return df
def load_stability_scores(key):
    stability_scores = pd.read_csv(
        INPUT_PATH
        / "global_analysis_of_protein_folding"
        / "aan0693_SI_datasets"
        / "stability_scores"
        / f"{key}_stability_scores",
        sep="\t",
    )
    stability_scores = remove_controls(stability_scores)

    for energy_function in ["talaris2013", "betanov15"]:
        rosetta_energies_file = (
            INPUT_PATH
            / "global_analysis_of_protein_folding"
            / "aan0693_SI_datasets"
            / "design_structural_metrics"
            / f"{key}_relax_scored_{'filtered_' if energy_function == 'betanov15' else ''}{energy_function}.sc"
        )
        if not rosetta_energies_file.is_file():
            print(f"Not loading Rosetta energies for {energy_function}!")
            continue

        before_ = len(stability_scores)
        relax_scored_filtered = pd.read_csv(
            rosetta_energies_file, sep="\t" if energy_function == "betanov15" else " +", engine="python"
        ).rename(columns={"description": "name", "total_score": f"{energy_function}_score"})
        stability_scores = stability_scores.merge(
            relax_scored_filtered[["name", f"{energy_function}_score"]], on="name", how="outer"
        )
#         assert len(stability_scores) == before_, (len(stability_scores), before_)

    stability_scores["library_name"] = key
    return stability_scores<jupyter_output><empty_output><jupyter_text>stability_scores<jupyter_code># stability_scores = {}
# for key in ["rd1", "rd2", "rd3", "rd4", "ssm2"]:
#     print(key)
#     stability_scores[key] = load_stability_scores(key)
# stability_scores["fig1"] = pd.read_csv(
#     INPUT_PATH / "global_analysis_of_protein_folding" / "aan0693_SI_datasets" / "fig1_thermodynamic_data.csv"
# ).assign(library_name="fig1")
# with open(NOTEBOOK_PATH.joinpath("stability_scores.raw.incremental.torch"), "wb") as fout:
#     torch.save(stability_scores, fout)
stability_scores = torch.load(NOTEBOOK_PATH.joinpath("stability_scores.torch"))
stability_scores["rd4"].head()<jupyter_output><empty_output><jupyter_text>Load model<jupyter_code>%run protein_train/{UNIQUE_ID}/model.py
batch_size = 1
num_features = 20
adj_input_size = 2
hidden_size = 128
frac_present = 0.5
frac_present_valid = frac_present
info_size= 1024
state_file = BEST_STATE_FILES[UNIQUE_ID]
state_file
net = Net(
    x_input_size=num_features + 1, adj_input_size=adj_input_size, hidden_size=hidden_size, output_size=num_features
)
net.load_state_dict(torch.load(state_file, map_location=device))
net.eval()
net = net.to(device)<jupyter_output><empty_output><jupyter_text>Structure probabilities<jupyter_code>structure_file = "/home/kimlab1/database_data/datapkg_input_dir/global_analysis_of_protein_folding/aan0693_designed-PDB-files/rd1/EEHEE_rd1_0004.pdb"
structure = PDB.load(structure_file)
pdata = proteinsolver.utils.extract_seq_and_adj(structure, 'A')
data = proteinsolver.datasets.protein.row_to_data(pdata)
data = proteinsolver.datasets.protein.transform_edge_attr(data)
data.to(device)

print(
    proteinsolver.utils.get_node_outputs(
        net, data.x, data.edge_index, data.edge_attr, 20, output_transform="logproba"
    )
)<jupyter_output>tensor([-2.1866, -2.8192, -0.8576, -3.1876, -6.9420, -1.0521, -0.1976, -1.1227,
        -1.4856, -2.1950, -2.0287, -0.3771, -2.8255, -1.1229, -2.2655, -3.1545,
        -0.8361, -0.3876, -1.8040, -1.5459, -2.0852, -0.1505, -1.7250, -2.2637,
        -3.0158, -0.1281, -0.4505, -3.4455, -3.2797, -0.2760, -1.1259, -2.8295,
        -2.0394, -0.4193, -4.0072, -1.1092, -2.4877, -1.2500, -3.2661, -2.8514,
        -3.5015])<jupyter_text>structure_proba<jupyter_code>def get_structure_proba(net, filename):
    try:
        library_name = filename.split("_")[1]
    except IndexError:
        return np.nan, np.nan
    filepath = (
        INPUT_PATH / "global_analysis_of_protein_folding" / "aan0693_designed-PDB-files" / library_name / filename
    )
    if not filepath.is_file():
        return np.nan, np.nan
    structure = PDB.load(filepath)
    pdata = proteinsolver.utils.extract_seq_and_adj(structure, "A")
    data = proteinsolver.datasets.protein.row_to_data(pdata)
    data = proteinsolver.datasets.protein.transform_edge_attr(data)
    data = data.to(device)
    #     log_prob = proteinsolver.utils.get_node_proba(net, data.x, data.edge_index, data.edge_attr, 20)
    log_prob = proteinsolver.utils.get_node_outputs(
        net, data.x, data.edge_index, data.edge_attr, 20, output_transform="logproba"
    )
    return log_prob.sum().item(), len(pdata.sequence)
def get_filepath(filename):
    try:
        library = filename.split("_")[1]
    except IndexError:
        #         print(f"Cound not guess library for file: {filename}!")
        return None
    filepath = INPUT_PATH / "global_analysis_of_protein_folding" / "aan0693_designed-PDB-files" / library / filename
    if not filepath.is_file():
        #         print(f"File {filename} does not exist!")
        return None
    return filepath


class DesignsDataset(Dataset):
    """Face Landmarks dataset."""

    @staticmethod
    def transform(pdata):
        data = proteinsolver.datasets.protein.row_to_data(pdata)
        data = proteinsolver.datasets.protein.transform_edge_attr(data)
        return data

    @staticmethod
    def add_domain(filepath, data):
        domain = filepath.name.split("_")[0]
        data.domain = domain
        return data

    def __init__(self, filenames, sequences, use_ref_graph=True):
        self.use_ref_graph = use_ref_graph

        ref_files = {
            "EEHEE": ("EEHEE_nmr_5up1.cif", "TTVKLGDIKVTFDNPEKAKKYAQKLAKIYQLTVHVHGDTIHVK"),
            "EHEE": ("EHEE_nmr_5up5.cif", "TQTQEFDNEEEARKAEKELRKENRRVTVTQENGRWRVTWD"),
            "HEEH": ("HEEH_nmr_5uyo.cif", "DVEEQIRRLEEVLKKNQPVTWNGTTYTDPNEIKKVIEELRKSM"),
            "HHH": ("HHH_nmr_5uoi.cif", "RKWEEIAERLREEFNINPEEAREAVEKAGGNEEEARRIVKKRL"),
        }
        domains = list({f.split("_")[0] for f in filenames})
        if set(domains) - set(ref_files):
            print(f"Do not have reference for the following domains: {set(domains) - set(ref_files)}!")
        pdata_ref = {}
        for domain, (filename, sequence) in ref_files.items():
            filepath = get_filepath(filename)
            structure = PDB.load(filepath)
            pdata = proteinsolver.utils.extract_seq_and_adj(structure, "A")
            assert pdata.sequence == sequence, (domain, filename, pdata.sequence, sequence)
            pdata_ref[domain] = pdata
        #         pdata_ref["EEHEE"] = pdata_ref["EEHEE"]._replace(sequence=pdata_ref["EEHEE"].sequence + "LESSGGS")
        self.data_ref = {domain: self.transform(pdata) for domain, pdata in pdata_ref.items()}

        self.filepaths, self.sequences = list(
            zip(
                *[
                    (f, s)
                    for f, s in itertools.zip_longest((get_filepath(f) for f in filenames), sequences)
                    if f is not None
                ]
            )
        )

    def __len__(self):
        return len(self.filepaths)

    def __getitem__(self, idx):
        filepath = self.filepaths[idx]
        sequence = self.sequences[idx]
        structure = PDB.load(filepath)
        pdata = proteinsolver.utils.extract_seq_and_adj(structure, "A")

        name = filepath.name
        domain = filepath.name.split("_")[0]

        start = None
        end = None
        if pd.isnull(sequence):
            pass
        elif pdata.sequence != sequence:
            mismatch_size = len(sequence) - len(pdata.sequence)
            if mismatch_size <= 0:
                print(f"Sequence mismatch of size '{mismatch_size}' for protein {filepath.name}!")
                return name, domain, None
            else:
                for start in range(mismatch_size + 1):
                    end = start + len(pdata.sequence)
                    sequence_slice = sequence[start:end]
                    frac_identical = sum(
                        a1 == a2 for (a1, a2) in itertools.zip_longest(pdata.sequence, sequence_slice)
                    ) / len(pdata.sequence)
                    if frac_identical >= 1.0:
                        break
                if frac_identical < 1.0:
                    print(f"Sequence mismatch for protein {filepath.name}: ({frac_identical:.2%} identical).")
                    return name, domain, None

        data = self.transform(pdata)

        if self.use_ref_graph and data.domain in self.data_ref:
            data_ref = self.data_ref[data.domain].clone()
            if start is not None:
                data_ref.x = data_ref.x[start:end]
                mask = ((data_ref.edge_index >= start) & (data_ref.edge_index < end)).all(dim=0)
                data_ref.edge_index = data_ref.edge_index[:, mask] - start
                data_ref.edge_attr = data_ref.edge_attr[mask, :]
                assert (data_ref.edge_index >= 0).all().item()
                assert (data_ref.edge_index < data_ref.x.size(0)).all().item()

            if data.x.size(0) != data_ref.x.size(0):
                print(f"Reference structure length mismatch for protein {filepath.name}.")
                data.edge_index_ref = None
                data.edge_attr_ref = None
                data.identity_to_ref = None
            else:
                assert data.x.size(0) == data_ref.x.size(0)
                assert data_ref.edge_index.min() >= 0
                assert data_ref.edge_index.max() < data.x.size(0)
                data.edge_index_ref = data_ref.edge_index
                data.edge_attr_ref = data_ref.edge_attr
                data.identity_to_ref = float((data.x == data_ref.x).sum().item()) / data.x.size(0)
        else:
            data.edge_index_ref = data.edge_index
            data.edge_attr_ref = data.edge_attr
            data.identity_to_ref = 1

        return name, domain, data
design_dataset = DesignsDataset(
    stability_scores["rd2"]["name"], stability_scores["rd2"]["sequence"], use_ref_graph=False
)
design_dataset[4000]
# for library in ["rd1", "rd2", "rd3", "rd4"]:
#     stability_scores[library]["network_score"], stability_scores[library]["sequence_length"] = np.array(
#         [get_structure_proba(net, name) for name in stability_scores[library]["name"]]
#     ).T
def apply_scores(df):
    results = []
    design_dataset = DesignsDataset(df["name"], df["sequence"], use_ref_graph=False)
    for i, (name, domain, data) in enumerate(
        DataLoader(design_dataset, batch_size=1, shuffle=False, num_workers=12, collate_fn=lambda lst: lst[0])
    ):
        if (i + 1) % 200 == 0:
            print(i + 1)

        if data is None:
            results.append((name, None, None, None, domain, None))
            continue

        data = data.to(device)
        network_score = proteinsolver.utils.get_node_outputs(
            net, data.x, data.edge_index, data.edge_attr, 20, output_transform="logproba"
        )
        if data.edge_index_ref is not None:
            network_score_wref = proteinsolver.utils.get_node_outputs(
                net, data.x, data.edge_index_ref, data.edge_attr_ref, 20, output_transform="logproba"
            )
        else:
            network_score_wref = None
        results.append(
            (
                name,
                network_score.mean().item(),
                network_score_wref.mean().item() if network_score_wref is not None else None,
                data.x.size(0),
                domain,
                data.identity_to_ref,
            )
        )
    columns = ["name", "network_score", "network_score_wref", "sequence_length", "domain", "identity_to_ref"]
    results_df = pd.DataFrame(results, columns=columns)
    return results_df
stability_scores["rd1"][stability_scores["rd1"]["name"] == "EEHEE_rd1_0005.pdb"]
# with open(NOTEBOOK_PATH.joinpath("stability_scores_wpreds.torch"), "wb") as fout:
#     torch.save(stability_scores_wpreds, fout)
stability_scores_wpreds = {}
for library in ["rd1", "rd2", "rd3", "rd4"]:
    print(library)
    if library in stability_scores_wpreds:
        continue
    df = apply_scores(stability_scores[library])
    stability_scores_wpreds[library] = stability_scores[library].merge(df, on="name", how="left", validate="1:1")
    display(stability_scores_wpreds[library].head())
torch.save(stability_scores_wpreds, NOTEBOOK_PATH.joinpath("stability_scores_for_selections.logproba.incremental.torch"))
# for library in ["rd1", "rd2", "rd3", "rd4"]:
#     stability_scores[library]["sum_log_prob"], stability_scores[library]["sequence_length"] = np.array(
#         [get_structure_proba(net, name) for name in stability_scores[library]["name"]]
#     ).T<jupyter_output><empty_output><jupyter_text>Correlations<jupyter_code>def get_corr_95conf(r, num):
    import math

    stderr = 1.0 / math.sqrt(num - 3)
    delta = 1.96 * stderr
    lower = math.tanh(math.atanh(r) - delta)
    upper = math.tanh(math.atanh(r) + delta)
    return r - lower, upper - r 
stability_scores["rd1"][stability_scores["rd1"]["name"] == "HHH_rd1_0995.pdb"]
results = []

for library in ["rd1", "rd2", "rd3", "rd4"]:
    df = (
        stability_scores_wpreds[library]
        .dropna(
            subset=[
                "network_score",
                "talaris2013_score",
                "betanov15_score",
                "domain",
                "stabilityscore_t",
                "stabilityscore_c",
            ]
        )
        .copy()
    )
    #     df["domain"] = df["name"].str.split("_").str[0]
    for domain, gp in df.groupby("domain"):
        for score in ["network_score", "talaris2013_score", "betanov15_score"]:
            corr_t = stats.spearmanr(gp[score] / gp["sequence_length"], gp["stabilityscore_t"])
            corr_c = stats.spearmanr(gp[score] / gp["sequence_length"], gp["stabilityscore_c"])
            corr = stats.spearmanr(
                gp[score] / gp["sequence_length"], gp[["stabilityscore_t", "stabilityscore_c"]].mean(axis=1)
            )
            corr_conf = get_corr_95conf(corr[0], len(gp[score]))
            results.append(
                (
                    library,
                    domain,
                    score,
                    corr_t[0],
                    corr_t[1],
                    corr_c[0],
                    corr_c[1],
                    corr[0],
                    corr[1],
                    corr_conf[0],
                    corr_conf[1],
                )
            )

results_df = pd.DataFrame(
    results,
    columns=[
        "library",
        "domain",
        "feature",
        "corr_t",
        "pvalue_t",
        "corr_c",
        "pvalue_c",
        "corr",
        "pvalue",
        "corr_conf_lower",
        "corr_conf_upper",
    ],
)
results_df.head()<jupyter_output><empty_output><jupyter_text>Final plot<jupyter_code>fg, axs = plt.subplots(4, 5, sharex=False, figsize=(12, 10))
domains = ["EEHEE", "EHEE", "HEEH", "HHH"]
for i, domain in enumerate(domains):
    for j, library in enumerate([None, "rd1", "rd2", "rd3", "rd4"]):
        ax = axs[i, j]
        if i == 0 and j > 0:
            ax.set_title(f"Round {j}")
        if j == 0:
            img = mpimg.imread(
                (
                    INPUT_PATH
                    / "global_analysis_of_protein_folding"
                    / "aan0693_designed-PDB-files"
                    / "nmr"
                    / f"{domain}.png"
                ).as_posix()
            )
            ax.imshow(img)
            ax.axis("off")
            ax.set_xlim(500, 3500)
            ax.set_ylim(4250, -400)
            ax.text(
                -0.25,
                0.5,
                domain,
                rotation=90,
                fontsize="x-large",
                horizontalalignment="center",
                verticalalignment="center",
                transform=ax.transAxes,
            )
        else:
            df = results_df[(results_df["domain"] == domain) & (results_df["library"] == library)]
            data = df["corr"].values
            data = [data[0]] + (-data[1:]).tolist()
            ax.bar([1, 2, 3], data, yerr=df[["corr_conf_upper", "corr_conf_lower"]].values.T)
            ax.set_ylim(-0.1, 0.5)
            ax.set_xticks([1, 2, 3])
            if j == 1:
                ax.set_ylabel("Spearman correlation")
            ax.set_xticklabels(["GCNN", "Rosetta\n(talaris)", "Rosetta\n(betanov)"], fontsize=9)
fg.tight_layout()
fg.savefig(NOTEBOOK_PATH.joinpath("correlations_with_designs.svg"), dpi=2000)
fg.savefig(NOTEBOOK_PATH.joinpath("correlations_with_designs.pdf"), dpi=2000)
fg.savefig(NOTEBOOK_PATH.joinpath("correlations_with_designs.png"), dpi=2000)
for library_name in list(stability_scores)[:4]:
    df = stability_scores[library_name][
        (stability_scores[library_name]["sum_log_prob"].isnull())
        & ((stability_scores[library_name]["talaris2013_score"].notnull()))
    ]
    print(library_name, len(df))
for library_name in list(stability_scores)[:4]:
    df = stability_scores[library_name][
        (stability_scores[library_name]["sum_log_prob"].notnull())
        & ((stability_scores[library_name]["talaris2013_score"].isnull()))
    ]
    print(library_name, len(df))
for library_name in list(stability_scores)[:4]:
    df = stability_scores[library_name][
        ((stability_scores[library_name]["talaris2013_score"].isnull()))
    ]
    print(library_name, len(df))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12567388871935714, 'fraction_numerical': 0.027616722612145254, 'mean_word_length': 2.7453333333333334, 'pattern_counts': {'":': 7, '<': 44, '<?xml version=': 0, '>': 44, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '280472', 'license': 'MIT', 'n_tokens_mistral': 7139, 'n_tokens_neox': 6603, 'n_words': 1084, 'path': 'notebooks/06_global_analysis_of_protein_folding_stability.ipynb', 'repo_name': 'ostrokach/proteinsolver', 'size': '19662'}