Add DistributedTopology2D

Co-authored-by: sriharshakandala <[email protected]>
CliMA · Jan 21, 2022 · 444fb20 · 444fb20
1 parent cbd90b6
commit 444fb20
Show file tree

Hide file tree

Showing 20 changed files with 1,598 additions and 168 deletions.
diff --git a/Project.toml b/Project.toml
@@ -7,6 +7,8 @@ version = "0.6.0"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
+ClimaCommsMPI = "5f86816e-8b66-43b2-912e-75384f99de49"
 CubedSphere = "7445602f-e544-4518-8976-18f8e8ae6cdb"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -36,4 +38,4 @@ Rotations = "1"
 Static = "0.4, 0.5"
 StaticArrays = "1"
 UnPack = "1"
-julia = "1.6"
+julia = "1.6"
diff --git a/examples/Manifest.toml b/examples/Manifest.toml
diff --git a/examples/Project.toml b/examples/Project.toml
@@ -1,7 +1,10 @@
 [deps]
+ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
+ClimaCommsMPI = "5f86816e-8b66-43b2-912e-75384f99de49"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaCorePlots = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"

diff --git a/examples/bickleyjet/bickleyjet_cg_invariant_hypervisc.jl b/examples/bickleyjet/bickleyjet_cg_invariant_hypervisc.jl
@@ -3,6 +3,14 @@ using LinearAlgebra
 import ClimaCore:
     Domains, Fields, Geometry, Meshes, Operators, Spaces, Topologies
 
+#=
+# set up communications for distributed topology
+using ClimaComms
+using ClimaCommsMPI
+const Context = ClimaCommsMPI.MPICommsContext
+const pid, nprocs = ClimaComms.init(Context)
+=#
+
 using OrdinaryDiffEq: ODEProblem, solve, SSPRK33
 
 import Logging
@@ -36,8 +44,11 @@ n1, n2 = 16, 16
 Nq = 4
 mesh = Meshes.RectilinearMesh(domain, n1, n2)
 grid_topology = Topologies.Topology2D(mesh)
+#grid_topology = Topologies.DistributedTopology2D(mesh, Context)
 quad = Spaces.Quadratures.GLL{Nq}()
+#const comms_ctx = Spaces.setup_comms(Context, grid_topology, quad, 1, 4) # last 2 parameters are levels and maximum number of field variables
 space = Spaces.SpectralElementSpace2D(grid_topology, quad)
+#space = Spaces.SpectralElementSpace2D(grid_topology, quad, comms_ctx)
 
 function init_state(local_geometry, p)
     coord = local_geometry.coordinates
@@ -77,6 +88,8 @@ function total_energy(y, parameters)
 end
 
 function rhs!(dydt, y, _, t)
+    #global comms_ctx
+
     D₄ = parameters.D₄
     g = parameters.g
 
@@ -94,6 +107,7 @@ function rhs!(dydt, y, _, t)
     @. dydt.ρθ = wdiv(grad(y.ρθ))
 
     Spaces.weighted_dss!(dydt)
+    #Spaces.weighted_dss!(dydt, comms_ctx)
 
     @. dydt.u =
         -D₄ * (
@@ -109,11 +123,13 @@ function rhs!(dydt, y, _, t)
         dydt.u += -grad(g * y.ρ + norm(y.u)^2 / 2) + y.u × curl(y.u)
         dydt.ρθ += -wdiv(y.ρθ * y.u)
     end
+
     Spaces.weighted_dss!(dydt)
+    #Spaces.weighted_dss!(dydt, comms_ctx)
+
     return dydt
 end
 
-
 dydt = similar(y0)
 rhs!(dydt, y0, nothing, 0.0)
 

diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl
@@ -291,6 +291,10 @@ function Spaces.weighted_dss!(field::Field)
     Spaces.weighted_dss!(field_values(field), axes(field))
     return field
 end
+function Spaces.weighted_dss!(field::Field, comms_ctx)
+    Spaces.weighted_dss!(field_values(field), axes(field), comms_ctx)
+    return field
+end
 
 
 function level(field::CenterExtrudedFiniteDifferenceField, v::Int)

diff --git a/src/Spaces/Spaces.jl b/src/Spaces/Spaces.jl
@@ -14,6 +14,8 @@
 """
 module Spaces
 
+using ClimaComms
+
 import ..slab, ..column, ..level
 import ..Utilities: PlusHalf
 import ..DataLayouts, ..Geometry, ..Domains, ..Meshes, ..Topologies