codebase / stimlib

#ifndef RTS_BEAM
#define RTS_BEAM

#include "../math/vec3.h"
#include "../optics/scalarwave.h"
#include "../math/bessel.h"
#include "../math/legendre.h"
#include <vector>

namespace stim{

		/// Function returns the value of the scalar field produced by a beam with the specified parameters

		template<typename T>
		std::vector< stim::vec3<T> > generate_focusing_vectors(size_t N, stim::vec3<T> d, T NA, T NA_in = 0){

			std::vector< stim::vec3<T> > dirs(N);					//allocate an array to store the focusing vectors

			///compute the rotation operator to transform (0, 0, 1) to k
			T cos_angle = d.dot(vec3<T>(0, 0, 1));
			stim::matrix<T, 3> rotation;

			//if the cosine of the angle is -1, the rotation is just a flip across the z axis
			if(cos_angle == -1){
				rotation(2, 2) = -1;
			}
			else if(cos_angle != 1.0)
			{
				vec3<T> r_axis = vec3<T>(0, 0, 1).cross(d).norm();	//compute the axis of rotation
				T angle = acos(cos_angle);							//compute the angle of rotation
				quaternion<T> quat;							//create a quaternion describing the rotation
				quat.CreateRotation(angle, r_axis);
				rotation = quat.toMatrix3();							//compute the rotation matrix
			}

			//find the phi values associated with the cassegrain ring
			T PHI[2];
			PHI[0] = (T)asin(NA);
			PHI[1] = (T)asin(NA_in);

			//calculate the z-axis cylinder coordinates associated with these angles
			T Z[2];
			Z[0] = cos(PHI[0]);
			Z[1] = cos(PHI[1]);
			T range = Z[0] - Z[1];

			//draw a distribution of random phi, z values
			T z, phi, theta;
			//T kmag = stim::TAU / lambda;
			for(int i=0; i<N; i++){								//for each sample
				z = (T)((double)rand() / (double)RAND_MAX) * range + Z[1];			//find a random position on the surface of a cylinder
				theta = (T)(((double)rand() / (double)RAND_MAX) * stim::TAU);
				phi = acos(z);													//project onto the sphere, computing phi in spherical coordinates

				//compute and store cartesian coordinates
				vec3<T> spherical(1, theta, phi);								//convert from spherical to cartesian coordinates
				vec3<T> cart = spherical.sph2cart();
				dirs[i] = rotation * cart;										//create a sample vector
			}
			return dirs;
		}
		
/// Class stim::beam represents a beam of light focused at a point and composed of several plane waves
template<typename T>
class scalarbeam
{
public:
	//enum beam_type {Uniform, Bartlett, Hamming, Hanning};

private:
	
	T NA[2];				//numerical aperature of the focusing optics	
	vec3<T> f;				//focal point
	vec3<T> d;				//propagation direction
	stim::complex<T> A;		//beam amplitude
	T lambda;				//beam wavelength
public:

	///constructor: build a default beam (NA=1.0)
	scalarbeam(T wavelength = 1, stim::complex<T> amplitude = 1, vec3<T> focal_point = vec3<T>(0, 0, 0), vec3<T> direction = vec3<T>(0, 0, 1), T numerical_aperture = 1, T center_obsc = 0){
		lambda = wavelength;
		A = amplitude;
		f = focal_point;
		d = direction.norm();					//make sure that the direction vector is normalized (makes calculations more efficient later on)
		NA[0] = numerical_aperture;
		NA[1] = center_obsc;
	}

	///Numerical Aperature functions
	void setNA(T na)
	{
		NA[0] = (T)0;
		NA[1] = na;
	}
	void setNA(T na0, T na1)
	{
		NA[0] = na0;
		NA[1] = na1;
	}

	//Monte-Carlo decomposition into plane waves
	std::vector< scalarwave<T> > mc(size_t N = 100000) const{

		std::vector< stim::vec3<T> > dirs = generate_focusing_vectors(N, d, NA[0], NA[1]);	//generate a random set of N vectors forming a focus
		std::vector< scalarwave<T> > samples(N);											//create a vector of plane waves
		T kmag = (T)stim::TAU / lambda;								//calculate the wavenumber
		stim::complex<T> apw;										//allocate space for the amplitude at the focal point
		stim::vec3<T> kpw;											//declare the new k-vector based on the focused plane wave direction
		for(size_t i=0; i<N; i++){										//for each sample
			kpw = dirs[i] * kmag;									//calculate the k-vector for the new plane wave
			apw = exp(stim::complex<T>(0, kpw.dot(-f)));				//calculate the amplitude for the new plane wave
			samples[i] = scalarwave<T>(kpw, apw);			//create a plane wave based on the direction
		}

		return samples;
	}

	/// Calculate the field at a given point
	/// @param x is the x-coordinate of the field point
	/// @O is the approximation accuracy
	stim::complex<T> field(T x, T y, T z, size_t O){
		std::vector< scalarwave<T> > W = mc(O);
		T result = 0;											//initialize the result to zero (0)
		for(size_t i = 0; i < O; i++){							//for each plane wave
			result += W[i].pos(x, y, z);
		}
		return result;
	}

	std::string str()
	{
		std::stringstream ss;
		ss<<"Beam:"<<std::endl;
		//ss<<"	Central Plane Wave: "<<beam::E0<<" e^i ( "<<beam::k<<" . r )"<<std::endl;
		ss<<"	Beam Direction: "<<d<<std::endl;
		if(NA[0] == 0)
			ss<<"	NA: "<<NA[1];
		else
			ss<<"	NA: "<<NA[0]<<" -- "<<NA[1];

		return ss.str();
	}



};			//end beam

/// Calculate the [0 Nl] terms for the aperture integral based on the give numerical aperture and center obscuration (optional)
/// @param C is a pointer to Nl + 1 values where the terms will be stored
template<typename T>
CUDA_CALLABLE void cpu_aperture_integral(T* C, size_t Nl, T NA, T NA_in = 0){

	size_t table_bytes = (Nl + 1) * sizeof(T);				//calculate the number of bytes required to store the terms
	T cos_alpha_1 = cos(asin(NA_in));						//calculate the cosine of the angle subtended by the central obscuration
	T cos_alpha_2 = cos(asin(NA));							//calculate the cosine of the angle subtended by the aperture

	// the aperture integral is computed using four individual Legendre polynomials, each a function of the angles subtended
	//		by the objective and central obscuration
	T* Pln_a1 = (T*) malloc(table_bytes);
	stim::legendre<T>(Nl-1, cos_alpha_1, &Pln_a1[1]);
	Pln_a1[0] = 1;

	T* Pln_a2 = (T*) malloc(table_bytes);
	stim::legendre<T>(Nl-1, cos_alpha_2, &Pln_a2[1]);
	Pln_a2[0] = 1;

	T* Plp_a1 = (T*) malloc(table_bytes+sizeof(T));
	stim::legendre<T>(Nl+1, cos_alpha_1, Plp_a1);

	T* Plp_a2 = (T*) malloc(table_bytes+sizeof(T));
	stim::legendre<T>(Nl+1, cos_alpha_2, Plp_a2);

	for(size_t l = 0; l <= Nl; l++){
		C[l] = Plp_a1[l+1] - Plp_a2[l+1] - Pln_a1[l] + Pln_a2[l];
	}

	free(Pln_a1);
	free(Pln_a2);
	free(Plp_a1);
	free(Plp_a2);
}

/// performs linear interpolation into a look-up table
template<typename T>
T lut_lookup(T* lut, T val, size_t N, T min_val, T delta, size_t stride = 0){
	size_t idx = (size_t)((val - min_val) / delta);
	T alpha = val - idx * delta + min_val;

	if(alpha == 0) return lut[idx];
	else return lut[idx * stride] * (1 - alpha) + lut[ (idx+1) * stride] * alpha;
}

template<typename T>
void cpu_scalar_psf(stim::complex<T>* F, size_t N, T* r, T* phi, T lambda, T A, stim::vec3<T> f, T NA, T NA_in, int Nl){
	T k = stim::TAU / lambda;

	T* C = (T*) malloc( (Nl + 1) * sizeof(T) );					//allocate space for the aperture integral terms
	cpu_aperture_integral(C, Nl, NA, NA_in);			//calculate the aperture integral terms
	memset(F, 0, N * sizeof(stim::complex<T>));
#ifdef NO_CUDA
	memset(F, 0, N * sizeof(stim::complex<T>));
	T jl, Pl, kr, cos_phi;

	double vm;
	double* jv = (double*) malloc( (Nl + 1) * sizeof(double) );
	double* yv = (double*) malloc( (Nl + 1) * sizeof(double) );
	double* djv= (double*) malloc( (Nl + 1) * sizeof(double) );
	double* dyv= (double*) malloc( (Nl + 1) * sizeof(double) );

	T* Pl_cos_phi = (T*) malloc((Nl + 1) * sizeof(T));

	for(size_t n = 0; n < N; n++){								//for each point in the field
		kr = k * r[n];											//calculate kr (the optical distance between the focal point and p)
		cos_phi = std::cos(phi[n]);								//calculate the cosine of phi
		stim::bessjyv_sph<double>(Nl, kr, vm, jv, yv, djv, dyv);		//compute the list of spherical bessel functions from [0 Nl]
		stim::legendre<T>(Nl, cos_phi, Pl_cos_phi);				//calculate the [0 Nl] legendre polynomials for this point

		for(int l = 0; l <= Nl; l++){
			jl = (T)jv[l];
			Pl = Pl_cos_phi[l];
			F[n] += pow(complex<T>(0, 1), l) * jl * Pl * C[l];
		}
		F[n] *= A * stim::TAU;
	}

	free(C);
	free(Pl_cos_phi);
#else
	T min_r = r[0];
	T max_r = r[0];
	for(size_t i = 0; i < N; i++){								//find the minimum and maximum values of r (min and max distance from the focal point)
		if(r[i] < min_r) min_r = r[i];
		if(r[i] > max_r) max_r = r[i];
	}
	T min_kr = k * min_r;
	T max_kr = k * max_r;

	//temporary variables
	double vm;
	double* jv = (double*) malloc( (Nl + 1) * sizeof(double) );
	double* yv = (double*) malloc( (Nl + 1) * sizeof(double) );
	double* djv= (double*) malloc( (Nl + 1) * sizeof(double) );
	double* dyv= (double*) malloc( (Nl + 1) * sizeof(double) );

	size_t Nlut = (size_t)sqrt(N) * 2;
	T* bessel_lut = (T*) malloc(sizeof(T) * (Nl+1) * Nlut);
	T delta_kr = (max_kr - min_kr) / (Nlut-1);
	for(size_t kri = 0; kri < Nlut; kri++){
		stim::bessjyv_sph<double>(Nl, min_kr + kri * delta_kr, vm, jv, yv, djv, dyv);		//compute the list of spherical bessel functions from [0 Nl]
		for(size_t l = 0; l <= Nl; l++){
			bessel_lut[kri * (Nl + 1) + l] = (T)jv[l];
		}
	}

	T* Pl_cos_phi = (T*) malloc((Nl + 1) * sizeof(T));
	T kr, cos_phi, jl, Pl;
	for(size_t n = 0; n < N; n++){								//for each point in the field
		kr = k * r[n];											//calculate kr (the optical distance between the focal point and p)
		cos_phi = std::cos(phi[n]);								//calculate the cosine of phi
		stim::legendre<T>(Nl, cos_phi, Pl_cos_phi);				//calculate the [0 Nl] legendre polynomials for this point

		for(int l = 0; l <= Nl; l++){
			jl = lut_lookup<T>(&bessel_lut[l], kr, Nlut, min_kr, delta_kr, Nl+1);
			Pl = Pl_cos_phi[l];
			F[n] += pow(complex<T>(0, 1), l) * jl * Pl * C[l];
		}
		F[n] *= A * stim::TAU;
	}
#endif
}


template<typename T>
void cpu_scalar_psf(stim::complex<T>* F, size_t N, T* x, T* y, T* z, T lambda, T A, stim::vec3<T> f, T NA, T NA_in, int Nl){
	T* r = (T*) malloc(N * sizeof(T));					//allocate space for p in spherical coordinates
	T* phi = (T*) malloc(N * sizeof(T));				//	only r and phi are necessary (the scalar PSF is symmetric about theta)

	stim::vec3<T> p, ps;
	for(size_t i = 0; i < N; i++){
		(x == NULL) ? p[0] = 0 : p[0] = x[i];	// test for NULL values and set positions
		(y == NULL) ? p[1] = 0 : p[1] = y[i];
		(z == NULL) ? p[2] = 0 : p[2] = z[i];

		ps = p.cart2sph();						//convert from cartesian to spherical coordinates
		r[i] = ps[0];							//store r
		phi[i] = ps[2];							//phi = [0 pi]
	}

	cpu_scalar_psf(F, N, r, phi, lambda, A, f, NA, NA_in, Nl);		//call the spherical coordinate CPU function

	free(r);
	free(phi);
}

}			//end namespace stim

#endif