Add zero element to DSL (#96)

* Add zero element to DSL

Closes #95

* Add operation overrides

* Cleanup division

* Finalize zero implementation

* Fix zero divided by zero

* Update chain.py

Author: cthoyt

src/y0/dsl.py

@@ -36,6 +36,7 @@
     "Fraction",
     "Expression",
     "One",
+    "Zero",
     "Q",
     "QFactor",
     "A",
@@ -636,7 +637,11 @@ def is_markov_kernel(self) -> bool:
         return self.distribution.is_markov_kernel()
 
     def __mul__(self, other: Expression) -> Expression:
-        if isinstance(other, Product):
+        if isinstance(other, Zero):
+            return other
+        elif isinstance(other, One):
+            return self
+        elif isinstance(other, Product):
             return Product((self, *other.expressions))
         elif isinstance(other, Fraction):
             return Fraction(self * other.numerator, other.denominator)
@@ -841,6 +846,8 @@ def to_latex(self):
         return " ".join(expression.to_latex() for expression in self.expressions)
 
     def __mul__(self, other: Expression):
+        if isinstance(other, Zero):
+            return other
         if isinstance(other, Product):
             return Product((*self.expressions, *other.expressions))
         elif isinstance(other, Fraction):
@@ -929,7 +936,9 @@ def to_y0(self):
         return f"Sum[{ranges}]({s})"
 
     def __mul__(self, expression: Expression):
-        if isinstance(expression, Product):
+        if isinstance(expression, Zero):
+            return expression
+        elif isinstance(expression, Product):
             return Product((self, *expression.expressions))
         else:
             return Product((self, expression))
@@ -969,6 +978,10 @@ class Fraction(Expression):
     #: The expression in the denominator of the fraction
     denominator: Expression
 
+    def __post_init__(self):
+        if isinstance(self.denominator, Zero):
+            raise ZeroDivisionError
+
     def to_text(self) -> str:
         """Output this fraction in the internal string format."""
         return f"frac_{{{self.numerator.to_text()}}}{{{self.denominator.to_text()}}}"
@@ -982,8 +995,10 @@ def to_y0(self, parens: bool = True) -> str:
         s = f"({self.numerator.to_y0()} / {self.denominator.to_y0()})"
         return f"({s})" if parens else s
 
-    def __mul__(self, expression: Expression) -> Fraction:
-        if isinstance(expression, Fraction):
+    def __mul__(self, expression: Expression) -> Expression:
+        if isinstance(expression, Zero):
+            return expression
+        elif isinstance(expression, Fraction):
             return Fraction(
                 self.numerator * expression.numerator,
                 self.denominator * expression.denominator,
@@ -1015,6 +1030,8 @@ def simplify(self) -> Expression:
         """Simplify this fraction."""
         if isinstance(self.denominator, One):
             return self.numerator
+        if isinstance(self.numerator, Zero):
+            return self.numerator
         if isinstance(self.numerator, One):
             if isinstance(self.denominator, Fraction):
                 return self.denominator.flip().simplify()
@@ -1111,6 +1128,40 @@ def _iter_variables(self) -> Iterable[Variable]:
         return iter([])
 
 
+class Zero(Expression):
+    """The additive identity (0)."""
+
+    def to_text(self) -> str:
+        """Output this identity variable in the internal string format."""
+        return "0"
+
+    def to_latex(self) -> str:
+        """Output this identity instance in the LaTeX string format."""
+        return "0"
+
+    def to_y0(self) -> str:
+        """Output this identity instance as y0 internal DSL code."""
+        return "Zero()"
+
+    def __rmul__(self, expression: Expression) -> Expression:
+        return self
+
+    def __mul__(self, expression: Expression) -> Expression:
+        return self
+
+    def __truediv__(self, other: Expression) -> Expression:
+        if isinstance(other, Zero):
+            raise ZeroDivisionError
+        return self
+
+    def __eq__(self, other):
+        return isinstance(other, Zero)  # all zeros are equal
+
+    def _iter_variables(self) -> Iterable[Variable]:
+        """Get the set of variables used in this expression."""
+        return iter([])
+
+
 class QBuilder(Protocol[T_co]):
     """A protocol for annotating the special class getitem functionality of the :class:`QFactor` class."""
 

tests/test_dsl.py

@@ -28,6 +28,7 @@
     X,
     Y,
     Z,
+    Zero,
 )
 from y0.parser import parse_y0
 
@@ -345,3 +346,41 @@ def test_product(self):
         self.assertEqual(Product((p,)), Product.safe({p}))
 
         self.assertEqual(Product((P(X), P(Y))), Product.safe(P(v) for v in [X, Y]))
+
+
+zero = Zero()
+
+
+class TestZero(unittest.TestCase):
+    """Tests for zero."""
+
+    exprs = [
+        One(),
+        Zero(),
+        P(A),
+        P(A) * P(B),
+        P(A) / P(B),
+        Sum.safe(P(A), [A]),
+    ]
+
+    def test_divide_failure(self):
+        """Test failure is thron on division by zero."""
+        for expr in self.exprs:
+            with self.subTest(expr=expr), self.assertRaises(ZeroDivisionError):
+                expr / zero
+
+    def test_identity(self):
+        """Test that zero divided by anything is zero."""
+        for expr in self.exprs:
+            if isinstance(expr, Zero):
+                continue  # would raise divides by zero
+            with self.subTest(expr=expr):
+                self.assertEqual(zero, zero / expr)
+
+    def test_multiply(self):
+        """Test other operations."""
+        for expr in self.exprs:
+            with self.subTest(expr=expr.to_y0(), direction="right"):
+                self.assertEqual(zero, zero * expr, msg=f"Got {zero * expr}")
+            with self.subTest(expr=expr.to_y0(), direction="left"):
+                self.assertEqual(zero, expr * zero, msg=f"Got {expr * zero}")