diff --git a/mbuild/compound.py b/mbuild/compound.py index e24bc8c05..5375d08e0 100644 --- a/mbuild/compound.py +++ b/mbuild/compound.py @@ -2,8 +2,6 @@ import itertools import logging -import os -import tempfile from collections import OrderedDict from collections.abc import Iterable from copy import deepcopy @@ -23,8 +21,7 @@ from mbuild.exceptions import MBuildError from mbuild.periodic_kdtree import PeriodicKDTree from mbuild.utils.geometry import bounding_box -from mbuild.utils.io import import_, run_from_ipython -from mbuild.utils.jsutils import overwrite_nglview_default +from mbuild.utils.io import import_ __all__ = ["clone", "Compound", "Particle"] @@ -1831,215 +1828,21 @@ def visualize( periodic_bond_opacity : bool, float, Optional, default=False Specify as a float from 0 to 1 to set the bond opacity for bonds that cross periodic boundaries. - """ - viz_pkg = { - "nglview": self._visualize_nglview, - "py3dmol": self._visualize_py3dmol, - } - if run_from_ipython(): - if backend.lower() in viz_pkg: - if backend.lower == "nglview": - return viz_pkg[backend.lower()](show_ports=show_ports) - else: - return viz_pkg[backend.lower()]( - show_ports=show_ports, - color_scheme=color_scheme, - bead_size=bead_size, - periodic_bond_opacity=periodic_bond_opacity, - ) - else: - raise RuntimeError( - f"Unsupported visualization backend ({backend}). " - "Currently supported backends include nglview and py3dmol" - ) - - else: - raise RuntimeError("Visualization is only supported in Jupyter Notebooks.") - - def _visualize_py3dmol( - self, - show_ports=False, - color_scheme={}, - bead_size=0.3, - periodic_bond_opacity=False, - ): - """Visualize the Compound using py3Dmol. - - Allows for visualization of a Compound within a Jupyter Notebook. - - Parameters - ---------- - show_ports : bool, optional, default=False - Visualize Ports in addition to Particles - color_scheme : dict, optional - Specify coloring for non-elemental particles - keys are strings of the particle names - values are strings of the colors - i.e. {'_CGBEAD': 'blue'} - bead_size : float, Optional, default=0.3 - Size of beads in visualization - periodic_bond_opacity : bool, float, Optional, default=False - Specify as a float from 0 to 1 to set the bond opacity - for bonds that cross periodic boundaries. - - Returns - ------- - view : py3Dmol.view - """ - py3Dmol = import_("py3Dmol") - - cloned = clone(self) - for edge in cloned.bond_graph.edges(data=True): - if edge[2]["bond_order"] == 0.0: - edge[2]["bond_order"] = 1.0 - - modified_color_scheme = {} - for name, color in color_scheme.items(): - # Py3dmol does some element string conversions, - # first character is as-is, rest of the characters are lowercase - new_name = name[0] + name[1:].lower() - modified_color_scheme[new_name] = color - modified_color_scheme[name] = color - - for particle in cloned.particles(): - if not particle.name: - particle.name = "UNK" - tmp_dir = tempfile.mkdtemp() - # bin bonds into periodic and aperiodic bonds - if isinstance(periodic_bond_opacity, float): - # save into two mol2 files, one with periodic bonds and one without - periodic_bonds, aperiodic_bonds = cloned._classify_periodic_bonds() - periodicGraph = nx.subgraph_view( - cloned.bond_graph, - filter_edge=lambda n1, n2: ( - (n1, n2) in periodic_bonds or (n2, n1) in periodic_bonds - ), - ) - aperiodicGraph = nx.subgraph_view( - cloned.bond_graph, - filter_edge=lambda n1, n2: ( - (n1, n2) in aperiodic_bonds or (n2, n1) in aperiodic_bonds - ), - ) - cpd1 = Compound.from_bondgraph(periodicGraph) - cpd2 = Compound.from_bondgraph(aperiodicGraph) - cpd1.save( - os.path.join(tmp_dir, "periodic.mol2"), - include_ports=show_ports, - ) - cpd2.save( - os.path.join(tmp_dir, "aperiodic.mol2"), - include_ports=show_ports, - ) - view = py3Dmol.view() - with open(os.path.join(tmp_dir, "periodic.mol2"), "r") as f: - view.addModel(f.read(), "mol2", keepH=True) - with open(os.path.join(tmp_dir, "aperiodic.mol2"), "r") as f: - view.addModel(f.read(), "mol2", keepH=True) - - view.setStyle( - {"model": 0}, - { - "stick": { - "radius": bead_size * 0.3, - "color": "grey", - "opacity": periodic_bond_opacity, - }, - "sphere": { - "scale": bead_size, - "colorscheme": modified_color_scheme, - }, - }, - ) - view.setStyle( - {"model": 1}, - { - "stick": {"radius": bead_size * 0.6, "color": "grey"}, - "sphere": { - "scale": bead_size, - "colorscheme": modified_color_scheme, - }, - }, - ) - view.zoomTo() - - else: - cloned.save( - os.path.join(tmp_dir, "tmp.mol2"), - include_ports=show_ports, - overwrite=True, - ) - - view = py3Dmol.view() - with open(os.path.join(tmp_dir, "tmp.mol2"), "r") as f: - view.addModel(f.read(), "mol2", keepH=True) - - view.setStyle( - { - "stick": {"radius": bead_size * 0.6, "color": "grey"}, - "sphere": { - "scale": bead_size, - "colorscheme": modified_color_scheme, - }, - } - ) - view.zoomTo() - - return view - def _visualize_nglview(self, show_ports=False): - """Visualize the Compound using nglview. - - Allows for visualization of a Compound within a Jupyter Notebook. - - Parameters - ---------- - show_ports : bool, optional, default=False - Visualize Ports in addition to Particles + Notes + ----- + See mbuild.utils.visualize.visualize_compound for more details. """ - nglview = import_("nglview") - mdtraj = import_("mdtraj") # noqa: F841 - from mdtraj.geometry.sasa import _ATOMIC_RADII - - def remove_digits(x): - return "".join(i for i in x if not i.isdigit() or i == "_") + from mbuild.utils.visualize import visualize_compound - for particle in self.particles(): - particle.name = remove_digits(particle.name).upper() - if not particle.name: - particle.name = "UNK" - tmp_dir = tempfile.mkdtemp() - self.save( - os.path.join(tmp_dir, "tmp.mol2"), - include_ports=show_ports, - overwrite=True, + return visualize_compound( + compound=self, + show_ports=show_ports, + backend=backend, + color_scheme=color_scheme, + bead_size=bead_size, + periodic_bond_opacity=periodic_bond_opacity, ) - widget = nglview.show_file(os.path.join(tmp_dir, "tmp.mol2")) - widget.clear() - widget.add_ball_and_stick(cylinderOnly=True) - elements = set([particle.name for particle in self.particles()]) - scale = 50.0 - for element in elements: - try: - widget.add_ball_and_stick( - f"_{element.upper()}", - aspect_ratio=_ATOMIC_RADII[element.title()] ** 1.5 * scale, - ) - except KeyError: - ids = [ - str(i) - for i, particle in enumerate(self.particles()) - if particle.name == element - ] - widget.add_ball_and_stick( - f"@{','.join(ids)}", - aspect_ratio=0.17**1.5 * scale, - color="grey", - ) - if show_ports: - widget.add_ball_and_stick("_VS", aspect_ratio=1.0, color="#991f00") - overwrite_nglview_default(widget) - return widget def condense(self, inplace=True): """Condense the hierarchical structure of the Compound to the level of molecules. diff --git a/mbuild/tests/test_compound.py b/mbuild/tests/test_compound.py index 2155cbf10..f8287e1a5 100644 --- a/mbuild/tests/test_compound.py +++ b/mbuild/tests/test_compound.py @@ -2280,12 +2280,16 @@ def test_get_boundingbox_error(self, bad_value): @pytest.mark.skipif(not has_py3Dmol, reason="Py3Dmol is not installed") def test_visualize_py3dmol(self, ethane): py3Dmol = import_("py3Dmol") - vis_object = ethane._visualize_py3dmol() + from mbuild.utils.visualize import _visualize_py3dmol + + vis_object = _visualize_py3dmol(ethane) assert isinstance(vis_object, py3Dmol.view) @pytest.mark.skipif(not has_py3Dmol, reason="Py3Dmol is not installed") def test_visualize_periodic_bonds_py3dmol(self): py3Dmol = import_("py3Dmol") + from mbuild.utils.visualize import _visualize_py3dmol + # create a periodic structure to test cpd = mb.load("CCCCCCCCCCCC", smiles=True) # position at left x wall @@ -2296,13 +2300,15 @@ def test_visualize_periodic_bonds_py3dmol(self): for particle in cpd.particles(): if particle.xyz[0][0] > cpd.box.Lz: particle.translate([-1 * cpd.box.Lx, 0, 0]) - vis_object = cpd._visualize_py3dmol(periodic_bond_opacity=0.2) + vis_object = _visualize_py3dmol(cpd, periodic_bond_opacity=0.2) assert isinstance(vis_object, py3Dmol.view) @pytest.mark.skipif(not has_nglview, reason="NGLView is not installed") def test_visualize_nglview(self, ethane): nglview = import_("nglview") - vis_object = ethane._visualize_nglview() + from mbuild.utils.visualize import _visualize_nglview + + vis_object = _visualize_nglview(ethane) assert isinstance(vis_object.component_0, nglview.component.ComponentViewer) def test_element(self): diff --git a/mbuild/tests/test_simulation.py b/mbuild/tests/test_simulation.py index f4a664784..0c35b5734 100644 --- a/mbuild/tests/test_simulation.py +++ b/mbuild/tests/test_simulation.py @@ -439,8 +439,12 @@ def test_path_per_type_bond_length(self): path = _make_two_type_path(n=6) pff = PathForcefield(radius=0.4, bond_length={"A-B": 0.25}) sim = HoomdSimulation(path, forcefield=pff, r_cut=0.5, run_on_gpu=False) - fene = sim.forces[1] - assert fene.params["A-B"]["sigma"] == pytest.approx(0.25) + harmonic = hoomd.md.bond.Harmonic() + for force in sim.forces: + if isinstance(force, hoomd.md.bond.Harmonic): + harmonic = force + break + assert harmonic.params["A-B"]["r0"] == pytest.approx(0.25) @pytest.mark.skipif(not has_hoomd, reason="hoomd is not installed") def test_path_auto_per_type(self): diff --git a/mbuild/utils/visualize.py b/mbuild/utils/visualize.py index 18db3bea6..9e08b8a76 100644 --- a/mbuild/utils/visualize.py +++ b/mbuild/utils/visualize.py @@ -1,11 +1,15 @@ """Methods for visualizing mBuild Compound and Path instances.""" +import os +import tempfile from copy import deepcopy +import networkx as nx import numpy as np from mbuild.path.formats import to_mol3000 -from mbuild.utils.io import import_ +from mbuild.utils.io import import_, run_from_ipython +from mbuild.utils.jsutils import overwrite_nglview_default def visualize_path(path, radius=0.1, hide_periodic_bonds=False): @@ -13,8 +17,10 @@ def visualize_path(path, radius=0.1, hide_periodic_bonds=False): Parameters ---------- - radius : float, default 0.06 + radius : float, default 0.1 Radius for sphere and stick representation + hide_periodic_bonds : bool, default False + If ``True`` bonds crossing periodic boundaries aren't shown. """ py3Dmol = import_("py3Dmol") @@ -64,9 +70,7 @@ def visualize_path(path, radius=0.1, hide_periodic_bonds=False): data = to_mol3000(path=path, G=G) view = py3Dmol.view(data=data) - # Select atoms by index rather than by element name. 3Dmol.js normalizes - # the SDF atom symbol to a 1-2 char element (e.g. "PEO" -> "Pe"), so an - # {"elem": name} selector only matches single/double-char bead names. + for i, name in enumerate(unique_names): color = colors[i % len(colors)] indices = [int(idx) for idx, bead in enumerate(path.beads) if bead == name] @@ -79,7 +83,254 @@ def visualize_path(path, radius=0.1, hide_periodic_bonds=False): ) view.setBackgroundColor("white") view.zoomTo() - # scale_factor = max(1, 5 - int(np.log10(len(path.coordinates)))) - # view.zoom(scale_factor) # helps zoom on smaller systems return view + + +def visualize_compound( + compound, + show_ports=False, + backend="py3dmol", + color_scheme={}, + bead_size=0.3, + periodic_bond_opacity=False, +): # pragma: no cover + """Visualize the Compound using py3dmol (default) or nglview. + + Allows for visualization of a Compound within a Jupyter Notebook. + + Parameters + ---------- + compound : mb.Compound + The compound to show. + show_ports : bool, optional, default=False + Visualize Ports in addition to Particles + backend : str, optional, default='py3dmol' + Specify the backend package to visualize compounds + Currently supported: py3dmol, nglview + color_scheme : dict, optional + Specify coloring for non-elemental particles + keys are strings of the particle names + values are strings of the colors + i.e. {'_CGBEAD': 'blue'} + bead_size : float, Optional, default=0.3 + Size of beads in visualization + periodic_bond_opacity : bool, float, Optional, default=False + Specify as a float from 0 to 1 to set the bond opacity + for bonds that cross periodic boundaries. + """ + viz_pkg = { + "nglview": _visualize_nglview, + "py3dmol": _visualize_py3dmol, + } + if run_from_ipython(): + if backend.lower() in viz_pkg: + if backend.lower() == "nglview": + return viz_pkg[backend.lower()](compound, show_ports=show_ports) + else: + return viz_pkg[backend.lower()]( + compound, + show_ports=show_ports, + color_scheme=color_scheme, + bead_size=bead_size, + periodic_bond_opacity=periodic_bond_opacity, + ) + else: + raise RuntimeError( + f"Unsupported visualization backend ({backend}). " + "Currently supported backends include nglview and py3dmol" + ) + + else: + raise RuntimeError("Visualization is only supported in Jupyter Notebooks.") + + +def _visualize_py3dmol( + compound, + show_ports=False, + color_scheme={}, + bead_size=0.3, + periodic_bond_opacity=False, +): + """Visualize the Compound using py3Dmol. + + Allows for visualization of a Compound within a Jupyter Notebook. + + Parameters + ---------- + show_ports : bool, optional, default=False + Visualize Ports in addition to Particles + color_scheme : dict, optional + Specify coloring for non-elemental particles + keys are strings of the particle names + values are strings of the colors + i.e. {'_CGBEAD': 'blue'} + bead_size : float, Optional, default=0.3 + Size of beads in visualization + periodic_bond_opacity : bool, float, Optional, default=False + Specify as a float from 0 to 1 to set the bond opacity + for bonds that cross periodic boundaries. + + Returns + ------- + view : py3Dmol.view + """ + py3Dmol = import_("py3Dmol") + + clone_of = {} + cloned = compound._clone(clone_of) + compound._clone_bonds(clone_of) + for edge in cloned.bond_graph.edges(data=True): + if edge[2]["bond_order"] == 0.0: + edge[2]["bond_order"] = 1.0 + + modified_color_scheme = {} + for name, color in color_scheme.items(): + # Py3dmol does some element string conversions, + # first character is as-is, rest of the characters are lowercase + new_name = name[0] + name[1:].lower() + modified_color_scheme[new_name] = color + modified_color_scheme[name] = color + + for particle in cloned.particles(): + if not particle.name: + particle.name = "UNK" + tmp_dir = tempfile.mkdtemp() + # bin bonds into periodic and aperiodic bonds + if isinstance(periodic_bond_opacity, float): + from mbuild import Compound + + # save into two mol2 files, one with periodic bonds and one without + periodic_bonds, aperiodic_bonds = cloned._classify_periodic_bonds() + periodicGraph = nx.subgraph_view( + cloned.bond_graph, + filter_edge=lambda n1, n2: ( + (n1, n2) in periodic_bonds or (n2, n1) in periodic_bonds + ), + ) + aperiodicGraph = nx.subgraph_view( + cloned.bond_graph, + filter_edge=lambda n1, n2: ( + (n1, n2) in aperiodic_bonds or (n2, n1) in aperiodic_bonds + ), + ) + cpd1 = Compound.from_bondgraph(periodicGraph) + cpd2 = Compound.from_bondgraph(aperiodicGraph) + cpd1.save( + os.path.join(tmp_dir, "periodic.mol2"), + include_ports=show_ports, + ) + cpd2.save( + os.path.join(tmp_dir, "aperiodic.mol2"), + include_ports=show_ports, + ) + view = py3Dmol.view() + with open(os.path.join(tmp_dir, "periodic.mol2"), "r") as f: + view.addModel(f.read(), "mol2", keepH=True) + with open(os.path.join(tmp_dir, "aperiodic.mol2"), "r") as f: + view.addModel(f.read(), "mol2", keepH=True) + + view.setStyle( + {"model": 0}, + { + "stick": { + "radius": bead_size * 0.3, + "color": "grey", + "opacity": periodic_bond_opacity, + }, + "sphere": { + "scale": bead_size, + "colorscheme": modified_color_scheme, + }, + }, + ) + view.setStyle( + {"model": 1}, + { + "stick": {"radius": bead_size * 0.6, "color": "grey"}, + "sphere": { + "scale": bead_size, + "colorscheme": modified_color_scheme, + }, + }, + ) + view.zoomTo() + + else: + cloned.save( + os.path.join(tmp_dir, "tmp.mol2"), + include_ports=show_ports, + overwrite=True, + ) + + view = py3Dmol.view() + with open(os.path.join(tmp_dir, "tmp.mol2"), "r") as f: + view.addModel(f.read(), "mol2", keepH=True) + + view.setStyle( + { + "stick": {"radius": bead_size * 0.6, "color": "grey"}, + "sphere": { + "scale": bead_size, + "colorscheme": modified_color_scheme, + }, + } + ) + view.zoomTo() + + return view + + +def _visualize_nglview(compound, show_ports=False): + """Visualize the Compound using nglview. + + Allows for visualization of a Compound within a Jupyter Notebook. + + Parameters + ---------- + show_ports : bool, optional, default=False + Visualize Ports in addition to Particles + """ + nglview = import_("nglview") + import_("mdtraj") # noqa: F841 + from mdtraj.geometry.sasa import _ATOMIC_RADII + + def remove_digits(x): + return "".join(i for i in x if not i.isdigit() or i == "_") + + for particle in compound.particles(): + particle.name = remove_digits(particle.name).upper() + if not particle.name: + particle.name = "UNK" + tmp_dir = tempfile.mkdtemp() + compound.save( + os.path.join(tmp_dir, "tmp.mol2"), + include_ports=show_ports, + overwrite=True, + ) + widget = nglview.show_file(os.path.join(tmp_dir, "tmp.mol2")) + widget.clear() + widget.add_ball_and_stick(cylinderOnly=True) + elements = set([particle.name for particle in compound.particles()]) + scale = 50.0 + for element in elements: + try: + widget.add_ball_and_stick( + f"_{element.upper()}", + aspect_ratio=_ATOMIC_RADII[element.title()] ** 1.5 * scale, + ) + except KeyError: + ids = [ + str(i) + for i, particle in enumerate(compound.particles()) + if particle.name == element + ] + widget.add_ball_and_stick( + f"@{','.join(ids)}", + aspect_ratio=0.17**1.5 * scale, + color="grey", + ) + if show_ports: + widget.add_ball_and_stick("_VS", aspect_ratio=1.0, color="#991f00") + overwrite_nglview_default(widget) + return widget