isolate area fraction logic to update_surface_fractions

juliasloan25 · juliasloan25 · commit dab0ae5bfbc0 · 2025-10-29T16:57:58.000-07:00
diff --git a/docs/src/fieldexchanger.md b/docs/src/fieldexchanger.md
@@ -39,6 +39,6 @@ the atmosphere and each surface model.
 
 ```@docs
     ClimaCoupler.FieldExchanger.combine_surfaces!
-    ClimaCoupler.FieldExchanger.resolve_ocean_ice_fractions!
+    ClimaCoupler.FieldExchanger.resolve_area_fractions!
     ClimaCoupler.FieldExchanger.import_atmos_fields!
 ```
diff --git a/experiments/ClimaEarth/components/ocean/oceananigans.jl b/experiments/ClimaEarth/components/ocean/oceananigans.jl
@@ -43,7 +43,7 @@ Specific details about the default model configuration
 can be found in the documentation for `ClimaOcean.ocean_simulation`.
 """
 function OceananigansSimulation(
-    area_fraction,
+    boundary_space,
     start_date,
     stop_date;
     output_dir,
@@ -141,7 +141,7 @@ function OceananigansSimulation(
     lat_cc = reshape(lat_cc, 1, length(lat_cc))
     target_points_cc = @. CC.Geometry.LatLongPoint(lat_cc, long_cc)
     # TODO: We can remove the `nothing` after CC > 0.14.33
-    remapper_cc = CC.Remapping.Remapper(axes(area_fraction), target_points_cc, nothing)
+    remapper_cc = CC.Remapping.Remapper(boundary_space, target_points_cc, nothing)
 
     # Construct two 2D Center/Center fields to use as scratch space while remapping
     scratch_cc1 = OC.Field{OC.Center, OC.Center, Nothing}(grid)
@@ -178,10 +178,13 @@ function OceananigansSimulation(
         ocean.output_writers[:diagnostics] = netcdf_writer
     end
 
-    # Initialize with 0 ice concentration; this will be updated in `resolve_ocean_ice_fractions!`
+    # Initialize with 0 ice concentration; this will be updated in `resolve_area_fractions!`
     # if the ocean is coupled to a non-prescribed sea ice model.
     ice_concentration = OC.Field{OC.Center, OC.Center, Nothing}(grid)
 
+    # Create a dummy area fraction that will get overwritten in `update_surface_fractions!`
+    area_fraction = ones(boundary_space)
+
     return OceananigansSimulation(
         ocean,
         area_fraction,
@@ -192,7 +195,7 @@ function OceananigansSimulation(
 end
 
 """
-    FieldExchanger.resolve_ocean_ice_fractions!(ocean_sim, ice_sim, land_fraction)
+    FieldExchanger.resolve_area_fractions!(ocean_sim, ice_sim, land_fraction)
 
 Ensure the ocean and ice area fractions are consistent with each other.
 This matters in the case of a LatitudeLongitudeGrid, which is only
@@ -206,7 +209,7 @@ and doesn't need to be set again since its fraction is static.
 This function also updates the ice concentration field in the ocean simulation
 so that it can be used for weighting flux updates.
 """
-function FieldExchanger.resolve_ocean_ice_fractions!(
+function FieldExchanger.resolve_area_fractions!(
     ocean_sim::OceananigansSimulation,
     ice_sim,
     land_fraction,
@@ -222,9 +225,9 @@ function FieldExchanger.resolve_ocean_ice_fractions!(
         polar_mask = CC.Fields.zeros(boundary_space)
         polar_mask .= abs.(lat) .>= FT(80)
 
-        # TODO do we want both to be 0 since we use capped lat/lon?
-        # Set ice fraction to 1 - land_fraction and ocean fraction to 0 where polar_mask is 1
-        @. ice_fraction = ifelse.(polar_mask == FT(1), FT(1) - land_fraction, ice_fraction)
+        # Set land fraction to 1 and ice/ocean fraction to 0 where polar_mask is 1
+        @. land_fraction = ifelse.(polar_mask == FT(1), FT(1), land_fraction)
+        @. ice_fraction = ifelse.(polar_mask == FT(1), FT(0), ice_fraction)
         @. ocean_fraction = ifelse.(polar_mask == FT(1), FT(0), ocean_fraction)
     end
 
diff --git a/experiments/ClimaEarth/components/ocean/prescr_ocean.jl b/experiments/ClimaEarth/components/ocean/prescr_ocean.jl
@@ -65,7 +65,6 @@ function PrescribedOceanSimulation(
     space,
     start_date,
     t_start,
-    area_fraction,
     thermo_params,
     comms_ctx;
     z0m = FT(5.8e-5),
@@ -113,7 +112,7 @@ function PrescribedOceanSimulation(
         α_diffuse = ones(space) .* α_diffuse_val,
         u_int = zeros(space),
         v_int = zeros(space),
-        area_fraction = area_fraction,
+        area_fraction = ones(space),
         phase = TD.Liquid(),
         thermo_params = thermo_params,
         SST_timevaryinginput = SST_timevaryinginput,
diff --git a/experiments/ClimaEarth/components/ocean/prescr_seaice.jl b/experiments/ClimaEarth/components/ocean/prescr_seaice.jl
@@ -119,12 +119,8 @@ function PrescribedIceSimulation(
     )
 
     # Get initial SIC values and use them to calculate ice fraction
-    SIC_init = CC.Fields.zeros(space)
-    evaluate!(SIC_init, SIC_timevaryinginput, tspan[1])
-
-    # Overwrite ice fraction with the static land area fraction anywhere we have nonzero land area
-    #  max needed to avoid Float32 errors (see issue #271; Heisenbug on HPC)
-    ice_fraction = @. max(min(SIC_init, FT(1) - land_fraction), FT(0))
+    ice_fraction = CC.Fields.zeros(space)
+    evaluate!(ice_fraction, SIC_timevaryinginput, tspan[1])
 
     params = IceSlabParameters{FT}()
 
diff --git a/experiments/ClimaEarth/components/ocean/slab_ocean.jl b/experiments/ClimaEarth/components/ocean/slab_ocean.jl
@@ -55,7 +55,7 @@ function slab_ocean_space_init(space, params)
 end
 
 """
-    SlabOceanSimulation(::Type{FT}; tspan, dt, saveat, space, area_fraction, stepper = CTS.RK4()) where {FT}
+    SlabOceanSimulation(::Type{FT}; tspan, dt, saveat, space, stepper = CTS.RK4()) where {FT}
 
 Initializes the `DiffEq` problem, and creates a Simulation-type object containing the necessary information for `step!` in the coupling loop.
 """
@@ -65,7 +65,6 @@ function SlabOceanSimulation(
     dt,
     saveat,
     space,
-    area_fraction,
     thermo_params,
     stepper = CTS.RK4(),
     evolving = true,
@@ -80,7 +79,7 @@ function SlabOceanSimulation(
         F_turb_energy = CC.Fields.zeros(space),
         SW_d = CC.Fields.zeros(space),
         LW_d = CC.Fields.zeros(space),
-        area_fraction = area_fraction,
+        area_fraction = CC.Fields.ones(space),
         thermo_params = thermo_params,
         α_direct = CC.Fields.ones(space) .* params.α,
         α_diffuse = CC.Fields.ones(space) .* params.α,
diff --git a/experiments/ClimaEarth/setup_run.jl b/experiments/ClimaEarth/setup_run.jl
@@ -307,16 +307,6 @@ function CoupledSimulation(config_dict::AbstractDict)
         # Determine whether to use a shared surface space
         shared_surface_space = share_surface_space ? boundary_space : nothing
         if land_model == "bucket"
-
-            # TODO move this into resolve_area_fractions and call after init, think about restarts
-            polar_mask = CC.Fields.zeros(boundary_space)
-            lat = CC.Fields.coordinate_field(boundary_space).lat
-            polar_mask .= abs.(lat) .>= FT(80)
-
-            # Set land fraction to 1 where polar_mask is 1
-            @. land_fraction = ifelse.(polar_mask == FT(1), FT(1), land_fraction)
-
-
             land_sim = BucketSimulation(
                 FT;
                 dt = component_dt_dict["dt_land"],
@@ -355,97 +345,51 @@ function CoupledSimulation(config_dict::AbstractDict)
             error("Invalid land model specified: $(land_model)")
         end
 
-        # TODO separate drivers to clean this up
-        ## sea ice model
-        if ice_model == "prescribed"
-            ice_sim = PrescribedIceSimulation(
-                FT;
-                tspan = tspan,
-                dt = component_dt_dict["dt_seaice"],
-                saveat = saveat,
-                space = boundary_space,
-                thermo_params = thermo_params,
-                comms_ctx,
-                start_date,
-                land_fraction,
-                sic_path = subseasonal_sic,
-            )
-            ice_fraction = Interfacer.get_field(ice_sim, Val(:area_fraction))
-        elseif ice_model == "clima_seaice"
-            @assert sim_mode <: CMIPMode
-
-            # TODO how should we initialize ocean fraction when using ClimaSeaIce?
-            # TODO init everything with AF=1, then resolve after getting sea ice
-            sic_data = try
-                joinpath(
-                    @clima_artifact("historical_sst_sic", comms_ctx),
-                    "MODEL.ICE.HAD187001-198110.OI198111-202206.nc",
-                )
-            catch error
-                @warn "Using lowres SIC. If you want the higher resolution version, you have to obtain it from ClimaArtifacts"
-                joinpath(
-                    @clima_artifact("historical_sst_sic_lowres", comms_ctx),
-                    "MODEL.ICE.HAD187001-198110.OI198111-202206_lowres.nc",
-                )
-            end
-            @info "Using initial condition prescribed SIC file: " sic_data
-
-            SIC_timevaryinginput = TimeVaryingInput(
-                sic_data,
-                "SEAICE",
-                boundary_space,
-                reference_date = start_date,
-                file_reader_kwargs = (; preprocess_func = (data) -> data / 100,), ## convert to fraction
-            )
-
-            # Get initial SIC values and use them to calculate ice fraction
-            ice_fraction = CC.Fields.zeros(boundary_space)
-            evaluate!(ice_fraction, SIC_timevaryinginput, tspan[1])
-        else
-            error("Invalid ice model specified: $(ice_model)")
-        end
-
-        ## ocean model using prescribed data
-        ice_fraction = Interfacer.get_field(ice_sim, Val(:area_fraction))
-        ocean_fraction = FT(1) .- ice_fraction .- land_fraction
-
+        ## ocean model
         if sim_mode <: CMIPMode
             stop_date = date(tspan[end] - tspan[begin])
             ocean_sim = OceananigansSimulation(
-                ocean_fraction,
+                boundary_space,
                 start_date,
                 stop_date;
                 output_dir = dir_paths.ocean_output_dir,
                 comms_ctx,
             )
-
-            if ice_model == "clima_seaice"
-                ice_sim = ClimaSeaIceSimulation(
-                    ice_fraction,
-                    ocean_sim;
-                    output_dir = dir_paths.ice_output_dir,
-                    start_date,
-                )
-                # TODO don't need to initialize ocean fraction correctly if we do this
-                # TODO can rename to `resolve_area_fractions!` and also make land_fraction cover poles here instead of in driver
-                FieldExchanger.resolve_ocean_ice_fractions!(
-                    ocean_sim,
-                    ice_sim,
-                    land_fraction,
-                )
-            end
         else
             ocean_sim = PrescribedOceanSimulation(
                 FT,
                 boundary_space,
                 start_date,
                 t_start,
-                ocean_fraction,
                 thermo_params,
                 comms_ctx;
                 sst_path = subseasonal_sst,
             )
         end
+        ## sea ice model
+        if ice_model == "clima_seaice"
+            ice_sim = ClimaSeaIceSimulation(
+                ice_fraction,
+                ocean_sim;
+                output_dir = dir_paths.ice_output_dir,
+                start_date,
+            )
+        elseif ice_model == "prescribed"
+            ice_sim = PrescribedIceSimulation(
+                FT;
+                tspan = tspan,
+                dt = component_dt_dict["dt_seaice"],
+                saveat = saveat,
+                space = boundary_space,
+                thermo_params = thermo_params,
+                comms_ctx,
+                start_date,
+                land_fraction,
+                sic_path = subseasonal_sic,
+            )
+        else
+            error("Invalid ice model specified: $(ice_model)")
+        end
 
     elseif (sim_mode <: AbstractSlabplanetSimulationMode)
 
@@ -477,7 +421,6 @@ function CoupledSimulation(config_dict::AbstractDict)
             dt = component_dt_dict["dt_ocean"],
             space = boundary_space,
             saveat = saveat,
-            area_fraction = (FT(1) .- land_fraction), ## NB: this ocean fraction includes areas covered by sea ice (unlike the one contained in the cs)
             thermo_params = thermo_params,
             evolving = evolving_ocean,
         )
@@ -598,8 +541,8 @@ function CoupledSimulation(config_dict::AbstractDict)
 
         The concrete steps for proper initialization are:
         =#
-
         # 1. Make sure surface model area fractions sum to 1 everywhere.
+        #  Note that ocean and ice fractions are not accurate until after this call.
         FieldExchanger.update_surface_fractions!(cs)
 
         # 2. Import atmospheric and surface fields into the coupler fields,
diff --git a/src/FieldExchanger.jl b/src/FieldExchanger.jl
@@ -16,7 +16,7 @@ export update_sim!,
     exchange!,
     set_caches!,
     update_surface_fractions!,
-    resolve_ocean_ice_fractions!
+    resolve_area_fractions!
 
 """
     update_surface_fractions!(cs::Interfacer.CoupledSimulation)
@@ -69,9 +69,9 @@ function update_surface_fractions!(cs::Interfacer.CoupledSimulation)
         )
         ocean_fraction = Interfacer.get_field(ocean_sim, Val(:area_fraction))
 
-        # ensure that ocean and ice fractions are consistent
+        # Apply any additional constraints on the ocean and ice fractions if necessary
         if haskey(cs.model_sims, :ice_sim)
-            resolve_ocean_ice_fractions!(ocean_sim, cs.model_sims.ice_sim, land_fraction)
+            resolve_area_fractions!(ocean_sim, cs.model_sims.ice_sim, land_fraction)
         end
     else
         cs.fields.scalar_temp1 .= 0
@@ -84,15 +84,15 @@ function update_surface_fractions!(cs::Interfacer.CoupledSimulation)
 end
 
 """
-    resolve_ocean_ice_fractions!(ocean_sim, ice_sim, land_fraction)
+    resolve_area_fractions!(ocean_sim, ice_sim, land_fraction)
 
 Ensure that the ocean and ice fractions are consistent with each other.
 For most ocean and ice models, this does nothing since the ocean fraction is
 defined as `1 - ice_fraction - land_fraction`. However, some models may have
 additional constraints on the ice and ocean fractions that need to be enforced.
 This function can be extended for such models.
 """
-function resolve_ocean_ice_fractions!(ocean_sim, ice_sim, land_fraction)
+function resolve_area_fractions!(ocean_sim, ice_sim, land_fraction)
     return nothing
 end
 
@@ -209,6 +209,11 @@ end
 Updates the surface component model cache with the current coupler fields
 *besides turbulent fluxes*, which are updated in `update_turbulent_fluxes`.
 
+Note that upwelling longwave and shortwave radiation are not computed here,
+and are expected to be computed internally by the surface model.
+Some component models extend this function and compute the upwelling longwave
+and shortwave radiation in their methods of `update_sim!`.
+
 # Arguments
 - `sim`: [Interfacer.SurfaceModelSimulation] containing a surface model simulation object.
 - `csf`: [NamedTuple] containing coupler fields.
@@ -217,7 +222,6 @@ function update_sim!(sim::Interfacer.SurfaceModelSimulation, csf)
     # radiative fluxes
     Interfacer.update_field!(sim, Val(:SW_d), csf.SW_d)
     Interfacer.update_field!(sim, Val(:LW_d), csf.LW_d)
-    # TODO need to compute SWU, LWU here too
 
     # precipitation
     Interfacer.update_field!(sim, Val(:liquid_precipitation), csf.P_liq)
