MultiRegion adaptation of the NonhydrostaticModel

examples/multi_regional_two_dimensional_turbulence.jl

@@ -0,0 +1,181 @@
+# # Two dimensional turbulence example
+#
+# In this example, we initialize a random velocity field and observe its turbulent decay
+# in a two-dimensional domain. This example demonstrates:
+#
+#   * How to run a model with no tracers and no buoyancy model.
+#   * How to use computed `Field`s to generate output.
+
+# ## Install dependencies
+#
+# First let's make sure we have all required packages installed.
+
+# ```julia
+# using Pkg
+# pkg"add Oceananigans, CairoMakie"
+# ```
+
+# ## Model setup
+
+# We instantiate the model with an isotropic diffusivity. We use a grid with 128² points,
+# a fifth-order advection scheme, third-order Runge-Kutta time-stepping,
+# and a small isotropic viscosity.  Note that we assign `Flat` to the `z` direction.
+
+using Oceananigans
+
+grid_base = RectilinearGrid(GPU(), size=(128, 128, 1), extent=(2π, 2π), topology=(Periodic, Periodic, Bounded))
+
+grid  = MultiRegionGrid(grid_base, devices = (0, 1))
+
+model = NonhydrostaticModel(; grid,
+                            timestepper = :RungeKutta3,
+                            advection = UpwindBiasedFifthOrder(),
+                            closure = ScalarDiffusivity(ν=1e-5))
+
+# ## Random initial conditions
+#
+# Our initial condition randomizes `model.velocities.u` and `model.velocities.v`.
+# We ensure that both have zero mean for aesthetic reasons.
+
+using Statistics
+
+u, v, w = model.velocities
+
+uᵢ = rand(128, 128, 1)
+vᵢ = rand(128, 128, 1)
+
+uᵢ .-= mean(uᵢ)
+vᵢ .-= mean(vᵢ)
+
+using Oceananigans.MultiRegion: multi_region_object_from_array
+
+u₀ = multi_region_object_from_array(uᵢ, grid)
+v₀ = multi_region_object_from_array(vᵢ, grid)
+
+set!(model, u=u₀, v=v₀)
+
+simulation = Simulation(model, Δt=0.2, stop_time=50)
+
+# ## Logging simulation progress
+#
+# We set up a callback that logs the simulation iteration and time every 100 iterations.
+
+progress(sim) = @info string("Iteration: ", iteration(sim), ", time: ", time(sim))
+simulation.callbacks[:progress] = Callback(progress, IterationInterval(100))
+
+# ## Output
+#
+# We set up an output writer for the simulation that saves vorticity and speed every 20 iterations.
+#
+# ### Computing vorticity and speed
+#
+# To make our equations prettier, we unpack `u`, `v`, and `w` from
+# the `NamedTuple` model.velocities:
+u, v, w = model.velocities
+
+# Next we create two `Field`s that calculate
+# _(i)_ vorticity that measures the rate at which the fluid rotates
+# and is defined as
+#
+# ```math
+# ω = ∂_x v - ∂_y u \, ,
+# ```
+
+ω = ∂x(v) - ∂y(u)
+
+# We also calculate _(ii)_ the _speed_ of the flow,
+#
+# ```math
+# s = \sqrt{u^2 + v^2} \, .
+# ```
+
+s = sqrt(u^2 + v^2)
+
+# We pass these operations to an output writer below to calculate and output them during the simulation.
+filename = "two_dimensional_turbulence"
+
+simulation.output_writers[:fields] = JLD2OutputWriter(model, (; ω, s),
+                                                      schedule = TimeInterval(0.6),
+                                                      filename = filename * ".jld2",
+                                                      overwrite_existing = true)
+
+# ## Running the simulation
+#
+# Pretty much just
+
+run!(simulation)
+
+# ## Visualizing the results
+#
+# We load the output.
+
+# Construct the ``x, y, z`` grid for plotting purposes,
+using Oceananigans.MultiRegion: reconstruct_global_field
+
+ω_global = reconstruct_global_field(Field(ω))
+s_global = reconstruct_global_field(Field(s))
+
+xω, yω, zω = nodes(ω_global)
+xs, ys, zs = nodes(s_global)
+nothing # hide
+
+ω_timeseries = []
+s_timeseries = []
+
+file = jldopen(filename * ".jld2")
+
+iterations = keys(file["timeseries/t"])
+
+for iter in iterations
+    push!(ω_timeseries, file["timeseries/ω/" * iter])
+    push!(s_timeseries, file["timeseries/s/" * iter])
+end
+
+# and animate the vorticity and fluid speed.
+
+using CairoMakie
+set_theme!(Theme(fontsize = 24))
+
+@info "Making a neat movie of vorticity and speed..."
+
+fig = Figure(resolution = (800, 500))
+
+axis_kwargs = (xlabel = "x",
+               ylabel = "y",
+               limits = ((0, 2π), (0, 2π)),
+               aspect = AxisAspect(1))
+
+ax_ω = Axis(fig[2, 1]; title = "Vorticity", axis_kwargs...)
+ax_s = Axis(fig[2, 2]; title = "Speed", axis_kwargs...)
+nothing #hide
+
+# We use Makie's `Observable` to animate the data. To dive into how `Observable`s work we
+# refer to [Makie.jl's Documentation](https://makie.juliaplots.org/stable/documentation/nodes/index.html).
+
+n = Observable(1)
+
+# Now let's plot the vorticity and speed.
+
+ω = @lift ω_timeseries[$n][:, :, 1]
+s = @lift s_timeseries[$n][:, :, 1]
+
+heatmap!(ax_ω, xω, yω, ω; colormap = :balance, colorrange = (-2, 2))
+heatmap!(ax_s, xs, ys, s; colormap = :speed, colorrange = (0, 0.2))
+
+# title = @lift "t = " * string(round(times[$n], digits=2))
+# Label(fig[1, 1:2], title, textsize=24, tellwidth=false)
+
+# Finally, we record a movie.
+
+frames = 1:length(s_timeseries)
+
+@info "Making a neat animation of vorticity and speed..."
+
+record(fig, filename * ".mp4", frames, framerate=24) do i
+    msg = string("Plotting frame ", i, " of ", frames[end])
+    print(msg * " \r")
+    n[] = i
+end
+nothing #hide
+
+# ![](two_dimensional_turbulence.mp4)

src/AbstractOperations/at.jl

@@ -58,49 +58,72 @@ indices(f::Function) = default_indices(3)
 indices(f::Number)   = default_indices(3)
 
 """
-    interpolate_indices(operands..., loc_operation = abstract_operation_location)
+    intersect_indices(loc, operands...)
 
-Utility to propagate operands' indices in `AbstractOperations`s with multiple operands 
-(`BinaryOperation`s and `MultiaryOperation`s).
+Utility to compute the intersection of `operands' indices.
 """
-function interpolate_indices(args...; loc_operation = (Center, Center, Center))
-    idxs = Any[:, :, :]
-    for i in 1:3
-        for arg in args
-            idxs[i] = interpolate_index(indices(arg)[i], idxs[i], location(arg)[i], loc_operation[i])
-        end
-    end
+function intersect_indices(loc, operands...)
 
-    return Tuple(idxs)
+    idx1 = compute_index_intersection(Colon(), loc[1], operands...; dim=1)
+    idx2 = compute_index_intersection(Colon(), loc[2], operands...; dim=2)
+    idx3 = compute_index_intersection(Colon(), loc[3], operands...; dim=3)
+
+    return (idx1, idx2, idx3)
 end
 
-interpolate_index(::Colon, ::Colon, args...)       = Colon()
-interpolate_index(::Colon, b::UnitRange, args...)  = b
+# Fallback for `KernelFunctionOperation`s with no argument 
+compute_index_intersection(::Colon, to_loc; kw...) = Colon()
 
-function interpolate_index(a::UnitRange, ::Colon, loc, new_loc)  
-    a = corrected_index(a, loc, new_loc)
+compute_index_intersection(to_idx, to_loc, op; dim) =
+    _compute_index_intersection(to_idx, indices(op, dim),
+                                to_loc, location(op, dim))
 
-    # Abstract operations that require an interpolation of a sliced fields are not supported!
-    first(a) > last(a) && throw(ArgumentError("Cannot interpolate a sliced field from $loc to $(new_loc)!"))
-    return a
+"""Compute index intersection recursively for `dim`ension ∈ (1, 2, 3)."""
+function compute_index_intersection(to_idx, to_loc, op1, op2, more_ops...; dim)
+    new_to_idx = _compute_index_intersection(to_idx, indices(op1, dim), to_loc, location(op1, dim))
+    return compute_index_intersection(new_to_idx, to_loc, op2, more_ops...; dim)
 end
 
-function interpolate_index(a::UnitRange, b::UnitRange, loc, new_loc)   
-    a = corrected_index(a, loc, new_loc)
+# Life is pretty simple in this case.
+_compute_index_intersection(to_idx::Colon, from_idx::Colon, args...) = Colon()
+
+# Because `from_idx` imposes no restrictions, we just return `to_idx`.
+_compute_index_intersection(to_idx::UnitRange, from_idx::Colon, args...) = to_idx
 
-    # Abstract operations that require an interpolation of a sliced fields are not supported!
-    first(a) > last(a) && throw(ArgumentError("Cannot interpolate a sliced field from $loc to $(new_loc)!"))
+# This time we account for the possible range-reducing effect of interpolation on `from_idx`.
+function _compute_index_intersection(to_idx::Colon, from_idx::UnitRange, to_loc, from_loc)
+    shifted_idx = restrict_index_for_interpolation(from_idx, from_loc, to_loc)
+    validate_shifted_index(shifted_idx)
+    return shifted_idx
+end
+
+# Compute the intersection of two index ranges
+function _compute_index_intersection(to_idx::UnitRange, from_idx::UnitRange, to_loc, from_loc)
+    shifted_idx = restrict_index_for_interpolation(from_idx, from_loc, to_loc)
+    validate_shifted_index(shifted_idx)
 
-    indices = UnitRange(max(first(a), first(b)), min(last(a), last(b)))
+    range_intersection = UnitRange(max(first(shifted_idx), first(to_idx)), min(last(shifted_idx), last(to_idx)))
 
-    # Abstract operations between parallel non-intersecating windowed fields are not supported
-    first(indices) > last(indices) && throw(ArgumentError("BinaryOperation operand indices $(a) and $(b) do not intersect!"))
-    return indices
+    # Check validity of the intersection index range
+    first(range_intersection) > last(range_intersection) &&
+        throw(ArgumentError("Indices $(from_idx) and $(to_idx) interpolated from $(from_loc) to $(to_loc) do not intersect!"))
+
+    return range_intersection
 end
 
-# Windowed fields interpolated from `Center`s to `Face`s lose the first index.
-# Viceverse, windowed fields interpolated from `Face`s to `Center`s lose the last index
-corrected_index(a, ::Type{Face},   ::Type{Face})   = UnitRange(first(a),   last(a))
-corrected_index(a, ::Type{Center}, ::Type{Center}) = UnitRange(first(a),   last(a))
-corrected_index(a, ::Type{Face},   ::Type{Center}) = UnitRange(first(a),   last(a)-1)
-corrected_index(a, ::Type{Center}, ::Type{Face})   = UnitRange(first(a)+1, last(a))
+validate_shifted_index(shifted_idx) = first(shifted_idx) > last(shifted_idx) &&
+    throw(ArgumentError("Cannot compute index intersection for indices $(from_idx) interpolating from $(from_loc) to $(to_loc)!"))
+
+"""
+    restrict_index_for_interpolation(from_idx, from_loc, to_loc)
+
+Return a "restricted" index range for the result of interpolating from
+`from_loc` to `to_loc`, over the index range `from_idx`:
+
+    * Windowed fields interpolated from `Center`s to `Face`s lose the first index.
+    * Conversely, windowed fields interpolated from `Face`s to `Center`s lose the last index
+"""
+restrict_index_for_interpolation(from_idx, ::Type{Face},   ::Type{Face})   = UnitRange(first(from_idx),   last(from_idx))
+restrict_index_for_interpolation(from_idx, ::Type{Center}, ::Type{Center}) = UnitRange(first(from_idx),   last(from_idx))
+restrict_index_for_interpolation(from_idx, ::Type{Face},   ::Type{Center}) = UnitRange(first(from_idx),   last(from_idx)-1)
+restrict_index_for_interpolation(from_idx, ::Type{Center}, ::Type{Face})   = UnitRange(first(from_idx)+1, last(from_idx))

src/AbstractOperations/binary_operations.jl

@@ -29,7 +29,7 @@ end
 # Recompute location of binary operation
 @inline at(loc, β::BinaryOperation) = β.op(loc, at(loc, β.a), at(loc, β.b))
 
-indices(β::BinaryOperation) = interpolate_indices(β.a, β.b; loc_operation = location(β))
+indices(β::BinaryOperation) = construct_regionally(intersect_indices, location(β), β.a, β.b)
 
 """Create a binary operation for `op` acting on `a` and `b` at `Lc`, where
 `a` and `b` have location `La` and `Lb`."""

src/AbstractOperations/computed_field.jl

@@ -44,7 +44,7 @@ function Field(operand::AbstractOperation;
     loc = location(operand)
     indices = validate_indices(indices, loc, grid)
 
-    boundary_conditions = FieldBoundaryConditions(indices, boundary_conditions)
+    @apply_regionally boundary_conditions = FieldBoundaryConditions(indices, boundary_conditions)
 
     if isnothing(data)
         data = new_data(grid, loc, indices)

src/AbstractOperations/kernel_function_operation.jl

@@ -62,7 +62,7 @@ function KernelFunctionOperation{LX, LY, LZ}(kernel_function, grid;
     return KernelFunctionOperation{LX, LY, LZ}(kernel_function, computed_dependencies, parameters, grid)
 end
 
-indices(κ::KernelFunctionOperation) = interpolate_indices(κ.computed_dependencies...; loc_operation = location(κ))
+indices(κ::KernelFunctionOperation) = construct_regionally(intersect_indices, location(κ), κ.computed_dependencies...)
 
 @inline Base.getindex(κ::KernelFunctionOperation, i, j, k) = κ.kernel_function(i, j, k, κ.grid, κ.computed_dependencies..., κ.parameters)
 @inline Base.getindex(κ::KernelFunctionOperation{LX, LY, LZ, <:Nothing}, i, j, k) where {LX, LY, LZ} = κ.kernel_function(i, j, k, κ.grid, κ.computed_dependencies...)

src/AbstractOperations/multiary_operations.jl

@@ -20,7 +20,7 @@ end
 ##### MultiaryOperation construction
 #####
 
-indices(Π::MultiaryOperation) = interpolate_indices(Π.args...; loc_operation = location(Π))
+indices(Π::MultiaryOperation) = construct_regionally(intersect_indices, location(Π), Π.args...)
 
 function _multiary_operation(L, op, args, Largs, grid)
     ▶ = Tuple(interpolation_operator(La, L) for La in Largs)

src/Architectures.jl

@@ -14,10 +14,10 @@ using OffsetArrays
 import CUDAKernels: next_stream
 
 if CUDA.has_cuda_gpu()     
-using CUDAKernels: STREAM_GC_LOCK
+    using CUDAKernels: STREAM_GC_LOCK
 
-    DEVICE_FREE_STREAMS = Tuple(CUDA.CuStream[] for dev in 1:length(CUDA.devices()))
-    DEVICE_STREAMS      = Tuple(CUDA.CuStream[] for dev in 1:length(CUDA.devices()))
+    DEVICE_FREE_STREAMS = Tuple(CUDA.CuStream[] for dev in 1:CUDA.ndevices())
+    DEVICE_STREAMS      = Tuple(CUDA.CuStream[] for dev in 1:CUDA.ndevices())
     const DEVICE_STREAM_GC_THRESHOLD = Ref{Int}(16)
 
     function next_stream()
@@ -123,7 +123,7 @@ function unified_array(::GPU, arr::AbstractArray)
     return vec
 end
 
-## Only for contiguous data!! (i.e. only if the offset for pointer(dst::CuArrat, offset::Int) is 1)
+## Only for contiguous data!! (i.e. only if the offset for pointer(dst::CuArray, offset::Int) is 1)
 @inline function device_copy_to!(dst::CuArray, src::CuArray; async::Bool = false) 
     n = length(src)
     context!(context(src)) do
@@ -133,8 +133,30 @@ end
     end
     return dst
 end
-
+
+@inline function device_copy_to!(dst::SubArray{T, N, <:CuArray}, src::SubArray{T, N, <:CuArray}; async::Bool = false) where {T, N}
+    dst_parent = parent(dst)
+    src_parent = parent(src)
+
+    n = size(src, 1) 
+
+    D = size(dst_parent)
+    S = size(src_parent)
+
+    context!(context(src_parent)) do
+        GC.@preserve src_parent dst_parent begin
+            for (jd, js) in zip(dst.indices[2], src.indices[2]), (kd, ks) in zip(dst.indices[3], src.indices[3])
+                dst_ptr = first(dst.indices[1]) + D[1] * (jd - 1 + D[2] * (kd - 1))
+                src_ptr = first(src.indices[1]) + S[1] * (js - 1 + S[2] * (ks - 1))
+                unsafe_copyto!(pointer(dst_parent, dst_ptr), pointer(src_parent, src_ptr), n; async)
+            end
+        end
+    end
+    return dst
+end
+
 @inline device_copy_to!(dst::Array, src::Array; kw...) = Base.copyto!(dst, src)
+@inline device_copy_to!(dst::SubArray{T, N, <:Array}, src::SubArray{T, N, <:Array}; async::Bool = false) where {T, N} = Base.copyto!(dst, src)
 
 device_event(arch) = Event(device(arch))
 

src/Fields/abstract_field.jl

@@ -37,7 +37,8 @@ Base.IndexStyle(::AbstractField) = IndexCartesian()
 "Returns the location `(LX, LY, LZ)` of an `AbstractField{LX, LY, LZ}`."
 @inline location(a) = (Nothing, Nothing, Nothing) # used in AbstractOperations for location inference
 @inline location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX, LY, LZ) # note no instantiation
-@inline location(f::AbstractField, i) = location(f)[i]
+@inline location(f, i) = location(f)[i]
+
 @inline instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ())
 
 "Returns the architecture of on which `f` is defined."