codebase / stimlib

#ifndef RTS_VECTOR_H
#define RTS_VECTOR_H

#include <iostream>
#include <cmath>
#include <sstream>
//#include "rts/math/point.h"
#include "../cuda/callable.h"


namespace stim
{



template <class T, int N=3>
struct vec
{
	T v[N];

	CUDA_CALLABLE vec()
	{
		//memset(v, 0, sizeof(T) * N);
		for(int i=0; i<N; i++)
			v[i] = 0;
	}

	//efficiency constructor, makes construction easier for 1D-4D vectors
	CUDA_CALLABLE vec(T rhs)
	{
		for(int i=0; i<N; i++)
			v[i] = rhs;
	}
	CUDA_CALLABLE vec(T x, T y)
	{
		v[0] = x;
		v[1] = y;
	}
	CUDA_CALLABLE vec(T x, T y, T z)
	{
		v[0] = x;
		v[1] = y;
		v[2] = z;
	}
	CUDA_CALLABLE vec(T x, T y, T z, T w)
	{
		v[0] = x;
		v[1] = y;
		v[2] = z;
		v[3] = w;
	}

	//copy constructor
	CUDA_CALLABLE vec( const vec<T, N>& other){
		for(int i=0; i<N; i++)
			v[i] = other.v[i];
	}

	
	template< typename U >
	CUDA_CALLABLE operator vec<U, N>(){
		vec<U, N> result;
		for(int i=0; i<N; i++)
			result.v[i] = v[i];

		return result;
	}

	//template<class U>
	//friend vec<U, N>::operator vec<T, N>();

	CUDA_CALLABLE T len() const
	{
        //compute and return the vector length
        T sum_sq = (T)0;
        for(int i=0; i<N; i++)
        {
            sum_sq += v[i] * v[i];
        }
        return sqrt(sum_sq);

	}

	CUDA_CALLABLE vec<T, N> cart2sph() const
	{
		//convert the vector from cartesian to spherical coordinates
		//x, y, z -> r, theta, phi (where theta = 0 to 2*pi)

		vec<T, N> sph;
		sph[0] = std::sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2]);
		sph[1] = std::atan2(v[1], v[0]);
		sph[2] = std::acos(v[2] / sph[0]);

		return sph;
	}

	CUDA_CALLABLE vec<T, N> sph2cart() const
	{
		//convert the vector from cartesian to spherical coordinates
		//r, theta, phi -> x, y, z (where theta = 0 to 2*pi)

		vec<T, N> cart;
		cart[0] = v[0] * std::cos(v[1]) * std::sin(v[2]);
		cart[1] = v[0] * std::sin(v[1]) * std::sin(v[2]);
		cart[2] = v[0] * std::cos(v[2]);

		return cart;
	}

	CUDA_CALLABLE vec<T, N> norm() const
	{
        //compute and return the vector norm
        vec<T, N> result;

        //compute the vector length
        T l = len();

        //normalize
        for(int i=0; i<N; i++)
        {
            result.v[i] = v[i] / l;
        }

        return result;
	}

	CUDA_CALLABLE vec<T, 3> cross(const vec<T, 3> rhs) const
	{
		vec<T, 3> result;

		//compute the cross product (only valid for 3D vectors)
		result[0] = v[1] * rhs.v[2] - v[2] * rhs.v[1];
		result[1] = v[2] * rhs.v[0] - v[0] * rhs.v[2];
		result[2] = v[0] * rhs.v[1] - v[1] * rhs.v[0];

		return result;
	}

    CUDA_CALLABLE T dot(vec<T, N> rhs) const
    {
        T result = (T)0;

        for(int i=0; i<N; i++)
            result += v[i] * rhs.v[i];

        return result;

    }

	//arithmetic
	CUDA_CALLABLE vec<T, N> operator+(vec<T, N> rhs) const
	{
		vec<T, N> result;

		for(int i=0; i<N; i++)
		    result.v[i] = v[i] + rhs.v[i];

		return result;
	}
	CUDA_CALLABLE vec<T, N> operator+(T rhs) const
	{
		vec<T, N> result;

		for(int i=0; i<N; i++)
		    result.v[i] = v[i] + rhs;

		return result;
	}
	CUDA_CALLABLE vec<T, N> operator-(vec<T, N> rhs) const
	{
        vec<T, N> result;

        for(int i=0; i<N; i++)
            result.v[i] = v[i] - rhs.v[i];

        return result;
	}
	CUDA_CALLABLE vec<T, N> operator*(T rhs) const
	{
        vec<T, N> result;

        for(int i=0; i<N; i++)
            result.v[i] = v[i] * rhs;

        return result;
	}
	CUDA_CALLABLE vec<T, N> operator/(T rhs) const
	{
        vec<T, N> result;

        for(int i=0; i<N; i++)
            result.v[i] = v[i] / rhs;

        return result;
	}
	CUDA_CALLABLE vec<T, N> operator*=(T rhs){
		for(int i=0; i<N; i++)
			v[i] = v[i] * rhs;
		return *this;
	}
	CUDA_CALLABLE vec<T, N> & operator=(T rhs){
		for(int i=0; i<N; i++)
			v[i] = rhs;
		return *this;
	}

	template<typename Y>
	CUDA_CALLABLE vec<T, N> & operator=(vec<Y, N> rhs){
		for(int i=0; i<N; i++)
			v[i] = rhs.v[i];
		return *this;
	}
	//unary minus
	CUDA_CALLABLE vec<T, N> operator-() const{
		vec<T, N> r;

		//negate the vector
		for(int i=0; i<N; i++)
		    r.v[i] = -v[i];

		return r;
	}

	CUDA_CALLABLE bool operator==(vec<T, N> rhs) const
	{
        if ( (rhs.v[0] == v[0]) && (rhs.v[1] == v[1]) && (rhs.v[2] == v[2]) )
            return true;

        return false;
	}

	std::string str() const
	{
		std::stringstream ss;

		ss<<"[";
		for(int i=0; i<N; i++)
		{
			ss<<v[i];
			if(i != N-1)
				ss<<", ";
		}
		ss<<"]";

		return ss.str();
	}

	//bracket operator - allows assignment to the vector
	CUDA_CALLABLE T& operator[](const unsigned int i)
	{
        return v[i];
    }

};


}	//end namespace rts

template <typename T, int N>
std::ostream& operator<<(std::ostream& os, stim::vec<T, N> v)
{
    os<<v.str();
    return os;
}



template <typename T, int N>
CUDA_CALLABLE stim::vec<T, N> operator*(T lhs, stim::vec<T, N> rhs)
{
	stim::vec<T, N> r;

    return rhs * lhs;
}

//#if __GNUC__ > 3 && __GNUC_MINOR__ > 7
//template<class T, int N> using rtsVector = rts::vector<T, N>;
//#endif

#endif