File tree Expand file tree Collapse file tree 1 file changed +44
-0
lines changed Expand file tree Collapse file tree 1 file changed +44
-0
lines changed Original file line number Diff line number Diff line change
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
You can’t perform that action at this time.
0 commit comments