Permit dist mixture syntax using or for nominal-continuous [fix #83].

src/distributions.py

@@ -4,7 +4,18 @@
 import scipy.stats
 import sympy
 
-class NominalDistribution():
+class Distribution():
+    def __rmul__(self, x):
+        from .sym_util import sympify_number
+        try:
+            x_val = sympify_number(x)
+            if not 0 < x_val < 1:
+                raise ValueError('invalid weight %s' % (str(x),))
+            return DistributionMix([self], [x])
+        except TypeError:
+            return NotImplemented
+
+class NominalDistribution(Distribution):
     def __init__(self, dist):
         self.dist = dict(dist)
     def __call__(self, symbol):
@@ -13,9 +24,9 @@ def __call__(self, symbol):
 
 choice = NominalDistribution
 
-# pylint: disable=not-callable
-# pylint: disable=multiple-statements
-class RealDistribution():
+class RealDistribution(Distribution):
+    # pylint: disable=not-callable
+    # pylint: disable=multiple-statements
     dist = None
     constructor = None
     def __init__(self, *args, **kwargs):
@@ -27,17 +38,7 @@ def __call__(self, symbol):
     def get_domain(self, **kwargs):
         raise NotImplementedError()
 
-    def __rmul__(self, x):
-        from .sym_util import sympify_number
-        try:
-            x_val = sympify_number(x)
-            if not 0 < x_val < 1:
-                raise ValueError('invalid weight %s' % (str(x),))
-            return RealDistributionMix([self], [x])
-        except TypeError:
-            return NotImplemented
-
-class RealDistributionMix():
+class DistributionMix():
     """Weighted mixture of SPNs that do not yet sum to unity."""
     def __init__(self, distributions, weights):
         self.distributions = distributions
@@ -50,13 +51,13 @@ def __call__(self, symbol):
         return SumSPN(distributions, weights)
 
     def __or__(self, x):
-        if not isinstance(x, RealDistributionMix):
+        if not isinstance(x, DistributionMix):
             return NotImplemented
         weights = self.weights + x.weights
         cumsum = float(sum(weights))
         assert 0 < cumsum <= 1
         distributions = self.distributions + x.distributions
-        return RealDistributionMix(distributions, weights)
+        return DistributionMix(distributions, weights)
 
 # ==============================================================================
 # ContinuousReal

tests/test_parse_distributions.py

@@ -5,20 +5,22 @@
 
 import pytest
 
-from spn.distributions import RealDistributionMix
+from spn.distributions import DistributionMix
 from spn.distributions import bernoulli
+from spn.distributions import choice
 from spn.distributions import norm
 from spn.distributions import poisson
 from spn.math_util import allclose
 from spn.spn import ContinuousLeaf
 from spn.spn import DiscreteLeaf
+from spn.spn import NominalLeaf
 from spn.spn import SumSPN
 from spn.transforms import Id
 
 X = Id('X')
 
-def test_simple_parse():
-    assert isinstance(.3*bernoulli(p=.1), RealDistributionMix)
+def test_simple_parse_real():
+    assert isinstance(.3*bernoulli(p=.1), DistributionMix)
     a = .3*bernoulli(p=.1) | .5 * norm() | .2*poisson(mu=7)
     spn = a(X)
     assert isinstance(spn, SumSPN)
@@ -27,6 +29,15 @@ def test_simple_parse():
     assert isinstance(spn.children[1], ContinuousLeaf)
     assert isinstance(spn.children[2], DiscreteLeaf)
 
+def test_simple_parse_nominal():
+    assert isinstance(.7 * choice({'a': .1, 'b': .9}), DistributionMix)
+    a = .3*bernoulli(p=.1) | .7*choice({'a': .1, 'b': .9})
+    spn = a(X)
+    assert isinstance(spn, SumSPN)
+    assert allclose(spn.weights, [log(.3), log(.7)])
+    assert isinstance(spn.children[0], DiscreteLeaf)
+    assert isinstance(spn.children[1], NominalLeaf)
+
 def test_error():
     with pytest.raises(TypeError):
         'a'*bernoulli(p=.1)