Revisit flux climatology

experiments/flux_climatology/flux_climatology.jl

@@ -1,6 +1,6 @@
 using ClimaOcean
 using ClimaOcean.ECCO
-using ClimaOcean.ECCO: all_ECCO_dates
+using Dates
 using Oceananigans
 using Oceananigans.Utils
 using Oceananigans.Fields: ZeroField, location, VelocityFields
@@ -19,61 +19,84 @@ import Oceananigans.Models: timestepper, update_model_field_time_series!
 import ClimaOcean.OceanSeaIceModels: reference_density, heat_capacity
 import Oceananigans.Architectures: on_architecture
 
+ClimaOcean.DataWrangling.dataset_defaults.FloatType = Float64
+
 #####
 ##### A Data structure that holds flux statistics
 #####
 
 struct FluxStatistics{F}
-    avg :: F
+    flux :: F
+    mean :: F
+    meansq :: F
     std :: F
     max :: F
     min :: F
 end
 
 function FluxStatistics(f::Field)
-    a = similar(f)
-    s = similar(f)
-    p = similar(f)
-    m = similar(f)
-
-    fill!(a, 0)
-    fill!(s, 0)
-    fill!(p, 0)
-    fill!(m, 0)
-
-    return FluxStatistics(a, s, p, m)
+    mean = similar(f)
+    meansq = similar(f)
+    std = similar(f)
+    max = similar(f)
+    min = similar(f)
+
+    fill!(mean, 0)
+    fill!(meansq, 0)
+    fill!(std, 0)
+    fill!(max, 0)
+    fill!(min, 0)
+
+    return FluxStatistics(f, mean, meansq, std, max, min)
 end
 
-Adapt.adapt_structure(to, f::FluxStatistics) = FluxStatistics(Adapt.adapt(to, f.avg),
+Adapt.adapt_structure(to, f::FluxStatistics) = FluxStatistics(Adapt.adapt(to, f.flux),
+                                                              Adapt.adapt(to, f.mean),
+                                                              Adapt.adapt(to, f.meansq),
                                                               Adapt.adapt(to, f.std),
                                                               Adapt.adapt(to, f.max),
                                                               Adapt.adapt(to, f.min))
 
-on_architecture(arch, f::FluxStatistics) = FluxStatistics(on_architecture(arch, f.avg),
+on_architecture(arch, f::FluxStatistics) = FluxStatistics(on_architecture(arch, f.flux),
+                                                          on_architecture(arch, f.mean),
+                                                          on_architecture(arch, f.meansq),
                                                           on_architecture(arch, f.std),
                                                           on_architecture(arch, f.max),
                                                           on_architecture(arch, f.min))
 
-function update_stats!(stats::FluxStatistics, flux, Δt, stop_time)
-    grid = flux.grid
+function update_stats!(stats::FluxStatistics, iteration)
+    grid = stats.flux.grid
     arch = architecture(grid)
-    launch!(arch, grid, :xy, _update_stats!, stats, flux, Δt, stop_time)
+    launch!(arch, grid, :xy, _update_stats!, stats, iteration)
+    return nothing
+end
+
+function update_stats!(stats::NamedTuple, iteration)
+    for stat in stats
+        update_stats!(stat, iteration)
+    end
+    return nothing
 end
 
-@kernel function _update_stats!(stats, flux, Δt, stop_time)
+@kernel function _update_stats!(stats, iteration)
     i, j = @index(Global, NTuple)
 
+    inverse_iteration = 1 / (iteration + 1)
+
+    # use iterative averaging via
+    # mean_n = (x1 + ... + xn) / n ->
+    # -> mean_{n+1} = mean_n * (1 - 1/(n+1)) * x_{n+1} / (n+1)
     @inbounds begin
-         f = flux[i, j, 1]
-        stats.avg[i, j, 1] += f * Δt / stop_time
-        stats.std[i, j, 1] += f^2 * Δt / stop_time
+        f = stats.flux[i, j, 1]
+        stats.mean[i, j, 1]   *= 1 - inverse_iteration
+        stats.mean[i, j, 1]   += f * inverse_iteration
+        stats.meansq[i, j, 1] *= 1 - inverse_iteration
+        stats.meansq[i, j, 1] += f^2 * inverse_iteration
         stats.max[i, j, 1] = max(stats.max[i, j, 1], f)
         stats.min[i, j, 1] = min(stats.min[i, j, 1], f)
     end
 end
 
-finalize_std!(f::FluxStatistics) = parent(f.std) .= sqrt.(parent(f.std) .- parent(f.avg).^2)
-
 #####
 ##### A function to compute flux statistics
 #####
@@ -85,37 +108,33 @@ function compute_flux_climatology(earth)
     Jᵀ = FluxStatistics(net_fluxes.T)
     Jˢ = FluxStatistics(net_fluxes.S)
 
-    stats = (; τx, τy, Jᵀ, Jˢ)
+    atmos_ocean_fluxes = earth.model.interfaces.atmosphere_ocean_interface.fluxes
+    Qc = FluxStatistics(atmos_ocean_fluxes.sensible_heat) 
+    Qv = FluxStatistics(atmos_ocean_fluxes.latent_heat)
+    Qu = FluxStatistics(atmos_ocean_fluxes.upwelling_longwave)
+    Qs = FluxStatistics(atmos_ocean_fluxes.downwelling_shortwave)
+    Qℓ = FluxStatistics(atmos_ocean_fluxes.downwelling_longwave)
 
-    function update_flux_stats!(earth)
-        Δt = earth.Δt
-        stop_time = earth.stop_time
-
-        update_stats!(τx, net_fluxes.u, Δt, stop_time)
-        update_stats!(τy, net_fluxes.v, Δt, stop_time)
-        update_stats!(Jᵀ, net_fluxes.T, Δt, stop_time)
-        update_stats!(Jˢ, net_fluxes.S, Δt, stop_time)
+    stats = (; τx, τy, Jᵀ, Jˢ, Qc, Qv, Qu, Qs, Qℓ)
 
+    function update_flux_stats!(earth)
+        iteration = earth.model.clock.iteration
+        update_stats!(stats, iteration)
         return nothing
     end
 
     add_callback!(earth, update_flux_stats!, IterationInterval(1))
 
     run!(earth)
 
-    finalize_std!(τx)
-    finalize_std!(τy)
-    finalize_std!(Jᵀ)
-    finalize_std!(Jˢ)
-
     return stats
 end
 
 #####
 ##### A prescribed ocean...
 #####
 
-struct PrescribedOcean{A, G, C, U, T, F} <: AbstractModel{Nothing}
+struct PrescribedOcean{A, G, C, U, T, F} <: AbstractModel{Nothing, A}
     architecture :: A
     grid :: G
     clock :: Clock{C}
@@ -142,28 +161,43 @@ function PrescribedOcean(timeseries;
 end
 
 # ...with prescribed velocity and tracer fields
-version = ECCO4Monthly()
-dates   = all_ECCO_dates(version)[1:24]
+dataset = ECCO4Monthly()
+arch    = CPU()
+
+start_date = DateTime(1992, 1, 1)
+end_date   = DateTime(1992, 1, 2)
 
-arch = CPU()
+stop_time = Day(end_date - start_date).value * Oceananigans.Units.days 
 
-u = ECCOFieldTimeSeries(:u_velocity,  version; time_indices_in_memory = 13, architecture = arch, dates)
-v = ECCOFieldTimeSeries(:v_velocity,  version; time_indices_in_memory = 13, architecture = arch, dates)
-T = ECCOFieldTimeSeries(:temperature, version; time_indices_in_memory = 13, architecture = arch, dates)
-S = ECCOFieldTimeSeries(:salinity,    version; time_indices_in_memory = 13, architecture = arch, dates)
+time_indices_in_memory = 13
 
-grid = ECCO.ECCO_immersed_grid(arch)
+u_meta = Metadata(:u_velocity;  start_date, end_date, dataset)
+v_meta = Metadata(:v_velocity;  start_date, end_date, dataset)
+T_meta = Metadata(:temperature; start_date, end_date, dataset)
+S_meta = Metadata(:salinity;    start_date, end_date, dataset)
+
+u = FieldTimeSeries(u_meta, arch; time_indices_in_memory)
+v = FieldTimeSeries(v_meta, arch; time_indices_in_memory)
+T = FieldTimeSeries(T_meta, arch; time_indices_in_memory)
+S = FieldTimeSeries(S_meta, arch; time_indices_in_memory)
+
+grid = u.grid
 
 ocean_model = PrescribedOcean((; u, v, T, S); grid)
-ocean = Simulation(ocean_model, Δt=3hour, stop_time=365days)
+ocean = Simulation(ocean_model; Δt=3hours, stop_time)
+
+set!(ocean_model.tracers.T,    first(T_meta))
+set!(ocean_model.tracers.S,    first(S_meta))
+set!(ocean_model.velocities.u, first(u_meta))
+set!(ocean_model.velocities.v, first(v_meta))
 
 #####
 ##### Need to extend a couple of methods
 #####
 
 function time_step!(model::PrescribedOcean, Δt; callbacks=[], euler=true)
     tick!(model.clock, Δt)
-    time = Time(model.clock.time)
+    time = Oceananigans.Units.Time(model.clock.time)
 
     possible_fts = merge(model.velocities, model.tracers)
     time_series_tuple = extract_field_time_series(possible_fts)
@@ -196,36 +230,68 @@ atmosphere = JRA55PrescribedAtmosphere(arch; backend=JRA55NetCDFBackend(1000))
 ##### A prescribed earth...
 #####
 
-earth_model = OceanSeaIceModel(ocean, nothing; atmosphere, radiation = Radiation(ocean_emissivity=0, ocean_albedo=1))
-earth = Simulation(earth_model, Δt=3hours, stop_time=365days)
+earth_model = OceanSeaIceModel(ocean, nothing; atmosphere, radiation = Radiation(arch))
+earth = Simulation(earth_model; Δt=3hours, stop_time)
 
 wall_time = Ref(time_ns())
 
 # Just checking that the ocean evolves as expected
 function progress(sim)
-    ocean = sim.model.ocean
-    u, v, w = ocean.model.velocities
-    T = ocean.model.tracers.T
+    net_fluxes = sim.model.interfaces.net_fluxes.ocean_surface
+    atmos_ocean_fluxes = sim.model.interfaces.atmosphere_ocean_interface.fluxes
 
-    Tmax = maximum(interior(T))
-    Tmin = minimum(interior(T))
+    Q = net_fluxes.Q
+    Qc = atmos_ocean_fluxes.sensible_heat
+    Qv = atmos_ocean_fluxes.latent_heat
 
-    umax = (maximum(abs, interior(u)),
-            maximum(abs, interior(v)))
+    τx = net_fluxes.u
+    τy = net_fluxes.v
+
+    Qmax = maximum(Q)
+    Qmin = minimum(Q)
+    Qcmax = maximum(Qc)
+    Qcmin = minimum(Qc)
+    Qvmax = maximum(Qv)
+    Qvmin = minimum(Qv)
+    τmax = (maximum(abs, τx), maximum(abs, τy))
 
     step_time = 1e-9 * (time_ns() - wall_time[])
 
-    msg = @sprintf("Iter: %d, time: %s, Δt: %s", iteration(sim), prettytime(sim), prettytime(sim.Δt))
-    msg *= @sprintf(", max|u|: (%.2e, %.2e) m s⁻¹, extrema(T): (%.2f, %.2f) ᵒC, wall time: %s",
-                    umax..., Tmax, Tmin, prettytime(step_time))
+    msg = @sprintf("Iter: %d, time: %s", iteration(sim), prettytime(sim))
+    msg *= @sprintf(", max|τ|: (%.2e, %.2e) m² s⁻², extrema(Q): (%.2f, %.2f) W m⁻², extrema(Qc): (%.2f, %.2f) W m⁻², extrema(Qv): (%.2f, %.2f) W m⁻², wall time: %s",
+                    τmax..., Qmin, Qmax, Qcmin, Qcmax, Qvmin, Qvmax, prettytime(step_time))
 
     @info msg
 
     wall_time[] = time_ns()
 end
 
-add_callback!(earth, progress, IterationInterval(10))
+add_callback!(earth, progress, IterationInterval(16))
 
 stats = compute_flux_climatology(earth)
 
-
+Qtecco = Metadata(:net_heat_flux; start_date, end_date, dataset)
+Qcecco = Metadata(:sensible_heat_flux; start_date, end_date, dataset)
+Qvecco = Metadata(:latent_heat_flux; start_date, end_date, dataset)
+Qℓecco = Metadata(:upwelling_longwave; start_date, end_date, dataset)
+Qsecco = Metadata(:downwelling_shortwave; start_date, end_date, dataset)
+
+Qℓecco = FieldTimeSeries(Qℓecco, arch; time_indices_in_memory)
+Qtecco = FieldTimeSeries(Qtecco, arch; time_indices_in_memory)
+Qcecco = FieldTimeSeries(Qcecco, arch; time_indices_in_memory)
+Qvecco = FieldTimeSeries(Qvecco, arch; time_indices_in_memory)
+Qsecco = FieldTimeSeries(Qsecco, arch; time_indices_in_memory)
+
+Q̄tecco = Field{Center, Center, Nothing}(Qtecco.grid)
+Q̄cecco = Field{Center, Center, Nothing}(Qtecco.grid)
+Q̄vecco = Field{Center, Center, Nothing}(Qtecco.grid)
+Q̄ℓecco = Field{Center, Center, Nothing}(Qtecco.grid)
+Q̄secco = Field{Center, Center, Nothing}(Qtecco.grid)
+
+for i in 1:length(Qtecco.times)
+    Q̄tecco .+= Qtecco[i] ./ length(Qtecco.times)
+    Q̄cecco .+= Qcecco[i] ./ length(Qcecco.times)
+    Q̄vecco .+= Qvecco[i] ./ length(Qvecco.times)
+    Q̄ℓecco .+= Qℓecco[i] ./ length(Qℓecco.times)
+    Q̄secco .+= Qsecco[i] ./ length(Qsecco.times)
+end