# Path: treecat/testutil.py
# def make_seed(*args):
#     """Create a random seed by fingerprinting args."""
#     digest = fingerprint(args)
#     bigseed = int('0x{}'.format(digest), 0)
#     seed = bigseed % (2**32)
#     return seed
# 
# Path: treecat/util.py
# def parallel_map(fun, args):
#     global POOL
#     args = list(args)
#     if len(args) < 2 or TREECAT_THREADS == 1:
#         return list(map(fun, args))
#     if POOL is None:
#         POOL = multiprocessing.Pool(TREECAT_THREADS)
#     return POOL.map(fun, args)
# 
# Path: treecat/util.py
# def quantize_from_probs2(probs, resolution):
#     """Quantize multiple non-normalized probs to given resolution.
# 
#     Args:
#         probs: An [N, M]-shaped numpy array of non-normalized probabilities.
# 
#     Returns:
#         An [N, M]-shaped array of quantized probabilities such that
#         np.all(result.sum(axis=1) == resolution).
#     """
#     assert len(probs.shape) == 2
#     N, M = probs.shape
#     probs = probs / probs.sum(axis=1, keepdims=True)
#     result = np.zeros(probs.shape, np.int8)
#     range_N = np.arange(N, dtype=np.int32)
#     for _ in range(resolution):
#         sample = probs.argmax(axis=1)
#         result[range_N, sample] += 1
#         probs[range_N, sample] -= 1.0 / resolution
#     return result
# 
# Path: treecat/util.py
# @jit(nopython=True, cache=True)
# def sample_from_probs(probs):
#     """Sample from a vector of non-normalized probabilitites.
# 
#     Args:
#         probs: An [M]-shaped numpy array of non-normalized probabilities.
# 
#     Returns:
#         An integer in range(M).
#     """
#     cdf = probs.cumsum()
#     return (np.random.rand() * cdf[-1] < cdf).argmax()
# 
# Path: treecat/util.py
# def sample_from_probs2(probs, out=None):
#     """Sample from multiple vectors of non-normalized probabilities.
# 
#     Args:
#         probs: An [N, M]-shaped numpy array of non-normalized probabilities.
#         out: An optional destination for the result.
# 
#     Returns:
#         An [N]-shaped numpy array of integers in range(M).
#     """
#     # Adapted from https://stackoverflow.com/questions/40474436
#     assert len(probs.shape) == 2
#     cdf = probs.cumsum(axis=1)
#     u = np.random.rand(probs.shape[0], 1) * cdf[:, -1, np.newaxis]
#     return (u < cdf).argmax(axis=1, out=out)
# 
# Path: treecat/util.py
# def set_random_seed(seed):
#     """Set random seeds for both numpy and numba."""
#     np.random.seed(seed)
#     jit_random_seed(seed)


import itertools
import numpy as np
import pytest
from goftests import multinomial_goodness_of_fit
from treecat.testutil import make_seed
from treecat.util import parallel_map
from treecat.util import quantize_from_probs2
from treecat.util import sample_from_probs
from treecat.util import sample_from_probs2
from treecat.util import set_random_seed
    assert 1e-2 < gof


@pytest.mark.parametrize('size,resolution', [
    (1, 0),
    (1, 1),
    (1, 2),
    (2, 0),
    (2, 1),
    (2, 2),
    (2, 3),
    (2, 4),
    (2, 5),
    (3, 0),
    (3, 1),
    (3, 2),
    (3, 3),
    (3, 4),
    (3, 5),
    (4, 0),
    (4, 1),
    (4, 2),
    (4, 3),
    (4, 4),
    (5, 0),
    (5, 1),
    (5, 2),
    (5, 3),
])
def test_quantize_from_probs2(size, resolution):

# Path: treecat/testutil.py
# def make_seed(*args):
#     """Create a random seed by fingerprinting args."""
#     digest = fingerprint(args)
#     bigseed = int('0x{}'.format(digest), 0)
#     seed = bigseed % (2**32)
#     return seed
# 
# Path: treecat/util.py
# def parallel_map(fun, args):
#     global POOL
#     args = list(args)
#     if len(args) < 2 or TREECAT_THREADS == 1:
#         return list(map(fun, args))
#     if POOL is None:
#         POOL = multiprocessing.Pool(TREECAT_THREADS)
#     return POOL.map(fun, args)
# 
# Path: treecat/util.py
# def quantize_from_probs2(probs, resolution):
#     """Quantize multiple non-normalized probs to given resolution.
# 
#     Args:
#         probs: An [N, M]-shaped numpy array of non-normalized probabilities.
# 
#     Returns:
#         An [N, M]-shaped array of quantized probabilities such that
#         np.all(result.sum(axis=1) == resolution).
#     """
#     assert len(probs.shape) == 2
#     N, M = probs.shape
#     probs = probs / probs.sum(axis=1, keepdims=True)
#     result = np.zeros(probs.shape, np.int8)
#     range_N = np.arange(N, dtype=np.int32)
#     for _ in range(resolution):
#         sample = probs.argmax(axis=1)
#         result[range_N, sample] += 1
#         probs[range_N, sample] -= 1.0 / resolution
#     return result
# 
# Path: treecat/util.py
# @jit(nopython=True, cache=True)
# def sample_from_probs(probs):
#     """Sample from a vector of non-normalized probabilitites.
# 
#     Args:
#         probs: An [M]-shaped numpy array of non-normalized probabilities.
# 
#     Returns:
#         An integer in range(M).
#     """
#     cdf = probs.cumsum()
#     return (np.random.rand() * cdf[-1] < cdf).argmax()
# 
# Path: treecat/util.py
# def sample_from_probs2(probs, out=None):
#     """Sample from multiple vectors of non-normalized probabilities.
# 
#     Args:
#         probs: An [N, M]-shaped numpy array of non-normalized probabilities.
#         out: An optional destination for the result.
# 
#     Returns:
#         An [N]-shaped numpy array of integers in range(M).
#     """
#     # Adapted from https://stackoverflow.com/questions/40474436
#     assert len(probs.shape) == 2
#     cdf = probs.cumsum(axis=1)
#     u = np.random.rand(probs.shape[0], 1) * cdf[:, -1, np.newaxis]
#     return (u < cdf).argmax(axis=1, out=out)
# 
# Path: treecat/util.py
# def set_random_seed(seed):
#     """Set random seeds for both numpy and numba."""
#     np.random.seed(seed)
#     jit_random_seed(seed)


# Path: treecat/util_test.py
import itertools
import numpy as np
import pytest
from goftests import multinomial_goodness_of_fit
from treecat.testutil import make_seed
from treecat.util import parallel_map
from treecat.util import quantize_from_probs2
from treecat.util import sample_from_probs
from treecat.util import sample_from_probs2
from treecat.util import set_random_seed

    assert 1e-2 < gof


@pytest.mark.parametrize('size,resolution', [
    (1, 0),
    (1, 1),
    (1, 2),
    (2, 0),
    (2, 1),
    (2, 2),
    (2, 3),
    (2, 4),
    (2, 5),
    (3, 0),
    (3, 1),
    (3, 2),
    (3, 3),
    (3, 4),
    (3, 5),
    (4, 0),
    (4, 1),
    (4, 2),
    (4, 3),
    (4, 4),
    (5, 0),
    (5, 1),
    (5, 2),
    (5, 3),
])
def test_quantize_from_probs2(size, resolution):

    set_random_seed(make_seed(size, resolution))