option to manually specify equation

3 files changed, 55 insertions, 16 deletions

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")
diff --git a/app.py b/app.py
index 161d429..39515e1 100644
--- a/app.py
+++ b/app.py
@@ -11,6 +11,8 @@ import matplotlib.font_manager as font_manager
 
 import re
 
+from sympy import symbols, sympify
+
 font_path = 'times_new_roman.ttf'
 times_new_roman = font_manager.FontProperties(fname=font_path, style='normal')
 
@@ -27,7 +29,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 = []
@@ -50,12 +52,24 @@ 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)
+                print("Valid Equation", p)
+            except ValueError:
+                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
@@ -212,6 +226,10 @@ def process_data():
     else:
         enable_trendline = False
 
+    trendline_equation = request.form.get('trendlineEquation', None)
+    if trendline_equation == "":
+        trendline_equation = None
+
     fig = plot_data(x_data, y_data, std_dev_data, color_picker, data_series_label, df, title=plot_title,
                     x_label=x_axis_label, y_label=y_axis_label,
                     plot_background_color=plot_background_color,
@@ -219,7 +237,8 @@ def process_data():
                     enable_trendline=enable_trendline,
                     trendline_color=color_picker_trendline,
                     x_axis_scale=x_axis_scale,
-                    y_axis_scale=y_axis_scale)
+                    y_axis_scale=y_axis_scale,
+                    trendline_equation=trendline_equation)
 
     # Return plot as image
     from io import BytesIO
diff --git a/templates/index.html b/templates/index.html
index 4fde9f2..75153ce 100644
--- a/templates/index.html
+++ b/templates/index.html
@@ -193,6 +193,11 @@
         </div>
     </div>
 
+    <div class="row mb-3">
+        <label for="trendlineEquation" class="form-label">Trendline Equation (Leave blank if you want it to be auto-calculated)</label>
+        <input type="text" name="trendlineEquation" id="trendlineEquation" class="form-control" placeholder="">
+    </div>
+
         <div class="mb-3 row">
             <label for="colorPickerPlotBackground" class="form-label">Plot Background Color</label>
             <input name="colorPickerPlotBackground" class="form-control" id="colorPickerPlotBackground" value="#ffffff" data-coloris>