# Path: group_testing/metrics.py
# def auc(y_true, y_pred):
# def kl(x, y):
# def binary_entropy(u):
# def entropy(table, axis=0):
#   def __init__(self,
#                workdir=None,
#                num_simulations=0,
#                num_cycles=0,
#                num_patients=0,
#                num_tests_per_cycle=0):
#   def update(self,
#              simulation,
#              cycle,
#              marginal,
#              ground_truth,
#              groups=None,
#              test_results=None,
#              lbp_convergence=None,
#              smc_convergence=None):
#   def load(self):
#   def export(self):
#   def _merge_one(self, other, name: str):
#   def merge(self, other):
#   def _extract_data(self, simulation=None, cycle: int = -1):
#   def roc_curve(self, simulation=None, cycle=-1):
#   def precision_recall_curve(self, simulation: int = None, cycle=-1):
# class Metrics:
# 
# 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/group_selectors/group_selector.py
# class GroupSelector:
#   NEEDS_POSTERIOR = False
#   def __call__(self, rng, state):
#   def get_groups(self, rng, state):


import functools
import gin
import jax
import jax.numpy as np
from absl import logging
from group_testing import metrics
from group_testing import utils
from group_testing.group_selectors import group_selector
  specificity / sensitivity setting. We just replace the vector by its value at
  the appropriate coordinate.

  The size of the group considered here will be the size of cur_group + 1 if
  going forward / -1 if backtracking.

  Args:
   particle_weights: weights of particles
   particles: particles summarizing belief about infection status
   cur_group: group currently considered to add to former groups.
   cur_positives: stores which particles would test positive w.r.t cur_group
   prob_prev_groups_of_particles_x_states: particles x test outcome
     probabilities
   previous_groups_cumcond_entropy: previous conditional entropies
   sensitivity: value (vector) of sensitivity(-ies depending on group size).
   specificity: value (vector) of specificity(-ies depending on group size).
   backtracking: (bool), True if removing rather than adding individuals.

  Returns:
    cur_group : group updated with best choice
    cur_positives : bool vector keeping trace of whether particles
                             would test or not positive
    new_objective : MI reached with this new group
    prob_particles_states : if cur_group were to be selected, this matrix
      would keep track of probability of seeing one of 2^j possible test
      outcomes across all particles.
    new_cond_entropy : if cur_group were to be selected, this constant would be
      added to store the conditional entropies of all tests carried out thusfar
  """
  group_size = np.atleast_1d(np.sum(cur_group) + 1 - 2 * backtracking)

# Path: group_testing/metrics.py
# def auc(y_true, y_pred):
# def kl(x, y):
# def binary_entropy(u):
# def entropy(table, axis=0):
#   def __init__(self,
#                workdir=None,
#                num_simulations=0,
#                num_cycles=0,
#                num_patients=0,
#                num_tests_per_cycle=0):
#   def update(self,
#              simulation,
#              cycle,
#              marginal,
#              ground_truth,
#              groups=None,
#              test_results=None,
#              lbp_convergence=None,
#              smc_convergence=None):
#   def load(self):
#   def export(self):
#   def _merge_one(self, other, name: str):
#   def merge(self, other):
#   def _extract_data(self, simulation=None, cycle: int = -1):
#   def roc_curve(self, simulation=None, cycle=-1):
#   def precision_recall_curve(self, simulation: int = None, cycle=-1):
# class Metrics:
# 
# 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/group_selectors/group_selector.py
# class GroupSelector:
#   NEEDS_POSTERIOR = False
#   def __call__(self, rng, state):
#   def get_groups(self, rng, state):


# Path: group_testing/group_selectors/mutual_information.py
import functools
import gin
import jax
import jax.numpy as np
from absl import logging
from group_testing import metrics
from group_testing import utils
from group_testing.group_selectors import group_selector

  specificity / sensitivity setting. We just replace the vector by its value at
  the appropriate coordinate.

  The size of the group considered here will be the size of cur_group + 1 if
  going forward / -1 if backtracking.

  Args:
   particle_weights: weights of particles
   particles: particles summarizing belief about infection status
   cur_group: group currently considered to add to former groups.
   cur_positives: stores which particles would test positive w.r.t cur_group
   prob_prev_groups_of_particles_x_states: particles x test outcome
     probabilities
   previous_groups_cumcond_entropy: previous conditional entropies
   sensitivity: value (vector) of sensitivity(-ies depending on group size).
   specificity: value (vector) of specificity(-ies depending on group size).
   backtracking: (bool), True if removing rather than adding individuals.

  Returns:
    cur_group : group updated with best choice
    cur_positives : bool vector keeping trace of whether particles
                             would test or not positive
    new_objective : MI reached with this new group
    prob_particles_states : if cur_group were to be selected, this matrix
      would keep track of probability of seeing one of 2^j possible test
      outcomes across all particles.
    new_cond_entropy : if cur_group were to be selected, this constant would be
      added to store the conditional entropies of all tests carried out thusfar
  """
  group_size = np.atleast_1d(np.sum(cur_group) + 1 - 2 * backtracking)

  sensitivity = utils.select_from_sizes(sensitivity, group_size)