From 70d645edf5d7afafa4824113ea35a3348d031f65 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Tue, 23 Jan 2024 08:27:59 -0800 Subject: [PATCH] WIP: TimeVaryingInputs2D [skip ci][ci skip] --- src/shared_utilities/FileReader.jl | 121 ++++++++++++++++++ src/shared_utilities/TimeVaryingInputs.jl | 22 +++- .../analytic_time_varying_input.jl | 15 ++- .../interpolating_time_varying_input0d.jl | 27 ++-- .../interpolating_time_varying_input2d.jl | 91 +++++++++++++ src/standalone/Bucket/Bucket.jl | 29 +++-- test/shared_utilities/time_varying_inputs.jl | 17 ++- 7 files changed, 289 insertions(+), 33 deletions(-) create mode 100644 src/shared_utilities/interpolating_time_varying_input2d.jl diff --git a/src/shared_utilities/FileReader.jl b/src/shared_utilities/FileReader.jl index 804b236bf6..1f0e18b917 100644 --- a/src/shared_utilities/FileReader.jl +++ b/src/shared_utilities/FileReader.jl @@ -470,4 +470,125 @@ exist in the DateTime calendar. """ to_datetime(date) = CFTime.reinterpret(DateTime, date) +struct TemporalDataHandler{PDT <: PrescribedDataTemporal, SPACE} + prescibed_data::PDT + space::SPACE +end + +""" + first_time(data_handler::TemporalDataHandler) + +Return the time of the first snapshot in seconds. +""" +function first_time(data_handler::TemporalDataHandler) + first_time_datetime = data_handler.prescibed_data.file_info.all_dates[begin] + date_ref = data_handler.prescribed_data.sim_info.date_ref + return Second(first_time_datetime - date_ref - t_start).value +end + +""" + last_time(data_handler::TemporalDataHandler) + +Return the time of the last snapshot in seconds. +""" +function last_time(data_handler::TemporalDataHandler) + last_time_datetime = data_handler.prescibed_data.file_info.all_dates[end] + date_ref = data_handler.prescribed_data.sim_info.date_ref + return Second(last_time_datetime - date_ref - t_start).value +end + +""" + previous_time(data_handler::TemporalDataHandler, time) + +Return the time in seconds of the snapshot before the given `time`. +""" +function previous_time(data_handler::TemporalDataHandler, time) + sim_info = data_handler.prescibed_data.sim_info + time_to_datetime = to_datetime( + sim_info.date_ref + + Second(round(sim_info.t_start)) + + Second(round(time)), + ) + return searchsortedfirst( + data_handler.prescibed_data.file_info.all_dates, + time_to_datetime, + ) - 1 +end + +""" + previous_time(data_handler::TemporalDataHandler, time) + +Return the time in seconds of the snapshot after the given `time`. +""" +function next_time(data_handler::TemporalDataHandler, time) + sim_info = data_handler.prescibed_data.sim_info + time_to_datetime = to_datetime( + sim_info.date_ref + + Second(round(sim_info.t_start)) + + Second(round(time)), + ) + return searchsortedfirst( + data_handler.prescibed_data.file_info.all_dates + time_to_datetime, + ) +end + +""" + previous_snapshot!(data_handler::TemporalDataHandler, time) + +Return the first data snapshot from `data_handler` before the given `time`. + +`previous_snapshot!` potentially modifies the internal state of `data_handler` and it might be a +very expensive operation. +""" +function previous_snapshot!(data_handler::TemporalDataHandler, time) + # Time in seconds + + # Get the current date from `time` + sim_info = data_handler.prescibed_data.sim_info + sim_date = to_datetime( + sim_info.date_ref + Second(round(sim_info.t_start)) + Second(round(t)), + ) + # Use next date if it's closest to current time + # This maintains `all_dates[date_idx]` <= `sim_date` < `all_dates[date_idx + 1]` + if sim_date >= to_datetime(next_date_in_file(prescibed_data.file_info)) + read_data_fields!( + data_handler.prescibed_data, + sim_date, + data_handler.space, + ) + end + + return prescribed_data.file_state.data_fields[1] +end + +""" + next_snapshot!(data_handler::TemporalDataHandler, time) + +Return the first data snapshot from `data_handler` after the given `time`. + +`next_snapshot!` potentially modifies the internal state of `data_handler` and it might be a +very expensive operation. +""" +function next_snapshot!(data_handler::TemporalDataHandler, time) + # Time in seconds + + # Get the current date from `time` + sim_info = data_handler.prescibed_data.sim_info + sim_date = to_datetime( + sim_info.date_ref + Second(round(sim_info.t_start)) + Second(round(t)), + ) + # Use next date if it's closest to current time + # This maintains `all_dates[date_idx]` <= `sim_date` < `all_dates[date_idx + 1]` + if sim_date >= to_datetime(next_date_in_file(model_albedo.albedo_info)) + read_data_fields!( + data_handler.prescibed_data, + sim_date, + data_handler.space, + ) + end + + return prescribed_data.file_state.data_fields[2] +end + end diff --git a/src/shared_utilities/TimeVaryingInputs.jl b/src/shared_utilities/TimeVaryingInputs.jl index 786ceed192..beef548203 100644 --- a/src/shared_utilities/TimeVaryingInputs.jl +++ b/src/shared_utilities/TimeVaryingInputs.jl @@ -80,19 +80,39 @@ When passing single-site data When a `times` and `vals` are passed, `times` have to be sorted and the two arrays have to have the same length. +======= +When the input is a function, the signature of the function can be `func(time, args...; +kwargs...)`. The function will be called with the additional arguments and keyword arguments +passed to `evaluate!`. This can be used to access the state and the cache and use those to +set the output field. + +For example: +```julia +CO2fromp(time, Y, p) = p.atmos.co2 +input = TimeVaryingInput(CO2fromY) +evaluate!(dest, input, t, Y, p) +``` """ function TimeVaryingInput end """ - evaluate!(dest, input, time) + evaluate!(dest, input, time, args...; kwargs...) Evaluate the `input` at the given `time`, writing the output in-place to `dest`. Depending on the details of `input`, this function might do I/O and communication. + +Extra arguments +================ + +`args` and `kwargs` are used only when the `input` is a non-interpolating function, e.g., +an analytic one. In that case, `args` and `kwargs` are passed down to the function itself. """ function evaluate! end include("analytic_time_varying_input.jl") +include("interpolating_time_varying_inputs.jl") # Shared stuff include("interpolating_time_varying_input0d.jl") +include("interpolating_time_varying_input2d.jl") end diff --git a/src/shared_utilities/analytic_time_varying_input.jl b/src/shared_utilities/analytic_time_varying_input.jl index acd7458d6c..d6cdd73f33 100644 --- a/src/shared_utilities/analytic_time_varying_input.jl +++ b/src/shared_utilities/analytic_time_varying_input.jl @@ -1,15 +1,22 @@ struct AnalyticTimeVaryingInput{F <: Function} <: AbstractTimeVaryingInput - # func here as to be GPU-compatible (e.g., splines are not) + # func here has to be GPU-compatible (e.g., splines are not) func::F end -function TimeVaryingInput(input::Function; method = nothing, device = nothing) +# _kwargs... is needed to seamlessly support the other TimeVaryingInputs. +function TimeVaryingInput(input::Function; _kwargs...) isnothing(method) || @warn "Interpolation method is ignored for analytical functions" return AnalyticTimeVaryingInput(input) end -function evaluate!(dest, input::AnalyticTimeVaryingInput, time) - dest .= input.func(time) +function evaluate!( + dest, + input::AnalyticTimeVaryingInput, + time, + args...; + kwargs..., +) + dest .= input.func(time, args...; kwargs...) return nothing end diff --git a/src/shared_utilities/interpolating_time_varying_input0d.jl b/src/shared_utilities/interpolating_time_varying_input0d.jl index 65762bc295..450f281f21 100644 --- a/src/shared_utilities/interpolating_time_varying_input0d.jl +++ b/src/shared_utilities/interpolating_time_varying_input0d.jl @@ -1,19 +1,3 @@ -import CUDA - -""" - NearestNeighbor - -Return the value corresponding to the point closest to the input time. -""" -struct NearestNeighbor <: AbstractInterpolationMethod end - -""" - LinearInterpolation - -Perform linear interpolation between the two neighboring points. -""" -struct LinearInterpolation <: AbstractInterpolationMethod end - """ InterpolatingTimeVaryingInput0D @@ -67,7 +51,13 @@ function Adapt.adapt_structure(to, itp::InterpolatingTimeVaryingInput0D) ) end -function evaluate!(destination, itp::InterpolatingTimeVaryingInput0D, time) +function evaluate!( + destination, + itp::InterpolatingTimeVaryingInput0D, + time, + args...; + kwargs..., +) time in itp || error("TimeVaryingInput does not cover time $time") if ClimaComms.device(itp.context) isa ClimaComms.CUDADevice CUDA.@cuda evaluate!(parent(destination), itp, time, itp.method) @@ -101,6 +91,7 @@ function TimeVaryingInput( ) end +<<<<<<< HEAD:src/shared_utilities/interpolating_time_varying_input0d.jl """ in(time, itp::InterpolatingTimeVaryingInput0D) @@ -111,6 +102,8 @@ function Base.in(time, itp::InterpolatingTimeVaryingInput0D) end +======= +>>>>>>> a55fa815 (WIP: TimeVaryingInputs2D):src/shared_utilities/interpolating_time_varying_input1d.jl function evaluate!( dest, itp::InterpolatingTimeVaryingInput0D, diff --git a/src/shared_utilities/interpolating_time_varying_input2d.jl b/src/shared_utilities/interpolating_time_varying_input2d.jl new file mode 100644 index 0000000000..2d3c5150dd --- /dev/null +++ b/src/shared_utilities/interpolating_time_varying_input2d.jl @@ -0,0 +1,91 @@ +""" + InterpolatingTimeVaryingInput2D + +The constructor for InterpolatingTimeVaryingInput2D is not supposed to be used directly, unless you +know what you are doing. The constructor does not perform any check and does not take care of +GPU compatibility. It is responsibility of the user-facing constructor TimeVaryingInput() to do so. +""" +struct InterpolatingTimeVaryingInput2D{ + DH <: TemporalDataHandler, + AA <: AbstractArray, + M <: AbstractInterpolationMethod, + CC <: ClimaComms.AbstractCommsContext, + R <: Tuple, +} <: AbstractTimeVaryingInput + """Object that has all the information on how to deal with files, data, and so on. + Having to deal with files, it lives on the CPU.""" + data_handler::DH + + """Interpolation method""" + method::M + + """ClimaComms context""" + context::CC + + """Range of times over which the interpolator is defined. range is always defined on the + CPU. Used by the in() function.""" + range::R +end + +function TimeVaryingInput( + data_handler::TemporalDataHandler, + method = LinearInterpolation(), + context = ClimaComms.context(), +) + range = (first_time(data_handler), last_time(data_handler)) + return InterpolatingTimeVaryingInput2D(data_handler, method, context, range) +end + +function evaluate!( + dest, + itp::InterpolatingTimeVaryingInput2D, + time, + args...; + kwargs..., +) + time in itp || error("TimeVaryingInput does not cover time $time") + evaluate!(parent(dest), itp, time, itp.method) + return nothing +end + +function evaluate!( + dest, + itp::InterpolatingTimeVaryingInput2D, + time, + ::NearestNeighbor, + args...; + kwargs..., +) + t0, t1 = + previous_time(itp.data_handler, time), next_time(itp.data_handler, time) + + # The closest snapshot could be either the previous or the next one + if (time - t0) < (t1 - time) + dest .= previous_snapshot!(itp.data_handler, time) + else + dest .= next_snapshot!(itp.data_handler, time) + end +end + +function evaluate!( + dest, + itp::InterpolatingTimeVaryingInput2D, + time, + ::LinearInterpolation, + args...; + kwargs..., +) + # Linear interpolation is: + # y = y0 + (y1 - y0) * (time - t0) / (t1 - t0) + # + # Define coeff = (time - t0) / (t1 - t0) + # + # y = (1 - coeff) * y0 + coeff * y1 + + t0, t1 = + previous_time(itp.data_handler, time), next_time(itp.data_handler, time) + coeff = (time - t0) / (t1 - t0) + dest .= + (1 - coeff) * previous_snapshot!(itp.data_handler, time) .+ + coeff * next_snapshot!(itp.data_handler, time) +end diff --git a/src/standalone/Bucket/Bucket.jl b/src/standalone/Bucket/Bucket.jl index 810fa43460..596ca4e289 100644 --- a/src/standalone/Bucket/Bucket.jl +++ b/src/standalone/Bucket/Bucket.jl @@ -67,7 +67,7 @@ abstract type AbstractBucketModel{FT} <: AbstractExpModel{FT} end abstract type AbstractLandAlbedoModel{FT <: AbstractFloat} end """ - BulkAlbedoFunction{FT, F <: FUnction} <: AbstractLandAlbedoModel + BulkAlbedoFunction{FT, F <: Function} <: AbstractLandAlbedoModel An albedo model where the albedo of different surface types is specified. Snow albedo is treated as constant across snow @@ -369,7 +369,7 @@ end Updates the spatially-varying but constant in time surface albedo stored in the auxiliary vector `p` in place, according to the -passed function of latitute and longitude stored in `albedo.α_sfc`. +passed function of latitude and longitude stored in `albedo.α_sfc`. """ function set_initial_parameter_field!( albedo::BulkAlbedoFunction{FT}, @@ -567,15 +567,22 @@ function make_update_aux(model::BucketModel{FT}) where {FT} p.bucket.R_n .= net_radiation(model.radiation, model, Y, p, t) # Surface albedo - p.bucket.α_sfc .= - next_albedo(model.parameters.albedo, model.parameters, Y, p, t) + next_albedo!( + p.bucket.α_sfc, + model.parameters.albedo, + model.parameters, + Y, + p, + t, + ) end return update_aux! end """ - next_albedo(model_albedo::Union{BulkAlbedoFunction{FT}, BulkAlbedoStatic{FT}}, - parameters, Y, p, t) + next_albedo!(next_α_sfc, + model_albedo::Union{BulkAlbedoFunction{FT}, BulkAlbedoStatic{FT}}, + parameters, Y, p, t) Update the surface albedo for time `t`. These albedo model types aren't explicitly dependent on `t`, but depend on quantities which may change over time. @@ -584,7 +591,8 @@ The albedo is calculated by linearly interpolating between the albedo of snow and of the surface, based on the snow water equivalent `S` relative to the parameter `S_c`. The linear interpolation is taken from Lague et al 2019. """ -function next_albedo( +function next_albedo!( + next_α_sfc, model_albedo::Union{BulkAlbedoFunction{FT}, BulkAlbedoStatic{FT}}, parameters, Y, @@ -595,16 +603,17 @@ function next_albedo( (; σS_c) = parameters α_sfc = p.bucket.α_sfc σS = Y.bucket.σS - return @. ((1 - σS / (σS + σS_c)) * α_sfc + σS / (σS + σS_c) * α_snow) + @. next_α_sfc = ((1 - σS / (σS + σS_c)) * α_sfc + σS / (σS + σS_c) * α_snow) end """ - next_albedo(model_albedo::BulkAlbedoTemporal{FT}, parameters, Y, p, t) + next_albedo!(next_α_sfc, model_albedo::BulkAlbedoTemporal{FT}, parameters, Y, p, t) Update the surface albedo for time t. For a file containing albedo information over time, this reads in the value for time t. """ -function next_albedo( +function next_albedo!( + next_α_sfc, model_albedo::BulkAlbedoTemporal{FT}, parameters, Y, diff --git a/test/shared_utilities/time_varying_inputs.jl b/test/shared_utilities/time_varying_inputs.jl index 153693439f..bec8ed631b 100644 --- a/test/shared_utilities/time_varying_inputs.jl +++ b/test/shared_utilities/time_varying_inputs.jl @@ -28,9 +28,16 @@ context = ClimaComms.SingletonCommsContext(device) TimeVaryingInputs.evaluate!(column_field, input, 10.0) @test Array(parent(column_field))[1] == fun(10.0) + + # Check with args and kwargs + fun2 = (x, y; z) -> 2x * y * z + input2 = TimeVaryingInputs.TimeVaryingInput(fun2) + + TimeVaryingInputs.evaluate!(column_field, input2, 10.0, 20.0; z = 30.0) + @test Array(parent(column_field))[1] == fun2(10.0, 20.0; z = 30.0) end -@testset "Temporal TimeVaryingInput 1D" begin +@testset "InteprolatingTimeVaryingInput0D" begin # Check times not sorted xs = [1.0, 0.0] ys = [1.0, 2.0] @@ -102,3 +109,11 @@ end end end end + +@testset "InteprolatingTimeVaryingInput2D" begin + for FT in (Float32, Float64) + radius, depth, nelements, npolynomial = FT(100), FT(3.5), (2, 10), 2 + domain = SphericalShell(; radius, depth, nelements, npolynomial) + + end +end