Skip to content

Commit

Permalink
add in-place mt_pgram!, and multitaper spectrogram, multitaper csd,…
Browse files Browse the repository at this point in the history 
… multitaper coherence (#401)

* Initial multitaper spectrogram code ported from a private Beacon repo

Co-authored-by: Alex Arslan <[email protected]>
Co-authored-by: Kalyan Palepu <[email protected]>
Co-authored-by: kolia <[email protected]>

* fix size of `tmp`, rename fields, return spectrogram object

* format

* port WIP coherence and cross spectral code

Co-authored-by: Alex Arslan <[email protected]>
Co-authored-by: Kalyan Palepu <[email protected]>
Co-authored-by: kolia <[email protected]>

wip

wip

wip

wip

add WIP cross-spectra code so I don't lose it

Co-authored-by: Alex Arslan <[email protected]>
Co-authored-by: Kalyan Palepu <[email protected]>
Co-authored-by: kolia <[email protected]>

fix bug, change syntax for 1.0

format

wip

wip

wip

fix

Rename `sample_rate` -> `fs`

fixes

* transfer tests

Co-authored-by: Alex Arslan <[email protected]>
Co-authored-by: Kalyan Palepu <[email protected]>
Co-authored-by: kolia <[email protected]>

* add some MATLAB reference tests

Co-authored-by: mich11 <[email protected]>

swap order to match Onda

wip

align better with existing api

fix test

Fix MATLAB reference test

Co-authored-by: mich11 <[email protected]>

* fix csd bugs, add tests

fix bugs

rename

add coherence reference test

* remove `d`s since these aren't normalized as densities

* whoops, rm extraneous file

* fix typo

* infer `MTSpectrogramConfig` eltype from `MTConfig` eltype

* rm `n_freq_bins`, `n_time_points`, add some checks to the MTSpectrogramConfig constructor

* try to fix bugs...

* fix find-and-replace bug

* reformat

* test magically works now (?)

* handle `n_samples < samples_per_window`

* tweak `coherence_from_cs!`

* improve test coverage

* tweak `MTSpectrogramConfig` API, add `mt_spectrogram` method, improve docstrings

* coherence should be 1 on the diagonal

* allow and test mixed-precision operations

* allocate less, use aligned memory

* make `demean=false` the default

* don't export `allocate_output`

* don't shadow `DSP.freq`

* don't average over frequencies in the coherence

* Update src/periodograms.jl

* add note crediting MNE-python

* change names, allow taper weight configuration

* update docstrings

* adjust docstring, make loadable on 1.0

* fixes

* add test and note that two-sided FFTs not supported in CSD

* add tests

* argh Julia 1.0

* increase codecov

* fix very broken tests, make them more robust

* add docs, cleanup docstrings

* rm old keyword argument from docstring

* rm old comment

* should be slightly faster to do the `real` inside rather than outside

Co-authored-by: Alex Arslan <[email protected]>
Co-authored-by: Kalyan Palepu <[email protected]>
Co-authored-by: kolia <[email protected]>
Co-authored-by: mich11 <[email protected]>
  • Loading branch information
@ericphanson @ararslan
@KalyanPalepu @mich11
5 people authored May 22, 2021
1 parent 2bd197d commit ea75bfb
Show file tree
Hide file tree
Showing 18 changed files with 17,171 additions and 70 deletions.
24 changes: 23 additions & 1 deletion docs/src/periodograms.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,33 @@
arraysplit
periodogram(s::AbstractVector{T}) where T <: Number
welch_pgram
mt_pgram
spectrogram
stft
periodogram(s::AbstractMatrix{T}) where T <: Real
freq
power
time
coherence
```

## Multitaper periodogram estimation

```@docs
mt_pgram
mt_pgram!
mt_spectrogram
mt_spectrogram!
mt_cross_power_spectra
mt_cross_power_spectra!
mt_coherence
mt_coherence!
```

### Configuration objects

```@docs
MTConfig
MTSpectrogramConfig
MTCrossSpectraConfig
MTCoherenceConfig
```
5 changes: 5 additions & 0 deletions src/DSP.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,11 @@ using IterTools: subsets

export conv, deconv, filt, filt!, xcorr

# This function has methods added in `peridograms` but is not exported,
# so we define it here so one can do `DSP.allocate_output` instead of
# `DSP.Periodograms.allocate_output`.
function allocate_output end

include("dspbase.jl")

include("util.jl")
Expand Down
825 changes: 825 additions & 0 deletions src/multitaper.jl
View file

Large diffs are not rendered by default.

91 changes: 27 additions & 64 deletions src/periodograms.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,17 @@
module Periodograms
using LinearAlgebra: mul!
using ..Util, ..Windows
export arraysplit, nextfastfft, periodogram, welch_pgram, mt_pgram,
spectrogram, power, freq, stft
using Statistics: mean!
export arraysplit, nextfastfft, periodogram, welch_pgram,
spectrogram, power, freq, stft,
MTConfig, mt_pgram, mt_pgram!,
MTSpectrogramConfig, mt_spectrogram, mt_spectrogram!,
MTCrossSpectraConfig, mt_cross_power_spectra, mt_cross_power_spectra!,
MTCoherenceConfig, mt_coherence, mt_coherence!,
coherence
import ..DSP: allocate_output
using FFTW
import FFTW: Frequencies, fftfreq, rfftfreq

## ARRAY SPLITTER

struct ArraySplit{T<:AbstractVector,S,W} <: AbstractVector{Vector{S}}
Expand All @@ -22,7 +28,7 @@ struct ArraySplit{T<:AbstractVector,S,W} <: AbstractVector{Vector{S}}
# n = noverlap is a problem - the algorithm will not terminate.
(0 ≤ noverlap < n) || error("noverlap must be between zero and n")
nfft >= n || error("nfft must be >= n")
new{Ti,Si,Wi}(s, zeros(Si, nfft), n, noverlap, window, div((length(s) - n), n - noverlap)+1)
new{Ti,Si,Wi}(s, zeros(Si, nfft), n, noverlap, window, length(s) >= n ? div((length(s) - n), n - noverlap)+1 : 0)
end
end
ArraySplit(s::AbstractVector, n, noverlap, nfft, window) =
Expand Down Expand Up @@ -62,7 +68,7 @@ Base.collect(x::ArraySplit) = collect(copy(a) for a in x)
## UTILITY FUNCTIONS

# Convert the output of an FFT to a PSD and add it to out
function fft2pow!(out::Array{T}, s_fft::Vector{Complex{T}}, nfft::Int, r::Real, onesided::Bool, offset::Int=0) where T
function fft2pow!(out::AbstractArray{T}, s_fft::AbstractVector{Complex{T}}, nfft::Int, r::Real, onesided::Bool, offset::Int=0) where T
m1 = convert(T, 1/r)
n = length(s_fft)
if onesided
Expand Down Expand Up @@ -184,34 +190,37 @@ end

## PERIODOGRAMS
abstract type TFR{T} end
struct Periodogram{T,F<:Union{Frequencies,AbstractRange}} <: TFR{T}
power::Vector{T}
struct Periodogram{T,F<:Union{Frequencies,AbstractRange}, V <: AbstractVector{T}} <: TFR{T}
power::V
freq::F
end
struct Periodogram2{T,F1<:Union{Frequencies,AbstractRange},F2<:Union{Frequencies,AbstractRange}} <: TFR{T}
power::Matrix{T}
struct Periodogram2{T,F1<:Union{Frequencies,AbstractRange},F2<:Union{Frequencies,AbstractRange}, M<:AbstractMatrix{T}} <: TFR{T}
power::M
freq1::F1
freq2::F2
end

"""
power(p)
For a Periodogram, returns the computed power at each frequency as
For a `Periodogram`, returns the computed power at each frequency as
a Vector.
For a Spectrogram, returns the computed power at each frequency and
For a `Spectrogram`, returns the computed power at each frequency and
time bin as a Matrix. Dimensions are frequency × time.
For a `CrossPowerSpectra`, returns the pairwise power between each pair
of channels at each frequency. Dimensions are channel x channel x frequency.
"""
power(p::TFR) = p.power

"""
freq(p)
Returns the frequency bin centers for a given Periodogram or
Spectrogram object.
Returns the frequency bin centers for a given `Periodogram`,
`Spectrogram`, `CrossPowerSpectra`, or `Coherence` object.
Returns a tuple of frequency bin centers for a given Periodogram2
Returns a tuple of frequency bin centers for a given `Periodogram2`
object.
"""
freq(p::TFR) = p.freq
Expand Down Expand Up @@ -381,62 +390,14 @@ function welch_pgram(s::AbstractVector{T}, n::Int=length(s)>>3, noverlap::Int=n>
Periodogram(out, onesided ? rfftfreq(nfft, fs) : fftfreq(nfft, fs))
end

"""
mt_pgram(s; onesided=eltype(s)<:Real, nfft=nextfastfft(n), fs=1, nw=4, ntapers=iceil(2nw)-1, window=dpss(length(s), nw, ntapers))
Computes the multitaper periodogram of a signal `s`.
If `window` is not specified, the signal is tapered with
`ntapers` discrete prolate spheroidal sequences with
time-bandwidth product `nw`. Each sequence is equally weighted;
adaptive multitaper is not (yet) supported.
If `window` is specified, each column is applied as a taper. The
sum of periodograms is normalized by the total sum of squares of
`window`.
See also: [`dpss`](@ref)
"""
function mt_pgram(s::AbstractVector{T}; onesided::Bool=eltype(s)<:Real,
nfft::Int=nextfastfft(length(s)), fs::Real=1,
nw::Real=4, ntapers::Int=ceil(Int, 2nw)-1,
window::Union{AbstractMatrix,Nothing}=nothing) where T<:Number
onesided && T <: Complex && error("cannot compute one-sided FFT of a complex signal")
nfft >= length(s) || error("nfft must be >= n")

if isa(window, Nothing)
window = dpss(length(s), nw, ntapers)
r::T = fs*ntapers
else
size(window, 1) == length(s) ||
error(DimensionMismatch("length of signal $(length(s)) must match first dimension of window $(size(window,1))"))
r = fs*sum(abs2, window)
end

out = zeros(fftabs2type(T), onesided ? (nfft >> 1)+1 : nfft)
input = zeros(fftintype(T), nfft)
tmp = Vector{fftouttype(T)}(undef, T<:Real ? (nfft >> 1)+1 : nfft)

plan = forward_plan(input, tmp)
for j = 1:size(window, 2)
for i = 1:size(window, 1)
@inbounds input[i] = window[i, j]*s[i]
end
mul!(tmp, plan, input)
fft2pow!(out, tmp, nfft, r, onesided)
end

Periodogram(out, onesided ? rfftfreq(nfft, fs) : fftfreq(nfft, fs))
end

## SPECTROGRAM

@static if isdefined(Base, :StepRangeLen)
const FloatRange{T} = StepRangeLen{T,Base.TwicePrecision{T},Base.TwicePrecision{T}}
end

struct Spectrogram{T,F<:Union{Frequencies,AbstractRange}} <: TFR{T}
power::Matrix{T}
struct Spectrogram{T,F<:Union{Frequencies,AbstractRange}, M<:AbstractMatrix{T}} <: TFR{T}
power::M
freq::F
time::FloatRange{Float64}
end
Expand Down Expand Up @@ -508,4 +469,6 @@ function stft(s::AbstractVector{T}, n::Int=length(s)>>3, noverlap::Int=n>>1,
out
end

include("multitaper.jl")

end # end module definition
Loading

0 comments on commit ea75bfb

Please sign in to comment.