Driving force as vector instead of function

1 files changed, 27 insertions, 14 deletions

diff --git a/src/main.c b/src/main.c
index 27ed314..f074720 100644
--- a/src/main.c
+++ b/src/main.c
@@ -32,6 +32,8 @@ int read_args(int argc, const char* argv[], char** params, unsigned int* driving
 int read_params(char* param_str, nstrophy_parameters* parameters);
 int set_parameter(char* lhs, char* rhs, nstrophy_parameters* parameters, bool* setN1, bool* setN2);
 
+// set driving force
+_Complex double* set_driving(unsigned int driving, nstrophy_parameters parameters);
 // set initial condition
 _Complex double* set_init(unsigned int init, nstrophy_parameters parameters);
 
@@ -53,11 +55,11 @@ int main (
 ){
   char* param_str=NULL;
   nstrophy_parameters parameters;
-  _Complex double (*g)(int,int);
   int ret;
   unsigned int driving,command,init;
   unsigned int nthreads=1;
   _Complex double* u0;
+  _Complex double *g;
 
   command=0;
   driving=0;
@@ -76,19 +78,7 @@ int main (
   }
 
   // set driving force
-  switch(driving){
-  case DRIVING_ZERO:
-    g=g_zero;
-    break;
-  case DRIVING_TEST:
-    g=g_test;
-    break;
-
-  default:
-    g=g_zero;
-    break;
-  }
-
+  g=set_driving(driving, parameters);
   // set initial condition
   u0=set_init(init, parameters);
 
@@ -110,6 +100,7 @@ int main (
     print_usage();
   }
 
+  free(g);
   free(u0);
 
   return(0);
@@ -441,6 +432,28 @@ int set_parameter(
   return(0);
 }
 
+// set driving force
+_Complex double* set_driving(
+  unsigned int driving,
+  nstrophy_parameters parameters
+){
+  _Complex double* g=calloc(sizeof(_Complex double),(2*parameters.K1+1)*(2*parameters.K2+1));
+
+  switch(driving){
+  case DRIVING_ZERO:
+    g_zero(g, parameters.K1, parameters.K2);
+    break;
+  case DRIVING_TEST:
+    g_test(g, parameters.K1, parameters.K2);
+    break;
+
+  default:
+    g_test(g, parameters.K1, parameters.K2);
+    break;
+  }
+   return g;
+}
+
 // set initial condition
 _Complex double* set_init(
   unsigned int init,