<jupyter_start><jupyter_text>Here we will show an example of how to run deepblast on simulation data.<jupyter_code>import os
from deepblast.sim import hmm_alignments
import argparse
import numpy as np<jupyter_output><empty_output><jupyter_text>We will first simulate multiple sequences from a single PFam family.The resulting generated alignments will be used to train the model.<jupyter_code>hmm = '../data/zf-C2H2.hmm'
n_alignments = 100
np.random.seed(0)
align_df = hmm_alignments(n=40, seed=0, n_alignments=n_alignments, hmmfile=hmm)

cols = [
    'chain1_name', 'chain2_name', 'tmscore1', 'tmscore2', 'rmsd',
    'chain1', 'chain2', 'alignment'
]
align_df.columns = cols<jupyter_output><empty_output><jupyter_text>The simulated alignments will be split into training / testing and validation.<jupyter_code>parts = n_alignments // 10
train_df = align_df.iloc[:parts * 8]
test_df = align_df.iloc[parts * 8:parts * 9]
valid_df = align_df.iloc[parts * 9:]

# save the files to disk.
train_df.to_csv('data/train.txt', sep='\t', index=None, header=None)
test_df.to_csv('data/test.txt', sep='\t', index=None, header=None)
valid_df.to_csv('data/valid.txt', sep='\t', index=None, header=None)<jupyter_output><empty_output><jupyter_text>We will prepare the environment to make sure that the appropriate output directories exist to store the results.<jupyter_code>from deepblast.trainer import LightningAligner
from pytorch_lightning import Trainer

output_dir = 'simulation_results'
if not os.path.exists(output_dir):
    os.mkdir(output_dir)
os.getcwd()<jupyter_output><empty_output><jupyter_text>We will now create the arguments.  Below is the way to create this in a python environment.This can also be recreated on a standard command line interface.<jupyter_code>args = [
    '--train-pairs', f'{os.getcwd()}/data/train.txt',
    '--test-pairs', f'{os.getcwd()}/data/test.txt',
    '--valid-pairs', f'{os.getcwd()}/data/valid.txt',
    '--output-directory', output_dir,
    '--epochs', '32',
    '--batch-size', '20',   
    '--num-workers', '30',
    '--learning-rate', '1e-4',
    '--layers', '2',
    '--visualization-fraction', '1',
    '--loss', 'cross_entropy',
    '--scheduler', 'cosine',
    '--gpus', '1'
]
parser = argparse.ArgumentParser(add_help=False)
parser = LightningAligner.add_model_specific_args(parser)
parser.add_argument('--num-workers', type=int)
parser.add_argument('--gpus', type=int)
args = parser.parse_args(args)
%tb<jupyter_output>No traceback available to show.<jupyter_text>We will then initialize the alignment model with the parameters we specified earlier.<jupyter_code>model = LightningAligner(args)<jupyter_output><empty_output><jupyter_text>We can now train the model.<jupyter_code>trainer = Trainer(
    max_epochs=args.epochs,
    gpus=args.gpus,
    check_val_every_n_epoch=1,
    # profiler=profiler,
    fast_dev_run=False,
    # auto_scale_batch_size='power'
)

trainer.fit(model)<jupyter_output>GPU available: True, used: True
TPU available: False, using: 0 TPU cores
CUDA_VISIBLE_DEVICES: [0]
/home/juermieboop/miniconda3/envs/pytorch/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:37: UserWarning: Could not log computational graph since the `model.example_input_array` attribute is not set or `input_array` was not given
  warnings.warn(*args, **kwargs)

  | Name    | Type                   | Params
---------------------------------------------------
0 | aligner | NeedlemanWunschAligner | 34 M<jupyter_text>The model diagnostics can be directly visualized in Tensorboard. Here, we show the losses, the accuracy and the alignment results.<jupyter_code>!ls lightning_logs
%load_ext tensorboard
%tensorboard --logdir lightning_logs<jupyter_output><empty_output><jupyter_text>And we did this with just a few million parameters<jupyter_code>model.aligner
!ls lightning_logs/version_29/checkpoints
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
checkpoint_dir = 'lightning_logs/version_29/checkpoints'
path = f'{checkpoint_dir}/epoch=25.ckpt'
model = LightningAligner.load_from_checkpoint(path)
from deepblast.dataset.dataset import states2matrix, tmstate_f
import numpy as np
import torch
from torch.nn.utils.rnn import pack_sequence
import matplotlib.pyplot as plt
import seaborn as sns
i = 9
test = align_df
x = str.encode(test.iloc[i]['chain1'])
y = str.encode(test.iloc[i]['chain2'])
states = test.iloc[i]['alignment']
x_ = torch.Tensor(model.tokenizer(x)).long().unsqueeze(0)
y_ = torch.Tensor(model.tokenizer(y)).long().unsqueeze(0)
x_ = pack_sequence(x_, enforce_sorted=False).cuda()
y_ = pack_sequence(y_, enforce_sorted=False).cuda()
states = list(map(tmstate_f, states))
A = states2matrix(states)
model = model.cuda()
aln, theta, gap = model.forward(x_, y_)
fig, ax = plt.subplots(1, 4, figsize=(16, 3))
sns.heatmap(A, ax=ax[0])
sns.heatmap(aln.cpu().detach().numpy().squeeze(), ax=ax[1])
sns.heatmap(np.log(theta.cpu().detach().numpy().squeeze()), ax=ax[2])
sns.heatmap(gap.cpu().detach().numpy().squeeze(), ax=ax[3])
from deepblast.dataset.alphabet import Uniprot21
from deepblast.dataset.dataset import decode
from deepblast.score import alignment_text, roc_edges
from deepblast.dataset.dataset import states2edges

#pred_edges = list(zip(pred_y, pred_x))
pred_edges = states2edges(pred_states)
true_edges = states2edges(truth_states)
stats = roc_edges(true_edges, pred_edges)
gen = model.aligner.traceback(x_, y_)
decoded, _ = next(gen)
pred_x, pred_y, pred_states = list(zip(*decoded))
pred_states = np.array(list(pred_states))
truth_states = np.array(list(states))                    
text = alignment_text(x.decode('utf-8'), y.decode('utf-8'), pred_states, truth_states, stats)
print(text)
px, py = zip(*pred_edges)
tx, ty = zip(*true_edges)

plt.plot(px, py, label='pred')
plt.plot(tx, ty, label='truth')
from deepblast.dataset.dataset import state_diff_f
def states2edges(states):
    """ Converts state string to bipartite matching. """
    prev_s, next_s = states[:-1], states[1:]
    transitions = list(zip(prev_s, next_s))
    state_diffs = np.array(list(map(state_diff_f, transitions)))
    coords = np.cumsum(state_diffs, axis=0).tolist()
    coords = [(0, 0)] + list(map(tuple, coords))
    return coords
pe = states2edges(list(pred_states))
te = states2edges(list(truth_states))
px, py = zip(*pe)
tx, ty = zip(*te)

plt.plot(px, py, label='pred')
plt.plot(tx, ty, label='truth')
pe
te
truth_states
pred_states<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14239281844915647, 'fraction_numerical': 0.012381984212970128, 'mean_word_length': 5.455544455544455, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 35, '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': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '193135', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 2348, 'n_tokens_neox': 2272, 'n_words': 496, 'path': 'ipynb/simulation-benchmark.ipynb', 'repo_name': 'wouterboomsma/deepblast', 'size': '6461'}