Propagate atom tags from selected image to packed box (#119)

* Propagate atom tags from selected image to packed box

Add logic to propagate atom tags from selected conformer images to packed box.

* test: verify atom tags propagate through pack

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Fabian Zills <fzills@icp.uni-stuttgart.de>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

src/molify/pack.py

@@ -137,6 +137,9 @@ def _extract_atom_arrays(
     charges = np.concatenate([atom.get_initial_charges() for atom in selected_images])
     if any(charge != 0 for charge in charges):
         packed_atoms.set_initial_charges(charges)
+    tags = np.concatenate([atom.get_tags() for atom in selected_images])
+    if any(tag != 0 for tag in tags):
+        packed_atoms.set_tags(tags)
     return packed_atoms
 
 

tests/test_solvate.py

@@ -109,6 +109,16 @@ def test_pack_charges():
     assert charges == [0, 0, 0, -1, 0, 1] + [1, 0, 0, -1, 0, 0, 0, 0, 0, 0]
 
 
+def test_pack_tags():
+    water = smiles2conformers("O", 1)
+    methane = smiles2conformers("C", 1)
+    water[0].set_tags(range(1, 4))  # O, H, H -> 1, 2, 3
+    methane[0].set_tags(range(4, 9))  # C, H, H, H, H -> 4, 5, 6, 7, 8
+
+    atoms = pack([water, methane], [2, 1], density=1000)
+    npt.assert_array_equal(atoms.get_tags(), [1, 2, 3, 1, 2, 3, 4, 5, 6, 7, 8])
+
+
 def test_pack_ratio():
     water = smiles2conformers("O", 1)
     # pack a cubic box