1 | from __future__ import division, print_function |
---|
2 | |
---|
3 | import time |
---|
4 | from copy import copy |
---|
5 | |
---|
6 | import numpy as np |
---|
7 | from numpy import pi, radians, sin, cos, sqrt |
---|
8 | from numpy.random import poisson, uniform |
---|
9 | from numpy.polynomial.legendre import leggauss |
---|
10 | from scipy.integrate import simps |
---|
11 | from scipy.special import j1 as J1 |
---|
12 | |
---|
13 | # Definition of rotation matrices comes from wikipedia: |
---|
14 | # https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations |
---|
15 | def Rx(angle): |
---|
16 | """Construct a matrix to rotate points about *x* by *angle* degrees.""" |
---|
17 | a = radians(angle) |
---|
18 | R = [[1, 0, 0], |
---|
19 | [0, +cos(a), -sin(a)], |
---|
20 | [0, +sin(a), +cos(a)]] |
---|
21 | return np.matrix(R) |
---|
22 | |
---|
23 | def Ry(angle): |
---|
24 | """Construct a matrix to rotate points about *y* by *angle* degrees.""" |
---|
25 | a = radians(angle) |
---|
26 | R = [[+cos(a), 0, +sin(a)], |
---|
27 | [0, 1, 0], |
---|
28 | [-sin(a), 0, +cos(a)]] |
---|
29 | return np.matrix(R) |
---|
30 | |
---|
31 | def Rz(angle): |
---|
32 | """Construct a matrix to rotate points about *z* by *angle* degrees.""" |
---|
33 | a = radians(angle) |
---|
34 | R = [[+cos(a), -sin(a), 0], |
---|
35 | [+sin(a), +cos(a), 0], |
---|
36 | [0, 0, 1]] |
---|
37 | return np.matrix(R) |
---|
38 | |
---|
39 | def rotation(theta, phi, psi): |
---|
40 | """ |
---|
41 | Apply the jitter transform to a set of points. |
---|
42 | |
---|
43 | Points are stored in a 3 x n numpy matrix, not a numpy array or tuple. |
---|
44 | """ |
---|
45 | return Rx(phi)*Ry(theta)*Rz(psi) |
---|
46 | |
---|
47 | def apply_view(points, view): |
---|
48 | """ |
---|
49 | Apply the view transform (theta, phi, psi) to a set of points. |
---|
50 | |
---|
51 | Points are stored in a 3 x n numpy array. |
---|
52 | |
---|
53 | View angles are in degrees. |
---|
54 | """ |
---|
55 | theta, phi, psi = view |
---|
56 | return np.asarray((Rz(phi)*Ry(theta)*Rz(psi))*np.matrix(points.T)).T |
---|
57 | |
---|
58 | |
---|
59 | def invert_view(qx, qy, view): |
---|
60 | """ |
---|
61 | Return (qa, qb, qc) for the (theta, phi, psi) view angle at detector |
---|
62 | pixel (qx, qy). |
---|
63 | |
---|
64 | View angles are in degrees. |
---|
65 | """ |
---|
66 | theta, phi, psi = view |
---|
67 | q = np.vstack((qx.flatten(), qy.flatten(), 0*qx.flatten())) |
---|
68 | return np.asarray((Rz(-psi)*Ry(-theta)*Rz(-phi))*np.matrix(q)) |
---|
69 | |
---|
70 | |
---|
71 | class Shape: |
---|
72 | rotation = np.matrix([[1., 0, 0], [0, 1, 0], [0, 0, 1]]) |
---|
73 | center = np.array([0., 0., 0.])[:, None] |
---|
74 | r_max = None |
---|
75 | |
---|
76 | def volume(self): |
---|
77 | # type: () -> float |
---|
78 | raise NotImplementedError() |
---|
79 | |
---|
80 | def sample(self, density): |
---|
81 | # type: (float) -> np.ndarray[N], np.ndarray[N, 3] |
---|
82 | raise NotImplementedError() |
---|
83 | |
---|
84 | def rotate(self, theta, phi, psi): |
---|
85 | self.rotation = rotation(theta, phi, psi) * self.rotation |
---|
86 | return self |
---|
87 | |
---|
88 | def shift(self, x, y, z): |
---|
89 | self.center = self.center + np.array([x, y, z])[:, None] |
---|
90 | return self |
---|
91 | |
---|
92 | def _adjust(self, points): |
---|
93 | points = np.asarray(self.rotation * np.matrix(points.T)) + self.center |
---|
94 | return points.T |
---|
95 | |
---|
96 | def r_bins(self, q, over_sampling=1, r_step=0.): |
---|
97 | r_max = min(2 * pi / q[0], self.r_max) |
---|
98 | if r_step == 0.: |
---|
99 | r_step = 2 * pi / q[-1] / over_sampling |
---|
100 | #r_step = 0.01 |
---|
101 | return np.arange(r_step, r_max, r_step) |
---|
102 | |
---|
103 | class Composite(Shape): |
---|
104 | def __init__(self, shapes, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
105 | self.shapes = shapes |
---|
106 | self.rotate(*orientation) |
---|
107 | self.shift(*center) |
---|
108 | |
---|
109 | # Find the worst case distance between any two points amongst a set |
---|
110 | # of shapes independent of orientation. This could easily be a |
---|
111 | # factor of two worse than necessary, e.g., a pair of thin rods |
---|
112 | # end-to-end vs the same pair side-by-side. |
---|
113 | distances = [((s1.r_max + s2.r_max)/2 |
---|
114 | + sqrt(np.sum((s1.center - s2.center)**2))) |
---|
115 | for s1 in shapes |
---|
116 | for s2 in shapes] |
---|
117 | self.r_max = max(distances + [s.r_max for s in shapes]) |
---|
118 | |
---|
119 | def volume(self): |
---|
120 | return sum(shape.volume() for shape in self.shapes) |
---|
121 | |
---|
122 | def sample(self, density): |
---|
123 | values, points = zip(*(shape.sample(density) for shape in self.shapes)) |
---|
124 | return np.hstack(values), self._adjust(np.vstack(points)) |
---|
125 | |
---|
126 | class Box(Shape): |
---|
127 | def __init__(self, a, b, c, |
---|
128 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
129 | self.value = np.asarray(value) |
---|
130 | self.rotate(*orientation) |
---|
131 | self.shift(*center) |
---|
132 | self.a, self.b, self.c = a, b, c |
---|
133 | self._scale = np.array([a/2, b/2, c/2])[None, :] |
---|
134 | self.r_max = sqrt(a**2 + b**2 + c**2) |
---|
135 | |
---|
136 | def volume(self): |
---|
137 | return self.a*self.b*self.c |
---|
138 | |
---|
139 | def sample(self, density): |
---|
140 | num_points = poisson(density*self.a*self.b*self.c) |
---|
141 | points = self._scale*uniform(-1, 1, size=(num_points, 3)) |
---|
142 | values = self.value.repeat(points.shape[0]) |
---|
143 | return values, self._adjust(points) |
---|
144 | |
---|
145 | class EllipticalCylinder(Shape): |
---|
146 | def __init__(self, ra, rb, length, |
---|
147 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
148 | self.value = np.asarray(value) |
---|
149 | self.rotate(*orientation) |
---|
150 | self.shift(*center) |
---|
151 | self.ra, self.rb, self.length = ra, rb, length |
---|
152 | self._scale = np.array([ra, rb, length/2])[None, :] |
---|
153 | self.r_max = sqrt(4*max(ra, rb)**2 + length**2) |
---|
154 | |
---|
155 | def volume(self): |
---|
156 | return pi*self.ra*self.rb*self.length |
---|
157 | |
---|
158 | def sample(self, density): |
---|
159 | # density of the bounding box |
---|
160 | num_points = poisson(density*4*self.ra*self.rb*self.length) |
---|
161 | points = uniform(-1, 1, size=(num_points, 3)) |
---|
162 | radius = points[:, 0]**2 + points[:, 1]**2 |
---|
163 | points = self._scale*points[radius <= 1] |
---|
164 | values = self.value.repeat(points.shape[0]) |
---|
165 | return values, self._adjust(points) |
---|
166 | |
---|
167 | class TriaxialEllipsoid(Shape): |
---|
168 | def __init__(self, ra, rb, rc, |
---|
169 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
170 | self.value = np.asarray(value) |
---|
171 | self.rotate(*orientation) |
---|
172 | self.shift(*center) |
---|
173 | self.ra, self.rb, self.rc = ra, rb, rc |
---|
174 | self._scale = np.array([ra, rb, rc])[None, :] |
---|
175 | self.r_max = 2*max(ra, rb, rc) |
---|
176 | |
---|
177 | def volume(self): |
---|
178 | return 4*pi/3 * self.ra * self.rb * self.rc |
---|
179 | |
---|
180 | def sample(self, density): |
---|
181 | # randomly sample from a box of side length 2*r, excluding anything |
---|
182 | # not in the ellipsoid |
---|
183 | num_points = poisson(density*8*self.ra*self.rb*self.rc) |
---|
184 | points = uniform(-1, 1, size=(num_points, 3)) |
---|
185 | radius = np.sum(points**2, axis=1) |
---|
186 | points = self._scale*points[radius <= 1] |
---|
187 | values = self.value.repeat(points.shape[0]) |
---|
188 | return values, self._adjust(points) |
---|
189 | |
---|
190 | class Helix(Shape): |
---|
191 | def __init__(self, helix_radius, helix_pitch, tube_radius, tube_length, |
---|
192 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
193 | self.value = np.asarray(value) |
---|
194 | self.rotate(*orientation) |
---|
195 | self.shift(*center) |
---|
196 | self.helix_radius, self.helix_pitch = helix_radius, helix_pitch |
---|
197 | self.tube_radius, self.tube_length = tube_radius, tube_length |
---|
198 | helix_length = helix_pitch * tube_length/sqrt(helix_radius**2 + helix_pitch**2) |
---|
199 | self.r_max = sqrt((2*helix_radius + 2*tube_radius)*2 |
---|
200 | + (helix_length + 2*tube_radius)**2) |
---|
201 | |
---|
202 | def volume(self): |
---|
203 | # small tube radius approximation; for larger tubes need to account |
---|
204 | # for the fact that the inner length is much shorter than the outer |
---|
205 | # length |
---|
206 | return pi*self.tube_radius**2*self.tube_length |
---|
207 | |
---|
208 | def points(self, density): |
---|
209 | num_points = poisson(density*4*self.tube_radius**2*self.tube_length) |
---|
210 | points = uniform(-1, 1, size=(num_points, 3)) |
---|
211 | radius = points[:, 0]**2 + points[:, 1]**2 |
---|
212 | points = points[radius <= 1] |
---|
213 | |
---|
214 | # Based on math stackexchange answer by Jyrki Lahtonen |
---|
215 | # https://math.stackexchange.com/a/461637 |
---|
216 | # with helix along z rather than x [so tuples in answer are (z, x, y)] |
---|
217 | # and with random points in the cross section (p1, p2) rather than |
---|
218 | # uniform points on the surface (cos u, sin u). |
---|
219 | a, R = self.tube_radius, self.helix_radius |
---|
220 | h = self.helix_pitch |
---|
221 | scale = 1/sqrt(R**2 + h**2) |
---|
222 | t = points[:, 3] * (self.tube_length * scale/2) |
---|
223 | cos_t, sin_t = cos(t), sin(t) |
---|
224 | |
---|
225 | # rx = R*cos_t |
---|
226 | # ry = R*sin_t |
---|
227 | # rz = h*t |
---|
228 | # nx = -a * cos_t * points[:, 1] |
---|
229 | # ny = -a * sin_t * points[:, 1] |
---|
230 | # nz = 0 |
---|
231 | # bx = (a * h/scale) * sin_t * points[:, 2] |
---|
232 | # by = (-a * h/scale) * cos_t * points[:, 2] |
---|
233 | # bz = a*R/scale |
---|
234 | # x = rx + nx + bx |
---|
235 | # y = ry + ny + by |
---|
236 | # z = rz + nz + bz |
---|
237 | u, v = (R - a*points[:, 1]), (a * h/scale)*points[:, 2] |
---|
238 | x = u * cos_t + v * sin_t |
---|
239 | y = u * sin_t - v * cos_t |
---|
240 | z = a*R/scale + h * t |
---|
241 | |
---|
242 | points = np.hstack((x, y, z)) |
---|
243 | values = self.value.repeat(points.shape[0]) |
---|
244 | return values, self._adjust(points) |
---|
245 | |
---|
246 | NUMBA = False |
---|
247 | if NUMBA: |
---|
248 | from numba import njit |
---|
249 | @njit("f8[:](f8[:],f8[:],f8[:],f8[:],f8[:],f8[:],f8[:])") |
---|
250 | def _Iqxy(values, x, y, z, qa, qb, qc): |
---|
251 | Iq = np.zeros_like(qa) |
---|
252 | for j in range(len(Iq)): |
---|
253 | total = 0. + 0j |
---|
254 | for k in range(len(Iq)): |
---|
255 | total += values[k]*np.exp(1j*(qa[j]*x[k] + qb[j]*y[k] + qc[j]*z[k])) |
---|
256 | Iq[j] = abs(total)**2 |
---|
257 | return Iq |
---|
258 | |
---|
259 | def calc_Iqxy(qx, qy, rho, points, volume=1.0, view=(0, 0, 0)): |
---|
260 | qx, qy = np.broadcast_arrays(qx, qy) |
---|
261 | qa, qb, qc = invert_view(qx, qy, view) |
---|
262 | rho, volume = np.broadcast_arrays(rho, volume) |
---|
263 | values = rho*volume |
---|
264 | x, y, z = points.T |
---|
265 | |
---|
266 | # I(q) = |sum V(r) rho(r) e^(1j q.r)|^2 / sum V(r) |
---|
267 | if NUMBA: |
---|
268 | Iq = _Iqxy(values, x, y, z, qa.flatten(), qb.flatten(), qc.flatten()) |
---|
269 | else: |
---|
270 | Iq = [abs(np.sum(values*np.exp(1j*(qa_k*x + qb_k*y + qc_k*z))))**2 |
---|
271 | for qa_k, qb_k, qc_k in zip(qa.flat, qb.flat, qc.flat)] |
---|
272 | return np.asarray(Iq).reshape(qx.shape) / np.sum(volume) |
---|
273 | |
---|
274 | def _calc_Pr_nonuniform(r, rho, points): |
---|
275 | # Make Pr a little be bigger than necessary so that only distances |
---|
276 | # min < d < max end up in Pr |
---|
277 | n_max = len(r)+1 |
---|
278 | extended_Pr = np.zeros(n_max+1, 'd') |
---|
279 | # r refers to bin centers; find corresponding bin edges |
---|
280 | bins = bin_edges(r) |
---|
281 | t_next = time.clock() + 3 |
---|
282 | for k, rho_k in enumerate(rho[:-1]): |
---|
283 | distance = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1)) |
---|
284 | weights = rho_k * rho[k+1:] |
---|
285 | index = np.searchsorted(bins, distance) |
---|
286 | # Note: indices may be duplicated, so "Pr[index] += w" will not work!! |
---|
287 | extended_Pr += np.bincount(index, weights, n_max+1) |
---|
288 | t = time.clock() |
---|
289 | if t > t_next: |
---|
290 | t_next = t + 3 |
---|
291 | print("processing %d of %d"%(k, len(rho)-1)) |
---|
292 | Pr = extended_Pr[1:-1] |
---|
293 | return Pr |
---|
294 | |
---|
295 | def _calc_Pr_uniform(r, rho, points): |
---|
296 | # Make Pr a little be bigger than necessary so that only distances |
---|
297 | # min < d < max end up in Pr |
---|
298 | dr, n_max = r[0], len(r) |
---|
299 | extended_Pr = np.zeros(n_max+1, 'd') |
---|
300 | t0 = time.clock() |
---|
301 | t_next = t0 + 3 |
---|
302 | for k, rho_k in enumerate(rho[:-1]): |
---|
303 | distances = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1)) |
---|
304 | weights = rho_k * rho[k+1:] |
---|
305 | index = np.minimum(np.asarray(distances/dr, 'i'), n_max) |
---|
306 | # Note: indices may be duplicated, so "Pr[index] += w" will not work!! |
---|
307 | extended_Pr += np.bincount(index, weights, n_max+1) |
---|
308 | t = time.clock() |
---|
309 | if t > t_next: |
---|
310 | t_next = t + 3 |
---|
311 | print("processing %d of %d"%(k, len(rho)-1)) |
---|
312 | #print("time py:", time.clock() - t0) |
---|
313 | Pr = extended_Pr[:-1] |
---|
314 | # Make Pr independent of sampling density. The factor of 2 comes because |
---|
315 | # we are only accumulating the upper triangular distances. |
---|
316 | #Pr = Pr * 2 / len(rho)**2 |
---|
317 | return Pr |
---|
318 | |
---|
319 | # Can get an additional 2x by going to C. Cuda/OpenCL will allow even |
---|
320 | # more speedup, though still bounded by the n^2 cose. |
---|
321 | """ |
---|
322 | void pdfcalc(int n, const double *pts, const double *rho, |
---|
323 | int nPr, double *Pr, double rstep) |
---|
324 | { |
---|
325 | int i,j; |
---|
326 | |
---|
327 | for (i=0; i<n-2; i++) { |
---|
328 | for (j=i+1; j<=n-1; j++) { |
---|
329 | const double dxx=pts[3*i]-pts[3*j]; |
---|
330 | const double dyy=pts[3*i+1]-pts[3*j+1]; |
---|
331 | const double dzz=pts[3*i+2]-pts[3*j+2]; |
---|
332 | const double d=sqrt(dxx*dxx+dyy*dyy+dzz*dzz); |
---|
333 | const int k=rint(d/rstep); |
---|
334 | if (k < nPr) Pr[k]+=rho[i]*rho[j]; |
---|
335 | } |
---|
336 | } |
---|
337 | } |
---|
338 | """ |
---|
339 | |
---|
340 | if NUMBA: |
---|
341 | @njit("f8[:](f8[:], f8[:], f8[:,:])") |
---|
342 | def _calc_Pr_uniform_jit(r, rho, points): |
---|
343 | dr = r[0] |
---|
344 | n_max = len(r) |
---|
345 | Pr = np.zeros_like(r) |
---|
346 | for j in range(len(rho) - 1): |
---|
347 | x, y, z = points[j, 0], points[j, 1], points[j, 2] |
---|
348 | for k in range(j+1, len(rho)): |
---|
349 | distance = sqrt((x - points[k, 0])**2 |
---|
350 | + (y - points[k, 1])**2 |
---|
351 | + (z - points[k, 2])**2) |
---|
352 | index = int(distance/dr) |
---|
353 | if index < n_max: |
---|
354 | Pr[index] += rho[j] * rho[k] |
---|
355 | # Make Pr independent of sampling density. The factor of 2 comes because |
---|
356 | # we are only accumulating the upper triangular distances. |
---|
357 | #Pr = Pr * 2 / len(rho)**2 |
---|
358 | return Pr |
---|
359 | |
---|
360 | |
---|
361 | def calc_Pr(r, rho, points): |
---|
362 | # P(r) with uniform steps in r is 3x faster; check if we are uniform |
---|
363 | # before continuing |
---|
364 | if np.max(np.abs(np.diff(r) - r[0])) > r[0]*0.01: |
---|
365 | Pr = _calc_Pr_nonuniform(r, rho, points) |
---|
366 | else: |
---|
367 | if NUMBA: |
---|
368 | Pr = _calc_Pr_uniform_jit(r, rho, points) |
---|
369 | else: |
---|
370 | Pr = _calc_Pr_uniform(r, rho, points) |
---|
371 | return Pr / Pr.max() |
---|
372 | |
---|
373 | |
---|
374 | def j0(x): |
---|
375 | return np.sinc(x/np.pi) |
---|
376 | |
---|
377 | def calc_Iq(q, r, Pr): |
---|
378 | Iq = np.array([simps(Pr * j0(qk*r), r) for qk in q]) |
---|
379 | #Iq = np.array([np.trapz(Pr * j0(qk*r), r) for qk in q]) |
---|
380 | Iq /= Iq[0] |
---|
381 | return Iq |
---|
382 | |
---|
383 | # NOTE: copied from sasmodels/resolution.py |
---|
384 | def bin_edges(x): |
---|
385 | """ |
---|
386 | Determine bin edges from bin centers, assuming that edges are centered |
---|
387 | between the bins. |
---|
388 | |
---|
389 | Note: this uses the arithmetic mean, which may not be appropriate for |
---|
390 | log-scaled data. |
---|
391 | """ |
---|
392 | if len(x) < 2 or (np.diff(x) < 0).any(): |
---|
393 | raise ValueError("Expected bins to be an increasing set") |
---|
394 | edges = np.hstack([ |
---|
395 | x[0] - 0.5*(x[1] - x[0]), # first point minus half first interval |
---|
396 | 0.5*(x[1:] + x[:-1]), # mid points of all central intervals |
---|
397 | x[-1] + 0.5*(x[-1] - x[-2]), # last point plus half last interval |
---|
398 | ]) |
---|
399 | return edges |
---|
400 | |
---|
401 | def plot_calc(r, Pr, q, Iq, theory=None): |
---|
402 | import matplotlib.pyplot as plt |
---|
403 | plt.subplot(211) |
---|
404 | plt.plot(r, Pr, '-', label="Pr") |
---|
405 | plt.xlabel('r (A)') |
---|
406 | plt.ylabel('Pr (1/A^2)') |
---|
407 | plt.subplot(212) |
---|
408 | plt.loglog(q, Iq, '-', label='from Pr') |
---|
409 | plt.xlabel('q (1/A') |
---|
410 | plt.ylabel('Iq') |
---|
411 | if theory is not None: |
---|
412 | plt.loglog(theory[0], theory[1], '-', label='analytic') |
---|
413 | plt.legend() |
---|
414 | |
---|
415 | def plot_calc_2d(qx, qy, Iqxy, theory=None): |
---|
416 | import matplotlib.pyplot as plt |
---|
417 | qx, qy = bin_edges(qx), bin_edges(qy) |
---|
418 | #qx, qy = np.meshgrid(qx, qy) |
---|
419 | if theory is not None: |
---|
420 | plt.subplot(121) |
---|
421 | plt.pcolormesh(qx, qy, np.log10(Iqxy)) |
---|
422 | plt.xlabel('qx (1/A)') |
---|
423 | plt.ylabel('qy (1/A)') |
---|
424 | if theory is not None: |
---|
425 | plt.subplot(122) |
---|
426 | plt.pcolormesh(qx, qy, np.log10(theory)) |
---|
427 | plt.xlabel('qx (1/A)') |
---|
428 | |
---|
429 | def plot_points(rho, points): |
---|
430 | import mpl_toolkits.mplot3d |
---|
431 | import matplotlib.pyplot as plt |
---|
432 | |
---|
433 | ax = plt.axes(projection='3d') |
---|
434 | try: |
---|
435 | ax.axis('square') |
---|
436 | except Exception: |
---|
437 | pass |
---|
438 | n = len(points) |
---|
439 | #print("len points", n) |
---|
440 | index = np.random.choice(n, size=500) if n > 500 else slice(None, None) |
---|
441 | ax.scatter(points[index, 0], points[index, 1], points[index, 2], c=rho[index]) |
---|
442 | #low, high = points.min(axis=0), points.max(axis=0) |
---|
443 | #ax.axis([low[0], high[0], low[1], high[1], low[2], high[2]]) |
---|
444 | ax.autoscale(True) |
---|
445 | |
---|
446 | def check_shape(shape, fn=None): |
---|
447 | rho_solvent = 0 |
---|
448 | q = np.logspace(-3, 0, 200) |
---|
449 | r = shape.r_bins(q, r_step=0.01) |
---|
450 | sampling_density = 6*5000 / shape.volume() |
---|
451 | rho, points = shape.sample(sampling_density) |
---|
452 | t0 = time.time() |
---|
453 | Pr = calc_Pr(r, rho-rho_solvent, points) |
---|
454 | print("calc Pr time", time.time() - t0) |
---|
455 | Iq = calc_Iq(q, r, Pr) |
---|
456 | theory = (q, fn(q)) if fn is not None else None |
---|
457 | |
---|
458 | import pylab |
---|
459 | #plot_points(rho, points); pylab.figure() |
---|
460 | plot_calc(r, Pr, q, Iq, theory=theory) |
---|
461 | pylab.show() |
---|
462 | |
---|
463 | def check_shape_2d(shape, fn=None, view=(0, 0, 0)): |
---|
464 | rho_solvent = 0 |
---|
465 | nq, qmax = 100, 1.0 |
---|
466 | qx = np.linspace(0.0, qmax, nq) |
---|
467 | qy = np.linspace(0.0, qmax, nq) |
---|
468 | Qx, Qy = np.meshgrid(qx, qy) |
---|
469 | sampling_density = 50000 / shape.volume() |
---|
470 | #t0 = time.time() |
---|
471 | rho, points = shape.sample(sampling_density) |
---|
472 | #print("sample time", time.time() - t0) |
---|
473 | t0 = time.time() |
---|
474 | Iqxy = calc_Iqxy(Qx, Qy, rho, points, view=view) |
---|
475 | print("calc time", time.time() - t0) |
---|
476 | theory = fn(Qx, Qy) if fn is not None else None |
---|
477 | Iqxy += 0.001 * Iqxy.max() |
---|
478 | if theory is not None: |
---|
479 | theory += 0.001 * theory.max() |
---|
480 | |
---|
481 | import pylab |
---|
482 | #plot_points(rho, points); pylab.figure() |
---|
483 | plot_calc_2d(qx, qy, Iqxy, theory=theory) |
---|
484 | pylab.show() |
---|
485 | |
---|
486 | def sas_sinx_x(x): |
---|
487 | with np.errstate(all='ignore'): |
---|
488 | retvalue = sin(x)/x |
---|
489 | retvalue[x == 0.] = 1. |
---|
490 | return retvalue |
---|
491 | |
---|
492 | def sas_2J1x_x(x): |
---|
493 | with np.errstate(all='ignore'): |
---|
494 | retvalue = 2*J1(x)/x |
---|
495 | retvalue[x == 0] = 1. |
---|
496 | return retvalue |
---|
497 | |
---|
498 | def sas_3j1x_x(x): |
---|
499 | """return 3*j1(x)/x""" |
---|
500 | with np.errstate(all='ignore'): |
---|
501 | retvalue = 3*(sin(x) - x*cos(x))/x**3 |
---|
502 | retvalue[x == 0.] = 1. |
---|
503 | return retvalue |
---|
504 | |
---|
505 | def cylinder_Iq(q, radius, length): |
---|
506 | z, w = leggauss(76) |
---|
507 | cos_alpha = (z+1)/2 |
---|
508 | sin_alpha = sqrt(1.0 - cos_alpha**2) |
---|
509 | Iq = np.empty_like(q) |
---|
510 | for k, qk in enumerate(q): |
---|
511 | qab, qc = qk*sin_alpha, qk*cos_alpha |
---|
512 | Fq = sas_2J1x_x(qab*radius) * j0(qc*length/2) |
---|
513 | Iq[k] = np.sum(w*Fq**2) |
---|
514 | Iq = Iq/Iq[0] |
---|
515 | return Iq |
---|
516 | |
---|
517 | def cylinder_Iqxy(qx, qy, radius, length, view=(0, 0, 0)): |
---|
518 | qa, qb, qc = invert_view(qx, qy, view) |
---|
519 | qab = np.sqrt(qa**2 + qb**2) |
---|
520 | Fq = sas_2J1x_x(qab*radius) * j0(qc*length/2) |
---|
521 | Iq = Fq**2 |
---|
522 | return Iq.reshape(qx.shape) |
---|
523 | |
---|
524 | def sphere_Iq(q, radius): |
---|
525 | Iq = sas_3j1x_x(q*radius)**2 |
---|
526 | return Iq/Iq[0] |
---|
527 | |
---|
528 | def csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core): |
---|
529 | z, w = leggauss(76) |
---|
530 | |
---|
531 | sld_solvent = 0 |
---|
532 | overlapping = False |
---|
533 | dr0 = sld_core - sld_solvent |
---|
534 | drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent |
---|
535 | tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc |
---|
536 | |
---|
537 | outer_sum = np.zeros_like(q) |
---|
538 | for cos_alpha, outer_w in zip((z+1)/2, w): |
---|
539 | sin_alpha = sqrt(1.0-cos_alpha*cos_alpha) |
---|
540 | qc = q*cos_alpha |
---|
541 | siC = c*j0(c*qc/2) |
---|
542 | siCt = tC*j0(tC*qc/2) |
---|
543 | inner_sum = np.zeros_like(q) |
---|
544 | for beta, inner_w in zip((z + 1)*pi/4, w): |
---|
545 | qa, qb = q*sin_alpha*sin(beta), q*sin_alpha*cos(beta) |
---|
546 | siA = a*j0(a*qa/2) |
---|
547 | siB = b*j0(b*qb/2) |
---|
548 | siAt = tA*j0(tA*qa/2) |
---|
549 | siBt = tB*j0(tB*qb/2) |
---|
550 | if overlapping: |
---|
551 | Fq = (dr0*siA*siB*siC |
---|
552 | + drA*(siAt-siA)*siB*siC |
---|
553 | + drB*siAt*(siBt-siB)*siC |
---|
554 | + drC*siAt*siBt*(siCt-siC)) |
---|
555 | else: |
---|
556 | Fq = (dr0*siA*siB*siC |
---|
557 | + drA*(siAt-siA)*siB*siC |
---|
558 | + drB*siA*(siBt-siB)*siC |
---|
559 | + drC*siA*siB*(siCt-siC)) |
---|
560 | inner_sum += inner_w * Fq**2 |
---|
561 | outer_sum += outer_w * inner_sum |
---|
562 | Iq = outer_sum / 4 # = outer*um*zm*8.0/(4.0*M_PI) |
---|
563 | return Iq/Iq[0] |
---|
564 | |
---|
565 | def csbox_Iqxy(qx, qy, a, b, c, da, db, dc, slda, sldb, sldc, sld_core, view=(0,0,0)): |
---|
566 | qa, qb, qc = invert_view(qx, qy, view) |
---|
567 | |
---|
568 | sld_solvent = 0 |
---|
569 | overlapping = False |
---|
570 | dr0 = sld_core - sld_solvent |
---|
571 | drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent |
---|
572 | tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc |
---|
573 | siA = a*sas_sinx_x(a*qa/2) |
---|
574 | siB = b*sas_sinx_x(b*qb/2) |
---|
575 | siC = c*sas_sinx_x(c*qc/2) |
---|
576 | siAt = tA*sas_sinx_x(tA*qa/2) |
---|
577 | siBt = tB*sas_sinx_x(tB*qb/2) |
---|
578 | siCt = tC*sas_sinx_x(tC*qc/2) |
---|
579 | Fq = (dr0*siA*siB*siC |
---|
580 | + drA*(siAt-siA)*siB*siC |
---|
581 | + drB*siA*(siBt-siB)*siC |
---|
582 | + drC*siA*siB*(siCt-siC)) |
---|
583 | Iq = Fq**2 |
---|
584 | return Iq.reshape(qx.shape) |
---|
585 | |
---|
586 | def check_cylinder(radius=25, length=125, rho=2.): |
---|
587 | shape = EllipticalCylinder(radius, radius, length, rho) |
---|
588 | fn = lambda q: cylinder_Iq(q, radius, length) |
---|
589 | check_shape(shape, fn) |
---|
590 | |
---|
591 | def check_cylinder_2d(radius=25, length=125, rho=2., view=(0, 0, 0)): |
---|
592 | shape = EllipticalCylinder(radius, radius, length, rho) |
---|
593 | fn = lambda qx, qy, view=view: cylinder_Iqxy(qx, qy, radius, length, view=view) |
---|
594 | check_shape_2d(shape, fn, view=view) |
---|
595 | |
---|
596 | def check_cylinder_2d_lattice(radius=25, length=125, rho=2., |
---|
597 | view=(0, 0, 0)): |
---|
598 | nx, dx = 1, 2*radius |
---|
599 | ny, dy = 30, 2*radius |
---|
600 | nz, dz = 30, length |
---|
601 | dx, dy, dz = 2*dx, 2*dy, 2*dz |
---|
602 | def center(*args): |
---|
603 | sigma = 0.333 |
---|
604 | space = 2 |
---|
605 | return [(space*n+np.random.randn()*sigma)*x for n, x in args] |
---|
606 | shapes = [EllipticalCylinder(radius, radius, length, rho, |
---|
607 | #center=(ix*dx, iy*dy, iz*dz) |
---|
608 | orientation=np.random.randn(3)*0, |
---|
609 | center=center((ix, dx), (iy, dy), (iz, dz)) |
---|
610 | ) |
---|
611 | for ix in range(nx) |
---|
612 | for iy in range(ny) |
---|
613 | for iz in range(nz)] |
---|
614 | shape = Composite(shapes) |
---|
615 | fn = lambda qx, qy, view=view: cylinder_Iqxy(qx, qy, radius, length, view=view) |
---|
616 | check_shape_2d(shape, fn, view=view) |
---|
617 | |
---|
618 | def check_sphere(radius=125, rho=2): |
---|
619 | shape = TriaxialEllipsoid(radius, radius, radius, rho) |
---|
620 | fn = lambda q: sphere_Iq(q, radius) |
---|
621 | check_shape(shape, fn) |
---|
622 | |
---|
623 | def check_csbox(a=10, b=20, c=30, da=1, db=2, dc=3, slda=1, sldb=2, sldc=3, sld_core=4): |
---|
624 | core = Box(a, b, c, sld_core) |
---|
625 | side_a = Box(da, b, c, slda, center=((a+da)/2, 0, 0)) |
---|
626 | side_b = Box(a, db, c, sldb, center=(0, (b+db)/2, 0)) |
---|
627 | side_c = Box(a, b, dc, sldc, center=(0, 0, (c+dc)/2)) |
---|
628 | side_a2 = copy(side_a).shift(-a-da, 0, 0) |
---|
629 | side_b2 = copy(side_b).shift(0, -b-db, 0) |
---|
630 | side_c2 = copy(side_c).shift(0, 0, -c-dc) |
---|
631 | shape = Composite((core, side_a, side_b, side_c, side_a2, side_b2, side_c2)) |
---|
632 | def fn(q): |
---|
633 | return csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core) |
---|
634 | #check_shape(shape, fn) |
---|
635 | |
---|
636 | view = (20, 30, 40) |
---|
637 | def fn_xy(qx, qy): |
---|
638 | return csbox_Iqxy(qx, qy, a, b, c, da, db, dc, |
---|
639 | slda, sldb, sldc, sld_core, view=view) |
---|
640 | check_shape_2d(shape, fn_xy, view=view) |
---|
641 | |
---|
642 | if __name__ == "__main__": |
---|
643 | check_cylinder(radius=10, length=40) |
---|
644 | #check_cylinder_2d(radius=10, length=40, view=(90,30,0)) |
---|
645 | #check_cylinder_2d_lattice(radius=10, length=50, view=(90,30,0)) |
---|
646 | #check_sphere() |
---|
647 | #check_csbox() |
---|
648 | #check_csbox(da=100, db=200, dc=300) |
---|