import numpy as np
import matplotlib
import matplotlib.pyplot as plt

import sys
sys.path.append("../")
from quelea import *

nx = 217
ny = 133

x0 = 0
x1 = 30 # lambdas
y0 = 0
y1 = 20 # lambdas

xs = np.linspace(x0, x1, nx)
ys = np.linspace(y0, y1, ny)

# 2d array of (x, y, z, t)
coords = np.array( [ [x, y, 0, 0] for x in xs for y in ys ] )
# for map_fields function this should be converted from 2D to 1D array
coords = coords.reshape((4 * nx * ny,))

ftype = 1 # plane wave
a0 = 1 # normalized field amplitude
omega = 1 # frequency
fparam = [a0, 1, 0, 0, 0, 1, 0, 0, omega] # parameters of the plane wave

ex, ey, ez, bx, by, bz = map_fields(coords, ftype, fparam)
# now convert to 2d arrays
ex = ex.reshape((nx, ny))
ey = ey.reshape((nx, ny))
ez = ez.reshape((nx, ny))
bx = bx.reshape((nx, ny))
by = by.reshape((nx, ny))
bz = bz.reshape((nx, ny))
ex = ex.transpose()
ey = ey.transpose()
ez = ez.transpose()
bx = bx.transpose()
by = by.transpose()
bz = bz.transpose()

plt.imshow(ey, cmap = 'RdYlBu', origin = 'lower', extent = [x0, x1, y0, y1])
plt.colorbar()
plt.clim(-a0, a0)

plt.savefig("map_fields.pdf")


[((6, 0, 6, 22), 'sys.path.append', 'sys.path.append', ({(6, 16, 6, 21): '"""../"""'}, {}), "('../')", False, 'import sys\n'), ((17, 5, 17, 28), 'numpy.linspace', 'np.linspace', ({(17, 17, 17, 19): 'x0', (17, 21, 17, 23): 'x1', (17, 25, 17, 27): 'nx'}, {}), '(x0, x1, nx)', True, 'import numpy as np\n'), ((18, 5, 18, 28), 'numpy.linspace', 'np.linspace', ({(18, 17, 18, 19): 'y0', (18, 21, 18, 23): 'y1', (18, 25, 18, 27): 'ny'}, {}), '(y0, y1, ny)', True, 'import numpy as np\n'), ((21, 9, 21, 61), 'numpy.array', 'np.array', ({(21, 19, 21, 59): '[[x, y, 0, 0] for x in xs for y in ys]'}, {}), '([[x, y, 0, 0] for x in xs for y in ys])', True, 'import numpy as np\n'), ((45, 0, 45, 76), 'matplotlib.pyplot.imshow', 'plt.imshow', (), '', True, 'import matplotlib.pyplot as plt\n'), ((46, 0, 46, 14), 'matplotlib.pyplot.colorbar', 'plt.colorbar', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((47, 0, 47, 17), 'matplotlib.pyplot.clim', 'plt.clim', ({(47, 9, 47, 12): '(-a0)', (47, 14, 47, 16): 'a0'}, {}), '(-a0, a0)', True, 'import matplotlib.pyplot as plt\n'), ((49, 0, 49, 29), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(49, 12, 49, 28): '"""map_fields.pdf"""'}, {}), "('map_fields.pdf')", True, 'import matplotlib.pyplot as plt\n')]