Add flat argument to AMPL.get_iis

CHANGELOG.md

@@ -2,6 +2,9 @@
 All notable changes to this project will be documented in this file.
 This project adheres to [Semantic Versioning](http://semver.org/).
 
+## 0.14.1 - 2024-05-##
+- Add flat argument to AMPL.get_iis.
+
 ## 0.14.0 - 2024-05-21
 - Allow assigning values to indexed sets from a dictionary with the lists of members
   for every index.

amplpy/__init__.py

@@ -42,7 +42,7 @@
 except Exception:
     pass
 
-__version__ = "0.14.0"
+__version__ = "0.14.1b0"
 
 
 def _list_aliases():

amplpy/ampl.py

@@ -25,6 +25,22 @@
 inf = float("inf")
 
 
+def nested_dict_of_suffixes(lst):
+    nested = {}
+    for name, value in lst:
+        if "[" not in name:
+            nested[name] = value
+        else:
+            p = name.find("[")
+            v, index = name[:p], literal_eval(f"({name[p+1:-1]},)")
+            if v not in nested:
+                nested[v] = {}
+            if len(index) == 1:
+                index = index[0]
+            nested[v][index] = value
+    return nested
+
+
 AMPL_NOT_FOUND_MESSAGE = """
 Please make sure that the AMPL directory is in the system search path, or
 add it before instantiating the AMPL object with:
@@ -892,10 +908,13 @@ def solve_result_num(self):
         """
         return self.get_value("solve_result_num")
 
-    def get_iis(self):
+    def get_iis(self, flat=True):
         """
         Get IIS attributes for all variables and constraints.
 
+        Args:
+            flat: Return flat dictionaries if set to True, or nested dictionaries otherwise.
+
         Returns:
             Tuple with a dictionary for variables in the IIS and another for the constraints.
 
@@ -919,16 +938,18 @@ def get_iis(self):
                 var_iis, con_iis = ampl.get_iis()
                 print(var_iis, con_iis)
         """
-        iis_var = dict(
-            self.get_data(
-                "{i in 1.._nvars: _var[i].iis != 'non'} (_varname[i], _var[i].iis)"
-            ).to_list(skip_index=True)
-        )
-        iis_con = dict(
-            self.get_data(
-                "{i in 1.._ncons: _con[i].iis != 'non'} (_conname[i], _con[i].iis)"
-            ).to_list(skip_index=True)
-        )
+        iis_var = self.get_data(
+            "{i in 1.._nvars: _var[i].iis != 'non'} (_varname[i], _var[i].iis)"
+        ).to_list(skip_index=True)
+        iis_con = self.get_data(
+            "{i in 1.._ncons: _con[i].iis != 'non'} (_conname[i], _con[i].iis)"
+        ).to_list(skip_index=True)
+        if flat is False:
+            iis_var = nested_dict_of_suffixes(iis_var)
+            iis_con = nested_dict_of_suffixes(iis_con)
+        else:
+            iis_var = dict(iis_var)
+            iis_con = dict(iis_con)
         return iis_var, iis_con
 
     def get_solution(self, flat=True, zeros=False):
@@ -955,22 +976,10 @@ def get_solution(self, flat=True, zeros=False):
             stmt = "{i in 1.._nvars} (_varname[i], _var[i].val)"
         else:
             stmt = "{i in 1.._nvars: _var[i].val != 0} (_varname[i], _var[i].val)"
-        flat_solution = self.get_data(stmt).to_list(skip_index=True)
+        lst_solution = self.get_data(stmt).to_list(skip_index=True)
         if flat:
-            return dict(flat_solution)
-        solution = {}
-        for varname, value in flat_solution:
-            if "[" not in varname:
-                solution[varname] = value
-            else:
-                p = varname.find("[")
-                v, index = varname[:p], literal_eval(f"({varname[p+1:-1]},)")
-                if v not in solution:
-                    solution[v] = {}
-                if len(index) == 1:
-                    index = index[0]
-                solution[v][index] = value
-        return solution
+            return dict(lst_solution)
+        return nested_dict_of_suffixes(lst_solution)
 
     def _start_recording(self, filename):
         """

amplpy/tests/test_ampl.py

@@ -426,16 +426,19 @@ def test_get_iis(self):
         ampl.eval(
             r"""
             var x >= 0;
-            var y >= 0;
-            maximize obj: x+y;
-            s.t. s: x+y <= -5;
+            var y{1..2} >= 0;
+            maximize obj: x+y[1]+y[2];
+            s.t. s: x+y[1] <= -5;
             """
         )
         ampl.option["presolve"] = 0
         ampl.solve(solver="gurobi", gurobi_options="outlev=1 iis=1")
         self.assertEqual(ampl.solve_result, "infeasible")
         var_iis, con_iis = ampl.get_iis()
-        self.assertEqual(var_iis, {"x": "low", "y": "low"})
+        self.assertEqual(var_iis, {"x": "low", "y[1]": "low"})
+        self.assertEqual(con_iis, {"s": "mem"})
+        var_iis, con_iis = ampl.get_iis(flat=False)
+        self.assertEqual(var_iis, {"x": "low", "y": {1: "low"}})
         self.assertEqual(con_iis, {"s": "mem"})
 
     def test_get_solution(self):

setup.py

@@ -165,7 +165,7 @@ def link_args():
 
 setup(
     name="amplpy",
-    version="0.14.0",
+    version="0.14.1b0",
     description="Python API for AMPL",
     long_description=__doc__,
     long_description_content_type="text/markdown",