fix: support atom_keys for ASE ingestion

atompack-py/python/atompack/__init__.pyi

@@ -1,6 +1,6 @@
 """Type stubs for atompack"""
 
-from typing import Any, Literal, Sequence, overload
+from typing import Any, Iterable, Literal, Sequence, overload
 
 import numpy as np
 import numpy.typing as npt
@@ -656,6 +656,7 @@ def from_ase(
     copy_info: bool = True,
     copy_arrays: bool = True,
     info: dict | None = None,
+    atom_keys: Iterable[str] | None = None,
 ) -> Molecule:
     """
     Convert an ASE Atoms object to an atompack Molecule.
@@ -691,6 +692,10 @@ def from_ase(
         Additional properties to store in the molecule. These will be added after
         copying atoms.info (if copy_info=True), so they can override atoms.info values.
         Supports the same types as copy_info.
+    atom_keys : iterable of str, optional
+        Custom property keys to store as atom-scope properties when found in
+        atoms.arrays, atoms.info, calculator results, or the explicit info override.
+        Unlisted keys remain molecule-scope properties.
 
     Returns
     -------
@@ -763,6 +768,7 @@ def add_ase_batch(
     copy_info: bool = True,
     copy_arrays: bool = True,
     info: dict | list[dict | None] | None = None,
+    atom_keys: Iterable[str] | None = None,
     batch_size: int = 512,
 ) -> None:
     """Write many ASE Atoms objects into an atompack database efficiently."""

atompack-py/python/atompack/ase_bridge.py

@@ -106,7 +106,30 @@ def _coerce_custom_property(key, value, source):
     return coerced
 
 
-def _merge_properties(properties, builtins, values, source):
+def _normalize_atom_keys(atom_keys):
+    if atom_keys is None:
+        return frozenset()
+    if isinstance(atom_keys, str):
+        raise TypeError("atom_keys must be an iterable of custom property names, not a string")
+    try:
+        normalized = frozenset(atom_keys)
+    except TypeError as exc:
+        raise TypeError("atom_keys must be an iterable of custom property names") from exc
+    invalid = [key for key in normalized if not isinstance(key, str)]
+    if invalid:
+        raise TypeError("atom_keys must contain only strings")
+    return normalized
+
+
+def _store_custom_property(properties, atom_properties, atom_keys, key, value, source):
+    coerced = _coerce_custom_property(key, value, source)
+    if key in atom_keys:
+        atom_properties[key] = coerced
+    else:
+        properties[key] = coerced
+
+
+def _merge_properties(properties, atom_properties, builtins, atom_keys, values, source):
     for key, value in values.items():
         if key in _BUILTIN_FIELDS:
             # Builtin keys in atoms.info / info-override go to the builtins
@@ -119,7 +142,7 @@ def _merge_properties(properties, builtins, values, source):
                 if arr.shape == (3, 3) and arr.dtype.kind == "f":
                     builtins["stress"] = arr.astype(np.float64, copy=False)
             continue
-        properties[key] = _coerce_custom_property(key, value, source)
+        _store_custom_property(properties, atom_properties, atom_keys, key, value, source)
 
 
 def _extract_ase_record(
@@ -134,7 +157,9 @@ def _extract_ase_record(
     copy_info=True,
     copy_arrays=True,
     info=None,
+    atom_keys=None,
 ):
+    atom_keys = _normalize_atom_keys(atom_keys)
     positions = np.asarray(atoms.get_positions(), dtype=np.float32)
     atomic_numbers = np.asarray(atoms.get_atomic_numbers(), dtype=np.uint8)
     n_atoms = len(atomic_numbers)
@@ -202,6 +227,7 @@ def _extract_ase_record(
         builtins["stress"] = _get_stress(atoms)
 
     properties = {}
+    atom_properties = {}
 
     arrays = getattr(atoms, "arrays", None)
     if copy_arrays and isinstance(arrays, dict):
@@ -212,26 +238,55 @@ def _extract_ase_record(
             # builtins["forces"] (from get_forces()) and properties["forces"].
             if key in _ASE_RESERVED_ARRAYS or key in _BUILTIN_FIELDS:
                 continue
-            properties[key] = _coerce_custom_property(key, value, "atoms.arrays")
+            _store_custom_property(
+                properties,
+                atom_properties,
+                atom_keys,
+                key,
+                value,
+                "atoms.arrays",
+            )
 
     calc = getattr(atoms, "calc", None)
     results = getattr(calc, "results", None)
     if isinstance(results, dict):
         for key, value in results.items():
             if key not in _BUILTIN_FIELDS:
-                properties[key] = _coerce_custom_property(key, value, "atoms.calc.results")
+                _store_custom_property(
+                    properties,
+                    atom_properties,
+                    atom_keys,
+                    key,
+                    value,
+                    "atoms.calc.results",
+                )
 
     if copy_info and getattr(atoms, "info", None):
-        _merge_properties(properties, builtins, atoms.info, "atoms.info")
+        _merge_properties(
+            properties,
+            atom_properties,
+            builtins,
+            atom_keys,
+            atoms.info,
+            "atoms.info",
+        )
     if info is not None:
-        _merge_properties(properties, builtins, info, "info override")
+        _merge_properties(
+            properties,
+            atom_properties,
+            builtins,
+            atom_keys,
+            info,
+            "info override",
+        )
 
     return {
         "positions": positions,
         "atomic_numbers": atomic_numbers,
         "n_atoms": n_atoms,
         "builtins": builtins,
         "properties": properties,
+        "atom_properties": atom_properties,
     }
 
 
@@ -250,6 +305,8 @@ def _record_to_molecule(record):
     )
     for key, value in record["properties"].items():
         mol.set_property(key, value)
+    for key, value in record["atom_properties"].items():
+        mol.set_property(key, value, scope="atom")
     return mol
 
 
@@ -682,11 +739,13 @@ def from_ase(
     copy_info=True,
     copy_arrays=True,
     info=None,
+    atom_keys=None,
 ):
     """Convert one ASE Atoms object to an atompack Molecule.
 
     Custom values from ``atoms.info``, ``atoms.arrays``, calculator results,
     and explicit ``info=`` overrides are stored as molecule-scope properties.
+    Keys listed in ``atom_keys`` are stored as atom-scope custom properties.
     Array shape is not used to infer atom-property scope during ingestion.
     """
     return _record_to_molecule(
@@ -701,6 +760,7 @@ def from_ase(
             copy_info=copy_info,
             copy_arrays=copy_arrays,
             info=info,
+            atom_keys=atom_keys,
         )
     )
 
@@ -712,13 +772,15 @@ def add_ase_batch(
     copy_info=True,
     copy_arrays=True,
     info=None,
+    atom_keys=None,
     batch_size=512,
 ):
     """Write many ASE Atoms objects efficiently, preserving supported metadata."""
     atoms_list = list(atoms_list)
     if not atoms_list:
         return
 
+    atom_keys = _normalize_atom_keys(atom_keys)
     info_overrides = _normalize_info_overrides(info, len(atoms_list))
     fast_key = None
     fast_records = []
@@ -742,8 +804,9 @@ def flush_slow():
             copy_info=copy_info,
             copy_arrays=copy_arrays,
             info=info_override,
+            atom_keys=atom_keys,
         )
-        if record["properties"]:
+        if record["properties"] or record["atom_properties"]:
             flush_fast()
             slow_records.append(_record_to_molecule(record))
             if len(slow_records) >= batch_size:

atompack-py/tests/test_from_ase.py

@@ -7,7 +7,9 @@
 import numpy as np
 import pytest
 
-ase = pytest.importorskip("ase")
+
+def requires_ase():
+    return pytest.importorskip("ase")
 
 
 @dataclass
@@ -201,6 +203,46 @@ def test_from_ase_custom_arrays_remain_molecule_properties() -> None:
     assert mol.has_property("descriptor", scope="atom") is False
 
 
+def test_from_ase_atom_keys_route_custom_values_to_atom_properties() -> None:
+    fixed_mask = np.array([0, 1], dtype=np.int32)
+    site_weight = np.array([0.25, 0.75], dtype=np.float64)
+    descriptor = np.arange(4, dtype=np.float32).reshape(2, 2)
+    atoms = FakeASEAtoms(
+        positions=np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64),
+        atomic_numbers=np.array([6, 8], dtype=np.int64),
+        pbc=np.array([False, False, False]),
+        arrays={"fixed_mask": fixed_mask, "descriptor": descriptor},
+        info={"site_weight": site_weight, "temperature": 300.0},
+    )
+
+    mol = atompack.from_ase(atoms, atom_keys=["fixed_mask", "site_weight"])
+
+    np.testing.assert_array_equal(mol.get_property("fixed_mask"), fixed_mask)
+    assert mol.has_property("fixed_mask", scope="atom") is True
+    assert mol.has_property("fixed_mask", scope="molecule") is False
+
+    np.testing.assert_allclose(mol.get_property("site_weight"), site_weight)
+    assert mol.has_property("site_weight", scope="atom") is True
+    assert mol.has_property("site_weight", scope="molecule") is False
+
+    np.testing.assert_allclose(mol.get_property("descriptor"), descriptor)
+    assert mol.has_property("descriptor", scope="molecule") is True
+    assert mol.has_property("descriptor", scope="atom") is False
+    assert mol.get_property("temperature") == pytest.approx(300.0)
+
+
+def test_from_ase_atom_keys_validate_first_dimension() -> None:
+    atoms = FakeASEAtoms(
+        positions=np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64),
+        atomic_numbers=np.array([6, 8], dtype=np.int64),
+        pbc=np.array([False, False, False]),
+        arrays={"bad_mask": np.array([0, 1, 0], dtype=np.int32)},
+    )
+
+    with pytest.raises(ValueError, match=r"bad_mask.*first dimension \(3\).*atom count \(2\)"):
+        atompack.from_ase(atoms, atom_keys=["bad_mask"])
+
+
 def test_from_ase_rejects_unsupported_enabled_custom_values_and_honors_optouts() -> None:
     atoms_with_bad_info = FakeASEAtoms(
         positions=np.array([[0.0, 0.0, 0.0]], dtype=np.float64),
@@ -352,6 +394,7 @@ def test_from_ase_info_override_kwarg_filters_builtins() -> None:
 
 
 def test_to_ase_does_not_duplicate_builtins_in_arrays() -> None:
+    requires_ase()
     # to_ase mirror of the from_ase fix: even if a user explicitly stuffs a
     # builtin name into custom properties via mol.set_property("forces", ...),
     # to_ase must NOT shovel that custom value into atoms.arrays["forces"]
@@ -420,7 +463,40 @@ def test_add_ase_batch_roundtrip(tmp_path) -> None:
     assert second.get_property("temperature") == pytest.approx(300.0)
 
 
+def test_add_ase_batch_atom_keys_allow_variable_atom_array_lengths(tmp_path) -> None:
+    path = tmp_path / "ase_variable_atom_keys.atp"
+    atoms_list = [
+        FakeASEAtoms(
+            positions=np.zeros((3, 3), dtype=np.float64),
+            atomic_numbers=np.ones(3, dtype=np.int64),
+            pbc=np.array([False, False, False]),
+            arrays={"fixed_mask": np.array([0, 1, 0], dtype=np.int32)},
+        ),
+        FakeASEAtoms(
+            positions=np.zeros((4, 3), dtype=np.float64),
+            atomic_numbers=np.ones(4, dtype=np.int64),
+            pbc=np.array([False, False, False]),
+            arrays={"fixed_mask": np.array([1, 0, 0, 1], dtype=np.int32)},
+        ),
+    ]
+
+    db = atompack.Database(str(path))
+    atompack.add_ase_batch(db, atoms_list, atom_keys=["fixed_mask"], batch_size=1)
+    db.flush()
+
+    reopened = atompack.Database.open(str(path))
+    flat = reopened.get_molecules_flat([0, 1])
+
+    np.testing.assert_array_equal(flat["n_atoms"], np.array([3, 4], dtype=np.uint32))
+    np.testing.assert_array_equal(
+        flat["atom_properties"]["fixed_mask"],
+        np.array([0, 1, 0, 1, 0, 0, 1], dtype=np.int32),
+    )
+    assert "properties" not in flat or "fixed_mask" not in flat["properties"]
+
+
 def test_to_ase_owned_maps_builtins_and_properties() -> None:
+    requires_ase()
     mol = atompack.Molecule.from_arrays(
         np.array([[0.0, 0.0, 0.0], [1.0, 0.5, 0.0]], dtype=np.float32),
         np.array([6, 8], dtype=np.uint8),
@@ -452,6 +528,7 @@ def test_to_ase_owned_maps_builtins_and_properties() -> None:
 
 @pytest.mark.parametrize("view_backed", [False, True])
 def test_to_ase_routes_tensor_properties_by_scope(tmp_path, view_backed: bool) -> None:
+    requires_ase()
     molecule_tensor = np.arange(8, dtype=np.float32).reshape(2, 2, 2)
     atom_tensor = np.arange(8, dtype=np.float64).reshape(2, 2, 2)
     mol = atompack.Molecule.from_arrays(
@@ -486,6 +563,7 @@ def test_to_ase_routes_tensor_properties_by_scope(tmp_path, view_backed: bool) -
 
 
 def test_to_ase_roundtrip_preserves_none_custom_property() -> None:
+    requires_ase()
     mol = atompack.Molecule.from_arrays(
         np.array([[0.0, 0.0, 0.0]], dtype=np.float32),
         np.array([1], dtype=np.uint8),
@@ -502,6 +580,7 @@ def test_to_ase_roundtrip_preserves_none_custom_property() -> None:
 
 
 def test_to_ase_calc_modes() -> None:
+    requires_ase()
     mol = atompack.Molecule.from_arrays(
         np.array([[0.0, 0.0, 0.0], [1.0, 0.5, 0.0]], dtype=np.float32),
         np.array([6, 8], dtype=np.uint8),
@@ -526,6 +605,7 @@ def test_to_ase_calc_modes() -> None:
 
 
 def test_to_ase_view_backed_molecule(tmp_path) -> None:
+    requires_ase()
     path = tmp_path / "to_ase_view.atp"
     mol = atompack.Molecule.from_arrays(
         np.array([[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], dtype=np.float32),
@@ -556,6 +636,7 @@ def test_to_ase_view_backed_molecule(tmp_path) -> None:
 
 @pytest.mark.parametrize("mmap", [True, False])
 def test_database_to_ase_batch_matches_per_molecule(tmp_path, mmap: bool) -> None:
+    requires_ase()
     path = tmp_path / "to_ase_batch.atp"
     positions = np.array(
         [
@@ -620,6 +701,7 @@ def test_database_to_ase_batch_matches_per_molecule(tmp_path, mmap: bool) -> Non
 
 
 def test_to_ase_batch_with_molecule_list_matches_individual() -> None:
+    requires_ase()
     molecules = [
         atompack.Molecule.from_arrays(
             np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float32),
@@ -651,6 +733,7 @@ def test_to_ase_batch_with_molecule_list_matches_individual() -> None:
 
 
 def test_to_ase_batch_nocopy_calc_mode(tmp_path) -> None:
+    requires_ase()
     path = tmp_path / "to_ase_batch_nocopy.atp"
     positions = np.array(
         [
@@ -691,6 +774,7 @@ def test_to_ase_batch_nocopy_calc_mode(tmp_path) -> None:
 
 
 def test_to_ase_batch_none_calc_mode_preserves_results(tmp_path) -> None:
+    requires_ase()
     path = tmp_path / "to_ase_batch_none.atp"
     positions = np.array(
         [