# Path: group_testing/policy.py
# class Policy(object):
# class MimaxPolicy(Policy):
# class OrigamiPolicy(Policy):
# class MezardPolicy(Policy):
# class Dorfman(Policy):
# class InformativeDorfmanPolicy(Policy):
#   def __init__(self, selectors=None):
#   def reset(self):
#   def get_selector(self, index=None):
#   def next_selector(self):
#   def act(self, rng, state):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, subsequent_selectors=None):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, second_split=None, **kwargs):
#   def __init__(self, **kwargs):
# 
# Path: group_testing/state.py
# class State:
#   def __init__(self, num_patients: int, num_tests_per_cycle: int,
#                max_group_size: int, prior_infection_rate: np.ndarray,
#                prior_specificity: np.ndarray, prior_sensitivity: np.ndarray):
#   def reset(self):
#   def add_test_results(self, test_results):
#   def add_groups_to_test(self,
#                          groups: np.ndarray,
#                          results_need_clearing: bool = False):
#   def add_past_groups(self, groups):
#   def update_to_clear_positives(self):
#   def num_groups_left_to_test(self):
#   def extra_tests_needed(self) -> int:
#   def next_groups_to_test(self) -> np.ndarray:
#   def update_particles(self, sampler):
#   def log_posterior_params(self,
#                            sampling_from_scratch=True,
#                            start_from_prior=False,
#                            sampled_up_to=0) -> Dict[str, np.ndarray]:
#   def log_base_measure_params(self,
#                               sampling_from_scratch=True,
#                               start_from_prior=False,
#                               sampled_up_to=0
#                              ) -> Dict[str, np.ndarray]:
# 
# Path: group_testing/group_selectors/random.py
# class RandomSelector(group_selector.GroupSelector):
# class Mezard(group_selector.GroupSelector):
# class RandomDisjoint(group_selector.GroupSelector):
#   def get_groups(self, rng, state):
#   def __init__(self, group_size: int = None):
#   def get_groups(self, rng, state):
# def macula_matrix(d: int, k: int, n: int):
# def sample_groups_of_size(shape: Tuple[int], group_size: int):
# def count_mean(values_arr: onp.ndarray):
# def count_min(values_arr: onp.ndarray):
# def eval_disjoint(mat: onp.ndarray, d: int, count_fn):
# def is_disjoint(mat: onp.ndarray, d: int):
# def sample_disjoint_matrix(num_cols: int,
#                            num_rows: int,
#                            n_max_test: int,
#                            d: int,
#                            max_tries=1e2):
# def sample_maxeval_disjoint_matrix(num_cols: int,
#                                    num_rows: int,
#                                    n_max_test: int,
#                                    d: int,
#                                    max_tries=100,
#                                    count_fn=count_mean):
#   def __init__(self, max_iter=1, method='count', count_fn=count_mean):
#   def get_groups(self, state):
# 
# Path: group_testing/group_selectors/split.py
# class SplitSelector(group_selector.GroupSelector):
# class SplitPositive(group_selector.GroupSelector):
# class TwoDDorfmanPostSelector(group_selector.GroupSelector):
#   def __init__(self, split_factor=None):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __init__(self, split_factor=2):
#   def _split_groups(self, groups):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __call__(self, rng, state):
#   NEEDS_POSTERIOR = False


from absl.testing import absltest
from group_testing import policy
from group_testing import state
from group_testing.group_selectors import random
from group_testing.group_selectors import split
import jax
import jax.numpy as np
import jax.test_util
# coding=utf-8
# Copyright 2020 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Tests for the Policy."""



class PolicyTest(jax.test_util.JaxTestCase):

  def setUp(self):
    super().setUp()
    self.rng = jax.random.PRNGKey(0)
    self.policy = policy.Policy(

# Path: group_testing/policy.py
# class Policy(object):
# class MimaxPolicy(Policy):
# class OrigamiPolicy(Policy):
# class MezardPolicy(Policy):
# class Dorfman(Policy):
# class InformativeDorfmanPolicy(Policy):
#   def __init__(self, selectors=None):
#   def reset(self):
#   def get_selector(self, index=None):
#   def next_selector(self):
#   def act(self, rng, state):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, subsequent_selectors=None):
#   def __init__(self, subsequent_selectors=None, **kwargs):
#   def __init__(self, second_split=None, **kwargs):
#   def __init__(self, **kwargs):
# 
# Path: group_testing/state.py
# class State:
#   def __init__(self, num_patients: int, num_tests_per_cycle: int,
#                max_group_size: int, prior_infection_rate: np.ndarray,
#                prior_specificity: np.ndarray, prior_sensitivity: np.ndarray):
#   def reset(self):
#   def add_test_results(self, test_results):
#   def add_groups_to_test(self,
#                          groups: np.ndarray,
#                          results_need_clearing: bool = False):
#   def add_past_groups(self, groups):
#   def update_to_clear_positives(self):
#   def num_groups_left_to_test(self):
#   def extra_tests_needed(self) -> int:
#   def next_groups_to_test(self) -> np.ndarray:
#   def update_particles(self, sampler):
#   def log_posterior_params(self,
#                            sampling_from_scratch=True,
#                            start_from_prior=False,
#                            sampled_up_to=0) -> Dict[str, np.ndarray]:
#   def log_base_measure_params(self,
#                               sampling_from_scratch=True,
#                               start_from_prior=False,
#                               sampled_up_to=0
#                              ) -> Dict[str, np.ndarray]:
# 
# Path: group_testing/group_selectors/random.py
# class RandomSelector(group_selector.GroupSelector):
# class Mezard(group_selector.GroupSelector):
# class RandomDisjoint(group_selector.GroupSelector):
#   def get_groups(self, rng, state):
#   def __init__(self, group_size: int = None):
#   def get_groups(self, rng, state):
# def macula_matrix(d: int, k: int, n: int):
# def sample_groups_of_size(shape: Tuple[int], group_size: int):
# def count_mean(values_arr: onp.ndarray):
# def count_min(values_arr: onp.ndarray):
# def eval_disjoint(mat: onp.ndarray, d: int, count_fn):
# def is_disjoint(mat: onp.ndarray, d: int):
# def sample_disjoint_matrix(num_cols: int,
#                            num_rows: int,
#                            n_max_test: int,
#                            d: int,
#                            max_tries=1e2):
# def sample_maxeval_disjoint_matrix(num_cols: int,
#                                    num_rows: int,
#                                    n_max_test: int,
#                                    d: int,
#                                    max_tries=100,
#                                    count_fn=count_mean):
#   def __init__(self, max_iter=1, method='count', count_fn=count_mean):
#   def get_groups(self, state):
# 
# Path: group_testing/group_selectors/split.py
# class SplitSelector(group_selector.GroupSelector):
# class SplitPositive(group_selector.GroupSelector):
# class TwoDDorfmanPostSelector(group_selector.GroupSelector):
#   def __init__(self, split_factor=None):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __init__(self, split_factor=2):
#   def _split_groups(self, groups):
#   def get_groups(self, rng, state):
#   def __call__(self, rng, state):
#   def __call__(self, rng, state):
#   NEEDS_POSTERIOR = False


# Path: group_testing/policy_test.py
from absl.testing import absltest
from group_testing import policy
from group_testing import state
from group_testing.group_selectors import random
from group_testing.group_selectors import split
import jax
import jax.numpy as np
import jax.test_util

# coding=utf-8
# Copyright 2020 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Tests for the Policy."""



class PolicyTest(jax.test_util.JaxTestCase):

  def setUp(self):
    super().setUp()
    self.rng = jax.random.PRNGKey(0)
    self.policy = policy.Policy(

        [random.RandomSelector(), split.SplitSelector()])