def test_cart2pol():
theta = np.pi/3.
phi = np.pi/3.
x = 1/2. * np.sqrt(3/4.)
y = 3/4.
z = 1/2.
res = cart2pol(np.array((x,y,z)))
assert_almost_equal(np.array((theta, phi)), np.array(res))
def test_coordinate_conversions():
n_tests = int(1e3)
for i in xrange(n_tests):
theta = np.random.uniform(low = 0., high = np.pi)
phi = np.random.uniform(low = 0., high = 2 * np.pi)
cart = pol2cart(np.array([theta, phi]))
theta_prime, phi_prime = cart2pol(cart)
assert_almost_equal(np.array((theta, phi % (2 * np.pi))),
np.array((theta_prime, phi_prime % (2 * np.pi))))
def _map_data(data, task, pd, resolution=100):
'''
Parameters
----------
data : ndarray
shape (ntime, ntask)
task : ndarray
shape (ntask, 6)
pd : ndarray
shape (3,)
Returns
-------
mapped_data : ndarray
shape (ntime, resolution, resolution)
Notes
-----
need to incorporate PD so that it lies at center of plot
'''
if data.shape[-1] != task.shape[0]:
raise ValueError('last axis of data (len %d) must be same length as\n'
'first axis of task (len %d)' % \
(data.shape[-1], data.shape[0]))
# generate theta, phi grid
# or rather map x,y grid to nearest of 26 targets
thetas = np.linspace(0., np.pi, resolution)
phis = np.linspace(0., 2 * np.pi, resolution)
theta_grid = np.tile(thetas[:,None], (1, resolution))
phi_grid = np.tile(phis[None,:], (resolution, 1))
# convert polar to cartesian co-ordinates
tp_grid = np.concatenate((theta_grid[None], phi_grid[None]), axis=0)
xyz_grid = pol2cart(tp_grid, axis=0)
# calculate direction of each task
start = task[:,:3]
stop = task[:,3:]
direction_task = unitvec(stop - start, axis=-1)
# rotate task directions until pd points towards 0,0,1
rotated_toz = rotate_targets(direction_task, cart2pol(pd))
# now rotate again to point pd towards 0,1,0
rotated_toy = rotate_targets(rotated_toz, np.array([np.pi/2., 0]))
# calculate angle between each grid square and each direction
angles = np.arccos(np.tensordot(rotated_toy, xyz_grid, [1,0]))
# get index of closest direction for each grid square
# = get argmin along 0th axis
nearest = np.argmin(angles, axis=0)
mapped_data = data[..., nearest]
return mapped_data
def plot_circle(self, theta, phi, angle, resolution=100.,
color='next', inc_color=True, **kwargs):
# calculate points in circle, in X,Y,Z
x, y, z = 0, 1, 2
rho, psi = 0, 1
P_c = generate_cone_circle(theta, phi, angle)
#P_c = P_c # [:-1]
if angle > np.pi/2.:
symbol = '--'
else:
symbol = '-'
if color == 'next':
color = self.next_colour(inc=inc_color)
# divide circle points into each hemisphere
top_hemisphere = P_c[..., z] >= 0
bot_hemisphere = P_c[..., z] <= 0
# convert P_c to 3d polar co-ordinates
P_p = cart2pol(P_c, 1)
#P_p = np.apply_along_axis(convert_cartesian_to_polar, 1, P_c)
# correct order of points must be maintained!!
# there are three cases: top only, bottom only, top and bottom
if np.all(top_hemisphere):
# top only
# flip hemisphere
P_flipped = np.apply_along_axis(self.flip_hemisphere_polar, 1, P_p)
# convert to Lambert projection
Q_p = np.apply_along_axis(self.project_polar, 1, P_flipped)
#if np.diff(Q_p[..., psi])[0] > 0.:
# Q_p = Q_p[::-1]
line = self.ax_top.plot(Q_p[..., psi], Q_p[..., rho],
symbol, color=color, zorder=0, **kwargs)
elif np.all(bot_hemisphere):
# bottom only
# convert P_p to Lambert projection
Q_p = np.apply_along_axis(self.project_polar, 1,
P_p % (2 * np.pi))
line = self.ax_bot.plot(Q_p[..., psi], Q_p[..., rho],
symbol, color=color, zorder=0, **kwargs)
#return P_c, P_p, Q_p
else:
# top and bottom
# rotate points list to a beginning
# (can rely on there being at most one contiguous
# region in each hemisphere)
switch_pt = np.nonzero(np.diff(top_hemisphere))[0][0] + 1
P_roll = np.roll(P_p, switch_pt, axis=0)
P_top_mask = np.roll(top_hemisphere, switch_pt, axis=0)
P_bot_mask = np.roll(bot_hemisphere, switch_pt, axis=0)
P_top = P_roll[P_top_mask]
P_bot = P_roll[P_bot_mask]
# flip top hemisphere
P_flipped = np.apply_along_axis(self.flip_hemisphere_polar,
1, P_top)
# convert to Lambert projection
Q_top = np.apply_along_axis(self.project_polar, 1,
P_flipped % (2 * np.pi))
Q_bot = np.apply_along_axis(self.project_polar, 1,
P_bot)
# plot each set of points
self.ax_top.plot(Q_top[..., psi], Q_top[..., rho],
symbol, color=color, zorder=0, **kwargs)
line = self.ax_bot.plot(Q_bot[..., psi], Q_bot[..., rho],
symbol, color=color, zorder=0, **kwargs)
plt.draw()
return line