source: sasmodels/explore/asymint.py @ ac995be

Last change on this file since ac995be was c462169, checked in by Paul Kienzle <pkienzle@…>, 7 years ago

show u-sub integration as used in models in explore/asymint.py
  • Property mode set to 100755
File size: 16.1 KB
Line 
1#!/usr/bin/env python
2"""
3Asymmetric shape integration
4
5Usage:
6
7    explore/asymint.py [MODEL] [q-value]
8
9Computes the numerical integral over theta and phi of the given model at a
10single point q using different algorithms or the same algorithm with different
11precision.  It also displays a 2-D image of the theta-phi surface that is
12being integrated.
13
14The available models are:
15
16    triaxial_ellipsoid, parallelpiped, paracrystal, cylinder, sphere
17
18Cylinder and sphere are included as simple checks on the integration
19algorithms. Cylinder is better investigated using 1-D integration methods in
20explore/symint.py.  Sphere has an easily computed analytic value which is
21identical for all theta-phi for a given q, so it is useful for checking
22that the normalization constants are correct for the different algorithms.
23"""
24
25from __future__ import print_function, division
26
27import os, sys
28sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
29
30import numpy as np
31import mpmath as mp
32from numpy import pi, sin, cos, sqrt, exp, expm1, degrees, log10, arccos
33from numpy.polynomial.legendre import leggauss
34from scipy.integrate import dblquad, simps, romb, romberg
35import pylab
36
37import sasmodels.special as sp
38
39# Need to parse shape early since it determines the kernel function
40# that will be used for the various integrators
41shape = 'parallelepiped' if len(sys.argv) < 2 else sys.argv[1]
42Qstr = '0.005' if len(sys.argv) < 3 else sys.argv[2]
43
44DTYPE = 'd'
45
46class MPenv:
47    sqrt = staticmethod(mp.sqrt)
48    exp = staticmethod(mp.exp)
49    expm1 = staticmethod(mp.expm1)
50    cos = staticmethod(mp.cos)
51    sin = staticmethod(mp.sin)
52    tan = staticmethod(mp.tan)
53    @staticmethod
54    def sas_3j1x_x(x):
55        return 3*(mp.sin(x)/x - mp.cos(x))/(x*x)
56    @staticmethod
57    def sas_2J1x_x(x):
58        return 2*mp.j1(x)/x
59    @staticmethod
60    def sas_sinx_x(x):
61        return mp.sin(x)/x
62    pi = mp.pi
63    mpf = staticmethod(mp.mpf)
64
65class NPenv:
66    sqrt = staticmethod(np.sqrt)
67    exp = staticmethod(np.exp)
68    expm1 = staticmethod(np.expm1)
69    cos = staticmethod(np.cos)
70    sin = staticmethod(np.sin)
71    tan = staticmethod(np.tan)
72    sas_3j1x_x = staticmethod(sp.sas_3j1x_x)
73    sas_2J1x_x = staticmethod(sp.sas_2J1x_x)
74    sas_sinx_x = staticmethod(sp.sas_sinx_x)
75    pi = np.pi
76    #mpf = staticmethod(float)
77    mpf = staticmethod(lambda x: np.array(x, DTYPE))
78
79SLD = 3
80SLD_SOLVENT = 6
81CONTRAST = SLD - SLD_SOLVENT
82
83# Carefully code models so that mpmath will use full precision.  That means:
84#    * wrap inputs in env.mpf
85#    * don't use floating point constants, only integers
86#    * for division, make sure the numerator or denominator is env.mpf
87#    * use env.pi, env.sas_sinx_x, etc. for functions
88def make_parallelepiped(a, b, c, env=NPenv):
89    a, b, c = env.mpf(a), env.mpf(b), env.mpf(c)
90    def Fq(qa, qb, qc):
91        siA = env.sas_sinx_x(a*qa/2)
92        siB = env.sas_sinx_x(b*qb/2)
93        siC = env.sas_sinx_x(c*qc/2)
94        return siA * siB * siC
95    Fq.__doc__ = "parallelepiped a=%g, b=%g c=%g"%(a, b, c)
96    volume = a*b*c
97    norm = CONTRAST**2*volume/10000
98    return norm, Fq
99
100def make_core_shell_parallelepiped(a, b, c, da, db, dc, slda, sldb, sldc, env=NPenv):
101    overlapping = False
102    a, b, c = env.mpf(a), env.mpf(b), env.mpf(c)
103    da, db, dc = env.mpf(da), env.mpf(db), env.mpf(dc)
104    slda, sldb, sldc = env.mpf(slda), env.mpf(sldb), env.mpf(sldc)
105    dr0 = CONTRAST
106    drA, drB, drC = slda-SLD_SOLVENT, sldb-SLD_SOLVENT, sldc-SLD_SOLVENT
107    tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc
108    def Fq(qa, qb, qc):
109        siA = a*env.sas_sinx_x(a*qa/2)
110        siB = b*env.sas_sinx_x(b*qb/2)
111        siC = c*env.sas_sinx_x(c*qc/2)
112        siAt = tA*env.sas_sinx_x(tA*qa/2)
113        siBt = tB*env.sas_sinx_x(tB*qb/2)
114        siCt = tC*env.sas_sinx_x(tC*qc/2)
115        if overlapping:
116            return (dr0*siA*siB*siC
117                    + drA*(siAt-siA)*siB*siC
118                    + drB*siAt*(siBt-siB)*siC
119                    + drC*siAt*siBt*(siCt-siC))
120        else:
121            return (dr0*siA*siB*siC
122                    + drA*(siAt-siA)*siB*siC
123                    + drB*siA*(siBt-siB)*siC
124                    + drC*siA*siB*(siCt-siC))
125    Fq.__doc__ = "core-shell parallelepiped a=%g, b=%g c=%g"%(a, b, c)
126    if overlapping:
127        volume = a*b*c + 2*da*b*c + 2*tA*db*c + 2*tA*tB*dc
128    else:
129        volume = a*b*c + 2*da*b*c + 2*a*db*c + 2*a*b*dc
130    norm = 1/(volume*10000)
131    return norm, Fq
132
133def make_triaxial_ellipsoid(a, b, c, env=NPenv):
134    a, b, c = env.mpf(a), env.mpf(b), env.mpf(c)
135    def Fq(qa, qb, qc):
136        qr = env.sqrt((a*qa)**2 + (b*qb)**2 + (c*qc)**2)
137        return env.sas_3j1x_x(qr)
138    Fq.__doc__ = "triaxial ellipsoid minor=%g, major=%g polar=%g"%(a, b, c)
139    volume = 4*env.pi*a*b*c/3
140    norm = CONTRAST**2*volume/10000
141    return norm, Fq
142
143def make_cylinder(radius, length, env=NPenv):
144    radius, length = env.mpf(radius), env.mpf(length)
145    def Fq(qa, qb, qc):
146        qab = env.sqrt(qa**2 + qb**2)
147        return env.sas_2J1x_x(qab*radius) * env.sas_sinx_x((qc*length)/2)
148    Fq.__doc__ = "cylinder radius=%g, length=%g"%(radius, length)
149    volume = env.pi*radius**2*length
150    norm = CONTRAST**2*volume/10000
151    return norm, Fq
152
153def make_sphere(radius, env=NPenv):
154    radius = env.mpf(radius)
155    def Fq(qa, qb, qc):
156        q = env.sqrt(qa**2 + qb**2 + qc**2)
157        return env.sas_3j1x_x(q*radius)
158    Fq.__doc__ = "sphere radius=%g"%(radius, )
159    volume = 4*pi*radius**3
160    norm = CONTRAST**2*volume/10000
161    return norm, Fq
162
163def make_paracrystal(radius, dnn, d_factor, lattice='bcc', env=NPenv):
164    radius, dnn, d_factor = env.mpf(radius), env.mpf(dnn), env.mpf(d_factor)
165    def sc(qa, qb, qc):
166        return qa, qb, qc
167    def bcc(qa, qb, qc):
168        a1 = (+qa + qb + qc)/2
169        a2 = (-qa - qb + qc)/2
170        a3 = (-qa + qb - qc)/2
171        return a1, a2, a3
172    def fcc(qa, qb, qc):
173        a1 = ( 0 + qb + qc)/2
174        a2 = (-qa + 0 + qc)/2
175        a3 = (-qa + qb + 0)/2
176        return a1, a2, a3
177    lattice_fn = {'sc': sc, 'bcc': bcc, 'fcc': fcc}[lattice]
178    radius, dnn, d_factor = env.mpf(radius), env.mpf(dnn), env.mpf(d_factor)
179    def Fq(qa, qb, qc):
180        a1, a2, a3 = lattice_fn(qa, qb, qc)
181        # Note: paper says that different directions can have different
182        # distoration factors.  Easy enough to add to the code.
183        arg = -(dnn*d_factor)**2*(a1**2 + a2**2 + a3**2)/2
184        exp_arg = env.exp(arg)
185        den = [((exp_arg - 2*env.cos(dnn*a))*exp_arg + 1) for a in (a1, a2, a3)]
186        Sq = -env.expm1(2*arg)**3/(den[0]*den[1]*den[2])
187
188        q = env.sqrt(qa**2 + qb**2 + qc**2)
189        Fq = env.sas_3j1x_x(q*radius)
190        # the caller computes F(q)**2, but we need it to compute S(q)*F(q)**2
191        return env.sqrt(Sq)*Fq
192    Fq.__doc__ = "%s paracrystal a=%g da=%g r=%g"%(lattice, dnn, d_factor, radius)
193    def sphere_volume(r): return 4*env.pi*r**3/3
194    Vf = {
195        'sc': sphere_volume(radius/dnn),
196        'bcc': 2*sphere_volume(env.sqrt(3)/2*radius/dnn),
197        'fcc': 4*sphere_volume(1/env.sqrt(2)*radius/dnn),
198    }[lattice]
199    volume = sphere_volume(radius)
200    norm = CONTRAST**2*volume/10000*Vf
201    return norm, Fq
202
203if shape == 'sphere':
204    RADIUS = 50  # integer for the sake of mpf
205    NORM, KERNEL = make_sphere(radius=RADIUS)
206    NORM_MP, KERNEL_MP = make_sphere(radius=RADIUS, env=MPenv)
207elif shape == 'cylinder':
208    #RADIUS, LENGTH = 10, 100000
209    RADIUS, LENGTH = 10, 300  # integer for the sake of mpf
210    NORM, KERNEL = make_cylinder(radius=RADIUS, length=LENGTH)
211    NORM_MP, KERNEL_MP = make_cylinder(radius=RADIUS, length=LENGTH, env=MPenv)
212elif shape == 'triaxial_ellipsoid':
213    #A, B, C = 4450, 14000, 47
214    A, B, C = 445, 140, 47  # integer for the sake of mpf
215    NORM, KERNEL = make_triaxial_ellipsoid(A, B, C)
216    NORM_MP, KERNEL_MP = make_triaxial_ellipsoid(A, B, C, env=MPenv)
217elif shape == 'parallelepiped':
218    #A, B, C = 4450, 14000, 47
219    A, B, C = 445, 140, 47  # integer for the sake of mpf
220    NORM, KERNEL = make_parallelepiped(A, B, C)
221    NORM_MP, KERNEL_MP = make_parallelepiped(A, B, C, env=MPenv)
222elif shape == 'core_shell_parallelepiped':
223    #A, B, C = 4450, 14000, 47
224    #A, B, C = 445, 140, 47  # integer for the sake of mpf
225    A, B, C = 114, 1380, 6800
226    DA, DB, DC = 21, 58, 2300
227    SLDA, SLDB, SLDC = "5", "-0.3", "11.5"
228    ## default parameters from sasmodels
229    #A,B,C,DA,DB,DC,SLDA,SLDB,SLDC = 400,75,35,10,10,10,2,4,2
230    ## swap A-B-C to C-B-A
231    #A, B, C, DA, DB, DC, SLDA, SLDB, SLDC = C, B, A, DC, DB, DA, SLDC, SLDB, SLDA
232    #A,B,C,DA,DB,DC,SLDA,SLDB,SLDC = 10,20,30,100,200,300,1,2,3
233    #SLD_SOLVENT,CONTRAST = 0, 4
234    if 1: # C shortest
235        B, C = C, B
236        DB, DC = DC, DB
237        SLDB, SLDC = SLDC, SLDB
238    elif 0: # C longest
239        A, C = C, A
240        DA, DC = DC, DA
241        SLDA, SLDC = SLDC, SLDA
242    #NORM, KERNEL = make_core_shell_parallelepiped(A, B, C, DA, DB, DC, SLDA, SLDB, SLDC)
243    NORM, KERNEL = make_core_shell_parallelepiped(A, B, C, DA, DB, DC, SLDA, SLDB, SLDC)
244    NORM_MP, KERNEL_MP = make_core_shell_parallelepiped(A, B, C, DA, DB, DC, SLDA, SLDB, SLDC, env=MPenv)
245elif shape == 'paracrystal':
246    LATTICE = 'bcc'
247    #LATTICE = 'fcc'
248    #LATTICE = 'sc'
249    DNN, D_FACTOR = 220, '0.06'  # mpmath needs to initialize floats from string
250    RADIUS = 40  # integer for the sake of mpf
251    NORM, KERNEL = make_paracrystal(
252        radius=RADIUS, dnn=DNN, d_factor=D_FACTOR, lattice=LATTICE)
253    NORM_MP, KERNEL_MP = make_paracrystal(
254        radius=RADIUS, dnn=DNN, d_factor=D_FACTOR, lattice=LATTICE, env=MPenv)
255else:
256    raise ValueError("Unknown shape %r"%shape)
257
258# Note: hardcoded in mp_quad
259THETA_LOW, THETA_HIGH = 0, pi
260PHI_LOW, PHI_HIGH = 0, 2*pi
261SCALE = 1
262
263# mathematica code for triaxial_ellipsoid (untested)
264_ = """
265R[theta_, phi_, a_, b_, c_] := Sqrt[(a Sin[theta]Cos[phi])^2 + (b Sin[theta]Sin[phi])^2 + (c Cos[theta])^2]
266Sphere[q_, r_] := 3 SphericalBesselJ[q r]/(q r)
267V[a_, b_, c_] := 4/3 pi a b c
268Norm[sld_, solvent_, a_, b_, c_] := V[a, b, c] (solvent - sld)^2
269F[q_, theta_, phi_, a_, b_, c_] := Sphere[q, R[theta, phi, a, b, c]]
270I[q_, sld_, solvent_, a_, b_, c_] := Norm[sld, solvent, a, b, c]/(4 pi) Integrate[F[q, theta, phi, a, b, c]^2 Sin[theta], {phi, 0, 2 pi}, {theta, 0, pi}]
271I[6/10^3, 63/10, 3, 445, 140, 47]
272"""
273
274# 2D integration functions
275def mp_quad_2d(q, shape):
276    evals = [0]
277    def integrand(theta, phi):
278        evals[0] += 1
279        qab = q*mp.sin(theta)
280        qa = qab*mp.cos(phi)
281        qb = qab*mp.sin(phi)
282        qc = q*mp.cos(theta)
283        Zq = KERNEL_MP(qa, qb, qc)**2
284        return Zq*mp.sin(theta)
285    ans = mp.quad(integrand, (0, mp.pi), (0, 2*mp.pi))
286    Iq = NORM_MP*ans/(4*mp.pi)
287    return evals[0], Iq
288
289def kernel_2d(q, theta, phi):
290    """
291    S(q) kernel for paracrystal forms.
292    """
293    qab = q*sin(theta)
294    qa = qab*cos(phi)
295    qb = qab*sin(phi)
296    qc = q*cos(theta)
297    return NORM*KERNEL(qa, qb, qc)**2
298
299def scipy_dblquad_2d(q):
300    """
301    Compute the integral using scipy dblquad.  This gets the correct answer
302    eventually, but it is slow.
303    """
304    evals = [0]
305    def integrand(phi, theta):
306        evals[0] += 1
307        Zq = kernel_2d(q, theta=theta, phi=phi)
308        return Zq*sin(theta)
309    ans = dblquad(integrand, THETA_LOW, THETA_HIGH, lambda x: PHI_LOW, lambda x: PHI_HIGH)[0]
310    return evals[0], ans*SCALE/(4*pi)
311
312
313def scipy_romberg_2d(q):
314    """
315    Compute the integral using romberg integration.  This function does not
316    complete in a reasonable time.  No idea if it is accurate.
317    """
318    evals = [0]
319    def inner(phi, theta):
320        evals[0] += 1
321        return kernel_2d(q, theta=theta, phi=phi)
322    def outer(theta):
323        Zq = romberg(inner, PHI_LOW, PHI_HIGH, divmax=100, args=(theta,))
324        return Zq*sin(theta)
325    ans = romberg(outer, THETA_LOW, THETA_HIGH, divmax=100)
326    return evals[0], ans*SCALE/(4*pi)
327
328
329def semi_romberg_2d(q, n=100):
330    """
331    Use 1D romberg integration in phi and regular simpsons rule in theta.
332    """
333    evals = [0]
334    def inner(phi, theta):
335        evals[0] += 1
336        return kernel_2d(q, theta=theta, phi=phi)
337    theta = np.linspace(THETA_LOW, THETA_HIGH, n)
338    Zq = [romberg(inner, PHI_LOW, PHI_HIGH, divmax=100, args=(t,)) for t in theta]
339    ans = simps(np.array(Zq)*sin(theta), dx=theta[1]-theta[0])
340    return evals[0], ans*SCALE/(4*pi)
341
342def gauss_quad_2d(q, n=150):
343    """
344    Compute the integral using gaussian quadrature for n = 20, 76 or 150.
345    """
346    z, w = leggauss(n)
347    theta = (THETA_HIGH-THETA_LOW)*(z + 1)/2 + THETA_LOW
348    phi = (PHI_HIGH-PHI_LOW)*(z + 1)/2 + PHI_LOW
349    Atheta, Aphi = np.meshgrid(theta, phi)
350    Aw = w[None, :] * w[:, None]
351    sin_theta = abs(sin(Atheta))
352    Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi)
353    # change from [-1,1] x [-1,1] range to [0, pi] x [0, 2 pi] range
354    dxdy_stretch = (THETA_HIGH-THETA_LOW)/2 * (PHI_HIGH-PHI_LOW)/2
355    Iq = np.sum(Zq*Aw*sin_theta)*SCALE/(4*pi) * dxdy_stretch
356    return n**2, Iq
357
358def gauss_quad_usub(q, n=150, dtype=DTYPE):
359    """
360    Compute the integral using gaussian quadrature for n = 20, 76 or 150.
361
362    Use *u = sin theta* substitution, and restrict integration over a single
363    quadrant for shapes that are mirror symmetric about AB, AC and BC planes.
364
365    Note that this doesn't work for fcc/bcc paracrystals, which instead step
366    over the entire 4 pi surface uniformly in theta-phi.
367    """
368    z, w = leggauss(n)
369    cos_theta = 0.5 * (z + 1)
370    theta = arccos(cos_theta)
371    phi = pi/2*(0.5 * (z + 1))
372    Atheta, Aphi = np.meshgrid(theta, phi)
373    Aw = w[None, :] * w[:, None]
374    q, Atheta, Aphi, Aw = [np.asarray(v, dtype=dtype) for v in (q, Atheta, Aphi, Aw)]
375    Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi)
376    Iq = np.sum(Zq*Aw)*0.25
377    return n**2, Iq
378
379def gridded_2d(q, n=300):
380    """
381    Compute the integral on a regular grid using rectangular, trapezoidal,
382    simpsons, and romberg integration.  Romberg integration requires that
383    the grid be of size n = 2**k + 1.
384    """
385    theta = np.linspace(THETA_LOW, THETA_HIGH, n)
386    phi = np.linspace(PHI_LOW, PHI_HIGH, n)
387    Atheta, Aphi = np.meshgrid(theta, phi)
388    Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi)
389    Zq *= abs(sin(Atheta))
390    dx, dy = theta[1]-theta[0], phi[1]-phi[0]
391    print("rect-%d"%n, n**2, np.sum(Zq)*dx*dy*SCALE/(4*pi))
392    print("trapz-%d"%n, n**2, np.trapz(np.trapz(Zq, dx=dx), dx=dy)*SCALE/(4*pi))
393    print("simpson-%d"%n, n**2, simps(simps(Zq, dx=dx), dx=dy)*SCALE/(4*pi))
394    print("romb-%d"%n, n**2, romb(romb(Zq, dx=dx), dx=dy)*SCALE/(4*pi))
395
396def plot_2d(q, n=300):
397    """
398    Plot the 2D surface that needs to be integrated in order to compute
399    the BCC S(q) at a particular q, dnn and d_factor.  *n* is the number
400    of points in the grid.
401    """
402    theta = np.linspace(THETA_LOW, THETA_HIGH, n)
403    phi = np.linspace(PHI_LOW, PHI_HIGH, n)
404    Atheta, Aphi = np.meshgrid(theta, phi)
405    Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi)
406    #Zq *= abs(sin(Atheta))
407    pylab.pcolor(degrees(theta), degrees(phi), log10(np.fmax(Zq, 1.e-6)))
408    pylab.axis('tight')
409    pylab.title("%s I(q,t) sin(t) for q=%g" % (KERNEL.__doc__, q))
410    pylab.xlabel("theta (degrees)")
411    pylab.ylabel("phi (degrees)")
412    cbar = pylab.colorbar()
413    cbar.set_label('log10 S(q)')
414    pylab.show()
415
416def main(Qstr):
417    Q = float(Qstr)
418    if shape == 'sphere':
419        print("exact", NORM*sp.sas_3j1x_x(Q*RADIUS)**2)
420    print("gauss-20", *gauss_quad_2d(Q, n=20))
421    print("gauss-76", *gauss_quad_2d(Q, n=76))
422    print("gauss-150", *gauss_quad_2d(Q, n=150))
423    print("gauss-500", *gauss_quad_2d(Q, n=500))
424    print("gauss-1025", *gauss_quad_2d(Q, n=1025))
425    print("gauss-2049", *gauss_quad_2d(Q, n=2049))
426    print("gauss-20 usub", *gauss_quad_usub(Q, n=20))
427    print("gauss-76 usub", *gauss_quad_usub(Q, n=76))
428    print("gauss-150 usub", *gauss_quad_usub(Q, n=150))
429    #gridded_2d(Q, n=2**8+1)
430    gridded_2d(Q, n=2**10+1)
431    #gridded_2d(Q, n=2**12+1)
432    #gridded_2d(Q, n=2**15+1)
433    if shape not in ('paracrystal', 'core_shell_parallelepiped'):
434        # adaptive forms on models for which the calculations are fast enough
435        print("dblquad", *scipy_dblquad_2d(Q))
436        print("semi-romberg-100", *semi_romberg_2d(Q, n=100))
437        print("romberg", *scipy_romberg_2d(Q))
438        with mp.workprec(100):
439            print("mpmath", *mp_quad_2d(mp.mpf(Qstr), shape))
440    plot_2d(Q, n=200)
441
442if __name__ == "__main__":
443    main(Qstr)
Note: See TracBrowser for help on using the repository browser.