Use CommonSpaces to create spaces

Author: nefrathenrici

.buildkite/Manifest-v1.11.toml

@@ -406,9 +406,9 @@ weakdeps = ["CUDA", "MPI"]
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
-git-tree-sha1 = "ca717446978d2815b4fa23a62a2131861e44d1e8"
+git-tree-sha1 = "a2acae071e36c1c69c94a83d1fb74b25e8b0fde0"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.14.42"
+version = "0.14.43"
 weakdeps = ["CUDA", "Krylov"]
 
     [deps.ClimaCore.extensions]
@@ -1864,7 +1864,7 @@ version = "2.5.5+0"
 [[deps.OpenLibm_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
-version = "0.8.5+0"
+version = "0.8.1+4"
 
 [[deps.OpenMPI_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"]

.buildkite/PrecompileCI/src/PrecompileCI.jl

@@ -3,7 +3,6 @@ module PrecompileCI
 using PrecompileTools, Logging
 import ClimaAtmos as CA
 import ClimaComms
-import ClimaCore: InputOutput, Meshes, Spaces, Quadratures
 import ClimaParams
 
 @compile_workload begin
@@ -14,34 +13,38 @@ import ClimaParams
         x_elem = y_elem = 2
         x_max = y_max = 1e8
         z_max = FT(30000.0)
-        dz_bottom = FT(500) # other values?
-        z_stretch = Meshes.HyperbolicTangentStretching(dz_bottom) # Meshes.Uniform()
-        bubble = true # false
-        parsed_args =
-            Dict{String, Any}("topography" => "NoWarp", "topo_smoothing" => false)
-        comms_ctx = ClimaComms.context(ClimaComms.CPUSingleThreaded())
-        deep = false
-
-        # constants
-        quad = Quadratures.GLL{4}()
+        dz_bottom = FT(500)
+        z_stretch = true
+        bubble = true
+        nh_poly = 3 # GLL{4} = nh_poly + 1
+        # TODO: compile CUDA methods as well
+        context = ClimaComms.context(ClimaComms.CPUSingleThreaded())
+        topography = CA.NoTopography()
         params = CA.ClimaAtmosParameters(FT)
         radius = CA.Parameters.planet_radius(params)
 
-        # Sphere
-        horz_mesh = CA.cubed_sphere_mesh(; radius, h_elem)
-        h_space = CA.make_horizontal_space(horz_mesh, quad, comms_ctx, bubble)
-        CA.make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args)
-
-        # box
-        horizontal_mesh = CA.periodic_rectangle_mesh(; x_max, y_max, x_elem, y_elem)
-        h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble)
-        # This is broken
-        # CA.make_hybrid_spaces(h_space, z_max, z_elem, z_stretch; parsed_args)
-
-        # plane
-        horizontal_mesh = CA.periodic_line_mesh(; x_max, x_elem)
-        h_space = CA.make_horizontal_space(horizontal_mesh, quad, comms_ctx, bubble)
+        sphere_grid = CA.SphereGrid(
+            FT;
+            context, radius, h_elem, nh_poly, z_elem, z_max, z_stretch,
+            dz_bottom, bubble,
+            topography,
+        )
+        box_grid = CA.BoxGrid(
+            FT; context, x_elem, x_max, y_elem, y_max, z_elem, z_max, nh_poly,
+            z_stretch, dz_bottom, bubble,
+            periodic_x = true, periodic_y = true, topography,
+        )
+        plane_grid = CA.PlaneGrid(
+            FT; context, x_elem, x_max, z_elem, z_max, nh_poly, z_stretch,
+            dz_bottom, bubble, periodic_x = true,
+            topography,
+        )
+        column_grid = CA.ColGrid(
+            FT; context, z_elem, z_max, z_stretch, dz_bottom,
+        )
+        all_grids = (sphere_grid, box_grid, plane_grid, column_grid)
+        foreach(CA.get_spaces, all_grids)
     end
 end
 
-end # module Precompile
+end # module PrecompileCI

.buildkite/pipeline.yml

@@ -1281,7 +1281,7 @@ steps:
           CLIMACOMMS_DEVICE: "CUDA"
           CLIMA_NAME_CUDA_KERNELS_FROM_STACK_TRACE: "true"
         agents:
-          slurm_mem: 24GB
+          slurm_mem: 28GB
           slurm_gpus: 1
 
   - group: "Flame graphs"

Project.toml

@@ -44,12 +44,12 @@ ArgParse = "1"
 Artifacts = "1"
 AtmosphericProfilesLibrary = "0.1.7"
 ClimaComms = "0.6.9"
-ClimaCore = "0.14.37"
+ClimaCore = "0.14.43"
 ClimaDiagnostics = "0.2.12"
 ClimaInterpolations = "0.1.0"
 ClimaParams = "1.0.2"
 ClimaTimeSteppers = "0.8.2"
-ClimaUtilities = "0.1.23"
+ClimaUtilities = "0.1.27"
 CloudMicrophysics = "0.28, 0.29"
 Dates = "1"
 ForwardDiff = "1"

calibration/experiments/gcm_driven_scm/helper_funcs.jl

@@ -23,8 +23,8 @@ CLIMADIAGNOSTICS_LES_NAME_MAP =
 """Get z cell centers coordinates for CA run, given config. """
 function get_z_grid(atmos_config; z_max = nothing)
     params = CA.ClimaAtmosParameters(atmos_config)
-    spaces =
-        CA.get_spaces(atmos_config.parsed_args, params, atmos_config.comms_ctx)
+    grid = CA.get_grid(atmos_config.parsed_args, params, atmos_config.comms_ctx)
+    spaces = CA.get_spaces(grid, atmos_config.comms_ctx)
     coord = CA.Fields.coordinate_field(spaces.center_space)
     z_vec = convert(Vector{Float64}, parent(coord.z)[:])
     if !isnothing(z_max)

docs/src/api.md

@@ -43,12 +43,21 @@ ClimaAtmos.InitialConditions.Bomex
 ClimaAtmos.InitialConditions.Soares
 ```
 
-### Helper
+## Helper
 
 ```@docs
 ClimaAtmos.InitialConditions.ColumnInterpolatableField
 ```
 
+## Grids
+
+```@docs
+ClimaAtmos.ColGrid
+ClimaAtmos.SphereGrid
+ClimaAtmos.PlaneGrid
+ClimaAtmos.BoxGrid
+```
+
 ## Jacobian
 
 ```@docs
@@ -68,6 +77,8 @@ ClimaAtmos.ScharTopography
 ClimaAtmos.EarthTopography
 ClimaAtmos.DCMIP200Topography
 ClimaAtmos.Hughes2023Topography
+ClimaAtmos.SLEVEWarp
+ClimaAtmos.LinearWarp
 ```
 
 ### Internals

post_processing/ci_plots.jl

@@ -437,8 +437,6 @@ make_plots_generic(
     simulation_path,
     vars,
     time = LAST_SNAP,
-    x = 0.0, # Our columns are still 3D objects...
-    y = 0.0,
     more_kwargs = YLINEARSCALE,
 )
 ```
@@ -453,8 +451,6 @@ make_plots_generic(
     simulation_path,
     vars,
     time = LAST_SNAP,
-    x = 0.0, # Our columns are still 3D objects...
-    y = 0.0,
     more_kwargs = YLINEARSCALE,
 )
 ```
@@ -589,14 +585,23 @@ ColumnPlots = Union{
 function make_plots(::ColumnPlots, output_paths::Vector{<:AbstractString})
     simdirs = SimDir.(output_paths)
     short_names = ["ta", "wa"]
-    vars = map_comparison(get, simdirs, short_names)
+    vars = map_comparison(simdirs, short_names) do simdir, short_name
+        var = get(simdir; short_name)
+        # For vertical-only (FiniteDifferenceGrid) spaces, the data may have 
+        # extra singleton dimensions. Check and squeeze if needed.
+        if haskey(var.dims, "x") && length(var.dims["x"]) == 1
+            var = slice(var; x = var.dims["x"][1])
+        end
+        if haskey(var.dims, "y") && length(var.dims["y"]) == 1
+            var = slice(var; y = var.dims["y"][1])
+        end
+        return var
+    end
 
     make_plots_generic(
         output_paths,
         vars,
         time = LAST_SNAP,
-        x = 0.0, # Our columns are still 3D objects...
-        y = 0.0,
         MAX_NUM_COLS = length(simdirs),
         more_kwargs = YLINEARSCALE,
     )
@@ -644,7 +649,7 @@ function make_plots(
     end
 
     vars = [
-        slice(get(simdir; short_name), x = 0.0, y = 0.0) for
+        get(simdir; short_name) for
         short_name in short_names
     ]
 
@@ -689,7 +694,7 @@ function make_plots(
     surface_precip = read_var(simdir.variable_paths["pr"]["inst"]["10s"])
     viz.line_plot1D!(
         fig,
-        slice(surface_precip, x = 0.0, y = 0.0);
+        surface_precip;
         p_loc = [pr_row, 1:3],
     )
 
@@ -1371,6 +1376,12 @@ EDMFBoxPlots = Union{
     Val{:diagnostic_edmfx_dycoms_rf01_box},
     Val{:diagnostic_edmfx_trmm_box_0M},
     Val{:diagnostic_edmfx_dycoms_rf01_explicit_box},
+    Val{:prognostic_edmfx_bomex_box},
+    Val{:rcemipii_box_diagnostic_edmfx},
+    Val{:diagnostic_edmfx_trmm_stretched_box},
+}
+
+EDMFColumnPlots = Union{
     Val{:prognostic_edmfx_adv_test_column},
     Val{:prognostic_edmfx_gabls_column},
     Val{:prognostic_edmfx_gabls_column_sparse_autodiff},
@@ -1386,13 +1397,10 @@ EDMFBoxPlots = Union{
     Val{:prognostic_edmfx_simpleplume_column},
     Val{:prognostic_edmfx_gcmdriven_column},
     Val{:prognostic_edmfx_tv_era5driven_column},
-    Val{:prognostic_edmfx_bomex_box},
-    Val{:rcemipii_box_diagnostic_edmfx},
     Val{:prognostic_edmfx_soares_column},
-    Val{:diagnostic_edmfx_trmm_stretched_box},
 }
 
-EDMFBoxPlotsWithPrecip = Union{
+EDMFColumnPlotsWithPrecip = Union{
     Val{:prognostic_edmfx_rico_column},
     Val{:prognostic_edmfx_rico_implicit_column},
     Val{:prognostic_edmfx_rico_column_2M},
@@ -1485,14 +1493,99 @@ end
 function make_plots(
     sim_type::Union{
         EDMFBoxPlots,
-        EDMFBoxPlotsWithPrecip,
         DiagEDMFBoxPlotsWithPrecip,
     },
     output_paths::Vector{<:AbstractString},
 )
     simdirs = SimDir.(output_paths)
 
-    if sim_type isa EDMFBoxPlotsWithPrecip
+    if sim_type isa DiagEDMFBoxPlotsWithPrecip
+        precip_names = ("husra", "hussn", "husraup", "hussnup")
+    else
+        precip_names = ()
+    end
+
+    short_names = [
+        "wa",
+        "waup",
+        "ta",
+        "taup",
+        "hus",
+        "husup",
+        "arup",
+        "tke",
+        "ua",
+        "thetaa",
+        "thetaaup",
+        "ha",
+        "haup",
+        "hur",
+        "hurup",
+        "lmix",
+        "cl",
+        "clw",
+        "clwup",
+        "cli",
+        "cliup",
+        precip_names...,
+    ]
+    reduction = "inst"
+
+    available_periods = ClimaAnalysis.available_periods(
+        simdirs[1];
+        short_name = short_names[1],
+        reduction,
+    )
+    # choose the shortest available period
+    available_periods = collect(available_periods) # ensure vector for indexing
+    period = available_periods[argmin(CA.time_to_seconds.(available_periods))]
+
+    short_name_tuples = pair_edmf_names(short_names)
+    var_groups_zt =
+        map_comparison(simdirs, short_name_tuples) do simdir, name_tuple
+            return [
+                slice(
+                    get(simdir; short_name, reduction, period),
+                    x = 0.0,
+                    y = 0.0,
+                ) for short_name in name_tuple
+            ]
+        end
+
+    var_groups_z = [
+        ([slice(v, time = LAST_SNAP) for v in group]...,) for
+        group in var_groups_zt
+    ]
+
+    tmp_file = make_plots_generic(
+        output_paths,
+        output_name = "tmp",
+        var_groups_z;
+        plot_fn = plot_edmf_vert_profile!,
+        MAX_NUM_COLS = 2,
+        MAX_NUM_ROWS = 4,
+    )
+
+    make_plots_generic(
+        output_paths,
+        vcat(var_groups_zt...),
+        plot_fn = plot_parsed_attribute_title!,
+        summary_files = [tmp_file],
+        MAX_NUM_COLS = 2,
+        MAX_NUM_ROWS = 4,
+    )
+end
+
+function make_plots(
+    sim_type::Union{
+        EDMFColumnPlots,
+        EDMFColumnPlotsWithPrecip,
+    },
+    output_paths::Vector{<:AbstractString},
+)
+    simdirs = SimDir.(output_paths)
+
+    if sim_type isa EDMFColumnPlotsWithPrecip
         if sim_type isa Val{:prognostic_edmfx_rico_column_2M}
             precip_names = (
                 "husra",
@@ -1512,8 +1605,6 @@ function make_plots(
             precip_names =
                 ("husra", "hussn", "husraup", "hussnup", "husraen", "hussnen")
         end
-    elseif sim_type isa DiagEDMFBoxPlotsWithPrecip
-        precip_names = ("husra", "hussn", "husraup", "hussnup")
     else
         precip_names = ()
     end
@@ -1557,11 +1648,7 @@ function make_plots(
     var_groups_zt =
         map_comparison(simdirs, short_name_tuples) do simdir, name_tuple
             return [
-                slice(
-                    get(simdir; short_name, reduction, period),
-                    x = 0.0,
-                    y = 0.0,
-                ) for short_name in name_tuple
+                get(simdir; short_name, reduction, period) for short_name in name_tuple
             ]
         end