Implementation of alternating direction implicit method

Question

I have written a program that implements the ADI method and Crank-Nicolson method for solving Schrodinger equations.

The program is working, but it takes a very long time to run. I am looking for tips on how to improve the performance of the program.

#define _USE_MATH_DEFINES
#include <cmath>
#include <vector>
#include <string.h>
#include <array>
#include <iostream>
#include <complex>
#include <stdlib.h>
#include <fstream>
#include <stdio.h>
#include <string>
#include <math.h>
#include <iomanip> 

#define PI M_PI
#define h_bar 1.0
#define M 1.0 
#define L .5
#define sig .05
#define grid_point 100
#define K1 1.0E6
#define K2 1.0E6
#define omega1 std::sqrt(K1/M)
#define omega2 std::sqrt(K2/M)


typedef std::complex<double> compx;
#define I compx(0.,1.)
#define one compx(1.,0.)
#define two compx(2.,0.)

class wave_function {
  public:
    wave_function();
    wave_function(bool);
    double dt = 4.2E-06;
    double dx = L / grid_point;
    std::vector<std::vector<compx>> value;
    void solve_triag();
    double potential(int, int);
    double real_space(int);
    double density_x1[grid_point];
    double density_x2[grid_point];
    compx sec_space_deriv(char, int, int);
    void normalize();
    void rho();
    void alpha_beta_solver(wave_function &, compx, std::vector<compx> &mid,    compx, std::vector<compx> &R, int, char);
    compx a = -h_bar / (two*M), b = -h_bar / (two*M);
    double r = dt / (dx*dx);
    compx A = (I*r*a / two), B = (I*r*b / two), C = I*dt / two;
};

wave_function::wave_function() {
  value.resize(grid_point);
  for (int l = 0; l < grid_point; l++) {
    value[l].resize(grid_point);
    for (int m = 0; m < grid_point; m++) {
        value[l][m] = compx(0., 0.);
    }
  }
}

wave_function::wave_function(bool init) {
  if (init) {
    compx psi00, psi10, psi12;
    value.resize(grid_point);
    for (int l = 0; l < grid_point; l++) {
        value[l].resize(grid_point);
        for (int m = 0; m < grid_point; m++) {
            /*value[l][m] = exp(I*compx(k1, 0)*compx(real_space(l), 0))*exp(-pow(real_space(l) - real_space(x_01), 2.) / (4.*sig*sig))
            *exp(I*compx(k2, 0)*compx(real_space(m), 0))*exp(-pow(real_space(m) - real_space(x_02), 2.) / (4.*sig*sig));*/
            psi00 = pow(M*omega1 / (PI*h_bar), .25)*exp(-M*omega1*pow(real_space(l), 2.) / (2.*h_bar))*pow(M*omega2 / (PI*h_bar), .25)*exp(-M*omega2*pow(real_space(m), 2.) / (2.*h_bar));
            psi10 = pow(M*omega1 / (PI*h_bar), .25)*pow(2., .5)*pow(M*omega1 / (h_bar), .5)*real_space(l)*exp(-M*omega1*pow(real_space(l), 2.) / (2.*h_bar))*pow(M*omega2 / (PI*h_bar), .25)*exp(-M*omega2*pow(real_space(m), 2.) / (2.*h_bar));
            value[l][m] = psi00+psi10;

        }
    }
    normalize();
  }
  else if (!init) {
    wave_function();
  }
}

void wave_function::solve_triag() {
  std::vector<compx> mid1, mid2, R1, R2;
  wave_function tmp;

  mid1.resize(grid_point);
  R1.resize(grid_point);
  for (int x2 = 0; x2 < grid_point; x2++) {
    for (int i = 0; i < grid_point; i++) {
        mid1[i] = one - two*A;
        R1[i] = (one - C*potential(i, x2))*value[i][x2] - B*sec_space_deriv('y', i, x2);
    }
    for (int x1 = 0; x1 < grid_point; x1++) {
        std::cout << x2 << " " << x1 << "  calling solver" << std::endl;
        alpha_beta_solver(tmp, A, mid1, A, R1, x2, 'x');
    }
  }
  tmp.normalize();

  mid2.resize(grid_point);
  R2.resize(grid_point);
  for (int x1 = 0; x1 < grid_point; x1++) {
    for (int i = 0; i < grid_point; i++) {
        mid2[i] = one - two*B + C*potential(x1, i);
        R2[i] = tmp.value[x1][i] - A*tmp.sec_space_deriv('x', x1, i);
    }
    for (int x2 = 0; x2 < grid_point; x2++) {
        alpha_beta_solver(*this, B, mid2, B, R2, x1, 'y');
    }
  }
  normalize();
}



void wave_function::alpha_beta_solver(wave_function &New, compx b4_mid, std::vector<compx> &mid, compx post_mid, std::vector<compx> &R, int coor, char x_or_y) {

  compx x_N, R_N;// new_mid[grid_point], new_R[grid_point];
  std::vector<compx> alpha(grid_point - 2);
  std::vector<compx> beta1(grid_point - 1);//for x_1
  std::vector<compx> beta2(grid_point - 1);//for x_2
  std::vector<compx> x_1(grid_point), x_2(grid_point);


  alpha[0] = post_mid / mid[0];
  beta1[0] = R[0] / mid[0];
  beta2[0] = -b4_mid / mid[0];

  //Forward run
  for (int l = 1; l < grid_point - 2; l++) {
    alpha[l] = post_mid / (mid[l] - b4_mid*alpha[l - 1]);
    beta1[l] = (R[l] - b4_mid*beta1[l - 1]) / (mid[l] - b4_mid*alpha[l - 1]);
    beta2[l] = (-b4_mid*beta2[l - 1]) / (mid[l] - b4_mid*alpha[l - 1]);
  }

  beta1[grid_point - 2] = (R[grid_point - 2] - b4_mid*beta1[grid_point - 3]) / (mid[grid_point - 2] - b4_mid*alpha[grid_point - 3]);
  beta2[grid_point - 2] = (-post_mid - b4_mid*beta2[grid_point - 3]) / (mid[grid_point - 2] - b4_mid*alpha[grid_point - 3]);

  //Backward run
  x_1[grid_point - 2] = beta1[grid_point - 2];
  x_2[grid_point - 2] = beta2[grid_point - 2];
  for (int l = grid_point - 3; l >= 0; l--) {
    x_1[l] = beta1[l] - alpha[l] * x_1[l + 1];
    x_2[l] = beta2[l] - alpha[l] * x_2[l + 1];
  }

  x_N = (R_N - post_mid*x_1[0] - b4_mid*x_1[grid_point - 2])
    / (mid[grid_point - 2] + post_mid*x_2[0] + b4_mid*x_2[grid_point - 2]);

  if (x_or_y == 'x') {
    value[grid_point - 1][coor] = x_N;
    for (int l = 0; l < grid_point - 1; l++) {
        value[l][coor] = x_1[l] + x_2[l] * x_N;
        //std::cout << x1 << " " << x2 << " " << l << " " << x_1[l] + x_2[l] * x_N << std::endl;
    }
  }
  else if (x_or_y == 'y') {
    value[coor][grid_point - 1] = x_N;
    for (int l = 0; l < grid_point - 1; l++) {
        value[coor][l] = x_1[l] + x_2[l] * x_N;
    }
  }

}

double wave_function::potential(int x1, int x2) {

  return .5*M*omega1*omega1*(real_space(x1))*(real_space(x1)) + .5*M*omega1*omega1*(real_space(x2))*(real_space(x2)) + 0.0*.5*M*omega2*omega2*pow(real_space(x2) - real_space(x1), 2.);
//square well
/*if (al - abs(real_space(x1) - real_space(x2)) > 0.)
return V_0;
else
return 0.0;*/

//gaussian
//return V_0*exp(-pow(abs(real_space(x1)-real_space(x2)),2.)/(2.*al*al));
}

double wave_function::real_space(int index) {
  return (-L / 2.0) + (index*dx);
}


//x stands for x1 or particle 1 and y stands for x2 or particle 2
compx wave_function::sec_space_deriv(char x1_or_x2, int l, int m) {
  if (x1_or_x2 == 'x') {
    if (l == 0)
        return value[l + 1][m] - two * value[l][m] + value[grid_point - 1][m];
    else if (l == grid_point - 1)
        return value[0][m] - two * value[l][m] + value[l - 1][m];
    else
        return value[l + 1][m] - two * value[l][m] + value[l - 1][m];
  }
  else if (x1_or_x2 == 'y') {
    if (m == 0)
        return value[l][m + 1] - two * value[l][m] + value[l][grid_point - 1];
    else if (m == grid_point - 1)
        return value[l][0] - two * value[l][m] + value[l][m - 1];
    else
        return value[l][m + 1] - two * value[l][m] + value[l][m - 1];
  }
}


template <typename NEW>
std::string to_string_with_precision(const NEW a_value, const int n = 6)
{
  std::ostringstream out;
  out << std::fixed << std::setprecision(n) << a_value;
  return out.str();
}


//normalization for the function above rho
void wave_function::normalize() {
  compx sum = 0;
  for (int i1 = 0; i1 < grid_point; i1++) {
    for (int i2 = 0; i2 < grid_point; i2++) {
        sum += pow(abs(value[i1][i2]), 2.0);
    }
  }
  sum *= dx * dx;
  compx amplitude = one / pow(sum, .5);
  for (int i1 = 0; i1 < grid_point; i1++) {
    for (int i2 = 0; i2 < grid_point; i2++) {
        value[i1][i2] *= amplitude;
    }
  }
}

void wave_function::rho() {
  for (int k = 0; k < grid_point; k++) {
    density_x1[k] = 0.;
    density_x2[k] = 0.;
  }

  for (int i = 0; i < grid_point; i++) {
    for (int j = 0; j < grid_point; j++) {
        density_x1[i] += pow(abs(value[i][j]), 2.)*dx;
        density_x2[i] += pow(abs(value[j][i]), 2.)*dx;
    }
  }
}

int main() {
  double check = 0.0;
  wave_function v(true);
  std::ofstream file1; //density of both particles
  std::ofstream file2;
  std::ofstream file5; //potential
  file5.open("potential.dat");
  for (int i = 0; i < grid_point; i++)
    for (int j = 0; j < grid_point; j++)
        file5 << v.real_space(i) << "\t" << v.real_space(j) << "\t" << v.potential(i, j) << std::endl;
  int index = 0;
  for (double k = v.dt; k <= 200 * v.dt; k += v.dt) {
    std::cout << index << std::endl;
    v.rho();
    if (index % 10 == 0) {
        file1.open("data_" + to_string_with_precision(index, 0) + ".dat");
        file1 << "Time   " << k - v.dt << std::endl
            << "x" << "\t" << "y" << "\t" << "imag" << "\t" << "real" << "\t" << "abs" << std::endl;
        file2.open("x_" + std::to_string(index) + ".dat");
        file2 << "Time   " << k - v.dt << std::endl
            << "Coordinate" << "\t" << "x1" << "\t" << "x2" << std::endl;
        for (int i = 0; i < grid_point; i++) {
            for (int j = 0; j < grid_point; j++) {
                file1 << v.real_space(i) << "\t"
                    << v.real_space(j) << "\t"
                    << imag(v.value[i][j]) << "\t"
                    << real(v.value[i][j]) << "\t"
                    << abs(v.value[i][j]) << std::endl;
            }
        }
        for (int k = 0; k < grid_point; k++) {
            file2 << v.real_space(k) << "\t"
                << abs(v.density_x1[k]) << "\t"
                << abs(v.density_x2[k]) << std::endl;
        }
        file2.close();
        file1.close();

    }
    std::cout << "About to call solve_triag()  " << std::endl;
    v.solve_triag();
    index++;
  }
  getchar();
}

Answer 1

I see a number of things that may help you improve your program.

Prefer constexpr to old-style #define

Rather than using a #define for all of the constants, the code could use a constexpr:

constexpr auto K2{1.0E6};
constexpr auto omega1{std::sqrt(K1/M)};

constexpr compx I{0.0, 1.0};
constexpr compx one{1.0,0.0};
constexpr compx two{2.0,0.0};

Generally the advantage is that the value has an associated type. Depending on compiler optimizations, it may also confer a speed advantage because you have explicitly told the compiler that the value is a constant known at compile time.

Fix your constructor

It appears that you want to have two kinds of contructors; one simply creates the data structures and the other does a more thorough initialization. Instead of providing two different constructors, I'd suggest instead that you create a single constructor that has a default value for the constant.

Eliminate unused variables

The variable check in your code is defined but never used. Similarly, the parameter New in alpha_beta_solver() is never used. Since unused variables are a sign of poor code quality, you should seek to eliminate them. Your compiler is probably smart enough to warn you about such things if you know how to ask it to do so.

Ensure every control path returns a proper value

The sec_space_deriv() function returns one calculated value if the first parameter was 'x', another if it was 'y' but then doesn't return anything at all otherwise. This is an error. The code could instead be written like this:

compx wave_function::sec_space_deriv(bool want_x1, int l, int m) const {
  if (want_x1)
    if (l == 0)
        return value[l + 1][m] - two * value[l][m] + value[grid_point - 1][m];
    else if (l == grid_point - 1)
        return value[0][m] - two * value[l][m] + value[l - 1][m];
    else
        return value[l + 1][m] - two * value[l][m] + value[l - 1][m];
  }
  if (m == 0)
      return value[l][m + 1] - two * value[l][m] + value[l][grid_point - 1];
  else if (m == grid_point - 1)
      return value[l][0] - two * value[l][m] + value[l][m - 1];
  return value[l][m + 1] - two * value[l][m] + value[l][m - 1];
}

Note that I've added the keyword const to that which brings me to the next suggestion.

Use const where practical

For a number of the operations, the wave_function data are unaltered. That's good design, but you should make it explicit by marking those member functions const. In particular, potential(), real_space() and sec_space_deriv() are all const functions. Similarly, the first and third arguments to alpha_beta_solver should be const references.

Prefer std::array to raw arrays

At the moment, the code for rho() begins like this:

for (int k = 0; k < grid_point; k++) {
  density_x1[k] = 0.;
  density_x2[k] = 0.;
}

If instead, you were to define density_x1 and density_x2 like this:

std::array<double, grid_point> density_x1;
std::array<double, grid_point> density_x2;

That part of rho() would look like this:

density_x1.fill(0);
density_x2.fill(0);

Use <cmath> instead of <math.h>

The difference between the two forms is that the former defines things within the std:: namespace versus into the global namespace. Language lawyers have lots of fun with this, but for daily use I'd recommend using <cmath>. See this SO question for details. There is never any need to include both.

Don't use std::endl if you don't really need it

The difference betweeen std::endl and '\n' is that '\n' just emits a newline character, while std::endl actually flushes the stream. This can be time-consuming in a program with a lot of I/O and is rarely actually needed. It's best to only use std::endl when you have some good reason to flush the stream and it's not very often needed for simple programs such as this one. Avoiding the habit of using std::endl when '\n' will do will pay dividends in the future as you write more complex programs with more I/O and where performance needs to be maximized.

Profile code performance

I profiled your code using gprof and noticed that most of the time (98.84%) was being spent in alpha_beta_solver(). Looking at the code, that function does two things: it initializes a number of local vectors and then applies that to the value stored in the wave_function class. However, all of the setup is actually constant, so we can split the function into two pieces. The first creates x_N, x_1 and x_2 and the second simply uses those to update values. Here's how I rewrote them:

compx wave_function::alpha_beta_solver(const compx &b4_mid, const std::array<compx, grid_point> &mid, const compx &post_mid, const std::array<compx, grid_point> &R, std::vector<compx> &x_1, std::vector<compx> &x_2) {

  compx R_N;// new_mid[grid_point], new_R[grid_point];
  std::vector<compx> alpha(grid_point - 2);
  std::vector<compx> beta1(grid_point - 1);//for x_1
  std::vector<compx> beta2(grid_point - 1);//for x_2

  alpha[0] = post_mid / mid[0];
  beta1[0] = R[0] / mid[0];
  beta2[0] = -b4_mid / mid[0];

  //Forward run
  for (int l = 1; l < grid_point - 2; l++) {
    alpha[l] = post_mid / (mid[l] - b4_mid*alpha[l - 1]);
    beta1[l] = (R[l] - b4_mid*beta1[l - 1]) / (mid[l] - b4_mid*alpha[l - 1]);
    beta2[l] = (-b4_mid*beta2[l - 1]) / (mid[l] - b4_mid*alpha[l - 1]);
  }

  beta1[grid_point - 2] = (R[grid_point - 2] - b4_mid*beta1[grid_point - 3]) / (mid[grid_point - 2] - b4_mid*alpha[grid_point - 3]);
  beta2[grid_point - 2] = (-post_mid - b4_mid*beta2[grid_point - 3]) / (mid[grid_point - 2] - b4_mid*alpha[grid_point - 3]);

  //Backward run
  x_1[grid_point - 2] = beta1[grid_point - 2];
  x_2[grid_point - 2] = beta2[grid_point - 2];
  for (int l = grid_point - 3; l >= 0; l--) {
    x_1[l] = beta1[l] - alpha[l] * x_1[l + 1];
    x_2[l] = beta2[l] - alpha[l] * x_2[l + 1];
  }

  return (R_N - post_mid*x_1[0] - b4_mid*x_1[grid_point - 2])
    / (mid[grid_point - 2] + post_mid*x_2[0] + b4_mid*x_2[grid_point - 2]);
}

void wave_function::update_value(int coor, bool want_x, const compx &x_N, const std::vector<compx> &x_1, const std::vector<compx> &x_2) {
  if (want_x) {
    value[grid_point - 1][coor] = x_N;
    for (int l = 0; l < grid_point - 1; l++) {
        value[l][coor] = x_1[l] + x_2[l] * x_N;
        //std::cout << x1 << " " << x2 << " " << l << " " << x_1[l] + x_2[l] * x_N << std::endl;
    }
  }
  else {
    value[coor][grid_point - 1] = x_N;
    for (int l = 0; l < grid_point - 1; l++) {
        value[coor][l] = x_1[l] + x_2[l] * x_N;
    }
  }

}

And here is the calling function, reword to use std::array and to take advantage of these split functions:

void wave_function::solve_triag() {
  std::array<compx, grid_point> mid1, mid2, R1, R2;
  std::vector<compx> x_1(grid_point), x_2(grid_point);

  mid1.fill(one - two*A);
  for (int x2 = 0; x2 < grid_point; x2++) {
    for (int i = 0; i < grid_point; i++) {
        R1[i] = (one - C*potential(i, x2))*value[i][x2] - B*sec_space_deriv(false, i, x2);
    }
    compx x_N = alpha_beta_solver(A, mid1, A, R1, x_1, x_2);
    for (int x1 = 0; x1 < grid_point; x1++) {
        update_value(x2, true, x_N, x_1, x_2);
    }
  }

  wave_function tmp;

  for (int x1 = 0; x1 < grid_point; x1++) {
    for (int i = 0; i < grid_point; i++) {
        mid2[i] = one - two*B + C*potential(x1, i);
        R2[i] = tmp.value[x1][i] - A*tmp.sec_space_deriv(true, x1, i);
    }
    compx x_N = alpha_beta_solver(B, mid2, B, R2, x_1, x_2);
    for (int x2 = 0; x2 < grid_point; x2++) {
        update_value(x1, false, x_N, x_1, x_2);
    }
  }
  normalize();
}

There are other optimizations possible here, but I wanted to illustrate the process rather than doing it all for you.

Results

The original version takes 16.48 seconds and the version with the modifications mentioned above takes 2.82 seconds on my machine, so over 5 times faster. I'm compiling and running on a 64-bit Linux machine. Here's the Makefile I created. The work directory has the original.cpp and schroedinger.cpp (the modified version). I then created directories named test and golden. I type make golden to create a reference version of all outputs in the golden directory and then make to test the new version and compare it to the original. Also, I am using my usual environment flags, so on my machine CXXFLAGS="-Wall -Wextra -pedantic -std=c++14". The profile.txt target uses gprof to generate a profile of the code after it has run.

Makefile

.PHONY:all
all: result.diff

profile.txt : result.diff
    mv test/gmon.out .
    gprof schroedinger >$@

schroedinger: schroedinger.cpp
    g++ -g -pg -O2 $(CXXFLAGS) $^ -o $@

original: original.cpp
    g++ -g -pg -O2 $(CXXFLAGS) $^ -o $@

.PHONY:golden
golden: original
    cd golden && time ../original >/dev/null

result.diff: schroedinger
    cd test && time ../schroedinger >/dev/null
    diff -Naur --exclude *.out golden test >$@

Answer 2

First, you should fix these compiler warnings:

157094.cpp: In constructor ‘wave_function::wave_function()’:
157094.cpp:53:1: warning: ‘wave_function::value’ should be initialized in the member initialization list [-Weffc++]
 wave_function::wave_function() {
 ^~~~~~~~~~~~~
157094.cpp: In constructor ‘wave_function::wave_function(bool)’:
157094.cpp:63:1: warning: ‘wave_function::value’ should be initialized in the member initialization list [-Weffc++]
 wave_function::wave_function(bool init) {
 ^~~~~~~~~~~~~
157094.cpp: In member function ‘void wave_function::alpha_beta_solver(wave_function&, compx, std::vector<std::complex<double> >&, compx, std::vector<std::complex<double> >&, int, char)’:
157094.cpp:119:54: warning: unused parameter ‘New’ [-Wunused-parameter]
 void wave_function::alpha_beta_solver(wave_function &New, compx b4_mid, std::vector<compx> &mid, compx post_mid, std::vector<compx> &R, int coor, char x_or_y) {
                                                      ^~~
157094.cpp: In function ‘int main()’:
157094.cpp:249:10: warning: unused variable ‘check’ [-Wunused-variable]
   double check = 0.0;
          ^~~~~
157094.cpp: In member function ‘compx wave_function::sec_space_deriv(char, int, int)’:
157094.cpp:205:1: warning: control reaches end of non-void function [-Wreturn-type]
 }
 ^

---

I'd advise that you use real, strongly-typed constants in preference to preprocessor substitution:

const double hbar{1};
const double L{0.5};
// etc.

const compx I{0,1};
// No need for complex values of 1 or 2, as doubles will promote okay

---

The following code doesn't do what you think it does:

wave_function::wave_function(bool init) {
  if (init) {
    // (snip)
  }
  else if (!init) {
    wave_function();
  }
}

In the else branch (with a redundant condition, but your compiler can probably eliminate that), you construct a new wave_function object and then discard it. wave_function(); here will not invoke the default constructor.

I think you want to turn this the other way around:

wave_function::wave_function()
    : wave_function(false)
{}

wave_function::wave_function(bool init) {
  if (init) {
      // "true" setup
  }
  else {
      // "false" setup
  }
}

---

The population of the initial values is hard to read, and repeats a lot of calculations. These can be avoided with the judicious use of temporary variables for the intermediate results. I believe the following is equivalent to your initialization loop:

    for (int l = 0; l < grid_point; l++) {
        value[l].resize(grid_point);
        for (int m = 0; m < grid_point; m++) {
            auto const rl = real_space(l);
            auto const rm = real_space(m);

            auto const k1 = M * omega1 / PI / h_bar;
            auto const k2 = M * omega2 / PI / h_bar;

            auto const e1 = std::exp(-M*omega1*rl*rl/h_bar/2);
            auto const e2 = std::exp(-M*omega2*rm*rm/h_bar/2);

            auto const psi00 = std::pow(k1*k2, .25) * e1 * e2;
            auto const psi10 = psi00 * std::sqrt(2*M*omega1/h_bar) * rl;
            value[l][m] = psi00+psi10;
        }
    }

I've used the identity pow(x,n)*pow(y,n) == pow(x*y,n) in a couple of places above, to reduce the number of calls.

---

Throughout your code, there are uses of std::pow with constant exponents. In general, you can expect x*x to be faster than std::pow(x, 2) and std::sqrt(x) to be faster than std::pow(x, 0.5) - and both are likely to be more accurate into the bargain.

---

(This is an incomplete answer - Real Life called me away. I was intending to return with additional review, but Edward has written an answer much better than this one was ever going to be.)