reduce parameter size of finite difference broadcasted objects

Author: simonbyrne

examples/column/hydrostatic_ekman.jl

@@ -131,7 +131,7 @@ function tendency!(dY, Y, _, t)
     # TODO!: Undesirable casting to vector required
     @. dρθ = -∂c(w * If(ρθ)) + ρ * ∂c(ν * ∂f(ρθ / ρ))
 
-    uv_1 = Operators.getidx(uv, Operators.Interior(), 1)
+    uv_1 = Operators.getidx(axes(uv), uv, Operators.Interior(), 1)
     u_wind = LinearAlgebra.norm(uv_1)
 
     A = Operators.AdvectionC2C(

src/DataLayouts/DataLayouts.jl

@@ -491,6 +491,10 @@ 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!(
         parent(data),
@@ -499,6 +503,19 @@ end
         CartesianIndex(i, 1, h),
     )
 end
+@inline function Base.setindex!(
+    data::IFH{S},
+    val,
+    I::CartesianIndex{5},
+) where {S}
+    i, _, _, _, h = I.I
+    @inbounds set_struct!(
+        parent(data),
+        convert(S, val),
+        Val(3),
+        CartesianIndex(i, 1, h),
+    )
+end
 
 # ======================
 # Data0D DataLayout
@@ -1291,6 +1308,19 @@ end
         CartesianIndex(v, i, 1, h),
     )
 end
+@inline function Base.setindex!(
+    data::VIFH{S},
+    val,
+    I::CartesianIndex{5},
+) where {S}
+    i, _, _, v, h = I.I
+    @inbounds set_struct!(
+        parent(data),
+        convert(S, val),
+        Val(3),
+        CartesianIndex(v, i, 1, h),
+    )
+end
 
 # =========================================
 # Special DataLayouts for regular gridding

src/Operators/common.jl

@@ -69,13 +69,50 @@ end
 
 # Functions for CUDASpectralStyle
 struct PlaceholderSpace <: Spaces.AbstractSpace end
+struct CenterPlaceholderSpace <: Spaces.AbstractSpace end
+struct FacePlaceholderSpace <: Spaces.AbstractSpace end
+
 
 placeholder_space(current_space::T, parent_space::T) where {T} =
     PlaceholderSpace()
 placeholder_space(current_space, parent_space) = current_space
-
-reconstruct_placeholder_space(::PlaceholderSpace, parent_space) = parent_space
-reconstruct_placeholder_space(current_space, parent_space) = current_space
+placeholder_space(
+    current_space::Spaces.CenterFiniteDifferenceSpace,
+    parent_space::Spaces.FaceFiniteDifferenceSpace,
+) = CenterPlaceholderSpace()
+placeholder_space(
+    current_space::Spaces.CenterExtrudedFiniteDifferenceSpace,
+    parent_space::Spaces.FaceExtrudedFiniteDifferenceSpace,
+) = CenterPlaceholderSpace()
+placeholder_space(
+    current_space::Spaces.FaceFiniteDifferenceSpace,
+    parent_space::Spaces.CenterFiniteDifferenceSpace,
+) = FacePlaceholderSpace()
+placeholder_space(
+    current_space::Spaces.FaceExtrudedFiniteDifferenceSpace,
+    parent_space::Spaces.CenterExtrudedFiniteDifferenceSpace,
+) = FacePlaceholderSpace()
+
+@inline reconstruct_placeholder_space(::PlaceholderSpace, parent_space) =
+    parent_space
+@inline reconstruct_placeholder_space(
+    ::CenterPlaceholderSpace,
+    parent_space::Spaces.FaceFiniteDifferenceSpace,
+) = Spaces.CenterFiniteDifferenceSpace(parent_space)
+@inline reconstruct_placeholder_space(
+    ::CenterPlaceholderSpace,
+    parent_space::Spaces.FaceExtrudedFiniteDifferenceSpace,
+) = Spaces.CenterExtrudedFiniteDifferenceSpace(parent_space)
+@inline reconstruct_placeholder_space(
+    ::FacePlaceholderSpace,
+    parent_space::Spaces.CenterFiniteDifferenceSpace,
+) = Spaces.FaceFiniteDifferenceSpace(parent_space)
+@inline reconstruct_placeholder_space(
+    ::FacePlaceholderSpace,
+    parent_space::Spaces.CenterExtrudedFiniteDifferenceSpace,
+) = Spaces.FaceExtrudedFiniteDifferenceSpace(parent_space)
+@inline reconstruct_placeholder_space(current_space, parent_space) =
+    current_space
 
 
 strip_space(obj, parent_space) = obj