def roots(a, b, c):
check = (b**2) - (4 * a * c)
if check > 0: # 2 real roots
root1 = ((-b) + math.sqrt(check)) / (a * 2)
root2 = ((-b) - math.sqrt(check)) / (a * 2)
tup = (root1, root2)
return tup
elif check < 0: # 2 complex root
root1 = Complex()
root2 = Complex()
root1.im = math.sqrt(abs(check)) / (a * 2)
root2.im = -(math.sqrt(abs(check)) / (a * 2))
root1.re = (-b) / (a * 2)
root2.re = (-b) / (a * 2)
tup = (root1, root2)
return tup
else: # 1 real root
root1 = (-b) / (a * 2)
tup = (root1, )
return (tup)
def test_sqrt(self):
a = Complex(3,2)
a = sqrt(a)
b = complex(3,2)
b = cmath.sqrt(b)
self.assertEqual(a.re,b.real)
self.assertEqual(a.im,b.imag)
a = Complex(-3,-2)
a = sqrt(a)
b = complex(-3,-2)
b = cmath.sqrt(b)
self.assertEqual(a.re,b.real)
self.assertEqual(a.im,b.imag)
a = Complex(3,0)
a = sqrt(a)
b = complex(3,0)
b = cmath.sqrt(b)
self.assertEqual(a.re,b.real)
self.assertEqual(a.im,b.imag)
a = Complex(0,3)
a = sqrt(a)
a.re = round(a.re,2)
a.im = round(a.im,2)
b = complex(0,3)
b = cmath.sqrt(b)
breal = round(b.real,2)
bimag = round(b.imag,2)
self.assertEqual(a.re,breal)
self.assertEqual(a.im,bimag)
def dft2(points):
complex_list = [Complex(p[0], p[1]) for p in points]
N = len(complex_list)
X = []
for k in range(N):
x = Complex(0, 0)
for n in range(N):
t = 2 * math.pi * k * n / N
c = Complex(math.cos(t), -math.sin(t))
x += c * complex_list[n]
x.re = x.re / N
x.im = x.im / N
radius = x.length()
phase = x.angle()
if radius > 1:
X.append((radius, k, phase))
return X