From 0f5781c927b271f00f6ce6d76404220832ae5b8a Mon Sep 17 00:00:00 2001
From: Martin Reinecke <martin@mpa-garching.mpg.de>
Date: Wed, 17 Jul 2019 19:57:38 +0200
Subject: [PATCH] fix normalization; add tests

---
 pocketfft_hdronly.h |  4 ++--
 pypocketfft.cc      |  4 ++++
 test.py             | 23 ++++++++++++++++++++++-
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/pocketfft_hdronly.h b/pocketfft_hdronly.h
index 1923c15..f7619a0 100644
--- a/pocketfft_hdronly.h
+++ b/pocketfft_hdronly.h
@@ -2344,8 +2344,8 @@ template<typename T0> class T_cosq
       if (N==2)
         {
         T TSQX = sqrt2*c[1];
-        c[1] = c[0]-TSQX;
-        c[0] = c[0]+TSQX;
+        c[1] = fct*(c[0]-TSQX);
+        c[0] = fct*(c[0]+TSQX);
         return;
         }
       size_t NS2 = (N+1)/2;
diff --git a/pypocketfft.cc b/pypocketfft.cc
index 608001e..503afdd 100644
--- a/pypocketfft.cc
+++ b/pypocketfft.cc
@@ -240,6 +240,8 @@ template<typename T> py::array r2r_dct23_internal(const py::array &in,
   {
   py::gil_scoped_release release;
   T fct = norm_fct<T>(inorm, dims, axes);
+  if (inorm==2) fct*=T(1/ldbl_t(2));
+  if (inorm==1) fct*=T(1/sqrt(ldbl_t(2)));
   pocketfft::r2r_dct23(dims, s_in, s_out, axes, forward, d_in, d_out, fct,
     nthreads);
   }
@@ -267,6 +269,8 @@ template<typename T> py::array r2r_dst23_internal(const py::array &in,
   {
   py::gil_scoped_release release;
   T fct = norm_fct<T>(inorm, dims, axes);
+  if (inorm==2) fct*=T(1/ldbl_t(2));
+  if (inorm==1) fct*=T(1/sqrt(ldbl_t(2)));
   pocketfft::r2r_dst23(dims, s_in, s_out, axes, forward, d_in, d_out, fct,
     nthreads);
   }
diff --git a/test.py b/test.py
index 5c43f82..17610ed 100644
--- a/test.py
+++ b/test.py
@@ -17,6 +17,13 @@ def _l2error(a, b):
     return np.sqrt(np.sum(np.abs(a-b)**2)/np.sum(np.abs(a)**2))
 
 
+def _assert_close(a, b, epsilon):
+    err = _l2error(a, b)
+    if (err >= epsilon):
+        print("Error: {} > {}".format(err, epsilon))
+    assert_(err<epsilon)
+
+
 def fftn(a, axes=None, inorm=0, out=None, nthreads=1):
     return pypocketfft.c2c(a, axes=axes, forward=True, inorm=inorm,
                            out=out, nthreads=nthreads)
@@ -55,7 +62,7 @@ def test1D(len, inorm):
     a = np.random.rand(len)-0.5 + 1j*np.random.rand(len)-0.5j
     b = a.astype(np.complex64)
     c = a.astype(np.complex256)
-    assert_(_l2error(a, ifftn(fftn(c, inorm=inorm), inorm=2-inorm)) < 1e-18)
+    _assert_close(a, ifftn(fftn(c, inorm=inorm), inorm=2-inorm), 1e-18)
     assert_(_l2error(a, ifftn(fftn(a, inorm=inorm), inorm=2-inorm)) < 1.5e-15)
     assert_(_l2error(a.real, ifftn(fftn(a.real, inorm=inorm), inorm=2-inorm))
             < 1.5e-15)
@@ -196,3 +203,17 @@ def test_genuine_hartley_2D(shp, axes):
     a = np.random.rand(*shp)-0.5
     assert_(_l2error(pypocketfft.genuine_hartley(pypocketfft.genuine_hartley(
         a, axes=axes), axes=axes, inorm=2), a) < 1e-15)
+
+
+@pmp("len", len1D)
+@pmp("inorm", [0, 1, 2])
+def testdcst1D(len, inorm):
+    a = np.random.rand(len)-0.5
+    b = a.astype(np.float32)
+    c = a.astype(np.float128)
+    _assert_close(a, pypocketfft.r2r_dct23(pypocketfft.r2r_dct23(c, inorm=inorm, forward=True), inorm=2-inorm, forward=False), 2e-18)
+    _assert_close(a, pypocketfft.r2r_dct23(pypocketfft.r2r_dct23(a, inorm=inorm, forward=True), inorm=2-inorm, forward=False), 1.5e-15)
+    _assert_close(b, pypocketfft.r2r_dct23(pypocketfft.r2r_dct23(b, inorm=inorm, forward=True), inorm=2-inorm, forward=False), 6e-7)
+    _assert_close(a, pypocketfft.r2r_dst23(pypocketfft.r2r_dst23(c, inorm=inorm, forward=True), inorm=2-inorm, forward=False), 2e-18)
+    _assert_close(a, pypocketfft.r2r_dst23(pypocketfft.r2r_dst23(a, inorm=inorm, forward=True), inorm=2-inorm, forward=False), 1.5e-15)
+    _assert_close(b, pypocketfft.r2r_dst23(pypocketfft.r2r_dst23(b, inorm=inorm, forward=True), inorm=2-inorm, forward=False), 6e-7)
-- 
GitLab