Heat equation tutorial

In this tutorial, we'll be solving the heat equation:

$\frac{∂ ρcT}{∂ t} + ∇ ⋅ (-α ∇ρcT) = 0$

where

• t is time
• α is the thermal diffusivity
• T is the temperature
• ρ is the density
• c is the heat capacity
• ρcT is the thermal energy

To put this in the form of ClimateMachine's BalanceLaw, we'll re-write the equation as:

$\frac{∂ ρcT}{∂ t} + ∇ ⋅ (F(α, ρcT, t)) = 0$

where

• $F(α, ρcT, t) = -α ∇ρcT$ is the second-order flux

with boundary conditions

• Fixed temperature $T_{surface}$ at $z_{min}$ (non-zero Dirichlet)
• No thermal flux at $z_{min}$ (zero Neumann)

Solving these equations is broken down into the following steps:

1. Preliminary configuration
2. PDEs
3. Space discretization
4. Time discretization / solver
5. Solver hooks / callbacks
6. Solve
7. Post-processing

Preliminary configuration

Loading code

First, we'll load our pre-requisites:

• load external packages:

using MPI
using OrderedCollections
using Plots
using StaticArrays
using OrdinaryDiffEq
using DiffEqBase

• load CLIMAParameters and set up to use it:

using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()

Main.##347.EarthParameterSet()

• load necessary ClimateMachine modules:

using ClimateMachine
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.BalanceLaws:
    BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux

using ClimateMachine.Mesh.Geometry: LocalGeometry
using ClimateMachine.MPIStateArrays
using ClimateMachine.GenericCallbacks
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.SingleStackUtils

• import necessary ClimateMachine modules: (importing enables us to

provide implementations of these structs/methods)

import ClimateMachine.BalanceLaws:
    vars_state,
    source!,
    flux_second_order!,
    flux_first_order!,
    compute_gradient_argument!,
    compute_gradient_flux!,
    nodal_update_auxiliary_state!,
    nodal_init_state_auxiliary!,
    init_state_prognostic!,
    BoundaryCondition,
    boundary_conditions,
    boundary_state!

import ClimateMachine.DGMethods: calculate_dt

Initialization

Define the float type (Float64 or Float32)

const FT = Float64;

Initialize ClimateMachine for CPU.

ClimateMachine.init(; disable_gpu = true);

const clima_dir = dirname(dirname(pathof(ClimateMachine)));

[1638959910.112310] [hpc-92-37:26672:0]       ib_verbs.h:84   UCX  ERROR ibv_exp_query_device(mlx5_0) returned 38: No space left on device

Load some helper functions for plotting

include(joinpath(clima_dir, "docs", "plothelpers.jl"));

Define the set of Partial Differential Equations (PDEs)

Define the model

Model parameters can be stored in the particular BalanceLaw, in this case, a HeatModel:

Base.@kwdef struct HeatModel{FT, APS} <: BalanceLaw
    "Parameters"
    param_set::APS
    "Heat capacity"
    ρc::FT = 1
    "Thermal diffusivity"
    α::FT = 0.01
    "Initial conditions for temperature"
    initialT::FT = 295.15
    "Bottom boundary value for temperature (Dirichlet boundary conditions)"
    T_bottom::FT = 300.0
    "Top flux (α∇ρcT) at top boundary (Neumann boundary conditions)"
    flux_top::FT = 0.0
end

Create an instance of the HeatModel:

m = HeatModel{FT, typeof(param_set)}(; param_set = param_set);

This model dictates the flow control, using Dynamic Multiple Dispatch, for which kernels are executed.

Define the variables

All of the methods defined in this section were imported in Loading code to let us provide implementations for our HeatModel as they will be used by the solver.

Specify auxiliary variables for HeatModel

vars_state(::HeatModel, ::Auxiliary, FT) = @vars(z::FT, T::FT);

Specify prognostic variables, the variables solved for in the PDEs, for HeatModel

vars_state(::HeatModel, ::Prognostic, FT) = @vars(ρcT::FT);

Specify state variables whose gradients are needed for HeatModel

vars_state(::HeatModel, ::Gradient, FT) = @vars(ρcT::FT);

Specify gradient variables for HeatModel

vars_state(::HeatModel, ::GradientFlux, FT) = @vars(α∇ρcT::SVector{3, FT});

Define the compute kernels

Specify the initial values in aux::Vars, which are available in init_state_prognostic!. Note that

• this method is only called at t=0
• aux.z and aux.T are available here because we've specified z and T

in vars_state given Auxiliary in vars_state

function nodal_init_state_auxiliary!(
    m::HeatModel,
    aux::Vars,
    tmp::Vars,
    geom::LocalGeometry,
)
    aux.z = geom.coord[3]
    aux.T = m.initialT
end;

Specify the initial values in state::Vars. Note that

• this method is only called at t=0
• state.ρcT is available here because we've specified ρcT in

vars_state given Prognostic

function init_state_prognostic!(
    m::HeatModel,
    state::Vars,
    aux::Vars,
    localgeo,
    t::Real,
)
    state.ρcT = m.ρc * aux.T
end;

The remaining methods, defined in this section, are called at every time-step in the solver by the BalanceLaw framework.

Compute/update all auxiliary variables at each node. Note that

• aux.T is available here because we've specified T in

vars_state given Auxiliary

function nodal_update_auxiliary_state!(
    m::HeatModel,
    state::Vars,
    aux::Vars,
    t::Real,
)
    aux.T = state.ρcT / m.ρc
end;

Since we have second-order fluxes, we must tell ClimateMachine to compute the gradient of ρcT. Here, we specify how ρcT is computed. Note that

• transform.ρcT is available here because we've specified ρcT in

vars_state given Gradient

function compute_gradient_argument!(
    m::HeatModel,
    transform::Vars,
    state::Vars,
    aux::Vars,
    t::Real,
)
    transform.ρcT = state.ρcT
end;

Specify where in diffusive::Vars to store the computed gradient from compute_gradient_argument!. Note that:

• diffusive.α∇ρcT is available here because we've specified α∇ρcT in

vars_state given Gradient

• ∇transform.ρcT is available here because we've specified ρcT in

vars_state given Gradient

function compute_gradient_flux!(
    m::HeatModel,
    diffusive::Vars,
    ∇transform::Grad,
    state::Vars,
    aux::Vars,
    t::Real,
)
    diffusive.α∇ρcT = -m.α * ∇transform.ρcT
end;

We have no sources, nor non-diffusive fluxes.

function source!(m::HeatModel, _...) end;
function flux_first_order!(m::HeatModel, _...) end;

Compute diffusive flux ($F(α, ρcT, t) = -α ∇ρcT$ in the original PDE). Note that:

• diffusive.α∇ρcT is available here because we've specified α∇ρcT in

vars_state given GradientFlux

function flux_second_order!(
    m::HeatModel,
    flux::Grad,
    state::Vars,
    diffusive::Vars,
    hyperdiffusive::Vars,
    aux::Vars,
    t::Real,
)
    flux.ρcT += diffusive.α∇ρcT
end;

Boundary conditions

Second-order terms in our equations, $∇⋅(F)$ where $F = -α∇ρcT$, are internally reformulated to first-order unknowns. Boundary conditions must be specified for all unknowns, both first-order and second-order unknowns which have been reformulated.

struct DirichletBC <: BoundaryCondition end;
struct NeumannBC <: BoundaryCondition end;

boundary_conditions(::HeatModel) = (DirichletBC(), NeumannBC())

boundary_conditions (generic function with 14 methods)

The boundary conditions for ρcT (first order unknown)

function boundary_state!(
    nf,
    bc::DirichletBC,
    m::HeatModel,
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    # Apply Dirichlet BCs
    state⁺.ρcT = m.ρc * m.T_bottom
end;
function boundary_state!(
    nf,
    bc::NeumannBC,
    m::HeatModel,
    state⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    nothing
end;

The boundary conditions for ρcT are specified here for second-order unknowns

function boundary_state!(
    nf,
    bc::DirichletBC,
    m::HeatModel,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    nothing
end;
function boundary_state!(
    nf,
    bc::NeumannBC,
    m::HeatModel,
    state⁺::Vars,
    diff⁺::Vars,
    hyperdiff⁺::Vars,
    aux⁺::Vars,
    n⁻,
    state⁻::Vars,
    diff⁻::Vars,
    hyperdiff⁻::Vars,
    aux⁻::Vars,
    t,
    _...,
)
    # Apply Neumann BCs
    diff⁺.α∇ρcT = n⁻ * m.flux_top
end;

Spatial discretization

Prescribe polynomial order of basis functions in finite elements

N_poly = 5;

Specify the number of vertical elements

nelem_vert = 10;

Specify the domain height

zmax = FT(1);

Establish a ClimateMachine single stack configuration

driver_config = ClimateMachine.SingleStackConfiguration(
    "HeatEquation",
    N_poly,
    nelem_vert,
    zmax,
    param_set,
    m,
    numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
);

ClimateMachine.array_type() = Array
┌ Info: Model composition
│     param_set = Main.##347.EarthParameterSet()
│     ρc = 1.0
│     α = 0.01
│     initialT = 295.15
│     T_bottom = 300.0
│     flux_top = 0.0
└ @ ClimateMachine /central/scratch/climaci/climatemachine-docs/1430/climatemachine-docs/src/Driver/driver_configs.jl:188
┌ Info: Defining `prognostic_vars` and
│ `eq_tends` for HeatModel will
│ enable printing a table of tendencies.
└ @ ClimateMachine.BalanceLaws /central/scratch/climaci/climatemachine-docs/1430/climatemachine-docs/src/BalanceLaws/show_tendencies.jl:65
┌ Info: Establishing single stack configuration for HeatEquation
│     precision               = Float64
│     horiz polynomial order  = 5
│     vert polynomial order   = 5
│     domain_min              = 0.00 m, 0.00 m, 0.00 m
│     domain_max              = 1.00 m, 1.00 m, 1.00 m
│     # vert elems            = 10
│     MPI ranks               = 1
│     min(Δ_horz)             = 0.12 m
│     min(Δ_vert)             = 0.01 m
└ @ ClimateMachine /central/scratch/climaci/climatemachine-docs/1430/climatemachine-docs/src/Driver/driver_configs.jl:612

Time discretization / solver

Specify simulation time (SI units)

t0 = FT(0)
timeend = FT(40)

40.0

In this section, we initialize the state vector and allocate memory for the solution in space (dg has the model m, which describes the PDEs as well as the function used for initialization). SolverConfiguration initializes the ODE solver, by default an explicit Low-Storage Runge-Kutta method. In this tutorial, we prescribe an option for an implicit Kvaerno3 method.

First, let's define how the time-step is computed, based on the Fourier number (i.e., diffusive Courant number) is defined. Because the HeatModel is a custom model, we must define how both are computed. First, we must define our own implementation of DGMethods.calculate_dt, (which we imported):

function calculate_dt(dg, model::HeatModel, Q, Courant_number, t, direction)
    Δt = one(eltype(Q))
    CFL = DGMethods.courant(diffusive_courant, dg, model, Q, Δt, t, direction)
    return Courant_number / CFL
end

calculate_dt (generic function with 5 methods)

Next, we'll define our implementation of diffusive_courant:

function diffusive_courant(
    m::HeatModel,
    state::Vars,
    aux::Vars,
    diffusive::Vars,
    Δx,
    Δt,
    t,
    direction,
)
    return Δt * m.α / (Δx * Δx)
end

diffusive_courant (generic function with 1 method)

Finally, we initialize the state vector and solver configuration based on the given Fourier number. Note that, we can use a much larger Fourier number for implicit solvers as compared to explicit solvers.

use_implicit_solver = false
if use_implicit_solver
    given_Fourier = FT(30)

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config;
        ode_solver_type = ImplicitSolverType(OrdinaryDiffEq.Kvaerno3(
            autodiff = false,
            linsolve = LinSolveGMRES(),
        )),
        Courant_number = given_Fourier,
        CFL_direction = VerticalDirection(),
    )
else
    given_Fourier = FT(0.7)

    solver_config = ClimateMachine.SolverConfiguration(
        t0,
        timeend,
        driver_config;
        Courant_number = given_Fourier,
        CFL_direction = VerticalDirection(),
    )
end;


grid = solver_config.dg.grid;
Q = solver_config.Q;
aux = solver_config.dg.state_auxiliary;

┌ Info: Initializing HeatEquation
└ @ ClimateMachine /central/scratch/climaci/climatemachine-docs/1430/climatemachine-docs/src/Driver/solver_configs.jl:185

Inspect the initial conditions

Let's export a plot of the initial state

output_dir = @__DIR__;

mkpath(output_dir);

z_scale = 100; # convert from meters to cm
z_key = "z";
z_label = "z [cm]";
z = get_z(grid; z_scale = z_scale, rm_dupes = true);

Create an array to store the solution:

dons_arr = Dict[dict_of_nodal_states(solver_config; interp = true)]  # store initial condition at ``t=0``
time_data = FT[0]                                      # store time data

export_plot(
    z,
    time_data,
    dons_arr,
    ("ρcT",),
    joinpath(output_dir, "initial_condition.png");
    xlabel = "ρcT",
    ylabel = z_label,
    xlims = (m.initialT - 1, m.T_bottom + 1),
);

It matches what we have in init_state_prognostic!(m::HeatModel, ...), so let's continue.

Solver hooks / callbacks

Define the number of outputs from t0 to timeend

const n_outputs = 5;

This equates to exports every ceil(Int, timeend/n_outputs) time-step:

const every_x_simulation_time = ceil(Int, timeend / n_outputs);

The ClimateMachine's time-steppers provide hooks, or callbacks, which allow users to inject code to be executed at specified intervals. In this callback, a dictionary of prognostic and auxiliary states are appended to dons_arr for time the callback is executed. In addition, time is collected and appended to time_data.

callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
    push!(dons_arr, dict_of_nodal_states(solver_config; interp = true))
    push!(time_data, gettime(solver_config.solver))
    nothing
end;

Solve

This is the main ClimateMachine solver invocation. While users do not have access to the time-stepping loop, code may be injected via user_callbacks, which is a Tuple of callbacks in GenericCallbacks.

ClimateMachine.invoke!(solver_config; user_callbacks = (callback,));

┌ Info: Starting HeatEquation
│     dt              = 9.65950e-03
│     timeend         =    40.00
│     number of steps = 4141
│     norm(Q)         = 2.9515000000000060e+02
└ @ ClimateMachine /central/scratch/climaci/climatemachine-docs/1430/climatemachine-docs/src/Driver/Driver.jl:802
┌ Info: Finished
│     norm(Q)            = 2.9853556724952733e+02
│     norm(Q) / norm(Q₀) = 1.0114706666085946e+00
│     norm(Q) - norm(Q₀) = 3.3855672495267299e+00
└ @ ClimateMachine /central/scratch/climaci/climatemachine-docs/1430/climatemachine-docs/src/Driver/Driver.jl:853

Append result at the end of the last time step:

push!(dons_arr, dict_of_nodal_states(solver_config; interp = true));
push!(time_data, gettime(solver_config.solver));

Post-processing

Our solution is stored in the array of dictionaries dons_arr whose keys are the output interval. The next level keys are the variable names, and the values are the values along the grid:

To get T at $t=0$, we can use T_at_t_0 = dons_arr[1]["T"][:]

@show keys(dons_arr[1])

Base.KeySet for a Dict{String,Array{Float64,1}} with 3 entries. Keys:
  "T"
  "z"
  "ρcT"

Let's plot the solution:

export_plot(
    z,
    time_data,
    dons_arr,
    ("ρcT",),
    joinpath(output_dir, "solution_vs_time.png");
    xlabel = "ρcT",
    ylabel = z_label,
);

export_contour(
    z,
    time_data,
    dons_arr,
    "ρcT",
    joinpath(output_dir, "solution_contour.png");
    ylabel = "z [cm]",
)

The results look as we would expect: a fixed temperature at the bottom is resulting in heat flux that propagates up the domain. To run this file, and inspect the solution in dons_arr, include this tutorial in the Julia REPL with:

include(joinpath("tutorials", "Land", "Heat", "heat_equation.jl"))

This page was generated using Literate.jl.