Skip to content

Commit 9dd3b63

Browse files
authored

File tree

1 file changed

+44
-0
lines changed

1 file changed

+44
-0
lines changed

dpdata/plugins/3dmol.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
from typing import Tuple
2+
import numpy as np
3+
4+
from dpdata.format import Format
5+
from dpdata.xyz.xyz import coord_to_xyz
6+
7+
8+
@Format.register("3dmol")
9+
class Py3DMolFormat(Format):
10+
"""3DMol format.
11+
12+
To use this format, py3Dmol should be installed in advance.
13+
"""
14+
def to_system(self,
15+
data: dict,
16+
f_idx: int = 0,
17+
size: Tuple[int] = (300,300),
18+
style: dict = {"stick":{}, "sphere":{"radius":0.4}},
19+
**kwargs):
20+
"""Show 3D structure of a frame in jupyter.
21+
22+
Parameters
23+
----------
24+
data : dict
25+
system data
26+
f_idx : int
27+
frame index to show
28+
size : tuple[int]
29+
(width, height) of the widget
30+
style : dict
31+
style of 3DMol. Read 3DMol documentation for details.
32+
33+
Examples
34+
--------
35+
>>> system.to_3dmol()
36+
"""
37+
import py3Dmol
38+
types = np.array(data['atom_names'])[data['atom_types']]
39+
xyz = coord_to_xyz(data['coords'][f_idx], types)
40+
viewer = py3Dmol.view(width=size[0], height=size[1])
41+
viewer.addModel(xyz, 'xyz')
42+
viewer.setStyle(style.copy())
43+
viewer.zoomTo()
44+
return viewer

0 commit comments

Comments
 (0)