Cleaning up integrated changes

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",

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):

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

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

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]

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

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 (

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():

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

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()
 

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,

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,