source: sasmodels/sasmodels/jitter.py @ 7d97437

core_shell_microgelsmagnetic_modelticket-1257-vesicle-productticket_1156ticket_1265_superballticket_822_more_unit_tests
Last change on this file since 7d97437 was 7d97437, checked in by Paul Kienzle <pkienzle@…>, 8 months ago

support for newer matplotlib

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