A simple Laplace solver for PHYS2231

induction.py

#python3 induction.py
import numpy as numpy
import math
import matplotlib.pyplot as plt

SIZE_X = 128
SIZE_Y = 256

LENGTH_SCALE = 0.01#meters/node

potential = numpy.zeros((SIZE_X,SIZE_Y))
boundary_conditions = numpy.zeros((SIZE_X,SIZE_Y))

forces = []

def gauss_seidel(potential,boundary_conditions):
	####Simple Gauss-Seidel Laplace solver####
	
	while(True):
		residual = potential.copy()
		for x in range(1,SIZE_X-1):
			for y in range(1,SIZE_Y-1):
			  if(not boundary_conditions[x,y]):
  				potential[x,y] = (potential[x+1,y]+potential[x-1,y]+
	  						              potential[x,y+1]+potential[x,y-1])/4.0
	  						              
		residual = residual-potential
		
		print("Residual: {}".format(numpy.linalg.norm(residual)))
		if(numpy.linalg.norm(residual) < 5): #tolerance
			break

def get_normal_vector(boundary_conditions,boundary_id,x,y):
      # Mainly handles corners.
      n_x = 0
      n_y = 0
      if(boundary_conditions[x+1,y] == boundary_id): n_x += 1
      if(boundary_conditions[x-1,y] == boundary_id): n_x -= 1
      if(boundary_conditions[x,y+1] == boundary_id): n_y += 1
      if(boundary_conditions[x,y-1] == boundary_id): n_y -= 1

      normal_mag = math.sqrt(n_x**2.0 + n_y**2.0)
      if(normal_mag != 0): #Convert to unit vector
        n_x /= normal_mag
        n_y /= normal_mag #horribly slow, but
        
      return (n_x, n_y)

def get_electric_field(potential,x,y,normal):
      # Obtains E from phi using central-difference method.
      #
      #   +---+
      #   | \
      #   |  \
      #   +---o
      #
      # You can see that the finite-difference step direction to determine electric field should change
      # to fit the normal direction.
      if(normal[0] > 0): # norm pointing right?
        Ex = (potential[x+1,y] - potential[x,y])/LENGTH_SCALE #get the field to the right. 
      else:
        Ex = (potential[x,y] - potential[x-1,y])/LENGTH_SCALE 

              
      if(normal[1] > 0): # norm pointing right?
        Ey = (potential[x,y+1] - potential[x,y])/LENGTH_SCALE #get the field to the right. 
      else:
        Ey = (potential[x,y] - potential[x,y-1])/LENGTH_SCALE #eyy lmao

      return (Ex,Ey)


def get_electric_force(boundary_id):
  total_x = 0
  total_y = 0
  force = numpy.zeros((SIZE_X,SIZE_Y))
  for x in range(1,SIZE_X-1):
    for y in range(1,SIZE_Y-1):
	    # We need the normal electric field to calculate the surface charge density.
	    # If this point is adjacent to one of the boundaries, 
	    # determine the normal vector via hard-coded LUT
	    # based on the local geometry.
	    # This won't be perfectly accurate, but it should be good enough for our purposes.
	    # ^E dot ^n / 4pi = sigma, surface charge density
	    # From that we can get the force:
	    # F = 2pi sigma ^2 dA
      n_x, n_y = get_normal_vector(boundary_conditions,boundary_id,x,y)
      Ex, Ey = get_electric_field(potential,x,y,(n_x,n_y))
      Fx = ((Ex*n_x / 4*math.pi)**2.0)*2.0*math.pi
      Fy = ((Ey*n_y / 4*math.pi)**2.0)*2.0*math.pi
      Fmag = math.sqrt(Fx**2.0 + Fy**2.0)
      force[x,y] = Fy
      total_x += Fx
      total_y += Fy
  print(total_x)
  print(total_y)
  return force,total_x,total_y
	    
      

def circle(potential,radius,center,voltage,boundary_id):
  #circles!
  for x in range(0,SIZE_X):
    for y in range(0,SIZE_Y):
      distance = math.sqrt(((x-center[0])**2)+((y-center[1])**2))
      if(distance < radius):
        potential[x,y] = voltage
        boundary_conditions[x,y] = boundary_id

circle_y = 50
distances = []
while True:
  potential = numpy.zeros((SIZE_X,SIZE_Y))
  boundary_conditions = numpy.zeros((SIZE_X,SIZE_Y))


  circle(potential,20,(64,circle_y),100,1)
  circle(potential,10,(64,200),100,2)
  circle_y+=30
  gauss_seidel(potential,boundary_conditions)

  force,total_x,total_y = get_electric_force(1)
  forces.append(total_y)
  distances.append(circle_y)
  
  plt.subplot(1, 2, 1)
  plt.gca().set_title('Electric Potential (x,y)')
  plt.imshow(potential)
  plt.subplot(1, 2, 2)
  plt.gca().set_title('Y force distribution (N)')
  plt.imshow(force)
#  plt.subplot(1, 3, 3)
#  plt.gca().set_title('Force / Distance')
#  plt.plot(distances,forces)
  plt.savefig(str(circle_y) + '.png')
  plt.draw()
  plt.pause(0.1)
  plt.clf()
  plt.cla()