Molecular Simulation: PBC using the Ewald Summation technique

Question

The following is a Rust implementation of polymer simulation where the PBC(periodic boundary condition) has been implemented using the Ewald Summation technique.

Can you review the code?

Is the PBC properly implemented using the Ewal Summation?

use rand::Rng;
use std::f64::consts::PI;
use std::f64::consts::E;
use plotters::prelude::*;

const N: usize = 9;
const SIGMA: f64 = 2.0;
const PERIODIC_BOUNDARY: f64 = 5.0;
const TEMPERATURE: f64 = 2.0;
const MIN_ATOM_DISTANCE: f64 = 1.0;
const MAX_ATOM_DISTANCE: f64 = 3.8;
const SIM_STEPS: usize = 10000;
const WRITE_STEPS: usize = 10;
const EWALD_ALPHA: f64 = 0.5;

type Point = [f64; 2];

struct Polymer {
    chain_vec: Vec<Point>,
}

impl Polymer {
    fn new() -> Self {
        Polymer {
            chain_vec: Vec::new(),
        }
    }

    fn plot_energies(&self, energies: &[f64]) {
        let root = BitMapBackend::new("energies.png", (800, 600)).into_drawing_area();
        root.fill(&WHITE).unwrap();
        let (min, max) = (0, energies.len());
        let (x_min, x_max) = (min as f64, max as f64);
        let (y_min, y_max) = (energies.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
                              energies.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)));
        let mut chart = ChartBuilder::on(&root)
            .x_label_area_size(40)
            .y_label_area_size(40)
            .margin(10)
            .caption("Energies", ("sans-serif", 40.0).into_font())
            .build_cartesian_2d(x_min..x_max, y_min..y_max)
            .unwrap();

        chart.configure_mesh().draw().unwrap();

        chart.draw_series(LineSeries::new(
            (min..max).map(|x| (x as f64, energies[x])),
            &RED,
        )).unwrap();

        println!("Check the output in file 'energies.png'");
    }

    fn apply_boundary_condition(&self, coord: Point) -> Point {
        let mut x = coord[0];
        let mut y = coord[1];
        if x > PERIODIC_BOUNDARY {
            x -= PERIODIC_BOUNDARY;
        } else if x < 0.0 {
            x += PERIODIC_BOUNDARY;
        }
        if y > PERIODIC_BOUNDARY {
            y -= PERIODIC_BOUNDARY;
        } else if y < 0.0 {
            y += PERIODIC_BOUNDARY;
        }
        [x, y]
    }

    fn get_distance(&self, coord1: Point, coord2: Point) -> f64 {
        let dx = coord1[0] - coord2[0];
        let dy = coord1[1] - coord2[1];
        let dx = dx - PERIODIC_BOUNDARY * (dx / PERIODIC_BOUNDARY).round();
        let dy = dy - PERIODIC_BOUNDARY * (dy / PERIODIC_BOUNDARY).round();
        (dx * dx + dy * dy).sqrt()
    }

    fn get_point_at_radius(&self, coord: Point, radius: f64) -> Point {
        let angle = rand::thread_rng().gen::<f64>() * 2.0 * PI;
        let x = (angle.cos() * radius) + coord[0];
        let y = (angle.sin() * radius) + coord[1];
        self.apply_boundary_condition([x, y])
    }

    fn morse_potential_func(&self, r: f64) -> f64 {
        (E.powf(-2.0 * SIGMA * (r - MIN_ATOM_DISTANCE)))
            - (2.0 * E.powf(-SIGMA * (r - MIN_ATOM_DISTANCE)))
    }

    fn ewald_short_range_potential(&self, coord: Point) -> f64 {
        let mut potential = 0.0;
        for (_i, &other_loc) in self.chain_vec.iter().enumerate() {
            if self.get_distance(coord, other_loc) < MIN_ATOM_DISTANCE {
                potential += self.morse_potential_func(MIN_ATOM_DISTANCE);
            }
        }
        potential
    }

    fn ewald_long_range_potential(&self, coord: Point) -> f64 {
        let mut potential = 0.0;
        for m in -5..=5 {
            for n in -5..=5 {
                let mut sum = 0.0;
                for &other_loc in self.chain_vec.iter() {
                    let dx = coord[0] - other_loc[0] + (m as f64) * PERIODIC_BOUNDARY;
                    let dy = coord[1] - other_loc[1] + (n as f64) * PERIODIC_BOUNDARY;
                    let r = (dx * dx + dy * dy).sqrt();
                    sum += self.morse_potential_func(r);
                }
                potential += sum;
            }
        }
        potential * PI / (EWALD_ALPHA * EWALD_ALPHA)
    }

    fn get_potential(&self, index: usize) -> f64 {
        let current_loc = self.chain_vec[index];
        let short_range_potential = self.ewald_short_range_potential(current_loc);
        let long_range_potential = self.ewald_long_range_potential(current_loc);
        short_range_potential + long_range_potential
    }

    fn get_total_potential(&self) -> f64 {
        let mut potential = 0.0;
        for i in 0..self.chain_vec.len() {
            potential += self.get_potential(i);
        }
        potential
    }

    fn initialize_polymer(&mut self) {
        let mut current_coord = [0.0, 0.0];
        self.chain_vec.push(current_coord);
        for _ in 0..N - 1 {
            let new_loc = self.get_point_at_radius(current_coord, MAX_ATOM_DISTANCE);
            self.chain_vec.push(new_loc);
            current_coord = new_loc;
        }
    }

    fn run_simulation(&mut self, steps: usize) {
        for _ in 0..steps {
            let rand_index = rand::thread_rng().gen_range(0..self.chain_vec.len());
            let old_loc = self.chain_vec[rand_index];
            let new_loc = self.get_point_at_radius(old_loc, MAX_ATOM_DISTANCE);
            self.chain_vec[rand_index] = new_loc;

            let old_potential = self.get_potential(rand_index);
            let new_potential = self.get_potential(rand_index);
            let delta_energy = new_potential - old_potential;

            if delta_energy > 0.0 {
                let acceptance_probability = (-delta_energy / TEMPERATURE).exp();
                if rand::thread_rng().gen::<f64>() > acceptance_probability {
                    self.chain_vec[rand_index] = old_loc;
                }
            }
        }
    }
}

fn main() {
    let mut polymer = Polymer::new();
    polymer.initialize_polymer();
    let mut energies = Vec::new();

    for step in 0..SIM_STEPS {
        polymer.run_simulation(WRITE_STEPS);
        let total_potential = polymer.get_total_potential();
        energies.push(total_potential);

        if step % WRITE_STEPS == 0 {
            println!("Step: {}, Energy: {}", step, total_potential);
        }
    }

    polymer.plot_energies(&energies);
}

Answer 1

I lack the domain-specific knowledge to currently check whether Ewald Summation is implemented correclty. Maybe somebody else can do that. I, however, have made some other observations about your code:

Separate main.rs and lib.rs

You currently have everything in one unit. That's fine. It may, however, be preferable to make your struct Polymer and the type Point public and oursource the library code, i.e. everything but main() into lib.rs in your src folder.

Make public methods public

Currently new() and simulate() are private methods of Polymer. Consider making them public to have access to them when you import them in other projects or your main.rs.

Separate new() and default()

Your current new() does not accept any of the struct's members. It should accept polymer_chain: Vec<Point>.
You can then also implement the Default trait to achieve what your current new() does.

Useless casting

Casting PERIODIC_BOUNDARY to f64 is useless. It is already declared as such.

Use linters

I discovered all of the above issues by just using cargo clippy on your project. You should use it, too.

Don't use unwrap()

...in production code. It might cause panics that will be hard to track for the end-user. Handle all errors that may arise.

Let the functions be free

Your struct Polymer currently has member functions that have nothing to do with the actual polymer. Consider making them free functions.

Use (Doc-)Tests

Testing your code for correct behaviour can make it unnecessary to ask for its correctness. Write tests for the functions that assert whether they produce the correct result. I recommend writing the tests before the functions (-> test-driven development / TDD) so that you let your design be guided by the outcomes you want to achieve.

lib.rs

use plotters::prelude::*;
use rand::Rng;
use std::f64::consts::PI;

const N: usize = 9;
const SIGMA: f64 = 2.0;
const PERIODIC_BOUNDARY: f64 = 5.0;
const TEMPERATURE: f64 = 2.0;
const MIN_ATOM_DISTANCE: f64 = 1.0;
const MAX_ATOM_DISTANCE: f64 = 3.8;
const K: f64 = 1.0;
const SIM_STEPS: usize = 10000;
const WRITE_STEPS: usize = 10;

pub type Point = [f64; 2];

pub struct Polymer {
    polymer_chain: Vec<Point>,
}

impl Polymer {
    #[must_use]
    pub fn new(polymer_chain: Vec<Point>) -> Self {
        Self { polymer_chain }
    }
    // Perform the simulation and plot the energies
    pub fn simulate(&mut self) {
        self.initialize_polymer();
        let mut total_pot_vec = Vec::new();

        for _ii in 1..=SIM_STEPS {
            self.run_simulation(WRITE_STEPS);
            let total_pot_double = self.get_total_potential().round();
            total_pot_vec.push(total_pot_double);
        }

        plot_energies(&total_pot_vec);
    }

    // Calculate the Ewald summation of the potential energy
    fn ewald_summation(&self, index: usize) -> f64 {
        let mut real_sum = 0.0;
        let current_loc = self.polymer_chain[index];

        // Calculate the real space summation
        for (i, &other_loc) in self.polymer_chain.iter().enumerate() {
            if i != index {
                let r = get_distance(current_loc, other_loc);
                real_sum += morse_potential(r);
            }
        }

        let mut image_vec = Vec::new();
        for i in -2..=2 {
            for j in -2..=2 {
                image_vec.push([
                    f64::from(i).mul_add(PERIODIC_BOUNDARY, current_loc[0]),
                    f64::from(j).mul_add(PERIODIC_BOUNDARY, current_loc[1]),
                ]);
            }
        }

        // Calculate the reciprocal space summation
        for (i, &image_loc) in image_vec.iter().enumerate() {
            for (j, &other_loc) in self.polymer_chain.iter().enumerate() {
                if j != index {
                    let r = get_distance(image_loc, other_loc);
                    real_sum += morse_potential(r);
                }
            }

            if i < image_vec.len() - 1 {
                let r = get_distance(image_loc, image_vec[i + 1]);
                real_sum += harmonic_energy(r);
            }
        }

        real_sum
    }

    // Calculate the total potential energy of the system
    fn get_total_potential(&self) -> f64 {
        let mut harmonic_pot = 0.0;
        let mut pair_pot = 0.0;

        for i in 0..self.polymer_chain.len() - 1 {
            let current_bead = self.polymer_chain[i];
            let next_bead = self.polymer_chain[i + 1];
            let r = get_distance(current_bead, next_bead);
            harmonic_pot += harmonic_energy(r);

            for j in i + 1..self.polymer_chain.len() {
                let bead = self.polymer_chain[j];
                let r = get_distance(current_bead, bead);
                pair_pot += morse_potential(r);
            }
        }

        harmonic_pot + pair_pot
    }

    // Initialize the polymer by generating random points
    fn initialize_polymer(&mut self) {
        let mut current_point = [0.0, 0.0];
        self.polymer_chain.push(current_point);

        for _ in 0..N - 1 {
            let new_loc = get_point_at_radius(current_point, MAX_ATOM_DISTANCE);
            self.polymer_chain.push(new_loc);
            current_point = new_loc;
        }
    }

    // Run the simulation for the specified number of steps
    fn run_simulation(&mut self, no_of_steps: usize) {
        for _ in 0..no_of_steps {
            let rand_index = rand::thread_rng().gen_range(0..self.polymer_chain.len());
            let before_loc = self.polymer_chain[rand_index];
            let before_pot = self.ewald_summation(rand_index);
            let new_loc = get_point_at_radius(before_loc, MAX_ATOM_DISTANCE);
            self.polymer_chain[rand_index] = new_loc;
            let after_pot = self.ewald_summation(rand_index);
            let pot_diff = after_pot - before_pot;

            if pot_diff < 0.0 {
                continue;
            }

            let rand = rand::thread_rng().gen::<f64>();

            if (-(pot_diff / TEMPERATURE)).exp() > rand {
                continue;
            }

            self.polymer_chain[rand_index] = before_loc;
        }
    }
}

impl Default for Polymer {
    fn default() -> Self {
        Self::new(Vec::new())
    }
}

// Plot the energies using the plotters library
fn plot_energies(energies: &[f64]) {
    // Create a drawing area for the plot
    let root = BitMapBackend::new("energies.png", (800, 600)).into_drawing_area();
    if let Err(error) = root.fill(&WHITE) {
        eprintln!("Error when filling root: {error}");
    }

    // Define the range of x and y axes for the plot
    let (min, max) = (0, energies.len());
    let (x_min, x_max) = (min as f64, max as f64);
    let (y_min, y_max) = (
        energies.iter().fold(f64::INFINITY, |a, &b| a.min(b)),
        energies.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b)),
    );

    // Create the chart with labels and captions
    ChartBuilder::on(&root)
        .x_label_area_size(40)
        .y_label_area_size(40)
        .margin(10)
        .caption("Energies", ("sans-serif", 40.0).into_font())
        .build_cartesian_2d(x_min..x_max, y_min..y_max)
        .map_or_else(
            |error| {
                eprintln!("Could not build chart: {error}");
            },
            |mut chart| {
                // Draw the mesh of the chart
                if let Err(error) = chart.configure_mesh().draw() {
                    eprintln!("Could not draw mesh: {error}");
                }

                // Draw the line series representing the energies
                if let Err(error) = chart.draw_series(LineSeries::new(
                    (min..max).map(|x| (x as f64, energies[x])),
                    &RED,
                )) {
                    eprintln!("Could not draw series: {error}");
                }

                println!("Check the output in file 'energies.png'");
            },
        );
}

// Apply periodic boundary conditions to a point
fn apply_boundary_condition(point: Point) -> Point {
    let mut x = point[0];
    let mut y = point[1];

    // Wrap x-coordinate within the periodic boundary
    if x > PERIODIC_BOUNDARY {
        x -= PERIODIC_BOUNDARY;
    } else if x < 0.0 {
        x += PERIODIC_BOUNDARY;
    }

    // Wrap y-coordinate within the periodic boundary
    if y > PERIODIC_BOUNDARY {
        y -= PERIODIC_BOUNDARY;
    } else if y < 0.0 {
        y += PERIODIC_BOUNDARY;
    }

    [x, y]
}

// Get a new point at a given radius from the current point
fn get_point_at_radius(curr_pos: Point, radius: f64) -> Point {
    let angle = rand::thread_rng().gen::<f64>() * 2.0 * PI;
    let x = angle.cos().mul_add(radius, curr_pos[0]);
    let y = angle.sin().mul_add(radius, curr_pos[1]);
    apply_boundary_condition([x, y])
}

// Calculate the distance between two points using Euclidean distance formula
fn get_distance(point_one: Point, point_two: Point) -> f64 {
    let dx = point_one[0] - point_two[0];
    let dy = point_one[1] - point_two[1];
    dx.hypot(dy)
}

// Calculate the Morse potential function given the distance
fn morse_potential(r: f64) -> f64 {
    2.0f64.mul_add(
        -(-SIGMA * (r - MIN_ATOM_DISTANCE)).exp(),
        (-2.0 * SIGMA * (r - MIN_ATOM_DISTANCE)).exp(),
    )
}

// Calculate the harmonic energy function given the distance
fn harmonic_energy(r: f64) -> f64 {
    K * (r - MAX_ATOM_DISTANCE).powi(2)
}

main.rs

use polymer::Polymer;

fn main() {
    let mut polymer = Polymer::default();
    polymer.simulate();
}