Best_EM.py

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from matplotlib.animation import FuncAnimation
from scipy import constants as si
from random import seed
from random import randint

"""
Seed(2) N_Electrons = 3 and N_Protons = 3 Cube = 5 is interesting

I Would advise to have no more than 12 Particle in a cube at once
"""


seed(21) #change number to randomize
N_Electrons = 3
N_Protons = 3
CubeSize = 5
Proton_yes = True# If false will create positrons

""" 
Blue represents Electrons 
Red represents Protons
It may appear as though the Electrons have no interactions with each other but it is evident that the electrons are 
interacting with each other when they are slowed down after being accelerated by a protons 
"""


v_0 = np.array([0,0,0]) # probably best to keep them all at rest from the start
a_0 = np.zeros_like(v_0) # starting acceleration is kinda useless but needed non the less

class Particle(object):
    number_of_particles = 0
    particle_color = []
    total_KE = 0
    def __init__(self, pos, vel, accl, radius=0.85e-15):
        self.mass = None
        self.pos = pos
        self.vel = vel
        self.accl = accl
        self.radius = radius
        self.KE = 0
        Particle.number_of_particles += 1

    def _r(self, P2):
        r = np.linalg.norm(self.pos - P2.pos)
        if r <= 8.7e-16: # Using the radius of the proton to simulate a collision to avoid a dividing by zero
            return 8.7e-16
        return r

    def r_hat(self, P2):
        return (self.pos - P2.pos) / self._r(P2)

    def force_columbs(self, P2):
        k = 1 / (4 * si.pi * si.epsilon_0)
        return self.charge * P2.charge * k * self.r_hat(P2) / (self._r(P2)) ** 2

    def force_magneto(self, P2):
        k = si.mu_0 / (4 * si.pi)
        B = P2.charge*k*np.cross(P2.vel, P2.r_hat(self)) / (P2._r(self)) ** 2 #magnetic field generated by particle 2, felt by particle 1
        return self.charge * np.cross(self.vel, B)

    def force(self, P2):
        return self.force_columbs(P2) + self.force_magneto(P2)

    def boundary_cross_check_and_update(self):  # checks to see if the x ,y z, positions are in the box of dimensionals maxsize
        if (self.pos[0] < 0 or self.pos[0] > CubeSize):
            self.pos[0] = CubeSize - (self.pos[0] % CubeSize)
            self.vel[0] = -self.vel[0]

        if (self.pos[1] < 0 or self.pos[1] > CubeSize):
            self.pos[1] = CubeSize - (self.pos[1] % CubeSize)
            self.vel[1] = -self.vel[1]

        if (self.pos[2] < 0 or self.pos[2] > CubeSize):
            self.pos[2] = CubeSize - (self.pos[2] % CubeSize)
            self.vel[2] = -self.vel[2]

    def update_accl(self,force):
        self.accl = force / self.mass

    def update_pos(self, dt):
        self.pos = self.pos + self.vel * dt + 1 / 2 * self.accl * dt ** 2
        self.vel = self.vel + self.accl * dt
        self.boundary_cross_check_and_update()
        if self.vel[0] > si.c or self.vel[1] > si.c or self.vel[2] > si.c:
            print("EINSTEIN WAS WRONG")
            self.vel = self.vel/1e9
        Particle.total_KE += 1/2*self.mass*self.vel**2



class Electron(Particle):
    def __init__(self, pos, vel, accl):
        super(Electron, self).__init__(pos, vel, accl)
        self.type = "Electron"
        self.mass = si.m_e
        self.charge = -1 * si.e
        self.color = "Blue"
        self.__class__.particle_color.append(self.color)


class Proton(Particle):
    def __init__(self, pos, vel, accl, ):
        super(Proton, self).__init__(pos, vel, accl)
        self.type = "Proton"
        self.mass = si.m_p  # change to si.m_e if you want a positron, alot more chaotic
        self.charge = si.e
        self.color = "Red"
        self.__class__.particle_color.append(self.color)


class Positron(Particle):
    def __init__(self, pos, vel, accl,):
        super(Positron, self).__init__(pos, vel, accl)
        self.type = "Proton"
        self.mass = si.m_e
        self.charge = si.e
        self.color = "Orange"
        self.__class__.particle_color.append(self.color)






class Simulation():
    def __init__(self, particles):
        self.particles = particles

    def step(self, dt=9e-4):
        coords = None
        for p1 in self.particles:
            sum_of_forces = np.zeros_like(p1.pos) # creates a array with the same dimensions as position
            for p2 in self.particles:
            # this for loop is for 1 particle to be updated, p1, it checks every other particle and sums up the forces p1 feels from all the other particles p2
                if p1 == p2:
                    continue
                sum_of_forces = np.add(sum_of_forces, p1.force(p2))
            p1.update_accl(sum_of_forces) #updates p1's acceleraton vector NOT POSITION
        #now that the acceleration vector of each particle is updated its time to actually move each particle
        #this ensures that a particles new position doesnt affect another particles from the previous time step
        for p1 in self.particles:
            p1.update_pos(dt)


            coords = np.vstack((coords, p1.pos)) if coords is not None else \
            p1.pos
        return coords


def main():
    particle_array=[]
    if Proton_yes:
        for k in range(N_Protons): # Generates randomized Protons
            if k == 0: # Gurantees a proton in the centre
                x=y=z = CubeSize/2
                pos_0 = np.array([x, y, z])
                particle_array.append(Proton(pos_0, v_0, a_0))
            else:
                x = CubeSize * randint(1, N_Protons-1) / (N_Protons+1)
                y = CubeSize * randint(1, N_Protons-1) / (N_Protons+1)
                z = CubeSize * randint(1, N_Protons-1) / (N_Protons+1)
                pos_0 = np.array([x,y,z])
                particle_array.append(Proton(pos_0, v_0, a_0))
    else:
        for k in range(N_Protons): # Generates randomized Protons
            if k == 0: # Gurantees a proton in the centre
                x=y=z = CubeSize/2
                pos_0 = np.array([x, y, z])
                particle_array.append(Positron(pos_0, v_0, a_0))
            else:
                x = CubeSize * randint(1, N_Protons-1) / (N_Protons+1)
                y = CubeSize * randint(1, N_Protons-1) / (N_Protons+1)
                z = CubeSize * randint(1, N_Protons-1) / (N_Protons+1)
                pos_0 = np.array([x,y,z])
                particle_array.append(Positron(pos_0, v_0, a_0))

    for k in range(N_Electrons): # Generates randomized Electrons
        x = CubeSize * randint(0, N_Electrons) / N_Electrons
        y = CubeSize * randint(0, N_Electrons) / N_Electrons
        z = CubeSize * randint(0, N_Electrons) / N_Electrons
        pos_0 = np.array([x, y, z])
        particle_array.append(Electron(pos_0, v_0, a_0))

    s = Simulation(particle_array)

    fig = plt.figure(1, figsize=(6, 6))
    ax = fig.add_subplot(111, projection='3d')
    ax.set_xticks([]), ax.set_yticks([]), ax.set_zticks([])
    plt.xlabel('x')
    plt.ylabel('y')
    ax.set_xlim3d(0, CubeSize)
    ax.set_ylim3d(0, CubeSize)
    ax.set_zlim3d(0, CubeSize)
    i = s.step()
    graph = ax.scatter(i[..., 0], i[..., 1], i[..., 2], '.', c=Particle.particle_color)



    def update(frame_number):
        i = s.step()
        graph._offsets3d = (i[..., 0], i[..., 1], i[..., 2])
        return graph,

    animation = FuncAnimation(fig, update, interval=1, blit=False)
    plt.show()


if __name__ == "__main__":
    main()