Aluode/PerceptionLabPortable

1"""2Demonstrates GLVolumeItem for displaying volumetric data.3"""4import sys5 6import numpy as np7 8import pyqtgraph as pg9from pyqtgraph.Qt import QtGui10import pyqtgraph.opengl as gl11from pyqtgraph import functions as fn12 13if 'darwin' in sys.platform:14    fmt = QtGui.QSurfaceFormat()15    fmt.setRenderableType(fmt.RenderableType.OpenGL)16    fmt.setProfile(fmt.OpenGLContextProfile.CoreProfile)17    fmt.setVersion(4, 1)18    QtGui.QSurfaceFormat.setDefaultFormat(fmt)19 20app = pg.mkQApp("GLVolumeItem Example")21w = gl.GLViewWidget()22w.show()23w.setWindowTitle('pyqtgraph example: GLVolumeItem')24w.setCameraPosition(distance=200)25 26g = gl.GLGridItem()27g.scale(10, 10, 1)28w.addItem(g)29 30## Hydrogen electron probability density31def psi(i, j, k, offset=(50,50,100)):32    x = i-offset[0]33    y = j-offset[1]34    z = k-offset[2]35    th = np.arctan2(z, np.hypot(x, y))36    r = np.sqrt(x**2 + y**2 + z **2)37    a0 = 238    return (39        (1.0 / 81.0)40        * 1.0 / (6.0 * np.pi) ** 0.541        * (1.0 / a0) ** (3 / 2)42        * (r / a0) ** 243        * np.exp(-r / (3 * a0))44        * (3 * np.cos(th) ** 2 - 1)45    )46 47 48data = np.fromfunction(psi, (100,100,200))49with np.errstate(divide = 'ignore'):50    positive = np.log(fn.clip_array(data, 0, data.max())**2)51    negative = np.log(fn.clip_array(-data, 0, -data.min())**2)52 53d2 = np.empty(data.shape + (4,), dtype=np.ubyte)54 55# Original Code56# d2[..., 0] = positive * (255./positive.max())57# d2[..., 1] = negative * (255./negative.max())58 59# Reformulated Code60# Both positive.max() and negative.max() are negative-valued.61# Thus the next 2 lines are _not_ bounded to [0, 255]62positive = positive * (255./positive.max())63negative = negative * (255./negative.max())64# When casting to ubyte, the original code relied on +Inf to be65# converted to 0. On arm64, it gets converted to 255.66# Thus the next 2 lines change +Inf explicitly to 0 instead.67positive[np.isinf(positive)] = 068negative[np.isinf(negative)] = 069# When casting to ubyte, the original code relied on the conversion70# to do modulo 256. The next 2 lines do it explicitly instead as71# documentation.72d2[..., 0] = positive.astype(int) % 25673d2[..., 1] = negative.astype(int) % 25674 75d2[..., 2] = d2[...,1]76d2[..., 3] = d2[..., 0]*0.3 + d2[..., 1]*0.377d2[..., 3] = (d2[..., 3].astype(float) / 255.) **2 * 25578 79d2[:, 0, 0] = [255,0,0,100]80d2[0, :, 0] = [0,255,0,100]81d2[0, 0, :] = [0,0,255,100]82 83v = gl.GLVolumeItem(d2)84v.translate(-50,-50,-100)85w.addItem(v)86 87ax = gl.GLAxisItem()88w.addItem(ax)89 90if __name__ == '__main__':91    pg.exec()92