From 7385be7bc92fcf712db9612b15ba5476468696e6 Mon Sep 17 00:00:00 2001
From: Gleb Belov <gleb@ampl.com>
Date: Thu, 22 Feb 2024 17:27:01 +1100
Subject: [PATCH] NLWPY: suffix output #30

---
 nl-writer2/include/mp/nl-model.h     | 10 ++++++----
 nl-writer2/nlwpy/src/nlw_bindings.cc | 19 +++++++++++--------
 nl-writer2/nlwpy/tests/test.py       | 18 ++++++++++++++++++
 3 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/nl-writer2/include/mp/nl-model.h b/nl-writer2/include/mp/nl-model.h
index 9bcc89384..03db817e4 100644
--- a/nl-writer2/include/mp/nl-model.h
+++ b/nl-writer2/include/mp/nl-model.h
@@ -65,8 +65,8 @@ struct NLSuffix {
 
   /// operator<
   bool operator<(const NLSuffix& s) const {
-    return std::make_pair(name_, kind_)
-        < std::make_pair(s.name_, s.kind_);
+    return std::make_pair(name_, kind_&3)
+        < std::make_pair(s.name_, s.kind_&3);
   }
 };
 
@@ -81,9 +81,10 @@ class NLSuffixSet : private std::set<NLSuffix> {
   { return this->insert(suf).second; }
 
   /// Find suffix.
-  /// @return NLSuffix*, nullptr if not found.
+  /// @param k: kind (only the first 2 bits are used.)
+  /// @return NLSuffix*, nullptr iff not found.
   const NLSuffix* Find(const std::string& nm, int k) const {
-    NLSuffix tmp {nm, {}, k};
+    NLSuffix tmp {nm, {}, k&3};
     auto it = this->find(tmp);
     return (this->end()!=it) ? &*it : nullptr;
   }
@@ -91,6 +92,7 @@ class NLSuffixSet : private std::set<NLSuffix> {
   /// Expose size, empty
   using Base::size;
   using Base::empty;
+  using Base::clear;
 
   /// Expose begin, end
   using Base::begin;
diff --git a/nl-writer2/nlwpy/src/nlw_bindings.cc b/nl-writer2/nlwpy/src/nlw_bindings.cc
index 885fbc4c8..9ccd99479 100644
--- a/nl-writer2/nlwpy/src/nlw_bindings.cc
+++ b/nl-writer2/nlwpy/src/nlw_bindings.cc
@@ -317,19 +317,22 @@ NLW2_NLSolver
       .def_readwrite("values_", &mp::NLSuffix::values_);
 
   /// NLSuffixSet
-  /*py::class_<mp::NLSuffixSet>(m, "NLW2_NLSuffixSet")
+  py::class_<mp::NLSuffixSet>(m, "NLW2_NLSuffixSet")
       .def("Find",           // Find(): return None if not found
            [=](mp::NLSuffixSet const& ss,
            std::string const& name, int kind) -> py::object {
     auto pelem = ss.Find(name, kind);
-    if (pelem) {
-      return py::cast(*pelem);
-    }
-    return py::object(py::cast(nullptr));
+    return py::cast(pelem);
   })
-      .def("begin", &mp::NLSuffixSet::begin)
-      .def("end", &mp::NLSuffixSet::end)
-      ;*/
+      .def("__len__", // &mp::NLSuffixSet::size - does not work)
+           [](const mp::NLSuffixSet &ss) { return ss.size(); })
+      .def("__iter__", [](const mp::NLSuffixSet &ss) {
+         return py::make_iterator(ss.begin(), ss.end());
+      }, py::keep_alive<0, 1>()) /* Keep vector alive while iterator is used */
+      .def("empty",
+           [](const mp::NLSuffixSet &ss) { return ss.empty(); })
+      .def("clear", [](mp::NLSuffixSet &ss) { ss.clear(); })
+      ;
 
   /// NLModel
   py::class_<NLWPY_NLModel>(m, "NLW2_NLModel")
diff --git a/nl-writer2/nlwpy/tests/test.py b/nl-writer2/nlwpy/tests/test.py
index 0e3a93fd3..babce431f 100644
--- a/nl-writer2/nlwpy/tests/test.py
+++ b/nl-writer2/nlwpy/tests/test.py
@@ -83,6 +83,24 @@ def Check(self, sol):
                i+1, sol.x_[i], self.x_ref_[i]))
         result = False
 
+    ### Printing suffixes.
+    ### TODO replace by checking debug suffixes and ini guesses.
+    suffixes = sol.suffixes_
+    print("Number of suffixes returned:", len(suffixes))
+    for suf in suffixes:
+      print("    SUFFIX '{}' [{}]".format(suf.name_, suf.kind_))
+      print("        Table:    ", suf.table_)
+      print("        Values:   ", *suf.values_)
+    sufMIPGapObj = suffixes.Find("relmipgap", 2) ## should be 2+4
+    if sufMIPGapObj is not None:
+      print("FOUND:    SUFFIX '{}' [{}].".format(
+        sufMIPGapObj.name_, sufMIPGapObj.kind_))
+    sufVarStatus = suffixes.Find("status", 0)
+    if sufVarStatus is not None:
+      print("FOUND:    SUFFIX '{}' [{}].".format(
+        sufVarStatus.name_, sufVarStatus.kind_))
+    suffixes.clear()
+
     print("MIQP 1: solution check {}, obj={:.17f}.".format(
       "ok" if result else "Failed", sol.obj_val_))
     return result