def random_nodes_generator(self, return_children: bool = False) -> NodeDataGenerator:
"""
Return a generator that yields the data of all the nodes sampling them at random.
Each value yielded corresponds to the data associated with one node and the \
edge connecting to one of its children. The edge data will be assigned \
to the parent node of the edge. This means the edge data of each node \
contains the data associated with the transition to a child.
Args:
return_children: If ``True`` the data corresponding to the child of \
each node in the path will also be returned.
Yields:
tuple of dictionaries containing the data for each node and edge \
sampled at random.
If ``return_children`` is ``False`` it will yield (node_data, \
edge_data).
If ``return_children`` is ``True`` it will yield (node_data, \
edge_data, next_node_data).
"""
with Backend.use_backend("numpy"):
permutaion = random_state.permutation(list(self.data.nodes))
for next_node in permutaion:
if next_node == self.root_id:
continue
node = self.get_parent(next_node)
yield self._one_node_tuples(node, next_node, return_children)
def iterate_batches(self, batch_size: int, as_dict: bool = True):
with Backend.use_backend("numpy"):
indexes = random_state.permutation(range(len(self)))
for i in range(0, len(self), batch_size):
batch_ix = indexes[i:i + batch_size]
data = tuple(
[getattr(self, val)[batch_ix] for val in self.names])
if as_dict:
yield dict(zip(self.names, data))
else:
yield data
def iterate_values(self,
randomize: bool = False) -> Iterable[Tuple[Tensor]]:
"""
Return a generator that yields a tuple containing the data of each state \
stored in the memory.
"""
indexes = range(len(self))
if randomize:
with Backend.use_backend("numpy"):
indexes = random_state.permutation(indexes)
for i in indexes:
yield tuple([getattr(self, val)[i] for val in self.names])
