rename residue to molecules, as to generalize the system to ligands

Author: lucidrains

README.md

@@ -38,13 +38,13 @@ alphafold3 = Alphafold3(
 # mock inputs
 
 seq_len = 16
-residue_atom_lens = torch.randint(1, 3, (2, seq_len))
-atom_seq_len = residue_atom_lens.sum(dim = -1).amax()
+molecule_atom_lens = torch.randint(1, 3, (2, seq_len))
+atom_seq_len = molecule_atom_lens.sum(dim = -1).amax()
 
 atom_inputs = torch.randn(2, atom_seq_len, 77)
 atompair_inputs = torch.randn(2, atom_seq_len, atom_seq_len, 5)
 
-additional_residue_feats = torch.randn(2, seq_len, 10)
+additional_molecule_feats = torch.randn(2, seq_len, 10)
 
 template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
 template_mask = torch.ones((2, 2)).bool()
@@ -55,7 +55,7 @@ msa_mask = torch.ones((2, 7)).bool()
 # required for training, but omitted on inference
 
 atom_pos = torch.randn(2, atom_seq_len, 3)
-residue_atom_indices = residue_atom_lens - 1 # last atom, as an example
+molecule_atom_indices = molecule_atom_lens - 1 # last atom, as an example
 
 distance_labels = torch.randint(0, 37, (2, seq_len, seq_len))
 pae_labels = torch.randint(0, 64, (2, seq_len, seq_len))
@@ -69,14 +69,14 @@ loss = alphafold3(
     num_recycling_steps = 2,
     atom_inputs = atom_inputs,
     atompair_inputs = atompair_inputs,
-    residue_atom_lens = residue_atom_lens,
-    additional_residue_feats = additional_residue_feats,
+    molecule_atom_lens = molecule_atom_lens,
+    additional_molecule_feats = additional_molecule_feats,
     msa = msa,
     msa_mask = msa_mask,
     templates = template_feats,
     template_mask = template_mask,
     atom_pos = atom_pos,
-    residue_atom_indices = residue_atom_indices,
+    molecule_atom_indices = molecule_atom_indices,
     distance_labels = distance_labels,
     pae_labels = pae_labels,
     pde_labels = pde_labels,
@@ -93,8 +93,8 @@ sampled_atom_pos = alphafold3(
     num_sample_steps = 16,
     atom_inputs = atom_inputs,
     atompair_inputs = atompair_inputs,
-    residue_atom_lens = residue_atom_lens,
-    additional_residue_feats = additional_residue_feats,
+    molecule_atom_lens = molecule_atom_lens,
+    additional_molecule_feats = additional_molecule_feats,
     msa = msa,
     msa_mask = msa_mask,
     templates = template_feats,