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