Skip to content

Commit

Permalink
Finishing touches
Browse files Browse the repository at this point in the history
  • Loading branch information
@jlk9
jlk9 committed Jun 11, 2024
1 parent a349b4a commit 96643aa
Showing 1 changed file with 29 additions and 29 deletions.
58 changes: 29 additions & 29 deletions ClimaIceOcean_tutorial/ice_ocean_notebook_tutorial.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -55,11 +55,13 @@ end
md"# Ocean Sea-Ice Inverse Problem"

# ╔═╡ 32388534-3ac7-497d-a933-086100ff0c20
md"As a (very!) simplified model, suppose we have a circlular floe of ice floating on top of a body of water. The water has a steady horizontal current represented by functions $u$ and $v$, and its initial temperature $T$ varies depending on depth and horizontal location. Both diffusion and advection will affect the water's temperature, which will in turn cause temperature changes and encourage growth or melt in the ice.
md"In this tutorial we'll set up and solve a basic inverse problem - reconstructing initial conditions in a coupled ocean and ice model - using a few packages available in the Julia programming language including Enzyme for automatic differentiation. Every part of this process will be modular, giving you, the user, flexibility to change it to fit your needs.
We'll implement this model using two Julia packages: `Oceananigans.jl` for ocean-flavored fluid dynamics, and `ClimaSeaIce.jl` for ice thermodynamics and dynamics. Both of these packages are designed for use with GPUs in mind, but for this tutorial we will only use CPUs since our problem will be fairly small. Both of these packages were written by scientists at the Climate Modeling Alliance (CliMA: https://clima.caltech.edu), and thus use very similar conventions.
As a (very!) simplified model, suppose we have a circular floe of ice floating on top of a body of water. The water has a steady-state horizontal current, and its initial temperature $T$ varies depending on depth and horizontal location. Both diffusion and advection will affect the water's temperature, which will in turn cause temperature changes and encourage growth or melt in the ice.
To differentiate this model, we'll use `Enzyme`, a tool that performs automatic differentiation (AD) of code. `Enzyme` operates at the LLVM level (low-level virtual machine), which makes it usable for several programming languages. Here we'll use its Julia bindings, from the package `Enzyme.jl`.
We'll implement this model using two Julia packages written by scientists at the Climate Modeling Alliance (CliMA: https://clima.caltech.edu): `Oceananigans.jl` for ocean-flavored fluid dynamics, and `ClimaSeaIce.jl` for ice thermodynamics and dynamics. Both of these packages are designed for use with GPUs in mind, but for this tutorial we will only use CPUs since our problem will be fairly small. They use a lot of the same conventions.
To differentiate this model, we'll use `Enzyme`, a tool that performs automatic differentiation (AD) of code. `Enzyme` operates at the LLVM level, which makes it usable for several programming languages. Here we'll use its Julia bindings, from the package `Enzyme.jl`.
The following block activates an environment using Oceananigans, ClimaSeaIce, and Enzyme - for this tutorial we'll use specific versions of each. You might also see references in the output to `KernelAbstractions.jl`, a library for writing code kernels that can be run on the CPU or GPU (both Oceananigans and ClimaSeaIce use this) as well as a file named `ice_ocean_interaction.jl`, which includes a helper function that implements energy transfer between water and ice."

Expand Down Expand Up @@ -103,7 +105,7 @@ begin
end

# ╔═╡ 9637e4af-6176-490d-9c96-c2d3b2a7b32d
md"Here we can see key information for both our ocean grid (labelled just `grid`) and our ice grid. Note the ocean grid is of size 8 $\times$ 8 $\times$ 16, while the ice grid is of size 8 $\times$ 8 $\times$ 1. We set the domain the be $[-\pi,\pi] \times [-\pi,\pi] \times [-0.5, 0.5]$, but any spatial coordinates can be used. All domain boundaries are either bounded or flat (in the case of the ice vertical dimension)."
md"Here we can see key information for both our ocean grid (labelled just `grid`) and our ice grid. Note the ocean grid is of size 16 $\times$ 16 $\times$ 16, while the ice grid is of size 16 $\times$ 16 $\times$ 1. We set the domain the be $[-\pi,\pi] \times [-\pi,\pi] \times [-1.0, 0]$, but any spatial coordinates can be used. All domain boundaries are either bounded or flat (in the case of the ice vertical dimension)."

# ╔═╡ 31ea7ca7-9e45-4722-9282-a636823f9a4e
@show grid
Expand All @@ -114,7 +116,7 @@ md"Here we can see key information for both our ocean grid (labelled just `grid`
# ╔═╡ 1819e8b4-2902-4d43-95ea-e2748513ba6b
md"### Step 2: Set up ocean and sea ice model objects
With our grids set up, we can now focus on the other aspects of the model. Our ocean model needs a diffusivity to handle the temperature tracer diffusing through the water - we'll use a vertical scalar diffusivity for this. We label this choice `diffusion`.
With our grids set up, we can now focus on the other aspects of the model. Our ocean model needs a turbulence closure to handle the effect of viscous dissipation and diffusion - we'll use a constant isotropic diffusivity for this, called `VerticalScalarDiffusivity` in Oceananigans.
We also need a velocity field for our water body. Here we'll use a time-invariant function, showing the water flows at a consistent speed and direction over time. We name our $x$-direction velocity `u` and our $y$-direction velocity `v`. We also isolate the ocean surface velocities with the `view` function so they can interact with the ice (`view` doesn't create a new instance of the array in its argument - it gives you an additional reference to that array that only tracks the slice passed to it).
Expand All @@ -128,15 +130,16 @@ Note that all of these variables use our `grid` or `ice_grid` objects, to determ
begin
# Then we set a maximal diffusivity and diffusion type:
const maximum_diffusivity = 100
diffusion = VerticalScalarDiffusivity(κ=0.1)
κ = 0.1
diffusion = VerticalScalarDiffusivity(κ=κ)

# For this problem we'll have a constant-values velocity field - the water will flow at a
# constant speed and direction over time. We have to assign the proper velocity values to
# every point in our grid:
u = XFaceField(grid)
v = YFaceField(grid)

U = 40
U = 2
u₀(x, y, z) = - U * cos(x + π/4) * sin(y)
v₀(x, y, z) = + U * 0.5

Expand All @@ -159,9 +162,13 @@ end
# ╔═╡ dd2b04b7-d35d-42ad-8887-206da0576688
md"#### With our required variables initialized, we can now create our ocean and ice `models`!
Our ocean `model` uses the hydrostatic approximation - hence we construct a `HydrostaticFreeSurfaceModel`. In addition to supplying the diffusivity and velocity fields above, we also add a tracer `T` to represent our temperature, and the WENO advection scheme to describe how temperature is moved by the ocean velocity field. Since we're assuming a constant salinity in our water, we won't track salinity as an additional tracer.
Our ocean `model` uses two common approximations:
- (*Hydrostatic*) The pressure of water at any point is only due to the weight of the water above it.
- (*Boussinesq*) Ignore density differences in the water except when a term is multiplied by $g$.
These are both used in Oceananigans' `HydrostaticFreeSurfaceModel`. In addition to supplying the diffusivity and velocity fields above, we also add a tracer `T` to represent our temperature, and the WENO advection scheme to describe how temperature is moved by the ocean velocity field. Since we're assuming a constant salinity in our water, we won't track salinity as an additional tracer.
Our ice `model` is a simple slab model with only one layer (it's often referred to as a zero-layer model, since most other ice models represent internal temperature or energy). In addition to the fluxes above, we also supply a starting salinity, internal heat flux, and top heat flux and boundary condition."
Our ice `model` is a simple zero-layer slab model (zero-layer since most other ice models represent internal temperature or energy). In addition to the fluxes above, we also supply a starting salinity, internal heat flux, and top heat flux and boundary condition."

# ╔═╡ a3efa97e-467b-43a2-87db-8c3bdc251d25
ocean_model = HydrostaticFreeSurfaceModel(; grid,
Expand Down Expand Up @@ -235,9 +242,7 @@ begin

# Do time-stepping
Nx, Ny, Nz = size(ocean_model.grid)
κ_max = maximum_diffusivity
Δz = 2π / Nz
Δt = 1e-1 * Δz^2 / κ_max
Δt = 0.0015
@show Δt

ocean_model.clock.time = 0
Expand Down Expand Up @@ -272,13 +277,10 @@ begin
end

# ╔═╡ 62b3adee-63b0-4f05-b304-d1d8b0d40ef9
md"Let's set diffusivity $\kappa = 1$ and our number of forward time steps $n_{\text{max}} = 100$:"
md"We set diffusivity $\kappa = 0.1$. Let our number of forward time steps $n_{\text{max}} = 100$:"

# ╔═╡ 7cc53014-2395-4357-bbbb-d2d7eff59e20
begin
κ = 1
n_max = 100
end
n_max = 100

# ╔═╡ 0b04432b-0e5d-4d99-9808-0a1ad7765f72
md"Then we can call `ice_ocean_data` and get the true initial water temperature and ice thickness $T_0$ and $h_0$, as well as the true final temperature and thickness $T_n$ and $h_n$. We can use the `@show` macro on each of these fields to see their shapes (which should align with their corresponding grids) and some statistics, and also plot them."
Expand Down Expand Up @@ -344,12 +346,12 @@ begin

axH0 = Axis(figh[1, 1], xlabel="x (m)", ylabel="y (m)", title="Initial ice thickness")
axHn = Axis(figh[1, 2], xlabel="x (m)", ylabel="y (m)", title="Final ice thickness")
Colorbar(figh[1, 3], limits = ice_color_range, colormap = :grays,
Colorbar(figh[1, 3], limits = ice_color_range, colormap = :viridis,
label = "Ice thickness (m)")


heatmap!(axH0, xpoints, ypoints, h₀.data[1:Nx,1:Ny], colormap=:grays, colorrange = ice_color_range)
heatmap!(axHn, xpoints, ypoints, hₙ.data[1:Nx,1:Ny], colormap=:grays, colorrange = ice_color_range)
heatmap!(axH0, xpoints, ypoints, h₀.data[1:Nx,1:Ny], colormap=:viridis, colorrange = ice_color_range)
heatmap!(axHn, xpoints, ypoints, hₙ.data[1:Nx,1:Ny], colormap=:viridis, colorrange = ice_color_range)

figh
end
Expand Down Expand Up @@ -379,9 +381,7 @@ begin

# Run the forward model:
Nx, Ny, Nz = size(ocean_model.grid)
κ_max = maximum_diffusivity
Δz = 2π / Nz
Δt = 1e-1 * Δz^2 / κ_max
Δt = 0.0015

ocean_model.clock.time = 0
ocean_model.clock.iteration = 0
Expand Down Expand Up @@ -532,12 +532,12 @@ begin

axH0i = Axis(fighi[1, 1], xlabel="x (m)", ylabel="y (m)", title="True final ice thickness")
axHi = Axis(fighi[1, 2], xlabel="x (m)", ylabel="y (m)", title="Inverted final ice thickness")
Colorbar(fighi[1, 3], limits = ice_color_range, colormap = :grays,
Colorbar(fighi[1, 3], limits = ice_color_range, colormap = :viridis,
label = "Ice thickness (m)")


heatmap!(axH0i, xpoints, ypoints, hₙ.data[1:Nx,1:Ny], colormap=:grays, colorrange = ice_color_range)
heatmap!(axHi, xpoints, ypoints, ice_model.ice_thickness.data[1:Nx,1:Ny], colormap=:grays, colorrange = ice_color_range)
heatmap!(axH0i, xpoints, ypoints, hₙ.data[1:Nx,1:Ny], colormap=:viridis, colorrange = ice_color_range)
heatmap!(axHi, xpoints, ypoints, ice_model.ice_thickness.data[1:Nx,1:Ny], colormap=:viridis, colorrange = ice_color_range)

fighi
end
Expand Down Expand Up @@ -593,7 +593,7 @@ But even with some limitations, this tutorial outlines a basic workflow for solv
# ╟─9afd85ee-b9de-4629-b9a8-3ce6ea0f10db
# ╟─3adeeda5-7db4-4d38-a0c4-2adee41c14e8
# ╠═071a8881-0be3-4052-9c28-bc76057b6b5a
# ╟─9637e4af-6176-490d-9c96-c2d3b2a7b32d
# ╠═9637e4af-6176-490d-9c96-c2d3b2a7b32d
# ╠═31ea7ca7-9e45-4722-9282-a636823f9a4e
# ╠═a30d3476-fe00-4e2c-a786-8673a64bc8cf
# ╟─1819e8b4-2902-4d43-95ea-e2748513ba6b
Expand All @@ -603,7 +603,7 @@ But even with some limitations, this tutorial outlines a basic workflow for solv
# ╠═63f3f573-d12f-4c22-aeec-275c33750ab9
# ╟─f8194371-97d4-45f4-8441-309bc535302c
# ╠═fc39f605-9635-402d-b60a-1c8c15a82c89
# ╟─62b3adee-63b0-4f05-b304-d1d8b0d40ef9
# ╠═62b3adee-63b0-4f05-b304-d1d8b0d40ef9
# ╠═7cc53014-2395-4357-bbbb-d2d7eff59e20
# ╠═0b04432b-0e5d-4d99-9808-0a1ad7765f72
# ╠═045b79be-be33-4f17-b142-5f5b82c9e1f1
Expand All @@ -625,6 +625,6 @@ But even with some limitations, this tutorial outlines a basic workflow for solv
# ╟─934fb784-be1b-4b9d-ae5e-57ad07a38cfd
# ╟─0bfdd29e-6ab3-406c-9109-40422d059f74
# ╟─222416ae-1aa0-4f48-9e8d-3aabc34e21dd
# ╟─8ac54b9e-7a50-49ce-99ad-60244bb12c9d
# ╠═8ac54b9e-7a50-49ce-99ad-60244bb12c9d
# ╟─9915a338-69d2-4fac-98d0-737bdfa69543
# ╟─ca0ca8a0-1692-4cc0-8726-26de146051b7

0 comments on commit 96643aa

Please sign in to comment.