initial PrescribedOceanSim

CliMA · Jan 18, 2025 · a0fabee · a0fabee
1 parent c3d396b
commit a0fabee
Show file tree

Hide file tree

Showing 2 changed files with 123 additions and 41 deletions.
diff --git a/experiments/ClimaEarth/components/ocean/prescr_ocean.jl b/experiments/ClimaEarth/components/ocean/prescr_ocean.jl
@@ -0,0 +1,116 @@
+# import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput, evaluate!
+# import ClimaUtilities.ClimaArtifacts: @clima_artifact
+import Thermodynamics as TD
+import ClimaCoupler: Checkpointer, FluxCalculator, Interfacer, Utilities
+
+"""
+    PrescribedOceanSimulation{P, Y, D, I}
+
+Sea surface temperature (SST) is prescribed and read in directly from a file
+at each timestep.
+
+Ocean roughness follows the https://github.com/NOAA-GFDL/ice_param/blob/main/ocean_rough.F90#L47.
+
+The cache is expected to contain the following variables:
+- `T_sfc` (surface temperature [K])
+- `ρ_sfc` (surface air density [kg / m3])
+- `z0m` (roughness length for momentum [m])
+- `z0b` (roughness length for tracers [m])
+- `beta` (evaporation scaling factor)
+- `α_direct` (direct albedo)
+- `α_diffuse` (diffuse albedo)
+- `area_fraction` (fraction of the grid cell covered by the ocean)
+- `phase` (phase of the water used to calculate surface humidity)
+- `thermo_params` (thermodynamic parameters)
+- `SST_timevaryinginput` (TimeVaryingInput object containing SST data)
+"""
+struct PrescribedOceanSimulation{I} <: Interfacer.SurfaceStub
+    cache::I
+end
+Interfacer.name(::PrescribedOceanSimulation) = "PrescribedOceanSimulation"
+
+"""
+    PrescribedOceanSimulation(
+        ::Type{FT},
+        space,
+        date0,
+        area_fraction,
+        thermo_params,
+        comms_ctx;
+        z0m = FT(5.8e-5),
+        z0b = FT(5.8e-5),
+        beta = FT(1),
+        α_direct_val = FT(0.06),
+        α_diffuse_val = FT(0.06),
+        phase = TD.Liquid(),
+    )
+
+Initialize the `PrescribedOceanSimulation` object with all required cache fields,
+and reading in prescribed SST data.
+
+"""
+function PrescribedOceanSimulation(
+    ::Type{FT},
+    space,
+    date0,
+    area_fraction,
+    thermo_params,
+    comms_ctx;
+    z0m = FT(5.8e-5),
+    z0b = FT(5.8e-5),
+    beta = FT(1),
+    α_direct_val = FT(0.06),
+    α_diffuse_val = FT(0.06),
+    phase = TD.Liquid(),
+) where {FT}
+    # Read in initial SST data
+    sst_data = try
+        joinpath(@clima_artifact("historical_sst_sic", comms_ctx), "MODEL.SST.HAD187001-198110.OI198111-202206.nc")
+    catch error
+        @warn "Using lowres SST. If you want the higher resolution version, you have to obtain it from ClimaArtifacts"
+        joinpath(
+            @clima_artifact("historical_sst_sic_lowres", comms_ctx),
+            "MODEL.SST.HAD187001-198110.OI198111-202206_lowres.nc",
+        )
+    end
+
+    SST_timevaryinginput = TimeVaryingInput(
+        sst_data,
+        "SST",
+        space,
+        reference_date = date0,
+        file_reader_kwargs = (; preprocess_func = (data) -> data + FT(273.15),), ## convert to Kelvin
+    )
+
+    SST_init = zeros(space)
+    evaluate!(SST_init, SST_timevaryinginput, t_start)
+
+    # Create the cache
+    cache = (;
+        T_sfc = SST_init,
+        ρ_sfc = zeros(space),
+        z0m = z0m,
+        z0b = z0b,
+        beta = beta,
+        α_direct = ones(space) .* α_direct_val,
+        α_diffuse = ones(space) .* α_diffuse_val,
+        area_fraction = area_fraction,
+        phase = phase,
+        thermo_params = thermo_params,
+        SST_timevaryinginput = SST_timevaryinginput,
+    )
+    return PrescribedOceanSimulation(cache)
+end
+
+"""
+    step!(sim::PrescribedOceanSimulation, t)
+
+Update the cached surface temperature field using the prescribed data
+at each timestep.
+
+This is the only function we extend explicitly for the `PrescribedOceanSimulation`;
+all other functions are inherited from `SurfaceStub`.
+"""
+function step!(sim::PrescribedOceanSimulation, t)
+    evaluate!(sim.cache.T_sfc, sim.cache.SST_timevaryinginput, t)
+end
diff --git a/experiments/ClimaEarth/run_amip.jl b/experiments/ClimaEarth/run_amip.jl
@@ -168,15 +168,9 @@ tspan = (t_start, t_end)
 #=
 ## Data File Paths
 =#
-sst_data, sic_data = try
-    joinpath(@clima_artifact("historical_sst_sic", comms_ctx), "MODEL.SST.HAD187001-198110.OI198111-202206.nc"),
+sic_data = try
     joinpath(@clima_artifact("historical_sst_sic", comms_ctx), "MODEL.ICE.HAD187001-198110.OI198111-202206.nc")
 catch error
-    @warn "Using lowres sst 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.SST.HAD187001-198110.OI198111-202206_lowres.nc",
-    ),
     joinpath(
         @clima_artifact("historical_sst_sic_lowres", comms_ctx),
         "MODEL.ICE.HAD187001-198110.OI198111-202206_lowres.nc",
@@ -276,31 +270,8 @@ if sim_mode <: AMIPMode
     )
 
     ## ocean stub
-    SST_timevaryinginput = TimeVaryingInput(
-        sst_data,
-        "SST",
-        boundary_space,
-        reference_date = date0,
-        file_reader_kwargs = (; preprocess_func = (data) -> data + FT(273.15),), ## convert to Kelvin
-    )
-
-    SST_init = zeros(boundary_space)
-    evaluate!(SST_init, SST_timevaryinginput, t_start)
-
-    ocean_sim = Interfacer.SurfaceStub((;
-        T_sfc = SST_init,
-        ρ_sfc = zeros(boundary_space),
-        # ocean roughness follows GFDL model
-        # (https://github.com/NOAA-GFDL/ice_param/blob/main/ocean_rough.F90#L47)
-        z0m = FT(5.8e-5),
-        z0b = FT(5.8e-5),
-        beta = FT(1),
-        α_direct = ones(boundary_space) .* FT(0.06),
-        α_diffuse = ones(boundary_space) .* FT(0.06),
-        area_fraction = (FT(1) .- land_area_fraction),
-        phase = TD.Liquid(),
-        thermo_params = thermo_params,
-    ))
+    ocean_sim =
+        PrescribedOceanSimulation(FT, boundary_space, date0, (FT(1) .- land_area_fraction), thermo_params, comms_ctx)
 
     ## sea ice model
     SIC_timevaryinginput = TimeVaryingInput(
@@ -340,12 +311,8 @@ if sim_mode <: AMIPMode
     evaluate!(CO2_init, CO2_timevaryinginput, t_start)
     CO2_field = Interfacer.update_field!(atmos_sim, Val(:co2), CO2_init)
 
-    mode_specifics = (;
-        type = sim_mode,
-        SST_timevaryinginput = SST_timevaryinginput,
-        SIC_timevaryinginput = SIC_timevaryinginput,
-        CO2_timevaryinginput = CO2_timevaryinginput,
-    )
+    mode_specifics =
+        (; type = sim_mode, SIC_timevaryinginput = SIC_timevaryinginput, CO2_timevaryinginput = CO2_timevaryinginput)
     Utilities.show_memory_usage()
 
 elseif (sim_mode <: AbstractSlabplanetSimulationMode) && !(sim_mode <: SlabplanetEisenmanMode)
@@ -400,7 +367,7 @@ elseif (sim_mode <: AbstractSlabplanetSimulationMode) && !(sim_mode <: Slabplane
         thermo_params = thermo_params,
     ))
 
-    mode_specifics = (; type = sim_mode, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
+    mode_specifics = (; type = sim_mode, SIC_timevaryinginput = nothing)
     Utilities.show_memory_usage()
 
 elseif sim_mode <: SlabplanetEisenmanMode
@@ -447,7 +414,7 @@ elseif sim_mode <: SlabplanetEisenmanMode
         thermo_params = thermo_params,
     )
 
-    mode_specifics = (; type = sim_mode, SST_timevaryinginput = nothing, SIC_timevaryinginput = nothing)
+    mode_specifics = (; type = sim_mode, SIC_timevaryinginput = nothing)
     Utilities.show_memory_usage()
 end
 
@@ -679,7 +646,6 @@ function solve_coupler!(cs)
         cs.dates.date[] = TimeManager.current_date(cs, t)
 
         if cs.mode.type <: AMIPMode
-            evaluate!(Interfacer.get_field(ocean_sim, Val(:surface_temperature)), cs.mode.SST_timevaryinginput, t)
             evaluate!(Interfacer.get_field(ice_sim, Val(:area_fraction)), cs.mode.SIC_timevaryinginput, t)
 
             # TODO: get_field with :co2 is not implemented, so this is a little awkward