fft.js

/**
 * WebGL 2D FFT実装
 * 波動光学計算用の高速フーリエ変換
 */

class WebGLFFT {
    constructor(webglUtils) {
        this.gl = webglUtils.gl;
        this.utils = webglUtils;
        this.size = 0;
        this.logSize = 0;
        this.programs = new Map();
        this.butterflyTexture = null;
        this.tempTextures = [];
        this.initialized = false;
    }

    /**
     * FFTシステムを初期化
     * @param {number} size FFTサイズ (2の冪乗) 
     */
    initialize(size) {
        if (!this.isPowerOfTwo(size)) {
            throw new Error('FFTサイズは2の冪乗である必要があります');
        }

        this.size = size;
        this.logSize = Math.log2(size);
        
        // バタフライテクスチャは不要 (シェーダーで直接計算) 
        this.createPrograms();
        this.createTempTextures();
        
        this.initialized = true;
        console.log(`WebGL FFT initialized for size ${size}x${size}`);
    }

    /**
     * 数値が2の冪乗かチェック
     * @param {number} n 
     * @returns {boolean}
     */
    isPowerOfTwo(n) {
        return n > 0 && (n & (n - 1)) === 0;
    }

    /**
     * バタフライテクスチャを作成
     * FFTの各段階で使用される係数を事前計算
     */
    createButterflyTexture() {
        const gl = this.gl;
        const data = new Float32Array(this.size * this.logSize * 4);
        
        let offset = 0;
        for (let stage = 0; stage < this.logSize; stage++) {
            const butterflyWidth = 1 << (stage + 1);
            const butterflyHeight = this.size / butterflyWidth;
            
            for (let y = 0; y < butterflyHeight; y++) {
                for (let x = 0; x < this.size; x++) {
                    const butterflyX = x % butterflyWidth;
                    const isPingPong = Math.floor(x / butterflyWidth) % 2;
                    
                    if (butterflyX < butterflyWidth / 2) {
                        // バタフライの左半分
                        const twiddle = -2.0 * Math.PI * butterflyX / butterflyWidth;
                        const twiddleReal = Math.cos(twiddle);
                        const twiddleImag = Math.sin(twiddle);
                        
                        const sourceX = isPingPong ? 
                            butterflyX + butterflyWidth / 2 : 
                            butterflyX;
                        
                        data[offset + 0] = twiddleReal;
                        data[offset + 1] = twiddleImag;
                        data[offset + 2] = sourceX;
                        data[offset + 3] = y;
                    } else {
                        // バタフライの右半分
                        const sourceX = isPingPong ? 
                            butterflyX - butterflyWidth / 2 : 
                            butterflyX;
                        
                        data[offset + 0] = 1.0;
                        data[offset + 1] = 0.0;
                        data[offset + 2] = sourceX;
                        data[offset + 3] = y;
                    }
                    
                    offset += 4;
                }
            }
        }

        this.butterflyTexture = this.utils.createTexture(
            this.size, this.logSize,
            gl.RGBA32F, gl.RGBA, gl.FLOAT,
            data, 'butterfly'
        );
    }

    /**
     * FFT用のシェーダープログラムを作成
     */
    createPrograms() {
        const gl = this.gl;

        // ビットリバーサルパス用
        const bitReversalFragment = `#version 300 es
            precision highp float;
            
            in vec2 v_texCoord;
            uniform sampler2D u_input;
            uniform int u_size;
            uniform bool u_horizontal;
            out vec4 fragColor;
            
            int bitReverse(int x, int logSize) {
                int result = 0;
                for (int i = 0; i < 16; i++) {
                    if (i >= logSize) break;
                    if ((x & (1 << i)) != 0) {
                        result |= (1 << (logSize - 1 - i));
                    }
                }
                return result;
            }
            
            void main() {
                ivec2 coord = ivec2(v_texCoord * float(u_size));
                int logSize = int(log2(float(u_size)));
                
                if (u_horizontal) {
                    int reversedX = bitReverse(coord.x, logSize);
                    fragColor = texelFetch(u_input, ivec2(reversedX, coord.y), 0);
                } else {
                    int reversedY = bitReverse(coord.y, logSize);
                    fragColor = texelFetch(u_input, ivec2(coord.x, reversedY), 0);
                }
            }
        `;

        // バタフライパス用
        const butterflyFragment = `#version 300 es
            precision highp float;
            
            in vec2 v_texCoord;
            uniform sampler2D u_input;
            uniform int u_stage;
            uniform int u_size;
            uniform bool u_inverse;
            out vec4 fragColor;
            
            ${CommonShaders.complexUtils}
            
            void main() {
                ivec2 coord = ivec2(v_texCoord * float(u_size));
                int logSize = int(log2(float(u_size)));
                
                // FFT段階のサイズを計算
                int stageSize = 1 << (u_stage + 1);
                int halfStageSize = stageSize >> 1;
                
                // バタフライの位置を計算
                int butterflyIndex = coord.x % stageSize;
                bool isUpperHalf = butterflyIndex < halfStageSize;
                
                // ペアの位置を計算
                int pairOffset = isUpperHalf ? halfStageSize : -halfStageSize;
                ivec2 pairCoord = ivec2(coord.x + pairOffset, coord.y);
                
                // 現在の値とペアの値を取得
                vec2 current = texture(u_input, v_texCoord).rg;
                vec2 pair = texelFetch(u_input, pairCoord, 0).rg;
                
                // ツイッドル因子を計算
                float angle = -2.0 * 3.14159265 * float(butterflyIndex % halfStageSize) / float(stageSize);
                if (u_inverse) {
                    angle = -angle;
                }
                vec2 twiddle = vec2(cos(angle), sin(angle));
                
                // バタフライ演算
                vec2 result;
                if (isUpperHalf) {
                    vec2 product = cmul(pair, twiddle);
                    result = current + product;
                } else {
                    vec2 product = cmul(current, twiddle);
                    result = pair - product;
                }
                
                fragColor = vec4(result.x, result.y, 0.0, 1.0);
            }
        `;

        // 正規化パス用
        const normalizeFragment = `#version 300 es
            precision highp float;
            
            in vec2 v_texCoord;
            uniform sampler2D u_input;
            uniform float u_scale;
            out vec4 fragColor;
            
            void main() {
                vec4 data = texture(u_input, v_texCoord);
                fragColor = data * u_scale;
            }
        `;

        // プログラム作成
        this.programs.set('bitReversal', 
            this.utils.createProgram(CommonShaders.vertexShader, bitReversalFragment));
        this.programs.set('butterfly', 
            this.utils.createProgram(CommonShaders.vertexShader, butterflyFragment));
        this.programs.set('normalize', 
            this.utils.createProgram(CommonShaders.vertexShader, normalizeFragment));
    }

    /**
     * 一時テクスチャを作成
     */
    createTempTextures() {
        const gl = this.gl;
        
        // ping-pong用に2つの複素テクスチャペアを作成
        for (let i = 0; i < 2; i++) {
            this.tempTextures.push({
                real: this.utils.createTexture(
                    this.size, this.size,
                    gl.R32F, gl.RED, gl.FLOAT,
                    null, `temp${i}_real`
                ),
                imag: this.utils.createTexture(
                    this.size, this.size,
                    gl.R32F, gl.RED, gl.FLOAT,
                    null, `temp${i}_imag`
                )
            });
        }
    }

    /**
     * 複素テクスチャペアをパックしてRGBAテクスチャに変換
     * @param {WebGLTexture} realTexture 
     * @param {WebGLTexture} imagTexture 
     * @returns {WebGLTexture}
     */
    packComplexTexture(realTexture, imagTexture) {
        const gl = this.gl;
        
        const packFragment = `#version 300 es
            precision highp float;
            
            in vec2 v_texCoord;
            uniform sampler2D u_real;
            uniform sampler2D u_imag;
            out vec4 fragColor;
            
            void main() {
                float real = texture(u_real, v_texCoord).r;
                float imag = texture(u_imag, v_texCoord).r;
                
                // RGBAにパック (実部をRに、虚部をGに) 
                fragColor = vec4(real, imag, 0.0, 1.0);
            }
        `;

        const program = this.utils.createProgram(CommonShaders.vertexShader, packFragment);
        const outputTexture = this.utils.createTexture(
            this.size, this.size,
            gl.RGBA32F, gl.RGBA, gl.FLOAT
        );
        const framebuffer = this.utils.createFramebuffer(outputTexture);
        const quad = this.utils.createFullscreenQuad();

        this.utils.renderPass(program, framebuffer, {
            u_real: realTexture,
            u_imag: imagTexture
        }, quad);

        return outputTexture;
    }

    /**
     * RGBAテクスチャを複素テクスチャペアにアンパック
     * @param {WebGLTexture} packedTexture 
     * @returns {{real: WebGLTexture, imag: WebGLTexture}}
     */
    unpackComplexTexture(packedTexture) {
        const gl = this.gl;
        
        const unpackFragment = `#version 300 es
            precision highp float;
            
            in vec2 v_texCoord;
            uniform sampler2D u_packed;
            uniform int u_component; // 0: real, 1: imag
            out vec4 fragColor;
            
            void main() {
                vec4 packed = texture(u_packed, v_texCoord);
                float value = (u_component == 0) ? packed.r : packed.g;
                fragColor = vec4(value, 0.0, 0.0, 1.0);
            }
        `;

        const program = this.utils.createProgram(CommonShaders.vertexShader, unpackFragment);
        const quad = this.utils.createFullscreenQuad();

        const realTexture = this.utils.createTexture(
            this.size, this.size, gl.R32F, gl.RED, gl.FLOAT);
        const imagTexture = this.utils.createTexture(
            this.size, this.size, gl.R32F, gl.RED, gl.FLOAT);

        const realFramebuffer = this.utils.createFramebuffer(realTexture);
        const imagFramebuffer = this.utils.createFramebuffer(imagTexture);

        // 実部を抽出
        this.utils.renderPass(program, realFramebuffer, {
            u_packed: packedTexture,
            u_component: 0
        }, quad);

        // 虚部を抽出
        this.utils.renderPass(program, imagFramebuffer, {
            u_packed: packedTexture,
            u_component: 1
        }, quad);

        return { real: realTexture, imag: imagTexture };
    }

    /**
     * 1D FFTを実行 (横方向または縦方向) 
     * @param {WebGLTexture} inputTexture 
     * @param {boolean} horizontal 
     * @param {boolean} inverse 
     * @returns {WebGLTexture}
     */
    fft1D(inputTexture, horizontal = true, inverse = false) {
        if (!this.initialized) {
            throw new Error('FFTが初期化されていません');
        }

        const gl = this.gl;
        const quad = this.utils.createFullscreenQuad();
        
        let currentTexture = inputTexture;
        let nextTexture = this.utils.createTexture(
            this.size, this.size, gl.RGBA32F, gl.RGBA, gl.FLOAT);
        
        // ビットリバーサル
        const bitReversalFB = this.utils.createFramebuffer(nextTexture);
        this.utils.renderPass(this.programs.get('bitReversal'), bitReversalFB, {
            u_input: currentTexture,
            u_size: this.size,
            u_horizontal: horizontal
        }, quad);

        [currentTexture, nextTexture] = [nextTexture, currentTexture];

        // バタフライパス
        for (let stage = 0; stage < this.logSize; stage++) {
            const framebuffer = this.utils.createFramebuffer(nextTexture);
            
            this.utils.renderPass(this.programs.get('butterfly'), framebuffer, {
                u_input: currentTexture,
                u_stage: stage,
                u_size: this.size,
                u_inverse: inverse
            }, quad);

            [currentTexture, nextTexture] = [nextTexture, currentTexture];
        }

        // 逆変換の場合は正規化
        if (inverse) {
            const normalizedTexture = this.utils.createTexture(
                this.size, this.size, gl.RGBA32F, gl.RGBA, gl.FLOAT);
            const normalizeFB = this.utils.createFramebuffer(normalizedTexture);
            
            this.utils.renderPass(this.programs.get('normalize'), normalizeFB, {
                u_input: currentTexture,
                u_scale: 1.0 / this.size
            }, quad);

            return normalizedTexture;
        }

        return currentTexture;
    }

    /**
     * 2D FFTを実行
     * @param {{real: WebGLTexture, imag: WebGLTexture}} inputPair 
     * @param {boolean} inverse 
     * @returns {{real: WebGLTexture, imag: WebGLTexture}}
     */
    fft2D(inputPair, inverse = false) {
        // 複素テクスチャペアをパック
        const packedInput = this.packComplexTexture(inputPair.real, inputPair.imag);
        
        // 横方向FFT
        const horizontalFFT = this.fft1D(packedInput, true, inverse);
        
        // 縦方向FFT
        const verticalFFT = this.fft1D(horizontalFFT, false, inverse);
        
        // 結果をアンパック
        return this.unpackComplexTexture(verticalFFT);
    }

    /**
     * 周波数シフト (FFTshift相当) 
     * @param {{real: WebGLTexture, imag: WebGLTexture}} inputPair 
     * @returns {{real: WebGLTexture, imag: WebGLTexture}}
     */
    fftShift(inputPair) {
        const gl = this.gl;
        
        const shiftFragment = `#version 300 es
            precision highp float;
            
            in vec2 v_texCoord;
            uniform sampler2D u_input;
            uniform int u_size;
            out vec4 fragColor;
            
            void main() {
                ivec2 coord = ivec2(v_texCoord * float(u_size));
                int halfSize = u_size / 2;
                
                ivec2 shiftedCoord = coord;
                shiftedCoord.x = (coord.x + halfSize) % u_size;
                shiftedCoord.y = (coord.y + halfSize) % u_size;
                
                fragColor = texelFetch(u_input, shiftedCoord, 0);
            }
        `;

        const program = this.utils.createProgram(CommonShaders.vertexShader, shiftFragment);
        const quad = this.utils.createFullscreenQuad();

        const realOutput = this.utils.createTexture(
            this.size, this.size, gl.R32F, gl.RED, gl.FLOAT);
        const imagOutput = this.utils.createTexture(
            this.size, this.size, gl.R32F, gl.RED, gl.FLOAT);

        const realFB = this.utils.createFramebuffer(realOutput);
        const imagFB = this.utils.createFramebuffer(imagOutput);

        this.utils.renderPass(program, realFB, {
            u_input: inputPair.real,
            u_size: this.size
        }, quad);

        this.utils.renderPass(program, imagFB, {
            u_input: inputPair.imag,
            u_size: this.size
        }, quad);

        return { real: realOutput, imag: imagOutput };
    }

    /**
     * リソースをクリーンアップ
     */
    cleanup() {
        const gl = this.gl;
        
        if (this.butterflyTexture) {
            gl.deleteTexture(this.butterflyTexture);
        }

        for (const texturePair of this.tempTextures) {
            gl.deleteTexture(texturePair.real);
            gl.deleteTexture(texturePair.imag);
        }

        for (const program of this.programs.values()) {
            gl.deleteProgram(program);
        }

        this.tempTextures = [];
        this.programs.clear();
        this.initialized = false;
    }
}