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

pepflow/__init__.py

@@ -35,6 +35,7 @@ from .pep_context import set_current_context as set_current_context
 # Function, Point, Scalar
 from .function import Function as Function
 from .function import SmoothConvexFunction as SmoothConvexFunction
+from .function import ConvexFunction as ConvexFunction
 from .function import Triplet as Triplet
 from .point import EvaluatedPoint as EvaluatedPoint
 from .point import Point as Point

pepflow/constraint_test.py

@@ -0,0 +1,71 @@
+# Copyright: 2025 The PEPFlow Developers
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# 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 pep_context as pc
+from pepflow import scalar, utils
+
+
+@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)
+
+
+def test_constraint(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
+
+    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(pep_context)
+
+    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

pepflow/e2e_test.py

@@ -10,8 +10,7 @@ def test_gd_e2e():
     eta = 1
     N = 9
 
-    f = pep_builder.declare_func(function.SmoothConvexFunction, L=1)
-    f.add_tag("f")
+    f = pep_builder.declare_func(function.SmoothConvexFunction, "f", L=1)
     x = pep_builder.set_init_point("x_0")
     x_star = f.add_stationary_point("x_star")
     pep_builder.set_initial_constraint(
@@ -32,3 +31,39 @@ def test_gd_e2e():
         result = pep_builder.solve()
         expected_opt_value = 1 / (4 * i + 2)
         assert math.isclose(result.primal_opt_value, expected_opt_value, rel_tol=1e-3)
+
+
+def test_pgm_e2e():
+    ctx = pc.PEPContext("pgm").set_as_current()
+    pep_builder = pep.PEPBuilder()
+    eta = 1
+    N = 1
+
+    f = pep_builder.declare_func(function.SmoothConvexFunction, "f", L=1)
+    g = pep_builder.declare_func(function.ConvexFunction, "g")
+
+    h = f + g
+
+    x = pep_builder.set_init_point("x_0")
+    x_star = h.add_stationary_point("x_star")
+    pep_builder.set_initial_constraint(
+        ((x - x_star) ** 2).le(1, name="initial_condition")
+    )
+
+    # We first build the algorithm with the largest number of iterations.
+    for i in range(N):
+        y = x - eta * f.gradient(x)
+        y.add_tag(f"y_{i + 1}")
+        x = g.proximal_step(y, eta)
+        x.add_tag(f"x_{i + 1}")
+
+    # To achieve the sweep, we can just update the performance_metric.
+    for i in range(1, N + 1):
+        p = ctx.get_by_tag(f"x_{i}")
+        pep_builder.set_performance_metric(
+            h.function_value(p) - h.function_value(x_star)
+        )
+
+        result = pep_builder.solve()
+        expected_opt_value = 1 / (4 * i)
+        assert math.isclose(result.primal_opt_value, expected_opt_value, rel_tol=1e-3)

pepflow/expression_manager.py

@@ -18,6 +18,7 @@
 # under the License.
 
 import functools
+import math
 
 import numpy as np
 
@@ -27,6 +28,17 @@ from pepflow import scalar as sc
 from pepflow import utils
 
 
+def tag_and_coef_to_str(tag: str, v: float) -> str:
+    coef = f"{abs(v):.3g}"
+    sign = "+" if v >= 0 else "-"
+    if math.isclose(abs(v), 1):
+        return f"{sign} {tag} "
+    elif math.isclose(v, 0):
+        return ""
+    else:
+        return f"{sign} {coef}*{tag} "
+
+
 class ExpressionManager:
     def __init__(self, pep_context: pc.PEPContext):
         self.context = pep_context
@@ -48,12 +60,18 @@ class ExpressionManager:
         self._num_basis_points = len(self._basis_points)
         self._num_basis_scalars = len(self._basis_scalars)
 
-    def get_index_of_basis_point(self, point: pt.Point):
+    def get_index_of_basis_point(self, point: pt.Point) -> int:
         return self._basis_point_uid_to_index[point.uid]
 
-    def get_index_of_basis_scalar(self, scalar: sc.Scalar):
+    def get_index_of_basis_scalar(self, scalar: sc.Scalar) -> int:
         return self._basis_scalar_uid_to_index[scalar.uid]
 
+    def get_tag_of_basis_point_index(self, index: int) -> str:
+        return self._basis_points[index].tag
+
+    def get_tag_of_basis_scalar_index(self, index: int) -> str:
+        return self._basis_scalars[index].tag
+
     @functools.cache
     def eval_point(self, point: pt.Point | float | int):
         if utils.is_numerical(point):
@@ -130,3 +148,55 @@ class ExpressionManager:
             ) / self.eval_scalar(scalar.eval_expression.right_scalar)
 
         raise ValueError("This should never happen!")
+
+    @functools.cache
+    def repr_point_by_basis(self, point: pt.Point) -> str:
+        assert isinstance(point, pt.Point)
+        repr_array = self.eval_point(point).vector
+
+        repr_str = ""
+        for i, v in enumerate(repr_array):
+            ith_tag = self.get_tag_of_basis_point_index(i)
+            repr_str += tag_and_coef_to_str(ith_tag, v)
+
+        # Post processing
+        if repr_str == "":
+            return "0"
+        if repr_str.startswith("+ "):
+            repr_str = repr_str[2:]
+        if repr_str.startswith("- "):
+            repr_str = "-" + repr_str[2:]
+        return repr_str.strip()
+
+    @functools.cache
+    def repr_scalar_by_basis(self, scalar: sc.Scalar) -> str:
+        assert isinstance(scalar, sc.Scalar)
+        evaluated_scalar = self.eval_scalar(scalar)
+
+        repr_str = ""
+        if not math.isclose(evaluated_scalar.constant, 0):
+            repr_str += f"{evaluated_scalar.constant:.3g}"
+
+        for i, v in enumerate(evaluated_scalar.vector):
+            # Note the tag is from scalar basis.
+            ith_tag = self.get_tag_of_basis_scalar_index(i)
+            repr_str += tag_and_coef_to_str(ith_tag, v)
+
+        for i in range(evaluated_scalar.matrix.shape[0]):
+            for j in range(i, evaluated_scalar.matrix.shape[0]):
+                ith_tag = self.get_tag_of_basis_point_index(i)
+                v = evaluated_scalar.matrix[i, j]
+                if i == j:
+                    repr_str += tag_and_coef_to_str(f"|{ith_tag}|^2", v)
+                    continue
+                jth_tag = self.get_tag_of_basis_point_index(j)
+                repr_str += tag_and_coef_to_str(f"<{ith_tag}, {jth_tag}>", 2 * v)
+
+        # Post processing
+        if repr_str == "":
+            return "0"
+        if repr_str.startswith("+ "):
+            repr_str = repr_str[2:]
+        if repr_str.startswith("- "):
+            repr_str = "-" + repr_str[2:]
+        return repr_str.strip()

pepflow/expression_manager_test.py

@@ -0,0 +1,116 @@
+# Copyright: 2025 The PEPFlow Developers
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# 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 function as fc
+from pepflow import pep as pep
+from pepflow import pep_context as pc
+from pepflow import point as pt
+
+
+@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)
+
+
+def test_repr_point_by_basis(pep_context: pc.PEPContext) -> None:
+    x = pt.Point(is_basis=True, tags=["x_0"])
+    f = fc.Function(is_basis=True, tags=["f"])
+    L = 0.5
+    for i in range(2):
+        x = x - L * f.gradient(x)
+        x.add_tag(f"x_{i + 1}")
+
+    em = exm.ExpressionManager(pep_context)
+    np.testing.assert_allclose(em.eval_point(x).vector, [1, -0.5, -0.5])
+    assert (
+        em.repr_point_by_basis(x) == "x_0 - 0.5*gradient_f(x_0) - 0.5*gradient_f(x_1)"
+    )
+
+
+def test_repr_point_by_basis_with_zero(pep_context: pc.PEPContext) -> None:
+    x = pt.Point(is_basis=True, tags=["x_0"])
+    _ = pt.Point(is_basis=True, tags=["x_unused"])  # Add this extra point.
+    f = fc.Function(is_basis=True, tags=["f"])
+    L = 0.5
+    for i in range(2):
+        x = x - L * f.gradient(x)
+        x.add_tag(f"x_{i + 1}")
+
+    em = exm.ExpressionManager(pep_context)
+    # Note the vector representation of point is different from previous case
+    # But the string representation is still the same.
+    np.testing.assert_allclose(em.eval_point(x).vector, [1, 0, -0.5, -0.5])
+    assert (
+        em.repr_point_by_basis(x) == "x_0 - 0.5*gradient_f(x_0) - 0.5*gradient_f(x_1)"
+    )
+
+
+def test_repr_point_by_basis_heavy_ball(pep_context: pc.PEPContext) -> None:
+    x_prev = pt.Point(is_basis=True, tags=["x_{-1}"])
+    x = pt.Point(is_basis=True, tags=["x_0"])
+    f = fc.Function(is_basis=True, tags=["f"])
+
+    beta = 0.5
+    for i in range(2):
+        x_next = x - f.gradient(x) + beta * (x - x_prev)
+        x_next.add_tag(f"x_{i + 1}")
+        x_prev = x
+        x = x_next
+
+    em = exm.ExpressionManager(pep_context)
+    np.testing.assert_allclose(em.eval_point(x).vector, [-0.75, 1.75, -1.5, -1])
+    assert (
+        em.repr_point_by_basis(x)
+        == "-0.75*x_{-1} + 1.75*x_0 - 1.5*gradient_f(x_0) - gradient_f(x_1)"
+    )
+
+
+def test_repr_scalar_by_basis(pep_context: pc.PEPContext) -> None:
+    x = pt.Point(is_basis=True, tags=["x"])
+    f = fc.Function(is_basis=True, tags=["f"])
+
+    s = f(x) + x * f.gradient(x)
+    em = exm.ExpressionManager(pep_context)
+    assert em.repr_scalar_by_basis(s) == "f(x) + <x, gradient_f(x)>"
+
+
+def test_repr_scalar_by_basis_interpolation(pep_context: pc.PEPContext) -> None:
+    xi = pt.Point(is_basis=True, tags=["x_i"])
+    xj = pt.Point(is_basis=True, tags=["x_j"])
+    f = fc.SmoothConvexFunction(is_basis=True, L=1)
+    f.add_tag("f")
+    fi = f(xi)  # noqa: F841
+    fj = f(xj)  # noqa: F841
+    interp_scalar = f.interpolate_ineq("x_i", "x_j")
+    em = exm.ExpressionManager(pep_context)
+    expected_repr = "-f(x_i) + f(x_j) + <x_i, gradient_f(x_j)> - <x_j, gradient_f(x_j)> + 0.5*|gradient_f(x_i)|^2 - <gradient_f(x_i), gradient_f(x_j)> + 0.5*|gradient_f(x_j)|^2"
+    assert em.repr_scalar_by_basis(interp_scalar) == expected_repr
+
+
+# TODO add more tests about repr_scalar_by_basis

pepflow/function.py

@@ -19,7 +19,9 @@
 
 from __future__ import annotations
 
+import numbers
 import uuid
+from typing import TYPE_CHECKING
 
 import attrs
 
@@ -28,6 +30,9 @@ from pepflow import point as pt
 from pepflow import scalar as sc
 from pepflow import utils
 
+if TYPE_CHECKING:
+    from pepflow.constraint import Constraint
+
 
 @attrs.frozen
 class Triplet:
@@ -57,7 +62,6 @@ class ScaledFunc:
 @attrs.mutable
 class Function:
     is_basis: bool
-    reuse_gradient: bool
 
     composition: AddedFunc | ScaledFunc | None = None
 
@@ -158,11 +162,19 @@ class Function:
 
     def add_stationary_point(self, name: str) -> pt.Point:
         # assert we can only add one stationary point?
+        pep_context = pc.get_current_context()
+        if pep_context is None:
+            raise RuntimeError("Did you forget to create a context?")
+        if len(pep_context.stationary_triplets[self]) > 0:
+            raise ValueError(
+                "You are trying to add a stationary point to a function that already has a stationary point."
+            )
         point = pt.Point(is_basis=True)
         point.add_tag(name)
         desired_grad = 0 * point
         desired_grad.add_tag(f"gradient_{self.tag}({name})")
-        self.add_point_with_grad_restriction(point, desired_grad)
+        triplet = self.add_point_with_grad_restriction(point, desired_grad)
+        pep_context.add_stationary_triplet(self, triplet)
         return point
 
         # The following the old gradient(opt) =  0 constraint style.
@@ -184,42 +196,22 @@ class Function:
             raise RuntimeError("Did you forget to create a context?")
 
         if self.is_basis:
-            generate_new_basis = True
-            instances_of_point = 0
             for triplet in pep_context.triplets[self]:
                 if triplet.point.uid == point.uid:
-                    instances_of_point += 1
-                    generate_new_basis = False
-                    previous_triplet = triplet
+                    return triplet
 
-            if generate_new_basis:
-                function_value = sc.Scalar(is_basis=True)
-                function_value.add_tag(f"{self.tag}({point.tag})")
-                gradient = pt.Point(is_basis=True)
-                gradient.add_tag(f"gradient_{self.tag}({point.tag})")
+            function_value = sc.Scalar(is_basis=True)
+            function_value.add_tag(f"{self.tag}({point.tag})")
+            gradient = pt.Point(is_basis=True)
+            gradient.add_tag(f"gradient_{self.tag}({point.tag})")
 
-                new_triplet = Triplet(
-                    point,
-                    function_value,
-                    gradient,
-                    name=f"{point.tag}_{function_value.tag}_{gradient.tag}",
-                )
-                self.add_triplet_to_func(new_triplet)
-            elif not generate_new_basis and self.reuse_gradient:
-                function_value = previous_triplet.function_value
-                gradient = previous_triplet.gradient
-            elif not generate_new_basis and not self.reuse_gradient:
-                function_value = previous_triplet.function_value
-                gradient = pt.Point(is_basis=True)
-                gradient.add_tag(f"gradient_{self.tag}({point.tag})")
-
-                new_triplet = Triplet(
-                    point,
-                    previous_triplet.function_value,
-                    gradient,
-                    name=f"{point.tag}_{function_value.tag}_{gradient.tag}_{instances_of_point}",
-                )
-                self.add_triplet_to_func(new_triplet)
+            new_triplet = Triplet(
+                point,
+                function_value,
+                gradient,
+                name=f"{point.tag}_{function_value.tag}_{gradient.tag}",
+            )
+            self.add_triplet_to_func(new_triplet)
         else:
             if isinstance(self.composition, AddedFunc):
                 left_triplet = self.composition.left_func.generate_triplet(point)
@@ -258,7 +250,6 @@ class Function:
         assert isinstance(other, Function)
         return Function(
             is_basis=False,
-            reuse_gradient=self.reuse_gradient and other.reuse_gradient,
             composition=AddedFunc(self, other),
             tags=[f"{self.tag}+{other.tag}"],
         )
@@ -270,7 +261,6 @@ class Function:
             tag_other = f"({other.tag})"
         return Function(
             is_basis=False,
-            reuse_gradient=self.reuse_gradient and other.reuse_gradient,
             composition=AddedFunc(self, -other),
             tags=[f"{self.tag}-{tag_other}"],
         )
@@ -282,7 +272,6 @@ class Function:
             tag_self = f"({self.tag})"
         return Function(
             is_basis=False,
-            reuse_gradient=self.reuse_gradient,
             composition=ScaledFunc(scale=other, base_func=self),
             tags=[f"{other:.4g}*{tag_self}"],
         )
@@ -294,7 +283,6 @@ class Function:
             tag_self = f"({self.tag})"
         return Function(
             is_basis=False,
-            reuse_gradient=self.reuse_gradient,
             composition=ScaledFunc(scale=other, base_func=self),
             tags=[f"{other:.4g}*{tag_self}"],
         )
@@ -305,7 +293,6 @@ class Function:
             tag_self = f"({self.tag})"
         return Function(
             is_basis=False,
-            reuse_gradient=self.reuse_gradient,
             composition=ScaledFunc(scale=-1, base_func=self),
             tags=[f"-{tag_self}"],
         )
@@ -317,7 +304,6 @@ class Function:
             tag_self = f"({self.tag})"
         return Function(
             is_basis=False,
-            reuse_gradient=self.reuse_gradient,
             composition=ScaledFunc(scale=1 / other, base_func=self),
             tags=[f"1/{other:.4g}*{tag_self}"],
         )
@@ -331,10 +317,96 @@ class Function:
         return self.uid == other.uid
 
 
+class ConvexFunction(Function):
+    def __init__(
+        self,
+        is_basis=True,
+        composition=None,
+    ):
+        super().__init__(
+            is_basis=is_basis,
+            composition=composition,
+        )
+
+    def convex_interpolability_constraints(
+        self, triplet_i: Triplet, triplet_j: Triplet
+    ) -> Constraint:
+        point_i = triplet_i.point
+        function_value_i = triplet_i.function_value
+
+        point_j = triplet_j.point
+        function_value_j = triplet_j.function_value
+        grad_j = triplet_j.gradient
+
+        func_diff = function_value_j - function_value_i
+        cross_term = grad_j * (point_i - point_j)
+
+        return (func_diff + cross_term).le(
+            0, name=f"{self.tag}:{point_i.tag},{point_j.tag}"
+        )
+
+    def get_interpolation_constraints(
+        self, pep_context: pc.PEPContext | None = None
+    ) -> list[Constraint]:
+        interpolation_constraints = []
+        if pep_context is None:
+            pep_context = pc.get_current_context()
+        if pep_context is None:
+            raise RuntimeError("Did you forget to create a context?")
+        for i in pep_context.triplets[self]:
+            for j in pep_context.triplets[self]:
+                if i == j:
+                    continue
+                interpolation_constraints.append(
+                    self.convex_interpolability_constraints(i, j)
+                )
+        return interpolation_constraints
+
+    def interpolate_ineq(
+        self, p1_tag: str, p2_tag: str, pep_context: pc.PEPContext | None = None
+    ) -> sc.Scalar:
+        """Generate the interpolation inequality scalar by tags."""
+        if pep_context is None:
+            pep_context = pc.get_current_context()
+        if pep_context is None:
+            raise RuntimeError("Did you forget to specify a context?")
+        # TODO: we definitely need a more robust tag system
+        x1 = pep_context.get_by_tag(p1_tag)
+        x2 = pep_context.get_by_tag(p2_tag)
+        f1 = pep_context.get_by_tag(f"{self.tag}({p1_tag})")
+        f2 = pep_context.get_by_tag(f"{self.tag}({p2_tag})")
+        g2 = pep_context.get_by_tag(f"gradient_{self.tag}({p2_tag})")
+        return f2 - f1 + g2 * (x1 - x2)
+
+    def proximal_step(self, x_0: pt.Point, stepsize: numbers.Number) -> pt.Point:
+        gradient = pt.Point(is_basis=True)
+        gradient.add_tag(
+            f"gradient_{self.tag}(prox_{{{stepsize}*{self.tag}}}({x_0.tag}))"
+        )
+        function_value = sc.Scalar(is_basis=True)
+        function_value.add_tag(f"{self.tag}(prox_{{{stepsize}*{self.tag}}}({x_0.tag}))")
+        x = x_0 - stepsize * gradient
+        x.add_tag(f"prox_{{{stepsize}*{self.tag}}}({x_0.tag})")
+        new_triplet = Triplet(
+            x,
+            function_value,
+            gradient,
+            name=f"{x.tag}_{function_value.tag}_{gradient.tag}",
+        )
+        self.add_triplet_to_func(new_triplet)
+        return x
+
+
 class SmoothConvexFunction(Function):
-    def __init__(self, L, is_basis=True, composition=None, reuse_gradient=True):
+    def __init__(
+        self,
+        L,
+        is_basis=True,
+        composition=None,
+    ):
         super().__init__(
-            is_basis=is_basis, composition=composition, reuse_gradient=reuse_gradient
+            is_basis=is_basis,
+            composition=composition,
         )
         self.L = L
 
@@ -372,7 +444,7 @@ class SmoothConvexFunction(Function):
 
     def interpolate_ineq(
         self, p1_tag: str, p2_tag: str, pep_context: pc.PEPContext | None = None
-    ) -> pt.Scalar:
+    ) -> sc.Scalar:
         """Generate the interpolation inequality scalar by tags."""
         if pep_context is None:
             pep_context = pc.get_current_context()