From 01f47ace6756c28deb9ea0daaee3904ffa5ce9e0 Mon Sep 17 00:00:00 2001
From: Ian Jauslin <jauslin@ias.edu>
Date: Thu, 11 Jan 2018 22:48:14 +0000
Subject: Initial commit

---
 src/main.c          | 282 +++++++++++++++++++++++++++++++++++++++++++++++++
 src/navier-stokes.c | 299 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/navier-stokes.h |  48 +++++++++
 3 files changed, 629 insertions(+)
 create mode 100644 src/main.c
 create mode 100644 src/navier-stokes.c
 create mode 100644 src/navier-stokes.h

(limited to 'src')

diff --git a/src/main.c b/src/main.c
new file mode 100644
index 0000000..9420e90
--- /dev/null
+++ b/src/main.c
@@ -0,0 +1,282 @@
+#define VERSION "0.0"
+
+#include <math.h>
+#include <complex.h>
+#include <fftw3.h>
+#include <string.h>
+#include <stdlib.h>
+#include "navier-stokes.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);
+
+// compute enstrophy as a function of time in the I-NS equation
+int enstrophy(ns_params params, unsigned int Nsteps);
+
+
+#define COMPUTATION_ENSTROPHY 1
+int main (int argc, const char* argv[]){
+  ns_params params;
+  unsigned int nsteps;
+  int ret;
+  unsigned int computation_nr;
+
+  // default computation: phase diagram
+  computation_nr=COMPUTATION_ENSTROPHY;
+
+  // read command line arguments
+  ret=read_args(argc, argv, &params, &nsteps, &computation_nr);
+  if(ret<0){
+    return(-1);
+  }
+  if(ret>0){
+    return(0);
+  }
+
+  // enstrophy
+  if(computation_nr==COMPUTATION_ENSTROPHY){
+    enstrophy(params, nsteps);
+  }
+
+  return(0);
+}
+
+// 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");
+  return(0);
+}
+
+// read command line arguments
+#define CP_FLAG_TIMESTEP 1
+#define CP_FLAG_NSTEPS 2
+#define CP_FLAG_MODES 2
+#define CP_FLAG_NU 3
+int read_args(int argc, const char* argv[], ns_params* params, unsigned int* nsteps, unsigned int* computation_nr){
+  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->h=6.103515625e-05;
+  params->K=16;
+  *nsteps=16777216;
+  params->nu=4.9632717887631524e-05;
+  */
+  params->h=1e-5;
+  params->K=16;
+  *nsteps=10000000;
+  params->nu=1e-4;
+
+  // loop over arguments
+  for(i=1;i<argc;i++){
+    // flag
+    if(argv[i][0]=='-'){
+      for(ptr=((char*)argv[i])+1;*ptr!='\0';ptr++){
+	switch(*ptr){
+	// timestep
+	case 'h':
+	  flag=CP_FLAG_TIMESTEP;
+	  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;
+	  break;
+	default:
+	  fprintf(stderr, "unrecognized option '-%c'\n", *ptr);
+	  print_usage();
+	  return(-1);
+	  break;
+	}
+      }
+    }
+    // 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;
+      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);
+      }
+      *nsteps=tmp_uint;
+      flag=0;
+    }
+    // nsteps
+    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_uint;
+      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]);
+	return(-1);
+      }
+      params->nu=tmp_double;
+      flag=0;
+    }
+    // computation to run
+    else{
+      if(strcmp(argv[i], "enstrophy")==0){
+	*computation_nr=COMPUTATION_ENSTROPHY;
+      }
+      else{
+	fprintf(stderr, "error: unrecognized computation: '%s'\n",argv[i]);
+	print_usage();
+	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;
+
+  // sizes
+  params.S=2*params.K+1;
+  params.N=4*params.K+1;
+
+  // 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);
+
+  // initial value
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      //u[KLOOKUP(kx,ky,params.S)]=kx*kx*ky*ky*exp(-(kx*kx+ky*ky));
+      if((kx==1 && ky==0) || (kx==-1 && ky==0)){
+	u[KLOOKUP(kx,ky,params.S)]=1;
+      }
+      else{
+	u[KLOOKUP(kx,ky,params.S)]=0;
+      }
+    }
+  }
+
+  // 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) || (kx==-2 && ky==1)){
+	params.g[KLOOKUP(kx,ky,params.S)]=1;
+      }
+      else{
+	params.g[KLOOKUP(kx,ky,params.S)]=0;
+      }
+    }
+  }
+
+
+  // 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)];
+      }
+    }
+
+    // running average
+    if(t>0){
+      avg=avg-(avg-alpha)/t;
+    }
+
+    if(t>0 && t%1000==0){
+      fprintf(stderr,"%8d % .8e % .8e % .8e % .8e\n",t, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
+      printf("%8d % .8e % .8e % .8e % .8e\n",t, __real__ avg, __imag__ avg, __real__ alpha, __imag__ alpha);
+    }
+  }
+
+  // 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);
+
+  return(0);
+}
diff --git a/src/navier-stokes.c b/src/navier-stokes.c
new file mode 100644
index 0000000..3def963
--- /dev/null
+++ b/src/navier-stokes.c
@@ -0,0 +1,299 @@
+#include "navier-stokes.h"
+#include <math.h>
+
+#define M_PI 3.14159265358979323846
+
+#define CHK 1
+
+// 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 kx,ky;
+
+  // k1
+  ins_rhs(tmp1, u, params, vects);
+  // 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)];
+    }
+  }
+
+  // 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)];
+    }
+  }
+  // k2
+  ins_rhs(tmp1, tmp2, params, vects);
+  // 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)];
+    }
+  }
+
+  // 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)];
+    }
+  }
+  // k3
+  ins_rhs(tmp1, tmp2, params, vects);
+  // 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)];
+    }
+  }
+
+  // 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)];
+    }
+  }
+  // k4
+  ins_rhs(tmp1, tmp2, params, vects);
+  // 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)];
+    }
+  }
+
+  return(0);
+}
+
+// right side of Irreversible Navier-Stokes equation
+int ins_rhs(_Complex double* out, _Complex double* u, ns_params params, fft_vects vects){
+  int kx,ky;
+
+#if CHK==0
+  // F(u/|p|)*F(q1*q2*u/|q|)
+  // init to 0
+  for(kx=0; kx<params.N*params.N; kx++){
+    vects.fft1[kx]=0;
+    vects.fft2[kx]=0;
+    vects.invfft[kx]=0;
+  }
+  // fill modes
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      if(kx!=0 || ky!=0){
+        vects.fft1[KLOOKUP(kx,ky,params.N)]=u[KLOOKUP(kx,ky,params.S)]/sqrt(kx*kx+ky*ky);
+        vects.fft2[KLOOKUP(kx,ky,params.N)]=kx*ky*u[KLOOKUP(kx,ky,params.S)]/sqrt(kx*kx+ky*ky);
+      }
+    }
+  }
+  // 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)];
+    }
+  }
+  // inverse fft
+  fftw_execute(vects.invfft_plan);
+  
+  // write out
+  for(kx=0; kx<params.S*params.S; kx++){
+    out[kx]=0;
+  }
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;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)]*(kx*kx-ky*ky)/params.N/params.N;
+      }
+    }
+  }
+
+  // F(u/|p|)*F((q1*q1-q2*q2)*u/|q|)
+  // init to 0
+  for(kx=0; kx<params.N*params.N; kx++){
+    vects.fft2[kx]=0;
+    vects.invfft[kx]=0;
+  }
+  // fill modes
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      if(kx!=0 || ky!=0){
+        vects.fft2[KLOOKUP(kx,ky,params.N)]=(kx*kx-ky*ky)*u[KLOOKUP(kx,ky,params.S)]/sqrt(kx*kx+ky*ky);
+      }
+    }
+  }
+  // fft
+  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)];
+    }
+  }
+  // inverse fft
+  fftw_execute(vects.invfft_plan);
+  
+  // write out
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      if(kx!=0 || ky!=0){
+        out[KLOOKUP(kx,ky,params.S)]=out[KLOOKUP(kx,ky,params.S)]-4*M_PI*M_PI/sqrt(kx*kx+ky*ky)*vects.invfft[KLOOKUP(kx,ky,params.N)]*(kx*ky)/params.N/params.N;
+      }
+    }
+  }
+#elif CHK == 1
+  // F(-p2/|p|*u)*F(q1*|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;
+  }
+  // fill modes
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;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)]=kx*sqrt(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)];
+      }
+    }
+  }
+  // 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)] - vects.fft1[KLOOKUP(ky,kx,params.N)]*vects.fft2[KLOOKUP(ky,kx,params.N)];
+    }
+  }
+  
+  // inverse fft
+  fftw_execute(vects.invfft_plan);
+  
+  
+  // write out
+  for(kx=0; kx<params.S*params.S; kx++){
+    out[kx]=0;
+  }
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      if(kx!=0 || ky!=0){
+        out[KLOOKUP(kx,ky,params.S)]=-4*M_PI*M_PI*params.nu/params.S*(kx*kx+ky*ky)*u[KLOOKUP(kx,ky,params.S)]+params.g[KLOOKUP(kx,ky,params.S)]+2*M_PI/sqrt(kx*kx+ky*ky)/params.S*vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N;
+      }
+    }
+  }
+
+#elif CHK==2
+  // F(u)*F(q1*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;
+  }
+  // fill modes
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      // u_k
+      vects.fft1[KLOOKUP(kx,ky,params.N)]=u[KLOOKUP(kx,ky,params.S)];
+      // 2i\pi k_x u_k
+      vects.fft2[KLOOKUP(kx,ky,params.N)]=2*M_PI*I*kx*u[KLOOKUP(kx,ky,params.S)];
+    }
+  }
+  // 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)];
+    }
+  }
+
+  // F(p1/p2*u)*F(q2*u)
+  // init to 0
+  for(kx=0; kx<params.N*params.N; kx++){
+    vects.fft1[kx]=0;
+    vects.fft2[kx]=0;
+  }
+  // fill modes
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      // k_x/k_y u_k
+      if(ky!=0){
+	vects.fft1[KLOOKUP(kx,ky,params.N)]=kx/ky*u[KLOOKUP(kx,ky,params.S)];
+      }
+      // 2i\pi k_y u_k
+      vects.fft2[KLOOKUP(kx,ky,params.N)]=2*M_PI*I*ky*u[KLOOKUP(kx,ky,params.S)];
+    }
+  }
+  // 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)];
+    }
+  }
+
+  // inverse fft
+  fftw_execute(vects.invfft_plan);
+
+  /*
+  // check: convolution instead of fft
+  for(kx=0; kx<params.S*params.S; kx++){
+    out[kx]=0;
+  }
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      for(px=-params.K;px<=params.K;px++){
+	for(py=-params.K;py<=params.K;py++){
+	  if(kx-px<=params.K && kx-px>=-params.K && ky-py<=params.K && ky-py>=-params.K){
+	    out[KLOOKUP(kx,ky,params.S)]+=2*M_PI*I*(u[KLOOKUP(px,py,params.S)]*(kx-px)*u[KLOOKUP(kx-px,ky-py,params.S)]-(py==0?0:px/py*u[KLOOKUP(px,py,params.S)]*(ky-py)*u[KLOOKUP(kx-px,ky-py,params.S)]));
+	  }
+	}
+      }
+      dd=(__real__ vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N-__real__ out[KLOOKUP(kx,ky,params.S)])*(__real__ vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N-__real__ out[KLOOKUP(kx,ky,params.S)])+(__imag__ vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N-__imag__ out[KLOOKUP(kx,ky,params.S)])*(__imag__ vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N-__imag__ out[KLOOKUP(kx,ky,params.S)]);
+      if(dd>1e-25){
+	printf("%d %d % .8e\n",kx,ky, dd);
+      }
+    }
+  }
+  */
+
+
+  // write out
+  for(kx=0; kx<params.S*params.S; kx++){
+    out[kx]=0;
+  }
+  for(kx=-params.K;kx<=params.K;kx++){
+    for(ky=-params.K;ky<=params.K;ky++){
+      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)]+vects.invfft[KLOOKUP(kx,ky,params.N)]/params.N/params.N;
+    }
+  }
+
+#endif
+  return(0);
+}
+
+
+// compute alpha
+_Complex double compute_alpha(_Complex double* u, ns_params params){
+  _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));
+    }
+  }
+
+  return(num/denom);
+}
diff --git a/src/navier-stokes.h b/src/navier-stokes.h
new file mode 100644
index 0000000..cd093d7
--- /dev/null
+++ b/src/navier-stokes.h
@@ -0,0 +1,48 @@
+#ifndef NAVIERSTOKES_H
+#define NAVIERSTOKES_H
+
+#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;
+
+// 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);
+
+// right side of Irreversible Navier-Stokes equation
+int ins_rhs(_Complex double* out, _Complex double* u, ns_params params, fft_vects vects);
+
+// compute alpha
+_Complex double compute_alpha(_Complex double* u, ns_params params);
+
+#endif
+
-- 
cgit v1.2.3-54-g00ecf