pepflow 0.1.4__py3-none-any.whl → 0.1.4a1__py3-none-any.whl

pepflow/pep.py

@@ -138,11 +138,19 @@ class PEPBuilder:
     def set_relaxed_constraints(self, relaxed_constraints: list[str]):
         self.relaxed_constraints.extend(relaxed_constraints)
 
-    def declare_func(self, function_class, **kwargs):
+    def declare_func(self, function_class: type[Function], tag: str, **kwargs):
         func = function_class(is_basis=True, composition=None, **kwargs)
+        func.add_tag(tag)
         self.functions.append(func)
         return func
 
+    def get_func_by_tag(self, tag: str):
+        # TODO: Add support to return composite functions as well. Right now we can only return base functions
+        for f in self.functions:
+            if tag in f.tags:
+                return f
+        raise ValueError("Cannot find the function of given tag.")
+
     def solve(self, context: pc.PEPContext | None = None, **kwargs):
         if context is None:
             context = pc.get_current_context()
@@ -152,7 +160,6 @@ class PEPBuilder:
         all_constraints: list[Constraint] = [*self.init_conditions]
         for f in self.functions:
             all_constraints.extend(f.get_interpolation_constraints())
-            all_constraints.extend(context.opt_conditions[f])
 
         # for now, we heavily rely on the CVX. We can make a wrapper class to avoid
         # direct dependency in the future.

pepflow/pep_context.py

@@ -26,7 +26,6 @@ import natsort
 import pandas as pd
 
 if TYPE_CHECKING:
-    from pepflow.constraint import Constraint
     from pepflow.function import Function, Triplet
     from pepflow.point import Point
     from pepflow.scalar import Scalar
@@ -53,7 +52,8 @@ class PEPContext:
         self.points: list[Point] = []
         self.scalars: list[Scalar] = []
         self.triplets: dict[Function, list[Triplet]] = defaultdict(list)
-        self.opt_conditions: dict[Function, list[Constraint]] = defaultdict(list)
+        # self.triplets will contain all stationary_triplets. They are not mutually exclusive.
+        self.stationary_triplets: dict[Function, list[Triplet]] = defaultdict(list)
 
         GLOBAL_CONTEXT_DICT[name] = self
 
@@ -70,8 +70,8 @@ class PEPContext:
     def add_triplet(self, function: Function, triplet: Triplet):
         self.triplets[function].append(triplet)
 
-    def add_opt_condition(self, function: Function, opt_condition: Constraint):
-        self.opt_conditions[function].append(opt_condition)
+    def add_stationary_triplet(self, function: Function, stationary_triplet: Triplet):
+        self.stationary_triplets[function].append(stationary_triplet)
 
     def get_by_tag(self, tag: str) -> Point | Scalar:
         for p in self.points:
@@ -80,13 +80,13 @@ class PEPContext:
         for s in self.scalars:
             if tag in s.tags:
                 return s
-        raise ValueError("Cannot find the point or scalar of given tag")
+        raise ValueError("Cannot find the point, scalar, or function of given tag.")
 
     def clear(self):
         self.points.clear()
         self.scalars.clear()
         self.triplets.clear()
-        self.opt_conditions.clear()
+        self.stationary_triplets.clear()
 
     def tracked_point(self, func: Function) -> list[Point]:
         return natsort.natsorted(

pepflow/pep_context_test.py

@@ -91,10 +91,14 @@ def test_get_by_tag(pep_context: pc.PEPContext):
     f = SmoothConvexFunction(L=1, is_basis=True)
     f.add_tag("f")
     p1 = Point(is_basis=True, tags=["x1"])
+    p2 = Point(is_basis=True, tags=["x2"])
+    p3 = p1 + p2
 
     triplet = f.generate_triplet(p1)
+    _ = f.generate_triplet(p2)
 
     assert pep_context.get_by_tag("x1") == p1
     assert pep_context.get_by_tag("f(x1)") == triplet.function_value
     assert pep_context.get_by_tag("gradient_f(x1)") == triplet.gradient
+    assert pep_context.get_by_tag("x1+x2") == p3
     pc.set_current_context(None)

pepflow/pep_test.py

@@ -19,6 +19,7 @@
 
 import pytest
 
+from pepflow import function as fc
 from pepflow import pep
 from pepflow import pep_context as pc
 
@@ -75,3 +76,10 @@ class TestPEPBuilder:
 
         with builder.make_context("test", override=True):
             pass
+
+    def test_get_func_by_tag(self) -> None:
+        builder = pep.PEPBuilder()
+        with builder.make_context("test"):
+            f = builder.declare_func(fc.SmoothConvexFunction, "f", L=1)
+
+            assert builder.get_func_by_tag("f") == f

pepflow/point_test.py

@@ -24,10 +24,9 @@ import numpy as np
 import pytest
 
 from pepflow import expression_manager as exm
-from pepflow import function as fc
 from pepflow import pep as pep
 from pepflow import pep_context as pc
-from pepflow import point, scalar
+from pepflow import point
 
 
 @pytest.fixture
@@ -103,12 +102,6 @@ def test_point_hash_different(pep_context: pc.PEPContext) -> None:
     assert p1.uid != p2.uid
 
 
-def test_scalar_hash_different(pep_context: pc.PEPContext) -> None:
-    s1 = scalar.Scalar(is_basis=True, eval_expression=None)
-    s2 = scalar.Scalar(is_basis=True, eval_expression=None)
-    assert s1.uid != s2.uid
-
-
 def test_point_tag(pep_context: pc.PEPContext) -> None:
     p1 = point.Point(is_basis=True, eval_expression=None)
     p1.add_tag(tag="my_tag")
@@ -123,202 +116,49 @@ def test_point_repr(pep_context: pc.PEPContext) -> None:
     assert str(p1) == "my_tag"
 
 
-def test_scalar_tag():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None)
-        s1.add_tag(tag="my_tag")
-    assert s1.tags == ["my_tag"]
-    assert s1.tag == "my_tag"
-
-
-def test_scalar_repr():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, tags=["s1"])
-        print(s1)  # it should be fine without tag
-        s1.add_tag("my_tag")
-        assert str(s1) == "my_tag"
-
-
-def test_point_in_a_list():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        p1 = point.Point(is_basis=True, eval_expression=None)
-        p2 = point.Point(is_basis=True, eval_expression=None)
-        p3 = point.Point(is_basis=True, eval_expression=None)
+def test_point_in_a_list(pep_context: pc.PEPContext):
+    p1 = point.Point(is_basis=True, eval_expression=None)
+    p2 = point.Point(is_basis=True, eval_expression=None)
+    p3 = point.Point(is_basis=True, eval_expression=None)
     assert p1 in [p1, p2]
     assert p3 not in [p1, p2]
 
 
-def test_scalar_in_a_list():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None)
-        s2 = scalar.Scalar(is_basis=True, eval_expression=None)
-        s3 = scalar.Scalar(is_basis=True, eval_expression=None)
-    assert s1 in [s1, s2]
-    assert s3 not in [s1, s2]
-
-
-def test_expression_manager_on_basis_point():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        p1 = point.Point(is_basis=True, eval_expression=None, tags=["p1"])
-        p2 = point.Point(is_basis=True, eval_expression=None, tags=["p2"])
-        pm = exm.ExpressionManager(ctx)
+def test_expression_manager_on_basis_point(pep_context: pc.PEPContext):
+    p1 = point.Point(is_basis=True, eval_expression=None, tags=["p1"])
+    p2 = point.Point(is_basis=True, eval_expression=None, tags=["p2"])
+    pm = exm.ExpressionManager(pep_context)
 
-        np.testing.assert_allclose(pm.eval_point(p1).vector, np.array([1, 0]))
-        np.testing.assert_allclose(pm.eval_point(p2).vector, np.array([0, 1]))
+    np.testing.assert_allclose(pm.eval_point(p1).vector, np.array([1, 0]))
+    np.testing.assert_allclose(pm.eval_point(p2).vector, np.array([0, 1]))
 
-        p3 = point.Point(is_basis=True, eval_expression=None, tags=["p3"])  # noqa: F841
-        pm = exm.ExpressionManager(ctx)
+    p3 = point.Point(is_basis=True, eval_expression=None, tags=["p3"])  # noqa: F841
+    pm = exm.ExpressionManager(pep_context)
 
     np.testing.assert_allclose(pm.eval_point(p1).vector, np.array([1, 0, 0]))
     np.testing.assert_allclose(pm.eval_point(p2).vector, np.array([0, 1, 0]))
 
 
-def test_expression_manager_on_basis_scalar():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s1"])
-        s2 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s2"])
-        pm = exm.ExpressionManager(ctx)
-
-        np.testing.assert_allclose(pm.eval_scalar(s1).vector, np.array([1, 0]))
-        np.testing.assert_allclose(pm.eval_scalar(s2).vector, np.array([0, 1]))
-
-        s3 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s3"])  # noqa: F841
-        pm = exm.ExpressionManager(ctx)
+def test_expression_manager_eval_point(pep_context: pc.PEPContext):
+    p1 = point.Point(is_basis=True, tags=["p1"])
+    p2 = point.Point(is_basis=True, tags=["p2"])
+    p3 = 2 * p1 + p2 / 4
+    p4 = p3 + p1
 
-        np.testing.assert_allclose(pm.eval_scalar(s1).vector, np.array([1, 0, 0]))
-        np.testing.assert_allclose(pm.eval_scalar(s2).vector, np.array([0, 1, 0]))
-
-
-def test_expression_manager_eval_point():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        p1 = point.Point(is_basis=True, tags=["p1"])
-        p2 = point.Point(is_basis=True, tags=["p2"])
-        p3 = 2 * p1 + p2 / 4
-        p4 = p3 + p1
-
-        pm = exm.ExpressionManager(ctx)
+    pm = exm.ExpressionManager(pep_context)
     np.testing.assert_allclose(pm.eval_point(p3).vector, np.array([2, 0.25]))
     np.testing.assert_allclose(pm.eval_point(p4).vector, np.array([3, 0.25]))
 
 
-def test_expression_manager_eval_point_large_scale():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        all_basis = [point.Point(is_basis=True, tags=[f"p_{i}"]) for i in range(100)]
-        p = all_basis[0]
-        for i in range(len(all_basis)):
-            for j in range(i + 1, len(all_basis)):
-                p += all_basis[i] * 2 + all_basis[j]
-        pm = exm.ExpressionManager(ctx)
-        t = time.time()
-        for pp in ctx.points:
-            pm.eval_point(pp)
-
-        assert (time.time() - t) < 0.5
-
-
-def test_function_generate_triplet():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        f = fc.Function(is_basis=True, reuse_gradient=True)
-        f.add_tag("f")
-        g = fc.Function(is_basis=True, reuse_gradient=True)
-        g.add_tag("g")
-        h = 5 * f + 5 * g
-        h.add_tag("h")
-
-        f1 = fc.Function(is_basis=True, reuse_gradient=False)
-        f1.add_tag("f1")
-        g1 = fc.Function(is_basis=True, reuse_gradient=False)
-        g1.add_tag("g1")
-        h1 = 5 * f1 + 5 * g1
-        h1.add_tag("h1")
-
-        p1 = point.Point(is_basis=True)
-        p1.add_tag("p1")
-        p1_triplet = h.generate_triplet(p1)
-        p1_triplet_1 = h.generate_triplet(p1)
-
-        p1_triplet_2 = h1.generate_triplet(p1)
-        p1_triplet_3 = h1.generate_triplet(p1)
-
-        pm = exm.ExpressionManager(ctx)
-
-    np.testing.assert_allclose(
-        pm.eval_point(p1).vector, np.array([1, 0, 0, 0, 0, 0, 0])
-    )
-
-    np.testing.assert_allclose(
-        pm.eval_point(p1_triplet.gradient).vector, np.array([0, 5, 5, 0, 0, 0, 0])
-    )
-    np.testing.assert_allclose(
-        pm.eval_scalar(p1_triplet.function_value).vector, np.array([5, 5, 0, 0])
-    )
-
-    np.testing.assert_allclose(
-        pm.eval_point(p1_triplet_1.gradient).vector, np.array([0, 5, 5, 0, 0, 0, 0])
-    )
-    np.testing.assert_allclose(
-        pm.eval_scalar(p1_triplet_1.function_value).vector, np.array([5, 5, 0, 0])
-    )
-
-    np.testing.assert_allclose(
-        pm.eval_point(p1_triplet_2.gradient).vector, np.array([0, 0, 0, 5, 5, 0, 0])
-    )
-    np.testing.assert_allclose(
-        pm.eval_scalar(p1_triplet_2.function_value).vector, np.array([0, 0, 5, 5])
-    )
-
-    np.testing.assert_allclose(
-        pm.eval_point(p1_triplet_3.gradient).vector, np.array([0, 0, 0, 0, 0, 5, 5])
-    )
-    np.testing.assert_allclose(
-        pm.eval_scalar(p1_triplet_3.function_value).vector, np.array([0, 0, 5, 5])
-    )
-
-
-def test_function_add_stationary_point():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        f = fc.Function(is_basis=True, reuse_gradient=True)
-        f.add_tag("f")
-        x_opt = f.add_stationary_point("x_opt")
-
-        pm = exm.ExpressionManager(ctx)
-
-        np.testing.assert_allclose(pm.eval_point(x_opt).vector, np.array([1]))
-
-
-def test_smooth_interpolability_constraints():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        f = fc.SmoothConvexFunction(L=1)
-        f.add_tag("f")
-        _ = f.add_stationary_point("x_opt")
-
-        x_0 = point.Point(is_basis=True)
-        x_0.add_tag("x_0")
-        _ = f.generate_triplet(x_0)
-
-        all_interpolation_constraints = f.get_interpolation_constraints()
-
-        pm = exm.ExpressionManager(ctx)
-
-    np.testing.assert_allclose(
-        pm.eval_scalar(all_interpolation_constraints[1].scalar).vector, [1, -1]
-    )
-    np.testing.assert_allclose(
-        pm.eval_scalar(all_interpolation_constraints[1].scalar).matrix,
-        [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.5]],
-    )
-
-    np.testing.assert_allclose(
-        pm.eval_scalar(all_interpolation_constraints[1].scalar).constant, 0
-    )
+def test_expression_manager_eval_point_large_scale(pep_context):
+    all_basis = [point.Point(is_basis=True, tags=[f"p_{i}"]) for i in range(100)]
+    p = all_basis[0]
+    for i in range(len(all_basis)):
+        for j in range(i + 1, len(all_basis)):
+            p += all_basis[i] * 2 + all_basis[j]
+    pm = exm.ExpressionManager(pep_context)
+    t = time.time()
+    for pp in pep_context.points:
+        pm.eval_point(pp)
+
+    assert (time.time() - t) < 0.5

pepflow/scalar_test.py

@@ -17,155 +17,148 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from typing import Iterator
 
 import numpy as np
+import pytest
 
 from pepflow import expression_manager as exm
 from pepflow import pep as pep
-from pepflow import point, scalar, utils
+from pepflow import pep_context as pc
+from pepflow import point, scalar
 
 
-def test_scalar_add_tag():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s1"])
-        s2 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s2"])
+@pytest.fixture
+def pep_context() -> Iterator[pc.PEPContext]:
+    """Prepare the pep context and reset the context to None at the end."""
+    ctx = pc.PEPContext("test").set_as_current()
+    yield ctx
+    pc.set_current_context(None)
 
-        s_add = s1 + s2
-        assert s_add.tag == "s1+s2"
 
-        s_add = s1 + 0.1
-        assert s_add.tag == "s1+0.1"
+def test_scalar_add_tag(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s1"])
+    s2 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s2"])
 
-        s_radd = 0.1 + s1
-        assert s_radd.tag == "0.1+s1"
+    s_add = s1 + s2
+    assert s_add.tag == "s1+s2"
 
-        s_sub = s1 - s2
-        assert s_sub.tag == "s1-s2"
+    s_add = s1 + 0.1
+    assert s_add.tag == "s1+0.1"
 
-        s_sub = s1 - (s2 + s1)
-        assert s_sub.tag == "s1-(s2+s1)"
+    s_radd = 0.1 + s1
+    assert s_radd.tag == "0.1+s1"
 
-        s_sub = s1 - (s2 - s1)
-        assert s_sub.tag == "s1-(s2-s1)"
+    s_sub = s1 - s2
+    assert s_sub.tag == "s1-s2"
 
-        s_sub = s1 - 0.1
-        assert s_sub.tag == "s1-0.1"
+    s_sub = s1 - (s2 + s1)
+    assert s_sub.tag == "s1-(s2+s1)"
 
-        s_rsub = 0.1 - s1
-        assert s_rsub.tag == "0.1-s1"
+    s_sub = s1 - (s2 - s1)
+    assert s_sub.tag == "s1-(s2-s1)"
 
+    s_sub = s1 - 0.1
+    assert s_sub.tag == "s1-0.1"
 
-def test_scalar_mul_tag():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s"])
+    s_rsub = 0.1 - s1
+    assert s_rsub.tag == "0.1-s1"
 
-        s_mul = s * 0.1
-        assert s_mul.tag == "s*0.1"
 
-        s_rmul = 0.1 * s
-        assert s_rmul.tag == "0.1*s"
+def test_scalar_mul_tag(pep_context: pc.PEPContext):
+    s = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s"])
 
-        s_neg = -s
-        assert s_neg.tag == "-s"
+    s_mul = s * 0.1
+    assert s_mul.tag == "s*0.1"
 
-        s_truediv = s / 0.1
-        assert s_truediv.tag == "1/0.1*s"
+    s_rmul = 0.1 * s
+    assert s_rmul.tag == "0.1*s"
 
+    s_neg = -s
+    assert s_neg.tag == "-s"
 
-def test_scalar_add_and_mul_tag():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s1"])
-        s2 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s2"])
+    s_truediv = s / 0.1
+    assert s_truediv.tag == "1/0.1*s"
 
-        s_add_mul = (s1 + s2) * 0.1
-        assert s_add_mul.tag == "(s1+s2)*0.1"
 
-        s_add_mul = s1 + s2 * 0.1
-        assert s_add_mul.tag == "s1+s2*0.1"
+def test_scalar_add_and_mul_tag(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s1"])
+    s2 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s2"])
 
-        s_neg_add = -(s1 + s2)
-        assert s_neg_add.tag == "-(s1+s2)"
+    s_add_mul = (s1 + s2) * 0.1
+    assert s_add_mul.tag == "(s1+s2)*0.1"
 
-        s_rmul_add = 0.1 * (s1 + s2)
-        assert s_rmul_add.tag == "0.1*(s1+s2)"
+    s_add_mul = s1 + s2 * 0.1
+    assert s_add_mul.tag == "s1+s2*0.1"
 
+    s_neg_add = -(s1 + s2)
+    assert s_neg_add.tag == "-(s1+s2)"
 
-def test_scalar_hash_different():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None)
-        s2 = scalar.Scalar(is_basis=True, eval_expression=None)
+    s_rmul_add = 0.1 * (s1 + s2)
+    assert s_rmul_add.tag == "0.1*(s1+s2)"
+
+
+def test_scalar_hash_different(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, eval_expression=None)
+    s2 = scalar.Scalar(is_basis=True, eval_expression=None)
     assert s1.uid != s2.uid
 
 
-def test_scalar_tag():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None)
-        s1.add_tag(tag="my_tag")
+def test_scalar_tag(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, eval_expression=None)
+    s1.add_tag(tag="my_tag")
     assert s1.tags == ["my_tag"]
     assert s1.tag == "my_tag"
 
 
-def test_scalar_repr():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, tags=["s1"])
-        print(s1)  # it should be fine without tag
-        s1.add_tag("my_tag")
-        assert str(s1) == "my_tag"
+def test_scalar_repr(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, tags=["s1"])
+    print(s1)  # it should be fine without tag
+    s1.add_tag("my_tag")
+    assert str(s1) == "my_tag"
 
 
-def test_scalar_in_a_list():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test"):
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None)
-        s2 = scalar.Scalar(is_basis=True, eval_expression=None)
-        s3 = scalar.Scalar(is_basis=True, eval_expression=None)
+def test_scalar_in_a_list(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, eval_expression=None)
+    s2 = scalar.Scalar(is_basis=True, eval_expression=None)
+    s3 = scalar.Scalar(is_basis=True, eval_expression=None)
     assert s1 in [s1, s2]
     assert s3 not in [s1, s2]
 
 
-def test_expression_manager_on_basis_scalar():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        s1 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s1"])
-        s2 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s2"])
-        pm = exm.ExpressionManager(ctx)
+def test_expression_manager_on_basis_scalar(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s1"])
+    s2 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s2"])
+    pm = exm.ExpressionManager(pep_context)
 
-        np.testing.assert_allclose(pm.eval_scalar(s1).vector, np.array([1, 0]))
-        np.testing.assert_allclose(pm.eval_scalar(s2).vector, np.array([0, 1]))
+    np.testing.assert_allclose(pm.eval_scalar(s1).vector, np.array([1, 0]))
+    np.testing.assert_allclose(pm.eval_scalar(s2).vector, np.array([0, 1]))
 
-        s3 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s3"])  # noqa: F841
-        pm = exm.ExpressionManager(ctx)
+    s3 = scalar.Scalar(is_basis=True, eval_expression=None, tags=["s3"])  # noqa: F841
+    pm = exm.ExpressionManager(pep_context)
 
-        np.testing.assert_allclose(pm.eval_scalar(s1).vector, np.array([1, 0, 0]))
-        np.testing.assert_allclose(pm.eval_scalar(s2).vector, np.array([0, 1, 0]))
+    np.testing.assert_allclose(pm.eval_scalar(s1).vector, np.array([1, 0, 0]))
+    np.testing.assert_allclose(pm.eval_scalar(s2).vector, np.array([0, 1, 0]))
 
 
-def test_expression_manager_eval_scalar():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        s1 = scalar.Scalar(is_basis=True, tags=["s1"])
-        s2 = scalar.Scalar(is_basis=True, tags=["s2"])
-        s3 = 2 * s1 + s2 / 4 + 5
-        s4 = s3 + s1
-        s5 = s4 + 5
+def test_expression_manager_eval_scalar(pep_context: pc.PEPContext):
+    s1 = scalar.Scalar(is_basis=True, tags=["s1"])
+    s2 = scalar.Scalar(is_basis=True, tags=["s2"])
+    s3 = 2 * s1 + s2 / 4 + 5
+    s4 = s3 + s1
+    s5 = s4 + 5
 
-        p1 = point.Point(is_basis=True, tags=["p1"])
-        p2 = point.Point(is_basis=True, tags=["p2"])
-        s6 = p1 * p2
+    p1 = point.Point(is_basis=True, tags=["p1"])
+    p2 = point.Point(is_basis=True, tags=["p2"])
+    s6 = p1 * p2
 
-        p3 = point.Point(is_basis=True, tags=["p3"])
-        p4 = point.Point(is_basis=True, tags=["p4"])
-        s7 = 5 * p3 * p4
+    p3 = point.Point(is_basis=True, tags=["p3"])
+    p4 = point.Point(is_basis=True, tags=["p4"])
+    s7 = 5 * p3 * p4
 
-        s8 = s6 + s7
+    s8 = s6 + s7
 
-        pm = exm.ExpressionManager(ctx)
+    pm = exm.ExpressionManager(pep_context)
 
     np.testing.assert_allclose(pm.eval_scalar(s3).vector, np.array([2, 0.25]))
     np.testing.assert_allclose(pm.eval_scalar(s3).constant, 5)
@@ -212,39 +205,3 @@ def test_expression_manager_eval_scalar():
             ]
         ),
     )
-
-
-def test_constraint():
-    pep_builder = pep.PEPBuilder()
-    with pep_builder.make_context("test") as ctx:
-        s1 = scalar.Scalar(is_basis=True, tags=["s1"])
-        s2 = scalar.Scalar(is_basis=True, tags=["s2"])
-        s3 = 2 * s1 + s2 / 4 + 5
-
-        c1 = s3.le(5, name="c1")
-        c2 = s3.lt(5, name="c2")
-        c3 = s3.ge(5, name="c3")
-        c4 = s3.gt(5, name="c4")
-        c5 = s3.eq(5, name="c5")
-
-        pm = exm.ExpressionManager(ctx)
-
-    np.testing.assert_allclose(pm.eval_scalar(c1.scalar).vector, np.array([2, 0.25]))
-    np.testing.assert_allclose(pm.eval_scalar(c1.scalar).constant, 0)
-    assert c1.comparator == utils.Comparator.LT
-
-    np.testing.assert_allclose(pm.eval_scalar(c2.scalar).vector, np.array([2, 0.25]))
-    np.testing.assert_allclose(pm.eval_scalar(c2.scalar).constant, 0)
-    assert c2.comparator == utils.Comparator.LT
-
-    np.testing.assert_allclose(pm.eval_scalar(c3.scalar).vector, np.array([2, 0.25]))
-    np.testing.assert_allclose(pm.eval_scalar(c3.scalar).constant, 0)
-    assert c3.comparator == utils.Comparator.GT
-
-    np.testing.assert_allclose(pm.eval_scalar(c4.scalar).vector, np.array([2, 0.25]))
-    np.testing.assert_allclose(pm.eval_scalar(c4.scalar).constant, 0)
-    assert c4.comparator == utils.Comparator.GT
-
-    np.testing.assert_allclose(pm.eval_scalar(c5.scalar).vector, np.array([2, 0.25]))
-    np.testing.assert_allclose(pm.eval_scalar(c5.scalar).constant, 0)
-    assert c5.comparator == utils.Comparator.EQ