source: sasmodels/explore/jitter.py @ 41917b8

Last change on this file since 41917b8 was 85190c2, checked in by Paul Kienzle <pkienzle@…>, 6 years ago

explore/jitter: remove spaces from ends of lines

  • Property mode set to 100644
File size: 7.4 KB
Line 
1"""
2Application to explore the difference between sasview 3.x orientation
3dispersity and possible replacement algorithms.
4"""
5import mpl_toolkits.mplot3d   # Adds projection='3d' option to subplot
6import matplotlib.pyplot as plt
7from matplotlib.widgets import Slider, CheckButtons
8from matplotlib import cm
9
10import numpy as np
11from numpy import pi, cos, sin, sqrt, exp, degrees, radians
12
13def draw_beam(ax):
14    #ax.plot([0,0],[0,0],[1,-1])
15    #ax.scatter([0]*100,[0]*100,np.linspace(1, -1, 100), alpha=0.8)
16
17    steps = 25
18    u = np.linspace(0, 2 * np.pi, steps)
19    v = np.linspace(-1, 1, steps)
20
21    r = 0.02
22    x = r*np.outer(np.cos(u), np.ones_like(v))
23    y = r*np.outer(np.sin(u), np.ones_like(v))
24    z = np.outer(np.ones_like(u), v)
25
26    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='y', alpha=0.5)
27
28def draw_shimmy(ax, theta, phi, psi, dtheta, dphi, dpsi):
29    size=[0.1, 0.4, 1.0]
30    view=[theta, phi, psi]
31    shimmy=[0,0,0]
32    #draw_shape = draw_parallelepiped
33    draw_shape = draw_ellipsoid
34
35    #np.random.seed(10)
36    #cloud = np.random.randn(10,3)
37    cloud = [
38        [-1, -1, -1],
39        [-1, -1,  0],
40        [-1, -1,  1],
41        [-1,  0, -1],
42        [-1,  0,  0],
43        [-1,  0,  1],
44        [-1,  1, -1],
45        [-1,  1,  0],
46        [-1,  1,  1],
47        [ 0, -1, -1],
48        [ 0, -1,  0],
49        [ 0, -1,  1],
50        [ 0,  0, -1],
51        [ 0,  0,  0],
52        [ 0,  0,  1],
53        [ 0,  1, -1],
54        [ 0,  1,  0],
55        [ 0,  1,  1],
56        [ 1, -1, -1],
57        [ 1, -1,  0],
58        [ 1, -1,  1],
59        [ 1,  0, -1],
60        [ 1,  0,  0],
61        [ 1,  0,  1],
62        [ 1,  1, -1],
63        [ 1,  1,  0],
64        [ 1,  1,  1],
65    ]
66    if dtheta == 0:
67        cloud = [v for v in cloud if v[0] == 0]
68    if dphi == 0:
69        cloud = [v for v in cloud if v[1] == 0]
70    if dpsi == 0:
71        cloud = [v for v in cloud if v[2] == 0]
72    draw_shape(ax, size, view, shimmy, steps=100, alpha=0.8)
73    for point in cloud:
74        shimmy=[dtheta*point[0], dphi*point[1], dpsi*point[2]]
75        draw_shape(ax, size, view, shimmy, alpha=0.8)
76    for v in 'xyz':
77        a, b, c = size
78        lim = np.sqrt(a**2+b**2+c**2)
79        getattr(ax, 'set_'+v+'lim')([-lim, lim])
80        getattr(ax, v+'axis').label.set_text(v)
81
82def draw_ellipsoid(ax, size, view, shimmy, steps=25, alpha=1):
83    a,b,c = size
84    theta, phi, psi = view
85    dtheta, dphi, dpsi = shimmy
86
87    u = np.linspace(0, 2 * np.pi, steps)
88    v = np.linspace(0, np.pi, steps)
89    x = a*np.outer(np.cos(u), np.sin(v))
90    y = b*np.outer(np.sin(u), np.sin(v))
91    z = c*np.outer(np.ones_like(u), np.cos(v))
92
93    shape = x.shape
94    points = np.matrix([x.flatten(),y.flatten(),z.flatten()])
95    points = Rz(dpsi)*Ry(dtheta)*Rx(dphi)*points
96    points = Rz(phi)*Ry(theta)*Rz(psi)*points
97    x,y,z = [v.reshape(shape) for v in points]
98
99    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w', alpha=alpha)
100
101def draw_parallelepiped(ax, size, view, shimmy, alpha=1):
102    a,b,c = size
103    theta, phi, psi = view
104    dtheta, dphi, dpsi = shimmy
105
106    x = a*np.array([ 1,-1, 1,-1, 1,-1, 1,-1])
107    y = b*np.array([ 1, 1,-1,-1, 1, 1,-1,-1])
108    z = c*np.array([ 1, 1, 1, 1,-1,-1,-1,-1])
109    tri = np.array([
110        # counter clockwise triangles
111        # z: up/down, x: right/left, y: front/back
112        [0,1,2], [3,2,1], # top face
113        [6,5,4], [5,6,7], # bottom face
114        [0,2,6], [6,4,0], # right face
115        [1,5,7], [7,3,1], # left face
116        [2,3,6], [7,6,3], # front face
117        [4,1,0], [5,1,4], # back face
118    ])
119
120    points = np.matrix([x,y,z])
121    points = Rz(dpsi)*Ry(dtheta)*Rx(dphi)*points
122    points = Rz(phi)*Ry(theta)*Rz(psi)*points
123
124    x,y,z = [np.array(v).flatten() for v in points]
125    ax.plot_trisurf(x, y, triangles=tri, Z=z, color='w', alpha=alpha)
126
127def draw_sphere(ax, radius=10., steps=100):
128    u = np.linspace(0, 2 * np.pi, steps)
129    v = np.linspace(0, np.pi, steps)
130
131    x = radius * np.outer(np.cos(u), np.sin(v))
132    y = radius * np.outer(np.sin(u), np.sin(v))
133    z = radius * np.outer(np.ones(np.size(u)), np.cos(v))
134    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w')
135
136def draw_mesh_new(ax, theta, dtheta, phi, dphi, flow, radius=10., dist='gauss'):
137    theta_center = radians(theta)
138    phi_center = radians(phi)
139    flow_center = radians(flow)
140    dtheta = radians(dtheta)
141    dphi = radians(dphi)
142
143    # 10 point 3-sigma gaussian weights
144    t = np.linspace(-3., 3., 11)
145    if dist == 'gauss':
146        weights = exp(-0.5*t**2)
147    elif dist == 'rect':
148        weights = np.ones_like(t)
149    else:
150        raise ValueError("expected dist to be 'gauss' or 'rect'")
151    theta = dtheta*t
152    phi = dphi*t
153
154    x = radius * np.outer(cos(phi), cos(theta))
155    y = radius * np.outer(sin(phi), cos(theta))
156    z = radius * np.outer(np.ones_like(phi), sin(theta))
157    #w = np.outer(weights, weights*abs(cos(dtheta*t)))
158    w = np.outer(weights, weights*abs(cos(theta)))
159
160    x, y, z, w = [v.flatten() for v in (x,y,z,w)]
161    x, y, z = rotate(x, y, z, phi_center, theta_center, flow_center)
162
163    ax.scatter(x, y, z, c=w, marker='o', vmin=0., vmax=1.)
164
165def rotate(x, y, z, phi, theta, psi):
166    R = Rz(phi)*Ry(theta)*Rz(psi)
167    p = np.vstack([x,y,z])
168    return R*p
169
170def Rx(angle):
171    a = radians(angle)
172    R = [[1., 0., 0.],
173         [0.,  cos(a), sin(a)],
174         [0., -sin(a), cos(a)]]
175    return np.matrix(R)
176
177def Ry(angle):
178    a = radians(angle)
179    R = [[cos(a), 0., -sin(a)],
180         [0., 1., 0.],
181         [sin(a), 0.,  cos(a)]]
182    return np.matrix(R)
183
184def Rz(angle):
185    a = radians(angle)
186    R = [[cos(a), -sin(a), 0.],
187         [sin(a),  cos(a), 0.],
188         [0., 0., 1.]]
189    return np.matrix(R)
190
191def main():
192    #plt.hold(True)
193    plt.set_cmap('gist_earth')
194    plt.clf()
195    #gs = gridspec.GridSpec(2,1,height_ratios=[4,1])
196    #ax = plt.subplot(gs[0], projection='3d')
197    ax = plt.axes([0.0, 0.2, 1.0, 0.8], projection='3d')
198
199    theta, dtheta = 70., 10.
200    phi, dphi = -45., 3.
201    psi, dpsi = -45., 3.
202    theta, phi, psi = 0, 0, 0
203    dtheta, dphi, dpsi = 0, 0, 0
204    #dist = 'rect'
205    dist = 'gauss'
206
207    axcolor = 'lightgoldenrodyellow'
208    axtheta  = plt.axes([0.1, 0.15, 0.45, 0.04], axisbg=axcolor)
209    axphi = plt.axes([0.1, 0.1, 0.45, 0.04], axisbg=axcolor)
210    axpsi = plt.axes([0.1, 0.05, 0.45, 0.04], axisbg=axcolor)
211    stheta = Slider(axtheta, 'Theta', -90, 90, valinit=theta)
212    sphi = Slider(axphi, 'Phi', -180, 180, valinit=phi)
213    spsi = Slider(axpsi, 'Psi', -180, 180, valinit=psi)
214    axdtheta  = plt.axes([0.75, 0.15, 0.15, 0.04], axisbg=axcolor)
215    axdphi = plt.axes([0.75, 0.1, 0.15, 0.04], axisbg=axcolor)
216    axdpsi= plt.axes([0.75, 0.05, 0.15, 0.04], axisbg=axcolor)
217    sdtheta = Slider(axdtheta, 'dTheta', 0, 30, valinit=dtheta)
218    sdphi = Slider(axdphi, 'dPhi', 0, 30, valinit=dphi)
219    sdpsi = Slider(axdpsi, 'dPsi', 0, 30, valinit=dphi)
220
221    def update(val, axis=None):
222        theta, phi, psi = stheta.val, sphi.val, spsi.val
223        dtheta, dphi, dpsi = sdtheta.val, sdphi.val, sdpsi.val
224        ax.cla()
225        draw_beam(ax)
226        draw_shimmy(ax, theta, phi, psi, dtheta, dphi, dpsi)
227        #if not axis.startswith('d'):
228        #    ax.view_init(elev=theta, azim=phi)
229        plt.gcf().canvas.draw()
230
231    stheta.on_changed(lambda v: update(v,'theta'))
232    sphi.on_changed(lambda v: update(v, 'phi'))
233    spsi.on_changed(lambda v: update(v, 'psi'))
234    sdtheta.on_changed(lambda v: update(v, 'dtheta'))
235    sdphi.on_changed(lambda v: update(v, 'dphi'))
236    sdpsi.on_changed(lambda v: update(v, 'dpsi'))
237
238    update(None, 'phi')
239
240    plt.show()
241
242if __name__ == "__main__":
243    main()
Note: See TracBrowser for help on using the repository browser.