From c22a010fce33e1975ef0f4ca193ac2fcf9967dec Mon Sep 17 00:00:00 2001 From: Ayush7614 Date: Wed, 29 Jul 2026 02:59:30 +0530 Subject: [PATCH] Harden folding input validation and fix DNA/RNA API inconsistencies. Reject malformed templates, empty ligands, and empty input dirs early; preserve raw RNA sequences when filling missing fields; expose DNA modifications as a property; and add unit tests covering these paths. --- src/alphafold3/common/folding_input.py | 62 ++++-- src/alphafold3/common/folding_input_test.py | 197 ++++++++++++++++++++ 2 files changed, 245 insertions(+), 14 deletions(-) create mode 100644 src/alphafold3/common/folding_input_test.py diff --git a/src/alphafold3/common/folding_input.py b/src/alphafold3/common/folding_input.py index 5ac92bf6..7f181086 100644 --- a/src/alphafold3/common/folding_input.py +++ b/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) diff --git a/src/alphafold3/common/folding_input_test.py b/src/alphafold3/common/folding_input_test.py new file mode 100644 index 00000000..f8894d54 --- /dev/null +++ b/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()