Include electric field (e37fadc6) · Commits · Cianciosa, Mark / graph_framework

graph_framework/equilibrium.hpp

+100 −12

Original line number	Diff line number	Diff line
		@@ -157,6 +157,17 @@ namespace equilibrium {
		graph::shared_leaf<T, SAFE_MATH> y,
		graph::shared_leaf<T, SAFE_MATH> z) = 0;

		//--------------------------------------------------------------------------
		/// @brief Get the electric field (default: zero — override if needed).
		//--------------------------------------------------------------------------
		virtual graph::shared_vector<T, SAFE_MATH>
		get_electric_field(graph::shared_leaf<T, SAFE_MATH> x,
		graph::shared_leaf<T, SAFE_MATH> y,
		graph::shared_leaf<T, SAFE_MATH> z) {
		auto zero = graph::zero<T, SAFE_MATH>();
		return graph::vector(zero, zero, zero);
		}

		//------------------------------------------------------------------------------
		/// @brief Get the characteristic field.
		///
		@@ -986,6 +997,14 @@ namespace equilibrium {
		/// Pressure scale factor.
		graph::shared_leaf<T, SAFE_MATH> pres_scale;


		// Electric‑potential (Φ) spline coefficients and scale ───────────────
		const backend::buffer<T> epsi_c0;
		const backend::buffer<T> epsi_c1;
		const backend::buffer<T> epsi_c2;
		const backend::buffer<T> epsi_c3;
		graph::shared_leaf<T, SAFE_MATH> epsi_scale;

		/// Minimum R.
		const T rmin;
		/// R grid spacing.
		@@ -1061,6 +1080,9 @@ namespace equilibrium {
		/// Cached magnetic field vector.
		graph::shared_vector<T, SAFE_MATH> b_cache;

		/// Cached electric field vector.
		graph::shared_vector<T, SAFE_MATH> e_cache;

		/// Cached magnetic flux.
		graph::shared_leaf<T, SAFE_MATH> psi_cache;

		@@ -1149,7 +1171,7 @@ namespace equilibrium {

		te_cache = te_scale*build_1D_spline({t0_temp, t1_temp, t2_temp, t3_temp}, psi_cache, dpsi, psimin);

		auto p0_temp = graph::piecewise_1D(pres_c0, psi_cache, dpsi, psimin) - psimin;
		auto p0_temp = graph::piecewise_1D(pres_c0, psi_cache, dpsi, psimin);
		auto p1_temp = graph::piecewise_1D(pres_c1, psi_cache, dpsi, psimin);
		auto p2_temp = graph::piecewise_1D(pres_c2, psi_cache, dpsi, psimin);
		auto p3_temp = graph::piecewise_1D(pres_c3, psi_cache, dpsi, psimin);
		@@ -1161,23 +1183,47 @@ namespace equilibrium {
		ni_cache = te_cache;
		ti_cache = (pressure - ne_cachete_cacheq)/(ni_cache*q);

		auto phi = graph::atan(x, y);

		// Calculate the F(psi) function [T m]
		auto f0 = graph::piecewise_1D(fpol_c0, psi_cache, dpsi, psimin);
		auto f1 = graph::piecewise_1D(fpol_c1, psi_cache, dpsi, psimin);
		auto f2 = graph::piecewise_1D(fpol_c2, psi_cache, dpsi, psimin);
		auto f3 = graph::piecewise_1D(fpol_c3, psi_cache, dpsi, psimin);
		// Toroidal magnetic field [T]
		auto bp = build_1D_spline({f0,f1,f2,f3},psi_cache, dpsi, psimin) / r;

		// Poloidal magnetic field (B_R and B_Z) [T]
		auto br = psi_cache->df(z)/r;
		auto bz = -psi_cache->df(r)/r;
		// Angle in the poloidal plane [rad]
		auto phi = graph::atan(x, y);
		auto cphi = graph::cos(phi);
		auto sphi = graph::sin(phi);
		// Magnetic field vector [T]
		b_cache = graph::vector(brcphi - bpsphi,
		brsphi + bpcphi,
		bz);

		auto cos = graph::cos(phi);
		auto sin = graph::sin(phi);
		// ───────────────────────── Electric field ─────────────────────────
		// 1‑D spline for dΦ/dψ [ V · Wb⁻¹ ]
		auto e0 = graph::piecewise_1D(epsi_c0, psi_cache, dpsi, psimin);
		auto e1 = graph::piecewise_1D(epsi_c1, psi_cache, dpsi, psimin);
		auto e2 = graph::piecewise_1D(epsi_c2, psi_cache, dpsi, psimin);
		auto e3 = graph::piecewise_1D(epsi_c3, psi_cache, dpsi, psimin);
		auto dphi_dpsi = epsi_scale *
		build_1D_spline({e0,e1,e2,e3},
		psi_cache, dpsi, psimin);

		b_cache = graph::vector(brcos - bpsin,
		brsin + bpcos,
		bz);
		// ∇ψ components
		auto dpsi_dr = psi_cache->df(r);
		auto dpsi_dz = psi_cache->df(z);

		// Cylindrical → Cartesian
		auto Er = -dphi_dpsi * dpsi_dr; // E = −∇Φ
		auto Ez = -dphi_dpsi * dpsi_dz;
		auto Ex = Er*cphi;
		auto Ey = Er*sphi;

		// Cache the result
		e_cache = graph::vector(Ex,Ey,Ez);
		}
		}

		@@ -1230,6 +1276,11 @@ namespace equilibrium {
		//------------------------------------------------------------------------------
		efit(const T psimin,
		const T dpsi,
		const backend::buffer<T> epsi_c0,
		const backend::buffer<T> epsi_c1,
		const backend::buffer<T> epsi_c2,
		const backend::buffer<T> epsi_c3,
		graph::shared_leaf<T, SAFE_MATH> epsi_scale,
		const backend::buffer<T> te_c0,
		const backend::buffer<T> te_c1,
		const backend::buffer<T> te_c2,
		@@ -1275,7 +1326,10 @@ namespace equilibrium {
		te_c0(te_c0), te_c1(te_c1), te_c2(te_c2), te_c3(te_c3), te_scale(te_scale),
		ne_c0(te_c0), ne_c1(te_c1), ne_c2(ne_c2), ne_c3(ne_c3), ne_scale(ne_scale),
		pres_c0(pres_c0), pres_c1(pres_c1), pres_c2(pres_c2), pres_c3(pres_c3),
		pres_scale(pres_scale), rmin(rmin), dr(dr), zmin(zmin), dz(dz),
		pres_scale(pres_scale),
		epsi_c0(epsi_c0), epsi_c1(epsi_c1), epsi_c2(epsi_c2), epsi_c3(epsi_c3),
		epsi_scale(epsi_scale),
		rmin(rmin), dr(dr), zmin(zmin), dz(dz),
		fpol_c0(fpol_c0), fpol_c1(fpol_c1), fpol_c2(fpol_c2), fpol_c3(fpol_c3),
		c00(c00), c01(c01), c02(c02), c03(c03),
		c10(c10), c11(c11), c12(c12), c13(c13),
		@@ -1371,6 +1425,14 @@ namespace equilibrium {
		return b_cache;
		}

		virtual graph::shared_vector<T, SAFE_MATH>
		get_electric_field(graph::shared_leaf<T, SAFE_MATH> x,
		graph::shared_leaf<T, SAFE_MATH> y,
		graph::shared_leaf<T, SAFE_MATH> z) {
		set_cache(x, y, z);
		return e_cache;
		}

		//------------------------------------------------------------------------------
		/// @brief Get the characteristic field.
		///
		@@ -1465,6 +1527,11 @@ namespace equilibrium {
		nc_inq_varid(ncid, "te_scale", &varid);
		nc_get_var(ncid, varid, &te_scale_value);

		double epsi_scale_value;
		nc_inq_varid(ncid, "epsi_scale", &varid);
		nc_get_var(ncid, varid, &epsi_scale_value);


		// Load 1D quantities.
		int dimid;

		@@ -1587,6 +1654,20 @@ namespace equilibrium {
		nc_inq_varid(ncid, "ne_c3", &varid);
		nc_get_var(ncid, varid, ne_c3_buffer.data());

		std::vector<double> epsi_c0_buffer(numr);
		std::vector<double> epsi_c1_buffer(numr);
		std::vector<double> epsi_c2_buffer(numr);
		std::vector<double> epsi_c3_buffer(numr);

		nc_inq_varid(ncid, "epsi_c0", &varid);
		nc_get_var(ncid, varid, epsi_c0_buffer.data());
		nc_inq_varid(ncid, "epsi_c1", &varid);
		nc_get_var(ncid, varid, epsi_c1_buffer.data());
		nc_inq_varid(ncid, "epsi_c2", &varid);
		nc_get_var(ncid, varid, epsi_c2_buffer.data());
		nc_inq_varid(ncid, "epsi_c3", &varid);
		nc_get_var(ncid, varid, epsi_c3_buffer.data());

		nc_close(ncid);
		sync.unlock();

		@@ -1599,6 +1680,7 @@ namespace equilibrium {
		auto pres_scale = graph::constant<T, SAFE_MATH> (static_cast<T> (pres_scale_value));
		auto ne_scale = graph::constant<T, SAFE_MATH> (static_cast<T> (ne_scale_value));
		auto te_scale = graph::constant<T, SAFE_MATH> (static_cast<T> (te_scale_value));
		const auto epsi_scale = graph::constant<T, SAFE_MATH> (static_cast<T> (epsi_scale_value));

		const auto fpol_c0 = backend::buffer(std::vector<T> (fpol_c0_buffer.begin(), fpol_c0_buffer.end()));
		const auto fpol_c1 = backend::buffer(std::vector<T> (fpol_c1_buffer.begin(), fpol_c1_buffer.end()));
		@@ -1637,7 +1719,13 @@ namespace equilibrium {
		const auto ne_c2 = backend::buffer(std::vector<T> (ne_c2_buffer.begin(), ne_c2_buffer.end()));
		const auto ne_c3 = backend::buffer(std::vector<T> (ne_c3_buffer.begin(), ne_c3_buffer.end()));

		const auto epsi_c0 = backend::buffer(std::vector<T> (epsi_c0_buffer.begin(), epsi_c0_buffer.end()));
		const auto epsi_c1 = backend::buffer(std::vector<T> (epsi_c1_buffer.begin(), epsi_c1_buffer.end()));
		const auto epsi_c2 = backend::buffer(std::vector<T> (epsi_c2_buffer.begin(), epsi_c2_buffer.end()));
		const auto epsi_c3 = backend::buffer(std::vector<T> (epsi_c3_buffer.begin(), epsi_c3_buffer.end()));

		return std::make_shared<efit<T, SAFE_MATH>> (psimin, dpsi,
		epsi_c0, epsi_c1, epsi_c2, epsi_c3, epsi_scale,
		te_c0, te_c1, te_c2, te_c3, te_scale,
		ne_c0, ne_c1, ne_c2, ne_c3, ne_scale,
		pres_c0, pres_c1, pres_c2, pres_c3, pres_scale,

graph_korc/xkorc.cpp

+45 −35

Original line number	Diff line number	Diff line
		//------------------------------------------------------------------------------
		// xkorc.cpp – Full-orbit Boris integrator (dimensionless, E = 0)
		// xkorc.cpp – Full-orbit Boris integrator (SI units)
		//
		// Reads from init_particles.nc
		// • x, y, z , ux, uy, uz (dimensionless)
		// • dt (dimensionless time-step)
		// • x, y, z , ux, uy, uz (m and m/s)
		// • dt (s)
		// • NSTEPS, DUMP_EVERY (run control)
		// • particle_mass, particle_charge (kg and C)
		//
		// Evolves {x,y,z,ux,uy,uz} with the Boris pusher including electric field.
		//
		// Evolves {x,y,z,ux,uy,uz} with the Boris pusher.
		// NEW: one dedicated writer thread per worker to handle NetCDF output.
		//------------------------------------------------------------------------------

		#include <cmath>
		@@ -39,9 +40,7 @@ void run_korc()
		{
		const timeing::measure_diagnostic t_total("Total Time");

		// ──────────────────────────────────────────────────────────────
		// 1. READ INITIAL DATA & RUN CONTROL
		// ──────────────────────────────────────────────────────────────
		// ─────────────────────────────────── 1. READ INITIAL DATA ────────────────
		int ncid;
		check_nc( nc_open("init_particles.nc", NC_NOWRITE, &ncid) );

		@@ -63,6 +62,8 @@ void run_korc()

		long long h_NSTEPS=0, h_DUMP=0;
		double dt_val=0.0;
		double m_val = 1.0; // defaults protect against missing attributes
		double q_val = 1.0;

		int vid;
		check_nc( nc_inq_varid(ncid,"NSTEPS", &vid) );
		@@ -72,14 +73,18 @@ void run_korc()
		check_nc( nc_inq_varid(ncid,"dt", &vid) );
		check_nc( nc_get_var_double (ncid,vid,&dt_val) );

		if (nc_inq_varid(ncid,"particle_mass",&vid)==NC_NOERR)
		check_nc( nc_get_var_double(ncid,vid,&m_val) );
		if (nc_inq_varid(ncid,"particle_charge",&vid)==NC_NOERR)
		check_nc( nc_get_var_double(ncid,vid,&q_val) );

		check_nc( nc_close(ncid) );

		const std::size_t NSTEPS = static_cast<std::size_t>(h_NSTEPS);
		const std::size_t DUMP_EVERY = static_cast<std::size_t>(h_DUMP );
		const double q_over_m_val = q_val / m_val;

		// ──────────────────────────────────────────────────────────────
		// 2. CREATE GRAPH VARIABLES
		// ──────────────────────────────────────────────────────────────
		// ─────────────────────────────────── 2. CREATE GRAPH VARS ────────────────
		auto x = graph::variable<T,false>(Np,"x");
		auto y = graph::variable<T,false>(Np,"y");
		auto z = graph::variable<T,false>(Np,"z");
		@@ -92,9 +97,7 @@ void run_korc()
		ux->set(i,h_ux[i]); uy->set(i,h_uy[i]); uz->set(i,h_uz[i]);
		}

		// ──────────────────────────────────────────────────────────────
		// 3. THREAD LAYOUT
		// ──────────────────────────────────────────────────────────────
		// ─────────────────────────────────── 3. THREAD LAYOUT ────────────────────
		unsigned n_thr = std::max(1u,
		std::min<unsigned>(jit::context<T>::max_concurrency(), static_cast<unsigned>(Np)));
		std::vector<std::thread> workers(n_thr);
		@@ -113,29 +116,35 @@ void run_korc()
		auto half = graph::constant<T,false>(0.5);
		auto one = graph::constant<T,false>(1.0);
		auto two = graph::constant<T,false>(2.0);
		auto qom = graph::constant<T,false>(static_cast<T>(q_over_m_val));
		auto dt_half = dt * half; // dt/2
		auto coeff = dt_half * qom; // (dt/2)*(q/m)

		/* geometry */
		auto eq = equilibrium::make_efit<T>(EFIT_FILE);
		auto b0 = eq->get_characteristic_field(tid);
		std::cout<<"Thread "<<tid<<" B0: "<<b0->evaluate().at(0)<<'\n';

		/* vector wrappers */
		auto pos = graph::vector(x ,y ,z );
		auto u_vec = graph::vector(ux,uy,uz);

		auto b_vec = eq->get_magnetic_field(pos->get_x(),
		pos->get_y(),
		pos->get_z());
		/* fields */
		auto b_vec = eq->get_magnetic_field(
		pos->get_x(), pos->get_y(), pos->get_z());
		auto e_vec = eq->get_electric_field(
		pos->get_x(), pos->get_y(), pos->get_z());

		/* Boris pusher algebra (symbolic) */
		auto t_vec = dt * half * b_vec;
		/* ── Boris pusher with E ─────────────────────────────────────── */
		auto u_minus = u_vec + coeff * e_vec; // first E-kick
		auto t_vec = coeff * b_vec; // t = (qB/m)*(dt/2)
		auto t_sq = t_vec->dot(t_vec);
		auto s_vec = (two * t_vec) / (one + t_sq);
		auto u_prime = u_vec;
		auto u_tilde = u_prime + u_prime->cross(t_vec);
		auto u_next = u_prime + u_tilde->cross(s_vec);
		auto pos_next= pos + dt * u_next;
		auto s_vec = (two * t_vec) / (one + t_sq); // s = 2t/(1+\|t\|²)
		auto u_prime = u_minus + u_minus->cross(t_vec);
		auto u_plus = u_minus + u_prime->cross(s_vec); // after B rotation
		auto u_next = u_plus + coeff * e_vec; // second E-kick

		auto pos_next = pos + dt * u_next; // advance position

		/* workflow manager */
		workflow::manager<T,false> w(tid);
		w.add_item({graph::variable_cast(x ),
		graph::variable_cast(y ),
		@@ -152,10 +161,10 @@ void run_korc()
		{u_next->get_y(), graph::variable_cast(uy)},
		{u_next->get_z(), graph::variable_cast(uz)}
		},
		"boris");
		"boris_EB");
		w.compile();

		/* set up output file */
		/* output file */
		std::ostringstream fn; fn<<"korc_"<<tid<<".nc";
		output::result_file of(fn.str(), Nloc);
		output::data_set<T> ds(of);
		@@ -167,7 +176,7 @@ void run_korc()
		ds.template create_variable<false>(of,"uz", uz, w.get_context());
		of.end_define_mode();

		/* ------------- one persistent writer thread ------------- */
		/* writer thread */
		std::queue<size_t> q;
		std::mutex mtx;
		std::condition_variable cv;
		@@ -191,7 +200,7 @@ void run_korc()
		const timeing::measure_diagnostic t_run("Run T"+std::to_string(tid));
		w.pre_run();

		// MAIN LOOP
		// ── MAIN LOOP ──────────────────────────────────────────────────
		for(size_t s=0; s<NSTEPS; ++s)
		{
		w.wait();
		@@ -220,6 +229,7 @@ void run_korc()
		t_total.print();
		}

		// ─────────────────────────────────────────────────────────────────────────────
		int main(int argc, const char* argv[])
		{
		START_GPU