# Source code for nltk.probability

```# -*- coding: utf-8 -*-
# Natural Language Toolkit: Probability and Statistics
#
# Copyright (C) 2001-2015 NLTK Project
# Author: Edward Loper <edloper@gmail.com>
#
# URL: <http://nltk.org/>

"""
Classes for representing and processing probabilistic information.

The ``FreqDist`` class is used to encode "frequency distributions",
which count the number of times that each outcome of an experiment
occurs.

The ``ProbDistI`` class defines a standard interface for "probability
distributions", which encode the probability of each outcome for an
experiment.  There are two types of probability distribution:

- "derived probability distributions" are created from frequency
distributions.  They attempt to model the probability distribution
that generated the frequency distribution.
- "analytic probability distributions" are created directly from
parameters (such as variance).

The ``ConditionalFreqDist`` class and ``ConditionalProbDistI`` interface
are used to encode conditional distributions.  Conditional probability
distributions can be derived or analytic; but currently the only
implementation of the ``ConditionalProbDistI`` interface is
``ConditionalProbDist``, a derived distribution.

"""
from __future__ import print_function, unicode_literals, division

import math
import random
import warnings
import array
from operator import itemgetter
from collections import defaultdict
from functools import reduce
from nltk import compat
from nltk.compat import Counter

from nltk.internals import raise_unorderable_types

_NINF = float('-1e300')

##//////////////////////////////////////////////////////
##  Frequency Distributions
##//////////////////////////////////////////////////////

@compat.python_2_unicode_compatible
[docs]class FreqDist(Counter):
"""
A frequency distribution for the outcomes of an experiment.  A
frequency distribution records the number of times each outcome of
an experiment has occurred.  For example, a frequency distribution
could be used to record the frequency of each word type in a
document.  Formally, a frequency distribution can be defined as a
function mapping from each sample to the number of times that
sample occurred as an outcome.

Frequency distributions are generally constructed by running a
number of experiments, and incrementing the count for a sample
every time it is an outcome of an experiment.  For example, the
following code will produce a frequency distribution that encodes
how often each word occurs in a text:

>>> from nltk.tokenize import word_tokenize
>>> from nltk.probability import FreqDist
>>> sent = 'This is an example sentence'
>>> fdist = FreqDist()
>>> for word in word_tokenize(sent):
...    fdist[word.lower()] += 1

An equivalent way to do this is with the initializer:

>>> fdist = FreqDist(word.lower() for word in word_tokenize(sent))

"""

[docs]    def __init__(self, samples=None):
"""
Construct a new frequency distribution.  If ``samples`` is
given, then the frequency distribution will be initialized
with the count of each object in ``samples``; otherwise, it
will be initialized to be empty.

In particular, ``FreqDist()`` returns an empty frequency
distribution; and ``FreqDist(samples)`` first creates an empty
frequency distribution, and then calls ``update`` with the
list ``samples``.

:param samples: The samples to initialize the frequency
distribution with.
:type samples: Sequence
"""
Counter.__init__(self, samples)

[docs]    def N(self):
"""
Return the total number of sample outcomes that have been
recorded by this FreqDist.  For the number of unique
sample values (or bins) with counts greater than zero, use
``FreqDist.B()``.

:rtype: int
"""
return sum(self.values())

[docs]    def B(self):
"""
Return the total number of sample values (or "bins") that
have counts greater than zero.  For the total
number of sample outcomes recorded, use ``FreqDist.N()``.
(FreqDist.B() is the same as len(FreqDist).)

:rtype: int
"""
return len(self)

[docs]    def hapaxes(self):
"""
Return a list of all samples that occur once (hapax legomena)

:rtype: list
"""
return [item for item in self if self[item] == 1]

[docs]    def Nr(self, r, bins=None):
return self.r_Nr(bins)[r]

[docs]    def r_Nr(self, bins=None):
"""
Return the dictionary mapping r to Nr, the number of samples with frequency r, where Nr > 0.

:type bins: int
:param bins: The number of possible sample outcomes.  ``bins``
is used to calculate Nr(0).  In particular, Nr(0) is
``bins-self.B()``.  If ``bins`` is not specified, it
defaults to ``self.B()`` (so Nr(0) will be 0).
:rtype: int
"""

_r_Nr = defaultdict(int)
for count in self.values():
_r_Nr[count] += 1

# Special case for Nr[0]:
_r_Nr[0] = bins - self.B() if bins is not None else 0

return _r_Nr

def _cumulative_frequencies(self, samples):
"""
Return the cumulative frequencies of the specified samples.
If no samples are specified, all counts are returned, starting
with the largest.

:param samples: the samples whose frequencies should be returned.
:type samples: any
:rtype: list(float)
"""
cf = 0.0
for sample in samples:
cf += self[sample]
yield cf

# slightly odd nomenclature freq() if FreqDist does counts and ProbDist does probs,
# here, freq() does probs
[docs]    def freq(self, sample):
"""
Return the frequency of a given sample.  The frequency of a
sample is defined as the count of that sample divided by the
total number of sample outcomes that have been recorded by
this FreqDist.  The count of a sample is defined as the
number of times that sample outcome was recorded by this
FreqDist.  Frequencies are always real numbers in the range
[0, 1].

:param sample: the sample whose frequency
should be returned.
:type sample: any
:rtype: float
"""
if self.N() == 0:
return 0
return self[sample] / self.N()

[docs]    def max(self):
"""
Return the sample with the greatest number of outcomes in this
frequency distribution.  If two or more samples have the same
number of outcomes, return one of them; which sample is
returned is undefined.  If no outcomes have occurred in this
frequency distribution, return None.

:return: The sample with the maximum number of outcomes in this
frequency distribution.
:rtype: any or None
"""
if len(self) == 0:
raise ValueError('A FreqDist must have at least one sample before max is defined.')
return self.most_common(1)[0][0]

[docs]    def plot(self, *args, **kwargs):
"""
Plot samples from the frequency distribution
displaying the most frequent sample first.  If an integer
parameter is supplied, stop after this many samples have been
plotted.  For a cumulative plot, specify cumulative=True.
(Requires Matplotlib to be installed.)

:param title: The title for the graph
:type title: str
:param cumulative: A flag to specify whether the plot is cumulative (default = False)
:type title: bool
"""
try:
from matplotlib import pylab
except ImportError:
raise ValueError('The plot function requires matplotlib to be installed.'
'See http://matplotlib.org/')

if len(args) == 0:
args = [len(self)]
samples = [item for item, _ in self.most_common(*args)]

cumulative = _get_kwarg(kwargs, 'cumulative', False)
if cumulative:
freqs = list(self._cumulative_frequencies(samples))
ylabel = "Cumulative Counts"
else:
freqs = [self[sample] for sample in samples]
ylabel = "Counts"
# percents = [f * 100 for f in freqs]  only in ProbDist?

pylab.grid(True, color="silver")
if not "linewidth" in kwargs:
kwargs["linewidth"] = 2
if "title" in kwargs:
pylab.title(kwargs["title"])
del kwargs["title"]
pylab.plot(freqs, **kwargs)
pylab.xticks(range(len(samples)), [compat.text_type(s) for s in samples], rotation=90)
pylab.xlabel("Samples")
pylab.ylabel(ylabel)
pylab.show()

[docs]    def tabulate(self, *args, **kwargs):
"""
Tabulate the given samples from the frequency distribution (cumulative),
displaying the most frequent sample first.  If an integer
parameter is supplied, stop after this many samples have been
plotted.

:param samples: The samples to plot (default is all samples)
:type samples: list
:param cumulative: A flag to specify whether the freqs are cumulative (default = False)
:type title: bool
"""
if len(args) == 0:
args = [len(self)]
samples = [item for item, _ in self.most_common(*args)]

cumulative = _get_kwarg(kwargs, 'cumulative', False)
if cumulative:
freqs = list(self._cumulative_frequencies(samples))
else:
freqs = [self[sample] for sample in samples]
# percents = [f * 100 for f in freqs]  only in ProbDist?

width = max(len("%s" % s) for s in samples)
width = max(width, max(len("%d" % f) for f in freqs))

for i in range(len(samples)):
print("%*s" % (width, samples[i]), end=' ')
print()
for i in range(len(samples)):
print("%*d" % (width, freqs[i]), end=' ')
print()

[docs]    def copy(self):
"""
Create a copy of this frequency distribution.

:rtype: FreqDist
"""
return self.__class__(self)

# Mathematical operatiors

"""

>>> FreqDist('abbb') + FreqDist('bcc')
FreqDist({'b': 4, 'c': 2, 'a': 1})

"""

def __sub__(self, other):
"""
Subtract count, but keep only results with positive counts.

>>> FreqDist('abbbc') - FreqDist('bccd')
FreqDist({'b': 2, 'a': 1})

"""
return self.__class__(super(FreqDist, self).__sub__(other))

def __or__(self, other):
"""
Union is the maximum of value in either of the input counters.

>>> FreqDist('abbb') | FreqDist('bcc')
FreqDist({'b': 3, 'c': 2, 'a': 1})

"""
return self.__class__(super(FreqDist, self).__or__(other))

def __and__(self, other):
"""
Intersection is the minimum of corresponding counts.

>>> FreqDist('abbb') & FreqDist('bcc')
FreqDist({'b': 1})

"""
return self.__class__(super(FreqDist, self).__and__(other))

def __le__(self, other):
if not isinstance(other, FreqDist):
raise_unorderable_types("<=", self, other)
return set(self).issubset(other) and all(self[key] <= other[key] for key in self)

# @total_ordering doesn't work here, since the class inherits from a builtin class
__ge__ = lambda self, other: not self <= other or self == other
__lt__ = lambda self, other: self <= other and not self == other
__gt__ = lambda self, other: not self <= other

def __repr__(self):
"""
Return a string representation of this FreqDist.

:rtype: string
"""
return self.pformat()

[docs]    def pprint(self, maxlen=10, stream=None):
"""
Print a string representation of this FreqDist to 'stream'

:param maxlen: The maximum number of items to print
:type maxlen: int
:param stream: The stream to print to. stdout by default
"""
print(self.pformat(maxlen=maxlen), file=stream)

[docs]    def pformat(self, maxlen=10):
"""
Return a string representation of this FreqDist.

:param maxlen: The maximum number of items to display
:type maxlen: int
:rtype: string
"""
items = ['{0!r}: {1!r}'.format(*item) for item in self.most_common(maxlen)]
if len(self) > maxlen:
items.append('...')
return 'FreqDist({{{0}}})'.format(', '.join(items))

def __str__(self):
"""
Return a string representation of this FreqDist.

:rtype: string
"""
return '<FreqDist with %d samples and %d outcomes>' % (len(self), self.N())

##//////////////////////////////////////////////////////
##  Probability Distributions
##//////////////////////////////////////////////////////

[docs]class ProbDistI(object):
"""
A probability distribution for the outcomes of an experiment.  A
probability distribution specifies how likely it is that an
experiment will have any given outcome.  For example, a
probability distribution could be used to predict the probability
that a token in a document will have a given type.  Formally, a
probability distribution can be defined as a function mapping from
samples to nonnegative real numbers, such that the sum of every
number in the function's range is 1.0.  A ``ProbDist`` is often
used to model the probability distribution of the experiment used
to generate a frequency distribution.
"""
SUM_TO_ONE = True
"""True if the probabilities of the samples in this probability
distribution will always sum to one."""

[docs]    def __init__(self):
if self.__class__ == ProbDistI:
raise NotImplementedError("Interfaces can't be instantiated")

[docs]    def prob(self, sample):
"""
Return the probability for a given sample.  Probabilities
are always real numbers in the range [0, 1].

:param sample: The sample whose probability
should be returned.
:type sample: any
:rtype: float
"""
raise NotImplementedError()

[docs]    def logprob(self, sample):
"""
Return the base 2 logarithm of the probability for a given sample.

:param sample: The sample whose probability
should be returned.
:type sample: any
:rtype: float
"""
# Default definition, in terms of prob()
p = self.prob(sample)
return (math.log(p, 2) if p != 0 else _NINF)

[docs]    def max(self):
"""
Return the sample with the greatest probability.  If two or
more samples have the same probability, return one of them;
which sample is returned is undefined.

:rtype: any
"""
raise NotImplementedError()

[docs]    def samples(self):
"""
Return a list of all samples that have nonzero probabilities.
Use ``prob`` to find the probability of each sample.

:rtype: list
"""
raise NotImplementedError()

# cf self.SUM_TO_ONE
[docs]    def discount(self):
"""
Return the ratio by which counts are discounted on average: c*/c

:rtype: float
"""
return 0.0

# Subclasses should define more efficient implementations of this,
# where possible.
[docs]    def generate(self):
"""
Return a randomly selected sample from this probability distribution.
The probability of returning each sample ``samp`` is equal to
``self.prob(samp)``.
"""
p = random.random()
p_init = p
for sample in self.samples():
p -= self.prob(sample)
if p <= 0: return sample
# allow for some rounding error:
if p < .0001:
return sample
# we *should* never get here
if self.SUM_TO_ONE:
warnings.warn("Probability distribution %r sums to %r; generate()"
" is returning an arbitrary sample." % (self, p_init-p))
return random.choice(list(self.samples()))

@compat.python_2_unicode_compatible
[docs]class UniformProbDist(ProbDistI):
"""
A probability distribution that assigns equal probability to each
sample in a given set; and a zero probability to all other
samples.
"""
[docs]    def __init__(self, samples):
"""
Construct a new uniform probability distribution, that assigns
equal probability to each sample in ``samples``.

:param samples: The samples that should be given uniform
probability.
:type samples: list
:raise ValueError: If ``samples`` is empty.
"""
if len(samples) == 0:
raise ValueError('A Uniform probability distribution must '+
'have at least one sample.')
self._sampleset = set(samples)
self._prob = 1.0/len(self._sampleset)
self._samples = list(self._sampleset)

[docs]    def prob(self, sample):
return (self._prob if sample in self._sampleset else 0)

[docs]    def max(self):
return self._samples[0]

[docs]    def samples(self):
return self._samples

def __repr__(self):
return '<UniformProbDist with %d samples>' % len(self._sampleset)

@compat.python_2_unicode_compatible
class RandomProbDist(ProbDistI):
"""
Generates a random probability distribution whereby each sample
will be between 0 and 1 with equal probability (uniform random distribution.
Also called a continuous uniform distribution).
"""
def __init__(self, samples):
if len(samples) == 0:
raise ValueError('A probability distribution must '+
'have at least one sample.')
self._probs = self.unirand(samples)
self._samples = list(self._probs.keys())

@classmethod
def unirand(cls, samples):
"""
The key function that creates a randomized initial distribution
that still sums to 1. Set as a dictionary of prob values so that
it can still be passed to MutableProbDist and called with identical
syntax to UniformProbDist
"""
randrow = [random.random() for i in range(len(samples))]
total = sum(randrow)
for i, x in enumerate(randrow):
randrow[i] = x/total

total = sum(randrow)
if total != 1:
#this difference, if present, is so small (near NINF) that it
#can be subtracted from any element without risking probs not (0 1)
randrow[-1] -= total - 1

return dict((s, randrow[i]) for i, s in enumerate(samples))

def prob(self, sample):
return self._probs.get(sample, 0)

def samples(self):
return self._samples

def __repr__(self):
return '<RandomUniformProbDist with %d samples>' %len(self._probs)

@compat.python_2_unicode_compatible
[docs]class DictionaryProbDist(ProbDistI):
"""
A probability distribution whose probabilities are directly
specified by a given dictionary.  The given dictionary maps
samples to probabilities.
"""
[docs]    def __init__(self, prob_dict=None, log=False, normalize=False):
"""
Construct a new probability distribution from the given
dictionary, which maps values to probabilities (or to log
probabilities, if ``log`` is true).  If ``normalize`` is
true, then the probability values are scaled by a constant
factor such that they sum to 1.

If called without arguments, the resulting probability
distribution assigns zero probability to all values.
"""

self._prob_dict = (prob_dict.copy() if prob_dict is not None else {})
self._log = log

# Normalize the distribution, if requested.
if normalize:
if len(prob_dict) == 0:
raise ValueError('A DictionaryProbDist must have at least one sample ' +
'before it can be normalized.')
if log:
value_sum = sum_logs(list(self._prob_dict.values()))
if value_sum <= _NINF:
logp = math.log(1.0/len(prob_dict), 2)
for x in prob_dict:
self._prob_dict[x] = logp
else:
for (x, p) in self._prob_dict.items():
self._prob_dict[x] -= value_sum
else:
value_sum = sum(self._prob_dict.values())
if value_sum == 0:
p = 1.0/len(prob_dict)
for x in prob_dict:
self._prob_dict[x] = p
else:
norm_factor = 1.0/value_sum
for (x, p) in self._prob_dict.items():
self._prob_dict[x] *= norm_factor

[docs]    def prob(self, sample):
if self._log:
return (2**(self._prob_dict[sample]) if sample in self._prob_dict else 0)
else:
return self._prob_dict.get(sample, 0)

[docs]    def logprob(self, sample):
if self._log:
return self._prob_dict.get(sample, _NINF)
else:
if sample not in self._prob_dict: return _NINF
elif self._prob_dict[sample] == 0: return _NINF
else: return math.log(self._prob_dict[sample], 2)

[docs]    def max(self):
if not hasattr(self, '_max'):
self._max = max((p,v) for (v,p) in self._prob_dict.items())[1]
return self._max
[docs]    def samples(self):
return self._prob_dict.keys()
def __repr__(self):
return '<ProbDist with %d samples>' % len(self._prob_dict)

@compat.python_2_unicode_compatible
[docs]class MLEProbDist(ProbDistI):
"""
The maximum likelihood estimate for the probability distribution
of the experiment used to generate a frequency distribution.  The
"maximum likelihood estimate" approximates the probability of
each sample as the frequency of that sample in the frequency
distribution.
"""
[docs]    def __init__(self, freqdist, bins=None):
"""
Use the maximum likelihood estimate to create a probability
distribution for the experiment used to generate ``freqdist``.

:type freqdist: FreqDist
:param freqdist: The frequency distribution that the
probability estimates should be based on.
"""
self._freqdist = freqdist

[docs]    def freqdist(self):
"""
Return the frequency distribution that this probability
distribution is based on.

:rtype: FreqDist
"""
return self._freqdist

[docs]    def prob(self, sample):
return self._freqdist.freq(sample)

[docs]    def max(self):
return self._freqdist.max()

[docs]    def samples(self):
return self._freqdist.keys()

def __repr__(self):
"""
:rtype: str
:return: A string representation of this ``ProbDist``.
"""
return '<MLEProbDist based on %d samples>' % self._freqdist.N()

@compat.python_2_unicode_compatible
[docs]class LidstoneProbDist(ProbDistI):
"""
The Lidstone estimate for the probability distribution of the
experiment used to generate a frequency distribution.  The
"Lidstone estimate" is parameterized by a real number *gamma*,
which typically ranges from 0 to 1.  The Lidstone estimate
approximates the probability of a sample with count *c* from an
experiment with *N* outcomes and *B* bins as
``c+gamma)/(N+B*gamma)``.  This is equivalent to adding
*gamma* to the count for each bin, and taking the maximum
likelihood estimate of the resulting frequency distribution.
"""
SUM_TO_ONE = False
[docs]    def __init__(self, freqdist, gamma, bins=None):
"""
Use the Lidstone estimate to create a probability distribution
for the experiment used to generate ``freqdist``.

:type freqdist: FreqDist
:param freqdist: The frequency distribution that the
probability estimates should be based on.
:type gamma: float
:param gamma: A real number used to parameterize the
estimate.  The Lidstone estimate is equivalent to adding
*gamma* to the count for each bin, and taking the
maximum likelihood estimate of the resulting frequency
distribution.
:type bins: int
:param bins: The number of sample values that can be generated
by the experiment that is described by the probability
distribution.  This value must be correctly set for the
probabilities of the sample values to sum to one.  If
``bins`` is not specified, it defaults to ``freqdist.B()``.
"""
if (bins == 0) or (bins is None and freqdist.N() == 0):
name = self.__class__.__name__[:-8]
raise ValueError('A %s probability distribution ' % name +
'must have at least one bin.')
if (bins is not None) and (bins < freqdist.B()):
name = self.__class__.__name__[:-8]
raise ValueError('\nThe number of bins in a %s distribution ' % name +
'(%d) must be greater than or equal to\n' % bins +
'the number of bins in the FreqDist used ' +
'to create it (%d).' % freqdist.B())

self._freqdist = freqdist
self._gamma = float(gamma)
self._N = self._freqdist.N()

if bins is None:
bins = freqdist.B()
self._bins = bins

self._divisor = self._N + bins * gamma
if self._divisor == 0.0:
# In extreme cases we force the probability to be 0,
# which it will be, since the count will be 0:
self._gamma = 0
self._divisor = 1

[docs]    def freqdist(self):
"""
Return the frequency distribution that this probability
distribution is based on.

:rtype: FreqDist
"""
return self._freqdist

[docs]    def prob(self, sample):
c = self._freqdist[sample]
return (c + self._gamma) / self._divisor

[docs]    def max(self):
# For Lidstone distributions, probability is monotonic with
# frequency, so the most probable sample is the one that
# occurs most frequently.
return self._freqdist.max()

[docs]    def samples(self):
return self._freqdist.keys()

[docs]    def discount(self):
gb = self._gamma * self._bins
return gb / (self._N + gb)

def __repr__(self):
"""
Return a string representation of this ``ProbDist``.

:rtype: str
"""
return '<LidstoneProbDist based on %d samples>' % self._freqdist.N()

@compat.python_2_unicode_compatible
[docs]class LaplaceProbDist(LidstoneProbDist):
"""
The Laplace estimate for the probability distribution of the
experiment used to generate a frequency distribution.  The
"Laplace estimate" approximates the probability of a sample with
count *c* from an experiment with *N* outcomes and *B* bins as
*(c+1)/(N+B)*.  This is equivalent to adding one to the count for
each bin, and taking the maximum likelihood estimate of the
resulting frequency distribution.
"""
[docs]    def __init__(self, freqdist, bins=None):
"""
Use the Laplace estimate to create a probability distribution
for the experiment used to generate ``freqdist``.

:type freqdist: FreqDist
:param freqdist: The frequency distribution that the
probability estimates should be based on.
:type bins: int
:param bins: The number of sample values that can be generated
by the experiment that is described by the probability
distribution.  This value must be correctly set for the
probabilities of the sample values to sum to one.  If
``bins`` is not specified, it defaults to ``freqdist.B()``.
"""
LidstoneProbDist.__init__(self, freqdist, 1, bins)

def __repr__(self):
"""
:rtype: str
:return: A string representation of this ``ProbDist``.
"""
return '<LaplaceProbDist based on %d samples>' % self._freqdist.N()

@compat.python_2_unicode_compatible
[docs]class ELEProbDist(LidstoneProbDist):
"""
The expected likelihood estimate for the probability distribution
of the experiment used to generate a frequency distribution.  The
"expected likelihood estimate" approximates the probability of a
sample with count *c* from an experiment with *N* outcomes and
*B* bins as *(c+0.5)/(N+B/2)*.  This is equivalent to adding 0.5
to the count for each bin, and taking the maximum likelihood
estimate of the resulting frequency distribution.
"""
[docs]    def __init__(self, freqdist, bins=None):
"""
Use the expected likelihood estimate to create a probability
distribution for the experiment used to generate ``freqdist``.

:type freqdist: FreqDist
:param freqdist: The frequency distribution that the
probability estimates should be based on.
:type bins: int
:param bins: The number of sample values that can be generated
by the experiment that is described by the probability
distribution.  This value must be correctly set for the
probabilities of the sample values to sum to one.  If
``bins`` is not specified, it defaults to ``freqdist.B()``.
"""
LidstoneProbDist.__init__(self, freqdist, 0.5, bins)

def __repr__(self):
"""
Return a string representation of this ``ProbDist``.

:rtype: str
"""
return '<ELEProbDist based on %d samples>' % self._freqdist.N()

@compat.python_2_unicode_compatible
[docs]class HeldoutProbDist(ProbDistI):
"""
The heldout estimate for the probability distribution of the
experiment used to generate two frequency distributions.  These
two frequency distributions are called the "heldout frequency
distribution" and the "base frequency distribution."  The
"heldout estimate" uses uses the "heldout frequency
distribution" to predict the probability of each sample, given its
frequency in the "base frequency distribution".

In particular, the heldout estimate approximates the probability
for a sample that occurs *r* times in the base distribution as
the average frequency in the heldout distribution of all samples
that occur *r* times in the base distribution.

This average frequency is *Tr[r]/(Nr[r].N)*, where:

- *Tr[r]* is the total count in the heldout distribution for
all samples that occur *r* times in the base distribution.
- *Nr[r]* is the number of samples that occur *r* times in
the base distribution.
- *N* is the number of outcomes recorded by the heldout
frequency distribution.

In order to increase the efficiency of the ``prob`` member
function, *Tr[r]/(Nr[r].N)* is precomputed for each value of *r*
when the ``HeldoutProbDist`` is created.

:type _estimate: list(float)
:ivar _estimate: A list mapping from *r*, the number of
times that a sample occurs in the base distribution, to the
probability estimate for that sample.  ``_estimate[r]`` is
calculated by finding the average frequency in the heldout
distribution of all samples that occur *r* times in the base
distribution.  In particular, ``_estimate[r]`` =
*Tr[r]/(Nr[r].N)*.
:type _max_r: int
:ivar _max_r: The maximum number of times that any sample occurs
in the base distribution.  ``_max_r`` is used to decide how
large ``_estimate`` must be.
"""
SUM_TO_ONE = False
[docs]    def __init__(self, base_fdist, heldout_fdist, bins=None):
"""
Use the heldout estimate to create a probability distribution
for the experiment used to generate ``base_fdist`` and
``heldout_fdist``.

:type base_fdist: FreqDist
:param base_fdist: The base frequency distribution.
:type heldout_fdist: FreqDist
:param heldout_fdist: The heldout frequency distribution.
:type bins: int
:param bins: The number of sample values that can be generated
by the experiment that is described by the probability
distribution.  This value must be correctly set for the
probabilities of the sample values to sum to one.  If
``bins`` is not specified, it defaults to ``freqdist.B()``.
"""

self._base_fdist = base_fdist
self._heldout_fdist = heldout_fdist

# The max number of times any sample occurs in base_fdist.
self._max_r = base_fdist[base_fdist.max()]

# Calculate Tr, Nr, and N.
Tr = self._calculate_Tr()
r_Nr = base_fdist.r_Nr(bins)
Nr = [r_Nr[r] for r in range(self._max_r+1)]
N = heldout_fdist.N()

# Use Tr, Nr, and N to compute the probability estimate for
# each value of r.
self._estimate = self._calculate_estimate(Tr, Nr, N)

def _calculate_Tr(self):
"""
Return the list *Tr*, where *Tr[r]* is the total count in
``heldout_fdist`` for all samples that occur *r*
times in ``base_fdist``.

:rtype: list(float)
"""
Tr = [0.0] * (self._max_r+1)
for sample in self._heldout_fdist:
r = self._base_fdist[sample]
Tr[r] += self._heldout_fdist[sample]
return Tr

def _calculate_estimate(self, Tr, Nr, N):
"""
Return the list *estimate*, where *estimate[r]* is the probability
estimate for any sample that occurs *r* times in the base frequency
distribution.  In particular, *estimate[r]* is *Tr[r]/(N[r].N)*.
In the special case that *N[r]=0*, *estimate[r]* will never be used;
so we define *estimate[r]=None* for those cases.

:rtype: list(float)
:type Tr: list(float)
:param Tr: the list *Tr*, where *Tr[r]* is the total count in
the heldout distribution for all samples that occur *r*
times in base distribution.
:type Nr: list(float)
:param Nr: The list *Nr*, where *Nr[r]* is the number of
samples that occur *r* times in the base distribution.
:type N: int
:param N: The total number of outcomes recorded by the heldout
frequency distribution.
"""
estimate = []
for r in range(self._max_r+1):
if Nr[r] == 0: estimate.append(None)
else: estimate.append(Tr[r]/(Nr[r]*N))
return estimate

[docs]    def base_fdist(self):
"""
Return the base frequency distribution that this probability
distribution is based on.

:rtype: FreqDist
"""
return self._base_fdist

[docs]    def heldout_fdist(self):
"""
Return the heldout frequency distribution that this
probability distribution is based on.

:rtype: FreqDist
"""
return self._heldout_fdist

[docs]    def samples(self):
return self._base_fdist.keys()

[docs]    def prob(self, sample):
# Use our precomputed probability estimate.
r = self._base_fdist[sample]
return self._estimate[r]

[docs]    def max(self):
# Note: the Heldout estimation is *not* necessarily monotonic;
# so this implementation is currently broken.  However, it
# should give the right answer *most* of the time. :)
return self._base_fdist.max()

[docs]    def discount(self):
raise NotImplementedError()

def __repr__(self):
"""
:rtype: str
:return: A string representation of this ``ProbDist``.
"""
s = '<HeldoutProbDist: %d base samples; %d heldout samples>'
return s % (self._base_fdist.N(), self._heldout_fdist.N())

@compat.python_2_unicode_compatible
[docs]class CrossValidationProbDist(ProbDistI):
"""
The cross-validation estimate for the probability distribution of
the experiment used to generate a set of frequency distribution.
The "cross-validation estimate" for the probability of a sample
is found by averaging the held-out estimates for the sample in
each pair of frequency distributions.
"""
SUM_TO_ONE = False
[docs]    def __init__(self, freqdists, bins):
"""
Use the cross-validation estimate to create a probability
distribution for the experiment used to generate
``freqdists``.

:type freqdists: list(FreqDist)
:param freqdists: A list of the frequency distributions
generated by the experiment.
:type bins: int
:param bins: The number of sample values that can be generated
by the experiment that is described by the probability
distribution.  This value must be correctly set for the
probabilities of the sample values to sum to one.  If
``bins`` is not specified, it defaults to ``freqdist.B()``.
"""
self._freqdists = freqdists

# Create a heldout probability distribution for each pair of
# frequency distributions in freqdists.
self._heldout_probdists = []
for fdist1 in freqdists:
for fdist2 in freqdists:
if fdist1 is not fdist2:
probdist = HeldoutProbDist(fdist1, fdist2, bins)
self._heldout_probdists.append(probdist)

[docs]    def freqdists(self):
"""
Return the list of frequency distributions that this ``ProbDist`` is based on.

:rtype: list(FreqDist)
"""
return self._freqdists

[docs]    def samples(self):
# [xx] nb: this is not too efficient
return set(sum([list(fd) for fd in self._freqdists], []))

[docs]    def prob(self, sample):
# Find the average probability estimate returned by each
# heldout distribution.
prob = 0.0
for heldout_probdist in self._heldout_probdists:
prob += heldout_probdist.prob(sample)
return prob/len(self._heldout_probdists)

[docs]    def discount(self):
raise NotImplementedError()

def __repr__(self):
"""
Return a string representation of this ``ProbDist``.

:rtype: str
"""
return '<CrossValidationProbDist: %d-way>' % len(self._freqdists)

@compat.python_2_unicode_compatible
[docs]class WittenBellProbDist(ProbDistI):
"""
The Witten-Bell estimate of a probability distribution. This distribution
allocates uniform probability mass to as yet unseen events by using the
number of events that have only been seen once. The probability mass
reserved for unseen events is equal to *T / (N + T)*
where *T* is the number of observed event types and *N* is the total
number of observed events. This equates to the maximum likelihood estimate
of a new type event occurring. The remaining probability mass is discounted
such that all probability estimates sum to one, yielding:

- *p = T / Z (N + T)*, if count = 0
- *p = c / (N + T)*, otherwise
"""

[docs]    def __init__(self, freqdist, bins=None):
"""
Creates a distribution of Witten-Bell probability estimates.  This
distribution allocates uniform probability mass to as yet unseen
events by using the number of events that have only been seen once. The
probability mass reserved for unseen events is equal to *T / (N + T)*
where *T* is the number of observed event types and *N* is the total
number of observed events. This equates to the maximum likelihood
estimate of a new type event occurring. The remaining probability mass
is discounted such that all probability estimates sum to one,
yielding:

- *p = T / Z (N + T)*, if count = 0
- *p = c / (N + T)*, otherwise

The parameters *T* and *N* are taken from the ``freqdist`` parameter
(the ``B()`` and ``N()`` values). The normalizing factor *Z* is
calculated using these values along with the ``bins`` parameter.

:param freqdist: The frequency counts upon which to base the
estimation.
:type freqdist: FreqDist
:param bins: The number of possible event types. This must be at least
as large as the number of bins in the ``freqdist``. If None, then
it's assumed to be equal to that of the ``freqdist``
:type bins: int
"""
assert bins is None or bins >= freqdist.B(),\
'bins parameter must not be less than %d=freqdist.B()' % freqdist.B()
if bins is None:
bins = freqdist.B()
self._freqdist = freqdist
self._T = self._freqdist.B()
self._Z = bins - self._freqdist.B()
self._N = self._freqdist.N()
# self._P0 is P(0), precalculated for efficiency:
if self._N==0:
# if freqdist is empty, we approximate P(0) by a UniformProbDist:
self._P0 = 1.0 / self._Z
else:
self._P0 = self._T / (self._Z * (self._N + self._T))

[docs]    def prob(self, sample):
# inherit docs from ProbDistI
c = self._freqdist[sample]
return (c / (self._N + self._T) if c != 0 else self._P0)

[docs]    def max(self):
return self._freqdist.max()

[docs]    def samples(self):
return self._freqdist.keys()

[docs]    def freqdist(self):
return self._freqdist

[docs]    def discount(self):
raise NotImplementedError()

def __repr__(self):
"""
Return a string representation of this ``ProbDist``.

:rtype: str
"""
return '<WittenBellProbDist based on %d samples>' % self._freqdist.N()

##//////////////////////////////////////////////////////
##  Good-Turing Probability Distributions
##//////////////////////////////////////////////////////

# Good-Turing frequency estimation was contributed by Alan Turing and
# his statistical assistant I.J. Good, during their collaboration in
# the WWII.  It is a statistical technique for predicting the
# probability of occurrence of objects belonging to an unknown number
# of species, given past observations of such objects and their
# species. (In drawing balls from an urn, the 'objects' would be balls
# and the 'species' would be the distinct colors of the balls (finite
# but unknown in number).
#
# Good-Turing method calculates the probability mass to assign to
# events with zero or low counts based on the number of events with
# higher counts. It does so by using the adjusted count *c\**:
#
#     - *c\* = (c + 1) N(c + 1) / N(c)*   for c >= 1
#     - *things with frequency zero in training* = N(1)  for c == 0
#
# where *c* is the original count, *N(i)* is the number of event types
# observed with count *i*. We can think the count of unseen as the count
# of frequency one (see Jurafsky & Martin 2nd Edition, p101).
#
# This method is problematic because the situation ``N(c+1) == 0``
# is quite common in the original Good-Turing estimation; smoothing or
# interpolation of *N(i)* values is essential in practice.
#
# Bill Gale and Geoffrey Sampson present a simple and effective approach,
# Simple Good-Turing.  As a smoothing curve they simply use a power curve:
#
#     Nr = a*r^b (with b < -1 to give the appropriate hyperbolic
#     relationship)
#
# They estimate a and b by simple linear regression technique on the
# logarithmic form of the equation:
#
#     log Nr = a + b*log(r)
#
# However, they suggest that such a simple curve is probably only
# appropriate for high values of r. For low values of r, they use the
# measured Nr directly.  (see M&S, p.213)
#
# Gale and Sampson propose to use r while the difference between r and
# r* is 1.96 greater than the standard deviation, and switch to r* if
# it is less or equal:
#
#     |r - r*| > 1.96 * sqrt((r + 1)^2 (Nr+1 / Nr^2) (1 + Nr+1 / Nr))
#
# The 1.96 coefficient correspond to a 0.05 significance criterion,
# some implementations can use a coefficient of 1.65 for a 0.1
# significance criterion.
#

##//////////////////////////////////////////////////////
##  Simple Good-Turing Probablity Distributions
##//////////////////////////////////////////////////////

@compat.python_2_unicode_compatible
[docs]class SimpleGoodTuringProbDist(ProbDistI):
"""
SimpleGoodTuring ProbDist approximates from frequency to frequency of
frequency into a linear line under log space by linear regression.
Details of Simple Good-Turing algorithm can be found in:

- Good Turing smoothing without tears" (Gale & Sampson 1995),
Journal of Quantitative Linguistics, vol. 2 pp. 217-237.
- "Speech and Language Processing (Jurafsky & Martin),
2nd Edition, Chapter 4.5 p103 (log(Nc) =  a + b*log(c))
- http://www.grsampson.net/RGoodTur.html

Given a set of pair (xi, yi),  where the xi denotes the frequency and
yi denotes the frequency of frequency, we want to minimize their
square variation. E(x) and E(y) represent the mean of xi and yi.

- slope: b = sigma ((xi-E(x)(yi-E(y))) / sigma ((xi-E(x))(xi-E(x)))
- intercept: a = E(y) - b.E(x)
"""
SUM_TO_ONE = False
[docs]    def __init__(self, freqdist, bins=None):
"""
:param freqdist: The frequency counts upon which to base the
estimation.
:type freqdist: FreqDist
:param bins: The number of possible event types. This must be
larger than the number of bins in the ``freqdist``. If None,
then it's assumed to be equal to ``freqdist``.B() + 1
:type bins: int
"""
assert bins is None or bins > freqdist.B(),\
'bins parameter must not be less than %d=freqdist.B()+1' % (freqdist.B()+1)
if bins is None:
bins = freqdist.B() + 1
self._freqdist = freqdist
self._bins = bins
r, nr = self._r_Nr()
self.find_best_fit(r, nr)
self._switch(r, nr)
self._renormalize(r, nr)

def _r_Nr_non_zero(self):
r_Nr = self._freqdist.r_Nr()
del r_Nr[0]
return r_Nr

def _r_Nr(self):
"""
Split the frequency distribution in two list (r, Nr), where Nr(r) > 0
"""
nonzero = self._r_Nr_non_zero()

if not nonzero:
return [], []
return zip(*sorted(nonzero.items()))

```