class_docstr: Find closest anchors using :mod:`torch_sparse` on a GPU.
class_docstr_tok: ["Find","closest","anchors","using",":","mod",":","`","torch_sparse","`","on","a","GPU","."]
docstr: Determine the candidate pool using breadth-first search.
  
  :param anchors: shape: (a,)
      the anchor node IDs
  :param edge_list: shape: (2, n)
      the edge list with symmetric edges and self-loops
  :param max_iter:
      the maximum number of hops to consider
  :param k:
      the minimum number of anchor nodes to reach
  :param device:
      the device on which the calculations are done
  
  :return: shape: (n, a)
      a boolean array indicating whether anchor $j$ is in the set of $k$ closest anchors for node $i$
  
  :raises ImportError:
      If :mod:`torch_sparse` is not installed
docstr_tok: ["Determine","the","candidate","pool","using","breadth-first","search",".",":","param","anchors",":","shape",":","(","a",",",")","the","anchor","node","IDs",":","param","edge_list",":","shape",":","(","2",",","n",")","the","edge","list","with","symmetric","edges","and","self-loops",":","param","max_iter",":","the","maximum","number","of","hops","to","consider",":","param","k",":","the","minimum","number","of","anchor","nodes","to","reach",":","param","device",":","the","device","on","which","the","calculations","are","done",":","return",":","shape",":","(","n",",","a",")","a","boolean","array","indicating","whether","anchor","$","j","$","is","in","the","set","of","$","k","$","closest","anchors","for","node","$","i","$",":","raises","ImportError",":","If",":","mod",":","`","torch_sparse","`","is","not","installed"]
code: def bfs(
      anchors: numpy.ndarray,
      edge_list: torch.tensor,
      max_iter: int,
      k: int,
      device: torch.device,
  ) -> numpy.ndarray:
      """
      Determine the candidate pool using breadth-first search.
  
      :param anchors: shape: (a,)
          the anchor node IDs
      :param edge_list: shape: (2, n)
          the edge list with symmetric edges and self-loops
      :param max_iter:
          the maximum number of hops to consider
      :param k:
          the minimum number of anchor nodes to reach
      :param device:
          the device on which the calculations are done
  
      :return: shape: (n, a)
          a boolean array indicating whether anchor $j$ is in the set of $k$ closest anchors for node $i$
  
      :raises ImportError:
          If :mod:`torch_sparse` is not installed
      """
      try:
          import torch_sparse
      except ImportError as err:
          raise ImportError(
              "Requires `torch_sparse` to be installed."
          ) from err
  
      num_entities = edge_list.max().item() + 1
      # for each entity, determine anchor pool by BFS
      num_anchors = len(anchors)
  
      anchors = torch.tensor(anchors, dtype=torch.long, device=device)
  
      # an array storing whether node i is reachable by anchor j
      reachable = torch.zeros(
          (num_entities, num_anchors), dtype=torch.bool, device=device
      )
      reachable[anchors] = torch.eye(
          num_anchors, dtype=torch.bool, device=device
      )
  
      # an array indicating whether a node is closed, i.e., has found at least $k$ anchors
      final = torch.zeros(
          (num_entities,), dtype=torch.bool, device=device
      )
  
      # the output that track the distance to each found anchor
      # dtype is unsigned int 8 bit, so we initialize the maximum distance to 255 (or max default)
      dtype = torch.uint8
      pool = torch.zeros(
          (num_entities, num_anchors), dtype=dtype, device=device
      ).fill_(torch.iinfo(dtype).max)
      # initial anchors are 0-hop away from themselves
      pool[
          anchors,
          torch.arange(len(anchors), dtype=torch.long, device=device),
      ] = 0
  
      edge_list = edge_list.to(device)
      values = torch.ones_like(
          edge_list[0], dtype=torch.bool, device=device
      )
  
      old_reachable = reachable
      for i in range(max_iter):
          # propagate one hop
          # TODO the float() trick for GPU result stability until the torch_sparse issue is resolved
          # https://github.com/rusty1s/pytorch_sparse/issues/243
          reachable = (
              torch_sparse.spmm(
                  index=edge_list,
                  value=values.float(),
                  m=num_entities,
                  n=num_entities,
                  matrix=reachable.float(),
              )
              > 0.0
          )
          # convergence check
          if (reachable == old_reachable).all():
              logger.warning(
                  f"Search converged after iteration {i} without all nodes being reachable."
              )
              break
          # newly reached is a mask that points to newly discovered anchors at this particular step
          # implemented as element-wise XOR (will only give True in 0 XOR 1 or 1 XOR 0)
          # in our case we enrich the set of found anchors, so we can only have values turning 0 to 1, eg 0 XOR 1
          newly_reached = reachable ^ old_reachable
          old_reachable = reachable
          # copy pool if we have seen enough anchors and have not yet stopped
          num_reachable = reachable.sum(axis=1)
          enough = num_reachable >= k
          logger.debug(
              f"Iteration {i}: {format_relative_comparison(enough.sum(), total=num_entities)} closed nodes.",
          )
          # update the value in the pool by the current hop value (we start from 0, so +1 be default)
          pool[newly_reached] = i + 1
          # stop once we have enough
          final |= enough
          if final.all():
              break
  
      return pool
  
imports: import logging
  from abc import ABC, abstractmethod
  from typing import Iterable, Optional
  import numpy
  import scipy.sparse
  import torch
  from class_resolver import ClassResolver, OptionalKwargs
  from torch_ppr import page_rank
  from torch_ppr.utils import edge_index_to_sparse_matrix, prepare_page_rank_adjacency, prepare_x0
  from tqdm.auto import tqdm
  from .utils import ensure_num_entities
  from ...typing import DeviceHint
  from ...utils import ExtraReprMixin, format_relative_comparison, resolve_device
label_desc: Use image node_id 4 for calling the SparseBFSSearcher obj's underlying member method code with example usage: obj.bfs(anchors, edge_list, max_iter, k, device) and returns: pool