source: sasmodels/sasmodels/jitter.py @ d86f0fc

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

lint reduction

  • Property mode set to 100755
File size: 32.1 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, CheckButtons
19from matplotlib import cm
20import numpy as np
21from numpy import pi, cos, sin, sqrt, exp, degrees, radians
22
23def draw_beam(ax, view=(0, 0)):
24    """
25    Draw the beam going from source at (0, 0, 1) to detector at (0, 0, -1)
26    """
27    #ax.plot([0,0],[0,0],[1,-1])
28    #ax.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    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='y', alpha=0.5)
46
47def draw_ellipsoid(ax, 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    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w', alpha=alpha)
58
59    draw_labels(ax, 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(ax, size, view, jitter, steps=None, alpha=1):
69    """Draw points for simple cubic paracrystal"""
70    atoms = _build_sc()
71    _draw_crystal(ax, size, view, jitter, atoms=atoms)
72
73def draw_fcc(ax, 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(ax, size, view, jitter, atoms=atoms)
87
88def draw_bcc(ax, 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(ax, size, view, jitter, atoms=atoms)
101
102def _draw_crystal(ax, size, view, jitter, steps=None, alpha=1, 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    ax.scatter([x[0]], [y[0]], [z[0]], c='yellow', marker='o')
107    ax.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(ax, 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    ax.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        ax.plot_trisurf(x, y, triangles=tri, Z=z, color=[1, 0.6, 0.6], alpha=alpha)
156
157    draw_labels(ax, 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(ax, 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    ax.plot_surface(x, y, z, rstride=4, cstride=4, color='w')
175
176def draw_jitter(ax, 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(ax, 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(ax, 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(ax, 'set_'+v+'lim')([-lim, lim])
232        getattr(ax, 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(ax, 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    t = np.linspace(-1, 1, n)
300    weights = np.ones_like(t)
301    if dist == 'gaussian':
302        t *= 3
303        weights = exp(-0.5*t**2)
304    elif dist == 'rectangle':
305        # Note: uses sasmodels ridiculous definition of rectangle width
306        t *= sqrt(3)
307    elif dist == 'uniform':
308        pass
309    else:
310        raise ValueError("expected dist to be gaussian, rectangle or uniform")
311
312    if projection == 'equirectangular':  #define PROJECTION 1
313        def _rotate(theta_i, phi_j):
314            return Rx(phi_j)*Ry(theta_i)
315        def _weight(theta_i, phi_j, wi, wj):
316            return wi*wj*abs(cos(radians(theta_i)))
317    elif projection == 'sinusoidal':  #define PROJECTION 2
318        def _rotate(theta_i, phi_j):
319            latitude = theta_i
320            scale = cos(radians(latitude))
321            longitude = phi_j/scale if abs(phi_j) < abs(scale)*180 else 0
322            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
323            return Rx(longitude)*Ry(latitude)
324        def _weight(theta_i, phi_j, wi, wj):
325            latitude = theta_i
326            scale = cos(radians(latitude))
327            w = 1 if abs(phi_j) < abs(scale)*180 else 0
328            return w*wi*wj
329    elif projection == 'guyou':  #define PROJECTION 3  (eventually?)
330        def _rotate(theta_i, phi_j):
331            from guyou import guyou_invert
332            #latitude, longitude = guyou_invert([theta_i], [phi_j])
333            longitude, latitude = guyou_invert([phi_j], [theta_i])
334            return Rx(longitude[0])*Ry(latitude[0])
335        def _weight(theta_i, phi_j, wi, wj):
336            return wi*wj
337    elif projection == 'azimuthal_equidistance':  # Note: Rz Ry, not Rx Ry
338        def _rotate(theta_i, phi_j):
339            latitude = sqrt(theta_i**2 + phi_j**2)
340            longitude = degrees(np.arctan2(phi_j, theta_i))
341            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
342            return Rz(longitude)*Ry(latitude)
343        def _weight(theta_i, phi_j, wi, wj):
344            # Weighting for each point comes from the integral:
345            #     \int\int I(q, lat, log) sin(lat) dlat dlog
346            # We are doing a conformal mapping from disk to sphere, so we need
347            # a change of variables g(theta, phi) -> (lat, long):
348            #     lat, long = sqrt(theta^2 + phi^2), arctan(phi/theta)
349            # giving:
350            #     dtheta dphi = det(J) dlat dlong
351            # where J is the jacobian from the partials of g. Using
352            #     R = sqrt(theta^2 + phi^2),
353            # then
354            #     J = [[x/R, Y/R], -y/R^2, x/R^2]]
355            # and
356            #     det(J) = 1/R
357            # with the final integral being:
358            #    \int\int I(q, theta, phi) sin(R)/R dtheta dphi
359            #
360            # This does approximately the right thing, decreasing the weight
361            # of each point as you go farther out on the disk, but it hasn't
362            # yet been checked against the 1D integral results. Prior
363            # to declaring this "good enough" and checking that integrals
364            # work in practice, we will examine alternative mappings.
365            #
366            # The issue is that the mapping does not support the case of free
367            # rotation about a single axis correctly, with a small deviation
368            # in the orthogonal axis independent of the first axis.  Like the
369            # usual polar coordiates integration, the integrated sections
370            # form wedges, though at least in this case the wedge cuts through
371            # the entire sphere, and treats theta and phi identically.
372            latitude = sqrt(theta_i**2 + phi_j**2)
373            w = sin(radians(latitude))/latitude if latitude != 0 else 1
374            return w*wi*wj if latitude < 180 else 0
375    elif projection == 'azimuthal_equal_area':
376        def _rotate(theta_i, phi_j):
377            R = min(1, sqrt(theta_i**2 + phi_j**2)/180)
378            latitude = 180-degrees(2*np.arccos(R))
379            longitude = degrees(np.arctan2(phi_j, theta_i))
380            #print("(%+7.2f, %+7.2f) => (%+7.2f, %+7.2f)"%(theta_i, phi_j, latitude, longitude))
381            return Rz(longitude)*Ry(latitude)
382        def _weight(theta_i, phi_j, wi, wj):
383            latitude = sqrt(theta_i**2 + phi_j**2)
384            w = sin(radians(latitude))/latitude if latitude != 0 else 1
385            return w*wi*wj if latitude < 180 else 0
386    else:
387        raise ValueError("unknown projection %r"%projection)
388
389    # mesh in theta, phi formed by rotating z
390    dtheta, dphi, dpsi = jitter
391    z = np.matrix([[0], [0], [radius]])
392    points = np.hstack([_rotate(theta_i, phi_j)*z
393                        for theta_i in dtheta*t
394                        for phi_j in dphi*t])
395    # select just the active points (i.e., those with phi < 180
396    w = np.array([_weight(theta_i, phi_j, wi, wj)
397                  for wi, theta_i in zip(weights, dtheta*t)
398                  for wj, 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    if 0: # Kent distribution
405        points = np.hstack([Rx(phi_j)*Ry(theta_i)*z for theta_i in 30*t for phi_j in 60*t])
406        xp, yp, zp = [np.array(v).flatten() for v in points]
407        kappa = max(1e6, radians(dtheta)/(2*pi))
408        beta = 1/max(1e-6, radians(dphi)/(2*pi))/kappa
409        w = exp(kappa*zp) #+ beta*(xp**2 + yp**2)
410        print(kappa, dtheta, radians(dtheta), min(w), max(w), sum(w))
411        #w /= abs(cos(radians(
412        #w /= sum(w)
413
414    # rotate relative to beam
415    points = orient_relative_to_beam(view, points)
416
417    x, y, z = [np.array(v).flatten() for v in points]
418    #plt.figure(2); plt.clf(); plt.hist(z, bins=np.linspace(-1, 1, 51))
419    ax.scatter(x, y, z, c=w, marker='o', vmin=0., vmax=1.)
420
421def draw_labels(ax, view, jitter, text):
422    """
423    Draw text at a particular location.
424    """
425    labels, locations, orientations = zip(*text)
426    px, py, pz = zip(*locations)
427    dx, dy, dz = zip(*orientations)
428
429    px, py, pz = transform_xyz(view, jitter, px, py, pz)
430    dx, dy, dz = transform_xyz(view, jitter, dx, dy, dz)
431
432    # TODO: zdir for labels is broken, and labels aren't appearing.
433    for label, p, zdir in zip(labels, zip(px, py, pz), zip(dx, dy, dz)):
434        zdir = np.asarray(zdir).flatten()
435        ax.text(p[0], p[1], p[2], label, zdir=zdir)
436
437# Definition of rotation matrices comes from wikipedia:
438#    https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
439def Rx(angle):
440    """Construct a matrix to rotate points about *x* by *angle* degrees."""
441    a = radians(angle)
442    R = [[1, 0, 0],
443         [0, +cos(a), -sin(a)],
444         [0, +sin(a), +cos(a)]]
445    return np.matrix(R)
446
447def Ry(angle):
448    """Construct a matrix to rotate points about *y* by *angle* degrees."""
449    a = radians(angle)
450    R = [[+cos(a), 0, +sin(a)],
451         [0, 1, 0],
452         [-sin(a), 0, +cos(a)]]
453    return np.matrix(R)
454
455def Rz(angle):
456    """Construct a matrix to rotate points about *z* by *angle* degrees."""
457    a = radians(angle)
458    R = [[+cos(a), -sin(a), 0],
459         [+sin(a), +cos(a), 0],
460         [0, 0, 1]]
461    return np.matrix(R)
462
463def transform_xyz(view, jitter, x, y, z):
464    """
465    Send a set of (x,y,z) points through the jitter and view transforms.
466    """
467    x, y, z = [np.asarray(v) for v in (x, y, z)]
468    shape = x.shape
469    points = np.matrix([x.flatten(), y.flatten(), z.flatten()])
470    points = apply_jitter(jitter, points)
471    points = orient_relative_to_beam(view, points)
472    x, y, z = [np.array(v).reshape(shape) for v in points]
473    return x, y, z
474
475def apply_jitter(jitter, points):
476    """
477    Apply the jitter 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    dtheta, dphi, dpsi = jitter
482    points = Rx(dphi)*Ry(dtheta)*Rz(dpsi)*points
483    return points
484
485def orient_relative_to_beam(view, points):
486    """
487    Apply the view transform to a set of points.
488
489    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
490    """
491    theta, phi, psi = view
492    points = Rz(phi)*Ry(theta)*Rz(psi)*points
493    return points
494
495# translate between number of dimension of dispersity and the number of
496# points along each dimension.
497PD_N_TABLE = {
498    (0, 0, 0): (0, 0, 0),     # 0
499    (1, 0, 0): (100, 0, 0),   # 100
500    (0, 1, 0): (0, 100, 0),
501    (0, 0, 1): (0, 0, 100),
502    (1, 1, 0): (30, 30, 0),   # 900
503    (1, 0, 1): (30, 0, 30),
504    (0, 1, 1): (0, 30, 30),
505    (1, 1, 1): (15, 15, 15),  # 3375
506}
507
508def clipped_range(data, portion=1.0, mode='central'):
509    """
510    Determine range from data.
511
512    If *portion* is 1, use full range, otherwise use the center of the range
513    or the top of the range, depending on whether *mode* is 'central' or 'top'.
514    """
515    if portion == 1.0:
516        return data.min(), data.max()
517    elif mode == 'central':
518        data = np.sort(data.flatten())
519        offset = int(portion*len(data)/2 + 0.5)
520        return data[offset], data[-offset]
521    elif mode == 'top':
522        data = np.sort(data.flatten())
523        offset = int(portion*len(data) + 0.5)
524        return data[offset], data[-1]
525
526def draw_scattering(calculator, ax, view, jitter, dist='gaussian'):
527    """
528    Plot the scattering for the particular view.
529
530    *calculator* is returned from :func:`build_model`.  *ax* are the 3D axes
531    on which the data will be plotted.  *view* and *jitter* are the current
532    orientation and orientation dispersity.  *dist* is one of the sasmodels
533    weight distributions.
534    """
535    if dist == 'uniform':  # uniform is not yet in this branch
536        dist, scale = 'rectangle', 1/sqrt(3)
537    else:
538        scale = 1
539
540    # add the orientation parameters to the model parameters
541    theta, phi, psi = view
542    theta_pd, phi_pd, psi_pd = [scale*v for v in jitter]
543    theta_pd_n, phi_pd_n, psi_pd_n = PD_N_TABLE[(theta_pd > 0, phi_pd > 0, psi_pd > 0)]
544    ## increase pd_n for testing jitter integration rather than simple viz
545    #theta_pd_n, phi_pd_n, psi_pd_n = [5*v for v in (theta_pd_n, phi_pd_n, psi_pd_n)]
546
547    pars = dict(
548        theta=theta, theta_pd=theta_pd, theta_pd_type=dist, theta_pd_n=theta_pd_n,
549        phi=phi, phi_pd=phi_pd, phi_pd_type=dist, phi_pd_n=phi_pd_n,
550        psi=psi, psi_pd=psi_pd, psi_pd_type=dist, psi_pd_n=psi_pd_n,
551    )
552    pars.update(calculator.pars)
553
554    # compute the pattern
555    qx, qy = calculator._data.x_bins, calculator._data.y_bins
556    Iqxy = calculator(**pars).reshape(len(qx), len(qy))
557
558    # scale it and draw it
559    Iqxy = np.log(Iqxy)
560    if calculator.limits:
561        # use limits from orientation (0,0,0)
562        vmin, vmax = calculator.limits
563    else:
564        vmax = Iqxy.max()
565        vmin = vmax*10**-7
566        #vmin, vmax = clipped_range(Iqxy, portion=portion, mode='top')
567    #print("range",(vmin,vmax))
568    #qx, qy = np.meshgrid(qx, qy)
569    if 0:
570        level = np.asarray(255*(Iqxy - vmin)/(vmax - vmin), 'i')
571        level[level < 0] = 0
572        colors = plt.get_cmap()(level)
573        ax.plot_surface(qx, qy, -1.1, rstride=1, cstride=1, facecolors=colors)
574    elif 1:
575        ax.contourf(qx/qx.max(), qy/qy.max(), Iqxy, zdir='z', offset=-1.1,
576                    levels=np.linspace(vmin, vmax, 24))
577    else:
578        ax.pcolormesh(qx, qy, Iqxy)
579
580def build_model(model_name, n=150, qmax=0.5, **pars):
581    """
582    Build a calculator for the given shape.
583
584    *model_name* is any sasmodels model.  *n* and *qmax* define an n x n mesh
585    on which to evaluate the model.  The remaining parameters are stored in
586    the returned calculator as *calculator.pars*.  They are used by
587    :func:`draw_scattering` to set the non-orientation parameters in the
588    calculation.
589
590    Returns a *calculator* function which takes a dictionary or parameters and
591    produces Iqxy.  The Iqxy value needs to be reshaped to an n x n matrix
592    for plotting.  See the :class:`sasmodels.direct_model.DirectModel` class
593    for details.
594    """
595    from sasmodels.core import load_model_info, build_model
596    from sasmodels.data import empty_data2D
597    from sasmodels.direct_model import DirectModel
598
599    model_info = load_model_info(model_name)
600    model = build_model(model_info) #, dtype='double!')
601    q = np.linspace(-qmax, qmax, n)
602    data = empty_data2D(q, q)
603    calculator = DirectModel(data, model)
604
605    # stuff the values for non-orientation parameters into the calculator
606    calculator.pars = pars.copy()
607    calculator.pars.setdefault('backgound', 1e-3)
608
609    # fix the data limits so that we can see if the pattern fades
610    # under rotation or angular dispersion
611    Iqxy = calculator(theta=0, phi=0, psi=0, **calculator.pars)
612    Iqxy = np.log(Iqxy)
613    vmin, vmax = clipped_range(Iqxy, 0.95, mode='top')
614    calculator.limits = vmin, vmax+1
615
616    return calculator
617
618def select_calculator(model_name, n=150, size=(10, 40, 100)):
619    """
620    Create a model calculator for the given shape.
621
622    *model_name* is one of sphere, cylinder, ellipsoid, triaxial_ellipsoid,
623    parallelepiped or bcc_paracrystal. *n* is the number of points to use
624    in the q range.  *qmax* is chosen based on model parameters for the
625    given model to show something intersting.
626
627    Returns *calculator* and tuple *size* (a,b,c) giving minor and major
628    equitorial axes and polar axis respectively.  See :func:`build_model`
629    for details on the returned calculator.
630    """
631    a, b, c = size
632    d_factor = 0.06  # for paracrystal models
633    if model_name == 'sphere':
634        calculator = build_model('sphere', n=n, radius=c)
635        a = b = c
636    elif model_name == 'sc_paracrystal':
637        a = b = c
638        dnn = c
639        radius = 0.5*c
640        calculator = build_model('sc_paracrystal', n=n, dnn=dnn,
641                                 d_factor=d_factor, radius=(1-d_factor)*radius,
642                                 background=0)
643    elif model_name == 'fcc_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(2)/4 * c
649        calculator = build_model('fcc_paracrystal', n=n, dnn=dnn,
650                                 d_factor=d_factor, radius=(1-d_factor)*radius,
651                                 background=0)
652    elif model_name == 'bcc_paracrystal':
653        a = b = c
654        # nearest neigbour distance dnn should be 2 radius, but I think the
655        # model uses lattice spacing rather than dnn in its calculations
656        dnn = 0.5*c
657        radius = sqrt(3)/2 * c
658        calculator = build_model('bcc_paracrystal', n=n, dnn=dnn,
659                                 d_factor=d_factor, radius=(1-d_factor)*radius,
660                                 background=0)
661    elif model_name == 'cylinder':
662        calculator = build_model('cylinder', n=n, qmax=0.3, radius=b, length=c)
663        a = b
664    elif model_name == 'ellipsoid':
665        calculator = build_model('ellipsoid', n=n, qmax=1.0,
666                                 radius_polar=c, radius_equatorial=b)
667        a = b
668    elif model_name == 'triaxial_ellipsoid':
669        calculator = build_model('triaxial_ellipsoid', n=n, qmax=0.5,
670                                 radius_equat_minor=a,
671                                 radius_equat_major=b,
672                                 radius_polar=c)
673    elif model_name == 'parallelepiped':
674        calculator = build_model('parallelepiped', n=n, a=a, b=b, c=c)
675    else:
676        raise ValueError("unknown model %s"%model_name)
677
678    return calculator, (a, b, c)
679
680SHAPES = [
681    'parallelepiped',
682    'sphere', 'ellipsoid', 'triaxial_ellipsoid',
683    'cylinder',
684    'fcc_paracrystal', 'bcc_paracrystal', 'sc_paracrystal',
685]
686
687DRAW_SHAPES = {
688    'fcc_paracrystal': draw_fcc,
689    'bcc_paracrystal': draw_bcc,
690    'sc_paracrystal': draw_sc,
691    'parallelepiped': draw_parallelepiped,
692}
693
694DISTRIBUTIONS = [
695    'gaussian', 'rectangle', 'uniform',
696]
697DIST_LIMITS = {
698    'gaussian': 30,
699    'rectangle': 90/sqrt(3),
700    'uniform': 90,
701}
702
703def run(model_name='parallelepiped', size=(10, 40, 100),
704        dist='gaussian', mesh=30,
705        projection='equirectangular'):
706    """
707    Show an interactive orientation and jitter demo.
708
709    *model_name* is one of: sphere, ellipsoid, triaxial_ellipsoid,
710    parallelepiped, cylinder, or sc/fcc/bcc_paracrystal
711
712    *size* gives the dimensions (a, b, c) of the shape.
713
714    *dist* is the type of dispersition: gaussian, rectangle, or uniform.
715
716    *mesh* is the number of points in the dispersion mesh.
717
718    *projection* is the map projection to use for the mesh: equirectangular,
719    sinusoidal, guyou, azimuthal_equidistance, or azimuthal_equal_area.
720    """
721    # projection number according to 1-order position in list, but
722    # only 1 and 2 are implemented so far.
723    from sasmodels import generate
724    generate.PROJECTION = PROJECTIONS.index(projection) + 1
725    if generate.PROJECTION > 2:
726        print("*** PROJECTION %s not implemented in scattering function ***"%projection)
727        generate.PROJECTION = 2
728
729    # set up calculator
730    calculator, size = select_calculator(model_name, n=150, size=size)
731    draw_shape = DRAW_SHAPES.get(model_name, draw_parallelepiped)
732
733    ## uncomment to set an independent the colour range for every view
734    ## If left commented, the colour range is fixed for all views
735    calculator.limits = None
736
737    ## initial view
738    #theta, dtheta = 70., 10.
739    #phi, dphi = -45., 3.
740    #psi, dpsi = -45., 3.
741    theta, phi, psi = 0, 0, 0
742    dtheta, dphi, dpsi = 0, 0, 0
743
744    ## create the plot window
745    #plt.hold(True)
746    plt.subplots(num=None, figsize=(5.5, 5.5))
747    plt.set_cmap('gist_earth')
748    plt.clf()
749    plt.gcf().canvas.set_window_title(projection)
750    #gs = gridspec.GridSpec(2,1,height_ratios=[4,1])
751    #ax = plt.subplot(gs[0], projection='3d')
752    ax = plt.axes([0.0, 0.2, 1.0, 0.8], projection='3d')
753    try:  # CRUFT: not all versions of matplotlib accept 'square' 3d projection
754        ax.axis('square')
755    except Exception:
756        pass
757
758    axcolor = 'lightgoldenrodyellow'
759
760    ## add control widgets to plot
761    axtheta = plt.axes([0.1, 0.15, 0.45, 0.04], axisbg=axcolor)
762    axphi = plt.axes([0.1, 0.1, 0.45, 0.04], axisbg=axcolor)
763    axpsi = plt.axes([0.1, 0.05, 0.45, 0.04], axisbg=axcolor)
764    stheta = Slider(axtheta, 'Theta', -90, 90, valinit=theta)
765    sphi = Slider(axphi, 'Phi', -180, 180, valinit=phi)
766    spsi = Slider(axpsi, 'Psi', -180, 180, valinit=psi)
767
768    axdtheta = plt.axes([0.75, 0.15, 0.15, 0.04], axisbg=axcolor)
769    axdphi = plt.axes([0.75, 0.1, 0.15, 0.04], axisbg=axcolor)
770    axdpsi = plt.axes([0.75, 0.05, 0.15, 0.04], axisbg=axcolor)
771    # Note: using ridiculous definition of rectangle distribution, whose width
772    # in sasmodels is sqrt(3) times the given width.  Divide by sqrt(3) to keep
773    # the maximum width to 90.
774    dlimit = DIST_LIMITS[dist]
775    sdtheta = Slider(axdtheta, 'dTheta', 0, 2*dlimit, valinit=dtheta)
776    sdphi = Slider(axdphi, 'dPhi', 0, 2*dlimit, valinit=dphi)
777    sdpsi = Slider(axdpsi, 'dPsi', 0, 2*dlimit, valinit=dpsi)
778
779
780    ## callback to draw the new view
781    def update(val, axis=None):
782        view = stheta.val, sphi.val, spsi.val
783        jitter = sdtheta.val, sdphi.val, sdpsi.val
784        # set small jitter as 0 if multiple pd dims
785        dims = sum(v > 0 for v in jitter)
786        limit = [0, 0, 0.5, 5][dims]
787        jitter = [0 if v < limit else v for v in jitter]
788        ax.cla()
789        draw_beam(ax, (0, 0))
790        draw_jitter(ax, view, jitter, dist=dist, size=size, draw_shape=draw_shape)
791        #draw_jitter(ax, view, (0,0,0))
792        draw_mesh(ax, view, jitter, dist=dist, n=mesh, projection=projection)
793        draw_scattering(calculator, ax, view, jitter, dist=dist)
794        plt.gcf().canvas.draw()
795
796    ## bind control widgets to view updater
797    stheta.on_changed(lambda v: update(v, 'theta'))
798    sphi.on_changed(lambda v: update(v, 'phi'))
799    spsi.on_changed(lambda v: update(v, 'psi'))
800    sdtheta.on_changed(lambda v: update(v, 'dtheta'))
801    sdphi.on_changed(lambda v: update(v, 'dphi'))
802    sdpsi.on_changed(lambda v: update(v, 'dpsi'))
803
804    ## initialize view
805    update(None, 'phi')
806
807    ## go interactive
808    plt.show()
809
810def main():
811    parser = argparse.ArgumentParser(
812        description="Display jitter",
813        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
814        )
815    parser.add_argument('-p', '--projection', choices=PROJECTIONS,
816                        default=PROJECTIONS[0],
817                        help='coordinate projection')
818    parser.add_argument('-s', '--size', type=str, default='10,40,100',
819                        help='a,b,c lengths')
820    parser.add_argument('-d', '--distribution', choices=DISTRIBUTIONS,
821                        default=DISTRIBUTIONS[0],
822                        help='jitter distribution')
823    parser.add_argument('-m', '--mesh', type=int, default=30,
824                        help='#points in theta-phi mesh')
825    parser.add_argument('shape', choices=SHAPES, nargs='?', default=SHAPES[0],
826                        help='oriented shape')
827    opts = parser.parse_args()
828    size = tuple(int(v) for v in opts.size.split(','))
829    run(opts.shape, size=size,
830        mesh=opts.mesh, dist=opts.distribution,
831        projection=opts.projection)
832
833if __name__ == "__main__":
834    main()
Note: See TracBrowser for help on using the repository browser.