improve PyMPL CmdCon command

pyvpsolver/pympl/cmd/model.py

@@ -38,11 +38,12 @@ def _evalcmd(self, name, typ="", lb=None, ub=None):
 class CmdCon(CmdBase):
     """Command for creating a new AMPL constraint."""
 
-    def _evalcmd(self, name, lincomb, sign, rhs):
+    def _evalcmd(self, name, left, sign, right):
         """Evalutates CMD[name](*args)."""
         match = utils.parse_symbname(name)
         assert match is not None
         name = match
+        lincomb, sign, rhs = utils.linear_constraint(left, sign, right)
         self._pyvars["_model"] += ampl_con(name, lincomb, sign, rhs)
 
 

pyvpsolver/pympl/test/unittests.py

@@ -87,12 +87,22 @@ def test_con(self):
         parser = PyMPL()
         parser.input = """
         $VAR[x1]{}; $VAR[x2]{}; $VAR[x3]{};
-        $CON[xyz]{[("x1",5),("x2",15),("x3",10)],">=",20};
-        $CON[^xyz2]{[("x1",5),("x2",15),("x3",10)],">=",20};
+        $CON[con1]{[("x1",5),("x2",15),("x3",10)],">=",20};
+        $CON[con2]{[("x1",5)],">=",[("x2",-15),("x3",-10),20]};
+        $CON[con3]{-20,">=",[("x1",-5),("x2",-15),("x3",-10)]};
+        $CON[con4]{-20,">=",[(-5, "x1"),("x2",-15),(-10, "x3")]};
+        $CON[con5]{[-20, "x1"],">=",[(-4, "x1"),("x2",-15),(-10, "x3")]};
+        $CON[con6]{"x1",">=",[(-4, "x1"),20,("x2",-15),(-10, "x3")]};
+        $CON[^xyz]{[("x1",5),("x2",15),("x3",10)],">=",20};
         """
         parser.parse(comment_cmds=False)
-        self.assertIn("s.t. xyz: +5*x1+15*x2+10*x3 >= 20;", parser.output)
-        self.assertNotIn("s.t. xyz2: +5*x1+15*x2+10*x3 >= 20;", parser.output)
+        self.assertIn("s.t. con1: +5*x1+15*x2+10*x3 >= 20;", parser.output)
+        self.assertIn("s.t. con2: +5*x1+15*x2+10*x3 >= 20;", parser.output)
+        self.assertIn("s.t. con3: +5*x1+15*x2+10*x3 >= 20;", parser.output)
+        self.assertIn("s.t. con4: +5*x1+15*x2+10*x3 >= 20;", parser.output)
+        self.assertIn("s.t. con5: +5*x1+15*x2+10*x3 >= 20;", parser.output)
+        self.assertIn("s.t. con6: +5*x1+15*x2+10*x3 >= 20;", parser.output)
+        self.assertNotIn("s.t. xyz: +5*x1+15*x2+10*x3 >= 20;", parser.output)
 
     def test_stmt(self):
         """Tests $STMT{stmt} calls"""

pyvpsolver/pympl/utils.py

@@ -21,6 +21,7 @@
 
 import re
 from copy import deepcopy
+from collections import Iterable, defaultdict
 
 t_SYMBNAME = r'\^?[a-zA-Z_][a-zA-Z0-9_]*'
 t_INDEX1 = r'(?:\s*\[[^\]]*\])?' # [] index
@@ -183,6 +184,51 @@ def format_entry(k, v):
         return defs, ""
 
 
+def linear_constraint(left, sign, right):
+    """Transforms (left, sign, right) constraints into (lincomb, sign, rhs)."""
+    assert (
+        not isinstance(left, (int, float))
+        or not isinstance(right, (int, float))
+    )
+    pairs = defaultdict(lambda: 0)
+    rhs = 0
+
+    def add_entry(e, signal):
+        assert isinstance(e, (int, float, str, Iterable))
+        if isinstance(e, (int, float)):
+            return -signal*e
+        elif isinstance(e, str):
+            pairs[e] += signal
+            return 0
+        elif isinstance(e, Iterable):
+            a, b = e
+            assert isinstance(a, int) or isinstance(b, int)
+            assert isinstance(a, str) or isinstance(b, str)
+            if isinstance(a, str):
+                pairs[a] += signal*b
+            else:
+                pairs[b] += signal*a
+            return 0
+
+    if isinstance(left, (int, float, str)):
+        rhs += add_entry(left, 1)
+    else:
+        for e in left:
+            rhs += add_entry(e, 1)
+
+    if isinstance(right, (int, float, str)):
+        rhs += add_entry(right, -1)
+    else:
+        for e in right:
+            rhs += add_entry(e, -1)
+
+    if sign in ("<", ">"):
+        sign += "="
+
+    lincomb = sorted(pairs.items())
+    return (lincomb, sign, rhs)
+
+
 def lincomb2str(lincomb):
     """Returns the linear combination as a string."""