#! /usr/bin/env python 

#################################################
# This plots the vector plot of the 
# Magnetic Field of a long thin needle      
#################################################

import matplotlib
import numpy as np
import matplotlib.pyplot as plt

# Size of NxN array note that N must be even 
N = 20
# Size of the two circular masks
R = 0.25

# Length of the rod
l = 1.0
l2 = l/2.0

# Define s and z
s_delt = 1.0/(N/2)
s_max =  1*(1.0 + 0.5*s_delt)
s_min =  -1*(1.0 - 0.5*s_delt)
# Square Axes
z_delt = 1.0/(N/2)
z_min = -1*(1.0 - 0.5*z_delt)
z_max =  1*(1.0 + 0.5*z_delt)

s, z = np.meshgrid(np.arange(s_min,s_max,s_delt),np.arange(z_min,z_max,z_delt))

# Make mask for both poles 
mask = np.logical_or(np.sqrt(s**2 + (z-l2)**2) < R,np.sqrt(s**2 + (z+l2)**2) < R)

# Find the solution for the Magnetic Field
# Find each term in turn
Bs = 1.0/np.sqrt(s**2 + (z - l2)**2)
Bs = Bs - 1.0/np.sqrt(s**2 + (z + l2)**2)
Bs = Bs + (z + l2)**2 / (s**2 + (z + l2)**2)**1.5
Bs = Bs - (z - l2)**2 / (s**2 + (z - l2)**2)**1.5
Bs = (1.0/(l*s))*Bs

Bz = (z + l2) / (s**2 + (z + l2)**2)**1.5
Bz = Bz - (z - l2) / (s**2 + (z - l2)**2)**1.5
Bz = (-1.0/l)*Bz

Bs = np.ma.masked_array(Bs, mask)
Bz = np.ma.masked_array(Bz, mask)

# Make the Plot
plt.close()
plt.box(on='on')
plt.axis('scaled')
plt.axis((-1.1, 1.1, -1.1, 1.1))
plt.title('Magnetic Field Surrounding a Thin Magnetic Needle')
plt.xlabel(r'$\bf s$', fontsize=20)
plt.ylabel(r'$\bf z$', fontsize=20)
plt.quiver(s,z,Bs,Bz,pivot='middle')
plt.savefig("Magnetic_Field_of_Needle.pdf", format="pdf", transparent=True, bbox_inches='tight')
plt.savefig("Magnetic_Field_of_Needle.png", format="png", transparent=True, bbox_inches='tight')