An application to visualize the field create by a list of coils.
This code is fairly complex, but it is actuallty a very rich application, and a full-blown example of what you might want to do
Python source code: coil_application.py
import numpy as np
from scipy import linalg, special
from traits.api import HasTraits, Array, CFloat, Str, List, \
Instance, on_trait_change
from traitsui.api import Item, View, HGroup, ListEditor, \
HSplit, VSplit, spring
from mayavi.core.ui.api import EngineView, MlabSceneModel, \
SceneEditor
##############################################################################
# A current loop
class Loop(HasTraits):
""" A current loop.
"""
direction = Array(float, value=(0, 0, 1), cols=3,
shape=(3,), desc='directing vector of the loop',
enter_set=True, auto_set=False)
# CFloat tries to convert automatically to floats
radius = CFloat(0.1, desc='radius of the loop',
enter_set=True, auto_set=False)
position = Array(float, value=(0, 0, 0), cols=3,
shape=(3,), desc='position of the center of the loop',
enter_set=True, auto_set=False)
plot = None
name = Str
view = View(HGroup(Item('name', style='readonly', show_label=False),
spring, 'radius'),
'position', 'direction', '_')
# For a Borg-like pattern
__shared_state = {'number':0}
def __init__(self, **traits):
HasTraits.__init__(self, **traits)
self.__shared_state['number'] += 1
self.name = 'Coil %i' % self.__shared_state['number']
def base_vectors(self):
""" Returns 3 orthognal base vectors, the first one colinear to
the axis of the loop.
"""
# normalize n
n = self.direction / (self.direction**2).sum(axis=-1)
# choose two vectors perpendicular to n
# choice is arbitrary since the coil is symetric about n
if np.abs(n[0])==1 :
l = np.r_[n[2], 0, -n[0]]
else:
l = np.r_[0, n[2], -n[1]]
l /= (l**2).sum(axis=-1)
m = np.cross(n, l)
return n, l, m
@on_trait_change('direction,radius,position')
def redraw(self):
if hasattr(self, 'app'):
self.mk_B_field()
if self.app.scene._renderer is not None:
self.display()
self.app.visualize_field()
def display(self, half=False):
"""
Display the coil in the 3D view.
If half is True, display only one half of the coil.
"""
n, l, m = self.base_vectors()
theta = np.linspace(0, (2-half)*np.pi, 30)
theta = theta[..., np.newaxis]
coil = self.radius*(np.sin(theta)*l + np.cos(theta)*m)
coil += self.position
coil_x, coil_y, coil_z = coil.T
if self.plot is None:
self.plot = self.app.scene.mlab.plot3d(coil_x, coil_y, coil_z,
tube_radius=0.007, color=(0, 0, 1),
name=self.name )
else:
self.plot.mlab_source.set(x=coil_x, y=coil_y, z=coil_z)
def mk_B_field(self):
"""
returns the magnetic field for the current loop calculated
from eqns (1) and (2) in Phys Rev A Vol. 35, N 4, pp. 1535-1546; 1987.
return:
B is a vector for the B field at point r in inverse units of
(mu I) / (2 pi d)
for I in amps and d in meters and mu = 4 pi * 10^-7 we get Tesla
"""
### Translate the coordinates in the coil's frame
n, l, m = self.base_vectors()
R = self.radius
r0 = self.position
r = np.c_[np.ravel(self.app.X), np.ravel(self.app.Y),
np.ravel(self.app.Z)]
# transformation matrix coil frame to lab frame
trans = np.vstack((l, m, n))
r -= r0 #point location from center of coil
r = np.dot(r, linalg.inv(trans) ) #transform vector to coil frame
#### calculate field
# express the coordinates in polar form
x = r[:, 0]
y = r[:, 1]
z = r[:, 2]
rho = np.sqrt(x**2 + y**2)
theta = np.arctan(x/y)
E = special.ellipe((4 * R * rho)/( (R + rho)**2 + z**2))
K = special.ellipk((4 * R * rho)/( (R + rho)**2 + z**2))
Bz = 1/np.sqrt((R + rho)**2 + z**2) * (
K
+ E * (R**2 - rho**2 - z**2)/((R - rho)**2 + z**2)
)
Brho = z/(rho*np.sqrt((R + rho)**2 + z**2)) * (
-K
+ E * (R**2 + rho**2 + z**2)/((R - rho)**2 + z**2)
)
# On the axis of the coil we get a divided by zero here. This returns a
# NaN, where the field is actually zero :
Brho[np.isnan(Brho)] = 0
B = np.c_[np.cos(theta)*Brho, np.sin(theta)*Brho, Bz ]
# Rotate the field back in the lab's frame
B = np.dot(B, trans)
Bx, By, Bz = B.T
Bx = np.reshape(Bx, self.app.X.shape)
By = np.reshape(By, self.app.X.shape)
Bz = np.reshape(Bz, self.app.X.shape)
Bnorm = np.sqrt(Bx**2 + By**2 + Bz**2)
# We need to threshold ourselves, rather than with VTK, to be able
# to use an ImageData
Bmax = 10 * np.median(Bnorm)
Bx[Bnorm > Bmax] = np.NAN
By[Bnorm > Bmax] = np.NAN
Bz[Bnorm > Bmax] = np.NAN
Bnorm[Bnorm > Bmax] = np.NAN
self.Bx = Bx
self.By = By
self.Bz = Bz
self.Bnorm = Bnorm
##############################################################################
# The application
class Application(HasTraits):
scene = Instance(MlabSceneModel, (), editor=SceneEditor())
# The mayavi engine view.
engine_view = Instance(EngineView)
# We use a traits List to be able to add coils to it
coils = List(Loop,
value=( Loop(position=(0, 0, -0.05), ),
Loop(position=(0, 0, 0.05), ), ),
editor=ListEditor(use_notebook=True, deletable=False,
style='custom'),
)
# The grid of points on which we want to evaluate the field
X, Y, Z = np.mgrid[-0.15:0.15:20j, -0.15:0.15:20j, -0.15:0.15:20j]
# Avoid rounding issues:
f = 1e4 # this gives the precision we are interested by :
X = np.round(X * f) / f
Y = np.round(Y * f) / f
Z = np.round(Z * f) / f
Bx = Array(value=np.zeros_like(X))
By = Array(value=np.zeros_like(X))
Bz = Array(value=np.zeros_like(X))
Bnorm = Array(value=np.zeros_like(X))
field = None
def __init__(self, **traits):
HasTraits.__init__(self, **traits)
self.engine_view = EngineView(engine=self.scene.engine)
@on_trait_change('scene.activated')
def init_view(self):
# This gets fired when the viewer of the scene is created
self.scene.scene_editor.background = (0, 0, 0)
for coil in self.coils:
coil.app = self
coil.display()
coil.mk_B_field()
self.visualize_field()
def visualize_field(self):
self.Bx = np.zeros_like(self.X)
self.By = np.zeros_like(self.X)
self.Bz = np.zeros_like(self.X)
self.Bnorm = np.zeros_like(self.X)
for coil in self.coils:
if hasattr(coil, 'Bx'):
self.Bx += coil.Bx
self.By += coil.By
self.Bz += coil.Bz
self.Bnorm += coil.Bnorm
if self.field is None:
self.field = self.scene.mlab.pipeline.vector_field(
self.X, self.Y, self.Z, self.Bx, self.By, self.Bz,
scalars = self.Bnorm,
name='B field')
vectors = self.scene.mlab.pipeline.vectors(self.field,
mode='arrow', resolution=10,
mask_points=6, colormap='YlOrRd',
scale_factor=2*np.abs(self.X[0,0,0]
-self.X[1,1,1]) )
vectors.module_manager.vector_lut_manager.reverse_lut = True
vectors.glyph.mask_points.random_mode = False
self.scene.mlab.axes()
self.scp = self.scene.mlab.pipeline.scalar_cut_plane(self.field,
colormap='hot')
else:
self.field.mlab_source.set(x=self.X, y=self.Y, z=self.Z,
u=self.Bx, v=self.By, w=self.Bz,
scalars=self.Bnorm)
view = View(HSplit(
VSplit(Item(name='engine_view',
style='custom',
resizable=True),
Item('coils', springy=True),
show_labels=False),
'scene',
show_labels=False),
resizable=True,
title='Coils...',
height=0.8,
width=0.8,
)
##############################################################################
if __name__ == '__main__':
app = Application()
app.configure_traits()