rust-matrix 1.0.0-alpha

use crate::{pow::Pow, sqrt::Sqrt};
use std::{
    fmt::{self, Display},
    ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub},
};

/// These are the traits required for `ComplexNumber`
/// to work with `Matrix`
/// These are identical to `MatrixElementRequiredTraits`
/// and should be combined at some point
/// Note the trait `From<u8>` is required in order to
/// create identity matrices
pub trait ComplexNumberRequiredTraits<T>:
    Add<Output = T>
    + Sub<Output = T>
    + Mul<Output = T>
    + Div<Output = T>
    + Add<Output = T>
    + AddAssign
    + MulAssign
    + DivAssign
    + Clone
    + Copy
    + Display
    + PartialOrd<T>
    + Neg<Output = T>
    + Default
    + Sqrt
    + Pow
    + From<u8>
{
}

impl<
        T: Add<Output = T>
            + Sub<Output = T>
            + Mul<Output = T>
            + Div<Output = T>
            + Add<Output = T>
            + AddAssign
            + MulAssign
            + DivAssign
            + Clone
            + Copy
            + Display
            + PartialOrd<T>
            + Neg<Output = T>
            + Default
            + Sqrt
            + Pow
            + From<u8>
            + ?Sized,
    > ComplexNumberRequiredTraits<T> for T
{
}

#[derive(PartialEq, Debug, PartialOrd, Copy, Clone)]
pub struct ComplexNumber<T: ComplexNumberRequiredTraits<T>> {
    pub real: T,
    pub complex: T,
}

impl<T: ComplexNumberRequiredTraits<T>> ComplexNumber<T> {
    pub fn new(real: T, complex: T) -> ComplexNumber<T> {
        ComplexNumber { real, complex }
    }

    /// Creates a new complex number that is the complex conjugate
    /// of the original
    /// ```
    /// use matrix::complex_number::ComplexNumber;
    /// use rand::prelude::*;
    /// let mut rng = rand::thread_rng();
    /// let real = rng.gen_range(1.0..=100.0);
    /// let complex = rng.gen_range(1.0..=100.0);
    /// let complex_number = ComplexNumber::new(real, complex);
    ///
    /// assert_eq!(
    ///     complex_number.complex_conjugate(),
    ///     ComplexNumber::new(real, -complex,)
    /// );
    /// ```
    pub fn complex_conjugate(&self) -> Self {
        ComplexNumber {
            real: self.real,
            complex: -self.complex,
        }
    }

    /// Calculates the cartesian magnitude of a complex number
    /// Requires T to implement the `sqrt` trait which
    /// exposes the `square_root()` method
    /// ```
    /// use matrix::complex_number::ComplexNumber;
    /// let test_complex_number = ComplexNumber::new(3.0, 4.0);
    /// assert_eq!(test_complex_number.magnitude(), 5.0);
    /// ```
    pub fn magnitude(&self) -> T {
        ((self.real * self.real) + (self.complex * self.complex)).square_root()
    }

    /// Creates a new complex number that is a normalised
    /// copy of the first.
    /// ```
    ///
    /// use matrix::complex_number::ComplexNumber;
    /// let test_complex_number = ComplexNumber::new(3.0, 4.0);
    /// let normalised = test_complex_number.normalized();
    /// assert_eq!(normalised.real, 0.6);
    /// assert_eq!(normalised.complex, 0.8);
    /// ```
    pub fn normalized(&self) -> Self {
        let magnitude = self.magnitude();

        ComplexNumber::new(self.real / magnitude, self.complex / magnitude)
    }

    /// Calculates the square root of a complex number
    /// ```
    /// use matrix::complex_number::ComplexNumber;
    /// let complex = ComplexNumber::new(3.0, 4.0);
    ///
    /// assert_eq!(complex.square_root(), ComplexNumber::new(2.0, 1.0));
    /// ```
    pub fn square_root(&self) -> Self {
        let magnitude = self.magnitude();
        let real = ((magnitude + self.real) / 2.into()).square_root();
        let complex = if self.complex < T::default() {
            -((magnitude - self.real) / 2.into()).square_root()
        } else {
            ((magnitude - self.real) / 2.into()).square_root()
        };
        ComplexNumber::new(real, complex)
    }

    /// Implementation of the Pow trait for ComplexNumber
    ///
    pub fn powf(&self, n: T) -> Self {
        if n == T::default() {
            return ComplexNumber::new(T::from(1), T::default());
        } else {
            *self * self.powf(n - T::from(1))
        }
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Neg for ComplexNumber<T> {
    type Output = Self;

    fn neg(self) -> Self::Output {
        ComplexNumber {
            real: -self.real,
            complex: -self.complex,
        }
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Default for ComplexNumber<T> {
    fn default() -> Self {
        Self {
            real: Default::default(),
            complex: Default::default(),
        }
    }
}

fn complex_number_add<T: ComplexNumberRequiredTraits<T>>(
    lhs: &ComplexNumber<T>,
    rhs: &ComplexNumber<T>,
) -> ComplexNumber<T> {
    ComplexNumber {
        real: lhs.real + rhs.real,
        complex: lhs.complex + rhs.complex,
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Add for ComplexNumber<T> {
    type Output = Self;

    fn add(self, other: Self) -> Self {
        complex_number_add(&self, &other)
    }
}

impl<T: ComplexNumberRequiredTraits<T>> From<u8> for ComplexNumber<T> {
    fn from(value: u8) -> Self {
        ComplexNumber {
            real: T::from(value),
            complex: T::from(0),
        }
    }
}

fn complex_number_sub<T: ComplexNumberRequiredTraits<T>>(
    lhs: &ComplexNumber<T>,
    rhs: &ComplexNumber<T>,
) -> ComplexNumber<T> {
    ComplexNumber {
        real: lhs.real - rhs.real,
        complex: lhs.complex - rhs.complex,
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Sub for ComplexNumber<T> {
    type Output = Self;

    fn sub(self, other: Self) -> Self {
        complex_number_sub(&self, &other)
    }
}

fn complex_number_multiply<T: ComplexNumberRequiredTraits<T>>(
    lhs: &ComplexNumber<T>,
    rhs: &ComplexNumber<T>,
) -> ComplexNumber<T> {
    ComplexNumber {
        real: (lhs.real * rhs.real) - (lhs.complex * rhs.complex),
        complex: (lhs.real * rhs.complex) + (rhs.real * lhs.complex),
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Mul for ComplexNumber<T> {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self {
        complex_number_multiply::<T>(&self, &rhs)
    }
}

fn complex_number_divide<T: ComplexNumberRequiredTraits<T>>(
    lhs: &ComplexNumber<T>,
    rhs: &ComplexNumber<T>,
) -> ComplexNumber<T> {
    let rhs_complex_conjugate = rhs.complex_conjugate();
    let numerator = *lhs * rhs_complex_conjugate;
    let denominator = *rhs * rhs_complex_conjugate;

    if denominator.complex != T::default() {
        panic!(
            "Something went wrong. Denominator has complex element: {}",
            denominator
        );
    }

    ComplexNumber {
        real: numerator.real / denominator.real,
        complex: numerator.complex / denominator.real,
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Div for ComplexNumber<T> {
    type Output = Self;

    fn div(self, rhs: Self) -> Self::Output {
        complex_number_divide::<T>(&self, &rhs)
    }
}

fn complex_number_add_assign<T: ComplexNumberRequiredTraits<T>>(
    lhs: &ComplexNumber<T>,
    rhs: &ComplexNumber<T>,
) -> ComplexNumber<T> {
    complex_number_add(lhs, rhs)
}

impl<T: ComplexNumberRequiredTraits<T>> AddAssign for ComplexNumber<T> {
    fn add_assign(&mut self, rhs: Self) {
        *self = complex_number_add_assign::<T>(self, &rhs);
    }
}

fn complex_number_mul_assign<T: ComplexNumberRequiredTraits<T>>(
    lhs: &ComplexNumber<T>,
    rhs: &ComplexNumber<T>,
) -> ComplexNumber<T> {
    complex_number_multiply(lhs, rhs)
}

impl<T: ComplexNumberRequiredTraits<T>> MulAssign for ComplexNumber<T> {
    fn mul_assign(&mut self, rhs: Self) {
        *self = complex_number_mul_assign::<T>(self, &rhs)
    }
}

fn complex_number_div_assign<T: ComplexNumberRequiredTraits<T>>(
    lhs: &ComplexNumber<T>,
    rhs: &ComplexNumber<T>,
) -> ComplexNumber<T> {
    complex_number_divide(lhs, rhs)
}

impl<T: ComplexNumberRequiredTraits<T>> DivAssign for ComplexNumber<T> {
    fn div_assign(&mut self, rhs: Self) {
        *self = complex_number_div_assign::<T>(self, &rhs)
    }
}

fn format_complex_number_component<T: ComplexNumberRequiredTraits<T>>(
    input: &T,
    complex: bool,
) -> String {
    if *input == T::default() {
        return "".to_owned();
    } else if complex {
        if *input > T::default() {
            return "+".to_owned() + &input.to_string() + "i";
        } else {
            return input.to_string() + "i";
        }
    } else {
        input.to_string()
    }
}

impl<T: ComplexNumberRequiredTraits<T>> fmt::Display for ComplexNumber<T> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        let display_string = format_complex_number_component(&self.real, false)
            + " "
            + &format_complex_number_component(&self.complex, true);
        let _ = fmt.write_str(&display_string);
        Ok(())
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Sqrt for ComplexNumber<T> {
    fn square_root(&self) -> Self {
        self.square_root()
    }
}

impl<T: ComplexNumberRequiredTraits<T>> Pow for ComplexNumber<T> {
    fn pow(&self, n: ComplexNumber<T>) -> Self {
        self.powf(n.real)
    }
}

#[cfg(test)]
mod tests {
    use crate::complex_number::ComplexNumber;
    use rand::prelude::*;

    #[test]
    fn test_complex_number_add() {
        let mut rng = rand::thread_rng();
        let real_one = rng.gen_range(1.0..=100.0);
        let complex_one = rng.gen_range(1.0..=100.0);
        let real_two = rng.gen_range(1.0..=100.0);
        let complex_two = rng.gen_range(1.0..=100.0);

        let lhs = ComplexNumber::new(real_one, complex_one);
        let rhs = ComplexNumber::new(real_two, complex_two);

        assert_eq!(
            lhs + rhs,
            ComplexNumber::new(real_one + real_two, complex_one + complex_two)
        );
    }

    #[test]
    fn test_complex_number_sub() {
        let mut rng = rand::thread_rng();
        let real_one = rng.gen_range(1.0..=100.0);
        let complex_one = rng.gen_range(1.0..=100.0);
        let real_two = rng.gen_range(1.0..=100.0);
        let complex_two = rng.gen_range(1.0..=100.0);

        let lhs = ComplexNumber::new(real_one, complex_one);
        let rhs = ComplexNumber::new(real_two, complex_two);

        assert_eq!(
            lhs - rhs,
            ComplexNumber::new(real_one - real_two, complex_one - complex_two)
        );
    }

    #[test]
    fn test_complex_number_multiply() {
        let mut rng = rand::thread_rng();
        let real_one = rng.gen_range(1.0..=100.0);
        let complex_one = rng.gen_range(1.0..=100.0);
        let real_two = rng.gen_range(1.0..=100.0);
        let complex_two = rng.gen_range(1.0..=100.0);

        let lhs = ComplexNumber::new(real_one, complex_one);
        let rhs = ComplexNumber::new(real_two, complex_two);

        assert_eq!(
            lhs * rhs,
            ComplexNumber::new(
                (real_one * real_two) - (complex_one * complex_two),
                (complex_one * real_two) + (complex_two * real_one)
            )
        );
    }

    #[test]
    fn test_complex_conjugate() {
        let mut rng = rand::thread_rng();
        let real = rng.gen_range(1.0..=100.0);
        let complex = rng.gen_range(1.0..=100.0);
        let complex_number = ComplexNumber::new(real, complex);

        assert_eq!(
            complex_number.complex_conjugate(),
            ComplexNumber::new(real, -complex,)
        );
    }

    #[test]
    fn test_complex_number_divide() {
        let mut rng = rand::thread_rng();
        let real_one = rng.gen_range(1.0..=100.0);
        let complex_one = rng.gen_range(1.0..=100.0);
        let real_two = rng.gen_range(1.0..=100.0);
        let complex_two = rng.gen_range(1.0..=100.0);

        let lhs = ComplexNumber::new(real_one, complex_one);
        let rhs = ComplexNumber::new(real_two, complex_two);

        assert_eq!(
            lhs / rhs,
            ComplexNumber::new(
                ((real_one * real_two) + (complex_one * complex_two))
                    / (real_two * real_two + complex_two * complex_two),
                ((complex_one * real_two) - (complex_two * real_one))
                    / (real_two * real_two + complex_two * complex_two),
            )
        );
    }

    #[test]
    fn test_complex_number_sqrt() {
        let complex = ComplexNumber::new(3.0, 4.0);

        assert_eq!(complex.square_root(), ComplexNumber::new(2.0, 1.0));
    }
}