diff --git a/integration_tests/utils/main.jl b/integration_tests/utils/main.jl index 0e4f14239a..955b3fe8f0 100644 --- a/integration_tests/utils/main.jl +++ b/integration_tests/utils/main.jl @@ -246,36 +246,64 @@ function run(sim::Simulation1d) grid = sim.grid state = sim.state TC.open_files(sim.Stats) # #removeVarsHack - while sim.TS.t <= sim.TS.t_max - TC.update(sim.Turb, grid, state, sim.GMV, sim.Case, sim.TS) - TC.update(sim.TS) - - if mod(iter, 100) == 0 - progress = sim.TS.t / sim.TS.t_max - @show progress - end - - if mod(round(Int, sim.TS.t), round(Int, sim.Stats.frequency)) == 0 - # TODO: is this the best location to call diagnostics? - TC.compute_diagnostics!(sim.Turb, sim.GMV, grid, state, sim.Case, sim.TS) - - # TODO: remove `vars` hack that avoids - # https://github.com/Alexander-Barth/NCDatasets.jl/issues/135 - # opening/closing files every step should be okay. #removeVarsHack - # TurbulenceConvection.io(sim) # #removeVarsHack - TC.write_simulation_time(sim.Stats, sim.TS.t) # #removeVarsHack - - TC.io(sim.io_nt.aux, sim.Stats) - TC.io(sim.io_nt.diagnostics, sim.Stats) - TC.io(sim.io_nt.prog, sim.Stats) - TC.io(sim.io_nt.tendencies, sim.Stats) - - TC.io(sim.GMV, grid, state, sim.Stats) # #removeVarsHack - TC.io(sim.Case, grid, state, sim.Stats) # #removeVarsHack - TC.io(sim.Turb, grid, state, sim.Stats, sim.TS, sim.param_set) # #removeVarsHack - end - iter += 1 + + t_span = (0.0, sim.TS.t_max) + params = (; + edmf = sim.Turb, + grid = grid, + state = state, + gm = sim.GMV, + Case = sim.Case, + TS = sim.TS, + Stats = sim.Stats, + io_nt = sim.io_nt, + ) + + condition_io(u, t, integrator) = mod(round(Int, t), round(Int, sim.Stats.frequency)) == 0 + condition_every_iter(u, t, integrator) = true + + function affect_io!(integrator) + UnPack.@unpack edmf, io_nt, grid, state, Case, gm, TS, Stats = integrator.p + parent(state.prog) .= integrator.u + param_set = TC.parameter_set(gm) + # TODO: is this the best location to call diagnostics? + TC.compute_diagnostics!(edmf, gm, grid, state, Case, TS) + + # TODO: remove `vars` hack that avoids + # https://github.com/Alexander-Barth/NCDatasets.jl/issues/135 + # opening/closing files every step should be okay. #removeVarsHack + # TurbulenceConvection.io(sim) # #removeVarsHack + TC.write_simulation_time(Stats, TS.t) # #removeVarsHack + + TC.io(io_nt.aux, Stats) + TC.io(io_nt.diagnostics, Stats) + TC.io(io_nt.prog, Stats) + TC.io(io_nt.tendencies, Stats) + + TC.io(gm, grid, state, Stats) # #removeVarsHack + TC.io(Case, grid, state, Stats) # #removeVarsHack + TC.io(edmf, grid, state, Stats, TS, param_set) # #removeVarsHack + end + + function affect_filter!(integrator) + UnPack.@unpack edmf, grid, state, gm, Case, TS = integrator.p + TC.affect_filter!(integrator.u, edmf, grid, state, gm, Case, TS) end + + callback_io = ODE.DiscreteCallback(condition_io, affect_io!) + callback_filters = ODE.DiscreteCallback(condition_every_iter, affect_filter!) + + prob = ODE.ODEProblem(TC.step!, state.prog, t_span, params; dt = sim.TS.dt) + sol = ODE.solve( + prob, + ODE.Euler(); + reltol = 1e-12, + abstol = 1e-12, + callback = ODE.CallbackSet(callback_filters, callback_io), + progress = true, + progress_message = (dt, u, p, t) -> t, + ) + TC.close_files(sim.Stats) # #removeVarsHack return end diff --git a/src/Fields.jl b/src/Fields.jl index ac6a552533..bc5f4458e6 100644 --- a/src/Fields.jl +++ b/src/Fields.jl @@ -89,3 +89,16 @@ function FieldFromNamedTuple(space, nt::NamedTuple) cmv(z) = nt return cmv.(CC.Fields.coordinate_field(space)) end + +# https://github.com/CliMA/ClimaCore.jl/issues/275 +transform_broadcasted(bc::Base.Broadcast.Broadcasted{CC.Fields.FieldVectorStyle}, symb, axes) = + Base.Broadcast.Broadcasted(bc.f, map(arg -> transform_broadcasted(arg, symb, axes), bc.args), axes) +transform_broadcasted(fv::CC.Fields.FieldVector, symb, axes) = parent(getproperty(fv, symb)) +transform_broadcasted(x, symb, axes) = x +@inline function Base.copyto!(dest::CC.Fields.FieldVector, bc::Base.Broadcast.Broadcasted{CC.Fields.FieldVectorStyle}) + for symb in propertynames(dest) + p = parent(getproperty(dest, symb)) + Base.copyto!(p, transform_broadcasted(bc, symb, axes(p))) + end + return dest +end diff --git a/src/Turbulence_PrognosticTKE.jl b/src/Turbulence_PrognosticTKE.jl index 6c3166d515..5185235799 100755 --- a/src/Turbulence_PrognosticTKE.jl +++ b/src/Turbulence_PrognosticTKE.jl @@ -230,8 +230,36 @@ function compute_diffusive_fluxes( return end +function affect_filter!(u, edmf, grid, state, gm, Case, TS) + # TODO: figure out why this filter kills the DryBubble results if called at t = 0. + if Case.casename == "DryBubble" && !(TS.t > 0) + return nothing + end + parent(state.prog) .= parent(u) + prog_en = center_prog_environment(state) + up = edmf.UpdVar + ### + ### Filters + ### + set_edmf_surface_bc(edmf, grid, state, up, Case.Sur) + filter_updraft_vars(edmf, grid, state, gm) + + @inbounds for k in real_center_indices(grid) + prog_en.ρatke[k] = max(prog_en.ρatke[k], 0.0) + prog_en.ρaHvar[k] = max(prog_en.ρaHvar[k], 0.0) + prog_en.ρaQTvar[k] = max(prog_en.ρaQTvar[k], 0.0) + prog_en.ρaHQTcov[k] = max(prog_en.ρaHQTcov[k], -sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k])) + prog_en.ρaHQTcov[k] = min(prog_en.ρaHQTcov[k], sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k])) + end + parent(u) .= parent(state.prog) +end + # Perform the update of the scheme -function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Case::CasesBase, TS::TimeStepping) +function step!(dY, Y, params, t) + UnPack.@unpack edmf, grid, state, gm, Case, TS = params + TS.t = t + affect_filter!(Y, edmf, grid, state, gm, Case, TS) + parent(state.prog) .= parent(Y) gm = gm up = edmf.UpdVar @@ -277,53 +305,7 @@ function update(edmf::EDMF_PrognosticTKE, grid, state, gm::GridMeanVariables, Ca compute_en_tendencies!(edmf, grid, state, param_set, TS, :QTvar, :ρaQTvar, n_updrafts) compute_en_tendencies!(edmf, grid, state, param_set, TS, :HQTcov, :ρaHQTcov, n_updrafts) - ### - ### update (to be removed) - ### - - Δt = TS.dt - prog_up = center_prog_updrafts(state) - prog_up_f = face_prog_updrafts(state) - prog_pr = center_prog_precipitation(state) - tendencies_pr = center_tendencies_precipitation(state) - - @inbounds for k in real_center_indices(grid) - prog_en.ρatke[k] += TS.dt * tendencies_en.ρatke[k] - prog_en.ρaHvar[k] += TS.dt * tendencies_en.ρaHvar[k] - prog_en.ρaQTvar[k] += TS.dt * tendencies_en.ρaQTvar[k] - prog_en.ρaHQTcov[k] += TS.dt * tendencies_en.ρaHQTcov[k] - prog_gm.u[k] += tendencies_gm.u[k] * TS.dt - prog_gm.v[k] += tendencies_gm.v[k] * TS.dt - prog_gm.θ_liq_ice[k] += tendencies_gm.θ_liq_ice[k] * TS.dt - prog_gm.q_tot[k] += tendencies_gm.q_tot[k] * TS.dt - @inbounds for i in 1:(up.n_updrafts) - prog_up[i].ρarea[k] += Δt * tendencies_up[i].ρarea[k] - prog_up[i].ρaθ_liq_ice[k] += Δt * tendencies_up[i].ρaθ_liq_ice[k] - prog_up[i].ρaq_tot[k] += Δt * tendencies_up[i].ρaq_tot[k] - end - if has_precip - prog_pr.qr[k] += tendencies_pr.qr[k] * TS.dt - end - end - - @inbounds for k in real_face_indices(grid) - @inbounds for i in 1:(up.n_updrafts) - prog_up_f[i].ρaw[k] += Δt * tendencies_up_f[i].ρaw[k] - end - end - - ### - ### Filters - ### - set_edmf_surface_bc(edmf, grid, state, up, Case.Sur) - filter_updraft_vars(edmf, grid, state, gm) - @inbounds for k in real_center_indices(grid) - prog_en.ρatke[k] = max(prog_en.ρatke[k], 0.0) - prog_en.ρaHvar[k] = max(prog_en.ρaHvar[k], 0.0) - prog_en.ρaQTvar[k] = max(prog_en.ρaQTvar[k], 0.0) - prog_en.ρaHQTcov[k] = max(prog_en.ρaHQTcov[k], -sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k])) - prog_en.ρaHQTcov[k] = min(prog_en.ρaHQTcov[k], sqrt(prog_en.ρaHvar[k] * prog_en.ρaQTvar[k])) - end + parent(dY) .= parent(state.tendencies) return end