pepflow 0.1.0__py3-none-any.whl → 0.1.3a1__py3-none-any.whl

pepflow/function_test.py

@@ -0,0 +1,134 @@
+# Copyright: 2025 The PEPFlow Developers
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import numpy as np
+
+from pepflow import expression_manager as exm
+from pepflow import function as fc
+from pepflow import pep as pep
+
+
+def test_function_repr():
+    pep_builder = pep.PEPBuilder()
+    with pep_builder.make_context("test"):
+        f = fc.Function(is_basis=True, reuse_gradient=False)
+        print(f)  # it should be fine without tag
+        f.add_tag("f")
+        assert str(f) == "f"
+
+
+def test_stationary_point():
+    pep_builder = pep.PEPBuilder()
+    with pep_builder.make_context("test") as ctx:
+        f = fc.Function(is_basis=True, reuse_gradient=False, tags=["f"])
+        f.add_stationary_point("x_star")
+
+        assert len(ctx.triplets) == 1
+        assert len(ctx.triplets[f]) == 1
+
+        f_triplet = ctx.triplets[f][0]
+        assert f_triplet.name == "x_star_f(x_star)_gradient_f(x_star)"
+        assert f_triplet.gradient.tag == "gradient_f(x_star)"
+        assert f_triplet.function_value.tag == "f(x_star)"
+
+        em = exm.ExpressionManager(ctx)
+        np.testing.assert_allclose(
+            em.eval_point(f_triplet.gradient).vector, np.array([0])
+        )
+        np.testing.assert_allclose(em.eval_point(f_triplet.point).vector, np.array([1]))
+
+
+def test_stationary_point_scaled():
+    pep_builder = pep.PEPBuilder()
+    with pep_builder.make_context("test") as ctx:
+        f = fc.Function(is_basis=True, reuse_gradient=False, tags=["f"])
+        g = 5 * f
+        g.add_stationary_point("x_star")
+
+        assert len(ctx.triplets) == 1
+        assert len(ctx.triplets[f]) == 1
+
+        f_triplet = ctx.triplets[f][0]
+        assert f_triplet.name == "x_star_f(x_star)_gradient_f(x_star)"
+        assert f_triplet.gradient.tag == "gradient_f(x_star)"
+        assert f_triplet.function_value.tag == "f(x_star)"
+
+        em = exm.ExpressionManager(ctx)
+        np.testing.assert_allclose(
+            em.eval_point(f_triplet.gradient).vector, np.array([0])
+        )
+        np.testing.assert_allclose(em.eval_point(f_triplet.point).vector, np.array([1]))
+
+
+def test_stationary_point_additive():
+    pep_builder = pep.PEPBuilder()
+    with pep_builder.make_context("test") as ctx:
+        f = fc.Function(is_basis=True, reuse_gradient=False)
+        f.add_tag("f")
+        g = fc.Function(is_basis=True, reuse_gradient=False)
+        g.add_tag("g")
+        h = f + g
+        h.add_tag("h")
+
+        h.add_stationary_point("x_star")
+        assert len(ctx.triplets) == 2
+        assert len(ctx.triplets[f]) == 1
+        assert len(ctx.triplets[g]) == 1
+
+        f_triplet = ctx.triplets[f][0]
+        g_triplet = ctx.triplets[g][0]
+        assert f_triplet.name == "x_star_f(x_star)_gradient_f(x_star)"
+        assert g_triplet.name == "x_star_g(x_star)_gradient_g(x_star)"
+
+        em = exm.ExpressionManager(ctx)
+        np.testing.assert_allclose(
+            em.eval_point(f_triplet.gradient).vector, np.array([0, 1])
+        )
+        np.testing.assert_allclose(
+            em.eval_point(g_triplet.gradient).vector, np.array([0, -1])
+        )
+
+
+def test_stationary_point_linear_combination():
+    pep_builder = pep.PEPBuilder()
+    with pep_builder.make_context("test") as ctx:
+        f = fc.Function(is_basis=True, reuse_gradient=False)
+        f.add_tag("f")
+        g = fc.Function(is_basis=True, reuse_gradient=False)
+        g.add_tag("g")
+        h = 3 * f + 2 * g
+        h.add_tag("h")
+
+        h.add_stationary_point("x_star")
+        assert len(ctx.triplets) == 2
+        assert len(ctx.triplets[f]) == 1
+        assert len(ctx.triplets[g]) == 1
+
+        f_triplet = ctx.triplets[f][0]
+        g_triplet = ctx.triplets[g][0]
+        assert f_triplet.name == "x_star_f(x_star)_gradient_f(x_star)"
+        assert g_triplet.name == "x_star_g(x_star)_gradient_g(x_star)"
+
+        em = exm.ExpressionManager(ctx)
+        np.testing.assert_allclose(
+            em.eval_point(f_triplet.gradient).vector, np.array([0, 1])
+        )
+        np.testing.assert_allclose(
+            em.eval_point(g_triplet.gradient).vector, np.array([0, -1.5])
+        )

pepflow/interactive_constraint.py

@@ -0,0 +1,264 @@
+# Copyright: 2025 The PEPFlow Developers
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+import dash
+import dash_bootstrap_components as dbc
+import numpy as np
+import pandas as pd
+import plotly
+import plotly.express as px
+import plotly.graph_objects as go
+from dash import Dash, Input, Output, State, dcc, html
+
+from pepflow.constants import PSD_CONSTRAINT
+
+if TYPE_CHECKING:
+    from pepflow.pep import PEPBuilder, PEPResult
+    from pepflow.pep_context import PEPContext
+
+
+plotly.io.renderers.default = "colab+vscode"
+plotly.io.templates.default = "plotly_white"
+
+
+def solve_prob_and_get_figure(
+    pep_builder: PEPBuilder, context: PEPContext
+) -> tuple[go.Figure, pd.DataFrame, PEPResult]:
+    assert len(context.triplets) == 1, "Support single function only for now"
+
+    result = pep_builder.solve(context=context)
+
+    df_dict, order_dict = context.triplets_to_df_and_order()
+    f = pep_builder.functions[0]
+
+    df = df_dict[f]
+    order = order_dict[f]
+
+    df["constraint"] = df.constraint_name.map(
+        lambda x: "inactive" if x in pep_builder.relaxed_constraints else "active"
+    )
+    df["dual_value"] = df.constraint_name.map(
+        lambda x: result.dual_var_manager.dual_value(x)
+    )
+    return processed_df_to_fig(df, order), df, result
+
+
+def processed_df_to_fig(df: pd.DataFrame, order: list[str]):
+    fig = px.scatter(
+        df,
+        x="row",
+        y="col",
+        color="dual_value",
+        symbol="constraint",
+        symbol_map={"inactive": "x-open", "active": "circle"},
+        custom_data="constraint_name",
+        color_continuous_scale="Viridis",
+    )
+    fig.update_layout(yaxis=dict(autorange="reversed"))
+    fig.update_traces(marker=dict(size=15))
+    fig.update_layout(
+        coloraxis_colorbar=dict(yanchor="top", y=1, x=1.3, ticks="outside")
+    )
+    fig.update_xaxes(tickmode="array", tickvals=list(range(len(order))), ticktext=order)
+    fig.update_yaxes(tickmode="array", tickvals=list(range(len(order))), ticktext=order)
+    return fig
+
+
+def get_matrix_of_dual_value(df: pd.DataFrame) -> np.ndarray:
+    # Check if we need to update the order.
+    return (
+        pd.pivot_table(
+            df, values="dual_value", index="row", columns="col", dropna=False
+        )
+        .fillna(0.0)
+        .to_numpy()
+        .T
+    )
+
+
+def launch(pep_builder: PEPBuilder, context: PEPContext):
+    app = Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])
+    fig, df, result = solve_prob_and_get_figure(pep_builder, context)
+
+    with np.printoptions(precision=3, linewidth=100, suppress=True):
+        dual_value_tab = html.Pre(
+            str(get_matrix_of_dual_value(df)),
+            id="dual-value-display",
+            style={
+                "border": "1px solid lightgrey",
+                "padding": "10px",
+                "height": "60vh",
+                "overflowY": "auto",
+            },
+        )
+    # Think how can we manipulate the pep_builder here.
+    display_row = dbc.Row(
+        [
+            dbc.Col(
+                [
+                    dbc.Button(
+                        "Relax All Constraints",
+                        id="relax-all-constraints-button",
+                        style={"margin-bottom": "5px", "margin-right": "5px"},
+                    ),
+                    dbc.Button(
+                        "Restore All Constraints",
+                        id="restore-all-constraints-button",
+                        style={"margin-bottom": "5px"},
+                        color="success",
+                    ),
+                    dbc.Tabs(
+                        [
+                            dbc.Tab(
+                                html.Div(
+                                    dcc.Graph(id="interactive-scatter", figure=fig),
+                                ),
+                                label="Interactive Heatmap",
+                                tab_id="interactive-scatter-tab",
+                            ),
+                            dbc.Tab(
+                                dual_value_tab,
+                                label="Dual Value Matrix",
+                                tab_id="dual_value_tab",
+                            ),
+                        ],
+                        active_tab="interactive-scatter-tab",
+                    ),
+                ],
+                width=5,
+            ),
+            # Column 2: The Selected Data Display (takes up 4 of 12 columns)
+            dbc.Col(
+                [
+                    dbc.Button(
+                        "Solve PEP Problem",
+                        id="solve-button",
+                        color="primary",
+                        className="me-1",
+                        style={"margin-bottom": "5px"},
+                    ),
+                    dcc.Loading(
+                        dbc.Card(
+                            id="result-card",
+                            style={"height": "60vh", "overflow-y": "auto"},
+                        )
+                    ),
+                ],
+                width=7,
+            ),
+        ],
+    )
+
+    # 3. Define the app layout
+    app.layout = html.Div(
+        [
+            html.H2("PEPFlow"),
+            display_row,
+            # Store the entire DataFrame as a dictionary in dcc.Store
+            dcc.Store(id="dataframe-store", data=df.to_dict("records")),
+        ]
+    )
+
+    @dash.callback(
+        Output("result-card", "children"),
+        Output("dual-value-display", "children"),
+        Output("interactive-scatter", "figure"),
+        Output("dataframe-store", "data"),
+        Input("solve-button", "n_clicks"),
+    )
+    def solve(_):
+        fig, df, result = solve_prob_and_get_figure(pep_builder, context)
+        with np.printoptions(precision=3, linewidth=100, suppress=True):
+            psd_dual_value = np.array(
+                result.dual_var_manager.dual_value(PSD_CONSTRAINT)
+            )
+            reslt_card = dbc.CardBody(
+                [
+                    html.H2(f"Optimal Value {result.primal_opt_value:.4g}"),
+                    html.H3(f"Solver Status: {result.solver_status}"),
+                    html.P("PSD Dual Variable:"),
+                    html.Pre(str(psd_dual_value)),
+                    html.P("Relaxed Constraints:"),
+                    html.Pre(json.dumps(pep_builder.relaxed_constraints, indent=2)),
+                ]
+            )
+            dual_value_display = str(get_matrix_of_dual_value(df))
+        return reslt_card, dual_value_display, fig, df.to_dict("records")
+
+    @dash.callback(
+        Output("interactive-scatter", "figure", allow_duplicate=True),
+        Output("dataframe-store", "data", allow_duplicate=True),
+        Input("restore-all-constraints-button", "n_clicks"),
+        State("dataframe-store", "data"),
+        prevent_initial_call=True,
+    )
+    def restore_all_constraints(_, previous_df):
+        nonlocal pep_builder
+        df_updated = pd.DataFrame(previous_df)
+        pep_builder.relaxed_constraints = []
+        df_updated["constraint"] = "active"
+        order = context.order_of_point(pep_builder.functions[0])
+        return processed_df_to_fig(df_updated, order), df_updated.to_dict("records")
+
+    @dash.callback(
+        Output("interactive-scatter", "figure", allow_duplicate=True),
+        Output("dataframe-store", "data", allow_duplicate=True),
+        Input("relax-all-constraints-button", "n_clicks"),
+        State("dataframe-store", "data"),
+        prevent_initial_call=True,
+    )
+    def relax_all_constraints(_, previous_df):
+        nonlocal pep_builder
+        df_updated = pd.DataFrame(previous_df)
+        pep_builder.relaxed_constraints = df_updated["constraint_name"].to_list()
+        df_updated["constraint"] = "inactive"
+        order = context.order_of_point(pep_builder.functions[0])
+        return processed_df_to_fig(df_updated, order), df_updated.to_dict("records")
+
+    @dash.callback(
+        Output("interactive-scatter", "figure", allow_duplicate=True),
+        Output("dataframe-store", "data", allow_duplicate=True),
+        Input("interactive-scatter", "clickData"),
+        State("dataframe-store", "data"),
+        prevent_initial_call=True,
+    )
+    def update_df_and_redraw(clickData, previous_df):
+        nonlocal pep_builder
+        if not clickData["points"][0]["customdata"]:
+            return dash.no_update, dash.no_update, dash.no_update
+
+        clicked_name = clickData["points"][0]["customdata"][0]
+        if clicked_name not in pep_builder.relaxed_constraints:
+            pep_builder.relaxed_constraints.append(clicked_name)
+        else:
+            pep_builder.relaxed_constraints.remove(clicked_name)
+
+        df_updated = pd.DataFrame(previous_df)
+        df_updated["constraint"] = df_updated.constraint_name.map(
+            lambda x: "inactive" if x in pep_builder.relaxed_constraints else "active"
+        )
+        order = context.order_of_point(pep_builder.functions[0])
+        return processed_df_to_fig(df_updated, order), df_updated.to_dict("records")
+
+    app.run(debug=True)

pepflow/pep.py

@@ -1,16 +1,40 @@
+# Copyright: 2025 The PEPFlow Developers
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
 from __future__ import annotations
 
 import contextlib
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Iterator
 
 import attrs
+import numpy as np
+import pandas as pd
 
 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.constants import PSD_CONSTRAINT
 
 if TYPE_CHECKING:
+    from pepflow.constraint import Constraint
+    from pepflow.function import Function
     from pepflow.solver import DualVariableManager
 
 
@@ -19,6 +43,33 @@ class PEPResult:
     primal_opt_value: float
     dual_var_manager: DualVariableManager
     solver_status: Any
+    context: pc.PEPContext
+
+    def get_function_dual_variables(self) -> dict[Function, np.ndarray]:
+        def get_matrix_of_dual_value(df: pd.DataFrame) -> np.ndarray:
+            # Check if we need to update the order.
+            return (
+                pd.pivot_table(
+                    df, values="dual_value", index="row", columns="col", dropna=False
+                )
+                .fillna(0.0)
+                .to_numpy()
+                .T
+            )
+
+        df_dict, _ = self.context.triplets_to_df_and_order()
+        df_dict_matrix = {}
+        for f in df_dict.keys():
+            df = df_dict[f]
+            df["dual_value"] = df.constraint_name.map(
+                lambda x: self.dual_var_manager.dual_value(x)
+            )
+            df_dict_matrix[f] = get_matrix_of_dual_value(df)
+
+        return df_dict_matrix
+
+    def get_psd_dual_matrix(self):
+        return np.array(self.dual_var_manager.dual_value(PSD_CONSTRAINT))
 
 
 class PEPBuilder:
@@ -37,7 +88,9 @@ class PEPBuilder:
         self.relaxed_constraints = []
 
     @contextlib.contextmanager
-    def make_context(self, name: str, override: bool = False) -> pc.PEPContext:
+    def make_context(
+        self, name: str, override: bool = False
+    ) -> Iterator[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:
@@ -63,7 +116,7 @@ class PEPBuilder:
     def clear_all_context(self) -> None:
         self.pep_context_dict.clear()
 
-    def set_init_point(self, tag: str | None = None) -> pt.Point:
+    def set_init_point(self, tag: str) -> pt.Point:
         point = pt.Point(is_basis=True)
         point.add_tag(tag)
         return point
@@ -74,6 +127,9 @@ class PEPBuilder:
     def set_performance_metric(self, metric: sc.Scalar):
         self.performance_metric = metric
 
+    def set_relaxed_constraints(self, relaxed_constraints: list[str]):
+        self.relaxed_constraints.extend(relaxed_constraints)
+
     def declare_func(self, function_class, **kwargs):
         func = function_class(is_basis=True, composition=None, **kwargs)
         self.functions.append(func)
@@ -85,11 +141,10 @@ class PEPBuilder:
         if context is None:
             raise RuntimeError("Did you forget to create a context?")
 
-        all_constraints = [*self.init_conditions]
+        all_constraints: list[Constraint] = [*self.init_conditions]
         for f in self.functions:
-            if f.is_basis:
-                all_constraints.extend(f.get_interpolation_constraints())
-            all_constraints.extend(f.constraints)
+            all_constraints.extend(f.get_interpolation_constraints())
+            all_constraints.extend(context.opt_conditions[f])
 
         # for now, we heavily rely on the CVX. We can make a wrapper class to avoid
         # direct dependency in the future.
@@ -106,4 +161,5 @@ class PEPBuilder:
             primal_opt_value=result,
             dual_var_manager=solver.dual_var_manager,
             solver_status=problem.status,
+            context=context,
         )

pepflow/pep_context.py

@@ -1,5 +1,36 @@
+# Copyright: 2025 The PEPFlow Developers
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
 from __future__ import annotations
 
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+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
+
 # A global variable for storing the current context that is used for points or scalars.
 CURRENT_CONTEXT: PEPContext | None = None
 
@@ -16,15 +47,85 @@ def set_current_context(ctx: PEPContext | None):
 
 class PEPContext:
     def __init__(self):
-        self.points = []
-        self.scalars = []
+        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)
 
-    def add_point(self, point):
+    def set_as_current(self) -> PEPContext:
+        set_current_context(self)
+        return self
+
+    def add_point(self, point: Point):
         self.points.append(point)
 
-    def add_scalar(self, scalar):
+    def add_scalar(self, scalar: Scalar):
         self.scalars.append(scalar)
 
+    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 get_by_tag(self, tag: str) -> Point | Scalar:
+        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")
+
     def clear(self):
-        self.points = []
-        self.scalars = []
+        self.points.clear()
+        self.scalars.clear()
+        self.triplets.clear()
+        self.opt_conditions.clear()
+
+    def tracked_point(self, func: Function) -> list[Point]:
+        return natsort.natsorted(
+            [t.point for t in self.triplets[func]], key=lambda x: x.tag
+        )
+
+    def tracked_grad(self, func: Function) -> list[Point]:
+        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]:
+        return natsort.natsorted(
+            [t.function_value for t in self.triplets[func]], key=lambda x: x.tag
+        )
+
+    def order_of_point(self, func: Function) -> list[str]:
+        return natsort.natsorted([t.point.tag for t in self.triplets[func]])
+
+    def triplets_to_df_and_order(
+        self,
+    ) -> tuple[dict[Function, pd.DataFrame], dict[Function, list[str]]]:
+        func_to_df: dict[Function, pd.DataFrame] = {}
+        func_to_order: dict[Function, list[str]] = {}
+
+        def name_to_point_tuple(c_name: str) -> list[str]:
+            _, points = c_name.split(":")
+            return points.split(",")
+
+        for func, triplets in self.triplets.items():
+            order = self.order_of_point(func)
+            df = pd.DataFrame(
+                [
+                    (
+                        constraint.name,
+                        *name_to_point_tuple(constraint.name),
+                    )
+                    for constraint in func.get_interpolation_constraints(self)
+                ],
+                columns=["constraint_name", "col_point", "row_point"],
+            )
+            df["row"] = df["row_point"].map(lambda x: order.index(x))
+            df["col"] = df["col_point"].map(lambda x: order.index(x))
+            func_to_df[func] = df
+            func_to_order[func] = order
+
+        return func_to_df, func_to_order