source: sasmodels/explore/asymint.py @ 5dd7cfb

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

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