|
''' |
|
Infers average ELBO values from DeepSequence VAE models. |
|
Based on open source code from DeepSequence repo. |
|
''' |
|
import argparse |
|
import numpy as np |
|
import os |
|
import pathlib |
|
from shutil import copyfile |
|
import sys |
|
import time |
|
|
|
import utils |
|
|
|
WORKING_DIR="/share/vault/Users/gz2294/combining-evolutionary-and-assay-labelled-data/DeepSequence/" |
|
N_ELBO_SAMPLES=400 |
|
|
|
module_path = os.path.abspath(WORKING_DIR) |
|
if module_path not in sys.path: |
|
sys.path.append(module_path) |
|
|
|
from DeepSequence.model import VariationalAutoencoder |
|
from DeepSequence import helper |
|
from DeepSequence import train |
|
|
|
model_params = { |
|
"bs" : 100, |
|
"encode_dim_zero" : 1500, |
|
"encode_dim_one" : 1500, |
|
"decode_dim_zero" : 100, |
|
"decode_dim_one" : 500, |
|
"n_latent" : 30, |
|
"logit_p" : 0.001, |
|
"sparsity" : "logit", |
|
"final_decode_nonlin": "sigmoid", |
|
"final_pwm_scale" : True, |
|
"n_pat" : 4, |
|
"r_seed" : 12345, |
|
"conv_pat" : True, |
|
"d_c_size" : 40 |
|
} |
|
|
|
|
|
def main(): |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument("model_prefix", type=str) |
|
parser.add_argument("fasta_file", type=str) |
|
parser.add_argument("wt_fasta_file", type=str) |
|
parser.add_argument("output_dir", type=str) |
|
args = parser.parse_args() |
|
data_helper = helper.DataHelper( |
|
working_dir=WORKING_DIR, |
|
alignment_file=args.fasta_file, |
|
calc_weights=False, |
|
) |
|
|
|
vae_model = VariationalAutoencoder(data_helper, |
|
batch_size = model_params["bs"], |
|
encoder_architecture = [model_params["encode_dim_zero"], |
|
model_params["encode_dim_one"]], |
|
decoder_architecture = [model_params["decode_dim_zero"], |
|
model_params["decode_dim_one"]], |
|
n_latent = model_params["n_latent"], |
|
logit_p = model_params["logit_p"], |
|
sparsity = model_params["sparsity"], |
|
encode_nonlinearity_type = "relu", |
|
decode_nonlinearity_type = "relu", |
|
final_decode_nonlinearity = model_params["final_decode_nonlin"], |
|
final_pwm_scale = model_params["final_pwm_scale"], |
|
conv_decoder_size = model_params["d_c_size"], |
|
convolve_patterns = model_params["conv_pat"], |
|
n_patterns = model_params["n_pat"], |
|
random_seed = model_params["r_seed"], |
|
working_dir = WORKING_DIR, |
|
) |
|
vae_model.load_parameters(args.model_prefix) |
|
|
|
pathlib.Path(args.output_dir).mkdir(parents=True, exist_ok=True) |
|
|
|
focuscols = set(data_helper.uniprot_focus_cols_list) |
|
|
|
seqs = utils.read_fasta(os.path.join(WORKING_DIR, "datasets", |
|
args.fasta_file)) |
|
wt, des = utils.read_fasta(args.wt_fasta_file, return_ids=True) |
|
wt = wt[0] |
|
des = des[0] |
|
offset = int(des.split('/')[-1].split('-')[0]) |
|
delta_elbos = np.zeros(len(seqs)) |
|
for i, s in enumerate(seqs): |
|
if i % 100 == 0: |
|
|
|
np.savetxt(os.path.join(args.output_dir, "elbo.npy"), delta_elbos) |
|
mut_tups = utils.seq2mutation_fromwt(s, wt, offset=offset) |
|
mut_tups = [t for t in mut_tups if t[0] in focuscols] |
|
delta_elbos[i] = data_helper.delta_elbo(vae_model, mut_tups, |
|
N_pred_iterations=N_ELBO_SAMPLES) |
|
np.savetxt(os.path.join(args.output_dir, "elbo.npy"), delta_elbos) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
