From 21e8dcdb8a4b975378dd453c6e610e32664c14cc Mon Sep 17 00:00:00 2001
From: Ian Jauslin <ian.jauslin@rutgers.edu>
Date: Sat, 1 Feb 2025 14:12:54 -0500
Subject: [PATCH] In RK algorithms: do not use kx,ky, use i

---
 src/navier-stokes.c | 269 +++++++++++++++++---------------------------
 1 file changed, 101 insertions(+), 168 deletions(-)

diff --git a/src/navier-stokes.c b/src/navier-stokes.c
index 0d32e89..253f1bf 100644
--- a/src/navier-stokes.c
+++ b/src/navier-stokes.c
@@ -23,6 +23,8 @@ limitations under the License.
 #include <stdlib.h>
 #include <string.h>
 
+#define USIZE (K1*(2*K2+1)+K2)
+
 // compute solution as a function of time
 int uk(
   int K1,
@@ -353,30 +355,30 @@ int ns_init_tmps(
   unsigned int algorithm
 ){
   // velocity field
-  *u=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
+  *u=calloc(USIZE,sizeof(_Complex double));
 
   // allocate tmp vectors for computation
   if(algorithm==ALGORITHM_RK2){
-    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
+    *tmp1=calloc(USIZE,sizeof(_Complex double));
+    *tmp2=calloc(USIZE,sizeof(_Complex double));
   } else if (algorithm==ALGORITHM_RK4){
-    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp3=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
+    *tmp1=calloc(USIZE,sizeof(_Complex double));
+    *tmp2=calloc(USIZE,sizeof(_Complex double));
+    *tmp3=calloc(USIZE,sizeof(_Complex double));
   } else if (algorithm==ALGORITHM_RKF45 || algorithm==ALGORITHM_RKDP54){
-    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp3=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp4=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp5=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp6=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp7=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
+    *tmp1=calloc(USIZE,sizeof(_Complex double));
+    *tmp2=calloc(USIZE,sizeof(_Complex double));
+    *tmp3=calloc(USIZE,sizeof(_Complex double));
+    *tmp4=calloc(USIZE,sizeof(_Complex double));
+    *tmp5=calloc(USIZE,sizeof(_Complex double));
+    *tmp6=calloc(USIZE,sizeof(_Complex double));
+    *tmp7=calloc(USIZE,sizeof(_Complex double));
   } else if (algorithm==ALGORITHM_RKBS32){
-    *tmp1=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp2=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp3=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp4=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
-    *tmp5=calloc(K1*(2*K2+1)+K2,sizeof(_Complex double));
+    *tmp1=calloc(USIZE,sizeof(_Complex double));
+    *tmp2=calloc(USIZE,sizeof(_Complex double));
+    *tmp3=calloc(USIZE,sizeof(_Complex double));
+    *tmp4=calloc(USIZE,sizeof(_Complex double));
+    *tmp5=calloc(USIZE,sizeof(_Complex double));
   } else {
     fprintf(stderr,"bug: unknown algorithm: %u, contact ian.jauslin@rutgers,edu\n",algorithm);
   };
@@ -462,7 +464,7 @@ int copy_u(
 ){
   int i;
 
-  for(i=0;i<K1*(2*K2+1)+K2;i++){
+  for(i=0;i<USIZE;i++){
     u[i]=u0[i];
   }
 
@@ -524,7 +526,6 @@ int ns_step(
   return (0);
 }
 
-// TODO: do not need to use klookup in any of the rk routines
 // RK 4 algorithm
 int ns_step_rk4(
   _Complex double* u,
@@ -546,69 +547,53 @@ int ns_step_rk4(
   bool keep_en_cst,
   double target_en
 ){
-  int kx,ky;
+  int i;
 
   // k1
   ns_rhs(tmp1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
   // add to output
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp3[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+delta/6*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp3[i]=u[i]+delta/6*tmp1[i];
   }
 
   // u+h*k1/2
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp2[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+delta/2*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp2[i]=u[i]+delta/2*tmp1[i];
   }
   // k2
   ns_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
   // add to output
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp3[klookup_sym(kx,ky,K2)]+=delta/3*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp3[i]+=delta/3*tmp1[i];
   }
 
   // u+h*k2/2
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp2[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+delta/2*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp2[i]=u[i]+delta/2*tmp1[i];
   }
   // k3
   ns_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
   // add to output
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp3[klookup_sym(kx,ky,K2)]+=delta/3*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp3[i]+=delta/3*tmp1[i];
   }
 
   // u+h*k3
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp2[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+delta*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp2[i]=u[i]+delta*tmp1[i];
   }
   // k4
   ns_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
   // add to output
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      u[klookup_sym(kx,ky,K2)]=tmp3[klookup_sym(kx,ky,K2)]+delta/6*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    u[i]=tmp3[i]+delta/6*tmp1[i];
   }
 
   // keep enstrophy constant
   if(keep_en_cst){
     double en=compute_enstrophy(u, K1, K2, L);
-    for(kx=0;kx<=K1;kx++){
-      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
-      }
+    for(i=0;i<USIZE;i++){
+      u[i]*=sqrt(target_en/en);
     }
   }
 
@@ -635,33 +620,27 @@ int ns_step_rk2(
   bool keep_en_cst,
   double target_en
 ){
-  int kx,ky;
+  int i;
 
   // k1
   ns_rhs(tmp1, u, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // u+h*k1/2
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp2[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+delta/2*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp2[i]=u[i]+delta/2*tmp1[i];
   }
   // k2
   ns_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
   // add to output
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      u[klookup_sym(kx,ky,K2)]+=delta*tmp1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    u[i]+=delta*tmp1[i];
   }
 
   // keep enstrophy constant
   if(keep_en_cst){
     double en=compute_enstrophy(u, K1, K2, L);
-    for(kx=0;kx<=K1;kx++){
-      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
-      }
+    for(i=0;i<USIZE;i++){
+      u[i]*=sqrt(target_en/en);
     }
   }
 
@@ -701,7 +680,7 @@ int ns_step_rkf45(
   // whether to compute k1 or leave it as is
   bool compute_k1
 ){
-  int kx,ky;
+  int i;
   double cost;
 
   // k1: u(t)
@@ -710,53 +689,41 @@ int ns_step_rkf45(
   }
 
   // k2 : u(t+1/4*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)/4*k1[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)/4*k1[i];
   }
   ns_rhs(k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k3 : u(t+3/8*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(3./32*k1[klookup_sym(kx,ky,K2)]+9./32*k2[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(3./32*k1[i]+9./32*k2[i]);
   }
   ns_rhs(k3, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k4 : u(t+12./13*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(1932./2197*k1[klookup_sym(kx,ky,K2)]-7200./2197*k2[klookup_sym(kx,ky,K2)]+7296./2197*k3[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(1932./2197*k1[i]-7200./2197*k2[i]+7296./2197*k3[i]);
   }
   ns_rhs(k4, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k5 : u(t+1*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(439./216*k1[klookup_sym(kx,ky,K2)]-8*k2[klookup_sym(kx,ky,K2)]+3680./513*k3[klookup_sym(kx,ky,K2)]-845./4104*k4[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(439./216*k1[i]-8*k2[i]+3680./513*k3[i]-845./4104*k4[i]);
   }
   ns_rhs(k5, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k6 : u(t+1./2*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(-8./27*k1[klookup_sym(kx,ky,K2)]+2*k2[klookup_sym(kx,ky,K2)]-3544./2565*k3[klookup_sym(kx,ky,K2)]+1859./4104*k4[klookup_sym(kx,ky,K2)]-11./40*k5[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(-8./27*k1[i]+2*k2[i]-3544./2565*k3[i]+1859./4104*k4[i]-11./40*k5[i]);
   }
   ns_rhs(k6, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
   
   // next step
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      // u
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(25./216*k1[klookup_sym(kx,ky,K2)]+1408./2565*k3[klookup_sym(kx,ky,K2)]+2197./4104*k4[klookup_sym(kx,ky,K2)]-1./5*k5[klookup_sym(kx,ky,K2)]);
-      // U: save to k6, which is no longer needed
-      k6[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(16./135*k1[klookup_sym(kx,ky,K2)]+6656./12825*k3[klookup_sym(kx,ky,K2)]+28561./56430*k4[klookup_sym(kx,ky,K2)]-9./50*k5[klookup_sym(kx,ky,K2)]+2./55*k6[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    // u
+    tmp[i]=u[i]+(*delta)*(25./216*k1[i]+1408./2565*k3[i]+2197./4104*k4[i]-1./5*k5[i]);
+    // U: save to k6, which is no longer needed
+    k6[i]=u[i]+(*delta)*(16./135*k1[i]+6656./12825*k3[i]+28561./56430*k4[i]-9./50*k5[i]+2./55*k6[i]);
   }
 
   // cost function
@@ -765,10 +732,8 @@ int ns_step_rkf45(
   // compare relative error with tolerance
   if(cost<tolerance){
     // copy to output
-    for(kx=0;kx<=K1;kx++){
-      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	u[klookup_sym(kx,ky,K2)]=tmp[klookup_sym(kx,ky,K2)];
-      }
+    for(i=0;i<USIZE;i++){
+      u[i]=tmp[i];
     }
     // next delta to use in future steps (do not exceed max_delta)
     *next_delta=fmin(max_delta, (*delta)*pow(tolerance/cost,0.2));
@@ -776,10 +741,8 @@ int ns_step_rkf45(
     // keep enstrophy constant
     if(keep_en_cst){
       double en=compute_enstrophy(u, K1, K2, L);
-      for(kx=0;kx<=K1;kx++){
-	for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	  u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
-	}
+      for(i=0;i<USIZE;i++){
+	u[i]*=sqrt(target_en/en);
       }
     }
   }
@@ -826,7 +789,7 @@ int ns_step_rkbs32(
   // whether to compute k1
   bool compute_k1
 ){
-  int kx,ky;
+  int i;
   double cost;
 
   // k1: u(t)
@@ -836,36 +799,28 @@ int ns_step_rkbs32(
   }
 
   // k2 : u(t+1/4*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)/2*(*k1)[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)/2*(*k1)[i];
   }
   ns_rhs(k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k3 : u(t+3/4*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(3./4*k2[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(3./4*k2[i]);
   }
   ns_rhs(k3, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k4 : u(t+delta)
   // tmp computed here is the next step
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(2./9*(*k1)[klookup_sym(kx,ky,K2)]+1./3*k2[klookup_sym(kx,ky,K2)]+4./9*k3[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(2./9*(*k1)[i]+1./3*k2[i]+4./9*k3[i]);
   }
   ns_rhs(*k4, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // next step
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      // U: store in k3, which is no longer needed
-      k3[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(7./24*(*k1)[klookup_sym(kx,ky,K2)]+1./4*k2[klookup_sym(kx,ky,K2)]+1./3*k3[klookup_sym(kx,ky,K2)]+1./8*(*k4)[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    // U: store in k3, which is no longer needed
+    k3[i]=u[i]+(*delta)*(7./24*(*k1)[i]+1./4*k2[i]+1./3*k3[i]+1./8*(*k4)[i]);
   }
   
   // compute cost
@@ -874,10 +829,8 @@ int ns_step_rkbs32(
   // compare cost with tolerance
   if(cost<tolerance){
     // add to output
-    for(kx=0;kx<=K1;kx++){
-      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	u[klookup_sym(kx,ky,K2)]=tmp[klookup_sym(kx,ky,K2)];
-      }
+    for(i=0;i<USIZE;i++){
+      u[i]=tmp[i];
     }
     // next delta to use in future steps (do not exceed max_delta)
     *next_delta=fmin(max_delta, (*delta)*pow(tolerance/cost,1./3));
@@ -890,10 +843,8 @@ int ns_step_rkbs32(
     // keep enstrophy constant
     if(keep_en_cst){
       double en=compute_enstrophy(u, K1, K2, L);
-      for(kx=0;kx<=K1;kx++){
-	for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	  u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
-	}
+      for(i=0;i<USIZE;i++){
+	u[i]*=sqrt(target_en/en);
       }
     }
   }
@@ -941,7 +892,7 @@ int ns_step_rkdp54(
   // whether to compute k1
   bool compute_k1
 ){
-  int kx,ky;
+  int i;
   double cost;
 
   // k1: u(t)
@@ -951,61 +902,47 @@ int ns_step_rkdp54(
   }
 
   // k2 : u(t+1/5*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)/5*(*k1)[klookup_sym(kx,ky,K2)];
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)/5*(*k1)[i];
   }
   ns_rhs(*k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k3 : u(t+3/10*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(3./40*(*k1)[klookup_sym(kx,ky,K2)]+9./40*(*k2)[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(3./40*(*k1)[i]+9./40*(*k2)[i]);
   }
   ns_rhs(k3, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k4 : u(t+4/5*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(44./45*(*k1)[klookup_sym(kx,ky,K2)]-56./15*(*k2)[klookup_sym(kx,ky,K2)]+32./9*k3[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(44./45*(*k1)[i]-56./15*(*k2)[i]+32./9*k3[i]);
   }
   ns_rhs(k4, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k5 : u(t+8/9*delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(19372./6561*(*k1)[klookup_sym(kx,ky,K2)]-25360./2187*(*k2)[klookup_sym(kx,ky,K2)]+64448./6561*k3[klookup_sym(kx,ky,K2)]-212./729*k4[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(19372./6561*(*k1)[i]-25360./2187*(*k2)[i]+64448./6561*k3[i]-212./729*k4[i]);
   }
   ns_rhs(k5, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k6 : u(t+delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(9017./3168*(*k1)[klookup_sym(kx,ky,K2)]-355./33*(*k2)[klookup_sym(kx,ky,K2)]+46732./5247*k3[klookup_sym(kx,ky,K2)]+49./176*k4[klookup_sym(kx,ky,K2)]-5103./18656*k5[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    tmp[i]=u[i]+(*delta)*(9017./3168*(*k1)[i]-355./33*(*k2)[i]+46732./5247*k3[i]+49./176*k4[i]-5103./18656*k5[i]);
   }
   ns_rhs(k6, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   // k7 : u(t+delta)
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      // tmp computed here is the step
-      tmp[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(35./384*(*k1)[klookup_sym(kx,ky,K2)]+500./1113*k3[klookup_sym(kx,ky,K2)]+125./192*k4[klookup_sym(kx,ky,K2)]-2187./6784*k5[klookup_sym(kx,ky,K2)]+11./84*k6[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    // tmp computed here is the step
+    tmp[i]=u[i]+(*delta)*(35./384*(*k1)[i]+500./1113*k3[i]+125./192*k4[i]-2187./6784*k5[i]+11./84*k6[i]);
   }
   // store in k2, which is not needed anymore
   ns_rhs(*k2, tmp, K1, K2, N1, N2, nu, L, g, fft1, fft2, ifft, irreversible);
 
   //next step
-  for(kx=0;kx<=K1;kx++){
-    for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-      // U: store in k6, which is not needed anymore
-      k6[klookup_sym(kx,ky,K2)]=u[klookup_sym(kx,ky,K2)]+(*delta)*(5179./57600*(*k1)[klookup_sym(kx,ky,K2)]+7571./16695*k3[klookup_sym(kx,ky,K2)]+393./640*k4[klookup_sym(kx,ky,K2)]-92097./339200*k5[klookup_sym(kx,ky,K2)]+187./2100*k6[klookup_sym(kx,ky,K2)]+1./40*(*k2)[klookup_sym(kx,ky,K2)]);
-    }
+  for(i=0;i<USIZE;i++){
+    // U: store in k6, which is not needed anymore
+    k6[i]=u[i]+(*delta)*(5179./57600*(*k1)[i]+7571./16695*k3[i]+393./640*k4[i]-92097./339200*k5[i]+187./2100*k6[i]+1./40*(*k2)[i]);
   }
   
   // compute cost
@@ -1014,10 +951,8 @@ int ns_step_rkdp54(
   // compare relative error with tolerance
   if(cost<tolerance){
     // add to output
-    for(kx=0;kx<=K1;kx++){
-      for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	u[klookup_sym(kx,ky,K2)]=tmp[klookup_sym(kx,ky,K2)];
-      }
+    for(i=0;i<USIZE;i++){
+      u[i]=tmp[i];
     }
     // next delta to use in future steps (do not exceed max_delta)
     *next_delta=fmin(max_delta, (*delta)*pow(tolerance/cost,0.2));
@@ -1030,10 +965,8 @@ int ns_step_rkdp54(
     // keep enstrophy constant
     if(keep_en_cst){
       double en=compute_enstrophy(u, K1, K2, L);
-      for(kx=0;kx<=K1;kx++){
-	for(ky=(kx>0 ? -K2 : 1);ky<=K2;ky++){
-	  u[klookup_sym(kx,ky,K2)]*=sqrt(target_en/en);
-	}
+      for(i=0;i<USIZE;i++){
+	u[i]*=sqrt(target_en/en);
       }
     }
   }
@@ -1138,7 +1071,7 @@ int ns_rhs(
     alpha=compute_alpha(u,K1,K2,g,L);
   }
 
-  for(i=0; i<K1*(2*K2+1)+K2; i++){
+  for(i=0; i<USIZE; i++){
     out[i]=0;
   }
   for(kx=0;kx<=K1;kx++){