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))))
Esempio n. 3
0
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