codebase / stimlib

#ifndef STIM_MATRIX_H
#define STIM_MATRIX_H

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

namespace stim{

	enum mat4Format {
		mat4_float64,
		mat4_float32,
		mat4_int32,
		mat4_int16,
		mat4_uint16,
		mat4_uint8,
		mat4_float						//floating point type, determined automatically
	};

	static size_t mat4Format_size(mat4Format f){
		switch(f){
			case mat4_float64: return 8;
			case mat4_float32:
			case mat4_int32:   return 4;
			case mat4_int16:
			case mat4_uint16:  return 2;
			case mat4_uint8:   return 1;
			default:           return 0;
		}
	}

	//class encapsulates a mat4 file, and can be used to write multiple matrices to a single mat4 file
	class mat4file {
		std::ofstream matfile;

	public:
		/// Constructor opens a mat4 file for writing
		mat4file(std::string filename) {
			matfile.open(filename, std::ios::binary);
		}

		bool is_open() {
			return matfile.is_open();
		}

		void close() {
			matfile.close();
		}

		bool writemat(char* data, std::string varname, size_t sx, size_t sy, mat4Format format) {
			//save the matrix file here (use the mat4 function above)
			//data format: https://maxwell.ict.griffith.edu.au/spl/matlab-page/matfile_format.pdf (page 32)

			int MOPT = 0;									//initialize the MOPT type value to zero
			int m = 0;										//little endian
			int o = 0;										//reserved, always 0
			int p = format;
			int t = 0;
			MOPT = m * 1000 + o * 100 + p * 10 + t;			//calculate the type value
			int mrows = (int)sx;
			int ncols = (int)sy;
			int imagf = 0;									//assume real (for now)
			varname.push_back('\0');									//add a null to the string
			int namlen = (int)varname.size();						//calculate the name size

			size_t bytes = sx * sy * mat4Format_size(format);
			matfile.write((char*)&MOPT, 4);
			matfile.write((char*)&mrows, 4);
			matfile.write((char*)&ncols, 4);
			matfile.write((char*)&imagf, 4);
			matfile.write((char*)&namlen, 4);
			matfile.write((char*)&varname[0], namlen);
			matfile.write((char*)data, bytes);				//write the matrix data
			return is_open();
		}
	};

	static void save_mat4(char* data, std::string filename, std::string varname, size_t sx, size_t sy, mat4Format format){
		mat4file outfile(filename);									//create a mat4 file object
		if (outfile.is_open()) {									//if the file is open
			outfile.writemat(data, varname, sx, sy, format);		//write the matrix
			outfile.close();										//close the file
		}		
	}

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

	size_t bytes() {
		return R * C * sizeof(T);		//return the number of bytes of matrix data
	}
	/*void init(size_t rows, size_t cols){
		R = rows;
		C = cols;
		if (R == 0 || C == 0) M = NULL;
		else
			M = (T*)malloc(R * C * sizeof(T));	//allocate space for the matrix
	}*/

	T get(const size_t row, const size_t col) const {
		if (row >= R || col >= C) {
			std::cout << "ERROR: row or column out of range." << std::endl;
			exit(1);
		}
		return M[col * R + row];
	}

	T& at(size_t row, size_t col){
		if (row >= R || col >= C) {
			std::cout << "ERROR: row or column out of range." << std::endl;
			exit(1);
		}
		return M[col * R + row];
	}

public:
	matrix() {
		R = 0;
		C = 0;
		M = NULL;
	}

	matrix(size_t rows, size_t cols) {
		R = rows;
		C = cols;
		M = NULL;
		if (R * C > 0) 
			M = (T*) malloc(R * C * sizeof(T));
	}

	matrix(size_t rows, size_t cols, const T* data) {
		R = rows;
		C = cols;
		M = NULL;
		if (R * C > 0)
			M = (T*)malloc(R * C * sizeof(T));
		memcpy(M, data, R * C * sizeof(T));
	}

	matrix(const matrix<T>& cpy){
		M = NULL;
		if (cpy.R * cpy.C > 0)
			M = (T*)malloc(cpy.R * cpy.C * sizeof(T));
		memcpy(M, cpy.M, cpy.R * cpy.C * sizeof(T));

		R = cpy.R;
		C = cpy.C;
	}

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

	size_t rows() const {
		return R;
	}

	size_t cols() const {
		return C;
	}

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

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

		return *this;
	}

	matrix<T>& operator=(const matrix<T>& rhs){
		if (this != &rhs) {											//if the matrix isn't self-assigned
			T* new_matrix = new T[rhs.R * rhs.C];					//allocate new resources
			memcpy(new_matrix, rhs.M, rhs.R * rhs.C * sizeof(T));	//copy the matrix

			delete[] M;												//delete the previous array
			M = new_matrix;
			R = rhs.R;
			C = rhs.C;
		}
		return *this;
	}
	
	//element-wise operations
	matrix<T> operator+(const T rhs) const {
		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+(const matrix<T> rhs) const {
		if (R != rhs.R || C != rhs.C) {
			std::cout << "ERROR: addition is only defined for matrices that are the same size." << std::endl;
			exit(1);
		}
		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.M[i];			//calculate the operation and assign to result

		return result;
	}

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

	matrix<T> operator-(const matrix<T> rhs) const {
		return operator+(-rhs);
	}

	matrix<T> operator-() const {
		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];			//calculate the operation and assign to result

		return result;
	}

	matrix<T> operator*(const T rhs) const {
		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/(const T rhs) const {
		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*(const matrix<T> rhs) const {
		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 += get(r, i) * rhs.get(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() const {
		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;
	}

	// Reshapes the matrix in place
	void reshape(size_t rows, size_t cols) {
		R = rows;
		C = cols;
	}

	///Calculate and return the determinant of the matrix
	T det() const {
		if (R != C) {
			std::cout << "ERROR: a determinant can only be calculated for a square matrix." << std::endl;
			exit(1);
		}
		if (R == 1) return M[0];			//if the matrix only contains one value, return it

		int r, c, ri, cia, cib;
		T a = 0;
		T b = 0;
		for (c = 0; c < (int)C; c++) {
			for (r = 0; r < R; r++) {
				ri = r;
				cia = (r + c) % (int)C;
				cib = ((int)C - 1 - r) % (int)C;
				a += get(ri, cia);
				b += get(ri, cib);
			}
		}
		return a - b;
	}

	/// Sum all elements in the matrix
	T sum() const {
		size_t N = R * C;								//calculate the number of elements in the matrix
		T s = (T)0;										//allocate a register to store the sum
		for (size_t n = 0; n < N; n++) s += M[n];		//perform the summation
		return s;
	}

	/// Sort rows of the matrix by the specified indices
	matrix<T> sort_rows(size_t* idx) const {
		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, size_t data_type = mat4_float) const {
		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;
	}

	/// Return the column specified by index i
	matrix<T> col(size_t i) {
		matrix<T> c(R, 1);										//create a single column matrix
		memcpy(c.data(), &data()[R*i], C * sizeof(T));				//copy the column
		return c;
	}

	/// Return the row specified by index i
	matrix<T> row(size_t i) {
		matrix<T> r(1, C);										//create a single row matrix
		for (size_t c = 0; c < C; c++)
			r(0, c) = at(i, c);
		return r;
	}

	std::string toStr() const {
		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();
	}

	void csv(std::ostream& out) const {
		//std::stringstream csvss;
		for (size_t i = 0; i < R; i++) {
			out << std::fixed << M[i];
			for (size_t j = 1; j < C; j++)
				out << ", " << std::fixed << M[j * R + i];
			out << std::endl;
		}
		//return csvss.str();
	}

	std::string csv() const {
		std::stringstream csvss;
		int digits = std::numeric_limits<double>::max_digits10;
		csvss.precision(digits);
		csv(csvss);
		return csvss.str();
	}



	//save the data as a CSV file
	void csv(std::string filename) const {
		std::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;
	}

	//loads a matrix from a stream in CSV format
	void csv(std::istream& in) {
		size_t c, r;
		T v;
		for (r = 0; r < R; r++) {
			for (c = 0; c < C; c++) {
				in >> v;
				if (in.peek() == ',') in.seekg(1, std::ios::cur);
				at(r, c) = v;;
			}
		}
	}

	void raw(std::string filename) {
		std::ofstream out(filename, std::ios::binary);
		if (out) {
			out.write((char*)data(), rows() * cols() * sizeof(T));
			out.close();
		}
	}

	void mat4(stim::mat4file& file, std::string name = std::string("unknown"), mat4Format format = mat4_float) {
		//make sure the matrix name is valid (only numbers and letters, with a letter at the beginning
		for (size_t c = 0; c < name.size(); c++) {
			if (name[c] < 48 ||												//if the character isn't a number or letter, replace it with '_'
				(name[c] > 57 && name[c] < 65) ||
				(name[c] > 90 && name[c] < 97) ||
				(name[c] > 122)) {
				name[c] = '_';
			}
		}
		if (name[0] < 65 ||
			(name[0] > 91 && name[0] < 97) ||
			name[0] > 122) {
			name = std::string("m") + name;
		}
		if (format == mat4_float) {
			if (sizeof(T) == 4) format = mat4_float32;
			else if (sizeof(T) == 8) format = mat4_float64;
			else {
				std::cout << "stim::matrix ERROR - incorrect format specified" << std::endl;
				exit(1);
			}
		}
		//the name is now valid

		//if the size of the array is more than 100,000,000 elements, the matrix isn't supported
		if (rows() * cols() > 100000000) {											//break the matrix up into multiple parts
			//mat4file out(filename);													//create a mat4 object to write the matrix
			if (file.is_open()) {
				if (rows() < 100000000) {												//if the size of the row is less than 100,000,000, split the matrix up by columns
					size_t ncols = 100000000 / rows();									//calculate the number of columns that can fit in one matrix
					size_t nmat = (size_t)std::ceil((double)cols() / (double)ncols);			//calculate the number of matrices required
					for (size_t m = 0; m < nmat; m++) {									//for each matrix
						std::stringstream ss;
						ss << name << "_part_" << m + 1;
						if (m == nmat - 1)
							file.writemat((char*)(data() + m * ncols * rows()), ss.str(), rows(), cols() - m * ncols, format);
						else
							file.writemat((char*)(data() + m * ncols * rows()), ss.str(), rows(), ncols, format);
					}
				}
			}
		}
		//call the mat4 subroutine
		else
			//stim::save_mat4((char*)M, filename, name, rows(), cols(), format);
			file.writemat((char*)data(), name, rows(), cols(), format);
	}

	// saves the matrix as a Level-4 MATLAB file
	void mat4(std::string filename, std::string name = std::string("unknown"), mat4Format format = mat4_float) {
		stim::mat4file matfile(filename);

		if (matfile.is_open()) {
			mat4(matfile, name, format);
			matfile.close();
		}
	}
};

}	//end namespace rts


#endif