# 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/simulator.py
# class Simulator(object):
#   def __init__(self,
#                workdir: str = None,
#                wetlab=wet_lab.WetLab(),
#                policy=group_policy.Dorfman(),
#                sampler=smc.SmcSampler(),
#                cheap_sampler=lbp.LbpSampler(),
#                num_simulations: int = 1,
#                num_tests_per_cycle: int = 10,
#                max_test_cycles: int = 5,
#                max_group_size: int = 8,
#                prior_specificity: Union[float, Iterable[float]] = 0.97,
#                prior_sensitivity: Union[float, Iterable[float]] = 0.85,
#                prior_infection_rate: Union[float, Iterable[float]] = 0.05,
#                metrics_cls=metrics.Metrics,
#                export_metrics_every: int = 5):
#   def reset(self, rng: np.ndarray):
#   def sampler(self):
#   def run(self, rngkey=None):
#   def is_over(self) -> bool:
#   def run_simulation(self, rng, sim_idx: int):
#   def run_one_cycle(self, rng, sim_idx: int):
#   def marginal(self):
#   def process_tests_results(self, rng, new_tests):
#   def get_new_groups(self, rng) -> np.ndarray:
#   def _resample(self, rng):
#   def _on_iteration_end(self,
#                         sim_idx: int,
#                         cycle_idx: int,
#                         new_groups=None,
#                         new_tests=None):
#   def interactive_loop(self, rngkey: int = None, show_fn=None):
# 
# Path: group_testing/wet_lab.py
# class WetLab:
# class WetLabCSV(WetLab):
#   def __init__(self,
#                num_patients: int = 2,
#                freeze_diseased: bool = True,
#                base_infection_rate: Union[float, Iterable[float]] = 0.05,
#                specificity: Union[float, Iterable[float]] = 0.97,
#                sensitivity: Union[float, Iterable[float]] = 0.85):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def __init__(self, filename: str = None):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def to_csv(self, groups: np.ndarray, marginals: np.ndarray):
#     def to_list(arr):
# 
# 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):


from absl.testing import absltest
from group_testing import policy
from group_testing import simulator
from group_testing import wet_lab
from group_testing.group_selectors import random
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 Simulator."""



class SimulatorTestCase(jax.test_util.JaxTestCase):

  def setUp(self):
    super().setUp()
    self.num_patients = 10

# 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/simulator.py
# class Simulator(object):
#   def __init__(self,
#                workdir: str = None,
#                wetlab=wet_lab.WetLab(),
#                policy=group_policy.Dorfman(),
#                sampler=smc.SmcSampler(),
#                cheap_sampler=lbp.LbpSampler(),
#                num_simulations: int = 1,
#                num_tests_per_cycle: int = 10,
#                max_test_cycles: int = 5,
#                max_group_size: int = 8,
#                prior_specificity: Union[float, Iterable[float]] = 0.97,
#                prior_sensitivity: Union[float, Iterable[float]] = 0.85,
#                prior_infection_rate: Union[float, Iterable[float]] = 0.05,
#                metrics_cls=metrics.Metrics,
#                export_metrics_every: int = 5):
#   def reset(self, rng: np.ndarray):
#   def sampler(self):
#   def run(self, rngkey=None):
#   def is_over(self) -> bool:
#   def run_simulation(self, rng, sim_idx: int):
#   def run_one_cycle(self, rng, sim_idx: int):
#   def marginal(self):
#   def process_tests_results(self, rng, new_tests):
#   def get_new_groups(self, rng) -> np.ndarray:
#   def _resample(self, rng):
#   def _on_iteration_end(self,
#                         sim_idx: int,
#                         cycle_idx: int,
#                         new_groups=None,
#                         new_tests=None):
#   def interactive_loop(self, rngkey: int = None, show_fn=None):
# 
# Path: group_testing/wet_lab.py
# class WetLab:
# class WetLabCSV(WetLab):
#   def __init__(self,
#                num_patients: int = 2,
#                freeze_diseased: bool = True,
#                base_infection_rate: Union[float, Iterable[float]] = 0.05,
#                specificity: Union[float, Iterable[float]] = 0.97,
#                sensitivity: Union[float, Iterable[float]] = 0.85):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def __init__(self, filename: str = None):
#   def reset(self, rng):
#   def group_tests_outputs(self, rng, groups):
#   def to_csv(self, groups: np.ndarray, marginals: np.ndarray):
#     def to_list(arr):
# 
# 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/simulator_test.py
from absl.testing import absltest
from group_testing import policy
from group_testing import simulator
from group_testing import wet_lab
from group_testing.group_selectors import random
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 Simulator."""



class SimulatorTestCase(jax.test_util.JaxTestCase):

  def setUp(self):
    super().setUp()
    self.num_patients = 10

    self.wetlab = wet_lab.WetLab(self.num_patients, freeze_diseased=True)