from IPython.display import clear_output
import time

from fipy import *

nx = 20
ny = nx
dx = 1.
dy = dx
L = dx * nx
mesh = Grid2D(dx=dx, dy=dy, nx=nx, ny=ny)

phi = CellVariable(name = "solution variable",
                   mesh = mesh,
                   value = 0.)

D = 1.
eq = TransientTerm() == DiffusionTerm(coeff=D)

valueTopLeft = 0
valueBottomRight = 1

X, Y = mesh.faceCenters
facesTopLeft = ((mesh.facesLeft & (Y > L / 2))
                | (mesh.facesTop & (X < L / 2)))
facesBottomRight = ((mesh.facesRight & (Y < L / 2))
                    | (mesh.facesBottom & (X > L / 2)))

phi.constrain(valueTopLeft, facesTopLeft)
phi.constrain(valueBottomRight, facesBottomRight)

viewer = Viewer(vars=phi, datamin=0., datamax=1.)

timeStepDuration = 10 * 0.9 * dx**2 / (2 * D)
steps = 10
for step in range(steps):
    eq.solve(var=phi,
             dt=timeStepDuration)
    viewer.axes.set_title("Solution variable (Step %d)" % (step + 1,))
    viewer.plot()
    time.sleep(1)
    clear_output()
    display(viewer.axes.get_figure())

DiffusionTerm().solve(var=phi)
viewer.plot()
viewer.axes.set_title("Solution variable (steady state)")
display(viewer.axes.get_figure())