DFT and IDFT nGraph Python API (#5387)

* Added nGraph Python API for operations DFT and IDFT.

* Written tests for the DFT Python API.

* Written tests for IDFT nGraph Python API.

* Small fixes.

* Started to add tests for the signal_size case.

* Written tests for signal_size case of DFT.

* Written tests for signal_size case of IDFT.

* Some code style fixes in IDFT nGraph Python API tests.

* Code style fixes in tests for DFT nGraph Python API.

* Now DFT nGraph Python API tests are used numpy FFT ressults as expected results of tests.

* Now IDFT nGraph Python API tests without signal_size are used numpy FFT result as input data.

* Now IDFT nGraph Python API tests use numpy IFFT as expected results for signal_size cases.

* Deleted redundant function.

* Formatting fix.

* Now test data for DFT and IDFT nGraph Python API are randomly generated.

* Added seed initialization.

ngraph/python/src/ngraph/__init__.py

@@ -52,6 +52,7 @@ from ngraph.opset7 import deformable_convolution
 from ngraph.opset7 import deformable_psroi_pooling
 from ngraph.opset7 import depth_to_space
 from ngraph.opset7 import detection_output
+from ngraph.opset7 import dft
 from ngraph.opset7 import divide
 from ngraph.opset7 import einsum
 from ngraph.opset7 import elu
@@ -80,6 +81,7 @@ from ngraph.opset7 import gru_sequence
 from ngraph.opset7 import hard_sigmoid
 from ngraph.opset7 import hsigmoid
 from ngraph.opset7 import hswish
+from ngraph.opset7 import idft
 from ngraph.opset7 import interpolate
 from ngraph.opset7 import less
 from ngraph.opset7 import less_equal

ngraph/python/src/ngraph/opset7/__init__.py

@@ -37,6 +37,7 @@ from ngraph.opset1.ops import deformable_convolution
 from ngraph.opset1.ops import deformable_psroi_pooling
 from ngraph.opset1.ops import depth_to_space
 from ngraph.opset1.ops import detection_output
+from ngraph.opset7.ops import dft
 from ngraph.opset1.ops import divide
 from ngraph.opset7.ops import einsum
 from ngraph.opset1.ops import elu
@@ -65,6 +66,7 @@ from ngraph.opset5.ops import gru_sequence
 from ngraph.opset1.ops import hard_sigmoid
 from ngraph.opset5.ops import hsigmoid
 from ngraph.opset4.ops import hswish
+from ngraph.opset7.ops import idft
 from ngraph.opset1.ops import interpolate
 from ngraph.opset1.ops import less
 from ngraph.opset1.ops import less_equal

ngraph/python/src/ngraph/opset7/ops.py

@@ -101,3 +101,45 @@ def roll(
     inputs = as_nodes(data, shift, axes)
 
     return _get_node_factory_opset7().create("Roll", inputs)
+
+
+@nameable_op
+def dft(
+        data: NodeInput,
+        axes: NodeInput,
+        signal_size: Optional[NodeInput] = None,
+) -> Node:
+    """Return a node which performs DFT operation.
+
+    @param data: Tensor with transformed data.
+    @param axes: Tensor with axes to transform.
+    @param signal_size: Tensor specifying signal size with respect to axes from the input 'axes'.
+    @return: The new node which performs DFT operation on the input data tensor.
+    """
+    if signal_size is None:
+        inputs = as_nodes(data, axes)
+    else:
+        inputs = as_nodes(data, axes, signal_size)
+
+    return _get_node_factory_opset7().create("DFT", inputs)
+
+
+@nameable_op
+def idft(
+        data: NodeInput,
+        axes: NodeInput,
+        signal_size: Optional[NodeInput] = None,
+) -> Node:
+    """Return a node which performs IDFT operation.
+
+    @param data: Tensor with transformed data.
+    @param axes: Tensor with axes to transform.
+    @param signal_size: Tensor specifying signal size with respect to axes from the input 'axes'.
+    @return: The new node which performs IDFT operation on the input data tensor.
+    """
+    if signal_size is None:
+        inputs = as_nodes(data, axes)
+    else:
+        inputs = as_nodes(data, axes, signal_size)
+
+    return _get_node_factory_opset7().create("IDFT", inputs)

ngraph/python/tests/__init__.py

@@ -159,4 +159,6 @@ xfail_issue_52463 = xfail_test(reason="test_operator_add_size1_singleton_broadca
 
 xfail_issue_49391 = xfail_test(reason="Roll is not implemented in CPU plugin.")
 
+xfail_issue_49359 = xfail_test(reason="DFT is not implemented in CPU plugin")
+xfail_issue_49375 = xfail_test(reason="IDFT is not implemented in CPU plugin")
 xfail_issue_45432 = xfail_test(reason="Einsum is not implemented in CPU plugin.")

ngraph/python/tests/test_ngraph/test_dft.py

@@ -0,0 +1,119 @@
+import ngraph as ng
+import numpy as np
+from tests import xfail_issue_49359
+from tests.runtime import get_runtime
+
+
+def build_fft_input_data():
+    np.random.seed(202104)
+    return np.random.uniform(0, 1, (2, 10, 10, 2)).astype(np.float32)
+
+
+@xfail_issue_49359
+def test_dft_1d():
+    runtime = get_runtime()
+    input_data = build_fft_input_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([2], dtype=np.int64))
+
+    dft_node = ng.dft(input_tensor, input_axes)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fft(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), axis=2)
+    expected_results = np_results.view(dtype=np.float32).reshape((2, 10, 10, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.00001)
+
+
+@xfail_issue_49359
+def test_dft_2d():
+    runtime = get_runtime()
+    input_data = build_fft_input_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([1, 2], dtype=np.int64))
+
+    dft_node = ng.dft(input_tensor, input_axes)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fft2(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), axes=[1, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((2, 10, 10, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.000062)
+
+
+@xfail_issue_49359
+def test_dft_3d():
+    runtime = get_runtime()
+    input_data = build_fft_input_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([0, 1, 2], dtype=np.int64))
+
+    dft_node = ng.dft(input_tensor, input_axes)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fftn(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), axes=[0, 1, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((2, 10, 10, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.0002)
+
+
+@xfail_issue_49359
+def test_dft_1d_signal_size():
+    runtime = get_runtime()
+    input_data = build_fft_input_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([-2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([20], dtype=np.int64))
+
+    dft_node = ng.dft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fft(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), n=20, axis=-2)
+    expected_results = np_results.view(dtype=np.float32).reshape((2, 20, 10, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.00001)
+
+
+@xfail_issue_49359
+def test_dft_2d_signal_size_1():
+    runtime = get_runtime()
+    input_data = build_fft_input_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([0, 2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([4, 5], dtype=np.int64))
+
+    dft_node = ng.dft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fft2(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), s=[4, 5], axes=[0, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((4, 10, 5, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.000062)
+
+
+@xfail_issue_49359
+def test_dft_2d_signal_size_2():
+    runtime = get_runtime()
+    input_data = build_fft_input_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([1, 2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([4, 5], dtype=np.int64))
+
+    dft_node = ng.dft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fft2(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), s=[4, 5], axes=[1, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((2, 4, 5, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.000062)
+
+
+@xfail_issue_49359
+def test_dft_3d_signal_size():
+    runtime = get_runtime()
+    input_data = build_fft_input_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([0, 1, 2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([4, 5, 16], dtype=np.int64))
+
+    dft_node = ng.dft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fftn(np.squeeze(input_data.view(dtype=np.complex64), axis=-1),
+                             s=[4, 5, 16], axes=[0, 1, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((4, 5, 16, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.0002)

ngraph/python/tests/test_ngraph/test_idft.py

@@ -0,0 +1,121 @@
+import ngraph as ng
+import numpy as np
+from tests import xfail_issue_49375
+from tests.runtime import get_runtime
+
+
+def get_data():
+    np.random.seed(202104)
+    return np.random.uniform(0, 1, (2, 10, 10, 2)).astype(np.float32)
+
+
+@xfail_issue_49375
+def test_idft_1d():
+    runtime = get_runtime()
+    expected_results = get_data()
+    complex_input_data = np.fft.fft(np.squeeze(expected_results.view(dtype=np.complex64), axis=-1), axis=2)
+    input_data = complex_input_data.view(dtype=np.float32).reshape((2, 10, 10, 2))
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([2], dtype=np.int64))
+
+    dft_node = ng.idft(input_tensor, input_axes)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    assert np.allclose(dft_results, expected_results, atol=0.000002)
+
+
+@xfail_issue_49375
+def test_idft_2d():
+    runtime = get_runtime()
+    expected_results = get_data()
+    complex_input_data = np.fft.fft2(np.squeeze(expected_results.view(dtype=np.complex64), axis=-1),
+                                     axes=[1, 2])
+    input_data = complex_input_data.view(dtype=np.float32).reshape((2, 10, 10, 2))
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([1, 2], dtype=np.int64))
+
+    dft_node = ng.idft(input_tensor, input_axes)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    assert np.allclose(dft_results, expected_results, atol=0.000002)
+
+
+@xfail_issue_49375
+def test_idft_3d():
+    runtime = get_runtime()
+    expected_results = get_data()
+    complex_input_data = np.fft.fft2(np.squeeze(expected_results.view(dtype=np.complex64), axis=-1),
+                                     axes=[0, 1, 2])
+    input_data = complex_input_data.view(dtype=np.float32).reshape((2, 10, 10, 2))
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([0, 1, 2], dtype=np.int64))
+
+    dft_node = ng.idft(input_tensor, input_axes)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    assert np.allclose(dft_results, expected_results, atol=0.000003)
+
+
+@xfail_issue_49375
+def test_idft_1d_signal_size():
+    runtime = get_runtime()
+    input_data = get_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([-2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([20], dtype=np.int64))
+
+    dft_node = ng.idft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.ifft(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), n=20, axis=-2)
+    expected_results = np_results.view(dtype=np.float32).reshape((2, 20, 10, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.000002)
+
+
+@xfail_issue_49375
+def test_idft_2d_signal_size_1():
+    runtime = get_runtime()
+    input_data = get_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([0, 2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([4, 5], dtype=np.int64))
+
+    dft_node = ng.idft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.ifft2(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), s=[4, 5], axes=[0, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((4, 10, 5, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.000002)
+
+
+@xfail_issue_49375
+def test_idft_2d_signal_size_2():
+    runtime = get_runtime()
+    input_data = get_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([1, 2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([4, 5], dtype=np.int64))
+
+    dft_node = ng.idft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.fft2(np.squeeze(input_data.view(dtype=np.complex64), axis=-1), s=[4, 5], axes=[1, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((2, 4, 5, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.000002)
+
+
+@xfail_issue_49375
+def test_idft_3d_signal_size():
+    runtime = get_runtime()
+    input_data = get_data()
+    input_tensor = ng.constant(input_data)
+    input_axes = ng.constant(np.array([0, 1, 2], dtype=np.int64))
+    input_signal_size = ng.constant(np.array([4, 5, 16], dtype=np.int64))
+
+    dft_node = ng.idft(input_tensor, input_axes, input_signal_size)
+    computation = runtime.computation(dft_node)
+    dft_results = computation()
+    np_results = np.fft.ifftn(np.squeeze(input_data.view(dtype=np.complex64), axis=-1),
+                              s=[4, 5, 16], axes=[0, 1, 2])
+    expected_results = np_results.view(dtype=np.float32).reshape((4, 5, 16, 2))
+    assert np.allclose(dft_results, expected_results, atol=0.000002)