scalarfield.h 4.45 KB

#include "../math/rect.h"
#include "../math/complex.h"

namespace stim{

	enum locationType {CPUmem, GPUmem};

	/// Class represents a scalar optical field.

	/// In general, this class is designed to operate between the CPU and GPU. So, make sure all functions have an option to create the output on either.
	///		The field is stored *either* on the GPU or host memory, but not both. This enforces that there can't be different copies of the same field.
	///		This class is designed to be included in all of the other scalar optics classes, allowing them to render output data so make sure to keep it general and compatible.

template<typename T>
class scalarfield : public rect<T>{

	stim::complex<T>* E;
	size_t R[2];
	locationType loc;



	CUDA_CALLABLE scalarfield(size_t X, size_t Y, T size = 1, T z_pos = 0) : rect<T>::rect(size, z_pos){
		R[0] = X;											//set the field resolution
		R[1] = Y;

		E = (stim::complex<T>*) malloc(sizeof(stim::complex<T>) * R[0] * R[1]);		//allocate in CPU memory
		loc = CPUmem;

	CUDA_CALLABLE ~scalarfield(){
		if(loc == CPUmem) free(E);
		else cudaFree(E);

	/// Returns the number of values in the field
	CUDA_CALLABLE size_t size(){
		return R[0] * R[1];

	CUDA_CALLABLE size_t grid_bytes(){
		return sizeof(stim::complex<T>) * R[0] * R[1];

	/// Calculates the distance between points on the grid
	T spacing(){
		T du = rect<T>::X.len() / R[0];
		T dv = rect<T>::Y.len() / R[1];
		return min<T>(du, dv);

	/// Copy the field array to the GPU, if it isn't already there
	void to_gpu(){
		if(loc == GPUmem) return;
			stim::complex<T>* dev_E;
			HANDLE_ERROR( cudaMalloc(&dev_E, e_bytes()) );								//allocate GPU memory
			HANDLE_ERROR( cudaMemcpy(dev_E, E, e_bytes(), cudaMemcpyHostToDevice) );	//copy the field to the GPU
			free(E);																	//free the CPU memory
			E = dev_E;																	//swap pointers

	/// Copy the field array to the CPU, if it isn't already there
	void to_cpu(){
		if(loc == CPUmem) return;
			stim::complex<T>* host_E = (stim::complex<T>*) malloc(grid_bytes());			//allocate space in main memory
			HANDLE_ERROR( cudaMemcpy(host_E, E, grid_bytes(), cudaMemcpyDeviceToHost) );	//copy from GPU to CPU
			HANDLE_ERROR( cudaFree(E) );												//free device memory
			E = host_E;																	//swap pointers

	std::string str(){
		std::stringstream ss;
		ss<<"[ "<<R[0]<<" x "<<R[1]<<" ]"<<std::endl;
		ss<<"location: ";
		if(loc == CPUmem) ss<<"CPU";
		else ss<<"GPU";

		return ss.str();

	stim::complex<T>* ptr(){
		return E;

	/// Evaluate the cartesian coordinates of each point in the field. The resulting arrays are allocated in the same memory where the field is stored.
	void meshgrid(T* X, T* Y, T* Z, locationType location){
		size_t array_size = sizeof(T) * R[0] * R[1];
		if(location == CPUmem){

			T du = 1.0 / (R[0] - 1);					//calculate the spacing between points in the grid
			T dv = 1.0 / (R[1] - 1);

			size_t ui, vi, i;
			stim::vec3<T> p;
			for(vi = 0; vi < R[1]; vi++){
				i = vi * R[0];
				for(ui = 0; ui < R[0]; ui++){
					p = rect<T>::p(ui * du, vi * dv);
					X[i] = p[0];
					Y[i] = p[1];
					Z[i] = p[2];
			stim::cpu2image(X, "X.bmp", R[0], R[1], stim::cmBrewer);
			stim::cpu2image(Y, "Y.bmp", R[0], R[1], stim::cmBrewer);
			stim::cpu2image(Z, "Z.bmp", R[0], R[1], stim::cmBrewer);
			std::cout<<"GPU allocation of a meshgrid isn't supported yet. You'll have to write kernels to do the calculation.";

	void image(std::string filename, stim::complexComponentType type = complexMag, stim::colormapType cmap = stim::cmBrewer){

		if(loc == GPUmem) to_cpu();									//if the field is in the GPU, move it to the CPU
		T* image = (T*) malloc( sizeof(T) * size() );				//allocate space for the real image

		switch(type){												//get the specified component from the complex value
		case complexMag:
			stim::abs(image, E, size());
		case complexReal:
			stim::real(image, E, size());
		case complexImaginary:
			stim::imag(image, E, size());
		stim::cpu2image(image, filename, R[0], R[1], cmap);			//save the resulting image
		free(image);												//free the real image

};				//end class scalarfield

//stream insertion operator
template<typename T>
std::ostream& operator<<(std::ostream& os, stim::scalarfield<T>& rhs){
	return os;