#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


from dolfin import *


# In[3]:


# Create mesh and define function space
mesh = UnitSquareMesh(8, 8)
V = FunctionSpace(mesh, "Lagrange", 1)


# In[4]:


# Define Dirichlet boundary (x = 0 or x = 1)
def boundary(x):
    return x[0] < DOLFIN_EPS or x[0] > 1.0 - DOLFIN_EPS


# In[5]:


from dolfin.common import plotting


# In[6]:


# Define boundary condition
u0 = Constant(0.0)
bc = DirichletBC(V, u0, boundary)


# In[7]:


# Define variational problem
u = TrialFunction(V)
v = TestFunction(V)
f = Expression("10*exp(-(pow(x[0] - 0.5, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2)
g = Expression("sin(5*x[0])", degree=2)
a = inner(grad(u), grad(v))*dx
L = f*v*dx + g*v*ds


# In[8]:


# Compute solution
w = Function(V)
solve(a == L, w, bc)


# In[9]:


import matplotlib.pyplot as plt
fig = plt.figure()
plot(w)


# In[10]:


_.collections


# In[11]:


from matplotlib import animation

def animate(i):
    print("frame %i" % i)
    f = Expression(f"10*exp(-(pow(x[0] - 0.05 * {i}, 2) + pow(x[1] - 0.5, 2)) / 0.02)", degree=2)
    L = f*v*dx + g*v*ds
    # Compute solution
    u = Function(V)
    solve(a == L, u, bc)
    plt.title("x-%.2f" % (.05 * i))
    return plot(u).collections


# In[12]:


# construct the animation
fig = plt.gcf()
anim = animation.FuncAnimation(fig, animate,
                               frames=20, interval=50)

# generate the HTML animation
anim.save("anim.html", writer="html")
plt.close(fig)

# display the resulting HTML
from IPython.display import HTML, display
with open("anim.html") as f:
    display(HTML(f.read()))