ctmc_simulation.cpp

#include <iostream>
#include <iomanip>
#include <boost/numeric/odeint.hpp>
#include <boost/math/constants/constants.hpp>
//#define _USE_MATH_DEFINES
#include <cmath>
#include <random>
#include <vector>
#include <fstream>
#include <ctime>
#include <chrono>
#include <cstdio>
#include <string>
#include "./toml11/toml.hpp"
#include "mpi.h"
using namespace boost::numeric::odeint;

typedef std::mt19937 MyRNG;
typedef std::vector< double > state_type;
typedef runge_kutta_cash_karp54< state_type > error_stepper_type;
typedef controlled_runge_kutta< error_stepper_type > controlled_stepper_type;
typedef decltype(std::random_device()()) seed_type;

struct Params {
	const double 	pi = boost::math::constants::pi<double>();
	const double 	ZZZ = 1.0;
	const int     	ncycle = 10;
	const int     	neqn = 7;
	double 	floq;
	double  ip;
	double	intensity;
	double  pola;
	double	cep;
	double	amt;
	double  tsts;
	double  T1;
	//const double  	waveno = floq / 137.0;
	double	alfaN; //static polarizability of Ar atom
	double	alfaI; //Ar cation
	//const double 	vmax = 1.5; // the maximal initial transverse velocity
	//const double 	wt_rel = std::pow(10, -4); //the maximal relative weight for each orbit
};


class Orbit {
public:
	Orbit(const Params& prms, const seed_type se) :params(prms), seed(se) { random_initialize(); }

	double envlp(double t) {
		if (t < 0.0) {
			return 0.0;
		} else if (t < params.ncycle * params.tsts) {
			return pow2(std::sin(t * params.pi / params.T1));
		} else {
			return 0.0;
		}
    }

	bool isInitialized(double& tt0, double& wt, state_type& y) {
		double ft0, epsi, epsiy, epsiall;		

		tt0 = uniform_dist(rng1) * params.T1;
		//ft0 = std::pow(std::sin(params.pi * tt0 / params.T1), 2);
		//auto prfl = params.pi * std::sin(2.0 * params.pi * tt0 / params.T1) / params.T1 / params.floq;

		//epsi = params.amt * (ft0 * std::cos( params.floq * tt0 + params.cep) + prfl * std::sin(params.floq * tt0 + params.cep));
		//epsiy = params.pola * params.amt * (ft0 * std::sin(params.floq * tt0 + params.cep) - prfl * std::cos(params.floq * tt0 + params.cep));  //params.pola * params.amt * ft0 * cos(params.floq * tt0 * params.nfrequency + params.alpha) / sqrt(1.0 + pow(params.pola, 2));
		epsi = params.amt * envlp(tt0) * std::cos(params.floq * tt0 + params.cep);
		epsiy = -params.pola * params.amt * envlp(tt0) * std::sin(params.floq * tt0 + params.cep);
		epsiall = std::sqrt(std::pow(epsi, 2) + std::pow(epsiy, 2));


		double theta, eta, zetax, zetay;

		theta = std::atan2(epsiy, epsi);
		/*theta = atan(abs(epsiy / epsi));
		if ((epsi < 0.0) and (epsiy >= 0.0)) theta = M_PI - theta;
		if ((epsi > 0.0) and (epsiy < 0.0)) theta = 2.0 * M_PI - theta;
		if ((epsi < 0.0) and (epsiy < 0.0)) theta = M_PI + theta;*/

		eta = r0(tt0, epsiall);
		zetax = -std::abs(eta / 2.0 * std::cos(theta));	//sometimes, abs(0.2) = 0  <<<<<<<<
		if (epsi < 0.0) zetax = -zetax;
		zetay = -std::abs(eta / 2.0 * std::sin(theta));
		if (epsiy < 0.0) zetay = -zetay;

		
		double vx0, vy0;
		

		auto kappa = std::sqrt(2.0 * starkIp(tt0, epsiall));
		auto kappa0 = std::sqrt(2.0 * params.ip);

		vx0 = std::sqrt(params.amt / kappa0) * normal_dist(rng2);
		vy0 = std::sqrt(params.amt / kappa0) * normal_dist(rng3);

		double pzmax, pzmax0;

		pzmax = std::sqrt(std::pow(2.0 * pow3(kappa0) / params.amt, 2.0/kappa0-1.0) * std::exp(-2.0 * pow3(kappa0) / 3.0 / params.amt) 
				* std::exp(-kappa0 * (pow2(vx0) + pow2(vy0)) / params.amt));
		pzmax = uniform_dist(rng4) * pzmax;
		pzmax0 = std::sqrt(std::pow(2.0 * pow3(kappa) / epsiall, 2.0/kappa-1.0) * std::exp(-2.0 * pow3(kappa) / 3.0 / epsiall) 
				* std::exp(-kappa * (pow2(vx0) + pow2(vy0)) / epsiall));

		if (pzmax > pzmax0) return false;
		
		y[0] = zetax;
		y[1] = zetay;
		y[2] = 0.0;

		y[3] = -vx0 * std::sin(theta);
		y[4] = vx0 * std::cos(theta);
		y[5] = vy0;

		y[6] = -(y[0] * y[3] + y[1] * y[4] + y[2] * y[5]) + starkIp(tt0, epsiall) * tt0;

		return true;
	}

	void trajectory(const state_type& y, state_type& yp, const double t) {

		double a_2 = 0.0;
		double b0 = 0.1;

		//ft = std::pow(std::sin(params.pi * t / params.T1), 2);
		auto r = std::sqrt(std::pow(y[0], 2) + std::pow(y[1], 2) + std::pow(y[2], 2));

		//auto prfl = params.pi * std::sin(2.0 * params.pi * t / params.T1) / params.T1 / params.floq;

		//epsix = params.amt * (ft * std::cos( params.floq * t + params.cep) + prfl * std::sin(params.floq * t + params.cep));
		//epsiy = params.pola * params.amt * (ft * std::sin(params.floq * t + params.cep) - prfl * std::cos(params.floq * t + params.cep));;  //params.pola * params.amt * ft * cos(params.floq * t * params.nfrequency + params.alpha) / sqrt(1.0 + pow(params.pola, 2));
		auto epsix = params.amt * envlp(t) * std::cos(params.floq * t + params.cep);
		auto epsiy = -params.pola * params.amt * envlp(t) * std::sin(params.floq * t + params.cep);

		//auto sft = std::exp(-b0 / (r + a_2)) * params.alfaI;
		//auto dp = epsix * y[0] + epsiy * y[1];
		auto r3 = std::pow(r + a_2, 3);
		//auto rv65 = b0 / std::pow(r + a_2, 6) - 3.0 / std::pow(r + a_2, 5);

		yp[0] = y[3];
		yp[1] = y[4];
		yp[2] = y[5];
		yp[3] = -params.ZZZ * y[0] / r3 - epsix;
		yp[4] = -params.ZZZ * y[1] / r3 - epsiy;
		yp[5] = -params.ZZZ * y[2] / r3;
		yp[6] = -((std::pow(y[3], 2) + std::pow(y[4], 2) + std::pow(y[5], 2)) / 2.0 - 2.0 * params.ZZZ / (r + a_2));
	}

	

private:
	void random_initialize() {
		srand(seed);
		rng1.seed(rand() % 1000);
		rng2.seed(rand() % 1000);
		rng3.seed(rand() % 1000);
		rng4.seed(rand() % 1000);
	}

	double starkIp(double t0, double epsiall) {
		return params.ip + 0.5 * (params.alfaN - params.alfaI) * std::pow(epsiall, 2);
	}

	double r0(double t0, double epsi) {
		//the largest solution to the equation: V(\eta) = - (1-sqrt(2*ip)/2)/2/\eta - \eta * epsi / 8 - (1-8*alfaI*epsi)/8/\eta^2=-ip/4
		auto alfa = 1.0;
		auto beta = 1.0 - std::sqrt(2 * starkIp(t0, epsi)) / 2.0;

		auto a = std::abs(epsi);
		auto b = -2.0 * starkIp(t0, epsi);
		auto c = 4.0 * beta;
		auto d = alfa;
		auto p = (3.0 * a * c - pow2(b)) / 9.0 / pow2(a); // p < 0.0
		auto q = (2 * pow3(b) / 27.0 / pow3(a) - b * c / 3.0 / pow2(a) + d / a) / 2.0; // q < 0.0
		auto r = (q < 0.0)? (-std::sqrt(std::abs(p))) : (std::sqrt(std::abs(p)));
		auto bigd = pow2(q) + pow3(p);
		double sol;

		if (bigd <= 0.0) {
			auto phi = std::acos(q / pow3(r));
			sol = -2.0 * r * std::cos(phi / 3.0);
		} else {
			auto phi = std::acosh(q / pow3(r));
			sol = -2.0 * r * std::cosh(phi / 3.0);
		}
		return sol - b / 3.0 / a;
	}

	double pow2(double a) {
		return std::pow(a, 2);
	}
	double pow3(double a) {
		return std::pow(a, 3);
	}

private:
	Params params;
	seed_type seed;
	MyRNG rng1, rng2, rng3, rng4;
	std::uniform_real_distribution<double> uniform_dist{ 0.0, 1.0 };
	std::normal_distribution<double> normal_dist{ 0.0, 1.0 };
};



int main(int argc, char** argv) {

	Params params;
	auto conf = toml::parse( "conf.toml" );

	params.floq = toml::find<double>(conf, "laser", "frequency");
	params.intensity = toml::find<double>(conf, "laser", "intensity");
	if (params.intensity == 0.0) std::cout << "laser intensity value Error!" << std::endl;
	params.pola = toml::find<double>(conf, "laser", "ellipticity");
	params.cep = toml::find<double>(conf, "laser", "cep");

	params.ip = toml::find<double>(conf, "atom", "ip");
	params.alfaN = toml::find<double>(conf, "atom", "alphaN");
	params.alfaI = toml::find<double>(conf, "atom", "alphaI");
	/*auto conf = toml::parse_file( "conf.toml" );

	params.floq = conf["laser"]["frequency"].value_or(0.0);
	params.intensity = conf["laser"]["intensity"].value_or(0.0);
	if (params.intensity == 0.0) std::cout << "laser intensity value Error!" << std::endl;
	params.pola = conf["laser"]["ellipticity"].value_or(0.0);
	params.cep = conf["laser"]["cep"].value_or(0.0);

	params.ip = conf["atom"]["ip"].value_or(0.0);
	params.alfaN = conf["atom"]["alphaN"].value_or(0.0);
	params.alfaI = conf["atom"]["alphaI"].value_or(0.0);*/
	
	params.amt = 0.169 * std::sqrt(params.intensity) / std::sqrt(1.0 + std::pow(params.pola, 2));
	params.tsts = 2.0 * params.pi / params.floq;
	params.T1 = params.ncycle * params.tsts;
	
	//std::cout << "laser intensity = " << params.intensity << std::endl;
	//std::cout << "laser electric field = " << params.amt << std::endl;



	int numprocs, myid;
	MPI_Init(NULL,NULL);
	MPI_Comm_size(MPI_COMM_WORLD, &numprocs);
	MPI_Comm_rank(MPI_COMM_WORLD, &myid);
	//myid = 1;



	double tt0, wt, rx0, ry0, rz0, px0, py0, pz0;
	double tfinal, tprint;
	double t;


	double abserr = 0.0, relerr = 1e-10;
	auto controlled_stepper = make_controlled< runge_kutta_dopri5<state_type> >(abserr, relerr);

	auto start = std::chrono::system_clock::now();
	time_t start_time = std::chrono::system_clock::to_time_t(start);
	std::cout << "started computation at " << ctime(&start_time);

	char file_postfix[10];
	std::string filename01("./t0rat");
	std::string filename02("./ini");
	std::string filename03("./electron");
	std::string filename04("./traj");
	std::string filename05("./counts");
	snprintf(file_postfix, sizeof(file_postfix), "%03d", myid);
	std::string file_path01 = filename01 + std::string(file_postfix) + std::string(".txt");
	std::string file_path02 = filename02 + std::string(file_postfix) + std::string(".txt");
	std::string file_path03 = filename03 + std::string(file_postfix) + std::string(".txt");
	std::string file_path04 = filename04 + std::string(file_postfix) + std::string(".txt");
	std::string file_path05 = filename05 + std::string(file_postfix) + std::string(".txt");


	std::ofstream file01;
	file01.open(file_path01);
	std::ofstream file02;
	file02.open(file_path02);
	std::ofstream file03;
	file03.open(file_path03);
	std::ofstream file04;
	file04.open(file_path04);
	std::ofstream file05;
    file05.open(file_path05);


	file01 << std::setprecision(10);
	file02 << std::setprecision(10);
	file03 << std::setprecision(10);
	file04 << std::setprecision(10);

	std::random_device rd;
	auto sid = rd();
	std::cout << "The seed is: " << sid << std::endl;

	//Params params;
	state_type y(params.neqn);
	Orbit orbit(params, sid);

	std::cout << "pi = " << params.pi << std::endl;
	std::cout << "laser intensity = " << params.intensity << std::endl;
	std::cout << "laser electric field = " << params.amt << std::endl;

	int count = 0;
	for (int i = 0; i != 100000000;) {


		if (not orbit.isInitialized(tt0, wt, y)) continue;

		rx0 = y[0];
		ry0 = y[1];
		rz0 = y[2];
		px0 = y[3];
		py0 = y[4];
		pz0 = y[5];



		t = tt0;
		tfinal = params.T1;
		tprint = 0.01;

		double dt;
		dt = tprint;

		bool first_too_close = false;// excluding some nearest orbits will induce hollows in the spectrum along the polarization axis
		auto write_state = [&first_too_close](const state_type& yy, double tt) {
			//file04 << tt << '\t' << yy[0] << '\t' << yy[1] << '\t' << yy[2] << '\t' << yy[3] << '\t' << yy[4] << std::endl;
			auto r = std::sqrt(yy[0] * yy[0] + yy[1] * yy[1]);
			if (!first_too_close) {
				if (r < 0.01) first_too_close = true;
			}
		};
		auto trajectory = [&orbit](const state_type& yy, state_type& yyp, const double tt) { orbit.trajectory(yy, yyp, tt); };
		integrate_adaptive(controlled_stepper, trajectory, y, t, tfinal, dt);
		//integrate_adaptive(controlled_stepper, trajectory, y, t, tfinal, dt, write_state);

		/*const double dt = (tfinal - tt0) / 1000.0;
		for (double t = tt0; t <= tfinal; t += dt) {
			integrate_adaptive(controlled_stepper, trajectory, y, t, t + dt, tprint);
		}

		for (double t = tt0; t <= tfinal;) {
			controlled_stepper.try_step(trajectory, y, t, dt)
		}*/
		//if (first_too_close) continue;


		auto rr1 = std::sqrt(std::pow(y[0], 2) + std::pow(y[1], 2) + std::pow(y[2], 2));
		auto pp2 = std::pow(y[3], 2) + std::pow(y[4], 2) + std::pow(y[5], 2);
		auto ee1 = -params.ZZZ / rr1 + pp2 / 2.0;

		++count;
		if (ee1 > 0.0) {
			++i;
			file01 << tt0 << '\t' << wt << std::endl;
			file02 << rx0 << '\t' << ry0 << '\t' << rz0 << '\t' << px0 << '\t' << py0 << '\t' << pz0 << std::endl;
			file03 << y[0] << '\t' << y[1] << '\t' << y[2] << '\t' << y[3] << '\t' << y[4] << '\t' << y[5] << '\t' << y[6] << std::endl;
		}
	}
	file05 << count << std::endl;

	file01.close();
	file02.close();
	file03.close();
	file04.close();
	file05.close();

	auto end = std::chrono::system_clock::now();
	std::chrono::duration<double> elapsed_seconds = end - start;
	time_t end_time = std::chrono::system_clock::to_time_t(end);
	std::cout << "finished computation at " << ctime(&end_time)
		<< "elapsed time: " << elapsed_seconds.count() << "s"
		<< std::endl;

	MPI_Finalize();
}