pepflow 0.1.0__py3-none-any.whl

pepflow/constraint.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import attrs
+
+if TYPE_CHECKING:
+    from pepflow.point import Scalar
+
+from pepflow import utils
+
+
+@attrs.frozen
+class Constraint:
+    """It represents `expression relation 0`."""
+
+    scalar: Scalar | float
+    comparator: utils.Comparator
+    name: str

pepflow/expression_manager.py

@@ -0,0 +1,113 @@
+import functools
+
+import numpy as np
+
+from pepflow import pep_context as pc
+from pepflow import point as pt
+from pepflow import scalar as sc
+from pepflow import utils
+
+
+class ExpressionManager:
+    def __init__(self, pep_context: pc.PEPContext):
+        self.context = pep_context
+        self._basis_points = []
+        self._basis_point_uid_to_index = {}
+        self._basis_scalars = []
+        self._basis_scalar_uid_to_index = {}
+        for point in self.context.points:
+            if point.is_basis:
+                self._basis_points.append(point)
+                self._basis_point_uid_to_index[point.uid] = len(self._basis_points) - 1
+        for scalar in self.context.scalars:
+            if scalar.is_basis:
+                self._basis_scalars.append(scalar)
+                self._basis_scalar_uid_to_index[scalar.uid] = (
+                    len(self._basis_scalars) - 1
+                )
+
+        self._num_basis_points = len(self._basis_points)
+        self._num_basis_scalars = len(self._basis_scalars)
+
+    def get_index_of_basis_point(self, point: pt.Point):
+        return self._basis_point_uid_to_index[point.uid]
+
+    def get_index_of_basis_scalar(self, scalar: sc.Scalar):
+        return self._basis_scalar_uid_to_index[scalar.uid]
+
+    @functools.cache
+    def eval_point(self, point: pt.Point | float | int):
+        if utils.is_numerical(point):
+            return point
+
+        array = np.zeros(self._num_basis_points)
+        if point.is_basis:
+            index = self.get_index_of_basis_point(point)
+            array[index] = 1
+            return pt.EvaluatedPoint(vector=array)
+
+        op = point.eval_expression.op
+        if op == utils.Op.ADD:
+            return self.eval_point(point.eval_expression.left_point) + self.eval_point(
+                point.eval_expression.right_point
+            )
+        if op == utils.Op.SUB:
+            return self.eval_point(point.eval_expression.left_point) - self.eval_point(
+                point.eval_expression.right_point
+            )
+        if op == utils.Op.MUL:
+            return self.eval_point(point.eval_expression.left_point) * self.eval_point(
+                point.eval_expression.right_point
+            )
+        if op == utils.Op.DIV:
+            return self.eval_point(point.eval_expression.left_point) / self.eval_point(
+                point.eval_expression.right_point
+            )
+
+        raise ValueError("This should never happen!")
+
+    @functools.cache
+    def eval_scalar(self, scalar: sc.Scalar | float | int):
+        if utils.is_numerical(scalar):
+            return scalar
+
+        array = np.zeros(self._num_basis_scalars)
+        if scalar.is_basis:
+            index = self.get_index_of_basis_scalar(scalar)
+            array[index] = 1
+            return sc.EvaluatedScalar(
+                vector=array,
+                matrix=np.zeros((self._num_basis_points, self._num_basis_points)),
+                constant=float(0.0),
+            )
+        op = scalar.eval_expression.op
+        if op == utils.Op.ADD:
+            return self.eval_scalar(
+                scalar.eval_expression.left_scalar
+            ) + self.eval_scalar(scalar.eval_expression.right_scalar)
+        if op == utils.Op.SUB:
+            return self.eval_scalar(
+                scalar.eval_expression.left_scalar
+            ) - self.eval_scalar(scalar.eval_expression.right_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)
+        if op == utils.Op.DIV:
+            return self.eval_scalar(
+                scalar.eval_expression.left_scalar
+            ) / self.eval_scalar(scalar.eval_expression.right_scalar)
+
+        raise ValueError("This should never happen!")

pepflow/function.py

@@ -0,0 +1,183 @@
+import uuid
+
+import attrs
+
+from pepflow import point as pt
+from pepflow import scalar as sc
+from pepflow import utils
+
+
+class Function:
+    def __init__(
+        self,
+        is_basis: bool,
+        reuse_gradient: bool,
+        composition: dict | None = None,
+        tag: str | None = None,
+    ):
+        self.is_basis = is_basis
+        self.reuse_gradient = reuse_gradient
+        self.tag = tag
+        self.uid = attrs.field(factory=uuid.uuid4, init=False)
+        self.triplets = []  #: list[("point", "scalar", "point")] = []
+        self.constraints = []  #: list["constraint"] = []
+
+        if is_basis:
+            assert composition is None
+            self.composition = {self: 1}
+        else:
+            assert isinstance(composition, dict)
+            self.composition = composition  #: dict[{"function": float})] = []
+
+    def add_tag(self, tag: str) -> None:
+        self.tag = tag
+        return None
+
+    def get_interpolation_constraints(self):
+        raise NotImplementedError(
+            "This method should be implemented in the children class."
+        )
+
+    def add_triplet(self, triplet: tuple) -> None:
+        return NotImplemented
+
+    def add_stationary_point(self) -> pt.Point:
+        point = pt.Point(is_basis=True)
+        _, _, grad = self.generate_triplet(point)
+        self.constraints.append(
+            (grad**2).eq(0, name=str(self.__hash__) + " stationary point")
+        )
+        return point
+
+    def generate_triplet(self, point: pt.Point) -> tuple:
+        func_value = 0
+        grad = 0
+
+        if self.is_basis:
+            generate_new_basis = True
+            for triplet in self.triplets:
+                if triplet[0].uid == point.uid and self.reuse_gradient:
+                    func_value = triplet[1]
+                    grad = triplet[2]
+                    generate_new_basis = False
+                    break
+                elif triplet[0].uid == point.uid and not self.reuse_gradient:
+                    func_value = triplet[1]
+                    grad = pt.Point(is_basis=True)
+                    generate_new_basis = False
+                    self.triplets.append((point, func_value, grad))
+                    break
+            if generate_new_basis:
+                func_value = sc.Scalar(is_basis=True)
+                grad = pt.Point(is_basis=True)
+                self.triplets.append((point, func_value, grad))
+        else:
+            for function, weights in self.composition.items():
+                _, func_value_slice, grad_slice = function.generate_triplet(point)
+                func_value += weights * func_value_slice
+                grad += weights * grad_slice
+
+        return (point, func_value, grad)
+
+    def gradient(self, point: pt.Point) -> pt.Point:
+        _, _, grad = self.generate_triplet(point)
+        return grad
+
+    def subgradient(self, point: pt.Point) -> pt.Point:
+        _, _, subgrad = self.generate_triplet(point)
+        return subgrad
+
+    def function_value(self, point: pt.Point) -> sc.Scalar:
+        _, func_value, _ = self.generate_triplet(point)
+        return func_value
+
+    def __add__(self, other):
+        assert isinstance(other, Function)
+        merged_composition = utils.merge_dict(self.composition, other.composition)
+        pruned_composition = utils.prune_dict(merged_composition)
+        return Function(
+            is_basis=False,
+            reuse_gradient=self.reuse_gradient and other.reuse_gradient,
+            composition=pruned_composition,
+            tag=None,
+        )
+
+    def __sub__(self, other):
+        return self.__add__(-other)
+
+    def __mul__(self, other):
+        return self.__rmul__(other=other)
+
+    def __rmul__(self, other):
+        assert utils.is_numerical(other)
+        scaled_composition = dict()
+        for key, value in self.composition.items():
+            scaled_composition[key] = value * other
+        pruned_composition = utils.prune_dict(scaled_composition)
+        return Function(
+            is_basis=False,
+            reuse_gradient=self.reuse_gradient,
+            composition=pruned_composition,
+            tag=None,
+        )
+
+    def __neg__(self):
+        return self.__mul__(other=-1)
+
+    def __truediv__(self, other):
+        assert utils.is_numerical(other)
+        scaled_composition = dict()
+        for key, value in self.composition.items():
+            scaled_composition[key] = value / other
+        pruned_composition = utils.prune_dict(scaled_composition)
+        return Function(
+            is_basis=False,
+            reuse_gradient=self.reuse_gradient,
+            composition=pruned_composition,
+            tag=None,
+        )
+
+    def __hash__(self):
+        return hash(self.uid)
+
+    def __eq__(self, other):
+        if not isinstance(other, Function):
+            return NotImplemented
+        return self.uid == other.uid
+
+
+class SmoothConvexFunction(Function):
+    def __init__(self, L, is_basis=True, composition=None, reuse_gradient=True):
+        super().__init__(
+            is_basis=is_basis, composition=composition, reuse_gradient=reuse_gradient
+        )
+        self.L = L
+
+    def smooth_convex_interpolability_constraints(self, triplet_i, triplet_j):
+        point_i, func_value_i, grad_i = triplet_i
+        point_j, func_value_j, grad_j = triplet_j
+        func_diff = func_value_j - func_value_i
+        cross_term = grad_j * (point_i - point_j)
+        quad_term = 1 / (2 * self.L) * (grad_i - grad_j) ** 2
+
+        return (func_diff + cross_term + quad_term).le(
+            0,
+            name=str(self.__hash__())
+            + ":"
+            + str(point_i.__hash__())
+            + ","
+            + str(point_j.__hash__()),
+        )
+
+    def get_interpolation_constraints(self):
+        interpolation_constraints = []
+        for i in range(len(self.triplets)):
+            for j in range(len(self.triplets)):
+                if i == j:
+                    continue
+                interpolation_constraints.append(
+                    self.smooth_convex_interpolability_constraints(
+                        self.triplets[i], self.triplets[j]
+                    )
+                )
+        return interpolation_constraints

pepflow/pep.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import contextlib
+from typing import TYPE_CHECKING, Any
+
+import attrs
+
+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
+
+if TYPE_CHECKING:
+    from pepflow.solver import DualVariableManager
+
+
+@attrs.frozen
+class PEPResult:
+    primal_opt_value: float
+    dual_var_manager: DualVariableManager
+    solver_status: Any
+
+
+class PEPBuilder:
+    """The main class for PEP primal formulation."""
+
+    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
+
+        # Contain the name for the constraints that should be removed.
+        # We should think about a better choice like manager.
+        self.relaxed_constraints = []
+
+    @contextlib.contextmanager
+    def make_context(self, name: str, override: bool = False) -> pc.PEPContext:
+        if not override and name in self.pep_context_dict:
+            raise KeyError(f"There is already a context {name} in the builder")
+        try:
+            ctx = pc.PEPContext()
+            self.pep_context_dict[name] = ctx
+            pc.set_current_context(ctx)
+            yield ctx
+        finally:
+            pc.set_current_context(None)
+
+    def get_context(self, name: str) -> pc.PEPContext:
+        if name not in self.pep_context_dict:
+            raise KeyError(f"Cannot find a context named {name} in the builder.")
+        ctx = self.pep_context_dict[name]
+        pc.set_current_context(ctx)
+        return ctx
+
+    def clear_context(self, name: str) -> None:
+        if name not in self.pep_context_dict:
+            raise KeyError(f"Cannot find a context named {name} in the builder.")
+        del self.pep_context_dict[name]
+
+    def clear_all_context(self) -> None:
+        self.pep_context_dict.clear()
+
+    def set_init_point(self, tag: str | None = None) -> pt.Point:
+        point = pt.Point(is_basis=True)
+        point.add_tag(tag)
+        return point
+
+    def set_initial_constraint(self, constraint):
+        self.init_conditions.append(constraint)
+
+    def set_performance_metric(self, metric: sc.Scalar):
+        self.performance_metric = metric
+
+    def declare_func(self, function_class, **kwargs):
+        func = function_class(is_basis=True, composition=None, **kwargs)
+        self.functions.append(func)
+        return func
+
+    def solve(self, context: pc.PEPContext | None = None, **kwargs):
+        if context is None:
+            context = pc.get_current_context()
+        if context is None:
+            raise RuntimeError("Did you forget to create a context?")
+
+        all_constraints = [*self.init_conditions]
+        for f in self.functions:
+            if f.is_basis:
+                all_constraints.extend(f.get_interpolation_constraints())
+            all_constraints.extend(f.constraints)
+
+        # for now, we heavily rely on the CVX. We can make a wrapper class to avoid
+        # direct dependency in the future.
+        solver = ps.CVXSolver(
+            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)
+        return PEPResult(
+            primal_opt_value=result,
+            dual_var_manager=solver.dual_var_manager,
+            solver_status=problem.status,
+        )

pepflow/pep_context.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+# A global variable for storing the current context that is used for points or scalars.
+CURRENT_CONTEXT: PEPContext | None = None
+
+
+def get_current_context() -> PEPContext | None:
+    return CURRENT_CONTEXT
+
+
+def set_current_context(ctx: PEPContext | None):
+    global CURRENT_CONTEXT
+    assert ctx is None or isinstance(ctx, PEPContext)
+    CURRENT_CONTEXT = ctx
+
+
+class PEPContext:
+    def __init__(self):
+        self.points = []
+        self.scalars = []
+
+    def add_point(self, point):
+        self.points.append(point)
+
+    def add_scalar(self, scalar):
+        self.scalars.append(scalar)
+
+    def clear(self):
+        self.points = []
+        self.scalars = []

pepflow/pep_test.py

@@ -0,0 +1,58 @@
+import pytest
+
+from pepflow import pep
+from pepflow import pep_context as pc
+
+
+class TestPEPBuilder:
+    def test_make_context(self) -> None:
+        builder = pep.PEPBuilder()
+        assert pc.get_current_context() is None
+
+        with builder.make_context("test") as ctx:
+            assert ctx is pc.get_current_context()
+
+        assert pc.get_current_context() is None
+
+    def test_get_context(self) -> None:
+        builder = pep.PEPBuilder()
+        with builder.make_context("test") as ctx:
+            prev_ctx = ctx
+
+        builder.get_context("test") is prev_ctx
+
+    def test_clear_context(self) -> None:
+        builder = pep.PEPBuilder()
+        with builder.make_context("test"):
+            pass
+
+        assert "test" in builder.pep_context_dict.keys()
+        builder.clear_context("test")
+        assert "test" not in builder.pep_context_dict.keys()
+
+    def test_clear_all_context(self) -> None:
+        builder = pep.PEPBuilder()
+        with builder.make_context("test"):
+            pass
+        with builder.make_context("test2"):
+            pass
+
+        assert len(builder.pep_context_dict) == 2
+        builder.clear_all_context()
+        assert len(builder.pep_context_dict) == 0
+
+    def test_make_context_twice(self) -> None:
+        builder = pep.PEPBuilder()
+        with builder.make_context("test"):
+            pass
+
+        assert "test" in builder.pep_context_dict.keys()
+
+        with pytest.raises(
+            KeyError, match="There is already a context test in the builder"
+        ):
+            with builder.make_context("test"):
+                pass
+
+        with builder.make_context("test", override=True):
+            pass

pepflow/point.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+import attrs
+import numpy as np
+
+from pepflow import pep_context as pc
+from pepflow import utils
+from pepflow.scalar import EvalExpressionScalar, Scalar
+
+
+def is_numerical_or_point(val: Any) -> bool:
+    return utils.is_numerical(val) or isinstance(val, Point)
+
+
+def is_numerical_or_evaluatedpoint(val: Any) -> bool:
+    return utils.is_numerical(val) or isinstance(val, EvaluatedPoint)
+
+
+@attrs.frozen
+class EvalExpressionPoint:
+    op: utils.Op
+    left_point: Point | float
+    right_point: Point | float
+
+
+@attrs.frozen
+class EvaluatedPoint:
+    vector: np.array
+
+    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)}"
+            )
+
+    def __radd__(self, other):
+        return self.__add__(other)
+
+    def __sub__(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 sub operation between EvaluatedPoint and {type(other)}"
+            )
+
+    def __rsub__(self, other):
+        if isinstance(other, EvaluatedPoint):
+            return EvaluatedPoint(vector=other.vector - self.vector)
+        elif utils.is_numerical(other):
+            return EvaluatedPoint(vector=other - self.vector)
+        else:
+            raise ValueError(
+                f"Unsupported sub operation between EvaluatedPoint and {type(other)}"
+            )
+
+    def __mul__(self, other):
+        assert utils.is_numerical(other)
+        return EvaluatedPoint(vector=self.vector * other)
+
+    def __rmul__(self, other):
+        assert utils.is_numerical(other)
+        return EvaluatedPoint(vector=other * self.vector)
+
+    def __truediv__(self, other):
+        assert utils.is_numerical(other)
+        return EvaluatedPoint(vector=self.vector / other)
+
+
+@attrs.frozen
+class Point:
+    # 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
+
+    # Human tagged value for the Point
+    tags: list[str] = attrs.field(factory=list)
+
+    # Generate an automatic id
+    uid: uuid.UUID = attrs.field(factory=uuid.uuid4, init=False)
+
+    def __attrs_post_init__(self):
+        if self.is_basis:
+            assert self.eval_expression is None
+        else:
+            assert self.eval_expression is not None
+
+        pep_context = pc.get_current_context()
+        if pep_context is None:
+            raise RuntimeError("Did you forget to create a context?")
+        pep_context.add_point(self)
+
+    def add_tag(self, tag: str) -> None:
+        self.tags.append(tag)
+
+    # TODO: add a validator that `is_basis` and `eval_expression` are properly setup.
+    def __add__(self, other):
+        assert is_numerical_or_point(other)
+        return Point(
+            is_basis=False,
+            eval_expression=EvalExpressionPoint(utils.Op.ADD, self, other),
+        )
+
+    def __radd__(self, other):
+        assert is_numerical_or_point(other)
+        return Point(
+            is_basis=False,
+            eval_expression=EvalExpressionPoint(utils.Op.ADD, other, self),
+        )
+
+    def __sub__(self, other):
+        assert is_numerical_or_point(other)
+        return Point(
+            is_basis=False,
+            eval_expression=EvalExpressionPoint(utils.Op.SUB, self, other),
+        )
+
+    def __rsub__(self, other):
+        assert is_numerical_or_point(other)
+        return Point(
+            is_basis=False,
+            eval_expression=EvalExpressionPoint(utils.Op.SUB, other, 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 utils.is_numerical(other):
+            return Point(
+                is_basis=False,
+                eval_expression=EvalExpressionPoint(utils.Op.MUL, self, other),
+            )
+        else:
+            return Scalar(
+                is_basis=False,
+                eval_expression=EvalExpressionScalar(utils.Op.MUL, self, other),  # TODO
+            )
+
+    def __rmul__(self, other):
+        # TODO allow the other to be point so that we return a scalar.
+        assert is_numerical_or_point(other)
+        if utils.is_numerical(other):
+            return Point(
+                is_basis=False,
+                eval_expression=EvalExpressionPoint(utils.Op.MUL, other, self),
+            )
+        else:
+            return Scalar(
+                is_basis=False,
+                eval_expression=EvalExpressionScalar(utils.Op.MUL, other, self),  # TODO
+            )
+
+    def __pow__(self, power):
+        assert power == 2
+        return self.__rmul__(self)
+
+    def __neg__(self):
+        return self.__rmul__(other=-1)
+
+    def __truediv__(self, other):
+        assert utils.is_numerical(other)
+        return Point(
+            is_basis=False,
+            eval_expression=EvalExpressionPoint(utils.Op.DIV, self, other),
+        )
+
+    def __hash__(self):
+        return hash(self.uid)
+
+    def __eq__(self, other):
+        if not isinstance(other, Point):
+            return NotImplemented
+        return self.uid == other.uid