Modify class

include/WaveSimulation.hpp

@@ -21,31 +21,96 @@
 class WaveSimulation {
 protected:
     uint32_t M;
+    int32_t N;
+    arma::cx_mat V;
+    arma::cx_mat U;
     arma::sp_cx_mat B;
+    arma::sp_cx_mat A;
     double h;
     double dt;
     double T;
 
-    void build_A();
-    void build_B();
+    /*  @brief Initialize the U matrix using an unormalized Gaussian wave
+     * packet.
+     *
+     *  @param x_c The center of the packet in the x direction.
+     *  @param y_c The center of the packet in the y direction.
+     *  @param sigma_x The The initial width in the x direction.
+     *  @param sigma_y The The initial width in the y direction.
+     *  @param p_x The wave packet momentum in the x direction.
+     *  @param p_y The wave packet momentum in the y direction.
+     * **/
     void initialize_U(double x_c, double y_c, double sigma_x, double sigma_y,
                       double p_x, double p_y);
-    void build_V(double thickness, double pos_x, double aperture_separation,
-                 double aperture, uint32_t slits);
+
+    /*  @brief Initialize the V matrix.
+     *
+     *  @param thickness The thickness of the wall in the x direction.
+     *  @param pos_x The center of the wall in the x direction.
+     *  @param ap_sep The separation between each aperture.
+     *  @param ap The aperture width.
+     *  @param slits The number of slits.
+     * **/
+    void initialize_V(double thickness, double pos_x,
+                      double aperture_separation, double aperture,
+                      uint32_t slits);
+
+    /*  @brief Initialize the V matrix with no wall.
+     * **/
+    void initialize_V();
+
+    /*  @brief Initialize the A matrix according to the Crank-Nicolson method
+     * **/
+    void initialize_A();
+
+    /*  @brief Initialize the B matrix according to the Crank-Nicolson method
+     * **/
+    void initialize_B();
+
 public:
-    int32_t N;
-    arma::cx_mat V;
-    arma::cx_mat U;
-    arma::sp_cx_mat A;
+    /*  @brief Constructor for the WaveSimulation class.
+     *
+     *  @param h The step size in the x and y direction.
+     *  @param dt The step size in the temporal dimension.
+     *  @param T The total time to simulate.
+     *  @param x_c The center of the packet in the x direction.
+     *  @param y_c The center of the packet in the y direction.
+     *  @param sigma_x The The initial width in the x direction.
+     *  @param sigma_y The The initial width in the y direction.
+     *  @param p_x The wave packet momentum in the x direction.
+     *  @param p_y The wave packet momentum in the y direction.
+     *  @param thickness The thickness of the wall in the x direction.
+     *  @param pos_x The center of the wall in the x direction.
+     *  @param ap_sep The separation between each aperture.
+     *  @param ap The aperture width.
+     *  @param slits The number of slits.
+     * **/
     WaveSimulation(double h, double dt, double T, double x_c, double y_c,
                    double sigma_x, double sigma_y, double p_x, double p_y,
                    double thickness, double pos_x, double ap_sep, double ap,
                    uint32_t slits);
+
+    /*  @brief Constructor for the WaveSimulation class with no wall.
+     *
+     *  @param h The step size in the x and y direction.
+     *  @param dt The step size in the temporal dimension.
+     *  @param T The total time to simulate.
+     *  @param x_c The center of the packet in the x direction.
+     *  @param y_c The center of the packet in the y direction.
+     *  @param sigma_x The The initial width in the x direction.
+     *  @param sigma_y The The initial width in the y direction.
+     *  @param p_x The wave packet momentum in the x direction.
+     *  @param p_y The wave packet momentum in the y direction.
+     * **/
     WaveSimulation(double h, double dt, double T, double x_c, double y_c,
                    double sigma_x, double sigma_y, double p_x, double p_y);
-    virtual void solve(std::ofstream &ofile);
-    void write_U(std::ofstream &ofile);
+
     void step();
+    void solve(std::string outfile, bool write_each_step = false);
+    void solve(std::string outfile, std::vector<double> &steps);
+    void probability_deviation(std::string outfile,
+                               bool write_each_step = false);
+    void write_U(std::ofstream &ofile);
 };
 
 #endif

lib/WaveSimulation.cpp

@@ -17,56 +17,113 @@
 #include <complex>
 #include <cstdint>
 #include <cstdlib>
+#include <vector>
 
-WaveSimulation::WaveSimulation(double h, double dt, double T, double x_c,
-                               double y_c, double sigma_x, double sigma_y,
-                               double p_x, double p_y, double thickness,
-                               double pos_x, double ap_sep, double ap,
-                               uint32_t slits)
+// Initializers
+void WaveSimulation::initialize_U(double x_c, double y_c, double sigma_x,
+                                  double sigma_y, double p_x, double p_y)
 {
-    this->dt = dt;
-    this->h = h;
-    this->T = T;
-    this->M = 1. / h;
-    this->N = M - 2;
-    this->build_V(thickness, pos_x, ap_sep, ap, slits);
-    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
-    this->build_A();
-    this->build_B();
+    this->U.set_size(this->N, this->N);
+    double x, y, diff_x, diff_y;
+    std::complex<double> sum = 0.;
+    for (size_t j = 0; j < this->U.n_cols; j++) {
+        x = j * h;
+        diff_x = x - x_c;
+        for (size_t i = 0; i < this->U.n_rows; i++) {
+            y = i * h;
+            diff_y = y - y_c;
+            this->U(i, j) =
+                std::exp(-(diff_x * diff_x) / (2. * sigma_x * sigma_x)
+                         - (diff_y * diff_y) / (2. * sigma_y * sigma_y)
+                         + p_x * x * 1._i + p_y * y * 1._i);
+            sum += this->U(i, j) * std::conj(this->U(i, j));
+        }
+    }
+
+    // The imaginary part of sum should be 0.
+    if (std::abs(sum.imag()) > 1e-7) {
+        abort();
+    }
+
+    double norm = 1. / std::sqrt(sum.real());
+
+    // Normalize each element
+    this->U.for_each([norm](std::complex<double> &el) { el *= norm; });
 }
 
-WaveSimulation::WaveSimulation(double h, double dt, double T, double x_c, double y_c,
-               double sigma_x, double sigma_y, double p_x, double p_y)
+void WaveSimulation::initialize_V()
 {
-    this->dt = dt;
-    this->h = h;
-    this->T = T;
-    this->M = 1. / h;
-    this->N = M - 2;
-    this->build_V(0.,0.,0.,0., 0);
-    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
-    this->build_A();
-    this->build_B();
-
+    this->V.set_size(this->N, this->N);
+    this->V.fill(0.);
 }
 
-void WaveSimulation::solve(std::ofstream &ofile)
+void WaveSimulation::initialize_V(double thickness, double pos_x, double ap_sep,
+                                  double ap, uint32_t slits)
 {
-    ofile << this->N << '\n';
-    uint32_t iterations = this->T / this->dt;
-    for (size_t i = 0; i < iterations; i++) {
-        this->write_U(ofile);
-        this->step();
+    this->V.set_size(this->N, this->N);
+    this->V.fill(0.);
+    if (slits == 0) {
+        return;
+    }
+
+    arma::cx_vec res;
+
+    // Differentiate between odd and even number of slits.
+    uint32_t ap_points, ap_sep_points;
+
+    if (slits % 2 == 0) {
+        ap_points = ap / this->h;
+        ap_sep_points = ap_sep / this->h
+                        + ((uint32_t)(ap_sep / this->h) % 2 != this->N % 2);
+
+        res = arma::cx_vec(ap_sep_points, arma::fill::value(1e10));
+
+        for (size_t i = 0; i < slits; i += 2) {
+            res = arma::join_cols(res,
+                                  arma::cx_vec(ap_points, arma::fill::zeros));
+            res = arma::join_cols(arma::cx_vec(ap_points, arma::fill::zeros),
+                                  res);
+            res = arma::join_cols(
+                res, arma::cx_vec(ap_sep_points, arma::fill::value(1e10)));
+            res = arma::join_cols(
+                arma::cx_vec(ap_sep_points, arma::fill::value(1e10)), res);
+        }
+    }
+    else {
+        ap_points =
+            ap / this->h + ((uint32_t)(ap / this->h) % 2 != this->N % 2);
+        ap_sep_points = ap_sep / this->h;
+
+        res = arma::cx_vec(ap_points, arma::fill::value(0));
+
+        for (size_t i = 0; i < slits - 1; i += 2) {
+            res = arma::join_cols(
+                res, arma::cx_vec(ap_sep_points, arma::fill::value(1e10)));
+            res = arma::join_cols(
+                arma::cx_vec(ap_sep_points, arma::fill::value(1e10)), res);
+            res = arma::join_cols(res,
+                                  arma::cx_vec(ap_points, arma::fill::zeros));
+            res = arma::join_cols(arma::cx_vec(ap_points, arma::fill::zeros),
+                                  res);
+        }
+    }
+
+    if (res.size() > this->N) {
+        abort();
+    }
+
+    uint32_t fill = (this->N - res.size()) / 2;
+    res = arma::join_cols(arma::cx_vec(fill, arma::fill::value(1e10)), res);
+    res = arma::join_cols(res, arma::cx_vec(fill + ((this->N - res.size()) % 2),
+                                            arma::fill::value(1e10)));
+
+    uint32_t start = pos_x / this->h - thickness / this->h / 2;
+    for (size_t i = 0; i < thickness / this->h; i++) {
+        this->V.col(start + i) = res;
     }
 }
 
-void WaveSimulation::step()
-{
-    arma::cx_vec tmp = this->B * this->U.as_col();
-    arma::spsolve(this->U, this->A, tmp);
-}
-
-void WaveSimulation::build_A()
+void WaveSimulation::initialize_A()
 {
     // Create the diagonal
     arma::cx_vec diagonal(this->N * this->N);
@@ -96,7 +153,7 @@ void WaveSimulation::build_A()
     this->A.diag(this->N).fill(-r);
 }
 
-void WaveSimulation::build_B()
+void WaveSimulation::initialize_B()
 {
     std::complex<double> r = (1._i * this->dt) / (2 * h * h);
 
@@ -126,92 +183,149 @@ void WaveSimulation::build_B()
     this->B.diag(this->N).fill(r);
 }
 
-void WaveSimulation::initialize_U(double x_c, double y_c, double sigma_x,
-                                  double sigma_y, double p_x, double p_y)
+// Constructors
+WaveSimulation::WaveSimulation(double h, double dt, double T, double x_c,
+                               double y_c, double sigma_x, double sigma_y,
+                               double p_x, double p_y, double thickness,
+                               double pos_x, double ap_sep, double ap,
+                               uint32_t slits)
 {
-    this->U.set_size(this->N, this->N);
-    double x, y, diff_x, diff_y;
-    std::complex<double> sum = 0.;
-    for (size_t j = 0; j < this->U.n_cols; j++) {
-        x = j * h;
-        diff_x = x - x_c;
-        for (size_t i = 0; i < this->U.n_rows; i++) {
-            y = i * h;
-            diff_y = y - y_c;
-            this->U(i, j) =
-                std::exp(-(diff_x * diff_x) / (2. * sigma_x * sigma_x)
-                         - (diff_y * diff_y) / (2. * sigma_y * sigma_y)
-                         + p_x * x * 1._i + p_y * y * 1._i);
-            sum += this->U(i, j) * std::conj(this->U(i, j));
-        }
-    }
-    if (std::abs(sum.imag()) > 1e-7) {
+    this->dt = dt;
+    this->h = h;
+    this->T = T;
+    this->M = 1. / h;
+    this->N = M - 2;
+    this->initialize_V(thickness, pos_x, ap_sep, ap, slits);
+    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
+    this->initialize_A();
+    this->initialize_B();
+}
+
+WaveSimulation::WaveSimulation(double h, double dt, double T, double x_c,
+                               double y_c, double sigma_x, double sigma_y,
+                               double p_x, double p_y)
+{
+    this->dt = dt;
+    this->h = h;
+    this->T = T;
+    this->M = 1. / h;
+    this->N = M - 2;
+    this->initialize_V();
+    this->initialize_U(x_c, y_c, sigma_x, sigma_y, p_x, p_y);
+    this->initialize_A();
+    this->initialize_B();
+}
+
+// Public methods
+void WaveSimulation::step()
+{
+    // Create the right hand side of Ax = b
+    arma::cx_vec tmp = this->B * this->U.as_col();
+
+    // Solve Ax = b
+    arma::spsolve(this->U, this->A, tmp);
+}
+
+void WaveSimulation::solve(std::string outfile, bool write_each_step)
+{
+    // Create path and proceed if successful.
+    if (!utils::mkpath(utils::dirname(outfile))) {
         abort();
     }
-    double norm = 1. / std::sqrt(sum.real());
 
-    this->U.for_each([norm](std::complex<double> &el) { el *= norm; });
+    // Open file
+    std::ofstream ofile;
+    ofile.open(outfile);
+
+    // Write the size of the matrix on the first line
+    ofile << this->N << '\n';
+
+    uint32_t iterations = this->T / this->dt;
+
+    if (write_each_step) {
+        for (size_t i = 0; i < iterations; i++) {
+            this->write_U(ofile);
+            this->step();
+        }
+    }
+    else {
+        for (size_t i = 0; i < iterations; i++) {
+            this->step();
+        }
+    }
+    this->write_U(ofile);
+    ofile.close();
+}
+void WaveSimulation::solve(std::string outfile, std::vector<double> &steps)
+{
+    // Create path and proceed if successful.
+    if (!utils::mkpath(utils::dirname(outfile))) {
+        abort();
+    }
+
+    // Open file
+    std::ofstream ofile;
+    ofile.open(outfile);
+
+    // Write the size of the matrix on the first line
+    ofile << this->N << '\n';
+    
+    uint32_t iterations;
+
+    for (size_t i=0; i < steps.size(); i++) {
+        if (i == 0) {
+            iterations = steps[i] / this->h;
+        }
+        else {
+            iterations = (steps[i] - steps[i-1]) / this->h;
+        }
+
+        for (size_t j=0; j < iterations; j++) {
+            this->step();
+        }
+        this->write_U(ofile);
+    }
+
+    ofile.close();
+}
+
+void WaveSimulation::probability_deviation(std::string outfile,
+                                           bool write_each_step)
+{
+    // Create path and proceed if successful.
+    if (!utils::mkpath(utils::dirname(outfile))) {
+        abort();
+    }
+
+    // Open file
+    std::ofstream ofile;
+    ofile.open(outfile);
+
+    uint32_t iterations = this->T / this->dt;
+
+    double sum;
+
+    if (write_each_step) {
+        for (size_t i = 0; i < iterations; i++) {
+            sum = arma::accu(this->U % arma::conj(this->U)).real();
+            ofile << i*this->dt << '\t' << 1. - sum << '\n';
+            this->step();
+        }
+    }
+    else {
+        for (size_t i = 0; i < iterations; i++) {
+            this->step();
+        }
+    }
+    sum = arma::accu(this->U % arma::conj(this->U)).real();
+    ofile << this->T << '\t' << 1. - sum << '\n';
+    ofile.close();
 }
 
 void WaveSimulation::write_U(std::ofstream &ofile)
 {
+    // Write each element to file in column-major order.
     this->U.for_each(
         [&ofile](std::complex<double> el) { ofile << el << '\t'; });
     ofile << '\n';
 }
-
-void WaveSimulation::build_V(double thickness, double pos_x,
-                             double aperture_separation, double aperture,
-                             uint32_t slits)
-{
-    this->V.set_size(this->N, this->N);
-    this->V.fill(0.);
-    if (slits == 0) {
-        return;
-    }
-    arma::cx_vec res;
-
-    if (slits % 2 == 0) {
-        res = arma::cx_vec(aperture_separation / this->h,
-                           arma::fill::value(1e10));
-        for (size_t i = 0; i < slits; i += 2) {
-            res = arma::join_cols(
-                res, arma::cx_vec(aperture / this->h, arma::fill::zeros));
-            res = arma::join_cols(
-                arma::cx_vec(aperture / this->h, arma::fill::zeros), res);
-            res =
-                arma::join_cols(res, arma::cx_vec(aperture_separation / this->h,
-                                                  arma::fill::value(1e10)));
-            res = arma::join_cols(arma::cx_vec(aperture_separation / this->h,
-                                               arma::fill::value(1e10)),
-                                  res);
-        }
-    }
-    else {
-        res = arma::cx_vec(aperture / this->h, arma::fill::value(0));
-        for (size_t i = 0; i < slits - 1; i += 2) {
-            res =
-                arma::join_cols(res, arma::cx_vec(aperture_separation / this->h,
-                                                  arma::fill::value(1e10)));
-            res = arma::join_cols(arma::cx_vec(aperture_separation / this->h,
-                                               arma::fill::value(1e10)),
-                                  res);
-            res = arma::join_cols(
-                res, arma::cx_vec(aperture / this->h, arma::fill::zeros));
-            res = arma::join_cols(
-                arma::cx_vec(aperture / this->h, arma::fill::zeros), res);
-        }
-    }
-    if (res.size() > this->N) {
-        abort();
-    }
-    uint32_t fill = (this->N - res.size()) / 2;
-    res = arma::join_cols(arma::cx_vec(fill, arma::fill::value(1e10)), res);
-    res = arma::join_cols(res, arma::cx_vec(fill + ((this->N - res.size()) % 2),
-                                            arma::fill::value(1e10)));
-
-    uint32_t start = pos_x / this->h - thickness / this->h / 2;
-    for (size_t i = 0; i < thickness / this->h; i++) {
-        this->V.col(start + i) = res;
-    }
-}