GitHub

Introduction to ClimaCore.jl

This tutorial is available as a Jupyter notebook.

What is ClimaCore?

A suite of tools for constructing spatial discretizations

  • primarily aimed at climate and weather models
  • initial aim:
    • spectral element discretization in the horizontal
    • staggered finite difference in the vertical
  • currently under development
using ClimaComms, ClimaCore, ClimaCorePlots, LinearAlgebra, IntervalSets, Plots

1. Constructing a discretization

1.1 Domains

A domain is a region of space (think of a mathematical domain).

column_domain = ClimaCore.Domains.IntervalDomain(
    ClimaCore.Geometry.ZPoint(0.0) .. ClimaCore.Geometry.ZPoint(10.0),
    boundary_names = (:bottom, :top),
)
IntervalDomain: z ∈ [0.0,10.0] (:bottom, :top)
rectangle_domain = ClimaCore.Domains.RectangleDomain(
    ClimaCore.Geometry.XPoint(-2π) .. ClimaCore.Geometry.XPoint(2π),
    ClimaCore.Geometry.YPoint(-2π) .. ClimaCore.Geometry.YPoint(2π),
    x1periodic = true,
    x2periodic = true,
)
RectangleDomain: x ∈ [-6.283185307179586,6.283185307179586] (periodic) × y ∈ [-6.283185307179586,6.283185307179586] (periodic)

1.2 Meshes

A mesh is a division of a domain into elements

column_mesh = ClimaCore.Meshes.IntervalMesh(column_domain, nelems = 32)
32-element IntervalMesh of IntervalDomain: z ∈ [0.0,10.0] (:bottom, :top)
rectangle_mesh = ClimaCore.Meshes.RectilinearMesh(rectangle_domain, 16, 16)
16×16-element RectilinearMesh of RectangleDomain: x ∈ [-6.283185307179586,6.283185307179586] (periodic) × y ∈ [-6.283185307179586,6.283185307179586] (periodic)

1.3 Topologies

A topology determines the ordering and connections between elements of a mesh

  • At the moment, this is only required for 2D meshes
device = ClimaComms.device()
rectangle_topology = ClimaCore.Topologies.Topology2D(
    ClimaComms.SingletonCommsContext(device),
    rectangle_mesh,
)
Topology2D
  context: SingletonCommsContext using CPUSingleThreaded
  mesh: 16×16-element RectilinearMesh of RectangleDomain: x ∈ [-6.283185307179586,6.283185307179586] (periodic) × y ∈ [-6.283185307179586,6.283185307179586] (periodic)
  elemorder: CartesianIndices((16, 16))

1.4 Spaces

A space represents a discretized function space over some domain. Currently two discretizations are supported.

1.4.1 Staggered finite difference discretization

This discretizes an interval domain by approximating the function by a value at either the center of each element (CenterFiniteDifferenceSpace), or the faces between elements (FaceFiniteDifferenceSpace).

You can construct either the center or face space from the mesh, then construct the opposite space from the original one (this is to avoid allocating additional memory).

column_center_space =
    ClimaCore.Spaces.CenterFiniteDifferenceSpace(device, column_mesh)
# construct the face space from the center one
column_face_space =
    ClimaCore.Spaces.FaceFiniteDifferenceSpace(column_center_space)
FaceFiniteDifferenceSpace:
  context: SingletonCommsContext using CPUSingleThreaded
  mesh: 32-element IntervalMesh of IntervalDomain: z ∈ [0.0,10.0] (:bottom, :top)

1.4.2 Spectral element discretization

A spectral element space approximates the function with polynomials in each element. The polynomials are represented using a nodal discretization, which stores the values of the polynomials at particular points in each element (termed nodes or degrees of freedom).

These nodes are chosen by a particular quadrature rule, which allows us to integrate functions over the domain. The only supported choice for now is a Gauss-Legendre-Lobatto rule.

# Gauss-Legendre-Lobatto quadrature with 4 nodes in each direction, so 16 in each element
quad = ClimaCore.Quadratures.GLL{4}()
rectangle_space =
    ClimaCore.Spaces.SpectralElementSpace2D(rectangle_topology, quad)
SpectralElementSpace2D:
  context: SingletonCommsContext using CPUSingleThreaded
  mesh: 16×16-element RectilinearMesh of RectangleDomain: x ∈ [-6.283185307179586,6.283185307179586] (periodic) × y ∈ [-6.283185307179586,6.283185307179586] (periodic)
  quadrature: 4-point Gauss-Legendre-Lobatto quadrature

1.5 Fields

Finally, we can construct a field: a function in a space. A field is simply a space and the values at each node in the space.

The easiest field to construct is the coordinate field

coord = ClimaCore.Fields.coordinate_field(rectangle_space)
ClimaCore.Geometry.XYPoint{Float64}-valued Field:
  x: [-6.28319, -6.06611, -5.71487, -5.49779, -6.28319, -6.06611, -5.71487, -5.49779, -6.28319, -6.06611  …  6.06611, 6.28319, 5.49779, 5.71487, 6.06611, 6.28319, 5.49779, 5.71487, 6.06611, 6.28319]
  y: [-6.28319, -6.28319, -6.28319, -6.28319, -6.06611, -6.06611, -6.06611, -6.06611, -5.71487, -5.71487  …  5.71487, 5.71487, 6.06611, 6.06611, 6.06611, 6.06611, 6.28319, 6.28319, 6.28319, 6.28319]

This is a struct-value field: it contains coordinates in a struct at each point. We can extract just the x coordinate, to get a scalar field:

x = coord.x
Float64-valued Field:
  [-6.28319, -6.06611, -5.71487, -5.49779, -6.28319, -6.06611, -5.71487, -5.49779, -6.28319, -6.06611  …  6.06611, 6.28319, 5.49779, 5.71487, 6.06611, 6.28319, 5.49779, 5.71487, 6.06611, 6.28319]

Although you can't index directly into a field, it can be used in some other ways similar to a Julia Array. For example, broadcasting can be used to define new fields in terms of other ones:

sinx = sin.(x)
Float64-valued Field:
  [2.44929e-16, 0.215378, 0.538216, 0.707107, 2.44929e-16, 0.215378, 0.538216, 0.707107, 2.44929e-16, 0.215378  …  -0.215378, -2.44929e-16, -0.707107, -0.538216, -0.215378, -2.44929e-16, -0.707107, -0.538216, -0.215378, -2.44929e-16]

Fields can be easily vizualized with Plots.jl:

import Plots
Plots.plot(sinx)
Example block output

If you're using the terminal, UnicodePlots is also supported.

This works similarly for finite difference discretizations

column_center_coords = ClimaCore.Fields.coordinate_field(column_center_space)
column_face_coords = ClimaCore.Fields.coordinate_field(column_face_space)
ClimaCore.Geometry.ZPoint{Float64}-valued Field:
  z: [0.0, 0.3125, 0.625, 0.9375, 1.25, 1.5625, 1.875, 2.1875, 2.5, 2.8125  …  7.1875, 7.5, 7.8125, 8.125, 8.4375, 8.75, 9.0625, 9.375, 9.6875, 10.0]
plot(sin.(column_center_coords.z), ylim = (0.0, 10.0))
plot!(cos.(column_face_coords.z), ylim = (0.0, 10.0))
Example block output

Reduction operations are defined anologously:

  • sum will give the integral of the function

\[\int_D f(x) dx\]

  • norm will give the L² function norm

\[\sqrt{\int_D |f(x)|^2 dx}\]

sum(sinx) ## integral
-1.609823385706477e-15
norm(sinx) ## L² norm
0.7071067811865474

1.6 Vectors and vector fields

A vector field is a field with vector-valued quantity, i.e. at every point in space, you have a vector.

However one of the key requirements of ClimaCore is to support vectors specified in curvilinear or non-Cartesian coordinates. We will discuss this in a bit further, but for now, you can define a 2-dimensional vector field using Geometry.UVVector:

v = ClimaCore.Geometry.UVVector.(coord.y, .-coord.x)
ClimaCore.Geometry.UVVector{Float64}-valued Field:
  components: 
    data: 
      1: [-6.28319, -6.28319, -6.28319, -6.28319, -6.06611, -6.06611, -6.06611, -6.06611, -5.71487, -5.71487  …  5.71487, 5.71487, 6.06611, 6.06611, 6.06611, 6.06611, 6.28319, 6.28319, 6.28319, 6.28319]
      2: [6.28319, 6.06611, 5.71487, 5.49779, 6.28319, 6.06611, 5.71487, 5.49779, 6.28319, 6.06611  …  -6.06611, -6.28319, -5.49779, -5.71487, -6.06611, -6.28319, -5.49779, -5.71487, -6.06611, -6.28319]

2. Operators

Operators can compute spatial derivative operations.

  • for performance reasons, we need to be able to "fuse" multiple operators and function applications
  • Julia provides a tool for this: broadcasting, with a very flexible API

Can think of operators are "pseudo-functions": can't be called directly, but act similar to functions in the context of broadcasting.

2.1 Spectral element operators

The Gradient operator takes the gradient of a scalar field, and returns a vector field.

grad = ClimaCore.Operators.Gradient()
∇sinx = grad.(sinx)
ClimaCore.Geometry.Covariant12Vector{Float64}-valued Field:
  components: 
    data: 
      1: [0.393185, 0.383236, 0.331262, 0.276966, 0.393185, 0.383236, 0.331262, 0.276966, 0.393185, 0.383236  …  0.383236, 0.393185, 0.276966, 0.331262, 0.383236, 0.393185, 0.276966, 0.331262, 0.383236, 0.393185]
      2: [1.72563e-31, 1.38778e-16, 3.33067e-16, 0.0, 0.0, -2.77556e-17, 5.55112e-17, 0.0, 0.0, 2.77556e-17  …  2.77556e-17, 0.0, 0.0, 5.55112e-17, -2.77556e-17, 0.0, 0.0, 4.44089e-16, 1.11022e-16, 1.97215e-31]
plot(∇sinx.components.data.:1, clim = (-1, 1))
Example block output

This returns the gradient in covariant coordinates

\[(\nabla f)_i = \frac{\partial f}{\partial \xi^i}\]

where $(\xi^1,\xi^2)$ are the coordinates in the reference element: a square $[-1,1]^2$.

This can be converted to a local orthogonal basis by multiplying by the partial derivative matrix

\[\frac{\partial \xi}{\partial x}\]

This can be done calling ClimaCore.Geometry.LocalVector:

∇sinx_cart = ClimaCore.Geometry.LocalVector.(∇sinx)
ClimaCore.Geometry.UVVector{Float64}-valued Field:
  components: 
    data: 
      1: [1.00124, 0.975903, 0.843551, 0.705289, 1.00124, 0.975903, 0.843551, 0.705289, 1.00124, 0.975903  …  0.975903, 1.00124, 0.705289, 0.843551, 0.975903, 1.00124, 0.705289, 0.843551, 0.975903, 1.00124]
      2: [4.39429e-31, 3.53395e-16, 8.48148e-16, 0.0, 0.0, -7.0679e-17, 1.41358e-16, 0.0, 0.0, 7.0679e-17  …  7.0679e-17, 0.0, 0.0, 1.41358e-16, -7.0679e-17, 0.0, 0.0, 1.13086e-15, 2.82716e-16, 5.02204e-31]
plot(∇sinx_cart.components.data.:1, clim = (-1, 1))
Example block output
plot(∇sinx_cart.components.data.:2, clim = (-1, 1))
Example block output
∇sinx_ref = ClimaCore.Geometry.UVVector.(cos.(x), 0.0)
norm(∇sinx_cart .- ∇sinx_ref)
0.0016360508523548926

Similarly, the Divergence operator takes the divergence of vector field, and returns a scalar field.

If we take the divergence of a gradient, we can get a Laplacian:

div = ClimaCore.Operators.Divergence()
∇²sinx = div.(grad.(sinx))
plot(∇²sinx)
Example block output

Note: In curvilinear coordinates, the divergence is defined in terms of the contravariant components $u^i$:

\[\nabla \cdot u = \frac{1}{J} \sum_i \frac{\partial}{\partial \xi^i} (J u^i)\]

The Divergence operator handles this conversion internally.

2.1.1 Direct stiffness summation

Spectral element operators only operate within a single element, and so the result may be discontinuous. To address this, the usual fix is direct stiffness summation (DSS), which averages the values at the element boundaries.

This corresponds to the $L^2$ projection onto the subset of continuous functions in our function space.

∇²sinx_dss = ClimaCore.Spaces.weighted_dss!(copy(∇²sinx))
plot(∇²sinx_dss)
Example block output
plot(∇²sinx_dss .- ∇²sinx)
Example block output

2.2 Finite difference operators

Finite difference operators are similar with some subtle differences:

  • they can change staggering (center to face, or vice versa)
  • they can span multiple elements
    • no DSS is required
    • boundary handling may be required

We use the following convention:

  • centers are indexed by integers 1, 2, ..., n
  • faces are indexed by half integers ½, 1+½, ..., n+½

Face to center gradient

An finite-difference operator defines a stencil. For example, the gradient operator

\[\nabla\theta[i] = \frac{\theta [i+\tfrac{1}{2}] - \theta[i-\tfrac{1}{2}]}{\Delta z}\]

(actually, a little more complicated as it gives a vector in a covariant basis)

        ...
      /
θ[2+½]
      \
        ∇θ[2]
      /
θ[1+½]
      \
        ∇θ[1]
      /
θ[½]

Every center value is well-defined, so boundary handling is optional.

cosz = cos.(column_face_coords.z)
gradf2c = ClimaCore.Operators.GradientF2C()
∇cosz = gradf2c.(cosz)
ClimaCore.Geometry.Covariant3Vector{Float64}-valued Field:
  components: 
    data: 
      1: [-0.0484321, -0.140605, -0.219158, -0.276483, -0.307026, -0.30783, -0.278816, -0.222794, -0.145192, -0.0535264  …  -0.211705, -0.271589, -0.305166, -0.309183, -0.283251, -0.229882, -0.154246, -0.0636694, 0.0330747, 0.126615]
plot(map(x -> x.w, ClimaCore.Geometry.WVector.(∇cosz)), ylim = (0, 10))
Example block output

Center to face gradient

Uses the same stencil, but doesn't work directly:

sinz = sin.(column_center_coords.z)
gradc2f = ClimaCore.Operators.GradientC2F()
# ∇sinz = gradc2f.(sinz) ## this would throw an error
ClimaCore.Operators.GradientC2F{@NamedTuple{}}(NamedTuple())

This throws an error because face values at the boundary are not well-defined:

...
      \
        ∇θ[2+½]
      /
θ[2]
      \
        ∇θ[1+½]
      /
θ[1]
      \
        ????

To handle boundaries we need to modify the stencil. Two options:

  • provide the value $\theta^*$ of $\theta$ at the boundary:

\[\nabla\theta[\tfrac{1}{2}] = \frac{\theta[1] - \theta^*}{\Delta z /2}\]

  • provide the gradient $\nabla\theta^*$ of $\theta$ at the boundary:

\[\nabla\theta[\tfrac{1}{2}] = \nabla\theta^*\]

These modified stencils are provided as keyword arguments to the operator (based on the boundary label names):

sinz = sin.(column_center_coords.z)
gradc2f = ClimaCore.Operators.GradientC2F(
    bottom = ClimaCore.Operators.SetValue(sin(0.0)),
    top = ClimaCore.Operators.SetGradient(
        ClimaCore.Geometry.WVector(cos(10.0)),
    ),
)
∇sinz = gradc2f.(sinz)
ClimaCore.Geometry.Covariant3Vector{Float64}-valued Field:
  components: 
    data: 
      1: [0.31123, 0.296156, 0.252396, 0.184187, 0.0981378, 0.00258204, -0.0932238, -0.18, -0.24934, -0.294528  …  0.19241, 0.107883, 0.0129066, -0.0833202, -0.171476, -0.243023, -0.291029, -0.310844, -0.300551, -0.26221]
plot(map(x -> x.w, ClimaCore.Geometry.WVector.(∇sinz)), ylim = (0, 10))
Example block output

As before, multiple operators (or functions) can be fused together with broadcasting.

One extra advantage of this is that boundaries of the inner operators only need to be specified if they would affect the final result.

Consider the center-to-center Laplacian:

...
      \       /
        ∇θ[2+½]
      /       \
θ[2]            ∇⋅∇θ[2]
      \       /
        ∇θ[1+½]
      /       \
θ[1]            ∇⋅∇θ[1]
              /
         ∇θ*
sinz = sin.(column_center_coords.z)
# we don't need to specify boundaries, as the stencil won't reach that far
gradc2f = ClimaCore.Operators.GradientC2F()
divf2c = ClimaCore.Operators.DivergenceF2C(
    bottom = ClimaCore.Operators.SetValue(ClimaCore.Geometry.WVector(cos(0.0))),
    top = ClimaCore.Operators.SetValue(ClimaCore.Geometry.WVector(cos(10.0))),
)
∇∇sinz = divf2c.(gradc2f.(sinz))
Float64-valued Field:
  [-0.167358, -0.448107, -0.698456, -0.881149, -0.978491, -0.981052, -0.888584, -0.710045, -0.462727, -0.170588  …  -0.674704, -0.865553, -0.972561, -0.985363, -0.902719, -0.732633, -0.491582, -0.202914, 0.105409, 0.39261]
plot(∇∇sinz, ylim = (0, 10))
Example block output

3. Solving PDEs

ClimaCore can be used for spatial discretizations of PDEs. For temporal discretization, we can use the OrdinaryDiffEq package, which we aim to be compatibile with.

import OrdinaryDiffEqSSPRK: ODEProblem, solve, SSPRK33

3.1 Heat equation using finite differences

We will use a cell-center discretization of the heat equation:

\[\frac{\partial y}{\partial t} = \alpha \nabla \cdot \nabla y\]

At the bottom we will use a Dirichlet condition $y(0) = 1$ at the bottom: since we don't actually have a value located at the bottom, we will use a SetValue boundary modifier on the inner gradient.

At the top we will use a Neumann condition $\frac{\partial y}{\partial z}(10) = 0$. We can do this two equivalent ways:

  • a SetGradient on the gradient operator
  • a SetValue on the divergence operator

either will work.

y0 = zeros(column_center_space)

# define the tendency function
function heat_fd_tendency!(dydt, y, α, t)
    gradc2f = ClimaCore.Operators.GradientC2F(
        bottom = ClimaCore.Operators.SetValue(1.0),
        top = ClimaCore.Operators.SetGradient(ClimaCore.Geometry.WVector(0.0)),
    )
    divf2c = ClimaCore.Operators.DivergenceF2C()
    # the @. macro "dots" the whole expression
    # i.e.  dydt .= α .* divf2c.(gradc2f.(y))
    @. dydt = α * divf2c(gradc2f(y))
end

heat_fd_prob = ODEProblem(heat_fd_tendency!, y0, (0.0, 5.0), 0.1)
heat_fd_sol = solve(heat_fd_prob, SSPRK33(), dt = 0.1, saveat = 0.25)
retcode: Success
Interpolation: 1st order linear
t: 21-element Vector{Float64}:
 0.0
 0.25
 0.5
 0.75
 1.0
 1.25
 1.5
 1.75
 2.0
 2.25
 ⋮
 3.0
 3.25
 3.5
 3.75
 4.0
 4.25
 4.5
 4.75
 5.0
u: 21-element Vector{ClimaCore.Fields.Field{ClimaCore.DataLayouts.VF{Float64, 32, Matrix{Float64}}, ClimaCore.Spaces.FiniteDifferenceSpace{ClimaCore.Grids.FiniteDifferenceGrid{ClimaCore.Topologies.IntervalTopology{ClimaComms.SingletonCommsContext{ClimaComms.CPUSingleThreaded}, ClimaCore.Meshes.IntervalMesh{ClimaCore.Domains.IntervalDomain{ClimaCore.Geometry.ZPoint{Float64}, Tuple{Symbol, Symbol}}, LinRange{ClimaCore.Geometry.ZPoint{Float64}, Int64}}, @NamedTuple{bottom::Int64, top::Int64}}, ClimaCore.Geometry.CartesianGlobalGeometry, ClimaCore.DataLayouts.VF{ClimaCore.Geometry.LocalGeometry{(3,), ClimaCore.Geometry.ZPoint{Float64}, Float64, StaticArraysCore.SMatrix{1, 1, Float64, 1}}, 32, Matrix{Float64}}, ClimaCore.DataLayouts.VF{ClimaCore.Geometry.LocalGeometry{(3,), ClimaCore.Geometry.ZPoint{Float64}, Float64, StaticArraysCore.SMatrix{1, 1, Float64, 1}}, 33, Matrix{Float64}}}, ClimaCore.Grids.CellCenter}}}:
 Float64-valued Field:
  [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
 Float64-valued Field:
  [0.361354, 0.0434878, 0.00369219, 0.000231351, 1.08527e-5, 3.9575e-7, 1.05494e-8, 2.03144e-10, 2.8656e-12, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
 Float64-valued Field:
  [0.540691, 0.122043, 0.0195753, 0.00241637, 0.000239543, 1.95657e-5, 1.33644e-6, 7.68579e-8, 3.72538e-9, 1.51359e-10  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
 Float64-valued Field:
  [0.638058, 0.199551, 0.0457887, 0.00822793, 0.00120846, 0.000149308, 1.58313e-5, 1.46111e-6, 1.18569e-7, 8.52068e-9  …  9.85532e-29, 5.25968e-31, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
 Float64-valued Field:
  [0.696003, 0.266731, 0.0775589, 0.0179644, 0.00343741, 0.000558021, 7.83785e-5, 9.66625e-6, 1.05847e-6, 1.0379e-7  …  1.29808e-24, 3.15045e-26, 6.6329e-28, 1.18493e-29, 1.7371e-31, 1.97805e-33, 1.57643e-35, 6.73777e-38, 7.65996e-54, 3.83841e-69]
 Float64-valued Field:
  [0.733759, 0.322823, 0.111157, 0.0310221, 0.0072294, 0.00144032, 0.000249871, 3.82933e-5, 5.24341e-6, 6.4733e-7  …  5.91602e-22, 2.46848e-23, 9.56159e-25, 3.43109e-26, 1.13754e-27, 3.4726e-29, 9.71884e-31, 2.48013e-32, 5.73082e-34, 1.21016e-35]
 Float64-valued Field:
  [0.760295, 0.369397, 0.144362, 0.0465007, 0.0126402, 0.00295846, 0.000606108, 0.000110158, 1.79549e-5, 2.64799e-6  …  3.84357e-20, 2.09678e-21, 1.0755e-22, 5.1847e-24, 2.34736e-25, 9.97082e-27, 3.96812e-28, 1.47701e-29, 5.13078e-31, 1.7085e-32]
 Float64-valued Field:
  [0.780136, 0.408345, 0.176016, 0.0635216, 0.0195479, 0.00521693, 0.00122513, 0.000256286, 4.82486e-5, 8.24493e-6  …  1.15548e-18, 7.81446e-20, 5.01303e-21, 3.05198e-22, 1.76395e-23, 9.68035e-25, 5.04432e-26, 2.49551e-27, 1.17183e-28, 5.43974e-30]
 Float64-valued Field:
  [0.79567, 0.441286, 0.205624, 0.0813748, 0.0277362, 0.00825645, 0.00217375, 0.00051177, 0.000108765, 2.10353e-5  …  1.79615e-17, 1.42207e-18, 1.07183e-19, 7.6964e-21, 5.26838e-22, 3.43961e-23, 2.14262e-24, 1.27378e-25, 7.22884e-27, 4.11706e-28]
 Float64-valued Field:
  [0.808272, 0.469502, 0.233053, 0.0995253, 0.0369525, 0.0120651, 0.00350186, 0.000912449, 0.000215283, 4.63418e-5  …  1.91179e-16, 1.73675e-17, 1.50662e-18, 1.2493e-19, 9.91056e-21, 7.52709e-22, 5.47689e-23, 3.81997e-24, 2.55568e-25, 1.74055e-26]
 ⋮
 Float64-valued Field:
  [0.835409, 0.534355, 0.303165, 0.152621, 0.0684478, 0.0275074, 0.00997066, 0.00328129, 0.000986593, 0.000272604  …  4.24639e-14, 5.23169e-15, 6.1801e-16, 7.00794e-17, 7.63647e-18, 8.00435e-19, 8.07754e-20, 7.85432e-21, 7.36918e-22, 7.24458e-23]
 Float64-valued Field:
  [0.842167, 0.551285, 0.323018, 0.169312, 0.0796018, 0.0337227, 0.0129436, 0.0045269, 0.00145073, 0.000428259  …  1.77326e-13, 2.37085e-14, 3.04205e-15, 3.75043e-16, 4.44764e-17, 5.0787e-18, 5.58927e-19, 5.93354e-20, 6.08614e-21, 6.59662e-22]
 Float64-valued Field:
  [0.848151, 0.566512, 0.341394, 0.185376, 0.0908525, 0.0403315, 0.0162908, 0.00601678, 0.00204203, 0.000639897  …  6.42152e-13, 9.23906e-14, 1.2765e-14, 1.69567e-15, 2.16808e-16, 2.67096e-17, 3.17344e-18, 3.63956e-19, 4.03688e-20, 4.76686e-21]
 Float64-valued Field:
  [0.853501, 0.580302, 0.358446, 0.200794, 0.102103, 0.047256, 0.0199814, 0.0077514, 0.00277088, 0.000916665  …  2.07924e-12, 3.20247e-13, 4.73968e-14, 6.74877e-15, 9.25548e-16, 1.22384e-16, 1.56178e-17, 1.92523e-18, 2.2975e-19, 2.94068e-20]
 Float64-valued Field:
  [0.85832, 0.592863, 0.374309, 0.215566, 0.113279, 0.0544266, 0.0239822, 0.00972635, 0.00364473, 0.00126682  …  6.07317e-12, 9.95514e-13, 1.56883e-13, 2.37971e-14, 3.4784e-15, 4.90448e-16, 6.67707e-17, 8.78533e-18, 1.11991e-18, 1.5416e-19]
 Float64-valued Field:
  [0.862692, 0.604367, 0.389106, 0.229705, 0.124325, 0.0617818, 0.0282585, 0.0119333, 0.00466831, 0.00169767  …  1.63045e-11, 2.83311e-12, 4.73519e-13, 7.62168e-14, 1.18274e-14, 1.77134e-15, 2.56278e-16, 3.58531e-17, 4.86365e-18, 7.17279e-19]
 Float64-valued Field:
  [0.866681, 0.614954, 0.402943, 0.243233, 0.1352, 0.0692684, 0.0327764, 0.0143611, 0.00584377, 0.00221526  …  4.04935e-11, 7.42659e-12, 1.31067e-12, 2.22849e-13, 3.65447e-14, 5.78602e-15, 8.85318e-16, 1.31036e-16, 1.88202e-17, 2.95712e-18]
 Float64-valued Field:
  [0.870341, 0.624737, 0.415913, 0.256173, 0.145873, 0.0768407, 0.0375031, 0.0169964, 0.00717101, 0.00282445  …  9.41822e-11, 1.8174e-11, 3.37615e-12, 6.04497e-13, 1.04434e-13, 1.74264e-14, 2.81132e-15, 4.38901e-16, 6.65448e-17, 1.11059e-17]
 Float64-valued Field:
  [0.873715, 0.633814, 0.4281, 0.268555, 0.156323, 0.0844602, 0.0424079, 0.0198249, 0.008648, 0.00352881  …  2.06147e-10, 4.17142e-11, 8.12922e-12, 1.52747e-12, 2.77027e-13, 4.85442e-14, 8.22681e-15, 1.34966e-15, 2.15194e-16, 3.79886e-17]
anim = Plots.@animate for u in heat_fd_sol.u
    plot(u, xlim = (0, 1), ylim = (0, 10))
end
mp4(anim)

3.2 Heat equation using continuous Galerkin (CG) spectral element

function heat_cg_tendency!(dydt, y, α, t)
    grad = ClimaCore.Operators.Gradient()
    wdiv = ClimaCore.Operators.WeakDivergence()
    # apply element operators
    @. dydt = α * wdiv(grad(y))

    # direct stiffness summation (DSS): project to continuous function space
    ClimaCore.Spaces.weighted_dss!(dydt)
    return dydt
end

y0 = exp.(.-(coord.y .^ 2 .+ coord.x .^ 2) ./ 2)

heat_cg_prob = ODEProblem(heat_cg_tendency!, y0, (0.0, 5.0), 0.1)
heat_cg_sol = solve(heat_cg_prob, SSPRK33(), dt = 0.1, saveat = 0.5)
retcode: Success
Interpolation: 1st order linear
t: 11-element Vector{Float64}:
 0.0
 0.5
 1.0
 1.5
 2.0
 2.5
 3.0
 3.5
 4.0
 4.5
 5.0
u: 11-element Vector{ClimaCore.Fields.Field{ClimaCore.DataLayouts.IJFH{Float64, 4, Array{Float64, 4}}, ClimaCore.Spaces.SpectralElementSpace2D{ClimaCore.Grids.SpectralElementGrid2D{ClimaCore.Topologies.Topology2D{ClimaComms.SingletonCommsContext{ClimaComms.CPUSingleThreaded}, ClimaCore.Meshes.RectilinearMesh{ClimaCore.Meshes.IntervalMesh{ClimaCore.Domains.IntervalDomain{ClimaCore.Geometry.XPoint{Float64}, Nothing}, LinRange{ClimaCore.Geometry.XPoint{Float64}, Int64}}, ClimaCore.Meshes.IntervalMesh{ClimaCore.Domains.IntervalDomain{ClimaCore.Geometry.YPoint{Float64}, Nothing}, LinRange{ClimaCore.Geometry.YPoint{Float64}, Int64}}}, CartesianIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, LinearIndices{2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, Vector{Tuple{Int64, Int64, Int64, Int64, Bool}}, Vector{Tuple{Int64, Int64, Int64, Int64, Bool}}, Vector{Tuple{Int64, Int64}}, Vector{Int64}, Vector{Tuple{Bool, Int64, Int64}}, Vector{Int64}, Vector{Int64}, @NamedTuple{}, Vector{Tuple{Int64, Int64}}}, ClimaCore.Quadratures.GLL{4}, ClimaCore.Geometry.CartesianGlobalGeometry, ClimaCore.DataLayouts.IJFH{ClimaCore.Geometry.LocalGeometry{(1, 2), ClimaCore.Geometry.XYPoint{Float64}, Float64, StaticArraysCore.SMatrix{2, 2, Float64, 4}}, 4, Array{Float64, 4}}, ClimaCore.DataLayouts.IJFH{Float64, 4, Array{Float64, 4}}, ClimaCore.DataLayouts.IFH{ClimaCore.Geometry.SurfaceGeometry{Float64, ClimaCore.Geometry.UVVector{Float64}}, 4, Array{Float64, 3}}, @NamedTuple{}}}}}:
 Float64-valued Field:
  [7.15717e-18, 2.7344e-17, 2.16475e-16, 7.31057e-16, 2.7344e-17, 1.04468e-16, 8.27047e-16, 2.79301e-15, 2.16475e-16, 8.27047e-16  …  8.27047e-16, 2.16475e-16, 2.79301e-15, 8.27047e-16, 1.04468e-16, 2.7344e-17, 7.31057e-16, 2.16475e-16, 2.7344e-17, 7.15717e-18]
 Float64-valued Field:
  [9.43714e-16, 1.65105e-15, 1.07955e-14, 3.19401e-14, 1.65105e-15, 2.88557e-15, 1.8897e-14, 5.57259e-14, 1.07955e-14, 1.8897e-14  …  1.8897e-14, 1.07955e-14, 5.57259e-14, 1.8897e-14, 2.88557e-15, 1.65105e-15, 3.19401e-14, 1.07955e-14, 1.65105e-15, 9.43714e-16]
 Float64-valued Field:
  [1.76354e-14, 2.8768e-14, 1.52544e-13, 4.13456e-13, 2.8768e-14, 4.6842e-14, 2.48378e-13, 6.73229e-13, 1.52544e-13, 2.48378e-13  …  2.48378e-13, 1.52544e-13, 6.73229e-13, 2.48378e-13, 4.6842e-14, 2.8768e-14, 4.13456e-13, 1.52544e-13, 2.8768e-14, 1.76354e-14]
 Float64-valued Field:
  [2.02605e-13, 3.12665e-13, 1.40719e-12, 3.53943e-12, 3.12665e-13, 4.82027e-13, 2.16942e-12, 5.45561e-12, 1.40719e-12, 2.16942e-12  …  2.16942e-12, 1.40719e-12, 5.45561e-12, 2.16942e-12, 4.82027e-13, 3.12665e-13, 3.53943e-12, 1.40719e-12, 3.12665e-13, 2.02605e-13]
 Float64-valued Field:
  [1.63452e-12, 2.40218e-12, 9.42221e-12, 2.21977e-11, 2.40218e-12, 3.53011e-12, 1.38442e-11, 3.26129e-11, 9.42221e-12, 1.38442e-11  …  1.38442e-11, 9.42221e-12, 3.26129e-11, 1.38442e-11, 3.53011e-12, 2.40218e-12, 2.21977e-11, 9.42221e-12, 2.40218e-12, 1.63452e-12]
 Float64-valued Field:
  [9.95014e-12, 1.40226e-11, 4.8862e-11, 1.08654e-10, 1.40226e-11, 1.97611e-11, 6.88498e-11, 1.53092e-10, 4.8862e-11, 6.88498e-11  …  6.88498e-11, 4.8862e-11, 1.53092e-10, 6.88498e-11, 1.97611e-11, 1.40226e-11, 1.08654e-10, 4.8862e-11, 1.40226e-11, 9.95014e-12]
 Float64-valued Field:
  [4.82014e-11, 6.55051e-11, 2.05863e-10, 4.34805e-10, 6.55051e-11, 8.90181e-11, 2.79733e-10, 5.90808e-10, 2.05863e-10, 2.79733e-10  …  2.79733e-10, 2.05863e-10, 5.90808e-10, 2.79733e-10, 8.90181e-11, 6.55051e-11, 4.34805e-10, 2.05863e-10, 6.55051e-11, 4.82014e-11]
 Float64-valued Field:
  [1.93409e-10, 2.54705e-10, 7.3086e-10, 1.47394e-9, 2.54705e-10, 3.35418e-10, 9.62408e-10, 1.94085e-9, 7.3086e-10, 9.62408e-10  …  9.62408e-10, 7.3086e-10, 1.94085e-9, 9.62408e-10, 3.35418e-10, 2.54705e-10, 1.47394e-9, 7.3086e-10, 2.54705e-10, 1.93409e-10]
 Float64-valued Field:
  [6.63222e-10, 8.4994e-10, 2.24969e-9, 4.3513e-9, 8.4994e-10, 1.08921e-9, 2.88289e-9, 5.57591e-9, 2.24969e-9, 2.88289e-9  …  2.88289e-9, 2.24969e-9, 5.57591e-9, 2.88289e-9, 1.08921e-9, 8.4994e-10, 4.3513e-9, 2.24969e-9, 8.4994e-10, 6.63222e-10]
 Float64-valued Field:
  [1.99234e-9, 2.49355e-9, 6.14103e-9, 1.1435e-8, 2.49355e-9, 3.12082e-9, 7.68562e-9, 1.4311e-8, 6.14103e-9, 7.68562e-9  …  7.68562e-9, 6.14103e-9, 1.4311e-8, 7.68562e-9, 3.12082e-9, 2.49355e-9, 1.1435e-8, 6.14103e-9, 2.49355e-9, 1.99234e-9]
 Float64-valued Field:
  [5.3484e-9, 6.55726e-9, 1.51373e-8, 2.72246e-8, 6.55726e-9, 8.03929e-9, 1.85581e-8, 3.33767e-8, 1.51373e-8, 1.85581e-8  …  1.85581e-8, 1.51373e-8, 3.33767e-8, 1.85581e-8, 8.03929e-9, 6.55726e-9, 2.72246e-8, 1.51373e-8, 6.55726e-9, 5.3484e-9]
anim = Plots.@animate for u in heat_cg_sol.u
    Plots.plot(u, c = :thermal)
end
mp4(anim)

3.3 Shallow water equations

The shallow water equations in vector invariant form can be written as

\[\begin{align*} \frac{\partial \rho}{\partial t} + \nabla \cdot (\rho u) &= 0\\ \frac{\partial u_i}{\partial t} + \nabla (\Phi + \tfrac{1}{2}\|u\|^2)_i &= J (u \times (\nabla \times u))_i \end{align*}\]

where $J$ is the Jacobian determinant, and $\Phi = g \rho$.

Note that the velocity $u$ is specified in covariant coordinates $u_i$.

For vizualization purposes, we can model a passive tracer $\theta$ as

\[\frac{\partial \rho \theta}{\partial t} + \nabla \cdot (\rho \theta u) = 0\]

using ClimaCore.Geometry

parameters = (
    ϵ = 0.1,  ## perturbation size for initial condition
    l = 0.5, ## Gaussian width
    k = 0.5, ## Sinusoidal wavenumber
    ρ₀ = 1.0, ## reference density
    c = 2,
    g = 10,
    D₄ = 1e-4, ## hyperdiffusion coefficient
)

function init_state(local_geometry, p)
    coord = local_geometry.coordinates
    (; x, y) = coord
    # set initial state
    ρ = p.ρ₀

    # set initial velocity
    U₁ = cosh(y)^(-2)

    # Ψ′ = exp(-(x2 + p.l / 10)^2 / 2p.l^2) * cos(p.k * x) * cos(p.k * y)
    # Vortical velocity fields (u₁′, u₂′) = (-∂²Ψ′, ∂¹Ψ′)
    ϕ = exp(-(y + p.l / 10)^2 / 2p.l^2)
    u₁′ = ϕ * (y + p.l / 10) / p.l^2 * cos(p.k * x) * cos(p.k * y)
    u₁′ += p.k * ϕ * cos(p.k * x) * sin(p.k * y)
    u₂′ = -p.k * ϕ * sin(p.k * x) * cos(p.k * y)

    u = Geometry.Covariant12Vector(
        Geometry.UVVector(U₁ + p.ϵ * u₁′, p.ϵ * u₂′),
        local_geometry,
    )

    # set initial tracer
    θ = sin(p.k * y)
    return (ρ = ρ, u = u, ρθ = ρ * θ)
end


y0 =
    init_state.(
        ClimaCore.Fields.local_geometry_field(rectangle_space),
        Ref(parameters),
    )

# plot initial tracer
Plots.plot(y0.ρθ)
Example block output
function shallow_water_tendency!(dydt, y, _, t)

    (; D₄, g) = parameters

    sdiv = ClimaCore.Operators.Divergence()
    wdiv = ClimaCore.Operators.WeakDivergence()
    grad = ClimaCore.Operators.Gradient()
    wgrad = ClimaCore.Operators.WeakGradient()
    curl = ClimaCore.Operators.Curl()
    wcurl = ClimaCore.Operators.WeakCurl()


    # compute hyperviscosity first
    @. dydt.u =
        wgrad(sdiv(y.u)) -
        Geometry.Covariant12Vector(wcurl(Geometry.Covariant3Vector(curl(y.u))))
    @. dydt.ρθ = wdiv(grad(y.ρθ))

    ClimaCore.Spaces.weighted_dss!(dydt)

    @. dydt.u =
        -D₄ * (
            wgrad(sdiv(dydt.u)) - Geometry.Covariant12Vector(
                wcurl(Geometry.Covariant3Vector(curl(dydt.u))),
            )
        )
    @. dydt.ρθ = -D₄ * wdiv(grad(dydt.ρθ))

    # comute rest of tendency
    @. begin
        dydt.ρ = -wdiv(y.ρ * y.u)
        dydt.u += -grad(g * y.ρ + norm(y.u)^2 / 2) + y.u × curl(y.u)
        dydt.ρθ += -wdiv(y.ρθ * y.u)
    end
    ClimaCore.Spaces.weighted_dss!(dydt)
    return dydt
end
shallow_water_tendency! (generic function with 1 method)
shallow_water_prob = ODEProblem(shallow_water_tendency!, y0, (0.0, 20.0))
@time shallow_water_sol =
    solve(shallow_water_prob, SSPRK33(), dt = 0.05, saveat = 1.0)
anim = Plots.@animate for u in shallow_water_sol.u
    Plots.plot(u.ρθ, clim = (-1, 1))
end
mp4(anim)

This page was generated using Literate.jl.