source: sasmodels/explore/realspace.py @ f354e46

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

play with numba JIT compiler

  • Property mode set to 100644
File size: 22.1 KB
Line 
1from __future__ import division, print_function
2
3import time
4from copy import copy
5
6import numpy as np
7from numpy import pi, radians, sin, cos, sqrt
8from numpy.random import poisson, uniform
9from numpy.polynomial.legendre import leggauss
10from scipy.integrate import simps
11from scipy.special import j1 as J1
12
13# Definition of rotation matrices comes from wikipedia:
14#    https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations
15def Rx(angle):
16    """Construct a matrix to rotate points about *x* by *angle* degrees."""
17    a = radians(angle)
18    R = [[1, 0, 0],
19         [0, +cos(a), -sin(a)],
20         [0, +sin(a), +cos(a)]]
21    return np.matrix(R)
22
23def Ry(angle):
24    """Construct a matrix to rotate points about *y* by *angle* degrees."""
25    a = radians(angle)
26    R = [[+cos(a), 0, +sin(a)],
27         [0, 1, 0],
28         [-sin(a), 0, +cos(a)]]
29    return np.matrix(R)
30
31def Rz(angle):
32    """Construct a matrix to rotate points about *z* by *angle* degrees."""
33    a = radians(angle)
34    R = [[+cos(a), -sin(a), 0],
35         [+sin(a), +cos(a), 0],
36         [0, 0, 1]]
37    return np.matrix(R)
38
39def rotation(theta, phi, psi):
40    """
41    Apply the jitter transform to a set of points.
42
43    Points are stored in a 3 x n numpy matrix, not a numpy array or tuple.
44    """
45    return Rx(phi)*Ry(theta)*Rz(psi)
46
47def apply_view(points, view):
48    """
49    Apply the view transform (theta, phi, psi) to a set of points.
50
51    Points are stored in a 3 x n numpy array.
52
53    View angles are in degrees.
54    """
55    theta, phi, psi = view
56    return np.asarray((Rz(phi)*Ry(theta)*Rz(psi))*np.matrix(points.T)).T
57
58
59def invert_view(qx, qy, view):
60    """
61    Return (qa, qb, qc) for the (theta, phi, psi) view angle at detector
62    pixel (qx, qy).
63
64    View angles are in degrees.
65    """
66    theta, phi, psi = view
67    q = np.vstack((qx.flatten(), qy.flatten(), 0*qx.flatten()))
68    return np.asarray((Rz(-psi)*Ry(-theta)*Rz(-phi))*np.matrix(q))
69
70
71class Shape:
72    rotation = np.matrix([[1., 0, 0], [0, 1, 0], [0, 0, 1]])
73    center = np.array([0., 0., 0.])[:, None]
74    r_max = None
75
76    def volume(self):
77        # type: () -> float
78        raise NotImplementedError()
79
80    def sample(self, density):
81        # type: (float) -> np.ndarray[N], np.ndarray[N, 3]
82        raise NotImplementedError()
83
84    def rotate(self, theta, phi, psi):
85        self.rotation = rotation(theta, phi, psi) * self.rotation
86        return self
87
88    def shift(self, x, y, z):
89        self.center = self.center + np.array([x, y, z])[:, None]
90        return self
91
92    def _adjust(self, points):
93        points = np.asarray(self.rotation * np.matrix(points.T)) + self.center
94        return points.T
95
96    def r_bins(self, q, over_sampling=1, r_step=0.):
97        r_max = min(2 * pi / q[0], self.r_max)
98        if r_step == 0.:
99            r_step = 2 * pi / q[-1] / over_sampling
100        #r_step = 0.01
101        return np.arange(r_step, r_max, r_step)
102
103class Composite(Shape):
104    def __init__(self, shapes, center=(0, 0, 0), orientation=(0, 0, 0)):
105        self.shapes = shapes
106        self.rotate(*orientation)
107        self.shift(*center)
108
109        # Find the worst case distance between any two points amongst a set
110        # of shapes independent of orientation.  This could easily be a
111        # factor of two worse than necessary, e.g., a pair of thin rods
112        # end-to-end vs the same pair side-by-side.
113        distances = [((s1.r_max + s2.r_max)/2
114                      + sqrt(np.sum((s1.center - s2.center)**2)))
115                     for s1 in shapes
116                     for s2 in shapes]
117        self.r_max = max(distances + [s.r_max for s in shapes])
118
119    def volume(self):
120        return sum(shape.volume() for shape in self.shapes)
121
122    def sample(self, density):
123        values, points = zip(*(shape.sample(density) for shape in self.shapes))
124        return np.hstack(values), self._adjust(np.vstack(points))
125
126class Box(Shape):
127    def __init__(self, a, b, c,
128                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
129        self.value = np.asarray(value)
130        self.rotate(*orientation)
131        self.shift(*center)
132        self.a, self.b, self.c = a, b, c
133        self._scale = np.array([a/2, b/2, c/2])[None, :]
134        self.r_max = sqrt(a**2 + b**2 + c**2)
135
136    def volume(self):
137        return self.a*self.b*self.c
138
139    def sample(self, density):
140        num_points = poisson(density*self.a*self.b*self.c)
141        points = self._scale*uniform(-1, 1, size=(num_points, 3))
142        values = self.value.repeat(points.shape[0])
143        return values, self._adjust(points)
144
145class EllipticalCylinder(Shape):
146    def __init__(self, ra, rb, length,
147                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
148        self.value = np.asarray(value)
149        self.rotate(*orientation)
150        self.shift(*center)
151        self.ra, self.rb, self.length = ra, rb, length
152        self._scale = np.array([ra, rb, length/2])[None, :]
153        self.r_max = sqrt(4*max(ra, rb)**2 + length**2)
154
155    def volume(self):
156        return pi*self.ra*self.rb*self.length
157
158    def sample(self, density):
159        # density of the bounding box
160        num_points = poisson(density*4*self.ra*self.rb*self.length)
161        points = uniform(-1, 1, size=(num_points, 3))
162        radius = points[:, 0]**2 + points[:, 1]**2
163        points = self._scale*points[radius <= 1]
164        values = self.value.repeat(points.shape[0])
165        return values, self._adjust(points)
166
167class TriaxialEllipsoid(Shape):
168    def __init__(self, ra, rb, rc,
169                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
170        self.value = np.asarray(value)
171        self.rotate(*orientation)
172        self.shift(*center)
173        self.ra, self.rb, self.rc = ra, rb, rc
174        self._scale = np.array([ra, rb, rc])[None, :]
175        self.r_max = 2*max(ra, rb, rc)
176
177    def volume(self):
178        return 4*pi/3 * self.ra * self.rb * self.rc
179
180    def sample(self, density):
181        # randomly sample from a box of side length 2*r, excluding anything
182        # not in the ellipsoid
183        num_points = poisson(density*8*self.ra*self.rb*self.rc)
184        points = uniform(-1, 1, size=(num_points, 3))
185        radius = np.sum(points**2, axis=1)
186        points = self._scale*points[radius <= 1]
187        values = self.value.repeat(points.shape[0])
188        return values, self._adjust(points)
189
190class Helix(Shape):
191    def __init__(self, helix_radius, helix_pitch, tube_radius, tube_length,
192                 value, center=(0, 0, 0), orientation=(0, 0, 0)):
193        self.value = np.asarray(value)
194        self.rotate(*orientation)
195        self.shift(*center)
196        self.helix_radius, self.helix_pitch = helix_radius, helix_pitch
197        self.tube_radius, self.tube_length = tube_radius, tube_length
198        helix_length = helix_pitch * tube_length/sqrt(helix_radius**2 + helix_pitch**2)
199        self.r_max = sqrt((2*helix_radius + 2*tube_radius)*2
200                          + (helix_length + 2*tube_radius)**2)
201
202    def volume(self):
203        # small tube radius approximation; for larger tubes need to account
204        # for the fact that the inner length is much shorter than the outer
205        # length
206        return pi*self.tube_radius**2*self.tube_length
207
208    def points(self, density):
209        num_points = poisson(density*4*self.tube_radius**2*self.tube_length)
210        points = uniform(-1, 1, size=(num_points, 3))
211        radius = points[:, 0]**2 + points[:, 1]**2
212        points = points[radius <= 1]
213
214        # Based on math stackexchange answer by Jyrki Lahtonen
215        #     https://math.stackexchange.com/a/461637
216        # with helix along z rather than x [so tuples in answer are (z, x, y)]
217        # and with random points in the cross section (p1, p2) rather than
218        # uniform points on the surface (cos u, sin u).
219        a, R = self.tube_radius, self.helix_radius
220        h = self.helix_pitch
221        scale = 1/sqrt(R**2 + h**2)
222        t = points[:, 3] * (self.tube_length * scale/2)
223        cos_t, sin_t = cos(t), sin(t)
224
225        # rx = R*cos_t
226        # ry = R*sin_t
227        # rz = h*t
228        # nx = -a * cos_t * points[:, 1]
229        # ny = -a * sin_t * points[:, 1]
230        # nz = 0
231        # bx = (a * h/scale) * sin_t * points[:, 2]
232        # by = (-a * h/scale) * cos_t * points[:, 2]
233        # bz = a*R/scale
234        # x = rx + nx + bx
235        # y = ry + ny + by
236        # z = rz + nz + bz
237        u, v = (R - a*points[:, 1]), (a * h/scale)*points[:, 2]
238        x = u * cos_t + v * sin_t
239        y = u * sin_t - v * cos_t
240        z = a*R/scale + h * t
241
242        points = np.hstack((x, y, z))
243        values = self.value.repeat(points.shape[0])
244        return values, self._adjust(points)
245
246NUMBA = False
247if NUMBA:
248    from numba import njit
249    @njit("f8[:](f8[:],f8[:],f8[:],f8[:],f8[:],f8[:],f8[:])")
250    def _Iqxy(values, x, y, z, qa, qb, qc):
251        Iq = np.zeros_like(qa)
252        for j in range(len(Iq)):
253            total = 0. + 0j
254            for k in range(len(Iq)):
255                total += values[k]*np.exp(1j*(qa[j]*x[k] + qb[j]*y[k] + qc[j]*z[k]))
256            Iq[j] = abs(total)**2
257        return Iq
258
259def calc_Iqxy(qx, qy, rho, points, volume=1.0, view=(0, 0, 0)):
260    qx, qy = np.broadcast_arrays(qx, qy)
261    qa, qb, qc = invert_view(qx, qy, view)
262    rho, volume = np.broadcast_arrays(rho, volume)
263    values = rho*volume
264    x, y, z = points.T
265
266    # I(q) = |sum V(r) rho(r) e^(1j q.r)|^2 / sum V(r)
267    if NUMBA:
268        Iq = _Iqxy(values, x, y, z, qa.flatten(), qb.flatten(), qc.flatten())
269    else:
270        Iq = [abs(np.sum(values*np.exp(1j*(qa_k*x + qb_k*y + qc_k*z))))**2
271              for qa_k, qb_k, qc_k in zip(qa.flat, qb.flat, qc.flat)]
272    return np.asarray(Iq).reshape(qx.shape) / np.sum(volume)
273
274def _calc_Pr_nonuniform(r, rho, points):
275    # Make Pr a little be bigger than necessary so that only distances
276    # min < d < max end up in Pr
277    n_max = len(r)+1
278    extended_Pr = np.zeros(n_max+1, 'd')
279    # r refers to bin centers; find corresponding bin edges
280    bins = bin_edges(r)
281    t_next = time.clock() + 3
282    for k, rho_k in enumerate(rho[:-1]):
283        distance = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1))
284        weights = rho_k * rho[k+1:]
285        index = np.searchsorted(bins, distance)
286        # Note: indices may be duplicated, so "Pr[index] += w" will not work!!
287        extended_Pr += np.bincount(index, weights, n_max+1)
288        t = time.clock()
289        if t > t_next:
290            t_next = t + 3
291            print("processing %d of %d"%(k, len(rho)-1))
292    Pr = extended_Pr[1:-1]
293    return Pr
294
295def _calc_Pr_uniform(r, rho, points):
296    # Make Pr a little be bigger than necessary so that only distances
297    # min < d < max end up in Pr
298    dr, n_max = r[0], len(r)
299    extended_Pr = np.zeros(n_max+1, 'd')
300    t0 = time.clock()
301    t_next = t0 + 3
302    for k, rho_k in enumerate(rho[:-1]):
303        distances = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1))
304        weights = rho_k * rho[k+1:]
305        index = np.minimum(np.asarray(distances/dr, 'i'), n_max)
306        # Note: indices may be duplicated, so "Pr[index] += w" will not work!!
307        extended_Pr += np.bincount(index, weights, n_max+1)
308        t = time.clock()
309        if t > t_next:
310            t_next = t + 3
311            print("processing %d of %d"%(k, len(rho)-1))
312    #print("time py:", time.clock() - t0)
313    Pr = extended_Pr[:-1]
314    # Make Pr independent of sampling density.  The factor of 2 comes because
315    # we are only accumulating the upper triangular distances.
316    #Pr = Pr * 2 / len(rho)**2
317    return Pr
318
319    # Can get an additional 2x by going to C.  Cuda/OpenCL will allow even
320    # more speedup, though still bounded by the n^2 cose.
321    """
322void pdfcalc(int n, const double *pts, const double *rho,
323             int nPr, double *Pr, double rstep)
324{
325  int i,j;
326
327  for (i=0; i<n-2; i++) {
328    for (j=i+1; j<=n-1; j++) {
329      const double dxx=pts[3*i]-pts[3*j];
330      const double dyy=pts[3*i+1]-pts[3*j+1];
331      const double dzz=pts[3*i+2]-pts[3*j+2];
332      const double d=sqrt(dxx*dxx+dyy*dyy+dzz*dzz);
333      const int k=rint(d/rstep);
334      if (k < nPr) Pr[k]+=rho[i]*rho[j];
335    }
336  }
337}
338"""
339
340if NUMBA:
341    @njit("f8[:](f8[:], f8[:], f8[:,:])")
342    def _calc_Pr_uniform_jit(r, rho, points):
343        dr = r[0]
344        n_max = len(r)
345        Pr = np.zeros_like(r)
346        for j in range(len(rho) - 1):
347            x, y, z = points[j, 0], points[j, 1], points[j, 2]
348            for k in range(j+1, len(rho)):
349                distance = sqrt((x - points[k, 0])**2
350                                + (y - points[k, 1])**2
351                                + (z - points[k, 2])**2)
352                index = int(distance/dr)
353                if index < n_max:
354                    Pr[index] += rho[j] * rho[k]
355        # Make Pr independent of sampling density.  The factor of 2 comes because
356        # we are only accumulating the upper triangular distances.
357        #Pr = Pr * 2 / len(rho)**2
358        return Pr
359
360
361def calc_Pr(r, rho, points):
362    # P(r) with uniform steps in r is 3x faster; check if we are uniform
363    # before continuing
364    if np.max(np.abs(np.diff(r) - r[0])) > r[0]*0.01:
365        Pr = _calc_Pr_nonuniform(r, rho, points)
366    else:
367        if NUMBA:
368            Pr = _calc_Pr_uniform_jit(r, rho, points)
369        else:
370            Pr = _calc_Pr_uniform(r, rho, points)
371    return Pr / Pr.max()
372
373
374def j0(x):
375    return np.sinc(x/np.pi)
376
377def calc_Iq(q, r, Pr):
378    Iq = np.array([simps(Pr * j0(qk*r), r) for qk in q])
379    #Iq = np.array([np.trapz(Pr * j0(qk*r), r) for qk in q])
380    Iq /= Iq[0]
381    return Iq
382
383# NOTE: copied from sasmodels/resolution.py
384def bin_edges(x):
385    """
386    Determine bin edges from bin centers, assuming that edges are centered
387    between the bins.
388
389    Note: this uses the arithmetic mean, which may not be appropriate for
390    log-scaled data.
391    """
392    if len(x) < 2 or (np.diff(x) < 0).any():
393        raise ValueError("Expected bins to be an increasing set")
394    edges = np.hstack([
395        x[0]  - 0.5*(x[1]  - x[0]),  # first point minus half first interval
396        0.5*(x[1:] + x[:-1]),        # mid points of all central intervals
397        x[-1] + 0.5*(x[-1] - x[-2]), # last point plus half last interval
398        ])
399    return edges
400
401def plot_calc(r, Pr, q, Iq, theory=None):
402    import matplotlib.pyplot as plt
403    plt.subplot(211)
404    plt.plot(r, Pr, '-', label="Pr")
405    plt.xlabel('r (A)')
406    plt.ylabel('Pr (1/A^2)')
407    plt.subplot(212)
408    plt.loglog(q, Iq, '-', label='from Pr')
409    plt.xlabel('q (1/A')
410    plt.ylabel('Iq')
411    if theory is not None:
412        plt.loglog(theory[0], theory[1], '-', label='analytic')
413        plt.legend()
414
415def plot_calc_2d(qx, qy, Iqxy, theory=None):
416    import matplotlib.pyplot as plt
417    qx, qy = bin_edges(qx), bin_edges(qy)
418    #qx, qy = np.meshgrid(qx, qy)
419    if theory is not None:
420        plt.subplot(121)
421    plt.pcolormesh(qx, qy, np.log10(Iqxy))
422    plt.xlabel('qx (1/A)')
423    plt.ylabel('qy (1/A)')
424    if theory is not None:
425        plt.subplot(122)
426        plt.pcolormesh(qx, qy, np.log10(theory))
427        plt.xlabel('qx (1/A)')
428
429def plot_points(rho, points):
430    import mpl_toolkits.mplot3d
431    import matplotlib.pyplot as plt
432
433    ax = plt.axes(projection='3d')
434    try:
435        ax.axis('square')
436    except Exception:
437        pass
438    n = len(points)
439    #print("len points", n)
440    index = np.random.choice(n, size=500) if n > 500 else slice(None, None)
441    ax.scatter(points[index, 0], points[index, 1], points[index, 2], c=rho[index])
442    #low, high = points.min(axis=0), points.max(axis=0)
443    #ax.axis([low[0], high[0], low[1], high[1], low[2], high[2]])
444    ax.autoscale(True)
445
446def check_shape(shape, fn=None):
447    rho_solvent = 0
448    q = np.logspace(-3, 0, 200)
449    r = shape.r_bins(q, r_step=0.01)
450    sampling_density = 6*5000 / shape.volume()
451    rho, points = shape.sample(sampling_density)
452    t0 = time.time()
453    Pr = calc_Pr(r, rho-rho_solvent, points)
454    print("calc Pr time", time.time() - t0)
455    Iq = calc_Iq(q, r, Pr)
456    theory = (q, fn(q)) if fn is not None else None
457
458    import pylab
459    #plot_points(rho, points); pylab.figure()
460    plot_calc(r, Pr, q, Iq, theory=theory)
461    pylab.show()
462
463def check_shape_2d(shape, fn=None, view=(0, 0, 0)):
464    rho_solvent = 0
465    nq, qmax = 100, 1.0
466    qx = np.linspace(0.0, qmax, nq)
467    qy = np.linspace(0.0, qmax, nq)
468    Qx, Qy = np.meshgrid(qx, qy)
469    sampling_density = 50000 / shape.volume()
470    #t0 = time.time()
471    rho, points = shape.sample(sampling_density)
472    #print("sample time", time.time() - t0)
473    t0 = time.time()
474    Iqxy = calc_Iqxy(Qx, Qy, rho, points, view=view)
475    print("calc time", time.time() - t0)
476    theory = fn(Qx, Qy) if fn is not None else None
477    Iqxy += 0.001 * Iqxy.max()
478    if theory is not None:
479        theory += 0.001 * theory.max()
480
481    import pylab
482    #plot_points(rho, points); pylab.figure()
483    plot_calc_2d(qx, qy, Iqxy, theory=theory)
484    pylab.show()
485
486def sas_sinx_x(x):
487    with np.errstate(all='ignore'):
488        retvalue = sin(x)/x
489    retvalue[x == 0.] = 1.
490    return retvalue
491
492def sas_2J1x_x(x):
493    with np.errstate(all='ignore'):
494        retvalue = 2*J1(x)/x
495    retvalue[x == 0] = 1.
496    return retvalue
497
498def sas_3j1x_x(x):
499    """return 3*j1(x)/x"""
500    with np.errstate(all='ignore'):
501        retvalue = 3*(sin(x) - x*cos(x))/x**3
502    retvalue[x == 0.] = 1.
503    return retvalue
504
505def cylinder_Iq(q, radius, length):
506    z, w = leggauss(76)
507    cos_alpha = (z+1)/2
508    sin_alpha = sqrt(1.0 - cos_alpha**2)
509    Iq = np.empty_like(q)
510    for k, qk in enumerate(q):
511        qab, qc = qk*sin_alpha, qk*cos_alpha
512        Fq = sas_2J1x_x(qab*radius) * j0(qc*length/2)
513        Iq[k] = np.sum(w*Fq**2)
514    Iq = Iq/Iq[0]
515    return Iq
516
517def cylinder_Iqxy(qx, qy, radius, length, view=(0, 0, 0)):
518    qa, qb, qc = invert_view(qx, qy, view)
519    qab = np.sqrt(qa**2 + qb**2)
520    Fq = sas_2J1x_x(qab*radius) * j0(qc*length/2)
521    Iq = Fq**2
522    return Iq.reshape(qx.shape)
523
524def sphere_Iq(q, radius):
525    Iq = sas_3j1x_x(q*radius)**2
526    return Iq/Iq[0]
527
528def csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core):
529    z, w = leggauss(76)
530
531    sld_solvent = 0
532    overlapping = False
533    dr0 = sld_core - sld_solvent
534    drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent
535    tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc
536
537    outer_sum = np.zeros_like(q)
538    for cos_alpha, outer_w in zip((z+1)/2, w):
539        sin_alpha = sqrt(1.0-cos_alpha*cos_alpha)
540        qc = q*cos_alpha
541        siC = c*j0(c*qc/2)
542        siCt = tC*j0(tC*qc/2)
543        inner_sum = np.zeros_like(q)
544        for beta, inner_w in zip((z + 1)*pi/4, w):
545            qa, qb = q*sin_alpha*sin(beta), q*sin_alpha*cos(beta)
546            siA = a*j0(a*qa/2)
547            siB = b*j0(b*qb/2)
548            siAt = tA*j0(tA*qa/2)
549            siBt = tB*j0(tB*qb/2)
550            if overlapping:
551                Fq = (dr0*siA*siB*siC
552                      + drA*(siAt-siA)*siB*siC
553                      + drB*siAt*(siBt-siB)*siC
554                      + drC*siAt*siBt*(siCt-siC))
555            else:
556                Fq = (dr0*siA*siB*siC
557                      + drA*(siAt-siA)*siB*siC
558                      + drB*siA*(siBt-siB)*siC
559                      + drC*siA*siB*(siCt-siC))
560            inner_sum += inner_w * Fq**2
561        outer_sum += outer_w * inner_sum
562    Iq = outer_sum / 4  # = outer*um*zm*8.0/(4.0*M_PI)
563    return Iq/Iq[0]
564
565def csbox_Iqxy(qx, qy, a, b, c, da, db, dc, slda, sldb, sldc, sld_core, view=(0,0,0)):
566    qa, qb, qc = invert_view(qx, qy, view)
567
568    sld_solvent = 0
569    overlapping = False
570    dr0 = sld_core - sld_solvent
571    drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent
572    tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc
573    siA = a*sas_sinx_x(a*qa/2)
574    siB = b*sas_sinx_x(b*qb/2)
575    siC = c*sas_sinx_x(c*qc/2)
576    siAt = tA*sas_sinx_x(tA*qa/2)
577    siBt = tB*sas_sinx_x(tB*qb/2)
578    siCt = tC*sas_sinx_x(tC*qc/2)
579    Fq = (dr0*siA*siB*siC
580          + drA*(siAt-siA)*siB*siC
581          + drB*siA*(siBt-siB)*siC
582          + drC*siA*siB*(siCt-siC))
583    Iq = Fq**2
584    return Iq.reshape(qx.shape)
585
586def check_cylinder(radius=25, length=125, rho=2.):
587    shape = EllipticalCylinder(radius, radius, length, rho)
588    fn = lambda q: cylinder_Iq(q, radius, length)
589    check_shape(shape, fn)
590
591def check_cylinder_2d(radius=25, length=125, rho=2., view=(0, 0, 0)):
592    shape = EllipticalCylinder(radius, radius, length, rho)
593    fn = lambda qx, qy, view=view: cylinder_Iqxy(qx, qy, radius, length, view=view)
594    check_shape_2d(shape, fn, view=view)
595
596def check_cylinder_2d_lattice(radius=25, length=125, rho=2.,
597                              view=(0, 0, 0)):
598    nx, dx = 1, 2*radius
599    ny, dy = 30, 2*radius
600    nz, dz = 30, length
601    dx, dy, dz = 2*dx, 2*dy, 2*dz
602    def center(*args):
603        sigma = 0.333
604        space = 2
605        return [(space*n+np.random.randn()*sigma)*x for n, x in args]
606    shapes = [EllipticalCylinder(radius, radius, length, rho,
607                                 #center=(ix*dx, iy*dy, iz*dz)
608                                 orientation=np.random.randn(3)*0,
609                                 center=center((ix, dx), (iy, dy), (iz, dz))
610                                )
611              for ix in range(nx)
612              for iy in range(ny)
613              for iz in range(nz)]
614    shape = Composite(shapes)
615    fn = lambda qx, qy, view=view: cylinder_Iqxy(qx, qy, radius, length, view=view)
616    check_shape_2d(shape, fn, view=view)
617
618def check_sphere(radius=125, rho=2):
619    shape = TriaxialEllipsoid(radius, radius, radius, rho)
620    fn = lambda q: sphere_Iq(q, radius)
621    check_shape(shape, fn)
622
623def check_csbox(a=10, b=20, c=30, da=1, db=2, dc=3, slda=1, sldb=2, sldc=3, sld_core=4):
624    core = Box(a, b, c, sld_core)
625    side_a = Box(da, b, c, slda, center=((a+da)/2, 0, 0))
626    side_b = Box(a, db, c, sldb, center=(0, (b+db)/2, 0))
627    side_c = Box(a, b, dc, sldc, center=(0, 0, (c+dc)/2))
628    side_a2 = copy(side_a).shift(-a-da, 0, 0)
629    side_b2 = copy(side_b).shift(0, -b-db, 0)
630    side_c2 = copy(side_c).shift(0, 0, -c-dc)
631    shape = Composite((core, side_a, side_b, side_c, side_a2, side_b2, side_c2))
632    def fn(q):
633        return csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core)
634    #check_shape(shape, fn)
635
636    view = (20, 30, 40)
637    def fn_xy(qx, qy):
638        return csbox_Iqxy(qx, qy, a, b, c, da, db, dc,
639                          slda, sldb, sldc, sld_core, view=view)
640    check_shape_2d(shape, fn_xy, view=view)
641
642if __name__ == "__main__":
643    check_cylinder(radius=10, length=40)
644    #check_cylinder_2d(radius=10, length=40, view=(90,30,0))
645    #check_cylinder_2d_lattice(radius=10, length=50, view=(90,30,0))
646    #check_sphere()
647    #check_csbox()
648    #check_csbox(da=100, db=200, dc=300)
Note: See TracBrowser for help on using the repository browser.