diff --git a/.dev/Manifest.toml b/.dev/Manifest.toml index ba35806a3a..2e73e1e3ba 100644 --- a/.dev/Manifest.toml +++ b/.dev/Manifest.toml @@ -79,9 +79,9 @@ version = "0.21.4" [[deps.JuliaFormatter]] deps = ["CSTParser", "CommonMark", "DataStructures", "Glob", "Pkg", "PrecompileTools", "Tokenize"] -git-tree-sha1 = "80031f6e58b09b0de4553bf63d9a36ec5db57967" +git-tree-sha1 = "0bac3374ff3aa798148669ecb5559ba20c8b0e73" uuid = "98e50ef6-434e-11e9-1051-2b60c6c9e899" -version = "1.0.39" +version = "1.0.40" [[deps.LibCURL]] deps = ["LibCURL_jll", "MozillaCACerts_jll"] diff --git a/.dev/up_deps.jl b/.dev/up_deps.jl index 338db084d7..7cb2edb256 100644 --- a/.dev/up_deps.jl +++ b/.dev/up_deps.jl @@ -12,6 +12,7 @@ dirs = ( joinpath(root, "perf"), joinpath(root, "docs"), joinpath(root, "test"), + joinpath(root, "benchmarks", "bickleyjet"), joinpath(root, "lib", "ClimaCoreMakie"), joinpath(root, "lib", "ClimaCorePlots"), joinpath(root, "lib", "ClimaCoreTempestRemap"), diff --git a/.github/workflows/ClimaCoreMakie.yml b/.github/workflows/ClimaCoreMakie.yml index 1d238be631..f16f4eeda9 100644 --- a/.github/workflows/ClimaCoreMakie.yml +++ b/.github/workflows/ClimaCoreMakie.yml @@ -11,7 +11,7 @@ on: jobs: lib-climacore-makie: runs-on: ubuntu-20.04 - timeout-minutes: 30 + timeout-minutes: 45 steps: - name: Checkout uses: actions/checkout@v2 diff --git a/Project.toml b/Project.toml index 18acf21e44..3e5125b839 100644 --- a/Project.toml +++ b/Project.toml @@ -19,6 +19,7 @@ HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Memoize = "c03570c3-d221-55d1-a50c-7939bbd78826" PkgVersion = "eebad327-c553-4316-9ea0-9fa01ccd7688" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Requires = "ae029012-a4dd-5104-9daa-d747884805df" @@ -28,6 +29,7 @@ Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" +WeakValueDicts = "897b6980-f191-5a31-bcb0-bf3c4585e0c1" [compat] Adapt = "3" @@ -43,6 +45,7 @@ GaussQuadrature = "0.5" GilbertCurves = "0.1" HDF5 = "0.16, 0.17" IntervalSets = "0.5, 0.6, 0.7" +Memoize = "0.4" PkgVersion = "0.1, 0.2, 0.3" RecursiveArrayTools = "2" Requires = "1" @@ -50,6 +53,7 @@ RootSolvers = "0.3, 0.4" Static = "0.4, 0.5, 0.6, 0.7, 0.8" StaticArrays = "1" UnPack = "1" +WeakValueDicts = "0.1" julia = "1.8" [extras] diff --git a/benchmarks/bickleyjet/Manifest.toml b/benchmarks/bickleyjet/Manifest.toml index fc0a8268ad..610ec216db 100644 --- a/benchmarks/bickleyjet/Manifest.toml +++ b/benchmarks/bickleyjet/Manifest.toml @@ -157,7 +157,7 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" version = "0.5.5" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"] +deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack", "WeakValueDicts"] path = "../.." uuid = "d414da3d-4745-48bb-8d80-42e94e092884" version = "0.10.55" @@ -784,6 +784,12 @@ git-tree-sha1 = "c13304c81eec1ed3af7fc20e75fb6b26092a1102" uuid = "442fdcdd-2543-5da2-b0f3-8c86c306513e" version = "0.3.2" +[[deps.Memoize]] +deps = ["MacroTools"] +git-tree-sha1 = "2b1dfcba103de714d31c033b5dacc2e4a12c7caa" +uuid = "c03570c3-d221-55d1-a50c-7939bbd78826" +version = "0.4.4" + [[deps.MicrosoftMPI_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] git-tree-sha1 = "a7023883872e52bc29bcaac74f19adf39347d2d5" @@ -1272,6 +1278,11 @@ git-tree-sha1 = "4528479aa01ee1b3b4cd0e6faef0e04cf16466da" uuid = "2381bf8a-dfd0-557d-9999-79630e7b1b91" version = "1.25.0+0" +[[deps.WeakValueDicts]] +git-tree-sha1 = "98528c2610a5479f091d470967a25becfd83edd0" +uuid = "897b6980-f191-5a31-bcb0-bf3c4585e0c1" +version = "0.1.0" + [[deps.XML2_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"] git-tree-sha1 = "24b81b59bd35b3c42ab84fa589086e19be919916" diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 622c3d83e4..db504d3f47 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -274,7 +274,7 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" version = "0.5.5" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"] +deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack", "WeakValueDicts"] path = ".." uuid = "d414da3d-4745-48bb-8d80-42e94e092884" version = "0.10.55" @@ -1310,9 +1310,9 @@ version = "0.1.12" [[deps.LinearSolve]] deps = ["ArrayInterface", "ConcreteStructs", "DocStringExtensions", "EnumX", "EnzymeCore", "FastLapackInterface", "GPUArraysCore", "InteractiveUtils", "KLU", "Krylov", "Libdl", "LinearAlgebra", "MKL_jll", "PrecompileTools", "Preferences", "RecursiveFactorization", "Reexport", "Requires", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Sparspak", "SuiteSparse", "UnPack"] -git-tree-sha1 = "158e45dd35cec1ecade0e554c0104ee89e772d82" +git-tree-sha1 = "3196f7408f7c6014dc4ed55260e6c5e1e01c58a3" uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" -version = "2.11.1" +version = "2.12.0" [deps.LinearSolve.extensions] LinearSolveBandedMatricesExt = "BandedMatrices" @@ -1477,6 +1477,12 @@ git-tree-sha1 = "c13304c81eec1ed3af7fc20e75fb6b26092a1102" uuid = "442fdcdd-2543-5da2-b0f3-8c86c306513e" version = "0.3.2" +[[deps.Memoize]] +deps = ["MacroTools"] +git-tree-sha1 = "2b1dfcba103de714d31c033b5dacc2e4a12c7caa" +uuid = "c03570c3-d221-55d1-a50c-7939bbd78826" +version = "0.4.4" + [[deps.MicrosoftMPI_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] git-tree-sha1 = "a7023883872e52bc29bcaac74f19adf39347d2d5" @@ -2543,6 +2549,11 @@ git-tree-sha1 = "4528479aa01ee1b3b4cd0e6faef0e04cf16466da" uuid = "2381bf8a-dfd0-557d-9999-79630e7b1b91" version = "1.25.0+0" +[[deps.WeakValueDicts]] +git-tree-sha1 = "98528c2610a5479f091d470967a25becfd83edd0" +uuid = "897b6980-f191-5a31-bcb0-bf3c4585e0c1" +version = "0.1.0" + [[deps.WoodburyMatrices]] deps = ["LinearAlgebra", "SparseArrays"] git-tree-sha1 = "de67fa59e33ad156a590055375a30b23c40299d3" diff --git a/docs/src/api.md b/docs/src/api.md index 9919baf1ed..93a75ba254 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -184,7 +184,6 @@ Spaces.Δz_metric_component ```@docs Spaces.SpectralElementSpace1D Spaces.SpectralElementSpace2D -Spaces.SpectralElementSpace2D(topology, quadrature_style; enable_bubble) Spaces.SpectralElementSpaceSlab ``` @@ -192,40 +191,39 @@ Spaces.SpectralElementSpaceSlab ```@docs -Spaces.Quadratures.QuadratureStyle -Spaces.Quadratures.GLL -Spaces.Quadratures.GL -Spaces.Quadratures.Uniform -Spaces.Quadratures.degrees_of_freedom -Spaces.Quadratures.polynomial_degree -Spaces.Quadratures.quadrature_points -Spaces.Quadratures.barycentric_weights -Spaces.Quadratures.interpolation_matrix -Spaces.Quadratures.differentiation_matrix -Spaces.Quadratures.orthonormal_poly +Quadratures.QuadratureStyle +Quadratures.GLL +Quadratures.GL +Quadratures.Uniform +Quadratures.degrees_of_freedom +Quadratures.polynomial_degree +Quadratures.quadrature_points +Quadratures.barycentric_weights +Quadratures.interpolation_matrix +Quadratures.differentiation_matrix +Quadratures.orthonormal_poly ``` #### Internals ```@docs -Spaces.dss_transform -Spaces.dss_transform! -Spaces.dss_untransform! -Spaces.dss_untransform -Spaces.dss_local_vertices! -Spaces.dss_local! -Spaces.dss_local_ghost! -Spaces.dss_ghost! -Spaces.create_dss_buffer -Spaces.fill_send_buffer! -Spaces.DSSBuffer -Spaces.create_ghost_buffer -Spaces.load_from_recv_buffer! +Topologies.dss_transform +Topologies.dss_transform! +Topologies.dss_untransform! +Topologies.dss_untransform +Topologies.dss_local_vertices! +Topologies.dss_local! +Topologies.dss_local_ghost! +Topologies.dss_ghost! +Topologies.create_dss_buffer +Topologies.fill_send_buffer! +Topologies.DSSBuffer +Topologies.load_from_recv_buffer! +Topologies.dss! Spaces.weighted_dss_start! Spaces.weighted_dss_internal! Spaces.weighted_dss_ghost! Spaces.weighted_dss! -Spaces.dss! Spaces.unique_nodes ``` diff --git a/examples/Manifest.toml b/examples/Manifest.toml index 56b04e6b5d..066c7fd3c0 100644 --- a/examples/Manifest.toml +++ b/examples/Manifest.toml @@ -238,7 +238,7 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" version = "0.5.5" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"] +deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack", "WeakValueDicts"] path = ".." uuid = "d414da3d-4745-48bb-8d80-42e94e092884" version = "0.10.55" @@ -1099,9 +1099,9 @@ version = "2.5.2" [[deps.LinearSolve]] deps = ["ArrayInterface", "ConcreteStructs", "DocStringExtensions", "EnumX", "EnzymeCore", "FastLapackInterface", "GPUArraysCore", "InteractiveUtils", "KLU", "Krylov", "Libdl", "LinearAlgebra", "MKL_jll", "PrecompileTools", "Preferences", "RecursiveFactorization", "Reexport", "Requires", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Sparspak", "SuiteSparse", "UnPack"] -git-tree-sha1 = "158e45dd35cec1ecade0e554c0104ee89e772d82" +git-tree-sha1 = "3196f7408f7c6014dc4ed55260e6c5e1e01c58a3" uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" -version = "2.11.1" +version = "2.12.0" [deps.LinearSolve.extensions] LinearSolveBandedMatricesExt = "BandedMatrices" @@ -1226,6 +1226,12 @@ git-tree-sha1 = "c13304c81eec1ed3af7fc20e75fb6b26092a1102" uuid = "442fdcdd-2543-5da2-b0f3-8c86c306513e" version = "0.3.2" +[[deps.Memoize]] +deps = ["MacroTools"] +git-tree-sha1 = "2b1dfcba103de714d31c033b5dacc2e4a12c7caa" +uuid = "c03570c3-d221-55d1-a50c-7939bbd78826" +version = "0.4.4" + [[deps.MicroCollections]] deps = ["BangBang", "InitialValues", "Setfield"] git-tree-sha1 = "629afd7d10dbc6935ec59b32daeb33bc4460a42e" @@ -2082,6 +2088,11 @@ git-tree-sha1 = "4528479aa01ee1b3b4cd0e6faef0e04cf16466da" uuid = "2381bf8a-dfd0-557d-9999-79630e7b1b91" version = "1.25.0+0" +[[deps.WeakValueDicts]] +git-tree-sha1 = "98528c2610a5479f091d470967a25becfd83edd0" +uuid = "897b6980-f191-5a31-bcb0-bf3c4585e0c1" +version = "0.1.0" + [[deps.WriteVTK]] deps = ["Base64", "CodecZlib", "FillArrays", "LightXML", "TranscodingStreams", "VTKBase"] git-tree-sha1 = "41f0dc2a8f6fd860c266b91fd5cdf4fead65ae69" diff --git a/examples/bickleyjet/bickleyjet_dg.jl b/examples/bickleyjet/bickleyjet_dg.jl index 30ed6ddd35..8d659aa161 100644 --- a/examples/bickleyjet/bickleyjet_dg.jl +++ b/examples/bickleyjet/bickleyjet_dg.jl @@ -208,7 +208,7 @@ function rhs!(dydt, y, (parameters, numflux), t) dydt_data .= RecursiveApply.rdiv.(dydt_data, space.local_geometry.WJ) M = Spaces.Quadratures.cutoff_filter_matrix( Float64, - space.quadrature_style, + Spaces.quadrature_style(space), 3, ) Operators.tensor_product!(dydt_data, M) diff --git a/examples/hybrid/hybrid3dcs_dss.jl b/examples/hybrid/hybrid3dcs_dss.jl index 5b9e68a0e1..ecffa77ada 100644 --- a/examples/hybrid/hybrid3dcs_dss.jl +++ b/examples/hybrid/hybrid3dcs_dss.jl @@ -85,10 +85,8 @@ function hybrid3dcubedsphere_dss_profiler( c = center_initial_condition(ᶜlocal_geometry), f = face_initial_condition(ᶠlocal_geometry), ) - ghost_buffer = ( - c = Spaces.create_ghost_buffer(Y.c), - f = Spaces.create_ghost_buffer(Y.f), - ) + ghost_buffer = + (c = Spaces.create_dss_buffer(Y.c), f = Spaces.create_dss_buffer(Y.f)) dss_buffer_f = Spaces.create_dss_buffer(Y.f) dss_buffer_c = Spaces.create_dss_buffer(Y.c) nsamples = 10000 @@ -96,7 +94,7 @@ function hybrid3dcubedsphere_dss_profiler( # precompile relevant functions space = axes(Y.c) - horizontal_topology = space.horizontal_space.topology + horizontal_topology = Spaces.topology(space) Spaces.weighted_dss_internal!(Y.c, ghost_buffer.c) weighted_dss_full!(Y.c, ghost_buffer.c) Spaces.fill_send_buffer!( diff --git a/examples/hybrid/sphere/deformation_flow.jl b/examples/hybrid/sphere/deformation_flow.jl index db995ae2f8..f56071e2ce 100644 --- a/examples/hybrid/sphere/deformation_flow.jl +++ b/examples/hybrid/sphere/deformation_flow.jl @@ -216,7 +216,7 @@ function run_deformation_flow(use_limiter, fct_op) Geometry.LatLongZPoint(ϕ_c, λ_c1, FT(0)), Geometry.LatLongZPoint(ϕ_c, λ_c2, FT(0)), ) - horz_geometry = horz_space.global_geometry + horz_geometry = Spaces.global_geometry(horz_space) rds = map(centers) do center Geometry.great_circle_distance(coord, center, horz_geometry) end diff --git a/examples/hybrid/sphere/hadley_circulation.jl b/examples/hybrid/sphere/hadley_circulation.jl index 441bcb326c..2b1db316e9 100644 --- a/examples/hybrid/sphere/hadley_circulation.jl +++ b/examples/hybrid/sphere/hadley_circulation.jl @@ -250,7 +250,7 @@ end # write out our cubed sphere mesh meshfile_cc = remap_tmpdir * "mesh_cubedsphere.g" -write_exodus(meshfile_cc, hv_center_space.horizontal_space.topology) +write_exodus(meshfile_cc, Spaces.topology(hv_center_space)) # write out RLL mesh nlat = 90 diff --git a/examples/hybrid/sphere/remap_pipeline.jl b/examples/hybrid/sphere/remap_pipeline.jl index 756e26111a..1fffecc833 100644 --- a/examples/hybrid/sphere/remap_pipeline.jl +++ b/examples/hybrid/sphere/remap_pipeline.jl @@ -53,7 +53,7 @@ function remap2latlon(filein, nc_dir, nlat, nlon) # reconstruct space, obtain Nq from space cspace = axes(Y.c) hspace = Spaces.horizontal_space(cspace) - Nq = Spaces.Quadratures.degrees_of_freedom(hspace.quadrature_style) + Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(hspace)) # create a temporary dir for intermediate data remap_tmpdir = nc_dir * "remaptmp/" @@ -97,7 +97,7 @@ function remap2latlon(filein, nc_dir, nlat, nlon) # write out our cubed sphere mesh meshfile_cc = remap_tmpdir * "mesh_cubedsphere.g" - write_exodus(meshfile_cc, hspace.topology) + write_exodus(meshfile_cc, Spaces.topology(hspace)) meshfile_rll = remap_tmpdir * "mesh_rll.g" rll_mesh(meshfile_rll; nlat = nlat, nlon = nlon) diff --git a/examples/sphere/limiters_advection.jl b/examples/sphere/limiters_advection.jl index db0e5bb23b..d5bc54e0d0 100644 --- a/examples/sphere/limiters_advection.jl +++ b/examples/sphere/limiters_advection.jl @@ -107,13 +107,13 @@ for (k, ne) in enumerate(ne_seq) Spaces.SpectralElementSpace2D(grid_topology, quad; enable_bubble = true) # Initialize variables needed for limiters - n_elems = Topologies.nlocalelems(space.topology) + n_elems = Topologies.nlocalelems(Spaces.topology(space)) min_q = zeros(n_elems) max_q = zeros(n_elems) coords = Fields.coordinate_field(space) Δh[k] = 2 * R / ne - global_geom = space.global_geometry + global_geom = Spaces.global_geometry(space) # Initialize state y0 = map(coords) do coord diff --git a/examples/sphere/shallow_water_cuda.jl b/examples/sphere/shallow_water_cuda.jl index 329e1782cb..407402e6d7 100644 --- a/examples/sphere/shallow_water_cuda.jl +++ b/examples/sphere/shallow_water_cuda.jl @@ -473,7 +473,7 @@ function rhs!(dYdt_fv, y_fv, parameters, t) wcurl(Geometry.Covariant3Vector(curl(y.u))), ) - NVTX.@range "dss" Spaces.weighted_dss2!(dYdt, ghost_buffer) + NVTX.@range "dss" Spaces.weighted_dss!(dYdt, ghost_buffer) end NVTX.@range "tendency" begin @@ -489,7 +489,7 @@ function rhs!(dYdt_fv, y_fv, parameters, t) dYdt.u += -grad(g * (y.h + h_s) + norm(y.u)^2 / 2) #+ dYdt.u += y.u × (f + curl(y.u)) end - NVTX.@range "dss" Spaces.weighted_dss2!(dYdt, ghost_buffer) + NVTX.@range "dss" Spaces.weighted_dss!(dYdt, ghost_buffer) end end return dYdt_fv diff --git a/examples/sphere/solidbody.jl b/examples/sphere/solidbody.jl index 2e1758c187..a7edc694d4 100644 --- a/examples/sphere/solidbody.jl +++ b/examples/sphere/solidbody.jl @@ -87,7 +87,7 @@ for (k, ne) in enumerate(ne_seq) Δh[k] = 2 * R / ne - global_geom = space.global_geometry + global_geom = Spaces.global_geometry(space) h_init = map(coords) do coord rd = Geometry.great_circle_distance(coord, center, global_geom) diff --git a/lib/ClimaCoreMakie/src/utils.jl b/lib/ClimaCoreMakie/src/utils.jl index 223dde88e3..4c183e7236 100644 --- a/lib/ClimaCoreMakie/src/utils.jl +++ b/lib/ClimaCoreMakie/src/utils.jl @@ -22,7 +22,7 @@ function plot_vertices( space, ClimaCore.Geometry.Cartesian123Point.( ClimaCore.Fields.coordinate_field(space), - Ref(space.global_geometry), + Ref(Spaces.global_geometry(space)), ), ) end diff --git a/lib/ClimaCorePlots/src/ClimaCorePlots.jl b/lib/ClimaCorePlots/src/ClimaCorePlots.jl index 00e9e23dba..989bfab9e6 100644 --- a/lib/ClimaCorePlots/src/ClimaCorePlots.jl +++ b/lib/ClimaCorePlots/src/ClimaCorePlots.jl @@ -31,7 +31,8 @@ RecipesBase.@recipe function f( coord_field = Fields.coordinate_field(space) lagrange_quad = Spaces.Quadratures.ClosedUniform{Nu}() - dataspace = Spaces.SpectralElementSpace1D(space.topology, lagrange_quad) + dataspace = + Spaces.SpectralElementSpace1D(Spaces.topology(space), lagrange_quad) interp = Operators.Interpolate(dataspace) ifield = interp.(field) @@ -60,7 +61,7 @@ RecipesBase.@recipe function f( (xcoords, xdata) end -RecipesBase.@recipe function f(space::Spaces.SpectralElementSpace2D;) +RecipesBase.@recipe function f(space::Spaces.RectilinearSpectralElementSpace2D;) quad = Spaces.quadrature_style(space) quad_name = Base.typename(typeof(quad)).name dof = Spaces.Quadratures.degrees_of_freedom(quad) @@ -98,7 +99,7 @@ RecipesBase.@recipe function f(space::Spaces.ExtrudedFiniteDifferenceSpace) #TODO: assumes VIFH layout @assert Nj == 1 "plotting only defined for 1D extruded fields" - hspace = space.horizontal_space + hspace = Spaces.horizontal_space(space) quad = Spaces.quadrature_style(hspace) quad_name = Base.typename(typeof(quad)).name @@ -151,7 +152,7 @@ RecipesBase.@recipe function f(field::Fields.FiniteDifferenceField) end RecipesBase.@recipe function f( - field::Fields.SpectralElementField2D; + field::Fields.RectilinearSpectralElementField2D; interpolate = 10, ) @assert interpolate ≥ 1 "number of element quadrature points for uniform interpolation must be ≥ 1" @@ -195,9 +196,9 @@ function _slice_triplot(field, hinterpolate, ncolors) Ni, Nj, _, Nv, Nh = size(data) space = axes(field) - htopology = Spaces.topology(space.horizontal_space) + htopology = Spaces.topology(space) hdomain = Topologies.domain(htopology) - vdomain = Topologies.domain(space.vertical_topology) + vdomain = Topologies.domain(Spaces.vertical_topology(space)) @assert Nj == 1 @@ -248,15 +249,10 @@ end # 2D hybrid plot RecipesBase.@recipe function f( - field::Fields.Field{<:Any, S}; + field::Fields.ExtrudedSpectralElementField2D; hinterpolate = 0, ncolors = 256, -) where { - S <: Spaces.ExtrudedFiniteDifferenceSpace{ - <:Any, - <:Spaces.IntervalSpectralElementSpace1D, - }, -} +) hcoord, vcoord, data = _slice_triplot(field, hinterpolate, ncolors) coord_symbols = propertynames(Fields.coordinate_field(axes(field))) @@ -282,7 +278,7 @@ function _slice_along(field, coord) ) end space = axes(field) - hspace = space.horizontal_space + hspace = Spaces.horizontal_space(space) htopo = ClimaCore.Spaces.topology(hspace) hmesh = htopo.mesh linear_idx = LinearIndices(ClimaCore.Meshes.elements(hmesh)) @@ -305,7 +301,7 @@ function _slice_along(field, coord) hidx = axis == 1 ? linear_idx[slice_h, 1] : linear_idx[1, slice_h] # find the node idx we want to slice along the given axis element - hcoord_data = hspace.local_geometry.coordinates + hcoord_data = Spaces.local_geometry_data(hspace).coordinates hdata = ClimaCore.slab(hcoord_data, hidx) hnode_idx = 1 for i in axes(hdata)[axis] @@ -324,8 +320,9 @@ function _slice_along(field, coord) htopo_ortho, ClimaCore.Spaces.quadrature_style(hspace), ) - vspace_ortho = - ClimaCore.Spaces.CenterFiniteDifferenceSpace(space.vertical_topology) + vspace_ortho = ClimaCore.Spaces.CenterFiniteDifferenceSpace( + Spaces.vertical_topology(space), + ) if space.staggering === ClimaCore.Spaces.CellFace vspace_ortho = ClimaCore.Spaces.FaceFiniteDifferenceSpace(slice_vspace) @@ -359,16 +356,11 @@ end # 3D hybrid plot RecipesBase.@recipe function f( - field::Fields.Field{<:Any, S}; + field::Fields.ExtrudedFiniteDifferenceField3D; slice = nothing, hinterpolate = 0, ncolors = 256, -) where { - S <: Spaces.ExtrudedFiniteDifferenceSpace{ - <:Any, - <:Spaces.RectilinearSpectralElementSpace2D, - }, -} +) if slice === nothing error("must specify coordinate axis slice for 3D hybrid plots") end @@ -503,15 +495,10 @@ RecipesBase.@recipe function f( end RecipesBase.@recipe function f( - field::Fields.Field{<:Any, S}; + field::Fields.ExtrudedCubedSphereSpectralElementField3D; level = nothing, hinterpolate = 10, -) where { - S <: Spaces.ExtrudedFiniteDifferenceSpace{ - <:Any, - <:Spaces.CubedSphereSpectralElementSpace2D, - }, -} +) @assert hinterpolate ≥ 1 "number of element quadrature points for uniform interpolation must be ≥ 1" space = axes(field) diff --git a/lib/ClimaCoreTempestRemap/src/netcdf.jl b/lib/ClimaCoreTempestRemap/src/netcdf.jl index d71386a0e2..c6c514aa19 100644 --- a/lib/ClimaCoreTempestRemap/src/netcdf.jl +++ b/lib/ClimaCoreTempestRemap/src/netcdf.jl @@ -53,7 +53,7 @@ def_space_coord( nc::NCDataset, space::Spaces.SpectralElementSpace2D; type = "dgll", -) = def_space_coord(nc, space, space.topology.mesh; type) +) = def_space_coord(nc, space, Spaces.topology(space).mesh; type) function def_space_coord( nc::NCDataset, @@ -224,13 +224,17 @@ end function def_space_coord( nc::NCDataset, - space::Spaces.ExtrudedFiniteDifferenceSpace{S}; + space::Spaces.ExtrudedFiniteDifferenceSpace; type = "dgll", -) where {S <: Spaces.Staggering} - hvar = def_space_coord(nc, space.horizontal_space; type = type) +) + staggering = Spaces.staggering(space) + hvar = def_space_coord(nc, Spaces.horizontal_space(space); type = type) vvar = def_space_coord( nc, - Spaces.FiniteDifferenceSpace{S}(space.vertical_topology), + Spaces.FiniteDifferenceSpace( + Spaces.vertical_topology(space), + staggering, + ), ) (hvar..., vvar...) end @@ -332,7 +336,7 @@ function Base.setindex!( ) nc = NCDataset(var) space = axes(field) - hspace = space.horizontal_space + hspace = Spaces.horizontal_space(space) if nc.attrib["node_type"] == "cgll" nodes = Spaces.unique_nodes(hspace) elseif nc.attrib["node_type"] == "dgll" diff --git a/lib/ClimaCoreTempestRemap/src/onlineremap.jl b/lib/ClimaCoreTempestRemap/src/onlineremap.jl index a2aad8748d..14cb831b00 100644 --- a/lib/ClimaCoreTempestRemap/src/onlineremap.jl +++ b/lib/ClimaCoreTempestRemap/src/onlineremap.jl @@ -71,9 +71,7 @@ function remap!( # using out_type == "cgll" if R.out_type == "cgll" topology = Spaces.topology(R.target_space) - hspace = Spaces.horizontal_space(R.target_space) - quadrature_style = hspace.quadrature_style - Spaces.dss2!(target, topology, quadrature_style) + Topologies.dss!(target, topology) end return target end @@ -131,13 +129,7 @@ function remap!(target::Fields.Field, R::LinearMap, source::Fields.Field) # using out_type == "cgll" if R.out_type == "cgll" topology = Spaces.topology(axes(target)) - hspace = Spaces.horizontal_space(axes(target)) - quadrature_style = hspace.quadrature_style - Spaces.dss2!( - Fields.field_values(target), - topology, - quadrature_style, - ) + Topologies.dss!(Fields.field_values(target), topology) end return target end @@ -170,15 +162,15 @@ function generate_map( out_type = "cgll", ) if (target_space_distr != nothing) - comms_ctx = target_space_distr.topology.context + comms_ctx = ClimaComms.context(target_space_distr) else - comms_ctx = target_space.topology.context + comms_ctx = ClimaComms.context(target_space) end if ClimaComms.iamroot(comms_ctx) # write meshes and generate weights on root process (using global indices) - write_exodus(meshfile_source, source_space.topology) - write_exodus(meshfile_target, target_space.topology) + write_exodus(meshfile_source, Spaces.topology(source_space)) + write_exodus(meshfile_target, Spaces.topology(target_space)) overlap_mesh(meshfile_overlap, meshfile_source, meshfile_target) remap_weights( weightfile, @@ -187,11 +179,11 @@ function generate_map( meshfile_overlap; in_type = in_type, in_np = Spaces.Quadratures.degrees_of_freedom( - source_space.quadrature_style, + Spaces.quadrature_style(source_space), ), out_type = out_type, out_np = Spaces.Quadratures.degrees_of_freedom( - target_space.quadrature_style, + Spaces.quadrature_style(target_space), ), ) diff --git a/lib/ClimaCoreTempestRemap/test/online_remap.jl b/lib/ClimaCoreTempestRemap/test/online_remap.jl index b91f2a4b18..d1aa95303d 100644 --- a/lib/ClimaCoreTempestRemap/test/online_remap.jl +++ b/lib/ClimaCoreTempestRemap/test/online_remap.jl @@ -33,8 +33,8 @@ function reshape_sparse_to_field!(field::Fields.Field, in_array::Array, R) # broadcast to the redundant nodes using unweighted dss topology = Spaces.topology(axes(field)) hspace = Spaces.horizontal_space(axes(field)) - quadrature_style = hspace.quadrature_style - Spaces.dss2!(Fields.field_values(field), topology, quadrature_style) + quadrature_style = Spaces.quadrature_style(hspace) + Topologies.dss!(Fields.field_values(field), topology) return field end diff --git a/lib/ClimaCoreVTK/src/ClimaCoreVTK.jl b/lib/ClimaCoreVTK/src/ClimaCoreVTK.jl index 57638b5545..e03b90e7a0 100644 --- a/lib/ClimaCoreVTK/src/ClimaCoreVTK.jl +++ b/lib/ClimaCoreVTK/src/ClimaCoreVTK.jl @@ -52,7 +52,7 @@ function WriteVTK.vtk_grid( coords = Geometry.Cartesian123Point.( Fields.coordinate_field(gridspace), - Ref(gridspace.global_geometry), + Ref(Spaces.global_geometry(gridspace)), ) coord_vecs = ( vec(parent(coords.x1)), diff --git a/lib/ClimaCoreVTK/src/space.jl b/lib/ClimaCoreVTK/src/space.jl index 884b31da72..e5ca6f25cb 100644 --- a/lib/ClimaCoreVTK/src/space.jl +++ b/lib/ClimaCoreVTK/src/space.jl @@ -83,13 +83,10 @@ Construct a vector of `MeshCell` objects representing the elements of `space` as an unstuctured mesh of Lagrange polynomial cells, suitable for passing to `vtk_grid`. """ -function vtk_cells_lagrange( - gspace::Spaces.SpectralElementSpace2D{ - T, - Spaces.Quadratures.ClosedUniform{Nq}, - }, -) where {T, Nq} - # TODO: this should depend on the backing DataLayouts (e.g. IJFH) +function vtk_cells_lagrange(gspace::Spaces.SpectralElementSpace2D) + quad = Spaces.quadrature_style(gspace) + @assert quad isa Quadratures.ClosedUniform + Nq = Quadratures.degrees_of_freedom(quad) # TODO: this should depend on the backing DataLayouts (e.g. IJFH) con_map = vtk_connectivity_map_lagrange(Nq, Nq) [ MeshCell( @@ -107,7 +104,9 @@ an unstuctured mesh of linear cells, suitable for passing to `vtk_grid`. """ function vtk_cells_linear(gridspace::Spaces.SpectralElementSpace2D) - Nq = Spaces.Quadratures.degrees_of_freedom(gridspace.quadrature_style) + Nq = Spaces.Quadratures.degrees_of_freedom( + Spaces.quadrature_style(gridspace), + ) Nh = Topologies.nlocalelems(gridspace) ind = LinearIndices((1:Nq, 1:Nq, 1:Nh)) cells = [ @@ -126,7 +125,7 @@ end function vtk_cells_linear(gridspace::Spaces.FaceExtrudedFiniteDifferenceSpace2D) hspace = Spaces.horizontal_space(gridspace) - Nq = Spaces.Quadratures.degrees_of_freedom(hspace.quadrature_style) + Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(hspace)) Nh = Topologies.nlocalelems(hspace) Nv = Spaces.nlevels(gridspace) ind = LinearIndices((1:Nv, 1:Nq, 1:Nh)) # assume VIFH @@ -145,7 +144,7 @@ function vtk_cells_linear(gridspace::Spaces.FaceExtrudedFiniteDifferenceSpace2D) end function vtk_cells_linear(gridspace::Spaces.FaceExtrudedFiniteDifferenceSpace3D) hspace = Spaces.horizontal_space(gridspace) - Nq = Spaces.Quadratures.degrees_of_freedom(hspace.quadrature_style) + Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(hspace)) Nh = Topologies.nlocalelems(hspace) Nv = Spaces.nlevels(gridspace) ind = LinearIndices((1:Nv, 1:Nq, 1:Nq, 1:Nh)) # assumes VIJFH @@ -180,25 +179,26 @@ This generally does two things: - Modifies the vertical space to be on the faces. """ function vtk_grid_space(space::Spaces.SpectralElementSpace1D) - if space.quadrature_style isa Spaces.Quadratures.ClosedUniform + if Spaces.quadrature_style(space) isa Spaces.Quadratures.ClosedUniform return space end - Nq = Spaces.Quadratures.degrees_of_freedom(space.quadrature_style) + Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(space)) lagrange_quad = Spaces.Quadratures.ClosedUniform{Nq}() - return Spaces.SpectralElementSpace1D(space.topology, lagrange_quad) + return Spaces.SpectralElementSpace1D(Spaces.topology(space), lagrange_quad) end function vtk_grid_space(space::Spaces.SpectralElementSpace2D) - if space.quadrature_style isa Spaces.Quadratures.ClosedUniform + if Spaces.quadrature_style(space) isa Spaces.Quadratures.ClosedUniform return space end - Nq = Spaces.Quadratures.degrees_of_freedom(space.quadrature_style) + Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(space)) lagrange_quad = Spaces.Quadratures.ClosedUniform{Nq}() - return Spaces.SpectralElementSpace2D(space.topology, lagrange_quad) + return Spaces.SpectralElementSpace2D(Spaces.topology(space), lagrange_quad) end function vtk_grid_space(space::Spaces.FaceExtrudedFiniteDifferenceSpace) # this will need to be updated for warped meshes horizontal_space = vtk_grid_space(Spaces.horizontal_space(space)) - vertical_space = Spaces.FaceFiniteDifferenceSpace(space.vertical_topology) + vertical_space = + Spaces.FaceFiniteDifferenceSpace(Spaces.vertical_topology(space)) return Spaces.ExtrudedFiniteDifferenceSpace( horizontal_space, vertical_space, @@ -226,22 +226,28 @@ This generally does two things: - Modifies the vertical space to be on the centers. """ function vtk_cell_space(gridspace::Spaces.SpectralElementSpace1D) - @assert gridspace.quadrature_style isa Spaces.Quadratures.ClosedUniform - Nq = Spaces.Quadratures.degrees_of_freedom(gridspace.quadrature_style) + @assert Spaces.quadrature_style(gridspace) isa + Spaces.Quadratures.ClosedUniform + Nq = Spaces.Quadratures.degrees_of_freedom( + Spaces.quadrature_style(gridspace), + ) quad = Spaces.Quadratures.Uniform{Nq - 1}() - return Spaces.SpectralElementSpace1D(gridspace.topology, quad) + return Spaces.SpectralElementSpace1D(Spaces.topology(gridspace), quad) end function vtk_cell_space(gridspace::Spaces.SpectralElementSpace2D) - @assert gridspace.quadrature_style isa Spaces.Quadratures.ClosedUniform - Nq = Spaces.Quadratures.degrees_of_freedom(gridspace.quadrature_style) + @assert Spaces.quadrature_style(gridspace) isa + Spaces.Quadratures.ClosedUniform + Nq = Spaces.Quadratures.degrees_of_freedom( + Spaces.quadrature_style(gridspace), + ) quad = Spaces.Quadratures.Uniform{Nq - 1}() - return Spaces.SpectralElementSpace2D(gridspace.topology, quad) + return Spaces.SpectralElementSpace2D(Spaces.topology(gridspace), quad) end function vtk_cell_space(gridspace::Spaces.FaceExtrudedFiniteDifferenceSpace) # this will need to be updated for warped meshes horizontal_space = vtk_cell_space(Spaces.horizontal_space(gridspace)) vertical_space = - Spaces.CenterFiniteDifferenceSpace(gridspace.vertical_topology) + Spaces.CenterFiniteDifferenceSpace(Spaces.vertical_topology(gridspace)) return Spaces.ExtrudedFiniteDifferenceSpace( horizontal_space, vertical_space, diff --git a/perf/Manifest.toml b/perf/Manifest.toml index 89288bb4d7..645f12352f 100644 --- a/perf/Manifest.toml +++ b/perf/Manifest.toml @@ -217,7 +217,7 @@ uuid = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" version = "0.5.5" [[deps.ClimaCore]] -deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack"] +deps = ["Adapt", "BandedMatrices", "BlockArrays", "CUDA", "ClimaComms", "CubedSphere", "DataStructures", "DocStringExtensions", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "LinearAlgebra", "Memoize", "PkgVersion", "RecursiveArrayTools", "Requires", "RootSolvers", "SparseArrays", "Static", "StaticArrays", "Statistics", "UnPack", "WeakValueDicts"] path = ".." uuid = "d414da3d-4745-48bb-8d80-42e94e092884" version = "0.10.55" @@ -1260,6 +1260,12 @@ git-tree-sha1 = "c13304c81eec1ed3af7fc20e75fb6b26092a1102" uuid = "442fdcdd-2543-5da2-b0f3-8c86c306513e" version = "0.3.2" +[[deps.Memoize]] +deps = ["MacroTools"] +git-tree-sha1 = "2b1dfcba103de714d31c033b5dacc2e4a12c7caa" +uuid = "c03570c3-d221-55d1-a50c-7939bbd78826" +version = "0.4.4" + [[deps.MicrosoftMPI_jll]] deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] git-tree-sha1 = "a7023883872e52bc29bcaac74f19adf39347d2d5" @@ -2128,6 +2134,11 @@ git-tree-sha1 = "4528479aa01ee1b3b4cd0e6faef0e04cf16466da" uuid = "2381bf8a-dfd0-557d-9999-79630e7b1b91" version = "1.25.0+0" +[[deps.WeakValueDicts]] +git-tree-sha1 = "98528c2610a5479f091d470967a25becfd83edd0" +uuid = "897b6980-f191-5a31-bcb0-bf3c4585e0c1" +version = "0.1.0" + [[deps.WidthLimitedIO]] deps = ["Unicode"] git-tree-sha1 = "71142739e695823729a335e9bc124ef41ec1433a" diff --git a/src/ClimaCore.jl b/src/ClimaCore.jl index 72ada48e34..97c94deaae 100644 --- a/src/ClimaCore.jl +++ b/src/ClimaCore.jl @@ -11,6 +11,8 @@ include("Geometry/Geometry.jl") include("Domains/Domains.jl") include("Meshes/Meshes.jl") include("Topologies/Topologies.jl") +include("Quadratures/Quadratures.jl") +include("Grids/Grids.jl") include("Spaces/Spaces.jl") include("Fields/Fields.jl") include("Operators/Operators.jl") diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl index ab2607817d..3405f7f4b0 100644 --- a/src/Fields/Fields.jl +++ b/src/Fields/Fields.jl @@ -5,6 +5,7 @@ import ..slab, ..slab_args, ..column, ..column_args, ..level import ..DataLayouts: DataLayouts, AbstractData, DataStyle import ..Domains import ..Topologies +import ..Grids: ColumnIndex import ..Spaces: Spaces, AbstractSpace, AbstractPointSpace import ..Geometry: Geometry, Cartesian12Vector import ..Utilities: PlusHalf, half @@ -38,13 +39,6 @@ Field(::Type{T}, space::S) where {T, S <: AbstractSpace} = ClimaComms.context(field::Field) = ClimaComms.context(axes(field)) -ClimaComms.context(space::Spaces.ExtrudedFiniteDifferenceSpace) = - ClimaComms.context(Spaces.horizontal_space(space)) -ClimaComms.context(space::Spaces.SpectralElementSpace2D) = - ClimaComms.context(space.topology) -ClimaComms.context(space::S) where {S <: Spaces.AbstractSpace} = - ClimaComms.context(space.topology) - ClimaComms.context(topology::Topologies.Topology2D) = topology.context ClimaComms.context(topology::T) where {T <: Topologies.AbstractTopology} = topology.context @@ -54,7 +48,7 @@ Adapt.adapt_structure(to, field::Field) = Field( Adapt.adapt(to, axes(field)), ) - +## aliases # Point Field const PointField{V, S} = Field{V, S} where {V <: AbstractData, S <: Spaces.PointSpace} @@ -83,6 +77,14 @@ const ExtrudedFiniteDifferenceField{V, S} = Field{ V, S, } where {V <: AbstractData, S <: Spaces.ExtrudedFiniteDifferenceSpace} +const ExtrudedFiniteDifferenceField2D{V, S} = Field{ + V, + S, +} where {V <: AbstractData, S <: Spaces.ExtrudedFiniteDifferenceSpace2D} +const ExtrudedFiniteDifferenceField3D{V, S} = Field{ + V, + S, +} where {V <: AbstractData, S <: Spaces.ExtrudedFiniteDifferenceSpace3D} const FaceExtrudedFiniteDifferenceField{V, S} = Field{ V, S, @@ -92,12 +94,40 @@ const CenterExtrudedFiniteDifferenceField{V, S} = Field{ S, } where {V <: AbstractData, S <: Spaces.CenterExtrudedFiniteDifferenceSpace} +# +const SpectralElementField1D{V, S} = + Field{V, S} where {V <: AbstractData, S <: Spaces.SpectralElementSpace1D} +const ExtrudedSpectralElementField2D{V, S} = Field{ + V, + S, +} where {V <: AbstractData, S <: Spaces.ExtrudedSpectralElementSpace2D} + +const RectilinearSpectralElementField2D{V, S} = Field{ + V, + S, +} where {V <: AbstractData, S <: Spaces.RectilinearSpectralElementSpace2D} +const ExtrudedRectilinearSpectralElementField3D{V, S} = Field{ + V, + S, +} where { + V <: AbstractData, + S <: Spaces.ExtrudedRectilinearSpectralElementSpace3D, +} + + # Cubed Sphere Fields + const CubedSphereSpectralElementField2D{V, S} = Field{ V, S, } where {V <: AbstractData, S <: Spaces.CubedSphereSpectralElementSpace2D} - +const ExtrudedCubedSphereSpectralElementField3D{V, S} = Field{ + V, + S, +} where { + V <: AbstractData, + S <: Spaces.ExtrudedCubedSphereSpectralElementSpace3D, +} Base.propertynames(field::Field) = propertynames(getfield(field, :values)) @inline field_values(field::Field) = getfield(field, :values) @@ -333,22 +363,12 @@ function interpcoord(elemrange, x::Real) return z, ξ end -""" - Spaces.variational_solve!(field) - -Divide `field` by the mass matrix. -""" -function Spaces.variational_solve!(field::Field) - Spaces.variational_solve!(field_values(field), axes(field)) - return field -end - """ Spaces.weighted_dss!(f::Field[, ghost_buffer = Spaces.create_dss_buffer(field)]) Apply weighted direct stiffness summation (DSS) to `f`. This operates in-place (i.e. it modifies the `f`). `ghost_buffer` contains the necessary information -for communication in a distributed setting, see [`Spaces.create_ghost_buffer`](@ref). +for communication in a distributed setting, see [`Spaces.create_dss_buffer`](@ref). This is a projection operation from the piecewise polynomial space ``\\mathcal{V}_0`` to the continuous space ``\\mathcal{V}_1 = \\mathcal{V}_0 @@ -389,13 +409,6 @@ Spaces.weighted_dss_internal!(field::Field, dss_buffer) = Spaces.weighted_dss_ghost!(field::Field, dss_buffer) = Spaces.weighted_dss_ghost!(field_values(field), axes(field), dss_buffer) -""" - Spaces.create_ghost_buffer(field::Field) - -Create a buffer for communicating neighbour information of `field`. -""" -Spaces.create_ghost_buffer(field::Field) = Spaces.create_dss_buffer(field) - """ Spaces.create_dss_buffer(field::Field) @@ -406,20 +419,6 @@ function Spaces.create_dss_buffer(field::Field) hspace = Spaces.horizontal_space(space) Spaces.create_dss_buffer(field_values(field), hspace) end -# Add definitions for backward compatibility -Spaces.weighted_dss2!( - field::Field, - dss_buffer = Spaces.create_dss_buffer(field), -) = Spaces.weighted_dss!(field, dss_buffer) - -Spaces.weighted_dss_start2!(field::Field, ghost_buffer) = - Spaces.weighted_dss_start!(field, ghost_buffer) - -Spaces.weighted_dss_internal2!(field::Field, ghost_buffer) = - Spaces.weighted_dss_internal!(field, ghost_buffer) - -Spaces.weighted_dss_ghost2!(field, ghost_buffer) = - Spaces.weighted_dss_ghost!(field, ghost_buffer) Base.@propagate_inbounds function level( field::Union{ diff --git a/src/Fields/broadcast.jl b/src/Fields/broadcast.jl index a5f5756373..f20235d255 100644 --- a/src/Fields/broadcast.jl +++ b/src/Fields/broadcast.jl @@ -153,31 +153,8 @@ end return dest end -allow_mismatched_diagonalized_spaces() = false - -@noinline function warn_mismatched_spaces( - space1::Type{S}, - space2::Type{S}, -) where {S <: AbstractSpace} - @warn "Broacasted spaces are the same ClimaCore.Spaces type but not the same instance" - return nothing -end - -is_diagonalized_spaces(::Type{S}, ::Type{S}) where {S <: AbstractSpace} = true - -is_diagonalized_spaces(::Type, ::Type) = false - -@noinline function error_mismatched_spaces( - space1::Type{S}, - space2::Type{S}, -) where {S <: AbstractSpace} - error( - "Broacasted spaces are the same ClimaCore.Spaces type but not the same instance", - ) -end - @noinline function error_mismatched_spaces(space1::Type, space2::Type) - error("Broacasted spaces are not the same ClimaCore.Spaces type") + error("Broacasted spaces are not the same.") end @inline function Base.Broadcast.broadcast_shape( @@ -185,14 +162,7 @@ end space2::AbstractSpace, ) if space1 !== space2 - if is_diagonalized_spaces(typeof(space1), typeof(space2)) && - allow_mismatched_diagonalized_spaces() && - ( - parent(Spaces.local_geometry_data(space1)) == - parent(Spaces.local_geometry_data(space2)) - ) - warn_mismatched_spaces(typeof(space1), typeof(space2)) - elseif Spaces.issubspace(space2, space1) + if Spaces.issubspace(space2, space1) return space1 elseif Spaces.issubspace(space1, space2) return space2 @@ -226,15 +196,8 @@ end space2::AbstractSpace, ) if space1 !== space2 - if is_diagonalized_spaces(typeof(space1), typeof(space2)) && - allow_mismatched_diagonalized_spaces() && - ( - parent(Spaces.local_geometry_data(space1)) == - parent(Spaces.local_geometry_data(space2)) - ) - warn_mismatched_spaces(typeof(space1), typeof(space2)) - elseif Spaces.issubspace(space2, space1) || - Spaces.issubspace(space1, space2) + if Spaces.issubspace(space2, space1) || + Spaces.issubspace(space1, space2) nothing else error_mismatched_spaces(typeof(space1), typeof(space2)) @@ -419,7 +382,7 @@ function Base.Broadcast.broadcasted( fs, V, arg, - tuple(space.global_geometry), + tuple(Spaces.global_geometry(space)), local_geometry_field(space), ) end @@ -449,3 +412,12 @@ function Base.Broadcast.copyto!(field::Field, nt::NamedTuple) ), ) end + + +# TODO: deprecate these + +allow_mismatched_diagonalized_spaces() = false + +is_diagonalized_spaces(::Type{S}, ::Type{S}) where {S <: AbstractSpace} = true + +is_diagonalized_spaces(::Type, ::Type) = false diff --git a/src/Fields/indices.jl b/src/Fields/indices.jl index 92bfcc2b5f..89cac213b7 100644 --- a/src/Fields/indices.jl +++ b/src/Fields/indices.jl @@ -1,19 +1,3 @@ -""" - ColumnIndex(ij,h) - -An index into a column of a field. This can be used as an argument to `getindex` -of a `Field`, to return a field on that column. - -# Example -```julia -colidx = ColumnIndex((1,1),1) -field[colidx] -``` -""" -struct ColumnIndex{N} - ij::NTuple{N, Int} - h::Int -end Base.@propagate_inbounds Base.getindex(field::Field, colidx::ColumnIndex) = column(field, colidx) @@ -145,7 +129,7 @@ bycolumn( fn, space::Spaces.ExtrudedFiniteDifferenceSpace, device::ClimaComms.AbstractCPUDevice, -) = bycolumn(fn, space.horizontal_space, device) +) = bycolumn(fn, Spaces.horizontal_space(space), device) function bycolumn(fn, space::AbstractSpace, ::ClimaComms.CUDADevice) @@ -164,7 +148,7 @@ Number of columns in a given space. ncolumns(field::Field) = ncolumns(axes(field)) ncolumns(space::Spaces.ExtrudedFiniteDifferenceSpace) = - ncolumns(space.horizontal_space) + ncolumns(Spaces.horizontal_space(space)) function ncolumns(space::Spaces.SpectralElementSpace1D) Nh = Topologies.nlocalelems(space) @@ -219,7 +203,7 @@ function byslab( ::ClimaComms.CPUSingleThreaded, space::Spaces.AbstractSpectralElementSpace, ) - Nh = Topologies.nlocalelems(space.topology)::Int + Nh = Topologies.nlocalelems(Spaces.topology(space))::Int @inbounds for h in 1:Nh fn(SlabIndex(nothing, h)) end @@ -230,7 +214,7 @@ function byslab( ::ClimaComms.CPUMultiThreaded, space::Spaces.AbstractSpectralElementSpace, ) - Nh = Topologies.nlocalelems(space.topology)::Int + Nh = Topologies.nlocalelems(Spaces.topology(space))::Int @inbounds begin Threads.@threads for h in 1:Nh fn(SlabIndex(nothing, h)) diff --git a/src/Fields/mapreduce_cuda.jl b/src/Fields/mapreduce_cuda.jl index 03ab059be0..61f0f7aaeb 100644 --- a/src/Fields/mapreduce_cuda.jl +++ b/src/Fields/mapreduce_cuda.jl @@ -220,15 +220,6 @@ end return (Nv, Nij, Nf, Nh) end -@inline function _get_idxs(Nv, Nij, Nf, Nh, fidx, gidx) - hidx = cld(gidx, Nv * Nij * Nij * Nf) - offset = ((hidx - 1) * Nf + (fidx - 1)) * Nv * Nij * Nij - jidx = cld(gidx - offset, Nv * Nij) - offset += (jidx - 1) * Nv * Nij - iidx = cld(gidx - offset, Nv) - return (iidx, jidx, hidx) -end - @inline function _cuda_reduce!(op, reduction, tidx, reduction_size, N) if reduction_size > N if tidx ≤ reduction_size - N diff --git a/src/Geometry/localgeometry.jl b/src/Geometry/localgeometry.jl index 98010215c6..098ef94f48 100644 --- a/src/Geometry/localgeometry.jl +++ b/src/Geometry/localgeometry.jl @@ -51,3 +51,89 @@ end undertype(::Type{LocalGeometry{I, C, FT, S}}) where {I, C, FT, S} = FT undertype(::Type{SurfaceGeometry{FT, N}}) where {FT, N} = FT + + +function blockmat( + a::Geometry.Axis2Tensor{ + FT, + Tuple{Geometry.UAxis, Geometry.Covariant1Axis}, + SMatrix{1, 1, FT, 1}, + }, + b::Geometry.Axis2Tensor{ + FT, + Tuple{Geometry.WAxis, Geometry.Covariant3Axis}, + SMatrix{1, 1, FT, 1}, + }, +) where {FT} + A = Geometry.components(a) + B = Geometry.components(b) + Geometry.AxisTensor( + (Geometry.UWAxis(), Geometry.Covariant13Axis()), + SMatrix{2, 2}(A[1, 1], zero(FT), zero(FT), B[1, 1]), + ) +end + +function blockmat( + a::Geometry.Axis2Tensor{ + FT, + Tuple{Geometry.VAxis, Geometry.Covariant2Axis}, + SMatrix{1, 1, FT, 1}, + }, + b::Geometry.Axis2Tensor{ + FT, + Tuple{Geometry.WAxis, Geometry.Covariant3Axis}, + SMatrix{1, 1, FT, 1}, + }, +) where {FT} + A = Geometry.components(a) + B = Geometry.components(b) + Geometry.AxisTensor( + (Geometry.VWAxis(), Geometry.Covariant23Axis()), + SMatrix{2, 2}(A[1, 1], zero(FT), zero(FT), B[1, 1]), + ) +end + +function blockmat( + a::Geometry.Axis2Tensor{ + FT, + Tuple{Geometry.UVAxis, Geometry.Covariant12Axis}, + SMatrix{2, 2, FT, 4}, + }, + b::Geometry.Axis2Tensor{ + FT, + Tuple{Geometry.WAxis, Geometry.Covariant3Axis}, + SMatrix{1, 1, FT, 1}, + }, +) where {FT} + A = Geometry.components(a) + B = Geometry.components(b) + Geometry.AxisTensor( + (Geometry.UVWAxis(), Geometry.Covariant123Axis()), + SMatrix{3, 3}( + A[1, 1], + A[2, 1], + zero(FT), + A[1, 2], + A[2, 2], + zero(FT), + zero(FT), + zero(FT), + B[1, 1], + ), + ) +end + +function product_geometry( + horizontal_local_geometry::Geometry.LocalGeometry, + vertical_local_geometry::Geometry.LocalGeometry, +) + coordinates = Geometry.product_coordinates( + horizontal_local_geometry.coordinates, + vertical_local_geometry.coordinates, + ) + J = horizontal_local_geometry.J * vertical_local_geometry.J + WJ = horizontal_local_geometry.WJ * vertical_local_geometry.WJ + ∂x∂ξ = + blockmat(horizontal_local_geometry.∂x∂ξ, vertical_local_geometry.∂x∂ξ) + return Geometry.LocalGeometry(coordinates, J, WJ, ∂x∂ξ) +end diff --git a/src/Grids/Grids.jl b/src/Grids/Grids.jl new file mode 100644 index 0000000000..7d594d0d8d --- /dev/null +++ b/src/Grids/Grids.jl @@ -0,0 +1,73 @@ +module Grids + +import ClimaComms, Adapt, ForwardDiff, LinearAlgebra +import LinearAlgebra: det, norm +import Memoize: @memoize +import WeakValueDicts: WeakValueDict +import ..DataLayouts, + ..Domains, ..Meshes, ..Topologies, ..Geometry, ..Quadratures +import ..Utilities: PlusHalf, half +import ..slab, ..column, ..level + +using StaticArrays + +""" + Grids.AbstractGrid + +Grids should define the following + + +- [`topology`](@ref): the topology of the grid +- [`mesh`](@ref): the mesh of the grid +- [`domain`](@ref): the domain of the grid +- `ClimaComms.context` +- `ClimaComms.device` + +- [`local_geometry_data`](@ref): the `DataLayout` object containing the local geometry data information + +""" +abstract type AbstractGrid end + + +""" + Grids.topology(grid::AbstractGrid) + +Get the topology of a grid. +""" +function topology end + +function vertical_topology end + +""" + Grids.local_geometry_data( + grid :: AbstractGrid, + staggering :: Union{Staggering, Nothing}, + ) + +Get the `DataLayout` object containing the local geometry data information of +the `grid` with staggering `staggering`. + +If the grid is not staggered, `staggering` should be `nothing`. +""" +function local_geometry_data end + +function local_dss_weights end +function quadrature_style end +function vertical_topology end + + + +ClimaComms.context(grid::AbstractGrid) = ClimaComms.context(topology(grid)) +ClimaComms.device(grid::AbstractGrid) = ClimaComms.device(topology(grid)) + +Meshes.domain(grid::AbstractGrid) = Meshes.domain(topology(grid)) + +include("finitedifference.jl") +include("spectralelement.jl") +include("extruded.jl") +include("column.jl") +include("level.jl") + + + +end # module diff --git a/src/Grids/column.jl b/src/Grids/column.jl new file mode 100644 index 0000000000..a85a940866 --- /dev/null +++ b/src/Grids/column.jl @@ -0,0 +1,51 @@ + + + + +""" + ColumnIndex(ij,h) + +An index into a column of a field. This can be used as an argument to `getindex` +of a `Field`, to return a field on that column. + +# Example +```julia +colidx = ColumnIndex((1,1),1) +field[colidx] +``` +""" +struct ColumnIndex{N} + ij::NTuple{N, Int} + h::Int +end + + +""" + ColumnGrid( + full_grid :: ExtrudedFiniteDifferenceGrid, + colidx :: ColumnIndex, + ) + +A view into a column of a `ExtrudedFiniteDifferenceGrid`. This can be used as an +""" +struct ColumnGrid{ + G <: AbstractExtrudedFiniteDifferenceGrid, + C <: ColumnIndex, +} <: AbstractFiniteDifferenceGrid + full_grid::G + colidx::C +end + + +column(grid::AbstractExtrudedFiniteDifferenceGrid, colidx::ColumnIndex) = + ColumnGrid(grid, colidx) + +topology(colgrid::ColumnGrid) = vertical_topology(colgrid.full_grid) +vertical_topology(colgrid::ColumnGrid) = vertical_topology(colgrid.full_grid) + +local_geometry_data(colgrid::ColumnGrid, staggering::Staggering) = column( + local_geometry_data(colgrid.full_grid, staggering::Staggering), + colgrid.colidx.ij..., + colgrid.colidx.h, +) +global_geometry(colgrid::ColumnGrid) = global_geometry(colgrid.full_grid) diff --git a/src/Grids/extruded.jl b/src/Grids/extruded.jl new file mode 100644 index 0000000000..9f10cef7a8 --- /dev/null +++ b/src/Grids/extruded.jl @@ -0,0 +1,120 @@ +##### +##### Hybrid mesh +##### + +abstract type HypsographyAdaption end + +""" + Flat() + +No surface hypsography. +""" +struct Flat <: HypsographyAdaption end + +abstract type AbstractExtrudedFiniteDifferenceGrid <: AbstractGrid end + +""" + ExtrudedFiniteDifferenceGrid( + horizontal_space::AbstractSpace, + vertical_space::FiniteDifferenceSpace, + hypsography::HypsographyAdaption = Flat(), + ) + +Construct an `ExtrudedFiniteDifferenceGrid` from the horizontal and vertical spaces. +""" +mutable struct ExtrudedFiniteDifferenceGrid{ + H <: AbstractGrid, + V <: FiniteDifferenceGrid, + A <: HypsographyAdaption, + GG <: Geometry.AbstractGlobalGeometry, + LG, +} <: AbstractExtrudedFiniteDifferenceGrid + horizontal_grid::H + vertical_grid::V + hypsography::A + global_geometry::GG + center_local_geometry::LG + face_local_geometry::LG +end + +@memoize WeakValueDict function ExtrudedFiniteDifferenceGrid( + horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D}, + vertical_grid::FiniteDifferenceGrid, + hypsography::Flat = Flat(), +) + global_geometry = horizontal_grid.global_geometry + center_local_geometry = + Geometry.product_geometry.( + horizontal_grid.local_geometry, + vertical_grid.center_local_geometry, + ) + face_local_geometry = + Geometry.product_geometry.( + horizontal_grid.local_geometry, + vertical_grid.face_local_geometry, + ) + + return ExtrudedFiniteDifferenceGrid( + horizontal_grid, + vertical_grid, + hypsography, + global_geometry, + center_local_geometry, + face_local_geometry, + ) +end + +topology(grid::ExtrudedFiniteDifferenceGrid) = topology(grid.horizontal_grid) + +vertical_topology(grid::ExtrudedFiniteDifferenceGrid) = + topology(grid.vertical_grid) + +local_dss_weights(grid::ExtrudedFiniteDifferenceGrid) = + local_dss_weights(grid.horizontal_grid) + + +local_geometry_data(grid::AbstractExtrudedFiniteDifferenceGrid, ::CellCenter) = + grid.center_local_geometry +local_geometry_data(grid::AbstractExtrudedFiniteDifferenceGrid, ::CellFace) = + grid.face_local_geometry +global_geometry(grid::AbstractExtrudedFiniteDifferenceGrid) = + grid.global_geometry + +quadrature_style(grid::ExtrudedFiniteDifferenceGrid) = + quadrature_style(grid.horizontal_grid) + + +## GPU compatibility +struct DeviceExtrudedFiniteDifferenceGrid{VT, Q, GG, LG} <: + AbstractExtrudedFiniteDifferenceGrid + vertical_topology::VT + quadrature_style::Q + global_geometry::GG + center_local_geometry::LG + face_local_geometry::LG +end + +Adapt.adapt_structure(to, grid::ExtrudedFiniteDifferenceGrid) = + DeviceExtrudedFiniteDifferenceGrid( + Adapt.adapt(to, vertical_topology(grid)), + 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), + ) + +quadrature_style(grid::DeviceExtrudedFiniteDifferenceGrid) = + grid.quadrature_style +vertical_topology(grid::DeviceExtrudedFiniteDifferenceGrid) = + grid.vertical_topology + +## aliases + +const ExtrudedSpectralElementGrid2D = + ExtrudedFiniteDifferenceGrid{<:SpectralElementGrid1D} +const ExtrudedSpectralElementGrid3D = + ExtrudedFiniteDifferenceGrid{<:SpectralElementGrid2D} +const ExtrudedRectilinearSpectralElementGrid3D = + ExtrudedFiniteDifferenceGrid{<:RectilinearSpectralElementGrid2D} +const ExtrudedCubedSphereSpectralElementGrid3D = + ExtrudedFiniteDifferenceGrid{<:CubedSphereSpectralElementGrid2D} diff --git a/src/Grids/finitedifference.jl b/src/Grids/finitedifference.jl new file mode 100644 index 0000000000..ce3204d07f --- /dev/null +++ b/src/Grids/finitedifference.jl @@ -0,0 +1,175 @@ + +abstract type Staggering end + +""" + CellCenter() + +Cell center location +""" +struct CellCenter <: Staggering end + +""" + CellFace() + +Cell face location +""" +struct CellFace <: Staggering end + + + + + + +abstract type AbstractFiniteDifferenceGrid <: AbstractGrid end + +""" + FiniteDifferenceGrid(topology::Topologies.IntervalTopology) + FiniteDifferenceGrid(mesh::Meshes.IntervalMesh) + +Construct a `FiniteDifferenceGrid` from an `IntervalTopology` (or an +`IntervalMesh`). + +This is an object which contains all the necessary geometric information. + +To avoid unnecessary duplication, we memoize the construction of the grid. +""" +mutable struct FiniteDifferenceGrid{ + T <: Topologies.AbstractIntervalTopology, + GG, + LG, +} <: AbstractFiniteDifferenceGrid + topology::T + global_geometry::GG + center_local_geometry::LG + face_local_geometry::LG +end + + +@memoize WeakValueDict function FiniteDifferenceGrid( + topology::Topologies.IntervalTopology, +) + global_geometry = Geometry.CartesianGlobalGeometry() + mesh = topology.mesh + CT = Meshes.coordinate_type(mesh) + AIdx = Geometry.coordinate_axis(CT) + # TODO: FD operators hardcoded to work over the 3-axis, need to generalize + # similar to spectral operators + @assert AIdx == (3,) "FiniteDifference operations only work over the 3-axis (ZPoint) domain" + FT = eltype(CT) + ArrayType = ClimaComms.array_type(topology) + face_coordinates = collect(mesh.faces) + LG = Geometry.LocalGeometry{AIdx, CT, FT, SMatrix{1, 1, FT, 1}} + nface = length(face_coordinates) - Topologies.isperiodic(topology) + ncent = length(face_coordinates) - 1 + # contstruct on CPU, copy to device at end + center_local_geometry = DataLayouts.VF{LG}(Array{FT}, ncent) + face_local_geometry = DataLayouts.VF{LG}(Array{FT}, nface) + for i in 1:ncent + # centers + coord⁻ = Geometry.component(face_coordinates[i], 1) + coord⁺ = Geometry.component(face_coordinates[i + 1], 1) + # at the moment we use a "discrete Jacobian" + # ideally we should use the continuous quantity via the derivative of the warp function + # could we just define this then as deriv on the mesh element coordinates? + coord = (coord⁺ + coord⁻) / 2 + Δcoord = coord⁺ - coord⁻ + J = Δcoord + WJ = Δcoord + ∂x∂ξ = SMatrix{1, 1}(J) + center_local_geometry[i] = Geometry.LocalGeometry( + CT(coord), + J, + WJ, + Geometry.AxisTensor( + (Geometry.LocalAxis{AIdx}(), Geometry.CovariantAxis{AIdx}()), + ∂x∂ξ, + ), + ) + end + for i in 1:nface + coord = Geometry.component(face_coordinates[i], 1) + if i == 1 + # bottom face + if Topologies.isperiodic(topology) + Δcoord⁺ = + Geometry.component(face_coordinates[2], 1) - + Geometry.component(face_coordinates[1], 1) + Δcoord⁻ = + Geometry.component(face_coordinates[end], 1) - + Geometry.component(face_coordinates[end - 1], 1) + J = (Δcoord⁺ + Δcoord⁻) / 2 + WJ = J + else + coord⁺ = Geometry.component(face_coordinates[2], 1) + J = coord⁺ - coord + WJ = J / 2 + end + elseif !Topologies.isperiodic(topology) && i == nface + # top face + coord⁻ = Geometry.component(face_coordinates[i - 1], 1) + J = coord - coord⁻ + WJ = J / 2 + else + coord⁺ = Geometry.component(face_coordinates[i + 1], 1) + coord⁻ = Geometry.component(face_coordinates[i - 1], 1) + J = (coord⁺ - coord⁻) / 2 + WJ = J + end + ∂x∂ξ = SMatrix{1, 1}(J) + ∂ξ∂x = SMatrix{1, 1}(inv(J)) + face_local_geometry[i] = Geometry.LocalGeometry( + CT(coord), + J, + WJ, + Geometry.AxisTensor( + (Geometry.LocalAxis{AIdx}(), Geometry.CovariantAxis{AIdx}()), + ∂x∂ξ, + ), + ) + end + return FiniteDifferenceGrid( + topology, + global_geometry, + Adapt.adapt(ArrayType, center_local_geometry), + Adapt.adapt(ArrayType, face_local_geometry), + ) +end + + +FiniteDifferenceGrid(mesh::Meshes.IntervalMesh) = + FiniteDifferenceGrid(Topologies.IntervalTopology(mesh)) + +# accessors +topology(grid::FiniteDifferenceGrid) = grid.topology +vertical_topology(grid::FiniteDifferenceGrid) = grid.topology + +local_geometry_data(grid::FiniteDifferenceGrid, ::CellCenter) = + grid.center_local_geometry +local_geometry_data(grid::FiniteDifferenceGrid, ::CellFace) = + grid.face_local_geometry +global_geometry(grid::FiniteDifferenceGrid) = grid.global_geometry + +## GPU compatibility +struct DeviceFiniteDifferenceGrid{T, GG, LG} <: AbstractFiniteDifferenceGrid + topology::T + global_geometry::GG + center_local_geometry::LG + face_local_geometry::LG +end + +Adapt.adapt_structure(to, grid::FiniteDifferenceGrid) = + DeviceFiniteDifferenceGrid( + Adapt.adapt(to, grid.topology), + Adapt.adapt(to, grid.global_geometry), + Adapt.adapt(to, grid.center_local_geometry), + Adapt.adapt(to, grid.face_local_geometry), + ) + +topology(grid::DeviceFiniteDifferenceGrid) = grid.topology +vertical_topology(grid::DeviceFiniteDifferenceGrid) = grid.topology + +local_geometry_data(grid::DeviceFiniteDifferenceGrid, ::CellCenter) = + grid.center_local_geometry +local_geometry_data(grid::DeviceFiniteDifferenceGrid, ::CellFace) = + grid.face_local_geometry +global_geometry(grid::DeviceFiniteDifferenceGrid) = grid.global_geometry diff --git a/src/Grids/level.jl b/src/Grids/level.jl new file mode 100644 index 0000000000..b6f5849b35 --- /dev/null +++ b/src/Grids/level.jl @@ -0,0 +1,32 @@ +struct LevelGrid{ + G <: AbstractExtrudedFiniteDifferenceGrid, + L <: Union{Int, PlusHalf{Int}}, +} <: AbstractGrid + full_grid::G + level::L +end + +quadrature_style(levelgrid::LevelGrid) = + quadrature_style(levelgrid.full_grid.horizontal_grid) + +level( + grid::AbstractExtrudedFiniteDifferenceGrid, + level::Union{Int, PlusHalf{Int}}, +) = LevelGrid(grid, level) + +topology(levelgrid::LevelGrid) = topology(levelgrid.full_grid) + +local_geometry_data(levelgrid::LevelGrid{<:Any, Int}, ::Nothing) = level( + local_geometry_data(levelgrid.full_grid, CellCenter()), + levelgrid.level, +) +local_geometry_data(levelgrid::LevelGrid{<:Any, PlusHalf{Int}}, ::Nothing) = + level( + local_geometry_data(levelgrid.full_grid, CellFace()), + levelgrid.level + half, + ) +global_geometry(levlgrid::LevelGrid) = global_geometry(levlgrid.full_grid) + +## GPU compatibility +Adapt.adapt_structure(to, grid::LevelGrid) = + LevelGrid(Adapt.adapt(to, grid.full_grid), grid.level) diff --git a/src/Grids/spectralelement.jl b/src/Grids/spectralelement.jl new file mode 100644 index 0000000000..d46cafe955 --- /dev/null +++ b/src/Grids/spectralelement.jl @@ -0,0 +1,501 @@ + + +abstract type AbstractSpectralElementGrid <: AbstractGrid end + +""" + SpectralElementGrid1D(mesh::Meshes.IntervalMesh, quadrature_style::Quadratures.QuadratureStyle) + +A one-dimensional space: within each element the space is represented as a polynomial. +""" +mutable struct SpectralElementGrid1D{ + T, + Q, + GG <: Geometry.AbstractGlobalGeometry, + LG, + D, +} <: AbstractSpectralElementGrid + topology::T + quadrature_style::Q + global_geometry::GG + local_geometry::LG + dss_weights::D +end + +@memoize WeakValueDict function SpectralElementGrid1D( + topology::Topologies.IntervalTopology, + quadrature_style::Quadratures.QuadratureStyle, +) + global_geometry = Geometry.CartesianGlobalGeometry() + CoordType = Topologies.coordinate_type(topology) + AIdx = Geometry.coordinate_axis(CoordType) + FT = eltype(CoordType) + nelements = Topologies.nlocalelems(topology) + Nq = Quadratures.degrees_of_freedom(quadrature_style) + + LG = Geometry.LocalGeometry{AIdx, CoordType, FT, SMatrix{1, 1, FT, 1}} + local_geometry = DataLayouts.IFH{LG, Nq}(Array{FT}, nelements) + quad_points, quad_weights = + Quadratures.quadrature_points(FT, quadrature_style) + + for elem in 1:nelements + local_geometry_slab = slab(local_geometry, elem) + for i in 1:Nq + ξ = quad_points[i] + # TODO: we need to massage the coordinate points because the grid is assumed 2D + vcoords = Topologies.vertex_coordinates(topology, elem) + x = Geometry.linear_interpolate(vcoords, ξ) + ∂x∂ξ = + ( + Geometry.component(vcoords[2], 1) - + Geometry.component(vcoords[1], 1) + ) / 2 + J = abs(∂x∂ξ) + WJ = J * quad_weights[i] + local_geometry_slab[i] = Geometry.LocalGeometry( + x, + J, + WJ, + Geometry.AxisTensor( + ( + Geometry.LocalAxis{AIdx}(), + Geometry.CovariantAxis{AIdx}(), + ), + ∂x∂ξ, + ), + ) + end + end + dss_weights = copy(local_geometry.J) + dss_weights .= one(FT) + Topologies.dss_1d!(topology, dss_weights) + dss_weights = one(FT) ./ dss_weights + + return SpectralElementGrid1D( + topology, + quadrature_style, + global_geometry, + local_geometry, + dss_weights, + ) +end + + + +""" + SpectralElementSpace2D <: AbstractSpace + +A two-dimensional space: within each element the space is represented as a polynomial. +""" +mutable struct SpectralElementGrid2D{ + T, + Q, + GG <: Geometry.AbstractGlobalGeometry, + LG, + D, + IS, + BS, +} <: AbstractSpectralElementGrid + topology::T + quadrature_style::Q + global_geometry::GG + local_geometry::LG + local_dss_weights::D + internal_surface_geometry::IS + boundary_surface_geometries::BS +end + + +""" + SpectralElementSpace2D(topology, quadrature_style; enable_bubble) + +Construct a `SpectralElementSpace2D` instance given a `topology` and `quadrature`. The +flag `enable_bubble` enables the `bubble correction` for more accurate element areas. + +# Input arguments: +- topology: Topology2D +- quadrature_style: QuadratureStyle +- enable_bubble: Bool + +The idea behind the so-called `bubble_correction` is that the numerical area +of the domain (e.g., the sphere) is given by the sum of nodal integration weights +times their corresponding Jacobians. However, this discrete sum is not exactly +equal to the exact geometric area (4pi*radius^2 for the sphere). To make these equal, +the "epsilon bubble" approach modifies the inner weights in each element so that +geometric and numerical areas of each element match. + +Let ``\\Delta A^e := A^e_{exact} - A^e_{approx}``, then, in +the case of linear elements, we correct ``W_{i,j} J^e_{i,j}`` by: +```math +\\widehat{W_{i,j} J^e}_{i,j} = W_{i,j} J^e_{i,j} + \\Delta A^e * W_{i,j} / Nq^2 . +``` +and the case of non linear elements, by +```math +\\widehat{W_{i,j} J^e}_{i,j} = W_{i,j} J^e_{i,j} \\left( 1 + \\tilde{A}^e \\right) , +``` +where ``\\tilde{A}^e`` is the approximated area given by the sum of the interior nodal integration weights. + +Note: This is accurate only for cubed-spheres of the [`Meshes.EquiangularCubedSphere`](@ref) and +[`Meshes.EquidistantCubedSphere`](@ref) type, not for [`Meshes.ConformalCubedSphere`](@ref). +""" +@memoize WeakValueDict function SpectralElementGrid2D( + topology, + quadrature_style; + enable_bubble = false, +) + + # 1. compute localgeom for local elememts + # 2. ghost exchange of localgeom + # 3. do a round of dss on WJs + # 4. compute dss weights (WJ ./ dss(WJ)) (local and ghost) + + # DSS on a field would consist of + # 1. copy to send buffers + # 2. start exchange + # 3. dss of internal connections + # - option for weighting and transformation + # 4. finish exchange + # 5. dss of ghost connections + + ### How to DSS multiple fields? + # 1. allocate buffers externally + DA = ClimaComms.array_type(topology) + domain = Topologies.domain(topology) + if domain isa Domains.SphereDomain + CoordType3D = Topologies.coordinate_type(topology) + FT = Geometry.float_type(CoordType3D) + CoordType2D = Geometry.LatLongPoint{FT} # Domains.coordinate_type(topology) + global_geometry = + Geometry.SphericalGlobalGeometry(topology.mesh.domain.radius) + else + CoordType2D = Topologies.coordinate_type(topology) + FT = Geometry.float_type(CoordType2D) + global_geometry = Geometry.CartesianGlobalGeometry() + end + AIdx = Geometry.coordinate_axis(CoordType2D) + nlelems = Topologies.nlocalelems(topology) + ngelems = Topologies.nghostelems(topology) + Nq = Quadratures.degrees_of_freedom(quadrature_style) + high_order_quadrature_style = Quadratures.GLL{Nq * 2}() + high_order_Nq = Quadratures.degrees_of_freedom(high_order_quadrature_style) + + LG = Geometry.LocalGeometry{AIdx, CoordType2D, FT, SMatrix{2, 2, FT, 4}} + + local_geometry = DataLayouts.IJFH{LG, Nq}(Array{FT}, nlelems) + + quad_points, quad_weights = + Quadratures.quadrature_points(FT, quadrature_style) + high_order_quad_points, high_order_quad_weights = + Quadratures.quadrature_points(FT, high_order_quadrature_style) + for (lidx, elem) in enumerate(Topologies.localelems(topology)) + elem_area = zero(FT) + high_order_elem_area = zero(FT) + Δarea = zero(FT) + interior_elem_area = zero(FT) + rel_interior_elem_area_Δ = zero(FT) + local_geometry_slab = slab(local_geometry, lidx) + # high-order quadrature loop for computing geometric element face area. + for i in 1:high_order_Nq, j in 1:high_order_Nq + ξ = SVector(high_order_quad_points[i], high_order_quad_points[j]) + u, ∂u∂ξ = + compute_local_geometry(global_geometry, topology, elem, ξ, AIdx) + J_high_order = det(Geometry.components(∂u∂ξ)) + WJ_high_order = + J_high_order * + high_order_quad_weights[i] * + high_order_quad_weights[j] + high_order_elem_area += WJ_high_order + end + # low-order quadrature loop for computing numerical element face area + for i in 1:Nq, j in 1:Nq + ξ = SVector(quad_points[i], quad_points[j]) + u, ∂u∂ξ = + compute_local_geometry(global_geometry, topology, elem, ξ, AIdx) + J = det(Geometry.components(∂u∂ξ)) + WJ = J * quad_weights[i] * quad_weights[j] + elem_area += WJ + if !enable_bubble + local_geometry_slab[i, j] = + Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) + end + end + + # If enabled, apply bubble correction + if enable_bubble + if abs(elem_area - high_order_elem_area) ≤ eps(FT) + for i in 1:Nq, j in 1:Nq + ξ = SVector(quad_points[i], quad_points[j]) + u, ∂u∂ξ = compute_local_geometry( + global_geometry, + topology, + elem, + ξ, + AIdx, + ) + J = det(Geometry.components(∂u∂ξ)) + WJ = J * quad_weights[i] * quad_weights[j] + local_geometry_slab[i, j] = + Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) + end + else + # The idea behind the so-called `bubble_correction` is that + # the numerical area of the domain (e.g., the sphere) is given by the sum + # of nodal integration weights times their corresponding Jacobians. However, + # this discrete sum is not exactly equal to the exact geometric area + # (4pi*radius^2 for the sphere). It is required that numerical area = geometric area. + # The "epsilon bubble" approach modifies the inner weights in each + # element so that geometric and numerical areas of each element match. + + # Compute difference between geometric area of an element and its approximate numerical area + Δarea = high_order_elem_area - elem_area + + # Linear elements: Nq == 2 (SpectralElementSpace2D cannot have Nq < 2) + # Use uniform bubble correction + if Nq == 2 + for i in 1:Nq, j in 1:Nq + ξ = SVector(quad_points[i], quad_points[j]) + u, ∂u∂ξ = compute_local_geometry( + global_geometry, + topology, + elem, + ξ, + AIdx, + ) + J = det(Geometry.components(∂u∂ξ)) + J += Δarea / Nq^2 + WJ = J * quad_weights[i] * quad_weights[j] + local_geometry_slab[i, j] = + Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) + end + else # Higher-order elements: Use HOMME bubble correction for the interior nodes + for i in 2:(Nq - 1), j in 2:(Nq - 1) + ξ = SVector(quad_points[i], quad_points[j]) + u, ∂u∂ξ = compute_local_geometry( + global_geometry, + topology, + elem, + ξ, + AIdx, + ) + J = det(Geometry.components(∂u∂ξ)) + WJ = J * quad_weights[i] * quad_weights[j] + interior_elem_area += WJ + end + # Check that interior_elem_area is not too small + if abs(interior_elem_area) ≤ sqrt(eps(FT)) + error( + "Bubble correction cannot be performed; sum of inner weights is too small.", + ) + end + rel_interior_elem_area_Δ = Δarea / interior_elem_area + + for i in 1:Nq, j in 1:Nq + ξ = SVector(quad_points[i], quad_points[j]) + u, ∂u∂ξ = compute_local_geometry( + global_geometry, + topology, + elem, + ξ, + AIdx, + ) + J = det(Geometry.components(∂u∂ξ)) + # Modify J only for interior nodes + if i != 1 && j != 1 && i != Nq && j != Nq + J *= (1 + rel_interior_elem_area_Δ) + end + WJ = J * quad_weights[i] * quad_weights[j] + # Finally allocate local geometry + local_geometry_slab[i, j] = + Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) + end + end + end + end + end + + # dss_weights = J ./ dss(J) + J = DataLayouts.rebuild(local_geometry.J, DA) + dss_local_weights = copy(J) + if quadrature_style isa Quadratures.GLL + Topologies.dss!(dss_local_weights, topology) + end + dss_local_weights .= J ./ dss_local_weights + + SG = Geometry.SurfaceGeometry{ + FT, + Geometry.AxisVector{FT, Geometry.LocalAxis{AIdx}, SVector{2, FT}}, + } + interior_faces = Array(Topologies.interior_faces(topology)) + + if quadrature_style isa Quadratures.GLL + internal_surface_geometry = + DataLayouts.IFH{SG, Nq}(Array{FT}, length(interior_faces)) + for (iface, (lidx⁻, face⁻, lidx⁺, face⁺, reversed)) in + enumerate(interior_faces) + internal_surface_geometry_slab = + slab(internal_surface_geometry, iface) + + local_geometry_slab⁻ = slab(local_geometry, lidx⁻) + local_geometry_slab⁺ = slab(local_geometry, lidx⁺) + + for q in 1:Nq + sgeom⁻ = compute_surface_geometry( + local_geometry_slab⁻, + quad_weights, + face⁻, + q, + false, + ) + sgeom⁺ = compute_surface_geometry( + local_geometry_slab⁺, + quad_weights, + face⁺, + q, + reversed, + ) + + @assert sgeom⁻.sWJ ≈ sgeom⁺.sWJ + @assert sgeom⁻.normal ≈ -sgeom⁺.normal + + internal_surface_geometry_slab[q] = sgeom⁻ + end + end + internal_surface_geometry = + DataLayouts.rebuild(internal_surface_geometry, DA) + + boundary_surface_geometries = + map(Topologies.boundary_tags(topology)) do boundarytag + boundary_faces = + Topologies.boundary_faces(topology, boundarytag) + boundary_surface_geometry = + DataLayouts.IFH{SG, Nq}(Array{FT}, length(boundary_faces)) + for (iface, (elem, face)) in enumerate(boundary_faces) + boundary_surface_geometry_slab = + slab(boundary_surface_geometry, iface) + local_geometry_slab = slab(local_geometry, elem) + for q in 1:Nq + boundary_surface_geometry_slab[q] = + compute_surface_geometry( + local_geometry_slab, + quad_weights, + face, + q, + false, + ) + end + end + DataLayouts.rebuild(boundary_surface_geometry, DA) + end + else + internal_surface_geometry = nothing + boundary_surface_geometries = nothing + end + return SpectralElementGrid2D( + topology, + quadrature_style, + global_geometry, + DataLayouts.rebuild(local_geometry, DA), + dss_local_weights, + internal_surface_geometry, + boundary_surface_geometries, + ) +end + +function compute_local_geometry( + global_geometry::Geometry.SphericalGlobalGeometry, + topology, + elem, + ξ, + AIdx, +) + x = Meshes.coordinates(topology.mesh, elem, ξ) + u = Geometry.LatLongPoint(x, global_geometry) + ∂x∂ξ = Geometry.AxisTensor( + (Geometry.Cartesian123Axis(), Geometry.CovariantAxis{AIdx}()), + ForwardDiff.jacobian(ξ) do ξ + Geometry.components(Meshes.coordinates(topology.mesh, elem, ξ)) + end, + ) + G = Geometry.local_to_cartesian(global_geometry, u) + ∂u∂ξ = Geometry.project(Geometry.LocalAxis{AIdx}(), G' * ∂x∂ξ) + + return u, ∂u∂ξ +end +function compute_local_geometry( + global_geometry::Geometry.AbstractGlobalGeometry, + topology, + elem, + ξ, + AIdx, +) + u = Meshes.coordinates(topology.mesh, elem, ξ) + ∂u∂ξ = Geometry.AxisTensor( + (Geometry.LocalAxis{AIdx}(), Geometry.CovariantAxis{AIdx}()), + ForwardDiff.jacobian(ξ) do ξ + Geometry.components(Meshes.coordinates(topology.mesh, elem, ξ)) + end, + ) + + return u, ∂u∂ξ +end + +function compute_surface_geometry( + local_geometry_slab, + quad_weights, + face, + q, + reversed = false, +) + Nq = length(quad_weights) + @assert size(local_geometry_slab) == (Nq, Nq, 1, 1, 1) + i, j = Topologies.face_node_index(face, Nq, q, reversed) + + local_geometry = local_geometry_slab[i, j] + (; J, ∂ξ∂x) = local_geometry + + # surface mass matrix + n = if face == 4 + -J * ∂ξ∂x[1, :] * quad_weights[j] + elseif face == 2 + J * ∂ξ∂x[1, :] * quad_weights[j] + elseif face == 1 + -J * ∂ξ∂x[2, :] * quad_weights[i] + elseif face == 3 + J * ∂ξ∂x[2, :] * quad_weights[i] + end + sWJ = norm(n) + n = n / sWJ + return Geometry.SurfaceGeometry(sWJ, n) +end + + +# accessors + +topology(grid::AbstractSpectralElementGrid) = grid.topology + +local_geometry_data(grid::AbstractSpectralElementGrid, ::Nothing) = + grid.local_geometry +global_geometry(grid::AbstractSpectralElementGrid) = grid.global_geometry + +quadrature_style(grid::AbstractSpectralElementGrid) = grid.quadrature_style +local_dss_weights(grid::SpectralElementGrid1D) = grid.dss_weights +local_dss_weights(grid::SpectralElementGrid2D) = grid.local_dss_weights + +## GPU compatibility +struct DeviceSpectralElementGrid2D{Q, GG, LG} <: AbstractSpectralElementGrid + quadrature_style::Q + global_geometry::GG + local_geometry::LG +end + +Adapt.adapt_structure(to, grid::SpectralElementGrid2D) = + DeviceSpectralElementGrid2D( + Adapt.adapt(to, grid.quadrature_style), + Adapt.adapt(to, grid.global_geometry), + Adapt.adapt(to, grid.local_geometry), + ) + +## aliases +const RectilinearSpectralElementGrid2D = + SpectralElementGrid2D{<:Topologies.RectilinearTopology2D} +const CubedSphereSpectralElementGrid2D = + SpectralElementGrid2D{<:Topologies.CubedSphereTopology2D} diff --git a/src/Hypsography/Hypsography.jl b/src/Hypsography/Hypsography.jl index 833e049edc..1261350675 100644 --- a/src/Hypsography/Hypsography.jl +++ b/src/Hypsography/Hypsography.jl @@ -1,10 +1,13 @@ module Hypsography import ..slab, ..column -import ..Geometry, ..Domains, ..Topologies, ..Spaces, ..Fields, ..Operators -import ..Spaces: ExtrudedFiniteDifferenceSpace, HypsographyAdaption, Flat +import ..Geometry, + ..Domains, ..Topologies, ..Grids, ..Spaces, ..Fields, ..Operators +import ..Spaces: ExtrudedFiniteDifferenceSpace -using StaticArrays, LinearAlgebra +import ..Grids: ExtrudedFiniteDifferenceGrid, HypsographyAdaption, Flat + +using StaticArrays, LinearAlgebra, Memoize, WeakValueDicts """ @@ -54,35 +57,35 @@ LinearAdaption() = LinearAdaption(nothing) ) # linear coordinates -function ExtrudedFiniteDifferenceSpace( - horizontal_space::Spaces.AbstractSpace, - vertical_space::Spaces.FiniteDifferenceSpace, +@memoize WeakValueDict function ExtrudedFiniteDifferenceGrid( + horizontal_grid::Grids.AbstractGrid, + vertical_grid::Grids.FiniteDifferenceGrid, adaption::HypsographyAdaption, ) if adaption isa LinearAdaption if isnothing(adaption.surface) error("LinearAdaption requires a Field argument") end - if axes(adaption.surface) !== horizontal_space + if axes(adaption.surface).grid !== horizontal_grid error("Terrain must be defined on the horizontal space") end end # construct initial flat space, then mutate - space = Spaces.ExtrudedFiniteDifferenceSpace( - horizontal_space, - vertical_space, + ref_grid = Grids.ExtrudedFiniteDifferenceGrid( + horizontal_grid, + vertical_grid, Flat(), ) - face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(space) - coord_type = eltype(Spaces.coordinates_data(face_space)) - z_ref = Spaces.coordinates_data(face_space).z + face_ref_space = Spaces.FaceExtrudedFiniteDifferenceSpace(ref_grid) + face_ref_coords = Spaces.coordinates_data(face_ref_space) + coord_type = eltype(face_ref_coords) + z_ref = face_ref_coords.z - vertical_domain = Topologies.domain(space.vertical_topology) + vertical_domain = Topologies.domain(vertical_grid) z_top = vertical_domain.coord_max.z grad = Operators.Gradient() - wdiv = Operators.WeakDivergence() z_surface = Fields.field_values(adaption.surface) FT = eltype(z_surface) @@ -90,7 +93,7 @@ function ExtrudedFiniteDifferenceSpace( # TODO: Function dispatch if adaption isa LinearAdaption fZ_data = @. z_ref + (1 - z_ref / z_top) * z_surface - fZ = Fields.Field(fZ_data, face_space) + fZ = Fields.Field(fZ_data, face_ref_space) elseif adaption isa SLEVEAdaption ηₕ = adaption.ηₕ s = adaption.s @@ -106,7 +109,7 @@ function ExtrudedFiniteDifferenceSpace( η * z_top + z_surface * (sinh((ηₕ - η) / s / ηₕ)) / (sinh(1 / s)), η * z_top, ) - fZ = Fields.Field(fZ_data, face_space) + fZ = Fields.Field(fZ_data, face_ref_space) end # Take the horizontal gradient for the Z surface field @@ -127,12 +130,15 @@ function ExtrudedFiniteDifferenceSpace( c∇Z = If2c.(f∇Z) Spaces.weighted_dss!(c∇Z) - Ni, Nj, _, Nv, Nh = size(space.center_local_geometry) + face_local_geometry = copy(ref_grid.face_local_geometry) + center_local_geometry = copy(ref_grid.center_local_geometry) + + Ni, Nj, _, Nv, Nh = size(center_local_geometry) for h in 1:Nh, j in 1:Nj, i in 1:Ni - face_column = column(space.face_local_geometry, i, j, h) + face_column = column(face_local_geometry, i, j, h) fZ_column = column(Fields.field_values(fZ), i, j, h) f∇Z_column = column(Fields.field_values(f∇Z), i, j, h) - center_column = column(space.center_local_geometry, i, j, h) + center_column = column(center_local_geometry, i, j, h) cZ_column = column(Fields.field_values(cZ), i, j, h) c∇Z_column = column(Fields.field_values(c∇Z), i, j, h) @@ -183,16 +189,13 @@ function ExtrudedFiniteDifferenceSpace( end end - return Spaces.ExtrudedFiniteDifferenceSpace( - space.staggering, - space.horizontal_space, - space.vertical_topology, + return ExtrudedFiniteDifferenceGrid( + horizontal_grid, + vertical_grid, adaption, - space.global_geometry, - space.center_local_geometry, - space.face_local_geometry, - space.center_ghost_geometry, - space.face_ghost_geometry, + ref_grid.global_geometry, + center_local_geometry, + face_local_geometry, ) end diff --git a/src/InputOutput/InputOutput.jl b/src/InputOutput/InputOutput.jl index 088dec93d7..8a54fae440 100644 --- a/src/InputOutput/InputOutput.jl +++ b/src/InputOutput/InputOutput.jl @@ -1,15 +1,18 @@ module InputOutput using HDF5, ClimaComms -import ..Geometry -import ..DataLayouts -import ..Domains -import ..Meshes -import ..Topologies -import ..Spaces -import ..Fields -import ..Hypsography +import ..Geometry, + ..DataLayouts, + ..Domains, + ..Meshes, + ..Topologies, + ..Quadratures, + ..Grids, + ..Spaces, + ..Fields, + ..Hypsography import ..VERSION +import ..Utilities: PlusHalf, half include("writers.jl") include("readers.jl") diff --git a/src/InputOutput/readers.jl b/src/InputOutput/readers.jl index 6801cfb29c..eaac23c4e5 100644 --- a/src/InputOutput/readers.jl +++ b/src/InputOutput/readers.jl @@ -14,10 +14,6 @@ using ..Meshes: using ..Topologies: Topologies, IntervalTopology using ..Spaces: Spaces, - Spaces.Quadratures, - Spaces.Quadratures.GLL, - Spaces.CellCenter, - Spaces.CellFace, SpectralElementSpace1D, SpectralElementSpace2D, CenterExtrudedFiniteDifferenceSpace, @@ -89,7 +85,7 @@ struct HDF5Reader{C <: ClimaComms.AbstractCommsContext} domain_cache::Dict{Any, Any} mesh_cache::Dict{Any, Any} topology_cache::Dict{Any, Any} - space_cache::Dict{Any, Any} + grid_cache::Dict{Any, Any} end @deprecate HDF5Reader(filename::AbstractString) HDF5Reader( @@ -130,7 +126,7 @@ function Base.close(hdfreader::HDF5Reader) empty!(hdfreader.domain_cache) empty!(hdfreader.mesh_cache) empty!(hdfreader.topology_cache) - empty!(hdfreader.space_cache) + empty!(hdfreader.grid_cache) close(hdfreader.file) return nothing end @@ -144,10 +140,10 @@ end function _scan_quadrature_style(quadraturestring::AbstractString, npts) @assert quadraturestring ∈ ("GLL", "GL", "Uniform", "ClosedUniform") - quadraturestring == "GLL" && return Spaces.Quadratures.GLL{npts}() - quadraturestring == "GL" && return Spaces.Quadratures.GL{npts}() - quadraturestring == "Uniform" && return Spaces.Quadratures.Uniform{npts}() - return Spaces.Quadratures.ClosedUniform{npts}() + quadraturestring == "GLL" && return Quadratures.GLL{npts}() + quadraturestring == "GL" && return Quadratures.GL{npts}() + quadraturestring == "Uniform" && return Quadratures.Uniform{npts}() + return Quadratures.ClosedUniform{npts}() end function _scan_data_layout(layoutstring::AbstractString) @@ -306,6 +302,56 @@ function read_topology_new(reader::HDF5Reader, name::AbstractString) end + +""" + read_grid(reader::AbstractReader, name) + +Reads a space named `name` from `reader`, or from the reader cache if it has +already been read. +""" +function read_grid(reader, name) + Base.get!(reader.grid_cache, name) do + read_grid_new(reader, name) + end +end + +function read_grid_new(reader, name) + group = reader.file["grids/$name"] + type = attrs(group)["type"] + if type in ("SpectralElementGrid1D", "SpectralElementGrid2D") + npts = attrs(group)["quadrature_num_points"] + quadrature_style = + _scan_quadrature_style(attrs(group)["quadrature_type"], npts) + topology = read_topology(reader, attrs(group)["topology"]) + if type == "SpectralElementGrid1D" + return Grids.SpectralElementGrid1D(topology, quadrature_style) + else + return Grids.SpectralElementGrid2D(topology, quadrature_style) + end + elseif type == "FiniteDifferenceGrid" + topology = read_topology(reader, attrs(group)["topology"]) + return Grids.FiniteDifferenceGrid(topology) + elseif type == "ExtrudedFiniteDifferenceGrid" + vertical_grid = read_grid(reader, attrs(group)["vertical_grid"]) + horizontal_grid = read_grid(reader, attrs(group)["horizontal_grid"]) + hypsography_type = get(attrs(group), "hypsography_type", "Flat") + if hypsography_type == "Flat" + hypsography = Grids.Flat() + elseif hypsography_type == "LinearAdaption" + hypsography = Hypsography.LinearAdaption( + read_field(reader, attrs(group)["hypsography_surface"]), + ) + else + error("Unsupported hypsography type $hypsography_type") + end + return Grids.ExtrudedFiniteDifferenceGrid( + horizontal_grid, + vertical_grid, + hypsography, + ) + end +end + """ read_space(reader::AbstractReader, name) @@ -342,7 +388,7 @@ function read_space_new(reader, name) read_space(reader, attrs(group)["horizontal_space"]) hypsography_type = get(attrs(group), "hypsography_type", "Flat") if hypsography_type == "Flat" - hypsography = Spaces.Flat() + hypsography = Grids.Flat() elseif hypsography_type == "LinearAdaption" hypsography = Hypsography.LinearAdaption( read_field(reader, attrs(group)["hypsography_surface"]), @@ -370,7 +416,18 @@ function read_field(reader::HDF5Reader, name::AbstractString) obj = reader.file["fields/$name"] type = attrs(obj)["type"] if type == "Field" - space = read_space(reader, attrs(obj)["space"]) + if haskey(attrs(obj), "grid") + grid = read_grid(reader, attrs(obj)["grid"]) + staggering = get(attrs(obj), "staggering", nothing) + if staggering == "CellCenter" + staggering = Grids.CellCenter() + elseif staggering == "CellFace" + staggering = Grids.CellFace() + end + space = Spaces.space(grid, staggering) + else + space = read_space(reader, attrs(obj)["space"]) + end topology = Spaces.topology(space) if topology isa Topologies.Topology2D nd = ndims(obj) diff --git a/src/InputOutput/writers.jl b/src/InputOutput/writers.jl index 6066550850..f972c73d57 100644 --- a/src/InputOutput/writers.jl +++ b/src/InputOutput/writers.jl @@ -295,16 +295,15 @@ function write_new!( return name end -# Spaces +# Grids # -defaultname(::Spaces.SpectralElementSpace1D) = "horizontal_space" -defaultname(::Spaces.SpectralElementSpace2D) = "horizontal_space" -defaultname(::Spaces.CenterExtrudedFiniteDifferenceSpace) = - "center_extruded_finite_difference_space" -defaultname(::Spaces.FaceExtrudedFiniteDifferenceSpace) = - "face_extruded_finite_difference_space" -defaultname(space::Spaces.FiniteDifferenceSpace) = defaultname(space.topology) - +defaultname(::Grids.SpectralElementGrid1D) = "horizontal_grid" +defaultname(::Grids.SpectralElementGrid2D) = "horizontal_grid" +defaultname(::Grids.ExtrudedFiniteDifferenceGrid) = + "extruded_finite_difference_grid" +defaultname(grid::Grids.FiniteDifferenceGrid) = defaultname(grid.topology) +defaultname(grid::Grids.LevelGrid) = + "$(defaultname(grid.full_grid)): level $(grid.level)" """ write_new!(writer, space, name) @@ -313,75 +312,70 @@ Write `SpectralElementSpace1D` data to HDF5. """ function write_new!( writer::HDF5Writer, - space::Spaces.SpectralElementSpace1D, + space::Grids.SpectralElementGrid1D, name::AbstractString = defaultname(space), ) - group = create_group(writer.file, "spaces/$name") - write_attribute(group, "type", "SpectralElementSpace1D") + group = create_group(writer.file, "grids/$name") + write_attribute(group, "type", "SpectralElementGrid1D") write_attribute( group, "quadrature_type", - string(nameof(typeof(space.quadrature_style))), + string(nameof(typeof(Spaces.quadrature_style(space)))), ) write_attribute( group, "quadrature_num_points", - Spaces.Quadratures.degrees_of_freedom(space.quadrature_style), + Quadratures.degrees_of_freedom(Spaces.quadrature_style(space)), ) - write_attribute(group, "topology", write!(writer, space.topology)) + write_attribute(group, "topology", write!(writer, Spaces.topology(space))) return name end function write_new!( writer::HDF5Writer, - space::Spaces.SpectralElementSpace2D, + space::Grids.SpectralElementGrid2D, name::AbstractString = defaultname(space), ) - group = create_group(writer.file, "spaces/$name") - write_attribute(group, "type", "SpectralElementSpace2D") + group = create_group(writer.file, "grids/$name") + write_attribute(group, "type", "SpectralElementGrid2D") write_attribute( group, "quadrature_type", - string(nameof(typeof(space.quadrature_style))), + string(nameof(typeof(Spaces.quadrature_style(space)))), ) write_attribute( group, "quadrature_num_points", - Spaces.Quadratures.degrees_of_freedom(space.quadrature_style), + Quadratures.degrees_of_freedom(Spaces.quadrature_style(space)), ) - write_attribute(group, "topology", write!(writer, space.topology)) + write_attribute(group, "topology", write!(writer, Spaces.topology(space))) return name end function write_new!( writer::HDF5Writer, - space::Spaces.FiniteDifferenceSpace, + space::Grids.FiniteDifferenceGrid, name::AbstractString = defaultname(space), ) - group = create_group(writer.file, "spaces/$name") - write_attribute(group, "type", "FiniteDifferenceSpace") - write_attribute(group, "topology", write!(writer, space.topology)) - return group + group = create_group(writer.file, "grids/$name") + write_attribute(group, "type", "FiniteDifferenceGrid") + write_attribute(group, "topology", write!(writer, Spaces.topology(space))) + return name end function write_new!( writer::HDF5Writer, - space::Spaces.CenterExtrudedFiniteDifferenceSpace, + space::Grids.ExtrudedFiniteDifferenceGrid, name::AbstractString = defaultname(space), ) - group = create_group(writer.file, "spaces/$name") - write_attribute(group, "type", "ExtrudedFiniteDifferenceSpace") - write_attribute(group, "staggering", "CellCenter") - write_attribute( - group, - "horizontal_space", - write!(writer, space.horizontal_space), - ) + group = create_group(writer.file, "grids/$name") + write_attribute(group, "type", "ExtrudedFiniteDifferenceGrid") write_attribute( group, - "vertical_topology", - write!(writer, space.vertical_topology), + "horizontal_grid", + write!(writer, space.horizontal_grid), ) + write_attribute(group, "vertical_grid", write!(writer, space.vertical_grid)) if space.hypsography isa Hypsography.LinearAdaption write_attribute(group, "hypsography_type", "LinearAdaption") write_attribute( @@ -393,28 +387,29 @@ function write_new!( return name end + function write_new!( writer::HDF5Writer, - space::Spaces.FaceExtrudedFiniteDifferenceSpace, + space::Grids.LevelGrid, name::AbstractString = defaultname(space), ) - group = create_group(writer.file, name) - group = create_group(writer.file, "spaces/$name") - write_attribute(group, "type", "ExtrudedFiniteDifferenceSpace") - write_attribute(group, "staggering", "CellFace") - write_attribute( - group, - "center_space", - write!(writer, Spaces.CenterExtrudedFiniteDifferenceSpace(space)), - ) + group = create_group(writer.file, "grids/$name") + write_attribute(group, "type", "LevelGrid") + write_attribute(group, "full_grid", write!(writer, space.full_grid)) + if space.level isa PlusHalf + write_attribute(group, "level + half", space.level - half) + else + write_attribute(group, "level", space.level) + end return name end - # write fields function write!(writer::HDF5Writer, field::Fields.Field, name::AbstractString) space = axes(field) - space_name = write!(writer, space) + staggering = Spaces.staggering(space) + grid = Spaces.grid(space) + grid_name = write!(writer, grid) array = parent(field) topology = Spaces.topology(space) @@ -449,7 +444,15 @@ function write!(writer::HDF5Writer, field::Fields.Field, name::AbstractString) string(nameof(typeof(Fields.field_values(field)))), ) write_attribute(dataset, "value_type", string(eltype(field))) - write_attribute(dataset, "space", space_name) + write_attribute(dataset, "grid", grid_name) + if !isnothing(staggering) + write_attribute( + dataset, + "staggering", + string(nameof(typeof(staggering))), + ) + end + return name end diff --git a/src/Limiters/quasimonotone.jl b/src/Limiters/quasimonotone.jl index 4dafc16c02..a259b8f97d 100644 --- a/src/Limiters/quasimonotone.jl +++ b/src/Limiters/quasimonotone.jl @@ -54,7 +54,7 @@ end function QuasiMonotoneLimiter(ρq::Fields.Field; rtol = eps(eltype(parent(ρq)))) q_bounds = make_q_bounds(Fields.field_values(ρq)) ghost_buffer = - Spaces.create_ghost_buffer(q_bounds, Spaces.topology(axes(ρq))) + Topologies.create_ghost_buffer(q_bounds, Spaces.topology(axes(ρq))) return QuasiMonotoneLimiter(q_bounds, similar(q_bounds), ghost_buffer, rtol) end @@ -191,7 +191,7 @@ function compute_neighbor_bounds_ghost!( ) q_bounds_nbr = limiter.q_bounds_nbr (_, _, _, Nv, Nh) = size(q_bounds_nbr) - if limiter.ghost_buffer isa Spaces.GhostBuffer + if limiter.ghost_buffer isa Topologies.GhostBuffer q_bounds_ghost = limiter.ghost_buffer.recv_data for h in 1:Nh @@ -232,7 +232,7 @@ function compute_bounds!( ρ::Fields.Field, ) compute_element_bounds!(limiter, ρq, ρ) - if limiter.ghost_buffer isa Spaces.GhostBuffer + if limiter.ghost_buffer isa Topologies.GhostBuffer Spaces.fill_send_buffer!( Spaces.topology(axes(ρq)), limiter.q_bounds, @@ -241,7 +241,7 @@ function compute_bounds!( ClimaComms.start(limiter.ghost_buffer.graph_context) end compute_neighbor_bounds_local!(limiter, ρ) - if limiter.ghost_buffer isa Spaces.GhostBuffer + if limiter.ghost_buffer isa Topologies.GhostBuffer ClimaComms.finish(limiter.ghost_buffer.graph_context) compute_neighbor_bounds_ghost!(limiter, Spaces.topology(axes(ρq))) end diff --git a/src/MatrixFields/MatrixFields.jl b/src/MatrixFields/MatrixFields.jl index 957c75de91..cbfa38dcb1 100644 --- a/src/MatrixFields/MatrixFields.jl +++ b/src/MatrixFields/MatrixFields.jl @@ -55,6 +55,7 @@ import ..RecursiveApply: import ..RecursiveApply: ⊠, ⊞, ⊟ import ..DataLayouts: AbstractData import ..Geometry +import ..Topologies import ..Spaces import ..Fields import ..Operators diff --git a/src/MatrixFields/single_field_solver.jl b/src/MatrixFields/single_field_solver.jl index bd02d67c2d..62ad93e8ee 100644 --- a/src/MatrixFields/single_field_solver.jl +++ b/src/MatrixFields/single_field_solver.jl @@ -50,7 +50,7 @@ single_field_solve!(::ClimaComms.AbstractCPUDevice, cache, x, A, b) = _single_field_solve!(cache, x, A, b) function single_field_solve!(::ClimaComms.CUDADevice, cache, x, A, b) Ni, Nj, _, _, Nh = size(Fields.field_values(A)) - nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh) + nthreads, nblocks = Topologies._configure_threadblock(Ni * Nj * Nh) CUDA.@cuda threads = nthreads blocks = nblocks single_field_solve_kernel!( cache, x, @@ -63,7 +63,7 @@ function single_field_solve_kernel!(cache, x, A, b) idx = CUDA.threadIdx().x + (CUDA.blockIdx().x - 1) * CUDA.blockDim().x Ni, Nj, _, _, Nh = size(Fields.field_values(A)) if idx <= Ni * Nj * Nh - i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx) + i, j, h = Topologies._get_idx((Ni, Nj, Nh), idx) _single_field_solve!( Spaces.column(cache, i, j, h), Spaces.column(x, i, j, h), diff --git a/src/Operators/Operators.jl b/src/Operators/Operators.jl index 973d59b842..b6673a3ca0 100644 --- a/src/Operators/Operators.jl +++ b/src/Operators/Operators.jl @@ -13,6 +13,7 @@ import ..Geometry: Geometry, Covariant12Vector, Contravariant12Vector, ⊗ import ..Spaces: Spaces, Quadratures, AbstractSpace import ..Topologies import ..Meshes +import ..Grids import ..Fields: Fields, Field using ..RecursiveApply diff --git a/src/Operators/finitedifference.jl b/src/Operators/finitedifference.jl index d2c8d44b9d..b13857f7d0 100644 --- a/src/Operators/finitedifference.jl +++ b/src/Operators/finitedifference.jl @@ -1,36 +1,36 @@ import ..Utilities: PlusHalf, half -left_idx( - space::Union{ - Spaces.CenterFiniteDifferenceSpace, - Spaces.CenterExtrudedFiniteDifferenceSpace, - }, -) = left_center_boundary_idx(space) -right_idx( - space::Union{ - Spaces.CenterFiniteDifferenceSpace, - Spaces.CenterExtrudedFiniteDifferenceSpace, - }, -) = right_center_boundary_idx(space) -left_idx( - space::Union{ - Spaces.FaceFiniteDifferenceSpace, - Spaces.FaceExtrudedFiniteDifferenceSpace, - }, -) = left_face_boundary_idx(space) -right_idx( - space::Union{ - Spaces.FaceFiniteDifferenceSpace, - Spaces.FaceExtrudedFiniteDifferenceSpace, - }, -) = right_face_boundary_idx(space) +const AllFiniteDifferenceSpace = + Union{Spaces.FiniteDifferenceSpace, Spaces.ExtrudedFiniteDifferenceSpace} +const AllFaceFiniteDifferenceSpace = Union{ + Spaces.FaceFiniteDifferenceSpace, + Spaces.FaceExtrudedFiniteDifferenceSpace, +} +const AllCenterFiniteDifferenceSpace = Union{ + Spaces.CenterFiniteDifferenceSpace, + Spaces.CenterExtrudedFiniteDifferenceSpace, +} + + -left_center_boundary_idx(space::Spaces.AbstractSpace) = 1 -right_center_boundary_idx(space::Spaces.AbstractSpace) = - size(space.center_local_geometry, 4) -left_face_boundary_idx(space::Spaces.AbstractSpace) = half -right_face_boundary_idx(space::Spaces.AbstractSpace) = - size(space.face_local_geometry, 4) - half +left_idx(space::AllCenterFiniteDifferenceSpace) = + left_center_boundary_idx(space) +right_idx(space::AllCenterFiniteDifferenceSpace) = + right_center_boundary_idx(space) +left_idx(space::AllFaceFiniteDifferenceSpace) = left_face_boundary_idx(space) +right_idx(space::AllFaceFiniteDifferenceSpace) = right_face_boundary_idx(space) + +left_center_boundary_idx(space::AllFiniteDifferenceSpace) = 1 +right_center_boundary_idx(space::AllFiniteDifferenceSpace) = size( + Spaces.local_geometry_data(Spaces.space(space, Spaces.CellCenter())), + 4, +) +left_face_boundary_idx(space::AllFiniteDifferenceSpace) = half +right_face_boundary_idx(space::AllFiniteDifferenceSpace) = + size( + Spaces.local_geometry_data(Spaces.space(space, Spaces.CellFace())), + 4, + ) - half left_face_boundary_idx(arg) = left_face_boundary_idx(axes(arg)) @@ -40,10 +40,7 @@ right_center_boundary_idx(arg) = right_center_boundary_idx(axes(arg)) # unlike getidx, we allow extracting the face local geometry from the center space, and vice-versa Base.@propagate_inbounds function Geometry.LocalGeometry( - space::Union{ - Spaces.FiniteDifferenceSpace, - Spaces.ExtrudedFiniteDifferenceSpace, - }, + space::AllFiniteDifferenceSpace, idx::Integer, hidx, ) @@ -52,13 +49,12 @@ Base.@propagate_inbounds function Geometry.LocalGeometry( v = mod1(v, length(space)) end i, j, h = hidx - return @inbounds space.center_local_geometry[CartesianIndex(i, j, 1, v, h)] + local_geom = + Grids.local_geometry_data(Spaces.grid(space), Grids.CellCenter()) + return @inbounds local_geom[CartesianIndex(i, j, 1, v, h)] end Base.@propagate_inbounds function Geometry.LocalGeometry( - space::Union{ - Spaces.FiniteDifferenceSpace, - Spaces.ExtrudedFiniteDifferenceSpace, - }, + space::AllFiniteDifferenceSpace, idx::PlusHalf, hidx, ) @@ -67,7 +63,8 @@ Base.@propagate_inbounds function Geometry.LocalGeometry( v = mod1(v, length(space)) end i, j, h = hidx - return @inbounds space.face_local_geometry[CartesianIndex(i, j, 1, v, h)] + local_geom = Grids.local_geometry_data(Spaces.grid(space), Grids.CellFace()) + return @inbounds local_geom[CartesianIndex(i, j, 1, v, h)] end @@ -354,12 +351,8 @@ struct InterpolateF2C{BCS} <: InterpolationOperator end InterpolateF2C(; kwargs...) = InterpolateF2C(NamedTuple(kwargs)) -return_space(::InterpolateF2C, space::Spaces.FaceFiniteDifferenceSpace) = - Spaces.CenterFiniteDifferenceSpace(space) -return_space( - ::InterpolateF2C, - space::Spaces.FaceExtrudedFiniteDifferenceSpace, -) = Spaces.CenterExtrudedFiniteDifferenceSpace(space) +return_space(::InterpolateF2C, space::AllFaceFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::InterpolateF2C, arg) = ((-half, half),) Base.@propagate_inbounds function stencil_interior( @@ -410,12 +403,8 @@ struct InterpolateC2F{BCS} <: InterpolationOperator end InterpolateC2F(; kwargs...) = InterpolateC2F(NamedTuple(kwargs)) -return_space(::InterpolateC2F, space::Spaces.CenterFiniteDifferenceSpace) = - Spaces.FaceFiniteDifferenceSpace(space) -return_space( - ::InterpolateC2F, - space::Spaces.CenterExtrudedFiniteDifferenceSpace, -) = Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::InterpolateC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::InterpolateC2F, arg) = ((-half, half),) Base.@propagate_inbounds function stencil_interior( @@ -539,12 +528,8 @@ struct LeftBiasedC2F{BCS} <: InterpolationOperator end LeftBiasedC2F(; kwargs...) = LeftBiasedC2F(NamedTuple(kwargs)) -return_space(::LeftBiasedC2F, space::Spaces.CenterFiniteDifferenceSpace) = - Spaces.FaceFiniteDifferenceSpace(space) -return_space( - ::LeftBiasedC2F, - space::Spaces.CenterExtrudedFiniteDifferenceSpace, -) = Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::LeftBiasedC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::LeftBiasedC2F, arg) = ((-half, -half),) Base.@propagate_inbounds stencil_interior( @@ -602,10 +587,8 @@ struct LeftBiasedF2C{BCS} <: InterpolationOperator end LeftBiasedF2C(; kwargs...) = LeftBiasedF2C(NamedTuple(kwargs)) -return_space(::LeftBiasedF2C, space::Spaces.FaceFiniteDifferenceSpace) = - Spaces.CenterFiniteDifferenceSpace(space) -return_space(::LeftBiasedF2C, space::Spaces.FaceExtrudedFiniteDifferenceSpace) = - Spaces.CenterExtrudedFiniteDifferenceSpace(space) +return_space(::LeftBiasedF2C, space::AllFaceFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::LeftBiasedF2C, arg) = ((-half, -half),) Base.@propagate_inbounds stencil_interior( @@ -664,14 +647,8 @@ struct LeftBiased3rdOrderC2F{BCS} <: InterpolationOperator end LeftBiased3rdOrderC2F(; kwargs...) = LeftBiased3rdOrderC2F(NamedTuple(kwargs)) -return_space( - ::LeftBiased3rdOrderC2F, - space::Spaces.CenterFiniteDifferenceSpace, -) = Spaces.FaceFiniteDifferenceSpace(space) -return_space( - ::LeftBiased3rdOrderC2F, - space::Spaces.CenterExtrudedFiniteDifferenceSpace, -) = Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::LeftBiased3rdOrderC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::LeftBiased3rdOrderC2F, arg) = ((-half - 1, half + 1),) Base.@propagate_inbounds stencil_interior( @@ -734,12 +711,8 @@ struct LeftBiased3rdOrderF2C{BCS} <: InterpolationOperator end LeftBiased3rdOrderF2C(; kwargs...) = LeftBiased3rdOrderF2C(NamedTuple(kwargs)) -return_space(::LeftBiased3rdOrderF2C, space::Spaces.FaceFiniteDifferenceSpace) = - Spaces.CenterFiniteDifferenceSpace(space) -return_space( - ::LeftBiased3rdOrderF2C, - space::Spaces.FaceExtrudedFiniteDifferenceSpace, -) = Spaces.CenterExtrudedFiniteDifferenceSpace(space) +return_space(::LeftBiased3rdOrderF2C, space::AllFaceFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::LeftBiased3rdOrderF2C, arg) = ((-half - 1, half + 1),) Base.@propagate_inbounds stencil_interior( @@ -802,12 +775,8 @@ struct RightBiasedC2F{BCS} <: InterpolationOperator end RightBiasedC2F(; kwargs...) = RightBiasedC2F(NamedTuple(kwargs)) -return_space(::RightBiasedC2F, space::Spaces.CenterFiniteDifferenceSpace) = - Spaces.FaceFiniteDifferenceSpace(space) -return_space( - ::RightBiasedC2F, - space::Spaces.CenterExtrudedFiniteDifferenceSpace, -) = Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::RightBiasedC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::RightBiasedC2F, arg) = ((half, half),) Base.@propagate_inbounds stencil_interior( @@ -865,12 +834,8 @@ struct RightBiasedF2C{BCS} <: InterpolationOperator end RightBiasedF2C(; kwargs...) = RightBiasedF2C(NamedTuple(kwargs)) -return_space(::RightBiasedF2C, space::Spaces.FaceFiniteDifferenceSpace) = - Spaces.CenterFiniteDifferenceSpace(space) -return_space( - ::RightBiasedF2C, - space::Spaces.FaceExtrudedFiniteDifferenceSpace, -) = Spaces.CenterExtrudedFiniteDifferenceSpace(space) +return_space(::RightBiasedF2C, space::AllFaceFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::RightBiasedF2C, arg) = ((half, half),) Base.@propagate_inbounds stencil_interior( @@ -931,14 +896,8 @@ struct RightBiased3rdOrderC2F{BCS} <: InterpolationOperator end RightBiased3rdOrderC2F(; kwargs...) = RightBiased3rdOrderC2F(NamedTuple(kwargs)) -return_space( - ::RightBiased3rdOrderC2F, - space::Spaces.CenterFiniteDifferenceSpace, -) = Spaces.FaceFiniteDifferenceSpace(space) -return_space( - ::RightBiased3rdOrderC2F, - space::Spaces.CenterExtrudedFiniteDifferenceSpace, -) = Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::RightBiased3rdOrderC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::RightBiased3rdOrderC2F, arg) = ((-half - 1, half + 1),) Base.@propagate_inbounds stencil_interior( @@ -989,14 +948,8 @@ struct RightBiased3rdOrderF2C{BCS} <: InterpolationOperator end RightBiased3rdOrderF2C(; kwargs...) = RightBiased3rdOrderF2C(NamedTuple(kwargs)) -return_space( - ::RightBiased3rdOrderF2C, - space::Spaces.FaceFiniteDifferenceSpace, -) = Spaces.CenterFiniteDifferenceSpace(space) -return_space( - ::RightBiased3rdOrderF2C, - space::Spaces.FaceExtrudedFiniteDifferenceSpace, -) = Spaces.CenterExtrudedFiniteDifferenceSpace(space) +return_space(::RightBiased3rdOrderF2C, space::AllFaceFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::RightBiased3rdOrderF2C, arg) = ((-half - 1, half + 1),) Base.@propagate_inbounds stencil_interior( @@ -1056,14 +1009,9 @@ WeightedInterpolateF2C(; kwargs...) = WeightedInterpolateF2C(NamedTuple(kwargs)) return_space( ::WeightedInterpolateF2C, - weight_space::Spaces.FaceFiniteDifferenceSpace, - arg_space::Spaces.FaceFiniteDifferenceSpace, -) = Spaces.CenterFiniteDifferenceSpace(arg_space) -return_space( - ::WeightedInterpolateF2C, - weight_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - arg_space::Spaces.FaceExtrudedFiniteDifferenceSpace, -) = Spaces.CenterExtrudedFiniteDifferenceSpace(arg_space) + weight_space::AllFaceFiniteDifferenceSpace, + arg_space::AllFaceFiniteDifferenceSpace, +) = Spaces.space(arg_space, Spaces.CellCenter()) stencil_interior_width(::WeightedInterpolateF2C, weight, arg) = ((-half, half), (-half, half)) @@ -1114,14 +1062,9 @@ WeightedInterpolateC2F(; kwargs...) = WeightedInterpolateC2F(NamedTuple(kwargs)) return_space( ::WeightedInterpolateC2F, - weight_space::Spaces.CenterFiniteDifferenceSpace, - arg_space::Spaces.CenterFiniteDifferenceSpace, -) = Spaces.FaceFiniteDifferenceSpace(arg_space) -return_space( - ::WeightedInterpolateC2F, - weight_space::Spaces.CenterExtrudedFiniteDifferenceSpace, - arg_space::Spaces.CenterExtrudedFiniteDifferenceSpace, -) = Spaces.FaceExtrudedFiniteDifferenceSpace(arg_space) + weight_space::AllCenterFiniteDifferenceSpace, + arg_space::AllCenterFiniteDifferenceSpace, +) = Spaces.space(arg_space, Spaces.CellFace()) stencil_interior_width(::WeightedInterpolateC2F, weight, arg) = ((-half, half), (-half, half)) @@ -1282,13 +1225,8 @@ return_eltype(::UpwindBiasedProductC2F, V, A) = return_space( ::UpwindBiasedProductC2F, - velocity_space::Spaces.FaceFiniteDifferenceSpace, - arg_space::Spaces.CenterFiniteDifferenceSpace, -) = velocity_space -return_space( - ::UpwindBiasedProductC2F, - velocity_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - arg_space::Spaces.CenterExtrudedFiniteDifferenceSpace, + velocity_space::AllFaceFiniteDifferenceSpace, + arg_space::AllCenterFiniteDifferenceSpace, ) = velocity_space function upwind_biased_product(v, a⁻, a⁺) @@ -1424,13 +1362,8 @@ return_eltype(::Upwind3rdOrderBiasedProductC2F, V, A) = return_space( ::Upwind3rdOrderBiasedProductC2F, - velocity_space::Spaces.FaceFiniteDifferenceSpace, - arg_space::Spaces.CenterFiniteDifferenceSpace, -) = velocity_space -return_space( - ::Upwind3rdOrderBiasedProductC2F, - velocity_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - arg_space::Spaces.CenterExtrudedFiniteDifferenceSpace, + velocity_space::AllFaceFiniteDifferenceSpace, + arg_space::AllCenterFiniteDifferenceSpace, ) = velocity_space function upwind_3rdorder_biased_product(v, a⁻, a⁻⁻, a⁺, a⁺⁺) @@ -1586,13 +1519,8 @@ return_eltype(::FCTBorisBook, V, A) = return_space( ::FCTBorisBook, - velocity_space::Spaces.FaceFiniteDifferenceSpace, - arg_space::Spaces.CenterFiniteDifferenceSpace, -) = velocity_space -return_space( - ::FCTBorisBook, - velocity_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - arg_space::Spaces.CenterExtrudedFiniteDifferenceSpace, + velocity_space::AllFaceFiniteDifferenceSpace, + arg_space::AllCenterFiniteDifferenceSpace, ) = velocity_space function fct_boris_book(v, a⁻⁻, a⁻, a⁺, a⁺⁺) @@ -1724,15 +1652,9 @@ return_eltype(::FCTZalesak, A, Φ, Φᵗᵈ) = return_space( ::FCTZalesak, - A_space::Spaces.FaceFiniteDifferenceSpace, - Φ_space::Spaces.CenterFiniteDifferenceSpace, - Φᵗᵈ_space::Spaces.CenterFiniteDifferenceSpace, -) = A_space -return_space( - ::FCTZalesak, - A_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - Φ_space::Spaces.CenterExtrudedFiniteDifferenceSpace, - Φᵗᵈ_space::Spaces.CenterExtrudedFiniteDifferenceSpace, + A_space::AllFaceFiniteDifferenceSpace, + Φ_space::AllCenterFiniteDifferenceSpace, + Φᵗᵈ_space::AllCenterFiniteDifferenceSpace, ) = A_space function fct_zalesak( @@ -1896,13 +1818,8 @@ AdvectionF2F(; kwargs...) = AdvectionF2F(NamedTuple(kwargs)) return_space( ::AdvectionF2F, - velocity_space::Spaces.FaceFiniteDifferenceSpace, - arg_space::Spaces.FaceFiniteDifferenceSpace, -) = arg_space -return_space( - ::AdvectionF2F, - velocity_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - arg_space::Spaces.FaceExtrudedFiniteDifferenceSpace, + velocity_space::AllFaceFiniteDifferenceSpace, + arg_space::AllFaceFiniteDifferenceSpace, ) = arg_space stencil_interior_width(::AdvectionF2F, velocity, arg) = ((0, 0), (-1, 1)) @@ -1958,13 +1875,8 @@ AdvectionC2C(; kwargs...) = AdvectionC2C(NamedTuple(kwargs)) return_space( ::AdvectionC2C, - velocity_space::Spaces.FaceFiniteDifferenceSpace, - arg_space::Spaces.CenterFiniteDifferenceSpace, -) = arg_space -return_space( - ::AdvectionC2C, - velocity_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - arg_space::Spaces.CenterExtrudedFiniteDifferenceSpace, + velocity_space::AllFaceFiniteDifferenceSpace, + arg_space::AllCenterFiniteDifferenceSpace, ) = arg_space stencil_interior_width(::AdvectionC2C, velocity, arg) = @@ -2096,13 +2008,8 @@ FluxCorrectionC2C(; kwargs...) = FluxCorrectionC2C(NamedTuple(kwargs)) return_space( ::FluxCorrectionC2C, - velocity_space::Spaces.FaceFiniteDifferenceSpace, - arg_space::Spaces.CenterFiniteDifferenceSpace, -) = arg_space -return_space( - ::FluxCorrectionC2C, - velocity_space::Spaces.FaceExtrudedFiniteDifferenceSpace, - arg_space::Spaces.CenterExtrudedFiniteDifferenceSpace, + velocity_space::AllFaceFiniteDifferenceSpace, + arg_space::AllCenterFiniteDifferenceSpace, ) = arg_space stencil_interior_width(::FluxCorrectionC2C, velocity, arg) = @@ -2181,13 +2088,8 @@ FluxCorrectionF2F(; kwargs...) = FluxCorrectionF2F(NamedTuple(kwargs)) return_space( ::FluxCorrectionF2F, - velocity_space::Spaces.CenterFiniteDifferenceSpace, - arg_space::Spaces.FaceFiniteDifferenceSpace, -) = arg_space -return_space( - ::FluxCorrectionF2F, - velocity_space::Spaces.CenterExtrudedFiniteDifferenceSpace, - arg_space::Spaces.FaceExtrudedFiniteDifferenceSpace, + velocity_space::AllCenterFiniteDifferenceSpace, + arg_space::AllFaceFiniteDifferenceSpace, ) = arg_space stencil_interior_width(::FluxCorrectionF2F, velocity, arg) = @@ -2273,12 +2175,7 @@ struct SetBoundaryOperator{BCS} <: BoundaryOperator end SetBoundaryOperator(; kwargs...) = SetBoundaryOperator(NamedTuple(kwargs)) -return_space(::SetBoundaryOperator, space::Spaces.FaceFiniteDifferenceSpace) = - space -return_space( - ::SetBoundaryOperator, - space::Spaces.FaceExtrudedFiniteDifferenceSpace, -) = space +return_space(::SetBoundaryOperator, space::AllFaceFiniteDifferenceSpace) = space stencil_interior_width(::SetBoundaryOperator, arg) = ((0, 0),) Base.@propagate_inbounds stencil_interior( @@ -2357,10 +2254,8 @@ struct GradientF2C{BCS} <: GradientOperator end GradientF2C(; kwargs...) = GradientF2C(NamedTuple(kwargs)) -return_space(::GradientF2C, space::Spaces.FaceFiniteDifferenceSpace) = - Spaces.CenterFiniteDifferenceSpace(space) -return_space(::GradientF2C, space::Spaces.FaceExtrudedFiniteDifferenceSpace) = - Spaces.CenterExtrudedFiniteDifferenceSpace(space) +return_space(::GradientF2C, space::AllFaceFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::GradientF2C, arg) = ((-half, half),) Base.@propagate_inbounds function stencil_interior( @@ -2472,10 +2367,8 @@ struct GradientC2F{BC} <: GradientOperator end GradientC2F(; kwargs...) = GradientC2F(NamedTuple(kwargs)) -return_space(::GradientC2F, space::Spaces.CenterFiniteDifferenceSpace) = - Spaces.FaceFiniteDifferenceSpace(space) -return_space(::GradientC2F, space::Spaces.CenterExtrudedFiniteDifferenceSpace) = - Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::GradientC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::GradientC2F, arg) = ((-half, half),) Base.@propagate_inbounds function stencil_interior( @@ -2598,10 +2491,8 @@ struct DivergenceF2C{BCS} <: DivergenceOperator end DivergenceF2C(; kwargs...) = DivergenceF2C(NamedTuple(kwargs)) -return_space(::DivergenceF2C, space::Spaces.FaceFiniteDifferenceSpace) = - Spaces.CenterFiniteDifferenceSpace(space) -return_space(::DivergenceF2C, space::Spaces.FaceExtrudedFiniteDifferenceSpace) = - Spaces.CenterExtrudedFiniteDifferenceSpace(space) +return_space(::DivergenceF2C, space::AllFaceFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::DivergenceF2C, arg) = ((-half, half),) Base.@propagate_inbounds function stencil_interior( @@ -2732,12 +2623,8 @@ struct DivergenceC2F{BC} <: DivergenceOperator end DivergenceC2F(; kwargs...) = DivergenceC2F(NamedTuple(kwargs)) -return_space(::DivergenceC2F, space::Spaces.CenterFiniteDifferenceSpace) = - Spaces.FaceFiniteDifferenceSpace(space) -return_space( - ::DivergenceC2F, - space::Spaces.CenterExtrudedFiniteDifferenceSpace, -) = Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::DivergenceC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::DivergenceC2F, arg) = ((-half, half),) Base.@propagate_inbounds function stencil_interior( @@ -2880,10 +2767,8 @@ struct CurlC2F{BC} <: CurlFiniteDifferenceOperator end CurlC2F(; kwargs...) = CurlC2F(NamedTuple(kwargs)) -return_space(::CurlC2F, space::Spaces.CenterFiniteDifferenceSpace) = - Spaces.FaceFiniteDifferenceSpace(space) -return_space(::CurlC2F, space::Spaces.CenterExtrudedFiniteDifferenceSpace) = - Spaces.FaceExtrudedFiniteDifferenceSpace(space) +return_space(::CurlC2F, space::AllCenterFiniteDifferenceSpace) = + Spaces.space(space, Spaces.CellFace()) fd3_curl(u₊::Geometry.Covariant1Vector, u₋::Geometry.Covariant1Vector, invJ) = Geometry.Contravariant2Vector((u₊.u₁ - u₋.u₁) * invJ) @@ -3190,13 +3075,7 @@ Base.Broadcast.BroadcastStyle( Base.eltype(bc::StencilBroadcasted) = return_eltype(bc.op, bc.args...) -function vidx( - space::Union{ - Spaces.FaceFiniteDifferenceSpace, - Spaces.FaceExtrudedFiniteDifferenceSpace, - }, - idx, -) +function vidx(space::AllFaceFiniteDifferenceSpace, idx) @assert idx isa PlusHalf v = idx + half if Topologies.isperiodic(Spaces.vertical_topology(space)) @@ -3204,13 +3083,7 @@ function vidx( end return v end -function vidx( - space::Union{ - Spaces.CenterFiniteDifferenceSpace, - Spaces.CenterExtrudedFiniteDifferenceSpace, - }, - idx, -) +function vidx(space::AllCenterFiniteDifferenceSpace, idx) @assert idx isa Integer v = idx if Topologies.isperiodic(Spaces.vertical_topology(space)) @@ -3459,7 +3332,7 @@ function Base.copyto!( Nv = ri - li + 1 max_threads = 256 nitems = Nv * Nq * Nq * Nh # # of independent items - (nthreads, nblocks) = Spaces._configure_threadblock(max_threads, nitems) + (nthreads, nblocks) = Topologies._configure_threadblock(max_threads, nitems) @cuda always_inline = true threads = (nthreads,) blocks = (nblocks,) copyto_stencil_kernel!( strip_space(out, space), strip_space(bc, space), @@ -3476,7 +3349,7 @@ function copyto_stencil_kernel!(out, bc, space, bds, Nq, Nh, Nv) gid = threadIdx().x + (blockIdx().x - 1) * blockDim().x if gid ≤ Nv * Nq * Nq * Nh (li, lw, rw, ri) = bds - (v, i, j, h) = Spaces._get_idx((Nv, Nq, Nq, Nh), gid) + (v, i, j, h) = Topologies._get_idx((Nv, Nq, Nq, Nh), gid) hidx = (i, j, h) idx = v - 1 + li window = diff --git a/src/Operators/integrals.jl b/src/Operators/integrals.jl index 2c1333733e..4406b9ef28 100644 --- a/src/Operators/integrals.jl +++ b/src/Operators/integrals.jl @@ -22,7 +22,7 @@ function column_integral_definite!( ) space = axes(∫field) Ni, Nj, _, _, Nh = size(Fields.field_values(∫field)) - nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh) + nthreads, nblocks = Topologies._configure_threadblock(Ni * Nj * Nh) @cuda threads = nthreads blocks = nblocks column_integral_definite_kernel!( strip_space(∫field, space), strip_space(ᶜfield, space), @@ -36,7 +36,7 @@ function column_integral_definite_kernel!( idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x Ni, Nj, _, _, Nh = size(Fields.field_values(ᶜfield)) if idx <= Ni * Nj * Nh - i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx) + i, j, h = Topologies._get_idx((Ni, Nj, Nh), idx) ∫field_column = Spaces.column(∫field, i, j, h) ᶜfield_column = Spaces.column(ᶜfield, i, j, h) _column_integral_definite!(∫field_column, ᶜfield_column) @@ -113,7 +113,7 @@ function column_integral_indefinite!( ᶜfield::Fields.Field, ) Ni, Nj, _, _, Nh = size(Fields.field_values(ᶠ∫field)) - nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh) + nthreads, nblocks = Topologies._configure_threadblock(Ni * Nj * Nh) @cuda threads = nthreads blocks = nblocks column_integral_indefinite_kernel!( ᶠ∫field, ᶜfield, @@ -127,7 +127,7 @@ function column_integral_indefinite_kernel!( idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x Ni, Nj, _, _, Nh = size(Fields.field_values(ᶜfield)) if idx <= Ni * Nj * Nh - i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx) + i, j, h = Topologies._get_idx((Ni, Nj, Nh), idx) ᶠ∫field_column = Spaces.column(ᶠ∫field, i, j, h) ᶜfield_column = Spaces.column(ᶜfield, i, j, h) _column_integral_indefinite!(ᶠ∫field_column, ᶜfield_column) @@ -289,7 +289,7 @@ function column_mapreduce_device!( fields::Fields.Field..., ) where {F, O} Ni, Nj, _, _, Nh = size(Fields.field_values(reduced_field)) - nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh) + nthreads, nblocks = Topologies._configure_threadblock(Ni * Nj * Nh) kernel! = if first(fields) isa Fields.ExtrudedFiniteDifferenceField column_mapreduce_kernel_extruded! else @@ -314,7 +314,7 @@ function column_mapreduce_kernel_extruded!( idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x Ni, Nj, _, _, Nh = size(Fields.field_values(reduced_field)) if idx <= Ni * Nj * Nh - i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx) + i, j, h = Topologies._get_idx((Ni, Nj, Nh), idx) reduced_field_column = Spaces.column(reduced_field, i, j, h) field_columns = map(field -> Spaces.column(field, i, j, h), fields) _column_mapreduce!(fn, op, reduced_field_column, field_columns...) @@ -387,7 +387,7 @@ Base.@propagate_inbounds data_level( ) = DataLayouts.level(data, v) Base.@propagate_inbounds data_level( data, - ::Spaces.FiniteDifferenceSpace{Spaces.CellCenter}, + ::Spaces.CenterFiniteDifferenceSpace, v::Int, ) = DataLayouts.level(data, v) @@ -398,7 +398,7 @@ Base.@propagate_inbounds data_level( ) = DataLayouts.level(data, v.i + 1) Base.@propagate_inbounds data_level( data, - ::Spaces.FiniteDifferenceSpace{Spaces.CellFace}, + ::Spaces.FaceFiniteDifferenceSpace, v::Utilities.PlusHalf, ) = DataLayouts.level(data, v.i + 1) @@ -407,9 +407,7 @@ right_boundary_idx(n, ::Operators.CenterPlaceholderSpace) = n left_boundary_idx(n, ::Operators.FacePlaceholderSpace) = Utilities.half right_boundary_idx(n, ::Operators.FacePlaceholderSpace) = n - Utilities.half -left_boundary_idx(n, ::Spaces.FiniteDifferenceSpace{Spaces.CellCenter}) = 1 -right_boundary_idx(n, ::Spaces.FiniteDifferenceSpace{Spaces.CellCenter}) = n -left_boundary_idx(n, ::Spaces.FiniteDifferenceSpace{Spaces.CellFace}) = - Utilities.half -right_boundary_idx(n, ::Spaces.FiniteDifferenceSpace{Spaces.CellFace}) = - n - Utilities.half +left_boundary_idx(n, ::Spaces.CenterFiniteDifferenceSpace) = 1 +right_boundary_idx(n, ::Spaces.CenterFiniteDifferenceSpace) = n +left_boundary_idx(n, ::Spaces.FaceFiniteDifferenceSpace) = Utilities.half +right_boundary_idx(n, ::Spaces.FaceFiniteDifferenceSpace) = n - Utilities.half diff --git a/src/Operators/numericalflux.jl b/src/Operators/numericalflux.jl index f0d9c0e3ca..d767c4cedc 100644 --- a/src/Operators/numericalflux.jl +++ b/src/Operators/numericalflux.jl @@ -24,8 +24,8 @@ See also: """ function add_numerical_flux_internal!(fn, dydt, args...) space = axes(dydt) - Nq = Spaces.Quadratures.degrees_of_freedom(space.quadrature_style) - topology = space.topology + Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(space)) + topology = Spaces.topology(space) for (iface, (elem⁻, face⁻, elem⁺, face⁺, reversed)) in enumerate(Topologies.interior_faces(topology)) @@ -100,8 +100,8 @@ end function add_numerical_flux_boundary!(fn, dydt, args...) space = axes(dydt) - Nq = Spaces.Quadratures.degrees_of_freedom(space.quadrature_style) - topology = space.topology + Nq = Spaces.Quadratures.degrees_of_freedom(Spaces.quadrature_style(space)) + topology = Spaces.topology(space) for (iboundary, boundarytag) in enumerate(Topologies.boundary_tags(topology)) diff --git a/src/Operators/remapping.jl b/src/Operators/remapping.jl index 0b4d37faca..42a864d6e1 100644 --- a/src/Operators/remapping.jl +++ b/src/Operators/remapping.jl @@ -67,28 +67,16 @@ end Computes local weights of the overlap mesh for `source` to `target` spaces. """ -function overlap( - target::T, - source::S, -) where { - T <: SpectralElementSpace1D{<:IntervalTopology}, - S <: SpectralElementSpace1D{<:IntervalTopology}, -} +function overlap(target::SpectralElementSpace1D, source::SpectralElementSpace1D) return x_overlap(target, source) end -function overlap( - target::T, - source::S, -) where { - T <: - SpectralElementSpace2D{<:Topology2D{<:ClimaComms.SingletonCommsContext}}, - S <: - SpectralElementSpace2D{<:Topology2D{<:ClimaComms.SingletonCommsContext}}, -} +function overlap(target::SpectralElementSpace2D, source::SpectralElementSpace2D) @assert ( - typeof(Spaces.topology(target).mesh) <: Meshes.RectilinearMesh && - typeof(Spaces.topology(source).mesh) <: Meshes.RectilinearMesh + ClimaComms.context(target) isa ClimaComms.SingletonCommsContext && + ClimaComms.context(source) isa ClimaComms.SingletonCommsContext && + Spaces.topology(target).mesh isa Meshes.RectilinearMesh && + Spaces.topology(source).mesh isa Meshes.RectilinearMesh ) X_ov = x_overlap(target, source) Y_ov = y_overlap(target, source) @@ -290,6 +278,6 @@ associated basis function. See [Ullrich2015] section 2. """ function local_weights(space::AbstractSpace) - wj = space.local_geometry.WJ + wj = Spaces.local_geometry_data(space).WJ return vec(parent(wj)) end diff --git a/src/Operators/spectralelement.jl b/src/Operators/spectralelement.jl index d0e953c06c..c9b926e849 100644 --- a/src/Operators/spectralelement.jl +++ b/src/Operators/spectralelement.jl @@ -56,7 +56,7 @@ operator_axes(space::Spaces.SpectralElementSpace2D) = (1, 2) operator_axes(space::Spaces.SpectralElementSpaceSlab1D) = (1,) operator_axes(space::Spaces.SpectralElementSpaceSlab2D) = (1, 2) operator_axes(space::Spaces.ExtrudedFiniteDifferenceSpace) = - operator_axes(space.horizontal_space) + operator_axes(Spaces.horizontal_space(space)) function node_indices(space::Spaces.SpectralElementSpace1D) @@ -70,7 +70,7 @@ function node_indices(space::Spaces.SpectralElementSpace2D) CartesianIndices((Nq, Nq)) end node_indices(space::Spaces.ExtrudedFiniteDifferenceSpace) = - node_indices(space.horizontal_space) + node_indices(Spaces.horizontal_space(space)) """ diff --git a/src/Operators/thomas_algorithm.jl b/src/Operators/thomas_algorithm.jl index 7154a78a6c..fa1a42e5c9 100644 --- a/src/Operators/thomas_algorithm.jl +++ b/src/Operators/thomas_algorithm.jl @@ -16,7 +16,7 @@ column_thomas_solve!(::ClimaComms.AbstractCPUDevice, A, b) = function column_thomas_solve!(::ClimaComms.CUDADevice, A, b) Ni, Nj, _, _, Nh = size(Fields.field_values(A)) - nthreads, nblocks = Spaces._configure_threadblock(Ni * Nj * Nh) + nthreads, nblocks = Topologies._configure_threadblock(Ni * Nj * Nh) @cuda threads = nthreads blocks = nblocks thomas_algorithm_kernel!(A, b) end @@ -27,7 +27,7 @@ function thomas_algorithm_kernel!( idx = threadIdx().x + (blockIdx().x - 1) * blockDim().x Ni, Nj, _, _, Nh = size(Fields.field_values(A)) if idx <= Ni * Nj * Nh - i, j, h = Spaces._get_idx((Ni, Nj, Nh), idx) + i, j, h = Topologies._get_idx((Ni, Nj, Nh), idx) thomas_algorithm!(Spaces.column(A, i, j, h), Spaces.column(b, i, j, h)) end return nothing @@ -38,10 +38,14 @@ thomas_algorithm_kernel!( b::Fields.FiniteDifferenceField, ) = thomas_algorithm!(A, b) -thomas_algorithm!( +function thomas_algorithm!( A::Fields.ExtrudedFiniteDifferenceField, b::Fields.ExtrudedFiniteDifferenceField, -) = Fields.bycolumn(colidx -> thomas_algorithm!(A[colidx], b[colidx]), axes(A)) +) + Fields.bycolumn(axes(A)) do colidx + thomas_algorithm!(A[colidx], b[colidx]) + end +end function thomas_algorithm!( A::Fields.FiniteDifferenceField, diff --git a/src/Quadratures/Quadratures.jl b/src/Quadratures/Quadratures.jl new file mode 100644 index 0000000000..5cacad9eb7 --- /dev/null +++ b/src/Quadratures/Quadratures.jl @@ -0,0 +1,285 @@ + +module Quadratures + +import GaussQuadrature +import StaticArrays: SVector, SMatrix, MMatrix +import LinearAlgebra: Diagonal + +export QuadratureStyle, + GLL, GL, polynomial_degree, degrees_of_freedom, quadrature_points + +""" + QuadratureStyle + +Quadrature style supertype. See sub-types: + - [`GLL`](@ref) + - [`GL`](@ref) + - [`Uniform`](@ref) +""" +abstract type QuadratureStyle end + +""" + polynomial_degree(QuadratureStyle) -> Int + +Returns the polynomial degree of the `QuadratureStyle` concrete type +""" +function polynomial_degree end + + +""" + degrees_of_freedom(QuadratureStyle) -> Int + +Returns the degrees_of_freedom of the `QuadratureStyle` concrete type +""" +function degrees_of_freedom end + +""" + points, weights = quadrature_points(::Type{FT}, quadrature_style) + +The points and weights of the quadrature rule in floating point type `FT`. +""" +function quadrature_points end + + +""" + GLL{Nq}() + +Gauss-Legendre-Lobatto quadrature using `Nq` quadrature points. +""" +struct GLL{Nq} <: QuadratureStyle end + +Base.show(io::IO, ::GLL{Nq}) where {Nq} = + print(io, Nq, "-point Gauss-Legendre-Lobatto quadrature") + +@inline polynomial_degree(::GLL{Nq}) where {Nq} = Int(Nq - 1) +@inline degrees_of_freedom(::GLL{Nq}) where {Nq} = Int(Nq) + +@generated function quadrature_points(::Type{FT}, ::GLL{Nq}) where {FT, Nq} + points, weights = GaussQuadrature.legendre(FT, Nq, GaussQuadrature.both) + :($(SVector{Nq}(points)), $(SVector{Nq}(weights))) +end + +""" + GL{Nq}() + +Gauss-Legendre quadrature using `Nq` quadrature points. +""" +struct GL{Nq} <: QuadratureStyle end +Base.show(io::IO, ::GL{Nq}) where {Nq} = + print(io, Nq, "-point Gauss-Legendre quadrature") + +@inline polynomial_degree(::GL{Nq}) where {Nq} = Int(Nq - 1) +@inline degrees_of_freedom(::GL{Nq}) where {Nq} = Int(Nq) + +@generated function quadrature_points(::Type{FT}, ::GL{Nq}) where {FT, Nq} + points, weights = GaussQuadrature.legendre(FT, Nq, GaussQuadrature.neither) + :($(SVector{Nq}(points)), $(SVector{Nq}(weights))) +end + +""" + Uniform{Nq}() + +Uniformly-spaced quadrature. +""" +struct Uniform{Nq} <: QuadratureStyle end + +@inline polynomial_degree(::Uniform{Nq}) where {Nq} = Int(Nq - 1) +@inline degrees_of_freedom(::Uniform{Nq}) where {Nq} = Int(Nq) + +@generated function quadrature_points(::Type{FT}, ::Uniform{Nq}) where {FT, Nq} + points = SVector{Nq}(range(-1 + 1 / Nq, step = 2 / Nq, length = Nq)) + weights = SVector{Nq}(ntuple(i -> 2 / Nq, Nq)) + :($points, $weights) +end + +""" + ClosedUniform{Nq}() + +Uniformly-spaced quadrature including boundary. +""" +struct ClosedUniform{Nq} <: QuadratureStyle end + +@inline polynomial_degree(::ClosedUniform{Nq}) where {Nq} = Int(Nq - 1) +@inline degrees_of_freedom(::ClosedUniform{Nq}) where {Nq} = Int(Nq) + +@generated function quadrature_points( + ::Type{FT}, + ::ClosedUniform{Nq}, +) where {FT, Nq} + points = SVector{Nq}(range(FT(-1), FT(1), length = Nq)) + weights = SVector{Nq}( + 1 / (Nq - 1), + ntuple(i -> 2 / (Nq - 1), Nq - 2)..., + 1 / (Nq - 1), + ) + :($points, $weights) +end + + +""" + barycentric_weights(x::SVector{Nq}) where {Nq} + +The barycentric weights associated with the array of point locations `x`: + +```math +w_j = \\frac{1}{\\prod_{k \\ne j} (x_i - x_j)} +``` + +See [Berrut2004](@cite), equation 3.2. +""" +function barycentric_weights(r::SVector{Nq, T}) where {Nq, T} + SVector{Nq}(ntuple(Nq) do i + w = one(T) + for j in 1:Nq + if j != i + w *= (r[j] - r[i]) + end + end + inv(w) + end) +end +@generated function barycentric_weights( + ::Type{FT}, + quadstyle::QuadratureStyle, +) where {FT} + barycentric_weights(quadrature_points(FT, quadstyle())[1]) +end + +""" + interpolation_matrix(x::SVector, r::SVector{Nq}) + +The matrix which interpolates the Lagrange polynomial of degree `Nq-1` through +the points `r`, to points `x`. The matrix coefficients are computed using the +Barycentric formula of [Berrut2004](@cite), section 4: +```math +I_{ij} = \\begin{cases} +1 & \\text{if } x_i = r_j, \\\\ +0 & \\text{if } x_i = r_k \\text{ for } k \\ne j, \\\\ +\\frac{\\displaystyle \\frac{w_j}{x_i - r_j}}{\\displaystyle \\sum_k \\frac{w_k}{x_i - r_k}} & \\text{otherwise,} +\\end{cases} +``` +where ``w_j`` are the barycentric weights, see [`barycentric_weights`](@ref). +""" +function interpolation_matrix( + points_to::SVector{Nto}, + points_from::SVector{Nfrom}, +) where {Nto, Nfrom} + T = eltype(points_to) + bw = barycentric_weights(points_from) + M = zeros(MMatrix{Nto, Nfrom, T, Nto * Nfrom}) + for i in 1:Nto + x_to = points_to[i] + skip_row = false + for j in 1:Nfrom + if x_to == points_from[j] + # assign to one to avoid singularity condition + M[i, j] = one(T) + # skip over the equal boundry condition + skip_row = true + end + skip_row && break + end + skip_row && continue + w = bw ./ (x_to .- points_from) + M[i, :] .= w ./ sum(w) + end + return SMatrix(M) +end + +@generated function interpolation_matrix( + ::Type{FT}, + quadto::QuadratureStyle, + quadfrom::QuadratureStyle, +) where {FT} + interpolation_matrix( + quadrature_points(FT, quadto())[1], + quadrature_points(FT, quadfrom())[1], + ) +end + +""" + V = orthonormal_poly(points, quad) + +`V_{ij}` contains the `j-1`th Legendre polynomial evaluated at `points[i]`. +i.e. it is the mapping from the modal to the nodal representation. +""" +function orthonormal_poly( + points::SVector{Np, FT}, + quad::GLL{Nq}, +) where {FT, Np, Nq} + N = Nq - 1 + a, b = GaussQuadrature.legendre_coefs(FT, N) + if N == 0 + return SMatrix{Np, 1}(ntuple(x -> b[1], Np)) + end + return SMatrix{Np, Nq}(GaussQuadrature.orthonormal_poly(points, a, b)) +end + +function spectral_filter_matrix( + quad::GLL{Nq}, + Σ::SVector{Nq, FT}, +) where {Nq, FT} + points, _ = quadrature_points(FT, quad) + V = orthonormal_poly(points, quad) + return V * Diagonal(Σ) / V +end + +function cutoff_filter_matrix( + ::Type{FT}, + quad::GLL{Nq}, + Nc::Integer, +) where {FT, Nq} + Σ = SVector(ntuple(i -> i <= Nc ? FT(1) : FT(0), Nq)) + return spectral_filter_matrix(quad, Σ) +end + +""" + differentiation_matrix(r::SVector{Nq, T}) where {Nq, T} + +The spectral differentiation matrix for the Lagrange polynomial of degree `Nq-1` +interpolating at points `r`. + +The matrix coefficients are computed using the [Berrut2004](@cite), section 9.3: +```math +D_{ij} = \\begin{cases} + \\displaystyle + \\frac{w_j}{w_i (x_i - x_j)} &\\text{ if } i \\ne j \\\\ + -\\sum_{k \\ne j} D_{kj} &\\text{ if } i = j +\\end{cases} +``` +where ``w_j`` are the barycentric weights, see [`barycentric_weights`](@ref). +""" +function differentiation_matrix(r::SVector{Nq, T}) where {Nq, T} + wb = barycentric_weights(r) + SMatrix{Nq, Nq, T, Nq * Nq}( + begin + if i == j + D = zero(T) + for l in 1:Nq + if l != i + D += one(T) / (r[i] - r[l]) + end + end + D + else + (wb[i] / wb[j]) / (r[j] - r[i]) + end + end for j in 1:Nq, i in 1:Nq + ) +end + +""" + differentiation_matrix(FT, quadstyle::QuadratureStyle) + +The spectral differentiation matrix at the quadrature points of `quadstyle`, +using floating point types `FT`. +""" +@generated function differentiation_matrix( + ::Type{FT}, + quadstyle::QuadratureStyle, +) where {FT} + differentiation_matrix(quadrature_points(FT, quadstyle())[1]) +end + + +end # module diff --git a/src/Remapping/Remapping.jl b/src/Remapping/Remapping.jl index 77ff1060e7..dab2c9f5c8 100644 --- a/src/Remapping/Remapping.jl +++ b/src/Remapping/Remapping.jl @@ -6,6 +6,7 @@ import ClimaComms import ..DataLayouts, ..Geometry, ..Spaces, + ..Grids, ..Topologies, ..Meshes, ..Operators, diff --git a/src/Remapping/interpolate_array.jl b/src/Remapping/interpolate_array.jl index 6a2467b30f..864138e870 100644 --- a/src/Remapping/interpolate_array.jl +++ b/src/Remapping/interpolate_array.jl @@ -238,7 +238,7 @@ function interpolate_column( # When we don't have hypsography, there is no notion of "interpolating hypsography". In # this case, the reference vertical points coincide with the physical ones. Setting # physical_z = false ensures that zpts_ref = zpts - if space.hypsography isa Spaces.Flat + if space.hypsography isa Grids.Flat physical_z = false end diff --git a/src/Spaces/Spaces.jl b/src/Spaces/Spaces.jl index 2ff34a2094..99ab5388e7 100644 --- a/src/Spaces/Spaces.jl +++ b/src/Spaces/Spaces.jl @@ -19,33 +19,78 @@ using Adapt using CUDA import ..slab, ..column, ..level -import ..Utilities: PlusHalf -import ..DataLayouts, ..Geometry, ..Domains, ..Meshes, ..Topologies +import ..Utilities: PlusHalf, half +import ..DataLayouts, + ..Geometry, ..Domains, ..Meshes, ..Topologies, ..Grids, ..Quadratures + +import ..Grids: + Staggering, + CellFace, + CellCenter, + topology, + vertical_topology, + local_geometry_data, + global_geometry, + local_dss_weights, + quadrature_style + import ClimaComms using StaticArrays, ForwardDiff, LinearAlgebra, UnPack, Adapt +""" + AbstractSpace + +Should define +- `grid` +- `staggering` + + +- `space` constructor + +""" abstract type AbstractSpace end +function grid end +function staggering end + + +ClimaComms.context(space::AbstractSpace) = ClimaComms.context(grid(space)) +ClimaComms.device(space::AbstractSpace) = ClimaComms.device(grid(space)) + +topology(space::AbstractSpace) = topology(grid(space)) +vertical_topology(space::AbstractSpace) = vertical_topology(grid(space)) + + +local_geometry_data(space::AbstractSpace) = + local_geometry_data(grid(space), staggering(space)) + +global_geometry(space::AbstractSpace) = global_geometry(grid(space)) + +space(refspace::AbstractSpace, staggering::Staggering) = + space(grid(refspace), staggering) + + + + + issubspace(::AbstractSpace, ::AbstractSpace) = false undertype(space::AbstractSpace) = Geometry.undertype(eltype(local_geometry_data(space))) coordinates_data(space::AbstractSpace) = local_geometry_data(space).coordinates - -include("quadrature.jl") -import .Quadratures +coordinates_data(grid::Grids.AbstractGrid) = + local_geometry_data(grid).coordinates +coordinates_data(staggering, grid::Grids.AbstractGrid) = + local_geometry_data(staggering, grid).coordinates include("pointspace.jl") include("spectralelement.jl") include("finitedifference.jl") include("extruded.jl") include("triangulation.jl") -include("dss_transform.jl") include("dss.jl") -horizontal_space(space::ExtrudedFiniteDifferenceSpace) = space.horizontal_space -horizontal_space(space::AbstractSpace) = space weighted_jacobian(space::Spaces.AbstractSpace) = local_geometry_data(space).WJ diff --git a/src/Spaces/dss.jl b/src/Spaces/dss.jl index ab448fae47..0e0abdab94 100644 --- a/src/Spaces/dss.jl +++ b/src/Spaces/dss.jl @@ -1,35 +1,21 @@ -using DocStringExtensions +import ..Topologies: + DSSBuffer, + create_dss_buffer, + assert_same_eltype, + dss_1d!, + dss_transform!, + dss_untransform!, + dss_local!, + dss_local_ghost!, + dss_ghost!, + fill_send_buffer!, + load_from_recv_buffer! + + +perimeter(space::AbstractSpectralElementSpace) = Topologies.Perimeter2D( + Quadratures.degrees_of_freedom(quadrature_style(space)), +) -""" - DSSBuffer{G, D, A, B} - -# Fields -$(DocStringExtensions.FIELDS) -""" -struct DSSBuffer{S, G, D, A, B, VI} - "ClimaComms graph context for communication" - graph_context::G - "Array for storing perimeter data" - perimeter_data::D - "send buffer" - send_data::A - "recv buffer" - recv_data::A - "indexing array for loading send buffer from `perimeter_data`" - send_buf_idx::B - "indexing array for loading (and summing) data from recv buffer to `perimeter_data`" - recv_buf_idx::B - "field id for all scalar fields stored in the `data` array" - scalarfidx::VI - "field id for all covariant12vector fields stored in the `data` array" - covariant12fidx::VI - "field id for all contravariant12vector fields stored in the `data` array" - contravariant12fidx::VI - "internal local elements (lidx)" - internal_elems::VI - "local elements (lidx) located on process boundary" - perimeter_elems::VI -end """ create_dss_buffer( @@ -41,174 +27,22 @@ Creates a [`DSSBuffer`](@ref) for the field data corresponding to `data` """ function create_dss_buffer( data::Union{DataLayouts.IJFH{S, Nij}, DataLayouts.VIJFH{S, Nij}}, - hspace::AbstractSpectralElementSpace, + hspace::SpectralElementSpace2D, ) where {S, Nij} - @assert hspace.quadrature_style isa Spaces.Quadratures.GLL "DSS2 is only compatible with GLL quadrature" - topology = hspace.topology - local_geometry = local_geometry_data(hspace) - local_weights = hspace.local_dss_weights - perimeter = Spaces.perimeter(hspace) - create_dss_buffer(data, topology, perimeter, local_geometry, local_weights) + create_dss_buffer( + data, + topology(hspace), + local_geometry_data(hspace), + local_dss_weights(hspace), + ) end function create_dss_buffer( - data::Union{DataLayouts.IJFH{S, Nij}, DataLayouts.VIJFH{S, Nij}}, - topology, - perimeter, - local_geometry = nothing, - local_weights = nothing, + data::Union{DataLayouts.IFH{S, Nij}, DataLayouts.VIFH{S, Nij}}, + hspace::SpectralElementSpace1D, ) where {S, Nij} - context = topology.context - DA = ClimaComms.array_type(topology) - convert_to_array = DA isa Array ? false : true - (_, _, _, Nv, nelems) = Base.size(data) - Np = Spaces.nperimeter(perimeter) - Nf = - length(parent(data)) == 0 ? 0 : - cld(length(parent(data)), (Nij * Nij * Nv * nelems)) - nfacedof = Nij - 2 - T = eltype(parent(data)) - TS = _transformed_type(data, local_geometry, local_weights, DA) # extract transformed type - # Add TS for Covariant123Vector - # For DSS of Covariant123Vector, the third component is treated like a scalar - # and is not transformed - if eltype(data) <: Geometry.Covariant123Vector - TS = Geometry.UVWVector{T} - end - perimeter_data = DataLayouts.VIFH{TS, Np}(DA{T}(undef, Nv, Np, Nf, nelems)) - if context isa ClimaComms.SingletonCommsContext - graph_context = ClimaComms.SingletonGraphContext(context) - send_data, recv_data = T[], T[] - send_buf_idx, recv_buf_idx = Int[], Int[] - send_data, recv_data = DA{T}(undef, 0), DA{T}(undef, 0) - send_buf_idx, recv_buf_idx = DA{Int}(undef, 0), DA{Int}(undef, 0) - internal_elems = DA{Int}(1:Topologies.nelems(topology)) - perimeter_elems = DA{Int}(undef, 0) - else - (; comm_vertex_lengths, comm_face_lengths) = topology - vertex_buffer_lengths = comm_vertex_lengths .* (Nv * Nf) - face_buffer_lengths = comm_face_lengths .* (Nv * Nf * nfacedof) - buffer_lengths = vertex_buffer_lengths .+ face_buffer_lengths - buffer_size = sum(buffer_lengths) - send_data = DA{T}(undef, buffer_size) - recv_data = DA{T}(undef, buffer_size) - neighbor_pids = topology.neighbor_pids - graph_context = ClimaComms.graph_context( - context, - send_data, - buffer_lengths, - neighbor_pids, - recv_data, - buffer_lengths, - neighbor_pids, - persistent = true, - ) - send_buf_idx, recv_buf_idx = - Topologies.compute_ghost_send_recv_idx(topology, Nij) - internal_elems = DA(topology.internal_elems) - perimeter_elems = DA(topology.perimeter_elems) - end - scalarfidx, covariant12fidx, contravariant12fidx = Int[], Int[], Int[] - supportedvectortypes = Union{ - Geometry.UVector, - Geometry.VVector, - Geometry.WVector, - Geometry.UVVector, - Geometry.UWVector, - Geometry.VWVector, - Geometry.UVWVector, - Geometry.Covariant12Vector, - Geometry.Covariant3Vector, - Geometry.Covariant123Vector, - Geometry.Contravariant12Vector, - Geometry.Contravariant3Vector, - } - - if S <: NamedTuple - for (i, fieldtype) in enumerate(S.parameters[2].types) - offset = DataLayouts.fieldtypeoffset(T, S, i) - ncomponents = DataLayouts.typesize(T, fieldtype) - if fieldtype <: Geometry.AxisVector # vector fields - if !(fieldtype <: supportedvectortypes) - @show fieldtype - @show supportedvectortypes - end - @assert fieldtype <: supportedvectortypes - if fieldtype <: Geometry.Covariant12Vector - push!(covariant12fidx, offset + 1) - elseif fieldtype <: Geometry.Covariant123Vector - push!(covariant12fidx, offset + 1) - push!(scalarfidx, offset + 3) - elseif fieldtype <: Geometry.Contravariant12Vector - push!(contravariant12fidx, offset + 1) - else - append!( - scalarfidx, - Vector((offset + 1):(offset + ncomponents)), - ) - end - elseif fieldtype <: NTuple # support a NTuple of primitive types - append!(scalarfidx, Vector((offset + 1):(offset + ncomponents))) - else # scalar fields - push!(scalarfidx, offset + 1) - end - end - else # deals with simple type, with single field (e.g: S = Float64, S = CovariantVector12, etc.) - ncomponents = DataLayouts.typesize(T, S) - if S <: Geometry.AxisVector # vector field - if !(S <: supportedvectortypes) - @show S - @show supportedvectortypes - end - @assert S <: supportedvectortypes - if S <: Geometry.Covariant12Vector - push!(covariant12fidx, 1) - elseif S <: Geometry.Covariant123Vector - push!(covariant12fidx, 1) - push!(scalarfidx, 3) - elseif S <: Geometry.Contravariant12Vector - push!(contravariant12fidx, 1) - else - append!(scalarfidx, Vector(1:ncomponents)) - end - elseif S <: NTuple # support a NTuple of primitive types - append!(scalarfidx, Vector(1:ncomponents)) - else # scalar field - push!(scalarfidx, 1) - end - end - scalarfidx = DA(scalarfidx) - covariant12fidx = DA(covariant12fidx) - contravariant12fidx = DA(contravariant12fidx) - G = typeof(graph_context) - D = typeof(perimeter_data) - A = typeof(send_data) - B = typeof(send_buf_idx) - VI = typeof(scalarfidx) - return DSSBuffer{S, G, D, A, B, VI}( - graph_context, - perimeter_data, - send_data, - recv_data, - send_buf_idx, - recv_buf_idx, - scalarfidx, - covariant12fidx, - contravariant12fidx, - internal_elems, - perimeter_elems, - ) -end - -Base.eltype(::DSSBuffer{S}) where {S} = S - -create_dss_buffer(data::DataLayouts.AbstractData, hspace) = nothing - -assert_same_eltype(::DataLayouts.AbstractData, ::DSSBuffer) = - error("Incorrect buffer eltype") -assert_same_eltype(::DataLayouts.AbstractData{S}, ::DSSBuffer{S}) where {S} = nothing -assert_same_eltype(::DataLayouts.AbstractData, ::Nothing) = nothing +end """ function weighted_dss!( @@ -251,6 +85,7 @@ function weighted_dss!( weighted_dss_ghost!(data, space, dss_buffer) end + """ weighted_dss_start!( data::Union{ @@ -292,24 +127,26 @@ weighted_dss_start!( dss_buffer::Union{DSSBuffer, Nothing}, ) = weighted_dss_start!(data, space, horizontal_space(space), dss_buffer) + + function weighted_dss_start!( data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, space::Union{ Spaces.SpectralElementSpace2D, Spaces.ExtrudedFiniteDifferenceSpace, }, - hspace::SpectralElementSpace2D{<:Topology2D}, + hspace::SpectralElementSpace2D, dss_buffer::DSSBuffer, ) assert_same_eltype(data, dss_buffer) length(parent(data)) == 0 && return nothing - device = ClimaComms.device(hspace.topology) + device = ClimaComms.device(topology(hspace)) dss_transform!( device, dss_buffer, data, local_geometry_data(space), - hspace.local_dss_weights, + local_dss_weights(hspace), Spaces.perimeter(hspace), dss_buffer.perimeter_elems, ) @@ -317,14 +154,20 @@ function weighted_dss_start!( device, dss_buffer.perimeter_data, Spaces.perimeter(hspace), - hspace.topology, + topology(hspace), ) fill_send_buffer!(device, dss_buffer) ClimaComms.start(dss_buffer.graph_context) return nothing end -weighted_dss_start!(data, space, hspace, dss_buffer) = nothing +weighted_dss_start!( + data, + space, + hspace::SpectralElementSpace1D, + dss_buffer::Nothing, +) = nothing + """ weighted_dss_internal!( data::Union{ @@ -358,6 +201,7 @@ weighted_dss_internal!( dss_buffer::Union{DSSBuffer, Nothing}, ) = weighted_dss_internal!(data, space, horizontal_space(space), dss_buffer) + function weighted_dss_internal!( data::Union{ DataLayouts.IFH, @@ -373,19 +217,19 @@ function weighted_dss_internal!( length(parent(data)) == 0 && return nothing if hspace isa SpectralElementSpace1D dss_1d!( - hspace.topology, + topology(hspace), data, local_geometry_data(space), - hspace.dss_weights, + local_dss_weights(space), ) else - device = ClimaComms.device(hspace.topology) + device = ClimaComms.device(topology(hspace)) dss_transform!( device, dss_buffer, data, local_geometry_data(space), - hspace.local_dss_weights, + local_dss_weights(space), Spaces.perimeter(hspace), dss_buffer.internal_elems, ) @@ -393,7 +237,7 @@ function weighted_dss_internal!( device, dss_buffer.perimeter_data, Spaces.perimeter(hspace), - hspace.topology, + topology(hspace), ) dss_untransform!( device, @@ -406,6 +250,8 @@ function weighted_dss_internal!( end return nothing end + + """ weighted_dss_ghost!( data::Union{ @@ -443,22 +289,24 @@ weighted_dss_ghost!( dss_buffer::Union{DSSBuffer, Nothing}, ) = weighted_dss_ghost!(data, space, horizontal_space(space), dss_buffer) + + function weighted_dss_ghost!( data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, space::Union{AbstractSpectralElementSpace, ExtrudedFiniteDifferenceSpace}, - hspace::SpectralElementSpace2D{<:Topology2D}, + hspace::SpectralElementSpace2D, dss_buffer::DSSBuffer, ) assert_same_eltype(data, dss_buffer) length(parent(data)) == 0 && return data - device = ClimaComms.device(hspace.topology) + device = ClimaComms.device(topology(hspace)) ClimaComms.finish(dss_buffer.graph_context) load_from_recv_buffer!(device, dss_buffer) dss_ghost!( device, dss_buffer.perimeter_data, Spaces.perimeter(hspace), - hspace.topology, + topology(hspace), ) dss_untransform!( device, @@ -472,675 +320,3 @@ function weighted_dss_ghost!( end weighted_dss_ghost!(data, space, hspace, dss_buffer) = data - -""" - function dss_transform!( - device::ClimaComms.AbstractDevice, - dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - weight::DataLayouts.IJFH, - perimeter::AbstractPerimeter, - localelems::Vector{Int}, - ) - -Transforms vectors from Covariant axes to physical (local axis), weights the data at perimeter nodes, -and stores result in the `perimeter_data` array. This function calls the appropriate version of -`dss_transform!` based on the data layout of the input arguments. - -Arguments: - -- `dss_buffer`: [`DSSBuffer`](@ref) generated by `create_dss_buffer` function for field data -- `data`: field data -- `local_geometry`: local metric information defined at each node -- `weight`: local dss weights for horizontal space -- `perimeter`: perimeter iterator -- `localelems`: list of local elements to perform transformation operations on - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function dss_transform!( - device::ClimaComms.AbstractDevice, - dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - weight::DataLayouts.IJFH, - perimeter::Perimeter2D, - localelems::AbstractVector{Int}, -) - if !isempty(localelems) - (; scalarfidx, covariant12fidx, contravariant12fidx, perimeter_data) = - dss_buffer - (; ∂ξ∂x, ∂x∂ξ) = local_geometry - dss_transform!( - device, - perimeter_data, - data, - ∂ξ∂x, - ∂x∂ξ, - weight, - perimeter, - scalarfidx, - covariant12fidx, - contravariant12fidx, - localelems, - ) - end - return nothing -end -""" - dss_untransform!( - device::ClimaComms.AbstractDevice, - dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - perimeter::AbstractPerimeter, - ) - -Transforms the DSS'd local vectors back to Covariant12 vectors, and copies the DSS'd data from the -`perimeter_data` to `data`. This function calls the appropriate version of `dss_transform!` function -based on the data layout of the input arguments. - -Arguments: - -- `dss_buffer`: [`DSSBuffer`](@ref) generated by `create_dss_buffer` function for field data -- `data`: field data -- `local_geometry`: local metric information defined at each node -- `perimeter`: perimeter iterator -- `localelems`: list of local elements to perform transformation operations on - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function dss_untransform!( - device::ClimaComms.AbstractDevice, - dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - perimeter::Perimeter2D, - localelems::AbstractVector{Int}, -) - (; scalarfidx, covariant12fidx, contravariant12fidx, perimeter_data) = - dss_buffer - (; ∂ξ∂x, ∂x∂ξ) = local_geometry - dss_untransform!( - device, - perimeter_data, - data, - ∂ξ∂x, - ∂x∂ξ, - perimeter, - scalarfidx, - covariant12fidx, - contravariant12fidx, - localelems, - ) - return nothing -end - -""" - function dss_transform!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - ∂ξ∂x::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - ∂x∂ξ::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - weight::DataLayouts.IJFH, - perimeter::AbstractPerimeter, - scalarfidx::Vector{Int}, - covariant12fidx::Vector{Int}, - contravariant12fidx::Vector{Int}, - localelems::Vector{Int}, - ) - -Transforms vectors from Covariant axes to physical (local axis), weights -the data at perimeter nodes, and stores result in the `perimeter_data` array. - -Arguments: - -- `perimeter_data`: contains the perimeter field data, represented on the physical axis, corresponding to the full field data in `data` -- `data`: field data -- `∂ξ∂x`: partial derivatives of the map from `x` to `ξ`: `∂ξ∂x[i,j]` is ∂ξⁱ/∂xʲ -- `weight`: local dss weights for horizontal space -- `perimeter`: perimeter iterator -- `scalarfidx`: field index for scalar fields in the data layout -- `covariant12fidx`: field index for Covariant12 vector fields in the data layout -- `localelems`: list of local elements to perform transformation operations on - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function dss_transform!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - ∂ξ∂x::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - ∂x∂ξ::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - weight::DataLayouts.IJFH, - perimeter::Perimeter2D{Nq}, - scalarfidx::Vector{Int}, - covariant12fidx::Vector{Int}, - contravariant12fidx::Vector{Int}, - localelems::Vector{Int}, -) where {Nq} - pdata = parent(data) - pweight = parent(weight) - p∂x∂ξ = parent(∂x∂ξ) - p∂ξ∂x = parent(∂ξ∂x) - pperimeter_data = parent(perimeter_data) - (nlevels, _, nfid, nelems) = size(pperimeter_data) - nmetric = cld(length(p∂ξ∂x), prod(size(∂ξ∂x))) - sizet_data = (nlevels, Nq, Nq, nfid, nelems) - sizet_wt = (Nq, Nq, 1, nelems) - sizet_metric = (nlevels, Nq, Nq, nmetric, nelems) - - @inbounds for elem in localelems - for (p, (ip, jp)) in enumerate(perimeter) - pw = pweight[_get_idx(sizet_wt, (ip, jp, 1, elem))] - - for fidx in scalarfidx, level in 1:nlevels - data_idx = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) - pperimeter_data[level, p, fidx, elem] = pdata[data_idx] * pw - end - - for fidx in covariant12fidx, level in 1:nlevels - data_idx1 = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) - data_idx2 = - _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) - (idx11, idx12, idx21, idx22) = - _get_idx_metric(sizet_metric, (level, ip, jp, elem)) - pperimeter_data[level, p, fidx, elem] = - ( - p∂ξ∂x[idx11] * pdata[data_idx1] + - p∂ξ∂x[idx12] * pdata[data_idx2] - ) * pw - pperimeter_data[level, p, fidx + 1, elem] = - ( - p∂ξ∂x[idx21] * pdata[data_idx1] + - p∂ξ∂x[idx22] * pdata[data_idx2] - ) * pw - end - - for fidx in contravariant12fidx, level in 1:nlevels - data_idx1 = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) - data_idx2 = - _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) - (idx11, idx12, idx21, idx22) = - _get_idx_metric(sizet_metric, (level, ip, jp, elem)) - pperimeter_data[level, p, fidx, elem] = - ( - p∂x∂ξ[idx11] * pdata[data_idx1] + - p∂x∂ξ[idx21] * pdata[data_idx2] - ) * pw - pperimeter_data[level, p, fidx + 1, elem] = - ( - p∂x∂ξ[idx12] * pdata[data_idx1] + - p∂x∂ξ[idx22] * pdata[data_idx2] - ) * pw - end - end - end - return nothing -end -""" - function dss_untransform!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - ∂ξ∂x::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - ∂x∂ξ::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - perimeter::AbstractPerimeter, - scalarfidx::Vector{Int}, - covariant12fidx::Vector{Int}, - contravariant12fidx::Vector{Int}, - localelems::Vector{Int}, - ) - -Transforms the DSS'd local vectors back to Covariant12 vectors, and copies the DSS'd data from the -`perimeter_data` to `data`. - -Arguments: - -- `perimeter_data`: contains the perimeter field data, represented on the physical axis, corresponding to the full field data in `data` -- `data`: field data -- `∂x∂ξ`: partial derivatives of the map from `ξ` to `x`: `∂x∂ξ[i,j]` is ∂xⁱ/∂ξʲ -- `perimeter`: perimeter iterator -- `scalarfidx`: field index for scalar fields in the data layout -- `covariant12fidx`: field index for Covariant12 vector fields in the data layout - -Part of [`Spaces.weighted_dss!`](@ref). -""" - -function dss_untransform!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - ∂ξ∂x::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - ∂x∂ξ::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, - perimeter::Perimeter2D{Nq}, - scalarfidx::Vector{Int}, - covariant12fidx::Vector{Int}, - contravariant12fidx::Vector{Int}, - localelems::Vector{Int}, -) where {Nq} - pdata = parent(data) - p∂x∂ξ = parent(∂x∂ξ) - p∂ξ∂x = parent(∂ξ∂x) - pperimeter_data = parent(perimeter_data) - (nlevels, _, nfid, nelems) = size(pperimeter_data) - nmetric = cld(length(p∂ξ∂x), prod(size(∂ξ∂x))) - sizet_data = (nlevels, Nq, Nq, nfid, nelems) - sizet_metric = (nlevels, Nq, Nq, nmetric, nelems) - - @inbounds for elem in localelems - for (p, (ip, jp)) in enumerate(perimeter) - for fidx in scalarfidx - for level in 1:nlevels - data_idx = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) - pdata[data_idx] = pperimeter_data[level, p, fidx, elem] - end - end - for fidx in covariant12fidx - for level in 1:nlevels - data_idx1 = - _get_idx(sizet_data, (level, ip, jp, fidx, elem)) - data_idx2 = - _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) - (idx11, idx12, idx21, idx22) = - _get_idx_metric(sizet_metric, (level, ip, jp, elem)) - pdata[data_idx1] = - p∂x∂ξ[idx11] * pperimeter_data[level, p, fidx, elem] + - p∂x∂ξ[idx12] * pperimeter_data[level, p, fidx + 1, elem] - pdata[data_idx2] = - p∂x∂ξ[idx21] * pperimeter_data[level, p, fidx, elem] + - p∂x∂ξ[idx22] * pperimeter_data[level, p, fidx + 1, elem] - end - end - for fidx in contravariant12fidx - for level in 1:nlevels - data_idx1 = - _get_idx(sizet_data, (level, ip, jp, fidx, elem)) - data_idx2 = - _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) - (idx11, idx12, idx21, idx22) = - _get_idx_metric(sizet_metric, (level, ip, jp, elem)) - pdata[data_idx1] = - p∂ξ∂x[idx11] * pperimeter_data[level, p, fidx, elem] + - p∂ξ∂x[idx21] * pperimeter_data[level, p, fidx + 1, elem] - pdata[data_idx2] = - p∂ξ∂x[idx12] * pperimeter_data[level, p, fidx, elem] + - p∂ξ∂x[idx22] * pperimeter_data[level, p, fidx + 1, elem] - end - end - end - end - return nothing -end - -function dss_load_perimeter_data!( - ::ClimaComms.AbstractCPUDevice, - dss_buffer::DSSBuffer, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - perimeter::Perimeter2D{Nq}, -) where {Nq} - pperimeter_data = parent(dss_buffer.perimeter_data) - pdata = parent(data) - (nlevels, _, nfid, nelems) = size(pperimeter_data) - sizet = (nlevels, Nq, Nq, nfid, nelems) - for elem in 1:nelems, (p, (ip, jp)) in enumerate(perimeter) - for fidx in 1:nfid, level in 1:nlevels - idx = _get_idx(sizet, (level, ip, jp, fidx, elem)) - pperimeter_data[level, p, fidx, elem] = pdata[idx] - end - end - return nothing -end - -function dss_unload_perimeter_data!( - ::ClimaComms.AbstractCPUDevice, - data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, - dss_buffer::DSSBuffer, - perimeter::Perimeter2D{Nq}, -) where {Nq} - pperimeter_data = parent(dss_buffer.perimeter_data) - pdata = parent(data) - (nlevels, _, nfid, nelems) = size(pperimeter_data) - sizet = (nlevels, Nq, Nq, nfid, nelems) - for elem in 1:nelems, (p, (ip, jp)) in enumerate(perimeter) - for fidx in 1:nfid, level in 1:nlevels - idx = _get_idx(sizet, (level, ip, jp, fidx, elem)) - pdata[idx] = pperimeter_data[level, p, fidx, elem] - end - end - return nothing -end - -""" - function dss_local!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - perimeter::AbstractPerimeter, - topology::Topologies.AbstractTopology, - ) - -Performs DSS on local vertices and faces. - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function dss_local!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - perimeter::Perimeter2D, - topology::Topologies.Topology2D, -) - dss_local_vertices!(perimeter_data, perimeter, topology) - dss_local_faces!(perimeter_data, perimeter, topology) - return nothing -end - -""" - dss_local_vertices!( - perimeter_data::DataLayouts.VIFH, - perimeter::Perimeter2D, - topology::Topologies.Topology2D, - ) - -Apply dss to local vertices. -""" -function dss_local_vertices!( - perimeter_data::DataLayouts.VIFH, - perimeter::Perimeter2D, - topology::Topologies.Topology2D, -) - Nv = size(perimeter_data, 4) - @inbounds for vertex in Topologies.local_vertices(topology) - # for each level - for level in 1:Nv - # gather: compute sum over shared vertices - sum_data = mapreduce( - ⊞, - vertex; - init = RecursiveApply.rzero(eltype(slab(perimeter_data, 1, 1))), - ) do (lidx, vert) - ip = Topologies.perimeter_vertex_node_index(vert) - perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] - end - # scatter: assign sum to shared vertices - for (lidx, vert) in vertex - perimeter_slab = slab(perimeter_data, level, lidx) - ip = Topologies.perimeter_vertex_node_index(vert) - perimeter_slab[ip] = sum_data - end - end - end - return nothing -end - -function dss_local_faces!( - perimeter_data::DataLayouts.VIFH, - perimeter::Perimeter2D, - topology::Topologies.Topology2D, -) - (Np, _, _, Nv, _) = size(perimeter_data) - nfacedof = div(Np - 4, 4) - - @inbounds for (lidx1, face1, lidx2, face2, reversed) in - Topologies.interior_faces(topology) - pr1 = Topologies.perimeter_face_indices(face1, nfacedof, false) - pr2 = Topologies.perimeter_face_indices(face2, nfacedof, reversed) - for level in 1:Nv - perimeter_slab1 = slab(perimeter_data, level, lidx1) - perimeter_slab2 = slab(perimeter_data, level, lidx2) - for (ip1, ip2) in zip(pr1, pr2) - val = perimeter_slab1[ip1] ⊞ perimeter_slab2[ip2] - perimeter_slab1[ip1] = val - perimeter_slab2[ip2] = val - end - end - end - return nothing -end -""" - function dss_local_ghost!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - perimeter::AbstractPerimeter, - topology::Topologies.AbstractTopology, - ) - -Computes the "local" part of ghost vertex dss. (i.e. it computes the summation of all the shared local -vertices of a unique ghost vertex and stores the value in each of the local vertex locations in -`perimeter_data`) - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function dss_local_ghost!( - ::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - perimeter::AbstractPerimeter, - topology::Topologies.AbstractTopology, -) - nghostvertices = length(topology.ghost_vertex_offset) - 1 - if nghostvertices > 0 - (Np, _, _, Nv, _) = size(perimeter_data) - @inbounds for vertex in Topologies.ghost_vertices(topology) - for level in 1:Nv - # gather: compute sum over shared vertices - sum_data = mapreduce( - ⊞, - vertex; - init = RecursiveApply.rzero( - eltype(slab(perimeter_data, 1, 1)), - ), - ) do (isghost, idx, vert) - ip = Topologies.perimeter_vertex_node_index(vert) - if !isghost - lidx = idx - perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] - else - RecursiveApply.rmap(zero, slab(perimeter_data, 1, 1)[1]) - end - end - for (isghost, idx, vert) in vertex - if !isghost - ip = Topologies.perimeter_vertex_node_index(vert) - lidx = idx - perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] = sum_data - end - end - end - end - end - return nothing -end -""" - dss_ghost!( - device::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - perimeter::AbstractPerimeter, - topology::Topologies.AbstractTopology, - ) - -Sets the value for all local vertices of each unique ghost vertex, in `perimeter_data`, to that of -the representative ghost vertex. - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function dss_ghost!( - device::ClimaComms.AbstractCPUDevice, - perimeter_data::DataLayouts.VIFH, - perimeter::AbstractPerimeter, - topology::Topologies.AbstractTopology, -) - nghostvertices = length(topology.ghost_vertex_offset) - 1 - if nghostvertices > 0 - nlevels = size(perimeter_data, 4) - perimeter_vertex_node_index = Topologies.perimeter_vertex_node_index - perimeter_face_indices = Topologies.perimeter_face_indices - (; repr_ghost_vertex) = topology - @inbounds for (i, vertex) in - enumerate(Topologies.ghost_vertices(topology)) - idxresult, lvertresult = repr_ghost_vertex[i] - ipresult = perimeter_vertex_node_index(lvertresult) - for level in 1:nlevels - result_slab = slab(perimeter_data, level, idxresult) - result = result_slab[ipresult] - for (isghost, idx, vert) in vertex - if !isghost - ip = perimeter_vertex_node_index(vert) - lidx = idx - perimeter_slab = slab(perimeter_data, level, lidx) - perimeter_slab[ip] = result - end - end - end - end - end - return nothing -end - -""" - fill_send_buffer!(::ClimaComms.AbstractCPUDevice, dss_buffer::DSSBuffer) - -Loads the send buffer from `perimeter_data`. For unique ghost vertices, only data from the -representative vertices which store result of "ghost local" DSS are loaded. - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function fill_send_buffer!( - ::ClimaComms.AbstractCPUDevice, - dss_buffer::DSSBuffer, -) - (; perimeter_data, send_buf_idx, send_data) = dss_buffer - (Np, _, _, Nv, nelems) = size(perimeter_data) - Nf = cld(length(parent(perimeter_data)), (Nv * Np * nelems)) - pdata = parent(perimeter_data) - nsend = size(send_buf_idx, 1) - ctr = 1 - @inbounds for i in 1:nsend - lidx = send_buf_idx[i, 1] - ip = send_buf_idx[i, 2] - for f in 1:Nf, v in 1:Nv - send_data[ctr] = pdata[v, ip, f, lidx] - ctr += 1 - end - end - return nothing -end -""" - load_from_recv_buffer!(::ClimaComms.AbstractCPUDevice, dss_buffer::DSSBuffer) - -Adds data from the recv buffer to the corresponding location in `perimeter_data`. -For ghost vertices, this data is added only to the representative vertices. The values are -then scattered to other local vertices corresponding to each unique ghost vertex in `dss_local_ghost`. - -Part of [`Spaces.weighted_dss!`](@ref). -""" -function load_from_recv_buffer!( - ::ClimaComms.AbstractCPUDevice, - dss_buffer::DSSBuffer, -) - (; perimeter_data, recv_buf_idx, recv_data) = dss_buffer - (Np, _, _, Nv, nelems) = size(perimeter_data) - Nf = cld(length(parent(perimeter_data)), (Nv * Np * nelems)) - pdata = parent(perimeter_data) - nrecv = size(recv_buf_idx, 1) - ctr = 1 - @inbounds for i in 1:nrecv - lidx = recv_buf_idx[i, 1] - ip = recv_buf_idx[i, 2] - for f in 1:Nf, v in 1:Nv - pdata[v, ip, f, lidx] += recv_data[ctr] - ctr += 1 - end - end - return nothing -end - -""" - dss!(data, topology, quadrature_style) - -Computed unweighted/pure DSS of `data`. -""" -function dss!(data, topology, quadrature_style) - length(parent(data)) == 0 && return nothing - device = ClimaComms.device(topology) - perimeter = Perimeter2D(Quadratures.degrees_of_freedom(quadrature_style)) - # create dss buffer - dss_buffer = create_dss_buffer(data, topology, perimeter) - # load perimeter data from data - dss_load_perimeter_data!(device, dss_buffer, data, perimeter) - # compute local dss for ghost dof - dss_local_ghost!(device, dss_buffer.perimeter_data, perimeter, topology) - # load send buffer - fill_send_buffer!(device, dss_buffer) - # initiate communication - ClimaComms.start(dss_buffer.graph_context) - # compute local dss - dss_local!(device, dss_buffer.perimeter_data, perimeter, topology) - # finish communication - ClimaComms.finish(dss_buffer.graph_context) - # load from receive buffer - load_from_recv_buffer!(device, dss_buffer) - # finish dss computation for ghost dof - dss_ghost!(device, dss_buffer.perimeter_data, perimeter, topology) - # load perimeter_data into data - dss_unload_perimeter_data!(device, data, dss_buffer, perimeter) - return nothing -end - -dss2!(data, topology, quadrature_style) = dss!(data, topology, quadrature_style) - -function dss_1d!( - htopology::Topologies.AbstractTopology, - data, - local_geometry_data = nothing, - dss_weights = nothing, -) - Nq = size(data, 1) - Nv = size(data, 4) - idx1 = CartesianIndex(1, 1, 1, 1, 1) - idx2 = CartesianIndex(Nq, 1, 1, 1, 1) - @inbounds for (elem1, face1, elem2, face2, reversed) in - Topologies.interior_faces(htopology) - for level in 1:Nv - @assert face1 == 1 && face2 == 2 && !reversed - local_geometry_slab1 = slab(local_geometry_data, level, elem1) - weight_slab1 = slab(dss_weights, level, elem1) - data_slab1 = slab(data, level, elem1) - - local_geometry_slab2 = slab(local_geometry_data, level, elem2) - weight_slab2 = slab(dss_weights, level, elem2) - data_slab2 = slab(data, level, elem2) - val = - dss_transform( - data_slab1, - local_geometry_slab1, - weight_slab1, - idx1, - ) ⊞ dss_transform( - data_slab2, - local_geometry_slab2, - weight_slab2, - idx2, - ) - - data_slab1[idx1] = dss_untransform( - eltype(data_slab1), - val, - local_geometry_slab1, - idx1, - ) - data_slab2[idx2] = dss_untransform( - eltype(data_slab2), - val, - local_geometry_slab2, - idx2, - ) - end - end - return data -end - -include("dss_cuda.jl") diff --git a/src/Spaces/extruded.jl b/src/Spaces/extruded.jl index 2bd93253b3..315e047215 100644 --- a/src/Spaces/extruded.jl +++ b/src/Spaces/extruded.jl @@ -1,96 +1,174 @@ -##### -##### Hybrid mesh -##### - -abstract type HypsographyAdaption end - -""" - Flat() - -No surface hypsography. -""" -struct Flat <: HypsographyAdaption end struct ExtrudedFiniteDifferenceSpace{ + G <: Grids.AbstractExtrudedFiniteDifferenceGrid, S <: Staggering, - H <: AbstractSpace, - T <: Topologies.AbstractIntervalTopology, - A <: HypsographyAdaption, - GG <: Geometry.AbstractGlobalGeometry, - LG, - LGG, } <: AbstractSpace + grid::G staggering::S - horizontal_space::H - vertical_topology::T - hypsography::A - global_geometry::GG - center_local_geometry::LG - face_local_geometry::LG - center_ghost_geometry::LGG - face_ghost_geometry::LGG end +space(grid::Grids.ExtrudedFiniteDifferenceGrid, staggering::Staggering) = + ExtrudedFiniteDifferenceSpace(grid, staggering) + +const FaceExtrudedFiniteDifferenceSpace{G} = + ExtrudedFiniteDifferenceSpace{G, CellFace} +const CenterExtrudedFiniteDifferenceSpace{G} = + ExtrudedFiniteDifferenceSpace{G, CellCenter} + +#= +ExtrudedFiniteDifferenceSpace{S}( + grid::Grids.ExtrudedFiniteDifferenceGrid, +) where {S <: Staggering} = ExtrudedFiniteDifferenceSpace(S(), grid) +ExtrudedFiniteDifferenceSpace{S}( + space::ExtrudedFiniteDifferenceSpace, +) where {S <: Staggering} = ExtrudedFiniteDifferenceSpace{S}(space.grid) +=# + +function ExtrudedFiniteDifferenceSpace( + horizontal_space::AbstractSpace, + vertical_space::FiniteDifferenceSpace, + hypsography::Grids.HypsographyAdaption = Grids.Flat(), +) + grid_space = Grids.ExtrudedFiniteDifferenceGrid( + grid(horizontal_space), + grid(vertical_space), + hypsography, + ) + return ExtrudedFiniteDifferenceSpace(grid_space, vertical_space.staggering) +end + +@inline function Base.getproperty( + space::ExtrudedFiniteDifferenceSpace, + name::Symbol, +) + if name == :horizontal_space + Base.depwarn( + "`space.horizontal_space` is deprecated, use `Spaces.horizontal_space(space)` instead", + :getproperty, + ) + return horizontal_space(space) + elseif name == :vertical_topology + Base.depwarn( + "`space.vertical_topology` is deprecated, use `Spaces.vertical_topology(space)` instead", + :getproperty, + ) + return vertical_topology(space) + elseif name == :hypsography + Base.depwarn( + "`space.hypsography` is deprecated, use `Spaces.grid(space).hypsography` instead", + :getproperty, + ) + return grid(space).hypsography + elseif name == :global_geometry + Base.depwarn( + "`space.global_geometry` is deprecated, use `Spaces.global_geometry(space)` instead", + :getproperty, + ) + return global_geometry(space) + elseif name == :center_local_geometry + Base.depwarn( + "`space.center_local_geometry` is deprecated, use `Spaces.local_geometry_data(grid(space), Grids.CellCenter())` instead", + :getproperty, + ) + return local_geometry_data(grid(space), Grids.CellCenter()) + elseif name == :face_local_geometry + Base.depwarn( + "`space.face_local_geometry` is deprecated, use `Spaces.local_geometry_data(grid(space), Grids.CellFace())` instead", + :getproperty, + ) + return local_geometry_data(grid(space), Grids.CellFace()) + elseif name == :center_ghost_geometry + Base.depwarn( + "`space.center_ghost_geometry` is deprecated, use `nothing` instead", + :getproperty, + ) + return nothing + elseif name == :face_ghost_geometry + Base.depwarn( + "`space.face_ghost_geometry` is deprecated, use `nothing` instead", + :getproperty, + ) + return nothing + end + return getfield(space, name) +end + +FaceExtrudedFiniteDifferenceSpace(grid::Grids.ExtrudedFiniteDifferenceGrid) = + ExtrudedFiniteDifferenceSpace(grid, CellFace()) +CenterExtrudedFiniteDifferenceSpace(grid::Grids.ExtrudedFiniteDifferenceGrid) = + ExtrudedFiniteDifferenceSpace(grid, CellCenter()) +FaceExtrudedFiniteDifferenceSpace(space::ExtrudedFiniteDifferenceSpace) = + ExtrudedFiniteDifferenceSpace(grid(space), CellFace()) +CenterExtrudedFiniteDifferenceSpace(space::ExtrudedFiniteDifferenceSpace) = + ExtrudedFiniteDifferenceSpace(grid(space), CellCenter()) + + +local_dss_weights(space::ExtrudedFiniteDifferenceSpace) = + local_dss_weights(grid(space)) + +staggering(space::ExtrudedFiniteDifferenceSpace) = getfield(space, :staggering) +grid(space::ExtrudedFiniteDifferenceSpace) = getfield(space, :grid) +space(space::ExtrudedFiniteDifferenceSpace, staggering::Staggering) = + ExtrudedFiniteDifferenceSpace(grid(space), staggering) + +#= + +ExtrudedFiniteDifferenceSpace{S}( + horizontal_space::AbstractSpace, + vertical_space::FiniteDifferenceSpace, + hypsography::HypsographyAdaption = Flat(), +) where {S <: Staggering} = ExtrudedFiniteDifferenceSpace{S}( + Grids.ExtrudedFiniteDifferenceGrid(horizontal_space, vertical_space, hypsography), +) +=# + function issubspace( hspace::AbstractSpectralElementSpace, extruded_space::ExtrudedFiniteDifferenceSpace, ) - ehspace = Spaces.horizontal_space(extruded_space) - if hspace === ehspace - return true - end - # TODO: improve level handling - return Spaces.topology(hspace) === Spaces.topology(ehspace) && - quadrature_style(hspace) === quadrature_style(ehspace) + return grid(hspace) === grid(extruded_space).horizontal_grid +end +function issubspace( + level_space::SpectralElementSpace2D{<:Grids.LevelGrid}, + extruded_space::ExtrudedFiniteDifferenceSpace, +) + return grid(level_space).full_grid === grid(extruded_space) end Adapt.adapt_structure(to, space::ExtrudedFiniteDifferenceSpace) = ExtrudedFiniteDifferenceSpace( - space.staggering, - Adapt.adapt(to, Spaces.horizontal_space(space)), - Adapt.adapt(to, space.vertical_topology), - Adapt.adapt(to, space.hypsography), - Adapt.adapt(to, space.global_geometry), - Adapt.adapt(to, space.center_local_geometry), - Adapt.adapt(to, space.face_local_geometry), - Adapt.adapt(to, space.center_ghost_geometry), - Adapt.adapt(to, space.face_ghost_geometry), + Adapt.adapt(to, grid(space)), + staggering(space), ) +const ExtrudedFiniteDifferenceSpace2D = ExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedFiniteDifferenceGrid{<:Grids.SpectralElementGrid1D}, +} +const ExtrudedFiniteDifferenceSpace3D = ExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedFiniteDifferenceGrid{<:Grids.SpectralElementGrid2D}, +} +const ExtrudedSpectralElementSpace2D = + ExtrudedFiniteDifferenceSpace{<:Grids.ExtrudedSpectralElementGrid2D} +const ExtrudedSpectralElementSpace3D = + ExtrudedFiniteDifferenceSpace{<:Grids.ExtrudedSpectralElementGrid3D} -const CenterExtrudedFiniteDifferenceSpace = - ExtrudedFiniteDifferenceSpace{CellCenter} - -const FaceExtrudedFiniteDifferenceSpace = - ExtrudedFiniteDifferenceSpace{CellFace} - -const FaceExtrudedFiniteDifferenceSpace2D = - ExtrudedFiniteDifferenceSpace{CellFace, <:SpectralElementSpace1D} -const FaceExtrudedFiniteDifferenceSpace3D = - ExtrudedFiniteDifferenceSpace{CellFace, <:SpectralElementSpace2D} const CenterExtrudedFiniteDifferenceSpace2D = - ExtrudedFiniteDifferenceSpace{CellCenter, <:SpectralElementSpace1D} + CenterExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedFiniteDifferenceGrid{<:Grids.SpectralElementGrid1D}, + } const CenterExtrudedFiniteDifferenceSpace3D = - ExtrudedFiniteDifferenceSpace{CellCenter, <:SpectralElementSpace2D} - + CenterExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedFiniteDifferenceGrid{<:Grids.SpectralElementGrid2D}, + } +const FaceExtrudedFiniteDifferenceSpace2D = FaceExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedFiniteDifferenceGrid{<:Grids.SpectralElementGrid1D}, +} +const FaceExtrudedFiniteDifferenceSpace3D = FaceExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedFiniteDifferenceGrid{<:Grids.SpectralElementGrid2D}, +} -function ExtrudedFiniteDifferenceSpace{S}( - space::ExtrudedFiniteDifferenceSpace, -) where {S <: Staggering} - ExtrudedFiniteDifferenceSpace( - S(), - Spaces.horizontal_space(space), - space.vertical_topology, - space.hypsography, - space.global_geometry, - space.center_local_geometry, - space.face_local_geometry, - space.center_ghost_geometry, - space.face_ghost_geometry, - ) -end function Base.show(io::IO, space::ExtrudedFiniteDifferenceSpace) indent = get(io, :indent, 0) iio = IOContext(io, :indent => indent + 2) @@ -103,81 +181,36 @@ function Base.show(io::IO, space::ExtrudedFiniteDifferenceSpace) ) print(iio, " "^(indent + 2), "context: ") hspace = Spaces.horizontal_space(space) - Topologies.print_context(iio, hspace.topology.context) + Topologies.print_context(iio, Spaces.topology(hspace).context) println(iio) println(iio, " "^(indent + 2), "horizontal:") - println(iio, " "^(indent + 4), "mesh: ", hspace.topology.mesh) - println(iio, " "^(indent + 4), "quadrature: ", hspace.quadrature_style) + println(iio, " "^(indent + 4), "mesh: ", Spaces.topology(hspace).mesh) + println(iio, " "^(indent + 4), "quadrature: ", quadrature_style(hspace)) println(iio, " "^(indent + 2), "vertical:") print(iio, " "^(indent + 4), "mesh: ", space.vertical_topology.mesh) end -local_geometry_data(space::CenterExtrudedFiniteDifferenceSpace) = - space.center_local_geometry -local_geometry_data(space::FaceExtrudedFiniteDifferenceSpace) = - space.face_local_geometry +quadrature_style(space::ExtrudedFiniteDifferenceSpace) = + quadrature_style(grid(space)) +topology(space::ExtrudedFiniteDifferenceSpace) = topology(grid(space)) -# TODO: will need to be defined for distributed -ghost_geometry_data(space::CenterExtrudedFiniteDifferenceSpace) = - space.center_ghost_geometry -ghost_geometry_data(space::FaceExtrudedFiniteDifferenceSpace) = - space.face_ghost_geometry -function ExtrudedFiniteDifferenceSpace( - horizontal_space::H, - vertical_space::V, - hypsography::Flat = Flat(), -) where {H <: AbstractSpace, V <: FiniteDifferenceSpace} - staggering = vertical_space.staggering - vertical_topology = vertical_space.topology - global_geometry = horizontal_space.global_geometry - center_local_geometry = - product_geometry.( - horizontal_space.local_geometry, - vertical_space.center_local_geometry, - ) - face_local_geometry = - product_geometry.( - horizontal_space.local_geometry, - vertical_space.face_local_geometry, - ) - if horizontal_space isa SpectralElementSpace2D - center_ghost_geometry = - product_geometry.( - horizontal_space.ghost_geometry, - vertical_space.center_local_geometry, - ) - face_ghost_geometry = - product_geometry.( - horizontal_space.ghost_geometry, - vertical_space.face_local_geometry, - ) - else - center_ghost_geometry = nothing - face_ghost_geometry = nothing - end - return ExtrudedFiniteDifferenceSpace( - staggering, - horizontal_space, - vertical_topology, - hypsography, - global_geometry, - center_local_geometry, - face_local_geometry, - center_ghost_geometry, - face_ghost_geometry, - ) +horizontal_space(full_space::ExtrudedFiniteDifferenceSpace) = + space(grid(full_space).horizontal_grid, nothing) + +vertical_topology(space::ExtrudedFiniteDifferenceSpace) = + vertical_topology(grid(space)) + + + +function column(space::ExtrudedFiniteDifferenceSpace, colidx::Grids.ColumnIndex) + column_grid = column(grid(space), colidx) + FiniteDifferenceSpace(column_grid, space.staggering) end -quadrature_style(space::ExtrudedFiniteDifferenceSpace) = - Spaces.horizontal_space(space).quadrature_style -topology(space::ExtrudedFiniteDifferenceSpace) = - Spaces.horizontal_space(space).topology -ClimaComms.device(space::ExtrudedFiniteDifferenceSpace) = - ClimaComms.device(topology(space)) -vertical_topology(space::ExtrudedFiniteDifferenceSpace) = - space.vertical_topology + + Base.@propagate_inbounds function slab( space::ExtrudedFiniteDifferenceSpace, @@ -190,193 +223,35 @@ Base.@propagate_inbounds function slab( ) end -Base.@propagate_inbounds function column( - space::ExtrudedFiniteDifferenceSpace, - i, - j, - h, -) - FiniteDifferenceSpace( - space.staggering, - space.vertical_topology, - Geometry.CartesianGlobalGeometry(), - column(space.center_local_geometry, i, j, h), - column(space.face_local_geometry, i, j, h), - ) -end -Base.@propagate_inbounds function column( - space::ExtrudedFiniteDifferenceSpace, - i, - h, -) - FiniteDifferenceSpace( - space.staggering, - space.vertical_topology, - Geometry.CartesianGlobalGeometry(), - column(space.center_local_geometry, i, h), - column(space.face_local_geometry, i, h), - ) -end +# TODO: deprecate these +column(space::ExtrudedFiniteDifferenceSpace, i, j, h) = + column(space, Grids.ColumnIndex((i, j), h)) +column(space::ExtrudedFiniteDifferenceSpace, i, h) = + column(space, Grids.ColumnIndex((i,), h)) -Base.@propagate_inbounds function level( - space::CenterExtrudedFiniteDifferenceSpace, - v::Integer, -) - horizontal_space = Spaces.horizontal_space(space) - if horizontal_space isa SpectralElementSpace1D - SpectralElementSpace1D( - horizontal_space.topology, - horizontal_space.quadrature_style, - horizontal_space.global_geometry, - level(space.center_local_geometry, v), - horizontal_space.dss_weights, - ) - elseif horizontal_space isa SpectralElementSpace2D - SpectralElementSpace2D( - horizontal_space.topology, - horizontal_space.quadrature_style, - horizontal_space.global_geometry, - level(space.center_local_geometry, v), - level(space.center_ghost_geometry, v), - horizontal_space.local_dss_weights, - horizontal_space.ghost_dss_weights, - horizontal_space.internal_surface_geometry, - horizontal_space.boundary_surface_geometries, - ) - else - error("Unsupported horizontal space") - end -end -Base.@propagate_inbounds function level( - space::FaceExtrudedFiniteDifferenceSpace, - v::PlusHalf, -) - horizontal_space = Spaces.horizontal_space(space) - if horizontal_space isa SpectralElementSpace1D - @inbounds SpectralElementSpace1D( - horizontal_space.topology, - horizontal_space.quadrature_style, - horizontal_space.global_geometry, - level(space.face_local_geometry, v.i + 1), - horizontal_space.dss_weights, - ) - elseif horizontal_space isa SpectralElementSpace2D - @inbounds SpectralElementSpace2D( - horizontal_space.topology, - horizontal_space.quadrature_style, - horizontal_space.global_geometry, - level(space.face_local_geometry, v.i + 1), - level(space.face_ghost_geometry, v.i + 1), - horizontal_space.local_dss_weights, - horizontal_space.ghost_dss_weights, - horizontal_space.internal_surface_geometry, - horizontal_space.boundary_surface_geometries, - ) - else - error("Unsupported horizontal space") - end -end +level(space::CenterExtrudedFiniteDifferenceSpace2D, v::Integer) = + SpectralElementSpace1D(level(grid(space), v)) +level(space::FaceExtrudedFiniteDifferenceSpace2D, v::PlusHalf) = + SpectralElementSpace1D(level(grid(space), v)) +level(space::CenterExtrudedFiniteDifferenceSpace3D, v::Integer) = + SpectralElementSpace2D(level(grid(space), v)) +level(space::FaceExtrudedFiniteDifferenceSpace3D, v::PlusHalf) = + SpectralElementSpace2D(level(grid(space), v)) -nlevels(space::CenterExtrudedFiniteDifferenceSpace) = - size(space.center_local_geometry, 4) -nlevels(space::FaceExtrudedFiniteDifferenceSpace) = - size(space.face_local_geometry, 4) +nlevels(space::ExtrudedFiniteDifferenceSpace) = + size(local_geometry_data(space), 4) function left_boundary_name(space::ExtrudedFiniteDifferenceSpace) - boundaries = Topologies.boundaries(space.vertical_topology) + boundaries = Topologies.boundaries(Spaces.vertical_topology(space)) propertynames(boundaries)[1] end function right_boundary_name(space::ExtrudedFiniteDifferenceSpace) - boundaries = Topologies.boundaries(space.vertical_topology) + boundaries = Topologies.boundaries(Spaces.vertical_topology(space)) propertynames(boundaries)[2] end -function blockmat( - a::Geometry.Axis2Tensor{ - FT, - Tuple{Geometry.UAxis, Geometry.Covariant1Axis}, - SMatrix{1, 1, FT, 1}, - }, - b::Geometry.Axis2Tensor{ - FT, - Tuple{Geometry.WAxis, Geometry.Covariant3Axis}, - SMatrix{1, 1, FT, 1}, - }, -) where {FT} - A = Geometry.components(a) - B = Geometry.components(b) - Geometry.AxisTensor( - (Geometry.UWAxis(), Geometry.Covariant13Axis()), - SMatrix{2, 2}(A[1, 1], zero(FT), zero(FT), B[1, 1]), - ) -end - -function blockmat( - a::Geometry.Axis2Tensor{ - FT, - Tuple{Geometry.VAxis, Geometry.Covariant2Axis}, - SMatrix{1, 1, FT, 1}, - }, - b::Geometry.Axis2Tensor{ - FT, - Tuple{Geometry.WAxis, Geometry.Covariant3Axis}, - SMatrix{1, 1, FT, 1}, - }, -) where {FT} - A = Geometry.components(a) - B = Geometry.components(b) - Geometry.AxisTensor( - (Geometry.VWAxis(), Geometry.Covariant23Axis()), - SMatrix{2, 2}(A[1, 1], zero(FT), zero(FT), B[1, 1]), - ) -end - -function blockmat( - a::Geometry.Axis2Tensor{ - FT, - Tuple{Geometry.UVAxis, Geometry.Covariant12Axis}, - SMatrix{2, 2, FT, 4}, - }, - b::Geometry.Axis2Tensor{ - FT, - Tuple{Geometry.WAxis, Geometry.Covariant3Axis}, - SMatrix{1, 1, FT, 1}, - }, -) where {FT} - A = Geometry.components(a) - B = Geometry.components(b) - Geometry.AxisTensor( - (Geometry.UVWAxis(), Geometry.Covariant123Axis()), - SMatrix{3, 3}( - A[1, 1], - A[2, 1], - zero(FT), - A[1, 2], - A[2, 2], - zero(FT), - zero(FT), - zero(FT), - B[1, 1], - ), - ) -end - -function product_geometry( - horizontal_local_geometry::Geometry.LocalGeometry, - vertical_local_geometry::Geometry.LocalGeometry, -) - coordinates = Geometry.product_coordinates( - horizontal_local_geometry.coordinates, - vertical_local_geometry.coordinates, - ) - J = horizontal_local_geometry.J * vertical_local_geometry.J - WJ = horizontal_local_geometry.WJ * vertical_local_geometry.WJ - ∂x∂ξ = - blockmat(horizontal_local_geometry.∂x∂ξ, vertical_local_geometry.∂x∂ξ) - return Geometry.LocalGeometry(coordinates, J, WJ, ∂x∂ξ) -end function eachslabindex(cspace::CenterExtrudedFiniteDifferenceSpace) h_iter = eachslabindex(Spaces.horizontal_space(cspace)) @@ -388,3 +263,12 @@ function eachslabindex(fspace::FaceExtrudedFiniteDifferenceSpace) Nv = size(fspace.face_local_geometry, 4) return Iterators.product(1:Nv, h_iter) end + + +## aliases +const ExtrudedRectilinearSpectralElementSpace3D = ExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedRectilinearSpectralElementGrid3D, +} +const ExtrudedCubedSphereSpectralElementSpace3D = ExtrudedFiniteDifferenceSpace{ + <:Grids.ExtrudedCubedSphereSpectralElementGrid3D, +} diff --git a/src/Spaces/finitedifference.jl b/src/Spaces/finitedifference.jl index 38302bcba6..7a7b465666 100644 --- a/src/Spaces/finitedifference.jl +++ b/src/Spaces/finitedifference.jl @@ -1,25 +1,35 @@ abstract type AbstractFiniteDifferenceSpace <: AbstractSpace end -abstract type Staggering end +""" + FiniteDifferenceSpace( + grid::Grids.FiniteDifferenceGrid, + staggering::Staggering, + ) -""" Cell center location """ -struct CellCenter <: Staggering end -""" Cell face location """ -struct CellFace <: Staggering end +""" struct FiniteDifferenceSpace{ + G <: Grids.AbstractFiniteDifferenceGrid, S <: Staggering, - T <: Topologies.AbstractIntervalTopology, - GG, - LG, } <: AbstractFiniteDifferenceSpace + grid::G staggering::S - topology::T - global_geometry::GG - center_local_geometry::LG - face_local_geometry::LG end +FiniteDifferenceSpace( + topology::Topologies.IntervalTopology, + staggering::Staggering, +) = FiniteDifferenceSpace(Grids.FiniteDifferenceGrid(topology), staggering) + + +const FaceFiniteDifferenceSpace{G} = FiniteDifferenceSpace{G, CellFace} +const CenterFiniteDifferenceSpace{G} = FiniteDifferenceSpace{G, CellCenter} + +grid(space::AbstractFiniteDifferenceSpace) = getfield(space, :grid) +staggering(space::FiniteDifferenceSpace) = getfield(space, :staggering) + +space(grid::Grids.AbstractFiniteDifferenceGrid, staggering::Staggering) = + FiniteDifferenceSpace(grid, staggering) function Base.show(io::IO, space::FiniteDifferenceSpace) indent = get(io, :indent, 0) @@ -31,142 +41,70 @@ function Base.show(io::IO, space::FiniteDifferenceSpace) ":", ) print(iio, " "^(indent + 2), "context: ") - Topologies.print_context(iio, space.topology.context) + Topologies.print_context(iio, ClimaComms.context(space)) println(iio) - print(iio, " "^(indent + 2), "mesh: ", space.topology.mesh) + print(iio, " "^(indent + 2), "mesh: ", topology(space).mesh) end -function FiniteDifferenceSpace{S}( - topology::Topologies.IntervalTopology, -) where {S <: Staggering} - global_geometry = Geometry.CartesianGlobalGeometry() - mesh = topology.mesh - CT = Meshes.coordinate_type(mesh) - AIdx = Geometry.coordinate_axis(CT) - # TODO: FD operators hardcoded to work over the 3-axis, need to generalize - # similar to spectral operators - @assert AIdx == (3,) "FiniteDifference operations only work over the 3-axis (ZPoint) domain" - FT = eltype(CT) - ArrayType = ClimaComms.array_type(topology) - face_coordinates = collect(mesh.faces) - LG = Geometry.LocalGeometry{AIdx, CT, FT, SMatrix{1, 1, FT, 1}} - nface = length(face_coordinates) - Topologies.isperiodic(topology) - ncent = length(face_coordinates) - 1 - # contstruct on CPU, copy to device at end - center_local_geometry = DataLayouts.VF{LG}(Array{FT}, ncent) - face_local_geometry = DataLayouts.VF{LG}(Array{FT}, nface) - for i in 1:ncent - # centers - coord⁻ = Geometry.component(face_coordinates[i], 1) - coord⁺ = Geometry.component(face_coordinates[i + 1], 1) - # at the moment we use a "discrete Jacobian" - # ideally we should use the continuous quantity via the derivative of the warp function - # could we just define this then as deriv on the mesh element coordinates? - coord = (coord⁺ + coord⁻) / 2 - Δcoord = coord⁺ - coord⁻ - J = Δcoord - WJ = Δcoord - ∂x∂ξ = SMatrix{1, 1}(J) - center_local_geometry[i] = Geometry.LocalGeometry( - CT(coord), - J, - WJ, - Geometry.AxisTensor( - (Geometry.LocalAxis{AIdx}(), Geometry.CovariantAxis{AIdx}()), - ∂x∂ξ, - ), +FaceFiniteDifferenceSpace(grid::Grids.AbstractFiniteDifferenceGrid) = + FiniteDifferenceSpace(grid, CellFace()) +CenterFiniteDifferenceSpace(grid::Grids.AbstractFiniteDifferenceGrid) = + FiniteDifferenceSpace(grid, CellCenter()) + +FaceFiniteDifferenceSpace(space::FiniteDifferenceSpace) = + FiniteDifferenceSpace(grid(space), CellFace()) +CenterFiniteDifferenceSpace(space::FiniteDifferenceSpace) = + FiniteDifferenceSpace(grid(space), CellCenter()) + +FaceFiniteDifferenceSpace(topology::Topologies.IntervalTopology) = + FiniteDifferenceSpace(Grids.FiniteDifferenceGrid(topology), CellFace()) +CenterFiniteDifferenceSpace(topology::Topologies.IntervalTopology) = + FiniteDifferenceSpace(Grids.FiniteDifferenceGrid(topology), CellCenter()) + +FaceFiniteDifferenceSpace(mesh::Meshes.IntervalMesh) = + FiniteDifferenceSpace(Grids.FiniteDifferenceGrid(mesh), CellFace()) +CenterFiniteDifferenceSpace(mesh::Meshes.IntervalMesh) = + FiniteDifferenceSpace(Grids.FiniteDifferenceGrid(mesh), CellCenter()) + +@inline function Base.getproperty(space::FiniteDifferenceSpace, name::Symbol) + if name == :topology + Base.depwarn( + "`space.topology` is deprecated, use `Spaces.topology(space)` instead", + :getproperty, ) - end - for i in 1:nface - coord = Geometry.component(face_coordinates[i], 1) - if i == 1 - # bottom face - if Topologies.isperiodic(topology) - Δcoord⁺ = - Geometry.component(face_coordinates[2], 1) - - Geometry.component(face_coordinates[1], 1) - Δcoord⁻ = - Geometry.component(face_coordinates[end], 1) - - Geometry.component(face_coordinates[end - 1], 1) - J = (Δcoord⁺ + Δcoord⁻) / 2 - WJ = J - else - coord⁺ = Geometry.component(face_coordinates[2], 1) - J = coord⁺ - coord - WJ = J / 2 - end - elseif !Topologies.isperiodic(topology) && i == nface - # top face - coord⁻ = Geometry.component(face_coordinates[i - 1], 1) - J = coord - coord⁻ - WJ = J / 2 - else - coord⁺ = Geometry.component(face_coordinates[i + 1], 1) - coord⁻ = Geometry.component(face_coordinates[i - 1], 1) - J = (coord⁺ - coord⁻) / 2 - WJ = J - end - ∂x∂ξ = SMatrix{1, 1}(J) - ∂ξ∂x = SMatrix{1, 1}(inv(J)) - face_local_geometry[i] = Geometry.LocalGeometry( - CT(coord), - J, - WJ, - Geometry.AxisTensor( - (Geometry.LocalAxis{AIdx}(), Geometry.CovariantAxis{AIdx}()), - ∂x∂ξ, - ), + return topology(space) + elseif name == :global_geometry + Base.depwarn( + "`space.global_geometry` is deprecated, use `Spaces.global_geometry(space)` instead", + :getproperty, + ) + return global_geometry(space) + elseif name == :center_local_geometry + Base.depwarn( + "`space.center_local_geometry` is deprecated, use `local_geometry_data(grid(space), Grids.CellCenter())` instead", + :getproperty, ) + return local_geometry_data(space, Grids.CellCenter()) + elseif name == :face_local_geometry + Base.depwarn( + "`space.face_local_geometry` is deprecated, use `local_geometry_data(grid(space), Grids.CellFace())` instead", + :getproperty, + ) + return local_geometry_data(space, Grids.CellFace()) end - return FiniteDifferenceSpace( - S(), - topology, - global_geometry, - Adapt.adapt(ArrayType, center_local_geometry), - Adapt.adapt(ArrayType, face_local_geometry), - ) + return getfield(space, name) end -FiniteDifferenceSpace{S}(mesh::Meshes.IntervalMesh) where {S <: Staggering} = - FiniteDifferenceSpace{S}(Topologies.IntervalTopology(mesh)) - -ClimaComms.device(space::FiniteDifferenceSpace) = - ClimaComms.device(space.topology) +Adapt.adapt_structure(to, space::FiniteDifferenceSpace) = + FiniteDifferenceSpace(Adapt.adapt(to, grid(space)), staggering(space)) -Adapt.adapt_structure(to, space::FiniteDifferenceSpace) = FiniteDifferenceSpace( - space.staggering, - Adapt.adapt(to, space.topology), - Adapt.adapt(to, space.global_geometry), - Adapt.adapt(to, space.center_local_geometry), - Adapt.adapt(to, space.face_local_geometry), -) -const CenterFiniteDifferenceSpace = FiniteDifferenceSpace{CellCenter} -const FaceFiniteDifferenceSpace = FiniteDifferenceSpace{CellFace} -function FiniteDifferenceSpace{S}( - space::FiniteDifferenceSpace, -) where {S <: Staggering} - FiniteDifferenceSpace( - S(), - space.topology, - space.global_geometry, - space.center_local_geometry, - space.face_local_geometry, - ) -end - -Base.length(space::FiniteDifferenceSpace) = length(coordinates_data(space)) - -topology(space::FiniteDifferenceSpace) = space.topology -vertical_topology(space::FiniteDifferenceSpace) = space.topology nlevels(space::FiniteDifferenceSpace) = length(space) +# TODO: deprecate? +Base.length(space::FiniteDifferenceSpace) = length(coordinates_data(space)) -local_geometry_data(space::CenterFiniteDifferenceSpace) = - space.center_local_geometry -local_geometry_data(space::FaceFiniteDifferenceSpace) = - space.face_local_geometry Base.@deprecate z_component(::Type{T}) where {T} Δz_metric_component(T) false @@ -200,13 +138,13 @@ function Δz_data(space::AbstractSpace) ) end -function left_boundary_name(space::FiniteDifferenceSpace) - boundaries = Topologies.boundaries(Spaces.topology(space)) +function left_boundary_name(space::AbstractSpace) + boundaries = Topologies.boundaries(Spaces.vertical_topology(space)) propertynames(boundaries)[1] end -function right_boundary_name(space::FiniteDifferenceSpace) - boundaries = Topologies.boundaries(Spaces.topology(space)) +function right_boundary_name(space::AbstractSpace) + boundaries = Topologies.boundaries(Spaces.vertical_topology(space)) propertynames(boundaries)[2] end diff --git a/src/Spaces/pointspace.jl b/src/Spaces/pointspace.jl index 4aee1a8ff4..10e0c36dfc 100644 --- a/src/Spaces/pointspace.jl +++ b/src/Spaces/pointspace.jl @@ -17,12 +17,11 @@ ClimaComms.device(space::PointSpace) = ClimaComms.context(space::PointSpace) = ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded()) -#= + PointSpace(x::Geometry.LocalGeometry) = PointSpace(ClimaComms.CPUSingleThreaded(), x) PointSpace(x::Geometry.AbstractPoint) = PointSpace(ClimaComms.CPUSingleThreaded(), x) -=# function PointSpace(device::ClimaComms.AbstractDevice, x) context = ClimaComms.SingletonCommsContext(device) diff --git a/src/Spaces/spectralelement.jl b/src/Spaces/spectralelement.jl index 2d7212befb..b44faf8072 100644 --- a/src/Spaces/spectralelement.jl +++ b/src/Spaces/spectralelement.jl @@ -3,615 +3,183 @@ abstract type AbstractSpectralElementSpace <: AbstractSpace end Topologies.nlocalelems(space::AbstractSpectralElementSpace) = Topologies.nlocalelems(Spaces.topology(space)) -local_geometry_data(space::AbstractSpectralElementSpace) = space.local_geometry -ghost_geometry_data(space::AbstractSpectralElementSpace) = space.ghost_geometry + + +quadrature_style(space::AbstractSpectralElementSpace) = + quadrature_style(grid(space)) +local_dss_weights(space::AbstractSpectralElementSpace) = + local_dss_weights(grid(space)) + +horizontal_space(space::AbstractSpectralElementSpace) = space +nlevels(space::AbstractSpectralElementSpace) = 1 eachslabindex(space::AbstractSpectralElementSpace) = 1:Topologies.nlocalelems(Spaces.topology(space)) +staggering(space::AbstractSpectralElementSpace) = nothing + function Base.show(io::IO, space::AbstractSpectralElementSpace) indent = get(io, :indent, 0) iio = IOContext(io, :indent => indent + 2) println(io, nameof(typeof(space)), ":") - if hasfield(typeof(space), :topology) + if hasfield(typeof(grid(space)), :topology) # some reduced spaces (like slab space) do not have topology print(iio, " "^(indent + 2), "context: ") - Topologies.print_context(iio, space.topology.context) + Topologies.print_context(iio, grid(space).topology.context) println(iio) - println(iio, " "^(indent + 2), "mesh: ", space.topology.mesh) - end - print(iio, " "^(indent + 2), "quadrature: ", space.quadrature_style) -end - -ClimaComms.device(space::AbstractSpectralElementSpace) = - ClimaComms.device(topology(space)) -ClimaComms.array_type(space::AbstractSpectralElementSpace) = - ClimaComms.array_type(ClimaComms.device(space)) -topology(space::AbstractSpectralElementSpace) = space.topology -quadrature_style(space::AbstractSpectralElementSpace) = space.quadrature_style - -abstract type AbstractPerimeter end - -""" - Perimeter2D <: AbstractPerimeter - -Iterate over the perimeter degrees of freedom of a 2D spectral element. -""" -struct Perimeter2D{Nq} <: AbstractPerimeter end - -""" - Perimeter2D(Nq) - -Construct a perimeter iterator for a 2D spectral element of degree `(Nq-1)`. -""" -Perimeter2D(Nq) = Perimeter2D{Nq}() -Adapt.adapt_structure(to, x::Perimeter2D) = x - -function Base.iterate(perimeter::Perimeter2D{Nq}, loc = 1) where {Nq} - if loc < 5 - return (Topologies.vertex_node_index(loc, Nq), loc + 1) - elseif loc ≤ nperimeter2d(Nq) - f = cld(loc - 4, Nq - 2) - n = mod(loc - 4, Nq - 2) == 0 ? (Nq - 2) : mod(loc - 4, Nq - 2) - return (Topologies.face_node_index(f, Nq, 1 + n), loc + 1) - else - return nothing + println(iio, " "^(indent + 2), "mesh: ", grid(space).topology.mesh) end + print(iio, " "^(indent + 2), "quadrature: ", grid(space).quadrature_style) end -function Base.getindex(perimeter::Perimeter2D{Nq}, loc = 1) where {Nq} - if loc < 1 || loc > nperimeter2d(Nq) - return (-1, -1) - elseif loc < 5 - return Topologies.vertex_node_index(loc, Nq) - else - f = cld(loc - 4, Nq - 2) - n = mod(loc - 4, Nq - 2) == 0 ? (Nq - 2) : mod(loc - 4, Nq - 2) - return Topologies.face_node_index(f, Nq, 1 + n) - end -end -nperimeter2d(Nq) = 4 + (Nq - 2) * 4 -nperimeter(::Perimeter2D{Nq}) where {Nq} = nperimeter2d(Nq) -Base.length(::Perimeter2D{Nq}) where {Nq} = nperimeter2d(Nq) +# 1D """ - SpectralElementSpace1D <: AbstractSpace - -A one-dimensional space: within each element the space is represented as a polynomial. + SpectralElementSpace1D(grid::SpectralElementGrid1D) """ -struct SpectralElementSpace1D{ - T, - Q, - GG <: Geometry.AbstractGlobalGeometry, - LG, - D, -} <: AbstractSpectralElementSpace - topology::T - quadrature_style::Q - global_geometry::GG - local_geometry::LG - dss_weights::D +struct SpectralElementSpace1D{G} <: AbstractSpectralElementSpace + grid::G end +space(grid::Grids.SpectralElementGrid1D, ::Nothing) = + SpectralElementSpace1D(grid) +grid(space::Spaces.SpectralElementSpace1D) = getfield(space, :grid) function SpectralElementSpace1D( topology::Topologies.IntervalTopology, - quadrature_style, + quadrature_style::Quadratures.QuadratureStyle, ) - global_geometry = Geometry.CartesianGlobalGeometry() - CoordType = Topologies.coordinate_type(topology) - AIdx = Geometry.coordinate_axis(CoordType) - FT = eltype(CoordType) - nelements = Topologies.nlocalelems(topology) - Nq = Quadratures.degrees_of_freedom(quadrature_style) - - LG = Geometry.LocalGeometry{AIdx, CoordType, FT, SMatrix{1, 1, FT, 1}} - local_geometry = DataLayouts.IFH{LG, Nq}(Array{FT}, nelements) - quad_points, quad_weights = - Quadratures.quadrature_points(FT, quadrature_style) - - for elem in 1:nelements - local_geometry_slab = slab(local_geometry, elem) - for i in 1:Nq - ξ = quad_points[i] - # TODO: we need to massage the coordinate points because the grid is assumed 2D - vcoords = Topologies.vertex_coordinates(topology, elem) - x = Geometry.linear_interpolate(vcoords, ξ) - ∂x∂ξ = - ( - Geometry.component(vcoords[2], 1) - - Geometry.component(vcoords[1], 1) - ) / 2 - J = abs(∂x∂ξ) - WJ = J * quad_weights[i] - local_geometry_slab[i] = Geometry.LocalGeometry( - x, - J, - WJ, - Geometry.AxisTensor( - ( - Geometry.LocalAxis{AIdx}(), - Geometry.CovariantAxis{AIdx}(), - ), - ∂x∂ξ, - ), - ) - end - end - dss_weights = copy(local_geometry.J) - dss_weights .= one(FT) - dss_1d!(topology, dss_weights) - dss_weights = one(FT) ./ dss_weights - - return SpectralElementSpace1D( - topology, - quadrature_style, - global_geometry, - local_geometry, - dss_weights, - ) + grid = Grids.SpectralElementGrid1D(topology, quadrature_style) + SpectralElementSpace1D(grid) end -nlevels(space::SpectralElementSpace1D) = 1 -const IntervalSpectralElementSpace1D = SpectralElementSpace1D{ - <:Topologies.IntervalTopology{ - <:ClimaComms.AbstractCommsContext, - <:Meshes.IntervalMesh, - }, -} +@inline function Base.getproperty(space::SpectralElementSpace1D, name::Symbol) + if name == :topology + Base.depwarn( + "`space.topology` is deprecated, use `Spaces.topology(space)` instead", + :getproperty, + ) + return topology(space) + elseif name == :quadrature_style + Base.depwarn( + "`space.quadrature_style` is deprecated, use `Spaces.quadrature_style(space)` instead", + :getproperty, + ) + return quadrature_style(space) + elseif name == :global_geometry + Base.depwarn( + "`space.global_geometry` is deprecated, use `Spaces.global_geometry(space)` instead", + :getproperty, + ) + return global_geometry(space) + elseif name == :local_geometry + Base.depwarn( + "`space.local_geometry` is deprecated, use `Spaces.local_geometry_data(space)` instead", + :getproperty, + ) + return local_geometry_data(space) + elseif name == :local_dss_weights + Base.depwarn( + "`space.local_dss_weights` is deprecated, use `Spaces.local_dss_weights(space)` instead", + :getproperty, + ) + return local_dss_weights(space) + end + return getfield(space, name) +end +# 2D """ - SpectralElementSpace2D <: AbstractSpace - -A two-dimensional space: within each element the space is represented as a polynomial. + SpectralElementSpace2D(grid::SpectralElementGrid1D) """ -struct SpectralElementSpace2D{ - T, - Q, - GG <: Geometry.AbstractGlobalGeometry, - LG, - D, - IS, - BS, -} <: AbstractSpectralElementSpace - topology::T - quadrature_style::Q - global_geometry::GG - local_geometry::LG - ghost_geometry::LG - local_dss_weights::D - ghost_dss_weights::D - internal_surface_geometry::IS - boundary_surface_geometries::BS +struct SpectralElementSpace2D{G} <: AbstractSpectralElementSpace + grid::G end +space(grid::Grids.SpectralElementGrid2D, ::Nothing) = + SpectralElementSpace2D(grid) +grid(space::Spaces.SpectralElementSpace2D) = getfield(space, :grid) -Adapt.adapt_structure(to, space::SpectralElementSpace2D) = - SpectralElementSpace2D( - nothing, # drop topology - Adapt.adapt(to, space.quadrature_style), - Adapt.adapt(to, space.global_geometry), - Adapt.adapt(to, space.local_geometry), - Adapt.adapt(to, space.ghost_geometry), - Adapt.adapt(to, space.local_dss_weights), - Adapt.adapt(to, space.ghost_dss_weights), - Adapt.adapt(to, space.internal_surface_geometry), - Adapt.adapt(to, space.boundary_surface_geometries), - ) - - - -""" - SpectralElementSpace2D(topology, quadrature_style; enable_bubble) - -Construct a `SpectralElementSpace2D` instance given a `topology` and `quadrature`. The -flag `enable_bubble` enables the `bubble correction` for more accurate element areas. - -# Input arguments: -- topology: Topology2D -- quadrature_style: QuadratureStyle -- enable_bubble: Bool - -The idea behind the so-called `bubble_correction` is that the numerical area -of the domain (e.g., the sphere) is given by the sum of nodal integration weights -times their corresponding Jacobians. However, this discrete sum is not exactly -equal to the exact geometric area (4pi*radius^2 for the sphere). To make these equal, -the "epsilon bubble" approach modifies the inner weights in each element so that -geometric and numerical areas of each element match. - -Let ``\\Delta A^e := A^e_{exact} - A^e_{approx}``, then, in -the case of linear elements, we correct ``W_{i,j} J^e_{i,j}`` by: -```math -\\widehat{W_{i,j} J^e}_{i,j} = W_{i,j} J^e_{i,j} + \\Delta A^e * W_{i,j} / Nq^2 . -``` -and the case of non linear elements, by -```math -\\widehat{W_{i,j} J^e}_{i,j} = W_{i,j} J^e_{i,j} \\left( 1 + \\tilde{A}^e \\right) , -``` -where ``\\tilde{A}^e`` is the approximated area given by the sum of the interior nodal integration weights. - -Note: This is accurate only for cubed-spheres of the [`Meshes.EquiangularCubedSphere`](@ref) and -[`Meshes.EquidistantCubedSphere`](@ref) type, not for [`Meshes.ConformalCubedSphere`](@ref). -""" function SpectralElementSpace2D( - topology, - quadrature_style; - enable_bubble = false, + topology::Topologies.Topology2D, + quadrature_style::Quadratures.QuadratureStyle; + kwargs..., ) + grid = Grids.SpectralElementGrid2D(topology, quadrature_style; kwargs...) + SpectralElementSpace2D(grid) +end - # 1. compute localgeom for local elememts - # 2. ghost exchange of localgeom - # 3. do a round of dss on WJs - # 4. compute dss weights (WJ ./ dss(WJ)) (local and ghost) - - # DSS on a field would consist of - # 1. copy to send buffers - # 2. start exchange - # 3. dss of internal connections - # - option for weighting and transformation - # 4. finish exchange - # 5. dss of ghost connections - - ### How to DSS multiple fields? - # 1. allocate buffers externally - DA = ClimaComms.array_type(topology) - domain = Topologies.domain(topology) - if domain isa Domains.SphereDomain - CoordType3D = Topologies.coordinate_type(topology) - FT = Geometry.float_type(CoordType3D) - CoordType2D = Geometry.LatLongPoint{FT} # Domains.coordinate_type(topology) - global_geometry = - Geometry.SphericalGlobalGeometry(topology.mesh.domain.radius) - else - CoordType2D = Topologies.coordinate_type(topology) - FT = Geometry.float_type(CoordType2D) - global_geometry = Geometry.CartesianGlobalGeometry() - end - AIdx = Geometry.coordinate_axis(CoordType2D) - nlelems = Topologies.nlocalelems(topology) - ngelems = Topologies.nghostelems(topology) - Nq = Quadratures.degrees_of_freedom(quadrature_style) - high_order_quadrature_style = Spaces.Quadratures.GLL{Nq * 2}() - high_order_Nq = Quadratures.degrees_of_freedom(high_order_quadrature_style) - - LG = Geometry.LocalGeometry{AIdx, CoordType2D, FT, SMatrix{2, 2, FT, 4}} - - local_geometry = DataLayouts.IJFH{LG, Nq}(Array{FT}, nlelems) - ghost_geometry = DataLayouts.IJFH{LG, Nq}(Array{FT}, ngelems) - - quad_points, quad_weights = - Quadratures.quadrature_points(FT, quadrature_style) - high_order_quad_points, high_order_quad_weights = - Quadratures.quadrature_points(FT, high_order_quadrature_style) - for (lidx, elem) in enumerate(Topologies.localelems(topology)) - elem_area = zero(FT) - high_order_elem_area = zero(FT) - Δarea = zero(FT) - interior_elem_area = zero(FT) - rel_interior_elem_area_Δ = zero(FT) - local_geometry_slab = slab(local_geometry, lidx) - # high-order quadrature loop for computing geometric element face area. - for i in 1:high_order_Nq, j in 1:high_order_Nq - ξ = SVector(high_order_quad_points[i], high_order_quad_points[j]) - u, ∂u∂ξ = - compute_local_geometry(global_geometry, topology, elem, ξ, AIdx) - J_high_order = det(Geometry.components(∂u∂ξ)) - WJ_high_order = - J_high_order * - high_order_quad_weights[i] * - high_order_quad_weights[j] - high_order_elem_area += WJ_high_order - end - # low-order quadrature loop for computing numerical element face area - for i in 1:Nq, j in 1:Nq - ξ = SVector(quad_points[i], quad_points[j]) - u, ∂u∂ξ = - compute_local_geometry(global_geometry, topology, elem, ξ, AIdx) - J = det(Geometry.components(∂u∂ξ)) - WJ = J * quad_weights[i] * quad_weights[j] - elem_area += WJ - if !enable_bubble - local_geometry_slab[i, j] = - Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) - end - end - - # If enabled, apply bubble correction - if enable_bubble - if abs(elem_area - high_order_elem_area) ≤ eps(FT) - for i in 1:Nq, j in 1:Nq - ξ = SVector(quad_points[i], quad_points[j]) - u, ∂u∂ξ = compute_local_geometry( - global_geometry, - topology, - elem, - ξ, - AIdx, - ) - J = det(Geometry.components(∂u∂ξ)) - WJ = J * quad_weights[i] * quad_weights[j] - local_geometry_slab[i, j] = - Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) - end - else - # The idea behind the so-called `bubble_correction` is that - # the numerical area of the domain (e.g., the sphere) is given by the sum - # of nodal integration weights times their corresponding Jacobians. However, - # this discrete sum is not exactly equal to the exact geometric area - # (4pi*radius^2 for the sphere). It is required that numerical area = geometric area. - # The "epsilon bubble" approach modifies the inner weights in each - # element so that geometric and numerical areas of each element match. - - # Compute difference between geometric area of an element and its approximate numerical area - Δarea = high_order_elem_area - elem_area - - # Linear elements: Nq == 2 (SpectralElementSpace2D cannot have Nq < 2) - # Use uniform bubble correction - if Nq == 2 - for i in 1:Nq, j in 1:Nq - ξ = SVector(quad_points[i], quad_points[j]) - u, ∂u∂ξ = compute_local_geometry( - global_geometry, - topology, - elem, - ξ, - AIdx, - ) - J = det(Geometry.components(∂u∂ξ)) - J += Δarea / Nq^2 - WJ = J * quad_weights[i] * quad_weights[j] - local_geometry_slab[i, j] = - Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) - end - else # Higher-order elements: Use HOMME bubble correction for the interior nodes - for i in 2:(Nq - 1), j in 2:(Nq - 1) - ξ = SVector(quad_points[i], quad_points[j]) - u, ∂u∂ξ = compute_local_geometry( - global_geometry, - topology, - elem, - ξ, - AIdx, - ) - J = det(Geometry.components(∂u∂ξ)) - WJ = J * quad_weights[i] * quad_weights[j] - interior_elem_area += WJ - end - # Check that interior_elem_area is not too small - if abs(interior_elem_area) ≤ sqrt(eps(FT)) - error( - "Bubble correction cannot be performed; sum of inner weights is too small.", - ) - end - rel_interior_elem_area_Δ = Δarea / interior_elem_area - - for i in 1:Nq, j in 1:Nq - ξ = SVector(quad_points[i], quad_points[j]) - u, ∂u∂ξ = compute_local_geometry( - global_geometry, - topology, - elem, - ξ, - AIdx, - ) - J = det(Geometry.components(∂u∂ξ)) - # Modify J only for interior nodes - if i != 1 && j != 1 && i != Nq && j != Nq - J *= (1 + rel_interior_elem_area_Δ) - end - WJ = J * quad_weights[i] * quad_weights[j] - # Finally allocate local geometry - local_geometry_slab[i, j] = - Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) - end - end - end - end - end - - # alternatively, we could do a ghost exchange here? - if topology isa Topologies.Topology2D - for (ridx, elem) in enumerate(Topologies.ghostelems(topology)) - ghost_geometry_slab = slab(ghost_geometry, ridx) - for i in 1:Nq, j in 1:Nq - ξ = SVector(quad_points[i], quad_points[j]) - u, ∂u∂ξ = compute_local_geometry( - global_geometry, - topology, - elem, - ξ, - AIdx, - ) - J = det(Geometry.components(∂u∂ξ)) - WJ = J * quad_weights[i] * quad_weights[j] - - ghost_geometry_slab[i, j] = - Geometry.LocalGeometry(u, J, WJ, ∂u∂ξ) - end - end - if !isnothing(ghost_geometry) && DA ≠ Array - ghost_geometry = DataLayouts.rebuild(ghost_geometry, DA) - end - end - # dss_weights = J ./ dss(J) - J = DataLayouts.rebuild(local_geometry.J, DA) - dss_local_weights = copy(J) - if quadrature_style isa Quadratures.GLL - dss!(dss_local_weights, topology, quadrature_style) - end - dss_local_weights .= J ./ dss_local_weights - dss_ghost_weights = copy(J) # not currently used - - SG = Geometry.SurfaceGeometry{ - FT, - Geometry.AxisVector{FT, Geometry.LocalAxis{AIdx}, SVector{2, FT}}, - } - interior_faces = Array(Topologies.interior_faces(topology)) - - if quadrature_style isa Quadratures.GLL - internal_surface_geometry = - DataLayouts.IFH{SG, Nq}(Array{FT}, length(interior_faces)) - for (iface, (lidx⁻, face⁻, lidx⁺, face⁺, reversed)) in - enumerate(interior_faces) - internal_surface_geometry_slab = - slab(internal_surface_geometry, iface) - - local_geometry_slab⁻ = slab(local_geometry, lidx⁻) - local_geometry_slab⁺ = slab(local_geometry, lidx⁺) - - for q in 1:Nq - sgeom⁻ = compute_surface_geometry( - local_geometry_slab⁻, - quad_weights, - face⁻, - q, - false, - ) - sgeom⁺ = compute_surface_geometry( - local_geometry_slab⁺, - quad_weights, - face⁺, - q, - reversed, - ) - - @assert sgeom⁻.sWJ ≈ sgeom⁺.sWJ - @assert sgeom⁻.normal ≈ -sgeom⁺.normal - - internal_surface_geometry_slab[q] = sgeom⁻ - end - end - internal_surface_geometry = - DataLayouts.rebuild(internal_surface_geometry, DA) - - boundary_surface_geometries = - map(Topologies.boundary_tags(topology)) do boundarytag - boundary_faces = - Topologies.boundary_faces(topology, boundarytag) - boundary_surface_geometry = - DataLayouts.IFH{SG, Nq}(Array{FT}, length(boundary_faces)) - for (iface, (elem, face)) in enumerate(boundary_faces) - boundary_surface_geometry_slab = - slab(boundary_surface_geometry, iface) - local_geometry_slab = slab(local_geometry, elem) - for q in 1:Nq - boundary_surface_geometry_slab[q] = - compute_surface_geometry( - local_geometry_slab, - quad_weights, - face, - q, - false, - ) - end - end - DataLayouts.rebuild(boundary_surface_geometry, DA) - end - else - internal_surface_geometry = nothing - boundary_surface_geometries = nothing +@inline function Base.getproperty(space::SpectralElementSpace2D, name::Symbol) + if name == :topology + Base.depwarn( + "`space.topology` is deprecated, use `Spaces.topology(space)` instead", + :getproperty, + ) + return topology(space) + elseif name == :quadrature_style + Base.depwarn( + "`space.quadrature_style` is deprecated, use `Spaces.quadrature_style(space)` instead", + :getproperty, + ) + return quadrature_style(space) + elseif name == :global_geometry + Base.depwarn( + "`space.global_geometry` is deprecated, use `Spaces.global_geometry(space)` instead", + :getproperty, + ) + return global_geometry(space) + elseif name == :local_geometry + Base.depwarn( + "`space.local_geometry` is deprecated, use `Spaces.local_geometry_data(space)` instead", + :getproperty, + ) + return local_geometry_data(space) + elseif name == :ghost_geometry + Base.depwarn( + "`space.ghost_geometry` is deprecated, use `nothing` instead", + :getproperty, + ) + return nothing + elseif name == :local_dss_weights + Base.depwarn( + "`space.local_dss_weights` is deprecated, use `Spaces.local_dss_weights(space)` instead", + :getproperty, + ) + return local_dss_weights(space) + elseif name == :ghost_dss_weights + Base.depwarn( + "`space.ghost_dss_weights` is deprecated, use `nothing` instead", + :getproperty, + ) + return nothing + elseif name == :internal_surface_geometry + Base.depwarn( + "`space.internal_surface_geometry` is deprecated, use `Spaces.grid(space).internal_surface_geometry` instead", + :getproperty, + ) + return grid(space).internal_surface_geometry + elseif name == :boundary_surface_geometries + Base.depwarn( + "`space.boundary_surface_geometries` is deprecated, use `Spaces.grid(space).boundary_surface_geometries` instead", + :getproperty, + ) + return grid(space).boundary_surface_geometries end - return SpectralElementSpace2D( - topology, - quadrature_style, - global_geometry, - DataLayouts.rebuild(local_geometry, DA), - ghost_geometry, - dss_local_weights, - dss_ghost_weights, - internal_surface_geometry, - boundary_surface_geometries, - ) + return getfield(space, name) end -nlevels(space::SpectralElementSpace2D) = 1 -perimeter(space::SpectralElementSpace2D) = - Perimeter2D(Quadratures.degrees_of_freedom(space.quadrature_style)) - -const RectilinearSpectralElementSpace2D = SpectralElementSpace2D{ - <:Topologies.Topology2D{ - <:ClimaComms.AbstractCommsContext, - <:Meshes.RectilinearMesh, - }, -} - -const CubedSphereSpectralElementSpace2D = SpectralElementSpace2D{ - <:Topologies.Topology2D{ - <:ClimaComms.AbstractCommsContext, - <:Meshes.AbstractCubedSphere, - }, -} - -function compute_local_geometry( - global_geometry::Geometry.SphericalGlobalGeometry, - topology, - elem, - ξ, - AIdx, -) - x = Meshes.coordinates(topology.mesh, elem, ξ) - u = Geometry.LatLongPoint(x, global_geometry) - ∂x∂ξ = Geometry.AxisTensor( - (Geometry.Cartesian123Axis(), Geometry.CovariantAxis{AIdx}()), - ForwardDiff.jacobian(ξ) do ξ - Geometry.components(Meshes.coordinates(topology.mesh, elem, ξ)) - end, - ) - G = Geometry.local_to_cartesian(global_geometry, u) - ∂u∂ξ = Geometry.project(Geometry.LocalAxis{AIdx}(), G' * ∂x∂ξ) - return u, ∂u∂ξ -end -function compute_local_geometry( - global_geometry::Geometry.AbstractGlobalGeometry, - topology, - elem, - ξ, - AIdx, -) - u = Meshes.coordinates(topology.mesh, elem, ξ) - ∂u∂ξ = Geometry.AxisTensor( - (Geometry.LocalAxis{AIdx}(), Geometry.CovariantAxis{AIdx}()), - ForwardDiff.jacobian(ξ) do ξ - Geometry.components(Meshes.coordinates(topology.mesh, elem, ξ)) - end, - ) +Adapt.adapt_structure(to, space::SpectralElementSpace2D) = + SpectralElementSpace2D(Adapt.adapt(to, grid(space))) - return u, ∂u∂ξ -end -function compute_surface_geometry( - local_geometry_slab, - quad_weights, - face, - q, - reversed = false, +function issubspace( + hspace::SpectralElementSpace2D{<:Grids.SpectralElementGrid2D}, + level_space::SpectralElementSpace2D{<:Grids.LevelGrid}, ) - Nq = length(quad_weights) - @assert size(local_geometry_slab) == (Nq, Nq, 1, 1, 1) - i, j = Topologies.face_node_index(face, Nq, q, reversed) - - local_geometry = local_geometry_slab[i, j] - @unpack J, ∂ξ∂x = local_geometry - - # surface mass matrix - n = if face == 4 - -J * ∂ξ∂x[1, :] * quad_weights[j] - elseif face == 2 - J * ∂ξ∂x[1, :] * quad_weights[j] - elseif face == 1 - -J * ∂ξ∂x[2, :] * quad_weights[i] - elseif face == 3 - J * ∂ξ∂x[2, :] * quad_weights[i] - end - sWJ = norm(n) - n = n / sWJ - return Geometry.SurfaceGeometry(sWJ, n) + return grid(hspace) === grid(level_space).full_grid.horizontal_grid end -function variational_solve!(data, space::AbstractSpace) - data .= RecursiveApply.rdiv.(data, space.local_geometry.WJ) -end """ SpectralElementSpaceSlab <: AbstractSpace @@ -632,13 +200,18 @@ const SpectralElementSpaceSlab2D = nlevels(space::SpectralElementSpaceSlab1D) = 1 nlevels(space::SpectralElementSpaceSlab2D) = 1 + + + + + Base.@propagate_inbounds function slab( space::AbstractSpectralElementSpace, v, h, ) SpectralElementSpaceSlab( - space.quadrature_style, + quadrature_style(space), slab(space.local_geometry, v, h), ) end @@ -657,40 +230,9 @@ Base.@propagate_inbounds function column(space::SpectralElementSpace2D, i, j, h) PointSpace(local_geometry) end -# XXX: this cannot take `space` as it must be constructed beforehand so -# that the `space` constructor can do DSS (to compute DSS weights) -function setup_comms( - Context::Type{<:ClimaComms.AbstractCommsContext}, - topology::Topologies.AbstractDistributedTopology, - quad_style::Spaces.Quadratures.QuadratureStyle, - Nv, - Nf = 2, -) - Ni = Quadratures.degrees_of_freedom(quad_style) - Nj = Ni - AT = Array # XXX: get this from `space`/`topology`? - FT = Geometry.float_type(Topologies.coordinate_type(topology)) - - # Determine send and receive buffer dimensions for each neighbor PID - # and add the neighbors in the same order as they are stored in - # `neighbor_pids`! - nbrs = ClimaComms.Neighbor[] - for (nidx, npid) in enumerate(Topologies.neighbors(topology)) - nse = Topologies.nsendelems(topology, nidx) - nge = Topologies.nghostelems(topology, nidx) - send_dims = (Nv, Ni, Nj, Nf, nse) - recv_dims = (Nv, Ni, Nj, Nf, nge) - push!( - nbrs, - ClimaComms.Neighbor(Context, npid, AT, FT, send_dims, recv_dims), - ) - end - return Context(nbrs) -end - function all_nodes(space::SpectralElementSpace2D) - Nq = Quadratures.degrees_of_freedom(space.quadrature_style) - nelem = Topologies.nlocalelems(space.topology) + Nq = Quadratures.degrees_of_freedom(quadrature_style(space)) + nelem = Topologies.nlocalelems(topology(space)) Iterators.product(Iterators.product(1:Nq, 1:Nq), 1:nelem) end @@ -703,7 +245,7 @@ first `((i,j), e)` triple. This function is experimental, and may change in future. """ unique_nodes(space::SpectralElementSpace2D) = - unique_nodes(space, space.quadrature_style) + unique_nodes(space, quadrature_style(space)) unique_nodes(space::SpectralElementSpace2D, quad::Quadratures.QuadratureStyle) = UniqueNodeIterator(space) @@ -718,8 +260,8 @@ Base.eltype(iter::UniqueNodeIterator{<:SpectralElementSpace2D}) = function Base.length(iter::UniqueNodeIterator{<:SpectralElementSpace2D}) space = iter.space - topology = space.topology - Nq = Quadratures.degrees_of_freedom(space.quadrature_style) + topology = Spaces.topology(space) + Nq = Quadratures.degrees_of_freedom(quadrature_style(space)) nelem = Topologies.nlocalelems(topology) nvert = length(Topologies.local_vertices(topology)) @@ -743,7 +285,7 @@ function Base.iterate( ((i, j), e), ) space = iter.space - Nq = Quadratures.degrees_of_freedom(space.quadrature_style) + Nq = Quadratures.degrees_of_freedom(quadrature_style(space)) while true # find next node i += 1 @@ -768,28 +310,28 @@ function Base.iterate( # this also doesn't deal with the case where eo == e if j == 1 # face 1 - eo, _, _ = Topologies.opposing_face(space.topology, e, 1) + eo, _, _ = Topologies.opposing_face(Spaces.topology(space), e, 1) if 0 < eo < e continue end end if i == Nq # face 2 - eo, _, _ = Topologies.opposing_face(space.topology, e, 2) + eo, _, _ = Topologies.opposing_face(Spaces.topology(space), e, 2) if 0 < eo < e continue end end if j == Nq # face 3 - eo, _, _ = Topologies.opposing_face(space.topology, e, 3) + eo, _, _ = Topologies.opposing_face(Spaces.topology(space), e, 3) if 0 < eo < e continue end end if i == 1 # face 4 - eo, _, _ = Topologies.opposing_face(space.topology, e, 4) + eo, _, _ = Topologies.opposing_face(Spaces.topology(space), e, 4) if 0 < eo < e continue end @@ -797,3 +339,9 @@ function Base.iterate( return ((i, j), e), ((i, j), e) end end + +## aliases +const RectilinearSpectralElementSpace2D = + SpectralElementSpace2D{<:Grids.RectilinearSpectralElementGrid2D} +const CubedSphereSpectralElementSpace2D = + SpectralElementSpace2D{<:Grids.CubedSphereSpectralElementGrid2D} diff --git a/src/Topologies/Topologies.jl b/src/Topologies/Topologies.jl index 958470637f..f0c7a92a3c 100644 --- a/src/Topologies/Topologies.jl +++ b/src/Topologies/Topologies.jl @@ -3,11 +3,17 @@ module Topologies using DocStringExtensions import ClimaComms, Adapt +using CUDA -import ClimaComms +import ..ClimaCore import ..Geometry import ..Domains: Domains, coordinate_type import ..Meshes: Meshes, domain, coordinates +import ..DataLayouts +import ..slab, ..column, ..level + + +using Memoize, WeakValueDicts """ AbstractTopology @@ -337,6 +343,9 @@ function boundary_faces end include("interval.jl") include("topology2d.jl") +include("dss_transform.jl") +include("dss.jl") + # deprecate @deprecate boundaries(topology::AbstractTopology) boundary_tags(topology) diff --git a/src/Topologies/dss.jl b/src/Topologies/dss.jl new file mode 100644 index 0000000000..fcd434a2de --- /dev/null +++ b/src/Topologies/dss.jl @@ -0,0 +1,876 @@ +using DocStringExtensions + +""" + DSSBuffer{G, D, A, B} + +# Fields +$(DocStringExtensions.FIELDS) +""" +struct DSSBuffer{S, G, D, A, B, VI} + "ClimaComms graph context for communication" + graph_context::G + "Array for storing perimeter data" + perimeter_data::D + "send buffer" + send_data::A + "recv buffer" + recv_data::A + "indexing array for loading send buffer from `perimeter_data`" + send_buf_idx::B + "indexing array for loading (and summing) data from recv buffer to `perimeter_data`" + recv_buf_idx::B + "field id for all scalar fields stored in the `data` array" + scalarfidx::VI + "field id for all covariant12vector fields stored in the `data` array" + covariant12fidx::VI + "field id for all contravariant12vector fields stored in the `data` array" + contravariant12fidx::VI + "internal local elements (lidx)" + internal_elems::VI + "local elements (lidx) located on process boundary" + perimeter_elems::VI +end + +""" + create_dss_buffer( + data::Union{DataLayouts.IJFH{S, Nij}, DataLayouts.VIJFH{S, Nij}}, + topology::Topology2D, + local_geometry = nothing, + local_weights = nothing, + ) where {S, Nij} + +Creates a [`DSSBuffer`](@ref) for the field data corresponding to `data` +""" +function create_dss_buffer( + data::Union{DataLayouts.IJFH{S, Nij}, DataLayouts.VIJFH{S, Nij}}, + topology::Topology2D, + local_geometry::Union{ + DataLayouts.IJFH{<:Any, Nij}, + DataLayouts.VIJFH{<:Any, Nij}, + Nothing, + } = nothing, + local_weights::Union{ + DataLayouts.IJFH{<:Any, Nij}, + DataLayouts.VIJFH{<:Any, Nij}, + Nothing, + } = nothing, +) where {S, Nij} + perimeter::Perimeter2D = Perimeter2D(Nij) + context = ClimaComms.context(topology) + DA = ClimaComms.array_type(topology) + convert_to_array = DA isa Array ? false : true + (_, _, _, Nv, Nh) = Base.size(data) + Np = length(perimeter) + Nf = + length(parent(data)) == 0 ? 0 : + cld(length(parent(data)), (Nij * Nij * Nv * Nh)) + nfacedof = Nij - 2 + T = eltype(parent(data)) + TS = _transformed_type(data, local_geometry, local_weights, DA) # extract transformed type + # Add TS for Covariant123Vector + # For DSS of Covariant123Vector, the third component is treated like a scalar + # and is not transformed + if eltype(data) <: Geometry.Covariant123Vector + TS = Geometry.UVWVector{T} + end + perimeter_data = DataLayouts.VIFH{TS, Np}(DA{T}(undef, Nv, Np, Nf, Nh)) + if context isa ClimaComms.SingletonCommsContext + graph_context = ClimaComms.SingletonGraphContext(context) + send_data, recv_data = T[], T[] + send_buf_idx, recv_buf_idx = Int[], Int[] + send_data, recv_data = DA{T}(undef, 0), DA{T}(undef, 0) + send_buf_idx, recv_buf_idx = DA{Int}(undef, 0), DA{Int}(undef, 0) + internal_elems = DA{Int}(1:nelems(topology)) + perimeter_elems = DA{Int}(undef, 0) + else + (; comm_vertex_lengths, comm_face_lengths) = topology + vertex_buffer_lengths = comm_vertex_lengths .* (Nv * Nf) + face_buffer_lengths = comm_face_lengths .* (Nv * Nf * nfacedof) + buffer_lengths = vertex_buffer_lengths .+ face_buffer_lengths + buffer_size = sum(buffer_lengths) + send_data = DA{T}(undef, buffer_size) + recv_data = DA{T}(undef, buffer_size) + neighbor_pids = topology.neighbor_pids + graph_context = ClimaComms.graph_context( + context, + send_data, + buffer_lengths, + neighbor_pids, + recv_data, + buffer_lengths, + neighbor_pids, + persistent = true, + ) + send_buf_idx, recv_buf_idx = compute_ghost_send_recv_idx(topology, Nij) + internal_elems = DA(topology.internal_elems) + perimeter_elems = DA(topology.perimeter_elems) + end + scalarfidx, covariant12fidx, contravariant12fidx = Int[], Int[], Int[] + supportedvectortypes = Union{ + Geometry.UVector, + Geometry.VVector, + Geometry.WVector, + Geometry.UVVector, + Geometry.UWVector, + Geometry.VWVector, + Geometry.UVWVector, + Geometry.Covariant12Vector, + Geometry.Covariant3Vector, + Geometry.Covariant123Vector, + Geometry.Contravariant12Vector, + Geometry.Contravariant3Vector, + } + + if S <: NamedTuple + for (i, fieldtype) in enumerate(S.parameters[2].types) + offset = DataLayouts.fieldtypeoffset(T, S, i) + ncomponents = DataLayouts.typesize(T, fieldtype) + if fieldtype <: Geometry.AxisVector # vector fields + if !(fieldtype <: supportedvectortypes) + @show fieldtype + @show supportedvectortypes + end + @assert fieldtype <: supportedvectortypes + if fieldtype <: Geometry.Covariant12Vector + push!(covariant12fidx, offset + 1) + elseif fieldtype <: Geometry.Covariant123Vector + push!(covariant12fidx, offset + 1) + push!(scalarfidx, offset + 3) + elseif fieldtype <: Geometry.Contravariant12Vector + push!(contravariant12fidx, offset + 1) + else + append!( + scalarfidx, + Vector((offset + 1):(offset + ncomponents)), + ) + end + elseif fieldtype <: NTuple # support a NTuple of primitive types + append!(scalarfidx, Vector((offset + 1):(offset + ncomponents))) + else # scalar fields + push!(scalarfidx, offset + 1) + end + end + else # deals with simple type, with single field (e.g: S = Float64, S = CovariantVector12, etc.) + ncomponents = DataLayouts.typesize(T, S) + if S <: Geometry.AxisVector # vector field + if !(S <: supportedvectortypes) + @show S + @show supportedvectortypes + end + @assert S <: supportedvectortypes + if S <: Geometry.Covariant12Vector + push!(covariant12fidx, 1) + elseif S <: Geometry.Covariant123Vector + push!(covariant12fidx, 1) + push!(scalarfidx, 3) + elseif S <: Geometry.Contravariant12Vector + push!(contravariant12fidx, 1) + else + append!(scalarfidx, Vector(1:ncomponents)) + end + elseif S <: NTuple # support a NTuple of primitive types + append!(scalarfidx, Vector(1:ncomponents)) + else # scalar field + push!(scalarfidx, 1) + end + end + scalarfidx = DA(scalarfidx) + covariant12fidx = DA(covariant12fidx) + contravariant12fidx = DA(contravariant12fidx) + G = typeof(graph_context) + D = typeof(perimeter_data) + A = typeof(send_data) + B = typeof(send_buf_idx) + VI = typeof(scalarfidx) + return DSSBuffer{S, G, D, A, B, VI}( + graph_context, + perimeter_data, + send_data, + recv_data, + send_buf_idx, + recv_buf_idx, + scalarfidx, + covariant12fidx, + contravariant12fidx, + internal_elems, + perimeter_elems, + ) +end + +Base.eltype(::DSSBuffer{S}) where {S} = S + +assert_same_eltype(::DataLayouts.AbstractData, ::DSSBuffer) = + error("Incorrect buffer eltype") +assert_same_eltype(::DataLayouts.AbstractData{S}, ::DSSBuffer{S}) where {S} = + nothing +assert_same_eltype(::DataLayouts.AbstractData, ::Nothing) = nothing + +""" + function dss_transform!( + device::ClimaComms.AbstractDevice, + dss_buffer::DSSBuffer, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + weight::DataLayouts.IJFH, + perimeter::AbstractPerimeter, + localelems::Vector{Int}, + ) + +Transforms vectors from Covariant axes to physical (local axis), weights the data at perimeter nodes, +and stores result in the `perimeter_data` array. This function calls the appropriate version of +`dss_transform!` based on the data layout of the input arguments. + +Arguments: + +- `dss_buffer`: [`DSSBuffer`](@ref) generated by `create_dss_buffer` function for field data +- `data`: field data +- `local_geometry`: local metric information defined at each node +- `weight`: local dss weights for horizontal space +- `perimeter`: perimeter iterator +- `localelems`: list of local elements to perform transformation operations on + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function dss_transform!( + device::ClimaComms.AbstractDevice, + dss_buffer::DSSBuffer, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + weight::DataLayouts.IJFH, + perimeter::Perimeter2D, + localelems::AbstractVector{Int}, +) + if !isempty(localelems) + (; scalarfidx, covariant12fidx, contravariant12fidx, perimeter_data) = + dss_buffer + (; ∂ξ∂x, ∂x∂ξ) = local_geometry + dss_transform!( + device, + perimeter_data, + data, + ∂ξ∂x, + ∂x∂ξ, + weight, + perimeter, + scalarfidx, + covariant12fidx, + contravariant12fidx, + localelems, + ) + end + return nothing +end +""" + dss_untransform!( + device::ClimaComms.AbstractDevice, + dss_buffer::DSSBuffer, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter::AbstractPerimeter, + ) + +Transforms the DSS'd local vectors back to Covariant12 vectors, and copies the DSS'd data from the +`perimeter_data` to `data`. This function calls the appropriate version of `dss_transform!` function +based on the data layout of the input arguments. + +Arguments: + +- `dss_buffer`: [`DSSBuffer`](@ref) generated by `create_dss_buffer` function for field data +- `data`: field data +- `local_geometry`: local metric information defined at each node +- `perimeter`: perimeter iterator +- `localelems`: list of local elements to perform transformation operations on + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function dss_untransform!( + device::ClimaComms.AbstractDevice, + dss_buffer::DSSBuffer, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + local_geometry::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter::Perimeter2D, + localelems::AbstractVector{Int}, +) + (; scalarfidx, covariant12fidx, contravariant12fidx, perimeter_data) = + dss_buffer + (; ∂ξ∂x, ∂x∂ξ) = local_geometry + dss_untransform!( + device, + perimeter_data, + data, + ∂ξ∂x, + ∂x∂ξ, + perimeter, + scalarfidx, + covariant12fidx, + contravariant12fidx, + localelems, + ) + return nothing +end + +""" + function dss_transform!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + ∂ξ∂x::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + ∂x∂ξ::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + weight::DataLayouts.IJFH, + perimeter::AbstractPerimeter, + scalarfidx::Vector{Int}, + covariant12fidx::Vector{Int}, + contravariant12fidx::Vector{Int}, + localelems::Vector{Int}, + ) + +Transforms vectors from Covariant axes to physical (local axis), weights +the data at perimeter nodes, and stores result in the `perimeter_data` array. + +Arguments: + +- `perimeter_data`: contains the perimeter field data, represented on the physical axis, corresponding to the full field data in `data` +- `data`: field data +- `∂ξ∂x`: partial derivatives of the map from `x` to `ξ`: `∂ξ∂x[i,j]` is ∂ξⁱ/∂xʲ +- `weight`: local dss weights for horizontal space +- `perimeter`: perimeter iterator +- `scalarfidx`: field index for scalar fields in the data layout +- `covariant12fidx`: field index for Covariant12 vector fields in the data layout +- `localelems`: list of local elements to perform transformation operations on + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function dss_transform!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + ∂ξ∂x::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + ∂x∂ξ::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + weight::DataLayouts.IJFH, + perimeter::Perimeter2D{Nq}, + scalarfidx::Vector{Int}, + covariant12fidx::Vector{Int}, + contravariant12fidx::Vector{Int}, + localelems::Vector{Int}, +) where {Nq} + pdata = parent(data) + pweight = parent(weight) + p∂x∂ξ = parent(∂x∂ξ) + p∂ξ∂x = parent(∂ξ∂x) + pperimeter_data = parent(perimeter_data) + (nlevels, _, nfid, nelems) = size(pperimeter_data) + nmetric = cld(length(p∂ξ∂x), prod(size(∂ξ∂x))) + sizet_data = (nlevels, Nq, Nq, nfid, nelems) + sizet_wt = (Nq, Nq, 1, nelems) + sizet_metric = (nlevels, Nq, Nq, nmetric, nelems) + + @inbounds for elem in localelems + for (p, (ip, jp)) in enumerate(perimeter) + pw = pweight[_get_idx(sizet_wt, (ip, jp, 1, elem))] + + for fidx in scalarfidx, level in 1:nlevels + data_idx = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) + pperimeter_data[level, p, fidx, elem] = pdata[data_idx] * pw + end + + for fidx in covariant12fidx, level in 1:nlevels + data_idx1 = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) + data_idx2 = + _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) + (idx11, idx12, idx21, idx22) = + _get_idx_metric(sizet_metric, (level, ip, jp, elem)) + pperimeter_data[level, p, fidx, elem] = + ( + p∂ξ∂x[idx11] * pdata[data_idx1] + + p∂ξ∂x[idx12] * pdata[data_idx2] + ) * pw + pperimeter_data[level, p, fidx + 1, elem] = + ( + p∂ξ∂x[idx21] * pdata[data_idx1] + + p∂ξ∂x[idx22] * pdata[data_idx2] + ) * pw + end + + for fidx in contravariant12fidx, level in 1:nlevels + data_idx1 = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) + data_idx2 = + _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) + (idx11, idx12, idx21, idx22) = + _get_idx_metric(sizet_metric, (level, ip, jp, elem)) + pperimeter_data[level, p, fidx, elem] = + ( + p∂x∂ξ[idx11] * pdata[data_idx1] + + p∂x∂ξ[idx21] * pdata[data_idx2] + ) * pw + pperimeter_data[level, p, fidx + 1, elem] = + ( + p∂x∂ξ[idx12] * pdata[data_idx1] + + p∂x∂ξ[idx22] * pdata[data_idx2] + ) * pw + end + end + end + return nothing +end +""" + function dss_untransform!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + ∂ξ∂x::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + ∂x∂ξ::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter::AbstractPerimeter, + scalarfidx::Vector{Int}, + covariant12fidx::Vector{Int}, + contravariant12fidx::Vector{Int}, + localelems::Vector{Int}, + ) + +Transforms the DSS'd local vectors back to Covariant12 vectors, and copies the DSS'd data from the +`perimeter_data` to `data`. + +Arguments: + +- `perimeter_data`: contains the perimeter field data, represented on the physical axis, corresponding to the full field data in `data` +- `data`: field data +- `∂x∂ξ`: partial derivatives of the map from `ξ` to `x`: `∂x∂ξ[i,j]` is ∂xⁱ/∂ξʲ +- `perimeter`: perimeter iterator +- `scalarfidx`: field index for scalar fields in the data layout +- `covariant12fidx`: field index for Covariant12 vector fields in the data layout + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" + +function dss_untransform!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + data::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + ∂ξ∂x::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + ∂x∂ξ::Union{DataLayouts.VIJFH, DataLayouts.IJFH}, + perimeter::Perimeter2D{Nq}, + scalarfidx::Vector{Int}, + covariant12fidx::Vector{Int}, + contravariant12fidx::Vector{Int}, + localelems::Vector{Int}, +) where {Nq} + pdata = parent(data) + p∂x∂ξ = parent(∂x∂ξ) + p∂ξ∂x = parent(∂ξ∂x) + pperimeter_data = parent(perimeter_data) + (nlevels, _, nfid, nelems) = size(pperimeter_data) + nmetric = cld(length(p∂ξ∂x), prod(size(∂ξ∂x))) + sizet_data = (nlevels, Nq, Nq, nfid, nelems) + sizet_metric = (nlevels, Nq, Nq, nmetric, nelems) + + @inbounds for elem in localelems + for (p, (ip, jp)) in enumerate(perimeter) + for fidx in scalarfidx + for level in 1:nlevels + data_idx = _get_idx(sizet_data, (level, ip, jp, fidx, elem)) + pdata[data_idx] = pperimeter_data[level, p, fidx, elem] + end + end + for fidx in covariant12fidx + for level in 1:nlevels + data_idx1 = + _get_idx(sizet_data, (level, ip, jp, fidx, elem)) + data_idx2 = + _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) + (idx11, idx12, idx21, idx22) = + _get_idx_metric(sizet_metric, (level, ip, jp, elem)) + pdata[data_idx1] = + p∂x∂ξ[idx11] * pperimeter_data[level, p, fidx, elem] + + p∂x∂ξ[idx12] * pperimeter_data[level, p, fidx + 1, elem] + pdata[data_idx2] = + p∂x∂ξ[idx21] * pperimeter_data[level, p, fidx, elem] + + p∂x∂ξ[idx22] * pperimeter_data[level, p, fidx + 1, elem] + end + end + for fidx in contravariant12fidx + for level in 1:nlevels + data_idx1 = + _get_idx(sizet_data, (level, ip, jp, fidx, elem)) + data_idx2 = + _get_idx(sizet_data, (level, ip, jp, fidx + 1, elem)) + (idx11, idx12, idx21, idx22) = + _get_idx_metric(sizet_metric, (level, ip, jp, elem)) + pdata[data_idx1] = + p∂ξ∂x[idx11] * pperimeter_data[level, p, fidx, elem] + + p∂ξ∂x[idx21] * pperimeter_data[level, p, fidx + 1, elem] + pdata[data_idx2] = + p∂ξ∂x[idx12] * pperimeter_data[level, p, fidx, elem] + + p∂ξ∂x[idx22] * pperimeter_data[level, p, fidx + 1, elem] + end + end + end + end + return nothing +end + +function dss_load_perimeter_data!( + ::ClimaComms.AbstractCPUDevice, + dss_buffer::DSSBuffer, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + perimeter::Perimeter2D{Nq}, +) where {Nq} + pperimeter_data = parent(dss_buffer.perimeter_data) + pdata = parent(data) + (nlevels, _, nfid, nelems) = size(pperimeter_data) + sizet = (nlevels, Nq, Nq, nfid, nelems) + for elem in 1:nelems, (p, (ip, jp)) in enumerate(perimeter) + for fidx in 1:nfid, level in 1:nlevels + idx = _get_idx(sizet, (level, ip, jp, fidx, elem)) + pperimeter_data[level, p, fidx, elem] = pdata[idx] + end + end + return nothing +end + +function dss_unload_perimeter_data!( + ::ClimaComms.AbstractCPUDevice, + data::Union{DataLayouts.IJFH, DataLayouts.VIJFH}, + dss_buffer::DSSBuffer, + perimeter::Perimeter2D{Nq}, +) where {Nq} + pperimeter_data = parent(dss_buffer.perimeter_data) + pdata = parent(data) + (nlevels, _, nfid, nelems) = size(pperimeter_data) + sizet = (nlevels, Nq, Nq, nfid, nelems) + for elem in 1:nelems, (p, (ip, jp)) in enumerate(perimeter) + for fidx in 1:nfid, level in 1:nlevels + idx = _get_idx(sizet, (level, ip, jp, fidx, elem)) + pdata[idx] = pperimeter_data[level, p, fidx, elem] + end + end + return nothing +end + +""" + function dss_local!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + perimeter::AbstractPerimeter, + topology::AbstractTopology, + ) + +Performs DSS on local vertices and faces. + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function dss_local!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + perimeter::Perimeter2D, + topology::Topology2D, +) + dss_local_vertices!(perimeter_data, perimeter, topology) + dss_local_faces!(perimeter_data, perimeter, topology) + return nothing +end + +""" + dss_local_vertices!( + perimeter_data::DataLayouts.VIFH, + perimeter::Perimeter2D, + topology::Topology2D, + ) + +Apply dss to local vertices. +""" +function dss_local_vertices!( + perimeter_data::DataLayouts.VIFH, + perimeter::Perimeter2D, + topology::Topology2D, +) + Nv = size(perimeter_data, 4) + @inbounds for vertex in local_vertices(topology) + # for each level + for level in 1:Nv + # gather: compute sum over shared vertices + sum_data = mapreduce( + ⊞, + vertex; + init = RecursiveApply.rzero(eltype(slab(perimeter_data, 1, 1))), + ) do (lidx, vert) + ip = perimeter_vertex_node_index(vert) + perimeter_slab = slab(perimeter_data, level, lidx) + perimeter_slab[ip] + end + # scatter: assign sum to shared vertices + for (lidx, vert) in vertex + perimeter_slab = slab(perimeter_data, level, lidx) + ip = perimeter_vertex_node_index(vert) + perimeter_slab[ip] = sum_data + end + end + end + return nothing +end + +function dss_local_faces!( + perimeter_data::DataLayouts.VIFH, + perimeter::Perimeter2D, + topology::Topology2D, +) + (Np, _, _, Nv, _) = size(perimeter_data) + nfacedof = div(Np - 4, 4) + + @inbounds for (lidx1, face1, lidx2, face2, reversed) in + interior_faces(topology) + pr1 = perimeter_face_indices(face1, nfacedof, false) + pr2 = perimeter_face_indices(face2, nfacedof, reversed) + for level in 1:Nv + perimeter_slab1 = slab(perimeter_data, level, lidx1) + perimeter_slab2 = slab(perimeter_data, level, lidx2) + for (ip1, ip2) in zip(pr1, pr2) + val = perimeter_slab1[ip1] ⊞ perimeter_slab2[ip2] + perimeter_slab1[ip1] = val + perimeter_slab2[ip2] = val + end + end + end + return nothing +end +""" + function dss_local_ghost!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + perimeter::AbstractPerimeter, + topology::AbstractTopology, + ) + +Computes the "local" part of ghost vertex dss. (i.e. it computes the summation of all the shared local +vertices of a unique ghost vertex and stores the value in each of the local vertex locations in +`perimeter_data`) + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function dss_local_ghost!( + ::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + perimeter::AbstractPerimeter, + topology::AbstractTopology, +) + nghostvertices = length(topology.ghost_vertex_offset) - 1 + if nghostvertices > 0 + (Np, _, _, Nv, _) = size(perimeter_data) + @inbounds for vertex in ghost_vertices(topology) + for level in 1:Nv + # gather: compute sum over shared vertices + sum_data = mapreduce( + ⊞, + vertex; + init = RecursiveApply.rzero( + eltype(slab(perimeter_data, 1, 1)), + ), + ) do (isghost, idx, vert) + ip = perimeter_vertex_node_index(vert) + if !isghost + lidx = idx + perimeter_slab = slab(perimeter_data, level, lidx) + perimeter_slab[ip] + else + RecursiveApply.rmap(zero, slab(perimeter_data, 1, 1)[1]) + end + end + for (isghost, idx, vert) in vertex + if !isghost + ip = perimeter_vertex_node_index(vert) + lidx = idx + perimeter_slab = slab(perimeter_data, level, lidx) + perimeter_slab[ip] = sum_data + end + end + end + end + end + return nothing +end +""" + dss_ghost!( + device::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + perimeter::AbstractPerimeter, + topology::AbstractTopology, + ) + +Sets the value for all local vertices of each unique ghost vertex, in `perimeter_data`, to that of +the representative ghost vertex. + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function dss_ghost!( + device::ClimaComms.AbstractCPUDevice, + perimeter_data::DataLayouts.VIFH, + perimeter::AbstractPerimeter, + topology::AbstractTopology, +) + nghostvertices = length(topology.ghost_vertex_offset) - 1 + if nghostvertices > 0 + nlevels = size(perimeter_data, 4) + (; repr_ghost_vertex) = topology + @inbounds for (i, vertex) in enumerate(ghost_vertices(topology)) + idxresult, lvertresult = repr_ghost_vertex[i] + ipresult = perimeter_vertex_node_index(lvertresult) + for level in 1:nlevels + result_slab = slab(perimeter_data, level, idxresult) + result = result_slab[ipresult] + for (isghost, idx, vert) in vertex + if !isghost + ip = perimeter_vertex_node_index(vert) + lidx = idx + perimeter_slab = slab(perimeter_data, level, lidx) + perimeter_slab[ip] = result + end + end + end + end + end + return nothing +end + +""" + fill_send_buffer!(::ClimaComms.AbstractCPUDevice, dss_buffer::DSSBuffer) + +Loads the send buffer from `perimeter_data`. For unique ghost vertices, only data from the +representative vertices which store result of "ghost local" DSS are loaded. + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function fill_send_buffer!( + ::ClimaComms.AbstractCPUDevice, + dss_buffer::DSSBuffer, +) + (; perimeter_data, send_buf_idx, send_data) = dss_buffer + (Np, _, _, Nv, nelems) = size(perimeter_data) + Nf = cld(length(parent(perimeter_data)), (Nv * Np * nelems)) + pdata = parent(perimeter_data) + nsend = size(send_buf_idx, 1) + ctr = 1 + @inbounds for i in 1:nsend + lidx = send_buf_idx[i, 1] + ip = send_buf_idx[i, 2] + for f in 1:Nf, v in 1:Nv + send_data[ctr] = pdata[v, ip, f, lidx] + ctr += 1 + end + end + return nothing +end +""" + load_from_recv_buffer!(::ClimaComms.AbstractCPUDevice, dss_buffer::DSSBuffer) + +Adds data from the recv buffer to the corresponding location in `perimeter_data`. +For ghost vertices, this data is added only to the representative vertices. The values are +then scattered to other local vertices corresponding to each unique ghost vertex in `dss_local_ghost`. + +Part of [`ClimaCore.Spaces.weighted_dss!`](@ref). +""" +function load_from_recv_buffer!( + ::ClimaComms.AbstractCPUDevice, + dss_buffer::DSSBuffer, +) + (; perimeter_data, recv_buf_idx, recv_data) = dss_buffer + (Np, _, _, Nv, nelems) = size(perimeter_data) + Nf = cld(length(parent(perimeter_data)), (Nv * Np * nelems)) + pdata = parent(perimeter_data) + nrecv = size(recv_buf_idx, 1) + ctr = 1 + @inbounds for i in 1:nrecv + lidx = recv_buf_idx[i, 1] + ip = recv_buf_idx[i, 2] + for f in 1:Nf, v in 1:Nv + pdata[v, ip, f, lidx] += recv_data[ctr] + ctr += 1 + end + end + return nothing +end + +""" + dss!(data, topology) + +Computed unweighted/pure DSS of `data`. +""" +function dss!( + data::Union{DataLayouts.IJFH{S, Nij}, DataLayouts.VIJFH{S, Nij}}, + topology::Topology2D, +) where {S, Nij} + length(parent(data)) == 0 && return nothing + device = ClimaComms.device(topology) + perimeter = Perimeter2D(Nij) + # create dss buffer + dss_buffer = create_dss_buffer(data, topology) + # load perimeter data from data + dss_load_perimeter_data!(device, dss_buffer, data, perimeter) + # compute local dss for ghost dof + dss_local_ghost!(device, dss_buffer.perimeter_data, perimeter, topology) + # load send buffer + fill_send_buffer!(device, dss_buffer) + # initiate communication + ClimaComms.start(dss_buffer.graph_context) + # compute local dss + dss_local!(device, dss_buffer.perimeter_data, perimeter, topology) + # finish communication + ClimaComms.finish(dss_buffer.graph_context) + # load from receive buffer + load_from_recv_buffer!(device, dss_buffer) + # finish dss computation for ghost dof + dss_ghost!(device, dss_buffer.perimeter_data, perimeter, topology) + # load perimeter_data into data + dss_unload_perimeter_data!(device, data, dss_buffer, perimeter) + return nothing +end + +function dss_1d!( + htopology::AbstractTopology, + data, + local_geometry_data = nothing, + dss_weights = nothing, +) + Nq = size(data, 1) + Nv = size(data, 4) + idx1 = CartesianIndex(1, 1, 1, 1, 1) + idx2 = CartesianIndex(Nq, 1, 1, 1, 1) + @inbounds for (elem1, face1, elem2, face2, reversed) in + interior_faces(htopology) + for level in 1:Nv + @assert face1 == 1 && face2 == 2 && !reversed + local_geometry_slab1 = slab(local_geometry_data, level, elem1) + weight_slab1 = slab(dss_weights, level, elem1) + data_slab1 = slab(data, level, elem1) + + local_geometry_slab2 = slab(local_geometry_data, level, elem2) + weight_slab2 = slab(dss_weights, level, elem2) + data_slab2 = slab(data, level, elem2) + val = + dss_transform( + data_slab1, + local_geometry_slab1, + weight_slab1, + idx1, + ) ⊞ dss_transform( + data_slab2, + local_geometry_slab2, + weight_slab2, + idx2, + ) + + data_slab1[idx1] = dss_untransform( + eltype(data_slab1), + val, + local_geometry_slab1, + idx1, + ) + data_slab2[idx2] = dss_untransform( + eltype(data_slab2), + val, + local_geometry_slab2, + idx2, + ) + end + end + return data +end + +include("dss_cuda.jl") diff --git a/src/Spaces/dss_cuda.jl b/src/Topologies/dss_cuda.jl similarity index 94% rename from src/Spaces/dss_cuda.jl rename to src/Topologies/dss_cuda.jl index 058186c344..8a5d19cbe9 100644 --- a/src/Spaces/dss_cuda.jl +++ b/src/Topologies/dss_cuda.jl @@ -21,7 +21,7 @@ function dss_load_perimeter_data!( (nlevels, nperimeter, nfid, nelems) = size(pperimeter_data) nitems = nlevels * nperimeter * nfid * nelems nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) dss_load_perimeter_data_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) dss_load_perimeter_data_kernel!( pperimeter_data, pdata, perimeter, @@ -58,7 +58,7 @@ function dss_unload_perimeter_data!( (nlevels, nperimeter, nfid, nelems) = size(pperimeter_data) nitems = nlevels * nperimeter * nfid * nelems nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) dss_unload_perimeter_data_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) dss_unload_perimeter_data_kernel!( pdata, pperimeter_data, perimeter, @@ -98,7 +98,7 @@ function dss_local!( nitems = nlevels * nfid * (nlocalfaces + nlocalvertices) nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) dss_local_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) dss_local_kernel!( pperimeter_data, topology.local_vertices, topology.local_vertex_offset, @@ -128,12 +128,12 @@ function dss_local_kernel!( local_vertex_offset[vertexid], local_vertex_offset[vertexid + 1] for idx in st:(en - 1) (lidx, vert) = local_vertices[idx] - ip = Topologies.perimeter_vertex_node_index(vert) + ip = perimeter_vertex_node_index(vert) sum_data += pperimeter_data[level, ip, fidx, lidx] end for idx in st:(en - 1) (lidx, vert) = local_vertices[idx] - ip = Topologies.perimeter_vertex_node_index(vert) + ip = perimeter_vertex_node_index(vert) pperimeter_data[level, ip, fidx, lidx] = sum_data end elseif gidx ≤ nlevels * nfidx * (nlocalvertices + nlocalfaces) # interior faces @@ -184,7 +184,7 @@ function dss_transform!( (nlevels, nperimeter, _, _) = size(pperimeter_data) nitems = nlevels * nperimeter * nlocalelems nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) dss_transform_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) dss_transform_kernel!( pperimeter_data, pdata, p∂ξ∂x, @@ -290,7 +290,7 @@ function dss_untransform!( (nlevels, nperimeter, _, _) = size(pperimeter_data) nitems = nlevels * nperimeter * nlocalelems nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) dss_untransform_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) dss_untransform_kernel!( pperimeter_data, pdata, p∂ξ∂x, @@ -376,7 +376,7 @@ function dss_local_ghost!( max_threads = 256 nitems = nlevels * nfid * nghostvertices nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) dss_local_ghost_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) dss_local_ghost_kernel!( pperimeter_data, topology.ghost_vertices, topology.ghost_vertex_offset, @@ -404,14 +404,14 @@ function dss_local_ghost_kernel!( for idx in st:(en - 1) isghost, lidx, vert = ghost_vertices[idx] if !isghost - ip = Topologies.perimeter_vertex_node_index(vert) + ip = perimeter_vertex_node_index(vert) sum_data += pperimeter_data[level, ip, fidx, lidx] end end for idx in st:(en - 1) isghost, lidx, vert = ghost_vertices[idx] if !isghost - ip = Topologies.perimeter_vertex_node_index(vert) + ip = perimeter_vertex_node_index(vert) pperimeter_data[level, ip, fidx, lidx] = sum_data end end @@ -427,7 +427,7 @@ function fill_send_buffer!(::ClimaComms.CUDADevice, dss_buffer::DSSBuffer) if nsend > 0 nitems = nsend * nlevels * nfid nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) fill_send_buffer_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) fill_send_buffer_kernel!( send_data, send_buf_idx, pperimeter_data, @@ -467,7 +467,7 @@ function load_from_recv_buffer!(::ClimaComms.CUDADevice, dss_buffer::DSSBuffer) if nrecv > 0 nitems = nrecv * nlevels * nfid nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) load_from_recv_buffer_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) load_from_recv_buffer_kernel!( pperimeter_data, recv_data, recv_buf_idx, @@ -511,7 +511,7 @@ function dss_ghost!( nlevels, _, nfidx, _ = size(pperimeter_data) nitems = nlevels * nfidx * nghostvertices nthreads, nblocks = _configure_threadblock(nitems) - @cuda threads = (nthreads) blocks = (nblocks) dss_ghost_kernel!( + CUDA.@cuda threads = (nthreads) blocks = (nblocks) dss_ghost_kernel!( pperimeter_data, topology.ghost_vertices, topology.ghost_vertex_offset, @@ -537,14 +537,14 @@ function dss_ghost_kernel!( (level, fidx, ghostvertexidx) = _get_idx((nlevels, nfidx, nghostvertices), gidx) idxresult, lvertresult = repr_ghost_vertex[ghostvertexidx] - ipresult = Topologies.perimeter_vertex_node_index(lvertresult) + ipresult = perimeter_vertex_node_index(lvertresult) result = pperimeter_data[level, ipresult, fidx, idxresult] st, en = ghost_vertex_offset[ghostvertexidx], ghost_vertex_offset[ghostvertexidx + 1] for vertexidx in st:(en - 1) isghost, eidx, lvert = ghost_vertices[vertexidx] if !isghost - ip = Topologies.perimeter_vertex_node_index(lvert) + ip = perimeter_vertex_node_index(lvert) pperimeter_data[level, ip, fidx, eidx] = result end end diff --git a/src/Spaces/dss_transform.jl b/src/Topologies/dss_transform.jl similarity index 92% rename from src/Spaces/dss_transform.jl rename to src/Topologies/dss_transform.jl index cebd52266a..466a6f641e 100644 --- a/src/Spaces/dss_transform.jl +++ b/src/Topologies/dss_transform.jl @@ -10,7 +10,7 @@ Transformations only apply to vector quantities. - `local_geometry[I...]` is the relevant `LocalGeometry` object. If it is `nothing`, then no transformation is performed - `weight[I...]` is the relevant DSS weights. If `weight` is `nothing`, then the result is simply summation. -See [`Spaces.weighted_dss!`](@ref). +See [`ClimaCore.Spaces.weighted_dss!`](@ref). """ Base.@propagate_inbounds dss_transform(arg, local_geometry, weight, i, j) = dss_transform(arg[i, j], local_geometry[i, j], weight[i, j]) @@ -146,7 +146,7 @@ end Transform `targ[I...]` back to a value of type `T` after performing direct stiffness summation (DSS). -See [`Spaces.weighted_dss!`](@ref). +See [`ClimaCore.Spaces.weighted_dss!`](@ref). """ Base.@propagate_inbounds dss_untransform( ::Type{T}, @@ -230,18 +230,6 @@ end Geometry.transform(ax, targ, local_geometry) end -function weighted_dss! end -function weighted_dss_start! end -function weighted_dss_internal! end -function weighted_dss_ghost! end - -# for backward compatibility -function weighted_dss2! end -function weighted_dss_start2! end -function weighted_dss_internal2! end -function weighted_dss_ghost2! end -function dss2! end - # helper functions for DSS2 function _get_idx(sizet::NTuple{5, Int}, loc::NTuple{5, Int}) (n1, n2, n3, n4, n5) = sizet @@ -286,7 +274,10 @@ function _get_idx_metric(sizet::NTuple{5, Int}, loc::NTuple{4, Int}) return nothing end -function _representative_slab(data, ::Type{DA}) where {DA} +function _representative_slab( + data::Union{DataLayouts.AbstractData, Nothing}, + ::Type{DA}, +) where {DA} rebuild_flag = DA isa Array ? false : true if isnothing(data) return nothing @@ -297,17 +288,23 @@ function _representative_slab(data, ::Type{DA}) where {DA} end end -_transformed_type(data, local_geometry, local_weights, ::Type{DA}) where {DA} = - typeof( - dss_transform( - _representative_slab(data, DA), - _representative_slab(local_geometry, DA), - _representative_slab(local_weights, DA), - 1, - 1, - ), - ) +_transformed_type( + data::DataLayouts.AbstractData, + local_geometry::Union{DataLayouts.AbstractData, Nothing}, + local_weights::Union{DataLayouts.AbstractData, Nothing}, + ::Type{DA}, +) where {DA} = typeof( + dss_transform( + _representative_slab(data, DA), + _representative_slab(local_geometry, DA), + _representative_slab(local_weights, DA), + 1, + 1, + ), +) +# currently only used in limiters (but not actually functional) +# see https://github.com/CliMA/ClimaCore.jl/issues/1511 struct GhostBuffer{G, D} graph_context::G send_data::D diff --git a/src/Topologies/interval.jl b/src/Topologies/interval.jl index bd571e8236..90beb4bc61 100644 --- a/src/Topologies/interval.jl +++ b/src/Topologies/interval.jl @@ -16,6 +16,7 @@ struct IntervalTopology{ boundaries::B end +## gpu struct DeviceIntervalTopology{B} <: AbstractIntervalTopology boundaries::B end diff --git a/src/Topologies/topology2d.jl b/src/Topologies/topology2d.jl index 3f858abce6..270c199283 100644 --- a/src/Topologies/topology2d.jl +++ b/src/Topologies/topology2d.jl @@ -23,7 +23,7 @@ Internally, we can refer to elements in several different ways: - `ridx`: "receive index": an index into the receive buffer of a ghost element. - `recv_elem_gidx[ridx] == gidx` """ -struct Topology2D{ +mutable struct Topology2D{ C <: ClimaComms.AbstractCommsContext, M <: Meshes.AbstractMesh{2}, EO, @@ -194,7 +194,7 @@ function simple_partition(nelems::Int, npart::Int) return partition, ranges end -function Topology2D( +@memoize WeakValueDict function Topology2D( context::ClimaComms.AbstractCommsContext, mesh::Meshes.AbstractMesh{2}, elemorder = Meshes.elements(mesh), @@ -718,3 +718,57 @@ boundary_tag(topology::Topology2D, boundary_name::Symbol) = findfirst(==(boundary_name), boundary_names(topology)) boundary_faces(topology::Topology2D, boundary) = topology.boundaries[boundary] + + + + +abstract type AbstractPerimeter end + +struct Perimeter2D{Nq} <: AbstractPerimeter end + +""" + Perimeter2D(Nq) + +Construct a perimeter iterator for a 2D spectral element with `Nq` nodes per +dimension (i.e. polynomial degree `Nq-1`). +""" +Perimeter2D(Nq) = Perimeter2D{Nq}() + +Adapt.adapt_structure(to, x::Perimeter2D) = x + +Base.IteratorEltype(::Type{Perimeter2D{Nq}}) where {Nq} = Base.HasEltype() +Base.eltype(::Type{Perimeter2D{Nq}}) where {Nq} = Tuple{Int, Int} + +Base.IteratorSize(::Type{Perimeter2D{Nq}}) where {Nq} = Base.HasLength() +Base.length(::Perimeter2D{Nq}) where {Nq} = 4 + (Nq - 2) * 4 + +function Base.iterate(perimeter::Perimeter2D{Nq}, loc = 1) where {Nq} + if loc <= 4 + return (vertex_node_index(loc, Nq), loc + 1) + elseif loc ≤ length(perimeter) + f = cld(loc - 4, Nq - 2) + n = mod(loc - 4, Nq - 2) == 0 ? (Nq - 2) : mod(loc - 4, Nq - 2) + return (face_node_index(f, Nq, 1 + n), loc + 1) + else + return nothing + end +end + +function Base.getindex(perimeter::Perimeter2D{Nq}, loc = 1) where {Nq} + if loc < 1 || loc > length(perimeter) + return (-1, -1) + elseif loc <= 4 + return vertex_node_index(loc, Nq) + else + f = cld(loc - 4, Nq - 2) + n = mod(loc - 4, Nq - 2) == 0 ? (Nq - 2) : mod(loc - 4, Nq - 2) + return face_node_index(f, Nq, 1 + n) + end +end + + +## aliases +const RectilinearTopology2D = + Topology2D{<:ClimaComms.AbstractCommsContext, <:Meshes.RectilinearMesh} +const CubedSphereTopology2D = + Topology2D{<:ClimaComms.AbstractCommsContext, <:Meshes.AbstractCubedSphere} diff --git a/test/Fields/field.jl b/test/Fields/field.jl index 8103456ede..c1218873aa 100644 --- a/test/Fields/field.jl +++ b/test/Fields/field.jl @@ -280,7 +280,7 @@ function call_getproperty(fv) end @testset "FieldVector getindex" begin cspace = TU.CenterExtrudedFiniteDifferenceSpace(Float32) - fspace = Spaces.ExtrudedFiniteDifferenceSpace{Spaces.CellFace}(cspace) + fspace = Spaces.FaceExtrudedFiniteDifferenceSpace(cspace) c = fill((a = Float32(1), b = Float32(2)), cspace) f = fill((x = Float32(1), y = Float32(2)), fspace) fv = Fields.FieldVector(; c, f) @@ -551,22 +551,10 @@ end nothing end -@testset "Broadcasting same spaces different instances" begin +@testset "Memoization of spaces" begin space1 = spectral_space_2D() space2 = spectral_space_2D() - field1 = ones(space1) - field2 = 2 .* ones(space2) - @test Fields.is_diagonalized_spaces(typeof(space1), typeof(space2)) - @test_throws ErrorException( - "Broacasted spaces are the same ClimaCore.Spaces type but not the same instance", - ) field1 .= field2 - - # turn warning on - Fields.allow_mismatched_diagonalized_spaces() = true - @test_warn "Broacasted spaces are the same ClimaCore.Spaces type but not the same instance" field1 .= - field2 - @test parent(field1) == parent(field2) - Fields.allow_mismatched_diagonalized_spaces() = false + @test space1 === space2 end struct InferenceFoo{FT} diff --git a/test/Operators/finitedifference/column.jl b/test/Operators/finitedifference/column.jl index 53ff7fef0c..34411617f9 100644 --- a/test/Operators/finitedifference/column.jl +++ b/test/Operators/finitedifference/column.jl @@ -291,7 +291,7 @@ convergence_rate(err, Δh) = cent_field = operator.(face_field) wcent_field = woperator.(face_J, face_field) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] err[k] = norm(cent_field .- cent_field_exact) werr[k] = norm(wcent_field .- cent_field_exact) end @@ -345,7 +345,7 @@ end face_field = operator.(cent_field) wface_field = woperator.(cent_J, cent_field) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] err[k] = norm(face_field .- face_field_exact) werr[k] = norm(wface_field .- face_field_exact) end @@ -395,7 +395,7 @@ end face_field .= operator.(cent_field) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] err[k] = norm(face_field .- face_field_exact) end conv = convergence_rate(err, Δh) @@ -443,7 +443,7 @@ end cent_field .= operator.(face_field) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] err[k] = norm(cent_field .- cent_field_exact) end conv = convergence_rate(err, Δh) @@ -543,7 +543,7 @@ end curlsinᶠ = curlᶠ.(Geometry.Covariant1Vector.(sin.(centers))) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Errors err_grad_sin_c[k] = norm(gradsinᶜ .- Geometry.WVector.(cos.(centers))) err_div_sin_c[k] = norm(divsinᶜ .- cos.(centers)) @@ -663,7 +663,7 @@ end divf2c = Operators.DivergenceF2C() adv_wc = divf2c.(third_order_fluxsinᶠ) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -714,7 +714,7 @@ end divf2c = Operators.DivergenceF2C() adv_wc = divf2c.(third_order_fluxsinᶠ) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Error err_adv_wc[k] = @@ -776,7 +776,7 @@ end adv_wc = divf2c.(third_order_fluxᶠ.(w, c)) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -831,7 +831,7 @@ end ) adv_wc = divf2c.(third_order_fluxᶠ.(w, c)) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Errors err_adv_wc[k] = norm(adv_wc .- ((cos.(centers)) .^ 2 .- (sin.(centers)) .^ 2)) @@ -886,7 +886,7 @@ end @. divf2c(C * (third_order_fluxsinᶠ - first_order_fluxsinᶠ)) adv_wc = @. divf2c.(first_order_fluxsinᶠ) + corrected_antidiff_flux - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -957,7 +957,7 @@ end adv_wc = @. divf2c.(first_order_fluxᶠ(w, c)) + corrected_antidiff_flux - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Error err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -1020,7 +1020,7 @@ end adv_wc = @. divf2c.(first_order_fluxᶠ(w, c)) + corrected_antidiff_flux - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] # Errors err_adv_wc[k] = norm(adv_wc .- cos.(centers)) @@ -1127,7 +1127,7 @@ end # Call the advection operator adv = advection(c, f, cs) - Δh[k] = cs.face_local_geometry.J[1] + Δh[k] = Spaces.local_geometry_data(fs).J[1] err[k] = norm(adv .- cos.(Fields.coordinate_field(cs).z)) end # AdvectionC2C convergence rate diff --git a/test/Operators/finitedifference/column_benchmark_utils.jl b/test/Operators/finitedifference/column_benchmark_utils.jl index 623b8e68d2..1254528a65 100644 --- a/test/Operators/finitedifference/column_benchmark_utils.jl +++ b/test/Operators/finitedifference/column_benchmark_utils.jl @@ -299,7 +299,7 @@ function benchmark_operators(z_elems, ::Type{FT}) where {FT} benchmark_operators_base(trials, t_ave, cfield, ffield) cspace = TU.CenterExtrudedFiniteDifferenceSpace(FT; zelem=z_elems) - fspace = Spaces.ExtrudedFiniteDifferenceSpace{Spaces.CellFace}(cspace) + fspace = Spaces.FaceExtrudedFiniteDifferenceSpace(cspace) cfield = fill(field_vars(FT), cspace) ffield = fill(field_vars(FT), fspace) benchmark_operators_base(trials, t_ave, cfield, ffield) diff --git a/test/Operators/hybrid/3d.jl b/test/Operators/hybrid/3d.jl index d7afe2bcb7..232c8b63b9 100644 --- a/test/Operators/hybrid/3d.jl +++ b/test/Operators/hybrid/3d.jl @@ -103,13 +103,13 @@ end @test parent(Fields.field_values(level(coord.x, half))) == parent( Fields.field_values( - Fields.coordinate_field(hv_face_space.horizontal_space).x, + Fields.coordinate_field(Spaces.horizontal_space(hv_face_space)).x, ), ) @test parent(Fields.field_values(level(coord.z, half))) == parent( Fields.field_values( - Fields.coordinate_field(hv_face_space.horizontal_space).x, + Fields.coordinate_field(Spaces.horizontal_space(hv_face_space)).x, ), ) .* 0 end @@ -208,7 +208,7 @@ end struct ZeroFieldFlux <: BCtag end function bc_divF2C_bottom!(::ZeroFieldFlux, dY, Y, p, t) - space = axes(Y.h).horizontal_space + space = Spaces.horizontal_space(axes(Y.h)) FT = Spaces.undertype(space) zeroflux = Fields.zeros(FT, space) return Operators.SetValue( diff --git a/test/Operators/remapping.jl b/test/Operators/remapping.jl index 880040565f..ac7b29ee61 100644 --- a/test/Operators/remapping.jl +++ b/test/Operators/remapping.jl @@ -174,14 +174,14 @@ end @test monotone(R) @testset "Scalar Remap Operator Application" begin - n = length(source.local_geometry) + n = length(Spaces.local_geometry_data(source)) source_field = Fields.Field(IJFH{FT, 1}(ones(1, 1, n, 1)), source) # test consistent remap target_field = remap(R, source_field) @test vec(parent(target_field)) ≈ - ones(length(target.local_geometry)) + ones(length(Spaces.local_geometry_data(target))) # test simple remap vec(parent(source_field)) .= [1.0; 2.0; 3.0; 4.0] @@ -245,14 +245,14 @@ end @test monotone(R) @testset "Scalar Remap Operator Application" begin - n = length(source.local_geometry) + n = length(Spaces.local_geometry_data(source)) source_field = Fields.Field(IJFH{FT, 1}(ones(1, 1, n, 1)), source) # test consistent remap target_field = remap(R, source_field) @test vec(parent(target_field)) ≈ - ones(length(target.local_geometry)) + ones(length(Spaces.local_geometry_data(target))) # test simple remap vec(parent(source_field)) .= [1.0; 2.0; 3.0; 4.0] @@ -299,14 +299,14 @@ end @test monotone(R) @testset "Scalar Remap Operator Application" begin - n = length(source.local_geometry) + n = length(Spaces.local_geometry_data(source)) source_field = Fields.Field(IJFH{FT, 1}(ones(1, 1, n, 1)), source) # test consistent remap target_field = remap(R, source_field) @test vec(parent(target_field)) ≈ - ones(length(target.local_geometry)) + ones(length(Spaces.local_geometry_data(target))) # test simple remap vec(parent(source_field)) .= [1.0; 2.0; 3.0; 4.0] @@ -374,14 +374,14 @@ end @test monotone(R) @testset "Scalar Remap Operator Application" begin - n = length(source.local_geometry) + n = length(Spaces.local_geometry_data(source)) source_field = Fields.Field(IJFH{FT, 1}(ones(1, 1, n, 1)), source) # test consistent remap target_field = remap(R, source_field) @test vec(parent(target_field)) ≈ - ones(length(target.local_geometry)) + ones(length(Spaces.local_geometry_data(target))) # test simple remap vec(parent(source_field)) .= [1.0; 2.0; 3.0; 4.0] diff --git a/test/Operators/spectralelement/rectilinear.jl b/test/Operators/spectralelement/rectilinear.jl index c90ebbc610..935412c208 100644 --- a/test/Operators/spectralelement/rectilinear.jl +++ b/test/Operators/spectralelement/rectilinear.jl @@ -47,20 +47,20 @@ ts_test_setup = (ts_topology, ts_space, ts_coords) interpolated_field = I.(f) Spaces.weighted_dss!(interpolated_field) - @test axes(interpolated_field).quadrature_style == Iquad - @test axes(interpolated_field).topology == topology + @test Spaces.quadrature_style(axes(interpolated_field)) == Iquad + @test Spaces.topology(axes(interpolated_field)) == topology restrict_field = R.(f) Spaces.weighted_dss!(restrict_field) - @test axes(restrict_field).quadrature_style == quad - @test axes(restrict_field).topology == topology + @test Spaces.quadrature_style(axes(restrict_field)) == quad + @test Spaces.topology(axes(restrict_field)) == topology interp_restrict_field = R.(I.(f)) Spaces.weighted_dss!(interp_restrict_field) - @test axes(interp_restrict_field).quadrature_style == quad - @test axes(interp_restrict_field).topology == topology + @test Spaces.quadrature_style(axes(interp_restrict_field)) == quad + @test Spaces.topology(axes(interp_restrict_field)) == topology @test norm(interp_restrict_field .- f) ≤ 3.0e-4 end diff --git a/test/Spaces/ddss1.jl b/test/Spaces/ddss1.jl index 908fe35d3e..05b63e678d 100644 --- a/test/Spaces/ddss1.jl +++ b/test/Spaces/ddss1.jl @@ -66,10 +66,14 @@ init_state_vector(local_geometry, p) = Geometry.Covariant12Vector(1.0, -1.0) @test Topologies.nlocalelems(Spaces.topology(space)) == 4 - @test Topologies.local_neighboring_elements(space.topology, 1) == [2, 4] - @test Topologies.local_neighboring_elements(space.topology, 2) == [1, 3] - @test Topologies.local_neighboring_elements(space.topology, 3) == [2, 4] - @test Topologies.local_neighboring_elements(space.topology, 4) == [1, 3] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 1) == + [2, 4] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 2) == + [1, 3] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 3) == + [2, 4] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 4) == + [1, 3] y0 = init_state_scalar.(Fields.local_geometry_field(space), Ref(nothing)) nel = Topologies.nlocalelems(Spaces.topology(space)) diff --git a/test/Spaces/distributed/ddss2.jl b/test/Spaces/distributed/ddss2.jl index f2ce3d7dbe..594b7c5f5a 100644 --- a/test/Spaces/distributed/ddss2.jl +++ b/test/Spaces/distributed/ddss2.jl @@ -17,11 +17,15 @@ include("ddss_setup.jl") @test Topologies.nlocalelems(Spaces.topology(space)) == 2 - @test Topologies.local_neighboring_elements(space.topology, 1) == [2] - @test Topologies.local_neighboring_elements(space.topology, 2) == [1] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 1) == + [2] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 2) == + [1] - @test Topologies.ghost_neighboring_elements(space.topology, 1) == [2] - @test Topologies.ghost_neighboring_elements(space.topology, 2) == [1] + @test Topologies.ghost_neighboring_elements(Spaces.topology(space), 1) == + [2] + @test Topologies.ghost_neighboring_elements(Spaces.topology(space), 2) == + [1] init_state(local_geometry, p) = (ρ = 1.0) y0 = init_state.(Fields.local_geometry_field(space), Ref(nothing)) diff --git a/test/Spaces/distributed/ddss3.jl b/test/Spaces/distributed/ddss3.jl index acdbf994b9..9981a03ab6 100644 --- a/test/Spaces/distributed/ddss3.jl +++ b/test/Spaces/distributed/ddss3.jl @@ -39,22 +39,31 @@ partition numbers @test Topologies.nlocalelems(Spaces.topology(space)) == (pid == 1 ? 6 : 5) if pid == 1 # gidx 1 - @test Topologies.local_neighboring_elements(space.topology, 1) == - [2, 5, 6] - @test Topologies.ghost_neighboring_elements(space.topology, 1) == [] + @test Topologies.local_neighboring_elements( + Spaces.topology(space), + 1, + ) == [2, 5, 6] + @test Topologies.ghost_neighboring_elements( + Spaces.topology(space), + 1, + ) == [] # gidx 6 - @test Topologies.local_neighboring_elements(space.topology, 6) == - [1, 2, 3, 5] - @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements( - space.topology, + @test Topologies.local_neighboring_elements( + Spaces.topology(space), + 6, + ) == [1, 2, 3, 5] + @test Spaces.topology(space).recv_elem_gidx[Topologies.ghost_neighboring_elements( + Spaces.topology(space), 6, )] == [7, 9, 10, 11] elseif pid == 2 # gidx 7 - @test Topologies.local_neighboring_elements(space.topology, 1) == - [2, 4, 5] - @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements( - space.topology, + @test Topologies.local_neighboring_elements( + Spaces.topology(space), + 1, + ) == [2, 4, 5] + @test Spaces.topology(space).recv_elem_gidx[Topologies.ghost_neighboring_elements( + Spaces.topology(space), 1, )] == [2, 3, 4, 6, 12] end diff --git a/test/Spaces/distributed_cuda/ddss2.jl b/test/Spaces/distributed_cuda/ddss2.jl index 8b26bad751..c7e5748c10 100644 --- a/test/Spaces/distributed_cuda/ddss2.jl +++ b/test/Spaces/distributed_cuda/ddss2.jl @@ -59,11 +59,15 @@ pid, nprocs = ClimaComms.init(context) @test Topologies.nlocalelems(Spaces.topology(space)) == 2 - @test Topologies.local_neighboring_elements(space.topology, 1) == [2] - @test Topologies.local_neighboring_elements(space.topology, 2) == [1] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 1) == + [2] + @test Topologies.local_neighboring_elements(Spaces.topology(space), 2) == + [1] - @test Topologies.ghost_neighboring_elements(space.topology, 1) == [2] - @test Topologies.ghost_neighboring_elements(space.topology, 2) == [1] + @test Topologies.ghost_neighboring_elements(Spaces.topology(space), 1) == + [2] + @test Topologies.ghost_neighboring_elements(Spaces.topology(space), 2) == + [1] init_state(local_geometry, p) = (ρ = 1.0) y0 = init_state.(Fields.local_geometry_field(space), Ref(nothing)) diff --git a/test/Spaces/distributed_cuda/ddss3.jl b/test/Spaces/distributed_cuda/ddss3.jl index fcf39a16f1..1bca8d268b 100644 --- a/test/Spaces/distributed_cuda/ddss3.jl +++ b/test/Spaces/distributed_cuda/ddss3.jl @@ -80,22 +80,31 @@ partition numbers @test Topologies.nlocalelems(Spaces.topology(space)) == (pid == 1 ? 6 : 5) if pid == 1 # gidx 1 - @test Topologies.local_neighboring_elements(space.topology, 1) == - [2, 5, 6] - @test Topologies.ghost_neighboring_elements(space.topology, 1) == [] + @test Topologies.local_neighboring_elements( + Spaces.topology(space), + 1, + ) == [2, 5, 6] + @test Topologies.ghost_neighboring_elements( + Spaces.topology(space), + 1, + ) == [] # gidx 6 - @test Topologies.local_neighboring_elements(space.topology, 6) == - [1, 2, 3, 5] - @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements( - space.topology, + @test Topologies.local_neighboring_elements( + Spaces.topology(space), + 6, + ) == [1, 2, 3, 5] + @test Spaces.topology(space).recv_elem_gidx[Topologies.ghost_neighboring_elements( + Spaces.topology(space), 6, )] == [7, 9, 10, 11] elseif pid == 2 # gidx 7 - @test Topologies.local_neighboring_elements(space.topology, 1) == - [2, 4, 5] - @test space.topology.recv_elem_gidx[Topologies.ghost_neighboring_elements( - space.topology, + @test Topologies.local_neighboring_elements( + Spaces.topology(space), + 1, + ) == [2, 4, 5] + @test Spaces.topology(space).recv_elem_gidx[Topologies.ghost_neighboring_elements( + Spaces.topology(space), 1, )] == [2, 3, 4, 6, 12] end diff --git a/test/Spaces/extruded_cuda.jl b/test/Spaces/extruded_cuda.jl index 52c10fe8bc..1f35c94e2b 100644 --- a/test/Spaces/extruded_cuda.jl +++ b/test/Spaces/extruded_cuda.jl @@ -42,8 +42,6 @@ end # Test that all geometries match with CPU version: @test compare(cpuspace, gpuspace, :center_local_geometry) @test compare(cpuspace, gpuspace, :face_local_geometry) - @test compare(cpuspace, gpuspace, :center_ghost_geometry) - @test compare(cpuspace, gpuspace, :face_ghost_geometry) space = gpuspace Y = Fields.Field(typeof((; v = FT(0))), space) @@ -65,7 +63,6 @@ end # Test that all geometries match with CPU version: @test compare(cpuspace, gpuspace, :local_geometry) - @test compare(cpuspace, gpuspace, :ghost_geometry) space = gpuspace Y = Fields.Field(typeof((; v = FT(0))), space) diff --git a/test/Spaces/quadrature.jl b/test/Spaces/quadrature.jl index cd2951e8cd..82916a8fc2 100644 --- a/test/Spaces/quadrature.jl +++ b/test/Spaces/quadrature.jl @@ -1,6 +1,6 @@ using Test using LinearAlgebra, StaticArrays -import ClimaCore.Spaces: Quadratures +import ClimaCore.Quadratures f(x) = x^3 + 2x^2 + 3x + 4 diff --git a/test/Spaces/spaces.jl b/test/Spaces/spaces.jl index fe151942c9..cdc820026c 100644 --- a/test/Spaces/spaces.jl +++ b/test/Spaces/spaces.jl @@ -2,9 +2,8 @@ using Test using ClimaComms using StaticArrays, IntervalSets, LinearAlgebra -import ClimaCore: slab, Domains, Meshes, Topologies, Spaces, Fields, DataLayouts - -import ClimaCore.Geometry: Geometry +import ClimaCore: + slab, Domains, Meshes, Topologies, Spaces, Fields, DataLayouts, Geometry import ClimaCore.DataLayouts: IJFH, VF @testset "1d domain space" begin @@ -36,8 +35,8 @@ import ClimaCore.DataLayouts: IJFH, VF @test coord_slab[1] == Geometry.XPoint{FT}(-3) @test coord_slab[4] == Geometry.XPoint{FT}(5) - local_geometry_slab = slab(space.local_geometry, 1) - dss_weights_slab = slab(space.dss_weights, 1) + local_geometry_slab = slab(Spaces.local_geometry_data(space), 1) + dss_weights_slab = slab(space.grid.dss_weights, 1) for i in 1:4 @test Geometry.components(local_geometry_slab[i].∂x∂ξ) ≈ @@ -178,8 +177,8 @@ end @test coord_slab[1, 4] ≈ Geometry.XYPoint{FT}(-3.0, 8.0) @test coord_slab[4, 4] ≈ Geometry.XYPoint{FT}(5.0, 8.0) - local_geometry_slab = slab(space.local_geometry, 1) - dss_weights_slab = slab(space.local_dss_weights, 1) + local_geometry_slab = slab(space.grid.local_geometry, 1) + dss_weights_slab = slab(space.grid.local_dss_weights, 1) for i in 1:4, j in 1:4 @@ -197,10 +196,10 @@ end end end - @test length(space.boundary_surface_geometries) == 2 - @test keys(space.boundary_surface_geometries) == (:south, :north) - @test sum(parent(space.boundary_surface_geometries.north.sWJ)) ≈ 8 - @test parent(space.boundary_surface_geometries.north.normal)[1, :, 1] ≈ + @test length(space.grid.boundary_surface_geometries) == 2 + @test keys(space.grid.boundary_surface_geometries) == (:south, :north) + @test sum(parent(space.grid.boundary_surface_geometries.north.sWJ)) ≈ 8 + @test parent(space.grid.boundary_surface_geometries.north.normal)[1, :, 1] ≈ [0.0, 1.0] point_space = Spaces.column(space, 1, 1, 1) diff --git a/test/TestUtilities/TestUtilities.jl b/test/TestUtilities/TestUtilities.jl index f9957fcace..e80e12a650 100644 --- a/test/TestUtilities/TestUtilities.jl +++ b/test/TestUtilities/TestUtilities.jl @@ -135,7 +135,7 @@ function FaceExtrudedFiniteDifferenceSpace( context = ClimaComms.SingletonCommsContext(), ) where {FT} cspace = CenterExtrudedFiniteDifferenceSpace(FT; zelem, context, helem) - return Spaces.ExtrudedFiniteDifferenceSpace{Spaces.CellFace}(cspace) + return Spaces.FaceExtrudedFiniteDifferenceSpace(cspace) end function all_spaces( diff --git a/test/Topologies/rectangle.jl b/test/Topologies/rectangle.jl index 6918c925f2..6e1fb0037d 100644 --- a/test/Topologies/rectangle.jl +++ b/test/Topologies/rectangle.jl @@ -5,6 +5,24 @@ import ClimaCore.Geometry: Geometry using StaticArrays using IntervalSets +@testset "Perimeter2D iterator" begin + @test collect(Topologies.Perimeter2D(4)) == [ + (1, 1) + (4, 1) + (4, 4) + (1, 4) + (2, 1) + (3, 1) + (4, 2) + (4, 3) + (3, 4) + (2, 4) + (1, 3) + (1, 2) + ] +end + + function rectangular_grid( n1, n2, @@ -311,3 +329,11 @@ end @test getfield(c2, 2) == 0.0 end end + +@testset "memoization" begin + topology1 = rectangular_grid(3, 3, true, true) + topology2 = rectangular_grid(3, 3, true, true) + topology3 = rectangular_grid(3, 3, true, false) + @test topology1 === topology2 + @test topology1 !== topology3 +end diff --git a/test/runtests.jl b/test/runtests.jl index fc6b304141..187070a4d1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -11,6 +11,7 @@ Stacktrace: [1] pkgerror(msg::String) =# if !Sys.iswindows() + #= @safetestset "Recursive" begin @time include("RecursiveApply/recursive_apply.jl") end @safetestset "PlusHalf" begin @time include("Utilities/plushalf.jl") end @@ -34,7 +35,7 @@ if !Sys.iswindows() @safetestset "Cubedsphere surface topology" begin @time include("Topologies/cubedsphere_sfc.jl") end # now part of buildkite # @safetestset "Distributed topology" begin @time include("Topologies/distributed.jl") end - +=# @safetestset "Quadrature" begin @time include("Spaces/quadrature.jl") end @safetestset "Spaces" begin @time include("Spaces/spaces.jl") end #= @@ -59,7 +60,7 @@ if !Sys.iswindows() @safetestset "Spectral elem - sphere diffusion vec" begin @time include("Operators/spectralelement/sphere_diffusion_vec.jl") end @safetestset "Spectral elem - sphere hyperdiffusion" begin @time include("Operators/spectralelement/sphere_hyperdiffusion.jl") end @safetestset "Spectral elem - sphere hyperdiffusion vec" begin @time include("Operators/spectralelement/sphere_hyperdiffusion_vec.jl") end - + @safetestset "FD ops - column" begin @time include("Operators/finitedifference/column.jl") end @safetestset "FD ops - opt" begin @time include("Operators/finitedifference/opt.jl") end @safetestset "FD ops - wfact" begin @time include("Operators/finitedifference/wfact.jl") end @@ -68,7 +69,6 @@ if !Sys.iswindows() # now part of buildkite # @time include("Operators/finitedifference/implicit_stencils.jl") # @time include("Operators/finitedifference/opt_implicit_stencils.jl") - @safetestset "Hybrid - 2D" begin @time include("Operators/hybrid/2d.jl") end @safetestset "Hybrid - 3D" begin @time include("Operators/hybrid/3d.jl") end @safetestset "Hybrid - dss opt" begin @time include("Operators/hybrid/dss_opt.jl") end