def testNonNegativeIntegersWithSum(self):
result = combinatorics.uniform_non_negative_integers_with_sum(1, 0)
self.assertEqual(result, [0])
result = combinatorics.uniform_non_negative_integers_with_sum(2, 0)
self.assertEqual(result, [0, 0])
result = combinatorics.uniform_non_negative_integers_with_sum(3, 10)
self.assertEqual(sum(result), 10)
def split(self, count):
"""Splits the entropy and module counts up.
Args:
count: Integer >= 1; the split size.
Returns:
List of `SampleArgs` of length `count`, to be passed to create child
entities.
Raises:
ValueError: If it was not possible to use up all the entropy, for example,
all requested types were `WithValue`.
"""
num_child_modules = self.num_modules - 1
# Sample module counts at random
module_counts = combinatorics.uniform_non_negative_integers_with_sum(
count, num_child_modules)
if num_child_modules == 0:
if self.entropy > 0:
raise ValueError('Unused entropy')
entropies = np.zeros(count)
else:
entropies = self.entropy * np.random.dirichlet(
np.maximum(1e-9, module_counts))
sample_args = []
for i, num_modules in enumerate(module_counts):
child_sample_args = SampleArgs(
num_modules=num_modules, entropy=entropies[i])
sample_args.append(child_sample_args)
return sample_args
def split(self, args):
"""Splits the entropy and op counts up."""
non_integer_count = sum(not arg.is_Integer for arg in args)
assert non_integer_count <= self.count - 1
count_split = combinatorics.uniform_non_negative_integers_with_sum(
len(args), (self.count - 1) - non_integer_count)
for i, arg in enumerate(args):
if not arg.is_Integer:
count_split[i] += 1
if all(count == 0 for count in count_split):
assert self.entropy == 0
entropies = np.zeros(len(count_split))
else:
entropies = (np.random.dirichlet(np.maximum(1e-9, count_split)) *
self.entropy)
return [
_SampleArgs(op_count, entropy)
for op_count, entropy in zip(count_split, entropies)
]
def expanded_coefficient_counts(length, is_zero):
"""Generates list of integers for number of terms of given power.
Args:
length: Integer >= `sum(is_zero)`.
is_zero: List of booleans.
Returns:
List of non-negative integers of length `is_zero`, summing to `length`,
such that if `is_zero[i]` then `return_value[i] != 1`.
Raises:
ValueError: If assignment not possible.
"""
if length == 1 and all(is_zero):
raise ValueError('length=1 and all zero')
counts = np.asarray([0 if zero else 1 for zero in is_zero])
extra_needed = (length - sum(counts))
if extra_needed < 0:
raise ValueError('length={} cannot handle is_zero={}'
.format(length, is_zero))
extra = combinatorics.uniform_non_negative_integers_with_sum(
count=len(is_zero), sum_=extra_needed)
counts += np.asarray(extra)
# Tweak so that no zeros get "1".
while True:
bad_zeros = [
i for i in range(len(is_zero)) if is_zero[i] and counts[i] == 1
]
if not bad_zeros:
break
take_from = random.choice(bad_zeros)
add_to = random.choice(
[i for i in range(len(is_zero)) if counts[i] >= 1 and i != take_from])
counts[take_from] -= 1
counts[add_to] += 1
return counts
def _level_set_event(values, length, verb):
"""Generates `LevelSetEvent`; see _generate_sequence_event."""
counts = combinatorics.uniform_non_negative_integers_with_sum(
len(values), length)
counts_dict = dict(list(zip(values, counts)))
event = probability.CountLevelSetEvent(counts_dict)
shuffled_values = list(values)
random.shuffle(shuffled_values)
counts_and_values = [
'{} {}'.format(counts_dict[value], value) for value in shuffled_values
if counts_dict[value] > 0
]
counts_and_values = _word_series(counts_and_values)
template = random.choice([
'{verbing} {counts_and_values}',
])
verbing = _GERUNDS[verb]
event_description = template.format(counts_and_values=counts_and_values,
verbing=verbing)
return event, event_description
