Retain atom index info in _meta (#470)

* add atom index to _meta

* Revert "add atom index to _meta"

This reverts commit 56497f2.

* add index to _meta

* remove empty line

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

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

* ruff format

* add retain_atom_idx

---------

Co-authored-by: Chiu Peter <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Patrick Huck <[email protected]>

Author: 4 people

crystal_toolkit/core/scene.py

@@ -214,6 +214,7 @@ def merge(cls, sphere_list):
             visible=sphere_list[0].visible,
             clickable=sphere_list[0].clickable,
             tooltip=sphere_list[0].tooltip,
+            _meta=sphere_list[0]._meta,
         )
 
 
@@ -320,13 +321,18 @@ def merge(cls, cylinder_list):
             chain.from_iterable([cylinder.positionPairs for cylinder in cylinder_list])
         )
 
+        new_meta_list = list(
+            chain.from_iterable([[cylinder._meta] for cylinder in cylinder_list])
+        )
+
         return cls(
             positionPairs=new_positionPairs,
             color=cylinder_list[0].color,
             radius=cylinder_list[0].radius,
             visible=cylinder_list[0].visible,
             clickable=cylinder_list[0].clickable,
             tooltip=cylinder_list[0].tooltip,
+            _meta=new_meta_list,
         )
 
     @property

crystal_toolkit/renderables/site.py

@@ -47,6 +47,7 @@ def get_site_scene(
     visualize_bond_orders: bool = False,
     magmom_scale: float = 1.0,
     legend: Legend | None = None,
+    retain_atom_idx: bool = False,
 ) -> Scene:
     """Get a Scene object for a Site.
 
@@ -70,6 +71,7 @@ def get_site_scene(
         visualize_bond_orders (bool, optional): Defaults to False.
         magmom_scale (float, optional): Defaults to 1.0.
         legend (Legend | None, optional): Defaults to None.
+        retain_atom_idx (bool, optional): Defaults to False.
 
     Returns:
         Scene: The scene object containing atoms, bonds, polyhedra, magmoms.
@@ -135,6 +137,7 @@ def get_site_scene(
                 phiEnd=phiEnd,
                 clickable=True,
                 tooltip=name,
+                _meta=[site_idx] if retain_atom_idx else None,
             )
             atoms.append(sphere)
 
@@ -207,6 +210,9 @@ def get_site_scene(
                                     radius=bond_radius / 2,
                                     clickable=True,
                                     tooltip=name_cyl,
+                                    _meta=[site_idx, connected_site.index]
+                                    if retain_atom_idx
+                                    else None,
                                 )
                             )
                             trans_vector = trans_vector + 0.25 * max_radius
@@ -218,6 +224,9 @@ def get_site_scene(
                             radius=bond_radius,
                             clickable=True,
                             tooltip=name_cyl,
+                            _meta=[site_idx, connected_site.index]
+                            if retain_atom_idx
+                            else None,
                         )
                         bonds.append(cylinder)
 
@@ -228,6 +237,7 @@ def get_site_scene(
                     radius=bond_radius,
                     clickable=True,
                     tooltip=name_cyl,
+                    _meta=[site_idx, connected_site.index] if retain_atom_idx else None,
                 )
                 bonds.append(cylinder)
             all_positions.append(connected_position.tolist())
@@ -251,6 +261,7 @@ def get_site_scene(
                     positionPairs=[[position, bond_midpoint.tolist()]],
                     color=color,
                     radius=bond_radius,
+                    _meta=[site_idx, connected_site.index] if retain_atom_idx else None,
                 )
                 bonds.append(cylinder)
                 all_positions.append(connected_position.tolist())

crystal_toolkit/renderables/structuregraph.py

@@ -197,6 +197,8 @@ def get_weight_color(weight):
             explicitly_calculate_polyhedra_hull=explicitly_calculate_polyhedra_hull,
             legend=legend,
             bond_radius=bond_radius,
+            site_idx=idx,
+            show_atom_idx=True,
             **(site_get_scene_kwargs or {}),
         )