Fix case where individual chains are sampled (#145)

* Update inputs.py

* Update mocks.py

Author: amorehead

alphafold3_pytorch/inputs.py

@@ -2142,10 +2142,16 @@ def pdb_input_to_molecule_input(
             chain_id: chain_idx
             for (chain_id, chain_idx) in zip(biomol.chain_id, biomol.chain_index)
         }
-        if exists(chain_id_1):
+        # NOTE: we have to manually nullify a chain ID value
+        # e.g., if an empty string is passed in as a "null" chain ID
+        if chain_id_1:
             chain_id_1 = chain_id_to_idx[chain_id_1]
-        if exists(chain_id_2):
+        else:
+            chain_id_1 = None
+        if chain_id_2:
             chain_id_2 = chain_id_to_idx[chain_id_2]
+        else:
+            chain_id_2 = None
         chains = (chain_id_1, chain_id_2)
 
     # crop the `Biomolecule` object during training only

alphafold3_pytorch/mocks.py

@@ -18,11 +18,13 @@ def __init__(
         data_length,
         max_seq_len = 16,
         atoms_per_window = 4,
+        dim_atom_inputs = 77,
         has_molecule_mods = True
     ):
         self.data_length = data_length
         self.max_seq_len = max_seq_len
         self.atoms_per_window = atoms_per_window
+        self.dim_atom_inputs = dim_atom_inputs
         self.has_molecule_mods = has_molecule_mods
 
     def __len__(self):
@@ -32,7 +34,7 @@ def __getitem__(self, idx):
         seq_len = randrange(1, self.max_seq_len)
         atom_seq_len = self.atoms_per_window * seq_len
 
-        atom_inputs = torch.randn(atom_seq_len, 77)
+        atom_inputs = torch.randn(atom_seq_len, self.dim_atom_inputs)
         atompair_inputs = torch.randn(atom_seq_len, atom_seq_len, 5)
 
         molecule_atom_lens = torch.randint(1, self.atoms_per_window, (seq_len,))