fix index for linking to terminal atom of previous biomolecule in chain

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -5,7 +5,6 @@
 from functools import partial, wraps
 from collections import namedtuple
 
-import einops
 import torch
 from torch import nn
 from torch import Tensor

alphafold3_pytorch/inputs.py

@@ -362,8 +362,10 @@ def molecule_to_atom_input(
                 _, last_atom_index = prev_src_tgt_atom_indices
                 first_atom_index, _ = src_tgt_atom_indices
 
-                src_atom_offset = offset + first_atom_index
-                tgt_atom_offset = offset - last_atom_index
+                last_atom_index_from_end = prev_mol.GetNumAtoms() - last_atom_index
+
+                src_atom_offset = offset - last_atom_index_from_end
+                tgt_atom_offset = offset + first_atom_index
 
                 atompair_ids[src_atom_offset, tgt_atom_offset] = 1
                 atompair_ids[tgt_atom_offset, src_atom_offset] = 1
@@ -671,8 +673,6 @@ def alphafold3_input_to_molecule_input(
         *metal_ions_pool_lens
     ]
 
-    total_atoms = sum(token_pool_lens)
-
     # construct the token bonds
 
     # will be linearly connected for proteins and nucleic acids

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.124"
+version = "0.1.125"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }