#!/usr/bin/env python3 class Atom: """Define details of atoms in our 1D simulation.""" def __init__(self, species, mass, position, velocity): """Initialise the atom class. The atomic species, mass (amu), position (Angstrom) and velocity (Angstrom/ps) must be specified. The force is always initialized to zero. """ self.species = species self.mass = mass self.position = position self.velocity = velocity self.force = 0.0 def update_position(self, timestep): """Update the position using its velocity and the timestep (ps).""" self.position = self.position + timestep * self.velocity def update_velocity(self, timestep): """Update the velocity given force on it and timestep.""" self.velocity = self.velocity + self.force/self.mass*timestep def set_force(self, atom_list, force_func): """Find the force on this atom. This is found from a list of atoms and a function for calculating the force between two atoms. Recall we can pass functions as arguments just as variables or any other objects. """ self.force = 0.0 for atom2 in atom_list: # We don't want to include a force from the atom on itself. if self != atom2: self.force = self.force + force_func(self.position, atom2.position) def harmonic_force(pos1, pos2, k=343415.05, equil=0.74): """Return a harmonic force on pos1 given pos2.""" # Two optional parameters above are specified by adding default values, so # they don't need to be specified when the function is called. # We use a very simple force function here. if pos1 < pos2: equil = - equil # This gives the force in the correct direction whether the atom is # to the left or right of the other. return -k * (pos1 - pos2 - equil) def get_atoms(): """Return a list of atom objects representing our simulated atoms.""" # In a real simulation this information would be read from a file. atom_list = [] # Lists can be composed of class instances. atom_list.append(Atom('H', 1.0, 0.0, 0.0)) # We could also refer to paramter names explicitly. The order doesn't # matter if we do so. atom_list.append(Atom(mass=1.0, species='H', position=0.9, velocity=0.0)) return atom_list def output_positions(atom_list, time): """Output all atomic positions along with the current time.""" # We use end=' ' to suppress the newline but add a space, so all positions # are on one line for each timestep. print(time, end=' ') for atom in atom_list: print(atom.position, end=' ') print() def run_simulation(atoms, timestep, nsteps, force_func): """Simulate the evolution of atomic positions under force_func,""" # Start by outputting the initial positions. time = 0.0 output_positions(atoms, time) for i_t in range(nsteps): time = time + timestep # First update all the velocities. for atom in atoms: atom.set_force(atoms, harmonic_force) atom.update_velocity(timestep) # Then update all the positions. for atom in atoms: atom.update_position(timestep) # And output the new positions. output_positions(atoms, time) def main(): atom_list = get_atoms() run_simulation(atoms=atom_list, timestep=0.0001, nsteps=1000, force_func=harmonic_force) if __name__ == '__main__': main()