"""
Supplementary material to:
Grain boundary diffusion of chromium in polycrystalline nickel studied by SIMS
Thomas Gheno, François Jomard, Clara Desgranges and Laure Martinelli
Materialia, 2019

This script determines the triple product for grain boundary diffusion,
P=s*delta*Dgb,
where s is the equilibrium segregation coefficient, delta the grain boundary
width and Dgb the grain boundary diffusion coefficient, from an experimental
concentration (or intensity) profile (user input provided as .txt file). The
evaluation is based on the general solution of the bicrystal grain boundary
diffusion problem. It requires volume diffusion parameters as user input.

Script written for Python 3.6 (for Python 2.7 just change the print 
statement in line 273)
"""

# =============================================================================
# Import useful functions

import numpy as np
from numpy import exp, sqrt, pi, log, vectorize, diff, mean, linspace
from scipy.special import erfc
from scipy.integrate import quad
from scipy.optimize import curve_fit, root

# Documentation for numpy and scipy functions can be found at:
# https://docs.scipy.org/doc/numpy/reference/
# https://docs.scipy.org/doc/scipy/reference/

# =============================================================================
# Define constants

R = 8.314 # ideal gas constant in J mol-1 K-1
delta = 0.5*1e-7 # grain boundary width in cm

# =============================================================================
# Define auxiliary functions

def Arrhenius(k0, Q, TC):
    """
    Calculates variable at given temperature from Arrhenius relationship.
    Input: pre-exponential factor k0, activation energy Q (kJ/mol) and
    temperature TC (Celsius).
    Output: quantity k in units of k0.
    """
    TK = TC + 273.15
    return k0*exp(-Q*1000/(R*TK))

def linear_model(x, a, b):
    """
    Model to be used in least-squares fitting of the diffusion tail.
    """
    return(a*x + b)
    
def linear_fit(z, I, z_min, z_max):
    """
    Performs linear least-squares fit of data in the form (z, ln(I)) within
    user-selected range.
    Input: depth z (1D array), intensity I (1D array), user-specified
    boundaries for the fit z_min and z_max (scalars)
    Output: adjusted slope in units of 1/z, scalar.
    """
    # define boolean array of index and selected data range
    idx = (z_min < z) & (z < z_max)
    z_sel = z[idx]
    I_sel = I[idx]
    
    # define x and y data to be fitted
    x = z_sel
    y = log(I_sel)
    
    # perform least-squares fit and return adjusted slope
    para, cov = curve_fit(linear_model, x, y)
        
    return para[0]

# =============================================================================
# Define general solution of the grain boundary diffusion problem,
# with constant concentration boundary conditions

def beta(t, Dv, P):
    """
    Input: time t, volume diffusion coefficient Dv, triple product P.
    Output: dimensionless parameter beta, scalar.
    """
    return 0.5*P/(Dv*sqrt(Dv*t))

def eta(z, t, Dv):
    """
    Input: depth z, time t, volume diffusion coefficient Dv.
    Output: reduced variable for depth (dimensionless), 1D array of same size
    as z.
    """
    return z/sqrt(Dv*t)
    
def Delta(Dv, P):
    """
    Input: volume diffusion coefficient Dv, triple product P.
    Output: dimensionless parameter Delta, scalar.
    """
    return P/(delta*Dv)

def C1(z, t, Dv):
    """
    Contribution from direct diffusion from the surface.
    Output: concentration (dimensionless), 1D array of same size as z.
    """
    return erfc(eta(z, t, Dv)/2)

def C2_scalar(z, t, Dv, P, L):
    """
    Contribution from lateral diffusion from the boundary into the grains.
    Input: L is the system size in the direction normal to the grain boundary
    plane.
    Note that the integration by numpy.quad cannot be broadcast across the z
    array; instead, z is provided as a scalar, quad returns a scalar, so does
    C2_scalar, and C2_scalar is later applied to the z array using
    numpy.vectorize.
    Output: concentration (dimensionless), scalar.
    """
    
    beta_ = beta(t, Dv, P)
    eta_ = eta(z, t, Dv)
    Delta_ = Delta(Dv, P)
    
    def integrand_C2(s):
        """
        auxiliary function to be integrated in C2.
        Input: s is an integration variable.
        Output: scalar.
        """
        Y = (s - 1)/(2*beta_)*sqrt((Delta_ - 1)/(Delta_ - s))
        return s**(-3/2)*exp(-eta_**2/(4*s))*sqrt((Delta_ - s)/(Delta_ - 1))*(exp(-Y**2)/sqrt(pi) - Y*erfc(Y))
    
    integral = quad(integrand_C2, 1, Delta_)[0]
    return 2*sqrt(Dv*t)*eta_/(sqrt(pi)*L)*integral

# apply C2_scalar to z array (returns 1D array of same size as z)
C2 = vectorize(C2_scalar)

def Cgb_scalar(z, t, Dv, P, L):
    """
    Contribution from tracer atoms within the grain boundary.
    Input: here, z is a scalar. The function is later applied to z arrays using
    vectorize.
    Output: concentration (dimensionless), scalar.
    """
    
    beta_ = beta(t, Dv, P)
    eta_ = eta(z, t, Dv)
    Delta_ = Delta(Dv, P)
    
    def integrand_gb(s):
        """
        auxiliary function to be integrated in Cgb.
        Input: s is an integration variable.
        Output: scalar.
        """
        Y = (s - 1)/(2*beta_)*sqrt((Delta_ - 1)/(Delta_ - s))
        return s**(-3/2)*exp(-eta_**2/(4*s))*erfc(Y)
    
    integral = quad(integrand_gb, 1, Delta_)[0]
    return delta*eta_/(2*sqrt(pi)*L)*integral

# apply Cgb_scalar to z array (returns 1D array of same size as z)
Cgb = vectorize(Cgb_scalar)

def C(z, t, Dv, P, L, C0):
    """
    Calculates concentration profile resulting from all contributions.
    Output: concentration in units of C0, 1D array of same size as z. 
    """
    return C0*(C1(z, t, Dv) + C2(z, t, Dv, P, L) + Cgb(z, t, Dv, P, L))

# =============================================================================
# Calculate triple product from fitted slope

def A1(t, Dv, P, L, C0, z_range):
    """
    Calculates A1 gradient in given z range.
    Output: 1D array of same size as z_range.
    """
    # initialize output array and compute gradient
    res = np.zeros(z_range.size)
    res[:-1] = diff(log(C(z_range, t, Dv, P, L, C0)))/diff(z_range)
    # diff returns arrays of length (z_range.size - 1). The end value is
    # repeated to match the size of z_range
    res[-1] = res[-2]
    return res

def P_calc(t, Dv, L, C0, z_min, z_max, n_points, P_init, slope):
    """
    Calculates triple product P.
    P is adjusted so that the average A1 gradient equals the measured slope.
    An initial estimate P_init is provided by the user.
    """
    # define z range with size n_points
    z_range = linspace(z_min, z_max, n_points) 
    
    def func(P):
        """
        Returns the difference between the computed and measured gradients as a
        function of P.
        """
        A1_mean = mean(A1(t, Dv, P, L, C0, z_range))
        return A1_mean - slope

    return root(func, P_init)['x'][0]

# =============================================================================
# Experimental data and conditions
# These are parameters to be entered by the user. An experimental profile and
# associated parameters from the paper are provided as an example. 

# Experimental profile provided as column vectors of depth (nm) and normalized
# intensity in attached text file
exp_profile = np.genfromtxt('experimental_profile.txt')
z_exp = exp_profile[:,0]*1e-7 # depth in cm, 1D array
I_exp = exp_profile[:,1] # normalized intensity, 1D array

t = 100800 # time in s
T = 563.1 # temperature in Celsius

# Arrhenius parameters for volume diffusion of Cr in Ni, from [T. Gheno et al.,
# Tracer diffusion of Cr in Ni and Ni-22Cr studied by SIMS, Materialia 2018]

D0_v = 0.2 # pre-exponential factor in cm2/s
Q_v = 260 # activation energy in kJ

# Arrhenius parameters for grain boundary diffusion of Cr in Ni (used as
# initial estimate in optimization), from [J. Čermák, Grain boundary self-
# diffusion of 51Cr and 59Fe in austenitic Ni-Fe-Cr alloys, Mater. Sci. Eng.
# A148 (1991) 279-287]

P0_gb = 1e-6 # pre-exponential factor in cm3/s
Q_gb = 190 # activation energy in kJ

# Values calculated at the temperature of interest
Dv = Arrhenius(D0_v, Q_v, T) # volume diffusion coefficient in cm2/s
P_init = Arrhenius(P0_gb, Q_gb, T) # estimate of triple product in cm3/s

# Width of analysis area in cm (the value doesn't affect the calculated
# diffusivity, so long as it is much larger than the diffusion distance,
# L >> sqrt(Dv*t))
L = 30*1e-4

# Depth range used in the linear fit of the diffusion tail, to be determined
# from visual examination of the experimental profile
z_min = 3e-5 # cm
z_max = 7e-5 # cm

# Surface concentration in arbitrary units. Since the optimization of P
# is based on the slope of the concentration profile, the value of C0 does not
# affect the result.
C0 = 1

# Number of points of the z array over which the average A1 gradient is
# computed. Should be determined based on the desired precision and the shape
# of the experimental profile. The optimized P generally converges rapidly
# when increasing n_points.
n_points = 10

# =============================================================================
# Determine slope of diffusion tail and the corresponding triple product

slope = linear_fit(z_exp, I_exp, z_min, z_max)

P_adjusted = P_calc(t, Dv, L, C0, z_min, z_max, n_points, P_init*0.1, slope)
# see note below regarding the initial estimate
print(f'P = {P_adjusted:.2e} cm3/s') # for Python 3.X
# print 'P = {0:.2e} cm3/s'.format(P_adjusted) # for Python 2.X

# Note: sometimes the initial P estimate is too far from the optimal P value,
# and the optimization fails. In this case, P_calc should be run again with a
# larger or smaller initial estimate, as was done above (P_init*0.1 was used
# instead of P_init). In most cases increasing or decreasing P_init by a factor
# of 10 is sufficient. The try-and-error procedure can be implemented as
# follows:
#
#try:
#    P_adjusted = P_calc(t, Dv, L, C0, z_min, z_max, n_points, P_init, slope)
#except TypeError:
#    try:
#        P_adjusted = P_calc(t, Dv, L, C0, z_min, z_max, n_points, P_init*10, slope)
#    except TypeError:
#        try:
#            P_adjusted = P_calc(t, Dv, L, C0, z_min, z_max, n_points, P_init*0.1, slope)
#        except TypeError:
#            pass
#            print('Optimization failed, try with different P_init')