Preparations to run inertial river routing on GPU

Wflow/Project.toml

@@ -5,13 +5,15 @@ version = "1.0.0-rc2-dev"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BasicModelInterface = "59605e27-edc0-445a-b93d-c09a3a50b330"
 CFTime = "179af706-886a-5703-950a-314cd64e0468"
 CompositionsBase = "a33af91c-f02d-484b-be07-31d278c5ca2b"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
@@ -24,6 +26,7 @@ TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 
 [compat]
 Accessors = "0.1"
+Adapt = "4.3.0"
 Aqua = "0.8"
 BasicModelInterface = "0.1"
 CFTime = "0.1, 0.2"
@@ -33,6 +36,7 @@ DelimitedFiles = "1"
 Downloads = "1"
 Glob = "1.3"
 Graphs = "1.4"
+KernelAbstractions = "0.9.36"
 LoggingExtras = "0.4.6,0.5,1"
 NCDatasets = "0.13.2, 0.14"
 Parameters = "0.12"

Wflow/src/Wflow.jl

@@ -1,8 +1,10 @@
 module Wflow
 
 import BasicModelInterface as BMI
+import KernelAbstractions as KA
 
 using Accessors: @optic, @reset, PropertyLens
+import Adapt: adapt, @adapt_structure, adapt_structure
 using Base.Threads: nthreads
 using CFTime: CFTime, monthday, dayofyear
 using CompositionsBase: decompose
@@ -59,6 +61,14 @@ const VERSION =
     VersionNumber(TOML.parsefile(joinpath(@__DIR__, "..", "Project.toml"))["version"])
 const ROUTING_OPTIONS = (("kinematic-wave", "local-inertial"))
 
+mutable struct Defaults
+    FloatType::DataType
+    Backend::KA.Backend
+end
+
+Defaults(; FloatType=Float64, Backend=KA.CPU()) = Defaults(FloatType, Backend)
+const defaults = Defaults()
+
 mutable struct Clock{T}
     time::T
     iteration::Int

Wflow/src/demand/water_demand.jl

@@ -438,12 +438,13 @@ function Demand(
 end
 
 "Struct to store river allocation model variables"
-@with_kw struct AllocationRiverVariables
-    act_surfacewater_abst::Vector{Float64}        # actual surface water abstraction [mm Δt⁻¹]
-    act_surfacewater_abst_vol::Vector{Float64}    # actual surface water abstraction [m³ Δt⁻¹]
-    available_surfacewater::Vector{Float64}       # available surface water [m³]
-    nonirri_returnflow::Vector{Float64}           # return flow from non irrigation [mm Δt⁻¹] 
+@with_kw struct AllocationRiverVariables{T <: AbstractArray{<:AbstractFloat}}
+    act_surfacewater_abst::T        # actual surface water abstraction [mm Δt⁻¹]
+    act_surfacewater_abst_vol::T    # actual surface water abstraction [m³ Δt⁻¹]
+    available_surfacewater::T       # available surface water [m³]
+    nonirri_returnflow::T           # return flow from non irrigation [mm Δt⁻¹] 
 end
+@adapt_structure AllocationRiverVariables
 
 "Initialize river allocation model variables"
 function AllocationRiverVariables(n::Int)
@@ -459,6 +460,7 @@ end
 @with_kw struct AllocationRiver <: AbstractAllocationModel
     variables::AllocationRiverVariables
 end
+@adapt_structure AllocationRiver
 
 get_nonirrigation_returnflow(model::AllocationRiver) = model.variables.nonirri_returnflow
 get_nonirrigation_returnflow(model::NoAllocationRiver) = 0.0

Wflow/src/io.jl

@@ -1426,8 +1426,14 @@ function read_hq_csv(path)
 end
 
 # these represent the type of the rating curve and specific storage data
-const SH = NamedTuple{(:H, :S), Tuple{Vector{Float64}, Vector{Float64}}}
-const HQ = NamedTuple{(:H, :Q), Tuple{Vector{Float64}, Matrix{Float64}}}
+const SH = NamedTuple{
+    (:H, :S),
+    Tuple{<:AbstractArray{<:AbstractFloat}, <:AbstractArray{<:AbstractFloat}},
+}
+const HQ = NamedTuple{
+    (:H, :Q),
+    Tuple{<:AbstractArray{<:AbstractFloat}, <:AbstractArray{<:AbstractFloat, 2}},
+}
 
 is_increasing(v) = last(v) > first(v)
 

Wflow/src/routing/reservoir.jl

@@ -1,38 +1,61 @@
 "Struct for storing reservoir model parameters"
-@with_kw struct ReservoirParameters
+@with_kw struct ReservoirParameters{
+    T <: AbstractArray{<:AbstractFloat},
+    I <: AbstractArray{Int},
+}
     # type of reservoir storage curve, 1: S = AH, 2: S = f(H) from reservoir data and
     # interpolation
-    storfunc::Vector{Int}
+    storfunc::I
     # type of reservoir rating curve, 1: Q = f(H) from reservoir data and interpolation, 2:
     # General Q = b(H - H₀)ᵉ, 3: Case of Puls Approach Q = b(H - H₀)², 4: Simple reservoir
-    outflowfunc::Vector{Int}
+    outflowfunc::I
     # reservoir area [m²]
-    area::Vector{Float64}
+    area::T
     # index of lower reservoir (linked reservoirs) [-]
-    lower_reservoir_ind::Vector{Int} = fill(0, length(area))
+    lower_reservoir_ind::I = fill(0, length(area))
     # reservoir maximum storage for rating curve types 1 and 4 [m³]
-    maxstorage::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    maxstorage::T = fill(MISSING_VALUE, length(area))
     # water level threshold H₀ [m] below that level outflow is zero
-    threshold::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    threshold::T = fill(MISSING_VALUE, length(area))
     # rating curve coefficient [m3/2 s-1] (if e=3/2)
-    b::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    b::T = fill(MISSING_VALUE, length(area))
     # rating curve exponent [-]
-    e::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    e::T = fill(MISSING_VALUE, length(area))
     # data for storage curve
     sh::Vector{Union{SH, Missing}} = Vector{Union{SH, Missing}}(missing, length(area))
     # data for rating curve
     hq::Vector{Union{HQ, Missing}} = Vector{Union{HQ, Missing}}(missing, length(area))
     # column index of rating curve data hq
-    col_index_hq::Vector{Int} = [1]
+    col_index_hq::I = [1]
     # maximum amount that can be released if below spillway for rating curve type 4 [m³ s⁻¹]
-    maxrelease::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    maxrelease::T = fill(MISSING_VALUE, length(area))
     # minimum (environmental) flow requirement downstream of the reservoir for rating curve
     # type 4 [m³ s⁻¹]
-    demand::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    demand::T = fill(MISSING_VALUE, length(area))
     # target minimum full fraction (of max storage) for rating curve type 4 [-]
-    targetminfrac::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    targetminfrac::T = fill(MISSING_VALUE, length(area))
     # target fraction full (of max storage) for rating curve type 4 [-]
-    targetfullfrac::Vector{Float64} = fill(MISSING_VALUE, length(area))
+    targetfullfrac::T = fill(MISSING_VALUE, length(area))
+end
+
+function adapt_structure(to, from::ReservoirParameters)
+    return ReservoirParameters(
+        adapt(to, from.storfunc),
+        adapt(to, from.outflowfunc),
+        adapt(to, from.area),
+        adapt(to, from.lower_reservoir_ind),
+        adapt(to, from.maxstorage),
+        adapt(to, from.threshold),
+        adapt(to, from.b),
+        adapt(to, from.e),
+        [adapt(to, sh) for sh in from.sh],  # TODO: might not be compatible with GPU computations
+        [adapt(to, hq) for hq in from.hq],
+        adapt(to, from.col_index_hq),
+        adapt(to, from.maxrelease),
+        adapt(to, from.demand),
+        adapt(to, from.targetminfrac),
+        adapt(to, from.targetfullfrac),
+    )
 end
 
 "Initialize reservoir model parameters"
@@ -181,26 +204,27 @@ function ReservoirParameters(dataset::NCDataset, config::Config, network::Networ
 end
 
 "Struct for storing reservoir model variables"
-@with_kw struct ReservoirVariables
+@with_kw struct ReservoirVariables{T <: AbstractArray{<:AbstractFloat}}
     # waterlevel H [m] of reservoir
-    waterlevel::Vector{Float64}
+    waterlevel::T
     # average waterlevel H [m] of reservoir for model timestep Δt
-    waterlevel_av::Vector{Float64} = fill(MISSING_VALUE, length(waterlevel))
+    waterlevel_av::T = fill(MISSING_VALUE, length(waterlevel))
     # reservoir storage [m³]
-    storage::Vector{Float64}
+    storage::T
     # average reservoir storage [m³] for model timestep Δt
-    storage_av::Vector{Float64} = fill(MISSING_VALUE, length(waterlevel))
+    storage_av::T = fill(MISSING_VALUE, length(waterlevel))
     # outflow from reservoir [m³ s⁻¹]
-    outflow::Vector{Float64} = fill(MISSING_VALUE, length(waterlevel))
+    outflow::T = fill(MISSING_VALUE, length(waterlevel))
     # average outflow from reservoir [m³ s⁻¹] for model timestep Δt
-    outflow_av::Vector{Float64} = fill(MISSING_VALUE, length(waterlevel))
+    outflow_av::T = fill(MISSING_VALUE, length(waterlevel))
     # average actual evaporation for reservoir area [mm Δt⁻¹]
-    actevap::Vector{Float64} = fill(MISSING_VALUE, length(waterlevel))
+    actevap::T = fill(MISSING_VALUE, length(waterlevel))
     # fraction full (of max storage) for rating curve type 4 [-]
-    percfull::Vector{Float64} = fill(MISSING_VALUE, length(waterlevel))
+    percfull::T = fill(MISSING_VALUE, length(waterlevel))
     # minimum (environmental) flow released from reservoir for rating curve type 4 [m³ s⁻¹]
-    demandrelease::Vector{Float64} = fill(MISSING_VALUE, length(waterlevel))
+    demandrelease::T = fill(MISSING_VALUE, length(waterlevel))
 end
+@adapt_structure ReservoirVariables
 
 "Initialize reservoir model variables"
 function ReservoirVariables(parameters::ReservoirParameters, waterlevel::Vector{Float64})
@@ -213,11 +237,12 @@ function ReservoirVariables(parameters::ReservoirParameters, waterlevel::Vector{
 end
 
 "Struct for storing reservoir model boundary conditions"
-@with_kw struct ReservoirBC
-    inflow::Vector{Float64}               # inflow into reservoir [m³ s⁻¹] for model timestep Δt
-    precipitation::Vector{Float64}        # average precipitation for reservoir area [mm Δt⁻¹]
-    evaporation::Vector{Float64}          # average potential evaporation for reservoir area [mm Δt⁻¹]
+@with_kw struct ReservoirBC{T <: AbstractArray{<:AbstractFloat}}
+    inflow::T               # inflow into reservoir [m³ s⁻¹] for model timestep Δt
+    precipitation::T        # average precipitation for reservoir area [mm Δt⁻¹]
+    evaporation::T          # average potential evaporation for reservoir area [mm Δt⁻¹]
 end
+@adapt_structure ReservoirBC
 
 "Initialize reservoir model boundary conditions"
 function ReservoirBC(n::Int)
@@ -235,6 +260,7 @@ end
     parameters::ReservoirParameters
     variables::ReservoirVariables
 end
+@adapt_structure Reservoir
 
 "Initialize reservoir model `SimpleReservoir`"
 function Reservoir(dataset::NCDataset, config::Config, network::NetworkReservoir)

Wflow/src/routing/surface_kinwave.jl

@@ -105,15 +105,17 @@ function RiverFlowParameters(dataset::NCDataset, config::Config, domain::DomainR
 end
 
 "Struct for storing river flow model boundary conditions"
-@with_kw struct RiverFlowBC{R}
-    inwater::Vector{Float64}                # Lateral inflow [m³ s⁻¹]
-    inflow::Vector{Float64}                 # External inflow (abstraction/supply/demand) [m³ s⁻¹]
-    actual_external_abstraction_av::Vector{Float64}  # Actual abstraction from external negative inflow [m³ s⁻¹]
-    inflow_reservoir::Vector{Float64}       # inflow reservoir from land part [m³ s⁻¹]
-    abstraction::Vector{Float64}            # Abstraction (computed as part of water demand and allocation) [m³ s⁻¹]
+@with_kw struct RiverFlowBC{T <: AbstractArray{<:AbstractFloat}, R}
+    inwater::T                              # Lateral inflow [m³ s⁻¹]
+    inflow::T                               # External inflow (abstraction/supply/demand) [m³ s⁻¹]
+    actual_external_abstraction_av::T       # Actual abstraction from external negative inflow [m³ s⁻¹]
+    inflow_reservoir::T                     # inflow reservoir from land part [m³ s⁻¹]
+    abstraction::T                          # Abstraction (computed as part of water demand and allocation) [m³ s⁻¹]
     reservoir::R                            # Reservoir model struct of arrays
 end
 
+@adapt_structure RiverFlowBC
+
 "Initialize river flow model boundary conditions"
 function RiverFlowBC(n::Int, reservoir::Union{Reservoir, Nothing})
     bc = RiverFlowBC(;