Merge #198

198: Remove FT values from structs r=juliasloan25 a=juliasloan25



Co-authored-by: Julia Sloan <[email protected]>

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,
 );

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,
 );
 
 

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)
 

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)

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)
 

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

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)

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))

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

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)

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)

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)
 

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