diff --git a/experiments/AMIP/moist_mpi_earth/coupler_driver.jl b/experiments/AMIP/moist_mpi_earth/coupler_driver.jl index bd930d4ff..4b51ba7ce 100644 --- a/experiments/AMIP/moist_mpi_earth/coupler_driver.jl +++ b/experiments/AMIP/moist_mpi_earth/coupler_driver.jl @@ -71,9 +71,15 @@ if isinteractive() parsed_args["vert_diff"] = true #hide parsed_args["rad"] = "gray" #hide parsed_args["microphy"] = "0M" #hide - parsed_args["energy_check"] = true - parsed_args["precip_model"] = "0M" #hide + parsed_args["energy_check"] = true #hide parsed_args["mode_name"] = "slabplanet" #hide + parsed_args["t_end"] = "10days" #hide + parsed_args["dt_save_to_sol"] = "3600secs" #hide + parsed_args["dt_cpl"] = 200 #hide + parsed_args["dt"] = "200secs" #hide + parsed_args["mono_surface"] = true #hide + parsed_args["h_elem"] = 4 #hide + parsed_args["precip_model"] = "0M" #hide end ## read in some parsed command line arguments @@ -269,19 +275,18 @@ monthly_2d_diags = (; ) ## coupler simulation -cs = CouplerSimulation( +cs = CouplerSimulation{FT}( comms_ctx, - FT(Δt_cpl), - integrator.t, tspan, dates, boundary_space, - FT, + parsed_args, + integrator.t, + FT(Δt_cpl), (; land = land_mask, ocean = zeros(boundary_space), ice = zeros(boundary_space)), coupler_fields, model_sims, mode_specifics, - parsed_args, monthly_3d_diags, monthly_2d_diags, ); diff --git a/experiments/AMIP/moist_mpi_earth/coupler_driver_modular.jl b/experiments/AMIP/moist_mpi_earth/coupler_driver_modular.jl index 336498dcf..86b6f5b48 100644 --- a/experiments/AMIP/moist_mpi_earth/coupler_driver_modular.jl +++ b/experiments/AMIP/moist_mpi_earth/coupler_driver_modular.jl @@ -90,10 +90,10 @@ date0 = date = DateTime(parsed_args["start_date"], dateformat"yyyymmdd") mono_surface = parsed_args["mono_surface"] import ClimaCoupler -import ClimaCoupler.Regridder: land_sea_mask, update_masks!, combine_surfaces! +import ClimaCoupler.Regridder: land_sea_mask, update_masks!, combine_surfaces!, dummmy_remap! import ClimaCoupler.ConservationChecker: EnergyConservationCheck, WaterConservationCheck, check_conservation!, plot_global_conservation -import ClimaCoupler.Utilities: CoupledSimulation +import ClimaCoupler.Utilities: CoupledSimulation, float_type, swap_space! pkg_dir = pkgdir(ClimaCoupler) COUPLER_OUTPUT_DIR = joinpath(pkg_dir, "experiments/AMIP/moist_mpi_earth/output", joinpath(mode_name, run_name)) @@ -111,9 +111,7 @@ mask_data = joinpath(mask_dataset_path(), "seamask.nc") # import coupler utils include("coupler_utils/flux_calculator.jl") -include("coupler_utils/regridder.jl") # update_midmonth_data! include("coupler_utils/calendar_timer.jl") -include("coupler_utils/general_helper.jl") include("coupler_utils/bcfile_reader.jl") include("coupler_utils/variable_definer.jl") include("coupler_utils/diagnostics_gatherer.jl") @@ -287,21 +285,20 @@ if energy_check end ## coupler simulation -cs = CoupledSimulation( - FT(Δt_cpl), - integrator.t, +cs = CoupledSimulation{FT}( tspan, dates, boundary_space, - FT, - (; land = land_mask, ocean = zeros(boundary_space), ice = zeros(boundary_space)), coupler_fields, + parsed_args, + conservation_checks, + integrator.t, + FT(Δt_cpl), + (; land = land_mask, ocean = zeros(boundary_space), ice = zeros(boundary_space)), model_sims, mode_specifics, - parsed_args, monthly_3d_diags, monthly_2d_diags, - conservation_checks, ); diff --git a/experiments/AMIP/moist_mpi_earth/coupler_utils/bcfile_reader.jl b/experiments/AMIP/moist_mpi_earth/coupler_utils/bcfile_reader.jl index ef20261aa..82ec59aac 100644 --- a/experiments/AMIP/moist_mpi_earth/coupler_utils/bcfile_reader.jl +++ b/experiments/AMIP/moist_mpi_earth/coupler_utils/bcfile_reader.jl @@ -7,41 +7,42 @@ Stores information specific to each boundary condition from a file and each variable. The inputs are: - FT::F # float type + comms_ctx::X # communications context used for MPI datafile_cgll::S # file containing all regridded fields varname::V # name of the variable all_dates::D # all dates contained in the original data file + monthly_fields::C # Tuple of the two monthly fields, that will be used for the daily interpolation + scaling_function::O # function that scales, offsets or transforms the raw variable + land_mask::M # mask with 1 = land, 0 = ocean / sea-ice segment_idx::Vector{Int} # index of the monthly data in the file segment_idx0::Vector{Int} # `segment_idx` of the file data that is closest to date0 - monthly_fields::C # Tuple of the two monthly fields, that will be used for the daily interpolation segment_length::Vector{Int} # length of each month segment (used in the daily interpolation) - scaling_function::O # function that scales, offsets or transforms the raw variable interpolate_daily::Bool # switch to trigger daily interpolation - land_mask::M # mask with 1 = land, 0 = ocean / sea-ice + """ -struct BCFileInfo{F, X, S, V, D, C, O, M} - FT::F +struct BCFileInfo{FT, X, S, V, D, C, O, M} comms_ctx::X hd_outfile_root::S varname::V all_dates::D + monthly_fields::C + scaling_function::O + land_mask::M segment_idx::Vector{Int} segment_idx0::Vector{Int} - monthly_fields::C segment_length::Vector{Int} - scaling_function::O interpolate_daily::Bool - land_mask::M end +BCFileInfo{FT}(args...) = BCFileInfo{FT, typeof.(args[1:7])...}(args...) + +float_type(::BCFileInfo{FT}) where {FT} = FT + """ - Bcfile_info_init(FT, comms_ctx, datafile_rll, varname, boundary_space; interpolate_daily = false, segment_idx0 = [Int(1)], scaling_function = false) + bcfile_info_init(FT, comms_ctx, datafile_rll, varname, boundary_space; interpolate_daily = false, segment_idx0 = [Int(1)], scaling_function = false) Regrids from lat-lon grid to cgll grid, saving the output in a new file, and returns the info packaged in a single struct """ -#FT, datafile_rll, varname,boundary_space, interpolate_daily, segment_idx0, scaling_function, land_mask, date0, mono = (FT, sst_data, "SST", boundary_space, true, nothing, clean_sst, land_mask, date0, true) - - function bcfile_info_init( FT, comms_ctx, @@ -80,19 +81,18 @@ function bcfile_info_init( segment_idx0 != nothing ? segment_idx0 : [argmin(abs.(parse(FT, datetime_to_strdate(date0)) .- parse.(FT, datetime_to_strdate.(data_dates[:]))))] - return BCFileInfo( - FT, + return BCFileInfo{FT}( comms_ctx, hd_outfile_root, varname, data_dates, + current_fields, + scaling_function, + land_mask, deepcopy(segment_idx0), segment_idx0, - current_fields, segment_length, - scaling_function, interpolate_daily, - land_mask, ) end @@ -101,14 +101,13 @@ no_scaling(x, _info) = swap_space!(x, axes(_info.land_mask)) # IO - monthly """ - update_midmonth_data!(bcf_info) + update_midmonth_data!(date, bcf_info) Extracts boundary condition data from regridded (to model grid) NetCDF files (which times, depends on the specifications in the `bcf_info` struct). """ function update_midmonth_data!(date, bcf_info) - # monthly count - + FT = float_type(bcf_info) all_dates = bcf_info.all_dates midmonth_idx = bcf_info.segment_idx[1] midmonth_idx0 = bcf_info.segment_idx0[1] @@ -118,7 +117,6 @@ function update_midmonth_data!(date, bcf_info) varname = bcf_info.varname interpolate_daily = bcf_info.interpolate_daily comms_ctx = bcf_info.comms_ctx - FT = bcf_info.FT ClimaComms.barrier(comms_ctx) diff --git a/experiments/AMIP/moist_mpi_earth/coupler_utils/general_helper.jl b/experiments/AMIP/moist_mpi_earth/coupler_utils/general_helper.jl index ecdfeb0f0..860b7a36d 100644 --- a/experiments/AMIP/moist_mpi_earth/coupler_utils/general_helper.jl +++ b/experiments/AMIP/moist_mpi_earth/coupler_utils/general_helper.jl @@ -1,23 +1,25 @@ # most of these functions are temporary helpers until upstream issues are resolved # TODO: unify with coupler interface -struct CouplerSimulation{C, I, F, S, D, B, T, P} +struct CouplerSimulation{FT, C, S, D, B, P} comms_ctx::C - Δt_cpl::I - t::F tspan::S dates::D boundary_space::B - FT::T + parsed_args::P + t::FT + Δt_cpl::FT surface_masks::NamedTuple fields::NamedTuple model_sims::NamedTuple mode::NamedTuple - parsed_args::P monthly_3d_diags::NamedTuple monthly_2d_diags::NamedTuple 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) diff --git a/experiments/AMIP/moist_mpi_earth/coupler_utils/variable_definer.jl b/experiments/AMIP/moist_mpi_earth/coupler_utils/variable_definer.jl index cb81c90f3..5de31157e 100644 --- a/experiments/AMIP/moist_mpi_earth/coupler_utils/variable_definer.jl +++ b/experiments/AMIP/moist_mpi_earth/coupler_utils/variable_definer.jl @@ -10,7 +10,7 @@ get_var(cs, ::Val{:T}) = air_temperature(cs) get_var(cs, ::Val{:u}) = ClimaCore.Geometry.UVVector.(cs.model_sims.atmos_sim.integrator.u.c.uₕ).components.data.:1 get_var(cs, ::Val{:q_tot}) = - cs.model_sims.atmos_sim.integrator.u.c.ρq_tot ./ cs.model_sims.atmos_sim.integrator.u.c.ρ .* cs.FT(1000) + 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) diff --git a/experiments/AMIP/moist_mpi_earth/push_pull.jl b/experiments/AMIP/moist_mpi_earth/push_pull.jl index 70e2527ee..ecbfe6de3 100644 --- a/experiments/AMIP/moist_mpi_earth/push_pull.jl +++ b/experiments/AMIP/moist_mpi_earth/push_pull.jl @@ -21,9 +21,9 @@ and hydrostatic balance assumptions. The land model does not compute the surface a reasonable stand-in. """ function land_pull!(cs) + FT = float_type(cs) land_sim = cs.model_sims.land_sim csf = cs.fields - FT = cs.FT land_mask = cs.surface_masks.land parent(land_sim.integrator.p.bucket.ρ_sfc) .= parent(csf.ρ_sfc) parent(land_sim.integrator.p.bucket.turbulent_energy_flux) .= @@ -59,6 +59,7 @@ In the current version, the sea ice has a prescribed thickness, and we assume th sublimating. That contribution has been zeroed out in the atmos fluxes. """ function ice_pull!(cs) + FT = float_type(cs) ice_sim = cs.model_sims.ice_sim csf = cs.fields ice_mask = cs.surface_masks.ice diff --git a/experiments/AMIP/moist_mpi_earth/slab_ice/slab_init.jl b/experiments/AMIP/moist_mpi_earth/slab_ice/slab_init.jl index ab850dc46..cbc6d7a61 100644 --- a/experiments/AMIP/moist_mpi_earth/slab_ice/slab_init.jl +++ b/experiments/AMIP/moist_mpi_earth/slab_ice/slab_init.jl @@ -70,7 +70,7 @@ end 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!(SIC, axes(_info.land_mask)) ./ _info.FT(100.0) +clean_sic(SIC, _info) = swap_space!(SIC, axes(_info.land_mask)) ./ float_type(_info)(100.0) # 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) diff --git a/experiments/AMIP/moist_mpi_earth/slab_ocean/slab_init.jl b/experiments/AMIP/moist_mpi_earth/slab_ocean/slab_init.jl index 882992801..f887f6a1c 100644 --- a/experiments/AMIP/moist_mpi_earth/slab_ocean/slab_init.jl +++ b/experiments/AMIP/moist_mpi_earth/slab_ocean/slab_init.jl @@ -72,4 +72,4 @@ end 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)) .+ _info.FT(273.15)) +clean_sst(SST, _info) = (swap_space!(SST, axes(_info.land_mask)) .+ float_type(_info)(273.15)) diff --git a/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_driver.jl b/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_driver.jl index 13cb6b8eb..5f147b1d9 100644 --- a/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_driver.jl +++ b/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_driver.jl @@ -76,7 +76,7 @@ else end # init coupler -coupler_sim = CouplerSimulation(FT(Δt_cpl), integrator.t, boundary_space, FT, mask) +coupler_sim = CouplerSimulation{FT}(boundary_space, mask, integrator.t, FT(Δt_cpl)) # init coupler's boundary fields for regridding (TODO: technically this can be bypassed by directly rigridding on model grids) T_S = ClimaCore.Fields.zeros(boundary_space) # temperature diff --git a/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_utils/general_helper.jl b/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_utils/general_helper.jl index ea5ee3872..340bb764a 100644 --- a/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_utils/general_helper.jl +++ b/experiments/AMIP/moist_mpi_earth_dynamical_sea_ice/coupler_utils/general_helper.jl @@ -1,14 +1,16 @@ # most of these functions are temporary helpers until upstream issues are resolved # TODO: unify with coupler interface -struct CouplerSimulation{I, F, B, T, M} - Δt::I - t::F +struct CouplerSimulation{FT, B, M} boundary_space::B - FT::T mask::M + t::FT + Δt::FT end +CouplerSimulation{FT}(args...) where {FT} = CouplerSimulation{FT, typeof.(args[1:2])...}(args...) +float_type(::CouplerSimulation{FT}) where {FT} = FT + get_u(sim, t) = Geometry.UVVector.(sim.integrator.sol.u[t].c.uₕ).components.data.:1 function swap_space!(field, new_space) diff --git a/src/ConservationChecker.jl b/src/ConservationChecker.jl index 32c9e40be..6865cd78c 100644 --- a/src/ConservationChecker.jl +++ b/src/ConservationChecker.jl @@ -273,11 +273,10 @@ end function plot_global_conservation( cc::WaterConservationCheck, - coupler_sim::CoupledSimulation; + coupler_sim::CoupledSimulation{FT}; figname1 = "total_water.png", figname2 = "total_water_log.png", -) - FT = coupler_sim.FT +) where {FT} times = collect(1:length(cc.ρq_tot_atmos)) * coupler_sim.Δt_cpl diff_ρe_tot_atmos = (cc.ρq_tot_atmos .- cc.ρq_tot_atmos[1]) diff_ρe_tot_slab = (cc.ρq_tot_land .- cc.ρq_tot_land[1]) * FT(1e3) diff --git a/src/Utilities.jl b/src/Utilities.jl index 133cf1c19..a2d2d65fb 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -9,28 +9,30 @@ module Utilities using ClimaCore: ClimaCore, Fields, Spaces, Domains, Meshes, Topologies using ClimaComms -export CoupledSimulation, heaviside, swap_space!, create_space +export CoupledSimulation, float_type, heaviside, swap_space!, create_space """ Stores information needed to run a simulation with the coupler. """ -struct CoupledSimulation{I, F, S, D, B, T, FV, P, E} - Δt_cpl::I - t::F +struct CoupledSimulation{FT, S, D, B, FV, P, E} tspan::S dates::D boundary_space::B - FT::T - surface_masks::NamedTuple fields::FV + parsed_args::P + conservation_checks::E + t::FT + Δt_cpl::FT + surface_masks::NamedTuple model_sims::NamedTuple mode::NamedTuple - parsed_args::P monthly_3d_diags::NamedTuple monthly_2d_diags::NamedTuple - conservation_checks::E end +CoupledSimulation{FT}(args...) where {FT} = CoupledSimulation{FT, typeof.(args[1:6])...}(args...) +float_type(::CoupledSimulation{FT}) where {FT} = FT + """ heaviside(var) diff --git a/test/conservation_checker_tests.jl b/test/conservation_checker_tests.jl index a3818d96f..693bf2988 100644 --- a/test/conservation_checker_tests.jl +++ b/test/conservation_checker_tests.jl @@ -63,21 +63,20 @@ function coupler_sim_from_file( boundary_space = axes(land_mask) FT = eltype(land_mask) - Utilities.CoupledSimulation( - Δt_cpl, - t, + Utilities.CoupledSimulation{FT}( tspan, dates, boundary_space, - FT, - (; land = land_mask, ocean = FT(1) .- land_mask, ice = land_mask .* FT(0)), coupler_fields, + parsed_args, + conservation_checks, + t, + Δt_cpl, + (; land = land_mask, ocean = FT(1) .- land_mask, ice = land_mask .* FT(0)), model_sims, mode_specifics, - parsed_args, monthly_3d_diags, monthly_2d_diags, - conservation_checks, ) end diff --git a/test/regridder_tests.jl b/test/regridder_tests.jl index cd1f3d6e4..1d93e1fe7 100644 --- a/test/regridder_tests.jl +++ b/test/regridder_tests.jl @@ -35,21 +35,20 @@ end parent(ice_d)[:, (n ÷ 2 + 1):n, :, :] .= FT(0.5) # Fill in only the necessary parts of the simulation - cs = Utilities.CoupledSimulation( - nothing, - nothing, - nothing, - nothing, - nothing, - FT, - (; land = land_mask, ice = Fields.zeros(test_space), ocean = Fields.zeros(test_space)), - (;), - (; ice_sim = (; integrator = (; p = (; ice_mask = ice_d)))), - (;), - nothing, - (;), - (;), - (;), + cs = Utilities.CoupledSimulation{FT}( + nothing, # tspan + nothing, # dates + nothing, # boundary_space + nothing, # fields + nothing, # parsed_args + nothing, # conservation_checks + FT(100), # t + FT(100), # Δt_cpl + (; land = land_mask, ice = Fields.zeros(test_space), ocean = Fields.zeros(test_space)), # surface_masks + (; ice_sim = (; integrator = (; p = (; ice_mask = ice_d)))), # model_sims + (;), # mode + (;), # monthly_3d_diags + (;), # monthly_2d_diags ) Regridder.update_masks!(cs)