Rewrite: change cli arguments handling

5 files changed, 644 insertions, 316 deletions

diff --git a/src/driving.c b/src/driving.c
new file mode 100644
index 0000000..56bea63
--- /dev/null
+++ b/src/driving.c
@@ -0,0 +1,17 @@
+#include "driving.h"
+#include <math.h>
+
+_Complex double g_test(
+  int kx,
+  int ky
+){
+  //return sqrt(kx*kx*ky*ky)*exp(-(kx*kx+ky*ky));
+  if(kx==2 && ky==-1){
+    return 0.5+sqrt(3)/2*I;
+  }
+  else if(kx==-2 && ky==1){
+     return 0.5-sqrt(3)/2*I;
+  }
+  return 0.;
+}
+
diff --git a/src/driving.h b/src/driving.h
new file mode 100644
index 0000000..cf2f3de
--- /dev/null
+++ b/src/driving.h
@@ -0,0 +1,8 @@
+#ifndef DRIVING_H
+#define DRIVING_H
+
+#include <complex.h>
+
+_Complex double g_test( int kx, int ky);
+
+#endif
diff --git a/src/main.c b/src/main.c
index bbdf143..3383f10 100644
--- a/src/main.c
+++ b/src/main.c
@@ -1,43 +1,77 @@
-#define VERSION "0.0"
+#define VERSION "0.1"
 
 #include <math.h>
 #include <complex.h>
 #include <fftw3.h>
 #include <string.h>
 #include <stdlib.h>
+#include <stdbool.h>
 #include "navier-stokes.h"
+#include "driving.h"
 
 // usage message
 int print_usage();
+
 // read command line arguments
-int read_args(int argc, const char* argv[], ns_params* params, unsigned int* nsteps, unsigned int* computation_nr);
+int read_args(int argc, const char* argv[], char** params, unsigned int* driving_force, unsigned int* command);
+int read_params(char* params, int* K1, int* K2, int* N1, int* N2, unsigned int* nsteps, double* nu, double* delta, unsigned int* print_freq);
+int set_parameter(char* lhs, char* rhs, int* K1, int* K2, int* N1, int* N2, unsigned int* nsteps, double* nu, double* delta, unsigned int* print_freq, bool* setN1, bool* setN2);
+
+
+#define COMMAND_UK 1
+#define COMMAND_ENSTROPHY 2
 
-// compute enstrophy as a function of time in the I-NS equation
-int enstrophy(ns_params params, unsigned int Nsteps);
+#define DRIVING_TEST 1
 
 
-#define COMPUTATION_ENSTROPHY 1
-int main (int argc, const char* argv[]){
-  ns_params params;
+int main (
+  int argc,
+  const char* argv[]
+){
+  char* params=NULL;
+  int K1,K2;
+  int N1,N2;
   unsigned int nsteps;
+  double nu,delta;
+  _Complex double (*g)(int,int);
   int ret;
-  unsigned int computation_nr;
+  unsigned int driving,command;
+  unsigned int print_freq;
 
-  // default computation: phase diagram
-  computation_nr=COMPUTATION_ENSTROPHY;
+  command=0;
+  driving=0;
 
   // read command line arguments
-  ret=read_args(argc, argv, &params, &nsteps, &computation_nr);
+  ret=read_args(argc, argv, &params, &driving, &command);
+  if(ret<0){
+    return(-1);
+  }
+  // read params
+  ret=read_params(params, &K1, &K2, &N1, &N2, &nsteps, &nu, &delta, &print_freq);
   if(ret<0){
     return(-1);
   }
-  if(ret>0){
-    return(0);
+  // set driving force
+  switch(driving){
+  case DRIVING_TEST:
+    g=g_test;
+    break;
+
+  default:
+    g=g_test;
+    break;
   }
 
-  // enstrophy
-  if(computation_nr==COMPUTATION_ENSTROPHY){
-    enstrophy(params, nsteps);
+  // run command
+  if (command==COMMAND_UK){
+    uk(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq);
+  }
+  else if(command==COMMAND_ENSTROPHY){
+    enstrophy(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq);
+  }
+  else if(command==0){
+    fprintf(stderr, "error: no command specified\n");
+    print_usage();
   }
 
   return(0);
@@ -45,44 +79,25 @@ int main (int argc, const char* argv[]){
 
 // usage message
 int print_usage(){
-  fprintf(stderr, "usage:\n       nstrophy enstrophy [-h timestep] [-K modes] [-v] [-N nsteps]\n\n       nstrophy -V [-v]\n\n");
+  fprintf(stderr, "usage:\n       nstrophy [-p parameters] [-g driving_force] <command>\n\n");
   return(0);
 }
 
 // read command line arguments
-#define CP_FLAG_TIMESTEP 1
-#define CP_FLAG_NSTEPS 2
-#define CP_FLAG_MODES 3
-#define CP_FLAG_NU 4
-int read_args(int argc, const char* argv[], ns_params* params, unsigned int* nsteps, unsigned int* computation_nr){
+#define CP_FLAG_PARAMS 1
+#define CP_FLAG_DRIVING 2
+int read_args(
+  int argc,
+  const char* argv[],
+  char** params,
+  unsigned int* driving_force,
+  unsigned int* command
+){
   int i;
-  int ret;
-  // temporary int
-  int tmp_int;
-  // temporary unsigned int
-  unsigned int tmp_uint;
-  // temporary double
-  double tmp_double;
   // pointers
   char* ptr;
   // flag that indicates what argument is being read
   int flag=0;
-  // print version and exit
-  char Vflag=0;
-
-  // defaults
-  /*
-  params->K=16;
-  params->h=1e-3/(2*params->K+1);
-  *nsteps=10000000;
-  params->nu=1./1024/(2*params->K+1);
-  */
-  params->K=16;
-  //h=2^-13
-  params->h=0.0001220703125;
-  //nu=2^-11
-  *nsteps=10000000;
-  params->nu=0.00048828125;
 
   // loop over arguments
   for(i=1;i<argc;i++){
@@ -91,24 +106,12 @@ int read_args(int argc, const char* argv[], ns_params* params, unsigned int* nst
       for(ptr=((char*)argv[i])+1;*ptr!='\0';ptr++){
 	switch(*ptr){
 	// timestep
-	case 'h':
-	  flag=CP_FLAG_TIMESTEP;
+	case 'p':
+	  flag=CP_FLAG_PARAMS;
 	  break;
 	// nsteps
-	case 'N':
-	  flag=CP_FLAG_NSTEPS;
-	  break;
-	// modes
-	case 'K':
-	  flag=CP_FLAG_MODES;
-	  break;
-	// friction
-	case 'n':
-	  flag=CP_FLAG_NU;
-	  break;
-	// print version
-	case 'V':
-	  Vflag=1;
+	case 'g':
+	  flag=CP_FLAG_DRIVING;
 	  break;
 	default:
 	  fprintf(stderr, "unrecognized option '-%c'\n", *ptr);
@@ -118,206 +121,241 @@ int read_args(int argc, const char* argv[], ns_params* params, unsigned int* nst
 	}
       }
     }
-    // timestep
-    else if(flag==CP_FLAG_TIMESTEP){
-      ret=sscanf(argv[i],"%lf",&tmp_double);
-      if(ret!=1){
-	fprintf(stderr, "error: '-h' should be followed by a double\n       got '%s'\n",argv[i]);
-	return(-1);
-      }
-      params->h=tmp_double;
+    // params
+    else if(flag==CP_FLAG_PARAMS){
+      *params=(char*)argv[i];
       flag=0;
     }
-    // nsteps
-    else if(flag==CP_FLAG_NSTEPS){
-      ret=sscanf(argv[i],"%u",&tmp_uint);
-      if(ret!=1){
-	fprintf(stderr, "error: '-N' should be followed by an unsigned int\n       got '%s'\n",argv[i]);
-	return(-1);
+    // driving force
+    else if(flag==CP_FLAG_DRIVING){
+      if (strcmp(argv[i],"test")==0){
+	*driving_force=DRIVING_TEST;
       }
-      *nsteps=tmp_uint;
-      flag=0;
-    }
-    // modes
-    else if(flag==CP_FLAG_MODES){
-      ret=sscanf(argv[i],"%d",&tmp_int);
-      if(ret!=1){
-	fprintf(stderr, "error: '-K' should be followed by an int\n       got '%s'\n",argv[i]);
-	return(-1);
-      }
-      params->K=tmp_int;
-      flag=0;
-    }
-    // friction
-    else if(flag==CP_FLAG_TIMESTEP){
-      ret=sscanf(argv[i],"%lf",&tmp_double);
-      if(ret!=1){
-	fprintf(stderr, "error: '-n' should be followed by a double\n       got '%s'\n",argv[i]);
+      else{
+	fprintf(stderr, "error: unrecognized driving force '%s'\n",argv[i]);
 	return(-1);
       }
-      params->nu=tmp_double;
       flag=0;
     }
     // computation to run
     else{
-      if(strcmp(argv[i], "enstrophy")==0){
-	*computation_nr=COMPUTATION_ENSTROPHY;
+      if(strcmp(argv[i], "uk")==0){
+	*command=COMMAND_UK;
+      }
+      else if(strcmp(argv[i], "enstrophy")==0){
+	*command=COMMAND_ENSTROPHY;
       }
       else{
-	fprintf(stderr, "error: unrecognized computation: '%s'\n",argv[i]);
-	print_usage();
+	fprintf(stderr, "error: unrecognized command: '%s'\n",argv[i]);
 	return(-1);
       }
       flag=0;
     }
   }
 
-  // print version and exit
-  if(Vflag==1){
-    printf("nstrophy " VERSION "\n");
-    return(1);
-  }
-
   return(0);
 }
 
-// compute enstrophy as a function of time in the I-NS equation
-int enstrophy(ns_params params, unsigned int Nsteps){
-  _Complex double* u;
-  _Complex double* tmp1;
-  _Complex double* tmp2;
-  _Complex double* tmp3;
-  _Complex double alpha;
-  _Complex double avg;
-  unsigned int t;
-  int kx,ky;
-  fft_vects fft_vects;
-  double rescale;
+// read parameters string
+int read_params(
+  char* params,
+  int* K1,
+  int* K2,
+  int* N1,
+  int* N2,
+  unsigned int* nsteps,
+  double* nu,
+  double* delta,
+  unsigned int* print_freq
+){
+  int ret;
+  // pointer in params
+  char* ptr;
+  // buffer and associated pointer
+  char *buffer_lhs, *lhs_ptr;
+  char *buffer_rhs, *rhs_ptr;
+  // whether N was set explicitly
+  bool setN1=false;
+  bool setN2=false;
+  // whether lhs (false is rhs)
+  bool lhs=true;
+
+  // defaults
+  *K1=16;
+  *K2=*K1;
+  //delta=2^-13
+  *delta=0.0001220703125;
+  //nu=2^-11
+  *nu=0.00048828125;
+  *nsteps=10000000;
+  *print_freq=1000;
+
+  if (params!=NULL){
+    // init
+    buffer_lhs=calloc(sizeof(char),strlen(params));
+    lhs_ptr=buffer_lhs;
+    *lhs_ptr='\0';
+    buffer_rhs=calloc(sizeof(char),strlen(params));
+    rhs_ptr=buffer_rhs;
+    *rhs_ptr='\0';
 
-  // sizes
-  params.S=2*params.K+1;
-  params.N=4*params.K+1;
+    for(ptr=params;*ptr!='\0';ptr++){
+      switch(*ptr){
+      case '=':
+	// reset buffer
+	rhs_ptr=buffer_rhs;
+	*rhs_ptr='\0';
+	lhs=false;
+	break;
+      case ';':
+	//set parameter
+	ret=set_parameter(buffer_lhs,buffer_rhs,K1,K2,N1,N2,nsteps,nu,delta,print_freq,&setN1,&setN2);
+	if(ret<0){
+	  return ret;
+	}
+	// reset buffer
+	lhs_ptr=buffer_lhs;
+	*lhs_ptr='\0';
+	lhs=true;
+	break;
+      default:
+	// add to buffer
+	if (lhs){
+	  *lhs_ptr=*ptr;
+	  lhs_ptr++;
+	  *lhs_ptr='\0';
+	}
+	else{
+	  *rhs_ptr=*ptr;
+	  rhs_ptr++;
+	  *rhs_ptr='\0';
+	}
+	break;
+      }
+    }
 
-  // velocity field
-  u=calloc(sizeof(_Complex double),params.S*params.S);
-  params.g=calloc(sizeof(_Complex double),params.S*params.S);
-  // allocate tmp vectors for computation
-  tmp1=calloc(sizeof(_Complex double),params.S*params.S);
-  tmp2=calloc(sizeof(_Complex double),params.S*params.S);
-  tmp3=calloc(sizeof(_Complex double),params.S*params.S);
+    // set last param
+    if (*params!='\0'){
+      ret=set_parameter(buffer_lhs,buffer_rhs,K1,K2,N1,N2,nsteps,nu,delta,print_freq,&setN1,&setN2);
+      if(ret<0){
+	return ret;
+      }
+    }
 
-  /*
-  srand(17);
+    // free vects
+    free(buffer_lhs);
+    free(buffer_rhs);
+  }
 
-  // initial value
-  for(ky=0;ky<=params.K;ky++){
-    u[KLOOKUP(0,ky,params.S)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
+  // if N not set explicitly, set it
+  if (!setN1){
+    *N1=4*(*K1)+1;
   }
-  for(kx=1;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      u[KLOOKUP(kx,ky,params.S)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
-    }
+  if (!setN2){
+    *N2=4*(*K2)+1;
   }
-  for(ky=-params.K;ky<=-1;ky++){
-    u[KLOOKUP(0,ky,params.S)]=conj(u[KLOOKUP(0,-ky,params.S)]);
+
+  return(0);
+}
+
+
+// set a parameter from the parameter string
+int set_parameter(
+  char* lhs,
+  char* rhs,
+  int* K1,
+  int* K2,
+  int* N1,
+  int* N2,
+  unsigned int* nsteps,
+  double* nu,
+  double* delta,
+  unsigned int* print_freq,
+  bool* setN1,
+  bool* setN2
+){
+  int ret;
+
+  if (strcmp(lhs,"K1")==0){
+    ret=sscanf(rhs,"%d",K1);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'K1' should be an integer\n       got '%s'\n",rhs);
+      return(-1);
+    }
   }
-  for(kx=-params.K;kx<=-1;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      u[KLOOKUP(kx,ky,params.S)]=conj(u[KLOOKUP(-kx,-ky,params.S)]);
+  else if (strcmp(lhs,"K2")==0){
+    ret=sscanf(rhs,"%d",K2);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'K2' should be an integer\n       got '%s'\n",rhs);
+      return(-1);
     }
   }
-  rescale=0;
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      rescale=rescale+((__real__ u[KLOOKUP(kx,ky,params.S)])*(__real__ u[KLOOKUP(kx,ky,params.S)])+(__imag__ u[KLOOKUP(kx,ky,params.S)])*(__imag__ u[KLOOKUP(kx,ky,params.S)]))*(kx*kx+ky*ky);
+  else if (strcmp(lhs,"K")==0){
+    ret=sscanf(rhs,"%d",K1);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'K' should be an integer\n       got '%s'\n",rhs);
+      return(-1);
     }
+    *K2=*K1;
   }
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      u[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]*sqrt(155.1/rescale);
+  else if (strcmp(lhs,"N1")==0){
+    ret=sscanf(rhs,"%d",N1);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'N1' should be an integer\n       got '%s'\n",rhs);
+      return(-1);
     }
+    *setN1=true;
   }
-  */
-  /*
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      u[KLOOKUP(kx,ky,params.S)]=1.;
+  else if (strcmp(lhs,"N2")==0){
+    ret=sscanf(rhs,"%d",N2);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'N2' should be an integer\n       got '%s'\n",rhs);
+      return(-1);
     }
+    *setN2=true;
   }
-  */
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      u[KLOOKUP(kx,ky,params.S)]=exp(-sqrt(kx*kx+ky*ky));
+  else if (strcmp(lhs,"N")==0){
+    ret=sscanf(rhs,"%d",N1);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'N' should be an integer\n       got '%s'\n",rhs);
+      return(-1);
     }
+    *N2=*N1;
+    *setN1=true;
+    *setN2=true;
   }
-
-
-  // driving force
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      //params.g[KLOOKUP(kx,ky,params.S)]=sqrt(kx*kx*ky*ky)*exp(-(kx*kx+ky*ky));
-      if(kx==2 && ky==-1){
-	params.g[KLOOKUP(kx,ky,params.S)]=0.5+sqrt(3)/2*I;
-      }
-      else if(kx==-2 && ky==1){
-	params.g[KLOOKUP(kx,ky,params.S)]=0.5-sqrt(3)/2*I;
-      }
-      else{
-	params.g[KLOOKUP(kx,ky,params.S)]=0;
-      }
+  else if (strcmp(lhs,"nsteps")==0){
+    ret=sscanf(rhs,"%u",nsteps);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'nsteps' should be an unsigned integer\n       got '%s'\n",rhs);
+      return(-1);
     }
   }
-
-
-  // prepare vectors for fft
-  fft_vects.fft1=fftw_malloc(sizeof(fftw_complex)*params.N*params.N);
-  fft_vects.fft1_plan=fftw_plan_dft_2d((int)params.N,(int)params.N, fft_vects.fft1, fft_vects.fft1, FFTW_FORWARD, FFTW_MEASURE);
-  fft_vects.fft2=fftw_malloc(sizeof(fftw_complex)*params.N*params.N);
-  fft_vects.fft2_plan=fftw_plan_dft_2d((int)params.N,(int)params.N, fft_vects.fft2, fft_vects.fft2, FFTW_FORWARD, FFTW_MEASURE);
-  fft_vects.invfft=fftw_malloc(sizeof(fftw_complex)*params.N*params.N);
-  fft_vects.invfft_plan=fftw_plan_dft_2d((int)params.N,(int)params.N, fft_vects.invfft, fft_vects.invfft, FFTW_BACKWARD, FFTW_MEASURE);
-
-  // init running average
-  avg=0;
-
-  // iterate
-  for(t=0;t<Nsteps;t++){
-    ins_step(u, params, fft_vects, tmp1, tmp2, tmp3);
-    alpha=compute_alpha(u, params);
-    
-    /*
-    // to avoid errors building up in imaginary part
-    for(kx=-params.K;kx<=params.K;kx++){
-      for(ky=-params.K;ky<=params.K;ky++){
-	u[KLOOKUP(kx,ky,params.S)]=__real__ u[KLOOKUP(kx,ky,params.S)];
-      }
+  else if (strcmp(lhs,"nu")==0){
+    ret=sscanf(rhs,"%lf",nu);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'nu' should be a double\n       got '%s'\n",rhs);
+      return(-1);
     }
-    */
-
-    // running average
-    if(t>0){
-      avg=avg-(avg-alpha)/t;
+  }
+  else if (strcmp(lhs,"delta")==0){
+    ret=sscanf(rhs,"%lf",delta);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'delta' should be a double\n       got '%s'\n",rhs);
+      return(-1);
     }
-
-    if(t>0 && t%1000==0){
-      fprintf(stderr,"% .15e % .15e % .15e % .15e % .15e\n",t*params.h, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
-      printf("% .15e % .15e % .15e % .15e % .15e\n",t*params.h, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
+  }
+  else if (strcmp(lhs,"print_freq")==0){
+    ret=sscanf(rhs,"%u",print_freq);
+    if(ret!=1){
+      fprintf(stderr, "error: parameter 'print_freq' should be an unsigned integer\n       got '%s'\n",rhs);
+      return(-1);
     }
   }
-
-  // free memory
-  fftw_destroy_plan(fft_vects.fft1_plan);
-  fftw_destroy_plan(fft_vects.fft2_plan);
-  fftw_destroy_plan(fft_vects.invfft_plan);
-  fftw_free(fft_vects.fft1);
-  fftw_free(fft_vects.fft2);
-  fftw_free(fft_vects.invfft);
-
-  free(tmp3);
-  free(tmp2);
-  free(tmp1);
-  free(params.g);
-  free(u);
+  else{
+    fprintf(stderr, "error: unrecognized parameter '%s'\n",lhs);
+    return(-1);
+  }
 
   return(0);
 }
diff --git a/src/navier-stokes.c b/src/navier-stokes.c
index 303392d..ab4803c 100644
--- a/src/navier-stokes.c
+++ b/src/navier-stokes.c
@@ -1,63 +1,311 @@
 #include "navier-stokes.h"
 #include <math.h>
+#include <stdlib.h>
 
 #define M_PI 3.14159265358979323846
 
+// compute solution as a function of time
+int uk(
+  int K1,
+  int K2,
+  int N1,
+  int N2,
+  unsigned int nsteps,
+  double nu,
+  double delta,
+  _Complex double (*g)(int,int),
+  unsigned int print_freq
+){
+  _Complex double* u;
+  _Complex double* tmp1;
+  _Complex double* tmp2;
+  _Complex double* tmp3;
+  unsigned int t;
+  fft_vect fft1;
+  fft_vect fft2;
+  fft_vect ifft;
+  int kx,ky;
+
+  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2);
+  ns_init_u(u, K1, K2);
+
+  // iterate
+  for(t=0;t<nsteps;t++){
+    ins_step(u, K1, K2, N1, N2, nu, delta, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
+    
+    if(t%print_freq==0){
+      fprintf(stderr,"% .8e ",t*delta);
+      printf("% .15e ",t*delta);
+
+      for(kx=-K1;kx<=K1;kx++){
+	for (ky=-K2;ky<=K2;ky++){
+	  if (kx*kx+ky*ky<=1){
+	    fprintf(stderr,"% .8e % .8e ",__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)], __imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]);
+	    
+	  }
+	  printf("% .8e % .8e ",__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)], __imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]);
+	}
+      }
+      fprintf(stderr,"\n");
+      printf("\n");
+    }
+  }
+
+  ns_free_tmps(u, tmp1, tmp2, tmp3, fft1, fft2, ifft);
+  return(0);
+}
+
+// compute enstrophy as a function of time in the I-NS equation
+int enstrophy(
+  int K1,
+  int K2,
+  int N1,
+  int N2,
+  unsigned int nsteps,
+  double nu,
+  double delta,
+  _Complex double (*g)(int,int),
+  unsigned int print_freq
+){
+  _Complex double* u;
+  _Complex double* tmp1;
+  _Complex double* tmp2;
+  _Complex double* tmp3;
+  _Complex double alpha;
+  _Complex double avg;
+  unsigned int t;
+  fft_vect fft1;
+  fft_vect fft2;
+  fft_vect ifft;
+
+  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2);
+  ns_init_u(u, K1, K2);
+
+
+  // init running average
+  avg=0;
+
+  // iterate
+  for(t=0;t<nsteps;t++){
+    ins_step(u, K1, K2, N1, N2, nu, delta, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
+    alpha=compute_alpha(u, K1, K2, g);
+    
+    // running average
+    if(t>0){
+      avg=avg-(avg-alpha)/t;
+    }
+
+    if(t>0 && t%print_freq==0){
+      fprintf(stderr,"% .15e % .15e % .15e % .15e % .15e\n",t*delta, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
+      printf("% .15e % .15e % .15e % .15e % .15e\n",t*delta, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
+    }
+  }
+
+  ns_free_tmps(u, tmp1, tmp2, tmp3, fft1, fft2, ifft);
+  return(0);
+}
+
+
+// initialize vectors for computation
+int  ns_init_tmps(
+  _Complex double ** u,
+  _Complex double ** tmp1,
+  _Complex double ** tmp2,
+  _Complex double ** tmp3,
+  fft_vect* fft1,
+  fft_vect* fft2,
+  fft_vect* ifft,
+  int K1,
+  int K2,
+  int N1,
+  int N2
+){
+  // velocity field
+  *u=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
+
+  // allocate tmp vectors for computation
+  *tmp1=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
+  *tmp2=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
+  *tmp3=calloc(sizeof(_Complex double),(2*K1+1)*(2*K2+1));
+
+  // prepare vectors for fft
+  fft1->fft=fftw_malloc(sizeof(fftw_complex)*N1*N2);
+  fft1->fft_plan=fftw_plan_dft_2d(N1,N2, fft1->fft, fft1->fft, FFTW_FORWARD, FFTW_MEASURE);
+  fft2->fft=fftw_malloc(sizeof(fftw_complex)*N1*N2);
+  fft2->fft_plan=fftw_plan_dft_2d(N1,N2, fft2->fft, fft2->fft, FFTW_FORWARD, FFTW_MEASURE);
+  ifft->fft=fftw_malloc(sizeof(fftw_complex)*N1*N2);
+  ifft->fft_plan=fftw_plan_dft_2d(N1,N2, ifft->fft, ifft->fft, FFTW_BACKWARD, FFTW_MEASURE);
+
+  return 0;
+}
+
+// release vectors
+int ns_free_tmps(
+  _Complex double* u,
+  _Complex double* tmp1,
+  _Complex double* tmp2,
+  _Complex double* tmp3,
+  fft_vect fft1,
+  fft_vect fft2,
+  fft_vect ifft
+){
+  // free memory
+  fftw_destroy_plan(fft1.fft_plan);
+  fftw_destroy_plan(fft2.fft_plan);
+  fftw_destroy_plan(ifft.fft_plan);
+  fftw_free(fft1.fft);
+  fftw_free(fft2.fft);
+  fftw_free(ifft.fft);
+
+  fftw_cleanup();
+
+  free(tmp3);
+  free(tmp2);
+  free(tmp1);
+
+  free(u);
+
+  return 0;
+}
+
+
+
+// initial value
+int ns_init_u(
+  _Complex double* u,
+  int K1,
+  int K2
+){
+  int kx,ky;
+  /*
+  double rescale;
+
+  srand(17);
+
+  // random init (set half, then the other half are the conjugates)
+  for(ky=0;ky<=K2;ky++){
+    u[klookup(0,ky,2*K1+1,2*K2+1)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
+  }
+  for(kx=1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      u[klookup(kx,ky,2*K1+1,2*K2+1)]=(-RAND_MAX*0.5+rand())*1.0/RAND_MAX+(-RAND_MAX*0.5+rand())*1.0/RAND_MAX*I;
+    }
+  }
+  // conjugates
+  for(ky=-K2;ky<=-1;ky++){
+    u[klookup(0,ky,2*K1+1,2*K2+1)]=conj(u[klookup(0,-ky,2*K1+1,2*K2+1)]);
+  }
+  for(kx=-K1;kx<=-1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      u[klookup(kx,ky,2*K1+1,2*K2+1)]=conj(u[klookup(-kx,-ky,2*K1+1,2*K2+1)]);
+    }
+  }
+
+  // rescale: 1/k decay
+  rescale=0;
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      rescale=rescale+((__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)])*(__real__ u[klookup(kx,ky,2*K1+1,2*K2+1)])+(__imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)])*(__imag__ u[klookup(kx,ky,2*K1+1,2*K2+1)]))*(kx*kx+ky*ky);
+    }
+  }
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      u[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]*sqrt(155.1/rescale);
+    }
+  }
+  */
+
+
+  /*
+  // constant init
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      u[klookup(kx,ky,2*K1+1,2*K2+1)]=1.;
+    }
+  }
+  */
+
+  // exponentially decaying init
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      u[klookup(kx,ky,2*K1+1,2*K2+1)]=exp(-sqrt(kx*kx+ky*ky));
+    }
+  }
+
+  return 0;
+}
+
+
 // next time step for Irreversible Navier-Stokes equation
-int ins_step(_Complex double* u, ns_params params, fft_vects vects, _Complex double* tmp1, _Complex double* tmp2, _Complex double* tmp3){
+int ins_step(
+  _Complex double* u,
+  int K1,
+  int K2,
+  int N1,
+  int N2,
+  double nu,
+  double delta,
+  _Complex double (*g)(int,int),
+  fft_vect fft1,
+  fft_vect fft2,
+  fft_vect ifft,
+  _Complex double* tmp1,
+  _Complex double* tmp2,
+  _Complex double* tmp3
+){
   int kx,ky;
 
   // k1
-  ins_rhs(tmp1, u, params, vects);
+  ins_rhs(tmp1, u, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
   // add to output
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      tmp3[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h/6*tmp1[KLOOKUP(kx,ky,params.S)];
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/6*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
     }
   }
 
   // u+h*k1/2
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      tmp2[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h/2*tmp1[KLOOKUP(kx,ky,params.S)];
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/2*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
     }
   }
   // k2
-  ins_rhs(tmp1, tmp2, params, vects);
+  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
   // add to output
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      tmp3[KLOOKUP(kx,ky,params.S)]+=params.h/3*tmp1[KLOOKUP(kx,ky,params.S)];
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+=delta/3*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
     }
   }
 
   // u+h*k2/2
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      tmp2[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h/2*tmp1[KLOOKUP(kx,ky,params.S)];
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/2*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
     }
   }
   // k3
-  ins_rhs(tmp1, tmp2, params, vects);
+  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
   // add to output
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      tmp3[KLOOKUP(kx,ky,params.S)]+=params.h/3*tmp1[KLOOKUP(kx,ky,params.S)];
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+=delta/3*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
     }
   }
 
   // u+h*k3
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      tmp2[KLOOKUP(kx,ky,params.S)]=u[KLOOKUP(kx,ky,params.S)]+params.h*tmp1[KLOOKUP(kx,ky,params.S)];
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      tmp2[klookup(kx,ky,2*K1+1,2*K2+1)]=u[klookup(kx,ky,2*K1+1,2*K2+1)]+delta*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
     }
   }
   // k4
-  ins_rhs(tmp1, tmp2, params, vects);
+  ins_rhs(tmp1, tmp2, K1, K2, N1, N2, nu, g, fft1, fft2, ifft);
   // add to output
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      u[KLOOKUP(kx,ky,params.S)]=tmp3[KLOOKUP(kx,ky,params.S)]+params.h/6*tmp1[KLOOKUP(kx,ky,params.S)];
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      u[klookup(kx,ky,2*K1+1,2*K2+1)]=tmp3[klookup(kx,ky,2*K1+1,2*K2+1)]+delta/6*tmp1[klookup(kx,ky,2*K1+1,2*K2+1)];
     }
   }
 
@@ -65,74 +313,82 @@ int ins_step(_Complex double* u, ns_params params, fft_vects vects, _Complex dou
 }
 
 // right side of Irreversible Navier-Stokes equation
-int ins_rhs(_Complex double* out, _Complex double* u, ns_params params, fft_vects vects){
+int ins_rhs(
+  _Complex double* out,
+  _Complex double* u,
+  int K1,
+  int K2,
+  int N1,
+  int N2,
+  double nu,
+  _Complex double (*g)(int,int),
+  fft_vect fft1,
+  fft_vect fft2,
+  fft_vect ifft
+){
   int kx,ky;
+  int i;
 
   // F(px/|p|*u)*F(qy*|q|*u)
   // init to 0
-  for(kx=0; kx<params.N*params.N; kx++){
-    vects.fft1[kx]=0;
-    vects.fft2[kx]=0;
-    vects.invfft[kx]=0;
+  for(i=0; i<N1*N2; i++){
+    fft1.fft[i]=0;
+    fft2.fft[i]=0;
+    ifft.fft[i]=0;
   }
   // fill modes
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
       if(kx!=0 || ky!=0){
-        vects.fft1[KLOOKUP(kx,ky,params.N)]=kx/sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
-        vects.fft2[KLOOKUP(kx,ky,params.N)]=ky*sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
+        fft1.fft[klookup(kx,ky,N1,N2)]=kx/sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)];
+        fft2.fft[klookup(kx,ky,N1,N2)]=ky*sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)];
       }
     }
   }
 
   // fft
-  fftw_execute(vects.fft1_plan);
-  
-  fftw_execute(vects.fft2_plan);
-  // write to invfft
-  for(kx=-2*params.K;kx<=2*params.K;kx++){
-    for(ky=-2*params.K;ky<=2*params.K;ky++){
-      vects.invfft[KLOOKUP(kx,ky,params.N)]=vects.fft1[KLOOKUP(kx,ky,params.N)]*vects.fft2[KLOOKUP(kx,ky,params.N)];
-    }
+  fftw_execute(fft1.fft_plan);
+  fftw_execute(fft2.fft_plan);
+  // write to ifft
+  for(i=0;i<N1*N2;i++){
+    ifft.fft[i]=fft1.fft[i]*fft2.fft[i];
   }
 
   // F(py/|p|*u)*F(qx*|q|*u)
   // init to 0
-  for(kx=0; kx<params.N*params.N; kx++){
-    vects.fft1[kx]=0;
-    vects.fft2[kx]=0;
+  for(i=0; i<N1*N2; i++){
+    fft1.fft[i]=0;
+    fft2.fft[i]=0;
   }
   // fill modes
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
       if(kx!=0 || ky!=0){
-        vects.fft1[KLOOKUP(kx,ky,params.N)]=ky/sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
-        vects.fft2[KLOOKUP(kx,ky,params.N)]=kx*sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
+        fft1.fft[klookup(kx,ky,N1,N2)]=ky/sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)];
+        fft2.fft[klookup(kx,ky,N1,N2)]=kx*sqrt(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)];
       }
     }
   }
 
   // fft
-  fftw_execute(vects.fft1_plan);
-  fftw_execute(vects.fft2_plan);
-  // write to invfft
-  for(kx=-2*params.K;kx<=2*params.K;kx++){
-    for(ky=-2*params.K;ky<=2*params.K;ky++){
-      vects.invfft[KLOOKUP(kx,ky,params.N)]=vects.invfft[KLOOKUP(kx,ky,params.N)]-vects.fft1[KLOOKUP(kx,ky,params.N)]*vects.fft2[KLOOKUP(kx,ky,params.N)];
-    }
+  fftw_execute(fft1.fft_plan);
+  fftw_execute(fft2.fft_plan);
+  // write to ifft
+  for(i=0;i<N1*N2;i++){
+    ifft.fft[i]=ifft.fft[i]-fft1.fft[i]*fft2.fft[i];
   }
 
   // inverse fft
-  fftw_execute(vects.invfft_plan);
+  fftw_execute(ifft.fft_plan);
 
   // write out
-  for(kx=0; kx<params.S*params.S; kx++){
-    out[kx]=0;
+  for(i=0; i<(2*K1+1)*(2*K2+1); i++){
+    out[i]=0;
   }
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
       if(kx!=0 || ky!=0){
-        out[KLOOKUP(kx,ky,params.S)]=-4*M_PI*M_PI*params.nu*(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)]+params.g[KLOOKUP(kx,ky,params.S)]+4*M_PI*M_PI/sqrt(kx*kx+ky*ky)*vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N;
+        out[klookup(kx,ky,2*K1+1,2*K2+1)]=-4*M_PI*M_PI*nu*(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]+(*g)(kx,ky)+4*M_PI*M_PI/sqrt(kx*kx+ky*ky)*ifft.fft[klookup(kx,ky,N1,N2)]/N1/N2;
       }
     }
   }
@@ -142,17 +398,33 @@ int ins_rhs(_Complex double* out, _Complex double* u, ns_params params, fft_vect
 
 
 // compute alpha
-_Complex double compute_alpha(_Complex double* u, ns_params params){
+_Complex double compute_alpha(
+  _Complex double* u,
+  int K1,
+  int K2,
+  _Complex double (*g)(int,int)
+){
   _Complex double num=0;
   _Complex double denom=0;
   int kx,ky;
 
-  for(kx=-params.K;kx<=params.K;kx++){
-    for(ky=-params.K;ky<=params.K;ky++){
-      denom+=(kx*kx+ky*ky)*(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)]*conj(u[KLOOKUP(kx,ky,params.S)])*(1+(ky!=0?kx*kx/ky/ky:0));
-      num+=(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)]*conj(params.g[KLOOKUP(kx,ky,params.S)])*(1+(ky!=0?kx*kx/ky/ky:0));
+  for(kx=-K1;kx<=K1;kx++){
+    for(ky=-K2;ky<=K2;ky++){
+      denom+=(kx*kx+ky*ky)*(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(u[klookup(kx,ky,2*K1+1,2*K2+1)])*(1+(ky!=0?kx*kx/ky/ky:0));
+      num+=(kx*kx+ky*ky)*u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj((*g)(kx,ky))*(1+(ky!=0?kx*kx/ky/ky:0));
     }
   }
 
   return(num/denom);
 }
+
+
+// get index for kx,ky in array of size S
+int klookup(
+  int kx,
+  int ky,
+  int S1,
+  int S2
+){
+  return (kx>=0 ? kx : S1+kx)*S2 + (ky>=0 ? ky : S2+ky);
+}
diff --git a/src/navier-stokes.h b/src/navier-stokes.h
index cd093d7..5ac8d7d 100644
--- a/src/navier-stokes.h
+++ b/src/navier-stokes.h
@@ -4,45 +4,38 @@
 #include <complex.h>
 #include <fftw3.h>
 
-// to extract elements from array of size S representing a function of momentum, use
-//   array[KEXTRACT(kx,ky,size)]
-#define KLOOKUP(X,Y,S) (X>=0?X:S+X)*S+(Y>=0?Y:S+Y)
-
-
-// parameters for the NS equation
-typedef struct ns_params {
-  // number of modes
-  int K;
-  // 2*K+1
-  int S;
-  // 4*K+1
-  int N;
-  // forcing term
-  _Complex double* g;
-  // time step
-  double h;
-  // friction
-  double nu;
-} ns_params;
 
 // arrays on which the ffts are performed
 typedef struct fft_vects {
-  fftw_complex* fft1;
-  fftw_complex* fft2;
-  fftw_complex* invfft;
-  fftw_plan fft1_plan;
-  fftw_plan fft2_plan;
-  fftw_plan invfft_plan;
-} fft_vects;
+  fftw_complex* fft;
+  fftw_plan fft_plan;
+} fft_vect;
+
+// compute u_k
+int uk( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int), unsigned int print_freq);
+
+// compute enstrophy
+int enstrophy( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, _Complex double (*g)(int,int), unsigned int print_freq);
+
+// initialize vectors for computation
+int  ns_init_tmps( _Complex double **u, _Complex double ** tmp1, _Complex double **tmp2, _Complex double **tmp3, fft_vect* fft1, fft_vect *fft2, fft_vect *ifft, int K1, int K2, int N1, int N2);
+// release vectors
+int ns_free_tmps( _Complex double* u, _Complex double* tmp1, _Complex double *tmp2,_Complex double *tmp3, fft_vect fft1, fft_vect fft2, fft_vect ifft);
+
+// initial value
+int ns_init_u( _Complex double* u, int K1, int K2);
 
 // next time step for Irreversible Navier-Stokes equation
-int ins_step(_Complex double* u, ns_params params, fft_vects vects, _Complex double* tmp1, _Complex double* tmp2, _Complex double* tmp3);
+int ins_step( _Complex double* u, int K1, int K2, int N1, int N2, double nu, double delta, _Complex double (*g)(int,int), fft_vect fft1, fft_vect fft2,fft_vect ifft, _Complex double* tmp1, _Complex double *tmp2, _Complex double *tmp3);
 
 // right side of Irreversible Navier-Stokes equation
-int ins_rhs(_Complex double* out, _Complex double* u, ns_params params, fft_vects vects);
+int ins_rhs( _Complex double* out, _Complex double* u, int K1, int K2, int N1, int N2, double nu, _Complex double (*g)(int,int), fft_vect fft1, fft_vect fft2, fft_vect ifft);
 
 // compute alpha
-_Complex double compute_alpha(_Complex double* u, ns_params params);
+_Complex double compute_alpha( _Complex double* u, int K1, int K2, _Complex double (*g)(int,int));
+
+// get index for kx,ky in array of size S
+int klookup( int kx, int ky, int S1, int S2);
 
 #endif