codebase / stimlib

  #ifndef	RTS_COMPLEXFIELD_H
  #define RTS_COMPLEXFIELD_H
  
  #include "cublas_v2.h"
  #include <cuda_runtime.h>
  
  #include "../math/field.cuh"
  #include "../math/complex.h"

  #include "../math/realfield.cuh"

  
  namespace rts{
  

  template<typename T>
  __global__ void gpu_complexfield_mag(T* dest, complex<T>* source, unsigned int r0, unsigned int r1){
  
  	int iu = blockIdx.x * blockDim.x + threadIdx.x;
  	int iv = blockIdx.y * blockDim.y + threadIdx.y;
  
  	//make sure that the thread indices are in-bounds
  	if(iu >= r0 || iv >= r1) return;
  
  	//compute the index into the field
  	int i = iv*r0 + iu;
  
  	//calculate and store the result
  	dest[i] = source[i].abs();
  }
  

  /*This class stores functions for saving images of complex fields
  */

  template<typename T, unsigned int D = 1>

  class complexfield : public field< rts::complex<T>, D >{
  	using field< rts::complex<T>, D >::R;
  	using field< rts::complex<T>, D >::X;
  	using field< rts::complex<T>, D >::shape;

  	using field< rts::complex<T>, D >::cuda_params;
  
  	

  
  public:
  
  	//find the maximum value of component n
  	rts::complex<T> find_max(unsigned int n){
  		cublasStatus_t stat;
  		cublasHandle_t handle;
  
  		//create a CUBLAS handle
  		stat = cublasCreate(&handle);
  		if(stat != CUBLAS_STATUS_SUCCESS){
  			std::cout<<"CUBLAS Error: initialization failed"<<std::endl;
  			exit(1);
  		}
  
  		int L = R[0] * R[1];    //compute the number of discrete points in a slice
  		int index;				//result of the max operation
  		rts::complex<T> result;
  

  		if(sizeof(T) == 8)

  			stat = cublasIcamax(handle, L, (const cuComplex*)X[n], 1, &index);
  		else
  			stat = cublasIzamax(handle, L, (const cuDoubleComplex*)X[n], 1, &index);
  
  		index -= 1;        //adjust for 1-based indexing
  
  		//if there was a GPU error, terminate
  		if(stat != CUBLAS_STATUS_SUCCESS){
  			std::cout<<"CUBLAS Error: failure finding maximum value."<<std::endl;
  			exit(1);
  		}
  
  		//retrieve the maximum value for this slice and store it in the maxVal array
  		std::cout<<X[n]<<std::endl;
  		HANDLE_ERROR(cudaMemcpy(&result, X[n] + index, sizeof(rts::complex<T>), cudaMemcpyDeviceToHost));
  		return result;
  	}
  
  public:
  

  	enum attribute {magnitude, real, imaginary};
  

  	//constructor (no parameters)
  	complexfield() : field<rts::complex<T>, D>(){};
  
  	//constructor (resolution specified)
  	complexfield(unsigned int r0, unsigned int r1) : field<rts::complex<T>, D>(r0, r1){};
  

  	//assignment from a field of complex values
  	complexfield & operator=(const field< rts::complex<T>, D > rhs){
  		field< complex<T>, D >::operator=(rhs);
  		return *this;
  	}
  

  	//assignment operator (scalar value)
  	complexfield & operator= (const complex<T> rhs){
  
  		field< complex<T>, D >::operator=(rhs);
  		return *this;
  	}
  
  	//assignment operator (vector value)
  	complexfield & operator= (const vec< complex<T>, D > rhs){
  
  		field< complex<T>, D >::operator=(rhs);
  		return *this;
  	}
  

  	//cropping
  	complexfield crop(unsigned int width, unsigned int height){
  
  		complexfield<T, D> result;
  		result = field< complex<T>, D>::crop(width, height);
  		return result;
  	}
  
  	void toImage(std::string filename, attribute type = magnitude, unsigned int n=0){
  

  		field<T, 1> rf(R[0], R[1]);

  
  		//get cuda parameters
  		dim3 blocks, grids;
  		cuda_params(grids, blocks);
  
  		if(type == magnitude){
  			gpu_complexfield_mag <<<grids, blocks>>> (rf.ptr(), X[n], R[0], R[1]);

  			rf.toImage(filename, n, true);

  		}

  
  	}
  
  
  };
  
  
  }	//end namespace rts
  
  

  #endif