1 #include <iostream>
  2 #include <vector>
  3 #include <cmath>  // for sqrt and pow
  4 #include <fstream>
  5 #include "typedefs.h"
  6 #include "ode.h"
  7 #include "vector_arithmetic.h"
  8 
  9 // Constructor
 10 ODE::ODE(std::vector<std::complex<Floattype> >
 11                 (*f_in)(const std::complex<Floattype> x,
 12                         const std::vector<std::complex<Floattype> > &y),
 13          const std::vector<std::complex<Floattype> > y_start,
 14          const std::complex<Floattype> a_in,
 15          const std::complex<Floattype> b_in,
 16          const Floattype h_start,
 17          const Inttype max_steps,
 18          const Floattype delta_in,
 19          const Floattype epsilon_in,
 20          const Floattype power_in,
 21          const Floattype safety_in)
 22   : f(f_in),
 23     a(a_in), b(b_in),
 24     h((b_in-a_in)*h_start),
 25     n_max(max_steps),
 26     delta(delta_in), epsilon(epsilon_in),
 27     power(power_in), safety(safety_in)
 28 {
 29   x_list.push_back(a);
 30   y_list.push_back(y_start);
 31 
 32   // Perform integration
 33   rkdriver();
 34 }
 35 
 36 
 37 // Estimate tolerance
 38 Floattype ODE::tau(const std::vector<std::complex<Floattype> > &y,
 39                    const std::complex<Floattype> h_i)
 40 {
 41   return abs((epsilon * cnorm(y) + delta) * sqrt(h_i/(b-a)));
 42 }
 43 
 44 // Runge-Kutta mid-point stepper
 45 void ODE::rkstep12(const std::complex<Floattype> x0,
 46                    const std::vector<std::complex<Floattype> > &y0,
 47                    std::vector<std::complex<Floattype> > &y1,
 48                    std::vector<std::complex<Floattype> > &dy)
 49 {
 50   // Denominator 2 used in arithmetic operations
 51   const std::complex<Floattype> den2 (2.0f,2.0f);
 52 
 53   // Evaluate function at two points
 54   (void)f(x0,y0);
 55   const std::vector<std::complex<Floattype> > k0  = f(x0,y0);
 56   const std::vector<std::complex<Floattype> > k12 = f(x0 + h/den2, y0 + k0*h/den2);
 57 
 58   // Write results to output vectors
 59   y1 = y0 + k12*h;
 60   dy = (k0 - k12) * h/den2;
 61 }
 62 
 63 
 64 // ODE driver with adaptive step size control
 65 void ODE::rkdriver()
 66 {
 67   std::vector<std::complex<Floattype> > dy(n_max);
 68   std::vector<std::complex<Floattype> > y1(n_max);
 69 
 70   std::complex<Floattype> x;
 71   Floattype err, tol;
 72   std::vector<std::complex<Floattype> > y;
 73 
 74   while (x_list.back().real() < b.real()
 75       || x_list.back().imag() < b.imag())
 76   {
 77     // Get values for this step
 78     x = x_list.back();
 79     y = y_list.back();
 80 
 81     // Make sure we don't step past the upper boundary
 82     if ((x + h).real() > b.real()
 83      || (x + h).imag() > b.imag())
 84       h = b - x;
 85 
 86     // Run Runge-Kutta mid-point stepper
 87     rkstep12(x, y, y1, dy);
 88 
 89     // Determine whether the step should be accepted
 90     err = cnorm(dy); // Error estimate
 91     tol = tau(y, h); // Tolerance
 92     if (err < tol) { // Step accepted
 93       x_list.push_back(x+h);
 94       y_list.push_back(y1);
 95     }
 96 
 97     // Check that we havn't hit the max. number of steps
 98     if (x_list.size() == n_max) {
 99       std::cerr << "Error, the max. number of steps "
100                 << "was insufficient\n"
101                 << "Try either increasing the max. number "
102                 << "of steps, or decreasing the precision "
103                 << "requirements\n";
104       return;
105     }
106 
107     // Determine new step size
108     std::complex<Floattype> multiplicator (2.0f, 2.0f);
109     if (err > 0.0f)
110       h = h*pow(tol/err, power) * safety;
111     else
112       h = multiplicator*h;
113   }
114 }
115 
116 
117 // Return the number of steps taken
118 Inttype ODE::steps()
119 {
120   return x_list.size();
121 }
122 
123 void ODE::print()
124 {
125   for (Inttype i=0; i<x_list.size(); ++i) {
126     std::cout << x_list[i] << '\t';
127     for (Inttype j=0; j<y_list[0].size(); ++j)
128       std::cout << y_list[i][j] << '\t';
129     std::cout << '\n';
130   }
131 }
132 
133 // Write the x- and y-values to file
134 void ODE::write(const char* filename)
135 {
136   std::ofstream outstream;
137 
138   // Open outfile for write
139   outstream.open(filename);
140   if (!outstream) {
141     std::cerr << "Error, can't open output file '"
142               << filename << "'.\n";
143     return;
144   }
145 
146   // Write output values
147   for (Inttype i=0; i<x_list.size(); ++i) {
148     outstream << x_list[i].real() << '\t';
149     outstream << x_list[i].imag() << '\t';
150     for (Inttype j=0; j<y_list[0].size(); ++j) {
151       outstream << y_list[i][j].real() << '\t';
152       outstream << y_list[i][j].imag() << '\t';
153     }
154     outstream << '\n';
155   }
156 
157   // Close file
158   outstream.close();
159 
160   if (verbose == true)
161     std::cout << "Output written in '" << filename << "'.\n";
162 }
163