onnx--onnx/onnx/reference/ops/op_dft.py

# Copyright (c) ONNX Project Contributors

# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import numpy as np

from onnx.reference.op_run import OpRun


def _fft(x: np.ndarray, fft_length: int, axis: int) -> np.ndarray:
    """Compute the FFT return the real representation of the complex result."""
    transformed = np.fft.fft(x, n=fft_length, axis=axis)
    real_frequencies = np.real(transformed)
    imaginary_frequencies = np.imag(transformed)
    return np.concatenate(
        (real_frequencies[..., np.newaxis], imaginary_frequencies[..., np.newaxis]),
        axis=-1,
    )


def _cfft(
    x: np.ndarray,
    fft_length: int,
    axis: int,
    onesided: bool,
    normalize: bool,
) -> np.ndarray:
    if x.shape[-1] == 1:
        # The input contains only the real part
        signal = x
    else:
        # The input is a real representation of a complex signal
        slices = [slice(0, x) for x in x.shape]
        slices[-1] = slice(0, x.shape[-1], 2)
        real = x[tuple(slices)]
        slices[-1] = slice(1, x.shape[-1], 2)
        imag = x[tuple(slices)]
        signal = real + 1j * imag

    complex_signals = np.squeeze(signal, -1)
    result = _fft(complex_signals, fft_length, axis=axis)
    # Post process the result based on arguments
    if onesided:
        slices = [slice(0, a) for a in result.shape]
        slices[axis] = slice(0, result.shape[axis] // 2 + 1)
        result = result[tuple(slices)]
    if normalize:
        result /= fft_length
    return result


def _ifft(x: np.ndarray, fft_length: int, axis: int) -> np.ndarray:
    """Standard IFFT: complex input -> complex output."""
    signals = np.fft.ifft(x, n=fft_length, axis=axis)
    real_signals = np.real(signals)
    imaginary_signals = np.imag(signals)
    return np.concatenate(
        (real_signals[..., np.newaxis], imaginary_signals[..., np.newaxis]),
        axis=-1,
    )


def _irfft(x: np.ndarray, fft_length: int, axis: int) -> np.ndarray:
    """IRFFT: one-sided complex input -> full real output."""
    signals = np.fft.irfft(x, n=fft_length, axis=axis)
    # Return real-valued output with last dimension = 1
    return signals[..., np.newaxis]


def _cifft(
    x: np.ndarray, fft_length: int, axis: int, onesided: bool = False
) -> np.ndarray:
    """Complex IFFT wrapper that handles both standard IFFT and IRFFT."""
    # Extract complex values from input
    if x.shape[-1] == 1:
        # Real input (shouldn't happen for IFFT, but handle it)
        frequencies = np.squeeze(x, -1)
    else:
        # Complex input: interleaved real/imaginary
        slices = [slice(0, dim) for dim in x.shape]
        slices[-1] = slice(0, x.shape[-1], 2)
        real = x[tuple(slices)]
        slices[-1] = slice(1, x.shape[-1], 2)
        imag = x[tuple(slices)]
        frequencies = np.squeeze(real, -1) + 1j * np.squeeze(imag, -1)

    if onesided:
        # IRFFT: one-sided complex input -> full real output
        return _irfft(frequencies, fft_length, axis=axis)
    # Standard IFFT: full complex input -> full complex output
    return _ifft(frequencies, fft_length, axis=axis)


class DFT_17(OpRun):
    def _run(
        self,
        x: np.ndarray,
        dft_length: int | None = None,
        axis: int = 1,
        inverse: bool = False,
        onesided: bool = False,
    ) -> tuple[np.ndarray]:
        # Convert to positive axis
        axis = axis % len(x.shape)

        # Set default dft_length based on operation type
        if dft_length is None:
            if inverse and onesided:
                # IRFFT: input is one-sided, default to even length
                dft_length = 2 * (x.shape[axis] - 1)
            else:
                dft_length = x.shape[axis]

        if inverse:
            result = _cifft(x, dft_length, axis=axis, onesided=onesided)
        else:
            result = _cfft(x, dft_length, axis=axis, onesided=onesided, normalize=False)
        return (result.astype(x.dtype),)


class DFT_20(OpRun):
    def _run(
        self,
        x: np.ndarray,
        dft_length: int | None = None,
        axis: int = -2,
        inverse: bool = False,
        onesided: bool = False,
    ) -> tuple[np.ndarray]:
        # Convert to positive axis
        axis = axis % len(x.shape)

        # Set default dft_length based on operation type
        if dft_length is None:
            if inverse and onesided:
                # IRFFT: input is one-sided, default to even length
                dft_length = 2 * (x.shape[axis] - 1)
            else:
                dft_length = x.shape[axis]

        if inverse:
            result = _cifft(x, dft_length, axis=axis, onesided=onesided)
        else:
            result = _cfft(x, dft_length, axis=axis, onesided=onesided, normalize=False)
        return (result.astype(x.dtype),)