Rewrite: change cli arguments handling

1 files changed, 340 insertions, 68 deletions

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);
+}