xpoints = arange(-5,5,0.05)
ypoints = arange(-5,5,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
    vmax = V.max()
    V /= vmax

    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)))

    vmax = 1
    
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')

#----------------------------------------
if n == 2:
    # quadrupole, draw some examples of force on charged particle

    # find kth level line on axes (x = 0 and y = 0)

    # the potential function is >= 0 on the x axis and <= 0 on the y axis
    # for a quadrupole
    lev = sorted(list(levels[levels >= 0]))[4]

    # find distance of this level on axis from origin
    # (exploit the 90 degree symmetry of the field)
    dist = fsolve(lambda x: func(x,0) / vmax - lev,3)[0]
    
    # rotation by +90 degrees
    rotMatrix = array([[0,-1],[1,0]])
    invRotMatrix = rotMatrix.T
    
    arrowLength = 1.5
    
    arrowStart = array([dist, 0])
    origArrowDir = array([0, arrowLength])
    
    bfieldLabelPosOffset = array([0.3, 0.5 * arrowLength])
    forceLabelPosOffset = array([0.5 * arrowLength, -0.3])
    
    for i in range(4):
        
        arrowDir = origArrowDir[:]
        for j in range(i):
            arrowDir = rotMatrix.dot(arrowDir)
        
        # take into account quadrupole structure
        arrowDir *= (-1)**i
        
        # draw arrow for the B field
        arrow(arrowStart[0],arrowStart[1], arrowDir[0], arrowDir[1],head_width=0.3, head_length=0.3, color = 'red')
    
        # add a label for the B field
        textPos = arrowStart + (-1)**i * bfieldLabelPosOffset
        text(textPos[0], 
             textPos[1], "B", size = 20, color = 'red',
             horizontalalignment='center',
             verticalalignment='center')

        # draw the arrow for the force
        arrowDir2 = invRotMatrix.dot(arrowDir)
    
        arrow(arrowStart[0], arrowStart[1], arrowDir2[0], arrowDir2[1],head_width=0.3, head_length=0.3, color = 'blue')
        
        # label for the force
        textPos = arrowStart + (-1)**i * forceLabelPosOffset
        text(textPos[0], 
             textPos[1], "F", size = 20, color = 'blue',
             horizontalalignment='center',
             verticalalignment='center')
        #----------
        # prepare next iteration
        arrowStart = rotMatrix.dot(arrowStart)
        
        bfieldLabelPosOffset = rotMatrix.dot(bfieldLabelPosOffset)
        forceLabelPosOffset = rotMatrix.dot(forceLabelPosOffset)