codebase / stimlib

#ifndef STIM_SPH_HARMONICS
#define STIM_SPH_HARMONICS

#include <complex>
#include <stim/math/vector.h>
#include <boost/math/special_functions/spherical_harmonic.hpp>
#include <stim/math/constants.h>
#include <stim/math/random.h>
#include <vector>

#define WIRE_SCALE 1.001
namespace stim{

template<class T>
class spharmonics{

public:
	std::vector<T> C;	//list of SH coefficients

protected:


	unsigned int mcN;	//number of Monte-Carlo samples

	//calculate the value of the SH basis function (l, m) at (theta, phi)
		//here, theta = [0, PI], phi = [0, 2*PI]
	double SH(int l, int m, double theta, double phi){
		//std::complex<T> result = boost::math::spherical_harmonic(l, m, phi, theta);
		//return result.imag() + result.real();

		//this calculation is based on calculating the real spherical harmonics:
		//		https://en.wikipedia.org/wiki/Spherical_harmonics#Addition_theorem
		if (m < 0) {
			return sqrt(2.0) * pow(-1, m) * boost::math::spherical_harmonic(l, abs(m), phi, theta).imag();
		}
		else if (m == 0) {
			return boost::math::spherical_harmonic(l, m, phi, theta).real();
		}
		else {
			return sqrt(2.0) * pow(-1, m) * boost::math::spherical_harmonic(l, m, phi, theta).real();
		}
	}

	unsigned int coeff_1d(unsigned int l, int m){
		return pow(l + 1, 2) - (l - m) - 1;
	}

	


public:
	spharmonics() {
		mcN = 0;
	}
	spharmonics(size_t c) : spharmonics() {
		resize(c);
	}

	void push(double c){
		C.push_back(c);
	}

	void resize(unsigned int n){
		C.resize(n);
	}

	void setc(unsigned int l, int m, T value){
		unsigned int c = coeff_1d(l, m);
		C[c] = value;
	}

	void setc(unsigned int c, T value){
		C[c] = value;
	}

	unsigned int getSize() const{
		return C.size();
	}

	std::vector<T> getC() const{
		return C;
	}

	

	/// Initialize Monte-Carlo sampling of a function using N spherical harmonics coefficients

	/// @param N is the number of spherical harmonics coefficients used to represent the user function
	void mcBegin(unsigned int coefficients){
		C.resize(coefficients, 0);
		mcN = 0;
	}

	void mcBegin(unsigned int l, int m){
		unsigned int c = pow(l + 1, 2) - (l - m);
		mcBegin(c);
	}

	void mcSample(double theta, double phi, double val){

		int l, m;
		double sh;

		l = m = 0;
		for(unsigned int i = 0; i < C.size(); i++){

			sh = SH(l, m, theta, phi);
			C[i] += sh * val;

			m++;			//increment m

			//if we're in a new tier, increment l and set m = -l
			if(m > l){		
				l++;
				m = -l;
			}
		}	//end for all coefficients

		//increment the number of samples
		mcN++;

	}	//end mcSample()

	void mcEnd(){

		//divide all coefficients by the number of samples
		for(unsigned int i = 0; i < C.size(); i++)
			C[i] /= mcN;
	}

	/// Generates a PDF describing the probability distribution of points on a spherical surface
	/// @param sph_pts is a list of points in spherical coordinates (theta, phi) where theta = [0, 2pi] and phi = [0, pi]
	/// @param l is the maximum degree of the spherical harmonic function
	/// @param m is the maximum order

	void pdf(std::vector<stim::vec<double> > sph_pts, unsigned int l, int m){
			
		mcBegin( l, m );		//begin spherical harmonic sampling

		unsigned int nP = sph_pts.size();

		for(unsigned int p = 0; p < nP; p++){
			mcSample(sph_pts[p][1], sph_pts[p][2], 1.0);
		}

		mcEnd();
	}

	/// Generates a PDF describing the probability distribution of points on a spherical surface
	/// @param sph_pts is a list of points in cartesian coordinates 
	/// @param l is the maximum degree of the spherical harmonic function
	/// @param m is the maximum order
	/// @param c is the centroid of the points in sph_pts. DEFAULT 0,0,0
	/// @param n is the number of points of the surface of the sphere used to create the PDF. DEFAULT 1000
	void pdf(std::vector<stim::vec3<double> > sph_pts, unsigned int l, int m, stim::vec3<double> c = stim::vec3<double>(0,0,0), unsigned int n = 1000)
	{
		std::vector<double> weights;		///the weight at each point on the surface of the sphere.
//		weights.resize(n);
		unsigned int nP = sph_pts.size();
		std::vector<stim::vec3<double> > sphere = stim::Random<double>::sample_sphere(n, 1.0, stim::TAU);
		for(int i = 0; i < n; i++)
		{
			double val = 0;
			for(int j = 0; j < nP; j++)
			{
				stim::vec3<double> temp = sph_pts[j] - c;
				if(temp.dot(sphere[i]) > 0)
					val += pow(temp.dot(sphere[i]),4);
			}
			weights.push_back(val);
		}
		
		
		mcBegin(l, m);		//begin spherical harmonic sampling
		
		for(unsigned int i = 0; i < n; i++)
		{
			stim::vec3<double> sph = sphere[i].cart2sph();
			mcSample(sph[1], sph[2], weights[i]);
		}

		mcEnd();
	}

	std::string str(){

		std::stringstream ss;

		int l, m;
		l = m = 0;
		for(unsigned int i = 0; i < C.size(); i++){
				
			ss<<C[i]<<'\t';

			m++;			//increment m

			//if we're in a new tier, increment l and set m = -l
			if(m > l){
				l++;
				m = -l;

				ss<<std::endl;
					
			}
		}

		return ss.str();


	}

	/// Returns the value of the function at the coordinate (theta, phi)

	/// @param theta = [0, 2pi]
	/// @param phi = [0, pi]
	double operator()(double theta, double phi){

		double fx = 0;

		int l = 0;
		int m = 0;
		for(unsigned int i = 0; i < C.size(); i++){
			fx += C[i] * SH(l, m, theta, phi);
			m++;
			if(m > l){
				l++;
				m = -l;					
			}

		}

		return fx;
	}

	/// Fill an NxN grid with the spherical function for theta = [0 2pi] and phi = [0 pi]
	void grid(T* data, size_t N){
		double dt = stim::TAU / (double)N;			//calculate the step size in each direction
		double dp = stim::PI / (double)(N - 1);
		for(size_t ti = 0; ti < N; ti++){
			for(size_t pi = 0; pi < N){
				data[pi * N + ti] = (*this)((double)ti * dt, (double)pi * dp);
			}
		}
	}
/*
	//overload arithmetic operations
	spharmonics<T> operator*(T rhs) const {
		spharmonics<T> result(C.size());				//create a new spherical harmonics object
		for (size_t c = 0; c < C.size(); c++)			//for each coefficient
			result.C[c] = C[c] * rhs;					//calculate the factor and store the result in the new spharmonics object
		return result;
	}

	spharmonics<T> operator+(spharmonics<T> rhs) {
		size_t low = std::min(C.size(), rhs.C.size());				//store the number of coefficients in the lowest object
		size_t high = std::max(C.size(), rhs.C.size());				//store the number of coefficients in the result
		bool rhs_lowest = false;								//true if rhs has the lowest number of coefficients
		if (rhs.C.size() < C.size()) rhs_lowest = true;			//if rhs has a lower number of coefficients, set the flag

		spharmonics<T> result(high);								//create a new object
		size_t c;
		for (c = 0; c < low; c++)						//perform the first batch of additions
			result.C[c] = C[c] + rhs.C[c];						//perform the addition

		for (c = low; c < high; c++) {
			if (rhs_lowest)
				result.C[c] = C[c];
			else
				result.C[c] = rhs.C[c];
		}
		return result;
	}

	spharmonics<T> operator-(spharmonics<T> rhs) {
		return (*this) + (rhs * (T)(-1));
	}
*/
};		//end class sph_harmonics




}


#endif