From b4c8b212980eacd3b32d50300dcc10880a123e3d Mon Sep 17 00:00:00 2001 From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com> Date: Tue, 9 Jul 2024 10:14:26 -0700 Subject: [PATCH] make aerosol radiation work on GPU --- config/model_configs/gpu_aquaplanet_dyamond.yml | 1 + src/callbacks/callbacks.jl | 5 ++--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/config/model_configs/gpu_aquaplanet_dyamond.yml b/config/model_configs/gpu_aquaplanet_dyamond.yml index 50c4a50869..c4e0430a78 100644 --- a/config/model_configs/gpu_aquaplanet_dyamond.yml +++ b/config/model_configs/gpu_aquaplanet_dyamond.yml @@ -18,5 +18,6 @@ surface_setup: "DefaultMoninObukhov" rayleigh_sponge: true dt: "100secs" t_end: "8hours" +aerosol_radiation: true prescribed_aerosols: ["CB1", "CB2", "DST01", "DST02", "DST03", "DST04", "OC1", "OC2", "SO4", "SOA", "SSLT01", "SSLT02", "SSLT03", "SSLT04"] toml: [toml/longrun_aquaplanet_dyamond.toml] diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl index 1da95fe6b2..0f919fb19f 100644 --- a/src/callbacks/callbacks.jl +++ b/src/callbacks/callbacks.jl @@ -262,10 +262,9 @@ function set_aerosol_type(; OC2 = 0, _..., ) - _, index = findmax((DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)) - return index + return argmax(n -> (DST01, SSLT01, SO4, CB1, CB2, OC1, OC2)[n], 1:7) end -set_aerosol_type(NT) = set_aerosol_type(; NT...) +set_aerosol_type(nt) = set_aerosol_type(; nt...) NVTX.@annotate function save_state_to_disk_func(integrator, output_dir) (; t, u, p) = integrator