math/complex.h

/*RTS Complex number class.  This class is CUDA compatible,
and can therefore be used in CUDA code and on CUDA devices.
*/
#ifndef RTS_COMPLEX
#define RTS_COMPLEX
#include "../cuda/callable.h"
#include <cmath>
#include <string>
#include <sstream>
#include <iostream>
namespace stim
{
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(int y)
	{
        return pow((double)y);
	}*/
	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;
	}
};
}	//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;
}
//trigonometric functions
//template<class A>
/*CUDA_CALLABLE static stim::complex<float> sinf(const stim::complex<float> x)
{
	stim::complex<float> result;
	result.r = sinf(x.r) * coshf(x.i);
	result.i = cosf(x.r) * sinhf(x.i);
	return result;
}*/
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;
}
//floating point template
//template<class A>
/*CUDA_CALLABLE static stim::complex<float> cosf(const stim::complex<float> x)
{
	stim::complex<float> result;
	result.r = cosf(x.r) * coshf(x.i);
	result.i = -(sinf(x.r) * sinhf(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
}
//#if __GNUC__ > 3 && __GNUC_MINOR__ > 7
//template<class T> using rtsComplex = stim::complex<T>;
//#endif
#endif