from fipy import *
import numpy as np

# Hey Daniel, this is the script I am using to perform piecewise
# linear interpolation of the diffusion as a function of
# concentration. I thought this was generic enough and uses
# current FiPy features. Is there a need to build the cases
# function you suggested? I comment below with extra details to
# explain what I do better.

# Molar concentration data points at which diffusion is known
molar=np.array([0.0,
0.902449063,
1.807061587,
2.713202143,
3.570465711,
4.425457264,
8.500000000]) 

# Diffusion coefficient at each point concentration in array molar

measDiff=np.array([1.052E-07,
3.636E-08,
2.394E-08,
3.661E-09,
2.391E-09,
2.045E-09,
0.000]) # assum high concentration limit diffusion is negligible

# Value of the concentration at the one of the boundaries
bcface=0.9*molar[-1]

# Obtains linear interpolation of diffusion coefficients as a
# function of concentration

# This computes the slobe between every two points
slopes=np.diff(measDiff)/np.diff(molar)
# This computes the intercepts between every two points
intercepts=slopes*-molar[:-1]+measDiff[:-1]

# Generate a mesh in which to solve the diffusion equation
# Dont worry about the actual numbers, just trying to show the principle
nx=100
dx=1./100./nx
dt=0.1
steps=1000
mesh = Grid1D(dx=dx,nx=nx)
sweeps = 5

# Generate a cell variable to represent concentration
phi = CellVariable(name="Concentration (M)",mesh=mesh,value=0.0,hasOld=1)

# Calculate the concentration at the phases by arithmetic averaging
# this will be used to calculate the diffusion
phifaces= phi.arithmeticFaceValue

# Generate a diffusion variable defined at the mesh faces
D=FaceVariable(mesh,value=1.052E-07)

# This is the key and I thought it worked. Please let me know if
# my logic is off. For every interval between two points I set D
# to depend on the slopes and intercepts calculated above. Note
# that we could use an arbitrary number of data points
for i in range(len(slopes)):
    D.setValue(phifaces*slopes[i]+intercepts[i],where=((
        phifaces>=molar[i]) & ( phifaces < molar[i+1] )))

# Generate the equation to solve
eq = TransientTerm() ==  DiffusionTerm(coeff=D)

# Set the boundary condition
phi.constrain(bcface,mesh.facesRight)

# Generate the viewer as in the examples
if __name__ == '__main__':
    vwrprof = Viewer(vars=phi)
    vwrprof.plot

# Iterate over the timesteps
for step in range(steps):
    phi.updateOld()
    for sweep in range(sweeps):
        res = eq.sweep(var=phi,dt=dt)

    # This is irrelevant to my point
    if __name__ == '__main__':
        if np.mod(step,10) == 0:
            vwrprof.plot(filename='concprof'+str(step).zfill(3)+'.png')
        else :
            vwrprof.plot()
