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, typename F = float>
class grid{

protected:

	size_t R[D];							//elements in each dimension
	F S[D];									//spacing between element samples
	T* ptr;									//pointer to the data (on the GPU or CPU)

	

	///Initializes a grid by allocating the necessary memory and setting all values to zero
	void init(){		
		ptr = NULL;										//initialize the data pointer to NULL
		memset(R, 0, sizeof(size_t) * D);				//set the resolution to zero
		for(size_t d = 0; d < D; d++) S[d] = (F)1.0;	//initialize the spacing to unity
	}

	void alloc(){
		if(ptr != NULL) free(ptr);						//if memory has already been allocated, free it
		size_t N = samples();							//calculate the total number of values		
		ptr = (T*)calloc(sizeof(T), N);					//allocate memory to store the grid
	}

public:

	///Default constructor doesn't do anything
	grid(){
		init();
	}

	///Constructor used to specify the grid size as a vector
	
	/// @param _R is a vector describing the grid resolution
	grid( stim::vec<size_t> _R){
		for (size_t d = 0; d < D; d++)
			R[d] = _R[d];
		init();		
	}

	///Return the total number of values in the binary file
	size_t samples(){
		size_t s = 1;
		for(size_t d = 0; d < D; d++)
			s *= R[d];
		return s;
	}

	///Return the number of bytes in the binary file
	size_t bytes(){
		return samples() * sizeof(T);
	}

	void
	setDim(stim::vec<float> s){
		for(size_t d = 0; d < D; d++)
			S[d] = s[d];
	}

	///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( size_t X0, ...){
		R[0] = X0;									//set the grid size of the first dimension
		va_list ap;									//get a variable list
		va_start(ap, X0);							//start the variable list at the first element
		for(size_t d = 1; d<D; d++)					//for each additional element
			R[d] = va_arg(ap, size_t);				//read the value from the variable list as a size_t
		va_end(ap);
		init();										//initialize the grid
	}*/

	///Set the spacing between grid sample points
	/// @param X0... is a list of values describing the grid sample spacing
	/*void spacing(F X0, ...) {
		S[0] = X0;											//set the grid size of the first dimension
		va_list ap;											//get a variable list
		va_start(ap, X0);									//start the variable list at the first element
		for (size_t d = 1; d<D; d++)						//for each additional element
			S[d] = va_arg(ap, F);						//read the value from the variable list as a size_t
		va_end(ap);
	}*/

	/// Set the spacing between grid sample points for the specified dimension
	void spacing(size_t d, F sp){
		if(d < D) S[d] = sp;
		else{
			std::cout<<"error in stim::grid::spacing() - insufficient dimensions"<<std::endl;
			exit(1);
		}
	}

	/// Return the spacing for a given dimension
	F spacing(size_t d){
		if(d < D) return S[d];
		else{
			std::cout<<"error in stim::grid::spacing() - insufficient dimensions"<<std::endl;
			exit(1);
		}
	}

	/// Get the sample spacing for the given dimension
	F get_spacing(size_t d) {
		return S[d];
	}

	/// Get the size of the grid along the specified dimension
	F size(size_t d){
		return (F)R[d] * S[d];
	}

	/// Return the number of samples
	size_t samples(size_t d){
		return R[d];
	}

	///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;		
		file.open(filename.c_str(), std::ios::out | std::ios::binary);		//open the file as binary for reading		
		file.write((char *)ptr, samples() * sizeof(T));						//write file to disk
	}

	///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<size_t> X, unsigned long header = 0){
		for(size_t d = 0; d < D; d++)
			R[d] = X[d];																//set the sample resolution		
		init();																//allocate space for the data
		std::fstream file;		
		file.open(filename.c_str(), std::ios::in | std::ios::binary);		//open the file as binary for writing		
		file.seekg(header, std::ios::beg);									//seek past the header		
		file.read((char *)ptr, samples() * sizeof(T));						//read the data
	}

	///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;									//create a variable list
		
		unsigned long F = 1;						//initialize the dimension size to 1
		unsigned long idx = x0;

		va_start(ap, x0);							//start a variable list
		for(unsigned int d = 1; d<D; d++){			//for each dimension
			F *= R[d-1];							//get the size of the first dimension
			idx += va_arg(ap, unsigned int) * F;	//increment the index
		}
		va_end(ap);

		return ptr[idx];							//access the appropriate element and return the value
	}*/

	///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;									//create a variable list		
		unsigned long F = 1;						//initialize the dimension counter to 1
		unsigned long idx = x0;						//initialize the index to the first variable

		va_start(ap, x0);							//start the variable list
		for(unsigned int d = 1; d<D; d++){			//for each dimension
			F *= R[d - 1];
			idx += va_arg(ap, unsigned int) * F;	//update the index
		}
		va_end(ap);
		ptr[idx] = value;							//set the value at the indexed location
	}*/


	///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 = std::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