calc/plotting.py at main · prettyMatthew/calc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
import matplotlib.pyplot as plt
import numpy as np
from sympy import symbols, sympify, lambdify, SympifyError
# 3D plotting specific import
from mpl_toolkits.mplot3d import Axes3D
from utils import show_error, show_info


def plot_functions(functions_str, var_str, x_range_str):
    try:
        x = symbols(var_str)

        functions = functions_str.split(';')

        x_start, x_end = map(float, x_range_str.split(','))
        x_vals = np.linspace(x_start, x_end, 500)  # 500 points for smooth curve

        plt.figure(figsize=(9, 7))
        plt.title("2D Function Plot")
        plt.xlabel(f"{var_str} axis")
        plt.ylabel("y axis")
        plt.grid(True)
        plt.axhline(0, color='gray', linewidth=0.8, linestyle='--')
        plt.axvline(0, color='gray', linewidth=0.8, linestyle='--')

        plotted_any = False
        for func_str in functions:
            func_str = func_str.strip()
            if not func_str:
                continue

            try:
                expr = sympify(func_str)
                f_numerical = lambdify(x, expr, 'numpy')
                y_vals = f_numerical(x_vals)

                y_vals = np.where(np.isinf(y_vals), np.nan, y_vals)

                plt.plot(x_vals, y_vals, label=f"y = {func_str}")
                plotted_any = True
            except SympifyError:
                show_error(f"Invalid function format: '{func_str}'. Skipping this function.")
            except Exception as e:
                show_error(f"Error plotting '{func_str}': {e}. Skipping this function.")

        if plotted_any:
            plt.legend()
            plt.tight_layout()
            plt.show()
        else:
            show_info("No valid functions were provided or could be plotted for 2D.")

    except ValueError:
        show_error("Invalid range format. Use 'start,end' (e.g., -5,5).")
    except Exception as e:
        show_error(f"An error occurred during 2D plotting setup: {e}")


def plot_3d_function(function_str, x_var_str, y_var_str, x_range_str, y_range_str, cmap='viridis'):
    try:
        # Define symbolic variables
        x_sym, y_sym = symbols(x_var_str + ' ' + y_var_str)

        # Sympyfy the expression
        expr = sympify(function_str)

        # Convert sympy expression to a numerical function for numpy
        f_numerical = lambdify((x_sym, y_sym), expr, 'numpy')

        # Parse ranges
        x_start, x_end = map(float, x_range_str.split(','))
        y_start, y_end = map(float, y_range_str.split(','))

        # Create meshgrid for x and y values
        X = np.linspace(x_start, x_end, 50)  # Use fewer points for 3D for performance
        Y = np.linspace(y_start, y_end, 50)
        X, Y = np.meshgrid(X, Y)

        # Calculate Z values
        Z = f_numerical(X, Y)

        # Create 3D plot
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection='3d')
        ax.plot_surface(X, Y, Z, cmap=cmap, edgecolor='none')

        ax.set_title(f"3D Plot of z = {function_str}")
        ax.set_xlabel(f"{x_var_str} axis")
        ax.set_ylabel(f"{y_var_str} axis")
        ax.set_zlabel("z axis")

        plt.tight_layout()
        plt.show()

    except SympifyError:
        show_error("Invalid function format. (e.g., x**2 + y**2, sin(x)*cos(y))")
    except ValueError:
        show_error("Invalid range format. Use 'start,end' (e.g., -5,5).")
    except Exception as e:
        show_error(f"An error occurred during 3D plotting: {e}")