Repository: openforcefield/openff-toolkit

Issue:
Add optional argument to Molecule.visualize() to suppress hydrogens
**Is your feature request related to a problem? Please describe.**
Frequently, `Molecule.visualize()` makes 2D diagrams that are too cluttered because explicit hydrogens are shown.
![image](https://user-images.githubusercontent.com/3656088/97918206-96b11c00-1d0a-11eb-983d-657ac6a027bf.png)

**Describe the solution you'd like**
In these circumstances, it would be useful to be able to provide an optional argument, such as `suppress_hydrogens=True` or `show_hydrogens=False`.

**Describe alternatives you've considered**
The alternative is to go through a toolkit, which is possible but less convenient.


Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/openforcefield/openff-toolkit/pull/1001?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openforcefield) Report
> Merging [#1001](https://codecov.io/gh/openforcefield/openff-toolkit/pull/1001?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openforcefield) (b763734) into [master](https://codecov.io/gh/openforcefield/openff-toolkit/commit/42a4764cd582a1b8201c301527d4382547e03855?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=openforcefield) (42a4764) will **decrease** coverage by `8.64%`.
> The diff coverage is `7.69%`.

> :exclamation: Current head b763734 differs from pull request most recent head f469453. Consider uploading reports for the commit f469453 to get more accurate results

j-wags: This looks really cool, and I'm glad we're closing that old issue. Based on our conversation in today's meeting, could you additionally:

* Rename the kwarg to `show_hydrogens`
* Make the default be to `show_hydrogens=True`
* Implement this behavior as well for the OE and NGLView backends

Also, I think @mattwthompson is more familiar with the drawing code than me, so he may be a better reviewer (once you're ready for someone to take a look!)
adalke: The drawing code looks fine to me.
j-wags: A clearer example of why this change is needed:

Using the current codebase

```python
from openff.toolkit.topology import Molecule
mol = Molecule.from_smiles('c1ccccc1COCc2ccccc2')
mol.generate_conformers()
mol.visualize(backend='rdkit')
```
![image](https://user-images.githubusercontent.com/16329175/124367581-42e63200-dc0d-11eb-96e9-4a4fd99f0573.png)

mattwthompson: I'm far from an expert in these features (when I worked on #560 etc. I was pretty much learning it at the same time) but I'd be happy to review once this is ready. (It might be now?)
j-wags: Thanks @mattwthompson. I'm going to take over finishing this so that @Yoshanuikabundi can move on to learning Interchange in a week or two. I'll still need to thread the `show_hydrogens` kwarg through to OETKW, but I _will_ ping you for a review once that's done :-) 
Yoshanuikabundi: @j-wags I've implemented the changes we discussed. I think this is ready for review! @mattwthompson 

- I haven't implemented `show_hydrogens` for nglview, because it would basically involving rewriting the code that nglview uses to construct the default representations. But I can!
- I renamed `show_hydrogens` to `show_all_hydrogens`, because both toolkits sometimes label terminal groups as "NH2" or whatever. `show_all_hydrogens=True` means that every hydrogen is in the picture, but `False` doesn't mean that none of them are
- I've made using ignored arguments an error (eg, `show_all_hydrogens` with the `"nglview"` backend). This is technically a breaking change, but Jeff I think you mentioned it in our call.
- I've added the ignored arguments error behaviour to the tests, but not anything else - I'm not sure how to test visualization.
- I've fixed a major bug in `Molecule.to_rdkit()`, where rdkit could add implicit hydrogens to molecules created from the toolkit. This was particularly pathological with radicals, which lost their radicalness on export. I've written up this change in commit a0263f9. The fix is one line, but it means that rdkit will not try to guess about the existence of any hydrogens not included explicitly.
mattwthompson: I think this is good to go once the single blocking comment is resolved. I did make some changes myself; feel free to poke through the "resolved" conversations if anything is disagreeable.

Below is the notebook I used while reviewing; everything looked good on this branch (including expected failures) and a couple of the RDKit cells looked weird using `upstream/master`, weird in similar ways to #759. I think that issue can be closed if this is merged; or at least I didn't see obvious things wrong with OpenEye-based visualization.
[1001-review.ipynb.txt](https://github.com/openforcefield/openff-toolkit/files/6994317/1001-review.ipynb.txt)


Patch:
diff --git a/docs/releasehistory.md b/docs/releasehistory.md
index 6e52d1c0..0b030490 100644
--- a/docs/releasehistory.md
+++ b/docs/releasehistory.md
@@ -14,6 +14,10 @@ Releases follow the `major.minor.micro` scheme recommended by [PEP440](https://w
   logic for library charges was updated to use only one unique match instead of
   enumerating all possible matches. This results in faster matching, particularly
   with larger moldules.
+- [PR #1001](https://github.com/openforcefield/openff-toolkit/pull/1001): Revamped the 
+  [`Molecule.visualize()`](openff.toolkit.topology.Molecule.visualize) method's `rdkit` 
+  backend for more pleasing and idiomatic 2D visualization by default.
+
 
 ### Bugfixes
 
@@ -28,6 +32,10 @@ Releases follow the `major.minor.micro` scheme recommended by [PEP440](https://w
   logic for library charges was updated to use only one unique match. No adverse side effects
   were found in testing, but could bad behavior may possibly exist in some unknown caes.
   Note that the default behavior for other parameter handlers was not updated.
+- [PR #1001](https://github.com/openforcefield/openff-toolkit/pull/1001): RDKit `Mol` objects 
+  created through the [`Molecule.to_rdkit()`](openff.toolkit.topology.Molecule.to_rdkit)
+  method have the `NoImplicit` property set to `True` on all atoms. This prevents RDKit from
+  incorrectly adding hydrogen atoms to to molecule. 
 
 ### Tests updated
 
diff --git a/openff/toolkit/topology/molecule.py b/openff/toolkit/topology/molecule.py
index b0aa96b6..5f620a6e 100644
--- a/openff/toolkit/topology/molecule.py
+++ b/openff/toolkit/topology/molecule.py
@@ -6290,7 +6290,13 @@ class Molecule(FrozenMolecule):
 
         return self._add_conformer(coordinates)
 
-    def visualize(self, backend="rdkit", width=500, height=300):
+    def visualize(
+        self,
+        backend="rdkit",
+        width=None,
+        height=None,
+        show_all_hydrogens=True,
+    ):
         """
         Render a visualization of the molecule in Jupyter
 
@@ -6305,17 +6311,21 @@ class Molecule(FrozenMolecule):
 
         width : int, optional, default=500
             Width of the generated representation (only applicable to
-            ``backend=openeye``)
+            ``backend=openeye`` or ``backend=rdkit``)
         height : int, optional, default=300
             Width of the generated representation (only applicable to
-            ``backend=openeye``)
+            ``backend=openeye`` or ``backend=rdkit``)
+        show_all_hydrogens : bool, optional, default=True
+            Whether to explicitly depict all hydrogen atoms. (only applicable to
+            ``backend=openeye`` or ``backend=rdkit``)
 
         Returns
         -------
         object
             Depending on the backend chosen:
 
-            - rdkit, openeye → IPython.display.Image
+            - rdkit → IPython.display.SVG
+            - openeye → IPython.display.Image
             - nglview → nglview.NGLWidget
 
         """
@@ -6328,6 +6338,19 @@ class Molecule(FrozenMolecule):
                 import nglview as nv
             except ImportError:
                 raise MissingDependencyError("nglview")
+
+            if width is not None or height is not None:
+                # TODO: More specific exception
+                raise ValueError(
+                    "The width, height, and show_all_hydrogens arguments do not apply to the nglview backend."
+                )
+            elif not show_all_hydrogens:
+                # TODO: More specific exception
+                # TODO: Implement this? Should be able to just strip hydrogens from the PDB
+                raise ValueError(
+                    "show_all_hydrogens=False is not supported by the nglview backend"
+                )
+
             if self.conformers:
                 from openff.toolkit.utils.viz import _OFFTrajectoryNGLView
 
@@ -6335,20 +6358,43 @@ class Molecule(FrozenMolecule):
                 widget = nv.NGLWidget(trajectory_like)
                 return widget
             else:
+                # TODO: More specific exception
                 raise ValueError(
                     "Visualizing with NGLview requires that the molecule has "
                     "conformers."
                 )
+
+        width = 500 if width is None else width
+        height = 300 if height is None else height
+        show_all_hydrogens = True if show_all_hydrogens is None else show_all_hydrogens
+
         if backend == "rdkit":
             if RDKIT_AVAILABLE:
-                from rdkit.Chem.Draw import IPythonConsole
+                from IPython.display import SVG
+                from rdkit.Chem.Draw import rdDepictor, rdMolDraw2D
+                from rdkit.Chem.rdmolops import RemoveHs
+
+                rdmol = self.to_rdkit()
+
+                if not show_all_hydrogens:
+                    # updateExplicitCount: Keep a record of the hydrogens we remove.
+                    # This is used in visualization to distinguish eg radicals from normal species
+                    rdmol = RemoveHs(rdmol, updateExplicitCount=True)
 
-                return self.to_rdkit()
+                rdDepictor.SetPreferCoordGen(True)
+                rdDepictor.Compute2DCoords(rdmol)
+                rdmol = rdMolDraw2D.PrepareMolForDrawing(rdmol)
+
+                drawer = rdMolDraw2D.MolDraw2DSVG(width, height)
+                drawer.DrawMolecule(rdmol)
+                drawer.FinishDrawing()
+
+                return SVG(drawer.GetDrawingText())
             else:
                 warnings.warn(
                     "RDKit was requested as a visualization backend but "
                     "it was not found to be installed. Falling back to "
-                    "trying to using OpenEye for visualization."
+                    "trying to use OpenEye for visualization."
                 )
                 backend = "openeye"
         if backend == "openeye":
@@ -6362,6 +6408,9 @@ class Molecule(FrozenMolecule):
                     width, height, oedepict.OEScale_AutoScale
                 )
 
+                if show_all_hydrogens:
+                    opts.SetHydrogenStyle(oedepict.OEHydrogenStyle_ImplicitAll)
+
                 oedepict.OEPrepareDepiction(oemol)
                 img = oedepict.OEImage(width, height)
                 display = oedepict.OE2DMolDisplay(oemol, opts)
@@ -6369,6 +6418,7 @@ class Molecule(FrozenMolecule):
                 png = oedepict.OEWriteImageToString("png", img)
                 return Image(png)
 
+        # TODO: More specific exception
         raise ValueError("Could not find an appropriate backend")
 
     def _ipython_display_(self):
diff --git a/openff/toolkit/utils/rdkit_wrapper.py b/openff/toolkit/utils/rdkit_wrapper.py
index 5f3ea704..3d889a98 100644
--- a/openff/toolkit/utils/rdkit_wrapper.py
+++ b/openff/toolkit/utils/rdkit_wrapper.py
@@ -1683,6 +1683,9 @@ class RDKitToolkitWrapper(base_wrapper.ToolkitWrapper):
             elif atom.stereochemistry == "R":
                 rdatom.SetChiralTag(Chem.CHI_TETRAHEDRAL_CCW)
 
+            # Stop rdkit from adding implicit hydrogens
+            rdatom.SetNoImplicit(True)
+
             rd_index = rdmol.AddAtom(rdatom)
 
             # Let's make sure al the atom indices in the two molecules


Test patch:
diff --git a/openff/toolkit/tests/test_molecule.py b/openff/toolkit/tests/test_molecule.py
index 55f5581b..3ed0ecfb 100644
--- a/openff/toolkit/tests/test_molecule.py
+++ b/openff/toolkit/tests/test_molecule.py
@@ -3678,15 +3678,17 @@ class TestMolecule:
         assert len([atom for atom in mol.atoms if atom.is_in_ring]) == n_atom_rings
         assert len([bond for bond in mol.bonds if bond.is_in_ring]) == n_bond_rings
 
-    @requires_pkg("ipython")
+
+class TestMoleculeVisualization:
+    @requires_pkg("IPython")
     @requires_rdkit
     def test_visualize_rdkit(self):
         """Test that the visualize method returns an expected object when using RDKit to generate a 2-D representation"""
-        import rdkit
+        import IPython
 
         mol = Molecule().from_smiles("CCO")
 
-        assert isinstance(mol.visualize(backend="rdkit"), rdkit.Chem.rdchem.Mol)
+        assert isinstance(mol.visualize(backend="rdkit"), IPython.core.display.SVG)
 
     @pytest.mark.skipif(
         has_pkg("rdkit"),
@@ -3720,7 +3722,15 @@ class TestMolecule:
         # Ensure an NGLView widget is returned
         assert isinstance(mol.visualize(backend="nglview"), nglview.NGLWidget)
 
-    @requires_pkg("ipython")
+        # Providing other arguments is an error
+        with pytest.raises(ValueError):
+            mol.visualize(backend="nglview", width=100)
+        with pytest.raises(ValueError):
+            mol.visualize(backend="nglview", height=100)
+        with pytest.raises(ValueError):
+            mol.visualize(backend="nglview", show_all_hydrogens=False)
+
+    @requires_pkg("IPython")
     @requires_openeye
     def test_visualize_openeye(self):
         """Test that the visualize method returns an expected object when using OpenEye to generate a 2-D representation"""