Refactor data selection to a common function

src/Telepathy.jl

@@ -35,11 +35,17 @@ include("io/events.jl")
 export find_events!
 
 include("channels/channels.jl")
-export channel_names, get_channels, set_type!
+export channel_names, set_type!
 
 include("channels/layout.jl")
 export read_layout, set_layout!
 
+include("data/selection.jl")
+export select
+
+include("data/aggregation.jl")
+export aggregate
+
 include("preprocessing/filter.jl")
 export filter_data, filter_data!, design_filter
 

src/channels/channels.jl

@@ -2,87 +2,32 @@ function channel_names(rec::Recording)
     return rec.chans.name
 end
 
-# Generic versions for use without EEG objects
-get_channels(data::AbstractArray, chanID::Integer) = get_channels(data, chanID:chanID)
-get_channels(data::AbstractArray, chanRange::UnitRange) = collect(intersect(1:size(data, 2), chanRange))
 
-# Selection based on integer indices
-get_channels(rec::Recording, chanID::Integer) = get_channels(rec, chanID:chanID) 
-get_channels(rec::Recording, chanRange::UnitRange) = collect(intersect(1:length(rec.chans.name), chanRange))
-get_channels(rec::Recording, chanRange::AbstractVector{<:Integer}) = intersect(chanRange, 1:length(rec.chans.name))
 
-# Selection based on string channel names
-# Selecting the first element make the slicing return a Vector rather than a Matrix
-function get_channels(rec::Recording, name::String)
-    idx = get_channels(rec, [name])
-    if isempty(idx)
-        error("Channel $name not found in data.")
-    else
-        return idx[1]
-    end
-end
-
-function get_channels(rec, names::Vector{String})
-    return findall(x -> x in names, rec.chans.name)
-end
-
-# Selection based on symbol channel type
-get_channels(rec::Recording, type::Symbol) = get_channels(rec, [type])
-function get_channels(rec, types::Vector{Symbol})
-    # We only eval the types, no function calls are made, so hopefully this reduces side effects
-    types = [eval(:(Telepathy.$type)) for type in types]
-    return findall(x -> typeof(x) in types, rec.chans.type)
-end
-
-# Selection based on direct channel type
-function get_channels(rec::Recording, type::Sensor)
-    return findall(x -> x==type, rec.chans.type)
-end
-get_channels(rec::Recording, channels::Colon) = 1:size(rec.data, 2)
-
-# TODO: Add info in docs that using floats needs to specify the step fine enough to get the desired decimal places
-get_times(raw::Raw, times::AbstractFloat) = get_times(raw, times-1:times)
-
-function get_times(raw::Raw, times::AbstractRange; anchor::Number=0)
-    if typeof(anchor) <: AbstractFloat
-        anchor = round(Int64, anchor*get_srate(raw))
-    elseif !(typeof(anchor) <: Integer)
-        error("Anchor must be an integer or float.")
-    end
-
-    start = round(Int64, times[begin]*get_srate(raw) + 1 + anchor)
-    finish = round(Int64, times[end]*get_srate(raw) + anchor)
-    return UnitRange(start, finish)
-end
-
-get_times(raw::Raw, times::Colon) = 1:size(raw.data, 1)
-
-# Convert to range to not loose precision
-get_times(raw::Raw, start::AbstractFloat, stop::AbstractFloat; kwargs...) = get_times(raw, start:(stop-start):stop; kwargs...)
-
-# Selection of data segments, if they exist
-get_segments(raw::Raw) = 1:1
-get_segments(epochs::Epochs) = 1:size(epochs.data, 3)
 
 # Separate type unions for times and channel selectors to check for input order in get_data
-rowTypes = Union{AbstractFloat, AbstractRange, Colon}
+rowTypes = Union{AbstractFloat, AbstractRange{<:AbstractFloat}}
 colTypes = Union{AbstractString, Symbol, Sensor, Integer,
-                 AbstractVector{<:Union{AbstractString, Symbol, Sensor, Integer}}, Colon}
+                 AbstractVector{<:Union{AbstractString, Symbol, Sensor, Integer}}}
+
+get_data(rec::Recording, times, chans) = rec.data[_get_data(rec, times, chans)...]
 
 # We are covering both orders of selector inputs for convenience
-get_data(raw::Raw, first::Colon) = error("Ambigous selection, please specify both times and channels.")
-get_data(raw::Raw, first::rowTypes) = get_data(raw, first, :)
-get_data(raw::Raw, first::colTypes) = get_data(raw, :, first)
-get_data(raw::Raw, first::colTypes, second::rowTypes) = get_data(raw, second, first)
-get_data(raw::Raw, first::Colon, second::Colon) = (first, second)
-function get_data(raw::Raw, first::rowTypes, second::colTypes)
-    return get_times(raw, first), get_channels(raw, second)
+_get_data(raw::Raw, first::Colon) = error("Ambigous selection, please specify both times and channels.")
+_get_data(raw::Raw, first::rowTypes) = _get_data(raw, first, :)
+_get_data(raw::Raw, first::colTypes) = _get_data(raw, :, first)
+_get_data(raw::Raw, first::colTypes, second::rowTypes) = _get_data(raw, second, first)
+_get_data(raw::Raw, first::Colon, second::Colon) = (first, second)
+_get_data(raw::Raw, first, second::Colon) = typeof(first) <: rowTypes ? (_get_times(raw, first), second) : (second, _get_channels(raw, first))
+_get_data(raw::Raw, first::Colon, second) = typeof(second) <: rowTypes ? (_get_times(raw, second), first) : (first, _get_channels(raw, second))
+function _get_data(raw::Raw, first::rowTypes, second::colTypes)
+    return _get_times(raw, first), _get_channels(raw, second)
 end
 
 set_type!(data, chans, type::Symbol) = set_type!(data, chans, eval(:(Telepathy.$type)))
 set_type!(data, chans, type::Type{<:Sensor}) = set_type!(data, chans, type())
 function set_type!(data, chans, type::Sensor)
-    chanIDs = get_channels(data, chans)
+    chanIDs = _get_channels(data, chans)
     for i in chanIDs
         data.chans.type[i] = type
     end

src/data/aggregation.jl

@@ -0,0 +1,77 @@
+# Generate an iterable that will allow to go through the data in specified segments
+function calculate_segments(data::Raw, timeSpan::Number; overlap=0)
+    nSamples, nChannels = size(data.data)
+    sRate = get_srate(data)
+
+    # If the timeSpan is 0, we return the whole data
+    timeSpan==0 && return 1:1, nSamples
+
+    offset = _get_sample(data, timeSpan-overlap)
+    sampleSpan = _get_sample(data, timeSpan)
+
+    startingPoints = 1:offset:cld(nSamples-sampleSpan, offset)*offset
+
+    return startingPoints, sampleSpan
+end
+
+function aggregate(raw::Raw, func::Function; mode=:channels, timeSpan=0., channels=:, overlap=0., threads=0)
+
+    # Get necessary indexes to slice the data into segments
+    segmentStarts, segmentSize = calculate_segments(raw, timeSpan, overlap=overlap)
+
+    # Aggregate either each channel separately or a whole segment
+    if mode == :channels
+        elements = _get_channels(raw, channels)
+        @debug "Aggregating $(length(segmentStarts)) segments of $(length(elements)) channels."
+    elseif mode == :segments
+        elements = [_get_channels(raw, channels)]
+        @debug "Aggregating $(length(segmentStarts)) segments reduced from $(length(elements[1])) channels."
+    else
+        error("Mode $mode not recognized.")
+    end
+
+    # Allocate memory for the aggregated data
+    aggregatedData = zeros(length(segmentStarts), length(elements))
+
+    Threads.@threads for (thrID, segmentIDs) in setup_workers(1:length(segmentStarts), threads)
+        @debug "Thread $thrID processing segments $segmentIDs"
+        for sID in segmentIDs
+            for (j, element) in enumerate(elements)
+                range = segmentStarts[sID]:segmentStarts[sID]+segmentSize-1
+                @views aggregatedData[sID, j] = func(raw.data[range, element])
+            end
+        end
+    end
+
+    return aggregatedData
+end
+
+function aggregate(epochs::Epochs, func::Function; mode=:channels, timeSpan=0, channels=:, segments=0, overlap=0., threads=0)
+
+    times, chans, segments = select(epochs, times=timeSpan, channels=channels, segments=segments, indices=true)
+
+    # Aggregate either each channel separately or a whole segment
+    if mode == :channels
+        elements = chans
+        @debug "Aggregating $(length(segments)) segments of $(length(elements)) channels."
+    elseif mode == :segments
+        elements = [chans]
+        @debug "Aggregating $(length(segments)) segments reduced from $(length(elements[1])) channels."
+    else
+        error("Mode $mode not recognized.")
+    end
+
+    # Allocate memory for the aggregated data
+    aggregatedData = zeros(length(segments), length(elements))
+
+    Threads.@threads for (thrID, segmentIDs) in setup_workers(1:length(segments), threads)
+        @debug "Thread $thrID processing segments $segmentIDs"
+        for sID in segmentIDs
+            for (j, element) in enumerate(elements)
+                @views aggregatedData[sID, j] = func(epochs.data[times, element, sID])
+            end
+        end
+    end
+
+    return aggregatedData
+end

src/data/selection.jl

@@ -0,0 +1,129 @@
+# SELECT DATA BASED ON TIME IN SECONDS
+# TODO: Add info in docs that using floats needs to specify the step fine enough to get the desired decimal places
+_get_times(rec::Recording, times::AbstractFloat) = _get_times(rec, times-1:times)
+
+_get_times(rec::Recording, times::Colon) = 1:size(rec.data, 1)
+
+# Convert to range to not loose precision
+_get_times(rec::Recording, start::AbstractFloat, stop::AbstractFloat; kwargs...) = _get_times(rec, start:(1/rec.chans.srate[1]):stop; kwargs...)
+_get_times(rec::Recording, times::Tuple{<:T, <:T}; kwargs...) where T <: AbstractFloat = _get_times(rec, times[1], times[2]; kwargs...)
+
+function _get_times(raw::Raw, times::StepRangeLen; anchor::Number=0)
+    # Check if range is not empty
+    if times.len == 1 
+        @warn """Time range contains only one point. 
+        In Julia, default step is 1, so you need to specify a smaller step 
+        to get ranges with distance different than a whole number, e.g. 0.2:0.1:0.7."""
+    end
+
+    # Check if anchor is an integer or float
+    if typeof(anchor) <: AbstractFloat
+        anchor = round(Int64, anchor*get_srate(raw))
+    elseif !(typeof(anchor) <: Integer)
+        error("Anchor must be an integer or float.")
+    end
+
+    start = round(Int64, times[begin]*get_srate(raw) + 1 + anchor)
+    finish = round(Int64, times[end]*get_srate(raw) + 1 + anchor)
+    return UnitRange(start, finish)
+end
+
+_get_sample(raw::Raw, time::Number) = round(Int64, time*get_srate(raw))
+
+function _offset_time(epochs::Epochs, time::AbstractFloat) 
+    return findfirst(x -> isapprox(x, time, atol=1/(epochs.chans.srate[1]*2)), epochs.times)
+end
+
+function _get_times(epochs::Epochs, times::StepRangeLen; anchor::Number=0.)
+    # Check if range is not empty
+    if times.len == 1 
+        @warn """Time range contains only one point. 
+        In Julia, default step is 1, so you need to specify a smaller step 
+        to get ranges with distance different than a whole number, e.g. 0.2:0.1:0.7."""
+    end
+
+    # Check if anchor is an integer or float
+    if typeof(anchor) <: Integer
+        anchor = epochs.times[anchor]
+    elseif typeof(anchor) <: AbstractFloat
+        if anchor != 0.
+            epochs.times[begin] <= anchor <= epochs.times[end] || error("Anchor is outside of the epoch range.")
+        end
+    else
+        error("Anchor must be an integer or float.")
+    end
+
+    start = _offset_time(epochs, times[begin] + anchor)
+    finish = _offset_time(epochs, times[end] + anchor)
+
+    isnothing(start) && error("Selection start is outside of the epoch range.")
+    isnothing(finish) && error("Selection end is outside of the epoch range.")
+    return UnitRange(start, finish)
+end
+
+# SELECT DATA BASED ON CHANNEL NAMES, TYPES OR INDICES
+# Generic versions for use without EEG objects
+_get_channels(data::AbstractArray, chanID::Integer) = _get_channels(data, chanID:chanID)
+_get_channels(data::AbstractArray, chanRange::UnitRange) = collect(intersect(1:size(data, 2), chanRange))
+
+# Selection based on integer indices
+_get_channels(rec::Recording, chanID::Integer) = _get_channels(rec, chanID:chanID) 
+_get_channels(rec::Recording, chanRange::UnitRange) = collect(intersect(1:length(rec.chans.name), chanRange))
+_get_channels(rec::Recording, chanRange::AbstractVector{<:Integer}) = intersect(chanRange, 1:length(rec.chans.name))
+
+# Selection based on string channel names
+# Selecting the first element make the slicing return a Vector rather than a Matrix
+function _get_channels(rec::Recording, name::String)
+    idx = _get_channels(rec, [name])
+    if isempty(idx)
+        error("Channel $name not found in data.")
+    else
+        return idx[1]
+    end
+end
+
+function _get_channels(rec, names::Vector{String})
+    return findall(x -> x in names, rec.chans.name)
+end
+
+function _get_channels(rec, regex::Regex)
+    return (1:length(rec.chans.name))[occursin.(regex, rec.chans.name)]
+end
+
+# Selection based on symbol channel type
+_get_channels(rec::Recording, type::Symbol) = _get_channels(rec, [type])
+function _get_channels(rec, types::Vector{Symbol})
+    # We only eval the types, no function calls are made, so hopefully this reduces side effects
+    types = [eval(:(Telepathy.$type)) for type in types]
+    return findall(x -> typeof(x) in types, rec.chans.type)
+end
+
+# Selection based on direct channel type
+function _get_channels(rec::Recording, type::Sensor)
+    return findall(x -> x==type, rec.chans.type)
+end
+_get_channels(rec::Recording, channels::Colon) = 1:size(rec.data, 2)
+
+
+# SELECT DATA BASED ON SEGMENT INDICES OR EVENTS
+# Selection of data segments, if they exist
+_get_segments(raw::Raw, args...) = 1
+_get_segments(epochs::Epochs) = _get_segments(epochs::Epochs, colon::Colon)
+_get_segments(epochs::Epochs, colon::Colon) = 1:size(epochs.data, 3)
+_get_segments(epochs::Epochs, segment::Integer) = _get_segments(epochs::Epochs, [segment])
+_get_segments(epochs::Epochs, segments::AbstractVector{<:Integer}) = intersect(segments, 1:size(epochs.data, 3))
+
+
+# PUBLIC INTERFACE FOR DATA SELECTION
+function select(rec::Recording; times=0, channels=0, segments=0, indices=false, anchor=0.)
+
+    times != 0 ? times = _get_times(rec, times, anchor=anchor) : times = _get_times(rec, :)
+    channels != 0 ? channels = _get_channels(rec, channels) : channels = _get_channels(rec, :)
+    segments != 0 ? segments = _get_segments(rec, segments) : segments = _get_segments(rec, :)
+
+    if indices
+        return times, channels, segments
+    else
+        return rec.data[times, channels, segments]
+    end
+end

src/io/load_data.jl

@@ -110,7 +110,7 @@ function parse_filters(flt::String)
 end
 
 function parse_status!(raw::Raw{BDF})
-    idx = get_channels(raw, "Status")
+    idx = _get_channels(raw, "Status")
     if length(idx) == 1
         raw.status["lowTrigger"], raw.status["highTrigger"], raw.status["status"] = parse_status(raw.data[:,idx[1]])
     elseif length(idx) == 0

src/preprocessing/filter.jl

@@ -293,7 +293,7 @@ function filter_data!(raw::Raw; highPass=0, lowPass=0,
     digFilter = design_filter(highPass, lowPass, srate, window, transition, passErr, stopErr)
 
     # Apply the filter to EEG channels
-    chans = get_channels(raw, :EEG)
+    chans = _get_channels(raw, :EEG)
     filter_data!(raw.data, digFilter, chans, srate, nThreads)
 
     # Update the channel information