Skip to content

Commit

Permalink
Merge branch 'main' into pre-commit-ci-update-config
Browse files Browse the repository at this point in the history
  • Loading branch information
daico007 authored Feb 1, 2024
2 parents aada462 + ea6e756 commit fe03c8f
Showing 1 changed file with 23 additions and 6 deletions.
29 changes: 23 additions & 6 deletions mbuild/compound.py
Original file line number Diff line number Diff line change
Expand Up @@ -1865,7 +1865,11 @@ def particles_in_range(
return particle_array[idxs]

def visualize(
self, show_ports=False, backend="py3dmol", color_scheme={}
self,
show_ports=False,
backend="py3dmol",
color_scheme={},
bead_size=0.3,
): # pragma: no cover
"""Visualize the Compound using py3dmol (default) or nglview.
Expand All @@ -1883,6 +1887,8 @@ def visualize(
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
"""
viz_pkg = {
"nglview": self._visualize_nglview,
Expand All @@ -1891,7 +1897,9 @@ def visualize(
if run_from_ipython():
if backend.lower() in viz_pkg:
return viz_pkg[backend.lower()](
show_ports=show_ports, color_scheme=color_scheme
show_ports=show_ports,
color_scheme=color_scheme,
bead_size=bead_size,
)
else:
raise RuntimeError(
Expand All @@ -1904,7 +1912,9 @@ def visualize(
"Visualization is only supported in Jupyter Notebooks."
)

def _visualize_py3dmol(self, show_ports=False, color_scheme={}):
def _visualize_py3dmol(
self, show_ports=False, color_scheme={}, bead_size=0.3
):
"""Visualize the Compound using py3Dmol.
Allows for visualization of a Compound within a Jupyter Notebook.
Expand All @@ -1918,6 +1928,8 @@ def _visualize_py3dmol(self, show_ports=False, color_scheme={}):
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
Returns
-------
Expand Down Expand Up @@ -1952,15 +1964,20 @@ def _visualize_py3dmol(self, show_ports=False, color_scheme={}):

view.setStyle(
{
"stick": {"radius": 0.2, "color": "grey"},
"sphere": {"scale": 0.3, "colorscheme": modified_color_scheme},
"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, color_scheme={}):
def _visualize_nglview(
self, show_ports=False, color_scheme={}, bead_size=0.3
):
"""Visualize the Compound using nglview.
Allows for visualization of a Compound within a Jupyter Notebook.
Expand Down

0 comments on commit fe03c8f

Please sign in to comment.