Merge pull request #276 from rdguha1995/main

First pass at bond order visualizations for molecules [WIP]

Author: rdguha1995

crystal_toolkit/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from pathlib import Path
 from importlib.metadata import PackageNotFoundError, version
+from pathlib import Path
 
 from monty.json import MSONable
 

crystal_toolkit/core/scene.py

@@ -316,7 +316,7 @@ class Cylinders(Primitive):
 
     @property
     def key(self):
-        return f"cylinder_{self.color}_{self.radius}_{self.reference}"
+        return f"cylinder_{self.color}_{self.radius}_{self.reference}_{self.clickable}_{self.tooltip}"
 
     @classmethod
     def merge(cls, cylinder_list):
@@ -330,6 +330,8 @@ def merge(cls, cylinder_list):
             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,
         )
 
     @property

crystal_toolkit/renderables/moleculegraph.py

@@ -3,6 +3,7 @@
 from collections import defaultdict
 
 from pymatgen.analysis.graphs import MoleculeGraph
+from pymatgen.analysis.local_env import OpenBabelNN
 
 from crystal_toolkit.core.legend import Legend
 from crystal_toolkit.core.scene import Scene
@@ -16,21 +17,47 @@ def get_molecule_graph_scene(
     explicitly_calculate_polyhedra_hull=False,
     legend=None,
     draw_polyhedra=False,
+    show_atom_idx=True,
+    show_atom_coord=True,
+    show_bond_order=True,
+    show_bond_length=False,
+    visualize_bond_orders=False,
 ) -> Scene:
 
+    """
+    Args:
+        show_atom_idx: Defaults to True, shows the site index of each atom in the molecule
+        show_atom_coord: Defaults to True, shows the 3D coordinates of each atom in the molecule
+        show_bond_order: Defaults to True, shows the calculated bond order in the chosen local environment strategy
+        show_bond_length: Defaults to False, shows the calculated length between two connected atoms
+        visualize_bpnd_orders: Defaults False, will show the 'integral' number of bonds calculated from the OpenBabelNN strategy in the Molecule Graph
+    Returns:
+        A Molecule Graph scene
+    """
+
+    vis_mol_graph = MoleculeGraph.with_local_env_strategy(self.molecule, OpenBabelNN())
     legend = legend or Legend(self.molecule)
 
     primitives: dict[str, list] = defaultdict(list)
 
     for idx, site in enumerate(self.molecule):
 
-        connected_sites = self.get_connected_sites(idx)
+        if visualize_bond_orders:
+            connected_sites = vis_mol_graph.get_connected_sites(idx)
+        else:
+            connected_sites = self.get_connected_sites(idx)
 
         site_scene = site.get_scene(
+            site_idx=idx,
             connected_sites=connected_sites,
             origin=origin,
             explicitly_calculate_polyhedra_hull=explicitly_calculate_polyhedra_hull,
             legend=legend,
+            show_atom_idx=show_atom_idx,
+            show_atom_coord=show_atom_coord,
+            show_bond_order=show_bond_order,
+            show_bond_length=show_bond_length,
+            visualize_bond_orders=visualize_bond_orders,
             draw_polyhedra=draw_polyhedra,
         )
         for scene in site_scene.contents:

crystal_toolkit/renderables/site.py

@@ -28,6 +28,7 @@ def get_site_scene(
     # connected_sites_to_draw,
     connected_sites_not_drawn: list[ConnectedSite] = None,
     hide_incomplete_edges: bool = False,
+    site_idx: int | None = 0,
     incomplete_edge_length_scale: float | None = 1.0,
     connected_sites_colors: list[str] | None = None,
     connected_sites_not_drawn_colors: list[str] | None = None,
@@ -36,24 +37,27 @@ def get_site_scene(
     explicitly_calculate_polyhedra_hull: bool = False,
     bond_radius: float = 0.1,
     draw_magmoms: bool = True,
+    show_atom_idx: bool = False,
+    show_atom_coord: bool = True,
+    show_bond_order: bool = True,
+    show_bond_length: bool = False,
+    visualize_bond_orders: bool = False,
     magmom_scale: float = 1.0,
     legend: Legend | None = None,
 ) -> Scene:
     """
-
     Args:
         connected_sites:
         connected_sites_not_drawn:
         hide_incomplete_edges:
+        site_idx:
         incomplete_edge_length_scale:
         connected_sites_colors:
         connected_sites_not_drawn_colors:
         origin:
         explicitly_calculate_polyhedra_hull:
         legend:
-
     Returns:
-
     """
 
     atoms = []
@@ -74,7 +78,6 @@ def get_site_scene(
     max_radius = float(min(radii))
 
     for sp, occu in self.species.items():
-
         if isinstance(sp, DummySpecie):
 
             cube = Cubes(
@@ -101,7 +104,12 @@ def get_site_scene(
             name = str(sp)
             if occu != 1.0:
                 name += f" ({occu}% occupancy)"
-            name += f" ({position[0]:.3f}, {position[1]:.3f}, {position[2]:.3f})"
+
+            if show_atom_coord:
+                name += f" ({position[0]:.3f}, {position[1]:.3f}, {position[2]:.3f})"
+
+            if show_atom_idx:
+                name += "\n" + "index:" + str(site_idx)
 
             if self.properties:
                 for k, v in self.properties.items():
@@ -157,9 +165,20 @@ def get_site_scene(
         # TODO: can cause a bug if all vertices almost co-planar
         # necessary to include center site in case it's outside polyhedra
         all_positions = [self.coords]
+        name_cyl = " "
 
         for idx, connected_site in enumerate(connected_sites):
 
+            if show_bond_order:
+                if connected_site.weight is not None:
+                    name_cyl = "bond order:" + str(f"{connected_site.weight:.2f}")
+
+            if show_bond_length:
+                if connected_site.dist is not None:
+                    name_cyl += (
+                        "\n" + "bond length:" + str(f"{connected_site.dist:.3f}")
+                    )
+
             connected_position = connected_site.site.coords
             bond_midpoint = np.add(position, connected_position) / 2
 
@@ -168,12 +187,47 @@ def get_site_scene(
             else:
                 color = site_color
 
-            cylinder = Cylinders(
-                positionPairs=[[position, bond_midpoint.tolist()]],
-                color=color,
-                radius=bond_radius,
-            )
-            bonds.append(cylinder)
+            if visualize_bond_orders:
+                cylinders = []
+
+                if connected_site.weight is not None:
+
+                    if connected_site.weight > 1:
+                        trans_vector = 0.0
+                        for _bond in range(connected_site.weight):
+                            pos_r_1 = [i + trans_vector for i in position]
+                            pos_r_2 = [i + trans_vector for i in bond_midpoint.tolist()]
+                            cylinders.append(
+                                Cylinders(
+                                    positionPairs=[[pos_r_1, pos_r_2]],
+                                    color=color,
+                                    radius=bond_radius / 2,
+                                    clickable=True,
+                                    tooltip=name_cyl,
+                                )
+                            )
+                            trans_vector = trans_vector + 0.25 * max_radius
+                        for cylinder in cylinders:
+                            bonds.append(cylinder)
+                    else:
+                        cylinder = Cylinders(
+                            positionPairs=[[position, bond_midpoint.tolist()]],
+                            color=color,
+                            radius=bond_radius,
+                            clickable=True,
+                            tooltip=name_cyl,
+                        )
+                        bonds.append(cylinder)
+
+            else:
+                cylinder = Cylinders(
+                    positionPairs=[[position, bond_midpoint.tolist()]],
+                    color=color,
+                    radius=bond_radius,
+                    clickable=True,
+                    tooltip=name_cyl,
+                )
+                bonds.append(cylinder)
             all_positions.append(connected_position.tolist())
 
         if connected_sites_not_drawn and not hide_incomplete_edges:

jupyterlab-extension/tsconfig.tsbuildinfo

@@ -1190,4 +1190,4 @@
     ]
   },
   "version": "4.1.3"
-}
+}