From 3a74ec6d78ca82b863e8027db6c071fc79ebd79a Mon Sep 17 00:00:00 2001 From: David Date: Wed, 15 Jun 2016 12:51:57 -0500 Subject: [PATCH] added a new scalarfield class for rendering and storing samples of EM fields --- stim/math/fft.h | 31 +++++++++++++++++++++++++++++++ stim/optics/lens.h | 148 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ stim/optics/scalarfield.h | 157 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 336 insertions(+), 0 deletions(-) create mode 100644 stim/math/fft.h create mode 100644 stim/optics/lens.h create mode 100644 stim/optics/scalarfield.h diff --git a/stim/math/fft.h b/stim/math/fft.h new file mode 100644 index 0000000..272ace7 --- /dev/null +++ b/stim/math/fft.h @@ -0,0 +1,31 @@ +#ifndef STIM_FFT_H +#define STIM_FFT_H + +namespace stim{ + + template + void circshift(T *out, const T *in, size_t xdim, size_t ydim, size_t xshift, size_t yshift){ + size_t i, j, ii, jj; + for (i =0; i < xdim; i++) { + ii = (i + xshift) % xdim; + for (j = 0; j < ydim; j++) { + jj = (j + yshift) % ydim; + out[ii * ydim + jj] = in[i * ydim + j]; + } + } + } + + template + void cpu_fftshift(T* out, T* in, size_t xdim, size_t ydim){ + circshift(out, in, xdim, ydim, xdim/2, ydim/2); + } + + template + void cpu_ifftshift(T* out, T* in, size_t xdim, size_t ydim){ + circshift(out, in, xdim, ydim, xdim/2, ydim/2); + } + + +} + +#endif \ No newline at end of file diff --git a/stim/optics/lens.h b/stim/optics/lens.h new file mode 100644 index 0000000..0e0dfcf --- /dev/null +++ b/stim/optics/lens.h @@ -0,0 +1,148 @@ +#ifndef STIM_LENS_H +#define STIM_LENS_H + +#include "scalarwave.h" +#include "../math/bessel.h" +#include "../cuda/cudatools/devices.h" +#include "../visualization/colormap.h" +#include "../math/fft.h" + +#include "cufft.h" + +#include + +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 + void cpu_scalar_to_kspace(stim::complex* K, T& kx, T& ky, stim::complex* 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* dev_FFT; + HANDLE_ERROR( cudaMalloc(&dev_FFT, sizeof(stim::complex) * nx * ny) ); //allocate space on the CUDA device for the output array + + stim::complex* dev_E; + HANDLE_ERROR( cudaMalloc(&dev_E, sizeof(stim::complex) * nx * ny) ); //allocate space for the field + HANDLE_ERROR( cudaMemcpy(dev_E, E, sizeof(stim::complex) * 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."<* fft = (stim::complex*) malloc(sizeof(stim::complex) * nx * ny); + HANDLE_ERROR( cudaMemcpy(fft, dev_FFT, sizeof(stim::complex) * nx * ny, cudaMemcpyDeviceToHost) ); + + stim::cpu_fftshift(K, fft, nx, ny); + } + + template + void cpu_scalar_from_kspace(stim::complex* E, T& x, T& y, stim::complex* 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* fft = (stim::complex*) malloc(sizeof(stim::complex) * nx * ny); + stim::cpu_ifftshift(fft, K, nx, ny); + + stim::complex* dev_FFT; + HANDLE_ERROR( cudaMalloc(&dev_FFT, sizeof(stim::complex) * nx * ny) ); //allocate space on the CUDA device for the output array + HANDLE_ERROR( cudaMemcpy(dev_FFT, fft, sizeof(stim::complex) * nx * ny, cudaMemcpyHostToDevice) ); //copy the field to GPU memory + + stim::complex* dev_E; + HANDLE_ERROR( cudaMalloc(&dev_E, sizeof(stim::complex) * 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."<) * nx * ny, cudaMemcpyDeviceToHost) ); + + + } + + /// Propagate a field slice along its orthogonal direction by a given distance z + /// @param Enew is the resulting propogated field + template + void cpu_scalar_propagate(stim::complex* Enew, stim::complex* E, T sx, T sy, T z, T k, size_t nx, size_t ny){ + + stim::complex* K = (stim::complex*) malloc( sizeof(stim::complex) * 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(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 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*k - kx * kx - ky * ky); //estimate kz using the Fresnel approximation + shift = -exp(stim::complex(0, kz * z)); + K[i] *= shift; + } + else{ + K[i] = 0; + } + } + } + + stim::abs(mag, K, nx * ny); + stim::cpu2image(mag, "kspace_post_shift.bmp", nx, ny, stim::cmBrewer); + + cpu_scalar_from_kspace(Enew, sx, sy, K, Kx, Ky, nx, ny); + } + +} + + +#endif \ No newline at end of file diff --git a/stim/optics/scalarfield.h b/stim/optics/scalarfield.h new file mode 100644 index 0000000..1c73002 --- /dev/null +++ b/stim/optics/scalarfield.h @@ -0,0 +1,157 @@ +#ifndef STIM_SCALARFIELD_H +#define STIM_SCALARFIELD_H + +#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 +class scalarfield : public rect{ + +protected: + stim::complex* E; + size_t R[2]; + locationType loc; + + + +public: + + CUDA_CALLABLE scalarfield(size_t X, size_t Y, T size = 1, T z_pos = 0) : rect::rect(size, z_pos){ + R[0] = X; //set the field resolution + R[1] = Y; + + E = (stim::complex*) malloc(sizeof(stim::complex) * 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) * R[0] * R[1]; + } + + /// Calculates the distance between points on the grid + T spacing(){ + T du = rect::X.len() / R[0]; + T dv = rect::Y.len() / R[1]; + return min(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* 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* host_E = (stim::complex*) malloc(e_bytes()); //allocate space in main memory + HANDLE_ERROR( cudaMemcpy(host_E, E, e_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<::str()<* 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 p; + for(vi = 0; vi < R[1]; vi++){ + i = vi * R[0]; + for(ui = 0; ui < R[0]; ui++){ + p = rect::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); + } + } + + 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 +std::ostream& operator<<(std::ostream& os, stim::scalarfield& rhs){ + os<