task: tests to cover #336

Author: jharlow-intel

mkl_fft/tests/test_dispatch_equivalence.py

@@ -180,6 +180,45 @@ def test_rfft_axis_3d(func, dtype, axis, order):
     _check(got, want, dtype)
 
 
+# norm="ortho" is unitary: it preserves the L2 norm. See gh-336.
+
+_UNITARY_CASES = [
+    ((16, 18), (0,)),  # gh-336
+    ((16, 18), (1,)),
+    ((16, 18), (0, 1)),
+    (_SHAPE_3D, (0,)),
+    (_SHAPE_3D, (1,)),
+    (_SHAPE_3D, (2,)),
+    (_SHAPE_3D, (0, 2)),
+    (_SHAPE_3D, None),
+]
+
+_UNITARY_RTOL = {"complex64": 1e-4, "complex128": 1e-10}
+
+
+@pytest.mark.parametrize("func", ["fftn", "ifftn"])
+@pytest.mark.parametrize("dtype", ["complex64", "complex128"])
+@pytest.mark.parametrize("shape,axes", _UNITARY_CASES)
+def test_ortho_preserves_norm(func, dtype, shape, axes):
+    x = _make(shape, dtype)
+    x /= np.sqrt((abs(x) ** 2).sum())
+    got = getattr(mkl_fft, func)(x, axes=axes, norm="ortho")
+    assert np.isclose(
+        (abs(got) ** 2).sum(), 1.0, rtol=_UNITARY_RTOL[dtype]
+    ), f"norm not preserved: {(abs(got) ** 2).sum()}"
+
+
+@pytest.mark.parametrize("shape,axes", _UNITARY_CASES)
+def test_ortho_roundtrip_is_identity(shape, axes):
+    """Insensitive to gh-336: the fwd and bwd scales cancel. Guards only that
+    the two compose to the identity.
+    """
+    x = _make(shape, "complex128")
+    fwd = mkl_fft.fftn(x, axes=axes, norm="ortho")
+    got = mkl_fft.ifftn(fwd, axes=axes, norm="ortho")
+    _check(got, x, "complex128")
+
+
 # ---------------------------------------------------------------------------
 # norm interacts with the scale factor applied at dispatch time
 # ---------------------------------------------------------------------------