particle_in_cell.hpp

Go to the documentation of this file.

//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
 
#ifndef particle_in_cell_h
#define particle_in_cell_h
 
#include <numbers>
#include <bit>
 
#include "piecewise.hpp"
#include "workflow.hpp"
#include "random.hpp"
#include "logical.hpp"
 
namespace graph {
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    class apply_u_node final : public no_derivative<T, false,
                                                    n_branch_node<7, T, false>,
                                                    7> {
    private:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string(std::array<shared_leaf<T>, 7> branches) {
            std::string s = "apply_u";
            for (auto &b : branches) {
                s += jit::format_to_string(reinterpret_cast<size_t> (b.get()));
            }
            return s;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static void func(backend::buffer<T> &x,
                         const backend::buffer<T> &i,
                         const backend::buffer<T> &rand,
                         const backend::buffer<T> &mof,
                         const backend::buffer<T> &tbnu_e_dt,
                         const backend::buffer<T> &A,
                         const backend::buffer<T> &B) {
            const size_t size = x.size();
            assert(size == i.size()         &&
                   size == rand.size()      &&
                   size == mof.size()       &&
                   size == tbnu_e_dt.size() &&
                   size == A.size()         &&
                   size == B.size()         &&
                   "Expected all arguments to have the same length.");
 
            for (size_t j = 0; j < size; j++) {
                const T mof_j = mof[j];
                const T a_j = A[j];
                const T b_j = B[j];
                T temp_x = x[j];
                const T tbnu_e_dt_j = tbnu_e_dt[j];
                if constexpr (std::same_as<T, float>) {
                    uint32_t rand_j = std::bit_cast<uint32_t> (rand[j]);
                    for (uint8_t k = 0, ke = i[j]; k < ke; k++, rand_j >>= 1) {
                        const T E0 = mof_j*temp_x;
                        const int8_t rm = 4*(rand_j & 1) - 2;
                        const T C = rm*std::sqrt(tbnu_e_dt_j*E0);
                        temp_x = (E0*a_j + b_j + C)/mof_j;
                    }
                } else {
                    uint64_t rand_j = std::bit_cast<uint64_t> (rand[j]);
                    for (uint8_t k = 0, ke = i[j]; k < ke; k++, rand_j >>= 1) {
                        const T E0 = mof_j*temp_x;
                        const int8_t rm = 4*(rand_j & 1) - 2;
                        const T C = rm*std::sqrt(tbnu_e_dt_j*E0);
                        temp_x = (E0*a_j + b_j + C)/mof_j;
                    }
                }
                x[j] = temp_x;
            }
        }
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        apply_u_node(shared_leaf<T> x,
                     shared_leaf<T> i,
                     shared_leaf<T> rand,
                     shared_leaf<T> mof,
                     shared_leaf<T> tbnu_e_dt,
                     shared_leaf<T> A,
                     shared_leaf<T> B) :
        no_derivative<T, false,
                      n_branch_node<7, T, false>, 7> ({x, i, rand, mof, tbnu_e_dt, A, B},
                                                      apply_u_node::to_string({x, i, rand, mof, tbnu_e_dt, A, B})) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            backend::buffer<T> x = this->branches[0]->evaluate();
            const backend::buffer<T> i = this->branches[1]->evaluate();
            const backend::buffer<T> rand = this->branches[2]->evaluate();
            const backend::buffer<T> mof = this->branches[3]->evaluate();
            const backend::buffer<T> tbnu_e_dt = this->branches[4]->evaluate();
            const backend::buffer<T> A = this->branches[5]->evaluate();
            const backend::buffer<T> B = this->branches[6]->evaluate();
 
            apply_u_node::func(x, i, rand, mof, tbnu_e_dt, A, B);
            return x;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> reduce() {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!pre_funcs.contains("apply_u")) {
                pre_funcs.insert("apply_u");
 
                jit::add_type<T> (stream);
                stream << " apply_u(const ";
                jit::add_type<T> (stream);
                stream << " x, const uint8_t i, uint32_t rand, const ";
                jit::add_type<T> (stream);
                stream << " mof, const ";
                jit::add_type<T> (stream);
                stream << " tbnu_e_dt, const ";
                jit::add_type<T> (stream);
                stream << " A, const ";
                jit::add_type<T> (stream);
                stream << " B) {" << std::endl
                       << "    ";
                jit::add_type<T> (stream);
                stream << " temp_x = x;" << std::endl
                       << "    for (uint8_t j = 0; j < i; j++, rand >>= 1) {" << std::endl
                       << "        const ";
                jit::add_type<T> (stream);
                stream << " E0 = mof*temp_x;" << std::endl
                       << "        const uint8_t rm = 4*(rand & 1) - 2;" << std::endl
                       << "        const ";
                jit::add_type<T> (stream);
                stream << " C = rm*sqrt(tbnu_e_dt*E0);" << std::endl
                       << "        temp_x = (E0*A + B + C)/mof;" << std::endl
                       << "    }" << std::endl
                       << "    return temp_x;" << std::endl
                       << "}" << std::endl;
            }
 
            if (!visited.contains(this)) {
                for (auto &b : this->branches) {
                    b->compile_preamble(stream, registers,
                                        visited, usage,
                                        textures1d, textures2d,
                                        pre_funcs, avail_const_mem);
                }
 
                visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> compile(std::ostringstream &stream,
                                       jit::register_map &registers,
                                       const jit::register_map &thread_mem,
                                       const jit::register_usage &usage) {
            if (registers.find(this) == registers.end()) {
                auto x = this->branches[0]->compile(stream, registers, thread_mem, usage);
                auto i = this->branches[1]->compile(stream, registers, thread_mem, usage);
                auto rand = this->branches[2]->compile(stream, registers, thread_mem, usage);
                auto mof = this->branches[3]->compile(stream, registers, thread_mem, usage);
                auto tbnu_e_dt = this->branches[4]->compile(stream, registers, thread_mem, usage);
                auto A = this->branches[5]->compile(stream, registers, thread_mem, usage);
                auto B = this->branches[6]->compile(stream, registers, thread_mem, usage);
 
                registers[this] = jit::to_string('r', this);
                stream << "        const ";
                jit::add_type<T> (stream);
                stream << " " << registers[this] << " = apply_u("
                       << registers[x.get()] << ", "
                       << registers[i.get()] << ", ";
                stream << registers[rand.get()] << ", "
                       << registers[mof.get()] << ", "
                       << registers[tbnu_e_dt.get()] << ", "
                       << registers[A.get()] << ", "
                       << registers[B.get()] << ")";
                this->endline(stream, usage);
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_match(shared_leaf<T> x) {
            if (this == x.get()) {
                return true;
            }
 
            auto x_cast = apply_u_cast(x);
            bool temp;
            if (x_cast.get()) {
                temp = this->branches[0]->is_match(x_cast->get_arg(0));
                for (size_t i = 1; i < 7 && temp; i++) {
                    temp = temp && this->branches[i]->is_match(x_cast->get_arg(i));
                }
            }
 
            return temp;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << "\\apply_u{\\left(";
            this->branches[0]->to_latex();
            for (uint8_t i = 1; i < 7; i++) {
                std::cout << ", ";
                this->branches[i]->to_latex();
            }
            std::cout << "\\right)}";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> remove_pseudo() {
            if (this->has_pseudo()) {
                return apply_u(this->branches[0]->remove_pseudo(),
                               this->branches[1]->remove_pseudo(),
                               this->branches[2]->remove_pseudo(),
                               this->branches[3]->remove_pseudo(),
                               this->branches[4]->remove_pseudo(),
                               this->branches[5]->remove_pseudo(),
                               this->branches[6]->remove_pseudo());
            }
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> to_vizgraph(std::stringstream &stream,
                                           jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('r', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"apply_u\", shape = oval, style = filled, fillcolor = blue, fontcolor = white];" << std::endl;
 
                for (auto &b : this->branches) {
                    auto temp = b->to_vizgraph(stream, registers);
                    stream << "    " << name << " -- " << registers[temp.get()] << ";" << std::endl;
                }
            }
 
            return this->shared_from_this();
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    shared_leaf<T> apply_u(shared_leaf<T> x,
                           shared_leaf<T> i,
                           shared_leaf<T> rand,
                           shared_leaf<T> mof,
                           shared_leaf<T> tbnu_e_dt,
                           shared_leaf<T> A,
                           shared_leaf<T> B) {
        auto temp = std::make_shared<apply_u_node<T>> (x, i, rand, mof,
                                                       tbnu_e_dt, A, B)->reduce();
//  Test for hash collisions.
        for (size_t i = temp->get_hash();
             i < std::numeric_limits<size_t>::max(); i++) {
            if (leaf_node<T>::caches.nodes.find(i) ==
                leaf_node<T>::caches.nodes.end()) {
                leaf_node<T>::caches.nodes[i] = temp;
                return temp;
            } else if (temp->is_match(leaf_node<T>::caches.nodes[i])) {
                return leaf_node<T>::caches.nodes[i];
            }
        }
#if defined(__clang__) || defined(__GNUC__)
        __builtin_unreachable();
#else
        assert(false && "Should never reach.");
#endif
    }
 
    template<std::floating_point T>
    using shared_apply_u = std::shared_ptr<apply_u_node<T>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    shared_apply_u<T> apply_u_cast(shared_leaf<T> x) {
        return std::dynamic_pointer_cast<apply_u_node<T>> (x);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    class apply_xi_node final : public no_derivative<T, false,
                                                     n_branch_node<4, T, false>,
                                                     4> {
    private:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string to_string(std::array<shared_leaf<T>, 4> branches) {
            std::string s = "apply_xi";
            for (auto &b : branches) {
                s += jit::format_to_string(reinterpret_cast<size_t> (b.get()));
            }
            return s;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static void func(backend::buffer<T> &x,
                         const backend::buffer<T> &i,
                         const backend::buffer<T> &rand,
                         const backend::buffer<T> &nu_D_dt) {
            const size_t size = x.size();
            assert(size == i.size()       &&
                   size == rand.size()    &&
                   size == nu_D_dt.size() &&
                   "Expected all arguments to have the same length.");
 
            for (size_t j = 0; j < size; j++) {
                const T nu_D_dt_j = nu_D_dt[j];
                T temp_x = x[j];
                if constexpr (std::same_as<T, float>) {
                    uint32_t rand_j = std::bit_cast<uint32_t> (rand[j]);
                    for (uint8_t k = 0, ke = i[j]; k < ke; k++, rand_j >>= 1) {
                        const T A = -temp_x*nu_D_dt_j;
                        const int8_t rm = 2*(rand_j & 1) - 1;
                        const T C = rm*std::sqrt((1 - temp_x*temp_x)*nu_D_dt_j);
                        temp_x += A + C;
                    }
                } else {
                    uint64_t rand_j = std::bit_cast<uint64_t> (rand[j]);
                    for (uint8_t k = 0, ke = i[j]; k < ke; k++, rand_j >>= 1) {
                        const T A = -temp_x*nu_D_dt_j;
                        const int8_t rm = 2*(rand_j & 1) - 1;
                        const T C = rm*std::sqrt((1 - temp_x*temp_x)*nu_D_dt_j);
                        temp_x += A + C;
                    }
                }
                x[j] = temp_x;
            }
        }
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        apply_xi_node(shared_leaf<T> x,
                      shared_leaf<T> i,
                      shared_leaf<T> rand,
                      shared_leaf<T> nu_D_dt) :
        no_derivative<T, false,
                      n_branch_node<4, T, false>, 4> ({x, i, rand, nu_D_dt},
                                                      apply_xi_node::to_string({x, i, rand, nu_D_dt})) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual backend::buffer<T> evaluate() {
            backend::buffer<T> x = this->branches[0]->evaluate();
            const backend::buffer<T> i = this->branches[1]->evaluate();
            const backend::buffer<T> rand = this->branches[2]->evaluate();
            const backend::buffer<T> nu_D_dt = this->branches[3]->evaluate();
 
            apply_xi_node::func(x, i, rand, nu_D_dt);
            return x;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> reduce() {
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void compile_preamble(std::ostringstream &stream,
                                      jit::register_map &registers,
                                      jit::visiter_map &visited,
                                      jit::register_usage &usage,
                                      jit::texture1d_list &textures1d,
                                      jit::texture2d_list &textures2d,
                                      jit::preamble_map &pre_funcs,
                                      int &avail_const_mem) {
            if (!pre_funcs.contains("apply_xi")) {
                pre_funcs.insert("apply_xi");
 
                jit::add_type<T> (stream);
                stream << " apply_xi(const ";
                jit::add_type<T> (stream);
                stream << " x, const uint8_t i, uint32_t rand, const ";
                jit::add_type<T> (stream);
                stream << " nu_D_dt) {" << std::endl
                       << "    ";
                jit::add_type<T> (stream);
                stream << " temp_x = x;" << std::endl
                       << "    for (uint8_t j = 0; j < i; j++, rand >>= 1) {" << std::endl
                       << "        const ";
                jit::add_type<T> (stream);
                stream << " A = -temp_x*nu_D_dt;" << std::endl
                       << "        const uint8_t rm = 2*(rand & 1) - 1;" << std::endl
                       << "        const ";
                jit::add_type<T> (stream);
                stream << " C = rm*sqrt((1 - temp_x*temp_x)*nu_D_dt);" << std::endl
                       << "        temp_x += A + C;" << std::endl
                       << "    }" << std::endl
                       << "    return temp_x;" << std::endl
                       << "}" << std::endl;
            }
 
            if (!visited.contains(this)) {
                for (auto &b : this->branches) {
                    b->compile_preamble(stream, registers,
                                        visited, usage,
                                        textures1d, textures2d,
                                        pre_funcs, avail_const_mem);
                }
 
                visited.insert(this);
#ifdef SHOW_USE_COUNT
                usage[this] = 1;
            } else {
                ++usage[this];
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> compile(std::ostringstream &stream,
                                       jit::register_map &registers,
                                       const jit::register_map &thread_mem,
                                       const jit::register_usage &usage) {
            if (registers.find(this) == registers.end()) {
                auto x = this->branches[0]->compile(stream, registers, thread_mem, usage);
                auto i = this->branches[1]->compile(stream, registers, thread_mem, usage);
                auto rand = this->branches[2]->compile(stream, registers, thread_mem, usage);
                auto nu_D_dt = this->branches[3]->compile(stream, registers, thread_mem, usage);
 
                registers[this] = jit::to_string('r', this);
                stream << "        const ";
                jit::add_type<T> (stream);
                stream << " " << registers[this] << " = apply_xi("
                       << registers[x.get()] << ", "
                       << registers[i.get()] << ", ";
                stream << registers[rand.get()] << ", "
                       << registers[nu_D_dt.get()] << ")";
                this->endline(stream, usage);
            }
 
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual bool is_match(shared_leaf<T> x) {
            if (this == x.get()) {
                return true;
            }
 
            auto x_cast = apply_xi_cast(x);
            bool temp;
            if (x_cast.get()) {
                temp = this->branches[0]->is_match(x_cast->get_arg(0));
                for (size_t i = 1; i < 4 && temp; i++) {
                    temp = temp && this->branches[i]->is_match(x_cast->get_arg(i));
                }
            }
 
            return temp;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual void to_latex() const {
            std::cout << "\\apply_xi{\\left(";
            this->branches[0]->to_latex();
            for (uint8_t i = 1; i < 4; i++) {
                std::cout << ", ";
                this->branches[i]->to_latex();
            }
            std::cout << "\\right)}";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> remove_pseudo() {
            if (this->has_pseudo()) {
                return apply_xi(this->branches[0]->remove_pseudo(),
                                this->branches[1]->remove_pseudo(),
                                this->branches[2]->remove_pseudo(),
                                this->branches[3]->remove_pseudo());
            }
            return this->shared_from_this();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual shared_leaf<T> to_vizgraph(std::stringstream &stream,
                                           jit::register_map &registers) {
            if (registers.find(this) == registers.end()) {
                const std::string name = jit::to_string('r', this);
                registers[this] = name;
                stream << "    " << name
                       << " [label = \"apply_xi\", shape = oval, style = filled, fillcolor = blue, fontcolor = white];" << std::endl;
 
                for (auto &b : this->branches) {
                    auto temp = b->to_vizgraph(stream, registers);
                    stream << "    " << name << " -- " << registers[temp.get()] << ";" << std::endl;
                }
            }
 
            return this->shared_from_this();
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    shared_leaf<T> apply_xi(shared_leaf<T> x,
                            shared_leaf<T> i,
                            shared_leaf<T> rand,
                            shared_leaf<T> nu_D_dt) {
        auto temp = std::make_shared<apply_xi_node<T>> (x, i, rand,
                                                        nu_D_dt)->reduce();
//  Test for hash collisions.
        for (size_t i = temp->get_hash();
             i < std::numeric_limits<size_t>::max(); i++) {
            if (leaf_node<T>::caches.nodes.find(i) ==
                leaf_node<T>::caches.nodes.end()) {
                leaf_node<T>::caches.nodes[i] = temp;
                return temp;
            } else if (temp->is_match(leaf_node<T>::caches.nodes[i])) {
                return leaf_node<T>::caches.nodes[i];
            }
        }
#if defined(__clang__) || defined(__GNUC__)
        __builtin_unreachable();
#else
        assert(false && "Should never reach.");
#endif
    }
 
    template<std::floating_point T>
    using shared_apply_xi = std::shared_ptr<apply_xi_node<T>>;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    shared_apply_xi<T> apply_xi_cast(shared_leaf<T> x) {
        return std::dynamic_pointer_cast<apply_xi_node<T>> (x);
    }
}
 
namespace pic {
// FIXME: This should be in a separate file of physics constants.
    template<std::floating_point T>
    constexpr T c = static_cast<T> (299792458.0);
    template<std::floating_point T>
    constexpr T epsilon0 = static_cast<T> (8.8541878188E-12);
    template<std::floating_point T>
    constexpr T q = static_cast<T> (1.602176634E-19);
    template<std::floating_point T>
    constexpr T m_hydrogen = static_cast<T> (1.67362192595E-27);
    template<std::floating_point T>
    constexpr T m_atomic = static_cast<T> (1.66053906892E-27);
    template<std::floating_point T>
    constexpr T m_electron = static_cast<T> (9.1093837139E-31);
    template<std::floating_point T>
    constexpr T kb = static_cast<T> (1.380649E-23);
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    class characteristics {
    private:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T make_m(const std::vector<T> &ion_masses) {
            T total_m = static_cast<T> (0);
            for (const T &mass : ion_masses) {
                total_m += mass;
            }
            return total_m/ion_masses.size();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T make_q(const std::vector<uint8_t> &ion_zs) {
            T total_q = static_cast<T> (0);
            for (const uint8_t &z : ion_zs) {
                total_q += z*pic::q<T>;
            }
            return total_q/ion_zs.size();
        }
 
    public:
        const T m;
        const T q;
        const T ne;
        const T wpe;
        const T t;
        const T l;
        const T v;
        const T te;
        const T efield;
        const T bfield;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        characteristics(const std::vector<T> &ion_masses,
                        const std::vector<uint8_t> &ion_zs,
                        const T ne) :
        m(make_m(ion_masses)), q(make_q(ion_zs)), ne(ne),
        wpe(std::sqrt(ne*q*q/(m*epsilon0<T>))),
        t(1/wpe), l(c<T>/wpe), v(c<T>), te(m*v*v/kb<T>), efield(m*c<T>/(q*t)),
        bfield(efield/c<T>) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T get_volume() const {
            return l*l*l;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    class parameters {
    public:
        const T b0;
        const T a0;
        const size_t filter_iterations;
        const T smoothing;
        const T dt;
        const T t_para;
        const T t_perp;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        parameters(const T b0, const T r1, const T r2,
                   const size_t filter_iterations,
                   const T smoothing, const T dt,
                   const T t_para, const T t_perp,
                   const characteristics<T> &norms) :
        b0(b0/norms.bfield),
        a0(std::numbers::pi_v<T>*(r2*r2 - r1*r1)/(norms.l*norms.l)),
        filter_iterations(filter_iterations), smoothing(smoothing),
        dt(dt/norms.t), t_para(t_para), t_perp(t_perp) {}
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    class ion {
    public:
        const uint8_t z;
        const T charge;
        const T mass;
        graph::shared_leaf<T> x;
        graph::shared_leaf<T> v_para;
        graph::shared_leaf<T> v_perp;
        std::array<graph::shared_leaf<T>, 3> weights;
        const T num_real;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        ion(const T mass,
            const uint8_t z,
            const size_t num_ions,
            const T num_real,
            const characteristics<T> &norms) :
        z(z), charge(z*pic::q<T>/norms.q),
        mass(mass), num_real(num_real),
        x(graph::variable<T> (num_ions, "x")),
        v_para(graph::variable<T> (num_ions, "v_{||}")),
        v_perp(graph::variable<T> (num_ions, "v_{\\perp}")),
        weights({
            graph::variable<T> (num_ions, "w_{0}"),
            graph::variable<T> (num_ions, "w_{1}"),
            graph::variable<T> (num_ions, "w_{2}")
        }) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_variable<T> get_x() const {
            return graph::variable_cast(x);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t size() const {
            return get_x()->size();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T *x_data() const {
            return get_x()->data();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_variable<T> get_v_para() const {
            return graph::variable_cast(v_para);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T *v_para_data() const {
            return get_v_para()->data();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_variable<T> get_v_perp() const {
            return graph::variable_cast(v_perp);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T *v_perp_data() const {
            return graph::variable_cast(v_perp)->data();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        T super_to_real() const {
            return num_real/size();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void define_variables(const output::result_file &file,
                              output::data_set<T> &data,
                              workflow::manager<T> &work,
                              const std::string tag) {
            data.create_variable(file, "x_" + tag, x,  work.get_context());
            data.create_variable(file, "vpara_" + tag, v_para,
                                 work.get_context());
            data.create_variable(file, "vperp_" + tag, v_perp,
                                 work.get_context());
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<T> build_profile(std::function<graph::shared_leaf<T>(graph::shared_leaf<T>)> func) const {
            return func(x);
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    class mesh {
    private:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        template<size_t I=0>
        graph::shared_leaf<T> build_y_index(graph::shared_leaf<T> x,
                                            const T scale,
                                            const size_t iterations=0) const {
            auto low = iterations ? build_y_index<I> (x - dx, scale, iterations - 1) :
            graph::index_1D(y[I], x, dx, xmin + dx);
            auto center = graph::index_1D(y[I], x, dx, xmin);
            auto high = iterations ? build_y_index<I> (x + dx, scale, iterations - 1) :
            graph::index_1D(y[I], x, dx, xmin - dx);
 
            const T center_w = static_cast<T> (0.5);
            const T side_w = static_cast<T> (0.25);
 
            auto b = center_w*center + side_w*low + side_w*high;
            return (1 - scale)*center + scale*b;
        }
 
    public:
        const T xmin;
        const T xmax;
        const T dx;
        std::array<graph::shared_leaf<T>, 4> y;
        enum offset {
            low,
            center,
            high
        };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        mesh(const T x_min,
             const T x_max,
             const size_t num,
             const characteristics<T> &norms) :
        y({
            graph::variable<T> (num, "y^{0}_{m}"),
            graph::variable<T> (num, "y^{1}_{m}"),
            graph::variable<T> (num, "y^{2}_{m}"),
            graph::variable<T> (num, "y^{3}_{m}")
        }), xmin(x_min/norms.l), xmax(x_max/norms.l),
        dx((xmax - xmin)/(num - 1)) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<T> build_x_index(graph::shared_leaf<T> x) const {
            return dx*graph::argument(x, dx, xmin, size()) + xmin;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<T> build_i_index(graph::shared_leaf<T> x) const {
            return graph::argument(x, dx, xmin, size());
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        template<size_t I=0, offset O=center>
        graph::shared_leaf<T> build_y_index(graph::shared_leaf<T> x,
                                            const parameters<T> &params) const {
            if constexpr (O == low) {
                return build_y_index<I> (x - dx, params.smoothing,
                                         params.filter_iterations - 1);
            } else if constexpr (O == center) {
                return build_y_index<I> (x, params.smoothing,
                                         params.filter_iterations - 1);
            } else {
                return build_y_index<I> (x + dx, params.smoothing,
                                         params.filter_iterations - 1);
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        template<size_t I=0, offset O=center>
        graph::shared_leaf<T> build_dydx_index(graph::shared_leaf<T> x,
                                               const parameters<T> &params) const {
            const T two = 2;
            if constexpr (O == low) {
                auto low = build_y_index<I> (x - two*dx, params.smoothing,
                                             params.filter_iterations);
                auto high = build_y_index<I> (x, params.smoothing,
                                              params.filter_iterations);
                return (high - low)/two;
            } else if constexpr (O == center) {
                auto low = build_y_index<I> (x - dx, params.smoothing,
                                             params.filter_iterations);
                auto high = build_y_index<I> (x + dx, params.smoothing,
                                              params.filter_iterations);
                return (high - low)/two;
            } else {
                auto low = build_y_index<I> (x, params.smoothing,
                                             params.filter_iterations);
                auto high = build_y_index<I> (x + two*dx, params.smoothing,
                                              params.filter_iterations);
                return (high - low)/two;
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::array<graph::shared_leaf<T>, 3> build_mesh_solve(const ion<T> &ion) const {
            auto weights = build_weights(ion.x);
            auto sum_low = graph::atomic_accumulate_1D(y[0], ion.x - dx,
                                                       dx, xmin, weights[0]);
            auto sum = graph::atomic_accumulate_1D(y[0], ion.x,
                                                   dx, xmin, weights[1]);
            auto sum_high = graph::atomic_accumulate_1D(y[0], ion.x + dx,
                                                        dx, xmin, weights[2]);
            return {sum_low, sum, sum_high};
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        size_t size() const {
            return graph::variable_cast(y[0])->size();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        template<size_t I=0>
        T *data() const {
            return graph::variable_cast(y[I])->data();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void define_variables(const output::result_file &file,
                              output::data_set<T> &data,
                              workflow::manager<T> &work) {
            data.create_variable(file, "y_0", y[0],  work.get_context());
            data.create_variable(file, "y_1", y[1],  work.get_context());
            data.create_variable(file, "y_2", y[2],  work.get_context());
            data.create_variable(file, "y_3", y[3],  work.get_context());
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::array<graph::shared_leaf<T>, 3> build_weights(graph::shared_leaf<T> x) const {
            auto x_off = build_x_index(x) - x;
            auto xnorm1 = static_cast<T> (1.5) + (x_off - dx)/dx;
            auto xnorm2 = x_off/dx;
            auto xnorm3 = static_cast<T> (1.5) - (x_off + dx)/dx;
 
            auto w0 = static_cast<T> (0.5)*xnorm1*xnorm1;
            auto w1 = static_cast<T> (0.75) - xnorm2*xnorm2;
            auto w2 = static_cast<T> (0.5)*xnorm3*xnorm3;
 
            return {w0, w1, w2};
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<T> build_interpolation(graph::shared_leaf<T> x) const {
            auto weights = build_weights(x);
 
            auto ymesh0 = graph::index_1D(y[0], x - dx, dx, xmin);
            auto ymesh1 = graph::index_1D(y[0], x,      dx, xmin);
            auto ymesh2 = graph::index_1D(y[0], x + dx, dx, xmin);
 
            return weights[0]*ymesh0 + weights[1]*ymesh1 + weights[2]*ymesh2;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    std::array<graph::shared_leaf<T>, 3> cartesian_to_sphereical(graph::shared_leaf<T> x,
                                                                 graph::shared_leaf<T> y) {
        auto w = graph::hypot(x, y);
        return {
            w, x/w,
            graph::none<T> ()*graph::copysign(static_cast<T> (1), y)
        };
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    std::array<graph::shared_leaf<T>, 2> sphereical_to_cartesian(graph::shared_leaf<T> w,
                                                                 graph::shared_leaf<T> xi,
                                                                 graph::shared_leaf<T> sinphi) {
        return {
            w*xi,
            graph::none<T> ()*w*graph::sqrt(static_cast<T> (1) - xi*xi)*sinphi
        };
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    enum model {
        hinton,
        chen
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T, model M>
    graph::shared_leaf<T> build_colision_rate(graph::shared_leaf<T> xab,
                                              const T mass_a,
                                              const T mass_b,
                                              graph::shared_leaf<T> gb,
                                              graph::shared_leaf<T> nuab0,
                                              graph::shared_leaf<T> erfp_xab) {
        auto mass_ratio = static_cast<T> (2)*mass_a/mass_b*gb;
        auto nu = nuab0/xab;
        if constexpr (M == model::hinton) {
            return nu*(mass_ratio - erfp_xab/xab);
        } else {
            return nu*mass_ratio;
        }
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T, model M>
    std::array<graph::shared_leaf<T>, 2> build_ion_ion_collision(const ion<T> &ion_a,
                                                                 const ion<T> &ion_b,
                                                                 const mesh<T> &mesh,
                                                                 const characteristics<T> &norms,
                                                                 const parameters<T> &params,
                                                                 graph::shared_leaf<T> total_density,
                                                                 graph::shared_leaf<T> total_flux,
                                                                 const graph::shared_random_state<T> state) {
        const T dt = params.dt*norms.t;
 
        const T mass_b = ion_b.mass*norms.m;
        const uint8_t zb2 = ion_b.z*ion_b.z;
 
        auto nb = build_density(ion_a.x, ion_b, mesh, norms, params)/norms.get_volume();
        auto tpara = ion_a.build_profile([](graph::shared_leaf<T> x) -> graph::shared_leaf<T> {
            return graph::one<T> ();
        });
        auto tperp = ion_a.build_profile([](graph::shared_leaf<T> x) -> graph::shared_leaf<T> {
            return graph::one<T> ();
        });
        auto tb = static_cast<T> (0.5)*(tpara + tperp)*norms.te*kb<T>/q<T>;
        auto nv = ion_a.build_profile([](graph::shared_leaf<T> x) -> graph::shared_leaf<T> {
            return graph::one<T> ();
        })*norms.v/norms.get_volume();
        auto uxb = nv/nb;
 
        total_density = total_density + nb;
        total_flux = total_flux + nv;
 
        return build_common_collision<T, M> (ion_a, norms, uxb, tb, nb, mass_b,
                                             zb2, dt, state);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T, model M>
    std::array<graph::shared_leaf<T>, 2> build_ion_electron_collision(const ion<T> &ion_a,
                                                                        const mesh<T> &mesh,
                                                                        const characteristics<T> &norms,
                                                                        const parameters<T> &params,
                                                                        graph::shared_leaf<T> total_density,
                                                                        graph::shared_leaf<T> total_flux,
                                                                        const graph::shared_random_state<T> state) {
        const T dt = params.dt*norms.t;
 
        const T mass_b = m_electron<T>;
        const uint8_t zb2 = 1;
 
        auto nb = total_density;
        auto tb = ion_a.build_profile([](graph::shared_leaf<T> x) -> graph::shared_leaf<T> {
            return graph::one<T> ();
        })*norms.te;
 
        auto uxb = total_flux/total_density;
 
        return build_common_collision<T, M> (ion_a, norms, uxb, tb, nb, mass_b,
                                             zb2, dt, state);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T, model M>
    std::array<graph::shared_leaf<T>, 2> build_common_collision(const ion<T> &ion_a,
                                                                  const characteristics<T> &norms,
                                                                  graph::shared_leaf<T> uxb,
                                                                  graph::shared_leaf<T> tb,
                                                                  graph::shared_leaf<T> nb,
                                                                  const T mass_b,
                                                                  const uint8_t zb2,
                                                                  const T dt,
                                                                  const graph::shared_random_state<T> state) {
        const T mass_a = ion_a.mass*norms.m;
        const uint8_t za2 = ion_a.z*ion_a.z;
 
//  Convert to b frame.
        auto wxa = ion_a.v_para*norms.v - uxb;
        auto wya = ion_a.v_perp*norms.v;
 
//  Convert to sphereical.
        auto sphere = cartesian_to_sphereical(wxa, wya);
 
//  Apply Monte-Carlo collision operator.
        auto wtb = graph::sqrt(static_cast<T> (2)*q<T>*tb/mass_b);
        auto xab = sphere[0]/wtb;
 
        auto erf_xab = graph::erf(xab);
        auto erfp_xab = erf_xab->df(xab);
        auto erfpp_xab = erfp_xab->df(xab);
 
        auto gb = (erf_xab - xab*erfp_xab)/(static_cast<T> (2)*xab*xab);
        auto logA = static_cast<T> (30) - static_cast<T> (0.5)*graph::log(nb/(tb*graph::sqrt(tb)));
        auto nuab0 = nb*q<T>*q<T>*q<T>*q<T>*static_cast<T> (za2*zb2)*logA
                   / (static_cast<T> (2)*std::numbers::pi_v<T>*mass_a*mass_a*epsilon0<T>*epsilon0<T>*wtb*wtb*wtb);
 
//  Velocity Scattering operator.
        auto nu_e_dt = build_colision_rate<T, M> (xab, mass_a, mass_b, gb, nuab0, erfp_xab)*dt;
        auto steps = graph::min(nu_e_dt*2.5, static_cast<T> (32));
 
        nu_e_dt = nu_e_dt*steps;
 
        auto E_nuE_d_nu_E_dE = static_cast<T> (0.5)*((static_cast<T> (3.0)*(xab*erfp_xab - erf_xab) - xab*xab*erfpp_xab)/(erf_xab - xab*erfp_xab));
 
        const T mof = mass_a/(static_cast<T> (2)*q<T>);
        auto A = static_cast<T> (1) - static_cast<T> (2)*nu_e_dt;
        auto tbnu_e_df = tb*nu_e_dt;
        auto B = static_cast<T> (2)*tbnu_e_df*(static_cast<T> (1.5) + E_nuE_d_nu_E_dE);
 
        auto rand1 = graph::random<T> (state);
        auto u_op = graph::apply_u(sphere[0]*sphere[0], steps, rand1,
                                   graph::constant(mof), tbnu_e_df, A, B);
        sphere[0] = graph::sqrt(u_op);
 
        auto nu_D_dt = nuab0*(erf_xab - gb)/(xab*xab*xab)*dt;
        steps = graph::min(nu_D_dt, static_cast<T> (32));
        nu_D_dt = nu_D_dt/steps;
 
        auto rand2 = graph::random<T> (state);
        sphere[1] = graph::apply_xi(sphere[1], steps, rand2, nu_D_dt);
 
        sphere[1] = graph::if_(sphere[1]*sphere[1] > static_cast<T> (1),
                               graph::copysign(static_cast<T> (1), sphere[1]) -
                               sphere[1] % graph::copysign(static_cast<T> (1), sphere[1]),
                               sphere[1]);
 
        auto cart = sphereical_to_cartesian(sphere[0], sphere[1], sphere[2]);
        return {(cart[0] + uxb)/norms.v, cart[1]/norms.v};
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    std::array<graph::shared_leaf<T>,3> build_initialization(const pic::ion<T> &ion,
                                                             const mesh<T> &mesh,
                                                             const characteristics<T> &norms,
                                                             const parameters<T> &params,
                                                             const graph::shared_random_state<T> state) {
//  The mesh is already normalized so position_dist will be a normalized quantity.
        auto position_dist = graph::uniform_random<T> (mesh.xmin, mesh.xmax,
                                                       state);
        auto phi_dist = graph::uniform_random<T> (static_cast<T> (0),
                                                  static_cast<T> (2)*std::numbers::pi_v<T>,
                                                  state);
        auto r_dist = graph::uniform_random<T> (std::numeric_limits<T>::min(),
                                                static_cast<T> (1),
                                                state);
 
        const T vtpara = std::sqrt(2*params.t_para*q<T>/ion.mass);
        auto vpara = vtpara*graph::sqrt(-graph::log(r_dist))
                   * graph::sin(phi_dist);
 
        phi_dist = graph::uniform_random<T> (static_cast<T> (0),
                                             static_cast<T> (2)*std::numbers::pi_v<T>,
                                             state);
        r_dist = graph::uniform_random<T> (std::numeric_limits<T>::min(),
                                           static_cast<T> (1),
                                           state);
 
        const T vtperp = std::sqrt(2*params.t_perp*q<T>/ion.mass);
        auto vperp1 = vtperp*graph::sqrt(-graph::log(r_dist))
                    * graph::cos(phi_dist);
        auto vperp2 = vtperp*graph::sqrt(-graph::log(r_dist))
                    * graph::sin(phi_dist);
        auto vperp = graph::sqrt(vperp1*vperp1 + vperp2*vperp2);
 
        return {position_dist, vpara/norms.v, vperp/norms.v};
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    std::array<graph::shared_leaf<T>, 3> build_reinjection(const ion<T> &ion,
                                                           const mesh<T> &mesh,
                                                           const characteristics<T> &norms,
                                                           const parameters<T> &params,
                                                           const graph::shared_random_state<T> state) {
        auto resampled = build_initialization(ion, mesh, norms, params, state);
        auto is_outside = ion.x <= mesh.xmin || ion.x >= mesh.xmax;
 
        auto reinject_x     = graph::if_(is_outside          ||
                                         graph::isnan(ion.x) ||
                                         graph::isinf(ion.x), resampled[0], ion.x);
        auto reinject_vpara = graph::if_(is_outside          ||
                                         graph::isnan(ion.x) ||
                                         graph::isinf(ion.x), resampled[1], ion.v_para);
        auto reinject_vperp = graph::if_(is_outside          ||
                                         graph::isnan(ion.x) ||
                                         graph::isinf(ion.x), resampled[2], ion.v_perp);
        return {reinject_x, reinject_vpara, reinject_vperp};
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    graph::shared_leaf<T> build_magnetic_field(graph::shared_leaf<T> x,
                                               const characteristics<T> &norms,
                                               const parameters<T> &params) {
        return (x*x*(norms.l*norms.l) + static_cast<T> (0.5))*params.b0;
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T, size_t I=0,
             mesh<T>::offset O=mesh<T>::center>
    graph::shared_leaf<T> build_density(graph::shared_leaf<T> x,
                                        const ion<T> &ion,
                                        const mesh<T> &mesh,
                                        const characteristics<T> &norms,
                                        const parameters<T> &params) {
        auto y = mesh.template build_y_index<I, O> (x, params);
 
//  Compression factor.
        auto cf = build_magnetic_field(x, norms, params)/params.b0;
//  Scale factor.
        const T sf = ion.super_to_real()/(params.a0*mesh.dx);
 
        return static_cast<T> (ion.z)*y*cf*sf;
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T, size_t I=0,
             mesh<T>::offset O=mesh<T>::center>
    graph::shared_leaf<T> build_density_gradient(graph::shared_leaf<T> x,
                                                 const ion<T> &ion,
                                                 const mesh<T> &mesh,
                                                 const characteristics<T> &norms,
                                                 const parameters<T> &params) {
        auto y = mesh.template build_dydx_index<I, O> (x, params);
 
//  Compression factor.
        auto cf = build_magnetic_field(x, norms, params)/params.b0;
//  Scale factor.
        const T sf = ion.super_to_real()/(params.a0*mesh.dx);
 
        return static_cast<T> (ion.z)*y*cf*sf;
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    graph::shared_leaf<T> build_electron_temperature(graph::shared_leaf<T> x,
                                                     const characteristics<T> &norms) {
        return graph::constant<T> (static_cast<T> (2.5)*q<T>/(norms.te*kb<T>));
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T, mesh<T>::offset O=mesh<T>::center>
    graph::shared_leaf<T> build_electric_efield(graph::shared_leaf<T> x,
                                                const ion<T> &ion,
                                                const mesh<T> &mesh,
                                                const characteristics<T> &norms,
                                                const parameters<T> &params) {
        auto n0 = build_density<T, 0, O> (x, ion, mesh, norms, params);
        auto n1 = build_density<T, 1, O> (x, ion, mesh, norms, params);
        auto n2 = build_density<T, 2, O> (x, ion, mesh, norms, params);
        auto n3 = build_density<T, 3, O> (x, ion, mesh, norms, params);
 
        auto dn0dx = build_density_gradient<T, 0, O> (x, ion, mesh, norms, params);
        auto dn1dx = build_density_gradient<T, 1, O> (x, ion, mesh, norms, params);
        auto dn2dx = build_density_gradient<T, 2, O> (x, ion, mesh, norms, params);
        auto dn3dx = build_density_gradient<T, 3, O> (x, ion, mesh, norms, params);
 
        auto n = (n0 + n1 + n2 + n3)/static_cast<T> (4);
        auto dndx = (dn0dx + dn1dx + dn2dx + dn3dx)/static_cast<T> (4);
 
        auto te = build_electron_temperature(x, norms);
        const T scale = norms.q/q<T>;
        auto pressure = te*n*scale;
 
        return graph::none<T> ()/n*(dndx*te*scale + pressure->df(x));
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    graph::shared_leaf<T> build_interpolate_efield(graph::shared_leaf<T> x,
                                                   const ion<T> &ion,
                                                   const mesh<T> &mesh,
                                                   const characteristics<T> &norms,
                                                   const parameters<T> &params) {
        auto weights = mesh.build_weights(x);
 
        auto ymesh0 = build_electric_efield<T, ::pic::mesh<T>::low> (x, ion, mesh, norms, params);
        auto ymesh1 = build_electric_efield<T, ::pic::mesh<T>::center> (x, ion, mesh, norms, params);
        auto ymesh2 = build_electric_efield<T, ::pic::mesh<T>::high> (x, ion, mesh, norms, params);
 
        return weights[0]*ymesh0 + weights[1]*ymesh1 + weights[2]*ymesh2;
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    std::array<graph::shared_leaf<T>, 3> build_F_expressions(const std::array<graph::shared_leaf<T>, 3> z,
                                                             const ion<T> &ion,
                                                             const mesh<T> &mesh,
                                                             const characteristics<T> &norms,
                                                             const parameters<T> &params) {
        auto bfield = build_magnetic_field<T> (z[0], norms, params);
        auto efield = build_interpolate_efield<T> (z[0], ion, mesh, norms, params);
        auto temp = 0.5*z[2]*z[1]*bfield->df(z[0])/bfield;
        return {
            z[1]*params.dt,
            temp*params.dt,
            (ion.charge/ion.mass*norms.m*efield - temp)*params.dt
        };
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<std::floating_point T>
    std::array<graph::shared_leaf<T>, 3> build_rk4_step(const ion<T> &ion,
                                                        const mesh<T> &mesh,
                                                        const characteristics<T> &norms,
                                                        const parameters<T> &params) {
//  Step 1
        std::array<graph::shared_leaf<T>, 3> Z1{ion.x, ion.v_para, ion.v_perp};
        std::array<graph::shared_leaf<T>, 3> dZ1(build_F_expressions<T> (Z1, ion, mesh, norms, params));
 
//  Step 2
        std::array<graph::shared_leaf<T>, 3> Z2{
            Z1[0] + dZ1[0]/static_cast<T> (2),
            Z1[1] + dZ1[1]/static_cast<T> (2),
            Z1[2] + dZ1[2]/static_cast<T> (2)
        };
        std::array<graph::shared_leaf<T>, 3> dZ2(build_F_expressions<T> (Z2, ion, mesh, norms, params));
 
//  Step 3
        std::array<graph::shared_leaf<T>, 3> Z3{
            Z1[0] + dZ2[0]/static_cast<T> (2),
            Z1[1] + dZ2[1]/static_cast<T> (2),
            Z1[2] + dZ2[2]/static_cast<T> (2)
        };
        std::array<graph::shared_leaf<T>, 3> dZ3(build_F_expressions<T> (Z3, ion, mesh, norms, params));
 
//  Step 4
        std::array<graph::shared_leaf<T>, 3> Z4{
            Z1[0] + dZ3[0],
            Z1[1] + dZ3[1],
            Z1[2] + dZ3[2]
        };
        std::array<graph::shared_leaf<T>, 3> dZ4(build_F_expressions<T> (Z4, ion, mesh, norms, params));
 
//  Rk4 Solution
        return {
            Z1[0] + (dZ1[0] + static_cast<T> (2)*(dZ2[0] + dZ3[0]) + dZ4[0])/static_cast<T> (6),
            Z1[1] + (dZ1[1] + static_cast<T> (2)*(dZ2[1] + dZ3[1]) + dZ4[1])/static_cast<T> (6),
            Z1[2] + (dZ1[2] + static_cast<T> (2)*(dZ2[2] + dZ3[2]) + dZ4[2])/static_cast<T> (6)
        };
    }
}
 
#endif /* particle_in_cell_h */