import numpy as np
import data

class _wca(object):
    def __init__(self, system):
        self.system = system
        self.coeff = data.pair_coeff(self.system)

    def get_rmin(self, i, j):
        raise Exception('programming error: potential minimum evaluator must be defined!')

    def energy(self, i, j):
        raise Exception('programming error: potential evaluator must be defined!')

    def U(self, i, j, r):
        if not self.coeff.verify():
            raise Exception('all required coefficients are not set!')

        upot = np.zeros_like(r)
        rs = np.fabs(np.array(r))

        rmin = self.get_rmin(i, j)
        flag_eval = rs >= rmin
        upot[~flag_eval] = self.energy(i, j, rmin)
        upot[flag_eval] = self.energy(i, j, rs[flag_eval])

        return upot

class lj(_wca):
    def __init__(self, system):
        _wca.__init__(self, system)
        self.coeff.require = ['sigma','epsilon','rcut','shift']

        # the default sigma is the hard sphere contact (i.e., arithmetic mixing)
        for i in self.system.types:
            for j in self.system.types:
                if self.system.sigma[i] is not None and self.system.sigma[j] is not None:
                    sigma_ij = 0.5*(self.system.sigma[i] + self.system.sigma[j])
                    self.coeff.set(i, j, sigma=sigma_ij)

    def get_rmin(self, i, j):
        return np.power(2.0,1.0/6.0) * self.coeff.get(i, j, 'sigma')

    def energy(self, i, j, r):
        U = np.zeros_like(r)
        rs = np.array(r)

        # pass on this type if it has no interaction
        rcut = self.coeff.get(i, j, 'rcut')
        if not rcut or not rcut > 0.0 or rcut is None:
            return U

        # potential is zero if either sigma or epsilon is close to zero
        sigma = self.coeff.get(i, j, 'sigma')
        eps = self.coeff.get(i, j, 'epsilon')
        if np.isclose(sigma, 0.0) or np.isclose(eps,0.0):
            return U

        shift = self.coeff.get(i, j,'shift')
        U_cut = 0.0
        if shift:
            U_cut = 4.*eps*((sigma/rcut)**12 - (sigma/rcut)**6)

        # only perform calculations for points inside the rcut
        flags = (np.fabs(rs) <= rcut)
        ri = sigma/rs[flags]
        r2i = ri*ri
        r4i = r2i*r2i
        r6i = r2i*r4i
        r12i = r6i*r6i
        U[flags] = 4.*eps*(r12i - r6i) - U_cut

        return U

@@ -2342,24 +2342,1264 @@
 U_cut%0A%0A        return U%0A
+%0Aclass depletion(object):%0A    def __init__(self, system):%0A        %22%22%22Depletion pair force between large hard spheres in a binary mixture%0A%0A        Roth, Evans, and Dietrich, Phys. Rev. E 2000, 62, 5360 - 5377%0A        %22%22%22%0A        self.system = system%0A        if len(self.system.types) %3E 2:%0A            raise Exception('Depletion pair force should only be used in binary systems')%0A%0A        self.coeff = data.pair_coeff(self.system, require=%5B'sigma', 'rho', 'rcut'%5D)%0A%0A    def U(self, i, j, r):%0A        if not self.coeff.verify():%0A            raise Exception('not all parameters are set!')%0A%0A        upot = np.zeros_like(r)%0A        # pass on this type if it has no interaction%0A        rcut = self.coeff.get(i, j, 'rcut')%0A        if not rcut or not rcut %3E 0.0 or rcut is None:%0A            return upot%0A%0A        sigma = self.coeff.get(i, j, 'sigma')%0A        rho = self.coeff.get(i, j, 'rho')%0A        if np.isclose(rho,0.0):%0A            return upot%0A%0A        rs = np.fabs(np.array(r))%0A%0A        Rs = 0.5 * rcut%0A        Rb = 0.5 * sigma%0A        flag_eval = np.logical_and(rs %3E= Rb, rs %3C= rcut)%0A        h = r%5Bflag_eval%5D - Rb%0A        upot%5Bflag_eval%5D = -rho * np.pi * (2.*Rs - h)*(Rs * (Rb + (2./3.)*Rs)-(h/2.)*(Rb + Rs/3.) - h**2/12.)%0A        return upot%0A

Add depletion pair force