R/wavelet-R-detection.R

In Teresa00/hfmAnnotation: HFM Internal Package for Annotations

# Teresa updated March 28th 2018

#' @importFrom assertthat assert_that
#' @importFrom utils tail
#' @importFrom magrittr %>% subtract add
#' @importFrom purrr as_vector discard partial reduce
#'
library("assertthat")
library("magrittr")
library("purrr")


find_modulus_maxima <- function(signal_, threshold, polarity) {
  assert_that(
    all(is.numeric(signal_)),
    threshold >= 0,
    polarity %in% c(-1, 0, 1)
  )

  signs <- sign(signal_)
  x <- abs(signal_)

  slice_start <- 2
  slice_end <- length(signal_)

  slice <- c(slice_start:(slice_end - 1))
  before <- c((slice_start - 1):(slice_end - 2))
  after <- c((slice_start + 1):slice_end)

  localmax <-
    matrix(
      (x[slice] >= x[before]) &
        (x[slice] > x[after]) &
        (x[slice] > threshold) &
        ((signs[slice] * polarity) >= 0)
    )

  extrema_indexes <- rep(0, length(signal_))
  extrema_indexes[slice] <- localmax

  which(extrema_indexes != 0)
}


index_diff_ <- function(signal_slice) {
  signal_diff <- diff(signal_slice)

  differences <-
    signal_diff[-length(signal_diff)] * signal_diff[-1]

  which(differences <= 0)[1]
}


peak_before <- function(signal_slice, index) {
  inverse_slice <- signal_slice[length(signal_slice):1]

  index - index_diff_(inverse_slice)
}


peak_after <- function(signal_slice, index) {
  index + index_diff_(signal_slice)
}


zerocross <- function(x) {
  m <- x[-1] * x[-length(x)]

  first_index <- which(m <= 0)[1]

  index <-
    if (!is.na(first_index)) {
      has_greater_amplitude <-
        abs(x[first_index]) > abs(x[first_index + 1])

      if (has_greater_amplitude) {
        Just(first_index + 1)
      } else {
        Just(first_index)
      }
    }else{
      Nothing()
    }
}


# Set modulus maximumlines for scale 3,2,1
find_neighbouring_maxima_ <-
  function(modulus_maxima,
             wavelets,
             thresholds,
             current_scale,
             next_scale,
             sampling_frequency = 250,
             peak_criterion = 1.2) {
    # Search for modulus maxima in next_scale
    # around the neighbour time of maxima in current_scale
    neighbourhood <- pracma::ceil(sampling_frequency * 0.025)
    current_maxima <- modulus_maxima$scale[current_scale, ]

    find_neighbour_indexes_ <-
      function(maxima_locations, location) {
        # Added July 24 for extra protection
        negative_index <-
          which(maxima_locations <= 0)

        maxima_locations <-
          if(length(negative_index) != 0){
            maxima_locations[-negative_index]
          }else{
            maxima_locations
          }
        ####
        number_of_maxima <- length(maxima_locations)

        if (number_of_maxima < 1) {
          0
        } else if (number_of_maxima == 1) {
          maxima_locations
        } else {
          neighbouring_maxima <- abs(wavelets[maxima_locations, next_scale])
          biggest_neighbours <- max(neighbouring_maxima) / neighbouring_maxima
          local_maxima <- abs(current_maxima[location] - maxima_locations)

          maxima_index <-
            if (!any(biggest_neighbours < peak_criterion)) {
              which.max(neighbouring_maxima)
            } else {
              which.min(local_maxima)
            }

          maxima_locations[maxima_index]
        }
      }

    for (location in 1:ncol(modulus_maxima$scale)) {
      slice_begin <- max(1, current_maxima[location] - neighbourhood)
      slice_end <- min(current_maxima[location] + neighbourhood, nrow(wavelets))

      neighbour_polarity <- modulus_maxima$polarities[location]

      maxima <-
        wavelets[slice_begin:slice_end, next_scale] %>%
        find_modulus_maxima(thresholds[next_scale], neighbour_polarity)

      modulus_maxima$scale[next_scale, location] <-
        current_maxima[location] %>%
        add(maxima - neighbourhood - 1) %>%
        find_neighbour_indexes_(location)
    }

    index <- which(modulus_maxima$scale[next_scale, ] <= 0)

    if (length(index) != 0) {
      list(
        scale = modulus_maxima$scale[, -index],
        polarities = modulus_maxima$polarities[-index]
      )
    } else {
      modulus_maxima
    }
  }


# Eliminate high-frequency noise in ECG signal
eliminate_noise_ <- function(modulus_maxima, scale3_wavelets) {
  alpha <- log(abs(scale3_wavelets[modulus_maxima$scale[3, ]]))

  threshold <- mean(alpha) - 0.6

  noise <- which(alpha <= threshold)

  if (length(noise) != 0) {
    list(
      scale = modulus_maxima$scale[, -noise],
      polarities = modulus_maxima$polarities[-noise]
    )
  } else {
    modulus_maxima
  }
}


eliminate_isolation_ <- function(modulus_maxima, sampling_frequency = 250) {
  interval_threshold <- pracma::ceil(sampling_frequency * 0.15)
  scale1_maxima <- modulus_maxima$scale[1, ]
  size <- ncol(modulus_maxima$scale)

  interval_before <-
    scale1_maxima[2:(size - 1)] - scale1_maxima[1:(size - 2)]

  interval_after <-
    scale1_maxima[3:size] - scale1_maxima[2:(size - 1)]

  beyond_threshold <-
    which(
      (interval_before > interval_threshold) &
        (interval_after > interval_threshold)
    )

  isolated_max <-
    if (size < 2) {
      1
    } else {
      beyond_threshold + 1
    }

  if (length(isolated_max) != 0) {
    list(
      scale = modulus_maxima$scale[, -isolated_max],
      polarities = modulus_maxima$polarities[-isolated_max]
    )
  } else {
    modulus_maxima
  }
}


# Eliminate redundant maxima
# ===========================================================
polarity_judge_ <- function(polarities, is_positive) {
  if (is_positive) {
    list(
      maxima_location = which(polarities > 0),
      polarity_criteria = polarities < 0
    )
  } else {
    list(
      maxima_location = which(polarities < 0),
      polarity_criteria = polarities > 0
    )
  }
}


remove_redundant_ <- function(redundant, modulus_maxima) {
  if (length(redundant) == 0) {
    modulus_maxima
  } else {
    list(
      scale = modulus_maxima$scale[, -redundant],
      polarities = modulus_maxima$polarities[-redundant]
    )
  }
}


find_indexes_ <-
  function(scale3_maxima, polarity_criteria, i, is_first, sampling_frequency = 250) {
    interval_threshold <-
      if (is_first) {
        pracma::ceil(sampling_frequency * 0.12)
      } else {
        1.5 * pracma::ceil(0.12 * sampling_frequency)
      }

    which(
      (scale3_maxima > (scale3_maxima[i] - interval_threshold)) &
        (scale3_maxima < (scale3_maxima[i] + interval_threshold)) &
        polarity_criteria[i]
    )
  }


first_elimination_ <- function(modulus_maxima, scale3_wavelets) {

  list_criterion <-
    polarity_judge_(modulus_maxima$polarities, TRUE)

  scale3_maxima <- modulus_maxima$scale[3, ]

  unique_location_ <-
    function(locations, location) {
      if (any(locations == location)) {
        locations
      } else {
        index <-
          find_indexes_(scale3_maxima,
                        list_criterion$polarity_criteria, location, TRUE)

        amp_max <- abs(scale3_wavelets[scale3_maxima[index]])

        append(locations, index[-which.max(amp_max)])
      }
    }

  list_criterion$maxima_locations %>%
    reduce(unique_location_, .init = list()) %>%
    remove_redundant_(modulus_maxima)
}


remaining_elimination_ <-
  function(modulus_maxima, scale3_wavelets, is_positive, peak_criterion = 1.2) {
    scale3_maxima <- modulus_maxima$scale[3, ]

    list_criterion <-
      polarity_judge_(modulus_maxima$polarities, is_positive)

    # TODO: Need a better name for this helper function...
    max_k_locations_ <- function(locations, location) {

      indexes <-
        find_indexes_(scale3_maxima,
                      list_criterion$polarity_criteria, location, FALSE)

      if (length(indexes) <= 1) {
        locations
      } else {
        amp_max <- abs(scale3_wavelets[scale3_maxima[indexes]])
        distance <- scale3_maxima[indexes] - scale3_maxima[location]

        k <- amp_max / distance
        max_k_index <- which(k == max(k))

        max_k_location <-
          if (all((max(k) / k[-max_k_index]) > peak_criterion)) {
            indexes[-max_k_index]
          } else {
            min_distance <-
              which.min(abs(scale3_maxima[location] - scale3_maxima[indexes]))

            indexes[-min_distance]
          }

        append(locations, max_k_location)
      }
    }

    list_criterion$maxima_location %>%
      reduce(max_k_locations_, .init = list()) %>%
      remove_redundant_(modulus_maxima)
  }


eliminate_all_redundant_ <- function(modulus_maxima, scale3_wavelets) {
  modulus_maxima %>%
    first_elimination_(scale3_wavelets) %>%
    remaining_elimination_(scale3_wavelets, TRUE) %>%
    remaining_elimination_(scale3_wavelets, FALSE)
}


# Extra protection
# ==============================================================
qrs_extra_protection_ <-
  function(modulus_maxima, scale2_wavelets, sampling_frequency = 250) {
    search_interval <- round(0.1 * sampling_frequency)
    scale2_maxima <- modulus_maxima$scale[2, ]

    discarded_locations_ <- function(locations, location) {
      index <- scale2_maxima[location]

      before <- max(1, index - search_interval)
      after <- min(length(scale2_wavelets), index + search_interval)

      peakbefore <- peak_before(scale2_wavelets[before:index], index)
      peakafter <- peak_after(scale2_wavelets[index:after], index)

      if (length(peakbefore) == 0 && length(peakafter) == 0) {
        append(locations, location)
      } else {
        locations
      }
    }

    discarded <-
      1:ncol(modulus_maxima$scale) %>%
      reduce(discarded_locations_, .init = list())

    if (length(discarded) == 0) {
      modulus_maxima
    } else {
      list(
        scale = modulus_maxima$scale[, -discarded],
        polarities = modulus_maxima$polarities[, -discarded]
      )
    }
  }


# Detect zerocrossing at scale 1
# ===============================================================
find_zerocross_ <-
  function(scale1_maxima, scale1_wavelets, k, sampling_frequency = 250) {
    window_overlap <- pracma::ceil(0.125 * sampling_frequency)

    slice_begin <- scale1_maxima[k]

    slice_end <-
      min((scale1_maxima[k] + window_overlap), length(scale1_wavelets))

    index <- zerocross(scale1_wavelets[slice_begin:slice_end])
  }


is_out_of_bounds_ <- function(scale1_maxima, i) {
  (i <= 1) || (i >= length(scale1_maxima))
}


distance_from_neighbour_ <- function(scale1_maxima, location) {
  scale1_maxima[location] - scale1_maxima[location - 1]
}


is_cross_behind_ <- function(modulus_maxima, i, sampling_frequency = 250) {
  scale1_maxima <- modulus_maxima$scale[1, ]

  if (scale1_maxima %>% is_out_of_bounds_(i) && i != 1) {
    FALSE
  } else {
    interval_behind <- scale1_maxima %>% distance_from_neighbour_(i + 1)
    interval_ahead <- scale1_maxima %>% distance_from_neighbour_(i)

    is_negative_ahead <- modulus_maxima$polarities[i + 1] < 0
    is_positive_here <- modulus_maxima$polarities[i] > 0

    if (i == 1) {
      is_negative_ahead
    } else {
      (interval_behind < interval_ahead) &&
        (is_negative_ahead && is_positive_here)
    }
  }
}


is_cross_ahead_ <- function(modulus_maxima, i) {
  if (modulus_maxima$scale[1, ] %>% is_out_of_bounds_(i)) {
    FALSE
  } else {
    modulus_maxima$polarities[i - 1] < 0
  }
}


amplitude_difference_ <- function(scale1_maxima, scale1_wavelets, k) {
  abs(scale1_wavelets[scale1_maxima[k]] - scale1_wavelets[scale1_maxima[k + 1]])
}


maybe_crossing_location_ <- function(modulus_maxima, location) {
  if (modulus_maxima %>% is_cross_behind_(location)) {
    Just(location)
  } else if (modulus_maxima %>% is_cross_ahead_(location)) {
    Just(location - 1)
  }else{
    Nothing()
  }
}


detect_zerocross_ <- function(modulus_maxima, scale1_wavelets) {
  scale1_maxima <- modulus_maxima$scale[1, ]
  zerocrossings <- list(time = c(), Amp_diff = c())

  collect_zerocross_ <- function(i) {
    difference <- amplitude_difference_(scale1_maxima, scale1_wavelets, i)

    add_crossing_ <- function(location, crossing) {
      list(
        times =
          append(zerocrossings$time, crossing + scale1_maxima[location] - 1),

        Amp_diff =
          append(zerocrossings$Amp_diff, difference)
      )
    }

    find_zerocross_(scale1_maxima, scale1_wavelets, i) %>%
      maybe_map(partial(add_crossing_, i)) %>%
      maybe_default(zerocrossings)
  }

  for (i in which(modulus_maxima$polarities > 0)) {
    zerocrossings <-
      modulus_maxima %>%
      maybe_crossing_location_(i) %>%
      maybe_map(collect_zerocross_) %>%
      maybe_default(zerocrossings)
  }

  return(zerocrossings)
}


# this belongs in a fp-utils.R file...
apply_if <- function(x, predicate, callback) {
  p <- purrr::as_mapper(predicate)
  f <- purrr::as_mapper(callback)

  if (p(x)) {
    f(x)
  } else {
    x
  }
}


# these belong in a list-utils.R file...
is_not_empty <- function(x) {
  length(x) != 0
}


is_empty <- function(x) {
  length(x) == 0
}


# Refractory period for T and P waves
# ==========================================================
refractory_period_ <-
  function(zerocrossings, w, alignments, sampling_frequency = 250) {
    refractory_period <- pracma::ceil(sampling_frequency * 0.275)

    pwave_threshold <- 0.25 * sampling_frequency
    twave_threshold <- 0.45 * sampling_frequency

    is_head_below_pwave_ <- function(x) {
      x[1] <= pwave_threshold
    }

    is_head_aligned_ <- function(x) {
      is_not_empty(x) && alignments[x[1]] <= 1
    }

    is_tail_diff_below_twave_ <- function(x) {
      is_not_empty(x) && ((nrow(w) - tail(x, 1)) < twave_threshold)
    }

    times <- zerocrossings$times
    amp_diffs <- zerocrossings$Amp_diff

    min_index_ <- function(i) {
      which.min(amp_diffs[i:(i + 1)])
    }

    collect_indexes_ <- function(indexes, i) {
      if (min_index_(i) == 2) {
        append(indexes, i + 1)
      } else {
        append(indexes, i)
      }
    }

    if (is_empty(times)) {
      c()
    } else {
      time_diff <- times[-1] - times[-length(times)]

      indexes <-
        which(time_diff < refractory_period) %>%
        reduce(collect_indexes_, .init = list()) %>%
        unlist()

      times %>%
        apply_if(~is_not_empty(indexes), ~.x[-indexes]) %>%
        apply_if(is_head_below_pwave_, ~.x[-1]) %>%
        apply_if(is_head_aligned_, ~.x[-1]) %>%
        apply_if(is_tail_diff_below_twave_, ~.x[-length(.x)])
    }
  }


slice_before <- function(index, window_overlap) {
  max(1, index - window_overlap)
}


slice_after <- function(index, window_overlap, w_scale) {
  min(length(w_scale), index + window_overlap)
}


slice_search_before_ <- function(w_scale, index, sampling_frequency = 250) {
  before_overlap <- round(0.08 * sampling_frequency)

  slice_begin <- slice_before(index, before_overlap)
  signal_slice <- w_scale[slice_begin:(index - 1)]

  signal_slice[length(signal_slice):1]
}


zerocross_search_before_ <- function(scale1_wavelets, peak_index) {
  scale1_wavelets %>%
    slice_search_before_(peak_index) %>%
    zerocross() %>%
    maybe_map(partial(subtract, peak_index))
}


# slice_search_after_ <- function(w_scale, index, sampling_frequency = 250) {
#   after_overlap <- round(0.13 * sampling_frequency)
#
#   slice_end <- slice_after(index, after_overlap, w_scale)
#
#   signal_slice <- w_scale[(index + 1):slice_end]
# }
#
# zerocross_search_after_ <- function(scale1_wavelets, peak_index) {
#   scale1_wavelets %>%
#     slice_search_after_(peak_index) %>%
#     zerocross() %>%
#     maybe_map(partial(add, peak_index))
# }


scale_modulus_maxima_ <- function(scale4_maxima) {
  m <- matrix(0, 4, length(scale4_maxima))
  m[4, ] <- scale4_maxima
  return(m)
}


R_peak_refractory_protection_ <- function(R_peaks, sampling_frequency = 250) {
  discard_R <-
    which(R_peaks < 0.25 * sampling_frequency)

  R_peaks <-
    if (is_not_empty(discard_R)) {
      R_peaks[-discard_R]
    } else {
      R_peaks
    }
}

#' Detect R peak in ECG signals based on wavelet transform
#' Return a vector of index for R_peaks.
#'
#' @param wavelets The n * 5 wavelet transform matrix of ECG signal
#' @param thresholds A length = 5 vector for finding modulus maxima
#' in 5 scales in wavelet matrix
#' @param alignments An index vector for aligning indexes
#' @param sampling_frequency signal sampling rate, in samples per second
#'
#' @export
#'
wavelet_R_detection <-
  function(wavelets, thresholds, alignments, sampling_frequency = 250) {
    scale1_wavelets <- wavelets[, 1]
    scale2_wavelets <- wavelets[, 2]
    scale3_wavelets <- wavelets[, 3]

    scale4_maxima <- find_modulus_maxima(wavelets[, 4], thresholds[4], 0)

    wavelet_qrs_position <-
      if (length(scale4_maxima) == 0) {
        c()
      } else {
        modulus_maxima <-
          list(
            scale = scale_modulus_maxima_(scale4_maxima),
            polarities = sign(wavelets[scale4_maxima, 4])
          )

        modulus_maxima %>%
          find_neighbouring_maxima_(wavelets, thresholds, 4, 3) %>%
          find_neighbouring_maxima_(wavelets, thresholds, 3, 2) %>%
          find_neighbouring_maxima_(wavelets, thresholds, 2, 1) %>%
          eliminate_noise_(scale3_wavelets) %>%
          eliminate_isolation_() %>%
          eliminate_all_redundant_(scale3_wavelets) %>%
          eliminate_isolation_() %>%
          qrs_extra_protection_(scale2_wavelets) %>%
          detect_zerocross_(scale1_wavelets) %>%
          refractory_period_(wavelets, alignments)
      }

    qrs_window <- round(0.1 * sampling_frequency)

    collect_R_waves_ <- function(r_waves, qrs) {
      before_qrs <- slice_before(qrs, qrs_window)
      after_qrs <- slice_after(qrs, qrs_window, scale2_wavelets)

      peak_before_qrs <- scale2_wavelets[before_qrs:qrs] %>% peak_before(qrs)
      peak_after_qrs <- scale2_wavelets[qrs:after_qrs] %>% peak_after(qrs)

      if((length(peak_before_qrs) == 0) |
         (length(peak_after_qrs) == 0) |
         is.na(peak_before_qrs) |
         is.na(peak_after_qrs)){
        r_wave <- qrs
      }else{
        amp_peak_before_qrs <- scale2_wavelets[peak_before_qrs]
        amp_peak_after_qrs <- scale2_wavelets[peak_after_qrs]

        r_wave <-
          if (amp_peak_before_qrs < 0 & amp_peak_after_qrs > 0) {
            scale1_wavelets %>%
              zerocross_search_before_(peak_before_qrs) %>%
              maybe_default(qrs)
            # r_replacement
          } else {
            qrs
          }
      }

      # TODO:simplify
      # Try new protection measure for R detection
      # zerocross_before_qrs <-
      #   scale1_wavelets %>%
      #   zerocross_search_before_(peak_before_qrs) %>%
      #   maybe_default(qrs)
      #
      # zerocross_after_qrs <-
      #   scale1_wavelets %>%
      #   zerocross_search_after_(peak_after_qrs) %>%
      #   maybe_default(qrs)
      #
      # r_replacement <-
      #   if(processed_ecg[zerocross_before_qrs] > processed_ecg[zerocross_after_qrs]){
      #     zerocross_before_qrs
      #   }else{
      #     zerocross_after_qrs
      #   }



      append(r_waves, r_wave)
    }

    wavelet_R_index <-
      wavelet_qrs_position %>%
      reduce(collect_R_waves_, .init = list()) %>%
      as_vector() %>%
      R_peak_refractory_protection_() %>%
      add(alignments[1] - 1)
  }