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

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
@@ -38,26 +37,54 @@ GLOBAL_CONTEXT_DICT: dict[str, PEPContext] = {}
 
 
 def get_current_context() -> PEPContext | None:
+    """
+    Return the current global :class:`PEPContext`.
+
+    Returns:
+        :class:`PEPContext`: The current global :class:`PEPContext`.
+    """
     return CURRENT_CONTEXT
 
 
 def set_current_context(ctx: PEPContext | None):
+    """
+    Change the current global :class:`PEPContext`.
+
+    Args:
+        ctx (:class:`PEPContext`): The :class:`PEPContext` to set as the new
+            global :class:`PEPContext`.
+    """
     global CURRENT_CONTEXT
     assert ctx is None or isinstance(ctx, PEPContext)
     CURRENT_CONTEXT = ctx
 
 
 class PEPContext:
+    """
+    A :class:`PEPContext` object is a context manager which maintains
+    the abstract mathematical objects of the Primal and Dual PEP.
+
+    Attributes:
+        name (str): The unique name of the :class:`PEPContext` object.
+    """
+
     def __init__(self, name: str):
         self.name = name
         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
 
     def set_as_current(self) -> PEPContext:
+        """
+        Set this :class:`PEPContext` object as the global context.
+
+        Returns:
+            :class:`PEPContext`: This :class:`PEPContext` object.
+        """
         set_current_context(self)
         return self
 
@@ -70,35 +97,98 @@ 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:
+        """
+        Under this :class:`PEPContext`, get the :class:`Point` or
+        :class:`Scalar` object associated with the provided `tag`.
+
+        Args:
+            tag (str): The tag of the :class:`Point` or :class:`Scalar` object
+                we want to retrieve.
+
+        Returns:
+            :class:`Point` | :class:`Scalar`: The :class:`Point` or
+            :class:`Scalar` object associated with the provided `tag`.
+        """
         for p in self.points:
             if tag in p.tags:
                 return p
         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):
+        """Reset this :class:`PEPContext` object."""
         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]:
+        """
+        Each function :math:`f` used in Primal and Dual PEP is associated with
+        a set of triplets :math:`\\{x_i, f(x_i), \\nabla f(x_i)\\}` visited by
+        the considered algorithm. We can also consider a subgradient
+        :math:`\\widetilde{\\nabla} f(x)` instead of the gradient. This
+        function returns a list of the visited points :math:`\\{x_i\\}` under
+        this :class:`PEPContext`.
+
+        Args:
+            func (:class:`Function`): The function associated with the set
+                of triplets :math:`\\{x_i, f(x_i), \\nabla f(x_i)\\}`.
+
+        Returns:
+            list[:class:`Point`]: The list of the visited points
+            :math:`\\{x_i\\}`.
+        """
         return natsort.natsorted(
             [t.point for t in self.triplets[func]], key=lambda x: x.tag
         )
 
     def tracked_grad(self, func: Function) -> list[Point]:
+        """
+        Each function :math:`f` used in Primal and Dual PEP is associated with
+        a set of triplets :math:`\\{x_i, f(x_i), \\nabla f(x_i)\\}` visited by
+        the considered algorithm. We can also consider a subgradient
+        :math:`\\widetilde{\\nabla} f(x)` instead of the gradient
+        :math:`\\nabla f(x_i)`. This function returns a list of the visited
+        gradients :math:`\\{\\nabla f(x_i)\\}` under this :class:`PEPContext`.
+
+        Args:
+            func (:class:`Function`): The function associated with the set
+                of triplets :math:`\\{x_i, f(x_i), \\nabla f(x_i)\\}`.
+
+        Returns:
+            list[:class:`Point`]: The list of the visited gradients
+            :math:`\\{\\nabla f(x_i)\\}`.
+
+        """
         return natsort.natsorted(
             [t.gradient for t in self.triplets[func]], key=lambda x: x.tag
         )
 
     def tracked_func_value(self, func: Function) -> list[Scalar]:
+        """
+        Each function :math:`f` used in Primal and Dual PEP is associated with
+        a set of triplets :math:`\\{x_i, f(x_i), \\nabla f(x_i)\\}` visited by
+        the considered algorithm. We can also consider a subgradient
+        :math:`\\widetilde{\\nabla} f(x)` instead of the gradient
+        :math:`\\nabla f(x_i)`. This function returns a list of the visited
+        function values :math:`\\{f(x_i)\\}` under this :class:`PEPContext`.
+
+        Args:
+            func (:class:`Function`): The function associated with the set of
+                triplets :math:`\\{x_i, f(x_i), \\nabla f(x_i)\\}`.
+
+        Returns:
+            list[:class:`Scalar`]: The list of the visited function values
+            :math:`\\{f(x_i)\\}`.
+
+        """
         return natsort.natsorted(
             [t.function_value for t in self.triplets[func]], key=lambda x: x.tag
         )
@@ -134,3 +224,29 @@ class PEPContext:
             func_to_order[func] = order
 
         return func_to_df, func_to_order
+
+    def basis_points(self) -> list[Point]:
+        """
+        Return a list of the basis :class:`Point` objects managed by this
+        :class:`PEPContext`.
+
+        Returns:
+            list[:class:`Point`]: A list of the basis :class:`Point` objects
+            managed by this :class:`PEPContext`.
+        """
+        return [
+            p for p in self.points if p.is_basis
+        ]  # Note the order is always the same as added time
+
+    def basis_scalars(self) -> list[Scalar]:
+        """
+        Return a list of the basis :class:`Scalar` objects managed by this
+        :class:`PEPContext`.
+
+        Returns:
+            list[:class:`Scalar`]: A list of the basis :class:`Scalar` objects
+            managed by this :class:`PEPContext`.
+        """
+        return [
+            s for s in self.scalars if s.is_basis
+        ]  # Note the order is always the same as added time

pepflow/pep_context_test.py

@@ -25,6 +25,7 @@ import pytest
 from pepflow import pep_context as pc
 from pepflow.function import SmoothConvexFunction
 from pepflow.point import Point
+from pepflow.scalar import Scalar
 
 
 @pytest.fixture
@@ -91,10 +92,34 @@ 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)
+
+
+def test_basis_points(pep_context: pc.PEPContext):
+    p1 = Point(is_basis=True, tags=["x1"])
+    p2 = Point(is_basis=True, tags=["x2"])
+    _ = p1 + p2  # not basis
+    ps = Point(is_basis=True, tags=["x_star"])
+    p0 = Point(is_basis=True, tags=["x0"])
+
+    assert pep_context.basis_points() == [p1, p2, ps, p0]
+
+
+def test_basis_scalars(pep_context: pc.PEPContext):
+    p1 = Point(is_basis=True, tags=["x1"])
+    p2 = Point(is_basis=True, tags=["x2"])
+    _ = p1 * p2  # not basis
+    s1 = Scalar(is_basis=True, tags=["s2"])
+    s2 = Scalar(is_basis=True, tags=["s1"])
+
+    assert pep_context.basis_scalars() == [s1, s2]

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.py

@@ -27,37 +27,65 @@ import numpy as np
 
 from pepflow import pep_context as pc
 from pepflow import utils
-from pepflow.scalar import EvalExpressionScalar, Scalar
+from pepflow.scalar import Scalar, ScalarRepresentation
 
 
 def is_numerical_or_point(val: Any) -> bool:
-    return utils.is_numerical(val) or isinstance(val, Point)
+    return utils.is_numerical_or_parameter(val) or isinstance(val, Point)
 
 
 def is_numerical_or_evaluatedpoint(val: Any) -> bool:
-    return utils.is_numerical(val) or isinstance(val, EvaluatedPoint)
+    return utils.is_numerical_or_parameter(val) or isinstance(val, EvaluatedPoint)
 
 
 @attrs.frozen
-class EvalExpressionPoint:
+class PointRepresentation:
     op: utils.Op
     left_point: Point | float
     right_point: Point | float
 
 
+@attrs.frozen
+class ZeroPoint:
+    """A special class to represent 0 in Point."""
+
+    pass
+
+
 @attrs.frozen
 class EvaluatedPoint:
+    """
+    The concrete representation of the abstract :class:`Point`.
+
+    Each abstract basis :class:`Point` object has a unique concrete
+    representation as a unit vector. The concrete representations of
+    linear combinations of abstract basis :class:`Point` objects are
+    linear combinations of the unit vectors. This information is stored
+    in the `vector` attribute.
+
+    :class:`EvaluatedPoint` objects can be constructed as linear combinations
+    of other :class:`EvaluatedPoint` objects. Let `a` and `b` be some numeric
+    data type. Let `x` and `y` be :class:`EvaluatedPoint` objects. Then, we
+    can form a new :class:`EvaluatedPoint` object: `a*x+b*y`.
+
+    Attributes:
+        vector (np.ndarray): The concrete representation of an
+            abstract :class:`Point`.
+    """
+
     vector: np.ndarray
 
+    @classmethod
+    def zero(cls, num_basis_points: int):
+        return EvaluatedPoint(vector=np.zeros(num_basis_points))
+
     def __add__(self, other):
         if isinstance(other, EvaluatedPoint):
             return EvaluatedPoint(vector=self.vector + other.vector)
         elif utils.is_numerical(other):
             return EvaluatedPoint(vector=self.vector + other)
         else:
-            raise ValueError(
-                f"Unsupported add operation between EvaluatedPoint and {type(other)}"
-            )
+            return NotImplemented
 
     def __radd__(self, other):
         return self.__add__(other)
@@ -68,9 +96,7 @@ class EvaluatedPoint:
         elif utils.is_numerical(other):
             return EvaluatedPoint(vector=self.vector - other)
         else:
-            raise ValueError(
-                f"Unsupported sub operation between EvaluatedPoint and {type(other)}"
-            )
+            return NotImplemented
 
     def __rsub__(self, other):
         if isinstance(other, EvaluatedPoint):
@@ -78,30 +104,51 @@ class EvaluatedPoint:
         elif utils.is_numerical(other):
             return EvaluatedPoint(vector=other - self.vector)
         else:
-            raise ValueError(
-                f"Unsupported sub operation between EvaluatedPoint and {type(other)}"
-            )
+            return NotImplemented
 
     def __mul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         return EvaluatedPoint(vector=self.vector * other)
 
     def __rmul__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         return EvaluatedPoint(vector=other * self.vector)
 
     def __truediv__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical(other):
+            return NotImplemented
         return EvaluatedPoint(vector=self.vector / other)
 
 
 @attrs.frozen
 class Point:
+    """
+    A :class:`Point` object represents an element of a pre-Hilbert space.
+    Examples include a point or a gradient.
+
+    :class:`Point` objects can be constructed as linear combinations of
+    other :class:`Point` objects. Let `a` and `b` be some numeric data type.
+    Let `x` and `y` be :class:`Point` objects. Then, we can form a new
+    :class:`Point` object: `a*x+b*y`.
+
+    The inner product of two :class:`Point` objects can also be taken.
+    Let `x` and `y` be :class:`Point` objects. Then, their inner product is
+    `x*y` and returns a :class:`Scalar` object.
+
+    Attributes:
+        is_basis (bool): `True` if this point is not formed through a linear
+            combination of other points. `False` otherwise.
+        tags (list[str]): A list that contains tags that can be used to
+            identify the :class:`Point` object. Tags should be unique.
+    """
+
     # If true, the point is the basis for the evaluations of G
     is_basis: bool
 
-    # How to evaluate the point.
-    eval_expression: EvalExpressionPoint | None = None
+    # The representation of point used for evaluation.
+    eval_expression: PointRepresentation | ZeroPoint | None = None
 
     # Human tagged value for the Point
     tags: list[str] = attrs.field(factory=list)
@@ -120,13 +167,27 @@ class Point:
             raise RuntimeError("Did you forget to create a context?")
         pep_context.add_point(self)
 
+    @staticmethod
+    def zero() -> Point:
+        return Point(is_basis=False, eval_expression=ZeroPoint(), tags=["0"])
+
     @property
     def tag(self):
+        """Returns the most recently added tag.
+
+        Returns:
+            str: The most recently added tag of this :class:`Point` object.
+        """
         if len(self.tags) == 0:
             raise ValueError("Point should have a name.")
         return self.tags[-1]
 
     def add_tag(self, tag: str) -> None:
+        """Add a new tag for this :class:`Point` object.
+
+        Args:
+            tag (str): The new tag to be added to the `tags` list.
+        """
         self.tags.append(tag)
 
     def __repr__(self):
@@ -135,18 +196,15 @@ class Point:
         return super().__repr__()
 
     def _repr_latex_(self):
-        s = repr(self)
-        s = s.replace("star", r"\star")
-        s = s.replace("gradient_", r"\nabla ")
-        s = s.replace("|", r"\|")
-        return rf"$\\displaystyle {s}$"
+        return utils.str_to_latex(repr(self))
 
     # TODO: add a validator that `is_basis` and `eval_expression` are properly setup.
     def __add__(self, other):
-        assert isinstance(other, Point)
+        if not isinstance(other, Point):
+            return NotImplemented
         return Point(
             is_basis=False,
-            eval_expression=EvalExpressionPoint(utils.Op.ADD, self, other),
+            eval_expression=PointRepresentation(utils.Op.ADD, self, other),
             tags=[f"{self.tag}+{other.tag}"],
         )
 
@@ -154,72 +212,78 @@ class Point:
         # TODO: come up with better way to handle this
         if other == 0:
             return self
-        assert isinstance(other, Point)
+        if not isinstance(other, Point):
+            return NotImplemented
         return Point(
             is_basis=False,
-            eval_expression=EvalExpressionPoint(utils.Op.ADD, other, self),
+            eval_expression=PointRepresentation(utils.Op.ADD, other, self),
             tags=[f"{other.tag}+{self.tag}"],
         )
 
     def __sub__(self, other):
-        assert isinstance(other, Point)
+        if not isinstance(other, Point):
+            return NotImplemented
         tag_other = utils.parenthesize_tag(other)
         return Point(
             is_basis=False,
-            eval_expression=EvalExpressionPoint(utils.Op.SUB, self, other),
+            eval_expression=PointRepresentation(utils.Op.SUB, self, other),
             tags=[f"{self.tag}-{tag_other}"],
         )
 
     def __rsub__(self, other):
-        assert isinstance(other, Point)
+        if not isinstance(other, Point):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
         return Point(
             is_basis=False,
-            eval_expression=EvalExpressionPoint(utils.Op.SUB, other, self),
+            eval_expression=PointRepresentation(utils.Op.SUB, other, self),
             tags=[f"{other.tag}-{tag_self}"],
         )
 
     def __mul__(self, other):
-        # TODO allow the other to be point so that we return a scalar.
-        assert is_numerical_or_point(other)
+        if not is_numerical_or_point(other):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
-        if utils.is_numerical(other):
+        if utils.is_numerical_or_parameter(other):
+            tag_other = utils.numerical_str(other)
             return Point(
                 is_basis=False,
-                eval_expression=EvalExpressionPoint(utils.Op.MUL, self, other),
-                tags=[f"{tag_self}*{other:.4g}"],
+                eval_expression=PointRepresentation(utils.Op.MUL, self, other),
+                tags=[f"{tag_self}*{tag_other}"],
             )
         else:
             tag_other = utils.parenthesize_tag(other)
             return Scalar(
                 is_basis=False,
-                eval_expression=EvalExpressionScalar(utils.Op.MUL, self, other),
+                eval_expression=ScalarRepresentation(utils.Op.MUL, self, other),
                 tags=[f"{tag_self}*{tag_other}"],
             )
 
     def __rmul__(self, other):
-        # TODO allow the other to be point so that we return a scalar.
-        assert is_numerical_or_point(other)
+        if not is_numerical_or_point(other):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
-        if utils.is_numerical(other):
+        if utils.is_numerical_or_parameter(other):
+            tag_other = utils.numerical_str(other)
             return Point(
                 is_basis=False,
-                eval_expression=EvalExpressionPoint(utils.Op.MUL, other, self),
-                tags=[f"{other:.4g}*{tag_self}"],
+                eval_expression=PointRepresentation(utils.Op.MUL, other, self),
+                tags=[f"{tag_other}*{tag_self}"],
             )
         else:
             tag_other = utils.parenthesize_tag(other)
             return Scalar(
                 is_basis=False,
-                eval_expression=EvalExpressionScalar(utils.Op.MUL, other, self),
+                eval_expression=ScalarRepresentation(utils.Op.MUL, other, self),
                 tags=[f"{tag_other}*{tag_self}"],
             )
 
     def __pow__(self, power):
-        assert power == 2
+        if power != 2:
+            return NotImplemented
         return Scalar(
             is_basis=False,
-            eval_expression=EvalExpressionScalar(utils.Op.MUL, self, self),
+            eval_expression=ScalarRepresentation(utils.Op.MUL, self, self),
             tags=[rf"|{self.tag}|^{power}"],
         )
 
@@ -227,17 +291,19 @@ class Point:
         tag_self = utils.parenthesize_tag(self)
         return Point(
             is_basis=False,
-            eval_expression=EvalExpressionPoint(utils.Op.MUL, -1, self),
+            eval_expression=PointRepresentation(utils.Op.MUL, -1, self),
             tags=[f"-{tag_self}"],
         )
 
     def __truediv__(self, other):
-        assert utils.is_numerical(other)
+        if not utils.is_numerical_or_parameter(other):
+            return NotImplemented
         tag_self = utils.parenthesize_tag(self)
+        tag_other = f"1/{utils.numerical_str(other)}"
         return Point(
             is_basis=False,
-            eval_expression=EvalExpressionPoint(utils.Op.DIV, self, other),
-            tags=[f"1/{other:.4g}*{tag_self}"],
+            eval_expression=PointRepresentation(utils.Op.DIV, self, other),
+            tags=[f"{tag_other}*{tag_self}"],
         )
 
     def __hash__(self):
@@ -249,6 +315,20 @@ class Point:
         return self.uid == other.uid
 
     def eval(self, ctx: pc.PEPContext | None = None) -> np.ndarray:
+        """
+        Return the concrete representation of this :class:`Point`.
+        Concrete representations of :class:`Point` objects are
+        :class:`EvaluatedPoint` objects.
+
+        Args:
+            ctx (:class:`PEPContext` | None): The :class:`PEPContext` object
+                we consider. `None` if we consider the current global
+                :class:`PEPContext` object.
+
+        Returns:
+            :class:`EvaluatedPoint`: The concrete representation of
+            this :class:`Point`.
+        """
         from pepflow.expression_manager import ExpressionManager
 
         # Note this can be inefficient.
@@ -258,3 +338,29 @@ class Point:
             raise RuntimeError("Did you forget to create a context?")
         em = ExpressionManager(ctx)
         return em.eval_point(self).vector
+
+    def repr_by_basis(self, ctx: pc.PEPContext | None = None) -> str:
+        """
+        Express this :class:`Point` object as the linear combination of
+        the basis :class:`Point` objects of the given :class:`PEPContext`.
+        This linear combination is expressed as a `str` where, to refer to
+        the basis :class:`Point` objects, we use their tags.
+
+        Args:
+            ctx (:class:`PEPContext`): The :class:`PEPContext` object
+                whose basis :class:`Point` objects we consider. `None` if
+                we consider the current global :class:`PEPContext` object.
+
+        Returns:
+            str: The representation of this :class:`Point` object in terms of
+            the basis :class:`Point` objects of the given :class:`PEPContext`.
+        """
+        from pepflow.expression_manager import ExpressionManager
+
+        # Note this can be inefficient.
+        if ctx is None:
+            ctx = pc.get_current_context()
+        if ctx is None:
+            raise RuntimeError("Did you forget to create a context?")
+        em = ExpressionManager(ctx)
+        return em.repr_point_by_basis(self)