Compute energy, enstrophy and alpha all at once

1 files changed, 52 insertions, 17 deletions

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