diff --git a/docs/src/APIs/common_grids_api.md b/docs/src/APIs/common_grids_api.md index 8cf57907f5..cdc8b1d486 100644 --- a/docs/src/APIs/common_grids_api.md +++ b/docs/src/APIs/common_grids_api.md @@ -12,4 +12,5 @@ CommonGrids.ColumnGrid CommonGrids.Box3DGrid CommonGrids.SliceXZGrid CommonGrids.RectangleXYGrid +CommonGrids.PointColumnEnsembleGrid ``` diff --git a/docs/src/APIs/common_spaces_api.md b/docs/src/APIs/common_spaces_api.md index bbcb4ca87a..7726c56136 100644 --- a/docs/src/APIs/common_spaces_api.md +++ b/docs/src/APIs/common_spaces_api.md @@ -12,4 +12,5 @@ CommonSpaces.ColumnSpace CommonSpaces.Box3DSpace CommonSpaces.SliceXZSpace CommonSpaces.RectangleXYSpace +CommonSpaces.PointColumnEnsembleSpace ``` diff --git a/ext/cuda/adapt.jl b/ext/cuda/adapt.jl index 14c0b645db..225e963ebd 100644 --- a/ext/cuda/adapt.jl +++ b/ext/cuda/adapt.jl @@ -7,6 +7,7 @@ Adapt.adapt_structure( grid::Grids.ExtrudedFiniteDifferenceGrid, ) = Grids.DeviceExtrudedFiniteDifferenceGrid( Adapt.adapt(to, Grids.vertical_topology(grid)), + grid.horizontal_grid isa Grids.PointCloudGrid ? nothing : Adapt.adapt(to, grid.horizontal_grid.quadrature_style), Adapt.adapt(to, grid.global_geometry), Adapt.adapt(to, grid.center_local_geometry), @@ -39,6 +40,13 @@ 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.PointCloudSpace) = + Spaces.PointCloudSpace( + ClimaCore.DeviceSideContext(), + Adapt.adapt(to, space.full_grid), + Adapt.adapt(to, Spaces.local_geometry_data(space)), + ) + Adapt.adapt_structure( to::CUDA.KernelAdaptor, topology::Topologies.IntervalTopology, diff --git a/src/CommonGrids/CommonGrids.jl b/src/CommonGrids/CommonGrids.jl index f5784f3dd6..5c285f218b 100644 --- a/src/CommonGrids/CommonGrids.jl +++ b/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,72 @@ 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(), + 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`, + - `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(), + stretch::Meshes.StretchingRule = Meshes.Uniform(), + z_mesh::Meshes.IntervalMesh = DefaultZMesh(FT; z_min, z_max, z_elem, stretch), +) where {FT} + h_grid = Grids.PointCloudGrid(points; radius, device) + 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 diff --git a/src/CommonSpaces/CommonSpaces.jl b/src/CommonSpaces/CommonSpaces.jl index 15ebb9048b..df87d76f66 100644 --- a/src/CommonSpaces/CommonSpaces.jl +++ b/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,56 @@ 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(), + 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.MultiColumnFiniteDifferenceSpace`) 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` + - `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 diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index efeb824844..1c75891d9c 100644 --- a/src/Fields/Fields.jl +++ b/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) diff --git a/src/Fields/indices.jl b/src/Fields/indices.jl index abf11e540d..b9471dca27 100644 --- a/src/Fields/indices.jl +++ b/src/Fields/indices.jl @@ -103,6 +103,28 @@ bycolumn( device::ClimaComms.AbstractCPUDevice, ) = bycolumn(fn, Spaces.horizontal_space(space), device) +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 """ diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index 55bb705d4d..22759fe63d 100644 --- a/src/Grids/Grids.jl +++ b/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) @@ -124,6 +126,7 @@ Retrieve the mask for the grid (defaults to DataLayouts.NoMask). """ get_mask(::AbstractGrid) = DataLayouts.NoMask() get_mask(grid::ExtrudedFiniteDifferenceGrid) = grid.horizontal_grid.mask +get_mask(::ExtrudedFiniteDifferenceGrid{<:PointCloudGrid}) = DataLayouts.NoMask() """ set_mask!(fn::Function, grid) diff --git a/src/Grids/extruded.jl b/src/Grids/extruded.jl index 7e57f3ee39..4c97b1a7a1 100644 --- a/src/Grids/extruded.jl +++ b/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::AbstractSpectralElementGrid, vertical_grid::FiniteDifferenceGrid, hypsography::HypsographyAdaption = Flat(); deep = false, @@ -70,7 +70,7 @@ end # memoized constructor function ExtrudedFiniteDifferenceGrid( - horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D}, + horizontal_grid::AbstractSpectralElementGrid, 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::AbstractSpectralElementGrid, vertical_grid::FiniteDifferenceGrid, hypsography::Flat, global_geometry::Geometry.AbstractGlobalGeometry, @@ -126,6 +126,11 @@ end topology(grid::ExtrudedFiniteDifferenceGrid) = topology(grid.horizontal_grid) +ClimaComms.context(grid::ExtrudedFiniteDifferenceGrid) = + ClimaComms.context(grid.horizontal_grid) +ClimaComms.device(grid::ExtrudedFiniteDifferenceGrid) = + ClimaComms.device(grid.horizontal_grid) + vertical_topology(grid::ExtrudedFiniteDifferenceGrid) = topology(grid.vertical_grid) diff --git a/src/Grids/level.jl b/src/Grids/level.jl index ba187bed15..470493a00c 100644 --- a/src/Grids/level.jl +++ b/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. diff --git a/src/Grids/pointcloud.jl b/src/Grids/pointcloud.jl new file mode 100644 index 0000000000..ac479649ee --- /dev/null +++ b/src/Grids/pointcloud.jl @@ -0,0 +1,121 @@ + +""" + PointCloudGrid( + context :: ClimaComms.AbstractCommsContext, + points :: AbstractVector{Geometry.LatLongPoint{FT}}, + ) + +A horizontal grid consisting of N arbitrary, disconnected (lat, long) locations +on a sphere. There is no connectivity between columns; no spectral element +basis, DSS, or horizontal operators are supported on this grid. + +This is the horizontal component used by a "point cloud" extruded space (N +independent columns at user-chosen sphere locations). + +The `local_geometry` is stored as an `IFH{LG, 1, N}` data layout, with each of +the `N` locations represented by an element with one nodal point. Based on the +[metric +tensor](https://en.wikipedia.org/wiki/Metric_tensor#The_round_metric_on_a_sphere) +of a sphere, the horizontal Jacobian `∂x∂ξ` is given by the diagonal matrix +`diag(R·π/180, R·cosd(lat)·π/180)`, with the determinant `J = +R²·cosd(lat)·(π/180)²`. +""" +struct PointCloudGrid{ + C <: ClimaComms.AbstractCommsContext, + GG <: Geometry.AbstractGlobalGeometry, + LG, +} <: AbstractSpectralElementGrid + context::C + global_geometry::GG + local_geometry::LG # IFH{LocalGeometry{(1,2), LatLongPoint{FT}, FT, Padded∂x∂ξ{FT}, PaddedContravariantMetric{FT}}, 1, N} +end + +Adapt.@adapt_structure PointCloudGrid + +local_geometry_type(::Type{PointCloudGrid{C, GG, LG}}) where {C, GG, LG} = + eltype(LG) + +ClimaComms.context(grid::PointCloudGrid) = grid.context +ClimaComms.device(grid::PointCloudGrid) = ClimaComms.device(grid.context) + +topology(::PointCloudGrid) = error( + "PointCloudGrid has no topology", +) + +local_geometry_data(grid::PointCloudGrid, ::Nothing) = grid.local_geometry + +global_geometry(grid::PointCloudGrid) = grid.global_geometry + +quadrature_style(::PointCloudGrid) = + error("PointCloudGrid has no quadrature style (no spectral element basis)") + +""" + PointCloudGrid( + points :: AbstractVector{Geometry.LatLongPoint{FT}}; + radius :: Real, + device :: ClimaComms.AbstractDevice = ClimaComms.device(), + ) + +Convenience constructor: build a `PointCloudGrid` from a vector of +`LatLongPoint`s and a sphere `radius`. The horizontal metric terms in +`local_geometry` are set from the sphere geometry at each point: +`∂x∂ξ = diag(R·π/180, R·cosd(lat)·π/180)`, `J = R²·cosd(lat)·(π/180)²`. +""" +function PointCloudGrid( + points::AbstractVector{Geometry.LatLongPoint{FT}}; + radius::Real, + device::ClimaComms.AbstractDevice = ClimaComms.device(), +) where {FT} + # PointCloudGrid is single-process only; the context is always a + # SingletonCommsContext built from the given device. + context = ClimaComms.SingletonCommsContext(device) + + N = length(points) + global_geometry = Geometry.SphericalGlobalGeometry(FT(radius)) + + AIdx = Geometry.coordinate_axis(Geometry.LatLongPoint{FT}) # (1, 2) + LG = Geometry.LocalGeometryType(Geometry.LatLongPoint{FT}, FT, AIdx) + + # Ni = 1 (one node per "column element"), Nh = N + local_geometry = DataLayouts.IFH{LG, 1}(Array{FT}, N) + + ∂x∂ξ_axes = ( + Geometry.LocalAxis{AIdx}(), + Geometry.CovariantAxis{AIdx}(), + ) + deg2rad = FT(π) / 180 + + for (h, pt) in enumerate(points) + # Sphere metric: arc length per degree in each coordinate direction. + # ∂x∂ξ is diagonal: (R·π/180) in the lat direction, + # (R·cosd(lat)·π/180) in the lon direction. + s_lat = FT(radius) * deg2rad + s_lon = FT(radius) * deg2rad * cosd(pt.lat) + J = s_lat * s_lon # det of diagonal Jacobian + ∂x∂ξ_mat = SMatrix{2, 2, FT, 4}(zero(FT), s_lat, s_lon, zero(FT)) + + lg_slab = slab(local_geometry, h) + lg_slab[slab_index(1)] = Geometry.LocalGeometry( + pt, + J, + J, # WJ — unit quadrature weight × J + Geometry.AxisTensor(∂x∂ξ_axes, ∂x∂ξ_mat), + ) + end + + DA = ClimaComms.array_type(device) + return PointCloudGrid( + context, + global_geometry, + DataLayouts.rebuild(local_geometry, DA), + ) +end + +function Base.show(io::IO, grid::PointCloudGrid) + indent = get(io, :indent, 0) + iio = IOContext(io, :indent => indent + 2) + println(io, nameof(typeof(grid)), ":") + # some reduced spaces (like slab space) do not have topology + println(iio, " "^(indent + 2), "horizontal:") + print(iio, " "^(indent + 4), "context: Nothing") +end diff --git a/src/MatrixFields/matrix_shape.jl b/src/MatrixFields/matrix_shape.jl index ab3c076270..483e27822f 100644 --- a/src/MatrixFields/matrix_shape.jl +++ b/src/MatrixFields/matrix_shape.jl @@ -27,7 +27,11 @@ matrix_shape(matrix_field, matrix_space) = _matrix_shape( function matrix_shape( matrix_field, - matrix_space::Union{Spaces.AbstractSpectralElementSpace, Spaces.PointSpace}, + matrix_space::Union{ + Spaces.AbstractSpectralElementSpace, + Spaces.PointSpace, + Spaces.PointCloudSpace, + }, ) @assert eltype(matrix_field) <: DiagonalMatrixRow Square() diff --git a/src/MatrixFields/operator_matrices.jl b/src/MatrixFields/operator_matrices.jl index f84d1ee0b7..e1001a5ef2 100644 --- a/src/MatrixFields/operator_matrices.jl +++ b/src/MatrixFields/operator_matrices.jl @@ -62,6 +62,15 @@ operator_input_space( space::Spaces.ExtrudedFiniteDifferenceSpace, ) = Spaces.FaceExtrudedFiniteDifferenceSpace(space) +operator_input_space( + ::FDOperatorWithCenterInput, + space::Spaces.MultiColumnFiniteDifferenceSpace, +) = Spaces.CenterMultiColumnFiniteDifferenceSpace(space) +operator_input_space( + ::FDOperatorWithFaceInput, + space::Spaces.MultiColumnFiniteDifferenceSpace, +) = Spaces.FaceMultiColumnFiniteDifferenceSpace(space) + has_affine_bc(op) = unrolled_any( bc -> bc isa Union{ diff --git a/src/Operators/common.jl b/src/Operators/common.jl index 98f76b94f8..b3bc35725b 100644 --- a/src/Operators/common.jl +++ b/src/Operators/common.jl @@ -85,6 +85,10 @@ placeholder_space( current_space::Spaces.CenterExtrudedFiniteDifferenceSpace, parent_space::Spaces.FaceExtrudedFiniteDifferenceSpace, ) = CenterPlaceholderSpace() +placeholder_space( + current_space::Spaces.CenterMultiColumnFiniteDifferenceSpace, + parent_space::Spaces.FaceMultiColumnFiniteDifferenceSpace, +) = CenterPlaceholderSpace() placeholder_space( current_space::Spaces.FaceFiniteDifferenceSpace, parent_space::Spaces.CenterFiniteDifferenceSpace, @@ -93,6 +97,10 @@ placeholder_space( current_space::Spaces.FaceExtrudedFiniteDifferenceSpace, parent_space::Spaces.CenterExtrudedFiniteDifferenceSpace, ) = FacePlaceholderSpace() +placeholder_space( + current_space::Spaces.FaceMultiColumnFiniteDifferenceSpace, + parent_space::Spaces.CenterMultiColumnFiniteDifferenceSpace, +) = FacePlaceholderSpace() @inline reconstruct_placeholder_space(current_space, parent_space) = current_space @@ -131,6 +139,23 @@ placeholder_space( }, ) = parent_space +@inline reconstruct_placeholder_space( + ::CenterPlaceholderSpace, + parent_space::Spaces.FaceMultiColumnFiniteDifferenceSpace, +) = Spaces.CenterMultiColumnFiniteDifferenceSpace(parent_space) +@inline reconstruct_placeholder_space( + ::FacePlaceholderSpace, + parent_space::Spaces.CenterMultiColumnFiniteDifferenceSpace, +) = Spaces.FaceMultiColumnFiniteDifferenceSpace(parent_space) +@inline reconstruct_placeholder_space( + ::FacePlaceholderSpace, + parent_space::Spaces.FaceMultiColumnFiniteDifferenceSpace, +) = parent_space +@inline reconstruct_placeholder_space( + ::CenterPlaceholderSpace, + parent_space::Spaces.CenterMultiColumnFiniteDifferenceSpace, +) = parent_space + strip_space(obj, parent_space) = obj function strip_space(field::Field, parent_space) diff --git a/src/Operators/finitedifference.jl b/src/Operators/finitedifference.jl index 5896625d6b..c8bf319a66 100644 --- a/src/Operators/finitedifference.jl +++ b/src/Operators/finitedifference.jl @@ -2,15 +2,20 @@ import ..Utilities: PlusHalf, half, unionall_type import ..DebugOnly: allow_mismatched_spaces_unsafe import UnrolledUtilities: unrolled_map -const AllFiniteDifferenceSpace = - Union{Spaces.FiniteDifferenceSpace, Spaces.ExtrudedFiniteDifferenceSpace} +const AllFiniteDifferenceSpace = Union{ + Spaces.FiniteDifferenceSpace, + Spaces.ExtrudedFiniteDifferenceSpace, + Spaces.MultiColumnFiniteDifferenceSpace, +} const AllFaceFiniteDifferenceSpace = Union{ Spaces.FaceFiniteDifferenceSpace, Spaces.FaceExtrudedFiniteDifferenceSpace, + Spaces.FaceMultiColumnFiniteDifferenceSpace, } const AllCenterFiniteDifferenceSpace = Union{ Spaces.CenterFiniteDifferenceSpace, Spaces.CenterExtrudedFiniteDifferenceSpace, + Spaces.CenterMultiColumnFiniteDifferenceSpace, } Topologies.isperiodic(space::AllFiniteDifferenceSpace) = diff --git a/src/Spaces/Spaces.jl b/src/Spaces/Spaces.jl index 5790e907a6..8466f71cce 100644 --- a/src/Spaces/Spaces.jl +++ b/src/Spaces/Spaces.jl @@ -99,6 +99,7 @@ include("pointspace.jl") include("spectralelement.jl") include("finitedifference.jl") include("extruded.jl") +include("multicolumn.jl") include("triangulation.jl") include("dss.jl") @@ -193,6 +194,7 @@ Returns a bool indicating that the space has a vertical grid. function has_vertical end has_vertical(::AbstractSpace) = false has_vertical(::ExtrudedFiniteDifferenceSpace) = true +has_vertical(::MultiColumnFiniteDifferenceSpace) = true has_vertical(::FiniteDifferenceSpace) = false """ diff --git a/src/Spaces/extruded.jl b/src/Spaces/extruded.jl index 56323d50d4..ef1bad872a 100644 --- a/src/Spaces/extruded.jl +++ b/src/Spaces/extruded.jl @@ -1,4 +1,3 @@ - """ ExtrudedFiniteDifferenceSpace(grid, staggering) diff --git a/src/Spaces/multicolumn.jl b/src/Spaces/multicolumn.jl new file mode 100644 index 0000000000..006983dda4 --- /dev/null +++ b/src/Spaces/multicolumn.jl @@ -0,0 +1,211 @@ +""" + PointCloudSpace + +A horizontal space of N independent points (lat, lon, z), produced by taking a +vertical level of a [`MultiColumnFiniteDifferenceSpace`](@ref). This is the +N-column analogue of [`PointSpace`](@ref), which is the single-column level +space. + +`full_grid` references the extruded grid the level was taken from (analogous +to `Grids.LevelGrid.full_grid`), and `local_geometry` holds an `IFH{LG, 1, N}` +data layout. +""" +struct PointCloudSpace{ + C <: ClimaComms.AbstractCommsContext, + G, + LG, +} <: AbstractSpace + context::C + full_grid::G # the AbstractExtrudedFiniteDifferenceGrid this level was sliced from + local_geometry::LG # IFH{FullLocalGeometry{(1,2,3), LatLongZPoint, ...}, 1, N} +end + +ClimaComms.context(space::PointCloudSpace) = space.context +ClimaComms.device(space::PointCloudSpace) = + ClimaComms.device(space.context) + +local_geometry_data(space::PointCloudSpace) = space.local_geometry +local_geometry_type(::Type{PointCloudSpace{C, G, LG}}) where {C, G, LG} = + eltype(LG) + +# PointCloudSpace has no grid/staggering; override the generic +# local_geometry_data(space) = local_geometry_data(grid(space), staggering(space)) +# so that it works without them. +local_geometry_data(space::PointCloudSpace, ::Nothing) = + space.local_geometry + +""" + MultiColumnFiniteDifferenceSpace + +A space of N independent vertical columns at arbitrary horizontal (lat, lon) +locations on a sphere. This is the N-column generalisation of +[`Spaces.FiniteDifferenceSpace`](@ref) (the single-column space): + +- The data layout is `VIFH{LG, Nv, 1, N}` (same vertical structure for every + column; full 3-D local geometry including lat/lon/z coordinates). +- [`Spaces.level`](@ref) returns a [`PointCloudSpace`](@ref) (N points + at that z-level) rather than a spectral-element horizontal space. +- [`Spaces.column`](@ref) returns a single-column + [`Spaces.FiniteDifferenceSpace`](@ref). +- [`Fields.bycolumn`](@ref) iterates over each column independently. + +There is no horizontal connectivity between columns; DSS and horizontal +spectral-element operators are not supported. +""" +struct MultiColumnFiniteDifferenceSpace{ + G <: Grids.AbstractExtrudedFiniteDifferenceGrid, + S <: Staggering, +} <: AbstractSpace + grid::G + staggering::S +end + +local_geometry_type(::Type{MultiColumnFiniteDifferenceSpace{G, S}}) where {G, S} = + local_geometry_type(G) + +grid(space::MultiColumnFiniteDifferenceSpace) = getfield(space, :grid) +staggering(space::MultiColumnFiniteDifferenceSpace) = getfield(space, :staggering) + +const FaceMultiColumnFiniteDifferenceSpace{G} = + MultiColumnFiniteDifferenceSpace{G, CellFace} +const CenterMultiColumnFiniteDifferenceSpace{G} = + MultiColumnFiniteDifferenceSpace{G, CellCenter} + +# Convenience constructors mirroring CenterExtrudedFiniteDifferenceSpace(space) +FaceMultiColumnFiniteDifferenceSpace(space::MultiColumnFiniteDifferenceSpace) = + MultiColumnFiniteDifferenceSpace(grid(space), CellFace()) +CenterMultiColumnFiniteDifferenceSpace(space::MultiColumnFiniteDifferenceSpace) = + MultiColumnFiniteDifferenceSpace(grid(space), CellCenter()) + +# Override the generic `space(refspace::AbstractSpace, staggering) = space(grid(refspace), staggering)` +# so that we return MultiColumnFiniteDifferenceSpace rather than ExtrudedFiniteDifferenceSpace. +space(refspace::MultiColumnFiniteDifferenceSpace, s::Staggering) = + MultiColumnFiniteDifferenceSpace(grid(refspace), s) + +function face_space(space::MultiColumnFiniteDifferenceSpace) + MultiColumnFiniteDifferenceSpace(grid(space), CellFace()) +end +function center_space(space::MultiColumnFiniteDifferenceSpace) + MultiColumnFiniteDifferenceSpace(grid(space), CellCenter()) +end + +Adapt.adapt_structure(to, space::MultiColumnFiniteDifferenceSpace) = + MultiColumnFiniteDifferenceSpace( + Adapt.adapt(to, grid(space)), + staggering(space), + ) + +# ---- column / level extraction ---------------------------------------- + +""" + column(space::MultiColumnFiniteDifferenceSpace, colidx::Grids.ColumnIndex) + +Return a single-column [`FiniteDifferenceSpace`](@ref) for column `colidx`. +""" +function column( + space::MultiColumnFiniteDifferenceSpace, + colidx::Grids.ColumnIndex, +) + column_grid = Grids.column(grid(space), colidx) + FiniteDifferenceSpace(column_grid, space.staggering) +end +column(space::MultiColumnFiniteDifferenceSpace, i, h) = + column(space, Grids.ColumnIndex((i,), h)) +column(space::MultiColumnFiniteDifferenceSpace, i, j, h) = + column(space, Grids.ColumnIndex((i,), h)) + +""" + column(space::PointCloudSpace, i, h) + +Return the single-point [`PointSpace`](@ref) for column `h` of a +[`PointCloudSpace`](@ref). This is the analogue of +`column(::SpectralElementSpace1D, i, h)` and enables `field[colidx]` indexing +inside [`Fields.bycolumn`](@ref) loops over a +[`MultiColumnFiniteDifferenceSpace`](@ref). +""" +Base.@propagate_inbounds function column(space::PointCloudSpace, i, h) + local_geometry = column(local_geometry_data(space), i, h) + PointSpace(ClimaComms.context(space), local_geometry) +end +Base.@propagate_inbounds column(space::PointCloudSpace, i, j, h) = + column(space, i, h) + +""" + level(space::MultiColumnFiniteDifferenceSpace, v) + +Return the [`PointCloudSpace`](@ref) (N-point horizontal slice) at +vertical level `v`. +""" +Base.@propagate_inbounds function level( + space::CenterMultiColumnFiniteDifferenceSpace, + v::Int, +) + data = level(local_geometry_data(space), v) + PointCloudSpace(ClimaComms.context(space), grid(space), data) +end +Base.@propagate_inbounds function level( + space::FaceMultiColumnFiniteDifferenceSpace, + v::PlusHalf, +) + data = level(local_geometry_data(space), v.i + 1) + PointCloudSpace(ClimaComms.context(space), grid(space), data) +end + +# ---- space properties -------------------------------------------------- + +ncolumns(space::MultiColumnFiniteDifferenceSpace) = + size(Grids.local_geometry_data(grid(space).horizontal_grid, nothing), 5) # Nh dim of IFH + +nlevels(space::MultiColumnFiniteDifferenceSpace) = + size(local_geometry_data(space), 4) # Nv dim of VIFH + +horizontal_space(space::MultiColumnFiniteDifferenceSpace) = + level(MultiColumnFiniteDifferenceSpace(grid(space), CellCenter()), 1) + +# No DSS / mask machinery needed. +get_mask(space::PointCloudSpace) = DataLayouts.NoMask() +get_mask(space::MultiColumnFiniteDifferenceSpace) = DataLayouts.NoMask() +set_mask!(::Any, ::MultiColumnFiniteDifferenceSpace) = nothing + +# A PointCloudSpace is a subspace of a MultiColumnFiniteDifferenceSpace when it +# was sliced from the same underlying grid. This enables broadcasting a level +# field (IFH) across a full multi-column field (VIFH), analogous to +# issubspace(::SpectralElementSpace2D{<:LevelGrid}, ::ExtrudedFiniteDifferenceSpace). +function issubspace( + level_space::PointCloudSpace, + full_space::MultiColumnFiniteDifferenceSpace, +) + return level_space.full_grid === grid(full_space) +end +function issubspace(space1::PointCloudSpace, space2::PointCloudSpace) + return space1.full_grid === space2.full_grid +end + +""" + obtain_surface_space(cs::CenterMultiColumnFiniteDifferenceSpace) + +Return the [`PointCloudSpace`](@ref) corresponding to the top face +(surface) of `cs`. Mirrors the single-column +`obtain_surface_space(::CenterFiniteDifferenceSpace)` pattern. +""" +function obtain_surface_space(cs::CenterMultiColumnFiniteDifferenceSpace) + fs = face_space(cs) + return level(fs, PlusHalf(nlevels(fs) - 1)) +end + +function Base.show(io::IO, space::MultiColumnFiniteDifferenceSpace) + indent = get(io, :indent, 0) + iio = IOContext(io, :indent => indent + 2) + println( + io, + space isa CenterMultiColumnFiniteDifferenceSpace ? + "CenterMultiColumnFiniteDifferenceSpace" : + "FaceMultiColumnFiniteDifferenceSpace", + ":", + ) + print(iio, " "^(indent + 2), "context: ") + Topologies.print_context(iio, ClimaComms.context(space)) + println(iio) + println(iio, " "^(indent + 2), "columns: ", ncolumns(space)) + println(iio, " "^(indent + 2), "levels: ", nlevels(space)) +end diff --git a/test/CommonGrids/CommonGrids.jl b/test/CommonGrids/CommonGrids.jl index 30d4ec0229..0be6b03812 100644 --- a/test/CommonGrids/CommonGrids.jl +++ b/test/CommonGrids/CommonGrids.jl @@ -5,7 +5,8 @@ using Revise; include(joinpath("test", "CommonGrids", "CommonGrids.jl")) import ClimaComms ClimaComms.@import_required_backends using ClimaCore.CommonGrids -import ClimaCore.CommonGrids: ExtrudedCubedSphereGrid, CubedSphereGrid, Box3DGrid +import ClimaCore.CommonGrids: + ExtrudedCubedSphereGrid, CubedSphereGrid, Box3DGrid, PointColumnEnsembleGrid using ClimaCore: Geometry, Hypsography, @@ -119,6 +120,32 @@ using Test ) @test grid isa Grids.SpectralElementGrid2D @test Grids.topology(grid).mesh isa Meshes.RectilinearMesh + + lats = [0.0, 3.0, 5.0] + longs = [0.0, 4.0, 7.0] + points = [ + Geometry.LatLongPoint(lat, long) for + (lat, long) in zip(lats, longs)] + radius = 100 + z_elem = 10 + z_min = -10.0 + z_max = 20.0 + grid = PointColumnEnsembleGrid(; + points, + z_elem, + z_min, + z_max, + radius, + ) + @test grid isa Grids.ExtrudedFiniteDifferenceGrid + @test grid.horizontal_grid isa Grids.PointCloudGrid + @test grid.horizontal_grid.global_geometry.radius == radius + @test grid.vertical_grid.topology.mesh.domain.coord_max == Geometry.ZPoint(z_max) + @test grid.vertical_grid.topology.mesh.domain.coord_min == Geometry.ZPoint(z_min) + @test Meshes.nelements(grid.vertical_grid.topology.mesh) == z_elem + @test vec(collect(parent(grid.horizontal_grid.local_geometry.coordinates.lat))) == lats + @test vec(collect(parent(grid.horizontal_grid.local_geometry.coordinates.long))) == + longs end @testset "Space-filling curve usage in CommonGrids" begin diff --git a/test/CommonSpaces/unit_common_spaces.jl b/test/CommonSpaces/unit_common_spaces.jl index a836423289..6056b66443 100644 --- a/test/CommonSpaces/unit_common_spaces.jl +++ b/test/CommonSpaces/unit_common_spaces.jl @@ -143,4 +143,34 @@ ClimaComms.init(ClimaComms.context()) grid = Spaces.grid(space) @test grid isa Grids.SpectralElementGrid2D @test Grids.topology(grid).mesh isa Meshes.RectilinearMesh + + lats = [0.0, 3.0, 5.0] + longs = [0.0, 4.0, 7.0] + points = [ + Geometry.LatLongPoint(lat, long) for + (lat, long) in zip(lats, longs)] + radius = 100 + z_elem = 10 + z_min = -10.0 + z_max = 20.0 + staggering = Grids.CellCenter() + space = PointColumnEnsembleSpace(; + points, + z_elem, + z_min, + z_max, + radius, + staggering, + ) + @test space.staggering isa Grids.CellCenter + grid = Spaces.grid(space) + @test grid isa Grids.ExtrudedFiniteDifferenceGrid + @test grid.horizontal_grid isa Grids.PointCloudGrid + @test grid.horizontal_grid.global_geometry.radius == radius + @test grid.vertical_grid.topology.mesh.domain.coord_max == Geometry.ZPoint(z_max) + @test grid.vertical_grid.topology.mesh.domain.coord_min == Geometry.ZPoint(z_min) + @test Meshes.nelements(grid.vertical_grid.topology.mesh) == z_elem + @test vec(collect(parent(grid.horizontal_grid.local_geometry.coordinates.lat))) == lats + @test vec(collect(parent(grid.horizontal_grid.local_geometry.coordinates.long))) == + longs end diff --git a/test/Fields/unit_field.jl b/test/Fields/unit_field.jl index 5bbe365b52..b031c17227 100644 --- a/test/Fields/unit_field.jl +++ b/test/Fields/unit_field.jl @@ -1181,6 +1181,7 @@ end FT = Float64 for space in TU.all_spaces(FT) TU.bycolumnable(space) || continue + space isa Spaces.MultiColumnFiniteDifferenceSpace && continue hspace = Spaces.horizontal_space(space) Nh = Topologies.nlocalelems(hspace) Nq = Quadratures.degrees_of_freedom(Spaces.quadrature_style(hspace)) diff --git a/test/Spaces/extruded_cuda.jl b/test/Spaces/extruded_cuda.jl index b06f766d4d..e657c3a307 100644 --- a/test/Spaces/extruded_cuda.jl +++ b/test/Spaces/extruded_cuda.jl @@ -25,7 +25,7 @@ compare(cpu, gpu, f) = all(parent(f(cpu)) .≈ Array(parent(f(gpu)))) context = SingletonCommsContext(device) collect(TU.all_spaces(Float64; zelem = 10, context)) # make sure we can construct spaces as = collect(TU.all_spaces(Float64; zelem = 10, context)) - @test length(as) == 8 + @test length(as) == 9 end @testset "copyto! with CuArray-backed extruded spaces" begin diff --git a/test/TestUtilities/TestUtilities.jl b/test/TestUtilities/TestUtilities.jl index 9625f38eed..5e6058e59a 100644 --- a/test/TestUtilities/TestUtilities.jl +++ b/test/TestUtilities/TestUtilities.jl @@ -20,6 +20,7 @@ import ClimaCore.Quadratures import ClimaCore.Geometry import ClimaCore.Meshes import ClimaCore.Spaces +import ClimaCore.CommonSpaces import ClimaCore.Topologies import ClimaCore.Domains import ClimaCore.Hypsography @@ -164,6 +165,26 @@ function FaceExtrudedFiniteDifferenceSpace(::Type{FT}; kwargs...) where {FT} return Spaces.FaceExtrudedFiniteDifferenceSpace(cspace) end +function PointColumnEnsembleSpace(::Type{FT}; context, zelem = 10, kwargs...) where {FT} + staggering = Spaces.CellCenter() + lats = FT.([0.0, 1.0, 2.0]) + longs = FT.([3.0, 4.0, 5.0]) + points = [Geometry.LatLongPoint(lat, long) for (lat, long) in zip(lats, longs)] + z_min = -10.0 + z_max = 10.0 + (; device) = context + return CommonSpaces.PointColumnEnsembleSpace( + FT; + z_elem = zelem, + staggering, + points, + z_min, + z_max, + device, + kwargs..., + ) +end + function all_spaces( ::Type{FT}; zelem = 10, @@ -179,6 +200,7 @@ function all_spaces( # SpectralElementFiniteDifferenceRectilinearSpace2D(FT; context), ColumnCenterFiniteDifferenceSpace(FT; zelem, context), ColumnFaceFiniteDifferenceSpace(FT; zelem, context), + PointColumnEnsembleSpace(FT; zelem, context), SphereSpectralElementSpace(FT; context), CenterExtrudedFiniteDifferenceSpace(FT; zelem, context, helem), FaceExtrudedFiniteDifferenceSpace(FT; zelem, context, helem), @@ -194,12 +216,14 @@ end bycolumnable(space) = ( space isa Spaces.ExtrudedFiniteDifferenceSpace || space isa Spaces.SpectralElementSpace1D || - space isa Spaces.SpectralElementSpace2D + space isa Spaces.SpectralElementSpace2D || + space isa Spaces.MultiColumnFiniteDifferenceSpace ) levelable(space) = ( space isa Spaces.ExtrudedFiniteDifferenceSpace || - space isa Spaces.FiniteDifferenceSpace + space isa Spaces.FiniteDifferenceSpace || + space isa Spaces.MultiColumnFiniteDifferenceSpace ) fc_index( @@ -207,6 +231,7 @@ fc_index( ::Union{ Spaces.FaceExtrudedFiniteDifferenceSpace, Spaces.FaceFiniteDifferenceSpace, + Spaces.FaceMultiColumnFiniteDifferenceSpace, }, ) = Utilities.PlusHalf(i) @@ -215,6 +240,7 @@ fc_index( ::Union{ Spaces.CenterExtrudedFiniteDifferenceSpace, Spaces.CenterFiniteDifferenceSpace, + Spaces.CenterMultiColumnFiniteDifferenceSpace, }, ) = i