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, data, 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()(size_t row, size_t col) {
		return at(row, col);
	}

	matrix<T> operator=(T rhs) {
		if (&rhs == this)
			return *this;
		init(R, C);
		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
		size_t c, r, i;
		for(c = 0; c < rhs.C; c++){
			for(r = 0; r < R; r++){
				inner = (T)0;
				for(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;
	}

	/// Sort rows of the matrix by the specified indices
	matrix<T> sort_rows(size_t* idx) {
		matrix<T> result(C, R);					//create the output matrix
		size_t r, c;
		for (c = 0; c < C; c++) {								//for each column
			for (r = 0; r < R; r++) {							//for each row element
				result.M[c * R + r] = M[c * R + idx[r]];		//copy each element of the row into its new position
			}
		}
		return result;
	}

	/// Sort columns of the matrix by the specified indices
	matrix<T> sort_cols(size_t* idx) {
		matrix<T> result(C, R);
		size_t c;
		for (c = 0; c < C; c++) {											//for each column
			memcpy(&result.M[c * R], &M[idx[c] * R], sizeof(T) * R);		//copy the entire column from this matrix to the appropriate location
		}
		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();
	}

	std::string csv() {
		std::stringstream csvss;
		for (size_t i = 0; i < R; i++) {
			csvss << M[i];
			for (size_t j = 0; j < C; j++) csvss << ", " << M[j * R + i];
			csvss << std::endl;
		}
		return csvss.str();
	}

	//save the data as a CSV file
	void csv(std::string filename) {
		ofstream basisfile(filename.c_str());
		basisfile << csv();
		basisfile.close();
	}

	static matrix<T> I(size_t N) {

		matrix<T> result(N, N);							//create the identity matrix
		memset(result.M, 0, N * N * sizeof(T));			//set the entire matrix to zero
		for (size_t n = 0; n < N; n++) {
			result(n, n) = (T)1;						//set the diagonal component to 1
		}
		return result;
	}

};

}	//end namespace rts


#endif