# Path: group_testing/utils.py
# def unique(rng, binary_vectors: np.ndarray) -> int:
# def select_from_sizes(values: np.ndarray, sizes: np.ndarray) -> np.ndarray:
# 
# Path: group_testing/samplers/sampler.py
# class Sampler:
#   def __init__(self):
#   def reset(self):
#   def is_cheap(self):
#   def marginal(self):
#   def reset_convergence_metric(self):


from absl import logging
from group_testing import utils
from group_testing.samplers import sampler
import gin
import jax
import jax.nn
import jax.numpy as np
import jax.scipy.special

  Args:
    tests : np.ndarray<bool>[n_groups] results stored as a vector of booleans
    groups : np.ndarray<bool>[n_groups, n_patients] matrix of groups
    base_infection_rate : np.ndarray<float> [1,] or [n_patients,] infection rate
    sensitivity : np.ndarray<float> [?,] of sensitivity per group size
    specificity : np.ndarray<float> [?,] of specificity per group size
    min_iterations: int, min number of belief propagation iterations
    max_iterations: int, max number of belief propagation iterations
    atol: float, elementwise tolerance for the difference between two
      consecutive iterations.

  Returns:
    two vectors of marginal probabilities for all n_patients, obtained
    as consecutive evaluations of the LBP algorithm after n_iter and n_iter+1
    iterations.
  """
  n_groups, n_patients = groups.shape
  if np.size(groups) == 0:
    if np.size(base_infection_rate) == 1:  # only one rate
      marginal = base_infection_rate * np.ones(n_patients)
      return marginal, 0
    elif np.size(base_infection_rate) == n_patients:
      return base_infection_rate, 0
    else:
      raise ValueError("Improper size for vector of base infection rates")

  mu = -jax.scipy.special.logit(base_infection_rate)

  groups_size = np.sum(groups, axis=1)

# Path: group_testing/utils.py
# def unique(rng, binary_vectors: np.ndarray) -> int:
# def select_from_sizes(values: np.ndarray, sizes: np.ndarray) -> np.ndarray:
# 
# Path: group_testing/samplers/sampler.py
# class Sampler:
#   def __init__(self):
#   def reset(self):
#   def is_cheap(self):
#   def marginal(self):
#   def reset_convergence_metric(self):


# Path: group_testing/samplers/loopy_belief_propagation.py
from absl import logging
from group_testing import utils
from group_testing.samplers import sampler
import gin
import jax
import jax.nn
import jax.numpy as np
import jax.scipy.special


  Args:
    tests : np.ndarray<bool>[n_groups] results stored as a vector of booleans
    groups : np.ndarray<bool>[n_groups, n_patients] matrix of groups
    base_infection_rate : np.ndarray<float> [1,] or [n_patients,] infection rate
    sensitivity : np.ndarray<float> [?,] of sensitivity per group size
    specificity : np.ndarray<float> [?,] of specificity per group size
    min_iterations: int, min number of belief propagation iterations
    max_iterations: int, max number of belief propagation iterations
    atol: float, elementwise tolerance for the difference between two
      consecutive iterations.

  Returns:
    two vectors of marginal probabilities for all n_patients, obtained
    as consecutive evaluations of the LBP algorithm after n_iter and n_iter+1
    iterations.
  """
  n_groups, n_patients = groups.shape
  if np.size(groups) == 0:
    if np.size(base_infection_rate) == 1:  # only one rate
      marginal = base_infection_rate * np.ones(n_patients)
      return marginal, 0
    elif np.size(base_infection_rate) == n_patients:
      return base_infection_rate, 0
    else:
      raise ValueError("Improper size for vector of base infection rates")

  mu = -jax.scipy.special.logit(base_infection_rate)

  groups_size = np.sum(groups, axis=1)

  sensitivity = utils.select_from_sizes(sensitivity, groups_size)