source: sasmodels/explore/asymint.py @ 1820208

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

asymint: fix sphere exact calc, and triaxial ellipsoid

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