scalarfield.h
4.45 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
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<typename T>
class scalarfield : public rect<T>{
protected:
stim::complex<T>* 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<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;
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
}
}
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 = 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];
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<typename T>
std::ostream& operator<<(std::ostream& os, stim::scalarfield<T>& rhs){
os<<rhs.str();
return os;
}
#endif