def expanded_samplespace(d, alphabets=None, union=True):
"""
Returns a new distribution with an expanded sample space.
Expand the sample space so that it is the Cartesian product of the
alphabets for each random variable. Note, only the effective alphabet of
each random variable is used. So if one index in an outcome only has the
value 1, then its alphabet is [1], and not [0,1] for example.
Parameters
----------
d : distribution
The distribution used to create the pruned distribution.
alphabets : list
A list of alphabets, with length equal to the outcome length in `d`.
Each alphabet specifies the alphabet to be used for a single index
random variable. The sample space of the new distribution will be the
Cartesian product of these alphabets.
union : bool
If True, then the alphabet for each random variable is unioned.
The unioned alphabet is then used for each random variable.
Returns
-------
ed : distribution
The distribution with an expanded sample space.
Notes
-----
The default constructor for Distribution will create a Cartesian product
sample space if not sample space is provided.
"""
joint = d.is_joint()
if alphabets is None:
# Note, we sort the alphabets now, so we are possibly changing the
# order of the original sample space.
alphabets = list(map(sorted, d.alphabet))
elif joint and len(alphabets) != d.outcome_length():
L = len(alphabets)
raise Exception("You need to provide {0} alphabets".format(L))
if joint and union:
alphabet = set.union(*map(set, alphabets))
alphabet = list(sorted(alphabet))
alphabets = [alphabet] * len(alphabets)
if joint:
sample_space = CartesianProduct(alphabets, d._product)
else:
sample_space = ScalarSampleSpace(alphabets)
ed = d.__class__(d.outcomes,
d.pmf,
sample_space=sample_space,
base=d.get_base())
return ed
def pruned_samplespace(d, sample_space=None):
"""
Returns a new distribution with pruned sample space.
The pruning is such that zero probability outcomes are removed.
Parameters
----------
d : distribution
The distribution used to create the pruned distribution.
sample_space : set
A list of outcomes with zero probability that should be kept in the
sample space. If `None`, then all outcomes with zero probability
will be removed.
Returns
-------
pd : distribution
The distribution with a pruned sample space.
"""
if sample_space is None:
sample_space = []
keep = set(sample_space)
outcomes = []
pmf = []
for o, p in d.zipped(mode='atoms'):
if not d.ops.is_null_exact(p) or o in keep:
outcomes.append(o)
pmf.append(p)
if d.is_joint():
sample_space = SampleSpace(outcomes)
else:
sample_space = ScalarSampleSpace(outcomes)
pd = d.__class__(outcomes,
pmf,
sample_space=sample_space,
base=d.get_base())
return pd