source: sasmodels/sasmodels/jitter.py @ 5f12750

core_shell_microgelsmagnetic_modelticket-1257-vesicle-productticket_1156ticket_1265_superballticket_822_more_unit_tests
Last change on this file since 5f12750 was 5f12750, checked in by Paul Kienzle <pkienzle@…>, 5 years ago

support ipyvolume for jitter in jupyter notebooks

  • Property mode set to 100755
File size: 46.9 KB
Line 
1#!/usr/bin/env python
2"""
3Jitter Explorer
4===============
5
6Application to explore orientation angle and angular dispersity.
7
8From the command line::
9
10    # Show docs
11    python -m sasmodels.jitter --help
12
13    # Guyou projection jitter, uniform over 20 degree theta and 10 in phi
14    python -m sasmodels.jitter --projection=guyou --dist=uniform --jitter=20,10,0
15
16From a jupyter cell::
17
18    import ipyvolume as ipv
19    from sasmodels import jitter
20    import importlib; importlib.reload(jitter)
21    jitter.set_plotter("ipv")
22
23    size = (10, 40, 100)
24    view = (20, 0, 0)
25
26    #size = (15, 15, 100)
27    #view = (60, 60, 0)
28
29    dview = (0, 0, 0)
30    #dview = (5, 5, 0)
31    #dview = (15, 180, 0)
32    #dview = (180, 15, 0)
33
34    projection = 'equirectangular'
35    #projection = 'azimuthal_equidistance'
36    #projection = 'guyou'
37    #projection = 'sinusoidal'
38    #projection = 'azimuthal_equal_area'
39
40    dist = 'uniform'
41    #dist = 'gaussian'
42
43    jitter.run(size=size, view=view, jitter=dview, dist=dist, projection=projection)
44    #filename = projection+('_theta' if dview[0] == 180 else '_phi' if dview[1] == 180 else '')
45    #ipv.savefig(filename+'.png')
46"""
47from __future__ import division, print_function
48
49import argparse
50
51import numpy as np
52from numpy import pi, cos, sin, sqrt, exp, degrees, radians
53
54def draw_beam(axes, view=(0, 0), alpha=0.5):
55    """
56    Draw the beam going from source at (0, 0, 1) to detector at (0, 0, -1)
57    """
58    #axes.plot([0,0],[0,0],[1,-1])
59    #axes.scatter([0]*100,[0]*100,np.linspace(1, -1, 100), alpha=alpha)
60
61    steps = [6, 6]
62    u = np.linspace(0, 2 * np.pi, steps[0])
63    v = np.linspace(-1, 1, steps[1])
64
65    r = 0.02
66    x = r*np.outer(np.cos(u), np.ones_like(v))
67    y = r*np.outer(np.sin(u), np.ones_like(v))
68    z = 1.3*np.outer(np.ones_like(u), v)
69
70    theta, phi = view
71    shape = x.shape
72    points = np.matrix([x.flatten(), y.flatten(), z.flatten()])
73    points = Rz(phi)*Ry(theta)*points
74    x, y, z = [v.reshape(shape) for v in points]
75
76    axes.plot_surface(x, y, z, color='yellow', alpha=alpha)
77
78def draw_ellipsoid(axes, size, view, jitter, steps=25, alpha=1):
79    """Draw an ellipsoid."""
80    a, b, c = size
81    u = np.linspace(0, 2 * np.pi, steps)
82    v = np.linspace(0, np.pi, steps)
83    x = a*np.outer(np.cos(u), np.sin(v))
84    y = b*np.outer(np.sin(u), np.sin(v))
85    z = c*np.outer(np.ones_like(u), np.cos(v))
86    x, y, z = transform_xyz(view, jitter, x, y, z)
87
88    axes.plot_surface(x, y, z, color='w', alpha=alpha)
89
90    draw_labels(axes, view, jitter, [
91        ('c+', [+0, +0, +c], [+1, +0, +0]),
92        ('c-', [+0, +0, -c], [+0, +0, -1]),
93        ('a+', [+a, +0, +0], [+0, +0, +1]),
94        ('a-', [-a, +0, +0], [+0, +0, -1]),
95        ('b+', [+0, +b, +0], [-1, +0, +0]),
96        ('b-', [+0, -b, +0], [-1, +0, +0]),
97    ])
98
99def draw_sc(axes, size, view, jitter, steps=None, alpha=1):
100    """Draw points for simple cubic paracrystal"""
101    atoms = _build_sc()
102    _draw_crystal(axes, size, view, jitter, atoms=atoms)
103
104def draw_fcc(axes, size, view, jitter, steps=None, alpha=1):
105    """Draw points for face-centered cubic paracrystal"""
106    # Build the simple cubic crystal
107    atoms = _build_sc()
108    # Define the centers for each face
109    # x planes at -1, 0, 1 have four centers per plane, at +/- 0.5 in y and z
110    x, y, z = (
111        [-1]*4 + [0]*4 + [1]*4,
112        ([-0.5]*2 + [0.5]*2)*3,
113        [-0.5, 0.5]*12,
114    )
115    # y and z planes can be generated by substituting x for y and z respectively
116    atoms.extend(zip(x+y+z, y+z+x, z+x+y))
117    _draw_crystal(axes, size, view, jitter, atoms=atoms)
118
119def draw_bcc(axes, size, view, jitter, steps=None, alpha=1):
120    """Draw points for body-centered cubic paracrystal"""
121    # Build the simple cubic crystal
122    atoms = _build_sc()
123    # Define the centers for each octant
124    # x plane at +/- 0.5 have four centers per plane at +/- 0.5 in y and z
125    x, y, z = (
126        [-0.5]*4 + [0.5]*4,
127        ([-0.5]*2 + [0.5]*2)*2,
128        [-0.5, 0.5]*8,
129    )
130    atoms.extend(zip(x, y, z))
131    _draw_crystal(axes, size, view, jitter, atoms=atoms)
132
133def _draw_crystal(axes, size, view, jitter, atoms=None):
134    atoms, size = np.asarray(atoms, 'd').T, np.asarray(size, 'd')
135    x, y, z = atoms*size[:, None]
136    x, y, z = transform_xyz(view, jitter, x, y, z)
137    axes.scatter([x[0]], [y[0]], [z[0]], c='yellow', marker='o')
138    axes.scatter(x[1:], y[1:], z[1:], c='r', marker='o')
139
140def _build_sc():
141    # three planes of 9 dots for x at -1, 0 and 1
142    x, y, z = (
143        [-1]*9 + [0]*9 + [1]*9,
144        ([-1]*3 + [0]*3 + [1]*3)*3,
145        [-1, 0, 1]*9,
146    )
147    atoms = list(zip(x, y, z))
148    #print(list(enumerate(atoms)))
149    # Pull the dot at (0, 0, 1) to the front of the list
150    # It will be highlighted in the view
151    index = 14
152    highlight = atoms[index]
153    del atoms[index]
154    atoms.insert(0, highlight)
155    return atoms
156
157def draw_parallelepiped(axes, size, view, jitter, steps=None, alpha=1):
158    """Draw a parallelepiped."""
159    a, b, c = size
160    x = a*np.array([+1, -1, +1, -1, +1, -1, +1, -1])
161    y = b*np.array([+1, +1, -1, -1, +1, +1, -1, -1])
162    z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1])
163    tri = np.array([
164        # counter clockwise triangles
165        # z: up/down, x: right/left, y: front/back
166        [0, 1, 2], [3, 2, 1], # top face
167        [6, 5, 4], [5, 6, 7], # bottom face
168        [0, 2, 6], [6, 4, 0], # right face
169        [1, 5, 7], [7, 3, 1], # left face
170        [2, 3, 6], [7, 6, 3], # front face
171        [4, 1, 0], [5, 1, 4], # back face
172    ])
173
174    x, y, z = transform_xyz(view, jitter, x, y, z)
175    color = [0.6, 1, 0.6]  # pale green
176    axes.plot_trisurf(x, y, triangles=tri, Z=z, color=color, alpha=alpha)
177
178    # Colour the c+ face of the box.
179    # Since I can't control face color, instead draw a thin box situated just
180    # in front of the "c+" face.  Use the c face so that rotations about psi
181    # rotate that face.
182    if 0:
183        #color = [1, 0.6, 0.6]  # pink
184        x = a*np.array([+1, -1, +1, -1, +1, -1, +1, -1])
185        y = b*np.array([+1, +1, -1, -1, +1, +1, -1, -1])
186        z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1])
187        x, y, z = transform_xyz(view, jitter, x, y, abs(z)+0.001)
188        axes.plot_trisurf(x, y, triangles=tri, Z=z, color=color, alpha=alpha)
189
190    draw_labels(axes, view, jitter, [
191        ('c+', [+0, +0, +c], [+1, +0, +0]),
192        ('c-', [+0, +0, -c], [+0, +0, -1]),
193        ('a+', [+a, +0, +0], [+0, +0, +1]),
194        ('a-', [-a, +0, +0], [+0, +0, -1]),
195        ('b+', [+0, +b, +0], [-1, +0, +0]),
196        ('b-', [+0, -b, +0], [-1, +0, +0]),
197    ])
198
199def draw_sphere(axes, radius=0.5, steps=25):
200    """Draw a sphere"""
201    u = np.linspace(0, 2 * np.pi, steps)
202    v = np.linspace(0, np.pi, steps)
203
204    x = radius * np.outer(np.cos(u), np.sin(v))
205    y = radius * np.outer(np.sin(u), np.sin(v))
206    z = radius * np.outer(np.ones(np.size(u)), np.cos(v))
207    axes.plot_surface(x, y, z, color='w')
208    #axes.plot_wireframe(x, y, z)
209
210def draw_person_on_sphere(axes, view, height=0.5, radius=0.5):
211    limb_offset = height * 0.05
212    head_radius = height * 0.10
213    head_height = height - head_radius
214    neck_length = head_radius * 0.50
215    shoulder_height = height - 2*head_radius - neck_length
216    torso_length = shoulder_height * 0.55
217    torso_radius = torso_length * 0.30
218    leg_length = shoulder_height - torso_length
219    arm_length = torso_length * 0.90
220
221    def _draw_part(x, z):
222        y = np.zeros_like(x)
223        xp, yp, zp = transform_xyz(view, None, x, y, z + radius)
224        axes.plot(xp, yp, zp, color='k')
225
226    # circle for head
227    u = np.linspace(0, 2 * np.pi, 40)
228    x = head_radius * np.cos(u)
229    z = head_radius * np.sin(u) + head_height
230    _draw_part(x, z)
231
232    # rectangle for body
233    x = np.array([-torso_radius, torso_radius, torso_radius, -torso_radius, -torso_radius])
234    z = np.array([0., 0, torso_length, torso_length, 0]) + leg_length
235    _draw_part(x, z)
236
237    # arms
238    x = np.array([-torso_radius - limb_offset, -torso_radius - limb_offset, -torso_radius])
239    z = np.array([shoulder_height - arm_length, shoulder_height, shoulder_height])
240    _draw_part(x, z)
241    _draw_part(-x, z)
242
243    # legs
244    x = np.array([-torso_radius + limb_offset, -torso_radius + limb_offset])
245    z = np.array([0, leg_length])
246    _draw_part(x, z)
247    _draw_part(-x, z)
248
249    limits = [-radius-height, radius+height]
250    axes.set_xlim(limits)
251    axes.set_ylim(limits)
252    axes.set_zlim(limits)
253    axes.set_axis_off()
254
255def draw_jitter(axes, view, jitter, dist='gaussian', 
256                size=(0.1, 0.4, 1.0),
257                draw_shape=draw_parallelepiped, 
258                projection='equirectangular',
259                alpha=0.8,
260                views=None):
261    """
262    Represent jitter as a set of shapes at different orientations.
263    """
264    project, weight = get_projection(projection)
265
266    # set max diagonal to 0.95
267    scale = 0.95/sqrt(sum(v**2 for v in size))
268    size = tuple(scale*v for v in size)
269
270    dtheta, dphi, dpsi = jitter
271    base = {'gaussian':3, 'rectangle':sqrt(3), 'uniform':1}[dist]
272    def steps(delta):
273        limit = base*delta
274        if views is None:
275            n = max(3, min(25, 2*int(base*delta/15)))
276        else:
277            n = views
278        s = base*delta*np.linspace(-1, 1, n) if delta > 0 else [0]
279        return s
280    stheta = steps(dtheta)
281    sphi = steps(dphi)
282    spsi = steps(dpsi)
283    #print(stheta, sphi, spsi)
284    for theta in stheta:
285        for phi in sphi:
286            for psi in spsi:
287                w = weight(theta, phi, 1.0, 1.0)
288                if w > 0:
289                    dview = project(theta, phi, psi)
290                    draw_shape(axes, size, view, dview, alpha=alpha)
291    for v in 'xyz':
292        a, b, c = size
293        lim = np.sqrt(a**2 + b**2 + c**2)
294        getattr(axes, 'set_'+v+'lim')([-lim, lim])
295        #getattr(axes, v+'axis').label.set_text(v)
296
297PROJECTIONS = [
298    # in order of PROJECTION number; do not change without updating the
299    # constants in kernel_iq.c
300    'equirectangular', 'sinusoidal', 'guyou', 'azimuthal_equidistance',
301    'azimuthal_equal_area',
302]
303def get_projection(projection):
304
305    """
306    jitter projections
307    <https://en.wikipedia.org/wiki/List_of_map_projections>
308
309    equirectangular (standard latitude-longitude mesh)
310        <https://en.wikipedia.org/wiki/Equirectangular_projection>
311        Allows free movement in phi (around the equator), but theta is
312        limited to +/- 90, and points are cos-weighted. Jitter in phi is
313        uniform in weight along a line of latitude.  With small theta and
314        phi ranging over +/- 180 this forms a wobbling disk.  With small
315        phi and theta ranging over +/- 90 this forms a wedge like a slice
316        of an orange.
317    azimuthal_equidistance (Postel)
318        <https://en.wikipedia.org/wiki/Azimuthal_equidistant_projection>
319        Preserves distance from center, and so is an excellent map for
320        representing a bivariate gaussian on the surface.  Theta and phi
321        operate identically, cutting wegdes from the antipode of the viewing
322        angle.  This unfortunately does not allow free movement in either
323        theta or phi since the orthogonal wobble decreases to 0 as the body
324        rotates through 180 degrees.
325    sinusoidal (Sanson-Flamsteed, Mercator equal-area)
326        <https://en.wikipedia.org/wiki/Sinusoidal_projection>
327        Preserves arc length with latitude, giving bad behaviour at
328        theta near +/- 90.  Theta and phi operate somewhat differently,
329        so a system with a-b-c dtheta-dphi-dpsi will not give the same
330        value as one with b-a-c dphi-dtheta-dpsi, as would be the case
331        for azimuthal equidistance.  Free movement using theta or phi
332        uniform over +/- 180 will work, but not as well as equirectangular
333        phi, with theta being slightly worse.  Computationally it is much
334        cheaper for wide theta-phi meshes since it excludes points which
335        lie outside the sinusoid near theta +/- 90 rather than packing
336        them close together as in equirectangle.  Note that the poles
337        will be slightly overweighted for theta > 90 with the circle
338        from theta at 90+dt winding backwards around the pole, overlapping
339        the circle from theta at 90-dt.
340    Guyou (hemisphere-in-a-square) **not weighted**
341        <https://en.wikipedia.org/wiki/Guyou_hemisphere-in-a-square_projection>
342        With tiling, allows rotation in phi or theta through +/- 180, with
343        uniform spacing.  Both theta and phi allow free rotation, with wobble
344        in the orthogonal direction reasonably well behaved (though not as
345        good as equirectangular phi). The forward/reverse transformations
346        relies on elliptic integrals that are somewhat expensive, so the
347        behaviour has to be very good to justify the cost and complexity.
348        The weighting function for each point has not yet been computed.
349        Note: run the module *guyou.py* directly and it will show the forward
350        and reverse mappings.
351    azimuthal_equal_area  **incomplete**
352        <https://en.wikipedia.org/wiki/Lambert_azimuthal_equal-area_projection>
353        Preserves the relative density of the surface patches.  Not that
354        useful and not completely implemented
355    Gauss-Kreuger **not implemented**
356        <https://en.wikipedia.org/wiki/Transverse_Mercator_projection#Ellipsoidal_transverse_Mercator>
357        Should allow free movement in theta, but phi is distorted.
358    """
359    # TODO: try Kent distribution instead of a gaussian warped by projection
360
361    if projection == 'equirectangular':  #define PROJECTION 1
362        def _project(theta_i, phi_j, psi):
363            latitude, longitude = theta_i, phi_j
364            return latitude, longitude, psi
365            #return Rx(phi_j)*Ry(theta_i)
366        def _weight(theta_i, phi_j, w_i, w_j):
367            return w_i*w_j*abs(cos(radians(theta_i)))
368    elif projection == 'sinusoidal':  #define PROJECTION 2
369        def _project(theta_i, phi_j, psi):
370            latitude = theta_i
371            scale = cos(radians(latitude))
372            longitude = phi_j/scale if abs(phi_j) < abs(scale)*180 else 0
373            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
374            return latitude, longitude, psi
375            #return Rx(longitude)*Ry(latitude)
376        def _project(theta_i, phi_j, w_i, w_j):
377            latitude = theta_i
378            scale = cos(radians(latitude))
379            active = 1 if abs(phi_j) < abs(scale)*180 else 0
380            return active*w_i*w_j
381    elif projection == 'guyou':  #define PROJECTION 3  (eventually?)
382        def _project(theta_i, phi_j, psi):
383            from .guyou import guyou_invert
384            #latitude, longitude = guyou_invert([theta_i], [phi_j])
385            longitude, latitude = guyou_invert([phi_j], [theta_i])
386            return latitude, longitude, psi
387            #return Rx(longitude[0])*Ry(latitude[0])
388        def _weight(theta_i, phi_j, w_i, w_j):
389            return w_i*w_j
390    elif projection == 'azimuthal_equidistance': 
391        # Note that calculates angles for Rz Ry rather than Rx Ry
392        def _project(theta_i, phi_j, psi):
393            latitude = sqrt(theta_i**2 + phi_j**2)
394            longitude = degrees(np.arctan2(phi_j, theta_i))
395            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
396            return latitude, longitude, psi-longitude, 'zyz'
397            #R = Rz(longitude)*Ry(latitude)*Rz(psi)
398            #return R_to_xyz(R)
399            #return Rz(longitude)*Ry(latitude)
400        def _weight(theta_i, phi_j, w_i, w_j):
401            # Weighting for each point comes from the integral:
402            #     \int\int I(q, lat, log) sin(lat) dlat dlog
403            # We are doing a conformal mapping from disk to sphere, so we need
404            # a change of variables g(theta, phi) -> (lat, long):
405            #     lat, long = sqrt(theta^2 + phi^2), arctan(phi/theta)
406            # giving:
407            #     dtheta dphi = det(J) dlat dlong
408            # where J is the jacobian from the partials of g. Using
409            #     R = sqrt(theta^2 + phi^2),
410            # then
411            #     J = [[x/R, Y/R], -y/R^2, x/R^2]]
412            # and
413            #     det(J) = 1/R
414            # with the final integral being:
415            #    \int\int I(q, theta, phi) sin(R)/R dtheta dphi
416            #
417            # This does approximately the right thing, decreasing the weight
418            # of each point as you go farther out on the disk, but it hasn't
419            # yet been checked against the 1D integral results. Prior
420            # to declaring this "good enough" and checking that integrals
421            # work in practice, we will examine alternative mappings.
422            #
423            # The issue is that the mapping does not support the case of free
424            # rotation about a single axis correctly, with a small deviation
425            # in the orthogonal axis independent of the first axis.  Like the
426            # usual polar coordiates integration, the integrated sections
427            # form wedges, though at least in this case the wedge cuts through
428            # the entire sphere, and treats theta and phi identically.
429            latitude = sqrt(theta_i**2 + phi_j**2)
430            weight = sin(radians(latitude))/latitude if latitude != 0 else 1
431            return weight*w_i*w_j if latitude < 180 else 0
432    elif projection == 'azimuthal_equal_area':
433        # Note that calculates angles for Rz Ry rather than Rx Ry
434        def _project(theta_i, phi_j, psi):
435            radius = min(1, sqrt(theta_i**2 + phi_j**2)/180)
436            latitude = 180-degrees(2*np.arccos(radius))
437            longitude = degrees(np.arctan2(phi_j, theta_i))
438            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
439            return latitude, longitude, psi, "zyz"
440            #R = Rz(longitude)*Ry(latitude)*Rz(psi)
441            #return R_to_xyz(R)
442            #return Rz(longitude)*Ry(latitude)
443        def _weight(theta_i, phi_j, w_i, w_j):
444            latitude = sqrt(theta_i**2 + phi_j**2)
445            weight = sin(radians(latitude))/latitude if latitude != 0 else 1
446            return weight*w_i*w_j if latitude < 180 else 0
447    else:
448        raise ValueError("unknown projection %r"%projection)
449
450    return _project, _weight
451
452def R_to_xyz(R):
453    """
454    Return phi, theta, psi Tait-Bryan angles corresponding to the given rotation matrix.
455
456    Extracting Euler Angles from a Rotation Matrix
457    Mike Day, Insomniac Games
458    https://d3cw3dd2w32x2b.cloudfront.net/wp-content/uploads/2012/07/euler-angles1.pdf
459    Based on: Shoemake’s “Euler Angle Conversion”, Graphics Gems IV, pp.  222-229
460    """
461    phi = np.arctan2(R[1, 2], R[2, 2])
462    theta = np.arctan2(-R[0, 2], np.sqrt(R[0, 0]**2 + R[0, 1]**2))
463    psi = np.arctan2(R[0, 1], R[0, 0])
464    return np.degrees(phi), np.degrees(theta), np.degrees(psi)
465
466def draw_mesh(axes, view, jitter, radius=1.2, n=11, dist='gaussian',
467              projection='equirectangular'):
468    """
469    Draw the dispersion mesh showing the theta-phi orientations at which
470    the model will be evaluated.
471    """
472
473    _project, _weight = get_projection(projection)
474    def _rotate(theta, phi, z):
475        dview = _project(theta, phi, 0.)
476        if len(dview) == 4: # hack for zyz coords
477            return Rz(dview[1])*Ry(dview[0])*z
478        else:
479            return Rx(dview[1])*Ry(dview[0])*z
480
481
482    dist_x = np.linspace(-1, 1, n)
483    weights = np.ones_like(dist_x)
484    if dist == 'gaussian':
485        dist_x *= 3
486        weights = exp(-0.5*dist_x**2)
487    elif dist == 'rectangle':
488        # Note: uses sasmodels ridiculous definition of rectangle width
489        dist_x *= sqrt(3)
490    elif dist == 'uniform':
491        pass
492    else:
493        raise ValueError("expected dist to be gaussian, rectangle or uniform")
494
495    # mesh in theta, phi formed by rotating z
496    dtheta, dphi, dpsi = jitter
497    z = np.matrix([[0], [0], [radius]])
498    points = np.hstack([_rotate(theta_i, phi_j, z)
499                        for theta_i in dtheta*dist_x
500                        for phi_j in dphi*dist_x])
501    dist_w = np.array([_weight(theta_i, phi_j, w_i, w_j)
502                       for w_i, theta_i in zip(weights, dtheta*dist_x)
503                       for w_j, phi_j in zip(weights, dphi*dist_x)])
504    #print(max(dist_w), min(dist_w), min(dist_w[dist_w > 0]))
505    points = points[:, dist_w > 0]
506    dist_w = dist_w[dist_w > 0]
507    dist_w /= max(dist_w)
508
509    # rotate relative to beam
510    points = orient_relative_to_beam(view, points)
511
512    x, y, z = [np.array(v).flatten() for v in points]
513    #plt.figure(2); plt.clf(); plt.hist(z, bins=np.linspace(-1, 1, 51))
514    axes.scatter(x, y, z, c=dist_w, marker='o', vmin=0., vmax=1.)
515
516def draw_labels(axes, view, jitter, text):
517    """
518    Draw text at a particular location.
519    """
520    labels, locations, orientations = zip(*text)
521    px, py, pz = zip(*locations)
522    dx, dy, dz = zip(*orientations)
523
524    px, py, pz = transform_xyz(view, jitter, px, py, pz)
525    dx, dy, dz = transform_xyz(view, jitter, dx, dy, dz)
526
527    # TODO: zdir for labels is broken, and labels aren't appearing.
528    for label, p, zdir in zip(labels, zip(px, py, pz), zip(dx, dy, dz)):
529        zdir = np.asarray(zdir).flatten()
530        axes.text(p[0], p[1], p[2], label, zdir=zdir)
531
532# Definition of rotation matrices comes from wikipedia:
533#    https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
534def Rx(angle):
535    """Construct a matrix to rotate points about *x* by *angle* degrees."""
536    angle = radians(angle)
537    rot = [[1, 0, 0],
538           [0, +cos(angle), -sin(angle)],
539           [0, +sin(angle), +cos(angle)]]
540    return np.matrix(rot)
541
542def Ry(angle):
543    """Construct a matrix to rotate points about *y* by *angle* degrees."""
544    angle = radians(angle)
545    rot = [[+cos(angle), 0, +sin(angle)],
546           [0, 1, 0],
547           [-sin(angle), 0, +cos(angle)]]
548    return np.matrix(rot)
549
550def Rz(angle):
551    """Construct a matrix to rotate points about *z* by *angle* degrees."""
552    angle = radians(angle)
553    rot = [[+cos(angle), -sin(angle), 0],
554           [+sin(angle), +cos(angle), 0],
555           [0, 0, 1]]
556    return np.matrix(rot)
557
558def transform_xyz(view, jitter, x, y, z):
559    """
560    Send a set of (x,y,z) points through the jitter and view transforms.
561    """
562    x, y, z = [np.asarray(v) for v in (x, y, z)]
563    shape = x.shape
564    points = np.matrix([x.flatten(), y.flatten(), z.flatten()])
565    points = apply_jitter(jitter, points)
566    points = orient_relative_to_beam(view, points)
567    x, y, z = [np.array(v).reshape(shape) for v in points]
568    return x, y, z
569
570def apply_jitter(jitter, points):
571    """
572    Apply the jitter transform to a set of points.
573
574    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
575    """
576    if jitter is None:
577        return points
578    # Hack to deal with the fact that azimuthal_equidistance uses euler angles
579    if len(jitter) == 4:
580        dtheta, dphi, dpsi, _ = jitter
581        points = Rz(dphi)*Ry(dtheta)*Rz(dpsi)*points
582    else:
583        dtheta, dphi, dpsi = jitter
584        points = Rx(dphi)*Ry(dtheta)*Rz(dpsi)*points
585    return points
586
587def orient_relative_to_beam(view, points):
588    """
589    Apply the view transform to a set of points.
590
591    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
592    """
593    theta, phi, psi = view
594    points = Rz(phi)*Ry(theta)*Rz(psi)*points # viewing angle
595    #points = Rz(psi)*Ry(pi/2-theta)*Rz(phi)*points # 1-D integration angles
596    #points = Rx(phi)*Ry(theta)*Rz(psi)*points  # angular dispersion angle
597    return points
598
599def orient_relative_to_beam_quaternion(view, points):
600    """
601    Apply the view transform to a set of points.
602
603    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
604
605    This variant uses quaternions rather than rotation matrices for the
606    computation.  It works but it is not used because it doesn't solve
607    any problems.  The challenge of mapping theta/phi/psi to SO(3) does
608    not disappear by calculating the transform differently.
609    """
610    theta, phi, psi = view
611    x, y, z = [1, 0, 0], [0, 1, 0], [0, 0, 1]
612    q = Quaternion(1, [0, 0, 0])
613    ## Compose a rotation about the three axes by rotating
614    ## the unit vectors before applying the rotation.
615    #q = Quaternion.from_angle_axis(theta, q.rot(x)) * q
616    #q = Quaternion.from_angle_axis(phi, q.rot(y)) * q
617    #q = Quaternion.from_angle_axis(psi, q.rot(z)) * q
618    ## The above turns out to be equivalent to reversing
619    ## the order of application, so ignore it and use below.
620    q = q * Quaternion.from_angle_axis(theta, x)
621    q = q * Quaternion.from_angle_axis(phi, y)
622    q = q * Quaternion.from_angle_axis(psi, z)
623    ## Reverse the order by post-multiply rather than pre-multiply
624    #q = Quaternion.from_angle_axis(theta, x) * q
625    #q = Quaternion.from_angle_axis(phi, y) * q
626    #q = Quaternion.from_angle_axis(psi, z) * q
627    #print("axes psi", q.rot(np.matrix([x, y, z]).T))
628    return q.rot(points)
629#orient_relative_to_beam = orient_relative_to_beam_quaternion
630 
631# Simple stand-alone quaternion class
632import numpy as np
633from copy import copy
634class Quaternion(object):
635    def __init__(self, w, r):
636         self.w = w
637         self.r = np.asarray(r, dtype='d')
638    @staticmethod
639    def from_angle_axis(theta, r):
640         theta = np.radians(theta)/2
641         r = np.asarray(r)
642         w = np.cos(theta)
643         r = np.sin(theta)*r/np.dot(r,r)
644         return Quaternion(w, r)
645    def __mul__(self, other):
646        if isinstance(other, Quaternion):
647            w = self.w*other.w - np.dot(self.r, other.r)
648            r = self.w*other.r + other.w*self.r + np.cross(self.r, other.r)
649            return Quaternion(w, r)
650    def rot(self, v):
651        v = np.asarray(v).T
652        use_transpose = (v.shape[-1] != 3)
653        if use_transpose: v = v.T
654        v = v + np.cross(2*self.r, np.cross(self.r, v) + self.w*v)
655        #v = v + 2*self.w*np.cross(self.r, v) + np.cross(2*self.r, np.cross(self.r, v))
656        if use_transpose: v = v.T
657        return v.T
658    def conj(self):
659        return Quaternion(self.w, -self.r)
660    def inv(self):
661        return self.conj()/self.norm()**2
662    def norm(self):
663        return np.sqrt(self.w**2 + np.sum(self.r**2))
664    def __str__(self):
665        return "%g%+gi%+gj%+gk"%(self.w, self.r[0], self.r[1], self.r[2])
666def test_qrot():
667    # Define rotation of 60 degrees around an axis in y-z that is 60 degrees from y.
668    # The rotation axis is determined by rotating the point [0, 1, 0] about x.
669    ax = Quaternion.from_angle_axis(60, [1, 0, 0]).rot([0, 1, 0])
670    q = Quaternion.from_angle_axis(60, ax)
671    # Set the point to be rotated, and its expected rotated position.
672    p = [1, -1, 2]
673    target = [(10+4*np.sqrt(3))/8, (1+2*np.sqrt(3))/8, (14-3*np.sqrt(3))/8]
674    #print(q, q.rot(p) - target)
675    assert max(abs(q.rot(p) - target)) < 1e-14 
676#test_qrot()
677#import sys; sys.exit()
678
679# translate between number of dimension of dispersity and the number of
680# points along each dimension.
681PD_N_TABLE = {
682    (0, 0, 0): (0, 0, 0),     # 0
683    (1, 0, 0): (100, 0, 0),   # 100
684    (0, 1, 0): (0, 100, 0),
685    (0, 0, 1): (0, 0, 100),
686    (1, 1, 0): (30, 30, 0),   # 900
687    (1, 0, 1): (30, 0, 30),
688    (0, 1, 1): (0, 30, 30),
689    (1, 1, 1): (15, 15, 15),  # 3375
690}
691
692def clipped_range(data, portion=1.0, mode='central'):
693    """
694    Determine range from data.
695
696    If *portion* is 1, use full range, otherwise use the center of the range
697    or the top of the range, depending on whether *mode* is 'central' or 'top'.
698    """
699    if portion == 1.0:
700        return data.min(), data.max()
701    elif mode == 'central':
702        data = np.sort(data.flatten())
703        offset = int(portion*len(data)/2 + 0.5)
704        return data[offset], data[-offset]
705    elif mode == 'top':
706        data = np.sort(data.flatten())
707        offset = int(portion*len(data) + 0.5)
708        return data[offset], data[-1]
709
710def draw_scattering(calculator, axes, view, jitter, dist='gaussian'):
711    """
712    Plot the scattering for the particular view.
713
714    *calculator* is returned from :func:`build_model`.  *axes* are the 3D axes
715    on which the data will be plotted.  *view* and *jitter* are the current
716    orientation and orientation dispersity.  *dist* is one of the sasmodels
717    weight distributions.
718    """
719    if dist == 'uniform':  # uniform is not yet in this branch
720        dist, scale = 'rectangle', 1/sqrt(3)
721    else:
722        scale = 1
723
724    # add the orientation parameters to the model parameters
725    theta, phi, psi = view
726    theta_pd, phi_pd, psi_pd = [scale*v for v in jitter]
727    theta_pd_n, phi_pd_n, psi_pd_n = PD_N_TABLE[(theta_pd > 0, phi_pd > 0, psi_pd > 0)]
728    ## increase pd_n for testing jitter integration rather than simple viz
729    #theta_pd_n, phi_pd_n, psi_pd_n = [5*v for v in (theta_pd_n, phi_pd_n, psi_pd_n)]
730
731    pars = dict(
732        theta=theta, theta_pd=theta_pd, theta_pd_type=dist, theta_pd_n=theta_pd_n,
733        phi=phi, phi_pd=phi_pd, phi_pd_type=dist, phi_pd_n=phi_pd_n,
734        psi=psi, psi_pd=psi_pd, psi_pd_type=dist, psi_pd_n=psi_pd_n,
735    )
736    pars.update(calculator.pars)
737
738    # compute the pattern
739    qx, qy = calculator._data.x_bins, calculator._data.y_bins
740    Iqxy = calculator(**pars).reshape(len(qx), len(qy))
741
742    # scale it and draw it
743    Iqxy = np.log(Iqxy)
744    if calculator.limits:
745        # use limits from orientation (0,0,0)
746        vmin, vmax = calculator.limits
747    else:
748        vmax = Iqxy.max()
749        vmin = vmax*10**-7
750        #vmin, vmax = clipped_range(Iqxy, portion=portion, mode='top')
751    #print("range",(vmin,vmax))
752    #qx, qy = np.meshgrid(qx, qy)
753    if 0:
754        level = np.asarray(255*(Iqxy - vmin)/(vmax - vmin), 'i')
755        level[level < 0] = 0
756        colors = plt.get_cmap()(level)
757        axes.plot_surface(qx, qy, -1.1, facecolors=colors)
758    elif 1:
759        axes.contourf(qx/qx.max(), qy/qy.max(), Iqxy, zdir='z', offset=-1.1,
760                      levels=np.linspace(vmin, vmax, 24))
761    else:
762        axes.pcolormesh(qx, qy, Iqxy)
763
764def build_model(model_name, n=150, qmax=0.5, **pars):
765    """
766    Build a calculator for the given shape.
767
768    *model_name* is any sasmodels model.  *n* and *qmax* define an n x n mesh
769    on which to evaluate the model.  The remaining parameters are stored in
770    the returned calculator as *calculator.pars*.  They are used by
771    :func:`draw_scattering` to set the non-orientation parameters in the
772    calculation.
773
774    Returns a *calculator* function which takes a dictionary or parameters and
775    produces Iqxy.  The Iqxy value needs to be reshaped to an n x n matrix
776    for plotting.  See the :class:`sasmodels.direct_model.DirectModel` class
777    for details.
778    """
779    from sasmodels.core import load_model_info, build_model as build_sasmodel
780    from sasmodels.data import empty_data2D
781    from sasmodels.direct_model import DirectModel
782
783    model_info = load_model_info(model_name)
784    model = build_sasmodel(model_info) #, dtype='double!')
785    q = np.linspace(-qmax, qmax, n)
786    data = empty_data2D(q, q)
787    calculator = DirectModel(data, model)
788
789    # stuff the values for non-orientation parameters into the calculator
790    calculator.pars = pars.copy()
791    calculator.pars.setdefault('backgound', 1e-3)
792
793    # fix the data limits so that we can see if the pattern fades
794    # under rotation or angular dispersion
795    Iqxy = calculator(theta=0, phi=0, psi=0, **calculator.pars)
796    Iqxy = np.log(Iqxy)
797    vmin, vmax = clipped_range(Iqxy, 0.95, mode='top')
798    calculator.limits = vmin, vmax+1
799
800    return calculator
801
802def select_calculator(model_name, n=150, size=(10, 40, 100)):
803    """
804    Create a model calculator for the given shape.
805
806    *model_name* is one of sphere, cylinder, ellipsoid, triaxial_ellipsoid,
807    parallelepiped or bcc_paracrystal. *n* is the number of points to use
808    in the q range.  *qmax* is chosen based on model parameters for the
809    given model to show something intersting.
810
811    Returns *calculator* and tuple *size* (a,b,c) giving minor and major
812    equitorial axes and polar axis respectively.  See :func:`build_model`
813    for details on the returned calculator.
814    """
815    a, b, c = size
816    d_factor = 0.06  # for paracrystal models
817    if model_name == 'sphere':
818        calculator = build_model('sphere', n=n, radius=c)
819        a = b = c
820    elif model_name == 'sc_paracrystal':
821        a = b = c
822        dnn = c
823        radius = 0.5*c
824        calculator = build_model('sc_paracrystal', n=n, dnn=dnn,
825                                 d_factor=d_factor, radius=(1-d_factor)*radius,
826                                 background=0)
827    elif model_name == 'fcc_paracrystal':
828        a = b = c
829        # nearest neigbour distance dnn should be 2 radius, but I think the
830        # model uses lattice spacing rather than dnn in its calculations
831        dnn = 0.5*c
832        radius = sqrt(2)/4 * c
833        calculator = build_model('fcc_paracrystal', n=n, dnn=dnn,
834                                 d_factor=d_factor, radius=(1-d_factor)*radius,
835                                 background=0)
836    elif model_name == 'bcc_paracrystal':
837        a = b = c
838        # nearest neigbour distance dnn should be 2 radius, but I think the
839        # model uses lattice spacing rather than dnn in its calculations
840        dnn = 0.5*c
841        radius = sqrt(3)/2 * c
842        calculator = build_model('bcc_paracrystal', n=n, dnn=dnn,
843                                 d_factor=d_factor, radius=(1-d_factor)*radius,
844                                 background=0)
845    elif model_name == 'cylinder':
846        calculator = build_model('cylinder', n=n, qmax=0.3, radius=b, length=c)
847        a = b
848    elif model_name == 'ellipsoid':
849        calculator = build_model('ellipsoid', n=n, qmax=1.0,
850                                 radius_polar=c, radius_equatorial=b)
851        a = b
852    elif model_name == 'triaxial_ellipsoid':
853        calculator = build_model('triaxial_ellipsoid', n=n, qmax=0.5,
854                                 radius_equat_minor=a,
855                                 radius_equat_major=b,
856                                 radius_polar=c)
857    elif model_name == 'parallelepiped':
858        calculator = build_model('parallelepiped', n=n, a=a, b=b, c=c)
859    else:
860        raise ValueError("unknown model %s"%model_name)
861
862    return calculator, (a, b, c)
863
864SHAPES = [
865    'parallelepiped',
866    'sphere', 'ellipsoid', 'triaxial_ellipsoid',
867    'cylinder',
868    'fcc_paracrystal', 'bcc_paracrystal', 'sc_paracrystal',
869]
870
871DRAW_SHAPES = {
872    'fcc_paracrystal': draw_fcc,
873    'bcc_paracrystal': draw_bcc,
874    'sc_paracrystal': draw_sc,
875    'parallelepiped': draw_parallelepiped,
876}
877
878DISTRIBUTIONS = [
879    'gaussian', 'rectangle', 'uniform',
880]
881DIST_LIMITS = {
882    'gaussian': 30,
883    'rectangle': 90/sqrt(3),
884    'uniform': 90,
885}
886
887
888def run(model_name='parallelepiped', size=(10, 40, 100), 
889        view=(0, 0, 0), jitter=(0, 0, 0),
890        dist='gaussian', mesh=30,
891        projection='equirectangular'):
892    """
893    Show an interactive orientation and jitter demo.
894
895    *model_name* is one of: sphere, ellipsoid, triaxial_ellipsoid,
896    parallelepiped, cylinder, or sc/fcc/bcc_paracrystal
897
898    *size* gives the dimensions (a, b, c) of the shape.
899
900    *view* gives the initial view (theta, phi, psi) of the shape.
901
902    *view* gives the initial jitter (dtheta, dphi, dpsi) of the shape.
903
904    *dist* is the type of dispersition: gaussian, rectangle, or uniform.
905
906    *mesh* is the number of points in the dispersion mesh.
907
908    *projection* is the map projection to use for the mesh: equirectangular,
909    sinusoidal, guyou, azimuthal_equidistance, or azimuthal_equal_area.
910    """
911    # projection number according to 1-order position in list, but
912    # only 1 and 2 are implemented so far.
913    from sasmodels import generate
914    generate.PROJECTION = PROJECTIONS.index(projection) + 1
915    if generate.PROJECTION > 2:
916        print("*** PROJECTION %s not implemented in scattering function ***"%projection)
917        generate.PROJECTION = 2
918
919    # set up calculator
920    calculator, size = select_calculator(model_name, n=150, size=size)
921    draw_shape = DRAW_SHAPES.get(model_name, draw_parallelepiped)
922    #draw_shape = draw_fcc
923
924    ## uncomment to set an independent the colour range for every view
925    ## If left commented, the colour range is fixed for all views
926    calculator.limits = None
927
928    PLOT_ENGINE(calculator, draw_shape, size, view, jitter, dist, mesh, projection)
929
930def mpl_plot(calculator, draw_shape, size, view, jitter, dist, mesh, projection):
931    import mpl_toolkits.mplot3d  # Adds projection='3d' option to subplot
932    import matplotlib.pyplot as plt
933    from matplotlib.widgets import Slider
934
935    ## create the plot window
936    #plt.hold(True)
937    plt.subplots(num=None, figsize=(5.5, 5.5))
938    plt.set_cmap('gist_earth')
939    plt.clf()
940    plt.gcf().canvas.set_window_title(projection)
941    #gs = gridspec.GridSpec(2,1,height_ratios=[4,1])
942    #axes = plt.subplot(gs[0], projection='3d')
943    axes = plt.axes([0.0, 0.2, 1.0, 0.8], projection='3d')
944    try:  # CRUFT: not all versions of matplotlib accept 'square' 3d projection
945        axes.axis('square')
946    except Exception:
947        pass
948
949    axcolor = 'lightgoldenrodyellow'
950
951    ## add control widgets to plot
952    axes_theta = plt.axes([0.1, 0.15, 0.45, 0.04], facecolor=axcolor)
953    axes_phi = plt.axes([0.1, 0.1, 0.45, 0.04], facecolor=axcolor)
954    axes_psi = plt.axes([0.1, 0.05, 0.45, 0.04], facecolor=axcolor)
955    stheta = Slider(axes_theta, 'Theta', -90, 90, valinit=0)
956    sphi = Slider(axes_phi, 'Phi', -180, 180, valinit=0)
957    spsi = Slider(axes_psi, 'Psi', -180, 180, valinit=0)
958
959    axes_dtheta = plt.axes([0.75, 0.15, 0.15, 0.04], facecolor=axcolor)
960    axes_dphi = plt.axes([0.75, 0.1, 0.15, 0.04], facecolor=axcolor)
961    axes_dpsi = plt.axes([0.75, 0.05, 0.15, 0.04], facecolor=axcolor)
962    # Note: using ridiculous definition of rectangle distribution, whose width
963    # in sasmodels is sqrt(3) times the given width.  Divide by sqrt(3) to keep
964    # the maximum width to 90.
965    dlimit = DIST_LIMITS[dist]
966    sdtheta = Slider(axes_dtheta, 'dTheta', 0, 2*dlimit, valinit=0)
967    sdphi = Slider(axes_dphi, 'dPhi', 0, 2*dlimit, valinit=0)
968    sdpsi = Slider(axes_dpsi, 'dPsi', 0, 2*dlimit, valinit=0)
969
970    ## initial view and jitter
971    theta, phi, psi = view
972    stheta.set_val(theta)
973    sphi.set_val(phi)
974    spsi.set_val(psi)
975    dtheta, dphi, dpsi = jitter
976    sdtheta.set_val(dtheta)
977    sdphi.set_val(dphi)
978    sdpsi.set_val(dpsi)
979
980    ## callback to draw the new view
981    def update(val, axis=None):
982        view = stheta.val, sphi.val, spsi.val
983        jitter = sdtheta.val, sdphi.val, sdpsi.val
984        # set small jitter as 0 if multiple pd dims
985        dims = sum(v > 0 for v in jitter)
986        limit = [0, 0.5, 5, 5][dims]
987        jitter = [0 if v < limit else v for v in jitter]
988        axes.cla()
989
990        ## Visualize as person on globe
991        #draw_sphere(axes)
992        #draw_person_on_sphere(axes, view)
993
994        ## Move beam instead of shape
995        #draw_beam(axes, -view[:2])
996        #draw_jitter(axes, (0,0,0), (0,0,0), views=3)
997
998        ## Move shape and draw scattering
999        draw_beam(axes, (0, 0))
1000        draw_jitter(axes, view, jitter, dist=dist, size=size, 
1001                    draw_shape=draw_shape, projection=projection, views=3)
1002        draw_mesh(axes, view, jitter, dist=dist, n=mesh, projection=projection)
1003        draw_scattering(calculator, axes, view, jitter, dist=dist)
1004
1005        plt.gcf().canvas.draw()
1006
1007    ## bind control widgets to view updater
1008    stheta.on_changed(lambda v: update(v, 'theta'))
1009    sphi.on_changed(lambda v: update(v, 'phi'))
1010    spsi.on_changed(lambda v: update(v, 'psi'))
1011    sdtheta.on_changed(lambda v: update(v, 'dtheta'))
1012    sdphi.on_changed(lambda v: update(v, 'dphi'))
1013    sdpsi.on_changed(lambda v: update(v, 'dpsi'))
1014
1015    ## initialize view
1016    update(None, 'phi')
1017
1018    ## go interactive
1019    plt.show()
1020
1021
1022def map_colors(z, kw):
1023    from matplotlib import cm
1024
1025    cmap = kw.pop('cmap', cm.coolwarm)
1026    alpha = kw.pop('alpha', None)
1027    vmin = kw.pop('vmin', z.min())
1028    vmax = kw.pop('vmax', z.max())
1029    c = kw.pop('c', None)
1030    color = kw.pop('color', c)
1031    if color is None:
1032        znorm = ((z - vmin) / (vmax - vmin)).clip(0, 1)
1033        color = cmap(znorm)
1034    elif isinstance(color, np.ndarray) and color.shape == z.shape:
1035        color = cmap(color)
1036    if alpha is None:
1037        if isinstance(color, np.ndarray):
1038            color = color[..., :3]
1039    else:
1040        color[..., 3] = alpha
1041    kw['color'] = color
1042
1043def make_vec(*args, flat=False):
1044    if flat:
1045        return [np.asarray(v, 'd').flatten() for v in args]
1046    else:
1047        return [np.asarray(v, 'd') for v in args]
1048
1049def make_image(z, kw):
1050    import PIL.Image
1051    from matplotlib import cm
1052
1053    cmap = kw.pop('cmap', cm.coolwarm)
1054
1055    znorm = (z-z.min())/z.ptp()
1056    c = cmap(znorm)
1057    c = c[..., :3]
1058    rgb = np.asarray(c*255, 'u1')
1059    image = PIL.Image.fromarray(rgb, mode='RGB')
1060    return image
1061
1062
1063_IPV_MARKERS = {
1064    'o': 'sphere',
1065}
1066_IPV_COLORS = {
1067    'w': 'white',
1068    'k': 'black',
1069    'c': 'cyan',
1070    'm': 'magenta',
1071    'y': 'yellow',
1072    'r': 'red',
1073    'g': 'green',
1074    'b': 'blue',
1075}
1076def ipv_fix_color(kw):
1077    kw.pop('alpha', None)
1078    color = kw.get('color', None)
1079    if isinstance(color, str):
1080        color = _IPV_COLORS.get(color, color)
1081        kw['color'] = color
1082
1083
1084def ipv_plot(calculator, draw_shape, size, view, jitter, dist, mesh, projection):
1085    import ipywidgets as widgets
1086    import ipyvolume as ipv
1087
1088    class Plotter:
1089        def plot(self, x, y, z, **kw):
1090            ipv_fix_color(kw)
1091            x, y, z = make_vec(x, y, z)
1092            ipv.plot(x, y, z, **kw)
1093        def plot_surface(self, x, y, z, **kw):
1094            ipv_fix_color(kw)
1095            x, y, z = make_vec(x, y, z)
1096            ipv.plot_surface(x, y, z, **kw)
1097        def plot_trisurf(self, x, y, triangles=None, Z=None, **kw):
1098            ipv_fix_color(kw)
1099            x, y, z = make_vec(x, y, Z)
1100            if triangles is not None:
1101                triangles = np.asarray(triangles)
1102            ipv.plot_trisurf(x, y, z, triangles=triangles, **kw)
1103        def scatter(self, x, y, z, **kw):
1104            x, y, z = make_vec(x, y, z)
1105            map_colors(z, kw)
1106            marker = kw.get('marker', None)
1107            kw['marker'] = _IPV_MARKERS.get(marker, marker)
1108            ipv.scatter(x, y, z, **kw)
1109        def contourf(self, x, y, v, zdir='z', offset=0, levels=None, **kw):
1110            # Don't use contour for now (although we might want to later)
1111            self.pcolor(x, y, v, zdir='z', offset=offset, **kw)
1112        def pcolor(self, x, y, v, zdir='z', offset=0, **kw):
1113            x, y, v = make_vec(x, y, v)
1114            image = make_image(v, kw)
1115            xmin, xmax = x.min(), x.max()
1116            ymin, ymax = y.min(), y.max()
1117            x = np.array([[xmin, xmax], [xmin, xmax]])
1118            y = np.array([[ymin, ymin], [ymax, ymax]])
1119            z = x*0 + offset
1120            u = np.array([[0., 1], [0, 1]])
1121            v = np.array([[0., 0], [1, 1]])
1122            ipv.plot_mesh(x, y, z, u=u, v=v, texture=image, wireframe=False)   
1123        def text(self, *args, **kw):
1124            pass
1125        def set_xlim(self, limits):
1126            ipv.xlim(*limits)
1127        def set_ylim(self, limits):
1128            ipv.ylim(*limits)
1129        def set_zlim(self, limits):
1130            ipv.zlim(*limits)
1131        def set_axes_on(self):
1132            ipv.style.axis_on()
1133        def set_axis_off(self):
1134            ipv.style.axes_off()
1135    axes = Plotter()
1136
1137
1138    def draw(view, jitter):
1139        camera = ipv.gcf().camera
1140        #print(ipv.gcf().__dict__.keys())
1141        #print(dir(ipv.gcf()))
1142        ipv.figure()
1143
1144        # set small jitter as 0 if multiple pd dims
1145        dims = sum(v > 0 for v in jitter)
1146        limit = [0, 0.5, 5, 5][dims]
1147        jitter = [0 if v < limit else v for v in jitter]
1148
1149        ## Visualize as person on globe
1150        #draw_beam(axes, (0, 0))
1151        #draw_sphere(axes)
1152        #draw_person_on_sphere(axes, view)
1153
1154        ## Move beam instead of shape
1155        #draw_beam(axes, view=(-view[0], -view[1]))
1156        #draw_jitter(axes, view=(0,0,0), jitter=(0,0,0))
1157
1158        ## Move shape and draw scattering
1159        draw_beam(axes, (0, 0))
1160        draw_jitter(axes, view, jitter, dist=dist, size=size, 
1161                    draw_shape=draw_shape, projection=projection)
1162        #draw_mesh(axes, view, jitter, dist=dist, n=mesh, projection=projection)
1163        #draw_scattering(calculator, axes, view, jitter, dist=dist)
1164   
1165        ipv.style.box_off()
1166        #ipv.style.axes_off()
1167        ipv.xyzlabel(" ", " ", " ")
1168
1169        ipv.gcf().camera = camera
1170        ipv.show()
1171
1172
1173    trange, prange = (-180., 180., 1.), (-180., 180., 1.)
1174    dtrange, dprange = (0., 180., 1.), (0., 180., 1.)
1175
1176    ## Super simple interfaca, but uses non-ascii variable namese
1177    # Ξ φ ψ Δξ Δφ Δψ
1178    #def update(**kw):
1179    #    view = kw['Ξ'], kw['φ'], kw['ψ']
1180    #    jitter = kw['Δξ'], kw['Δφ'], kw['Δψ']
1181    #    draw(view, jitter)
1182    #widgets.interact(update, Ξ=trange, φ=prange, ψ=prange, Δξ=dtrange, Δφ=dprange, Δψ=dprange)
1183
1184    def update(theta, phi, psi, dtheta, dphi, dpsi):
1185        draw(view=(theta, phi, psi), jitter=(dtheta, dphi, dpsi))
1186
1187    def slider(name, slice, init=0.):
1188        return widgets.FloatSlider(
1189            value=init,
1190            min=slice[0],
1191            max=slice[1],
1192            step=slice[2],
1193            description=name,
1194            disabled=False,
1195            continuous_update=False,
1196            orientation='horizontal',
1197            readout=True,
1198            readout_format='.1f',
1199            )
1200    theta = slider('Ξ', trange, view[0])
1201    phi = slider('φ', prange, view[1])
1202    psi = slider('ψ', prange, view[2])
1203    dtheta = slider('Δξ', dtrange, jitter[0])
1204    dphi = slider('Δφ', dprange, jitter[1])
1205    dpsi = slider('Δψ', dprange, jitter[2])
1206    fields = {
1207        'theta': theta, 'phi': phi, 'psi': psi,
1208        'dtheta': dtheta, 'dphi': dphi, 'dpsi': dpsi,
1209    }
1210    ui = widgets.HBox([
1211        widgets.VBox([theta, phi, psi]),
1212        widgets.VBox([dtheta, dphi, dpsi])
1213    ])
1214
1215    out = widgets.interactive_output(update, fields)
1216    display(ui, out)
1217
1218
1219_ENGINES = {
1220    "matplotlib": mpl_plot,
1221    "mpl": mpl_plot,
1222    #"plotly": plotly_plot,
1223    "ipvolume": ipv_plot,
1224    "ipv": ipv_plot,
1225}
1226PLOT_ENGINE = _ENGINES["matplotlib"]
1227def set_plotter(name):
1228    global PLOT_ENGINE
1229    PLOT_ENGINE = _ENGINES[name]
1230
1231def main():
1232    parser = argparse.ArgumentParser(
1233        description="Display jitter",
1234        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
1235        )
1236    parser.add_argument('-p', '--projection', choices=PROJECTIONS,
1237                        default=PROJECTIONS[0],
1238                        help='coordinate projection')
1239    parser.add_argument('-s', '--size', type=str, default='10,40,100',
1240                        help='a,b,c lengths')
1241    parser.add_argument('-v', '--view', type=str, default='0,0,0',
1242                        help='initial view angles')
1243    parser.add_argument('-j', '--jitter', type=str, default='0,0,0',
1244                        help='initial angular dispersion')
1245    parser.add_argument('-d', '--distribution', choices=DISTRIBUTIONS,
1246                        default=DISTRIBUTIONS[0],
1247                        help='jitter distribution')
1248    parser.add_argument('-m', '--mesh', type=int, default=30,
1249                        help='#points in theta-phi mesh')
1250    parser.add_argument('shape', choices=SHAPES, nargs='?', default=SHAPES[0],
1251                        help='oriented shape')
1252    opts = parser.parse_args()
1253    size = tuple(float(v) for v in opts.size.split(','))
1254    view = tuple(float(v) for v in opts.view.split(','))
1255    jitter = tuple(float(v) for v in opts.jitter.split(','))
1256    run(opts.shape, size=size, view=view, jitter=jitter,
1257        mesh=opts.mesh, dist=opts.distribution,
1258        projection=opts.projection)
1259
1260if __name__ == "__main__":
1261    main()
Note: See TracBrowser for help on using the repository browser.