"""
simulates heat equation on rectangle returning a heat map at a number of times
boundary and initial conditions are 0, source represents burner on a stove
This program is based on the script FEniCS tutorial demo program: Diffusion of a Gaussian hill.     
  u'= Laplace(u) + f  in a square domain
  u = u_D = 0            on the boundary
  u = u_0 = 0            at t = 0
  u_D = f = stove burner flame
This program succesfully runs in the fenics docker image, see the book Solving PDEs in Python.
to animate: convert -delay 4 -loop 100 heatequation10*.png heatstovelinn.gif
to crop:convert heatstovelinn.gif -coalesce -repage 0x0 -crop 810x810+95+15 +repage heatstovelin.gif
"""

from fenics import *
import time
import matplotlib.pyplot as plt
from matplotlib import cm

# Create mesh and define function space
nx = ny = 100
mesh = RectangleMesh(Point(-2, -2), Point(2, 2), nx, ny)
V = FunctionSpace(mesh, 'P', 1)

# Define boundary, source, initial
def boundary(x, on_boundary):
    return on_boundary
bc = DirichletBC(V, Constant(0), boundary)
u_0 = interpolate(Constant(0), V)
f = Expression('exp(-sqrt(pow((a*pow(x[0], 2) + a*pow(x[1], 2)-a*1),2)))', degree=2, a=5) #steep guassian centred on the unit sphere

final_time = 0.035
num_pics = 72
for i in range(num_pics):
    T =   final_time*(i+1.0)/(num_pics+1)      #solve time even space
    #T = final_time*1.1**(i-num_pics+1)        #solve time log  space
    num_steps = 30
    dt = T / num_steps # time step size

    # Define variational problem
    u = TrialFunction(V)
    v = TestFunction(V)
    F = u*v*dx + dt*dot(grad(u), grad(v))*dx - (u_0 + dt*f)*v*dx
    a, L = lhs(F), rhs(F)
    
    # Time-stepping
    u = Function(V)
    t = 0
    for n in range(num_steps):
        t += dt              #step
        solve(a == L, u, bc) #solve
        u_0.assign(u)        #update
        
    #plot solution
    plot(u,cmap=cm.hot,vmin=0,vmax=0.07)
    plt.axis('off')
    plt.savefig('heatequation10%s.png'%(i+10),figsize=(8, 8), dpi=220,bbox_inches='tight', pad_inches=0,transparent=True)