stim/math/vec3.h

#ifndef STIM_VEC3_H
#define STIM_VEC3_H
#include <stim/cuda/cudatools/callable.h>
#include <cmath>
namespace stim{
/// A class designed to act as a 3D vector with CUDA compatibility
template<typename T>
class vec3{
protected:
	T ptr[3];
public:
	CUDA_CALLABLE vec3(){}
	CUDA_CALLABLE vec3(T v){
		ptr[0] = ptr[1] = ptr[2] = v;
	}
	CUDA_CALLABLE vec3(T x, T y, T z){
		ptr[0] = x;
		ptr[1] = y;
		ptr[2] = z;
	}
	//copy constructor
	CUDA_CALLABLE vec3( const vec3<T>& other){
		ptr[0] = other.ptr[0];
		ptr[1] = other.ptr[1];
		ptr[2] = other.ptr[2];
	}
	//access an element using an index
	CUDA_CALLABLE T& operator[](size_t idx){
		return ptr[idx];
	}
	CUDA_CALLABLE T* data(){
		return ptr;
	}
/// Casting operator. Creates a new vector with a new type U.
	template< typename U >
	CUDA_CALLABLE operator vec3<U>(){
		vec3<U> result;
		result.ptr[0] = (U)ptr[0];
		result.ptr[1] = (U)ptr[1];
		result.ptr[2] = (U)ptr[2];
		return result;
	}
	// computes the squared Euclidean length (useful for several operations where only >, =, or < matter)
	CUDA_CALLABLE T len_sq() const{
		return ptr[0] * ptr[0] + ptr[1] * ptr[1] + ptr[2] * ptr[2];
	}
	/// computes the Euclidean length of the vector
	CUDA_CALLABLE T len() const{
		return sqrt(len_sq());
	}
	
	/// Convert the vector from cartesian to spherical coordinates (x, y, z -> r, theta, phi where theta = [0, 2*pi])
	CUDA_CALLABLE vec3<T> cart2sph() const{
		vec3<T> sph;
		sph.ptr[0] = len();
		sph.ptr[1] = std::atan2(ptr[1], ptr[0]);
		if(sph.ptr[0] == 0)
			sph.ptr[2] = 0;
		else
			sph.ptr[2] = std::acos(ptr[2] / sph.ptr[0]);
		return sph;
	}
	/// Convert the vector from cartesian to spherical coordinates (r, theta, phi -> x, y, z where theta = [0, 2*pi])
	CUDA_CALLABLE vec3<T> sph2cart() const{
		vec3<T> cart;
		cart.ptr[0] = ptr[0] * std::cos(ptr[1]) * std::sin(ptr[2]);
		cart.ptr[1] = ptr[0] * std::sin(ptr[1]) * std::sin(ptr[2]);
		cart.ptr[2] = ptr[0] * std::cos(ptr[2]);
		return cart;
	}
	/// Computes the normalized vector (where each coordinate is divided by the L2 norm)
	CUDA_CALLABLE vec3<T> norm() const{
        vec3<T> result;
        T l = len();						//compute the vector length
        return (*this) / l;
	}
	/// Computes the cross product of a 3-dimensional vector
	CUDA_CALLABLE vec3<T> cross(const vec3<T> rhs) const{
		vec3<T> result;
		result[0] = (ptr[1] * rhs.ptr[2] - ptr[2] * rhs.ptr[1]);
		result[1] = (ptr[2] * rhs.ptr[0] - ptr[0] * rhs.ptr[2]);
		result[2] = (ptr[0] * rhs.ptr[1] - ptr[1] * rhs.ptr[0]);
		return result;
	}
	/// Compute the Euclidean inner (dot) product
    CUDA_CALLABLE T dot(vec3<T> rhs) const{
        return ptr[0] * rhs.ptr[0] + ptr[1] * rhs.ptr[1] + ptr[2] * rhs.ptr[2];
    }
	/// Arithmetic addition operator
    /// @param rhs is the right-hand-side operator for the addition
	CUDA_CALLABLE vec3<T> operator+(vec3<T> rhs) const{
		vec3<T> result;
		result.ptr[0] = ptr[0] + rhs[0];
		result.ptr[1] = ptr[1] + rhs[1];
		result.ptr[2] = ptr[2] + rhs[2];
		return result;
	}
	/// Arithmetic addition to a scalar
	/// @param rhs is the right-hand-side operator for the addition
	CUDA_CALLABLE vec3<T> operator+(T rhs) const{
		vec3<T> result;
		result.ptr[0] = ptr[0] + rhs;
		result.ptr[1] = ptr[1] + rhs;
		result.ptr[2] = ptr[2] + rhs;
		return result;
	}
	/// Arithmetic subtraction operator
	/// @param rhs is the right-hand-side operator for the subtraction
	CUDA_CALLABLE vec3<T> operator-(vec3<T> rhs) const{
		vec3<T> result;
		result.ptr[0] = ptr[0] - rhs[0];
		result.ptr[1] = ptr[1] - rhs[1];
		result.ptr[2] = ptr[2] - rhs[2];
		return result;
	}
	/// Arithmetic subtraction to a scalar
	/// @param rhs is the right-hand-side operator for the addition
	CUDA_CALLABLE vec3<T> operator-(T rhs) const{
		vec3<T> result;
		result.ptr[0] = ptr[0] - rhs;
		result.ptr[1] = ptr[1] - rhs;
		result.ptr[2] = ptr[2] - rhs;
		return result;
	}
	/// Arithmetic scalar multiplication operator
	/// @param rhs is the right-hand-side operator for the subtraction
	CUDA_CALLABLE vec3<T> operator*(T rhs) const{
		vec3<T> result;
		result.ptr[0] = ptr[0] * rhs;
		result.ptr[1] = ptr[1] * rhs;
		result.ptr[2] = ptr[2] * rhs;
		return result;
	}
	/// Arithmetic scalar division operator
	/// @param rhs is the right-hand-side operator for the subtraction
	CUDA_CALLABLE vec3<T> operator/(T rhs) const{
		return (*this) * ((T)1.0/rhs);
	}
	/// Multiplication by a scalar, followed by assignment
	CUDA_CALLABLE vec3<T> operator*=(T rhs){
		ptr[0] = ptr[0] * rhs;
		ptr[1] = ptr[1] * rhs;
		ptr[2] = ptr[2] * rhs;
		return *this;
	}
	/// Addition and assignment
	CUDA_CALLABLE vec3<T> operator+=(vec3<T> rhs){
		ptr[0] = ptr[0] + rhs;
		ptr[1] = ptr[1] + rhs;
		ptr[2] = ptr[2] + rhs;
		return *this;
	}
	/// Assign a scalar to all values
	CUDA_CALLABLE vec3<T> & operator=(T rhs){
		ptr[0] = ptr[0] = rhs;
		ptr[1] = ptr[1] = rhs;
		ptr[2] = ptr[2] = rhs;
		return *this;
	}
	/// Casting and assignment
	template<typename Y>
	CUDA_CALLABLE vec3<T> & operator=(vec3<Y> rhs){
		ptr[0] = (T)rhs.ptr[0];
		ptr[1] = (T)rhs.ptr[1];
		ptr[2] = (T)rhs.ptr[2];
		return *this;
	}
	/// Unary minus (returns the negative of the vector)
	CUDA_CALLABLE vec3<T> operator-() const{
		vec3<T> result;
		result.ptr[0] = -ptr[0];
		result.ptr[1] = -ptr[1];
		result.ptr[2] = -ptr[2];
		return result;
	}
//#ifndef __NVCC__
	/// Outputs the vector as a string
	std::string str() const{
		std::stringstream ss;
		const size_t N = 3;
		ss<<"[";
		for(size_t i=0; i<N; i++)
		{
			ss<<ptr[i];
			if(i != N-1)
				ss<<", ";
		}
		ss<<"]";
		return ss.str();
	}
//#endif
	size_t size(){ return 3; }
	};						//end class vec3
}							//end namespace stim
/// Multiply a vector by a constant when the vector is on the right hand side
template <typename T>
stim::vec3<T> operator*(T lhs, stim::vec3<T> rhs){
    return rhs * lhs;
}
//stream operator
template<typename T>
std::ostream& operator<<(std::ostream& os, stim::vec3<T> const& rhs){
	os<<rhs.str();
	return os;
}
#endif