Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make swap_spaces not allocating and more GPU compatible #249

Merged
merged 1 commit into from
Feb 23, 2023
Merged

Make swap_spaces not allocating and more GPU compatible #249

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions experiments/AMIP/moist_mpi_aquaplanet/coupler_driver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ include("coupler_utils/conservation_checker.jl")
include("coupler_utils/regridder.jl")

# initiate spatial and temporal info
t_end = 2592000 # 100e2
t_end = 2592000 # 100e2
tspan = (0, t_end)
Δt_cpl = 2e2
saveat = Δt_cpl * 1000
Expand Down Expand Up @@ -51,7 +51,7 @@ walltime = @elapsed for t in (tspan[1]:Δt_cpl:tspan[end])
parent(atmos_sim.integrator.p.dif_flux_energy) .= FT(0)
calculate_surface_fluxes_atmos_grid!(atmos_sim.integrator, T_S)

# run
# run
step!(atmos_sim.integrator, t - atmos_sim.integrator.t, true) # NOTE: use (t - integ_atm.t) here instead of Δt_cpl to avoid accumulating roundoff error in our timestepping.

## Slab
Expand Down
36 changes: 18 additions & 18 deletions experiments/AMIP/moist_mpi_earth/coupler_driver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,16 +5,16 @@ include("mpi/mpi_init.jl") # setup MPI context for distributed runs #hide
#=
## Overview

AMIP is a standard experimental protocol of the Program for Climate Model Diagnosis & Intercomparison (PCMDI).
AMIP is a standard experimental protocol of the Program for Climate Model Diagnosis & Intercomparison (PCMDI).
It is used as a model benchmark for the atmospheric and land model components, while sea-surface temperatures (SST) and sea-ice concentration (SIC)
are prescribed using time-interpolations between monthly observed data. We use standard data files with original sources:
- SST and SIC: https://gdex.ucar.edu/dataset/158_asphilli.html
- SST and SIC: https://gdex.ucar.edu/dataset/158_asphilli.html
- land-sea mask: https://www.ncl.ucar.edu/Applications/Data/#cdf

For more information, see the PCMDI's specifications for [AMIP I](https://pcmdi.github.io/mips/amip/) and [AMIP II](https://pcmdi.github.io/mips/amip2/).
For more information, see the PCMDI's specifications for [AMIP I](https://pcmdi.github.io/mips/amip/) and [AMIP II](https://pcmdi.github.io/mips/amip2/).

This driver contains two modes. The full `AMIP` mode and a `SlabPlanet` (all surfaces are thermal slabs) mode. Since `AMIP` is not a closed system, the
`SlabPlanet` mode is useful for checking conservation properties of the coupling.
This driver contains two modes. The full `AMIP` mode and a `SlabPlanet` (all surfaces are thermal slabs) mode. Since `AMIP` is not a closed system, the
`SlabPlanet` mode is useful for checking conservation properties of the coupling.

=#

Expand All @@ -41,7 +41,7 @@ you can run directly from the Julia REPL. The latter is recommended for debuggin
with threading enabled:
```julia
julia --project --threads 8
```
```
=#

#=
Expand Down Expand Up @@ -84,7 +84,7 @@ if isinteractive()
parsed_args["albedo_from_file"] = true #hide
end

## read in some parsed command line arguments
## read in some parsed command line arguments
mode_name = parsed_args["mode_name"]
run_name = parsed_args["run_name"]
energy_check = parsed_args["energy_check"]
Expand Down Expand Up @@ -140,9 +140,9 @@ include("atmos/atmos_init.jl")
atmos_sim = atmos_init(FT, Y, integrator, params = params);

#=
We use a common `Space` for all global surfaces. This enables the MPI processes to operate on the same columns in both
the atmospheric and surface components, so exchanges are parallelized. Note this is only possible when the
atmosphere and surface are of the same horizontal resolution.
We use a common `Space` for all global surfaces. This enables the MPI processes to operate on the same columns in both
the atmospheric and surface components, so exchanges are parallelized. Note this is only possible when the
atmosphere and surface are of the same horizontal resolution.
=#
## init a 2D bounary space at the surface
boundary_space = atmos_sim.domain.face_space.horizontal_space
Expand Down Expand Up @@ -176,9 +176,9 @@ land_sim = bucket_init(
#=
### Ocean and Sea Ice
In the `AMIP` mode, all ocean properties are prescribed from a file, while sea-ice temperatures are calculated using observed
SIC and assuming a 2m thickness of the ice.
SIC and assuming a 2m thickness of the ice.

In the `SlabPlanet` mode, all ocean and sea ice are dynamical models, namely thermal slabs, with different parameters.
In the `SlabPlanet` mode, all ocean and sea ice are dynamical models, namely thermal slabs, with different parameters.
=#

@info mode_name
Expand Down Expand Up @@ -252,7 +252,7 @@ end
#=
## Coupler Initialization
The coupler needs to contain exchange information, manage the calendar and be able to access all component models. It can also optionally
save online diagnostics. These are all initialized here and saved in a global `CouplerSimulation` struct, `cs`.
save online diagnostics. These are all initialized here and saved in a global `CouplerSimulation` struct, `cs`.
=#

## coupler exchange fields
Expand Down Expand Up @@ -305,7 +305,7 @@ cs = CouplerSimulation{FT}(
);

#=
## Initial States Exchange
## Initial States Exchange
=#
## share states between models
include("./push_pull.jl")
Expand Down Expand Up @@ -343,7 +343,7 @@ function solve_coupler!(cs, energy_check)

cs.dates.date[1] = current_date(cs, t) # if not global, `date` is not updated.

## print date on the first of month
## print date on the first of month
@calendar_callback :(@show(cs.dates.date[1])) cs.dates.date[1] cs.dates.date1[1]

if cs.mode.name == "amip"
Expand Down Expand Up @@ -434,8 +434,8 @@ end
solve_coupler!(cs, energy_check);

#=
## Postprocessing
Currently all postprocessing is performed using the root process only.
## Postprocessing
Currently all postprocessing is performed using the root process only.
=#

if ClimaComms.iamroot(comms_ctx)
Expand Down Expand Up @@ -518,7 +518,7 @@ if ClimaComms.iamroot(comms_ctx)
end

#=
## Temporary Unit Tests
## Temporary Unit Tests
To be moved to `test/`
=#
if !is_distributed && cs.mode.name == "amip"
Expand Down
2 changes: 1 addition & 1 deletion experiments/AMIP/moist_mpi_earth/coupler_utils/bcfile_reader.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ function bcfile_info_init(

end

no_scaling(x, _info) = swap_space!(x, axes(_info.land_mask))
no_scaling(x, _info) = swap_space!(zeros(axes(_info.land_mask)), x)

# IO - monthly
"""
Expand Down
7 changes: 4 additions & 3 deletions experiments/AMIP/moist_mpi_earth/coupler_utils/conservation_checker.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,9 @@ function check_conservation(cc::OnlineConservationCheck, coupler_sim)
get_land_energy(land_sim, e_per_area_land)
end

u_ocn = ocean_sim !== nothing ? swap_space!(ocean_sim.integrator.u.T_sfc, coupler_sim.boundary_space) : nothing
u_ice = ice_sim !== nothing ? swap_space!(ice_sim.integrator.u.T_sfc, coupler_sim.boundary_space) : nothing
u_ocn =
ocean_sim !== nothing ? swap_space!(zeros(coupler_sim.boundary_space), ocean_sim.integrator.u.T_sfc) : nothing
u_ice = ice_sim !== nothing ? swap_space!(zeros(coupler_sim.boundary_space), ice_sim.integrator.u.T_sfc) : nothing

# global sums
atmos_e = sum(u_atm)
Expand Down Expand Up @@ -101,7 +102,7 @@ function check_conservation(cc::OnlineConservationCheck, coupler_sim)
push!(cc.ρe_tot_ocean, ocean_e)

# save surface friction sink
push!(cc.friction_sink, sum((ke_dissipation(atmos_sim)) .* coupler_sim.Δt_cpl)) # ρ d ke_friction / dt
push!(cc.friction_sink, sum((ke_dissipation(atmos_sim)) .* coupler_sim.Δt_cpl)) # ρ d ke_friction / dt

end
function ke_dissipation(sim)
Expand Down
12 changes: 6 additions & 6 deletions experiments/AMIP/moist_mpi_earth/coupler_utils/flux_calculator.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,15 @@ import ClimaAtmos: get_surface_fluxes!
"""
set_ρ_sfc!(ρ_sfc, T_S, integrator)

sets the value of the ρ_sfc field based on the temperature of the surface,
sets the value of the ρ_sfc field based on the temperature of the surface,
the temperature of the atmosphere at the lowest level, and the heigh
of the lowest level.
"""
function set_ρ_sfc!(ρ_sfc, T_S, integrator)
ts = integrator.p.ᶜts
thermo_params = CAP.thermodynamics_params(integrator.p.params)
ts_int = Spaces.level(ts, 1)
parent(ρ_sfc) .= parent(ρ_sfc_at_point.(thermo_params, ts_int, swap_space!(T_S, axes(ts_int))))
parent(ρ_sfc) .= parent(ρ_sfc_at_point.(thermo_params, ts_int, swap_space!(zeros(axes(ts_int)), T_S)))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of creating a new field here, can we create it as part of the CouplerSimulation struct? Similar to how we treat boundary_space. @charleskawczynski is this what you meant?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure. Here is the MWE gist where I compare all the different approaches. Creating it as part of the CoupledSimulation struct is fine, but we need to be careful to create a temp copy, wherever we have repeated usages, such as in these conservation checks. I can proceed with the proposed changes (Case 4 in the MWE gist) if you agree that's the best option.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It actually seems like Case 2 in the gist is the best option for our purposes, so I will leave the current version of the code. I am going to squash/rebase and you're good to merge. Thanks for the review!

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Output of the gist at first run:

[ Info: Case 1:
alloc_u_ocn = 19493429
u_ocn = Float64-valued Field:
  [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0  …  2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
alloc_u_ice = 0
u_ice = Float64-valued Field:
  [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0  …  2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
[ Info: Case 2:
alloc_u_ocn = 12800
u_ocn = Float64-valued Field:
  [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0  …  1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
alloc_u_ice = 12800
u_ice = Float64-valued Field:
  [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0  …  2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
[ Info: Case 3:
alloc_u_ocn = 33705
u_ocn = Float64-valued Field:
  [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0  …  2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
alloc_u_ice = 1008
u_ice = Float64-valued Field:
  [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0  …  2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
[ Info: Case 4:
alloc_u_ocn = 465631
u_ocn = Float64-valued Field:
  [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0  …  1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
alloc_u_ice = 1008
u_ice = Float64-valued Field:
  [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0  …  2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]
Float64-valued Field:
  [2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0  …  2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]

end

"""
Expand All @@ -23,7 +23,7 @@ Computes the surface density at a point given the atmospheric state
at the lowest level, the surface temperature, and the assumption of
an ideal gas and hydrostatic balance.

Required because the surface models do not compute air density as a
Required because the surface models do not compute air density as a
variable.
"""
function ρ_sfc_at_point(params, ts_int, T_sfc)
Expand All @@ -37,9 +37,9 @@ end
"""
calculate_surface_fluxes_atmos_grid!(integrator)

Calculates surface fluxes using adapter function `get_surface_fluxes!`
from ClimaAtmos that calls `SurfaceFluxes.jl`. The coupler updates in
atmos model cache fluxes at each coupling timestep.
Calculates surface fluxes using adapter function `get_surface_fluxes!`
from ClimaAtmos that calls `SurfaceFluxes.jl`. The coupler updates in
atmos model cache fluxes at each coupling timestep.

- TODO: generalize interface for regridding and take land state out of atmos's integrator.p
"""
Expand Down
5 changes: 2 additions & 3 deletions experiments/AMIP/moist_mpi_earth/coupler_utils/general_helper.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,8 @@ end
CouplerSimulation{FT}(args...) where {FT} = CouplerSimulation{FT, typeof.(args[1:5])...}(args...)
float_type(::CouplerSimulation{FT}) where {FT} = FT

function swap_space!(field, new_space)
field_out = zeros(new_space)
parent(field_out) .= parent(field)
function swap_space!(field_out, field_in)
parent(field_out) .= parent(field_in)
return field_out
end

Expand Down
2 changes: 1 addition & 1 deletion experiments/AMIP/moist_mpi_earth/coupler_utils/masker.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ function land_sea_mask(
ClimaComms.barrier(comms_ctx)
file_dates = load(joinpath(REGRID_DIR, outfile_root * "_times.jld2"), "times")
mask = hdread_regridfile(comms_ctx, outfile_root, file_dates[1], varname)
mask = swap_space!(mask, boundary_space) # needed if we are reading from previous run
mask = swap_space!(zeros(boundary_space), mask) # needed if we are reading from previous run
return mono ? mask : binary_mask.(mask, threshold = threshold)
end

Expand Down
6 changes: 3 additions & 3 deletions experiments/AMIP/moist_mpi_earth/coupler_utils/variable_definer.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,16 +12,16 @@ get_var(cs, ::Val{:u}) = ClimaCore.Geometry.UVVector.(cs.model_sims.atmos_sim.in
get_var(cs, ::Val{:q_tot}) =
cs.model_sims.atmos_sim.integrator.u.c.ρq_tot ./ cs.model_sims.atmos_sim.integrator.u.c.ρ .* float_type(cs)(1000)

get_var(cs, ::Val{:toa}) = swap_space!(toa_fluxes(cs), cs.boundary_space)
get_var(cs, ::Val{:toa}) = swap_space!(zeros(cs.boundary_space), toa_fluxes(cs))

get_var(cs, ::Val{:precipitation}) =
.-swap_space!(
zeros(cs.boundary_space),
cs.model_sims.atmos_sim.integrator.p.col_integrated_rain .+
cs.model_sims.atmos_sim.integrator.p.col_integrated_snow,
cs.boundary_space,
)

get_var(cs, ::Val{:T_sfc}) = swap_space!(cs.fields.T_S, cs.boundary_space)
get_var(cs, ::Val{:T_sfc}) = swap_space!(zeros(cs.boundary_space), cs.fields.T_S)

# more complex calculations
function zonal_wind(cs)
Expand Down
15 changes: 7 additions & 8 deletions experiments/AMIP/moist_mpi_earth/slab_ice/slab_init.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ struct IceSlabParameters{FT <: AbstractFloat}
T_base::FT
z0m::FT
z0b::FT
T_freeze::FT # temperature at freezing point [K]
T_freeze::FT # temperature at freezing point [K]
k_ice::FT # thermal conductivity of ice [W / m / K]
α::FT # albedo
end
Expand Down Expand Up @@ -38,7 +38,7 @@ end
"""
ice_init(::Type{FT}; tspan, dt, saveat, space, ice_mask, stepper = Euler()) where {FT}

Initializes the `DiffEq` problem, and creates a Simulation-type object containing the necessary information for `step!` in the coupling loop.
Initializes the `DiffEq` problem, and creates a Simulation-type object containing the necessary information for `step!` in the coupling loop.
"""
function ice_init(::Type{FT}; tspan, saveat, dt, space, ice_mask, stepper = Euler()) where {FT}

Expand Down Expand Up @@ -66,17 +66,16 @@ end

# file-specific
"""
clean_sic(SIC, _info)
Ensures that the space of the SIC struct matches that of the mask, and converts the units from area % to area fraction.
clean_sic(SIC, _info)
Ensures that the space of the SIC struct matches that of the mask, and converts the units from area % to area fraction.
"""
clean_sic(SIC, _info) = swap_space!(zeros(axes(_info.land_mask)), SIC) ./ float_type_bcf(_info)(100.0)

clean_sic(SIC, _info) = swap_space!(SIC, axes(_info.land_mask)) ./ float_type_bcf(_info)(100.0)

# setting that SIC < 0.5 os counted as ocean if binary remapping of landsea mask.
# setting that SIC < 0.5 os counted as ocean if binary remapping of landsea mask.
get_ice_mask(h_ice::FT, mono, threshold = 0.5) where {FT} = mono ? h_ice : binary_mask.(h_ice, threshold = threshold)

"""
apply_mask(mask, condition, yes, no, value = 0.5)
apply_mask(mask, condition, yes, no, value = 0.5)
- apply mask mased on a threshold value in the mask
"""
apply_mask(mask, condition, yes, no, value) = condition(mask, value) ? yes : no
6 changes: 3 additions & 3 deletions experiments/AMIP/moist_mpi_earth/slab_ocean/slab_init.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ end
"""
ocean_init(::Type{FT}; tspan, dt, saveat, space, land_mask, stepper = Euler()) where {FT}

Initializes the `DiffEq` problem, and creates a Simulation-type object containing the necessary information for `step!` in the coupling loop.
Initializes the `DiffEq` problem, and creates a Simulation-type object containing the necessary information for `step!` in the coupling loop.
"""
function ocean_init(::Type{FT}; tspan, dt, saveat, space, ocean_mask, stepper = Euler()) where {FT}

Expand All @@ -69,7 +69,7 @@ end

# file specific
"""
clean_sst(SST::FT, _info)
clean_sst(SST::FT, _info)
Ensures that the space of the SST struct matches that of the mask, and converts the units to Kelvin (N.B.: this is dataset specific)
"""
clean_sst(SST, _info) = (swap_space!(SST, axes(_info.land_mask)) .+ float_type_bcf(_info)(273.15))
clean_sst(SST, _info) = (swap_space!(zeros(axes(_info.land_mask)), SST) .+ float_type_bcf(_info)(273.15))
6 changes: 3 additions & 3 deletions experiments/AMIP/moist_mpi_earth/user_diagnostics.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ function get_var(cs::CoupledSimulation, ::Val{:toa})
)

radiation_sources = @. -(LWd_TOA + SWd_TOA - LWu_TOA - SWu_TOA)
swap_space!(radiation_sources, cs.boundary_space)
swap_space!(zeros(cs.boundary_space), radiation_sources)
end

"""
Expand All @@ -70,9 +70,9 @@ Precipitation (m m⁻²).
"""
get_var(cs::CoupledSimulation, ::Val{:precipitation}) =
.-swap_space!(
zeros(cs.boundary_space),
cs.model_sims.atmos_sim.integrator.p.col_integrated_rain .+
cs.model_sims.atmos_sim.integrator.p.col_integrated_snow,
cs.boundary_space,
)

# coupler diagnotics
Expand All @@ -81,6 +81,6 @@ get_var(cs::CoupledSimulation, ::Val{:precipitation}) =

Combined surface temperature (K).
"""
get_var(cs::CoupledSimulation, ::Val{:T_sfc}) = swap_space!(cs.fields.T_S, cs.boundary_space)
get_var(cs::CoupledSimulation, ::Val{:T_sfc}) = swap_space!(zeros(cs.boundary_space), cs.fields.T_S)

# land diagnotics
6 changes: 3 additions & 3 deletions experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_driver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ include("slab_ocean/slab_init.jl")
prescribed_sst = false
if prescribed_sst == true
SST = ncreader_rll_to_cgll_from_space(FT, "data/sst.nc", "SST", boundary_space) # a sample SST field from https://gdex.ucar.edu/dataset/158_asphilli.html
SST = swap_space!(SST, axes(mask)) .* (abs.(mask .- 1)) .+ FT(273.15) # TODO: avoids the "space not the same instance" error
SST = swap_space!(zeros(axes(mask)), SST) .* (abs.(mask .- 1)) .+ FT(273.15) # TODO: avoids the "space not the same instance" error
ocean_params = OceanSlabParameters(FT(20), FT(1500.0), FT(800.0), FT(280.0), FT(1e-3), FT(1e-5))
else
slab_ocean_sim = slab_ocean_init(FT, tspan, dt = Δt_cpl, space = boundary_space, saveat = saveat, mask = mask)
Expand All @@ -60,7 +60,7 @@ prescribed_sic = false
if prescribed_sic == true
# sample SST field
SIC = ncreader_rll_to_cgll_from_space(FT, "data/sic.nc", "SEAICE", boundary_space)
SIC = swap_space!(SIC, axes(mask)) .* (abs.(mask .- 1))
SIC = swap_space!(zeros(axes(mask)), SIC) .* (abs.(mask .- 1))
slab_ice_sim = slab_ice_init(
FT,
tspan,
Expand Down Expand Up @@ -179,7 +179,7 @@ walltime = @elapsed for t in (tspan[1]:Δt_cpl:tspan[end])

atmos_sim.integrator.p.rrtmgp_model.surface_temperature .= field2array(T_S) # supplied to atmos for radiation

# run
# run
step!(atmos_sim.integrator, t - atmos_sim.integrator.t, true) # NOTE: instead of Δt_cpl, to avoid accumulating roundoff error

#clip TODO: this is bad!! > limiters
Expand Down
17 changes: 12 additions & 5 deletions experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_utils/conservation_checker.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,14 @@ function check_conservation(
z = parent(Fields.coordinate_field(face_space).z)
Δz_bot = FT(0.5) * (z[2, 1, 1, 1, 1] - z[1, 1, 1, 1, 1])

boundary_space_field = zeros(coupler_sim.boundary_space)

u_atm = atmos_sim !== nothing ? atmos_sim.integrator.u.c.ρe : nothing
u_lnd = land_sim !== nothing ? swap_space!(land_sim.integrator.u.T_sfc, coupler_sim.boundary_space) : nothing
u_ocn = ocean_sim !== nothing ? swap_space!(ocean_sim.integrator.u.T_sfc, coupler_sim.boundary_space) : nothing
u_ice = seaice_sim !== nothing ? swap_space!(seaice_sim.integrator.u.T_sfc, coupler_sim.boundary_space) : nothing
u_lnd = land_sim !== nothing ? swap_space!(zeros(coupler_sim.boundary_space), land_sim.integrator.u.T_sfc) : nothing
u_ocn =
ocean_sim !== nothing ? swap_space!(zeros(coupler_sim.boundary_space), ocean_sim.integrator.u.T_sfc) : nothing
u_ice =
seaice_sim !== nothing ? swap_space!(zeros(coupler_sim.boundary_space), seaice_sim.integrator.u.T_sfc) : nothing

# global sums
land_e = land_sim !== nothing ? sum(get_slab_energy(land_sim, u_lnd)) ./ Δz_bot : FT(0)
Expand Down Expand Up @@ -89,13 +93,16 @@ function check_conservation(
if (prescribed_sic == false)

# ocean (from T_ml)
u_ocn = seaice_sim !== nothing ? swap_space!(seaice_sim.integrator.u.T_ml, coupler_sim.boundary_space) : nothing
u_ocn =
seaice_sim !== nothing ? swap_space!(zeros(coupler_sim.boundary_space), seaice_sim.integrator.u.T_ml) :
nothing
parent(u_ocn) .= abs.(parent(mask) .- FT(1)) .* parent(u_ocn)
ocean_e = ocean_sim !== nothing ? sum(get_ml_energy(seaice_sim, u_ocn)) ./ Δz_bot : FT(0)

# ice (from h_ice)
u_ice =
seaice_sim !== nothing ? swap_space!(seaice_sim.integrator.u.h_ice, coupler_sim.boundary_space) : nothing
seaice_sim !== nothing ? swap_space!(zeros(coupler_sim.boundary_space), seaice_sim.integrator.u.h_ice) :
nothing
parent(u_ice) .= abs.(parent(mask) .- FT(1)) .* parent(u_ice)
seaice_e = ocean_sim !== nothing ? sum(get_dyn_ice_energy(seaice_sim, u_ice)) ./ Δz_bot : FT(0)

Expand Down
Loading