compNum.py

# Jack Eckert

# dependencies
import math

'''This basic program was created for practice, and is intented to be used as a module that
offers support for complex numbers, including many of the operations that can be performed on
them as taught in an intro to complex analysis class'''

class cmpx():
    def __init__(self, n=0, z=0):
        self.real = n # real component
        self.imaginary = z # imaginary component
        self.mod = math.sqrt(n ** 2 + z ** 2) # modulus or absolute value of the number
        try: # argument or angle the number makes with the real axis
            if self.real < 0 and self.imaginary == 0:
                self.arg = math.pi
            else: self.arg = math.atan(self.imaginary / self.real)
        except ZeroDivisionError:
            if self.imaginary > 0:
                self.arg = math.pi / 2
            elif self.imaginary < 0:
                self.arg = (math.pi / 2) * 3

    # formats complex numbers as a string
    def __str__(self):
        if self.imaginary >= 0:
            return f'{round(self.real, 6)} + {round(self.imaginary, 6)}i'
        else:
            return f'{round(self.real, 6)} - {round(-self.imaginary, 6)}i'
        
    # defines equality for complex numbers
    def __eq__(self, Z):
        if self.real == Z.real and self.imaginary == Z.imaginary:
            return True
        else:
            return False
        
    # defines negation for complex numbers
    def __neg__(self):
        return cmpx(-self.real, -self.imaginary)

    # defines addition for complex numbers
    def __add__(self, Z):
        if type(Z) == cmpx:
            return cmpx(self.real + Z.real, self.imaginary + Z.imaginary)
        else:
            return cmpx(self.real + Z, self.imaginary)
        
    def __radd__(self, Z):
        return cmpx(self.real + Z, self.imaginary)
    
    # defines subtraction for complex numbers
    def __sub__(self, Z):
        if type(Z) == cmpx:
            return cmpx(self.real - Z.real, self.imaginary - Z.imaginary)
        else:
            return cmpx(self.real - Z, self.imaginary)
        
    def __rsub__(self, Z):
        return cmpx(self.real - Z, self.imaginary)
    
    # defines multiplication for complex numbers
    def __mul__(self, Z):
        if type(Z) == cmpx:
            return cmpx(self.real * Z.real - self.imaginary * Z.imaginary, self.real * Z.imaginary + self.imaginary * Z.real)
        else:
            return cmpx(self.real * Z, self.imaginary * Z)
        
    def __rmul__(self, Z):
        return cmpx(self.real * Z, self.imaginary * Z)
    
    # defines division for complex numbers
    def __truediv__(self, Z):
        if type(Z) == cmpx:
            num = self * Z.conjugate()
            return cmpx(num.real / self.mod, num.imaginary / self.mod)
        else:
            return cmpx(self.real / Z, self.imaginary / Z)
        
    def __rtruediv__(self, Z):
        return math.pow(self, -1) * Z
        
    # returns the modulus when the abs() function is called on a complex number
    def __abs__(self):
        return self.mod
    
    # rounds each component of the complex number to the specified number of decimal places
    def __round__(self, n=0):
        return cmpx(round(self.real, n), round(self.imaginary, n))

    # defines exponentiation for complex numbers
    def __pow__(self, Z):
        if type(Z) == cmpx:
            rePow = self ** Z.real
            ang = ln(self) * Z.imaginary
            return (cos(ang) + cmpx(0, 1) * sin(ang)) * rePow
        elif type(Z) == int:
            if Z == 0:
                return 1
            if Z > 0:
                self.result = self
                for i in range(Z - 1):
                    self.result *= self
                return self.result
            if Z < 0:
                return cmpx(0, 1) / math.pow(self, -Z)
        elif type(Z) == float:
            if Z == 0:
                return 1
            if Z > 0:
                itn = self ** math.floor(Z)
                dec = Z - math.floor(Z)
                outAbs = math.pow(self.mod, Z)
                return itn * cmpx(outAbs * math.cos(self.arg * dec), outAbs * math.sin(self.arg * dec))
            if Z < 0:
                return cmpx(0, 1) / (self ** -Z)

    def __rpow__(self, x):
        ang = math.log(x) * self.imaginary
        return (math.cos(ang) + cmpx(0, 1) * math.sin(ang)) * math.pow(x, self.real)
  
    # Gives a tuple of all of the roots of a complex number, where "n" is the power of the polynomial
    def roots(self, n):
        root = self ** (1 / n)
        outList = [root]
        ang = root.arg
        for i in range(n - 1):
            ang += (2 * math.pi) / (n)
            outList.append(cmpx(root.mod * math.cos(ang), root.mod * math.sin(ang)))
        return tuple(outList)
    
    # Returns the complex conjugate of the complex number
    def conj(self):
        return cmpx(self.real, -self.imaginary)

# redefines math.log to work with complex numbers
def ln(Z):
    if type(Z) == cmpx:
        return cmpx(math.log(Z.mod), Z.arg)
    else:
        return math.log(Z)
    
# redefines math.sin to work with complex numbers
def sin(Z):
    if type(Z) == cmpx:
        return ((math.e ** (-Z * cmpx(0, 1))) -(math.e ** (Z * cmpx(0, 1)))) / 2 * cmpx(0, 1)
    else:
        return math.sin(Z)
    
# redefines math.cos to work with complex numbers
def cos(Z):
    if type(Z) == cmpx:
        return ((math.e ** (Z * cmpx(0, 1))) + (math.e ** (-Z * cmpx(0, 1)))) / 2
    else:
        return math.cos(Z)

# redefines math.tan to work with complex numbers
def tan(Z):
    return sin(Z) / cos(Z)
    
# finds the roots of unity given a degree of polynomial
def uniRoots(n):
    return cmpx(1, 0).roots(n)

# finds the complex number with a modulus of 1 and an argument of the angle given
def unitZ(theta):
    return cmpx(cos(theta), sin(theta))