From 877aac1851bdf884c3c1e2e1b21df913d10f0a3e Mon Sep 17 00:00:00 2001
From: Jeremy Sadler <53983960+jezsadler@users.noreply.github.com>
Date: Thu, 3 Oct 2024 01:14:21 +0000
Subject: [PATCH] Cleaning up integrated changes

---
 src/omlt/base/__init__.py                  |  5 +----
 src/omlt/base/constraint.py                |  4 ++--
 src/omlt/base/pyomo.py                     | 14 ++++++++++++--
 src/omlt/block.py                          |  3 +--
 src/omlt/formulation.py                    | 10 +++++++---
 src/omlt/linear_tree/lt_formulation.py     |  6 ++----
 src/omlt/neuralnet/activations/__init__.py |  2 --
 tests/base/test_block.py                   |  5 ++---
 tests/base/test_expression.py              |  3 ++-
 tests/base/test_var.py                     | 12 ++++++++----
 tests/linear_tree/test_lt_formulation.py   |  2 --
 tests/neuralnet/test_onnx.py               |  1 -
 12 files changed, 37 insertions(+), 30 deletions(-)

diff --git a/src/omlt/base/__init__.py b/src/omlt/base/__init__.py
index 95b95fce..53f0ccc4 100644
--- a/src/omlt/base/__init__.py
+++ b/src/omlt/base/__init__.py
@@ -1,8 +1,4 @@
 from omlt.base.language import DEFAULT_MODELING_LANGUAGE
-from omlt.dependencies import julia_available
-
-# if julia_available:
-#     from omlt.base.julia import jl, jump
 
 from omlt.base.constraint import (
     OmltConstraint,
@@ -19,6 +15,7 @@
     OmltScalarPyomo,
 )
 from omlt.base.var import OmltIndexed, OmltScalar, OmltVar, OmltVarFactory
+from omlt.dependencies import julia_available
 
 __all__ = [
     "DEFAULT_MODELING_LANGUAGE",
diff --git a/src/omlt/base/constraint.py b/src/omlt/base/constraint.py
index cd881854..6a3bbd75 100644
--- a/src/omlt/base/constraint.py
+++ b/src/omlt/base/constraint.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from typing import Any
 
 import pyomo.environ as pyo
@@ -9,7 +9,7 @@
 from omlt.base import DEFAULT_MODELING_LANGUAGE
 
 
-class OmltConstraint(ABC):
+class OmltConstraint:
 
     @property
     def ctype(self):
diff --git a/src/omlt/base/pyomo.py b/src/omlt/base/pyomo.py
index 5a12d9d6..c3d9a095 100644
--- a/src/omlt/base/pyomo.py
+++ b/src/omlt/base/pyomo.py
@@ -98,6 +98,7 @@ def __init__(self, *indexes: Any, **kwargs: Any):
         self._pyovar = pyo.Var(*indexes, **kwargs)
         self._name = None
         self._parent = None
+        self._pyovar._parent = None
         self._constructed = self._pyovar._constructed
         self._index_set = self._pyovar._index_set
         self._rule_init = self._pyovar._rule_init
@@ -106,7 +107,6 @@ def __init__(self, *indexes: Any, **kwargs: Any):
         self._dense = self._pyovar._dense
         self._data = self._pyovar._data
         self._units = self._pyovar._units
-        self._implicit_subsets = self._pyovar._implicit_subsets
         self.doc = self._pyovar.doc
         self._ctype = pyo.Var
         self.bounds = (None, None)
@@ -166,6 +166,16 @@ def setlb(self, value):
         for vardata in self.values():
             vardata.lb = value
 
+    @property
+    def _parent(self):
+        return self._pyovar._parent
+
+    @_parent.setter
+    def _parent(self, value):
+        self._pyovar._parent = value
+        if self.is_constructed():
+            for idx in self.keys():
+                self[idx]._parent = value
 
 # Constraints
 
@@ -289,7 +299,7 @@ def doc(self):
 class OmltExprScalarPyomo(OmltExpr, pyo.Expression):
     format = "pyomo"
 
-    def __init__(self, expr=None, **kwargs):
+    def __init__(self, expr=None):
         self._index_set = {}
         if isinstance(expr, OmltExprScalarPyomo):
             self._expression = expr._expression
diff --git a/src/omlt/block.py b/src/omlt/block.py
index 217e9338..d57ec62d 100644
--- a/src/omlt/block.py
+++ b/src/omlt/block.py
@@ -24,7 +24,6 @@ class is used in combination with a formulation object to construct the
         pyo.assert_optimal_termination(status)
 """
 
-import warnings
 
 import pyomo.environ as pyo
 from pyomo.core.base.block import BlockData, declare_custom_block
@@ -109,7 +108,7 @@ def build_formulation(self, formulation, lang=None):
 
 
 @declare_custom_block(name="OmltBlock")
-class OmltBlockData(_BlockData, OmltBlockCore):
+class OmltBlockData(BlockData, OmltBlockCore):
     def __init__(self, component):
         super().__init__(component)
         self.__formulation = None
diff --git a/src/omlt/formulation.py b/src/omlt/formulation.py
index 12112f64..e6f4416a 100644
--- a/src/omlt/formulation.py
+++ b/src/omlt/formulation.py
@@ -72,8 +72,10 @@ def _set_block(self, block):
 
     @property
     def block(self):
-        """The underlying block containing the constraints / variables for this
-        formulation."""
+        """Block.
+
+        The underlying block containing the constraints/variables for this formulation.
+        """
         return self.__block()
 
 
@@ -134,7 +136,9 @@ def _setup_scaled_inputs_outputs(block, scaler=None, scaled_input_bounds=None):
         )
     constraint_factory = OmltConstraintFactory()
 
-    block._scale_input_constraint = constraint_factory.new_constraint(block.inputs_set, lang=block._format)
+    block._scale_input_constraint = constraint_factory.new_constraint(
+        block.inputs_set, lang=block._format
+    )
     for idx in block.inputs_set:
         block._scale_input_constraint[idx] = (
             block.scaled_inputs[idx] == input_scaling_expressions[idx]
diff --git a/src/omlt/linear_tree/lt_formulation.py b/src/omlt/linear_tree/lt_formulation.py
index 115d67d1..904b36f6 100644
--- a/src/omlt/linear_tree/lt_formulation.py
+++ b/src/omlt/linear_tree/lt_formulation.py
@@ -2,7 +2,6 @@
 import pyomo.environ as pe
 from pyomo.gdp import Disjunct
 
-from omlt.base import OmltConstraintFactory, OmltVarFactory
 from omlt.formulation import _PyomoFormulation, _setup_scaled_inputs_outputs
 
 
@@ -267,9 +266,8 @@ def _add_gdp_formulation_to_block(  # noqa: PLR0913
     block.scaled_outputs.setub(output_bounds[1])
     block.scaled_outputs.setlb(output_bounds[0])
 
-    var_factory = OmltVarFactory()
-    block.intermediate_output = var_factory.new_var(
-        tree_ids, lang=block._format, bounds=(output_bounds[0], output_bounds[1])
+    block.intermediate_output = pe.Var(
+        tree_ids, bounds=(output_bounds[0], output_bounds[1])
     )
 
     # Create a disjunct for each leaf containing the bound constraints
diff --git a/src/omlt/neuralnet/activations/__init__.py b/src/omlt/neuralnet/activations/__init__.py
index 07e1e7b9..740022ad 100644
--- a/src/omlt/neuralnet/activations/__init__.py
+++ b/src/omlt/neuralnet/activations/__init__.py
@@ -7,8 +7,6 @@
 """
 from typing import Any
 
-from typing import Any
-
 from .linear import linear_activation_constraint, linear_activation_function
 from .relu import ComplementarityReLUActivation, bigm_relu_activation_constraint
 from .smooth import (
diff --git a/tests/base/test_block.py b/tests/base/test_block.py
index f439bbe3..e7de62bd 100644
--- a/tests/base/test_block.py
+++ b/tests/base/test_block.py
@@ -2,7 +2,6 @@
 import pytest
 from omlt import OmltBlock
 
-
 INPUTS_LENGTH = 3
 OUTPUTS_LENGTH = 2
 
@@ -47,8 +46,8 @@ def test_block():
     m.b.build_formulation(formulation, lang="pyomo")
 
     assert m.b._OmltBlockCore__formulation is formulation
-    assert [k for k in m.b.inputs] == ["A", "C", "D"]
-    assert [k for k in m.b.outputs] == [(0, 0), (0, 1), (1, 0), (1, 1)]
+    assert list(m.b.inputs) == ["A", "C", "D"]
+    assert list(m.b.outputs) == [(0, 0), (0, 1), (1, 0), (1, 1)]
 
 
 def test_input_output_auto_creation():
diff --git a/tests/base/test_expression.py b/tests/base/test_expression.py
index 34e1d3e7..1d33ed9d 100644
--- a/tests/base/test_expression.py
+++ b/tests/base/test_expression.py
@@ -5,6 +5,7 @@
 VAR1_VALUE = 6
 VAR2_VALUE = 3
 CONST_VALUE = 4
+NUM_ARGS = 2
 
 var_factory = OmltVarFactory()
 expr_factory = OmltExprFactory()
@@ -71,7 +72,7 @@ def test_init_scalar_expression():
     assert e2.is_potentially_variable()
     assert not e2.is_indexed()
 
-    assert e2.nargs() == 2
+    assert e2.nargs() == NUM_ARGS
     assert e2.args[1] == CONST_VALUE
     assert e2.arg(1) == CONST_VALUE
     assert len(e2) == 1
diff --git a/tests/base/test_var.py b/tests/base/test_var.py
index cf9a9fa3..0a4f06d5 100644
--- a/tests/base/test_var.py
+++ b/tests/base/test_var.py
@@ -3,6 +3,10 @@
 from omlt.base import OmltVarFactory
 from omlt.dependencies import julia_available
 
+VAR_VALUE = 3
+FIX_VALUE = 2
+UPPER_BOUND = 5
+
 var_factory = OmltVarFactory()
 
 
@@ -20,12 +24,12 @@ def _test_scalar_var(lang):
     assert v.is_constructed()
 
     v.value = 3
-    assert v.value == 3
+    assert v.value == VAR_VALUE
 
     v.fix(2, skip_validation=True)
     v.bounds = (0, 5)
     assert v.lb == 0
-    assert v.ub == 5
+    assert v.ub == UPPER_BOUND
     v.lb = 1
     v.ub = 3
     assert v.bounds == (1, 3)
@@ -62,11 +66,11 @@ def _test_indexed_var(lang):
     assert v.is_constructed()
 
     v.value = 3
-    assert v.value == 3
+    assert v.value == VAR_VALUE
 
     v.fix(2, skip_validation=True)
     for e in v:
-        assert v[e].value == 2
+        assert v[e].value == FIX_VALUE
 
     v.fix()
 
diff --git a/tests/linear_tree/test_lt_formulation.py b/tests/linear_tree/test_lt_formulation.py
index e9959133..a3082c01 100644
--- a/tests/linear_tree/test_lt_formulation.py
+++ b/tests/linear_tree/test_lt_formulation.py
@@ -5,8 +5,6 @@
 
 if lineartree_available:
     from lineartree import LinearTreeRegressor
-    from sklearn.linear_model import LinearRegression
-
     from omlt.linear_tree import (
         LinearTreeDefinition,
         LinearTreeGDPFormulation,
diff --git a/tests/neuralnet/test_onnx.py b/tests/neuralnet/test_onnx.py
index b2ee298e..7d33675f 100644
--- a/tests/neuralnet/test_onnx.py
+++ b/tests/neuralnet/test_onnx.py
@@ -7,7 +7,6 @@
 
 if onnx_available:
     import onnxruntime as ort
-
     from omlt.io.onnx import (
         load_onnx_neural_network,
         load_onnx_neural_network_with_bounds,