train.py

from argparse import ArgumentParser
import os
from pathlib import Path
import random
import string
import io
import pickle as pkl

import boto3
import numpy as np
import matplotlib.pyplot as plt
import pytorch_lightning as pl
import torch
import wandb

from mogwai.data_loading import MSADataModule, MSDataModule
from mogwai.parsing import read_contacts
from mogwai import models
from mogwai.utils.functional import apc
from mogwai.metrics import contact_auc
from mogwai.plotting import (
    plot_colored_preds_on_trues,
    plot_precision_vs_length,
)
from mogwai.vocab import FastaVocab

from loggers import WandbLoggerFrozenVal

s3_client = boto3.client("s3")
s3_bucket = "songlabdata"



def torch_to_numpy(state_dict, keys=['weight', 'bias', '_true_contacts', '_max_auc']):
    numpy_dict = dict()
    for key in keys:
        numpy_dict[key] = state_dict[key].numpy()
    return numpy_dict


def train():
    # Initialize parser
    parser = ArgumentParser()
    parser.add_argument(
        "--model",
        default="gremlin",
        choices=models.MODELS.keys(),
        help="Which model to train.",
    )
    parser.add_argument(
        "--train_unaligned",
        action="store_true",
        help="Whether to train unaligned instead.",
    )
    model_name = parser.parse_known_args()[0].model
    train_unaligned = parser.parse_known_args()[0].train_unaligned
    parser.add_argument(
        "--save_model_s3",
        action="store_true",
        help="Whether to save the model state dict.",
    )
    parser.add_argument(
        "--wandb_project",
        type=str,
        default="synthetic-protein-landscapes",
        help="W&B project used for logging.",
    )
    parser.add_argument(
        "--pdb",
        type=str,
        help="PDB id for training",
    )

    if train_unaligned:
        parser = MSDataModule.add_args(parser)
    else:
        parser = MSADataModule.add_args(parser)

    parser = pl.Trainer.add_argparse_args(parser)
    parser.set_defaults(
        gpus=1,
        min_steps=50,
        max_steps=1000,
        log_every_n_steps=10,
    )
    model_type = models.get(model_name)
    model_type.add_args(parser)
    args = parser.parse_args()

    # Modify name
    pdb = args.pdb
    args.data = "data/npz/" + pdb + ".npz"

    # Load ms(a)
    if train_unaligned:
        msa_dm = MSDataModule.from_args(args)
    else:
        msa_dm = MSADataModule.from_args(args)
    msa_dm.setup()

    # Load contacts
    true_contacts = torch.from_numpy(read_contacts(args.data))

    # Initialize model
    num_seqs, msa_length, msa_counts = msa_dm.get_stats()
    model = model_type.from_args(
        args,
        num_seqs=num_seqs,
        msa_length=msa_length,
        msa_counts=msa_counts,
        vocab_size=len(FastaVocab),
        pad_idx=FastaVocab.pad_idx,
        true_contacts=true_contacts,
    )

    kwargs = {}
    randstring = "".join(random.choice(string.ascii_lowercase) for i in range(6))
    run_name = "_".join([args.model, pdb, randstring])
    logger = WandbLoggerFrozenVal(project=args.wandb_project, name=run_name)
    logger.log_hyperparams(args)
    logger.log_hyperparams(
        {
            "pdb": pdb,
            "num_seqs": num_seqs,
            "msa_length": msa_length,
        }
    )
    kwargs["logger"] = logger

    # Initialize Trainer
    trainer = pl.Trainer.from_argparse_args(args, checkpoint_callback=False, **kwargs)

    trainer.fit(model, msa_dm)

    # Log and print some metrics after training.
    contacts = model.get_contacts()
    apc_contacts = apc(contacts)

    auc = contact_auc(contacts, true_contacts).item()
    auc_apc = contact_auc(apc_contacts, true_contacts).item()
    print(f"AUC: {auc:0.3f}, AUC_APC: {auc_apc:0.3f}")

    filename = "top_L_contacts.png"
    plot_colored_preds_on_trues(contacts, true_contacts, point_size=5, cutoff=1)
    plt.title(f"Top L no APC {model.get_precision(do_apc=False)}")
    logger.log_metrics({filename: wandb.Image(plt)})
    plt.close()

    filename = "top_L_contacts_apc.png"
    plot_colored_preds_on_trues(apc_contacts, true_contacts, point_size=5, cutoff=1)
    plt.title(f"Top L APC {model.get_precision(do_apc=True)}")
    logger.log_metrics({filename: wandb.Image(plt)})
    plt.close()

    filename = "top_L_5_contacts.png"
    plot_colored_preds_on_trues(contacts, true_contacts, point_size=5, cutoff=5)
    plt.title(f"Top L/5 no APC {model.get_precision(do_apc=False, cutoff=5)}")
    logger.log_metrics({filename: wandb.Image(plt)})
    plt.close()

    filename = "top_L_5_contacts_apc.png"
    plot_colored_preds_on_trues(apc_contacts, true_contacts, point_size=5, cutoff=5)
    plt.title(f"Top L/5 APC {model.get_precision(do_apc=True, cutoff=5)}")
    logger.log_metrics({filename: wandb.Image(plt)})
    plt.close()

    filename = "precision_vs_L.png"
    plot_precision_vs_length(apc_contacts, true_contacts)
    logger.log_metrics({filename: wandb.Image(plt)})
    plt.close()

    if args.save_model_s3:
        bytestream = io.BytesIO()
        model_dict = torch_to_numpy(model.state_dict())
        model_dict['precision_at_l'] = model.get_precision(do_apc=True).numpy()

        # add query sequence
        with open(args.data, 'rb') as f:
            msa_raw = np.load(f)['msa']
        model_dict['query_seq'] = msa_raw[0]

        np.savez(bytestream, **model_dict)
        bytestream.seek(0)

        key = os.path.join(
            "proteindata", "synthetic-protein-landscapes", wandb.run.path, "{}_model_state_dict.npz".format(pdb)
        )
        response = s3_client.put_object(
            Bucket=s3_bucket, Body=bytestream, Key=key, ACL="public-read"
        )
        print(f"uploaded state dict to s3://{s3_bucket}/{key}")

if __name__ == "__main__":
    train()