updated pairwise forces to new interface; added lennard-jones forces

stoked/__init__.py

@@ -17,6 +17,7 @@
 from .hydrodynamics import grand_mobility_matrix
 from .constraints import constrain_position, constrain_rotation
 from .forces import pairwise_central_force, pairwise_force
+from .common import lennard_jones
 
 from .analysis import msd
 from .utility import quaternion_to_angles

stoked/collisions.py

@@ -1,46 +1,22 @@
 import numpy as np
 from stoked import interactions
+from stoked.forces import pairwise_central_force
 
-class collisions_sphere(interactions):
-    """
-    Collision between spheres using the force model F = kn*overlap^1.5
-    """
-    def __init__(self, radii, kn):
-        """
-        Arguments:
-            radii   particle radii
-            kn      force constant
-        """
-        # self.radii = stoked.array(radii)
-        self.radii = radii
-        self.kn = kn
-
-    def force(self):
-        Nparticles = len(self.position)
-        if np.isscalar(self.radii):
-            rad = np.full(Nparticles, self.radii, dtype=float)
-        else:
-            rad = np.asarray(self.radii, dtype=float)
+def collisions_sphere(radius, kn):
+    def F(r):
+        nonlocal radius
 
-        F = np.zeros_like(self.position)
+        if np.isscalar(radius):
+            Nparticles = len(r)
+            radius = np.full(Nparticles, radius, dtype=float)
 
-        for i in range(0,Nparticles):
-            for j in range(i+1,Nparticles):
-                d = self.position[i] - self.position[j]
-                r = np.linalg.norm(d)
-                if r < rad[i] + rad[j]:
-                    overlap = abs(r - rad[i] + rad[j])
+        T1 = np.add.outer(radius, radius)
 
-                    r_hat = d/r
-                    Fn = r_hat*self.kn*overlap**1.5
-
-                    F[i] += Fn
-                    F[j] -= Fn
-        return F
-
-    def torque(self):
-        return np.zeros_like(self.position)
+        overlap = np.abs(r - T1)
+        overlap[overlap>T1] = 0
+        return kn*overlap**1.5
 
+    return pairwise_central_force(F)
 
 class collisions_sphere_interface(interactions):
     """

stoked/common.py

@@ -0,0 +1,14 @@
+from stoked.forces import pairwise_central_force
+
+def lennard_jones(rmin, eps):
+    """Lennard-Jones pairwise interactions
+
+    Arguments:
+        rmin       distance where potential is a minimum 
+        eps        depth of the potential well
+    """
+
+    def F(r):
+        return 12*eps*(rmin**12/r**13 - rmin**6/r**7)
+
+    return pairwise_central_force(F)

stoked/electrostatics.py

@@ -1,81 +1,55 @@
 import numpy as np
 from scipy import constants
 from stoked import interactions
+from stoked.forces import pairwise_central_force
 
-class electrostatics(interactions):
+def electrostatics(charges):
     """
     Free-space point electrostatic interactions
     """
-    def __init__(self, charges):
-        # self.charges = stoked.array(charges)
-        self.charges = charges
+    ke = 1/(4*np.pi*constants.epsilon_0)
 
-    def force(self):
-        Nparticles = len(self.position)
-        if np.isscalar(self.charges):
-            Q = np.full(Nparticles, self.charges, dtype=float)
+    def F(r):
+        if np.isscalar(charges):
+            Nparticles = len(r)
+            Q = np.full(Nparticles, charges, dtype=float)
         else:
-            Q = np.asarray(self.charges, dtype=float)
+            Q = np.asarray(charges, dtype=float)
 
-        ke = 1/(4*np.pi*constants.epsilon_0)
+        return ke*np.outer(Q, Q)/r**2
 
-        r_ijx = self.position[:,np.newaxis,:] - self.position[np.newaxis,...]
-        with np.errstate(divide='ignore'):
-            F_ijx = ke*np.einsum('i,j,ij,ijx->ijx', Q, Q, 1/np.sum(np.abs(r_ijx + 1e-20)**3, axis=-1), r_ijx)
-        np.einsum('iix->x', F_ijx)[...] = 0
-
-        F = np.sum(F_ijx, axis=1)
+    return pairwise_central_force(F)
 
-        return F
-
-    def torque(self):
-        return np.zeros_like(self.position)
-
-
-class double_layer_sphere(interactions):
+def double_layer_sphere(radius, potential, temperature=300, debye=27.6e-9, eps_m=80.4, zp=1):
     """
     Electrostatic interactions in a fluid medium for spheres
     """
-    def __init__(self, radius, potential, debye=27.6e-9, eps_m=80.4, zp=1):
-        self.radius = np.asarray(radius, dtype=float)
-        self.potential = np.asarray(potential, dtype=float)
-        self.debye = debye
-        self.eps_m = eps_m
-        self.zp = zp
+    radius = np.asarray(radius, dtype=float)
+    potential = np.asarray(potential, dtype=float)
 
-    def force(self):
-        Nparticles = len(self.position)
-        if not np.ndim(self.radius):
-            radius = np.full(Nparticles, self.radius, dtype=float)
-        else:
-            radius = self.radius
-        if not np.ndim(self.potential):
-            potential = np.full(Nparticles, self.potential, dtype=float)
-        else:
-            potential = self.potential
-
-        factor = 16*np.pi*constants.epsilon_0*self.eps_m/self.debye \
-                 *(constants.k*self.temperature/(self.zp*constants.elementary_charge))**2
+    def F(r):
+        nonlocal radius, potential
+        Nparticles = len(r)
+        if not np.ndim(radius):
+            radius = np.full(Nparticles, radius, dtype=float)
 
-        r_ijx = self.position[:,np.newaxis,:] - self.position[np.newaxis,...]
-        r_ij = np.linalg.norm(r_ijx, axis=-1)
+        if not np.ndim(potential):
+            potential = np.full(Nparticles, potential, dtype=float)
 
+        factor = 16*np.pi*constants.epsilon_0*eps_m/debye \
+                 *(constants.k*temperature/(zp*constants.elementary_charge))**2
+
         T1 = np.add.outer(radius, radius)
-        if self.temperature == 0:
+        if temperature == 0:
             T2 = np.ones_like(potential)
         else:
-            T2 = np.tanh(self.zp*constants.elementary_charge*potential/(4*constants.k*self.temperature))
-        Q = factor*T1*np.multiply.outer(T2, T2)*np.exp(-(r_ij - T1)/self.debye)
-
-        F_ijx = np.einsum('ij,ij,ijx->ijx', Q, 1/(r_ij + 1e-20), r_ijx)
+            T2 = np.tanh(zp*constants.elementary_charge*potential/(4*constants.k*temperature))
+        T2 = np.outer(T2, T2)
 
-        np.einsum('iix->x', F_ijx)[...] = 0
-        F = np.sum(F_ijx, axis=1)
+        Q = factor*T1*T2*np.exp(-(r - T1)/debye)
+        return Q
 
-        return F
-
-    def torque(self):
-        return np.zeros_like(self.position)
+    return pairwise_central_force(F)
 
 class double_layer_sphere_interface(interactions):
     """

stoked/forces.py

@@ -16,7 +16,8 @@ def force(self):
         Nparticles = len(self.position)
 
         r_ij = self.position[:,np.newaxis] - self.position[np.newaxis]   # N x N x 3
-        F_ij = self.force_func(r_ij)   # N x N x 3
+        with np.errstate(divide='ignore'):
+            F_ij = self.force_func(r_ij)   # N x N x 3
         np.einsum('iix->ix', F_ij)[...] = 0
         F_i = np.sum(F_ij, axis=1)
 
@@ -26,16 +27,44 @@ def torque(self):
         T_i = np.zeros_like(self.position)
         return T_i
 
-def pairwise_central_force(force_func):
-    def F(rvec):
-        r = np.linalg.norm(rvec, axis=-1)
-        f = force_func(r)
-        r_inv = 1/r
-        np.einsum('ii->i', r_inv)[...] = 0
+    def __add__(self, other):
+        def new_func(r):
+            return self.force_func(r) + other.force_func(r)
 
-        return np.einsum('ij,ijx,ij->ijx', f, rvec, r_inv)
+        return pairwise_force(new_func)
 
-    return pairwise_force(F)
+class pairwise_central_force(interactions):
+    def __init__(self, force_func):
+        """
+        Arguments:
+            force_func     force function of the form F(r[dim]) -> [dim]
+        """
+        self.force_func = force_func
+
+    def force(self):
+        Nparticles = len(self.position)
+
+        r_ij = self.position[:,np.newaxis] - self.position[np.newaxis]   # N x N x 3
+        r = np.linalg.norm(r_ij, axis=-1)
+        with np.errstate(divide='ignore'):
+            F = self.force_func(r)
+            r_inv = 1/r
+
+        F_ij = np.einsum('ij,ijx,ij->ijx', F, r_ij, r_inv)
+        np.einsum('iix->ix', F_ij)[...] = 0
+        F_i = np.sum(F_ij, axis=1)
+
+        return F_i
+
+    def torque(self):
+        T_i = np.zeros_like(self.position)
+        return T_i
+
+    def __add__(self, other):
+        def new_func(r):
+            return self.force_func(r) + other.force_func(r)
+
+        return pairwise_central_force(new_func)
 
 def pairwise_potential(potential_func):
     pass

stoked/solver.py

@@ -111,6 +111,7 @@ def __init__(self, *, temperature, dt, position, drag, orientation=None,
             self.interactions = [interactions]
         else:
             self.interactions = interactions
+        self._combine_pairwise_interactions()
 
         if constraint is None:
             self.constraint = []
@@ -362,6 +363,20 @@ def _total_torque(self, time, position, orientation):
                 T += I.torque()
         return T
 
+    def _combine_pairwise_interactions(self):
+        from stoked.forces import pairwise_central_force
+
+        if self.interactions is not None:
+            pairwise_forces = []
+            for i,I in enumerate(self.interactions):
+                if isinstance(I, pairwise_central_force):
+                    pairwise_forces.append(I)
+
+            if len(pairwise_forces) > 1:
+                for pf in pairwise_forces:
+                    self.interactions.remove(pf)
+                self.interactions.append(sum(pairwise_forces[1:], start=pairwise_forces[0]))
+
 def brownian_dynamics(*, temperature, dt, position, drag, orientation=None, 
                  force=None, torque=None, interactions=None, constraint=None):
     """

stoked/van_der_waals.py

@@ -1,36 +1,26 @@
 import numpy as np
 from stoked import interactions
+from stoked.forces import pairwise_central_force
 
-class van_der_waals_sphere(interactions):
-    def __init__(self, radius, hamaker):
-        self.radius = np.asarray(radius, dtype=float)
-        self.hamaker = hamaker
+def van_der_waals_sphere(radius, hamaker):
+    radius = np.asarray(radius, dtype=float)
 
-    def force(self):
-        Nparticles = len(self.position)
-        if not np.ndim(self.radius):
-            radius = np.full(Nparticles, self.radius, dtype=float)
-        else:
-            radius = self.radius
+    def F(r):
+        nonlocal radius
 
-        factor = 2*self.hamaker/3
-        r_ijx = self.position[:,np.newaxis,:] - self.position[np.newaxis,...]
-        r_ij = np.linalg.norm(r_ijx, axis=-1)
+        Nparticles = len(r)
+        if not np.ndim(radius):
+            radius = np.full(Nparticles, radius, dtype=float)
 
-        T1 = np.multiply.outer(radius, radius)
+        factor = 2*self.hamaker/3
+        T1 = np.outer(radius, radius)
         T2 = np.add.outer(radius, radius)**2
         T3 = np.substract.outer(radius, radius)**2
-        Q = r_ij*T1*(1/(r_ij**2 - T2**2) - 1 /(r_ij**2 - T3**2))**2
+        Q = r*T1*(1/(r**2 - T2**2) - 1 /(r**2 - T3**2))**2
 
-        F_ijx = np.einsum('ij,ij,ijx->ijx', -factor*Q, 1/(r_ij + 1e-20), r_ijx)
-        np.einsum('iix->x', F_ijx)[...] = 0
-        F = np.sum(F_ijx, axis=1)
-
-        return F
-
-    def torque(self):
-        return np.zeros_like(self.position)
+        return Q
 
+    return F
 
 class van_der_waals_sphere_interface(interactions):
     def __init__(self, radius, hamaker, zpos=0):