# Copyright 2018 DeepMind Technologies Limited.
#
# 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.
"""Probability questions (sampling, independence, expectations, ...)."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import collections
import functools
import random
import string
# Dependency imports
from mathematics_dataset import example
from mathematics_dataset.modules import train_test_split
from mathematics_dataset.util import combinatorics
from mathematics_dataset.util import composition
from mathematics_dataset.util import display
from mathematics_dataset.util import probability
import numpy as np
from six.moves import range
from six.moves import zip
_LETTERS = string.ascii_lowercase
_MAX_FRAC_TRIVIAL_PROB = 0.1
# Maximum number of colours and objects in a bag.
_MAX_DISTINCT_LETTERS = 6
_MAX_TOTAL_LETTERS = 20
_MAX_LETTER_REPEAT = 10
_SWR_SAMPLE_COUNT = [2, 4]
_SWR_SAMPLE_COUNT_EXTRAPOLATE = [5, 5]
_GERUNDS = {
'pick': 'picking',
}
def _make_modules(is_train):
"""Returns modules, with split based on the boolean `is_train`."""
return {
'swr_p_sequence': functools.partial(
swr_prob_sequence, is_train=is_train, sample_range=_SWR_SAMPLE_COUNT),
'swr_p_level_set': functools.partial(
swr_prob_level_set, is_train=is_train,
sample_range=_SWR_SAMPLE_COUNT),
}
def train(entropy_fn):
"""Returns dict of training modules."""
del entropy_fn # unused
return _make_modules(is_train=True)
def test():
"""Returns dict of testing modules."""
return _make_modules(is_train=False)
def test_extra():
"""Returns dict of extrapolation testing modules."""
return {
'swr_p_sequence_more_samples': functools.partial(
swr_prob_sequence, is_train=None,
sample_range=_SWR_SAMPLE_COUNT_EXTRAPOLATE),
'swr_p_level_set_more_samples': functools.partial(
swr_prob_level_set, is_train=None,
sample_range=_SWR_SAMPLE_COUNT_EXTRAPOLATE),
}
def _sequence_event(values, length, verb):
"""Returns sequence (finite product) event.
Args:
values: List of values to sample from.
length: Length of the sequence to generate.
verb: Verb in infinitive form.
Returns:
Instance of `probability.FiniteProductEvent`, together with a text
description.
"""
del verb # unused
samples = [random.choice(values) for _ in range(length)]
events = [probability.DiscreteEvent([sample]) for sample in samples]
event = probability.FiniteProductEvent(events)
sequence = ''.join(str(sample) for sample in samples)
event_description = 'sequence {sequence}'.format(sequence=sequence)
return event, event_description
def _word_series(words, conjunction='and'):
"""Combines the words using commas and the final conjunction."""
len_words = len(words)
if len_words == 0:
return ''
if len_words == 1:
return words[0]
return '{} {} {}'.format(', '.join(words[:-1]), conjunction, words[-1])
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
LetterBag = collections.namedtuple(
'LetterBag',
('weights', 'random_variable', 'letters_distinct', 'bag_contents'))
def _sample_letter_bag(is_train, min_total):
"""Samples a "container of letters" and returns info on it."""
while True:
num_distinct_letters = random.randint(1, _MAX_DISTINCT_LETTERS)
num_letters_total = random.randint(
max(num_distinct_letters, min_total),
min(_MAX_TOTAL_LETTERS, num_distinct_letters * _MAX_LETTER_REPEAT))
letter_counts = combinatorics.uniform_positive_integers_with_sum(
num_distinct_letters, num_letters_total)
# Test/train split.
if (is_train is None
or train_test_split.is_train(sorted(letter_counts)) == is_train):
break
letters_distinct = random.sample(_LETTERS, num_distinct_letters)
weights = {i: 1 for i in range(num_letters_total)}
letters_with_repetition = []
for letter, count in zip(letters_distinct, letter_counts):
letters_with_repetition += [letter] * count
random.shuffle(letters_with_repetition)
random_variable = probability.DiscreteRandomVariable(
{i: letter for i, letter in enumerate(letters_with_repetition)})
if random.choice([False, True]):
bag_contents = ''.join(letters_with_repetition)
else:
letters_and_counts = [
'{}: {}'.format(letter, count)
for letter, count in zip(letters_distinct, letter_counts)]
bag_contents = '{' + ', '.join(letters_and_counts) + '}'
return LetterBag(
weights=weights,
random_variable=random_variable,
letters_distinct=letters_distinct,
bag_contents=bag_contents)
def _swr_space(is_train, sample_range):
"""Returns probability space for sampling without replacement."""
num_sampled = random.randint(*sample_range)
sample = _sample_letter_bag(is_train=is_train, min_total=num_sampled)
space = probability.SampleWithoutReplacementSpace(sample.weights, num_sampled)
random_variable = probability.FiniteProductRandomVariable(
[sample.random_variable] * num_sampled)
random_variable.description = (
str(display.StringNumber(num_sampled))
+ ' letters picked without replacement from '
+ sample.bag_contents)
return sample.letters_distinct, space, random_variable
def _sample_without_replacement_probability_question(
is_train, event_fn, sample_range):
"""Question for prob of some event when sampling without replacement."""
def too_big(event_in_space):
if isinstance(event_in_space, probability.SequenceEvent):
size = len(event_in_space.all_sequences())
else:
assert isinstance(event_in_space, probability.FiniteProductEvent)
size = np.prod([len(event.values) for event in event_in_space.events])
return size > int(2e5)
allow_trivial_prob = random.random() < _MAX_FRAC_TRIVIAL_PROB
while True:
distinct_letters, space, random_variable = _swr_space(
is_train, sample_range)
event, event_description = event_fn(
values=distinct_letters, length=space.n_samples, verb='pick')
event_in_space = random_variable.inverse(event)
if too_big(event_in_space):
continue
answer = space.probability(event_in_space)
if answer not in [0, 1] or allow_trivial_prob:
break
context = composition.Context()
template = random.choice([
'{random_variable_capitalize}. What is prob of {event}?',
'{random_variable_capitalize}. Give prob of {event}.',
'What is prob of {event} when {random_variable}?',
'Calculate prob of {event} when {random_variable}.',
])
question = example.question(
context,
template,
random_variable=random_variable.description,
random_variable_capitalize=(
str(random_variable.description).capitalize()),
event=event_description)
return example.Problem(question, answer)
def swr_prob_sequence(is_train, sample_range):
"""Probability of given sequence when sampling without replacement."""
return _sample_without_replacement_probability_question(
is_train=is_train, event_fn=_sequence_event, sample_range=sample_range)
def swr_prob_level_set(is_train, sample_range):
"""Probability of given level set when sampling without replacement."""
return _sample_without_replacement_probability_question(
is_train=is_train, event_fn=_level_set_event, sample_range=sample_range)
