add Wflow ext

.gitignore

@@ -52,3 +52,5 @@ backup/
 
 *.html
 *.py
+
+.ipynb_checkpoints

Project.toml

@@ -28,12 +28,14 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 GeoDataFrames = "62cb38b5-d8d2-4862-a48e-6a340996859f"
 NaNStatistics = "b946abbf-3ea7-4610-9019-9858bfdeaf2d"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Wflow = "d48b7d99-76e7-47ae-b1d5-ff0c1cf9a818"
 
 [extensions]
 IpaperArchGDALExt = "ArchGDAL"
 IpaperNaNExt = "NaNStatistics"
-IpaperPlotExt = "Plots"
 IpaperSlopeExt = "Distributions"
+IpaperPlotExt = "Plots"
+IpaperWflowExt = "Wflow"
 
 [compat]
 CategoricalArrays = "0.10"

ext/IpaperArchGDALExt/IpaperArchGDALExt.jl

@@ -24,6 +24,7 @@ include("gdal_polygonize.jl")
 export gdal_polygonize
 export nband, nlayer
 export bandnames, set_bandnames
+export gdal_nodata
 export gdal_info, ogr_info
 export write_gdal
 

ext/IpaperArchGDALExt/gdal_basic.jl

@@ -77,6 +77,19 @@ function set_bandnames(f, bandnames)
   nothing
 end
 
+# gdal_nodata(band::ArchGDAL.GDALRasterBand) = ArchGDAL.GDAL.gdalgetrasternodatavalue(band.ptr)
+function gdal_nodata(f)
+  n = nband(f)
+  gdal_open(f, GDAL.GA_Update) do ds
+    map(iband -> begin
+        band = GDAL.gdalgetrasterband(ds, iband)
+        nodata = GDAL.gdalgetrasternodatavalue(band, Ref(Cint(0)))
+        Type = ArchGDAL.pixeltype(band)
+        Type(nodata)
+      end, 1:n)
+  end
+end
+
 function gdal_info(f)
   run(`$(gdalinfo_path()) $f`)
   nothing

ext/IpaperWflowExt.jl

@@ -0,0 +1,49 @@
+export IpaperWflowExt
+module IpaperWflowExt
+
+
+import Ipaper: flowdir_drop_missing!
+import Wflow.Graphs: outneighbors, topological_sort_by_dfs
+import Wflow: stream_order, subbasins, fillnodata_upstream
+stream_link = subbasins
+# stream_link(g, streamorder, toposort, min_sto)
+
+
+function reverse_index(z::AbstractVector{T}, inds, inds_rev; mv::T) where {T}
+  R = zeros(T, size(inds_rev)) .+ mv
+  @inbounds for (i, I) in enumerate(inds)
+    R[I] = z[i]
+  end
+  R
+end
+
+"""
+- `min_sto`: 过小的sto不认为是stream link
+"""
+function SubBasins(A_fdir::AbstractMatrix; min_sto::Int=4)
+  ldd_mv = UInt8(99)
+  # _drop_missing(A) = drop_missing(A, ldd_mv)
+  flowdir_drop_missing!(A_fdir)
+  # ldd_2d = Array(nc["wflow_ldd"])[:, end:-1:1] |> _drop_missing
+
+  inds, inds_rev = Wflow.active_indices(A_fdir, ldd_mv)
+  ldd = A_fdir[inds]
+
+  g = Wflow.flowgraph(ldd, inds, Wflow.pcr_dir)
+  toposort = topological_sort_by_dfs(g)
+
+  strord = stream_order(g, toposort)
+  strord_2d = reverse_index(strord, inds, inds_rev; mv=-1)
+
+  links = stream_link(g, strord, toposort, min_sto)
+  links_2d = reverse_index(links, inds, inds_rev; mv=0)
+
+  basinId_fill = fillnodata_upstream(g, toposort, links, 0)
+  basinId_2d = reverse_index(basinId_fill, inds, inds_rev; mv=0)
+  strord_2d, links_2d, basinId_2d
+end
+
+
+export SubBasins
+
+end

src/Ipaper.jl

@@ -54,6 +54,7 @@ include("precompile.jl")
 include("tools_plot.jl")
 
 include("apply.jl")
+
 include("hydro/Hydro.jl")
 
 
@@ -73,6 +74,12 @@ function __init__()
     ext = load_ext(:IpaperArchGDALExt)
     @reexport using .ext
   end
+
+  @require Wflow = "d48b7d99-76e7-47ae-b1d5-ff0c1cf9a818" begin
+    # include("../ext/IpaperArchGDALExt/IpaperArchGDALExt.jl")
+    ext = load_ext(:IpaperWflowExt)
+    @reexport using .ext
+  end
 end
 
 

src/hydro/Hydro.jl

@@ -2,5 +2,7 @@
 using DataStructures: PriorityQueue, dequeue!
 using ProgressMeter
 
+using .sf
+
 include("utils.jl")
 include("FlowDirection.jl")

src/hydro/utils.jl

@@ -2,13 +2,28 @@
 # 1. The dimemsion of `dem`: [lon, lat], which is different from ArcGIS. Hence, DIR
 # is different.
 # 2. `lat` in the reverse order.
-const _GIS = UInt8.([1, 2, 4, 8, 16, 32, 64, 128])
-const _TAU = UInt8.(1:8)
+const DIR_GIS = UInt8.([1, 2, 4, 8, 16, 32, 64, 128])
+const DIR_TAU = UInt8.(1:8)
+const DIR_WFLOW = UInt8.([6, 3, 2, 1, 4, 7, 8, 9])
+# # 按照这种方法
+# const pcr_dir = [
+#     CartesianIndex(-1, -1),  # 1, 8
+#     CartesianIndex(0, -1),   # 2, 4
+#     CartesianIndex(1, -1),   # 3, 2
+#     CartesianIndex(-1, 0),   # 4, 16
+#     CartesianIndex(0, 0),    # 5, 0
+#     CartesianIndex(1, 0),    # 6, 1
+#     CartesianIndex(-1, 1),   # 7, 32
+#     CartesianIndex(0, 1),    # 8, 64
+#     CartesianIndex(1, 1),    # 9, 128
+# ]
+
 const NODATA = 0x00
 const DY = [0, 1, 1, 1, 0, -1, -1, -1]
 const DX = [1, 1, 0, -1, -1, -1, 0, 1]
 const DIR = [1, 2, 4, 8, 16, 32, 64, 128]
 const DIR_INV = [16, 32, 64, 128, 1, 2, 4, 8]
+
 # DIV = [
 #   32 64 128
 #   16 0  1
@@ -22,22 +37,48 @@ const DIR_INV = [16, 32, 64, 128, 1, 2, 4, 8]
 # const DX = [0, 1, 1, 1, 0, -1, -1, -1]
 # const DY = [1, 1, 0, -1, -1, -1, 0, 1]
 
+# 0~9
+function gis2wflow(A::AbstractArray)
+  R = copy(A)
+  for i in 1:8
+    replace!(R, DIR_GIS[i] => DIR_WFLOW[i])
+  end
+  replace!(R, UInt8(0) => UInt8(5))
+  R
+end
+
 function gis2tau(A::AbstractArray)
   R = copy(A)
   for i in 1:8
-    replace!(R, _GIS[i] => _TAU[i])
+    replace!(R, DIR_GIS[i] => DIR_TAU[i])
   end
   R
 end
 
 function tau2gis(A::AbstractArray)
   R = copy(A)
   for i in 1:8
-    replace!(R, _TAU[i] => _GIS[i])
+    replace!(R, DIR_TAU[i] => DIR_GIS[i])
   end
   R
 end
 
+
+function flowdir_drop_missing!(A::AbstractMatrix{T}; mv=T(99)) where {T<:Real}
+  replace!(A, missing => mv)
+  replace!(A, 0 => mv)
+end
+
+function read_flowdir_wflow(f::String)
+  A_gis = read_gdal(f)[:, end:-1:1] # 修正颠倒的lat
+  A = gis2wflow(A_gis)
+
+  nodata = gdal_nodata(f)[1]
+  replace!(A, nodata => 0) # replace missing value with 0
+  A
+end
+
+
 """
   row_goto(row::Int, idir::Int)::Int
   col_goto(col::Int, idir::Int)::Int
@@ -88,4 +129,5 @@ end
 
 
 export InGrid, DirLength, NextCell, row_goto, col_goto
-export gis2tau, tau2gis
+export gis2wflow, gis2tau, tau2gis
+export read_flowdir_wflow, flowdir_drop_missing!