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

pepflow/function.py

@@ -19,7 +19,9 @@
 
 from __future__ import annotations
 
+import numbers
 import uuid
+from typing import TYPE_CHECKING
 
 import attrs
 
@@ -28,9 +30,25 @@ 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:
+    """
+    A data class that represents, for some given function :math:`f`,
+    the tuple :math:`\\{x, f(x), \\nabla f(x)\\}`. We can also consider
+    a subgradient :math:`\\widetilde{\\nabla} f(x)` instead of the gradient.
+
+    Attributes:
+        point (:class:`Point`): The point :math:`x`.
+        function_value (:class:`Scalar`): The function value :math:`f(x)`.
+        gradient (:class:`Point`): The gradient :math:`\\nabla f(x)` or
+            a subgradient :math:`\\widetilde{\\nabla} f(x)`.
+        name (str): The unique name of the :class:`Triplet` object.
+    """
+
     point: pt.Point
     function_value: sc.Scalar
     gradient: pt.Point
@@ -48,7 +66,7 @@ class AddedFunc:
 
 @attrs.frozen
 class ScaledFunc:
-    """Represents scale * base_func."""
+    """Represents scalar * base_func."""
 
     scale: float
     base_func: Function
@@ -56,8 +74,26 @@ class ScaledFunc:
 
 @attrs.mutable
 class Function:
+    """A :class:`Function` object represents a function.
+
+    :class:`Function` objects can be constructed as linear combinations
+    of other :class:`Function` objects. Let `a` and `b` be some numeric
+    data type. Let `f` and `g` be :class:`Function` objects. Then, we
+    can form a new :class:`Function` object: `a*f+b*g`.
+
+    A :class:`Function` object should never be explicitly constructed. Only
+    children of :class:`Function` such as :class:`ConvexFunction` or
+    :class:`SmoothConvexFunction` should be constructed. See their respective
+    documentation to see how.
+
+    Attributes:
+        is_basis (bool): `True` if this function is not formed through a linear
+            combination of other functions. `False` otherwise.
+        tags (list[str]): A list that contains tags that can be used to
+            identify the :class:`Function` object. Tags should be unique.
+    """
+
     is_basis: bool
-    reuse_gradient: bool
 
     composition: AddedFunc | ScaledFunc | None = None
 
@@ -75,11 +111,21 @@ class Function:
 
     @property
     def tag(self):
+        """Returns the most recently added tag.
+
+        Returns:
+            str: The most recently added tag of this :class:`Function` object.
+        """
         if len(self.tags) == 0:
             raise ValueError("Function should have a name.")
         return self.tags[-1]
 
     def add_tag(self, tag: str) -> None:
+        """Add a new tag for this :class:`Function` object.
+
+        Args:
+            tag (str): The new tag to be added to the `tags` list.
+        """
         self.tags.append(tag)
 
     def __repr__(self):
@@ -89,7 +135,7 @@ class Function:
 
     def _repr_latex_(self):
         s = repr(self)
-        return rf"$\\displaystyle {s}$"
+        return rf"$\displaystyle {s}$"
 
     def get_interpolation_constraints(self):
         raise NotImplementedError(
@@ -157,12 +203,32 @@ class Function:
         return triplet
 
     def add_stationary_point(self, name: str) -> pt.Point:
+        """
+        Return a stationary point for this :class:`Function` object.
+        A :class:`Function` object can only have one stationary point.
+
+        Args:
+            name (str): The tag for the :class:`Point` object which
+                 will serve as the stationary point.
+
+        Returns:
+            :class:`Point`: The stationary point for this :class:`Function`
+            object.
+        """
         # 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 = pt.Point.zero()  # zero 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.
@@ -183,43 +249,26 @@ class Function:
         if pep_context is None:
             raise RuntimeError("Did you forget to create a context?")
 
+        if not isinstance(point, pt.Point):
+            raise ValueError("The Function can only take point as input.")
+
         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)
@@ -240,14 +289,50 @@ class Function:
         return Triplet(point, function_value, gradient, name=None)
 
     def gradient(self, point: pt.Point) -> pt.Point:
+        """
+        Returns a :class:`Point` object that is the gradient of the
+        :class:`Function` at the given :class:`Point`.
+
+        Args:
+            point (:class:`Point`): Any :class:`Point`.
+
+        Returns:
+            :class:`Point`: The gradient of the :class:`Function` at the
+            given :class:`Point`.
+        """
         triplet = self.generate_triplet(point)
         return triplet.gradient
 
     def subgradient(self, point: pt.Point) -> pt.Point:
+        """
+        Returns a :class:`Point` object that is the subgradient of the
+        :class:`Function` at the given :class:`Point`.
+
+        Args:
+            point (:class:`Point`): Any :class:`Point`.
+
+        Returns:
+            :class:`Point`: The subgradient of the :class:`Function` at the
+            given :class:`Point`.
+
+        Note:
+            The method `gradient` is exactly the same.
+        """
         triplet = self.generate_triplet(point)
         return triplet.gradient
 
     def function_value(self, point: pt.Point) -> sc.Scalar:
+        """
+        Returns a :class:`Scalar` object that is the function value of the
+        :class:`Function` at the given :class:`Point`.
+
+        Args:
+            point (:class:`Point`): Any :class:`Point`.
+
+        Returns:
+            :class:`Point`: The function value of the :class:`Function` at the
+            given :class:`Point`.
+        """
         triplet = self.generate_triplet(point)
         return triplet.function_value
 
@@ -255,46 +340,46 @@ class Function:
         return self.function_value(point)
 
     def __add__(self, other):
-        assert isinstance(other, Function)
+        if not isinstance(other, Function):
+            return NotImplemented
         return Function(
             is_basis=False,
-            reuse_gradient=self.reuse_gradient and other.reuse_gradient,
             composition=AddedFunc(self, other),
             tags=[f"{self.tag}+{other.tag}"],
         )
 
     def __sub__(self, other):
-        assert isinstance(other, Function)
+        if not isinstance(other, Function):
+            return NotImplemented
         tag_other = other.tag
         if isinstance(other.composition, AddedFunc):
             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}"],
         )
 
     def __mul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         tag_self = self.tag
         if isinstance(self.composition, AddedFunc):
             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}"],
         )
 
     def __rmul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         tag_self = self.tag
         if isinstance(self.composition, AddedFunc):
             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,19 +390,18 @@ 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}"],
         )
 
     def __truediv__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         tag_self = self.tag
         if isinstance(self.composition, AddedFunc):
             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 +415,159 @@ class Function:
         return self.uid == other.uid
 
 
+class ConvexFunction(Function):
+    """
+    The :class:`ConvexFunction` class is a child of :class:`Function.`
+    The :class:`ConvexFunction` class represents a closed, convex, and
+    proper (CCP) function, i.e., a convex function whose epigraph is a
+    non-empty closed set.
+
+    A CCP function typically has no parameters. We can instantiate a
+    :class:`ConvexFunction` object as follows:
+
+    Example:
+        >>> import pepflow as pf
+        >>> ctx = pf.PEPContext("example").set_as_current()
+        >>> pep_builder = pf.PEPBuilder()
+        >>> g = pep_builder.declare_func(pf.ConvexFunction, "g")
+    """
+
+    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 :class:`Scalar` by tags.
+        The interpolation inequality between two points :math:`p_1, p_2` for a
+        CCP function :math:`f` is
+
+        .. math:: f(p_2) - f(p_1) + \\langle \\nabla f(p_2), p_1 - p_2 \\rangle.
+
+        Args:
+            p1_tag (str): A tag of the :class:`Point` :math:`p_1`.
+            p2_tag (str): A tag of the :class:`Point` :math:`p_2`.
+        """
+        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:
+        """ Define the proximal operator as
+
+        .. math:: \\text{prox}_{\\gamma f}(x_0) := \\arg\\min_x \\left\\{ \\gamma f(x) + \\frac{1}{2} \\|x - x_0\\|^2 \\right\\}.
+
+        This function performs a proximal step with respect to some
+        :class:`Function` :math:`f` on the :class:`Point` :math:`x_0`
+        with stepsize :math:`\\gamma`: 
+
+        .. math::
+            :nowrap:
+
+            \\begin{eqnarray}
+                x := \\text{prox}_{\\gamma f}(x_0) & := & \\arg\\min_x \\left\\{ \\gamma f(x) + \\frac{1}{2} \\|x - x_0\\|^2 \\right\\}, \\\\
+                & \\Updownarrow & \\\\
+                0 & = & \\gamma \\partial f(x) + x - x_0,\\\\
+                & \\Updownarrow & \\\\
+                x & = & x_0 - \\gamma \\widetilde{\\nabla} f(x) \\text{ where } \\widetilde{\\nabla} f(x)\\in\\partial f(x).
+            \\end{eqnarray}
+
+        Args:
+            x_0 (:class:`Point`): The initial point.
+            stepsize (int | float): The stepsize.
+        """
+        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):
+    """
+    The :class:`SmoothConvexFunction` class is a child of :class:`Function.`
+    The :class:`SmoothConvexFunction` class represents a smooth,
+    convex function.
+
+    A smooth, convex function has a smoothness parameter :math:`L`.
+    We can instantiate a :class:`SmoothConvexFunction` object as follows:
+
+    Example:
+        >>> import pepflow as pf
+        >>> ctx = pf.PEPContext("example").set_as_current()
+        >>> pep_builder = pf.PEPBuilder()
+        >>> f = pep_builder.declare_func(pf.SmoothConvexFunction, "f", L=1)
+    """
+
+    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,8 +605,17 @@ class SmoothConvexFunction(Function):
 
     def interpolate_ineq(
         self, p1_tag: str, p2_tag: str, pep_context: pc.PEPContext | None = None
-    ) -> pt.Scalar:
-        """Generate the interpolation inequality scalar by tags."""
+    ) -> sc.Scalar:
+        """Generate the interpolation inequality :class:`Scalar` by tags.
+        The interpolation inequality between two points :math:`p_1, p_2` for a
+        smooth, convex function :math:`f` is
+
+        .. math:: f(p_2) - f(p_1) + \\langle \\nabla f(p_2), p_1 - p_2 \\rangle + \\tfrac{1}{2} \\lVert \\nabla f(p_1) - \\nabla f(p_2) \\rVert^2.
+
+        Args:
+            p1_tag (str): A tag of the :class:`Point` :math:`p_1`.
+            p2_tag (str): A tag of the :class:`Point` :math:`p_2`.
+        """
         if pep_context is None:
             pep_context = pc.get_current_context()
         if pep_context is None: