codebase / stimlib

/// CUDA compatible complex number class

#ifndef STIM_COMPLEX
#define STIM_COMPLEX

#include "../cuda/cudatools/callable.h"
#include <cmath>
#include <string>
#include <sstream>
#include <iostream>

namespace stim
{
    enum complexComponentType {complexFull, complexReal, complexImaginary, complexMag, complexIntensity};

template <class T>
struct complex
{
    T r, i;

    //default constructor
    CUDA_CALLABLE complex()
    {
        r = 0;
	   i = 0;
    }

    //constructor when given real and imaginary values
    CUDA_CALLABLE complex(T r, T i = 0)
    {
        this->r = r;
        this->i = i;
    }

	//access methods
	CUDA_CALLABLE T real()
	{
		return r;
	}

	CUDA_CALLABLE T real(T r_val)
	{
		r = r_val;
		return r_val;
	}

	CUDA_CALLABLE T imag()
	{
		return i;
	}
	CUDA_CALLABLE T imag(T i_val)
	{
		i = i_val;
		return i_val;
	}

    

    //return the current value multiplied by i
    CUDA_CALLABLE complex<T> imul()
    {
        complex<T> result;
        result.r = -i;
        result.i = r;

        return result;
    }

    //returns the complex signum (-1, 0, 1)
    CUDA_CALLABLE int sgn(){
        if(r > 0) return 1;
        else if(r < 0) return -1;
        else return (0 < i - i < 0);
    }

	//ARITHMETIC OPERATORS--------------------

    //binary + operator (returns the result of adding two complex values)
    CUDA_CALLABLE complex<T> operator+ (const complex<T> rhs) const
    {
        complex<T> result;
        result.r = r + rhs.r;
        result.i = i + rhs.i;
        return result;
    }

	CUDA_CALLABLE complex<T> operator+ (const T rhs) const
    {
        complex<T> result;
        result.r = r + rhs;
        result.i = i;
        return result;
    }

    //binary - operator (returns the result of adding two complex values)
    CUDA_CALLABLE complex<T> operator- (const complex<T> rhs) const
    {
        complex<T> result;
        result.r = r - rhs.r;
        result.i = i - rhs.i;
        return result;
    }

    //binary - operator (returns the result of adding two complex values)
    CUDA_CALLABLE complex<T> operator- (const T rhs)
    {
        complex<T> result;
        result.r = r - rhs;
        result.i = i;
        return result;
    }

    //binary MULTIPLICATION operators (returns the result of multiplying complex values)
    CUDA_CALLABLE complex<T> operator* (const complex<T> rhs) const
    {
        complex<T> result;
        result.r = r * rhs.r - i * rhs.i;
        result.i = r * rhs.i + i * rhs.r;
        return result;
    }
    CUDA_CALLABLE complex<T> operator* (const T rhs)
    {
        return complex<T>(r * rhs, i * rhs);
    }

    //binary DIVISION operators (returns the result of dividing complex values)
    CUDA_CALLABLE complex<T> operator/ (const complex<T> rhs) const
    {
        complex<T> result;
        T denom = rhs.r * rhs.r + rhs.i * rhs.i;
        result.r = (r * rhs.r + i * rhs.i) / denom;
        result.i = (- r * rhs.i + i * rhs.r) / denom;

        return result;
    }
    CUDA_CALLABLE complex<T> operator/ (const T rhs)
    {
        return complex<T>(r / rhs, i / rhs);
    }

    //ASSIGNMENT operators-----------------------------------
    CUDA_CALLABLE complex<T> & operator=(const complex<T> &rhs)
    {
        //check for self-assignment
        if(this != &rhs)
        {
            this->r = rhs.r;
            this->i = rhs.i;
        }
        return *this;
    }
    CUDA_CALLABLE complex<T> & operator=(const T &rhs)
    {
        	this->r = rhs;
        	this->i = 0;

		return *this;
    }

    //arithmetic assignment operators
    CUDA_CALLABLE complex<T> operator+=(const complex<T> &rhs)
    {
		*this = *this + rhs;
        	return *this;
    }
    CUDA_CALLABLE complex<T> operator+=(const T &rhs)
    {
		*this = *this + rhs;
        	return *this;
    }

	CUDA_CALLABLE complex<T> operator-=(const complex<T> &rhs)
    {
		*this = *this - rhs;
        	return *this;
    }
    CUDA_CALLABLE complex<T> operator-=(const T &rhs)
    {
		*this = *this - rhs;
        	return *this;
    }

    CUDA_CALLABLE complex<T> operator*=(const complex<T> &rhs)
    {
		*this = *this * rhs;
        	return *this;
    }
	CUDA_CALLABLE complex<T> operator*=(const T &rhs)
    {
		*this = *this * rhs;
        	return *this;
    }
	//divide and assign
	CUDA_CALLABLE complex<T> operator/=(const complex<T> &rhs)
    {
		*this = *this / rhs;
        	return *this;
    }
    CUDA_CALLABLE complex<T> operator/=(const T &rhs)
    {
		*this = *this / rhs;
        	return *this;
    }

    //absolute value operator (returns the absolute value of the complex number)
	CUDA_CALLABLE T abs()
	{
		return std::sqrt(r * r + i * i);
	}

	CUDA_CALLABLE complex<T> log()
	{
		complex<T> result;
		result.r = (T)std::log(std::sqrt(r * r + i * i));
		result.i = (T)std::atan2(i, r);


		return result;
	}

	CUDA_CALLABLE complex<T> exp()
	{
		complex<T> result;

		T e_r = std::exp(r);
		result.r = e_r * (T)std::cos(i);
		result.i = e_r * (T)std::sin(i);

		return result;
	}

	CUDA_CALLABLE complex<T> pow(T y)
	{
		complex<T> result;

		result = log() * y;

		return result.exp();
	}

	CUDA_CALLABLE complex<T> sqrt()
	{
		complex<T> result;

		//convert to polar coordinates
		T a = std::sqrt(r*r + i*i);
		T theta = std::atan2(i, r);

		//find the square root
		T a_p = std::sqrt(a);
		T theta_p = theta/2.0f;

		//convert back to cartesian coordinates
		result.r = a_p * std::cos(theta_p);
		result.i = a_p * std::sin(theta_p);

		return result;
	}

	std::string str()
	{
		std::stringstream ss;
		ss<<"("<<r<<","<<i<<")";

		return ss.str();
	}

	//COMPARISON operators
	CUDA_CALLABLE bool operator==(complex<T> rhs)
	{
        if(r == rhs.r && i == rhs.i)
            return true;
        return false;
    }

    CUDA_CALLABLE bool operator==(T rhs)
	{
        if(r == rhs && i == 0)
            return true;
        return false;
    }

    CUDA_CALLABLE bool operator!=(T rhs)
    {
        if(r != rhs || i != 0)
            return true;
        return false;
    }

    CUDA_CALLABLE bool operator<(complex<T> rhs){
    	return abs() < rhs.abs();
    }
    CUDA_CALLABLE bool operator<=(complex<T> rhs){
    	return abs() <= rhs.abs();
    }
    CUDA_CALLABLE bool operator>(complex<T> rhs){
    	return abs() > rhs.abs();
    }
    CUDA_CALLABLE bool operator >=(complex<T> rhs){
    	return abs() >= rhs.abs();
    }

	//CASTING operators
	template < typename otherT >
	operator complex<otherT>()
	{
		complex<otherT> result((otherT)r, (otherT)i);
		return result;
	}
	template< typename otherT >
	complex( const complex<otherT> &rhs)
	{
		r = (T)rhs.r;
		i = (T)rhs.i;
	}
	template< typename otherT >
	complex& operator=(const complex<otherT> &rhs)
	{
		r = (T)rhs.r;
		i = (T)rhs.i;
		return *this;
	}

	

};

/// Cast an array of complex values to an array of real values
template<typename T>
static void real(T* r, complex<T>* c, size_t n){
	for(size_t i = 0; i < n; i++)
		r[i] = c[i].real();
}

/// Cast an array of complex values to an array of real values
template<typename T>
static void imag(T* r, complex<T>* c, size_t n){
	for(size_t i = 0; i < n; i++)
		r[i] = c[i].imag();
}

/// Calculate the magnitude of an array of complex values
template<typename T>
static void abs(T* m, complex<T>* c, size_t n){
	for(size_t i = 0; i < n; i++)
		m[i] = c[i].abs();
}

/// Calculate the intensity of an array of complex values
template<typename T>
static void intensity(T* m, complex<T>* c, size_t n){
	for(size_t i = 0; i < n; i++)
		m[i] = pow(c[i].abs(), 2);
}

}	//end RTS namespace

//addition
template<typename T>
CUDA_CALLABLE static stim::complex<T> operator+(const double a, const stim::complex<T> b)
{
    return stim::complex<T>((T)a + b.r, b.i);
}

//subtraction with a real value
template<typename T>
CUDA_CALLABLE static stim::complex<T> operator-(const double a, const stim::complex<T> b)
{
    return stim::complex<T>((T)a - b.r, -b.i);
}

//minus sign
template<typename T>
CUDA_CALLABLE static stim::complex<T> operator-(const stim::complex<T> &rhs)
{
    return stim::complex<T>(-rhs.r, -rhs.i);
}

//multiply a T value by a complex value
template<typename T>
CUDA_CALLABLE static stim::complex<T> operator*(const double a, const stim::complex<T> b)
{
    return stim::complex<T>((T)a * b.r, (T)a * b.i);
}

//divide a T value by a complex value
template<typename T>
CUDA_CALLABLE static stim::complex<T> operator/(const double a, const stim::complex<T> b)
{
    stim::complex<T> result;

    T denom = b.r * b.r + b.i * b.i;

    result.r = ((T)a * b.r) / denom;
    result.i = -((T)a * b.i) / denom;

    return result;
}


template<typename T>
CUDA_CALLABLE static stim::complex<T> pow(stim::complex<T> x, T y)
{
	return x.pow(y);
}
template<typename T>
CUDA_CALLABLE static stim::complex<T> pow(stim::complex<T> x, int y)
{
	return x.pow(y);
}

//log function
template<typename T>
CUDA_CALLABLE static stim::complex<T> log(stim::complex<T> x)
{
	return x.log();
}

//exp function
template<typename T>
CUDA_CALLABLE static stim::complex<T> exp(stim::complex<T> x)
{
	return x.exp();
}

//sqrt function
template<typename T>
CUDA_CALLABLE static stim::complex<T> sqrt(stim::complex<T> x)
{
	return x.sqrt();
}


template <typename T>
CUDA_CALLABLE static T abs(stim::complex<T> a)
{
    return a.abs();
}

template <typename T>
CUDA_CALLABLE static T real(stim::complex<T> a)
{
    return a.r;
}

//template <typename T>
CUDA_CALLABLE static float real(float a)
{
    return a;
}

template <typename T>
CUDA_CALLABLE static T imag(stim::complex<T> a)
{
    return a.i;
}

template<class A>
CUDA_CALLABLE stim::complex<A> sin(const stim::complex<A> x)
{
	stim::complex<A> result;
	result.r = (A)std::sin(x.r) * (A)std::cosh(x.i);
	result.i = (A)std::cos(x.r) * (A)std::sinh(x.i);

	return result;
}

template<class A>
CUDA_CALLABLE stim::complex<A> cos(const stim::complex<A> x)
{
	stim::complex<A> result;
	result.r = (A)std::cos(x.r) * (A)std::cosh(x.i);
	result.i = -((A)std::sin(x.r) * (A)std::sinh(x.i));

	return result;
}


template<class A>
std::ostream& operator<<(std::ostream& os, stim::complex<A> x)
{
    os<<x.str();
    return os;
}

template<class A>
std::istream& operator>>(std::istream& is, stim::complex<A>& x)
{
    A r, i;
	r = i = 0;		//initialize the real and imaginary parts to zero
    is>>r;			//parse
    is>>i;

    x.real(r);		//assign the parsed values to x
    x.imag(i);

    return is;		//return the stream
}

#endif