Add space with multiple columns

Author: ph-kev

docs/src/APIs/common_grids_api.md

@@ -12,4 +12,5 @@ CommonGrids.ColumnGrid
 CommonGrids.Box3DGrid
 CommonGrids.SliceXZGrid
 CommonGrids.RectangleXYGrid
+CommonGrids.PointColumnEnsembleGrid
 ```

docs/src/APIs/common_spaces_api.md

@@ -12,4 +12,5 @@ CommonSpaces.ColumnSpace
 CommonSpaces.Box3DSpace
 CommonSpaces.SliceXZSpace
 CommonSpaces.RectangleXYSpace
+CommonSpaces.PointColumnEnsembleSpace
 ```

ext/cuda/adapt.jl

@@ -13,6 +13,19 @@ Adapt.adapt_structure(
     Adapt.adapt(to, grid.face_local_geometry),
 )
 
+# PointCloudGrid has no quadrature_style field; pass nothing so that FD
+# operators on PointColumnEnsembleSpace work on CUDA.
+Adapt.adapt_structure(
+    to::CUDA.KernelAdaptor,
+    grid::Grids.ExtrudedFiniteDifferenceGrid{<:Grids.PointCloudGrid},
+) = Grids.DeviceExtrudedFiniteDifferenceGrid(
+    Adapt.adapt(to, Grids.vertical_topology(grid)),
+    nothing,
+    Adapt.adapt(to, grid.global_geometry),
+    Adapt.adapt(to, grid.center_local_geometry),
+    Adapt.adapt(to, grid.face_local_geometry),
+)
+
 Adapt.adapt_structure(
     to::CUDA.KernelAdaptor,
     grid::Grids.FiniteDifferenceGrid,
@@ -39,6 +52,12 @@ Adapt.adapt_structure(to::CUDA.KernelAdaptor, space::Spaces.PointSpace) =
         Adapt.adapt(to, Spaces.local_geometry_data(space)),
     )
 
+Adapt.adapt_structure(to::CUDA.KernelAdaptor, space::Spaces.PointCloudLevelSpace) =
+    Spaces.PointCloudLevelSpace(
+        ClimaCore.DeviceSideContext(),
+        Adapt.adapt(to, Spaces.local_geometry_data(space)),
+    )
+
 Adapt.adapt_structure(
     to::CUDA.KernelAdaptor,
     topology::Topologies.IntervalTopology,

src/CommonGrids/CommonGrids.jl

@@ -67,7 +67,12 @@ grid = ExtrudedCubedSphereGrid(;
 module CommonGrids
 
 export ExtrudedCubedSphereGrid,
-    CubedSphereGrid, ColumnGrid, Box3DGrid, SliceXZGrid, RectangleXYGrid
+    CubedSphereGrid,
+    ColumnGrid,
+    Box3DGrid,
+    SliceXZGrid,
+    RectangleXYGrid,
+    PointColumnEnsembleGrid
 
 import ClimaComms
 import ..DataLayouts,
@@ -706,4 +711,77 @@ function RectangleXYGrid(
     )
 end
 
+"""
+    PointColumnEnsembleGrid(
+        ::Type{<:AbstractFloat}; # defaults to Float64
+        points::AbstractVector{Geometry.LatLongPoint{FT}},
+        z_elem::Integer,
+        z_min::Real,
+        z_max::Real,
+        radius::Real,
+        device::ClimaComms.AbstractDevice = ClimaComms.device(),
+        context::ClimaComms.AbstractCommsContext = ClimaComms.SingletonCommsContext(device),
+        stretch::Meshes.StretchingRule = Meshes.Uniform(),
+        z_mesh::Meshes.IntervalMesh = DefaultZMesh(FT; z_min, z_max, z_elem, stretch),
+    )
+
+A convenience constructor that builds an
+[`Grids.ExtrudedFiniteDifferenceGrid`](@ref) for N independent columns at
+arbitrary (lat, lon) locations on a sphere, given:
+
+ - `FT` the floating-point type (defaults to `Float64`) [`Float32`, `Float64`],
+ - `points` a vector of `Geometry.LatLongPoint` specifying each column
+   location,
+ - `z_elem` the number of z-points,
+ - `z_min` the domain minimum along the z-direction,
+ - `z_max` the domain maximum along the z-direction,
+ - `radius` the radius of the sphere,
+ - `device` the `ClimaComms.device`,
+ - `context` the `ClimaComms.context` (must be a `SingletonCommsContext`),
+ - `stretch` the mesh `Meshes.StretchingRule` (defaults to
+   [`Meshes.Uniform`](@ref)),
+ - `z_mesh` the vertical mesh, defaults to an `Meshes.IntervalMesh` along `z`
+   with given `stretch`.
+
+There is no horizontal connectivity between columns. Horizontal operators are
+not supported. Use [`ClimaCore.Fields.bycolumn`](@ref) to iterate over columns.
+
+# Example usage
+
+```julia
+using ClimaCore.CommonGrids, ClimaCore.Geometry
+points = [LatLongPoint(0.0, 0.0), LatLongPoint(10.0, 20.0), LatLongPoint(-5.0, 90.0)]
+grid = PointColumnEnsembleGrid(;
+    points  = points,
+    z_elem  = 10,
+    z_min   = 0,
+    z_max   = 10_000,
+    radius  = 6.371229e6,
+)
+```
+"""
+PointColumnEnsembleGrid(; kwargs...) = PointColumnEnsembleGrid(Float64; kwargs...)
+function PointColumnEnsembleGrid(
+    ::Type{FT};
+    points::AbstractVector{Geometry.LatLongPoint{FT}},
+    z_elem::Integer,
+    z_min::Real,
+    z_max::Real,
+    radius::Real = 6.371229e6,
+    device::ClimaComms.AbstractDevice = ClimaComms.device(),
+    context::ClimaComms.AbstractCommsContext = ClimaComms.SingletonCommsContext(device),
+    stretch::Meshes.StretchingRule = Meshes.Uniform(),
+    z_mesh::Meshes.IntervalMesh = DefaultZMesh(FT; z_min, z_max, z_elem, stretch),
+) where {FT}
+    @assert context isa ClimaComms.SingletonCommsContext "PointColumnEnsembleGrid only supports SingletonCommsContext."
+    @assert ClimaComms.device(context) == device "The given device and context device do not match."
+    h_grid = Grids.PointCloudGrid(points; radius, device, context)
+    z_topology = Topologies.IntervalTopology(
+        ClimaComms.SingletonCommsContext(device),
+        z_mesh,
+    )
+    z_grid = Grids.FiniteDifferenceGrid(z_topology)
+    return Grids.ExtrudedFiniteDifferenceGrid(h_grid, z_grid)
+end
+
 end # module

src/CommonSpaces/CommonSpaces.jl

@@ -13,6 +13,7 @@ export ExtrudedCubedSphereSpace,
     Box3DSpace,
     SliceXZSpace,
     RectangleXYSpace,
+    PointColumnEnsembleSpace,
     CellCenter,
     CellFace,
     face_space,
@@ -32,7 +33,8 @@ import ..CommonGrids:
     ColumnGrid,
     Box3DGrid,
     SliceXZGrid,
-    RectangleXYGrid
+    RectangleXYGrid,
+    PointColumnEnsembleGrid
 import ..Spaces: face_space, center_space
 
 
@@ -426,4 +428,58 @@ RectangleXYSpace(; kwargs...) = RectangleXYSpace(Float64; kwargs...)
 RectangleXYSpace(::Type{FT}; kwargs...) where {FT} =
     Spaces.SpectralElementSpace2D(RectangleXYGrid(FT; kwargs...))
 
+"""
+	PointColumnEnsembleSpace(
+        ::Type{<:AbstractFloat}; # defaults to Float64
+        points::AbstractVector{Geometry.LatLongPoint{FT}},
+        z_elem::Integer,
+        z_min::Real,
+        z_max::Real,
+        device::ClimaComms.AbstractDevice = ClimaComms.device(),
+        context::ClimaComms.AbstractCommsContext = ClimaComms.SingletonCommsContext(device),
+        stretch::Meshes.StretchingRule = Meshes.Uniform(),
+        z_mesh::Meshes.IntervalMesh = DefaultZMesh(FT; z_min, z_max, z_elem, stretch),
+        staggering::Staggering,
+    )
+
+Construct a [`Spaces.ExtrudedFiniteDifferenceSpace`](@ref) (aliased as
+`Spaces.PointCloudSpace`) for N independent columns at arbitrary (lat, lon)
+locations on a sphere, given:
+
+ - `FT` the floating-point type (defaults to `Float64`) [`Float32`, `Float64`]
+ - `points` a vector of `Geometry.LatLongPoint` specifying each column location
+ - `z_elem` the number of z-points
+ - `z_min` the domain minimum along the z-direction
+ - `z_max` the domain maximum along the z-direction
+ - `device` the `ClimaComms.device`
+ - `context` the `ClimaComms.context` (must be a `SingletonCommsContext`)
+ - `stretch` the mesh `Meshes.StretchingRule` (defaults to [`Meshes.Uniform`](@ref))
+ - `z_mesh` the vertical mesh, defaults to an `Meshes.IntervalMesh` along `z` with given `stretch`
+ - `staggering` vertical staggering, can be one of [[`Grids.CellFace`](@ref), [`Grids.CellCenter`](@ref)]
+
+Note that these arguments are all the same as [`CommonGrids.PointColumnEnsembleGrid`](@ref),
+except for `staggering`.
+
+# Example usage
+
+```julia
+using ClimaCore.CommonSpaces, ClimaCore.Geometry
+points = [LatLongPoint(0.0, 0.0), LatLongPoint(10.0, 20.0), LatLongPoint(-5.0, 90.0)]
+space = PointColumnEnsembleSpace(;
+    points  = points,
+    z_elem  = 10,
+    z_min   = 0,
+    z_max   = 10_000,
+    staggering = CellCenter()
+)
+```
+"""
+function PointColumnEnsembleSpace end
+PointColumnEnsembleSpace(; kwargs...) = PointColumnEnsembleSpace(Float64; kwargs...)
+PointColumnEnsembleSpace(::Type{FT}; staggering::Staggering, kwargs...) where {FT} =
+    Spaces.MultiColumnFiniteDifferenceSpace(
+        PointColumnEnsembleGrid(FT; kwargs...),
+        staggering,
+    )
+
 end # module

src/Fields/Fields.jl

@@ -116,6 +116,22 @@ const CenterExtrudedFiniteDifferenceField{V, S} = Field{
     S,
 } where {V <: AbstractData, S <: Spaces.CenterExtrudedFiniteDifferenceSpace}
 
+const MultiColumnFiniteDifferenceField{V, S} = Field{
+    V,
+    S,
+} where {V <: AbstractData, S <: Spaces.MultiColumnFiniteDifferenceSpace}
+const FaceMultiColumnFiniteDifferenceField{V, S} = Field{
+    V,
+    S,
+} where {V <: AbstractData, S <: Spaces.FaceMultiColumnFiniteDifferenceSpace}
+const CenterMultiColumnFiniteDifferenceField{V, S} = Field{
+    V,
+    S,
+} where {
+    V <: AbstractData,
+    S <: Spaces.CenterMultiColumnFiniteDifferenceSpace,
+}
+
 #
 const SpectralElementField1D{V, S} =
     Field{V, S} where {V <: AbstractData, S <: Spaces.SpectralElementSpace1D}
@@ -497,6 +513,7 @@ Base.@propagate_inbounds function level(
     field::Union{
         CenterFiniteDifferenceField,
         CenterExtrudedFiniteDifferenceField,
+        CenterMultiColumnFiniteDifferenceField,
     },
     v::Int,
 )
@@ -505,7 +522,11 @@ Base.@propagate_inbounds function level(
     Field(data, hspace)
 end
 Base.@propagate_inbounds function level(
-    field::Union{FaceFiniteDifferenceField, FaceExtrudedFiniteDifferenceField},
+    field::Union{
+        FaceFiniteDifferenceField,
+        FaceExtrudedFiniteDifferenceField,
+        FaceMultiColumnFiniteDifferenceField,
+    },
     v::PlusHalf,
 )
     hspace = level(axes(field), v)

src/Fields/indices.jl

@@ -103,6 +103,51 @@ bycolumn(
     device::ClimaComms.AbstractCPUDevice,
 ) = bycolumn(fn, Spaces.horizontal_space(space), device)
 
+function bycolumn(
+    fn,
+    space::Spaces.PointCloudSpace,
+    ::ClimaComms.CPUSingleThreaded,
+)
+    N = Spaces.ncolumns(space)
+    @inbounds for h in 1:N
+        fn(ColumnIndex((1,), h))
+    end
+    return nothing
+end
+function bycolumn(
+    fn,
+    space::Spaces.PointCloudSpace,
+    ::ClimaComms.CPUMultiThreaded,
+)
+    N = Spaces.ncolumns(space)
+    @inbounds Threads.@threads for h in 1:N
+        fn(ColumnIndex((1,), h))
+    end
+    return nothing
+end
+
+function bycolumn(
+    fn,
+    space::Spaces.MultiColumnFiniteDifferenceSpace,
+    ::ClimaComms.CPUSingleThreaded,
+)
+    N = Spaces.ncolumns(space)
+    @inbounds for h in 1:N
+        fn(ColumnIndex((1,), h))
+    end
+    return nothing
+end
+function bycolumn(
+    fn,
+    space::Spaces.MultiColumnFiniteDifferenceSpace,
+    ::ClimaComms.CPUMultiThreaded,
+)
+    N = Spaces.ncolumns(space)
+    @inbounds Threads.@threads for h in 1:N
+        fn(ColumnIndex((1,), h))
+    end
+    return nothing
+end
 
 
 """

src/Grids/Grids.jl

@@ -71,6 +71,7 @@ Meshes.domain(grid::AbstractGrid) = Meshes.domain(topology(grid))
 
 include("finitedifference.jl")
 include("spectralelement.jl")
+include("pointcloud.jl")
 include("extruded.jl")
 include("column.jl")
 include("level.jl")
@@ -106,6 +107,7 @@ has_horizontal(::ExtrudedFiniteDifferenceGrid) = true
 has_horizontal(::DeviceSpectralElementGrid2D) = true
 has_horizontal(::SpectralElementGrid2D) = true
 has_horizontal(::SpectralElementGrid1D) = true
+has_horizontal(::PointCloudGrid) = true
 
 """
     has_vertical(::AbstractGrid)

src/Grids/extruded.jl

@@ -45,7 +45,7 @@ local_geometry_type(
 ) where {H, V, A, GG, CLG, FLG} = eltype(CLG) # calls eltype from DataLayouts
 
 function ExtrudedFiniteDifferenceGrid(
-    horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D},
+    horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D, PointCloudGrid},
     vertical_grid::FiniteDifferenceGrid,
     hypsography::HypsographyAdaption = Flat();
     deep = false,
@@ -70,7 +70,7 @@ end
 
 # memoized constructor
 function ExtrudedFiniteDifferenceGrid(
-    horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D},
+    horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D, PointCloudGrid},
     vertical_grid::FiniteDifferenceGrid,
     hypsography::HypsographyAdaption,
     global_geometry::Geometry.AbstractGlobalGeometry,
@@ -96,7 +96,7 @@ end
 
 # Non-memoized constructor. Should not generally be called, but can be defined for other Hypsography types
 function _ExtrudedFiniteDifferenceGrid(
-    horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D},
+    horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D, PointCloudGrid},
     vertical_grid::FiniteDifferenceGrid,
     hypsography::Flat,
     global_geometry::Geometry.AbstractGlobalGeometry,
@@ -126,6 +126,13 @@ end
 
 topology(grid::ExtrudedFiniteDifferenceGrid) = topology(grid.horizontal_grid)
 
+# For PointCloudGrid there is no horizontal topology; route context/device
+# through the horizontal grid directly (which stores its own context).
+ClimaComms.context(grid::ExtrudedFiniteDifferenceGrid{<:PointCloudGrid}) =
+    ClimaComms.context(grid.horizontal_grid)
+ClimaComms.device(grid::ExtrudedFiniteDifferenceGrid{<:PointCloudGrid}) =
+    ClimaComms.device(grid.horizontal_grid)
+
 vertical_topology(grid::ExtrudedFiniteDifferenceGrid) =
     topology(grid.vertical_grid)
 

src/Grids/level.jl

@@ -15,6 +15,9 @@ level(
 
 topology(levelgrid::LevelGrid) = topology(levelgrid.full_grid)
 
+ClimaComms.context(levelgrid::LevelGrid) = ClimaComms.context(levelgrid.full_grid)
+ClimaComms.device(levelgrid::LevelGrid) = ClimaComms.device(levelgrid.full_grid)
+
 # The DSS weights for extruded spaces are currently the same as the weights for
 # horizontal spaces. If we ever need to use extruded weights, this method will
 # need to extract the weights at a particular level.