synthesis bugfix

Project.toml

@@ -1,13 +1,13 @@
 name = "STFT"
 uuid = "58bb99bf-048b-48b7-93e7-1cbf3ee61509"
-authors = ["Szymon M. Woźniak"]
-version = "1.0.0"
+authors = ["Szymon M. Woźniak <s@zymon.org>"]
+version = "1.3.0"
 
 [deps]
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
-
+LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 
 [compat]
-
-FFTW = "1.4"
+FFTW = "1.7"
+LoopVectorization = "0.12.101"
 julia = "1.6"

README.md

@@ -12,16 +12,20 @@ Check the [documentation](https://docs.zymon.org/STFT.jl/) for more insights.
 
 ## Installation
 
-The package is currently not available in [General](https://github.com/JuliaRegistries/General), the default Julia package registry.
-
-To install it, use the following command in Julia package manager:
+The package is currently available in General, the default Julia package registry.
+To install this package from General registry, use the following command in Julia REPL:
 ```julia
-pkg> add https://github.com/s-zymon/STFT.jl
+] add STFT
 ```
+Alternatively, directly via repository:
+```julia
+pkg> add https://codeberg.org/zymon/STFT.jl
+```
+
 
 ## Examples
 
-Below you can find a few standalone example with basic usage of the package.
+Below you can find a few standalone examples with basic usage of the package.
 
 ### Show spectrogram
 

src/STFT.jl

@@ -1,12 +1,25 @@
 module STFT
 
+using LoopVectorization
 using FFTW
 
 export stft, istft
 
-_fft(x::AbstractMatrix{<:Real}, d) = rfft(x, d)
-_fft(x::AbstractMatrix{<:Complex}, d) = fft(x, d)
+_fft(x::AbstractArray{<:Real}, d) = rfft(x, d)
+_fft(x::AbstractArray{<:Complex}, d) = fft(x, d)
 
+function malloc_stft(
+    x::M,
+    w::V,
+    L::I = zero(I),
+    N::I = length(w);
+) where {T<:Number, I<:Integer, V<:AbstractVector{T}, M<:AbstractMatrix{T}}
+    X, K = size(x)      # Length of the signal in samples
+    W = length(w)       # Length of the window in samples
+    S = (X-L) ÷ (W - L) # Number of segments
+    N = N < W ? W : N   # DFT size
+    zeros(T, (N, S, K)) # Allocate container for signal segments
+end
 
 
 doc_analysis = """
@@ -94,24 +107,31 @@ function analysis() end
 stft(x, w, L=0, N=length(w)) = analysis(x, w, L, N)
 
 function analysis(
-    x::AbstractVector{T},
-    w::AbstractVector{T},
-    L::Integer = 0,
-    N::Integer = length(w);
-)::AbstractMatrix{<:Complex} where {T<:Number}
-    X = length(x)       # Length of the signal in samples
-    W = length(w)       # Length of the window in samples
-    H = W - L           # Hop
-    S = (X-L) ÷ H       # Number of segments
-    N = N < W ? W : N   # DFT size
-    sc = zeros(T, N, S) # Allocate container for signal segments
+    x::M,
+    w::V,
+    L::I = zero(I),
+    N::I = length(w);
+) where {T<:Number, I<:Integer, V<:AbstractVector{T}, M<:AbstractMatrix{T}}
+    sc = malloc_stft(x, w, L, N)
+    N, S, K = sc |> size
+    W = w |> length
 
-    for s ∈ 1:S, n ∈ 1:W # Slice the signal
-        sc[n, s] = w[n] * x[(s-1)*H+n]
+    @turbo for k ∈ 1:K, s ∈ 1:S, n ∈ 1:W
+        sc[n, s, k] = w[n] * x[(s-1)*(W-L)+n, k]
     end
     _fft(sc, 1) # Convert segments to frequency-domain
 end
 
+function analysis(
+    x::V,
+    w::V,
+    L::I = zero(I),
+    N::I = length(w);
+)::Matrix{T |> complex} where {T<:Number, I<:Integer, V<:AbstractVector{T}}
+    xx = @view x[:,:]
+    @view analysis(xx, w, L, N)[:, :, 1]
+end
+
 function analysis(
     xs::AbstractArray{<:AbstractVector{T}},
     w::AbstractVector{T},
@@ -211,7 +231,7 @@ xr = istft(X, w, L) # Synthesis
 function synthesis() end
 
 "$doc_synthesis"
-isftf(X, w, L=0, N=length(w)) = synthesis(X, w, L, N)
+istft(X, w, L=0, N=length(w)) = synthesis(X, w, L, N)
 
 function synthesis(
     X::AbstractMatrix{<:Complex},
@@ -229,14 +249,32 @@ function synthesis(
 
     xs = irfft(X, N, 1) # Convert segments to time-domain
 
-    for s ∈ 1:S
-        ss = (s-1)*H # Segment start
+    @turbo for s ∈ 0:(S-1)
+        ss = s*H # Segment shift
         for n = 1:W
-            xn[ss+n] += xs[n, s] * w[n]
+            xn[ss+n] += xs[n, s+1] * w[n]
             xd[ss+n] += w²[n]
         end
     end
     xn ./ xd # Normalize
 end
 
+"""
+Real-value signal STFT with constant input size.
+"""
+function rSTFTm(A, w, L, N=length(w))
+    mem = STFT.malloc_stft(A, w, L, N)
+    N, S, K = mem |> size
+    W = w |> length
+    P = plan_rfft(mem, 1)
+    function f(x)
+        @turbo for k ∈ 1:K, s ∈ 1:S, n ∈ 1:W
+            mem[n, s, k] = w[n] * x[(s-1)*(W-L)+n, k]
+        end
+        P * mem
+    end
+    return f
+end
+
+
 end # module