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

pepflow/__init__.py

@@ -21,6 +21,7 @@
 from .constants import PSD_CONSTRAINT as PSD_CONSTRAINT
 from .constraint import Constraint as Constraint
 from .expression_manager import ExpressionManager as ExpressionManager
+from .expression_manager import represent_matrix_by_basis as represent_matrix_by_basis
 
 # interactive_constraint
 from .interactive_constraint import launch as launch
@@ -28,6 +29,7 @@ from .interactive_constraint import launch as launch
 # pep
 from .pep import PEPBuilder as PEPBuilder
 from .pep import PEPResult as PEPResult
+from .pep import DualPEPResult as DualPEPResult
 from .pep_context import PEPContext as PEPContext
 from .pep_context import get_current_context as get_current_context
 from .pep_context import set_current_context as set_current_context
@@ -43,8 +45,10 @@ from .scalar import EvaluatedScalar as EvaluatedScalar
 from .scalar import Scalar as Scalar
 
 # Solver
-from .solver import CVXSolver as CVXSolver
+from .solver import CVXPrimalSolver as CVXPrimalSolver
+from .solver import CVXDualSolver as CVXDualSolver
 from .solver import DualVariableManager as DualVariableManager
+from .solver import PrimalVariableManager as PrimalVariableManager
 
 # Others
 from .utils import SOP as SOP

pepflow/constraint.py

@@ -31,8 +31,65 @@ from pepflow import utils
 
 @attrs.frozen
 class Constraint:
-    """It represents `expression relation 0`."""
+    """A :class:`Constraint` object that represents inequalities and
+    equalities of :class:`Scalar` objects.
+
+    Denote an arbitrary :class:`Scalar` object as `x`. Constraints represent:
+    `x <= 0`, `x >= 0`, and `x = 0`.
+
+    Attributes:
+        scalar (:class:`Scalar`): The :class:`Scalar` object involved in
+            the inequality or equality.
+        comparator (:class:`Comparator`): :class:`Comparator` is an enumeration
+            that can be either `GT`, `LT`, or `EQ`. They represent `>=`, `<=`,
+            and `=` respectively.
+        name (str): The unique name of the :class:`Comparator` object.
+        associated_dual_var_constraints (list[tuple[utils.Comparator, float]]):
+            A list of all the constraints imposed on the associated dual
+            variable of this :class:`Constraint` object.
+    """
 
     scalar: Scalar | float
     comparator: utils.Comparator
     name: str
+
+    # Used to represent the constraint on primal variable in dual PEP.
+    associated_dual_var_constraints: list[tuple[utils.Comparator, float]] = attrs.field(
+        factory=list
+    )
+
+    def dual_lt(self, val: float) -> None:
+        """
+        Denote the associated dual variable of this constraint as `lambd`.
+        This generates a relation of the form `lambd <= val`.
+
+        Args:
+            val (float): The other object in the relation.
+        """
+        if not utils.is_numerical(val):
+            raise ValueError(f"The input {val=} must be a numerical value")
+        self.associated_dual_var_constraints.append((utils.Comparator.LT, val))
+
+    def dual_gt(self, val: float) -> None:
+        """
+        Denote the associated dual variable of this constraint as `lambd`.
+        This generates a relation of the form `lambd >= val`.
+
+        Args:
+            val (float): The other object in the relation.
+        """
+        if not utils.is_numerical(val):
+            raise ValueError(f"The input {val=} must be a numerical value")
+        self.associated_dual_var_constraints.append((utils.Comparator.GT, val))
+
+    def dual_eq(self, val: float) -> None:
+        """
+        Denote the associated dual variable of this constraint as `lambd`.
+        This generates a relation of the form `lambd = val`.
+
+        Args:
+            val (float): The other object in the relation.
+        """
+        if not utils.is_numerical(val):
+            raise ValueError(f"The input {val=} must be a numerical value")
+        self.associated_dual_var_constraint.append((utils.Comparator.EQ, val))

pepflow/e2e_test.py

@@ -1,6 +1,8 @@
 import math
 
-from pepflow import function, pep
+from pepflow import function
+from pepflow import parameter as pm
+from pepflow import pep
 from pepflow import pep_context as pc
 
 
@@ -28,10 +30,47 @@ def test_gd_e2e():
         pep_builder.set_performance_metric(
             f.function_value(p) - f.function_value(x_star)
         )
-        result = pep_builder.solve()
+        result = pep_builder.solve_primal()
         expected_opt_value = 1 / (4 * i + 2)
         assert math.isclose(result.primal_opt_value, expected_opt_value, rel_tol=1e-3)
 
+        dual_result = pep_builder.solve_dual()
+        assert math.isclose(
+            dual_result.dual_opt_value, expected_opt_value, rel_tol=1e-3
+        )
+
+
+def test_gd_diff_stepsize_e2e():
+    pc.PEPContext("gd").set_as_current()
+    pep_builder = pep.PEPBuilder()
+    eta = 1 / pm.Parameter(name="L")
+    N = 4
+
+    f = pep_builder.declare_func(
+        function.SmoothConvexFunction, "f", L=pm.Parameter(name="L")
+    )
+    x = pep_builder.set_init_point("x_0")
+    x_star = f.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):
+        x = x - eta * f.gradient(x)
+        x.add_tag(f"x_{i + 1}")
+    pep_builder.set_performance_metric(f(x) - f(x_star))
+
+    for l_val in [1, 4, 0.25]:
+        result = pep_builder.solve_primal(resolve_parameters={"L": l_val})
+        expected_opt_value = l_val / (4 * N + 2)
+        assert math.isclose(result.primal_opt_value, expected_opt_value, rel_tol=1e-3)
+
+        dual_result = pep_builder.solve_dual(resolve_parameters={"L": l_val})
+        assert math.isclose(
+            dual_result.dual_opt_value, expected_opt_value, rel_tol=1e-3
+        )
+
 
 def test_pgm_e2e():
     ctx = pc.PEPContext("pgm").set_as_current()
@@ -64,6 +103,11 @@ def test_pgm_e2e():
             h.function_value(p) - h.function_value(x_star)
         )
 
-        result = pep_builder.solve()
+        result = pep_builder.solve_primal()
         expected_opt_value = 1 / (4 * i)
         assert math.isclose(result.primal_opt_value, expected_opt_value, rel_tol=1e-3)
+
+        dual_result = pep_builder.solve_dual()
+        assert math.isclose(
+            dual_result.dual_opt_value, expected_opt_value, rel_tol=1e-3
+        )

pepflow/expression_manager.py

@@ -22,6 +22,7 @@ import math
 
 import numpy as np
 
+from pepflow import parameter as pm
 from pepflow import pep_context as pc
 from pepflow import point as pt
 from pepflow import scalar as sc
@@ -29,23 +30,41 @@ from pepflow import utils
 
 
 def tag_and_coef_to_str(tag: str, v: float) -> str:
-    coef = f"{abs(v):.3g}"
+    coef = utils.numerical_str(abs(v))
     sign = "+" if v >= 0 else "-"
     if math.isclose(abs(v), 1):
         return f"{sign} {tag} "
-    elif math.isclose(v, 0):
+    elif math.isclose(v, 0, abs_tol=1e-5):
         return ""
     else:
         return f"{sign} {coef}*{tag} "
 
 
 class ExpressionManager:
-    def __init__(self, pep_context: pc.PEPContext):
+    """
+    A class that handles the concrete representations of abstract
+    :class:`Point` and :class:`Scalar` objects managed by a particular
+    :class:`PEPContext` object.
+
+    Attributes:
+        context (:class:`PEPContext`): The :class:`PEPContext` object which
+            manages the abstract :class:`Point` and :class:`Scalar` objects
+            of interest.
+        resolve_parameters (dict[str, :class:`NUMERICAL_TYPE`]): A dictionary that
+            maps the name of parameters to the numerical values.
+    """
+
+    def __init__(
+        self,
+        pep_context: pc.PEPContext,
+        resolve_parameters: dict[str, utils.NUMERICAL_TYPE] | None = None,
+    ):
         self.context = pep_context
         self._basis_points = []
         self._basis_point_uid_to_index = {}
         self._basis_scalars = []
         self._basis_scalar_uid_to_index = {}
+        self.resolve_parameters = resolve_parameters or {}
         for point in self.context.points:
             if point.is_basis:
                 self._basis_points.append(point)
@@ -74,88 +93,161 @@ class ExpressionManager:
 
     @functools.cache
     def eval_point(self, point: pt.Point | float | int):
+        """
+        Return the concrete representation of the given :class:`Point`,
+        `float`, or `int`. Concrete representations of :class:`Point` objects
+        are :class:`EvaluatedPoint` objects. Concrete representations of
+        `float` or `int` arguments are themselves.
+
+        Args:
+            point (:class:`Point`, float, int): The abstract :class:`Point`,
+                `float`, or `int` object whose concrete representation we want
+                to find.
+
+        Returns:
+            :class:`EvaluatedPoint` | float | int: The concrete representation
+            the `point` argument.
+        """
         if utils.is_numerical(point):
             return point
 
-        array = np.zeros(self._num_basis_points)
+        if isinstance(point, pm.Parameter):
+            return point.get_value(self.resolve_parameters)
+
         if point.is_basis:
             index = self.get_index_of_basis_point(point)
+            array = np.zeros(self._num_basis_points)
             array[index] = 1
             return pt.EvaluatedPoint(vector=array)
 
+        if isinstance(point.eval_expression, pt.ZeroPoint):
+            return pt.EvaluatedPoint.zero(num_basis_points=self._num_basis_points)
+
         op = point.eval_expression.op
+        left_evaled_point = self.eval_point(point.eval_expression.left_point)
+        right_evaled_point = self.eval_point(point.eval_expression.right_point)
         if op == utils.Op.ADD:
-            return self.eval_point(point.eval_expression.left_point) + self.eval_point(
-                point.eval_expression.right_point
-            )
+            return left_evaled_point + right_evaled_point
         if op == utils.Op.SUB:
-            return self.eval_point(point.eval_expression.left_point) - self.eval_point(
-                point.eval_expression.right_point
-            )
+            return left_evaled_point - right_evaled_point
         if op == utils.Op.MUL:
-            return self.eval_point(point.eval_expression.left_point) * self.eval_point(
-                point.eval_expression.right_point
-            )
+            return left_evaled_point * right_evaled_point
         if op == utils.Op.DIV:
-            return self.eval_point(point.eval_expression.left_point) / self.eval_point(
-                point.eval_expression.right_point
-            )
+            return left_evaled_point / right_evaled_point
 
-        raise ValueError("This should never happen!")
+        raise ValueError(f"Encountered unknown {op=} when evaluation the point.")
 
     @functools.cache
     def eval_scalar(self, scalar: sc.Scalar | float | int):
+        """
+        Return the concrete representation of the given :class:`Scalar`,
+        `float`, or `int`. Concrete representations of :class:`Scalar` objects
+        are :class:`EvaluatedScalar` objects. Concrete representations of
+        `float` or `int` arguments are themselves.
+
+        Args:
+            scalar (:class:`Point`, float, int): The abstract :class:`Scalar`,
+                `float`, or `int` object whose concrete representation we want
+                to find.
+
+        Returns:
+            :class:`EvaluatedScalar` | float | int: The concrete representation
+            the `scalar` argument.
+        """
         if utils.is_numerical(scalar):
             return scalar
+        if isinstance(scalar, pm.Parameter):
+            return scalar.get_value(self.resolve_parameters)
 
-        array = np.zeros(self._num_basis_scalars)
         if scalar.is_basis:
             index = self.get_index_of_basis_scalar(scalar)
+            array = np.zeros(self._num_basis_scalars)
             array[index] = 1
+            matrix = np.zeros((self._num_basis_points, self._num_basis_points))
             return sc.EvaluatedScalar(
                 vector=array,
-                matrix=np.zeros((self._num_basis_points, self._num_basis_points)),
+                matrix=matrix,
                 constant=float(0.0),
             )
+
+        if isinstance(scalar.eval_expression, sc.ZeroScalar):
+            return sc.EvaluatedScalar.zero(
+                num_basis_scalars=self._num_basis_scalars,
+                num_basis_points=self._num_basis_points,
+            )
+
         op = scalar.eval_expression.op
+        # The special inner product usage.
+        if (
+            op == utils.Op.MUL
+            and isinstance(scalar.eval_expression.left_scalar, pt.Point)
+            and isinstance(scalar.eval_expression.right_scalar, pt.Point)
+        ):
+            array = np.zeros(self._num_basis_scalars)
+            return sc.EvaluatedScalar(
+                vector=array,
+                matrix=utils.SOP(
+                    self.eval_point(scalar.eval_expression.left_scalar).vector,
+                    self.eval_point(scalar.eval_expression.right_scalar).vector,
+                ),
+                constant=float(0.0),
+            )
+
+        left_evaled_scalar = self.eval_scalar(scalar.eval_expression.left_scalar)
+        right_evaled_scalar = self.eval_scalar(scalar.eval_expression.right_scalar)
         if op == utils.Op.ADD:
-            return self.eval_scalar(
-                scalar.eval_expression.left_scalar
-            ) + self.eval_scalar(scalar.eval_expression.right_scalar)
+            return left_evaled_scalar + right_evaled_scalar
         if op == utils.Op.SUB:
-            return self.eval_scalar(
-                scalar.eval_expression.left_scalar
-            ) - self.eval_scalar(scalar.eval_expression.right_scalar)
+            return left_evaled_scalar - right_evaled_scalar
         if op == utils.Op.MUL:
-            if isinstance(scalar.eval_expression.left_scalar, pt.Point) and isinstance(
-                scalar.eval_expression.right_scalar, pt.Point
-            ):
-                return sc.EvaluatedScalar(
-                    vector=np.zeros(self._num_basis_scalars),
-                    matrix=utils.SOP(
-                        self.eval_point(scalar.eval_expression.left_scalar).vector,
-                        self.eval_point(scalar.eval_expression.right_scalar).vector,
-                    ),
-                    constant=float(0.0),
-                )
-            else:
-                return self.eval_scalar(
-                    scalar.eval_expression.left_scalar
-                ) * self.eval_scalar(scalar.eval_expression.right_scalar)
+            return left_evaled_scalar * right_evaled_scalar
         if op == utils.Op.DIV:
-            return self.eval_scalar(
-                scalar.eval_expression.left_scalar
-            ) / self.eval_scalar(scalar.eval_expression.right_scalar)
+            return left_evaled_scalar / right_evaled_scalar
 
-        raise ValueError("This should never happen!")
+        raise ValueError(f"Encountered unknown {op=} when evaluation the scalar.")
 
     @functools.cache
     def repr_point_by_basis(self, point: pt.Point) -> str:
+        """
+        Express the given :class:`Point` object as the linear combination of
+        the basis :class:`Point` objects of the :class:`PEPContext` associated
+        with this :class:`ExpressionManager`. This linear combination is
+        expressed as a `str` where, to refer to the basis :class:`Point`
+        objects, we use their tags.
+
+        Args:
+            point (:class:`Point`): The :class:`Point` object which we want
+                to express in terms of the basis :class:`Point` objects.
+
+        Returns:
+            str: The representation of `point` in terms of the basis
+            :class:`Point` objects of the :class:`PEPContext` associated
+            with this :class:`ExpressionManager`.
+        """
         assert isinstance(point, pt.Point)
-        repr_array = self.eval_point(point).vector
+        evaluated_point = self.eval_point(point)
+        return self.repr_evaluated_point_by_basis(evaluated_point)
+
+    def repr_evaluated_point_by_basis(self, evaluated_point: pt.EvaluatedPoint) -> str:
+        """
+        Express the given :class:`EvaluatedPoint` object as the linear
+        combination of the basis :class:`Point` objects of the
+        :class:`PEPContext` associated with this :class:`ExpressionManager`.
+        This linear combination is expressed as a `str` where, to refer to the
+        basis :class:`Point` objects, we use their tags.
 
+        Args:
+            evaluated_point (:class:`EvaluatedPoint`): The
+                :class:`EvaluatedPoint` object which we want to express in
+                terms of the basis :class:`Point` objects.
+
+        Returns:
+            str: The representation of `evaluated_point` in terms of
+            the basis :class:`Point` objects of the :class:`PEPContext`
+            associated with this :class:`ExpressionManager`.
+        """
         repr_str = ""
-        for i, v in enumerate(repr_array):
+        for i, v in enumerate(evaluated_point.vector):
             ith_tag = self.get_tag_of_basis_point_index(i)
             repr_str += tag_and_coef_to_str(ith_tag, v)
 
@@ -169,28 +261,113 @@ class ExpressionManager:
         return repr_str.strip()
 
     @functools.cache
-    def repr_scalar_by_basis(self, scalar: sc.Scalar) -> str:
+    def repr_scalar_by_basis(
+        self, scalar: sc.Scalar, greedy_square: bool = True
+    ) -> str:
+        """Express the given :class:`Scalar` object in terms of the basis
+        :class:`Point` and :class:`Scalar` objects of the :class:`PEPContext`
+        associated with this :class:`ExpressionManager`.
+
+        A :class:`Scalar` can be formed by linear combinations of basis
+        :class:`Scalar` objects. A :class:`Scalar` can also be formed through
+        the inner product of two basis :class:`Point` objects. This function
+        returns the representation of this :class:`Scalar` object in terms of
+        the basis :class:`Point` and :class:`Scalar` objects as a `str` where,
+        to refer to the basis :class:`Point` and :class:`Scalar` objects,
+        we use their tags.
+
+        Args:
+            scalar (:class:`Scalar`): The :class:`Scalar` object which we want
+                to express in terms of the basis :class:`Point` and
+                :class:`Scalar` objects.
+            greedy_square (bool): If `greedy_square` is true, the function will
+                try to return :math:`\\|a-b\\|^2` whenever possible. If not,
+                the function will return
+                :math:`\\|a\\|^2 - 2 * \\langle a, b \\rangle + \\|b\\|^2` instead.
+                `True` by default.
+
+        Returns:
+            str: The representation of `scalar` in terms of the basis
+            :class:`Point` and :class:`Scalar` objects of the
+            :class:`PEPContext` associated with this
+            :class:`ExpressionManager`.
+        """
         assert isinstance(scalar, sc.Scalar)
         evaluated_scalar = self.eval_scalar(scalar)
+        return self.repr_evaluated_scalar_by_basis(
+            evaluated_scalar, greedy_square=greedy_square
+        )
+
+    def repr_evaluated_scalar_by_basis(
+        self, evaluated_scalar: sc.EvaluatedScalar, greedy_square: bool = True
+    ) -> str:
+        """Express the given :class:`EvaluatedScalar` object in terms of the
+        basis :class:`Point` and :class:`Scalar` objects of the
+        :class:`PEPContext` associated with this :class:`ExpressionManager`.
 
+        A :class:`Scalar` can be formed by linear combinations of basis
+        :class:`Scalar` objects. A :class:`Scalar` can also be formed through
+        the inner product of two basis :class:`Point` objects. This function
+        returns the representation of this :class:`Scalar` object in terms of
+        the basis :class:`Point` and :class:`Scalar` objects as a `str` where,
+        to refer to the basis :class:`Point` and :class:`Scalar` objects,
+        we use their tags.
+
+        Args:
+            evaluated_scalar (:class:`EvaluatedScalar`): The
+                :class:`EvaluatedScalar` object which we want to express in
+                terms of the basis :class:`Point` and :class:`Scalar` objects.
+            greedy_square (bool): If `greedy_square` is true, the function will
+                try to return :math:`\\|a-b\\|^2` whenever possible. If not,
+                the function will return
+                :math:`\\|a\\|^2 - 2 * \\langle a, b \\rangle + \\|b\\|^2` instead.
+                `True` by default.
+
+        Returns:
+            str: The representation of `evaluated_scalar` in terms of
+            the basis :class:`Point` and :class:`Scalar` objects of the
+            :class:`PEPContext` associated with this :class:`ExpressionManager`.
+        """
         repr_str = ""
-        if not math.isclose(evaluated_scalar.constant, 0):
-            repr_str += f"{evaluated_scalar.constant:.3g}"
+        if not math.isclose(evaluated_scalar.constant, 0, abs_tol=1e-5):
+            repr_str += utils.numerical_str(evaluated_scalar.constant)
 
         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]):
+        if greedy_square:
+            diag_elem = np.diag(evaluated_scalar.matrix).copy()
+            for i in range(evaluated_scalar.matrix.shape[0]):
+                ith_tag = self.get_tag_of_basis_point_index(i)
+                # j starts from i+1 since we want to handle the diag elem at last.
+                for j in range(i + 1, evaluated_scalar.matrix.shape[0]):
+                    jth_tag = self.get_tag_of_basis_point_index(j)
+                    v = float(evaluated_scalar.matrix[i, j])
+                    # We want to minimize the diagonal elements to zero greedily.
+                    if diag_elem[i] * v > 0:  # same sign with diagonal elem
+                        diag_elem[i] -= v
+                        diag_elem[j] -= v
+                        repr_str += tag_and_coef_to_str(f"|{ith_tag}+{jth_tag}|^2", v)
+                    else:  # different sign
+                        diag_elem[i] += v
+                        diag_elem[j] += v
+                        repr_str += tag_and_coef_to_str(f"|{ith_tag}-{jth_tag}|^2", -v)
+                # Handle the diagonal elements
+                repr_str += tag_and_coef_to_str(f"|{ith_tag}|^2", diag_elem[i])
+        else:
+            for i in range(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)
+                for j in range(i, evaluated_scalar.matrix.shape[0]):
+                    jth_tag = self.get_tag_of_basis_point_index(j)
+                    v = float(evaluated_scalar.matrix[i, j])
+                    if i == j:
+                        repr_str += tag_and_coef_to_str(f"|{ith_tag}|^2", v)
+                    else:
+                        repr_str += tag_and_coef_to_str(
+                            f"<{ith_tag}, {jth_tag}>", 2 * v
+                        )
 
         # Post processing
         if repr_str == "":
@@ -200,3 +377,41 @@ class ExpressionManager:
         if repr_str.startswith("- "):
             repr_str = "-" + repr_str[2:]
         return repr_str.strip()
+
+
+def represent_matrix_by_basis(matrix: np.ndarray, ctx: pc.PEPContext) -> str:
+    """Express the given matrix in terms of the basis :class:`Point` objects
+    of the given :class:`PEPContext`.
+
+    The concrete representation of the inner product of two abstract
+    basis :class:`Point` objects is a matrix (the outer product of the
+    basis vectors corresponding to the concrete representations of the abstract
+    basis :class:`Point` objects). The given matrix can then be expressed
+    as a linear combination of the inner products of abstract basis
+    :class:`Point` objects. This is provided as a `str` where, to refer to
+    the basis :class:`Point` objects, we use their tags.
+
+    Args:
+        matrix (np.ndarray): The matrix which we want to express in terms of
+            the basis :class:`Point` objects of the given :class:`PEPContext`.
+        ctx (:class:`PEPContext`): The :class:`PEPContext` whose basis
+            :class:`Point` objects we consider.
+
+    Returns:
+        str: The representation of `matrix` in terms of the basis
+        :class:`Point` objects of `ctx`.
+    """
+    em = ExpressionManager(ctx)
+    matrix_shape = (len(em._basis_points), len(em._basis_points))
+    if matrix.shape != matrix_shape:
+        raise ValueError(
+            "The valid matrix for given context should have shape {matrix_shape}"
+        )
+    if not np.allclose(matrix, matrix.T):
+        raise ValueError("Input matrix must be symmetric.")
+
+    return em.repr_evaluated_scalar_by_basis(
+        sc.EvaluatedScalar(
+            vector=np.zeros(len(em._basis_scalars)), matrix=matrix, constant=0.0
+        )
+    )

pepflow/expression_manager_test.py

@@ -97,7 +97,28 @@ def test_repr_scalar_by_basis(pep_context: pc.PEPContext) -> None:
 
     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)>"
+    assert (
+        em.repr_scalar_by_basis(s, greedy_square=False) == "f(x) + <x, gradient_f(x)>"
+    )
+    assert (
+        em.repr_scalar_by_basis(s)
+        == "f(x) - 0.5*|x-gradient_f(x)|^2 + 0.5*|x|^2 + 0.5*|gradient_f(x)|^2"
+    )
+
+
+def test_repr_scalar_by_basis2(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, greedy_square=False) == "f(x) - <x, gradient_f(x)>"
+    )
+    assert (
+        em.repr_scalar_by_basis(s)
+        == "f(x) + 0.5*|x-gradient_f(x)|^2 - 0.5*|x|^2 - 0.5*|gradient_f(x)|^2"
+    )
 
 
 def test_repr_scalar_by_basis_interpolation(pep_context: pc.PEPContext) -> None:
@@ -110,7 +131,20 @@ def test_repr_scalar_by_basis_interpolation(pep_context: pc.PEPContext) -> None:
     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
+    assert em.repr_scalar_by_basis(interp_scalar, greedy_square=False) == expected_repr
+    expected_square_repr = "-f(x_i) + f(x_j) - 0.5*|x_i-gradient_f(x_j)|^2 + 0.5*|x_i|^2 + 0.5*|x_j-gradient_f(x_j)|^2 - 0.5*|x_j|^2 + 0.5*|gradient_f(x_i)-gradient_f(x_j)|^2"
+    assert em.repr_scalar_by_basis(interp_scalar) == expected_square_repr
 
 
 # TODO add more tests about repr_scalar_by_basis
+
+
+def test_represent_matrix_by_basis(pep_context: pc.PEPContext) -> None:
+    _ = pt.Point(is_basis=True, tags=["x_1"])
+    _ = pt.Point(is_basis=True, tags=["x_2"])
+    _ = pt.Point(is_basis=True, tags=["x_3"])
+    matrix = np.array([[0.5, 0.5, 0], [0.5, 2, 0], [0, 0, 3]])
+    assert (
+        exm.represent_matrix_by_basis(matrix, pep_context)
+        == "0.5*|x_1+x_2|^2 + 1.5*|x_2|^2 + 3*|x_3|^2"
+    )