[20fe0cd] | 1 | #!/usr/bin/env python |
---|
[4f611f1] | 2 | """ |
---|
| 3 | Asymmetric shape integration |
---|
| 4 | """ |
---|
| 5 | |
---|
| 6 | from __future__ import print_function, division |
---|
| 7 | |
---|
| 8 | import os, sys |
---|
[20fe0cd] | 9 | sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) |
---|
[4f611f1] | 10 | |
---|
| 11 | import numpy as np |
---|
[31eea1f] | 12 | import mpmath as mp |
---|
[4f611f1] | 13 | from numpy import pi, sin, cos, sqrt, exp, expm1, degrees, log10 |
---|
[31eea1f] | 14 | from numpy.polynomial.legendre import leggauss |
---|
[4f611f1] | 15 | from scipy.integrate import dblquad, simps, romb, romberg |
---|
| 16 | import pylab |
---|
| 17 | |
---|
[20fe0cd] | 18 | import sasmodels.special as sp |
---|
[4f611f1] | 19 | |
---|
[20fe0cd] | 20 | # Need to parse shape early since it determines the kernel function |
---|
| 21 | # that will be used for the various integrators |
---|
| 22 | shape = 'parallelepiped' if len(sys.argv) < 2 else sys.argv[1] |
---|
| 23 | Qstr = '0.005' if len(sys.argv) < 3 else sys.argv[2] |
---|
| 24 | |
---|
| 25 | class 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 | |
---|
| 41 | class 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 |
---|
[31eea1f] | 53 | |
---|
| 54 | SLD = 3 |
---|
| 55 | SLD_SOLVENT = 6 |
---|
[4f611f1] | 56 | CONTRAST = SLD - SLD_SOLVENT |
---|
| 57 | |
---|
[20fe0cd] | 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 |
---|
| 63 | def make_parallelepiped(a, b, c, env=NPenv): |
---|
| 64 | a, b, c = env.mpf(a), env.mpf(b), env.mpf(c) |
---|
[31eea1f] | 65 | def Fq(qa, qb, qc): |
---|
[20fe0cd] | 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) |
---|
[31eea1f] | 69 | return siA * siB * siC |
---|
[20fe0cd] | 70 | Fq.__doc__ = "parallelepiped a=%g, b=%g c=%g"%(a, b, c) |
---|
[31eea1f] | 71 | volume = a*b*c |
---|
[20fe0cd] | 72 | norm = CONTRAST**2*volume/10000 |
---|
[31eea1f] | 73 | return norm, Fq |
---|
| 74 | |
---|
[20fe0cd] | 75 | def make_triaxial_ellipsoid(a, b, c, env=NPenv): |
---|
| 76 | a, b, c = env.mpf(a), env.mpf(b), env.mpf(c) |
---|
[31eea1f] | 77 | def Fq(qa, qb, qc): |
---|
[20fe0cd] | 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 |
---|
[31eea1f] | 83 | return norm, Fq |
---|
| 84 | |
---|
[20fe0cd] | 85 | def make_cylinder(radius, length, env=NPenv): |
---|
| 86 | radius, length = env.mpf(radius), env.mpf(length) |
---|
[31eea1f] | 87 | def Fq(qa, qb, qc): |
---|
[20fe0cd] | 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 |
---|
[31eea1f] | 93 | return norm, Fq |
---|
| 94 | |
---|
[20fe0cd] | 95 | def make_sphere(radius, env=NPenv): |
---|
| 96 | radius = env.mpf(radius) |
---|
[31eea1f] | 97 | def Fq(qa, qb, qc): |
---|
[20fe0cd] | 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 |
---|
[31eea1f] | 103 | return norm, Fq |
---|
| 104 | |
---|
[20fe0cd] | 105 | def 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) |
---|
[31eea1f] | 121 | def Fq(qa, qb, qc): |
---|
[20fe0cd] | 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]) |
---|
[31eea1f] | 130 | |
---|
[20fe0cd] | 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 |
---|
[31eea1f] | 144 | return norm, Fq |
---|
[4f611f1] | 145 | |
---|
[20fe0cd] | 146 | if 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) |
---|
| 150 | elif shape == 'cylinder': |
---|
[31eea1f] | 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) |
---|
[20fe0cd] | 154 | NORM_MP, KERNEL_MP = make_cylinder(radius=RADIUS, length=LENGTH, env=MPenv) |
---|
| 155 | elif shape == 'triaxial_ellipsoid': |
---|
[31eea1f] | 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) |
---|
[20fe0cd] | 159 | NORM_MP, KERNEL_MP = make_triellip(A, B, C, env=MPenv) |
---|
[31eea1f] | 160 | elif 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) |
---|
[20fe0cd] | 164 | NORM_MP, KERNEL_MP = make_parallelepiped(A, B, C, env=MPenv) |
---|
| 165 | elif 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) |
---|
[31eea1f] | 175 | else: |
---|
| 176 | raise ValueError("Unknown shape %r"%shape) |
---|
[4f611f1] | 177 | |
---|
[20fe0cd] | 178 | # Note: hardcoded in mp_quad |
---|
[4f611f1] | 179 | THETA_LOW, THETA_HIGH = 0, pi |
---|
| 180 | PHI_LOW, PHI_HIGH = 0, 2*pi |
---|
| 181 | SCALE = 1 |
---|
| 182 | |
---|
[31eea1f] | 183 | # mathematica code for triaxial_ellipsoid (untested) |
---|
| 184 | _ = """ |
---|
| 185 | R[theta_, phi_, a_, b_, c_] := Sqrt[(a Sin[theta]Cos[phi])^2 + (b Sin[theta]Sin[phi])^2 + (c Cos[theta])^2] |
---|
| 186 | Sphere[q_, r_] := 3 SphericalBesselJ[q r]/(q r) |
---|
| 187 | V[a_, b_, c_] := 4/3 pi a b c |
---|
| 188 | Norm[sld_, solvent_, a_, b_, c_] := V[a, b, c] (solvent - sld)^2 |
---|
| 189 | F[q_, theta_, phi_, a_, b_, c_] := Sphere[q, R[theta, phi, a, b, c]] |
---|
| 190 | I[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}] |
---|
| 191 | I[6/10^3, 63/10, 3, 445, 140, 47] |
---|
| 192 | """ |
---|
| 193 | |
---|
[20fe0cd] | 194 | # 2D integration functions |
---|
| 195 | def mp_quad_2d(q, shape): |
---|
[31eea1f] | 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 |
---|
[4f611f1] | 208 | |
---|
[20fe0cd] | 209 | def kernel_2d(q, theta, phi): |
---|
[4f611f1] | 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 | |
---|
[20fe0cd] | 219 | def scipy_dblquad_2d(q): |
---|
[4f611f1] | 220 | """ |
---|
| 221 | Compute the integral using scipy dblquad. This gets the correct answer |
---|
| 222 | eventually, but it is slow. |
---|
| 223 | """ |
---|
| 224 | evals = [0] |
---|
[31eea1f] | 225 | def integrand(phi, theta): |
---|
[4f611f1] | 226 | evals[0] += 1 |
---|
[20fe0cd] | 227 | Zq = kernel_2d(q, theta=theta, phi=phi) |
---|
[4f611f1] | 228 | return Zq*sin(theta) |
---|
[31eea1f] | 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) |
---|
[4f611f1] | 231 | |
---|
| 232 | |
---|
| 233 | def 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 | """ |
---|
[31eea1f] | 238 | evals = [0] |
---|
[4f611f1] | 239 | def inner(phi, theta): |
---|
[31eea1f] | 240 | evals[0] += 1 |
---|
[20fe0cd] | 241 | return kernel_2d(q, theta=theta, phi=phi) |
---|
[4f611f1] | 242 | def outer(theta): |
---|
[31eea1f] | 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) |
---|
[4f611f1] | 247 | |
---|
| 248 | |
---|
[20fe0cd] | 249 | def semi_romberg_2d(q, n=100): |
---|
[4f611f1] | 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 |
---|
[20fe0cd] | 256 | return kernel_2d(q, theta=theta, phi=phi) |
---|
[4f611f1] | 257 | theta = np.linspace(THETA_LOW, THETA_HIGH, n) |
---|
[31eea1f] | 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) |
---|
[4f611f1] | 261 | |
---|
[20fe0cd] | 262 | def gauss_quad_2d(q, n=150): |
---|
[4f611f1] | 263 | """ |
---|
| 264 | Compute the integral using gaussian quadrature for n = 20, 76 or 150. |
---|
| 265 | """ |
---|
[20fe0cd] | 266 | z, w = leggauss(n) |
---|
[4f611f1] | 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] |
---|
[31eea1f] | 271 | sin_theta = abs(sin(Atheta)) |
---|
[20fe0cd] | 272 | Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi) |
---|
[31eea1f] | 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 |
---|
[4f611f1] | 277 | |
---|
[20fe0cd] | 278 | def gridded_2d(q, n=300): |
---|
[4f611f1] | 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) |
---|
[20fe0cd] | 287 | Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi) |
---|
[4f611f1] | 288 | Zq *= abs(sin(Atheta)) |
---|
| 289 | dx, dy = theta[1]-theta[0], phi[1]-phi[0] |
---|
[31eea1f] | 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)) |
---|
[4f611f1] | 294 | |
---|
[20fe0cd] | 295 | def plot_2d(q, n=300): |
---|
[4f611f1] | 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) |
---|
[20fe0cd] | 304 | Zq = kernel_2d(q=q, theta=Atheta, phi=Aphi) |
---|
[4f611f1] | 305 | #Zq *= abs(sin(Atheta)) |
---|
| 306 | pylab.pcolor(degrees(theta), degrees(phi), log10(np.fmax(Zq, 1.e-6))) |
---|
| 307 | pylab.axis('tight') |
---|
[20fe0cd] | 308 | pylab.title("%s I(q,t) sin(t) for q=%g" % (KERNEL.__doc__, q)) |
---|
[4f611f1] | 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 | |
---|
[20fe0cd] | 315 | def main(Qstr): |
---|
[31eea1f] | 316 | Q = float(Qstr) |
---|
| 317 | if shape == 'sphere': |
---|
| 318 | print("exact", NORM*sas_3j1x_x(Q*RADIUS)**2) |
---|
[20fe0cd] | 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 | |
---|
| 335 | if __name__ == "__main__": |
---|
| 336 | main(Qstr) |
---|