unify cal_mTRS_full

src/Climate/Climate.jl

@@ -1,6 +1,5 @@
 include("base.jl")
 include("threshold.jl")
-include("threshold_3d.jl")
 include("climatology.jl")
 include("anomaly.jl")
 include("warming_level.jl")

src/Climate/anomaly.jl

@@ -36,7 +36,7 @@ function _cal_anomaly(
 
   else option == 2
     res = zeros(dtype, size(arr))
-    nlon, nlat, ntime = size(arr)
+    nlon, nlat, _ = size(arr)
 
     years = year.(dates)
     year_grps = unique(years) |> sort

src/Climate/climatology.jl

@@ -1,20 +1,16 @@
 """
 $(TYPEDSIGNATURES)
 """
-function cal_climatology_base!(Q::AbstractArray{T,3}, data::AbstractArray{T,3}, dates;
+function cal_climatology_base3!(Q::AbstractArray{T,3}, data::AbstractArray{T,3}, mmdd;
   use_mov=true, halfwin::Int=7,
   parallel::Bool=true, fun=nanmean,
   ignored...) where {T<:Real}
 
-
-  mmdd = factor(format_md.(dates)).refs
-  mds = mmdd |> unique_sort
-  doy_max = length(mds)
+  doy_max = maximum(mmdd)
 
   nlon, nlat, _ = size(data)
   @inbounds @par parallel for doy = 1:doy_max
     ind = filter_mds(mmdd, doy; doy_max, halfwin, use_mov)
-    # q = @view Q[:, :, doy, :]
     for i = 1:nlon, j = 1:nlat
       x = @view data[i, j, ind]
       Q[i, j, doy] = fun(x)
@@ -28,82 +24,17 @@ end
 $(TYPEDSIGNATURES)
 """
 function cal_climatology_base(arr::AbstractArray{<:Real,3}, dates;
-  # probs::Vector=[0.90, 0.95, 0.99, 0.999, 0.9999],
-  type=nothing,
-  p1::Int=1961, p2::Int=1990, kw...)
+  (fun!)=cal_climatology_base3!, kw...)
 
-  mmdd = format_md.(dates)
-  doy_max = length_unique(mmdd)
-
-  type = type === nothing ? eltype(arr) : type
-  dim = size(arr)
-  Q = zeros(type, dim[1:2]..., doy_max)
-
-  # constrain date in [p1, p2]
-  years = year.(dates)
-  ind = findall(p1 .<= years .<= p2)
-
-  _data = selectdim(arr, 3, ind)
-  _dates = @view dates[ind]
-  cal_climatology_base!(Q, _data, _dates; kw...)
+  cal_mTRS_base(arr, dates; use_quantile=false, (fun!), kw...)
 end
 
 
-"""
-  cal_climatology_full
-  
+"""  
 $(TYPEDSIGNATURES)
 """
 function cal_climatology_full(arr::AbstractArray{T}, dates;
-  width=15, verbose=true, use_mov=true,
-  kw...) where {T<:Real}
-
-  years = year.(dates)
-  grps = unique(years)
-
-  YEAR_MIN = minimum(grps)
-  YEAR_MAX = maximum(grps)
-
-  mmdd = format_md.(dates)
-  mds = unique(mmdd) |> sort
-  doy_max = length(mds)
-
-  # 滑动平均两种不同的做法
-  if !use_mov
-    printstyled("running: 15d moving average first ... ")
-    @time arr = movmean(arr, 7; dims=3, FT=Float32)
-  end
+  (fun!)=cal_climatology_base3!, kw...) where {T<:Real}
 
-  dim = size(arr)
-  mTRS_full = zeros(T, dim[1:3]...)
-  mTRS = zeros(T, dim[1:2]..., doy_max)
-
-  TRS_head = cal_climatology_base(arr, dates; p1=YEAR_MIN, p2=YEAR_MIN + width * 2, use_mov, kw...)
-  TRS_tail = cal_climatology_base(arr, dates; p1=YEAR_MAX - width * 2, p2=YEAR_MAX, use_mov, kw...)
-
-  for year = grps
-    verbose && mod(year, 5) == 0 && println("running [year=$year]")
-
-    inds_year = years .== year
-    md = @view mmdd[inds_year]
-    ind = findall(r_in(mds, md))
-
-    year_beg = max(year - width, YEAR_MIN)
-    year_end = min(year + width, YEAR_MAX)
-
-    if year <= YEAR_MIN + width
-      _mTRS = TRS_head
-    elseif year >= YEAR_MAX - width
-      _mTRS = TRS_tail
-    else
-      inds_data = @.(years >= year_beg && year <= year_end)
-      _data = selectdim(arr, 3, inds_data)
-      _dates = @view dates[inds_data]
-
-      cal_climatology_base!(mTRS, _data, _dates; use_mov, kw...)
-      _mTRS = mTRS
-    end
-    @views copy!(mTRS_full[:, :, inds_year], _mTRS[:, :, ind])
-  end
-  mTRS_full
+  cal_mTRS_full(arr, dates; use_quantile=false, (fun!), kw...)
 end

src/Climate/threshold.jl

@@ -1,20 +1,50 @@
-export Threshold
+export cal_mTRS_base, cal_mTRS_full
 
-module Threshold
 
-export cal_mTRS_base
-using Ipaper
+"""
+$(TYPEDSIGNATURES)
+"""
+function cal_mTRS_base3!(Q::AbstractArray, data::AbstractArray{T,3}, mmdd;
+  dims=3,
+  probs::Vector=[0.90, 0.95, 0.99, 0.999, 0.9999],
+  parallel::Bool=true, halfwin::Int=7, use_mov=true,
+  method_q="base", na_rm=false,
+  ignore...) where {T<:Real}
+
+  doy_max = maximum(mmdd)
+
+  @inbounds @par parallel for doy = 1:doy_max
+    ind = filter_mds(mmdd, doy; doy_max, halfwin, use_mov)
+
+    q = @view Q[:, :, doy, :]
+    x = @view data[:, :, ind] # 这一步耗费内存
+    if method_q == "base"
+      NanQuantile_3d!(q, x; probs, dims, na_rm)
+      # NanQuantile_low!(q, x; probs, dims=3, na_rm)
+    elseif method_q == "mapslices"
+      q .= NanQuantile(x; probs, dims, na_rm) # mapslices is suppressed for 3d `NanQuantile`
+    end
+  end
+  Q
+end
+
 
+"""
+# Arguments
 
-function cal_mTRS_base!(Q::AbstractArray{T}, 
+- `method_q`: method to calculate quantile, one of `base`, `mapslices`.
+  `base` is about 3 times faster and reduce used memory in 20 times. 
+"""
+function cal_mTRS_base!(Q::AbstractArray{T},
   arr::AbstractArray{T,N}, mmdd;
   dims=N,
+  fun=NanQuantile,
   probs::Vector=[0.90, 0.95, 0.99, 0.999, 0.9999],
   parallel=true, halfwin=7, use_mov=true,
-  kw...) where {T<:Real, N}
+  ignore...) where {T<:Real,N}
 
   doy_max = maximum(mmdd)
-  
+
   @inbounds @par parallel for doy = 1:doy_max
     ind = filter_mds(mmdd, doy; doy_max, halfwin, use_mov)
     # idx = ntuple(i -> i == dims ? ind : Colon(), N)
@@ -23,16 +53,35 @@ function cal_mTRS_base!(Q::AbstractArray{T},
     # q = @view(Q[ridx...])
     x = selectdim(arr, dims, ind)
     q = selectdim(Q, dims, doy)
-    q .= NanQuantile(x; probs, dims) # mapslices 
+    q .= fun(x; probs, dims) # NanQuantile, mapslices 
   end
   Q
 end
 
 
+
+"""
+Moving Threshold for Heatwaves Definition
+
+  $(TYPEDSIGNATURES)
+
+# Arguments
+- `type`: The matching type of the moving `doys`, "md" (default) or "doy".
+
+# Return
+- `TRS`: in the dimension of `[nlat, nlon, ndoy, nprob]`
+
+# References
+1. Vogel, M. M., Zscheischler, J., Fischer, E. M., & Seneviratne, S. I. (2020).
+  Development of Future Heatwaves for Different Hazard Thresholds. Journal of
+  Geophysical Research: Atmospheres, 125(9).
+  https://doi.org/10.1029/2019JD032070
+"""
 function cal_mTRS_base(arr::AbstractArray{T,N}, dates;
   dims=N,
+  use_quantile=true, (fun!)=cal_mTRS_base!,
   probs::Vector=[0.90, 0.95, 0.99, 0.999, 0.9999],
-  dtype=nothing, 
+  dtype=nothing,
   p1::Int=1961, p2::Int=1990, kw...) where {T<:Real,N}
 
   mmdd = factor(format_md.(dates)).refs
@@ -41,24 +90,50 @@ function cal_mTRS_base(arr::AbstractArray{T,N}, dates;
   nprob = length(probs)
   dtype = dtype === nothing ? eltype(arr) : dtype
 
-  dims_r = map(d -> d in dims ? [doy_max, nprob] : size(arr, d), 1:N)
-  dims_r = cat(dims_r..., dims=1)
+  if use_quantile
+    dims_r = map(d -> d in dims ? [doy_max, nprob] : size(arr, d), 1:N) |> x -> vcat(x...)
+  else
+    dims_r = map(d -> d in dims ? doy_max : size(arr, d), 1:N) |> x -> vcat(x...)
+  end
   Q = zeros(dtype, dims_r...)
 
   # constrain date in [p1, p2]
-  years = year.(dates)
-  ind = findall(p1 .<= years .<= p2)
-  _data = selectdim(arr, dims, ind)
-  _mmdd = @view mmdd[ind]
+  inds = p1 .<= year.(dates) .<= p2
+  _data = selectdim(arr, dims, inds)
+  _mmdd = @view mmdd[inds]
 
-  cal_mTRS_base!(Q, _data, _mmdd; dims, probs, kw...)
+  fun!(Q, _data, _mmdd; dims, probs, kw...)
 end
 
+cal_mTRS = cal_mTRS_base
+
+
+
+"""
+Moving Threshold for Heatwaves Definition
+
+$(TYPEDSIGNATURES)
 
+# Arguments
 
+- `use_mov`: Boolean (default true). 
+  + if `true`, 31*15 values will be used to calculate threshold for each grid; 
+  + if `false`, the input `arr` is smoothed first, then only 15 values will be 
+    used to calculate threshold.
+
+!!! 必须是完整的年份，不然会出错
+
+# References
+1. Vogel, M. M., Zscheischler, J., Fischer, E. M., & Seneviratne, S. I. (2020).
+  Development of Future Heatwaves for Different Hazard Thresholds. Journal of
+  Geophysical Research: Atmospheres, 125(9).
+  https://doi.org/10.1029/2019JD032070
+"""
 function cal_mTRS_full(arr::AbstractArray{T,N}, dates;
   dims=N,
-  width=15, verbose=true, use_mov=true,
+  width=15, verbose=true,
+  use_quantile=true, (fun!)=cal_mTRS_base!,
+  use_mov=true,
   probs=[0.90, 0.95, 0.99, 0.999, 0.9999], kw...) where {T<:Real,N}
 
   years = year.(dates)
@@ -70,24 +145,35 @@ function cal_mTRS_full(arr::AbstractArray{T,N}, dates;
   mmdd = factor(format_md.(dates)).refs
   mds = unique(mmdd) |> sort
   doy_max = length(mds)
-  
+
   # 滑动平均两种不同的做法
   if !use_mov
     printstyled("running: 15d moving average first ... ")
     @time arr = movmean(arr, 7; dims, FT=Float32)
   end
 
-  nprob = length(probs)
+  if use_quantile
+    nprob = length(probs)
+    dim_full = map(d -> d in dims ? [size(arr, d), nprob] : size(arr, d), 1:N) |> x -> vcat(x...)
+    dim_mTRS = map(d -> d in dims ? [doy_max, nprob] : size(arr, d), 1:N) |> x -> vcat(x...)
+  else
+    dim_full = size(arr)
+    dim_mTRS = map(d -> d in dims ? doy_max : size(arr, d), 1:N) |> x -> vcat(x...)
 
-  dim_full = map(d -> d in dims ? [size(arr, d), nprob] : size(arr, d), 1:N) |> x -> vcat(x...)
-  dim_mTRS = map(d -> d in dims ? [doy_max, nprob] : size(arr, d), 1:N) |> x -> vcat(x...)
+    @show(dim_full)
+    @show(dim_mTRS)
+  end
 
   mTRS_full = zeros(T, dim_full...)
-  mTRS      = zeros(T, dim_mTRS...)
-
-  TRS_head = cal_mTRS_base(arr, dates; p1=YEAR_MIN, p2=YEAR_MIN + width * 2, use_mov, probs, kw...)
-  TRS_tail = cal_mTRS_base(arr, dates; p1=YEAR_MAX - width * 2, p2=YEAR_MAX, use_mov, probs, kw...)
-
+  mTRS = zeros(T, dim_mTRS...)
+
+  TRS_head = zeros(T, dim_mTRS...)
+  TRS_tail = zeros(T, dim_mTRS...)
+  inds_head = findall(YEAR_MIN .<= years .<= (YEAR_MIN + width * 2))
+  inds_tail = findall((YEAR_MAX - width * 2) .<= years .<= YEAR_MAX)
+  fun!(TRS_head, selectdim(arr, dims, inds_head), @view(mmdd[inds_head]); use_mov, probs, kw...)
+  fun!(TRS_tail, selectdim(arr, dims, inds_tail), @view(mmdd[inds_tail]); use_mov, probs, kw...)
+
   for year = grps
     verbose && mod(year, 20) == 0 && println("running [year=$year]")
 
@@ -97,17 +183,17 @@ function cal_mTRS_full(arr::AbstractArray{T,N}, dates;
 
     year_beg = max(year - width, YEAR_MIN)
     year_end = min(year + width, YEAR_MAX)
-    # @show year, YEAR_MIN + width, YEAR_MAX - width
+
     if year <= YEAR_MIN + width
       _mTRS = TRS_head
     elseif year >= YEAR_MAX - width
       _mTRS = TRS_tail
     else
-      inds_data = @.(years >= year_beg && year <= year_end)
+      inds_data = year_beg .<= years .<= year_end
       _data = selectdim(arr, dims, inds_data)
       _mmdd = @view mmdd[inds_data]
-      
-      cal_mTRS_base!(mTRS, _data, _mmdd; use_mov, probs, kw...)
+
+      fun!(mTRS, _data, _mmdd; use_mov, probs, kw...)
       _mTRS = mTRS # 366, 后面统一取ind
     end
     idx_year = ntuple(d -> d in dims ? inds_year : Colon(), N + 1)
@@ -118,6 +204,3 @@ function cal_mTRS_full(arr::AbstractArray{T,N}, dates;
   end
   mTRS_full
 end
-
-
-end