Use AtmosGrids in type_getters

src/ClimaAtmos.jl

@@ -6,7 +6,6 @@ include(joinpath("parameters", "Parameters.jl"))
 import .Parameters as CAP
 
 include(joinpath("utils", "abbreviations.jl"))
-include(joinpath("utils", "common_spaces.jl"))
 include(joinpath("solver", "types.jl"))
 include(joinpath("solver", "cli_options.jl"))
 include(joinpath("utils", "utilities.jl"))
@@ -113,6 +112,9 @@ include(joinpath("callbacks", "callbacks.jl"))
 include(joinpath("diagnostics", "Diagnostics.jl"))
 import .Diagnostics as CAD
 
+include(joinpath("simulation", "AtmosGrids.jl"))
+import .Grid as CAG
+
 include(joinpath("solver", "model_getters.jl")) # high-level (using parsed_args) model getters
 include(joinpath("time_stepper", "time_stepper.jl"))
 include(joinpath("solver", "type_getters.jl"))
@@ -121,6 +123,4 @@ include(joinpath("solver", "solve.jl"))
 
 include(joinpath("parameters", "create_parameters.jl"))
 
-include(joinpath("simulation", "AtmosGrids.jl"))
-
 end # module

src/simulation/atmos_grids.jl

@@ -264,8 +264,8 @@ function VerticallyUniformBoxGrid(; nh_poly,
                                     y_max,
                                     z_elem,
                                     z_max,
-                                      topography = nothing,
-                                      topo_smoothing = false,
+                                    topography = nothing,
+                                    topo_smoothing = false,
                                     enable_bubble = false,
                                     comms_ctx = ClimaComms.context(),
                                     float_type = Float64)
@@ -461,15 +461,15 @@ const Sphere = VerticallyStretchedSphereGrid
 
 """
 function VerticallyUniformSphereGrid(; nh_poly,
-                                         h_elem,
-                                         radius,
-                                         z_elem,
-                                         z_max,
-                                         topography = nothing,
-                                         topo_smoothing = false,
-                                         enable_bubble = false,
-                                         comms_ctx = ClimaComms.context(),
-                                         float_type = Float64)
+                                       h_elem,
+                                       radius,
+                                       z_elem,
+                                       z_max,
+                                       topography = nothing,
+                                       topo_smoothing = false,
+                                       enable_bubble = false,
+                                       comms_ctx = ClimaComms.context(),
+                                       float_type = Float64)
 
 
 Construct an `SphereGrid` for a cubed sphere with columns with uniform resolution.
@@ -563,14 +563,14 @@ end
 
 """
 function VerticallyStretchedPlaneGrid(; nh_poly,
-                                      x_elem,
-                                      x_max,
-                                      z_elem,
-                                      z_max,
-                                      dz_bottom,
-                                      dz_top,
-                                      comms_ctx = ClimaComms.context(),
-                                      float_type = Float64)
+                                        x_elem,
+                                        x_max,
+                                        z_elem,
+                                        z_max,
+                                        dz_bottom,
+                                        dz_top,
+                                        comms_ctx = ClimaComms.context(),
+                                        float_type = Float64)
 
 Construct a `PlaneGrid` for a periodic linear domain with columns with
 non-uniform resolution (as prescribed by
@@ -638,14 +638,14 @@ const Plane = VerticallyStretchedPlaneGrid
 
 """
 function VerticallyUniformPlaneGrid(; nh_poly,
-                                    x_elem,
-                                    x_max,
-                                    y_elem,
-                                    y_max,
-                                    z_elem,
-                                    z_max,
-                                    comms_ctx = ClimaComms.context(),
-                                    float_type = Float64)
+                                      x_elem,
+                                      x_max,
+                                      y_elem,
+                                      y_max,
+                                      z_elem,
+                                      z_max,
+                                      comms_ctx = ClimaComms.context(),
+                                      float_type = Float64)
 
 Construct an `PlaneGrid` for a periodic box with columns with uniform resolution.
 

src/solver/type_getters.jl

@@ -104,157 +104,81 @@ end
 function get_spaces(parsed_args, params, comms_ctx)
 
     FT = eltype(params)
-    z_elem = Int(parsed_args["z_elem"])
-    z_max = FT(parsed_args["z_max"])
-    dz_bottom = FT(parsed_args["dz_bottom"])
-    dz_top = FT(parsed_args["dz_top"])
-    topography = parsed_args["topography"]
-    bubble = parsed_args["bubble"]
-
-    @assert topography in ("NoWarp", "DCMIP200", "Earth", "Agnesi", "Schar")
-    if topography == "DCMIP200"
-        warp_function = topography_dcmip200
-    elseif topography == "Agnesi"
-        warp_function = topography_agnesi
-    elseif topography == "Schar"
-        warp_function = topography_schar
-    elseif topography == "NoWarp"
-        warp_function = nothing
-    elseif topography == "Earth"
-        data_path = joinpath(topo_elev_dataset_path(), "ETOPO1_coarse.nc")
-        earth_spline = NCDataset(data_path) do data
-            zlevels = data["elevation"][:]
-            lon = data["longitude"][:]
-            lat = data["latitude"][:]
-            # Apply Smoothing
-            smooth_degree = Int(parsed_args["smoothing_order"])
-            esmth = imfilter(zlevels, Kernel.gaussian(smooth_degree))
-            linear_interpolation(
-                (lon, lat),
-                esmth,
-                extrapolation_bc = (Periodic(), Flat()),
-            )
-        end
-        @info "Generated interpolation stencil"
-        warp_function = generate_topography_warp(earth_spline)
-    end
-    @info "Topography" topography
-
-
-    h_elem = parsed_args["h_elem"]
-    radius = CAP.planet_radius(params)
-    center_space, face_space = if parsed_args["config"] == "sphere"
-        nh_poly = parsed_args["nh_poly"]
-        quad = Spaces.Quadratures.GLL{nh_poly + 1}()
-        horizontal_mesh = cubed_sphere_mesh(; radius, h_elem)
-        h_space =
-            make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble)
-        z_stretch = if parsed_args["z_stretch"]
-            Meshes.GeneralizedExponentialStretching(dz_bottom, dz_top)
-        else
-            Meshes.Uniform()
-        end
-        if warp_function == nothing
-            make_hybrid_spaces(h_space, z_max, z_elem, z_stretch)
-        else
-            make_hybrid_spaces(
-                h_space,
-                z_max,
-                z_elem,
-                z_stretch;
-                surface_warp = warp_function,
-                topo_smoothing = parsed_args["topo_smoothing"],
-            )
-        end
-    elseif parsed_args["config"] == "column" # single column
-        @warn "perturb_initstate flag is ignored for single column configuration"
-        FT = eltype(params)
-        Δx = FT(1) # Note: This value shouldn't matter, since we only have 1 column.
-        quad = Spaces.Quadratures.GL{1}()
-        horizontal_mesh = periodic_rectangle_mesh(;
-            x_max = Δx,
-            y_max = Δx,
-            x_elem = 1,
-            y_elem = 1,
-        )
-        if bubble
-            @warn "Bubble correction not compatible with single column configuration. It will be switched off."
-            bubble = false
-        end
-        h_space =
-            make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble)
-        z_stretch = if parsed_args["z_stretch"]
-            Meshes.GeneralizedExponentialStretching(dz_bottom, dz_top)
-        else
-            Meshes.Uniform()
-        end
-        make_hybrid_spaces(h_space, z_max, z_elem, z_stretch)
-    elseif parsed_args["config"] == "box"
-        FT = eltype(params)
-        nh_poly = parsed_args["nh_poly"]
-        quad = Spaces.Quadratures.GLL{nh_poly + 1}()
-        x_elem = Int(parsed_args["x_elem"])
-        x_max = FT(parsed_args["x_max"])
-        y_elem = Int(parsed_args["y_elem"])
-        y_max = FT(parsed_args["y_max"])
-        horizontal_mesh = periodic_rectangle_mesh(;
-            x_max = x_max,
-            y_max = y_max,
-            x_elem = x_elem,
-            y_elem = y_elem,
-        )
-        h_space =
-            make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble)
-        z_stretch = if parsed_args["z_stretch"]
-            Meshes.GeneralizedExponentialStretching(dz_bottom, dz_top)
-        else
-            Meshes.Uniform()
-        end
-        make_hybrid_spaces(
-            h_space,
-            z_max,
-            z_elem,
-            z_stretch;
-            surface_warp = warp_function,
-        )
-    elseif parsed_args["config"] == "plane"
-        FT = eltype(params)
-        nh_poly = parsed_args["nh_poly"]
-        quad = Spaces.Quadratures.GLL{nh_poly + 1}()
-        x_elem = Int(parsed_args["x_elem"])
-        x_max = FT(parsed_args["x_max"])
-        horizontal_mesh =
-            periodic_line_mesh(; x_max = x_max, x_elem = x_elem)
-        h_space =
-            make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble)
-        z_stretch = if parsed_args["z_stretch"]
-            Meshes.GeneralizedExponentialStretching(dz_bottom, dz_top)
-        else
-            Meshes.Uniform()
-        end
-        make_hybrid_spaces(
-            h_space,
-            z_max,
-            z_elem,
-            z_stretch;
-            surface_warp = warp_function,
+    z_stretch = parsed_args["z_stretch"]
+
+    available_grids = Dict(
+        "sphere" =>
+            z_stretch ? CAG.Sphere : CAG.VerticallyUniformSphereGrid,
+        "box" => z_stretch ? CAG.Box : CAG.VerticallyUniformBoxGrid,
+        "column" =>
+            z_stretch ? CAG.StretchedColumnGrid : CAG.UniformColumnGrid,
+        "plane" => z_stretch ? CAG.Plane : CAG.VerticallyUniformPlaneGrid,
+    )
+
+    GridType = available_grids[parsed_args["config"]]
+
+    # The Grid objects are such that they mimic the parsed_args interface (for the most
+    # part), so we can directly set them up from the parsed_args.
+    #
+    # Base.kwarg_decl. returns the list of keyword arguments for GridType as symbols
+    #
+    # The main differences are that `topography` takes a function, `bubble` is renamed to
+    # `enable_bubble`. We also have some additional parameters such as radius and comms_ctx.
+    # To take care of this, we will defined a "fixed" parsed_args which has the correct
+    # values
+    constructor_arguments_needed = Base.kwarg_decl.(methods(GridType))[1]
+
+    parsed_args_fixed = copy(parsed_args)
+    parsed_args_fixed["enable_bubble"] = parsed_args_fixed["bubble"]
+    parsed_args_fixed["radius"] = CAP.planet_radius(params)
+    parsed_args_fixed["comms_ctx"] = comms_ctx
+    parsed_args_fixed["float_type"] = FT
+
+    if :topography in constructor_arguments_needed
+        available_topographies = Dict(
+            "NoWarp" => nothing,
+            "DCMIP200" => topography_dcmip200,
+            "Earth" => generate_topography_earth(
+                smooth_degree = Int(parsed_args["smoothing_order"]),
+            ),
+            "Agnesi" => topography_agnesi,
+            "Schar" => topography_schar,
         )
+
+        topography = parsed_args["topography"]
+
+        haskey(available_topographies, topography) ||
+            error("Topography $topography not available")
+        @info "Topography" topography
+
+        warp_function = available_topographies[topography]
+        parsed_args_fixed["topography"] = warp_function
+    else
+        parsed_args_fixed["topography"] = nothing
     end
-    ncols = Fields.ncolumns(center_space)
-    ndofs_total = ncols * z_elem
-    hspace = Spaces.horizontal_space(center_space)
+
+    kargs = Dict(
+        arg => parsed_args_fixed[String(arg)] for
+        arg in constructor_arguments_needed
+    )
+
+    grid = GridType(; kargs...)
+
+    ncols = Fields.ncolumns(grid.center_space)
+    ndofs_total = ncols * grid.z_elem
+    hspace = Spaces.horizontal_space(grid.center_space)
     quad_style = Spaces.quadrature_style(hspace)
     Nq = Spaces.Quadratures.degrees_of_freedom(quad_style)
 
     @info "Resolution stats: " Nq h_elem z_elem ncols ndofs_total
     return (;
-        center_space,
-        face_space,
-        horizontal_mesh,
-        quad,
-        z_max,
-        z_elem,
-        z_stretch,
+        grid.center_space,
+        grid.face_space,
+        hspace.topology.mesh,
+        quad_style,
+        grid.z_max,
+        grid.z_elem,
+        grid.z_stretch,
     )
 end