From b775b69bc41a0a6ecbf9ae744289c781678b4546 Mon Sep 17 00:00:00 2001 From: Moritz Ertelt <59534445+MoritzErtelt@users.noreply.github.com> Date: Mon, 3 Aug 2026 13:29:23 +0200 Subject: [PATCH 1/3] Add regression tests for ion feature metadata --- tests/test_ion_feature_metadata.py | 96 ++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 tests/test_ion_feature_metadata.py diff --git a/tests/test_ion_feature_metadata.py b/tests/test_ion_feature_metadata.py new file mode 100644 index 0000000..54d5988 --- /dev/null +++ b/tests/test_ion_feature_metadata.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 Aureka AI Research + +import numpy as np +import pytest +from biotite.structure import AtomArray + +from opendde.data.msa.msa_featurizer import InferenceMSAFeaturizer +from opendde.data.template.template_featurizer import InferenceTemplateFeaturizer + + +PROTEIN_SEQUENCE = "ACDEF" + + +def _make_input(entity_order: tuple[str, ...]) -> tuple[list[dict], AtomArray]: + bioassembly = [] + asym_ids = [] + res_ids = [] + chain_ids = [] + + for asym_id, entity_type in enumerate(entity_order): + chain_id = chr(ord("A") + asym_id) + if entity_type == "protein": + bioassembly.append( + { + "proteinChain": { + "sequence": PROTEIN_SEQUENCE, + "count": 1, + } + } + ) + entity_size = len(PROTEIN_SEQUENCE) + entity_res_ids = range(1, entity_size + 1) + elif entity_type == "ion": + bioassembly.append({"ion": {"ion": "ZN", "count": 1}}) + entity_size = 1 + entity_res_ids = [1] + else: + raise ValueError(f"Unsupported test entity type: {entity_type}") + + asym_ids.extend([asym_id] * entity_size) + res_ids.extend(entity_res_ids) + chain_ids.extend([chain_id] * entity_size) + + atom_array = AtomArray(len(asym_ids)) + atom_array.res_id[:] = res_ids + atom_array.chain_id[:] = chain_ids + atom_array.set_annotation("asym_id_int", np.asarray(asym_ids, dtype=np.int64)) + atom_array.set_annotation( + "centre_atom_mask", np.ones(len(asym_ids), dtype=np.int8) + ) + return bioassembly, atom_array + + +@pytest.mark.parametrize( + "entity_order", + [ + ("protein", "ion"), + ("ion", "protein"), + ], +) +def test_msa_feature_metadata_includes_ion(entity_order: tuple[str, ...]) -> None: + bioassembly, atom_array = _make_input(entity_order) + + features = InferenceMSAFeaturizer.make_msa_feature( + bioassembly=bioassembly, + atom_array=atom_array, + msa_pair_as_unpair=False, + use_rna_msa=False, + ) + + assert features["msa"].shape[1] == len(atom_array) + assert features["profile"].shape[0] == len(atom_array) + + +@pytest.mark.parametrize( + "entity_order", + [ + ("protein", "ion"), + ("ion", "protein"), + ], +) +def test_template_feature_metadata_includes_ion( + entity_order: tuple[str, ...], +) -> None: + bioassembly, atom_array = _make_input(entity_order) + + features = InferenceTemplateFeaturizer.make_template_feature( + bioassembly=bioassembly, + atom_array=atom_array, + use_template=False, + online_template_featurizer=None, + ) + + assert features["template_aatype"].shape[1] == len(atom_array) + assert features["template_atom_mask"].shape[1] == len(atom_array) From dbe4a7ab3084b8f4d3a61632981217e76fbba5ea Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 11:32:13 +0000 Subject: [PATCH 2/3] Handle ions in inference feature metadata --- opendde/data/msa/msa_featurizer.py | 5 +++-- opendde/data/template/template_featurizer.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/opendde/data/msa/msa_featurizer.py b/opendde/data/msa/msa_featurizer.py index ab62346..f78bf4f 100644 --- a/opendde/data/msa/msa_featurizer.py +++ b/opendde/data/msa/msa_featurizer.py @@ -361,9 +361,10 @@ def make_msa_feature( elif "dnaSequence" in info: c = info["dnaSequence"] seq, count, ctype = c["sequence"], c["count"], DNA_CHAIN - elif "ligand" in info: + elif "ligand" in info or "ion" in info: + entity_key = "ligand" if "ligand" in info else "ion" count, ctype, seq = ( - info["ligand"]["count"], + info[entity_key]["count"], LIGAND_CHAIN_TYPES, "X" * (atom_array.asym_id_int == curr_aid).sum(), ) diff --git a/opendde/data/template/template_featurizer.py b/opendde/data/template/template_featurizer.py index 200efed..1809c49 100644 --- a/opendde/data/template/template_featurizer.py +++ b/opendde/data/template/template_featurizer.py @@ -283,8 +283,9 @@ def make_template_feature( elif "dnaSequence" in info: c = info["dnaSequence"] seq, count, ctype = c["sequence"], c["count"], DNA_CHAIN - elif "ligand" in info: - count, ctype = info["ligand"]["count"], LIGAND_CHAIN_TYPES + elif "ligand" in info or "ion" in info: + entity_key = "ligand" if "ligand" in info else "ion" + count, ctype = info[entity_key]["count"], LIGAND_CHAIN_TYPES seq = "X" * (atom_array.asym_id_int == curr_asym_id).sum() templates = [] From 08f36546bf9a4edcaddb011617c33dbe1ffb7f4f Mon Sep 17 00:00:00 2001 From: MoritzErtelt Date: Mon, 3 Aug 2026 13:50:15 +0200 Subject: [PATCH 3/3] Strengthen ion feature metadata handling --- opendde/data/msa/msa_featurizer.py | 15 +++-- opendde/data/template/template_featurizer.py | 15 +++-- tests/test_ion_feature_metadata.py | 59 ++++++++++++-------- 3 files changed, 58 insertions(+), 31 deletions(-) diff --git a/opendde/data/msa/msa_featurizer.py b/opendde/data/msa/msa_featurizer.py index f78bf4f..1098f9b 100644 --- a/opendde/data/msa/msa_featurizer.py +++ b/opendde/data/msa/msa_featurizer.py @@ -361,13 +361,14 @@ def make_msa_feature( elif "dnaSequence" in info: c = info["dnaSequence"] seq, count, ctype = c["sequence"], c["count"], DNA_CHAIN - elif "ligand" in info or "ion" in info: - entity_key = "ligand" if "ligand" in info else "ion" + elif "ligand" in info: count, ctype, seq = ( - info[entity_key]["count"], + info["ligand"]["count"], LIGAND_CHAIN_TYPES, "X" * (atom_array.asym_id_int == curr_aid).sum(), ) + elif "ion" in info: + count, ctype = info["ion"]["count"], LIGAND_CHAIN_TYPES p_a3m = ensure_ends_with_newline(p_a3m) u_a3m = ensure_ends_with_newline(u_a3m) @@ -379,10 +380,16 @@ def make_msa_feature( for c_idx in range(count): aid = curr_aid + c_idx + chain_seq = seq + if "ion" in info: + chain_seq = "X" * np.count_nonzero( + (atom_array.asym_id_int == aid) + & atom_array.centre_atom_mask.astype(bool) + ) meta[aid] = { "entity_id": eid, "chain_id": atom_array.chain_id[atom_array.asym_id_int == aid][0], - "sequence": seq, + "sequence": chain_seq, "paired_msa": p_a3m or "", "unpaired_msa": u_a3m or "", "chain_entity_type": ctype, diff --git a/opendde/data/template/template_featurizer.py b/opendde/data/template/template_featurizer.py index 1809c49..acaa627 100644 --- a/opendde/data/template/template_featurizer.py +++ b/opendde/data/template/template_featurizer.py @@ -283,10 +283,11 @@ def make_template_feature( elif "dnaSequence" in info: c = info["dnaSequence"] seq, count, ctype = c["sequence"], c["count"], DNA_CHAIN - elif "ligand" in info or "ion" in info: - entity_key = "ligand" if "ligand" in info else "ion" - count, ctype = info[entity_key]["count"], LIGAND_CHAIN_TYPES + elif "ligand" in info: + count, ctype = info["ligand"]["count"], LIGAND_CHAIN_TYPES seq = "X" * (atom_array.asym_id_int == curr_asym_id).sum() + elif "ion" in info: + count, ctype = info["ion"]["count"], LIGAND_CHAIN_TYPES templates = [] if t_path and use_template and online_template_featurizer: @@ -314,10 +315,16 @@ def make_template_feature( for i in range(count): aid = curr_asym_id + i + chain_seq = seq + if "ion" in info: + chain_seq = "X" * np.count_nonzero( + (atom_array.asym_id_int == aid) + & atom_array.centre_atom_mask.astype(bool) + ) template_meta_infos[aid] = { "entity_id": eid, "chain_id": atom_array.chain_id[atom_array.asym_id_int == aid][0], - "sequence": seq, + "sequence": chain_seq, "chain_entity_type": ctype, "templates": templates, } diff --git a/tests/test_ion_feature_metadata.py b/tests/test_ion_feature_metadata.py index 54d5988..805745e 100644 --- a/tests/test_ion_feature_metadata.py +++ b/tests/test_ion_feature_metadata.py @@ -10,57 +10,66 @@ PROTEIN_SEQUENCE = "ACDEF" +EntitySpec = tuple[str, int] -def _make_input(entity_order: tuple[str, ...]) -> tuple[list[dict], AtomArray]: +def _make_input(entity_specs: tuple[EntitySpec, ...]) -> tuple[list[dict], AtomArray]: bioassembly = [] asym_ids = [] res_ids = [] chain_ids = [] + curr_asym_id = 0 - for asym_id, entity_type in enumerate(entity_order): - chain_id = chr(ord("A") + asym_id) + for entity_type, count in entity_specs: if entity_type == "protein": bioassembly.append( { "proteinChain": { "sequence": PROTEIN_SEQUENCE, - "count": 1, + "count": count, } } ) entity_size = len(PROTEIN_SEQUENCE) - entity_res_ids = range(1, entity_size + 1) elif entity_type == "ion": - bioassembly.append({"ion": {"ion": "ZN", "count": 1}}) + bioassembly.append({"ion": {"ion": "ZN", "count": count}}) entity_size = 1 - entity_res_ids = [1] + elif entity_type == "ligand": + bioassembly.append({"ligand": {"ligand": "CCD_ATP", "count": count}}) + entity_size = 2 else: raise ValueError(f"Unsupported test entity type: {entity_type}") - asym_ids.extend([asym_id] * entity_size) - res_ids.extend(entity_res_ids) - chain_ids.extend([chain_id] * entity_size) + for _ in range(count): + chain_id = chr(ord("A") + curr_asym_id) + asym_ids.extend([curr_asym_id] * entity_size) + res_ids.extend(range(1, entity_size + 1)) + chain_ids.extend([chain_id] * entity_size) + curr_asym_id += 1 atom_array = AtomArray(len(asym_ids)) atom_array.res_id[:] = res_ids atom_array.chain_id[:] = chain_ids atom_array.set_annotation("asym_id_int", np.asarray(asym_ids, dtype=np.int64)) - atom_array.set_annotation( - "centre_atom_mask", np.ones(len(asym_ids), dtype=np.int8) - ) + atom_array.set_annotation("centre_atom_mask", np.ones(len(asym_ids), dtype=np.int8)) return bioassembly, atom_array @pytest.mark.parametrize( - "entity_order", + "entity_specs", [ - ("protein", "ion"), - ("ion", "protein"), + (("protein", 1), ("ion", 1)), + (("ion", 1), ("protein", 1)), + (("protein", 1), ("ion", 2)), + (("protein", 1), ("ligand", 1), ("ion", 1)), + (("ion", 1), ("ion", 1), ("protein", 1)), + (("protein", 1), ("ion", 1), ("protein", 1)), ], ) -def test_msa_feature_metadata_includes_ion(entity_order: tuple[str, ...]) -> None: - bioassembly, atom_array = _make_input(entity_order) +def test_msa_feature_metadata_includes_ion( + entity_specs: tuple[EntitySpec, ...], +) -> None: + bioassembly, atom_array = _make_input(entity_specs) features = InferenceMSAFeaturizer.make_msa_feature( bioassembly=bioassembly, @@ -74,16 +83,20 @@ def test_msa_feature_metadata_includes_ion(entity_order: tuple[str, ...]) -> Non @pytest.mark.parametrize( - "entity_order", + "entity_specs", [ - ("protein", "ion"), - ("ion", "protein"), + (("protein", 1), ("ion", 1)), + (("ion", 1), ("protein", 1)), + (("protein", 1), ("ion", 2)), + (("protein", 1), ("ligand", 1), ("ion", 1)), + (("ion", 1), ("ion", 1), ("protein", 1)), + (("protein", 1), ("ion", 1), ("protein", 1)), ], ) def test_template_feature_metadata_includes_ion( - entity_order: tuple[str, ...], + entity_specs: tuple[EntitySpec, ...], ) -> None: - bioassembly, atom_array = _make_input(entity_order) + bioassembly, atom_array = _make_input(entity_specs) features = InferenceTemplateFeaturizer.make_template_feature( bioassembly=bioassembly,