Merge pull request #3 from dwieker/nkeim/add-gofr-rebased

Finished 3D g(r)

Author: dwieker

trackpy/static.py

@@ -4,7 +4,6 @@
 from warnings import warn
 
 
-
 def pairCorrelationKDTree2D(feat, cutoff, fraction = 1., dr = .5, p_indices = None, ndensity=None, boundary = None,
                             handle_edge=True):
     """   
@@ -172,12 +171,8 @@ def pairCorrelationKDTree3D(feat, cutoff, fraction = 1., dr = .5, p_indices = No
     for idx in p_indices:
         dist, idxs = ckdtree.query(points[idx], k=max_p_count, distance_upper_bound=cutoff)
         dist = dist[dist > 0] # We don't want to count the same particle
-        #print dist.shape
-        #print dist[dist.shape[0] - 10:]
-
+     
         area = (4./3.) * np.pi * (np.arange(dr, cutoff + 2*dr, dr)**3 - np.arange(0, cutoff + dr, dr)**3)
-
-        #print area
 
         if handle_edge:
             # Find the number of edge collisions at each radii
@@ -188,58 +183,40 @@ def pairCorrelationKDTree3D(feat, cutoff, fraction = 1., dr = .5, p_indices = No
 
                 # Use analyitcal solution to find area of disks cut off by one wall.
                 # Grab the distance to the closest wall
-                #d = _distances_to_wall3D(points[idx], xmin, xmax, ymin, ymax, zmin, zmax).min()
-
-                #inx = np.where(collisions == 1)[0]
-
-                #theta = np.arccos(d / (r_edges[inx] + dr/2))
-                #area[inx] *= 1 - 2*np.pi*(1 - np.cos(theta)) / (4*np.pi)
+                d = _distances_to_wall3D(points[idx], xmin, xmax, ymin, ymax, zmin, zmax).min()
+                inx = np.where(collisions == 1)[0]
+                theta = np.arccos(d / (r_edges[inx] + dr/2))
+                area[inx] *= 1 - 2*np.pi*(1 - np.cos(theta)) / (4*np.pi)
 
                 # If shell is cutoff by 2 or more walls, generate a bunch of points and use a mask to
                 # estimate the area within the boundaries
-                inx = np.where(collisions >= 1)[0]
+                inx = np.where(collisions >= 2)[0]
                 x = refx[inx] + points[idx,0]
                 y = refy[inx] + points[idx,1]
                 z = refz[inx] + points[idx,2]
                 mask = (x >= xmin) & (x <= xmax) & (y >= ymin) & (y <= ymax) & (z >= zmin) & (z <= zmax)
                 area[inx] *= mask.sum(axis=1, dtype='float') / len(refx[0])
 
-                print points[idx]
-                print mask.sum(axis=1, dtype='float') / len(refx[0])
-
         g_r +=  np.histogram(dist, bins = r_edges)[0] / area[:-1]
 
     g_r /= (ndensity * len(p_indices))
     return r_edges, g_r
 
 def _num_wall_collisions2D(point, radius, xmin, xmax, ymin, ymax):
+    """Returns the number of walls a shell of a certain radius and position collides with.
+       Wall boundaries specified by min, max parameters"""
     collisions = (point[0] + radius >= xmax).astype(int) + (point[0] - radius <= xmin).astype(int) + \
                  (point[1] + radius >= ymax).astype(int) + (point[1] - radius <= ymin).astype(int)
 
     return collisions
 
 def _distances_to_wall2D(point, xmin, xmax, ymin, ymax): 
+    """Returns the distance of a paritlce a position 'point' to the nearest wall"""
     return np.array([point[0]-xmin, xmax-point[0], point[1]-ymin, ymax-point[1]])
 
 def _points_ring2D(r_edges, dr, n):
     """Returns x, y array of points comprising shells extending from r to r_dr.
-
-    layers determines how many concentric layers are in each shell,
-        and n determines the number of points in each layer"""
-
-    """
-    refx=np.empty((len(r_edges), n*layers))
-    refy=refx.copy()
-    for index, r in enumerate(r_edges): 
-        theta = np.linspace(0, 2*np.pi, n)
-        theta = theta.repeat(layers).reshape((len(theta), layers))
-        x = np.cos(theta) * np.linspace(r, r+dr, layers)
-        y = np.sin(theta) * np.linspace(r, r+dr, layers)
-        refx[index] = x.reshape(n*layers)
-        refy[index] = y.reshape(n*layers)
-
-    return refx, refy
-    """
+       n determines the number of points in each ring"""
 
     refx_all, refy_all = [],[]
     for r in r_edges:
@@ -255,17 +232,21 @@ def _points_ring2D(r_edges, dr, n):
 
 
 def _num_wall_collisions3D(point, radius, xmin, xmax, ymin, ymax, zmin, zmax):
+    """Returns the number of walls a shell of a certain radius and position collides with.
+       Wall boundaries specified by min, max parameters"""
     collisions = (point[0] + radius >= xmax).astype(int) + (point[0] - radius <= xmin).astype(int) + \
                  (point[1] + radius >= ymax).astype(int) + (point[1] - radius <= ymin).astype(int) + \
                  (point[2] + radius >= zmax).astype(int) + (point[2] - radius <= zmin).astype(int) 
 
     return collisions
 
 def _distances_to_wall3D(point, xmin, xmax, ymin, ymax, zmin, zmax): 
+    """Returns the distance of a paritlce a position 'point' to the nearest wall"""
     return np.array([point[0]-xmin, xmax-point[0], point[1]-ymin, ymax-point[1], point[2]-zmin, zmax-point[2]])
 
 def _points_ring3D(r_edges, dr, n):
-    """Returns x, y, z arrays of points comprising shells extending from r to r_dr. n determines the density of the shells"""
+    """Returns x, y, z arrays of points comprising shells extending from r to r_dr. 
+       n determines the number of particles in each of the shells"""
 
     refx_all, refy_all, refz_all = [],[],[]
     for r in r_edges:

trackpy/tests/test_static.py

@@ -10,9 +10,9 @@
 class TestPairCorrelation(unittest.TestCase):
 
 
-    def test_correlation2D(self):
-
-        ##### 2D TEST ######
+    def test_correlation2D_lattice(self):
+        ### Lattice Test
+        # With proper edge handling, g(r) of the particle at the center should be the same as g(r) for all particles.
         lattice = self._lattice2D()
 
         # Calculate g_r on the center particle only (index 210)
@@ -34,6 +34,10 @@ def test_correlation2D(self):
         self.assertFalse(np.allclose(g_r_all, g_r_no_edge, atol=.02))
 
 
+    def test_correlation2D_ring(self):
+        # Ring test
+        # Generate a series of concentric shells, each with the same number of particles.
+        # The peaks in g(r) should decay as 1/r.
         ring = self._rings2D()
 
         edges, g_r = pairCorrelationKDTree2D(ring, dr=.1, cutoff=10, p_indices=[0], boundary = (-10., 10., -10., 10.))
@@ -53,65 +57,40 @@ def test_correlation3D_ring(self):
         # Generate a series of concentric shells, each with the same number of particles.
         # The peaks in g(r) should decay as 1/r^2.
         ring = self._rings3D()
-
         edges, g_r = pairCorrelationKDTree3D(ring, dr=.1, cutoff=10, p_indices=[len(ring) - 1], boundary = (-10., 10., -10., 10., -10., 10.), handle_edge=True)
         g_r /= np.linalg.norm(g_r)
         peaks = g_r[g_r > 0]
-
         assert len(peaks) == 9
-
         x = np.arange(1,10,1)
         r = peaks.max() * 1/x**2
-
         self.assertTrue( np.allclose(peaks, r, atol=.02) )
 
 
     def test_correlation3D_lattice(self):
         ### Lattice Test
         # With proper edge handling, g(r) of the particle at the center should be the same as g(r) for all particles.
-        lattice = self._lattice3D(n = 10)
-        
-        print lattice.iloc[444]
+        lattice = self._lattice3D(n = 20)
+
         # Calculate g_r on the center particle only (index 210)
-        edges, g_r_one = pairCorrelationKDTree3D(lattice, dr=.1, cutoff=4, p_indices=[444])
+        edges, g_r_one = pairCorrelationKDTree3D(lattice, dr=.1, cutoff=7, p_indices=[4649])
         g_r_one /= np.linalg.norm(g_r_one) #We care about the relative difference of g_r in this case, so let's normalize both.
 
         # Calculate g_r on all particles
-        edges, g_r_all = pairCorrelationKDTree3D(lattice, dr=.1, cutoff=4)
+        edges, g_r_all = pairCorrelationKDTree3D(lattice, dr=.1, cutoff=7)
         g_r_all /= np.linalg.norm(g_r_all)
 
         # Calculate g_r on all particles
-        edges, g_r_no_edge = pairCorrelationKDTree3D(lattice, dr=.1, cutoff=4, handle_edge=False)
+        edges, g_r_no_edge = pairCorrelationKDTree3D(lattice, dr=.1, cutoff=7, handle_edge=False)
         g_r_no_edge /= np.linalg.norm(g_r_no_edge)
 
-        
-        
-        plt.plot(edges[:-1], g_r_one, label='one')
-        plt.plot(edges[:-1], g_r_all, label='all')
-        plt.plot(edges[:-1], g_r_no_edge, label='all, no edge')
-        plt.legend(loc='best')
-        plt.show()
-
         # Assert the functions are essentially the same
-        self.assertTrue(np.allclose(g_r_all, g_r_one, atol=.02))
+        self.assertTrue(np.allclose(g_r_all, g_r_one, atol=.04))
 
         # Turning off edge handling should give incorrect result
-        self.assertFalse(np.allclose(g_r_all, g_r_no_edge, atol=.02))
-
-    def test_sphere_mask(self):
-        x,y,z = _points_ring3D([1], 0, 10000)
-        #plt.scatter(x,y)
-        #plt.show()
-        
-        mask= (x >= 0) & (x <= 1) & (y >= 0) & (y <= 1)# & (z >= 0) & (z <= 1) 
-        print mask.sum() / len(x)
-        #plt.scatter(x[mask], y[mask])
-        #plt.show()
-
-
-
+        self.assertFalse(np.allclose(g_r_all, g_r_no_edge, atol=.04))
 
     def _lattice2D(self, n = 20):
+        #Generates 2D lattice, spacing = 1
         x,y = [],[]
         epsilon = 0.0
         for i in range(n):
@@ -123,6 +102,7 @@ def _lattice2D(self, n = 20):
 
 
     def _rings2D(self):
+        #Generates concentric rings, with a particle at the center
         theta = np.linspace(0, 2*np.pi, 10)
         points = np.zeros((100,2))
 
@@ -138,6 +118,7 @@ def _rings2D(self):
 
 
     def _lattice3D(self, n = 20):
+        #Generates 3D lattice, spacing = 1
         x,y,z = [],[],[]
         for i in range(n):
             for j in range(n):
@@ -150,6 +131,7 @@ def _lattice3D(self, n = 20):
 
 
     def _rings3D(self):
+        #Generates concentric spherical shells, with a particle at the center
         epsilon = .02
         r = np.arange(1, 10, 1) + epsilon
         refx, refy, refz = _points_ring3D(r, 0, 500)