Update marginalize and conditional functions in DSL (#100)

Co-Authored-By: Jeremy Zucker <[email protected]>

Author: cthoyt

README.md

@@ -60,7 +60,8 @@ It can also be used to manipulate expressions:
 ```python
 from y0.dsl import P, A, B, Sum
 
-P(A, B).marginalize(A) == P(A, B) / Sum[A](P(A, B))
+P(A, B).marginalize(A) == Sum[A](P(A, B))
+P(A, B).conditional(A) == P(A, B) / Sum[A](P(A, B))
 ```
 
 DSL objects can be converted into strings with `str()` and parsed back

src/y0/algorithm/identify/id_c.py

@@ -24,7 +24,7 @@ def idc(identification: Identification) -> Expression:
             return idc(identification.exchange_observation_with_action(condition))
 
     # Run ID algorithm
-    return identify(identification.uncondition()).marginalize(identification.outcomes)
+    return identify(identification.uncondition()).conditional(identification.outcomes)
 
 
 def rule_2_of_do_calculus_applies(identification: Identification, condition: Variable) -> bool:

src/y0/dsl.py

@@ -585,16 +585,29 @@ def __truediv__(self, expression: Expression) -> Expression:
         else:
             return Fraction(self, expression)
 
-    def marginalize(self, ranges: VariableHint) -> Fraction:
-        """Return this expression, marginalized by the given variables.
+    def conditional(self, ranges: VariableHint) -> Fraction:
+        """Return this expression, conditioned by the given variables.
 
         :param ranges: A variable or list of variables over which to marginalize this expression
         :returns: A fraction in which the denominator is represents the sum over the given ranges
 
         >>> from y0.dsl import P, A, B
-        >>> assert P(A, B).marginalize(A) == P(A, B) / Sum[A](P(A, B))
+        >>> assert P(A, B).conditional(A) == P(A, B) / Sum[A](P(A, B))
+        >>> assert P(A, B, C).conditional([A, B]) == P(A, B, C) / Sum[A, B](P(A, B, C))
         """
-        return Fraction(self, Sum(expression=self, ranges=_upgrade_variables(ranges)))
+        return Fraction(self, self.marginalize(ranges))
+
+    def marginalize(self, ranges: VariableHint) -> Sum:
+        """Return this expression, marginalizing out the given variables.
+
+        :param ranges: A variable or list of variables over which to marginalize this expression
+        :returns: The expression but summed over the given variables
+
+        >>> from y0.dsl import P, A, B, C
+        >>> assert P(A, B).marginalize(A) == Sum[A](P(A, B))
+        >>> assert P(A, B, C).marginalize([A, B]) == Sum[A, B](P(A, B, C))
+        """
+        return Sum(expression=self, ranges=_upgrade_variables(ranges))
 
 
 @dataclass(frozen=True, repr=False)

src/y0/mutate/chain.py

@@ -102,4 +102,4 @@ def bayes_expand(p: Probability) -> Fraction:
 
     .. note:: This expansion will create a different but equal expression to :func:`fraction_expand`.
     """
-    return p.uncondition().marginalize(p.children)
+    return p.uncondition().conditional(p.children)