add ZSI

Project.toml

@@ -4,9 +4,13 @@ authors = ["Dongdong Kong <kongdd@users.noreply.github.com>"]
 version = "0.1.1"
 
 [deps]
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Ipaper = "e58298cb-69f7-4186-aecd-5834b6793426"
+NaNStatistics = "b946abbf-3ea7-4610-9019-9858bfdeaf2d"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

src/SPEI.jl

@@ -24,5 +24,6 @@ include("lmoments.jl")
 
 include("DIST/DIST.jl")
 include("main_spei.jl")
+include("ZSI.jl")
 
 end # module SPEI

src/ZSI.jl

@@ -0,0 +1,52 @@
+using Ipaper
+using Dates
+using NaNStatistics
+
+function get_dn(date; delta=8)
+  days = Dates.dayofyear(date)
+  cld(days, delta) # int
+end
+
+
+"""
+    drought_ZSI(A::AbstractArray{T,3}, dates;
+        ref=(2000, 2020),
+        fun_date=get_dn, delta::Int=8,
+        progress=true, parallel=true) where {T<:Real}
+
+Calculate the ZSI (Z-Score Index) of the input data `A` along the time dimension.
+> Extremely fast!
+"""
+function drought_ZSI(A::AbstractArray{T,3}, dates;
+  ref=(2000, 2020),
+  fun_date=get_dn, delta::Int=8,
+  progress=true, parallel=true) where {T<:Real}
+
+  by = fun_date.(dates; delta)
+  grps = unique_sort(by)
+
+  nlon, nlat, ntime = size(A)
+  R = zeros(T, nlon, nlat, ntime)
+
+  p = Progress(length(grps))
+  @inbounds @par parallel for grp in grps
+    progress && next!(p)
+
+    I = findall(by .== grp)
+    I_ref = I[findall(ref[1] .<= year.(dates[I]) .<= ref[2])]
+
+    for j = 1:nlat, i = 1:nlon
+      _x_ref = @view A[i, j, I_ref] # 两次嵌套会导致速度变慢
+      μ = nanmean(_x_ref)
+      sd = nanstd(_x_ref)
+
+      for _i in I
+        R[i, j, _i] = (A[i, j, _i] - μ) / sd
+      end
+    end
+  end
+  R
+end
+
+
+export get_dn, drought_ZSI

test/runtests.jl

@@ -2,6 +2,8 @@ using Test
 using SPEI
 using Serialization
 
+
+include("test-ZSI.jl")
 nanmaximum(x) = maximum(x[.!isnan.(x)])
 
 wb = deserialize("../data/wb")

test/test-ZSI.jl

@@ -0,0 +1,21 @@
+using Statistics, Test
+
+@testset "drought_ZSI" begin
+  dates = Date(2000):Day(8):Date(2010, 12, 31) |> collect
+  A = rand(Float32, 100, 100, length(dates))
+  ref = (2001, 2005)
+  @time R = drought_ZSI(A, dates; ref)
+
+  by = get_dn.(dates; delta=8)
+  grps = unique_sort(by)
+
+  grp = grps[1]
+  I = findall(by .== grp)
+  I_ref = I[findall(ref[1] .<= year.(dates[I]) .<= ref[2])]
+
+  x = A[1, 1, I]
+  x_ref = A[1, 1, I_ref]
+  μ = mean(x_ref)
+  sd = std(x_ref)
+  @test R[1, 1, I] ≈ (x .- μ) ./ sd
+end