Harden folding input validation and fix DNA/RNA API inconsistencies

src/alphafold3/common/folding_input.py

@@ -109,6 +109,8 @@ def __init__(self, *, mmcif: str, query_to_template_map: Mapping[int, int]):
       query_to_template_map: A mapping from query residue index to template
         residue index.
     """
+    if not isinstance(mmcif, str) or not mmcif:
+      raise ValueError('Template mmcif must be a non-empty string.')
     self._mmcif = mmcif
     # Needed to make the Template class hashable.
     self._query_to_template = tuple(query_to_template_map.items())
@@ -381,13 +383,22 @@ def from_dict(
         mmcif_path = raw_template.get('mmcifPath', None)
         if mmcif and mmcif_path:
           raise ValueError('Only one of mmcif/mmcifPath can be set.')
+        if not mmcif and not mmcif_path:
+          raise ValueError(
+              'Template must set exactly one of "mmcif" or "mmcifPath".'
+          )
         if mmcif and len(mmcif) < 256 and epath.Path(mmcif).exists():
           raise ValueError('Set the template path using the "mmcifPath" field.')
         if mmcif_path:
           mmcif = _read_file(path=mmcif_path, json_path=json_path)
-        query_to_template_map = dict(
-            zip(raw_template['queryIndices'], raw_template['templateIndices'])
-        )
+        query_indices = raw_template['queryIndices']
+        template_indices = raw_template['templateIndices']
+        if len(query_indices) != len(template_indices):
+          raise ValueError(
+              'queryIndices and templateIndices must have the same length,'
+              f' got {len(query_indices)} and {len(template_indices)}.'
+          )
+        query_to_template_map = dict(zip(query_indices, template_indices))
         templates.append(
             Template(mmcif=mmcif, query_to_template_map=query_to_template_map)
         )
@@ -656,9 +667,11 @@ def fill_missing_fields(self) -> Self:
     """Fill missing MSA fields with default values."""
     return RnaChain(  # pyrefly: ignore[bad-return]
         id=self.id,
-        sequence=self.sequence,
-        modifications=self.modifications,
-        description=self.description,
+        # Use the raw sequence, not the CCD-remapped `.sequence` property,
+        # which replaces unknown/modified bases with 'N'.
+        sequence=self._sequence,
+        modifications=self._modifications,
+        description=self._description,
         unpaired_msa=self._unpaired_msa or '',
     )
 
@@ -719,6 +732,10 @@ def sequence(self) -> str:
   def description(self) -> str | None:
     return self._description
 
+  @property
+  def modifications(self) -> Sequence[tuple[str, int]]:
+    return self._modifications
+
   def __len__(self) -> int:
     return len(self._sequence)
 
@@ -737,9 +754,6 @@ def __hash__(self) -> int:
         (self._id, self._sequence, self._modifications, self._description)
     )
 
-  def modifications(self) -> Sequence[tuple[str, int]]:
-    return self._modifications
-
   def hash_without_id(self) -> int:
     """Returns a hash ignoring the ID - useful for deduplication."""
     return hash((self._sequence, self._modifications, self._description))
@@ -829,7 +843,12 @@ def __post_init__(self):
     if (self.ccd_ids is None) == (self.smiles is None):
       raise ValueError('Ligand must have one of CCD ID or SMILES set.')
 
+    if self.ccd_ids is not None and not self.ccd_ids:
+      raise ValueError('Ligand ccd_ids must be a non-empty sequence.')
+
     if self.smiles is not None:
+      if not self.smiles:
+        raise ValueError('Ligand SMILES must be a non-empty string.')
       mol = rd_chem.MolFromSmiles(self.smiles)
       if not mol:
         raise ValueError(f'Unable to make RDKit Mol from SMILES: {self.smiles}')
@@ -884,15 +903,20 @@ def from_dict(
             'CCD codes must be a list of strings, got '
             f'{type(ccd_codes).__name__} instead: {ccd_codes}'
         )
+      if not ccd_codes:
+        raise ValueError('Ligand ccdCodes must be a non-empty list.')
       return cls(
           id=seq_id or json_dict['id'],
           ccd_ids=ccd_codes,
           description=json_dict.get('description', None),
       )
     elif 'smiles' in json_dict:
+      smiles = json_dict['smiles']
+      if not smiles:
+        raise ValueError('Ligand smiles must be a non-empty string.')
       return cls(
           id=seq_id or json_dict['id'],
-          smiles=json_dict['smiles'],
+          smiles=smiles,
           description=json_dict.get('description', None),
       )
     else:
@@ -1552,10 +1576,20 @@ def load_fold_inputs_from_dir(input_dir: epath.PathLike) -> Iterator[Input]:
 
   Yields:
     The fold inputs from all JSON files in the input directory.
+
+  Raises:
+    ValueError: If the directory contains no JSON files.
   """
   input_dir = epath.Path(input_dir)
-  for file_path in sorted(input_dir.glob('*.json')):
-    if not file_path.is_file():
-      continue
-
+  json_files = [
+      file_path
+      for file_path in sorted(input_dir.glob('*.json'))
+      if file_path.is_file()
+  ]
+  if not json_files:
+    raise ValueError(
+        f'No JSON files found in input directory: {input_dir}. Provide at'
+        ' least one *.json fold input, or use --json_path for a single file.'
+    )
+  for file_path in json_files:
     yield from load_fold_inputs_from_path(file_path)

src/alphafold3/common/folding_input_test.py

@@ -0,0 +1,197 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# AlphaFold 3 source code is licensed under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with the
+# License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for folding_input validation and chain API consistency."""
+
+import json
+import pathlib
+import tempfile
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from alphafold3.common import folding_input
+
+
+_MINIMAL_MMCIF = """\
+data_test
+_entry.id test
+#
+loop_
+_atom_site.group_PDB
+_atom_site.id
+_atom_site.type_symbol
+_atom_site.label_atom_id
+_atom_site.label_alt_id
+_atom_site.label_comp_id
+_atom_site.label_asym_id
+_atom_site.label_entity_id
+_atom_site.label_seq_id
+_atom_site.pdbx_PDB_ins_code
+_atom_site.Cartn_x
+_atom_site.Cartn_y
+_atom_site.Cartn_z
+_atom_site.occupancy
+_atom_site.B_iso_or_equiv
+_atom_site.auth_seq_id
+_atom_site.auth_asym_id
+_atom_site.pdbx_PDB_model_num
+ATOM 1 C CA . ALA A 1 1 ? 0.000 0.000 0.000 1.00 0.00 1 A 1
+#
+"""
+
+
+def _protein_json(
+    *,
+    templates: list[dict] | None = None,
+    extra_sequences: list[dict] | None = None,
+) -> str:
+  protein = {
+      'id': 'A',
+      'sequence': 'ACDE',
+      'modifications': [],
+      'unpairedMsa': '',
+      'pairedMsa': '',
+  }
+  if templates is not None:
+    protein['templates'] = templates
+  sequences = [{'protein': protein}]
+  if extra_sequences:
+    sequences.extend(extra_sequences)
+  return json.dumps({
+      'name': 'test',
+      'sequences': sequences,
+      'modelSeeds': [1],
+      'dialect': 'alphafold3',
+      'version': 4,
+  })
+
+
+class TemplateValidationTest(parameterized.TestCase):
+
+  def test_rejects_mismatched_query_and_template_indices(self):
+    with self.assertRaisesRegex(ValueError, 'same length'):
+      folding_input.Input.from_json(
+          _protein_json(
+              templates=[{
+                  'mmcif': _MINIMAL_MMCIF,
+                  'queryIndices': [0, 1],
+                  'templateIndices': [0],
+              }]
+          )
+      )
+
+  def test_rejects_template_without_mmcif_or_path(self):
+    with self.assertRaisesRegex(ValueError, 'mmcif|mmcifPath'):
+      folding_input.Input.from_json(
+          _protein_json(
+              templates=[{
+                  'queryIndices': [0],
+                  'templateIndices': [0],
+              }]
+          )
+      )
+
+  def test_rejects_empty_template_mmcif(self):
+    with self.assertRaisesRegex(ValueError, 'non-empty'):
+      folding_input.Template(mmcif='', query_to_template_map={0: 0})
+
+  def test_accepts_matching_template_indices(self):
+    fold_input = folding_input.Input.from_json(
+        _protein_json(
+            templates=[{
+                'mmcif': _MINIMAL_MMCIF,
+                'queryIndices': [0, 1],
+                'templateIndices': [10, 11],
+            }]
+        )
+    )
+    self.assertEqual(
+        fold_input.protein_chains[0].templates[0].query_to_template_map,
+        {0: 10, 1: 11},
+    )
+
+
+class RnaChainFillMissingFieldsTest(absltest.TestCase):
+
+  def test_preserves_raw_sequence_with_unknown_modification(self):
+    # Unknown CCD codes map to 'N' via the `.sequence` property. Filling missing
+    # fields must keep the original residue letters in `_sequence`.
+    chain = folding_input.RnaChain(
+        id='R',
+        sequence='ACGU',
+        modifications=[('1MA', 2)],
+        unpaired_msa=None,
+    )
+    filled = chain.fill_missing_fields()
+    self.assertEqual(filled._sequence, 'ACGU')
+    self.assertEqual(filled.unpaired_msa, '')
+    self.assertEqual(filled.modifications, (('1MA', 2),))
+
+
+class DnaChainModificationsPropertyTest(absltest.TestCase):
+
+  def test_modifications_is_a_property(self):
+    chain = folding_input.DnaChain(
+        id='D', sequence='ACGT', modifications=[('6OG', 1)]
+    )
+    mods = chain.modifications
+    self.assertEqual(mods, (('6OG', 1),))
+    self.assertFalse(callable(mods))
+    # Iterable like RnaChain.modifications.
+    self.assertEqual(list(chain.modifications), [('6OG', 1)])
+
+
+class LigandValidationTest(parameterized.TestCase):
+
+  def test_rejects_empty_ccd_codes_list(self):
+    with self.assertRaisesRegex(ValueError, 'non-empty'):
+      folding_input.Ligand(id='L', ccd_ids=[])
+
+  def test_rejects_empty_ccd_codes_in_json(self):
+    payload = _protein_json(
+        extra_sequences=[{
+            'ligand': {'id': 'L', 'ccdCodes': []},
+        }]
+    )
+    with self.assertRaisesRegex(ValueError, 'non-empty'):
+      folding_input.Input.from_json(payload)
+
+  def test_rejects_empty_smiles(self):
+    with self.assertRaisesRegex(ValueError, 'non-empty'):
+      folding_input.Ligand(id='L', smiles='')
+
+  def test_accepts_valid_ccd_ligand(self):
+    ligand = folding_input.Ligand(id='L', ccd_ids=['ATP'])
+    self.assertEqual(ligand.ccd_ids, ('ATP',))
+
+
+class LoadFoldInputsFromDirTest(absltest.TestCase):
+
+  def test_empty_directory_raises(self):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+      with self.assertRaisesRegex(ValueError, 'No JSON files'):
+        list(folding_input.load_fold_inputs_from_dir(tmp_dir))
+
+  def test_loads_json_files(self):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+      path = pathlib.Path(tmp_dir) / 'fold.json'
+      path.write_text(_protein_json())
+      inputs = list(folding_input.load_fold_inputs_from_dir(tmp_dir))
+      self.assertLen(inputs, 1)
+      self.assertEqual(inputs[0].name, 'test')
+
+
+if __name__ == '__main__':
+  absltest.main()