Merge pull request #500 from avcopan/dev

Normal mode visualization + electron count

automol/extern/py3dmol_.py

@@ -1,5 +1,6 @@
 """ py3Dmol interface
 """
+
 from typing import Tuple
 
 import numpy
@@ -17,6 +18,33 @@ def create_view(image_size: int = 400) -> py3Dmol.view:
     return py3Dmol.view(width=image_size, height=image_size)
 
 
+def view_molecule_from_xyz(
+    xyz_str: str, view: py3Dmol.view = None, image_size: int = 400, vib: bool = False
+) -> py3Dmol.view:
+    """Visualize a molecule in a 3D view, using a MOLFile string
+
+    :param xyz_str: MOLFile block string
+    :param view: An existing 3D view to append to, defaults to None
+    :param image_size: The image size, if creating a new view, defaults to 400
+    :param vib: Animate vibrations for the molecule?
+    :return: A 3D view containing the molecule
+    :rtype: py3Dmol.view
+    """
+    if view is None:
+        view = create_view(image_size=image_size)
+
+    options = {"vibrate": {"frames": 10, "amplitude": 1}} if vib else {}
+
+    view.addModel(xyz_str, "xyz", options)
+    view.setStyle({"stick": {}, "sphere": {"radius": 0.3}})
+    view.setBackgroundColor("0xeeeeee")
+    if vib:
+        view.animate({"loop": "backAndForth"})
+
+    view.zoomTo()
+    return view
+
+
 def view_molecule_from_molfile(
     mfl: str, view: py3Dmol.view = None, image_size: int = 400
 ) -> py3Dmol.view:

automol/geom/__init__.py

@@ -43,6 +43,7 @@
 from automol.geom.base._0core import is_diatomic
 from automol.geom.base._0core import is_linear
 from automol.geom.base._0core import atom_count
+from automol.geom.base._0core import electron_count
 from automol.geom.base._0core import atom_indices
 from automol.geom.base._0core import dummy_atom_indices
 from automol.geom.base._0core import masses
@@ -177,6 +178,7 @@
     'is_diatomic',
     'is_linear',
     'atom_count',
+    'electron_count',
     'atom_indices',
     'dummy_atom_indices',
     'masses',

automol/geom/_conv.py

@@ -7,8 +7,8 @@
 
 import numpy
 import pyparsing as pp
+from numpy.typing import ArrayLike
 from pyparsing import pyparsing_common as ppc
-from phydat import phycon
 
 from automol import vmat
 from automol.extern import molfile, py3dmol_, rdkit_
@@ -28,11 +28,13 @@
     subgeom,
     symbols,
     translate,
+    xyz_string,
 )
 from automol.graph import base as graph_base
 from automol.inchi import base as inchi_base
 from automol.util import ZmatConv, dict_, heuristic, vector, zmat_conv
 from automol.zmat import base as zmat_base
+from phydat import phycon
 
 
 # # conversions
@@ -547,26 +549,29 @@ def rdkit_molecule(geo, gra=None, stereo=True):
     return rdkit_.from_geometry_with_graph(geo, gra)
 
 
-def py3dmol_view(geo, gra=None, view=None, image_size=400):
+def py3dmol_view(
+    geo, gra=None, view=None, image_size: int = 400, mode: Optional[ArrayLike] = None
+):
     """Get a py3DMol view of this molecular geometry
 
     :param geo: molecular geometry
     :type geo: automol geometry data structure
     :param gra: A molecular graph, describing the connectivity, defaults to None
     :type gra: automol graph data structure, optional
-    :param image_size: The image size, if creating a new view, defaults to 400
-    :type image_size: int, optional
     :param view: An existing 3D view to append to, defaults to None
     :type view: py3Dmol.view, optional
+    :param image_size: The image size, if creating a new view, defaults to 400
+    :param mode: A vibrational mode or molecular motion to visualize
     :return: A 3D view containing the molecule
     :rtype: py3Dmol.view
     """
     if gra is not None and graph_base.is_ts_graph(gra):
         gra = graph_base.ts.reactants_graph(gra, stereo=False, dummy=True)
 
-    rdm = rdkit_molecule(geo, gra=gra, stereo=False)
-    mlf = rdkit_.to_molfile(rdm)
-    return py3dmol_.view_molecule_from_molfile(mlf, view=view, image_size=image_size)
+    xyz_str = xyz_string(geo, mode=mode)
+    return py3dmol_.view_molecule_from_xyz(
+        xyz_str, view=view, image_size=image_size, vib=(mode is not None)
+    )
 
 
 def display(
@@ -575,18 +580,19 @@ def display(
     view=None,
     image_size=400,
     vis_bkeys: Optional[tuple[tuple[int, int]]] = None,
+    mode: Optional[ArrayLike] = None,
 ):
     """Display molecule to IPython using the RDKit visualizer
 
     :param geo: molecular geometry
     :type geo: automol geometry data structure
     :param gra: A molecular graph, describing the connectivity
     :type gra: automol graph data structure
-    :param image_size: The image size, if creating a new view, defaults to 400
-    :type image_size: int, optional
-    :param vis_bkeys: Only visualize these bonds, by key
     :param view: An existing 3D view to append to, defaults to None
     :type view: py3Dmol.view, optional
+    :param image_size: The image size, if creating a new view, defaults to 400
+    :param vis_bkeys: Only visualize these bonds, by key
+    :param mode: A vibrational mode or molecular motion to visualize
     """
     ts_ = gra is not None and graph_base.is_ts_graph(gra)
     if ts_:
@@ -600,12 +606,12 @@ def display(
         excl_bkeys = graph_base.bond_keys(gra) - set(map(frozenset, vis_bkeys))
         gra = graph_base.remove_bonds(gra, excl_bkeys, stereo=False, check=False)
 
-    view = py3dmol_view(geo, gra=gra, view=view, image_size=image_size)
+    view = py3dmol_view(geo, gra=gra, view=view, image_size=image_size, mode=mode)
 
     if ts_:
         for frm_bkey in graph_base.ts.forming_bond_keys(tsg):
             fidx1, fidx2 = frm_bkey
-            fxyz1, fxyz2 = coordinates(geo, idxs=(fidx1, fidx2))
+            fxyz1, fxyz2 = coordinates(geo, idxs=(fidx1, fidx2), angstrom=True)
             rvec1 = ts_reacting_electron_direction(geo, tsg, fidx1)
             rvec2 = ts_reacting_electron_direction(geo, tsg, fidx2)
             view = py3dmol_.view_vector(rvec1, orig_xyz=fxyz1, view=view)

automol/geom/base/_0core.py

@@ -1,13 +1,15 @@
 """
     Core functions defining the geometry data type
 """
+
 import itertools
-from typing import List
+from typing import List, Optional
 
 import more_itertools as mit
 import numpy
 import pyparsing as pp
 from pyparsing import pyparsing_common as ppc
+from numpy.typing import ArrayLike
 from phydat import phycon, ptab
 
 from automol import form, util
@@ -16,8 +18,10 @@
 
 CHAR = pp.Char(pp.alphas)
 SYMBOL = pp.Combine(CHAR + pp.Opt(CHAR))
-XYZ_LINE = pp.Group(SYMBOL + pp.Group(ppc.fnumber * 3))
-XYZ_LINES = pp.delimitedList(XYZ_LINE, delim=pp.lineEnd())
+XYZ_LINE = pp.Group(
+    SYMBOL + pp.Group(ppc.fnumber * 3) + pp.Suppress(... + pp.LineEnd())
+)
+XYZ_LINES = pp.delimitedList(XYZ_LINE, delim=pp.LineStart())
 
 
 # # constructors
@@ -144,46 +148,51 @@ def set_coordinates(geo, xyz_dct, angstrom=False):
 
 
 # # I/O
-def string(geo, angstrom=True):
+def string(geo, angstrom: bool = True, mode: Optional[ArrayLike] = None):
     """Write a molecular geometry to a string:
        symb1  xyz1 xyz2 xyz3
        symbn  xyzn xyzn xyzn
 
     :param geo: molecular geometry
     :type geo: automol molecular geometry data structure
     :param angstrom: parameter to control coordinate conversion to Angstrom
-    :type angstrom: bool
+    :param mode: A vibrational mode or molecular motion to visualize
     :rtype: str
     """
 
     symbs = symbols(geo)
     xyzs = coordinates(geo, angstrom=angstrom)
 
     natms = len(symbs)
-    assert len(xyzs) == natms
 
-    geo_str = "\n".join(
-        f"{symb:2s} {xyz[0]:10.6f} {xyz[1]:10.6f} {xyz[2]:10.6f}"
-        for symb, xyz in zip(symbs, xyzs)
-    )
+    symb_strs = [f"{s:2s}" for s in symbs]
+    xyz_strs = [f"{x[0]:10.6f} {x[1]:10.6f} {x[2]:10.6f}" for x in xyzs]
+    lines = [" ".join([s, x]) for s, x in zip(symb_strs, xyz_strs)]
+
+    # If requested, include the mode in the string
+    if mode is not None:
+        assert len(lines) == natms, f"Invalid mode for geometry:\n{mode}\n{geo}"
+        mode_strs = [f"{m[0]:10.6f} {m[1]:10.6f} {m[2]:10.6f}" for m in mode]
+        lines = ["  ".join([s, x]) for s, x in zip(lines, mode_strs)]
+
+    geo_str = "\n".join(lines)
 
     return geo_str
 
 
-def xyz_string(geo, comment=""):
+def xyz_string(geo, comment="", mode: Optional[ArrayLike] = None) -> str:
     """Write a molecular geometry to a string:
        natom
        comment
        symb1  xyz1 xyz2 xyz3
        symbn  xyzn xyzn xyzn
 
     :param geo: molecular geometry
-    :type geo: automol molecular geometry data structure
     :param comment: string to place in the comment line of string
-    :type comment: str
+    :param mode: A vibrational mode or molecular motion to visualize
     :rtype: str
     """
-    geo_str = string(geo, angstrom=True)
+    geo_str = string(geo, angstrom=True, mode=mode)
     xyz_str = f" {count(geo):d}\n{comment:s}\n{geo_str:s}"
     return xyz_str
 
@@ -434,6 +443,15 @@ def atom_count(geo, symb, match=True):
     return len(atom_indices(geo, symb, match=match))
 
 
+def electron_count(geo) -> int:
+    """Count the number of electrons in the geometry
+
+    :param geo: A molecular geometry
+    :return: The number of electrons
+    """
+    return form.electron_count(formula(geo))
+
+
 def atom_indices(geo, symb, match=True):
     """Obtain the indices of a atoms of a particular type in the geometry.
 

automol/geom/base/__init__.py

@@ -35,6 +35,7 @@
 from automol.geom.base._0core import is_diatomic
 from automol.geom.base._0core import is_linear
 from automol.geom.base._0core import atom_count
+from automol.geom.base._0core import electron_count
 from automol.geom.base._0core import atom_indices
 from automol.geom.base._0core import dummy_atom_indices
 from automol.geom.base._0core import masses
@@ -120,6 +121,7 @@
     'is_diatomic',
     'is_linear',
     'atom_count',
+    'electron_count',
     'atom_indices',
     'dummy_atom_indices',
     'masses',

automol/tests/test_geom.py

@@ -790,9 +790,9 @@ def test__repulsion_energy():
     # test__chi_with_sort()
     # test__argunique_coulomb_spectrum()
     # test__rotation_properties()
-    test__insert_dummies()
+    # test__insert_dummies()
     # test__hydrogen_bonded_structure()
-    # test__from_string()
+    test__from_string()
     # test__from_xyz_string()
     # test__insert_dummies()
     # test__closest_unbonded_atoms()