#!/usr/bin/env python
from numpy import *
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt


mu0 = 1e3*4*pi*1e-7
I = .01 #current, amps
n_wire_pts = 50 #points along wire
wire_d = .04 #diameter of wire
wire_sep = .07 #wire separation

gain = 90 #mv/mT
lsb_per_gauss = 6842

#for finite wire diameter use a single hexagonal packing -- 7 filaments, each owning a circle of diameter wire_d/3
filament_positions = vstack((
						array([[0,0]]), 
						(wire_d/3)*array([[cos(t*pi/6),sin(t*pi/6)] for t in range(12)]),
						(wire_d/6)*array([[cos(t*pi/3+pi/6),sin(t*pi/3+pi/6)] for t in range(6)]),
						))
n_filaments = shape(filament_positions)[0]


def dots(a,b):
	return sum(a*b,axis=-1)[...,None]
def magnitude(a):
	return sqrt(dots(a,a)[...,0])

def B(I,ds,r):
	m = sqrt(dots(r,r))
	return mu0*I/4/pi *sum( cross(ds,r)/m/m/m, axis=0) #sum over wire pts

#current carrying wire along x axis
#arrays have axes: (wire pt index ,pt z index,pt x index,pt y index, coords)
x = linspace(-3*wire_sep,3*wire_sep,2)
y = linspace(-1.5*wire_sep,1.5*wire_sep,100)
z = linspace(-1.5*wire_sep,1.5*wire_sep,100)

pts = stack(meshgrid(x,y,z)).T
print shape(pts)
p_wire1 = dstack((linspace(-2,2,n_wire_pts), (-wire_sep/2)*ones(n_wire_pts), zeros(n_wire_pts)) )[0] #points along wire
p_wire2 = dstack((linspace(-2,2,n_wire_pts)[::-1], (wire_sep/2)*ones(n_wire_pts), zeros(n_wire_pts)) )[0] #points along wire
p_wire = []
for f in filament_positions:
	p_wire.extend( p_wire1 + array([0,f[0],f[1]]) )
	p_wire.extend( p_wire2 + array([0,f[0],f[1]]) )
p_wire = asarray(p_wire)
print shape(p_wire)

ds = p_wire[1:]-p_wire[:-1] #differential elements of wire
m = (p_wire[1:]+p_wire[:-1])/2 #midpoints of differential elements

ds = ds.reshape(-1,1,1,1,3)
m = m.reshape(-1,1,1,1,3)

b = B( I/n_filaments, ds, pts-m )
print shape(b)

fig = plt.figure()
#ax = fig.gca(projection='3d')
ax = fig.gca()
print shape(dots(b[:,0,:],b[:,0,:]))
field_strength = magnitude(b[:,0,:])
lw = field_strength
strm = ax.streamplot(pts[:,0,:,1],pts[:,0,:,2], b[:,0,:,1],b[:,0,:,2], 
	color=lsb_per_gauss*1e4*lw, 
	linewidth=1, cmap=plt.cm.autumn)
#C = [5,8,10,12,16,20,32,50]
C = [500,800,1000,1200,1600,2000,3200,5000]
CS = plt.contour(pts[:,0,:,1],pts[:,0,:,2], lsb_per_gauss*1e4*lw, C,colors='k')
plt.clabel(CS, fontsize=9, inline=1, fmt='%d')

#strm = ax.streamplot(pts[plane,1],pts[plane,2], b[plane,1],b[plane,2], color=dots(b[plane],b[plane]), linewidth=2, cmap=plt.cm.autumn)
#fig.colorbar(strm.lines).set_label('Hall effect voltage in uV (with %.1f mV/mT gain)'%(gain))
fig.colorbar(strm.lines).set_label('Magnetometer bits (with %d bits per Gauss)'%(lsb_per_gauss))

#plot wire
plt.plot( wire_sep/2 + filament_positions[...,0], filament_positions[...,1], c='g', ls='',marker='o'  )
plt.plot( -wire_sep/2 + filament_positions[...,0], filament_positions[...,1], c='b', ls='',marker='o'  )
t = linspace(0,2*pi,100)
plt.plot( wire_sep/2 +(wire_d/2)*cos(t), (wire_d/2)*sin(t), lw=2, c='g' )
plt.plot( -wire_sep/2 +(wire_d/2)*cos(t), (wire_d/2)*sin(t), lw=2, c='b' )
plt.title("Magnetic field around pair of current carrying conductors (%.1f mA)"%(I*1e3))

#ax.quiver( pts[plane,1],pts[plane,2], b[plane,1],b[plane,2]) #quiver uses constant length - fix this!
#ax.quiver( pts[:,0],pts[:,1],pts[:,2], b[...,0],b[...,1],b[...,2],length=.1 ) #quiver uses constant length - fix this!
plt.xlim( [amin(y),amax(y)] )
plt.ylim( [amin(z),amax(z)] )
plt.show()