restructured _filter_constant_condition in sympy, removed FIXMEs

LKlinke · Dec 11, 2024 · 3387c43 · 3387c43
1 parent 11c9fcc
commit 3387c43
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 48 deletions.
diff --git a/prodigy/distribution/generating_function.py b/prodigy/distribution/generating_function.py
@@ -175,51 +175,54 @@ def factory() -> Type[CommonDistributionsFactory]:
 
     def _filter_constant_condition(self,
                                    condition: Expr) -> GeneratingFunction:
+        # TODO refactor this in superclass?
         # Normalize the condition into the format _var_ (< | <= | =) const. I.e., having the variable on the lhs.
 
-        # FIXME variable comparison does not work as condition.rhs.var is str and self._var of type
-        #   set[sympy.Symbol]
-        if isinstance(condition.rhs,
-                      VarExpr) and condition.rhs.var not in self._variables:
-            if condition.operator == Binop.LEQ:
-                return self.filter(
-                    UnopExpr(operator=Unop.NEG,
-                             expr=BinopExpr(operator=Binop.LT,
-                                            lhs=condition.rhs,
-                                            rhs=condition.lhs)))
-            elif condition.operator == Binop.LT:
-                return self.filter(
-                    UnopExpr(operator=Unop.NEG,
-                             expr=BinopExpr(operator=Binop.LEQ,
-                                            lhs=condition.rhs,
-                                            rhs=condition.lhs)))
-            elif condition.operator == Binop.GT:
-                return self.filter(
-                    BinopExpr(operator=Binop.LT,
-                              lhs=condition.rhs,
-                              rhs=condition.lhs))
-            elif condition.operator == Binop.GEQ:
-                return self.filter(
-                    BinopExpr(operator=Binop.LEQ,
-                              lhs=condition.rhs,
-                              rhs=condition.lhs))
-
-        # FIXME variable comparison does not work as condition.rhs.var is str and self._var of type
-        #   set[sympy.Symbol]
-        if isinstance(condition.lhs,
-                      VarExpr) and condition.lhs.var not in self._variables:
-            if condition.operator == Binop.GT:
-                return self.filter(
-                    UnopExpr(operator=Unop.NEG,
-                             expr=BinopExpr(operator=Binop.LEQ,
-                                            lhs=condition.lhs,
-                                            rhs=condition.rhs)))
-
-            elif condition.operator == Binop.GEQ:
-                return self.filter(
-                    BinopExpr(operator=Binop.LEQ,
-                              lhs=condition.rhs,
-                              rhs=condition.lhs))
+        # Mapping of operator from expression
+        #   <const> <op> <var>
+        # to
+        #   <var> <op> <const>
+        flipped_operators = {
+            Binop.LT: Binop.GT,
+            Binop.LEQ: Binop.GEQ,
+            Binop.GT: Binop.LT,
+            Binop.GEQ: Binop.LEQ,
+            Binop.EQ: Binop.EQ
+        }
+
+        # Mapping of operator from
+        #   <var> <op> <const>
+        # to
+        #   ¬(<var> <op> <const>)
+        neg_operator = {
+            Binop.GT: Binop.LEQ,
+            Binop.GEQ: Binop.LT
+        }
+
+        # If condition of form
+        #   <const> <op> <var>
+        # flip operator and move to
+        #   <var> <op> <const>
+        if isinstance(condition.rhs, VarExpr):
+            return self.filter(
+                BinopExpr(
+                    operator=flipped_operators[condition.operator],
+                    rhs=condition.lhs,
+                    lhs=condition.rhs
+                )
+            )
+
+        # Now we have the form <var> <op> <const>
+        # As we only allow <op> \in {<, <=, =}, we need to remove the other operator
+        # by using negation
+        if condition.operator in neg_operator:
+            return self.filter(
+                UnopExpr(operator=Unop.NEG, expr=BinopExpr(
+                    operator=neg_operator[condition.operator],
+                    rhs=condition.rhs,
+                    lhs=condition.lhs
+                ))
+            )
 
         # Now we have an expression of the form _var_ (< | <=, =) _const_.
         variable = _sympy_symbol(str(condition.lhs))

diff --git a/prodigy/distribution/symengine_distribution.py b/prodigy/distribution/symengine_distribution.py
@@ -313,11 +313,6 @@ def _exhaustive_search(self, condition: Expr) -> SymengineDist:
         )
 
     def _filter_constant_condition(self, condition: Expr) -> SymengineDist:
-        # FIXME
-        #   calling
-        #       python prodigy/cli.py --engine symengine main pgfexamples/sequential_loops_second_inv.pgcl
-        #   leads to an error with the expression "m > 0", where m \in self._variables
-
         # Mapping of operator from expression
         #   <const> <op> <var>
         # to