feat(dsp): add Zak transform

src/dsp.jl

@@ -6,7 +6,7 @@ import Optim: optimize, minimizer, BFGS
 export fir, removedc, removedc!, demon
 export upconvert, downconvert, rrcosfir, rcosfir
 export mseq, gmseq, circconv, goertzel, pll, hadamard
-export mfilter, findsignal
+export mfilter, findsignal, zak, izak
 export istft, whiten, filt, filtfilt, resample, delay!, compose
 
 """
@@ -793,3 +793,63 @@ function compose(r, t, a; duration=duration(r)+maximum(t), fs=framerate(r))
   end
   signal(isanalytic(r) ? x : √2 * real(x), fs)
 end
+
+"""
+    zak(x, L, K)
+    zak(x, L)
+
+Compute the Zak transform of signal `x` with `L` delay bins and `K` Doppler bins.
+The length of signal `x` must be equal to `LK`. If `K` is not specified, it is
+assumed to be the length of `x` divided by `L`. Returns a `K × L` complex matrix.
+
+If the frame rate of `x` if `fs` Sa/s, the delay bins are spaced at `1/fs` seconds
+and the Doppler bins are spaced at `fs/LK` Hz. The Zak transform is scaled such
+that the energy of the signal is preserved, i.e., `sum(abs2, x) ≈ sum(abs2, X)`.
+
+For efficient computation of the Zak transform, `K` should product of small primes.
+
+# Examples:
+```julia-repl
+julia> x = randn(ComplexF64, 4096)
+4096-element Vector{ComplexF64}:
+  :
+julia> X = zak(x, 64)
+64×64 Matrix{ComplexF64}:
+  :
+```
+"""
+function zak(x::AbstractVector, L::Int, K::Int)
+  length(x) == L * K || throw(ArgumentError("Length of x must be L * K"))
+  X = complex.(collect(transpose(reshape(x, L, K)))) ./ sqrt(K)
+  fft!(X, 1)
+end
+
+function zak(x::AbstractVector, L::Int)
+  length(x) % L == 0 || throw(ArgumentError("Length of x must be a multiple of L"))
+  zak(x, L, length(x) ÷ L)
+end
+
+"""
+$(SIGNATURES)
+Compute the inverse Zak transform of 2D `K × L` complex signal `X` with `L`
+delay bins and `K` Doppler bins. The length of the returned signal is `LK`.
+
+See [`zak`](@ref) for more details.
+
+# Examples:
+```julia-repl
+julia> x = randn(ComplexF64, 4096)
+4096-element Vector{ComplexF64}:
+  :
+julia> X = zak(x, 64)
+64×64 Matrix{ComplexF64}:
+  :
+julia> izak(X) ≈ x
+true
+```
+"""
+function izak(X::AbstractMatrix)
+  X = ifft(complex.(X), 1)
+  X .*= sqrt(size(X, 1))
+  collect(vec(transpose(X)))
+end

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2"
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"

test/runtests.jl

@@ -1,4 +1,4 @@
-using Test, Statistics, LinearAlgebra, DSP, DSP.Util
+using Test, Statistics, LinearAlgebra, DSP, DSP.Util, FFTW
 using Plots
 using SignalAnalysis
 using SignalAnalysis.Units

test/tests-core.jl

@@ -694,6 +694,21 @@ function test_dsp()
   delay!(x, -10)
   @test samples(x) ≈ y atol=1e-3
 
+  x = rand(rng, ComplexF64, 4096)
+  @test vec(zak(x, length(x))) == x
+  @test zak(x, 1) ≈ fft(x) / sqrt(length(x))
+  Zx = zak(x, 64)
+  @test Zx == zak(x, 64, 64)
+  @test izak(Zx) ≈ x
+  @test sum(abs2, x) ≈ sum(abs2, Zx)
+  @test vec(sum(Zx; dims=1)) / sqrt(64) ≈ x[1:64]
+  @test zak(4.2x, 64) ≈ 4.2Zx
+  @test abs.(zak(circshift(x, 10), 64)) ≈ abs.(circshift(Zx, (0, 10)))
+  y = rand(rng, ComplexF64, 4096)
+  Zy = zak(y, 64)
+  @test vec(Zx)' * vec(Zy) ≈ x' * y
+  @test zak(4.2x + 2.7y, 64) ≈ 4.2Zx + 2.7Zy
+
 end
 
 function test_rand()