# 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 gin
import jax
import jax.numpy as np
import numpy as onp
from group_testing import utils
from group_testing.group_selectors import group_selector

    super().__init__()
    self.modified = modified
    self.cut_off_high = cut_off_high
    self.cut_off_low = cut_off_low

  def __call__(self, rng, state):
    """Produces new groups and adds them to state's stack."""
    p_weights, particles = state.particle_weights, state.particles
    marginal = onp.array(np.sum(p_weights[:, np.newaxis] * particles, axis=0))
    marginal = onp.squeeze(marginal)
    not_cut_ids, = onp.where(np.logical_and(
        marginal < self.cut_off_high, marginal > self.cut_off_low))
    marginal = marginal[not_cut_ids]
    sorted_ids = onp.argsort(marginal)
    sorted_marginal = onp.array(marginal[sorted_ids])
    n_p = 0
    n_r = marginal.size
    if n_r == 0:  # no one left to test in between thresholds
      state.all_cleared = True
      return state

    all_new_groups = np.empty((0, state.num_patients), dtype=bool)
    while n_p < marginal.size:
      index_max = onp.amin((n_r, state.max_group_size))
      group_sizes = onp.arange(1, index_max + 1)
      cum_prod_prob = onp.cumprod(1 - sorted_marginal[n_p:(n_p + index_max)])
      # formula below is only valid for group_size > 1,
      # corrected below for a group of size 1.
      sensitivity = onp.array(

# 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/informative_dorfman.py
import gin
import jax
import jax.numpy as np
import numpy as onp
from group_testing import utils
from group_testing.group_selectors import group_selector


    super().__init__()
    self.modified = modified
    self.cut_off_high = cut_off_high
    self.cut_off_low = cut_off_low

  def __call__(self, rng, state):
    """Produces new groups and adds them to state's stack."""
    p_weights, particles = state.particle_weights, state.particles
    marginal = onp.array(np.sum(p_weights[:, np.newaxis] * particles, axis=0))
    marginal = onp.squeeze(marginal)
    not_cut_ids, = onp.where(np.logical_and(
        marginal < self.cut_off_high, marginal > self.cut_off_low))
    marginal = marginal[not_cut_ids]
    sorted_ids = onp.argsort(marginal)
    sorted_marginal = onp.array(marginal[sorted_ids])
    n_p = 0
    n_r = marginal.size
    if n_r == 0:  # no one left to test in between thresholds
      state.all_cleared = True
      return state

    all_new_groups = np.empty((0, state.num_patients), dtype=bool)
    while n_p < marginal.size:
      index_max = onp.amin((n_r, state.max_group_size))
      group_sizes = onp.arange(1, index_max + 1)
      cum_prod_prob = onp.cumprod(1 - sorted_marginal[n_p:(n_p + index_max)])
      # formula below is only valid for group_size > 1,
      # corrected below for a group of size 1.
      sensitivity = onp.array(

          utils.select_from_sizes(state.prior_sensitivity, group_sizes))