# source:sasmodels/sasmodels/jitter.py@802c412

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

lint reduction

• Property mode set to 100755
File size: 31.9 KB
Line
1#!/usr/bin/env python
2"""
3Jitter Explorer
4===============
5
6Application to explore orientation angle and angular dispersity.
7"""
8from __future__ import division, print_function
9
10import argparse
11
12try: # CRUFT: travis-ci does not support mpl_toolkits.mplot3d
13    import mpl_toolkits.mplot3d  # Adds projection='3d' option to subplot
14except ImportError:
15    pass
16
17import matplotlib.pyplot as plt
18from matplotlib.widgets import Slider
19import numpy as np
20from numpy import pi, cos, sin, sqrt, exp, degrees, radians
21
22def draw_beam(axes, view=(0, 0)):
23    """
24    Draw the beam going from source at (0, 0, 1) to detector at (0, 0, -1)
25    """
26    #axes.plot([0,0],[0,0],[1,-1])
27    #axes.scatter([0]*100,[0]*100,np.linspace(1, -1, 100), alpha=0.8)
28
29    steps = 25
30    u = np.linspace(0, 2 * np.pi, steps)
31    v = np.linspace(-1, 1, steps)
32
33    r = 0.02
34    x = r*np.outer(np.cos(u), np.ones_like(v))
35    y = r*np.outer(np.sin(u), np.ones_like(v))
36    z = 1.3*np.outer(np.ones_like(u), v)
37
38    theta, phi = view
39    shape = x.shape
40    points = np.matrix([x.flatten(), y.flatten(), z.flatten()])
41    points = Rz(phi)*Ry(theta)*points
42    x, y, z = [v.reshape(shape) for v in points]
43
44    axes.plot_surface(x, y, z, rstride=4, cstride=4, color='y', alpha=0.5)
45
46def draw_ellipsoid(axes, size, view, jitter, steps=25, alpha=1):
47    """Draw an ellipsoid."""
48    a, b, c = size
49    u = np.linspace(0, 2 * np.pi, steps)
50    v = np.linspace(0, np.pi, steps)
51    x = a*np.outer(np.cos(u), np.sin(v))
52    y = b*np.outer(np.sin(u), np.sin(v))
53    z = c*np.outer(np.ones_like(u), np.cos(v))
54    x, y, z = transform_xyz(view, jitter, x, y, z)
55
56    axes.plot_surface(x, y, z, rstride=4, cstride=4, color='w', alpha=alpha)
57
58    draw_labels(axes, view, jitter, [
59        ('c+', [+0, +0, +c], [+1, +0, +0]),
60        ('c-', [+0, +0, -c], [+0, +0, -1]),
61        ('a+', [+a, +0, +0], [+0, +0, +1]),
62        ('a-', [-a, +0, +0], [+0, +0, -1]),
63        ('b+', [+0, +b, +0], [-1, +0, +0]),
64        ('b-', [+0, -b, +0], [-1, +0, +0]),
65    ])
66
67def draw_sc(axes, size, view, jitter, steps=None, alpha=1):
68    """Draw points for simple cubic paracrystal"""
69    atoms = _build_sc()
70    _draw_crystal(axes, size, view, jitter, atoms=atoms)
71
72def draw_fcc(axes, size, view, jitter, steps=None, alpha=1):
73    """Draw points for face-centered cubic paracrystal"""
74    # Build the simple cubic crystal
75    atoms = _build_sc()
76    # Define the centers for each face
77    # x planes at -1, 0, 1 have four centers per plane, at +/- 0.5 in y and z
78    x, y, z = (
79        [-1]*4 + [0]*4 + [1]*4,
80        ([-0.5]*2 + [0.5]*2)*3,
81        [-0.5, 0.5]*12,
82    )
83    # y and z planes can be generated by substituting x for y and z respectively
84    atoms.extend(zip(x+y+z, y+z+x, z+x+y))
85    _draw_crystal(axes, size, view, jitter, atoms=atoms)
86
87def draw_bcc(axes, size, view, jitter, steps=None, alpha=1):
88    """Draw points for body-centered cubic paracrystal"""
89    # Build the simple cubic crystal
90    atoms = _build_sc()
91    # Define the centers for each octant
92    # x plane at +/- 0.5 have four centers per plane at +/- 0.5 in y and z
93    x, y, z = (
94        [-0.5]*4 + [0.5]*4,
95        ([-0.5]*2 + [0.5]*2)*2,
96        [-0.5, 0.5]*8,
97    )
98    atoms.extend(zip(x, y, z))
99    _draw_crystal(axes, size, view, jitter, atoms=atoms)
100
101def _draw_crystal(axes, size, view, jitter, steps=None, alpha=1, atoms=None):
102    atoms, size = np.asarray(atoms, 'd').T, np.asarray(size, 'd')
103    x, y, z = atoms*size[:, None]
104    x, y, z = transform_xyz(view, jitter, x, y, z)
105    axes.scatter([x[0]], [y[0]], [z[0]], c='yellow', marker='o')
106    axes.scatter(x[1:], y[1:], z[1:], c='r', marker='o')
107
108def _build_sc():
109    # three planes of 9 dots for x at -1, 0 and 1
110    x, y, z = (
111        [-1]*9 + [0]*9 + [1]*9,
112        ([-1]*3 + [0]*3 + [1]*3)*3,
113        [-1, 0, 1]*9,
114    )
115    atoms = list(zip(x, y, z))
116    #print(list(enumerate(atoms)))
117    # Pull the dot at (0, 0, 1) to the front of the list
118    # It will be highlighted in the view
119    index = 14
120    highlight = atoms[index]
121    del atoms[index]
122    atoms.insert(0, highlight)
123    return atoms
124
125def draw_parallelepiped(axes, size, view, jitter, steps=None, alpha=1):
126    """Draw a parallelepiped."""
127    a, b, c = size
128    x = a*np.array([+1, -1, +1, -1, +1, -1, +1, -1])
129    y = b*np.array([+1, +1, -1, -1, +1, +1, -1, -1])
130    z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1])
131    tri = np.array([
132        # counter clockwise triangles
133        # z: up/down, x: right/left, y: front/back
134        [0, 1, 2], [3, 2, 1], # top face
135        [6, 5, 4], [5, 6, 7], # bottom face
136        [0, 2, 6], [6, 4, 0], # right face
137        [1, 5, 7], [7, 3, 1], # left face
138        [2, 3, 6], [7, 6, 3], # front face
139        [4, 1, 0], [5, 1, 4], # back face
140    ])
141
142    x, y, z = transform_xyz(view, jitter, x, y, z)
143    axes.plot_trisurf(x, y, triangles=tri, Z=z, color='w', alpha=alpha)
144
145    # Draw pink face on box.
146    # Since I can't control face color, instead draw a thin box situated just
147    # in front of the "c+" face.  Use the c face so that rotations about psi
148    # rotate that face.
149    if 1:
150        x = a*np.array([+1, -1, +1, -1, +1, -1, +1, -1])
151        y = b*np.array([+1, +1, -1, -1, +1, +1, -1, -1])
152        z = c*np.array([+1, +1, +1, +1, -1, -1, -1, -1])
153        x, y, z = transform_xyz(view, jitter, x, y, abs(z)+0.001)
154        axes.plot_trisurf(x, y, triangles=tri, Z=z, color=[1, 0.6, 0.6], alpha=alpha)
155
156    draw_labels(axes, view, jitter, [
157        ('c+', [+0, +0, +c], [+1, +0, +0]),
158        ('c-', [+0, +0, -c], [+0, +0, -1]),
159        ('a+', [+a, +0, +0], [+0, +0, +1]),
160        ('a-', [-a, +0, +0], [+0, +0, -1]),
161        ('b+', [+0, +b, +0], [-1, +0, +0]),
162        ('b-', [+0, -b, +0], [-1, +0, +0]),
163    ])
164
166    """Draw a sphere"""
167    u = np.linspace(0, 2 * np.pi, steps)
168    v = np.linspace(0, np.pi, steps)
169
170    x = radius * np.outer(np.cos(u), np.sin(v))
171    y = radius * np.outer(np.sin(u), np.sin(v))
172    z = radius * np.outer(np.ones(np.size(u)), np.cos(v))
173    axes.plot_surface(x, y, z, rstride=4, cstride=4, color='w')
174
175def draw_jitter(axes, view, jitter, dist='gaussian', size=(0.1, 0.4, 1.0),
176                draw_shape=draw_parallelepiped):
177    """
178    Represent jitter as a set of shapes at different orientations.
179    """
180    # set max diagonal to 0.95
181    scale = 0.95/sqrt(sum(v**2 for v in size))
182    size = tuple(scale*v for v in size)
183
184    #np.random.seed(10)
185    #cloud = np.random.randn(10,3)
186    cloud = [
187        [-1, -1, -1],
188        [-1, -1, +0],
189        [-1, -1, +1],
190        [-1, +0, -1],
191        [-1, +0, +0],
192        [-1, +0, +1],
193        [-1, +1, -1],
194        [-1, +1, +0],
195        [-1, +1, +1],
196        [+0, -1, -1],
197        [+0, -1, +0],
198        [+0, -1, +1],
199        [+0, +0, -1],
200        [+0, +0, +0],
201        [+0, +0, +1],
202        [+0, +1, -1],
203        [+0, +1, +0],
204        [+0, +1, +1],
205        [+1, -1, -1],
206        [+1, -1, +0],
207        [+1, -1, +1],
208        [+1, +0, -1],
209        [+1, +0, +0],
210        [+1, +0, +1],
211        [+1, +1, -1],
212        [+1, +1, +0],
213        [+1, +1, +1],
214    ]
215    dtheta, dphi, dpsi = jitter
216    if dtheta == 0:
217        cloud = [v for v in cloud if v[0] == 0]
218    if dphi == 0:
219        cloud = [v for v in cloud if v[1] == 0]
220    if dpsi == 0:
221        cloud = [v for v in cloud if v[2] == 0]
222    draw_shape(axes, size, view, [0, 0, 0], steps=100, alpha=0.8)
223    scale = {'gaussian':1, 'rectangle':1/sqrt(3), 'uniform':1/3}[dist]
224    for point in cloud:
225        delta = [scale*dtheta*point[0], scale*dphi*point[1], scale*dpsi*point[2]]
226        draw_shape(axes, size, view, delta, alpha=0.8)
227    for v in 'xyz':
228        a, b, c = size
229        lim = np.sqrt(a**2 + b**2 + c**2)
230        getattr(axes, 'set_'+v+'lim')([-lim, lim])
231        getattr(axes, v+'axis').label.set_text(v)
232
233PROJECTIONS = [
234    # in order of PROJECTION number; do not change without updating the
235    # constants in kernel_iq.c
236    'equirectangular', 'sinusoidal', 'guyou', 'azimuthal_equidistance',
237    'azimuthal_equal_area',
238]
239def draw_mesh(axes, view, jitter, radius=1.2, n=11, dist='gaussian',
240              projection='equirectangular'):
241    """
242    Draw the dispersion mesh showing the theta-phi orientations at which
243    the model will be evaluated.
244
245    jitter projections
246    <https://en.wikipedia.org/wiki/List_of_map_projections>
247
248    equirectangular (standard latitude-longitude mesh)
249        <https://en.wikipedia.org/wiki/Equirectangular_projection>
250        Allows free movement in phi (around the equator), but theta is
251        limited to +/- 90, and points are cos-weighted. Jitter in phi is
252        uniform in weight along a line of latitude.  With small theta and
253        phi ranging over +/- 180 this forms a wobbling disk.  With small
254        phi and theta ranging over +/- 90 this forms a wedge like a slice
255        of an orange.
256    azimuthal_equidistance (Postel)
257        <https://en.wikipedia.org/wiki/Azimuthal_equidistant_projection>
258        Preserves distance from center, and so is an excellent map for
259        representing a bivariate gaussian on the surface.  Theta and phi
260        operate identically, cutting wegdes from the antipode of the viewing
261        angle.  This unfortunately does not allow free movement in either
262        theta or phi since the orthogonal wobble decreases to 0 as the body
263        rotates through 180 degrees.
264    sinusoidal (Sanson-Flamsteed, Mercator equal-area)
265        <https://en.wikipedia.org/wiki/Sinusoidal_projection>
266        Preserves arc length with latitude, giving bad behaviour at
267        theta near +/- 90.  Theta and phi operate somewhat differently,
268        so a system with a-b-c dtheta-dphi-dpsi will not give the same
269        value as one with b-a-c dphi-dtheta-dpsi, as would be the case
270        for azimuthal equidistance.  Free movement using theta or phi
271        uniform over +/- 180 will work, but not as well as equirectangular
272        phi, with theta being slightly worse.  Computationally it is much
273        cheaper for wide theta-phi meshes since it excludes points which
274        lie outside the sinusoid near theta +/- 90 rather than packing
275        them close together as in equirectangle.  Note that the poles
276        will be slightly overweighted for theta > 90 with the circle
277        from theta at 90+dt winding backwards around the pole, overlapping
278        the circle from theta at 90-dt.
279    Guyou (hemisphere-in-a-square) **not weighted**
280        <https://en.wikipedia.org/wiki/Guyou_hemisphere-in-a-square_projection>
281        With tiling, allows rotation in phi or theta through +/- 180, with
282        uniform spacing.  Both theta and phi allow free rotation, with wobble
283        in the orthogonal direction reasonably well behaved (though not as
284        good as equirectangular phi). The forward/reverse transformations
285        relies on elliptic integrals that are somewhat expensive, so the
286        behaviour has to be very good to justify the cost and complexity.
287        The weighting function for each point has not yet been computed.
288        Note: run the module *guyou.py* directly and it will show the forward
289        and reverse mappings.
290    azimuthal_equal_area  **incomplete**
291        <https://en.wikipedia.org/wiki/Lambert_azimuthal_equal-area_projection>
292        Preserves the relative density of the surface patches.  Not that
293        useful and not completely implemented
294    Gauss-Kreuger **not implemented**
295        <https://en.wikipedia.org/wiki/Transverse_Mercator_projection#Ellipsoidal_transverse_Mercator>
296        Should allow free movement in theta, but phi is distorted.
297    """
298    # TODO: try Kent distribution instead of a gaussian warped by projection
299
300    t = np.linspace(-1, 1, n)
301    weights = np.ones_like(t)
302    if dist == 'gaussian':
303        t *= 3
304        weights = exp(-0.5*t**2)
305    elif dist == 'rectangle':
306        # Note: uses sasmodels ridiculous definition of rectangle width
307        t *= sqrt(3)
308    elif dist == 'uniform':
309        pass
310    else:
311        raise ValueError("expected dist to be gaussian, rectangle or uniform")
312
313    if projection == 'equirectangular':  #define PROJECTION 1
314        def _rotate(theta_i, phi_j):
315            return Rx(phi_j)*Ry(theta_i)
316        def _weight(theta_i, phi_j, w_i, w_j):
318    elif projection == 'sinusoidal':  #define PROJECTION 2
319        def _rotate(theta_i, phi_j):
320            latitude = theta_i
322            longitude = phi_j/scale if abs(phi_j) < abs(scale)*180 else 0
323            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
324            return Rx(longitude)*Ry(latitude)
325        def _weight(theta_i, phi_j, w_i, w_j):
326            latitude = theta_i
328            active = 1 if abs(phi_j) < abs(scale)*180 else 0
329            return active*w_i*w_j
330    elif projection == 'guyou':  #define PROJECTION 3  (eventually?)
331        def _rotate(theta_i, phi_j):
332            from .guyou import guyou_invert
333            #latitude, longitude = guyou_invert([theta_i], [phi_j])
334            longitude, latitude = guyou_invert([phi_j], [theta_i])
335            return Rx(longitude[0])*Ry(latitude[0])
336        def _weight(theta_i, phi_j, w_i, w_j):
337            return w_i*w_j
338    elif projection == 'azimuthal_equidistance':  # Note: Rz Ry, not Rx Ry
339        def _rotate(theta_i, phi_j):
340            latitude = sqrt(theta_i**2 + phi_j**2)
341            longitude = degrees(np.arctan2(phi_j, theta_i))
342            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
343            return Rz(longitude)*Ry(latitude)
344        def _weight(theta_i, phi_j, w_i, w_j):
345            # Weighting for each point comes from the integral:
346            #     \int\int I(q, lat, log) sin(lat) dlat dlog
347            # We are doing a conformal mapping from disk to sphere, so we need
348            # a change of variables g(theta, phi) -> (lat, long):
349            #     lat, long = sqrt(theta^2 + phi^2), arctan(phi/theta)
350            # giving:
351            #     dtheta dphi = det(J) dlat dlong
352            # where J is the jacobian from the partials of g. Using
353            #     R = sqrt(theta^2 + phi^2),
354            # then
355            #     J = [[x/R, Y/R], -y/R^2, x/R^2]]
356            # and
357            #     det(J) = 1/R
358            # with the final integral being:
359            #    \int\int I(q, theta, phi) sin(R)/R dtheta dphi
360            #
361            # This does approximately the right thing, decreasing the weight
362            # of each point as you go farther out on the disk, but it hasn't
363            # yet been checked against the 1D integral results. Prior
364            # to declaring this "good enough" and checking that integrals
365            # work in practice, we will examine alternative mappings.
366            #
367            # The issue is that the mapping does not support the case of free
368            # rotation about a single axis correctly, with a small deviation
369            # in the orthogonal axis independent of the first axis.  Like the
370            # usual polar coordiates integration, the integrated sections
371            # form wedges, though at least in this case the wedge cuts through
372            # the entire sphere, and treats theta and phi identically.
373            latitude = sqrt(theta_i**2 + phi_j**2)
374            weight = sin(radians(latitude))/latitude if latitude != 0 else 1
375            return weight*w_i*w_j if latitude < 180 else 0
376    elif projection == 'azimuthal_equal_area':
377        def _rotate(theta_i, phi_j):
378            radius = min(1, sqrt(theta_i**2 + phi_j**2)/180)
380            longitude = degrees(np.arctan2(phi_j, theta_i))
381            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
382            return Rz(longitude)*Ry(latitude)
383        def _weight(theta_i, phi_j, w_i, w_j):
384            latitude = sqrt(theta_i**2 + phi_j**2)
385            weight = sin(radians(latitude))/latitude if latitude != 0 else 1
386            return weight*w_i*w_j if latitude < 180 else 0
387    else:
388        raise ValueError("unknown projection %r"%projection)
389
390    # mesh in theta, phi formed by rotating z
391    dtheta, dphi, dpsi = jitter
392    z = np.matrix([[0], [0], [radius]])
393    points = np.hstack([_rotate(theta_i, phi_j)*z
394                        for theta_i in dtheta*t
395                        for phi_j in dphi*t])
396    w = np.array([_weight(theta_i, phi_j, w_i, w_j)
397                  for w_i, theta_i in zip(weights, dtheta*t)
398                  for w_j, phi_j in zip(weights, dphi*t)])
399    #print(max(w), min(w), min(w[w>0]))
400    points = points[:, w > 0]
401    w = w[w > 0]
402    w /= max(w)
403
404    # rotate relative to beam
405    points = orient_relative_to_beam(view, points)
406
407    x, y, z = [np.array(v).flatten() for v in points]
408    #plt.figure(2); plt.clf(); plt.hist(z, bins=np.linspace(-1, 1, 51))
409    axes.scatter(x, y, z, c=w, marker='o', vmin=0., vmax=1.)
410
411def draw_labels(axes, view, jitter, text):
412    """
413    Draw text at a particular location.
414    """
415    labels, locations, orientations = zip(*text)
416    px, py, pz = zip(*locations)
417    dx, dy, dz = zip(*orientations)
418
419    px, py, pz = transform_xyz(view, jitter, px, py, pz)
420    dx, dy, dz = transform_xyz(view, jitter, dx, dy, dz)
421
422    # TODO: zdir for labels is broken, and labels aren't appearing.
423    for label, p, zdir in zip(labels, zip(px, py, pz), zip(dx, dy, dz)):
424        zdir = np.asarray(zdir).flatten()
425        axes.text(p[0], p[1], p[2], label, zdir=zdir)
426
427# Definition of rotation matrices comes from wikipedia:
428#    https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
429def Rx(angle):
430    """Construct a matrix to rotate points about *x* by *angle* degrees."""
432    R = [[1, 0, 0],
433         [0, +cos(angle), -sin(angle)],
434         [0, +sin(angle), +cos(angle)]]
435    return np.matrix(R)
436
437def Ry(angle):
438    """Construct a matrix to rotate points about *y* by *angle* degrees."""
440    R = [[+cos(angle), 0, +sin(angle)],
441         [0, 1, 0],
442         [-sin(angle), 0, +cos(angle)]]
443    return np.matrix(R)
444
445def Rz(angle):
446    """Construct a matrix to rotate points about *z* by *angle* degrees."""
448    R = [[+cos(angle), -sin(angle), 0],
449         [+sin(angle), +cos(angle), 0],
450         [0, 0, 1]]
451    return np.matrix(R)
452
453def transform_xyz(view, jitter, x, y, z):
454    """
455    Send a set of (x,y,z) points through the jitter and view transforms.
456    """
457    x, y, z = [np.asarray(v) for v in (x, y, z)]
458    shape = x.shape
459    points = np.matrix([x.flatten(), y.flatten(), z.flatten()])
460    points = apply_jitter(jitter, points)
461    points = orient_relative_to_beam(view, points)
462    x, y, z = [np.array(v).reshape(shape) for v in points]
463    return x, y, z
464
465def apply_jitter(jitter, points):
466    """
467    Apply the jitter transform to a set of points.
468
469    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
470    """
471    dtheta, dphi, dpsi = jitter
472    points = Rx(dphi)*Ry(dtheta)*Rz(dpsi)*points
473    return points
474
475def orient_relative_to_beam(view, points):
476    """
477    Apply the view transform to a set of points.
478
479    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
480    """
481    theta, phi, psi = view
482    points = Rz(phi)*Ry(theta)*Rz(psi)*points
483    return points
484
485# translate between number of dimension of dispersity and the number of
486# points along each dimension.
487PD_N_TABLE = {
488    (0, 0, 0): (0, 0, 0),     # 0
489    (1, 0, 0): (100, 0, 0),   # 100
490    (0, 1, 0): (0, 100, 0),
491    (0, 0, 1): (0, 0, 100),
492    (1, 1, 0): (30, 30, 0),   # 900
493    (1, 0, 1): (30, 0, 30),
494    (0, 1, 1): (0, 30, 30),
495    (1, 1, 1): (15, 15, 15),  # 3375
496}
497
498def clipped_range(data, portion=1.0, mode='central'):
499    """
500    Determine range from data.
501
502    If *portion* is 1, use full range, otherwise use the center of the range
503    or the top of the range, depending on whether *mode* is 'central' or 'top'.
504    """
505    if portion == 1.0:
506        return data.min(), data.max()
507    elif mode == 'central':
508        data = np.sort(data.flatten())
509        offset = int(portion*len(data)/2 + 0.5)
510        return data[offset], data[-offset]
511    elif mode == 'top':
512        data = np.sort(data.flatten())
513        offset = int(portion*len(data) + 0.5)
514        return data[offset], data[-1]
515
516def draw_scattering(calculator, axes, view, jitter, dist='gaussian'):
517    """
518    Plot the scattering for the particular view.
519
520    *calculator* is returned from :func:build_model.  *axes* are the 3D axes
521    on which the data will be plotted.  *view* and *jitter* are the current
522    orientation and orientation dispersity.  *dist* is one of the sasmodels
523    weight distributions.
524    """
525    if dist == 'uniform':  # uniform is not yet in this branch
526        dist, scale = 'rectangle', 1/sqrt(3)
527    else:
528        scale = 1
529
530    # add the orientation parameters to the model parameters
531    theta, phi, psi = view
532    theta_pd, phi_pd, psi_pd = [scale*v for v in jitter]
533    theta_pd_n, phi_pd_n, psi_pd_n = PD_N_TABLE[(theta_pd > 0, phi_pd > 0, psi_pd > 0)]
534    ## increase pd_n for testing jitter integration rather than simple viz
535    #theta_pd_n, phi_pd_n, psi_pd_n = [5*v for v in (theta_pd_n, phi_pd_n, psi_pd_n)]
536
537    pars = dict(
538        theta=theta, theta_pd=theta_pd, theta_pd_type=dist, theta_pd_n=theta_pd_n,
539        phi=phi, phi_pd=phi_pd, phi_pd_type=dist, phi_pd_n=phi_pd_n,
540        psi=psi, psi_pd=psi_pd, psi_pd_type=dist, psi_pd_n=psi_pd_n,
541    )
542    pars.update(calculator.pars)
543
544    # compute the pattern
545    qx, qy = calculator._data.x_bins, calculator._data.y_bins
546    Iqxy = calculator(**pars).reshape(len(qx), len(qy))
547
548    # scale it and draw it
549    Iqxy = np.log(Iqxy)
550    if calculator.limits:
551        # use limits from orientation (0,0,0)
552        vmin, vmax = calculator.limits
553    else:
554        vmax = Iqxy.max()
555        vmin = vmax*10**-7
556        #vmin, vmax = clipped_range(Iqxy, portion=portion, mode='top')
557    #print("range",(vmin,vmax))
558    #qx, qy = np.meshgrid(qx, qy)
559    if 0:
560        level = np.asarray(255*(Iqxy - vmin)/(vmax - vmin), 'i')
561        level[level < 0] = 0
562        colors = plt.get_cmap()(level)
563        axes.plot_surface(qx, qy, -1.1, rstride=1, cstride=1, facecolors=colors)
564    elif 1:
565        axes.contourf(qx/qx.max(), qy/qy.max(), Iqxy, zdir='z', offset=-1.1,
566                    levels=np.linspace(vmin, vmax, 24))
567    else:
568        axes.pcolormesh(qx, qy, Iqxy)
569
570def build_model(model_name, n=150, qmax=0.5, **pars):
571    """
572    Build a calculator for the given shape.
573
574    *model_name* is any sasmodels model.  *n* and *qmax* define an n x n mesh
575    on which to evaluate the model.  The remaining parameters are stored in
576    the returned calculator as *calculator.pars*.  They are used by
577    :func:draw_scattering to set the non-orientation parameters in the
578    calculation.
579
580    Returns a *calculator* function which takes a dictionary or parameters and
581    produces Iqxy.  The Iqxy value needs to be reshaped to an n x n matrix
582    for plotting.  See the :class:sasmodels.direct_model.DirectModel class
583    for details.
584    """
585    from sasmodels.core import load_model_info, build_model
586    from sasmodels.data import empty_data2D
587    from sasmodels.direct_model import DirectModel
588
590    model = build_model(model_info) #, dtype='double!')
591    q = np.linspace(-qmax, qmax, n)
592    data = empty_data2D(q, q)
593    calculator = DirectModel(data, model)
594
595    # stuff the values for non-orientation parameters into the calculator
596    calculator.pars = pars.copy()
597    calculator.pars.setdefault('backgound', 1e-3)
598
599    # fix the data limits so that we can see if the pattern fades
600    # under rotation or angular dispersion
601    Iqxy = calculator(theta=0, phi=0, psi=0, **calculator.pars)
602    Iqxy = np.log(Iqxy)
603    vmin, vmax = clipped_range(Iqxy, 0.95, mode='top')
604    calculator.limits = vmin, vmax+1
605
606    return calculator
607
608def select_calculator(model_name, n=150, size=(10, 40, 100)):
609    """
610    Create a model calculator for the given shape.
611
612    *model_name* is one of sphere, cylinder, ellipsoid, triaxial_ellipsoid,
613    parallelepiped or bcc_paracrystal. *n* is the number of points to use
614    in the q range.  *qmax* is chosen based on model parameters for the
615    given model to show something intersting.
616
617    Returns *calculator* and tuple *size* (a,b,c) giving minor and major
618    equitorial axes and polar axis respectively.  See :func:build_model
619    for details on the returned calculator.
620    """
621    a, b, c = size
622    d_factor = 0.06  # for paracrystal models
623    if model_name == 'sphere':
624        calculator = build_model('sphere', n=n, radius=c)
625        a = b = c
626    elif model_name == 'sc_paracrystal':
627        a = b = c
628        dnn = c
630        calculator = build_model('sc_paracrystal', n=n, dnn=dnn,
632                                 background=0)
633    elif model_name == 'fcc_paracrystal':
634        a = b = c
635        # nearest neigbour distance dnn should be 2 radius, but I think the
636        # model uses lattice spacing rather than dnn in its calculations
637        dnn = 0.5*c
638        radius = sqrt(2)/4 * c
639        calculator = build_model('fcc_paracrystal', n=n, dnn=dnn,
641                                 background=0)
642    elif model_name == 'bcc_paracrystal':
643        a = b = c
644        # nearest neigbour distance dnn should be 2 radius, but I think the
645        # model uses lattice spacing rather than dnn in its calculations
646        dnn = 0.5*c
647        radius = sqrt(3)/2 * c
648        calculator = build_model('bcc_paracrystal', n=n, dnn=dnn,
650                                 background=0)
651    elif model_name == 'cylinder':
652        calculator = build_model('cylinder', n=n, qmax=0.3, radius=b, length=c)
653        a = b
654    elif model_name == 'ellipsoid':
655        calculator = build_model('ellipsoid', n=n, qmax=1.0,
657        a = b
658    elif model_name == 'triaxial_ellipsoid':
659        calculator = build_model('triaxial_ellipsoid', n=n, qmax=0.5,
663    elif model_name == 'parallelepiped':
664        calculator = build_model('parallelepiped', n=n, a=a, b=b, c=c)
665    else:
666        raise ValueError("unknown model %s"%model_name)
667
668    return calculator, (a, b, c)
669
670SHAPES = [
671    'parallelepiped',
672    'sphere', 'ellipsoid', 'triaxial_ellipsoid',
673    'cylinder',
674    'fcc_paracrystal', 'bcc_paracrystal', 'sc_paracrystal',
675]
676
677DRAW_SHAPES = {
678    'fcc_paracrystal': draw_fcc,
679    'bcc_paracrystal': draw_bcc,
680    'sc_paracrystal': draw_sc,
681    'parallelepiped': draw_parallelepiped,
682}
683
684DISTRIBUTIONS = [
685    'gaussian', 'rectangle', 'uniform',
686]
687DIST_LIMITS = {
688    'gaussian': 30,
689    'rectangle': 90/sqrt(3),
690    'uniform': 90,
691}
692
693def run(model_name='parallelepiped', size=(10, 40, 100),
694        dist='gaussian', mesh=30,
695        projection='equirectangular'):
696    """
697    Show an interactive orientation and jitter demo.
698
699    *model_name* is one of: sphere, ellipsoid, triaxial_ellipsoid,
700    parallelepiped, cylinder, or sc/fcc/bcc_paracrystal
701
702    *size* gives the dimensions (a, b, c) of the shape.
703
704    *dist* is the type of dispersition: gaussian, rectangle, or uniform.
705
706    *mesh* is the number of points in the dispersion mesh.
707
708    *projection* is the map projection to use for the mesh: equirectangular,
709    sinusoidal, guyou, azimuthal_equidistance, or azimuthal_equal_area.
710    """
711    # projection number according to 1-order position in list, but
712    # only 1 and 2 are implemented so far.
713    from sasmodels import generate
714    generate.PROJECTION = PROJECTIONS.index(projection) + 1
715    if generate.PROJECTION > 2:
716        print("*** PROJECTION %s not implemented in scattering function ***"%projection)
717        generate.PROJECTION = 2
718
719    # set up calculator
720    calculator, size = select_calculator(model_name, n=150, size=size)
721    draw_shape = DRAW_SHAPES.get(model_name, draw_parallelepiped)
722
723    ## uncomment to set an independent the colour range for every view
724    ## If left commented, the colour range is fixed for all views
725    calculator.limits = None
726
727    ## initial view
728    #theta, dtheta = 70., 10.
729    #phi, dphi = -45., 3.
730    #psi, dpsi = -45., 3.
731    theta, phi, psi = 0, 0, 0
732    dtheta, dphi, dpsi = 0, 0, 0
733
734    ## create the plot window
735    #plt.hold(True)
736    plt.subplots(num=None, figsize=(5.5, 5.5))
737    plt.set_cmap('gist_earth')
738    plt.clf()
739    plt.gcf().canvas.set_window_title(projection)
740    #gs = gridspec.GridSpec(2,1,height_ratios=[4,1])
741    #axes = plt.subplot(gs[0], projection='3d')
742    axes = plt.axes([0.0, 0.2, 1.0, 0.8], projection='3d')
743    try:  # CRUFT: not all versions of matplotlib accept 'square' 3d projection
744        axes.axis('square')
745    except Exception:
746        pass
747
748    axcolor = 'lightgoldenrodyellow'
749
750    ## add control widgets to plot
751    axes_theta = plt.axes([0.1, 0.15, 0.45, 0.04], axisbg=axcolor)
752    axes_phi = plt.axes([0.1, 0.1, 0.45, 0.04], axisbg=axcolor)
753    axes_psi = plt.axes([0.1, 0.05, 0.45, 0.04], axisbg=axcolor)
754    stheta = Slider(axes_theta, 'Theta', -90, 90, valinit=theta)
755    sphi = Slider(axes_phi, 'Phi', -180, 180, valinit=phi)
756    spsi = Slider(axes_psi, 'Psi', -180, 180, valinit=psi)
757
758    axes_dtheta = plt.axes([0.75, 0.15, 0.15, 0.04], axisbg=axcolor)
759    axes_dphi = plt.axes([0.75, 0.1, 0.15, 0.04], axisbg=axcolor)
760    axes_dpsi = plt.axes([0.75, 0.05, 0.15, 0.04], axisbg=axcolor)
761    # Note: using ridiculous definition of rectangle distribution, whose width
762    # in sasmodels is sqrt(3) times the given width.  Divide by sqrt(3) to keep
763    # the maximum width to 90.
764    dlimit = DIST_LIMITS[dist]
765    sdtheta = Slider(axes_dtheta, 'dTheta', 0, 2*dlimit, valinit=dtheta)
766    sdphi = Slider(axes_dphi, 'dPhi', 0, 2*dlimit, valinit=dphi)
767    sdpsi = Slider(axes_dpsi, 'dPsi', 0, 2*dlimit, valinit=dpsi)
768
769
770    ## callback to draw the new view
771    def update(val, axis=None):
772        view = stheta.val, sphi.val, spsi.val
773        jitter = sdtheta.val, sdphi.val, sdpsi.val
774        # set small jitter as 0 if multiple pd dims
775        dims = sum(v > 0 for v in jitter)
776        limit = [0, 0, 0.5, 5][dims]
777        jitter = [0 if v < limit else v for v in jitter]
778        axes.cla()
779        draw_beam(axes, (0, 0))
780        draw_jitter(axes, view, jitter, dist=dist, size=size, draw_shape=draw_shape)
781        #draw_jitter(axes, view, (0,0,0))
782        draw_mesh(axes, view, jitter, dist=dist, n=mesh, projection=projection)
783        draw_scattering(calculator, axes, view, jitter, dist=dist)
784        plt.gcf().canvas.draw()
785
786    ## bind control widgets to view updater
787    stheta.on_changed(lambda v: update(v, 'theta'))
788    sphi.on_changed(lambda v: update(v, 'phi'))
789    spsi.on_changed(lambda v: update(v, 'psi'))
790    sdtheta.on_changed(lambda v: update(v, 'dtheta'))
791    sdphi.on_changed(lambda v: update(v, 'dphi'))
792    sdpsi.on_changed(lambda v: update(v, 'dpsi'))
793
794    ## initialize view
795    update(None, 'phi')
796
797    ## go interactive
798    plt.show()
799
800def main():
801    parser = argparse.ArgumentParser(
802        description="Display jitter",
803        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
804        )
806                        default=PROJECTIONS[0],
807                        help='coordinate projection')
809                        help='a,b,c lengths')
811                        default=DISTRIBUTIONS[0],
812                        help='jitter distribution')