from pylab import *

xpoints = arange(-6,6,0.05)
ypoints = arange(-6,6,0.05)
X,Y = meshgrid(xpoints, ypoints)

circularMask = False

areaRadius = 4

# order of the magnet
n = 2

def func(x,y):
    # the function to draw
    return ((x + 1j * y)**(n)).real

func = vectorize(func)

V = func(X,Y)


# mask points which we don't want to draw
if circularMask:
    # circular mask
    distance = sqrt(X**2 + Y**2)
    V = ma.masked_where(distance > areaRadius, V)
else:
    # polygonal mask

    # principal directions are at  (i + 0.5) / (2n) * 2pi
    # 
    for i in range(2*n):
        angle = (i + 0.5) / float(2*n) * 2*pi
    
        # define a straight angle perpendicular to angle
        # mask all points on one side of this line
        anchor_x = areaRadius * cos(angle)
        anchor_y = areaRadius * sin(angle)

        normal_x = cos(angle)
        normal_y = sin(angle)
        
        def acceptFunc(x,y):
            value = (x - anchor_x) * normal_x + (y - anchor_y) * normal_y
            return value > 0
        
        acceptFunc = vectorize(acceptFunc)
        
        V = ma.masked_where(acceptFunc(X,Y), V)        

    
    
if True:
    # levels equidistant in function value

    V /= V.max()

    levels = arange(-2,2,0.05)

else:
    # levels equidistant on x and y axis

    # determine the levels to draw from values on one of the axes

    levels = [ float(func(x,0)) for x in arange(min(xpoints), max(xpoints),0.50) ] + \
        [ float(func(0,y)) for y in arange(min(ypoints), max(ypoints),0.50) ]
    levels = sorted(list(set(levels)))

    
figure(figsize=(6,6)); 
Q = contour(X,Y, V, colors=  'black', linestyles = 'solid', 
    levels = levels
)
axis([-5,5,-5,5])
xlabel("x coordinate")
ylabel("y coordinate")

# mask points which we don't want to draw
if not circularMask:
    # polygonal mask

    # principal directions are at  (i + 0.5) / (2n) * 2pi
    # 
    for i in range(2*n):
        angle = (i + 0.5) / float(2*n) * 2*pi
        
        if i % 2:
            label = "N"
            color = 'red'
        else:
            label = "S"
            color = 'green'
        
        anchor_x = 1.1 * areaRadius * cos(angle)
        anchor_y = 1.1 * areaRadius * sin(angle)

        text(anchor_x, anchor_y, label, size = 20, color = color,
             horizontalalignment='center',
             verticalalignment='center')