#!/usr/bin/env python
from pylab import *
from scipy.optimize import fsolve, minimize_scalar
import time
import warnings
warnings.filterwarnings("ignore")
ion()

"""

Please first rename this file from OZ3.txt to OZ3.py. 

This program contains the numerical models presented in the associated article. It computes the equation of state of 
a colloidal suspension of uniformly charged hard spheres in a 1:1 electrolyte in thermodynamic equilibrium with a 
reservoir of salt ions. Available methods are the cell model, the OZ numerical resolution with different closures, the 
MSA and RMSA (semi)analytical solutions, and a perturbation method. 

The file is divided in sections. Here is a brief summary. 
- SECTION 1: definition of constants. Here for water at 293K. 
- SECTION 2: This is a numerical solver for the Ornstein-Zernike equation. The potential is of hard-sphere Yukawa type. 
             Available closure relations are Hypernetted Chain (HNC), Percus-Yevick (PY), or Rogers-Young (RY). The use 
             of this solver is detailed below. The (R)MSA and perturbation schemes are also included in this object. 
- SECTION 3: This is a numerical solver for the Poisson-Boltzmann equation in the cell model. It also performs EPC 
             renormalization. 
- SECTION 4: This is a small program showing the use of these solvers. It reproduces some results of Fig. 1 from the article. 

Here is how to use this program. 
The parameters in section 1 may be altered. 
Sections 2, 3 are the numerical methods and should not be altered unless you know precisely whaat you are doing. 
Section 4 will usually be the only thing to modify. See the comments in this section for a step by step explanation. 
You need to have python installed, with its modules matplotlib, scipy and numpy. 
Then I would recommend using the python shell called ipython (not compulsory). Inside this shell, type: run OZ3.py

Please let me know if you have any problem to run this code, if you find any bug, or if you improve it !
Contact: yannick (dot) hallez (at) univ-tlse3.fr

"""

##########################################
# SECTION 1
# Physical constants:
# -------------------
# Boltzmann const x Temperature
kT  = 1.3806503e-23*293
# electron charge
elc = 1.60217646e-19
# Avogadro number
Na  = 6.0221415e23
# eps0
eps0 = 8.854187817620e-12
# Eps_r x Eps_0 for water
eps = 78*eps0
# Bjerum length
lambdaB = elc**2/(4.0*pi*kT*eps)
lB = lambdaB 
##########################################



####################################################################################
# SECTION 2
# Define the OZ solver as an object
class OZSimulation():
  def __init__(self):
    #parameters
    self.potential = "yukawa"
    self.closure   = "HNC"
    self.phi       =  0.15      #   0.08       #     0.30           
    self.kappa     = 1./2.8e-8  #   1./7.5e-9  #     1./100.e-9      
    self.a         = 14.e-9     #   7.5e-9     #     500.e-9         
    self.a_pot     = self.a
    self.Z         = 9.573763490998747       #   25.44      #     1500.           
    self.delta     = 0.
    self.Hamaker   = 0.
    self.dc        = 0.16e-9
    self.lB        = lB
    self.PEL       = 0.
    self.quiet     = True 
    self.RYQuiet   = True 
    self.printpel  = True  
    self.FirstCall = True 
    # NUM
    self.N            = 1024 
    self.Ndmax        = 20.  
    self.alphad       = 0.6
    self.alphamin     = -1
    self.alphamax     = -1
    self.relax        = 0.1
    self.tol          = 1.0e-12
    self.log          = "stdout"
    self.Ng           = False
    self.res_start_Ng = 1.0e12
    self.Ng_renorm    = False 
    self.solver       = "Bisection"
    # variables 
    self.Ng_step = -1
    self.relaxfactor = 1.00

  def UpdateParameters(self):
    if self.FirstCall:
      # Store some stuff
      self.a_m          = self.a     # real particle radius in meters 
      self.kappa_stored = self.kappa # value of kappa at first call 
      self.delta_stored = self.delta # value of delta at first call 
    self.d   = 2*self.a    
    self.HSd = self.d #+ 2*self.delta
    self.Vp  = np.pi*self.HSd**3/6.0
    self.rho = self.phi / self.Vp 
    self.dr = self.Ndmax*self.d/(self.N-1) #d/10.
    self.Rmax = (self.N-1)*self.dr
    self.it = 0
    # define distances of interest
    self.r = np.linspace(0, self.Rmax, self.N)
    self.alpha = self.alphad/self.HSd  # if close 0, PY, or if infinity, HNC 
    self.iHSd    = 0
    while self.r[self.iHSd]<self.HSd: self.iHSd+=1
    # POTENTIAL
    self.a_pot = self.a - self.delta
    ka = self.kappa*self.a_pot
    kr = self.kappa*self.r

    # ---------- HS Yukawa -----------
    self.beta_u    = zeros(len(self.r))
    self.beta_u = np.exp(-kr)/(self.r+1.0e-20) \
        * self.Z**2*self.lB*(np.exp(ka)/(1.0+ka))**2
    self.beta_u[self.r<self.HSd]=1.0e8 *ones(len(self.r[self.r<self.HSd]))

    # This is not the first call anymore 
    self.FirstCall = False   
    # 
    self.g   = zeros(self.N)
    self.k   = zeros(self.N)
    self.S   = zeros(self.N)
    self.Gam = zeros(self.N)
 

  def update_yukawa(self,rho,g,ka,kr):
    r = kr/self.kappa
    extra = 0.
    beta_u = np.exp(-kr)/(r+1.0e-40) \
                  * self.Z**2*self.lB*(np.exp(ka)/(1.0+ka+extra))**2
    beta_u[r<self.HSd]=1.0e12*ones(len(r[r<self.HSd]))
    return beta_u

  # Define sine and Fourier-Bessel transforms and their inverse
  def dst(self,r,y):
    """
    The direct sine transform is -Imag(FFT)/2 for an odd function
    """
    N = len(y)
    dr = r[1]-r[0]
    # build odd function from y between 0 and infty, and the abscissa associated
    y2 = np.empty(2*N,float)
    y2[0] = y2[N] = 0.0
    y2[1:N] = y[1:]
    y2[:N:-1] = -y[1:]
    r2 = np.empty(2*N,float)
    r2[0]=0.
    r2[1:N] = r[1:]
    r2[:N:-1] = -r[1:]
    # get frequency
    f = rfftfreq(len(y2),dr)[:N]
    # and wavenumber 2pi/lambda=2pif
    k = 2*np.pi*f
    # get ft
    a = -np.imag(rfft(y2))[:N] *dr/2.0
    a[0] = 0.0
    return k, a
  
  def idst(self,w,a):
      N = len(a)
      dk = w[1]-w[0]
      df = dk/(2*pi)
      c = empty(N+1,complex)
      c[0] = c[N] = 0.0
      c[1:N] = -1j*a[1:]
      y = irfft(c)[:N] *df * 4.* N #*df for the mathematical integral, and 4N to scale the ft(ift(f))=f
      y[0] = 0.0
  
      return y
  
  def fbt(self,r,y):
    # compute Fourier-Bessel transform of y(r)
    k, y_hat   = self.dst(r,y*r)
    y_hat[1:] *= 1./k[1:]
    y_hat[0 ]  = 0. #reg function has to be odd anyway
    return k, y_hat
  
  def ifbt(self,k,y_hat):
    # compute inverse Fourier-Bessel transform of y_hat(k)
    y     = self.idst(k,y_hat*k)
    dr    = 1.*np.pi/((k[1]-k[0])*len(k))
    r     = linspace(0,(len(k)-1)*dr,len(k))
    y[1:]*= 1./r[1:]
    y[0]  = inf
    return r,y


  def compute_c_from_closure(self, Gamma): 
    # compute direct correlation functin c(r) in real space
    # from knowledge of Gamma function
    if self.closure == "HNC":
      if self.Ng_renorm:
        self.c = exp(-self.beta_u+self.beta_ucorr+Gamma)-Gamma-1.0
      else:
        self.c = exp(-self.beta_u+Gamma)-Gamma-1.0 
    elif self.closure == "PY":
      self.c = (Gamma+1.0)*(exp(-self.beta_u)-1.0) 
    elif self.closure == "RY":
      # PY  when fry=0, i.e. when alpha.r<<.01  whatever r (say 10a is enough), so alpha<.001/a is full PY
      # HNC when fry=1, i.e. when alpha.r>>4    whatever r (r=2a is minimum),   so alpha>2/a    is full HNC
      fry = 1.0 - exp(-self.alpha*self.r)
      self.c = exp(-self.beta_u)*(1.0+(exp(Gamma*fry)-1.0)/fry)-1.0-Gamma
    else:
      print "Unknown closure. Choose among HNC, PY, RY"
      sys.exit()
   

  def SolveOZ(self,gam_init=-1):
        # This function is the one called from outside to trigger the resolution 
    if self.closure=="RY":
        self.g, self.k, self.S, gam = self.FindRYalpha(gam_init=gam_init)
    else:
        self.g, self.k, self.S, self.Gam = self.SolveOZInternal(self.rho,gam_init=gam_init,quiet=self.quiet)
    # Fix S[0]
    coeff = polyfit(self.k[1:6],self.S[1:6],4)
    self.S[0] = coeff[-1]  
    return self.g, self.k, self.S, self.Gam


  def SolveOZInternal(self,rho,gam_init=-1,quiet=False):
    # need initializing gamma function with self.Gam, which will be 
    # either 0 et first call or the previous result at another call:
    if type(gam_init)==ndarray:
      self.fn = gam_init
    else:
      self.fn = zeros(self.N)

    # usefull f function
    #f = exp(-self.beta_u)-1.0
    self.res = 1.0
    if (not quiet) and (self.log=="X"):
      prevtime = time.time()
      counttab = []
      restab = []
      logplot, = plot(counttab,restab,'-r', lw=2)
      yscale('log')
      autoscale(True)
    #for i in range(20): 
    self.it = -1
    while self.res>self.tol:
      self.it += 1

      # Compute gn
      self.gn = self.gn_from_fn(rho,self.fn)

      # We are ready to define the input for the next iteration. 
      fnew = self.Ng5()

      # Compare gn and fn to know if convergence was reached 
      # actually it is before Ng update...
      self.res = sqrt( sum((fnew-self.fn)**2) )
      if not quiet: 
	if self.log=="stdout":
	  print "residual", self.it,  self.res
	elif self.log=="X":
	  if time.time()-prevtime > 0.2 or self.res<=self.tol: 
	    prevtime = time.time()
	    restab.append(self.res)
	    counttab.append(self.it)
	    logplot.set_xdata(counttab)
	    logplot.set_ydata(restab)
	    gca().set_xlim(0,self.it*1.3)
	    gca().set_ylim(min(restab)*0.8,max(restab)*1.2)
	    if self.res<=self.tol:
	      logplot.set_color('black')
	      logplot.set_linewidth(1)
	    draw()
	else:
	  print "Unknown log output. Choose between stdout and X."
	  sys.exit()
      if isnan(self.res):
	print "blow up"
	sys.exit()

      # place input for next iteration, fnew, in self.fn    
      self.fn = fnew     


    # Iteration is finished, compute useful results from our 
    # brand new gamma in self.fn !
    # Get direct correlation function c :        # it is c_short if Ng_renorm
    self.compute_c_from_closure(self.fn)
    # Get total correlation function
    self.h = self.fn + self.c                    # same relation whether Ng_renorm or not
    # Get structure factor : 
    self.k, self.h_hat = self.fbt(self.r,self.h)
    self.h_hat *= 4.*pi
    self.S = 1 + rho*self.h_hat
    # Get radial distribution function
    self.g = self.h + 1
    return self.g, self.k, self.S, self.fn 

  def gn_from_fn(self,rho,fn):
    # compute direct correlation functin c(r) in real space
    # from knowledge for initial guess Gamma (fn)
    self.compute_c_from_closure(fn)

    # compute Fourier-Bessel transform of c.
    self.k, self.c_hat = self.fbt(self.r, self.c)
    #scale
    self.c_hat *= rho

    # solve OZ in fourier space to get output for FT(Gamma) (self.gn_hat)
    if self.Ng_renorm:
      # Ng renormalization: here c_hat is the short-ranged c, Gam_hat also is short-ranged
      #self.beta_ucorr_hat = self.Z**2*self.lB*\
      #                      (np.exp(self.kappa*self.a_pot)/(1.0+self.kappa*self.a_pot))**2/\
      #                      (self.kappa**2+self.k**2) 
      self.beta_ucorr_hat     = zeros(len(self.k))       
      self.beta_ucorr_hat[1:] = self.Z**2*self.lB*\
                                (np.exp(self.kappa*self.a_pot)/(1.0+self.kappa*self.a_pot))**2*\
                                exp(-(self.k[1:]*self.a_pot)**2/(4.*1.08**2))/(self.a_pot*self.k[1:])**2 
      self.beta_ucorr_hat[0]  = 10.0

    else:
      # No Ng renormalization
      self.beta_ucorr     = 0.
      self.beta_ucorr_hat = 0.
    #sys.exit()
    self.Gam_hat = (self.c_hat - self.beta_ucorr_hat)/(1-4*pi*(self.c_hat-self.beta_ucorr_hat)) - self.c_hat
    self.Gam_hat /= rho

    # Transform it back to real space to get output of OZ (self.gn)
    dummy, gn = self.ifbt(self.k, self.Gam_hat)
    gn[0]  = 0.

    return gn


  def Ng2(self,gn,fn,gn1,fn1,gn2,fn2):
    # 2 parameter Ng acceleration
    dn  = gn -fn
    dn1 = gn1-fn1
    dn2 = gn2-fn2
    d01 = dn-dn1
    d02 = dn-dn2
    d01d01 = trapz(d01*d01,self.r)#/self.N**2  
    d01d02 = trapz(d01*d02,self.r)#/self.N**2
    d02d02 = trapz(d02*d02,self.r)#/self.N**2
    dnd01  = trapz(dn *d01,self.r)#/self.N**2
    dnd02  = trapz(dn *d02,self.r)#/self.N**2
    det = d01d01*d02d02-d01d02**2
    if not det==0:
      self.c1 = (dnd01*d02d02-dnd02*d01d02) / det
      self.c2 = (d01d01*dnd02-d01d02*dnd01) / det
      fnp1 = (1.-self.c1-self.c2)*gn + self.c1*gn1 + self.c2*gn2
    else:
      fnp1 = gn # just do one picard step
    return fnp1

  def Ng5(self):
    # 5 parameter Ng acceleration, if possible
    dn  = self.gn -self.fn

    if (self.it >=10  and self.Ng and self.res<self.res_start_Ng ) or self.Ng_step>=0: # two parameter Ng

      # five parameter Ng
      d0 = zeros((5,self.N))
      d0[0,:] = dn-(self.gn1-self.fn1)
      d0[1,:] = dn-(self.gn2-self.fn2)
      d0[2,:] = dn-(self.gn3-self.fn3)
      d0[3,:] = dn-(self.gn4-self.fn4)
      d0[4,:] = dn-(self.gn5-self.fn5)
      AA = zeros((5,5))
      BB = zeros(5)
      for i in range(5):
        BB[i] = trapz(dn*d0[i,:],self.r)
        for j in range(5):
          AA[i,j] = trapz(d0[i,:]*d0[j,:],self.r)
      detAA = det(AA)	
      if not detAA==0.0 :
        # Fine, let's solve              
	AA = AA / detAA 
	BB = BB / detAA
        ci = solve(AA,BB)
        fnp1 = (1.-sum(ci))*self.gn + ci[0]*self.gn1 + ci[1]*self.gn2+ ci[2]*self.gn3 + ci[3]*self.gn4+ ci[4]*self.gn5
      else:
        # matrix is singular, switch to Ng2 
        print "matrix singular, det=", det(AA)
        print AA
        # Use Ng2 instead 
        fnp1 = self.Ng2(self.gn,self.fn,self.gn1,self.fn1,self.gn2,self.fn2)
        #fnp1 = gn # this will lead to a simple picard step

      ## Ng2
      #fnp1 = self.Ng2(self.gn,self.fn,self.gn1,self.fn1,self.gn2,self.fn2)
      
      if self.res > 1.0e-2: #mix ng and picard slowly
        fnp1 = (1-self.relax)*self.fn + self.relax*fnp1    

    else: # picard
      fnp1 = (1-self.relax)*self.fn + self.relax*self.gn
      self.relax *= self.relaxfactor

    # Update input/output functions of OZ
    if self.it>=4:
      self.fn5 = self.fn4; self.gn5 = self.gn4     
    if self.it>=3:
      self.fn4 = self.fn3; self.gn4 = self.gn3     
    if self.it>=2:
      self.fn3 = self.fn2; self.gn3 = self.gn2     
    if self.it>=1:
      self.fn2 = self.fn1; self.gn2 = self.gn1     
    self.fn1 = self.fn; self.gn1 = self.gn

    return fnp1

  def GetCompressibilitiesRY(self,gam_init=-1):
    # Solve OZ for two rhos close enough to assess dP/drho. Zeff must be the same for the two.
    self.chi_S = 1. ; self.chi_P = 0.

    if not type(gam_init)==ndarray:
      gam_init = zeros(self.N)
    else:
      #print "using provided gam_init", gam_init
      pass

    rho1 = self.rho*1.01
    self.g1, k1, self.S1, self.gam1 = self.SolveOZInternal(rho1,gam_init=gam_init,quiet=self.RYQuiet )
    self.betaP1 = self.ComputePressure(rho1,self.g1,self.beta_u)

    rho2 = self.rho*0.99
    self.g2, k2, self.S2, self.gam2 = self.SolveOZInternal(rho2,gam_init=self.gam1,quiet=self.RYQuiet ) 
    self.betaP2 = self.ComputePressure(rho2,self.g2,self.beta_u)

    # Compressibility as derivative of rho vs betaP
    self.chi_P = (rho1-rho2)/(self.betaP1-self.betaP2)

    # Compressibility as S(0)
    self.g , k , self.S , self.gam  = self.SolveOZInternal(self.rho,gam_init=self.gam1,quiet=self.RYQuiet )

    coeff = polyfit(self.k[1:6],self.S[1:6],4)
    self.chi_S = coeff[-1]

    return self.chi_S, self.chi_P, self.g, self.S, self.gam


  def FindRYalpha(self,gam_init=-1):

    print "--------------------- RY iteration start --------------------"
    print "    alpha           chi_S           chi_P         chi_S-chi_P"

    if self.alphamax==-1:
      self.alphamax = 2.0/self.a   # Full HNC
    self.alpha = self.alphamax
    chi_S_alphamax, chi_P_alphamax, g, S, gam = self.GetCompressibilitiesRY(gam_init=gam_init)
    fmax = chi_S_alphamax-chi_P_alphamax
    print    format(self.alphamax,'+4.6e')+'   '+format(chi_S_alphamax,'+4.6e')+'   '+format(chi_P_alphamax,'+4.6e')+'   '+format(fmax,'+4.6e')

    if self.alphamin==-1:
      # if not defined, assume alpha is between 0.001 and 2/d
      self.alphamin = .001/self.a  # Full PY
    self.alpha = self.alphamin
    chi_S_alphamin, chi_P_alphamin, g, S, gam = self.GetCompressibilitiesRY(gam_init=gam_init)
    fmin = chi_S_alphamin-chi_P_alphamin
    print    format(self.alphamin,'+4.6e')+'   '+format(chi_S_alphamin,'+4.6e')+'   '+format(chi_P_alphamin,'+4.6e')+'   '+format(fmin,'+4.6e')


    if self.solver == "Bisection":
      self.alphamid = 0.5*(self.alphamin+self.alphamax)
    elif self.solver == "RegulaFalsi": 
      self.alphamid = self.alphamax - fmax*(self.alphamax-self.alphamin)/(fmax-fmin)
      if (not self.alphamid < self.alphamax) or (not self.alphamin<self.alphamid): 
          print "bounds not wide enough"
          sys.exit()  
    else:
      print "Unknown solver: select Bisection or RegulaFalsi"
      sys.exit()
    self.alpha = self.alphamid
    chi_S_alphamid, chi_P_alphamid, g, S, gam = self.GetCompressibilitiesRY(gam_init=gam_init)
    fmid = chi_S_alphamid-chi_P_alphamid
    print    format(self.alphamid,'+4.6e')+'   '+format(chi_S_alphamid,'+4.6e')+'   '+format(chi_P_alphamid,'+4.6e')+'   '+format(fmid,'+4.6e')



    # Get direct correlation function c :
    self.compute_c_from_closure(gam)
    # Get total correlation function
    self.h = gam + self.c
    # Get structure factor : 
    self.k, self.h_hat = self.fbt(self.r,self.h)
    self.h_hat *= 4.*pi
    self.S = 1 + self.rho*self.h_hat
    # Get radial distribution function
    self.g = self.h + 1

    if fmin*fmax>0 and abs(fmid)>(chi_S_alphamid+chi_P_alphamid)/2 * 0.001:
      print "stop. alpha bounds are not wide enough"
      sys.exit(0)

    if fmid*fmax>0:
      fmax = fmid
      self.alphamax = self.alphamid
      last_side_retained = "left"
    else:
      fmin = fmid
      self.alphamin = self.alphamid
      last_side_retained = "right"



    while abs(fmid)>(chi_S_alphamid+chi_P_alphamid)/2 * 0.001:

      if self.solver == "Bisection":
        self.alphamid = 0.5*(self.alphamin+self.alphamax)
      elif self.solver == "RegulaFalsi": # Regula Falsi
        self.alphamid = self.alphamax - fmax*(self.alphamax-self.alphamin)/(fmax-fmin)
      else:
        print "Unknown solver: select Bisection or RegulaFalsi"
        sys.exit()
      self.alpha = self.alphamid
      chi_S_alphamid, chi_P_alphamid, g, S, gam = self.GetCompressibilitiesRY(gam_init=gam_init)
      fmid = chi_S_alphamid-chi_P_alphamid

      if fmid*fmax>0:
        m    = 1.0 - fmid/fmax              # m /= 0.5 is Anderson-Bjork version 
        fmax = fmid
	self.alphamax = self.alphamid
        if last_side_retained == "left":
          if m<0: m = 0.5                     # 0.5 is Illinois version
          fmin *= m  
        last_side_retained = "left"
      else:
        m    = 1.0 - fmid/fmin              # m /= 0.5 is Anderson-Bjork version 
        fmin = fmid
	self.alphamin = self.alphamid
        if last_side_retained == "right":
          if m<0: m = 0.5                     # 0.5 is Illinois version
          fmax *= m  
        last_side_retained = "right"

      #print "using alpha=", alphamid
      #print "and getting chi_S, chi_P, err =", chi_S_alphamid, chi_P_alphamid, abs(fmid)
      print    format(self.alphamid,'+4.6e')+'   '+format(chi_S_alphamid,'+4.6e')+'   '+format(chi_P_alphamid,'+4.6e')+'   '+format(fmid,'+4.6e')

      # Get direct correlation function c :
      self.compute_c_from_closure(gam)
      # Get total correlation function
      self.h = gam + self.c
      # Get structure factor : 
      self.k, self.h_hat = self.fbt(self.r,self.h)
      self.h_hat *= 4.*pi
      self.S = 1 + self.rho*self.h_hat
      # Get radial distribution function
      self.g = self.h + 1
      # Compute pressure
      self.ComputePressure(self.rho,self.g,self.beta_u)
    # Fix S[0]
    coeff = polyfit(self.k[1:6],self.S[1:6],4)
    self.S[0] = coeff[-1]  
    return self.g, self.k, self.S, gam

  def ComputePressure(self,rho,g,beta_u):
    # Pressure from g
    kr = self.kappa*self.r
    # soft part of the yukawa pot
    integ = trapz(-g[self.r>self.HSd]*beta_u[self.r>self.HSd]*(kr[self.r>self.HSd]+1.)*4*pi*self.r[self.r>self.HSd]**2,self.r[self.r>self.HSd])

    # Electrostatic contribution
    self.betaPEL =  -rho**2/6.*integ

    # add hard core part: compute g(2)
    i = 0
    while self.r[i]<self.HSd: i+=1
    g2 = 2*g[i] - g[i+1]

    self.betaP = rho + 4*rho**2*(pi*self.HSd**3/6.)*g2 + self.betaPEL 

    return self.betaP

  def ExactSolutionPY(self,k):
    l1 = (1+2*self.phi)**2/(1-self.phi)**4
    l2 = -(1+.5*self.phi)**2/(1-self.phi)**4
    kd = k*self.a*2
    NC  = -24*self.phi*(  l1*(sin(kd)-kd*cos(kd))/kd**3 \
	            -6*self.phi*l2*(kd**2*cos(kd)-2*kd*sin(kd)-2*cos(kd)+2)/kd**4 \
		     -0.5*self.phi*l1*(kd**4*cos(kd)-4*kd**3*sin(kd)-12*kd**2*cos(kd)+24*kd*sin(kd)+24*cos(kd)-24)/kd**6)
    self.S = 1./(1-NC)
    self.h_hat = (self.S - 1.0) / self.rho
    self.r, self.h = self.ifbt(self.k,self.h_hat)
    self.h /= (2*pi**2)
    self.fn = 0.
    self.g = self.h + 1
    return self.g, self.k, self.S, self.fn 


  def HayterPenfold(self,eta=-1,sigma=-1,k=-1,gamma=-1,MSAquiet=True):

    # Used for validation against Hayter Penfold 1981
    #sigma = 5.0e-9
    #eta = 0.3
    #k   = 2.0
    #gamma = exp(k)*pi*eps*sigma*(0.078)**2/kT

    N = 2**14

    delta = 1-eta
    alpha1 = -(2*eta+1.0)*delta/k
    alpha2 = (14.0*eta**2-4*eta-1.0)/k**2
    alpha3 = 36.0*eta**2/k**4
    beta1  = -(eta**2+7.0*eta+1.0)*delta/k
    beta2  = 9*eta*(eta**2+4.0*eta-2.0)/k**2
    beta3  = 12.0*eta*(2*eta**2+8*eta-1)/k**4
    nu1    = -(eta**3+3*eta**2+45*eta+5)*delta/k
    nu2    = (2*eta**3+3*eta**2+42*eta-20)/k**2
    nu3    = (2*eta**3+30*eta-5)/k**4
    nu4    = nu1 + 24*eta*k*nu3
    nu5    = 6*eta*(nu2+4*nu3)
    phi1   = 6*eta/k
    phi2   = delta - 12*eta/k**2
    tau1   = (eta+5)/(5*k)
    tau2   = (eta+2)/k**2
    tau3   = -12*eta*gamma*exp(-k)*(tau1+tau2)
    tau4   = 3*eta*k**2*(tau1**2-tau2**2)
    tau5   = 3*eta*(eta+8)/10.-2*(2*eta+1)**2/k**2
    a1     = (24*eta*gamma*exp(-k)*(alpha1+alpha2+(1+k)*alpha3)-(2*eta+1)**2)/delta**4
    a2     = 24*eta*(alpha3*(sinh(k)-k*cosh(k))+alpha2*sinh(k)-alpha1*cosh(k))/delta**4
    a3     = 24*eta*((2*eta+1)**2/k**2-delta**2/2.0+alpha3*(cosh(k)-1-k*sinh(k))-alpha1*sinh(k)+alpha2*cosh(k))/delta**4
    b1     = (3*eta*(eta+2)**2/2.0-12*eta*gamma*exp(-k)*(beta1+beta2+(1+k)*beta3))/delta**4
    b2     = 12*eta*(beta3*(k*cosh(k)-sinh(k))-beta2*sinh(k)+beta1*cosh(k))/delta**4
    b3     = 12*eta*(delta**2*(eta+2)/2.0-3*eta*(eta+2)**2/k**2-beta3*(cosh(k)-1-k*sinh(k))+beta1*sinh(k)-beta2*cosh(k))/delta**4
    v1     = ((2*eta+1)*(eta**2-2*eta+10)/4.0-gamma*exp(-k)*(nu4+nu5))/(5*delta**4)
    v2     = (nu4*cosh(k)-nu5*sinh(k))/(5*delta**4)
    v3     = ((eta**3-6*eta**2+5)*delta-6*eta*(2*eta**3-3*eta**2+18*eta+10)/k**2+24*eta*nu3+nu4*sinh(k)-nu5*cosh(k))/(5*delta**4)
    p1     = (gamma*exp(-k)*(phi1-phi2)**2-(eta+2)/2.0)/delta**2
    p2     = ((phi1**2+phi2**2)*sinh(k)+2*phi1*phi2*cosh(k))/delta**2
    p3     = ((phi1**2+phi2**2)*cosh(k)+2*phi1*phi2*sinh(k)+phi1**2-phi2**2)/delta**2
    t1     = tau3+tau4*a1+tau5*b1
    t2     = tau4*a2+tau5*b2+12*eta*(tau1*cosh(k)-tau2*sinh(k))
    t3     = tau4*a3+tau5*b3+12*eta*(tau1*sinh(k)-tau2*(cosh(k)-1))-2*eta*(eta+10)/5.0-1
    mu1    = t2*a2-12*eta*v2**2
    mu2    = t1*a2+t2*a1-24*eta*v1*v2
    mu3    = t2*a3+t3*a2-24*eta*v2*v3
    mu4    = t1*a1-12*eta*v1**2
    mu5    = t1*a3+t3*a1-24*eta*v1*v3
    mu6    = t3*a3-12*eta*v3**2
    lambda1 = 12*eta*p2**2
    lambda2 = 24*eta*p1*p2-2*b2
    lambda3 = 24*eta*p2*p3
    lambda4 = 12*eta*p1**2-2*b1
    lambda5 = 24*eta*p1*p3-2*b3-k**2
    lambda6 = 12*eta*p3**2
    omega12 = mu1*lambda2-mu2*lambda1
    omega13 = mu1*lambda3-mu3*lambda1
    omega14 = mu1*lambda4-mu4*lambda1
    omega15 = mu1*lambda5-mu5*lambda1
    omega16 = mu1*lambda6-mu6*lambda1
    omega24 = mu2*lambda4-mu4*lambda2
    omega25 = mu2*lambda5-mu5*lambda2
    omega26 = mu2*lambda6-mu6*lambda2
    omega34 = mu3*lambda4-mu4*lambda3
    omega35 = mu3*lambda5-mu5*lambda3
    omega36 = mu3*lambda6-mu6*lambda3
    omega4 = omega16**2-omega13*omega36
    omega3 = 2*omega16*omega15-omega13*(omega35+omega26)-omega12*omega36
    omega2 = omega15**2+2*omega16*omega14-omega13*(omega34+omega25)-omega12*(omega35+omega26)
    omega1 = 2*omega15*omega14-omega13*omega24-omega12*(omega34+omega25)
    omega0 = omega14**2-omega12*omega24
    Ftab = roots([omega4,omega3,omega2,omega1,omega0])   
    g = -1
    ming2 = 1000
    for F in Ftab:
      if not iscomplex(F):
        C = -(omega16*F**2+omega15*F+omega14)/(omega13*F+omega12)
	if not MSAquiet: print -(p1+p2*C+p3*F)
        B = b1+b2*C+b3*F
        A = a1+a2*C+a3*F
        V = (v1+v2*C+v3*F)/A
        U = (k**3*C-24*eta*gamma)/(24*eta*gamma*k**2)
        K = linspace(0,1000,N)
        aK = A*(sin(K)-K*cos(K))/K**3 \
                + B*((2/K**2-1)*K*cos(K)+2*sin(K)-2/K)/K**3 \
                + eta*A*(24/K**3+4*(1-6/K**2)*sin(K) \
                   - (1-12/K**2+24/K**4)*K*cos(K))/(2*K**3) \
                + C*(k*cosh(k)*sin(K)-K*sinh(k)*cos(K))/(K*(K**2+k**2)) \
                + F*(k*sinh(k)*sin(K)-K*(cosh(k)*cos(K)-1))/(K*(K**2+k**2)) \
                + F*(cos(K)-1)/K**2 \
                - gamma*exp(-k)*(k*sin(K)+K*cos(K))/(K*(K**2+k**2))

        S = 1.0/(1.0-24*eta*aK)
        S[0] = -1./A         

        # compute one value of g
        #g14 = trapz(K*(S-1)*sin(K*0.25),K)/(12*pi*eta*0.25)

        ## Explicit FBT to check 
        #x = linspace(0.1,10,1000)
        #g = 0.*x
        #for i in range(len(x)):
        #  g[i] = 1+1/(12*pi*eta*x[i])*trapz( (S-1)*K*sin(K*x[i]) , K)    
        #figure(2)
        #plot(x,g,'o')

        #x, myifbt = self.ifbt(K,(S-1)**2/S)
        #beta_u_star = gamma*exp(-k*x)/x#; beta_u_star[x<1]=0
        #g3 = 1 - beta_u_star + myifbt / (24*eta)
        #plot(x,g2)
        #plot(x,g3,'r')
        #plot(x,-beta_u_star,'g')
        #sys.exit()

        # works for RMSA
        x, myifbt = self.ifbt(K,S-1)  
        g2 = 1 + myifbt/(24*eta)

	# return c if needed
	c2        = - gamma * exp( -k*x ) / x       # only for x>1
	c2[x<1.0] = A + B*x[x<1] + 0.5*eta*A*x[x<1]**3 + C*sinh(k*x[x<1])/x[x<1] + F*(cosh(k*x[x<1])-1.0)/x[x<1]   # for x<1

        #print abs(mean(g2[x<0.5][1:]))
        if ming2 > abs(mean(g2[x<0.5][1:])): 
          ming2 = abs(mean(g2[x<0.5][1:]))
          g = g2
	  c = c2

    if ming2 > 0.1:
      print "ming2 was ", ming2, "in MSA calculation. Something may be wrong.\n"
      sys.exit()

    # K=k*sigma and x=r/sigma
    return  K, S, x, g, c

  def RMSAHSYSolution(self,eta=-999999, sigma=-999999, k=-999999, gamma=-999999, RMSAquiet=True):

    """
    RMSA Following Heinen 2011 / Hansen 1983
    """

    tol = 1.0e-6

    # K is -gam.exp(-z) 
    # k is z_Hoye

    def gc(s):
      eta_prime   = eta/s**3
      eta_star    = eta_prime  
      sigma_prime = sigma/s
      sigma_star  = sigma_prime
      gamma_prime = gamma*s
      gamma_star  = gamma_prime 
      k_prime     = k_mod / s
      k_star      = k_prime
      # solve MSA with star variables	
      q, S, r, g, cMSA  = self.HayterPenfold(eta=eta_star,sigma=sigma_star,k=k_star,gamma=gamma_star)
      # here r is actually defined by sigma_star, so it's r'


      #global coeffs
      #coeffs = polyfit(r[ r>1 ][10:30],g[ r>1 ][10:30],1)
      g_cont      = g[r>1][0] #coeffs[0]*1 + coeffs[1] 
      if not RMSAquiet: print "s in = ", s, g_cont,  abs(g[ r>sigma_star ])[0:10]
      if s<=0 or s>1: g_cont=1000
      return g_cont, q, S, r, g, cMSA


    # First try without rescaling (step 2)
    gamma_star  = gamma
    k_mod       = k
    # get MSA solution for gamma_star (PB) and k_mod (M)
    q, S, r, g, cMSA  = self.HayterPenfold(eta=eta,sigma=sigma,k=k_mod,gamma=gamma_star)
    g_cont      = g[r>1][0]

    # if MSA restult is not satisfactory...
    if g_cont < 0:
    
      # perform RMSA by solving for s such that g_cont=0
      s, info, ier, msg = fsolve( lambda s: abs(gc(s[0])[0]),0.95,xtol=1.0e-4, full_output=True)
      if not ier==1:
        print "Beware in RMSA."
        print "last s = ", s
        print info
        print "ier = ", ier
        print msg
      s = s[0]

      ## Same with dichotomy
      #b = 1.; fb = g_cont # obtained for s=1
      #a = 0.5; fa = gc(a)[0]
      #if fa*fb>0: 
      #  print "fail", fb, fa
      #  sys.exit()
      #fc = 1
      #while abs(fc)>1.0e-2:
      #  c = 0.5*(a+b); fc = gc(c)[0]
      #  if fa*fc<0:
      #    b = c; fb = fc
      #  else:
      #    a = c; fa = fc
      #  print c  
      #s = c	

      # one more time with correct "s" to get the g and S
      g_cont, qstar, S, rstar, g, cMSA = gc(s)

      # scale back
      q = qstar*s 
      r = rstar/s

    # let's compute the pressure with the virial theorem
    beta_u = gamma*exp(-k*r)/r
    rho = eta/(pi*sigma**3/6.)
    PEL = - kT*rho**2/6*trapz(-g[r>1]*beta_u[r>1]*(k*r[r>1]+1.)*4*pi*(r[r>1]*sigma)**2,r[r>1]*sigma)
    Pcont = 4*rho**2*kT*(pi*sigma**3/6.)*g_cont 
    Pvir = rho*kT + Pcont + PEL 
    # no contact part by construction of RMSA

    return q, S, r, g, Pvir, cMSA

  def HardSpheres_rdf(self,r,eta,sigma):
    """
    From Trokhymchuk2005, and erratum. 
    r     : distance between centers
    eta   : volume fraction
    sigma : diameter
    """
    g = 0.*r
    d = (2*eta*(eta**2-3*eta-3.+np.sqrt(3.*(eta**4-2.0*eta**3+eta**2+6.0*eta+3.0))))**(1./3.)
    beta0sigma = 2*eta/(1.0-eta)*np.sqrt(3.)*(-d/(4*eta)-eta/(2*d))
    alpha0sigma = 2*eta/(1.0-eta)*(-1.0+d/(4*eta)-eta/(2*d))
    musigma = 2*eta/(1.0-eta)*(-1.0-d/(2.*eta)+eta/d) # cf erratum
    mu0sigma = musigma 
    beta0 = beta0sigma/sigma
    alpha0 = alpha0sigma/sigma
    mu0 = mu0sigma/sigma 
    mu = mu0
    gamma = np.arctan(-1.0/beta0*(sigma*(alpha0*(alpha0**2+beta0**2)-mu0*(alpha0**2-beta0**2))*(1+0.5*eta)+(alpha0**2+beta0**2-mu0*alpha0)*(1+2.0*eta))/(sigma*(alpha0**2+beta0**2-2.0*mu0*alpha0)*(1.0+0.5*eta)-mu0*(1+2.0*eta)))
    gamma0 = gamma
    omega = (-0.682*np.exp(-24.697*eta)+4.720+4.450*eta)/sigma
    kappa = (4.674*np.exp(-3.935*eta)+3.536*np.exp(-56.270*eta))/sigma
    alpha = (44.554+79.868*eta+116.432*eta**2-44.652*np.exp(2*eta))/sigma
    beta  = (-5.022+5.857*eta+5.089*np.exp(-4.0*eta))/sigma
    rstar = sigma*(2.0116-1.0647*eta+0.0538*eta**2)
    gm = 1.0286-0.6095*eta+3.5781*eta**2-21.3651*eta**3+42.6344*eta**4-33.8485*eta**5
    gsigmaexpt = 1./(4*eta)*((1+eta+eta**2-2./3.*eta**3-2./3.*eta**4)/(1.0-eta)**3-1.0)
    B = (gm-(sigma*gsigmaexpt/rstar)*np.exp(mu*(rstar-sigma)))/(np.cos(beta*(rstar-sigma)+gamma)*np.exp(alpha*(rstar-sigma))-np.cos(gamma)*np.exp(mu*(rstar-sigma)))*rstar
    A = sigma*gsigmaexpt-B*np.cos(gamma)
    delta = -omega*rstar-np.arctan((kappa*rstar+1.0)/(omega*rstar))
    C = (rstar*(gm-1.0)*np.exp(kappa*rstar))/np.cos(omega*rstar+delta)
    g = A/r*np.exp(mu*(r-sigma))+B/r*np.cos(beta*(r-sigma)+gamma)*np.exp(alpha*(r-sigma)) 
    g[r<sigma] = 0.  
    g[r>=rstar] = 1.0+C/r[r>=rstar]*np.cos(omega*r[r>=rstar]+delta)*np.exp(-kappa*r[r>=rstar])
    return g
  
  def HardSpheres_beta_f(self,eta):
    """
    Excess helmholtz free energy from CS EOS
    """
    return (4.0*eta-3.0*eta**2)/(1.0-eta)**2


  def FexcPerturb(self, rho, N=-1, tol=1.0e-20):
  
    if N==-1: N = self.N

    # redefine a more refined r and beta_u
    r = np.linspace(0, 20./self.kappa, N)
    beta_u = self.update_yukawa(rho,self.g,self.kappa*self.a,self.kappa*r)

    #from scipy.integrate import quadrature
    def fun( d ):
      if d<2*self.a or d>(6./(pi*rho))**(1./3.):
	return 1.0e12
      # maintain rho but eta varies with effective diameter d  
      eta = rho * pi*d**3/6.
      val = self.HardSpheres_beta_f(eta)  + trapz( rho*2*pi*r**2*self.HardSpheres_rdf(r,eta,d)*beta_u , r )
      if isnan(val):
	return 1.0e12
      return val
  
    b1 = self.a*2.0 
    b2 = self.a*3.2
    result = minimize_scalar( fun, bracket=(b1, b2) , method='brent' , tol=tol )
    if abs(result.x-b2)<1.0e-12:
      print "Bounded optimization ended up on upper bound"
      sys.exit()
    Fexc_s_NkT = result.fun
  
    return Fexc_s_NkT 

  def ComputePressurePerturb(self, N=-1, tol=1.0e-20):

    # Following Zoetekouw & van Roij 2006 / Gibbs-Bogoliubov inequality
    F = []
    rho = []
    for i in range(5):
      rho.append(self.rho*(0.98+i/100.))
      F.append(self.FexcPerturb( rho[-1] , N, tol))
    rho = array(rho)
    F = array(F)
    #figure()
    #plot(rho,F,'o')
    c=polyfit(rho,F,1)
    #plot(rho,c[0]*rho+c[1],'r')
    # c[0] is d(Fex/NkT)/drho and Piex is rho**2.kT.d(Fex/NkT)/drho
    betaPi_ex = c[0]*self.rho**2
    betaPi = self.rho + betaPi_ex
    return betaPi


  def ValidationHS(self):
    # Comparison with Hard Sphere Solution   
    self.a = 100.0e-9
    Vp = 4./3*pi*self.a**3
    self.phi = 0.41887902047863906 #0.40 
    self.potential = "HS"
    self.relax = 0.01
    self.closure = "PY"
    self.quiet = False
    self.N     = 8192 
    self.Ndmax = 20.  
    self.log = "X"
    self.Ng = True
    self.UpdateParameters()
    g, k, S, gam = self.SolveOZ()  # initialize with c at previous phi
    l1 = (1+2*self.phi)**2/(1-self.phi)**4
    l2 = -(1+.5*self.phi)**2/(1-self.phi)**4
    kd = k*self.a*2
    NC  = -24*self.phi*(  l1*(sin(kd)-kd*cos(kd))/kd**3 \
	            -6*self.phi*l2*(kd**2*cos(kd)-2*kd*sin(kd)-2*cos(kd)+2)/kd**4 \
		     -0.5*self.phi*l1*(kd**4*cos(kd)-4*kd**3*sin(kd)-12*kd**2*cos(kd)+24*kd*sin(kd)+24*cos(kd)-24)/kd**6)
    Sth = 1./(1-NC)
    figure()
    plot(self.k*self.a,S,'o')
    plot(self.k*self.a,Sth,'r',lw=2)
    print "Pressure from OZ / Pressure from CS (ok, we don't know u' for HS interactions): ", self.ComputePressure(self.rho,self.g,self.beta_u)*kT, self.rho*kT*(1+self.phi+self.phi**2-self.phi**3)/(1-self.phi)**3
    print "Compress from OZ / Compress from CS : ", self.S[1], Sth[1]
    xlim(0,35)

    # MSA solution (should be same as PY)
    K,S,x,g,_ = self.HayterPenfold(eta=self.phi,sigma=2*self.a,k=1,gamma=0.01)
    plot(K/2,S,'--g')
    figure()
    plot(x,g)
    K, S, x, gRMSA, PRMSA, cRMSA = self.RMSAHSYSolution(eta=self.phi,sigma=2*self.a,k=1,gamma=.01)
    print "RMSA pressure = ", PRMSA

####################################################################################


####################################################################################
# Section 3
# Define the cell model solver as an object
class CellModelSimulation():
  def __init__(self):
    self.adim  = False 
    self.eps   = 78.
    self.a     = -1. 
    self.I     = -1.
    self.sigma = -1.
    self.phi   = -1. 
    self.quiet = True
    self.resolution = 400
    self.n_0   = -1.
    self.Pi    = 0.
    self.gamma0 = 0.
    self.psi    = 0.


  def SetParameters(self,eps=78, a=-1, I=-1, sigma=-1, phi=-1, quiet=True, res=400, adim=False):
    """
    Set the CM simulation parameters. eps is the relative permittivity of the medium, I the ionic strength in mol/L, sigma \
            the surface charge density in e/nm2, phi the volume fraction
    """
    self.eps = eps
    self.a = a
    self.I = I
    self.sigma = sigma
    self.phi = phi
    self.quiet = quiet
    self.resolution = res
    self.ComputeConstants()
    self.BuildMesh()

  def ComputeConstants(self):
    # Check 
    if self.eps<0 or self.a<0 or self.I<0 or self.phi<0:
      print "Parameters need being set: eps, a, I, sigma, phi"
      sys.exit()
    # Bulk number density:
    self.n_0 = self.I*1000*Na
    S = 4*pi*self.a**2  # surface in nm**2
    self.Z = self.sigma * S  # nb charges
    # Debye length
    self.lambdaD = sqrt(eps0*self.eps*kT/(2*self.n_0*elc**2))
    # Non-dimensional particle radius (a comes in nm)
    self.a = self.a*1e-9/self.lambdaD
    # Radius of the cell (non-dimensional)
    self.R = self.a/self.phi**(1./3.)
    self.V = 4./3.*pi*self.R**3
    # Non-dimensional charge (sigma comes in e/nm**2)
    self.sigmascale = sqrt( 2*self.n_0*eps0*self.eps*kT)
    self.sigma     = self.sigma    *elc*1.0e18/self.sigmascale
    # Bjerrum length
    self.lambdaB = elc**2/(4*pi*kT*eps0*self.eps)

  def BuildMesh(self):  
    self.dr = (self.R-self.a)/self.resolution;
    self.r = arange(self.a, self.R, self.dr)
    return

  def f_const_charge(self, psi):
    out = zeros(len(self.r))
    out[0]    = sinh(psi[   0]) - \
                1./self.r[0   ]**2*( (self.r[   0]+self.dr*0.5)**2*(psi[1   ]-psi[0   ])/self.dr - \
	                             (self.r[   0]-self.dr*0.5)**2*(psi[0   ]-psi[1   ]-self.sigma*2*self.dr)/self.dr \
	         		     )/self.dr
    out[1:-1] = sinh(psi[1:-1]) - \
	        1./self.r[1:-1]**2*( (self.r[1:-1]+self.dr*0.5)**2*(psi[2:  ]-psi[1:-1])/self.dr - \
		                     (self.r[1:-1]-self.dr*0.5)**2*(psi[1:-1]-psi[0:-2])/self.dr \
				     )/self.dr
    out[-1]   = sinh(psi[  -1]) - \
	        1./self.r[  -1]**2*( (self.r[  -1]+self.dr*0.5)**2*(psi[  -2]-psi[  -1])/self.dr - \
	                             (self.r[  -1]-self.dr*0.5)**2*(psi[  -1]-psi[  -2])/self.dr \
				     )/self.dr

    return out

  def solve(self):
    # Solve the PB equation with BC to obtain the potential psi.     
    initpsi = 0.
    ier = 0
    while not ier == 1:
      initpsi += 1
      if not self.quiet: print "Trying CM resolution with initial potential = ", initpsi
      psi_0 = initpsi*ones(len(self.r))
      x, infodict, ier, mesg = fsolve( lambda pot: self.f_const_charge(pot), psi_0,full_output=True,xtol=1.0e-12, maxfev=100000000)
    self.psi = x
    # The osmotic pressure is given by the ion concentration at the cell boundary. 
    self.Pi =  ( cosh(self.psi[-1]) - 1. ) * 2*self.n_0*kT

  def renormalize(self,type='SCR'):
    """
    Renormalize following Trizac et al. (2003)'s Alexander's prescription revisited.
    """
    # Solve the non-linear problem for constant charge 
    self.lin = False
    self.solve()
    # Get the potential at the cell boundary
    psi_R = self.psi[-1]
    # Compute the inverse screening length at cell boundary 
    self.kappa_PB = sqrt(cosh(psi_R))/self.lambdaD
    self.gamma0 = tanh(psi_R)
    kR = self.kappa_PB*self.R*self.lambdaD
    ka = self.kappa_PB*self.a*self.lambdaD
    # Define the EPC value for the charge Q
    self.Q = self.gamma0/(self.kappa_PB*self.lambdaB)*( kR*cosh(kR)-sinh(kR))
    # Transform it into a valency to be used in a Yukawa potential for a sphere with dimensionless radius ka.
    self.Zeff = (1+ka)*exp(-ka)*self.Q
    return self.Zeff, self.kappa_PB

####################################################################################



####################################################################################
# SECTION 4
# A program using the different solvers built above. Please, read the comments in this section to know how to use these models. 


# Lucas Goehring's data for Ludox HS40 compression at 5mM salt. It will be used to plot results 
# [Goehring et al., 2017 Drying paint: from micro-scale dynamics to mechanical instabilities. 
#  Phil. Trans. R. Soc. A 375: 20160161. http://dx.doi.org/10.1098/rsta.2016.0161]
phi_exp = array([3.79616151,4.67357825,5.79318335,6.65619683,6.67927608,8.53266975,10.05356483,11.22118409,11.89143766,13.21947069,16.17580735,18.7878745,19.20054415,20.4057481,21.06399596,21.72774869,23.35989238,23.53080911,23.63279646,24.04738243,26.12949411,28.03540504,30.5309703,31.35535706,31.37158484,37.22989083,37.51347239,37.85196735,38.67532035,41.09033148,42.93076791,45.80763375,46.66192941,46.92620051,48.10554804,48.43404538,48.57527021,15.66,16.28,17.55,18.14,18.75,19.08,19.85,20.66,21.27,21.76,21.82,21.93,22.77,23.33,23.54,24.57,25.47,25.57,26.03,27.1,27.53,27.75,28.81]) / 100.0
Pi_exp = array([35.62049902,28.34058135,364.721004,1374.829083,1374.829083,2736.954254,4631.80577,7102.441356,7102.441356,10185.19755,18311.64992,29233.42371,29233.42371,29233.42371,40113.01829,40113.01829,51271.0941,51271.0941,43142.2229,51271.0941,69978.47398,69978.47398,104445.7197,92078.31285,104445.7197,163099.4158,131902.2819,327838.3478,327838.3478,217191.7064,632653.7888,217191.7064,562280.8036,632653.7888,465721.4862,465721.4862,465721.4862,18311.64992,20772.12998,23410.58496,26230.03855,29233.42371,29233.42371,32423.59143,35803.31832,39375.3132,43142.2229,43142.2229,43142.2229,47106.63736,47106.63736,51271.0941,55638.08226,60210.04621,60210.04621,64989.38873,69978.47398,80595.15144,75179.63005,80595.15144])


# Set physical parameters for the simulation. You can modify these values.  
a       = 8.0e-9                                 # sphere radius in meters 
I       = 5.0e-3                                 # Ionic strength in mol/L
sig     = 0.5                                    # Bare surface charge density in e/nm2
phi_tab = linspace(0.02,0.40,10)                 # Volume fractions for the simulation. The arguments of linspace 
                                                 # are (initial value, final value, number of points)

# Physical constants following from input parameters. You should not modify these lines. 
kappa = sqrt(8*pi*lB*I*1000*Na)                  # Inverse screening length (1/m)
lD    = 1./kappa                                 # Debye length in meters
ka    = kappa*a                                  # Scaled radius
Z     = sig * 4 * pi * (a*1.0e9)**2              # Colloid number of charges
n0    = kappa**2/(8*pi*lB)                       # Ion density in the reservoir (m-3)


# Define a Cell Model object. The parameters will be defined later. 
simCM = CellModelSimulation()


# Define an OZ numerical solver. The physical parameters are given from what's above. The numerical parameters can be changed here.       
sim           = OZSimulation()
# Physical parameters
sim.potential = "yukawa"                 # Cannot be changed here. 
sim.kappa     = kappa      
sim.a         = a     
sim.Z         = Z 
#sim.phi       = phi_tab[0]
# Numerical parameters... 
# ... iterative method for OZ equation 
sim.relax     = 0.05                     # <1. Relaxation factor of the Picard iteration in the iterative OZ solver. Large values make the resolution 
                                         # faster but less stable, and vice versa. Usually between 0.001 and 0.5. 
sim.tol       = 1.0e-11                  # Residual value used to consider the OZ convergence achived. Machine accuracy usually makes values 
                                         # below 1.0e-12 or 1.0e-13 impossible to reach. 
sim.Ng        = True                     # False/True. Activate Ng 5 acceleration. Much faster if it's True, but it can make things unstable. 
sim.res_start_Ng = 1                     # If Ng acceleration is activated, it will *not* start before the standard Picard iteration brings 
                                         # the residual below this value. Use a large value > 1 to start Ng acceleration immediately. 
                                         # Besides, when the residual is still large, the Ng estimation of the next g(r) is mixed slowly with the 
                                         # standard Picard estimation to help convergence. Here for residual > 0.01. This value can be changed at line 
                                         # 389 of this file in the Ng5 function. Use a smaller value for a slower/more stable code. A higher value 
                                         # can result in an unstable code, but will make things much faster (for better of for worse!). 
# ... Spatial discretization of g(r)                                          
sim.N         = 32768                    # Number of points in the discretization of g(r). This value is large and should be sufficient. 
sim.Ndmax     = 120.                     # Largest value of r. 
# ... closure for OZ equation 
sim.closure   = "HNC"                    # Choose HNC, PY or RY closure here. 
# ... for RY closure only                                         
sim.solver    = "RegulaFalsi"            # Numerical method used to determine the interpolation coefficient in the RY closure. 
                                         # If left empty, the bisection method is used. It is very safe but slow. Another choice 
                                         # is "RegulaFalsi", which is usually faster but occasionnally unstable.
sim.alphamin  = 1e7                      # Lower bound for the interpolation parameter in the RY closure 
sim.alphamax  = 2.5e8                    # Upper bound for the interpolation parameter in the RY closure
# ... convergence monitoring 
sim.quiet     = False                    # False/True. If True, convergence will not be reported for a faster resolution. (for HNC and PY)
sim.RYQuiet   = False                    # False/True. Same for RY closure. 
sim.log       = "X"                      # Gives residuals for the convergence of the OZ solver. If left empty, no convergence report, 
                                         # if "stdout" residuals are printed in the terminal, if "X" residuals are plotted as a function 
                                         # of iteration number in figure 0.
#


# Please, don't modify this block. It generates figure(0) for convergence monitoring and figure(1) for result monitoring. 
# In OZ convergence monitoring, the current simulation appears in red and the former ones in black.  
if sim.log=="X":
    convf = figure(0)
# Plot experimental data first 
debf = figure(1)
yscale('log')
xlim(0,0.4)
ylim(1000,200000)
plot(phi_exp,Pi_exp,'or',ms=8, label='Experiments')
xlabel('$\\phi$',fontsize=20)
ylabel('$\\Pi\\,(Pa)$',fontsize=20)
tight_layout()
figure(0)

# Prepare data structures for results
PiCMplot = []
PiCSplot = []
Pimodplot = []
Pi_perturb_all = []
Pi_OZ_all = []
Pi_RMSA_all = []
Pi_CSCM_all = []

print "         phi             Pi_OZ          Pi_perturb        Pi_RMSA    "

# Loop over all the volume fractions asked in the parameter section. 
i=-1
for phi in phi_tab:
  i+=1

  # Set the volume fraction is the OZ solver
  sim.phi = phi 

  # Perform renormalization based on the cell model and EPC ...
  # ... set cell model parameters with this command
  simCM.SetParameters(a=a*1.0e9,I=I,res=800,phi=phi,sigma=sig) # parameters are: colloid radius in nm, 
                                                               #ionic strength in mol/L, numer of points in the spatial 
                                                               #discretization of the cell model, volume fraction, surface charge density in e/nm2
  # ... this line directly calls the cell model resolution and returns the EPC renormalization parameters. The cell model osmotic pressure is 
  # also computed and stored in simCM.Pi
  sim.Z, sim.kappa = simCM.renormalize() # use the CM to define the effective charge and inverse screening length in the OZ solver (sim object). 

  # The volume term can be defined now
  volume_term = kT*sim.kappa**2/(8*pi*lB)*(1-(kappa/sim.kappa)**2)**2


  # RMSA solution: it is embeded in the OZ object
  sigma = 2*sim.a
  k = sim.kappa*sigma 
  gamma = sim.Z**2*lB*exp(k)/(1+sim.kappa*sim.a)**2/sigma
  # The resolution is performed with the next call. It returns the wave svector amplitude K, the structure factor S, the radial coordinate scaled by 
  # the particle radius x, the radial distribution function gRMSA, the pressure PRMSA, the direct correlation function cMSA. 
  # These quantities can be plotted for example with plot(x,gRMSA) or plot(K,S)
  K, S, x, gRMSA, PRMSA, cMSA = sim.RMSAHSYSolution(eta=phi,sigma=sigma,k=k,gamma=gamma) 
  # Store the pressure in a table                 
  Pi_RMSA_all.append(PRMSA + volume_term)


  # Numerical OZ solution with renormalized parameters. The closure was set above.
  # ... In case we use RY closure, the min/max values of the interpolation parameter alpha have been set to span a large range 
  # which should be ok for various physical  parameters. When changing the volume fraction slowly like in this loop, the suspension 
  # structure is not drastically changed compared to the one at the previous volume fraction, so that alpha should not be different from 
  # the previous one. In the next lines, we use an initial guess for the alpha range that is probably better than the standard range above. 
  # If alpha ends up being out of these bounds, the code will tell it so you can change these values. The next lines are expected to make 
  # RY faster, but if they cause problems, set alphamin and alphamax to -1 in order to keep the initial values for all volume fractions. 
  if (not phi==phi_tab[0]) and sim.closure=="RY":
    sim.alphamin=sim.alphamid/3 # or -1
    sim.alphamax=sim.alphamid*3 # or -1
  # Update physical and numerical parameters of OZ simulation object  
  sim.UpdateParameters()
  # Solve OZ equation numerically. It returns the rdf g, the wave vector k, the structure factor S, the indirect correlation function gam.  
  g, k, S, gam = sim.SolveOZ()  
  # here pressure computation has to be called explicitely. It is stored in sim.betaP
  sim.ComputePressure(sim.rho,sim.g,sim.beta_u)    
  P_OCM_OZ = sim.betaP*kT
  Pi_OZ_all.append(P_OCM_OZ + volume_term)

  # Perturbation theory. Numerical parameters may be changed if necessary in the code above. 
  betaPiperturb = sim.ComputePressurePerturb(sim.N)         # parameter = number of points in the spatial discretization of the integral 
  Pi_perturb_all.append(betaPiperturb*kT + volume_term)

  # Carnahan-Starling pressure 
  P_CS = sim.rho*kT*(1+phi+phi**2-phi**3)/(1-phi)**3

  # CSCM model is juste the summ of CS pressure and CM pressure stored in simCM object
  Pi_CSCM_all.append(P_CS + simCM.Pi)

  # Print osmotic pressure estimations in Pascals. 
  print format(phi,'16.4e'), format(Pi_OZ_all[-1],'16.4e'), format(Pi_perturb_all[-1],'16.4e'), format(Pi_RMSA_all[-1],'16.4e')

  # Plot the equation of state
  figure(1)
  plot(phi_tab[:i+1], Pi_OZ_all    , '-k', lw=3, label='$\Pi_{OZ}$')
  plot(phi_tab[:i+1], Pi_perturb_all        , '--+m', lw=3, ms=10, label='$\Pi_{perturb}$')
  plot(phi_tab[:i+1], Pi_RMSA_all        , '--xg', lw=3, ms=10, label='$\Pi_{RMSA}$')
  plot(phi_tab[:i+1], Pi_CSCM_all        , ':b', lw=3, ms=10, label='$\Pi_{CSCM}$')
  if phi==phi_tab[0]: legend(loc='lower right',fontsize=20)
  draw()
  figure(0)