From ea6e756a6bdf72f43dac724554d23455b2d12edf Mon Sep 17 00:00:00 2001
From: Jennifer A Clark <jennifer.a.clark13@gmail.com>
Date: Wed, 31 Jan 2024 21:59:12 -0500
Subject: [PATCH] Added ability to adjust bead_size (#1158)

---
 mbuild/compound.py | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)

diff --git a/mbuild/compound.py b/mbuild/compound.py
index 30d2f27f7..5ca296d9d 100644
--- a/mbuild/compound.py
+++ b/mbuild/compound.py
@@ -1864,7 +1864,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.
 
@@ -1882,6 +1886,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,
@@ -1890,7 +1896,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(
@@ -1903,7 +1911,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.
@@ -1917,6 +1927,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
         -------
@@ -1951,15 +1963,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.