-
Notifications
You must be signed in to change notification settings - Fork 0
/
pyquartic_nonumba.py
132 lines (111 loc) · 3.5 KB
/
pyquartic_nonumba.py
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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
"""
Modified Ferrari's quartic solver and modified Cardano's cubic solver for Python (4th and 3rd order polynomials)
Using modified algorithms from: https://quarticequations.com
Functions can optionally be 'numbarized' by just installing numba
Python implementation by: mzivic7
Source code: https://github.com/mzivic7/pyquartic
Issues: https://github.com/mzivic7/pyquartic/issues
Licence: GNU General Public License v3.0
"""
import math
import cmath
# constants
sq3 = math.sqrt(3)
pi23 = 2 * math.pi / 3
def solve_cubic(a, b, c, d):
"""Solves cubic equation with modified Cardano's method.
az^3 + bz^2 + cz + d = 0
Uses modified Cardano's method from 'Numerical Recipes' and Viète’s trigonometric method to avoid large error in ceratin cases.
https://quarticequations.com/Cubic.pdf"""
# convert to depressed form
a2, a1, a0 = b/a, c/a, d/a
q = a1/3 - a2**2/9
r = (a1*a2-3*a0)/6 - a2**3/27
rq = r**2 + q**3
if rq > 0:
# Numerical Recipes algorithm
aa = (abs(r) + math.sqrt(rq))**(1/3)
if r >= 0:
t = aa - q / aa
else:
t = q/aa - aa
z1 = t - a2/3
x = -t/2 - a2/3
y = (sq3/2) * (aa + (q/aa))
z2 = complex(x, y)
z3 = complex(x, -y)
else:
# Viète algorithm
if q == 0:
theta = 0
elif q < 0:
theta = math.acos(r/(-q)**(3/2))
# repeating stuff
m = 2 * math.sqrt(-q)
n = a2/3
# solutions
phi1 = theta / 3
phi2 = phi1 - pi23
phi3 = phi1 + pi23
z1 = complex(m * math.cos(phi1) - n)
z2 = complex(m * math.cos(phi2) - n)
z3 = complex(m * math.cos(phi3) - n)
return (z1, z2, z3)
def solve_cubic_one(a, b, c):
"""Calculates only one real root for depressed cubic equation.
z^3 + az^2 + bz + c = 0
Uses modified Cardano's method from 'Numerical Recipes' and Viète’s trigonometric method to avoid large error in ceratin cases.
https://quarticequations.com/Cubic.pdf"""
q = b/3 - a**2/9
r = (b*a-3*c)/6 - a**3/27
rq = r**2 + q**3
if rq > 0:
# Numerical Recipes algorithm
aa = (abs(r) + math.sqrt(rq))**(1/3)
if r >= 0:
t = aa - q/aa
else:
t = q/aa - aa
z1 = t - a/3
else:
# Viete algorithm
if q == 0:
theta = 0
elif q < 0:
theta = math.acos(r/(-q)**(3/2))
fi = theta/3
z1 = 2 * math.sqrt(-q) * math.cos(fi) - a/3
return z1
def solve_quartic(a, b, c, d, e):
"""Solves quartic equation with modified Ferrari's method.
az^4 + bz^3 + cz^2 + dz + e = 0
https://quarticequations.com/Quartic2.pdf"""
# convert to depressed form
a3, a2, a1, a0 = b/a, c/a, d/a, e/a
cc = a3/4
b2 = a2 - 6*cc**2
b1 = a1 - 2*a2*cc + 8*cc**3
b0 = a0 - a1*cc + a2*cc**2 - 3*cc**4
# one real root of Ferrari's resolvent cubic
y = solve_cubic_one(b2, b2**2/4 - b0, -b1**2/8)
if y < 0:
y = 0
s = y**2 + b2*y + b2**2/4 - b0
# protection from root of negative number
if s > 0:
if b1 < 0:
r = -math.sqrt(s)
else:
r = math.sqrt(s)
else:
r = float("nan")
# repeating stuff
p = cmath.sqrt(y/2) - cc
p1 = -cmath.sqrt(y/2) - cc
q = -y/2 - b2/2
# solutions to depressed quartic equation
z1 = p + cmath.sqrt(q - r)
z2 = p - cmath.sqrt(q - r)
z3 = p1 + cmath.sqrt(q + r)
z4 = p1 - cmath.sqrt(q + r)
return (z1, z2, z3, z4)