# paste this code at the end of VectorFieldPlot 3.1
# https://commons.wikimedia.org/wiki/User:Geek3/VectorFieldPlot
doc = FieldplotDocument('VFPt_tilted-magnets-array_potential+contour',
    commons=True, width=800, height=600)

x0, y0 = 0, -1.7
phi = pi/4
dx = 2
R = 0.4
L2 = 1.3
m = 1
Nmag = 101
xarr = x0 + sc.arange(-(Nmag//2)*dx, ((Nmag+1)//2)*dx, dx)

discs = []
Q = m / (2 * L2)
for x in xarr:
    if fabs(x) <= 10:
        p0 = array([x, y0]) + rot([-L2,R], phi)
        p1 = array([x, y0]) + rot([-L2,-R], phi)
        discs.append(['charged_disc', {'x0':p0[0], 'y0':p0[1], 'x1':p1[0], 'y1':p1[1], 'Q':-Q}])
        p0 = array([x, y0]) + rot([L2,R], phi)
        p1 = array([x, y0]) + rot([L2,-R], phi)
        discs.append(['charged_disc', {'x0':p0[0], 'y0':p0[1], 'x1':p1[0], 'y1':p1[1], 'Q':Q}])
    else:
        # save computing time using simpler pole model for remote magnets
        p0 = array([x, y0]) + rot([-L2, 0], phi)
        discs.append(['monopole', {'x':p0[0], 'y':p0[1], 'Q':-Q}])
        p1 = array([x, y0]) + rot([L2, 0], phi)
        discs.append(['monopole', {'x':p1[0], 'y':p1[1], 'Q':Q}])

fieldH = Field(discs)
fieldB = Field([ ['coil', {'x':x, 'y':y0, 'phi':phi, 'R':R, 'Lhalf':L2,
    'I':m/(R**2*pi)}] for x in xarr])
field_symbols = Field([ ['coil', {'x':x, 'y':y0, 'phi':phi, 'R':R, 'Lhalf':L2,
    'I':m/(R**2*pi)}] for x in xarr if fabs(x) < 4 + L2])

doc.draw_magnets(field_symbols)
U0 = fieldH.V(array([x0, y0]) + rot([L2, 0], phi))

def bounds(xy):
    dmax = -1
    for i in range(Nmag):
        r = xy - array([xarr[i], y0])
        r = rot(r, -phi)
        dmax = max(dmax, min(1-fabs(r[0]/L2), 1-fabs(r[1]/R)))
    return dmax

nlines = 6
xoff = 0.1
for iline in range(Nmag * nlines):
    for y, di, s in (4, 'backward', 1), (2*y0-4, 'forward', -1):
        xstart = x0 + s * xoff + dx * (iline / nlines - Nmag // 2)
        if fabs(xstart) < 4.5:
            p0 = [xstart, y]
            line = FieldLine(fieldH, p0, directions=di, maxr=8.,
                bounds_func=bounds)
            doc.draw_line(line, linewidth=2.4, arrows_style=
                {'at_potentials':[-0.4 * U0, 0.23 * U0], 'potential':fieldH.V})

nlines2 = 12
for imag in range(Nmag):
    xmag = dx * (imag - Nmag // 2)
    for iline in range(nlines2):
        a = (iline + 0.5) / nlines2
        a += -0.4 * (((2 * a - 1)**3 + 1) / 2 - a)
        p1 = rot([-0.36*L2, -R], phi)
        p2 = array([dx, 0]) + rot([0.36*L2, R], phi)
        xstart = xmag + p1[0] + a * (p2[0] - p1[0])
        ystart = y0 + p1[1] + a * (p2[1] - p1[1])
        if fabs(xstart) < 4.5:
            line = FieldLine(fieldH, [xstart, ystart], directions='both', maxr=2*L2,
                stop_funcs=2*[bounds])
            doc.draw_line(line, linewidth=2.4, arrows_style=
                {'max_arrows':1, 'min_arrows':1})

print('computing scalar field.')
doc.draw_scalar_field(func=fieldH.V, cmap=doc.cmap_AqYlFs, vmin=-U0, vmax=U0)
doc.draw_contours(func=fieldH.V, linewidth=1, linecolor='#111111',
    levels=sc.linspace(-U0, U0, 17)[1:-1], attributes={'opacity':'0.7'})

doc.write()