From 3b8b9ce4db68750fec47e220f49d8f621191ada6 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 19 Dec 2024 22:18:46 +0100 Subject: [PATCH 01/31] Disambiguate `diffusive_flux_*` methods for `NoDiffusionISSD` (#4016) * disambiguate methods * ok let's go --- .../isopycnal_skew_symmetric_diffusivity.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl index 2a6b2f676b..aaedf4545e 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl @@ -212,9 +212,9 @@ end end # Make sure we do not need to perform heavy calculations if we really do not need to -@inline diffusive_flux_x(i, j, k, grid, ::NoDiffusionISSD, args...) = zero(grid) -@inline diffusive_flux_y(i, j, k, grid, ::NoDiffusionISSD, args...) = zero(grid) -@inline diffusive_flux_z(i, j, k, grid, ::NoDiffusionISSD, args...) = zero(grid) +@inline diffusive_flux_x(i, j, k, grid, ::NoDiffusionISSD, K, ::Val{tracer_index}, args...) where tracer_index = zero(grid) +@inline diffusive_flux_y(i, j, k, grid, ::NoDiffusionISSD, K, ::Val{tracer_index}, args...) where tracer_index = zero(grid) +@inline diffusive_flux_z(i, j, k, grid, ::NoDiffusionISSD, K, ::Val{tracer_index}, args...) where tracer_index = zero(grid) # Diffusive fluxes @inline get_tracer_κ(κ::NamedTuple, grid, tracer_index) = @inbounds κ[tracer_index] From 0cf7878f96cf4267cb93572611b977aa6080f213 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 19 Dec 2024 14:19:03 -0700 Subject: [PATCH 02/31] Bump actions/cache from 1 to 4 (#3995) Bumps [actions/cache](https://github.com/actions/cache) from 1 to 4. - [Release notes](https://github.com/actions/cache/releases) - [Changelog](https://github.com/actions/cache/blob/main/RELEASES.md) - [Commits](https://github.com/actions/cache/compare/v1...v4) --- updated-dependencies: - dependency-name: actions/cache dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Simone Silvestri --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) mode change 100755 => 100644 .github/workflows/ci.yml diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml old mode 100755 new mode 100644 index 99bbd42978..c9fc7c4260 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -20,7 +20,7 @@ jobs: with: version: ${{ matrix.version }} arch: ${{ matrix.arch }} - - uses: actions/cache@v1 + - uses: actions/cache@v4 env: cache-name: cache-artifacts with: From 46a38a7e280145e3960b3127bc582df0063f7b93 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 19 Dec 2024 14:19:18 -0700 Subject: [PATCH 03/31] Bump julia-actions/setup-julia from 1 to 2 (#3996) Bumps [julia-actions/setup-julia](https://github.com/julia-actions/setup-julia) from 1 to 2. - [Release notes](https://github.com/julia-actions/setup-julia/releases) - [Changelog](https://github.com/julia-actions/setup-julia/blob/master/devdocs/making_a_new_release.md) - [Commits](https://github.com/julia-actions/setup-julia/compare/v1...v2) --- updated-dependencies: - dependency-name: julia-actions/setup-julia dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Simone Silvestri --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c9fc7c4260..f786218b00 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,7 +16,7 @@ jobs: - x64 steps: - uses: actions/checkout@v2 - - uses: julia-actions/setup-julia@v1 + - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} arch: ${{ matrix.arch }} From 325866be76d42fe4190b8de147b48d7632223fc7 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Fri, 20 Dec 2024 09:02:51 +0100 Subject: [PATCH 04/31] Bump actions/checkout from 2 to 4 (#3994) Bumps [actions/checkout](https://github.com/actions/checkout) from 2 to 4. - [Release notes](https://github.com/actions/checkout/releases) - [Changelog](https://github.com/actions/checkout/blob/main/CHANGELOG.md) - [Commits](https://github.com/actions/checkout/compare/v2...v4) --- updated-dependencies: - dependency-name: actions/checkout dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f786218b00..7cb2227a84 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -15,7 +15,7 @@ jobs: arch: - x64 steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v2 with: version: ${{ matrix.version }} From da6ee8745fda308d75a7d8388a47992ebb7f1cc7 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Fri, 20 Dec 2024 14:03:18 -0700 Subject: [PATCH 05/31] More JLD2 shenanigans for reconstructing FieldTimeSeries (#4020) * More JLD2 shenanigans for reconstructing FieldTimeSeries * Delete reconstructed_grid_type for now --- src/OutputReaders/field_time_series.jl | 33 +++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/src/OutputReaders/field_time_series.jl b/src/OutputReaders/field_time_series.jl index 6967d845a3..9d77c161c3 100644 --- a/src/OutputReaders/field_time_series.jl +++ b/src/OutputReaders/field_time_series.jl @@ -392,6 +392,31 @@ end isreconstructed(grid::JLD2.ReconstructedStatic) = true isreconstructed(grid::AbstractGrid) = false isreconstructed(grid::ImmersedBoundaryGrid) = isreconstructed(grid.underlying_grid) +reconstructed_name(::JLD2.ReconstructedStatic{N}) where N = string(N) + +function reconstructed_topology(grid::JLD2.ReconstructedStatic) + name = reconstructed_name(grid) + firstcurly = findfirst('{', name) + grid_type = name[1:firstcurly-1] + + type_parameters = name[firstcurly+1:end-1] + parameter_list = split(type_parameters, ',') + + FTstr = parameter_list[1] + TXstr = parameter_list[2] + TYstr = parameter_list[3] + TZstr = parameter_list[4] + + TXsym = Symbol(TXstr) + TYsym = Symbol(TYstr) + TZsym = Symbol(TZstr) + + TX = eval(:($(TXsym))) + TY = eval(:($(TYsym))) + TZ = eval(:($(TZsym))) + + return (TX, TY, TZ) +end """ FieldTimeSeries{LX, LY, LZ}(grid::AbstractGrid [, times=()]; kwargs...) @@ -509,7 +534,13 @@ function FieldTimeSeries(path::String, name::String; Hz = file["$address/Hz"] zᵃᵃᶠ = file["$address/zᵃᵃᶠ"] z = file["$address/Δzᵃᵃᶠ"] isa Number ? (zᵃᵃᶠ[1], zᵃᵃᶠ[Nz+1]) : zᵃᵃᶠ[1:Nz+1] - topo = topology(grid) + + if isibg + topo = topology(grid) + else + topo = reconstructed_topology(grid) + end + size = (Nx, Ny, Nz) halo = (Hx, Hy, Hz) From d8273e7a07e8f53f9772f87f67f6eca8d17cb98c Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Sun, 22 Dec 2024 11:31:21 +1100 Subject: [PATCH 06/31] add paper + urge users for julia 1.10 --- docs/src/index.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index 628b4eebf1..4ac15e0404 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -24,8 +24,8 @@ julia> using Pkg julia> Pkg.add("Oceananigans") ``` -!!! compat "Julia 1.9 is required" - Oceananigans requires Julia 1.9 or later. +!!! compat "Julia 1.9 or later is required; Julia 1.10 is suggested" + Oceananigans requires Julia 1.9 or later. We strongly urge you to use Julia 1.10. If you're [new to Julia](https://docs.julialang.org/en/v1/manual/getting-started/) and its [wonderful `Pkg` manager](https://docs.julialang.org/en/v1/stdlib/Pkg/), the [Oceananigans wiki](https://github.com/CliMA/Oceananigans.jl/wiki) provides [more detailed installation instructions](https://github.com/CliMA/Oceananigans.jl/wiki/Installation-and-getting-started-with-Oceananigans). @@ -79,7 +79,9 @@ cite our work and mention Oceananigans by name. If you have work using Oceananigans that you would like to have listed here, please open a pull request to add it or let us know! -1. Bisits, J. I., Zika, J. D., and Evans, D. G. (2024) [Does cabbeling shape the thermohaline structure of high-latitude oceans?](https://doi.org/10.1175/JPO-D-24-0061.1) _Journal of Physical Oceanography_, in press. DOI: [10.1175/JPO-D-24-0061.1](https://doi.org/10.1175/JPO-D-24-0061.1) +1. Abbott, K. and Mahadevan, A. (2024). [Why is the monsoon coastal upwelling signal subdued in the Bay of Bengal?](https://doi.org/10.1029/2024JC022023), _Journal of Geophysical Research: Oceans_, **129**, e2024JC022023. DOI: [10.1029/2024JC022023](https://doi.org/10.1029/2024JC022023) + +1. Bisits, J. I., Zika, J. D., and Evans, D. G. (2024) [Does cabbeling shape the thermohaline structure of high-latitude oceans?](https://doi.org/10.1175/JPO-D-24-0061.1), _Journal of Physical Oceanography_, in press. DOI: [10.1175/JPO-D-24-0061.1](https://doi.org/10.1175/JPO-D-24-0061.1) 1. Strong-Wright J. and Taylor, J. R. (2024) [A model of tidal flow and tracer release in a giant kelp forest](https://doi.org/10.1017/flo.2024.13), _Flow_, **4**, E21. DOI: [10.1017/flo.2024.13](https://doi.org/10.1017/flo.2024.13) From 4669c0a8495cff5b5706ef31e49155d9fb40318a Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 24 Dec 2024 03:09:09 +0100 Subject: [PATCH 07/31] Support interpolating onto `Windowed` fields (#4021) * support interpolating * fix params * correct interior_view_indices + add a constructor * import * probably better like this * disambiguate * remove space * a little more verbose * bugfix * add a test * bugfix * this should work now * this should work --- src/AbstractOperations/at.jl | 72 ++++++---------------------- src/Fields/Fields.jl | 1 + src/Fields/abstract_field.jl | 4 +- src/Fields/field.jl | 13 ++--- src/Fields/field_indices.jl | 89 +++++++++++++++++++++++++++++++++++ src/Fields/interpolate.jl | 4 +- src/Utils/kernel_launching.jl | 3 ++ test/test_field.jl | 10 ++++ 8 files changed, 127 insertions(+), 69 deletions(-) create mode 100644 src/Fields/field_indices.jl diff --git a/src/AbstractOperations/at.jl b/src/AbstractOperations/at.jl index 18157fcab4..c39e90e77b 100644 --- a/src/AbstractOperations/at.jl +++ b/src/AbstractOperations/at.jl @@ -1,5 +1,5 @@ using Oceananigans -using Oceananigans.Fields: default_indices +using Oceananigans.Fields: default_indices, compute_index_intersection """ insert_location(ex::Expr, location) @@ -55,8 +55,8 @@ macro at(location, abstract_operation) end # Numbers and functions do not have index restrictions -indices(f::Function) = default_indices(3) -indices(f::Number) = default_indices(3) +indices(::Function) = default_indices(3) +indices(::Number) = default_indices(3) """ intersect_indices(loc, operands...) @@ -65,66 +65,22 @@ Utility to compute the intersection of `operands' indices. """ function intersect_indices(loc, operands...) - 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) + idx1 = compute_operand_intersection(Colon(), loc[1], operands...; dim=1) + idx2 = compute_operand_intersection(Colon(), loc[2], operands...; dim=2) + idx3 = compute_operand_intersection(Colon(), loc[3], operands...; dim=3) return (idx1, idx2, idx3) end # Fallback for `KernelFunctionOperation`s with no argument -compute_index_intersection(::Colon, to_loc; kw...) = Colon() +compute_operand_intersection(::Colon, to_loc; kw...) = Colon() -compute_index_intersection(to_idx, to_loc, op; dim) = - _compute_index_intersection(to_idx, indices(op)[dim], - to_loc, location(op, dim)) +compute_operand_intersection(to_idx, to_loc, op; dim) = + compute_index_intersection(to_idx, indices(op)[dim], + to_loc, location(op, dim)) """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 - -# 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 - -# 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) - - range_intersection = UnitRange(max(first(shifted_idx), first(to_idx)), min(last(shifted_idx), last(to_idx))) - - # 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 - -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)) +function compute_operand_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_operand_intersection(new_to_idx, to_loc, op2, more_ops...; dim) +end \ No newline at end of file diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index f7f1130141..f1a7e3db18 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -27,6 +27,7 @@ include("constant_field.jl") include("function_field.jl") include("field_boundary_buffers.jl") include("field.jl") +include("field_indices.jl") include("scans.jl") include("regridding_fields.jl") include("field_tuples.jl") diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl index 72a8e78540..bdc5add9e1 100644 --- a/src/Fields/abstract_field.jl +++ b/src/Fields/abstract_field.jl @@ -6,11 +6,11 @@ using Statistics using Oceananigans.Architectures using Oceananigans.Utils -using Oceananigans.Grids: interior_indices, interior_parent_indices +using Oceananigans.Grids: interior_parent_indices import Base: minimum, maximum, extrema import Oceananigans.Architectures: architecture, child_architecture -import Oceananigans.Grids: interior_x_indices, interior_y_indices, interior_z_indices +import Oceananigans.Grids: interior_x_indices, interior_y_indices, interior_z_indices, interior_indices import Oceananigans.Grids: total_size, topology, nodes, xnodes, ynodes, znodes, rnodes, node, xnode, ynode, znode, rnode import Oceananigans.Utils: datatuple diff --git a/src/Fields/field.jl b/src/Fields/field.jl index f1aabf32b0..ef53e387ee 100644 --- a/src/Fields/field.jl +++ b/src/Fields/field.jl @@ -360,18 +360,14 @@ boundary_conditions(not_field) = nothing immersed_boundary_condition(f::Field) = f.boundary_conditions.immersed data(field::Field) = field.data -indices(obj, i=default_indices(3)) = i -indices(f::Field, i=default_indices(3)) = f.indices -indices(a::SubArray, i=default_indices(ndims(a))) = a.indices -indices(a::OffsetArray, i=default_indices(ndims(a))) = indices(parent(a), i) - -"""Return indices that create a `view` over the interior of a Field.""" -interior_view_indices(field_indices, interior_indices) = Colon() -interior_view_indices(::Colon, interior_indices) = interior_indices instantiate(T::Type) = T() instantiate(t) = t +"""Return indices that create a `view` over the interior of a Field.""" +interior_view_indices(field_indices, interior_indices) = Colon() +interior_view_indices(::Colon, interior_indices) = interior_indices + function interior(a::OffsetArray, Loc::Tuple, Topo::Tuple, @@ -383,6 +379,7 @@ function interior(a::OffsetArray, topo = map(instantiate, Topo) i_interior = map(interior_parent_indices, loc, topo, sz, halo_sz) i_view = map(interior_view_indices, ind, i_interior) + return view(parent(a), i_view...) end diff --git a/src/Fields/field_indices.jl b/src/Fields/field_indices.jl new file mode 100644 index 0000000000..544561a7a9 --- /dev/null +++ b/src/Fields/field_indices.jl @@ -0,0 +1,89 @@ +indices(obj, i=default_indices(3)) = i +indices(f::Field, i=default_indices(3)) = f.indices +indices(a::SubArray, i=default_indices(ndims(a))) = a.indices +indices(a::OffsetArray, i=default_indices(ndims(a))) = indices(parent(a), i) + +function interior_x_indices(f::Field) + loc = map(instantiate, location(f)) + interior_indices = interior_x_indices(f.grid, loc) + return compute_index_intersection(interior_indices, f.indices[1]) +end + +function interior_y_indices(f::Field) + loc = map(instantiate, location(f)) + interior_indices = interior_y_indices(f.grid, loc) + return compute_index_intersection(interior_indices, f.indices[2]) +end + +function interior_z_indices(f::Field) + loc = map(instantiate, location(f)) + interior_indices = interior_z_indices(f.grid, loc) + return compute_index_intersection(interior_indices, f.indices[3]) +end + +# Interior indices for a field with a given location and topology +function interior_indices(f::Field) + ind_x = interior_x_indices(f) + ind_y = interior_y_indices(f) + ind_z = interior_z_indices(f) + return (ind_x, ind_y, ind_z) +end + +# 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 + +# In case of no locations specified, Because `to_idx` imposes no restrictions, we just return `from_idx`. +compute_index_intersection(to_idx::Colon, from_idx::UnitRange) = from_idx + +# 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_on_location(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_on_location(from_idx, from_loc, to_loc) + validate_shifted_index(shifted_idx) + + range_intersection = UnitRange(max(first(shifted_idx), first(to_idx)), min(last(shifted_idx), last(to_idx))) + + # 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 + +# Compute the intersection of two index ranges where the location is the same +function compute_index_intersection(to_idx::UnitRange, from_idx::UnitRange) + range_intersection = UnitRange(max(first(from_idx), first(to_idx)), + min(last(from_idx), last(to_idx))) + + # Check validity of the intersection index range + first(range_intersection) > last(range_intersection) && + throw(ArgumentError("Indices $(from_idx) and $(to_idx) do not intersect!")) + + return range_intersection +end + +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_on_location(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_on_location(from_idx, ::Type{Face}, ::Type{Face}) = UnitRange(first(from_idx), last(from_idx)) +restrict_index_on_location(from_idx, ::Type{Center}, ::Type{Center}) = UnitRange(first(from_idx), last(from_idx)) +restrict_index_on_location(from_idx, ::Type{Face}, ::Type{Center}) = UnitRange(first(from_idx), last(from_idx)-1) +restrict_index_on_location(from_idx, ::Type{Center}, ::Type{Face}) = UnitRange(first(from_idx)+1, last(from_idx)) diff --git a/src/Fields/interpolate.jl b/src/Fields/interpolate.jl index 95de911e8d..856a121a9a 100644 --- a/src/Fields/interpolate.jl +++ b/src/Fields/interpolate.jl @@ -352,7 +352,9 @@ function interpolate!(to_field::Field, from_field::AbstractField) from_location = Tuple(L() for L in location(from_field)) to_location = Tuple(L() for L in location(to_field)) - launch!(to_arch, to_grid, size(to_field), + params = KernelParameters(interior_indices(to_field)) + + launch!(to_arch, to_grid, params, _interpolate!, to_field, to_grid, to_location, from_field, from_grid, from_location) diff --git a/src/Utils/kernel_launching.jl b/src/Utils/kernel_launching.jl index 097ce8fe00..b1d539c21c 100644 --- a/src/Utils/kernel_launching.jl +++ b/src/Utils/kernel_launching.jl @@ -77,6 +77,9 @@ function KernelParameters(r1::UnitRange, r2::UnitRange, r3::UnitRange) return KernelParameters(size, offsets) end +# Convenience `Tuple`d constructor +KernelParameters(args::Tuple) = KernelParameters(args...) + contiguousrange(range::NTuple{N, Int}, offset::NTuple{N, Int}) where N = Tuple(1+o:r+o for (r, o) in zip(range, offset)) flatten_reduced_dimensions(worksize, dims) = Tuple(d ∈ dims ? 1 : worksize[d] for d = 1:3) diff --git a/test/test_field.jl b/test/test_field.jl index 6759eb5870..2169ba94c3 100644 --- a/test/test_field.jl +++ b/test/test_field.jl @@ -184,6 +184,16 @@ function run_field_interpolation_tests(grid) end end + # Check interpolation on Windowed fields + wf = ZFaceField(grid; indices=(:, :, grid.Nz+1)) + If = Field{Center, Center, Nothing}(grid) + set!(If, (x, y)-> x * y) + interpolate!(wf, If) + + CUDA.@allowscalar begin + @test all(interior(wf) .≈ interior(If)) + end + return nothing end From 7c0fb644b44aca31f834d3c6773c5e1e7f628ba0 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 24 Dec 2024 09:12:00 +0100 Subject: [PATCH 08/31] Correct `tupled_halo_filling` for distributed fields (#4017) * add correction * add a test * comment * probably better solution? * import catke * add some shear * fix tests * test CATKE only on a non-immersed grid * fix tests * just on the underlying grid * correct tests * test only this for the moment * test this at the moment * remove the show --- .../DistributedComputations.jl | 2 + .../distributed_architectures.jl | 13 ++ .../distributed_fields.jl | 2 - .../distributed_on_architecture.jl | 3 - .../halo_communication.jl | 11 +- .../compute_slow_tendencies.jl | 2 - test/test_distributed_hydrostatic_model.jl | 130 +++++++++++++----- .../distributed_hydrostatic_simulation.jl | 2 +- 8 files changed, 116 insertions(+), 49 deletions(-) diff --git a/src/DistributedComputations/DistributedComputations.jl b/src/DistributedComputations/DistributedComputations.jl index cd22c122b7..7595cd84b6 100644 --- a/src/DistributedComputations/DistributedComputations.jl +++ b/src/DistributedComputations/DistributedComputations.jl @@ -10,6 +10,8 @@ using MPI using Oceananigans.Utils using Oceananigans.Grids +using OffsetArrays +using CUDA: CuArray include("distributed_architectures.jl") include("partition_assemble.jl") diff --git a/src/DistributedComputations/distributed_architectures.jl b/src/DistributedComputations/distributed_architectures.jl index eca877fc8c..b00ef07a3a 100644 --- a/src/DistributedComputations/distributed_architectures.jl +++ b/src/DistributedComputations/distributed_architectures.jl @@ -311,6 +311,19 @@ array_type(arch::Distributed) = array_type(child_architecture(arch)) sync_device!(arch::Distributed) = sync_device!(arch.child_architecture) convert_args(arch::Distributed, arg) = convert_args(child_architecture(arch), arg) +# Switch to a synchronized architecture +synchronized(arch) = arch +synchronized(arch::SynchronizedDistributed) = arch +synchronized(arch::Distributed) = Distributed{true}(child_architecture(arch), + arch.partition, + arch.ranks, + arch.local_rank, + arch.local_index, + arch.connectivity, + arch.communicator, + arch.mpi_requests, + arch.mpi_tag) + cpu_architecture(arch::DistributedCPU) = arch cpu_architecture(arch::Distributed{A, S}) where {A, S} = Distributed{S}(CPU(), diff --git a/src/DistributedComputations/distributed_fields.jl b/src/DistributedComputations/distributed_fields.jl index 17f5ef930e..946df63e26 100644 --- a/src/DistributedComputations/distributed_fields.jl +++ b/src/DistributedComputations/distributed_fields.jl @@ -1,5 +1,3 @@ -using CUDA: CuArray -using OffsetArrays: OffsetArray using Oceananigans.Grids: topology using Oceananigans.Fields: validate_field_data, indices, validate_boundary_conditions using Oceananigans.Fields: validate_indices, recv_from_buffers!, set_to_array!, set_to_field! diff --git a/src/DistributedComputations/distributed_on_architecture.jl b/src/DistributedComputations/distributed_on_architecture.jl index 75b936fb9f..6905f29b33 100644 --- a/src/DistributedComputations/distributed_on_architecture.jl +++ b/src/DistributedComputations/distributed_on_architecture.jl @@ -1,6 +1,3 @@ -using CUDA: CuArray -using OffsetArrays - import Oceananigans.Architectures: on_architecture # We do not support switching from distributed and serial through `on_architecture`. diff --git a/src/DistributedComputations/halo_communication.jl b/src/DistributedComputations/halo_communication.jl index 011bb12dbe..c8833d6725 100644 --- a/src/DistributedComputations/halo_communication.jl +++ b/src/DistributedComputations/halo_communication.jl @@ -1,10 +1,10 @@ using KernelAbstractions: @kernel, @index -using OffsetArrays: OffsetArray using Oceananigans.Fields: fill_send_buffers!, recv_from_buffers!, reduced_dimensions, - instantiated_location + instantiated_location, + produce_ordinary_fields import Oceananigans.Fields: tupled_fill_halo_regions! @@ -79,8 +79,11 @@ end ##### function tupled_fill_halo_regions!(fields, grid::DistributedGrid, args...; kwargs...) - for field in fields - fill_halo_regions!(field, args...; kwargs...) + ordinary_fields = produce_ordinary_fields(fields, args...; kwargs) + + for field in ordinary_fields + # Make sure we are filling a `Field` type. + field isa Field && fill_halo_regions!(field, args...; kwargs...) end end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl index 30f327ed20..67302844fa 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl @@ -91,8 +91,6 @@ end end @inline function compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, ::Val{2}) - FT = eltype(GUⁿ) - @inbounds GUⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center()) @inbounds GVⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center()) diff --git a/test/test_distributed_hydrostatic_model.jl b/test/test_distributed_hydrostatic_model.jl index c11d3c5248..e5ab4bd461 100644 --- a/test/test_distributed_hydrostatic_model.jl +++ b/test/test_distributed_hydrostatic_model.jl @@ -26,7 +26,8 @@ MPI.Initialized() || MPI.Init() # to initialize MPI. using Oceananigans.Operators: hack_cosd -using Oceananigans.DistributedComputations: partition, all_reduce, cpu_architecture, reconstruct_global_grid +using Oceananigans.DistributedComputations: ranks, partition, all_reduce, cpu_architecture, reconstruct_global_grid, synchronized +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity function Δ_min(grid) Δx_min = minimum_xspacing(grid, Center(), Center(), Center()) @@ -36,25 +37,32 @@ end @inline Gaussian(x, y, L) = exp(-(x^2 + y^2) / L^2) -function solid_body_rotation_test(grid) +function rotation_with_shear_test(grid, closure=nothing) - free_surface = SplitExplicitFreeSurface(grid; substeps = 5, gravitational_acceleration = 1) + free_surface = SplitExplicitFreeSurface(grid; substeps = 8, gravitational_acceleration = 1) coriolis = HydrostaticSphericalCoriolis(rotation_rate = 1) + tracers = if closure isa CATKEVerticalDiffusivity + (:c, :b, :e) + else + (:c, :b) + end + model = HydrostaticFreeSurfaceModel(; grid, momentum_advection = VectorInvariant(), free_surface = free_surface, coriolis = coriolis, - tracers = :c, + closure, + tracers, tracer_advection = WENO(), - buoyancy = nothing, - closure = nothing) + buoyancy = BuoyancyTracer()) g = model.free_surface.gravitational_acceleration R = grid.radius Ω = model.coriolis.rotation_rate - uᵢ(λ, φ, z) = 0.1 * cosd(φ) * sind(λ) + # Add some shear on the velocity field + uᵢ(λ, φ, z) = 0.1 * cosd(φ) * sind(λ) + 0.05 * z ηᵢ(λ, φ, z) = (R * Ω * 0.1 + 0.1^2 / 2) * sind(φ)^2 / g * sind(λ) # Gaussian leads to values with O(1e-60), @@ -64,13 +72,13 @@ function solid_body_rotation_test(grid) set!(model, u=uᵢ, η=ηᵢ, c=cᵢ) - @show Δt_local = 0.1 * Δ_min(grid) / sqrt(g * grid.Lz) - @show Δt = all_reduce(min, Δt_local, architecture(grid)) + Δt_local = 0.1 * Δ_min(grid) / sqrt(g * grid.Lz) + Δt = all_reduce(min, Δt_local, architecture(grid)) - simulation = Simulation(model; Δt, stop_iteration = 10) + simulation = Simulation(model; Δt, stop_iteration = 10, verbose = false) run!(simulation) - return merge(model.velocities, model.tracers, (; η = model.free_surface.η)) + return model end Nx = 32 @@ -78,7 +86,7 @@ Ny = 32 for arch in archs @testset "Testing distributed solid body rotation" begin - underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 1), + underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 3), halo = (4, 4, 4), latitude = (-80, 80), longitude = (-160, 160), @@ -95,41 +103,89 @@ for arch in archs global_immersed_grid = ImmersedBoundaryGrid(global_underlying_grid, GridFittedBottom(bottom)) for (grid, global_grid) in zip((underlying_grid, immersed_grid, immersed_active_grid), (global_underlying_grid, global_immersed_grid, global_immersed_grid)) + if arch.local_rank == 0 + @info " Testing distributed solid body rotation with $(ranks(arch)) ranks on $(typeof(grid).name.wrapper)" + end - # "s" for "serial" computation - us, vs, ws, cs, ηs = solid_body_rotation_test(global_grid) + # "s" for "serial" computation, "p" for parallel + ms = rotation_with_shear_test(global_grid) + mp = rotation_with_shear_test(grid) - us = interior(on_architecture(CPU(), us)) - vs = interior(on_architecture(CPU(), vs)) - ws = interior(on_architecture(CPU(), ws)) - cs = interior(on_architecture(CPU(), cs)) - ηs = interior(on_architecture(CPU(), ηs)) - - @info " Testing distributed solid body rotation with architecture $arch on $(typeof(grid).name.wrapper)" - u, v, w, c, η = solid_body_rotation_test(grid) + us = interior(on_architecture(CPU(), ms.velocities.u)) + vs = interior(on_architecture(CPU(), ms.velocities.v)) + ws = interior(on_architecture(CPU(), ms.velocities.w)) + cs = interior(on_architecture(CPU(), ms.tracers.c)) + ηs = interior(on_architecture(CPU(), ms.free_surface.η)) cpu_arch = cpu_architecture(arch) - u = interior(on_architecture(cpu_arch, u)) - v = interior(on_architecture(cpu_arch, v)) - w = interior(on_architecture(cpu_arch, w)) - c = interior(on_architecture(cpu_arch, c)) - η = interior(on_architecture(cpu_arch, η)) + up = interior(on_architecture(cpu_arch, mp.velocities.u)) + vp = interior(on_architecture(cpu_arch, mp.velocities.v)) + wp = interior(on_architecture(cpu_arch, mp.velocities.w)) + cp = interior(on_architecture(cpu_arch, mp.tracers.c)) + ηp = interior(on_architecture(cpu_arch, mp.free_surface.η)) - us = partition(us, cpu_arch, size(u)) - vs = partition(vs, cpu_arch, size(v)) - ws = partition(ws, cpu_arch, size(w)) - cs = partition(cs, cpu_arch, size(c)) - ηs = partition(ηs, cpu_arch, size(η)) + us = partition(us, cpu_arch, size(up)) + vs = partition(vs, cpu_arch, size(vp)) + ws = partition(ws, cpu_arch, size(wp)) + cs = partition(cs, cpu_arch, size(cp)) + ηs = partition(ηs, cpu_arch, size(ηp)) atol = eps(eltype(grid)) rtol = sqrt(eps(eltype(grid))) - @test all(isapprox(u, us; atol, rtol)) - @test all(isapprox(v, vs; atol, rtol)) - @test all(isapprox(w, ws; atol, rtol)) - @test all(isapprox(c, cs; atol, rtol)) - @test all(isapprox(η, ηs; atol, rtol)) + @test all(isapprox(up, us; atol, rtol)) + @test all(isapprox(vp, vs; atol, rtol)) + @test all(isapprox(wp, ws; atol, rtol)) + @test all(isapprox(cp, cs; atol, rtol)) + @test all(isapprox(ηp, ηs; atol, rtol)) end + + # # CATKE works only with synchronized communication at the moment + # arch = synchronized(arch) + # closure = CATKEVerticalDiffusivity() + + # underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 3), + # halo = (4, 4, 4), + # latitude = (-80, 80), + # longitude = (-160, 160), + # z = (-1, 0), + # radius = 1, + # topology=(Bounded, Bounded, Bounded)) + + # bottom(λ, φ) = -30 < λ < 30 && -40 < φ < 20 ? 0 : - 1 + + # # "s" for "serial" computation, "p" for parallel + # ms = rotation_with_shear_test(global_underlying_grid, closure) + # mp = rotation_with_shear_test(underlying_grid, closure) + + # us = interior(on_architecture(CPU(), ms.velocities.u)) + # vs = interior(on_architecture(CPU(), ms.velocities.v)) + # ws = interior(on_architecture(CPU(), ms.velocities.w)) + # cs = interior(on_architecture(CPU(), ms.tracers.c)) + # ηs = interior(on_architecture(CPU(), ms.free_surface.η)) + + # cpu_arch = cpu_architecture(arch) + + # up = interior(on_architecture(cpu_arch, mp.velocities.u)) + # vp = interior(on_architecture(cpu_arch, mp.velocities.v)) + # wp = interior(on_architecture(cpu_arch, mp.velocities.w)) + # cp = interior(on_architecture(cpu_arch, mp.tracers.c)) + # ηp = interior(on_architecture(cpu_arch, mp.free_surface.η)) + + # us = partition(us, cpu_arch, size(up)) + # vs = partition(vs, cpu_arch, size(vp)) + # ws = partition(ws, cpu_arch, size(wp)) + # cs = partition(cs, cpu_arch, size(cp)) + # ηs = partition(ηs, cpu_arch, size(ηp)) + + # atol = eps(eltype(immersed_active_grid)) + # rtol = sqrt(eps(eltype(immersed_active_grid))) + + # @test all(isapprox(up, us; atol, rtol)) + # @test all(isapprox(vp, vs; atol, rtol)) + # @test all(isapprox(wp, ws; atol, rtol)) + # @test all(isapprox(cp, cs; atol, rtol)) + # @test all(isapprox(ηp, ηs; atol, rtol)) end end diff --git a/validation/distributed_simulations/distributed_scaling/distributed_hydrostatic_simulation.jl b/validation/distributed_simulations/distributed_scaling/distributed_hydrostatic_simulation.jl index 87621c1bff..022ac00038 100644 --- a/validation/distributed_simulations/distributed_scaling/distributed_hydrostatic_simulation.jl +++ b/validation/distributed_simulations/distributed_scaling/distributed_hydrostatic_simulation.jl @@ -9,7 +9,7 @@ using Oceananigans.Utils: prettytime using Oceananigans.DistributedComputations using Oceananigans.Grids: node using Oceananigans.Advection: cell_advection_timescale -using Oceananigans.TurbulenceClosures.CATKEVerticalDiffusivities: CATKEVerticalDiffusivity +using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity using Oceananigans.Units using SeawaterPolynomials.TEOS10: TEOS10EquationOfState From 786bce808d3c1145f38d7a3ac0dd1fd619283a5c Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 25 Dec 2024 06:26:25 +1100 Subject: [PATCH 09/31] allow finding axis labels for immersed boundaries and Observables (#4023) --- ext/OceananigansMakieExt.jl | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ext/OceananigansMakieExt.jl b/ext/OceananigansMakieExt.jl index 60c8bbf413..05c9d766cd 100644 --- a/ext/OceananigansMakieExt.jl +++ b/ext/OceananigansMakieExt.jl @@ -7,6 +7,7 @@ using Oceananigans.AbstractOperations: AbstractOperation using Oceananigans.Architectures: on_architecture using Oceananigans.ImmersedBoundaries: mask_immersed_field! +using Makie: Observable using MakieCore: AbstractPlot import MakieCore: convert_arguments, _create_plot import Makie: args_preferred_axis @@ -46,6 +47,7 @@ end axis_str(::RectilinearGrid, dim) = ("x", "y", "z")[dim] axis_str(::LatitudeLongitudeGrid, dim) = ("Longitude (deg)", "Latitude (deg)", "z")[dim] +axis_str(grid::ImmersedBoundaryGrid, dim) = axis_str(grid.underlying_grid, dim) function _create_plot(F::Function, attributes::Dict, f::Field) converted_args = convert_field_argument(f) @@ -81,6 +83,9 @@ function _create_plot(F::Function, attributes::Dict, op::AbstractOperation) return _create_plot(F::Function, attributes::Dict, f) end +_create_plot(F::Function, attributes::Dict, f::Observable{<:Field}) = + _create_plot(F, attributes, f[]) + convert_arguments(pl::Type{<:AbstractPlot}, f::Field) = convert_arguments(pl, convert_field_argument(f)...) @@ -183,4 +188,3 @@ function convert_arguments(pl::Type{<:AbstractPlot}, ξ1::AbstractArray, ξ2::Ab end end # module - From 01cad9efb409f900e749bc6ea9a471dd921c73f7 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 26 Dec 2024 05:56:54 +1100 Subject: [PATCH 10/31] Add info for Julia version used for testing (#4022) * update index.md * punctuate * update paper list * Update docs/src/index.md Co-authored-by: Gregory L. Wagner * Update docs/src/index.md Co-authored-by: Gregory L. Wagner --------- Co-authored-by: Gregory L. Wagner --- docs/src/index.md | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/docs/src/index.md b/docs/src/index.md index 4ac15e0404..9116971878 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -24,8 +24,11 @@ julia> using Pkg julia> Pkg.add("Oceananigans") ``` -!!! compat "Julia 1.9 or later is required; Julia 1.10 is suggested" - Oceananigans requires Julia 1.9 or later. We strongly urge you to use Julia 1.10. +!!! compat "Julia 1.9 or later is required" + Oceananigans requires Julia 1.9 or later. + +!!! info "Tested Julia versions" + Oceananigans is currently tested on Julia 1.10. If you're [new to Julia](https://docs.julialang.org/en/v1/manual/getting-started/) and its [wonderful `Pkg` manager](https://docs.julialang.org/en/v1/stdlib/Pkg/), the [Oceananigans wiki](https://github.com/CliMA/Oceananigans.jl/wiki) provides [more detailed installation instructions](https://github.com/CliMA/Oceananigans.jl/wiki/Installation-and-getting-started-with-Oceananigans). @@ -93,10 +96,10 @@ If you have work using Oceananigans that you would like to have listed here, ple 1. Silvestri, S., Wagner, G. L., Constantinou, N. C., Hill, C., Campin, J.-M., Souza, A., Bishnu, S., Churavy, V., Marshall, J., and Ferrari, R. (2024) [A GPU-based ocean dynamical core for routine mesoscale-resolving climate simulations](https://doi.org/10.22541/essoar.171708158.82342448/v1), _ESS Open Archive_. DOI: [10.22541/essoar.171708158.82342448/v1](https://doi.org/10.22541/essoar.171708158.82342448/v1) -1. Silvestri, S., Wagner, G. L., Campin, J.-M., Constantinou, N. C., Hill, C., Souza, A., and Ferrari, R. (2024). [A new WENO-based momentum advection scheme for simulations of ocean mesoscale turbulence](https://doi.org/10.22541/essoar.170110657.76489860/v2), _ESS Open Archive_. DOI: [10.22541/essoar.170110657.76489860/v2](https://doi.org/10.22541/essoar.170110657.76489860/v2) - 1. Whitley V. and Wenegrat, J. O. (2024) [Breaking internal waves on sloping topography: connecting parcel displacements to overturn size, interior-boundary exchanges, and mixing](https://doi.org/10.31223/X5PM5Q), _Earth Arxiv_. DOI: [10.31223/X5PM5Q](https://doi.org/10.31223/X5PM5Q) +1. Silvestri, S., Wagner, G. L., Campin, J.-M., Constantinou, N. C., Hill, C., Souza, A., and Ferrari, R. (2024). [A new WENO-based momentum advection scheme for simulations of ocean mesoscale turbulence](https://doi.org/10.1029/2023MS004130), _Journal of Advances in Modeling Earth Systems_, **16(7)**, e2023MS004130. DOI: [10.1029/2023MS004130](https://doi.org/10.1029/2023MS004130) + 1. Chen S., Strong-Wright J., and Taylor, J. R. (2024) [Modeling carbon dioxide removal via sinking of particulate organic carbon from macroalgae cultivation](https://doi.org/10.3389/fmars.2024.1359614), _Frontiers in Marine Science_, **11**, 1359614. DOI: [10.3389/fmars.2024.1359614](https://doi.org/10.3389/fmars.2024.1359614) 1. Gupta, M., Gürcan, E., and Thompson, A. F. (2024). [Eddy-induced dispersion of sea ice floes at the marginal ice zone](https://doi.org/10.1029/2023GL105656), _Geophysical Research Letters_, **51**, e2023GL105656. DOI: [10.1029/2023GL105656](https://doi.org/10.1029/2023GL105656) From 753c672e8c5bb345d2e960c78863530f40e01f8b Mon Sep 17 00:00:00 2001 From: William Moses Date: Mon, 30 Dec 2024 15:17:33 -0500 Subject: [PATCH 11/31] Enzyme: mark prettytime as non-differentiable (#4026) --- ext/OceananigansEnzymeExt.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/ext/OceananigansEnzymeExt.jl b/ext/OceananigansEnzymeExt.jl index e1b7d74a1f..82350e23a9 100644 --- a/ext/OceananigansEnzymeExt.jl +++ b/ext/OceananigansEnzymeExt.jl @@ -17,6 +17,7 @@ EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Base.:(==)), ::Oceananigans.Abstr EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.validate_grid), x...) = nothing EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.AbstractOperations.metric_function), x...) = nothing EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.flatten_reduced_dimensions), x...) = nothing +EnzymeCore.EnzymeRules.inactive_noinl(::typeof(Oceananigans.Utils.prettytime), x...) = nothing EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.Grids.total_size), x...) = nothing EnzymeCore.EnzymeRules.inactive(::typeof(Oceananigans.BoundaryConditions.parent_size_and_offset), x...) = nothing @inline EnzymeCore.EnzymeRules.inactive_type(v::Type{Oceananigans.Utils.KernelParameters}) = true From aa5f2393b2bd82b8bc13e217e9dd62590c9663e4 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Tue, 31 Dec 2024 17:24:07 +1100 Subject: [PATCH 12/31] Makie extension: tweaks the xticks for domains that wraps around the globe (#4025) * longitude ticks multiple of 60 * remove extra space * Update ext/OceananigansMakieExt.jl Co-authored-by: Gregory L. Wagner * Update OceananigansMakieExt.jl --------- Co-authored-by: Gregory L. Wagner --- ext/OceananigansMakieExt.jl | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/ext/OceananigansMakieExt.jl b/ext/OceananigansMakieExt.jl index 05c9d766cd..42c6d32ef8 100644 --- a/ext/OceananigansMakieExt.jl +++ b/ext/OceananigansMakieExt.jl @@ -1,7 +1,7 @@ module OceananigansMakieExt using Oceananigans -using Oceananigans.Grids: OrthogonalSphericalShellGrid +using Oceananigans.Grids: OrthogonalSphericalShellGrid, topology using Oceananigans.Fields: AbstractField using Oceananigans.AbstractOperations: AbstractOperation using Oceananigans.Architectures: on_architecture @@ -49,10 +49,12 @@ axis_str(::RectilinearGrid, dim) = ("x", "y", "z")[dim] axis_str(::LatitudeLongitudeGrid, dim) = ("Longitude (deg)", "Latitude (deg)", "z")[dim] axis_str(grid::ImmersedBoundaryGrid, dim) = axis_str(grid.underlying_grid, dim) +const LLGOrIBLLG = Union{LatitudeLongitudeGrid, ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:LatitudeLongitudeGrid}} + function _create_plot(F::Function, attributes::Dict, f::Field) converted_args = convert_field_argument(f) - if !(:axis ∈ keys(attributes)) # Let's try to add labels automatically + if !(:axis ∈ keys(attributes)) # Let's try to automatically add labels and ticks d1, d2, D = deduce_dimensionality(f) grid = f.grid @@ -71,6 +73,11 @@ function _create_plot(F::Function, attributes::Dict, f::Field) throw(ArgumentError("Cannot create axis labels for a 3D field!")) end + # if longitude wraps around the globe then adjust the longitude ticks + if grid isa LLGOrIBLLG && grid.Lx == 360 && topology(grid, 1) == Periodic + axis = merge(axis, (xticks = -360:60:360,)) + end + attributes[:axis] = axis end From 0345a2cb58a71df62471a2a721765c07caa13c46 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Tue, 7 Jan 2025 11:38:41 -0700 Subject: [PATCH 13/31] Add default Cb for CATKE (#4031) This PR changes the default Cb parameter for CATKE to yield a Von Karman constant of 0.4 in near-bottom similarity layers. Closes #4015 --- .../TKEBasedVerticalDiffusivities/catke_mixing_length.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl index dce68696b6..c409159ba8 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/TKEBasedVerticalDiffusivities/catke_mixing_length.jl @@ -14,7 +14,7 @@ Contains mixing length parameters for CATKE vertical diffusivity. """ Base.@kwdef struct CATKEMixingLength{FT} Cˢ :: FT = 1.131 # Surface distance coefficient for shear length scale - Cᵇ :: FT = Inf # Bottom distance coefficient for shear length scale + Cᵇ :: FT = 0.28 # Bottom distance coefficient for shear length scale Cˢᵖ :: FT = 0.505 # Sheared convective plume coefficient CRiᵟ :: FT = 1.02 # Stability function width CRi⁰ :: FT = 0.254 # Stability function lower Ri From 19baa398e793b64ca182ac6c1de5599034f188fa Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 7 Jan 2025 23:02:42 -0100 Subject: [PATCH 14/31] Inference problem on the CPU (#4032) * fix the problem * bugfix --- ...ute_hydrostatic_free_surface_tendencies.jl | 27 ++++++++++--------- ..._free_surface_tendency_kernel_functions.jl | 6 ++--- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl index 9e21089913..29303f6581 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl @@ -126,6 +126,9 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para u_immersed_bc = immersed_boundary_condition(velocities.u) v_immersed_bc = immersed_boundary_condition(velocities.v) + u_forcing = model.forcing.u + v_forcing = model.forcing.v + start_momentum_kernel_args = (model.advection.momentum, model.coriolis, model.closure) @@ -139,17 +142,17 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para model.auxiliary_fields, model.clock) - u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args...) - v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args...) + u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args..., u_forcing) + v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args..., v_forcing) launch!(arch, grid, kernel_parameters, compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid, - active_cells_map, u_kernel_args, model.forcing.u; + active_cells_map, u_kernel_args; active_cells_map) launch!(arch, grid, kernel_parameters, compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid, - active_cells_map, v_kernel_args, model.forcing.v; + active_cells_map, v_kernel_args; active_cells_map) return nothing @@ -179,27 +182,27 @@ end ##### """ Calculate the right-hand-side of the u-velocity equation. """ -@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args) i, j, k = @index(Global, NTuple) - @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...) end -@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args) idx = @index(Global, Linear) i, j, k = active_linear_index_to_tuple(idx, active_cells_map) - @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...) end """ Calculate the right-hand-side of the v-velocity equation. """ -@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args) i, j, k = @index(Global, NTuple) - @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...) end -@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args, forcing) +@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args) idx = @index(Global, Linear) i, j, k = active_linear_index_to_tuple(idx, active_cells_map) - @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing) + @inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...) end ##### diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl index c5e3e9c1ab..28cb289e58 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -38,7 +38,7 @@ implicitly during time-stepping. diffusivities, hydrostatic_pressure_anomaly, auxiliary_fields, - clock, + clock, forcing) model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) @@ -49,7 +49,7 @@ implicitly during time-stepping. - ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy) - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields) - + forcing(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers))) + + forcing(i, j, k, grid, clock, model_fields)) end """ @@ -82,7 +82,7 @@ implicitly during time-stepping. model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields) - return ( - U_dot_∇v(i, j, k, grid, advection, velocities) + return ( - U_dot_∇v(i, j, k, grid, advection, velocities) - explicit_barotropic_pressure_y_gradient(i, j, k, grid, free_surface) - y_f_cross_U(i, j, k, grid, coriolis, velocities) - ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) From e02564df8ab90fb8f994fddf465b10d123fdfc3a Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Sat, 11 Jan 2025 04:41:36 -0100 Subject: [PATCH 15/31] (0.95.6) Pin GPUArrays version + fix GPU testing pipeline (#4037) * pin cuda * pin previous version * remove file that entered branch by mistake * Update pipeline.yml --------- Co-authored-by: Gregory L. Wagner --- .buildkite/pipeline.yml | 2 +- Project.toml | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 2ac6d44c86..8a162418cd 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,7 +1,7 @@ env: JULIA_VERSION: "1.10.6" JULIA_MINOR_VERSION: "1.10" - SVERDRUP_HOME: "/data5/glwagner" + SVERDRUP_HOME: "/home/users/glwagner" TARTARUS_HOME: "/storage5/buildkite-agent" JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager JULIA_NUM_PRECOMPILE_TASKS: 8 diff --git a/Project.toml b/Project.toml index 69f41143dd..734e3e1ac1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -version = "0.95.5" +version = "0.95.6" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -12,6 +12,7 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" Glob = "c27321d9-0574-5035-807b-f59d2c89b15c" IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" @@ -44,7 +45,7 @@ OceananigansMakieExt = ["MakieCore", "Makie"] [compat] Adapt = "4.1.1" -CUDA = "4.1.1, 5" +CUDA = "5" Crayons = "4" CubedSphere = "0.2, 0.3" Dates = "1.9" @@ -53,6 +54,7 @@ DocStringExtensions = "0.8, 0.9" Enzyme = "0.13.14" FFTW = "1" Glob = "1.3" +GPUArrays = "10.3" IncompleteLU = "0.2" InteractiveUtils = "1.9" IterativeSolvers = "0.9" From 31e62c6599583c4ae4af2349c8068c39fceba3de Mon Sep 17 00:00:00 2001 From: William Moses Date: Sun, 12 Jan 2025 17:56:44 -0600 Subject: [PATCH 16/31] Allows properties of `StaticVerticalCoordinate` to have different types (#4029) * Fix adapt of StaticVerticalCoordinate if subadaptations have distinct types * fixed * Fix definition of static coordinate * Change name to vertical_coordinates * fix docstring --------- Co-authored-by: Gregory Wagner Co-authored-by: Navid C. Constantinou --- src/Grids/Grids.jl | 2 +- ..._coordinate.jl => vertical_coordinates.jl} | 56 ++++++++++--------- 2 files changed, 32 insertions(+), 26 deletions(-) rename src/Grids/{vertical_coordinate.jl => vertical_coordinates.jl} (66%) diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 792572b291..320e9035cc 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -118,7 +118,7 @@ struct ZDirection <: AbstractDirection end struct NegativeZDirection <: AbstractDirection end include("abstract_grid.jl") -include("vertical_coordinate.jl") +include("vertical_coordinates.jl") include("grid_utils.jl") include("nodes_and_spacings.jl") include("zeros_and_ones.jl") diff --git a/src/Grids/vertical_coordinate.jl b/src/Grids/vertical_coordinates.jl similarity index 66% rename from src/Grids/vertical_coordinate.jl rename to src/Grids/vertical_coordinates.jl index 47447746c1..3028fb013b 100644 --- a/src/Grids/vertical_coordinate.jl +++ b/src/Grids/vertical_coordinates.jl @@ -3,48 +3,54 @@ #### # This file implements everything related to vertical coordinates in Oceananigans. -# Vertical coordinates are independent of the underlying grid type as we support grids that are -# "unstructured" or "curvilinear" only in the horizontal direction. -# For this reason the vertical coodinate is _special_, and it can be implemented once for all grid types. +# Vertical coordinates are independent of the underlying grid type since only grids that are +# "unstructured" or "curvilinear" in the horizontal directions are supported in Oceananigans. +# Thus the vertical coordinate is _special_, and it can be implemented once for all grid types. abstract type AbstractVerticalCoordinate end -# Represents a static one-dimensional vertical coordinate. -# -# # Fields -# - `cᶜ::C`: Cell-centered coordinate. -# - `cᶠ::C`: Face-centered coordinate. -# - `Δᶜ::D`: Cell-centered grid spacing. -# - `Δᶠ::D`: Face-centered grid spacing. -struct StaticVerticalCoordinate{C, D} <: AbstractVerticalCoordinate +""" + struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate + +Represent a static one-dimensional vertical coordinate. + +Fields +====== + +- `cᶜ::C`: Cell-centered coordinate. +- `cᶠ::D`: Face-centered coordinate. +- `Δᶜ::E`: Cell-centered grid spacing. +- `Δᶠ::F`: Face-centered grid spacing. +""" +struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate cᵃᵃᶠ :: C - cᵃᵃᶜ :: C - Δᵃᵃᶠ :: D - Δᵃᵃᶜ :: D + cᵃᵃᶜ :: D + Δᵃᵃᶠ :: E + Δᵃᵃᶜ :: F end #### -#### Some usefull aliases +#### Some useful aliases #### -const RegularVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Number} +const RegularVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Any, <:Number, <:Number} const RegularVerticalGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate} #### #### Adapt and on_architecture #### -Adapt.adapt_structure(to, coord::StaticVerticalCoordinate) = - StaticVerticalCoordinate(Adapt.adapt(to, coord.cᵃᵃᶠ), - Adapt.adapt(to, coord.cᵃᵃᶜ), - Adapt.adapt(to, coord.Δᵃᵃᶠ), - Adapt.adapt(to, coord.Δᵃᵃᶜ)) +Adapt.adapt_structure(to, coord::StaticVerticalCoordinate) = + StaticVerticalCoordinate(Adapt.adapt(to, coord.cᵃᵃᶠ), + Adapt.adapt(to, coord.cᵃᵃᶜ), + Adapt.adapt(to, coord.Δᵃᵃᶠ), + Adapt.adapt(to, coord.Δᵃᵃᶜ)) on_architecture(arch, coord::StaticVerticalCoordinate) = - StaticVerticalCoordinate(on_architecture(arch, coord.cᵃᵃᶠ), - on_architecture(arch, coord.cᵃᵃᶜ), - on_architecture(arch, coord.Δᵃᵃᶠ), - on_architecture(arch, coord.Δᵃᵃᶜ)) + StaticVerticalCoordinate(on_architecture(arch, coord.cᵃᵃᶠ), + on_architecture(arch, coord.cᵃᵃᶜ), + on_architecture(arch, coord.Δᵃᵃᶠ), + on_architecture(arch, coord.Δᵃᵃᶜ)) ##### ##### Nodes and spacings (common to every grid)... From 2b601f1ab052b6b257c5f273511efe46e20db5f3 Mon Sep 17 00:00:00 2001 From: William Moses Date: Mon, 13 Jan 2025 09:17:43 -0600 Subject: [PATCH 17/31] Generalize ntuple type constraints to be tuple{vararg} (#4035) --- src/BoundaryConditions/fill_halo_regions.jl | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/BoundaryConditions/fill_halo_regions.jl b/src/BoundaryConditions/fill_halo_regions.jl index 5eb29f687e..88886fd569 100644 --- a/src/BoundaryConditions/fill_halo_regions.jl +++ b/src/BoundaryConditions/fill_halo_regions.jl @@ -44,7 +44,7 @@ end @inline extract_bc(bc, ::Val{:bottom_and_top}) = (extract_bottom_bc(bc), extract_top_bc(bc)) # Finally, the true fill_halo! -const MaybeTupledData = Union{OffsetArray, NTuple{<:Any, OffsetArray}} +const MaybeTupledData = Union{OffsetArray, Tuple{Vararg{OffsetArray}}} "Fill halo regions in ``x``, ``y``, and ``z`` for a given field's data." function fill_halo_regions!(c::MaybeTupledData, boundary_conditions, indices, loc, grid, args...; @@ -145,9 +145,9 @@ split_halo_filling(::DCBC, bcs2) = true ##### Halo filling order ##### -const PBCT = Union{PBC, NTuple{<:Any, <:PBC}} -const MCBCT = Union{MCBC, NTuple{<:Any, <:MCBC}} -const DCBCT = Union{DCBC, NTuple{<:Any, <:DCBC}} +const PBCT = Union{PBC, Tuple{Vararg{PBC}}} +const MCBCT = Union{MCBC, Tuple{Vararg{MCBC}}} +const DCBCT = Union{DCBC, Tuple{Vararg{DCBC}}} # Distributed halos have to be filled last to allow the # possibility of asynchronous communication: @@ -259,7 +259,7 @@ end # support tupled halo filling import Oceananigans.Utils: @constprop -@kernel function _fill_west_and_east_halo!(c::NTuple, west_bc, east_bc, loc, grid, args) +@kernel function _fill_west_and_east_halo!(c::Tuple, west_bc, east_bc, loc, grid, args) j, k = @index(Global, NTuple) ntuple(Val(length(west_bc))) do n Base.@_inline_meta @@ -271,7 +271,7 @@ import Oceananigans.Utils: @constprop end end -@kernel function _fill_south_and_north_halo!(c::NTuple, south_bc, north_bc, loc, grid, args) +@kernel function _fill_south_and_north_halo!(c::Tuple, south_bc, north_bc, loc, grid, args) i, k = @index(Global, NTuple) ntuple(Val(length(south_bc))) do n Base.@_inline_meta @@ -283,7 +283,7 @@ end end end -@kernel function _fill_bottom_and_top_halo!(c::NTuple, bottom_bc, top_bc, loc, grid, args) +@kernel function _fill_bottom_and_top_halo!(c::Tuple, bottom_bc, top_bc, loc, grid, args) i, j = @index(Global, NTuple) ntuple(Val(length(bottom_bc))) do n Base.@_inline_meta From c3d14a9105172369c0f53e84318e83ee75aa15f3 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 14 Jan 2025 11:43:23 -0500 Subject: [PATCH 18/31] Optimize masking xy planes (#4040) * fix the problem * bugfix * improve mask immersed field * remove all splatting * make sure everything is inferrable * provide fallbacks * mask only the 3D fields * Update update_hydrostatic_free_surface_model_state.jl * remove the mask function --- src/ImmersedBoundaries/mask_immersed_field.jl | 100 ++++++++---------- ...te_hydrostatic_free_surface_model_state.jl | 6 +- 2 files changed, 48 insertions(+), 58 deletions(-) diff --git a/src/ImmersedBoundaries/mask_immersed_field.jl b/src/ImmersedBoundaries/mask_immersed_field.jl index 586f6054c3..7e1e03c68a 100644 --- a/src/ImmersedBoundaries/mask_immersed_field.jl +++ b/src/ImmersedBoundaries/mask_immersed_field.jl @@ -7,22 +7,23 @@ using Oceananigans.Fields: ConstantField, OneField, ZeroField instantiate(T::Type) = T() instantiate(t) = t -# No masking for constant fields -mask_immersed_field!(::OneField, args...) = nothing -mask_immersed_field!(::ZeroField, args...) = nothing -mask_immersed_field!(::ConstantField, args...) = nothing - -# No masking for constant fields +# No masking for constant fields, numbers or nothing +mask_immersed_field!(::OneField, args...; kw...) = nothing +mask_immersed_field!(::ZeroField, args...; kw...) = nothing +mask_immersed_field!(::ConstantField, args...; kw...) = nothing +mask_immersed_field!(::Number, args...; kw...) = nothing +mask_immersed_field!(::Nothing, args...; kw...) = nothing + +# No masking for constant fields, numbers or nothing mask_immersed_field_xy!(::OneField, args...; kw...) = nothing mask_immersed_field_xy!(::ZeroField, args...; kw...) = nothing mask_immersed_field_xy!(::ConstantField, args...; kw...) = nothing +mask_immersed_field_xy!(::Number, args...; kw...) = nothing +mask_immersed_field_xy!(::Nothing, args...; kw...) = nothing -mask_immersed_field!(field, grid, loc, value) = nothing mask_immersed_field!(field::Field, value=zero(eltype(field.grid))) = mask_immersed_field!(field, field.grid, location(field), value) -mask_immersed_field!(::Number, args...) = nothing - function mask_immersed_field!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(+)}, value=zero(eltype(bop))) a_value = ifelse(bop.b isa Number, -bop.b, value) mask_immersed_field!(bop.a, a_value) @@ -41,6 +42,9 @@ function mask_immersed_field!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(- return nothing end +# Fallback +mask_immersed_field!(field, grid, loc, value) = nothing + """ mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, value) @@ -48,56 +52,56 @@ masks `field` defined on `grid` with a value `val` at locations where `periphera """ function mask_immersed_field!(field::Field, grid::ImmersedBoundaryGrid, loc, value) arch = architecture(field) - loc = instantiate.(loc) + loc = instantiate.(loc) launch!(arch, grid, :xyz, _mask_immersed_field!, field, loc, grid, value) return nothing end -@kernel function _mask_immersed_field!(field, loc, grid, value) +@kernel function _mask_immersed_field!(field, (ℓx, ℓy, ℓz), grid, value) i, j, k = @index(Global, NTuple) - @inbounds field[i, j, k] = scalar_mask(i, j, k, grid, grid.immersed_boundary, loc..., value, field) + masked = immersed_peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) + @inbounds field[i, j, k] = ifelse(masked, value, field[i, j, k]) end -mask_immersed_field_xy!(field, args...; kw...) = nothing -mask_immersed_field_xy!(::Nothing, args...; kw...) = nothing -mask_immersed_field_xy!(field, value=zero(eltype(field.grid)); k, mask = peripheral_node) = - mask_immersed_field_xy!(field, field.grid, location(field), value; k, mask) - -mask_immersed_field_xy!(::Number, args...) = nothing +mask_immersed_field_xy!(field, value=zero(eltype(field.grid)); k) = + mask_immersed_field_xy!(field, field.grid, location(field), value, k) -function mask_immersed_field_xy!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(+)}, value=zero(eltype(bop))) +function mask_immersed_field_xy!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(+)}, value=zero(eltype(bop)); k) a_value = ifelse(bop.b isa Number, -bop.b, value) - mask_immersed_field_xy!(bop.a, a_value) + mask_immersed_field_xy!(bop.a, a_value; k) b_value = ifelse(bop.a isa Number, -bop.a, value) - mask_immersed_field_xy!(bop.b, b_value) + mask_immersed_field_xy!(bop.b, b_value; k) return nothing end -function mask_immersed_field_xy!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(-)}, value=zero(eltype(bop))) +function mask_immersed_field_xy!(bop::BinaryOperation{<:Any, <:Any, <:Any, typeof(-)}, value=zero(eltype(bop)); k) a_value = ifelse(bop.b isa Number, bop.b, value) - mask_immersed_field_xy!(bop.a, a_value) + mask_immersed_field_xy!(bop.a, a_value; k) b_value = ifelse(bop.a isa Number, bop.a, value) - mask_immersed_field_xy!(bop.b, b_value) + mask_immersed_field_xy!(bop.b, b_value; k) return nothing end +# Fallback +mask_immersed_field_xy!(field, grid, loc, value, k) = nothing + """ - mask_immersed_field_xy!(field::Field, grid::ImmersedBoundaryGrid, loc, value; k, mask=peripheral_node) + mask_immersed_field_xy!(field::Field, grid::ImmersedBoundaryGrid, loc, value; k) -Mask `field` on `grid` with a `value` on the slices `[:, :, k]` where `mask` is `true`. +Mask `field` on `grid` with a `value` on the slices `[:, :, k]` where `immersed_peripheral_node` returns `true`. """ -function mask_immersed_field_xy!(field::Field, grid::ImmersedBoundaryGrid, loc, value; k, mask) +function mask_immersed_field_xy!(field::Field, grid::ImmersedBoundaryGrid, loc, value, k) arch = architecture(field) - loc = instantiate.(loc) - return launch!(arch, grid, :xy, - _mask_immersed_field_xy!, field, loc, grid, value, k, mask) + loc = instantiate.(loc) + return launch!(arch, grid, :xy, _mask_immersed_field_xy!, field, loc, grid, value, k) end -@kernel function _mask_immersed_field_xy!(field, loc, grid, value, k, mask) +@kernel function _mask_immersed_field_xy!(field, (ℓx, ℓy, ℓz), grid, value, k) i, j = @index(Global, NTuple) - @inbounds field[i, j, k] = scalar_mask(i, j, k, grid, grid.immersed_boundary, loc..., value, field, mask) + masked = immersed_peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) + @inbounds field[i, j, k] = ifelse(masked, value, field[i, j, k]) end ##### @@ -106,7 +110,6 @@ end # We mask a `ReducedField` if the entire reduced direction is immersed. # This requires a sweep over the reduced direction - function mask_immersed_field!(field::ReducedField, grid::ImmersedBoundaryGrid, loc, value) loc = instantiate.(loc) dims = reduced_dimensions(field) @@ -114,30 +117,25 @@ function mask_immersed_field!(field::ReducedField, grid::ImmersedBoundaryGrid, l return nothing end -@kernel function _mask_immersed_reduced_field!(field, dims, loc, grid, value) - i, j, k = @index(Global, NTuple) - mask = inactive_dimensions(i, j, k, grid, dims, loc) - @inbounds field[i, j, k] = ifelse(mask, value, field[i, j, k]) -end - -@inline inactive_search_range(i, grid, dim, dims) = ifelse(dim ∈ dims, 1:size(grid, dim), i:i) - -@inline function inactive_dimensions(i₀, j₀, k₀, grid, dims, loc) - mask = true +@kernel function _mask_immersed_reduced_field!(field, dims, (ℓx, ℓy, ℓz), grid, value) + i₀, j₀, k₀ = @index(Global, NTuple) + masked = true irange = inactive_search_range(i₀, grid, 1, dims) jrange = inactive_search_range(j₀, grid, 2, dims) krange = inactive_search_range(k₀, grid, 3, dims) # The loop activates over the whole direction only if reduced directions for i in irange, j in jrange, k in krange - mask = mask & peripheral_node(i, j, k, grid, loc...) + masked = masked & immersed_peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz) end - return mask + @inbounds field[i₀, j₀, k₀] = ifelse(masked, value, field[i₀, j₀, k₀]) end +@inline inactive_search_range(i, grid, dim, dims) = ifelse(dim ∈ dims, 1:size(grid, dim), i:i) + ### -### Efficient masking for `OnlyZReducedField` and an `AbstractGridFittedBoundary` +### Efficient masking for `OnlyZReducedField` and an `AbstractGridFittedBottom` ### const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom} @@ -147,12 +145,4 @@ const OnlyZReducedField = Field{<:CenterOrFace, <:CenterOrFace, Nothing} # Does not require a sweep mask_immersed_field!(field::OnlyZReducedField, grid::AGFBIBG, loc, value) = - mask_immersed_field_xy!(field, grid, loc, value; k=size(grid, 3), mask=peripheral_node) - -##### -##### Masking for GridFittedBoundary -##### - -@inline scalar_mask(i, j, k, grid, ::AbstractGridFittedBoundary, LX, LY, LZ, value, field, mask=peripheral_node) = - @inbounds ifelse(mask(i, j, k, grid, LX, LY, LZ), value, field[i, j, k]) - + mask_immersed_field_xy!(field, grid, loc, value, size(grid, 3)) \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index e96f53cb45..00a7dedb47 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -58,7 +58,7 @@ end # Mask immersed fields function mask_immersed_model_fields!(model, grid) - η = displacement(model.free_surface) + η = displacement(model.free_surface) fields_to_mask = merge(model.auxiliary_fields, prognostic_fields(model)) foreach(fields_to_mask) do field @@ -66,8 +66,8 @@ function mask_immersed_model_fields!(model, grid) mask_immersed_field!(field) end end - mask_immersed_field_xy!(η, k=size(grid, 3)+1, mask = inactive_node) - + mask_immersed_field_xy!(η, k=size(grid, 3)+1) + return nothing end From 78a874cfcd4ba679e6cf0d31dcb72fda5eda4f96 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Wed, 15 Jan 2025 14:09:40 +1100 Subject: [PATCH 19/31] add compat entries for Makie/MakieCore (#4045) --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 734e3e1ac1..abebca1b28 100644 --- a/Project.toml +++ b/Project.toml @@ -63,8 +63,8 @@ KernelAbstractions = "0.9.21" LinearAlgebra = "1.9" Logging = "1.9" MPI = "0.16, 0.17, 0.18, 0.19, 0.20" -Makie = "0.21" -MakieCore = "0.7, 0.8" +Makie = "0.21, 0.22" +MakieCore = "0.7, 0.8, 0.9" NCDatasets = "0.12.10, 0.13.1, 0.14" OffsetArrays = "1.4" OrderedCollections = "1.1" From c0a35e85585bce4779fb50a75432ab8cfc80b59a Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 17 Jan 2025 23:17:34 -0500 Subject: [PATCH 20/31] (0.95.7) Yet another ZStar implementation (#3956) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * static_column_depth * Update src/Grids/grid_utils.jl Co-authored-by: Gregory L. Wagner * address comments * new comment * another name change * AGFBIBG istead of AGFBIB and z_bottom only in TurbulenceClosures * some bugfixes * adapting to new interface * should change these operators * remove orthogonalsphericalshellgrids for the moment * better bottom height * bugfix in bottom height * better definition of bottom height * fixed partial cells? * fixed partial cell * fixed partial cells * fix split explicit tests * check if these pass on a GPU for the moment * remove OrthogonalSphericalShellGrids while we decide what to do * these files shouldn't go here * nothing was happening * correct coordinate test * new operators * fix pipeline * mpi test and gpu test * do we need to precompile it inside? * precompile inside the node * try previous climacommon version * go even more back * use the ClimaOcean implementation * using the ClimaOcean implementation * try removing boundary conditoin on barotropic velocity * found the bug * remove the show * check where it is failing * fairer comparison * see if this test passes * this was removed * maybe precompiling before... * double O0 * back to previous clima_common * another quick test * change environment * correct the utils * this should load mpitrampoline * Fix formatting * Go back to latest climacommon * try adding Manifest * Manifest from julia 1.10 * we probably need to initialize on a GPU * these options should not create problems * restart the tests * let's see if this differs * just version infos * fiddling with O0 * why are we using 8 threads? * memory requirements are not this huge * update enzyme * speed up the precompilation a bit, to revert later * might this be the culprit? * revert to 8 tasks to precompile * final version? * return to previous state of affairs * bugfix * cuda runtime version * bugfix * some bugfixes * more fixes * just need to fix immersed boundaries * new zstar * much better! * this should work nicely * more working * fx grid generation * delete zstar * should work like this * going on * some simplification * this to finish * should compile * better * add new definitions * should work? * This works * This seems to work * make sure implicit operator works * do not change this * changes * some changes * more changes * more changes * add some more tests * better * make it compile * make sure adapts work * some more corrections * a comment * remove duplicate field * where is zspacing??? * works * add cpu face constructor * correct aliases * at least some tests should pass * sloggin along * another bugfix * most of the tests should pass * remove partial cells from this business * transport the correct term * some more comment * better comment * some more bugfixes * add z face constructor for nothing * a couple of bugfixes * another bugfix * it works for all free surfaces!! * it works for all free surfaces * back to previous cpu face constructor * change vector invariant upwinding * no need for a loop * simplify * couple of fixes * some improvements * regression tests should pass * rework the constructor * simplify a bit * change face constructor * fix distributed tests * still debugging * fix gpu scripts * correct the face constructor * comment * all tests should be fixed now except for the distributed * last bugfix * now it should work * check which is the test that errors * remove stackoverflow * fix tests * revert file * revert file * removed rk3 by mistake * revert files * revert to rk3 * make sure we do not need to use regular grids * change e₃ to σ * chenga name * more corrections * add a test for active_cells_map * make the test a bit smaller * improve test * xor instead of or * correct all tests * on architecture * removed the drift in immersed boundary grids * better lock_release validation * remove extra validation examples * update tests * increase the time step * works for all free surfaces * test also rectilinear grids * better info msg * back to all the grids * deepcopy before integrating * make sure tests do not crash * remove space * use correct timesteps * zstar following coordinate * skip the test * chnage emojii * update to new syntax * align * skip the correct tests * remove correct tests * make sure we don't hit nans * bugfix * Update src/Grids/grid_generation.jl Co-authored-by: Gregory L. Wagner * Update src/Grids/grid_generation.jl Co-authored-by: Gregory L. Wagner * new name * corrections * Update hydrostatic_free_surface_ab2_step.jl * Update step_split_explicit_free_surface.jl * bugfix * add function support * add function support * remove convenience functionality * Update grid_generation.jl * add a docstring * Update test/test_active_cells_map.jl Co-authored-by: Gregory L. Wagner * Update test/test_split_explicit_vertical_integrals.jl Co-authored-by: Gregory L. Wagner * Update test/test_split_explicit_vertical_integrals.jl Co-authored-by: Gregory L. Wagner * start implementing suggestions * continue changing comments * simplify unscale tracer * all the refactor * some bugfix * remove HydrostaticFreeSurfaceModel * scale tracers for all moving coordinates * should work now * bump version * name changes, remove exporting and add test * make sure that the test is appropriate --------- Co-authored-by: Gregory L. Wagner --- .buildkite/pipeline.yml | 36 ++++ Project.toml | 2 +- src/Advection/Advection.jl | 1 + src/Advection/vector_invariant_advection.jl | 19 +- .../vector_invariant_cross_upwinding.jl | 26 ++- .../vector_invariant_self_upwinding.jl | 15 +- src/Grids/Grids.jl | 4 +- src/Grids/grid_generation.jl | 55 ++++- src/Grids/grid_utils.jl | 10 +- src/Grids/vertical_coordinates.jl | 109 ---------- src/Grids/vertical_discretization.jl | 181 ++++++++++++++++ src/ImmersedBoundaries/ImmersedBoundaries.jl | 1 + .../immersed_grid_metrics.jl | 4 +- .../mutable_immersed_grid.jl | 58 +++++ .../HydrostaticFreeSurfaceModels.jl | 9 +- .../SplitExplicitFreeSurfaces.jl | 6 + .../barotropic_split_explicit_corrector.jl | 48 +++-- .../initialize_split_explicit_substepping.jl | 2 +- .../split_explicit_timesteppers.jl | 4 +- .../step_split_explicit_free_surface.jl | 12 +- ...ute_hydrostatic_free_surface_tendencies.jl | 1 + .../compute_w_from_continuity.jl | 30 ++- .../hydrostatic_free_surface_ab2_step.jl | 34 ++- .../hydrostatic_free_surface_model.jl | 52 +++-- ..._free_surface_tendency_kernel_functions.jl | 4 + .../prescribed_hydrostatic_velocity_fields.jl | 1 + ...te_hydrostatic_free_surface_model_state.jl | 4 + .../z_star_vertical_spacing.jl | 158 ++++++++++++++ src/Models/Models.jl | 4 +- src/Operators/Operators.jl | 4 + .../spacings_and_areas_and_volumes.jl | 22 +- src/Operators/time_variable_grid_operators.jl | 57 +++++ .../vertically_implicit_diffusion_solver.jl | 2 +- test/runtests.jl | 9 +- test/test_active_cells_map.jl | 124 +++++++++++ .../test_split_explicit_vertical_integrals.jl | 30 +-- test/test_zstar_coordinate.jl | 201 ++++++++++++++++++ validation/z_star_coordinate/lock_release.jl | 80 +++++++ 38 files changed, 1199 insertions(+), 220 deletions(-) delete mode 100644 src/Grids/vertical_coordinates.jl create mode 100644 src/Grids/vertical_discretization.jl create mode 100644 src/ImmersedBoundaries/mutable_immersed_grid.jl create mode 100644 src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl create mode 100644 src/Operators/time_variable_grid_operators.jl create mode 100644 test/test_active_cells_map.jl create mode 100644 test/test_zstar_coordinate.jl create mode 100644 validation/z_star_coordinate/lock_release.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 8a162418cd..910d23d885 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -649,6 +649,42 @@ steps: limit: 1 depends_on: "init_cpu" +##### +##### Vertical Coordinates tests +##### + + - label: "🥑 gpu vertical coordinate" + env: + JULIA_DEPOT_PATH: "$SVERDRUP_HOME/.julia-$BUILDKITE_BUILD_NUMBER" + TEST_GROUP: "vertical_coordinate" + GPU_TEST: "true" + commands: + - "$SVERDRUP_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" + agents: + queue: Oceananigans + architecture: GPU + retry: + automatic: + - exit_status: 1 + limit: 1 + depends_on: "init_gpu" + + - label: "🥒 cpu vertical coordinate" + env: + JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER" + TEST_GROUP: "vertical_coordinate" + CUDA_VISIBLE_DEVICES: "-1" + commands: + - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'" + agents: + queue: Oceananigans + architecture: CPU + retry: + automatic: + - exit_status: 1 + limit: 1 + depends_on: "init_cpu" + ##### ##### Enzyme extension tests ##### diff --git a/Project.toml b/Project.toml index abebca1b28..a1eb9fd78c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -version = "0.95.6" +version = "0.95.7" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index 2a46564dd8..d528aa8610 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -34,6 +34,7 @@ using Oceananigans.Grids: with_halo using Oceananigans.Architectures: architecture, CPU using Oceananigans.Operators +using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ, ∂t_σ import Base: show, summary import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl index 5d5becb2f5..7f02532200 100644 --- a/src/Advection/vector_invariant_advection.jl +++ b/src/Advection/vector_invariant_advection.jl @@ -1,6 +1,3 @@ -using Oceananigans.Operators -using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ - # These are also used in Coriolis/hydrostatic_spherical_coriolis.jl struct EnergyConserving{FT} <: AbstractAdvectionScheme{1, FT} end struct EnstrophyConserving{FT} <: AbstractAdvectionScheme{1, FT} end @@ -168,7 +165,7 @@ Base.show(io::IO, a::VectorInvariant{N, FT}) where {N, FT} = ##### Convenience for WENO Vector Invariant ##### -nothing_to_default(user_value; default) = isnothing(user_value) ? default : user_value +nothing_to_default(user_value; default = nothing) = isnothing(user_value) ? default : user_value """ WENOVectorInvariant(FT = Float64; @@ -221,14 +218,14 @@ function WENOVectorInvariant(FT::DataType = Float64; default_upwinding = OnlySelfUpwinding(cross_scheme = divergence_scheme) upwinding = nothing_to_default(upwinding; default = default_upwinding) - N = max(required_halo_size_x(vorticity_scheme), - required_halo_size_y(vorticity_scheme), - required_halo_size_x(divergence_scheme), - required_halo_size_y(divergence_scheme), - required_halo_size_x(kinetic_energy_gradient_scheme), - required_halo_size_y(kinetic_energy_gradient_scheme), - required_halo_size_z(vertical_scheme)) + schemes = (vorticity_scheme, vertical_scheme, kinetic_energy_gradient_scheme, divergence_scheme) + + NX = maximum(required_halo_size_x(s) for s in schemes) + NY = maximum(required_halo_size_y(s) for s in schemes) + NZ = maximum(required_halo_size_z(s) for s in schemes) + N = max(NX, NY, NZ) + FT = eltype(vorticity_scheme) # assumption return VectorInvariant{N, FT, multi_dimensional_stencil}(vorticity_scheme, diff --git a/src/Advection/vector_invariant_cross_upwinding.jl b/src/Advection/vector_invariant_cross_upwinding.jl index c3f2f88fd9..045094bad3 100644 --- a/src/Advection/vector_invariant_cross_upwinding.jl +++ b/src/Advection/vector_invariant_cross_upwinding.jl @@ -17,20 +17,38 @@ ##### Cross and Self Upwinding of the Divergence flux ##### +# If the grid is moving, the discrete continuity equation is calculated as: +# +# ωᵏ⁺¹ - ωᵏ δx(Ax u) + δy(Ay v) Δrᶜᶜᶜ ∂t_σ +# ---------- = - --------------------- - ------------- +# Δzᶜᶜᶜ Vᶜᶜᶜ Δzᶜᶜᶜ +# +# Where ω is the vertical velocity with respect to a moving grid. +# We upwind the discrete divergence `δx(Ax u) + δy(Ay v)` and then divide by the volume, +# therefore, the correct term to be added to the divergence transport due to the moving grid is: +# +# Azᶜᶜᶜ Δrᶜᶜᶜ ∂t_σ +# +# which represents the static volume times the time derivative of the vertical grid scaling. +# If the grid is stationary, ∂t_σ evaluates to zero, so this term disappears from the divergence flux. +@inline Az_Δr_∂t_σ(i, j, k, grid) = Azᶜᶜᶜ(i, j, k, grid) * Δrᶜᶜᶜ(i, j, k, grid) * ∂t_σ(i, j, k, grid) + @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantCrossVerticalUpwinding, u, v) @inbounds û = u[i, j, k] δ_stencil = scheme.upwinding.divergence_stencil - δᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + δᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + ∂t_σ = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, Az_Δr_∂t_σ) - return û * δᴿ + return û * (δᴿ + ∂t_σ) # For static grids, ∂t_σ == 0 end @inline function upwinded_divergence_flux_Vᶜᶠᶜ(i, j, k, grid, scheme::VectorInvariantCrossVerticalUpwinding, u, v) @inbounds v̂ = v[i, j, k] δ_stencil = scheme.upwinding.divergence_stencil - δᴿ = _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + δᴿ = _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), flux_div_xyᶜᶜᶜ, δ_stencil, u, v) + ∂t_σ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, Az_Δr_∂t_σ) - return v̂ * δᴿ + return v̂ * (δᴿ + ∂t_σ) # For static grids, ∂t_σ == 0 end diff --git a/src/Advection/vector_invariant_self_upwinding.jl b/src/Advection/vector_invariant_self_upwinding.jl index e8ee512cfd..ad1a253832 100644 --- a/src/Advection/vector_invariant_self_upwinding.jl +++ b/src/Advection/vector_invariant_self_upwinding.jl @@ -2,15 +2,20 @@ ##### Self Upwinding of Divergence Flux, the best option! ##### -@inline δx_U(i, j, k, grid, u, v) = δxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, u) -@inline δy_V(i, j, k, grid, u, v) = δyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, v) +@inline δx_U(i, j, k, grid, u, v) = δxᶜᶜᶜ(i, j, k, grid, Ax_qᶠᶜᶜ, u) +@inline δy_V(i, j, k, grid, u, v) = δyᶜᶜᶜ(i, j, k, grid, Ay_qᶜᶠᶜ, v) + +# For moving grids, we include the time-derivative of the grid scaling in the divergence flux. +# If the grid is stationary, `Az_Δr_∂t_σ` evaluates to zero. +@inline δx_U_plus_∂t_σ(i, j, k, grid, u, v) = δxᶜᶜᶜ(i, j, k, grid, Ax_qᶠᶜᶜ, u) + Az_Δr_∂t_σ(i, j, k, grid) +@inline δy_V_plus_∂t_σ(i, j, k, grid, u, v) = δyᶜᶜᶜ(i, j, k, grid, Ay_qᶜᶠᶜ, v) + Az_Δr_∂t_σ(i, j, k, grid) # Velocity smoothness for divergence upwinding @inline U_smoothness(i, j, k, grid, u, v) = ℑxᶜᵃᵃ(i, j, k, grid, Ax_qᶠᶜᶜ, u) @inline V_smoothness(i, j, k, grid, u, v) = ℑyᵃᶜᵃ(i, j, k, grid, Ay_qᶜᶠᶜ, v) # Divergence smoothness for divergence upwinding -@inline divergence_smoothness(i, j, k, grid, u, v) = δx_U(i, j, k, grid, u, v) + δy_V(i, j, k, grid, u, v) +@inline divergence_smoothness(i, j, k, grid, u, v) = δx_U(i, j, k, grid, u, v) + δy_V(i, j, k, grid, u, v) @inline function upwinded_divergence_flux_Uᶠᶜᶜ(i, j, k, grid, scheme::VectorInvariantSelfVerticalUpwinding, u, v) @@ -18,7 +23,7 @@ cross_scheme = scheme.upwinding.cross_scheme @inbounds û = u[i, j, k] - δvˢ = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, δy_V, u, v) + δvˢ = _symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, cross_scheme, δy_V_plus_∂t_σ, u, v) δuᴿ = _biased_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(û), δx_U, δU_stencil, u, v) return û * (δvˢ + δuᴿ) @@ -30,7 +35,7 @@ end cross_scheme = scheme.upwinding.cross_scheme @inbounds v̂ = v[i, j, k] - δuˢ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, δx_U, u, v) + δuˢ = _symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, cross_scheme, δx_U_plus_∂t_σ, u, v) δvᴿ = _biased_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme, scheme.divergence_scheme, bias(v̂), δy_V, δV_stencil, u, v) return v̂ * (δuˢ + δvᴿ) diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 320e9035cc..4b331048d3 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -11,6 +11,7 @@ export XRegularRG, YRegularRG, ZRegularRG, XYRegularRG, XYZRegularRG export LatitudeLongitudeGrid, XRegularLLG, YRegularLLG, ZRegularLLG export OrthogonalSphericalShellGrid, ConformalCubedSphereGrid, ZRegOrthogonalSphericalShellGrid export conformal_cubed_sphere_panel +export MutableVerticalDiscretization export node, nodes export ξnode, ηnode, rnode export xnode, ynode, znode, λnode, φnode @@ -19,6 +20,7 @@ export spacings export xspacings, yspacings, zspacings, λspacings, φspacings, rspacings export minimum_xspacing, minimum_yspacing, minimum_zspacing export static_column_depthᶜᶜᵃ, static_column_depthᶠᶜᵃ, static_column_depthᶜᶠᵃ, static_column_depthᶠᶠᵃ +export column_depthᶜᶜᵃ, column_depthᶠᶜᵃ, column_depthᶜᶠᵃ, column_depthᶠᶠᵃ export offset_data, new_data export on_architecture @@ -118,7 +120,7 @@ struct ZDirection <: AbstractDirection end struct NegativeZDirection <: AbstractDirection end include("abstract_grid.jl") -include("vertical_coordinates.jl") +include("vertical_discretization.jl") include("grid_utils.jl") include("nodes_and_spacings.jl") include("zeros_and_ones.jl") diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl index fc90eeba92..e3c54e323b 100644 --- a/src/Grids/grid_generation.jl +++ b/src/Grids/grid_generation.jl @@ -88,7 +88,7 @@ function generate_coordinate(FT, topo::AT, N, H, node_generator, coordinate_name C = OffsetArray(on_architecture(arch, C.parent), C.offsets...) if coordinate_name == :z - return L, StaticVerticalCoordinate(F, C, Δᶠ, Δᶜ) + return L, StaticVerticalDiscretization(F, C, Δᶠ, Δᶜ) else return L, F, C, Δᶠ, Δᶜ end @@ -125,7 +125,7 @@ function generate_coordinate(FT, topo::AT, N, H, node_interval::Tuple{<:Number, C = OffsetArray(C, -H) if coordinate_name == :z - return FT(L), StaticVerticalCoordinate(F, C, FT(Δᶠ), FT(Δᶜ)) + return FT(L), StaticVerticalDiscretization(F, C, FT(Δᶠ), FT(Δᶜ)) else return FT(L), F, C, FT(Δᶠ), FT(Δᶜ) end @@ -134,7 +134,7 @@ end # Flat domains function generate_coordinate(FT, ::Flat, N, H, c::Number, coordinate_name, arch) if coordinate_name == :z - return FT(1), StaticVerticalCoordinate(range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1)) + return FT(1), StaticVerticalDiscretization(range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1)) else return FT(1), range(FT(c), FT(c), length=N), range(FT(c), FT(c), length=N), FT(1), FT(1) end @@ -145,8 +145,55 @@ end # FT(1), c, c, FT(1), FT(1) function generate_coordinate(FT, ::Flat, N, H, ::Nothing, coordinate_name, arch) if coordinate_name == :z - return FT(1), StaticVerticalCoordinate(nothing, nothing, FT(1), FT(1)) + return FT(1), StaticVerticalDiscretization(nothing, nothing, FT(1), FT(1)) else return FT(1), nothing, nothing, FT(1), FT(1) end end + +##### +##### MutableVerticalDiscretization +##### + +generate_coordinate(FT, ::Periodic, N, H, ::MutableVerticalDiscretization, coordinate_name, arch, args...) = + throw(ArgumentError("Periodic domains are not supported for MutableVerticalDiscretization")) + +# Generate a vertical coordinate with a scaling (`σ`) with respect to a reference coordinate `r` with spacing `Δr`. +# The grid might move with time, so the coordinate includes the time-derivative of the scaling `∂t_σ`. +# The value of the vertical coordinate at `Nz+1` is saved in `ηⁿ`. +function generate_coordinate(FT, topo, size, halo, coordinate::MutableVerticalDiscretization, coordinate_name, dim::Int, arch) + + Nx, Ny, Nz = size + Hx, Hy, Hz = halo + + if dim != 3 + msg = "MutableVerticalDiscretization is supported only in the third dimension (z)" + throw(ArgumentError(msg)) + end + + if coordinate_name != :z + msg = "MutableVerticalDiscretization is supported only for the z-coordinate" + throw(ArgumentError(msg)) + end + + r_faces = coordinate.cᵃᵃᶠ + + Lr, rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ = generate_coordinate(FT, topo[3](), Nz, Hz, r_faces, :r, arch) + + args = (topo, (Nx, Ny, Nz), (Hx, Hy, Hz)) + + σᶜᶜ⁻ = new_data(FT, arch, (Center, Center, Nothing), args...) + σᶜᶜⁿ = new_data(FT, arch, (Center, Center, Nothing), args...) + σᶠᶜⁿ = new_data(FT, arch, (Face, Center, Nothing), args...) + σᶜᶠⁿ = new_data(FT, arch, (Center, Face, Nothing), args...) + σᶠᶠⁿ = new_data(FT, arch, (Face, Face, Nothing), args...) + ηⁿ = new_data(FT, arch, (Center, Center, Nothing), args...) + ∂t_σ = new_data(FT, arch, (Center, Center, Nothing), args...) + + # Fill all the scalings with one for now (i.e. z == r) + for σ in (σᶜᶜ⁻, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ) + fill!(σ, 1) + end + + return Lr, MutableVerticalDiscretization(rᵃᵃᶠ, rᵃᵃᶜ, Δrᵃᵃᶠ, Δrᵃᵃᶜ, ηⁿ, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ∂t_σ) +end diff --git a/src/Grids/grid_utils.jl b/src/Grids/grid_utils.jl index 92c3aec5b8..7bf74fbdeb 100644 --- a/src/Grids/grid_utils.jl +++ b/src/Grids/grid_utils.jl @@ -299,7 +299,7 @@ end function dimension_summary(topo, name, dom, z::AbstractVerticalCoordinate, pad_domain=0) prefix = domain_summary(topo, name, dom) padding = " "^(pad_domain+1) - return string(prefix, padding, coordinate_summary(topo, z.Δᵃᵃᶜ, name)) + return string(prefix, padding, coordinate_summary(topo, z, name)) end function dimension_summary(topo, name, dom, spacing, pad_domain=0) @@ -317,7 +317,7 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) = name, prettysummary(maximum(parent(Δ)))) ##### -##### Static column depths +##### Static and Dynamic column depths ##### @inline static_column_depthᶜᶜᵃ(i, j, grid) = grid.Lz @@ -325,6 +325,12 @@ coordinate_summary(topo, Δ::Union{AbstractVector, AbstractMatrix}, name) = @inline static_column_depthᶠᶜᵃ(i, j, grid) = grid.Lz @inline static_column_depthᶠᶠᵃ(i, j, grid) = grid.Lz +# Will be extended in the `ImmersedBoundaries` module for a ``mutable'' grid type +@inline column_depthᶜᶜᵃ(i, j, k, grid, η) = static_column_depthᶜᶜᵃ(i, j, grid) +@inline column_depthᶠᶜᵃ(i, j, k, grid, η) = static_column_depthᶠᶜᵃ(i, j, grid) +@inline column_depthᶜᶠᵃ(i, j, k, grid, η) = static_column_depthᶜᶠᵃ(i, j, grid) +@inline column_depthᶠᶠᵃ(i, j, k, grid, η) = static_column_depthᶠᶠᵃ(i, j, grid) + ##### ##### Spherical geometry ##### diff --git a/src/Grids/vertical_coordinates.jl b/src/Grids/vertical_coordinates.jl deleted file mode 100644 index 3028fb013b..0000000000 --- a/src/Grids/vertical_coordinates.jl +++ /dev/null @@ -1,109 +0,0 @@ -#### -#### Vertical coordinates -#### - -# This file implements everything related to vertical coordinates in Oceananigans. -# Vertical coordinates are independent of the underlying grid type since only grids that are -# "unstructured" or "curvilinear" in the horizontal directions are supported in Oceananigans. -# Thus the vertical coordinate is _special_, and it can be implemented once for all grid types. - -abstract type AbstractVerticalCoordinate end - -""" - struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate - -Represent a static one-dimensional vertical coordinate. - -Fields -====== - -- `cᶜ::C`: Cell-centered coordinate. -- `cᶠ::D`: Face-centered coordinate. -- `Δᶜ::E`: Cell-centered grid spacing. -- `Δᶠ::F`: Face-centered grid spacing. -""" -struct StaticVerticalCoordinate{C, D, E, F} <: AbstractVerticalCoordinate - cᵃᵃᶠ :: C - cᵃᵃᶜ :: D - Δᵃᵃᶠ :: E - Δᵃᵃᶜ :: F -end - -#### -#### Some useful aliases -#### - -const RegularVerticalCoordinate = StaticVerticalCoordinate{<:Any, <:Any, <:Number, <:Number} -const RegularVerticalGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate} - -#### -#### Adapt and on_architecture -#### - -Adapt.adapt_structure(to, coord::StaticVerticalCoordinate) = - StaticVerticalCoordinate(Adapt.adapt(to, coord.cᵃᵃᶠ), - Adapt.adapt(to, coord.cᵃᵃᶜ), - Adapt.adapt(to, coord.Δᵃᵃᶠ), - Adapt.adapt(to, coord.Δᵃᵃᶜ)) - -on_architecture(arch, coord::StaticVerticalCoordinate) = - StaticVerticalCoordinate(on_architecture(arch, coord.cᵃᵃᶠ), - on_architecture(arch, coord.cᵃᵃᶜ), - on_architecture(arch, coord.Δᵃᵃᶠ), - on_architecture(arch, coord.Δᵃᵃᶜ)) - -##### -##### Nodes and spacings (common to every grid)... -##### - -AUG = AbstractUnderlyingGrid - -@inline rnode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(k, grid, ℓz) -@inline rnode(k, grid, ::Center) = getnode(grid.z.cᵃᵃᶜ, k) -@inline rnode(k, grid, ::Face) = getnode(grid.z.cᵃᵃᶠ, k) - -# These will be extended in the Operators module -@inline znode(k, grid, ℓz) = rnode(k, grid, ℓz) -@inline znode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(i, j, k, grid, ℓx, ℓy, ℓz) - -@inline rnodes(grid::AUG, ℓz::Face; with_halos=false) = _property(grid.z.cᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) -@inline rnodes(grid::AUG, ℓz::Center; with_halos=false) = _property(grid.z.cᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) -@inline rnodes(grid::AUG, ℓx, ℓy, ℓz; with_halos=false) = rnodes(grid, ℓz; with_halos) - -rnodes(grid::AUG, ::Nothing; kwargs...) = 1:1 -znodes(grid::AUG, ::Nothing; kwargs...) = 1:1 - -# TODO: extend in the Operators module -@inline znodes(grid::AUG, ℓz; kwargs...) = rnodes(grid, ℓz; kwargs...) -@inline znodes(grid::AUG, ℓx, ℓy, ℓz; kwargs...) = rnodes(grid, ℓx, ℓy, ℓz; kwargs...) - -function rspacings end -function zspacings end - -@inline rspacings(grid, ℓz) = rspacings(grid, nothing, nothing, ℓz) -@inline zspacings(grid, ℓz) = zspacings(grid, nothing, nothing, ℓz) - -#### -#### `z_domain` and `cpu_face_constructor_z` -#### - -z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.z.cᵃᵃᶠ) - -@inline cpu_face_constructor_r(grid::RegularVerticalGrid) = z_domain(grid) - -@inline function cpu_face_constructor_r(grid) - Nz = size(grid, 3) - nodes = rnodes(grid, Face(); with_halos=true) - cpu_nodes = on_architecture(CPU(), nodes) - return cpu_nodes[1:Nz+1] -end - -@inline cpu_face_constructor_z(grid) = cpu_face_constructor_r(grid) - -#### -#### Utilities -#### - -# Summaries -coordinate_summary(::Bounded, z::AbstractVerticalCoordinate, name) = - @sprintf("Free-surface following with Δ%s=%s", name, prettysummary(z.Δᵃᵃᶜ)) diff --git a/src/Grids/vertical_discretization.jl b/src/Grids/vertical_discretization.jl new file mode 100644 index 0000000000..d622d58ce8 --- /dev/null +++ b/src/Grids/vertical_discretization.jl @@ -0,0 +1,181 @@ +#### +#### Vertical coordinates +#### + +# This file implements everything related to vertical coordinates in Oceananigans. +# Vertical coordinates are independent of the underlying grid type since only grids that are +# "unstructured" or "curvilinear" in the horizontal directions are supported in Oceananigans. +# Thus the vertical coordinate is _special_, and it can be implemented once for all grid types. + +abstract type AbstractVerticalCoordinate end + +""" + struct StaticVerticalDiscretization{C, D, E, F} <: AbstractVerticalCoordinate + +Represent a static one-dimensional vertical coordinate. + +Fields +====== + +- `cᶜ::C`: Cell-centered coordinate. +- `cᶠ::D`: Face-centered coordinate. +- `Δᶜ::E`: Cell-centered grid spacing. +- `Δᶠ::F`: Face-centered grid spacing. +""" +struct StaticVerticalDiscretization{C, D, E, F} <: AbstractVerticalCoordinate + cᵃᵃᶠ :: C + cᵃᵃᶜ :: D + Δᵃᵃᶠ :: E + Δᵃᵃᶜ :: F +end + +struct MutableVerticalDiscretization{C, D, E, F, H, CC, FC, CF, FF} <: AbstractVerticalCoordinate + cᵃᵃᶠ :: C + cᵃᵃᶜ :: D + Δᵃᵃᶠ :: E + Δᵃᵃᶜ :: F + ηⁿ :: H + σᶜᶜⁿ :: CC + σᶠᶜⁿ :: FC + σᶜᶠⁿ :: CF + σᶠᶠⁿ :: FF + σᶜᶜ⁻ :: CC + ∂t_σ :: CC +end + +""" + MutableVerticalDiscretization(r_faces) + +Construct a `MutableVerticalDiscretization` from `r_faces` that can be a `Tuple`, a function of an index `k`, +or an `AbstractArray`. A `MutableVerticalDiscretization` defines a vertical coordinate that might evolve in time +following certain rules. Examples of `MutableVerticalDiscretization`s are free-surface following coordinates, +or sigma coordinates. +""" +MutableVerticalDiscretization(r_faces) = MutableVerticalDiscretization(r_faces, r_faces, [nothing for i in 1:9]...) + +#### +#### Some useful aliases +#### + +const RegularStaticVerticalDiscretization = StaticVerticalDiscretization{<:Any, <:Any, <:Number} +const RegularMutableVerticalDiscretization = MutableVerticalDiscretization{<:Any, <:Any, <:Number} + +const RegularVerticalCoordinate = Union{RegularStaticVerticalDiscretization, RegularMutableVerticalDiscretization} + +const AbstractMutableGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Bounded, <:MutableVerticalDiscretization} +const AbstractStaticGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:StaticVerticalDiscretization} +const RegularVerticalGrid = AbstractUnderlyingGrid{<:Any, <:Any, <:Any, <:Any, <:RegularVerticalCoordinate} + +#### +#### Adapt and on_architecture +#### + +Adapt.adapt_structure(to, coord::StaticVerticalDiscretization) = + StaticVerticalDiscretization(Adapt.adapt(to, coord.cᵃᵃᶠ), + Adapt.adapt(to, coord.cᵃᵃᶜ), + Adapt.adapt(to, coord.Δᵃᵃᶠ), + Adapt.adapt(to, coord.Δᵃᵃᶜ)) + +on_architecture(arch, coord::StaticVerticalDiscretization) = + StaticVerticalDiscretization(on_architecture(arch, coord.cᵃᵃᶠ), + on_architecture(arch, coord.cᵃᵃᶜ), + on_architecture(arch, coord.Δᵃᵃᶠ), + on_architecture(arch, coord.Δᵃᵃᶜ)) + +Adapt.adapt_structure(to, coord::MutableVerticalDiscretization) = + MutableVerticalDiscretization(Adapt.adapt(to, coord.cᵃᵃᶠ), + Adapt.adapt(to, coord.cᵃᵃᶜ), + Adapt.adapt(to, coord.Δᵃᵃᶠ), + Adapt.adapt(to, coord.Δᵃᵃᶜ), + Adapt.adapt(to, coord.ηⁿ), + Adapt.adapt(to, coord.σᶜᶜⁿ), + Adapt.adapt(to, coord.σᶠᶜⁿ), + Adapt.adapt(to, coord.σᶜᶠⁿ), + Adapt.adapt(to, coord.σᶠᶠⁿ), + Adapt.adapt(to, coord.σᶜᶜ⁻), + Adapt.adapt(to, coord.∂t_σ)) + +on_architecture(arch, coord::MutableVerticalDiscretization) = + MutableVerticalDiscretization(on_architecture(arch, coord.cᵃᵃᶠ), + on_architecture(arch, coord.cᵃᵃᶜ), + on_architecture(arch, coord.Δᵃᵃᶠ), + on_architecture(arch, coord.Δᵃᵃᶜ), + on_architecture(arch, coord.ηⁿ), + on_architecture(arch, coord.σᶜᶜⁿ), + on_architecture(arch, coord.σᶠᶜⁿ), + on_architecture(arch, coord.σᶜᶠⁿ), + on_architecture(arch, coord.σᶠᶠⁿ), + on_architecture(arch, coord.σᶜᶜ⁻), + on_architecture(arch, coord.∂t_σ)) + +##### +##### Nodes and spacings (common to every grid)... +##### + +AUG = AbstractUnderlyingGrid + +@inline rnode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(k, grid, ℓz) +@inline rnode(k, grid, ::Center) = getnode(grid.z.cᵃᵃᶜ, k) +@inline rnode(k, grid, ::Face) = getnode(grid.z.cᵃᵃᶠ, k) + +# These will be extended in the Operators module +@inline znode(k, grid, ℓz) = rnode(k, grid, ℓz) +@inline znode(i, j, k, grid, ℓx, ℓy, ℓz) = rnode(i, j, k, grid, ℓx, ℓy, ℓz) + +@inline rnodes(grid::AUG, ℓz::Face; with_halos=false) = _property(grid.z.cᵃᵃᶠ, ℓz, topology(grid, 3), size(grid, 3), with_halos) +@inline rnodes(grid::AUG, ℓz::Center; with_halos=false) = _property(grid.z.cᵃᵃᶜ, ℓz, topology(grid, 3), size(grid, 3), with_halos) +@inline rnodes(grid::AUG, ℓx, ℓy, ℓz; with_halos=false) = rnodes(grid, ℓz; with_halos) + +rnodes(grid::AUG, ::Nothing; kwargs...) = 1:1 +znodes(grid::AUG, ::Nothing; kwargs...) = 1:1 + +# TODO: extend in the Operators module +@inline znodes(grid::AUG, ℓz; kwargs...) = rnodes(grid, ℓz; kwargs...) +@inline znodes(grid::AUG, ℓx, ℓy, ℓz; kwargs...) = rnodes(grid, ℓx, ℓy, ℓz; kwargs...) + +function rspacings end +function zspacings end + +@inline rspacings(grid, ℓz) = rspacings(grid, nothing, nothing, ℓz) +@inline zspacings(grid, ℓz) = zspacings(grid, nothing, nothing, ℓz) + +#### +#### `z_domain` and `cpu_face_constructor_z` +#### + +z_domain(grid) = domain(topology(grid, 3)(), grid.Nz, grid.z.cᵃᵃᶠ) + +@inline cpu_face_constructor_r(grid::RegularVerticalGrid) = z_domain(grid) + +@inline function cpu_face_constructor_r(grid) + Nz = size(grid, 3) + nodes = rnodes(grid, Face(); with_halos=true) + cpu_nodes = on_architecture(CPU(), nodes) + return cpu_nodes[1:Nz+1] +end + +@inline cpu_face_constructor_z(grid) = cpu_face_constructor_r(grid) +@inline cpu_face_constructor_z(grid::AbstractMutableGrid) = MutableVerticalDiscretization(cpu_face_constructor_r(grid)) + +#### +#### Utilities +#### + +function validate_dimension_specification(T, ξ::MutableVerticalDiscretization, dir, N, FT) + cᶠ = validate_dimension_specification(T, ξ.cᵃᵃᶠ, dir, N, FT) + cᶜ = validate_dimension_specification(T, ξ.cᵃᵃᶜ, dir, N, FT) + args = Tuple(getproperty(ξ, prop) for prop in propertynames(ξ)) + + return MutableVerticalDiscretization(cᶠ, cᶜ, args[3:end]...) +end + +# Summaries +coordinate_summary(topo, z::StaticVerticalDiscretization, name) = coordinate_summary(topo, z.Δᵃᵃᶜ, name) + +coordinate_summary(::Bounded, z::RegularMutableVerticalDiscretization, name) = + @sprintf("regularly spaced with Δr=%s (mutable)", prettysummary(z.Δᵃᵃᶜ)) + +coordinate_summary(::Bounded, z::MutableVerticalDiscretization, name) = + @sprintf("variably spaced with min(Δr)=%s, max(Δr)=%s (mutable)", + prettysummary(minimum(z.Δᵃᵃᶜ)), + prettysummary(maximum(z.Δᵃᵃᶜ))) diff --git a/src/ImmersedBoundaries/ImmersedBoundaries.jl b/src/ImmersedBoundaries/ImmersedBoundaries.jl index 2b45f4af9a..87b655861f 100644 --- a/src/ImmersedBoundaries/ImmersedBoundaries.jl +++ b/src/ImmersedBoundaries/ImmersedBoundaries.jl @@ -44,5 +44,6 @@ include("immersed_boundary_condition.jl") include("conditional_differences.jl") include("mask_immersed_field.jl") include("immersed_reductions.jl") +include("mutable_immersed_grid.jl") end # module diff --git a/src/ImmersedBoundaries/immersed_grid_metrics.jl b/src/ImmersedBoundaries/immersed_grid_metrics.jl index 54cd4b4eb9..7857d20016 100644 --- a/src/ImmersedBoundaries/immersed_grid_metrics.jl +++ b/src/ImmersedBoundaries/immersed_grid_metrics.jl @@ -1,6 +1,6 @@ using Oceananigans.AbstractOperations: GridMetricOperation -import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ +import Oceananigans.Operators: Δrᵃᵃᶠ, Δrᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ import Oceananigans.Operators: Δxᶠᵃᵃ, Δxᶜᵃᵃ, Δxᶠᶜᵃ, Δxᶜᶠᵃ, Δxᶠᶠᵃ, Δxᶜᶜᵃ import Oceananigans.Operators: Δyᵃᶠᵃ, Δyᵃᶜᵃ, Δyᶠᶜᵃ, Δyᶜᶠᵃ, Δyᶠᶠᵃ, Δyᶜᶜᵃ import Oceananigans.Operators: Azᶠᶜᵃ, Azᶜᶠᵃ, Azᶠᶠᵃ, Azᶜᶜᵃ @@ -22,8 +22,8 @@ import Oceananigans.Operators: intrinsic_vector, extrinsic_vector @inline Δrᵃᵃᶠ(i, j, k, ibg::IBG) = Δrᵃᵃᶠ(i, j, k, ibg.underlying_grid) @inline Δrᵃᵃᶜ(i, j, k, ibg::IBG) = Δrᵃᵃᶜ(i, j, k, ibg.underlying_grid) -@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid) @inline Δzᵃᵃᶜ(i, j, k, ibg::IBG) = Δzᵃᵃᶜ(i, j, k, ibg.underlying_grid) +@inline Δzᵃᵃᶠ(i, j, k, ibg::IBG) = Δzᵃᵃᶠ(i, j, k, ibg.underlying_grid) # 1D Horizontal spacings diff --git a/src/ImmersedBoundaries/mutable_immersed_grid.jl b/src/ImmersedBoundaries/mutable_immersed_grid.jl new file mode 100644 index 0000000000..c4b9f73b57 --- /dev/null +++ b/src/ImmersedBoundaries/mutable_immersed_grid.jl @@ -0,0 +1,58 @@ +using Oceananigans.Grids: AbstractMutableGrid +using Oceananigans.Operators +using Oceananigans.Operators: MRG, MLLG, MOSG, superscript_location + +import Oceananigans.Grids: column_depthᶜᶜᵃ, + column_depthᶜᶠᵃ, + column_depthᶠᶜᵃ, + column_depthᶠᶠᵃ + +import Oceananigans.Operators: σⁿ, σ⁻, ∂t_σ + +const MutableImmersedGrid = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:AbstractMutableGrid} +const MutableGridOfSomeKind = Union{MutableImmersedGrid, AbstractMutableGrid} + +@inline column_depthᶜᶜᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶜᶜᵃ(i, j, grid) + @inbounds η[i, j, k] +@inline column_depthᶠᶜᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶠᶜᵃ(i, j, grid) + ℑxᶠᵃᵃ(i, j, k, grid, η) +@inline column_depthᶜᶠᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶜᶠᵃ(i, j, grid) + ℑyᵃᶠᵃ(i, j, k, grid, η) +@inline column_depthᶠᶠᵃ(i, j, k, grid::MutableGridOfSomeKind, η) = static_column_depthᶠᶠᵃ(i, j, grid) + ℑxyᶠᶠᵃ(i, j, k, grid, η) + +# Convenience methods +@inline column_depthᶜᶜᵃ(i, j, grid) = static_column_depthᶜᶜᵃ(i, j, grid) +@inline column_depthᶜᶠᵃ(i, j, grid) = static_column_depthᶜᶠᵃ(i, j, grid) +@inline column_depthᶠᶜᵃ(i, j, grid) = static_column_depthᶠᶜᵃ(i, j, grid) +@inline column_depthᶠᶠᵃ(i, j, grid) = static_column_depthᶠᶠᵃ(i, j, grid) + +@inline column_depthᶜᶜᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶜᶜᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline column_depthᶜᶠᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶜᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline column_depthᶠᶜᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶠᶜᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline column_depthᶠᶠᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) + +# Fallbacks +@inline σⁿ(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = σⁿ(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) +@inline σ⁻(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = σ⁻(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz) + +@inline ∂t_σ(i, j, k, ibg::IBG) = ∂t_σ(i, j, k, ibg.underlying_grid) + +# Extend the 3D vertical spacing operators on an Immersed Mutable grid +const IMRG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MRG} +const IMLLG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MLLG} +const IMOSG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:MOSG} + +for LX in (:ᶠ, :ᶜ), LY in (:ᶠ, :ᶜ), LZ in (:ᶠ, :ᶜ) + zspacing = Symbol(:Δz, LX, LY, LZ) + rspacing = Symbol(:Δr, LX, LY, LZ) + + ℓx = superscript_location(LX) + ℓy = superscript_location(LY) + ℓz = superscript_location(LZ) + + @eval begin + using Oceananigans.Operators: $rspacing + import Oceananigans.Operators: $zspacing + + @inline $zspacing(i, j, k, grid::IMRG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::IMLLG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::IMOSG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + end +end \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index 2c17fc47d5..7ab1fc8d08 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -3,7 +3,7 @@ module HydrostaticFreeSurfaceModels export HydrostaticFreeSurfaceModel, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface, - PrescribedVelocityFields + PrescribedVelocityFields, ZStar, ZCoordinate using KernelAbstractions: @index, @kernel using KernelAbstractions.Extras.LoopInfo: @unroll @@ -23,6 +23,9 @@ using Oceananigans.TimeSteppers: SplitRungeKutta3TimeStepper, QuasiAdamsBashfort abstract type AbstractFreeSurface{E, G} end +struct ZCoordinate end +struct ZStar end + # This is only used by the cubed sphere for now. fill_horizontal_velocity_halos!(args...) = nothing @@ -57,6 +60,10 @@ include("SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl") using .SplitExplicitFreeSurfaces +# ZStar implementation +include("z_star_vertical_spacing.jl") + +# Hydrostatic model implementation include("hydrostatic_free_surface_field_tuples.jl") include("hydrostatic_free_surface_model.jl") include("show_hydrostatic_free_surface_model.jl") diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl index 6ce20fd0b7..b2408e695e 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl @@ -11,6 +11,7 @@ using Oceananigans.Utils using Oceananigans.Grids using Oceananigans.Operators using Oceananigans.BoundaryConditions +using Oceananigans.ImmersedBoundaries using Oceananigans.Grids: AbstractGrid, topology using Oceananigans.ImmersedBoundaries: active_linear_index_to_tuple, mask_immersed_field! using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface, @@ -21,6 +22,11 @@ using Base using KernelAbstractions: @index, @kernel using KernelAbstractions.Extras.LoopInfo: @unroll +using Oceananigans.Grids: column_depthᶜᶜᵃ, + column_depthᶜᶠᵃ, + column_depthᶠᶜᵃ, + column_depthᶠᶠᵃ + import Oceananigans.Models.HydrostaticFreeSurfaceModels: initialize_free_surface!, materialize_free_surface, step_free_surface!, diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl index 2b7f7845f6..5b31e5f98c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl @@ -1,39 +1,54 @@ # Kernels to compute the vertical integral of the velocities -@kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v) +@kernel function _barotropic_mode_kernel!(U̅, V̅, grid, ::Nothing, u, v, η) i, j = @index(Global, NTuple) - barotropic_mode_kernel!(U, V, i, j, grid, u, v) + barotropic_mode_kernel!(U̅, V̅, i, j, grid, u, v, η) end -@kernel function _barotropic_mode_kernel!(U, V, grid, active_cells_map, u, v) +@kernel function _barotropic_mode_kernel!(U̅, V̅, grid, active_cells_map, u, v, η) idx = @index(Global, Linear) i, j = active_linear_index_to_tuple(idx, active_cells_map) - barotropic_mode_kernel!(U, V, i, j, grid, u, v) + barotropic_mode_kernel!(U̅, V̅, i, j, grid, u, v, η) end -@inline function barotropic_mode_kernel!(U, V, i, j, grid, u, v) - @inbounds U[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] - @inbounds V[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] +@inline function barotropic_mode_kernel!(U̅, V̅, i, j, grid, u, v, η) + k_top = size(grid, 3) + 1 + + hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) + hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) + + # If the static depths are zero (i.e. the column is immersed), + # we set the grid scaling factor to 1 + # (There is no free surface on an immersed column (η == 0)) + σᶠᶜ = ifelse(hᶠᶜ == 0, one(grid), Hᶠᶜ / hᶠᶜ) + σᶜᶠ = ifelse(hᶜᶠ == 0, one(grid), Hᶜᶠ / hᶜᶠ) + + @inbounds U̅[i, j, 1] = Δrᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1] * σᶠᶜ + @inbounds V̅[i, j, 1] = Δrᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1] * σᶜᶠ for k in 2:grid.Nz - @inbounds U[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k] - @inbounds V[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k] + @inbounds U̅[i, j, 1] += Δrᶠᶜᶜ(i, j, k, grid) * u[i, j, k] * σᶠᶜ + @inbounds V̅[i, j, 1] += Δrᶜᶠᶜ(i, j, k, grid) * v[i, j, k] * σᶜᶠ end return nothing end -@inline function compute_barotropic_mode!(U, V, grid, u, v) +@inline function compute_barotropic_mode!(U̅, V̅, grid, u, v, η) active_cells_map = retrieve_surface_active_cells_map(grid) - launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map) + launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U̅, V̅, grid, active_cells_map, u, v, η; active_cells_map) return nothing end -@kernel function _barotropic_split_explicit_corrector!(u, v, U, V, U̅, V̅, grid) +@kernel function _barotropic_split_explicit_corrector!(u, v, U, V, U̅, V̅, η, grid) i, j, k = @index(Global, NTuple) + k_top = size(grid, 3) + 1 @inbounds begin - Hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) - Hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) u[i, j, k] = u[i, j, k] + (U[i, j, 1] - U̅[i, j, 1]) / Hᶠᶜ v[i, j, k] = v[i, j, k] + (V[i, j, 1] - V̅[i, j, 1]) / Hᶜᶠ @@ -43,6 +58,7 @@ end # Correcting `u` and `v` with the barotropic mode computed in `free_surface` function barotropic_split_explicit_corrector!(u, v, free_surface, grid) state = free_surface.filtered_state + η = free_surface.η U, V = free_surface.barotropic_velocities U̅, V̅ = state.U, state.V arch = architecture(grid) @@ -50,11 +66,11 @@ function barotropic_split_explicit_corrector!(u, v, free_surface, grid) # NOTE: the filtered `U̅` and `V̅` have been copied in the instantaneous `U` and `V`, # so we use the filtered velocities as "work arrays" to store the vertical integrals # of the instantaneous velocities `u` and `v`. - compute_barotropic_mode!(U̅, V̅, grid, u, v) + compute_barotropic_mode!(U̅, V̅, grid, u, v, η) # add in "good" barotropic mode launch!(arch, grid, :xyz, _barotropic_split_explicit_corrector!, - u, v, U, V, U̅, V̅, grid) + u, v, U, V, U̅, V̅, η, grid) return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl index 4c2976591f..ea49c3591c 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl @@ -14,7 +14,7 @@ using Oceananigans.Operators: Δz # from the initial velocity conditions. function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velocities) barotropic_velocities = sefs.barotropic_velocities - @apply_regionally compute_barotropic_mode!(barotropic_velocities.U, barotropic_velocities.V, grid, velocities.u, velocities.v) + @apply_regionally compute_barotropic_mode!(barotropic_velocities.U, barotropic_velocities.V, grid, velocities.u, velocities.v, sefs.η) fill_halo_regions!((barotropic_velocities.U, barotropic_velocities.V)) return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl index e624054d5e..a450e99de6 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl @@ -135,10 +135,10 @@ function initialize_free_surface_timestepper!(timestepper::AdamsBashforth3Scheme end # The functions `η★` `U★` and `V★` represent the value of free surface, barotropic zonal and meridional velocity at time step m+1/2 -@inline U★(i, j, k, grid, t::ForwardBackwardScheme, Uᵐ) = @inbounds Uᵐ[i, j, k] +@inline U★(i, j, k, grid, ::ForwardBackwardScheme, Uᵐ) = @inbounds Uᵐ[i, j, k] @inline U★(i, j, k, grid, t::AdamsBashforth3Scheme, Uᵐ) = @inbounds t.α * Uᵐ[i, j, k] + t.θ * t.Uᵐ⁻¹[i, j, k] + t.β * t.Uᵐ⁻²[i, j, k] -@inline η★(i, j, k, grid, t::ForwardBackwardScheme, ηᵐ⁺¹) = @inbounds ηᵐ⁺¹[i, j, k] +@inline η★(i, j, k, grid, ::ForwardBackwardScheme, ηᵐ⁺¹) = @inbounds ηᵐ⁺¹[i, j, k] @inline η★(i, j, k, grid, t::AdamsBashforth3Scheme, ηᵐ⁺¹) = @inbounds t.δ * ηᵐ⁺¹[i, j, k] + t.μ * t.ηᵐ[i, j, k] + t.γ * t.ηᵐ⁻¹[i, j, k] + t.ϵ * t.ηᵐ⁻²[i, j, k] @inline store_previous_velocities!(::ForwardBackwardScheme, i, j, k, U) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl index 655e2e9068..21c3b11a40 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl @@ -1,3 +1,5 @@ +using Oceananigans.ImmersedBoundaries: MutableGridOfSomeKind + # Evolution Kernels # # ∂t(η) = -∇⋅U @@ -26,8 +28,8 @@ end store_previous_velocities!(timestepper, i, j, 1, U) store_previous_velocities!(timestepper, i, j, 1, V) - Hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) - Hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) @inbounds begin # ∂τ(U) = - ∇η + G @@ -164,5 +166,11 @@ function step_free_surface!(free_surface::SplitExplicitFreeSurface, model, baroc mask_immersed_field!(model.velocities.v) end + # Needed for Mutable to compute the barotropic correction. + # TODO: Would it be possible to remove it in some way? + if model.grid isa MutableGridOfSomeKind + fill_halo_regions!(η) + end + return nothing end diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl index 29303f6581..481b962df0 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl @@ -140,6 +140,7 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para model.diffusivity_fields, model.pressure.pHY′, model.auxiliary_fields, + model.vertical_coordinate, model.clock) u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args..., u_forcing) diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl index 927489cf50..4e633d9499 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_w_from_continuity.jl @@ -1,7 +1,8 @@ using Oceananigans.Architectures: device using Oceananigans.Grids: halo_size, topology using Oceananigans.Grids: XFlatGrid, YFlatGrid -using Oceananigans.Operators: div_xyᶜᶜᶜ, Δzᶜᶜᶜ +using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, div_xyᶜᶜᶜ, Δzᶜᶜᶜ +using Oceananigans.ImmersedBoundaries: immersed_cell """ compute_w_from_continuity!(model) @@ -18,12 +19,37 @@ compute_w_from_continuity!(model; kwargs...) = compute_w_from_continuity!(velocities, arch, grid; parameters = w_kernel_parameters(grid)) = launch!(arch, grid, parameters, _compute_w_from_continuity!, velocities, grid) +# If the grid is following the free surface, then the derivative of the moving grid is: +# +# δx(Δy U) + δy(Δx V) ∇ ⋅ U +# ∂t_σ = - --------------------- = - -------- +# Az ⋅ H H +# +# The discrete divergence is then calculated as: +# +# wᵏ⁺¹ - wᵏ δx(Ax u) + δy(Ay v) Δr ∂t_σ +# ---------- = - --------------------- - ---------- +# Δz vol Δz +# +# This makes sure that summing up till the top of the domain, results in: +# +# ∇ ⋅ U +# wᴺᶻ⁺¹ = w⁰ - ------- - ∂t_σ ≈ 0 (if w⁰ == 0) +# H +# +# If the grid is static, then ∂t_σ = 0 and the moving grid contribution is equal to zero @kernel function _compute_w_from_continuity!(U, grid) i, j = @index(Global, NTuple) @inbounds U.w[i, j, 1] = 0 for k in 2:grid.Nz+1 - @inbounds U.w[i, j, k] = U.w[i, j, k-1] - Δzᶜᶜᶜ(i, j, k-1, grid) * div_xyᶜᶜᶜ(i, j, k-1, grid, U.u, U.v) + δh_u = flux_div_xyᶜᶜᶜ(i, j, k-1, grid, U.u, U.v) / Azᶜᶜᶜ(i, j, k-1, grid) + ∂tσ = Δrᶜᶜᶜ(i, j, k-1, grid) * ∂t_σ(i, j, k-1, grid) + + immersed = immersed_cell(i, j, k-1, grid) + Δw = δh_u + ifelse(immersed, zero(grid), ∂tσ) # We do not account for grid changes in immersed cells + + @inbounds U.w[i, j, k] = U.w[i, j, k-1] - Δw end end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl index ad0dbe784a..1f7fa9a40f 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl @@ -90,8 +90,7 @@ function ab2_step_tracers!(tracers, model, Δt, χ) tracer_field = tracers[tracer_name] closure = model.closure - launch!(model.architecture, model.grid, :xyz, - ab2_step_field!, tracer_field, Δt, χ, Gⁿ, G⁻) + ab2_step_tracer_field!(tracer_field, model.grid, Δt, χ, Gⁿ, G⁻) implicit_step!(tracer_field, model.timestepper.implicit_solver, @@ -106,3 +105,34 @@ function ab2_step_tracers!(tracers, model, Δt, χ) return nothing end +ab2_step_tracer_field!(tracer_field, grid, Δt, χ, Gⁿ, G⁻) = + launch!(architecture(grid), grid, :xyz, _ab2_step_tracer_field!, tracer_field, grid, Δt, χ, Gⁿ, G⁻) + +##### +##### Tracer update in mutable vertical coordinates +##### + +# σθ is the evolved quantity. Once σⁿ⁺¹ is known we can retrieve θⁿ⁺¹ +# with the `unscale_tracers!` function +@kernel function _ab2_step_tracer_field!(θ, grid, Δt, χ, Gⁿ, G⁻) + i, j, k = @index(Global, NTuple) + + FT = eltype(χ) + α = convert(FT, 1.5) + χ + β = convert(FT, 0.5) + χ + + σᶜᶜⁿ = σⁿ(i, j, k, grid, Center(), Center(), Center()) + σᶜᶜ⁻ = σ⁻(i, j, k, grid, Center(), Center(), Center()) + + @inbounds begin + ∂t_σθ = α * σᶜᶜⁿ * Gⁿ[i, j, k] - β * σᶜᶜ⁻ * G⁻[i, j, k] + + # We store temporarily σθ in θ. + # The unscaled θ will be retrieved with `unscale_tracers!` + θ[i, j, k] = σᶜᶜⁿ * θ[i, j, k] + convert(FT, Δt) * ∂t_σθ + end +end + +# Fallback! We need to unscale the tracers only in case of +# a grid with a mutable vertical coordinate, i.e. where `σ != 1` +unscale_tracers!(tracers, grid; kwargs...) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 354b273017..675db7f6ba 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -26,25 +26,26 @@ const ParticlesOrNothing = Union{Nothing, AbstractLagrangianParticles} const AbstractBGCOrNothing = Union{Nothing, AbstractBiogeochemistry} mutable struct HydrostaticFreeSurfaceModel{TS, E, A<:AbstractArchitecture, S, - G, T, V, B, R, F, P, BGC, U, C, Φ, K, AF} <: AbstractModel{TS} - - architecture :: A # Computer `Architecture` on which `Model` is run - grid :: G # Grid of physical points on which `Model` is solved - clock :: Clock{T} # Tracks iteration number and simulation time of `Model` - advection :: V # Advection scheme for tracers - buoyancy :: B # Set of parameters for buoyancy model - coriolis :: R # Set of parameters for the background rotation rate of `Model` - free_surface :: S # Free surface parameters and fields - forcing :: F # Container for forcing functions defined by the user - closure :: E # Diffusive 'turbulence closure' for all model fields - particles :: P # Particle set for Lagrangian tracking - biogeochemistry :: BGC # Biogeochemistry for Oceananigans tracers - velocities :: U # Container for velocity fields `u`, `v`, and `w` - tracers :: C # Container for tracer fields - pressure :: Φ # Container for hydrostatic pressure - diffusivity_fields :: K # Container for turbulent diffusivities - timestepper :: TS # Object containing timestepper fields and parameters - auxiliary_fields :: AF # User-specified auxiliary fields for forcing functions and boundary conditions + G, T, V, B, R, F, P, BGC, U, C, Φ, K, AF, Z} <: AbstractModel{TS} + + architecture :: A # Computer `Architecture` on which `Model` is run + grid :: G # Grid of physical points on which `Model` is solved + clock :: Clock{T} # Tracks iteration number and simulation time of `Model` + advection :: V # Advection scheme for tracers + buoyancy :: B # Set of parameters for buoyancy model + coriolis :: R # Set of parameters for the background rotation rate of `Model` + free_surface :: S # Free surface parameters and fields + forcing :: F # Container for forcing functions defined by the user + closure :: E # Diffusive 'turbulence closure' for all model fields + particles :: P # Particle set for Lagrangian tracking + biogeochemistry :: BGC # Biogeochemistry for Oceananigans tracers + velocities :: U # Container for velocity fields `u`, `v`, and `w` + tracers :: C # Container for tracer fields + pressure :: Φ # Container for hydrostatic pressure + diffusivity_fields :: K # Container for turbulent diffusivities + timestepper :: TS # Object containing timestepper fields and parameters + auxiliary_fields :: AF # User-specified auxiliary fields for forcing functions and boundary conditions + vertical_coordinate :: Z # Rulesets that define the time-evolution of the grid end default_free_surface(grid::XYRegularRG; gravitational_acceleration=g_Earth) = @@ -71,7 +72,8 @@ default_free_surface(grid; gravitational_acceleration=g_Earth) = velocities = nothing, pressure = nothing, diffusivity_fields = nothing, - auxiliary_fields = NamedTuple()) + auxiliary_fields = NamedTuple(), + vertical_coordinate = ZCoordinate()) Construct a hydrostatic model with a free surface on `grid`. @@ -103,6 +105,7 @@ Keyword arguments - `pressure`: Hydrostatic pressure field. Default: `nothing`. - `diffusivity_fields`: Diffusivity fields. Default: `nothing`. - `auxiliary_fields`: `NamedTuple` of auxiliary fields. Default: `nothing`. + - `vertical_coordinate`: Rulesets that define the time-evolution of the grid (ZStar/ZCoordinate). Default: `ZCoordinate()`. """ function HydrostaticFreeSurfaceModel(; grid, clock = Clock{eltype(grid)}(time = 0), @@ -121,11 +124,16 @@ function HydrostaticFreeSurfaceModel(; grid, velocities = nothing, pressure = nothing, diffusivity_fields = nothing, - auxiliary_fields = NamedTuple()) + auxiliary_fields = NamedTuple(), + vertical_coordinate = ZCoordinate()) # Check halos and throw an error if the grid's halo is too small @apply_regionally validate_model_halo(grid, momentum_advection, tracer_advection, closure) + if !(grid isa MutableGridOfSomeKind) && (vertical_coordinate isa ZStar) + error("The grid does not support ZStar vertical coordinates. Use a `MutableVerticalDiscretization` to allow the use of ZStar (see `MutableVerticalDiscretization`).") + end + # Validate biogeochemistry (add biogeochemical tracers automagically) tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example) biogeochemical_fields = merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)) @@ -208,7 +216,7 @@ function HydrostaticFreeSurfaceModel(; grid, model = HydrostaticFreeSurfaceModel(arch, grid, clock, advection, buoyancy, coriolis, free_surface, forcing, closure, particles, biogeochemistry, velocities, tracers, - pressure, diffusivity_fields, timestepper, auxiliary_fields) + pressure, diffusivity_fields, timestepper, auxiliary_fields, vertical_coordinate) update_state!(model; compute_tendencies = false) diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl index 28cb289e58..9447b3effd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -38,6 +38,7 @@ implicitly during time-stepping. diffusivities, hydrostatic_pressure_anomaly, auxiliary_fields, + ztype, clock, forcing) @@ -47,6 +48,7 @@ implicitly during time-stepping. - explicit_barotropic_pressure_x_gradient(i, j, k, grid, free_surface) - x_f_cross_U(i, j, k, grid, coriolis, velocities) - ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) + - grid_slope_contribution_x(i, j, k, grid, buoyancy, ztype, model_fields) - ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy) - immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields) + forcing(i, j, k, grid, clock, model_fields)) @@ -77,6 +79,7 @@ implicitly during time-stepping. diffusivities, hydrostatic_pressure_anomaly, auxiliary_fields, + ztype, clock, forcing) @@ -86,6 +89,7 @@ implicitly during time-stepping. - explicit_barotropic_pressure_y_gradient(i, j, k, grid, free_surface) - y_f_cross_U(i, j, k, grid, coriolis, velocities) - ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure_anomaly) + - grid_slope_contribution_y(i, j, k, grid, buoyancy, ztype, model_fields) - ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy) - immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields) + forcing(i, j, k, grid, clock, model_fields)) diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index 1e44b9bea1..b63b797321 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -100,6 +100,7 @@ extract_boundary_conditions(::PrescribedVelocityFields) = NamedTuple() free_surface_displacement_field(::PrescribedVelocityFields, ::Nothing, grid) = nothing HorizontalVelocityFields(::PrescribedVelocityFields, grid) = nothing, nothing +materialize_free_surface(::Nothing, velocities, grid) = nothing materialize_free_surface(::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing materialize_free_surface(::ImplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing materialize_free_surface(::SplitExplicitFreeSurface, ::PrescribedVelocityFields, grid) = nothing diff --git a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl index 00a7dedb47..7ab34ff634 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl @@ -84,6 +84,10 @@ function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = P = model.pressure.pHY′ arch = architecture(grid) + # Update the grid and unscale the tracers + update_grid!(model, grid, model.vertical_coordinate; parameters = w_parameters) + unscale_tracers!(tracers, grid; parameters = w_parameters) + # Advance diagnostic quantities compute_w_from_continuity!(model; parameters = w_parameters) update_hydrostatic_pressure!(P, arch, grid, buoyancy, tracers; parameters = p_parameters) diff --git a/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl b/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl new file mode 100644 index 0000000000..9f69f148ca --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/z_star_vertical_spacing.jl @@ -0,0 +1,158 @@ +using Oceananigans.Grids +using Oceananigans.ImmersedBoundaries: MutableGridOfSomeKind + +##### +##### Mutable-specific vertical spacings update +##### + +# The easy case +barotropic_velocities(free_surface::SplitExplicitFreeSurface) = free_surface.barotropic_velocities + +# The "harder" case, barotropic velocities are computed on the fly +barotropic_velocities(free_surface) = nothing, nothing + +# Fallback +update_grid!(model, grid, ztype; parameters) = nothing + +function update_grid!(model, grid::MutableGridOfSomeKind, ::ZStar; parameters = :xy) + + # Scalings and free surface + σᶜᶜ⁻ = grid.z.σᶜᶜ⁻ + σᶜᶜⁿ = grid.z.σᶜᶜⁿ + σᶠᶜⁿ = grid.z.σᶠᶜⁿ + σᶜᶠⁿ = grid.z.σᶜᶠⁿ + σᶠᶠⁿ = grid.z.σᶠᶠⁿ + ∂t_σ = grid.z.∂t_σ + ηⁿ = grid.z.ηⁿ + η = model.free_surface.η + + launch!(architecture(grid), grid, parameters, _update_grid_scaling!, + σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ηⁿ, grid, η) + + # the barotropic velocities are retrieved from the free surface model for a + # SplitExplicitFreeSurface and are calculated for other free surface models + U, V = barotropic_velocities(model.free_surface) + u, v, _ = model.velocities + + # Update the time derivative of the vertical spacing, + # No need to fill the halo as the scaling is updated _IN_ the halos + launch!(architecture(grid), grid, parameters, _update_grid_vertical_velocity!, ∂t_σ, grid, U, V, u, v) + + return nothing +end + +@kernel function _update_grid_scaling!(σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ηⁿ, grid, η) + i, j = @index(Global, NTuple) + k_top = size(grid, 3) + 1 + + hᶜᶜ = static_column_depthᶜᶜᵃ(i, j, grid) + hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid) + hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid) + hᶠᶠ = static_column_depthᶠᶠᵃ(i, j, grid) + + Hᶜᶜ = column_depthᶜᶜᵃ(i, j, k_top, grid, η) + Hᶠᶜ = column_depthᶠᶜᵃ(i, j, k_top, grid, η) + Hᶜᶠ = column_depthᶜᶠᵃ(i, j, k_top, grid, η) + Hᶠᶠ = column_depthᶠᶠᵃ(i, j, k_top, grid, η) + + @inbounds begin + σᶜᶜ = ifelse(hᶜᶜ == 0, one(grid), Hᶜᶜ / hᶜᶜ) + σᶠᶜ = ifelse(hᶠᶜ == 0, one(grid), Hᶠᶜ / hᶠᶜ) + σᶜᶠ = ifelse(hᶜᶠ == 0, one(grid), Hᶜᶠ / hᶜᶠ) + σᶠᶠ = ifelse(hᶠᶠ == 0, one(grid), Hᶠᶠ / hᶠᶠ) + + # Update previous scaling + σᶜᶜ⁻[i, j, 1] = σᶜᶜⁿ[i, j, 1] + + # update current scaling + σᶜᶜⁿ[i, j, 1] = σᶜᶜ + σᶠᶜⁿ[i, j, 1] = σᶠᶜ + σᶜᶠⁿ[i, j, 1] = σᶜᶠ + σᶠᶠⁿ[i, j, 1] = σᶠᶠ + + # Update η in the grid + ηⁿ[i, j, 1] = η[i, j, k_top] + end +end + +@kernel function _update_grid_vertical_velocity!(∂t_σ, grid, U, V, u, v) + i, j = @index(Global, NTuple) + kᴺ = size(grid, 3) + + hᶜᶜ = static_column_depthᶜᶜᵃ(i, j, grid) + + # ∂(η / H)/∂t = - ∇ ⋅ ∫udz / H + δx_U = δxᶜᶜᶜ(i, j, kᴺ, grid, Δy_qᶠᶜᶜ, barotropic_U, U, u) + δy_V = δyᶜᶜᶜ(i, j, kᴺ, grid, Δx_qᶜᶠᶜ, barotropic_V, V, v) + + δh_U = (δx_U + δy_V) / Azᶜᶜᶜ(i, j, kᴺ, grid) + + @inbounds ∂t_σ[i, j, 1] = ifelse(hᶜᶜ == 0, zero(grid), - δh_U / hᶜᶜ) +end + +# If U and V exist, we just take them +@inline barotropic_U(i, j, k, grid, U, u) = @inbounds U[i, j, k] +@inline barotropic_V(i, j, k, grid, V, v) = @inbounds V[i, j, k] + +# If U and V are not available, we compute them +@inline function barotropic_U(i, j, k, grid, ::Nothing, u) + U = 0 + for k in 1:size(grid, 3) + U += u[i, j, k] * Δzᶠᶜᶜ(i, j, k, grid) + end + return U +end + +@inline function barotropic_V(i, j, k, grid, ::Nothing, v) + V = 0 + for k in 1:size(grid, 3) + V += v[i, j, k] * Δzᶜᶠᶜ(i, j, k, grid) + end + return V +end + +##### +##### ZStar-specific implementation of the additional terms to be included in the momentum equations +##### + +# Fallbacks +@inline grid_slope_contribution_x(i, j, k, grid, buoyancy, ztype, model_fields) = zero(grid) +@inline grid_slope_contribution_y(i, j, k, grid, buoyancy, ztype, model_fields) = zero(grid) + +@inline grid_slope_contribution_x(i, j, k, grid::MutableGridOfSomeKind, ::Nothing, ::ZStar, model_fields) = zero(grid) +@inline grid_slope_contribution_y(i, j, k, grid::MutableGridOfSomeKind, ::Nothing, ::ZStar, model_fields) = zero(grid) + +@inline ∂x_z(i, j, k, grid) = @inbounds ∂xᶠᶜᶜ(i, j, k, grid, znode, Center(), Center(), Center()) +@inline ∂y_z(i, j, k, grid) = @inbounds ∂yᶜᶠᶜ(i, j, k, grid, znode, Center(), Center(), Center()) + +@inline grid_slope_contribution_x(i, j, k, grid::MutableGridOfSomeKind, buoyancy, ::ZStar, model_fields) = + ℑxᶠᵃᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.formulation, model_fields) * ∂x_z(i, j, k, grid) + +@inline grid_slope_contribution_y(i, j, k, grid::MutableGridOfSomeKind, buoyancy, ::ZStar, model_fields) = + ℑyᵃᶠᵃ(i, j, k, grid, buoyancy_perturbationᶜᶜᶜ, buoyancy.formulation, model_fields) * ∂y_z(i, j, k, grid) + +#### +#### Removing the scaling of the vertical coordinate from the tracer fields +#### + +const EmptyTuples = Union{NamedTuple{(), Tuple{}}, Tuple{}} + +unscale_tracers!(::EmptyTuples, ::MutableGridOfSomeKind; kwargs...) = nothing + +function unscale_tracers!(tracers, grid::MutableGridOfSomeKind; parameters = :xy) + + for tracer in tracers + launch!(architecture(grid), grid, parameters, _unscale_tracer!, + tracer, grid, Val(grid.Hz), Val(grid.Nz)) + end + + return nothing +end + +@kernel function _unscale_tracer!(tracer, grid, ::Val{Hz}, ::Val{Nz}) where {Hz, Nz} + i, j = @index(Global, NTuple) + + @unroll for k in -Hz+1:Nz+Hz + tracer[i, j, k] /= σⁿ(i, j, k, grid, Center(), Center(), Center()) + end +end diff --git a/src/Models/Models.jl b/src/Models/Models.jl index b7dbdf0702..268927dbe9 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -3,7 +3,7 @@ module Models export NonhydrostaticModel, ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation, - HydrostaticFreeSurfaceModel, + HydrostaticFreeSurfaceModel, ZStar, ZCoordinate, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface, PrescribedVelocityFields, PressureField, LagrangianParticles, @@ -105,7 +105,7 @@ using .NonhydrostaticModels: NonhydrostaticModel, PressureField using .HydrostaticFreeSurfaceModels: HydrostaticFreeSurfaceModel, ExplicitFreeSurface, ImplicitFreeSurface, SplitExplicitFreeSurface, - PrescribedVelocityFields + PrescribedVelocityFields, ZStar, ZCoordinate using .ShallowWaterModels: ShallowWaterModel, ConservativeFormulation, VectorInvariantFormulation diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index ffd9eeb1e2..7be1572ccf 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -77,6 +77,9 @@ export ∂xTᶠᶜᶠ, ∂yTᶜᶠᶠ # Reference frame conversion export intrinsic_vector, extrinsic_vector +# Variable grid operators +export σⁿ, σ⁻, ∂t_σ + using Oceananigans.Grids ##### @@ -119,6 +122,7 @@ include("topology_aware_operators.jl") include("vorticity_operators.jl") include("laplacian_operators.jl") +include("time_variable_grid_operators.jl") include("vector_rotation_operators.jl") @inline xspacing(args...) = Δx(args...) diff --git a/src/Operators/spacings_and_areas_and_volumes.jl b/src/Operators/spacings_and_areas_and_volumes.jl index a7df404aca..9f8f866a3e 100644 --- a/src/Operators/spacings_and_areas_and_volumes.jl +++ b/src/Operators/spacings_and_areas_and_volumes.jl @@ -35,17 +35,6 @@ The operators in this file fall into three categories: ##### ##### -const ZRG = Union{LLGZ, RGZ, OSSGZ} - -@inline getspacing(k, Δz::AbstractVector) = @inbounds Δz[k] -@inline getspacing(k, Δz::Number) = @inbounds Δz - -@inline Δrᵃᵃᶜ(i, j, k, grid::AbstractGrid) = getspacing(k, grid.z.Δᵃᵃᶜ) -@inline Δrᵃᵃᶠ(i, j, k, grid::AbstractGrid) = getspacing(k, grid.z.Δᵃᵃᶠ) - -@inline Δzᵃᵃᶜ(i, j, k, grid::AbstractGrid) = Δrᵃᵃᶜ(i, j, k, grid) -@inline Δzᵃᵃᶠ(i, j, k, grid::AbstractGrid) = Δrᵃᵃᶠ(i, j, k, grid) - # This metaprogramming loop defines all possible combinations of locations for spacings. # The general 2D and 3D spacings are reconducted to their one - dimensional counterparts. # Grids that do not have a specific one - dimensional spacing for a given location need to @@ -114,11 +103,14 @@ end ##### One - dimensional Vertical spacing (same for all grids) ##### -@inline Δzᵃᵃᶜ(i, j, k, grid) = @inbounds grid.z.Δᵃᵃᶜ[k] -@inline Δzᵃᵃᶠ(i, j, k, grid) = @inbounds grid.z.Δᵃᵃᶠ[k] +@inline getspacing(k, Δz::Number) = Δz +@inline getspacing(k, Δz::AbstractVector) = @inbounds Δz[k] + +@inline Δrᵃᵃᶜ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶜ) +@inline Δrᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ) -@inline Δzᵃᵃᶜ(i, j, k, grid::ZRG) = grid.z.Δᵃᵃᶜ -@inline Δzᵃᵃᶠ(i, j, k, grid::ZRG) = grid.z.Δᵃᵃᶠ +@inline Δzᵃᵃᶜ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶜ) +@inline Δzᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ) ##### ##### diff --git a/src/Operators/time_variable_grid_operators.jl b/src/Operators/time_variable_grid_operators.jl new file mode 100644 index 0000000000..0dbdf1a959 --- /dev/null +++ b/src/Operators/time_variable_grid_operators.jl @@ -0,0 +1,57 @@ +import Oceananigans.Grids: znode, AbstractMutableGrid + +##### +##### MutableVerticalDiscretization-specific vertical spacing functions +##### + +const C = Center +const F = Face + +const AMG = AbstractMutableGrid + +# Fallbacks +@inline σⁿ(i, j, k, grid, ℓx, ℓy, ℓz) = one(grid) +@inline σ⁻(i, j, k, grid, ℓx, ℓy, ℓz) = one(grid) + +@inline σⁿ(i, j, k, grid::AMG, ::C, ::C, ℓz) = @inbounds grid.z.σᶜᶜⁿ[i, j, 1] +@inline σⁿ(i, j, k, grid::AMG, ::F, ::C, ℓz) = @inbounds grid.z.σᶠᶜⁿ[i, j, 1] +@inline σⁿ(i, j, k, grid::AMG, ::C, ::F, ℓz) = @inbounds grid.z.σᶜᶠⁿ[i, j, 1] +@inline σⁿ(i, j, k, grid::AMG, ::F, ::F, ℓz) = @inbounds grid.z.σᶠᶠⁿ[i, j, 1] + +# σ⁻ is needed only at centers +@inline σ⁻(i, j, k, grid::AMG, ::C, ::C, ℓz) = @inbounds grid.z.σᶜᶜ⁻[i, j, 1] + +@inline ∂t_σ(i, j, k, grid) = zero(grid) +@inline ∂t_σ(i, j, k, grid::AMG) = @inbounds grid.z.∂t_σ[i, j, 1] + +#### +#### Vertical spacing functions +#### + +const MRG = RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalDiscretization} +const MLLG = LatitudeLongitudeGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalDiscretization} +const MOSG = OrthogonalSphericalShellGrid{<:Any, <:Any, <:Any, <:Any, <:MutableVerticalDiscretization} + +superscript_location(s::Symbol) = s == :ᶜ ? :Center : :Face + +for LX in (:ᶠ, :ᶜ), LY in (:ᶠ, :ᶜ), LZ in (:ᶠ, :ᶜ) + zspacing = Symbol(:Δz, LX, LY, LZ) + rspacing = Symbol(:Δr, LX, LY, LZ) + + ℓx = superscript_location(LX) + ℓy = superscript_location(LY) + ℓz = superscript_location(LZ) + + @eval begin + @inline $zspacing(i, j, k, grid::MRG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::MLLG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + @inline $zspacing(i, j, k, grid::MOSG) = $rspacing(i, j, k, grid) * σⁿ(i, j, k, grid, $ℓx(), $ℓy(), $ℓz()) + end +end + +# znode for an AbstractMutableGrid grid is the reference node (`rnode`) scaled by the derivative with respect to the reference (σⁿ) +# added to the surface value of `z` (which we here call ηⁿ) +@inline znode(i, j, k, grid::AMG, ::C, ::C, ℓz) = rnode(i, j, k, grid, C(), C(), ℓz) * σⁿ(i, j, k, grid, C(), C(), ℓz) + @inbounds grid.z.ηⁿ[i, j, 1] +@inline znode(i, j, k, grid::AMG, ::F, ::C, ℓz) = rnode(i, j, k, grid, F(), C(), ℓz) * σⁿ(i, j, k, grid, F(), C(), ℓz) + ℑxᶠᵃᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline znode(i, j, k, grid::AMG, ::C, ::F, ℓz) = rnode(i, j, k, grid, C(), F(), ℓz) * σⁿ(i, j, k, grid, C(), F(), ℓz) + ℑyᵃᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) +@inline znode(i, j, k, grid::AMG, ::F, ::F, ℓz) = rnode(i, j, k, grid, F(), F(), ℓz) * σⁿ(i, j, k, grid, F(), F(), ℓz) + ℑxyᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ) diff --git a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl index 86795b9e24..3aaa551622 100644 --- a/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl +++ b/src/TurbulenceClosures/vertically_implicit_diffusion_solver.jl @@ -46,7 +46,7 @@ const c = Center() const f = Face() # The vertical spacing used here is Δz for velocities and Δr for tracers, since the -# implicit solver operator is applied to the scaled tracer e₃θ instead of just θ +# implicit solver operator is applied to the scaled tracer σθ instead of just θ @inline vertical_spacing(i, j, k, grid, ℓx, ℓy, ℓz) = Δz(i, j, k, grid, ℓx, ℓy, ℓz) @inline vertical_spacing(i, j, k, grid, ::Center, ::Center, ℓz) = Δr(i, j, k, grid, c, c, ℓz) diff --git a/test/runtests.jl b/test/runtests.jl index a47f0e3d3d..fbfbe6dac7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -41,7 +41,7 @@ CUDA.allowscalar() do CUDA.versioninfo() catch; end end - + # Core Oceananigans if group == :unit || group == :all @testset "Unit tests" begin @@ -120,6 +120,7 @@ CUDA.allowscalar() do @testset "Model and time stepping tests (part 1)" begin include("test_nonhydrostatic_models.jl") include("test_time_stepping.jl") + include("test_active_cells_map.jl") end end @@ -219,6 +220,12 @@ CUDA.allowscalar() do end end + if group == :vertical_coordinate || group == :all + @testset "Vertical coordinate tests" begin + include("test_zstar_coordinate.jl") + end + end + # Tests for Enzyme extension if group == :enzyme || group == :all @testset "Enzyme extension tests" begin diff --git a/test/test_active_cells_map.jl b/test/test_active_cells_map.jl new file mode 100644 index 0000000000..d1533b9451 --- /dev/null +++ b/test/test_active_cells_map.jl @@ -0,0 +1,124 @@ +include("dependencies_for_runtests.jl") + +using Oceananigans.Operators: hack_cosd +using Oceananigans.ImmersedBoundaries: retrieve_surface_active_cells_map, + retrieve_interior_active_cells_map, + immersed_cell + +function Δ_min(grid) + Δx_min = minimum_xspacing(grid, Center(), Center(), Center()) + Δy_min = minimum_yspacing(grid, Center(), Center(), Center()) + return min(Δx_min, Δy_min) +end + +@inline Gaussian(x, y, L) = exp(-(x^2 + y^2) / L^2) + +function solid_body_rotation_test(grid) + + free_surface = SplitExplicitFreeSurface(grid; substeps = 10, gravitational_acceleration = 1) + coriolis = HydrostaticSphericalCoriolis(rotation_rate = 1) + + model = HydrostaticFreeSurfaceModel(; grid, + momentum_advection = VectorInvariant(), + free_surface = free_surface, + coriolis = coriolis, + tracers = :c, + tracer_advection = WENO(), + buoyancy = nothing, + closure = nothing) + + g = model.free_surface.gravitational_acceleration + R = grid.radius + Ω = model.coriolis.rotation_rate + + uᵢ(λ, φ, z) = 0.1 * cosd(φ) * sind(λ) + ηᵢ(λ, φ, z) = (R * Ω * 0.1 + 0.1^2 / 2) * sind(φ)^2 / g * sind(λ) + + # Gaussian leads to values with O(1e-60), + # too small for repetible testing. We cap it at 0.1 + cᵢ(λ, φ, z) = max(Gaussian(λ, φ - 5, 10), 0.1) + vᵢ(λ, φ, z) = 0.1 + + set!(model, u=uᵢ, η=ηᵢ, c=cᵢ) + + Δt = 0.1 * Δ_min(grid) / sqrt(g * grid.Lz) + + simulation = Simulation(model; Δt, stop_iteration = 10) + run!(simulation) + + return merge(model.velocities, model.tracers, (; η = model.free_surface.η)) +end + +Nx = 16 +Ny = 16 +Nz = 10 + +@testset "Active cells map" begin + for arch in archs + underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, Nz), + halo = (4, 4, 4), + latitude = (-80, 80), + longitude = (-160, 160), + z = (-10, 0), + radius = 1, + topology = (Bounded, Bounded, Bounded)) + + # Make sure the bottom is the same + bottom_height = zeros(Nx, Ny) + for i in 1:Nx, j in 1:Ny + bottom_height[i, j] = - rand() * 5 - 5 + end + + bottom_height = on_architecture(arch, bottom_height) + immersed_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height)) + immersed_active_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height); active_cells_map = true) + + @testset "Active cells map construction" begin + surface_active_cells_map = retrieve_surface_active_cells_map(immersed_active_grid) + interior_active_cells_map = retrieve_interior_active_cells_map(immersed_active_grid, Val(:interior)) + + surface_active_cells_map = on_architecture(CPU(), surface_active_cells_map) + interior_active_cells_map = on_architecture(CPU(), interior_active_cells_map) + grid = on_architecture(CPU(), immersed_grid) + + for i in 1:Nx, j in 1:Ny, k in 1:Nz + immersed = immersed_cell(i, j, k, grid) + active = (i, j, k) ∈ interior_active_cells_map + @test immersed ⊻ active + end + + for i in 1:Nx, j in 1:Ny + immersed = all(immersed_cell(i, j, k, grid) for k in 1:Nz) + active = (i, j) ∈ surface_active_cells_map + @test immersed ⊻ active + end + end + + @testset "Active cells map solid body rotation" begin + + ua, va, wa, ca, ηa = solid_body_rotation_test(immersed_active_grid) + u, v, w, c, η = solid_body_rotation_test(immersed_grid) + + ua = interior(on_architecture(CPU(), ua)) + va = interior(on_architecture(CPU(), va)) + wa = interior(on_architecture(CPU(), wa)) + ca = interior(on_architecture(CPU(), ca)) + ηa = interior(on_architecture(CPU(), ηa)) + + u = interior(on_architecture(CPU(), u)) + v = interior(on_architecture(CPU(), v)) + w = interior(on_architecture(CPU(), w)) + c = interior(on_architecture(CPU(), c)) + η = interior(on_architecture(CPU(), η)) + + atol = eps(eltype(immersed_grid)) + rtol = sqrt(eps(eltype(immersed_grid))) + + @test all(isapprox(u, ua; atol, rtol)) + @test all(isapprox(v, va; atol, rtol)) + @test all(isapprox(w, wa; atol, rtol)) + @test all(isapprox(c, ca; atol, rtol)) + @test all(isapprox(η, ηa; atol, rtol)) + end + end +end \ No newline at end of file diff --git a/test/test_split_explicit_vertical_integrals.jl b/test/test_split_explicit_vertical_integrals.jl index 045531ed67..62f9acf6f6 100644 --- a/test/test_split_explicit_vertical_integrals.jl +++ b/test/test_split_explicit_vertical_integrals.jl @@ -35,7 +35,9 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces @testset "Average to zero" begin # set equal to something else - η̅ .= U̅ .= V̅ .= 1.0 + η̅ .= 1 + U̅ .= 1 + V̅ .= 1 # now set equal to zero initialize_free_surface_state!(sefs, sefs.timestepper, sefs.timestepper, Val(1)) @@ -47,8 +49,8 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces # check @test all(Array(η̅.data.parent) .== 0.0) - @test all(Array(U̅.data.parent .== 0.0)) - @test all(Array(V̅.data.parent .== 0.0)) + @test all(Array(U̅.data.parent) .== 0.0) + @test all(Array(V̅.data.parent) .== 0.0) end @testset "Inexact integration" begin @@ -61,7 +63,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(u, set_u_check) exact_U = similar(U) set!(exact_U, set_U_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) tolerance = 1e-3 @test all((Array(interior(U) .- interior(exact_U))) .< tolerance) @@ -70,7 +72,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(v, set_v_check) exact_V = similar(V) set!(exact_V, set_V_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all((Array(interior(V) .- interior(exact_V))) .< tolerance) end @@ -78,14 +80,14 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces Δz = zeros(Nz) Δz .= grid.z.Δᵃᵃᶜ - u .= 0.0 - U .= 1.0 - compute_barotropic_mode!(U, V, grid, u, v) - @test all(Array(U.data.parent) .== 0.0) + set!(u, 0) + set!(U, 1) + compute_barotropic_mode!(U, V, grid, u, v, η̅) + @test all(Array(interior(U)) .== 0.0) - u .= 1.0 - U .= 1.0 - compute_barotropic_mode!(U, V, grid, u, v) + set!(u, 1) + set!(U, 1) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all(Array(interior(U)) .≈ Lz) set_u_check(x, y, z) = sin(x) @@ -93,7 +95,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(u, set_u_check) exact_U = similar(U) set!(exact_U, set_U_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all(Array(interior(U)) .≈ Array(interior(exact_U))) set_v_check(x, y, z) = sin(x) * z * cos(y) @@ -101,7 +103,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces set!(v, set_v_check) exact_V = similar(V) set!(exact_V, set_V_check) - compute_barotropic_mode!(U, V, grid, u, v) + compute_barotropic_mode!(U, V, grid, u, v, η̅) @test all(Array(interior(V)) .≈ Array(interior(exact_V))) end diff --git a/test/test_zstar_coordinate.jl b/test/test_zstar_coordinate.jl new file mode 100644 index 0000000000..d60f75fb88 --- /dev/null +++ b/test/test_zstar_coordinate.jl @@ -0,0 +1,201 @@ +include("dependencies_for_runtests.jl") + +using Random +using Oceananigans: initialize! +using Oceananigans.ImmersedBoundaries: PartialCellBottom +using Oceananigans.Grids: MutableVerticalDiscretization +using Oceananigans.Models: ZStar + +function test_zstar_coordinate(model, Ni, Δt) + + bᵢ = deepcopy(model.tracers.b) + cᵢ = deepcopy(model.tracers.c) + + ∫bᵢ = Field(Integral(bᵢ)) + ∫cᵢ = Field(Integral(cᵢ)) + w = model.velocities.w + Nz = model.grid.Nz + + for _ in 1:Ni + time_step!(model, Δt) + end + + ∫b = Field(Integral(model.tracers.b)) + ∫c = Field(Integral(model.tracers.c)) + + # Testing that: + # (1) tracers are conserved down to machine precision + # (2) vertical velocities are zero at the top surface + @test interior(∫b, 1, 1, 1) ≈ interior(∫bᵢ, 1, 1, 1) + @test interior(∫c, 1, 1, 1) ≈ interior(∫cᵢ, 1, 1, 1) + @test maximum(abs, interior(w, :, :, Nz+1)) < eps(eltype(w)) + + return nothing +end + +function info_message(grid, free_surface) + msg1 = "$(architecture(grid)) " + msg2 = string(getnamewrapper(grid)) + msg3 = grid isa ImmersedBoundaryGrid ? " on a " * string(getnamewrapper(grid.underlying_grid)) : "" + msg4 = grid.z.Δᵃᵃᶠ isa Number ? " with uniform spacing" : " with stretched spacing" + msg5 = grid isa ImmersedBoundaryGrid ? " and $(string(getnamewrapper(grid.immersed_boundary))) immersed boundary" : "" + msg6 = " using a " * string(getnamewrapper(free_surface)) + return msg1 * msg2 * msg3 * msg4 * msg5 * msg6 +end + +const C = Center +const F = Face + +@testset "ZStar coordinate scaling tests" begin + @info "testing the ZStar coordinate scalings" + + z = MutableVerticalDiscretization((-20, 0)) + + grid = RectilinearGrid(size = (2, 2, 20), + x = (0, 2), + y = (0, 1), + z = z, + topology = (Bounded, Periodic, Bounded)) + + grid = ImmersedBoundaryGrid(grid, GridFittedBottom((x, y) -> -10)) + + model = HydrostaticFreeSurfaceModel(; grid, + free_surface = SplitExplicitFreeSurface(grid; substeps = 20), + vertical_coordinate = ZStar()) + + @test znode(1, 1, 21, grid, C(), C(), F()) == 0 + @test column_depthᶜᶜᵃ(1, 1, grid) == 10 + @test static_column_depthᶜᶜᵃ(1, 1, grid) == 10 + + set!(model, η = [1 1; 2 2]) + set!(model, u = (x, y, z) -> x, v = (x, y, z) -> y) + update_state!(model) + + @test σⁿ(1, 1, 1, grid, C(), C(), C()) == 11 / 10 + @test σⁿ(2, 1, 1, grid, C(), C(), C()) == 12 / 10 + + @test znode(1, 1, 21, grid, C(), C(), F()) == 1 + @test znode(2, 1, 21, grid, C(), C(), F()) == 2 + @test rnode(1, 1, 21, grid, C(), C(), F()) == 0 + @test column_depthᶜᶜᵃ(1, 1, grid) == 11 + @test column_depthᶜᶜᵃ(2, 1, grid) == 12 + @test static_column_depthᶜᶜᵃ(1, 1, grid) == 10 + @test static_column_depthᶜᶜᵃ(2, 1, grid) == 10 +end + +@testset "MutableVerticalDiscretization tests" begin + @info "testing the MutableVerticalDiscretization in ZCoordinate mode" + + z = MutableVerticalDiscretization((-20, 0)) + + # A mutable immersed grid + mutable_grid = RectilinearGrid(size=(2, 2, 20), x=(0, 2), y=(0, 1), z=z) + mutable_grid = ImmersedBoundaryGrid(mutable_grid, GridFittedBottom((x, y) -> -10)) + + # A static immersed grid + static_grid = RectilinearGrid(size=(2, 2, 20), x=(0, 2), y=(0, 1), z=(-20, 0)) + static_grid = ImmersedBoundaryGrid(static_grid, GridFittedBottom((x, y) -> -10)) + + # Make sure a model with a MutableVerticalDiscretization but ZCoordinate still runs and + # the results are the same as a model with a static vertical discretization. + mutable_model = HydrostaticFreeSurfaceModel(; grid=mutable_grid, free_surface=ImplicitFreeSurface()) + static_model = HydrostaticFreeSurfaceModel(; grid=static_grid, free_surface=ImplicitFreeSurface()) + + uᵢ = rand(size(mutable_model.velocities.u)...) + vᵢ = rand(size(mutable_model.velocities.v)...) + + set!(mutable_model; u=uᵢ, v=vᵢ) + set!(static_model; u=uᵢ, v=vᵢ) + + static_sim = Simulation(static_model; Δt=1e-3, stop_iteration=100) + mutable_sim = Simulation(mutable_model; Δt=1e-3, stop_iteration=100) + + run!(mutable_sim) + run!(static_sim) + + # Check that fields are the same + um, vm, wm = mutable_model.velocities + us, vs, ws = static_model.velocities + + @test all(um.data .≈ us.data) + @test all(vm.data .≈ vs.data) + @test all(wm.data .≈ ws.data) + @test all(um.data .≈ us.data) +end + +@testset "ZStar coordinate simulation testset" begin + z_uniform = MutableVerticalDiscretization((-20, 0)) + z_stretched = MutableVerticalDiscretization(collect(-20:0)) + topologies = ((Periodic, Periodic, Bounded), + (Periodic, Bounded, Bounded), + (Bounded, Periodic, Bounded), + (Bounded, Bounded, Bounded)) + + for arch in archs + for topology in topologies + Random.seed!(1234) + + rtg = RectilinearGrid(arch; size = (10, 10, 20), x = (0, 100kilometers), y = (-10kilometers, 10kilometers), topology, z = z_uniform) + rtgv = RectilinearGrid(arch; size = (10, 10, 20), x = (0, 100kilometers), y = (-10kilometers, 10kilometers), topology, z = z_stretched) + + irtg = ImmersedBoundaryGrid(rtg, GridFittedBottom((x, y) -> rand() - 10)) + irtgv = ImmersedBoundaryGrid(rtgv, GridFittedBottom((x, y) -> rand() - 10)) + prtg = ImmersedBoundaryGrid(rtg, PartialCellBottom((x, y) -> rand() - 10)) + prtgv = ImmersedBoundaryGrid(rtgv, PartialCellBottom((x, y) -> rand() - 10)) + + if topology[2] == Bounded + llg = LatitudeLongitudeGrid(arch; size = (10, 10, 20), latitude = (0, 1), longitude = (0, 1), topology, z = z_uniform) + llgv = LatitudeLongitudeGrid(arch; size = (10, 10, 20), latitude = (0, 1), longitude = (0, 1), topology, z = z_stretched) + + illg = ImmersedBoundaryGrid(llg, GridFittedBottom((x, y) -> rand() - 10)) + illgv = ImmersedBoundaryGrid(llgv, GridFittedBottom((x, y) -> rand() - 10)) + pllg = ImmersedBoundaryGrid(llg, PartialCellBottom((x, y) -> rand() - 10)) + pllgv = ImmersedBoundaryGrid(llgv, PartialCellBottom((x, y) -> rand() - 10)) + + # TODO: Partial cell bottom are broken at the moment and do not account for the Δz in the volumes + # and vertical areas (see https://github.com/CliMA/Oceananigans.jl/issues/3958) + # When this is issue is fixed we can add the partial cells to the testing. + grids = [llg, rtg, llgv, rtgv, illg, irtg, illgv, irtgv] # , pllg, prtg, pllgv, prtgv] + else + grids = [rtg, rtgv, irtg, irtgv] #, prtg, prtgv] + end + + for grid in grids + split_free_surface = SplitExplicitFreeSurface(grid; cfl = 0.75) + implicit_free_surface = ImplicitFreeSurface() + explicit_free_surface = ExplicitFreeSurface() + + for free_surface in [split_free_surface, implicit_free_surface, explicit_free_surface] + + # TODO: There are parameter space issues with ImplicitFreeSurface and a immersed LatitudeLongitudeGrid + # For the moment we are skipping these tests. + if (arch isa GPU) && + (free_surface isa ImplicitFreeSurface) && + (grid isa ImmersedBoundaryGrid) && + (grid.underlying_grid isa LatitudeLongitudeGrid) + + @info " Skipping $(info_message(grid, free_surface)) because of parameter space issues" + continue + end + + info_msg = info_message(grid, free_surface) + @testset "$info_msg" begin + @info " Testing a $info_msg" + model = HydrostaticFreeSurfaceModel(; grid, + free_surface, + tracers = (:b, :c), + buoyancy = BuoyancyTracer(), + vertical_coordinate = ZStar()) + + bᵢ(x, y, z) = x < grid.Lx / 2 ? 0.06 : 0.01 + + set!(model, c = (x, y, z) -> rand(), b = bᵢ) + + Δt = free_surface isa ExplicitFreeSurface ? 10 : 2minutes + test_zstar_coordinate(model, 100, Δt) + end + end + end + end + end +end \ No newline at end of file diff --git a/validation/z_star_coordinate/lock_release.jl b/validation/z_star_coordinate/lock_release.jl new file mode 100644 index 0000000000..d4c625ff1f --- /dev/null +++ b/validation/z_star_coordinate/lock_release.jl @@ -0,0 +1,80 @@ +using Oceananigans +using Oceananigans.Grids +using Oceananigans.Units +using Oceananigans.Utils: prettytime +using Oceananigans.Advection: WENOVectorInvariant +using Oceananigans.AbstractOperations: GridMetricOperation +using Printf + +z_faces = MutableVerticalDiscretization(-20, 0) + +grid = RectilinearGrid(size = (128, 20), + x = (0, 64kilometers), + z = z_faces, + halo = (6, 6), + topology = (Bounded, Flat, Bounded)) + +grid = ImmersedBoundaryGrid(grid, GridFittedBottom(x -> - (64kilometers - x) / 64kilometers * 20)) + +model = HydrostaticFreeSurfaceModel(; grid, + momentum_advection = WENO(order = 5), + tracer_advection = WENO(order = 5), + buoyancy = BuoyancyTracer(), + closure = nothing, + tracers = :b, + free_surface = SplitExplicitFreeSurface(grid; substeps = 10)) + +g = model.free_surface.gravitational_acceleration + +bᵢ(x, z) = x < 32kilometers ? 0.06 : 0.01 + +set!(model, b = bᵢ) + +Δt = 10 + +@info "the time step is $Δt" + +simulation = Simulation(model; Δt, stop_iteration = 10000, stop_time = 17hours) + +Δz = zspacings(grid, Center(), Center(), Center()) +∫b_init = sum(model.tracers.b * Δz) / sum(Δz) + +field_outputs = merge(model.velocities, model.tracers, (; Δz)) + +simulation.output_writers[:other_variables] = JLD2OutputWriter(model, field_outputs, + overwrite_existing = true, + schedule = IterationInterval(100), + filename = "zstar_model") + +function progress(sim) + w = interior(sim.model.velocities.w, :, :, sim.model.grid.Nz+1) + u = sim.model.velocities.u + ∫b = sum(model.tracers.b * Δz) / sum(Δz) + + msg0 = @sprintf("Time: %s iteration %d ", prettytime(sim.model.clock.time), sim.model.clock.iteration) + msg1 = @sprintf("extrema w: %.2e %.2e ", maximum(w), minimum(w)) + msg2 = @sprintf("extrema u: %.2e %.2e ", maximum(u), minimum(u)) + msg3 = @sprintf("drift b: %6.3e ", ∫b - ∫b_init) + msg4 = @sprintf("extrema Δz: %.2e %.2e ", maximum(Δz), minimum(Δz)) + @info msg0 * msg1 * msg2 * msg3 * msg4 + + return nothing +end + +simulation.callbacks[:progress] = Callback(progress, IterationInterval(100)) + +run!(simulation) + +using Oceananigans.Fields: OneField + +# # Check tracer conservation +b = FieldTimeSeries("zstar_model.jld2", "b") +dz = FieldTimeSeries("zstar_model.jld2", "Δz") + +init = sum(dz[1] * b[1]) / sum(dz[1]) +drift = [] + +for t in 1:length(b.times) + push!(drift, sum(dz[t] * b[t]) / sum(dz[t]) - init) +end + From 234bc988bebb5ba92f59252da1489e86285b429d Mon Sep 17 00:00:00 2001 From: Tomas Chor Date: Sat, 18 Jan 2025 19:18:25 -0500 Subject: [PATCH 21/31] Fix bug in validation of `MutableVerticalDiscretization()` (#4050) * fix validation * pass vertical_coordinate to model --- validation/z_star_coordinate/lock_release.jl | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/validation/z_star_coordinate/lock_release.jl b/validation/z_star_coordinate/lock_release.jl index d4c625ff1f..072388bff4 100644 --- a/validation/z_star_coordinate/lock_release.jl +++ b/validation/z_star_coordinate/lock_release.jl @@ -6,7 +6,7 @@ using Oceananigans.Advection: WENOVectorInvariant using Oceananigans.AbstractOperations: GridMetricOperation using Printf -z_faces = MutableVerticalDiscretization(-20, 0) +z_faces = MutableVerticalDiscretization((-20, 0)) grid = RectilinearGrid(size = (128, 20), x = (0, 64kilometers), @@ -22,6 +22,7 @@ model = HydrostaticFreeSurfaceModel(; grid, buoyancy = BuoyancyTracer(), closure = nothing, tracers = :b, + vertical_coordinate = Oceananigans.Models.ZStar(), free_surface = SplitExplicitFreeSurface(grid; substeps = 10)) g = model.free_surface.gravitational_acceleration @@ -78,3 +79,5 @@ for t in 1:length(b.times) push!(drift, sum(dz[t] * b[t]) / sum(dz[t]) - init) end +using CairoMakie +lines(drift) From 8b4c1ad4f683960a4a6f74813e97b6c9f23ea625 Mon Sep 17 00:00:00 2001 From: William Moses Date: Mon, 20 Jan 2025 12:38:10 -0600 Subject: [PATCH 22/31] More generalization of tuples (#4049) Co-authored-by: Gregory L. Wagner Co-authored-by: Navid C. Constantinou --- .../fill_halo_regions_periodic.jl | 17 ++++++++++------- ...ic_free_surface_tendency_kernel_functions.jl | 2 +- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/BoundaryConditions/fill_halo_regions_periodic.jl b/src/BoundaryConditions/fill_halo_regions_periodic.jl index edd97528fa..549926f5e3 100644 --- a/src/BoundaryConditions/fill_halo_regions_periodic.jl +++ b/src/BoundaryConditions/fill_halo_regions_periodic.jl @@ -7,7 +7,7 @@ using KernelAbstractions.Extras.LoopInfo: @unroll @inline parent_size_and_offset(c, dim1, dim2, size, offset) = (parent(c), size, fix_halo_offsets.(offset, c.offsets[[dim1, dim2]])) @inline parent_size_and_offset(c, dim1, dim2, ::Symbol, offset) = (parent(c), size(parent(c))[[dim1, dim2]], (0, 0)) -@inline function parent_size_and_offset(c::NTuple, dim1, dim2, ::Symbol, offset) +@inline function parent_size_and_offset(c::Tuple, dim1, dim2, ::Symbol, offset) p = parent.(c) p_size = (minimum([size(t, dim1) for t in p]), minimum([size(t, dim2) for t in p])) return p, p_size, (0, 0) @@ -71,9 +71,10 @@ end #### Tupled periodic boundary condition #### -@kernel function fill_periodic_west_and_east_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H} +@kernel function fill_periodic_west_and_east_halo!(c::Tuple, ::Val{H}, N) where {H} j, k = @index(Global, NTuple) - @unroll for n = 1:M + ntuple(Val(length(c))) do n + Base.@_inline_meta @unroll for i = 1:H @inbounds begin c[n][i, j, k] = c[n][N+i, j, k] # west @@ -83,9 +84,10 @@ end end end -@kernel function fill_periodic_south_and_north_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H} +@kernel function fill_periodic_south_and_north_halo!(c::Tuple, ::Val{H}, N) where {H} i, k = @index(Global, NTuple) - @unroll for n = 1:M + ntuple(Val(length(c))) do n + Base.@_inline_meta @unroll for j = 1:H @inbounds begin c[n][i, j, k] = c[n][i, N+j, k] # south @@ -95,9 +97,10 @@ end end end -@kernel function fill_periodic_bottom_and_top_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H} +@kernel function fill_periodic_bottom_and_top_halo!(c::Tuple, ::Val{H}, N) where {H} i, j = @index(Global, NTuple) - @unroll for n = 1:M + ntuple(Val(length(c))) do n + Base.@_inline_meta @unroll for k = 1:H @inbounds begin c[n][i, j, k] = c[n][i, j, N+k] # top diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl index 9447b3effd..0e0fe6b20a 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -137,4 +137,4 @@ where `c = C[tracer_index]`. - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields) + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields) + forcing(i, j, k, grid, clock, model_fields)) -end \ No newline at end of file +end From 373647faf63a09caddf92b1b3bbb260f86c4ef86 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 23 Jan 2025 07:51:52 +0100 Subject: [PATCH 23/31] Remove all hardcoded Float64 from the `Advection` module (#4053) * remove all float64 * small fix * disambiguate * add a sum to one * compute coefficients in high precision * Fix doctests * implement suggestions --- src/Advection/Advection.jl | 2 +- src/Advection/centered_reconstruction.jl | 14 +- src/Advection/reconstruction_coefficients.jl | 96 ++++----- src/Advection/upwind_biased_reconstruction.jl | 42 ++-- src/Advection/vector_invariant_advection.jl | 46 ++-- src/Advection/weno_interpolants.jl | 201 +++++++++--------- .../compute_slow_tendencies.jl | 10 +- .../split_explicit_free_surface.jl | 2 +- .../hydrostatic_free_surface_field_tuples.jl | 8 +- ..._free_surface_tendency_kernel_functions.jl | 2 +- src/Oceananigans.jl | 4 + 11 files changed, 215 insertions(+), 212 deletions(-) diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index d528aa8610..1e92f54c9b 100644 --- a/src/Advection/Advection.jl +++ b/src/Advection/Advection.jl @@ -28,10 +28,10 @@ using DocStringExtensions using Base: @propagate_inbounds using Adapt using OffsetArrays - using Oceananigans.Grids using Oceananigans.Grids: with_halo using Oceananigans.Architectures: architecture, CPU +using Oceananigans: supported_float_types using Oceananigans.Operators using Oceananigans.Operators: flux_div_xyᶜᶜᶜ, Γᶠᶠᶜ, ∂t_σ diff --git a/src/Advection/centered_reconstruction.jl b/src/Advection/centered_reconstruction.jl index c0dbd4a22c..ed5d3d80c0 100644 --- a/src/Advection/centered_reconstruction.jl +++ b/src/Advection/centered_reconstruction.jl @@ -98,16 +98,16 @@ const ACAS = AbstractCenteredAdvectionScheme @inline biased_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme::ACAS, bias, c, args...) = symmetric_interpolate_zᵃᵃᶜ(i, j, k, grid, scheme, c, args...) # uniform centered reconstruction -for buffer in advection_buffers +for buffer in advection_buffers, FT in supported_float_types @eval begin - @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, <:Nothing}, ψ, idx, loc, args...) where FT = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :x, false)) - @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, <:Nothing}, ψ::Function, idx, loc, args...) where FT = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :x, true)) + @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, <:Nothing}, ψ, idx, loc, args...) = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x, false)) + @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, <:Nothing}, ψ::Function, idx, loc, args...) = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :x, true)) - @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, <:Nothing}, ψ, idx, loc, args...) where {FT, XT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :y, false)) - @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, <:Nothing}, ψ::Function, idx, loc, args...) where {FT, XT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :y, true)) + @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, <:Nothing}, ψ, idx, loc, args...) where XT = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y, false)) + @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, <:Nothing}, ψ::Function, idx, loc, args...) where XT = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :y, true)) - @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, YT, <:Nothing}, ψ, idx, loc, args...) where {FT, XT, YT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :z, false)) - @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, FT, XT, YT, <:Nothing}, ψ::Function, idx, loc, args...) where {FT, XT, YT} = @inbounds $(calc_reconstruction_stencil(buffer, :symmetric, :z, true)) + @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, YT, <:Nothing}, ψ, idx, loc, args...) where {XT, YT} = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z, false)) + @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, scheme::Centered{$buffer, $FT, XT, YT, <:Nothing}, ψ::Function, idx, loc, args...) where {XT, YT} = @inbounds $(calc_reconstruction_stencil(FT, buffer, :symmetric, :z, true)) end end diff --git a/src/Advection/reconstruction_coefficients.jl b/src/Advection/reconstruction_coefficients.jl index 13aec9f8b4..a78b63213a 100644 --- a/src/Advection/reconstruction_coefficients.jl +++ b/src/Advection/reconstruction_coefficients.jl @@ -97,8 +97,8 @@ Positional Arguments On a uniform `grid`, the coefficients are independent of the `xr` and `xi` values. """ -@inline function stencil_coefficients(i, r, xr, xi; shift = 0, op = Base.:(-), order = 3, der = nothing) - coeffs = zeros(order) +@inline function stencil_coefficients(FT, i, r, xr, xi; shift = 0, op = Base.:(-), order = 3, der = nothing) + coeffs = zeros(BigFloat, order) @inbounds begin for j in 0:order-1 for m in j+1:order @@ -109,51 +109,38 @@ On a uniform `grid`, the coefficients are independent of the `xr` and `xi` value end end - return tuple(coeffs...) + coeffs = FT.(coeffs)[1:end-1] + + return tuple(coeffs..., 1-sum(coeffs)) # Coefficients should sum to 1! end """ - Coefficients for uniform centered and upwind schemes + uniform_reconstruction_coefficients(FT, Val(bias), buffer) + +Returns coefficients for finite volume reconstruction used in linear advection schemes (`Centered` and `UpwindBiased`). +`FT` is the floating type (e.g. `Float32`, `Float64`), `bias` is either `:symmetric`, `:left`, or `:right`, +and `buffer` is the buffer size which determines the order of the reconstruction. -symmetric coefficients are for centered reconstruction (dispersive, even order), -left and right are for upwind biased (diffusive, odd order) examples: -julia> using Oceananigans.Advection: coeff2_symmetric, coeff3_left, coeff3_right, coeff4_symmetric, coeff5_left +```julia +julia> using Oceananigans.Advection: uniform_reconstruction_coefficients -julia> coeff2_symmetric +julia> uniform_reconstruction_coefficients(Float64, Val(:symmetric), 1) (0.5, 0.5) -julia> coeff3_left, coeff3_right -((0.33333333333333337, 0.8333333333333334, -0.16666666666666666), (-0.16666666666666669, 0.8333333333333333, 0.3333333333333333)) - -julia> coeff4_symmetric -(-0.08333333333333333, 0.5833333333333333, 0.5833333333333333, -0.08333333333333333) +julia> uniform_reconstruction_coefficients(Float32, Val(:left), 3) +(-0.05f0, 0.45f0, 0.78333336f0, -0.21666667f0, 0.033333335f0) -julia> coeff5_left -(-0.049999999999999926, 0.45000000000000007, 0.7833333333333333, -0.21666666666666667, 0.03333333333333333) +julia> uniform_reconstruction_coefficients(Float16, Val(:right), 4) +(Float16(-0.00714), Float16(0.0595), Float16(-0.2405), Float16(0.76), Float16(0.51), Float16(-0.09045), Float16(0.00952)) +``` """ -const coeff1_left = 1.0 -const coeff1_right = 1.0 - -# buffer in [1:6] allows up to Centered(order = 12) and UpwindBiased(order = 11) -for buffer in advection_buffers - order_bias = 2buffer - 1 - order_symm = 2buffer - - coeff_symm = Symbol(:coeff, order_symm, :_symmetric) - coeff_left = Symbol(:coeff, order_bias, :_left) - coeff_right = Symbol(:coeff, order_bias, :_right) - @eval begin - const $coeff_symm = stencil_coefficients(50, $(buffer - 1), collect(1:100), collect(1:100); order = $order_symm) - if $order_bias > 1 - const $coeff_left = stencil_coefficients(50, $(buffer - 2), collect(1:100), collect(1:100); order = $order_bias) - const $coeff_right = stencil_coefficients(50, $(buffer - 1), collect(1:100), collect(1:100); order = $order_bias) - end - end -end +uniform_reconstruction_coefficients(FT, ::Val{:symmetric}, buffer) = stencil_coefficients(FT, 50, buffer-1, collect(1:100), collect(1:100); order = 2buffer) +uniform_reconstruction_coefficients(FT, ::Val{:left}, buffer) = buffer==1 ? (one(FT),) : stencil_coefficients(FT, 50, buffer-2, collect(1:100), collect(1:100); order = 2buffer-1) +uniform_reconstruction_coefficients(FT, ::Val{:right}, buffer) = buffer==1 ? (one(FT),) : stencil_coefficients(FT, 50, buffer-1, collect(1:100), collect(1:100); order = 2buffer-1) """ - calc_reconstruction_stencil(buffer, shift, dir, func::Bool = false) + calc_reconstruction_stencil(FT, buffer, shift, dir, func::Bool = false) Stencils for reconstruction calculations (note that WENO has its own reconstruction stencils) @@ -167,23 +154,23 @@ Examples ```jldoctest julia> using Oceananigans.Advection: calc_reconstruction_stencil -julia> calc_reconstruction_stencil(1, :right, :x) -:(+(convert(FT, coeff1_right[1]) * ψ[i + 0, j, k])) +julia> calc_reconstruction_stencil(Float32, 1, :right, :x) +:(+(1.0f0 * ψ[i + 0, j, k])) -julia> calc_reconstruction_stencil(1, :left, :x) -:(+(convert(FT, coeff1_left[1]) * ψ[i + -1, j, k])) +julia> calc_reconstruction_stencil(Float64, 1, :left, :x) +:(+(1.0 * ψ[i + -1, j, k])) -julia> calc_reconstruction_stencil(1, :symmetric, :x) -:(convert(FT, coeff2_symmetric[2]) * ψ[i + -1, j, k] + convert(FT, coeff2_symmetric[1]) * ψ[i + 0, j, k]) +julia> calc_reconstruction_stencil(Float64, 1, :symmetric, :y) +:(0.5 * ψ[i, j + -1, k] + 0.5 * ψ[i, j + 0, k]) -julia> calc_reconstruction_stencil(2, :symmetric, :x) -:(convert(FT, coeff4_symmetric[4]) * ψ[i + -2, j, k] + convert(FT, coeff4_symmetric[3]) * ψ[i + -1, j, k] + convert(FT, coeff4_symmetric[2]) * ψ[i + 0, j, k] + convert(FT, coeff4_symmetric[1]) * ψ[i + 1, j, k]) +julia> calc_reconstruction_stencil(Float32, 2, :symmetric, :x) +:(-0.083333254f0 * ψ[i + -2, j, k] + 0.5833333f0 * ψ[i + -1, j, k] + 0.5833333f0 * ψ[i + 0, j, k] + -0.083333336f0 * ψ[i + 1, j, k]) -julia> calc_reconstruction_stencil(3, :left, :x) -:(convert(FT, coeff5_left[5]) * ψ[i + -3, j, k] + convert(FT, coeff5_left[4]) * ψ[i + -2, j, k] + convert(FT, coeff5_left[3]) * ψ[i + -1, j, k] + convert(FT, coeff5_left[2]) * ψ[i + 0, j, k] + convert(FT, coeff5_left[1]) * ψ[i + 1, j, k]) +julia> calc_reconstruction_stencil(Float32, 3, :left, :x) +:(0.0333333f0 * ψ[i + -3, j, k] + -0.21666667f0 * ψ[i + -2, j, k] + 0.78333336f0 * ψ[i + -1, j, k] + 0.45f0 * ψ[i + 0, j, k] + -0.05f0 * ψ[i + 1, j, k]) ``` """ -@inline function calc_reconstruction_stencil(buffer, shift, dir, func::Bool = false) +@inline function calc_reconstruction_stencil(FT, buffer, shift, dir, func::Bool = false) N = buffer * 2 order = shift == :symmetric ? N : N - 1 if shift != :symmetric @@ -194,21 +181,22 @@ julia> calc_reconstruction_stencil(3, :left, :x) rng = rng .+ 1 end stencil_full = Vector(undef, N) - coeff = Symbol(:coeff, order, :_, shift) + coeff = uniform_reconstruction_coefficients(FT, Val(shift), buffer) for (idx, n) in enumerate(rng) c = n - buffer - 1 + C = coeff[order - idx + 1] if func stencil_full[idx] = dir == :x ? - :(convert(FT, $coeff[$(order - idx + 1)]) * ψ(i + $c, j, k, grid, args...)) : + :($C * ψ(i + $c, j, k, grid, args...)) : dir == :y ? - :(convert(FT, $coeff[$(order - idx + 1)]) * ψ(i, j + $c, k, grid, args...)) : - :(convert(FT, $coeff[$(order - idx + 1)]) * ψ(i, j, k + $c, grid, args...)) + :($C * ψ(i, j + $c, k, grid, args...)) : + :($C * ψ(i, j, k + $c, grid, args...)) else stencil_full[idx] = dir == :x ? - :(convert(FT, $coeff[$(order - idx + 1)]) * ψ[i + $c, j, k]) : + :($C * ψ[i + $c, j, k]) : dir == :y ? - :(convert(FT, $coeff[$(order - idx + 1)]) * ψ[i, j + $c, k]) : - :(convert(FT, $coeff[$(order - idx + 1)]) * ψ[i, j, k + $c]) + :($C * ψ[i, j + $c, k]) : + :($C * ψ[i, j, k + $c]) end end return Expr(:call, :+, stencil_full...) @@ -329,7 +317,7 @@ end stencil = NTuple{order, FT}[] @inbounds begin for i = 0:N+1 - push!(stencil, stencil_coefficients(i, r, cpu_coord, cpu_coord; order)) + push!(stencil, stencil_coefficients(FT, i, r, cpu_coord, cpu_coord; order)) end end return OffsetArray(on_architecture(arch, stencil), -1) diff --git a/src/Advection/upwind_biased_reconstruction.jl b/src/Advection/upwind_biased_reconstruction.jl index 303d32aab7..e8412bfdab 100644 --- a/src/Advection/upwind_biased_reconstruction.jl +++ b/src/Advection/upwind_biased_reconstruction.jl @@ -111,31 +111,31 @@ const UY{N, FT} = UpwindBiased{N, FT, <:Any, <:Nothing} where {N, FT} const UZ{N, FT} = UpwindBiased{N, FT, <:Any, <:Any, <:Nothing} where {N, FT} # Uniform upwind biased reconstruction -for buffer in advection_buffers +for buffer in advection_buffers, FT in supported_float_types @eval begin - @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, FT}, bias, ψ, idx, loc, args...) where FT = - @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :x, false)), - $(calc_reconstruction_stencil(buffer, :right, :x, false))) + @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, $FT}, bias, ψ, idx, loc, args...) = + @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :x, false)), + $(calc_reconstruction_stencil(FT, buffer, :right, :x, false))) - @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT = - @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :x, true)), - $(calc_reconstruction_stencil(buffer, :right, :x, true))) + @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, ::UX{$buffer, $FT}, bias, ψ::Function, idx, loc, args...) = + @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :x, true)), + $(calc_reconstruction_stencil(FT, buffer, :right, :x, true))) - @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, FT}, bias, ψ, idx, loc, args...) where FT = - @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :y, false)), - $(calc_reconstruction_stencil(buffer, :right, :y, false))) + @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, $FT}, bias, ψ, idx, loc, args...) = + @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :y, false)), + $(calc_reconstruction_stencil(FT, buffer, :right, :y, false))) - @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT = - @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :y, true)), - $(calc_reconstruction_stencil(buffer, :right, :y, true))) + @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, ::UY{$buffer, $FT}, bias, ψ::Function, idx, loc, args...) = + @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :y, true)), + $(calc_reconstruction_stencil(FT, buffer, :right, :y, true))) - @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, FT}, bias, ψ, idx, loc, args...) where FT = - @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :z, false)), - $(calc_reconstruction_stencil(buffer, :right, :z, false))) + @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, $FT}, bias, ψ, idx, loc, args...) = + @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :z, false)), + $(calc_reconstruction_stencil(FT, buffer, :right, :z, false))) - @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT = - @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(buffer, :left, :z, true)), - $(calc_reconstruction_stencil(buffer, :right, :z, true))) + @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, ::UZ{$buffer, $FT}, bias, ψ::Function, idx, loc, args...) = + @inbounds ifelse(bias isa LeftBias, $(calc_reconstruction_stencil(FT, buffer, :left, :z, true)), + $(calc_reconstruction_stencil(FT, buffer, :right, :z, true))) end end @@ -147,11 +147,11 @@ for (dir, ξ, val) in zip((:xᶠᵃᵃ, :yᵃᶠᵃ, :zᵃᵃᶠ), (:x, :y, :z), @eval begin @inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer, FT}, bias, ψ, idx, loc, args...) where FT = @inbounds ifelse(bias isa LeftBias, sum($(reconstruction_stencil(buffer, :left, ξ, false)) .* retrieve_coeff(scheme, Val(1), Val($val), idx, loc)), - sum($(reconstruction_stencil(buffer, :right, ξ, false)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc))) + sum($(reconstruction_stencil(buffer, :right, ξ, false)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc))) @inline $stencil(i, j, k, grid, scheme::UpwindBiased{$buffer, FT}, bias, ψ::Function, idx, loc, args...) where FT = @inbounds ifelse(bias isa LeftBias, sum($(reconstruction_stencil(buffer, :left, ξ, true)) .* retrieve_coeff(scheme, Val(1), Val($val), idx, loc)), - sum($(reconstruction_stencil(buffer, :right, ξ, true)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc))) + sum($(reconstruction_stencil(buffer, :right, ξ, true)) .* retrieve_coeff(scheme, Val(2), Val($val), idx, loc))) end end end diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl index 7f02532200..f599b110c1 100644 --- a/src/Advection/vector_invariant_advection.jl +++ b/src/Advection/vector_invariant_advection.jl @@ -426,33 +426,33 @@ end @inline U_dot_∇u(i, j, k, grid::AbstractGrid{FT}, scheme::Nothing, U) where FT = zero(FT) @inline U_dot_∇v(i, j, k, grid::AbstractGrid{FT}, scheme::Nothing, U) where FT = zero(FT) -const UB{N} = UpwindBiased{N} -const UBX{N} = UpwindBiased{N, <:Any, <:Nothing} -const UBY{N} = UpwindBiased{N, <:Any, <:Any, <:Nothing} -const UBZ{N} = UpwindBiased{N, <:Any, <:Any, <:Any, <:Nothing} +const UB{N, FT} = UpwindBiased{N, FT} +const UBX{N, FT} = UpwindBiased{N, FT, <:Nothing} +const UBY{N, FT} = UpwindBiased{N, FT, <:Any, <:Nothing} +const UBZ{N, FT} = UpwindBiased{N, FT, <:Any, <:Any, <:Nothing} -const C{N} = Centered{N, <:Any} -const CX{N} = Centered{N, <:Any, <:Nothing} -const CY{N} = Centered{N, <:Any, <:Any, <:Nothing} -const CZ{N} = Centered{N, <:Any, <:Any, <:Any, <:Nothing} +const C{N, FT} = Centered{N, FT} +const CX{N, FT} = Centered{N, FT, <:Nothing} +const CY{N, FT} = Centered{N, FT, <:Any, <:Nothing} +const CZ{N, FT} = Centered{N, FT, <:Any, <:Any, <:Nothing} const AS = AbstractSmoothnessStencil # To adapt passing smoothness stencils to upwind biased schemes and centered schemes (not WENO) -for b in 1:6 +for b in advection_buffers, FT in supported_float_types @eval begin - @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s::C{$b}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s, f, idx, loc, args...) - @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s::C{$b}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s, f, idx, loc, args...) - @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s::C{$b}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s, f, idx, loc, args...) - @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s::CX{$b}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s, f, idx, loc, args...) - @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s::CY{$b}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s, f, idx, loc, args...) - @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s::CZ{$b}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s, f, idx, loc, args...) - - @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s::UB{$b}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) - @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s::UB{$b}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) - @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s::UB{$b}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s, bias, f, idx, loc, args...) - @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s::UBX{$b}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) - @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s::UBY{$b}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) - @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s::UBZ{$b}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s, bias, f, idx, loc, args...) + @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s::C{$b, $FT}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s, f, idx, loc, args...) + @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s::C{$b, $FT}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s, f, idx, loc, args...) + @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s::C{$b, $FT}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s, f, idx, loc, args...) + @inline inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s::CX{$b, $FT}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_xᶠᵃᵃ(i, j, k, grid, s, f, idx, loc, args...) + @inline inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s::CY{$b, $FT}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_yᵃᶠᵃ(i, j, k, grid, s, f, idx, loc, args...) + @inline inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s::CZ{$b, $FT}, f::Function, idx, loc, ::AS, args...) = inner_symmetric_interpolate_zᵃᵃᶠ(i, j, k, grid, s, f, idx, loc, args...) + + @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s::UB{$b, $FT}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) + @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s::UB{$b, $FT}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) + @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s::UB{$b, $FT}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s, bias, f, idx, loc, args...) + @inline inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s::UBX{$b, $FT}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_xᶠᵃᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) + @inline inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s::UBY{$b, $FT}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_yᵃᶠᵃ(i, j, k, grid, s, bias, f, idx, loc, args...) + @inline inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s::UBZ{$b, $FT}, bias, f::Function, idx, loc, ::AS, args...) = inner_biased_interpolate_zᵃᵃᶠ(i, j, k, grid, s, bias, f, idx, loc, args...) end -end +end \ No newline at end of file diff --git a/src/Advection/weno_interpolants.jl b/src/Advection/weno_interpolants.jl index 513039b9f1..f49ef86c1e 100644 --- a/src/Advection/weno_interpolants.jl +++ b/src/Advection/weno_interpolants.jl @@ -67,53 +67,60 @@ end Base.show(io::IO, a::FunctionStencil) = print(io, "FunctionStencil f = $(a.func)") -const ƞ = Int32(2) # WENO exponent -const ε = 1e-8 +const ε = 1f-8 # Optimal values for finite volume reconstruction of order `WENO{order}` and stencil `Val{stencil}` from # Balsara & Shu, "Monotonicity Preserving Weighted Essentially Non-oscillatory Schemes with Inceasingly High Order of Accuracy" -@inline C★(::WENO{2}, ::Val{0}) = 2/3 -@inline C★(::WENO{2}, ::Val{1}) = 1/3 - -@inline C★(::WENO{3}, ::Val{0}) = 3/10 -@inline C★(::WENO{3}, ::Val{1}) = 3/5 -@inline C★(::WENO{3}, ::Val{2}) = 1/10 - -@inline C★(::WENO{4}, ::Val{0}) = 4/35 -@inline C★(::WENO{4}, ::Val{1}) = 18/35 -@inline C★(::WENO{4}, ::Val{2}) = 12/35 -@inline C★(::WENO{4}, ::Val{3}) = 1/35 - -@inline C★(::WENO{5}, ::Val{0}) = 5/126 -@inline C★(::WENO{5}, ::Val{1}) = 20/63 -@inline C★(::WENO{5}, ::Val{2}) = 10/21 -@inline C★(::WENO{5}, ::Val{3}) = 10/63 -@inline C★(::WENO{5}, ::Val{4}) = 1/126 - -@inline C★(::WENO{6}, ::Val{0}) = 1/77 -@inline C★(::WENO{6}, ::Val{1}) = 25/154 -@inline C★(::WENO{6}, ::Val{2}) = 100/231 -@inline C★(::WENO{6}, ::Val{3}) = 25/77 -@inline C★(::WENO{6}, ::Val{4}) = 5/77 -@inline C★(::WENO{6}, ::Val{5}) = 1/462 + +for FT in supported_float_types + @eval begin + @inline C★(::WENO{2, $FT}, ::Val{0}) = $(FT(2//3)) + @inline C★(::WENO{2, $FT}, ::Val{1}) = $(FT(1//3)) + + @inline C★(::WENO{3, $FT}, ::Val{0}) = $(FT(3//10)) + @inline C★(::WENO{3, $FT}, ::Val{1}) = $(FT(3//5)) + @inline C★(::WENO{3, $FT}, ::Val{2}) = $(FT(1//10)) + + @inline C★(::WENO{4, $FT}, ::Val{0}) = $(FT(4//35)) + @inline C★(::WENO{4, $FT}, ::Val{1}) = $(FT(18//35)) + @inline C★(::WENO{4, $FT}, ::Val{2}) = $(FT(12//35)) + @inline C★(::WENO{4, $FT}, ::Val{3}) = $(FT(1//35)) + + @inline C★(::WENO{5, $FT}, ::Val{0}) = $(FT(5//126)) + @inline C★(::WENO{5, $FT}, ::Val{1}) = $(FT(20//63)) + @inline C★(::WENO{5, $FT}, ::Val{2}) = $(FT(10//21)) + @inline C★(::WENO{5, $FT}, ::Val{3}) = $(FT(10//63)) + @inline C★(::WENO{5, $FT}, ::Val{4}) = $(FT(1//126)) + + @inline C★(::WENO{6, $FT}, ::Val{0}) = $(FT(1//77)) + @inline C★(::WENO{6, $FT}, ::Val{1}) = $(FT(25//154)) + @inline C★(::WENO{6, $FT}, ::Val{2}) = $(FT(100//231)) + @inline C★(::WENO{6, $FT}, ::Val{3}) = $(FT(25//77)) + @inline C★(::WENO{6, $FT}, ::Val{4}) = $(FT(5//77)) + @inline C★(::WENO{6, $FT}, ::Val{5}) = $(FT(1//462)) + end +end # ENO reconstruction procedure per stencil -for buffer in [2, 3, 4, 5, 6] +for buffer in advection_buffers[2:end] # WENO{<:Any, 1} does not exist for stencil in collect(0:1:buffer-1) - - # ENO coefficients for uniform direction (when T<:Nothing) and stretched directions (when T<:Any) + for FT in supported_float_types + # ENO coefficients for uniform direction (when T<:Nothing) and stretched directions (when T<:Any) + @eval begin + """ + coeff_p(::WENO{buffer, FT}, bias, ::Val{stencil}, T, args...) + + Reconstruction coefficients for the stencil number `stencil` of a WENO reconstruction + of order `buffer * 2 - 1`. Uniform coefficients (i.e. when `T == Nothing`) are independent on the + `bias` of the reconstruction (either `LeftBias` or `RightBias`), while stretched coeffiecients are + retrieved from the precomputed coefficients via the `retrieve_coeff` function + """ + @inline coeff_p(::WENO{$buffer, $FT}, bias, ::Val{$stencil}, ::Type{Nothing}, args...) = + @inbounds $(stencil_coefficients(FT, 50, stencil, collect(1:100), collect(1:100); order=buffer)) + end + end + @eval begin - """ - coeff_p(::WENO{buffer, FT}, bias, ::Val{stencil}, T, args...) - - Reconstruction coefficients for the stencil number `stencil` of a WENO reconstruction - of order `buffer * 2 - 1`. Uniform coefficients (i.e. when `T == Nothing`) are independent on the - `bias` of the reconstruction (either `LeftBias` or `RightBias`), while stretched coeffiecients are - retrieved from the precomputed coefficients via the `retrieve_coeff` function - """ - @inline coeff_p(::WENO{$buffer, FT}, bias, ::Val{$stencil}, ::Type{Nothing}, args...) where FT = - @inbounds map(FT, $(stencil_coefficients(50, stencil, collect(1:100), collect(1:100); order = buffer))) - # stretched coefficients are retrieved from precalculated coefficients @inline coeff_p(scheme::WENO{$buffer}, bias, ::Val{$stencil}, T, dir, i, loc) = ifelse(bias isa LeftBias, retrieve_coeff(scheme, $stencil, dir, i, loc), @@ -142,54 +149,58 @@ end # _UNIFORM_ smoothness coefficients (stretched smoothness coefficients are to be fixed!) -""" - smoothness_coefficients(::Val{buffer}, ::Val{stencil}) - -Return the coefficients used to calculate the smoothness indicators for the stencil -number `stencil` of a WENO reconstruction of order `buffer * 2 - 1`. The coefficients -are ordered in such a way to calculate the smoothness in the following fashion: - -```julia -buffer = 4 -stencil = 0 - -ψ = # The stencil corresponding to S₀ with buffer 4 (7th order WENO) - -C = smoothness_coefficients(Val(buffer), Val(0)) - -# The smoothness indicator -β = ψ[1] * (C[1] * ψ[1] + C[2] * ψ[2] + C[3] * ψ[3] + C[4] * ψ[4]) + - ψ[2] * (C[5] * ψ[2] + C[6] * ψ[3] + C[7] * ψ[4]) + - ψ[3] * (C[8] * ψ[3] + C[9] * ψ[4]) - ψ[4] * (C[10] * ψ[4]) -``` - -This last operation is metaprogrammed in the function `metaprogrammed_smoothness_operation` -""" -@inline smoothness_coefficients(::Val{2}, ::Val{0}) = :((1, -2, 1)) -@inline smoothness_coefficients(::Val{2}, ::Val{1}) = :((1, -2, 1)) - -@inline smoothness_coefficients(::Val{3}, ::Val{0}) = :((10, -31, 11, 25, -19, 4)) -@inline smoothness_coefficients(::Val{3}, ::Val{1}) = :((4, -13, 5, 13, -13, 4)) -@inline smoothness_coefficients(::Val{3}, ::Val{2}) = :((4, -19, 11, 25, -31, 10)) - -@inline smoothness_coefficients(::Val{4}, ::Val{0}) = :((2.107, -9.402, 7.042, -1.854, 11.003, -17.246, 4.642, 7.043, -3.882, 0.547)) -@inline smoothness_coefficients(::Val{4}, ::Val{1}) = :((0.547, -2.522, 1.922, -0.494, 3.443, - 5.966, 1.602, 2.843, -1.642, 0.267)) -@inline smoothness_coefficients(::Val{4}, ::Val{2}) = :((0.267, -1.642, 1.602, -0.494, 2.843, - 5.966, 1.922, 3.443, -2.522, 0.547)) -@inline smoothness_coefficients(::Val{4}, ::Val{3}) = :((0.547, -3.882, 4.642, -1.854, 7.043, -17.246, 7.042, 11.003, -9.402, 2.107)) - -@inline smoothness_coefficients(::Val{5}, ::Val{0}) = :((1.07918, -6.49501, 7.58823, -4.11487, 0.86329, 10.20563, -24.62076, 13.58458, -2.88007, 15.21393, -17.04396, 3.64863, 4.82963, -2.08501, 0.22658)) -@inline smoothness_coefficients(::Val{5}, ::Val{1}) = :((0.22658, -1.40251, 1.65153, -0.88297, 0.18079, 2.42723, -6.11976, 3.37018, -0.70237, 4.06293, -4.64976, 0.99213, 1.38563, -0.60871, 0.06908)) -@inline smoothness_coefficients(::Val{5}, ::Val{2}) = :((0.06908, -0.51001, 0.67923, -0.38947, 0.08209, 1.04963, -2.99076, 1.79098, -0.38947, 2.31153, -2.99076, 0.67923, 1.04963, -0.51001, 0.06908)) -@inline smoothness_coefficients(::Val{5}, ::Val{3}) = :((0.06908, -0.60871, 0.99213, -0.70237, 0.18079, 1.38563, -4.64976, 3.37018, -0.88297, 4.06293, -6.11976, 1.65153, 2.42723, -1.40251, 0.22658)) -@inline smoothness_coefficients(::Val{5}, ::Val{4}) = :((0.22658, -2.08501, 3.64863, -2.88007, 0.86329, 4.82963, -17.04396, 13.58458, -4.11487, 15.21393, -24.62076, 7.58823, 10.20563, -6.49501, 1.07918)) - -@inline smoothness_coefficients(::Val{6}, ::Val{0}) = :((0.6150211, -4.7460464, 7.6206736, -6.3394124, 2.7060170, -0.4712740, 9.4851237, -31.1771244, 26.2901672, -11.3206788, 1.9834350, 26.0445372, -44.4003904, 19.2596472, -3.3918804, 19.0757572, -16.6461044, 2.9442256, 3.6480687, -1.2950184, 0.1152561)) -@inline smoothness_coefficients(::Val{6}, ::Val{1}) = :((0.1152561, -0.9117992, 1.4742480, -1.2183636, 0.5134574, -0.0880548, 1.9365967, -6.5224244, 5.5053752, -2.3510468, 0.4067018, 5.6662212, -9.7838784, 4.2405032, -0.7408908, 4.3093692, -3.7913324, 0.6694608, 0.8449957, -0.3015728, 0.0271779)) -@inline smoothness_coefficients(::Val{6}, ::Val{2}) = :((0.0271779, -0.2380800, 0.4086352, -0.3462252, 0.1458762, -0.0245620, 0.5653317, -2.0427884, 1.7905032, -0.7727988, 0.1325006, 1.9510972, -3.5817664, 1.5929912, -0.2792660, 1.7195652, -1.5880404, 0.2863984, 0.3824847, -0.1429976, 0.0139633)) -@inline smoothness_coefficients(::Val{6}, ::Val{3}) = :((0.0139633, -0.1429976, 0.2863984, -0.2792660, 0.1325006, -0.0245620, 0.3824847, -1.5880404, 1.5929912, -0.7727988, 0.1458762, 1.7195652, -3.5817664, 1.7905032, -0.3462252, 1.9510972, -2.0427884, 0.4086352, 0.5653317, -0.2380800, 0.0271779)) -@inline smoothness_coefficients(::Val{6}, ::Val{4}) = :((0.0271779, -0.3015728, 0.6694608, -0.7408908, 0.4067018, -0.0880548, 0.8449957, -3.7913324, 4.2405032, -2.3510468, 0.5134574, 4.3093692, -9.7838784, 5.5053752, -1.2183636, 5.6662212, -6.5224244, 1.4742480, 1.9365967, -0.9117992, 0.1152561)) -@inline smoothness_coefficients(::Val{6}, ::Val{5}) = :((0.1152561, -1.2950184, 2.9442256, -3.3918804, 1.9834350, -0.4712740, 3.6480687, -16.6461044, 19.2596472, -11.3206788, 2.7060170, 19.0757572, -44.4003904, 26.2901672, -6.3394124, 26.0445372, -31.1771244, 7.6206736, 9.4851237, -4.7460464, 0.6150211)) +for FT in supported_float_types + @eval begin + """ + smoothness_coefficients(::Val{FT}, ::Val{buffer}, ::Val{stencil}) + + Return the coefficients used to calculate the smoothness indicators for the stencil + number `stencil` of a WENO reconstruction of order `buffer * 2 - 1`. The coefficients + are ordered in such a way to calculate the smoothness in the following fashion: + + ```julia + buffer = 4 + stencil = 0 + + ψ = # The stencil corresponding to S₀ with buffer 4 (7th order WENO) + + C = smoothness_coefficients(Val(buffer), Val(0)) + + # The smoothness indicator + β = ψ[1] * (C[1] * ψ[1] + C[2] * ψ[2] + C[3] * ψ[3] + C[4] * ψ[4]) + + ψ[2] * (C[5] * ψ[2] + C[6] * ψ[3] + C[7] * ψ[4]) + + ψ[3] * (C[8] * ψ[3] + C[9] * ψ[4]) + ψ[4] * (C[10] * ψ[4]) + ``` + + This last operation is metaprogrammed in the function `metaprogrammed_smoothness_operation` + """ + @inline smoothness_coefficients(::Val{$FT}, ::Val{2}, ::Val{0}) = $(FT.((1, -2, 1))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{2}, ::Val{1}) = $(FT.((1, -2, 1))) + + @inline smoothness_coefficients(::Val{$FT}, ::Val{3}, ::Val{0}) = $(FT.((10, -31, 11, 25, -19, 4))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{3}, ::Val{1}) = $(FT.((4, -13, 5, 13, -13, 4))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{3}, ::Val{2}) = $(FT.((4, -19, 11, 25, -31, 10))) + + @inline smoothness_coefficients(::Val{$FT}, ::Val{4}, ::Val{0}) = $(FT.((2.107, -9.402, 7.042, -1.854, 11.003, -17.246, 4.642, 7.043, -3.882, 0.547))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{4}, ::Val{1}) = $(FT.((0.547, -2.522, 1.922, -0.494, 3.443, - 5.966, 1.602, 2.843, -1.642, 0.267))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{4}, ::Val{2}) = $(FT.((0.267, -1.642, 1.602, -0.494, 2.843, - 5.966, 1.922, 3.443, -2.522, 0.547))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{4}, ::Val{3}) = $(FT.((0.547, -3.882, 4.642, -1.854, 7.043, -17.246, 7.042, 11.003, -9.402, 2.107))) + + @inline smoothness_coefficients(::Val{$FT}, ::Val{5}, ::Val{0}) = $(FT.((1.07918, -6.49501, 7.58823, -4.11487, 0.86329, 10.20563, -24.62076, 13.58458, -2.88007, 15.21393, -17.04396, 3.64863, 4.82963, -2.08501, 0.22658))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{5}, ::Val{1}) = $(FT.((0.22658, -1.40251, 1.65153, -0.88297, 0.18079, 2.42723, -6.11976, 3.37018, -0.70237, 4.06293, -4.64976, 0.99213, 1.38563, -0.60871, 0.06908))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{5}, ::Val{2}) = $(FT.((0.06908, -0.51001, 0.67923, -0.38947, 0.08209, 1.04963, -2.99076, 1.79098, -0.38947, 2.31153, -2.99076, 0.67923, 1.04963, -0.51001, 0.06908))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{5}, ::Val{3}) = $(FT.((0.06908, -0.60871, 0.99213, -0.70237, 0.18079, 1.38563, -4.64976, 3.37018, -0.88297, 4.06293, -6.11976, 1.65153, 2.42723, -1.40251, 0.22658))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{5}, ::Val{4}) = $(FT.((0.22658, -2.08501, 3.64863, -2.88007, 0.86329, 4.82963, -17.04396, 13.58458, -4.11487, 15.21393, -24.62076, 7.58823, 10.20563, -6.49501, 1.07918))) + + @inline smoothness_coefficients(::Val{$FT}, ::Val{6}, ::Val{0}) = $(FT.((0.6150211, -4.7460464, 7.6206736, -6.3394124, 2.7060170, -0.4712740, 9.4851237, -31.1771244, 26.2901672, -11.3206788, 1.9834350, 26.0445372, -44.4003904, 19.2596472, -3.3918804, 19.0757572, -16.6461044, 2.9442256, 3.6480687, -1.2950184, 0.1152561))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{6}, ::Val{1}) = $(FT.((0.1152561, -0.9117992, 1.4742480, -1.2183636, 0.5134574, -0.0880548, 1.9365967, -6.5224244, 5.5053752, -2.3510468, 0.4067018, 5.6662212, -9.7838784, 4.2405032, -0.7408908, 4.3093692, -3.7913324, 0.6694608, 0.8449957, -0.3015728, 0.0271779))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{6}, ::Val{2}) = $(FT.((0.0271779, -0.2380800, 0.4086352, -0.3462252, 0.1458762, -0.0245620, 0.5653317, -2.0427884, 1.7905032, -0.7727988, 0.1325006, 1.9510972, -3.5817664, 1.5929912, -0.2792660, 1.7195652, -1.5880404, 0.2863984, 0.3824847, -0.1429976, 0.0139633))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{6}, ::Val{3}) = $(FT.((0.0139633, -0.1429976, 0.2863984, -0.2792660, 0.1325006, -0.0245620, 0.3824847, -1.5880404, 1.5929912, -0.7727988, 0.1458762, 1.7195652, -3.5817664, 1.7905032, -0.3462252, 1.9510972, -2.0427884, 0.4086352, 0.5653317, -0.2380800, 0.0271779))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{6}, ::Val{4}) = $(FT.((0.0271779, -0.3015728, 0.6694608, -0.7408908, 0.4067018, -0.0880548, 0.8449957, -3.7913324, 4.2405032, -2.3510468, 0.5134574, 4.3093692, -9.7838784, 5.5053752, -1.2183636, 5.6662212, -6.5224244, 1.4742480, 1.9365967, -0.9117992, 0.1152561))) + @inline smoothness_coefficients(::Val{$FT}, ::Val{6}, ::Val{5}) = $(FT.((0.1152561, -1.2950184, 2.9442256, -3.3918804, 1.9834350, -0.4712740, 3.6480687, -16.6461044, 19.2596472, -11.3206788, 2.7060170, 19.0757572, -44.4003904, 26.2901672, -6.3394124, 26.0445372, -31.1771244, 7.6206736, 9.4851237, -4.7460464, 0.6150211))) + end +end # The rule for calculating smoothness indicators is the following (example WENO{4} which is seventh order) # ψ[1] (C[1] * ψ[1] + C[2] * ψ[2] + C[3] * ψ[3] + C[4] * ψ[4]) + @@ -255,12 +266,12 @@ while for `buffer == 4` unrolls into @inline smoothness_indicator(ψ, args...) = zero(ψ[1]) # This is a fallback method, here only for documentation purposes # Smoothness indicators for stencil `stencil` for left and right biased reconstruction -for buffer in [2, 3, 4, 5, 6] +for buffer in advection_buffers[2:end] # WENO{<:Any, 1} does not exist @eval @inline smoothness_operation(scheme::WENO{$buffer}, ψ, C) = @inbounds $(metaprogrammed_smoothness_operation(buffer)) - for stencil in 0:buffer-1 - @eval @inline smoothness_indicator(ψ, scheme::WENO{$buffer, FT}, ::Val{$stencil}) where FT = - smoothness_operation(scheme, ψ, map(FT, $(smoothness_coefficients(Val(buffer), Val(stencil))))) + for stencil in 0:buffer-1, FT in supported_float_types + @eval @inline smoothness_indicator(ψ, scheme::WENO{$buffer, $FT}, ::Val{$stencil}) = + smoothness_operation(scheme, ψ, $(smoothness_coefficients(Val(FT), Val(buffer), Val(stencil)))) end end @@ -288,13 +299,13 @@ end @inline function metaprogrammed_zweno_alpha_loop(buffer) elem = Vector(undef, buffer) for stencil = 1:buffer - elem[stencil] = :(convert(FT, C★(scheme, Val($(stencil-1)))) * (1 + (τ / (β[$stencil] + FT(ε)))^ƞ)) + elem[stencil] = :(C★(scheme, Val($(stencil-1))) * (1 + (τ / (β[$stencil] + ε))^2)) end return :($(elem...),) end -for buffer in [2, 3, 4, 5, 6] +for buffer in advection_buffers[2:end] @eval begin @inline beta_sum(scheme::WENO{$buffer, FT}, β₁, β₂) where FT = @inbounds $(metaprogrammed_beta_sum(buffer)) @inline beta_loop(scheme::WENO{$buffer, FT}, ψ) where FT = @inbounds $(metaprogrammed_beta_loop(buffer)) @@ -495,7 +506,7 @@ Here, [`biased_p`](@ref) is the function that computes the linear reconstruction @inline weno_reconstruction(scheme, bias, ψ, args...) = zero(ψ[1][1]) # Fallback only for documentation purposes # Calculation of WENO reconstructed value v⋆ = ∑ᵣ(wᵣv̂ᵣ) -for buffer in [2, 3, 4, 5, 6] +for buffer in advection_buffers[2:end] @eval @inline weno_reconstruction(scheme::WENO{$buffer}, bias, ψ, ω, cT, val, idx, loc) = @inbounds $(metaprogrammed_weno_reconstruction(buffer)) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl index 67302844fa..b87d1a09be 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl @@ -27,9 +27,9 @@ end end end -@inline function ab2_step_G(i, j, k, grid, ℓx, ℓy, ℓz, G⁻, Gⁿ, χ::FT) where FT - C₁ = convert(FT, 3/2) + χ - C₂ = convert(FT, 1/2) + χ +@inline function ab2_step_G(i, j, k, grid, ℓx, ℓy, ℓz, G⁻, Gⁿ, χ) + C₁ = 3 * one(grid) / 2 + χ + C₂ = one(grid) / 2 + χ # multiply G⁻ by false if C₂ is zero to # prevent propagationg possible NaNs @@ -104,8 +104,8 @@ end GUi = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center()) GVi = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center()) - @inbounds GUⁿ[i, j, 1] = convert(FT, 2/3) * GUi + GU⁻[i, j, 1] - @inbounds GVⁿ[i, j, 1] = convert(FT, 2/3) * GVi + GV⁻[i, j, 1] + @inbounds GUⁿ[i, j, 1] = 2 * GUi / 3 + GU⁻[i, j, 1] + @inbounds GVⁿ[i, j, 1] = 2 * GVi / 3 + GV⁻[i, j, 1] end @inline function compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl index 9052c1eb7f..7ebb0a6334 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl @@ -256,7 +256,7 @@ end averaging_weights = averaging_weights[1:idx] averaging_weights ./= sum(averaging_weights) - return Δτ, tuple(averaging_weights...) + return Δτ, map(FT, tuple(averaging_weights...)) end Base.summary(s::FixedTimeStepSize) = string("FixedTimeStepSize($(prettytime(s.Δt_barotropic)))") diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl index 7ac5d8e2c4..c3d435e2db 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl @@ -27,8 +27,8 @@ end function hydrostatic_tendency_fields(velocities, free_surface::SplitExplicitFreeSurface, grid, tracer_names) u = XFaceField(grid) v = YFaceField(grid) - U = deepcopy(free_surface.barotropic_velocities.U) - V = deepcopy(free_surface.barotropic_velocities.V) + U = similar(free_surface.barotropic_velocities.U) + V = similar(free_surface.barotropic_velocities.V) tracers = TracerFields(tracer_names, grid) return merge((u=u, v=v, U=U, V=V), tracers) end @@ -37,7 +37,7 @@ previous_hydrostatic_tendency_fields(::Val{:QuasiAdamsBashforth2}, args...) = hy previous_hydrostatic_tendency_fields(::Val{:SplitRungeKutta3}, args...) = nothing function previous_hydrostatic_tendency_fields(::Val{:SplitRungeKutta3}, velocities, free_surface::SplitExplicitFreeSurface, args...) - U = deepcopy(free_surface.barotropic_velocities.U) - V = deepcopy(free_surface.barotropic_velocities.V) + U = similar(free_surface.barotropic_velocities.U) + V = similar(free_surface.barotropic_velocities.V) return (; U=U, V=V) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl index 0e0fe6b20a..9447b3effd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl @@ -137,4 +137,4 @@ where `c = C[tracer_index]`. - immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields) + biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields) + forcing(i, j, k, grid, clock, model_fields)) -end +end \ No newline at end of file diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index cf245a5b6e..8159c3fbac 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -132,6 +132,10 @@ function __init__() end end +# Supported types for the Advection module. +# To support new floating point types, add them here +const supported_float_types = [Float32, Float64] + ##### ##### Abstract types ##### From 19caa11ab54b402ff30180f65f7731e75766c5ba Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 23 Jan 2025 12:29:47 +0100 Subject: [PATCH 24/31] Disambiguate interpolation of `::Number` on grids with `Flat` dimensions (#4051) * fix issue * restart the tests --- src/Operators/interpolation_operators.jl | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/Operators/interpolation_operators.jl b/src/Operators/interpolation_operators.jl index dd495a39e9..0e7548ac29 100644 --- a/src/Operators/interpolation_operators.jl +++ b/src/Operators/interpolation_operators.jl @@ -86,12 +86,21 @@ using Oceananigans.Grids: XFlatGrid, YFlatGrid, ZFlatGrid @inline ℑxᶜᵃᵃ(i, j, k, grid::XFlatGrid, u) = @inbounds u[i, j, k] @inline ℑxᶠᵃᵃ(i, j, k, grid::XFlatGrid, c) = @inbounds c[i, j, k] -@inline ℑyᵃᶜᵃ(i, j, k, grid::YFlatGrid, w) = @inbounds w[i, j, k] +@inline ℑyᵃᶜᵃ(i, j, k, grid::YFlatGrid, v) = @inbounds v[i, j, k] @inline ℑyᵃᶠᵃ(i, j, k, grid::YFlatGrid, c) = @inbounds c[i, j, k] @inline ℑzᵃᵃᶜ(i, j, k, grid::ZFlatGrid, w) = @inbounds w[i, j, k] @inline ℑzᵃᵃᶠ(i, j, k, grid::ZFlatGrid, c) = @inbounds c[i, j, k] +@inline ℑxᶜᵃᵃ(i, j, k, grid::XFlatGrid, u::Number) = u +@inline ℑxᶠᵃᵃ(i, j, k, grid::XFlatGrid, c::Number) = c + +@inline ℑyᵃᶜᵃ(i, j, k, grid::YFlatGrid, v::Number) = v +@inline ℑyᵃᶠᵃ(i, j, k, grid::YFlatGrid, c::Number) = c + +@inline ℑzᵃᵃᶜ(i, j, k, grid::ZFlatGrid, w::Number) = w +@inline ℑzᵃᵃᶠ(i, j, k, grid::ZFlatGrid, c::Number) = c + @inline ℑxᶜᵃᵃ(i, j, k, grid::XFlatGrid, f::F, args...) where {F<:Function} = f(i, j, k, grid, args...) @inline ℑxᶠᵃᵃ(i, j, k, grid::XFlatGrid, f::F, args...) where {F<:Function} = f(i, j, k, grid, args...) From 3e89def32055c3a2cc635160080a74d469e24e5d Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 23 Jan 2025 23:22:31 +0100 Subject: [PATCH 25/31] Add a fallback `set_timestepper!` (#4059) * correct the bug * fallback --- src/OutputWriters/checkpointer.jl | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/OutputWriters/checkpointer.jl b/src/OutputWriters/checkpointer.jl index 879d7d89e3..083c9790a1 100644 --- a/src/OutputWriters/checkpointer.jl +++ b/src/OutputWriters/checkpointer.jl @@ -3,7 +3,7 @@ using Glob using Oceananigans using Oceananigans: fields, prognostic_fields using Oceananigans.Fields: offset_data -using Oceananigans.TimeSteppers: RungeKutta3TimeStepper, QuasiAdamsBashforth2TimeStepper +using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper import Oceananigans.Fields: set! @@ -280,7 +280,9 @@ function set_time_stepper_tendencies!(timestepper, file, model_fields, addr) return nothing end -set_time_stepper!(timestepper::RungeKutta3TimeStepper, args...) = nothing +# For self-starting timesteppers like RK3 we do nothing +set_time_stepper!(timestepper, args...) = nothing + set_time_stepper!(timestepper::QuasiAdamsBashforth2TimeStepper, args...) = set_time_stepper_tendencies!(timestepper, args...) From 12ee14ed5409841ea56268e2b028e6ec58048cc5 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 24 Jan 2025 10:34:26 +0100 Subject: [PATCH 26/31] Bugfix for advection on distributed topologies (#4058) * first change * add a test * bugfix * works * don't need this * valid partitions * add comment * simplify tests --- ...topologically_conditional_interpolation.jl | 80 ++++++--- .../distributed_immersed_boundaries.jl | 23 +-- ...ompute_hydrostatic_free_surface_buffers.jl | 24 ++- .../update_nonhydrostatic_model_state.jl | 3 +- test/test_distributed_hydrostatic_model.jl | 153 +++++++++--------- test/utils_for_runtests.jl | 12 +- 6 files changed, 162 insertions(+), 133 deletions(-) diff --git a/src/Advection/topologically_conditional_interpolation.jl b/src/Advection/topologically_conditional_interpolation.jl index 9c623c6b57..884396d359 100644 --- a/src/Advection/topologically_conditional_interpolation.jl +++ b/src/Advection/topologically_conditional_interpolation.jl @@ -10,18 +10,26 @@ ##### close to the boundary, or a second-order interpolation if i is close to a boundary. ##### -using Oceananigans.Grids: AbstractUnderlyingGrid, Bounded +using Oceananigans.Grids: AbstractUnderlyingGrid, + Bounded, + RightConnected, + LeftConnected, + topology, + architecture const AUG = AbstractUnderlyingGrid +# topologies bounded at least on one side +const BT = Union{Bounded, RightConnected, LeftConnected} + # Bounded underlying Grids -const AUGX = AUG{<:Any, <:Bounded} -const AUGY = AUG{<:Any, <:Any, <:Bounded} -const AUGZ = AUG{<:Any, <:Any, <:Any, <:Bounded} -const AUGXY = AUG{<:Any, <:Bounded, <:Bounded} -const AUGXZ = AUG{<:Any, <:Bounded, <:Any, <:Bounded} -const AUGYZ = AUG{<:Any, <:Any, <:Bounded, <:Bounded} -const AUGXYZ = AUG{<:Any, <:Bounded, <:Bounded, <:Bounded} +const AUGX = AUG{<:Any, <:BT} +const AUGY = AUG{<:Any, <:Any, <:BT} +const AUGZ = AUG{<:Any, <:Any, <:Any, <:BT} +const AUGXY = AUG{<:Any, <:BT, <:BT} +const AUGXZ = AUG{<:Any, <:BT, <:Any, <:BT} +const AUGYZ = AUG{<:Any, <:Any, <:BT, <:BT} +const AUGXYZ = AUG{<:Any, <:BT, <:BT, <:BT} # Left-biased buffers are smaller by one grid point on the right side; vice versa for right-biased buffers # Center interpolation stencil look at i + 1 (i.e., require one less point on the left) @@ -29,20 +37,40 @@ const AUGXYZ = AUG{<:Any, <:Bounded, <:Bounded, <:Bounded} for dir in (:x, :y, :z) outside_symmetric_haloᶠ = Symbol(:outside_symmetric_halo_, dir, :ᶠ) outside_symmetric_haloᶜ = Symbol(:outside_symmetric_halo_, dir, :ᶜ) - outside_biased_haloᶠ = Symbol(:outside_biased_halo_, dir, :ᶠ) - outside_biased_haloᶜ = Symbol(:outside_biased_halo_, dir, :ᶜ) - required_halo_size = Symbol(:required_halo_size_, dir) + outside_biased_haloᶠ = Symbol(:outside_biased_halo_, dir, :ᶠ) + outside_biased_haloᶜ = Symbol(:outside_biased_halo_, dir, :ᶜ) + required_halo_size = Symbol(:required_halo_size_, dir) @eval begin - @inline $outside_symmetric_haloᶠ(i, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - $required_halo_size(adv)) - @inline $outside_symmetric_haloᶜ(i, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv)) - - @inline $outside_biased_haloᶠ(i, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias - (i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv)) # Right bias - @inline $outside_biased_haloᶜ(i, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias - (i >= $required_halo_size(adv) - 1) & (i <= N + 1 - $required_halo_size(adv)) # Right bias + # Bounded topologies + @inline $outside_symmetric_haloᶠ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - $required_halo_size(adv)) + @inline $outside_symmetric_haloᶜ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv)) + + @inline $outside_biased_haloᶠ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias + (i >= $required_halo_size(adv)) & (i <= N + 1 - $required_halo_size(adv)) # Right bias + @inline $outside_biased_haloᶜ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv)) & (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias + (i >= $required_halo_size(adv) - 1) & (i <= N + 1 - $required_halo_size(adv)) # Right bias + + # Right connected topologies (only test the left side, i.e. the bounded side) + @inline $outside_symmetric_haloᶠ(i, ::Type{RightConnected}, N, adv) = i >= $required_halo_size(adv) + 1 + @inline $outside_symmetric_haloᶜ(i, ::Type{RightConnected}, N, adv) = i >= $required_halo_size(adv) + + @inline $outside_biased_haloᶠ(i, ::Type{RightConnected}, N, adv) = (i >= $required_halo_size(adv) + 1) & # Left bias + (i >= $required_halo_size(adv)) # Right bias + @inline $outside_biased_haloᶜ(i, ::Type{RightConnected}, N, adv) = (i >= $required_halo_size(adv)) & # Left bias + (i >= $required_halo_size(adv) - 1) # Right bias + + # Left bounded topologies (only test the right side, i.e. the bounded side) + @inline $outside_symmetric_haloᶠ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - $required_halo_size(adv)) + @inline $outside_symmetric_haloᶜ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - $required_halo_size(adv)) + + @inline $outside_biased_haloᶠ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias + (i <= N + 1 - $required_halo_size(adv)) # Right bias + @inline $outside_biased_haloᶜ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) & # Left bias + (i <= N + 1 - $required_halo_size(adv)) # Right bias end end + # Separate High order advection from low order advection const HOADV = Union{WENO, Tuple(Centered{N} for N in advection_buffers[2:end])..., @@ -75,22 +103,22 @@ for bias in (:symmetric, :biased) # Conditional high-order interpolation in Bounded directions if ξ == :x @eval begin - @inline $alt1_interp(i, j, k, grid::AUGX, scheme::HOADV, args...) = - ifelse($outside_buffer(i, grid.Nx, scheme), - $interp(i, j, k, grid, scheme, args...), - $alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...)) + @inline $alt1_interp(i, j, k, grid::AUGX, scheme::HOADV, args...) = + ifelse($outside_buffer(i, topology(grid, 1), grid.Nx, scheme), + $interp(i, j, k, grid, scheme, args...), + $alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...)) end elseif ξ == :y @eval begin @inline $alt1_interp(i, j, k, grid::AUGY, scheme::HOADV, args...) = - ifelse($outside_buffer(j, grid.Ny, scheme), + ifelse($outside_buffer(j, topology(grid, 2), grid.Ny, scheme), $interp(i, j, k, grid, scheme, args...), $alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...)) end elseif ξ == :z @eval begin @inline $alt1_interp(i, j, k, grid::AUGZ, scheme::HOADV, args...) = - ifelse($outside_buffer(k, grid.Nz, scheme), + ifelse($outside_buffer(k, topology(grid, 3), grid.Nz, scheme), $interp(i, j, k, grid, scheme, args...), $alt2_interp(i, j, k, grid, scheme.buffer_scheme, args...)) end @@ -100,11 +128,11 @@ for bias in (:symmetric, :biased) end @inline _multi_dimensional_reconstruction_x(i, j, k, grid::AUGX, scheme, interp, args...) = - ifelse(outside_symmetric_bufferᶜ(i, grid.Nx, scheme), + ifelse(outside_symmetric_bufferᶜ(i, topology(grid, 1), grid.Nx, scheme), multi_dimensional_reconstruction_x(i, j, k, grid::AUGX, scheme, interp, args...), interp(i, j, k, grid, scheme, args...)) @inline _multi_dimensional_reconstruction_y(i, j, k, grid::AUGY, scheme, interp, args...) = - ifelse(outside_symmetric_bufferᶜ(j, grid.Ny, scheme), + ifelse(outside_symmetric_bufferᶜ(j, topology(grid, 2), grid.Ny, scheme), multi_dimensional_reconstruction_y(i, j, k, grid::AUGY, scheme, interp, args...), interp(i, j, k, grid, scheme, args...)) diff --git a/src/DistributedComputations/distributed_immersed_boundaries.jl b/src/DistributedComputations/distributed_immersed_boundaries.jl index 01ffb3d67b..3395d01fd2 100644 --- a/src/DistributedComputations/distributed_immersed_boundaries.jl +++ b/src/DistributedComputations/distributed_immersed_boundaries.jl @@ -121,22 +121,25 @@ function map_interior_active_cells(ibg::ImmersedBoundaryGrid{<:Any, <:Any, <:Any Nx, Ny, Nz = size(ibg) Hx, Hy, _ = halo_size(ibg) - x_boundary = (Hx, Ny, Nz) - y_boundary = (Nx, Hy, Nz) - - left_offsets = (0, 0, 0) - right_x_offsets = (Nx-Hx, 0, 0) - right_y_offsets = (0, Ny-Hy, 0) + west_boundary = (1:Hx, 1:Ny, 1:Nz) + east_boundary = (Nx-Hx+1:Nx, 1:Ny, 1:Nz) + south_boundary = (1:Nx, 1:Hy, 1:Nz) + north_boundary = (1:Nx, Ny-Hy+1:Ny, 1:Nz) include_west = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected) include_east = !isa(ibg, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected) include_south = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected) include_north = !isa(ibg, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected) - west_halo_dependent_cells = include_west ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, left_offsets)) : nothing - east_halo_dependent_cells = include_east ? interior_active_indices(ibg; parameters = KernelParameters(x_boundary, right_x_offsets)) : nothing - south_halo_dependent_cells = include_south ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, left_offsets)) : nothing - north_halo_dependent_cells = include_north ? interior_active_indices(ibg; parameters = KernelParameters(y_boundary, right_y_offsets)) : nothing + west_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(west_boundary...)) + east_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(east_boundary...)) + south_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(south_boundary...)) + north_halo_dependent_cells = interior_active_indices(ibg; parameters = KernelParameters(north_boundary...)) + + west_halo_dependent_cells = ifelse(include_west, west_halo_dependent_cells, nothing) + east_halo_dependent_cells = ifelse(include_east, east_halo_dependent_cells, nothing) + south_halo_dependent_cells = ifelse(include_south, south_halo_dependent_cells, nothing) + north_halo_dependent_cells = ifelse(include_north, north_halo_dependent_cells, nothing) nx = Rx == 1 ? Nx : (Tx == RightConnected || Tx == LeftConnected ? Nx - Hx : Nx - 2Hx) ny = Ry == 1 ? Ny : (Ty == RightConnected || Ty == LeftConnected ? Ny - Hy : Ny - 2Hy) diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl index 7dc110d662..05b42eb03f 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl @@ -35,17 +35,16 @@ end function compute_buffer_tendency_contributions!(grid::DistributedActiveCellsIBG, arch, model) maps = grid.interior_active_cells - for (name, map) in zip(keys(maps), maps) - - # If there exists a buffer map, then we compute the buffer contributions. If not, the - # buffer contributions have already been calculated. We exclude the interior because it has - # already been calculated - compute_buffer = (name != :interior) && !isnothing(map) - - if compute_buffer - active_cells_map = retrieve_interior_active_cells_map(grid, Val(name)) - compute_hydrostatic_free_surface_tendency_contributions!(model, :xyz; active_cells_map) - end + for name in (:west_halo_dependent_cells, + :east_halo_dependent_cells, + :south_halo_dependent_cells, + :north_halo_dependent_cells) + + active_cells_map = @inbounds maps[name] + + # If the map == nothing, we don't need to compute the buffer because + # the buffer is not adjacent to a processor boundary + !isnothing(map) && compute_hydrostatic_free_surface_tendency_contributions!(model, :xyz; active_cells_map) end return nothing @@ -56,9 +55,6 @@ function buffer_w_kernel_parameters(grid, arch) Nx, Ny, _ = size(grid) Hx, Hy, _ = halo_size(grid) - Sx = (Hx, Ny+2) - Sy = (Nx+2, Hy) - # Offsets in tangential direction are == -1 to # cover the required corners param_west = (-Hx+2:1, 0:Ny+1) diff --git a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl index 705ca507ac..260fd0e0ad 100644 --- a/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl +++ b/src/Models/NonhydrostaticModels/update_nonhydrostatic_model_state.jl @@ -41,7 +41,8 @@ function update_state!(model::NonhydrostaticModel, callbacks=[]; compute_tendenc # Calculate diffusivities and hydrostatic pressure @apply_regionally compute_auxiliaries!(model) - fill_halo_regions!(model.diffusivity_fields; only_local_halos = true) + + fill_halo_regions!(model.diffusivity_fields; only_local_halos=true) for callback in callbacks callback.callsite isa UpdateStateCallsite && callback(model) diff --git a/test/test_distributed_hydrostatic_model.jl b/test/test_distributed_hydrostatic_model.jl index e5ab4bd461..68ce34424a 100644 --- a/test/test_distributed_hydrostatic_model.jl +++ b/test/test_distributed_hydrostatic_model.jl @@ -49,12 +49,12 @@ function rotation_with_shear_test(grid, closure=nothing) end model = HydrostaticFreeSurfaceModel(; grid, - momentum_advection = VectorInvariant(), + momentum_advection = WENOVectorInvariant(order=3), free_surface = free_surface, coriolis = coriolis, closure, tracers, - tracer_advection = WENO(), + tracer_advection = WENO(order=3), buoyancy = BuoyancyTracer()) g = model.free_surface.gravitational_acceleration @@ -75,41 +75,89 @@ function rotation_with_shear_test(grid, closure=nothing) Δt_local = 0.1 * Δ_min(grid) / sqrt(g * grid.Lz) Δt = all_reduce(min, Δt_local, architecture(grid)) - simulation = Simulation(model; Δt, stop_iteration = 10, verbose = false) - run!(simulation) + for _ in 1:10 + time_step!(model, Δt) + end return model end Nx = 32 -Ny = 32 +Ny = 32 for arch in archs - @testset "Testing distributed solid body rotation" begin - underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 3), - halo = (4, 4, 4), - latitude = (-80, 80), - longitude = (-160, 160), - z = (-1, 0), - radius = 1, - topology=(Bounded, Bounded, Bounded)) - - bottom(λ, φ) = -30 < λ < 30 && -40 < φ < 20 ? 0 : - 1 - - immersed_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom)) - immersed_active_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom); active_cells_map = true) - - global_underlying_grid = reconstruct_global_grid(underlying_grid) - global_immersed_grid = ImmersedBoundaryGrid(global_underlying_grid, GridFittedBottom(bottom)) - - for (grid, global_grid) in zip((underlying_grid, immersed_grid, immersed_active_grid), (global_underlying_grid, global_immersed_grid, global_immersed_grid)) - if arch.local_rank == 0 - @info " Testing distributed solid body rotation with $(ranks(arch)) ranks on $(typeof(grid).name.wrapper)" + + # We do not test on `Fractional` partitions where we cannot easily ensure that H ≤ N + # which would lead to different advection schemes for partitioned and non-partitioned grids. + # `Fractional` is, however, tested in regression tests where the horizontal dimensions are larger. + valid_x_partition = !(arch.partition.x isa Fractional) + valid_y_partition = !(arch.partition.y isa Fractional) + valid_z_partition = !(arch.partition.z isa Fractional) + + if valid_x_partition & valid_y_partition & valid_z_partition + @testset "Testing distributed solid body rotation" begin + underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 3), + halo = (4, 4, 3), + latitude = (-80, 80), + longitude = (-160, 160), + z = (-1, 0), + radius = 1, + topology=(Bounded, Bounded, Bounded)) + + bottom(λ, φ) = -30 < λ < 30 && -40 < φ < 20 ? 0 : - 1 + + immersed_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom)) + immersed_active_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom); active_cells_map = true) + + global_underlying_grid = reconstruct_global_grid(underlying_grid) + global_immersed_grid = ImmersedBoundaryGrid(global_underlying_grid, GridFittedBottom(bottom)) + + for (grid, global_grid) in zip((underlying_grid, immersed_grid, immersed_active_grid), (global_underlying_grid, global_immersed_grid, global_immersed_grid)) + if arch.local_rank == 0 + @info " Testing distributed solid body rotation with $(ranks(arch)) ranks on $(typeof(grid).name.wrapper)" + end + + # "s" for "serial" computation, "p" for parallel + ms = rotation_with_shear_test(global_grid) + mp = rotation_with_shear_test(grid) + + us = interior(on_architecture(CPU(), ms.velocities.u)) + vs = interior(on_architecture(CPU(), ms.velocities.v)) + ws = interior(on_architecture(CPU(), ms.velocities.w)) + cs = interior(on_architecture(CPU(), ms.tracers.c)) + ηs = interior(on_architecture(CPU(), ms.free_surface.η)) + + cpu_arch = cpu_architecture(arch) + + up = interior(on_architecture(cpu_arch, mp.velocities.u)) + vp = interior(on_architecture(cpu_arch, mp.velocities.v)) + wp = interior(on_architecture(cpu_arch, mp.velocities.w)) + cp = interior(on_architecture(cpu_arch, mp.tracers.c)) + ηp = interior(on_architecture(cpu_arch, mp.free_surface.η)) + + us = partition(us, cpu_arch, size(up)) + vs = partition(vs, cpu_arch, size(vp)) + ws = partition(ws, cpu_arch, size(wp)) + cs = partition(cs, cpu_arch, size(cp)) + ηs = partition(ηs, cpu_arch, size(ηp)) + + atol = eps(eltype(grid)) + rtol = sqrt(eps(eltype(grid))) + + @test all(isapprox(up, us; atol, rtol)) + @test all(isapprox(vp, vs; atol, rtol)) + @test all(isapprox(wp, ws; atol, rtol)) + @test all(isapprox(cp, cs; atol, rtol)) + @test all(isapprox(ηp, ηs; atol, rtol)) end + # CATKE works only with synchronized communication at the moment + arch = synchronized(arch) + closure = CATKEVerticalDiffusivity() + # "s" for "serial" computation, "p" for parallel - ms = rotation_with_shear_test(global_grid) - mp = rotation_with_shear_test(grid) + ms = rotation_with_shear_test(global_underlying_grid, closure) + mp = rotation_with_shear_test(underlying_grid, closure) us = interior(on_architecture(CPU(), ms.velocities.u)) vs = interior(on_architecture(CPU(), ms.velocities.v)) @@ -131,8 +179,8 @@ for arch in archs cs = partition(cs, cpu_arch, size(cp)) ηs = partition(ηs, cpu_arch, size(ηp)) - atol = eps(eltype(grid)) - rtol = sqrt(eps(eltype(grid))) + atol = eps(eltype(global_underlying_grid)) + rtol = sqrt(eps(eltype(global_underlying_grid))) @test all(isapprox(up, us; atol, rtol)) @test all(isapprox(vp, vs; atol, rtol)) @@ -140,52 +188,5 @@ for arch in archs @test all(isapprox(cp, cs; atol, rtol)) @test all(isapprox(ηp, ηs; atol, rtol)) end - - # # CATKE works only with synchronized communication at the moment - # arch = synchronized(arch) - # closure = CATKEVerticalDiffusivity() - - # underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 3), - # halo = (4, 4, 4), - # latitude = (-80, 80), - # longitude = (-160, 160), - # z = (-1, 0), - # radius = 1, - # topology=(Bounded, Bounded, Bounded)) - - # bottom(λ, φ) = -30 < λ < 30 && -40 < φ < 20 ? 0 : - 1 - - # # "s" for "serial" computation, "p" for parallel - # ms = rotation_with_shear_test(global_underlying_grid, closure) - # mp = rotation_with_shear_test(underlying_grid, closure) - - # us = interior(on_architecture(CPU(), ms.velocities.u)) - # vs = interior(on_architecture(CPU(), ms.velocities.v)) - # ws = interior(on_architecture(CPU(), ms.velocities.w)) - # cs = interior(on_architecture(CPU(), ms.tracers.c)) - # ηs = interior(on_architecture(CPU(), ms.free_surface.η)) - - # cpu_arch = cpu_architecture(arch) - - # up = interior(on_architecture(cpu_arch, mp.velocities.u)) - # vp = interior(on_architecture(cpu_arch, mp.velocities.v)) - # wp = interior(on_architecture(cpu_arch, mp.velocities.w)) - # cp = interior(on_architecture(cpu_arch, mp.tracers.c)) - # ηp = interior(on_architecture(cpu_arch, mp.free_surface.η)) - - # us = partition(us, cpu_arch, size(up)) - # vs = partition(vs, cpu_arch, size(vp)) - # ws = partition(ws, cpu_arch, size(wp)) - # cs = partition(cs, cpu_arch, size(cp)) - # ηs = partition(ηs, cpu_arch, size(ηp)) - - # atol = eps(eltype(immersed_active_grid)) - # rtol = sqrt(eps(eltype(immersed_active_grid))) - - # @test all(isapprox(up, us; atol, rtol)) - # @test all(isapprox(vp, vs; atol, rtol)) - # @test all(isapprox(wp, ws; atol, rtol)) - # @test all(isapprox(cp, cs; atol, rtol)) - # @test all(isapprox(ηp, ηs; atol, rtol)) end end diff --git a/test/utils_for_runtests.jl b/test/utils_for_runtests.jl index 11dc6ecd5e..36fb02ce6a 100644 --- a/test/utils_for_runtests.jl +++ b/test/utils_for_runtests.jl @@ -43,12 +43,12 @@ function test_architectures() # `Partition(x = 2, y = 2)`, and different fractional subdivisions in x, y and xy if mpi_test if MPI.Initialized() && MPI.Comm_size(MPI.COMM_WORLD) == 4 - return (Distributed(child_arch; partition = Partition(4)), - Distributed(child_arch; partition = Partition(1, 4)), - Distributed(child_arch; partition = Partition(2, 2)), - Distributed(child_arch; partition = Partition(x = Fractional(1, 2, 3, 4))), - Distributed(child_arch; partition = Partition(y = Fractional(1, 2, 3, 4))), - Distributed(child_arch; partition = Partition(x = Fractional(1, 2), y = Equal()))) + return (Distributed(child_arch; partition=Partition(4)), + Distributed(child_arch; partition=Partition(1, 4)), + Distributed(child_arch; partition=Partition(2, 2)), + Distributed(child_arch; partition=Partition(x = Fractional(1, 2, 3, 4))), + Distributed(child_arch; partition=Partition(y = Fractional(1, 2, 3, 4))), + Distributed(child_arch; partition=Partition(x = Fractional(1, 2), y = Equal()))) else return throw("The MPI partitioning is not correctly configured.") end From fda82754a39246a081377936add6449eba48e46e Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 24 Jan 2025 17:04:13 +0100 Subject: [PATCH 27/31] add a better comment (#4056) --- src/Oceananigans.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 8159c3fbac..b7f2e3d1be 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -132,9 +132,10 @@ function __init__() end end -# Supported types for the Advection module. -# To support new floating point types, add them here -const supported_float_types = [Float32, Float64] +# List of fully-supported floating point types where applicable. +# Currently used only in the Advection module to specialize +# reconstruction schemes (WENO, UpwindBiased, and Centered). +const supported_float_types = (Float32, Float64) ##### ##### Abstract types From 260eb47dfae7f95bfc006207865f4a2afc21e182 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Fri, 24 Jan 2025 14:13:19 -0500 Subject: [PATCH 28/31] Use 1.10.8 in tests (#4062) Might speed things up a bit Co-authored-by: Simone Silvestri --- .buildkite/pipeline.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 910d23d885..848fda0b85 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,5 +1,5 @@ env: - JULIA_VERSION: "1.10.6" + JULIA_VERSION: "1.10.8" JULIA_MINOR_VERSION: "1.10" SVERDRUP_HOME: "/home/users/glwagner" TARTARUS_HOME: "/storage5/buildkite-agent" From 2b3f7ba3b23d2547ef0e203ea47329e344386c36 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Sat, 25 Jan 2025 02:55:29 -0500 Subject: [PATCH 29/31] Bump to 0.95.8 (#4060) To capture #4049 Co-authored-by: Simone Silvestri --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index a1eb9fd78c..83eb9ec8e9 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Oceananigans" uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" authors = ["Climate Modeling Alliance and contributors"] -version = "0.95.7" +version = "0.95.8" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From 06a7c1b6bd34ba585854843adee211102b7e2dcb Mon Sep 17 00:00:00 2001 From: Tomas Chor Date: Tue, 28 Jan 2025 11:12:09 -0500 Subject: [PATCH 30/31] Correct indentation in two test files (#4066) * corrected indents in test_netcdf file * fix indent in test_jld2 file --- test/test_jld2_output_writer.jl | 64 +++++++++++------------ test/test_netcdf_output_writer.jl | 86 ++++++++++++++----------------- 2 files changed, 69 insertions(+), 81 deletions(-) diff --git a/test/test_jld2_output_writer.jl b/test/test_jld2_output_writer.jl index 325a1cf994..f98c677178 100644 --- a/test/test_jld2_output_writer.jl +++ b/test/test_jld2_output_writer.jl @@ -17,14 +17,13 @@ function jld2_sliced_field_output(model, outputs=model.velocities) simulation = Simulation(model, Δt=1.0, stop_iteration=1) - simulation.output_writers[:velocities] = - JLD2OutputWriter(model, outputs, - schedule = TimeInterval(1), - indices = (1:2, 1:4, :), - with_halos = false, - dir = ".", - filename = "test.jld2", - overwrite_existing = true) + simulation.output_writers[:velocities] = JLD2OutputWriter(model, outputs, + schedule = TimeInterval(1), + indices = (1:2, 1:4, :), + with_halos = false, + dir = ".", + filename = "test.jld2", + overwrite_existing = true) run!(simulation) @@ -227,14 +226,13 @@ function test_jld2_time_averaging(arch) jld2_outputs = Dict("c1" => ∫c1_dxdy, "c2" => ∫c2_dxdy) horizontal_average_jld2_filepath = "decay_averaged_field_test.jld2" - simulation.output_writers[:horizontal_average] = JLD2OutputWriter( - model, - jld2_outputs, - schedule = TimeInterval(10Δt), - dir = ".", - with_halos = false, - filename = horizontal_average_jld2_filepath, - overwrite_existing = true) + simulation.output_writers[:horizontal_average] = JLD2OutputWriter(model, + jld2_outputs, + schedule = TimeInterval(10Δt), + dir = ".", + with_halos = false, + filename = horizontal_average_jld2_filepath, + overwrite_existing = true) multiple_time_average_jld2_filepath = "decay_windowed_time_average_test.jld2" single_time_average_jld2_filepath = "single_decay_windowed_time_average_test.jld2" @@ -242,23 +240,21 @@ function test_jld2_time_averaging(arch) single_jld2_output = Dict("c1" => ∫c1_dxdy) - simulation.output_writers[:single_output_time_average] = JLD2OutputWriter( - model, - single_jld2_output, - schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), - dir = ".", - with_halos = false, - filename = single_time_average_jld2_filepath, - overwrite_existing = true) - - simulation.output_writers[:multiple_output_time_average] = JLD2OutputWriter( - model, - jld2_outputs, - schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), - dir = ".", - with_halos = false, - filename = multiple_time_average_jld2_filepath, - overwrite_existing = true) + simulation.output_writers[:single_output_time_average] = JLD2OutputWriter(model, + single_jld2_output, + schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), + dir = ".", + with_halos = false, + filename = single_time_average_jld2_filepath, + overwrite_existing = true) + + simulation.output_writers[:multiple_output_time_average] = JLD2OutputWriter(model, + jld2_outputs, + schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), + dir = ".", + with_halos = false, + filename = multiple_time_average_jld2_filepath, + overwrite_existing = true) run!(simulation) @@ -471,4 +467,4 @@ for arch in archs ##### test_jld2_time_averaging(arch) end -end \ No newline at end of file +end diff --git a/test/test_netcdf_output_writer.jl b/test/test_netcdf_output_writer.jl index 079d7a5962..ff13867e78 100644 --- a/test/test_netcdf_output_writer.jl +++ b/test/test_netcdf_output_writer.jl @@ -468,13 +468,12 @@ function test_netcdf_function_output(arch) nc_filepath = "test_function_outputs_$(typeof(arch)).nc" - simulation.output_writers[:food] = - NetCDFOutputWriter(model, outputs; global_attributes, output_attributes, - filename = nc_filepath, - schedule = TimeInterval(Δt), - dimensions = dims, - array_type = Array{Float64}, - verbose=true) + simulation.output_writers[:food] = NetCDFOutputWriter(model, outputs; global_attributes, output_attributes, + filename = nc_filepath, + schedule = TimeInterval(Δt), + dimensions = dims, + array_type = Array{Float64}, + verbose=true) run!(simulation) @@ -679,14 +678,13 @@ function test_netcdf_time_averaging(arch) horizontal_average_nc_filepath = "decay_averaged_field_test.nc" - simulation.output_writers[:horizontal_average] = NetCDFOutputWriter( - model, - nc_outputs, - array_type = Array{Float64}, - verbose = true, - filename = horizontal_average_nc_filepath, - schedule = TimeInterval(10Δt), - dimensions = nc_dimensions) + simulation.output_writers[:horizontal_average] = NetCDFOutputWriter(model, + nc_outputs, + array_type = Array{Float64}, + verbose = true, + filename = horizontal_average_nc_filepath, + schedule = TimeInterval(10Δt), + dimensions = nc_dimensions) multiple_time_average_nc_filepath = "decay_windowed_time_average_test.nc" single_time_average_nc_filepath = "single_decay_windowed_time_average_test.nc" @@ -695,23 +693,21 @@ function test_netcdf_time_averaging(arch) single_nc_output = Dict("c1" => ∫c1_dxdy) single_nc_dimension = Dict("c1" => ("zC",)) - simulation.output_writers[:single_output_time_average] = NetCDFOutputWriter( - model, - single_nc_output, - array_type = Array{Float64}, - verbose = true, - filename = single_time_average_nc_filepath, - schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), - dimensions = single_nc_dimension) - - simulation.output_writers[:multiple_output_time_average] = NetCDFOutputWriter( - model, - nc_outputs, - array_type = Array{Float64}, - verbose = true, - filename = multiple_time_average_nc_filepath, - schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), - dimensions = nc_dimensions) + simulation.output_writers[:single_output_time_average] = NetCDFOutputWriter(model, + single_nc_output, + array_type = Array{Float64}, + verbose = true, + filename = single_time_average_nc_filepath, + schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), + dimensions = single_nc_dimension) + + simulation.output_writers[:multiple_output_time_average] = NetCDFOutputWriter(model, + nc_outputs, + array_type = Array{Float64}, + verbose = true, + filename = multiple_time_average_nc_filepath, + schedule = AveragedTimeInterval(10Δt, window = window, stride = stride), + dimensions = nc_dimensions) run!(simulation) @@ -749,8 +745,8 @@ function test_netcdf_time_averaging(arch) single_ds = NCDataset(single_time_average_nc_filepath) attribute_names = ("schedule", "interval", "output time interval", - "time_averaging_window", "time averaging window", - "time_averaging_stride", "time averaging stride") + "time_averaging_window", "time averaging window", + "time_averaging_stride", "time averaging stride") for name in attribute_names @test haskey(single_ds.attrib, name) && !isnothing(single_ds.attrib[name]) @@ -798,14 +794,12 @@ function test_netcdf_output_alignment(arch) simulation = Simulation(model, Δt=0.2, stop_time=40) test_filename1 = "test_output_alignment1.nc" - simulation.output_writers[:stuff] = - NetCDFOutputWriter(model, model.velocities, filename=test_filename1, - schedule=TimeInterval(7.3)) + simulation.output_writers[:stuff] = NetCDFOutputWriter(model, model.velocities, + filename=test_filename1, schedule=TimeInterval(7.3)) test_filename2 = "test_output_alignment2.nc" - simulation.output_writers[:something] = - NetCDFOutputWriter(model, model.tracers, filename=test_filename2, - schedule=TimeInterval(3.0)) + simulation.output_writers[:something] = NetCDFOutputWriter(model, model.tracers, + filename=test_filename2, schedule=TimeInterval(3.0)) run!(simulation) @@ -838,13 +832,11 @@ function test_netcdf_vertically_stretched_grid_output(arch) nc_filepath = "test_netcdf_vertically_stretched_grid_output_$(typeof(arch)).nc" - simulation.output_writers[:fields] = - NetCDFOutputWriter(model, merge(model.velocities, model.tracers), - filename = nc_filepath, - schedule = IterationInterval(1), - array_type = Array{Float64}, - verbose = true) - + simulation.output_writers[:fields] = NetCDFOutputWriter(model, merge(model.velocities, model.tracers), + filename = nc_filepath, + schedule = IterationInterval(1), + array_type = Array{Float64}, + verbose = true) run!(simulation) grid = model.grid From a8cf967376d9bc13351a57035649824e4bbbae77 Mon Sep 17 00:00:00 2001 From: "Gregory L. Wagner" Date: Wed, 29 Jan 2025 13:52:05 -0500 Subject: [PATCH 31/31] Fix Corsican bug in validation test (#4070) * Fix a bug in validation test * AdvectiveSkew -> AdvectiveFormulation * Fix wall clock definition * fix typo * Replace WENO5 with WINO(order=5) * an other little fix --------- Co-authored-by: Jean-Michel Campin --- .../coarse_baroclinic_adjustment.jl | 8 ++--- ... immersed_coarse_baroclinic_adjustment.jl} | 34 ++++++++----------- 2 files changed, 19 insertions(+), 23 deletions(-) rename validation/mesoscale_turbulence/{immersed_corse_baroclinic_adjustment.jl => immersed_coarse_baroclinic_adjustment.jl} (89%) diff --git a/validation/mesoscale_turbulence/coarse_baroclinic_adjustment.jl b/validation/mesoscale_turbulence/coarse_baroclinic_adjustment.jl index cd26a180dc..22e57aa4bd 100644 --- a/validation/mesoscale_turbulence/coarse_baroclinic_adjustment.jl +++ b/validation/mesoscale_turbulence/coarse_baroclinic_adjustment.jl @@ -22,7 +22,7 @@ stop_time = 30days grid = RectilinearGrid(architecture; topology = (Bounded, Bounded, Bounded), - size = (Ny, Ny, Nz), + size = (Ny, Ny, Nz), x = (-Ly/2, Ly/2), y = (-Ly/2, Ly/2), z = (-Lz, 0), @@ -49,8 +49,8 @@ model = HydrostaticFreeSurfaceModel(grid = grid, buoyancy = BuoyancyTracer(), closure = closures, tracers = (:b, :c), - momentum_advection = WENO5(), - tracer_advection = WENO5(), + momentum_advection = WENO(order=5), + tracer_advection = WENO(order=5), free_surface = ImplicitFreeSurface()) @info "Built $model." @@ -140,7 +140,7 @@ bt = FieldTimeSeries(filepath, "b") ct = FieldTimeSeries(filepath, "c") # Build coordinates, rescaling the vertical coordinate -x, y, z = nodes((Center, Center, Center), grid) +x, y, z = nodes(bt) zscale = 1 z = z .* zscale diff --git a/validation/mesoscale_turbulence/immersed_corse_baroclinic_adjustment.jl b/validation/mesoscale_turbulence/immersed_coarse_baroclinic_adjustment.jl similarity index 89% rename from validation/mesoscale_turbulence/immersed_corse_baroclinic_adjustment.jl rename to validation/mesoscale_turbulence/immersed_coarse_baroclinic_adjustment.jl index d923d376a5..f94cfb6c43 100644 --- a/validation/mesoscale_turbulence/immersed_corse_baroclinic_adjustment.jl +++ b/validation/mesoscale_turbulence/immersed_coarse_baroclinic_adjustment.jl @@ -4,10 +4,11 @@ using Random using Oceananigans using Oceananigans.Units using GLMakie -using Oceananigans.TurbulenceClosures: IsopycnalSkewSymmetricDiffusivity -using Oceananigans.TurbulenceClosures: FluxTapering, DiffusiveFormulation, AdvectiveSkewClosure +using Oceananigans.TurbulenceClosures: IsopycnalSkewSymmetricDiffusivity, SkewAdvectionISSD +using Oceananigans.TurbulenceClosures: FluxTapering, DiffusiveFormulation, AdvectiveFormulation filename = "coarse_baroclinic_adjustment" +wall_clock = Ref(time_ns()) function progress(sim) @printf("[%05.2f%%] i: %d, t: %s, wall time: %s, max(u): (%6.3e, %6.3e, %6.3e) m/s, next Δt: %s\n", @@ -21,7 +22,7 @@ function progress(sim) prettytime(sim.Δt)) wall_clock[] = time_ns() - + return nothing end @@ -37,7 +38,7 @@ save_fields_interval = 2day stop_time = 300days grid = RectilinearGrid(topology = (Flat, Bounded, Bounded), - size = (Ny, Nz), + size = (Ny, Nz), y = (-Ly/2, Ly/2), z = (-Lz, 0), halo = (5, 5)) @@ -50,14 +51,13 @@ grid = ImmersedBoundaryGrid(grid, GridFittedBottom(y -> y > 0 ? -Lz : -Lz/2)) slope_limiter = FluxTapering(1000) # Allow very steep slopes adv_closure = IsopycnalSkewSymmetricDiffusivity(; κ_skew, slope_limiter) -dif_closure = IsopycnalSkewSymmetricDiffusivity(; κ_skew, slope_limiter, skew_fluxes_formulation = DiffusiveFormulation()) +dif_closure = IsopycnalSkewSymmetricDiffusivity(; κ_skew, slope_limiter, skew_flux_formulation = DiffusiveFormulation()) function run_simulation(closure, grid) - model = HydrostaticFreeSurfaceModel(; grid, + model = HydrostaticFreeSurfaceModel(; grid, closure, coriolis = FPlane(latitude = -45), buoyancy = BuoyancyTracer(), tracer_advection = WENO(order=7), - closure = adv_closure, tracers = (:b, :c)) @info "Built $model." @@ -85,14 +85,10 @@ function run_simulation(closure, grid) ##### simulation = Simulation(model; Δt, stop_time) + add_callback!(simulation, progress, IterationInterval(144)) + suffix = closure isa SkewAdvectionISSD ? "advective" : "diffusive" - wall_clock = Ref(time_ns()) - - add_callback!(simulation, progress, IterationInterval(10)) - - suffix = closure isa AdvectiveSkewClosure ? "advective" : "diffusive" - - simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers), + simulation.output_writers[:fields] = JLD2OutputWriter(model, merge(model.velocities, model.tracers), schedule = TimeInterval(save_fields_interval), filename = filename * "_fields_" * suffix, overwrite_existing = true) @@ -145,10 +141,10 @@ cd = @lift ctd[$n] fig = Figure(size=(1800, 700)) -axua = Axis(fig[2, 1], xlabel="y (km)", ylabel="z (km)", title="Zonal velocity") -axca = Axis(fig[3, 1], xlabel="y (km)", ylabel="z (km)", title="Tracer concentration") -axud = Axis(fig[2, 2], xlabel="y (km)", ylabel="z (km)", title="Zonal velocity") -axcd = Axis(fig[3, 2], xlabel="y (km)", ylabel="z (km)", title="Tracer concentration") +axua = Axis(fig[2, 1], xlabel="y (km)", ylabel="z (km)", title="GM_Adv: Zonal velocity") +axca = Axis(fig[3, 1], xlabel="y (km)", ylabel="z (km)", title="GM_Adv: Tracer concentration") +axud = Axis(fig[2, 2], xlabel="y (km)", ylabel="z (km)", title="GM_Dif: Zonal velocity") +axcd = Axis(fig[3, 2], xlabel="y (km)", ylabel="z (km)", title="GM_Dif: Tracer concentration") levels = [-0.0015 + 0.0005 * i for i in 0:19] @@ -174,4 +170,4 @@ display(fig) record(fig, filename * ".mp4", 1:Nt, framerate=8) do i @info "Plotting frame $i of $Nt" n[] = i -end \ No newline at end of file +end