From 915584406879780f52c42f3bad1f50700802e1cb Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Sun, 24 Dec 2023 23:59:10 -0800
Subject: [PATCH 01/12] initial commit

---
 maraboupy/Marabou.py                   |  16 ++
 maraboupy/MarabouNetworkComposition.py | 284 +++++++++++++++++++++++++
 maraboupy/MarabouNetworkONNX.py        |  20 +-
 3 files changed, 317 insertions(+), 3 deletions(-)
 create mode 100644 maraboupy/MarabouNetworkComposition.py

diff --git a/maraboupy/Marabou.py b/maraboupy/Marabou.py
index 2ba07abe7f..cbf4c59d74 100644
--- a/maraboupy/Marabou.py
+++ b/maraboupy/Marabou.py
@@ -3,6 +3,7 @@
     - Christopher Lazarus
     - Kyle Julian
     - Andrew Wu
+    - Teruhiro Tagomori
     
 This file is part of the Marabou project.
 Copyright (c) 2017-2019 by the authors listed in the file AUTHORS
@@ -29,6 +30,7 @@
     from maraboupy.MarabouNetworkONNX import *
 except ImportError:
     warnings.warn("ONNX parser is unavailable because onnx or onnxruntime packages are not installed")
+from maraboupy.MarabouNetworkComposition import MarabouNetworkComposition
 
 def read_nnet(filename, normalize=False):
     """Constructs a MarabouNetworkNnet object from a .nnet file
@@ -74,6 +76,20 @@ def read_onnx(filename, inputNames=None, outputNames=None, reindexOutputVars=Tru
     """
     return MarabouNetworkONNX(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars)
 
+def read_onnx_with_threshould(filename, inputNames=None, outputNames=None, reindexOutputVars=True, threshold=None):
+    """Constructs a MarabouNetworkComposition object from an ONNX file
+
+    Args:
+        filename (str): Path to the ONNX file
+        inputNames (list of str, optional): List of node names corresponding to inputs
+        outputNames (list of str, optional): List of node names corresponding to outputs
+        threshold (int, optional): Threshold for the number of linear equations
+
+    Returns:
+        :class:`~maraboupy.MarabouNetworkONNX.MarabouNetworkONNX`
+    """
+    return MarabouNetworkComposition(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars, threshold=threshold)
+
 def load_query(filename):
     """Load the serialized inputQuery from the given filename
 
diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
new file mode 100644
index 0000000000..ec5b087e8c
--- /dev/null
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -0,0 +1,284 @@
+'''
+Top contributors (to current version):
+    - Teruhiro Tagomori
+
+This file is part of the Marabou project.
+Copyright (c) 2017-2019 by the authors listed in the file AUTHORS
+in the top-level source directory) and their institutional affiliations.
+All rights reserved. See the file COPYING in the top-level source
+directory for licensing information.
+
+MarabouNetworkComposition represents neural networks with piecewise linear constraints derived from the ONNX format
+'''
+
+import numpy as np
+import onnx
+import sys
+import pathlib
+import os
+sys.path.insert(0, os.path.join(str(pathlib.Path(__file__).parent.absolute()), "../"))
+
+from maraboupy import Marabou
+from maraboupy import MarabouCore
+from maraboupy import MarabouNetwork
+from maraboupy import MarabouNetworkONNX
+
+class MarabouNetworkComposition(MarabouNetwork.MarabouNetwork):
+    """Constructs a MarabouNetworkComposition object from an ONNX file
+
+    Args:
+        filename (str): Path to the ONNX file
+        inputNames: (list of str, optional): List of node names corresponding to inputs
+        outputNames: (list of str, optional): List of node names corresponding to outputs
+        reindexOutputVars: (bool, optional): Reindex the variables so that the output variables are immediate after input variables.
+        threshold: (float, optional): Threshold for equation functions. If the absolute value of the lower bound of a variable is less than this value, the variable is considered to be 0.
+
+    Returns:
+        :class:`~maraboupy.Marabou.MarabouNetworkComposition.MarabouNetworkComposition`
+    """
+    def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVars=True, threshold=None):
+        super().__init__()
+        self.shapeMap = {}
+        self.madeGraphEquations = []
+        self.ipqs = []
+        self.ipqToInVars = {}
+        self.ipqToOutVars = {}
+        self.inputVars, self.outputVars = self.getInputOutputVars(filename, inputNames, outputNames, reindexOutputVars=True)
+
+        network = MarabouNetworkONNX.MarabouNetworkONNX(filename, reindexOutputVars=reindexOutputVars, threshold=threshold)
+        print(f'TG: input vars: {network.inputVars}')
+        print(f'TG: output vars: {network.outputVars}')
+
+        network.saveQuery('q1.ipq')
+        self.ipqs.append('q1.ipq')
+        self.ipqToInVars['q1.ipq'] = network.inputVars
+        self.ipqToOutVars['q1.ipq'] = network.outputVars
+
+        index = 2
+        # while if post_split.onnx exists
+        while os.path.exists('post_split.onnx'):
+            # delete network
+            del network
+            network = MarabouNetworkONNX.MarabouNetworkONNX('post_split.onnx', reindexOutputVars=reindexOutputVars, threshold=threshold)
+            network.saveQuery(f'q{index}.ipq')
+            self.ipqs.append(f'q{index}.ipq')
+            self.ipqToInVars[f'q{index}.ipq'] = network.inputVars
+            self.ipqToOutVars[f'q{index}.ipq'] = network.outputVars
+            # self.ipq_to_inVars['q{index}.ipq'] = network.inputVars
+            index += 1
+            print(network.inputVars)
+            print(network.outputVars)
+        print(self.ipqs)
+        # print(inputQuery.inputVars)
+        # MarabouNetworkONNX.readONNX(filename, reindexOutputVars=reindexOutputVars, threshold=threshold)
+        # network, post_network_file = MarabouNetworkONNX.readONNX(filename, reindexOutputVars=reindexOutputVars, threshold=threshold)
+
+    def solve(self):
+        # https://github.com/wu-haoze/Marabou/blob/1a3ca6010b51bba792ef8ddd5e1ccf9119121bd8/resources/runVerify.py#L200-L225
+        options = Marabou.createOptions(verbosity = 1)
+        for i, ipqFile in enumerate(self.ipqs):
+            # ipq = Marabou.load_query(ipqFile)
+            # load inputquery
+            ipq = Marabou.loadQuery(ipqFile)
+
+            if i == 0:
+                self.encodeInput(ipq)
+
+            if i > 0:
+                self.encodeCalculateInputBounds(ipq, i, bounds)
+                # print(bounds)
+            
+            if i == len(self.ipqs) - 1:
+                self.encodeOutput(ipq, i)
+                # exitCode, vals, _ = MarabouCore.solve(ipq, options)
+                # print(f'TG: exit code: {exitCode}')
+                # if exitCode == "sat":
+                #     for j in range(len(self.inputVars)):
+                #         for i in range(self.inputVars[j].size):
+                #             print("input {} = {}".format(i, vals[self.inputVars[j].item(i)]))
+
+                #     for j in range(len(self.outputVars)): #TG: Original のものを表示してしまっているので修正
+                #         for i in range(self.outputVars[j].size):
+                #             print("output {} = {}".format(i, vals[self.outputVars[j].item(i)]))
+                ret, bounds, stats = MarabouCore.calculateBounds(ipq, options)
+                # print(f'TG: bounds: {bounds}')
+            else:
+                ret, bounds, stats = MarabouCore.calculateBounds(ipq, options)
+                print(f'TG: bounds: {bounds}')
+        print(f'TG: bounds: {bounds}')
+
+    def encodeCalculateInputBounds(self, ipq, i, bounds):
+        print('TG: encodeCalculateInputBounds')
+        previousOutputVars = self.ipqToOutVars[f'q{i}.ipq']
+        currentInputVars = self.ipqToInVars[f'q{i+1}.ipq']
+        print('TG: ', self.ipqToOutVars[f'q{i+1}.ipq'])
+        print(f'TG: previous output vars: {previousOutputVars}')
+        print(f'TG: current input vars: {currentInputVars}')      
+        
+        for previousOutputVar, currentInputVar in zip(previousOutputVars, currentInputVars):
+            for previousOutputVarElement, currentInputVarElement in zip(previousOutputVar.flatten(), currentInputVar.flatten()):
+                print(f'TG: previous output var element: {previousOutputVarElement}')
+                print(f'TG: current input var element: {currentInputVarElement}')
+                print(f'TG: bounds: {bounds[previousOutputVarElement]}')
+                ipq.setLowerBound(currentInputVarElement, bounds[previousOutputVarElement][0])
+                ipq.setUpperBound(currentInputVarElement, bounds[previousOutputVarElement][1])
+
+    def encodeInput(self, ipq):
+        inputVars = self.ipqToInVars['q1.ipq']
+        # print('TG: ', self.ipqToOutVars[f'q1.ipq'])
+        # print('TG: ', self.ipqToInVars[f'q2.ipq'])
+        for array in inputVars:
+            for var in array.flatten():
+                ipq.setLowerBound(var, self.lowerBounds[var])
+                ipq.setUpperBound(var, self.upperBounds[var])
+
+    def encodeOutput(self, ipq, i):
+        outputVars = self.ipqToOutVars[f'q{i+1}.ipq']
+        originalOutputVars = self.outputVars
+        print(f'TG: original output vars: {originalOutputVars}')
+        print(f'TG: output vars: {outputVars}')
+        
+        for originalOutputVar, outputVar in zip(originalOutputVars, outputVars):
+            for originalOutputVarElement, outputVarElement in zip(originalOutputVar.flatten(), outputVar.flatten()):                
+                if originalOutputVarElement in self.lowerBounds:
+                    ipq.setLowerBound(outputVarElement, self.lowerBounds[originalOutputVarElement])
+                if originalOutputVarElement in self.upperBounds:
+                    ipq.setUpperBound(outputVarElement, self.upperBounds[originalOutputVarElement])
+
+
+    # def getInputOutVars(self, filename, inputNames, outputNames, reindexOutputVars=True, threshold=None):
+    def getInputOutputVars(self, filename, inputNames, outputNames, reindexOutputVars=True):
+        """Read an ONNX file and create a MarabouNetworkONNX object
+
+        Args:
+            filename: (str): Path to the ONNX file
+            inputNames: (list of str): List of node names corresponding to inputs
+            outputNames: (list of str): List of node names corresponding to outputs
+            reindexOutputVars: (bool): Reindex the variables so that the output variables are immediate after input variables.
+
+        :meta private:
+        """
+        self.filename = filename
+        self.graph = onnx.load(filename).graph
+
+        # Get default inputs/outputs if no names are provided
+        if not inputNames:
+            assert len(self.graph.input) >= 1
+            initNames = [node.name for node in self.graph.initializer]
+            inputNames = [inp.name for inp in self.graph.input if inp.name not in initNames]
+        if not outputNames:
+            assert len(self.graph.output) >= 1
+            initNames = [node.name for node in self.graph.initializer]
+            outputNames = [out.name for out in self.graph.output if out.name not in initNames]
+        elif isinstance(outputNames, str):
+            outputNames = [outputNames]
+
+        # Check that input/outputs are in the graph
+        for name in inputNames:
+            if not len([nde for nde in self.graph.node if name in nde.input]):
+                raise RuntimeError("Input %s not found in graph!" % name)
+        for name in outputNames:
+            if not len([nde for nde in self.graph.node if name in nde.output]):
+                raise RuntimeError("Output %s not found in graph!" % name)
+
+        self.inputNames = inputNames
+        self.outputNames = outputNames
+
+        # Process the shapes and values of the graph while making Marabou equations and constraints
+        self.foundnInputFlags = 0
+
+        # Add shapes for the graph's inputs
+        inputVars = []
+        for node in self.graph.input:
+            self.shapeMap[node.name] = list([dim.dim_value if dim.dim_value > 0 else 1 for dim in node.type.tensor_type.shape.dim])
+
+            # If we find one of the specified inputs, create new variables
+            if node.name in self.inputNames:
+                self.madeGraphEquations += [node.name]
+                self.foundnInputFlags += 1
+                v = self.makeNewVariables(node.name)
+                # self.inputVars += [np.array(self.varMap[node.name])]
+                inputVars += [v]
+        print(f'TG: get input vars: {inputVars}')
+
+        # Add shapes for the graph's outputs
+        outputVars = []
+        for node in self.graph.output:
+            self.shapeMap[node.name] = list([dim.dim_value if dim.dim_value > 0 else 1 for dim in node.type.tensor_type.shape.dim])
+
+            # If we find one of the specified inputs, create new variables
+            if node.name in self.outputNames:
+                self.madeGraphEquations += [node.name]
+                self.foundnInputFlags += 1
+                v = self.makeNewVariables(node.name)
+                outputVars += [v]
+        print(f'TG: get ouput vars: {outputVars}')
+        
+        return inputVars, outputVars
+
+        # # Recursively create remaining shapes and equations as needed
+        # for outputName in self.outputNames: # TG: 
+        #     print(f'TG: outputName: {outputName}')
+        #     self.makeGraphEquations(outputName, True)
+    
+
+
+        # If the given inputNames/outputNames specify only a portion of the network, then we will have
+        # shape information saved not relevant to the portion of the network. Remove extra shapes.
+        # self.cleanShapes() # TG: needed?
+
+        # TG: needed?
+
+        # if reindexOutputVars:
+        #     # Other Marabou input parsers assign output variables immediately after input variables and before any
+        #     # intermediate variables. This function reassigns variable numbering to match other parsers.
+        #     # If this is skipped, the output variables will be the last variables defined.
+        #     self.reassignOutputVariables()
+        # else:
+        #     self.outputVars = [self.varMap[outputName] for outputName in self.outputNames]
+        
+
+
+    def makeNewVariables(self, nodeName):
+        """Assuming the node's shape is known, return a set of new variables in the same shape
+
+        Args:
+            nodeName (str): Name of node
+
+        Returns:
+            (numpy array): Array of variable numbers
+
+        :meta private:
+        """
+        # assert nodeName not in self.varMap
+        shape = self.shapeMap[nodeName]
+        size = np.prod(shape)
+        v = np.array([self.getNewVariable() for _ in range(size)]).reshape(shape)
+        # self.varMap[nodeName] = v
+        assert all([np.equal(np.mod(i, 1), 0) for i in v.reshape(-1)]) # check if integers
+        return v
+
+    def setLowerBound(self, x, v):
+        """Function to set lower bound for variable
+
+        Args:
+            x (int): Variable number to set
+            v (float): Value representing lower bound
+        """
+        if any(x in arr for arr in self.inputVars) or any(x in arr for arr in self.outputVars):
+            self.lowerBounds[x] = v
+        else:
+            raise RuntimeError("Cannot set bounds on input or output variables")
+        
+
+    def setUpperBound(self, x, v):
+        """Function to set upper bound for variable
+
+        Args:
+            x (int): Variable number to set
+            v (float): Value representing upper bound
+        """
+        if any(x in arr for arr in self.inputVars) or any(x in arr for arr in self.outputVars):
+            self.upperBounds[x] = v
+        else:
+            raise RuntimeError("Cannot set bounds on input or output variables")
\ No newline at end of file
diff --git a/maraboupy/MarabouNetworkONNX.py b/maraboupy/MarabouNetworkONNX.py
index 73003181d8..fbf6755062 100644
--- a/maraboupy/MarabouNetworkONNX.py
+++ b/maraboupy/MarabouNetworkONNX.py
@@ -37,9 +37,9 @@ class MarabouNetworkONNX(MarabouNetwork.MarabouNetwork):
     Returns:
         :class:`~maraboupy.Marabou.marabouNetworkONNX.marabouNetworkONNX`
     """
-    def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVars=True):
+    def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVars=True, threshold=None):
         super().__init__()
-        self.readONNX(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars)
+        self.readONNX(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars, threshold=threshold)
 
     def clear(self):
         """Reset values to represent empty network
@@ -67,7 +67,7 @@ def shallowClear(self):
         self.outputNames = None
         self.graph = None
 
-    def readONNX(self, filename, inputNames, outputNames, reindexOutputVars=True):
+    def readONNX(self, filename, inputNames, outputNames, reindexOutputVars=True, threshold=None):
         """Read an ONNX file and create a MarabouNetworkONNX object
 
         Args:
@@ -80,6 +80,11 @@ def readONNX(self, filename, inputNames, outputNames, reindexOutputVars=True):
         """
         self.filename = filename
         self.graph = onnx.load(filename).graph
+        self.threshold = threshold
+        self.thresholdReached = False
+
+        if os.path.exists('post_split.onnx'):
+            os.remove('post_split.onnx')
 
         # Get default inputs/outputs if no names are provided
         if not inputNames:
@@ -185,6 +190,8 @@ def processGraph(self):
         # Recursively create remaining shapes and equations as needed
         for outputName in self.outputNames:
             self.makeGraphEquations(outputName, True)
+            # if self.thresholdReached:
+            #     return
 
     def makeGraphEquations(self, nodeName, makeEquations):
         """Recursively populates self.shapeMap, self.varMap, and self.constantMap while adding equations and constraints
@@ -217,8 +224,15 @@ def makeGraphEquations(self, nodeName, makeEquations):
             raise RuntimeError(err_msg)
 
         # Compute node's shape and create Marabou equations as needed
+        if self.thresholdReached:
+            return
         self.makeMarabouEquations(nodeName, makeEquations)
 
+        if self.threshold is not None:
+            if not self.thresholdReached and len(self.equList) > self.threshold:
+                if self.splitNetworkAtNode(nodeName, networkNamePostSplit='post_split.onnx'):
+                    self.thresholdReached = True
+
         # Create new variables when we find one of the inputs
         if nodeName in self.inputNames:
             self.makeNewVariables(nodeName)

From db5ff08a42eebf0475da7580f687419734ad239d Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Wed, 27 Dec 2023 12:00:19 -0800
Subject: [PATCH 02/12] refactor

---
 maraboupy/MarabouNetworkComposition.py | 76 +++-----------------------
 1 file changed, 9 insertions(+), 67 deletions(-)

diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
index ec5b087e8c..8e96a46e95 100644
--- a/maraboupy/MarabouNetworkComposition.py
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -43,11 +43,9 @@ def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVar
         self.ipqs = []
         self.ipqToInVars = {}
         self.ipqToOutVars = {}
-        self.inputVars, self.outputVars = self.getInputOutputVars(filename, inputNames, outputNames, reindexOutputVars=True)
+        self.inputVars, self.outputVars = self.getInputOutputVars(filename, inputNames, outputNames)
 
         network = MarabouNetworkONNX.MarabouNetworkONNX(filename, reindexOutputVars=reindexOutputVars, threshold=threshold)
-        print(f'TG: input vars: {network.inputVars}')
-        print(f'TG: output vars: {network.outputVars}')
 
         network.saveQuery('q1.ipq')
         self.ipqs.append('q1.ipq')
@@ -55,7 +53,7 @@ def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVar
         self.ipqToOutVars['q1.ipq'] = network.outputVars
 
         index = 2
-        # while if post_split.onnx exists
+
         while os.path.exists('post_split.onnx'):
             # delete network
             del network
@@ -64,21 +62,13 @@ def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVar
             self.ipqs.append(f'q{index}.ipq')
             self.ipqToInVars[f'q{index}.ipq'] = network.inputVars
             self.ipqToOutVars[f'q{index}.ipq'] = network.outputVars
-            # self.ipq_to_inVars['q{index}.ipq'] = network.inputVars
             index += 1
-            print(network.inputVars)
-            print(network.outputVars)
-        print(self.ipqs)
-        # print(inputQuery.inputVars)
-        # MarabouNetworkONNX.readONNX(filename, reindexOutputVars=reindexOutputVars, threshold=threshold)
-        # network, post_network_file = MarabouNetworkONNX.readONNX(filename, reindexOutputVars=reindexOutputVars, threshold=threshold)
 
     def solve(self):
         # https://github.com/wu-haoze/Marabou/blob/1a3ca6010b51bba792ef8ddd5e1ccf9119121bd8/resources/runVerify.py#L200-L225
-        options = Marabou.createOptions(verbosity = 1)
+        options = Marabou.createOptions(verbosity = 1) # TG: Option を引数でとる
         for i, ipqFile in enumerate(self.ipqs):
-            # ipq = Marabou.load_query(ipqFile)
-            # load inputquery
+            # load input query
             ipq = Marabou.loadQuery(ipqFile)
 
             if i == 0:
@@ -86,7 +76,6 @@ def solve(self):
 
             if i > 0:
                 self.encodeCalculateInputBounds(ipq, i, bounds)
-                # print(bounds)
             
             if i == len(self.ipqs) - 1:
                 self.encodeOutput(ipq, i)
@@ -101,32 +90,20 @@ def solve(self):
                 #         for i in range(self.outputVars[j].size):
                 #             print("output {} = {}".format(i, vals[self.outputVars[j].item(i)]))
                 ret, bounds, stats = MarabouCore.calculateBounds(ipq, options)
-                # print(f'TG: bounds: {bounds}')
             else:
                 ret, bounds, stats = MarabouCore.calculateBounds(ipq, options)
-                print(f'TG: bounds: {bounds}')
-        print(f'TG: bounds: {bounds}')
 
     def encodeCalculateInputBounds(self, ipq, i, bounds):
-        print('TG: encodeCalculateInputBounds')
         previousOutputVars = self.ipqToOutVars[f'q{i}.ipq']
-        currentInputVars = self.ipqToInVars[f'q{i+1}.ipq']
-        print('TG: ', self.ipqToOutVars[f'q{i+1}.ipq'])
-        print(f'TG: previous output vars: {previousOutputVars}')
-        print(f'TG: current input vars: {currentInputVars}')      
+        currentInputVars = self.ipqToInVars[f'q{i+1}.ipq']    
         
         for previousOutputVar, currentInputVar in zip(previousOutputVars, currentInputVars):
             for previousOutputVarElement, currentInputVarElement in zip(previousOutputVar.flatten(), currentInputVar.flatten()):
-                print(f'TG: previous output var element: {previousOutputVarElement}')
-                print(f'TG: current input var element: {currentInputVarElement}')
-                print(f'TG: bounds: {bounds[previousOutputVarElement]}')
                 ipq.setLowerBound(currentInputVarElement, bounds[previousOutputVarElement][0])
                 ipq.setUpperBound(currentInputVarElement, bounds[previousOutputVarElement][1])
 
     def encodeInput(self, ipq):
         inputVars = self.ipqToInVars['q1.ipq']
-        # print('TG: ', self.ipqToOutVars[f'q1.ipq'])
-        # print('TG: ', self.ipqToInVars[f'q2.ipq'])
         for array in inputVars:
             for var in array.flatten():
                 ipq.setLowerBound(var, self.lowerBounds[var])
@@ -135,8 +112,6 @@ def encodeInput(self, ipq):
     def encodeOutput(self, ipq, i):
         outputVars = self.ipqToOutVars[f'q{i+1}.ipq']
         originalOutputVars = self.outputVars
-        print(f'TG: original output vars: {originalOutputVars}')
-        print(f'TG: output vars: {outputVars}')
         
         for originalOutputVar, outputVar in zip(originalOutputVars, outputVars):
             for originalOutputVarElement, outputVarElement in zip(originalOutputVar.flatten(), outputVar.flatten()):                
@@ -145,16 +120,13 @@ def encodeOutput(self, ipq, i):
                 if originalOutputVarElement in self.upperBounds:
                     ipq.setUpperBound(outputVarElement, self.upperBounds[originalOutputVarElement])
 
-
-    # def getInputOutVars(self, filename, inputNames, outputNames, reindexOutputVars=True, threshold=None):
-    def getInputOutputVars(self, filename, inputNames, outputNames, reindexOutputVars=True):
-        """Read an ONNX file and create a MarabouNetworkONNX object
+    def getInputOutputVars(self, filename, inputNames, outputNames):
+        """Get input and output variables of an original network
 
         Args:
             filename: (str): Path to the ONNX file
             inputNames: (list of str): List of node names corresponding to inputs
             outputNames: (list of str): List of node names corresponding to outputs
-            reindexOutputVars: (bool): Reindex the variables so that the output variables are immediate after input variables.
 
         :meta private:
         """
@@ -197,9 +169,7 @@ def getInputOutputVars(self, filename, inputNames, outputNames, reindexOutputVar
                 self.madeGraphEquations += [node.name]
                 self.foundnInputFlags += 1
                 v = self.makeNewVariables(node.name)
-                # self.inputVars += [np.array(self.varMap[node.name])]
                 inputVars += [v]
-        print(f'TG: get input vars: {inputVars}')
 
         # Add shapes for the graph's outputs
         outputVars = []
@@ -212,33 +182,8 @@ def getInputOutputVars(self, filename, inputNames, outputNames, reindexOutputVar
                 self.foundnInputFlags += 1
                 v = self.makeNewVariables(node.name)
                 outputVars += [v]
-        print(f'TG: get ouput vars: {outputVars}')
-        
         return inputVars, outputVars
 
-        # # Recursively create remaining shapes and equations as needed
-        # for outputName in self.outputNames: # TG: 
-        #     print(f'TG: outputName: {outputName}')
-        #     self.makeGraphEquations(outputName, True)
-    
-
-
-        # If the given inputNames/outputNames specify only a portion of the network, then we will have
-        # shape information saved not relevant to the portion of the network. Remove extra shapes.
-        # self.cleanShapes() # TG: needed?
-
-        # TG: needed?
-
-        # if reindexOutputVars:
-        #     # Other Marabou input parsers assign output variables immediately after input variables and before any
-        #     # intermediate variables. This function reassigns variable numbering to match other parsers.
-        #     # If this is skipped, the output variables will be the last variables defined.
-        #     self.reassignOutputVariables()
-        # else:
-        #     self.outputVars = [self.varMap[outputName] for outputName in self.outputNames]
-        
-
-
     def makeNewVariables(self, nodeName):
         """Assuming the node's shape is known, return a set of new variables in the same shape
 
@@ -250,11 +195,9 @@ def makeNewVariables(self, nodeName):
 
         :meta private:
         """
-        # assert nodeName not in self.varMap
         shape = self.shapeMap[nodeName]
         size = np.prod(shape)
         v = np.array([self.getNewVariable() for _ in range(size)]).reshape(shape)
-        # self.varMap[nodeName] = v
         assert all([np.equal(np.mod(i, 1), 0) for i in v.reshape(-1)]) # check if integers
         return v
 
@@ -268,9 +211,8 @@ def setLowerBound(self, x, v):
         if any(x in arr for arr in self.inputVars) or any(x in arr for arr in self.outputVars):
             self.lowerBounds[x] = v
         else:
-            raise RuntimeError("Cannot set bounds on input or output variables")
+            raise RuntimeError("Can set bounds only on either input or output variables")
         
-
     def setUpperBound(self, x, v):
         """Function to set upper bound for variable
 
@@ -281,4 +223,4 @@ def setUpperBound(self, x, v):
         if any(x in arr for arr in self.inputVars) or any(x in arr for arr in self.outputVars):
             self.upperBounds[x] = v
         else:
-            raise RuntimeError("Cannot set bounds on input or output variables")
\ No newline at end of file
+            raise RuntimeError("Can set bounds only on either input or output variables")
\ No newline at end of file

From 0f6d38a0529cdcf627ae3f2b28ddea5c9d90b17d Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Wed, 27 Dec 2023 12:04:54 -0800
Subject: [PATCH 03/12] change args of solve mthod

---
 maraboupy/MarabouNetworkComposition.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
index 8e96a46e95..741708aada 100644
--- a/maraboupy/MarabouNetworkComposition.py
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -64,13 +64,26 @@ def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVar
             self.ipqToOutVars[f'q{index}.ipq'] = network.outputVars
             index += 1
 
-    def solve(self):
+    def solve(self, filename="", verbose=True, options=None):
+        """Function to solve query represented by this network
+
+        Args:
+            filename (string): Path for redirecting output
+            verbose (bool): If true, print out solution after solve finishes
+            options (:class:`~maraboupy.MarabouCore.Options`): Object for specifying Marabou options, defaults to None
+
+        Returns:
+            (tuple): tuple containing:
+                - exitCode (str): A string representing the exit code (sat/unsat/TIMEOUT/ERROR/UNKNOWN/QUIT_REQUESTED).
+                - vals (Dict[int, float]): Empty dictionary if UNSAT, otherwise a dictionary of SATisfying values for variables
+                - stats (:class:`~maraboupy.MarabouCore.Statistics`): A Statistics object to how Marabou performed
+        """
         # https://github.com/wu-haoze/Marabou/blob/1a3ca6010b51bba792ef8ddd5e1ccf9119121bd8/resources/runVerify.py#L200-L225
-        options = Marabou.createOptions(verbosity = 1) # TG: Option を引数でとる
+        if options == None:
+            options = MarabouCore.Options()        
         for i, ipqFile in enumerate(self.ipqs):
             # load input query
             ipq = Marabou.loadQuery(ipqFile)
-
             if i == 0:
                 self.encodeInput(ipq)
 

From e614fb13151ebef7ec38fd189112334ce4aa8147 Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Thu, 28 Dec 2023 23:27:49 -0800
Subject: [PATCH 04/12] change the logic for the last calcbounds

---
 maraboupy/MarabouNetworkComposition.py | 24 +++++++++---------------
 1 file changed, 9 insertions(+), 15 deletions(-)

diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
index 741708aada..0057d0df6e 100644
--- a/maraboupy/MarabouNetworkComposition.py
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -68,7 +68,7 @@ def solve(self, filename="", verbose=True, options=None):
         """Function to solve query represented by this network
 
         Args:
-            filename (string): Path for redirecting output
+            filename (string): Path for redirecting output (Only for the last subnet)
             verbose (bool): If true, print out solution after solve finishes
             options (:class:`~maraboupy.MarabouCore.Options`): Object for specifying Marabou options, defaults to None
 
@@ -76,9 +76,8 @@ def solve(self, filename="", verbose=True, options=None):
             (tuple): tuple containing:
                 - exitCode (str): A string representing the exit code (sat/unsat/TIMEOUT/ERROR/UNKNOWN/QUIT_REQUESTED).
                 - vals (Dict[int, float]): Empty dictionary if UNSAT, otherwise a dictionary of SATisfying values for variables
-                - stats (:class:`~maraboupy.MarabouCore.Statistics`): A Statistics object to how Marabou performed
+                - stats (:class:`~maraboupy.MarabouCore.Statistics`): A Statistics object to how Marabou performed (Only for the last subnet)
         """
-        # https://github.com/wu-haoze/Marabou/blob/1a3ca6010b51bba792ef8ddd5e1ccf9119121bd8/resources/runVerify.py#L200-L225
         if options == None:
             options = MarabouCore.Options()        
         for i, ipqFile in enumerate(self.ipqs):
@@ -92,19 +91,14 @@ def solve(self, filename="", verbose=True, options=None):
             
             if i == len(self.ipqs) - 1:
                 self.encodeOutput(ipq, i)
-                # exitCode, vals, _ = MarabouCore.solve(ipq, options)
-                # print(f'TG: exit code: {exitCode}')
-                # if exitCode == "sat":
-                #     for j in range(len(self.inputVars)):
-                #         for i in range(self.inputVars[j].size):
-                #             print("input {} = {}".format(i, vals[self.inputVars[j].item(i)]))
-
-                #     for j in range(len(self.outputVars)): #TG: Original のものを表示してしまっているので修正
-                #         for i in range(self.outputVars[j].size):
-                #             print("output {} = {}".format(i, vals[self.outputVars[j].item(i)]))
-                ret, bounds, stats = MarabouCore.calculateBounds(ipq, options)
+                exitCode, bounds, stats = MarabouCore.calculateBounds(ipq, options, str(filename))
+                if exitCode == "":
+                    exitCode = "UNKNOWN"
+                if verbose:
+                    print(exitCode)
+                return [exitCode, {}, stats]
             else:
-                ret, bounds, stats = MarabouCore.calculateBounds(ipq, options)
+                _, bounds, _ = MarabouCore.calculateBounds(ipq, options)
 
     def encodeCalculateInputBounds(self, ipq, i, bounds):
         previousOutputVars = self.ipqToOutVars[f'q{i}.ipq']

From a12ddc6c43655af356765f1751d14769824ddf02 Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Mon, 1 Jan 2024 23:15:32 -0800
Subject: [PATCH 05/12] Change threshold name to maxNumberOfLinearEquations and
 remove reindexOutputVars from MarabouNetworkComposition

---
 maraboupy/Marabou.py                   | 11 +++++++----
 maraboupy/MarabouNetworkComposition.py | 11 ++++++-----
 maraboupy/MarabouNetworkONNX.py        | 17 +++++++++--------
 3 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/maraboupy/Marabou.py b/maraboupy/Marabou.py
index cbf4c59d74..92a815e8c3 100644
--- a/maraboupy/Marabou.py
+++ b/maraboupy/Marabou.py
@@ -76,19 +76,22 @@ def read_onnx(filename, inputNames=None, outputNames=None, reindexOutputVars=Tru
     """
     return MarabouNetworkONNX(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars)
 
-def read_onnx_with_threshould(filename, inputNames=None, outputNames=None, reindexOutputVars=True, threshold=None):
+def read_onnx_with_threshold(filename, inputNames=None, outputNames=None, maxNumberOfLinearEquations=None):
     """Constructs a MarabouNetworkComposition object from an ONNX file
 
     Args:
         filename (str): Path to the ONNX file
         inputNames (list of str, optional): List of node names corresponding to inputs
         outputNames (list of str, optional): List of node names corresponding to outputs
-        threshold (int, optional): Threshold for the number of linear equations
+        maxNumberOfLinearEquations (int, optional): Threshold for the number of linear equations.
+                                                    If the number of linear equations is greater than this threshold,
+                                                    the network will be split into two networks. Defaults to None.
 
     Returns:
-        :class:`~maraboupy.MarabouNetworkONNX.MarabouNetworkONNX`
+        :class:`~maraboupy.MarabouNetworkComposition.MarabouNetworkComposition`
     """
-    return MarabouNetworkComposition(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars, threshold=threshold)
+    return MarabouNetworkComposition(filename, inputNames, outputNames,
+                                        maxNumberOfLinearEquations=maxNumberOfLinearEquations)
 
 def load_query(filename):
     """Load the serialized inputQuery from the given filename
diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
index 0057d0df6e..d4ef95564c 100644
--- a/maraboupy/MarabouNetworkComposition.py
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -30,13 +30,14 @@ class MarabouNetworkComposition(MarabouNetwork.MarabouNetwork):
         filename (str): Path to the ONNX file
         inputNames: (list of str, optional): List of node names corresponding to inputs
         outputNames: (list of str, optional): List of node names corresponding to outputs
-        reindexOutputVars: (bool, optional): Reindex the variables so that the output variables are immediate after input variables.
-        threshold: (float, optional): Threshold for equation functions. If the absolute value of the lower bound of a variable is less than this value, the variable is considered to be 0.
+        maxNumberOfLinearEquations (int, optional): Threshold for the number of linear equations.
+                                                    If the number of linear equations is greater than this threshold,
+                                                    the network will be split into two networks. Defaults to None.
 
     Returns:
         :class:`~maraboupy.Marabou.MarabouNetworkComposition.MarabouNetworkComposition`
     """
-    def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVars=True, threshold=None):
+    def __init__(self, filename, inputNames=None, outputNames=None, maxNumberOfLinearEquations=None):
         super().__init__()
         self.shapeMap = {}
         self.madeGraphEquations = []
@@ -45,7 +46,7 @@ def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVar
         self.ipqToOutVars = {}
         self.inputVars, self.outputVars = self.getInputOutputVars(filename, inputNames, outputNames)
 
-        network = MarabouNetworkONNX.MarabouNetworkONNX(filename, reindexOutputVars=reindexOutputVars, threshold=threshold)
+        network = MarabouNetworkONNX.MarabouNetworkONNX(filename, maxNumberOfLinearEquations=maxNumberOfLinearEquations)
 
         network.saveQuery('q1.ipq')
         self.ipqs.append('q1.ipq')
@@ -57,7 +58,7 @@ def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVar
         while os.path.exists('post_split.onnx'):
             # delete network
             del network
-            network = MarabouNetworkONNX.MarabouNetworkONNX('post_split.onnx', reindexOutputVars=reindexOutputVars, threshold=threshold)
+            network = MarabouNetworkONNX.MarabouNetworkONNX('post_split.onnx', maxNumberOfLinearEquations=maxNumberOfLinearEquations)
             network.saveQuery(f'q{index}.ipq')
             self.ipqs.append(f'q{index}.ipq')
             self.ipqToInVars[f'q{index}.ipq'] = network.inputVars
diff --git a/maraboupy/MarabouNetworkONNX.py b/maraboupy/MarabouNetworkONNX.py
index fbf6755062..3fe72b4c15 100644
--- a/maraboupy/MarabouNetworkONNX.py
+++ b/maraboupy/MarabouNetworkONNX.py
@@ -37,9 +37,9 @@ class MarabouNetworkONNX(MarabouNetwork.MarabouNetwork):
     Returns:
         :class:`~maraboupy.Marabou.marabouNetworkONNX.marabouNetworkONNX`
     """
-    def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVars=True, threshold=None):
+    def __init__(self, filename, inputNames=None, outputNames=None, reindexOutputVars=True, maxNumberOfLinearEquations=None):
         super().__init__()
-        self.readONNX(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars, threshold=threshold)
+        self.readONNX(filename, inputNames, outputNames, reindexOutputVars=reindexOutputVars, maxNumberOfLinearEquations=maxNumberOfLinearEquations)
 
     def clear(self):
         """Reset values to represent empty network
@@ -67,7 +67,7 @@ def shallowClear(self):
         self.outputNames = None
         self.graph = None
 
-    def readONNX(self, filename, inputNames, outputNames, reindexOutputVars=True, threshold=None):
+    def readONNX(self, filename, inputNames, outputNames, reindexOutputVars=True, maxNumberOfLinearEquations=None):
         """Read an ONNX file and create a MarabouNetworkONNX object
 
         Args:
@@ -75,12 +75,15 @@ def readONNX(self, filename, inputNames, outputNames, reindexOutputVars=True, th
             inputNames: (list of str): List of node names corresponding to inputs
             outputNames: (list of str): List of node names corresponding to outputs
             reindexOutputVars: (bool): Reindex the variables so that the output variables are immediate after input variables.
+            maxNumberOfLinearEquations (int, optional): Threshold for the number of linear equations.
+                                                        If the number of linear equations is greater than this threshold,
+                                                        the network will be split into two networks. Defaults to None.
 
         :meta private:
         """
         self.filename = filename
         self.graph = onnx.load(filename).graph
-        self.threshold = threshold
+        self.maxNumberOfLinearEquations = maxNumberOfLinearEquations
         self.thresholdReached = False
 
         if os.path.exists('post_split.onnx'):
@@ -190,8 +193,6 @@ def processGraph(self):
         # Recursively create remaining shapes and equations as needed
         for outputName in self.outputNames:
             self.makeGraphEquations(outputName, True)
-            # if self.thresholdReached:
-            #     return
 
     def makeGraphEquations(self, nodeName, makeEquations):
         """Recursively populates self.shapeMap, self.varMap, and self.constantMap while adding equations and constraints
@@ -228,8 +229,8 @@ def makeGraphEquations(self, nodeName, makeEquations):
             return
         self.makeMarabouEquations(nodeName, makeEquations)
 
-        if self.threshold is not None:
-            if not self.thresholdReached and len(self.equList) > self.threshold:
+        if self.maxNumberOfLinearEquations is not None:
+            if not self.thresholdReached and len(self.equList) > self.maxNumberOfLinearEquations:
                 if self.splitNetworkAtNode(nodeName, networkNamePostSplit='post_split.onnx'):
                     self.thresholdReached = True
 

From 1a4717018629bf80d099126d71ffd3c995db55bc Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Mon, 1 Jan 2024 23:20:33 -0800
Subject: [PATCH 06/12] add class description to MarabouNetworkComposition

---
 maraboupy/MarabouNetworkComposition.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
index d4ef95564c..4fafaed558 100644
--- a/maraboupy/MarabouNetworkComposition.py
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -25,6 +25,8 @@
 
 class MarabouNetworkComposition(MarabouNetwork.MarabouNetwork):
     """Constructs a MarabouNetworkComposition object from an ONNX file
+        This class splits into subnets every time the number of linear equations reaches maxNumberOfLinearEquations.
+        It provides the function to propagate bounds for each subnet.
 
     Args:
         filename (str): Path to the ONNX file

From aa9c509e26954538456a52796fe4e97ca2010b3b Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Tue, 2 Jan 2024 23:25:30 -0800
Subject: [PATCH 07/12] add comments

---
 maraboupy/MarabouNetworkComposition.py | 84 +++++++++++++++++++++++---
 1 file changed, 74 insertions(+), 10 deletions(-)

diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
index 4fafaed558..598f359fd4 100644
--- a/maraboupy/MarabouNetworkComposition.py
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -48,23 +48,39 @@ def __init__(self, filename, inputNames=None, outputNames=None, maxNumberOfLinea
         self.ipqToOutVars = {}
         self.inputVars, self.outputVars = self.getInputOutputVars(filename, inputNames, outputNames)
 
+        # Instantiate the first subnet
         network = MarabouNetworkONNX.MarabouNetworkONNX(filename, maxNumberOfLinearEquations=maxNumberOfLinearEquations)
 
-        network.saveQuery('q1.ipq')
-        self.ipqs.append('q1.ipq')
-        self.ipqToInVars['q1.ipq'] = network.inputVars
-        self.ipqToOutVars['q1.ipq'] = network.outputVars
+        savedInputQueryName = 'q1.ipq'
+        network.saveQuery(savedInputQueryName)
+        self.ipqs.append(savedInputQueryName)
+        self.ipqToInVars[savedInputQueryName] = network.inputVars
+        self.ipqToOutVars[savedInputQueryName] = network.outputVars
 
+        # index of ipq file
         index = 2
 
         while os.path.exists('post_split.onnx'):
             # delete network
             del network
-            network = MarabouNetworkONNX.MarabouNetworkONNX('post_split.onnx', maxNumberOfLinearEquations=maxNumberOfLinearEquations)
-            network.saveQuery(f'q{index}.ipq')
-            self.ipqs.append(f'q{index}.ipq')
-            self.ipqToInVars[f'q{index}.ipq'] = network.inputVars
-            self.ipqToOutVars[f'q{index}.ipq'] = network.outputVars
+
+            # Instantiate the next subnet
+            network = MarabouNetworkONNX.MarabouNetworkONNX('post_split.onnx',
+                                                            maxNumberOfLinearEquations=maxNumberOfLinearEquations)
+            # name of the input query file
+            savedInputQueryName = f'q{index}.ipq'
+
+            # save input query
+            network.saveQuery(savedInputQueryName)
+
+            # append input query to the lsit
+            self.ipqs.append(savedInputQueryName)
+
+            # save input and output variables so that this can map them to the next input query
+            self.ipqToInVars[savedInputQueryName] = network.inputVars
+            self.ipqToOutVars[savedInputQueryName] = network.outputVars
+            
+            # increment index
             index += 1
 
     def solve(self, filename="", verbose=True, options=None):
@@ -77,7 +93,7 @@ def solve(self, filename="", verbose=True, options=None):
 
         Returns:
             (tuple): tuple containing:
-                - exitCode (str): A string representing the exit code (sat/unsat/TIMEOUT/ERROR/UNKNOWN/QUIT_REQUESTED).
+                - exitCode (str): A string representing the exit code (unsat/TIMEOUT/ERROR/UNKNOWN/QUIT_REQUESTED).
                 - vals (Dict[int, float]): Empty dictionary if UNSAT, otherwise a dictionary of SATisfying values for variables
                 - stats (:class:`~maraboupy.MarabouCore.Statistics`): A Statistics object to how Marabou performed (Only for the last subnet)
         """
@@ -86,14 +102,20 @@ def solve(self, filename="", verbose=True, options=None):
         for i, ipqFile in enumerate(self.ipqs):
             # load input query
             ipq = Marabou.loadQuery(ipqFile)
+
+            # If the first subnet, encode input variables with the input bounds of the original network
             if i == 0:
                 self.encodeInput(ipq)
 
+            # If not the first subnet, encode input variables with the output bounds of the previous subwork
             if i > 0:
                 self.encodeCalculateInputBounds(ipq, i, bounds)
             
             if i == len(self.ipqs) - 1:
+                # If the last subnet, encode output variables with the output bounds of the original network
                 self.encodeOutput(ipq, i)
+
+                # If the last subnet, propagate bounds and return the exit code, values, and statistics
                 exitCode, bounds, stats = MarabouCore.calculateBounds(ipq, options, str(filename))
                 if exitCode == "":
                     exitCode = "UNKNOWN"
@@ -101,28 +123,70 @@ def solve(self, filename="", verbose=True, options=None):
                     print(exitCode)
                 return [exitCode, {}, stats]
             else:
+                # If not the last subnet, propagate bounds
                 _, bounds, _ = MarabouCore.calculateBounds(ipq, options)
 
     def encodeCalculateInputBounds(self, ipq, i, bounds):
+        """Function to encode input variables and calculate bounds for the next subnet
+        
+        Args:
+            ipq (:class:`~maraboupy.MarabouCore.InputQuery`): InputQuery object to encode input variables
+            i (int): Index of the subnet
+            bounds (dict): Dictionary containing bounds for variables
+
+        Returns:
+            None
+
+        :meta private:
+        """
+        # Output variables of the previous subnet
         previousOutputVars = self.ipqToOutVars[f'q{i}.ipq']
+
+        # Input variables of the current subnet
         currentInputVars = self.ipqToInVars[f'q{i+1}.ipq']    
         
+        # Set bounds for the current subnet
         for previousOutputVar, currentInputVar in zip(previousOutputVars, currentInputVars):
             for previousOutputVarElement, currentInputVarElement in zip(previousOutputVar.flatten(), currentInputVar.flatten()):
                 ipq.setLowerBound(currentInputVarElement, bounds[previousOutputVarElement][0])
                 ipq.setUpperBound(currentInputVarElement, bounds[previousOutputVarElement][1])
 
     def encodeInput(self, ipq):
+        """Function to encode input variables
+        
+        Args:
+            ipq (:class:`~maraboupy.MarabouCore.InputQuery`): InputQuery object to encode input variables
+        Returns:
+            None
+
+        :meta private:
+        """
         inputVars = self.ipqToInVars['q1.ipq']
+
+        # Set bounds for the first subnet
         for array in inputVars:
             for var in array.flatten():
                 ipq.setLowerBound(var, self.lowerBounds[var])
                 ipq.setUpperBound(var, self.upperBounds[var])
 
     def encodeOutput(self, ipq, i):
+        """Function to encode output variables
+        Args:
+            ipq:    (:class:`~maraboupy.MarabouCore.InputQuery`): InputQuery object to encode output variables
+            i:      (int): Index of the subnet
+        
+        Returns:
+            None
+
+        :meta private:
+        """
+        # Output variables of the current subnet
         outputVars = self.ipqToOutVars[f'q{i+1}.ipq']
+
+        # Set bounds for the current subnet
         originalOutputVars = self.outputVars
         
+        # Set bounds for the last subnet
         for originalOutputVar, outputVar in zip(originalOutputVars, outputVars):
             for originalOutputVarElement, outputVarElement in zip(originalOutputVar.flatten(), outputVar.flatten()):                
                 if originalOutputVarElement in self.lowerBounds:

From f59d01d0e8d8b1f5a53df34ab647fa3ee8da09d8 Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Wed, 3 Jan 2024 21:08:01 -0800
Subject: [PATCH 08/12] add test code

---
 maraboupy/test/test_network_composition.py | 121 +++++++++++++++++++++
 1 file changed, 121 insertions(+)
 create mode 100644 maraboupy/test/test_network_composition.py

diff --git a/maraboupy/test/test_network_composition.py b/maraboupy/test/test_network_composition.py
new file mode 100644
index 0000000000..f4faf945c9
--- /dev/null
+++ b/maraboupy/test/test_network_composition.py
@@ -0,0 +1,121 @@
+# Tests MarabouNetwork features not tested by it's subclasses
+from maraboupy import Marabou
+from maraboupy import MarabouCore
+import os
+import numpy as np
+
+# Global settings
+OPT = Marabou.createOptions(verbosity = 0) # Turn off printing
+TOL = 1e-6                                 # Set tolerance for checking Marabou evaluations
+NETWORK_FOLDER = "../../resources/nnet/"   # Folder for test networks
+NETWORK_ONNX_FOLDER = "../../resources/onnx/"   # Folder for test networks in onnx format
+
+def test_zero_split_unknown():
+    """
+    Tests that a network with no splits is correctly read and solved as unknown
+    """
+    filename = 'fc1.onnx'
+    filename = os.path.join(os.path.dirname(__file__), NETWORK_ONNX_FOLDER, filename)
+    network = Marabou.read_onnx_with_threshold(filename, maxNumberOfLinearEquations=100)
+
+    # check that the network has one split
+    assert len(network.ipqs) == 1
+
+    network.setLowerBound(network.inputVars[0][0][0], 3)
+    network.setUpperBound(network.inputVars[0][0][0], 5)
+    network.setLowerBound(network.inputVars[0][0][1], 3)
+    network.setUpperBound(network.inputVars[0][0][1], 10)
+
+    exitCode, _, _ = network.solve(options=OPT)
+
+    assert exitCode == "UNKNOWN"
+    
+def test_zero_split_unsat():
+    """ 
+    Tests that a network with no splits is correctly read and solved as unsat
+    """
+    filename = 'fc1.onnx'
+    filename = os.path.join(os.path.dirname(__file__), NETWORK_ONNX_FOLDER, filename)
+    network = Marabou.read_onnx_with_threshold(filename, maxNumberOfLinearEquations=100)
+
+    # check that the network has no splits
+    assert len(network.ipqs) == 1
+
+    network.setLowerBound(network.inputVars[0][0][0], 3)
+    network.setUpperBound(network.inputVars[0][0][0], 5)
+    network.setLowerBound(network.inputVars[0][0][1], 3)
+    network.setUpperBound(network.inputVars[0][0][1], 10)
+
+    network.setLowerBound(network.outputVars[0][0][0], 100)
+
+    exitCode, _, _ = network.solve(options=OPT)
+
+    assert exitCode == "unsat"
+    
+    network = Marabou.read_onnx(filename)
+    network.setLowerBound(network.inputVars[0][0][0], 3)
+    network.setUpperBound(network.inputVars[0][0][0], 5)
+    network.setLowerBound(network.inputVars[0][0][1], 3)
+    network.setUpperBound(network.inputVars[0][0][1], 10)
+
+    network.setLowerBound(network.outputVars[0][0][0], 100)
+
+    exitCode2, _, _ = network.calculateBounds(options=OPT)
+
+    # exitCode2 should be also unsat
+    assert exitCode == exitCode2
+
+def test_one_split_unknown():
+    """
+    Tests that a network with one split is correctly read and solved as unknown
+    """
+    filename = 'fc1.onnx'
+    filename = os.path.join(os.path.dirname(__file__), NETWORK_ONNX_FOLDER, filename)
+    network = Marabou.read_onnx_with_threshold(filename, maxNumberOfLinearEquations=50)
+    
+    # check that the network has one split
+    assert len(network.ipqs) == 2
+
+    network.setLowerBound(network.inputVars[0][0][0], 3)
+    network.setUpperBound(network.inputVars[0][0][0], 5)
+    network.setLowerBound(network.inputVars[0][0][1], 3)
+    network.setUpperBound(network.inputVars[0][0][1], 10)
+
+    exitCode, _, _ = network.solve(options=OPT)
+
+    assert exitCode == "UNKNOWN"
+
+def test_one_split_unsat():
+    """ 
+    Tests that a network with one split is correctly read and solved as unsat
+    """
+    filename = 'fc1.onnx'
+    filename = os.path.join(os.path.dirname(__file__), NETWORK_ONNX_FOLDER, filename)
+    network = Marabou.read_onnx_with_threshold(filename, maxNumberOfLinearEquations=50)
+
+    # check that the network has one split
+    assert len(network.ipqs) == 2
+
+    network.setLowerBound(network.inputVars[0][0][0], 3)
+    network.setUpperBound(network.inputVars[0][0][0], 5)
+    network.setLowerBound(network.inputVars[0][0][1], 3)
+    network.setUpperBound(network.inputVars[0][0][1], 10)
+
+    network.setLowerBound(network.outputVars[0][0][0], 100)
+
+    exitCode, _, _ = network.solve(options=OPT)
+
+    assert exitCode == "unsat"
+    
+    network = Marabou.read_onnx(filename)
+    network.setLowerBound(network.inputVars[0][0][0], 3)
+    network.setUpperBound(network.inputVars[0][0][0], 5)
+    network.setLowerBound(network.inputVars[0][0][1], 3)
+    network.setUpperBound(network.inputVars[0][0][1], 10)
+
+    network.setLowerBound(network.outputVars[0][0][0], 100)
+
+    exitCode2, _, _ = network.calculateBounds(options=OPT)
+
+    # exitCode2 should be also unsat
+    assert exitCode == exitCode2
\ No newline at end of file

From 1aca0c8c4d93f54f4996029142f65c027e80d450 Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Wed, 3 Jan 2024 23:11:12 -0800
Subject: [PATCH 09/12] cosmetic change

---
 maraboupy/MarabouNetworkComposition.py     | 17 +++++++++--------
 maraboupy/test/test_network_composition.py |  5 +----
 2 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/maraboupy/MarabouNetworkComposition.py b/maraboupy/MarabouNetworkComposition.py
index 598f359fd4..4e924ab65b 100644
--- a/maraboupy/MarabouNetworkComposition.py
+++ b/maraboupy/MarabouNetworkComposition.py
@@ -8,7 +8,7 @@
 All rights reserved. See the file COPYING in the top-level source
 directory for licensing information.
 
-MarabouNetworkComposition represents neural networks with piecewise linear constraints derived from the ONNX format
+MarabouNetworkComposition represents split subnets of a neural network with piecewise linear constraints derived from the ONNX format
 '''
 
 import numpy as np
@@ -94,11 +94,12 @@ def solve(self, filename="", verbose=True, options=None):
         Returns:
             (tuple): tuple containing:
                 - exitCode (str): A string representing the exit code (unsat/TIMEOUT/ERROR/UNKNOWN/QUIT_REQUESTED).
-                - vals (Dict[int, float]): Empty dictionary if UNSAT, otherwise a dictionary of SATisfying values for variables
+                - vals (Dict[int, float]): Empty dictionary. This is for compatibility with MarabouNetwork.
                 - stats (:class:`~maraboupy.MarabouCore.Statistics`): A Statistics object to how Marabou performed (Only for the last subnet)
         """
         if options == None:
             options = MarabouCore.Options()        
+        
         for i, ipqFile in enumerate(self.ipqs):
             # load input query
             ipq = Marabou.loadQuery(ipqFile)
@@ -127,12 +128,12 @@ def solve(self, filename="", verbose=True, options=None):
                 _, bounds, _ = MarabouCore.calculateBounds(ipq, options)
 
     def encodeCalculateInputBounds(self, ipq, i, bounds):
-        """Function to encode input variables and calculate bounds for the next subnet
-        
+        """Function to encode input variables and set bounds for the current subnet
+
         Args:
             ipq (:class:`~maraboupy.MarabouCore.InputQuery`): InputQuery object to encode input variables
-            i (int): Index of the subnet
-            bounds (dict): Dictionary containing bounds for variables
+            i (int): Index of the previous subnet
+            bounds (dict): Dictionary containing bounds for variables of the previous subnet
 
         Returns:
             None
@@ -173,7 +174,7 @@ def encodeOutput(self, ipq, i):
         """Function to encode output variables
         Args:
             ipq:    (:class:`~maraboupy.MarabouCore.InputQuery`): InputQuery object to encode output variables
-            i:      (int): Index of the subnet
+            i:      (int): Index of the previous subnet
         
         Returns:
             None
@@ -297,4 +298,4 @@ def setUpperBound(self, x, v):
         if any(x in arr for arr in self.inputVars) or any(x in arr for arr in self.outputVars):
             self.upperBounds[x] = v
         else:
-            raise RuntimeError("Can set bounds only on either input or output variables")
\ No newline at end of file
+            raise RuntimeError("Can set bounds only on either input or output variables")
diff --git a/maraboupy/test/test_network_composition.py b/maraboupy/test/test_network_composition.py
index f4faf945c9..f7d49d3e3d 100644
--- a/maraboupy/test/test_network_composition.py
+++ b/maraboupy/test/test_network_composition.py
@@ -1,8 +1,5 @@
-# Tests MarabouNetwork features not tested by it's subclasses
 from maraboupy import Marabou
-from maraboupy import MarabouCore
 import os
-import numpy as np
 
 # Global settings
 OPT = Marabou.createOptions(verbosity = 0) # Turn off printing
@@ -118,4 +115,4 @@ def test_one_split_unsat():
     exitCode2, _, _ = network.calculateBounds(options=OPT)
 
     # exitCode2 should be also unsat
-    assert exitCode == exitCode2
\ No newline at end of file
+    assert exitCode == exitCode2

From e9daad8be308d47d85232b8186185b7e50c0576e Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Sun, 25 Feb 2024 21:11:54 -0800
Subject: [PATCH 10/12] indentation

---
 maraboupy/Marabou.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/maraboupy/Marabou.py b/maraboupy/Marabou.py
index 2d90251285..66e6bc648b 100644
--- a/maraboupy/Marabou.py
+++ b/maraboupy/Marabou.py
@@ -92,7 +92,7 @@ def read_onnx_with_threshold(filename, inputNames=None, outputNames=None, maxNum
         :class:`~maraboupy.MarabouNetworkComposition.MarabouNetworkComposition`
     """
     return MarabouNetworkComposition(filename, inputNames, outputNames,
-                                        maxNumberOfLinearEquations=maxNumberOfLinearEquations)
+                                     maxNumberOfLinearEquations=maxNumberOfLinearEquations)
 
 def load_query(filename):
     """Load the serialized inputQuery from the given filename

From 1ea36498f54732f79270f713f8c288912c68af52 Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Wed, 28 Feb 2024 22:33:46 -0800
Subject: [PATCH 11/12] address the problems after mergning ONNSParser

---
 maraboupy/MarabouNetworkONNX.py |  4 +---
 maraboupy/parsers/ONNXParser.py | 17 +++++++++++------
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/maraboupy/MarabouNetworkONNX.py b/maraboupy/MarabouNetworkONNX.py
index ecb0d99cb6..6cadd71280 100644
--- a/maraboupy/MarabouNetworkONNX.py
+++ b/maraboupy/MarabouNetworkONNX.py
@@ -55,8 +55,6 @@ def readONNX(self, filename, inputNames=None, outputNames=None, preserveExisting
 
         self.filename = filename
         self.graph = onnx.load(filename).graph
-        self.maxNumberOfLinearEquations = maxNumberOfLinearEquations
-        self.thresholdReached = False
 
         if os.path.exists('post_split.onnx'):
             os.remove('post_split.onnx')
@@ -90,7 +88,7 @@ def readONNX(self, filename, inputNames=None, outputNames=None, preserveExisting
             initNames = [node.name for node in self.graph.initializer]
             self.outputNames = [out.name for out in self.graph.output if out.name not in initNames]
 
-        ONNXParser.parse(self, self.graph, self.inputNames, self.outputNames)
+        ONNXParser.parse(self, self.graph, self.inputNames, self.outputNames, maxNumberOfLinearEquations=maxNumberOfLinearEquations)
 
     def getNode(self, nodeName):
         """Find the node in the graph corresponding to the given name
diff --git a/maraboupy/parsers/ONNXParser.py b/maraboupy/parsers/ONNXParser.py
index 0d11aee152..d2881e8556 100644
--- a/maraboupy/parsers/ONNXParser.py
+++ b/maraboupy/parsers/ONNXParser.py
@@ -32,7 +32,7 @@ class ONNXParser:
     """
 
     @staticmethod
-    def parse(query:InputQueryBuilder, graph, inputNames:List[str], outputNames:List[str]):
+    def parse(query:InputQueryBuilder, graph, inputNames:List[str], outputNames:List[str], maxNumberOfLinearEquations=None):
         """
         Parses the provided ONNX graph into constraints which are stored in the query argument.
 
@@ -45,11 +45,11 @@ def parse(query:InputQueryBuilder, graph, inputNames:List[str], outputNames:List
         Returns:
             :class:`~maraboupy.Marabou.marabouNetworkONNX.marabouNetworkONNX`
         """
-        parser = ONNXParser(query, graph, inputNames, outputNames)
+        parser = ONNXParser(query, graph, inputNames, outputNames, maxNumberOfLinearEquations=maxNumberOfLinearEquations)
         parser.parseGraph()
 
 
-    def __init__(self, query:InputQueryBuilder, graph, inputNames, outputNames):
+    def __init__(self, query:InputQueryBuilder, graph, inputNames, outputNames, maxNumberOfLinearEquations=None):
         """
         Should not be called directly. Use `ONNXParser.parse` instead.
 
@@ -66,6 +66,9 @@ def __init__(self, query:InputQueryBuilder, graph, inputNames, outputNames):
         self.constantMap = dict()
         self.shapeMap = dict()
 
+        self.maxNumberOfLinearEquations = maxNumberOfLinearEquations
+        self.thresholdReached = False
+
 
     def parseGraph(self):
         """Read an ONNX file and create a MarabouNetworkONNX object
@@ -84,7 +87,9 @@ def parseGraph(self):
         for outputName in self.outputNames:
             if outputName in self.constantMap:
                 raise RuntimeError("Output variable %s is a constant, not the output of equations!" % outputName)
-        self.query.outputVars.extend([self.varMap[outputName] for outputName in self.outputNames])
+            # If maxNumberOfLinearEquations is reached, the network is split and the outputVars are not set
+            if outputName in self.varMap:
+                self.query.outputVars.extend([self.varMap[outputName]])
 
     def processGraph(self):
         """Processes the ONNX graph to produce Marabou equations
@@ -147,8 +152,8 @@ def makeGraphEquations(self, nodeName, makeEquations):
         self.makeMarabouEquations(nodeName, makeEquations)
 
         if self.maxNumberOfLinearEquations is not None:
-            if not self.thresholdReached and len(self.equList) > self.maxNumberOfLinearEquations:
-                if self.splitNetworkAtNode(nodeName, networkNamePostSplit='post_split.onnx'):
+            if not self.thresholdReached and len(self.query.equList) > self.maxNumberOfLinearEquations:
+                if self.query.splitNetworkAtNode(nodeName, networkNamePostSplit='post_split.onnx'):
                     self.thresholdReached = True
 
         # Create new variables when we find one of the inputs

From 81a75cbf5991ae8b0b1cf9cb16066433b7d239bb Mon Sep 17 00:00:00 2001
From: tagomaru <tagomaru@users.noreply.github.com>
Date: Wed, 28 Feb 2024 22:45:31 -0800
Subject: [PATCH 12/12] fix the bug of extention of outputVars

---
 maraboupy/parsers/ONNXParser.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/maraboupy/parsers/ONNXParser.py b/maraboupy/parsers/ONNXParser.py
index d2881e8556..f66d35c133 100644
--- a/maraboupy/parsers/ONNXParser.py
+++ b/maraboupy/parsers/ONNXParser.py
@@ -87,6 +87,8 @@ def parseGraph(self):
         for outputName in self.outputNames:
             if outputName in self.constantMap:
                 raise RuntimeError("Output variable %s is a constant, not the output of equations!" % outputName)
+
+        for outputName in self.outputNames:
             # If maxNumberOfLinearEquations is reached, the network is split and the outputVars are not set
             if outputName in self.varMap:
                 self.query.outputVars.extend([self.varMap[outputName]])