Merge pull request #131 from JuliaGeodynamics/bk-waterflow

Waterflow routing algorithm

.typos.toml

@@ -2,6 +2,7 @@
 MOR = "MOR"
 dum = "dum"
 Shepard = "Shepard"
+nin = "nin"
 
 [files]
 extend-exclude = ["tutorials/*.pvsm"]

Project.toml

@@ -1,14 +1,15 @@
 name = "GeophysicalModelGenerator"
 uuid = "3700c31b-fa53-48a6-808a-ef22d5a84742"
 authors = ["Boris Kaus", "Marcel Thielmann"]
-version = "0.7.4"
+version = "0.7.5"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 FFMPEG = "c87230d0-a227-11e9-1b43-d7ebe4e7570a"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+GDAL_jll = "a7073274-a066-55f0-b90d-d619367d196c"
 GeoParams = "e018b62d-d9de-4a26-8697-af89c310ae38"
 Geodesy = "0ef565a4-170c-5f04-8de2-149903a85f3d"
 GeometryBasics = "5c1252a2-5f33-56bf-86c9-59e7332b4326"
@@ -26,6 +27,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+WhereTheWaterFlows = "ea906314-1493-4d22-a0af-f886a20c9fba"
 WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
 
 [weakdeps]
@@ -43,8 +45,9 @@ Colors = "0.9 - 0.12"
 Downloads = "1"
 FFMPEG = "0.4"
 FileIO = "1"
+GDAL_jll = "300.900.0 - 301.900.0"
 GLMakie = "0.8, 0.9, 0.10"
-GMT = "1"
+GMT = "1.0 - 1.14"
 GeoParams = "0.2 - 0.6"
 Geodesy = "1"
 GeometryBasics = "0.1 - 0.4"
@@ -69,4 +72,4 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "GMT", "StableRNGs","GridapGmsh"]
+test = ["Test", "GMT", "StableRNGs", "GridapGmsh"]

ext/GLMakie_Visualisation.jl

@@ -3,7 +3,7 @@ module GLMakie_Visualisation
 
 using Statistics
 using GeophysicalModelGenerator: lonlatdepth_grid, GeoData, CartData, km, AbstractGeneralGrid
-import GeophysicalModelGenerator: visualise
+import GeophysicalModelGenerator: visualise, ustrip
 
 # We do not check `isdefined(Base, :get_extension)` as recommended since
 # Julia v1.9.0 does not load package extensions when their dependency is
@@ -14,7 +14,9 @@ else
     using ..GLMakie
 end
 
-export visualise
+import GLMakie: heatmap!, heatmap
+
+export visualise, heatmap, heatmap!
 
 println("Loading GLMakie extensions for GMG")
 
@@ -271,5 +273,112 @@ function visualise(Data::AbstractGeneralGrid; Topography=nothing, Topo_range=not
     return nothing
 end
 
+"""
+    heatmap(x::GeoData, args...; field=:Topography, kwargs...)
+heatmap for a 2D GeoData object (surface)
+"""
+function heatmap(x::GeoData, args...; field=:Topography, kwargs...)
+    @assert size(x.depth.val,3)==1
+    @assert hasfield(typeof(x.fields), field)
+
+    heatmap(x.lon.val[:,1], x.lat.val[1,:], ustrip.(x.fields[field][:,:,1]), args...; kwargs...)
+
+end
+
+"""
+    heatmap(x::GeoData, a::Array{_T,N}, args...; kwargs...)
+in-place heatmap for a 2D GeoData object (surface) with an array `a`. 
+"""
+function heatmap(x::GeoData, a::Array{_T,N}, args...; kwargs...) where{_T,N}
+    @assert size(x.depth.val,3)==1
+
+    if N==3
+        heatmap(x.lon.val[:,1], x.lat.val[1,:], ustrip.(a[:,:,1]), args...; kwargs...)
+    elseif N==2
+        heatmap(x.lon.val[:,1], x.lat.val[1,:], ustrip.(a), args...; kwargs...)
+    end
+
+end
+
+"""
+    heatmap(x::CartData, args...; field=:Topography, kwargs...)
+heatmap for a 2D CartData object (surface)
+"""
+function heatmap(x::CartData, args...; field=:Topography, kwargs...)
+    @assert size(x.z.val,3)==1
+    @assert hasfield(typeof(x.fields), field)
+
+    heatmap(x.x.val[:,1], x.y.val[1,:], ustrip.(x.fields[field][:,:,1]), args...; kwargs...)
+
+end
+
+"""
+    heatmap(x::CartData, a::Array{_T,N}, args...; kwargs...)
+in-place heatmap for a 2D CartData object (surface) with an array `a`. 
+"""
+function heatmap(x::CartData, a::Array{_T,N}, args...; kwargs...) where{_T,N}
+    @assert size(x.z.val,3)==1
+
+    if N==3
+        heatmap(x.x.val[:,1], x.y.val[1,:], ustrip.(a[:,:,1]), args...; kwargs...)
+    elseif N==2
+        heatmap(x.x.val[:,1], x.y.val[1,:], ustrip.(a), args...; kwargs...)
+    end
+
+end
+
+
+"""
+    heatmap!(x::GeoData, args...; field=:Topography, kwargs...)
+in-place heatmap for a 2D GeoData object (surface), 
+"""
+function heatmap!(x::GeoData, args...; field=:Topography, kwargs...)
+    @assert size(x.z.val,3)==1
+    @assert hasfield(typeof(x.fields), field)
+
+    heatmap(x.lon.val[:,1], x.lat.val[1,:], ustrip.(x.fields[field][:,:,1]), args...; kwargs...)
+
+end
+
+"""
+    heatmap!(x::GeoData, a::Array{_T,N}, args...; kwargs...)
+in-place heatmap for a 2D GeoData object (surface) with an array `a`. 
+"""
+function heatmap!(x::GeoData, a::Array{_T,N}, args...; kwargs...) where{_T,N}
+    @assert size(x.z.val,3)==1
+
+    if N==3
+        heatmap!(x.lon.val[:,1], x.lat.val[1,:], ustrip.(a[:,:,1]), args...; kwargs...)
+    elseif N==2
+        heatmap!(x.lon.val[:,1], x.lat.val[1,:], ustrip.(a), args...; kwargs...)
+    end
+
+end
+
+"""
+    heatmap!(x::CartData, args...; field=:Topography, kwargs...)
+in-place heatmap for a 2D CartData object (surface)
+"""
+function heatmap!(x::CartData, args...; field=:Topography, kwargs...)
+    @assert size(x.z.val,3)==1
+    @assert hasfield(typeof(x.fields), field)
+
+    heatmap!(x.x.val[:,1], x.y.val[1,:], ustrip.(x.fields[field][:,:,1]), args...; kwargs...)
+end
+
+"""
+    heatmap!(x::CartData, a::Array{_T,N}, args...; kwargs...)
+in-place heatmap for a 2D CartData object (surface) with an array `a`. 
+"""
+function heatmap!(x::CartData, a::Array{_T,N}, args...; kwargs...) where{_T,N}
+    @assert size(x.z.val,3)==1
+
+    if N==3
+        heatmap!(x.x.val[:,1], x.y.val[1,:], ustrip.(a[:,:,1]), args...; kwargs...)
+    elseif N==2
+        heatmap!(x.x.val[:,1], x.y.val[1,:], ustrip.(a[:,:]), args...; kwargs...)
+    end
+
+end
 
 end

src/GeophysicalModelGenerator.jl

@@ -48,6 +48,7 @@ include("event_counts.jl")
 include("surface_functions.jl")
 include("movies_from_pics.jl")
 include("sea_lvl.jl")
+include("WaterFlow.jl")
 
 # Add optional routines (only activated when the packages are loaded)
 

src/WaterFlow.jl

@@ -0,0 +1,102 @@
+# These are routes to perform waterflow calculations (upstream area etc.) on a DEM
+
+using WhereTheWaterFlows
+import WhereTheWaterFlows: waterflows
+
+export waterflows
+
+"""
+    dlon, dlat = spacing(lon,lat)
+
+Computes the spacing with central differences
+"""
+function spacing(lon,lat)
+    dlon = zeros(size(lon.val)[1:2])
+    dlat = zeros(size(lat.val)[1:2])
+    @views dlon[2:end-1,:] = (lon.val[3:end,:,1] - lon.val[1:end-2,:,1])/2
+     dlon[1,:] = dlon[2,:]
+     dlon[end,:] = dlon[end-1,:]
+     dlat[:,2:end-1] = (lat.val[:,3:end,1] - lat.val[:,1:end-2,1])/2
+     dlat[:,1] = dlat[:,2]
+     dlat[:,end] = dlat[:,end-1]
+
+    return dlon, dlat
+end
+
+"""
+    area_m2 = cell_area(Topo::GeoData)
+Returns the cell area for a Topographic dataset in m² (required for upstream area calculation)
+"""
+function cell_area(Topo::GeoData)
+
+    proj = ProjectionPoint(Lon=mean(Topo.lon.val[:]), Lat=mean(Topo.lat.val[:]))
+    Topo_cart = convert2CartData(Topo, proj)
+    dx, dy = spacing(Topo_cart.x, Topo_cart.y)
+
+    area_m2 = dx.*dy*1e6
+    return area_m2
+end
+
+
+"""
+    Topo_water, sinks, pits, bnds  = waterflows(Topo::GeoData; 
+        flowdir_fn=WhereTheWaterFlows.d8dir_feature, feedback_fn=nothing, drain_pits=true, bnd_as_sink=true,
+        rainfall = nothing,
+        minsize=300)
+
+Takes a GMG GeoData object of a topographic map and routes water through the grid. Optionally,
+you can specify `rainfall` in which case we accumulate the rain as specified in this 2D array instead of the cellarea. 
+This allows you to, for example, sum, up water if you have variable rainfall in the area.
+The other options are as in the `waterflows` function of the package `WhereTheWaterFlows`.
+
+Example
+===
+```julia
+# Download some topographic data
+julia> Topo = import_topo([6.5,7.3,50.2,50.6], file="@earth_relief_03s");
+
+# Flow the water through the area:
+julia> Topo_water, sinks, pits, bnds  = waterflows(Topo)
+julia> Topo_water
+GeoData 
+  size      : (961, 481, 1)
+  lon       ϵ [ 6.5 : 7.3]
+  lat       ϵ [ 50.2 : 50.59999999999999]
+  depth     ϵ [ 0.045 : 0.724]
+  fields    : (:Topography, :area, :slen, :dir, :nout, :nin, :c)
+
+```
+
+"""
+function waterflows(Topo::GeoData, flowdir_fn= WhereTheWaterFlows.d8dir_feature; feedback_fn=nothing, drain_pits=true, bnd_as_sink=true, rainfall=nothing, minsize=300) 
+
+    cellarea = cell_area(Topo)
+    cellarea_m2 = cellarea
+    if !isnothing(rainfall)
+        @assert typeof(rainfall) == Array{Float64,2}
+        cellarea = rainfall
+    end
+
+    dem = Topo.depth.val[:,:,1]
+
+    ni      = size(Topo.depth.val)
+    area    = zeros(ni)
+    slen    = zeros(Int64, ni)
+    dir     = zeros(Int8, ni)
+    nout    = zeros(Int8, ni)
+    nin     = zeros(Int8, ni)
+    c       = zeros(Int64, ni)
+
+    area[:,:,1], slen[:,:,1], dir[:,:,1], nout[:,:,1], nin[:,:,1], sinks, pits, c[:,:,1], bnds = waterflows(dem, cellarea, flowdir_fn;
+                        feedback_fn=feedback_fn, drain_pits=drain_pits, bnd_as_sink=bnd_as_sink)
+
+    catchment_large =  prune_catchments(c, minsize; val=0)
+    largest_catchment = catchment_large .== maximum(catchment_large)
+    largest_area = copy(area)
+    largest_area[.!largest_catchment] .= NaN
+
+    log10_area = log10.(area)
+
+    Topo_water =  addfield(Topo,(;area, slen, dir, nout, nin, c, cellarea_m2, catchment_large, log10_area, largest_catchment, largest_area))                        
+    return Topo_water, sinks, pits, bnds 
+end

test/runtests.jl

@@ -73,6 +73,10 @@ end
 end
 
 
+@testset "Waterflow" begin
+    include("test_WaterFlow.jl")
+end
+
 # Cleanup
 foreach(rm, filter(endswith(".vts"), readdir()))
 foreach(rm, filter(endswith(".vtu"), readdir()))

test/test_WaterFlow.jl

@@ -0,0 +1,20 @@
+using Test, GMT
+
+
+# Download some topographic data
+Topo = import_topo([6.5,7.3,50.2,50.6], file="@earth_relief_03s");
+
+# Flow the water through the area:
+Topo_water, sinks, pits, bnds  = waterflows(Topo)
+
+@test maximum(Topo_water.fields.area) ≈ 9.309204547276944e8
+@test sum(Topo_water.fields.c) == 834501044
+@test sum(Topo_water.fields.nin) == 459361
+@test sum(Topo_water.fields.dir) == 2412566
+
+# With rain in m3/s per cell
+rainfall = ones(size(Topo.lon.val[:,:,1]))*1e-3 # 2D array with rainfall per cell area
+Topo_water1, sinks, pits, bnds  = waterflows(Topo, rainfall=rainfall)
+
+@test maximum(Topo_water1.fields.area) ≈ 169.79800000000208
+