source: sasmodels/explore/asymint.py @ 49eb251

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

add core_shell_parallelepiped to asymint (and fix parallelepiped)

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