fix ion feature metadata

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,

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,
                 }

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)