atom_class.py 3.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96
#!/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()