Add column mesh sketch

src/data/data.jl

@@ -233,9 +233,15 @@ end
 # TODO: should this return a S or a 0-d box containing S?
 #  - perhaps the latter, as then it is mutable?
 
-function column(ijfh::IJFH{S}, i::Integer, j::Integer, h) where {S}
-    get_struct(view(parent(ijfh), i, j, :, h), S)
+# function column(ijfh::IJFH{S}, i::Integer, j::Integer, h) where {S}
+#     get_struct(view(parent(ijfh), i, j, :, h), S)
+# end
+
+# Dummy temporary
+struct ColumnData{A}
+    array::A
 end
+column(array::A) where {A} = ColumnData(array)
 
 @inline function slab(ijfh::IJFH{S, Nij}, h::Integer) where {S, Nij} # k,v are unused
     @boundscheck (1 <= h <= length(ijfh)) || throw(BoundsError(ijfh, (h,)))

src/mesh/column_mesh.jl

@@ -0,0 +1,242 @@
+#####
+##### Column mesh
+#####
+
+export n_cells, n_faces
+export ColumnMesh, coords
+export AbstractDataLocation, CellCent, CellFace
+export All, Ghost, Boundary, Interior
+export VerticalExtrema, VerticalMin, VerticalMax
+export n_hat, binary
+
+"""
+    AbstractDataLocation
+
+Subtypes are used for dispatching
+on the data location
+"""
+abstract type AbstractDataLocation end
+
+""" cell center location """
+struct CellCent <: AbstractDataLocation end
+
+""" cell face location """
+struct CellFace <: AbstractDataLocation end
+
+struct CollocatedCoordinates{DL <: AbstractDataLocation, A}
+    h::A
+    Δh::A
+end
+
+Base.length(c::CollocatedCoordinates) = length(c.h)
+
+function Base.show(io::IO, c::CollocatedCoordinates)
+    print(io, "\n")
+    println(io, "  size(h) = ", length(c))
+    println(io, "  h       = ", c.h)
+    println(io, "  Δh      = ", c.Δh)
+end
+
+"""
+    LocalStencil
+
+A local stencil, used for interpolations or
+finite-differences. This is effectively a
+single row of the global operator.
+"""
+struct LocalStencil{T}
+    A::T
+end
+
+"""
+    ColumnMesh
+
+Column coordinates containing a collocated
+grid at cell centers (`cent`) and cell faces
+(`face`).
+"""
+struct ColumnMesh{A, VecLocalStencil} <: AbstractMesh
+    face::CollocatedCoordinates{CellFace, A}
+    cent::CollocatedCoordinates{CellCent, A}
+    interp_stencil::VecLocalStencil
+    ∇_cent_to_face::VecLocalStencil
+    ∇_face_to_cent::VecLocalStencil
+end
+
+coords(c::ColumnMesh, ::CellCent) = c.cent
+coords(c::ColumnMesh, ::CellFace) = c.face
+
+undertype(mesh::ColumnMesh) = undertype(eltype(mesh.cent.h))
+Base.length(c::ColumnMesh, ::CellCent) = length(c.cent)
+Base.length(c::ColumnMesh, ::CellFace) = length(c.face)
+
+n_cells(cm::ColumnMesh) = length(cm.cent)
+n_faces(cm::ColumnMesh) = length(cm.face)
+
+abstract type VerticalExtrema end
+
+"""
+    VerticalMin
+
+A type for dispatching on the minimum of the vertical domain.
+"""
+struct VerticalMin <: VerticalExtrema end
+
+"""
+    VerticalMax
+
+A type for dispatching on the maximum of the vertical domain.
+"""
+struct VerticalMax <: VerticalExtrema end
+
+"""
+    n_hat(::VerticalExtrema)
+
+The outward normal vector to the boundary
+"""
+n_hat(::VerticalMin) = -1
+n_hat(::VerticalMax) = 1
+
+"""
+    binary(::VerticalExtrema)
+
+Returns 0 for `VerticalMin` and 1 for `VerticalMax`
+"""
+binary(::VerticalMin) = 0
+binary(::VerticalMax) = 1
+
+""" A super-type for dispatching on parts of the domain """
+abstract type DomainDecomp end
+struct All <: DomainDecomp end
+struct Boundary <: DomainDecomp end
+struct Interior <: DomainDecomp end
+struct Ghost <: DomainDecomp end
+
+function pad!(h, n_cells_ghost::Int) # TODO: use ghost cells
+    pushfirst!(h, h[1] - (h[2] - h[1]))
+    push!(h, h[end] + (h[end] - h[end - 1]))
+end
+
+function ColumnMesh(
+    h_min::FT,
+    h_max::FT,
+    n_cells_real::Int;
+    order = OrderOfAccuracy{2}(),
+    warpfun::F = nothing,
+    n_cells_ghost::Int = 1,
+    kwargs...,
+) where {FT, F <: Union{Nothing, Function}}
+    n_faces = n_cells_real + 3 # TODO: use n_cells_ghost
+    n_cells = n_cells_real + 2 # TODO: use n_cells_ghost
+
+    # Construct coordinates from faces
+    if !(warpfun == nothing)
+        h_face = warpfun(h_min, h_max, n_cells_real; kwargs...)
+    else # uniform grid
+        # TODO: use n_cells_ghost
+        h_face = map(
+            i -> h_min + FT(i - 1) * (h_max - h_min) / FT(n_cells_real),
+            1:(n_cells_real + 1),
+        )
+        h_face = collect(h_face) # TODO: use ArrayType
+    end
+    pad!(h_face, n_cells_ghost)
+    @assert length(h_face) == n_cells_real + 3 # TODO: use n_cells_ghost
+    Δh_face = map(i -> h_face[i + 1] - h_face[i], 1:(n_faces - 1))
+    coords_face =
+        CollocatedCoordinates{CellFace, typeof(h_face)}(h_face, Δh_face)
+
+    # Construct cell centers (mid-points between faces):
+    h_cent = map(i -> h_face[i] + Δh_face[i] / 2, 1:n_cells)
+    @assert length(h_cent) == n_cells_real + 2 # TODO: use n_cells_ghost
+    Δh_cent = map(i -> h_cent[i + 1] - h_cent[i], 1:(n_cells - 1))
+    coords_cent =
+        CollocatedCoordinates{CellCent, typeof(h_cent)}(h_cent, Δh_cent)
+
+    if !(all(isfinite.(h_face)) && all(isfinite.(h_cent)))
+        error("ColumnMesh is not finite.")
+    end
+
+    # Construct interpolation and derivative stencils:
+    interp_stencil = interpolation_stencil(order, h_face, h_cent)
+    ∇_cent_to_face = ∇_cent_to_face_stencil(order, h_cent)
+    ∇_face_to_cent = ∇_face_to_cent_stencil(order, h_face)
+
+    return ColumnMesh(
+        coords_face,
+        coords_cent,
+        interp_stencil,
+        ∇_cent_to_face,
+        ∇_face_to_cent,
+    )
+end
+
+#####
+##### Order of accuracy type
+#####
+
+struct OrderOfAccuracy{N} end
+
+#####
+##### Interpolation
+#####
+
+"""
+    interpolation_stencil(::OrderOfAccuracy{2}, h_face::AbstractVector, h_cent::AbstractVector)
+
+A second-order interpolation stencil, used
+to operate on fields that live on cell centers.
+The interpolation of these fields live on cell faces.
+
+!!! note
+    interpolation stencils for `face -> cent` is not needed
+    because cell centers are exactly half-way between faces.
+"""
+function interpolation_stencil(
+    ::OrderOfAccuracy{2},
+    h_face::AbstractVector,
+    h_cent::AbstractVector,
+)
+    n_cells = length(h_cent)
+    map(1:(n_cells - 1)) do i
+        diag = (h_face[i + 1] - h_cent[i]) / (h_cent[i + 1] - h_cent[i])
+        upper_diag = 1 - diag
+        LocalStencil(SVector(diag, upper_diag))
+    end
+end
+
+#####
+##### Derivative operators
+#####
+
+"""
+    ∇_cent_to_face_stencil(::OrderOfAccuracy{2}, Δh_cent::AbstractVector)
+
+A second-order finite difference stencil, used
+to operate on fields that live on cell centers.
+The gradient of these fields live on cell faces.
+"""
+function ∇_cent_to_face_stencil(::OrderOfAccuracy{2}, Δh_cent::AbstractVector)
+    n_cells = length(Δh_cent) + 1
+    map(1:(n_cells - 1)) do i
+        diag = -1 / Δh_cent[i]
+        upper_diag = 1 / Δh_cent[i]
+        LocalStencil(SVector(diag, upper_diag))
+    end
+end
+
+"""
+    ∇_face_to_cent_stencil(::OrderOfAccuracy{2}, Δh_face::AbstractVector)
+
+A second-order finite difference stencil, used
+to operate on fields that live on cell faces.
+The gradient of these fields live on cell centers.
+"""
+function ∇_face_to_cent_stencil(::OrderOfAccuracy{2}, Δh_face::AbstractVector)
+    n_faces = length(Δh_face) + 1
+    map(1:(n_faces - 1)) do i
+        diag = -1 / Δh_face[i]
+        upper_diag = 1 / Δh_face[i]
+        LocalStencil(SVector(diag, upper_diag))
+    end
+end

src/mesh/composite_mesh.jl

@@ -0,0 +1,8 @@
+#####
+##### Composite mesh
+#####
+
+struct Composite3DMesh{M2D, MCOL}
+    mesh2D::M2D
+    mesh_col::MCOL
+end

src/mesh/mesh.jl

@@ -149,4 +149,8 @@ struct LocalGeometry{FT}
     J::FT
 end
 =#
+
+include("column_mesh.jl")
+include("composite_mesh.jl")
+
 end # module

test/mesh.jl

@@ -1,4 +1,5 @@
 using Test
+using LinearAlgebra
 using StaticArrays
 import ClimateMachineCore: slab, Domains, Topologies, Meshes
 import ClimateMachineCore.Geometry: Cartesian2DPoint
@@ -37,3 +38,22 @@ import ClimateMachineCore.Geometry: Cartesian2DPoint
         @test local_geometry_slab[i, j].invM * local_geometry_slab[i, j].WJ ≈ 1
     end
 end
+
+@testset "ColumnMesh" begin
+    for FT in (Float32, Float64)
+        a = FT(0.0)
+        b = FT(1.0)
+        n = 10
+        c = Meshes.ColumnMesh(a, b, n)
+        @test c.cent.Δh[1] ≈ FT(1 / 10)
+        @test c.face.Δh[1] ≈ FT(1 / 10)
+        sprint(show, c)
+        @test c.cent == Meshes.coords(c, Meshes.CellCent())
+        @test c.face == Meshes.coords(c, Meshes.CellFace())
+    end
+    @test Meshes.n_hat(Meshes.VerticalMin()) == -1
+    @test Meshes.binary(Meshes.VerticalMin()) == 0
+
+    @test Meshes.n_hat(Meshes.VerticalMax()) == 1
+    @test Meshes.binary(Meshes.VerticalMax()) == 1
+end