codebase / stimlib

#ifndef RTS_MATRIX_H
#define RTS_MATRIX_H

//#include "rts/vector.h"
#include <string.h>
#include <iostream>
#include <stim/math/vector.h>
#include <stim/math/vec3.h>
#include <stim/cuda/cudatools/callable.h>

namespace stim{

template <class T>
class matrix {
	//the matrix will be stored in column-major order (compatible with OpenGL)
	T* M;								//pointer to the matrix data
	size_t R;							//number of rows
	size_t C;							//number of colums

	void init(size_t rows, size_t cols){
		R = rows;
		C = cols;
		M = (T*) malloc (R * C * sizeof(T));	//allocate space for the matrix
	}

	T& at(size_t row, size_t col){
		return M[col * R + row];
	}

public:
	matrix(size_t rows, size_t cols) {
		init(rows, cols);									//initialize memory
	}

	matrix(size_t rows, size_t cols, T* data) {
		init(rows, cols);
		memcpy(M, rhs, R * C * sizeof(T));
	}

	matrix(const matrix<T>& cpy){
		init(cpy.R, cpy.C);
		memcpy(M, cpy.M, R * C * sizeof(T));
	}

	~matrix() {
		R = C = 0;
		if(M) free(M);
	}

	size_t rows(){
		return R;
	}

	size_t cols(){
		return C;
	}

	T& operator()(int row, int col) {
		return at(row, col);
	}

	matrix<T> operator=(T rhs) {
		size_t N = R * C;
		for(size_t n=0; n<N; n++)
			M[n] = rhs;

		return *this;
	}

	matrix<T>& operator=(matrix<T> rhs){
		init(rhs.R, rhs.C);
		memcpy(M, rhs.M, R * C * sizeof(T));
		return *this;
	}
	
	//element-wise operations
	matrix<T> operator+(T rhs) {
		matrix<T> result(R, C);					//create a result matrix
		size_t N = R * C;

		for(int i=0; i<N; i++)
			result.M[i] = M[i] + rhs;			//calculate the operation and assign to result

		return result;
	}

	matrix<T> operator-(T rhs) {
		return operator+(-rhs);					//add the negative of rhs
	}

	matrix<T> operator*(T rhs) {
		matrix<T> result(R, C);					//create a result matrix
		size_t N = R * C;

		for(int i=0; i<N; i++)
			result.M[i] = M[i] * rhs;			//calculate the operation and assign to result

		return result;
	}

	matrix<T> operator/(T rhs) {
		matrix<T> result(R, C);					//create a result matrix
		size_t N = R * C;

		for(int i=0; i<N; i++)
			result.M[i] = M[i] / rhs;			//calculate the operation and assign to result

		return result;
	}

	//matrix multiplication
	matrix<T> operator*(matrix<T> rhs){
		if(C != rhs.R){
			std::cout<<"ERROR: matrix multiplication is undefined for matrices of size ";
			std::cout<<"[ "<<R<<" x "<<C<<" ] and [ "<<rhs.R<<" x "<<rhs.C<<"]"<<std::endl;
			exit(1);
		}

		matrix<T> result(R, rhs.C);				//create the output matrix
		T inner;								//stores the running inner product
		for(size_t c = 0; c < rhs.C; c++){
			for(size_t r = 0; r < R; r++){
				inner = (T)0;
				for(size_t i = 0; i < C; i++){
					inner += at(r, i) * rhs.at(i, c);
				}
				result.M[c * R + r] = inner;
			}
		}
		return result;
	}

	//returns a pointer to the raw matrix data (in column major format)
	T* data(){
		return M;
	}

	//return a transposed matrix
	matrix<T> transpose(){
		matrix<T> result(C, R);
		size_t c, r;
		for(c = 0; c < C; c++){
			for(r = 0; r < R; r++){
				result.M[r * C + c] = M[c * R + r];
			}
		}
		return result;
	}

	std::string toStr()
	{
		std::stringstream ss;

		for(int r = 0; r < R; r++)
		{
			ss << "| ";
			for(int c=0; c<C; c++)
			{
				ss << M[c * R + r] << " ";
			}
			ss << "|" << std::endl;
		}

		return ss.str();
	}

};

}	//end namespace rts


#endif