[cfa28d3] | 1 | from __future__ import division, print_function |
---|
| 2 | |
---|
| 3 | import time |
---|
| 4 | from copy import copy |
---|
[7d0afa3] | 5 | import os |
---|
| 6 | import argparse |
---|
| 7 | from collections import OrderedDict |
---|
[cfa28d3] | 8 | |
---|
| 9 | import numpy as np |
---|
| 10 | from numpy import pi, radians, sin, cos, sqrt |
---|
[5778279] | 11 | from numpy.random import poisson, uniform, randn, rand, randint |
---|
[a1c32c2] | 12 | from numpy.polynomial.legendre import leggauss |
---|
[cfa28d3] | 13 | from scipy.integrate import simps |
---|
| 14 | from scipy.special import j1 as J1 |
---|
| 15 | |
---|
[7d0afa3] | 16 | try: |
---|
| 17 | import numba |
---|
| 18 | USE_NUMBA = True |
---|
| 19 | except ImportError: |
---|
| 20 | USE_NUMBA = False |
---|
| 21 | |
---|
[cfa28d3] | 22 | # Definition of rotation matrices comes from wikipedia: |
---|
| 23 | # https://en.wikipedia.org/wiki/Rotation_matrix#Basic_rotations |
---|
| 24 | def Rx(angle): |
---|
| 25 | """Construct a matrix to rotate points about *x* by *angle* degrees.""" |
---|
| 26 | a = radians(angle) |
---|
| 27 | R = [[1, 0, 0], |
---|
| 28 | [0, +cos(a), -sin(a)], |
---|
| 29 | [0, +sin(a), +cos(a)]] |
---|
| 30 | return np.matrix(R) |
---|
| 31 | |
---|
| 32 | def Ry(angle): |
---|
| 33 | """Construct a matrix to rotate points about *y* by *angle* degrees.""" |
---|
| 34 | a = radians(angle) |
---|
| 35 | R = [[+cos(a), 0, +sin(a)], |
---|
| 36 | [0, 1, 0], |
---|
| 37 | [-sin(a), 0, +cos(a)]] |
---|
| 38 | return np.matrix(R) |
---|
| 39 | |
---|
| 40 | def Rz(angle): |
---|
| 41 | """Construct a matrix to rotate points about *z* by *angle* degrees.""" |
---|
| 42 | a = radians(angle) |
---|
| 43 | R = [[+cos(a), -sin(a), 0], |
---|
| 44 | [+sin(a), +cos(a), 0], |
---|
| 45 | [0, 0, 1]] |
---|
| 46 | return np.matrix(R) |
---|
| 47 | |
---|
| 48 | def rotation(theta, phi, psi): |
---|
| 49 | """ |
---|
| 50 | Apply the jitter transform to a set of points. |
---|
| 51 | |
---|
| 52 | Points are stored in a 3 x n numpy matrix, not a numpy array or tuple. |
---|
| 53 | """ |
---|
| 54 | return Rx(phi)*Ry(theta)*Rz(psi) |
---|
| 55 | |
---|
[226473d] | 56 | def apply_view(points, view): |
---|
| 57 | """ |
---|
| 58 | Apply the view transform (theta, phi, psi) to a set of points. |
---|
| 59 | |
---|
| 60 | Points are stored in a 3 x n numpy array. |
---|
| 61 | |
---|
| 62 | View angles are in degrees. |
---|
| 63 | """ |
---|
| 64 | theta, phi, psi = view |
---|
| 65 | return np.asarray((Rz(phi)*Ry(theta)*Rz(psi))*np.matrix(points.T)).T |
---|
| 66 | |
---|
| 67 | |
---|
| 68 | def invert_view(qx, qy, view): |
---|
| 69 | """ |
---|
| 70 | Return (qa, qb, qc) for the (theta, phi, psi) view angle at detector |
---|
| 71 | pixel (qx, qy). |
---|
| 72 | |
---|
| 73 | View angles are in degrees. |
---|
| 74 | """ |
---|
| 75 | theta, phi, psi = view |
---|
| 76 | q = np.vstack((qx.flatten(), qy.flatten(), 0*qx.flatten())) |
---|
| 77 | return np.asarray((Rz(-psi)*Ry(-theta)*Rz(-phi))*np.matrix(q)) |
---|
| 78 | |
---|
| 79 | |
---|
[2a39ca4] | 80 | I3 = np.matrix([[1., 0, 0], [0, 1, 0], [0, 0, 1]]) |
---|
| 81 | |
---|
[cfa28d3] | 82 | class Shape: |
---|
[2a39ca4] | 83 | rotation = I3 |
---|
[cfa28d3] | 84 | center = np.array([0., 0., 0.])[:, None] |
---|
| 85 | r_max = None |
---|
[5778279] | 86 | lattice_size = np.array((1, 1, 1)) |
---|
| 87 | lattice_spacing = np.array((1., 1., 1.)) |
---|
| 88 | lattice_distortion = 0.0 |
---|
| 89 | lattice_rotation = 0.0 |
---|
| 90 | lattice_type = "" |
---|
[cfa28d3] | 91 | |
---|
| 92 | def volume(self): |
---|
| 93 | # type: () -> float |
---|
| 94 | raise NotImplementedError() |
---|
| 95 | |
---|
| 96 | def sample(self, density): |
---|
| 97 | # type: (float) -> np.ndarray[N], np.ndarray[N, 3] |
---|
| 98 | raise NotImplementedError() |
---|
| 99 | |
---|
[032646d] | 100 | def dims(self): |
---|
| 101 | # type: () -> float, float, float |
---|
| 102 | raise NotImplementedError() |
---|
| 103 | |
---|
[cfa28d3] | 104 | def rotate(self, theta, phi, psi): |
---|
[2a39ca4] | 105 | if theta != 0. or phi != 0. or psi != 0.: |
---|
| 106 | self.rotation = rotation(theta, phi, psi) * self.rotation |
---|
[cfa28d3] | 107 | return self |
---|
| 108 | |
---|
| 109 | def shift(self, x, y, z): |
---|
| 110 | self.center = self.center + np.array([x, y, z])[:, None] |
---|
| 111 | return self |
---|
| 112 | |
---|
[5778279] | 113 | def lattice(self, size=(1, 1, 1), spacing=(2, 2, 2), type="sc", |
---|
| 114 | distortion=0.0, rotation=0.0): |
---|
| 115 | self.lattice_size = np.asarray(size, 'i') |
---|
| 116 | self.lattice_spacing = np.asarray(spacing, 'd') |
---|
| 117 | self.lattice_type = type |
---|
| 118 | self.lattice_distortion = distortion |
---|
| 119 | self.lattice_rotation = rotation |
---|
| 120 | |
---|
[cfa28d3] | 121 | def _adjust(self, points): |
---|
[2a39ca4] | 122 | if self.rotation is I3: |
---|
| 123 | points = points.T + self.center |
---|
| 124 | else: |
---|
| 125 | points = np.asarray(self.rotation * np.matrix(points.T)) + self.center |
---|
[5778279] | 126 | if self.lattice_type: |
---|
| 127 | points = self._apply_lattice(points) |
---|
[cfa28d3] | 128 | return points.T |
---|
| 129 | |
---|
[5778279] | 130 | def r_bins(self, q, over_sampling=10, r_step=0.): |
---|
| 131 | if self.lattice_type: |
---|
| 132 | r_max = np.sqrt(np.sum(self.lattice_size*self.lattice_spacing*self.dims)**2)/2 |
---|
| 133 | else: |
---|
| 134 | r_max = self.r_max |
---|
| 135 | #r_max = min(2 * pi / q[0], r_max) |
---|
[cfa28d3] | 136 | if r_step == 0.: |
---|
| 137 | r_step = 2 * pi / q[-1] / over_sampling |
---|
| 138 | #r_step = 0.01 |
---|
| 139 | return np.arange(r_step, r_max, r_step) |
---|
| 140 | |
---|
[5778279] | 141 | def _apply_lattice(self, points): |
---|
| 142 | """Spread points to different lattice positions""" |
---|
| 143 | size = self.lattice_size |
---|
| 144 | spacing = self.lattice_spacing |
---|
| 145 | shuffle = self.lattice_distortion |
---|
| 146 | rotate = self.lattice_rotation |
---|
| 147 | lattice = self.lattice_type |
---|
| 148 | |
---|
| 149 | if rotate != 0: |
---|
| 150 | # To vectorize the rotations we will need to unwrap the matrix multiply |
---|
| 151 | raise NotImplementedError("don't handle rotations yet") |
---|
| 152 | |
---|
| 153 | # Determine the number of lattice points in the lattice |
---|
| 154 | shapes_per_cell = 2 if lattice == "bcc" else 4 if lattice == "fcc" else 1 |
---|
| 155 | number_of_lattice_points = np.prod(size) * shapes_per_cell |
---|
| 156 | |
---|
| 157 | # For each point in the original shape, figure out which lattice point |
---|
| 158 | # to translate it to. This is both cell index (i*ny*nz + j*nz + k) as |
---|
| 159 | # well as the point in the cell (corner, body center or face center). |
---|
| 160 | nsamples = points.shape[1] |
---|
| 161 | lattice_point = randint(number_of_lattice_points, size=nsamples) |
---|
| 162 | |
---|
| 163 | # Translate the cell index into the i,j,k coordinates of the senter |
---|
| 164 | cell_index = lattice_point // shapes_per_cell |
---|
| 165 | center = np.vstack((cell_index//(size[1]*size[2]), |
---|
| 166 | (cell_index%(size[1]*size[2]))//size[2], |
---|
| 167 | cell_index%size[2])) |
---|
| 168 | center = np.asarray(center, dtype='d') |
---|
| 169 | if lattice == "bcc": |
---|
| 170 | center[:, lattice_point % shapes_per_cell == 1] += [[0.5], [0.5], [0.5]] |
---|
| 171 | elif lattice == "fcc": |
---|
| 172 | center[:, lattice_point % shapes_per_cell == 1] += [[0.0], [0.5], [0.5]] |
---|
| 173 | center[:, lattice_point % shapes_per_cell == 2] += [[0.5], [0.0], [0.5]] |
---|
| 174 | center[:, lattice_point % shapes_per_cell == 3] += [[0.5], [0.5], [0.0]] |
---|
| 175 | |
---|
| 176 | # Each lattice point has its own displacement from the ideal position. |
---|
| 177 | # Not checking that shapes do not overlap if displacement is too large. |
---|
| 178 | offset = shuffle*(randn(3, number_of_lattice_points) if shuffle < 0.3 |
---|
| 179 | else rand(3, number_of_lattice_points)) |
---|
| 180 | center += offset[:, cell_index] |
---|
| 181 | |
---|
| 182 | # Each lattice point has its own rotation. Rotate the point prior to |
---|
| 183 | # applying any displacement. |
---|
| 184 | # rotation = rotate*(randn(size=(shapes, 3)) if shuffle < 30 else rand(size=(nsamples, 3))) |
---|
| 185 | # for k in shapes: points[k] = rotation[k]*points[k] |
---|
| 186 | points += center*(np.array([spacing])*np.array(self.dims)).T |
---|
| 187 | return points |
---|
| 188 | |
---|
[cfa28d3] | 189 | class Composite(Shape): |
---|
| 190 | def __init__(self, shapes, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
| 191 | self.shapes = shapes |
---|
| 192 | self.rotate(*orientation) |
---|
| 193 | self.shift(*center) |
---|
| 194 | |
---|
| 195 | # Find the worst case distance between any two points amongst a set |
---|
| 196 | # of shapes independent of orientation. This could easily be a |
---|
| 197 | # factor of two worse than necessary, e.g., a pair of thin rods |
---|
| 198 | # end-to-end vs the same pair side-by-side. |
---|
| 199 | distances = [((s1.r_max + s2.r_max)/2 |
---|
| 200 | + sqrt(np.sum((s1.center - s2.center)**2))) |
---|
| 201 | for s1 in shapes |
---|
| 202 | for s2 in shapes] |
---|
| 203 | self.r_max = max(distances + [s.r_max for s in shapes]) |
---|
[032646d] | 204 | self.volume = sum(shape.volume for shape in self.shapes) |
---|
[cfa28d3] | 205 | |
---|
| 206 | def sample(self, density): |
---|
| 207 | values, points = zip(*(shape.sample(density) for shape in self.shapes)) |
---|
| 208 | return np.hstack(values), self._adjust(np.vstack(points)) |
---|
| 209 | |
---|
| 210 | class Box(Shape): |
---|
| 211 | def __init__(self, a, b, c, |
---|
| 212 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
| 213 | self.value = np.asarray(value) |
---|
| 214 | self.rotate(*orientation) |
---|
| 215 | self.shift(*center) |
---|
| 216 | self.a, self.b, self.c = a, b, c |
---|
| 217 | self._scale = np.array([a/2, b/2, c/2])[None, :] |
---|
| 218 | self.r_max = sqrt(a**2 + b**2 + c**2) |
---|
[032646d] | 219 | self.dims = a, b, c |
---|
| 220 | self.volume = a*b*c |
---|
[cfa28d3] | 221 | |
---|
| 222 | def sample(self, density): |
---|
[032646d] | 223 | num_points = poisson(density*self.volume) |
---|
[cfa28d3] | 224 | points = self._scale*uniform(-1, 1, size=(num_points, 3)) |
---|
| 225 | values = self.value.repeat(points.shape[0]) |
---|
| 226 | return values, self._adjust(points) |
---|
| 227 | |
---|
| 228 | class EllipticalCylinder(Shape): |
---|
| 229 | def __init__(self, ra, rb, length, |
---|
| 230 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
| 231 | self.value = np.asarray(value) |
---|
| 232 | self.rotate(*orientation) |
---|
| 233 | self.shift(*center) |
---|
| 234 | self.ra, self.rb, self.length = ra, rb, length |
---|
| 235 | self._scale = np.array([ra, rb, length/2])[None, :] |
---|
| 236 | self.r_max = sqrt(4*max(ra, rb)**2 + length**2) |
---|
[032646d] | 237 | self.dims = 2*ra, 2*rb, length |
---|
| 238 | self.volume = pi*ra*rb*length |
---|
[cfa28d3] | 239 | |
---|
| 240 | def sample(self, density): |
---|
[032646d] | 241 | # randomly sample from a box of side length 2*r, excluding anything |
---|
| 242 | # not in the cylinder |
---|
[cfa28d3] | 243 | num_points = poisson(density*4*self.ra*self.rb*self.length) |
---|
| 244 | points = uniform(-1, 1, size=(num_points, 3)) |
---|
| 245 | radius = points[:, 0]**2 + points[:, 1]**2 |
---|
[893bea2] | 246 | points = points[radius <= 1] |
---|
[cfa28d3] | 247 | values = self.value.repeat(points.shape[0]) |
---|
[893bea2] | 248 | return values, self._adjust(self._scale*points) |
---|
| 249 | |
---|
| 250 | class EllipticalBicelle(Shape): |
---|
| 251 | def __init__(self, ra, rb, length, |
---|
| 252 | thick_rim, thick_face, |
---|
| 253 | value_core, value_rim, value_face, |
---|
| 254 | center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
| 255 | self.rotate(*orientation) |
---|
| 256 | self.shift(*center) |
---|
| 257 | self.value = value_core |
---|
| 258 | self.ra, self.rb, self.length = ra, rb, length |
---|
| 259 | self.thick_rim, self.thick_face = thick_rim, thick_face |
---|
| 260 | self.value_rim, self.value_face = value_rim, value_face |
---|
| 261 | |
---|
| 262 | # reset cylinder to outer dimensions for calculating scale, etc. |
---|
| 263 | ra = self.ra + self.thick_rim |
---|
| 264 | rb = self.rb + self.thick_rim |
---|
| 265 | length = self.length + 2*self.thick_face |
---|
| 266 | self._scale = np.array([ra, rb, length/2])[None, :] |
---|
| 267 | self.r_max = sqrt(4*max(ra, rb)**2 + length**2) |
---|
| 268 | self.dims = 2*ra, 2*rb, length |
---|
| 269 | self.volume = pi*ra*rb*length |
---|
| 270 | |
---|
| 271 | def sample(self, density): |
---|
| 272 | # randomly sample from a box of side length 2*r, excluding anything |
---|
| 273 | # not in the cylinder |
---|
| 274 | ra = self.ra + self.thick_rim |
---|
| 275 | rb = self.rb + self.thick_rim |
---|
| 276 | length = self.length + 2*self.thick_face |
---|
| 277 | num_points = poisson(density*4*ra*rb*length) |
---|
| 278 | points = uniform(-1, 1, size=(num_points, 3)) |
---|
| 279 | radius = points[:, 0]**2 + points[:, 1]**2 |
---|
| 280 | points = points[radius <= 1] |
---|
| 281 | # set all to core value first |
---|
| 282 | values = np.ones_like(points[:, 0])*self.value |
---|
| 283 | # then set value to face value if |z| > face/(length/2)) |
---|
| 284 | values[abs(points[:, 2]) > self.length/(self.length + 2*self.thick_face)] = self.value_face |
---|
| 285 | # finally set value to rim value if outside the core ellipse |
---|
[362a66f] | 286 | radius = (points[:, 0]**2*(1 + self.thick_rim/self.ra)**2 |
---|
| 287 | + points[:, 1]**2*(1 + self.thick_rim/self.rb)**2) |
---|
[893bea2] | 288 | values[radius>1] = self.value_rim |
---|
| 289 | return values, self._adjust(self._scale*points) |
---|
[cfa28d3] | 290 | |
---|
| 291 | class TriaxialEllipsoid(Shape): |
---|
| 292 | def __init__(self, ra, rb, rc, |
---|
| 293 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
| 294 | self.value = np.asarray(value) |
---|
| 295 | self.rotate(*orientation) |
---|
| 296 | self.shift(*center) |
---|
| 297 | self.ra, self.rb, self.rc = ra, rb, rc |
---|
| 298 | self._scale = np.array([ra, rb, rc])[None, :] |
---|
| 299 | self.r_max = 2*max(ra, rb, rc) |
---|
[032646d] | 300 | self.dims = 2*ra, 2*rb, 2*rc |
---|
| 301 | self.volume = 4*pi/3 * ra * rb * rc |
---|
[cfa28d3] | 302 | |
---|
| 303 | def sample(self, density): |
---|
| 304 | # randomly sample from a box of side length 2*r, excluding anything |
---|
| 305 | # not in the ellipsoid |
---|
| 306 | num_points = poisson(density*8*self.ra*self.rb*self.rc) |
---|
| 307 | points = uniform(-1, 1, size=(num_points, 3)) |
---|
| 308 | radius = np.sum(points**2, axis=1) |
---|
| 309 | points = self._scale*points[radius <= 1] |
---|
| 310 | values = self.value.repeat(points.shape[0]) |
---|
| 311 | return values, self._adjust(points) |
---|
| 312 | |
---|
| 313 | class Helix(Shape): |
---|
| 314 | def __init__(self, helix_radius, helix_pitch, tube_radius, tube_length, |
---|
| 315 | value, center=(0, 0, 0), orientation=(0, 0, 0)): |
---|
| 316 | self.value = np.asarray(value) |
---|
| 317 | self.rotate(*orientation) |
---|
| 318 | self.shift(*center) |
---|
[032646d] | 319 | helix_length = helix_pitch * tube_length/sqrt(helix_radius**2 + helix_pitch**2) |
---|
| 320 | total_radius = self.helix_radius + self.tube_radius |
---|
[cfa28d3] | 321 | self.helix_radius, self.helix_pitch = helix_radius, helix_pitch |
---|
| 322 | self.tube_radius, self.tube_length = tube_radius, tube_length |
---|
[032646d] | 323 | self.r_max = sqrt(4*total_radius + (helix_length + 2*tube_radius)**2) |
---|
| 324 | self.dims = 2*total_radius, 2*total_radius, helix_length |
---|
[cfa28d3] | 325 | # small tube radius approximation; for larger tubes need to account |
---|
| 326 | # for the fact that the inner length is much shorter than the outer |
---|
| 327 | # length |
---|
[032646d] | 328 | self.volume = pi*self.tube_radius**2*self.tube_length |
---|
[cfa28d3] | 329 | |
---|
| 330 | def points(self, density): |
---|
| 331 | num_points = poisson(density*4*self.tube_radius**2*self.tube_length) |
---|
| 332 | points = uniform(-1, 1, size=(num_points, 3)) |
---|
| 333 | radius = points[:, 0]**2 + points[:, 1]**2 |
---|
| 334 | points = points[radius <= 1] |
---|
| 335 | |
---|
| 336 | # Based on math stackexchange answer by Jyrki Lahtonen |
---|
| 337 | # https://math.stackexchange.com/a/461637 |
---|
| 338 | # with helix along z rather than x [so tuples in answer are (z, x, y)] |
---|
| 339 | # and with random points in the cross section (p1, p2) rather than |
---|
| 340 | # uniform points on the surface (cos u, sin u). |
---|
| 341 | a, R = self.tube_radius, self.helix_radius |
---|
| 342 | h = self.helix_pitch |
---|
| 343 | scale = 1/sqrt(R**2 + h**2) |
---|
| 344 | t = points[:, 3] * (self.tube_length * scale/2) |
---|
| 345 | cos_t, sin_t = cos(t), sin(t) |
---|
| 346 | |
---|
| 347 | # rx = R*cos_t |
---|
| 348 | # ry = R*sin_t |
---|
| 349 | # rz = h*t |
---|
| 350 | # nx = -a * cos_t * points[:, 1] |
---|
| 351 | # ny = -a * sin_t * points[:, 1] |
---|
| 352 | # nz = 0 |
---|
| 353 | # bx = (a * h/scale) * sin_t * points[:, 2] |
---|
| 354 | # by = (-a * h/scale) * cos_t * points[:, 2] |
---|
| 355 | # bz = a*R/scale |
---|
| 356 | # x = rx + nx + bx |
---|
| 357 | # y = ry + ny + by |
---|
| 358 | # z = rz + nz + bz |
---|
| 359 | u, v = (R - a*points[:, 1]), (a * h/scale)*points[:, 2] |
---|
| 360 | x = u * cos_t + v * sin_t |
---|
| 361 | y = u * sin_t - v * cos_t |
---|
| 362 | z = a*R/scale + h * t |
---|
| 363 | |
---|
| 364 | points = np.hstack((x, y, z)) |
---|
| 365 | values = self.value.repeat(points.shape[0]) |
---|
| 366 | return values, self._adjust(points) |
---|
| 367 | |
---|
[7d0afa3] | 368 | def csbox(a=10, b=20, c=30, da=1, db=2, dc=3, slda=1, sldb=2, sldc=3, sld_core=4): |
---|
| 369 | core = Box(a, b, c, sld_core) |
---|
| 370 | side_a = Box(da, b, c, slda, center=((a+da)/2, 0, 0)) |
---|
| 371 | side_b = Box(a, db, c, sldb, center=(0, (b+db)/2, 0)) |
---|
| 372 | side_c = Box(a, b, dc, sldc, center=(0, 0, (c+dc)/2)) |
---|
| 373 | side_a2 = copy(side_a).shift(-a-da, 0, 0) |
---|
| 374 | side_b2 = copy(side_b).shift(0, -b-db, 0) |
---|
| 375 | side_c2 = copy(side_c).shift(0, 0, -c-dc) |
---|
| 376 | shape = Composite((core, side_a, side_b, side_c, side_a2, side_b2, side_c2)) |
---|
[032646d] | 377 | shape.dims = 2*da+a, 2*db+b, 2*dc+c |
---|
[7d0afa3] | 378 | return shape |
---|
| 379 | |
---|
| 380 | def _Iqxy(values, x, y, z, qa, qb, qc): |
---|
| 381 | """I(q) = |sum V(r) rho(r) e^(1j q.r)|^2 / sum V(r)""" |
---|
| 382 | Iq = [abs(np.sum(values*np.exp(1j*(qa_k*x + qb_k*y + qc_k*z))))**2 |
---|
| 383 | for qa_k, qb_k, qc_k in zip(qa.flat, qb.flat, qc.flat)] |
---|
| 384 | return Iq |
---|
| 385 | |
---|
| 386 | if USE_NUMBA: |
---|
| 387 | # Override simple numpy solution with numba if available |
---|
[8fb2a94] | 388 | from numba import njit |
---|
| 389 | @njit("f8[:](f8[:],f8[:],f8[:],f8[:],f8[:],f8[:],f8[:])") |
---|
[8bd379a] | 390 | def _Iqxy(values, x, y, z, qa, qb, qc): |
---|
[8fb2a94] | 391 | Iq = np.zeros_like(qa) |
---|
| 392 | for j in range(len(Iq)): |
---|
| 393 | total = 0. + 0j |
---|
[8bd379a] | 394 | for k in range(len(values)): |
---|
[8fb2a94] | 395 | total += values[k]*np.exp(1j*(qa[j]*x[k] + qb[j]*y[k] + qc[j]*z[k])) |
---|
| 396 | Iq[j] = abs(total)**2 |
---|
| 397 | return Iq |
---|
| 398 | |
---|
| 399 | def calc_Iqxy(qx, qy, rho, points, volume=1.0, view=(0, 0, 0)): |
---|
[226473d] | 400 | qx, qy = np.broadcast_arrays(qx, qy) |
---|
| 401 | qa, qb, qc = invert_view(qx, qy, view) |
---|
| 402 | rho, volume = np.broadcast_arrays(rho, volume) |
---|
| 403 | values = rho*volume |
---|
[8fb2a94] | 404 | x, y, z = points.T |
---|
[8bd379a] | 405 | values, x, y, z, qa, qb, qc = [np.asarray(v, 'd') |
---|
| 406 | for v in (values, x, y, z, qa, qb, qc)] |
---|
[226473d] | 407 | |
---|
| 408 | # I(q) = |sum V(r) rho(r) e^(1j q.r)|^2 / sum V(r) |
---|
[7d0afa3] | 409 | Iq = _Iqxy(values, x, y, z, qa.flatten(), qb.flatten(), qc.flatten()) |
---|
[8fb2a94] | 410 | return np.asarray(Iq).reshape(qx.shape) / np.sum(volume) |
---|
[226473d] | 411 | |
---|
[cfa28d3] | 412 | def _calc_Pr_nonuniform(r, rho, points): |
---|
| 413 | # Make Pr a little be bigger than necessary so that only distances |
---|
| 414 | # min < d < max end up in Pr |
---|
| 415 | n_max = len(r)+1 |
---|
| 416 | extended_Pr = np.zeros(n_max+1, 'd') |
---|
| 417 | # r refers to bin centers; find corresponding bin edges |
---|
| 418 | bins = bin_edges(r) |
---|
| 419 | t_next = time.clock() + 3 |
---|
| 420 | for k, rho_k in enumerate(rho[:-1]): |
---|
| 421 | distance = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1)) |
---|
| 422 | weights = rho_k * rho[k+1:] |
---|
| 423 | index = np.searchsorted(bins, distance) |
---|
| 424 | # Note: indices may be duplicated, so "Pr[index] += w" will not work!! |
---|
| 425 | extended_Pr += np.bincount(index, weights, n_max+1) |
---|
| 426 | t = time.clock() |
---|
| 427 | if t > t_next: |
---|
| 428 | t_next = t + 3 |
---|
| 429 | print("processing %d of %d"%(k, len(rho)-1)) |
---|
| 430 | Pr = extended_Pr[1:-1] |
---|
[8fb2a94] | 431 | return Pr |
---|
[cfa28d3] | 432 | |
---|
[8fb2a94] | 433 | def _calc_Pr_uniform(r, rho, points): |
---|
[cfa28d3] | 434 | # Make Pr a little be bigger than necessary so that only distances |
---|
| 435 | # min < d < max end up in Pr |
---|
[8fb2a94] | 436 | dr, n_max = r[0], len(r) |
---|
[cfa28d3] | 437 | extended_Pr = np.zeros(n_max+1, 'd') |
---|
| 438 | t0 = time.clock() |
---|
| 439 | t_next = t0 + 3 |
---|
| 440 | for k, rho_k in enumerate(rho[:-1]): |
---|
| 441 | distances = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1)) |
---|
| 442 | weights = rho_k * rho[k+1:] |
---|
[8fb2a94] | 443 | index = np.minimum(np.asarray(distances/dr, 'i'), n_max) |
---|
[cfa28d3] | 444 | # Note: indices may be duplicated, so "Pr[index] += w" will not work!! |
---|
| 445 | extended_Pr += np.bincount(index, weights, n_max+1) |
---|
| 446 | t = time.clock() |
---|
| 447 | if t > t_next: |
---|
| 448 | t_next = t + 3 |
---|
| 449 | print("processing %d of %d"%(k, len(rho)-1)) |
---|
| 450 | #print("time py:", time.clock() - t0) |
---|
| 451 | Pr = extended_Pr[:-1] |
---|
| 452 | # Make Pr independent of sampling density. The factor of 2 comes because |
---|
| 453 | # we are only accumulating the upper triangular distances. |
---|
| 454 | #Pr = Pr * 2 / len(rho)**2 |
---|
[8fb2a94] | 455 | return Pr |
---|
[cfa28d3] | 456 | |
---|
| 457 | # Can get an additional 2x by going to C. Cuda/OpenCL will allow even |
---|
| 458 | # more speedup, though still bounded by the n^2 cose. |
---|
| 459 | """ |
---|
| 460 | void pdfcalc(int n, const double *pts, const double *rho, |
---|
| 461 | int nPr, double *Pr, double rstep) |
---|
| 462 | { |
---|
| 463 | int i,j; |
---|
| 464 | |
---|
| 465 | for (i=0; i<n-2; i++) { |
---|
| 466 | for (j=i+1; j<=n-1; j++) { |
---|
| 467 | const double dxx=pts[3*i]-pts[3*j]; |
---|
| 468 | const double dyy=pts[3*i+1]-pts[3*j+1]; |
---|
| 469 | const double dzz=pts[3*i+2]-pts[3*j+2]; |
---|
| 470 | const double d=sqrt(dxx*dxx+dyy*dyy+dzz*dzz); |
---|
| 471 | const int k=rint(d/rstep); |
---|
| 472 | if (k < nPr) Pr[k]+=rho[i]*rho[j]; |
---|
| 473 | } |
---|
| 474 | } |
---|
| 475 | } |
---|
| 476 | """ |
---|
| 477 | |
---|
[7d0afa3] | 478 | if USE_NUMBA: |
---|
| 479 | # Override simple numpy solution with numba if available |
---|
[8fb2a94] | 480 | @njit("f8[:](f8[:], f8[:], f8[:,:])") |
---|
[7d0afa3] | 481 | def _calc_Pr_uniform(r, rho, points): |
---|
[8fb2a94] | 482 | dr = r[0] |
---|
| 483 | n_max = len(r) |
---|
| 484 | Pr = np.zeros_like(r) |
---|
| 485 | for j in range(len(rho) - 1): |
---|
| 486 | x, y, z = points[j, 0], points[j, 1], points[j, 2] |
---|
| 487 | for k in range(j+1, len(rho)): |
---|
| 488 | distance = sqrt((x - points[k, 0])**2 |
---|
| 489 | + (y - points[k, 1])**2 |
---|
| 490 | + (z - points[k, 2])**2) |
---|
| 491 | index = int(distance/dr) |
---|
| 492 | if index < n_max: |
---|
| 493 | Pr[index] += rho[j] * rho[k] |
---|
| 494 | # Make Pr independent of sampling density. The factor of 2 comes because |
---|
| 495 | # we are only accumulating the upper triangular distances. |
---|
| 496 | #Pr = Pr * 2 / len(rho)**2 |
---|
| 497 | return Pr |
---|
| 498 | |
---|
| 499 | |
---|
| 500 | def calc_Pr(r, rho, points): |
---|
| 501 | # P(r) with uniform steps in r is 3x faster; check if we are uniform |
---|
| 502 | # before continuing |
---|
[8bd379a] | 503 | r, rho, points = [np.asarray(v, 'd') for v in (r, rho, points)] |
---|
[8fb2a94] | 504 | if np.max(np.abs(np.diff(r) - r[0])) > r[0]*0.01: |
---|
| 505 | Pr = _calc_Pr_nonuniform(r, rho, points) |
---|
| 506 | else: |
---|
[7d0afa3] | 507 | Pr = _calc_Pr_uniform(r, rho, points) |
---|
[8fb2a94] | 508 | return Pr / Pr.max() |
---|
| 509 | |
---|
| 510 | |
---|
[cfa28d3] | 511 | def j0(x): |
---|
| 512 | return np.sinc(x/np.pi) |
---|
| 513 | |
---|
| 514 | def calc_Iq(q, r, Pr): |
---|
| 515 | Iq = np.array([simps(Pr * j0(qk*r), r) for qk in q]) |
---|
| 516 | #Iq = np.array([np.trapz(Pr * j0(qk*r), r) for qk in q]) |
---|
| 517 | Iq /= Iq[0] |
---|
| 518 | return Iq |
---|
| 519 | |
---|
| 520 | # NOTE: copied from sasmodels/resolution.py |
---|
| 521 | def bin_edges(x): |
---|
| 522 | """ |
---|
| 523 | Determine bin edges from bin centers, assuming that edges are centered |
---|
| 524 | between the bins. |
---|
| 525 | |
---|
| 526 | Note: this uses the arithmetic mean, which may not be appropriate for |
---|
| 527 | log-scaled data. |
---|
| 528 | """ |
---|
| 529 | if len(x) < 2 or (np.diff(x) < 0).any(): |
---|
| 530 | raise ValueError("Expected bins to be an increasing set") |
---|
| 531 | edges = np.hstack([ |
---|
| 532 | x[0] - 0.5*(x[1] - x[0]), # first point minus half first interval |
---|
| 533 | 0.5*(x[1:] + x[:-1]), # mid points of all central intervals |
---|
| 534 | x[-1] + 0.5*(x[-1] - x[-2]), # last point plus half last interval |
---|
| 535 | ]) |
---|
| 536 | return edges |
---|
| 537 | |
---|
[7d0afa3] | 538 | # -------------- plotters ---------------- |
---|
[3db96b0] | 539 | def plot_calc(r, Pr, q, Iq, theory=None, title=None): |
---|
[cfa28d3] | 540 | import matplotlib.pyplot as plt |
---|
| 541 | plt.subplot(211) |
---|
| 542 | plt.plot(r, Pr, '-', label="Pr") |
---|
| 543 | plt.xlabel('r (A)') |
---|
| 544 | plt.ylabel('Pr (1/A^2)') |
---|
[3db96b0] | 545 | if title is not None: |
---|
| 546 | plt.title(title) |
---|
[cfa28d3] | 547 | plt.subplot(212) |
---|
| 548 | plt.loglog(q, Iq, '-', label='from Pr') |
---|
| 549 | plt.xlabel('q (1/A') |
---|
| 550 | plt.ylabel('Iq') |
---|
| 551 | if theory is not None: |
---|
[7d0afa3] | 552 | plt.loglog(theory[0], theory[1]/theory[1][0], '-', label='analytic') |
---|
[cfa28d3] | 553 | plt.legend() |
---|
| 554 | |
---|
[3db96b0] | 555 | def plot_calc_2d(qx, qy, Iqxy, theory=None, title=None): |
---|
[226473d] | 556 | import matplotlib.pyplot as plt |
---|
| 557 | qx, qy = bin_edges(qx), bin_edges(qy) |
---|
| 558 | #qx, qy = np.meshgrid(qx, qy) |
---|
| 559 | if theory is not None: |
---|
| 560 | plt.subplot(121) |
---|
[3db96b0] | 561 | #plt.pcolor(qx, qy, np.log10(Iqxy)) |
---|
| 562 | extent = [qx[0], qx[-1], qy[0], qy[-1]] |
---|
| 563 | plt.imshow(np.log10(Iqxy), extent=extent, interpolation="nearest", |
---|
| 564 | origin='lower') |
---|
[226473d] | 565 | plt.xlabel('qx (1/A)') |
---|
| 566 | plt.ylabel('qy (1/A)') |
---|
[3db96b0] | 567 | plt.axis('equal') |
---|
| 568 | plt.axis(extent) |
---|
| 569 | #plt.grid(True) |
---|
| 570 | if title is not None: |
---|
| 571 | plt.title(title) |
---|
[226473d] | 572 | if theory is not None: |
---|
| 573 | plt.subplot(122) |
---|
[3db96b0] | 574 | plt.imshow(np.log10(theory), extent=extent, interpolation="nearest", |
---|
| 575 | origin='lower') |
---|
| 576 | plt.axis('equal') |
---|
| 577 | plt.axis(extent) |
---|
[226473d] | 578 | plt.xlabel('qx (1/A)') |
---|
| 579 | |
---|
[cfa28d3] | 580 | def plot_points(rho, points): |
---|
| 581 | import mpl_toolkits.mplot3d |
---|
| 582 | import matplotlib.pyplot as plt |
---|
| 583 | |
---|
| 584 | ax = plt.axes(projection='3d') |
---|
| 585 | try: |
---|
| 586 | ax.axis('square') |
---|
| 587 | except Exception: |
---|
| 588 | pass |
---|
| 589 | n = len(points) |
---|
[2ab331f] | 590 | #print("len points", n) |
---|
| 591 | index = np.random.choice(n, size=500) if n > 500 else slice(None, None) |
---|
[cfa28d3] | 592 | ax.scatter(points[index, 0], points[index, 1], points[index, 2], c=rho[index]) |
---|
[4f6f9431] | 593 | # make square axes |
---|
| 594 | minmax = np.array([points.min(), points.max()]) |
---|
| 595 | ax.scatter(minmax, minmax, minmax, c='w') |
---|
[cfa28d3] | 596 | #low, high = points.min(axis=0), points.max(axis=0) |
---|
| 597 | #ax.axis([low[0], high[0], low[1], high[1], low[2], high[2]]) |
---|
[893bea2] | 598 | ax.set_xlabel("x") |
---|
| 599 | ax.set_ylabel("y") |
---|
| 600 | ax.set_zlabel("z") |
---|
[cfa28d3] | 601 | ax.autoscale(True) |
---|
| 602 | |
---|
[7d0afa3] | 603 | # ----------- Analytic models -------------- |
---|
[97be877] | 604 | def sas_sinx_x(x): |
---|
| 605 | with np.errstate(all='ignore'): |
---|
| 606 | retvalue = sin(x)/x |
---|
| 607 | retvalue[x == 0.] = 1. |
---|
| 608 | return retvalue |
---|
| 609 | |
---|
[cfa28d3] | 610 | def sas_2J1x_x(x): |
---|
| 611 | with np.errstate(all='ignore'): |
---|
| 612 | retvalue = 2*J1(x)/x |
---|
| 613 | retvalue[x == 0] = 1. |
---|
| 614 | return retvalue |
---|
| 615 | |
---|
| 616 | def sas_3j1x_x(x): |
---|
| 617 | """return 3*j1(x)/x""" |
---|
| 618 | with np.errstate(all='ignore'): |
---|
| 619 | retvalue = 3*(sin(x) - x*cos(x))/x**3 |
---|
| 620 | retvalue[x == 0.] = 1. |
---|
| 621 | return retvalue |
---|
| 622 | |
---|
| 623 | def cylinder_Iq(q, radius, length): |
---|
[a1c32c2] | 624 | z, w = leggauss(76) |
---|
| 625 | cos_alpha = (z+1)/2 |
---|
| 626 | sin_alpha = sqrt(1.0 - cos_alpha**2) |
---|
[cfa28d3] | 627 | Iq = np.empty_like(q) |
---|
| 628 | for k, qk in enumerate(q): |
---|
[a1c32c2] | 629 | qab, qc = qk*sin_alpha, qk*cos_alpha |
---|
[7d0afa3] | 630 | Fq = sas_2J1x_x(qab*radius) * sas_sinx_x(qc*length/2) |
---|
[a1c32c2] | 631 | Iq[k] = np.sum(w*Fq**2) |
---|
[032646d] | 632 | Iq = Iq |
---|
[cfa28d3] | 633 | return Iq |
---|
| 634 | |
---|
[226473d] | 635 | def cylinder_Iqxy(qx, qy, radius, length, view=(0, 0, 0)): |
---|
| 636 | qa, qb, qc = invert_view(qx, qy, view) |
---|
[7d0afa3] | 637 | qab = sqrt(qa**2 + qb**2) |
---|
| 638 | Fq = sas_2J1x_x(qab*radius) * sas_sinx_x(qc*length/2) |
---|
[226473d] | 639 | Iq = Fq**2 |
---|
| 640 | return Iq.reshape(qx.shape) |
---|
| 641 | |
---|
[cfa28d3] | 642 | def sphere_Iq(q, radius): |
---|
| 643 | Iq = sas_3j1x_x(q*radius)**2 |
---|
[032646d] | 644 | return Iq |
---|
[cfa28d3] | 645 | |
---|
[7d0afa3] | 646 | def box_Iq(q, a, b, c): |
---|
| 647 | z, w = leggauss(76) |
---|
| 648 | outer_sum = np.zeros_like(q) |
---|
| 649 | for cos_alpha, outer_w in zip((z+1)/2, w): |
---|
| 650 | sin_alpha = sqrt(1.0-cos_alpha*cos_alpha) |
---|
| 651 | qc = q*cos_alpha |
---|
| 652 | siC = c*sas_sinx_x(c*qc/2) |
---|
| 653 | inner_sum = np.zeros_like(q) |
---|
| 654 | for beta, inner_w in zip((z + 1)*pi/4, w): |
---|
| 655 | qa, qb = q*sin_alpha*sin(beta), q*sin_alpha*cos(beta) |
---|
| 656 | siA = a*sas_sinx_x(a*qa/2) |
---|
| 657 | siB = b*sas_sinx_x(b*qb/2) |
---|
| 658 | Fq = siA*siB*siC |
---|
| 659 | inner_sum += inner_w * Fq**2 |
---|
| 660 | outer_sum += outer_w * inner_sum |
---|
| 661 | Iq = outer_sum / 4 # = outer*um*zm*8.0/(4.0*M_PI) |
---|
[032646d] | 662 | return Iq |
---|
[7d0afa3] | 663 | |
---|
| 664 | def box_Iqxy(qx, qy, a, b, c, view=(0, 0, 0)): |
---|
| 665 | qa, qb, qc = invert_view(qx, qy, view) |
---|
| 666 | sia = sas_sinx_x(qa*a/2) |
---|
| 667 | sib = sas_sinx_x(qb*b/2) |
---|
| 668 | sic = sas_sinx_x(qc*c/2) |
---|
| 669 | Fq = sia*sib*sic |
---|
| 670 | Iq = Fq**2 |
---|
| 671 | return Iq.reshape(qx.shape) |
---|
| 672 | |
---|
[cfa28d3] | 673 | def csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core): |
---|
[a1c32c2] | 674 | z, w = leggauss(76) |
---|
| 675 | |
---|
[cfa28d3] | 676 | sld_solvent = 0 |
---|
| 677 | overlapping = False |
---|
| 678 | dr0 = sld_core - sld_solvent |
---|
| 679 | drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent |
---|
| 680 | tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc |
---|
| 681 | |
---|
[a1c32c2] | 682 | outer_sum = np.zeros_like(q) |
---|
| 683 | for cos_alpha, outer_w in zip((z+1)/2, w): |
---|
[cfa28d3] | 684 | sin_alpha = sqrt(1.0-cos_alpha*cos_alpha) |
---|
| 685 | qc = q*cos_alpha |
---|
[7d0afa3] | 686 | siC = c*sas_sinx_x(c*qc/2) |
---|
| 687 | siCt = tC*sas_sinx_x(tC*qc/2) |
---|
[a1c32c2] | 688 | inner_sum = np.zeros_like(q) |
---|
| 689 | for beta, inner_w in zip((z + 1)*pi/4, w): |
---|
[cfa28d3] | 690 | qa, qb = q*sin_alpha*sin(beta), q*sin_alpha*cos(beta) |
---|
[7d0afa3] | 691 | siA = a*sas_sinx_x(a*qa/2) |
---|
| 692 | siB = b*sas_sinx_x(b*qb/2) |
---|
| 693 | siAt = tA*sas_sinx_x(tA*qa/2) |
---|
| 694 | siBt = tB*sas_sinx_x(tB*qb/2) |
---|
[cfa28d3] | 695 | if overlapping: |
---|
[a1c32c2] | 696 | Fq = (dr0*siA*siB*siC |
---|
| 697 | + drA*(siAt-siA)*siB*siC |
---|
| 698 | + drB*siAt*(siBt-siB)*siC |
---|
| 699 | + drC*siAt*siBt*(siCt-siC)) |
---|
[cfa28d3] | 700 | else: |
---|
[a1c32c2] | 701 | Fq = (dr0*siA*siB*siC |
---|
| 702 | + drA*(siAt-siA)*siB*siC |
---|
| 703 | + drB*siA*(siBt-siB)*siC |
---|
| 704 | + drC*siA*siB*(siCt-siC)) |
---|
| 705 | inner_sum += inner_w * Fq**2 |
---|
| 706 | outer_sum += outer_w * inner_sum |
---|
| 707 | Iq = outer_sum / 4 # = outer*um*zm*8.0/(4.0*M_PI) |
---|
[cfa28d3] | 708 | return Iq/Iq[0] |
---|
| 709 | |
---|
[97be877] | 710 | def csbox_Iqxy(qx, qy, a, b, c, da, db, dc, slda, sldb, sldc, sld_core, view=(0,0,0)): |
---|
| 711 | qa, qb, qc = invert_view(qx, qy, view) |
---|
[226473d] | 712 | |
---|
[97be877] | 713 | sld_solvent = 0 |
---|
| 714 | overlapping = False |
---|
| 715 | dr0 = sld_core - sld_solvent |
---|
| 716 | drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent |
---|
| 717 | tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc |
---|
| 718 | siA = a*sas_sinx_x(a*qa/2) |
---|
| 719 | siB = b*sas_sinx_x(b*qb/2) |
---|
| 720 | siC = c*sas_sinx_x(c*qc/2) |
---|
| 721 | siAt = tA*sas_sinx_x(tA*qa/2) |
---|
| 722 | siBt = tB*sas_sinx_x(tB*qb/2) |
---|
| 723 | siCt = tC*sas_sinx_x(tC*qc/2) |
---|
| 724 | Fq = (dr0*siA*siB*siC |
---|
| 725 | + drA*(siAt-siA)*siB*siC |
---|
| 726 | + drB*siA*(siBt-siB)*siC |
---|
| 727 | + drC*siA*siB*(siCt-siC)) |
---|
| 728 | Iq = Fq**2 |
---|
| 729 | return Iq.reshape(qx.shape) |
---|
[226473d] | 730 | |
---|
[893bea2] | 731 | def sasmodels_Iq(kernel, q, pars): |
---|
| 732 | from sasmodels.data import empty_data1D |
---|
| 733 | from sasmodels.direct_model import DirectModel |
---|
| 734 | data = empty_data1D(q) |
---|
| 735 | calculator = DirectModel(data, kernel) |
---|
| 736 | Iq = calculator(**pars) |
---|
| 737 | return Iq |
---|
| 738 | |
---|
| 739 | def sasmodels_Iqxy(kernel, qx, qy, pars, view): |
---|
| 740 | from sasmodels.data import Data2D |
---|
| 741 | from sasmodels.direct_model import DirectModel |
---|
| 742 | Iq = 100 * np.ones_like(qx) |
---|
| 743 | data = Data2D(x=qx, y=qy, z=Iq, dx=None, dy=None, dz=np.sqrt(Iq)) |
---|
[2a39ca4] | 744 | data.x_bins = qx[0, :] |
---|
| 745 | data.y_bins = qy[:, 0] |
---|
[893bea2] | 746 | data.filename = "fake data" |
---|
| 747 | |
---|
| 748 | calculator = DirectModel(data, kernel) |
---|
| 749 | pars_plus_view = pars.copy() |
---|
| 750 | pars_plus_view.update(theta=view[0], phi=view[1], psi=view[2]) |
---|
| 751 | Iqxy = calculator(**pars_plus_view) |
---|
| 752 | return Iqxy.reshape(qx.shape) |
---|
| 753 | |
---|
| 754 | def wrap_sasmodel(name, **pars): |
---|
| 755 | from sasmodels.core import load_model |
---|
| 756 | kernel = load_model(name) |
---|
| 757 | fn = lambda q: sasmodels_Iq(kernel, q, pars) |
---|
| 758 | fn_xy = lambda qx, qy, view: sasmodels_Iqxy(kernel, qx, qy, pars, view) |
---|
| 759 | return fn, fn_xy |
---|
| 760 | |
---|
| 761 | |
---|
[7d0afa3] | 762 | # --------- Test cases ----------- |
---|
[cfa28d3] | 763 | |
---|
[032646d] | 764 | def build_cylinder(radius=25, length=125, rho=2.): |
---|
[226473d] | 765 | shape = EllipticalCylinder(radius, radius, length, rho) |
---|
[7d0afa3] | 766 | fn = lambda q: cylinder_Iq(q, radius, length)*rho**2 |
---|
| 767 | fn_xy = lambda qx, qy, view: cylinder_Iqxy(qx, qy, radius, length, view=view)*rho**2 |
---|
| 768 | return shape, fn, fn_xy |
---|
[226473d] | 769 | |
---|
[2a39ca4] | 770 | DEFAULT_SPHERE_RADIUS = 125 |
---|
| 771 | DEFAULT_SPHERE_CONTRAST = 2 |
---|
| 772 | def build_sphere(radius=DEFAULT_SPHERE_RADIUS, rho=DEFAULT_SPHERE_CONTRAST): |
---|
[cfa28d3] | 773 | shape = TriaxialEllipsoid(radius, radius, radius, rho) |
---|
[032646d] | 774 | fn = lambda q: sphere_Iq(q, radius)*rho**2 |
---|
[7d0afa3] | 775 | fn_xy = lambda qx, qy, view: sphere_Iq(np.sqrt(qx**2+qy**2), radius)*rho**2 |
---|
| 776 | return shape, fn, fn_xy |
---|
| 777 | |
---|
[032646d] | 778 | def build_box(a=10, b=20, c=30, rho=2.): |
---|
[7d0afa3] | 779 | shape = Box(a, b, c, rho) |
---|
| 780 | fn = lambda q: box_Iq(q, a, b, c)*rho**2 |
---|
| 781 | fn_xy = lambda qx, qy, view: box_Iqxy(qx, qy, a, b, c, view=view)*rho**2 |
---|
| 782 | return shape, fn, fn_xy |
---|
| 783 | |
---|
[032646d] | 784 | def build_csbox(a=10, b=20, c=30, da=1, db=2, dc=3, slda=1, sldb=2, sldc=3, sld_core=4): |
---|
[7d0afa3] | 785 | shape = csbox(a, b, c, da, db, dc, slda, sldb, sldc, sld_core) |
---|
| 786 | fn = lambda q: csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core) |
---|
| 787 | fn_xy = lambda qx, qy, view: csbox_Iqxy(qx, qy, a, b, c, da, db, dc, |
---|
| 788 | slda, sldb, sldc, sld_core, view=view) |
---|
| 789 | return shape, fn, fn_xy |
---|
| 790 | |
---|
[893bea2] | 791 | def build_ellcyl(ra=25, rb=50, length=125, rho=2.): |
---|
| 792 | shape = EllipticalCylinder(ra, rb, length, rho) |
---|
| 793 | fn, fn_xy = wrap_sasmodel( |
---|
| 794 | 'elliptical_cylinder', |
---|
| 795 | scale=1, |
---|
| 796 | background=0, |
---|
| 797 | radius_minor=ra, |
---|
| 798 | axis_ratio=rb/ra, |
---|
| 799 | length=length, |
---|
| 800 | sld=rho, |
---|
| 801 | sld_solvent=0, |
---|
| 802 | ) |
---|
| 803 | return shape, fn, fn_xy |
---|
| 804 | |
---|
| 805 | def build_cscyl(ra=30, rb=90, length=30, thick_rim=8, thick_face=14, |
---|
| 806 | sld_core=4, sld_rim=1, sld_face=7): |
---|
| 807 | shape = EllipticalBicelle( |
---|
| 808 | ra=ra, rb=rb, length=length, |
---|
| 809 | thick_rim=thick_rim, thick_face=thick_face, |
---|
| 810 | value_core=sld_core, value_rim=sld_rim, value_face=sld_face, |
---|
| 811 | ) |
---|
| 812 | fn, fn_xy = wrap_sasmodel( |
---|
| 813 | 'core_shell_bicelle_elliptical', |
---|
| 814 | scale=1, |
---|
| 815 | background=0, |
---|
[4f6f9431] | 816 | radius=ra, |
---|
| 817 | x_core=rb/ra, |
---|
[893bea2] | 818 | length=length, |
---|
[362a66f] | 819 | thick_rim=thick_rim, |
---|
| 820 | thick_face=thick_face, |
---|
[893bea2] | 821 | sld_core=sld_core, |
---|
| 822 | sld_face=sld_face, |
---|
| 823 | sld_rim=sld_rim, |
---|
| 824 | sld_solvent=0, |
---|
| 825 | ) |
---|
| 826 | return shape, fn, fn_xy |
---|
| 827 | |
---|
[2a39ca4] | 828 | def build_sc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2, |
---|
| 829 | shuffle=0, rotate=0): |
---|
[032646d] | 830 | a, b, c = shape.dims |
---|
[2a39ca4] | 831 | corners= [copy(shape) |
---|
[032646d] | 832 | .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a, |
---|
| 833 | (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b, |
---|
| 834 | (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c) |
---|
| 835 | .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate)) |
---|
| 836 | for ix in range(nx) |
---|
| 837 | for iy in range(ny) |
---|
| 838 | for iz in range(nz)] |
---|
[2a39ca4] | 839 | lattice = Composite(corners) |
---|
[032646d] | 840 | return lattice |
---|
| 841 | |
---|
[2a39ca4] | 842 | def build_bcc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2, |
---|
| 843 | shuffle=0, rotate=0): |
---|
| 844 | a, b, c = shape.dims |
---|
| 845 | corners = [copy(shape) |
---|
| 846 | .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a, |
---|
| 847 | (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b, |
---|
| 848 | (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c) |
---|
| 849 | .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate)) |
---|
| 850 | for ix in range(nx) |
---|
| 851 | for iy in range(ny) |
---|
| 852 | for iz in range(nz)] |
---|
| 853 | centers = [copy(shape) |
---|
| 854 | .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a, |
---|
| 855 | (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b, |
---|
| 856 | (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c) |
---|
| 857 | .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate)) |
---|
| 858 | for ix in range(nx) |
---|
| 859 | for iy in range(ny) |
---|
| 860 | for iz in range(nz)] |
---|
| 861 | lattice = Composite(corners + centers) |
---|
| 862 | return lattice |
---|
| 863 | |
---|
| 864 | def build_fcc_lattice(shape, nx=1, ny=1, nz=1, dx=2, dy=2, dz=2, |
---|
| 865 | shuffle=0, rotate=0): |
---|
| 866 | a, b, c = shape.dims |
---|
| 867 | corners = [copy(shape) |
---|
| 868 | .shift((ix+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a, |
---|
| 869 | (iy+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b, |
---|
| 870 | (iz+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c) |
---|
| 871 | .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate)) |
---|
| 872 | for ix in range(nx) |
---|
| 873 | for iy in range(ny) |
---|
| 874 | for iz in range(nz)] |
---|
| 875 | faces_a = [copy(shape) |
---|
| 876 | .shift((ix+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a, |
---|
| 877 | (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b, |
---|
| 878 | (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c) |
---|
| 879 | .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate)) |
---|
| 880 | for ix in range(nx) |
---|
| 881 | for iy in range(ny) |
---|
| 882 | for iz in range(nz)] |
---|
| 883 | faces_b = [copy(shape) |
---|
| 884 | .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a, |
---|
| 885 | (iy+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b, |
---|
| 886 | (iz+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c) |
---|
| 887 | .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate)) |
---|
| 888 | for ix in range(nx) |
---|
| 889 | for iy in range(ny) |
---|
| 890 | for iz in range(nz)] |
---|
| 891 | faces_c = [copy(shape) |
---|
| 892 | .shift((ix+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dx*a, |
---|
| 893 | (iy+0.5+(randn() if shuffle < 0.3 else rand())*shuffle)*dy*b, |
---|
| 894 | (iz+0.0+(randn() if shuffle < 0.3 else rand())*shuffle)*dz*c) |
---|
| 895 | .rotate(*((randn(3) if rotate < 30 else rand(3))*rotate)) |
---|
| 896 | for ix in range(nx) |
---|
| 897 | for iy in range(ny) |
---|
| 898 | for iz in range(nz)] |
---|
| 899 | lattice = Composite(corners + faces_a + faces_b + faces_c) |
---|
| 900 | return lattice |
---|
[7d0afa3] | 901 | |
---|
| 902 | SHAPE_FUNCTIONS = OrderedDict([ |
---|
[893bea2] | 903 | ("cyl", build_cylinder), |
---|
| 904 | ("ellcyl", build_ellcyl), |
---|
[032646d] | 905 | ("sphere", build_sphere), |
---|
| 906 | ("box", build_box), |
---|
| 907 | ("csbox", build_csbox), |
---|
[893bea2] | 908 | ("cscyl", build_cscyl), |
---|
[7d0afa3] | 909 | ]) |
---|
| 910 | SHAPES = list(SHAPE_FUNCTIONS.keys()) |
---|
[2a39ca4] | 911 | LATTICE_FUNCTIONS = OrderedDict([ |
---|
| 912 | ("sc", build_sc_lattice), |
---|
| 913 | ("bcc", build_bcc_lattice), |
---|
| 914 | ("fcc", build_fcc_lattice), |
---|
| 915 | ]) |
---|
| 916 | LATTICE_TYPES = list(LATTICE_FUNCTIONS.keys()) |
---|
[7d0afa3] | 917 | |
---|
| 918 | def check_shape(title, shape, fn=None, show_points=False, |
---|
| 919 | mesh=100, qmax=1.0, r_step=0.01, samples=5000): |
---|
| 920 | rho_solvent = 0 |
---|
[032646d] | 921 | qmin = qmax/100. |
---|
[7d0afa3] | 922 | q = np.logspace(np.log10(qmin), np.log10(qmax), mesh) |
---|
| 923 | r = shape.r_bins(q, r_step=r_step) |
---|
[032646d] | 924 | sampling_density = samples / shape.volume |
---|
[2a39ca4] | 925 | print("sampling points") |
---|
[7d0afa3] | 926 | rho, points = shape.sample(sampling_density) |
---|
[2a39ca4] | 927 | print("calculating Pr") |
---|
[7d0afa3] | 928 | t0 = time.time() |
---|
| 929 | Pr = calc_Pr(r, rho-rho_solvent, points) |
---|
| 930 | print("calc Pr time", time.time() - t0) |
---|
| 931 | Iq = calc_Iq(q, r, Pr) |
---|
| 932 | theory = (q, fn(q)) if fn is not None else None |
---|
| 933 | |
---|
| 934 | import pylab |
---|
| 935 | if show_points: |
---|
[5778279] | 936 | plot_points(rho, points); pylab.figure() |
---|
[3db96b0] | 937 | plot_calc(r, Pr, q, Iq, theory=theory, title=title) |
---|
[7d0afa3] | 938 | pylab.gcf().canvas.set_window_title(title) |
---|
| 939 | pylab.show() |
---|
| 940 | |
---|
| 941 | def check_shape_2d(title, shape, fn=None, view=(0, 0, 0), show_points=False, |
---|
| 942 | mesh=100, qmax=1.0, samples=5000): |
---|
| 943 | rho_solvent = 0 |
---|
[893bea2] | 944 | #qx = np.linspace(0.0, qmax, mesh) |
---|
| 945 | #qy = np.linspace(0.0, qmax, mesh) |
---|
| 946 | qx = np.linspace(-qmax, qmax, mesh) |
---|
| 947 | qy = np.linspace(-qmax, qmax, mesh) |
---|
[7d0afa3] | 948 | Qx, Qy = np.meshgrid(qx, qy) |
---|
[032646d] | 949 | sampling_density = samples / shape.volume |
---|
[2a39ca4] | 950 | print("sampling points") |
---|
[032646d] | 951 | t0 = time.time() |
---|
[7d0afa3] | 952 | rho, points = shape.sample(sampling_density) |
---|
[032646d] | 953 | print("point generation time", time.time() - t0) |
---|
[7d0afa3] | 954 | t0 = time.time() |
---|
| 955 | Iqxy = calc_Iqxy(Qx, Qy, rho, points, view=view) |
---|
| 956 | print("calc Iqxy time", time.time() - t0) |
---|
[893bea2] | 957 | t0 = time.time() |
---|
[7d0afa3] | 958 | theory = fn(Qx, Qy, view) if fn is not None else None |
---|
[893bea2] | 959 | print("calc theory time", time.time() - t0) |
---|
[7d0afa3] | 960 | Iqxy += 0.001 * Iqxy.max() |
---|
| 961 | if theory is not None: |
---|
| 962 | theory += 0.001 * theory.max() |
---|
| 963 | |
---|
| 964 | import pylab |
---|
| 965 | if show_points: |
---|
| 966 | plot_points(rho, points); pylab.figure() |
---|
[3db96b0] | 967 | plot_calc_2d(qx, qy, Iqxy, theory=theory, title=title) |
---|
[7d0afa3] | 968 | pylab.gcf().canvas.set_window_title(title) |
---|
| 969 | pylab.show() |
---|
| 970 | |
---|
| 971 | def main(): |
---|
| 972 | parser = argparse.ArgumentParser( |
---|
| 973 | description="Compute scattering from realspace sampling", |
---|
| 974 | formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
---|
| 975 | ) |
---|
[032646d] | 976 | parser.add_argument('-d', '--dim', type=int, default=1, |
---|
| 977 | help='dimension 1 or 2') |
---|
| 978 | parser.add_argument('-m', '--mesh', type=int, default=100, |
---|
| 979 | help='number of mesh points') |
---|
| 980 | parser.add_argument('-s', '--samples', type=int, default=5000, |
---|
| 981 | help="number of sample points") |
---|
| 982 | parser.add_argument('-q', '--qmax', type=float, default=0.5, |
---|
| 983 | help='max q') |
---|
| 984 | parser.add_argument('-v', '--view', type=str, default='0,0,0', |
---|
| 985 | help='theta,phi,psi angles') |
---|
| 986 | parser.add_argument('-n', '--lattice', type=str, default='1,1,1', |
---|
| 987 | help='lattice size') |
---|
| 988 | parser.add_argument('-z', '--spacing', type=str, default='2,2,2', |
---|
[2a39ca4] | 989 | help='lattice spacing (relative to shape)') |
---|
| 990 | parser.add_argument('-t', '--type', choices=LATTICE_TYPES, |
---|
| 991 | default=LATTICE_TYPES[0], |
---|
| 992 | help='lattice type') |
---|
[032646d] | 993 | parser.add_argument('-r', '--rotate', type=float, default=0., |
---|
| 994 | help="rotation relative to lattice, gaussian < 30 degrees, uniform otherwise") |
---|
| 995 | parser.add_argument('-w', '--shuffle', type=float, default=0., |
---|
| 996 | help="position relative to lattice, gaussian < 0.3, uniform otherwise") |
---|
| 997 | parser.add_argument('-p', '--plot', action='store_true', |
---|
| 998 | help='plot points') |
---|
| 999 | parser.add_argument('shape', choices=SHAPES, nargs='?', default=SHAPES[0], |
---|
| 1000 | help='oriented shape') |
---|
[7d0afa3] | 1001 | parser.add_argument('pars', type=str, nargs='*', help='shape parameters') |
---|
| 1002 | opts = parser.parse_args() |
---|
| 1003 | pars = {key: float(value) for p in opts.pars for key, value in [p.split('=')]} |
---|
[032646d] | 1004 | nx, ny, nz = [int(v) for v in opts.lattice.split(',')] |
---|
| 1005 | dx, dy, dz = [float(v) for v in opts.spacing.split(',')] |
---|
[5778279] | 1006 | distortion, rotation = opts.shuffle, opts.rotate |
---|
[7d0afa3] | 1007 | shape, fn, fn_xy = SHAPE_FUNCTIONS[opts.shape](**pars) |
---|
[2a39ca4] | 1008 | view = tuple(float(v) for v in opts.view.split(',')) |
---|
[5778279] | 1009 | # If comparing a sphere in a cubic lattice, compare against the |
---|
| 1010 | # corresponding paracrystalline model. |
---|
| 1011 | if opts.shape == "sphere" and dx == dy == dz and nx*ny*nz > 1: |
---|
| 1012 | radius = pars.get('radius', DEFAULT_SPHERE_RADIUS) |
---|
| 1013 | model_name = opts.type + "_paracrystal" |
---|
| 1014 | model_pars = { |
---|
| 1015 | "scale": 1., |
---|
| 1016 | "background": 0., |
---|
| 1017 | "lattice_spacing": 2*radius*dx, |
---|
| 1018 | "lattice_distortion": distortion, |
---|
| 1019 | "radius": radius, |
---|
| 1020 | "sld": pars.get('rho', DEFAULT_SPHERE_CONTRAST), |
---|
| 1021 | "sld_solvent": 0., |
---|
| 1022 | "theta": view[0], |
---|
| 1023 | "phi": view[1], |
---|
| 1024 | "psi": view[2], |
---|
| 1025 | } |
---|
| 1026 | fn, fn_xy = wrap_sasmodel(model_name, **model_pars) |
---|
| 1027 | if nx*ny*nz > 1: |
---|
| 1028 | if rotation != 0: |
---|
| 1029 | print("building %s lattice"%opts.type) |
---|
| 1030 | build_lattice = LATTICE_FUNCTIONS[opts.type] |
---|
| 1031 | shape = build_lattice(shape, nx, ny, nz, dx, dy, dz, |
---|
| 1032 | distortion, rotation) |
---|
| 1033 | else: |
---|
| 1034 | shape.lattice(size=(nx, ny, nz), spacing=(dx, dy, dz), |
---|
| 1035 | type=opts.type, |
---|
| 1036 | rotation=rotation, distortion=distortion) |
---|
| 1037 | |
---|
[7d0afa3] | 1038 | title = "%s(%s)" % (opts.shape, " ".join(opts.pars)) |
---|
| 1039 | if opts.dim == 1: |
---|
| 1040 | check_shape(title, shape, fn, show_points=opts.plot, |
---|
| 1041 | mesh=opts.mesh, qmax=opts.qmax, samples=opts.samples) |
---|
| 1042 | else: |
---|
| 1043 | check_shape_2d(title, shape, fn_xy, view=view, show_points=opts.plot, |
---|
| 1044 | mesh=opts.mesh, qmax=opts.qmax, samples=opts.samples) |
---|
[cfa28d3] | 1045 | |
---|
[97be877] | 1046 | |
---|
[cfa28d3] | 1047 | if __name__ == "__main__": |
---|
[7d0afa3] | 1048 | main() |
---|