HPS.java

import org.jtransforms.fft.FloatFFT_1D;

import javafx.application.Application;
import javafx.scene.Scene;
import javafx.scene.chart.Chart;
import javafx.scene.chart.LineChart;
import javafx.scene.chart.NumberAxis;
import javafx.scene.chart.XYChart;
import javafx.scene.control.ScrollPane;
import javafx.scene.layout.VBox;
import javafx.stage.Stage;

public class Main extends Application {
	public static void main(String[] args) {
		launch(args);
	}
	
	private Chart createSourceDataChart(float[] data, String title) {
		final NumberAxis xAxis = new NumberAxis();
        final NumberAxis yAxis = new NumberAxis();
        
        xAxis.setLabel("Sample");
        yAxis.setLabel("Level");
        
		LineChart<Number, Number> c = new LineChart<Number, Number>(xAxis, yAxis);
		c.setCreateSymbols(false); //hide dots
		XYChart.Series<Number, Number> series = new XYChart.Series<>();
        series.setName(title);
        
        for (int i=0; i < data.length; i+=4) {
        	series.getData().add(new XYChart.Data<Number, Number>(i, data[i]));
        }

        c.getData().add(series);

        return c;
	}
	
	private Chart createFFTChart(float[] data, String title, int num) {
		final NumberAxis xAxis = new NumberAxis();
        final NumberAxis yAxis = new NumberAxis();
        
        xAxis.setLabel("Bin");
        yAxis.setLabel("Power");
        
		LineChart<Number, Number> c = new LineChart<Number, Number>(xAxis, yAxis);
		c.setCreateSymbols(false); //hide dots
		XYChart.Series<Number, Number> series = new XYChart.Series<Number, Number>();
        series.setName(title);
        
        int limit = num != -1 ? num : data.length;
        for (int i=0; i < limit; i+=2) {
        	series.getData().add(new XYChart.Data<Number, Number>(i, data[i]));
        }

        c.getData().add(series);

        return c;
	}
	
	/**
     * Gets the next value of the hamming function.
     *
     * @param i the index
     * @param size the total size
     * @return the hamming value
     */
    private float getHammingValue(int i, int size) {
        return (float) (0.54 - 0.46 * Math.cos((2 * Math.PI * i) / (size - 1)));
    }

	private void applyWindow(float[] from, float[] to) {
		int M = from.length;
		for (int n = 0; n < M; n++) {
			to[n] = from[n] * getHammingValue(n, M);
		}		
	}
		
	private float[] powerSpectrum(float[] window) {
		float[] powerSpectrum = new float[window.length];
		float[] fftBuffer = new float[window.length * 2 + 1];
		System.arraycopy(window, 0, fftBuffer, 0, window.length);
		FloatFFT_1D fft = new FloatFFT_1D(window.length);
		fft.realForward(fftBuffer);
		
		for (int i=0; i < fftBuffer.length / 2 - 1; i++) {
			float real = fftBuffer[2*i];
			float imag = fftBuffer[2*i+1];
			
			powerSpectrum[i] = (float)Math.sqrt(real*real + imag * imag); 
		}
		
		return powerSpectrum;
	}
	
	@Override
	public void start(Stage s) throws Exception {

		ScrollPane sp = new ScrollPane();
		sp.setFitToHeight(false);
		sp.setFitToWidth(true);
		VBox bp = new VBox();
		sp.setContent(bp);
		
		// 1. Generate Input
		float[] sine = SineWaveGenerator.createSinWaveBuffer(3600, 160, 1000);
		bp.getChildren().add(createSourceDataChart(sine, "Sine Wave"));
		
		// 2. Apply Window Function
		float[] windowed = new float[sine.length];		
		applyWindow(sine, windowed);
		bp.getChildren().add(createSourceDataChart(windowed, "Windowed Signal"));
		
		// 3. Calculate Power Spectrum (using FFT)
		float[] spectrum = powerSpectrum(windowed);
		bp.getChildren().add(createFFTChart(spectrum, "FFT Power Spectrum", -1));
		
		// 4. Compress
		float[] spectrumCopy = new float[spectrum.length];
		System.arraycopy(spectrum, 0, spectrumCopy, 0, spectrum.length);		
		for (int compression = 2; compression < 4; compression++) {
			for (int i = 1; i < spectrum.length; i++) {
				spectrum[i] = spectrum[i] * getCompressedSample(spectrumCopy, 1, compression, i);
			}
		}

		// 5. HPS
		bp.getChildren().add(createFFTChart(spectrum, "HPS", -1));
		
		// 6. Find Peak
		int maxBin = 0;
		float maxVal = Float.NEGATIVE_INFINITY;
		
		for (int i=1; i < spectrum.length; i++) {
			float val = spectrum[i];
			if (val > maxVal) {
				maxVal = val;
				maxBin = i;
			}
		}
		System.out.println("MAXBIN: " + maxBin);
		
		// 7. Interpolate
		float mid = spectrum[maxBin];
		float left  = spectrum[maxBin- 1];
		float right = spectrum[maxBin + 1];
		float shift = 0.5f*(right-left) / ( 2.0f*mid - left - right );
		float pEst = maxBin + shift;
		System.out.println("MAXBIN AFTER INTERPOLATION: " + (int)pEst);
		
		// 8. Convert to frequency
		float freq = (float) getFrequencyForIndex((int)pEst, spectrum.length, 8000);
		System.out.println("FOUND FREQUENCY: " + freq);
		
		s.setMinWidth(1024);
		s.setScene(new Scene(sp));
		s.show();		
	}
	
	private float getCompressedSample(float[] buffer, int offset, int compression, int loc) {

		if (offset + loc * compression < buffer.length) {
			return buffer[offset + loc * compression];
		}
		
		return 0;
	}
	
	private float getFrequencyForIndex(int index, int size, int rate) {
		float freq = (float)index * (float)rate / (float)size;
		return freq;
	}
}