source: sasmodels/explore/realspace.py @ 6868151

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

realspace: fix sampling for cscyl rim

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