From 9a73b9ed8f7232700f91e5ab0b5ca92d060dfe1f Mon Sep 17 00:00:00 2001 From: ph-kev <98072684+ph-kev@users.noreply.github.com> Date: Mon, 8 Jun 2026 16:17:29 -0700 Subject: [PATCH 1/3] Add space with multiple columns --- docs/src/APIs/common_grids_api.md | 1 + docs/src/APIs/common_spaces_api.md | 1 + ext/cuda/adapt.jl | 19 +++ src/CommonGrids/CommonGrids.jl | 80 ++++++++- src/CommonSpaces/CommonSpaces.jl | 58 ++++++- src/Fields/Fields.jl | 23 ++- src/Fields/indices.jl | 22 +++ src/Grids/Grids.jl | 2 + src/Grids/extruded.jl | 13 +- src/Grids/level.jl | 3 + src/Grids/pointcloud.jl | 118 +++++++++++++ src/MatrixFields/matrix_shape.jl | 6 +- src/MatrixFields/operator_matrices.jl | 9 + src/Operators/common.jl | 25 +++ src/Operators/finitedifference.jl | 9 +- src/Spaces/Spaces.jl | 2 + src/Spaces/extruded.jl | 1 - src/Spaces/multicolumn.jl | 209 ++++++++++++++++++++++++ test/CommonGrids/CommonGrids.jl | 29 +++- test/CommonSpaces/unit_common_spaces.jl | 30 ++++ test/Fields/unit_field.jl | 1 + test/Spaces/extruded_cuda.jl | 2 +- test/TestUtilities/TestUtilities.jl | 30 +++- 23 files changed, 679 insertions(+), 14 deletions(-) create mode 100644 src/Grids/pointcloud.jl create mode 100644 src/Spaces/multicolumn.jl 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..7f7415352c 100644 --- a/ext/cuda/adapt.jl +++ b/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, diff --git a/src/CommonGrids/CommonGrids.jl b/src/CommonGrids/CommonGrids.jl index f5784f3dd6..7908214169 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,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 diff --git a/src/CommonSpaces/CommonSpaces.jl b/src/CommonSpaces/CommonSpaces.jl index 15ebb9048b..e3b02c735a 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,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.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` + - `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 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..d25432d65e 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) diff --git a/src/Grids/extruded.jl b/src/Grids/extruded.jl index 7e57f3ee39..bc28b66729 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::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) 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..be918aca8d --- /dev/null +++ b/src/Grids/pointcloud.jl @@ -0,0 +1,118 @@ + +""" + 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 — one node per +"element" (`Ni = 1`), N "elements" (`Nh = N`). The horizontal Jacobian `∂x∂ξ` is the +diagonal matrix `diag(R·π/180, R·cosd(lat)·π/180)` reflecting the sphere metric at +each point, and `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) + +# No topology — return nothing. Callers that need a topology (e.g. DSS) should +# not be called on a PointCloudGrid. +topology(::PointCloudGrid) = nothing + +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(), + context :: ClimaComms.AbstractCommsContext = ClimaComms.context(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(), + context::ClimaComms.AbstractCommsContext = ClimaComms.SingletonCommsContext(device), +) where {FT} + @assert context isa ClimaComms.SingletonCommsContext "PointCloudGrid only supports SingletonCommsContext." + + 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}(0, s_lat, s_lon, 0) + + 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..5758fb5da3 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.PointCloudLevelSpace, + }, ) @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..58bfd557d2 --- /dev/null +++ b/src/Spaces/multicolumn.jl @@ -0,0 +1,209 @@ +""" + PointCloudLevelSpace + +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. + +`local_geometry` holds an `IFH{LG, 1, N}` data layout. +""" +struct PointCloudLevelSpace{ + C <: ClimaComms.AbstractCommsContext, + LG, +} <: AbstractSpace + context::C + local_geometry::LG # IFH{FullLocalGeometry{(1,2,3), LatLongZPoint, ...}, 1, N} +end + +ClimaComms.context(space::PointCloudLevelSpace) = space.context +ClimaComms.device(space::PointCloudLevelSpace) = + ClimaComms.device(space.context) + +local_geometry_data(space::PointCloudLevelSpace) = space.local_geometry +local_geometry_type(::Type{PointCloudLevelSpace{C, LG}}) where {C, LG} = + eltype(LG) + +# PointCloudLevelSpace 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::PointCloudLevelSpace, ::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 [`PointCloudLevelSpace`](@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::PointCloudLevelSpace, i, h) + +Return the single-point [`PointSpace`](@ref) for column `h` of a +[`PointCloudLevelSpace`](@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::PointCloudLevelSpace, i, h) + local_geometry = column(local_geometry_data(space), i, h) + PointSpace(ClimaComms.context(space), local_geometry) +end +Base.@propagate_inbounds column(space::PointCloudLevelSpace, i, j, h) = + column(space, i, h) + +""" + level(space::MultiColumnFiniteDifferenceSpace, v) + +Return the [`PointCloudLevelSpace`](@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) + PointCloudLevelSpace(ClimaComms.context(space), data) +end +Base.@propagate_inbounds function level( + space::FaceMultiColumnFiniteDifferenceSpace, + v::PlusHalf, +) + data = level(local_geometry_data(space), v.i + 1) + PointCloudLevelSpace(ClimaComms.context(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::PointCloudLevelSpace) = DataLayouts.NoMask() +get_mask(space::MultiColumnFiniteDifferenceSpace) = DataLayouts.NoMask() +set_mask!(fn, space::MultiColumnFiniteDifferenceSpace) = + set_mask!(fn, grid(space).horizontal_grid) + +# A PointCloudLevelSpace is a subspace of a MultiColumnFiniteDifferenceSpace +# when they share the same underlying grid. This enables broadcasting a level +# field (IFH) across a full multi-column field (VIFH), analogous to how a +# PointSpace field can be broadcast with a CenterFiniteDifferenceSpace field. +function issubspace( + level_space::PointCloudLevelSpace, + full_space::MultiColumnFiniteDifferenceSpace, +) + hgrid = grid(full_space).horizontal_grid + # local_geometry of the level space is an IFH slice of the full VIFH; + # both share the same horizontal grid data array. + return size(level_space.local_geometry, 5) == size( + Grids.local_geometry_data(hgrid, nothing), 5) +end + +""" + obtain_surface_space(cs::CenterMultiColumnFiniteDifferenceSpace) + +Return the [`PointCloudLevelSpace`](@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 From cd184bdc3ead85a98bf27a1f215127730ccca393 Mon Sep 17 00:00:00 2001 From: ph-kev <98072684+ph-kev@users.noreply.github.com> Date: Wed, 22 Jul 2026 13:31:34 -0700 Subject: [PATCH 2/3] Address comments --- ext/cuda/adapt.jl | 19 +++------- src/CommonGrids/CommonGrids.jl | 7 +--- src/CommonSpaces/CommonSpaces.jl | 2 -- src/Grids/Grids.jl | 1 + src/Grids/extruded.jl | 12 +++---- src/Grids/pointcloud.jl | 39 ++++++++++---------- src/MatrixFields/matrix_shape.jl | 2 +- src/Spaces/multicolumn.jl | 61 ++++++++++++++++---------------- 8 files changed, 63 insertions(+), 80 deletions(-) diff --git a/ext/cuda/adapt.jl b/ext/cuda/adapt.jl index 7f7415352c..225e963ebd 100644 --- a/ext/cuda/adapt.jl +++ b/ext/cuda/adapt.jl @@ -7,25 +7,13 @@ 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), 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, @@ -52,9 +40,10 @@ 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( +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)), ) diff --git a/src/CommonGrids/CommonGrids.jl b/src/CommonGrids/CommonGrids.jl index 7908214169..5c285f218b 100644 --- a/src/CommonGrids/CommonGrids.jl +++ b/src/CommonGrids/CommonGrids.jl @@ -720,7 +720,6 @@ end 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), ) @@ -737,7 +736,6 @@ arbitrary (lat, lon) locations on a sphere, given: - `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` @@ -769,13 +767,10 @@ function PointColumnEnsembleGrid( 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) + h_grid = Grids.PointCloudGrid(points; radius, device) z_topology = Topologies.IntervalTopology( ClimaComms.SingletonCommsContext(device), z_mesh, diff --git a/src/CommonSpaces/CommonSpaces.jl b/src/CommonSpaces/CommonSpaces.jl index e3b02c735a..df87d76f66 100644 --- a/src/CommonSpaces/CommonSpaces.jl +++ b/src/CommonSpaces/CommonSpaces.jl @@ -436,7 +436,6 @@ RectangleXYSpace(::Type{FT}; kwargs...) where {FT} = 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, @@ -452,7 +451,6 @@ locations on a sphere, given: - `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)] diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl index d25432d65e..22759fe63d 100644 --- a/src/Grids/Grids.jl +++ b/src/Grids/Grids.jl @@ -126,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 bc28b66729..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, PointCloudGrid}, + 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, PointCloudGrid}, + 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, PointCloudGrid}, + horizontal_grid::AbstractSpectralElementGrid, vertical_grid::FiniteDifferenceGrid, hypsography::Flat, global_geometry::Geometry.AbstractGlobalGeometry, @@ -126,11 +126,9 @@ 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::ExtrudedFiniteDifferenceGrid) = ClimaComms.context(grid.horizontal_grid) -ClimaComms.device(grid::ExtrudedFiniteDifferenceGrid{<:PointCloudGrid}) = +ClimaComms.device(grid::ExtrudedFiniteDifferenceGrid) = ClimaComms.device(grid.horizontal_grid) vertical_topology(grid::ExtrudedFiniteDifferenceGrid) = diff --git a/src/Grids/pointcloud.jl b/src/Grids/pointcloud.jl index be918aca8d..ac479649ee 100644 --- a/src/Grids/pointcloud.jl +++ b/src/Grids/pointcloud.jl @@ -5,17 +5,20 @@ 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 — one node per -"element" (`Ni = 1`), N "elements" (`Nh = N`). The horizontal Jacobian `∂x∂ξ` is the -diagonal matrix `diag(R·π/180, R·cosd(lat)·π/180)` reflecting the sphere metric at -each point, and `J = R²·cosd(lat)·(π/180)²`. +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, @@ -35,9 +38,9 @@ local_geometry_type(::Type{PointCloudGrid{C, GG, LG}}) where {C, GG, LG} = ClimaComms.context(grid::PointCloudGrid) = grid.context ClimaComms.device(grid::PointCloudGrid) = ClimaComms.device(grid.context) -# No topology — return nothing. Callers that need a topology (e.g. DSS) should -# not be called on a PointCloudGrid. -topology(::PointCloudGrid) = nothing +topology(::PointCloudGrid) = error( + "PointCloudGrid has no topology", +) local_geometry_data(grid::PointCloudGrid, ::Nothing) = grid.local_geometry @@ -51,7 +54,6 @@ quadrature_style(::PointCloudGrid) = points :: AbstractVector{Geometry.LatLongPoint{FT}}; radius :: Real, device :: ClimaComms.AbstractDevice = ClimaComms.device(), - context :: ClimaComms.AbstractCommsContext = ClimaComms.context(device), ) Convenience constructor: build a `PointCloudGrid` from a vector of @@ -63,9 +65,10 @@ function PointCloudGrid( points::AbstractVector{Geometry.LatLongPoint{FT}}; radius::Real, device::ClimaComms.AbstractDevice = ClimaComms.device(), - context::ClimaComms.AbstractCommsContext = ClimaComms.SingletonCommsContext(device), ) where {FT} - @assert context isa ClimaComms.SingletonCommsContext "PointCloudGrid only supports SingletonCommsContext." + # 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)) @@ -89,7 +92,7 @@ function PointCloudGrid( 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}(0, s_lat, s_lon, 0) + ∂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( diff --git a/src/MatrixFields/matrix_shape.jl b/src/MatrixFields/matrix_shape.jl index 5758fb5da3..483e27822f 100644 --- a/src/MatrixFields/matrix_shape.jl +++ b/src/MatrixFields/matrix_shape.jl @@ -30,7 +30,7 @@ function matrix_shape( matrix_space::Union{ Spaces.AbstractSpectralElementSpace, Spaces.PointSpace, - Spaces.PointCloudLevelSpace, + Spaces.PointCloudSpace, }, ) @assert eltype(matrix_field) <: DiagonalMatrixRow diff --git a/src/Spaces/multicolumn.jl b/src/Spaces/multicolumn.jl index 58bfd557d2..058d832177 100644 --- a/src/Spaces/multicolumn.jl +++ b/src/Spaces/multicolumn.jl @@ -1,33 +1,37 @@ """ - PointCloudLevelSpace + 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. -`local_geometry` holds an `IFH{LG, 1, N}` data layout. +`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 PointCloudLevelSpace{ +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::PointCloudLevelSpace) = space.context -ClimaComms.device(space::PointCloudLevelSpace) = +ClimaComms.context(space::PointCloudSpace) = space.context +ClimaComms.device(space::PointCloudSpace) = ClimaComms.device(space.context) -local_geometry_data(space::PointCloudLevelSpace) = space.local_geometry -local_geometry_type(::Type{PointCloudLevelSpace{C, LG}}) where {C, LG} = +local_geometry_data(space::PointCloudSpace) = space.local_geometry +local_geometry_type(::Type{PointCloudSpace{C, G, LG}}) where {C, G, LG} = eltype(LG) -# PointCloudLevelSpace has no grid/staggering; override the generic +# 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::PointCloudLevelSpace, ::Nothing) = +local_geometry_data(space::PointCloudSpace, ::Nothing) = space.local_geometry """ @@ -39,7 +43,7 @@ locations on a sphere. This is the N-column generalisation of - 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 [`PointCloudLevelSpace`](@ref) (N points +- [`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). @@ -111,25 +115,25 @@ column(space::MultiColumnFiniteDifferenceSpace, i, j, h) = column(space, Grids.ColumnIndex((i,), h)) """ - column(space::PointCloudLevelSpace, i, h) + column(space::PointCloudSpace, i, h) Return the single-point [`PointSpace`](@ref) for column `h` of a -[`PointCloudLevelSpace`](@ref). This is the analogue of +[`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::PointCloudLevelSpace, i, h) +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::PointCloudLevelSpace, i, j, h) = +Base.@propagate_inbounds column(space::PointCloudSpace, i, j, h) = column(space, i, h) """ level(space::MultiColumnFiniteDifferenceSpace, v) -Return the [`PointCloudLevelSpace`](@ref) (N-point horizontal slice) at +Return the [`PointCloudSpace`](@ref) (N-point horizontal slice) at vertical level `v`. """ Base.@propagate_inbounds function level( @@ -137,14 +141,14 @@ Base.@propagate_inbounds function level( v::Int, ) data = level(local_geometry_data(space), v) - PointCloudLevelSpace(ClimaComms.context(space), data) + 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) - PointCloudLevelSpace(ClimaComms.context(space), data) + PointCloudSpace(ClimaComms.context(space), grid(space), data) end # ---- space properties -------------------------------------------------- @@ -159,30 +163,25 @@ horizontal_space(space::MultiColumnFiniteDifferenceSpace) = level(MultiColumnFiniteDifferenceSpace(grid(space), CellCenter()), 1) # No DSS / mask machinery needed. -get_mask(space::PointCloudLevelSpace) = DataLayouts.NoMask() +get_mask(space::PointCloudSpace) = DataLayouts.NoMask() get_mask(space::MultiColumnFiniteDifferenceSpace) = DataLayouts.NoMask() -set_mask!(fn, space::MultiColumnFiniteDifferenceSpace) = - set_mask!(fn, grid(space).horizontal_grid) +set_mask!(::Any, ::MultiColumnFiniteDifferenceSpace) = nothing -# A PointCloudLevelSpace is a subspace of a MultiColumnFiniteDifferenceSpace -# when they share the same underlying grid. This enables broadcasting a level -# field (IFH) across a full multi-column field (VIFH), analogous to how a -# PointSpace field can be broadcast with a CenterFiniteDifferenceSpace field. +# 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::PointCloudLevelSpace, + level_space::PointCloudSpace, full_space::MultiColumnFiniteDifferenceSpace, ) - hgrid = grid(full_space).horizontal_grid - # local_geometry of the level space is an IFH slice of the full VIFH; - # both share the same horizontal grid data array. - return size(level_space.local_geometry, 5) == size( - Grids.local_geometry_data(hgrid, nothing), 5) + return level_space.full_grid === grid(full_space) end """ obtain_surface_space(cs::CenterMultiColumnFiniteDifferenceSpace) -Return the [`PointCloudLevelSpace`](@ref) corresponding to the top face +Return the [`PointCloudSpace`](@ref) corresponding to the top face (surface) of `cs`. Mirrors the single-column `obtain_surface_space(::CenterFiniteDifferenceSpace)` pattern. """ From d7f8be338429b9971a3d4187f701c8b92d2d70e0 Mon Sep 17 00:00:00 2001 From: ph-kev <98072684+ph-kev@users.noreply.github.com> Date: Fri, 24 Jul 2026 10:06:17 -0700 Subject: [PATCH 3/3] Add another subspace method --- src/Spaces/multicolumn.jl | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/Spaces/multicolumn.jl b/src/Spaces/multicolumn.jl index 058d832177..006983dda4 100644 --- a/src/Spaces/multicolumn.jl +++ b/src/Spaces/multicolumn.jl @@ -167,9 +167,9 @@ 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 +# 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, @@ -177,6 +177,9 @@ function issubspace( ) 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)