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

pepflow/function_test.py

@@ -38,8 +38,8 @@ def pep_context() -> Iterator[pc.PEPContext]:
 
 
 def test_function_add_tag(pep_context: pc.PEPContext) -> None:
-    f1 = fc.Function(is_basis=True, reuse_gradient=False, tags=["f1"])
-    f2 = fc.Function(is_basis=True, reuse_gradient=False, tags=["f2"])
+    f1 = fc.Function(is_basis=True, tags=["f1"])
+    f2 = fc.Function(is_basis=True, tags=["f2"])
 
     f_add = f1 + f2
     assert f_add.tag == "f1+f2"
@@ -55,7 +55,7 @@ def test_function_add_tag(pep_context: pc.PEPContext) -> None:
 
 
 def test_function_mul_tag(pep_context: pc.PEPContext) -> None:
-    f = fc.Function(is_basis=True, reuse_gradient=False, tags=["f"])
+    f = fc.Function(is_basis=True, tags=["f"])
 
     f_mul = f * 0.1
     assert f_mul.tag == "0.1*f"
@@ -71,8 +71,8 @@ def test_function_mul_tag(pep_context: pc.PEPContext) -> None:
 
 
 def test_function_add_and_mul_tag(pep_context: pc.PEPContext) -> None:
-    f1 = fc.Function(is_basis=True, reuse_gradient=False, tags=["f1"])
-    f2 = fc.Function(is_basis=True, reuse_gradient=False, tags=["f2"])
+    f1 = fc.Function(is_basis=True, tags=["f1"])
+    f2 = fc.Function(is_basis=True, tags=["f2"])
 
     f_add_mul = (f1 + f2) * 0.1
     assert f_add_mul.tag == "0.1*(f1+f2)"
@@ -94,116 +94,182 @@ def test_function_add_and_mul_tag(pep_context: pc.PEPContext) -> None:
 
 
 def test_function_call(pep_context: pc.PEPContext) -> None:
-    f = fc.Function(is_basis=True, reuse_gradient=False, tags=["f"])
+    f = fc.Function(is_basis=True, tags=["f"])
     x = point.Point(is_basis=True, eval_expression=None, tags=["x"])
     assert f.function_value(x) == f(x)
 
 
-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])
-        )
+def test_function_repr(pep_context: pc.PEPContext):
+    f = fc.Function(
+        is_basis=True,
+    )
+    print(f)  # it should be fine without tag
+    f.add_tag("f")
+    assert str(f) == "f"
+
+
+def test_stationary_point(pep_context: pc.PEPContext):
+    f = fc.Function(
+        is_basis=True,
+        tags=["f"],
+    )
+    f.add_stationary_point("x_star")
+
+    assert len(pep_context.triplets) == 1
+    assert len(pep_context.triplets[f]) == 1
+
+    f_triplet = pep_context.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(pep_context)
+    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_context: pc.PEPContext):
+    f = fc.Function(
+        is_basis=True,
+        tags=["f"],
+    )
+    g = 5 * f
+    g.add_stationary_point("x_star")
+
+    assert len(pep_context.triplets) == 1
+    assert len(pep_context.triplets[f]) == 1
+
+    f_triplet = pep_context.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(pep_context)
+    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_context: pc.PEPContext):
+    f = fc.Function(is_basis=True)
+    f.add_tag("f")
+    g = fc.Function(is_basis=True)
+    g.add_tag("g")
+    h = f + g
+    h.add_tag("h")
+
+    h.add_stationary_point("x_star")
+    assert len(pep_context.triplets) == 2
+    assert len(pep_context.triplets[f]) == 1
+    assert len(pep_context.triplets[g]) == 1
+
+    f_triplet = pep_context.triplets[f][0]
+    g_triplet = pep_context.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(pep_context)
+    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_context: pc.PEPContext):
+    f = fc.Function(
+        is_basis=True,
+    )
+    f.add_tag("f")
+    g = fc.Function(
+        is_basis=True,
+    )
+    g.add_tag("g")
+    h = 3 * f + 2 * g
+    h.add_tag("h")
+
+    h.add_stationary_point("x_star")
+    assert len(pep_context.triplets) == 2
+    assert len(pep_context.triplets[f]) == 1
+    assert len(pep_context.triplets[g]) == 1
+
+    f_triplet = pep_context.triplets[f][0]
+    g_triplet = pep_context.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(pep_context)
+    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])
+    )
+
+
+def test_function_generate_triplet(pep_context: pc.PEPContext):
+    f = fc.Function(is_basis=True)
+    f.add_tag("f")
+    g = fc.Function(is_basis=True)
+    g.add_tag("g")
+    h = 5 * f + 5 * g
+    h.add_tag("h")
+
+    p1 = point.Point(is_basis=True)
+    p1.add_tag("p1")
+    p1_triplet = h.generate_triplet(p1)
+    p1_triplet_1 = h.generate_triplet(p1)
+
+    pm = exm.ExpressionManager(pep_context)
+
+    np.testing.assert_allclose(pm.eval_point(p1).vector, np.array([1, 0, 0]))
+
+    np.testing.assert_allclose(
+        pm.eval_point(p1_triplet.gradient).vector, np.array([0, 5, 5])
+    )
+    np.testing.assert_allclose(
+        pm.eval_scalar(p1_triplet.function_value).vector, np.array([5, 5])
+    )
+
+    np.testing.assert_allclose(
+        pm.eval_point(p1_triplet_1.gradient).vector, np.array([0, 5, 5])
+    )
+    np.testing.assert_allclose(
+        pm.eval_scalar(p1_triplet_1.function_value).vector, np.array([5, 5])
+    )
+
+
+def test_function_add_stationary_point(pep_context: pc.PEPContext):
+    f = fc.Function(is_basis=True)
+    f.add_tag("f")
+    x_opt = f.add_stationary_point("x_opt")
+
+    pm = exm.ExpressionManager(pep_context)
+
+    np.testing.assert_allclose(pm.eval_point(x_opt).vector, np.array([1]))
+
+
+def test_smooth_interpolability_constraints(pep_context: pc.PEPContext):
+    f = fc.SmoothConvexFunction(L=1)
+    f.add_tag("f")
+    _ = f.add_stationary_point("x_opt")
+
+    x_0 = point.Point(is_basis=True)
+    x_0.add_tag("x_0")
+    _ = f.generate_triplet(x_0)
+
+    all_interpolation_constraints = f.get_interpolation_constraints()
+
+    pm = exm.ExpressionManager(pep_context)
+
+    np.testing.assert_allclose(
+        pm.eval_scalar(all_interpolation_constraints[1].scalar).vector, [1, -1]
+    )
+    np.testing.assert_allclose(
+        pm.eval_scalar(all_interpolation_constraints[1].scalar).matrix,
+        [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.5]],
+    )
+
+    np.testing.assert_allclose(
+        pm.eval_scalar(all_interpolation_constraints[1].scalar).constant, 0
+    )

pepflow/interactive_constraint.py

@@ -22,6 +22,7 @@ from __future__ import annotations
 import json
 from typing import TYPE_CHECKING
 
+import attrs
 import dash
 import dash_bootstrap_components as dbc
 import numpy as np
@@ -29,11 +30,12 @@ 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 dash import ALL, MATCH, Dash, Input, Output, State, dcc, html
 
 from pepflow.constants import PSD_CONSTRAINT
 
 if TYPE_CHECKING:
+    from pepflow.function import Function
     from pepflow.pep import PEPBuilder, PEPResult
     from pepflow.pep_context import PEPContext
 
@@ -42,26 +44,84 @@ plotly.io.renderers.default = "colab+vscode"
 plotly.io.templates.default = "plotly_white"
 
 
+@attrs.frozen
+class PlotData:
+    dataframe: pd.DataFrame
+    figure: go.Figure
+    function: Function
+
+    def dual_matrix_to_tab(self) -> html.Pre:
+        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
+            )
+
+        with np.printoptions(precision=3, linewidth=100, suppress=True):
+            dual_value_tab = html.Pre(
+                str(get_matrix_of_dual_value(self.dataframe)),
+                id={"type": "dual-value-display", "index": self.function.tag},
+                style={
+                    "border": "1px solid lightgrey",
+                    "padding": "10px",
+                    "height": "60vh",
+                    "overflowY": "auto",
+                },
+            )
+        return dual_value_tab
+
+    def plot_data_to_tab(self) -> dbc.Tab:
+        tab = dbc.Tab(
+            html.Div(
+                [
+                    html.P("Interactive Heat Map:"),
+                    dcc.Graph(
+                        id={
+                            "type": "interactive-scatter",
+                            "index": self.function.tag,
+                        },
+                        figure=self.figure,
+                    ),
+                    html.P("Dual Value Matrix:"),
+                    self.dual_matrix_to_tab(),
+                ]
+            ),
+            label=f"{self.function.tag}-Interpolation Conditions",
+            tab_id=f"{self.function.tag}-interactive-scatter-tab",
+        )
+        return tab
+
+
 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"
+) -> tuple[list[PlotData], PEPResult]:
+    plot_data_list: list[PlotData] = []
 
     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]
+    for f in context.triplets.keys():
+        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
+        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)
+        )
+
+        plot_data_list.append(
+            PlotData(dataframe=df, figure=processed_df_to_fig(df, order), function=f)
+        )
+
+    return plot_data_list, result
 
 
 def processed_df_to_fig(df: pd.DataFrame, order: list[str]):
@@ -99,19 +159,7 @@ def get_matrix_of_dual_value(df: pd.DataFrame) -> np.ndarray:
 
 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",
-            },
-        )
+    plot_data_list, result = solve_prob_and_get_figure(pep_builder, context)
     # Think how can we manipulate the pep_builder here.
     display_row = dbc.Row(
         [
@@ -129,21 +177,8 @@ def launch(pep_builder: PEPBuilder, context: PEPContext):
                         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",
+                        [plot_data.plot_data_to_tab() for plot_data in plot_data_list],
+                        active_tab=f"{plot_data_list[0].function.tag}-interactive-scatter-tab",
                     ),
                 ],
                 width=5,
@@ -175,25 +210,37 @@ def launch(pep_builder: PEPBuilder, context: PEPContext):
         [
             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")),
+            # For each function, store the corresponding DataFrame as a dictionary in dcc.Store
+            *[
+                dcc.Store(
+                    id={
+                        "type": "dataframe-store",
+                        "index": plot_data.function.tag,
+                    },
+                    data=(
+                        plot_data.function.tag,
+                        plot_data.dataframe.to_dict("records"),
+                    ),
+                )
+                for plot_data in plot_data_list
+            ],
         ]
     )
 
     @dash.callback(
         Output("result-card", "children"),
-        Output("dual-value-display", "children"),
-        Output("interactive-scatter", "figure"),
-        Output("dataframe-store", "data"),
+        Output({"type": "dual-value-display", "index": ALL}, "children"),
+        Output({"type": "interactive-scatter", "index": ALL}, "figure"),
+        Output({"type": "dataframe-store", "index": ALL}, "data"),
         Input("solve-button", "n_clicks"),
     )
     def solve(_):
-        fig, df, result = solve_prob_and_get_figure(pep_builder, context)
+        plot_data_list, 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(
+            result_card = dbc.CardBody(
                 [
                     html.H2(f"Optimal Value {result.primal_opt_value:.4g}"),
                     html.H3(f"Solver Status: {result.solver_status}"),
@@ -203,50 +250,89 @@ def launch(pep_builder: PEPBuilder, context: PEPContext):
                     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")
+            dual_value_displays = [
+                str(get_matrix_of_dual_value(plot_data.dataframe))
+                for plot_data in plot_data_list
+            ]
+        figs = [plot_data.figure for plot_data in plot_data_list]
+
+        df_data = [
+            (plot_data.function.tag, plot_data.dataframe.to_dict("records"))
+            for plot_data in plot_data_list
+        ]
+
+        return result_card, dual_value_displays, figs, df_data
 
     @dash.callback(
-        Output("interactive-scatter", "figure", allow_duplicate=True),
-        Output("dataframe-store", "data", allow_duplicate=True),
+        Output(
+            {"type": "interactive-scatter", "index": ALL},
+            "figure",
+            allow_duplicate=True,
+        ),
+        Output({"type": "dataframe-store", "index": ALL}, "data", allow_duplicate=True),
         Input("restore-all-constraints-button", "n_clicks"),
-        State("dataframe-store", "data"),
+        State({"type": "dataframe-store", "index": ALL}, "data"),
         prevent_initial_call=True,
     )
-    def restore_all_constraints(_, previous_df):
+    def restore_all_constraints(_, list_previous_df_tuples):
         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")
+        updated_figs = []
+        df_data = []
+        for previous_df_tuple in list_previous_df_tuples:
+            tag, previous_df = previous_df_tuple
+            df_updated = pd.DataFrame(previous_df)
+            df_updated["constraint"] = "active"
+            order = context.order_of_point(pep_builder.get_func_by_tag(tag))
+            updated_figs.append(processed_df_to_fig(df_updated, order))
+            df_data.append((tag, df_updated.to_dict("records")))
+        return updated_figs, df_data
 
     @dash.callback(
-        Output("interactive-scatter", "figure", allow_duplicate=True),
-        Output("dataframe-store", "data", allow_duplicate=True),
+        Output(
+            {"type": "interactive-scatter", "index": ALL},
+            "figure",
+            allow_duplicate=True,
+        ),
+        Output({"type": "dataframe-store", "index": ALL}, "data", allow_duplicate=True),
         Input("relax-all-constraints-button", "n_clicks"),
-        State("dataframe-store", "data"),
+        State({"type": "dataframe-store", "index": ALL}, "data"),
         prevent_initial_call=True,
     )
-    def relax_all_constraints(_, previous_df):
+    def relax_all_constraints(_, list_previous_df_tuples):
         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")
+        pep_builder.relaxed_constraints = []
+        updated_figs = []
+        df_data = []
+        for previous_df_tuple in list_previous_df_tuples:
+            tag, previous_df = previous_df_tuple
+            df_updated = pd.DataFrame(previous_df)
+            pep_builder.relaxed_constraints.extend(
+                df_updated["constraint_name"].to_list()
+            )
+            df_updated["constraint"] = "inactive"
+            order = context.order_of_point(pep_builder.get_func_by_tag(tag))
+            updated_figs.append(processed_df_to_fig(df_updated, order))
+            df_data.append((tag, df_updated.to_dict("records")))
+        return updated_figs, df_data
 
     @dash.callback(
-        Output("interactive-scatter", "figure", allow_duplicate=True),
-        Output("dataframe-store", "data", allow_duplicate=True),
-        Input("interactive-scatter", "clickData"),
-        State("dataframe-store", "data"),
+        Output(
+            {"type": "interactive-scatter", "index": MATCH},
+            "figure",
+            allow_duplicate=True,
+        ),
+        Output(
+            {"type": "dataframe-store", "index": MATCH}, "data", allow_duplicate=True
+        ),
+        Input({"type": "interactive-scatter", "index": MATCH}, "clickData"),
+        State({"type": "dataframe-store", "index": MATCH}, "data"),
         prevent_initial_call=True,
     )
-    def update_df_and_redraw(clickData, previous_df):
+    def update_df_and_redraw(clickData, previous_df_tuple):
         nonlocal pep_builder
         if not clickData["points"][0]["customdata"]:
-            return dash.no_update, dash.no_update, dash.no_update
+            return dash.no_update, dash.no_update
 
         clicked_name = clickData["points"][0]["customdata"][0]
         if clicked_name not in pep_builder.relaxed_constraints:
@@ -254,11 +340,15 @@ def launch(pep_builder: PEPBuilder, context: PEPContext):
         else:
             pep_builder.relaxed_constraints.remove(clicked_name)
 
+        tag, previous_df = previous_df_tuple
         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")
+        order = context.order_of_point(pep_builder.get_func_by_tag(tag))
+        return processed_df_to_fig(df_updated, order), (
+            tag,
+            df_updated.to_dict("records"),
+        )
 
     app.run(debug=True)