init

clean

Author: LenkaNovak

experiments/ClimaEarth/run_cloudless_aquaplanet.jl

@@ -0,0 +1,374 @@
+# # Cloudless Aquaplanet
+
+redirect_stderr(IOContext(stderr, :stacktrace_types_limited => Ref(false)))
+
+#=
+## Configuration initialization
+=#
+
+#=
+### Package Import
+=#
+
+## standard packages
+import Dates
+import YAML
+
+# ## ClimaESM packages
+import ClimaAtmos as CA
+import ClimaComms
+import ClimaCore as CC
+
+# ## Coupler specific imports
+import ClimaCoupler
+import ClimaCoupler:
+    BCReader,
+    ConservationChecker,
+    Checkpointer,
+    Diagnostics,
+    FieldExchanger,
+    FluxCalculator,
+    Interfacer,
+    Regridder,
+    TimeManager,
+    Utilities
+
+pkg_dir = pkgdir(ClimaCoupler)
+
+#=
+### Helper Functions
+These will be eventually moved to their respective component model and diagnostics packages, and so they should not
+contain any internals of the ClimaCoupler source code, except extensions to the Interfacer functions.
+=#
+
+## helpers for component models
+include("components/atmosphere/climaatmos.jl")
+include("components/ocean/slab_ocean.jl")
+
+## helpers for user-specified IO
+include("user_io/user_diagnostics.jl")
+include("user_io/user_logging.jl")
+
+include("user_io/io_helpers.jl")
+
+#=
+### Setup simulation parameters
+Here we follow ClimaCore's dry Held-Suarez `held_suarez_rhoe` example.
+=#
+
+## run names
+run_name = "cloudless_aquaplanet"
+coupler_output_dir = "$run_name"
+const FT = Float64
+restart_dir = "unspecified"
+restart_t = Int(0)
+
+## coupler simulation specific configuration
+Δt_cpl = Float64(400)
+t_end = "1000days"
+tspan = (Float64(0.0), Float64(time_to_seconds(t_end)))
+start_date = "19790301"
+hourly_checkpoint = true
+
+## namelist
+config_dict = Dict(
+    # general
+    "FLOAT_TYPE" => string(FT),
+    # file paths
+    "atmos_config_file" => nothing,
+    "coupler_toml_file" => nothing,
+    "coupler_output_dir" => coupler_output_dir,
+    "mode_name" => "",
+    "run_name" => run_name,
+    "atmos_config_repo" => "ClimaAtmos",
+    # timestepping
+    "dt" => "$(Δt_cpl)secs",
+    "dt_save_to_sol" => "1days",
+    "t_end" => t_end,
+    "start_date" => "19790301",
+    # domain
+    "h_elem" => 4,
+    "z_elem" => 10,
+    "z_max" => 30000.0, # semi-high top
+    "dz_bottom" => 300.0,
+    "nh_poly" => 4,
+    # output
+    "dt_save_to_sol" => "1days",
+    # numerics
+    "apply_limiter" => false,
+    "viscous_sponge" => false,
+    "rayleigh_sponge" => false,
+    "vert_diff" => "true",
+    "hyperdiff" => "ClimaHyperdiffusion",
+    # run
+    "job_id" => run_name,
+    "surface_setup" => "PrescribedSurface",
+    # diagnostic (nested with period and short_name)
+    "output_default_diagnostics" => false,
+    "diagnostics" => [
+        Dict(
+            "short_name" => [
+                "mse",
+                "lr",
+                "mass_strf",
+                "stab",
+                "vt",
+                "egr",
+                "ua",
+                "va",
+                "wa",
+                "ta",
+                "rhoa",
+                "pfull",
+            ],
+            "period" => "6hours",
+            "reduction" => "inst",
+        ),
+    ],
+    # held-suarez specific
+    "precip_model" => "0M",
+    "moist" => "equil",
+    "prognostic_surface" => "PrescribedSurfaceTemperature",
+    "turb_flux_partition" => "CombinedStateFluxes",
+    "rad" => "gray",
+)
+# TODO: increase diffusion coef?
+
+## merge dictionaries of command line arguments, coupler dictionary and component model dictionaries
+atmos_config_dict, config_dict = get_atmos_config_dict(config_dict)
+atmos_config_object = CA.AtmosConfig(atmos_config_dict)
+
+#=
+## Setup Communication Context
+=#
+
+comms_ctx = Utilities.get_comms_context(Dict("device" => "auto"))
+ClimaComms.init(comms_ctx)
+
+#=
+### I/O Directory Setup
+=#
+
+dir_paths = setup_output_dirs(output_dir = coupler_output_dir, comms_ctx = comms_ctx)
+ClimaComms.iamroot(comms_ctx) ? @info(config_dict) : nothing
+
+#=
+## Component Model Initialization
+=#
+
+#=
+### Atmosphere
+This uses the `ClimaAtmos.jl` model, with parameterization options specified in the `config_dict_atmos` dictionary.
+=#
+
+## init atmos model component
+atmos_sim = atmos_init(atmos_config_object);
+thermo_params = get_thermo_params(atmos_sim)
+
+#=
+### Boundary Space
+=#
+
+## init a 2D boundary space at the surface
+boundary_space = CC.Spaces.horizontal_space(atmos_sim.domain.face_space) # TODO: specify this in the coupler and pass it to all component models #665
+
+#=
+### Surface Model: Slab Ocean
+=#
+ocean_sim = ocean_init(
+    FT;
+    tspan = tspan,
+    dt = Δt_cpl,
+    space = boundary_space,
+    saveat = saveat,
+    area_fraction = ones(boundary_space),
+    thermo_params = thermo_params,
+    evolving = evolving_ocean,
+)
+
+#=
+## Coupler Initialization
+=#
+
+## coupler exchange fields
+coupler_field_names = (
+    :T_S,
+    :z0m_S,
+    :z0b_S,
+    :ρ_sfc,
+    :q_sfc,
+    :surface_direct_albedo,
+    :surface_diffuse_albedo,
+    :beta,
+    :F_turb_energy,
+    :F_turb_moisture,
+    :F_turb_ρτxz,
+    :F_turb_ρτyz,
+    :F_radiative,
+    :P_liq,
+    :P_snow,
+    :radiative_energy_flux_toa,
+    :P_net,
+    :temp1,
+    :temp2,
+)
+coupler_fields =
+    NamedTuple{coupler_field_names}(ntuple(i -> CC.Fields.zeros(boundary_space), length(coupler_field_names)))
+Utilities.show_memory_usage(comms_ctx)
+
+## model simulations
+model_sims = (atmos_sim = atmos_sim, ocean_sim = ocean_sim);
+
+## dates
+dates = (; date = [date], date0 = [date0], date1 = [Dates.firstdayofmonth(date0)], new_month = [false])
+
+#=
+## Initialize Callbacks
+=#
+checkpoint_cb = TimeManager.HourlyCallback(
+    dt = hourly_checkpoint_dt,
+    func = checkpoint_sims,
+    ref_date = [dates.date[1]],
+    active = hourly_checkpoint,
+) # 20 days
+update_firstdayofmonth!_cb = TimeManager.MonthlyCallback(
+    dt = FT(1),
+    func = TimeManager.update_firstdayofmonth!,
+    ref_date = [dates.date1[1]],
+    active = true,
+)
+dt_water_albedo = parse(FT, filter(x -> !occursin(x, "hours"), dt_rad))
+albedo_cb = TimeManager.HourlyCallback(
+    dt = dt_water_albedo,
+    func = FluxCalculator.water_albedo_from_atmosphere!,
+    ref_date = [dates.date[1]],
+    active = mode_name == "amip",
+)
+callbacks =
+    (; checkpoint = checkpoint_cb, update_firstdayofmonth! = update_firstdayofmonth!_cb, water_albedo = albedo_cb)
+
+#=
+## Initialize turbulent fluxes
+
+Decide on the type of turbulent flux partition (see `FluxCalculator` documentation for more details).
+=#
+turbulent_fluxes = nothing
+if config_dict["turb_flux_partition"] == "CombinedStateFluxes"
+    turbulent_fluxes = FluxCalculator.CombinedStateFluxes()
+else
+    error("turb_flux_partition must be either PartitionedStateFluxes or CombinedStateFluxes")
+end
+
+#=
+## Initialize Coupled Simulation
+=#
+
+cs = Interfacer.CoupledSimulation{FT}(
+    comms_ctx,
+    dates,
+    boundary_space,
+    coupler_fields,
+    config_dict,
+    conservation_checks,
+    [tspan[1], tspan[2]],
+    atmos_sim.integrator.t,
+    Δt_cpl,
+    (; land = land_area_fraction, ocean = zeros(boundary_space), ice = zeros(boundary_space)),
+    model_sims,
+    mode_specifics,
+    diagnostics,
+    callbacks,
+    dir_paths,
+    turbulent_fluxes,
+    thermo_params,
+);
+
+#=
+## Restart component model states if specified in the config_dict
+=#
+
+if restart_dir !== "unspecified"
+    for sim in cs.model_sims
+        if Checkpointer.get_model_prog_state(sim) !== nothing
+            Checkpointer.restart_model_state!(sim, comms_ctx, restart_t; input_dir = restart_dir)
+        end
+    end
+end
+
+#=
+## Initialize Component Model Exchange
+=#
+
+# 1.coupler updates surface model area fractions
+Regridder.update_surface_fractions!(cs)
+
+# 2.surface density (`ρ_sfc`): calculated by the coupler by adiabatically extrapolating atmospheric thermal state to the surface.
+# For this, we need to import surface and atmospheric fields. The model sims are then updated with the new surface density.
+FieldExchanger.import_combined_surface_fields!(cs.fields, cs.model_sims, cs.turbulent_fluxes)
+FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes)
+FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
+
+# 3.surface vapor specific humidity (`q_sfc`): step surface models with the new surface density to calculate their respective `q_sfc` internally
+Interfacer.step!(ocean_sim, Δt_cpl)
+
+# 4.turbulent fluxes
+## import the new surface properties into the coupler (note the atmos state was also imported in step 3.)
+FieldExchanger.import_combined_surface_fields!(cs.fields, cs.model_sims, cs.turbulent_fluxes) # i.e. T_sfc, albedo, z0, beta, q_sfc
+## calculate turbulent fluxes inside the atmos cache based on the combined surface state in each grid box
+FluxCalculator.combined_turbulent_fluxes!(cs.model_sims, cs.fields, cs.turbulent_fluxes) # this updates the atmos thermo state, sfc_ts
+
+# 5.reinitialize models + radiative flux: prognostic states and time are set to their initial conditions. For atmos, this also triggers the callbacks and sets a nonzero radiation flux (given the new sfc_conditions)
+FieldExchanger.reinit_model_sims!(cs.model_sims)
+
+# 6.update all fluxes: coupler re-imports updated atmos fluxes
+FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes)
+FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
+
+#=
+## Coupling Loop
+=#
+
+function solve_coupler!(cs)
+    (; Δt_cpl, tspan, comms_ctx) = cs
+
+    ClimaComms.iamroot(comms_ctx) && @info("Starting coupling loop")
+    ## step in time
+    for t in ((tspan[begin] + Δt_cpl):Δt_cpl:tspan[end])
+
+        cs.dates.date[1] = TimeManager.current_date(cs, t)
+
+        ## print date on the first of month
+        if cs.dates.date[1] >= cs.dates.date1[1]
+            ClimaComms.iamroot(comms_ctx) && @show(cs.dates.date[1])
+        end
+
+        ClimaComms.barrier(comms_ctx)
+
+        ## update water albedo from wind at dt_water_albedo (this will be extended to a radiation callback from the coupler)
+        TimeManager.trigger_callback!(cs, cs.callbacks.water_albedo)
+
+        ## run component models sequentially for one coupling timestep (Δt_cpl)
+        FieldExchanger.update_model_sims!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
+
+        ## step sims
+        FieldExchanger.step_model_sims!(cs.model_sims, t)
+
+        ## exchange combined fields and (if specified) calculate fluxes using combined states
+        FieldExchanger.import_combined_surface_fields!(cs.fields, cs.model_sims, cs.turbulent_fluxes) # i.e. T_sfc, surface_albedo, z0, beta
+        FluxCalculator.combined_turbulent_fluxes!(cs.model_sims, cs.fields, cs.turbulent_fluxes)
+
+        FieldExchanger.import_atmos_fields!(cs.fields, cs.model_sims, cs.boundary_space, cs.turbulent_fluxes) # radiative and/or turbulent
+
+        ## callback to update the fist day of month if needed (for BCReader)
+        TimeManager.trigger_callback!(cs, cs.callbacks.update_firstdayofmonth!)
+
+        ## callback to checkpoint model state
+        TimeManager.trigger_callback!(cs, cs.callbacks.checkpoint)
+
+    end
+
+    return nothing
+end
+
+
+solve_coupler!(cs)