# 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:

- Preliminary configuration
- PDEs
- Space discretization
- Time discretization / solver
- Solver hooks / callbacks
- Solve
- 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: (
`import`

ing 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 `import`

ed 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.*