Merge pull request #12 from comecattin/9-potential-and-kinetic-energies

Compute kinetic, potential and total energies

.gitignore

@@ -39,4 +39,6 @@ bin/
 *.arc
 trajectories.dat
 input.dat
+energies.dat
+energies.pdf
 md_simulation

README.md

@@ -41,3 +41,9 @@ An example is given in the `examples/example_input.dat` file.
 By default, positions along time steps are outputted in the `trajectories.dat` file.
 
 To plot an animation of the trajectories, run `python plot_trajectories.py` in the directory where the `trajectories.dat` file is written.
+
+## Output and plot the energies
+
+Kinetic, potential and total energies are by default computed at each time step and saved under the `energies.dat` file.
+
+To plot theses energies along the time, run `python plot_energies.py energies.dat`. An `energies.pdf` file will be written.

src/energies.f90

@@ -0,0 +1,48 @@
+module energies_module
+    implicit none
+    real(8), parameter :: epsilon = 1.0 ! LJ potential well depth
+    real(8), parameter :: sigma = 1.0 ! LJ distance where potential is zero
+
+contains 
+    subroutine compute_energies(pos, vel, n, ke, pe, te)
+        real(8), dimension(n, 3), intent(in) :: pos, vel
+        integer, intent(in) :: n
+        real(8), intent(out) :: ke, pe, te
+
+        call compute_kinetic_energy(vel, n, ke)
+        call compute_potential_energy(pos, n, pe)
+        te = ke + pe
+
+        print *, "Kinetic energy: ", ke
+        print *, "Potential energy: ", pe
+        print *, "Total energy: ", te
+    end subroutine compute_energies
+
+    subroutine compute_kinetic_energy(vel, n, ke)
+        real(8), dimension(n, 3), intent(in) :: vel
+        integer, intent(in) :: n
+        real(8), intent(out) :: ke
+        integer :: i
+
+        ke = 0.0
+        do i = 1, n
+            ke = ke + 0.5 * sum(vel(i, :) ** 2)
+        end do
+    end subroutine compute_kinetic_energy
+
+    subroutine compute_potential_energy(pos, n, pe)
+        real(8), dimension(n, 3), intent(in) :: pos
+        integer, intent(in) :: n
+        real(8), intent(out) :: pe
+        integer :: i, j
+        real(8) :: r
+
+        pe = 0.0
+        do i = 1, n
+            do j = i + 1, n
+                r = sqrt(sum((pos(i, :) - pos(j, :)) ** 2))
+                pe = pe + 4.0 * epsilon * ((sigma / r) ** 12 - (sigma / r) ** 6)
+            end do
+        end do
+    end subroutine compute_potential_energy   
+end module energies_module

src/main.f90

@@ -1,19 +1,23 @@
 program md_simulation
     use initialization_module, only : initialize
     use forces_module, only : compute_forces
+    use energies_module
     use integration_module, only: integrate
-    use output_module, only: output_positions
+    use output_module, only: output_positions, output_energies
     implicit none
 
     ! Define parameters
     integer :: n_atoms, n_steps, step, num_args
     real(8) :: dt, box_length
-    character(len=100) :: input_file, output_file
+    character(len=100) :: input_file, output_file, output_file_energies
     logical :: file_exists
 
     ! Define position, velocity and forces arrays
     real(8), allocatable:: positions(:,:), velocities(:,:), forces(:,:)
 
+    ! Kinetic, potential and total energies
+    real(8) :: ke, pe, te
+
     num_args = command_argument_count()
     if (num_args /= 1) then
         print *, 'Usage: ./md_simulation input_file'
@@ -48,13 +52,19 @@ program md_simulation
 
     ! Open the output file
     output_file = 'trajectories.dat'
+    output_file_energies = 'energies.dat'
     ! Check if the file already exists
     inquire(file=output_file, exist=file_exists)
     if (file_exists) then
         print *, 'The file ', output_file, ' already exists. Removing it and the .arc file.'
         call system('rm ' //trim(output_file))
         call system('rm ' //trim(output_file(1:index(output_file, '.')-1)) // '.arc')
     end if
+    inquire(file=output_file_energies, exist=file_exists)
+    if (file_exists) then
+        print *, 'The file ', output_file_energies, ' already exists. Removing it'
+        call system('rm ' //trim(output_file_energies))
+    end if
 
     ! Perform the molecular dynamics simulation
     do step = 1, n_steps
@@ -65,9 +75,13 @@ program md_simulation
 
         ! Integrate positions and velocities
         call integrate(positions, velocities, forces, dt, n_atoms, box_length)
+
+        ! Compute the energies
+        call compute_energies(positions, velocities, n_atoms, ke, pe, te)
 
-        ! Output the positions
+        ! Output the positions and the energies
         call output_positions(step, positions, n_atoms, output_file)
+        call output_energies(step, ke, pe, te, output_file_energies)
 
     end do
 

src/makefile

@@ -1,5 +1,5 @@
 FC = gfortran
-OBJ = initialization.o forces.o integration.o output_module.o  main.o
+OBJ = initialization.o forces.o integration.o output_module.o  energies.o main.o
 
 md_simulation: $(OBJ)
 	$(FC) -o md_simulation $(OBJ)
@@ -16,6 +16,9 @@ forces.o: forces.f90
 integration.o: integration.f90
 	$(FC) -c integration.f90
 
+energies.o: energies.f90
+	$(FC) -c energies.f90
+
 output_module.o: output_module.f90
 	$(FC) -c output_module.f90
 

src/output_module.f90

@@ -23,5 +23,19 @@ subroutine output_positions(step, pos, n, filename)
         close(11)
 
     end subroutine output_positions
+
+    subroutine output_energies(step, ek, ep, et, filename)
+        integer, intent(in) :: step
+        real(8), intent(in) :: ek, ep, et
+        character(len=*), intent(in) :: filename
+
+        open(unit=10, file=filename, status='unknown', action='write', position='append')
+        write(10, '(A, I6)') 'Step: ', step
+        write(10, '(A, F10.5)') 'Kinetic Energy: ', ek
+        write(10, '(A, F10.5)') 'Potential Energy: ', ep
+        write(10, '(A, F10.5)') 'Total Energy: ', et
+        close(10)
+
+    end subroutine output_energies
 
 end module output_module

src/plot_energies.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+"""Plot the energies of the system as a function of time."""
+import matplotlib.pyplot as plt
+import numpy as np
+import argparse
+
+def load_data(filename):
+    """Load the data from the file."""
+    with open(filename, 'r') as f:
+        data = f.readlines()
+    for i, line in enumerate(data):
+        data[i] = float(line.replace('Step:', '').replace('Kinetic Energy:', '').replace('Potential Energy:', '').replace('Total Energy:', '').replace('\n', '').replace(' ', ''))
+    step = np.array(data[0::4], dtype=int)
+    kinetic_energy = np.array(data[1::4])
+    potential_energy = np.array(data[2::4])
+    total_energy = np.array(data[3::4])
+    return kinetic_energy, potential_energy, total_energy, step
+
+def plot_energies(kinetic_energy, potential_energy, total_energy, step):
+    """Plot the energies as a function of time."""
+    plt.plot(step, kinetic_energy, label='Kinetic Energy')
+    plt.plot(step, potential_energy, label='Potential Energy')
+    plt.plot(step, total_energy, label='Total Energy')
+    plt.xlabel('Time step')
+    plt.ylabel('Energy')
+    plt.legend()
+    plt.savefig('energies.pdf', bbox_inches='tight', dpi=300)
+    plt.show()
+
+def main():
+    """Read the filename from the command line and plot the energies."""
+    parser = argparse.ArgumentParser(
+        description='Plot the energies of the system as a function of time.'
+    )
+    parser.add_argument(
+        'filename',
+        help='The name of the file containing the energies.'
+    )
+    args = parser.parse_args()
+    kinetic_energy, potential_energy, total_energy, step = load_data(args.filename)
+    plot_energies(kinetic_energy, potential_energy, total_energy, step)
+
+if __name__ == "__main__":
+    main()