source: sasmodels/explore/realspace.py @ 2a39ca4

ticket_1156
Last change on this file since 2a39ca4 was 2a39ca4, checked in by Paul Kienzle <pkienzle@…>, 6 years ago

sample from sc/bcc/fcc lattice in real space calculation

  • Property mode set to 100644
File size: 36.0 KB
Line 
1from __future__ import division, print_function
2
3import time
4from copy import copy
5import os
6import argparse
7from collections import OrderedDict
8
9import numpy as np
10from numpy import pi, radians, sin, cos, sqrt
11from numpy.random import poisson, uniform, randn, rand
12from numpy.polynomial.legendre import leggauss
13from scipy.integrate import simps
14from scipy.special import j1 as J1
15
16try:
17    import numba
18    USE_NUMBA = True
19except ImportError:
20    USE_NUMBA = False
21
22# Definition of rotation matrices comes from wikipedia:
23#    https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
24def Rx(angle):
25    """Construct a matrix to rotate points about *x* by *angle* degrees."""
26    a = radians(angle)
27    R = [[1, 0, 0],
28         [0, +cos(a), -sin(a)],
29         [0, +sin(a), +cos(a)]]
30    return np.matrix(R)
31
32def Ry(angle):
33    """Construct a matrix to rotate points about *y* by *angle* degrees."""
34    a = radians(angle)
35    R = [[+cos(a), 0, +sin(a)],
36         [0, 1, 0],
37         [-sin(a), 0, +cos(a)]]
38    return np.matrix(R)
39
40def Rz(angle):
41    """Construct a matrix to rotate points about *z* by *angle* degrees."""
42    a = radians(angle)
43    R = [[+cos(a), -sin(a), 0],
44         [+sin(a), +cos(a), 0],
45         [0, 0, 1]]
46    return np.matrix(R)
47
48def rotation(theta, phi, psi):
49    """
50    Apply the jitter transform to a set of points.
51
52    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
53    """
54    return Rx(phi)*Ry(theta)*Rz(psi)
55
56def apply_view(points, view):
57    """
58    Apply the view transform (theta, phi, psi) to a set of points.
59
60    Points are stored in a 3 x n numpy array.
61
62    View angles are in degrees.
63    """
64    theta, phi, psi = view
65    return np.asarray((Rz(phi)*Ry(theta)*Rz(psi))*np.matrix(points.T)).T
66
67
68def invert_view(qx, qy, view):
69    """
70    Return (qa, qb, qc) for the (theta, phi, psi) view angle at detector
71    pixel (qx, qy).
72
73    View angles are in degrees.
74    """
75    theta, phi, psi = view
76    q = np.vstack((qx.flatten(), qy.flatten(), 0*qx.flatten()))
77    return np.asarray((Rz(-psi)*Ry(-theta)*Rz(-phi))*np.matrix(q))
78
79
80I3 = np.matrix([[1., 0, 0], [0, 1, 0], [0, 0, 1]])
81
82class Shape:
83    rotation = I3
84    center = np.array([0., 0., 0.])[:, None]
85    r_max = None
86
87    def volume(self):
88        # type: () -> float
89        raise NotImplementedError()
90
91    def sample(self, density):
92        # type: (float) -> np.ndarray[N], np.ndarray[N, 3]
93        raise NotImplementedError()
94
95    def dims(self):
96        # type: () -> float, float, float
97        raise NotImplementedError()
98
99    def rotate(self, theta, phi, psi):
100        if theta != 0. or phi != 0. or psi != 0.:
101            self.rotation = rotation(theta, phi, psi) * self.rotation
102        return self
103
104    def shift(self, x, y, z):
105        self.center = self.center + np.array([x, y, z])[:, None]
106        return self
107
108    def _adjust(self, points):
109        if self.rotation is I3:
110            points = points.T + self.center
111        else:
112            points = np.asarray(self.rotation * np.matrix(points.T)) + self.center
113        return points.T
114
115    def r_bins(self, q, over_sampling=1, r_step=0.):
116        r_max = min(2 * pi / q[0], self.r_max)
117        if r_step == 0.:
118            r_step = 2 * pi / q[-1] / over_sampling
119        #r_step = 0.01
120        return np.arange(r_step, r_max, r_step)
121
122class Composite(Shape):
123    def __init__(self, shapes, center=(0, 0, 0), orientation=(0, 0, 0)):
124        self.shapes = shapes
125        self.rotate(*orientation)
126        self.shift(*center)
127
128        # Find the worst case distance between any two points amongst a set
129        # of shapes independent of orientation.  This could easily be a
130        # factor of two worse than necessary, e.g., a pair of thin rods
131        # end-to-end vs the same pair side-by-side.
132        distances = [((s1.r_max + s2.r_max)/2
133                      + sqrt(np.sum((s1.center - s2.center)**2)))
134                     for s1 in shapes
135                     for s2 in shapes]
136        self.r_max = max(distances + [s.r_max for s in shapes])
137        self.volume = sum(shape.volume for shape in self.shapes)
138
139    def sample(self, density):
140        values, points = zip(*(shape.sample(density) for shape in self.shapes))
141        return np.hstack(values), self._adjust(np.vstack(points))
142
143class Box(Shape):
144    def __init__(self, a, b, c,
145                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
146        self.value = np.asarray(value)
147        self.rotate(*orientation)
148        self.shift(*center)
149        self.a, self.b, self.c = a, b, c
150        self._scale = np.array([a/2, b/2, c/2])[None, :]
151        self.r_max = sqrt(a**2 + b**2 + c**2)
152        self.dims = a, b, c
153        self.volume = a*b*c
154
155    def sample(self, density):
156        num_points = poisson(density*self.volume)
157        points = self._scale*uniform(-1, 1, size=(num_points, 3))
158        values = self.value.repeat(points.shape[0])
159        return values, self._adjust(points)
160
161class EllipticalCylinder(Shape):
162    def __init__(self, ra, rb, length,
163                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
164        self.value = np.asarray(value)
165        self.rotate(*orientation)
166        self.shift(*center)
167        self.ra, self.rb, self.length = ra, rb, length
168        self._scale = np.array([ra, rb, length/2])[None, :]
169        self.r_max = sqrt(4*max(ra, rb)**2 + length**2)
170        self.dims = 2*ra, 2*rb, length
171        self.volume = pi*ra*rb*length
172
173    def sample(self, density):
174        # randomly sample from a box of side length 2*r, excluding anything
175        # not in the cylinder
176        num_points = poisson(density*4*self.ra*self.rb*self.length)
177        points = uniform(-1, 1, size=(num_points, 3))
178        radius = points[:, 0]**2 + points[:, 1]**2
179        points = points[radius <= 1]
180        values = self.value.repeat(points.shape[0])
181        return values, self._adjust(self._scale*points)
182
183class EllipticalBicelle(Shape):
184    def __init__(self, ra, rb, length,
185                 thick_rim, thick_face,
186                 value_core, value_rim, value_face,
187                 center=(0, 0, 0), orientation=(0, 0, 0)):
188        self.rotate(*orientation)
189        self.shift(*center)
190        self.value = value_core
191        self.ra, self.rb, self.length = ra, rb, length
192        self.thick_rim, self.thick_face = thick_rim, thick_face
193        self.value_rim, self.value_face = value_rim, value_face
194
195        # reset cylinder to outer dimensions for calculating scale, etc.
196        ra = self.ra + self.thick_rim
197        rb = self.rb + self.thick_rim
198        length = self.length + 2*self.thick_face
199        self._scale = np.array([ra, rb, length/2])[None, :]
200        self.r_max = sqrt(4*max(ra, rb)**2 + length**2)
201        self.dims = 2*ra, 2*rb, length
202        self.volume = pi*ra*rb*length
203
204    def sample(self, density):
205        # randomly sample from a box of side length 2*r, excluding anything
206        # not in the cylinder
207        ra = self.ra + self.thick_rim
208        rb = self.rb + self.thick_rim
209        length = self.length + 2*self.thick_face
210        num_points = poisson(density*4*ra*rb*length)
211        points = uniform(-1, 1, size=(num_points, 3))
212        radius = points[:, 0]**2 + points[:, 1]**2
213        points = points[radius <= 1]
214        # set all to core value first
215        values = np.ones_like(points[:, 0])*self.value
216        # then set value to face value if |z| > face/(length/2))
217        values[abs(points[:, 2]) > self.length/(self.length + 2*self.thick_face)] = self.value_face
218        # finally set value to rim value if outside the core ellipse
219        radius = (points[:, 0]**2*(1 + self.thick_rim/self.ra)**2
220                  + points[:, 1]**2*(1 + self.thick_rim/self.rb)**2)
221        values[radius>1] = self.value_rim
222        return values, self._adjust(self._scale*points)
223
224class TriaxialEllipsoid(Shape):
225    def __init__(self, ra, rb, rc,
226                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
227        self.value = np.asarray(value)
228        self.rotate(*orientation)
229        self.shift(*center)
230        self.ra, self.rb, self.rc = ra, rb, rc
231        self._scale = np.array([ra, rb, rc])[None, :]
232        self.r_max = 2*max(ra, rb, rc)
233        self.dims = 2*ra, 2*rb, 2*rc
234        self.volume = 4*pi/3 * ra * rb * rc
235
236    def sample(self, density):
237        # randomly sample from a box of side length 2*r, excluding anything
238        # not in the ellipsoid
239        num_points = poisson(density*8*self.ra*self.rb*self.rc)
240        points = uniform(-1, 1, size=(num_points, 3))
241        radius = np.sum(points**2, axis=1)
242        points = self._scale*points[radius <= 1]
243        values = self.value.repeat(points.shape[0])
244        return values, self._adjust(points)
245
246class Helix(Shape):
247    def __init__(self, helix_radius, helix_pitch, tube_radius, tube_length,
248                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
249        self.value = np.asarray(value)
250        self.rotate(*orientation)
251        self.shift(*center)
252        helix_length = helix_pitch * tube_length/sqrt(helix_radius**2 + helix_pitch**2)
253        total_radius = self.helix_radius + self.tube_radius
254        self.helix_radius, self.helix_pitch = helix_radius, helix_pitch
255        self.tube_radius, self.tube_length = tube_radius, tube_length
256        self.r_max = sqrt(4*total_radius + (helix_length + 2*tube_radius)**2)
257        self.dims = 2*total_radius, 2*total_radius, helix_length
258        # small tube radius approximation; for larger tubes need to account
259        # for the fact that the inner length is much shorter than the outer
260        # length
261        self.volume = pi*self.tube_radius**2*self.tube_length
262
263    def points(self, density):
264        num_points = poisson(density*4*self.tube_radius**2*self.tube_length)
265        points = uniform(-1, 1, size=(num_points, 3))
266        radius = points[:, 0]**2 + points[:, 1]**2
267        points = points[radius <= 1]
268
269        # Based on math stackexchange answer by Jyrki Lahtonen
270        #     https://math.stackexchange.com/a/461637
271        # with helix along z rather than x [so tuples in answer are (z, x, y)]
272        # and with random points in the cross section (p1, p2) rather than
273        # uniform points on the surface (cos u, sin u).
274        a, R = self.tube_radius, self.helix_radius
275        h = self.helix_pitch
276        scale = 1/sqrt(R**2 + h**2)
277        t = points[:, 3] * (self.tube_length * scale/2)
278        cos_t, sin_t = cos(t), sin(t)
279
280        # rx = R*cos_t
281        # ry = R*sin_t
282        # rz = h*t
283        # nx = -a * cos_t * points[:, 1]
284        # ny = -a * sin_t * points[:, 1]
285        # nz = 0
286        # bx = (a * h/scale) * sin_t * points[:, 2]
287        # by = (-a * h/scale) * cos_t * points[:, 2]
288        # bz = a*R/scale
289        # x = rx + nx + bx
290        # y = ry + ny + by
291        # z = rz + nz + bz
292        u, v = (R - a*points[:, 1]), (a * h/scale)*points[:, 2]
293        x = u * cos_t + v * sin_t
294        y = u * sin_t - v * cos_t
295        z = a*R/scale + h * t
296
297        points = np.hstack((x, y, z))
298        values = self.value.repeat(points.shape[0])
299        return values, self._adjust(points)
300
301def csbox(a=10, b=20, c=30, da=1, db=2, dc=3, slda=1, sldb=2, sldc=3, sld_core=4):
302    core = Box(a, b, c, sld_core)
303    side_a = Box(da, b, c, slda, center=((a+da)/2, 0, 0))
304    side_b = Box(a, db, c, sldb, center=(0, (b+db)/2, 0))
305    side_c = Box(a, b, dc, sldc, center=(0, 0, (c+dc)/2))
306    side_a2 = copy(side_a).shift(-a-da, 0, 0)
307    side_b2 = copy(side_b).shift(0, -b-db, 0)
308    side_c2 = copy(side_c).shift(0, 0, -c-dc)
309    shape = Composite((core, side_a, side_b, side_c, side_a2, side_b2, side_c2))
310    shape.dims = 2*da+a, 2*db+b, 2*dc+c
311    return shape
312
313def _Iqxy(values, x, y, z, qa, qb, qc):
314    """I(q) = |sum V(r) rho(r) e^(1j q.r)|^2 / sum V(r)"""
315    Iq = [abs(np.sum(values*np.exp(1j*(qa_k*x + qb_k*y + qc_k*z))))**2
316            for qa_k, qb_k, qc_k in zip(qa.flat, qb.flat, qc.flat)]
317    return Iq
318
319if USE_NUMBA:
320    # Override simple numpy solution with numba if available
321    from numba import njit
322    @njit("f8[:](f8[:],f8[:],f8[:],f8[:],f8[:],f8[:],f8[:])")
323    def _Iqxy(values, x, y, z, qa, qb, qc):
324        Iq = np.zeros_like(qa)
325        for j in range(len(Iq)):
326            total = 0. + 0j
327            for k in range(len(values)):
328                total += values[k]*np.exp(1j*(qa[j]*x[k] + qb[j]*y[k] + qc[j]*z[k]))
329            Iq[j] = abs(total)**2
330        return Iq
331
332def calc_Iqxy(qx, qy, rho, points, volume=1.0, view=(0, 0, 0)):
333    qx, qy = np.broadcast_arrays(qx, qy)
334    qa, qb, qc = invert_view(qx, qy, view)
335    rho, volume = np.broadcast_arrays(rho, volume)
336    values = rho*volume
337    x, y, z = points.T
338    values, x, y, z, qa, qb, qc = [np.asarray(v, 'd')
339                                   for v in (values, x, y, z, qa, qb, qc)]
340
341    # I(q) = |sum V(r) rho(r) e^(1j q.r)|^2 / sum V(r)
342    Iq = _Iqxy(values, x, y, z, qa.flatten(), qb.flatten(), qc.flatten())
343    return np.asarray(Iq).reshape(qx.shape) / np.sum(volume)
344
345def _calc_Pr_nonuniform(r, rho, points):
346    # Make Pr a little be bigger than necessary so that only distances
347    # min < d < max end up in Pr
348    n_max = len(r)+1
349    extended_Pr = np.zeros(n_max+1, 'd')
350    # r refers to bin centers; find corresponding bin edges
351    bins = bin_edges(r)
352    t_next = time.clock() + 3
353    for k, rho_k in enumerate(rho[:-1]):
354        distance = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1))
355        weights = rho_k * rho[k+1:]
356        index = np.searchsorted(bins, distance)
357        # Note: indices may be duplicated, so "Pr[index] += w" will not work!!
358        extended_Pr += np.bincount(index, weights, n_max+1)
359        t = time.clock()
360        if t > t_next:
361            t_next = t + 3
362            print("processing %d of %d"%(k, len(rho)-1))
363    Pr = extended_Pr[1:-1]
364    return Pr
365
366def _calc_Pr_uniform(r, rho, points):
367    # Make Pr a little be bigger than necessary so that only distances
368    # min < d < max end up in Pr
369    dr, n_max = r[0], len(r)
370    extended_Pr = np.zeros(n_max+1, 'd')
371    t0 = time.clock()
372    t_next = t0 + 3
373    for k, rho_k in enumerate(rho[:-1]):
374        distances = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1))
375        weights = rho_k * rho[k+1:]
376        index = np.minimum(np.asarray(distances/dr, 'i'), n_max)
377        # Note: indices may be duplicated, so "Pr[index] += w" will not work!!
378        extended_Pr += np.bincount(index, weights, n_max+1)
379        t = time.clock()
380        if t > t_next:
381            t_next = t + 3
382            print("processing %d of %d"%(k, len(rho)-1))
383    #print("time py:", time.clock() - t0)
384    Pr = extended_Pr[:-1]
385    # Make Pr independent of sampling density.  The factor of 2 comes because
386    # we are only accumulating the upper triangular distances.
387    #Pr = Pr * 2 / len(rho)**2
388    return Pr
389
390    # Can get an additional 2x by going to C.  Cuda/OpenCL will allow even
391    # more speedup, though still bounded by the n^2 cose.
392    """
393void pdfcalc(int n, const double *pts, const double *rho,
394             int nPr, double *Pr, double rstep)
395{
396  int i,j;
397
398  for (i=0; i<n-2; i++) {
399    for (j=i+1; j<=n-1; j++) {
400      const double dxx=pts[3*i]-pts[3*j];
401      const double dyy=pts[3*i+1]-pts[3*j+1];
402      const double dzz=pts[3*i+2]-pts[3*j+2];
403      const double d=sqrt(dxx*dxx+dyy*dyy+dzz*dzz);
404      const int k=rint(d/rstep);
405      if (k < nPr) Pr[k]+=rho[i]*rho[j];
406    }
407  }
408}
409"""
410
411if USE_NUMBA:
412    # Override simple numpy solution with numba if available
413    @njit("f8[:](f8[:], f8[:], f8[:,:])")
414    def _calc_Pr_uniform(r, rho, points):
415        dr = r[0]
416        n_max = len(r)
417        Pr = np.zeros_like(r)
418        for j in range(len(rho) - 1):
419            x, y, z = points[j, 0], points[j, 1], points[j, 2]
420            for k in range(j+1, len(rho)):
421                distance = sqrt((x - points[k, 0])**2
422                                + (y - points[k, 1])**2
423                                + (z - points[k, 2])**2)
424                index = int(distance/dr)
425                if index < n_max:
426                    Pr[index] += rho[j] * rho[k]
427        # Make Pr independent of sampling density.  The factor of 2 comes because
428        # we are only accumulating the upper triangular distances.
429        #Pr = Pr * 2 / len(rho)**2
430        return Pr
431
432
433def calc_Pr(r, rho, points):
434    # P(r) with uniform steps in r is 3x faster; check if we are uniform
435    # before continuing
436    r, rho, points = [np.asarray(v, 'd') for v in (r, rho, points)]
437    if np.max(np.abs(np.diff(r) - r[0])) > r[0]*0.01:
438        Pr = _calc_Pr_nonuniform(r, rho, points)
439    else:
440        Pr = _calc_Pr_uniform(r, rho, points)
441    return Pr / Pr.max()
442
443
444def j0(x):
445    return np.sinc(x/np.pi)
446
447def calc_Iq(q, r, Pr):
448    Iq = np.array([simps(Pr * j0(qk*r), r) for qk in q])
449    #Iq = np.array([np.trapz(Pr * j0(qk*r), r) for qk in q])
450    Iq /= Iq[0]
451    return Iq
452
453# NOTE: copied from sasmodels/resolution.py
454def bin_edges(x):
455    """
456    Determine bin edges from bin centers, assuming that edges are centered
457    between the bins.
458
459    Note: this uses the arithmetic mean, which may not be appropriate for
460    log-scaled data.
461    """
462    if len(x) < 2 or (np.diff(x) < 0).any():
463        raise ValueError("Expected bins to be an increasing set")
464    edges = np.hstack([
465        x[0]  - 0.5*(x[1]  - x[0]),  # first point minus half first interval
466        0.5*(x[1:] + x[:-1]),        # mid points of all central intervals
467        x[-1] + 0.5*(x[-1] - x[-2]), # last point plus half last interval
468        ])
469    return edges
470
471# -------------- plotters ----------------
472def plot_calc(r, Pr, q, Iq, theory=None, title=None):
473    import matplotlib.pyplot as plt
474    plt.subplot(211)
475    plt.plot(r, Pr, '-', label="Pr")
476    plt.xlabel('r (A)')
477    plt.ylabel('Pr (1/A^2)')
478    if title is not None:
479        plt.title(title)
480    plt.subplot(212)
481    plt.loglog(q, Iq, '-', label='from Pr')
482    plt.xlabel('q (1/A')
483    plt.ylabel('Iq')
484    if theory is not None:
485        plt.loglog(theory[0], theory[1]/theory[1][0], '-', label='analytic')
486        plt.legend()
487
488def plot_calc_2d(qx, qy, Iqxy, theory=None, title=None):
489    import matplotlib.pyplot as plt
490    qx, qy = bin_edges(qx), bin_edges(qy)
491    #qx, qy = np.meshgrid(qx, qy)
492    if theory is not None:
493        plt.subplot(121)
494    #plt.pcolor(qx, qy, np.log10(Iqxy))
495    extent = [qx[0], qx[-1], qy[0], qy[-1]]
496    plt.imshow(np.log10(Iqxy), extent=extent, interpolation="nearest",
497               origin='lower')
498    plt.xlabel('qx (1/A)')
499    plt.ylabel('qy (1/A)')
500    plt.axis('equal')
501    plt.axis(extent)
502    #plt.grid(True)
503    if title is not None:
504        plt.title(title)
505    if theory is not None:
506        plt.subplot(122)
507        plt.imshow(np.log10(theory), extent=extent, interpolation="nearest",
508                   origin='lower')
509        plt.axis('equal')
510        plt.axis(extent)
511        plt.xlabel('qx (1/A)')
512
513def plot_points(rho, points):
514    import mpl_toolkits.mplot3d
515    import matplotlib.pyplot as plt
516
517    ax = plt.axes(projection='3d')
518    try:
519        ax.axis('square')
520    except Exception:
521        pass
522    n = len(points)
523    #print("len points", n)
524    index = np.random.choice(n, size=500) if n > 500 else slice(None, None)
525    ax.scatter(points[index, 0], points[index, 1], points[index, 2], c=rho[index])
526    # make square axes
527    minmax = np.array([points.min(), points.max()])
528    ax.scatter(minmax, minmax, minmax, c='w')
529    #low, high = points.min(axis=0), points.max(axis=0)
530    #ax.axis([low[0], high[0], low[1], high[1], low[2], high[2]])
531    ax.set_xlabel("x")
532    ax.set_ylabel("y")
533    ax.set_zlabel("z")
534    ax.autoscale(True)
535
536# ----------- Analytic models --------------
537def sas_sinx_x(x):
538    with np.errstate(all='ignore'):
539        retvalue = sin(x)/x
540    retvalue[x == 0.] = 1.
541    return retvalue
542
543def sas_2J1x_x(x):
544    with np.errstate(all='ignore'):
545        retvalue = 2*J1(x)/x
546    retvalue[x == 0] = 1.
547    return retvalue
548
549def sas_3j1x_x(x):
550    """return 3*j1(x)/x"""
551    with np.errstate(all='ignore'):
552        retvalue = 3*(sin(x) - x*cos(x))/x**3
553    retvalue[x == 0.] = 1.
554    return retvalue
555
556def cylinder_Iq(q, radius, length):
557    z, w = leggauss(76)
558    cos_alpha = (z+1)/2
559    sin_alpha = sqrt(1.0 - cos_alpha**2)
560    Iq = np.empty_like(q)
561    for k, qk in enumerate(q):
562        qab, qc = qk*sin_alpha, qk*cos_alpha
563        Fq = sas_2J1x_x(qab*radius) * sas_sinx_x(qc*length/2)
564        Iq[k] = np.sum(w*Fq**2)
565    Iq = Iq
566    return Iq
567
568def cylinder_Iqxy(qx, qy, radius, length, view=(0, 0, 0)):
569    qa, qb, qc = invert_view(qx, qy, view)
570    qab = sqrt(qa**2 + qb**2)
571    Fq = sas_2J1x_x(qab*radius) * sas_sinx_x(qc*length/2)
572    Iq = Fq**2
573    return Iq.reshape(qx.shape)
574
575def sphere_Iq(q, radius):
576    Iq = sas_3j1x_x(q*radius)**2
577    return Iq
578
579def box_Iq(q, a, b, c):
580    z, w = leggauss(76)
581    outer_sum = np.zeros_like(q)
582    for cos_alpha, outer_w in zip((z+1)/2, w):
583        sin_alpha = sqrt(1.0-cos_alpha*cos_alpha)
584        qc = q*cos_alpha
585        siC = c*sas_sinx_x(c*qc/2)
586        inner_sum = np.zeros_like(q)
587        for beta, inner_w in zip((z + 1)*pi/4, w):
588            qa, qb = q*sin_alpha*sin(beta), q*sin_alpha*cos(beta)
589            siA = a*sas_sinx_x(a*qa/2)
590            siB = b*sas_sinx_x(b*qb/2)
591            Fq = siA*siB*siC
592            inner_sum += inner_w * Fq**2
593        outer_sum += outer_w * inner_sum
594    Iq = outer_sum / 4  # = outer*um*zm*8.0/(4.0*M_PI)
595    return Iq
596
597def box_Iqxy(qx, qy, a, b, c, view=(0, 0, 0)):
598    qa, qb, qc = invert_view(qx, qy, view)
599    sia = sas_sinx_x(qa*a/2)
600    sib = sas_sinx_x(qb*b/2)
601    sic = sas_sinx_x(qc*c/2)
602    Fq = sia*sib*sic
603    Iq = Fq**2
604    return Iq.reshape(qx.shape)
605
606def csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core):
607    z, w = leggauss(76)
608
609    sld_solvent = 0
610    overlapping = False
611    dr0 = sld_core - sld_solvent
612    drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent
613    tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc
614
615    outer_sum = np.zeros_like(q)
616    for cos_alpha, outer_w in zip((z+1)/2, w):
617        sin_alpha = sqrt(1.0-cos_alpha*cos_alpha)
618        qc = q*cos_alpha
619        siC = c*sas_sinx_x(c*qc/2)
620        siCt = tC*sas_sinx_x(tC*qc/2)
621        inner_sum = np.zeros_like(q)
622        for beta, inner_w in zip((z + 1)*pi/4, w):
623            qa, qb = q*sin_alpha*sin(beta), q*sin_alpha*cos(beta)
624            siA = a*sas_sinx_x(a*qa/2)
625            siB = b*sas_sinx_x(b*qb/2)
626            siAt = tA*sas_sinx_x(tA*qa/2)
627            siBt = tB*sas_sinx_x(tB*qb/2)
628            if overlapping:
629                Fq = (dr0*siA*siB*siC
630                      + drA*(siAt-siA)*siB*siC
631                      + drB*siAt*(siBt-siB)*siC
632                      + drC*siAt*siBt*(siCt-siC))
633            else:
634                Fq = (dr0*siA*siB*siC
635                      + drA*(siAt-siA)*siB*siC
636                      + drB*siA*(siBt-siB)*siC
637                      + drC*siA*siB*(siCt-siC))
638            inner_sum += inner_w * Fq**2
639        outer_sum += outer_w * inner_sum
640    Iq = outer_sum / 4  # = outer*um*zm*8.0/(4.0*M_PI)
641    return Iq/Iq[0]
642
643def csbox_Iqxy(qx, qy, a, b, c, da, db, dc, slda, sldb, sldc, sld_core, view=(0,0,0)):
644    qa, qb, qc = invert_view(qx, qy, view)
645
646    sld_solvent = 0
647    overlapping = False
648    dr0 = sld_core - sld_solvent
649    drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent
650    tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc
651    siA = a*sas_sinx_x(a*qa/2)
652    siB = b*sas_sinx_x(b*qb/2)
653    siC = c*sas_sinx_x(c*qc/2)
654    siAt = tA*sas_sinx_x(tA*qa/2)
655    siBt = tB*sas_sinx_x(tB*qb/2)
656    siCt = tC*sas_sinx_x(tC*qc/2)
657    Fq = (dr0*siA*siB*siC
658          + drA*(siAt-siA)*siB*siC
659          + drB*siA*(siBt-siB)*siC
660          + drC*siA*siB*(siCt-siC))
661    Iq = Fq**2
662    return Iq.reshape(qx.shape)
663
664def sasmodels_Iq(kernel, q, pars):
665    from sasmodels.data import empty_data1D
666    from sasmodels.direct_model import DirectModel
667    data = empty_data1D(q)
668    calculator = DirectModel(data, kernel)
669    Iq = calculator(**pars)
670    return Iq
671
672def sasmodels_Iqxy(kernel, qx, qy, pars, view):
673    from sasmodels.data import Data2D
674    from sasmodels.direct_model import DirectModel
675    Iq = 100 * np.ones_like(qx)
676    data = Data2D(x=qx, y=qy, z=Iq, dx=None, dy=None, dz=np.sqrt(Iq))
677    data.x_bins = qx[0, :]
678    data.y_bins = qy[:, 0]
679    data.filename = "fake data"
680
681    calculator = DirectModel(data, kernel)
682    pars_plus_view = pars.copy()
683    pars_plus_view.update(theta=view[0], phi=view[1], psi=view[2])
684    Iqxy = calculator(**pars_plus_view)
685    return Iqxy.reshape(qx.shape)
686
687def wrap_sasmodel(name, **pars):
688    from sasmodels.core import load_model
689    kernel = load_model(name)
690    fn = lambda q: sasmodels_Iq(kernel, q, pars)
691    fn_xy = lambda qx, qy, view: sasmodels_Iqxy(kernel, qx, qy, pars, view)
692    return fn, fn_xy
693
694
695# --------- Test cases -----------
696
697def build_cylinder(radius=25, length=125, rho=2.):
698    shape = EllipticalCylinder(radius, radius, length, rho)
699    fn = lambda q: cylinder_Iq(q, radius, length)*rho**2
700    fn_xy = lambda qx, qy, view: cylinder_Iqxy(qx, qy, radius, length, view=view)*rho**2
701    return shape, fn, fn_xy
702
703DEFAULT_SPHERE_RADIUS = 125
704DEFAULT_SPHERE_CONTRAST = 2
705def build_sphere(radius=DEFAULT_SPHERE_RADIUS, rho=DEFAULT_SPHERE_CONTRAST):
706    shape = TriaxialEllipsoid(radius, radius, radius, rho)
707    fn = lambda q: sphere_Iq(q, radius)*rho**2
708    fn_xy = lambda qx, qy, view: sphere_Iq(np.sqrt(qx**2+qy**2), radius)*rho**2
709    return shape, fn, fn_xy
710
711def build_box(a=10, b=20, c=30, rho=2.):
712    shape = Box(a, b, c, rho)
713    fn = lambda q: box_Iq(q, a, b, c)*rho**2
714    fn_xy = lambda qx, qy, view: box_Iqxy(qx, qy, a, b, c, view=view)*rho**2
715    return shape, fn, fn_xy
716
717def build_csbox(a=10, b=20, c=30, da=1, db=2, dc=3, slda=1, sldb=2, sldc=3, sld_core=4):
718    shape = csbox(a, b, c, da, db, dc, slda, sldb, sldc, sld_core)
719    fn = lambda q: csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core)
720    fn_xy = lambda qx, qy, view: csbox_Iqxy(qx, qy, a, b, c, da, db, dc,
721                                            slda, sldb, sldc, sld_core, view=view)
722    return shape, fn, fn_xy
723
724def build_ellcyl(ra=25, rb=50, length=125, rho=2.):
725    shape = EllipticalCylinder(ra, rb, length, rho)
726    fn, fn_xy = wrap_sasmodel(
727        'elliptical_cylinder',
728        scale=1,
729        background=0,
730        radius_minor=ra,
731        axis_ratio=rb/ra,
732        length=length,
733        sld=rho,
734        sld_solvent=0,
735    )
736    return shape, fn, fn_xy
737
738def build_cscyl(ra=30, rb=90, length=30, thick_rim=8, thick_face=14,
739                sld_core=4, sld_rim=1, sld_face=7):
740    shape = EllipticalBicelle(
741        ra=ra, rb=rb, length=length,
742        thick_rim=thick_rim, thick_face=thick_face,
743        value_core=sld_core, value_rim=sld_rim, value_face=sld_face,
744        )
745    fn, fn_xy = wrap_sasmodel(
746        'core_shell_bicelle_elliptical',
747        scale=1,
748        background=0,
749        radius=ra,
750        x_core=rb/ra,
751        length=length,
752        thick_rim=thick_rim,
753        thick_face=thick_face,
754        sld_core=sld_core,
755        sld_face=sld_face,
756        sld_rim=sld_rim,
757        sld_solvent=0,
758    )
759    return shape, fn, fn_xy
760
761def build_sc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2,
762                        shuffle=0, rotate=0):
763    a, b, c = shape.dims
764    corners= [copy(shape)
765              .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
766                     (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
767                     (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
768              .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
769              for ix in range(nx)
770              for iy in range(ny)
771              for iz in range(nz)]
772    lattice = Composite(corners)
773    return lattice
774
775def build_bcc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2,
776                      shuffle=0, rotate=0):
777    a, b, c = shape.dims
778    corners = [copy(shape)
779               .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
780                      (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
781                      (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
782               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
783               for ix in range(nx)
784               for iy in range(ny)
785               for iz in range(nz)]
786    centers = [copy(shape)
787               .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
788                      (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
789                      (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
790               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
791               for ix in range(nx)
792               for iy in range(ny)
793               for iz in range(nz)]
794    lattice = Composite(corners + centers)
795    return lattice
796
797def build_fcc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2,
798                      shuffle=0, rotate=0):
799    a, b, c = shape.dims
800    corners = [copy(shape)
801               .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
802                      (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
803                      (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
804               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
805               for ix in range(nx)
806               for iy in range(ny)
807               for iz in range(nz)]
808    faces_a = [copy(shape)
809               .shift((ix+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
810                      (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
811                      (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
812               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
813               for ix in range(nx)
814               for iy in range(ny)
815               for iz in range(nz)]
816    faces_b = [copy(shape)
817               .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
818                      (iy+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
819                      (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
820               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
821               for ix in range(nx)
822               for iy in range(ny)
823               for iz in range(nz)]
824    faces_c = [copy(shape)
825               .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a,
826                      (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b,
827                      (iz+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c)
828               .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate))
829               for ix in range(nx)
830               for iy in range(ny)
831               for iz in range(nz)]
832    lattice = Composite(corners + faces_a + faces_b + faces_c)
833    return lattice
834
835SHAPE_FUNCTIONS = OrderedDict([
836    ("cyl", build_cylinder),
837    ("ellcyl", build_ellcyl),
838    ("sphere", build_sphere),
839    ("box", build_box),
840    ("csbox", build_csbox),
841    ("cscyl", build_cscyl),
842])
843SHAPES = list(SHAPE_FUNCTIONS.keys())
844LATTICE_FUNCTIONS = OrderedDict([
845    ("sc", build_sc_lattice),
846    ("bcc", build_bcc_lattice),
847    ("fcc", build_fcc_lattice),
848])
849LATTICE_TYPES = list(LATTICE_FUNCTIONS.keys())
850
851def check_shape(title, shape, fn=None, show_points=False,
852                mesh=100, qmax=1.0, r_step=0.01, samples=5000):
853    rho_solvent = 0
854    qmin = qmax/100.
855    q = np.logspace(np.log10(qmin), np.log10(qmax), mesh)
856    r = shape.r_bins(q, r_step=r_step)
857    sampling_density = samples / shape.volume
858    print("sampling points")
859    rho, points = shape.sample(sampling_density)
860    print("calculating Pr")
861    t0 = time.time()
862    Pr = calc_Pr(r, rho-rho_solvent, points)
863    print("calc Pr time", time.time() - t0)
864    Iq = calc_Iq(q, r, Pr)
865    theory = (q, fn(q)) if fn is not None else None
866
867    import pylab
868    if show_points:
869         plot_points(rho, points); pylab.figure()
870    plot_calc(r, Pr, q, Iq, theory=theory, title=title)
871    pylab.gcf().canvas.set_window_title(title)
872    pylab.show()
873
874def check_shape_2d(title, shape, fn=None, view=(0, 0, 0), show_points=False,
875                   mesh=100, qmax=1.0, samples=5000):
876    rho_solvent = 0
877    #qx = np.linspace(0.0, qmax, mesh)
878    #qy = np.linspace(0.0, qmax, mesh)
879    qx = np.linspace(-qmax, qmax, mesh)
880    qy = np.linspace(-qmax, qmax, mesh)
881    Qx, Qy = np.meshgrid(qx, qy)
882    sampling_density = samples / shape.volume
883    print("sampling points")
884    t0 = time.time()
885    rho, points = shape.sample(sampling_density)
886    print("point generation time", time.time() - t0)
887    t0 = time.time()
888    Iqxy = calc_Iqxy(Qx, Qy, rho, points, view=view)
889    print("calc Iqxy time", time.time() - t0)
890    t0 = time.time()
891    theory = fn(Qx, Qy, view) if fn is not None else None
892    print("calc theory time", time.time() - t0)
893    Iqxy += 0.001 * Iqxy.max()
894    if theory is not None:
895        theory += 0.001 * theory.max()
896
897    import pylab
898    if show_points:
899        plot_points(rho, points); pylab.figure()
900    plot_calc_2d(qx, qy, Iqxy, theory=theory, title=title)
901    pylab.gcf().canvas.set_window_title(title)
902    pylab.show()
903
904def main():
905    parser = argparse.ArgumentParser(
906        description="Compute scattering from realspace sampling",
907        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
908        )
909    parser.add_argument('-d', '--dim', type=int, default=1,
910                        help='dimension 1 or 2')
911    parser.add_argument('-m', '--mesh', type=int, default=100,
912                        help='number of mesh points')
913    parser.add_argument('-s', '--samples', type=int, default=5000,
914                        help="number of sample points")
915    parser.add_argument('-q', '--qmax', type=float, default=0.5,
916                        help='max q')
917    parser.add_argument('-v', '--view', type=str, default='0,0,0',
918                        help='theta,phi,psi angles')
919    parser.add_argument('-n', '--lattice', type=str, default='1,1,1',
920                        help='lattice size')
921    parser.add_argument('-z', '--spacing', type=str, default='2,2,2',
922                        help='lattice spacing (relative to shape)')
923    parser.add_argument('-t', '--type', choices=LATTICE_TYPES,
924                        default=LATTICE_TYPES[0],
925                        help='lattice type')
926    parser.add_argument('-r', '--rotate', type=float, default=0.,
927                        help="rotation relative to lattice, gaussian < 30 degrees, uniform otherwise")
928    parser.add_argument('-w', '--shuffle', type=float, default=0.,
929                        help="position relative to lattice, gaussian < 0.3, uniform otherwise")
930    parser.add_argument('-p', '--plot', action='store_true',
931                        help='plot points')
932    parser.add_argument('shape', choices=SHAPES, nargs='?', default=SHAPES[0],
933                        help='oriented shape')
934    parser.add_argument('pars', type=str, nargs='*', help='shape parameters')
935    opts = parser.parse_args()
936    pars = {key: float(value) for p in opts.pars for key, value in [p.split('=')]}
937    nx, ny, nz = [int(v) for v in opts.lattice.split(',')]
938    dx, dy, dz = [float(v) for v in opts.spacing.split(',')]
939    shuffle, rotate = opts.shuffle, opts.rotate
940    shape, fn, fn_xy = SHAPE_FUNCTIONS[opts.shape](**pars)
941    view = tuple(float(v) for v in opts.view.split(','))
942    if nx > 1 or ny > 1 or nz > 1:
943        print("building %s lattice"%opts.type)
944        lattice = LATTICE_FUNCTIONS[opts.type]
945        shape = lattice(shape, nx, ny, nz, dx, dy, dz, shuffle, rotate)
946        # If comparing a sphere in a cubic lattice, compare against the
947        # corresponding paracrystalline model.
948        if opts.shape == "sphere" and dx == dy == dz:
949            radius = pars.get('radius', DEFAULT_SPHERE_RADIUS)
950            model_name = opts.type + "_paracrystal"
951            model_pars = {
952                "scale": 1.,
953                "background": 0.,
954                "lattice_spacing": 2*radius*dx,
955                "lattice_distortion": shuffle,
956                "radius": radius,
957                "sld": pars.get('rho', DEFAULT_SPHERE_CONTRAST),
958                "sld_solvent": 0.,
959                "theta": view[0],
960                "phi": view[1],
961                "psi": view[2],
962            }
963            fn, fn_xy = wrap_sasmodel(model_name, **model_pars)
964    title = "%s(%s)" % (opts.shape, " ".join(opts.pars))
965    if opts.dim == 1:
966        check_shape(title, shape, fn, show_points=opts.plot,
967                    mesh=opts.mesh, qmax=opts.qmax, samples=opts.samples)
968    else:
969        check_shape_2d(title, shape, fn_xy, view=view, show_points=opts.plot,
970                       mesh=opts.mesh, qmax=opts.qmax, samples=opts.samples)
971
972
973if __name__ == "__main__":
974    main()
Note: See TracBrowser for help on using the repository browser.