improve spei

.gitignore

@@ -3,4 +3,4 @@ deps
 .vscode
 
 Manifest.toml
-temp
+# temp

src/DIST/logLogistic.jl

@@ -1,25 +1,24 @@
-function _fit_logLogistic(beta::AbstractVector)
-  # params = zeros(3)
+function _fit_logLogistic(β1::T, β2::T, β3::T) where {T<:Real}
   g1 = 0.0
   g2 = 0.0
 
   # estimate gamma parameter
-  γ = (2 * beta[2] - beta[1]) / (6 * beta[2] - beta[1] - 6 * beta[3])
+  γ = (2 * β2 - β1) / (6 * β2 - β1 - 6 * β3)
   g1 = exp(loggamma(1 + 1 / γ))
   g2 = exp(loggamma(1 - 1 / γ))
 
   # estimate alpha parameter
-  α = (beta[1] - 2 * beta[2]) * γ / (g1 * g2) # params[2]
+  α = (β1 - 2 * β2) * γ / (g1 * g2) # params[2]
   # estimate beta parameter
-  β = beta[1] - α * g1 * g2                   # params[1]
+  β = β1 - α * g1 * g2                   # params[1]
   β, α, γ
 end
 
 
 function fit_logLogistic(x::AbstractVector)
   x2 = x[.!isnan.(x)] |> sort
-  beta = pwm(x2, 0.0, 0.0, 0)
-  params = _fit_logLogistic(beta)
+  β1, β2, β3 = pwm(x2, 0.0, 0.0, 0)
+  params = _fit_logLogistic(β1, β2, β3)
   params
 end
 
@@ -33,7 +32,7 @@ end
 const _C = [2.515517, 0.802853, 0.010328]
 const _d = [0.0, 1.432788, 0.189269, 0.001308]
 
-function invcdf_standardGaussian(prob::Real)
+function invcdf_standardGaussian(prob::T) where {T<:Real}
   W = prob <= 0.5 ? sqrt(-2 * log(prob)) : sqrt(-2 * log(1 - prob))
   WW = W * W
   WWW = WW * W

src/drought_SPEI.jl

@@ -16,13 +16,7 @@ function drought_SPEI(A::AbstractArray{T,3}, dates;
 
   nlon, nlat, ntime = size(A)
   A2 = similar(A) .* 0
-  @inbounds @par for k = scale:ntime
-    for j = 1:nlon, i = 1:nlat
-      _x = @view A[i, j, k-scale+1:k]
-      A2[i, j, k] = nanmean(_x)
-    end
-  end
-  A2[:, :, 1:scale-1] .= T(NaN)
+  accu_scale!(A2, A, scale)
 
   by = fun_date.(dates; delta)
   grps = unique_sort(by)
@@ -48,5 +42,17 @@ function drought_SPEI(A::AbstractArray{T,3}, dates;
   R
 end
 
+function accu_scale!(A2::AbstractArray{<:Real,3}, A::AbstractArray{<:Real,3}, scale::Int)
+  nlon, nlat, ntime = size(A)
+  @inbounds @par for k = scale:ntime
+    for j = 1:nlon, i = 1:nlat
+      _x = @view A[i, j, k-scale+1:k]
+      A2[i, j, k] = nanmean(_x)
+    end
+  end
+  A2[:, :, 1:scale-1] .= T(NaN)
+  A2
+end
+
 
 export drought_SPEI

src/lmoments.jl

@@ -5,52 +5,45 @@ A<=B<=0).
 
 This are alpha PWMs, following Rao & Hamed 2000, eqs. 3.1.4 and 3.1.6
 """
-function pwm!(acum::AbstractVector{T}, pwms::AbstractVector{T},
-  series::AbstractVector{T}, A::Float64, B::Float64, isBeta::Int) where {T<:Real}
-
+function pwm(series::AbstractVector{T}, A::Float64, B::Float64, isBeta::Int) where {T<:Real}
   F = 0.0
   n = length(series)
+  acum1 = acum2 = acum3 = T(0)
 
   if A == 0 && B == 0  # use unbiased estimator
     for i = 1:n
-      acum[1] += series[i]
+      acum1 += series[i]
       if isBeta == 0  # compute alpha PWMs
-        acum[2] += series[i] * (n - i)
-        acum[3] += series[i] * (n - i) * (n - i - 1)
+        acum2 += series[i] * (n - i)
+        acum3 += series[i] * (n - i) * (n - i - 1)
       end
       if isBeta == 1  # compute beta PWMs
-        acum[2] += series[i] * (i - 1)
-        acum[3] += series[i] * (i - 1) * (i - 2)
+        acum2 += series[i] * (i - 1)
+        acum3 += series[i] * (i - 1) * (i - 2)
       end
     end
   else  # use plotting-position (biased) estimator
     for i = 1:n
-      acum[1] += series[i]
+      acum1 += series[i]
       F = (i + A) / (n + B)
       if isBeta == 0  # compute alpha PWMs
-        acum[2] += series[i] * (1 - F)
-        acum[3] += series[i] * (1 - F) * (1 - F)
+        acum2 += series[i] * (1 - F)
+        acum3 += series[i] * (1 - F) * (1 - F)
       end
       if isBeta == 1  # compute beta PWMs
-        acum[2] += series[i] * F
-        acum[3] += series[i] * F * F
+        acum2 += series[i] * F
+        acum3 += series[i] * F * F
       end
     end
   end
 
-  pwms[1] = acum[1] / n
-  pwms[2] = acum[2] / n / (n - 1)
-  pwms[3] = acum[3] / n / ((n - 1) * (n - 2))
-  pwms
-end
-
-function pwm(series::AbstractVector{T}, A::Float64, B::Float64, isBeta::Int) where {T<:Real}
-  acum = zeros(T, 3)
-  pwms = zeros(T, 3)
-  pwm!(acum, pwms, series, A, B, isBeta)
+  β1 = acum1 / n
+  β2 = acum2 / n / (n - 1)
+  β3 = acum3 / n / ((n - 1) * (n - 2))
+  β1, β2, β3
 end
 
-function lmoments(series::AbstractVector{T}, A::Float64, B::Float64) where { T<:Real }
+function lmoments(series::AbstractVector{T}, A::Float64, B::Float64) where {T<:Real}
   # alpha = zeros(Float64, 3)
   # Calculate the first three PWMs
   lmom = zeros(T, 3)

src/main_spei.jl

@@ -3,8 +3,8 @@ function param_spei(x_ref::AbstractVector)
   x2 = x_ref[lgl]
   sort!(x2) # have to sort
 
-  beta = pwm(x2, 0.0, 0.0, 0)
-  params = _fit_logLogistic(beta)
+  β1, β2, β3 = pwm(x2, 0.0, 0.0, 0) # improve at here
+  params = _fit_logLogistic(β1, β2, β3)
   params
 end
 
@@ -26,29 +26,39 @@ function param_spi(x_ref::AbstractVector; fit="lmom")
 end
 
 
-"""
-    spei(x::AbstractVector)
-"""
-function spei(x::AbstractVector, x_ref::AbstractVector=x)
+function spei!(z::AbstractVector, x::AbstractVector, x_ref::AbstractVector=x)
   params = param_spei(x_ref)
-  z = zeros(Float64, length(x))
-
-  for i in eachindex(z)
+  @inbounds for i in eachindex(z)
     prob = cdf_logLogistic(x[i], params)
     z[i] = -invcdf_standardGaussian(prob)
   end
-  (; z, coef=params)
+  params
 end
 
-function spei!(z::AbstractVector, x::AbstractVector, x_ref::AbstractVector=x)
-  params = param_spei(x_ref)
+function spi!(z::AbstractVector, x::AbstractVector, x_ref::AbstractVector=x; fit="lmom")
+  param, q = param_spi(x_ref; fit)
+  D = Gamma(param...)
+
   @inbounds for i in eachindex(z)
-    prob = cdf_logLogistic(x[i], params)
-    z[i] = -invcdf_standardGaussian(prob)
+    if !isnan(x[i])
+      p = q .+ (1 - q) * cdf(D, x[i])
+      z[i] = qnorm(p)
+      x[i] < 0 && (z[i] = -Inf)
+    end
   end
+  params(D)
+end
+
+
+"""
+    spei(x::AbstractVector)
+"""
+function spei(x::AbstractVector, x_ref::AbstractVector=x)
+  z = zeros(Float64, length(x))
+  params = spei!(z, x, x_ref)
+  (; z, coef=params)
 end
 
-# 这里需要提供一个reference period
 """
     spi(x::AbstractVector; fit="lmom")
 
@@ -58,27 +68,10 @@ end
   + `mle`: maximum likelihood estimation, julia solver
 """
 function spi(x::AbstractVector, x_ref::AbstractVector=x; fit="lmom")
-  param, q = param_spi(x_ref; fit)
-  D = Gamma(param...)
-
   z = fill(NaN, length(x))
-  spi!(z, x, x_ref; fit)
-  (; z, coef=params(D))
+  params = spi!(z, x, x_ref; fit)
+  (; z, coef=params)
 end
 
-function spi!(z::AbstractVector, x::AbstractVector, x_ref::AbstractVector=x; fit="lmom")
-  param, q = param_spi(x_ref; fit)
-  D = Gamma(param...)
-
-  @inbounds for i in eachindex(z)
-    if !isnan(x[i])
-      p = q .+ (1 - q) * cdf(D, x[i])
-      z[i] = qnorm(p)
-      x[i] < 0 && (z[i] = -Inf)
-    end
-  end
-  return z
-end
 
-# TODO: 需要添加scale参数
 export spei!

test/runtests.jl

@@ -6,7 +6,8 @@ using Serialization
 include("test-ZSI.jl")
 nanmaximum(x) = maximum(x[.!isnan.(x)])
 
-wb = deserialize("../data/wb")
+# wb = deserialize("../data/wb")
+wb = deserialize("data/wb")
 x = wb[wb .> 0]
 
 @testset "gamma and loglogistic" begin