From 5c07518e0e4fb12c47de70928f63690be6280046 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sun, 25 Sep 2022 03:05:03 -0400
Subject: [PATCH 1/5] add 3dmol format

---
 dpdata/plugins/3dmol.py | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)
 create mode 100644 dpdata/plugins/3dmol.py

diff --git a/dpdata/plugins/3dmol.py b/dpdata/plugins/3dmol.py
new file mode 100644
index 000000000..ccc3e9db4
--- /dev/null
+++ b/dpdata/plugins/3dmol.py
@@ -0,0 +1,18 @@
+from dpdata.format import Format
+from dpdata.xyz.xyz import coord_to_xyz
+
+
+@Format.register("3dmol")
+class AmberMDFormat(Format):
+    """See 3D structure of a frame in jupyter.
+    
+    py3Dmol should be installed in advance.
+    """
+    def to_system(self, data: dict, f_idx: int = 0, **kwargs):
+        import py3Dmol
+        types = np.array(data['atom_names'])[data['atom_types']]
+        xyz = coord_to_xyz(data['coords'][f_idx], types)
+        viewer = py3Dmol.view(width=300, height=300)
+        viewer.addModel(xyz, 'xyz')
+        viewer.zoomTo()
+        return viewer

From 0289421daa0541debfc45e3a2d57cc064545a54e Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sun, 25 Sep 2022 03:07:09 -0400
Subject: [PATCH 2/5] import numpy

---
 dpdata/plugins/3dmol.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/dpdata/plugins/3dmol.py b/dpdata/plugins/3dmol.py
index ccc3e9db4..8d7b8e8a2 100644
--- a/dpdata/plugins/3dmol.py
+++ b/dpdata/plugins/3dmol.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 from dpdata.format import Format
 from dpdata.xyz.xyz import coord_to_xyz
 

From b5b6e55953602313e573d65cadb1aff33e76ed26 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sun, 25 Sep 2022 03:08:29 -0400
Subject: [PATCH 3/5] set style

---
 dpdata/plugins/3dmol.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/dpdata/plugins/3dmol.py b/dpdata/plugins/3dmol.py
index 8d7b8e8a2..ab1bef8d0 100644
--- a/dpdata/plugins/3dmol.py
+++ b/dpdata/plugins/3dmol.py
@@ -16,5 +16,6 @@ def to_system(self, data: dict, f_idx: int = 0, **kwargs):
         xyz = coord_to_xyz(data['coords'][f_idx], types)
         viewer = py3Dmol.view(width=300, height=300)
         viewer.addModel(xyz, 'xyz')
+        viewer.setStyle({"stick": {}})
         viewer.zoomTo()
         return viewer

From e2e22a7338766e592ec8fd4e6c8ca8486c00003e Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sun, 25 Sep 2022 03:19:36 -0400
Subject: [PATCH 4/5] show atoms

---
 dpdata/plugins/3dmol.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dpdata/plugins/3dmol.py b/dpdata/plugins/3dmol.py
index ab1bef8d0..25760f443 100644
--- a/dpdata/plugins/3dmol.py
+++ b/dpdata/plugins/3dmol.py
@@ -16,6 +16,6 @@ def to_system(self, data: dict, f_idx: int = 0, **kwargs):
         xyz = coord_to_xyz(data['coords'][f_idx], types)
         viewer = py3Dmol.view(width=300, height=300)
         viewer.addModel(xyz, 'xyz')
-        viewer.setStyle({"stick": {}})
+        viewer.setStyle({"stick":{}, "sphere":{"radius":0.4}})
         viewer.zoomTo()
         return viewer

From 10b73feebdb6436d5eaa3ea225ffdebd93e89f42 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sun, 25 Sep 2022 03:27:15 -0400
Subject: [PATCH 5/5] add parameters

---
 dpdata/plugins/3dmol.py | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/dpdata/plugins/3dmol.py b/dpdata/plugins/3dmol.py
index 25760f443..c30893295 100644
--- a/dpdata/plugins/3dmol.py
+++ b/dpdata/plugins/3dmol.py
@@ -1,3 +1,4 @@
+from typing import Tuple
 import numpy as np
 
 from dpdata.format import Format
@@ -5,17 +6,39 @@
 
 
 @Format.register("3dmol")
-class AmberMDFormat(Format):
-    """See 3D structure of a frame in jupyter.
-    
-    py3Dmol should be installed in advance.
+class Py3DMolFormat(Format):
+    """3DMol format.
+
+    To use this format,  py3Dmol should be installed in advance.
     """
-    def to_system(self, data: dict, f_idx: int = 0, **kwargs):
+    def to_system(self,
+                  data: dict,
+                  f_idx: int = 0,
+                  size: Tuple[int] = (300,300),
+                  style: dict = {"stick":{}, "sphere":{"radius":0.4}},
+                  **kwargs):
+        """Show 3D structure of a frame in jupyter.
+
+        Parameters
+        ----------
+        data : dict
+            system data
+        f_idx : int
+            frame index to show
+        size : tuple[int]
+            (width, height) of the widget
+        style : dict
+            style of 3DMol. Read 3DMol documentation for details.
+
+        Examples
+        --------
+        >>> system.to_3dmol()
+        """
         import py3Dmol
         types = np.array(data['atom_names'])[data['atom_types']]
         xyz = coord_to_xyz(data['coords'][f_idx], types)
-        viewer = py3Dmol.view(width=300, height=300)
+        viewer = py3Dmol.view(width=size[0], height=size[1])
         viewer.addModel(xyz, 'xyz')
-        viewer.setStyle({"stick":{}, "sphere":{"radius":0.4}})
+        viewer.setStyle(style.copy())
         viewer.zoomTo()
         return viewer