annotate yeti/cqtkernel.yeti @ 12:765bf12dcafb

Return the atom count & spacing as well

module cqtkernel;

vec = load may.vector;
bf = load may.vector.blockfuncs;
complex = load may.complex;
window = load may.signal.window;
fft = load may.transform.fft;
pl = load may.plot;
cm = load may.matrix.complex;

{ pow, round, floor, ceil, nextPowerOfTwo } = load may.mathmisc;

makeKernel { sampleRate, maxFreq, binsPerOctave } =
   (q = 1;
    atomHopFactor = 0.25;
    thresh = 0.0005;
    minFreq = (maxFreq/2) * (pow 2 (1/binsPerOctave));
    bigQ = q / ((pow 2 (1/binsPerOctave)) - 1);

    maxNK = round(bigQ * sampleRate / minFreq);
    minNK = round(bigQ * sampleRate /
                  (minFreq * (pow 2 ((binsPerOctave-1) / binsPerOctave))));

    atomHop = round(minNK * atomHopFactor);
    
    firstCentre = atomHop * (ceil ((ceil (maxNK/2)) / atomHop));
    
    fftSize = nextPowerOfTwo (firstCentre + ceil (maxNK/2));
    
    println "sampleRate = \(sampleRate), maxFreq = \(maxFreq), binsPerOctave = \(binsPerOctave), q = \(q), atomHopFactor = \(atomHopFactor), thresh = \(thresh)";
    println "minFreq = \(minFreq), bigQ = \(bigQ), maxNK = \(maxNK), minNK = \(minNK), atomHop = \(atomHop), firstCentre = \(firstCentre), fftSize = \(fftSize)";
    
    winNr = floor((fftSize - ceil(maxNK/2) - firstCentre) / atomHop) + 1;
    
    lastCentre = firstCentre + (winNr - 1) * atomHop;
    
    fftHop = (lastCentre + atomHop) - firstCentre;
    
    println "winNr = \(winNr), lastCentre = \(lastCentre), fftHop = \(fftHop)";
    
    fftFunc = fft.forward fftSize;
    
    // Note the MATLAB uses exp(2*pi*1i*x) for a complex generating
    // function. We can't do that here; we need to generate real and imag
    // parts separately as real = cos(2*pi*x), imag = sin(2*pi*x).
    
    kernels = map do k:
    
        nk = round(bigQ * sampleRate / (minFreq * (pow 2 ((k-1)/binsPerOctave))));
    
        // the cq MATLAB toolbox uses a symmetric window for
        // blackmanharris -- which is odd because it uses a periodic one
        // for other types. Oh well
        win = bf.divideBy nk
           (bf.sqrt
               (window.windowFunction (BlackmanHarris ()) [Symmetric true] nk));
        
        fk = minFreq * (pow 2 ((k-1)/binsPerOctave));
    
        genKernel f = bf.multiply win
           (vec.fromList
               (map do i: f (2 * pi * fk * i / sampleRate) done [0..nk-1]));
    
        reals = genKernel cos;
        imags = genKernel sin;
    
        atomOffset = firstCentre - ceil(nk/2);
    
        map do i:
    
            shift = vec.zeros (atomOffset + ((i-1) * atomHop));
    
            specKernel = fftFunc
               (complex.complexArray
                   (vec.concat [shift, reals])
                   (vec.concat [shift, imags]));
    
            map do c:
                if complex.magnitude c <= thresh then complex.zero else c fi
                done specKernel;
    
        done [1..winNr];
    
    done [1..binsPerOctave];
    
    kmat = cm.toSparse
       (cm.scaled (1/fftSize)
           (cm.newComplexMatrix (RowMajor()) (concat kernels)));
    
    println "density = \(cm.density kmat)";
    
    // Normalisation
    
    wx1 = bf.maxindex (complex.magnitudes (cm.getRow 0 kmat));
    wx2 = bf.maxindex (complex.magnitudes (cm.getRow (cm.height kmat - 1) kmat));
    
    subset = cm.columnSlice kmat wx1 (wx2+1);
    square = cm.product (cm.conjugateTransposed subset) subset;
    diag = complex.magnitudes (cm.getDiagonal 0 square);
    wK = vec.slice diag (round(1/q)) (vec.length diag - round(1/q) - 2);
    
    weight = (fftHop / fftSize) / (bf.mean (bf.abs wK));
    weight = sqrt(weight);

    {
        kernel = cm.scaled weight kmat,
        fftSize,
        fftHop,
        binsPerOctave,
        atomsPerFrame = winNr,
        atomSpacing = atomHop,
        maxFreq,
        minFreq,
        bigQ
    });

{
    makeKernel
}