improve anomaly

src/Climate/anomaly.jl

@@ -1,96 +1,104 @@
 ## 计算anomaly的3种方法，考虑每年的升温幅度
+export _cal_anomaly_3d, _cal_anomaly
 
 # some CMIP6 model not use 366 calendar |> "solved"
 # DateType = Union{Date,DateTime,AbstractCFDateTime,Nothing}
-_gte(x::T, y::T) where {T<:Real} = x >= y
-_gt(x::T, y::T) where {T<:Real} = x > y
-_exceed(x::T, y::T) where {T<:Real} = x - y
+_gte(x::T, trs::T, wl::T=T(0)) where {T<:Real} = x >= trs + wl
+_gt(x::T, trs::T, wl::T=T(0)) where {T<:Real} = x > trs + wl
+_exceed(x::T, trs::T, wl::T=T(0)) where {T<:Real} = x - trs + wl
 
 
 """
 $(TYPEDSIGNATURES)
 """
-function _cal_anomaly(
-  arr::AbstractArray{T,3},
+function _cal_anomaly_3d(
+  A::AbstractArray{T,3},
   TRS::AbstractArray{T,3},
   dates;
-  T_wl::Union{AbstractArray{T,3}, Nothing}=nothing,
-  parallel::Bool=false, 
-  option = 2,
-  # ΔTRS=nothing,
-  dtype=nothing,
-  fun_anom=_exceed, verbose=false
+  T_wl::Union{AbstractArray{T,3},Nothing}=nothing,
+  fun_anom=_exceed, ignored...
 ) where {T<:Real}
 
-  dtype === nothing && (dtype = T)
-  # res = BitArray(undef, size(arr))
-
   mmdd = format_md.(dates)
   mds = mmdd |> unique |> sort
-  # doy_max = length(mds)
-
-  if option == 1
-    # 方案1：无法添加T_wl
-    idxs = indexin(mmdd, mds)
-    fun_anom.(arr, TRS[:, :, idxs])
-    
-  else option == 2
-    res = zeros(dtype, size(arr))
-    nlon, nlat, _ = size(arr)
-
-    years = year.(dates)
-    year_grps = unique(years) |> sort
-
-    # @timeit_all 
-    @par parallel for iy in 1:length(year_grps)
-      year = year_grps[iy]
-      
-      verbose && println("year = $year")
-      ind_y = findall(years .== year)
-      ind_d = indexin(mmdd[ind_y], mds) # allow some doys missing
-
-      @inbounds for i = 1:nlon, j = 1:nlat, k = eachindex(ind_y)
-        x = arr[i, j, ind_y[k]]
-        y = TRS[i, j, ind_d[k]] 
-
-        if T_wl !== nothing
-          y += T_wl[i, j, iy]
-        end
-        res[i, j, ind_y[k]] = fun_anom(x, y)
-      end
-    end
-    res    
-  end
+
+  years = year.(dates)
+  year_grps = years |> unique_sort
+
+  ind_d = indexin(mmdd, mds)
+  ind_y = indexin(years, year_grps)
+
+  _wl = T_wl === nothing ? T(0) : T_wl[:, :, ind_y]
+  fun_anom.(A, TRS[:, :, ind_d], _wl)
+end
+
+
+function _cal_anomaly(
+  A::AbstractArray{T,Na},
+  TRS::AbstractArray{T,Nt},
+  dates;
+  T_wl::Union{AbstractArray{T,Na},Nothing}=nothing,
+  fun_anom=_exceed,
+  deep=true,
+  ignored...
+) where {T<:Real,Na,Nt}
+
+  mmdd = format_md.(dates)
+  mds = mmdd |> unique_sort
+
+  years = year.(dates)
+  year_grps = years |> unique_sort
+
+  ind_d = indexin(mmdd, mds)
+  ind_y = indexin(years, year_grps)
+
+  _wl = T_wl === nothing ? T(0) : selectdim_deep(T_wl, Na, ind_y; deep)
+  fun_anom.(A, selectdim_deep(TRS, Na, ind_d; deep), _wl)
+  # broadcast(fun_anom, A, deepcopy(selectdim(TRS, Na, ind_d)), _wl)
 end
 
 
-## 更加傻瓜的版本
 """
 $(TYPEDSIGNATURES)
 
+Calculate the anomaly of a 3D array of temperature data.
+
 # Arguments
 
-- `method`: one of `["full", "season", "base"]`, see Vogel2020 for details
+- `arr`      : the 3D array of temperature data
+- `dates`    : an array of dates corresponding to the temperature data
+- `parallel` : whether to use parallel processing (default `false`)
+- `use_mov`  : whether to use a moving window to calculate the threshold (default `true`)
+- `method`   : the method to use for calculating the threshold, one of `["full", "season", "base"]` (default `"full"`)
+- `probs`    : default `[0.5]`
+- `p1`       : the start year for the reference period (default `1981`)
+- `p2`       : the end year for the reference period (default `2010`)
+- `fun_clim` : the function to use for calculating the climate state, one of `nanmean` or `nanmedian` (default `nanmean`)
+
+# Returns
+
+An array of the same shape as `arr` containing the temperature anomaly.
 
 # 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_anomaly_quantile(
-  arr::AbstractArray{T,3}, dates;
-  parallel::Bool=false, 
+  arr::AbstractArray{T}, dates;
+  parallel::Bool=false,
   use_mov=true, na_rm=true,
-  method = "full", 
+  method="full",
   p1=1981, p2=2010,
-  fun = _exceed, 
-  probs=0.5, 
+  fun=_exceed,
+  probs=[0.5],
   options...
 ) where {T<:Real}
-  
-  kw = (;probs=[probs], use_mov, na_rm, parallel, options...)
+
+  kw = (; probs, use_mov, na_rm, parallel, options...)
   # TODO: 多个阈值，需要再嵌套for循环了
   if method == "base"
     mTRS = cal_mTRS_base(arr, dates; p1, p2, kw...) |> squeeze_tail
@@ -108,11 +116,9 @@ end
 
 
 # 气候态用的是median的方法，如果想计算均值则需要另一套独立的方法
-"""
-$(TYPEDSIGNATURES)
-"""
+
 function cal_anomaly(
-  arr::AbstractArray{T,3},
+  arr::AbstractArray{T},
   dates;
   parallel::Bool=false,
   use_mov=true,
@@ -123,14 +129,14 @@ function cal_anomaly(
 ) where {T<:Real}
 
   kw = (; use_mov, parallel, fun=fun_clim)
-  
+
   if method == "base"
     mTRS = cal_climatology_base(arr, dates; p1, p2, kw...) |> squeeze_tail
-    anom = _cal_anomaly(arr, mTRS, dates; option=1, fun_anom)
+    anom = _cal_anomaly(arr, mTRS, dates; fun_anom)
   elseif method == "season"
     mTRS = cal_climatology_base(arr, dates; p1, p2, kw...) |> squeeze_tail
     T_wl = cal_warming_level(arr, dates; p1, p2)
-    anom = _cal_anomaly(arr, mTRS, dates; option=2, T_wl, fun_anom)
+    anom = _cal_anomaly(arr, mTRS, dates; T_wl, fun_anom)
   elseif method == "full"
     TRS_full = cal_climatology_full(arr, dates; kw...) |> squeeze_tail
     anom = fun_anom.(arr, TRS_full)
@@ -139,14 +145,13 @@ function cal_anomaly(
 end
 
 
-
-function cal_anomaly_quantile(arr::AbstractVector{<:Real}, dates; kw...)
-  arr = reshape(arr, 1, 1, length(arr))
-  cal_anomaly_quantile(arr, dates; kw...)[1, 1, :]
-end
+# function cal_anomaly_quantile(arr::AbstractVector{<:Real}, dates; kw...)
+#   arr = reshape(arr, 1, 1, length(arr))
+#   cal_anomaly_quantile(arr, dates; kw...)[1, 1, :]
+# end
 
 
-function cal_anomaly(arr::AbstractVector{<:Real}, dates; kw...)
-  arr = reshape(arr, 1, 1, length(arr))
-  cal_anomaly(arr, dates; kw...)[1, 1, :]
-end
+# function cal_anomaly(arr::AbstractVector{<:Real}, dates; kw...)
+#   arr = reshape(arr, 1, 1, length(arr))
+#   cal_anomaly(arr, dates; kw...)[1, 1, :]
+# end

src/Climate/warming_level.jl

@@ -1,45 +1,62 @@
 """
-Seasonally moving thresholds
+Calculate yearly air temperature.
 
+# Description
 we use the fixed thresholds and add the seasonal warming signal. 
-
 Thus, thresholds are defined as a fixed baseline (such as for the fixed threshold) plus
 seasonally moving mean warming of the corresponding future climate based on the
 31-year moving mean of the warmest three months.
+
+# Details
+This function calculates the yearly air temperature based on the input temperature data and dates. If `only_summer` is true, it only calculates the temperature for summer months. The function applies the calculation along the specified dimensions.
+
+# Arguments
+- `A::AbstractArray{T,N}`: input array of temperature data.
+- `dates`: array of dates corresponding to the temperature data.
+- `dims=N`: dimensions to apply the function along.
+- `only_summer=false`: if true, only calculate temperature for summer months.
+
+# Returns
+- `T_year`: array of yearly temperature data.
 """
-function cal_yearly_Tair(arr::AbstractArray{<:Real, 3}, dates; only_summer=false)
+function cal_yearly_Tair(A::AbstractArray{T,N}, dates;
+  dims=N, only_summer=false) where {T<:Real,N}
+
   if only_summer
     yms = format.(dates, "yyyy-mm")
     ys = SubString.(unique(yms), 1, 4)
-    T_mon = apply(arr, 3; by=yms)
-    T_mon = movmean(T_mon, 1; dims=3) #3个月滑动平均
-    T_year = apply(T_mon, 3; by=ys, fun=maximum) # 最热的3个月，作为每年的升温幅度
+    T_mon = apply(A, dims; by=yms)
+    T_mon = movmean(T_mon, 1; dims) #3个月滑动平均
+
+    T_year = apply(T_mon, dims; by=ys, fun=maximum) # 最热的3个月，作为每年的升温幅度
   else
     ys = format.(dates, "yyyy")
-    T_year = apply(arr, 3; by=ys)
+    T_year = apply(A, dims; by=ys)
   end
   T_year
 end
 
-cal_climatology_season(arr::AbstractArray, dates) = cal_yearly_Tair(arr, dates; only_summer=false)
+cal_climatology_season(A::AbstractArray, dates) = cal_yearly_Tair(A, dates; only_summer=false)
 
-cal_mTRS_season(arr::AbstractArray, dates) = cal_yearly_Tair(arr, dates; only_summer=true)
+cal_mTRS_season(A::AbstractArray, dates) = cal_yearly_Tair(A, dates; only_summer=true)
 
 
 """
 $(TYPEDSIGNATURES)
 """
-function cal_warming_level(arr::AbstractArray{<:Real, 3}, dates; 
-  p1=1981, p2=2010, only_summer=false)
-  
-  T_year = cal_yearly_Tair(arr, dates; only_summer)
-  
+function cal_warming_level(A::AbstractArray{T,N}, dates;
+  p1=1981, p2=2010, dims=N, only_summer=false) where {T<:Real,N}
+
+  T_year = cal_yearly_Tair(A, dates; dims, only_summer)
+
   # yms = format.(dates, "yyyy-mm")
   # ys = SubString.(unique(yms), 1, 4)
   # grps = unique_sort(ys)
-  grps = year.(dates) |> unique_sort  
+  grps = year.(dates) |> unique_sort
   inds_clim = @.(p1 <= grps <= p2)
-  T_clim = apply(@view(T_year[:, :, inds_clim]), 3; fun=nanmean)
+
+  T_year_clim = selectdim(T_year, dims, inds_clim)
+  T_clim = apply(T_year_clim, dims; fun=nanmean)
 
   T_year .- T_clim
 end

src/tools.jl

@@ -125,12 +125,23 @@ function abind(x::AbstractVector, dim=3)
 end
 
 
+function selectdim_deep(A, dims::Integer, i; deep=true)
+  if deep
+    inds = ntuple(d -> d in dims ? i : (:), ndims(A))
+    A[inds...]
+  else
+    selectdim(A, dims, i)
+  end
+end
+
+
 export which_isna, which_notna,
   is_empty, not_empty,
   mean, weighted_mean, weighted_sum,
   seq_along, seq_len,
   r_range,
   nth, 
+  selectdim_deep,
   length_unique, unique_sort, 
   squeeze, squeeze_tail, squeeze_head,
   abind,