From bcf1306884908d30262c77f6902f7d5a7152c71a Mon Sep 17 00:00:00 2001
From: Gleb Belov <gleb@ampl.com>
Date: Wed, 20 Mar 2024 13:24:52 +1100
Subject: [PATCH] ModelExplorer: NL SOS cons #232

---
 support/modelexplore/scripts/python/model.py | 23 ++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/support/modelexplore/scripts/python/model.py b/support/modelexplore/scripts/python/model.py
index 62880ca4d..225fe42a3 100644
--- a/support/modelexplore/scripts/python/model.py
+++ b/support/modelexplore/scripts/python/model.py
@@ -24,7 +24,8 @@ def __init__(self):
       "Nonlinear" : [],
       "Linear" : [],
       "Logical": [],
-      "SOS": []}
+      "SOS1": [],
+      "SOS2": []}
     self._cons_Flat = {}
     self._cons_Flat_Group = {}
     self._objs_NL = []
@@ -50,9 +51,17 @@ def UpdateNLCon(self, type, idx, data):
     elif "lin"==type:
       self._updateNodeData(self._cons_NL["Linear"],
       len(self._cons_NL["Linear"]), data)
-    else:  # "logical"==type:
+    elif "logical"==type:
       self._updateNodeData(self._cons_NL["Logical"],
       len(self._cons_NL["Logical"]), data)
+    elif "_sos1"==type:
+      self._updateNodeData(self._cons_NL["SOS1"],
+      len(self._cons_NL["SOS1"]), data)
+    elif "_sos2"==type:
+      self._updateNodeData(self._cons_NL["SOS2"],
+      len(self._cons_NL["SOS2"]), data)
+    else:
+      raise Exception("Unknown NL constraint type: "+type)
 
   def UpdateFlatConGroup(self, type, data):
     self._cons_Flat_Group[type] = data
@@ -61,6 +70,10 @@ def UpdateFlatCon(self, type, idx, data):
     if type not in self._cons_Flat:
       self._cons_Flat[type] = []
     self._updateNodeData(self._cons_Flat[type], idx, data)
+    if 0==data["depth"] \
+        and type.startswith('_sos') \
+        and "printed" in data:   ## we need the final status
+      self.UpdateNLCon(type, 0, data)
 
   def _updateNodeData(self, specnodecnt, idx, data):
     data1, upd = self._updateItemData(specnodecnt, idx, data)
@@ -94,8 +107,10 @@ def MatchOrigModel(self, keyw):
       = self._matchRecords(self._cons_NL.get("Linear"), keyw)
     result["NL Logical Constraints"] \
       = self._matchRecords(self._cons_NL.get("Logical"), keyw)
-    result["NL SOS Constraints"] \
-      = self._matchRecords(self._cons_NL.get("SOS"), keyw)
+    result["NL SOS1 Constraints"] \
+      = self._matchRecords(self._cons_NL.get("SOS1"), keyw)
+    result["NL SOS2 Constraints"] \
+      = self._matchRecords(self._cons_NL.get("SOS2"), keyw)
     return result
 
   # Match keyword to the final model