NLWPY: suffix output #30

ampl · Feb 22, 2024 · 7385be7 · 7385be7
1 parent 8bdf2f5
commit 7385be7
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 12 deletions.
diff --git 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,16 +81,18 @@ 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;
   }
 
   /// 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
@@ -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
@@ -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