Slicing for Zygote

New to Julia
,
1

Hello,

Do you remember this issue Understanding Flux and Zygote … I got stuck with slicing again. And since I do not know how to use Zygote properly I couldn’t make progress.

The original function istft . And my objective is to convert this function to handle 3d-arrays. To do so first I wanted to see how the first for loop can be converted; since it has mutation and slicing.

using PencilFFTs
using AbstractFFTs: rfft, irfft


function hanning(n::Integer)
    [0.5*(1-cos(2*pi*k/(n-1))) for k=0:n-1]
end


function istft(spectrogram::Matrix{T},
                           framelen::Int=1024,
                           hopsize::Int=div(framelen,2),
                           window=hanning(framelen)) where T <: Complex

    numframes = size(spectrogram, 2)

    expectedlen = framelen + (numframes-1)*hopsize
    reconstructed = zeros(expectedlen)
    windowsum = zeros(expectedlen)
    windowsquare = window .* window

    # Overlapping addition
    for i=1:numframes
        s, e = (i-1)*hopsize+1, (i-1)*hopsize+framelen
        r = irfft(spectrogram[:,i], framelen)
        reconstructed[s:e] += r .* window # how can it be converted w/o slicing? 
        windowsum[s:e] += windowsquare # how can it be converted w/o slicing? 
    end

    # Normalized by window
    for i=1:reconstructed[end]
        # avoid zero division
        if windowsum[i] > 1.0e-7
            reconstructed[i] /= windowsum[i]
        end
    end

    return reconstructed
end

The good part is that I can convert most part of the code with directly vector operations. But I could not handle the slicing part where it overlaps and sums.


function my_istft(spectrogram::Matrix{T},
                            framelen::Int=1024,
                            hopsize::Int=div(framelen, 2),
                            window=hanning(framelen)) where T <: Complex

    numframes = size(spectrogram, 2);
    

    expectedlen = framelen + (numframes-1)*hopsize
    reconstructed = zeros(expectedlen)
    windowsum = zeros(expectedlen)

    se = [(i-1)*hopsize+1:(i-1)*hopsize+framelen for i in 1:numframes]; # starting and ending points of overlaps

    # overlapping
    reconst = vec(irfft(spectrogram, framelen, 1) .* window)
    winsum  = vec(ones(framelen, numframes) .* (window .^2))
    idx = winsum .> 1.0e-7 

    reconstructed[s:e] # how can I convert this ? 

    reconst[idx] = reconst[idx] ./ winsum[idx]
    return reconstructed

end

My overall objective is to convert this method for n-dimensional arrays. But since the 2-d arrays are more simple . I chose to translate first with this.

Could you help me to convert the function ?

B.R.

p.s. The first istft is not mine I cloned it from another repo and did some minor changes for Julia 1.5 support.

2

Without checking anything too carefully, can you do something like this?

# test data:
spectrogram = rand(ComplexF32, 5,16)
framelen = 8
hopsize = 4
window=hanning(framelen) # real 8-vec
numframes = size(spectrogram, 2) # 16
expectedlen = framelen + (numframes-1)*hopsize # 68

# even/odd & combine:
wr_odd = window .* irfft(spectrogram[:,1:2:end], framelen, 1)
wr_even = window .* irfft(spectrogram[:,2:2:end], framelen, 1)
recon5 = vcat(vec(wr_odd), zeros(hopsize)) .+ vcat(zeros(hopsize), vec(wr_even))

reconstructed ≈ recon5 # after stage "Overlapping addition"
2 Likes
3

@mcabbott thank you. Now, it works on 4 dimensional arrays. In addition, the original function is now faster :slight_smile: