compactify

pocketfft.cc

@@ -317,6 +317,8 @@ template<typename T> struct cmplx {
   T r, i;
   cmplx() {}
   cmplx(T r_, T i_) : r(r_), i(i_) {}
+  void Set(T r_, T i_) { r=r_; i=i_; }
+  void Set(T r_) { r=r_; i=T(0); }
   cmplx &operator+= (const cmplx &other)
     { r+=other.r; i+=other.i; return *this; }
   template<typename T2>cmplx &operator*= (T2 other)
@@ -2122,10 +2124,7 @@ template<typename T> void pocketfft_general_c(const shape_t &shape,
               : plan->backward((vtype *)tdatav, fct);
       for (size_t i=0; i<it_out.length(); ++i)
         for (size_t j=0; j<vlen; ++j)
-          {
-          data_out[p_o[j]+i*it_out.stride()].r = tdatav[i].r[j];
-          data_out[p_o[j]+i*it_out.stride()].i = tdatav[i].i[j];
-          }
+          data_out[p_o[j]+i*it_out.stride()].Set(tdatav[i].r[j],tdatav[i].i[j]);
       }
 #endif
     while (it_in.remaining()>0)
@@ -2251,28 +2250,19 @@ template<typename T> void pocketfft_general_r2c(const shape_t &shape,
         tdatav[i][j] = data_in[p_i[j]+i*it_in.stride()];
     plan.forward((vtype *)tdatav, fct);
     for (size_t j=0; j<vlen; ++j)
-      {
-      data_out[p_o[j]].r = tdatav[0][j];
-      data_out[p_o[j]].i = 0.;
-      }
+      data_out[p_o[j]].Set(tdatav[0][j]);
     size_t i;
     for (i=1; i<len-1; i+=2)
       {
       size_t io = (i+1)/2;
       for (size_t j=0; j<vlen; ++j)
-        {
-        data_out[p_o[j]+io*it_out.stride()].r = tdatav[i][j];
-        data_out[p_o[j]+io*it_out.stride()].i = tdatav[i+1][j];
-        }
+        data_out[p_o[j]+io*it_out.stride()].Set(tdatav[i][j], tdatav[i+1][j]);
       }
     if (i<len)
       {
       size_t io = (i+1)/2;
       for (size_t j=0; j<vlen; ++j)
-        {
-        data_out[p_o[j]+io*it_out.stride()].r = tdatav[i][j];
-        data_out[p_o[j]+io*it_out.stride()].i = 0.;
-        }
+        data_out[p_o[j]+io*it_out.stride()].Set(tdatav[i][j]);
       }
     }
 #endif
@@ -2283,21 +2273,12 @@ template<typename T> void pocketfft_general_r2c(const shape_t &shape,
     for (size_t i=0; i<len; ++i)
       tdata[i] = d_i[i*s_i];
     plan.forward(tdata, fct);
-    d_o[0].r = tdata[0];
-    d_o[0].i = 0.;
+    d_o[0].Set(tdata[0]);
     size_t i;
     for (i=1; i<len-1; i+=2)
-      {
-      size_t io = (i+1)/2;
-      d_o[io*s_o].r = tdata[i];
-      d_o[io*s_o].i = tdata[i+1];
-      }
+      d_o[((i+1)/2)*s_o].Set(tdata[i], tdata[i+1]);
     if (i<len)
-      {
-      size_t io = (i+1)/2;
-      d_o[io*s_o].r = tdata[i];
-      d_o[io*s_o].i = 0.;
-      }
+      d_o[((i+1)/2)*s_o].Set(tdata[i]);
     it_in.advance();
     it_out.advance();
     }
@@ -2364,10 +2345,7 @@ template<typename T> void pocketfft_general_c2r(const shape_t &shape_out,
       tdata[i+1] = d_i[ii*s_i].i;
       }
     if (i<len)
-      {
-      size_t ii = (i+1)/2;
-      tdata[i] = d_i[ii*s_i].r;
-      }
+      tdata[i] = d_i[((i+1)/2)*s_i].r;
     plan.backward(tdata, fct);
     for (size_t i=0; i<len; ++i)
       d_o[i*s_o] = tdata[i];