FiPy solving PDE on 1D cylindrical coordinates - Stackoverflow question

conv-diff-1D

import fipy as fp  ## finite volume PDE solver
from fipy.tools import numerix  ## requirement for FiPy, in practice same as numpy
import copy  ## we need the deepcopy() function because some FiPy objects are mutable
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import math

## numeric implementation of Dirac delta function as mentioned in the equation above
def delta_func(x, epsilon, coeff):
    return ((x < epsilon) & (x > -epsilon)) * \
        (coeff * (1 + numerix.cos(numerix.pi * x / epsilon)) / (2 * epsilon))

## Logarithmic data plotting with logaritmic time scale (base=10)
## plotting function requires a numpy.array with dimensions nt*nr*2, 
## ticks sets the number of ticks on the colorbar (recommendended<=10)
## levels sets the number of color levels on the contour plot
## logdiff sets the interesting data interval width from the maximum down by 10^logdiff
def plot_solution(solution_array,ticks=10,levels=300,logdiff=10,figsize=(8,7)):
    sol_min = solution_array[:,:,1].min()
    sol_max = solution_array[:,:,1].max()
    logmax = math.ceil(np.log10(sol_max))
    logmin = logmax - logdiff
    numofticks = ticks
    div = logdiff // numofticks
    power = np.arange((logmax - (numofticks * div)), logmax, div)
    array1 = np.zeros(len(power)) + 10.
    ticks1 = np.power(array1, power)
    levels1 = np.logspace(logmin,logmax,levels, base=10.0)
    norm = matplotlib.colors.LogNorm(vmin=(10 ** logmin),vmax=sol_max)
    formatter = matplotlib.ticker.LogFormatter(10, labelOnlyBase=False)
    fig = plt.figure(figsize=figsize)
    plot = plt.contourf(solution_array[0,:,0], np.linspace(dt,duration,nt), solution_array[:,:,1],
                        levels=levels1, norm=norm, cmap=plt.cm.jet)
    axes = plt.gca()
    axes.set_facecolor((0.,0.,0.25)) ## dark blue will display where the values are too small
    axes.set_yscale('log')
    cbar = plt.colorbar(ticks=ticks1, format=formatter)
    cbar.set_label(r'Density (1/m$^3$)')
    plt.xlabel('radius (m)')
    plt.ylabel(r'log$_{10}$[time (s)]')
    plt.title(r'Density of runaway electrons (1/m$^3$)')
    return

#####################################################################

## diffCoeff=100, convCoeff=1000

R_from = 0.7  ## inner radius in meters
R_to = 1.  ## outer radius in meters
nr = 1000  ## number of mesh cells
dr = (R_to - R_from) / nr  ## distance between the centers of the mesh cells
duration = 0.001  ## length of examined time evolution in seconds
nt = 1000  ## number of timesteps
dt = duration / nt  ## length of one timestep

## 3D array for storing the density with the correspondant radius values
## the density values corresponding to the n-th timestep will be in the n-th line 
solution = np.zeros((nt,nr,2))
## loading the radial coordinates into the array
for j in range(nr):
    solution[:,j,0] = (j * dr) + (dr / 2) + R_from

mesh = fp.CylindricalGrid1D(dx=dr, nx=nr)  ## 1D mesh based on the radial coordinates 
mesh = mesh + (0.7,)  ## translation of the mesh to r=0.7
n = fp.CellVariable(mesh=mesh)  ## fipy.CellVariable for the density solution in each timestep

diracLoc = 0.8  ## location of the middle of the Dirac delta
diracCoeff = 1.  ## Dirac delta coefficient ("height")
diracPercentage = 2  ## width of Dirac delta (full width from 0 to 0) in percentage of full examined radius
diracWidth = int((nr / 100) * diracPercentage)

## diffusion coefficient
diffCoeff = fp.CellVariable(mesh=mesh, value=100.)
## convection coefficient - must be a vector
# cC = (numerix.sin(mesh.x)) * 2
convCoeff = fp.CellVariable(mesh=mesh, value=(1000.,))

## boundary conditions
gradLeft = (0.,)  ## density gradient (at the "left side of the radius") - must be a vector
valueRight = 0.  ## density value (at the "right end of the radius")
n.faceGrad.constrain(gradLeft, where=mesh.facesLeft)  ## applying Neumann boundary condition
n.constrain(valueRight, mesh.facesRight)  ## applying Dirichlet boundary condition
convCoeff.setValue(0, where=mesh.x<(R_from + dr))  ## convection coefficient 0 at the inner edge

## applying initial conditions
n.setValue(delta_func(mesh.x - diracLoc, diracWidth * dr, diracCoeff))
## the PDE
eq = (fp.TransientTerm() == fp.DiffusionTerm(coeff=diffCoeff)
      - fp.ConvectionTerm(coeff=convCoeff))

## Solving the PDE and storing the data
for i in range(nt):
    eq.solve(var=n, dt=dt)
    solution[i,0:nr,1]=copy.deepcopy(n.value)

## Plotting the results
plot_solution(solution)