Example205_LowLevelSpaceTimePoisson.jl

#=

# 205 : Space-Time FEM for Heat Equation
([source code](@__SOURCE_URL__))

This example computes the solution ``u`` of the
two-dimensional heat equation
```math
\begin{aligned}
u_t - \Delta u & = f \quad \text{in } \Omega
\end{aligned}
```
with a (possibly space- and time-depepdent)
right-hand side ``f`` and homogeneous Dirichlet boundary
and initial conditions
on the unit square domain ``\Omega`` on a given grid
with space-time finite element methods based on
tensorized ansatz functions.

![](example205.png)

=#

module Example205_LowLevelSpaceTimePoisson

using ExtendableFEMBase
using ExtendableGrids
using ExtendableSparse
using GridVisualize
using UnicodePlots
using Test #

## data for Poisson problem
const μ = (t) -> 1.0e-1 * t + 1 * max(0, (1 - 2 * t))
const f = (x, t) -> sin(3 * pi * x[1]) * 4 * t - cos(3 * pi * x[2]) * 4 * (1 - t)

function main(; dt = 0.01, Tfinal = 1, level = 5, order = 1, Plotter = nothing, produce_movie = false)

    ## Finite element type
    FEType_time = H1Pk{1, 1, order}
    FEType_space = H1Pk{1, 2, order}

    ## time grid
    T = LinRange(0, Tfinal, round(Int, Tfinal / dt + 1))
    grid_time = simplexgrid(T)

    ## space grid
    X = LinRange(0, 1, 2^level + 1)
    grid_space = simplexgrid(X, X)

    ## FESpaces for time and space
    FES_time = FESpace{FEType_time}(grid_time)
    FES_space = FESpace{FEType_space}(grid_space)

    ## solve
    sol = solve_poisson_lowlevel(FES_time, FES_space, μ, f)

    ## visualize
    if produce_movie
        @info "Producing movie..."
        vis = GridVisualizer(Plotter = Plotter)
        movie(vis, file = "example205_video.mp4") do vis
            for tj in 2:length(T)
                t = T[tj]
                first = (tj - 1) * FES_space.ndofs + 1
                last = tj * FES_space.ndofs
                scalarplot!(vis, grid_space, view(sol, first:last), title = "t = $(Float16(t))")
                reveal(vis)
            end
        end
        return sol, vis
    else
        @info "Plotting at five times..."
        plot_timesteps = [2, round(Int, length(T) / 4 + 0.25), round(Int, length(T) / 2 + 0.5), round(Int, length(T) - length(T) / 4), FES_time.ndofs]
        plt = GridVisualizer(; Plotter = Plotter, layout = (1, length(plot_timesteps)), clear = true, resolution = (200 * length(plot_timesteps), 200))
        for tj in 1:length(plot_timesteps)
            t = plot_timesteps[tj]
            first = (t - 1) * FES_space.ndofs + 1
            last = t * FES_space.ndofs
            scalarplot!(plt[1, tj], grid_space, view(sol, first:last), title = "t = $(T[t])")
        end
        return sol, plt
    end

end


function solve_poisson_lowlevel(FES_time, FES_space, μ, f)
    ndofs_time = FES_time.ndofs
    ndofs_space = FES_space.ndofs
    ndofs_total = ndofs_time * ndofs_space
    sol = zeros(Float64, ndofs_total)

    A = ExtendableSparseMatrix{Float64, Int64}(ndofs_total, ndofs_total)
    b = zeros(Float64, ndofs_total)

    println("Assembling...")
    time_assembly = @elapsed @time begin
        loop_allocations = assemble!(A, b, FES_time, FES_space, f, μ)

        ## fix homogeneous boundary dofs
        bdofs = boundarydofs(FES_space)
        for sdof in bdofs
            for dof_t in 1:ndofs_time
                dof = (dof_t - 1) * ndofs_space + sdof
                A[dof, dof] = 1.0e60
                b[dof] = 0
            end
        end

        ## fix initial value by zero
        for j in 1:ndofs_space
            A[j, j] = 1.0e60
            b[j] = 0
        end
        ExtendableSparse.flush!(A)
    end

    @info ".... spy plot of system matrix:\n$(UnicodePlots.spy(sparse(A.cscmatrix)))"

    ## solve
    println("Solving linear system...")
    time_solve = @elapsed @time copyto!(sol, A \ b)

    ## compute linear residual
    @show norm(A * sol - b)

    return sol
end

function assemble!(A::ExtendableSparseMatrix, b::Vector, FES_time, FES_space, f, μ = 1)

    ## get space and time grids
    grid_time = FES_time.xgrid
    grid_space = FES_space.xgrid

    ## get number of degrees of freedom
    ndofs_time = FES_time.ndofs
    ndofs_space = FES_space.ndofs

    ## get local to global maps
    EG_time = grid_time[UniqueCellGeometries][1]
    EG_space = grid_space[UniqueCellGeometries][1]
    L2G_time = L2GTransformer(EG_time, grid_time, ON_CELLS)
    L2G_space = L2GTransformer(EG_space, grid_space, ON_CELLS)

    ## get finite element types
    FEType_time = eltype(FES_time)
    FEType_space = eltype(FES_space)

    ## quadrature formula in space
    qf_space = QuadratureRule{Float64, EG_space}(2 * (get_polynomialorder(FEType_space, EG_space) - 1))
    weights_space::Vector{Float64} = qf_space.w
    xref_space::Vector{Vector{Float64}} = qf_space.xref
    nweights_space::Int = length(weights_space)
    cellvolumes_space = grid_space[CellVolumes]

    ## quadrature formula in time
    qf_time = QuadratureRule{Float64, EG_time}(2 * (get_polynomialorder(FEType_time, EG_time) - 1))
    weights_time::Vector{Float64} = qf_time.w
    xref_time::Vector{Vector{Float64}} = qf_time.xref
    nweights_time::Int = length(weights_time)
    cellvolumes_time = grid_time[CellVolumes]

    ## FE basis evaluators and dofmap for space elements
    FEBasis_space_∇ = FEEvaluator(FES_space, Gradient, qf_space)
    ∇vals_space = FEBasis_space_∇.cvals
    FEBasis_space_id = FEEvaluator(FES_space, Identity, qf_space)
    idvals_space = FEBasis_space_id.cvals
    celldofs_space = FES_space[ExtendableFEMBase.CellDofs]

    ## FE basis evaluators and dofmap for time elements
    FEBasis_time_∇ = FEEvaluator(FES_time, Gradient, qf_time)
    ∇vals_time = FEBasis_time_∇.cvals
    FEBasis_time_id = FEEvaluator(FES_time, Identity, qf_time)
    idvals_time = FEBasis_time_id.cvals
    celldofs_time = FES_time[ExtendableFEMBase.CellDofs]

    ## ASSEMBLY LOOP
    loop_allocations = 0
    function barrier(EG_time, EG_space, L2G_time::L2GTransformer, L2G_space::L2GTransformer)
        ## barrier function to avoid allocations by type dispatch

        ndofs4cell_time::Int = get_ndofs(ON_CELLS, FEType_time, EG_time)
        ndofs4cell_space::Int = get_ndofs(ON_CELLS, FEType_space, EG_space)
        Aloc = zeros(Float64, ndofs4cell_space, ndofs4cell_space)
        Mloc = zeros(Float64, ndofs4cell_time, ndofs4cell_time)
        ncells_space::Int = num_cells(grid_space)
        ncells_time::Int = num_cells(grid_time)
        x::Vector{Float64} = zeros(Float64, 2)
        t::Vector{Float64} = zeros(Float64, 1)

        ## assemble Laplacian
        loop_allocations += @allocated for cell in 1:ncells_space
            ## update FE basis evaluators for space
            FEBasis_space_∇.citem[] = cell
            update_basis!(FEBasis_space_∇)

            ## assemble local stiffness matrix in space
            for j in 1:ndofs4cell_space, k in 1:ndofs4cell_space
                temp = 0
                for qp in 1:nweights_space
                    temp += weights_space[qp] * dot(view(∇vals_space, :, j, qp), view(∇vals_space, :, k, qp))
                end
                Aloc[j, k] = temp
            end
            Aloc .*= cellvolumes_space[cell]

            ## add local matrix to global matrix
            for time_cell in 1:ncells_time
                update_trafo!(L2G_time, time_cell)
                for jT in 1:ndofs4cell_time, kT in 1:ndofs4cell_time
                    dofTj = celldofs_time[jT, time_cell]
                    dofTk = celldofs_time[kT, time_cell]
                    for qpT in 1:nweights_time
                        ## evaluate time coordinate and μ
                        eval_trafo!(t, L2G_time, xref_time[qpT])
                        factor = μ(t[1]) * weights_time[qpT] * idvals_time[1, jT, qpT] * idvals_time[1, kT, qpT] * cellvolumes_time[time_cell]
                        for j in 1:ndofs4cell_space
                            dof_j = celldofs_space[j, cell] + (dofTj - 1) * ndofs_space
                            for k in 1:ndofs4cell_space
                                dof_k = celldofs_space[k, cell] + (dofTk - 1) * ndofs_space
                                if abs(Aloc[j, k]) > 1.0e-15
                                    ## write into sparse matrix, only lines with allocations
                                    rawupdateindex!(A, +, Aloc[j, k] * factor, dof_j, dof_k)
                                end
                            end
                        end
                    end
                end
            end
            fill!(Aloc, 0)

            ## assemble right-hand side
            update_trafo!(L2G_space, cell)
            for qp in 1:nweights_space
                ## evaluate coordinates of quadrature point in space
                eval_trafo!(x, L2G_space, xref_space[qp])
                for time_cell in 1:ncells_time
                    update_trafo!(L2G_time, time_cell)
                    for qpT in 1:nweights_time
                        ## evaluate time coordinate
                        eval_trafo!(t, L2G_time, xref_time[qpT])

                        ## evaluate right-hand side in x and t
                        fval = f(x, t[1])

                        ## multiply with test function and add to right-hand side
                        for j in 1:ndofs4cell_space
                            temp = weights_time[qpT] * weights_space[qp] * idvals_space[1, j, qp] * fval * cellvolumes_space[cell] * cellvolumes_time[time_cell]

                            ## write into global vector
                            for jT in 1:ndofs4cell_time
                                dof_j = celldofs_space[j, cell] + (celldofs_time[jT, time_cell] - 1) * ndofs_space
                                b[dof_j] += temp * idvals_time[1, jT, qpT]
                            end
                        end
                    end
                end
            end
        end

        ## assemble time derivative
        return loop_allocations += @allocated for time_cell in 1:ncells_time
            ## update FE basis evaluators for time derivative
            FEBasis_time_∇.citem[] = time_cell
            update_basis!(FEBasis_time_∇)

            ## assemble local convection term in time
            for j in 1:ndofs4cell_time, k in 1:ndofs4cell_time
                temp = 0
                for qpT in 1:nweights_time
                    temp += weights_time[qpT] * dot(view(∇vals_time, :, j, qpT), view(∇vals_time, :, k, qpT))
                end
                Mloc[j, k] = temp
            end
            Mloc .*= cellvolumes_time[time_cell]

            ## add local matrix to global matrix
            for cell in 1:ncells_space
                for jX in 1:ndofs4cell_space, kX in 1:ndofs4cell_space
                    dofXj = celldofs_space[jX, cell]
                    dofXk = celldofs_space[kX, cell]
                    for qpX in 1:nweights_space
                        factor = weights_space[qpX] * idvals_space[1, jX, qpX] * idvals_space[1, kX, qpX] * cellvolumes_space[cell]
                        for j in 1:ndofs4cell_time
                            dof_j = dofXj + (celldofs_time[j, time_cell] - 1) * ndofs_space
                            for k in 1:ndofs4cell_time
                                dof_k = dofXk + (celldofs_time[k, time_cell] - 1) * ndofs_space
                                if abs(Mloc[j, k]) > 1.0e-15
                                    ## write into sparse matrix, only lines with allocations
                                    rawupdateindex!(A, +, Mloc[j, k] * factor, dof_j, dof_k)
                                end
                            end
                        end
                    end
                end
            end
            fill!(Mloc, 0)
        end
    end
    barrier(EG_time, EG_space, L2G_time, L2G_space)
    flush!(A)
    return loop_allocations
end

function generateplots(dir = pwd(); Plotter = nothing, kwargs...)
    ~, plt = main(; Plotter = Plotter, kwargs...)
    scene = GridVisualize.reveal(plt)
    return GridVisualize.save(joinpath(dir, "example205.png"), scene; Plotter = Plotter)
end

end #module