From d80df0115644ac65c747e57054fe54c7ef899bfa Mon Sep 17 00:00:00 2001 From: Navan Chauhan Date: Wed, 27 Mar 2024 12:16:46 -0600 Subject: option to manually specify equation --- api.py | 31 +++++++++++++++++++++++-------- 1 file changed, 23 insertions(+), 8 deletions(-) (limited to 'api.py') diff --git a/api.py b/api.py index a52345f..b331879 100644 --- a/api.py +++ b/api.py @@ -8,6 +8,8 @@ import pandas as pd from pydantic import BaseModel, Field +from sympy import symbols, sympify + font_path = 'times_new_roman.ttf' times_new_roman = font_manager.FontProperties(fname=font_path, style='normal') @@ -29,6 +31,8 @@ class PlotDataRequest(BaseModel): x_axis_scale: str = Field("linear", title="X Axis Scale", description="Scale of the X axis") y_axis_scale: str = Field("linear", title="Y Axis Scale", description="Scale of the Y axis") + trendline_equation: str = Field(None, title="Trendline Equation", description="Manually specify the equation for the trendline") + constant_line_vals: list[float] = Field([], title="Constant Line Values", description="List of values for the constant line") constant_line_name: list[str] = Field([], title="Constant Line Names", description="List of names for the constant line") @@ -38,7 +42,7 @@ def plot_data(x_data, y_data, std_dev_data, color_picker, labels, df, title = "Plot", x_label = "X Axis", y_label = "Y Axis", plot_background_color="#ffffff", constant_line=[], enable_trendline=True, enable_grid=False, - trendline_color="#000000", x_axis_scale="linear", y_axis_scale="linear"): + trendline_color="#000000", x_axis_scale="linear", y_axis_scale="linear", trendline_equation=None): fig, ax = plt.subplots(dpi=300) plots = [] @@ -61,12 +65,23 @@ def plot_data(x_data, y_data, std_dev_data, color_picker, labels, df, handles = plots if enable_trendline: - x = df[x_data[0]].astype(float) - y = df[y_data[0]].astype(float) - z = np.polyfit(x, y, 2) - p = np.poly1d(z) - h, = ax.plot(x,p(x), linestyle="dashed", label="Trendline", color=trendline_color) - handles.append(h) + if trendline_equation != None: + try: + x = symbols('x') + p = sympify(trendline_equation) + x_range = np.linspace(df[x_data[0]].astype(float).min(), df[x_data[0]].astype(float).max(), 100) + y_range = [p.subs(x, i) for i in x_range] + h, = ax.plot(x_range, y_range, linestyle="dashed", label="Trendline", color=trendline_color) + handles.append(h) + except: + print("Invalid Equation") + else: + x = df[x_data[0]].astype(float) + y = df[y_data[0]].astype(float) + z = np.polyfit(x, y, 2) + p = np.poly1d(z) + h, = ax.plot(x,p(x), linestyle="dashed", label="Trendline", color=trendline_color) + handles.append(h) light_grey = 0.9 dar_grey = 0.4 @@ -140,7 +155,7 @@ async def create_plot(request: PlotDataRequest, data: Request): title=request.title, x_label=request.x_label, y_label=request.y_label, enable_trendline=request.enable_trendline, enable_grid=request.enable_grid, constant_line=constant_line, - x_axis_scale=request.x_axis_scale, y_axis_scale=request.y_axis_scale) + x_axis_scale=request.x_axis_scale, y_axis_scale=request.y_axis_scale, enable_trendline=request.enable_trendline) buf = BytesIO() fig.savefig(buf, format="png") -- cgit v1.2.3