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

pepflow/pep.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 import contextlib
+from collections import defaultdict
 from typing import TYPE_CHECKING, Any, Iterator
 
 import attrs
@@ -30,22 +31,50 @@ from pepflow import pep_context as pc
 from pepflow import point as pt
 from pepflow import scalar as sc
 from pepflow import solver as ps
+from pepflow import utils
 from pepflow.constants import PSD_CONSTRAINT
 
 if TYPE_CHECKING:
     from pepflow.constraint import Constraint
     from pepflow.function import Function
-    from pepflow.solver import DualVariableManager
+    from pepflow.solver import DualVariableManager, PrimalVariableManager
 
 
 @attrs.frozen
 class PEPResult:
+    """
+    A data class object that contains the results of solving the Primal
+    PEP.
+
+    Attributes:
+        primal_opt_value (float): The objective value of the solved Primal PEP.
+        dual_var_manager (:class:`DualVariableManager`): A manager object which
+            provides access to the dual variables associated with the
+            constraints of the Primal PEP.
+        solver_status (Any): States whether the solver managed to solve the
+            Primal PEP successfully.
+        context (:class:`PEPContext`): The :class:`PEPContext` object used to
+            solve the Primal PEP.
+
+    """
+
     primal_opt_value: float
     dual_var_manager: DualVariableManager
     solver_status: Any
     context: pc.PEPContext
 
     def get_function_dual_variables(self) -> dict[Function, np.ndarray]:
+        """
+        Return a dictionary which contains the associated dual variables of the
+        interpolation constraints for Primal PEP.
+
+        Returns:
+            dict[:class:`Function`, np.ndarray]: A dictionary where the keys
+            are :class:`Function` objects, and the values are the dual
+            variables associated with the interpolation constraints of the
+            :class:`Function` key.
+        """
+
         def get_matrix_of_dual_value(df: pd.DataFrame) -> np.ndarray:
             # Check if we need to update the order.
             return (
@@ -68,31 +97,90 @@ class PEPResult:
 
         return df_dict_matrix
 
-    def get_psd_dual_matrix(self):
+    def get_psd_dual_matrix(self) -> np.ndarray:
+        """
+        Return the PSD dual variable matrix associated with the constraint
+        that the Primal PEP decision variable :math:`G` is PSD.
+
+        Returns:
+            np.ndarray: The PSD dual variable matrix associated with the
+            constraint that the Primal PEP decision variable :math:`G` is PSD.
+        """
         return np.array(self.dual_var_manager.dual_value(PSD_CONSTRAINT))
 
 
+@attrs.frozen
+class DualPEPResult:
+    """
+    A data class object that contains the results of solving the Dual
+    PEP.
+
+    Attributes:
+        dual_opt_value (float): The objective value of the solved Dual PEP.
+        primal_var_manager (:class:`PrimalVariableManager`): A manager object
+            which provides access to the primal variables of Dual PEP.
+        solver_status (Any): States whether the solver managed to solve the
+            Dual PEP successfully.
+        context (:class:`PEPContext`): The :class:`PEPContext` object used to
+            solve the Dual PEP.
+
+    """
+
+    dual_opt_value: float
+    primal_var_manager: PrimalVariableManager
+    solver_status: Any
+    context: pc.PEPContext
+
+
 class PEPBuilder:
-    """The main class for PEP primal formulation."""
+    """The main class for Primal and Dual PEP formulation.
+
+    Attributes:
+        init_conditions (list[:class:`Constraint`]): A list of all the initial
+            conditions associated with this PEP.
+        functions (list[:class:`Function`]): A list of all the functions
+            associated with this PEP.
+        performance_metric (:class:`Scalar`): The performance metric for this
+            PEP.
+        relaxed_constraints (list[str]): A list of names of the constraints
+            that will be ignored when the Primal or Dual PEP is constructed.
+        dual_val_constraint (dict[str, list[tuple[str, float]]]): A dictionary
+            of the form `{constraint_name: [op, val]}`. The `constraint_name`
+            is the name of the constraint the dual variable is associated with.
+            The `op` is a string for the type of relation, i.e., `lt`, `gt`,
+            or `eq`. The `val` is the value for the other side of the
+            constraint. For example, consider `{"f:x_1,x_0", [("eq", 0)]}`.
+            Denote the associated dual variable as :math:`\\lambda_{1,0}`.
+            Then, this means to add a constraint of the form
+            :math:`\\lambda_{1,0} = 0` to the Dual PEP. Because it is hard to
+            judge if the constraint associated with `constraint_name` is
+            active, we suggest to not add dual variable constraints manually
+            but instead use the interactive dashboard.
+    """
 
     def __init__(self):
         self.pep_context_dict: dict[str, pc.PEPContext] = {}
 
-        self.init_conditions = []  #: list["constraint"] =[]
-        self.functions = []  #: list["function"] = []
-        self.interpolation_constraints = []  #: list["constraint"] = []
-        self.performance_metric = None  # scalar
+        self.init_conditions: list[Constraint] = []
+        self.functions: list[Function] = []
+        self.performance_metric: sc.Scalar | None = None
 
         # Contain the name for the constraints that should be removed.
         # We should think about a better choice like manager.
-        self.relaxed_constraints = []
+        self.relaxed_constraints: list[str] = []
+
+        # `dual_val_constraint` has the data structure: {constraint_name: [op, val]}.
+        # Because it is hard to judge if the dual_val_constraint is applied or not,
+        # we recommend that do not use manually but through the interactive dashboard.
+        self.dual_val_constraint: dict[str, list[tuple[str, float]]] = defaultdict(list)
 
     def clear_setup(self):
+        """Resets the :class:`PEPBuilder` object."""
         self.init_conditions.clear()
         self.functions.clear()
-        self.interpolation_constraints.clear()
         self.performance_metric = None
         self.relaxed_constraints.clear()
+        self.dual_val_constraint.clear()
 
     @contextlib.contextmanager
     def make_context(
@@ -130,28 +218,116 @@ class PEPBuilder:
         return point
 
     def set_initial_constraint(self, constraint):
+        """
+        Set an initial condition.
+
+        Args:
+            constraint (:class:`Constraint`): A :class:`Constraint` object that
+                represents the desired initial condition.
+        """
         self.init_conditions.append(constraint)
 
     def set_performance_metric(self, metric: sc.Scalar):
+        """
+        Set the performance metric.
+
+        Args:
+            metric (:class:`Scalar`): A :class:`Scalar` object that
+                represents the desired performance metric.
+        """
         self.performance_metric = metric
 
     def set_relaxed_constraints(self, relaxed_constraints: list[str]):
+        """
+        Set the constraints that will be ignored.
+
+        Args:
+            relaxed_constraints (list[str]): A list of names of constraints
+                that will be ignored.
+        """
         self.relaxed_constraints.extend(relaxed_constraints)
 
+    def add_dual_val_constraint(
+        self, constraint_name: str, op: str, val: float
+    ) -> None:
+        if op not in ["lt", "gt", "eq"]:
+            raise ValueError(f"op must be one of `lt`, `gt`, or `eq` but get {op}")
+        if not utils.is_numerical(val):
+            raise ValueError("Value must be some numerical value.")
+
+        self.dual_val_constraint[constraint_name].append((op, val))
+
     def declare_func(self, function_class: type[Function], tag: str, **kwargs):
+        """
+        Declare a function.
+
+        Args:
+            function_class (type[:class:`Function`]): The type of function we want to
+                declare. Examples include :class:`ConvexFunction` or
+                :class:`SmoothConvexFunction`.
+            tag (str): A tag that will be added to the :class:`Function`'s
+                `tags` list. It can be used to identify the :class:`Function`
+                object.
+            **kwargs: The other parameters needed to declare the function. For
+                example, :class:`SmoothConvexFunction` will require a
+                smoothness parameter `L`.
+        """
         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):
+        """
+        Return the :class:`Function` object associated with the provided `tag`.
+
+        Args:
+            tag (str): The `tag` of the :class:`Function` object we want to
+                retrieve.
+
+        Returns:
+            :class:`Function`: The :class:`Function` object associated with
+            the `tag`.
+
+        Note:
+            Currently, only basis :class:`Function` objects can be retrieved.
+            This will be updated eventually.
+        """
         # 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):
+    def solve(
+        self,
+        context: pc.PEPContext | None = None,
+        resolve_parameters: dict[str, utils.NUMERICAL_TYPE] | None = None,
+    ):
+        return self.solve_primal(context, resolve_parameters=resolve_parameters)
+
+    def solve_primal(
+        self,
+        context: pc.PEPContext | None = None,
+        resolve_parameters: dict[str, utils.NUMERICAL_TYPE] | None = None,
+    ):
+        """
+        Solve the Primal PEP associated with this :class:`PEPBuilder` object
+        using the given :class:`PEPContext` object.
+
+        Args:
+            context (:class:`PEPContext`): The :class:`PEPContext` object used
+                to solve the Primal PEP associated with this
+                :class:`PEPBuilder` object. `None` if we consider the current
+                global :class:`PEPContext` object.
+            resolve_parameters (dict[str, :class:`NUMERICAL_TYPE`]): A dictionary that
+                maps the name of parameters to the numerical values.
+
+        Returns:
+            :class:`PEPResult`: A :class:`PEPResult` object that contains the
+            information obtained after solving the Primal PEP associated with
+            this :class:`PEPBuilder` object.
+        """
         if context is None:
             context = pc.get_current_context()
         if context is None:
@@ -159,22 +335,86 @@ class PEPBuilder:
 
         all_constraints: list[Constraint] = [*self.init_conditions]
         for f in self.functions:
-            all_constraints.extend(f.get_interpolation_constraints())
+            all_constraints.extend(f.get_interpolation_constraints(context))
 
         # for now, we heavily rely on the CVX. We can make a wrapper class to avoid
         # direct dependency in the future.
-        solver = ps.CVXSolver(
+        solver = ps.CVXPrimalSolver(
             perf_metric=self.performance_metric,
             constraints=[
                 c for c in all_constraints if c.name not in self.relaxed_constraints
             ],
             context=context,
         )
-        problem = solver.build_problem()
-        result = problem.solve(**kwargs)
+        problem = solver.build_problem(resolve_parameters=resolve_parameters)
+        result = problem.solve()
         return PEPResult(
             primal_opt_value=result,
             dual_var_manager=solver.dual_var_manager,
             solver_status=problem.status,
             context=context,
         )
+
+    def solve_dual(
+        self,
+        context: pc.PEPContext | None = None,
+        resolve_parameters: dict[str, utils.NUMERICAL_TYPE] | None = None,
+    ):
+        """
+        Solve the Dual PEP associated with this :class:`PEPBuilder` object
+        using the given :class:`PEPContext` object.
+
+        Args:
+            context (:class:`PEPContext`): The :class:`PEPContext` object used
+                to solve the Dual PEP associated with this :class:`PEPBuilder`
+                object. `None` if we consider the current global
+                :class:`PEPContext` object.
+            resolve_parameters (dict[str, :class:`NUMERICAL_TYPE`]): A dictionary that
+                maps the name of parameters to the numerical values.
+
+        Returns:
+            :class:`DualPEPResult`: A :class:`DualPEPResult` object that
+            contains the information obtained after solving the Dual PEP
+            associated with this :class:`PEPBuilder` object.
+        """
+        if context is None:
+            context = pc.get_current_context()
+        if context is None:
+            raise RuntimeError("Did you forget to create a context?")
+
+        all_constraints: list[Constraint] = [*self.init_conditions]
+        for f in self.functions:
+            all_constraints.extend(f.get_interpolation_constraints(context))
+
+        # TODO: Consider a better API and interface to adding constraint for primal
+        # variable of dual problem. We can add `extra_primal_val_constraints` to add
+        # more constraints on primal var in dual PEP, i.e. dual var in primal PEP.
+        constraints = []
+        for c in all_constraints:
+            if c.name in self.relaxed_constraints:
+                continue
+            for op, val in self.dual_val_constraint[c.name]:
+                if op == "lt":
+                    c.dual_lt(val)
+                elif op == "gt":
+                    c.dual_gt(val)
+                elif op == "eq":
+                    c.dual_eq(val)
+                else:
+                    raise ValueError(f"Unknown op when construct the {c}")
+            constraints.append(c)
+
+        dual_solver = ps.CVXDualSolver(
+            perf_metric=self.performance_metric,
+            constraints=constraints,
+            context=context,
+        )
+        problem = dual_solver.build_problem(resolve_parameters=resolve_parameters)
+        result = problem.solve()
+
+        return DualPEPResult(
+            dual_opt_value=result,
+            primal_var_manager=dual_solver.primal_var_manager,
+            solver_status=problem.status,
+            context=context,
+        )

pepflow/pep_context.py

@@ -37,16 +37,37 @@ 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] = []
@@ -58,6 +79,12 @@ class PEPContext:
         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
 
@@ -74,6 +101,18 @@ class PEPContext:
         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
@@ -83,22 +122,73 @@ class PEPContext:
         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.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
@@ -102,3 +103,23 @@ def test_get_by_tag(pep_context: pc.PEPContext):
     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]