lucidrains/alphafold2

To eventually become an unofficial Pytorch implementation / replication of Alphafold2, as details of the architecture get released

alphafold2's Issues

Question: Why are lengths of seq and msa same and different?

Hello, the lengths of seq and msa are same and different in different examples in the readme.
Am I not understanding something?

Thanks,

MSA attention gated problem

Hi Phil:

Compare the 'gating' unit with the original algorithm:

alphafold2/alphafold2_pytorch/alphafold2.py

Line 147 in 586792d

gates = self.gating(x)

I thought gating here just missed a 'sigmoid()', please have a check.

Best,

Can AlphaFold2 predict the structure of cyclic peptides?

I have a cyclic peptide sequence. I put it into Alphafold2 Colab, but I didn't get a cyclic peptide structure. What should I do to connect the C-terminal and N-terminal for the next dynamic simulation (GROMACS)?

Should I do some processing on the structure obtained by Alphafold2 Colab for Gromacs dynamics? Or input sequence to Alphafold2 Colab need for some pre-processing? Or can Alphafold2 not predict the sequence of cyclic peptides?

It's really important for my research, thanks for any help.

Typo?

Hi Phil and Eric @hypnopump

I'm not familiar with the operations on protein structure or whatever :(, but here I wonder is it a mistake when indexing:

alphafold2/alphafold2_pytorch/utils.py

Line 702 in 5214156

new_coords[:, :-padding or None, i] = coords[:, i]

Perhaps it should be new_coords[s,:-padding or None, i] since in the inner loop it just constructs the s-th sequence sidechain's coords.

Best,

center_distogram_torch breaks with center = 'median'

the dimensions of coordinates in CASP12

@lucidrains Hi! In the test_coor, there is code like: assert coords.shape == (2, 16 * 14, 3), 'must output coordinates' . What does “14“ here mean。In the model, num_backbone_atoms = 3, shouldn’t the coordinate dimension be（2,16 * 3, 3）。

[todo] embeddings

need to build the embedding module

[todo] recycling

recycling needs to be built

ModuleNotFoundError: No module named 'invariant_point_attention'

ModuleNotFoundError: No module named 'invariant_point_attention'

Should be first MSA-attention then Pairwise-Attention

Here, original evoformer should be first MSA-attention then Pairwise-Attention:

alphafold2/alphafold2_pytorch/alphafold2.py

Lines 422 to 428 in 2c57480

	for _ in range(depth):
	self.layers.append(nn.ModuleList([
	PairwiseAttentionBlock(dim = dim, seq_len = seq_len, heads = heads, dim_head = dim_head, dropout = attn_dropout, global_column_attn = global_column_attn),
	FeedForward(dim = dim, dropout = ff_dropout),
	MsaAttentionBlock(dim = dim, seq_len = seq_len, heads = heads, dim_head = dim_head, dropout = attn_dropout),
	FeedForward(dim = dim, dropout = ff_dropout),
	]))

plz have a check

Best

How can I get the coordinates of all the atoms?

@lucidrains ,HI,
I read this quote in Readme, set the output atom type by setting the" Atoms" .

Valid choices for atoms include:
backbone - 3 backbone atoms (C, Ca, N) [default]
backbone-with-cbeta - 3 backbone atoms and C beta
backbone-with-oxygen - 3 backbone atoms and oxygen from carboxyl
backbone-with-cbeta-and-oxygen - 3 backbone atoms with C beta and oxygen
all - backbone and all other atoms from sidechain
You can also pass in a tensor of shape (14,) defining which atoms you would like to include

When I set this parameter, I got an error like that,

model = Alphafold2(
    dim = 256,
    depth = 1,
    heads = 8,
    dim_head = 64,
    predict_coords = True,
    # structure_module_dim = 8,
    structure_module_depth = 2,
    structure_module_heads = 4,
    structure_module_dim_head = 16,
    atoms = "all"
    # structure_module_refinement_iters = 2
).to(DEVICE)

Traceback (most recent call last):
  File "/tmp/pycharm_project_152/train_end2end.py", line 99, in <module>
    model = Alphafold2(
TypeError: __init__() got an unexpected keyword argument 'atoms'

Process finished with exit code 1

When I remove this parameter, I can only get the coordinates of one type atom in the sequence in refined.

    refined = model(
            seq,
            msa = msa,
            embedds = embedds,
            mask = mask
        )

Is this atom a CA atom?
#How can I get the coordinates of all the atoms, what parameters should I set, or which version of the previous version can I get ？

Missing adding the `relpos` to pairwise-repr `x`

Hi,

I thought here might miss adding the rel_pos_emb to x

alphafold2/alphafold2_pytorch/alphafold2.py

Line 705 in c42d432

rel_pos_emb = self.pos_emb(seq_rel_dist)

Best,

How datasets are load ? And how to calculate GDT ?

Output pdb file

What can I do if I want to input protein sequences and output pdb files based on your code?

read me dimensions are off. Example to run alpha fold 2 msa, and seq lengths should be the same. One is 128, the other is 120.

Add PL Lightning to Enable Distributed Training and Deep Speed

This work is awesome I see that PL is already being used for some of the DataModules it would be awesome to see Lightning Module Integration to make training more robust. https://pytorch-lightning.readthedocs.io/en/stable/starter/converting.html

Basic pretrained models + quickstart instructions?

Hi there,

Practitioner of several years with deep experience (not much GH footprint in terms of recent experience) in model architecture selection, training, etc. This looks like a fun enough problem for me to work on, not definitely confirmed on it, but looking on potentially picking this up as a side project from when I have spare 'problem-solving' cycles to burn on things for funsies.

I guess in looking through everything there's a possibility to kick everything off from scratch, but I'm a fan of pretrained models, or at least partially-trained models in some form. Are there any S3 dumps of model checkpoints we can finetune from?

Also, in scanning the README, my old training setup is rather aged -- 1070, so the TPU/V100 space on CoLab is probably my go to for now. Do you have a local setup, conda env, docker container, etc that you're working in? Looking to see if there's anything I need to do to standardize the path on my end.

Thanks for putting all of this together. Been working on a few semi-ish related things (not too close, but semi) for a little while, and been watching this repo. Glad to see it's alive and still moving.

If it's of any value, are there branch privs on this repo, and is there any process I should go through to getting those? I'm not looking for anything beyond being able to create/push to individual remote branches, I'd like the PR process and everything that comes with it.

Many thanks, hope you are well,

EF

How

How can one man make such advanced implementations at such speed?

pairwise_repr should be used in rowAttention

Hi, Phil

alphafold2/alphafold2_pytorch/alphafold2.py

Lines 400 to 401 in a783528

	x = self.row_attn(x, mask = mask) + x
	x = self.col_attn(x, mask = mask, edges = pairwise_repr) + x

according to the original paper, (Algorithm 6, line2-4), the pairwise_repr(bias) used in MSARowAttentionWithPairBias instead of MSAColumnAttention. plz have a check

Best,

outerMean module question

Hi, I found that the outerMean module is different from Deepmind alphafold, expecially here.
In DeepMind code, the output is like [seq,seq,dim]has no msa_num, but in your code the result maybe [batch, msa_num, seq,seq,dim]. Is this a bug?

[todo] structure module

in the structure module, the FAPE loss needs to be applied on every iteration, while the rotation needs to have a stop-gradient for all but the last iteration

Example script of an end-to-end prediction from FASTA sequence to 3D pdb

Thanks for your codes. I just wonder if you could provide a script of the prediction process with a FASTA sequence as input and predicted PDB file as output; if not just give me some hint please. Thanks in advance.

the sidechainnet pkg updates, the load function doesnot has parameter return_masks and the format of casp12 data from scn has changed, which is not compatible with train_end2end.py

alphafold2/train_end2end.py

Line 71 in b7ac212

return_masks = True

How to adjust the recycling times in model?

I just found the function of "add recyclables"(line 728 in alphafold.py), but how to adjust the recycling times during the forward function?

Maybe duplicate residual addition

Hi, Phil

I checked the code part here, but I found maybe it adds more than need.

Original:

alphafold2/alphafold2_pytorch/alphafold2.py

Lines 437 to 448 in a783528

	for attn, ff, msa_attn, msa_ff in self.layers:
	# msa attention and transition
	m = msa_attn(m, mask = msa_mask, pairwise_repr = x) + m
	m = msa_ff(m) + m
	# pairwise attention and transition
	x = attn(x, mask = mask) + x
	x = ff(x) + x
	return x, m

I thought it should be:
m = msa_attn(m, mask = msa_mask, pairwise_repr = x)
m = msa_ff(m) + m
x = attn(x, mask = mask, msa_repr = m)
x = ff(x) + x

Since, in the attn_block, the return tensors have already been added with the residual part.

plz have a check

Best

[todo] positional encoding

either use the RPE scheme as in the paper, or see if one can use rotary embeddings where possible, and T5 RPE bias in other places

Architecture

Share what you think you know about the architecture here

MSA tensor format

In the Usage, there is code like
seq = torch.randint(0, 21, (1, 128)).cuda()
msa = torch.randint(0, 21, (1, 5, 64)).cuda().
If I have a a3m msa file, how to encode the file to this tensor? And why the seq length is 128 but the msa is 5 times 64 (5 timeshalf the length of seq?).
Could you give an example of how to use that or how to generate that msa tensor?

test error in utils

MDScaling isn't working right!

You can easily find the issue when rerun the "structure_utils_tests.ipynb". You can't reconstruct the distance matrix.

[todo] FAPE loss

fape loss needs to be added

no coords_3d in train_end2end.py

AttributeError : Constant object has no attribute ‘kind’

• I used covid protein ‘nsp2.pdb’ file.(In this “ structure_utils_tests.ipynb” test file)
• In this part occured error

N_mask  = torch.tensor( prot_traj.topology.select("name == N and backbone") ).unsqueeze(0)
CA_mask = torch.tensor( prot_traj.topology.select("name == CA and backbone") ).unsqueeze(0) 
C_mask  = torch.tensor( prot_traj.topology.select("name == C and backbone")  ).unsqueeze(0)
CA_mask.shape, N_mask.shape, C_mask.shape

• Error Log

   549             self.write("...")
    550         else:
--> 551             if t.kind == "u":
    552                 self.write("u")
    553             self._write_constant(t.value)

AttributeError: 'Constant' object has no attribute 'kind'

• Please give me solutions.

The definition of row-wise attention and col-wise attention

Hi, Luci:
Sorry, it's me again.
Here, I was confused by the definition of row-wise and col-wise attention.

alphafold2/alphafold2_pytorch/alphafold2.py

Lines 208 to 209 in 586792d

	if self.row_attn:
	w_x = rearrange(x, 'b h w d -> (b w) h d')

alphafold2/alphafold2_pytorch/alphafold2.py

Lines 219 to 220 in 586792d

	if self.col_attn:
	h_x = rearrange(x, 'b h w d -> (b h) w d')

Based on what I thought, the row w_x should be represented by (b h) w d, since once fetch one row, each row should have w(width) units.

So, maybe here need an inverse of the above definition?

__init__() got an unexpected keyword argument 'structure_module_type'

When I run

import torch
from alphafold2_pytorch import Alphafold2

model = Alphafold2(
    dim = 256,
    depth = 2,
    heads = 8,
    dim_head = 64,
    predict_coords = True,
    structure_module_type = 'se3',          # use SE3 Transformer - if set to False, will use E(n)-Transformer, Victor and Max Welling's new paper
    structure_module_dim = 4,               # se3 transformer dimension
    structure_module_depth = 1,             # depth
    structure_module_heads = 1,             # heads
    structure_module_dim_head = 16,         # dimension of heads
    structure_module_refinement_iters = 2,  # number of equivariant coordinate refinement iterations
    structure_num_global_nodes = 1          # number of global nodes for the structure module, only works with SE3 transformer
)

I will get ERROR that TypeError: __init__() got an unexpected keyword argument 'structure_module_type'.
and I found that MSA have to be the same length as primary sequence.

OuterMean block's typo

Hi, Luci:
Here, the original paper uses outer-product, but here uses outer-sum

alphafold2/alphafold2_pytorch/alphafold2.py

Line 316 in 586792d

outer = rearrange(left, 'b m i d -> b m i () d') + rearrange(right, 'b m j d -> b m () j d')

plz have a check, many thanks!

[todo] ensembling

TypeError: __init__() got an unexpected keyword argument 'num_backbone_atoms'

alphafold2_pytorch-0.4.15

center_distogram_torch with center = 'median' breaks when tensors are on CUDA device

I want use casp13 dataset.

• I want use casp13 dataset.
• And i want use covid-13 protein pdb file.(In this project)
• Please give me solutions.

[todo] BERT loss on MSA

needs to be added

Link to discord?

Hi, I would like to follow the development of this project on discord but the link does not seem to be working.
Could someone verify?
Thank you.

[todo] checkpointing

make sure checkpointing works, or see if evoformer can be made reversible

Maybe Inverse of the outgoing-attention and ingoing-attention here

Hi Phil:

Haha, one similar issue here, I thought the outgoing-attention[ = paper's triangular gated self-attn around starting node] should use row-attention since we start from a node i and propagate information to j and k. So, it's ij,ik ->jk, maybe we should use row-attention.
Likewise, in the ingoing-attention, we should use col-attention (ij,kj->ki)

alphafold2/alphafold2_pytorch/alphafold2.py

Lines 361 to 362 in a783528

	self.triangle_attention_ingoing = AxialAttention(dim = dim, heads = heads, dim_head = dim_head, row_attn = True, col_attn = False, accept_edges = True)
	self.triangle_attention_outgoing = AxialAttention(dim = dim, heads = heads, dim_head = dim_head, row_attn = False, col_attn = True, accept_edges = True, global_query_attn = global_column_attn)

Best,

How to use casp13 dataset & covid-13 protein sequence on this project.

• I want use casp13 dataset.
• And i want use covid-13 protein pdb file.(In this project)
• Please give me solutions.

problem with alphafold2/train_pre.py

here is the problem when running alphafold2/train_pre.py:

/usr/local/lib/python3.7/dist-packages/torch/_tensor.py:575: UserWarning: floor_divide is deprecated, and will be removed in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values.
To keep the current behavior, use torch.div(a, b, rounding_mode='trunc'), or for actual floor division, use torch.div(a, b, rounding_mode='floor'). (Triggered internally at /pytorch/aten/src/ATen/native/BinaryOps.cpp:467.)
return torch.floor_divide(self, other)

RuntimeError Traceback (most recent call last)
in ()
114 distogram,
115 discretized_distances,
--> 116 ignore_index = IGNORE_INDEX
117 )
118

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
2822 if size_average is not None or reduce is not None:
2823 reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 2824 return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
2825
2826

RuntimeError: input and target batch or spatial sizes don't match: target [1 x 52 x 52], input [1 x 256 x 156 x 156] at /pytorch/aten/src/THCUNN/generic/SpatialClassNLLCriterion.cu:23

CUBLAS ERROR while running example

When running the first example, prediction to distogram, I get the following error.
Is this error reproducible for you guys?
It looks like its crashing in the trunk part for the attention qkv remapping.

torch 1.8.0
cuda 10.2
python 3.6.9

Traceback (most recent call last):
  File "test.py", line 25, in <module>
    msa_mask = msa_mask
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ychnh/af/alphafold2_pytorch/alphafold2.py", line 840, in forward
    msa_mask = msa_mask
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ychnh/af/alphafold2_pytorch/alphafold2.py", line 511, in forward
    x = attn(x, shape = seq_shape, mask = mask) + x
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ychnh/af/alphafold2_pytorch/alphafold2.py", line 45, in forward
    return self.fn(x, *args, **kwargs)
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ychnh/af/alphafold2_pytorch/alphafold2.py", line 107, in forward
    return attn(x, *args, shape = shape, mask = mask, **kwargs)
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ychnh/af/alphafold2_pytorch/alphafold2.py", line 411, in forward
    w_out = self.attn_width(w_x, mask = w_mask)
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ychnh/af/alphafold2_pytorch/alphafold2.py", line 235, in forward
    q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim = -1))
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/modules/linear.py", line 94, in forward
    return F.linear(input, self.weight, self.bias)
  File "/home/ychnh/.local/lib/python3.6/site-packages/torch/nn/functional.py", line 1753, in linear
    return torch._C._nn.linear(input, weight, bias)

RuntimeError: CUDA error: CUBLAS_STATUS_INTERNAL_ERROR when calling `cublasCreate(handle)`

A bug……

When I test the example code of usage, there is a "RuntimeError", it says "input value can be a scalar only when boundaries tensor dimension is 1, but we got boundaries tensor dim(4) and input value's dim(0) numel(1).", which confuses me.

Does anyone have a similar problem, if anyone can help me, I would be very grateful.