为R语言wd函数添加湿季重置赤字条件的技术求助

需求说明 #

原wd函数逻辑：通过WB（水量平衡）列计算赤字，当WB为负时，deficit列记录赤字并随时间累积；WB为正时重置deficit为0。
新增需求：当日期处于年度最湿润月份时，无论WB数值如何，都需将deficit重置为0，其余代码逻辑保持不变。

之前尝试的问题 #

曾尝试修改原条件if (df[[ varname_wbal ]][idx]<0)，添加&& df$is_wettest_month[idx] == TRUE或|| df$is_wettest_month[idx] == TRUE，但均未达到预期效果：

• 赤字事件追踪字段iinst、持续时长dday记录出现异常
• 最湿润月份的deficit未成功重置为0

修改思路 #

问题核心在于：原逻辑仅通过WB值控制赤字累积，我们需要将最湿润月份的条件作为强制终止/跳过累积的触发条件，而非修改进入累积的判断条件。需在两个关键节点处理：

1. 外层循环中，先检查当前日期是否属于最湿润月份，若是则直接重置deficit为0，不进入赤字累积逻辑
2. 内层赤字累积循环中，每一步都检查是否进入最湿润月份，若是则立即终止当前累积事件，重置deficit为0

修改后的完整代码 #

wd <- function(df, varname_wbal, varname_date, thresh_terminate = 0.0, thresh_drop = 0.9){
  
  if (thresh_terminate > thresh_drop){
    warning("Aborting. thresh_terminate must be smaller or equal thresh_drop. Setting it equal.")
    thresh_terminate <- thresh_drop
  }
  
  inst <- tibble()
  idx <- 0
  iinst <- 1
  
  df <- df |>
    ungroup() |>
    mutate(iinst = NA, dday = NA, deficit = 0)
  
  # 搜索所有日期
  while (idx <= (nrow(df)-1)){
    idx <- idx + 1
    
    # 新增：如果当前日期是年度最湿润月份，直接重置deficit为0，跳过累积逻辑
    if (df$is_wettest_month[idx]) {
      df$deficit[idx] <- 0
      next
    }
    
    # 原逻辑：当水量平衡为负时，开始累积赤字
    if (df[[ varname_wbal ]][idx]<0 ){
      
      dday <- 0
      deficit <- 0
      max_deficit <- 0
      iidx <- idx
      done_finding_dropday <- FALSE
      
      # 修改内层循环条件：新增检查是否进入最湿润月份，若是则终止累积
      while (iidx <= (nrow(df)-1) && 
             (deficit - df[[ varname_wbal ]][iidx] > thresh_terminate * max_deficit) &&
             !df$is_wettest_month[iidx]) { # 新增：最湿润月份则终止循环
        dday <- dday + 1
        deficit <- deficit - df[[ varname_wbal ]][iidx]
        
        # 记录当前事件的最大赤字
        if (deficit > max_deficit){
          max_deficit <- deficit
          done_finding_dropday <- FALSE
        }
        
        # 记录赤字降至最大赤字thresh_drop倍时的日期
        if (deficit < (max_deficit * thresh_drop) && !done_finding_dropday){
          iidx_drop <- iidx
          done_finding_dropday <- TRUE
        }
        
        # 赤字降至阈值后，后续日期标记为NA
        if (done_finding_dropday){
          df$iinst[iidx] <- NA
          df$dday[iidx]  <- NA
        } else {
          df$iinst[iidx] <- iinst
          df$dday[iidx]  <- dday
          iidx_drop <- iidx
        }
        
        df$deficit[iidx] <- deficit
        
        iidx <- iidx + 1
        
      }
      
      # 处理循环终止时的最后一个日期：如果是最湿润月份，重置deficit为0
      if (iidx <= nrow(df) && df$is_wettest_month[iidx]) {
        df$deficit[iidx] <- 0
      }
      
      # 记录当前赤字事件
      this_inst <- tibble( idx_start = idx, len = iidx_drop-idx, iinst = iinst, date_start=df[[varname_date]][idx], date_end = df[[varname_date]][iidx_drop-1], deficit = max_deficit )
      inst <- inst |> bind_rows(this_inst)
      
      # 更新事件编号和索引
      iinst <- iinst + 1
      dday <- 0
      idx <- iidx
    }
    
  }
  return( list(inst=inst, df=df))
}

验证要点 #

修改后需确认以下内容：

• 所有标记为is_wettest_month == TRUE的日期，deficit均为0
• 非最湿润月份的赤字累积逻辑与原函数完全一致
• 赤字事件的iinst编号、dday持续时长记录正常，无异常中断或重复

内容的提问来源于stack exchange，提问作者Shunrei