diff --git a/opendde/data/msa/msa_featurizer.py b/opendde/data/msa/msa_featurizer.py index ab62346..1098f9b 100644 --- a/opendde/data/msa/msa_featurizer.py +++ b/opendde/data/msa/msa_featurizer.py @@ -367,6 +367,8 @@ def make_msa_feature( 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) @@ -378,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 200efed..acaa627 100644 --- a/opendde/data/template/template_featurizer.py +++ b/opendde/data/template/template_featurizer.py @@ -286,6 +286,8 @@ def make_template_feature( 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: @@ -313,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 new file mode 100644 index 0000000..805745e --- /dev/null +++ b/tests/test_ion_feature_metadata.py @@ -0,0 +1,109 @@ +# 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" +EntitySpec = tuple[str, int] + + +def _make_input(entity_specs: tuple[EntitySpec, ...]) -> tuple[list[dict], AtomArray]: + bioassembly = [] + asym_ids = [] + res_ids = [] + chain_ids = [] + curr_asym_id = 0 + + for entity_type, count in entity_specs: + if entity_type == "protein": + bioassembly.append( + { + "proteinChain": { + "sequence": PROTEIN_SEQUENCE, + "count": count, + } + } + ) + entity_size = len(PROTEIN_SEQUENCE) + elif entity_type == "ion": + bioassembly.append({"ion": {"ion": "ZN", "count": count}}) + entity_size = 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}") + + 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)) + return bioassembly, atom_array + + +@pytest.mark.parametrize( + "entity_specs", + [ + (("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_specs: tuple[EntitySpec, ...], +) -> None: + bioassembly, atom_array = _make_input(entity_specs) + + 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_specs", + [ + (("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_specs: tuple[EntitySpec, ...], +) -> None: + bioassembly, atom_array = _make_input(entity_specs) + + 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)