codebase / stimlib

#ifndef STIM_SCALARFIELD_H
#define STIM_SCALARFIELD_H


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

namespace stim{

	/// Perform a k-space transform of a scalar field (FFT). The given field has a width of x and the calculated momentum space has a
	///		width of kx (in radians).
	/// @param K is a pointer to the output array of all plane waves in the field
	/// @param kx is the width of the frame in momentum space
	/// @param ky is the height of the frame in momentum space
	/// @param E is the field to be transformed
	/// @param x is the width of the field in the spatial domain
	/// @param y is the height of the field in the spatial domain
	/// @param nx is the number of pixels representing the field in the x (and kx) direction
	/// @param ny is the number of pixels representing the field in the y (and ky) direction
	template<typename T>
	void cpu_scalar_to_kspace(stim::complex<T>* K, T& kx, T& ky, stim::complex<T>* E, T x, T y, size_t nx, size_t ny){

		kx = stim::TAU * nx / x;			//calculate the width of the momentum space
		ky = stim::TAU * ny / y;

		stim::complex<T>* dev_FFT;
		HANDLE_ERROR( cudaMalloc(&dev_FFT, sizeof(stim::complex<T>) * nx * ny) );		//allocate space on the CUDA device for the output array

		stim::complex<T>* dev_E;
		HANDLE_ERROR( cudaMalloc(&dev_E, sizeof(stim::complex<T>) * nx * ny) );		//allocate space for the field
		HANDLE_ERROR( cudaMemcpy(dev_E, E, sizeof(stim::complex<T>) * nx * ny, cudaMemcpyHostToDevice) );	//copy the field to GPU memory

		cufftResult result;
		cufftHandle plan;
		result = cufftPlan2d(&plan, nx, ny, CUFFT_C2C);
		if(result != CUFFT_SUCCESS){
			std::cout<<"Error creating cuFFT plan."<<std::endl;
			exit(1);
		}

		result = cufftExecC2C(plan, (cufftComplex*)dev_E, (cufftComplex*)dev_FFT, CUFFT_FORWARD);
		if(result != CUFFT_SUCCESS){
			std::cout<<"Error using cuFFT to perform a forward Fourier transform of the field."<<std::endl;
			exit(1);
		}

		stim::complex<T>* fft = (stim::complex<T>*) malloc(sizeof(stim::complex<T>) * nx * ny);
		HANDLE_ERROR( cudaMemcpy(fft, dev_FFT, sizeof(stim::complex<T>) * nx * ny, cudaMemcpyDeviceToHost) );

		stim::cpu_fftshift(K, fft, nx, ny);
		//memcpy(K, fft, sizeof(stim::complex<T>) * nx * ny);
	}

	template<typename T>
	void cpu_scalar_from_kspace(stim::complex<T>* E, T& x, T& y, stim::complex<T>* K, T kx, T ky, size_t nx, size_t ny){

		x = stim::TAU * nx / kx;			//calculate the width of the momentum space
		y = stim::TAU * ny / ky;
		
		stim::complex<T>* fft = (stim::complex<T>*) malloc(sizeof(stim::complex<T>) * nx * ny);
		stim::cpu_ifftshift(fft, K, nx, ny);
		//memcpy(fft, K, sizeof(stim::complex<T>) * nx * ny);

		stim::complex<T>* dev_FFT;
		HANDLE_ERROR( cudaMalloc(&dev_FFT, sizeof(stim::complex<T>) * nx * ny) );		//allocate space on the CUDA device for the output array
		HANDLE_ERROR( cudaMemcpy(dev_FFT, fft, sizeof(stim::complex<T>) * nx * ny, cudaMemcpyHostToDevice) );	//copy the field to GPU memory

		stim::complex<T>* dev_E;
		HANDLE_ERROR( cudaMalloc(&dev_E, sizeof(stim::complex<T>) * nx * ny) );		//allocate space for the field

		cufftResult result;
		cufftHandle plan;
		result = cufftPlan2d(&plan, nx, ny, CUFFT_C2C);
		if(result != CUFFT_SUCCESS){
			std::cout<<"Error creating cuFFT plan."<<std::endl;
			exit(1);
		}

		result = cufftExecC2C(plan, (cufftComplex*)dev_FFT, (cufftComplex*)dev_E, CUFFT_INVERSE);
		if(result != CUFFT_SUCCESS){
			std::cout<<"Error using cuFFT to perform a forward Fourier transform of the field."<<std::endl;
			exit(1);
		}

		HANDLE_ERROR( cudaMemcpy(E, dev_E, sizeof(stim::complex<T>) * nx * ny, cudaMemcpyDeviceToHost) );

		
	}

	/// Propagate a field slice along its orthogonal direction by a given distance z
	/// @param Enew is the resulting propogated field
	/// @param E is the field to be propogated
	/// @param sx is the size of the field in the lateral x direction
	/// @param sy is the size of the field in the lateral y direction
	/// @param z is the distance to be propagated
	/// @param k is the wavenumber 2*pi/lambda
	/// @param nx is the number of samples in the field along the lateral x direction
	/// @param ny is the number of samples in the field along the lateral y direction
	template<typename T>
	void cpu_scalar_propagate(stim::complex<T>* Enew, stim::complex<T>* E, T sx, T sy, T z, T k, size_t nx, size_t ny){
		
		stim::complex<T>* K = (stim::complex<T>*) malloc( sizeof(stim::complex<T>) * nx * ny );

		T Kx, Ky;											//width and height in k space
		cpu_scalar_to_kspace(K, Kx, Ky, E ,sx, sy, nx, ny);

		T* mag = (T*) malloc( sizeof(T) * nx * ny );
		stim::abs(mag, K, nx * ny);
		stim::cpu2image<float>(mag, "kspace_pre_shift.bmp", nx, ny, stim::cmBrewer);
		
		size_t kxi, kyi;
		size_t i;
		T kx, kx_sq, ky, ky_sq, k_sq;
		T kz;
		stim::complex<T> shift;
		T min_kx = -Kx / 2;
		T dkx = Kx / (nx);

		T min_ky = -Ky / 2;
		T dky = Ky / (ny);

		for(kyi = 0; kyi < ny; kyi++){						//for each plane wave in the ky direction
			for(kxi = 0; kxi < nx; kxi++){					//for each plane wave in the ky direction
				i = kyi * nx + kxi;

				kx = min_kx + kxi * dkx;					//calculate the position of the current plane wave
				ky = min_ky + kyi * dky;

				kx_sq = kx * kx;
				ky_sq = ky * ky;
				k_sq = k*k;
				
				if(kx_sq + ky_sq < k_sq){
					kz = sqrt(k_sq - kx_sq - ky_sq);			//estimate kz using the Fresnel approximation				
					shift = -exp(stim::complex<T>(0, kz * z));
					K[i] *= shift;
					K[i] /= (nx*ny);							//normalize the DFT
				}
				else{
					K[i] = 0;
				}
			}
		}
		
		stim::abs(mag, K, nx * ny);
		stim::cpu2image<float>(mag, "kspace_post_shift.bmp", nx, ny, stim::cmBrewer);
		
		cpu_scalar_from_kspace(Enew, sx, sy, K, Kx, Ky, nx, ny);
	}

	/// Apply a lowpass filter to a field slice
	/// @param Enew is the resulting propogated field
	/// @param E is the field to be propogated
	/// @param sx is the size of the field in the lateral x direction
	/// @param sy is the size of the field in the lateral y direction
	/// @param highest is the highest spatial frequency that can pass through the filter
	/// @param nx is the number of samples in the field along the lateral x direction
	/// @param ny is the number of samples in the field along the lateral y direction
	template<typename T>
	void cpu_scalar_lowpass(stim::complex<T>* Enew, stim::complex<T>* E, T sx, T sy, T highest, size_t nx, size_t ny){
		
		stim::complex<T>* K = (stim::complex<T>*) malloc( sizeof(stim::complex<T>) * nx * ny );

		T Kx, Ky;											//width and height in k space
		cpu_scalar_to_kspace(K, Kx, Ky, E ,sx, sy, nx, ny);

		T* mag = (T*) malloc( sizeof(T) * nx * ny );
		stim::abs(mag, K, nx * ny);
		stim::cpu2image<float>(mag, "kspace_pre_lowpass.bmp", nx, ny, stim::cmBrewer);
		
		size_t kxi, kyi;
		size_t i;
		T kx, kx_sq, ky, ky_sq, k_sq;
		T kz;
		stim::complex<T> shift;
		T min_kx = -Kx / 2;
		T dkx = Kx / (nx);

		T min_ky = -Ky / 2;
		T dky = Ky / (ny);

		T highest_sq = highest * highest;

		for(kyi = 0; kyi < ny; kyi++){						//for each plane wave in the ky direction
			for(kxi = 0; kxi < nx; kxi++){					//for each plane wave in the ky direction
				i = kyi * nx + kxi;

				kx = min_kx + kxi * dkx;					//calculate the position of the current plane wave
				ky = min_ky + kyi * dky;

				kx_sq = kx * kx;
				ky_sq = ky * ky;
				
				if(kx_sq + ky_sq > highest_sq){
					K[i] = 0;
				}
				else
					K[i] /= nx * ny;						//normalize the DFT
			}
		}
		
		stim::abs(mag, K, nx * ny);
		stim::cpu2image<float>(mag, "kspace_post_lowpass.bmp", nx, ny, stim::cmBrewer);
		
		cpu_scalar_from_kspace(Enew, sx, sy, K, Kx, Ky, nx, ny);
	}

	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>{

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

	/// Convert the field to a k-space representation (do an FFT)
	void to_kspace(T& kx, T& ky){
		cpu_scalar_to_kspace(E, kx, ky, E, X.len(), Y.len(), R[0], R[1]);
	}

	void from_kspace(){
		kx = stim::TAU * R[0] / X.len();			//calculate the width of the momentum space
		ky = stim::TAU * R[1] / Y.len();
		T x, y;
		cpu_scalar_from_kspace(E, x, y, E, kx, ky, R[0], R[1]);
	}

public:

	/// 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];
	}

	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(grid_bytes());		//allocate in CPU memory
		memset(E, 0, grid_bytes());
		loc = CPUmem;
	}

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

	/// 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;
		else{
			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;
		else{
			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
		}
	}

	/// Propagate the field along its orthogonal direction by a distance d
	void propagate(T d, T k){
		cpu_scalar_propagate(E, E, X.len(), Y.len(), d, k, R[0], R[1]);
	}

	/// Propagate the field along its orthogonal direction by a distance d
	void lowpass(T highest){
		cpu_scalar_lowpass(E, E, X.len(), Y.len(), highest, R[0], R[1]);
	}

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

		ss<<endl;
		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 = (T)1.0 / (R[0] - 1);					//calculate the spacing between points in the grid
			T dv = (T)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];
					i++;					
				}
			}
			//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);
		}
		else{
			std::cout<<"GPU allocation of a meshgrid isn't supported yet. You'll have to write kernels to do the calculation.";
			exit(1);
		}
	}

	//clear the field, setting all values to zero
	void clear(){
		if(loc == GPUmem)
			HANDLE_ERROR(cudaMemset(E, 0, grid_bytes()));
		else
			memset(E, 0, grid_bytes());
	}

	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());
			break;
		case complexReal:
			stim::real(image, E, size());
			break;
		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){
	os<<rhs.str();
	return os;
}


#endif