def assembleMvv(A, B, M, P):
#A ?
#B ?
#M Mesh object
#P Physics object
nodes = 4
sa0, ma0 = localStiffness(M)
etabflat = M.etab.reshape(M.nelx * M.nelz, nodes)
Aflat = A.flatten()[pl.newaxis].T
Bflat = B.flatten()[pl.newaxis].T
inew = pl.kron(etabflat, pl.ones((1, nodes))).flatten()
jnew = pl.kron(etabflat, pl.ones((nodes, 1))).flatten()
dnew = pl.dot(Aflat, sa0.reshape(
(1, nodes * nodes))).flatten() - P.k0**2 * pl.dot(
Bflat, ma0.reshape((1, nodes * nodes))).flatten()
i = list(inew)
j = list(jnew)
d = list(dnew)
#Periodic Bloch-Floquet BC:
pen = BCpen
n1 = list(range(0, (M.nelz + 1) * (M.nelx + 1), M.nelx + 1))
n2 = list(range(M.nelx, (M.nelz + 1) * (M.nelx + 1) + M.nelx, M.nelx + 1))
i += n1 + n2 + n1 + n2
j += n1 + n2 + n2 + n1
h = [pen]*2*len(n1)+\
[-pen*exp(1.j*P.kInx*M.lx)]*len(n1)+\
[-pen*exp(1.j*P.kInx*M.lx).conj()]*len(n1)
d += [pen]*2*len(n1)+\
[-pen*exp(1.j*P.kInx*M.lx)]*len(n1)+\
[-pen*exp(1.j*P.kInx*M.lx).conj()]*len(n1)
#Calculate the matrices given in Eq (43) in Dossou2006. See Fuchi2010 for a "nicer" way
# of writing them. The elements used are the same as in Andreassen2011
Mvv = coo_matrix((d, (i, j)), shape=(M.ndof, M.ndof),
dtype='complex').toarray()
return Mvv
def make_tec_screen_plots(pp, tec_screen, residuals, station_positions,
source_names, times, height, order, beta_val, r_0, prefix = 'frame_',
remove_gradient=True, show_source_names=False, min_tec=None, max_tec=None):
"""Makes plots of TEC screens
Keyword arguments:
pp -- array of piercepoint locations
tec_screen -- array of TEC screen values at the piercepoints
residuals -- array of TEC screen residuals at the piercepoints
source_names -- array of source names
times -- array of times
height -- height of screen (m)
order -- order of screen (e.g., number of KL base vectors to keep)
r_0 -- scale size of phase fluctuations (m)
beta_val -- power-law index for phase structure function (5/3 =>
pure Kolmogorov turbulence)
prefix -- prefix for output file names
remove_gradient -- fit and remove a gradient from each screen
show_source_names -- label sources on screen plots
min_tec -- minimum TEC value for plot range
max_tec -- maximum TEC value for plot range
"""
from pylab import kron, concatenate, pinv, norm, newaxis, normalize
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import numpy as np
import os
from operations.tecscreen import calc_piercepoint
import progressbar
root_dir = os.path.dirname(prefix)
if root_dir == '':
root_dir = './'
prestr = os.path.basename(prefix) + 'screen_'
try:
os.makedirs(root_dir)
except OSError:
pass
N_stations = station_positions.shape[0]
N_sources = len(source_names)
N_times = len(times)
A = concatenate([kron(np.eye(N_sources),
np.ones((N_stations,1))), kron(np.ones((N_sources,1)),
np.eye(N_stations))], axis=1)
N_piercepoints = N_sources * N_stations
P = np.eye(N_piercepoints) - np.dot(np.dot(A, pinv(np.dot(A.T, A))), A.T)
x, y, z = station_positions[0, :]
east = np.array([-y, x, 0])
east = east / norm(east)
north = np.array([ -x, -y, (x*x + y*y)/z])
north = north / norm(north)
up = np.array([x ,y, z])
up = up / norm(up)
T = concatenate([east[:, newaxis], north[:, newaxis]], axis=1)
pp1 = np.dot(pp[0, :, :], T)
lower = np.amin(pp1, axis=0)
upper = np.amax(pp1, axis=0)
extent = upper - lower
lower = lower - 0.05 * extent
upper = upper + 0.05 * extent
extent = upper - lower
Nx = 25
Ny = int(extent[1] / extent[0] * np.float(Nx))
xr = np.arange(lower[0], upper[0], extent[0]/Nx)
yr = np.arange(lower[1], upper[1], extent[1]/Ny)
screen = np.zeros((Nx, Ny, N_times))
gradient = np.zeros((Nx, Ny, N_times))
residuals = residuals.transpose([0, 2, 1]).reshape(N_piercepoints, N_times)
fitted_tec1 = tec_screen.transpose([0, 2, 1]).reshape(N_piercepoints, N_times) + residuals
logging.info('Calculating TEC screen images...')
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
D = np.resize(pp[k, :, :], (N_piercepoints, N_piercepoints, 3))
D = np.transpose(D, (1, 0, 2)) - D
D2 = np.sum(D**2, axis=2)
C = -(D2 / r_0**2)**(beta_val / 2.0) / 2.0
f = np.dot(pinv(C), tec_screen[:, k, :].reshape(N_piercepoints))
for i, x in enumerate(xr[0: Nx]):
for j, y in enumerate(yr[0: Ny]):
p = calc_piercepoint(np.dot(np.array([x, y]), np.array([east, north])), up, height)
d2 = np.sum(np.square(pp[k, :, :] - p[0]), axis=1)
c = -(d2 / ( r_0**2 ))**(beta_val / 2.0) / 2.0
screen[i, j, k] = np.dot(c, f)
# Fit and remove a gradient.
# Plot gradient in lower-left corner with its own color bar?
if remove_gradient:
xs, ys = np.indices(screen.shape[0:2])
zs = screen[:, :, k]
XYZ = []
for xf, yf in zip(xs.flatten().tolist(), ys.flatten().tolist()):
XYZ.append([xf, yf, zs[xf, yf]])
XYZ = np.array(XYZ)
a, b, c = fitPLaneLTSQ(XYZ)
grad_plane = a * xs + b * ys + c
gradient[:, :, k] = grad_plane
screen[:, :, k] = screen[:, :, k] - grad_plane
screen[:, :, k] = screen[:, :, k] - np.mean(screen[:, :, k])
for t in range(fitted_tec1.shape[0]):
xs_pt = np.where(np.array(xr) > pp1[t, 0])[0][0]
ys_pt = np.where(np.array(yr) > pp1[t, 1])[0][0]
grad_plane_pt = a * xs_pt + b * ys_pt + c
fitted_tec1[t, k] = fitted_tec1[t, k] - grad_plane_pt
fitted_tec1[:, k] = fitted_tec1[:, k] - np.mean(fitted_tec1[:, k])
pbar.update(ipbar)
ipbar += 1
pbar.finish()
if min_tec is None:
vmin = np.min([np.amin(screen), np.amin(fitted_tec1)])
else:
vmin = min_tec
if max_tec is None:
vmax = np.max([np.amax(screen), np.amax(fitted_tec1)])
else:
vmax = max_tec
logging.info('Plotting TEC screens...')
fig, ax = plt.subplots(figsize=[7, 7])
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
plt.clf()
im = plt.imshow(screen.transpose([1, 0, 2])[:, :, k],
cmap = plt.cm.jet,
origin = 'lower',
interpolation = 'nearest',
extent = (xr[0]/1000.0, xr[-1]/1000.0, yr[0]/1000.0, yr[-1]/1000.0),
vmin=vmin, vmax=vmax)
sm = plt.cm.ScalarMappable(cmap=plt.cm.jet,
norm=normalize(vmin=vmin, vmax=vmax))
sm._A = []
cbar = plt.colorbar(im)
cbar.set_label('TECU', rotation=270)
x = []
y = []
s = []
c = []
for j in range(fitted_tec1.shape[0]):
x.append(pp1[j, 0] / 1000.0)
y.append(pp1[j, 1] / 1000.0)
xs = np.where(np.array(xr) > pp1[j, 0])[0][0]
ys = np.where(np.array(yr) > pp1[j, 1])[0][0]
fit_screen_diff = abs(fitted_tec1[j, k] - screen[xs, ys, k])
s.append(max(20*fit_screen_diff/0.01, 10))
c.append(sm.to_rgba(fitted_tec1[j, k]))
plt.scatter(x, y, s=s, c=c)
if show_source_names:
labels = source_names
for label, xl, yl in zip(labels, x[0::N_stations], y[0::N_stations]):
plt.annotate(
label,
xy = (xl, yl), xytext = (-2, 2),
textcoords = 'offset points', ha = 'right', va = 'bottom')
plt.title('Screen {0}'.format(k))
plt.xlim(xr[-1]/1000.0, xr[0]/1000.0)
plt.ylim(yr[0]/1000.0, yr[-1]/1000.0)
plt.xlabel('Projected Distance along RA (km)')
plt.ylabel('Projected Distance along Dec (km)')
if remove_gradient:
axins = inset_axes(ax, width="15%", height="10%", loc=2)
axins.imshow(gradient.transpose([1, 0, 2])[:, : ,k],
cmap = plt.cm.jet,
origin = 'lower',
interpolation = 'nearest',
extent = (xr[0]/1000.0, xr[-1]/1000.0, yr[0]/1000.0, yr[-1]/1000.0),
vmin=vmin, vmax=vmax)
plt.xticks(visible=False)
plt.yticks(visible=False)
axins.set_xlim(xr[-1]/1000.0, xr[0]/1000.0)
axins.set_ylim(yr[0]/1000.0, yr[-1]/1000.0)
plt.savefig(root_dir+'/'+prestr+'frame%0.3i.png' % k)
pbar.update(ipbar)
ipbar += 1
pbar.finish()
plt.close(fig)
def fit_screen_to_tec(station_names, source_names, pp, airmass, rr, times,
height, order, r_0, beta):
"""
Fits a screen to given TEC values using Karhunen-Lo`eve base vectors
Keyword arguments:
station_names -- array of station names
source_names -- array of source names
pp -- array of piercepoint locations
airmass -- array of airmass values
rr -- array of TEC solutions
times -- array of times
height -- height of screen (m)
order -- order of screen (i.e., number of KL base vectors to keep)
r_0 -- scale size of phase fluctuations (m)
beta -- power-law index for phase structure function (5/3 =>
pure Kolmogorov turbulence)
"""
import numpy as np
from pylab import kron, concatenate, pinv, norm, newaxis, find, amin, svd, eye
try:
import progressbar
except ImportError:
import losoto.progressbar as progressbar
logging.info('Fitting screens to TEC values...')
N_stations = len(station_names)
N_sources = len(source_names)
N_times = len(times)
tec_fit_all = np.zeros((N_times, N_sources, N_stations))
residual_all = np.zeros((N_times, N_sources, N_stations))
A = concatenate([
kron(eye(N_sources), np.ones((N_stations, 1))),
kron(np.ones((N_sources, 1)), eye(N_stations))
],
axis=1)
N_piercepoints = N_sources * N_stations
P = eye(N_piercepoints) - np.dot(np.dot(A, pinv(np.dot(A.T, A))), A.T)
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
try:
D = np.resize(pp[k, :, :], (N_piercepoints, N_piercepoints, 3))
D = np.transpose(D, (1, 0, 2)) - D
D2 = np.sum(D**2, axis=2)
C = -(D2 / r_0**2)**(beta / 2.0) / 2.0
P1 = eye(N_piercepoints) - np.ones(
(N_piercepoints, N_piercepoints)) / N_piercepoints
C1 = np.dot(np.dot(P1, C), P1)
U, S, V = svd(C1)
B = np.dot(P, np.dot(np.diag(airmass[k, :]), U[:, :order]))
pinvB = pinv(B, rcond=1e-3)
rr1 = np.dot(P, rr[:, k])
tec_fit = np.dot(U[:, :order], np.dot(pinvB, rr1))
tec_fit_all[k, :, :] = tec_fit.reshape((N_sources, N_stations))
residual = rr1 - np.dot(P, tec_fit)
residual_all[k, :, :] = residual.reshape((N_sources, N_stations))
except:
# Set screen to zero if fit did not work
logging.debug('Tecscreen fit failed for timeslot {0}'.format(k))
tec_fit_all[k, :, :] = np.zeros((N_sources, N_stations))
residual_all[k, :, :] = np.ones((N_sources, N_stations))
pbar.update(ipbar)
ipbar += 1
pbar.finish()
return tec_fit_all, residual_all
def fit_screen_to_tec(station_names, source_names, pp, airmass, rr, times,
height, order, r_0, beta):
"""
Fits a screen to given TEC values using Karhunen-Lo`eve base vectors
Keyword arguments:
station_names -- array of station names
source_names -- array of source names
pp -- array of piercepoint locations
airmass -- array of airmass values
rr -- array of TEC solutions
times -- array of times
height -- height of screen (m)
order -- order of screen (i.e., number of KL base vectors to keep)
r_0 -- scale size of phase fluctuations (m)
beta -- power-law index for phase structure function (5/3 =>
pure Kolmogorov turbulence)
"""
import numpy as np
from pylab import kron, concatenate, pinv, norm, newaxis, find, amin, svd, eye
try:
import progressbar
except ImportError:
import losoto.progressbar as progressbar
logging.info('Fitting screens to TEC values...')
N_stations = len(station_names)
N_sources = len(source_names)
N_times = len(times)
tec_fit_all = np.zeros((N_times, N_sources, N_stations))
residual_all = np.zeros((N_times, N_sources, N_stations))
A = concatenate([kron(eye(N_sources), np.ones((N_stations, 1))),
kron(np.ones((N_sources, 1)), eye(N_stations))], axis=1)
N_piercepoints = N_sources * N_stations
P = eye(N_piercepoints) - np.dot(np.dot(A, pinv(np.dot(A.T, A))), A.T)
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
try:
D = np.resize(pp[k, :, :], (N_piercepoints, N_piercepoints, 3))
D = np.transpose(D, (1, 0, 2)) - D
D2 = np.sum(D**2, axis=2)
C = -(D2 / r_0**2)**(beta / 2.0) / 2.0
P1 = eye(N_piercepoints) - np.ones((N_piercepoints, N_piercepoints)) / N_piercepoints
C1 = np.dot(np.dot(P1, C), P1)
U, S, V = svd(C1)
B = np.dot(P, np.dot(np.diag(airmass[k, :]), U[:, :order]))
pinvB = pinv(B, rcond=1e-3)
rr1 = np.dot(P, rr[:, k])
tec_fit = np.dot(U[:, :order], np.dot(pinvB, rr1))
tec_fit_all[k, :, :] = tec_fit.reshape((N_sources, N_stations))
residual = rr1 - np.dot(P, tec_fit)
residual_all[k, :, :] = residual.reshape((N_sources, N_stations))
except:
# Set screen to zero if fit did not work
logging.debug('Tecscreen fit failed for timeslot {0}'.format(k))
tec_fit_all[k, :, :] = np.zeros((N_sources, N_stations))
residual_all[k, :, :] = np.ones((N_sources, N_stations))
pbar.update(ipbar)
ipbar += 1
pbar.finish()
return tec_fit_all, residual_all
def make_tec_screen_plots(pp,
tec_screen,
residuals,
station_positions,
source_names,
times,
height,
order,
beta_val,
r_0,
prefix='frame_',
remove_gradient=True,
show_source_names=False,
min_tec=None,
max_tec=None):
"""Makes plots of TEC screens
Keyword arguments:
pp -- array of piercepoint locations
tec_screen -- array of TEC screen values at the piercepoints
residuals -- array of TEC screen residuals at the piercepoints
source_names -- array of source names
times -- array of times
height -- height of screen (m)
order -- order of screen (e.g., number of KL base vectors to keep)
r_0 -- scale size of phase fluctuations (m)
beta_val -- power-law index for phase structure function (5/3 =>
pure Kolmogorov turbulence)
prefix -- prefix for output file names
remove_gradient -- fit and remove a gradient from each screen
show_source_names -- label sources on screen plots
min_tec -- minimum TEC value for plot range
max_tec -- maximum TEC value for plot range
"""
from pylab import kron, concatenate, pinv, norm, newaxis, normalize
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import numpy as np
import os
from operations.tecscreen import calc_piercepoint
import progressbar
root_dir = os.path.dirname(prefix)
if root_dir == '':
root_dir = './'
prestr = os.path.basename(prefix) + 'screen_'
try:
os.makedirs(root_dir)
except OSError:
pass
N_stations = station_positions.shape[0]
N_sources = len(source_names)
N_times = len(times)
A = concatenate([
kron(np.eye(N_sources), np.ones((N_stations, 1))),
kron(np.ones((N_sources, 1)), np.eye(N_stations))
],
axis=1)
N_piercepoints = N_sources * N_stations
P = np.eye(N_piercepoints) - np.dot(np.dot(A, pinv(np.dot(A.T, A))), A.T)
x, y, z = station_positions[0, :]
east = np.array([-y, x, 0])
east = east / norm(east)
north = np.array([-x, -y, (x * x + y * y) / z])
north = north / norm(north)
up = np.array([x, y, z])
up = up / norm(up)
T = concatenate([east[:, newaxis], north[:, newaxis]], axis=1)
pp1 = np.dot(pp[0, :, :], T)
lower = np.amin(pp1, axis=0)
upper = np.amax(pp1, axis=0)
extent = upper - lower
lower = lower - 0.05 * extent
upper = upper + 0.05 * extent
extent = upper - lower
Nx = 25
Ny = int(extent[1] / extent[0] * np.float(Nx))
xr = np.arange(lower[0], upper[0], extent[0] / Nx)
yr = np.arange(lower[1], upper[1], extent[1] / Ny)
screen = np.zeros((Nx, Ny, N_times))
gradient = np.zeros((Nx, Ny, N_times))
residuals = residuals.transpose([0, 2, 1]).reshape(N_piercepoints, N_times)
fitted_tec1 = tec_screen.transpose([0, 2, 1]).reshape(
N_piercepoints, N_times) + residuals
logging.info('Calculating TEC screen images...')
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
D = np.resize(pp[k, :, :], (N_piercepoints, N_piercepoints, 3))
D = np.transpose(D, (1, 0, 2)) - D
D2 = np.sum(D**2, axis=2)
C = -(D2 / r_0**2)**(beta_val / 2.0) / 2.0
f = np.dot(pinv(C), tec_screen[:, k, :].reshape(N_piercepoints))
for i, x in enumerate(xr[0:Nx]):
for j, y in enumerate(yr[0:Ny]):
p = calc_piercepoint(
np.dot(np.array([x, y]), np.array([east, north])), up,
height)
d2 = np.sum(np.square(pp[k, :, :] - p[0]), axis=1)
c = -(d2 / (r_0**2))**(beta_val / 2.0) / 2.0
screen[i, j, k] = np.dot(c, f)
# Fit and remove a gradient.
# Plot gradient in lower-left corner with its own color bar?
if remove_gradient:
xs, ys = np.indices(screen.shape[0:2])
zs = screen[:, :, k]
XYZ = []
for xf, yf in zip(xs.flatten().tolist(), ys.flatten().tolist()):
XYZ.append([xf, yf, zs[xf, yf]])
XYZ = np.array(XYZ)
a, b, c = fitPLaneLTSQ(XYZ)
grad_plane = a * xs + b * ys + c
gradient[:, :, k] = grad_plane
screen[:, :, k] = screen[:, :, k] - grad_plane
screen[:, :, k] = screen[:, :, k] - np.mean(screen[:, :, k])
for t in range(fitted_tec1.shape[0]):
xs_pt = np.where(np.array(xr) > pp1[t, 0])[0][0]
ys_pt = np.where(np.array(yr) > pp1[t, 1])[0][0]
grad_plane_pt = a * xs_pt + b * ys_pt + c
fitted_tec1[t, k] = fitted_tec1[t, k] - grad_plane_pt
fitted_tec1[:, k] = fitted_tec1[:, k] - np.mean(fitted_tec1[:, k])
pbar.update(ipbar)
ipbar += 1
pbar.finish()
if min_tec is None:
vmin = np.min([np.amin(screen), np.amin(fitted_tec1)])
else:
vmin = min_tec
if max_tec is None:
vmax = np.max([np.amax(screen), np.amax(fitted_tec1)])
else:
vmax = max_tec
logging.info('Plotting TEC screens...')
fig, ax = plt.subplots(figsize=[7, 7])
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
plt.clf()
im = plt.imshow(screen.transpose([1, 0, 2])[:, :, k],
cmap=plt.cm.jet,
origin='lower',
interpolation='nearest',
extent=(xr[0] / 1000.0, xr[-1] / 1000.0,
yr[0] / 1000.0, yr[-1] / 1000.0),
vmin=vmin,
vmax=vmax)
sm = plt.cm.ScalarMappable(cmap=plt.cm.jet,
norm=normalize(vmin=vmin, vmax=vmax))
sm._A = []
cbar = plt.colorbar(im)
cbar.set_label('TECU', rotation=270)
x = []
y = []
s = []
c = []
for j in range(fitted_tec1.shape[0]):
x.append(pp1[j, 0] / 1000.0)
y.append(pp1[j, 1] / 1000.0)
xs = np.where(np.array(xr) > pp1[j, 0])[0][0]
ys = np.where(np.array(yr) > pp1[j, 1])[0][0]
fit_screen_diff = abs(fitted_tec1[j, k] - screen[xs, ys, k])
s.append(max(20 * fit_screen_diff / 0.01, 10))
c.append(sm.to_rgba(fitted_tec1[j, k]))
plt.scatter(x, y, s=s, c=c)
if show_source_names:
labels = source_names
for label, xl, yl in zip(labels, x[0::N_stations],
y[0::N_stations]):
plt.annotate(label,
xy=(xl, yl),
xytext=(-2, 2),
textcoords='offset points',
ha='right',
va='bottom')
plt.title('Screen {0}'.format(k))
plt.xlim(xr[-1] / 1000.0, xr[0] / 1000.0)
plt.ylim(yr[0] / 1000.0, yr[-1] / 1000.0)
plt.xlabel('Projected Distance along RA (km)')
plt.ylabel('Projected Distance along Dec (km)')
if remove_gradient:
axins = inset_axes(ax, width="15%", height="10%", loc=2)
axins.imshow(gradient.transpose([1, 0, 2])[:, :, k],
cmap=plt.cm.jet,
origin='lower',
interpolation='nearest',
extent=(xr[0] / 1000.0, xr[-1] / 1000.0,
yr[0] / 1000.0, yr[-1] / 1000.0),
vmin=vmin,
vmax=vmax)
plt.xticks(visible=False)
plt.yticks(visible=False)
axins.set_xlim(xr[-1] / 1000.0, xr[0] / 1000.0)
axins.set_ylim(yr[0] / 1000.0, yr[-1] / 1000.0)
plt.savefig(root_dir + '/' + prestr + 'frame%0.3i.png' % k)
pbar.update(ipbar)
ipbar += 1
pbar.finish()
plt.close(fig)
def make_tec_screen_plots(pp, rr, tec_fit_white, tec_fit, station_positions,
source_names, times, height, order, beta_val, r_0, prefix = 'frame_',
remove_gradient=True):
"""Makes plots of TEC screens"""
import pylab
import numpy as np
import os
from operations.tecscreen import calc_piercepoint
import progressbar
root_dir = os.path.dirname(prefix)
if root_dir == '':
root_dir = './'
prestr = os.path.basename(prefix)
try:
os.makedirs(root_dir)
except OSError:
pass
N_stations = station_positions.shape[0]
N_sources = len(source_names)
N_times = len(times)
A = pylab.concatenate([pylab.kron(np.eye(N_sources),
np.ones((N_stations,1))), pylab.kron(np.ones((N_sources,1)),
np.eye(N_stations))], axis=1)
N_piercepoints = N_sources * N_stations
P = np.eye(N_piercepoints) - np.dot(np.dot(A, pylab.pinv(np.dot(A.T, A))), A.T)
x,y,z = station_positions[0, :]
east = np.array([-y, x, 0])
east = east / pylab.norm(east)
north = np.array([ -x, -y, (x*x + y*y)/z])
north = north / pylab.norm(north)
up = np.array([x,y,z])
up = up / pylab.norm(up)
T = pylab.concatenate([east[:, pylab.newaxis], north[:, pylab.newaxis]], axis=1)
pp1 = np.dot(pp[0, :, :], T)
lower = np.amin(pp1, axis=0)
upper = np.amax(pp1, axis=0)
extent = upper - lower
lower = lower - 0.05 * extent
upper = upper + 0.05 * extent
extent = upper - lower
N = 50
xr = np.arange(lower[0], upper[0], extent[0]/N)
yr = np.arange(lower[1], upper[1], extent[1]/N)
screen = np.zeros((N, N, N_times))
fitted_tec1 = tec_fit.transpose([0, 2, 1]).reshape(N_piercepoints, N_times) + np.dot(P, rr-tec_fit.transpose([0, 2, 1]).reshape(N_piercepoints, N_times))
logging.info('Calculating TEC screen images...')
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
f = tec_fit_white[:, k, :]
for i, x in enumerate(xr[0: N]):
for j, y in enumerate(yr[0: N]):
p = calc_piercepoint(np.dot(np.array([x, y]), np.array([east, north])), up, height)
d2 = np.sum(np.square(pp[k, :, :] - p[0]), axis=1)
c = -(d2 / ( r_0**2 ) )**( beta_val / 2.0 ) / 2.0
screen[j, i, k] = np.dot(c, f.reshape(N_piercepoints))
# Fit and remove a gradient
if remove_gradient:
xs, ys = np.indices(screen.shape[0:2])
zs = screen[:, :, k]
XYZ = []
for xf, yf in zip(xs.flatten().tolist(), ys.flatten().tolist()):
XYZ.append([xf, yf, zs[xf, yf]])
XYZ = np.array(XYZ)
a, b, c = fitPLaneLTSQ(XYZ)
grad_plane = a * xs + b * ys + c
screen[:, :, k] = screen[:, :, k] - grad_plane
for t in range(fitted_tec1.shape[0]):
xs_pt = np.where(np.array(xr) > pp1[t, 0])[0][0]
ys_pt = np.where(np.array(yr) > pp1[t, 1])[0][0]
grad_plane_pt = a * ys_pt + b * xs_pt + c
fitted_tec1[t, k] = fitted_tec1[t, k] - grad_plane_pt
pbar.update(ipbar)
ipbar += 1
pbar.finish()
vmin = np.min([np.amin(screen), np.amin(fitted_tec1)])
vmax = np.max([np.amax(screen), np.amax(fitted_tec1)])
logging.info('Plotting TEC screens...')
fig1 = pylab.figure(figsize = (7, 7))
pbar = progressbar.ProgressBar(maxval=N_times).start()
ipbar = 0
for k in range(N_times):
pylab.clf()
im = pylab.imshow(screen[:, :, k],
cmap = pylab.cm.jet,
origin = 'lower',
interpolation = 'nearest',
extent = (xr[0], xr[-1], yr[0], yr[-1]), vmin = vmin, vmax = vmax)
sm = pylab.cm.ScalarMappable(cmap = pylab.cm.jet,
norm = pylab.normalize(vmin = vmin, vmax=vmax))
sm._A = []
pylab.title(str(i))
cbar = pylab.colorbar()
cbar.set_label('TECU', rotation=270)
x = []
y = []
s = []
c = []
for j in range(fitted_tec1.shape[0]):
x.append(pp1[j,0])
y.append(pp1[j,1])
xs = np.where(np.array(xr) > pp1[j,0])[0][0]
ys = np.where(np.array(yr) > pp1[j,1])[0][0]
s.append(max(2400*abs(fitted_tec1[j, k] - screen[ys, xs, k]), 10))
c.append(sm.to_rgba(fitted_tec1[j, k]))
pylab.scatter(x, y, s=s, c=c)
labels = source_names
for label, xl, yl in zip(labels, x[0::N_stations], y[0::N_stations]):
pylab.annotate(
label,
xy = (xl, yl), xytext = (-20, 20),
textcoords = 'offset points', ha = 'right', va = 'bottom',
bbox = dict(boxstyle = 'round,pad=0.5', fc = 'gray', alpha = 0.5),
arrowprops = dict(arrowstyle = '->', connectionstyle = 'arc3,rad=0'))
pylab.title('Time {0}'.format(k))
pylab.xlim(xr[-1], xr[0])
pylab.ylim(yr[0], yr[-1])
pylab.xlabel('Projected Distance (m)')
pylab.ylabel('Projected Distance (m)')
pylab.savefig(root_dir+'/'+prestr+'frame%0.3i.png' % k)
pbar.update(ipbar)
ipbar += 1
pbar.finish()
pylab.close(fig1)