Heat Equation

Introduction

The heat equation is the "Hello, world!" equation of finite elements. Here we solve the equation on a unit square, with a uniform internal source. The strong form of the (linear) heat equation is given by

$$ -\nabla \cdot (k \nabla u) = f \quad x \in \Omega, $$

where $u$ is the unknown temperature field, $k$ the heat conductivity, $f$ the heat source and $\Omega$ the domain. For simplicity we set $f = 1$ and $k = 1$. We will consider homogeneous Dirichlet boundary conditions such that $$ u(x) = 0 \quad x \in \partial \Omega, $$ where $\partial \Omega$ denotes the boundary of $\Omega$.

The resulting weak form is given by $$ \int_{\Omega} \nabla v \cdot \nabla u \ d\Omega = \int_{\Omega} v \ d\Omega, $$ where $v$ is a suitable test function.

Commented Program

Now we solve the problem in JuAFEM. What follows is a program spliced with comments.

First we load JuAFEM, and some other packages we need

In [1]:
using JuAFEM, SparseArrays

We start generating a simple grid with 20x20 quadrilateral elements using generate_grid. The generator defaults to the unit square, so we don't need to specify the corners of the domain.

In [2]:
grid = generate_grid(Quadrilateral, (20, 20));

Trial and test functions

A CellValues facilitates the process of evaluating values and gradients of test and trial functions (among other things). Since the problem is a scalar problem we will use a CellScalarValues object. To define this we need to specify an interpolation space for the shape functions. We use Lagrange functions (both for interpolating the function and the geometry) based on the reference "cube". We also define a quadrature rule based on the same reference cube. We combine the interpolation and the quadrature rule to a CellScalarValues object.

In [3]:
dim = 2
ip = Lagrange{dim, RefCube, 1}()
qr = QuadratureRule{dim, RefCube}(2)
cellvalues = CellScalarValues(qr, ip);

Degrees of freedom

Next we need to define a DofHandler, which will take care of numbering and distribution of degrees of freedom for our approximated fields. We create the DofHandler and then add a single field called u. Lastly we close! the DofHandler, it is now that the dofs are distributed for all the elements.

In [4]:
dh = DofHandler(grid)
push!(dh, :u, 1)
close!(dh);

Now that we have distributed all our dofs we can create our tangent matrix, using create_sparsity_pattern. This function returns a sparse matrix with the correct elements stored.

In [5]:
K = create_sparsity_pattern(dh);

We can inspect the pattern using the spy function from UnicodePlots.jl. By default the stored values are set to $0$, so we first need to fill the stored values, e.g. K.nzval with something meaningful.

In [6]:
using UnicodePlots
fill!(K.nzval, 1.0)
spy(K; height = 15)
Out[6]:
              Sparsity Pattern
       ┌──────────────────────────────┐    
     1⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ > 0
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ < 0
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
      ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
   441⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀    
       └──────────────────────────────┘    
       1                            441
                  nz = 3721

Boundary conditions

In JuAFEM constraints like Dirichlet boundary conditions are handled by a ConstraintHandler.

In [7]:
ch = ConstraintHandler(dh);

Next we need to add constraints to ch. For this problem we define homogeneous Dirichlet boundary conditions on the whole boundary, i.e. the union of all the face sets on the boundary.

In [8]:
∂Ω = union(getfaceset.((grid, ), ["left", "right", "top", "bottom"])...);

Now we are set up to define our constraint. We specify which field the condition is for, and our combined face set ∂Ω. The last argument is a function which takes the spatial coordinate $x$ and the current time $t$ and returns the prescribed value. In this case it is trivial -- no matter what $x$ and $t$ we return $0$. When we have specified our constraint we add! it to ch.

In [9]:
dbc = Dirichlet(:u, ∂Ω, (x, t) -> 0)
add!(ch, dbc);

We also need to close! and update! our boundary conditions. When we call close! the dofs which will be constrained by the boundary conditions are calculated and stored in our ch object. Since the boundary conditions are, in this case, independent of time we can update! them directly with e.g. $t = 0$.

In [10]:
close!(ch)
update!(ch, 0.0);

Assembling the linear system

Now we have all the pieces needed to assemble the linear system, $K u = f$. We define a function, doassemble to do the assembly, which takes our cellvalues, the sparse matrix and our DofHandler as input arguments. The function returns the assembled stiffness matrix, and the force vector.

In [11]:
function doassemble(cellvalues::CellScalarValues{dim}, K::SparseMatrixCSC, dh::DofHandler) where {dim}
    # We allocate the element stiffness matrix and element force vector
    # just once before looping over all the cells instead of allocating
    # them every time in the loop.
    n_basefuncs = getnbasefunctions(cellvalues)
    Ke = zeros(n_basefuncs, n_basefuncs)
    fe = zeros(n_basefuncs)
    # Next we define the global force vector `f` and use that and
    # the stiffness matrix `K` and create an assembler. The assembler
    # is just a thin wrapper around `f` and `K` and some extra storage
    # to make the assembling faster.
    f = zeros(ndofs(dh))
    assembler = start_assemble(K, f)
    # It is now time to loop over all the cells in our grid. We do this by iterating
    # over a `CellIterator`. The iterator caches some useful things for us, for example
    # the nodal coordinates for the cell, and the local degrees of freedom.
    @inbounds for cell in CellIterator(dh)
        # Always remember to reset the element stiffness matrix and
        # force vector since we reuse them for all elements.
        fill!(Ke, 0)
        fill!(fe, 0)
        # For each cell we also need to reinitialize the cached values in `cellvalues`.
        reinit!(cellvalues, cell)
        # It is now time to loop over all the quadrature points in the cell and
        # assemble the contribution to `Ke` and `fe`. The integration weight
        # can be queried from `cellvalues` by `getdetJdV`.
        for q_point in 1:getnquadpoints(cellvalues)
             = getdetJdV(cellvalues, q_point)
            # For each quadrature point we loop over all the (local) shape functions.
            # We need the value and gradient of the testfunction `v` and also the gradient
            # of the trial function `u`. We get all of these from `cellvalues`.
            for i in 1:n_basefuncs
                v  = shape_value(cellvalues, q_point, i)
                ∇v = shape_gradient(cellvalues, q_point, i)
                fe[i] += v * 
                for j in 1:n_basefuncs
                    ∇u = shape_gradient(cellvalues, q_point, j)
                    Ke[i, j] += (∇v  ∇u) * 
                end
            end
        end
        # The last step in the element loop is to assemble `Ke` and `fe`
        # into the global `K` and `f` with `assemble!`.
        assemble!(assembler, celldofs(cell), fe, Ke)
    end
    return K, f
end
Out[11]:
doassemble (generic function with 1 method)

Solution of the system

The last step is to solve the system. First we call doassemble to obtain the global stiffness matrix K and force vector f.

In [12]:
K, f = doassemble(cellvalues, K, dh);

To account for the boundary conditions we use the apply! function. This modifies elements in K and f respectively, such that we can get the correct solution vector u by using \.

In [13]:
apply!(K, f, ch)
u = K \ f;

Exporting to VTK

To visualize the result we export the grid and our field u to a VTK-file, which can be viewed in e.g. ParaView.

In [14]:
vtk_grid("heat_equation", dh) do vtk
    vtk_point_data(vtk, dh, u)
end
Out[14]:
1-element Array{String,1}:
 "heat_equation.vtu"

This notebook was generated using Literate.jl.