diff --git a/ext/cuda/limiters.jl b/ext/cuda/limiters.jl index bf2c80e901..7511c279f4 100644 --- a/ext/cuda/limiters.jl +++ b/ext/cuda/limiters.jl @@ -5,6 +5,7 @@ import ClimaCore.Limiters: apply_limiter! import ClimaCore.Fields import ClimaCore: DataLayouts, Spaces, Topologies, Fields +import ClimaCore.DataLayouts: slab_index using CUDA function config_threadblock(Nv, Nh) @@ -63,7 +64,7 @@ function compute_element_bounds_kernel!( slab_ρ = slab(ρ, v, h) for j in 1:Nj for i in 1:Ni - q = rdiv(slab_ρq[i, j], slab_ρ[i, j]) + q = rdiv(slab_ρq[slab_index(i, j)], slab_ρ[slab_index(i, j)]) if i == 1 && j == 1 q_min = q q_max = q @@ -74,8 +75,8 @@ function compute_element_bounds_kernel!( end end slab_q_bounds = slab(q_bounds, v, h) - slab_q_bounds[1] = q_min - slab_q_bounds[2] = q_max + slab_q_bounds[slab_index(1)] = q_min + slab_q_bounds[slab_index(2)] = q_max end return nothing end @@ -123,18 +124,18 @@ function compute_neighbor_bounds_local_kernel!( (v, h) = kernel_indexes(tidx, n).I (; q_bounds, q_bounds_nbr, ghost_buffer, rtol) = limiter slab_q_bounds = slab(q_bounds, v, h) - q_min = slab_q_bounds[1] - q_max = slab_q_bounds[2] + q_min = slab_q_bounds[slab_index(1)] + q_max = slab_q_bounds[slab_index(2)] for lne in local_neighbor_elem_offset[h]:(local_neighbor_elem_offset[h + 1] - 1) h_nbr = local_neighbor_elem[lne] slab_q_bounds = slab(q_bounds, v, h_nbr) - q_min = rmin(q_min, slab_q_bounds[1]) - q_max = rmax(q_max, slab_q_bounds[2]) + q_min = rmin(q_min, slab_q_bounds[slab_index(1)]) + q_max = rmax(q_max, slab_q_bounds[slab_index(2)]) end slab_q_bounds_nbr = slab(q_bounds_nbr, v, h) - slab_q_bounds_nbr[1] = q_min - slab_q_bounds_nbr[2] = q_max + slab_q_bounds_nbr[slab_index(1)] = q_min + slab_q_bounds_nbr[slab_index(2)] = q_max end return nothing end diff --git a/ext/cuda/matrix_fields_single_field_solve.jl b/ext/cuda/matrix_fields_single_field_solve.jl index f746ca4b9d..60dddb6f40 100644 --- a/ext/cuda/matrix_fields_single_field_solve.jl +++ b/ext/cuda/matrix_fields_single_field_solve.jl @@ -7,6 +7,7 @@ import ClimaCore.Fields import ClimaCore.Spaces import ClimaCore.Topologies import ClimaCore.MatrixFields +import ClimaCore.DataLayouts: vindex import ClimaCore.MatrixFields: single_field_solve! import ClimaCore.MatrixFields: _single_field_solve! import ClimaCore.MatrixFields: band_matrix_solve!, unzip_tuple_field_values @@ -71,7 +72,7 @@ function _single_field_solve!( b_data = Fields.field_values(b) Nv = DataLayouts.nlevels(x_data) @inbounds for v in 1:Nv - x_data[v] = inv(A.λ) ⊠ b_data[v] + x_data[vindex(v)] = inv(A.λ) ⊠ b_data[vindex(v)] end end @@ -98,6 +99,7 @@ function band_matrix_solve_local_mem!( Nv = DataLayouts.nlevels(x) Ux, U₊₁ = cache A₋₁, A₀, A₊₁ = Aⱼs + vi = vindex Ux_local = MArray{Tuple{Nv}, eltype(Ux)}(undef) U₊₁_local = MArray{Tuple{Nv}, eltype(U₊₁)}(undef) @@ -107,16 +109,16 @@ function band_matrix_solve_local_mem!( A₊₁_local = MArray{Tuple{Nv}, eltype(A₊₁)}(undef) b_local = MArray{Tuple{Nv}, eltype(b)}(undef) @inbounds for v in 1:Nv - A₋₁_local[v] = A₋₁[v] - A₀_local[v] = A₀[v] - A₊₁_local[v] = A₊₁[v] - b_local[v] = b[v] + A₋₁_local[v] = A₋₁[vi(v)] + A₀_local[v] = A₀[vi(v)] + A₊₁_local[v] = A₊₁[vi(v)] + b_local[v] = b[vi(v)] end cache_local = (Ux_local, U₊₁_local) Aⱼs_local = (A₋₁, A₀, A₊₁) - band_matrix_solve!(t, cache_local, x_local, Aⱼs_local, b_local) + band_matrix_solve!(t, cache_local, x_local, Aⱼs_local, b_local, vi) @inbounds for v in 1:Nv - x[v] = x_local[v] + x[vi(v)] = x_local[v] end return nothing end @@ -128,6 +130,7 @@ function band_matrix_solve_local_mem!( Aⱼs, b, ) + vi = vindex Nv = DataLayouts.nlevels(x) Ux, U₊₁, U₊₂ = cache A₋₂, A₋₁, A₀, A₊₁, A₊₂ = Aⱼs @@ -142,18 +145,18 @@ function band_matrix_solve_local_mem!( A₊₂_local = MArray{Tuple{Nv}, eltype(A₊₂)}(undef) b_local = MArray{Tuple{Nv}, eltype(b)}(undef) @inbounds for v in 1:Nv - A₋₂_local[v] = A₋₂[v] - A₋₁_local[v] = A₋₁[v] - A₀_local[v] = A₀[v] - A₊₁_local[v] = A₊₁[v] - A₊₂_local[v] = A₊₂[v] - b_local[v] = b[v] + A₋₂_local[v] = A₋₂[vi(v)] + A₋₁_local[v] = A₋₁[vi(v)] + A₀_local[v] = A₀[vi(v)] + A₊₁_local[v] = A₊₁[vi(v)] + A₊₂_local[v] = A₊₂[vi(v)] + b_local[v] = b[vi(v)] end cache_local = (Ux_local, U₊₁_local, U₊₂_local) Aⱼs_local = (A₋₂, A₋₁, A₀, A₊₁, A₊₂) - band_matrix_solve!(t, cache_local, x_local, Aⱼs_local, b_local) + band_matrix_solve!(t, cache_local, x_local, Aⱼs_local, b_local, vi) @inbounds for v in 1:Nv - x[v] = x_local[v] + x[vi(v)] = x_local[v] end return nothing end @@ -168,7 +171,7 @@ function band_matrix_solve_local_mem!( Nv = DataLayouts.nlevels(x) (A₀,) = Aⱼs @inbounds for v in 1:Nv - x[v] = inv(A₀[v]) ⊠ b[v] + x[vindex(v)] = inv(A₀[vindex(v)]) ⊠ b[vindex(v)] end return nothing end diff --git a/ext/cuda/remapping_distributed.jl b/ext/cuda/remapping_distributed.jl index 5dcf5532b3..f5fc20183b 100644 --- a/ext/cuda/remapping_distributed.jl +++ b/ext/cuda/remapping_distributed.jl @@ -59,7 +59,7 @@ function set_interpolated_values_kernel!( totalThreadsZ = gridDim().z * blockDim().z _, Nq = size(I1) - + CI = CartesianIndex for i in hindex:totalThreadsX:num_horiz h = local_horiz_indices[i] for j in vindex:totalThreadsY:num_vert @@ -73,8 +73,8 @@ function set_interpolated_values_kernel!( I1[i, t] * I2[i, s] * ( - A * field_values[k][t, s, nothing, v_lo, h] + - B * field_values[k][t, s, nothing, v_hi, h] + A * field_values[k][CI(t, s, 1, v_lo, h)] + + B * field_values[k][CI(t, s, 1, v_hi, h)] ) end end @@ -107,7 +107,7 @@ function set_interpolated_values_kernel!( totalThreadsZ = gridDim().z * blockDim().z _, Nq = size(I) - + CI = CartesianIndex for i in hindex:totalThreadsX:num_horiz h = local_horiz_indices[i] for j in vindex:totalThreadsY:num_vert @@ -121,10 +121,8 @@ function set_interpolated_values_kernel!( I[i, t] * I[i, s] * ( - A * - field_values[k][t, nothing, nothing, v_lo, h] + - B * - field_values[k][t, nothing, nothing, v_hi, h] + A * field_values[k][CI(t, 1, 1, v_lo, h)] + + B * field_values[k][CI(t, 1, 1, v_hi, h)] ) end end @@ -199,7 +197,7 @@ function set_interpolated_values_kernel!( out[i, k] += I1[i, t] * I2[i, s] * - field_values[k][t, s, nothing, nothing, h] + field_values[k][CartesianIndex(t, s, 1, 1, h)] end end end @@ -232,8 +230,7 @@ function set_interpolated_values_kernel!( out[i, k] = 0 for t in 1:Nq, s in 1:Nq out[i, k] += - I[i, i] * - field_values[k][t, nothing, nothing, nothing, h] + I[i, i] * field_values[k][CartesianIndex(t, 1, 1, 1, h)] end end end diff --git a/lib/ClimaCorePlots/src/ClimaCorePlots.jl b/lib/ClimaCorePlots/src/ClimaCorePlots.jl index 35580b7007..a86b4544c8 100644 --- a/lib/ClimaCorePlots/src/ClimaCorePlots.jl +++ b/lib/ClimaCorePlots/src/ClimaCorePlots.jl @@ -4,6 +4,7 @@ import RecipesBase import TriplotBase import ClimaComms +import ClimaCore.DataLayouts: slab_index import ClimaCore: ClimaCore, DataLayouts, @@ -308,7 +309,7 @@ function _slice_along(field, coord) hdata = ClimaCore.slab(hcoord_data, hidx) hnode_idx = 1 for i in axes(hdata)[axis] - pt = axis == 1 ? hdata[i, 1] : hdata[1, i] + pt = axis == 1 ? hdata[slab_index(i, 1)] : hdata[slab_index(1, i)] axis_value = Geometry.component(pt, axis) coord_value = Geometry.component(coord, 1) if axis_value > coord_value @@ -353,8 +354,9 @@ function _slice_along(field, coord) islab = ClimaCore.slab(ortho_data, v, i) # copy the nodal data for ni in 1:size(islab)[1] - islab[ni] = - axis == 1 ? ijslab[hnode_idx, ni] : ijslab[ni, hnode_idx] + islab[slab_index(ni)] = + axis == 1 ? ijslab[slab_index(hnode_idx, ni)] : + ijslab[slab_index(ni, hnode_idx)] end end end diff --git a/lib/ClimaCoreTempestRemap/src/netcdf.jl b/lib/ClimaCoreTempestRemap/src/netcdf.jl index eb15c7e4bb..24fd815c51 100644 --- a/lib/ClimaCoreTempestRemap/src/netcdf.jl +++ b/lib/ClimaCoreTempestRemap/src/netcdf.jl @@ -1,5 +1,6 @@ import CommonDataModel import ClimaCore: slab, column +import ClimaCore.DataLayouts: slab_index """ def_time_coord(nc::NCDataset, length=Inf, eltype=Float64; @@ -97,7 +98,7 @@ function def_space_coord( coords = Spaces.coordinates_data(space) for (col, ((i, j), e)) in enumerate(nodes) - coord = slab(coords, e)[i, j] + coord = slab(coords, e)[slab_index(i, j)] X[col] = coord.x Y[col] = coord.y end @@ -149,7 +150,7 @@ function def_space_coord( coords = Spaces.coordinates_data(space) for (col, ((i, j), e)) in enumerate(nodes) - coord = slab(coords, e)[i, j] + coord = slab(coords, e)[slab_index(i, j)] lon[col] = coord.long lat[col] = coord.lat end @@ -328,7 +329,7 @@ function Base.setindex!( end data = Fields.field_values(field) for (col, ((i, j), e)) in enumerate(nodes) - var[col, extraidx...] = slab(data, e)[i, j] + var[col, extraidx...] = slab(data, e)[slab_index(i, j)] end return var end diff --git a/src/DataLayouts/DataLayouts.jl b/src/DataLayouts/DataLayouts.jl index c8f0cb0b16..2d65ef6c9a 100644 --- a/src/DataLayouts/DataLayouts.jl +++ b/src/DataLayouts/DataLayouts.jl @@ -387,32 +387,23 @@ end data::IJFH{S}, I::CartesianIndex{5}, ) where {S} - i, j, _, _, h = I.I - @inbounds get_struct(parent(data), S, Val(3), CartesianIndex(i, j, 1, h)) + @inbounds get_struct( + parent(data), + S, + Val(field_dim(data)), + to_data_specific(data, I), + ) end @propagate_inbounds function Base.setindex!( data::IJFH{S}, val, I::CartesianIndex{5}, ) where {S} - i, j, _, _, h = I.I - @inbounds set_struct!( - parent(data), - convert(S, val), - Val(3), - CartesianIndex(i, j, 1, h), - ) -end - -@inline function Base.getindex(data::IJFH{S}, i, j, _, _, h) where {S} - @inbounds get_struct(parent(data), S, Val(3), CartesianIndex(i, j, 1, h)) -end -@inline function Base.setindex!(data::IJFH{S}, val, i, j, _, _, h) where {S} @inbounds set_struct!( parent(data), convert(S, val), - Val(3), - CartesianIndex(i, j, 1, h), + Val(field_dim(data)), + to_data_specific(data, I), ) end @@ -579,19 +570,12 @@ end IFH{SS, Ni, Nh}(dataview) end -@inline function Base.getindex(data::IFH{S}, i, _, _, _, h) where {S} - @inbounds get_struct(parent(data), S, Val(2), CartesianIndex(i, 1, h)) -end @inline function Base.getindex(data::IFH{S}, I::CartesianIndex{5}) where {S} - i, _, _, _, h = I.I - @inbounds get_struct(parent(data), S, Val(2), CartesianIndex(i, 1, h)) -end -@inline function Base.setindex!(data::IFH{S}, val, i, _, _, _, h) where {S} - @inbounds set_struct!( + @inbounds get_struct( parent(data), - convert(S, val), - Val(2), - CartesianIndex(i, 1, h), + S, + Val(field_dim(data)), + to_data_specific(data, I), ) end @inline function Base.setindex!( @@ -599,12 +583,11 @@ end val, I::CartesianIndex{5}, ) where {S} - i, _, _, _, h = I.I @inbounds set_struct!( parent(data), convert(S, val), - Val(2), - CartesianIndex(i, 1, h), + Val(field_dim(data)), + to_data_specific(data, I), ) end @@ -677,7 +660,12 @@ end end Base.@propagate_inbounds function Base.getindex(data::DataF{S}) where {S} - @inbounds get_struct(parent(data), S, Val(1), CartesianIndex(1)) + @inbounds get_struct( + parent(data), + S, + Val(field_dim(data)), + CartesianIndex(1), + ) end @propagate_inbounds function Base.getindex(col::Data0D, I::CartesianIndex{5}) @@ -688,7 +676,7 @@ Base.@propagate_inbounds function Base.setindex!(data::DataF{S}, val) where {S} @inbounds set_struct!( parent(data), convert(S, val), - Val(1), + Val(field_dim(data)), CartesianIndex(1), ) end @@ -707,21 +695,6 @@ Base.copy(data::DataF{S}) where {S} = DataF{S}(copy(parent(data))) # DataSlab2D DataLayout # ====================== -@propagate_inbounds function Base.getindex( - slab::DataSlab2D{S}, - I::CartesianIndex, -) where {S} - @inbounds slab[I[1], I[2]] -end - -@propagate_inbounds function Base.setindex!( - slab::DataSlab2D{S}, - val, - I::CartesianIndex, -) where {S} - @inbounds slab[I[1], I[2]] = val -end - Base.size(::DataSlab2D{S, Nij}) where {S, Nij} = (Nij, Nij, 1, 1, 1) Base.axes(::DataSlab2D{S, Nij}) where {S, Nij} = (SOneTo(Nij), SOneTo(Nij)) @@ -805,29 +778,33 @@ end @inline function Base.getindex( data::IJF{S, Nij}, - i::Integer, - j::Integer, - k = nothing, - v = nothing, - h = nothing, + I::CartesianIndex, ) where {S, Nij} + i = I.I[1] + j = I.I[2] @boundscheck (1 <= i <= Nij && 1 <= j <= Nij) || throw(BoundsError(data, (i, j))) - @inbounds get_struct(parent(data), S, Val(3), CartesianIndex(i, j, 1)) + @inbounds get_struct( + parent(data), + S, + Val(field_dim(data)), + CartesianIndex(i, j, 1), + ) end @inline function Base.setindex!( data::IJF{S, Nij}, val, - i::Integer, - j::Integer, + I::CartesianIndex, ) where {S, Nij} + i = I.I[1] + j = I.I[2] @boundscheck (1 <= i <= Nij && 1 <= j <= Nij) || throw(BoundsError(data, (i, j))) @inbounds set_struct!( parent(data), convert(S, val), - Val(3), + Val(field_dim(data)), CartesianIndex(i, j, 1), ) end @@ -843,18 +820,6 @@ end # DataSlab1D DataLayout # ====================== -@propagate_inbounds function Base.getindex(slab::DataSlab1D, I::CartesianIndex) - @inbounds slab[I[1]] -end - -@propagate_inbounds function Base.setindex!( - slab::DataSlab1D, - val, - I::CartesianIndex, -) - @inbounds slab[I[1]] = val -end - function Base.size(::DataSlab1D{<:Any, Ni}) where {Ni} return (Ni, 1, 1, 1, 1) end @@ -932,24 +897,28 @@ end IF{SS, Ni}(dataview) end -@inline function Base.getindex( - data::IF{S, Ni}, - i::Integer, - j = nothing, - k = nothing, - v = nothing, - h = nothing, -) where {S, Ni} +@inline function Base.getindex(data::IF{S, Ni}, I::CartesianIndex) where {S, Ni} + i = I.I[1] @boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,))) - @inbounds get_struct(parent(data), S, Val(2), CartesianIndex(i, 1)) + @inbounds get_struct( + parent(data), + S, + Val(field_dim(data)), + CartesianIndex(i, 1), + ) end -@inline function Base.setindex!(data::IF{S, Ni}, val, i::Integer) where {S, Ni} +@inline function Base.setindex!( + data::IF{S, Ni}, + val, + I::CartesianIndex, +) where {S, Ni} + i = I.I[1] @boundscheck (1 <= i <= Ni) || throw(BoundsError(data, (i,))) @inbounds set_struct!( parent(data), convert(S, val), - Val(2), + Val(field_dim(data)), CartesianIndex(i, 1), ) end @@ -1031,32 +1000,24 @@ end Base.@propagate_inbounds Base.getproperty(data::VF, i::Integer) = _property_view(data, Val(i)) -@inline function Base.getindex(data::VF{S, Nv}, v::Integer) where {S, Nv} +@inline function Base.getindex(data::VF{S, Nv}, I::CartesianIndex) where {S, Nv} + v = I.I[4] @boundscheck 1 <= v <= nlevels(data) || throw(BoundsError(data, (v,))) - @inbounds get_struct(parent(data), S, Val(2), CartesianIndex(v, 1)) -end - -@propagate_inbounds function Base.getindex( - col::DataColumn, - I::CartesianIndex{5}, -) - @inbounds col[I[4]] -end - -@propagate_inbounds function Base.setindex!( - col::DataColumn, - val, - I::CartesianIndex{5}, -) - @inbounds col[I[4]] = val + @inbounds get_struct( + parent(data), + S, + Val(field_dim(data)), + CartesianIndex(v, 1), + ) end -@inline function Base.setindex!(data::VF{S}, val, v::Integer) where {S} +@inline function Base.setindex!(data::VF{S}, val, I::CartesianIndex) where {S} + v = I.I[4] @boundscheck (1 <= v <= nlevels(data)) || throw(BoundsError(data, (v,))) @inbounds set_struct!( parent(data), convert(S, val), - Val(2), + Val(field_dim(data)), CartesianIndex(v, 1), ) end @@ -1206,8 +1167,12 @@ end data::VIJFH{S}, I::CartesianIndex{5}, ) where {S} - i, j, _, v, h = I.I - @inbounds get_struct(parent(data), S, Val(4), CartesianIndex(v, i, j, 1, h)) + @inbounds get_struct( + parent(data), + S, + Val(field_dim(data)), + to_data_specific(data, I), + ) end @propagate_inbounds function Base.setindex!( @@ -1215,12 +1180,11 @@ end val, I::CartesianIndex{5}, ) where {S} - i, j, _, v, h = I.I @inbounds set_struct!( parent(data), convert(S, val), - Val(4), - CartesianIndex(v, i, j, 1, h), + Val(field_dim(data)), + to_data_specific(data, I), ) end @@ -1237,18 +1201,6 @@ function gather( end end -@inline function Base.getindex(data::VIJFH{S}, i, j, _, v, h) where {S} - @inbounds get_struct(parent(data), S, Val(4), CartesianIndex(v, i, j, 1, h)) -end -@inline function Base.setindex!(data::VIJFH{S}, val, i, j, _, v, h) where {S} - @inbounds set_struct!( - parent(data), - convert(S, val), - Val(4), - CartesianIndex(v, i, j, 1, h), - ) -end - # ====================== # Data1DX DataLayout # ====================== @@ -1300,15 +1252,6 @@ Base.size(data::VIFH{<:Any, Nv, Ni, Nh}) where {Nv, Ni, Nh} = (Ni, 1, 1, Nv, Nh) Base.length(data::VIFH) = nlevels(data) * size(parent(data), 4) -@propagate_inbounds function Base.getindex( - data::VIFH{S}, - I::CartesianIndex{5}, -) where {S} - i, _, _, v, h = I.I - @inbounds get_struct(parent(data), S, Val(3), CartesianIndex(v, i, 1, h)) -end - - @generated function _property_view( data::VIFH{S, Nv, Ni, Nh, A}, ::Val{Idx}, @@ -1386,17 +1329,19 @@ end IFH{S, Nij, Nh}(dataview) end -@inline function Base.getindex(data::VIFH{S}, i, _, _, v, h) where {S} - @inbounds get_struct(parent(data), S, Val(3), CartesianIndex(v, i, 1, h)) -end -@inline function Base.setindex!(data::VIFH{S}, val, i, _, _, v, h) where {S} - @inbounds set_struct!( +@propagate_inbounds function Base.getindex( + data::VIFH{S}, + I::CartesianIndex{5}, +) where {S} + i, _, _, v, h = I.I + @inbounds get_struct( parent(data), - convert(S, val), - Val(3), + S, + Val(field_dim(data)), CartesianIndex(v, i, 1, h), ) end + @inline function Base.setindex!( data::VIFH{S}, val, @@ -1406,7 +1351,7 @@ end @inbounds set_struct!( parent(data), convert(S, val), - Val(3), + Val(field_dim(data)), CartesianIndex(v, i, 1, h), ) end @@ -1568,6 +1513,24 @@ get_Nij(::IFH{S, Nij}) where {S, Nij} = Nij get_Nij(::IJF{S, Nij}) where {S, Nij} = Nij get_Nij(::IF{S, Nij}) where {S, Nij} = Nij +@inline field_dim(::IJKFVH) = 4 +@inline field_dim(::IJFH) = 3 +@inline field_dim(::IFH) = 2 +@inline field_dim(::DataF) = 1 +@inline field_dim(::IJF) = 3 +@inline field_dim(::IF) = 2 +@inline field_dim(::VF) = 2 +@inline field_dim(::VIJFH) = 4 +@inline field_dim(::VIFH) = 3 + +#! format: off +@inline to_data_specific(::IJFH, I::CartesianIndex) = CartesianIndex(I.I[1], I.I[2], 1, I.I[5]) +@inline to_data_specific(::IFH, I::CartesianIndex) = CartesianIndex(I.I[1], 1, I.I[5]) +@inline to_data_specific(::VIJFH, I::CartesianIndex) = CartesianIndex(I.I[4], I.I[1], I.I[2], 1, I.I[5]) +@inline to_data_specific(::VIFH, I::CartesianIndex) = CartesianIndex(I.I[4]I.I[1], 1, I.I[5]) +@inline to_data_specific(::DataSlab1D, I::CartesianIndex) = CartesianIndex(I.I[1]I.I[1], 1, I.I[5]) +#! format: on + Base.ndims(data::AbstractData) = Base.ndims(typeof(data)) Base.ndims(::Type{T}) where {T <: AbstractData} = Base.ndims(parent_array_type(T)) @@ -1642,4 +1605,8 @@ include("fused_copyto.jl") include("fill.jl") include("mapreduce.jl") +slab_index(i, j) = CartesianIndex(i, j, 1, 1, 1) +slab_index(i) = CartesianIndex(i, 1, 1, 1, 1) +vindex(v) = CartesianIndex(1, 1, 1, v, 1) + end # module diff --git a/src/DataLayouts/fill.jl b/src/DataLayouts/fill.jl index 1e3105810e..c942b0c959 100644 --- a/src/DataLayouts/fill.jl +++ b/src/DataLayouts/fill.jl @@ -21,14 +21,14 @@ end function Base.fill!(data::IJF{S, Nij}, val, ::ToCPU) where {S, Nij} @inbounds for j in 1:Nij, i in 1:Nij - data[i, j] = val + data[CartesianIndex(i, j, 1, 1, 1)] = val end return data end function Base.fill!(data::IF{S, Ni}, val, ::ToCPU) where {S, Ni} @inbounds for i in 1:Ni - data[i] = val + data[CartesianIndex(i, 1, 1, 1, 1)] = val end return data end @@ -36,7 +36,7 @@ end function Base.fill!(data::VF, val, ::ToCPU) Nv = nlevels(data) @inbounds for v in 1:Nv - data[v] = val + data[CartesianIndex(1, 1, 1, v, 1)] = val end return data end diff --git a/src/Fields/broadcast.jl b/src/Fields/broadcast.jl index b7ce970794..eadf4913cf 100644 --- a/src/Fields/broadcast.jl +++ b/src/Fields/broadcast.jl @@ -46,7 +46,7 @@ function _first end # we're trying to do is throw a more helpful # error message. So, let's throw it here instead. _first(bc, ::Any) = throw(BroadcastInferenceError(bc)) -_first_data_layout(data::DataLayouts.VF) = data[1] +_first_data_layout(data::DataLayouts.VF) = data[CartesianIndex(1, 1, 1, 1, 1)] _first_data_layout(data::DataLayouts.DataF) = data[] _first(bc, x::Real) = x _first(bc, x::Geometry.LocalGeometry) = x diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 528fc62073..39f0f235d2 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -2,6 +2,7 @@ module Grids import ClimaComms, Adapt, ForwardDiff, LinearAlgebra import LinearAlgebra: det, norm +import ..DataLayouts: slab_index, vindex import ..DataLayouts, ..Domains, ..Meshes, ..Topologies, ..Geometry, ..Quadratures import ..Utilities: PlusHalf, half, Cache diff --git a/src/Grids/finitedifference.jl b/src/Grids/finitedifference.jl index aa7f739551..8f594ec417 100644 --- a/src/Grids/finitedifference.jl +++ b/src/Grids/finitedifference.jl @@ -59,7 +59,7 @@ function _FiniteDifferenceGrid(topology::Topologies.IntervalTopology) # construct on CPU, adapt to GPU face_coordinates = DataLayouts.VF{CT, Nv_face}(Array{FT}, Nv_face) for v in 1:Nv_face - face_coordinates[v] = mesh.faces[v] + face_coordinates[vindex(v)] = mesh.faces[v] end center_local_geometry, face_local_geometry = fd_geometry_data(face_coordinates, Val(Topologies.isperiodic(topology))) diff --git a/src/Grids/spectralelement.jl b/src/Grids/spectralelement.jl index 87434aecc5..3913c52913 100644 --- a/src/Grids/spectralelement.jl +++ b/src/Grids/spectralelement.jl @@ -70,7 +70,7 @@ function _SpectralElementGrid1D( ) / 2 J = abs(∂x∂ξ) WJ = J * quad_weights[i] - local_geometry_slab[i] = Geometry.LocalGeometry( + local_geometry_slab[slab_index(i)] = Geometry.LocalGeometry( x, J, WJ, @@ -266,7 +266,7 @@ function _SpectralElementGrid2D( WJ = J * quad_weights[i] * quad_weights[j] elem_area += WJ if !enable_bubble - local_geometry_slab[i, j] = + local_geometry_slab[slab_index(i, j)] = Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) end end @@ -285,7 +285,7 @@ function _SpectralElementGrid2D( ) J = det(Geometry.components(∂u∂ξ)) WJ = J * quad_weights[i] * quad_weights[j] - local_geometry_slab[i, j] = + local_geometry_slab[slab_index(i, j)] = Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) end else @@ -315,7 +315,7 @@ function _SpectralElementGrid2D( J = det(Geometry.components(∂u∂ξ)) J += Δarea / Nq^2 WJ = J * quad_weights[i] * quad_weights[j] - local_geometry_slab[i, j] = + local_geometry_slab[slab_index(i, j)] = Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) end else # Higher-order elements: Use HOMME bubble correction for the interior nodes @@ -356,7 +356,7 @@ function _SpectralElementGrid2D( end WJ = J * quad_weights[i] * quad_weights[j] # Finally allocate local geometry - local_geometry_slab[i, j] = + local_geometry_slab[slab_index(i, j)] = Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) end end @@ -408,7 +408,7 @@ function _SpectralElementGrid2D( @assert sgeom⁻.sWJ ≈ sgeom⁺.sWJ @assert sgeom⁻.normal ≈ -sgeom⁺.normal - internal_surface_geometry_slab[q] = sgeom⁻ + internal_surface_geometry_slab[slab_index(q)] = sgeom⁻ end end internal_surface_geometry = @@ -425,7 +425,7 @@ function _SpectralElementGrid2D( slab(boundary_surface_geometry, iface) local_geometry_slab = slab(local_geometry, elem) for q in 1:Nq - boundary_surface_geometry_slab[q] = + boundary_surface_geometry_slab[slab_index(q)] = compute_surface_geometry( local_geometry_slab, quad_weights, @@ -501,7 +501,7 @@ function compute_surface_geometry( @assert size(local_geometry_slab) == (Nq, Nq, 1, 1, 1) i, j = Topologies.face_node_index(face, Nq, q, reversed) - local_geometry = local_geometry_slab[i, j] + local_geometry = local_geometry_slab[slab_index(i, j)] (; J, ∂ξ∂x) = local_geometry # surface mass matrix diff --git a/src/Limiters/quasimonotone.jl b/src/Limiters/quasimonotone.jl index 24fcc32060..0fa200b171 100644 --- a/src/Limiters/quasimonotone.jl +++ b/src/Limiters/quasimonotone.jl @@ -1,6 +1,7 @@ import ClimaComms import ..Operators import ..RecursiveApply: ⊠, ⊞, ⊟, rmap, rzero, rdiv +import ..DataLayouts: slab_index import Adapt """ @@ -111,7 +112,10 @@ function compute_element_bounds!( local q_min, q_max for j in 1:Nj for i in 1:Ni - q = rdiv(slab_ρq[i, j], slab_ρ[i, j]) + q = rdiv( + slab_ρq[slab_index(i, j)], + slab_ρ[slab_index(i, j)], + ) if i == 1 && j == 1 q_min = q q_max = q @@ -122,8 +126,8 @@ function compute_element_bounds!( end end slab_q_bounds = slab(q_bounds, v, h) - slab_q_bounds[1] = q_min - slab_q_bounds[2] = q_max + slab_q_bounds[slab_index(1)] = q_min + slab_q_bounds[slab_index(2)] = q_max end end return nothing @@ -152,16 +156,16 @@ function compute_neighbor_bounds_local!( for h in 1:Nh for v in 1:Nv slab_q_bounds = slab(q_bounds, v, h) - q_min = slab_q_bounds[1] - q_max = slab_q_bounds[2] + q_min = slab_q_bounds[slab_index(1)] + q_max = slab_q_bounds[slab_index(2)] for h_nbr in Topologies.local_neighboring_elements(topology, h) slab_q_bounds = slab(q_bounds, v, h_nbr) - q_min = rmin(q_min, slab_q_bounds[1]) - q_max = rmax(q_max, slab_q_bounds[2]) + q_min = rmin(q_min, slab_q_bounds[slab_index(1)]) + q_max = rmax(q_max, slab_q_bounds[slab_index(2)]) end slab_q_bounds_nbr = slab(q_bounds_nbr, v, h) - slab_q_bounds_nbr[1] = q_min - slab_q_bounds_nbr[2] = q_max + slab_q_bounds_nbr[slab_index(1)] = q_min + slab_q_bounds_nbr[slab_index(2)] = q_max end end return nothing @@ -187,16 +191,16 @@ function compute_neighbor_bounds_ghost!( for h in 1:Nh for v in 1:Nv slab_q_bounds = slab(q_bounds_nbr, v, h) - q_min = slab_q_bounds[1] - q_max = slab_q_bounds[2] + q_min = slab_q_bounds[slab_index(1)] + q_max = slab_q_bounds[slab_index(2)] for gidx in Topologies.ghost_neighboring_elements(topology, h) ghost_slab_q_bounds = slab(q_bounds_ghost, v, gidx) - q_min = rmin(q_min, ghost_slab_q_bounds[1]) - q_max = rmax(q_max, ghost_slab_q_bounds[2]) + q_min = rmin(q_min, ghost_slab_q_bounds[slab_index(1)]) + q_max = rmax(q_max, ghost_slab_q_bounds[slab_index(2)]) end slab_q_bounds_nbr = slab(q_bounds_nbr, v, h) - slab_q_bounds_nbr[1] = q_min - slab_q_bounds_nbr[2] = q_max + slab_q_bounds_nbr[slab_index(1)] = q_min + slab_q_bounds_nbr[slab_index(2)] = q_max end end end diff --git a/src/MatrixFields/MatrixFields.jl b/src/MatrixFields/MatrixFields.jl index 64f06402a2..1a63772d3f 100644 --- a/src/MatrixFields/MatrixFields.jl +++ b/src/MatrixFields/MatrixFields.jl @@ -58,6 +58,7 @@ import ..RecursiveApply: rmap, rmaptype, rpromote_type, rzero, rconvert, radd, rsub, rmul, rdiv import ..RecursiveApply: ⊠, ⊞, ⊟ import ..DataLayouts: AbstractData +import ..DataLayouts: vindex import ..Geometry import ..Topologies import ..Spaces diff --git a/src/MatrixFields/single_field_solver.jl b/src/MatrixFields/single_field_solver.jl index e785c7cee0..362639f6d1 100644 --- a/src/MatrixFields/single_field_solver.jl +++ b/src/MatrixFields/single_field_solver.jl @@ -1,3 +1,4 @@ +import .DataLayouts: vindex dual_type(::Type{A}) where {A} = typeof(Geometry.dual(A.instance)) inv_return_type(::Type{X}) where {X} = error( @@ -110,6 +111,7 @@ function _single_field_solve_col!( Fields.field_values(x), unzip_tuple_field_values(Fields.field_values(A.entries)), Fields.field_values(b), + vindex, ) elseif A isa UniformScaling x .= inv(A.λ) .* b @@ -121,11 +123,18 @@ end unzip_tuple_field_values(data) = ntuple(i -> data.:($i), Val(length(propertynames(data)))) -function band_matrix_solve!(::Type{<:DiagonalMatrixRow}, _, x, Aⱼs, b) +function band_matrix_solve!( + ::Type{<:DiagonalMatrixRow}, + _, + x, + Aⱼs, + b, + vi = identity, +) (A₀,) = Aⱼs n = length(x) @inbounds for i in 1:n - x[i] = inv(A₀[i]) ⊠ b[i] + x[vi(i)] = inv(A₀[vi(i)]) ⊠ b[vi(i)] end end @@ -143,34 +152,41 @@ Transforms the tri-diagonal matrix into a unit upper bi-diagonal matrix, then solves the resulting system using back substitution. The order of multiplications has been modified in order to handle block vectors/matrices. =# -function band_matrix_solve!(::Type{<:TridiagonalMatrixRow}, cache, x, Aⱼs, b) +function band_matrix_solve!( + ::Type{<:TridiagonalMatrixRow}, + cache, + x, + Aⱼs, + b, + vi = identity, +) A₋₁, A₀, A₊₁ = Aⱼs Ux, U₊₁ = cache n = length(x) @inbounds begin - inv_D₀ = inv(A₀[1]) - U₊₁ᵢ₋₁ = inv_D₀ ⊠ A₊₁[1] - Uxᵢ₋₁ = inv_D₀ ⊠ b[1] - Ux[1] = Uxᵢ₋₁ - U₊₁[1] = U₊₁ᵢ₋₁ + inv_D₀ = inv(A₀[vi(1)]) + U₊₁ᵢ₋₁ = inv_D₀ ⊠ A₊₁[vi(1)] + Uxᵢ₋₁ = inv_D₀ ⊠ b[vi(1)] + Ux[vi(1)] = Uxᵢ₋₁ + U₊₁[vi(1)] = U₊₁ᵢ₋₁ for i in 2:n - A₋₁ᵢ = A₋₁[i] - inv_D₀ = inv(A₀[i] ⊟ A₋₁ᵢ ⊠ U₊₁ᵢ₋₁) - Uxᵢ₋₁ = inv_D₀ ⊠ (b[i] ⊟ A₋₁ᵢ ⊠ Uxᵢ₋₁) - Ux[i] = Uxᵢ₋₁ + A₋₁ᵢ = A₋₁[vi(i)] + inv_D₀ = inv(A₀[vi(i)] ⊟ A₋₁ᵢ ⊠ U₊₁ᵢ₋₁) + Uxᵢ₋₁ = inv_D₀ ⊠ (b[vi(i)] ⊟ A₋₁ᵢ ⊠ Uxᵢ₋₁) + Ux[vi(i)] = Uxᵢ₋₁ if i < n - U₊₁ᵢ₋₁ = inv_D₀ ⊠ A₊₁[i] # U₊₁[n] is outside the matrix. - U₊₁[i] = U₊₁ᵢ₋₁ + U₊₁ᵢ₋₁ = inv_D₀ ⊠ A₊₁[vi(i)] # U₊₁[n] is outside the matrix. + U₊₁[vi(i)] = U₊₁ᵢ₋₁ end end - x[n] = Ux[n] + x[vi(n)] = Ux[vi(n)] # Avoid steprange on GPU: https://cuda.juliagpu.org/stable/tutorials/performance/#Avoiding-StepRange i = (n - 1) # for i in (n - 1):-1:1 while i ≥ 1 - x[i] = Ux[i] ⊟ U₊₁[i] ⊠ x[i + 1] + x[vi(i)] = Ux[vi(i)] ⊟ U₊₁[vi(i)] ⊠ x[vi(i + 1)] i -= 1 end end @@ -195,36 +211,49 @@ Transforms the penta-diagonal matrix into a unit upper tri-diagonal matrix, then solves the resulting system using back substitution. The order of multiplications has been modified in order to handle block vectors/matrices. =# -function band_matrix_solve!(::Type{<:PentadiagonalMatrixRow}, cache, x, Aⱼs, b) +function band_matrix_solve!( + ::Type{<:PentadiagonalMatrixRow}, + cache, + x, + Aⱼs, + b, + vi = identity, +) A₋₂, A₋₁, A₀, A₊₁, A₊₂ = Aⱼs Ux, U₊₁, U₊₂ = cache n = length(x) @inbounds begin - inv_D₀ = inv(A₀[1]) - Ux[1] = inv_D₀ ⊠ b[1] - U₊₁[1] = inv_D₀ ⊠ A₊₁[1] - U₊₂[1] = inv_D₀ ⊠ A₊₂[1] + inv_D₀ = inv(A₀[vi(1)]) + Ux[vi(1)] = inv_D₀ ⊠ b[vi(1)] + U₊₁[vi(1)] = inv_D₀ ⊠ A₊₁[vi(1)] + U₊₂[vi(1)] = inv_D₀ ⊠ A₊₂[vi(1)] - inv_D₀ = inv(A₀[2] ⊟ A₋₁[2] ⊠ U₊₁[1]) - Ux[2] = inv_D₀ ⊠ (b[2] ⊟ A₋₁[2] ⊠ Ux[1]) - U₊₁[2] = inv_D₀ ⊠ (A₊₁[2] ⊟ A₋₁[2] ⊠ U₊₂[1]) - U₊₂[2] = inv_D₀ ⊠ A₊₂[2] + inv_D₀ = inv(A₀[vi(2)] ⊟ A₋₁[vi(2)] ⊠ U₊₁[vi(1)]) + Ux[vi(2)] = inv_D₀ ⊠ (b[vi(2)] ⊟ A₋₁[vi(2)] ⊠ Ux[vi(1)]) + U₊₁[vi(2)] = inv_D₀ ⊠ (A₊₁[vi(2)] ⊟ A₋₁[vi(2)] ⊠ U₊₂[vi(1)]) + U₊₂[vi(2)] = inv_D₀ ⊠ A₊₂[vi(2)] for i in 3:n - L₋₁ = A₋₁[i] ⊟ A₋₂[i] ⊠ U₊₁[i - 2] - inv_D₀ = inv(A₀[i] ⊟ L₋₁ ⊠ U₊₁[i - 1] ⊟ A₋₂[i] ⊠ U₊₂[i - 2]) - Ux[i] = inv_D₀ ⊠ (b[i] ⊟ L₋₁ ⊠ Ux[i - 1] ⊟ A₋₂[i] ⊠ Ux[i - 2]) - i < n && (U₊₁[i] = inv_D₀ ⊠ (A₊₁[i] ⊟ L₋₁ ⊠ U₊₂[i - 1])) - i < n - 1 && (U₊₂[i] = inv_D₀ ⊠ A₊₂[i]) + L₋₁ = A₋₁[vi(i)] ⊟ A₋₂[vi(i)] ⊠ U₊₁[vi(i - 2)] + inv_D₀ = inv( + A₀[vi(i)] ⊟ L₋₁ ⊠ U₊₁[vi(i - 1)] ⊟ A₋₂[vi(i)] ⊠ U₊₂[vi(i - 2)], + ) + Ux[vi(i)] = + inv_D₀ ⊠ + (b[vi(i)] ⊟ L₋₁ ⊠ Ux[vi(i - 1)] ⊟ A₋₂[vi(i)] ⊠ Ux[vi(i - 2)]) + i < n && (U₊₁[vi(i)] = inv_D₀ ⊠ (A₊₁[vi(i)] ⊟ L₋₁ ⊠ U₊₂[vi(i - 1)])) + i < n - 1 && (U₊₂[vi(i)] = inv_D₀ ⊠ A₊₂[vi(i)]) end - x[n] = Ux[n] - x[n - 1] = Ux[n - 1] ⊟ U₊₁[n - 1] ⊠ x[n] + x[vi(n)] = Ux[vi(n)] + x[vi(n - 1)] = Ux[vi(n - 1)] ⊟ U₊₁[vi(n - 1)] ⊠ x[vi(n)] # Avoid steprange on GPU: https://cuda.juliagpu.org/stable/tutorials/performance/#Avoiding-StepRange # for i in (n - 2):-1:1 i = (n - 2) while i ≥ 1 - x[i] = Ux[i] ⊟ U₊₁[i] ⊠ x[i + 1] ⊟ U₊₂[i] ⊠ x[i + 2] + x[vi(i)] = + Ux[vi(i)] ⊟ U₊₁[vi(i)] ⊠ x[vi(i + 1)] ⊟ + U₊₂[vi(i)] ⊠ x[vi(i + 2)] i -= 1 end end diff --git a/src/Operators/finitedifference.jl b/src/Operators/finitedifference.jl index 207808a514..a0872452bf 100644 --- a/src/Operators/finitedifference.jl +++ b/src/Operators/finitedifference.jl @@ -3144,7 +3144,7 @@ Base.@propagate_inbounds function getidx( field_data = Fields.field_values(bc) space = reconstruct_placeholder_space(axes(bc), parent_space) v = vidx(space, idx) - return @inbounds field_data[v] + return @inbounds field_data[CartesianIndex(1, 1, 1, v, 1)] end Base.@propagate_inbounds function getidx( parent_space, diff --git a/src/Operators/numericalflux.jl b/src/Operators/numericalflux.jl index aa3d38fa69..40fdeb2bd4 100644 --- a/src/Operators/numericalflux.jl +++ b/src/Operators/numericalflux.jl @@ -1,3 +1,4 @@ +import .DataLayouts: slab_index """ add_numerical_flux_internal!(fn, dydt, args...) @@ -40,7 +41,7 @@ function add_numerical_flux_internal!(fn, dydt, args...) dydt_slab⁺ = slab(Fields.field_values(dydt), elem⁺) for q in 1:Nq - sgeom⁻ = internal_surface_geometry_slab[q] + sgeom⁻ = internal_surface_geometry_slab[slab_index(q)] i⁻, j⁻ = Topologies.face_node_index(face⁻, Nq, q, false) i⁺, j⁺ = Topologies.face_node_index(face⁺, Nq, q, reversed) @@ -48,17 +49,21 @@ function add_numerical_flux_internal!(fn, dydt, args...) numflux⁻ = fn( sgeom⁻.normal, map( - slab -> slab isa DataSlab2D ? slab[i⁻, j⁻] : slab, + slab -> + slab isa DataSlab2D ? slab[slab_index(i⁻, j⁻)] : slab, arg_slabs⁻, ), map( - slab -> slab isa DataSlab2D ? slab[i⁺, j⁺] : slab, + slab -> + slab isa DataSlab2D ? slab[slab_index(i⁺, j⁺)] : slab, arg_slabs⁺, ), ) - dydt_slab⁻[i⁻, j⁻] = dydt_slab⁻[i⁻, j⁻] ⊟ (sgeom⁻.sWJ ⊠ numflux⁻) - dydt_slab⁺[i⁺, j⁺] = dydt_slab⁺[i⁺, j⁺] ⊞ (sgeom⁻.sWJ ⊠ numflux⁻) + dydt_slab⁻[slab_index(i⁻, j⁻)] = + dydt_slab⁻[slab_index(i⁻, j⁻)] ⊟ (sgeom⁻.sWJ ⊠ numflux⁻) + dydt_slab⁺[slab_index(i⁺, j⁺)] = + dydt_slab⁺[slab_index(i⁺, j⁺)] ⊞ (sgeom⁻.sWJ ⊠ numflux⁻) end end end @@ -115,17 +120,19 @@ function add_numerical_flux_boundary!(fn, dydt, args...) arg_slabs⁻ = map(arg -> slab(Fields.todata(arg), elem⁻), args) dydt_slab⁻ = slab(Fields.field_values(dydt), elem⁻) for q in 1:Nq - sgeom⁻ = boundary_surface_geometry_slab[q] + sgeom⁻ = boundary_surface_geometry_slab[slab_index(q)] i⁻, j⁻ = Topologies.face_node_index(face⁻, Nq, q, false) numflux⁻ = fn( sgeom⁻.normal, map( - slab -> slab isa DataSlab2D ? slab[i⁻, j⁻] : slab, + slab -> + slab isa DataSlab2D ? slab[slab_index(i⁻, j⁻)] : + slab, arg_slabs⁻, ), ) - dydt_slab⁻[i⁻, j⁻] = - dydt_slab⁻[i⁻, j⁻] ⊟ (sgeom⁻.sWJ ⊠ numflux⁻) + dydt_slab⁻[slab_index(i⁻, j⁻)] = + dydt_slab⁻[slab_index(i⁻, j⁻)] ⊟ (sgeom⁻.sWJ ⊠ numflux⁻) end end end diff --git a/src/Operators/spectralelement.jl b/src/Operators/spectralelement.jl index 6e54f8057f..39b3eef714 100644 --- a/src/Operators/spectralelement.jl +++ b/src/Operators/spectralelement.jl @@ -335,16 +335,17 @@ Base.@propagate_inbounds function get_node( space = reconstruct_placeholder_space(axes(field), parent_space) i, = Tuple(ij) if space isa Spaces.FaceExtrudedFiniteDifferenceSpace - v = slabidx.v + half + _v = slabidx.v + half elseif space isa Spaces.CenterExtrudedFiniteDifferenceSpace || space isa Spaces.AbstractSpectralElementSpace - v = slabidx.v + _v = slabidx.v else error("invalid space") end h = slabidx.h fv = Fields.field_values(field) - return fv[i, nothing, nothing, v, h] + v = isnothing(_v) ? 1 : _v + return fv[CartesianIndex(i, 1, 1, v, h)] end Base.@propagate_inbounds function get_node( parent_space, @@ -355,16 +356,17 @@ Base.@propagate_inbounds function get_node( space = reconstruct_placeholder_space(axes(field), parent_space) i, j = Tuple(ij) if space isa Spaces.FaceExtrudedFiniteDifferenceSpace - v = slabidx.v + half + _v = slabidx.v + half elseif space isa Spaces.CenterExtrudedFiniteDifferenceSpace || space isa Spaces.AbstractSpectralElementSpace - v = slabidx.v + _v = slabidx.v else error("invalid space") end h = slabidx.h fv = Fields.field_values(field) - return fv[i, j, nothing, v, h] + v = isnothing(_v) ? 1 : _v + return fv[CartesianIndex(i, j, 1, v, h)] end @@ -411,12 +413,13 @@ Base.@propagate_inbounds function get_local_geometry( i, = Tuple(ij) h = slabidx.h if space isa Spaces.FaceExtrudedFiniteDifferenceSpace - v = slabidx.v + half + _v = slabidx.v + half else - v = slabidx.v + _v = slabidx.v end lgd = Spaces.local_geometry_data(space) - return lgd[i, nothing, nothing, v, h] + v = isnothing(_v) ? 1 : _v + return lgd[CartesianIndex(i, 1, 1, v, h)] end Base.@propagate_inbounds function get_local_geometry( space::Union{ @@ -429,12 +432,13 @@ Base.@propagate_inbounds function get_local_geometry( i, j = Tuple(ij) h = slabidx.h if space isa Spaces.FaceExtrudedFiniteDifferenceSpace - v = slabidx.v + half + _v = slabidx.v + half else - v = slabidx.v + _v = slabidx.v end + v = isnothing(_v) ? 1 : _v lgd = Spaces.local_geometry_data(space) - return lgd[i, j, nothing, v, h] + return lgd[CartesianIndex(i, j, 1, v, h)] end Base.@propagate_inbounds function set_node!( @@ -446,13 +450,14 @@ Base.@propagate_inbounds function set_node!( ) i, = Tuple(ij) if space isa Spaces.FaceExtrudedFiniteDifferenceSpace - v = slabidx.v + half + _v = slabidx.v + half else - v = slabidx.v + _v = slabidx.v end h = slabidx.h fv = Fields.field_values(field) - fv[i, nothing, nothing, v, h] = val + v = isnothing(_v) ? 1 : _v + fv[CartesianIndex(i, 1, 1, v, h)] = val end Base.@propagate_inbounds function set_node!( space, @@ -463,13 +468,14 @@ Base.@propagate_inbounds function set_node!( ) i, j = Tuple(ij) if space isa Spaces.FaceExtrudedFiniteDifferenceSpace - v = slabidx.v + half + _v = slabidx.v + half else - v = slabidx.v + _v = slabidx.v end h = slabidx.h fv = Fields.field_values(field) - fv[i, j, nothing, v, h] = val + v = isnothing(_v) ? 1 : _v + fv[CartesianIndex(i, j, 1, v, h)] = val end Base.Broadcast.BroadcastStyle( @@ -540,13 +546,14 @@ function apply_operator(op::Divergence{(1,)}, space, slabidx, arg) v, ) for ii in 1:Nq - out[ii] = out[ii] ⊞ (D[ii, i] ⊠ Jv¹) + out[slab_index(ii)] = out[slab_index(ii)] ⊞ (D[ii, i] ⊠ Jv¹) end end @inbounds for i in 1:Nq ij = CartesianIndex((i,)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i] = RecursiveApply.rmul(out[i], local_geometry.invJ) + out[slab_index(i)] = + RecursiveApply.rmul(out[slab_index(i)], local_geometry.invJ) end return Field(SArray(out), space) end @@ -575,7 +582,7 @@ Base.@propagate_inbounds function apply_operator( v, ) for ii in 1:Nq - out[ii, j] = out[ii, j] ⊞ (D[ii, i] ⊠ Jv¹) + out[slab_index(ii, j)] = out[slab_index(ii, j)] ⊞ (D[ii, i] ⊠ Jv¹) end Jv² = local_geometry.J ⊠ RecursiveApply.rmap( @@ -583,13 +590,14 @@ Base.@propagate_inbounds function apply_operator( v, ) for jj in 1:Nq - out[i, jj] = out[i, jj] ⊞ (D[jj, j] ⊠ Jv²) + out[slab_index(i, jj)] = out[slab_index(i, jj)] ⊞ (D[jj, j] ⊠ Jv²) end end @inbounds for j in 1:Nq, i in 1:Nq ij = CartesianIndex((i, j)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i, j] = RecursiveApply.rmul(out[i, j], local_geometry.invJ) + out[slab_index(i, j)] = + RecursiveApply.rmul(out[slab_index(i, j)], local_geometry.invJ) end return Field(SArray(out), space) end @@ -657,13 +665,14 @@ function apply_operator(op::WeakDivergence{(1,)}, space, slabidx, arg) v, ) for ii in 1:Nq - out[ii] = out[ii] ⊞ (D[i, ii] ⊠ WJv¹) + out[slab_index(ii)] = out[slab_index(ii)] ⊞ (D[i, ii] ⊠ WJv¹) end end @inbounds for i in 1:Nq ij = CartesianIndex((i,)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i] = RecursiveApply.rdiv(out[i], ⊟(local_geometry.WJ)) + out[slab_index(i)] = + RecursiveApply.rdiv(out[slab_index(i)], ⊟(local_geometry.WJ)) end return Field(SArray(out), space) end @@ -687,7 +696,7 @@ function apply_operator(op::WeakDivergence{(1, 2)}, space, slabidx, arg) v, ) for ii in 1:Nq - out[ii, j] = out[ii, j] ⊞ (D[i, ii] ⊠ WJv¹) + out[slab_index(ii, j)] = out[slab_index(ii, j)] ⊞ (D[i, ii] ⊠ WJv¹) end WJv² = local_geometry.WJ ⊠ RecursiveApply.rmap( @@ -695,13 +704,14 @@ function apply_operator(op::WeakDivergence{(1, 2)}, space, slabidx, arg) v, ) for jj in 1:Nq - out[i, jj] = out[i, jj] ⊞ (D[j, jj] ⊠ WJv²) + out[slab_index(i, jj)] = out[slab_index(i, jj)] ⊞ (D[j, jj] ⊠ WJv²) end end @inbounds for j in 1:Nq, i in 1:Nq ij = CartesianIndex((i, j)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i, j] = RecursiveApply.rdiv(out[i, j], ⊟(local_geometry.WJ)) + out[slab_index(i, j)] = + RecursiveApply.rdiv(out[slab_index(i, j)], ⊟(local_geometry.WJ)) end return Field(SArray(out), space) end @@ -749,7 +759,7 @@ function apply_operator(op::Gradient{(1,)}, space, slabidx, arg) x = get_node(space, arg, ij, slabidx) for ii in 1:Nq ∂f∂ξ = Geometry.Covariant1Vector(D[ii, i]) ⊗ x - out[ii] += ∂f∂ξ + out[slab_index(ii)] += ∂f∂ξ end end return Field(SArray(out), space) @@ -775,11 +785,11 @@ Base.@propagate_inbounds function apply_operator( x = get_node(space, arg, ij, slabidx) for ii in 1:Nq ∂f∂ξ₁ = Geometry.Covariant12Vector(D[ii, i], zero(eltype(D))) ⊗ x - out[ii, j] = out[ii, j] ⊞ ∂f∂ξ₁ + out[slab_index(ii, j)] = out[slab_index(ii, j)] ⊞ ∂f∂ξ₁ end for jj in 1:Nq ∂f∂ξ₂ = Geometry.Covariant12Vector(zero(eltype(D)), D[jj, j]) ⊗ x - out[i, jj] = out[i, jj] ⊞ ∂f∂ξ₂ + out[slab_index(i, jj)] = out[slab_index(i, jj)] ⊞ ∂f∂ξ₂ end end return Field(SArray(out), space) @@ -842,14 +852,14 @@ function apply_operator(op::WeakGradient{(1,)}, space, slabidx, arg) Wx = W ⊠ get_node(space, arg, ij, slabidx) for ii in 1:Nq Dᵀ₁Wf = Geometry.Covariant1Vector(D[i, ii]) ⊗ Wx - out[ii] = out[ii] ⊟ Dᵀ₁Wf + out[slab_index(ii)] = out[slab_index(ii)] ⊟ Dᵀ₁Wf end end @inbounds for i in 1:Nq ij = CartesianIndex((i,)) local_geometry = get_local_geometry(space, ij, slabidx) W = local_geometry.WJ * local_geometry.invJ - out[i] = RecursiveApply.rdiv(out[i], W) + out[slab_index(i)] = RecursiveApply.rdiv(out[slab_index(i)], W) end return Field(SArray(out), space) end @@ -871,18 +881,18 @@ function apply_operator(op::WeakGradient{(1, 2)}, space, slabidx, arg) Wx = W ⊠ get_node(space, arg, ij, slabidx) for ii in 1:Nq Dᵀ₁Wf = Geometry.Covariant12Vector(D[i, ii], zero(eltype(D))) ⊗ Wx - out[ii, j] = out[ii, j] ⊟ Dᵀ₁Wf + out[slab_index(ii, j)] = out[slab_index(ii, j)] ⊟ Dᵀ₁Wf end for jj in 1:Nq Dᵀ₂Wf = Geometry.Covariant12Vector(zero(eltype(D)), D[j, jj]) ⊗ Wx - out[i, jj] = out[i, jj] ⊟ Dᵀ₂Wf + out[slab_index(i, jj)] = out[slab_index(i, jj)] ⊟ Dᵀ₂Wf end end @inbounds for j in 1:Nq, i in 1:Nq ij = CartesianIndex((i, j)) local_geometry = get_local_geometry(space, ij, slabidx) W = local_geometry.WJ * local_geometry.invJ - out[i, j] = RecursiveApply.rdiv(out[i, j], W) + out[slab_index(i, j)] = RecursiveApply.rdiv(out[slab_index(i, j)], W) end return Field(SArray(out), space) end @@ -954,7 +964,8 @@ function apply_operator(op::Curl{(1,)}, space, slabidx, arg) v₃ = Geometry.covariant3(v, local_geometry) for ii in 1:Nq D₁v₃ = D[ii, i] ⊠ v₃ - out[ii] = out[ii] ⊞ Geometry.Contravariant2Vector(⊟(D₁v₃)) + out[slab_index(ii)] = + out[slab_index(ii)] ⊞ Geometry.Contravariant2Vector(⊟(D₁v₃)) end end elseif RT <: Geometry.Contravariant3Vector @@ -965,7 +976,8 @@ function apply_operator(op::Curl{(1,)}, space, slabidx, arg) v₂ = Geometry.covariant2(v, local_geometry) for ii in 1:Nq D₁v₂ = D[ii, i] ⊠ v₂ - out[ii] = out[ii] ⊞ Geometry.Contravariant3Vector(D₁v₂) + out[slab_index(ii)] = + out[slab_index(ii)] ⊞ Geometry.Contravariant3Vector(D₁v₂) end end elseif RT <: Geometry.Contravariant23Vector @@ -978,8 +990,9 @@ function apply_operator(op::Curl{(1,)}, space, slabidx, arg) for ii in 1:Nq D₁v₃ = D[ii, i] ⊠ v₃ D₁v₂ = D[ii, i] ⊠ v₂ - out[ii] = - out[ii] ⊞ Geometry.Contravariant23Vector(⊟(D₁v₃), D₁v₂) + out[slab_index(ii)] = + out[slab_index(ii)] ⊞ + Geometry.Contravariant23Vector(⊟(D₁v₃), D₁v₂) end end else @@ -988,7 +1001,8 @@ function apply_operator(op::Curl{(1,)}, space, slabidx, arg) @inbounds for i in 1:Nq ij = CartesianIndex((i,)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i] = RecursiveApply.rmul(out[i], local_geometry.invJ) + out[slab_index(i)] = + RecursiveApply.rmul(out[slab_index(i)], local_geometry.invJ) end return Field(SArray(out), space) end @@ -1012,12 +1026,15 @@ function apply_operator(op::Curl{(1, 2)}, space, slabidx, arg) v₁ = Geometry.covariant1(v, local_geometry) for jj in 1:Nq D₂v₁ = D[jj, j] ⊠ v₁ - out[i, jj] = out[i, jj] ⊞ Geometry.Contravariant3Vector(⊟(D₂v₁)) + out[slab_index(i, jj)] = + out[slab_index(i, jj)] ⊞ + Geometry.Contravariant3Vector(⊟(D₂v₁)) end v₂ = Geometry.covariant2(v, local_geometry) for ii in 1:Nq D₁v₂ = D[ii, i] ⊠ v₂ - out[ii, j] = out[ii, j] ⊞ Geometry.Contravariant3Vector(D₁v₂) + out[slab_index(ii, j)] = + out[slab_index(ii, j)] ⊞ Geometry.Contravariant3Vector(D₁v₂) end end # input data is a Covariant3Vector field @@ -1029,14 +1046,14 @@ function apply_operator(op::Curl{(1, 2)}, space, slabidx, arg) v₃ = Geometry.covariant3(v, local_geometry) for ii in 1:Nq D₁v₃ = D[ii, i] ⊠ v₃ - out[ii, j] = - out[ii, j] ⊞ + out[slab_index(ii, j)] = + out[slab_index(ii, j)] ⊞ Geometry.Contravariant12Vector(zero(D₁v₃), ⊟(D₁v₃)) end for jj in 1:Nq D₂v₃ = D[jj, j] ⊠ v₃ - out[i, jj] = - out[i, jj] ⊞ + out[slab_index(i, jj)] = + out[slab_index(i, jj)] ⊞ Geometry.Contravariant12Vector(D₂v₃, zero(D₂v₃)) end end @@ -1051,15 +1068,15 @@ function apply_operator(op::Curl{(1, 2)}, space, slabidx, arg) for ii in 1:Nq D₁v₃ = D[ii, i] ⊠ v₃ D₁v₂ = D[ii, i] ⊠ v₂ - out[ii, j] = - out[ii, j] ⊞ + out[slab_index(ii, j)] = + out[slab_index(ii, j)] ⊞ Geometry.Contravariant123Vector(zero(D₁v₃), ⊟(D₁v₃), D₁v₂) end for jj in 1:Nq D₂v₃ = D[jj, j] ⊠ v₃ D₂v₁ = D[jj, j] ⊠ v₁ - out[i, jj] = - out[i, jj] ⊞ + out[slab_index(i, jj)] = + out[slab_index(i, jj)] ⊞ Geometry.Contravariant123Vector(D₂v₃, zero(D₂v₃), ⊟(D₂v₁)) end end @@ -1069,7 +1086,8 @@ function apply_operator(op::Curl{(1, 2)}, space, slabidx, arg) @inbounds for j in 1:Nq, i in 1:Nq ij = CartesianIndex((i, j)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i, j] = RecursiveApply.rmul(out[i, j], local_geometry.invJ) + out[slab_index(i, j)] = + RecursiveApply.rmul(out[slab_index(i, j)], local_geometry.invJ) end return Field(SArray(out), space) end @@ -1136,7 +1154,8 @@ function apply_operator(op::WeakCurl{(1,)}, space, slabidx, arg) Wv₃ = W ⊠ Geometry.covariant3(v, local_geometry) for ii in 1:Nq Dᵀ₁Wv₃ = D[i, ii] ⊠ Wv₃ - out[ii] = out[ii] ⊞ Geometry.Contravariant2Vector(Dᵀ₁Wv₃) + out[slab_index(ii)] = + out[slab_index(ii)] ⊞ Geometry.Contravariant2Vector(Dᵀ₁Wv₃) end end elseif RT <: Geometry.Contravariant3Vector @@ -1148,7 +1167,9 @@ function apply_operator(op::WeakCurl{(1,)}, space, slabidx, arg) Wv₂ = W ⊠ Geometry.covariant2(v, local_geometry) for ii in 1:Nq Dᵀ₁Wv₂ = D[i, ii] ⊠ Wv₂ - out[ii] = out[ii] ⊞ Geometry.Contravariant3Vector(⊟(Dᵀ₁Wv₂)) + out[slab_index(ii)] = + out[slab_index(ii)] ⊞ + Geometry.Contravariant3Vector(⊟(Dᵀ₁Wv₂)) end end elseif RT <: Geometry.Contravariant23Vector @@ -1162,8 +1183,9 @@ function apply_operator(op::WeakCurl{(1,)}, space, slabidx, arg) for ii in 1:Nq Dᵀ₁Wv₃ = D[i, ii] ⊠ Wv₃ Dᵀ₁Wv₂ = D[i, ii] ⊠ Wv₂ - out[ii] = - out[ii] ⊞ Geometry.Contravariant23Vector(Dᵀ₁Wv₃, ⊟(Dᵀ₁Wv₂)) + out[slab_index(ii)] = + out[slab_index(ii)] ⊞ + Geometry.Contravariant23Vector(Dᵀ₁Wv₃, ⊟(Dᵀ₁Wv₂)) end end else @@ -1172,7 +1194,8 @@ function apply_operator(op::WeakCurl{(1,)}, space, slabidx, arg) @inbounds for i in 1:Nq ij = CartesianIndex((i,)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i] = RecursiveApply.rdiv(out[i], local_geometry.WJ) + out[slab_index(i)] = + RecursiveApply.rdiv(out[slab_index(i)], local_geometry.WJ) end return Field(SArray(out), space) end @@ -1197,13 +1220,16 @@ function apply_operator(op::WeakCurl{(1, 2)}, space, slabidx, arg) Wv₁ = W ⊠ Geometry.covariant1(v, local_geometry) for jj in 1:Nq Dᵀ₂Wv₁ = D[j, jj] ⊠ Wv₁ - out[i, jj] = out[i, jj] ⊞ Geometry.Contravariant3Vector(Dᵀ₂Wv₁) + out[slab_index(i, jj)] = + out[slab_index(i, jj)] ⊞ + Geometry.Contravariant3Vector(Dᵀ₂Wv₁) end Wv₂ = W ⊠ Geometry.covariant2(v, local_geometry) for ii in 1:Nq Dᵀ₁Wv₂ = D[i, ii] ⊠ Wv₂ - out[ii, j] = - out[ii, j] ⊞ Geometry.Contravariant3Vector(⊟(Dᵀ₁Wv₂)) + out[slab_index(ii, j)] = + out[slab_index(ii, j)] ⊞ + Geometry.Contravariant3Vector(⊟(Dᵀ₁Wv₂)) end end # input data is a Covariant3Vector field @@ -1216,14 +1242,14 @@ function apply_operator(op::WeakCurl{(1, 2)}, space, slabidx, arg) Wv₃ = W ⊠ Geometry.covariant3(v, local_geometry) for ii in 1:Nq Dᵀ₁Wv₃ = D[i, ii] ⊠ Wv₃ - out[ii, j] = - out[ii, j] ⊞ + out[slab_index(ii, j)] = + out[slab_index(ii, j)] ⊞ Geometry.Contravariant12Vector(zero(Dᵀ₁Wv₃), Dᵀ₁Wv₃) end for jj in 1:Nq Dᵀ₂Wv₃ = D[j, jj] ⊠ Wv₃ - out[i, jj] = - out[i, jj] ⊞ + out[slab_index(i, jj)] = + out[slab_index(i, jj)] ⊞ Geometry.Contravariant12Vector(⊟(Dᵀ₂Wv₃), zero(Dᵀ₂Wv₃)) end end @@ -1239,8 +1265,8 @@ function apply_operator(op::WeakCurl{(1, 2)}, space, slabidx, arg) for ii in 1:Nq Dᵀ₁Wv₃ = D[i, ii] ⊠ Wv₃ Dᵀ₁Wv₂ = D[i, ii] ⊠ Wv₂ - out[ii, j] = - out[ii, j] ⊞ Geometry.Contravariant123Vector( + out[slab_index(ii, j)] = + out[slab_index(ii, j)] ⊞ Geometry.Contravariant123Vector( zero(Dᵀ₁Wv₃), Dᵀ₁Wv₃, ⊟(Dᵀ₁Wv₂), @@ -1249,8 +1275,8 @@ function apply_operator(op::WeakCurl{(1, 2)}, space, slabidx, arg) for jj in 1:Nq Dᵀ₂Wv₃ = D[j, jj] ⊠ Wv₃ Dᵀ₂Wv₁ = D[j, jj] ⊠ Wv₁ - out[i, jj] = - out[i, jj] ⊞ Geometry.Contravariant123Vector( + out[slab_index(i, jj)] = + out[slab_index(i, jj)] ⊞ Geometry.Contravariant123Vector( ⊟(Dᵀ₂Wv₃), zero(Dᵀ₂Wv₃), Dᵀ₂Wv₁, @@ -1263,7 +1289,8 @@ function apply_operator(op::WeakCurl{(1, 2)}, space, slabidx, arg) for j in 1:Nq, i in 1:Nq ij = CartesianIndex((i, j)) local_geometry = get_local_geometry(space, ij, slabidx) - out[i, j] = RecursiveApply.rdiv(out[i, j], local_geometry.WJ) + out[slab_index(i, j)] = + RecursiveApply.rdiv(out[slab_index(i, j)], local_geometry.WJ) end return Field(SArray(out), space) end @@ -1316,7 +1343,7 @@ function apply_operator(op::Interpolate{(1,)}, space_out, slabidx, arg) r, ) end - out[i] = r + out[slab_index(i)] = r end return Field(SArray(out), space_out) end @@ -1346,10 +1373,10 @@ function apply_operator(op::Interpolate{(1, 2)}, space_out, slabidx, arg) r, ) end - temp[i, j] = r + temp[slab_index(i, j)] = r end @inbounds for j in 1:Nq_out, i in 1:Nq_out - out[i, j] = RecursiveApply.rmatmul2(Imat, temp, i, j) + out[slab_index(i, j)] = rmatmul2(Imat, temp, i, j) end return Field(SArray(out), space_out) end @@ -1409,7 +1436,7 @@ function apply_operator(op::Restrict{(1,)}, space_out, slabidx, arg) end ij_out = CartesianIndex((i,)) WJ_out = get_local_geometry(space_out, ij_out, slabidx).WJ - out[i] = RecursiveApply.rdiv(r, WJ_out) + out[slab_index(i)] = RecursiveApply.rdiv(r, WJ_out) end return Field(SArray(out), space_out) end @@ -1440,15 +1467,13 @@ function apply_operator(op::Restrict{(1, 2)}, space_out, slabidx, arg) r, ) end - temp[i, j] = r + temp[slab_index(i, j)] = r end @inbounds for j in 1:Nq_out, i in 1:Nq_out ij_out = CartesianIndex((i, j)) WJ_out = get_local_geometry(space_out, ij_out, slabidx).WJ - out[i, j] = RecursiveApply.rdiv( - RecursiveApply.rmatmul2(ImatT, temp, i, j), - WJ_out, - ) + out[slab_index(i, j)] = + RecursiveApply.rdiv(rmatmul2(ImatT, temp, i, j), WJ_out) end return Field(SArray(out), space_out) end @@ -1475,11 +1500,11 @@ function tensor_product!( in_slab = slab(indata, v, h) out_slab = slab(out, v, h) for i in 1:Ni_out - r = M[i, 1] ⊠ in_slab[1] + r = M[i, 1] ⊠ in_slab[slab_index(1)] for ii in 2:Ni_in - r = RecursiveApply.rmuladd(M[i, ii], in_slab[ii], r) + r = RecursiveApply.rmuladd(M[i, ii], in_slab[slab_index(ii)], r) end - out_slab[i] = r + out_slab[slab_index(i)] = r end end return out @@ -1501,10 +1526,10 @@ function tensor_product!( in_slab = slab(indata, h) out_slab = slab(out, h) for j in 1:Nij_in, i in 1:Nij_out - temp[i, j] = RecursiveApply.rmatmul1(M, in_slab, i, j) + temp[slab_index(i, j)] = rmatmul1(M, in_slab, i, j) end for j in 1:Nij_out, i in 1:Nij_out - out_slab[i, j] = RecursiveApply.rmatmul2(M, temp, i, j) + out_slab[slab_index(i, j)] = rmatmul2(M, temp, i, j) end end return out @@ -1518,10 +1543,10 @@ function tensor_product!( # temporary storage temp = MArray{Tuple{Nij_out, Nij_in}, S, 2, Nij_out * Nij_in}(undef) @inbounds for j in 1:Nij_in, i in 1:Nij_out - temp[i, j] = RecursiveApply.rmatmul1(M, in_slab, i, j) + temp[slab_index(i, j)] = rmatmul1(M, in_slab, i, j) end @inbounds for j in 1:Nij_out, i in 1:Nij_out - out_slab[i, j] = RecursiveApply.rmatmul2(M, temp, i, j) + out_slab[slab_index(i, j)] = rmatmul2(M, temp, i, j) end return out_slab end @@ -1581,3 +1606,38 @@ Returns a 2D Matrix for plotting / visualizing 2D Fields. """ matrix_interpolate(field::Field, Nu::Integer) = matrix_interpolate(field, Quadratures.Uniform{Nu}()) + +import .DataLayouts: slab_index + +""" + rmatmul1(W, S, i, j) + +Recursive matrix product along the 1st dimension of `S`. Equivalent to: + + mapreduce(⊠, ⊞, W[i,:], S[:,j]) + +""" +function rmatmul1(W, S, i, j) + Nq = size(W, 2) + @inbounds r = W[i, 1] ⊠ S[slab_index(1, j)] + @inbounds for ii in 2:Nq + r = RecursiveApply.rmuladd(W[i, ii], S[slab_index(ii, j)], r) + end + return r +end + +""" + rmatmul2(W, S, i, j) + +Recursive matrix product along the 2nd dimension `S`. Equivalent to: + + mapreduce(⊠, ⊞, W[j,:], S[i, :]) +""" +function rmatmul2(W, S, i, j) + Nq = size(W, 2) + @inbounds r = W[j, 1] ⊠ S[slab_index(i, 1)] + @inbounds for jj in 2:Nq + r = RecursiveApply.rmuladd(W[j, jj], S[slab_index(i, jj)], r) + end + return r +end diff --git a/src/RecursiveApply/RecursiveApply.jl b/src/RecursiveApply/RecursiveApply.jl index dfbbe79e5b..4fcc817f8f 100755 --- a/src/RecursiveApply/RecursiveApply.jl +++ b/src/RecursiveApply/RecursiveApply.jl @@ -212,38 +212,4 @@ rmuladd(X, w::Number, Y) = rmap((x, y) -> muladd(x, w, y), X, Y) rmuladd(X::Number, w::Number, Y) = rmap((x, y) -> muladd(x, w, y), X, Y) rmuladd(w::Number, x::Number, y::Number) = muladd(w, x, y) -""" - rmatmul1(W, S, i, j) - -Recursive matrix product along the 1st dimension of `S`. Equivalent to: - - mapreduce(⊠, ⊞, W[i,:], S[:,j]) - -""" -function rmatmul1(W, S, i, j) - Nq = size(W, 2) - @inbounds r = W[i, 1] ⊠ S[1, j] - @inbounds for ii in 2:Nq - r = rmuladd(W[i, ii], S[ii, j], r) - end - return r -end - -""" - rmatmul2(W, S, i, j) - -Recursive matrix product along the 2nd dimension `S`. Equivalent to: - - mapreduce(⊠, ⊞, W[j,:], S[i, :]) - -""" -function rmatmul2(W, S, i, j) - Nq = size(W, 2) - @inbounds r = W[j, 1] ⊠ S[i, 1] - @inbounds for jj in 2:Nq - r = rmuladd(W[j, jj], S[i, jj], r) - end - return r -end - end # module diff --git a/src/Remapping/distributed_remapping.jl b/src/Remapping/distributed_remapping.jl index 8f958c8627..5f22c9636d 100644 --- a/src/Remapping/distributed_remapping.jl +++ b/src/Remapping/distributed_remapping.jl @@ -418,8 +418,8 @@ function set_interpolated_values_cpu_kernel!( if prev_lidx != h || prev_vindex != vindex for j in 1:Nq, i in 1:Nq scratch_field_values[i, j] = ( - A * field_values[i, j, nothing, v_lo, h] + - B * field_values[i, j, nothing, v_hi, h] + A * field_values[CartesianIndex(i, j, 1, v_lo, h)] + + B * field_values[CartesianIndex(i, j, 1, v_hi, h)] ) end prev_vindex, prev_lidx = vindex, h @@ -467,8 +467,8 @@ function set_interpolated_values_cpu_kernel!( if prev_lidx != h || prev_vindex != vindex for i in 1:Nq scratch_field_values[i] = ( - A * field_values[i, nothing, nothing, v_lo, h] + - B * field_values[i, nothing, nothing, v_hi, h] + A * field_values[CartesianIndex(i, 1, 1, v_lo, h)] + + B * field_values[CartesianIndex(i, 1, 1, v_hi, h)] ) end prev_vindex, prev_lidx = vindex, h @@ -577,13 +577,13 @@ function _set_interpolated_values_device!( out[out_index, field_index] += local_horiz_interpolation_weights[1][out_index, i] * local_horiz_interpolation_weights[2][out_index, j] * - field_values[i, j, nothing, nothing, h] + field_values[CartesianIndex(i, j, 1, 1, h)] end elseif hdims == 1 for i in 1:Nq out[out_index, field_index] += local_horiz_interpolation_weights[1][out_index, i] * - field_values[i, nothing, nothing, nothing, h] + field_values[CartesianIndex(i, 1, 1, 1, h)] end end end diff --git a/src/Topologies/Topologies.jl b/src/Topologies/Topologies.jl index 2827781ae9..77dd9b20d6 100644 --- a/src/Topologies/Topologies.jl +++ b/src/Topologies/Topologies.jl @@ -10,6 +10,7 @@ import ..Geometry import ..Domains: Domains, coordinate_type import ..Meshes: Meshes, domain, coordinates import ..DataLayouts +import ..DataLayouts: slab_index import ..slab, ..column, ..level import ..DeviceSideDevice, ..DeviceSideContext diff --git a/src/Topologies/dss.jl b/src/Topologies/dss.jl index 5d2ba785a2..4432ffc0b6 100644 --- a/src/Topologies/dss.jl +++ b/src/Topologies/dss.jl @@ -786,13 +786,13 @@ function dss_local_vertices!( ) do (lidx, vert) ip = perimeter_vertex_node_index(vert) perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] + perimeter_slab[slab_index(ip)] end # scatter: assign sum to shared vertices for (lidx, vert) in vertex perimeter_slab = slab(perimeter_data, level, lidx) ip = perimeter_vertex_node_index(vert) - perimeter_slab[ip] = sum_data + perimeter_slab[slab_index(ip)] = sum_data end end end @@ -815,9 +815,11 @@ function dss_local_faces!( perimeter_slab1 = slab(perimeter_data, level, lidx1) perimeter_slab2 = slab(perimeter_data, level, lidx2) for (ip1, ip2) in zip(pr1, pr2) - val = perimeter_slab1[ip1] ⊞ perimeter_slab2[ip2] - perimeter_slab1[ip1] = val - perimeter_slab2[ip2] = val + val = + perimeter_slab1[slab_index(ip1)] ⊞ + perimeter_slab2[slab_index(ip2)] + perimeter_slab1[slab_index(ip1)] = val + perimeter_slab2[slab_index(ip2)] = val end end end @@ -860,9 +862,12 @@ function dss_local_ghost!( if !isghost lidx = idx perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] + perimeter_slab[slab_index(ip)] else - RecursiveApply.rmap(zero, slab(perimeter_data, 1, 1)[1]) + RecursiveApply.rmap( + zero, + slab(perimeter_data, 1, 1)[slab_index(1)], + ) end end for (isghost, idx, vert) in vertex @@ -870,7 +875,7 @@ function dss_local_ghost!( ip = perimeter_vertex_node_index(vert) lidx = idx perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] = sum_data + perimeter_slab[slab_index(ip)] = sum_data end end end @@ -906,13 +911,13 @@ function dss_ghost!( ipresult = perimeter_vertex_node_index(lvertresult) for level in 1:nlevels result_slab = slab(perimeter_data, level, idxresult) - result = result_slab[ipresult] + result = result_slab[slab_index(ipresult)] for (isghost, idx, vert) in vertex if !isghost ip = perimeter_vertex_node_index(vert) lidx = idx perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] = result + perimeter_slab[slab_index(ip)] = result end end end diff --git a/src/Topologies/dss_transform.jl b/src/Topologies/dss_transform.jl index 75fb5f199c..1fec9d587a 100644 --- a/src/Topologies/dss_transform.jl +++ b/src/Topologies/dss_transform.jl @@ -2,47 +2,30 @@ import ..Topologies: Topology2D using ..RecursiveApply """ - dss_transform(arg, local_geometry, weight, I...) + dss_transform(arg, local_geometry, weight, I) -Transfrom `arg[I...]` to a basis for direct stiffness summation (DSS). +Transfrom `arg[I]` to a basis for direct stiffness summation (DSS). Transformations only apply to vector quantities. -- `local_geometry[I...]` is the relevant `LocalGeometry` object. If it is `nothing`, then no transformation is performed -- `weight[I...]` is the relevant DSS weights. If `weight` is `nothing`, then the result is simply summation. +- `local_geometry[I]` is the relevant `LocalGeometry` object. If it is `nothing`, then no transformation is performed +- `weight[I]` is the relevant DSS weights. If `weight` is `nothing`, then the result is simply summation. See [`ClimaCore.Spaces.weighted_dss!`](@ref). """ -Base.@propagate_inbounds dss_transform(arg, local_geometry, weight, i, j) = - dss_transform(arg[i, j], local_geometry[i, j], weight[i, j]) +Base.@propagate_inbounds dss_transform(arg, local_geometry, weight, I) = + dss_transform(arg[I], local_geometry[I], weight[I]) Base.@propagate_inbounds dss_transform( arg, local_geometry, weight::Nothing, - i, - j, -) = dss_transform(arg[i, j], local_geometry[i, j], 1) + I, +) = dss_transform(arg[I], local_geometry[I], 1) Base.@propagate_inbounds dss_transform( arg, local_geometry::Nothing, weight::Nothing, - i, - j, -) = arg[i, j] - -Base.@propagate_inbounds dss_transform(arg, local_geometry, weight, i) = - dss_transform(arg[i], local_geometry[i], weight[i]) -Base.@propagate_inbounds dss_transform( - arg, - local_geometry, - weight::Nothing, - i, -) = dss_transform(arg[i], local_geometry[i], 1) -Base.@propagate_inbounds dss_transform( - arg, - local_geometry::Nothing, - weight::Nothing, - i, -) = arg[i] + I, +) = arg[I] @inline function dss_transform( arg::Tuple{}, @@ -152,23 +135,9 @@ Base.@propagate_inbounds dss_untransform( ::Type{T}, targ, local_geometry, - i, - j, -) where {T} = dss_untransform(T, targ, local_geometry[i, j]) -Base.@propagate_inbounds dss_untransform( - ::Type{T}, - targ, - local_geometry::Nothing, - i, - j, -) where {T} = dss_untransform(T, targ, local_geometry) -Base.@propagate_inbounds dss_untransform( - ::Type{T}, - targ, - local_geometry, - i, -) where {T} = dss_untransform(T, targ, local_geometry[i]) -@inline dss_untransform(::Type{T}, targ, local_geometry::Nothing, i) where {T} = + I, +) where {T} = dss_untransform(T, targ, local_geometry[I]) +@inline dss_untransform(::Type{T}, targ, local_geometry::Nothing, I) where {T} = dss_untransform(T, targ, local_geometry) @inline function dss_untransform( ::Type{NamedTuple{names, T}}, @@ -268,9 +237,12 @@ function _representative_slab( if isnothing(data) return nothing elseif rebuild_flag - return DataLayouts.rebuild(slab(data, 1, 1), Array) + return DataLayouts.rebuild( + slab(data, CartesianIndex(1, 1, 1, 1, 1)), + Array, + ) else - return slab(data, 1, 1) + return slab(data, CartesianIndex(1, 1, 1, 1, 1)) end end @@ -284,8 +256,7 @@ _transformed_type( _representative_slab(data, DA), _representative_slab(local_geometry, DA), _representative_slab(local_weights, DA), - 1, - 1, + CartesianIndex(1, 1, 1, 1, 1), ), ) diff --git a/test/DataLayouts/cuda.jl b/test/DataLayouts/cuda.jl index 7d97158bb7..b954117903 100644 --- a/test/DataLayouts/cuda.jl +++ b/test/DataLayouts/cuda.jl @@ -7,6 +7,7 @@ using ClimaComms using CUDA ClimaComms.@import_required_backends using ClimaCore.DataLayouts +using ClimaCore.DataLayouts: slab_index function knl_copy!(dst, src) i = threadIdx().x @@ -17,7 +18,7 @@ function knl_copy!(dst, src) p_dst = slab(dst, h) p_src = slab(src, h) - @inbounds p_dst[i, j] = p_src[i, j] + @inbounds p_dst[slab_index(i, j)] = p_src[slab_index(i, j)] return nothing end diff --git a/test/DataLayouts/data1d.jl b/test/DataLayouts/data1d.jl index b666635c92..b425bb2965 100644 --- a/test/DataLayouts/data1d.jl +++ b/test/DataLayouts/data1d.jl @@ -7,7 +7,7 @@ using JET using ClimaCore.DataLayouts using StaticArrays -using ClimaCore.DataLayouts: get_struct, set_struct! +using ClimaCore.DataLayouts: get_struct, set_struct!, vindex TestFloatTypes = (Float32, Float64) @@ -21,7 +21,7 @@ TestFloatTypes = (Float32, Float64) @test getfield(data.:1, :array) == @view(array[:, 1:2]) # test tuple assignment - data[1] = (Complex{FT}(-1.0, -2.0), FT(-3.0)) + data[vindex(1)] = (Complex{FT}(-1.0, -2.0), FT(-3.0)) @test array[1, 1] == -1.0 @test array[1, 2] == -2.0 @test array[1, 3] == -3.0 @@ -47,9 +47,9 @@ end Nv = 4 array = zeros(Float64, Nv, 3) data = VF{S, Nv}(array) - @test data[1][2] == zero(Float64) - @test_throws BoundsError data[-1] - @test_throws BoundsError data[5] + @test data[vindex(1)][2] == zero(Float64) + @test_throws BoundsError data[vindex(-1)] + @test_throws BoundsError data[vindex(5)] end @testset "VF type safety" begin @@ -63,11 +63,11 @@ end data = VF{typeof(SA), Nv}(array) ret = begin - data[1] = SA + data[vindex(1)] = SA end @test ret === SA - @test data[1] isa typeof(SA) - @test_throws MethodError data[1] = SB + @test data[vindex(1)] isa typeof(SA) + @test_throws MethodError data[vindex(1)] = SB end @testset "VF broadcasting between 1D data objects and scalars" begin diff --git a/test/DataLayouts/data1dx.jl b/test/DataLayouts/data1dx.jl index 5812bfbe41..a14c70c02e 100644 --- a/test/DataLayouts/data1dx.jl +++ b/test/DataLayouts/data1dx.jl @@ -4,7 +4,7 @@ using Revise; include(joinpath("test", "DataLayouts", "data1dx.jl")) =# using Test using ClimaCore.DataLayouts -import ClimaCore.DataLayouts: VIFH, slab, column, VF, IFH +import ClimaCore.DataLayouts: VIFH, slab, column, VF, IFH, vindex, slab_index @testset "VIFH" begin TestFloatTypes = (Float32, Float64) @@ -28,14 +28,14 @@ import ClimaCore.DataLayouts: VIFH, slab, column, VF, IFH # test tuple assignment on columns val = (Complex{FT}(-1.0, -2.0), FT(-3.0)) - column(data, 1, 1)[1] = val + column(data, 1, 1)[vindex(1)] = val @test array[1, 1, 1, 1] == -1.0 @test array[1, 1, 2, 1] == -2.0 @test array[1, 1, 3, 1] == -3.0 # test value of assing tuple on slab sdata = slab(data, 1, 1) - @test sdata[1] == val + @test sdata[slab_index(1)] == val # sum of all the first field elements @test sum(data.:1) ≈ @@ -68,8 +68,8 @@ end @test_throws BoundsError slab(data, 1, 3) sdata = slab(data, 1, 1) - @test_throws BoundsError sdata[-1] - @test_throws BoundsError sdata[2] + @test_throws BoundsError sdata[slab_index(-1)] + @test_throws BoundsError sdata[slab_index(2)] @test_throws BoundsError column(data, -1, 1) @test_throws BoundsError column(data, -1, 1, 1) @@ -91,13 +91,13 @@ end data = VIFH{typeof(SA), Nv, Ni, Nh}(array) cdata = column(data, 1, 1) - cdata[1] = SA - @test cdata[1] isa typeof(SA) - @test_throws MethodError cdata[1] = SB + cdata[slab_index(1)] = SA + @test cdata[slab_index(1)] isa typeof(SA) + @test_throws MethodError cdata[slab_index(1)] = SB sdata = slab(data, 1, 1) - @test sdata[1] isa typeof(SA) - @test_throws MethodError sdata[1] = SB + @test sdata[slab_index(1)] isa typeof(SA) + @test_throws MethodError sdata[slab_index(1)] = SB end @testset "broadcasting between VIFH data object + scalars" begin diff --git a/test/DataLayouts/data2d.jl b/test/DataLayouts/data2d.jl index e459164045..3399e4d5d0 100644 --- a/test/DataLayouts/data2d.jl +++ b/test/DataLayouts/data2d.jl @@ -5,7 +5,7 @@ using Revise; include(joinpath("test", "DataLayouts", "data2d.jl")) using Test using ClimaCore.DataLayouts using StaticArrays -using ClimaCore.DataLayouts: check_basetype, get_struct, set_struct! +using ClimaCore.DataLayouts: check_basetype, get_struct, set_struct!, slab_index @testset "check_basetype" begin @test_throws Exception check_basetype(Real, Float64) @@ -49,16 +49,16 @@ end data = IJFH{S, 2, Nh}(array) @test getfield(data.:1, :array) == @view(array[:, :, 1:2, :]) data_slab = slab(data, 1) - @test data_slab[2, 1] == + @test data_slab[slab_index(2, 1)] == (Complex(array[2, 1, 1, 1], array[2, 1, 2, 1]), array[2, 1, 3, 1]) - data_slab[2, 1] = (Complex(-1.0, -2.0), -3.0) + data_slab[slab_index(2, 1)] = (Complex(-1.0, -2.0), -3.0) @test array[2, 1, 1, 1] == -1.0 @test array[2, 1, 2, 1] == -2.0 @test array[2, 1, 3, 1] == -3.0 subdata_slab = data_slab.:2 - @test subdata_slab[2, 1] == -3.0 - subdata_slab[2, 1] = -5.0 + @test subdata_slab[slab_index(2, 1)] == -3.0 + subdata_slab[slab_index(2, 1)] = -5.0 @test array[2, 1, 3, 1] == -5.0 @test sum(data.:1) ≈ Complex(sum(array[:, :, 1, :]), sum(array[:, :, 2, :])) atol = @@ -89,10 +89,10 @@ end # 2D Slab boundscheck sdata = slab(data, 1) - @test_throws BoundsError sdata[-1, 1] - @test_throws BoundsError sdata[1, -1] - @test_throws BoundsError sdata[2, 1] - @test_throws BoundsError sdata[1, 2] + @test_throws BoundsError sdata[slab_index(-1, 1)] + @test_throws BoundsError sdata[slab_index(1, -1)] + @test_throws BoundsError sdata[slab_index(2, 1)] + @test_throws BoundsError sdata[slab_index(1, 2)] end @testset "IJFH type safety" begin @@ -107,11 +107,11 @@ end data = IJFH{typeof(SA), Nij, Nh}(array) data_slab = slab(data, 1) ret = begin - data_slab[1, 1] = SA + data_slab[slab_index(1, 1)] = SA end @test ret === SA - @test data_slab[1, 1] isa typeof(SA) - @test_throws MethodError data_slab[1, 1] = SB + @test data_slab[slab_index(1, 1)] isa typeof(SA) + @test_throws MethodError data_slab[slab_index(1, 1)] = SB end @testset "2D slab broadcasting" begin diff --git a/test/DataLayouts/data2dx.jl b/test/DataLayouts/data2dx.jl index 3abbfd5291..c9924f2ae1 100644 --- a/test/DataLayouts/data2dx.jl +++ b/test/DataLayouts/data2dx.jl @@ -4,7 +4,7 @@ using Revise; include(joinpath("test", "DataLayouts", "data2dx.jl")) =# using Test using ClimaCore.DataLayouts -import ClimaCore.DataLayouts: VF, IJFH, VIJFH, slab, column +import ClimaCore.DataLayouts: VF, IJFH, VIJFH, slab, column, slab_index, vindex @testset "VIJFH" begin Nv = 10 # number of vertical levels @@ -29,14 +29,14 @@ import ClimaCore.DataLayouts: VF, IJFH, VIJFH, slab, column # test tuple assignment on columns val = (Complex{FT}(-1.0, -2.0), FT(-3.0)) - column(data, 1, 2, 1)[1] = val + column(data, 1, 2, 1)[vindex(1)] = val @test array[1, 1, 2, 1, 1] == -1.0 @test array[1, 1, 2, 2, 1] == -2.0 @test array[1, 1, 2, 3, 1] == -3.0 # test value of assing tuple on slab sdata = slab(data, 1, 1) - @test sdata[1, 2] == val + @test sdata[slab_index(1, 2)] == val # sum of all the first field elements @test sum(data.:1) ≈ @@ -91,13 +91,13 @@ end data = VIJFH{typeof(SA), Nv, Nij, Nh}(array) cdata = column(data, 1, 2, 1) - cdata[1] = SA - @test cdata[1] isa typeof(SA) - @test_throws MethodError cdata[1] = SB + cdata[vindex(1)] = SA + @test cdata[vindex(1)] isa typeof(SA) + @test_throws MethodError cdata[vindex(1)] = SB sdata = slab(data, 1, 1) - @test sdata[1, 2] isa typeof(SA) - @test_throws MethodError sdata[1] = SB + @test sdata[slab_index(1, 2)] isa typeof(SA) + @test_throws MethodError sdata[slab_index(1)] = SB end @testset "broadcasting between VIJFH data object + scalars" begin diff --git a/test/Limiters/limiter.jl b/test/Limiters/limiter.jl index 8b679420e2..7e03d28c07 100644 --- a/test/Limiters/limiter.jl +++ b/test/Limiters/limiter.jl @@ -15,9 +15,11 @@ using ClimaCore: Limiters, Quadratures using ClimaCore.RecursiveApply +import ClimaCore.DataLayouts: slab_index using ClimaCore: slab using Test +si = slab_index # 2D mesh setup function rectangular_mesh_space( n1, @@ -139,10 +141,10 @@ end (h1, h2, slab(limiter.q_bounds, h1 + n1 * (h2 - 1))) end ClimaComms.allowscalar(device) do - @test all(map(T -> T[3][1].x ≈ 2 * (T[1] - 1), S)) # q_min - @test all(map(T -> T[3][1].y ≈ 3 * (T[2] - 1), S)) # q_min - @test all(map(T -> T[3][2].x ≈ 2 * T[1], S)) # q_max - @test all(map(T -> T[3][2].y ≈ 3 * T[2], S)) # q_max + @test all(map(T -> T[3][si(1)].x ≈ 2 * (T[1] - 1), S)) # q_min + @test all(map(T -> T[3][si(1)].y ≈ 3 * (T[2] - 1), S)) # q_min + @test all(map(T -> T[3][si(2)].x ≈ 2 * T[1], S)) # q_max + @test all(map(T -> T[3][si(2)].y ≈ 3 * T[2], S)) # q_max end Limiters.compute_neighbor_bounds_local!(limiter, ρ) @@ -150,10 +152,10 @@ end (h1, h2, slab(limiter.q_bounds_nbr, h1 + n1 * (h2 - 1))) end ClimaComms.allowscalar(device) do - @test all(map(T -> T[3][1].x ≈ 2 * max(T[1] - 2, 0), SN)) # q_min - @test all(map(T -> T[3][1].y ≈ 3 * max(T[2] - 2, 0), SN)) # q_min - @test all(map(T -> T[3][2].x ≈ 2 * min(T[1] + 1, n1), SN)) # q_max - @test all(map(T -> T[3][2].y ≈ 3 * min(T[2] + 1, n2), SN)) # q_max + @test all(map(T -> T[3][si(1)].x ≈ 2 * max(T[1] - 2, 0), SN)) # q_min + @test all(map(T -> T[3][si(1)].y ≈ 3 * max(T[2] - 2, 0), SN)) # q_min + @test all(map(T -> T[3][si(2)].x ≈ 2 * min(T[1] + 1, n1), SN)) # q_max + @test all(map(T -> T[3][si(2)].y ≈ 3 * min(T[2] + 1, n2), SN)) # q_max end end end @@ -170,8 +172,8 @@ end q_min = (FT(3.2), FT(3.0)) q_max = (FT(5.2), FT(5.0)) q_bounds = DataLayouts.IF{Tuple{FT, FT}, 2}(zeros(FT, 2, 2)) - q_bounds[1] = q_min - q_bounds[2] = q_max + q_bounds[si(1)] = q_min + q_bounds[si(2)] = q_max ρq_new = deepcopy(ρq) @@ -179,8 +181,8 @@ end q_new = RecursiveApply.rdiv.(ρq_new, ρ) for j in 1:5, i in 1:5 - @test q_min[1] <= q_new[i, j][1] <= q_max[1] - @test q_min[2] <= q_new[i, j][2] <= q_max[2] + @test q_min[1] <= q_new[si(i, j)][1] <= q_max[1] + @test q_min[2] <= q_new[si(i, j)][2] <= q_max[2] end @test sum(ρq_new.:1 .* WJ) ≈ sum(ρq.:1 .* WJ) @test sum(ρq_new.:2 .* WJ) ≈ sum(ρq.:2 .* WJ) diff --git a/test/Operators/finitedifference/convergence_column.jl b/test/Operators/finitedifference/convergence_column.jl index 53ed67aee2..d6bcc889f9 100644 --- a/test/Operators/finitedifference/convergence_column.jl +++ b/test/Operators/finitedifference/convergence_column.jl @@ -5,6 +5,7 @@ using ClimaComms ClimaComms.@import_required_backends import ClimaCore: slab, Domains, Meshes, Topologies, Spaces, Fields, Operators import ClimaCore.Domains: Geometry +import ClimaCore.DataLayouts: vindex device = ClimaComms.device() @@ -62,7 +63,7 @@ convergence_rate(err, Δh) = cent_field = operator.(face_field) wcent_field = woperator.(face_J, face_field) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] err[k] = norm(cent_field .- cent_field_exact) werr[k] = norm(wcent_field .- cent_field_exact) end @@ -116,7 +117,7 @@ end face_field = operator.(cent_field) wface_field = woperator.(cent_J, cent_field) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] err[k] = norm(face_field .- face_field_exact) werr[k] = norm(wface_field .- face_field_exact) end @@ -167,7 +168,7 @@ end face_field .= operator.(cent_field) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] err[k] = norm(face_field .- face_field_exact) end conv = convergence_rate(err, Δh) @@ -216,7 +217,7 @@ end cent_field .= operator.(face_field) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] err[k] = norm(cent_field .- cent_field_exact) end conv = convergence_rate(err, Δh) @@ -317,7 +318,7 @@ end curlsinᶠ = curlᶠ.(Geometry.Covariant1Vector.(sin.(centers))) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Errors err_grad_sin_c[k] = norm(gradsinᶜ .- Geometry.WVector.(cos.(centers))) err_div_sin_c[k] = norm(divsinᶜ .- cos.(centers)) @@ -439,7 +440,7 @@ end divf2c = Operators.DivergenceF2C() adv_wc = divf2c.(third_order_fluxsinᶠ) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -491,7 +492,7 @@ end divf2c = Operators.DivergenceF2C() adv_wc = divf2c.(third_order_fluxsinᶠ) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Error err_adv_wc[k] = @@ -554,7 +555,7 @@ end adv_wc = divf2c.(third_order_fluxᶠ.(w, c)) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -610,7 +611,7 @@ end ) adv_wc = divf2c.(third_order_fluxᶠ.(w, c)) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Errors err_adv_wc[k] = norm(adv_wc .- ((cos.(centers)) .^ 2 .- (sin.(centers)) .^ 2)) @@ -666,7 +667,7 @@ end @. divf2c(C * (third_order_fluxsinᶠ - first_order_fluxsinᶠ)) adv_wc = @. divf2c.(first_order_fluxsinᶠ) + corrected_antidiff_flux - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -738,7 +739,7 @@ end adv_wc = @. divf2c.(first_order_fluxᶠ(w, c)) + corrected_antidiff_flux - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -802,7 +803,7 @@ end adv_wc = @. divf2c.(first_order_fluxᶠ(w, c)) + corrected_antidiff_flux - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] # Errors err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -854,7 +855,7 @@ end # Call the advection operator adv = advection(c, f, cs) - Δh[k] = Spaces.local_geometry_data(fs).J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[vindex(1)] err[k] = norm(adv .- cos.(Fields.coordinate_field(cs).z)) end # AdvectionC2C convergence rate diff --git a/test/Spaces/opt_spaces.jl b/test/Spaces/opt_spaces.jl index c21bde5cd9..f58c53f9e7 100644 --- a/test/Spaces/opt_spaces.jl +++ b/test/Spaces/opt_spaces.jl @@ -44,10 +44,10 @@ end else test_n_failures(0, TU.PointSpace, context) test_n_failures(137, TU.SpectralElementSpace1D, context) - test_n_failures(295, TU.SpectralElementSpace2D, context) + test_n_failures(298, TU.SpectralElementSpace2D, context) test_n_failures(118, TU.ColumnCenterFiniteDifferenceSpace, context) test_n_failures(118, TU.ColumnFaceFiniteDifferenceSpace, context) - test_n_failures(301, TU.SphereSpectralElementSpace, context) + test_n_failures(304, TU.SphereSpectralElementSpace, context) test_n_failures(321, TU.CenterExtrudedFiniteDifferenceSpace, context) test_n_failures(321, TU.FaceExtrudedFiniteDifferenceSpace, context) diff --git a/test/Spaces/unit_spaces.jl b/test/Spaces/unit_spaces.jl index 68a8884450..2b1596a9ea 100644 --- a/test/Spaces/unit_spaces.jl +++ b/test/Spaces/unit_spaces.jl @@ -21,7 +21,7 @@ import ClimaCore: DeviceSideContext, DeviceSideDevice -import ClimaCore.DataLayouts: IJFH, VF +import ClimaCore.DataLayouts: IJFH, VF, slab_index on_gpu = ClimaComms.device() isa ClimaComms.CUDADevice @@ -55,23 +55,23 @@ on_gpu = ClimaComms.device() isa ClimaComms.CUDADevice array = parent(Spaces.coordinates_data(space)) @test size(array) == (4, 1, 1) coord_slab = slab(Spaces.coordinates_data(space), 1) - @test coord_slab[1] == Geometry.XPoint{FT}(-3) - @test coord_slab[4] == Geometry.XPoint{FT}(5) + @test coord_slab[slab_index(1)] == Geometry.XPoint{FT}(-3) + @test coord_slab[slab_index(4)] == Geometry.XPoint{FT}(5) local_geometry_slab = slab(Spaces.local_geometry_data(space), 1) dss_weights_slab = slab(space.grid.dss_weights, 1) for i in 1:4 - @test Geometry.components(local_geometry_slab[i].∂x∂ξ) ≈ + @test Geometry.components(local_geometry_slab[slab_index(i)].∂x∂ξ) ≈ @SMatrix [8 / 2] - @test Geometry.components(local_geometry_slab[i].∂ξ∂x) ≈ + @test Geometry.components(local_geometry_slab[slab_index(i)].∂ξ∂x) ≈ @SMatrix [2 / 8] - @test local_geometry_slab[i].J ≈ (8 / 2) - @test local_geometry_slab[i].WJ ≈ (8 / 2) * weights[i] + @test local_geometry_slab[slab_index(i)].J ≈ (8 / 2) + @test local_geometry_slab[slab_index(i)].WJ ≈ (8 / 2) * weights[i] if i in (1, 4) - @test dss_weights_slab[i] ≈ 1 / 2 + @test dss_weights_slab[slab_index(i)] ≈ 1 / 2 else - @test dss_weights_slab[i] ≈ 1 + @test dss_weights_slab[slab_index(i)] ≈ 1 end end @@ -217,10 +217,10 @@ end array = parent(coord_data) @test size(array) == (4, 4, 2, 1) coord_slab = slab(coord_data, 1) - @test coord_slab[1, 1] ≈ Geometry.XYPoint{FT}(-3.0, -2.0) - @test coord_slab[4, 1] ≈ Geometry.XYPoint{FT}(5.0, -2.0) - @test coord_slab[1, 4] ≈ Geometry.XYPoint{FT}(-3.0, 8.0) - @test coord_slab[4, 4] ≈ Geometry.XYPoint{FT}(5.0, 8.0) + @test coord_slab[slab_index(1, 1)] ≈ Geometry.XYPoint{FT}(-3.0, -2.0) + @test coord_slab[slab_index(4, 1)] ≈ Geometry.XYPoint{FT}(5.0, -2.0) + @test coord_slab[slab_index(1, 4)] ≈ Geometry.XYPoint{FT}(-3.0, 8.0) + @test coord_slab[slab_index(4, 4)] ≈ Geometry.XYPoint{FT}(5.0, 8.0) @test Spaces.local_geometry_type(typeof(space)) <: Geometry.LocalGeometry local_geometry_slab = slab(Spaces.local_geometry_data(space), 1) @@ -233,17 +233,17 @@ end end for i in 1:4, j in 1:4 - @test Geometry.components(local_geometry_slab[i, j].∂x∂ξ) ≈ + @test Geometry.components(local_geometry_slab[slab_index(i, j)].∂x∂ξ) ≈ @SMatrix [8/2 0; 0 10/2] - @test Geometry.components(local_geometry_slab[i, j].∂ξ∂x) ≈ + @test Geometry.components(local_geometry_slab[slab_index(i, j)].∂ξ∂x) ≈ @SMatrix [2/8 0; 0 2/10] - @test local_geometry_slab[i, j].J ≈ (10 / 2) * (8 / 2) - @test local_geometry_slab[i, j].WJ ≈ + @test local_geometry_slab[slab_index(i, j)].J ≈ (10 / 2) * (8 / 2) + @test local_geometry_slab[slab_index(i, j)].WJ ≈ (10 / 2) * (8 / 2) * weights[i] * weights[j] if i in (1, 4) - @test dss_weights_slab[i, j] ≈ 1 / 2 + @test dss_weights_slab[slab_index(i, j)] ≈ 1 / 2 else - @test dss_weights_slab[i, j] ≈ 1 + @test dss_weights_slab[slab_index(i, j)] ≈ 1 end end