dispersion.hpp

Go to the documentation of this file.

//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
 
#ifndef dispersion_h
#define dispersion_h
 
#include "newton.hpp"
#include "equilibrium.hpp"
 
namespace dispersion {
//******************************************************************************
//  Z Function interface.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class z_function {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        Z(graph::shared_leaf<T, SAFE_MATH> zeta) = 0;
 
        typedef T base;
        static constexpr bool safe_math = SAFE_MATH;
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::complex_scalar T, bool SAFE_MATH=false>
    class z_power_series final : public z_function<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        Z(graph::shared_leaf<T, SAFE_MATH> zeta) {
            auto zeta2 = zeta*zeta;
            auto zeta4 = zeta2*zeta2;
            auto zeta6 = zeta4*zeta2;
            return graph::i<T>*std::sqrt(M_PI)/graph::exp(zeta2) -
                   2.0*(1.0 - 2.0/3.0*zeta2
                            + 4.0/15.0*zeta4
                            - 8.0/105.0*zeta6)*zeta;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::complex_scalar T, bool SAFE_MATH=false>
    class z_erfi final : public z_function<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        Z(graph::shared_leaf<T, SAFE_MATH> zeta) {
            return -std::sqrt(M_PI)*graph::exp(-zeta*zeta)*(graph::erfi(zeta) -
                                                            graph::i<T>);
        }
    };
 
    template<class Z>
    concept z_func = std::is_base_of<z_function<typename Z::base, Z::safe_math>, Z>::value;
 
//******************************************************************************
//  Common physics expressions.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    static constexpr graph::shared_leaf<T, SAFE_MATH>
    build_plasma_frequency(graph::shared_leaf<T, SAFE_MATH> n,
                           const T q,
                           const T m,
                           const T c,
                           const T epsilon0) {
        return n*q*q/(epsilon0*m*c*c);
    }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    static constexpr graph::shared_leaf<T, SAFE_MATH>
    build_cyclotron_frequency(const T q,
                              graph::shared_leaf<T, SAFE_MATH> b,
                              const T m,
                              const T c) {
        return q*b/(m*c);
    }
 
//******************************************************************************
//  Dispersion function.
//******************************************************************************
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class dispersion_function {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) = 0;
 
        typedef T base;
        static constexpr bool safe_math = SAFE_MATH;
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class stiff final : public dispersion_function<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
            return (1.0E3*(x - graph::exp(-t)) - graph::exp(-t))*k_vec->get_x() + w;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class simple final : public dispersion_function<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
            const T c = 1.0;
 
            auto npar2 = k_vec->get_z()*k_vec->get_z()*c*c/(w*w);
            auto nperp2 = (k_vec->get_x()*k_vec->get_x() +
                           k_vec->get_y()*k_vec->get_y())*c*c/(w*w);
            return npar2 + nperp2 - c;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class physics : public dispersion_function<T, SAFE_MATH> {
    protected:
//  Define some common constants.
        const T epsilon0 = 8.8541878138E-12;
        const T mu0 = M_PI*4.0E-7;
        const T q = 1.602176634E-19;
        const T me = 9.1093837015E-31;
        const T c = static_cast<T> (1.0)/std::sqrt(epsilon0*mu0);
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class bohm_gross final : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
 
//  Equilibrium quantities.
            auto ne = eq->get_electron_density(x, y, z);
            auto wpe2 = build_plasma_frequency(ne, physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
            auto te = eq->get_electron_temperature(x, y, z);
//  2*1.602176634E-19 to convert eV to J.
 
            auto vterm2 = static_cast<T> (2.0)*physics<T, SAFE_MATH>::q*te
                        / (physics<T, SAFE_MATH>::me *
                           physics<T, SAFE_MATH>::c *
                           physics<T, SAFE_MATH>::c);
 
//  Wave numbers should be parallel to B if there is a magnetic field. Otherwise
//  B should be zero.
            auto b_vec = eq->get_magnetic_field(x, y, z);
            graph::shared_leaf<T, SAFE_MATH> kpara2;
            if (b_vec->length()->is_match(graph::zero<T, SAFE_MATH> ())) {
                kpara2 = k_vec->dot(k_vec);
            } else {
                auto b_hat = b_vec->unit();
                auto kpara = b_hat->dot(k_vec);
                kpara2 = kpara*kpara;
            }
            
            return wpe2 + 3.0/2.0*kpara2*vterm2 - w*w;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class light_wave final : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
 
//  Equilibrium quantities.
            auto ne = eq->get_electron_density(x, y, z);
            auto wpe2 = build_plasma_frequency(ne,
                                               physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
 
//  Wave numbers should be parallel to B if there is a magnetic field. Otherwise
//  B should be zero.
            assert(eq->get_magnetic_field(x, y, z)->length()->is_match(graph::zero<T, SAFE_MATH> ()) &&
                   "Expected equilibrium with no magnetic field.");
 
            auto k2 = k_vec->dot(k_vec);
            
            return wpe2 + k2 - w*w;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class acoustic_wave final : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
 
//  Equilibrium quantities.
            const T mi = eq->get_ion_mass(0);
            auto te = eq->get_electron_temperature(x, y, z);
            auto ti = eq->get_ion_temperature(0, x, y, z);
            const T gamma = 3.0;
            auto vs2 = (physics<T, SAFE_MATH>::q*te + gamma*physics<T, SAFE_MATH>::q*ti)
                     / (mi*physics<T, SAFE_MATH>::c*physics<T, SAFE_MATH>::c);
 
//  Wave numbers should be parallel to B if there is a magnetic field. Otherwise
//  B should be zero.
            auto b_vec = eq->get_magnetic_field(x, y, z);
            graph::shared_leaf<T, SAFE_MATH> kpara2;
            if (b_vec->length()->is_match(graph::zero<T, SAFE_MATH> ())) {
                kpara2 = k_vec->dot(k_vec);
            } else {
                auto b_hat = b_vec->unit();
                auto kpara = b_hat->dot(k_vec);
                kpara2 = kpara*kpara;
            }
            
            return kpara2*vs2 - w*w;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class gaussian_well final : public dispersion_function<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
            const T c = 1.0;
            auto well = c - 0.5*exp(-(x*x + y*y)/0.1);
            auto npar2 = k_vec->get_z()*k_vec->get_z()*c*c/(w*w);
            auto nperp2 = (k_vec->get_x()*k_vec->get_x() +
                           k_vec->get_y()*k_vec->get_y())*c*c/(w*w);
            return npar2 + nperp2 - well;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class ion_cyclotron final : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
//  Equilibrium quantities.
            const T mi = eq->get_ion_mass(0);
 
            auto te = eq->get_electron_temperature(x, y, z);
            auto ti = eq->get_ion_temperature(0, x, y, z);
            const T gamma = 3.0;
            auto vs2 = (physics<T, SAFE_MATH>::q*te +
                        gamma*physics<T, SAFE_MATH>::q*ti)
                     / (mi*physics<T, SAFE_MATH>::c *
                        physics<T, SAFE_MATH>::c);
            
            auto b_vec = eq->get_magnetic_field(x, y, z);
            auto wce = build_cyclotron_frequency(-physics<T, SAFE_MATH>::q,
                                                 b_vec->length(),
                                                 physics<T, SAFE_MATH>::me,
                                                 physics<T, SAFE_MATH>::c);
 
//  Wave numbers.
            auto b_hat = b_vec->unit();
            auto kperp = b_hat->cross(k_vec)->length();
            auto kperp2 = kperp*kperp;
 
            auto w2 = w*w;
                        
            return wce - kperp2*vs2 - w*w;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class ordinary_wave final : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
//  Equilibrium quantities.
            auto ne = eq->get_electron_density(x, y, z);
            auto wpe2 = build_plasma_frequency(ne,
                                               physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
 
//  Wave numbers.
            auto n = k_vec/w;
            auto b_vec = eq->get_magnetic_field(x, y, z);
            auto b_hat = b_vec->unit();
            auto nperp = b_hat->cross(n);
            auto nperp2 = nperp->dot(nperp);
 
            auto w2 = w*w;
 
            return 1.0 - wpe2/w2 - nperp2;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class extra_ordinary_wave final : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
//  Equilibrium quantities.
            auto ne = eq->get_electron_density(x, y, z);
            auto wpe2 = build_plasma_frequency(ne,
                                               physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
            
            auto b_vec = eq->get_magnetic_field(x, y, z);
            auto b_len = b_vec->length();
            auto wec = build_cyclotron_frequency(-physics<T, SAFE_MATH>::q,
                                                 b_len,
                                                 physics<T, SAFE_MATH>::me,
                                                 physics<T, SAFE_MATH>::c);
            
//  Wave numbers.
            auto n = k_vec/w;
            auto b_hat = b_vec->unit();
            auto nperp = b_hat->cross(n);
            auto nperp2 = nperp->dot(nperp);
 
            auto wh = wpe2 + wec*wec;
            
            auto w2 = w*w;
            
            return 1.0 - wpe2/(w2)*(w2 - wpe2)/(w2 - wh) - nperp2;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class cold_plasma : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
//  Dielectric terms.
//  Frequencies
            auto ne = eq->get_electron_density(x, y, z);
            auto wpe2 = build_plasma_frequency(ne,
                                               physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
            auto b_vec = eq->get_magnetic_field(x, y, z);
            auto b_len = b_vec->length();
            auto ec = build_cyclotron_frequency(-physics<T, SAFE_MATH>::q,
                                                b_len,
                                                physics<T, SAFE_MATH>::me,
                                                physics<T, SAFE_MATH>::c);
 
            auto w2 = w*w;
            auto denome = 1.0 - ec*ec/w2;
            auto e11 = 1.0 - (wpe2/w2)/denome;
            auto e12 = ((ec/w)*(wpe2/w2))/denome;
            auto e33 = wpe2;
 
            for (size_t i = 0, ie = eq->get_num_ion_species(); i < ie; i++) {
                const T mi = eq->get_ion_mass(i);
                const T charge = static_cast<T> (eq->get_ion_charge(i))
                               * physics<T, SAFE_MATH>::q;
 
                auto ni = eq->get_ion_density(i, x, y, z);
                auto wpi2 = build_plasma_frequency(ni, charge, mi,
                                                   physics<T, SAFE_MATH>::c,
                                                   physics<T, SAFE_MATH>::epsilon0);
                auto ic = build_cyclotron_frequency(charge, b_len, mi,
                                                    physics<T, SAFE_MATH>::c);
 
                auto denomi = 1.0 - ic*ic/w2;
                e11 = e11 - (wpi2/w2)/denomi;
                e12 = e12 + ((ic/w)*(wpi2/w2))/denomi;
                e33 = e33 + wpi2;
            }
 
            e12 = -1.0*e12;
            e33 = 1.0 - e33/w2;
 
//  Wave numbers.
            auto n = k_vec/w;
            auto b_hat = b_vec->unit();
 
            auto npara = b_hat->dot(n);
            auto npara2 = npara*npara;
            auto nperp = b_hat->cross(n)->length();
            auto nperp2 = nperp*nperp;
            
//  Determinate matrix elements
            auto m11 = e11 - npara2;
            auto m12 = e12;
            auto m13 = npara*nperp;
            auto m22 = e11 - npara2 - nperp2;
            auto m33 = e33 - nperp2;
 
            return (m11*m22 - m12*m12)*m33 - m22*(m13*m13);
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, bool SAFE_MATH=false>
    class cold_plasma_expansion : public physics<T, SAFE_MATH> {
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
//  Setup plasma parameters.
            auto b_vec = eq->get_magnetic_field(x, y, z);
            auto b_len = b_vec->length();
            auto b_hat = b_vec/b_len;
            auto ne = eq->get_electron_density(x, y, z);
            auto te = eq->get_electron_temperature(x, y, z);
 
            auto ve = graph::sqrt(static_cast<T> (2.0)*physics<T, SAFE_MATH>::q*te /
                                  physics<T, SAFE_MATH>::me)
                    / physics<T, SAFE_MATH>::c;
 
//  Setup characteristic frequencies.
            auto ec = build_cyclotron_frequency(physics<T, SAFE_MATH>::q,
                                                b_len,
                                                physics<T, SAFE_MATH>::me,
                                                physics<T, SAFE_MATH>::c);
            auto wpe2 = build_plasma_frequency(ne, physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
 
//  Dispersion quantities.
            auto P = wpe2/(w*w);
            auto q = P/(2.0*(1.0 + ec/w));
 
            auto n = k_vec/w;
            auto n2 = n->dot(n);
            auto npara = n->dot(b_hat);
            auto npara2 = npara*npara;
            auto nperp = b_hat->cross(n);
            auto nperp2 = nperp->dot(nperp);
            auto n2nperp2 = n2*nperp2;
 
            auto q_func = 1.0 - 2.0*q;
            auto n_func = n2 + npara2;
            auto p_func = 1.0 - P;
 
            auto gamma1 = (1.0 - q)*n2nperp2
                        + p_func*(n2*npara2 - (1.0 - q)*n_func)
                        + q_func*(p_func - nperp2);
            auto gamma0 = nperp2*(n2 - 2.0*q_func) + p_func*(2.0*q_func - n_func);
 
            return -P/2.0*(1.0 + ec/w)*gamma0 + (1.0 - ec*ec/(w*w))*gamma1;
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::complex_scalar T, z_func Z, bool SAFE_MATH=false>
    class hot_plasma final : public physics<T, SAFE_MATH> {
    private:
        Z z;
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
//  Setup plasma parameters.
            auto b_vec = eq->get_magnetic_field(x, y, z);
            auto b_len = b_vec->length();
            auto b_hat = b_vec/b_len;
            auto ne = eq->get_electron_density(x, y, z);
            auto te = eq->get_electron_temperature(x, y, z);
            
            auto ve = graph::sqrt(static_cast<T> (2.0)*physics<T, SAFE_MATH>::q*te /
                                  physics<T, SAFE_MATH>::me)
                    / physics<T, SAFE_MATH>::c;
 
//  Setup characteristic frequencies.
            auto ec = build_cyclotron_frequency(physics<T, SAFE_MATH>::q,
                                                b_len,
                                                physics<T, SAFE_MATH>::me,
                                                physics<T, SAFE_MATH>::c);
            auto wpe2 = build_plasma_frequency(ne, physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
 
//  Dispersion quantities.
            auto P = wpe2/(w*w);
            auto q = P/(2.0*(1.0 + ec/w));
 
            auto n = k_vec/w;
            auto n2 = n->dot(n);
            auto npara = n->dot(b_hat);
            auto npara2 = npara*npara;
            auto nperp = b_hat->cross(n);
            auto nperp2 = nperp->dot(nperp);
 
            auto zeta = (1.0 - ec/w)/(npara*ve);
            auto Z_func = this->z.Z(zeta);
            auto zeta_func = 1.0 + zeta*Z_func;
            auto F = ve*zeta*w/(2.0*npara*ec);
            auto isigma = P*Z_func/(2.0*npara*ve);
 
            auto q_func = 1.0 - 2.0*q;
            auto n_func = n2 + npara2;
            auto p_func = 1.0 - P;
 
            auto gamma5 = n2*npara2 - (1.0 - q)*n_func + q_func;
            auto gamma2 = (n2 - q_func)
                        + P*w/(4.0*ec*npara2)*(n_func - 2.0*q_func);
            auto gamma1 = nperp2*((1.0 - q)*n2 - q_func)
                        + p_func*(n2*npara2 - (1.0 - q)*n_func + q_func);
            auto gamma0 = nperp2*(n2 - 2.0*q_func) + p_func*(2.0*q_func - n_func);
 
            return isigma*gamma0 + gamma1 + nperp2*P*w/ec*zeta_func*(gamma2 + gamma5*F);
        }
    };
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<jit::float_scalar T, class Z, bool SAFE_MATH=false>
    class hot_plasma_expansion final : public physics<T, SAFE_MATH> {
    private:
        Z z;
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        virtual graph::shared_leaf<T, SAFE_MATH>
        D(graph::shared_leaf<T, SAFE_MATH> w,
          graph::shared_vector<T, SAFE_MATH> k_vec,
          graph::shared_leaf<T, SAFE_MATH> x,
          graph::shared_leaf<T, SAFE_MATH> y,
          graph::shared_leaf<T, SAFE_MATH> z,
          graph::shared_leaf<T, SAFE_MATH> t,
          equilibrium::shared<T, SAFE_MATH> &eq) {
//  Setup plasma parameters.
            auto b_vec = eq->get_magnetic_field(x, y, z);
            auto b_hat = b_vec->unit();
            auto b_len = b_vec->length();
            auto ne = eq->get_electron_density(x, y, z);
            auto te = eq->get_electron_temperature(x, y, z);
 
            auto ve = graph::sqrt(static_cast<T> (2.0)*physics<T, SAFE_MATH>::q*te /
                                  physics<T, SAFE_MATH>::me);
 
//  Setup characteristic frequencies.
            auto ec = build_cyclotron_frequency(physics<T, SAFE_MATH>::q, b_len,
                                                physics<T, SAFE_MATH>::me,
                                                physics<T, SAFE_MATH>::c);
            auto wpe2 = build_plasma_frequency(ne, physics<T, SAFE_MATH>::q,
                                               physics<T, SAFE_MATH>::me,
                                               physics<T, SAFE_MATH>::c,
                                               physics<T, SAFE_MATH>::epsilon0);
    
//  Dispersion quantities.
            auto P = wpe2/(w*w);
            auto q = P/(2.0*(1.0 + ec/w));
 
            auto n = k_vec/w;
            auto n2 = n->dot(n);
            auto npara = b_hat->dot(n);
            auto npara2 = npara*npara;
            auto nperp = b_hat->cross(n);
            auto nperp2 = nperp->dot(nperp);
    
            auto vtnorm = ve/physics<T, SAFE_MATH>::c;
 
            auto zeta = (1.0 - ec/w)/(npara*vtnorm);
            auto Z_func = this->z.Z(zeta);
 
            auto q_func = 1.0 - 2.0*q;
            auto n_func = n2 + npara2;
            auto n2nperp2 = n2*nperp2;
            auto p_func = 1.0 - P;
 
            auto gamma5 = P*(n2*npara2 - (1.0 - q)*n_func + q_func);
            auto gamma2 = P*w/ec*nperp2*(n2 - q_func)
                        + P*P*w*w/(4.0*ec*ec)*(n_func - 2.0*q_func)*nperp2/npara2;
            auto gamma1 = (1.0 - q)*n2nperp2
                        + p_func*(n2*npara2 - (1.0 - q)*n_func)
                        + q_func*(p_func - nperp2);
 
            return -(1.0 + ec/w)*npara*vtnorm *
                   (gamma1 + gamma2 + nperp2/(2.0*npara)*(w*w/(ec*ec))*vtnorm*zeta*gamma5)*(1.0/Z_func + zeta);
        }
    };
 
//******************************************************************************
//  Dispersion interface.
//******************************************************************************
    template<class D>
    concept function = std::is_base_of<dispersion_function<typename D::base, D::safe_math>, D>::value;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<function DISPERSION_FUNCTION>
    class dispersion_interface {
    protected:
        graph::shared_vector<typename DISPERSION_FUNCTION::base,
                             DISPERSION_FUNCTION::safe_math> k_vec;
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> D;
 
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> dxdt;
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> dydt;
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> dzdt;
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> dkxdt;
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> dkydt;
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> dkzdt;
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math> dsdt;
 
    public:
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        dispersion_interface(graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> w,
                             graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> kx,
                             graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> ky,
                             graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> kz,
                             graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> x,
                             graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> y,
                             graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> z,
                             graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                                DISPERSION_FUNCTION::safe_math> t,
                             equilibrium::shared<typename DISPERSION_FUNCTION::base,
                                                 DISPERSION_FUNCTION::safe_math> &eq) :
        k_vec(kx*eq->get_esup1(x, y, z) +
              ky*eq->get_esup2(x, y, z) +
              kz*eq->get_esup3(x, y, z)),
        D(DISPERSION_FUNCTION().D(w, k_vec, x, y, z, t, eq)) {
//  Add correction term to deal with k_vec being a function of the coordinates.
            auto dkdx = k_vec->df(x);
            auto dkdy = k_vec->df(y);
            auto dkdz = k_vec->df(z);
 
            auto dDdk_vec = graph::vector<typename DISPERSION_FUNCTION::base,
                                          DISPERSION_FUNCTION::safe_math> (
                this->D->df(k_vec->get_x()),
                this->D->df(k_vec->get_y()),
                this->D->df(k_vec->get_z())
            );
 
            auto dDdw = this->D->df(w);
            auto dDdkx = this->D->df(kx);
            auto dDdky = this->D->df(ky);
            auto dDdkz = this->D->df(kz);
            auto dDdx = this->D->df(x);
            auto dDdy = this->D->df(y);
            auto dDdz = this->D->df(z);
 
            if (graph::pseudo_variable_cast(x).get()) {
                dkdx = dkdx->remove_pseudo();
                dkdy = dkdy->remove_pseudo();
                dkdz = dkdz->remove_pseudo();
                dDdk_vec = dDdk_vec->remove_pseudo();
 
                dDdw = dDdw->remove_pseudo();
                dDdkx = dDdkx->remove_pseudo();
                dDdky = dDdky->remove_pseudo();
                dDdkz = dDdkz->remove_pseudo();
                dDdx = dDdx->remove_pseudo();
                dDdy = dDdy->remove_pseudo();
                dDdz = dDdz->remove_pseudo();
            }
 
            dxdt = -dDdkx/dDdw;
            dydt = -dDdky/dDdw;
            dzdt = -dDdkz/dDdw;
            dkxdt = (dDdx - dDdk_vec->dot(dkdx))/dDdw;
            dkydt = (dDdy - dDdk_vec->dot(dkdy))/dDdw;
            dkzdt = (dDdz - dDdk_vec->dot(dkdz))/dDdw;
 
            dsdt = graph::vector(dxdt, dydt, dzdt)->length();
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        solve(graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                                 DISPERSION_FUNCTION::safe_math> x,
              graph::input_nodes<typename DISPERSION_FUNCTION::base,
                                 DISPERSION_FUNCTION::safe_math> inputs,
              const size_t index=0,
              const typename DISPERSION_FUNCTION::base tolerance = 1.0E-30,
              const size_t max_iterations = 1000) {
            auto x_var = graph::variable_cast(x);
 
            workflow::manager<typename DISPERSION_FUNCTION::base,
                              DISPERSION_FUNCTION::safe_math> work(index);
 
            solver::newton(work, {x}, inputs, this->D,
                           graph::shared_random_state<typename DISPERSION_FUNCTION::base,
                                                      DISPERSION_FUNCTION::safe_math> (),
                           tolerance, max_iterations);
 
            work.compile();
            work.run();
 
            work.copy_to_host(x, x_var->data());
 
            return this->D*this->D;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_residual() {
            return this->D*this->D;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_d() {
            return this->D;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_dsdt() {
            return this->dsdt;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_dxdt() {
            return this->dxdt;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_dydt() {
            return this->dydt;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_dzdt() {
            return this->dzdt;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_dkxdt() {
            return this->dkxdt;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_dkydt() {
            return this->dkydt;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        graph::shared_leaf<typename DISPERSION_FUNCTION::base,
                           DISPERSION_FUNCTION::safe_math>
        get_dkzdt() {
            return this->dkzdt;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_dispersion() {
            D->to_latex();
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_dkxdt() {
            get_dkxdt()->to_latex();
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_dkydt() {
            get_dkydt()->to_latex();
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_dkzdt() {
            get_dkzdt()->to_latex();
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_dxdt() {
            get_dxdt()->to_latex();
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_dydt() {
            get_dydt()->to_latex();
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_dzdt() {
            get_dzdt()->to_latex();
            std::cout << std::endl;
        }
    };
}
 
#endif /* dispersion_h */