Analytic methods

squigglepy/distributions.py

@@ -12,6 +12,8 @@
 from collections.abc import Iterable
 
 from abc import ABC, abstractmethod
+from functools import reduce
+from numbers import Real
 
 
 class BaseDistribution(ABC):
@@ -172,6 +174,43 @@ def __rpow__(self, dist):
     def __hash__(self):
         return hash(repr(self))
 
+    def simplify(self):
+        """Simplify a distribution by evaluating all operations that can be
+        performed analytically, for example by reducing a sum of normal
+        distributions into a single normal distribution. Return a new
+        ``OperableDistribution`` that represents the simplified distribution.
+
+        Possible simplifications:
+            any - any = any + (-any)
+            any / constant = any * (1 / constant)
+
+            -Normal = Normal
+
+            Normal + Normal = Normal
+            constant + Normal = Normal
+            Bernoulli + Bernoulli = Binomial
+            Bernoulli + Binomial = Binomial
+            Binomial + Binomial = Binomial
+            Chi-Square + Chi-Square = Chi-Square
+            Exponential + Exponential = Exponential
+            Exponential + Gamma = Gamma
+            Gamma + Gamma = Gamma
+            Poisson + Poisson = Poisson
+
+            constant * Normal = Normal
+            Lognormal * Lognormal = Lognormal
+            constant * Lognormal = Lognormal
+            constant * Exponential = Exponential
+            constant * Gamma = Gamma
+
+            Lognormal / Lognormal = Lognormal
+            constant / Lognormal = Lognormal
+
+            Lognormal ** constant = Lognormal
+
+        """
+        return self
+
 
 # Distribution are either discrete, continuous, or composite
 
@@ -242,6 +281,9 @@ def __str__(self):
             raise ValueError
         return out
 
+    def simplify(self):
+        return FlatTree.build(self).simplify()
+
 
 def _get_fname(f, name):
     if name is None:
@@ -1699,3 +1741,321 @@ def geometric(p):
     <Distribution> geometric(0.1)
     """
     return GeometricDistribution(p=p)
+
+
+class FlatTree:
+    """Helper class for simplifying analytic expressions. A ``FlatTree`` is
+    sort of like a ``ComplexDistribution`` except that it flattens
+    commutative/associative operations onto a single object instead of having
+    one object per binary operation.
+
+    This class operates in two phases.
+
+    Phase 1: Generate a ``FlatTree`` object from a :ref:``BaseDistribution`` by
+    calling ``FlatTree.build(dist)``. This generates a tree where any series of
+    a single commutative/associative operation done repeatedly is flattened
+    onto a single ``FlatTree`` node. It also converts operations into a
+    normalized form, for example converting ``a - b`` into ``a + (-b)``.
+
+    Phase 2: Generate a simplified ``Distribution`` by calling
+    :ref:``simplify``. This works by combing through each flat list of
+    distributions to find which ones can be analytically simplified (for
+    example, converting a sum of normal distributions into a single normal
+    distribution).
+
+    """
+
+    COMMUTABLE_OPERATIONS = set([operator.add, operator.mul])
+
+    def __init__(self, dist=None, fn=None, fn_str=None, children=None, is_unary=False, infix=None):
+        self.dist = dist
+        self.fn = fn
+        self.fn_str = fn_str
+        self.children = children
+        self.is_unary = is_unary
+        self.infix = infix
+        if dist is not None:
+            self.is_leaf = True
+        elif fn is not None and children is not None:
+            self.is_leaf = False
+        else:
+            raise ValueError("Missing arguments to FlatTree constructor")
+
+    def __str__(self):
+        if self.is_leaf:
+            return f"FlatTree({self.dist})"
+        else:
+            return "FlatTree({})[{}]".format(self.fn_str, ", ".join(map(str, self.children)))
+
+    def __repr__(self):
+        return str(self)
+
+    @classmethod
+    def build(cls, dist):
+        if dist is None:
+            return None
+        if isinstance(dist, Real):
+            return cls(dist=dist)
+        if not isinstance(dist, BaseDistribution):
+            raise ValueError(f"dist must be a BaseDistribution or numeric type, not {type(dist)}")
+        if not isinstance(dist, ComplexDistribution):
+            return cls(dist=dist)
+
+        is_unary = dist.right is None
+        if is_unary and dist.right is not None:
+            raise ValueError(f"Multiple arguments provided for unary operator {dist.fn}")
+
+        # Convert x - y into x + (-y)
+        if dist.fn == operator.sub:
+            return cls.build(
+                ComplexDistribution(
+                    dist.left,
+                    ComplexDistribution(dist.right, right=None, fn=operator.neg, fn_str="-"),
+                    fn=operator.add,
+                    fn_str="+",
+                )
+            )
+
+        # If the denominator is a constant, replace division by constant
+        # with multiplication by the reciprocal of the constant
+        if dist.fn == operator.truediv and isinstance(dist.right, Real):
+            if dist.right == 0:
+                raise ZeroDivisionError("Division by zero in ComplexDistribution: {dist}")
+            return cls.build(
+                ComplexDistribution(
+                    dist.left,
+                    1 / dist.right,
+                    fn=operator.mul,
+                    fn_str="*",
+                )
+            )
+        if dist.fn == operator.truediv and isinstance(dist.right, LognormalDistribution):
+            return cls.build(
+                ComplexDistribution(
+                    dist.left,
+                    LognormalDistribution(
+                        norm_mean=-dist.right.norm_mean, norm_sd=dist.right.norm_sd
+                    ),
+                    fn=operator.mul,
+                    fn_str="*",
+                )
+            )
+
+        left_tree = cls.build(dist.left)
+        right_tree = cls.build(dist.right)
+
+        # Make a list of possibly-joinable distributions, plus a list of
+        # children as trees who could not be simplified at this level
+        children = []
+
+        # If the child nodes use the same commutable operation as ``dist``, add
+        # their flattened ``children`` lists to ``children``. Otherwise, put
+        # the whole node in ``children``.
+        if left_tree.is_leaf:
+            children.append(left_tree.dist)
+        elif left_tree.fn == dist.fn and dist.fn in cls.COMMUTABLE_OPERATIONS:
+            children.extend(left_tree.children)
+        else:
+            children.append(left_tree)
+        if right_tree is not None:
+            if right_tree.is_leaf:
+                children.append(right_tree.dist)
+            elif right_tree.fn == dist.fn and dist.fn in cls.COMMUTABLE_OPERATIONS:
+                children.extend(right_tree.children)
+            else:
+                children.append(right_tree)
+
+        return cls(
+            fn=dist.fn, fn_str=dist.fn_str, children=children, is_unary=is_unary, infix=dist.infix
+        )
+
+    def _join_dists(self, left_type, right_type, join_fn, commutative=True, condition=None):
+        simplified_dists = []
+        acc = None
+        acc_index = None
+        acc_is_left = True
+        for i, x in enumerate(self.children):
+            if isinstance(x, BaseDistribution) and (x.lclip is not None or x.rclip is not None):
+                # We can't simplify a clipped distribution
+                simplified_dists.append(x)
+            elif acc is None and isinstance(x, left_type):
+                acc = x
+                acc_index = i
+            elif (
+                acc is not None
+                and isinstance(x, right_type)
+                and acc_is_left
+                and (condition is None or condition(acc, x))
+            ):
+                acc = join_fn(acc, x)
+            elif commutative and acc is None and isinstance(x, right_type):
+                acc = x
+                acc_index = i
+                acc_is_left = False
+            elif (
+                commutative
+                and acc is not None
+                and isinstance(x, left_type)
+                and not acc_is_left
+                and (condition is None or condition(x, acc))
+            ):
+                acc = join_fn(x, acc)
+            else:
+                simplified_dists.append(x)
+
+        if acc is not None:
+            simplified_dists.insert(acc_index, acc)
+        self.children = simplified_dists
+
+    @classmethod
+    def _lognormal_times_const(cls, norm_mean, norm_sd, k):
+        if k == 0:
+            return 0
+        elif k > 0:
+            return LognormalDistribution(norm_mean=norm_mean + np.log(k), norm_sd=norm_sd)
+        else:
+            return -LognormalDistribution(norm_mean=norm_mean + np.log(-k), norm_sd=norm_sd)
+
+    def simplify(self):
+        """Convert a FlatTree back into a Distribution, simplifying as much as
+        possible."""
+        if self.is_leaf:
+            return self.dist
+
+        for i in range(len(self.children)):
+            if isinstance(self.children[i], FlatTree):
+                self.children[i] = self.children[i].simplify()
+
+        # Simplify unary operations
+        if len(self.children) == 1:
+            child = self.children[0]
+            if self.fn == operator.neg:
+                if isinstance(child, Real):
+                    return -child
+                if isinstance(child, NormalDistribution):
+                    return NormalDistribution(mean=-child.mean, sd=child.sd)
+
+            return ComplexDistribution(
+                child, right=None, fn=self.fn, fn_str=self.fn_str, infix=self.infix
+            )
+
+        if self.fn == operator.add:
+            self._join_dists(
+                NormalDistribution,
+                NormalDistribution,
+                lambda x, y: NormalDistribution(
+                    mean=x.mean + y.mean, sd=np.sqrt(x.sd**2 + y.sd**2)
+                ),
+            )
+            self._join_dists(
+                NormalDistribution, Real, lambda x, y: NormalDistribution(mean=x.mean + y, sd=x.sd)
+            )
+            self._join_dists(
+                BernoulliDistribution,
+                BernoulliDistribution,
+                lambda x, y: BinomialDistribution(n=2, p=x.p),
+                condition=lambda x, y: x.p == y.p,
+            )
+            self._join_dists(
+                BinomialDistribution,
+                BernoulliDistribution,
+                lambda x, y: BinomialDistribution(n=x.n + 1, p=x.p),
+                condition=lambda x, y: x.p == y.p,
+            )
+            self._join_dists(
+                BinomialDistribution,
+                BinomialDistribution,
+                lambda x, y: BinomialDistribution(n=x.n + y.n, p=x.p),
+                condition=lambda x, y: x.p == y.p,
+            )
+            self._join_dists(
+                ChiSquareDistribution,
+                ChiSquareDistribution,
+                lambda x, y: ChiSquareDistribution(df=x.df + y.df),
+            )
+            self._join_dists(
+                ExponentialDistribution,
+                ExponentialDistribution,
+                lambda x, y: GammaDistribution(shape=2, scale=x.scale),
+                condition=lambda x, y: x.scale == y.scale,
+            )
+            self._join_dists(
+                ExponentialDistribution,
+                GammaDistribution,
+                lambda x, y: GammaDistribution(shape=y.shape + 1, scale=x.scale),
+                condition=lambda x, y: x.scale == y.scale,
+            )
+            self._join_dists(
+                GammaDistribution,
+                GammaDistribution,
+                lambda x, y: GammaDistribution(shape=x.shape + y.shape, scale=x.scale),
+                condition=lambda x, y: x.scale == y.scale,
+            )
+            self._join_dists(
+                PoissonDistribution,
+                PoissonDistribution,
+                lambda x, y: PoissonDistribution(lam=x.lam + y.lam),
+            )
+
+        elif self.fn == operator.mul:
+            self._join_dists(
+                NormalDistribution,
+                Real,
+                lambda x, y: NormalDistribution(mean=x.mean * y, sd=x.sd * y),
+            )
+            self._join_dists(
+                LognormalDistribution,
+                LognormalDistribution,
+                lambda x, y: LognormalDistribution(
+                    norm_mean=x.norm_mean + y.norm_mean,
+                    norm_sd=np.sqrt(x.norm_sd**2 + y.norm_sd**2),
+                ),
+            )
+            self._join_dists(
+                LognormalDistribution,
+                Real,
+                lambda x, y: self._lognormal_times_const(x.norm_mean, x.norm_sd, y),
+            )
+            self._join_dists(
+                ExponentialDistribution,
+                Real,
+                lambda x, y: ExponentialDistribution(scale=x.scale * y),
+            )
+            self._join_dists(
+                GammaDistribution,
+                Real,
+                lambda x, y: GammaDistribution(shape=x.shape, scale=x.scale * y),
+            )
+
+        elif self.fn == operator.truediv:
+            self._join_dists(
+                LognormalDistribution,
+                LognormalDistribution,
+                lambda x, y: LognormalDistribution(
+                    norm_mean=x.norm_mean - y.norm_mean,
+                    norm_sd=np.sqrt(x.norm_sd**2 + y.norm_sd**2),
+                ),
+                commutative=False,
+            )
+            self._join_dists(
+                Real,
+                LognormalDistribution,
+                lambda x, y: self._lognormal_times_const(-y.norm_mean, y.norm_sd, x),
+                commutative=False,
+            )
+
+        elif self.fn == operator.pow:
+            self._join_dists(
+                LognormalDistribution,
+                Real,
+                lambda x, y: LognormalDistribution(
+                    norm_mean=x.norm_mean * y, norm_sd=x.norm_sd * y
+                ),
+                commutative=False,
+                condition=lambda x, y: y > 0,
+            )
+
+        return reduce(
+            lambda acc, x: ComplexDistribution(acc, x, fn=self.fn, fn_str=self.fn_str),
+            self.children,
+        )