From 3e39e7a4757b187344ca1097168c6ebe49a71d82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=B0=A2=E8=90=A7=E6=B6=AF?= Date: Mon, 2 Sep 2024 09:25:15 -0700 Subject: [PATCH] foo1 --- .buildkite/pipeline.yml | 15 + Project.toml | 1 + ...ngle_column_nonorographic_gravity_wave.yml | 2 +- examples/Manifest.toml | 15 +- examples/Project.toml | 1 + perf/Manifest.toml | 15 +- perf/Project.toml | 1 + .../non_orographic_gravity_wave.jl | 324 ++++++++++-------- .../nogw_test_3d.jl | 5 +- .../nogw_test_mima.jl | 3 +- .../nogw_test_single_column.jl | 83 ++++- 11 files changed, 299 insertions(+), 166 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index e0d0ceac9b5..95bd03f9876 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -25,17 +25,20 @@ steps: - echo "--- Instantiate project" - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + - "julia --project -e 'using Pkg; Pkg.add(PackageSpec(name=\"UnrolledUtilities\", rev=\"dy/lazy_and_low_storage\"))'" - "julia --project -e 'using Pkg; Pkg.precompile()'" - "julia --project -e 'using Pkg; Pkg.status()'" - echo "--- Instantiate examples" - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + - "julia --project=examples -e 'using Pkg; Pkg.develop(path=\".\")'" - "julia --project=examples -e 'using Pkg; Pkg.precompile()'" - "julia --project=examples -e 'using CUDA; CUDA.precompile_runtime()'" - "julia --project=examples -e 'using Pkg; Pkg.status()'" - echo "--- Instantiate perf" - "julia --project=perf -e 'using Pkg; Pkg.instantiate(;verbose=true)'" + - "julia --project=examples -e 'using Pkg; Pkg.develop(path=\".\")'" - "julia --project=perf -e 'using Pkg; Pkg.precompile()'" - "julia --project=perf -e 'using Pkg; Pkg.status()'" @@ -761,6 +764,18 @@ steps: - group: "GPU" steps: + - label: "GPU:Gravity waves" + command: > + julia --color=yes --project=examples examples/hybrid/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_wave2/*" + env: + CLIMACOMMS_DEVICE: "CUDA" + agents: + slurm_gpus: 1 + slurm_mem: 16G + - label: "GPU: baroclinic wave" key: "sphere_baroclinic_wave_rhoe_gpu" command: > diff --git a/Project.toml b/Project.toml index 5f9b9d4ada5..010fad02164 100644 --- a/Project.toml +++ b/Project.toml @@ -33,6 +33,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f" Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" +UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b" YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6" [compat] diff --git a/config/model_configs/single_column_nonorographic_gravity_wave.yml b/config/model_configs/single_column_nonorographic_gravity_wave.yml index d36e27440a7..7f3fe517260 100644 --- a/config/model_configs/single_column_nonorographic_gravity_wave.yml +++ b/config/model_configs/single_column_nonorographic_gravity_wave.yml @@ -1,7 +1,7 @@ dt_save_state_to_disk: "400secs" initial_condition: "IsothermalProfile" t_end: "1500secs" -config: "column" +# config: "column" hyperdiff: false dt: "400secs" non_orographic_gravity_wave: true diff --git a/examples/Manifest.toml b/examples/Manifest.toml index 0b0edeb55db..02109ea7360 100644 --- a/examples/Manifest.toml +++ b/examples/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.4" manifest_format = "2.0" -project_hash = "c1104a002a5c87ddfd3855d6cdf90ee1b0892b0a" +project_hash = "ce341806d75ec75e183549eadf24687e56ec64f2" [[deps.ADTypes]] git-tree-sha1 = "99a6f5d0ce1c7c6afdb759daa30226f71c54f6b0" @@ -321,7 +321,7 @@ version = "0.5.7" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"] path = ".." uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717" version = "0.27.4" @@ -2723,6 +2723,17 @@ git-tree-sha1 = "6cc9d682755680e0f0be87c56392b7651efc2c7b" uuid = "9602ed7d-8fef-5bc8-8597-8f21381861e8" version = "0.1.5" +[[deps.UnrolledUtilities]] +git-tree-sha1 = "2a272e181393584f697a45a371cca200358eda31" +repo-rev = "dy/lazy_and_low_storage" +repo-url = "https://github.com/CliMA/UnrolledUtilities.jl.git" +uuid = "0fe1646c-419e-43be-ac14-22321958931b" +version = "0.1.2" +weakdeps = ["StaticArrays"] + + [deps.UnrolledUtilities.extensions] + UnrolledUtilitiesStaticArraysExt = "StaticArrays" + [[deps.UnsafeAtomics]] git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278" uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f" diff --git a/examples/Project.toml b/examples/Project.toml index c89bea32734..e373918093f 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -62,6 +62,7 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Tar = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" +UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b" YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6" [compat] diff --git a/perf/Manifest.toml b/perf/Manifest.toml index 004d5af5686..378b7c3f4b1 100644 --- a/perf/Manifest.toml +++ b/perf/Manifest.toml @@ -2,7 +2,7 @@ julia_version = "1.10.4" manifest_format = "2.0" -project_hash = "ddcac0ed38fc33b661f9c30ef50fd3e33f33a3e0" +project_hash = "0b89db3c2616b7468c64bb528229695548afef20" [[deps.ADTypes]] git-tree-sha1 = "99a6f5d0ce1c7c6afdb759daa30226f71c54f6b0" @@ -332,7 +332,7 @@ version = "0.5.6" GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6" [[deps.ClimaAtmos]] -deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "YAML"] +deps = ["Adapt", "ArgParse", "ArtifactWrappers", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "DiffEqBase", "FastGaussQuadrature", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "Pkg", "RRTMGP", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"] path = ".." uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717" version = "0.27.4" @@ -2833,6 +2833,17 @@ git-tree-sha1 = "6cc9d682755680e0f0be87c56392b7651efc2c7b" uuid = "9602ed7d-8fef-5bc8-8597-8f21381861e8" version = "0.1.5" +[[deps.UnrolledUtilities]] +git-tree-sha1 = "2e8613c8c08f67005eb7f92151cfcc78ca5a841d" +repo-rev = "dy/lazy_and_low_storage" +repo-url = "https://github.com/CliMA/UnrolledUtilities.jl.git" +uuid = "0fe1646c-419e-43be-ac14-22321958931b" +version = "0.1.2" +weakdeps = ["StaticArrays"] + + [deps.UnrolledUtilities.extensions] + UnrolledUtilitiesStaticArraysExt = "StaticArrays" + [[deps.UnsafeAtomics]] git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278" uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f" diff --git a/perf/Project.toml b/perf/Project.toml index 8d62abf7d99..3bb62935010 100644 --- a/perf/Project.toml +++ b/perf/Project.toml @@ -66,6 +66,7 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c" URIs = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4" +UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b" YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6" [compat] diff --git a/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl b/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl index 14ee475fd87..443bc39f9b3 100644 --- a/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl +++ b/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl @@ -1,12 +1,18 @@ ##### ##### Non-orographic gravity wave parameterization ##### - +using UnrolledUtilities import ClimaCore.Spaces as Spaces import ClimaCore.Fields as Fields import ClimaCore.Geometry as Geometry import ClimaCore.Operators as Operators +Operators.placeholder_space(current_space, parent_space) = + current_space isa Spaces.AbstractSpace && + parent_space isa Spaces.AbstractSpace && + Spaces.issubspace(current_space, parent_space) ? + Operators.PlaceholderSpace() : current_space + non_orographic_gravity_wave_cache(Y, atmos::AtmosModel) = non_orographic_gravity_wave_cache( Y, @@ -27,7 +33,7 @@ function non_orographic_gravity_wave_cache( (; source_height, Bw, Bn, Bt_0, dc, cmax, c0, nk, cw, cn) = gw nc = Int(floor(FT(2 * cmax / dc + 1))) - c = [FT((n - 1) * dc - cmax) for n in 1:nc] + c = ntuple(n -> FT((n - 1) * dc - cmax), Val(nc)) source_ρ_z_u_v_level = similar(Fields.level(Y.c.ρ, 1), Tuple{FT, FT, FT, FT, FT}) ᶜlevel = similar(Y.c.ρ, FT) @@ -40,7 +46,6 @@ function non_orographic_gravity_wave_cache( gw_source_ampl = Bt_0 .* ones(FT, axes(Fields.level(Y.c.ρ, 1))), gw_Bw = Bw .* ones(FT, axes(Fields.level(Y.c.ρ, 1))), gw_Bn = Bn .* ones(FT, axes(Fields.level(Y.c.ρ, 1))), - gw_B0 = similar(c), gw_c = c, gw_dc = dc, gw_cmax = cmax, @@ -74,7 +79,8 @@ function non_orographic_gravity_wave_cache( (; ϕ0_s, ϕ0_n, dϕ_n, dϕ_s, dc, cmax, c0, nk, cw, cw_tropics, cn) = gw nc = Int(floor(FT(2 * cmax / dc + 1))) - c = [FT((n - 1) * dc - cmax) for n in 1:nc] + c = ntuple(n -> FT((n - 1) * dc - cmax), Val(nc)) + # c = [FT((n - 1) * dc - cmax) for n in 1:nc] ᶜlocal_geometry = Fields.local_geometry_field(Fields.level(Y.c, 1)) lat = ᶜlocal_geometry.coordinates.lat @@ -121,7 +127,6 @@ function non_orographic_gravity_wave_cache( gw_source_ampl = source_ampl, gw_Bw = gw_Bw, gw_Bn = gw_Bn, - gw_B0 = similar(c), gw_c = c, gw_cw = gw_cw, gw_cn = gw_cn, @@ -169,13 +174,6 @@ function non_orographic_gravity_wave_tendency!( ) = p.non_orographic_gravity_wave (; model_config) = p.atmos - if model_config isa SingleColumnModel - (; gw_source_height, source_ρ_z_u_v_level) = - p.non_orographic_gravity_wave - elseif model_config isa SphericalModel - (; gw_source_pressure, gw_damp_pressure, source_p_ρ_z_u_v_level) = - p.non_orographic_gravity_wave - end ᶜρ = Y.c.ρ ᶜz = Fields.coordinate_field(Y.c).z FT = Spaces.undertype(axes(Y.c)) @@ -202,6 +200,8 @@ function non_orographic_gravity_wave_tendency!( if model_config isa SingleColumnModel # source level: the index of the level that is closest to the source height + (; gw_source_height, source_ρ_z_u_v_level) = + p.non_orographic_gravity_wave input = Base.Broadcast.broadcasted(tuple, ᶜρ, ᶜz, ᶜu, ᶜv, ᶜlevel) Operators.column_reduce!( @@ -224,6 +224,8 @@ function non_orographic_gravity_wave_tendency!( elseif model_config isa SphericalModel (; ᶜp) = p.precomputed + (; gw_source_pressure, gw_damp_pressure, source_p_ρ_z_u_v_level) = + p.non_orographic_gravity_wave # source level: the index of the highest level whose pressure is higher than source pressure input = Base.Broadcast.broadcasted(tuple, ᶜp, ᶜρ, ᶜz, ᶜu, ᶜv, ᶜlevel) @@ -327,7 +329,6 @@ function non_orographic_gravity_wave_forcing( ᶜu_p1 = p.scratch.ᶜtemp_scalar_3 ᶜv_p1 = p.scratch.ᶜtemp_scalar_4 ᶜbf_p1 = p.scratch.ᶜtemp_scalar_5 - nci = get_nc(gw_ncval) FT = eltype(ᶜρ) ρ_endlevel = Fields.level(ᶜρ, Spaces.nlevels(axes(ᶜρ))) ρ_endlevel_m1 = Fields.level(ᶜρ, Spaces.nlevels(axes(ᶜρ)) - 1) @@ -368,9 +369,10 @@ function non_orographic_gravity_wave_forcing( ) field_shiftlevel_up!(ᶜz, ᶜz_p1, Boundary_value) - mask = BitVector(ones(nc)) + mask_u = StaticBitVector{nc}(_ -> true) + mask_v = StaticBitVector{nc}(_ -> true) level_end = Spaces.nlevels(axes(ᶜρ)) - B1 = ntuple(i -> 0.0, Val(nc)) + c = gw_c for ink in 1:gw_nk @@ -389,13 +391,11 @@ function non_orographic_gravity_wave_forcing( gw_Bn, gw_cw, gw_cn, - gw_c, gw_flag, ᶜlevel, gw_source_ampl, ) - input_v = Base.Broadcast.broadcasted( tuple, ᶜv_p1, @@ -411,7 +411,6 @@ function non_orographic_gravity_wave_forcing( gw_Bn, gw_cw, gw_cn, - gw_c, gw_flag, ᶜlevel, gw_source_ampl, @@ -420,9 +419,9 @@ function non_orographic_gravity_wave_forcing( Operators.column_accumulate!( u_waveforcing, input_u; - init = (FT(0.0), mask, 0.0, B1), + init = (FT(0.0), mask_u, FT(NaN), ntuple(i -> FT(NaN), Val(nc))), transform = first, - ) do (wave_forcing, mask, Bsum, B0), + ) do (wave_forcing_u, mask_u, Bsum_or_NaN_u, B0_or_NaNs_u), ( u_kp1, u_source, @@ -437,74 +436,88 @@ function non_orographic_gravity_wave_forcing( Bn, cw, cn, - c, flag, level, source_ampl, ) - if level >= (source_level - 1) - FT1 = typeof(u_kp1) - kwv = 2.0 * π / ((30.0 * (10.0^ink)) * 1.e3) - k2 = kwv * kwv - fac = FT1(0.5) * (ρ_kp1 / ρ_source) * kwv / bf_kp1 - Hb = (z_kp1 - z_k) / log(ρ_k / ρ_kp1) # density scale height - alp2 = FT1(0.25) / (Hb * Hb) - ω_r = sqrt((bf_kp1 * bf_kp1 * k2) / (k2 + alp2)) # omc: (critical frequency that marks total internal reflection) - fm = 0.0 - if level == (source_level - 1) - mask .= 1 - Bsum = 0.0 - B0 = ntuple( - n -> - sign(c[n] - u_source) * ( - Bw * exp( - -log(2.0) * + + FT1 = typeof(u_kp1) + kwv = 2.0 * π / ((30.0 * (10.0^ink)) * 1.e3) + k2 = kwv * kwv + + # loop over all wave lengths + + + # here ᶜu has one additional level above model top + fac = FT1(0.5) * (ρ_kp1 / ρ_source) * kwv / bf_kp1 + Hb = (z_kp1 - z_k) / log(ρ_k / ρ_kp1) # density scale height + alp2 = FT1(0.25) / (Hb * Hb) + ω_r = sqrt((bf_kp1 * bf_kp1 * k2) / (k2 + alp2)) # omc: (critical frequency that marks total internal reflection) + + fm = FT1(0.0) + B0_u, Bsum_u = if level == 1 + mask_u = StaticBitVector{nc}(_ -> true) + B1 = ntuple( + n -> + sign((c[n] - u_source)) * ( + Bw * exp( + -log(2.0f0) * + ( ( - ( - c[n] * flag + - (c[n] - u_source) * (1 - flag) - - gw_c0 - ) / cw - )^2, - ) + - Bn * exp( - -log(2.0) * + c[n] * flag + + (c[n] - u_source) * (1 - flag) - + gw_c0 + ) / cw + )^2, + ) + + Bn * exp( + -log(2.0f0) * + ( ( - ( - c[n] * flag + - (c[n] - u_source) * (1 - flag) - - gw_c0 - ) / cn - )^2, - ) - ), - Val(nc), - ) - Bsum = sum(abs.(B0)) - end - for n in 1:nci - # check only those waves which are still propagating, i.e., mask = 1.0 - if (mask[n]) == 1 + c[n] * flag + + (c[n] - u_source) * (1 - flag) - + gw_c0 + ) / cn + )^2, + ) + ), + Val(nc), + ) + Bsum1 = sum(abs, B1) + B1, Bsum1 + else + B0_or_NaNs_u, Bsum_or_NaN_u + end + + if level >= source_level - 1 + (mask_u, fm) = unrolled_reduce( + StaticOneTo{nc}(), + (mask_u, fm), + ) do (mask, fm), (n) + if (mask[n]) == true c_hat = c[n] - u_kp1 # c0mu # f phase speed matches the wind speed, remove c(n) from the set of propagating waves. if c_hat == 0.0 - mask[n] = 0 + mask = Base.setindex(mask, false, n) + #mask = unrolled_setindex(mask,false,n) else c_hat0 = c[n] - u_source # define the criterion which determines if wave is reflected at this level (test). test = abs(c_hat) * kwv - ω_r if test >= 0.0 # wave has undergone total internal reflection. remove it from the propagating set. - mask[n] = 0 + mask = Base.setindex(mask, false, n) + #mask = unrolled_setindex(mask,false,n) else if level == level_end # this is added in MiMA implementation: # all momentum flux that escapes across the model top # is deposited to the extra level being added so that # momentum flux is conserved - mask[n] = 0 + mask = Base.setindex(mask, false, n) + #mask = unrolled_setindex(mask,false,n) if level >= source_level - fm = fm + B0[n] + fm = fm + B0_u[n] end else # if wave is not reflected at this level, determine if it is @@ -515,46 +528,49 @@ function non_orographic_gravity_wave_forcing( # this level. # set mask=0.0 to remove phase speed band c[n] from the set of active # waves moving upwards to the next level. - Foc = B0[n] / FT1((c_hat)^3) - fac + Foc = B0_u[n] / (c_hat)^3 - fac if Foc >= 0.0 || (c_hat0 * c_hat <= 0.0) - mask[n] = 0 + mask = Base.setindex(mask, false, n) if level >= source_level - fm = fm + B0[n] + fm = fm + B0_u[n] end end end end # (test >= 0.0) + end #(c_hat == 0.0) end # mask = 0 - - end # nc: phase speed loop + return (mask, fm) + end # compute the gravity wave momentum flux forcing # obtained across the entire wave spectrum at this level. - eps = calc_intermitency(ρ_source, source_ampl, gw_nk, FT1(Bsum)) + eps = + calc_intermitency(ρ_source, source_ampl, gw_nk, FT1(Bsum_u)) if level >= source_level rbh = sqrt(ρ_k * ρ_kp1) - wave_forcing = + wave_forcing_u = (ρ_source / rbh) * FT1(fm) * eps / (z_kp1 - z_k) else - wave_forcing = FT1(0.0) + wave_forcing_u = FT1(0.0) end - end - return (wave_forcing, mask, Bsum, B0) - end + fm = FT1(0.0) + end + return (wave_forcing_u, mask_u, Bsum_u, B0_u) + end Operators.column_accumulate!( v_waveforcing, input_v; - init = (FT(0.0), mask, 0.0, B1), + init = (FT(0.0), mask_u, FT(NaN), ntuple(i -> FT(NaN), Val(nc))), transform = first, - ) do (wave_forcing, mask, Bsum, B0), + ) do (wave_forcing_v, mask_v, Bsum_or_NaN_v, B0_or_NaNs_v), ( - u_kp1, - u_source, + v_kp1, + v_source, bf_kp1, ρ_k, ρ_kp1, @@ -566,76 +582,85 @@ function non_orographic_gravity_wave_forcing( Bn, cw, cn, - c, flag, level, source_ampl, ) - if level >= (source_level - 1) - FT2 = typeof(u_kp1) - kwv = 2.0 * π / ((30.0 * (10.0^ink)) * 1.e3) - k2 = kwv * kwv - fac = FT2(0.5) * (ρ_kp1 / ρ_source) * kwv / bf_kp1 - Hb = (z_kp1 - z_k) / log(ρ_k / ρ_kp1) # density scale height - alp2 = FT2(0.25) / (Hb * Hb) - ω_r = sqrt((bf_kp1 * bf_kp1 * k2) / (k2 + alp2)) # omc: (critical frequency that marks total internal reflection) - - fm = 0.0 - if level == (source_level - 1) - mask .= 1 - Bsum = 0.0 - B0 = ntuple( - n -> - sign((c[n] - u_source)) * ( - Bw * exp( - -log(2.0) * + FT1 = typeof(v_kp1) + kwv = 2.0 * π / ((30.0 * (10.0^ink)) * 1.e3) + k2 = kwv * kwv + + # loop over all wave lengths + + # here ᶜu has one additional level above model top + fac = FT1(0.5) * (ρ_kp1 / ρ_source) * kwv / bf_kp1 + Hb = (z_kp1 - z_k) / log(ρ_k / ρ_kp1) # density scale height + alp2 = FT1(0.25) / (Hb * Hb) + ω_r = sqrt((bf_kp1 * bf_kp1 * k2) / (k2 + alp2)) # omc: (critical frequency that marks total internal reflection) + + B0_v, Bsum_v = if level == 1 + mask_v = StaticBitVector{nc}(_ -> true) + B1 = ntuple( + n -> + sign((c[n] - v_source)) * ( + Bw * exp( + -log(2.0f0) * + ( ( - ( - c[n] * flag + - (c[n] - u_source) * (1 - flag) - - gw_c0 - ) / cw - )^2, - ) + - Bn * exp( - -log(2.0) * + c[n] * flag + + (c[n] - v_source) * (1 - flag) - + gw_c0 + ) / cw + )^2, + ) + + Bn * exp( + -log(2.0f0) * + ( ( - ( - c[n] * flag + - (c[n] - u_source) * (1 - flag) - - gw_c0 - ) / cn - )^2, - ) - ), - Val(nc), - ) - Bsum = sum(abs.(B0)) - end - for n in 1:nci - # check only those waves which are still propagating, i.e., mask = 1.0 - if (mask[n]) == 1 - c_hat = c[n] - u_kp1 # c0mu + c[n] * flag + + (c[n] - v_source) * (1 - flag) - + gw_c0 + ) / cn + )^2, + ) + ), + Val(nc), + ) + Bsum1 = sum(abs, B1) + B1, Bsum1 + else + B0_or_NaNs_v, Bsum_or_NaN_v + end + + if level >= source_level - 1 + fm = FT1(0.0) + (mask_v, fm) = unrolled_reduce( + StaticOneTo{nc}(), + (mask_v, fm), + ) do (mask, fm), (n) + if (mask[n]) == true + c_hat = c[n] - v_kp1 # c0mu # f phase speed matches the wind speed, remove c(n) from the set of propagating waves. if c_hat == 0.0 - mask[n] = 0 + mask = Base.setindex(mask, false, n) + #mask = unrolled_setindex(mask,false,n) else - c_hat0 = c[n] - u_source + c_hat0 = c[n] - v_source # define the criterion which determines if wave is reflected at this level (test). test = abs(c_hat) * kwv - ω_r if test >= 0.0 # wave has undergone total internal reflection. remove it from the propagating set. - mask[n] = 0 + mask = Base.setindex(mask, false, n) else if level == level_end # this is added in MiMA implementation: # all momentum flux that escapes across the model top # is deposited to the extra level being added so that # momentum flux is conserved - mask[n] = 0 + mask = Base.setindex(mask, false, n) if level >= source_level - fm = fm + B0[n] + fm = fm + B0_v[n] end else # if wave is not reflected at this level, determine if it is @@ -646,35 +671,39 @@ function non_orographic_gravity_wave_forcing( # this level. # set mask=0.0 to remove phase speed band c[n] from the set of active # waves moving upwards to the next level. - Foc = B0[n] / FT2((c_hat)^3) - fac + Foc = B0_v[n] / (c_hat)^3 - fac if Foc >= 0.0 || (c_hat0 * c_hat <= 0.0) - mask[n] = 0 + mask = Base.setindex(mask, false, n) if level >= source_level - fm = fm + B0[n] + fm = fm + B0_v[n] end end end end # (test >= 0.0) + end #(c_hat == 0.0) end # mask = 0 - - end # nc: phase speed loop + return (mask, fm) + end # compute the gravity wave momentum flux forcing # obtained across the entire wave spectrum at this level. - eps = calc_intermitency(ρ_source, source_ampl, gw_nk, FT2(Bsum)) + eps = + calc_intermitency(ρ_source, source_ampl, gw_nk, FT1(Bsum_v)) if level >= source_level rbh = sqrt(ρ_k * ρ_kp1) - wave_forcing = - (ρ_source / rbh) * FT2(fm) * eps / (z_kp1 - z_k) + wave_forcing_v = + (ρ_source / rbh) * FT1(fm) * eps / (z_kp1 - z_k) else - wave_forcing = FT2(0.0) + wave_forcing_v = FT1(0.0) end + end - return (wave_forcing, mask, Bsum, B0) + return (wave_forcing_v, mask_v, Bsum_v, B0_v) end + u_waveforcing_top = p.scratch.temp_field_level copyto!( Fields.field_values(u_waveforcing_top), @@ -690,9 +719,6 @@ function non_orographic_gravity_wave_forcing( 0, ) - # v_waveforcing_top = similar( - # Fields.level(v_waveforcing, Spaces.nlevels(axes(v_waveforcing))), - # ) v_waveforcing_top = p.scratch.temp_field_level copyto!( Fields.field_values(v_waveforcing_top), @@ -732,7 +758,6 @@ function non_orographic_gravity_wave_forcing( end return nothing end - # calculate the intermittency factor eps -> assuming constant Δc. function calc_intermitency(ρ_source_level, source_ampl, nk, Bsum) @@ -740,10 +765,10 @@ function calc_intermitency(ρ_source_level, source_ampl, nk, Bsum) end function gw_average!(wave_forcing, wave_forcing_m1) - L1 = Operators.LeftBiasedC2F(; bottom = Operators.SetValue(0.0)) + L1 = Operators.LeftBiasedC2F(; bottom = Operators.SetValue(0.0f0)) L2 = Operators.LeftBiasedF2C(;) wave_forcing_m1 .= L2.(L1.(wave_forcing)) - @. wave_forcing = 0.5 * (wave_forcing + wave_forcing_m1) + @. wave_forcing = 0.5f0 * (wave_forcing + wave_forcing_m1) end @@ -764,3 +789,10 @@ function field_shiftlevel_up!(ᶜexample_field, ᶜshifted_field, Boundary_value R2 = Operators.RightBiasedF2C(;) ᶜshifted_field .= R2.(R1.(ᶜexample_field)) end + +function first_two_elements(t::Tuple) + if length(t) < 2 + error("Tuple must have at least two elements") + end + return (t[1], t[2]) +end diff --git a/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_3d.jl b/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_3d.jl index bee52d3ccd0..3730fb39235 100644 --- a/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_3d.jl +++ b/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_3d.jl @@ -153,7 +153,7 @@ gw_ncval = Val(500) ᶜv = copy(ᶜz) ᶜbf = copy(ᶜz) ᶜlevel = similar(ᶜρ, FT) -u_waveforcing = similar(ᶜu) +u_waveforcing = similar(ᶜv) v_waveforcing = similar(ᶜv) for i in 1:Spaces.nlevels(axes(ᶜρ)) fill!(Fields.level(ᶜlevel, i), i) @@ -206,8 +206,7 @@ for j in 1:length(lat) uforcing, vforcing, gw_ncval, - u_waveforcing, - v_waveforcing, + waveforcing, params, ) Jan_uforcing[j, :] = parent(uforcing) diff --git a/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_mima.jl b/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_mima.jl index 24f39cbd9c5..9fec90aed0c 100644 --- a/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_mima.jl +++ b/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_mima.jl @@ -162,8 +162,9 @@ gw_ncval = Val(333) ᶜv = copy(ᶜz) ᶜbf = copy(ᶜz) ᶜlevel = similar(ᶜρ, FT) +# waveforcing = similar(ᶜu, Tuple{FT, FT}) u_waveforcing = similar(ᶜu) -v_waveforcing = similar(ᶜv) +v_waveforcing = similar(ᶜu) for i in 1:Spaces.nlevels(axes(ᶜρ)) fill!(Fields.level(ᶜlevel, i), i) end diff --git a/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_single_column.jl b/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_single_column.jl index b419b743da1..b33a9e0e366 100644 --- a/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_single_column.jl +++ b/test/parameterized_tendencies/gravity_wave/non_orographic_gravity_wave/nogw_test_single_column.jl @@ -11,6 +11,8 @@ import ClimaCore.Spaces as Spaces import ClimaCore.Fields as Fields import ClimaCore.Geometry as Geometry import ClimaCore.Operators as Operators +import ClimaCore.Domains as Domains +import ClimaCore: InputOutput, Meshes, Spaces, Quadratures include("../gw_plotutils.jl") @@ -135,19 +137,73 @@ center_ρ_mean = mean(center_ρ, dims = 1)[1, :, :] # monthly ave Jan, April, July, Oct -column_domain = ClimaCore.Domains.IntervalDomain( - ClimaCore.Geometry.ZPoint(0.0) .. ClimaCore.Geometry.ZPoint(50000.0), +# column_domain = ClimaCore.Domains.IntervalDomain( +# ClimaCore.Geometry.ZPoint(0.0) .. ClimaCore.Geometry.ZPoint(50000.0), +# boundary_names = (:bottom, :top), +# ) + +# column_mesh = ClimaCore.Meshes.IntervalMesh(column_domain, nelems = 50) + +# column_center_space = ClimaCore.Spaces.CenterFiniteDifferenceSpace(column_mesh) +# # construct the face space from the center one +# column_face_space = +# ClimaCore.Spaces.FaceFiniteDifferenceSpace(column_center_space) + +# coord = ClimaCore.Fields.coordinate_field(column_center_space) + +Δx = FT(1) # Note: This value shouldn't matter, since we only have 1 column. +quad = Quadratures.GL{1}() + +x_domain = Domains.IntervalDomain( + Geometry.XPoint(zero(Δx)), + Geometry.XPoint(Δx); + periodic = true, +) +y_domain = Domains.IntervalDomain( + Geometry.YPoint(zero(Δx)), + Geometry.YPoint(Δx); + periodic = true, +) +domain = Domains.RectangleDomain(x_domain, y_domain) +horizontal_mesh = Meshes.RectilinearMesh(domain, 1, 1) + +comms_ctx = ClimaComms.SingletonCommsContext{ClimaComms.CPUSingleThreaded}( + ClimaComms.CPUSingleThreaded(), +) +topology = ClimaCore.Topologies.Topology2D( + comms_ctx, + horizontal_mesh, + ClimaCore.Topologies.spacefillingcurve(horizontal_mesh), +) +h_space = Spaces.SpectralElementSpace2D( + topology, + quad; + # enable_bubble = true, +) + +h_grid = Spaces.grid(h_space) +z_domain = Domains.IntervalDomain( + Geometry.ZPoint(0.0), + Geometry.ZPoint(50000.0); boundary_names = (:bottom, :top), ) +z_stretch = Meshes.Uniform() +z_mesh = Meshes.IntervalMesh(z_domain, z_stretch; nelems = 50) -column_mesh = ClimaCore.Meshes.IntervalMesh(column_domain, nelems = 50) +device = ClimaComms.device(h_space) +z_topology = ClimaCore.Topologies.IntervalTopology( + ClimaComms.SingletonCommsContext(device), + z_mesh, +) +z_grid = ClimaCore.Grids.FiniteDifferenceGrid(z_topology) +hypsography = ClimaCore.Hypsography.Flat() +grid = + ClimaCore.Grids.ExtrudedFiniteDifferenceGrid(h_grid, z_grid, hypsography;) -column_center_space = ClimaCore.Spaces.CenterFiniteDifferenceSpace(column_mesh) -# construct the face space from the center one -column_face_space = - ClimaCore.Spaces.FaceFiniteDifferenceSpace(column_center_space) +center_space = Spaces.CenterExtrudedFiniteDifferenceSpace(grid) +face_space = Spaces.FaceExtrudedFiniteDifferenceSpace(grid) -coord = ClimaCore.Fields.coordinate_field(column_center_space) +coord = ClimaCore.Fields.coordinate_field(center_space) gw_ncval = Val(501) ᶜz = coord.z @@ -156,8 +212,9 @@ gw_ncval = Val(501) ᶜv = copy(ᶜz) ᶜbf = copy(ᶜz) ᶜlevel = similar(ᶜρ, FT) +# waveforcing = similar(ᶜu, Tuple{FT, FT}) u_waveforcing = similar(ᶜu) -v_waveforcing = similar(ᶜv) +v_waveforcing = similar(ᶜu) for i in 1:Spaces.nlevels(axes(ᶜρ)) fill!(Fields.level(ᶜlevel, i), i) end @@ -200,6 +257,10 @@ Jan_uforcing = similar(ᶜρ, FT) Jan_uforcing .= 0 Jan_vforcing = similar(ᶜρ, FT) Jan_vforcing .= 0 +source_leveli = copy(ᶜρ_source) +source_leveli .= 10 +damp_leveli = copy(ᶜρ_source) +damp_leveli .= 50 CA.non_orographic_gravity_wave_forcing( ᶜu, @@ -208,8 +269,8 @@ CA.non_orographic_gravity_wave_forcing( ᶜρ, ᶜz, ᶜlevel, - source_level, - damp_level, + source_leveli, + damp_leveli, ᶜρ_source, ᶜu_source, ᶜv_source,