source: sasmodels/explore/asymint.py @ 20fe0cd

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

move paracrystal integration tests with the rest of the non-rotationally symmetric tests

  • Property mode set to 100755
File size: 11.6 KB
Line 
1#!/usr/bin/env python
2"""
3Asymmetric shape integration
4"""
5
6from __future__ import print_function, division
7
8import os, sys
9sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
10
11import numpy as np
12import mpmath as mp
13from numpy import pi, sin, cos, sqrt, exp, expm1, degrees, log10
14from numpy.polynomial.legendre import leggauss
15from scipy.integrate import dblquad, simps, romb, romberg
16import pylab
17
18import sasmodels.special as sp
19
20# Need to parse shape early since it determines the kernel function
21# that will be used for the various integrators
22shape = 'parallelepiped' if len(sys.argv) < 2 else sys.argv[1]
23Qstr = '0.005' if len(sys.argv) < 3 else sys.argv[2]
24
25class MPenv:
26    sqrt = mp.sqrt
27    exp = mp.exp
28    expm1 = mp.expm1
29    cos = mp.cos
30    sin = mp.sin
31    tan = mp.tan
32    def sas_3j1x_x(x):
33        return 3*(mp.sin(x)/x - mp.cos(x))/(x*x)
34    def sas_2J1x_x(x):
35        return 2*mp.j1(x)/x
36    def sas_sinx_x(x):
37        return mp.sin(x)/x
38    pi = mp.pi
39    mpf = mp.mpf
40
41class NPenv:
42    sqrt = np.sqrt
43    exp = np.exp
44    expm1 = np.expm1
45    cos = np.cos
46    sin = np.sin
47    tan = np.tan
48    sas_3j1x_x = sp.sas_3j1x_x
49    sas_2J1x_x = sp.sas_2J1x_x
50    sas_sinx_x = sp.sas_sinx_x
51    pi = np.pi
52    mpf = float
53
54SLD = 3
55SLD_SOLVENT = 6
56CONTRAST = SLD - SLD_SOLVENT
57
58# Carefully code models so that mpmath will use full precision.  That means:
59#    * wrap inputs in env.mpf
60#    * don't use floating point constants, only integers
61#    * for division, make sure the numerator or denominator is env.mpf
62#    * use env.pi, env.sas_sinx_x, etc. for functions
63def make_parallelepiped(a, b, c, env=NPenv):
64    a, b, c = env.mpf(a), env.mpf(b), env.mpf(c)
65    def Fq(qa, qb, qc):
66        siA = env.sas_sinx_x(0.5*a*qa/2)
67        siB = env.sas_sinx_x(0.5*b*qb/2)
68        siC = env.sas_sinx_x(0.5*c*qc/2)
69        return siA * siB * siC
70    Fq.__doc__ = "parallelepiped a=%g, b=%g c=%g"%(a, b, c)
71    volume = a*b*c
72    norm = CONTRAST**2*volume/10000
73    return norm, Fq
74
75def make_triaxial_ellipsoid(a, b, c, env=NPenv):
76    a, b, c = env.mpf(a), env.mpf(b), env.mpf(c)
77    def Fq(qa, qb, qc):
78        qr = env.sqrt((a*qa)**2 + (b*qb)**2 + (c*qc)**2)
79        return env.sas_3j1x_x(qr)
80    Fq.__doc__ = "triaxial ellipse minor=%g, major=%g polar=%g"%(a, b, c)
81    volume = 4*env.pi*a*b*c/3
82    norm = CONTRAST**2*volume/10000
83    return norm, Fq
84
85def make_cylinder(radius, length, env=NPenv):
86    radius, length = env.mpf(radius), env.mpf(length)
87    def Fq(qa, qb, qc):
88        qab = env.sqrt(qa**2 + qb**2)
89        return env.sas_2J1x_x(qab*radius) * env.sas_sinx_x((qc*length)/2)
90    Fq.__doc__ = "cylinder radius=%g, length=%g"%(radius, length)
91    volume = env.pi*radius**2*length
92    norm = CONTRAST**2*volume/10000
93    return norm, Fq
94
95def make_sphere(radius, env=NPenv):
96    radius = env.mpf(radius)
97    def Fq(qa, qb, qc):
98        q = env.sqrt(qa**2 + qb**2 + qc**2)
99        return env.sas_3j1x_x(q*radius)
100    Fq.__doc__ = "sphere radius=%g"%(radius, )
101    volume = 4*pi*radius**3
102    norm = CONTRAST**2*volume/10000
103    return norm, Fq
104
105def make_paracrystal(radius, dnn, d_factor, lattice='bcc', env=NPenv):
106    radius, dnn, d_factor = env.mpf(radius), env.mpf(dnn), env.mpf(d_factor)
107    def sc(qa, qb, qc):
108        return qa, qb, qc
109    def bcc(qa, qb, qc):
110        a1 = (+qa + qb + qc)/2
111        a2 = (-qa - qb + qc)/2
112        a3 = (-qa + qb - qc)/2
113        return a1, a2, a3
114    def fcc(qa, qb, qc):
115        a1 = ( 0. + qb + qc)/2
116        a2 = (-qa + 0. + qc)/2
117        a3 = (-qa + qb + 0.)/2
118        return a1, a2, a3
119    lattice_fn = {'sc': sc, 'bcc': bcc, 'fcc': fcc}[lattice]
120    radius, dnn, d_factor = env.mpf(radius), env.mpf(dnn), env.mpf(d_factor)
121    def Fq(qa, qb, qc):
122        a1, a2, a3 = lattice_fn(qa, qb, qc)
123        # Note: paper says that different directions can have different
124        # distoration factors.  Easy enough to add to the code.
125        # This would definitely break 8-fold symmetry.
126        arg = -(dnn*d_factor)**2*(a1**2 + a2**2 + a3**2)/2
127        exp_arg = env.exp(arg)
128        den = [((exp_arg - 2*env.cos(dnn*a))*exp_arg + 1) for a in (a1, a2, a3)]
129        Sq = -env.expm1(2*arg)**3/(den[0]*den[1]*den[2])
130
131        q = env.sqrt(qa**2 + qb**2 + qc**2)
132        Fq = env.sas_3j1x_x(q*radius)
133        # the kernel computes F(q)**2, but we need S(q)*F(q)**2
134        return env.sqrt(Sq)*Fq
135    Fq.__doc__ = "%s paracrystal a=%g da=%g r=%g"%(lattice, dnn, d_factor, radius)
136    def sphere_volume(r): return 4*env.pi*r**3/3
137    Vf = {
138        'sc': sphere_volume(radius/dnn),
139        'bcc': 2*sphere_volume(env.sqrt(3)/2*radius/dnn),
140        'fcc': 4*sphere_volume(radius/dnn/env.sqrt(2)),
141    }[lattice]
142    volume = sphere_volume(radius)
143    norm = CONTRAST**2*volume*Vf/10000
144    return norm, Fq
145
146if shape == 'sphere':
147    RADIUS = 50  # integer for the sake of mpf
148    NORM, KERNEL = make_sphere(radius=RADIUS)
149    NORM_MP, KERNEL_MP = make_sphere(radius=RADIUS, env=MPenv)
150elif shape == 'cylinder':
151    #RADIUS, LENGTH = 10, 100000
152    RADIUS, LENGTH = 10, 300  # integer for the sake of mpf
153    NORM, KERNEL = make_cylinder(radius=RADIUS, length=LENGTH)
154    NORM_MP, KERNEL_MP = make_cylinder(radius=RADIUS, length=LENGTH, env=MPenv)
155elif shape == 'triaxial_ellipsoid':
156    #A, B, C = 4450, 14000, 47
157    A, B, C = 445, 140, 47  # integer for the sake of mpf
158    NORM, KERNEL = make_triellip(A, B, C)
159    NORM_MP, KERNEL_MP = make_triellip(A, B, C, env=MPenv)
160elif shape == 'parallelepiped':
161    #A, B, C = 4450, 14000, 47
162    A, B, C = 445, 140, 47  # integer for the sake of mpf
163    NORM, KERNEL = make_parallelepiped(A, B, C)
164    NORM_MP, KERNEL_MP = make_parallelepiped(A, B, C, env=MPenv)
165elif shape == 'paracrystal':
166    LATTICE = 'bcc'
167    #LATTICE = 'fcc'
168    #LATTICE = 'sc'
169    DNN, D_FACTOR = 220, '0.06'  # mpmath needs to initialize floats from string
170    RADIUS = 40  # integer for the sake of mpf
171    NORM, KERNEL = make_paracrystal(
172        radius=RADIUS, dnn=DNN, d_factor=D_FACTOR, lattice=LATTICE)
173    NORM_MP, KERNEL_MP = make_paracrystal(
174        radius=RADIUS, dnn=DNN, d_factor=D_FACTOR, lattice=LATTICE, env=MPenv)
175else:
176    raise ValueError("Unknown shape %r"%shape)
177
178# Note: hardcoded in mp_quad
179THETA_LOW, THETA_HIGH = 0, pi
180PHI_LOW, PHI_HIGH = 0, 2*pi
181SCALE = 1
182
183# mathematica code for triaxial_ellipsoid (untested)
184_ = """
185R[theta_, phi_, a_, b_, c_] := Sqrt[(a Sin[theta]Cos[phi])^2 + (b Sin[theta]Sin[phi])^2 + (c Cos[theta])^2]
186Sphere[q_, r_] := 3 SphericalBesselJ[q r]/(q r)
187V[a_, b_, c_] := 4/3 pi a b c
188Norm[sld_, solvent_, a_, b_, c_] := V[a, b, c] (solvent - sld)^2
189F[q_, theta_, phi_, a_, b_, c_] := Sphere[q, R[theta, phi, a, b, c]]
190I[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}]
191I[6/10^3, 63/10, 3, 445, 140, 47]
192"""
193
194# 2D integration functions
195def mp_quad_2d(q, shape):
196    evals = [0]
197    def integrand(theta, phi):
198        evals[0] += 1
199        qab = q*mp.sin(theta)
200        qa = qab*mp.cos(phi)
201        qb = qab*mp.sin(phi)
202        qc = q*mp.cos(theta)
203        Zq = KERNEL_MP(qa, qb, qc)**2
204        return Zq*mp.sin(theta)
205    ans = mp.quad(integrand, (0, mp.pi), (0, 2*mp.pi))
206    Iq = NORM_MP*ans/(4*mp.pi)
207    return evals[0], Iq
208
209def kernel_2d(q, theta, phi):
210    """
211    S(q) kernel for paracrystal forms.
212    """
213    qab = q*sin(theta)
214    qa = qab*cos(phi)
215    qb = qab*sin(phi)
216    qc = q*cos(theta)
217    return NORM*KERNEL(qa, qb, qc)**2
218
219def scipy_dblquad_2d(q):
220    """
221    Compute the integral using scipy dblquad.  This gets the correct answer
222    eventually, but it is slow.
223    """
224    evals = [0]
225    def integrand(phi, theta):
226        evals[0] += 1
227        Zq = kernel_2d(q, theta=theta, phi=phi)
228        return Zq*sin(theta)
229    ans = dblquad(integrand, THETA_LOW, THETA_HIGH, lambda x: PHI_LOW, lambda x: PHI_HIGH)[0]
230    return evals[0], ans*SCALE/(4*pi)
231
232
233def scipy_romberg_2d(q):
234    """
235    Compute the integral using romberg integration.  This function does not
236    complete in a reasonable time.  No idea if it is accurate.
237    """
238    evals = [0]
239    def inner(phi, theta):
240        evals[0] += 1
241        return kernel_2d(q, theta=theta, phi=phi)
242    def outer(theta):
243        Zq = romberg(inner, PHI_LOW, PHI_HIGH, divmax=100, args=(theta,))
244        return Zq*sin(theta)
245    ans = romberg(outer, THETA_LOW, THETA_HIGH, divmax=100)
246    return evals[0], ans*SCALE/(4*pi)
247
248
249def semi_romberg_2d(q, n=100):
250    """
251    Use 1D romberg integration in phi and regular simpsons rule in theta.
252    """
253    evals = [0]
254    def inner(phi, theta):
255        evals[0] += 1
256        return kernel_2d(q, theta=theta, phi=phi)
257    theta = np.linspace(THETA_LOW, THETA_HIGH, n)
258    Zq = [romberg(inner, PHI_LOW, PHI_HIGH, divmax=100, args=(t,)) for t in theta]
259    ans = simps(np.array(Zq)*sin(theta), dx=theta[1]-theta[0])
260    return evals[0], ans*SCALE/(4*pi)
261
262def gauss_quad_2d(q, n=150):
263    """
264    Compute the integral using gaussian quadrature for n = 20, 76 or 150.
265    """
266    z, w = leggauss(n)
267    theta = (THETA_HIGH-THETA_LOW)*(z + 1)/2 + THETA_LOW
268    phi = (PHI_HIGH-PHI_LOW)*(z + 1)/2 + PHI_LOW
269    Atheta, Aphi = np.meshgrid(theta, phi)
270    Aw = w[None, :] * w[:, None]
271    sin_theta = abs(sin(Atheta))
272    Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi)
273    # change from [-1,1] x [-1,1] range to [0, pi] x [0, 2 pi] range
274    dxdy_stretch = (THETA_HIGH-THETA_LOW)/2 * (PHI_HIGH-PHI_LOW)/2
275    Iq = np.sum(Zq*Aw*sin_theta)*SCALE/(4*pi) * dxdy_stretch
276    return n**2, Iq
277
278def gridded_2d(q, n=300):
279    """
280    Compute the integral on a regular grid using rectangular, trapezoidal,
281    simpsons, and romberg integration.  Romberg integration requires that
282    the grid be of size n = 2**k + 1.
283    """
284    theta = np.linspace(THETA_LOW, THETA_HIGH, n)
285    phi = np.linspace(PHI_LOW, PHI_HIGH, n)
286    Atheta, Aphi = np.meshgrid(theta, phi)
287    Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi)
288    Zq *= abs(sin(Atheta))
289    dx, dy = theta[1]-theta[0], phi[1]-phi[0]
290    print("rect-%d"%n, n**2, np.sum(Zq)*dx*dy*SCALE/(4*pi))
291    print("trapz-%d"%n, n**2, np.trapz(np.trapz(Zq, dx=dx), dx=dy)*SCALE/(4*pi))
292    print("simpson-%d"%n, n**2, simps(simps(Zq, dx=dx), dx=dy)*SCALE/(4*pi))
293    print("romb-%d"%n, n**2, romb(romb(Zq, dx=dx), dx=dy)*SCALE/(4*pi))
294
295def plot_2d(q, n=300):
296    """
297    Plot the 2D surface that needs to be integrated in order to compute
298    the BCC S(q) at a particular q, dnn and d_factor.  *n* is the number
299    of points in the grid.
300    """
301    theta = np.linspace(THETA_LOW, THETA_HIGH, n)
302    phi = np.linspace(PHI_LOW, PHI_HIGH, n)
303    Atheta, Aphi = np.meshgrid(theta, phi)
304    Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi)
305    #Zq *= abs(sin(Atheta))
306    pylab.pcolor(degrees(theta), degrees(phi), log10(np.fmax(Zq, 1.e-6)))
307    pylab.axis('tight')
308    pylab.title("%s I(q,t) sin(t) for q=%g" % (KERNEL.__doc__, q))
309    pylab.xlabel("theta (degrees)")
310    pylab.ylabel("phi (degrees)")
311    cbar = pylab.colorbar()
312    cbar.set_label('log10 S(q)')
313    pylab.show()
314
315def main(Qstr):
316    Q = float(Qstr)
317    if shape == 'sphere':
318        print("exact", NORM*sas_3j1x_x(Q*RADIUS)**2)
319    print("gauss-20", *gauss_quad_2d(Q, n=20))
320    print("gauss-76", *gauss_quad_2d(Q, n=76))
321    print("gauss-150", *gauss_quad_2d(Q, n=150))
322    print("gauss-500", *gauss_quad_2d(Q, n=500))
323    #gridded_2d(Q, n=2**8+1)
324    gridded_2d(Q, n=2**10+1)
325    #gridded_2d(Q, n=2**13+1)
326    #gridded_2d(Q, n=2**15+1)
327    if shape != 'paracrystal':  # adaptive forms are too slow!
328        print("dblquad", *scipy_dblquad_2d(Q))
329        print("semi-romberg-100", *semi_romberg_2d(Q, n=100))
330        print("romberg", *scipy_romberg_2d(Q))
331        with mp.workprec(100):
332            print("mpmath", *mp_quad_2d(mp.mpf(Qstr), shape))
333    plot_2d(Q, n=200)
334
335if __name__ == "__main__":
336    main(Qstr)
Note: See TracBrowser for help on using the repository browser.