codebase / stimlib

#ifndef STIM_GRID_H
#define STIM_GRID_H

#include <vector>
#include <string>
#include <sstream>
#include <fstream>
#include <cstdarg> 

#include <stim/math/vector.h>

namespace stim{

/**This object describes a generic D-dimensional grid containing data of type T.
	Functions are provided for saving and loading binary data.

**/
template<typename T, unsigned int D = 1>
class grid{

protected:

	stim::vec<unsigned long> R;		//elements in each dimension
	stim::vec<float> S;
	T* ptr;									//pointer to the data (on the GPU or CPU)

	///Return the total number of values in the binary file
	unsigned long samples(){

		unsigned long s = 1;
		for(unsigned int d = 0; d < D; d++)
			s *= R[d];

		return s;
	}

	///Initializes a grid by allocating the necessary memory and setting all values to zero
	void init(){

		//calculate the total number of values
		unsigned long S = samples();

		//allocate memory to store the grid
		ptr = (T*)malloc(sizeof(T) * S);

		//initialize the memory to zero
		memset(ptr, 0, sizeof(T) * S);

	}

public:

	///Default constructor doesn't do anything
	grid(){
		ptr = NULL;			//set the pointer to NULL so that we know nothing is allocated
	}

	///Constructor used to specify the grid size as a vector
	
	/// @param _R is a vector describing the grid resolution
	grid( stim::vec<unsigned long> _R){
		
		//set the grid resolution
		R = _R;

		init();		
	}

	void
	setDim(stim::vec<float> s)
	{
		S = s;
	}

	///Constructor used to specify the grid size as a set of parameters

	/// @param X0... is a list of values describing the grid size along each dimension
	grid( unsigned long X0, ...){

		R[0] = X0;

		va_list ap;
		va_start(ap, X0);
		for(unsigned int d = 1; d<D; d++)
			R[d] = va_arg(ap, unsigned long);
		va_end(ap);

		init();

	}

	///Writes the binary data to disk

	/// @param filename is the name of the binary file to be written
	void write(std::string filename){

		std::fstream file;

		//open the file as binary for reading
		file.open(filename.c_str(), std::ios::out | std::ios::binary);

		//write file to disk
		file.write((char *)ptr, samples() * sizeof(T));
	}

	///Loads a binary file from disk

	/// @param filename is the name of the file containing the binary data
	/// @param S is the size of the binary file along each dimension
	/// @param header is the size of the header in bytes
	void read(std::string filename, stim::vec<unsigned long> S, unsigned long header = 0){

		R = S;	//set the sample resolution

		//allocate space for the data
		init();

		std::fstream file;

		//open the file as binary for writing
		file.open(filename.c_str(), std::ios::in | std::ios::binary);

		//seek past the header
		file.seekg(header, std::ios::beg);


		//read the data
		file.read((char *)ptr, samples() * sizeof(T));
	}

	///Gets a single value from the grid given a set of coordinates

	/// @param x0... is a list of coordinates specifying the desired value
	T get(unsigned long x0, ...){

		va_list ap;
		
		unsigned long F = 1;
		unsigned long p = x0;

		va_start(ap, x0);
		for(unsigned int d = 1; d<D; d++){
			F *= R[d-1];
			p += va_arg(ap, unsigned int) * F;
		}
		va_end(ap);

		return ptr[p];
	}

	///Sets a value in the grid

	/// @param value is the grid point value
	/// @x0... is the coordinate of the value to be set
	void set(T value, unsigned long x0, ...){

		va_list ap;
		
		unsigned long F = 1;
		unsigned long p = x0;

		va_start(ap, x0);
		for(unsigned int d = 1; d<D; d++){
			F *= R[d-1];
			p += va_arg(ap, unsigned int) * F;
		}
		va_end(ap);

		ptr[p] = value;
	}


	///Outputs grid data as a string
	std::string str(){

		std::stringstream result;

		result<<"stim::grid structure of size [";
		for(unsigned int d = 0; d<D; d++){
			if(d!=0) result<<", ";
			result<<R[d];

		}

		result<<"]"<<std::endl;

		//calculate the number of values to output
		unsigned long nV = min((unsigned long long)R[0], (unsigned long long)10);

		for(unsigned long v = 0; v<nV; v++){
			result<<ptr[v];
			if(v != nV-1) result<<", ";
		}

		return result.str();
	}
};

}


#endif