Changeset 8fb2a94 in sasmodels for explore/realspace.py

Timestamp:

Jan 11, 2018 1:51:11 PM (6 years ago)

Author:

Paul Kienzle <pkienzle@…>

Branches:

master, core_shell_microgels, magnetic_model, ticket-1257-vesicle-product, ticket_1156, ticket_1265_superball, ticket_822_more_unit_tests

Children:

c11d09f

Parents:

97be877

Message:

play with numba JIT compiler

File:

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

explore/realspace.py

@@ -244 +244 @@
         return values, self._adjust(points)
 
-def calc_Iqxy(qx, qy, rho, points, volume=1, view=(0, 0, 0)):
-    x, y, z = points.T
+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)
-    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.array(Iq).reshape(qx.shape) / np.sum(volume)
+    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):
@@ -275 +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)
@@ -305 +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
@@ -327 +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):
@@ -404 +448 @@
     q = np.logspace(-3, 0, 200)
     r = shape.r_bins(q, r_step=0.01)
-    sampling_density = 50000 / shape.volume()
+    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
@@ -586 +632 @@
     def fn(q):
         return csbox_Iq(q, a, b, c, da, db, dc, slda, sldb, sldc, sld_core)
-    check_shape(shape, fn)
+    #check_shape(shape, fn)
 
     view = (20, 30, 40)
@@ -592 +638 @@
         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)
+    check_shape_2d(shape, fn_xy, view=view)
 
 if __name__ == "__main__":