Compute energy

3 files changed, 60 insertions, 1 deletions

diff --git a/src/main.c b/src/main.c
index 5ffc672..27bfb7a 100644
--- a/src/main.c
+++ b/src/main.c
@@ -21,6 +21,7 @@ int set_parameter(char* lhs, char* rhs, int* K1, int* K2, int* N1, int* N2, unsi
 #define COMMAND_UK 1
 #define COMMAND_ENSTROPHY 2
 #define COMMAND_QUIET 3
+#define COMMAND_ENERGY 4
 
 #define DRIVING_ZERO 1
 #define DRIVING_TEST 2
@@ -72,8 +73,11 @@ int main (
   if (command==COMMAND_UK){
     uk(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq, nthreads);
   }
+  else if (command==COMMAND_ENERGY){
+    energy(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq, nthreads);
+  }
   else if(command==COMMAND_ENSTROPHY){
-    enstrophy(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq, nthreads);
+    energy(K1, K2, N1, N2, nsteps, nu, delta, g, print_freq, nthreads);
   }
   else if(command==COMMAND_QUIET){
     quiet(K1, K2, N1, N2, nsteps, nu, delta, g, nthreads);
@@ -173,6 +177,9 @@ int read_args(
       else if(strcmp(argv[i], "quiet")==0){
 	*command=COMMAND_QUIET;
       }
+      else if(strcmp(argv[i], "energy")==0){
+	*command=COMMAND_ENERGY;
+      }
       else{
 	fprintf(stderr, "error: unrecognized command: '%s'\n",argv[i]);
 	return(-1);
diff --git a/src/navier-stokes.c b/src/navier-stokes.c
index 18c2fc5..0a2842c 100644
--- a/src/navier-stokes.c
+++ b/src/navier-stokes.c
@@ -68,6 +68,55 @@ int uk(
   return(0);
 }
 
+// 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,
+  _Complex double (*g)(int,int),
+  unsigned int print_freq,
+  unsigned int nthreads
+){
+  _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;
+  double energy;
+
+  ns_init_tmps(&u, &tmp1, &tmp2, &tmp3, &fft1, &fft2, &ifft, K1, K2, N1, N2, nthreads);
+  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){
+
+      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)]));
+	}
+      }
+
+      fprintf(stderr,"%d % .8e % .8e\n",t,t*delta, energy);
+      printf("%8d % .15e % .15e\n",t,t*delta,energy);
+    }
+  }
+
+  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,
diff --git a/src/navier-stokes.h b/src/navier-stokes.h
index 707dc52..421342a 100644
--- a/src/navier-stokes.h
+++ b/src/navier-stokes.h
@@ -14,6 +14,9 @@ typedef struct fft_vects {
 // 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, unsigned int nthreads);
 
+// 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, _Complex double (*g)(int,int), unsigned int print_freq, unsigned int nthreads);
+
 // 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, unsigned int nthreads);