Reformulate IMEX ARK and SSPRK timestepping schemes

Author: dennisYatunin

.buildkite/pipeline.yml

@@ -2,7 +2,6 @@ agents:
   queue: new-central
   slurm_mem: 8G
   modules: climacommon/2024_05_27
-  partition: expansion
 
 env:
   JULIA_LOAD_PATH: "${JULIA_LOAD_PATH}:${BUILDKITE_BUILD_CHECKOUT_PATH}/.buildkite"
@@ -24,40 +23,38 @@ steps:
       - echo "--- Instantiate perf"
       - "julia --project=perf -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true)'"
 
+      - echo "--- Instantiate docs"
+      - "julia --project=docs -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true)'"
+
       - echo "--- Package status"
       - "julia --project -e 'using Pkg; Pkg.status()'"
-
-
-  # - label: "init gpu env"
-  #   key: "init_gpu_env"
-  #   command:
-  #     - echo "--- Configure MPI"
-  #     - julia -e 'using Pkg; Pkg.add("MPIPreferences"); using MPIPreferences; use_system_binary()'
-
-  #     - echo "--- Instantiate project"
-  #     - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true)'"
-
-  #     - echo "--- Instantiate test"
-  #     - "julia --project=test -e 'using Pkg; Pkg.develop(path=\".\"); Pkg.instantiate(;verbose=true); Pkg.precompile()'"
-
-  #     - echo "--- Initialize CUDA runtime"
-  #     - "julia --project -e 'using CUDA; CUDA.precompile_runtime(); CUDA.versioninfo()'"
-
-  #     - echo "--- Package status"
-  #     - "julia --project -e 'using Pkg; Pkg.status()'"
-  #     slurm_gres: "gpu:1"
+    agents:
+      slurm_gpus: 1
 
   - wait
 
   - label: "CPU tests"
     command: "julia --project=test --check-bounds=yes test/runtests.jl"
     artifact_paths: "output/*"
 
-  # - label: "GPU tests"
-  #   command:
-  #     - "julia --project=test --check-bounds=yes test/runtests.jl CuArray"
-  #   artifact_paths: "output/*"
-  #     slurm_gres: "gpu:1"
+  - label: "Deformational flow limiter test (CPU)"
+    command:
+      - "julia --project=docs --check-bounds=yes docs/src/test_deformational_flow.jl CPU"
+    artifact_paths: "output_CPU/*"
+
+  - label: "Deformational flow limiter test (GPU)"
+    command:
+      - "julia --project=docs docs/src/test_deformational_flow.jl GPU"
+    artifact_paths: "output_GPU/*"
+    agents:
+      slurm_gpus: 1
+
+  - label: "Deformational flow limiter test (GPU w/ check-bounds)"
+    command:
+      - "julia --project=docs --check-bounds=yes docs/src/test_deformational_flow.jl GPU_checkbounds"
+    artifact_paths: "output_GPU_checkbounds/*"
+    agents:
+      slurm_gpus: 1
 
   - label: "Flame graph (1D diffusion)"
     command: "julia --project=perf perf/flame.jl --job_id diffusion_1D"

Project.toml

@@ -16,6 +16,7 @@ LinearOperators = "5c8ed15e-5a4c-59e4-a42b-c7e8811fb125"
 NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [weakdeps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
@@ -45,4 +46,5 @@ NVTX = "0.3"
 SciMLBase = "1, 2"
 StaticArrays = "1"
 StatsBase = "0.33, 0.34"
+UnrolledUtilities = "0.1"
 julia = "1.8"

docs/Manifest.toml

@@ -249,7 +249,7 @@ uuid = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 version = "0.2.10"
 
 [[deps.ClimaTimeSteppers]]
-deps = ["ClimaComms", "Colors", "DataStructures", "DiffEqBase", "DiffEqCallbacks", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays"]
+deps = ["ClimaComms", "Colors", "DataStructures", "DiffEqBase", "DiffEqCallbacks", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays", "UnrolledUtilities"]
 path = ".."
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 version = "0.7.31"
@@ -2123,6 +2123,11 @@ git-tree-sha1 = "6cc9d682755680e0f0be87c56392b7651efc2c7b"
 uuid = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
 version = "0.1.5"
 
+[[deps.UnrolledUtilities]]
+git-tree-sha1 = "b73f7a7c25a2618c5052c80ed32b07e471cc6cb0"
+uuid = "0fe1646c-419e-43be-ac14-22321958931b"
+version = "0.1.2"
+
 [[deps.UnsafeAtomics]]
 git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278"
 uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f"

docs/src/api/ode_solvers.md

@@ -4,19 +4,39 @@
 CurrentModule = ClimaTimeSteppers
 ```
 
-## Interface
+## Tableau Interface
 
 ```@docs
-AbstractAlgorithmConstraint
-Unconstrained
-SSP
-IMEXTableau
-IMEXAlgorithm
-ExplicitTableau
-ExplicitAlgorithm
+ClimaTimeSteppers.RKTableau
+ClimaTimeSteppers.ButcherTableau
+ClimaTimeSteppers.ShuOsherTableau
+ClimaTimeSteppers.PaddedTableau
+ClimaTimeSteppers.ARKTableau
+ClimaTimeSteppers.is_ERK
+ClimaTimeSteppers.is_DIRK
 ```
 
-## IMEX Algorithm Names
+## Algorithm Interface
+
+```@docs
+AbstractAlgorithmName
+ClimaTimeSteppers.RKAlgorithmName
+ClimaTimeSteppers.SSPRKAlgorithmName
+ClimaTimeSteppers.ARKAlgorithmName
+ClimaTimeSteppers.IMEXSSPRKAlgorithmName
+RKAlgorithm
+ARKAlgorithm
+```
+
+## RK Algorithm Names
+
+```@docs
+SSP22Heuns
+SSP33ShuOsher
+RK4
+```
+
+## ARK Algorithm Names
 
 ```@docs
 ARS111
@@ -54,14 +74,6 @@ ARK548L2SA2
 SSPKnoth
 ```
 
-## Explicit Algorithm Names
-
-```@docs
-SSP22Heuns
-SSP33ShuOsher
-RK4
-```
-
 ## Old LSRK Interface
 
 ```@docs

docs/src/dev/report_gen.jl

@@ -8,66 +8,53 @@ using InteractiveUtils: subtypes
 
 ENV["GKSwstype"] = "nul" # avoid displaying plots
 
-include(joinpath(@__DIR__, "..", "plotting_utils.jl"))
 include(joinpath(pkgdir(ClimaTimeSteppers), "test", "problems.jl"))
+include(joinpath(@__DIR__, "..", "plotting_utils.jl"))
 
 all_subtypes(::Type{T}) where {T} = isabstracttype(T) ? vcat(all_subtypes.(subtypes(T))...) : [T]
 
 let # Convergence
     title = "All Algorithms"
     algorithm_names = map(T -> T(), all_subtypes(ClimaTimeSteppers.AbstractAlgorithmName))
-    algorithm_names = filter(name -> !(name isa ARK437L2SA1 || name isa ARK548L2SA2), algorithm_names) # too high order
 
-    # NOTE: Some imperfections in the convergence order for SSPKnoth are to be
-    # expected because we are not using the exact Jacobian
+    verify_convergence(title, algorithm_names, ark_analytic_nonlin_test_cts(Float64), 300)
+    verify_convergence(title, algorithm_names, ark_analytic_sys_test_cts(Float64), 350)
+    verify_convergence(title, algorithm_names, onewaycouple_mri_test_cts(Float64), 2000)
+    verify_convergence(
+        title,
+        algorithm_names,
+        ark_analytic_test_cts(Float64),
+        16650;
+        num_test_points = 6,
+        num_steps_scaling_factor = 23,
+        super_convergence = (ARS121(),),
+    )
 
-    verify_convergence(title, algorithm_names, ark_analytic_nonlin_test_cts(Float64), 200)
-    verify_convergence(title, algorithm_names, ark_analytic_sys_test_cts(Float64), 400)
-    verify_convergence(title, algorithm_names, ark_analytic_test_cts(Float64), 40000; super_convergence = (ARS121(),))
-    verify_convergence(title, algorithm_names, onewaycouple_mri_test_cts(Float64), 10000; num_steps_scaling_factor = 5)
     verify_convergence(
         title,
         algorithm_names,
         climacore_1Dheat_test_cts(Float64),
-        400;
-        num_steps_scaling_factor = 4,
-        numerical_reference_algorithm_name = ARS343(),
+        40;
+        numerical_reference_algorithm_name = ARK548L2SA2(),
+        numerical_reference_num_steps = 500000,
     )
-    rosenbrock_schems = filter(name -> name isa ClimaTimeSteppers.RosenbrockAlgorithmName, algorithm_names)
-    verify_convergence(title, rosenbrock_schems, climacore_1Dheat_test_implicit_cts(Float64), 400)
     verify_convergence(
         title,
         algorithm_names,
         climacore_2Dheat_test_cts(Float64),
-        600;
-        num_steps_scaling_factor = 4,
-        numerical_reference_algorithm_name = ARS343(),
+        40;
+        numerical_reference_algorithm_name = ARK548L2SA2(),
+        numerical_reference_num_steps = 500000,
     )
-end
 
-let # Unconstrained vs SSP results without limiters
-    algorithm_names = map(T -> T(), all_subtypes(ClimaTimeSteppers.IMEXSSPRKAlgorithmName))
-    for (test_case, num_steps) in (
-        (ark_analytic_nonlin_test_cts(Float64), 200),
-        (ark_analytic_sys_test_cts(Float64), 400),
-        (ark_analytic_test_cts(Float64), 40000),
-        (onewaycouple_mri_test_cts(Float64), 10000),
-        (climacore_1Dheat_test_cts(Float64), 200),
-        (climacore_2Dheat_test_cts(Float64), 200),
+    verify_convergence(
+        title,
+        algorithm_names,
+        climacore_1Dheat_test_implicit_cts(Float64),
+        60;
+        num_test_points = 4,
+        num_steps_scaling_factor = 8,
+        numerical_reference_algorithm_name = ARK548L2SA2(),
+        numerical_reference_num_steps = 500000,
     )
-        prob = test_case.split_prob
-        dt = test_case.t_end / num_steps
-        newtons_method = NewtonsMethod(; max_iters = test_case.linear_implicit ? 1 : 2)
-        for algorithm_name in algorithm_names
-            algorithm = IMEXAlgorithm(algorithm_name, newtons_method)
-            reference_algorithm = IMEXAlgorithm(algorithm_name, newtons_method, Unconstrained())
-            solution = solve(deepcopy(prob), algorithm; dt).u[end]
-            reference_solution = solve(deepcopy(prob), reference_algorithm; dt).u[end]
-            if norm(solution .- reference_solution) / norm(reference_solution) > 30 * eps(Float64)
-                alg_str = string(nameof(typeof(algorithm_name)))
-                @warn "Unconstrained and SSP versions of $alg_str \
-                       give different results for $(test_case.test_name)"
-            end
-        end
-    end
 end

docs/src/dev/report_gen.md

@@ -1,6 +1,6 @@
 # Verifying Correctness
 
-The `IMEXAlgorithm` supports problems that specify any combination of the following: an implicit tendency `T_imp!`, an explicit tendency `T_exp!`, a limited tendency `T_lim!`, a function `dss!` that applies a direct stiffness summation, and a function `lim!` that applies a monotonicity-preserving limiter.
+The `ARKAlgorithm` supports problems that specify any combination of the following: an implicit tendency `T_imp!`, an explicit tendency `T_exp!`, a limited tendency `T_lim!`, a function `dss!` that applies a direct stiffness summation, and a function `lim!` that applies a monotonicity-preserving limiter.
 
 ## Convergence without a Limiter
 
@@ -12,7 +12,9 @@ The test cases we use for this analysis are:
     - `ark_analytic`, which uses a nonlinear `T_exp!` and a linear `T_imp!`
     - `ark_analytic_sys` and `ark_onewaycouple_mri`, which use a linear `T_imp!`
     - `ark_analytic_nonlin`, which uses a nonlinear `T_imp!`
-    - `1d_heat_equation` and `2d_heat_equation`, which use a nonlinear `T_exp!` and `dss!`, where the spatial discretization is implemented using `ClimaCore`
+    - `1d_heat_equation`, which uses a nonlinear `T_exp!` implemented with `ClimaCore`
+    - `2d_heat_equation`, which uses a nonlinear `T_exp!` and `dss!` implemented with `ClimaCore`
+    - `1d_heat_equation_implicit`, which uses a nonlinear `T_imp!` implemented with `ClimaCore`
 
 ```@example
 include("report_gen.jl")
@@ -23,6 +25,7 @@ include("report_gen.jl")
  ![](output/convergence_ark_analytic_nonlin_all_algorithms.png)
  ![](output/convergence_1d_heat_equation_all_algorithms.png)
  ![](output/convergence_2d_heat_equation_all_algorithms.png)
+ ![](output/convergence_1d_heat_equation_implicit_all_algorithms.png)
 
 ## Errors with a Limiter
 
@@ -40,8 +43,8 @@ limiter_summary(Float64, [SSP333(), ARS343()], horizontal_deformational_flow_tes
 ```
 
 Plots of the tracer specific humidities that were used to compute this table are shown below.
- ![](output/limiter_summary_SSP333.png)
- ![](output/limiter_summary_ARS343.png)
+ ![](output/horizontal_deformational_flow_limiter_summary_SSP333.png)
+ ![](output/horizontal_deformational_flow_limiter_summary_ARS343.png)
 
 ## References