R/wavelet-P-detection.R

In Teresa00/hfmAnnotation: HFM Internal Package for Annotations

# Teresa updated March 29th 2018

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

# some convenience constants for describing search directions...
BEFORE <-
  TRUE

AFTER <-
  FALSE

POSITIVE <-
  TRUE

NEGATIVE <-
  FALSE


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

  if (length(signal_) > 2) {
    find_modulus_maxima(signal_, threshold, polarity)
  } else {
    list()
  }
}


is_onset_extrema_ <-
  function(window_begin, window_end, polarity) {
    is_onset_ <-
      ifelse(polarity, `>=`, `<=`)

    has_expected_polarity_ <-
      partial(ifelse(polarity, `>`, `<`), 0)

    is_onset_(window_begin, window_end) &&
      has_expected_polarity_(window_begin)
  }


is_begin_extrema_ <-
  function(w_scale, search_window, polarity) {
    is_onset_extrema_(
      w_scale[search_window$begin],
      w_scale[search_window$end],
      polarity
    )
  }


is_end_extrema_ <-
  function(w_scale, search_window, polarity) {
    is_onset_extrema_(
      w_scale[search_window$end],
      w_scale[search_window$begin],
      polarity
    )
  }


define_extrema_ <- function(w_scale, extrema_indexes, is_max) {
  assert_that(!is.null(extrema_indexes))

  amplitudes <-
    w_scale[extrema_indexes]

  extrema_location_ <-
    function(a) {
      ifelse(is_max, which.max, which.min)(a) %>%
        partial(magrittr::extract, extrema_indexes)()
    }

  extrema_amplitude_ <-
    function(a) {
      ifelse(is_max, max, min)(a) %>% abs()
    }

  list(
    location =
      extrema_location_(amplitudes),

    amplitude =
      extrema_amplitude_(amplitudes)
  )
}


find_extrema_ <-
  function(extremas, w_scale, search_window, is_positive) {
    if (is_not_empty(extremas)) {
      define_extrema_(w_scale, search_window$begin + extremas, is_positive)
    } else if (is_begin_extrema_(w_scale, search_window, is_positive)) {
      list(location = search_window$begin, amplitude = abs(w_scale[search_window$begin]))
    } else if (is_end_extrema_(w_scale, search_window, is_positive)) {
      list(location = search_window$end, amplitude = abs(w_scale[search_window$end]))
    } else {
      list(location = 0, amplitude = 0)
    }
  }


PT_slopes_searching_ <-
  function(w_scale, R_peaks, R, search_window, search_scale, sampling_frequency = 250) {
    search_slice <-
      w_scale[(search_window$begin + 1):search_window$end]

    positive_maximums <-
      find_modulus_maxima_PT(search_slice, 0, +1)

    negative_minimums <-
      find_modulus_maxima_PT(search_slice, 0, -1)

    find_extreme_ <-
      function(extremas, search_direction) {
        find_extrema_(extremas, w_scale, search_window, search_direction)
      }

    highest <-
      find_extreme_(positive_maximums, POSITIVE)

    lowest <-
      find_extreme_(negative_minimums, NEGATIVE)

    interval <-
      0.11 * sampling_frequency

    is_PT_within_interval <-
      abs(highest$location - lowest$location) < interval

    before_R <-
      ifelse(R == 1, 1, R - 1)

    rpeaks_slice <-
      R_peaks[before_R]:R_peaks[R]

    amplitude_threshold <-
      0.02 * sqrt(mean(w_scale[rpeaks_slice]^2))

    is_PT_above_threshold <-
      highest$amplitude > amplitude_threshold &&
        lowest$amplitude > amplitude_threshold

    does_PT_exist <-
      is_PT_within_interval && is_PT_above_threshold

    list(
      positive_maximum =
        highest$location,

      negative_maximum =
        lowest$location,

      abs_max_amp =
        highest$amplitude,

      abs_min_amp =
        lowest$amplitude,

      does_wave_exist =
        does_PT_exist,

      search_scale =
        search_scale
    )
  }


extrema_around_index_ <-
  function(w_scale, pivot_index, search_window, search_direction, amplitude_threshold) {
    polarity <-
      ifelse(search_direction$is_positive, +1, -1)

    slice_start <-
      ifelse(search_direction$is_before, search_window$begin, pivot_index)

    slice_end <-
      ifelse(search_direction$is_before, pivot_index, search_window$end)

    wavelet_slice <-
      w_scale[(slice_start + 1):slice_end]

    extrema_set <-
      find_modulus_maxima_PT(wavelet_slice, 0, polarity)

    extrema <-
      if (is_not_empty(extrema_set)) {
        define_extrema_(w_scale, slice_start + extrema_set, search_direction$is_positive)
      } else {
        list(location = 0, amplitude = 0)
      }

    # Amplitude Protection
    ifelse(extrema$amplitude < amplitude_threshold, 0, extrema$location)
  }


detect_zerocross_PT_ <-
  function(w_zerocross, slice_begin, slice_end) {
    slice_length <-
      slice_end - slice_begin

    if (slice_length <= 2) {
      slice_begin
    } else {
      w_zerocross %>%
        extract(slice_begin:slice_end) %>%
        zerocross() %>%
        maybe_default(NaN) %>%
        add(slice_begin - 1)
    }
  }


define_morphology_Pwave_ <-
  function(w_scale, w_zerocross, search_window, p_slopes) {
    ps <-
      p_slopes

    PM <-
      ps$positive_maximum

    NM <-
      ps$negative_maximum

    is_minima <-
      ps$abs_max_amp > ps$abs_min_amp

    is_maxima <-
      !is_minima

    is_min_before <-
      PM < NM

    is_min_after <-
      PM >= NM

    is_max_before <-
      NM < PM

    is_max_after <-
      NM >= PM

    amplitude_threshold <-
      0.125 * ifelse(is_minima, ps$abs_max_amp, ps$abs_min_amp)

    find_min_around_pm_ <-
      function(is_before) {
        search_direction <-
          list(is_before = is_before, is_positive = FALSE)

        extrema_around_index_(
          w_scale,
          PM,
          search_window,
          search_direction,
          amplitude_threshold
        )
      }

    find_max_around_nm_ <-
      function(is_before) {
        search_direction <-
          list(is_before = is_before, is_positive = TRUE)

        extrema_around_index_(
          w_scale,
          NM,
          search_window,
          search_direction,
          amplitude_threshold
        )
      }

    min_before_pm <-
      ifelse(is_minima && is_min_before, find_min_around_pm_(BEFORE), 0)

    min_after_pm <-
      ifelse(is_minima && is_min_after, find_min_around_pm_(AFTER), 0)

    max_before_nm <-
      ifelse(is_maxima && is_max_before, find_max_around_nm_(BEFORE), 0)

    max_after_nm <-
      ifelse(is_maxima && is_max_after, find_max_around_nm_(AFTER), 0)

    around_extrema <-
      c(min_before_pm, min_after_pm, max_before_nm, max_after_nm)

    all_zero_around_extrema <-
      all(around_extrema == 0)

    NORMAL_P_WAVE <-
      all_zero_around_extrema && PM < NM

    INVERTED_P_WAVE <-
      all_zero_around_extrema && PM >= NM

    BIPHASIC_PLUS_MINUS <-
      is_maxima && !all_zero_around_extrema

    BIPHASIC_MINUS_PLUS <-
      is_minima && !all_zero_around_extrema

    extremas <-
      if (NORMAL_P_WAVE) {
        c(PM, NM)
      } else if (INVERTED_P_WAVE) {
        c(NM, PM)
      } else if (BIPHASIC_PLUS_MINUS) {
        sort(c(PM, NM, max_before_nm, max_after_nm)) %>%
          keep(~.x != 0)
      } else if (BIPHASIC_MINUS_PLUS) {
        sort(c(PM, NM, min_before_pm, min_after_pm)) %>%
          keep(~.x != 0)
      } else {
        sort(c(PM, NM, min_before_pm, min_after_pm, max_before_nm, max_after_nm)) %>%
          keep(~.x != 0)
      }

    onset <-
      ifelse(BIPHASIC_MINUS_PLUS, 2, 1)

    list(
      p =
        detect_zerocross_PT_(w_zerocross, extremas[onset], extremas[onset + 1]),

      type =
        if (NORMAL_P_WAVE) {
          1
        } else if (INVERTED_P_WAVE) {
          2
        } else if (BIPHASIC_PLUS_MINUS) {
          3
        } else if (BIPHASIC_MINUS_PLUS) {
          4
        } else { # no-type
          0
        },

      peak_onset =
        extremas[onset]
    )
  }


#' Search for P wave in ECG signals based on wavelet transform
#' Return a vector of P wave locations (based on processed ECG signal)
#' and corresponding type of morphology
#' 1 - normal; 2 - inverse; 3 - biphasic +/-; 4 - biphasic -/+; 0 - no type
#'
#' TODO: Add details, perhaps better description
#'
#' @param R_peaks An index vector of R peaks
#' @param wavelets The n * 5 wavelet transform matrix of ECG signal
#' @param alignments An index vector for aligning indexes
#' @param sampling_frequency The sampling frequency in samples/second (default 250)
#'
#' @export
#'
wavelet_P_detection <-
  function(R_peaks, wavelets, alignments, sampling_frequency = 250) {
    fourth_scale <-
      4

    fifth_scale <-
      5

    # For aligning the indexes in wavelet transform vectors
    rpeaks <-
      R_peaks %>%
        add(1 - alignments[1])

    begin_window_margin <-
      round(0.25 * sampling_frequency)

    end_window_margin <-
      floor(0.065 * sampling_frequency)

    detect_P_wave_near_R_ <-
      function(R) {
        search_window <-
          list(
            begin =
              rpeaks[R] - begin_window_margin,

            end =
              rpeaks[R] - end_window_margin
          )

        fourth_scale_slopes <-
          wavelets[, fourth_scale] %>%
            PT_slopes_searching_(rpeaks, R, search_window, fourth_scale)

        p_slopes <-
          if (fourth_scale_slopes$does_wave_exist) {
            fourth_scale_slopes
          } else {
            wavelets[, fifth_scale] %>%
              PT_slopes_searching_(rpeaks, R, search_window, fifth_scale)
          }

        search_scale <-
          p_slopes$search_scale

        search_zerocross_at_ <-
          function(next_scale) {
            define_morphology_Pwave_(
              wavelets[, search_scale],
              wavelets[, next_scale],
              search_window,
              p_slopes
            )
          }

        if (!p_slopes$does_wave_exist) {
          list(p = NaN, type = NaN, peak_onset = NaN)
        } else {
          # search zerocross at the scale below the slope search scale
          zerocross_pwave <-
            search_zerocross_at_(search_scale - 1)

          if (!is.nan(zerocross_pwave$p)) {
            zerocross_pwave
          } else {
            # search zerocross at the same scale with the slope search scale
            search_zerocross_at_(search_scale)
          }
        }
      }

    collect_element_ <-
      function(p_waves, key) {
        sapply(p_waves, function(p) {
          extract2(p, key)
        })
      }

    p_waves <-
      lapply(seq_along(rpeaks), detect_P_wave_near_R_)

    realignment <-
      alignments[1] - 1

    list(
      P =
        p_waves %>%
          collect_element_("p") %>%
          add(realignment),

      P_onset =
        p_waves %>%
          collect_element_("peak_onset") %>%
          add(realignment),

      type_of_P =
        p_waves %>%
          collect_element_("type")
    )
  }