Loss Function for AminoBret

you can test it by running following

vocab_size = 24
bs = 20
logit_out = torch.rand(bs, 10, vocab_size)
logit_chunk_perm = torch.rand(bs, 2)
target = torch.randint(1, 20, (bs, 10))
chunk_perm = torch.randint(0, 2, (bs,))

loss_func = AminoBretLoss(vocab_size=vocab_size)

loss = loss_func(logit_out, logit_chunk_perm, target, chunk_perm)
print(loss)

Adding two functions for protein structure refinement in PyRosetta:

quick_refine: refines full-atom structure in cartesian coordinates using MinMover

• followed by idealization of bond lengths and angles

relax_refine: refines full-atom structure in both cartesian and internal coordinates using FastRelax

• structures are minimized before and after relaxation with MinMover (similar to quick_refine)
• followed by idealization of bond lengths and angles

To use these functions, obtain a license and install PyRosetta from here. Before calling either function, you must initialize PyRosetta like:

import pyrosetta
pyrosetta.init("-mute all")

PyRosetta is quite verbose so I usually mute outputs, but that can be removed for debugging.

I want to train the rgn2-replica model with the following command:

export WANDB_MODE=offline
python scripts/train_rgn2.py \
    --device cuda:1 \
    --wb_entity xxxxx \
    --wb_proj rgn2 \
    --run_name RGN2X_vanillaLSTM_full_run \
    --min_len_valid 0 
    --xray 1

But an AttributeError occurred:

Traceback (most recent call last):
  File "scripts/train_rgn2.py", line 345, in <module>
    init_and_train(args)
  File "scripts/train_rgn2.py", line 134, in init_and_train
    results = run_train_schedule(dataloaders, embedder, config, args)
  File "scripts/train_rgn2.py", line 213, in run_train_schedule
    config=config,
  File "/home/lipan/usr/rgn2-replica/rgn2_replica/rgn2_trainers.py", line 432, in train
    config=config,
  File "/home/lipan/usr/rgn2-replica/rgn2_replica/rgn2_trainers.py", line 62, in batched_inference
    coords_rebuilt = mp_nerf.proteins.ca_bb_fold( ca_trace ) # beware extremes
AttributeError: module 'mp_nerf.proteins' has no attribute 'ca_bb_fold'

It seems that mp_nerf.proteins.ca_bb_fold is not defined in https://github.com/eleutherAI/mp_nerf

What?

Instead of En-transformer refiner, IPA (https://github.com/lucidrains/invariant-point-attention) refiner has been integrated.

Why?

To see the effectiveness of different experiments. IPA has lied in the core of AF2 structure module. It is worthy to give a try.

How?

Simply follow existing RGN2_transformer implementation. The difference is output of coordinates instead of angles. However, the compatibility of angles (used by losses) has been done via discrete Frenet-Serret equations.

Testing?

dRMSD can reach around 5. Experimental hyper-para's include "structure_module_depth" and other IPA parameters. Running script has been included.

Screenshots

At the moment the xtension branch of mp_nerf repo has been integrated via source code. Some functions (e.g. mp_nerf.ml_utils.backbone_forcefield()) has not been implemented. Not a big deal. I leave it for further improvement.

Anything Else?

Integrate LSTM output of 4 angles has been discussed but not implemented. LSTM can give nice prediction of secondary structures. Worthy to give a trial.

Usage example:

sys.argv = [
    'scripts/rgn2_predict_fold.py',
    
    '--pdb_input', '/home/hypnopump/bruba/rgn2-replica/data/input/0_lolasso.pdb',
    # --- OR ---
#     '--input', '/home/hypnopump/own_research/notebooks/test_predictor/proteins.fasta',
#     '--model', '/home/hypnopump/own_research/notebooks/rgn2_models/restart_improve_vanillaLSTM_improve@_150K.pt',
    
    '--bidirectional', 'True',
    '--output_path', '/home/hypnopump/bruba/rgn2-replica/data/output',
    '--device', 'cpu',
    '--af2_refine', 'True'
]

from scripts.rgn2_predict_fold import predict_refine
result = predict_refine()

• Prediction of angular features:

  • See this: https://discord.com/channels/729741769192767510/797547607345201162/874987318069592095

• Types of recurrent layers

  • multiplicativeLSTM: https://florianwilhelm.info/2018/08/multiplicative_LSTM_for_sequence_based_recos/
  • layerNormLSTM: https://github.com/exe1023/LSTM_LN/blob/master/lstm.py
  • multiplicative Integration: https://papers.nips.cc/paper/2016/file/f69e505b08403ad2298b9f262659929a-Paper.pdf
  • other variants: https://github.com/FlorianMuellerklein/death_metal_lstm/blob/master/scripts/models.py
  • Novel Arch: https://arxiv.org/pdf/1911.11033.pdf
  • SRU (https://arxiv.org/abs/1709.02755): https://github.com/bamtercelboo/pytorch_SRU
  • peephole: https://towardsdatascience.com/building-a-lstm-by-hand-on-pytorch-59c02a4ec091
  • uRNN: https://web.stanford.edu/class/archive/cs/cs224n/cs224n.1194/reports/custom/15842215.pdf
  • delayedLSTM: https://arxiv.org/pdf/1909.00021.pdf
  • osiparc: ?