Merge pull request #534 from CliMA/js/ft32

add Float32 compatibility

.buildkite/pipeline.yml

@@ -21,7 +21,7 @@ env:
   PERF_CONFIG_PATH: "config/perf_configs"
   MPI_CONFIG_PATH: "config/mpi_configs"
 
-timeout_in_minutes: 1440
+timeout_in_minutes: 240
 
 steps:
   - label: "init environment :computer:"
@@ -143,6 +143,13 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: "Slabplanet: default with Float32"
+        key: "slabplanet_ft32"
+        command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/slabplanet_ft32.yml"
+        artifact_paths: "experiments/AMIP/modular/output/slabplanet/slabplanet_ft32_artifacts/*"
+        agents:
+          slurm_mem: 20GB
+
       - label: "Slabplanet: partitioned turbulent fluxes"
         key: "slabplanet_partitioned_fluxes"
         command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/slabplanet_partitioned_fluxes.yml"
@@ -295,9 +302,17 @@ steps:
           slurm_ntasks: 1
           slurm_mem: 20GB
 
+      - label: "AMIP - modular, Float32 test"
+        command: "julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/coarse_single_modular_ft32.yml"
+        artifact_paths: "experiments/AMIP/modular/output/amip/coarse_single_modular_ft32_artifacts/*"
+        agents:
+          slurm_ntasks: 1
+          slurm_mem: 20GB
+
       - label: "MPI AMIP"
         command: "srun julia --color=yes --project=experiments/AMIP/modular/ experiments/AMIP/modular/coupler_driver_modular.jl --config_file $CONFIG_PATH/coarse_mpi_n2.yml"
         artifact_paths: "experiments/AMIP/modular/output/amip/coarse_mpi_n2_artifacts/*"
+        timeout_in_minutes: 240
         env:
           CLIMACORE_DISTRIBUTED: "MPI"
         agents:

config/model_configs/coarse_single_modular_ft32.yml

@@ -0,0 +1,18 @@
+run_name: "coarse_single_modular_ft32"
+moist: "equil"
+vert_diff: "true"
+rad: "gray"
+energy_check: false
+mode_name: "amip"
+anim: true
+t_end: "10days"
+dt_save_to_sol: "100days"
+dt_cpl: 200
+dt: "200secs"
+mono_surface: true
+h_elem: 6
+dt_save_restart: "10days"
+precip_model: "0M"
+apply_limiter: false
+FLOAT_TYPE: "Float32"
+job_id: "coarse_single_modular_ft32"

config/model_configs/slabplanet_ft32.yml

@@ -0,0 +1,17 @@
+run_name: "slabplanet_ft32"
+moist: "equil"
+vert_diff: "true"
+rad: "gray"
+energy_check: true
+mode_name: "slabplanet"
+t_end: "10days"
+dt_save_to_sol: "9days"
+dt_cpl: 200
+dt: "200secs"
+mono_surface: true
+h_elem: 4
+precip_model: "0M"
+anim: true
+apply_limiter: false
+FLOAT_TYPE: "Float32"
+job_id: "slabplanet_ft32"

experiments/AMIP/modular/components/land/bucket_init.jl

@@ -163,19 +163,19 @@ Initializes the bucket model variables.
 """
 function bucket_init(
     ::Type{FT},
-    tspan::Tuple{FT, FT},
+    tspan::Tuple{Float64, Float64},
     config::String,
     albedo_type::String,
     land_temperature_anomaly::String,
     comms_ctx::AbstractCommsContext,
     regrid_dirpath::String;
     space,
-    dt::FT,
-    saveat::FT,
+    dt::Float64,
+    saveat::Float64,
     area_fraction,
     stepper = CTS.RK4(),
     date_ref::DateTime,
-    t_start::FT,
+    t_start::Float64,
 ) where {FT}
     if config != "sphere"
         println(
@@ -210,7 +210,7 @@ function bucket_init(
     z_0b = FT(1e-3)
     κ_soil = FT(0.7)
     ρc_soil = FT(2e8)
-    t_crit = dt # This is the timescale on which snow exponentially damps to zero, in the case where all
+    t_crit = FT(dt) # This is the timescale on which snow exponentially damps to zero, in the case where all
     # the snow would melt in time t_crit. It prevents us from having to specially time step in cases where
     # all the snow melts in a single timestep.
     params = BucketModelParameters(κ_soil, ρc_soil, albedo, σS_c, W_f, z_0m, z_0b, t_crit, earth_param_set)

experiments/AMIP/modular/components/ocean/eisenman_seaice_init.jl

@@ -103,13 +103,15 @@ function solve_eisenman_model!(Y, Ya, p, thermo_params, Δt)
 
     # ice thickness and mixed layer temperature changes due to atmosphereic and ocean fluxes
     ice_covered = parent(h_ice)[1] > 0
+
+    FT = eltype(T_ml)
     if ice_covered # ice-covered
         F_base = @. C0_base * (T_ml - T_base)
-        ΔT_ml = @. -(F_base - ocean_qflux) * Δt / (hρc_ml)
-        Δh_ice = @. (F_atm - F_base - ocean_qflux) * Δt / L_ice
-        @. e_base .+= F_base * Δt
+        ΔT_ml = @. -(F_base - ocean_qflux) * FT(Δt) / (hρc_ml)
+        Δh_ice = @. (F_atm - F_base - ocean_qflux) * FT(Δt) / L_ice
+        @. e_base .+= F_base * FT(Δt)
     else # ice-free
-        ΔT_ml = @. -(F_atm - ocean_qflux) * Δt / (hρc_ml)
+        ΔT_ml = @. -(F_atm - ocean_qflux) * FT(Δt) / (hρc_ml)
         Δh_ice = 0
     end
 
@@ -317,11 +319,13 @@ function get_field(sim::EisenmanIceSimulation, ::Val{:energy})
     e_base = cache.Ya.e_base
     ocean_qflux = cache.Ya.ocean_qflux
 
+    FT = eltype(sim.integrator.u)
+
     hρc_ml = p_o.h * p_o.ρ * p_o.c
 
     e_ml = @. p_o.h * p_o.ρ * p_o.c * sim.integrator.u.T_ml # heat
     e_ice = @. p_i.L_ice * sim.integrator.u.h_ice # phase
-    e_qflux = @. ocean_qflux * sim.integrator.t
+    e_qflux = @. ocean_qflux * FT(sim.integrator.t)
 
     return @. e_ml + e_ice + e_qflux + e_base
 

experiments/AMIP/modular/components/ocean/prescr_seaice_init.jl

@@ -79,7 +79,7 @@ function ice_rhs!(du, u, p, _)
 
     # do not count tendencies that lead to temperatures above freezing, and mask out no-ice areas
     area_mask = Regridder.binary_mask.(area_fraction, threshold = eps(FT))
-    unphysical = @. Regridder.binary_mask.(T_freeze - (Y.T_sfc + FT(rhs) * p.dt), threshold = FT(0)) .* area_mask
+    unphysical = @. Regridder.binary_mask.(T_freeze - (Y.T_sfc + FT(rhs) * FT(p.dt)), threshold = FT(0)) .* area_mask
     parent(dY.T_sfc) .= parent(rhs .* unphysical)
 
     @. p.q_sfc = TD.q_vap_saturation_generic.(p.thermo_params, Y.T_sfc, p.ρ_sfc, TD.Ice())

experiments/AMIP/modular/coupler_driver_modular.jl

@@ -214,7 +214,17 @@ boundary_space = Spaces.horizontal_space(atmos_sim.domain.face_space)
 
 # init land-sea fraction
 land_fraction =
-    Regridder.land_fraction(FT, REGRID_DIR, comms_ctx, land_mask_data, "LSMASK", boundary_space, mono = mono_surface)
+    FT.(
+        Regridder.land_fraction(
+            FT,
+            REGRID_DIR,
+            comms_ctx,
+            land_mask_data,
+            "LSMASK",
+            boundary_space,
+            mono = mono_surface,
+        )
+    )
 
 #=
 ### Ocean and Sea Ice

src/ConservationChecker.jl

@@ -100,10 +100,10 @@ function check_conservation!(
             parent(coupler_sim.fields.F_radiative_TOA) .= parent(Interfacer.get_field(sim, Val(:F_radiative_TOA)))
 
             if isempty(ccs.toa_net_source)
-                radiation_sources_accum = sum(coupler_sim.fields.F_radiative_TOA .* coupler_sim.Δt_cpl) # ∫ J / m^2 dA
+                radiation_sources_accum = sum(coupler_sim.fields.F_radiative_TOA .* FT(coupler_sim.Δt_cpl)) # ∫ J / m^2 dA
             else
                 radiation_sources_accum =
-                    sum(coupler_sim.fields.F_radiative_TOA .* coupler_sim.Δt_cpl) .+ ccs.toa_net_source[end] # ∫ J / m^2 dA
+                    sum(coupler_sim.fields.F_radiative_TOA .* FT(coupler_sim.Δt_cpl)) .+ ccs.toa_net_source[end] # ∫ J / m^2 dA
             end
             push!(ccs.toa_net_source, radiation_sources_accum)
 
@@ -211,7 +211,8 @@ function surface_water_gain_from_rates(cs::Interfacer.CoupledSimulation)
     evaporation = cs.fields.F_turb_moisture # kg / m^2 / s / layer depth
     precipitation_l = cs.fields.P_liq
     precipitation_s = cs.fields.P_snow
-    @. -(evaporation + precipitation_l + precipitation_s) * cs.Δt_cpl # kg / m^2 / layer depth
+    FT = eltype(evaporation)
+    @. -(evaporation + precipitation_l + precipitation_s) * FT(cs.Δt_cpl) # kg / m^2 / layer depth
 end
 
 # setup the GKS socket application environment as nul for better performance and to avoid GKS connection errors while plotting

src/Diagnostics.jl

@@ -6,7 +6,7 @@ This module contains functions for defining, gathering and outputting online mod
 module Diagnostics
 
 using ClimaCore: Spaces, Fields, InputOutput
-using ClimaCoupler.Interfacer: CoupledSimulation
+using ClimaCoupler.Interfacer: CoupledSimulation, float_type
 using Dates
 using ClimaCoupler.TimeManager: AbstractFrequency, Monthly, EveryTimestep, trigger_callback
 using ClimaComms
@@ -78,12 +78,12 @@ end
 Collects diagnostics in diags names.
 """
 function collect_diags(cs::CoupledSimulation, dg::DiagnosticsGroup)
-
+    FT = float_type(cs)
     diags = (;)
 
     diag_names = propertynames(dg.field_vector)
     for name in diag_names
-        diags = (; diags..., zip((name,), (get_var(cs, Val(name)),))...)
+        diags = (; diags..., zip((name,), (FT.(get_var(cs, Val(name))),))...)
     end
 
     return Fields.FieldVector(; diags...)

src/FieldExchanger.jl

@@ -121,7 +121,6 @@ Updates the surface component model cache with the current coupler fields of F_t
 - `csf`: [NamedTuple] containing coupler fields.
 """
 function update_sim!(sim::Interfacer.SurfaceModelSimulation, csf, turbulent_fluxes, area_fraction = nothing)
-
     FT = eltype(area_fraction)
 
     # atmospheric surface density
@@ -130,24 +129,23 @@ function update_sim!(sim::Interfacer.SurfaceModelSimulation, csf, turbulent_flux
     # turbulent fluxes
     # when PartitionedStateFluxes, turbulent fluxes are updated during the flux calculation
     if turbulent_fluxes isa FluxCalculator.CombinedStateFluxes
-
         Interfacer.update_field!(
             sim,
             Val(:turbulent_energy_flux),
-            Regridder.binary_mask.(area_fraction, threshold = eps(FT)) .* csf.F_turb_energy,
+            FT.(Regridder.binary_mask.(area_fraction, threshold = eps(FT)) .* csf.F_turb_energy),
         )
         Interfacer.update_field!(
             sim,
             Val(:turbulent_moisture_flux),
-            Regridder.binary_mask.(area_fraction, threshold = eps(FT)) .* csf.F_turb_moisture,
+            FT.(Regridder.binary_mask.(area_fraction, threshold = eps(FT)) .* csf.F_turb_moisture),
         )
     end
 
     # radiative fluxes
     Interfacer.update_field!(
         sim,
         Val(:radiative_energy_flux),
-        Regridder.binary_mask.(area_fraction, threshold = eps(FT)) .* csf.F_radiative,
+        FT.(Regridder.binary_mask.(area_fraction, threshold = eps(FT)) .* csf.F_radiative),
     )
 
     # precipitation

test/component_model_tests/bucket_tests.jl

@@ -6,12 +6,7 @@ using ClimaCore: Fields, Spaces
 include(pkgdir(ClimaCoupler, "experiments/AMIP/modular/components/land/bucket_init.jl"))
 include(pkgdir(ClimaCoupler, "experiments/AMIP/modular/components/land/bucket_utils.jl"))
 
-# struct DummySimulationBucket{I} <: BucketSimulation
-#     integrator::I
-# end
-
-# TODO bucket doesn't currently work with Float32, but we want to eventually test with both FTs
-for FT in (Float64,)
+for FT in (Float32, Float64)
     @testset "dss_state! BucketSimulation for FT=$FT" begin
         # use TestHelper to create space, extract surface space
         subsurface_space = create_space(FT, nz = 2)
@@ -35,7 +30,6 @@ for FT in (Float64,)
             t = FT(0),
         )
         integrator_copy = deepcopy(integrator)
-        # sim = DummySimulationBucket(integrator)
         sim = BucketSimulation(nothing, nothing, nothing, integrator, nothing)
 
         # make fields non-constant to check the impact of the dss step