Compute energy, enstrophy and alpha all at once

3 files changed, 60 insertions, 21 deletions

diff --git a/src/main.c b/src/main.c
index 9999de5..60698ef 100644
--- a/src/main.c
+++ b/src/main.c
@@ -97,7 +97,7 @@ int main (
     energy(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.print_freq, nthreads);
   }
   else if(command==COMMAND_ALPHA){
-    alpha(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.print_freq, nthreads);
+    eea(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.print_freq, nthreads);
   }
   else if(command==COMMAND_QUIET){
     quiet(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, nthreads);
diff --git a/src/navier-stokes.c b/src/navier-stokes.c
index 531398d..588dcf4 100644
--- a/src/navier-stokes.c
+++ b/src/navier-stokes.c
@@ -92,8 +92,8 @@ int energy(
   fft_vect fft1;
   fft_vect fft2;
   fft_vect ifft;
-  int kx,ky;
   double energy;
+  double avg;
 
   ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads);
   // copy initial condition
@@ -107,12 +107,7 @@ int energy(
     ins_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
     
     if(t%print_freq==0){
-      energy=0.;
-      for(kx=-K1;kx<=K1;kx++){
-	for (ky=-K2;ky<=K2;ky++){
-	  energy+=__real__ (u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(u[klookup(kx,ky,2*K1+1,2*K2+1)]));
-	}
-      }
+      energy=compute_energy(u, K1, K2);
 
       // running average
       if(t>0){
@@ -128,8 +123,8 @@ int energy(
   return(0);
 }
 
-// compute alpha as a function of time in the I-NS equation
-int alpha(
+// compute energy, enstrophy, alpha as a function of time in the I-NS equation
+int eea(
   int K1,
   int K2,
   int N1,
@@ -147,8 +142,8 @@ int alpha(
   _Complex double* tmp1;
   _Complex double* tmp2;
   _Complex double* tmp3;
-  _Complex double alpha;
-  _Complex double avg;
+  double alpha, energy, enstrophy;
+  double avg_e,avg_a,avg_en;
   unsigned int t;
   fft_vect fft1;
   fft_vect fft2;
@@ -160,21 +155,28 @@ int alpha(
 
 
   // init running average
-  avg=0;
+  avg_e=0;
+  avg_a=0;
+  avg_en=0;
 
   // iterate
   for(t=0;t<nsteps;t++){
     ins_step(u, K1, K2, N1, N2, nu, delta, L, g, fft1, fft2, ifft, tmp1, tmp2, tmp3);
+
+    energy=compute_energy(u, K1, K2);
     alpha=compute_alpha(u, K1, K2, g);
+    enstrophy=compute_enstrophy(u, K1, K2, L);
     
     // running average
     if(t>0){
-      avg=avg-(avg-alpha)/t;
+      avg_e=avg_e-(avg_e-energy)/t;
+      avg_a=avg_a-(avg_a-alpha)/t;
+      avg_en=avg_en-(avg_en-enstrophy)/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);
+      fprintf(stderr,"%d % .8e % .8e % .8e % .8e % .8e % .8e % .8e\n",t,t*delta, avg_a, avg_e, avg_en, alpha, energy, enstrophy);
+      printf("%8d % .15e % .15e % .15e % .15e % .15e % .15e % .15e\n",t,t*delta, avg_a, avg_e, avg_en, alpha, energy, enstrophy);
     }
   }
 
@@ -539,7 +541,7 @@ int ns_T_nofft(
 }
 
 // compute alpha
-_Complex double compute_alpha(
+double compute_alpha(
   _Complex double* u,
   int K1,
   int K2,
@@ -556,7 +558,40 @@ _Complex double compute_alpha(
     }
   }
 
-  return(num/denom);
+  return __real__ (num/denom);
+}
+
+// compute energy
+double compute_energy(
+  _Complex double* u,
+  int K1,
+  int K2
+){
+  int kx,ky;
+  double out=0.;
+  for(kx=-K1;kx<=K1;kx++){
+    for (ky=-K2;ky<=K2;ky++){
+      out+=__real__ (u[klookup(kx,ky,2*K1+1,2*K2+1)]*conj(u[klookup(kx,ky,2*K1+1,2*K2+1)]));
+    }
+  }
+  return out;
+}
+
+// compute enstrophy
+double compute_enstrophy(
+  _Complex double* u,
+  int K1,
+  int K2,
+  double L
+){
+  int kx,ky;
+  double out=0.;
+  for(kx=-K1;kx<=K1;kx++){
+    for (ky=-K2;ky<=K2;ky++){
+      out+=__real__ (4*M_PI*M_PI/L/L*(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)]));
+    }
+  }
+  return out;
 }
 
 
diff --git a/src/navier-stokes.h b/src/navier-stokes.h
index 07b7550..7841f8f 100644
--- a/src/navier-stokes.h
+++ b/src/navier-stokes.h
@@ -18,8 +18,8 @@ int uk( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double d
 // compute the energy as a function of time
 int energy( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, unsigned int print_freq, unsigned int nthreads);
 
-// compute alpha
-int alpha( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, unsigned int print_freq, unsigned int nthreads);
+// compute energy, enstrophy and alpha
+int eea( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, unsigned int print_freq, unsigned int nthreads);
 
 // compute solution as a function of time, but do not print anything (useful for debugging)
 int quiet( int K1, int K2, int N1, int N2, unsigned int nsteps, double nu, double delta, double L, _Complex double* u0, _Complex double* g, unsigned int nthreads);
@@ -46,7 +46,11 @@ int ns_T( _Complex double* u, int K1, int K2, int N1, int N2, fft_vect fft1, fft
 int ns_T_nofft( _Complex double* out, _Complex double* u, int K1, int K2, int N1, int N2);
 
 // compute alpha
-_Complex double compute_alpha( _Complex double* u, int K1, int K2, _Complex double* g);
+double compute_alpha( _Complex double* u, int K1, int K2, _Complex double* g);
+// compute energy
+double compute_energy( _Complex double* u, int K1, int K2);
+// compute enstrophy
+double compute_enstrophy( _Complex double* u, int K1, int K2, double L);
 
 // get index for kx,ky in array of size S
 int klookup( int kx, int ky, int S1, int S2);