diff-eq-solver/solver.py at main · erichanwang/diff-eq-solver · 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
101
102
103
104
105
106
107
108
109
110
111
112
import sympy

def solve_homogeneous_ode(coeffs, initial_conditions):
    """
    Solves a homogeneous linear ordinary differential equation with constant coefficients.

    Args:
        coeffs (dict): A dictionary of coefficients for the derivatives of y with respect to t.
                       Example: {'3': 1, '2': -2, '1': -1, '0': 2} for y'''(t) - 2y''(t) - y'(t) + 2y(t) = 0
        initial_conditions (dict): A dictionary of initial conditions.
                                   Example: {'0': 4, '1': 6, '2': 10} for y(0)=4, y'(0)=6, y''(0)=10

    Returns:
        dict: A dictionary containing the solution in both 'latex' and 'plain' formats, or an 'error' key.
    """
    t = sympy.symbols('t')
    y = sympy.Function('y')(t)

    # Construct the differential equation
    diffeq = 0
    for order, coeff in coeffs.items():
        diffeq += coeff * sympy.diff(y, t, int(order))

    equation = sympy.Eq(diffeq, 0)

    # Construct the initial conditions dictionary for dsolve
    ics = {}
    if initial_conditions:
        for order_str, value in initial_conditions.items():
            order = int(order_str)
            if order == 0:
                ics[y.subs(t, 0)] = value
            else:
                ics[sympy.diff(y, t, order).subs(t, 0)] = value

    # Solve the differential equation
    try:
        if ics:
            solution = sympy.dsolve(equation, y, ics=ics)
        else:
            solution = sympy.dsolve(equation, y)

        # Prepare both LaTeX and plain string solutions
        latex_solution = sympy.latex(solution.rhs)
        plain_solution = str(solution.rhs)

        return {'latex': latex_solution, 'plain': plain_solution}
    except Exception as e:
        return {'error': f"An error occurred while solving: {e}"}

def solve_nonhomogeneous_ode(coeffs, rhs, initial_conditions):
    """
    Solves a nonhomogeneous linear ordinary differential equation with constant coefficients.

    Args:
        coeffs (dict): A dictionary of coefficients for the derivatives of y with respect to t.
                       Example: {'2': 1, '0': 1} for y''(t) + y(t) = sin(t)
        rhs (str): The right-hand side function as a string, e.g., 'sin(t)' or 't**2'.
        initial_conditions (dict): A dictionary of initial conditions.
                                   Example: {'0': 0, '1': 1} for y(0)=0, y'(0)=1

    Returns:
        dict: A dictionary containing the solution in both 'latex' and 'plain' formats, or an 'error' key.
    """
    t = sympy.symbols('t')
    y = sympy.Function('y')(t)

    # Construct the differential equation
    diffeq = 0
    for order, coeff in coeffs.items():
        diffeq += coeff * sympy.diff(y, t, int(order))

    # Parse the RHS
    try:
        rhs_expr = sympy.sympify(rhs)
    except (sympy.SympifyError, TypeError):
        return {'error': f"Invalid right-hand side expression: {rhs}"}

    equation = sympy.Eq(diffeq, rhs_expr)

    # Construct the initial conditions dictionary for dsolve
    ics = {}
    if initial_conditions:
        for order_str, value in initial_conditions.items():
            order = int(order_str)
            if order == 0:
                ics[y.subs(t, 0)] = value
            else:
                ics[sympy.diff(y, t, order).subs(t, 0)] = value

    # Solve the differential equation
    try:
        if ics:
            solution = sympy.dsolve(equation, y, ics=ics)
        else:
            solution = sympy.dsolve(equation, y)

        # Prepare both LaTeX and plain string solutions
        latex_solution = sympy.latex(solution.rhs)
        plain_solution = str(solution.rhs)

        return {'latex': latex_solution, 'plain': plain_solution}
    except Exception as e:
        return {'error': f"An error occurred while solving: {e}"}

if __name__ == '__main__':
    # Example: y'''(t) - 2y''(t) - y'(t) + 2y(t) = 0
    # y(0) = 4, y'(0) = 6, y''(0) = 10
    coeffs_example = {'3': 1, '2': -2, '1': -1, '0': 2}
    ics_example = {'0': 4, '1': 6, '2': 10}
    solution = solve_homogeneous_ode(coeffs_example, ics_example)
    print(f"The solution is: {solution}")