-
Notifications
You must be signed in to change notification settings - Fork 49
/
ipython.py
87 lines (67 loc) · 2.33 KB
/
ipython.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
"""Ipython integration
"""
from __future__ import absolute_import
from io import BytesIO
from chemlab.graphics.renderers import AtomRenderer, BoxRenderer, BallAndStickRenderer
from chemlab.graphics.qtviewer import QtViewer
from chemlab.core import Molecule, System
from OpenGL.GL import *
from IPython.display import Image as ipy_Image
from PIL import Image as pil_Image
def load_ipython_extension(ipython):
# The `ipython` argument is the currently active `InteractiveShell`
# instance, which can be used in any way. This allows you to register
# new magics, plugins or aliases, for example.
png_formatter = ipython.display_formatter.formatters['image/png']
png_formatter.for_type(Molecule, mol_to_png)
png_formatter.for_type(System, sys_to_png)
def unload_ipython_extension(ipython):
# If you want your extension to be unloadable, put that logic here.
return
def showmol(mol, style='ball-and-stick',
width=300, height=300):
v = QtViewer()
w = v.widget
w.initializeGL()
if style == 'ball-and-stick':
bs = v.add_renderer(BallAndStickRenderer,
mol.r_array,
mol.type_array,
mol.bonds)
elif style == 'vdw':
sr = v.add_renderer(AtomRenderer, mol.r_array, mol.type_array,
backend='impostors')
w.camera.autozoom(mol.r_array)
image = w.toimage(width, height)
b = BytesIO()
image.save(b, format='png')
data = b.getvalue()
# Cleanup
del v
del w
# Save as png
return ipy_Image(data=data)
def mol_to_png(mol):
return showmol(mol)._repr_png_()
def showsys(sys, width=400, height=400):
v = QtViewer()
w = v.widget
w.initializeGL()
sr = v.add_renderer(AtomRenderer, sys.r_array, sys.type_array,
backend='impostors')
if sys.box_vectors is not None:
v.add_renderer(BoxRenderer, sys.box_vectors)
w.camera.autozoom(sys.r_array)
w.camera.orbit_y(-3.14/4)
w.camera.orbit_x(-3.14/4)
image = w.toimage(width, height)
b = BytesIO()
image.save(b, format='png')
data = b.getvalue()
# Cleanup
del v
del w
# Save as png
return ipy_Image(data=data)
def sys_to_png(syst):
return showsys(syst)._repr_png_()