diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 27f552d174..572cc75eee 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -36,18 +36,18 @@ steps: - "julia --project=experiments/ClimaEarth/ -e 'using Pkg; Pkg.precompile()'" - "julia --project=experiments/ClimaEarth/ -e 'using Pkg; Pkg.status()'" - - echo "--- Instantiate ClimaCore experiments env" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.add(\"MPI\")'" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.precompile()'" - - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.status()'" - - - echo "--- Instantiate test env" - - "julia --project=test/ -e 'using Pkg; Pkg.develop(path=\".\")'" - - "julia --project=test/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - - "julia --project=test/ -e 'using Pkg; Pkg.add(\"MPI\")'" - - "julia --project=test/ -e 'using Pkg; Pkg.precompile()'" - - "julia --project=test/ -e 'using Pkg; Pkg.status()'" + # - echo "--- Instantiate ClimaCore experiments env" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.add(\"MPI\")'" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.precompile()'" + # - "julia --project=experiments/ClimaCore/ -e 'using Pkg; Pkg.status()'" + + # - echo "--- Instantiate test env" + # - "julia --project=test/ -e 'using Pkg; Pkg.develop(path=\".\")'" + # - "julia --project=test/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + # - "julia --project=test/ -e 'using Pkg; Pkg.add(\"MPI\")'" + # - "julia --project=test/ -e 'using Pkg; Pkg.precompile()'" + # - "julia --project=test/ -e 'using Pkg; Pkg.status()'" concurrency: 1 concurrency_group: 'depot/climacoupler-ci' @@ -60,288 +60,288 @@ steps: - wait - - group: "Unit Tests" - steps: - - - label: "MPI Utilities unit tests" - key: "utilities_mpi_tests" - command: "srun julia --color=yes --project=test/ test/utilities_tests.jl" - timeout_in_minutes: 5 - env: - CLIMACOMMS_CONTEXT: "MPI" - agents: - slurm_ntasks: 2 - slurm_mem: 16GB - - - label: "MPI Interfacer unit tests" - key: "interfacer_mpi_tests" - command: "srun julia --color=yes --project=test/ test/interfacer_tests.jl" - timeout_in_minutes: 5 - env: - CLIMACOMMS_CONTEXT: "MPI" - agents: - slurm_ntasks: 2 - slurm_mem: 16GB - - - group: "GPU: unit tests" - steps: - - label: "GPU runtests" - command: "julia --color=yes --project=test/ test/runtests.jl" - timeout_in_minutes: 10 - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_ntasks: 1 - slurm_gres: "gpu:1" - slurm_mem: 24GB - - - group: "ClimaEarth tests" - steps: - - label: "ClimaEarth runtests" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/runtests.jl" - agents: - slurm_mem: 16GB - - - label: "MPI restarts" - key: "mpi_restarts" - command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" - env: - CLIMACOMMS_CONTEXT: "MPI" - timeout_in_minutes: 50 - soft_fail: - - exit_status: -1 - - exit_status: 255 - agents: - slurm_ntasks: 2 - slurm_mem: 32GB - - - label: "GPU restarts" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_ntasks: 1 - slurm_gres: "gpu:1" - slurm_mem: 32GB - - - label: "GPU restarts" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart_state_only.jl" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_ntasks: 1 - slurm_gres: "gpu:1" - slurm_mem: 32GB - - - group: "Integration Tests" - steps: - # SLABPLANET EXPERIMENTS - - # Slabplanet default: - # - this is the most lightweight example with conservation and visual checks, with CLI specification as follows - # - numerics: dt = dt_cpl = 200s, nelem = 4 - # - physics: bulk aerodynamic surface fluxes, gray radiation, idealized insolation, equil moisture model, 0-moment microphysics - # - input data: monotonous remapping (land mask, SST, SIC) - # - slurm: unthreaded, 1 ntask - # - diagnostics: check and plot energy conservation, output plots after 9 days - - label: "Slabplanet: default" - key: "slabplanet_default" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_default.yml --job_id slabplanet_default" - artifact_paths: "output/slabplanet_default/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet: dry, no radiation, fixed ocean T" - key: "slabplanet_dry_norad" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_dry_norad.yml --job_id slabplanet_dry_norad" - artifact_paths: "output/slabplanet_dry_norad/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet: extra atmos diagnostics" - key: "slabplanet_atmos_diags" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_atmos_diags.yml --job_id slabplanet_atmos_diags" - artifact_paths: "output/slabplanet_atmos_diags/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet terra: atmos and bucket" - key: "slabplanet_terra" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_terra.yml --job_id slabplanet_terra" - artifact_paths: "output/slabplanet_terra/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "Slabplanet aqua: atmos and slab ocean" - key: "slabplanet_aqua" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_aqua.yml --job_id slabplanet_aqua" - artifact_paths: "output/slabplanet_aqua/artifacts/*" - agents: - slurm_mem: 20GB - - # AMIP EXPERIMENTS - - # Test default behavior with no config file or job ID provided - - label: "AMIP: default" - key: "amip_default" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl" - artifact_paths: "output/amip_default/artifacts/*" - agents: - slurm_mem: 20GB - - - label: "AMIP: integrated land non-spun up initial condition test" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_land_ic.yml --job_id amip_land_ic" - artifact_paths: "output/amip_land_ic/artifacts/*" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "AMIP - Float64 + hourly checkpoint" - key: "amip" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_ft64_hourly_checkpoints.yml --job_id amip_coarse_ft64_hourly_checkpoints" - artifact_paths: "output/amip_coarse_ft64_hourly_checkpoints/artifacts/*" - env: - FLAME_PLOT: "" - BUILD_HISTORY_HANDLE: "" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "AMIP - Component dts test" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_component_dts.yml --job_id target_amip_component_dts" - artifact_paths: "output/target_amip_component_dts/artifacts/*" - agents: - slurm_ntasks: 1 - slurm_mem: 20GB - - - label: "MPI AMIP" - command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_mpi.yml --job_id amip_coarse_mpi" - artifact_paths: "output/amip_coarse_mpi/artifacts/*" - timeout_in_minutes: 30 - env: - CLIMACOMMS_CONTEXT: "MPI" - agents: - slurm_ntasks: 4 - slurm_mem_per_cpu: 12GB - - # short high-res performance test - - label: "Unthreaded AMIP FINE" # also reported by longruns with a flame graph - key: "unthreaded_amip_fine" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_n1_shortrun.yml --job_id target_amip_n1_shortrun" - artifact_paths: "output/target_amip_n1_shortrun/artifacts/*" - env: - BUILD_HISTORY_HANDLE: "" - agents: - slurm_mem: 20GB - - # CLIMACORE EXPERIMENTS - - - label: "sea_breeze" - command: "julia --color=yes --project=experiments/ClimaCore experiments/ClimaCore/sea_breeze/run.jl" - artifact_paths: "experiments/ClimaCore/sea_breeze/output/*" - agents: - slurm_mem: 20GB - - - label: "heat-diffusion" - command: "julia --color=yes --project=experiments/ClimaCore/ experiments/ClimaCore/heat-diffusion/run.jl" - artifact_paths: "experiments/ClimaCore/output/heat-diffusion/artifacts/*" - agents: - slurm_mem: 20GB - - - group: "GPU integration tests" - steps: - # GPU RUNS: slabplanet - - label: "GPU Slabplanet: albedo from function" - key: "gpu_slabplanet_albedo_function" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_albedo_function.yml --job_id gpu_slabplanet_albedo_function" - artifact_paths: "output/gpu_slabplanet_albedo_function/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - # GPU RUNS: AMIP - - label: "GPU AMIP: ED only + integrated land" - key: "gpu_amip_edonly_integrated_land" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_integrated_land.yml --job_id gpu_amip_edonly_integrated_land" - artifact_paths: "output/gpu_amip_edonly_integrated_land/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: ED only + bucket" - key: "gpu_amip_edonly_bucket" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_bucket.yml --job_id gpu_amip_edonly_bucket" - artifact_paths: "output/gpu_amip_edonly_bucket/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: diag. EDMF + integrated land" - key: "gpu_amip_diagedmf_integrated_land" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_integrated_land.yml --job_id gpu_amip_diagedmf_integrated_land" - artifact_paths: "output/gpu_amip_diagedmf_integrated_land/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: diag. EDMF + bucket" - key: "gpu_amip_diagedmf_bucket" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_bucket.yml --job_id gpu_amip_diagedmf_bucket" - artifact_paths: "output/gpu_amip_diagedmf_bucket/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: prog. EDMF + bucket" - key: "gpu_amip_progedmf_bucket" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_progedmf_bucket.yml --job_id gpu_amip_progedmf_bucket" - artifact_paths: "output/gpu_amip_progedmf_bucket/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - - label: "GPU AMIP test: albedo from function" - key: "gpu_amip_albedo_function" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_function.yml --job_id gpu_amip_albedo_function" - artifact_paths: "output/gpu_amip_albedo_function/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: albedo from temporal map + 0M" - key: "gpu_amip_albedo_temporal_map" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map.yml --job_id gpu_amip_albedo_temporal_map" - artifact_paths: "output/gpu_amip_albedo_temporal_map/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 - - - label: "GPU AMIP: albedo from temporal map + 1M" - key: "gpu_amip_albedo_temporal_map_1M" - command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map_1M.yml --job_id gpu_amip_albedo_temporal_map_1M" - artifact_paths: "output/gpu_amip_albedo_temporal_map_1M/artifacts/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - agents: - slurm_mem: 20GB - slurm_gpus: 1 + # - group: "Unit Tests" + # steps: + + # - label: "MPI Utilities unit tests" + # key: "utilities_mpi_tests" + # command: "srun julia --color=yes --project=test/ test/utilities_tests.jl" + # timeout_in_minutes: 5 + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # agents: + # slurm_ntasks: 2 + # slurm_mem: 16GB + + # - label: "MPI Interfacer unit tests" + # key: "interfacer_mpi_tests" + # command: "srun julia --color=yes --project=test/ test/interfacer_tests.jl" + # timeout_in_minutes: 5 + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # agents: + # slurm_ntasks: 2 + # slurm_mem: 16GB + + # - group: "GPU: unit tests" + # steps: + # - label: "GPU runtests" + # command: "julia --color=yes --project=test/ test/runtests.jl" + # timeout_in_minutes: 10 + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_ntasks: 1 + # slurm_gres: "gpu:1" + # slurm_mem: 24GB + + # - group: "ClimaEarth tests" + # steps: + # - label: "ClimaEarth runtests" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/runtests.jl" + # agents: + # slurm_mem: 16GB + + # - label: "MPI restarts" + # key: "mpi_restarts" + # command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # timeout_in_minutes: 50 + # soft_fail: + # - exit_status: -1 + # - exit_status: 255 + # agents: + # slurm_ntasks: 2 + # slurm_mem: 32GB + + # - label: "GPU restarts" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart.jl" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_ntasks: 1 + # slurm_gres: "gpu:1" + # slurm_mem: 32GB + + # - label: "GPU restarts" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/test/restart_state_only.jl" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_ntasks: 1 + # slurm_gres: "gpu:1" + # slurm_mem: 32GB + + # - group: "Integration Tests" + # steps: + # # SLABPLANET EXPERIMENTS + + # # Slabplanet default: + # # - this is the most lightweight example with conservation and visual checks, with CLI specification as follows + # # - numerics: dt = dt_cpl = 200s, nelem = 4 + # # - physics: bulk aerodynamic surface fluxes, gray radiation, idealized insolation, equil moisture model, 0-moment microphysics + # # - input data: monotonous remapping (land mask, SST, SIC) + # # - slurm: unthreaded, 1 ntask + # # - diagnostics: check and plot energy conservation, output plots after 9 days + # - label: "Slabplanet: default" + # key: "slabplanet_default" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_default.yml --job_id slabplanet_default" + # artifact_paths: "output/slabplanet_default/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet: dry, no radiation, fixed ocean T" + # key: "slabplanet_dry_norad" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_dry_norad.yml --job_id slabplanet_dry_norad" + # artifact_paths: "output/slabplanet_dry_norad/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet: extra atmos diagnostics" + # key: "slabplanet_atmos_diags" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_atmos_diags.yml --job_id slabplanet_atmos_diags" + # artifact_paths: "output/slabplanet_atmos_diags/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet terra: atmos and bucket" + # key: "slabplanet_terra" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_terra.yml --job_id slabplanet_terra" + # artifact_paths: "output/slabplanet_terra/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "Slabplanet aqua: atmos and slab ocean" + # key: "slabplanet_aqua" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_aqua.yml --job_id slabplanet_aqua" + # artifact_paths: "output/slabplanet_aqua/artifacts/*" + # agents: + # slurm_mem: 20GB + + # # AMIP EXPERIMENTS + + # # Test default behavior with no config file or job ID provided + # - label: "AMIP: default" + # key: "amip_default" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl" + # artifact_paths: "output/amip_default/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - label: "AMIP: integrated land non-spun up initial condition test" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_land_ic.yml --job_id amip_land_ic" + # artifact_paths: "output/amip_land_ic/artifacts/*" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "AMIP - Float64 + hourly checkpoint" + # key: "amip" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_ft64_hourly_checkpoints.yml --job_id amip_coarse_ft64_hourly_checkpoints" + # artifact_paths: "output/amip_coarse_ft64_hourly_checkpoints/artifacts/*" + # env: + # FLAME_PLOT: "" + # BUILD_HISTORY_HANDLE: "" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "AMIP - Component dts test" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_component_dts.yml --job_id target_amip_component_dts" + # artifact_paths: "output/target_amip_component_dts/artifacts/*" + # agents: + # slurm_ntasks: 1 + # slurm_mem: 20GB + + # - label: "MPI AMIP" + # command: "srun julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_coarse_mpi.yml --job_id amip_coarse_mpi" + # artifact_paths: "output/amip_coarse_mpi/artifacts/*" + # timeout_in_minutes: 30 + # env: + # CLIMACOMMS_CONTEXT: "MPI" + # agents: + # slurm_ntasks: 4 + # slurm_mem_per_cpu: 12GB + + # # short high-res performance test + # - label: "Unthreaded AMIP FINE" # also reported by longruns with a flame graph + # key: "unthreaded_amip_fine" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_n1_shortrun.yml --job_id target_amip_n1_shortrun" + # artifact_paths: "output/target_amip_n1_shortrun/artifacts/*" + # env: + # BUILD_HISTORY_HANDLE: "" + # agents: + # slurm_mem: 20GB + + # # CLIMACORE EXPERIMENTS + + # - label: "sea_breeze" + # command: "julia --color=yes --project=experiments/ClimaCore experiments/ClimaCore/sea_breeze/run.jl" + # artifact_paths: "experiments/ClimaCore/sea_breeze/output/*" + # agents: + # slurm_mem: 20GB + + # - label: "heat-diffusion" + # command: "julia --color=yes --project=experiments/ClimaCore/ experiments/ClimaCore/heat-diffusion/run.jl" + # artifact_paths: "experiments/ClimaCore/output/heat-diffusion/artifacts/*" + # agents: + # slurm_mem: 20GB + + # - group: "GPU integration tests" + # steps: + # # GPU RUNS: slabplanet + # - label: "GPU Slabplanet: albedo from function" + # key: "gpu_slabplanet_albedo_function" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/slabplanet_albedo_function.yml --job_id gpu_slabplanet_albedo_function" + # artifact_paths: "output/gpu_slabplanet_albedo_function/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # # GPU RUNS: AMIP + # - label: "GPU AMIP: ED only + integrated land" + # key: "gpu_amip_edonly_integrated_land" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_integrated_land.yml --job_id gpu_amip_edonly_integrated_land" + # artifact_paths: "output/gpu_amip_edonly_integrated_land/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: ED only + bucket" + # key: "gpu_amip_edonly_bucket" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_edonly_bucket.yml --job_id gpu_amip_edonly_bucket" + # artifact_paths: "output/gpu_amip_edonly_bucket/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: diag. EDMF + integrated land" + # key: "gpu_amip_diagedmf_integrated_land" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_integrated_land.yml --job_id gpu_amip_diagedmf_integrated_land" + # artifact_paths: "output/gpu_amip_diagedmf_integrated_land/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: diag. EDMF + bucket" + # key: "gpu_amip_diagedmf_bucket" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_diagedmf_bucket.yml --job_id gpu_amip_diagedmf_bucket" + # artifact_paths: "output/gpu_amip_diagedmf_bucket/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: prog. EDMF + bucket" + # key: "gpu_amip_progedmf_bucket" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_progedmf_bucket.yml --job_id gpu_amip_progedmf_bucket" + # artifact_paths: "output/gpu_amip_progedmf_bucket/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + + # - label: "GPU AMIP test: albedo from function" + # key: "gpu_amip_albedo_function" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_function.yml --job_id gpu_amip_albedo_function" + # artifact_paths: "output/gpu_amip_albedo_function/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: albedo from temporal map + 0M" + # key: "gpu_amip_albedo_temporal_map" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map.yml --job_id gpu_amip_albedo_temporal_map" + # artifact_paths: "output/gpu_amip_albedo_temporal_map/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 + + # - label: "GPU AMIP: albedo from temporal map + 1M" + # key: "gpu_amip_albedo_temporal_map_1M" + # command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_temporal_map_1M.yml --job_id gpu_amip_albedo_temporal_map_1M" + # artifact_paths: "output/gpu_amip_albedo_temporal_map_1M/artifacts/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # agents: + # slurm_mem: 20GB + # slurm_gpus: 1 - group: "CMIP" steps: @@ -366,28 +366,28 @@ steps: slurm_mem: 20GB slurm_gpus: 1 - - group: "Calibration experiments" - steps: - - label: "Perfect model calibration test" - key: "amip_pm_calibration" - command: - - "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/run_calibration.jl" - artifact_paths: "experiments/calibration/output/*" - env: - CLIMACOMMS_DEVICE: "CUDA" - CLIMACOMMS_CONTEXT: "SINGLETON" - SHORT_RUN: "" - agents: - slurm_mem: 64GB - slurm_ntasks: 3 - slurm_gpus_per_task: 1 - slurm_cpus_per_task: 4 - - - wait - - # plot job performance history - - label: ":chart_with_downwards_trend: build history" - command: - - build_history staging # name of branch to plot - artifact_paths: - - "build_history.html" + # - group: "Calibration experiments" + # steps: + # - label: "Perfect model calibration test" + # key: "amip_pm_calibration" + # command: + # - "julia --color=yes --project=experiments/ClimaEarth experiments/calibration/run_calibration.jl" + # artifact_paths: "experiments/calibration/output/*" + # env: + # CLIMACOMMS_DEVICE: "CUDA" + # CLIMACOMMS_CONTEXT: "SINGLETON" + # SHORT_RUN: "" + # agents: + # slurm_mem: 64GB + # slurm_ntasks: 3 + # slurm_gpus_per_task: 1 + # slurm_cpus_per_task: 4 + + # - wait + + # # plot job performance history + # - label: ":chart_with_downwards_trend: build history" + # command: + # - build_history staging # name of branch to plot + # artifact_paths: + # - "build_history.html" diff --git a/config/ci_configs/cmip_oceananigans_climaseaice.yml b/config/ci_configs/cmip_oceananigans_climaseaice.yml index abbc951299..fc628d0246 100644 --- a/config/ci_configs/cmip_oceananigans_climaseaice.yml +++ b/config/ci_configs/cmip_oceananigans_climaseaice.yml @@ -1,4 +1,4 @@ -FLOAT_TYPE: "Float32" +FLOAT_TYPE: "Float64" albedo_model: "CouplerAlbedo" atmos_config_file: "config/atmos_configs/climaatmos_edonly.yml" coupler_toml: ["toml/amip_edonly.toml"] @@ -9,16 +9,18 @@ dt_ocean: "1800secs" # 30 minutes dt_seaice: "1800secs" # 30 minutes dz_bottom: 100.0 energy_check: false -h_elem: 8 +h_elem: 32 ice_model: "clima_seaice" land_model: "integrated" land_spun_up_ic: false mode_name: "cmip" netcdf_output_at_levels: true +nh_poly: 1 ocean_model: "oceananigans" output_default_diagnostics: true radiation_reset_rng_seed: true rayleigh_sponge: true +simple_ocean: true start_date: "20100101" surface_setup: "PrescribedSurface" t_end: "1days" diff --git a/config/ci_configs/cmip_oceananigans_climaseaice_bucket.yml b/config/ci_configs/cmip_oceananigans_climaseaice_bucket.yml index d02a009ac8..4dfd02bbd7 100644 --- a/config/ci_configs/cmip_oceananigans_climaseaice_bucket.yml +++ b/config/ci_configs/cmip_oceananigans_climaseaice_bucket.yml @@ -1,4 +1,4 @@ -FLOAT_TYPE: "Float32" +FLOAT_TYPE: "Float64" albedo_model: "CouplerAlbedo" atmos_config_file: "config/atmos_configs/climaatmos_edonly.yml" bucket_albedo_type: "map_temporal" @@ -10,14 +10,16 @@ dt_ocean: "1800secs" # 30 minutes dt_seaice: "1800secs" # 30 minutes dz_bottom: 100.0 energy_check: false -h_elem: 8 +h_elem: 32 ice_model: "clima_seaice" mode_name: "cmip" netcdf_output_at_levels: true +nh_poly: 1 ocean_model: "oceananigans" output_default_diagnostics: true radiation_reset_rng_seed: true rayleigh_sponge: true +simple_ocean: true start_date: "20100101" surface_setup: "PrescribedSurface" t_end: "1days" diff --git a/config/longrun_configs/cmip_edonly_bucket.yml b/config/longrun_configs/cmip_edonly_bucket.yml index 018686137d..c2a3219eda 100644 --- a/config/longrun_configs/cmip_edonly_bucket.yml +++ b/config/longrun_configs/cmip_edonly_bucket.yml @@ -1,4 +1,4 @@ -FLOAT_TYPE: "Float32" +FLOAT_TYPE: "Float64" aerosol_radiation: true atmos_config_file: "config/atmos_configs/climaatmos_edonly.yml" bucket_albedo_type: "map_temporal" @@ -11,9 +11,11 @@ dt_ocean: "120secs" # 2 minutes dt_rad: "1hours" dt_seaice: "120secs" # 2 minutes energy_check: false +h_elem: 64 ice_model: "clima_seaice" insolation: "timevarying" mode_name: "cmip" +nh_poly: 1 ocean_model: "oceananigans" prescribed_aerosols: ["CB1", "CB2", "DST01", "OC1", "OC2", "SO4", "SSLT01"] rmse_check: true diff --git a/config/longrun_configs/cmip_edonly_land.yml b/config/longrun_configs/cmip_edonly_land.yml index ac1bb78130..4576e454d9 100644 --- a/config/longrun_configs/cmip_edonly_land.yml +++ b/config/longrun_configs/cmip_edonly_land.yml @@ -1,4 +1,4 @@ -FLOAT_TYPE: "Float32" +FLOAT_TYPE: "Float64" aerosol_radiation: true atmos_config_file: "config/atmos_configs/climaatmos_edonly.yml" checkpoint_dt: "720hours" @@ -10,11 +10,13 @@ dt_ocean: "120secs" # 2 minutes dt_rad: "1hours" dt_seaice: "120secs" # 2 minutes energy_check: false +h_elem: 64 ice_model: "clima_seaice" insolation: "timevarying" land_model: "integrated" land_spun_up_ic: false mode_name: "cmip" +nh_poly: 1 ocean_model: "oceananigans" prescribed_aerosols: ["CB1", "CB2", "DST01", "OC1", "OC2", "SO4", "SSLT01"] rmse_check: true diff --git a/experiments/ClimaEarth/Manifest-v1.11.toml b/experiments/ClimaEarth/Manifest-v1.11.toml index 9de8c5620c..5bb56dff3c 100644 --- a/experiments/ClimaEarth/Manifest-v1.11.toml +++ b/experiments/ClimaEarth/Manifest-v1.11.toml @@ -2,7 +2,7 @@ julia_version = "1.11.8" manifest_format = "2.0" -project_hash = "88cc89bf1807e712c8314e95f2147d5f7248896f" +project_hash = "63a52083e7157c640c34a856b0951a8330688e84" [[deps.ADTypes]] git-tree-sha1 = "f7304359109c768cf32dc5fa2d371565bb63b68a" @@ -434,6 +434,11 @@ git-tree-sha1 = "34d9873079e4cb3d0c62926a225136824677073f" uuid = "55437552-ac27-4d47-9aa3-63184e8fd398" version = "1.0.0" +[[deps.ChunkSplitters]] +git-tree-sha1 = "63a3903063d035260f0f6eab00f517471c5dc784" +uuid = "ae650224-84b6-46f8-82ea-d812ca08434e" +version = "3.1.2" + [[deps.ClimaAnalysis]] deps = ["Artifacts", "Dates", "Interpolations", "NCDatasets", "NaNStatistics", "OrderedCollections", "Reexport", "Statistics", "Unitful"] git-tree-sha1 = "6becd117214e883e0069455303c43060a4d5eed1" @@ -739,6 +744,30 @@ git-tree-sha1 = "d9d26935a0bcffc87d2613ce14c527c99fc543fd" uuid = "f0e56b4a-5159-44fe-b623-3e5288b988bb" version = "2.5.0" +[[deps.ConservativeRegridding]] +deps = ["ChunkSplitters", "DocStringExtensions", "Extents", "GeoInterface", "GeometryOps", "GeometryOpsCore", "LinearAlgebra", "ProgressMeter", "SciMLPublic", "SortTileRecursiveTree", "SparseArrays", "StableTasks", "StaticArrays"] +git-tree-sha1 = "4392da3295bd4bf5f6c43cb1ee5da678a76bfb35" +repo-rev = "js/climacore-fix" +repo-url = "https://github.com/JuliaGeo/ConservativeRegridding.jl.git" +uuid = "8e50ac2c-eb48-49bc-a402-07c87b949343" +version = "0.2.0" + + [deps.ConservativeRegridding.extensions] + ConservativeRegriddingClimaCoreExt = "ClimaCore" + ConservativeRegriddingHealpixExt = "Healpix" + ConservativeRegriddingInterfacesExt = "Interfaces" + ConservativeRegriddingOceananigansExt = "Oceananigans" + ConservativeRegriddingRingGridsExt = "RingGrids" + ConservativeRegriddingSpeedyWeatherExt = ["RingGrids", "SpeedyWeather"] + + [deps.ConservativeRegridding.weakdeps] + ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884" + Healpix = "9f4e344d-96bc-545a-84a3-ae6b9e1b672b" + Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87" + Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" + RingGrids = "d1845624-ad4f-453b-8ff4-a8db365bf3a7" + SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" + [[deps.ConstructionBase]] git-tree-sha1 = "b4b092499347b18a015186eae3042f72267106cb" uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" diff --git a/experiments/ClimaEarth/Manifest.toml b/experiments/ClimaEarth/Manifest.toml index d32cc8186b..ab6f5efbdc 100644 --- a/experiments/ClimaEarth/Manifest.toml +++ b/experiments/ClimaEarth/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.9" manifest_format = "2.0" -project_hash = "b6758675118500db19bc41873b8c8741b74a3997" +project_hash = "6b7bac012d444043544ee6ea533d5321afdc5bde" [[deps.ADTypes]] git-tree-sha1 = "f7304359109c768cf32dc5fa2d371565bb63b68a" @@ -431,6 +431,11 @@ git-tree-sha1 = "34d9873079e4cb3d0c62926a225136824677073f" uuid = "55437552-ac27-4d47-9aa3-63184e8fd398" version = "1.0.0" +[[deps.ChunkSplitters]] +git-tree-sha1 = "63a3903063d035260f0f6eab00f517471c5dc784" +uuid = "ae650224-84b6-46f8-82ea-d812ca08434e" +version = "3.1.2" + [[deps.ClimaAnalysis]] deps = ["Artifacts", "Dates", "Interpolations", "NCDatasets", "NaNStatistics", "OrderedCollections", "Reexport", "Statistics", "Unitful"] git-tree-sha1 = "6becd117214e883e0069455303c43060a4d5eed1" @@ -738,6 +743,30 @@ git-tree-sha1 = "d9d26935a0bcffc87d2613ce14c527c99fc543fd" uuid = "f0e56b4a-5159-44fe-b623-3e5288b988bb" version = "2.5.0" +[[deps.ConservativeRegridding]] +deps = ["ChunkSplitters", "DocStringExtensions", "Extents", "GeoInterface", "GeometryOps", "GeometryOpsCore", "LinearAlgebra", "ProgressMeter", "SciMLPublic", "SortTileRecursiveTree", "SparseArrays", "StableTasks", "StaticArrays"] +git-tree-sha1 = "14be129cf6c44210906eef3fd017fd9ba57acd43" +repo-rev = "main" +repo-url = "https://github.com/JuliaGeo/ConservativeRegridding.jl.git" +uuid = "8e50ac2c-eb48-49bc-a402-07c87b949343" +version = "0.2.0" + + [deps.ConservativeRegridding.extensions] + ConservativeRegriddingClimaCoreExt = "ClimaCore" + ConservativeRegriddingHealpixExt = "Healpix" + ConservativeRegriddingInterfacesExt = "Interfaces" + ConservativeRegriddingOceananigansExt = "Oceananigans" + ConservativeRegriddingRingGridsExt = "RingGrids" + ConservativeRegriddingSpeedyWeatherExt = ["RingGrids", "SpeedyWeather"] + + [deps.ConservativeRegridding.weakdeps] + ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884" + Healpix = "9f4e344d-96bc-545a-84a3-ae6b9e1b672b" + Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87" + Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" + RingGrids = "d1845624-ad4f-453b-8ff4-a8db365bf3a7" + SpeedyWeather = "9e226e20-d153-4fed-8a5b-493def4f21a9" + [[deps.ConstructionBase]] git-tree-sha1 = "b4b092499347b18a015186eae3042f72267106cb" uuid = "187b0558-2788-49d3-abe0-74a17ed4e7c9" diff --git a/experiments/ClimaEarth/Project.toml b/experiments/ClimaEarth/Project.toml index 8aa4b6b47c..5d8b7deb5f 100644 --- a/experiments/ClimaEarth/Project.toml +++ b/experiments/ClimaEarth/Project.toml @@ -1,4 +1,5 @@ [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" Cairo = "159f3aea-2a34-519c-b102-8c37f9878175" @@ -17,6 +18,7 @@ ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c" ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4" ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79" ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513" +ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343" EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d" GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55" GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6" @@ -45,7 +47,9 @@ ClimaParams = "1.0" ClimaSeaIce = "0.4" ClimaTimeSteppers = "0.7, 0.8" ClimaUtilities = "0.1" +ConservativeRegridding = "0.2" Insolation = "0.10.2" Interpolations = "0.14, 0.15" +JLD2 = "0.4, 0.5, 0.6" Poppler_jll = "24" StaticArrays = "1" diff --git a/experiments/ClimaEarth/components/ocean/clima_seaice.jl b/experiments/ClimaEarth/components/ocean/clima_seaice.jl index a4de1feece..34d638fe7b 100644 --- a/experiments/ClimaEarth/components/ocean/clima_seaice.jl +++ b/experiments/ClimaEarth/components/ocean/clima_seaice.jl @@ -207,7 +207,7 @@ function ClimaSeaIceSimulation( # Get the initial area fraction from the fractional ice concentration boundary_space = axes(ocean.area_fraction) - area_fraction = Interfacer.remap(boundary_space, ice.model.ice_concentration) + area_fraction = Interfacer.remap(boundary_space, ice.model.ice_concentration, remapping) sim = ClimaSeaIceSimulation( ice, @@ -332,23 +332,13 @@ function FluxCalculator.update_turbulent_fluxes!(sim::ClimaSeaIceSimulation, fie F_turb_ρτxz_uv = sim.remapping.temp_uv_vec.components.data.:1 F_turb_ρτyz_uv = sim.remapping.temp_uv_vec.components.data.:2 - # Remap momentum fluxes onto reduced 2D Center, Center fields using scratch arrays and fields - CC.Remapping.interpolate!( - sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, - F_turb_ρτxz_uv, - ) - OC.set!(sim.remapping.scratch_cc1, sim.remapping.scratch_arr1) # zonal momentum flux - CC.Remapping.interpolate!( - sim.remapping.scratch_arr2, - sim.remapping.remapper_cc, - F_turb_ρτyz_uv, - ) - OC.set!(sim.remapping.scratch_cc2, sim.remapping.scratch_arr2) # meridional momentum flux + # Remap momentum fluxes onto reduced 2D Center, Center fields + Interfacer.remap!(sim.remapping.scratch_field_oc1, F_turb_ρτxz, sim.remapping) # zonal momentum flux + Interfacer.remap!(sim.remapping.scratch_field_oc2, F_turb_ρτyz, sim.remapping) # meridional momentum flux # Rename for clarity; these are now Center, Center Oceananigans fields - F_turb_ρτxz_cc = sim.remapping.scratch_cc1 - F_turb_ρτyz_cc = sim.remapping.scratch_cc2 + F_turb_ρτxz_cc = sim.remapping.scratch_field_oc1 + F_turb_ρτyz_cc = sim.remapping.scratch_field_oc2 # Set the momentum flux BCs at the correct locations using the remapped scratch fields # Note that this requires the sea ice model to always be run with dynamics turned on @@ -362,12 +352,12 @@ function FluxCalculator.update_turbulent_fluxes!(sim::ClimaSeaIceSimulation, fie ) # Remap the latent and sensible heat fluxes using scratch arrays - CC.Remapping.interpolate!(sim.remapping.scratch_arr1, sim.remapping.remapper_cc, F_lh) # latent heat flux - CC.Remapping.interpolate!(sim.remapping.scratch_arr2, sim.remapping.remapper_cc, F_sh) # sensible heat flux + Interfacer.remap!(sim.remapping.scratch_field_oc1, F_lh, sim.remapping) # latent heat flux + Interfacer.remap!(sim.remapping.scratch_field_oc2, F_sh, sim.remapping) # sensible heat flux # Rename for clarity; recall F_turb_energy = F_lh + F_sh - remapped_F_lh = sim.remapping.scratch_arr1 - remapped_F_sh = sim.remapping.scratch_arr2 + remapped_F_lh = OC.interior(sim.remapping.scratch_field_oc1, :, :, 1) + remapped_F_sh = OC.interior(sim.remapping.scratch_field_oc2, :, :, 1) # Update the sea ice only where the concentration is greater than zero. si_flux_heat = sim.ice.model.external_heat_fluxes.top[1] @@ -405,19 +395,11 @@ function FieldExchanger.update_sim!(sim::ClimaSeaIceSimulation, csf) ice_concentration = sim.ice.model.ice_concentration # Remap radiative flux onto scratch array; rename for clarity - CC.Remapping.interpolate!( - sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, - csf.SW_d, - ) - remapped_SW_d = sim.remapping.scratch_arr1 + Interfacer.remap!(sim.remapping.scratch_field_oc1, csf.SW_d, sim.remapping) # shortwave radiation + remapped_SW_d = OC.interior(sim.remapping.scratch_field_oc1, :, :, 1) - CC.Remapping.interpolate!( - sim.remapping.scratch_arr2, - sim.remapping.remapper_cc, - csf.LW_d, - ) - remapped_LW_d = sim.remapping.scratch_arr2 + Interfacer.remap!(sim.remapping.scratch_field_oc2, csf.LW_d, sim.remapping) # longwave radiation + remapped_LW_d = OC.interior(sim.remapping.scratch_field_oc2, :, :, 1) # Update only the part due to radiative fluxes. For the full update, the component due # to latent and sensible heat is missing and will be updated in update_turbulent_fluxes. @@ -432,122 +414,6 @@ function FieldExchanger.update_sim!(sim::ClimaSeaIceSimulation, csf) return nothing end -""" - ocean_seaice_fluxes!(ocean_sim::OceananigansSimulation, ice_sim::ClimaSeaIceSimulation) - -Compute the fluxes between the ocean and sea ice, storing them in the `ocean_ice_fluxes` -fields of the ocean and sea ice simulations. - -This function assumes both simulations share the same grid, so no remapping is done. - -Both simulations have had their atmospheric fluxes updated already in this timestep -(see `update_sim!` and `update_turbulent_fluxes!`), so we add the contributions from the -ocean-sea ice interactions to the existing fluxes, rather than overwriting all fluxes. - -!!! note - This function must be called after the turbulent fluxes have been updated in both - simulations. Here only the contributions from the sea ice/ocean interactions - are added to the fluxes. -""" -function FluxCalculator.ocean_seaice_fluxes!( - ocean_sim::OceananigansSimulation, - ice_sim::ClimaSeaIceSimulation, -) - ocean_properties = ocean_sim.ocean_properties - ice_concentration = Interfacer.get_field(ice_sim, Val(:ice_concentration)) - - # Update the sea ice concentration in the ocean simulation - ocean_sim.ice_concentration .= ice_concentration - - # Compute the fluxes and store them in the both simulations - CO.OceanSeaIceModels.InterfaceComputations.compute_sea_ice_ocean_fluxes!( - ice_sim.ocean_ice_interface, - ocean_sim.ocean, - ice_sim.ice, - ocean_properties, - ) - - ## Update the internals of the sea ice model - # Set the bottom heat flux to the sum of the frazil and interface heat fluxes - bottom_heat_flux = ice_sim.ice.model.external_heat_fluxes.bottom - - Qf = ice_sim.ocean_ice_interface.fluxes.frazil_heat # frazil heat flux - Qi = ice_sim.ocean_ice_interface.fluxes.interface_heat # interfacial heat flux - bottom_heat_flux .= Qf .+ Qi - - ## Update the internals of the ocean model - ρₒ⁻¹ = 1 / ocean_sim.ocean_properties.reference_density - cₒ = ocean_sim.ocean_properties.heat_capacity - - # Compute fluxes for u, v, T, and S from momentum, heat, and freshwater fluxes - oc_flux_u = surface_flux(ocean_sim.ocean.model.velocities.u) - oc_flux_v = surface_flux(ocean_sim.ocean.model.velocities.v) - - ρτxio = ice_sim.ocean_ice_interface.fluxes.x_momentum # sea_ice - ocean zonal momentum flux - ρτyio = ice_sim.ocean_ice_interface.fluxes.y_momentum # sea_ice - ocean meridional momentum flux - - # Update the momentum flux contributions from ocean/sea ice fluxes - grid = ocean_sim.ocean.model.grid - arch = OC.Architectures.architecture(grid) - OC.Utils.launch!( - arch, - grid, - :xy, - _add_ocean_ice_stress!, - oc_flux_u, - oc_flux_v, - grid, - ρτxio, - ρτyio, - ρₒ⁻¹, - ice_concentration, - ) - - oc_flux_T = surface_flux(ocean_sim.ocean.model.tracers.T) - OC.interior(oc_flux_T, :, :, 1) .+= - OC.interior(ice_concentration, :, :, 1) .* OC.interior(Qi, :, :, 1) .* ρₒ⁻¹ ./ cₒ - - oc_flux_S = surface_flux(ocean_sim.ocean.model.tracers.S) - OC.interior(oc_flux_S, :, :, 1) .+= - OC.interior(ice_concentration, :, :, 1) .* - OC.interior(ice_sim.ocean_ice_interface.fluxes.salt, :, :, 1) - - return nothing -end - -""" - _add_ocean_ice_stress!(oc_flux_u, oc_flux_v, grid, ρτxio, ρτyio, ρₒ⁻¹, ice_concentration) - -Add the contribution from the ocean-ice stress to the surface fluxes for each -component of the ocean velocity (u and v). - -Arguments: -- `oc_flux_u`: [Field] the surface flux for the ocean zonal velocity. -- `oc_flux_v`: [Field] the surface flux for the ocean meridional velocity. -- `grid`: [Grid] the grid used by ocean and ice. -- `ρτxio`: [Field] the ice-ocean zonal momentum flux. -- `ρτyio`: [Field] the ice-ocean meridional momentum flux. -- `ρₒ⁻¹`: [Float] the inverse of the ocean reference density. -- `ice_concentration`: [Field] the sea ice concentration. -""" -@kernel function _add_ocean_ice_stress!( - oc_flux_u, - oc_flux_v, - grid, - ρτxio, - ρτyio, - ρₒ⁻¹, - ice_concentration, -) - i, j = @index(Global, NTuple) - - # ℑxᶠᵃᵃ: interpolate faces to centers - oc_flux_u[i, j, 1] += - ρτxio[i, j, 1] * ρₒ⁻¹ * OC.Operators.ℑxᶠᵃᵃ(i, j, 1, grid, ice_concentration) - oc_flux_v[i, j, 1] += - ρτyio[i, j, 1] * ρₒ⁻¹ * OC.Operators.ℑyᵃᶠᵃ(i, j, 1, grid, ice_concentration) -end - """ get_model_prog_state(sim::ClimaSeaIceSimulation) diff --git a/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl b/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl index dde6ec4472..bf93444ceb 100644 --- a/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl +++ b/experiments/ClimaEarth/components/ocean/climaocean_helpers.jl @@ -1,46 +1,3 @@ -""" - to_node(pt::CC.Geometry.LatLongPoint) - -Transform `LatLongPoint` into a tuple (long, lat, 0), where the 0 is needed because we only -care about the surface. -""" -@inline to_node(pt::CC.Geometry.LatLongPoint) = pt.long, pt.lat, zero(pt.lat) -# This next one is needed if we have "LevelGrid" -@inline to_node(pt::CC.Geometry.LatLongZPoint) = pt.long, pt.lat, zero(pt.lat) - -""" - map_interpolate!(target_field::CC.Fields.Field, source_field::OC.Field) - -Interpolate the given 3D field onto the target field, modifying the target field in place. - -If the underlying grid does not contain a given point, writes 0 instead. - -Note: `map_interpolate!` does not support interpolation from `Field`s defined on -`OrthogononalSphericalShellGrids` such as the `TripolarGrid`. -""" -function map_interpolate!(target_field::CC.Fields.Field, source_field::OC.Field) - points = CC.Fields.coordinate_field(axes(target_field)) - loc = map(L -> L(), OC.Fields.location(source_field)) - grid = source_field.grid - data = source_field.data - - # TODO: There has to be a better way - min_lat, max_lat = extrema(OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center())) - - # Use map! on the field directly, which handles complex data layouts correctly - map!(target_field, points) do pt - FT = eltype(pt) - - # The oceananigans grid does not cover the entire globe, so we should not - # interpolate outside of its latitude bounds. Instead we return 0 - min_lat < pt.lat < max_lat || return FT(0) - - fᵢ = OC.Fields.interpolate(to_node(pt), data, loc, grid) - return convert(FT, fᵢ)::FT - end - return nothing -end - """ surface_flux(f::OC.AbstractField) @@ -55,35 +12,6 @@ function surface_flux(f::OC.AbstractField) end end -function Interfacer.remap!(target_field::CC.Fields.Field, source_field::OC.Field) - map_interpolate!(target_field, source_field) - return nothing -end - -function Interfacer.remap!( - target_field::CC.Fields.Field, - operation::OC.AbstractOperations.AbstractOperation, -) - evaluated_field = OC.Field(operation) - OC.compute!(evaluated_field) - Interfacer.remap!(target_field, evaluated_field) - return nothing -end - -function Interfacer.remap(target_space, field::OC.Field) - # Allocate target field and call remap! - target_field = CC.Fields.zeros(target_space) - Interfacer.remap!(target_field, field) - return target_field -end - -function Interfacer.remap(target_space, operation::OC.AbstractOperations.AbstractOperation) - # Allocate target field and call remap! - target_field = CC.Fields.zeros(target_space) - Interfacer.remap!(target_field, operation) - return target_field -end - """ set_from_extrinsic_vector!(vector, grid, u_cc, v_cc) diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl index bf251a9c4f..ec3a748062 100644 --- a/experiments/ClimaEarth/components/ocean/oceananigans.jl +++ b/experiments/ClimaEarth/components/ocean/oceananigans.jl @@ -7,6 +7,8 @@ import Thermodynamics as TD import ClimaParams as CP import ClimaOcean.EN4: download_dataset using KernelAbstractions: @kernel, @index, @inbounds +import ConservativeRegridding as CR +import Adapt # for ConservativeRegridding include("climaocean_helpers.jl") @@ -55,7 +57,7 @@ dispatch in coupling. - `start_date`: Start date for the simulation - `stop_date`: Stop date for the simulation - `output_dir`: Directory for output files -- `ice_model`: Ice model type (Val type) +- `simple_ocean`: Whether to use a simpler ocean model setup (default: `false`) # Optional keyword arguments - `dt`: Time step (default: `nothing`) @@ -71,19 +73,18 @@ function OceananigansSimulation( start_date, tspan, output_dir, - ice_model, + simple_ocean = false, dt = nothing, comms_ctx = ClimaComms.context(), coupled_param_dict = CP.create_toml_dict(FT), extra_kwargs..., ) where {FT} arch = comms_ctx.device isa ClimaComms.CUDADevice ? OC.GPU() : OC.CPU() - OC.Oceananigans.defaults.FloatType = FT # Compute stop_date for oceananigans (needed for EN4 data retrieval) stop_date = start_date + Dates.Second(float(tspan[2] - tspan[1])) - # Use Float64 for the ocean to avoid precision issues + # Use Float64 for the ocean to avoid memory access issues since ClimaSeaIce hardcodes Float64 FT_ocean = Float64 OC.Oceananigans.defaults.FloatType = FT_ocean @@ -95,116 +96,81 @@ function OceananigansSimulation( download_dataset(en4_temperature) download_dataset(en4_salinity) - # Set up ocean grid (1 degree) - resolution_points = (360, 160, 32) - Nz = last(resolution_points) - depth = 4000 # meters - z = OC.ExponentialDiscretization(Nz, -depth, 0; scale = 0.85 * depth) + # Set up tripolar ocean grid (1 degree) + Nx = 720 + Ny = 360 + Nz = 100 + depth = 6000 # meters + z = OC.ExponentialDiscretization(Nz, -depth, 0; scale = 1800) # Regular LatLong because we know how to do interpolation there - underlying_grid = OC.LatitudeLongitudeGrid( - arch; - size = resolution_points, - longitude = (-180, 180), - latitude = (-80, 80), # NOTE: Don't goo to high up when using LatLongGrid, or the cells will be too small - z, - halo = (7, 7, 7), - ) + underlying_grid = OC.TripolarGrid(arch; size = (Nx, Ny, Nz), halo = (7, 7, 7), z) bottom_height = CO.regrid_bathymetry( underlying_grid; - minimum_depth = 30, - interpolation_passes = 20, + minimum_depth = 20, + interpolation_passes = 2, major_basins = 1, ) - grid = OC.ImmersedBoundaryGrid( underlying_grid, OC.GridFittedBottom(bottom_height); active_cells_map = true, ) - # Restore the ocean to the EN4 state periodically if running for more than one month and not using ClimaSeaIce - use_restoring = - start_date + Dates.Month(1) < stop_date && ice_model != Val(:clima_seaice) - - if use_restoring - # When we use EN4 data, the forcing takes care of everything, including - # the initial conditions - restoring_rate = 1 / (3 * 86400) - mask = CO.LinearlyTaperedPolarMask( - southern = (-80, -70), - northern = (70, 90), - z = (z(1), 0), + # Create ocean simulation + if !simple_ocean + free_surface = OC.SplitExplicitFreeSurface(grid; substeps = 150) + momentum_advection = OC.WENOVectorInvariant(order = 5) + tracer_advection = OC.WENO(order = 7) + eddy_closure = OC.TurbulenceClosures.IsopycnalSkewSymmetricDiffusivity( + κ_skew = 500, + κ_symmetric = 100, ) + @inline νhb(i, j, k, grid, ℓx, ℓy, ℓz, clock, fields, λ) = + OC.Operators.Az(i, j, k, grid, ℓx, ℓy, ℓz)^2 / λ - forcing_T = CO.DatasetRestoring(en4_temperature, grid; mask, rate = restoring_rate) - forcing_S = CO.DatasetRestoring(en4_salinity, grid; mask, rate = restoring_rate) - forcing = (T = forcing_T, S = forcing_S) + horizontal_viscosity = OC.HorizontalScalarBiharmonicDiffusivity( + ν = νhb, + discrete_form = true, + parameters = 40 * 24 * 60 * 60, # 40 days + ) + vertical_closure = OC.VerticalScalarDiffusivity(ν = 1e-5, κ = 2e-6) + catke_closure = CO.Oceans.default_ocean_closure() + closure = (catke_closure, eddy_closure, horizontal_viscosity, vertical_closure) else - forcing = (;) - end - - # Create ocean simulation - free_surface = OC.SplitExplicitFreeSurface(grid; substeps = 70) - momentum_advection = OC.VectorInvariant() - tracer_advection = OC.WENO(order = 5) - horizontal_viscosity = OC.HorizontalScalarBiharmonicDiffusivity(ν = 1e11) + # Simpler setup + @info "Using simpler ocean setup; to be used for software testing only." + free_surface = OC.SplitExplicitFreeSurface(grid; substeps = 70) + momentum_advection = OC.VectorInvariant() + tracer_advection = OC.WENO(order = 5) + horizontal_viscosity = OC.HorizontalScalarBiharmonicDiffusivity(ν = 1e11) + vertical_mixing = OC.ConvectiveAdjustmentVerticalDiffusivity( + background_κz = 1e-5, + convective_κz = 0.1, + background_νz = 1e-4, + convective_νz = 0.1, + ) - # Use Float32 for the vertical mixing parameters to avoid parameter memory limits - vertical_mixing = OC.CATKEVerticalDiffusivity(Float32) + closure = (horizontal_viscosity, vertical_mixing) + end Δt = isnothing(dt) ? CO.OceanSimulations.estimate_maximum_Δt(grid) : dt ocean = CO.ocean_simulation( grid; Δt, - forcing, momentum_advection, tracer_advection, + timestepper = :SplitRungeKutta3, free_surface, - closure = (horizontal_viscosity, vertical_mixing), + closure, ) # Set initial condition to EN4 state estimate at start_date OC.set!(ocean.model, T = en4_temperature[1], S = en4_salinity[1]) - long_cc = OC.λnodes(grid, OC.Center(), OC.Center(), OC.Center()) - lat_cc = OC.φnodes(grid, OC.Center(), OC.Center(), OC.Center()) - - # TODO: Go from 0 to Nx+1, Ny+1 (for halos) (for LatLongGrid) - - # Construct a remapper from the exchange grid to `Center, Center` fields - long_cc = reshape(long_cc, length(long_cc), 1) - lat_cc = reshape(lat_cc, 1, length(lat_cc)) - target_points_cc = @. CC.Geometry.LatLongPoint(lat_cc, long_cc) - # TODO: We can remove the `nothing` after CC > 0.14.33 - remapper_cc = CC.Remapping.Remapper(boundary_space, target_points_cc, nothing) - - # Construct two 2D Center/Center fields to use as scratch space while remapping - scratch_cc1 = OC.Field{OC.Center, OC.Center, Nothing}(grid) - scratch_cc2 = OC.Field{OC.Center, OC.Center, Nothing}(grid) - - # Construct two scratch arrays to use while remapping - # We get the array type and dimensions from the remapper object to maintain consistency - ArrayType = ClimaComms.array_type(remapper_cc.space) - interpolated_values_dim..., _buffer_length = size(remapper_cc._interpolated_values) - - scratch_arr1 = ArrayType(zeros(FT, interpolated_values_dim...)) - scratch_arr2 = ArrayType(zeros(FT, interpolated_values_dim...)) - scratch_arr3 = ArrayType(zeros(FT, interpolated_values_dim...)) - - # Allocate space for a Field of UVVectors, which we need for remapping momentum fluxes - temp_uv_vec = CC.Fields.Field(CC.Geometry.UVVector{FT}, boundary_space) - - remapping = (; - remapper_cc, - scratch_cc1, - scratch_cc2, - scratch_arr1, - scratch_arr2, - scratch_arr3, - temp_uv_vec, - ) + # Construct the remapper object and allocate scratch space + remapping = construct_remappers(grid, boundary_space) # Get some ocean properties and parameters ocean_properties = (; @@ -214,44 +180,44 @@ function OceananigansSimulation( C_to_K = coupled_param_dict["temperature_water_freeze"], ) - # Before version 0.96.22, the NetCDFWriter was broken on GPU - if arch isa OC.CPU - # Save all tracers and velocities to a NetCDF file at daily frequency - outputs = merge(ocean.model.tracers, ocean.model.velocities) - surface_writer = OC.NetCDFWriter( - ocean.model, - outputs; - schedule = OC.TimeInterval(86400), # Daily output - filename = joinpath(output_dir, "ocean_diagnostics.nc"), - indices = (:, :, grid.Nz), - overwrite_existing = true, - array_type = Array{FT}, - ) - free_surface_writer = OC.NetCDFWriter( - ocean.model, - (; η = ocean.model.free_surface.displacement); - schedule = OC.TimeInterval(3600), # hourly snapshots - filename = joinpath(output_dir, "ocean_free_surface.nc"), - overwrite_existing = true, - array_type = Array{FT}, - ) - Tflux = ocean.model.tracers.T.boundary_conditions.top.condition - Sflux = ocean.model.tracers.S.boundary_conditions.top.condition - uflux = ocean.model.velocities.u.boundary_conditions.top.condition - vflux = ocean.model.velocities.v.boundary_conditions.top.condition - fluxes_writer = OC.NetCDFWriter( - ocean.model, - (; Tflux, Sflux, uflux, vflux); - schedule = OC.TimeInterval(3600), # hourly snapshots - filename = joinpath(output_dir, "ocean_fluxes.nc"), - overwrite_existing = true, - array_type = Array{FT}, - ) + # Save all tracers and velocities to a JLD2 file at daily frequency + outputs = merge(ocean.model.tracers, ocean.model.velocities) + surface_writer = OC.JLD2Writer( + ocean.model, + outputs; + schedule = OC.TimeInterval(86400), # Daily output + filename = joinpath(output_dir, "ocean_diagnostics.jld2"), + indices = (:, :, grid.Nz), + overwrite_existing = true, + array_type = Array{FT}, + ) - ocean.output_writers[:surface] = surface_writer - ocean.output_writers[:free_surface] = free_surface_writer - ocean.output_writers[:fluxes] = fluxes_writer - end + # Save free surface to a JLD2 file at hourly frequency + free_surface_writer = OC.JLD2Writer( + ocean.model, + (; η = ocean.model.free_surface.displacement); + schedule = OC.TimeInterval(3600), # hourly snapshots + filename = joinpath(output_dir, "ocean_free_surface.jld2"), + overwrite_existing = true, + array_type = Array{FT}, + ) + + # Save fluxes to a JLD2 file at hourly frequency + Tflux = ocean.model.tracers.T.boundary_conditions.top.condition + Sflux = ocean.model.tracers.S.boundary_conditions.top.condition + uflux = ocean.model.velocities.u.boundary_conditions.top.condition + vflux = ocean.model.velocities.v.boundary_conditions.top.condition + fluxes_writer = OC.JLD2Writer( + ocean.model, + (; Tflux, Sflux, uflux, vflux); + schedule = OC.TimeInterval(3600), # hourly snapshots + filename = joinpath(output_dir, "ocean_fluxes.jld2"), + overwrite_existing = true, + array_type = Array{FT}, + ) + ocean.output_writers[:surface] = surface_writer + ocean.output_writers[:free_surface] = free_surface_writer + ocean.output_writers[:fluxes] = fluxes_writer # Initialize with 0 ice concentration; this will be updated in `resolve_area_fractions!` # if the ocean is coupled to a non-prescribed sea ice model. @@ -269,6 +235,58 @@ function OceananigansSimulation( ) end +include("remapping.jl") +""" + construct_remappers(grid_oc, boundary_space) + +Given an Oceananigans grid and a ClimaCore boundary space, construct the +remappers needed to remap between the two grids in both directions. + +Returns a remapper from the Oceananigans grid to the ClimaCore boundary space. +To regrid from Oceananigans to ClimaCore, use `CR.regrid!(dest_vector, remapper_oc_to_cc, src_vector)`. +To regrid from ClimaCore to Oceananigans, use `CR.regrid!(dest_vector, transpose(remapper_oc_to_cc), src_vector)`. +""" +function construct_remappers(grid_oc, boundary_space) + # Move grids to CPU since ConservativeRegridding doesn't support GPU grids yet + grid_oc_underlying_cpu = OC.on_architecture(OC.CPU(), grid_oc.underlying_grid) + boundary_space_cpu = Adapt.adapt_structure(Array, boundary_space) + + # Create the remapper from the Oceananigans grid to the ClimaCore boundary space + remapper_oc_to_cc = CR.Regridder( + boundary_space_cpu, + grid_oc_underlying_cpu; + normalize = false, + threaded = false, + ) + + # Move remapper to GPU if needed + remapper_oc_to_cc = OC.on_architecture(OC.architecture(grid_oc), remapper_oc_to_cc) + + # Create a field of ones on the boundary space so we can compute element areas + field_ones_cc = CC.Fields.ones(boundary_space) + + # Allocate a vector with length equal to the number of elements in the target space + # To be used as a temp field for remapping + FT = CC.Spaces.undertype(boundary_space) + ArrayType = ClimaComms.array_type(boundary_space) + value_per_element_cc = ArrayType(zeros(FT, CC.Meshes.nelements(boundary_space.grid.topology.mesh))) + + # Construct two 2D Oceananigans Center/Center fields to use as scratch space while remapping + scratch_field_oc1 = OC.Field{OC.Center, OC.Center, Nothing}(grid_oc) + scratch_field_oc2 = OC.Field{OC.Center, OC.Center, Nothing}(grid_oc) + + # Allocate space for a Field of UVVectors, which we need for remapping momentum fluxes + temp_uv_vec = CC.Fields.Field(CC.Geometry.UVVector{FT}, boundary_space) + return (; + remapper_oc_to_cc, + field_ones_cc, + value_per_element_cc, + scratch_field_oc1, + scratch_field_oc2, + temp_uv_vec, + ) +end + """ FieldExchanger.resolve_area_fractions!(ocean_sim, ice_sim, land_fraction) @@ -368,30 +386,20 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi F_turb_ρτxz_uv = sim.remapping.temp_uv_vec.components.data.:1 F_turb_ρτyz_uv = sim.remapping.temp_uv_vec.components.data.:2 - # Remap momentum fluxes onto reduced 2D Center, Center fields using scratch arrays and fields - CC.Remapping.interpolate!( - sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, - F_turb_ρτxz_uv, - ) - OC.set!(sim.remapping.scratch_cc1, sim.remapping.scratch_arr1) # zonal momentum flux - CC.Remapping.interpolate!( - sim.remapping.scratch_arr2, - sim.remapping.remapper_cc, - F_turb_ρτyz_uv, - ) - OC.set!(sim.remapping.scratch_cc2, sim.remapping.scratch_arr2) # meridional momentum flux + # Remap momentum fluxes onto reduced 2D Center, Center fields + Interfacer.remap!(sim.remapping.scratch_field_oc1, F_turb_ρτxz, sim.remapping) # zonal momentum flux + Interfacer.remap!(sim.remapping.scratch_field_oc2, F_turb_ρτyz, sim.remapping) # meridional momentum flux # Rename for clarity; these are now Center, Center Oceananigans fields - F_turb_ρτxz_cc = sim.remapping.scratch_cc1 - F_turb_ρτyz_cc = sim.remapping.scratch_cc2 + oc_F_turb_ρτxz = sim.remapping.scratch_field_oc1 + oc_F_turb_ρτyz = sim.remapping.scratch_field_oc2 # Weight by (1 - sea ice concentration) - OC.interior(F_turb_ρτxz_cc, :, :, 1) .= - OC.interior(F_turb_ρτxz_cc, :, :, 1) .* (1.0 .- ice_concentration) ./ + OC.interior(oc_F_turb_ρτxz, :, :, 1) .= + OC.interior(oc_F_turb_ρτxz, :, :, 1) .* (1.0 .- ice_concentration) ./ reference_density - OC.interior(F_turb_ρτyz_cc, :, :, 1) .= - OC.interior(F_turb_ρτyz_cc, :, :, 1) .* (1.0 .- ice_concentration) ./ + OC.interior(oc_F_turb_ρτyz, :, :, 1) .= + OC.interior(oc_F_turb_ρτyz, :, :, 1) .* (1.0 .- ice_concentration) ./ reference_density # Set the momentum flux BCs at the correct locations using the remapped scratch fields @@ -400,17 +408,17 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi set_from_extrinsic_vector!( (; u = oc_flux_u, v = oc_flux_v), grid, - F_turb_ρτxz_cc, - F_turb_ρτyz_cc, + oc_F_turb_ρτxz, + oc_F_turb_ρτyz, ) # Remap the latent and sensible heat fluxes using scratch arrays - CC.Remapping.interpolate!(sim.remapping.scratch_arr1, sim.remapping.remapper_cc, F_lh) # latent heat flux - CC.Remapping.interpolate!(sim.remapping.scratch_arr2, sim.remapping.remapper_cc, F_sh) # sensible heat flux + Interfacer.remap!(sim.remapping.scratch_field_oc1, F_lh, sim.remapping) # latent heat flux + Interfacer.remap!(sim.remapping.scratch_field_oc2, F_sh, sim.remapping) # sensible heat flux # Rename for clarity; recall F_turb_energy = F_lh + F_sh - remapped_F_lh = sim.remapping.scratch_arr1 - remapped_F_sh = sim.remapping.scratch_arr2 + remapped_F_lh = OC.interior(sim.remapping.scratch_field_oc1, :, :, 1) + remapped_F_sh = OC.interior(sim.remapping.scratch_field_oc2, :, :, 1) # TODO: Note, SW radiation penetrates the surface. Right now, we just put # everything on the surface, but later we will need to account for this. @@ -423,12 +431,9 @@ function FluxCalculator.update_turbulent_fluxes!(sim::OceananigansSimulation, fi # Add the part of the salinity flux that comes from the moisture flux, we also need to # add the component due to precipitation (that was done with the radiative fluxes) - CC.Remapping.interpolate!( - sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, - F_turb_moisture, - ) - moisture_fresh_water_flux = sim.remapping.scratch_arr1 ./ reference_density + Interfacer.remap!(sim.remapping.scratch_field_oc1, F_turb_moisture, sim.remapping) # moisture flux + moisture_fresh_water_flux = + OC.interior(sim.remapping.scratch_field_oc1, :, :, 1) ./ reference_density oc_flux_S = surface_flux(sim.ocean.model.tracers.S) surface_salinity = OC.interior(sim.ocean.model.tracers.S, :, :, grid.Nz) OC.interior(oc_flux_S, :, :, 1) .= @@ -467,19 +472,11 @@ function FieldExchanger.update_sim!(sim::OceananigansSimulation, csf) Nz = sim.ocean.model.grid.Nz # Remap radiative flux onto scratch array; rename for clarity - CC.Remapping.interpolate!( - sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, - csf.SW_d, - ) - remapped_SW_d = sim.remapping.scratch_arr1 + Interfacer.remap!(sim.remapping.scratch_field_oc1, csf.SW_d, sim.remapping) # shortwave radiation + remapped_SW_d = OC.interior(sim.remapping.scratch_field_oc1, :, :, 1) - CC.Remapping.interpolate!( - sim.remapping.scratch_arr2, - sim.remapping.remapper_cc, - csf.LW_d, - ) - remapped_LW_d = sim.remapping.scratch_arr2 + Interfacer.remap!(sim.remapping.scratch_field_oc2, csf.LW_d, sim.remapping) # longwave radiation + remapped_LW_d = OC.interior(sim.remapping.scratch_field_oc2, :, :, 1) # Update only the part due to radiative fluxes. For the full update, the component due # to latent and sensible heat is missing and will be updated in update_turbulent_fluxes. @@ -496,19 +493,12 @@ function FieldExchanger.update_sim!(sim::OceananigansSimulation, csf) ) ) ./ (reference_density * heat_capacity) - # Remap precipitation fields onto scratch arrays; rename for clarity - CC.Remapping.interpolate!( - sim.remapping.scratch_arr1, - sim.remapping.remapper_cc, - csf.P_liq, - ) - CC.Remapping.interpolate!( - sim.remapping.scratch_arr2, - sim.remapping.remapper_cc, - csf.P_snow, - ) - remapped_P_liq = sim.remapping.scratch_arr1 - remapped_P_snow = sim.remapping.scratch_arr2 + # Remap precipitation fields onto scratch fields; rename for clarity + Interfacer.remap!(sim.remapping.scratch_field_oc1, csf.P_liq, sim.remapping) # liquid precipitation + remapped_P_liq = OC.interior(sim.remapping.scratch_field_oc1, :, :, 1) + + Interfacer.remap!(sim.remapping.scratch_field_oc2, csf.P_snow, sim.remapping) # snow precipitation + remapped_P_snow = OC.interior(sim.remapping.scratch_field_oc2, :, :, 1) # Virtual salt flux oc_flux_S = surface_flux(sim.ocean.model.tracers.S) @@ -519,6 +509,122 @@ function FieldExchanger.update_sim!(sim::OceananigansSimulation, csf) return nothing end +""" + ocean_seaice_fluxes!(ocean_sim::OceananigansSimulation, ice_sim::ClimaSeaIceSimulation) + +Compute the fluxes between the ocean and sea ice, storing them in the `ocean_ice_fluxes` +fields of the ocean and sea ice simulations. + +This function assumes both simulations share the same grid, so no remapping is done. + +Both simulations have had their atmospheric fluxes updated already in this timestep +(see `update_sim!` and `update_turbulent_fluxes!`), so we add the contributions from the +ocean-sea ice interactions to the existing fluxes, rather than overwriting all fluxes. + +!!! note + This function must be called after the turbulent fluxes have been updated in both + simulations. Here only the contributions from the sea ice/ocean interactions + are added to the fluxes. +""" +function FluxCalculator.ocean_seaice_fluxes!( + ocean_sim::OceananigansSimulation, + ice_sim::ClimaSeaIceSimulation, +) + ocean_properties = ocean_sim.ocean_properties + ice_concentration = Interfacer.get_field(ice_sim, Val(:ice_concentration)) + + # Update the sea ice concentration in the ocean simulation + ocean_sim.ice_concentration .= ice_concentration + + # Compute the fluxes and store them in the both simulations + CO.OceanSeaIceModels.InterfaceComputations.compute_sea_ice_ocean_fluxes!( + ice_sim.ocean_ice_interface, + ocean_sim.ocean, + ice_sim.ice, + ocean_properties, + ) + + ## Update the internals of the sea ice model + # Set the bottom heat flux to the sum of the frazil and interface heat fluxes + bottom_heat_flux = ice_sim.ice.model.external_heat_fluxes.bottom + + Qf = ice_sim.ocean_ice_interface.fluxes.frazil_heat # frazil heat flux + Qi = ice_sim.ocean_ice_interface.fluxes.interface_heat # interfacial heat flux + bottom_heat_flux .= Qf .+ Qi + + ## Update the internals of the ocean model + ρₒ⁻¹ = 1 / ocean_sim.ocean_properties.reference_density + cₒ = ocean_sim.ocean_properties.heat_capacity + + # Compute fluxes for u, v, T, and S from momentum, heat, and freshwater fluxes + oc_flux_u = surface_flux(ocean_sim.ocean.model.velocities.u) + oc_flux_v = surface_flux(ocean_sim.ocean.model.velocities.v) + + ρτxio = ice_sim.ocean_ice_interface.fluxes.x_momentum # sea_ice - ocean zonal momentum flux + ρτyio = ice_sim.ocean_ice_interface.fluxes.y_momentum # sea_ice - ocean meridional momentum flux + + # Update the momentum flux contributions from ocean/sea ice fluxes + grid = ocean_sim.ocean.model.grid + arch = OC.Architectures.architecture(grid) + OC.Utils.launch!( + arch, + grid, + :xy, + _add_ocean_ice_stress!, + oc_flux_u, + oc_flux_v, + grid, + ρτxio, + ρτyio, + ρₒ⁻¹, + ice_concentration, + ) + + oc_flux_T = surface_flux(ocean_sim.ocean.model.tracers.T) + OC.interior(oc_flux_T, :, :, 1) .+= + OC.interior(ice_concentration, :, :, 1) .* OC.interior(Qi, :, :, 1) .* ρₒ⁻¹ ./ cₒ + + oc_flux_S = surface_flux(ocean_sim.ocean.model.tracers.S) + OC.interior(oc_flux_S, :, :, 1) .+= + OC.interior(ice_concentration, :, :, 1) .* + OC.interior(ice_sim.ocean_ice_interface.fluxes.salt, :, :, 1) + + return nothing +end + +""" + _add_ocean_ice_stress!(oc_flux_u, oc_flux_v, grid, ρτxio, ρτyio, ρₒ⁻¹, ice_concentration) + +Add the contribution from the ocean-ice stress to the surface fluxes for each +component of the ocean velocity (u and v). + +Arguments: +- `oc_flux_u`: [Field] the surface flux for the ocean zonal velocity. +- `oc_flux_v`: [Field] the surface flux for the ocean meridional velocity. +- `grid`: [Grid] the grid used by ocean and ice. +- `ρτxio`: [Field] the ice-ocean zonal momentum flux. +- `ρτyio`: [Field] the ice-ocean meridional momentum flux. +- `ρₒ⁻¹`: [Float] the inverse of the ocean reference density. +- `ice_concentration`: [Field] the sea ice concentration. +""" +@kernel function _add_ocean_ice_stress!( + oc_flux_u, + oc_flux_v, + grid, + ρτxio, + ρτyio, + ρₒ⁻¹, + ice_concentration, +) + i, j = @index(Global, NTuple) + + # ℑxᶠᵃᵃ: interpolate faces to centers + oc_flux_u[i, j, 1] += + ρτxio[i, j, 1] * ρₒ⁻¹ * OC.Operators.ℑxᶠᵃᵃ(i, j, 1, grid, ice_concentration) + oc_flux_v[i, j, 1] += + ρτyio[i, j, 1] * ρₒ⁻¹ * OC.Operators.ℑyᵃᶠᵃ(i, j, 1, grid, ice_concentration) +end + """ get_model_prog_state(sim::OceananigansSimulation) diff --git a/experiments/ClimaEarth/components/ocean/remapping.jl b/experiments/ClimaEarth/components/ocean/remapping.jl new file mode 100644 index 0000000000..31ab1aeb37 --- /dev/null +++ b/experiments/ClimaEarth/components/ocean/remapping.jl @@ -0,0 +1,112 @@ +const ConservativeRegriddingCCExt = + Base.get_extension(CR, :ConservativeRegriddingClimaCoreExt) + +### Extensions of Interfacer.jl functions for Oceananigans fields/grids +# Non-allocating ClimaCore -> Oceananigans remap +function Interfacer.remap!(target_field::OC.Field, source_field::CC.Fields.Field, remapping) + ConservativeRegriddingCCExt.get_value_per_element!( + remapping.value_per_element_cc, + source_field, + remapping.field_ones_cc, + ) + + # Get the index of the top level (surface); 1 for 2D fields, Nz for 3D fields + z = size(target_field, 3) + dst = vec(OC.interior(target_field, :, :, z)) + src = remapping.value_per_element_cc + + # Regrid the source field to the target field + CR.regrid!(dst, transpose(remapping.remapper_oc_to_cc), src) + return nothing +end +# Allocating ClimaCore -> Oceananigans remap +function Interfacer.remap( + target_space::Union{ + OC.OrthogonalSphericalShellGrid, + OC.ImmersedBoundaryGrid, + OC.LatitudeLongitudeGrid, + }, + source_field::CC.Fields.Field, + remapping, +) + target_field = OC.Field{OC.Center, OC.Center, Nothing}(target_space) + Interfacer.remap!(target_field, source_field, remapping) + return target_field +end + +# Non-allocating Oceananigans Field -> ClimaCore remap +function Interfacer.remap!(target_field::CC.Fields.Field, source_field::OC.Field, remapping) + # Get the index of the top level (surface); 1 for 2D fields, Nz for 3D fields + z = size(source_field, 3) + src = vec(OC.interior(source_field, :, :, z)) + + # Store the remapped FV values in a vector of length equal to the number of elements in the target space + dst = remapping.value_per_element_cc + + # Regrid the source field to the target field + CR.regrid!(dst, remapping.remapper_oc_to_cc, src) + + # Convert the vector of remapped values to a ClimaCore Field with one value per element + ConservativeRegriddingCCExt.set_value_per_element!(target_field, dst) + return nothing +end +# Allocating Oceananigans Field -> ClimaCore remap +function Interfacer.remap( + target_space::CC.Spaces.AbstractSpace, + source_field::OC.Field, + remapping, +) + target_field = CC.Fields.zeros(target_space) + Interfacer.remap!(target_field, source_field, remapping) + return target_field +end + +# Non-allocating Oceananigans operation -> ClimaCore remap +function Interfacer.remap!( + target_field::CC.Fields.Field, + operation::OC.AbstractOperations.AbstractOperation, + remapping, +) + evaluated_field = OC.Field(operation) + OC.compute!(evaluated_field) + Interfacer.remap!(target_field, evaluated_field, remapping) + return nothing +end +# Allocating Oceananigans operation -> ClimaCore remap +function Interfacer.remap( + target_space::CC.Spaces.AbstractSpace, + operation::OC.AbstractOperations.AbstractOperation, + remapping, +) + target_field = CC.Fields.zeros(target_space) + Interfacer.remap!(target_field, operation, remapping) + return target_field +end + +# Handle the case of remapping a scalar number to a ClimaCore space +Interfacer.remap!(target_field::CC.Fields.Field, source_field::Number, remapping) = + Interfacer.remap!(target_field, source_field) +Interfacer.remap(target_space::CC.Spaces.AbstractSpace, source_num::Number, remapping) = + Interfacer.remap(target_space, source_num, remapping) + +# Handle the case of remapping the area fraction field, which is a ClimaCore Field +Interfacer.remap!(target_field::CC.Fields.Field, source_field::CC.Fields.Field, remapping) = + Interfacer.remap!(target_field, source_field) + +# Extend Interfacer.get_field to allow automatic remapping to the target space +function Interfacer.get_field!(target_field, sim::Union{OceananigansSimulation, ClimaSeaIceSimulation}, quantity) + Interfacer.remap!(target_field, Interfacer.get_field(sim, quantity), sim.remapping) + return nothing +end +# TODO see if we can remove this allocating version +function Interfacer.get_field( + target_space::CC.Spaces.AbstractSpace, + sim::OceananigansSimulation, + quantity, +) + return Interfacer.remap( + target_space, + Interfacer.get_field(sim, quantity), + sim.remapping, + ) +end diff --git a/experiments/ClimaEarth/setup_run.jl b/experiments/ClimaEarth/setup_run.jl index 9d047398c4..8e18491007 100644 --- a/experiments/ClimaEarth/setup_run.jl +++ b/experiments/ClimaEarth/setup_run.jl @@ -72,8 +72,8 @@ include("components/land/climaland_integrated.jl") include("components/ocean/slab_ocean.jl") include("components/ocean/prescr_ocean.jl") include("components/ocean/prescr_seaice.jl") -include("components/ocean/oceananigans.jl") include("components/ocean/clima_seaice.jl") +include("components/ocean/oceananigans.jl") #= ### Configuration Dictionaries @@ -137,6 +137,7 @@ function CoupledSimulation(config_dict::AbstractDict) parameter_files, era5_initial_condition_dir, ocean_model, + simple_ocean, ice_model, land_fraction_source, binary_area_fraction, @@ -231,7 +232,7 @@ function CoupledSimulation(config_dict::AbstractDict) boundary_space = CC.Spaces.horizontal_space(atmos_sim.domain.face_space) else h_elem = config_dict["h_elem"] - n_quad_points = 4 + n_quad_points = config_dict["nh_poly"] + 1 radius = coupled_param_dict["planet_radius"] # in meters boundary_space = CC.CommonSpaces.CubedSphereSpace(FT; radius, n_quad_points, h_elem) end @@ -326,7 +327,7 @@ function CoupledSimulation(config_dict::AbstractDict) boundary_space, # Arguments used by Oceananigans output_dir = dir_paths.ocean_output_dir, - ice_model, + simple_ocean, # Arguments used by prescribed ocean sst_path = subseasonal_sst, # Arguments used by slab ocean diff --git a/src/Input.jl b/src/Input.jl index 2ff1aa368b..4ce2032151 100644 --- a/src/Input.jl +++ b/src/Input.jl @@ -120,6 +120,10 @@ function argparse_settings() help = "Number of horizontal elements to use for the boundary space [16 (default)]" arg_type = Int default = 16 + "--nh_poly" + help = "Polynomial order to use for the boundary space [3 (default)]" + arg_type = Int + default = 3 "--share_surface_space" help = "Boolean flag indicating whether to share the surface space between the surface models, atmosphere, and boundary [`true` (default), `false`]" arg_type = Bool @@ -229,6 +233,10 @@ function argparse_settings() help = "Ocean model to use. [`prescribed` (default), `oceananigans`, `slab`, `nothing`]" arg_type = String default = "prescribed" + "--simple_ocean" + help = "Boolean flag indicating whether to use a simpler ocean model setup with Oceananigans [`false` (default), `true`]" + arg_type = Bool + default = false # Ice model specific "--ice_model" help = "Sea ice model to use. [`prescribed` (default), `clima_seaice`, `nothing`]" @@ -424,6 +432,7 @@ function get_coupler_args(config_dict::Dict) # Ocean model-specific information ocean_model = Val(Symbol(config_dict["ocean_model"])) + simple_ocean = config_dict["simple_ocean"] # Ice model-specific information ice_model = Val(Symbol(config_dict["ice_model"])) @@ -472,6 +481,7 @@ function get_coupler_args(config_dict::Dict) parameter_files, era5_initial_condition_dir, ocean_model, + simple_ocean, ice_model, land_fraction_source, binary_area_fraction, diff --git a/test/input_tests.jl b/test/input_tests.jl index e589c40f05..fa80d3ae3b 100644 --- a/test/input_tests.jl +++ b/test/input_tests.jl @@ -89,6 +89,7 @@ end "coupler_toml" => [], "era5_initial_condition_dir" => nothing, "ocean_model" => "prescribed", + "simple_ocean" => false, "ice_model" => "prescribed", "land_fraction_source" => "etopo", "binary_area_fraction" => true,