source: sasmodels/sasmodels/jitter.py @ f4ae8c4

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

jitter: move tracking dot for paracrystal to 0,0,1

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