Changes in explore/realspace.py [a1c32c2:8fb2a94] in sasmodels

File:

• explore/realspace.py (modified) (11 diffs)

explore/realspace.py

@@ -44 +44 @@
     """
     return Rx(phi)*Ry(theta)*Rz(psi)
+
+def apply_view(points, view):
+    """
+    Apply the view transform (theta, phi, psi) to a set of points.
+
+    Points are stored in a 3 x n numpy array.
+
+    View angles are in degrees.
+    """
+    theta, phi, psi = view
+    return np.asarray((Rz(phi)*Ry(theta)*Rz(psi))*np.matrix(points.T)).T
+
+
+def invert_view(qx, qy, view):
+    """
+    Return (qa, qb, qc) for the (theta, phi, psi) view angle at detector
+    pixel (qx, qy).
+
+    View angles are in degrees.
+    """
+    theta, phi, psi = view
+    q = np.vstack((qx.flatten(), qy.flatten(), 0*qx.flatten()))
+    return np.asarray((Rz(-psi)*Ry(-theta)*Rz(-phi))*np.matrix(q))
+
 
 class Shape:
@@ -219 +243 @@
         values = self.value.repeat(points.shape[0])
         return values, self._adjust(points)
+
+NUMBA = False
+if NUMBA:
+    from numba import njit
+    @njit("f8[:](f8[:],f8[:],f8[:],f8[:],f8[:],f8[:],f8[:])")
+    def _Iqxy(values, x, y, z, qa, qb, qc):
+        Iq = np.zeros_like(qa)
+        for j in range(len(Iq)):
+            total = 0. + 0j
+            for k in range(len(Iq)):
+                total += values[k]*np.exp(1j*(qa[j]*x[k] + qb[j]*y[k] + qc[j]*z[k]))
+            Iq[j] = abs(total)**2
+        return Iq
+
+def calc_Iqxy(qx, qy, rho, points, volume=1.0, view=(0, 0, 0)):
+    qx, qy = np.broadcast_arrays(qx, qy)
+    qa, qb, qc = invert_view(qx, qy, view)
+    rho, volume = np.broadcast_arrays(rho, volume)
+    values = rho*volume
+    x, y, z = points.T
+
+    # I(q) = |sum V(r) rho(r) e^(1j q.r)|^2 / sum V(r)
+    if NUMBA:
+        Iq = _Iqxy(values, x, y, z, qa.flatten(), qb.flatten(), qc.flatten())
+    else:
+        Iq = [abs(np.sum(values*np.exp(1j*(qa_k*x + qb_k*y + qc_k*z))))**2
+              for qa_k, qb_k, qc_k in zip(qa.flat, qb.flat, qc.flat)]
+    return np.asarray(Iq).reshape(qx.shape) / np.sum(volume)
 
 def _calc_Pr_nonuniform(r, rho, points):
@@ -239 +291 @@
             print("processing %d of %d"%(k, len(rho)-1))
     Pr = extended_Pr[1:-1]
-    return Pr / Pr.max()
-
-def calc_Pr(r, rho, points):
-    # P(r) with uniform steps in r is 3x faster; check if we are uniform
-    # before continuing
-    if np.max(np.abs(np.diff(r) - r[0])) > r[0]*0.01:
-        return _calc_Pr_nonuniform(r, rho, points)
-
+    return Pr
+
+def _calc_Pr_uniform(r, rho, points):
     # Make Pr a little be bigger than necessary so that only distances
     # min < d < max end up in Pr
-    n_max = len(r)
+    dr, n_max = r[0], len(r)
     extended_Pr = np.zeros(n_max+1, 'd')
     t0 = time.clock()
     t_next = t0 + 3
-    row_zero = np.zeros(len(rho), 'i')
     for k, rho_k in enumerate(rho[:-1]):
         distances = sqrt(np.sum((points[k] - points[k+1:])**2, axis=1))
         weights = rho_k * rho[k+1:]
-        index = np.minimum(np.asarray(distances/r[0], 'i'), n_max)
+        index = np.minimum(np.asarray(distances/dr, 'i'), n_max)
         # Note: indices may be duplicated, so "Pr[index] += w" will not work!!
         extended_Pr += np.bincount(index, weights, n_max+1)
@@ -269 +315 @@
     # we are only accumulating the upper triangular distances.
     #Pr = Pr * 2 / len(rho)**2
-    return Pr / Pr.max()
+    return Pr
 
     # Can get an additional 2x by going to C.  Cuda/OpenCL will allow even
@@ -291 +337 @@
 }
 """
+
+if NUMBA:
+    @njit("f8[:](f8[:], f8[:], f8[:,:])")
+    def _calc_Pr_uniform_jit(r, rho, points):
+        dr = r[0]
+        n_max = len(r)
+        Pr = np.zeros_like(r)
+        for j in range(len(rho) - 1):
+            x, y, z = points[j, 0], points[j, 1], points[j, 2]
+            for k in range(j+1, len(rho)):
+                distance = sqrt((x - points[k, 0])**2
+                                + (y - points[k, 1])**2
+                                + (z - points[k, 2])**2)
+                index = int(distance/dr)
+                if index < n_max:
+                    Pr[index] += rho[j] * rho[k]
+        # Make Pr independent of sampling density.  The factor of 2 comes because
+        # we are only accumulating the upper triangular distances.
+        #Pr = Pr * 2 / len(rho)**2
+        return Pr
+
+
+def calc_Pr(r, rho, points):
+    # P(r) with uniform steps in r is 3x faster; check if we are uniform
+    # before continuing
+    if np.max(np.abs(np.diff(r) - r[0])) > r[0]*0.01:
+        Pr = _calc_Pr_nonuniform(r, rho, points)
+    else:
+        if NUMBA:
+            Pr = _calc_Pr_uniform_jit(r, rho, points)
+        else:
+            Pr = _calc_Pr_uniform(r, rho, points)
+    return Pr / Pr.max()
+
 
 def j0(x):
@@ -333 +413 @@
         plt.legend()
 
+def plot_calc_2d(qx, qy, Iqxy, theory=None):
+    import matplotlib.pyplot as plt
+    qx, qy = bin_edges(qx), bin_edges(qy)
+    #qx, qy = np.meshgrid(qx, qy)
+    if theory is not None:
+        plt.subplot(121)
+    plt.pcolormesh(qx, qy, np.log10(Iqxy))
+    plt.xlabel('qx (1/A)')
+    plt.ylabel('qy (1/A)')
+    if theory is not None:
+        plt.subplot(122)
+        plt.pcolormesh(qx, qy, np.log10(theory))
+        plt.xlabel('qx (1/A)')
+
 def plot_points(rho, points):
     import mpl_toolkits.mplot3d
@@ -343 +437 @@
         pass
     n = len(points)
-    index = np.random.choice(n, size=1000) if n > 1000 else slice(None, None)
+    #print("len points", n)
+    index = np.random.choice(n, size=500) if n > 500 else slice(None, None)
     ax.scatter(points[index, 0], points[index, 1], points[index, 2], c=rho[index])
     #low, high = points.min(axis=0), points.max(axis=0)
     #ax.axis([low[0], high[0], low[1], high[1], low[2], high[2]])
     ax.autoscale(True)
+
+def check_shape(shape, fn=None):
+    rho_solvent = 0
+    q = np.logspace(-3, 0, 200)
+    r = shape.r_bins(q, r_step=0.01)
+    sampling_density = 6*5000 / shape.volume()
+    rho, points = shape.sample(sampling_density)
+    t0 = time.time()
+    Pr = calc_Pr(r, rho-rho_solvent, points)
+    print("calc Pr time", time.time() - t0)
+    Iq = calc_Iq(q, r, Pr)
+    theory = (q, fn(q)) if fn is not None else None
+
+    import pylab
+    #plot_points(rho, points); pylab.figure()
+    plot_calc(r, Pr, q, Iq, theory=theory)
+    pylab.show()
+
+def check_shape_2d(shape, fn=None, view=(0, 0, 0)):
+    rho_solvent = 0
+    nq, qmax = 100, 1.0
+    qx = np.linspace(0.0, qmax, nq)
+    qy = np.linspace(0.0, qmax, nq)
+    Qx, Qy = np.meshgrid(qx, qy)
+    sampling_density = 50000 / shape.volume()
+    #t0 = time.time()
+    rho, points = shape.sample(sampling_density)
+    #print("sample time", time.time() - t0)
+    t0 = time.time()
+    Iqxy = calc_Iqxy(Qx, Qy, rho, points, view=view)
+    print("calc time", time.time() - t0)
+    theory = fn(Qx, Qy) if fn is not None else None
+    Iqxy += 0.001 * Iqxy.max()
+    if theory is not None:
+        theory += 0.001 * theory.max()
+
+    import pylab
+    #plot_points(rho, points); pylab.figure()
+    plot_calc_2d(qx, qy, Iqxy, theory=theory)
+    pylab.show()
+
+def sas_sinx_x(x):
+    with np.errstate(all='ignore'):
+        retvalue = sin(x)/x
+    retvalue[x == 0.] = 1.
+    return retvalue
 
 def sas_2J1x_x(x):
@@ -373 +514 @@
     Iq = Iq/Iq[0]
     return Iq
+
+def cylinder_Iqxy(qx, qy, radius, length, view=(0, 0, 0)):
+    qa, qb, qc = invert_view(qx, qy, view)
+    qab = np.sqrt(qa**2 + qb**2)
+    Fq = sas_2J1x_x(qab*radius) * j0(qc*length/2)
+    Iq = Fq**2
+    return Iq.reshape(qx.shape)
 
 def sphere_Iq(q, radius):
@@ -415 +563 @@
     return Iq/Iq[0]
 
-def check_shape(shape, fn=None):
-    rho_solvent = 0
-    q = np.logspace(-3, 0, 200)
-    r = shape.r_bins(q, r_step=0.01)
-    sampling_density = 15000 / shape.volume()
-    rho, points = shape.sample(sampling_density)
-    Pr = calc_Pr(r, rho-rho_solvent, points)
-    Iq = calc_Iq(q, r, Pr)
-    theory = (q, fn(q)) if fn is not None else None
-
-    import pylab
-    #plot_points(rho, points); pylab.figure()
-    plot_calc(r, Pr, q, Iq, theory=theory)
-    pylab.show()
+def csbox_Iqxy(qx, qy, a, b, c, da, db, dc, slda, sldb, sldc, sld_core, view=(0,0,0)):
+    qa, qb, qc = invert_view(qx, qy, view)
+
+    sld_solvent = 0
+    overlapping = False
+    dr0 = sld_core - sld_solvent
+    drA, drB, drC = slda-sld_solvent, sldb-sld_solvent, sldc-sld_solvent
+    tA, tB, tC = a + 2*da, b + 2*db, c + 2*dc
+    siA = a*sas_sinx_x(a*qa/2)
+    siB = b*sas_sinx_x(b*qb/2)
+    siC = c*sas_sinx_x(c*qc/2)
+    siAt = tA*sas_sinx_x(tA*qa/2)
+    siBt = tB*sas_sinx_x(tB*qb/2)
+    siCt = tC*sas_sinx_x(tC*qc/2)
+    Fq = (dr0*siA*siB*siC
+          + drA*(siAt-siA)*siB*siC
+          + drB*siA*(siBt-siB)*siC
+          + drC*siA*siB*(siCt-siC))
+    Iq = Fq**2
+    return Iq.reshape(qx.shape)
 
 def check_cylinder(radius=25, length=125, rho=2.):
@@ -434 +588 @@
     fn = lambda q: cylinder_Iq(q, radius, length)
     check_shape(shape, fn)
+
+def check_cylinder_2d(radius=25, length=125, rho=2., view=(0, 0, 0)):
+    shape = EllipticalCylinder(radius, radius, length, rho)
+    fn = lambda qx, qy, view=view: cylinder_Iqxy(qx, qy, radius, length, view=view)
+    check_shape_2d(shape, fn, view=view)
+
+def check_cylinder_2d_lattice(radius=25, length=125, rho=2.,
+                              view=(0, 0, 0)):
+    nx, dx = 1, 2*radius
+    ny, dy = 30, 2*radius
+    nz, dz = 30, length
+    dx, dy, dz = 2*dx, 2*dy, 2*dz
+    def center(*args):
+        sigma = 0.333
+        space = 2
+        return [(space*n+np.random.randn()*sigma)*x for n, x in args]
+    shapes = [EllipticalCylinder(radius, radius, length, rho,
+                                 #center=(ix*dx, iy*dy, iz*dz)
+                                 orientation=np.random.randn(3)*0,
+                                 center=center((ix, dx), (iy, dy), (iz, dz))
+                                )
+              for ix in range(nx)
+              for iy in range(ny)
+              for iz in range(nz)]
+    shape = Composite(shapes)
+    fn = lambda qx, qy, view=view: cylinder_Iqxy(qx, qy, radius, length, view=view)
+    check_shape_2d(shape, fn, view=view)
 
 def check_sphere(radius=125, rho=2):
@@ -449 +630 @@
     side_c2 = copy(side_c).shift(0, 0, -c-dc)
     shape = Composite((core, side_a, side_b, side_c, side_a2, side_b2, side_c2))
-    fn = lambda q: csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core)
-    check_shape(shape, fn)
+    def fn(q):
+        return csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core)
+    #check_shape(shape, fn)
+
+    view = (20, 30, 40)
+    def fn_xy(qx, qy):
+        return csbox_Iqxy(qx, qy, a, b, c, da, db, dc,
+                          slda, sldb, sldc, sld_core, view=view)
+    check_shape_2d(shape, fn_xy, view=view)
 
 if __name__ == "__main__":
     check_cylinder(radius=10, length=40)
+    #check_cylinder_2d(radius=10, length=40, view=(90,30,0))
+    #check_cylinder_2d_lattice(radius=10, length=50, view=(90,30,0))
     #check_sphere()
     #check_csbox()