option to manually specify equation

1 files changed, 27 insertions, 8 deletions

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