Counting Apneas in EDF Annotations

tutorial_05.rb

# tutorial_05.rb
# sleepdata.org
# author: @remomueller
#
# Required Gems:
#
# gem install colorize nsrr xml-simple --no-document
#
# To Run Script:
#
# ruby tutorial_05.rb

require 'rubygems'
require 'colorize'
require 'csv'
require 'xmlsimple'
require 'nsrr'
require 'nsrr/commands/download'

# Take the XML files in this folder `/annotation-testing` and answer two
# questions:
#   1) How many epochs/minutes was the subject asleep?
#   2) How many obstructive apnea events were marked on the recording?
# Asleep is defined as SleepStage = 1,2,3,4,5

AWAKE_STAGES = [0]
SLEEP_STAGES = [1, 2, 3, 4, 5]
NREM_STAGES = [1, 2, 3, 4]
REM_STAGES = [5]

# Simplifies filtering an array of objects by attributes.
class Array
  def where(filters)
    results = dup
    filters.each do |method, value|
      if value.is_a?(Array)
        results.select! { |object| value.include?(object.send(method)) }
      else
        results.select! { |object| object.send(method) == value }
      end
    end
    results
  end
end

# Represents a single scored event in an EDF, along with information present in
# the signal at the location of the event.
class ScoredEvent
  attr_accessor :name, :lowest_spo2, :desaturation, :start, :duration, :input,
                :stage

  def initialize(hash = {})
    @name = hash[:name]
    @lowest_spo2 = hash[:lowest_spo2]
    @desaturation = hash[:desaturation]
    @start = hash[:start]
    @duration = hash[:duration]
    @input = hash[:input]
    @stage = nil
  end

  def set_stage(sleep_stages, epoch_length)
    return if @start.nil?
    index = (@start / epoch_length).floor
    @stage = sleep_stages[index]
  end
end

# Represents a set of Annotation of an EDF, including sleep stages, epoch
# length, and scored events.
class Annotation
  attr_accessor :filename, :xml, :epoch_length, :sleep_stages, :scored_events

  def initialize(filename)
    @filename = filename
    @xml = XmlSimple.xml_in(filename)
    @epoch_length = parse_epoch_length
    @sleep_stages = parse_sleep_stages
    @scored_events = parse_scored_events
  end

  private

  def parse_epoch_length
    @xml['EpochLength'][0].to_i
  end

  def parse_sleep_stages
    @xml['SleepStages'][0]['SleepStage'].collect(&:to_i)
  end

  def parse_scored_events
    @xml['ScoredEvents'][0]['ScoredEvent'].collect do |hash|
      hash = {
        name: extract(hash, 'Name'),
        lowest_spo2: extract(hash, 'LowestSpO2', convert: :to_f),
        desaturation: extract(hash, 'Desaturation', convert: :to_f),
        start: extract(hash, 'Start', convert: :to_f),
        duration: extract(hash, 'Duration', convert: :to_f),
        input: extract(hash, 'Input')
      }
      scored_event = ScoredEvent.new(hash)
      scored_event.set_stage(@sleep_stages, @epoch_length)
      scored_event
    end
  end

  def extract(hash, key, convert: nil)
    value = hash[key] ? hash[key][0] : nil
    convert && value ? value.send(convert) : value
  end
end

def annotations(recursive: true)
  return enum_for(:annotations, recursive: recursive) unless block_given?
  annotation_paths(recursive: recursive).each do |file_path|
    yield Annotation.new(file_path)
  end
end

def annotation_paths(recursive: true)
  path = "#{'**/' if recursive}*.xml"
  Dir.glob(path, File::FNM_CASEFOLD)
end

Nsrr::Commands::Download.run(%w(download learn/polysomnography/annotations-events-profusion/learn-nsrr01-profusion.xml --token=public))
Nsrr::Commands::Download.run(%w(download learn/polysomnography/annotations-events-profusion/learn-nsrr02-profusion.xml --token=public))
Nsrr::Commands::Download.run(%w(download learn/polysomnography/annotations-events-profusion/learn-nsrr03-profusion.xml --token=public))

total_annotations = annotation_paths.count
csv_path = 'annotations.csv'

CSV.open(csv_path, 'wb') do |csv|
  csv << [
    'Annotation', 'Asleep (min)', 'Recording Time (min)',
    'OAs', 'REM OAs', 'NREM OAs',
    'CAs', 'REM CAs', 'NREM CAs',
    'HYs', 'REM HYs', 'NREM HYs'
  ]
  annotations.each_with_index do |annotation, index|
    asleep_stages = annotation.sleep_stages.where(itself: SLEEP_STAGES)
    minutes_asleep = asleep_stages.count * annotation.epoch_length / 60.0
    minutes_total = annotation.sleep_stages.count * annotation.epoch_length / 60.0

    oa_events = annotation.scored_events.where(name: 'Obstructive Apnea')
    rem_oa_events = oa_events.where(stage: REM_STAGES)
    nrem_oa_events = oa_events.where(stage: NREM_STAGES)

    ca_events = annotation.scored_events.where(name: 'Central Apnea')
    rem_ca_events = ca_events.where(stage: REM_STAGES)
    nrem_ca_events = ca_events.where(stage: NREM_STAGES)

    hy_events = annotation.scored_events.where(name: 'Hypopnea')
    rem_hy_events = hy_events.where(stage: REM_STAGES)
    nrem_hy_events = hy_events.where(stage: NREM_STAGES)

    # puts message
    print "\rFile #{index + 1} of #{total_annotations}: #{annotation.filename}         "

    csv << [
      annotation.filename, minutes_asleep, minutes_total,
      oa_events.count, rem_oa_events.count, nrem_oa_events.count,
      ca_events.count, rem_ca_events.count, nrem_ca_events.count,
      hy_events.count, rem_hy_events.count, nrem_hy_events.count
    ]
  end
end

puts "\n\nAnnotation apnea counts exported to " + "#{csv_path}".colorize(:white) + "."

puts "\nCongrats! You have completed your fifth tutorial!".colorize(:green).on_white