Follow-up division improvements

Author: cthoyt

src/y0/dsl.py

@@ -555,9 +555,14 @@ class Expression(Element, ABC):
     def __mul__(self, other):
         pass
 
-    def __truediv__(self, expression: Expression) -> Fraction:
+    def __truediv__(self, expression: Expression) -> Expression:
         """Divide this expression by another and create a fraction."""
-        return Fraction(self, expression)
+        if isinstance(expression, One):
+            return self
+        elif isinstance(expression, Fraction):
+            return Fraction(self * expression.denominator, expression.numerator)
+        else:
+            return Fraction(self, expression)
 
     def marginalize(self, ranges: VariableHint) -> Fraction:
         """Return this expression, marginalized by the given variables.
@@ -987,13 +992,13 @@ def __mul__(self, expression: Expression) -> Fraction:
             return Fraction(self.numerator * expression, self.denominator)
 
     def __truediv__(self, expression: Expression) -> Fraction:
-        if isinstance(expression, Fraction):
+        if isinstance(expression, One):
+            return self
+        elif isinstance(expression, Fraction):
             return Fraction(
                 self.numerator * expression.denominator,
                 self.denominator * expression.numerator,
             )
-        elif isinstance(expression, One):
-            return self
         else:
             return Fraction(self.numerator, self.denominator * expression)
 
@@ -1002,12 +1007,19 @@ def _iter_variables(self) -> Iterable[Variable]:
         yield from self.numerator._iter_variables()
         yield from self.denominator._iter_variables()
 
+    def flip(self) -> Fraction:
+        """Exchange the numerator and denominator."""
+        return Fraction(self.denominator, self.numerator)
+
     def simplify(self) -> Expression:
         """Simplify this fraction."""
         if isinstance(self.denominator, One):
             return self.numerator
         if isinstance(self.numerator, One):
-            return self
+            if isinstance(self.denominator, Fraction):
+                return self.denominator.flip().simplify()
+            else:
+                return self
         if self.numerator == self.denominator:
             return One()
         if isinstance(self.numerator, Product) and isinstance(self.denominator, Product):
@@ -1183,12 +1195,6 @@ def __mul__(self, other: Expression):
         else:
             return Product((self, other))
 
-    def __truediv__(self, expression: Expression) -> Fraction:
-        if isinstance(expression, Fraction):
-            return Fraction(self * expression.denominator, expression.numerator)
-        else:
-            return super().__truediv__(expression)
-
     def _iter_variables(self) -> Iterable[Variable]:
         yield from self.codomain
         yield from self.domain

src/y0/mutate/chain.py

@@ -87,7 +87,7 @@ def fraction_expand(p: Probability) -> Fraction:
     .. math::
         P(Y_1,\dots,Y_n | X_1, \dots, X_m) = \frac{P(Y_1,\dots,Y_n,X_1,\dots,X_m)}{P(X_1,\dots,X_m)}
     """
-    return p.uncondition() / P(p.parents)
+    return Fraction(p.uncondition(), P(p.parents))
 
 
 def bayes_expand(p: Probability) -> Fraction:

tests/test_mutate/test_simplify.py

@@ -4,7 +4,7 @@
 
 import unittest
 
-from y0.dsl import A, B, One, P, Sum
+from y0.dsl import A, B, Fraction, One, P, Sum
 
 one = One()
 
@@ -15,12 +15,13 @@ class TestCancel(unittest.TestCase):
     def test_simple_identity(self):
         """Test cancelling when the numerator and denominator are the same."""
         for label, frac in [
-            ("one", one / one),
+            ("one", Fraction(one, one)),
             ("prob", P(A) / P(A)),
             ("sum", Sum(P(A)) / Sum(P(A))),
             ("product", (P(A) * P(B)) / (P(A) * P(B))),
         ]:
             with self.subTest(type=label):
+                self.assertIsInstance(frac, Fraction)
                 self.assertEqual(one, frac.simplify(), msg=f"\n\nActual:{frac}")
 
     def test_fraction_simplify(self):
@@ -29,6 +30,13 @@ def test_fraction_simplify(self):
             ("leave num.", P(B), (P(A) * P(B)) / P(A)),
             ("leave den.", one / P(B), P(A) / (P(A) * P(B))),
             ("unordered", one, (P(A) * P(B)) / (P(B) * P(A))),
+            ("canonical", one / P(A), one / P(A)),
+            ("flipper", P(A), Fraction(one, Fraction(one, P(A)))),
+            ("prob-redundant-one", P(A), Fraction(P(A), one)),
+            ("sum-redundant-one", Sum(P(A)), Fraction(Sum(P(A)), one)),
+            ("frac-redundant-one", P(A) / P(B), Fraction(Fraction(P(A), P(B)), one)),
+            ("prod-redundant-one", P(A) * P(B), Fraction(P(A) * P(B), one)),
         ]:
             with self.subTest(type=label):
+                self.assertIsInstance(frac, Fraction)
                 self.assertEqual(expected, frac.simplify(), msg=f"\n\nActual:{frac}")