codebase / stimlib

#ifndef STIM_SPH_HARMONICS
#define STIM_SPH_HARMONICS

#include <complex>
#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
		unsigned int coeff_1d(unsigned int l, int m) {			//convert (l,m) to i (1D coefficient index)
			return pow(l + 1, 2) - (l - m) - 1;
		}
		void coeff_2d(size_t c, unsigned int& l, int& m) {		//convert a 1D coefficient index into (l, m)
			l = (unsigned int)ceil(sqrt((double)c + 1)) - 1;		//the major index is equal to sqrt(c) - 1
			m = (int)(c - (size_t)(l * l)) - (int)l;			//the minor index is calculated by finding the difference
		}

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

		void push(T 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;
		}

		T getc(unsigned int l, int m) {
			unsigned int c = coeff_1d(l, m);
			return C[c];
		}

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

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

		std::vector<T> getC() const {
			return C;
		}
		//calculate the value of the SH basis function (l, m) at (theta, phi)
		//here, theta = [0, PI], phi = [0, 2*PI]
		T SH(unsigned int l, int m, T theta, T 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();
			}
		}

		/// Calculate the spherical harmonic result given a 1D coefficient index
		T SH(size_t c, T theta, T phi) {
			unsigned int l;
			int m;
			coeff_2d(c, l, m);
			return SH(l, m, theta, phi);
		}



		/// 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(T theta, T phi, T val) {

			int l, m;
			T 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::vec3<T> > 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();
		}

		void pdf(std::vector<stim::vec3<T> > sph_pts, size_t c) {
			unsigned int l;
			int m;
			coeff_2d(c, l, m);
			pdf(sph_pts, l, m);
		}

		/// Project a set of samples onto a spherical harmonic basis
		void project(std::vector<stim::vec3<T> > 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], sph_pts[p][0]);
			}
			mcEnd();
		}
		void project(std::vector<stim::vec3<T> > sph_pts, size_t c) {
			unsigned int l;
			int m;
			coeff_2d(c, l, m);
			project(sph_pts, l, m);
		}

		/// Generates a PDF describing the density distribution of points on a sphere
		/// @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
		/// @param norm, a boolean that sets where the output vectors will be normalized between 0 and 1.
		/// @param 
		void pdf(std::vector<stim::vec3<T> > sph_pts, unsigned int l, int m, stim::vec3<T> c = stim::vec3<T>(0, 0, 0), unsigned int n = 1000, bool norm = true, std::vector<T> w = std::vector<T>())
		{
			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<T> > sphere = stim::Random<T>::sample_sphere(n, 1.0, stim::TAU);
			if (w.size() < nP)
				w = std::vector<T>(nP, 1.0);

			for (int i = 0; i < n; i++)
			{
				T val = 0;
				for (int j = 0; j < nP; j++)
				{
					stim::vec3<T> temp = sph_pts[j] - c;
					if (temp.dot(sphere[i]) > 0)
						val += pow(temp.dot(sphere[i]), 4)*w[j];
				}
				weights.push_back(val);
			}

			mcBegin(l, m);		//begin spherical harmonic sampling

			if (norm)
			{
				T min = *std::min_element(weights.begin(), weights.end());
				T max = *std::max_element(weights.begin(), weights.end());
				for (unsigned int i = 0; i < n; i++)
				{
					stim::vec3<T> sph = sphere[i].cart2sph();
					mcSample(sph[1], sph[2], (weights[i] - min) / (max - min));
				}

			}
			else {
				for (unsigned int i = 0; i < n; i++)
				{
					stim::vec3<T> 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 coordinate (theta, phi)
		T p(T theta, T phi) {
			T 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;
		}

		/// Returns the derivative of the spherical function with respect to theta
		///		return value is in cartesian coordinates
		vec3<T> dtheta(T theta, T phi, T d = 0.01) {
			T r = p(theta, phi);											//calculate the value of the spherical function at three points
			T rt = p(theta + d, phi);
			//double rp = p(theta, phi + d);

			vec3<T> s(r, theta, phi);										//get the spherical coordinate position for all three points
			vec3<T> st(rt, theta + d, phi);
			//vec3<double> sp(rp, theta, phi + d);

			vec3<T> c = s.sph2cart();
			vec3<T> ct = st.sph2cart();
			//vec3<double> cp = sp.sph2cart();

			vec3<T> dt = (ct - c)/d;									//calculate the derivative
			return dt;
		}

		/// Returns the derivative of the spherical function with respect to phi
		///		return value is in cartesian coordinates
		vec3<T> dphi(T theta, T phi, T d = 0.01) {
			T r = p(theta, phi);											//calculate the value of the spherical function at three points
			//double rt = p(theta + d, phi);
			T rp = p(theta, phi + d);

			vec3<T> s(r, theta, phi);										//get the spherical coordinate position for all three points
			//vec3<double> st(rt, theta + d, phi);
			vec3<T> sp(rp, theta, phi + d);

			vec3<T> c = s.sph2cart();
			//vec3<double> ct = st.sph2cart();
			vec3<T> cp = sp.sph2cart();

			vec3<T> dp = (cp - c) / d;									//calculate the derivative
			return dp;
		}
		
		/// Returns the value of the function at the coordinate (theta, phi)
		/// @param theta = [0, 2pi]
		/// @param phi = [0, pi]
		T operator()(T theta, T phi) {
			return p(theta, phi);			
		}

		//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));
		}
		/// Fill an NxN grid with the spherical function for theta = [0 2pi] and phi = [0 pi]
		void get_func(T* data, size_t X, size_t Y) {
			T dt = stim::TAU / (T)X;			//calculate the step size in each direction
			T dp = stim::PI / (T)(Y - 1);
			for (size_t ti = 0; ti < X; ti++) {
				for (size_t pi = 0; pi < Y; pi++) {
					data[pi * X + ti] = (*this)((T)ti * dt, (T)pi * dp);
				}
			}
		}

		/// Project a spherical function onto the basis using C coefficients
		/// @param data is a pointer to the function values in (theta, phi) coordinates
		/// @param N is the number of samples along each axis, where theta = [0 2pi), phi = [0 pi]
		void project(T* data, size_t x, size_t y, size_t nc) {
			stim::cpu2image(data, "test.ppm", x, y, stim::cmBrewer);
			C.resize(nc, 0);													//resize the coefficient array to store the necessary coefficients
			T dtheta = stim::TAU / (T)(x - 1);									//calculate the grid spacing along theta
			T dphi = stim::PI / (T)y;											//calculate the grid spacing along phi
			T theta, phi;
			for (size_t c = 0; c < nc; c++) {									//for each coefficient
				for (size_t theta_i = 0; theta_i < x; theta_i++) {				//for each coordinate in the provided array
					theta = theta_i * dtheta;									//calculate theta
					for (size_t phi_i = 0; phi_i < y; phi_i++) {
						phi = phi_i * dphi;										//calculate phi
						C[c] += data[phi_i * x + theta_i] * SH(c, theta, phi) * dtheta * dphi * sin(phi);
					}
				}
			}
		}
	};		//end class sph_harmonics




}


#endif