diff --git a/atompack-py/python/atompack/__init__.pyi b/atompack-py/python/atompack/__init__.pyi index d2fbece..b6610da 100644 --- a/atompack-py/python/atompack/__init__.pyi +++ b/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.""" diff --git a/atompack-py/python/atompack/ase_bridge.py b/atompack-py/python/atompack/ase_bridge.py index b9ab31d..d2279e5 100644 --- a/atompack-py/python/atompack/ase_bridge.py +++ b/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,19 +238,47 @@ 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, @@ -232,6 +286,7 @@ def _extract_ase_record( "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,6 +772,7 @@ 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.""" @@ -719,6 +780,7 @@ def add_ase_batch( 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: diff --git a/atompack-py/tests/test_from_ase.py b/atompack-py/tests/test_from_ase.py index a244322..e7e2140 100644 --- a/atompack-py/tests/test_from_ase.py +++ b/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( [