Visualize fixes

docs/releasehistory.md

@@ -35,7 +35,8 @@ Releases follow the `major.minor.micro` scheme recommended by [PEP440](https://w
 
 ### New features
 
-- [PR #1731](https://github.com/openforcefield/openff-toolkit/pull/1731): Suppot SMIRNOFF vdW version 0.5.
+- [PR #1731](https://github.com/openforcefield/openff-toolkit/pull/1731): Support SMIRNOFF vdW version 0.5.
+- [PR #1751](https://github.com/openforcefield/openff-toolkit/pull/1751): Improve visualization API docs and support multiple bonds in `Topology.visualize()`
 
 ### Improved documentation and warnings
 

openff/toolkit/topology/molecule.py

@@ -87,7 +87,7 @@
 from openff.toolkit.utils.utils import get_data_file_path, requires_package
 
 if TYPE_CHECKING:
-    import IPython.display
+    import IPython.core.display
     import nglview
 
     from openff.toolkit.topology._mm_molecule import _SimpleAtom, _SimpleMolecule
@@ -5394,61 +5394,65 @@ def add_conformer(self, coordinates: Quantity) -> int:
     @overload
     def visualize(
         self,
-        backend: Literal["rdkit"],
-    ) -> "IPython.display.SVG":
+        backend: Literal["rdkit"] = ...,
+        width: int = ...,
+        height: int = ...,
+        show_all_hydrogens: bool = ...,
+    ) -> "IPython.core.display.SVG":
         ...
 
     @overload
     def visualize(
         self,
         backend: Literal["openeye"],
-    ) -> "IPython.display.Image":
+        width: int = ...,
+        height: int = ...,
+        show_all_hydrogens: bool = ...,
+    ) -> "IPython.core.display.Image":
         ...
 
     @overload
     def visualize(
         self,
         backend: Literal["nglview"],
+        *,
+        show_all_hydrogens: bool = ...,
     ) -> "nglview.NGLWidget":
         ...
 
     def visualize(
         self,
-        backend: str = "rdkit",
-        width: int = 500,
-        height: int = 300,
-        show_all_hydrogens: bool = True,
-    ) -> Union["IPython.display.SVG", "IPython.display.Image", "nglview.NGLWidget"]:
+        backend="rdkit",
+        width=500,
+        height=300,
+        show_all_hydrogens=True,
+    ):
         """
-        Render a visualization of the molecule in Jupyter
+        Render a visualization of the molecule in Jupyter.
+
+        Note that the ``"nglview"`` backend may, in strained conformations,
+        include bonds not present in the topology.
 
         Parameters
         ----------
-        backend : str, optional, default='rdkit'
+        backend
             The visualization engine to use. Choose from:
 
-            - ``"rdkit"``
+            - ``"rdkit"`` (default)
             - ``"openeye"``
             - ``"nglview"`` (requires conformers)
 
-        width : int, default=500
-            Width of the generated representation (only applicable to
+        width
+            Width of the generated representation in pixels (only applicable to
             ``backend="openeye"`` or ``backend="rdkit"``)
-        height : int, default=300
-            Width of the generated representation (only applicable to
-            ``backend="openeye"`` or ``backend="rdkit"``)
-        show_all_hydrogens : bool, default=True
-            Whether to explicitly depict all hydrogen atoms. (only applicable to
+        height
+            Width of the generated representation in pixels (only applicable to
             ``backend="openeye"`` or ``backend="rdkit"``)
+        show_all_hydrogens
+            Whether to explicitly depict all hydrogen atoms.
 
         Returns
         -------
-        object
-            Depending on the backend chosen:
-
-            - rdkit → IPython.display.SVG
-            - openeye → IPython.display.Image
-            - nglview → nglview.NGLWidget
 
         """
         import inspect
@@ -5469,47 +5473,46 @@ def visualize(
             ):
                 warnings.warn(
                     f"Arguments `width` and `height` are ignored with {backend=}."
-                    f"Found non-default values {width=} and {height=}",
+                    + f"Found non-default values {width=} and {height=}",
                     stacklevel=2,
                 )
 
             if self.conformers is None:
                 raise MissingConformersError(
                     "Visualizing with NGLview requires that the molecule has "
-                    f"conformers, found {self.conformers=}"
+                    + f"conformers, found {self.conformers=}"
                 )
 
-            else:
-                from openff.toolkit.utils._viz import MoleculeNGLViewTrajectory
+            from openff.toolkit.utils._viz import MoleculeNGLViewTrajectory
 
-                try:
-                    widget = nv.NGLWidget(
-                        MoleculeNGLViewTrajectory(
-                            molecule=self,
-                            ext="MOL2",
-                        )
+            try:
+                widget = nv.NGLWidget(
+                    MoleculeNGLViewTrajectory(
+                        molecule=self,
+                        ext="MOL2",
                     )
-                except ValueError:
-                    widget = nv.NGLWidget(
-                        MoleculeNGLViewTrajectory(
-                            molecule=self,
-                            ext="PDB",
-                        )
+                )
+            except ValueError:
+                widget = nv.NGLWidget(
+                    MoleculeNGLViewTrajectory(
+                        molecule=self,
+                        ext="PDB",
                     )
-
-                widget.clear_representations()
-                widget.add_representation(
-                    "licorice",
-                    sele="*" if show_all_hydrogens else "NOT hydrogen",
-                    radius=0.25,
-                    multipleBond=True,
                 )
 
-                return widget
+            widget.clear_representations()
+            widget.add_representation(
+                "licorice",
+                sele="*" if show_all_hydrogens else "NOT hydrogen",
+                radius=0.25,
+                multipleBond=True,
+            )
+
+            return widget
 
         if backend == "rdkit":
             if RDKIT_AVAILABLE:
-                from IPython.display import SVG
+                from IPython.core.display import SVG
                 from rdkit.Chem.Draw import (  # type: ignore[import-untyped]
                     rdDepictor,
                     rdMolDraw2D,
@@ -5535,14 +5538,14 @@ def visualize(
             else:
                 warnings.warn(
                     "RDKit was requested as a visualization backend but "
-                    "it was not found to be installed. Falling back to "
-                    "trying to use OpenEye for visualization.",
+                    + "it was not found to be installed. Falling back to "
+                    + "trying to use OpenEye for visualization.",
                     stacklevel=2,
                 )
                 backend = "openeye"
         if backend == "openeye":
             if OPENEYE_AVAILABLE:
-                from IPython.display import Image
+                from IPython.core.display import Image
                 from openeye import oedepict
 
                 oemol = self.to_openeye()

openff/toolkit/topology/topology.py

@@ -2402,24 +2402,16 @@ def is_constrained(self, iatom, jatom):
             return False
 
     @requires_package("nglview")
-    def visualize(self, ensure_correct_connectivity: bool = False) -> "NGLWidget":
+    def visualize(self) -> "NGLWidget":
         """
         Visualize with NGLView.
 
-        Requires all molecules in this topology have positions.
+        Requires that all molecules in this topology have positions.
 
         NGLView is a 3D molecular visualization library for use in Jupyter
-        notebooks. Note that for performance reasons, by default the
-        visualized connectivity is inferred from positions and may not reflect
-        the connectivity in the ``Topology``.
-
-        Parameters
-        ==========
-
-        ensure_correct_connectivity: bool, default=False
-            If ``True``, the visualization will be guaranteed to reflect the
-            connectivity in the ``Topology``. Note that this will severely
-            degrade performance, especially for topologies with many atoms.
+        notebooks. Note that the visualized molecule may, in strained
+        conformations, include bonds not present in the topology. Rendering
+        multiple bonds additionally requires RDKit.
 
         Examples
         ========
@@ -2437,22 +2429,14 @@ def visualize(self, ensure_correct_connectivity: bool = False) -> "NGLWidget":
 
         from openff.toolkit.utils._viz import TopologyNGLViewStructure
 
-        if ensure_correct_connectivity:
-            raise ValueError(
-                "`ensure_correct_connectivity` not (yet) implemented "
-                "(requires passing multi-molecule SDF files to NGLview)"
-            )
-
         if self.get_positions() is None:
             raise MissingConformersError(
-                "All molecules in this topology must have positions for it to be visualized in a widget."
+                "All molecules in this topology must have positions for it to "
+                + "be visualized in a widget."
             )
 
         widget = nglview.NGLWidget(
-            TopologyNGLViewStructure(
-                topology=self,
-                ext="pdb",
-            ),
+            TopologyNGLViewStructure(topology=self),
             representations=[
                 dict(type="unitcell", params=dict()),
             ],
@@ -2465,7 +2449,7 @@ def visualize(self, ensure_correct_connectivity: bool = False) -> "NGLWidget":
             "licorice",
             sele="not water and not ion and not protein",
             radius=0.25,
-            multipleBond=bool(ensure_correct_connectivity),
+            multipleBond=True,
         )
 
         return widget

openff/toolkit/utils/_viz.py

@@ -1,9 +1,12 @@
+import math
 import uuid
 from io import StringIO
+from typing import TextIO
 
 from nglview.base_adaptor import Structure, Trajectory
 
 from openff.toolkit import Molecule, Topology, unit
+from openff.toolkit.utils.toolkits import RDKIT_AVAILABLE
 
 MOLECULE_DEFAULT_REPS = [
     dict(type="licorice", params=dict(radius=0.25, multipleBond=True))
@@ -22,7 +25,7 @@ class MoleculeNGLViewTrajectory(Structure, Trajectory):
     Parameters
     ----------
     molecule
-        The `Molecule` object to display.
+        The ``Molecule`` object to display.
     ext
         The file extension to use to communicate with NGLView. The format must
         be supported for export by the Toolkit via the `Molecule.to_file()
@@ -70,16 +73,14 @@ class TopologyNGLViewStructure(Structure):
     """
     OpenFF Topology adaptor.
 
+    Communicates with NGLView via PDB, using RDKit to write redundant CONECT
+    records indicating multiple bonds. If RDKit is unavailable, falls back
+    to ``Topology.to_file``.
+
     Parameters
     ----------
     topology
-        The `Topology` object to display.
-    ext
-        The file extension to use to communicate with NGLView. The format must
-        be supported for export by the Toolkit via the `Topology.to_file()
-        <openff.toolkit.topology.Topology.to_file>` method, and import by
-        NGLView. File formats supported by NGLView can be found at
-        https://nglviewer.org/ngl/api/manual/file-formats.html
+        The ``Topology`` object to display.
 
     Example
     -------
@@ -92,15 +93,54 @@ class TopologyNGLViewStructure(Structure):
     def __init__(
         self,
         topology: Topology,
-        ext: str = "PDB",
     ):
         self.topology = topology
-        self.ext = ext.lower()
+        self.ext = "pdb"
         self.params: dict = dict()
         self.id = str(uuid.uuid4())
 
     def get_structure_string(self):
         with StringIO() as f:
-            self.topology.to_file(f, file_format=self.ext)
-            structure_string = f.getvalue()
+            if RDKIT_AVAILABLE:
+                from rdkit.Chem.rdmolfiles import PDBWriter  # type: ignore
+
+                write_box_vectors(f, self.topology)
+
+                writer: PDBWriter = PDBWriter(f)
+                for mol in self.topology.molecules:
+                    writer.write(mol.to_rdkit())
+
+                writer.close()
+
+                structure_string = f.getvalue()
+            else:
+                self.topology.to_file(f, file_format="pdb")
+                structure_string = f.getvalue()
+
         return structure_string
+
+
+def write_box_vectors(file_obj: TextIO, topology: Topology):
+    if topology.box_vectors is not None:
+        a, b, c = topology.box_vectors.m_as(unit.nanometer)
+        a_length = a.norm()
+        b_length = b.norm()
+        c_length = c.norm()
+
+        alpha = math.acos(b.dot(c) / (b_length * c_length))
+        beta = math.acos(c.dot(a) / (c_length * a_length))
+        gamma = math.acos(a.dot(b) / (a_length * b_length))
+
+        RAD_TO_DEG = 180 / math.pi
+        print(
+            "CRYST1%9.3f%9.3f%9.3f%7.2f%7.2f%7.2f P 1           1 "
+            % (
+                a_length * 10,
+                b_length * 10,
+                c_length * 10,
+                alpha * RAD_TO_DEG,
+                beta * RAD_TO_DEG,
+                gamma * RAD_TO_DEG,
+            ),
+            file=file_obj,
+        )