diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml index 032d028974..fb9587b760 100644 --- a/.buildkite/Manifest-v1.11.toml +++ b/.buildkite/Manifest-v1.11.toml @@ -352,7 +352,7 @@ version = "0.5.17" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "NullBroadcasts", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"] +deps = ["Adapt", "ArgParse", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "ForwardDiff", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "NullBroadcasts", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"] path = ".." uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717" version = "0.30.0" diff --git a/.buildkite/ci_driver.jl b/.buildkite/ci_driver.jl index 3207bcd3a0..308cddec38 100644 --- a/.buildkite/ci_driver.jl +++ b/.buildkite/ci_driver.jl @@ -41,6 +41,7 @@ import Base.Filesystem: rm import Statistics: mean import LinearAlgebra: norm_sqr include(joinpath(pkgdir(CA), "post_processing", "ci_plots.jl")) +include(joinpath(pkgdir(CA), "post_processing", "jacobian_plots.jl")) ref_job_id = config.parsed_args["reference_job_id"] reference_job_id = isnothing(ref_job_id) ? simulation.job_id : ref_job_id @@ -187,6 +188,7 @@ if ClimaComms.iamroot(config.comms_ctx) "reproducibility_utils.jl", ), ) + @info "Plotting" paths = latest_comparable_dirs() # __build__ path (not job path) if isempty(paths) @@ -218,6 +220,17 @@ if ClimaComms.iamroot(config.comms_ctx) end make_plots(Val(Symbol(reference_job_id)), paths) end + if ( + isnothing(config.parsed_args["plot_jacobian"]) ? + config.parsed_args["debug_approximate_jacobian"] : + config.parsed_args["plot_jacobian"] + ) + make_jacobian_plots( + simulation.output_dir, + integrator.u, + float(integrator.dt), + ) + end @info "Plotting done" if islink(simulation.output_dir) diff --git a/.buildkite/longruns_gpu/pipeline.yml b/.buildkite/longruns_gpu/pipeline.yml index d752c2a90d..eaf2f8eb85 100644 --- a/.buildkite/longruns_gpu/pipeline.yml +++ b/.buildkite/longruns_gpu/pipeline.yml @@ -25,6 +25,8 @@ steps: - julia --project=.buildkite -e 'using Pkg; Pkg.precompile()' - julia --project=.buildkite -e 'using CUDA; CUDA.precompile_runtime()' - julia --project=.buildkite -e 'using Pkg; Pkg.status()' + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"Thermodynamics\", rev=\"dy/more_autodiff_fixes\"))'" + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"CloudMicrophysics\", rev=\"main\"))'" agents: slurm_gpus: 1 @@ -35,192 +37,214 @@ steps: - wait - - group: "helem 30 dycore" - steps: - - - label: ":computer: hydrostatic balance" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 12:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_hydrostatic_balance" - - - label: ":computer: dry baroclinic wave" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 12:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_dry_baroclinic_wave" - - - label: ":computer: dry baroclinic wave high res" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 12:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_dry_baroclinic_wave_he60" - - - label: ":computer: baroclinic wave equilmoist" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 12:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_moist_baroclinic_wave" - - - label: ":computer: baroclinic wave equilmoist high res" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_moist_baroclinic_wave_he60" - - - label: ":computer: dry held-suarez" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 12:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_dry_held_suarez" - - - label: ":computer: held-suarez, equilmoist" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_moist_held_suarez" - - - group: "helem 16 aquaplanet" - steps: - - - label: ":computer: aquaplanet equilmoist allsky radiation + 0M microphysics" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_aquaplanet_allsky_0M" - - - label: ":computer: aquaplanet equilmoist allsky radiation + diagnostic edmf + 0M microphysics" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_aquaplanet_allsky_diagedmf_0M" - - - label: ":computer: aquaplanet equilmoist allsky radiation + prognostic edmf + 0M microphysics" - command: - - srun julia --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_aquaplanet_allsky_progedmf_0M" - - - label: ":computer: aquaplanet equilmoist allsky radiation + 0M microphysics + earth topography" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_aquaplanet_allsky_0M_earth" - - - label: ":umbrella: aquaplanet equilmoist allsky radiation + 1M microphysics" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_aquaplanet_allsky_1M" - - - label: ":computer: aquaplanet equilmoist allsky radiation + time-varying insolation + 0M microphysics + slab ocean" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_aquaplanet_allsky_tvinsol_0M_slabocean" - - - group: "DYAMOND" - - steps: - - - label: ":computer: aquaplanet dyamond" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "longrun_aquaplanet_dyamond" + # - group: "helem 30 dycore" + # steps: + + # - label: ":computer: hydrostatic balance" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 12:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_hydrostatic_balance" + + # - label: ":computer: dry baroclinic wave" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 12:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_dry_baroclinic_wave" + + # - label: ":computer: dry baroclinic wave high res" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 12:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_dry_baroclinic_wave_he60" + + # - label: ":computer: baroclinic wave equilmoist" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 12:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_moist_baroclinic_wave" + + # - label: ":computer: baroclinic wave equilmoist high res" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_moist_baroclinic_wave_he60" + + # - label: ":computer: dry held-suarez" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 12:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_dry_held_suarez" + + # - label: ":computer: held-suarez, equilmoist" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_moist_held_suarez" + + # - group: "helem 16 aquaplanet" + # steps: + + # - label: ":computer: aquaplanet equilmoist allsky radiation + 0M microphysics" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_aquaplanet_allsky_0M" + + # - label: ":computer: aquaplanet equilmoist allsky radiation + diagnostic edmf + 0M microphysics" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_aquaplanet_allsky_diagedmf_0M" + + # - label: ":computer: aquaplanet equilmoist allsky radiation + prognostic edmf + 0M microphysics" + # command: + # - srun julia --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_aquaplanet_allsky_progedmf_0M" + + # - label: ":computer: aquaplanet equilmoist allsky radiation + 0M microphysics + earth topography" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_aquaplanet_allsky_0M_earth" + + # - label: ":umbrella: aquaplanet equilmoist allsky radiation + 1M microphysics" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_aquaplanet_allsky_1M" + + # - label: ":computer: aquaplanet equilmoist allsky radiation + time-varying insolation + 0M microphysics + slab ocean" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_aquaplanet_allsky_tvinsol_0M_slabocean" + + # - group: "DYAMOND" + + # steps: + + # - label: ":computer: aquaplanet dyamond" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "longrun_aquaplanet_dyamond" - group: "atmos-only coupler runs" steps: - - label: ":computer: amip target diagnostic edmf" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "amip_target_diagedmf" - - - label: ":computer: amip target edonly edmf" - command: - - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME - artifact_paths: "$$JOB_NAME/output_active/*" - agents: - slurm_gpus: 1 - slurm_time: 24:00:00 - env: - CLIMACOMMS_DEVICE: "CUDA" - JOB_NAME: "amip_target_edonly" + # - label: ":computer: amip target diagnostic edmf" + # command: + # - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + # artifact_paths: "$$JOB_NAME/output_active/*" + # agents: + # slurm_gpus: 1 + # slurm_time: 24:00:00 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # JOB_NAME: "amip_target_diagedmf" + + - label: ":computer: amip target edonly edmf" + command: + - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + artifact_paths: "$$JOB_NAME/output_active/*" + agents: + slurm_gpus: 1 + slurm_time: 24:00:00 + env: + CLIMACOMMS_DEVICE: "CUDA" + JOB_NAME: "amip_target_edonly" + + - label: ":computer: amip target edonly edmf nonequil (dt = 50, slow relaxation)" + command: + - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + artifact_paths: "$$JOB_NAME/output_active/*" + agents: + slurm_gpus: 1 + slurm_time: 24:00:00 + env: + CLIMACOMMS_DEVICE: "CUDA" + JOB_NAME: "amip_target_edonly_nonequil_50_slow" + + - label: ":computer: amip target edonly edmf nonequil (dt = 50, fast relaxation)" + command: + - srun julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/$$JOB_NAME.yml --job_id $$JOB_NAME + artifact_paths: "$$JOB_NAME/output_active/*" + agents: + slurm_gpus: 1 + slurm_time: 24:00:00 + env: + CLIMACOMMS_DEVICE: "CUDA" + JOB_NAME: "amip_target_edonly_nonequil_50_fast" diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 93e05a210d..09b19b8f1c 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -31,6 +31,9 @@ steps: - echo "--- Instantiate .buildkite" - "julia --project=.buildkite -e 'using Pkg; Pkg.instantiate(;verbose=true); Pkg.precompile(;strict=true); using CUDA; CUDA.precompile_runtime(); Pkg.status()'" + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaCore\", rev=\"dy/disable_shmem\"))'" + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"Thermodynamics\", rev=\"dy/more_autodiff_fixes\"))'" + - "julia --project=.buildkite -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"CloudMicrophysics\", rev=\"main\"))'" agents: slurm_cpus_per_task: 8 @@ -119,7 +122,7 @@ steps: julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/single_column_nonorographic_gravity_wave.yml --job_id single_column_nonorographic_gravity_wave - artifact_paths: "single_column_nonorographic_gravity_wave/*" + artifact_paths: "single_column_nonorographic_gravity_wave/output_active/*" - group: "Column Examples" steps: @@ -153,6 +156,8 @@ steps: --config_file $CONFIG_PATH/box_density_current_test.yml --job_id box_density_current_test artifact_paths: "box_density_current_test/output_active/*" + agents: + slurm_mem: 20GB - label: ":computer: Box rcemipii with diagnostic edmf" command: > @@ -168,7 +173,7 @@ steps: julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/les_isdac_box.yml --job_id les_isdac_box - artifact_paths: "les_isdac_box/*" + artifact_paths: "les_isdac_box/output_active/*" env: CLIMACOMMS_DEVICE: "CUDA" agents: @@ -859,7 +864,7 @@ steps: julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/prognostic_edmfx_bomex_box.yml --job_id prognostic_edmfx_bomex_box - artifact_paths: "prognostic_edmfx_bomex_box/*" + artifact_paths: "prognostic_edmfx_bomex_box/output_active/*" agents: slurm_mem: 20GB @@ -885,7 +890,7 @@ steps: julia --color=yes --project=.buildkite .buildkite/ci_driver.jl --config_file $CONFIG_PATH/sphere_nonorographic_gravity_wave.yml --job_id sphere_nonorographic_gravity_wave - artifact_paths: "sphere_nonorographic_gravity_wave/*" + artifact_paths: "sphere_nonorographic_gravity_wave/output_active/*" env: CLIMACOMMS_DEVICE: "CUDA" agents: diff --git a/Project.toml b/Project.toml index 8555b91630..33ec400cdb 100644 --- a/Project.toml +++ b/Project.toml @@ -16,6 +16,7 @@ ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79" ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513" CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3" @@ -48,6 +49,7 @@ ClimaTimeSteppers = "0.8.2" ClimaUtilities = "0.1.22" CloudMicrophysics = "0.22.9" Dates = "1" +ForwardDiff = "0.10.36" Insolation = "0.9.2" Interpolations = "0.15.1" LazyArtifacts = "1" diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index c26aa4c9ed..542289cdfc 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -62,9 +62,6 @@ newton_rtol: surface_thermo_state_type: help: "Surface thermo state type [`GCMSurfaceThermoState` (default), `PrescribedThermoState`]" value: "GCMSurfaceThermoState" -split_ode: - help: "Use split of ODE problem. Examples: [`true` (implicit, default), `false` (explicit)]" - value: true use_krylov_method: help: "Whether to use a Krylov method to solve the linear system in Newton's method (only for ClimaTimeSteppers.jl)" value: false @@ -86,6 +83,24 @@ eisenstat_walker_forcing_alpha: jvp_step_adjustment: help: "Amount by which the step size of the forward difference approximation of the Jacobian-vector product in the Krylov method should be scaled (only used if `use_krylov_method` is `true`)" value: 1.0 +use_exact_jacobian: + help: "Whether to solve the linear system using the exact Jacobian, which is computed with forward-mode automatic differentiation once per `dt_update_exact_jacobian` (default is `false`)" + value: false +debug_approximate_jacobian: + help: "Whether to compare the approximate Jacobian against the exact Jacobian once per `dt_update_exact_jacobian` (default is `false`)" + value: true # TODO: Change this to false +only_debug_first_column_jacobian: + help: "Whether to debug only the Jacobian in the first column when `debug_approximate_jacobian` is `true` and `use_exact_jacobian` is `false` (default is equal to `n_cols > 10000`)" + value: ~ +plot_jacobian: + help: "Whether to save and plot the Jacobian, and also the difference between the approximate Jacobian and exact Jacobian when `debug_approximate_jacobian` is `true` (default is equal to `debug_approximate_jacobian`)" + value: ~ +dt_save_jacobian: + help: "Time between snapshots of the Jacobian when `plot_jacobian` is `true` (default is equal to `t_span / 5`)" + value: ~ +dt_update_exact_jacobian: + help: "Time between updates to the exact Jacobian (when either `use_exact_jacobian` or `debug_approximate_jacobian` is `true`), with any value smaller than `dt` corresponding to an update before each linear solve (default is equal to `use_exact_jacobian ? 0 : dt_save_jacobian`)" + value: ~ # Radiation rad: help: "Radiation model [`nothing` (default), `gray`, `clearsky`, `allsky`, `allskywithclear`]" diff --git a/config/longrun_configs/amip_target_edonly_nonequil_50_fast.yml b/config/longrun_configs/amip_target_edonly_nonequil_50_fast.yml new file mode 100644 index 0000000000..5d699c4fa3 --- /dev/null +++ b/config/longrun_configs/amip_target_edonly_nonequil_50_fast.yml @@ -0,0 +1,34 @@ +h_elem: 9 +z_elem: 39 +z_max: 60000.0 +dz_bottom: 100.0 +deep_atmosphere: false +rayleigh_sponge: true +viscous_sponge: true +dt_save_state_to_disk: "30days" +cloud_model: "grid_scale" +moist: "nonequil" +implicit_noneq_cloud_formation: true +precip_model: "1M" +rad: "allskywithclear" +dt_rad: "1hours" +dt_cloud_fraction: "1hours" +insolation: "timevarying" +co2_model: maunaloa +prescribe_ozone: true +aerosol_radiation: true +edmfx_sgs_diffusive_flux: true +prescribed_aerosols: ["CB1", "CB2", "DST01", "DST02", "DST03", "DST04", "DST05", "OC1", "OC2", "SO4", "SSLT01", "SSLT02", "SSLT03", "SSLT04", "SSLT05"] +surface_setup: "DefaultMoninObukhov" +turbconv: "edonly_edmfx" +implicit_diffusion: true +approximate_linear_solve_iters: 2 +ode_algo: "ARS111" +max_newton_iters_ode: 10 +dt: "50secs" +t_end: "5hours" +toml: [toml/aquaplanet_nonequil_12_120.toml] +diagnostics: + - short_name: [ta, ua, wa, va, rhoa, hur, hus, hussn, husra, clw, cli] + period: 5days +netcdf_output_at_levels: true \ No newline at end of file diff --git a/config/longrun_configs/amip_target_edonly_nonequil_50_slow.yml b/config/longrun_configs/amip_target_edonly_nonequil_50_slow.yml new file mode 100644 index 0000000000..0bc430143f --- /dev/null +++ b/config/longrun_configs/amip_target_edonly_nonequil_50_slow.yml @@ -0,0 +1,34 @@ +h_elem: 9 +z_elem: 39 +z_max: 60000.0 +dz_bottom: 100.0 +deep_atmosphere: false +rayleigh_sponge: true +viscous_sponge: true +dt_save_state_to_disk: "30days" +cloud_model: "grid_scale" +moist: "nonequil" +implicit_noneq_cloud_formation: true +precip_model: "1M" +rad: "allskywithclear" +dt_rad: "1hours" +dt_cloud_fraction: "1hours" +insolation: "timevarying" +co2_model: maunaloa +prescribe_ozone: true +aerosol_radiation: true +edmfx_sgs_diffusive_flux: true +prescribed_aerosols: ["CB1", "CB2", "DST01", "DST02", "DST03", "DST04", "DST05", "OC1", "OC2", "SO4", "SSLT01", "SSLT02", "SSLT03", "SSLT04", "SSLT05"] +surface_setup: "DefaultMoninObukhov" +turbconv: "edonly_edmfx" +implicit_diffusion: true +ode_algo: "ARS111" +approximate_linear_solve_iters: 2 +max_newton_iters_ode: 10 +dt: "50secs" +t_end: "5hours" +toml: [toml/aquaplanet_nonequil_120_1200.toml] +diagnostics: + - short_name: [ta, ua, wa, va, rhoa, hur, hus, hussn, husra, clw, cli] + period: 5days +netcdf_output_at_levels: true \ No newline at end of file diff --git a/docs/src/api.md b/docs/src/api.md index cff24a866f..28f678d47f 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -46,6 +46,10 @@ ClimaAtmos.InitialConditions.Soares ### Implicit Solver ```@docs +ClimaAtmos.JacobianAlgorithm +ClimaAtmos.ApproxJacobian +ClimaAtmos.ExactJacobian +ClimaAtmos.DebugJacobian ClimaAtmos.ImplicitEquationJacobian ``` diff --git a/post_processing/jacobian_plots.jl b/post_processing/jacobian_plots.jl new file mode 100644 index 0000000000..072e764b37 --- /dev/null +++ b/post_processing/jacobian_plots.jl @@ -0,0 +1,375 @@ +import ClimaComms +import ClimaCore.InputOutput: HDF5Reader +import LinearAlgebra: diag + +function make_jacobian_plots(output_path, Yₜ_end, dt) + context = ClimaComms.context(Yₜ_end.c) + file_paths = filter(endswith(".hdf5"), readdir(output_path; join = true)) + for file_path in file_paths + file_name = split(basename(file_path), '.')[1] + contains(file_name, "jacobian") || continue + time_series = read_time_series(file_path, context) + plot_jacobian(output_path, time_series, file_name, Yₜ_end, dt) + if startswith(file_name, "approx_jacobian") + exact_file_path = joinpath( + output_path, + replace(basename(file_path), "approx" => "exact"), + ) + if exact_file_path in file_paths + exact_time_series = read_time_series(exact_file_path, context) + error_time_series = map( + time_series, + exact_time_series, + ) do approx_data, exact_data + (t, approx_title, approx_∂Yₜ_∂Y, approx_rescaled_∂Yₜ_∂Y, Yₜ) = approx_data + (_, _, exact_∂Yₜ_∂Y, exact_rescaled_∂Yₜ_∂Y, _) = + exact_data + @assert exact_data[1] == t + @assert exact_data[5] == Yₜ + error_title = "Error of $approx_title" + error_∂Yₜ_∂Y = approx_∂Yₜ_∂Y .- exact_∂Yₜ_∂Y + error_rescaled_∂Yₜ_∂Y = + approx_rescaled_∂Yₜ_∂Y .- exact_rescaled_∂Yₜ_∂Y + (t, error_title, error_∂Yₜ_∂Y, error_rescaled_∂Yₜ_∂Y, Yₜ) + end + error_file_name = + replace(file_name, "jacobian" => "jacobian_error") + plot_jacobian( + output_path, + error_time_series, + error_file_name, + Yₜ_end, + dt, + ) + end + end + end +end + +function read_time_series(file_path, context) + time_series = nothing + HDF5Reader(file_path, context) do reader + t_strings = + sort(keys(reader.file); by = t_string -> parse(Float64, t_string)) + time_series = map(t_strings) do t_string + data_at_t = reader.file[t_string] + title = data_at_t["title"][] + ∂Yₜ_∂Y = data_at_t["∂Yₜ_∂Y"][] + Yₜ = data_at_t["Yₜ"][] + + safe_inv(x) = iszero(x) || issubnormal(x) ? zero(x) : inv(x) + rescaled_∂Yₜ_∂Y = safe_inv.(Yₜ) .* ∂Yₜ_∂Y .* Yₜ' + + # Take the transpose of each matrix so that its rows are plotted + # along the y-axis and its columns are plotted along the x-axis. + (parse(Float64, t_string), title, ∂Yₜ_∂Y', rescaled_∂Yₜ_∂Y', Yₜ) + end + end + return time_series +end + +function plot_jacobian(output_path, time_series, file_name, Yₜ_end, dt) + times = getindex.(time_series, 1) + titles = getindex.(time_series, 2) + ∂Yₜ_∂Ys = getindex.(time_series, 3) + rescaled_∂Yₜ_∂Ys = getindex.(time_series, 4) + Yₜs = getindex.(time_series, 5) + FT = eltype(Yₜ_end) + + field_names = collect(CA.scalar_field_names(Yₜ_end)) + tick_labels = map(field_names) do field_name + replace( + string(field_name), + "@name(" => "", + ".components.data" => "", + ":(" => "", + ')' => "", + ) + end + + index_ranges = collect(CA.scalar_field_index_ranges(Yₜ_end)) + first_tick_positions = first.(index_ranges) + last_tick_positions = last.(index_ranges) + center_tick_positions = (first_tick_positions .+ last_tick_positions) ./ 2 + boundary_tick_positions = [FT(0.5), (last_tick_positions .+ FT(0.5))...] + limit_padding = length(index_ranges[1]) / 20 + limit_positions = + extrema(boundary_tick_positions) .+ (-limit_padding, limit_padding) + + sign_or_nan(x) = iszero(x) ? FT(NaN) : sign(x) + logabs_or_nan(x) = iszero(x) ? FT(NaN) : log10(abs(x)) + function block_bandwidth(block) + band_indices = (1 - size(block, 1)):(size(block, 2) - 1) + nonempty_band_indices = filter(band_indices) do band_index + any(!iszero, diag(block, band_index)) + end + main_diagonal_index = (band_indices[1] + band_indices[end]) / 2 + return maximum(nonempty_band_indices; init = 0) do band_index + 2 * abs(band_index - main_diagonal_index) + 1 + end + end + + entry_sign_transform(matrix) = sign_or_nan.(matrix) + entry_logabs_transform(matrix) = logabs_or_nan.(matrix) + function block_logabs_transform(matrix) + block_max_matrix = similar(matrix) + for row_index_range in index_ranges, col_index_range in index_ranges + block_max_matrix[row_index_range, col_index_range] .= + CA.smooth_maximum(matrix[row_index_range, col_index_range]) + end + return entry_logabs_transform(block_max_matrix) + end + function block_row_logabs_transform(matrix) + row_max_matrix = similar(matrix) + for row_index_range in index_ranges, col in axes(matrix, 1) + row_max_matrix[row_index_range, col] .= + CA.smooth_maximum(matrix[row_index_range, col]) + end + return entry_logabs_transform(row_max_matrix) + end + function block_bandwidth_transform(matrix) + block_bandwidth_matrix = similar(matrix, Int) + for row_index_range in index_ranges, col_index_range in index_ranges + block_bandwidth_matrix[row_index_range, col_index_range] .= + block_bandwidth(matrix[row_index_range, col_index_range]) + end + return block_bandwidth_matrix + end + + bandwidth_matrices = map(block_bandwidth_transform, ∂Yₜ_∂Ys) + max_bandwidth = maximum(maximum, bandwidth_matrices) + sign_matrices = map(entry_sign_transform, ∂Yₜ_∂Ys) + + categorical(colormap, n) = CairoMakie.cgrad(colormap, n; categorical = true) + main_colormap = categorical(:tol_iridescent, 21) + bandwidth_colors = [CairoMakie.RGB(1, 1, 1), main_colormap[1:(end - 1)]...] + bandwidth_colormap = + categorical(bandwidth_colors, min(max_bandwidth, 9) + 1) + sign_colormap = categorical(:RdBu_5, 2) + rescaling_colors = + setindex!(CairoMakie.to_colormap(:RdBu_11), CairoMakie.RGB(1, 1, 1), 6) + rescaling_colormap = categorical(rescaling_colors, 21) + + value_for_min(x) = isnan(x) ? FT(Inf) : x + value_for_max(x) = isnan(x) ? FT(-Inf) : x + + abs_Yₜs = map(Yₜs) do Yₜ + @. ifelse(iszero(Yₜ) || issubnormal(Yₜ), FT(NaN), abs(Yₜ)) + end + if all(abs_Yₜ -> all(isnan, abs_Yₜ), abs_Yₜs) + min_abs_Yₜ = FT(1e-9) + max_abs_Yₜ = FT(1) + else + min_abs_Yₜ = minimum(abs_Yₜ -> minimum(value_for_min, abs_Yₜ), abs_Yₜs) + max_abs_Yₜ = maximum(abs_Yₜ -> maximum(value_for_max, abs_Yₜ), abs_Yₜs) + end + + figure_kwargs = (; + size = (2.2, cld(length(times), 2)) .* 5000, + figure_padding = 300, + fontsize = 150, + ) + colorbar_kwargs = (; + size = 150, + labelpadding = 60, + ticklabelpad = 50, + ticksize = 50, + tickwidth = 10, + spinewidth = 10, + ) + axis_kwargs = (; + titlegap = 100, + xlabelpadding = 60, + ylabelpadding = 60, + xticks = (center_tick_positions, tick_labels), + xticksvisible = false, + xticklabelrotation = pi / 4, + xticklabelpad = 50, + yticklabelpad = 80, + xminorticks = boundary_tick_positions, + xminorticksvisible = true, + xminorticksize = 50, + yminorticksize = 50, + xminortickwidth = 10, + yminortickwidth = 10, + xgridvisible = false, + xminorgridvisible = true, + xminorgridwidth = 10, + yminorgridwidth = 10, + spinewidth = 10, + ) + ∂Yₜ_∂Y_axis_kwargs = (; + axis_kwargs..., + limits = (limit_positions, limit_positions), + xlabel = "Y index", + ylabel = "Yₜ index", + yreversed = true, + yticks = (center_tick_positions, tick_labels), + yticksvisible = false, + yminorticks = boundary_tick_positions, + yminorticksvisible = true, + ygridvisible = false, + yminorgridvisible = true, + ) # Flip the y-axis so that the diagonal runs from top-left to bottom-right. + Yₜ_axis_kwargs = (; + axis_kwargs..., + limits = (limit_positions, (min_abs_Yₜ / 2, max_abs_Yₜ * 2)), + xlabel = "Yₜ index", + ylabel = "Yₜ magnitude", + yscale = log10, + yticksize = 50, + ) + + page_file_paths = map(n -> joinpath(output_path, "$file_name $n.pdf"), 1:6) + full_file_path = joinpath(output_path, "$file_name.pdf") + + figure = CairoMakie.Figure(; figure_kwargs...) + for (t_index, (t, title, bandwidth_matrix, sign_matrix)) in + enumerate(zip(times, titles, bandwidth_matrices, sign_matrices)) + grid_position = figure[cld(t_index, 2), (t_index - 1) % 2 + 1] + grid_layout = CairoMakie.GridLayout(grid_position) + axis = CairoMakie.Axis(grid_layout[1, 1]; title, ∂Yₜ_∂Y_axis_kwargs...) + bandwidth_plot = CairoMakie.heatmap!( + axis, + boundary_tick_positions, + boundary_tick_positions, + bandwidth_matrix[first_tick_positions, first_tick_positions]; + colormap = bandwidth_colormap, + colorrange = (-0.5, min(max_bandwidth, 9) + 0.5), + (max_bandwidth > 9 ? (; highclip = :black) : (;))..., + ) + CairoMakie.translate!(bandwidth_plot, 0, 0, -100) # Put plot under grid. + sign_plot = CairoMakie.heatmap!( + axis, + sign_matrix; + colormap = sign_colormap, + colorrange = (-1, 1), + ) + grid_sublayout = CairoMakie.GridLayout(grid_layout[1, 2]) + CairoMakie.Colorbar( + grid_sublayout[1, 1], + bandwidth_plot; + label = "Bandwidth of matrix block", + colorbar_kwargs..., + (max_bandwidth < 4 ? (; ticks = [(0:max_bandwidth)...]) : (;))..., + ) + CairoMakie.Colorbar( + grid_sublayout[2, 1], + sign_plot; + label = "Sign of matrix entry", + ticks = ([-0.5, 0.5], Makie.latexstring.(["-", "+"])), + colorbar_kwargs..., + labelpadding = 27, + ) + CairoMakie.colgap!(grid_layout, CairoMakie.Relative(0.05)) + CairoMakie.rowgap!(grid_sublayout, CairoMakie.Relative(0.05)) + end + CairoMakie.colsize!(figure.layout, 1, CairoMakie.Aspect(1, 1)) + CairoMakie.colsize!(figure.layout, 2, CairoMakie.Aspect(1, 1)) + CairoMakie.colgap!(figure.layout, CairoMakie.Relative(0.08)) + CairoMakie.rowgap!(figure.layout, CairoMakie.Relative(0.05)) + CairoMakie.save(page_file_paths[1], figure) + + for (page_index, (is_rescaled, transform, transform_string)) in enumerate(( + (true, block_logabs_transform, "block"), + (true, block_row_logabs_transform, "block row"), + (true, entry_logabs_transform, "entry"), + (false, entry_logabs_transform, "entry"), + )) + matrices = map(transform, is_rescaled ? rescaled_∂Yₜ_∂Ys : ∂Yₜ_∂Ys) + max_logabs = + all(matrix -> all(isnan, matrix), matrices[2:end]) ? FT(0) : + maximum(matrix -> maximum(value_for_max, matrix), matrices[2:end]) + if is_rescaled + dt_tick_value = log10(1 / dt) + min_logabs = min(max_logabs - 9, dt_tick_value - FT(0.5)) + else + min_logabs = max_logabs - 25 + end + colorbar_range = (min_logabs, max_logabs) + wilkinson_kwargs = (; k_min = 4, k_max = 6, niceness_weight = 1) + colorbar_ticks = Makie.WilkinsonTicks(5; wilkinson_kwargs...) + colorbar_tick_values = + Makie.get_tickvalues(colorbar_ticks, colorbar_range...) + colorbar_tick_labels = map(colorbar_tick_values) do tick_value + Makie.latexstring("10^{$(round(Int, tick_value))}") + end + if is_rescaled + dt_tick_label = Makie.latexstring("\\mathbf{Δt^{-1}}") + min_tick_value_distance, colorbar_tick_index = + findmin(colorbar_tick_values) do colorbar_tick_value + abs(colorbar_tick_value - dt_tick_value) + end + if min_tick_value_distance < 0.5 + colorbar_tick_values[colorbar_tick_index] = dt_tick_value + colorbar_tick_labels[colorbar_tick_index] = dt_tick_label + else + push!(colorbar_tick_values, dt_tick_value) + push!(colorbar_tick_labels, dt_tick_label) + end + end + figure = CairoMakie.Figure(; figure_kwargs...) + for (t_index, (t, title, matrix)) in + enumerate(zip(times, titles, matrices)) + grid_position = figure[cld(t_index, 2), (t_index - 1) % 2 + 1] + grid_layout = CairoMakie.GridLayout(grid_position) + axis = + CairoMakie.Axis(grid_layout[1, 1]; title, ∂Yₜ_∂Y_axis_kwargs...) + plot_args = if transform == block_logabs_transform + submatrix = matrix[first_tick_positions, first_tick_positions] + (boundary_tick_positions, boundary_tick_positions, submatrix) + elseif transform == block_row_logabs_transform + submatrix = matrix[first_tick_positions, :] + (boundary_tick_positions, 1:last_tick_positions[end], submatrix) + else + (matrix,) + end + plot = CairoMakie.heatmap!( + axis, + plot_args...; + colormap = main_colormap, + colorrange = colorbar_range, + lowclip = main_colormap[1], + highclip = sign_colormap[1], + ) # Only color in the top 9 orders of magnitude for all t > 0. + CairoMakie.translate!(plot, 0, 0, -100) # Put plot under grid. + label = "Magnitude of$(is_rescaled ? " rescaled" : "") \ + $(transform_string != "entry" ? "max of " : "")matrix \ + $transform_string$(is_rescaled ? " [s⁻¹]" : "")" + CairoMakie.Colorbar( + grid_layout[1, 2], + plot; + label, + ticks = (colorbar_tick_values, colorbar_tick_labels), + colorbar_kwargs..., + ) + CairoMakie.colgap!(grid_layout, CairoMakie.Relative(0.05)) + end + CairoMakie.colsize!(figure.layout, 1, CairoMakie.Aspect(1, 1)) + CairoMakie.colsize!(figure.layout, 2, CairoMakie.Aspect(1, 1)) + CairoMakie.colgap!(figure.layout, CairoMakie.Relative(0.08)) + CairoMakie.rowgap!(figure.layout, CairoMakie.Relative(0.05)) + CairoMakie.save(page_file_paths[page_index + 1], figure) + end + + figure = CairoMakie.Figure(; figure_kwargs...) + for (t_index, (t, title, abs_Yₜ)) in enumerate(zip(times, titles, abs_Yₜs)) + grid_position = figure[cld(t_index, 2), (t_index - 1) % 2 + 1] + Yₜ_title = "Yₜ, avg over all columns," * split(title, ',')[end] + axis = + CairoMakie.Axis(grid_position; title = Yₜ_title, Yₜ_axis_kwargs...) + CairoMakie.lines!(axis, axes(abs_Yₜ, 1), abs_Yₜ; linewidth = 10) + end + CairoMakie.colsize!(figure.layout, 1, CairoMakie.Aspect(1, 1)) + CairoMakie.colsize!(figure.layout, 2, CairoMakie.Aspect(1, 1)) + CairoMakie.colgap!(figure.layout, CairoMakie.Relative(0.12)) + CairoMakie.rowgap!(figure.layout, CairoMakie.Relative(0.05)) + CairoMakie.save(page_file_paths[end], figure) + + pdfunite() do unite + run(Cmd([unite, page_file_paths..., full_file_path])) + end + for page_file_path in page_file_paths + Filesystem.rm(page_file_path; force = true) + end +end diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 749be06168..b7653264a0 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -2,6 +2,7 @@ module ClimaAtmos using NVTX import Adapt +import LinearAlgebra import NullBroadcasts: NullBroadcasted import LazyBroadcast import LazyBroadcast: lazy @@ -56,6 +57,10 @@ include(joinpath("prognostic_equations", "zero_velocity.jl")) include(joinpath("prognostic_equations", "implicit", "implicit_tendency.jl")) include(joinpath("prognostic_equations", "implicit", "implicit_solver.jl")) +include(joinpath("prognostic_equations", "implicit", "approx_jacobian.jl")) +include(joinpath("prognostic_equations", "implicit", "exact_jacobian.jl")) +include(joinpath("prognostic_equations", "implicit", "debug_jacobian.jl")) +include(joinpath("prognostic_equations", "implicit", "dual_fixes.jl")) include(joinpath("prognostic_equations", "water_advection.jl")) include(joinpath("prognostic_equations", "remaining_tendency.jl")) diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl index e5ee222d8f..0f19c06aee 100644 --- a/src/callbacks/callbacks.jl +++ b/src/callbacks/callbacks.jl @@ -551,6 +551,56 @@ function gc_func(integrator) return nothing end +# TODO: Move these functions to ClimaTimeSteppers.jl +is_implicit_timestepper(integrator) = + integrator.alg isa CTS.RosenbrockAlgorithm || ( + integrator.alg.name isa CTS.IMEXARKAlgorithmName && + !isnothing(integrator.cache.newtons_method_cache) + ) +jacobian(integrator) = + if integrator.alg isa CTS.RosenbrockAlgorithm + integrator.cache.W + elseif ( + integrator.alg.name isa CTS.IMEXARKAlgorithmName && + !isnothing(integrator.cache.newtons_method_cache) + ) + integrator.cache.newtons_method_cache.j + else + error("Jacobian is not available") + end +function dtγ(integrator) + is_implicit_timestepper(integrator) || error("dtγ is not available") + (; dt, alg) = integrator + tableau_coefficients = + alg isa CTS.RosenbrockAlgorithm ? alg.tableau.Γ : alg.tableau.a_imp + γs = unique(filter(!iszero, LinearAlgebra.diag(tableau_coefficients))) + length(γs) == 1 || error( + "The exact Jacobian must be updated on every Newton iteration, rather \ + than on every timestep (or every N steps), because the specified IMEX \ + algorithm has implicit stages with distinct tableau coefficients \ + (i.e., it is not an SDIRK algorithm).", + ) + return float(dt) * γs[1] +end + +function update_jacobian!(integrator) + is_implicit_timestepper(integrator) || return # TODO: Remove this check. + (; u, p, t) = integrator + update_jacobian!(jacobian(integrator), u, p, dtγ(integrator), t) +end + +function update_and_check_jacobian!(integrator) + is_implicit_timestepper(integrator) || return # TODO: Remove this check. + (; u, p, t) = integrator + update_and_check_jacobian!(jacobian(integrator), u, p, dtγ(integrator), t) +end + +function save_jacobian!(integrator) + is_implicit_timestepper(integrator) || return # TODO: Remove this check. + (; u, p, t) = integrator + save_jacobian!(jacobian(integrator), u, p, dtγ(integrator), t) +end + """ maybe_graceful_exit(integrator) diff --git a/src/callbacks/get_callbacks.jl b/src/callbacks/get_callbacks.jl index e389b8e2cd..1a5a411e37 100644 --- a/src/callbacks/get_callbacks.jl +++ b/src/callbacks/get_callbacks.jl @@ -188,7 +188,6 @@ function get_diagnostics(parsed_args, atmos_model, Y, p, sim_info, output_dir) diagnostics..., ] end - diagnostics = collect(diagnostics) periods_reductions = Set() for diag in diagnostics @@ -381,5 +380,48 @@ function get_callbacks(config, sim_info, atmos, params, Y, p) callbacks = (callbacks..., call_every_dt(nogw_model_callback!, dt_nogw)) end + use_exact_jacobian = parsed_args["use_exact_jacobian"] + debug_approximate_jacobian = parsed_args["debug_approximate_jacobian"] + need_to_update_exact_jacobian_in_callback = + ( + (debug_approximate_jacobian && !use_exact_jacobian) && + isnothing(parsed_args["dt_update_exact_jacobian"]) + ) || ( + (debug_approximate_jacobian || use_exact_jacobian) && + !isnothing(parsed_args["dt_update_exact_jacobian"]) && + time_to_seconds(parsed_args["dt_update_exact_jacobian"]) >= + time_to_seconds(parsed_args["dt"]) + ) + need_to_save_jacobian = + isnothing(parsed_args["plot_jacobian"]) ? debug_approximate_jacobian : + parsed_args["plot_jacobian"] + dt_save_jacobian = if isnothing(parsed_args["dt_save_jacobian"]) + (t_end - t_start) ÷ 5 + else + dt isa ITime ? + ITime(time_to_seconds(parsed_args["dt_save_jacobian"])) : + FT(time_to_seconds(parsed_args["dt_save_jacobian"])) + end + dt_update_jacobian = if isnothing(parsed_args["dt_update_exact_jacobian"]) + dt_save_jacobian + else + dt isa ITime ? + ITime(time_to_seconds(parsed_args["dt_update_exact_jacobian"])) : + FT(time_to_seconds(parsed_args["dt_update_exact_jacobian"])) + end + if need_to_update_exact_jacobian_in_callback + update_jacobian_in_callback! = + debug_approximate_jacobian ? update_and_check_jacobian! : + update_jacobian! + callbacks = ( + callbacks..., + call_every_dt(update_jacobian_in_callback!, dt_update_jacobian), + ) + end + if need_to_save_jacobian + callbacks = + (callbacks..., call_every_dt(save_jacobian!, dt_save_jacobian)) + end + return callbacks end diff --git a/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl index 76d8181401..8cbb8092b2 100644 --- a/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl +++ b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl @@ -86,11 +86,12 @@ function cloud_sources( S = CMNe.conv_q_vap_to_q_liq_ice_MM2015(cm_params, thp, q, ρ, Tₐ(thp, ts)) # keeping the same limiter for now - return ifelse( - S > FT(0), - min(S, limit(qᵥ(thp, ts), dt, 2)), - -min(abs(S), limit(qₗ(thp, ts, qₚ(qᵣ)), dt, 2)), - ) + #return ifelse( + # S > FT(0), + # min(S, limit(qᵥ(thp, ts), dt, 2)), + # -min(abs(S), limit(qₗ(thp, ts, qₚ(qᵣ)), dt, 2)), + #) + return S end function cloud_sources(cm_params::CMP.CloudIce{FT}, thp, ts, qₛ, dt) where {FT} @@ -100,11 +101,12 @@ function cloud_sources(cm_params::CMP.CloudIce{FT}, thp, ts, qₛ, dt) where {FT S = CMNe.conv_q_vap_to_q_liq_ice_MM2015(cm_params, thp, q, ρ, Tₐ(thp, ts)) # keeping the same limiter for now - return ifelse( - S > FT(0), - min(S, limit(qᵥ(thp, ts), dt, 2)), - -min(abs(S), limit(qᵢ(thp, ts, qₚ(qₛ)), dt, 2)), - ) + #return ifelse( + # S > FT(0), + # min(S, limit(qᵥ(thp, ts), dt, 2)), + # -min(abs(S), limit(qᵢ(thp, ts, qₚ(qₛ)), dt, 2)), + #) + return S end """ diff --git a/src/prognostic_equations/edmfx_closures.jl b/src/prognostic_equations/edmfx_closures.jl index 6b4c1421f3..5ef1212358 100644 --- a/src/prognostic_equations/edmfx_closures.jl +++ b/src/prognostic_equations/edmfx_closures.jl @@ -225,7 +225,7 @@ function mixing_length( l_smin = lamb_smooth_minimum(l, smin_ub, smin_rm) l_limited = max(l_smag, min(l_smin, l_z)) - return MixingLength{FT}(l_limited, l_W, l_TKE, l_N) + return MixingLength(l_limited, l_W, l_TKE, l_N) end """ diff --git a/src/prognostic_equations/gm_sgs_closures.jl b/src/prognostic_equations/gm_sgs_closures.jl index d1cecfcf74..30105cf122 100644 --- a/src/prognostic_equations/gm_sgs_closures.jl +++ b/src/prognostic_equations/gm_sgs_closures.jl @@ -19,7 +19,7 @@ function smagorinsky_lilly_length(c_smag, N_eff, dz, Pr, ϵ_st) return N_eff > FT(0) ? c_smag * dz * - max(0, 1 - N_eff^2 / Pr / 2 / max(ϵ_st, eps(FT)))^(1 / 4) : + max(0, 1 - N_eff^2 / Pr / 2 / max(ϵ_st, eps(FT)))^(1 / FT(4)) : c_smag * dz end diff --git a/src/prognostic_equations/implicit/approx_jacobian.jl b/src/prognostic_equations/implicit/approx_jacobian.jl new file mode 100644 index 0000000000..bca7b4eb78 --- /dev/null +++ b/src/prognostic_equations/implicit/approx_jacobian.jl @@ -0,0 +1,1396 @@ +abstract type DerivativeFlag end +struct UseDerivative <: DerivativeFlag end +struct IgnoreDerivative <: DerivativeFlag end + +DerivativeFlag(value) = value ? UseDerivative() : IgnoreDerivative() +DerivativeFlag(mode::AbstractTimesteppingMode) = + DerivativeFlag(mode == Implicit()) + +use_derivative(::UseDerivative) = true +use_derivative(::IgnoreDerivative) = false + +""" + ApproxJacobian( + topography_flag, + diffusion_flag, + sgs_advection_flag, + sgs_entr_detr_flag, + sgs_mass_flux_flag, + sgs_nh_pressure_flag, + noneq_cloud_formation_flag, + approximate_solve_iters, + ) + +A `JacobianAlgorithm` that approximates the `ImplicitEquationJacobian` using +analytically derived tendency derivatives and inverts it using a specialized +nested linear solver. Certain groups of derivatives can be toggled on or off by +setting their `DerivativeFlag`s to either `UseDerivative` or `IgnoreDerivative`. + +# Arguments + +- `topography_flag::DerivativeFlag`: whether the derivative of vertical + contravariant velocity with respect to horizontal covariant velocity should be + computed +- `diffusion_flag::DerivativeFlag`: whether the derivatives of the grid-scale + diffusion tendency should be computed +- `sgs_advection_flag::DerivativeFlag`: whether the derivatives of the + subgrid-scale advection tendency should be computed +- `sgs_entr_detr_flag::DerivativeFlag`: whether the derivatives of the + subgrid-scale entrainment and detrainment tendencies should be computed +- `sgs_mass_flux_flag::DerivativeFlag`: whether the derivatives of the + subgrid-scale mass flux tendency should be computed +- `sgs_nh_pressure_flag::DerivativeFlag`: whether the derivatives of the + subgrid-scale non-hydrostatic pressure drag tendency should be computed +- `noneq_cloud_formation_flag::DerivativeFlag`: whether the derivatives + of the nonequilibrium cloud formation sources should be computed +- `approximate_solve_iters::Int`: number of iterations to take for the + approximate linear solve required when the `diffusion_flag` is `UseDerivative` +""" +struct ApproxJacobian{F1, F2, F3, F4, F5, F6, F7} <: JacobianAlgorithm + topography_flag::F1 + diffusion_flag::F2 + sgs_advection_flag::F3 + sgs_entr_detr_flag::F4 + sgs_mass_flux_flag::F5 + sgs_nh_pressure_flag::F6 + noneq_cloud_formation_flag::F7 + approximate_solve_iters::Int +end + +function jacobian_cache(alg::ApproxJacobian, Y, atmos) + (; + topography_flag, + diffusion_flag, + sgs_advection_flag, + sgs_mass_flux_flag, + noneq_cloud_formation_flag, + ) = alg + FT = Spaces.undertype(axes(Y.c)) + CTh = CTh_vector_type(axes(Y.c)) + + DiagonalRow = DiagonalMatrixRow{FT} + TridiagonalRow = TridiagonalMatrixRow{FT} + BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} + TridiagonalRow_ACTh = TridiagonalMatrixRow{Adjoint{FT, CTh{FT}}} + BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} + BidiagonalRow_C3xACTh = + BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CTh{FT})')} + DiagonalRow_C3xACT3 = + DiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} + TridiagonalRow_C3xACT3 = + TridiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} + + is_in_Y(name) = MatrixFields.has_field(Y, name) + + ρq_tot_if_available = is_in_Y(@name(c.ρq_tot)) ? (@name(c.ρq_tot),) : () + ρatke_if_available = + is_in_Y(@name(c.sgs⁰.ρatke)) ? (@name(c.sgs⁰.ρatke),) : () + sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : () + + condensate_names = + (@name(c.ρq_liq), @name(c.ρq_ice), @name(c.ρq_rai), @name(c.ρq_sno)) + available_condensate_names = + MatrixFields.unrolled_filter(is_in_Y, condensate_names) + available_tracer_names = + (ρq_tot_if_available..., available_condensate_names...) + + sgs_tracer_names = ( + @name(c.sgsʲs.:(1).q_tot), + @name(c.sgsʲs.:(1).q_liq), + @name(c.sgsʲs.:(1).q_ice), + @name(c.sgsʲs.:(1).q_rai), + @name(c.sgsʲs.:(1).q_sno), + ) + available_sgs_tracer_names = + MatrixFields.unrolled_filter(is_in_Y, sgs_tracer_names) + + sgs_scalar_names = + (sgs_tracer_names..., @name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).ρa)) + available_sgs_scalar_names = + MatrixFields.unrolled_filter(is_in_Y, sgs_scalar_names) + + sgs_u³_if_available = + is_in_Y(@name(f.sgsʲs.:(1).u₃)) ? (@name(f.sgsʲs.:(1).u₃),) : () + + # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0, + # which means that multiplying inv(-1) by a Float32 will yield a Float64. + identity_blocks = MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + (@name(c.ρ), sfc_if_available...), + ) + + active_scalar_names = (@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...) + advection_blocks = ( + ( + use_derivative(topography_flag) ? + MatrixFields.unrolled_map( + name -> + (name, @name(c.uₕ)) => + similar(Y.c, TridiagonalRow_ACTh), + active_scalar_names, + ) : () + )..., + MatrixFields.unrolled_map( + name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), + active_scalar_names, + )..., + MatrixFields.unrolled_map( + name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3), + active_scalar_names, + )..., + (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACTh), + (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), + ) + + condensate_blocks = + if atmos.moisture_model isa NonEquilMoistModel && + use_derivative(noneq_cloud_formation_flag) + ( + (@name(c.ρq_liq), @name(c.ρq_tot)) => similar(Y.c, DiagonalRow), + (@name(c.ρq_ice), @name(c.ρq_tot)) => similar(Y.c, DiagonalRow), + ) + else + () + end + + diffused_scalar_names = (@name(c.ρe_tot), available_tracer_names...) + diffusion_blocks = if use_derivative(diffusion_flag) + ( + MatrixFields.unrolled_map( + name -> (name, @name(c.ρ)) => similar(Y.c, TridiagonalRow), + (diffused_scalar_names..., ρatke_if_available...), + )..., + MatrixFields.unrolled_map( + name -> (name, name) => similar(Y.c, TridiagonalRow), + (diffused_scalar_names..., ρatke_if_available...), + )..., + ( + is_in_Y(@name(c.ρq_tot)) ? + ( + (@name(c.ρe_tot), @name(c.ρq_tot)) => + similar(Y.c, TridiagonalRow), + ) : () + )..., + (@name(c.uₕ), @name(c.uₕ)) => + !isnothing(atmos.turbconv_model) || + !disable_momentum_vertical_diffusion(atmos.vert_diff) ? + similar(Y.c, TridiagonalRow) : FT(-1) * I, + ) + elseif atmos.moisture_model isa DryModel + MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + (diffused_scalar_names..., ρatke_if_available..., @name(c.uₕ)), + ) + else + ( + MatrixFields.unrolled_map( + name -> (name, name) => similar(Y.c, TridiagonalRow), + diffused_scalar_names, + )..., + (@name(c.ρe_tot), @name(c.ρq_tot)) => + similar(Y.c, TridiagonalRow), + MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + (ρatke_if_available..., @name(c.uₕ)), + )..., + ) + + end + + sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX + @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1 + if use_derivative(sgs_advection_flag) + ( + MatrixFields.unrolled_map( + name -> (name, name) => similar(Y.c, TridiagonalRow), + available_sgs_scalar_names, + )..., + (@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)) => + similar(Y.c, DiagonalRow), + (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)) => + similar(Y.c, TridiagonalRow), + (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)) => + similar(Y.c, TridiagonalRow), + (@name(c.sgsʲs.:(1).ρa), @name(f.sgsʲs.:(1).u₃)) => + similar(Y.c, BidiagonalRow_ACT3), + (@name(c.sgsʲs.:(1).mse), @name(f.sgsʲs.:(1).u₃)) => + similar(Y.c, BidiagonalRow_ACT3), + (@name(c.sgsʲs.:(1).q_tot), @name(f.sgsʲs.:(1).u₃)) => + similar(Y.c, BidiagonalRow_ACT3), + (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)) => + similar(Y.f, BidiagonalRow_C3), + (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)) => + similar(Y.f, BidiagonalRow_C3), + (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => + similar(Y.f, TridiagonalRow_C3xACT3), + ) + else + ( + MatrixFields.unrolled_map( + name -> (name, name) => FT(-1) * I, + available_sgs_scalar_names, + )..., + (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => + !isnothing(atmos.rayleigh_sponge) ? + similar(Y.f, DiagonalRow_C3xACT3) : FT(-1) * I, + ) + end + else + () + end + + sgs_massflux_blocks = if atmos.turbconv_model isa PrognosticEDMFX + @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1 + if use_derivative(sgs_mass_flux_flag) + ( + (@name(c.ρe_tot), @name(c.sgsʲs.:(1).mse)) => + similar(Y.c, TridiagonalRow), + (@name(c.ρq_tot), @name(c.sgsʲs.:(1).q_tot)) => + similar(Y.c, TridiagonalRow), + (@name(c.ρe_tot), @name(f.sgsʲs.:(1).u₃)) => + similar(Y.c, BidiagonalRow_ACT3), + (@name(c.ρq_tot), @name(f.sgsʲs.:(1).u₃)) => + similar(Y.c, BidiagonalRow_ACT3), + (@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρa)) => + similar(Y.c, TridiagonalRow), + (@name(c.ρq_tot), @name(c.sgsʲs.:(1).ρa)) => + similar(Y.c, TridiagonalRow), + ) + else + () + end + else + () + end + + matrix = MatrixFields.FieldMatrix( + identity_blocks..., + sgs_advection_blocks..., + advection_blocks..., + condensate_blocks..., + diffusion_blocks..., + sgs_massflux_blocks..., + ) + + mass_and_surface_names = (@name(c.ρ), sfc_if_available...) + available_scalar_names = ( + mass_and_surface_names..., + ρq_tot_if_available..., + @name(c.ρe_tot), + available_condensate_names..., + ρatke_if_available..., + available_sgs_scalar_names..., + ) + + velocity_alg = MatrixFields.BlockLowerTriangularSolve( + @name(c.uₕ), + sgs_u³_if_available..., + ) + full_alg = + if use_derivative(diffusion_flag) || + use_derivative(sgs_advection_flag) || + !(atmos.moisture_model isa DryModel) + gs_scalar_subalg = if !(atmos.moisture_model isa DryModel) + MatrixFields.BlockLowerTriangularSolve(@name(c.ρq_tot)) + else + MatrixFields.BlockDiagonalSolve() + end + scalar_subalg = + if atmos.turbconv_model isa PrognosticEDMFX && + use_derivative(sgs_advection_flag) + MatrixFields.BlockLowerTriangularSolve( + available_sgs_tracer_names...; + alg₂ = MatrixFields.BlockLowerTriangularSolve( + @name(c.sgsʲs.:(1).mse); + alg₂ = MatrixFields.BlockLowerTriangularSolve( + @name(c.sgsʲs.:(1).ρa); + alg₂ = gs_scalar_subalg, + ), + ), + ) + else + gs_scalar_subalg + end + scalar_alg = MatrixFields.BlockLowerTriangularSolve( + mass_and_surface_names...; + alg₂ = scalar_subalg, + ) + MatrixFields.ApproximateBlockArrowheadIterativeSolve( + available_scalar_names...; + alg₁ = scalar_alg, + alg₂ = velocity_alg, + P_alg₁ = MatrixFields.MainDiagonalPreconditioner(), + n_iters = alg.approximate_solve_iters, + ) + else + MatrixFields.BlockArrowheadSolve( + available_scalar_names...; + alg₂ = velocity_alg, + ) + end + + temp_matrix = (matrix .+ identity_matrix(matrix, Y)) ./ FT(1) + temp_matrix_column = similar(first_column(temp_matrix)) + + return (; + matrix = MatrixFields.FieldMatrixWithSolver(matrix, Y, full_alg), + temp_matrix, + temp_matrix_column, + ) +end + +# TODO: Replace some scalar matrix entries with tensor entries so that we can +# use MatrixFields.identity_field_matrix(Y) instead of identity_matrix(Y). +function identity_matrix(matrix, Y) + I_matrix = MatrixFields.identity_field_matrix(Y) + new_pairs = MatrixFields.unrolled_map(pairs(I_matrix)) do (key, value) + replace_tensor_value_with_scalar_value = + key == (@name(c.uₕ), @name(c.uₕ)) || ( + key == (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) && + matrix[key] isa LinearAlgebra.UniformScaling + ) + key => (replace_tensor_value_with_scalar_value ? I : value) + end + return MatrixFields.replace_name_tree( + MatrixFields.FieldMatrix(new_pairs...), + MatrixFields.FieldNameTree(Y), + ) +end + +function update_jacobian!(alg::ApproxJacobian, cache, Y, p, dtγ, t) + (; + topography_flag, + diffusion_flag, + sgs_advection_flag, + sgs_entr_detr_flag, + sgs_mass_flux_flag, + sgs_nh_pressure_flag, + noneq_cloud_formation_flag, + ) = alg + (; matrix) = cache + (; params, dt) = p + (; ᶜΦ, ᶠgradᵥ_ᶜΦ) = p.core + (; ᶜspecific, ᶠu³, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed + (; + ∂ᶜK_∂ᶜuₕ, + ∂ᶜK_∂ᶠu₃, + ᶠp_grad_matrix, + ᶜadvection_matrix, + ᶜdiffusion_h_matrix, + ᶜdiffusion_h_matrix_scaled, + ᶜdiffusion_u_matrix, + ᶠbidiagonal_matrix_ct3, + ᶠbidiagonal_matrix_ct3_2, + ᶠtridiagonal_matrix_c3, + ) = p.scratch + rs = p.atmos.rayleigh_sponge + + FT = Spaces.undertype(axes(Y.c)) + CTh = CTh_vector_type(axes(Y.c)) + one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' + + cv_d = FT(CAP.cv_d(params)) + Δcv_v = FT(CAP.cv_v(params)) - cv_d + T_0 = FT(CAP.T_0(params)) + R_d = FT(CAP.R_d(params)) + ΔR_v = FT(CAP.R_v(params)) - R_d + cp_d = FT(CAP.cp_d(params)) + Δcp_v = FT(CAP.cp_v(params)) - cp_d + # This term appears a few times in the Jacobian, and is technically + # minus ∂e_int_∂q_tot + ∂e_int_∂q_tot = T_0 * (Δcv_v - R_d) - FT(CAP.e_int_v0(params)) + thermo_params = CAP.thermodynamics_params(params) + + ᶜρ = Y.c.ρ + ᶜuₕ = Y.c.uₕ + ᶠu₃ = Y.f.u₃ + ᶜJ = Fields.local_geometry_field(Y.c).J + ᶠJ = Fields.local_geometry_field(Y.f).J + ᶜgⁱʲ = Fields.local_geometry_field(Y.c).gⁱʲ + ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ + ᶠz = Fields.coordinate_field(Y.f).z + zmax = z_max(axes(Y.f)) + + ᶜkappa_m = p.scratch.ᶜtemp_scalar + @. ᶜkappa_m = + TD.gas_constant_air(thermo_params, ᶜts) / TD.cv_m(thermo_params, ᶜts) + + ᶜ∂kappa_m∂q_tot = p.scratch.ᶜtemp_scalar_2 + # Using abs2 because ^2 results in allocation + @. ᶜ∂kappa_m∂q_tot = + ( + ΔR_v * TD.cv_m(thermo_params, ᶜts) - + Δcv_v * TD.gas_constant_air(thermo_params, ᶜts) + ) / abs2(TD.cv_m(thermo_params, ᶜts)) + + if use_derivative(topography_flag) + @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow( + adjoint(CTh(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ), + ) + else + @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ))) + end + @. ∂ᶜK_∂ᶠu₃ = + ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) + + DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix() + + @. ᶠp_grad_matrix = DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() + + @. ᶜadvection_matrix = + -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) + + if use_derivative(topography_flag) + ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)] + @. ∂ᶜρ_err_∂ᶜuₕ = + dtγ * ᶜadvection_matrix ⋅ ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ + DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) + end + ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] + @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) + + tracer_info = ( + (@name(c.ρe_tot), @name(ᶜh_tot)), + (@name(c.ρq_tot), @name(ᶜspecific.q_tot)), + ) + MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name) + MatrixFields.has_field(Y, ρχ_name) || return + ᶜχ = MatrixFields.get_field(p.precomputed, χ_name) + if use_derivative(topography_flag) + ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] + end + ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)] + use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ = + dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ)) ⋅ + ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) + @. ∂ᶜρχ_err_∂ᶠu₃ = + dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ) * g³³(ᶠgⁱʲ)) + end + + ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)] + ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] + @. ∂ᶠu₃_err_∂ᶜρ = + dtγ * ( + ᶠp_grad_matrix ⋅ + DiagonalMatrixRow(ᶜkappa_m * (T_0 * cp_d - ᶜK - ᶜΦ)) + + DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / abs2(ᶠinterp(ᶜρ))) ⋅ + ᶠinterp_matrix() + ) + @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) + if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] + @. ∂ᶠu₃_err_∂ᶜρq_tot = + dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(( + ᶜkappa_m * ∂e_int_∂q_tot + + ᶜ∂kappa_m∂q_tot * ( + cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * ᶜspecific.q_tot + ) + )) + end + + ∂ᶠu₃_err_∂ᶜuₕ = matrix[@name(f.u₃), @name(c.uₕ)] + ∂ᶠu₃_err_∂ᶠu₃ = matrix[@name(f.u₃), @name(f.u₃)] + I_u₃ = DiagonalMatrixRow(one_C3xACT3) + @. ∂ᶠu₃_err_∂ᶜuₕ = + dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶜuₕ + if rs isa RayleighSponge + @. ∂ᶠu₃_err_∂ᶠu₃ = + dtγ * ( + ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ + ∂ᶜK_∂ᶠu₃ + + DiagonalMatrixRow(-β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,)) + ) - (I_u₃,) + else + @. ∂ᶠu₃_err_∂ᶠu₃ = + dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ + ∂ᶜK_∂ᶠu₃ - (I_u₃,) + end + + tracer_info = ( + (@name(c.ρq_liq), @name(q_liq), @name(ᶜwₗ)), + (@name(c.ρq_ice), @name(q_ice), @name(ᶜwᵢ)), + (@name(c.ρq_rai), @name(q_rai), @name(ᶜwᵣ)), + (@name(c.ρq_sno), @name(q_sno), @name(ᶜwₛ)), + ) + if !(p.atmos.moisture_model isa DryModel) || use_derivative(diffusion_flag) + ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)] + @. ∂ᶜρe_tot_err_∂ᶜρe_tot = zero(typeof(∂ᶜρe_tot_err_∂ᶜρe_tot)) - (I,) + end + + if !(p.atmos.moisture_model isa DryModel) + #TODO: tetsing explicit vs implicit + #@. ∂ᶜρe_tot_err_∂ᶜρe_tot += + # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ + # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ + # DiagonalMatrixRow( + # -(1 + ᶜkappa_m) / ᶜρ * ifelse( + # ᶜh_tot == 0, + # (Geometry.WVector(FT(0)),), + # p.precomputed.ᶜwₕhₜ / ᶜh_tot, + # ), + # ) + + ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] + @. ∂ᶜρe_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρe_tot_err_∂ᶜρq_tot)) + #TODO: tetsing explicit vs implicit + #@. ∂ᶜρe_tot_err_∂ᶜρq_tot = + # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ + # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ + # DiagonalMatrixRow( + # -(ᶜkappa_m) * ∂e_int_∂q_tot / ᶜρ * ifelse( + # ᶜh_tot == 0, + # (Geometry.WVector(FT(0)),), + # p.precomputed.ᶜwₕhₜ / ᶜh_tot, + # ), + # ) + + ∂ᶜρq_tot_err_∂ᶜρq_tot = matrix[@name(c.ρq_tot), @name(c.ρq_tot)] + @. ∂ᶜρq_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρq_tot_err_∂ᶜρq_tot)) - (I,) + #TODO: tetsing explicit vs implicit + #@. ∂ᶜρq_tot_err_∂ᶜρq_tot = + # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ + # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ + # DiagonalMatrixRow( + # -1 / ᶜρ * ifelse( + # ᶜspecific.q_tot == 0, + # (Geometry.WVector(FT(0)),), + # p.precomputed.ᶜwₜqₜ / ᶜspecific.q_tot, + # ), + # ) - (I,) + + MatrixFields.unrolled_foreach(tracer_info) do (ρqₚ_name, _, wₚ_name) + MatrixFields.has_field(Y, ρqₚ_name) || return + ∂ᶜρqₚ_err_∂ᶜρqₚ = matrix[ρqₚ_name, ρqₚ_name] + ᶜwₚ = MatrixFields.get_field(p.precomputed, wₚ_name) + @. ∂ᶜρqₚ_err_∂ᶜρqₚ = + dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ + ᶠright_bias_matrix() ⋅ + DiagonalMatrixRow(-Geometry.WVector(ᶜwₚ) / ᶜρ) - (I,) + end + + end + + if p.atmos.moisture_model isa NonEquilMoistModel && + use_derivative(noneq_cloud_formation_flag) + p_vapₛₗ(tps, ts) = TD.saturation_vapor_pressure(tps, ts, TD.Liquid()) + p_vapₛᵢ(tps, ts) = TD.saturation_vapor_pressure(tps, ts, TD.Ice()) + + function ∂p_vapₛₗ_∂T(tps, ts) + T = TD.air_temperature(tps, ts) + Rᵥ = TD.Parameters.R_v(tps) + Lᵥ = TD.latent_heat_vapor(tps, ts) + return p_vapₛₗ(tps, ts) * Lᵥ / (Rᵥ * T^2) + end + function ∂p_vapₛᵢ_∂T(tps, ts) + T = TD.air_temperature(tps, ts) + Rᵥ = TD.Parameters.R_v(tps) + Lₛ = TD.latent_heat_sublim(tps, ts) + return p_vapₛᵢ(tps, ts) * Lₛ / (Rᵥ * T^2) + end + + function ∂qₛₗ_∂T(tps, ts) + T = TD.air_temperature(tps, ts) + Rᵥ = TD.Parameters.R_v(tps) + Lᵥ = TD.latent_heat_vapor(tps, ts) + qᵥ_sat_liq = TD.q_vap_saturation_liquid(tps, ts) + return qᵥ_sat_liq * (Lᵥ / (Rᵥ * T^2) - 1 / T) + end + function ∂qₛᵢ_∂T(tps, ts) + T = TD.air_temperature(tps, ts) + Rᵥ = TD.Parameters.R_v(tps) + Lₛ = TD.latent_heat_sublim(tps, ts) + qᵥ_sat_ice = TD.q_vap_saturation_ice(tps, ts) + return qᵥ_sat_ice * (Lₛ / (Rᵥ * T^2) - 1 / T) + end + + function Γₗ(tps, ts) + cₚ_air = TD.cp_m(tps, ts) + Lᵥ = TD.latent_heat_vapor(tps, ts) + return 1 + (Lᵥ / cₚ_air) * ∂qₛₗ_∂T(tps, ts) + end + function Γᵢ(tps, ts) + cₚ_air = TD.cp_m(tps, ts) + Lₛ = TD.latent_heat_sublim(tps, ts) + return 1 + (Lₛ / cₚ_air) * ∂qₛᵢ_∂T(tps, ts) + end + + cmc = CAP.microphysics_cloud_params(params) + τₗ = cmc.liquid.τ_relax + τᵢ = cmc.ice.τ_relax + function limit(q, dt, n::Int) + return q / float(dt) / n + end + + function ∂ρqₗ_err_∂ρqᵪ(tps, ts, deriv) + FT_inner = eltype(tps) + q = TD.PhasePartition(tps, ts) + ρ = TD.air_density(tps, ts) + + if TD.vapor_specific_humidity(q) + TD.liquid_specific_humidity(q) > FT_inner(0) + return deriv + else + return FT_inner(0) + end + end + + function ∂ρqᵢ_err_∂ρqᵪ(tps, ts, deriv) + FT_inner = eltype(tps) + q = TD.PhasePartition(tps, ts) + ρ = TD.air_density(tps, ts) + + if TD.vapor_specific_humidity(q) + TD.ice_specific_humidity(q) > FT_inner(0) + return deriv + else + return FT_inner(0) + end + end + + ∂ᶜρqₗ_err_∂ᶜρqₗ = matrix[@name(c.ρq_liq), @name(c.ρq_liq)] + ∂ᶜρqᵢ_err_∂ᶜρqᵢ = matrix[@name(c.ρq_ice), @name(c.ρq_ice)] + + ∂ᶜρqₗ_err_∂ᶜρqₜ = matrix[@name(c.ρq_liq), @name(c.ρq_tot)] + ∂ᶜρqᵢ_err_∂ᶜρqₜ = matrix[@name(c.ρq_ice), @name(c.ρq_tot)] + + #@. ∂ᶜρqₗ_err_∂ᶜρqₗ -= + # DiagonalMatrixRow(1 / (τₗ * Γₗ(thermo_params, ᶜts))) + @. ∂ᶜρqₗ_err_∂ᶜρqₗ += + DiagonalMatrixRow( + ∂ρqₗ_err_∂ρqᵪ( + thermo_params, ᶜts, (-1 / (τₗ * Γₗ(thermo_params, ᶜts))), + ) + ) + + #@. ∂ᶜρqᵢ_err_∂ᶜρqᵢ -= + # DiagonalMatrixRow(1 / (τᵢ * Γᵢ(thermo_params, ᶜts))) + + @. ∂ᶜρqᵢ_err_∂ᶜρqᵢ += + DiagonalMatrixRow( + ∂ρqᵢ_err_∂ρqᵪ( + thermo_params, ᶜts, (-1 / (τᵢ * Γᵢ(thermo_params, ᶜts))), + ) + ) + + ᶜp = @. lazy(TD.air_pressure(thermo_params, ᶜts)) + ᶜ∂T_∂p = @. lazy(1 / (ᶜρ * TD.gas_constant_air(thermo_params, ᶜts))) + + # qₛₗ = p_vapₛₗ / p, qₛᵢ = p_vapₛᵢ / p + ᶜ∂qₛₗ_∂p = @. lazy( + -p_vapₛₗ(thermo_params, ᶜts) / ᶜp^2 + + ∂p_vapₛₗ_∂T(thermo_params, ᶜts) * ᶜ∂T_∂p / ᶜp, + ) + ᶜ∂qₛᵢ_∂p = @. lazy( + -p_vapₛᵢ(thermo_params, ᶜts) / ᶜp^2 + + ∂p_vapₛᵢ_∂T(thermo_params, ᶜts) * ᶜ∂T_∂p / ᶜp, + ) + + ᶜ∂p_∂ρqₜ = @. lazy( + ᶜkappa_m * ∂e_int_∂q_tot + + ᶜ∂kappa_m∂q_tot * ( + cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * ᶜspecific.q_tot + ), + ) + + #@. ∂ᶜρqₗ_err_∂ᶜρqₜ = DiagonalMatrixRow( + # (1 - ᶜρ * ᶜ∂qₛₗ_∂p * ᶜ∂p_∂ρqₜ) / (τₗ * Γₗ(thermo_params, ᶜts)), + #) + @. ∂ᶜρqₗ_err_∂ᶜρqₜ = DiagonalMatrixRow( + ∂ρqₗ_err_∂ρqᵪ( + thermo_params, ᶜts, ((1 - ᶜρ * ᶜ∂qₛₗ_∂p * ᶜ∂p_∂ρqₜ) / (τₗ * Γₗ(thermo_params, ᶜts))), + ) + ) + + #@. ∂ᶜρqᵢ_err_∂ᶜρqₜ = DiagonalMatrixRow( + # (1 - ᶜρ * ᶜ∂qₛᵢ_∂p * ᶜ∂p_∂ρqₜ) / (τᵢ * Γᵢ(thermo_params, ᶜts)), + #) + @. ∂ᶜρqᵢ_err_∂ᶜρqₜ = DiagonalMatrixRow( + ∂ρqᵢ_err_∂ρqᵪ( + thermo_params, ᶜts, ((1 - ᶜρ * ᶜ∂qₛᵢ_∂p * ᶜ∂p_∂ρqₜ) / (τᵢ * Γᵢ(thermo_params, ᶜts))), + ) + ) + end + + if use_derivative(diffusion_flag) + α_vert_diff_tracer = CAP.α_vert_diff_tracer(params) + (; ᶜK_h, ᶜK_u) = p.precomputed + @. ᶜdiffusion_h_matrix = + ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h)) ⋅ + ᶠgradᵥ_matrix() + @. ᶜdiffusion_h_matrix_scaled = + ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( + ᶠinterp(ᶜρ) * ᶠinterp(α_vert_diff_tracer * ᶜK_h), + ) ⋅ ᶠgradᵥ_matrix() + if ( + MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) || + !isnothing(p.atmos.turbconv_model) || + !disable_momentum_vertical_diffusion(p.atmos.vert_diff) + ) + @. ᶜdiffusion_u_matrix = + ᶜadvdivᵥ_matrix() ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u)) ⋅ ᶠgradᵥ_matrix() + end + + ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] + @. ∂ᶜρe_tot_err_∂ᶜρ = + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( + ( + -(1 + ᶜkappa_m) * ᶜspecific.e_tot - + ᶜkappa_m * ∂e_int_∂q_tot * ᶜspecific.q_tot + ) / ᶜρ, + ) + @. ∂ᶜρe_tot_err_∂ᶜρe_tot += + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) + + if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] + ∂ᶜρq_tot_err_∂ᶜρ = matrix[@name(c.ρq_tot), @name(c.ρ)] + @. ∂ᶜρe_tot_err_∂ᶜρq_tot += + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(( + ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + + ᶜ∂kappa_m∂q_tot * ( + cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * ᶜspecific.q_tot + ) + )) + @. ∂ᶜρq_tot_err_∂ᶜρ = + dtγ * ᶜdiffusion_h_matrix ⋅ + DiagonalMatrixRow(-(ᶜspecific.q_tot) / ᶜρ) + @. ∂ᶜρq_tot_err_∂ᶜρq_tot += + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) + end + + MatrixFields.unrolled_foreach(tracer_info) do (ρq_name, q_name, _) + MatrixFields.has_field(Y, ρq_name) || return + ᶜq = MatrixFields.get_field(ᶜspecific, q_name) + ∂ᶜρq_err_∂ᶜρ = matrix[ρq_name, @name(c.ρ)] + ∂ᶜρq_err_∂ᶜρq = matrix[ρq_name, ρq_name] + ᶜtridiagonal_matrix_scalar = ifelse( + q_name in (@name(q_rai), @name(q_sno)), + ᶜdiffusion_h_matrix_scaled, + ᶜdiffusion_h_matrix, + ) + @. ∂ᶜρq_err_∂ᶜρ = + dtγ * ᶜtridiagonal_matrix_scalar ⋅ DiagonalMatrixRow(-(ᶜq) / ᶜρ) + @. ∂ᶜρq_err_∂ᶜρq += + dtγ * ᶜtridiagonal_matrix_scalar ⋅ DiagonalMatrixRow(1 / ᶜρ) + end + + if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) + turbconv_params = CAP.turbconv_params(params) + c_d = CAP.tke_diss_coeff(turbconv_params) + #(; dt) = p + (; ᶜtke⁰, ᶜmixing_length) = p.precomputed + ᶜρa⁰ = + p.atmos.turbconv_model isa PrognosticEDMFX ? + p.precomputed.ᶜρa⁰ : ᶜρ + ᶜρatke⁰ = Y.c.sgs⁰.ρatke + + @inline dissipation_rate(tke⁰, mixing_length) = + tke⁰ >= 0 ? c_d * sqrt(tke⁰) / max(mixing_length, 1) : + 1 / float(dt) + @inline ∂dissipation_rate_∂tke⁰(tke⁰, mixing_length) = + tke⁰ > 0 ? c_d / (2 * max(mixing_length, 1) * sqrt(tke⁰)) : + typeof(tke⁰)(0) + + ᶜdissipation_matrix_diagonal = p.scratch.ᶜtemp_scalar + @. ᶜdissipation_matrix_diagonal = + ᶜρatke⁰ * ∂dissipation_rate_∂tke⁰(ᶜtke⁰, ᶜmixing_length) + + ∂ᶜρatke⁰_err_∂ᶜρ = matrix[@name(c.sgs⁰.ρatke), @name(c.ρ)] + ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = + matrix[@name(c.sgs⁰.ρatke), @name(c.sgs⁰.ρatke)] + @. ∂ᶜρatke⁰_err_∂ᶜρ = + dtγ * ( + ᶜdiffusion_u_matrix - + DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) + ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰) / ᶜρa⁰) + @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = + dtγ * ( + ( + ᶜdiffusion_u_matrix - + DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) + ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰) - + DiagonalMatrixRow(dissipation_rate(ᶜtke⁰, ᶜmixing_length)) + ) - (I,) + end + + if ( + !isnothing(p.atmos.turbconv_model) || + !disable_momentum_vertical_diffusion(p.atmos.vert_diff) + ) + ∂ᶜuₕ_err_∂ᶜuₕ = matrix[@name(c.uₕ), @name(c.uₕ)] + @. ∂ᶜuₕ_err_∂ᶜuₕ = + dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_u_matrix - (I,) + end + + end + + if p.atmos.turbconv_model isa PrognosticEDMFX + if use_derivative(sgs_advection_flag) + (; ᶜgradᵥ_ᶠΦ) = p.core + (; ᶜρʲs, ᶠu³ʲs, ᶜtsʲs, ᶜKʲs, bdmr_l, bdmr_r, bdmr) = p.precomputed + is_third_order = + p.atmos.numerics.edmfx_upwinding == Val(:third_order) + ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1 + ᶠset_upwind_bcs = Operators.SetBoundaryOperator(; + top = Operators.SetValue(zero(CT3{FT})), + bottom = Operators.SetValue(zero(CT3{FT})), + ) # Need to wrap ᶠupwind in this for well-defined boundaries. + UpwindMatrixRowType = + is_third_order ? QuaddiagonalMatrixRow : BidiagonalMatrixRow + ᶠupwind_matrix = is_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix + ᶠset_upwind_matrix_bcs = Operators.SetBoundaryOperator(; + top = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), + bottom = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), + ) # Need to wrap ᶠupwind_matrix in this for well-defined boundaries. + + ᶠu³ʲ_data = ᶠu³ʲs.:(1).components.data.:1 + + ᶜkappa_mʲ = p.scratch.ᶜtemp_scalar + @. ᶜkappa_mʲ = + TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) / + TD.cv_m(thermo_params, ᶜtsʲs.:(1)) + + # Note this is the derivative of R_m / cp_m with respect to q_tot + # but we call it ∂kappa_m∂q_totʲ + ᶜ∂kappa_m∂q_totʲ = p.scratch.ᶜtemp_scalar_2 + @. ᶜ∂kappa_m∂q_totʲ = + ( + ΔR_v * TD.cp_m(thermo_params, ᶜtsʲs.:(1)) - + Δcp_v * TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) + ) / abs2(TD.cp_m(thermo_params, ᶜtsʲs.:(1))) + + ∂ᶜq_totʲ_err_∂ᶜq_totʲ = + matrix[@name(c.sgsʲs.:(1).q_tot), @name(c.sgsʲs.:(1).q_tot)] + @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ = + dtγ * ( + DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - + ᶜadvdivᵥ_matrix() ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) + ) - (I,) + ∂ᶜq_totʲ_err_∂ᶠu₃ʲ = + matrix[@name(c.sgsʲs.:(1).q_tot), @name(f.sgsʲs.:(1).u₃)] + @. ∂ᶜq_totʲ_err_∂ᶠu₃ʲ = + dtγ * ( + -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠset_upwind_bcs( + ᶠupwind(CT3(sign(ᶠu³ʲ_data)), Y.c.sgsʲs.:(1).q_tot), + ) * adjoint(C3(sign(ᶠu³ʲ_data))), + ) + + DiagonalMatrixRow(Y.c.sgsʲs.:(1).q_tot) ⋅ ᶜadvdivᵥ_matrix() + ) ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) + + ∂ᶜmseʲ_err_∂ᶜq_totʲ = + matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)] + @. ∂ᶜmseʲ_err_∂ᶜq_totʲ = + dtγ * ( + -DiagonalMatrixRow( + adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * ᶜgradᵥ_ᶠΦ * Y.c.ρ / ᶜp * + ( + (ᶜkappa_mʲ / (ᶜkappa_mʲ + 1) * ∂e_int_∂q_tot) + + ᶜ∂kappa_m∂q_totʲ * ( + Y.c.sgsʲs.:(1).mse - ᶜΦ + + cp_d * T_0 + + ∂e_int_∂q_tot * Y.c.sgsʲs.:(1).q_tot + ) + ), + ) + ) + ∂ᶜmseʲ_err_∂ᶜmseʲ = + matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).mse)] + @. ∂ᶜmseʲ_err_∂ᶜmseʲ = + dtγ * ( + DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - + ᶜadvdivᵥ_matrix() ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) - + DiagonalMatrixRow( + adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * + ᶜgradᵥ_ᶠΦ * + Y.c.ρ * + ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + ) - (I,) + ∂ᶜmseʲ_err_∂ᶠu₃ʲ = + matrix[@name(c.sgsʲs.:(1).mse), @name(f.sgsʲs.:(1).u₃)] + @. ∂ᶜmseʲ_err_∂ᶠu₃ʲ = + dtγ * ( + -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠset_upwind_bcs( + ᶠupwind(CT3(sign(ᶠu³ʲ_data)), Y.c.sgsʲs.:(1).mse), + ) * adjoint(C3(sign(ᶠu³ʲ_data))), + ) + + DiagonalMatrixRow(Y.c.sgsʲs.:(1).mse) ⋅ ᶜadvdivᵥ_matrix() + ) ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) + + ∂ᶜρaʲ_err_∂ᶜq_totʲ = + matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)] + @. ᶠbidiagonal_matrix_ct3 = + DiagonalMatrixRow( + ᶠset_upwind_bcs( + ᶠupwind( + ᶠu³ʲs.:(1), + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ), + ) / ᶠJ, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( + ᶜJ * (ᶜρʲs.:(1))^2 / ᶜp * ( + ᶜkappa_mʲ / (ᶜkappa_mʲ + 1) * ∂e_int_∂q_tot + + ᶜ∂kappa_m∂q_totʲ * ( + Y.c.sgsʲs.:(1).mse - ᶜΦ + + cp_d * T_0 + + ∂e_int_∂q_tot * Y.c.sgsʲs.:(1).q_tot + ) + ), + ) + @. ᶠbidiagonal_matrix_ct3_2 = + DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ + DiagonalMatrixRow( + Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp) * + ∂e_int_∂q_tot, + ) + @. ∂ᶜρaʲ_err_∂ᶜq_totʲ = + dtγ * ᶜadvdivᵥ_matrix() ⋅ + (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) + + ∂ᶜρaʲ_err_∂ᶜmseʲ = + matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)] + @. ᶠbidiagonal_matrix_ct3 = + DiagonalMatrixRow( + ᶠset_upwind_bcs( + ᶠupwind( + ᶠu³ʲs.:(1), + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ), + ) / ᶠJ, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( + ᶜJ * ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + @. ᶠbidiagonal_matrix_ct3_2 = + DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ + DiagonalMatrixRow( + Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + @. ∂ᶜρaʲ_err_∂ᶜmseʲ = + dtγ * ᶜadvdivᵥ_matrix() ⋅ + (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) + + ∂ᶜρaʲ_err_∂ᶜρaʲ = + matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).ρa)] + @. ᶜadvection_matrix = + -(ᶜadvdivᵥ_matrix()) ⋅ + DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) + @. ∂ᶜρaʲ_err_∂ᶜρaʲ = + dtγ * ᶜadvection_matrix ⋅ + ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ + DiagonalMatrixRow(1 / ᶜρʲs.:(1)) - (I,) + + ∂ᶜρaʲ_err_∂ᶠu₃ʲ = + matrix[@name(c.sgsʲs.:(1).ρa), @name(f.sgsʲs.:(1).u₃)] + @. ∂ᶜρaʲ_err_∂ᶠu₃ʲ = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ * + ᶠset_upwind_bcs( + ᶠupwind( + CT3(sign(ᶠu³ʲ_data)), + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ), + ) * + adjoint(C3(sign(ᶠu³ʲ_data))) * + g³³(ᶠgⁱʲ), + ) + + turbconv_params = CAP.turbconv_params(params) + α_b = CAP.pressure_normalmode_buoy_coeff1(turbconv_params) + ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)] + @. ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = + dtγ * DiagonalMatrixRow( + (1 - α_b) * ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / + (ᶠinterp(ᶜρʲs.:(1)))^2, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( + (ᶜρʲs.:(1))^2 / ᶜp * ( + ᶜkappa_mʲ / (ᶜkappa_mʲ + 1) * ∂e_int_∂q_tot + + ᶜ∂kappa_m∂q_totʲ * ( + Y.c.sgsʲs.:(1).mse - ᶜΦ + + cp_d * T_0 + + ∂e_int_∂q_tot * Y.c.sgsʲs.:(1).q_tot + ) + ), + ) + ∂ᶠu₃ʲ_err_∂ᶜmseʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)] + @. ∂ᶠu₃ʲ_err_∂ᶜmseʲ = + dtγ * DiagonalMatrixRow( + (1 - α_b) * ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / + (ᶠinterp(ᶜρʲs.:(1)))^2, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( + ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), + ) + + ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] + ᶜu₃ʲ = p.scratch.ᶜtemp_C3 + @. ᶜu₃ʲ = ᶜinterp(Y.f.sgsʲs.:(1).u₃) + @. bdmr_l = convert(BidiagonalMatrixRow{FT}, ᶜleft_bias_matrix()) + @. bdmr_r = convert(BidiagonalMatrixRow{FT}, ᶜright_bias_matrix()) + @. bdmr = ifelse(ᶜu₃ʲ.components.data.:1 > 0, bdmr_l, bdmr_r) + @. ᶠtridiagonal_matrix_c3 = -(ᶠgradᵥ_matrix()) ⋅ bdmr + if rs isa RayleighSponge + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + dtγ * ( + ᶠtridiagonal_matrix_c3 ⋅ + DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - + DiagonalMatrixRow( + β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,), + ) + ) - (I_u₃,) + else + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + dtγ * ᶠtridiagonal_matrix_c3 ⋅ + DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - (I_u₃,) + end + + # entrainment and detrainment (rates are treated explicitly) + if use_derivative(sgs_entr_detr_flag) + (; ᶜentrʲs, ᶜdetrʲs, ᶜturb_entrʲs) = p.precomputed + @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ -= + dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1)) + @. ∂ᶜmseʲ_err_∂ᶜmseʲ -= + dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1)) + @. ∂ᶜρaʲ_err_∂ᶜρaʲ += + dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) - ᶜdetrʲs.:(1)) + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ -= + dtγ * (DiagonalMatrixRow( + (ᶠinterp(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1))) * + (one_C3xACT3,), + )) + end + + # non-hydrostatic pressure drag + # (quadratic drag term treated implicitly, buoyancy term explicitly) + if use_derivative(sgs_nh_pressure_flag) + (; ᶠu₃⁰) = p.precomputed + α_d = CAP.pressure_normalmode_drag_coeff(turbconv_params) + scale_height = + CAP.R_d(params) * CAP.T_surf_ref(params) / CAP.grav(params) + H_up_min = CAP.min_updraft_top(turbconv_params) + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ -= + dtγ * (DiagonalMatrixRow( + 2 * α_d * norm(Y.f.sgsʲs.:(1).u₃ - ᶠu₃⁰) / + max(scale_height, H_up_min) * (one_C3xACT3,), + )) + end + + # add updraft mass flux contributions to grid-mean + if use_derivative(sgs_mass_flux_flag) + # Jacobian contributions of updraft massflux to grid-mean + ∂ᶜupdraft_mass_flux_∂ᶜscalar = ᶠbidiagonal_matrix_ct3 + @. ∂ᶜupdraft_mass_flux_∂ᶜscalar = + DiagonalMatrixRow( + (ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) * (ᶠu³ʲs.:(1) - ᶠu³), + ) ⋅ ᶠinterp_matrix() ⋅ + DiagonalMatrixRow(Y.c.sgsʲs.:(1).ρa / ᶜρʲs.:(1)) + + # Derivative of total energy tendency with respect to updraft MSE + ## grid-mean ρe_tot + ᶜkappa_m = p.scratch.ᶜtemp_scalar + @. ᶜkappa_m = + TD.gas_constant_air(thermo_params, ᶜts) / + TD.cv_m(thermo_params, ᶜts) + + ᶜ∂kappa_m∂q_tot = p.scratch.ᶜtemp_scalar_2 + @. ᶜ∂kappa_m∂q_tot = + ( + ΔR_v * TD.cv_m(thermo_params, ᶜts) - + Δcv_v * TD.gas_constant_air(thermo_params, ᶜts) + ) / abs2(TD.cv_m(thermo_params, ᶜts)) + + @. ∂ᶜρe_tot_err_∂ᶜρ += + dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ + DiagonalMatrixRow( + ( + -(1 + ᶜkappa_m) * ᶜspecific.e_tot - + ᶜkappa_m * ∂e_int_∂q_tot * ᶜspecific.q_tot + ) / ᶜρ, + ) + + @. ∂ᶜρe_tot_err_∂ᶜρq_tot += + dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ + DiagonalMatrixRow(( + ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + + ᶜ∂kappa_m∂q_tot * ( + cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * ᶜspecific.q_tot + ) + )) + + @. ∂ᶜρe_tot_err_∂ᶜρe_tot += + dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ + DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) + + ∂ᶜρe_tot_err_∂ᶜmseʲ = + matrix[@name(c.ρe_tot), @name(c.sgsʲs.:(1).mse)] + @. ∂ᶜρe_tot_err_∂ᶜmseʲ = + -(dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar) + + ## grid-mean ρq_tot + @. ∂ᶜρq_tot_err_∂ᶜρ += + dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ + DiagonalMatrixRow(-(ᶜspecific.q_tot) / ᶜρ) + + @. ∂ᶜρq_tot_err_∂ᶜρq_tot += + dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ + DiagonalMatrixRow(1 / ᶜρ) + + ∂ᶜρq_tot_err_∂ᶜq_totʲ = + matrix[@name(c.ρq_tot), @name(c.sgsʲs.:(1).q_tot)] + @. ∂ᶜρq_tot_err_∂ᶜq_totʲ = + -(dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar) + + # grid-mean ∂/∂(u₃ʲ) + ∂ᶜρe_tot_err_∂ᶠu₃ = matrix[@name(c.ρe_tot), @name(f.u₃)] + @. ∂ᶜρe_tot_err_∂ᶠu₃ += + dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( + ᶠinterp( + (Y.c.sgsʲs.:(1).mse + ᶜKʲs.:(1) - ᶜh_tot) * + ᶜρʲs.:(1) * + ᶜJ * + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ) / ᶠJ * (g³³(ᶠgⁱʲ)), + ) + + ∂ᶜρe_tot_err_∂ᶠu₃ʲ = + matrix[@name(c.ρe_tot), @name(f.sgsʲs.:(1).u₃)] + @. ∂ᶜρe_tot_err_∂ᶠu₃ʲ = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠinterp( + (Y.c.sgsʲs.:(1).mse + ᶜKʲs.:(1) - ᶜh_tot) * + ᶜρʲs.:(1) * + ᶜJ * + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ) / ᶠJ * (g³³(ᶠgⁱʲ)), + ) + + ∂ᶜρq_tot_err_∂ᶠu₃ = matrix[@name(c.ρq_tot), @name(f.u₃)] + @. ∂ᶜρq_tot_err_∂ᶠu₃ += + dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( + ᶠinterp( + (Y.c.sgsʲs.:(1).q_tot - ᶜspecific.q_tot) * + ᶜρʲs.:(1) * + ᶜJ * + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ) / ᶠJ * (g³³(ᶠgⁱʲ)), + ) + + ∂ᶜρq_tot_err_∂ᶠu₃ʲ = + matrix[@name(c.ρq_tot), @name(f.sgsʲs.:(1).u₃)] + @. ∂ᶜρq_tot_err_∂ᶠu₃ʲ = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + ᶠinterp( + (Y.c.sgsʲs.:(1).q_tot - ᶜspecific.q_tot) * + ᶜρʲs.:(1) * + ᶜJ * + draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), + ) / ᶠJ * (g³³(ᶠgⁱʲ)), + ) + + # grid-mean ∂/∂(rho*a) + ∂ᶜρe_tot_err_∂ᶜρa = + matrix[@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρa)] + @. ∂ᶜρe_tot_err_∂ᶜρa = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + (ᶠu³ʲs.:(1) - ᶠu³) * + ᶠinterp((Y.c.sgsʲs.:(1).mse + ᶜKʲs.:(1) - ᶜh_tot)) / ᶠJ, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ) + + ∂ᶜρq_tot_err_∂ᶜρa = + matrix[@name(c.ρq_tot), @name(c.sgsʲs.:(1).ρa)] + @. ∂ᶜρq_tot_err_∂ᶜρa = + dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( + (ᶠu³ʲs.:(1) - ᶠu³) * + ᶠinterp((Y.c.sgsʲs.:(1).q_tot - ᶜspecific.q_tot)) / ᶠJ, + ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ) + end + elseif rs isa RayleighSponge + ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] + @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = + dtγ * + -DiagonalMatrixRow( + β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,), + ) - (I_u₃,) + end + end + + # NOTE: All velocity tendency derivatives should be set BEFORE this call. + zero_velocity_jacobian!(matrix, Y, p, t) +end + +invert_jacobian!(::ApproxJacobian, cache, ΔY, R) = + LinearAlgebra.ldiv!(ΔY, cache.matrix, R) + +# TODO: Rewrite the plotting infrastructure to handle `FieldMatrix`, so that we +# can avoid inefficiently converting the approximate Jacobian to a dense matrix. +function save_jacobian!(alg::ApproxJacobian, cache, Y, dtγ, t) + (; matrix, temp_matrix, temp_matrix_column, column_matrix) = cache + n_columns = Fields.ncolumns(Y.c) + + # TODO: Fix bug in ClimaCore's column function, so that we can use + # @. lazy((matrix + I_matrix) / dtγ) instead of caching this FieldMatrix. + temp_matrix .= (matrix .+ identity_matrix(matrix, Y)) ./ dtγ + + field_matrix_to_dense_matrix!(column_matrix, first_column(temp_matrix), Y) + file_name = "approx_jacobian_first" + description = + "Approx ∂Yₜ/∂Y" * (n_columns == 1 ? "" : " at $(first_column_str(Y))") + save_cached_column_matrix_and_vector!(cache, file_name, description, t) + + if n_columns > 1 + level_mapreduce_matrix!(abs, max, temp_matrix_column, temp_matrix) + field_matrix_to_dense_matrix!(column_matrix, temp_matrix_column, Y) + file_name = "approx_jacobian_max" + description = "Approx ∂Yₜ/∂Y, max over all columns" + save_cached_column_matrix_and_vector!(cache, file_name, description, t) + + level_mapreduce_matrix!(abs, +, temp_matrix_column, temp_matrix) + field_matrix_to_dense_matrix!(column_matrix, temp_matrix_column, Y) + column_matrix ./= n_columns + file_name = "approx_jacobian_avg" + description = "Approx ∂Yₜ/∂Y, avg over all columns" + save_cached_column_matrix_and_vector!(cache, file_name, description, t) + end +end + +# TODO: Remove all of the following code after extending ClimaCore.MatrixFields. + +function level_mapreduce_field!(f::F, op::O, field_column, field) where {F, O} + (Nv, Nf) = size(parent(field_column)) + parent_dimensions = length(size(parent(field))) + @assert parent_dimensions in (4, 5) + reshaped_size = parent_dimensions == 4 ? (Nv, 1, Nf, 1) : (Nv, 1, 1, Nf, 1) + reshaped_column_parent = reshape(parent(field_column), reshaped_size...) + if op == + + sum!(f, reshaped_column_parent, parent(field)) + elseif op == max + maximum!(f, reshaped_column_parent, parent(field)) + elseif op == min + minimum!(f, reshaped_column_parent, parent(field)) + else + error("level_mapreduce_field! has not been defined for op = $op") + end +end +level_mapreduce_matrix!(f::F, op::O, matrix_column, matrix) where {F, O} = + foreach(keys(matrix)) do key + matrix[key] isa Fields.Field || return + level_mapreduce_field!(f, op, matrix_column[key], matrix[key]) + end + +tensor_axes_tuple(::Type{T}) where {T} = + T <: Geometry.AxisTensor ? + map(axis -> typeof(axis).parameters[1], axes(T)) : () + +primitive_value_at_index(value, (row_axes, col_axes)) = + if isprimitivetype(typeof(value)) # same as a LinearAlgebra.UniformScaling + row_axes == col_axes ? value : zero(value) + elseif value isa Geometry.AxisVector + @assert isprimitivetype(eltype(value)) + @assert length(row_axes) == 1 && length(col_axes) == 0 + value_axes = tensor_axes_tuple(typeof(value)) + row_axis_index = findfirst(==(row_axes[1]), value_axes[1]) + isnothing(row_axis_index) ? zero(eltype(value)) : value[row_axis_index] + elseif value isa Geometry.AxisTensor + @assert isprimitivetype(eltype(value)) + @assert length(row_axes) == 1 && length(col_axes) == 1 + value_axes = tensor_axes_tuple(typeof(value)) + row_axis_index = findfirst(==(row_axes[1]), value_axes[1]) + col_axis_index = findfirst(==(col_axes[1]), value_axes[2]) + isnothing(row_axis_index) || isnothing(col_axis_index) ? + zero(eltype(value)) : value[row_axis_index, col_axis_index] + elseif value isa LinearAlgebra.Adjoint + primitive_value_at_index(parent(value), (col_axes, row_axes)) + else + sub_names = fieldnames(typeof(value)) + sub_values = + MatrixFields.unrolled_map(Base.Fix1(getfield, value), sub_names) + nonempty_sub_values = + MatrixFields.unrolled_filter(x -> sizeof(x) > 0, sub_values) + @assert length(nonempty_sub_values) == 1 + primitive_value_at_index(nonempty_sub_values[1], (row_axes, col_axes)) + end + +@static if hasfield(Method, :recursion_relation) + for method in methods(primitive_value_at_index) + method.recursion_relation = Returns(true) + end +end + +function field_matrix_to_dense_matrix!(out, matrix, Y) + device = ClimaComms.device(Y.c) # ClimaComms.device(Y) + field_names = scalar_field_names(Y) + index_ranges = scalar_field_index_ranges(Y) + out .= 0 + + for ((block_row, block_col), matrix_block) in matrix + is_child_name_of_row = Base.Fix2(MatrixFields.is_child_name, block_row) + is_child_name_of_col = Base.Fix2(MatrixFields.is_child_name, block_col) + subblock_row_indices = findall(is_child_name_of_row, field_names) + subblock_col_indices = findall(is_child_name_of_col, field_names) + block_row_field = MatrixFields.get_field(Y, block_row) + block_col_field = MatrixFields.get_field(Y, block_col) + + for (sub_row, subblock_row_index) in enumerate(subblock_row_indices) + for (sub_col, subblock_col_index) in enumerate(subblock_col_indices) + row_index_range = index_ranges[subblock_row_index] + col_index_range = index_ranges[subblock_col_index] + out_subblock = view(out, row_index_range, col_index_range) + + if matrix_block isa LinearAlgebra.UniformScaling + view(out_subblock, LinearAlgebra.diagind(out_subblock)) .= + sub_row == sub_col ? matrix_block.λ : + zero(matrix_block.λ) + else + subblock_row_axes = map( + Base.Fix2(getindex, sub_row), + tensor_axes_tuple(eltype(block_row_field)), + ) + subblock_col_axes = map( + Base.Fix2(getindex, sub_col), + tensor_axes_tuple(eltype(block_col_field)), + ) + @assert length(subblock_row_axes) in (0, 1) + @assert length(subblock_col_axes) in (0, 1) + get_value_in_subblock = Base.Fix2( + primitive_value_at_index, + (subblock_row_axes, subblock_col_axes), + ) + # TODO: Get rid of this allocation. + subblock_values = map.(get_value_in_subblock, matrix_block) + + banded_matrix_transpose = + MatrixFields.column_field2array_view(subblock_values)' + band_indices = + (-banded_matrix_transpose.u):(banded_matrix_transpose.l) + for band_index in band_indices + out_subblock_indices = diagind(out_subblock, band_index) + transposed_band = MatrixFields.band(-band_index) + band_wrapper = + view(banded_matrix_transpose, transposed_band) + view(out_subblock, out_subblock_indices) .= + dataview(band_wrapper) + end # TODO: Fuse this loop. + end + end + end + end +end + +# TODO: Import BandedMatrices.dataview. +function dataview(V) + A = parent(parent(V)) + b = MatrixFields.band(V) + m, n = size(A) + return view(A.data, A.u - b + 1, (max(b, 0) + 1):min(n, m + b)) +end diff --git a/src/prognostic_equations/implicit/debug_jacobian.jl b/src/prognostic_equations/implicit/debug_jacobian.jl new file mode 100644 index 0000000000..1b6f93c48e --- /dev/null +++ b/src/prognostic_equations/implicit/debug_jacobian.jl @@ -0,0 +1,73 @@ +""" + DebugJacobian( + approx_jacobian_algorithm, + use_exact_jacobian, + only_debug_first_column_jacobian, + ) + +A `JacobianAlgorithm` that simultaneously computes an `ExactJacobian` and an +`ApproxJacobian`, so that the quality of the approximation can be assessed. The +`use_exact_jacobian` flag controls whether the exact Jacobian is used in the +implicit solver instead of the approximation, and, when `use_exact_jacobian` is +`false`, the `only_debug_first_column_jacobian` flag controls whether the exact +Jacobian is only evaluated in the first column. +""" +struct DebugJacobian{A <: ApproxJacobian, EJ, FC} <: JacobianAlgorithm + approx_jacobian_algorithm::A +end +DebugJacobian( + approx_jacobian_algorithm, + use_exact_jacobian, + only_debug_first_column_jacobian, +) = DebugJacobian{ + typeof(approx_jacobian_algorithm), + use_exact_jacobian, + only_debug_first_column_jacobian, +}( + approx_jacobian_algorithm, +) + +use_exact_jacobian(::DebugJacobian{A, EJ, FC}) where {A, EJ, FC} = EJ + +only_debug_first_column_jacobian(::DebugJacobian{A, EJ, FC}) where {A, EJ, FC} = + FC + +required_columns_for_exact_jacobian(alg, x) = + use_exact_jacobian(alg) || !only_debug_first_column_jacobian(alg) ? x : + first_column(x) + +function jacobian_cache(alg::DebugJacobian, Y, atmos) + Y_for_exact = required_columns_for_exact_jacobian(alg, Y) + exact_cache = jacobian_cache(ExactJacobian(), Y_for_exact, atmos) + approx_cache = jacobian_cache(alg.approx_jacobian_algorithm, Y, atmos) + return (; exact_cache..., approx_cache...) +end + +update_jacobian!(alg::DebugJacobian, cache, Y, p, dtγ, t) = + use_exact_jacobian(alg) ? + update_jacobian!(ExactJacobian(), cache, Y, p, dtγ, t) : + update_jacobian!(alg.approx_jacobian_algorithm, cache, Y, p, dtγ, t) + +function update_and_check_jacobian!(alg::DebugJacobian, cache, Y, p, dtγ, t) + update_for_exact! = + use_exact_jacobian(alg) ? update_jacobian! : + update_jacobian_skip_factorizing! + Y_for_exact = required_columns_for_exact_jacobian(alg, Y) + p_for_exact = required_columns_for_exact_jacobian(alg, p) + update_for_exact!(ExactJacobian(), cache, Y_for_exact, p_for_exact, dtγ, t) + update_jacobian!(alg.approx_jacobian_algorithm, cache, Y, p, dtγ, t) + # TODO: Add a quantitative check of the Jacobian approximation. +end + +invert_jacobian!(alg::DebugJacobian, cache, ΔY, R) = + use_exact_jacobian(alg) ? invert_jacobian!(ExactJacobian(), cache, ΔY, R) : + invert_jacobian!(alg.approx_jacobian_algorithm, cache, ΔY, R) + +function save_jacobian!(alg::DebugJacobian, cache, Y, dtγ, t) + Y_for_exact = required_columns_for_exact_jacobian(alg, Y) + save_jacobian!(ExactJacobian(), cache, Y_for_exact, dtγ, t) + save_jacobian!(alg.approx_jacobian_algorithm, cache, Y, dtγ, t) + # TODO: Save the average/maximum difference between the approximate and + # exact Jacobians, instead of just computing the difference between their + # averages/maxima when plotting. +end diff --git a/src/prognostic_equations/implicit/dual_fixes.jl b/src/prognostic_equations/implicit/dual_fixes.jl new file mode 100644 index 0000000000..fc95139e6a --- /dev/null +++ b/src/prognostic_equations/implicit/dual_fixes.jl @@ -0,0 +1,165 @@ +# TODO: Move all of the following to ClimaCore. +# Method definitions that break precompilation have been commented out. + +# Missing methods: + +using ClimaCore: Adapt, Grids +Adapt.@adapt_structure Grids.ColumnGrid + +Spaces.ncolumns(::Spaces.FiniteDifferenceSpace) = 1 + +# ClimaComms.device(fv::Fields.FieldVector) = +# ClimaComms.device(first(Fields._values(fv))) +# ClimaComms.context(fv::Fields.FieldVector) = +# ClimaComms.context(first(Fields._values(fv))) + +# Bugfix for column views into Field broadcasts: + +# column_style(::Type{S}) where {DS, S <: Fields.FieldStyle{DS}} = +# Fields.FieldStyle{column_style(DS)} +# column_style( +# ::Type{S}, +# ) where {Nv, Ni, A, S <: DataLayouts.VIJFHStyle{Nv, Ni, A}} = +# DataLayouts.VFStyle{Nv, A} +# column_style(::Type{S}) where {Ni, A, S <: DataLayouts.IJFHStyle{Ni, A}} = +# DataLayouts.DataFStyle{A} +# Base.@propagate_inbounds function Fields.column( +# bc::Base.Broadcast.Broadcasted{Style}, +# i, +# j, +# h, +# ) where {Style <: Fields.AbstractFieldStyle} +# _args = Fields.column_args(bc.args, i, j, h) +# _axes = Fields.column(axes(bc), i, j, h) +# Base.Broadcast.Broadcasted{column_style(Style)}(bc.f, _args, _axes) +# end +# Base.@propagate_inbounds function Fields.column( +# bc::DataLayouts.NonExtrudedBroadcasted{Style}, +# i, +# j, +# h, +# ) where {Style <: Fields.AbstractFieldStyle} +# _args = Fields.column_args(bc.args, i, j, h) +# _axes = Fields.column(axes(bc), i, j, h) +# DataLayouts.NonExtrudedBroadcasted{column_style(Style)}(bc.f, _args, _axes) +# end + +# Wrapping arrays in FieldVectors: + +import ClimaCore.DataLayouts: parent_array_type, device_dispatch + +parent_array_type( + ::Type{<:Base.ReshapedArray{T, N, P, MI}}, +) where {T, N, P, MI} = parent_array_type(P) +parent_array_type( + ::Type{<:PermutedDimsArray{T, N, P, IP, A}}, +) where {T, N, P, IP, A} = parent_array_type(A) + +device_dispatch(x::PermutedDimsArray) = device_dispatch(parent(x)) + +# TODO: Reshape column_vectors from (Ni * Nj * Nh) × (Nv * Nf) to the transpose. +function column_vectors_to_field_vector(column_vectors, example_field_vector) + example_fields = values(Fields._values(example_field_vector)) + example_column_fields = unrolled_map(first_column, example_fields) + column_lengths = unrolled_map(length ∘ parent, example_column_fields) + column_range_starts = unrolled_cumsum((1, column_lengths[1:(end - 1)]...)) + column_range_ends = column_range_starts .+ column_lengths .- 1 + column_ranges = + unrolled_map(UnitRange, column_range_starts, column_range_ends) + new_fields = unrolled_map( + example_fields, + column_ranges, + ) do example_field, column_range + new_data_layout = column_vectors_to_data_layout( + view(column_vectors, :, column_range), + Fields.field_values(example_field), + ) + Fields.Field(new_data_layout, axes(example_field)) + end + return Fields.FieldVector{eltype(column_vectors)}(new_fields) +end + +import ClimaCore: DataLayouts +import ClimaCore.DataLayouts: replace_basetype, union_all, singleton +import ClimaCore.DataLayouts: type_params, farray_size, universal_size + +# This is no longer needed, but it would still be good to fix. +# function array2data(array::AbstractArray, data::DataLayouts.AbstractData) +# T = replace_basetype(eltype(parent(data)), eltype(array), eltype(data)) +# return union_all(singleton(data)){T, Base.tail(type_params(data))...}( +# reshape(array, farray_size(data)...), +# ) +# end + +function column_vectors_to_data_layout(array, data) + T = replace_basetype(eltype(parent(data)), eltype(array), eltype(data)) + return union_all(singleton(data)){T, Base.tail(type_params(data))...}( + reshaped_column_vectors(array, data), + ) +end + +reshaped_column_vectors(array, data::DataLayouts.VF) = + reshape(array, get_Nv(data), :) +reshaped_column_vectors(array, data::DataLayouts.IHF) = + reshape(array, get_Ni(data), get_Nh(data), :) +reshaped_column_vectors(array, data::DataLayouts.IFH) = + PermutedDimsArray(reshape(array, get_Ni(data), get_Nh(data), :), (1, 3, 2)) +reshaped_column_vectors(array, data::DataLayouts.IJHF) = + reshape(array, get_Ni(data), get_Nj(data), get_Nh(data), :) +reshaped_column_vectors(array, data::DataLayouts.IJFH) = PermutedDimsArray( + reshape(array, get_Ni(data), get_Nj(data), get_Nh(data), :), + (1, 2, 4, 3), +) +reshaped_column_vectors(array, data::DataLayouts.VIHF) = PermutedDimsArray( + reshape(array, get_Ni(data), get_Nh(data), get_Nv(data), :), + (3, 1, 2, 4), +) +reshaped_column_vectors(array, data::DataLayouts.VIFH) = PermutedDimsArray( + reshape(array, get_Ni(data), get_Nh(data), get_Nv(data), :), + (3, 1, 4, 2), +) +reshaped_column_vectors(array, data::DataLayouts.VIJHF) = PermutedDimsArray( + reshape(array, get_Ni(data), get_Nj(data), get_Nh(data), get_Nv(data), :), + (4, 1, 2, 3, 5), +) +reshaped_column_vectors(array, data::DataLayouts.VIJFH) = PermutedDimsArray( + reshape(array, get_Ni(data), get_Nj(data), get_Nh(data), get_Nv(data), :), + (4, 1, 2, 5, 3), +) +get_Ni(data) = universal_size(data)[1] +get_Nj(data) = universal_size(data)[2] +get_Nv(data) = universal_size(data)[4] +get_Nh(data) = universal_size(data)[5] + +# Iterating over levels of scalars: + +scalar_field_names(field_vector) = + MatrixFields.filtered_names(field_vector) do x + x isa Fields.Field && eltype(x) == eltype(field_vector) + end + +function scalar_level_iterator(field_vector) + Iterators.flatmap(scalar_field_names(field_vector)) do name + field = MatrixFields.get_field(field_vector, name) + if field isa Fields.SpectralElementField + (field,) + else + Iterators.map(1:Spaces.nlevels(axes(field))) do v + Fields.level(field, v - 1 + Operators.left_idx(axes(field))) + end + end + end +end + +function scalar_field_index_ranges(field_vector) + field_names = scalar_field_names(field_vector) + last_level_indices = accumulate(field_names; init = 0) do index, name + field = MatrixFields.get_field(field_vector, name) + n_levels = + field isa Fields.SpectralElementField ? 1 : + Spaces.nlevels(axes(field)) + index + n_levels + end + first_level_indices = (1, (last_level_indices[1:(end - 1)] .+ 1)...) + return map(UnitRange, first_level_indices, last_level_indices) +end diff --git a/src/prognostic_equations/implicit/exact_jacobian.jl b/src/prognostic_equations/implicit/exact_jacobian.jl new file mode 100644 index 0000000000..bad6754a28 --- /dev/null +++ b/src/prognostic_equations/implicit/exact_jacobian.jl @@ -0,0 +1,196 @@ +""" + ExactJacobian() + +A `JacobianAlgorithm` that computes the `ImplicitEquationJacobian` using +forward-mode automatic differentiation and inverts it using LU factorization. +""" +struct ExactJacobian <: JacobianAlgorithm end + +exact_jacobian_batch_size() = 32 # number of derivatives computed simultaneously + +replace_parent_type(x::Fields.FieldVector, ::Type{T}) where {T} = similar(x, T) +replace_parent_type(x::Fields.Field, ::Type{T}) where {T} = + Fields.Field(replace_parent_type(Fields.field_values(x), T), axes(x)) +replace_parent_type(x::DataLayouts.AbstractData, ::Type{T}) where {T} = + DataLayouts.replace_basetype(x, T) +replace_parent_type(x::Union{Tuple, NamedTuple}, ::Type{T}) where {T} = + unrolled_map(Base.Fix2(replace_parent_type, T), x) + +function jacobian_cache(::ExactJacobian, Y, atmos) + FT = eltype(Y) + DA = ClimaComms.array_type(Y) + + FT_dual = ForwardDiff.Dual{ExactJacobian, FT, exact_jacobian_batch_size()} + Y_dual = replace_parent_type(Y, FT_dual) + Yₜ_dual = similar(Y_dual) + precomputed_dual = + replace_parent_type(implicit_precomputed_quantities(Y, atmos), FT_dual) + scratch_dual = replace_parent_type(temporary_quantities(Y, atmos), FT_dual) + + n_columns = Fields.ncolumns(Y.c) + n_εs = length(first_column(Y)) + column_matrices = DA{FT}(undef, n_columns, n_εs, n_εs) + column_lu_factors = copy(column_matrices) + column_lu_solve_vectors = DA{FT}(undef, n_columns, n_εs) + lu_cache = DA{FT}(undef, n_columns) + + # LinearAlgebra.I does not support broadcasting, so we need a workaround. + I_column_matrix = DA{FT}(undef, n_εs, n_εs) + I_column_matrix .= 0 + I_column_matrix[diagind(I_column_matrix)] .= 1 + I_matrix = reshape(I_column_matrix, 1, n_εs, n_εs) + + return (; + Y_dual, + Yₜ_dual, + precomputed_dual, + scratch_dual, + column_matrices, + column_lu_factors, + column_lu_solve_vectors, + lu_cache, + I_matrix, + ) +end + +function update_jacobian_skip_factorizing!(::ExactJacobian, cache, Y, p, dtγ, t) + (; Y_dual, Yₜ_dual, precomputed_dual, scratch_dual, column_matrices) = cache + batch_size = exact_jacobian_batch_size() + + p_dual_args = ntuple(Val(fieldcount(typeof(p)))) do cache_field_index + cache_field_name = fieldname(typeof(p), cache_field_index) + if cache_field_name == :precomputed + (; p.precomputed..., precomputed_dual...) + elseif cache_field_name == :scratch + scratch_dual + else + getfield(p, cache_field_index) + end + end + p_dual = AtmosCache(p_dual_args...) + + Y_dual_scalar_levels = scalar_level_iterator(Y_dual) + level_batches = + Iterators.partition(enumerate(Y_dual_scalar_levels), batch_size) + + Y_dual .= Y + for level_batch in level_batches + for (partial_index, (_, Y_dual_scalar_level)) in enumerate(level_batch) + partials = ntuple(i -> i == partial_index ? 1 : 0, Val(batch_size)) + parent(Y_dual_scalar_level) .+= + ForwardDiff.Dual{ExactJacobian}(0, partials...) + end # Add a unique ε to Y_dual for each combination of scalar and level. + set_implicit_precomputed_quantities!(Y_dual, p_dual, t) # Compute ∂p/∂Y. + implicit_tendency!(Yₜ_dual, Y_dual, p_dual, t) # Compute ∂Yₜ/∂Y. + for (partial_index, (ε_index, _)) in enumerate(level_batch) + ε_column_vectors = view(column_matrices, :, :, ε_index) + column_vectors_to_field_vector(ε_column_vectors, Y) .= + getindex.(ForwardDiff.partials.(Yₜ_dual), partial_index) + end # Copy the new values of ∂Yₜ/∂Y into column_matrices. + Y_dual .= ForwardDiff.value.(Y_dual) # Drop this batch's εs from Y_dual. + end +end + +function update_jacobian!(::ExactJacobian, cache, Y, p, dtγ, t) + (; column_matrices, column_lu_factors, lu_cache, I_matrix) = cache + update_jacobian_skip_factorizing!(ExactJacobian(), cache, Y, p, dtγ, t) + column_lu_factors .= dtγ .* column_matrices .- I_matrix + parallel_lu_factorize!(column_lu_factors, lu_cache) +end + +function invert_jacobian!(::ExactJacobian, cache, ΔY, R) + (; column_lu_solve_vectors, column_lu_factors, lu_cache) = cache + column_vectors_to_field_vector(column_lu_solve_vectors, R) .= R + parallel_lu_solve!(column_lu_solve_vectors, column_lu_factors, lu_cache) + ΔY .= column_vectors_to_field_vector(column_lu_solve_vectors, ΔY) +end + +function save_jacobian!(::ExactJacobian, cache, Y, dtγ, t) + (; column_matrices, column_matrix) = cache + (n_columns, n_εs, _) = size(column_matrices) + + column_matrix .= view(column_matrices, 1, :, :) + file_name = "exact_jacobian_first" + title = "Exact ∂Yₜ/∂Y$(n_columns == 1 ? "" : " at $(first_column_str(Y))")" + save_cached_column_matrix_and_vector!(cache, file_name, title, t) + + if n_columns > 1 + maximum!(abs, reshape(column_matrix, 1, n_εs, n_εs), column_matrices) + file_name = "exact_jacobian_max" + title = "Exact ∂Yₜ/∂Y, max over all columns" + save_cached_column_matrix_and_vector!(cache, file_name, title, t) + + sum!(abs, reshape(column_matrix, 1, n_εs, n_εs), column_matrices) + column_matrix ./= n_columns + file_name = "exact_jacobian_avg" + title = "Exact ∂Yₜ/∂Y, avg over all columns" + save_cached_column_matrix_and_vector!(cache, file_name, title, t) + end +end + +# Set the derivative of `sqrt(x)` to `iszero(x) ? zero(x) : inv(2 * sqrt(x))` in +# order to properly handle derivatives of `x * sqrt(x)`. Without this change, +# the derivative of `x * sqrt(x)` is `NaN` when `x` is zero. This method +# specializes on the tag `ExactJacobian` because not specializing on any tag +# overwrites the generic method for `Dual` in `ForwardDiff` and breaks +# precompilation, while specializing on the default tag `Nothing` causes the +# type piracy Aqua test to fail. +@inline function Base.sqrt(d::ForwardDiff.Dual{ExactJacobian}) + tag = Val{ExactJacobian}() + x = ForwardDiff.value(d) + partials = ForwardDiff.partials(d) + val = sqrt(x) + deriv = iszero(x) ? zero(x) : inv(2 * val) + return ForwardDiff.dual_definition_retval(tag, val, deriv, partials) +end + +# TODO: Reshape column_matrices and turn this into a single kernel launch. +function parallel_lu_factorize!(column_matrices, temporary_vector) + @assert ndims(column_matrices) == 3 + n = size(column_matrices, 2) + @assert size(column_matrices, 3) == n + @inbounds for k in 1:n + all(!isnan, view(column_matrices, :, k, k)) || + error("LU error: NaN on diagonal") + all(!iszero, view(column_matrices, :, k, k)) || + error("LU error: 0 on diagonal") + temporary_vector .= inv.(view(column_matrices, :, k, k)) + for i in (k + 1):n + view(column_matrices, :, i, k) .*= temporary_vector + end + for j in (k + 1):n + for i in (k + 1):n + view(column_matrices, :, i, j) .-= + view(column_matrices, :, i, k) .* + view(column_matrices, :, k, j) + end + end + end +end + +# TODO: Reshape column_matrices and turn this into a single kernel launch. +function parallel_lu_solve!(column_vectors, column_matrices, temporary_vector) + @assert ndims(column_vectors) == 2 && ndims(column_matrices) == 3 + n = size(column_matrices, 2) + @assert size(column_vectors, 2) == n && size(column_matrices, 3) == n + @inbounds begin + for i in 2:n + temporary_vector .= zero(eltype(column_vectors)) + for j in 1:(i - 1) + temporary_vector .+= + view(column_matrices, :, i, j) .* view(column_vectors, :, j) + end + view(column_vectors, :, i) .-= temporary_vector + end + view(column_vectors, :, n) ./= view(column_matrices, :, n, n) + for i in (n - 1):-1:1 + temporary_vector .= zero(eltype(column_vectors)) + for j in (i + 1):n + temporary_vector .+= + view(column_matrices, :, i, j) .* view(column_vectors, :, j) + end + view(column_vectors, :, i) .-= temporary_vector + view(column_vectors, :, i) ./= view(column_matrices, :, i, i) + end + end +end diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index f5c2b61ce5..c56bacd498 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -1,28 +1,29 @@ -import LinearAlgebra: I, Adjoint, ldiv! -import ClimaCore.MatrixFields: @name -using ClimaCore.MatrixFields +import ForwardDiff +import LinearAlgebra: I, Adjoint, diagind +import ClimaCore.InputOutput: HDF5, HDF5Writer -abstract type DerivativeFlag end -struct UseDerivative <: DerivativeFlag end -struct IgnoreDerivative <: DerivativeFlag end -use_derivative(::UseDerivative) = true -use_derivative(::IgnoreDerivative) = false +using ClimaCore.MatrixFields +import ClimaCore.MatrixFields: @name """ - ImplicitEquationJacobian( - Y, atmos; - approximate_solve_iters, diffusion_flag, topography_flag, sgs_advection_flag, transform_flag - ) - -A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the -implicit step with the state ``Y``. + JacobianAlgorithm + +A description of how to compute the matrix ``∂R/∂Y``, where ``R(Y)`` denotes the +residual of an implicit step with the state ``Y``. Concrete implementations of +this abstract type should define 4 methods: + - `jacobian_cache(alg::JacobianAlgorithm, Y, atmos)` + - `update_jacobian!(alg::JacobianAlgorithm, cache, Y, p, dtγ, t)` + - `invert_jacobian!(alg::JacobianAlgorithm, cache, ΔY, R)` + - `save_jacobian!(alg::JacobianAlgorithm, cache, Y, p, dtγ, t)` +An additional method can also be defined to enable debugging of the Jacobian: + - `update_and_check_jacobian!(alg::JacobianAlgorithm, cache, Y, p, dtγ, t)` # Background When we use an implicit or split implicit-explicit (IMEX) timestepping scheme, -we end up with a nonlinear equation of the form ``E(Y) = 0``, where +we end up with a nonlinear equation of the form ``R(Y) = 0``, where ```math - E(Y) = Y_{imp}(Y) - Y = \\hat{Y} + Δt * T_{imp}(Y) - Y. + R(Y) = Y_{imp}(Y) - Y = \\hat{Y} + Δt * T_{imp}(Y) - Y. ``` In this expression, ``Y_{imp}(Y)`` denotes the state at some time ``t + Δt``. This can be expressed as the sum of ``\\hat{Y}``, the contribution from the @@ -38,1295 +39,182 @@ divided into several sub-steps or "stages", where the duration of stage ``i`` is ``Δt * γ_i`` for some constant ``γ_i`` between 0 and 1. In order to solve this equation using Newton's method, we must specify the -derivative ``∂E/∂Y``. Since ``\\hat{Y}`` does not depend on ``Y`` (it is only a +derivative ``∂R/∂Y``. Since ``\\hat{Y}`` does not depend on ``Y`` (it is only a function of the state at or before time ``t``), this derivative is ```math - E'(Y) = Δt * T_{imp}'(Y) - I. + R'(Y) = Δt * T_{imp}'(Y) - I. ``` -In addition, we must specify how to divide ``E(Y)`` by this derivative, i.e., +In addition, we must specify how to divide ``R(Y)`` by this derivative, i.e., how to solve the linear equation ```math - E'(Y) * ΔY = E(Y). + R'(Y) * ΔY = R(Y). ``` Note: This equation comes from assuming that there is some ``ΔY`` such that -``E(Y - ΔY) = 0`` and making the first-order approximation +``R(Y - ΔY) = 0`` and making the first-order approximation ```math - E(Y - ΔY) \\approx E(Y) - E'(Y) * ΔY. + R(Y - ΔY) \\approx R(Y) - R'(Y) * ΔY. ``` After initializing ``Y`` to ``Y[0] = \\hat{Y}``, Newton's method executes the following steps: -- Compute the derivative ``E'(Y[0])``. -- Compute the implicit tendency ``T_{imp}(Y[0])`` and use it to get ``E(Y[0])``. -- Solve the linear equation ``E'(Y[0]) * ΔY[0] = E(Y[0])`` for ``ΔY[0]``. +- Compute the derivative ``R'(Y[0])``. +- Compute the implicit tendency ``T_{imp}(Y[0])`` and use it to get ``R(Y[0])``. +- Solve the linear equation ``R'(Y[0]) * ΔY[0] = R(Y[0])`` for ``ΔY[0]``. - Update ``Y`` to ``Y[1] = Y[0] - ΔY[0]``. If the number of Newton iterations is limited to 1, this new value of ``Y`` is taken to be the solution of the implicit equation. Otherwise, this sequence of steps is repeated, i.e., ``ΔY[1]`` is computed and used to update ``Y`` to ``Y[2] = Y[1] - ΔY[1]``, then ``ΔY[2]`` is computed and used to update ``Y`` to ``Y[3] = Y[2] - ΔY[2]``, and so on. The iterative process is terminated either -when the error ``E(Y)`` is sufficiently close to 0 (according to the convergence -condition passed to Newton's method), or when the maximum number of iterations -is reached. - -# Arguments - -- `Y::FieldVector`: the state of the simulation -- `atmos::AtmosModel`: the model configuration -- `approximate_solve_iters::Int`: number of iterations to take for the - approximate linear solve required when `diffusion_flag = UseDerivative()` -- `diffusion_flag::DerivativeFlag`: whether the derivative of the - diffusion tendency with respect to the quantities being diffused should be - computed or approximated as 0; must be either `UseDerivative()` or - `IgnoreDerivative()` instead of a `Bool` to ensure type-stability -- `topography_flag::DerivativeFlag`: whether the derivative of vertical - contravariant velocity with respect to horizontal covariant velocity should be - computed or approximated as 0; must be either `UseDerivative()` or - `IgnoreDerivative()` instead of a `Bool` to ensure type-stability -- `sgs_advection_flag::DerivativeFlag`: whether the derivative of the - subgrid-scale advection tendency with respect to the updraft quantities should be - computed or approximated as 0; must be either `UseDerivative()` or - `IgnoreDerivative()` instead of a `Bool` to ensure type-stability -- `transform_flag::Bool`: whether the error should be transformed from ``E(Y)`` - to ``E(Y)/Δt``, which is required for non-Rosenbrock timestepping schemes from - OrdinaryDiffEq.jl +when the residual ``R(Y)`` is sufficiently close to 0 (according to the +convergence condition passed to Newton's method), or when the maximum number of +iterations is reached. """ -struct ImplicitEquationJacobian{ - M <: MatrixFields.FieldMatrix, - S <: MatrixFields.FieldMatrixSolver, - F1 <: DerivativeFlag, - F2 <: DerivativeFlag, - F3 <: DerivativeFlag, - F4 <: DerivativeFlag, - F5 <: DerivativeFlag, - F6 <: DerivativeFlag, - T <: Fields.FieldVector, - R <: Base.RefValue, -} - # stores the matrix E'(Y) = Δt * T_imp'(Y) - I - matrix::M - - # solves the linear equation E'(Y) * ΔY = E(Y) for ΔY - solver::S +abstract type JacobianAlgorithm end - # flags that determine how E'(Y) is approximated - diffusion_flag::F1 - topography_flag::F2 - sgs_advection_flag::F3 - sgs_entr_detr_flag::F4 - sgs_nh_pressure_flag::F5 - sgs_mass_flux_flag::F6 - - # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs - temp_b::T - temp_x::T +""" + ImplicitEquationJacobian(alg, Y, atmos, output_dir) - # required by OrdinaryDiffEq.jl to run non-Rosenbrock timestepping schemes - transform_flag::Bool - dtγ_ref::R +Wrapper for a `JacobianAlgorithm` and its cache, which it uses to update and +invert the Jacobian. The `output_dir` is the directory used for saving plots. +""" +struct ImplicitEquationJacobian{A <: JacobianAlgorithm, C} + alg::A + cache::C end - -function Base.zero(jac::ImplicitEquationJacobian) - return ImplicitEquationJacobian( - Base.zero(jac.matrix), - jac.solver, - jac.diffusion_flag, - jac.topography_flag, - jac.sgs_advection_flag, - jac.sgs_entr_detr_flag, - jac.sgs_nh_pressure_flag, - jac.sgs_mass_flux_flag, - jac.temp_b, - jac.temp_x, - jac.transform_flag, - jac.dtγ_ref, +function ImplicitEquationJacobian(alg, Y, atmos, output_dir) + FT = eltype(Y) + DA = ClimaComms.array_type(Y) + Yₜ = similar(Y) + + n_columns = Fields.ncolumns(Y.c) + n_εs = length(first_column(Y)) + column_vectors = DA{FT}(undef, n_columns, n_εs) + column_matrix = DA{FT}(undef, n_εs, n_εs) + column_vector = DA{FT}(undef, n_εs) + + is_cpu = DA <: Array + cpu_column_matrix = is_cpu ? column_matrix : Array(column_matrix) + cpu_column_vector = is_cpu ? column_vector : Array(column_vector) + + plot_cache = (; + Yₜ, + column_vectors, + column_matrix, + column_vector, + cpu_column_matrix, + cpu_column_vector, + output_dir, ) + krylov_cache = (; ΔY_krylov = similar(Y), R_krylov = similar(Y)) + cache = (; jacobian_cache(alg, Y, atmos)..., plot_cache..., krylov_cache...) + return ImplicitEquationJacobian(alg, cache) end -function ImplicitEquationJacobian( - Y, - atmos; - approximate_solve_iters = 1, - diffusion_flag = atmos.diff_mode == Implicit() ? UseDerivative() : - IgnoreDerivative(), - topography_flag = has_topography(axes(Y.c)) ? UseDerivative() : - IgnoreDerivative(), - sgs_advection_flag = atmos.sgs_adv_mode == Implicit() ? UseDerivative() : - IgnoreDerivative(), - sgs_entr_detr_flag = atmos.sgs_entr_detr_mode == Implicit() ? - UseDerivative() : IgnoreDerivative(), - sgs_nh_pressure_flag = atmos.sgs_nh_pressure_mode == Implicit() ? - UseDerivative() : IgnoreDerivative(), - sgs_mass_flux_flag = atmos.sgs_mf_mode == Implicit() ? UseDerivative() : - IgnoreDerivative(), - transform_flag = false, -) - FT = Spaces.undertype(axes(Y.c)) - CTh = CTh_vector_type(axes(Y.c)) - - DiagonalRow = DiagonalMatrixRow{FT} - TridiagonalRow = TridiagonalMatrixRow{FT} - BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}} - TridiagonalRow_ACTh = TridiagonalMatrixRow{Adjoint{FT, CTh{FT}}} - BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}} - BidiagonalRow_C3xACTh = - BidiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CTh{FT})')} - DiagonalRow_C3xACT3 = - DiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} - TridiagonalRow_C3xACT3 = - TridiagonalMatrixRow{typeof(zero(C3{FT}) * zero(CT3{FT})')} - - is_in_Y(name) = MatrixFields.has_field(Y, name) - - ρq_tot_if_available = is_in_Y(@name(c.ρq_tot)) ? (@name(c.ρq_tot),) : () - ρatke_if_available = - is_in_Y(@name(c.sgs⁰.ρatke)) ? (@name(c.sgs⁰.ρatke),) : () - sfc_if_available = is_in_Y(@name(sfc)) ? (@name(sfc),) : () - - tracer_names = ( - @name(c.ρq_tot), - @name(c.ρq_liq), - @name(c.ρq_ice), - @name(c.ρq_rai), - @name(c.ρq_sno), - ) - available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names) - - # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0, - # which means that multiplying inv(-1) by a Float32 will yield a Float64. - identity_blocks = MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - (@name(c.ρ), sfc_if_available...), - ) - - active_scalar_names = (@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...) - advection_blocks = ( - ( - use_derivative(topography_flag) ? - MatrixFields.unrolled_map( - name -> - (name, @name(c.uₕ)) => - similar(Y.c, TridiagonalRow_ACTh), - active_scalar_names, - ) : () - )..., - MatrixFields.unrolled_map( - name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3), - active_scalar_names, - )..., - MatrixFields.unrolled_map( - name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3), - active_scalar_names, - )..., - (@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACTh), - (@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3), - ) - - diffused_scalar_names = (@name(c.ρe_tot), available_tracer_names...) - diffusion_blocks = if use_derivative(diffusion_flag) - ( - MatrixFields.unrolled_map( - name -> (name, @name(c.ρ)) => similar(Y.c, TridiagonalRow), - (diffused_scalar_names..., ρatke_if_available...), - )..., - MatrixFields.unrolled_map( - name -> (name, name) => similar(Y.c, TridiagonalRow), - (diffused_scalar_names..., ρatke_if_available...), - )..., - ( - is_in_Y(@name(c.ρq_tot)) ? - ( - (@name(c.ρe_tot), @name(c.ρq_tot)) => - similar(Y.c, TridiagonalRow), - ) : () - )..., - (@name(c.uₕ), @name(c.uₕ)) => - !isnothing(atmos.turbconv_model) || - !disable_momentum_vertical_diffusion(atmos.vert_diff) ? - similar(Y.c, TridiagonalRow) : FT(-1) * I, - ) - elseif atmos.moisture_model isa DryModel - MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - (diffused_scalar_names..., ρatke_if_available..., @name(c.uₕ)), - ) - else - ( - MatrixFields.unrolled_map( - name -> (name, name) => similar(Y.c, TridiagonalRow), - diffused_scalar_names, - )..., - (@name(c.ρe_tot), @name(c.ρq_tot)) => - similar(Y.c, TridiagonalRow), - MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - (ρatke_if_available..., @name(c.uₕ)), - )..., - ) - end - - sgs_scalar_names = ( - @name(c.sgsʲs.:(1).q_tot), - @name(c.sgsʲs.:(1).q_liq), - @name(c.sgsʲs.:(1).q_ice), - @name(c.sgsʲs.:(1).q_rai), - @name(c.sgsʲs.:(1).q_sno), - @name(c.sgsʲs.:(1).mse), - @name(c.sgsʲs.:(1).ρa) +# ClimaTimeSteppers.jl calls zero(jac_prototype) to initialize the Jacobian, but +# we don't need to allocate a second Jacobian for this (in particular, the exact +# Jacobian can be very expensive to allocate). +Base.zero(jacobian::ImplicitEquationJacobian) = jacobian + +# These are either called by ClimaTimeSteppers.jl before each linear solve, or +# by a callback once every dt_update_exact_jacobian. +NVTX.@annotate update_jacobian!(jacobian, Y, p, dtγ, t) = + update_jacobian!(jacobian.alg, jacobian.cache, Y, p, eltype(Y)(dtγ), t) +NVTX.@annotate update_and_check_jacobian!(jacobian, Y, p, dtγ, t) = + update_and_check_jacobian!( + jacobian.alg, + jacobian.cache, + Y, + p, + eltype(Y)(dtγ), + t, ) - available_sgs_scalar_names = - MatrixFields.unrolled_filter(is_in_Y, sgs_scalar_names) - - sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX - @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1 - - if use_derivative(sgs_advection_flag) - ( - MatrixFields.unrolled_map( - name -> (name, name) => similar(Y.c, TridiagonalRow), - available_sgs_scalar_names, - )..., - (@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)) => - similar(Y.c, DiagonalRow), - (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)) => - similar(Y.c, TridiagonalRow), - (@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)) => - similar(Y.c, TridiagonalRow), - (@name(c.sgsʲs.:(1).ρa), @name(f.sgsʲs.:(1).u₃)) => - similar(Y.c, BidiagonalRow_ACT3), - (@name(c.sgsʲs.:(1).mse), @name(f.sgsʲs.:(1).u₃)) => - similar(Y.c, BidiagonalRow_ACT3), - (@name(c.sgsʲs.:(1).q_tot), @name(f.sgsʲs.:(1).u₃)) => - similar(Y.c, BidiagonalRow_ACT3), - (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)) => - similar(Y.f, BidiagonalRow_C3), - (@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)) => - similar(Y.f, BidiagonalRow_C3), - (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => - similar(Y.f, TridiagonalRow_C3xACT3), - ) - else - ( - MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - available_sgs_scalar_names, - )..., - (@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)) => - !isnothing(atmos.rayleigh_sponge) ? - similar(Y.f, DiagonalRow_C3xACT3) : FT(-1) * I, - ) - end - else - () - end - sgs_massflux_blocks = if atmos.turbconv_model isa PrognosticEDMFX - @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1 - if use_derivative(sgs_mass_flux_flag) - ( - (@name(c.ρe_tot), @name(c.sgsʲs.:(1).mse)) => - similar(Y.c, TridiagonalRow), - (@name(c.ρq_tot), @name(c.sgsʲs.:(1).q_tot)) => - similar(Y.c, TridiagonalRow), - (@name(c.ρe_tot), @name(f.sgsʲs.:(1).u₃)) => - similar(Y.c, BidiagonalRow_ACT3), - (@name(c.ρq_tot), @name(f.sgsʲs.:(1).u₃)) => - similar(Y.c, BidiagonalRow_ACT3), - (@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρa)) => - similar(Y.c, TridiagonalRow), - (@name(c.ρq_tot), @name(c.sgsʲs.:(1).ρa)) => - similar(Y.c, TridiagonalRow), - ) - else - () - end - else - () - end - - matrix = MatrixFields.FieldMatrix( - identity_blocks..., - sgs_advection_blocks..., - advection_blocks..., - diffusion_blocks..., - sgs_massflux_blocks..., - ) - - sgs_u³_names_if_available = if atmos.turbconv_model isa PrognosticEDMFX - (@name(f.sgsʲs.:(1).u₃),) - else - () - end - - names₁_group₁ = (@name(c.ρ), sfc_if_available...) - names₁_group₂ = (available_tracer_names..., ρatke_if_available...) - names₁_group₃ = (@name(c.ρe_tot),) - names₁ = ( - names₁_group₁..., - available_sgs_scalar_names..., - names₁_group₂..., - names₁_group₃..., - ) - - alg₂ = MatrixFields.BlockLowerTriangularSolve( - @name(c.uₕ), - sgs_u³_names_if_available..., - ) - alg = - if use_derivative(diffusion_flag) || - use_derivative(sgs_advection_flag) || - !(atmos.moisture_model isa DryModel) - alg₁_subalg₂ = - if atmos.turbconv_model isa PrognosticEDMFX && - use_derivative(sgs_advection_flag) - diff_subalg = - use_derivative(diffusion_flag) ? - (; - alg₂ = MatrixFields.BlockLowerTriangularSolve( - names₁_group₂..., - ) - ) : (;) - (; - alg₂ = MatrixFields.BlockLowerTriangularSolve( - # TODO: What needs to be changed here for 1M? - @name(c.sgsʲs.:(1).q_tot); - alg₂ = MatrixFields.BlockLowerTriangularSolve( - @name(c.sgsʲs.:(1).mse); - alg₂ = MatrixFields.BlockLowerTriangularSolve( - @name(c.sgsʲs.:(1).ρa); - diff_subalg..., - ), - ), - ) - ) - else - is_in_Y(@name(c.ρq_tot)) ? - (; - alg₂ = MatrixFields.BlockLowerTriangularSolve( - names₁_group₂..., - ) - ) : (;) - end - alg₁ = MatrixFields.BlockLowerTriangularSolve( - names₁_group₁...; - alg₁_subalg₂..., - ) - MatrixFields.ApproximateBlockArrowheadIterativeSolve( - names₁...; - alg₁, - alg₂, - P_alg₁ = MatrixFields.MainDiagonalPreconditioner(), - n_iters = approximate_solve_iters, - ) - else - MatrixFields.BlockArrowheadSolve(names₁...; alg₂) - end - - return ImplicitEquationJacobian( - matrix, - MatrixFields.FieldMatrixSolver(alg, matrix, Y), - diffusion_flag, - topography_flag, - sgs_advection_flag, - sgs_entr_detr_flag, - sgs_nh_pressure_flag, - sgs_mass_flux_flag, - similar(Y), - similar(Y), - transform_flag, - Ref{FT}(), - ) -end - -# We only use A, but ClimaTimeSteppers.jl require us to -# pass jac_prototype and then call similar(jac_prototype) to -# obtain A. This is a workaround to avoid unnecessary allocations. -Base.similar(A::ImplicitEquationJacobian) = A - -# This method specifies how to solve the equation E'(Y) * ΔY = E(Y) for ΔY. -NVTX.@annotate function ldiv!( - x::Fields.FieldVector, - A::ImplicitEquationJacobian, - b::Fields.FieldVector, -) - MatrixFields.field_matrix_solve!(A.solver, x, A.matrix, b) - if A.transform_flag - @. x *= -A.dtγ_ref[] - end -end - -# This method for ldiv! is called by Krylov.jl from inside ClimaTimeSteppers.jl. -# See https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a -# related issue that requires the same workaround. -NVTX.@annotate function ldiv!( - x::AbstractVector, - A::ImplicitEquationJacobian, - b::AbstractVector, +# This is called by ClimaTimeSteppers.jl before each linear solve. +NVTX.@annotate LinearAlgebra.ldiv!( + ΔY::Fields.FieldVector, + jacobian::ImplicitEquationJacobian, + R::Fields.FieldVector, +) = invert_jacobian!(jacobian.alg, jacobian.cache, ΔY, R) + +# This is called by Krylov.jl from inside ClimaTimeSteppers.jl. See +# https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/605 for a related +# issue that requires the same workaround. +function LinearAlgebra.ldiv!( + ΔY::AbstractVector, + jacobian::ImplicitEquationJacobian, + R::AbstractVector, ) - A.temp_b .= b - ldiv!(A.temp_x, A, A.temp_b) - x .= A.temp_x + (; ΔY_krylov, R_krylov) = jacobian.cache + R_krylov .= R + LinearAlgebra.ldiv!(ΔY_krylov, jacobian, R_krylov) + ΔY .= ΔY_krylov end -# This function is used by DiffEqBase.jl instead of ldiv!. -linsolve!(::Type{Val{:init}}, f, u0; kwargs...) = _linsolve! -_linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b) +# This is called by a callback once every dt_save_jacobian. +NVTX.@annotate function save_jacobian!(jacobian, Y, p, dtγ, t) + (; Yₜ, column_vectors, column_vector) = jacobian.cache -# This method specifies how to compute E'(Y), which is referred to as "Wfact" in -# DiffEqBase.jl. -NVTX.@annotate function Wfact!(A, Y, p, dtγ, t) - # Remove unnecessary values from p to avoid allocations in bycolumn. - p′ = (; - p.precomputed.ᶜspecific, - p.precomputed.ᶜK, - p.precomputed.ᶠu³, - p.precomputed.ᶜts, - p.precomputed.ᶜp, - p.precomputed.ᶜwₜqₜ, - p.precomputed.ᶜwₕhₜ, - ( - p.atmos.moisture_model isa NonEquilMoistModel ? - (; p.precomputed.ᶜwₗ, p.precomputed.ᶜwᵢ) : (;) - )..., - ( - p.atmos.precip_model isa Microphysics1Moment ? - (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;) - )..., - p.precomputed.ᶜh_tot, - ( - use_derivative(A.diffusion_flag) ? - (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;) - )..., - ( - use_derivative(A.diffusion_flag) && - p.atmos.turbconv_model isa AbstractEDMF ? - (; p.precomputed.ᶜtke⁰, p.precomputed.ᶜmixing_length) : (;) - )..., - ( - use_derivative(A.diffusion_flag) && - p.atmos.turbconv_model isa PrognosticEDMFX ? - (; p.precomputed.ᶜρa⁰) : (;) - )..., - ( - use_derivative(A.sgs_advection_flag) && - p.atmos.turbconv_model isa PrognosticEDMFX ? - (; - p.core.ᶜgradᵥ_ᶠΦ, - p.precomputed.ᶜρʲs, - p.precomputed.ᶠu³ʲs, - p.precomputed.ᶜtsʲs, - p.precomputed.bdmr_l, - p.precomputed.bdmr_r, - p.precomputed.bdmr, - ) : (;) - )..., - ( - use_derivative(A.sgs_entr_detr_flag) && - p.atmos.turbconv_model isa PrognosticEDMFX ? - (; - p.precomputed.ᶜentrʲs, - p.precomputed.ᶜdetrʲs, - p.precomputed.ᶜturb_entrʲs, - ) : (;) - )..., - ( - use_derivative(A.sgs_nh_pressure_flag) && - p.atmos.turbconv_model isa PrognosticEDMFX ? - (; p.precomputed.ᶠu₃⁰,) : (;) - )..., - ( - use_derivative(A.sgs_mass_flux_flag) && - p.atmos.turbconv_model isa PrognosticEDMFX ? - (; p.precomputed.ᶜKʲs) : (;) - )..., - p.core.ᶜΦ, - p.core.ᶠgradᵥ_ᶜΦ, - p.scratch.ᶜtemp_scalar, - p.scratch.ᶜtemp_scalar_2, - p.scratch.ᶜtemp_C3, - p.scratch.ᶠtemp_CT3, - p.scratch.∂ᶜK_∂ᶜuₕ, - p.scratch.∂ᶜK_∂ᶠu₃, - p.scratch.ᶠp_grad_matrix, - p.scratch.ᶜadvection_matrix, - p.scratch.ᶜdiffusion_h_matrix, - p.scratch.ᶜdiffusion_h_matrix_scaled, - p.scratch.ᶜdiffusion_u_matrix, - p.scratch.ᶠbidiagonal_matrix_ct3, - p.scratch.ᶠbidiagonal_matrix_ct3_2, - p.scratch.ᶠtridiagonal_matrix_c3, - p.dt, - p.params, - p.atmos, - ) + # TODO: Add support for MPI reductions, instead of only saving from root. + ClimaComms.iamroot(ClimaComms.context(Y.c)) || return - # Convert dtγ from a Float64 to an FT. - FT = Spaces.undertype(axes(Y.c)) - dtγ′ = FT(float(dtγ)) + implicit_tendency!(Yₜ, Y, p, t) + column_vectors_to_field_vector(column_vectors, Yₜ) .= Yₜ + sum!(abs, reshape(column_vector, 1, :), column_vectors) + column_vector ./= Fields.ncolumns(Y.c) - A.dtγ_ref[] = dtγ′ - update_implicit_equation_jacobian!(A, Y, p′, dtγ′, t) + save_jacobian!(jacobian.alg, jacobian.cache, Y, eltype(Y)(dtγ), t) end -function update_implicit_equation_jacobian!(A, Y, p, dtγ, t) - dtγ = float(dtγ) - (; - matrix, - diffusion_flag, - sgs_advection_flag, - sgs_entr_detr_flag, - sgs_nh_pressure_flag, - topography_flag, - sgs_mass_flux_flag, - ) = A - (; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜh_tot) = p - (; - ᶜtemp_C3, - ∂ᶜK_∂ᶜuₕ, - ∂ᶜK_∂ᶠu₃, - ᶠp_grad_matrix, - ᶜadvection_matrix, - ᶠbidiagonal_matrix_ct3, - ᶠbidiagonal_matrix_ct3_2, - ᶠtridiagonal_matrix_c3, - ) = p - (; - ᶜdiffusion_h_matrix, - ᶜdiffusion_h_matrix_scaled, - ᶜdiffusion_u_matrix, - params, - ) = p - (; edmfx_upwinding) = p.atmos.numerics - - FT = Spaces.undertype(axes(Y.c)) - CTh = CTh_vector_type(axes(Y.c)) - one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' - - cv_d = FT(CAP.cv_d(params)) - Δcv_v = FT(CAP.cv_v(params)) - cv_d - T_0 = FT(CAP.T_0(params)) - R_d = FT(CAP.R_d(params)) - ΔR_v = FT(CAP.R_v(params)) - R_d - cp_d = FT(CAP.cp_d(params)) - Δcp_v = FT(CAP.cp_v(params)) - cp_d - # This term appears a few times in the Jacobian, and is technically - # minus ∂e_int_∂q_tot - ∂e_int_∂q_tot = T_0 * (Δcv_v - R_d) - FT(CAP.e_int_v0(params)) - thermo_params = CAP.thermodynamics_params(params) - - ᶜρ = Y.c.ρ - ᶜuₕ = Y.c.uₕ - ᶠu₃ = Y.f.u₃ - ᶜJ = Fields.local_geometry_field(Y.c).J - ᶠJ = Fields.local_geometry_field(Y.f).J - ᶜgⁱʲ = Fields.local_geometry_field(Y.c).gⁱʲ - ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ - ᶠlg = Fields.local_geometry_field(Y.f) - - ᶜkappa_m = p.ᶜtemp_scalar - @. ᶜkappa_m = - TD.gas_constant_air(thermo_params, ᶜts) / TD.cv_m(thermo_params, ᶜts) - - ᶜ∂kappa_m∂q_tot = p.ᶜtemp_scalar_2 - # Using abs2 because ^2 results in allocation - @. ᶜ∂kappa_m∂q_tot = - ( - ΔR_v * TD.cv_m(thermo_params, ᶜts) - - Δcv_v * TD.gas_constant_air(thermo_params, ᶜts) - ) / abs2(TD.cv_m(thermo_params, ᶜts)) - - if use_derivative(topography_flag) - @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow( - adjoint(CTh(ᶜuₕ)) + adjoint(ᶜinterp(ᶠu₃)) * g³ʰ(ᶜgⁱʲ), - ) - else - @. ∂ᶜK_∂ᶜuₕ = DiagonalMatrixRow(adjoint(CTh(ᶜuₕ))) - end - @. ∂ᶜK_∂ᶠu₃ = - ᶜinterp_matrix() ⋅ DiagonalMatrixRow(adjoint(CT3(ᶠu₃))) + - DiagonalMatrixRow(adjoint(CT3(ᶜuₕ))) ⋅ ᶜinterp_matrix() - - @. ᶠp_grad_matrix = DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() - - @. ᶜadvection_matrix = - -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) - - if use_derivative(topography_flag) - ∂ᶜρ_err_∂ᶜuₕ = matrix[@name(c.ρ), @name(c.uₕ)] - @. ∂ᶜρ_err_∂ᶜuₕ = - dtγ * ᶜadvection_matrix ⋅ ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ - DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) +contains_any_fields(::Union{Fields.Field, Fields.FieldVector}) = true +contains_any_fields(x::T) where {T} = + fieldcount(T) == 0 ? false : unrolled_any(StaticOneTo(fieldcount(T))) do i + contains_any_fields(getfield(x, i)) end - ∂ᶜρ_err_∂ᶠu₃ = matrix[@name(c.ρ), @name(f.u₃)] - @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) - tracer_info = ( - (@name(c.ρe_tot), @name(ᶜh_tot)), - (@name(c.ρq_tot), @name(ᶜspecific.q_tot)), +first_column(x::MatrixFields.FieldNameDict) = + MatrixFields.FieldNameDict(x.keys, first_column(x.entries)) +first_column(x::Union{Fields.Field, Fields.FieldVector}) = + Fields.column(x, 1, 1, 1) +first_column(x::Union{Tuple, NamedTuple}) = unrolled_map(first_column, x) +first_column(x::T) where {T} = + fieldcount(T) == 0 || !contains_any_fields(x) ? x : + T.name.wrapper( + ntuple(i -> first_column(getfield(x, i)), Val(fieldcount(T)))..., ) - MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name) - MatrixFields.has_field(Y, ρχ_name) || return - ᶜχ = MatrixFields.get_field(p, χ_name) - if use_derivative(topography_flag) - ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] - end - ∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)] - use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ = - dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ)) ⋅ - ᶠwinterp_matrix(ᶜJ * ᶜρ) ⋅ DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ)) - @. ∂ᶜρχ_err_∂ᶠu₃ = - dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(ᶠinterp(ᶜχ) * g³³(ᶠgⁱʲ)) - end - - ∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)] - ∂ᶠu₃_err_∂ᶜρe_tot = matrix[@name(f.u₃), @name(c.ρe_tot)] - @. ∂ᶠu₃_err_∂ᶜρ = - dtγ * ( - ᶠp_grad_matrix ⋅ - DiagonalMatrixRow(ᶜkappa_m * (T_0 * cp_d - ᶜK - ᶜΦ)) + - DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / abs2(ᶠinterp(ᶜρ))) ⋅ - ᶠinterp_matrix() - ) - @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) - if MatrixFields.has_field(Y, @name(c.ρq_tot)) - ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] - @. ∂ᶠu₃_err_∂ᶜρq_tot = - dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(( - ᶜkappa_m * ∂e_int_∂q_tot + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * ᶜspecific.q_tot - ) - )) - end - ∂ᶠu₃_err_∂ᶜuₕ = matrix[@name(f.u₃), @name(c.uₕ)] - ∂ᶠu₃_err_∂ᶠu₃ = matrix[@name(f.u₃), @name(f.u₃)] - I_u₃ = DiagonalMatrixRow(one_C3xACT3) - @. ∂ᶠu₃_err_∂ᶜuₕ = - dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ ∂ᶜK_∂ᶜuₕ - rs = p.atmos.rayleigh_sponge - ᶠz = Fields.coordinate_field(Y.f).z - zmax = z_max(axes(Y.f)) - if rs isa RayleighSponge - @. ∂ᶠu₃_err_∂ᶠu₃ = - dtγ * ( - ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ - ∂ᶜK_∂ᶠu₃ + - DiagonalMatrixRow(-β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,)) - ) - (I_u₃,) +function first_column_str(Y) + coord_field = + Fields.coordinate_field(Fields.level(Fields.column(Y.c, 1, 1, 1), 1)) + coord = + ClimaComms.allowscalar(getindex, ClimaComms.device(Y.c), coord_field) + round_value(value) = round(value; sigdigits = 3) + return if coord isa Geometry.XZPoint + "x = $(round_value(coord.x)) Meters" + elseif coord isa Geometry.XYZPoint + "x = $(round_value(coord.x)) Meters, y = $(round_value(coord.y)) Meters" + elseif coord isa Geometry.LatLongZPoint + "lat = $(round_value(coord.lat))°, long = $(round_value(coord.long))°" else - @. ∂ᶠu₃_err_∂ᶠu₃ = - dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(-(ᶜkappa_m) * ᶜρ) ⋅ - ∂ᶜK_∂ᶠu₃ - (I_u₃,) + error("Unrecognized coordinate type $(typeof(coord))") end +end - - tracer_info = ( - (@name(c.ρq_liq), @name(q_liq), @name(ᶜwₗ)), - (@name(c.ρq_ice), @name(q_ice), @name(ᶜwᵢ)), - (@name(c.ρq_rai), @name(q_rai), @name(ᶜwᵣ)), - (@name(c.ρq_sno), @name(q_sno), @name(ᶜwₛ)), - ) - if !(p.atmos.moisture_model isa DryModel) || use_derivative(diffusion_flag) - ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)] - @. ∂ᶜρe_tot_err_∂ᶜρe_tot = zero(typeof(∂ᶜρe_tot_err_∂ᶜρe_tot)) - (I,) +function save_cached_column_matrix_and_vector!(cache, file_name, title, t) + (; Yₜ, column_matrix, column_vector, output_dir) = cache + (; cpu_column_matrix, cpu_column_vector) = cache + if !(column_matrix isa Array) + copyto!(cpu_column_matrix, column_matrix) + copyto!(cpu_column_vector, column_vector) end - - if !(p.atmos.moisture_model isa DryModel) - #TODO: tetsing explicit vs implicit - #@. ∂ᶜρe_tot_err_∂ᶜρe_tot += - # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ - # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ - # DiagonalMatrixRow( - # -(1 + ᶜkappa_m) / ᶜρ * ifelse( - # ᶜh_tot == 0, - # (Geometry.WVector(FT(0)),), - # p.ᶜwₕhₜ / ᶜh_tot, - # ), - # ) - - ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] - @. ∂ᶜρe_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρe_tot_err_∂ᶜρq_tot)) - #TODO: tetsing explicit vs implicit - #@. ∂ᶜρe_tot_err_∂ᶜρq_tot = - # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ - # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ - # DiagonalMatrixRow( - # -(ᶜkappa_m) * ∂e_int_∂q_tot / ᶜρ * ifelse( - # ᶜh_tot == 0, - # (Geometry.WVector(FT(0)),), - # p.ᶜwₕhₜ / ᶜh_tot, - # ), - # ) - - ∂ᶜρq_tot_err_∂ᶜρq_tot = matrix[@name(c.ρq_tot), @name(c.ρq_tot)] - @. ∂ᶜρq_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρq_tot_err_∂ᶜρq_tot)) - (I,) - #TODO: tetsing explicit vs implicit - #@. ∂ᶜρq_tot_err_∂ᶜρq_tot = - # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ - # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ - # DiagonalMatrixRow( - # -1 / ᶜρ * ifelse( - # ᶜspecific.q_tot == 0, - # (Geometry.WVector(FT(0)),), - # p.ᶜwₜqₜ / ᶜspecific.q_tot, - # ), - # ) - (I,) - - MatrixFields.unrolled_foreach(tracer_info) do (ρqₚ_name, _, wₚ_name) - MatrixFields.has_field(Y, ρqₚ_name) || return - ∂ᶜρqₚ_err_∂ᶜρqₚ = matrix[ρqₚ_name, ρqₚ_name] - ᶜwₚ = MatrixFields.get_field(p, wₚ_name) - @. ∂ᶜρqₚ_err_∂ᶜρqₚ = - dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ - ᶠright_bias_matrix() ⋅ - DiagonalMatrixRow(-Geometry.WVector(ᶜwₚ) / ᶜρ) - (I,) - end - + file_name = joinpath(output_dir, file_name * ".hdf5") + HDF5Writer(file_name, ClimaComms.context(Yₜ.c); overwrite = false) do writer + key = string(float(t)) + haskey(writer.file, key) && return + group = HDF5.create_group(writer.file, key) + group["title"] = "$title, t = $(time_and_units_str(float(t)))" + group["∂Yₜ_∂Y"] = cpu_column_matrix + group["Yₜ"] = cpu_column_vector end - - if use_derivative(diffusion_flag) - α_vert_diff_tracer = CAP.α_vert_diff_tracer(params) - (; ᶜK_h, ᶜK_u) = p - @. ᶜdiffusion_h_matrix = - ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h)) ⋅ - ᶠgradᵥ_matrix() - @. ᶜdiffusion_h_matrix_scaled = - ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( - ᶠinterp(ᶜρ) * ᶠinterp(α_vert_diff_tracer * ᶜK_h), - ) ⋅ ᶠgradᵥ_matrix() - if ( - MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) || - !isnothing(p.atmos.turbconv_model) || - !disable_momentum_vertical_diffusion(p.atmos.vert_diff) - ) - @. ᶜdiffusion_u_matrix = - ᶜadvdivᵥ_matrix() ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u)) ⋅ ᶠgradᵥ_matrix() - end - - ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] - @. ∂ᶜρe_tot_err_∂ᶜρ = - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( - ( - -(1 + ᶜkappa_m) * ᶜspecific.e_tot - - ᶜkappa_m * ∂e_int_∂q_tot * ᶜspecific.q_tot - ) / ᶜρ, - ) - @. ∂ᶜρe_tot_err_∂ᶜρe_tot += - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) - - if MatrixFields.has_field(Y, @name(c.ρq_tot)) - ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] - ∂ᶜρq_tot_err_∂ᶜρ = matrix[@name(c.ρq_tot), @name(c.ρ)] - @. ∂ᶜρe_tot_err_∂ᶜρq_tot += - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(( - ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * ᶜspecific.q_tot - ) - )) - @. ∂ᶜρq_tot_err_∂ᶜρ = - dtγ * ᶜdiffusion_h_matrix ⋅ - DiagonalMatrixRow(-(ᶜspecific.q_tot) / ᶜρ) - @. ∂ᶜρq_tot_err_∂ᶜρq_tot += - dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) - end - - - MatrixFields.unrolled_foreach(tracer_info) do (ρq_name, q_name, _) - MatrixFields.has_field(Y, ρq_name) || return - ᶜq = MatrixFields.get_field(ᶜspecific, q_name) - ∂ᶜρq_err_∂ᶜρ = matrix[ρq_name, @name(c.ρ)] - ∂ᶜρq_err_∂ᶜρq = matrix[ρq_name, ρq_name] - ᶜtridiagonal_matrix_scalar = ifelse( - q_name in (@name(q_rai), @name(q_sno)), - ᶜdiffusion_h_matrix_scaled, - ᶜdiffusion_h_matrix, - ) - @. ∂ᶜρq_err_∂ᶜρ = - dtγ * ᶜtridiagonal_matrix_scalar ⋅ DiagonalMatrixRow(-(ᶜq) / ᶜρ) - @. ∂ᶜρq_err_∂ᶜρq += - dtγ * ᶜtridiagonal_matrix_scalar ⋅ DiagonalMatrixRow(1 / ᶜρ) - end - - if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) - turbconv_params = CAP.turbconv_params(params) - c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜtke⁰, ᶜmixing_length, dt) = p - ᶜρa⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? p.ᶜρa⁰ : ᶜρ - ᶜρatke⁰ = Y.c.sgs⁰.ρatke - - @inline dissipation_rate(tke⁰, mixing_length) = - tke⁰ >= 0 ? c_d * sqrt(tke⁰) / max(mixing_length, 1) : - 1 / float(dt) - @inline ∂dissipation_rate_∂tke⁰(tke⁰, mixing_length) = - tke⁰ > 0 ? c_d / (2 * max(mixing_length, 1) * sqrt(tke⁰)) : - typeof(tke⁰)(0) - - ᶜdissipation_matrix_diagonal = p.ᶜtemp_scalar - @. ᶜdissipation_matrix_diagonal = - ᶜρatke⁰ * ∂dissipation_rate_∂tke⁰(ᶜtke⁰, ᶜmixing_length) - - ∂ᶜρatke⁰_err_∂ᶜρ = matrix[@name(c.sgs⁰.ρatke), @name(c.ρ)] - ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = - matrix[@name(c.sgs⁰.ρatke), @name(c.sgs⁰.ρatke)] - @. ∂ᶜρatke⁰_err_∂ᶜρ = - dtγ * ( - ᶜdiffusion_u_matrix - - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) - ) ⋅ DiagonalMatrixRow(-(ᶜtke⁰) / ᶜρa⁰) - @. ∂ᶜρatke⁰_err_∂ᶜρatke⁰ = - dtγ * ( - ( - ᶜdiffusion_u_matrix - - DiagonalMatrixRow(ᶜdissipation_matrix_diagonal) - ) ⋅ DiagonalMatrixRow(1 / ᶜρa⁰) - - DiagonalMatrixRow(dissipation_rate(ᶜtke⁰, ᶜmixing_length)) - ) - (I,) - end - - if ( - !isnothing(p.atmos.turbconv_model) || - !disable_momentum_vertical_diffusion(p.atmos.vert_diff) - ) - ∂ᶜuₕ_err_∂ᶜuₕ = matrix[@name(c.uₕ), @name(c.uₕ)] - @. ∂ᶜuₕ_err_∂ᶜuₕ = - dtγ * DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜdiffusion_u_matrix - (I,) - end - - end - - if p.atmos.turbconv_model isa PrognosticEDMFX - - if use_derivative(sgs_advection_flag) - (; ᶜgradᵥ_ᶠΦ, ᶜρʲs, ᶠu³ʲs, ᶜtsʲs) = p - (; bdmr_l, bdmr_r, bdmr) = p - is_third_order = edmfx_upwinding == Val(:third_order) - ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1 - ᶠset_upwind_bcs = Operators.SetBoundaryOperator(; - top = Operators.SetValue(zero(CT3{FT})), - bottom = Operators.SetValue(zero(CT3{FT})), - ) # Need to wrap ᶠupwind in this for well-defined boundaries. - UpwindMatrixRowType = - is_third_order ? QuaddiagonalMatrixRow : BidiagonalMatrixRow - ᶠupwind_matrix = is_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix - ᶠset_upwind_matrix_bcs = Operators.SetBoundaryOperator(; - top = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), - bottom = Operators.SetValue(zero(UpwindMatrixRowType{CT3{FT}})), - ) # Need to wrap ᶠupwind_matrix in this for well-defined boundaries. - - ᶠu³ʲ_data = ᶠu³ʲs.:(1).components.data.:1 - ᶜkappa_mʲ = p.ᶜtemp_scalar - @. ᶜkappa_mʲ = - TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) / - TD.cv_m(thermo_params, ᶜtsʲs.:(1)) - - # Note this is the derivative of R_m / cp_m with respect to q_tot - # but we call it ∂kappa_m∂q_totʲ - ᶜ∂kappa_m∂q_totʲ = p.ᶜtemp_scalar_2 - @. ᶜ∂kappa_m∂q_totʲ = - ( - ΔR_v * TD.cp_m(thermo_params, ᶜtsʲs.:(1)) - - Δcp_v * TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) - ) / abs2(TD.cp_m(thermo_params, ᶜtsʲs.:(1))) - - turbconv_params = CAP.turbconv_params(params) - α_b = CAP.pressure_normalmode_buoy_coeff1(turbconv_params) - - ∂ᶜq_totʲ_err_∂ᶜq_totʲ = - matrix[@name(c.sgsʲs.:(1).q_tot), @name(c.sgsʲs.:(1).q_tot)] - @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ = - dtγ * ( - DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - - ᶜadvdivᵥ_matrix() ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) - ) - (I,) - - ∂ᶜq_totʲ_err_∂ᶠu₃ʲ = - matrix[@name(c.sgsʲs.:(1).q_tot), @name(f.sgsʲs.:(1).u₃)] - - @. ∂ᶜq_totʲ_err_∂ᶠu₃ʲ = - dtγ * ( - -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - ᶠset_upwind_bcs( - ᶠupwind(CT3(sign(ᶠu³ʲ_data)), Y.c.sgsʲs.:(1).q_tot), - ) * adjoint(C3(sign(ᶠu³ʲ_data))), - ) + - DiagonalMatrixRow(Y.c.sgsʲs.:(1).q_tot) ⋅ ᶜadvdivᵥ_matrix() - ) ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) - - - ∂ᶜmseʲ_err_∂ᶜq_totʲ = - matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot)] - @. ∂ᶜmseʲ_err_∂ᶜq_totʲ = - dtγ * ( - -DiagonalMatrixRow( - adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * ᶜgradᵥ_ᶠΦ * Y.c.ρ / ᶜp * - ( - (ᶜkappa_mʲ / (ᶜkappa_mʲ + 1) * ∂e_int_∂q_tot) + - ᶜ∂kappa_m∂q_totʲ * ( - Y.c.sgsʲs.:(1).mse - ᶜΦ + - cp_d * T_0 + - ∂e_int_∂q_tot * Y.c.sgsʲs.:(1).q_tot - ) - ), - ) - ) - - ∂ᶜmseʲ_err_∂ᶠu₃ʲ = - matrix[@name(c.sgsʲs.:(1).mse), @name(f.sgsʲs.:(1).u₃)] - @. ∂ᶜmseʲ_err_∂ᶠu₃ʲ = - dtγ * ( - -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - ᶠset_upwind_bcs( - ᶠupwind(CT3(sign(ᶠu³ʲ_data)), Y.c.sgsʲs.:(1).mse), - ) * adjoint(C3(sign(ᶠu³ʲ_data))), - ) + - DiagonalMatrixRow(Y.c.sgsʲs.:(1).mse) ⋅ ᶜadvdivᵥ_matrix() - ) ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) - - ∂ᶜmseʲ_err_∂ᶜmseʲ = - matrix[@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).mse)] - @. ∂ᶜmseʲ_err_∂ᶜmseʲ = - dtγ * ( - DiagonalMatrixRow(ᶜadvdivᵥ(ᶠu³ʲs.:(1))) - - ᶜadvdivᵥ_matrix() ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) - - DiagonalMatrixRow( - adjoint(ᶜinterp(ᶠu³ʲs.:(1))) * - ᶜgradᵥ_ᶠΦ * - Y.c.ρ * - ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - ) - (I,) - - ∂ᶜρaʲ_err_∂ᶜq_totʲ = - matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).q_tot)] - @. ᶠbidiagonal_matrix_ct3 = - DiagonalMatrixRow( - ᶠset_upwind_bcs( - ᶠupwind( - ᶠu³ʲs.:(1), - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ), - ) / ᶠJ, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( - ᶜJ * (ᶜρʲs.:(1))^2 / ᶜp * ( - ᶜkappa_mʲ / (ᶜkappa_mʲ + 1) * ∂e_int_∂q_tot + - ᶜ∂kappa_m∂q_totʲ * ( - Y.c.sgsʲs.:(1).mse - ᶜΦ + - cp_d * T_0 + - ∂e_int_∂q_tot * Y.c.sgsʲs.:(1).q_tot - ) - ), - ) - @. ᶠbidiagonal_matrix_ct3_2 = - DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ - DiagonalMatrixRow( - Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp) * - ∂e_int_∂q_tot, - ) - - @. ∂ᶜρaʲ_err_∂ᶜq_totʲ = - dtγ * ᶜadvdivᵥ_matrix() ⋅ - (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) - ∂ᶜρaʲ_err_∂ᶜmseʲ = - matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).mse)] - @. ᶠbidiagonal_matrix_ct3 = - DiagonalMatrixRow( - ᶠset_upwind_bcs( - ᶠupwind( - ᶠu³ʲs.:(1), - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ), - ) / ᶠJ, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( - ᶜJ * ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - @. ᶠbidiagonal_matrix_ct3_2 = - DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ - DiagonalMatrixRow( - Y.c.sgsʲs.:(1).ρa * ᶜkappa_mʲ / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - @. ∂ᶜρaʲ_err_∂ᶜmseʲ = - dtγ * ᶜadvdivᵥ_matrix() ⋅ - (ᶠbidiagonal_matrix_ct3 - ᶠbidiagonal_matrix_ct3_2) - ∂ᶜρaʲ_err_∂ᶜρaʲ = - matrix[@name(c.sgsʲs.:(1).ρa), @name(c.sgsʲs.:(1).ρa)] - @. ᶜadvection_matrix = - -(ᶜadvdivᵥ_matrix()) ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ) - @. ∂ᶜρaʲ_err_∂ᶜρaʲ = - dtγ * ᶜadvection_matrix ⋅ - ᶠset_upwind_matrix_bcs(ᶠupwind_matrix(ᶠu³ʲs.:(1))) ⋅ - DiagonalMatrixRow(1 / ᶜρʲs.:(1)) - (I,) - - - ∂ᶜρaʲ_err_∂ᶠu₃ʲ = - matrix[@name(c.sgsʲs.:(1).ρa), @name(f.sgsʲs.:(1).u₃)] - @. ∂ᶜρaʲ_err_∂ᶠu₃ʲ = - dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - ᶠinterp(ᶜρʲs.:(1) * ᶜJ) / ᶠJ * - ᶠset_upwind_bcs( - ᶠupwind( - CT3(sign(ᶠu³ʲ_data)), - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ), - ) * - adjoint(C3(sign(ᶠu³ʲ_data))) * - g³³(ᶠgⁱʲ), - ) - - ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).q_tot)] - @. ∂ᶠu₃ʲ_err_∂ᶜq_totʲ = - dtγ * DiagonalMatrixRow( - (1 - α_b) * ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / - (ᶠinterp(ᶜρʲs.:(1)))^2, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( - (ᶜρʲs.:(1))^2 / ᶜp * ( - ᶜkappa_mʲ / (ᶜkappa_mʲ + 1) * ∂e_int_∂q_tot + - ᶜ∂kappa_m∂q_totʲ * ( - Y.c.sgsʲs.:(1).mse - ᶜΦ + - cp_d * T_0 + - ∂e_int_∂q_tot * Y.c.sgsʲs.:(1).q_tot - ) - ), - ) - ∂ᶠu₃ʲ_err_∂ᶜmseʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(c.sgsʲs.:(1).mse)] - @. ∂ᶠu₃ʲ_err_∂ᶜmseʲ = - dtγ * DiagonalMatrixRow( - (1 - α_b) * ᶠgradᵥ_ᶜΦ * ᶠinterp(Y.c.ρ) / - (ᶠinterp(ᶜρʲs.:(1)))^2, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow( - ᶜkappa_mʲ * (ᶜρʲs.:(1))^2 / ((ᶜkappa_mʲ + 1) * ᶜp), - ) - ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] - ᶜu₃ʲ = ᶜtemp_C3 - @. ᶜu₃ʲ = ᶜinterp(Y.f.sgsʲs.:(1).u₃) - - @. bdmr_l = convert(BidiagonalMatrixRow{FT}, ᶜleft_bias_matrix()) - @. bdmr_r = convert(BidiagonalMatrixRow{FT}, ᶜright_bias_matrix()) - @. bdmr = ifelse(ᶜu₃ʲ.components.data.:1 > 0, bdmr_l, bdmr_r) - - @. ᶠtridiagonal_matrix_c3 = -(ᶠgradᵥ_matrix()) ⋅ bdmr - if rs isa RayleighSponge - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - dtγ * ( - ᶠtridiagonal_matrix_c3 ⋅ - DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - - DiagonalMatrixRow( - β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,), - ) - ) - (I_u₃,) - else - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - dtγ * ᶠtridiagonal_matrix_c3 ⋅ - DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - (I_u₃,) - end - - # entrainment and detrainment - if use_derivative(sgs_entr_detr_flag) - (; ᶜentrʲs, ᶜdetrʲs, ᶜturb_entrʲs) = p - # This assumes entrainment and detrainment rates are constant in the Jacobian - @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ -= - dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1)) - @. ∂ᶜmseʲ_err_∂ᶜmseʲ -= - dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1)) - @. ∂ᶜρaʲ_err_∂ᶜρaʲ += - dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) - ᶜdetrʲs.:(1)) - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ -= - dtγ * (DiagonalMatrixRow( - (ᶠinterp(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1))) * - (one_C3xACT3,), - )) - end - - # non-hydrostatic pressure drag - # Only the quadratic drag term is considered in the Jacobian, the buoyancy term is ignored - if use_derivative(sgs_nh_pressure_flag) - (; ᶠu₃⁰) = p - turbconv_params = CAP.turbconv_params(params) - α_d = CAP.pressure_normalmode_drag_coeff(turbconv_params) - scale_height = - CAP.R_d(params) * CAP.T_surf_ref(params) / CAP.grav(params) - H_up_min = CAP.min_updraft_top(turbconv_params) - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ -= - dtγ * (DiagonalMatrixRow( - 2 * - α_d * - Geometry._norm((Y.f.sgsʲs.:(1).u₃ - ᶠu₃⁰), ᶠlg) / - max(scale_height, H_up_min) * (one_C3xACT3,), - )) - end - - # add updraft mass flux contributions to grid-mean - if use_derivative(sgs_mass_flux_flag) - - (; ᶜgradᵥ_ᶠΦ, ᶜρʲs, ᶠu³ʲs, ᶜtsʲs, ᶠu³, ᶜKʲs) = p - (; bdmr_l, bdmr_r, bdmr) = p - is_third_order = edmfx_upwinding == Val(:third_order) - ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1 - ᶠset_upwind_bcs = Operators.SetBoundaryOperator(; - top = Operators.SetValue(zero(CT3{FT})), - bottom = Operators.SetValue(zero(CT3{FT})), - ) # Need to wrap ᶠupwind in this for well-defined boundaries. - UpwindMatrixRowType = - is_third_order ? QuaddiagonalMatrixRow : BidiagonalMatrixRow - ᶠupwind_matrix = - is_third_order ? ᶠupwind3_matrix : ᶠupwind1_matrix - - # Jacobian contributions of updraft massflux to grid-mean - - ∂ᶜupdraft_mass_flux_∂ᶜscalar = ᶠbidiagonal_matrix_ct3 - @. ∂ᶜupdraft_mass_flux_∂ᶜscalar = - DiagonalMatrixRow( - (ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) * (ᶠu³ʲs.:(1) - ᶠu³), - ) ⋅ ᶠinterp_matrix() ⋅ - DiagonalMatrixRow(Y.c.sgsʲs.:(1).ρa / ᶜρʲs.:(1)) - - # Derivative of total energy tendency with respect to updraft MSE - ## grid-mean ρe_tot - ᶜkappa_m = p.ᶜtemp_scalar - @. ᶜkappa_m = - TD.gas_constant_air(thermo_params, ᶜts) / - TD.cv_m(thermo_params, ᶜts) - - ᶜ∂kappa_m∂q_tot = p.ᶜtemp_scalar_2 - @. ᶜ∂kappa_m∂q_tot = - ( - ΔR_v * TD.cv_m(thermo_params, ᶜts) - - Δcv_v * TD.gas_constant_air(thermo_params, ᶜts) - ) / abs2(TD.cv_m(thermo_params, ᶜts)) - - @. ∂ᶜρe_tot_err_∂ᶜρ += - dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow( - ( - -(1 + ᶜkappa_m) * ᶜspecific.e_tot - - ᶜkappa_m * ∂e_int_∂q_tot * ᶜspecific.q_tot - ) / ᶜρ, - ) - - @. ∂ᶜρe_tot_err_∂ᶜρq_tot += - dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(( - ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + ᶜspecific.e_tot - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * ᶜspecific.q_tot - ) - )) - - @. ∂ᶜρe_tot_err_∂ᶜρe_tot += - dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) - - ∂ᶜρe_tot_err_∂ᶜmseʲ = - matrix[@name(c.ρe_tot), @name(c.sgsʲs.:(1).mse)] - @. ∂ᶜρe_tot_err_∂ᶜmseʲ = - -(dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar) - - ## grid-mean ρq_tot - @. ∂ᶜρq_tot_err_∂ᶜρ += - dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(-(ᶜspecific.q_tot) / ᶜρ) - - @. ∂ᶜρq_tot_err_∂ᶜρq_tot += - dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(1 / ᶜρ) - - ∂ᶜρq_tot_err_∂ᶜq_totʲ = - matrix[@name(c.ρq_tot), @name(c.sgsʲs.:(1).q_tot)] - @. ∂ᶜρq_tot_err_∂ᶜq_totʲ = - -(dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar) - - # grid-mean ∂/∂(u₃ʲ) - ∂ᶜρe_tot_err_∂ᶠu₃ = matrix[@name(c.ρe_tot), @name(f.u₃)] - @. ∂ᶜρe_tot_err_∂ᶠu₃ += - dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( - ᶠinterp( - (Y.c.sgsʲs.:(1).mse + ᶜKʲs.:(1) - ᶜh_tot) * - ᶜρʲs.:(1) * - ᶜJ * - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ) / ᶠJ * (g³³(ᶠgⁱʲ)), - ) - - ∂ᶜρe_tot_err_∂ᶠu₃ʲ = - matrix[@name(c.ρe_tot), @name(f.sgsʲs.:(1).u₃)] - @. ∂ᶜρe_tot_err_∂ᶠu₃ʲ = - dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - ᶠinterp( - (Y.c.sgsʲs.:(1).mse + ᶜKʲs.:(1) - ᶜh_tot) * - ᶜρʲs.:(1) * - ᶜJ * - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ) / ᶠJ * (g³³(ᶠgⁱʲ)), - ) - - ∂ᶜρq_tot_err_∂ᶠu₃ = matrix[@name(c.ρq_tot), @name(f.u₃)] - @. ∂ᶜρq_tot_err_∂ᶠu₃ += - dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( - ᶠinterp( - (Y.c.sgsʲs.:(1).q_tot - ᶜspecific.q_tot) * - ᶜρʲs.:(1) * - ᶜJ * - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ) / ᶠJ * (g³³(ᶠgⁱʲ)), - ) - - ∂ᶜρq_tot_err_∂ᶠu₃ʲ = - matrix[@name(c.ρq_tot), @name(f.sgsʲs.:(1).u₃)] - @. ∂ᶜρq_tot_err_∂ᶠu₃ʲ = - dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - ᶠinterp( - (Y.c.sgsʲs.:(1).q_tot - ᶜspecific.q_tot) * - ᶜρʲs.:(1) * - ᶜJ * - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), - ) / ᶠJ * (g³³(ᶠgⁱʲ)), - ) - - # grid-mean ∂/∂(rho*a) - ∂ᶜρe_tot_err_∂ᶜρa = - matrix[@name(c.ρe_tot), @name(c.sgsʲs.:(1).ρa)] - @. ∂ᶜρe_tot_err_∂ᶜρa = - dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - (ᶠu³ʲs.:(1) - ᶠu³) * - ᶠinterp((Y.c.sgsʲs.:(1).mse + ᶜKʲs.:(1) - ᶜh_tot)) / ᶠJ, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ) - - ∂ᶜρq_tot_err_∂ᶜρa = - matrix[@name(c.ρq_tot), @name(c.sgsʲs.:(1).ρa)] - @. ∂ᶜρq_tot_err_∂ᶜρa = - dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - (ᶠu³ʲs.:(1) - ᶠu³) * - ᶠinterp((Y.c.sgsʲs.:(1).q_tot - ᶜspecific.q_tot)) / ᶠJ, - ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ) - - end - - elseif rs isa RayleighSponge - ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - matrix[@name(f.sgsʲs.:(1).u₃), @name(f.sgsʲs.:(1).u₃)] - @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ = - dtγ * - -DiagonalMatrixRow( - β_rayleigh_w(rs, ᶠz, zmax) * (one_C3xACT3,), - ) - (I_u₃,) - end - end - - # NOTE: All velocity tendency derivatives should be set BEFORE this call. - zero_velocity_jacobian!(matrix, Y, p, t) end diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index e45de3e9cf..05cb2180cb 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -32,8 +32,8 @@ function vertical_diffusion_boundary_layer_tendency!( end ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C( - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(C3(FT(0))), + top = Operators.SetValue(C3(0)), + bottom = Operators.SetValue(C3(0)), ) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠgradᵥ(ᶜh_tot))) @@ -48,8 +48,8 @@ function vertical_diffusion_boundary_layer_tendency!( @. ᶜK_h_scaled = ᶜK_h end ᶜdivᵥ_ρχ = Operators.DivergenceF2C( - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(C3(FT(0))), + top = Operators.SetValue(C3(0)), + bottom = Operators.SetValue(C3(0)), ) @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρχ(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h_scaled) * ᶠgradᵥ(ᶜχ))) diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 91c0a72341..0bf4229c45 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -6,7 +6,6 @@ import ClimaUtilities.OutputPathGenerator import ClimaCore: InputOutput, Meshes, Spaces, Quadratures import ClimaAtmos.RRTMGPInterface as RRTMGPI import ClimaAtmos as CA -import LinearAlgebra import ClimaCore.Fields import ClimaTimeSteppers as CTS import Logging @@ -415,74 +414,52 @@ function get_surface_setup(parsed_args) return getproperty(SurfaceConditions, Symbol(parsed_args["surface_setup"]))() end -is_explicit_CTS_algo_type(alg_or_tableau) = - alg_or_tableau <: CTS.ERKAlgorithmName - -is_imex_CTS_algo_type(alg_or_tableau) = - alg_or_tableau <: CTS.IMEXARKAlgorithmName - -is_implicit_type(alg_or_tableau) = is_imex_CTS_algo_type(alg_or_tableau) - -is_imex_CTS_algo(::CTS.IMEXAlgorithm) = true -is_imex_CTS_algo(::CTS.RosenbrockAlgorithm) = true -is_imex_CTS_algo(::SciMLBase.AbstractODEAlgorithm) = false - -is_implicit(ode_algo) = is_imex_CTS_algo(ode_algo) - -is_rosenbrock(::SciMLBase.AbstractODEAlgorithm) = false -use_transform(ode_algo) = - !(is_imex_CTS_algo(ode_algo) || is_rosenbrock(ode_algo)) - -additional_integrator_kwargs(::SciMLBase.AbstractODEAlgorithm) = (; - adaptive = false, - progress = isinteractive(), - progress_steps = isinteractive() ? 1 : 1000, -) -additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (; - kwargshandle = CTS.DiffEqBase.KeywordArgSilent, # allow custom kwargs - adjustfinal = true, - # TODO: enable progress bars in ClimaTimeSteppers -) - -is_cts_algo(::SciMLBase.AbstractODEAlgorithm) = false -is_cts_algo(::CTS.DistributedODEAlgorithm) = true - -function jac_kwargs(ode_algo, Y, atmos, parsed_args) - if is_implicit(ode_algo) - A = ImplicitEquationJacobian( - Y, - atmos; - approximate_solve_iters = parsed_args["approximate_linear_solve_iters"], - transform_flag = use_transform(ode_algo), +get_jacobian(ode_algo, Y, atmos, parsed_args, output_dir) = + if ode_algo isa Union{CTS.IMEXAlgorithm, CTS.RosenbrockAlgorithm} + approx_jacobian_algorithm = ApproxJacobian( + DerivativeFlag(has_topography(axes(Y.c))), + DerivativeFlag(atmos.diff_mode), + DerivativeFlag(atmos.sgs_adv_mode), + DerivativeFlag(atmos.sgs_entr_detr_mode), + DerivativeFlag(atmos.sgs_mf_mode), + DerivativeFlag(atmos.sgs_nh_pressure_mode), + DerivativeFlag(atmos.noneq_cloud_formation_mode), + parsed_args["approximate_linear_solve_iters"], ) - if use_transform(ode_algo) - return (; jac_prototype = A, Wfact_t = Wfact!) + use_exact_jacobian = parsed_args["use_exact_jacobian"] + debug_approximate_jacobian = parsed_args["debug_approximate_jacobian"] + jacobian_algorithm = if debug_approximate_jacobian + # Debug one column when there are many columns. + only_debug_first_column_jacobian = + isnothing(parsed_args["only_debug_first_column_jacobian"]) ? + Fields.ncolumns(axes(Y.c)) > 10000 : + parsed_args["only_debug_first_column_jacobian"] + DebugJacobian( + approx_jacobian_algorithm, + use_exact_jacobian, + only_debug_first_column_jacobian, + ) + elseif use_exact_jacobian + ExactJacobian() else - return (; jac_prototype = A, Wfact = Wfact!) + approx_jacobian_algorithm end + @info "Jacobian algorithm: $(summary_string(jacobian_algorithm))" + ImplicitEquationJacobian(jacobian_algorithm, Y, atmos, output_dir) else - return NamedTuple() + nothing end -end -#= - ode_configuration(Y, parsed_args) - -Returns the ode algorithm -=# function ode_configuration(::Type{FT}, parsed_args) where {FT} ode_name = parsed_args["ode_algo"] - alg_or_tableau = getproperty(CTS, Symbol(ode_name)) - @info "Using ODE config: `$alg_or_tableau`" - if ode_name == "SSPKnoth" - return CTS.RosenbrockAlgorithm(CTS.tableau(CTS.SSPKnoth())) - end - - if is_explicit_CTS_algo_type(alg_or_tableau) - return CTS.ExplicitAlgorithm(alg_or_tableau()) - elseif !is_implicit_type(alg_or_tableau) - return alg_or_tableau() - elseif is_imex_CTS_algo_type(alg_or_tableau) + ode_algo_name = getproperty(CTS, Symbol(ode_name)) + @info "Using ODE config: `$ode_algo_name`" + return if ode_algo_name <: CTS.RosenbrockAlgorithmName + CTS.RosenbrockAlgorithm(CTS.tableau(ode_algo_name())) + elseif ode_algo_name <: CTS.ERKAlgorithmName + CTS.ExplicitAlgorithm(ode_algo_name()) + else + @assert ode_algo_name <: CTS.IMEXARKAlgorithmName newtons_method = CTS.NewtonsMethod(; max_iters = parsed_args["max_newton_iters_ode"], krylov_method = if parsed_args["use_krylov_method"] @@ -511,9 +488,7 @@ function ode_configuration(::Type{FT}, parsed_args) where {FT} nothing end, ) - return CTS.IMEXAlgorithm(alg_or_tableau(), newtons_method) - else - return alg_or_tableau(; linsolve = linsolve!) + CTS.IMEXAlgorithm(ode_algo_name(), newtons_method) end end @@ -656,40 +631,64 @@ function get_sim_info(config::AtmosConfig) return sim end -function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback) - (; atmos, dt) = p - - s = @timed_str begin - func = if parsed_args["split_ode"] - implicit_func = SciMLBase.ODEFunction( - implicit_tendency!; - jac_kwargs(ode_algo, Y, atmos, parsed_args)..., - tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0), - ) - if is_cts_algo(ode_algo) - CTS.ClimaODEFunction(; - T_exp_T_lim! = remaining_tendency!, - T_imp! = implicit_func, - # Can we just pass implicit_tendency! and jac_prototype etc.? - lim! = limiters_func!, - dss!, - cache! = set_precomputed_quantities!, - cache_imp! = set_implicit_precomputed_quantities!, - ) - else - SciMLBase.SplitFunction(implicit_func, remaining_tendency!) - end +function args_integrator(parsed_args, Y, p, sim_info, ode_algo, callback) + (; dt, atmos) = p + (; t_start, t_end, output_dir) = sim_info + use_exact_jacobian = parsed_args["use_exact_jacobian"] + debug_approximate_jacobian = parsed_args["debug_approximate_jacobian"] + need_to_update_exact_jacobian_before_each_solve = + ( + use_exact_jacobian && + isnothing(parsed_args["dt_update_exact_jacobian"]) + ) || ( + (debug_approximate_jacobian || use_exact_jacobian) && + !isnothing(parsed_args["dt_update_exact_jacobian"]) && + time_to_seconds(parsed_args["dt_update_exact_jacobian"]) < + time_to_seconds(parsed_args["dt"]) + ) + update_jacobian_before_each_solve! = + if need_to_update_exact_jacobian_before_each_solve + debug_approximate_jacobian ? update_and_check_jacobian! : + update_jacobian! + elseif !use_exact_jacobian + update_jacobian! else - remaining_tendency! # should be total_tendency! + Returns(nothing) # Update the exact Jacobian in a separate callback. end + s = @timed_str begin + T_imp! = SciMLBase.ODEFunction( + implicit_tendency!; + jac_prototype = get_jacobian( + ode_algo, + Y, + atmos, + parsed_args, + output_dir, + ), + Wfact = update_jacobian_before_each_solve!, + ) + tendency_function = CTS.ClimaODEFunction(; + T_exp_T_lim! = remaining_tendency!, + T_imp!, + lim! = limiters_func!, + dss!, + cache! = set_precomputed_quantities!, + cache_imp! = set_implicit_precomputed_quantities!, + ) end @info "Define ode function: $s" - problem = SciMLBase.ODEProblem(func, Y, tspan, p) - t_begin, t_end, _ = promote(tspan[1], tspan[2], p.dt) + problem = SciMLBase.ODEProblem(tendency_function, Y, (t_start, t_end), p) + t_start, t_end, _ = promote(t_start, t_end, dt) # Save solution to integrator.sol at the beginning and end - saveat = [t_begin, t_end] + saveat = [t_start, t_end] args = (problem, ode_algo) - kwargs = (; saveat, callback, dt, additional_integrator_kwargs(ode_algo)...) + kwargs = (; + saveat, + callback, + dt, + kwargshandle = CTS.DiffEqBase.KeywordArgSilent, # allow custom kwargs + adjustfinal = true, + ) return (args, kwargs) end @@ -856,13 +855,12 @@ function get_simulation(config::AtmosConfig) @info "n_steps_per_cycle_per_cb (non diagnostics): $steps_cycle_non_diag" @info "n_steps_per_cycle (non diagnostics): $steps_cycle" - tspan = (sim_info.t_start, sim_info.t_end) s = @timed_str begin integrator_args, integrator_kwargs = args_integrator( config.parsed_args, Y, p, - tspan, + sim_info, ode_algo, all_callbacks, ) diff --git a/src/solver/types.jl b/src/solver/types.jl index da459220da..1ea966e593 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -344,6 +344,7 @@ struct MixingLength{FT} tke::FT buoy::FT end +MixingLength(args...) = MixingLength(promote(args...)...) abstract type AbstractEDMF end diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl index d89cfc8797..424b83dae6 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -4,6 +4,7 @@ import ClimaComms import ClimaCore: Spaces, Topologies, Fields, Geometry import LinearAlgebra: norm_sqr +import Statistics: quantile is_energy_var(symbol) = symbol in (:ρe_tot, :ρae_tot) is_momentum_var(symbol) = symbol in (:uₕ, :ρuₕ, :u₃, :ρw) @@ -292,28 +293,39 @@ end import Dates """ - time_and_units_str(x::Real) - -Returns a truncated string of time and units, -given a time `x` in Seconds. -""" -time_and_units_str(x::Real) = - trunc_time(string(compound_period(x, Dates.Second))) + time_and_units_str(seconds) + +Converts the given number of seconds into a string that contains a numerical +value (rounded to the nearest 2 decimal digits) and its corresponding units. +""" +function time_and_units_str(seconds) + # Convert the time to a Float64 in order to avoid round-off errors. + full_period = compound_period(Float64(seconds)) + isempty(Dates.periods(full_period)) && return "0 Seconds" + whole_period = Dates.periods(full_period)[1] + whole_period_value = Dates.value(whole_period) + remaining_period = full_period - whole_period + no_remaining_digits = isempty(Dates.periods(remaining_period)) + period_units = + string(typeof(whole_period).name.name) * + (whole_period_value == 1 && no_remaining_digits ? "" : "s") + no_remaining_digits && return "$whole_period_value $period_units" + remaining_period_value = + Dates.value(convert(Dates.Nanosecond, remaining_period)) / + Dates.value(convert(Dates.Nanosecond, typeof(whole_period)(1))) + remaining_period_value < 0.01 && return "$whole_period_value $period_units" + remaining_digits = lpad(round(Int, remaining_period_value * 100), 2, '0') + return "$whole_period_value.$remaining_digits $period_units" +end """ - compound_period(x::Real, ::Type{T}) where {T <: Dates.Period} + compound_period(seconds) -A canonicalized `Dates.CompoundPeriod` given a real value -`x`, and its units via the period type `T`. +A `Dates.CompoundPeriod` that represents the given number of seconds (rounded to +the nearest nanosecond). """ -function compound_period(x::Real, ::Type{T}) where {T <: Dates.Period} - nf = Dates.value(convert(Dates.Nanosecond, T(1))) - ns = Dates.Nanosecond(ceil(x * nf)) - return Dates.canonicalize(Dates.CompoundPeriod(ns)) -end - -trunc_time(s::String) = count(',', s) > 1 ? join(split(s, ",")[1:2], ",") : s - +compound_period(seconds) = + Dates.canonicalize(Dates.Nanosecond(round(Int, seconds * 10^9))) function prettymemory(b) if b < 1024 @@ -341,6 +353,24 @@ macro timed_str(ex) end end +""" + summary_string(x) + +Returns a string that is similar to the output of `dump(x)`, but without any +type parameters. +""" +summary_string(x) = summary_string(x, 0) +summary_string(x, depth) = + fieldcount(typeof(x)) == 0 ? repr(x) : + (string(nameof(typeof(x))) * '(') * + mapreduce(*, 1:fieldcount(typeof(x))) do i + field = + x isa Tuple ? ':' * string(i) : string(fieldname(typeof(x), i)) + ('\n' * " "^(depth + 1) * field * " = ") * + (summary_string(getfield(x, i), depth + 1) * ',') + end * + ('\n' * " "^depth * ')') + struct AllNothing end const all_nothing = AllNothing() Base.getproperty(::AllNothing, ::Symbol) = nothing @@ -361,6 +391,23 @@ function horizontal_integral_at_boundary(f::Fields.Field) sum(f ./ Fields.Δz_field(axes(f)) .* 2) # TODO: is there a way to ensure this is derived from face z? The 2d topology doesn't contain this info end +""" + smooth_maximum(array, [percentile_rank]) + +Filters out all non-negligible values in the given `array` and computes the +specified percentile of the result. + +Specifically, this assumes that only the top 9 orders of magnitude of nonzero +data are significant, and, by default, it roughly selects the 90th percentile +along each dimension of the filtered data. +""" +function smooth_maximum(array, percentile_rank = 1 - 0.1^ndims(array)) + minimum_filtered_value = maximum(abs, array) / 1e9 + filtered_array_values = filter(>(minimum_filtered_value), map(abs, array)) + isempty(filtered_array_values) && return zero(eltype(array)) + return quantile(filtered_array_values, percentile_rank) +end + """ isdivisible(dt_large::Dates.Period, dt_small::Dates.Period) diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 485b21599c..f528bfb9a0 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -62,9 +62,9 @@ function as (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2`. """ sgs_weight_function(a, a_half) = - if a < 0 + if a <= 0 # autodiff generates NaNs when a is 0 zero(a) - elseif a > 1 + elseif a > min(1, 42 * a_half) # autodiff generates NaNs when a is large one(a) else (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2 diff --git a/test/dependencies.jl b/test/dependencies.jl index 69f1f81186..2aca57482a 100644 --- a/test/dependencies.jl +++ b/test/dependencies.jl @@ -44,6 +44,7 @@ known_dependencies = Set([ "ClimaUtilities", "CloudMicrophysics", "Dates", + "ForwardDiff", # for automatic differentiation in the implicit solver "Insolation", "Interpolations", "LazyBroadcast", # for https://github.com/CliMA/ClimaAtmos.jl/issues/3594 diff --git a/toml/aquaplanet_nonequil_120_1200.toml b/toml/aquaplanet_nonequil_120_1200.toml new file mode 100644 index 0000000000..f2eca4e552 --- /dev/null +++ b/toml/aquaplanet_nonequil_120_1200.toml @@ -0,0 +1,23 @@ +[condensation_evaporation_timescale] +value = 120 + +[sublimation_deposition_timescale] +value = 1200 + +[tracer_vertical_diffusion_factor] +value = 0 + +[snow_autoconversion_timescale] +value = 1e4 + +[rain_autoconversion_timescale] +value = 150 + +[zd_rayleigh] +value = 40000.0 + +[zd_viscous] +value = 40000.0 + +[precipitation_timescale] +value = 120 \ No newline at end of file diff --git a/toml/aquaplanet_nonequil_12_120.toml b/toml/aquaplanet_nonequil_12_120.toml new file mode 100644 index 0000000000..736858ec5e --- /dev/null +++ b/toml/aquaplanet_nonequil_12_120.toml @@ -0,0 +1,23 @@ +[condensation_evaporation_timescale] +value = 12 + +[sublimation_deposition_timescale] +value = 120 + +[tracer_vertical_diffusion_factor] +value = 0 + +[snow_autoconversion_timescale] +value = 1e4 + +[rain_autoconversion_timescale] +value = 150 + +[zd_rayleigh] +value = 40000.0 + +[zd_viscous] +value = 40000.0 + +[precipitation_timescale] +value = 120 \ No newline at end of file