# -*- coding: utf-8 -*-
"""
Created on Sat May 25 12:22:58 2024

@author: b.janotta
"""

"""
Supporing Information for:
Activity Coefficients for Lithium Salts in Aqueous and Nonaqueous Solvents and
in Solvent Mixtures

Authors: Andrew R. Crothers, Clayton J. Radke, and John M. Prausnitz*

*Corresponding author




File contains Python script of classes used to calculate the mean ionic 
activity coefficients for LiCl in mixture of water and methanol. This script 
is called by the file "mixed_solvent_activity_meoh_licl.py" to make 
calculations. The equations used here are provided completely in the main text.
"""



#additional packages
import scipy as sp   #SciPy python package
import numpy as np   #NumPy python package
import scipy.special
import scipy.optimize
#special functions
from numpy import exp, sqrt, pi, sinh, log
#charge of electron, Avagadro's number, vacuume dielectric constant
from scipy.constants import e, N_A, epsilon_0

#Faraday constant
F = sp.constants.physical_constants['Faraday constant'][0]
#gas constant
R = sp.constants.R
#Boltzmann constant
kb = sp.constants.k
#temeprature
T = 298.15        #K

class Activity_calculation:
    """ class to calculate activity coefficients"""
    
    
    def __init__(self, sigma, epsr_b):
        """initiate activity calculation
        sigma: distance of closest approach in m
        epsr_b: bulk solution dielectric constant"""
        
        self.epsr_b = epsr_b
        self.sigma = sigma
        #Bjerrum length over sigma
        self.b = e**2/(4*pi*epsr_b*epsilon_0*kb*T*self.sigma)

    def kappa(self, c):
        """returns Debye-Huckel screening parameter
        c: salt concentration in mol/m^3"""
        #Debye-Huckel screening parameter
        kappa = sqrt(N_A*e**2*2*c/(self.epsr_b*epsilon_0*kb*T))
        return kappa

    
    def gam(self, c, delta, ig):   #BiMSA parameter
        """returns BiMSA screening constant in 1/m
        c: salt concentration in mol/m^3
        delta: fraction dissociation
        ig: initial guess on gamma_B"""
        #Solving Equation 20
        solution = sp.optimize.root(lambda x: 4*(x**2)*(1 + x*self.sigma)**2 - 
                                    self.kappa(c)**2 * (delta + x*self.sigma)/
                                    (1 + x*self.sigma), ig)
        if solution.success==False:
            #check that solver was successful
            print('Problem solving gam!')
        return solution.x
    
    def delta(self, c, gam, ig):
        """returns fraction cations associated
        c: salt concentration in mol/m^3
        gam: BiMSA screening constant in 1/m
        ig: initial guess on fraction ions associated"""
        #second factor in Equation 22
        el_part = exp(self.b/(1 + gam*self.sigma)**2) 
        #Third factor in Equation 22
        hs_part = (1 + 0.5*self.eta(c))/(1 - self.eta(c))**3
        #solution to Equation 22 and 21
        solution = sp.optimize.root(lambda x: (1 - x)/x**2 - c* self.K()*exp(-self.b) * hs_part * el_part, ig)
        if solution.success==False:
            #check that solver was successful
            print('Problem solving delta!')
        return solution.x    
    
    def K(self):
        """Equilibrium constant at infinite dilution"""
        #Equation 23
        return 8*pi*N_A*self.sigma**3*np.sum([self.b**(2*m) / (sp.special.factorial(2*m)*(2*m - 3)) for m in np.arange(2,10,1)])
    
    def eta(self, c):
        """returns packing factor
        c: salt concentration in mol/m^3"""
        #equation 18
        return pi*self.sigma**3/6*(2*c*N_A)
    
    def log_y_hs(self, c):
        """returns hard-sphere contribution to activity coefficient
        c: salt concentration in mol/m^3"""
        eta = self.eta(c)
        #equation 19
        return (8*eta - 9*eta**2 + 3*eta**3)/(1 - eta)**3  #Carnahan-Starling
    
    def log_y_el(self, c, gam):
        """electrostatic contribution to the activity coefficient
        c: salt concentration in mol/m^3
        gam: BiMSA screening constant in 1/m"""
        #Equation 19
        return -self.b*gam*self.sigma/(1 + gam*self.sigma)
    
    def log_y_mal(self, c, delta):
        """ion pairing contribution to the activity coefficient
        c: salt concentration in mol/m^3
        delta: fraction of ions associated"""
        eta = self.eta(c)
        #Equation 24
        return log(delta) - 1/4*(1 - delta)*(5*eta - 2*eta**2) / ((1 - eta)*(1 - 0.5*eta))
    
    def log_y_tot(self, c, delta, gam):
        """sum of all activity coefficient contributions
        c: salt concentration in mol/m^3
        delta: fraction of ions associated
        gam: BiMSA screening constant in 1/m"""
        #Equation 16
        return self.log_y_hs(c) + self.log_y_el(c, gam) + self.log_y_mal(c, delta)

    def delta_gam_solve(self, c, igs):
        """solving Equation 21 and 20 at the same time
        return screening parameter and delta
        c: salt concentration in mol/m^3
        ig: tuple with initial guess on fraction ions associated and BiMSA
        screening parameter in 1/m"""
        def gam_delta(x): #x = [delta, gam]
            
            #second factor in Equation 22
            el_part = exp(self.b/(1 + x[1]*self.sigma)**2)
            #Third factor in Equation 22
            hs_part = (1 - 0.5*self.eta(c))/(1 - self.eta(c))**3
            
            #Equation 21
            equ1 = (1 - x[0])/x[0]**2 - c* self.K()*exp(-self.b) * hs_part * el_part
            #Equation 20
            equ2 = 4*(x[1]**2)*(1 + x[1]*self.sigma)**2 - self.kappa(c)**2 * (x[0] + x[1]*self.sigma)/(1 + x[1]*self.sigma)
            return [equ1/igs[0], equ2/igs[1]] #error on fits
        solution = sp.optimize.root(gam_delta, igs) #minimize error
        if solution.success==False: #error if cannot solve
            print('Problem at solving gam_delta!')
        return solution.x
    
class Solvent_Partitioning:
    """Equations from Zhuang and Wang 2016"""
    
    def __init__(self, deltas, mus, vs, epsrs, R_ion):
        """initiation function. takes in vectors of properties of each species
        and the ionic radius"""
        self.deltas = deltas  #F*m^2,  polarizability
        self.mus = mus        #m*C,   permenent dipole moment
        self.vs = vs          #1/m^3, molecular volume
        self.epsrs = epsrs    #relative permetivity
        self.R_ion = R_ion    #m, radius of central ion


    def massfrac_to_volfrac(self, wa, Ma, Mb, Va, Vb):
        """helper function to calculate volume fraction from mass fraction
        assuming ideal mixing"""
        return Mb*Va*wa/(Ma*Vb + Mb*Va*wa - Ma*Vb*wa)

    def phis_fun(self, Efield, u, phis_inf):
        """calculates fraction of solvent at surface
        Efield: electric field in V/m
        u: incompressibility field
        phis_inf: solvent volume fraction in bulk"""
        phis = phis_inf*exp(-u*self.vs/(kb*T) + self.deltas/(2*kb*T)*Efield**2)*(sinh(self.mus*Efield/(kb*T))/(self.mus*Efield/(kb*T)))
        return phis
       
    def calculations(self, phis_inf, u_ig=0):
        """numerical solver
        phis_inf: solvent volume fraction in bulk
        u_ig: initial guess for the incompressibility field"""
        #electric field at cation surface Equation I2
        Efield_ig = e/(4*pi*self.R_ion**2*epsilon_0*np.sum(self.epsrs*phis_inf))
        
        #numerical solver
        def Efield_solve(Efield):
            #solve Equation I3
            u = sp.optimize.root(lambda u_: -1 + np.sum(self.phis_fun(Efield, u_, phis_inf)), u_ig).x
            phis = self.phis_fun(Efield, u, phis_inf)
            epsr = np.sum(phis*self.epsrs)
            #return error on Equation I2
            return epsr*Efield - e/(4*pi*self.R_ion**2*epsilon_0)
        #solve Equation I2
        Efield = sp.optimize.root(Efield_solve, Efield_ig).x 
        #find incompressibility field
        u = sp.optimize.root(lambda u_: -1 + np.sum(self.phis_fun(Efield, u_, phis_inf)), u_ig).x
        #calculate local solvent volume fraction
        phis_local = self.phis_fun(Efield, u, phis_inf)
        return Efield, phis_local, u