Enable AveragedSpecifiedTimes to be used in output writers

src/OutputWriters/windowed_time_average.jl

@@ -2,6 +2,7 @@ using Oceananigans.Diagnostics: AbstractDiagnostic
 using Oceananigans.OutputWriters: fetch_output
 using Oceananigans.Utils: AbstractSchedule, prettytime
 using Oceananigans.TimeSteppers: Clock
+using Dates: Period, Second, value
 
 import Oceananigans: run_diagnostic!
 import Oceananigans.Utils: TimeInterval, SpecifiedTimes
@@ -12,11 +13,11 @@ import Oceananigans.Fields: location, indices, set!
 
 Container for parameters that configure and handle time-averaged output.
 """
-mutable struct AveragedTimeInterval <: AbstractSchedule
-    interval :: Float64
-    window :: Float64
+mutable struct AveragedTimeInterval{I, T} <: AbstractSchedule
+    interval :: I
+    window :: I
     stride :: Int
-    first_actuation_time :: Float64
+    first_actuation_time :: T
     actuations :: Int
     collecting :: Bool
 end
@@ -91,30 +92,84 @@ function (sch::AveragedTimeInterval)(model)
     return scheduled
 end
 initialize_schedule!(sch::AveragedTimeInterval, clock) = nothing
-outside_window(sch::AveragedTimeInterval, clock) = clock.time <=  next_actuation_time(sch) - sch.window
+outside_window(sch::AveragedTimeInterval, clock) = clock.time <= next_actuation_time(sch) - sch.window
 end_of_window(sch::AveragedTimeInterval, clock) = clock.time >= next_actuation_time(sch)
 
 TimeInterval(sch::AveragedTimeInterval) = TimeInterval(sch.interval)
 Base.copy(sch::AveragedTimeInterval) = AveragedTimeInterval(sch.interval, window=sch.window, stride=sch.stride)
 
-
-
 """
     mutable struct AveragedSpecifiedTimes <: AbstractSchedule
 
 A schedule for averaging over windows that precede SpecifiedTimes.
 """
-mutable struct AveragedSpecifiedTimes <: AbstractSchedule
-    specified_times :: SpecifiedTimes
-    window :: Float64
+mutable struct AveragedSpecifiedTimes{S, W} <: AbstractSchedule
+    specified_times :: S
+    window :: W
     stride :: Int
     collecting :: Bool
 end
 
+const VaryingWindowAveragedSpecifiedTimes = AveragedSpecifiedTimes{<:Any, <:Vector}
+
 AveragedSpecifiedTimes(specified_times::SpecifiedTimes; window, stride=1) =
     AveragedSpecifiedTimes(specified_times, window, stride, false)
 
-AveragedSpecifiedTimes(times; kw...) = AveragedSpecifiedTimes(SpecifiedTimes(times); kw...)
+AveragedSpecifiedTimes(times; window, kw...) = AveragedSpecifiedTimes(times, window; kw...)
+
+determine_epsilon(eltype) = 0
+determine_epsilon(::Type{T}) where T <: AbstractFloat = eps(T)
+determine_epsilon(::Type{<:Period}) = Second(0)
+
+const NumberTypeWindows = Union{Number, Vector{<:Number}}
+const PeriodTypeWindows = Union{Period, Vector{<:Period}}
+
+
+validate_windows(times, window) = nothing  # Fallback method
+
+function validate_windows(times, window::NumberTypeWindows)
+    tol = 10 * determine_epsilon(eltype(times))
+
+    gaps = diff(vcat(0, times))  # Prepend 0 to check first window against t=0
+    any(gaps .- window .< -tol) && throw(ArgumentError("Averaging windows overlap: some gaps between specified times are less than the window size."))
+end
+
+function validate_windows(times, window::PeriodTypeWindows)
+    if length(times) >= 2
+        window_starts = times .- window
+        prev_window_ends = times[1:end-1]
+        any(window_starts[2:end] .< prev_window_ends) && throw(ArgumentError("Averaging windows overlap: some gaps between specified times are less than the window size."))
+    end
+
+    # Note: We cannot check if the first window extends before the simulation start
+    # because the model clock is not available at construction time
+end
+
+function AveragedSpecifiedTimes(times, window::Vector; kw...)
+    length(window) == length(times) || throw(ArgumentError("When providing a vector of windows, its length $(length(window)) must match the number of specified times $(length(times))."))
+    perm = sortperm(times)
+    sorted_times = times[perm]
+    sorted_window = window[perm]
+
+    # Check for overlapping windows
+    validate_windows(sorted_times, sorted_window)
+
+    return AveragedSpecifiedTimes(SpecifiedTimes(sorted_times); window=sorted_window, kw...)
+end
+
+function AveragedSpecifiedTimes(times, window; kw...)
+    specified_times = SpecifiedTimes(times)
+
+    # Check for overlapping windows (scalar window case)
+    if length(specified_times.times) > 1
+        validate_windows(specified_times.times, window)
+    end
+
+    return AveragedSpecifiedTimes(specified_times; window, kw...)
+end
+
+get_next_window(schedule::VaryingWindowAveragedSpecifiedTimes) = schedule.window[schedule.specified_times.previous_actuation + 1]
+get_next_window(schedule) = schedule.window
 
 function (schedule::AveragedSpecifiedTimes)(model)
     time = model.clock.time
@@ -123,7 +178,7 @@ function (schedule::AveragedSpecifiedTimes)(model)
     next > length(schedule.specified_times.times) && return false
 
     next_time = schedule.specified_times.times[next]
-    window = schedule.window
+    window = get_next_window(schedule)
 
     schedule.collecting || time >= next_time - window
 end
@@ -134,7 +189,8 @@ function outside_window(schedule::AveragedSpecifiedTimes, clock)
     next = schedule.specified_times.previous_actuation + 1
     next > length(schedule.specified_times.times) && return true
     next_time = schedule.specified_times.times[next]
-    return clock.time < next_time - schedule.window
+    window = get_next_window(schedule)
+    return clock.time <= next_time - window
 end
 
 function end_of_window(schedule::AveragedSpecifiedTimes, clock)
@@ -144,22 +200,27 @@ function end_of_window(schedule::AveragedSpecifiedTimes, clock)
     return clock.time >= next_time
 end
 
+TimeInterval(sch::AveragedSpecifiedTimes) = TimeInterval(sch.specified_times.times)
+Base.copy(sch::AveragedSpecifiedTimes) = AveragedSpecifiedTimes(copy(sch.specified_times); window=sch.window, stride=sch.stride)
+
+next_actuation_time(sch::AveragedSpecifiedTimes) = Oceananigans.Utils.next_actuation_time(sch.specified_times)
+
 #####
 ##### WindowedTimeAverage
 #####
 
-mutable struct WindowedTimeAverage{OP, R, S} <: AbstractDiagnostic
+mutable struct WindowedTimeAverage{OP, R, T, S} <: AbstractDiagnostic
                       result :: R
                      operand :: OP
-           window_start_time :: Float64
+           window_start_time :: T
       window_start_iteration :: Int
-    previous_collection_time :: Float64
+    previous_collection_time :: T
                     schedule :: S
                fetch_operand :: Bool
 end
 
-const IntervalWindowedTimeAverage = WindowedTimeAverage{<:Any, <:Any, <:AveragedTimeInterval}
-const SpecifiedWindowedTimeAverage = WindowedTimeAverage{<:Any, <:Any, <:AveragedSpecifiedTimes}
+const IntervalWindowedTimeAverage = WindowedTimeAverage{<:Any, <:Any, <:Any, <:AveragedTimeInterval}
+const SpecifiedWindowedTimeAverage = WindowedTimeAverage{<:Any, <:Any, <:Any, <:AveragedSpecifiedTimes}
 
 stride(wta::IntervalWindowedTimeAverage) = wta.schedule.stride
 stride(wta::SpecifiedWindowedTimeAverage) = wta.schedule.stride
@@ -168,7 +229,7 @@ stride(wta::SpecifiedWindowedTimeAverage) = wta.schedule.stride
     WindowedTimeAverage(operand, model=nothing; schedule)
 
 Returns an object for computing running averages of `operand` over `schedule.window` and
-recurring on `schedule.interval`, where `schedule` is an `AveragedTimeInterval`.
+recurring on `schedule.interval`, where `schedule` is an `AveragedTimeInterval` or `AveragedSpecifiedTimes`.
 During the collection period, averages are computed every `schedule.stride` iteration.
 
 `operand` may be a `Oceananigans.Field` or a function that returns an array or scalar.
@@ -186,9 +247,15 @@ function WindowedTimeAverage(operand, model=nothing; schedule, fetch_operand=tru
         result .= operand
     end
 
-    return WindowedTimeAverage(result, operand, 0.0, 0, 0.0, schedule, fetch_operand)
+    time = isnothing(model) ? get_default_time(schedule) : model.clock.time
+
+    return WindowedTimeAverage(result, operand, time, 0, time, schedule, fetch_operand)
 end
 
+# Helper functions to get default time based on schedule type
+get_default_time(schedule::AveragedTimeInterval) = zero(typeof(schedule.interval))
+get_default_time(schedule::AveragedSpecifiedTimes) = zero(eltype(schedule.specified_times.times))
+
 # Time-averaging doesn't change spatial location
 location(wta::WindowedTimeAverage) = location(wta.operand)
 indices(wta::WindowedTimeAverage) = indices(wta.operand)
@@ -213,12 +280,15 @@ function accumulate_result!(wta, model)
     return accumulate_result!(wta, model.clock, integrand)
 end
 
+period_to_number(p::Period) = value(p)
+period_to_number(n::Number) = n
+
 function accumulate_result!(wta, clock::Clock, integrand=wta.operand)
     # Time increment:
-    Δt = clock.time - wta.previous_collection_time
+    Δt = period_to_number(clock.time - wta.previous_collection_time)
     # Time intervals:
-    T_current = clock.time - wta.window_start_time
-    T_previous = wta.previous_collection_time - wta.window_start_time
+    T_current = period_to_number(clock.time - wta.window_start_time)
+    T_previous = period_to_number(wta.previous_collection_time - wta.window_start_time)
 
     # Accumulate left Riemann sum
     @. wta.result = (wta.result * T_previous + integrand * Δt) / T_current
@@ -261,6 +331,40 @@ function advance_time_average!(wta::WindowedTimeAverage, model)
     return nothing
 end
 
+function advance_time_average!(wta::SpecifiedWindowedTimeAverage, model)
+
+    unscheduled = model.clock.iteration == 0 || outside_window(wta.schedule, model.clock)
+    if !(unscheduled)
+        if !(wta.schedule.collecting)
+            # Zero out result to begin new accumulation window
+            wta.result .= 0
+
+            # Begin collecting window-averaged increments
+            wta.schedule.collecting = true
+
+            wta.window_start_time = next_actuation_time(wta.schedule) - get_next_window(wta.schedule)
+            wta.previous_collection_time = wta.window_start_time
+            wta.window_start_iteration = model.clock.iteration - 1
+            # @info "t $(prettytime(model.clock.time)), next actuation time: $(prettytime(next_actuation_time(wta.schedule))), window $(prettytime(wta.schedule.window))"
+        end
+
+        if end_of_window(wta.schedule, model.clock)
+            accumulate_result!(wta, model)
+            # Save averaging start time and the initial data collection time
+            wta.schedule.collecting = false
+            wta.schedule.specified_times.previous_actuation += 1
+
+        elseif mod(model.clock.iteration - wta.window_start_iteration, stride(wta)) == 0
+            accumulate_result!(wta, model)
+        else
+            # Off stride, so do nothing.
+        end
+
+    end
+    return nothing
+end
+
+
 # So it can be used as a Diagnostic
 run_diagnostic!(wta::WindowedTimeAverage, model) = advance_time_average!(wta, model)
 
@@ -271,8 +375,14 @@ Base.summary(schedule::AveragedTimeInterval) = string("AveragedTimeInterval(",
                                                       "stride=", schedule.stride, ", ",
                                                       "interval=", prettytime(schedule.interval),  ")")
 
+Base.summary(schedule::AveragedSpecifiedTimes) = string("AveragedSpecifiedTimes(",
+                                                        "window=", prettytime(schedule.window), ", ",
+                                                        "stride=", schedule.stride, ", ",
+                                                        "times=", schedule.specified_times,  ")")
+
 show_averaging_schedule(schedule) = ""
 show_averaging_schedule(schedule::AveragedTimeInterval) = string(" averaged on ", summary(schedule))
+show_averaging_schedule(schedule::AveragedSpecifiedTimes) = string(" averaged on ", summary(schedule))
 
 output_averaging_schedule(output::WindowedTimeAverage) = output.schedule
 
@@ -282,6 +392,8 @@ output_averaging_schedule(output::WindowedTimeAverage) = output.schedule
 
 time_average_outputs(schedule, outputs, model) = schedule, outputs # fallback
 
+const AveragedTimeSchedule = Union{AveragedTimeInterval, AveragedSpecifiedTimes}
+
 """
     time_average_outputs(schedule::AveragedTimeInterval, outputs, model, field_slicer)
 
@@ -290,17 +402,16 @@ Wrap each `output` in a `WindowedTimeAverage` on the time-averaged `schedule` an
 Returns the `TimeInterval` associated with `schedule` and a `NamedTuple` or `Dict` of the wrapped
 outputs.
 """
-function time_average_outputs(schedule::AveragedTimeInterval, outputs::Dict, model)
+function time_average_outputs(schedule::AveragedTimeSchedule, outputs::Dict, model)
     averaged_outputs = Dict(name => WindowedTimeAverage(output, model; schedule=copy(schedule))
                             for (name, output) in outputs)
 
     return TimeInterval(schedule), averaged_outputs
 end
 
-function time_average_outputs(schedule::AveragedTimeInterval, outputs::NamedTuple, model)
+function time_average_outputs(schedule::AveragedTimeSchedule, outputs::NamedTuple, model)
     averaged_outputs = NamedTuple(name => WindowedTimeAverage(outputs[name], model; schedule=copy(schedule))
                                   for name in keys(outputs))
 
     return TimeInterval(schedule), averaged_outputs
-end
-
+end

src/Utils/prettytime.jl

@@ -64,3 +64,4 @@ function prettytimeunits(t, longform=true)
 end
 
 prettytime(dt::AbstractTime) = "$dt"
+prettytime(t::Array) = prettytime.(t)

src/Utils/schedules.jl

@@ -242,6 +242,8 @@ function specified_times_str(st)
     return string(str, "]")
 end
 
+Base.copy(st::SpecifiedTimes) = SpecifiedTimes(copy(st.times), st.previous_actuation)
+
 #####
 ##### ConsecutiveIterations
 #####

src/Utils/times_and_datetimes.jl

@@ -29,11 +29,23 @@ end
 @inline add_time_interval(base::AbstractTime, interval::Number, count=1) = base + seconds_to_nanosecond(interval * count)
 @inline add_time_interval(base::AbstractTime, interval::Period, count=1) = base + count * interval
 
+@inline add_time_interval(base, interval::Array{<:Number}, count=1) = interval[count]
+@inline add_time_interval(base, interval::Array{<:Dates.Period}, count=1) = seconds_to_nanosecond(interval[count])
+@inline add_time_interval(base, interval::Array{Dates.DateTime}, count=1) = interval[count]
+
 function period_type(interval::Number)
     FT = Oceananigans.defaults.FloatType
     return FT
 end
 
+function period_type(interval::Array{<:Number})
+    FT = Oceananigans.defaults.FloatType
+    return Array{FT, 1}
+end
+
 period_type(interval::Dates.Period) = typeof(interval)
+period_type(interval) = typeof(interval)
+
 time_type(interval::Number) = typeof(interval)
 time_type(interval::Dates.Period) = Dates.DateTime
+time_type(interval::Array) = eltype(interval)

test/runtests.jl

@@ -96,6 +96,7 @@ CUDA.allowscalar() do
             include("test_implicit_diffusion_diagnostic.jl")
             include("test_output_writers.jl")
             include("test_output_readers.jl")
+            include("test_averaged_specified_times.jl")
             include("test_set_field_time_series.jl")
         end
     end