<s> """<STR_LIT>""" <EOL> import numpy as np <EOL> from scipy import special <EOL> from mayavi import mlab <EOL> def base_vectors ( n ) : <EOL> """<STR_LIT>""" <EOL> n = n / ( n ** <NUM_LIT:2> ) . sum ( axis = - <NUM_LIT:1> ) <EOL> if np . abs ( n [ <NUM_LIT:0> ] ) == <NUM_LIT:1> : <EOL> l = np . r_ [ n [ <NUM_LIT:2> ] , <NUM_LIT:0> , - n [ <NUM_LIT:0> ] ] <EOL> else : <EOL> l = np . r_ [ <NUM_LIT:0> , n [ <NUM_LIT:2> ] , - n [ <NUM_LIT:1> ] ] <EOL> l = l / ( l ** <NUM_LIT:2> ) . sum ( axis = - <NUM_LIT:1> ) <EOL> m = np . cross ( n , l ) <EOL> return n , l , m <EOL> def magnetic_field ( r , n , r0 , R ) : <EOL> """<STR_LIT>""" <EOL> n , l , m = base_vectors ( n ) <EOL> trans = np . vstack ( ( l , m , n ) ) <EOL> inv_trans = np . linalg . inv ( trans ) <EOL> r = r - r0 <EOL> r = np . dot ( r , inv_trans ) <EOL> x = r [ : , <NUM_LIT:0> ] <EOL> y = r [ : , <NUM_LIT:1> ] <EOL> z = r [ : , <NUM_LIT:2> ] <EOL> rho = np . sqrt ( x ** <NUM_LIT:2> + y ** <NUM_LIT:2> ) <EOL> theta = np . arctan ( x / y ) <EOL> theta [ y == <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> E = special . ellipe ( ( <NUM_LIT:4> * R * rho ) / ( ( R + rho ) ** <NUM_LIT:2> + z ** <NUM_LIT:2> ) ) <EOL> K = special . ellipk ( ( <NUM_LIT:4> * R * rho ) / ( ( R + rho ) ** <NUM_LIT:2> + z ** <NUM_LIT:2> ) ) <EOL> Bz = <NUM_LIT:1> / np . sqrt ( ( R + rho ) ** <NUM_LIT:2> + z ** <NUM_LIT:2> ) * ( <EOL> K <EOL> + E * ( R ** <NUM_LIT:2> - rho ** <NUM_LIT:2> - z ** <NUM_LIT:2> ) / ( ( R - rho ) ** <NUM_LIT:2> + z ** <NUM_LIT:2> ) <EOL> ) <EOL> Brho = z / ( rho * np . sqrt ( ( R + rho ) ** <NUM_LIT:2> + z ** <NUM_LIT:2> ) ) * ( <EOL> - K <EOL> + E * ( R ** <NUM_LIT:2> + rho ** <NUM_LIT:2> + z ** <NUM_LIT:2> ) / ( ( R - rho ) ** <NUM_LIT:2> + z ** <NUM_LIT:2> ) <EOL> ) <EOL> Brho [ np . isnan ( Brho ) ] = <NUM_LIT:0> <EOL> Brho [ np . isinf ( Brho ) ] = <NUM_LIT:0> <EOL> Bz [ np . isnan ( Bz ) ] = <NUM_LIT:0> <EOL> Bz [ np . isinf ( Bz ) ] = <NUM_LIT:0> <EOL> B = np . c_ [ np . cos ( theta ) * Brho , np . sin ( theta ) * Brho , Bz ] <EOL> B = np . dot ( B , trans ) <EOL> return B <EOL> def display_coil ( n , r0 , R , half = False ) : <EOL> """<STR_LIT>""" <EOL> n , l , m = base_vectors ( n ) <EOL> theta = np . linspace ( <NUM_LIT:0> , ( <NUM_LIT:2> - half ) * np . pi , <NUM_LIT:30> ) <EOL> theta = theta [ ... , np . newaxis ] <EOL> coil = np . atleast_1d ( R ) * ( np . sin ( theta ) * l + np . cos ( theta ) * m ) <EOL> coil += r0 <EOL> coil_x = coil [ : , <NUM_LIT:0> ] <EOL> coil_y = coil [ : , <NUM_LIT:1> ] <EOL> coil_z = coil [ : , <NUM_LIT:2> ] <EOL> mlab . plot3d ( coil_x , coil_y , coil_z , <EOL> tube_radius = <NUM_LIT> , <EOL> name = '<STR_LIT>' % display_coil . num , <EOL> color = ( <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ) ) <EOL> display_coil . num += <NUM_LIT:1> <EOL> return coil_x , coil_y , coil_z <EOL> display_coil . num = <NUM_LIT:0> <EOL> X , Y , Z = np . mgrid [ - <NUM_LIT> : <NUM_LIT> : <NUM_LIT> , - <NUM_LIT> : <NUM_LIT> : <NUM_LIT> , - <NUM_LIT> : <NUM_LIT> : <NUM_LIT> ] <EOL> f = <NUM_LIT> <EOL> X = np . round ( X * f ) / f <EOL> Y = np . round ( Y * f ) / f <EOL> Z = np . round ( Z * f ) / f <EOL> r = np . c_ [ X . ravel ( ) , Y . ravel ( ) , Z . ravel ( ) ] <EOL> r0 = np . r_ [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0.1> ] <EOL> n = np . r_ [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:1> ] <EOL> R = <NUM_LIT:0.1> <EOL> r0 = np . vstack ( ( r0 , - r0 ) ) <EOL> R = np . r_ [ R , R ] <EOL> n = np . vstack ( ( n , n ) ) <EOL> B = np . empty_like ( r ) <EOL> for this_n , this_r0 , this_R in zip ( n , r0 , R ) : <EOL> this_n = np . array ( this_n ) <EOL> this_r0 = np . array ( this_r0 ) <EOL> this_R = np . array ( this_R ) <EOL> B += magnetic_field ( r , this_n , this_r0 , this_R ) <EOL> Bx = B [ : , <NUM_LIT:0> ] <EOL> By = B [ : , <NUM_LIT:1> ] <EOL> Bz = B [ : , <NUM_LIT:2> ] <EOL> Bx . shape = X . shape <EOL> By . shape = Y . shape <EOL> Bz . shape = Z . shape <EOL> Bnorm = np . sqrt ( Bx ** <NUM_LIT:2> + By ** <NUM_LIT:2> + Bz ** <NUM_LIT:2> ) <EOL> Bmax = <NUM_LIT:100> <EOL> Bx [ Bnorm > Bmax ] = <NUM_LIT:0> <EOL> By [ Bnorm > Bmax ] = <NUM_LIT:0> <EOL> Bz [ Bnorm > Bmax ] = <NUM_LIT:0> <EOL> Bnorm [ Bnorm > Bmax ] = Bmax <EOL> mlab . figure ( <NUM_LIT:1> , bgcolor = ( <NUM_LIT:1> , <NUM_LIT:1> , <NUM_LIT:1> ) , fgcolor = ( <NUM_LIT:0.5> , <NUM_LIT:0.5> , <NUM_LIT:0.5> ) , <EOL> size = ( <NUM_LIT> , <NUM_LIT> ) ) <EOL> mlab . clf ( ) <EOL> for this_n , this_r0 , this_R in zip ( n , r0 , R ) : <EOL> display_coil ( this_n , this_r0 , this_R ) <EOL> field = mlab . pipeline . vector_field ( X , Y , Z , Bx , By , Bz , <EOL> scalars = Bnorm , name = '<STR_LIT>' ) <EOL> vectors = mlab . pipeline . vectors ( field , <EOL> scale_factor = ( X [ <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:0> ] - X [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] ) , <EOL> ) <EOL> vectors . glyph . mask_input_points = True <EOL> vectors . glyph . mask_points . on_ratio = <NUM_LIT:6> <EOL> vcp = mlab . pipeline . vector_cut_plane ( field ) <EOL> vcp . glyph . glyph . scale_factor = <NUM_LIT:5> * ( X [ <NUM_LIT:1> , <NUM_LIT:0> , <NUM_LIT:0> ] - X [ <NUM_LIT:0> , <NUM_LIT:0> , <NUM_LIT:0> ] ) <EOL> iso = mlab . pipeline . iso_surface ( field , <EOL> contours = [ <NUM_LIT:0.1> * Bmax , <NUM_LIT> * Bmax ] , <EOL> opacity = <NUM_LIT> , <EOL> colormap = '<STR_LIT>' ) <EOL> iso . actor . property . frontface_culling = True <EOL> mlab . view ( <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , [ <NUM_LIT> , <NUM_LIT> , - <NUM_LIT> ] ) <EOL> mlab . show ( ) </s>
