Python numpy.ravel_multi_index() Examples

def raw_valid_fn_vec(self, xyt):
    """Returns if the given set of nodes is valid or not."""
    height = self.traversible.shape[0]
    width = self.traversible.shape[1]
    x = np.round(xyt[:,[0]]).astype(np.int32)
    y = np.round(xyt[:,[1]]).astype(np.int32)
    is_inside = np.all(np.concatenate((x >= 0, y >= 0,
                                       x < width, y < height), axis=1), axis=1)
    x = np.minimum(np.maximum(x, 0), width-1)
    y = np.minimum(np.maximum(y, 0), height-1)
    ind = np.ravel_multi_index((y,x), self.traversible.shape)
    is_traversible = self.traversible.ravel()[ind]

    is_valid = np.all(np.concatenate((is_inside[:,np.newaxis], is_traversible),
                                     axis=1), axis=1)
    return is_valid 

def valid_fn_vec(self, pqr):
    """Returns if the given set of nodes is valid or not."""
    xyt = self.to_actual_xyt_vec(np.array(pqr))
    height = self.traversible.shape[0]
    width = self.traversible.shape[1]
    x = np.round(xyt[:,[0]]).astype(np.int32)
    y = np.round(xyt[:,[1]]).astype(np.int32)
    is_inside = np.all(np.concatenate((x >= 0, y >= 0,
                                       x < width, y < height), axis=1), axis=1)
    x = np.minimum(np.maximum(x, 0), width-1)
    y = np.minimum(np.maximum(y, 0), height-1)
    ind = np.ravel_multi_index((y,x), self.traversible.shape)
    is_traversible = self.traversible.ravel()[ind]

    is_valid = np.all(np.concatenate((is_inside[:,np.newaxis], is_traversible),
                                     axis=1), axis=1)
    return is_valid 

def _sparse_equal(A, B, atol=1e-8):
    """ NOTE: same as matrixtools.sparse_equal - but can't import that here """
    if _np.array_equal(A.shape, B.shape) == 0:
        return False

    r1, c1 = A.nonzero()
    r2, c2 = B.nonzero()

    lidx1 = _np.ravel_multi_index((r1, c1), A.shape)
    lidx2 = _np.ravel_multi_index((r2, c2), B.shape)
    sidx1 = lidx1.argsort()
    sidx2 = lidx2.argsort()

    index_match = _np.array_equal(lidx1[sidx1], lidx2[sidx2])
    if index_match == 0:
        return False
    else:
        v1 = A.data
        v2 = B.data
        V1 = v1[sidx1]
        V2 = v2[sidx2]
    return _np.allclose(V1, V2, atol=atol) 

def test_big_indices(self):
        # ravel_multi_index for big indices (issue #7546)
        if np.intp == np.int64:
            arr = ([1, 29], [3, 5], [3, 117], [19, 2],
                   [2379, 1284], [2, 2], [0, 1])
            assert_equal(
                np.ravel_multi_index(arr, (41, 7, 120, 36, 2706, 8, 6)),
                [5627771580, 117259570957])

        # test overflow checking for too big array (issue #7546)
        dummy_arr = ([0],[0])
        half_max = np.iinfo(np.intp).max // 2
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2)), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2))
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2), order='F'), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2), order='F') 

def get_grad_operator(mask):
    """Returns sparse matrix computing horizontal, vertical, and two diagonal gradients."""
    horizontal_left = np.ravel_multi_index(np.nonzero(mask[:, :-1] | mask[:, 1:]), mask.shape)
    horizontal_right = horizontal_left + 1

    vertical_top = np.ravel_multi_index(np.nonzero(mask[:-1, :] | mask[1:, :]), mask.shape)
    vertical_bottom = vertical_top + mask.shape[1]

    diag_main_1 = np.ravel_multi_index(np.nonzero(mask[:-1, :-1] | mask[1:, 1:]), mask.shape)
    diag_main_2 = diag_main_1 + mask.shape[1] + 1

    diag_sub_1 = np.ravel_multi_index(np.nonzero(mask[:-1, 1:] | mask[1:, :-1]), mask.shape) + 1
    diag_sub_2 = diag_sub_1 + mask.shape[1] - 1

    indices = np.stack((
        np.concatenate((horizontal_left, vertical_top, diag_main_1, diag_sub_1)),
        np.concatenate((horizontal_right, vertical_bottom, diag_main_2, diag_sub_2))
    ), axis=-1)
    return scipy.sparse.coo_matrix(
        (np.tile([-1, 1], len(indices)), (np.arange(indices.size) // 2, indices.flatten())),
        shape=(len(indices), mask.size)) 

def figure_3_2_linear_system():
    '''
    Here we solve the linear system of equations to find the exact solution.
    We do this by filling the coefficients for each of the states with their respective right side constant.
    '''
    A = -1 * np.eye(WORLD_SIZE * WORLD_SIZE)
    b = np.zeros(WORLD_SIZE * WORLD_SIZE)
    for i in range(WORLD_SIZE):
        for j in range(WORLD_SIZE):
            s = [i, j]  # current state
            index_s = np.ravel_multi_index(s, (WORLD_SIZE, WORLD_SIZE))
            for a in ACTIONS:
                s_, r = step(s, a)
                index_s_ = np.ravel_multi_index(s_, (WORLD_SIZE, WORLD_SIZE))

                A[index_s, index_s_] += ACTION_PROB * DISCOUNT
                b[index_s] -= ACTION_PROB * r

    x = np.linalg.solve(A, b)
    draw_image(np.round(x.reshape(WORLD_SIZE, WORLD_SIZE), decimals=2))
    plt.savefig('../images/figure_3_2_linear_system.png')
    plt.close() 

def offset_to_index(self, index, offset):
    """Calculate the index of another box at offset w.r.t.

    current index.

    Args:
      index: the current flat index from which to calculate the offset index.
      offset: the xyz offset from current index at which to calculate the new
        index.

    Returns:
      The flat index at offset from current index, or None if the given offset
      goes beyond the range of sub-boxes.

    This is usually used to calculate the boxes that neighbor the current box.
    """
    coords = np.unravel_index(index, self.total_sub_boxes_xyz, order='F')
    offset_coords = np.array(coords) + offset
    if np.any(offset_coords < 0) or np.any(
        offset_coords >= self.total_sub_boxes_xyz):
      return None
    return np.ravel_multi_index(
        offset_coords, self.total_sub_boxes_xyz, order='F') 

def __init__(self):
        self.shape = (4, 12)

        nS = np.prod(self.shape)
        nA = 4

        # Cliff Location
        self._cliff = np.zeros(self.shape, dtype=np.bool)
        self._cliff[3, 1:-1] = True

        # Calculate transition probabilities
        P = {}
        for s in range(nS):
            position = np.unravel_index(s, self.shape)
            P[s] = { a : [] for a in range(nA) }
            P[s][UP] = self._calculate_transition_prob(position, [-1, 0])
            P[s][RIGHT] = self._calculate_transition_prob(position, [0, 1])
            P[s][DOWN] = self._calculate_transition_prob(position, [1, 0])
            P[s][LEFT] = self._calculate_transition_prob(position, [0, -1])

        # We always start in state (3, 0)
        isd = np.zeros(nS)
        isd[np.ravel_multi_index((3,0), self.shape)] = 1.0

        super(CliffWalkingEnv, self).__init__(nS, nA, P, isd) 

def test_big_indices(self):
        # ravel_multi_index for big indices (issue #7546)
        if np.intp == np.int64:
            arr = ([1, 29], [3, 5], [3, 117], [19, 2],
                   [2379, 1284], [2, 2], [0, 1])
            assert_equal(
                np.ravel_multi_index(arr, (41, 7, 120, 36, 2706, 8, 6)),
                [5627771580, 117259570957])

        # test overflow checking for too big array (issue #7546)
        dummy_arr = ([0],[0])
        half_max = np.iinfo(np.intp).max // 2
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2)), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2))
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2), order='F'), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2), order='F') 

def test_big_indices(self):
        # ravel_multi_index for big indices (issue #7546)
        if np.intp == np.int64:
            arr = ([1, 29], [3, 5], [3, 117], [19, 2],
                   [2379, 1284], [2, 2], [0, 1])
            assert_equal(
                np.ravel_multi_index(arr, (41, 7, 120, 36, 2706, 8, 6)),
                [5627771580, 117259570957])

        # test overflow checking for too big array (issue #7546)
        dummy_arr = ([0],[0])
        half_max = np.iinfo(np.intp).max // 2
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2)), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2))
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2), order='F'), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2), order='F') 

def _calculate_transition_prob(self, current, delta):
        """
        Determine the outcome for an action. Transition Prob is always 1.0. 
        :param current: Current position on the grid as (row, col) 
        :param delta: Change in position for transition
        :return: (1.0, new_state, reward, done)
        """
        new_position = np.array(current) + np.array(delta)
        new_position = self._limit_coordinates(new_position).astype(int)
        new_state = np.ravel_multi_index(tuple(new_position), self.shape)
        if self._cliff[tuple(new_position)]:
            return [(1.0, self.start_state_index, -100, False)]

        terminal_state = (self.shape[0] - 1, self.shape[1] - 1)
        is_done = tuple(new_position) == terminal_state
        return [(1.0, new_state, -1, is_done)] 

def test_dtypes(self):
        # Test with different data types
        for dtype in [np.int16, np.uint16, np.int32,
                      np.uint32, np.int64, np.uint64]:
            coords = np.array([[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0]], dtype=dtype)
            shape = (5, 8)
            uncoords = 8*coords[0]+coords[1]
            assert_equal(np.ravel_multi_index(coords, shape), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape))
            uncoords = coords[0]+5*coords[1]
            assert_equal(np.ravel_multi_index(coords, shape, order='F'), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape, order='F'))

            coords = np.array([[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0], [1, 3, 1, 0, 9, 5]],
                              dtype=dtype)
            shape = (5, 8, 10)
            uncoords = 10*(8*coords[0]+coords[1])+coords[2]
            assert_equal(np.ravel_multi_index(coords, shape), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape))
            uncoords = coords[0]+5*(coords[1]+8*coords[2])
            assert_equal(np.ravel_multi_index(coords, shape, order='F'), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) 

def test_big_indices(self):
        # ravel_multi_index for big indices (issue #7546)
        if np.intp == np.int64:
            arr = ([1, 29], [3, 5], [3, 117], [19, 2],
                   [2379, 1284], [2, 2], [0, 1])
            assert_equal(
                np.ravel_multi_index(arr, (41, 7, 120, 36, 2706, 8, 6)),
                [5627771580, 117259570957])

        # test overflow checking for too big array (issue #7546)
        dummy_arr = ([0],[0])
        half_max = np.iinfo(np.intp).max // 2
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2)), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2))
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2), order='F'), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2), order='F') 

def raw_valid_fn_vec(self, xyt):
    """Returns if the given set of nodes is valid or not."""
    height = self.traversible.shape[0]
    width = self.traversible.shape[1]
    x = np.round(xyt[:,[0]]).astype(np.int32)
    y = np.round(xyt[:,[1]]).astype(np.int32)
    is_inside = np.all(np.concatenate((x >= 0, y >= 0,
                                       x < width, y < height), axis=1), axis=1)
    x = np.minimum(np.maximum(x, 0), width-1)
    y = np.minimum(np.maximum(y, 0), height-1)
    ind = np.ravel_multi_index((y,x), self.traversible.shape)
    is_traversible = self.traversible.ravel()[ind]

    is_valid = np.all(np.concatenate((is_inside[:,np.newaxis], is_traversible),
                                     axis=1), axis=1)
    return is_valid 

def valid_fn_vec(self, pqr):
    """Returns if the given set of nodes is valid or not."""
    xyt = self.to_actual_xyt_vec(np.array(pqr))
    height = self.traversible.shape[0]
    width = self.traversible.shape[1]
    x = np.round(xyt[:,[0]]).astype(np.int32)
    y = np.round(xyt[:,[1]]).astype(np.int32)
    is_inside = np.all(np.concatenate((x >= 0, y >= 0,
                                       x < width, y < height), axis=1), axis=1)
    x = np.minimum(np.maximum(x, 0), width-1)
    y = np.minimum(np.maximum(y, 0), height-1)
    ind = np.ravel_multi_index((y,x), self.traversible.shape)
    is_traversible = self.traversible.ravel()[ind]

    is_valid = np.all(np.concatenate((is_inside[:,np.newaxis], is_traversible),
                                     axis=1), axis=1)
    return is_valid 

def test_big_indices(self):
        # ravel_multi_index for big indices (issue #7546)
        if np.intp == np.int64:
            arr = ([1, 29], [3, 5], [3, 117], [19, 2],
                   [2379, 1284], [2, 2], [0, 1])
            assert_equal(
                np.ravel_multi_index(arr, (41, 7, 120, 36, 2706, 8, 6)),
                [5627771580, 117259570957])

        # test overflow checking for too big array (issue #7546)
        dummy_arr = ([0],[0])
        half_max = np.iinfo(np.intp).max // 2
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2)), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2))
        assert_equal(
            np.ravel_multi_index(dummy_arr, (half_max, 2), order='F'), [0])
        assert_raises(ValueError,
            np.ravel_multi_index, dummy_arr, (half_max+1, 2), order='F') 

def pack_samples(self, samples, dtype=None):
        """Pack samples into one integer per sample

        Store one sample in a single integer instead of a list of
        integers with length `len(self.nsoutdims)`. Example:

        >>> p = pauli_mpp(nr_sites=2, local_dim=2)
        >>> p.outdims
        (6, 6)
        >>> p.pack_samples(np.array([[0, 1], [1, 0], [1, 2], [5, 5]]))
        array([ 1,  6,  8, 35])

        """
        assert samples.ndim == 2
        assert samples.shape[1] == len(self.nsoutdims)
        samples = np.ravel_multi_index(samples.T, self.nsoutdims)
        if dtype not in (True, False, None) and issubclass(dtype, np.integer):
            info = np.iinfo(dtype)
            assert samples.min() >= info.min
            assert samples.max() <= info.max
            samples = samples.astype(dtype)
        return samples 

def convert_traversible_to_graph(traversible, ff_cost=1., fo_cost=1.,
                                 oo_cost=1., connectivity=4):
  assert(connectivity == 4 or connectivity == 8)

  sz_x = traversible.shape[1]
  sz_y = traversible.shape[0]
  g, nodes = generate_lattice(sz_x, sz_y)

  # Assign costs.
  edge_wts = g.new_edge_property('float')
  g.edge_properties['wts'] = edge_wts
  wts = np.ones(g.num_edges(), dtype=np.float32)
  edge_wts.get_array()[:] = wts

  if connectivity == 8:
    add_diagonal_edges(g, nodes, sz_x, sz_y, np.sqrt(2.))

  se = np.array([[int(e.source()), int(e.target())] for e in g.edges()])
  s_xy = nodes[se[:,0]]
  t_xy = nodes[se[:,1]]
  s_t = np.ravel_multi_index((s_xy[:,1], s_xy[:,0]), traversible.shape)
  t_t = np.ravel_multi_index((t_xy[:,1], t_xy[:,0]), traversible.shape)
  s_t = traversible.ravel()[s_t]
  t_t = traversible.ravel()[t_t]

  wts = np.zeros(g.num_edges(), dtype=np.float32)
  wts[np.logical_and(s_t == True, t_t == True)] = ff_cost
  wts[np.logical_and(s_t == False, t_t == False)] = oo_cost
  wts[np.logical_xor(s_t, t_t)] = fo_cost

  edge_wts = g.edge_properties['wts']
  for i, e in enumerate(g.edges()):
    edge_wts[e] = edge_wts[e] * wts[i]
  # d = edge_wts.get_array()*1.
  # edge_wts.get_array()[:] = d*wts 
  return g, nodes 

def label_nodes_with_class(nodes_xyt, class_maps, pix):
  """
  Returns: 
    class_maps__: one-hot class_map for each class.
    node_class_label: one-hot class_map for each class, nodes_xyt.shape[0] x n_classes
  """
  # Assign each pixel to a node.
  selem = skimage.morphology.disk(pix)
  class_maps_ = class_maps*1.
  for i in range(class_maps.shape[2]):
    class_maps_[:,:,i] = skimage.morphology.dilation(class_maps[:,:,i]*1, selem)
  class_maps__ = np.argmax(class_maps_, axis=2)
  class_maps__[np.max(class_maps_, axis=2) == 0] = -1

  # For each node pick out the label from this class map.
  x = np.round(nodes_xyt[:,[0]]).astype(np.int32)
  y = np.round(nodes_xyt[:,[1]]).astype(np.int32)
  ind = np.ravel_multi_index((y,x), class_maps__.shape)
  node_class_label = class_maps__.ravel()[ind][:,0]

  # Convert to one hot versions.
  class_maps_one_hot = np.zeros(class_maps.shape, dtype=np.bool)
  node_class_label_one_hot = np.zeros((node_class_label.shape[0], class_maps.shape[2]), dtype=np.bool)
  for i in range(class_maps.shape[2]):
    class_maps_one_hot[:,:,i] = class_maps__ == i 
    node_class_label_one_hot[:,i] = node_class_label == i
  return class_maps_one_hot, node_class_label_one_hot 

def _diagonal_idx_array(batch_size, n):
    idx_offsets = np.arange(
        start=0, stop=batch_size * n * n, step=n * n, dtype=np.int32).reshape(
        (batch_size, 1))
    idx = np.ravel_multi_index(
        np.diag_indices(n), (n, n)).reshape((1, n)).astype(np.int32)
    return cuda.to_gpu(idx + idx_offsets) 

def _non_diagonal_idx_array(batch_size, n):
    idx_offsets = np.arange(
        start=0, stop=batch_size * n * n, step=n * n, dtype=np.int32).reshape(
        (batch_size, 1))
    idx = np.ravel_multi_index(
        np.tril_indices(n, -1), (n, n)).reshape((1, -1)).astype(np.int32)
    return cuda.to_gpu(idx + idx_offsets) 

def unique_rows_counts(a):
    lidx = np.ravel_multi_index(a.T, a.max(0) + 1)
    _, unq_idx, counts = np.unique(lidx, return_index=True, return_counts=True)
    return a[unq_idx], counts 

def sparse_equal(A, B, atol=1e-8):
    """
    Checks whether two Scipy sparse matrices are (almost) equal.

    Parameters
    ----------
    A, B : scipy.sparse matrix
        The two matrices to compare.

    atol : float, optional
        The tolerance to use, passed to `numpy.allclose`, when comparing
        the elements of `A` and `B`.

    Returns
    -------
    bool
    """
    if _np.array_equal(A.shape, B.shape) == 0:
        return False

    r1, c1 = A.nonzero()
    r2, c2 = B.nonzero()

    lidx1 = _np.ravel_multi_index((r1, c1), A.shape)
    lidx2 = _np.ravel_multi_index((r2, c2), B.shape)
    sidx1 = lidx1.argsort()
    sidx2 = lidx2.argsort()

    index_match = _np.array_equal(lidx1[sidx1], lidx2[sidx2])
    if index_match == 0:
        return False
    else:
        v1 = A.data
        v2 = B.data
        V1 = v1[sidx1]
        V2 = v2[sidx2]
    return _np.allclose(V1, V2, atol=atol) 

def __getitem__(self, item):
        """Access electrode(s) in the grid

        Parameters
        ----------
        item : index, string, tuple, or list thereof
            An electrode in the grid can be accessed in three ways:

            *  by name, e.g. grid['A1']
            *  by index into the flattened array, e.g. grid[0]
            *  by (row, column) index into the 2D grid, e.g. grid[0, 0]

            You can also pass a list or NumPy array of the above.

        Returns
        -------
        electrode : `~pulse2percept.implants.Electrode`, list thereof, or None
            Returns the corresponding `~pulse2percept.implants.Electrode`
            object or ``None`` if index is not valid.
        """
        if isinstance(item, (list, np.ndarray)):
            # Recursive call for list items:
            return [self.__getitem__(i) for i in item]
        try:
            # Access by key into OrderedDict, e.g. grid['A1']:
            return self.electrodes[item]
        except (KeyError, TypeError):
            # Access by index into flattened array, e.g. grid[0]:
            try:
                return list(self.electrodes.values())[item]
            except (IndexError, KeyError, TypeError):
                # Access by [r, c] into 2D grid, e.g. grid[0, 3]:
                try:
                    idx = np.ravel_multi_index(item, self.shape)
                    return list(self.electrodes.values())[idx]
                except (KeyError, TypeError, ValueError):
                    # Index not found:
                    return None 

def test_dtypes(self):
        # Test with different data types
        for dtype in [np.int16, np.uint16, np.int32,
                      np.uint32, np.int64, np.uint64]:
            coords = np.array(
                [[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0]], dtype=dtype)
            shape = (5, 8)
            uncoords = 8*coords[0]+coords[1]
            assert_equal(np.ravel_multi_index(coords, shape), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape))
            uncoords = coords[0]+5*coords[1]
            assert_equal(
                np.ravel_multi_index(coords, shape, order='F'), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape, order='F'))

            coords = np.array(
                [[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0], [1, 3, 1, 0, 9, 5]],
                dtype=dtype)
            shape = (5, 8, 10)
            uncoords = 10*(8*coords[0]+coords[1])+coords[2]
            assert_equal(np.ravel_multi_index(coords, shape), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape))
            uncoords = coords[0]+5*(coords[1]+8*coords[2])
            assert_equal(
                np.ravel_multi_index(coords, shape, order='F'), uncoords)
            assert_equal(coords, np.unravel_index(uncoords, shape, order='F')) 

def test_clipmodes(self):
        # Test clipmodes
        assert_equal(
            np.ravel_multi_index([5, 1, -1, 2], (4, 3, 7, 12), mode='wrap'),
            np.ravel_multi_index([1, 1, 6, 2], (4, 3, 7, 12)))
        assert_equal(np.ravel_multi_index([5, 1, -1, 2], (4, 3, 7, 12),
                                          mode=(
                                              'wrap', 'raise', 'clip', 'raise')),
                     np.ravel_multi_index([1, 1, 0, 2], (4, 3, 7, 12)))
        assert_raises(
            ValueError, np.ravel_multi_index, [5, 1, -1, 2], (4, 3, 7, 12)) 

def _lookup(self, xyz):
        """
        Performs a 3D lookup from real world coordinates to the flat indices in the volume
        defined in the BrainCoordinates object
        :param xyz: [n, 3] array of coordinates
        :return: n array of label values
        """
        idims = np.split(self.bc.xyz2i(xyz)[:, self.xyz2dims], [1, 2], axis=-1)
        inds = np.ravel_multi_index(idims, self.bc.nxyz[self.xyz2dims])
        return inds.squeeze() 

def test_tiny_cycle():
    g, idxs, degimg = csr.skeleton_to_csgraph(tinycycle)
    expected_indptr = [0, 0, 2, 4, 6, 8]
    expected_indices = [2, 3, 1, 4, 1, 4, 2, 3]
    expected_data = np.sqrt(2)

    assert_equal(g.indptr, expected_indptr)
    assert_equal(g.indices, expected_indices)
    assert_almost_equal(g.data, expected_data)

    expected_degrees = np.array([[0, 2, 0], [2, 0, 2], [0, 2, 0]])
    assert_equal(degimg, expected_degrees)
    assert_equal(np.ravel_multi_index(idxs.astype(int).T, tinycycle.shape),
                 [0, 1, 3, 5, 7]) 

def _calculate_transition_prob(self, current, delta):
        new_position = np.array(current) + np.array(delta)
        new_position = self._limit_coordinates(new_position).astype(int)
        new_state = np.ravel_multi_index(tuple(new_position), self.shape)

        # Newer version of rewards/costs from G-learning paper
        # reward = -100.0 if self._cliff[tuple(new_position)] else -1.0
        reward = -1.0
        if self._cliff[tuple(new_position)]:
            reward = -100.0
        elif tuple(new_position) == (3,11):
            reward = 0.0

        is_done = self._cliff[tuple(new_position)] or (tuple(new_position) == (3,11))
        return [(1.0, new_state, reward, is_done)] 

def _calculate_transition_prob(self, current, delta, winds):
        new_position = np.array(current) + np.array(delta) + np.array([-1, 0]) * winds[tuple(current)]
        new_position = self._limit_coordinates(new_position).astype(int)
        new_state = np.ravel_multi_index(tuple(new_position), self.shape)
        is_done = tuple(new_position) == (3, 7)
        return [(1.0, new_state, -1.0, is_done)]