ASCOT5 is a high-performance orbit-following code for fusion plasma physics and engineering
"""Package for processing ASCOT5 data and generating inputs.
"""
import os
import numpy as np
import unyt

from .ascot5io import Ascot5IO
from .ascotpy  import Ascotpy

from .ascotpy.libascot import _LIBASCOT
from .exceptions import *
from .routines.plotting import openfigureifnoaxes, line2d

# Define the unit system ascot uses and add our own unit types
unyt.define_unit("markers", 1*unyt.Dimensionless)
unyt.define_unit("particles", 1*unyt.Dimensionless)
unyt.define_unit("e", -unyt.electron_charge)
u=unyt.UnitSystem("ascot", "m", "atomic_mass_unit", "s", angle_unit="deg")
u["energy"] = "eV"
u["charge"] = "e"
u["magnetic_field"] = "T"

class Ascot(Ascotpy):
    """Primary tool for processing ASCOT5 data.

    Ascot object represents the data in a single HDF5 file. The class provides
    methods to:

    - add, modify, and delete the data in HDF5 and set descriptions.
    - access to input and output data, and premade plots for visualization.
    - calculate derived quantities from the data.
    - Python interface to ASCOT5 C routines to evaluate data as ASCOT5 does.
    - capability to simulate individual markers using the Python interface.

    In a nutshell, this object is all you need to work with ASCOT5 data unless
    one is doing something exotic.

    The data is accessed via :attr:`~a5py.Ascot.data` attribute and its
    attributes that form a treeview. For example, active magnetic field is
    accessed as :code:`Ascot.data.bfield.active`. See
    :class:`.Ascot5IO` for details.

    The active run is accessed as `Ascot5.data.active` and its methods contain
    most of the plotting and data access routines.

    Attributes
    ----------
    data : :class:`.Ascot5IO`
        Container for the HDF5 data.
    """

    def __init__(self, inputfile=None, create=False, mute="err"):
        """Initialize Ascot instance.

        Parameters
        ----------
        inputfile : str, optional
            Name of the HDF5 file or `None` to create an empty instance.
        create : bool, optional
            Create a new HDF5 file with given name.
        mute : {"yes", "no", "err"}, optional
            Mute output from libascot.so, i.e. don't display the
            "Reading input X" stuff.

            Possible values are: "yes" - both stdout and stderr are muted,
            "no" - nothing is muted, "err" - stderr is shown.
        """
        super().__init__()

        self._inputfile = None
        self.data       = None
        if mute not in ["yes", "no", "err"]:
            raise ValueError("mute must be either \"yes\", \"no\" or \"err\".")

        if create:
            if os.path.isfile(inputfile):
                raise AscotIOException("Cannot create file: file exists.")
            Ascot5IO._create_file(inputfile)

        self.file_load(inputfile)
        self._mute = mute

    def file_getpath(self):
        """Return name of the HDF5 file from which this instance reads the data.

        Returns
        -------
        path : str
            Absolute path or `None` if this is an empty instance.
        """
        return self._inputfile

    def file_load(self, inputfile):
        """Open a new HDF5 file and free resources from the old one.

        Parameters
        ----------
        inputfile : str
            Path to the new HDF5 file.
        """
        # Free any resources that are currently being used (by the old file)
        if _LIBASCOT: self.input_free()

        if inputfile == None:
            # No input file so this is now an empty instance
            self._inputfile = None
            self.data = Ascot5IO(self)
            if _LIBASCOT:
                self._sim.hdf5_in = "".encode('UTF-8')
        else:
            inputfile = os.path.abspath(inputfile)
            try:
                self._inputfile = inputfile
                self.data = Ascot5IO(self)
                if _LIBASCOT:
                    self._sim.hdf5_in = inputfile.encode('UTF-8')
            except:
                self.inputfile = None
                self.data = Ascot5IO(self)
                if _LIBASCOT:
                    self._sim.hdf5_in = "".encode('UTF-8')

                # Pass exception from failed RootNode initialization
                raise

    def input_init(self, run=False, bfield=False, efield=False, plasma=False,
                   wall=False, neutral=False, boozer=False, mhd=False,
                   asigma=False, switch=False):
        """Initialize input so it can be accessed via Python interface.

        This method can be used in three ways:

        1. Without any arguments in which case all inputs are initialized. The
           initialized input is the one that has been set as active.
        2. Provide one of the input arguments (bfield, efield, plasma, wall,
           neutral, boozer, mhd) in which case only that input is initialized.
           The argument can either be QID of the input to be initialized, or
           `True` in which case the active input of that type is initialized.
        3. Set run = `True` or QID of a run group, in which case the inputs of
           that run group are initialized. If you only want a specific input
           to be initialized, set that argument as True, e.g. bfield = True
           would only initialize the magnetic field used by the given run.

        If the input is already initialized, nothing is done. In case there is
        already different input of same type initialized, an error is raised.

        If the input argument is a dict, it is used instead of reading the data from hdf5.

        Parameters
        ----------
        run     : str or bool, optional
            The run group or True to use the one that is active.
        bfield  : str or bool or dict, optional
            The magnetic field input or True to use the one that is active.
        efield  : str or bool or dict, optional
            The electric field input or True to use the one that is active.
        plasma  : str or bool or dict, optional
            The plasma input or True to use the one that is active.
        wall    : str or bool or dict, optional
            The wall input or True to use the one that is active.
        neutral : str or bool or dict, optional
            The neutral input or True to use the one that is active.
        boozer  : str or bool or dict, optional
            The boozer input or True to use the one that is active.
        mhd     : str or bool or dict, optional
            The MHD input or True to use the one that is active.
        asigma  : str or bool or dict, optional
            The atomic data input or True to use the one that is active.
        switch  : bool, optional
            If ``True``, no error is raised if there is already a different
            input initialized and instead the old input is free'd and the new
            one initialized.
        """
        if not _LIBASCOT:
            raise AscotInitException(
                "Python interface disabled as libascot.so is not found")

        args = locals()
        if run:
            # Init inputs using the run; first check whether to init everything
            if not ( isinstance(bfield, bool)  and isinstance(efield, bool) and
                     isinstance(plasma, bool)  and isinstance(wall, bool)   and
                     isinstance(neutral, bool) and isinstance(boozer, bool) and
                     isinstance(mhd, bool)     and isinstance(asigma, bool)):
                raise ValueError(
                    "Cannot specify input explicitly when run=True")
            initall = not any([bfield, efield, plasma, wall, neutral, boozer,
                               mhd, asigma])

            # Init either the active run or that which was requested
            run = self.data.active if isinstance(run, bool) else\
                  self.data["q"+run]

            # Find QIDs of the inputs to be initialized
            for inp in ["bfield", "efield", "plasma", "wall", "neutral",
                        "boozer", "mhd", "asigma"]:
                if args[inp] and not inp in self.data:
                    raise AscotIOException("Input \"" + inp + "\" not present.")
                elif args[inp]:
                    args[inp] = getattr(run, inp).get_qid()
                elif initall and inp in self.data:
                    args[inp] = getattr(run, inp).get_qid()
                else:
                    args[inp] = None

        else:
            # Init given inputs; first check whether to init everything
            initall = not any([bfield, efield, plasma, wall, neutral, boozer,
                               mhd, asigma])

            for inp in ["bfield", "efield", "plasma", "wall", "neutral",
                        "boozer", "mhd", "asigma"]:
                if isinstance(args[inp], dict):
                    # Argument is a dictionary, presumably in the correct format
                    # It is simply passed forward to _init()
                    pass
                elif args[inp] and not inp in self.data:
                    # Requested data not present
                    raise AscotIOException("Input \"" + inp + "\" not present.")
                elif initall and inp in self.data:
                    # Init all inputs that are present
                    args[inp] = getattr(self.data, inp).active.get_qid()
                elif isinstance(args[inp], str):
                    # Argument is QID already but check that it is in the data
                    parent = getattr(self.data, inp)
                    if not ("q"+args[inp]) in parent:
                        raise AscotIOException(
                            "Input \"" + inp + "/" + args[inp]
                            + "\" not present.")
                elif args[inp]:
                    # Argument is True, and data is present
                    args[inp] = getattr(self.data, inp).active.get_qid()
                else:
                    args[inp] = None

        self._init(
            self.data, bfield=args["bfield"], efield=args["efield"],
            plasma=args["plasma"], wall=args["wall"], neutral=args["neutral"],
            boozer=args["boozer"], mhd=args["mhd"], asigma=args["asigma"],
            switch=switch)

    def input_free(self, bfield=False, efield=False, plasma=False, wall=False,
                   neutral=False, boozer=False, mhd=False, asigma=False):
        """Free input used by the Python interface.

        Arguments toggle which input fields are free'd. If called without
        providing any arguments, then all input fields are free'd.

        Parameters
        ----------
        bfield  : bool, optional
            Toggle to free the magnetic field.
        efield  : bool, optional
            Toggle to free the electric field input.
        plasma  : bool, optional
            Toggle to free the plasma input.
        wall    : bool, optional
            Toggle to free the wall input.
        neutral : bool, optional
            Toggle to free the neutral input.
        boozer  : bool, optional
            Toggle to free the boozer input.
        mhd     : bool, optional
            Toggle to free the MHD input.
        asigma  : bool, optional
            Toggle to free the atomic data input.
        """
        if not _LIBASCOT:
            raise AscotInitException(
                "Python interface disabled as libascot.so is not found")
        freeall = not any([bfield, efield, plasma, wall, neutral, boozer,
                           mhd, asigma])
        if freeall:
            self._free(bfield=True, efield=True, plasma=True, wall=True,
                       neutral=True, boozer=True, mhd=True, asigma=True)
        else:
            self._free(bfield=bfield, efield=efield, plasma=plasma, wall=wall,
                       neutral=neutral, boozer=boozer, mhd=mhd, asigma=asigma)

    def preflight_inputspresent(self):
        """Check required inputs are present for this run.
        """

        # Determine what we are simulating.
        runtype = "ascot"

        msg = []
        if runtype == "ascot":
            for inp in ["bfield", "efield", "plasma", "neutral", "options",
                        "marker", "wall", "boozer", "mhd", "asigma"]:
                if not hasattr(self.data, inp):
                    msg += ["Error: %s data is missing" % inp]

        return msg

    def preflight_optionsconsistent(self):
        """Check that options are consistent with inputs and each other.
        """
        msg = []
        opt = self.data.options.active

        #msg0 = opt.validatevalues()
        msg0 = []
        if len(msg0) > 0:
            msg += ["Error: following options had invalid parameters:"]
            msg += msg0

        opt = opt.read()
        if opt["ENABLE_MHD"] == 1 and \
           self.data.bfield.active.get_type() == "B_STS":
            msg += ["Error: cannot enable MHD for stellarators"]

        if opt["SIM_MODE"] in [1,2,3] and \
           self.data.marker.active.get_type() == "mrk_fl":
            msg += ["Error: SIM_MODE invalid for field-line markers"]

        # Make rough estimates on how much memory is consumed by diagnostics
        high_memory_consumption = 1e9 # In bits
        # Orbit = Npoint * Nfields (~10) * Nmrk * 8 bit
        orb_mem = opt["ORBITWRITE_NPOINT"] * 10 \
            * self.data.marker.active.read()["n"] * 8
        # Distributions = NR * Nz * ... * 8 bit
        rzp = opt["DIST_NBIN_R"] * opt["DIST_NBIN_Z"] * opt["DIST_NBIN_PHI"]
        rtp = opt["DIST_NBIN_RHO"] * opt["DIST_NBIN_THETA"] * opt["DIST_NBIN_PHI"]
        p2d = opt["DIST_NBIN_PPA"] * opt["DIST_NBIN_PPE"]
        p3d = opt["DIST_NBIN_PR"] * opt["DIST_NBIN_PZ"] * opt["DIST_NBIN_PPHI"]

        if opt["ENABLE_ORBITWRITE"] == 1 and orb_mem > high_memory_consumption:
            msg += ["Warning: orbit diagnostic memory consumption high (~" +
                    str(int(orb_mem / 1e9)) + "Gb)"]

        if opt["ENABLE_DIST_5D"] == 1 and rzp * p2d * 8 > high_memory_consumption:
            msg += ["Warning: 5D distribution memory consumption high (~" +
                    str(int(rzp * p2d * 8 / 1e9)) + "Gb)"]

        if opt["ENABLE_DIST_6D"] == 1 and rzp * p3d * 8 > high_memory_consumption:
            msg += ["Warning: 6D distribution memory consumption high (~" +
                    str(int(rzp * p3d * 8 / 1e9)) + "Gb)"]

        if opt["ENABLE_DIST_RHO5D"] == 1 and rtp * p2d * 8 > high_memory_consumption:
            msg += ["Warning: rho5D distribution memory consumption high (~" +
                    str(int(rtp * p2d * 8 / 1e9)) + "Gb)"]

        if opt["ENABLE_DIST_RHO6D"] == 1 and rtp * p3d * 8 > high_memory_consumption:
            msg += ["Warning: rho6D distribution memory consumption high (~" +
                    str(int(rtp * p3d * 8 / 1e9)) + "Gb)"]

        return msg

    def preflight_bfieldpsi0(self):
        """Checks whether psi0 given in input is actually extreme value.

        Because psi is interpolated with splines, there might be numerical error
        that causes psi (near the axis) to have more extreme value than psi0
        which is given in input. This leads to imaginary rho and termination
        of the simulation if marker ends up there.

        This check uses Monte Carlo method to 1. Draw phi 2. Evaluate axis (R,z)
        3. Draw random (R,z) coordinates within 10 cm of the axis. 4. Evaluate psi
        at that point and compare to psi0. Process repeats N times. Check passes
        if all evaluations are valid.
        """

        data = self.data.bfield.active.read()
        if "psi0" not in data:
            return []

        psi0 = data["psi0"]
        psi1 = data["psi1"]
        psi0 = -7.1

        N     = 10000
        phi   = np.random.rand(N,) * 360
        theta = np.random.rand(N,) * 2 * np.pi

        axis = self.input_eval(1, phi, 0, 0, "axis")
        z0 = axis["axisz"]
        r0 = axis["axisr"]

        R = 0.1 # 10 cm
        r = R * np.cos(theta) + r0
        z = R * np.sin(theta) + z0
        psi = self.input_eval(r, phi, z, 0, "psi")

        if psi0 < psi1 and any(psi < psi0):
            return ["Error: psi0 = %.2e but we found near axis that psi = %.2e"\
                    % (psi0, np.amin(psi)) ]
        elif psi0 > psi1 and any(psi > psi0):
            return ["Error: psi0 = %.2e but we found near axis that psi = %.2e"\
                    % (psi0, np.amax(psi)) ]
        return []

    @openfigureifnoaxes(projection=None)
    def preflight_plottopview(self, axes=None):
        """Plot top view of the machine showing Ip, Bphi, and markers.

        Assumes bfield is initialized in ascotpy.
        """
        r0, z0 = self.input_eval(1, 0, 0, 0, "axisr", "axisz")
        #z0 = axis["axisz"]
        #r0 = axis["axisr"]

        rmin = r0-r0/2
        rmax = r0+r0/2
        dphi = 10 * np.pi/180 # So that b and j quivers dont overlap

        r   = np.linspace(rmin, rmax, 10)
        phi = np.linspace(0, 360, 18, endpoint=False) * np.pi/180
        phi, r = np.meshgrid(phi, r)
        r   = r.ravel()
        phi = phi.ravel()

        br   = np.squeeze(self.input_eval(r, phi, z0, 0, "br"))
        bphi = np.squeeze(self.input_eval(r, phi, z0, 0, "bphi"))
        jr   = np.squeeze(self.input_eval(r, phi + dphi, z0, 0, "jr"))
        jphi = np.squeeze(self.input_eval(r, phi + dphi, z0, 0, "jphi"))

        x  = np.cos(phi) * r
        y  = np.sin(phi) * r
        bx = np.cos(phi) * br - np.sin(phi) * bphi
        by = np.sin(phi) * br + np.cos(phi) * bphi
        bnorm = np.sqrt(bx**2 + by**2)

        xj = np.cos(phi+dphi) * r
        yj = np.sin(phi+dphi) * r
        jx = np.cos(phi+dphi) * jr - np.sin(phi+dphi) * jphi
        jy = np.sin(phi+dphi) * jr + np.cos(phi+dphi) * jphi
        jnorm = np.sqrt(jx**2 + jy**2)

        axes.quiver(x,y, bx/bnorm, by/bnorm,
                    color="blue", scale=40)
        axes.quiver(xj,yj, jx/jnorm, jy/jnorm,
                    color="red", scale=40)

        axes.set_aspect("equal", adjustable="box")

        marker = self.data.marker.active.read()
        x  = np.cos(marker["phi"] * np.pi/180) * marker["r"]
        y  = np.sin(marker["phi"] * np.pi/180) * marker["r"]

        axes.scatter(x,y, s=1, c="black", zorder=-2)

        axes.legend(("Magnetic field", "Plasma current", "Markers"))

    def input_plotwallcontour(self, phi=0*unyt.deg, axes=None):
        """Plot intersection of the wall and the poloidal plane at the given
        toroidal angle.

        Parameters
        ----------
        phi : float
            Toroidal angle of the poloidal plane.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        """
        ls = self.data.wall.active.getwallcontour(phi=phi)
        line2d(ls[:,:,0], ls[:,:,1], c="black", axesequal=True, axes=axes)

a5py/a5py/__init__.py
ascot4fusion/ascot5
"""AFSI5: Versatile fusion source integrator AFSI for fast ion and neutron
studies in fusion devices

ASCOT Fusion Source Integrator (AFSI) is a tool for calculating fusion rates
and fusion products for arbitrary particle populations. It can be used either
as a preprocessing tool e.g. to generate source distribution of alphas or as
a postprocessing tool e.g. to evaluate neutron source from beam-thermal fusion.

AFSI can be used in three ways:

  - thermal: Calculates fusion rates between two Maxwellian populations.
  - beam-thermal: Calculates fusion rates between a Maxwellian and arbitrary
    population e.g. beam ions as obtained from ASCOT5 simulation.
  - beam-beam: Calculates fusion rates between two arbitrary populations.

Possible fusion reactions are listed below in a format
reactant 1 + reactant 2 -> product1 + product2:

  1. D + T   -> He4 + n
  2. D + He3 -> He4 + p
  3. D + D   -> T   + p
  4. D + D   -> He3 + n

Reference:
https://iopscience.iop.org/article/10.1088/1741-4326/aa92e9
"""
import sys
import ctypes
import copy
import numpy as np
import numpy.ctypeslib as npctypes
from a5py.ascot5io.ascot5 import Ascot
from a5py.ascotpy import Ascotpy

from a5py.ascotpy import ascotpy2

def thermal(fn, reaction, minr, maxr, nr, minz, maxz, nz,
            minphi=0, maxphi=2*np.pi, nphi=1, nmc=1000, mult=1.0,
            ispecies1=1, ispecies2=1,
            minppara=-1.3e-19, maxppara=1.3e-19, nppara=80,
            minpperp=0, maxpperp=1.3e-19, npperp=40):
    """Calculate thermonuclear fusion between two thermal (Maxwellian)
    species.

    Parameters
    ----------
    fn : str
        Ascot5 HDF5 file name where plasma data is located.
    reaction : int
        Fusion reaction index.
    minr : float
        Minimum R value in the output distribution.
    maxr : float
        Maximum R value in the output distribution.
    nr : int
        Number of R bins in the output distribution.
    minz : float
        Minimum z value in the output distribution.
    maxz : float
        Maximum z value in the output distribution.
    nz : int
        Number of z bins in the output distribution.
    minphi : float, optional
        Minimum phi value in the output distribution.
    maxphi : float, optional
        Maximum phi value in the output distribution.
    nphi : int, optional
        Number of phi bins in the output distribution.
    nmc : int, optional
        Number of MC samples used in each (R, phi, z) bin.
    mult : float, optional
        Multiplier for the fusion rate.

        Use 0.5 if the population is interacting with itself to avoid double
        counting.
    ispecies1 : int, optional
        Ion species index (as they are listed in the plasma input) for the first
        reactant.
    ispecies2 : int, optional
        Ion species index (as they are listed in the plasma input) for
        the second reactant.
    minppara : float, optional
        Minimum ppara value in the output distribution.
    maxppara : float, optional
        Maximum ppara value in the output distribution.
    nppara : int, optional
        Number of pperp bins in the output distribution.
    minpperp : float, optional
        Minimum pperp value in the output distribution.
    maxpperp : float, optional
        Maximum ppara value in the output distribution.
    npperp : int, optional
        Number of pperp bins in the output distribution.

    Returns
    -------
    prod1 : array_like
        Fusion product 1 distribution.
    prod2 : array_like
        Fusion product 2 distribution.
    """
    time = 0
    r_edges   = np.linspace(minr, maxr, nr+1)
    phi_edges = np.linspace(minphi, maxphi, nphi+1)
    z_edges   = np.linspace(minz, maxz, nz+1)
    r   = (r_edges[1:]   + r_edges[:-1])   / 2
    phi = (phi_edges[1:] + phi_edges[:-1]) / 2
    z   = (z_edges[1:]   + z_edges[:-1])   / 2
    temp  = np.zeros((nr, nphi, nz))
    dens1 = np.zeros((nr, nphi, nz))
    dens2 = np.zeros((nr, nphi, nz))

    apy = Ascotpy(fn)
    apy.init(bfield=True, plasma=True)
    for j in range(nphi):
        for i in range(nr):
            for k in range(nz):
                temp[i,j,k]  = apy.evaluate(r[i], phi[j], z[k], time, "ti1")
                dens1[i,j,k] = apy.evaluate(r[i], phi[j], z[k], time,
                                            "ni"+str(ispecies1))
                dens2[i,j,k] = apy.evaluate(r[i], phi[j], z[k], time,
                                            "ni"+str(ispecies2))
    apy.free(bfield=True,plasma=True)

    thermal1 = _init_thermal_data(
        minr, maxr, nr, minphi, maxphi, nphi, minz, maxz, nz, temp, dens1)
    thermal2 = _init_thermal_data(
        minr, maxr, nr, minphi, maxphi, nphi, minz, maxz, nz, temp, dens2)

    react1 = _init_afsi_data(dist_thermal=thermal1)
    react2 = _init_afsi_data(dist_thermal=thermal2)

    d = {
        "nr" : nr, "nphi" : nphi, "nz" : nz, "nppar" : 1, "npperp" : 1,
        "ntime" : 1, "ncharge" : 1,
        "r_edges"      : r_edges,
        "phi_edges"    : phi_edges,
        "z_edges"      : z_edges,
        "ppar_edges"   : np.linspace(0, 1, 2),
        "pperp_edges"  : np.linspace(0, 1, 2),
        "charge_edges" : np.linspace(0, 2, 2),
        "time_edges"   : np.linspace(0, 1, 2),
        "histogram"    : np.zeros((nr, nphi, nz, 1, 1, 1, 1))
    }
    temp = _init_dist_5d(d)

    q1, q2 = _product_charge(reaction)
    prod1 = _init_product_dist(temp, q1, minppara, maxppara, nppara,
                               minpperp, maxpperp, npperp)
    prod2 = _init_product_dist(temp, q2, minppara, maxppara, nppara,
                               minpperp, maxpperp, npperp)

    ascotpy2.afsi_run(reaction, nmc, react1, react2, prod1, prod2)
    prod1.dist *= mult
    prod2.dist *= mult

    return prod1, prod2

def beamthermal(fn, reaction, it, beam, nmc=1000, mult=1.0, ispecies=1,
                swap=False, minppara=-1.3e-19, maxppara=1.3e-19, nppara=80,
                minpperp=0, maxpperp=1.3e-19, npperp=40):
    """Calculate beam-thermal fusion.

    Parameters
    ----------
    fn : str
        Ascot5 HDF5 file name where plasma data is located.
    reaction : int
        Fusion reaction index
    it : int
        Time index from 5D distribution
    beam1 : dict
        Beam distribution that acts as reactant 1.
    nmc : int, optional
        Number of MC samples used in each (R, phi, z) bin.
    mult : float, optional
        Multiplier for the fusion rate.

        Use 0.5 if the population is interacting with itself to avoid double
        counting.
    ispecies : int, optional
        Ion species index (as they are listed in the plasma input) for thermal
        reactant.
    swap : bool, optional
        Swap reactants so that the beam becomes reactant 2 and thermal species
        reactant 1.
    minppara : float, optional
        Minimum ppara value in the output distribution.
    maxppara : float, optional
        Maximum ppara value in the output distribution.
    nppara : int, optional
        Number of pperp bins in the output distribution.
    minpperp : float, optional
        Minimum pperp value in the output distribution.
    maxpperp : float, optional
        Maximum ppara value in the output distribution.
    npperp : int, optional
        Number of pperp bins in the output distribution.

    Returns
    -------
    prod1 : array_like
        Fusion product 1 distribution.
    prod2 : array_like
        Fusion product 2 distribution.
    """
    time   = 0
    nr     = beam["nr"]
    nphi   = beam["nphi"]
    nz     = beam["nz"]
    r      = beam["r"]
    phi    = beam["phi"]
    z      = beam["z"]
    minr   = beam["r_edges"][0]
    maxr   = beam["r_edges"][-1]
    minphi = beam["phi_edges"][0]
    maxphi = beam["phi_edges"][-1]
    minz   = beam["z_edges"][0]
    maxz   = beam["z_edges"][-1]

    apy = Ascotpy(fn)
    apy.init(bfield=True,plasma=True)
    temp = np.zeros((nr, nphi, nz))
    dens = np.zeros((nr, nphi, nz))
    for j in range(nphi):
        for i in range(nr):
            for k in range(nz):
                temp[i,j,k] = apy.evaluate(r[i], phi[j], z[k], time, "ti1")
                dens[i,j,k] = apy.evaluate(r[i], phi[j], z[k], time,
                                           "ni"+str(ispecies))
    apy.free(bfield=True,plasma=True)

    thermal = _init_thermal_data(
        minr, maxr, nr, minphi, maxphi, nphi, minz, maxz, nz, temp, dens)

    dist = _init_dist_5d(beam)
    react1 = _init_afsi_data(dist_thermal=thermal)
    react2 = _init_afsi_data(dist_5D=dist)

    q1, q2 = _product_charge(reaction)
    prod1 = _init_product_dist(dist, q1, minppara, maxppara, nppara,
                               minpperp, maxpperp, npperp)
    prod2 = _init_product_dist(dist, q2, minppara, maxppara, nppara,
                               minpperp, maxpperp, npperp)

    if swap:
        ascotpy2.afsi_run(reaction, nmc, react1, react2, prod1, prod2)
    else:
        ascotpy2.afsi_run(reaction, nmc, react2, react1, prod1, prod2)
    prod1.dist *= mult
    prod2.dist *= mult

    return prod1, prod2

def beambeam(reaction, it, beam1, beam2=None, nmc=1000, mult=1.0,
             minppara=-1.3e-19, maxppara=1.3e-19, nppara=80,
             minpperp=0, maxpperp=1.3e-19, npperp=40):
    """Calculate beam-beam fusion.

    Parameters
    ----------
    reaction : int
        Fusion reaction index.
    it : int
        Time index from 5D distribution.
    beam1 : dict
        Beam1 distribution.
    beam2 : dict, optional
        Beam2 distribution or None to calculate fusion generation with
        beam1 itself.
    nmc : int, optional
        Number of MC samples used in each (R, phi, z) bin.
    mult : float, optional
        Multiplier for the fusion rate.

        Use 0.5 if the population is interacting with itself to avoid double
        counting.
    minppara : float, optional
        Minimum ppara value in the output distribution.
    maxppara : float, optional
        Maximum ppara value in the output distribution.
    nppara : int, optional
        Number of pperp bins in the output distribution.
    minpperp : float, optional
        Minimum pperp value in the output distribution.
    maxpperp : float, optional
        Maximum ppara value in the output distribution.
    npperp : int, optional
        Number of pperp bins in the output distribution.

    Returns
    -------
    prod1 : array_like
        Fusion product 1 distribution.
    prod2 : array_like
        Fusion product 2 distribution.
    """
    beam1["ntime"] = 1
    beam1["time_edges"] = beam1["time_edges"][it:it+1]
    beam1["histogram"]  = beam1["histogram"][:,:,:,:,:,[it],:]
    dist1 = _init_dist_5d(beam1)

    react1 = _init_afsi_data(dist_5D=dist1)
    if beam2 is not None:
        beam2["ntime"] = 1
        beam2["time_edges"] = beam2["time_edges"][it:it+1]
        beam2["histogram"]  = beam2["histogram"][:,:,:,:,:,[it],:]
        dist2  = _init_dist_5d(beam2)
        react2 = _init_afsi_data(dist_5D=dist2)
    else:
        react2 = react1

    q1, q2 = _product_charge(reaction)
    prod1 = _init_product_dist(dist1, q1, minppara, maxppara, nppara,
                               minpperp, maxpperp, npperp)
    prod2 = _init_product_dist(dist1, q2, minppara, maxppara, nppara,
                               minpperp, maxpperp, npperp)

    ascotpy2.afsi_run(reaction, nmc, react1, react2, prod1, prod2)
    prod1.dist *= mult
    prod2.dist *= mult

    return prod1, prod2

def _init_afsi_data(dist_5D=None, dist_thermal=None):
    afsidata = ascotpy2.afsi_data()
    if dist_5D is not None:
        afsidata.type = 1
        afsidata.dist_5D = ctypes.pointer(dist_5D)
    elif dist_thermal is not None:
        afsidata.type = 2
        afsidata.dist_thermal = ctypes.pointer(dist_thermal)
    else:
        afsidata.type = 0
    return afsidata

def _init_thermal_data(minr, maxr, nr, minphi, maxphi, nphi, minz, maxz, nz,
                       temperature, density):
    thermaldata         = ascotpy2.afsi_thermal_data()
    thermaldata.n_r     = nr
    thermaldata.min_r   = minr
    thermaldata.max_r   = maxr
    thermaldata.n_phi   = nphi
    thermaldata.min_phi = minphi
    thermaldata.max_phi = maxphi
    thermaldata.n_z     = nz
    thermaldata.min_z   = minz
    thermaldata.max_z   = maxz
    thermaldata.temperature = npctypes.as_ctypes(
        np.ascontiguousarray(temperature.ravel(), dtype="f8"))
    thermaldata.density = npctypes.as_ctypes(
        np.ascontiguousarray(density.ravel(), dtype="f8"))
    return thermaldata

def _init_dist_5d(dist):
    data           = ascotpy2.struct_c__SA_dist_5D_data()
    data.n_r       = dist["nr"]
    data.min_r     = dist["r_edges"][0]
    data.max_r     = dist["r_edges"][-1]
    data.n_phi     = dist["nphi"]
    data.min_phi   = dist["phi_edges"][0]
    data.max_phi   = dist["phi_edges"][-1]
    data.n_z       = dist["nz"]
    data.min_z     = dist["z_edges"][0]
    data.max_z     = dist["z_edges"][-1]
    data.n_ppara   = dist["nppar"]
    data.min_ppara = dist["ppar_edges"][0]
    data.max_ppara = dist["ppar_edges"][-1]
    data.n_pperp   = dist["npperp"]
    data.min_pperp = dist["pperp_edges"][0]
    data.max_pperp = dist["pperp_edges"][-1]
    data.n_time    = dist["ntime"]
    data.min_time  = dist["time_edges"][0]
    data.max_time  = dist["time_edges"][-1]
    data.n_q       = dist["ncharge"]
    data.min_q     = dist["charge_edges"][0]
    data.max_q     = dist["charge_edges"][-1]
    data.histogram = npctypes.as_ctypes(np.ascontiguousarray(
        dist["histogram"].ravel(), dtype="f8"))
    return data

def _init_product_dist(react, charge, minppara, maxppara, nppara,
                       minpperp, maxpperp, npperp):
    prod            = ascotpy2.struct_c__SA_dist_5D_data()
    prod.n_r        = react.n_r
    prod.min_r      = react.min_r
    prod.max_r      = react.max_r
    prod.n_phi      = react.n_phi
    prod.min_phi    = react.min_phi
    prod.max_phi    = react.max_phi
    prod.n_z        = react.n_z
    prod.min_z      = react.min_z
    prod.max_z      = react.max_z
    prod.n_ppara    = nppara
    prod.min_ppara  = minppara
    prod.max_ppara  = maxppara
    prod.n_pperp    = npperp
    prod.min_pperp  = minpperp
    prod.max_pperp  = maxpperp
    prod.n_time     = react.n_time
    prod.min_time   = react.min_time
    prod.max_time   = react.max_time
    prod.n_charge   = 1
    prod.min_charge = charge - 1
    prod.max_charge = charge + 1

    distsize = prod.n_r * prod.n_phi * prod.n_z * prod.n_ppara * prod.n_pperp
    prod.histogram  = npctypes.as_ctypes(
        np.ascontiguousarray(np.zeros(distsize), dtype="f8") )
    shape = [prod.n_r, prod.n_phi, prod.n_z, prod.n_ppara, prod.n_pperp]
    prod.dist = npctypes.as_array(prod.histogram, shape=shape)
    return prod

def _product_charge(reaction):
    if reaction == 1:
        return 2, 0
    elif reaction == 2:
        return 2, 1
    elif reaction == 3:
        return 1, 1
    elif reaction == 4:
        return 2, 0

a5py/a5py/afsi/afsi5.py
ascot4fusion/ascot5
import numpy as np

def read_erad(fn):

    with open(fn,'r') as f:
        data = {'comm1' : f.readline()} #Skip comment line
        data['n_rho'] = int(float(f.readline().split()[0]))

        h5data = np.loadtxt(f)

        data['rho'] = h5data[:,0]
        data['dV_drho'] = h5data[:,1]
        # For data in format dV/rho, we can ignore effective minor radius

    # Make sure the input is linearly spaced. If not, interpolate
    tol = 1.0001
    diff = np.diff(data['rho'])
    if ( max(diff)/min(diff) > tol):
        print("Warning! Interpolating dV_drho to uniform grid")
        new_rho = np.linspace(np.amin(data['rho']),
                              np.amax(data['rho']),
                              data['n_rho'])
        data['dV_drho'] = np.interp(new_rho, data['rho'],
                                    data['dV_drho'])
        data['rho'] = new_rho

    return data

def write_ascot4_input_erad(filename,plasmaData,comment):
    '''
    Write the radial electric field to input.erad
    @param plasmaData: dictionary with keys:
        'rho','Er'
    '''

    header=comment+ '\n'

    Nrad = len(plasmaData['rho'])
    header = header + ' {} # Nradial points\n'.format(Nrad)


    D=np.vstack( (plasmaData['rho'],plasmaData['Er']) )
    np.savetxt(filename,np.transpose(D),delimiter=' ',comments='',header=header)

a5py/a5py/ascot4interface/erad.py
ascot4fusion/ascot5
"""
Converting ASCOT4 input files to ASCOT5 input HDF5.
"""
import os.path

import numpy as np
import h5py

import a5py.ascot4interface.markers  as a4markers
import a5py.ascot4interface.magn_bkg as a4magn_bkg
import a5py.ascot4interface.plasma   as a4plasma
import a5py.ascot4interface.erad     as a4erad
import a5py.ascot4interface.neutral  as a4neutral
import a5py.ascot4interface.wall_2d  as a4wall_2d
import a5py.ascot4interface.wall_3d  as a4wall_3d
import a5py.ascot4interface.mhdinput as a4mhdinput

import a5py.ascot5io.B_2DS       as B_2DS
import a5py.ascot5io.B_3DS       as B_3DS
import a5py.ascot5io.B_STS       as B_STS
import a5py.ascot5io.N0_1D       as N0_1D
import a5py.ascot5io.N0_3D       as N0_3D
import a5py.ascot5io.options     as options
import a5py.ascot5io.plasma_1D   as plasma_1D
import a5py.ascot5io.mrk_prt     as mrk_prt
import a5py.ascot5io.mrk_gc      as mrk_gc
import a5py.ascot5io.E_TC        as E_TC
import a5py.ascot5io.E_1DS       as E_1DS
import a5py.ascot5io.wall_2D     as wall_2D
import a5py.ascot5io.wall_3D     as wall_3D
import a5py.ascot5io.boozer      as boozer
import a5py.ascot5io.mhd         as mhd
import a5py.ascot5io.asigma_loc  as asigma_loc
import a5py.ascot5io.ascot5tools as a5tools
import a5py.testascot.helpers    as helpers

from a5py.preprocessing.boozermaps import Boozermaps
from a5py.postprocessing.physicslib import guessMass


import a5py.preprocessing.psilims as psilims

def read_markers(a4folder, h5fn):
    fname = a4folder + "input.particles"
    if (os.path.isfile(fname)):
        data = a4markers.read_particles(fname)
        if 'vphi' in data['fieldNames']:
            # We have particles
            print("Warning! Forcing time to zero for all markers.")
            mrk_prt.write_hdf5(
                fn=h5fn, n=data["fields"]["id"].size, ids=data["fields"]["id"],
                mass=data["fields"]["mass"], charge=data["fields"]["charge"],
                r=data["fields"]["Rprt"], phi=data["fields"]["phiprt"],
                z=data["fields"]["zprt"],
                vr=data["fields"]["vR"], vphi=data["fields"]["vphi"],
                vz=data["fields"]["vz"],
                anum=data["fields"]['Anum'], znum=data["fields"]['Znum'],
                weight=data["fields"]["weight"],
                time=data["fields"]["weight"]*0 )
        elif 'energy' in data['fieldNames']:
            # We have guiding centers (theta is random)
            print("Warning! Forcing time to zero and "
                  "randomizing zeta for all markers.")
            zeta = 2*np.pi*np.random.rand(data["fields"]["id"].size)
            mrk_gc.write_hdf5(
                fn=h5fn, n=data["fields"]["id"].size, ids=data["fields"]["id"],
                mass=data["fields"]["mass"], charge=data["fields"]["charge"],
                r=data["fields"]["R"], phi=data["fields"]["phi"],
                z=data["fields"]["z"],
                energy=data["fields"]["energy"], pitch=data["fields"]["pitch"],
                zeta=zeta,
                anum=data["fields"]['Anum'], znum=data["fields"]['Znum'],
                weight=data["fields"]["weight"],
                time=data["fields"]["weight"]*0 )

def read_bfield(a4folder, h5fn):
    fnamebkg = a4folder + "input.magn_bkg"
    fnamehdr = a4folder + "input.magn_header"
    fnameh5  = a4folder + "input.h5"
    if (os.path.isfile(fnamebkg)) and (os.path.isfile(fnamehdr)):
        data = a4magn_bkg.read_magn_bkg(fnamebkg, fnamehdr)
        if data["nPhi"] > 1:
            B_3DS.write_hdf5(
                fn=h5fn,
                b_rmin=data['r'][0], b_rmax=data['r'][-1], b_nr=data['r'].size,
                b_zmin=data['z'][0], b_zmax=data['z'][-1], b_nz=data['z'].size,
                b_phimin=0, b_phimax=360, b_nphi=data['nPhi'],
                axisr=data['axis_r'], axisz=data['axis_z'],
                psi=data['psi']/(2*np.pi),
                psi0=data['psi0']/(2*np.pi), psi1=data['psi1']/(2*np.pi),
                br=data['br'], bphi=data['bphi'], bz=data['bz'])
        else:
            B_2DS.write_hdf5(
                fn=h5fn,
                rmin=data['r'][0], rmax=data['r'][-1], nr=data['r'].size,
                zmin=data['z'][0], zmax=data['z'][-1], nz=data['z'].size,
                axisr=data['axis_r'], axisz=data['axis_z'],
                psi=data['psi']/(2*np.pi),
                psi0=data['psi0']/(2*np.pi), psi1=data['psi1']/(2*np.pi),
                br=data['br']*0, bphi=data['bphi'], bz=data['bz']*0)

    elif os.path.isfile(fnameh5):
        with h5py.File(fnameh5, 'r') as f:
            if (not "/bfield" in f):
                return
        data = a4magn_bkg.read_magn_bkg_stellarator(fnameh5)
        psilim = [0, 1]
        temp_B_name = B_STS.write_hdf5(
            fn=h5fn,
            b_rmin=data['r'][0], b_rmax=data['r'][-1], b_nr=data['r'].size,
            b_zmin=data['z'][0], b_zmax=data['z'][-1], b_nz=data['z'].size,
            b_phimin=data['phi'][0], b_phimax=data['phi'][-1],
            b_nphi=data['phi'].size - 1,
            psi0=psilim[0], psi1=psilim[1],
            br=data['br'], bphi=data['bphi'], bz=data['bz'], psi=data['s'],
            axis_phimin=data['axis_phi'][0], axis_phimax=data['axis_phi'][-1],
            axis_nphi=data['axis_phi'].size-1,
            axisr=data['axis_r'], axisz=data['axis_z'])
        print("Searching for psiaxis and psisepx.")
        try:
            psilim = psilims.bfield_psi_lims(data, h5fn)
        except OSError:
            print("Error: Ascotpy initialization failed. "
                  "Is libascot.so is in current folder?")
            return
        # psi1 > 1 breaks plasma evaluation, so we only keep the lower limit
        psilim = [psilim[0], 1]
        print("New limits: [" + str(psilim[0]) + ", " + str(psilim[1]) + "]")
        a5tools.removegroup(h5fn, temp_B_name)
        B_STS.write_hdf5(
            fn=h5fn,
            b_rmin=data['r'][0], b_rmax=data['r'][-1], b_nr=data['r'].size,
            b_zmin=data['z'][0], b_zmax=data['z'][-1], b_nz=data['z'].size,
            b_phimin=data['phi'][0], b_phimax=data['phi'][-1],
            b_nphi=data['phi'].size - 1,
            psi0=psilim[0], psi1=psilim[1],
            br=data['br'], bphi=data['bphi'], bz=data['bz'], psi=data['s'],
            axis_phimin=data['axis_phi'][0], axis_phimax=data['axis_phi'][-1],
            axis_nphi=data['axis_phi'].size-1,
            axisr=data['axis_r'], axisz=data['axis_z'])

def read_plasma(a4folder, h5fn):
    fname1d = a4folder + "input.plasma_1d"
    fname2d = a4folder + "input.plasma_2d"
    if (os.path.isfile(fname1d)):
        data = a4plasma.read_plasma(fname1d)
        plasma_1D.write_hdf5(
            fn=h5fn, nrho=data['nrho'], nion=data['nion'],
            anum=data['anum'], znum=data['znum'],
            mass=data['anum'], charge=data['znum'],
            rho=data['rho'], edensity=data['ne'], etemperature=data['te'],
            idensity=data['ni'], itemperature=data['ti'])
    if (os.path.isfile(fname2d)):
        data = a4plasma.read_plasma(fname2d)
        dens_i = np.array([data['ni'+str(i)] for i in range(1,data['nion']+1)])
        print("2D plasma data not yet implemented for ASCOT4. "
              "Skipping 2D plasma input.")
        # plasma_2D.write_hdf5(h5fn, data['nrho'], data['nion'], data['znum'],
        #                      data['anum'],
        #                      data['rho'], np.zeros(data['rho'].shape),
        #                      np.zeros(data['rho'].shape),
        #                      data['ne'], data['te'], dens_i, data['ti1'])

def read_erad(a4folder, h5fn):
    fname = a4folder + "input.erad"
    if (os.path.isfile(fname)):
        data = a4erad.read_erad(fname)
        E_1DS.write_hdf5(
            fn=h5fn, nrho=data['n_rho'],
            rhomin=np.amin(data['rho']), rhomax=np.amax(data['rho']),
            dvdrho=data['dV_drho'], reff=1.0)
    else:
        E = np.array([0, 0, 0])
        E_TC.write_hdf5(fn=h5fn, exyz=E)

def read_neutral(a4folder, h5fn, anum_pls, znum_pls):
    fname1d = a4folder + "input.neutral_1d"
    fname3d = a4folder + "input.neutral_3d"
    if (os.path.isfile(fname1d)):
        data = a4neutral.read_neutral(fname1d)
        N0_1D.write_hdf5(fn=h5fn,
                         rhomin=data['rho'][0],
                         rhomax=data['rho'][data['nrho']-1],
                         nrho=data['nrho'], nspecies=data['nspecies'],
                         anum=anum_pls, znum=znum_pls,
                         density=data['dens'], temperature=data['temp'],
                         maxwellian=1, desc=None)
    else:
        print("No 1D neutral data found. Writing dummy data.")
        N0_1D.write_hdf5_dummy(h5fn)
    if (os.path.isfile(fname3d)):
        print("3D neutral data not yet implemented for ASCOT4. "
              "Skipping 3D neutral input.")

def read_wall(a4folder, h5fn):
    fname = a4folder + "input.wall_2d"
    if (os.path.isfile(fname)):
        data = a4wall_2d.read_wall_2d(fname)
        wall_2D.write_hdf5(h5fn, data['r'].size, data['r'], data['z'])
    fname   = a4folder + "input.wall_3d"
    fnameh5 = a4folder + "input.h5"
    if (os.path.isfile(fname)):
        data = a4wall_3d.read_wall_3d(fname)
        wall_3D.write_hdf5(
            fn=h5fn, nelements=data["flag"].size,
            x1x2x3=data['x1x2x3'], y1y2y3=data['y1y2y3'], z1z2z3=data['z1z2z3'],
            desc='fromASCOT4',
            flag=np.reshape(data['flag'],(data["flag"].size,1)))
    elif (os.path.isfile(fnameh5)):
        with h5py.File(fnameh5, 'r') as f:
            if (not "/wall" in f):
                return
        data = a4wall_3d.read_wall_3d_hdf5(fnameh5)
        wall_3D.write_hdf5(
            fn=h5fn, nelements=data["flag"].size,
            x1x2x3=data['x1x2x3'], y1y2y3=data['y1y2y3'], z1z2z3=data['z1z2z3'],
            desc='fromASCOT4',
            flag=np.reshape(data['flag'],(data["flag"].size,1)))


def read_boozer(a4folder, h5fn):
    fname = a4folder + "boozer_maps.out"
    if (os.path.isfile(fname)):
        b = Boozermaps(fname)
        b.write_hdf5(h5fn)
    else:
        boozer.write_hdf5_dummy(h5fn)


def read_mhd(a4folder, h5fn):
    fname = a4folder + "input.alfven"
    if (os.path.isfile(fname)):
        data = a4mhdinput.read_alfven(fname)
        print("No MHD phase data in Ascot4. Assuming phase = 0.")
        mhd.write_hdf5(
            fn=h5fn,
            nmode     = data["nmode"],
            nmodes    = data["nmodes"],
            mmodes    = data["mmodes"],
            amplitude = data["amplitude"],
            omega     = data["omega"],
            phase     = np.zeros(data["omega"].shape),
            alpha     = data["alpha"],
            phi       = data["phi"],
            nrho      = data["nrho"],
            rhomin    = data["rhomin"],
            rhomax    = data["rhomax"])
    else:
        mhd.write_hdf5_dummy(h5fn)


def run(a4folder, h5fn, overwrite=True):
    """
    Convert ASCOT4 input files to ASCOT5 input HDF5 file.

    Parameters
    ----------

    a4folder : str
        Path to folder where ASCOT4 input files are located. Input
        files should be named "input.magn_bkg", "input.magn_header",
        "input.plasma_1d",    "input.wall_2d",  "input.wall_3d",
        "input.particles",    "input.h5"    and "input.neutral_1d".

    h5fn : str
        Full path to output HDF5 file.
    overwrite : bool, optional
        Indicates whether any existing fields within HDF5 file
        will bee overwritten. Default is true.

    Notes
    -----

    Standard ASCOT4 input does not contain electric field so a dummy
    field is written instead (if radial electric field is not given).

    Avoid having multiple files for same input, e.g. 3D wall both in
    input.h5 and input.wall_3d as this is not supported. Having both
    input.wall_2d and input.wall_3d is okay as these are different
    inputs.
    """

    if not os.path.isfile(h5fn):
        with h5py.File(h5fn, 'a') as f:
            pass

    if a4folder[-1] != "/":
        a4folder += "/"

    with h5py.File(h5fn, 'r') as f:
        groups = f.keys()

    if overwrite or (not "markers" in groups):
        read_markers(a4folder, h5fn)

    # Magnetic field.
    if overwrite or (not "bfield" in groups):
        read_bfield(a4folder, h5fn)

    # Plasma profiles.
    if overwrite or (not "plasma" in groups):
        read_plasma(a4folder, h5fn)

    # Electric field.
    if overwrite or (not "efield" in groups):
        read_erad(a4folder, h5fn)

    # Neutral profiles
    if overwrite or (not "neutral" in groups):
        # Since ASCOT4 neutral input does not specify the species,
        # this information must be extracted from the plasma input
        qid_pls = a5tools.call_ascot5file(h5fn, "get_activeqid", "plasma")
        pls = plasma_1D.read_hdf5(fn=h5fn, qid=qid_pls)
        read_neutral(a4folder, h5fn, pls["anum"][0], pls["znum"][0])

    # Boozer data
    if overwrite or (not "boozer" in groups):
        read_boozer(a4folder, h5fn)

    # MHD input
    if overwrite or (not "mhd" in groups):
        read_mhd(a4folder, h5fn)

    # Wall.
    if overwrite or (not "wall" in groups):
        read_wall(a4folder, h5fn)

    # Atomic data
    if overwrite or (not "asigma" in groups):
        asigma_loc.write_hdf5_empty(h5fn)

    # Options
    if overwrite or (not "options" in groups):
        odict = options.generateopt()
        helpers.clean_opt(odict)
        #GCF
        odict["SIM_MODE"]                  = 2
        odict["FIXEDSTEP_USE_USERDEFINED"] = 1
        odict["FIXEDSTEP_USERDEFINED"]     = 1e-8
        odict["ENDCOND_SIMTIMELIM"]        = 1
        odict["ENDCOND_MAX_SIMTIME"]       = 5e-6
        odict["ENABLE_ORBIT_FOLLOWING"]    = 1
        odict["ENABLE_MHD"]                = 1
        odict["ENABLE_COULOMB_COLLISIONS"] = 1

        options.write_hdf5(h5fn, odict)

a5py/a5py/ascot4interface/interface.py
ascot4fusion/ascot5
import numpy as np
import h5py
from scipy.interpolate import RegularGridInterpolator as rgi
from a5py.ascotpy import Ascotpy
from scipy.optimize import fmin

def read_magn_bkg(fn,hdrfn):
    data = dict()
    with open(fn) as fh:
        tmp = list(map(float,fh.readline().split()))
        data['phi0']       = tmp[0]
        data['nSector']    = int(tmp[1])
        data['nPhi']       = int(tmp[2])
        data['nCoil']      = int(tmp[3])
        data['zeroAtCoil'] = int(tmp[4])

        r1,r2,nr = [float(number) for number in fh.readline().split()]
        nr = int(nr)
        z1,z2,nz = [float(number) for number in fh.readline().split()]
        nz = int(nz)

        data['r'] = np.linspace(r1,r2,nr)
        data['z'] = np.linspace(z1,z2,nz)
        if(data['nPhi'] > 1):
            dphi = 360.0 / ( data['nSector'] * data['nPhi'] )
            data['phi'] = ( np.linspace( data['phi0'] + dphi * 0.5,
                                     data['phi0'] + 360.0 / data['nSector'] - dphi * 0.5,
                                     data['nPhi'] ) )

        data['phimap_tor'] = np.array([int(float(number)) for number in fh.readline().split()])
        data['phimap_pol'] = np.array([int(float(number)) for number in fh.readline().split()])

        h5data = np.array(fh.read().split(), dtype=float).flatten()

        sz2d = nr*nz
        sz3d = nr*nz*data['nPhi']
        data['psi']  = h5data[:sz2d].reshape(nz,nr)
        print(np.shape(h5data),nz,data['nPhi'],nr,sz2d,sz3d)
        data['br']   = h5data[sz2d+0*sz3d:sz2d+1*sz3d].reshape(nz,data['nPhi'],nr).squeeze()
        data['bphi'] = h5data[sz2d+1*sz3d:sz2d+2*sz3d].reshape(nz,data['nPhi'],nr).squeeze()
        data['bz']   = h5data[sz2d+2*sz3d:sz2d+3*sz3d].reshape(nz,data['nPhi'],nr).squeeze()

        data['psi']  = np.transpose(data['psi'])
        data['br']   = np.transpose(data['br'])
        data['bphi'] = np.transpose(data['bphi'])
        data['bz']   = np.transpose(data['bz'])

    read_magn_header(hdrfn,data)
    return data

def read_magn_header(fn,data):
    with open(fn) as fh:
        # first four lines contain nothing interesting
        fh.readline()
        fh.readline()
        fh.readline()
        fh.readline()

        # Next three lines contain axis psi, R, and z values
        tmp = [float(number) for number in fh.readline().split()]
        data['psi0']       = tmp[0]
        data['psi1']       = tmp[1]

        tmp = [float(number) for number in fh.readline().split()]
        data['axis_r']       = tmp[0]

        tmp = [float(number) for number in fh.readline().split()]
        data['axis_z']       = tmp[0]

    return data

def read_magn_bkg_stellarator(fn):

    with h5py.File(fn, 'r') as f: # Open for reading
        data = dict()
        bstr = 'bfield/stellarator/'

        data['r']    = f[bstr+'r'][:]
        data['phi']  = f[bstr+'phi'][:]
        data['z']    = f[bstr+'z'][:]

        data['br']   = f[bstr+'br'][:]
        data['bphi'] = f[bstr+'bphi'][:]
        data['bz']   = f[bstr+'bz'][:]
        data['s']    = f[bstr+'s'][:]

        data['axis_r']   = f[bstr+'axis_R'][:]
        data['axis_phi'] = f[bstr+'axis_phi'][:]
        data['axis_z']   = f[bstr+'axis_z'][:]

        data['n_periods'] = f[bstr+'toroidalPeriods'][:]
        try:
            data['symmetrymode'] = f[bstr+'symmetrymode'][:]
        except KeyError:
            print("Warning! No symmetry mode specified in input.h5/bfield")
            print("Defaulting to stellarator symmetry")
            data['symmetrymode'] = 0
    philim = 360/data['n_periods']/(2 if data['symmetrymode'] == 0 else 1)
    if(data['symmetrymode'] == 0):
        print("Converting stellarator symmetric input to periodic.")
        data = stellarator_bfield_sector2full(data)
    elif (data['phi'][0] == np.mod(data['phi'][-1],philim)):
        print("Warning! Removing duplicate bfield data point.")
        data = bfield_remove_duplicate_phi(data)
    if (data['axis_phi'][0] == np.mod(data['axis_phi'][-1],360)):
        print("Warning! Removing duplicated axis datapoint.")
        data['axis_r'] = data['axis_r'][0:-1]
        data['axis_z'] = data['axis_z'][0:-1]
    # Transpose to ascot5io format
    data['br']   = np.transpose(data['br'],   (2,1,0))
    data['bphi'] = np.transpose(data['bphi'], (2,1,0))
    data['bz']   = np.transpose(data['bz'],   (2,1,0))
    data['s']    = np.transpose(data['s'],    (2,1,0))

    return data

def stellarator_bfield_sector2full(data):
    out = data
    # Data is in the format f(z, phi, r)
    out['r'] = data['r']
    out['z'] = data['z']
    out['phi'] = np.concatenate( (data['phi'][:-1],
                                  data['phi'][:] + data['phi'][-1]) )
    # br
    br = data['br'][:, :-1, :]
    br_sym = - data['br'][-1::-1, -1:0:-1, :]
    out['br'] = np.concatenate((br, br_sym),1)
    # bphi
    bphi = data['bphi'][:, :-1, :]
    bphi_sym = data['bphi'][-1::-1, -1:0:-1, :]
    out['bphi'] = np.concatenate((bphi, bphi_sym),1)
    # bz
    bz = data['bz'][:, :-1, :]
    bz_sym = data['bz'][-1::-1, -1:0:-1, :]
    out['bz'] = np.concatenate((bz, bz_sym),1)
    # s
    s = data['s'][:, :-1, :]
    s_sym = data['s'][-1::-1, -1:0:-1, :]
    out['s'] = np.concatenate((s, s_sym),1)
    out['symmetrymode'] = 1
    return out

def bfield_remove_duplicate_phi(data):
    out = data
    out['br'] = data['br'][:, :-1, :]
    out['bphi'] = data['bphi'][:, :-1, :]
    out['bz'] = data['bz'][:, :-1, :]
    out['s'] = data['s'][:, :-1, :]
    return out

a5py/a5py/ascot4interface/magn_bkg.py
ascot4fusion/ascot5
import numpy as np
from a5py.postprocessing.physicslib import guessMass
import a5py.ascot5io.mrk_gc
import a5py.ascot5io.mrk_prt
import time

def read_particles(fname):
    '''
    Read ASCOT4 input.particles file
    '''
    import numpy

    # Read Header part
    #-----------------
    with open(fname,'r') as f:
        headerLength = 0

        data={'filename':fname,'fields':{}}

        line = f.readline()
        headerLength += 1
        if line[:-1] != ' PARTICLE INITIAL DATA FOR ASCOT' and line[:-1] != ' PARTICLE OUTPUT DATA FROM ASCOT':
            print( '"'+line[:-1]+'"')
            raise NameError('Unrecognized first line in "'+fname+'".')

        line = f.readline()
        headerLength += 1
        if line.split()[0] != '4':
            print( '"'+line+'"')
            raise NameError('Bad file version in "'+fname+'".')

        line = f.readline()
        headerLength += 1

        line = f.readline()
        headerLength += 1
        nComments = int(float(line.split()[0]))

        comments=[]
        for i in range(0,nComments):
            comments.append(f.readline()[:-1])
            headerLength += 1
        data['comments']=comments

        line = f.readline()
        headerLength += 1

        line = f.readline()
        headerLength += 1
        nParticles = int(float(line.split()[0]))

        line = f.readline()
        headerLength += 1

        line = f.readline()
        headerLength += 1
        nFields = int(float(line.split()[0]))

        fields = []
        for i in range(0,nFields):
            fields.append(f.readline()[:10].strip())
            headerLength += 1
        data['fieldNames']=fields

        line = f.readline()
        headerLength += 1

    # Read the actual data table
    # --------------------------

    print('Reading', nFields, 'fields for', nParticles, 'particles.')
    columns = numpy.loadtxt(fname,skiprows=headerLength)
    
    # -1 means unkown number of lines.
    if nParticles == -1:
        nParticles = columns.shape[0]

    print('Read ', nParticles, 'particles')

    for i in range(0,nFields):
        data['fields'][data['fieldNames'][i]]=-999.0 * numpy.ones(nParticles)

    for i in range(0,nParticles):
        for j in range(0,nFields):
            data['fields'][data['fieldNames'][j]][i] = columns[i,j]

    if 'charge' not in data['fieldNames'] and 'Znum' in data['fieldNames']:
        print("Converting Znum to charge.")
        data["fields"]['charge'] = data["fields"]['Znum'].astype('float')
    if 'mass' not in data['fieldNames'] and 'Anum' in data['fieldNames']:
        print("Converting Anum to mass.")
        data["fields"]['mass'] = np.array(
            list(map(guessMass,
                     data["fields"]['Anum'],
                     data["fields"]['Znum'],
                     data["fields"]['charge'])))
    if 'id' not in data['fieldNames']:
        print("Generating unique ids.")
        data["fields"]['id'] = np.array(
            range(1,data["fields"]['charge'].size + 1))
    if (min(data["fields"]["id"]) <= 0):
        zero_ind = np.where(data["fields"]["id"] == 0)[0];
        data["fields"]["id"][zero_ind] = max(data["fields"]["id"] ) + 1
        print("Converting id 0 to new unique id: "
              + str(int(max(data["fields"]["id"]))))

    return data



def write_particles(filename,dataObject):
    ''' 
        Write ASCOT4 input.particles 
        @param filename: The file to write
        @param dataObject: The particle data to write, e.g. a5py.ascot5io.mrk_gc.mrk_gc 
    '''
        
    if isinstance(dataObject, a5py.ascot5io.mrk_gc.mrk_gc):
        print('Writing guiding centres')
        
        d=dataObject.read()
        # ascot5 name, ascot4 name, conversion factor
        fieldTuples=[('anum','Anum',    1.0),
                     ('charge','charge',1.0),
                     ('mass','mass',    1.0),
                     ('znum','Znum',    1.0),
                     ('r','R',          1.0),
                     ('phi','phi',      1.0),
                     ('z','z',          1.0),
                     ('pitch','pitch',  1.0),
                     ('energy','energy',1.0),
                     ('time','time',    1.0),
                     ('weight','weight',1.0),
                     ('ids','id',       1.0),   ]
        
        data={'fields':{}}
        for t in fieldTuples:
            data['fields'][t[1]] = t[2] * d[t[0]]
            #print(t[0], d[t[0]].shape)
        
    elif isinstance(dataObject, a5py.ascot5io.mrk_prt.mrk_prt):
        raise ValueError('Unsupported dataObject to write.')
        print('Writing particles')
        
    else:
        raise ValueError('Unknown dataObject to write.')
        
        
        
    data['comments']=[time.strftime('Written with a5py on %Y-%m-%d at %H:%M:%S%z.'),
                      'The input file and group are "{}" : "{}"'.format(dataObject._file, dataObject._path),
                      'Desc: "{}"'.format(dataObject.get_desc())]

    
    
    
    nFields   = len(data['fields'])
    fields    = list(data['fields'])
    nParticles= len(data['fields'][fields[0]])
    
    header=''' PARTICLE INITIAL DATA FOR ASCOT
 4 VERSION =====================
 
 {0}  # Number of comment lines, max length 256 char, (defined in prtRead_lineLength)
'''.format(len(data['comments']))

    # ADD THe COMMENTS
    
    for i in range(0,len(data['comments'])):
        header= header + data['comments'][i] + '\n'

    # ADD NUMBER OF PARTICLES
        
    header = header + '\n{0} # Number of particles (-1 means unknown number)\n\n'.format(nParticles)

    
    # ADD THE FIELD LIST
    
    header =  header + '{0} # Number of different fields for each particle [10 first letters are significant]\n'.format(nFields)

    fmt = []
    D   = np.empty((nParticles,0))
    
    descs={ 'Anum'   : 'mass number of particle        (integer)',
            'mass'   : 'mass of the particle           (amu)',
            'Znum'   : 'charge number of particle      (integer)',
            'charge' : 'charge of the particle         (elemental charge)',
            'energy' : 'kinetic energy of particle     (eV)',
            'rho'    : 'normalized poloidal flux       (real)',
            'phi'    : 'toroidal angle of particle     (deg)',
            'R'      : 'R of particle                  (m)',
            'z'      : 'z of particle                  (m)',
            'pitch'  : 'pitch angle cosine of particle (vpar/vtot)',
            'origin' : 'origin of the particle         ()',
            'weight' : 'weight factor of particle      (particle/second)',
            'id'     : 'unique identifier of particle  (integer)',
            'Tmax'   : 'maximum time to follow the prt (s)',
            'time'   : 'elapsed time of the particle   (s)',
            }

    
    #( D.shape, np.transpose(np.array([data['fields'][fields[0]]])).shape)
    integerFields = ['Anum','Znum','id','origin']
    for i in range(0,nFields):
        description = descs.get(fields[i],' ---')
        header = header + '{:10s}- {}\n'.format(fields[i],description)
        if any( fields[i] in s   for s in integerFields ):
            fmt.append('%5d')
        else:
            fmt.append('%20.10e')
        D = np.hstack((D,np.transpose(np.array([data['fields'][fields[i]]]))))
    
    footer = '#EOF'
    
    
    #print '"' + header + '"' 
    #print fmt
    #print D.shape
    #print footer
    print ('writing', nFields, 'fields for', nParticles,'particles to "'+filename+'".')
    np.savetxt(filename,D,fmt=fmt,delimiter=' ',comments='',header=header,footer=footer)



a5py/a5py/ascot4interface/markers.py
ascot4fusion/ascot5
import numpy as np
from scipy import interpolate

def read_alfven(fn):
    data = {}
    with open(fn) as fh:

        # Skip first three lines
        fh.readline()
        fh.readline()
        fh.readline()

        # AEs on (1) or off (0)
        data["ae"] = int(fh.readline().split()[0]) == 1

        # Total number of different modes
        data["nmode"] = int(fh.readline().split()[0])

        # Number of radial points for profiles
        data["nrho"] = int(fh.readline().split()[0])

        # Skip empty line
        fh.readline()

        # Poloidal mode numbers
        data["mmodes"] = np.array(fh.readline().split()[:data["nmode"]])
        data["mmodes"] = data["mmodes"].astype(int)

        # Toroidal mode numbers
        data["nmodes"] = np.array(fh.readline().split()[:data["nmode"]])
        data["nmodes"] = data["nmodes"].astype(int)

        # Amplitudes
        data["amplitude"] = np.array(fh.readline().split()[:data["nmode"]])
        data["amplitude"] = data["amplitude"].astype(np.float)

        # Angular frequencies (omega) [rad/s]
        data["omega"] = np.array(fh.readline().split()[:data["nmode"]])
        data["omega"] = data["omega"].astype(np.float)

        # Phase not given, fix it at zero
        data["phase"] = data["omega"]*0

        # psin, alpha profile, phi profile, each line corresponds to one psislot
        line = fh.readline()
        psi   = np.zeros( (data["nrho"],1) )
        alpha = np.zeros( (data["nrho"],data["nmode"]) )
        phi   = np.zeros( (data["nrho"],data["nmode"]) )
        line = fh.readline()

        # File ends in #EOF
        for i in range(data["nrho"]):
            line = line.split()
            psi[i]     = line[0]
            alpha[i,:] = line[1:data["nmode"]+1]
            phi[i,:]   = line[data["nmode"]+1:2*data["nmode"]+1]
            line = fh.readline()

        # Set data on uniform grid
        data["rhomin"] = psi[0]
        data["rhomax"] = psi[-1]

        psiq = np.linspace(data["rhomin"], data["rhomax"], data["nrho"])

        data["alpha"] = np.zeros( (data["nrho"],data["nmode"]) )
        data["phi"]   = np.zeros( (data["nrho"],data["nmode"]) )
        for i in range(data["nmode"]):
            f = interpolate.interp1d(psi.ravel(), alpha[:,i])
            data["alpha"][:,i] = f(psiq).ravel()

            f = interpolate.interp1d(psi.ravel(), phi[:,i])
            data["phi"][:,i] = f(psiq).ravel()

        for i in range(data["nmode"]):
            f = interpolate.interp1d(psiq.ravel(), data["alpha"][:,i])
            data["alpha"][:,i] = f(psiq*psiq).ravel()

            f = interpolate.interp1d(psiq.ravel(), data["phi"][:,i])
            data["phi"][:,i] = f(psiq*psiq).ravel()

    return data

a5py/a5py/ascot4interface/mhdinput.py
ascot4fusion/ascot5
import numpy as np

def read_neutral(fn):
    with open(fn,'r') as f:
        if(fn[-2:] == '1d'):
            f.seek(0)
            return read_1d(f)
        elif(fn[-2:] == '3d'):
            print("3D neutral data has not yet been implemented for ASCOT4, "
                  "and neither has a method for reading it here.")
        else:
            print('Unrecognized suffix in filename')

def read_1d(fh):
    data = {'comm1' : fh.readline(),
            'comm2' : fh.readline(),
            'comm3' : fh.readline()}
    nrho = int(float(fh.readline()))
    nspecies = 1 # ASCOT4 does not support several neutral species
    data['nrho'] = nrho
    data['nspecies'] = nspecies
    fh.readline() # ignore headers
    h5data = np.loadtxt(fh)
    data['rho'] = np.array(h5data[:,0])
    for i in range(0, data['nspecies']):
        data['dens'+str(i+1)] = np.array(h5data[:,1+i])
    data['temp'] = np.array(h5data[:,2])
    # Make sure the input is linearly spaced. If not, interpolate
    tol = 1.0001
    diff = np.diff(data['rho'])
    if ( max(diff)/min(diff) > tol):
        print("Warning! Interpolating neutral data to uniform grid")
        new_rho = np.linspace(np.amin(data['rho']),
                              np.amax(data['rho']),
                              data['nrho'])
        for i in range(0, data['nspecies']):
            data['dens'+str(i+1)] = np.interp(new_rho, data['rho'],
                                              data['dens'+str(i+1)])
        data['temp'] = np.interp(new_rho, data['rho'], data['temp'])
        data['rho'] = new_rho
    data['dens'] = np.array([data['dens'+str(i+1)] for i in range(0, data['nspecies'])])
    data['dens'] = np.transpose(data['dens'])
    # Add extra data point outside rho=1 to avoid out of data range errors
    if ( np.amax(data['rho']) <= 1.0 ):
        print("Warning! Adding small datapoint outside rho=1.0")
        data['nrho'] = data['nrho'] + 1
        data['rho'] = np.append(data['rho'], 2*data['rho'][-1]-data['rho'][-2])
        data['dens'] = np.append(data['dens'],
                                 np.expand_dims(data['dens'][-1,:]*1e-10, 1).T,
                                 axis=0)
        data['temp'] = np.append(data['temp'], data['temp'][-1])
    return data

a5py/a5py/ascot4interface/neutral.py
ascot4fusion/ascot5
import numpy as np

def read_plasma(fn):
    with open(fn,'r') as f:
        c1 = f.readline()
        c2 = f.readline()
        c3 = f.readline()
        l = f.readline().split()
        if(fn[-2:] == '1d'):
            f.seek(0)
            return read_1d(f)
        elif(fn[-2:] == '2d'):
            f.seek(0)
            return read_2d(f)
        else:
            print('unrecognized first line after comments')
            print(l)

def read_1d(fh):
    data = {'comm1' : fh.readline(),'comm2' : fh.readline(),'comm3' : fh.readline()}
    nrho,nion = [int(float(number)) for number in fh.readline().split()[:2]]
    data['znum'] = np.array([int(float(znum)) for znum in fh.readline().split()[:nion]])
    data['anum'] = np.array([int(float(anum)) for anum in fh.readline().split()[:nion]])
    data['coll'] = np.array([int(float(coll)) for coll in fh.readline().split()[:nion+1]])
    data['nrho'] = nrho
    data['nion'] = nion
    fh.readline() # ignore headers
    h5data = np.loadtxt(fh)
    data['rho'] = np.array(h5data[:,0])
    data['te'] = np.array(h5data[:,1])
    data['ne'] = np.array(h5data[:,2])
    data['vtor'] = np.array(h5data[:,3])
    data['ti'] = np.array(h5data[:,4])
    for i in range(0,nion):
        data['ni'+str(i+1)] = np.array(h5data[:,5+i])
    # Make sure the input is linearly spaced. If not, interpolate
    tol = 1.0001
    diff = np.diff(data['rho'])
    if ( max(diff)/min(diff) > tol):
        print("Warning! Interpolating plasma data to uniform grid")
        new_rho = np.linspace(np.amin(data['rho']),
                              np.amax(data['rho']),
                              data['nrho'])
        data['ne'] = np.interp(new_rho, data['rho'], data['ne'])
        data['te'] = np.interp(new_rho, data['rho'], data['te'])
        for i in range(1, data['nion']+1):
            data['ni'+str(i)] = np.interp(new_rho, data['rho'],
                                          data['ni'+str(i)])
        data['ti'] = np.interp(new_rho, data['rho'], data['ti'])
        data['rho'] = new_rho
    data['ni'] = np.array([data['ni'+str(i)] for i in range(1,data['nion']+1)])
    data['ni'] = np.transpose(data['ni'])
    # Add extra data point outside rho=1 to avoid out of data range errors
    if ( np.amax(data['rho']) <= 1.0 ):
        print("Warning! Adding small datapoint outside rho=1.0")
        data['nrho'] = data['nrho'] + 1
        data['rho'] = np.append(data['rho'], 2*data['rho'][-1]-data['rho'][-2])
        data['ne'] = np.append(data['ne'], data['ne'][-1]*1e-10)
        data['te'] = np.append(data['te'], data['te'][-1])
        data['ni'] = np.append(data['ni'],
                               np.expand_dims(data['ni'][-1,:]*1e-10, 1).T,
                               axis=0)
        data['ti'] = np.append(data['ti'], data['ti'][-1])
    return data

def read_2d(fh):
    data = {'comm1' : fh.readline(),'comm2' : fh.readline(),'comm3' : fh.readline()}
    nr,nz = [int(float(number)) for number in fh.readline().split()[:2]]
    rmin,rmax,zmin,zmax = [float(number) for number in fh.readline().split()[:4]]
    nion,rho2d = [float(number) for number in fh.readline().split()[:2]]
    nion = int(nion)
    data['znum'] = np.array([float(znum) for znum in fh.readline().split()[:nion]])
    data['anum'] = np.array([float(anum) for anum in fh.readline().split()[:nion]])
    data['coll'] = np.array([float(coll) for coll in fh.readline().split()[:nion+1]])
    fieldnames = fh.readline().split()[0:-1:2]
    h5data = loadtxt(fh)
    data['r'] = linspace(rmin,rmax,nr)
    data['z'] = linspace(zmin,zmax,nz)
    for i,name in enumerate(fieldnames):
        data[name.lower()] = h5data[:,i].reshape(nr,nz).T

    if(i+1 < shape(h5data)[1]):
        print('Not all data fields assigned to struct')
        print('Something is probably wrong with the data file')

    return data

def write_plasma_1d(fn, p):
    dens_i = np.array([p['ni'+str(i)] for i in range(1,p['nion']+1)])

    raise RuntimeError('Bug on the next line, znum used twice, anum never. And p1D undefined.')
    p1D.write_hdf5(fn, p['nrho'], p['nion'], p['znum'], p['znum'], p['rho'],
                   np.zeros(p['rho'].shape), np.zeros(p['rho'].shape),
                   p['ne'], p['te'], dens_i, p['ti'])


def write_ascot4_plasma_1d(filename,plasmaData,Anum,Znum,comment1,comment2,comment3):
    '''
    This routine writes the plasma_1d in ASCOT4 format.


    @param plasmaData: A dictionary with the following keys:
        'rho', 'te', 'ne', 'ti', 'ni'
    '''

    header=comment1 + '\n' + comment2 + '\n' + comment3 + '\n'

    Nrad = len(plasmaData['rho'])
    Nion = len(Anum)
    header = header + ' {} {} # Nrad,Nion\n'.format(Nrad,Nion)

    for i in range(0,Nion):
        header = header + ' {}'.format(Znum[i])
    header = header + ' # ion Znum\n'

    for i in range(0,Nion):
        header = header + ' {}'.format(Anum[i])
    header = header + ' # ion Anum\n'

    header = header + ' {}'.format(1)
    for i in range(0,Nion):
        header = header + ' {}'.format(1)
    header = header + ' # OBSOLETE VALUES. PUT 1\n'
    header = header + 'RHO (pol)     Te (eV)         Ne (1/m3)       Vtor_I (rad/s)  Ti1 (eV)        Ni1 (1/m3)      Ni2 (1/m3) ...'

    D=np.vstack( (plasmaData['rho'],plasmaData['te'],plasmaData['ne'],
                  -999.0*np.ones_like(plasmaData['rho']),plasmaData['ti'],plasmaData['ni']) )
    np.savetxt(filename,np.transpose(D),delimiter=' ',comments='',header=header)


a5py/a5py/ascot4interface/plasma.py
ascot4fusion/ascot5
import numpy as np

def read_wall_2d(fn):
    with open(fn,'r') as f:
        h5data = np.loadtxt(f,skiprows=1)
    data = {'r' : h5data[:,0], 'z' : h5data[:,1], 'divflag' : h5data[:,2]}
    return data

a5py/a5py/ascot4interface/wall_2d.py
ascot4fusion/ascot5
import h5py
import numpy as np

# The input file contains an integer giving the number of wall triangles
# followed by a list of coordinates for the corners of each triangle
# (x1,y1,z1,x2,y2,z2,x3,y3,z3).

def read_wall_3d(fn):

    # Crudely check number of lines to get maximum size for data
    num_lines = sum(1 for line in open(fn))
    copysector = False

    with open(fn,'r') as f:

        # fprintf(fid,'%d wall sectors; number, type, min. & max. phi(deg): \n\n',w.nsector);
        data = {}
        data['n_sectors'] = int(f.readline().split()[0])

        # Read sector ids
        data['ids'] = {}
        f.readline()
        for sector in range(data['n_sectors']):
            line = f.readline().split()
            data['ids'][int(float(line[0]))] = int(float(line[1]))

        # Read data for each sector
        n_read = 0

        data['x1x2x3'] = np.array(np.zeros((num_lines, 3)))
        data['y1y2y3'] = np.array(np.zeros((num_lines, 3)))
        data['z1z2z3'] = np.array(np.zeros((num_lines, 3)))
        data['id'] = np.array(np.zeros((num_lines, 1)))
        data['flag'] = np.array(np.zeros((num_lines, 1))) # This is not filled (for lack of test-files)
        for sector in range(1,data['n_sectors']+1):
            f.readline() # Skip empty line
            try:
                n_elements = int(float(f.readline().split()[0]))
            except Exception:
                print("Warning: There are multiple sectors but\n"
                      + "data only for one sector. The sector is\n"
                      + "copied and rotated." )
                copysector = True
                break
            print(f.readline())
            for i in range(n_elements):
                f.readline() # Skip empty line
                f.readline() # Skip element info
                for j in range(3):
                    line = f.readline().split()
                    data['x1x2x3'][n_read,j] = line[0]
                    data['y1y2y3'][n_read,j] = line[1]
                    data['z1z2z3'][n_read,j] = line[2]
                data['id'][n_read] = data['ids'][sector]
                n_read = n_read + 1

                # Some data may contain rectangles instead of triangles.
                # This just means we read the "extra" vertice and make another
                # triangle from it.
                pos = f.tell()
                line = f.readline()
                if len(line) == 0 or len(line) == 1:
                    # This was an empty line (no vertice data). Return to
                    # previous line.
                    f.seek(pos)
                else:
                    line = line.split()
                    data['x1x2x3'][n_read,0] = line[0]
                    data['y1y2y3'][n_read,0] = line[1]
                    data['z1z2z3'][n_read,0] = line[2]
                    data['x1x2x3'][n_read,1] = data['x1x2x3'][n_read-1,0]
                    data['y1y2y3'][n_read,1] = data['y1y2y3'][n_read-1,0]
                    data['z1z2z3'][n_read,1] = data['z1z2z3'][n_read-1,0]
                    data['x1x2x3'][n_read,2] = data['x1x2x3'][n_read-1,2]
                    data['y1y2y3'][n_read,2] = data['y1y2y3'][n_read-1,2]
                    data['z1z2z3'][n_read,2] = data['z1z2z3'][n_read-1,2]
                    data['id'][n_read] = data['ids'][sector]
                    n_read = n_read + 1

    # Remove extra zeros
    data['x1x2x3'] = data['x1x2x3'][0:n_read,:]
    data['y1y2y3'] = data['y1y2y3'][0:n_read,:]
    data['z1z2z3'] = data['z1z2z3'][0:n_read,:]
    data['id'] = data['id'][0:n_read,:]

    if copysector:
        r1r2r3 = np.sqrt(data['x1x2x3']*data['x1x2x3'] + data['y1y2y3']*data['y1y2y3'])
        t1t2t3 = np.arctan2(data['y1y2y3'], data['x1x2x3'])
        z1z2z3 = data['z1z2z3']
        data['x1x2x3'] = np.zeros( (n_read*data['n_sectors'],3) )
        data['y1y2y3'] = np.zeros( (n_read*data['n_sectors'],3) )
        data['z1z2z3'] = np.zeros( (n_read*data['n_sectors'],3) )

        for i in range(data['n_sectors']):
            x1x2x3 = r1r2r3 * np.cos(t1t2t3 + i*2*np.pi/data['n_sectors'])
            y1y2y3 = r1r2r3 * np.sin(t1t2t3 + i*2*np.pi/data['n_sectors'])
            data['x1x2x3'][n_read*i:n_read*(i+1),:] = x1x2x3
            data['y1y2y3'][n_read*i:n_read*(i+1),:] = y1y2y3
            data['z1z2z3'][n_read*i:n_read*(i+1),:] = z1z2z3

        data['id'] = np.linspace(1, n_read*data['n_sectors'],n_read*data['n_sectors'] )

    return data

def read_wall_3d_hdf5(fname):

    data = dict()
    
    with h5py.File(fname, 'r') as f: # Open for reading
        data['x1x2x3'] = f['/wall/3d/triangles_x1x2x3'][:]
        data['y1y2y3'] = f['/wall/3d/triangles_y1y2y3'][:]
        data['z1z2z3'] = f['/wall/3d/triangles_z1z2z3'][:]
        data[ "flag" ] = f['/wall/3d/triangles_flag'  ][:]
        
    data['n']=[[ data['x1x2x3'].shape[0] ]]

    return data

a5py/a5py/ascot4interface/wall_3d.py
ascot4fusion/ascot5
"""Interface for accessing data in ASCOT5 HDF5 files.
"""
from .bfield  import B_TC, B_GS, B_2DS, B_3DS, B_3DST, B_STS
from .efield  import E_TC, E_1DS, E_3D, E_3DS, E_3DST
from .marker  import Marker, Prt, GC, FL
from .plasma  import plasma_1D, plasma_1DS, plasma_1Dt
from .wall    import wall_2D, wall_3D
from .neutral import N0_1D, N0_3D
from .boozer  import Boozer
from .mhd     import MHD_STAT, MHD_NONSTAT
from .asigma  import Asigma_loc
from .options import Opt
from .nbi     import NBI

from a5py.ascot5io.state      import State
from a5py.ascot5io.orbits     import Orbits
from a5py.ascot5io.transcoef  import Transcoef
from .dist import Dist_5D, Dist_6D, Dist_rho5D, Dist_rho6D, Dist_COM
from .dist import Dist

from .coreio.fileapi import INPUTGROUPS
from .coreio.treeview import RootNode, InputNode, ResultNode
from .coreio.treedata import DataGroup
from a5py.routines.runmixin import RunMixin
from a5py.templates import InputFactory

HDF5TOOBJ = {
    "B_TC" : B_TC, "B_GS" : B_GS, "B_2DS" : B_2DS, "B_3DS" : B_3DS,
    "B_3DST" : B_3DST, "B_STS" : B_STS,
    "E_TC" : E_TC, "E_1DS" : E_1DS, "E_3D" : E_3D, "E_3DS" : E_3DS,
    "E_3DST" : E_3DST,
    "prt" : Prt, "gc" : GC, "fl" : FL,
    "wall_2D" : wall_2D, "wall_3D" : wall_3D,
    "plasma_1D" : plasma_1D, "plasma_1DS" : plasma_1DS,
    "plasma_1Dt" : plasma_1Dt,
    "N0_1D" : N0_1D, "N0_3D" : N0_3D,
    "Boozer" : Boozer, "MHD_STAT" : MHD_STAT, "MHD_NONSTAT" : MHD_NONSTAT,
    "asigma_loc" : Asigma_loc,
    "opt" : Opt,
    "nbi" : NBI,
    "inistate" : State,
    "endstate" : State,
    "orbit" : Orbits,
    "dist5d" : Dist_5D,
    "dist6d" : Dist_6D,
    "distrho5d" : Dist_rho5D,
    "distrho6d" : Dist_rho6D,
    "distcom" : Dist_COM,
    "transcoef": Transcoef
}
"""Dictionary connecting group names in HDF5 file to corresponding data objects.
"""

class Ascot5IO(RootNode):
    """Entry node for accessing data in the HDF5 file.

    Initializing this node builds rest of the tree. This object and its child
    nodes act as container objects for the data in the HDF5 file. At top level
    (this node) run nodes containing simulation results can be accessed as well
    as the parent nodes (bfield, efield, etc.) that in turn contain the actual
    input groups.

    The data can be accessed as

    .. code-block:: python

       data.bfield.B_2DS_1234567890

    or, equivalently,

    .. code-block:: python

       data["bfield"]["B_2DS_1234567890"]

    In each input group, one input is always set as "active" (meaning it would
    be used for the next simulation) and it can be accessed as

    .. code-block:: python

       data.bfield.active

    QID can be used as a reference as well

    .. code-block:: python

       data.bfield.q1234567890

    Run groups are accessed in a similar fashion, e.g.

    .. code-block:: python

       data.run_1234567890 .

    However, the easiest way to access the simulation output is via the methods
    in the RunNode. The active run (by default the most recent simulation) can
    be accessed with

    .. code-block:: python

       data.active

    and its inputs as

    .. code-block:: python

       data.active.bfield

    Most of these examples also work with dictionary-like reference but here we
    use only the attribute-like referencing for brevity.

    You can even use user defined tag taken from description to refer to it

    .. code-block:: python

       data.THATPRLRUN

    However, there are few rules to this:

    - Tag is the first word in description converted to all caps for brevity.
    - Maximum length is ten characters.
    - Only letters and numbers are allowed (all special characters are removed).
    - First character must be a letter.
    - If two fields would have identical tags, they are changed to format
      <tag>_<i>, where i is running index starting from zero (corresponding
      to most recent data).

    Finally, you can print the contents of a node with

    .. code-block:: python

       data.ls()
    """

    def _create_inputgroup(self, path, h5):
        """Create an input group to be added to the treeview.

        Parameters
        ----------
        path : str
            Path to the input group within the HDF5 file.
        h5 : `h5py.File`
            The HDF5 file from which the tree is constructed.

        Returns
        -------
        group : `InputGroup`
            The input group that was created
        """
        return InputGroup(self, path, h5)

    def _create_resultgroup(self, path, h5, inputqids):
        """Create a result group to be added to the treeview.

        Parameters
        ----------
        path : str
            Path to the result node within the HDF5 file.
        h5 : `h5py.File`
            The HDF5 file from which the tree is constructed.
        inputqids : dict [str, str]
            Dictionary containing the name of the input parent group
            (e.g. "bfield") and the QID of the input used for this result.

        Returns
        -------
        group : `ResultGroup`
            The result group that was created.
        """
        return RunGroup(self, path, h5, inputqids)

    def _create_datagroup(self, grouptype, path):
        """Create data group based on the given type.

        Parameters
        ----------
        grouptype : str
            Type of the group as it appears in `HDF5TOOBJ`.
        path : str
            Path to the data in the HDF5 file that corresponds to the group.

        Returns
        -------
        `DataContainer`
            The data group that was created.
        """
        return HDF5TOOBJ[grouptype](self, path)

    def create_input(self, inp, desc=None, activate=None, dryrun=False,
                     **kwargs):
        """Create input and write the data to the HDF5 file.

        This method can be used in two ways to write input data.

        1. From :mod:`ascot5io` choose the input type and find what is requested
           by the corresponding ``write_hdf5`` function to write inputs
           explicitly.

        2. From :mod:`template` choose a template or import, and find what
           parameters are requested to write premade inputs or imported data.

        Parameters
        ----------
        inp : str
            Type of the input e.g. "B_2DS" or name of the template/import
            e.g. "analytical ITER-like".
        desc : str, optional
            Input description.
        activate : bool, optional
            Set created input as active.
        dryrun : bool, optional
            If True, the data is not written to HDF5 but instead it is returned
            as a dictionary.
        **kwargs : dict
            The parameters that are needed to create the requested input type
            or template.

            If input type is given but kwargs is empty, dummy input of given
            type is created.

        Returns
        -------
        data : str or dict
            Name, i.e. "<type>_<qid>", of the new input that was written or the
            data that was created but not written if ``dryrun`` is True.
        """
        if inp in HDF5TOOBJ.keys():
            if len(kwargs) == 0:
                name = HDF5TOOBJ[inp].write_hdf5_dummy(
                self._ascot.file_getpath())
            else:
                name = HDF5TOOBJ[inp].write_hdf5(
                    self._ascot.file_getpath(), **kwargs)
        else:
            gtype, data = InputFactory(self._ascot).construct(inp, **kwargs)
            if dryrun:
                return data
            else:
                name = HDF5TOOBJ[gtype].write_hdf5(
                    self._ascot.file_getpath(), **data)

        self._build(self._ascot.file_getpath())
        if activate:
            self._activate_group(name)
        if desc is not None:
            self._get_group(name).set_desc(desc)
        return name

class InputGroup(InputNode):
    """Node containing input data groups.
    """
    pass

class RunGroup(ResultNode, RunMixin):
    """Node containing results and methods to process them.
    """
    pass

a5py/a5py/ascot5io/__init__.py
ascot4fusion/ascot5
"""Atomic reaction data HDF5 IO.

The data consists either of cross-sections or rate coefficients.
"""
import h5py
import numpy as np

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class Asigma_loc(DataGroup):
    """Local atomic data.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key in ["nreac"]:
                    out[key] = int(out[key])

        return out

    @staticmethod
    def write_hdf5(fn, nreac, z1, a1, z2, a2, reactype, nenergy, energymin,
                   energymax, ndensity, densitymin, densitymax, ntemperature,
                   temperaturemin, temperaturemax, sigma, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Path to hdf5 file.
        nreac : int
            Number of available atomic reactions.
        z1 : array_like (nreac,1)
            Atomic number of test particle.
        a1 : array_like (nreac,1)
            Atomic mass number of test particle.
        z2 : array_like (nreac,1)
            Atomic number of bulk particle.
        a2 : array_like (nreac,1)
            Atomic mass number of bulk particle.
        reactype : array_like (nreac,1)
            Type of atomic reaction.
        nenergy : array_like (nreac,1)
            Number of energy grid points.
        energymin : array_like (nreac,1)
            Energy grid minimum edge [eV].
        energymax : array_like (nreac,1)
            Energy grid maximum edge [eV].
        ndensity : array_like (nreac,1)
            Number of density grid points.
        densitymin : array_like (nreac,1)
            Density grid minimum edge [m^-3].
        densitymax : array_like (nreac,1)
            Density grid maximum edge [m^-3].
        ntemperature : array_like (nreac,1)
            Number of temperature grid points.
        temperaturemin : array_like (nreac,1)
            Temperature grid minimum edge [eV].
        temperaturemax : array_like (nreac,1)
            Temperature grid maximum edge [eV].
        sigma : array_like (1,sum(nenergy[i]*ndensity[i]*ntemperature[i]))
            Reaction cross-section or other probability data [cm^2 or other].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if z1.size != nreac:
            raise ValueError("Invalid number of reactions.")
        n = np.zeros(nreac, dtype=int)
        ntot = 0
        for i in range(0, nreac):
            n[i] = nenergy[i]*ndensity[i]*ntemperature[i]
            ntot += n[i]
        if sigma.shape != (1,ntot):
            raise ValueError("Invalid size for sigma.")
        parent = "asigma"
        group  = "asigma_loc"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset('nreac',          (1,1), data=nreac,      dtype='i4')
            g.create_dataset('z1',         (nreac,1), data=z1,         dtype='i4')
            g.create_dataset('a1',         (nreac,1), data=a1,         dtype='i4')
            g.create_dataset('z2',         (nreac,1), data=z2,         dtype='i4')
            g.create_dataset('a2',         (nreac,1), data=a2,         dtype='i4')
            g.create_dataset('reactype',   (nreac,1), data=reactype,   dtype='i4')
            g.create_dataset('nenergy',    (nreac,1), data=nenergy,    dtype='i4')
            g.create_dataset('energymin',  (nreac,1), data=energymin,  dtype='f8')
            g.create_dataset('energymax',  (nreac,1), data=energymax,  dtype='f8')
            g.create_dataset('ndensity',   (nreac,1), data=ndensity,   dtype='i4')
            g.create_dataset('densitymin', (nreac,1), data=densitymin, dtype='f8')
            g.create_dataset('densitymax', (nreac,1), data=densitymax, dtype='f8')
            g.create_dataset('ntemperature',   (nreac,1), data=ntemperature,
                             dtype='i4')
            g.create_dataset('temperaturemin', (nreac,1), data=temperaturemin,
                             dtype='f8')
            g.create_dataset('temperaturemax', (nreac,1), data=temperaturemax,
                             dtype='f8')
            g.create_dataset('sigma',      (1,ntot),  data=sigma,      dtype='f8')

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        N_reac    = 0
        z_1       = np.zeros(N_reac, dtype=int)
        a_1       = np.zeros(N_reac, dtype=int)
        z_2       = np.zeros(N_reac, dtype=int)
        a_2       = np.zeros(N_reac, dtype=int)
        reac_type = np.zeros(N_reac, dtype=int)
        N_E       = np.zeros(N_reac, dtype=int)
        E_min     = np.zeros(N_reac, dtype=float)
        E_max     = np.zeros(N_reac, dtype=float)
        N_n       = np.zeros(N_reac, dtype=int)
        n_min     = np.zeros(N_reac, dtype=float)
        n_max     = np.zeros(N_reac, dtype=float)
        N_T       = np.zeros(N_reac, dtype=int)
        T_min     = np.zeros(N_reac, dtype=float)
        T_max     = np.zeros(N_reac, dtype=float)
        sigma = np.zeros((1,0))
        return Asigma_loc.write_hdf5(fn, N_reac, z_1, a_1, z_2, a_2, reac_type,
                                     N_E, E_min, E_max, N_n, n_min, n_max,
                                     N_T, T_min, T_max, sigma, desc="DUMMY")

a5py/a5py/ascot5io/asigma.py
ascot4fusion/ascot5
"""Magnetic field input.

Magnetic field is present in every simulation and plays for an important role
in accuracy of the results. High-fidelity magnetic field is essential. Most of
the time magnetic field is the most memory-intensive input (and also most of
the CPU time is spent interpolating it).

The poloidal flux psi is also used as a radial coordinate in simulations e.g. to
interpolate 1D data. More precisely,
rho = sqrt( (psi - psi0) / ( psi1 - psi0 ) ), where psi0 is psi on axis and psi1
is psi on separatrix, is used. Note that psi0 and psi1 must be provided
explicitly in all magnetic field inputs. It is a good idea to include some
padding to psi0, as otherwise rho might have a complex value near separatrix
(which leads to markers being aborted there) if the given psi0 differs from the
actual extrema of the interpolation scheme.
"""
import h5py
import numpy as np

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

import a5py.physlib.analyticequilibrium as psifun

class B_TC(DataGroup):
    """Magnetic field in Cartesian basis for testing purposes.

    This input defines the magnetic field vector on the Cartesian basis, and
    uses constant Jacobian to make the field non-uniform if needed. This field
    is only used to test the validity of the orbit-integrators.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]
        return out

    @staticmethod
    def write_hdf5(fn, bxyz, jacobian, rhoval, psival=None, axisr=1, axisz=0,
                   desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        bxyz : array_like (3,1)
            Magnetic field in cartesian coordinates at origo.
        jacobian : array_like (3,3)
            Magnetic field Jacobian, jacobian[i,j] = dB_i/dx_j.
        rhoval: float
            Constant rho value.
        psival: float, optional
            Constant psi value. If None, same as rhoval.
        axisr: float, optional
            Magnetic axis R coordinate.
        axisz: real, optional
            Magnetic axis z coordinate.
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if bxyz.shape != (3,1) and bxyz.shape != (3,):
            raise ValueError("Invalid shape for Bxyz.")

        parent = "bfield"
        group  = "B_TC"
        gname  = ""

        if psival is None:
            psival = rhoval

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("bxyz",     (3,1), data=bxyz,     dtype="f8")
            g.create_dataset("jacobian", (3,3), data=jacobian, dtype="f8")
            g.create_dataset("rhoval",   (1,),  data=rhoval,   dtype="f8")
            g.create_dataset("psival",   (1,),  data=psival,   dtype="f8")
            g.create_dataset("axisr",    (1,),  data=axisr,    dtype="f8")
            g.create_dataset("axisz",    (1,),  data=axisz,    dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        bxyz = np.array([1,0,3])
        jac  = np.array([ [0,0,0.01], [0,0,0], [0,0,0] ])
        return B_TC.write_hdf5(fn=fn, bxyz=bxyz, jacobian=jac, rhoval=0.5,
                               psival=1.5, axisr=6, axisz=0, desc="DUMMY")

class B_GS(DataGroup):
    """Analytical tokamak field.

    This field can either be axisymmetric or it can include a ripple-like
    perturbation. The field is very fast to interpolate and it is not memory
    intensive. However, it is still intended mostly for testing purposes.

    The axisymmetric field is divergence-free whereas the 3D field is not
    as it included only the toroidal component of the perturbation.

    See:
    A.J. Cerfon, J.P. Freidberg. "One size fits all" analytic solutions to
    the Grad-Shafranov equation. Physics of Plasmas 17 (3) (2010) 032502.
    http://scitation.aip.org/content/aip/journal/pop/17/3/10.1063/1.3328818
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]

        return out

    @staticmethod
    def write_hdf5(fn, r0, z0, bphi0, psimult, coefficients, psi0=None, psi1=0,
                   raxis=None, zaxis=None, nripple=0, a0=2, alpha0=2,
                   delta0=0.05, desc=None):
        """Write input data to the HDF5 file.

        If psi0 is None, magnetic axis location is found numerically and psi0
        value is evaluated at that point. The analytical field is defined by
        assuming midplane is at z=0, but it is moved to z=z0 here.

        Setting number of toroidal field coils to non-zero makes the field 3D.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        r0 : float
            Major radius R coordinate [m].
        z0 : float
            Distance by which midplane is moved from z=0 plane [m].
        bphi0 : float
            Toroidal field at axis [T].
        psimult : float
            Scaling factor for psi.
        coefficients : array_like (13,1)
            Coefficients defining psi [c0, c1, ..., c11, A].
        psi0 : float, optional
            Poloidal flux at magnetic axis.
        psi1 : float, optional
            Poloidal flux at the separatrix.
        raxis : float, optional
            Magnetic axis R coordinate [m].
        zaxis : float, optional
            Magnetic axis z coordinate [m].
        nripple : float, optional
            Number of TF coils.
        a0 : float, optional
            Minor radius. [m]
        alpha0 : float, optional
            Ripple penetration.
        delta0 : float, optional
            Ripple (maximum) strength.
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if coefficients.size != 13:
            raise ValueError("Coefficients invalid.")

        parent = "bfield"
        group  = "B_GS"
        gname  = ""

        c = coefficients # For shorter notation.

        # Search for magnetic axis psi
        if psi0 is None:
            x = psifun.find_axis(r0, z0, c[0], c[1], c[2], c[3], c[4], c[5],
                                 c[6], c[7], c[8], c[9], c[10], c[11], c[12])
            psi0 = psifun.psi0(x[0], x[1], c[0], c[1], c[2], c[3], c[4],
                               c[5], c[6], c[7], c[8], c[9], c[10], c[11],
                               c[12]) * psimult
            psi1  = 0
            raxis = x[0]*r0
            zaxis = x[1]*r0
            psi0 = psi0 - 1e-8 if psi0 < psi1 else psi0 + 1e-8

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("r0",           (1,),   data=r0,      dtype='f8')
            g.create_dataset("z0",           (1,),   data=z0,      dtype='f8')
            g.create_dataset("raxis",        (1,),   data=raxis,   dtype='f8')
            g.create_dataset("zaxis",        (1,),   data=zaxis,   dtype='f8')
            g.create_dataset("bphi0",        (1,),   data=bphi0,   dtype='f8')
            g.create_dataset("psi0",         (1,),   data=psi0,    dtype='f8')
            g.create_dataset("psi1",         (1,),   data=psi1,    dtype='f8')
            g.create_dataset("psimult",      (1,),   data=psimult, dtype='f8')
            g.create_dataset("nripple",      (1,),   data=nripple, dtype='i8')
            g.create_dataset("a0",           (1,),   data=a0,      dtype='f8')
            g.create_dataset("alpha0",       (1,),   data=alpha0,  dtype='f8')
            g.create_dataset("delta0",       (1,),   data=delta0,  dtype='f8')
            g.create_dataset("coefficients", (13,1), data=coefficients,
                             dtype='f8')

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        coefficients = np.array([ 2.218e-02, -1.288e-01, -4.177e-02, -6.227e-02,
                                  6.200e-03, -1.205e-03, -3.701e-05,  0,
                                  0,          0,          0,          0,
                                  -0.155])
        return B_GS.write_hdf5(fn, r0=6.2, z0=0, bphi0=5.3, psimult=200,
                               coefficients=coefficients, nripple=1, a0=2,
                               alpha0=2, delta0=0.05, desc="DUMMY")


class B_2DS(DataGroup):
    """Axisymmetric tokamak field interpolated with splines.

    This field is suitable for simulations where toroidal ripple or other
    3D perturbations are not relevant (for MHD one can use the dedicated input).
    Using axisymmetric field is **much** faster than 3D fields, so using this
    input when applicable is strongly recommended.

    The input consists of BR, Bphi, Bz, and psi tabulated on an uniform
    (R,z) grid which are interpolated with cubic splines in simulation.
    During the simulation, the total BR (and Bz) are computed as a sum of input
    BR (Bz) and the component calculated from the gradient of psi. In other
    words, the equilibrium BR and Bz are already contained in psi so in most
    cases the input BR and Bz should be set to zero. If for some reason psi has
    a poor quality, it can be scaled to an insignificant value and the field can
    be provided explicitly in BR and Bz. However, in such case the field is no
    longer divergence-free.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]

        out["psi"]  = np.transpose(out["psi"])
        out["br"]   = np.transpose(out["br"])
        out["bphi"] = np.transpose(out["bphi"])
        out["bz"]   = np.transpose(out["bz"])
        return out

    @staticmethod
    def write_hdf5(fn, rmin, rmax, nr, zmin, zmax, nz,
                   axisr, axisz, psi, psi0, psi1,
                   br, bphi, bz, desc=None):
        """Write input data to the HDF5 file.

        Note that br and bz should not include the equilibrium component of the
        magnetic field as that is calculated from psi by ASCOT5 during the
        simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        rmin : float
            R grid min edge [m].
        rmax : float
            R grid max edge [m].
        nr : int
            Number of R grid points.
        zmin : float
            z grid min edge [m].
        zmax : float
            z grid max edge [m].
        nz : int
            Number of z grid points.
        axisr : float
            Magnetic axis R coordinate [m].
        axisz : float
            Magnetic axis z coordinate [m].
        psi0 : float
            On-axis poloidal flux value [Vs/m].
        psi1 : float
            Separatrix poloidal flux value [Vs/m].
        psi : array_like (nr, nz)
            Poloidal flux values on the Rz grid [Vs/m].
        br : array_like (nr,nz)
            Magnetic field R component (excl. equilibrium comp.) on Rz grid [T].
        bphi : array_like (nr,nz)
            Magnetic field phi component on Rz grid [T].
        bz : array_like (nr,nz)
            Magnetic field z component (excl. equilibrium comp.) onRz grid [T].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """

        if psi.shape  != (nr,nz):
            raise ValueError("Inconsistent shape for psi.")
        if br.shape   != (nr,nz):
            raise ValueError("Inconsistent shape for br.")
        if bphi.shape != (nr,nz):
            raise ValueError("Inconsistent shape for bphi.")
        if bz.shape   != (nr,nz):
            raise ValueError("Inconsistent shape for bz.")

        psi  = np.transpose(psi)
        br   = np.transpose(br)
        bphi = np.transpose(bphi)
        bz   = np.transpose(bz)

        parent = "bfield"
        group  = "B_2DS"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("rmin",  (1,), data=rmin,  dtype="f8")
            g.create_dataset("rmax",  (1,), data=rmax,  dtype="f8")
            g.create_dataset("nr",    (1,), data=nr,    dtype="i4")
            g.create_dataset("zmin",  (1,), data=zmin,  dtype="f8")
            g.create_dataset("zmax",  (1,), data=zmax,  dtype="f8")
            g.create_dataset("nz",    (1,), data=nz,    dtype="i4")
            g.create_dataset("axisr", (1,), data=axisr, dtype="f8")
            g.create_dataset("axisz", (1,), data=axisz, dtype="f8")
            g.create_dataset("psi0",  (1,), data=psi0,  dtype="f8")
            g.create_dataset("psi1",  (1,), data=psi1,  dtype="f8")

            g.create_dataset("psi",  (nz, nr), data=psi,  dtype="f8")
            g.create_dataset("br",   (nz, nr), data=br,   dtype="f8")
            g.create_dataset("bphi", (nz, nr), data=bphi, dtype="f8")
            g.create_dataset("bz",   (nz, nr), data=bz,   dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output is an ITER-like but circular equilibrium.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        # ITER-like but circular equilibrium
        coefficients = np.array([ 2.218e-02, -1.288e-01, -4.177e-02, -6.227e-02,
                                  6.200e-03, -1.205e-03, -3.701e-05,  0,
                                  0,          0,          0,          0,
                                  -0.155])
        gs = {"rmin":1, "rmax":6, "nr":50, "zmin":-4, "zmax":4, "nz":100,
              "r0":6.2, "z0":0, "bphi0":5.3, "psimult":200,
              "coefficients":coefficients}
        return B_2DS.write_hdf5(fn=fn, desc="DUMMY", **B_2DS.convert_B_GS(**gs))

    @staticmethod
    def convert_B_GS(rmin, rmax, nr, zmin, zmax, nz, **kwargs):
        """Convert :class:`B_GS` input to `B_2DS` input.

        Parameters
        ----------
        rmin : float
            R grid min edge [m].
        rmax : float
            R grid max edge [m].
        nr : int
            Number of R grid points.
        zmin : float
            z grid min edge [m].
        zmax : float
            z grid max edge [m].
        nz : int
            Number of z grid points.
        **kwargs
            Arguments passed to :meth:`B_GS.write_hdf5` excluding ``fn`` and
            ``desc``.

        Returns
        -------
        out : dict
            :class:`B_GS` converted as an input for :meth:`write_hdf5`.
        """
        rgrid = np.linspace(rmin, rmax, nr)
        zgrid = np.linspace(zmin, zmax, nz)
        zg, rg = np.meshgrid(zgrid, rgrid);

        c = kwargs["coefficients"] # For shorter notation.
        psirz = kwargs["psimult"] * psifun.psi0(
            rg/kwargs["r0"], zg/kwargs["r0"], c[0], c[1], c[2], c[3], c[4],
            c[5], c[6], c[7], c[8], c[9], c[10], c[11], c[12])

        br = np.zeros((nr,nz))
        bz = np.zeros((nr,nz))
        bphi = ( kwargs["r0"] / rg ) * kwargs["bphi0"]

        # search for magnetic axis if not given
        if not "psi0" in kwargs:
            x = psifun.find_axis(kwargs["r0"], kwargs["z0"], c[0], c[1], c[2],
                                 c[3], c[4], c[5],
                                 c[6], c[7], c[8], c[9], c[10], c[11], c[12])
            psi0 = kwargs["psimult"]*psifun.psi0(
                x[0], x[1], c[0], c[1], c[2], c[3], c[4],
                c[5], c[6], c[7], c[8], c[9], c[10], c[11], c[12])
            r0 = x[0]*kwargs["r0"]
            z0 = x[1]*kwargs["r0"]
            psi1 = 0
            # Padding
            psi0 = psi0 - 1e-4 if psi0 < psi1 else psi0 + 1e-4

        return {"rmin" : rmin, "rmax" : rmax, "nr" : nr, "zmin" : zmin,
                "zmax" : zmax, "nz" : nz, "axisr" : r0, "axisz" : z0,
                "psi" : psirz, "psi0" : psi0, "psi1" : psi1, "br" : br,
                "bphi" : bphi, "bz" : bz}

    @staticmethod
    def convert_B_3DS(**kwargs):
        """Convert :class:`B_3DS` input to `B_2DS` input.

        This function takes a toroidal average of Bphi and sets BR and Bz to
        zero.

        Parameters
        ----------
        **kwargs
            Arguments passed to :meth:`B_3DS.write_hdf5` excluding ``fn`` and
            ``desc``.

        Returns
        -------
        out : dict
            :class:`B_3DS` converted as an input for :meth:`write_hdf5`.
        """
        bphi = np.mean(kwargs["bphi"], axis=1)
        br   = np.mean(kwargs["bphi"], axis=1) * 0
        bz   = np.mean(kwargs["bphi"], axis=1) * 0
        return {
            "rmin":kwargs["b_rmin"], "rmax":kwargs["b_rmax"],
            "nr":kwargs["b_nr"], "zmin":kwargs["b_zmin"],
            "zmax":kwargs["b_zmax"], "nz":kwargs["b_nz"],
            "axisr":kwargs["axisr"], "axisz":kwargs["axisz"],
            "psi0":kwargs["psi0"], "psi1":kwargs["psi1"],
            "psi":kwargs["psi"], "br":br, "bz":bz,
            "bphi":bphi}

class B_3DS(DataGroup):
    """Non-axisymmetric tokamak field.

    The input consists of BR, Bphi, and Bz tabulated on an uniform cylindrical
    grid and psi tabulated on (R,z) grid. Data is interpolated with cubic
    splines during the simulation, and the total BR (and Bz) are computed as
    a sum of input BR (Bz) and the component calculated from the gradient of
    psi. In other words, the equilibrium BR and Bz are already contained in psi
    so BR and Bz should only contain the components from 3D perturbation.

    This input is suitable for artificial 3D perturbations. For MHD, consider
    the dedicated input.

    Since this input interpolates BR and Bz with splines, the resulting B is
    not divergence free. This can be amended by using a dense grid for magnetic
    field evaluation. Note however that it can become memory-intensive.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]

        out["psi"]  = np.transpose(out["psi"])
        out["br"]   = np.transpose(out["br"],   (2,1,0))
        out["bphi"] = np.transpose(out["bphi"], (2,1,0))
        out["bz"]   = np.transpose(out["bz"],   (2,1,0))
        return out

    @staticmethod
    def write_hdf5(fn, b_rmin, b_rmax, b_nr, b_zmin, b_zmax, b_nz,
                   b_phimin, b_phimax, b_nphi,
                   axisr, axisz, psi, psi0, psi1, br, bphi, bz,
                   psi_rmin=None, psi_rmax=None, psi_nr=None,
                   psi_zmin=None, psi_zmax=None, psi_nz=None, desc=None):
        """Write input data to the HDF5 file.

        It is possible to use different (R,z) grids for psi and magnetic field
        components by giving the (R,z)-grid for psi separately. 3D data can be
        memory intensive which necessitates sparser grid for B components, but
        psi can still be evaluated on a dense grid.

        The toroidal angle phi is treated as a periodic coordinate, meaning
        ``A(phi=phimin) == A(phi=phimax)``. However, the phi grid, where input
        arrays are tabulated, is ``linspace(phimin, phimax, nphi+1)[:-1]``
        to avoid storing duplicate data.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        b_rmin : float
            Magnetic field data R grid min edge [m].
        b_rmax : float
            Magnetic field data R grid max edge [m].
        b_nr : int
            Number of R grid points in magnetic field data.
        b_zmin : float
            Magnetic field data z grid min edge [m].
        b_zmax : float
            Magnetic field data z grid max edge [m].
        b_nz : int
            Number of z grid points in magnetic field data.
        b_phimin : float
            Beginning of the toroidal period [deg].
        b_phimax : float
            End of the toroidal period [deg].
        b_nphi : int
            Number of phi grid points in magnetic field data.
        axisr : float
            Magnetic axis R coordinate [m].
        axisz : float
            Magnetic axis z coordinate [m].
        psi0 : float
            On-axis poloidal flux value [Vs/m].
        psi1 : float
            Separatrix poloidal flux value [Vs/m].
        psi : array_like (nr, nz)
            Poloidal flux values on the Rz grid [Vs/m].
        br : array_like (nr,nphi,nz)
            Magnetic field R component (excl. equilibrium comp.) on Rz grid [T].
        bphi : array_like (nr,nphi,nz)
            Magnetic field phi component on Rz grid [T].
        bz : array_like (nr,nphi,nz)
            Magnetic field z component (excl. equilibrium comp.) onRz grid [T].
        psi_rmin : float, optional
            Psi data R grid min edge [m].
        psi_rmax : float, optional
            Psi data R grid max edge [m].
        psi_nr : int, optional
            Number of R grid points in psi data.
        psi_zmin : float, optional
            Psi data z grid min edge [m].
        psi_zmax : float, optional
            Psi data z grid max edge [m].
        psi_nz : int, optional
            Number of z grid points in psi data.
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        parent = "bfield"
        group  = "B_3DS"
        gname  = ""

        # Define psigrid to be same as Bgrid if not stated otherwise.
        if(psi_rmin is None or psi_rmax is None or psi_nr is None or
           psi_zmin is None or psi_zmax is None or psi_nz is None):
            psi_rmin = b_rmin
            psi_rmax = b_rmax
            psi_nr   = b_nr
            psi_zmin = b_zmin
            psi_zmax = b_zmax
            psi_nz   = b_nz

        if psi.shape  != (psi_nr,psi_nz):
            raise ValueError("Inconsistent shape for psi.")
        if br.shape   != (b_nr,b_nphi,b_nz):
            raise ValueError("Inconsistent shape for br.")
        if bphi.shape != (b_nr,b_nphi,b_nz):
            raise ValueError("Inconsistent shape for bphi.")
        if bz.shape   != (b_nr,b_nphi,b_nz):
            raise ValueError("Inconsistent shape for bz.")

        psi  = np.transpose(psi)
        br   = np.transpose(br,   (2,1,0))
        bphi = np.transpose(bphi, (2,1,0))
        bz   = np.transpose(bz,   (2,1,0))

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("b_rmin",   (1,), data=b_rmin,   dtype="f8")
            g.create_dataset("b_rmax",   (1,), data=b_rmax,   dtype="f8")
            g.create_dataset("b_nr",     (1,), data=b_nr,     dtype="i4")
            g.create_dataset("b_phimin", (1,), data=b_phimin, dtype="f8")
            g.create_dataset("b_phimax", (1,), data=b_phimax, dtype="f8")
            g.create_dataset("b_nphi",   (1,), data=b_nphi,   dtype="i4")
            g.create_dataset("b_zmin",   (1,), data=b_zmin,   dtype="f8")
            g.create_dataset("b_zmax",   (1,), data=b_zmax,   dtype="f8")
            g.create_dataset("b_nz",     (1,), data=b_nz,     dtype="i4")
            g.create_dataset("psi_rmin", (1,), data=psi_rmin, dtype="f8")
            g.create_dataset("psi_rmax", (1,), data=psi_rmax, dtype="f8")
            g.create_dataset("psi_nr",   (1,), data=psi_nr,   dtype="i4")
            g.create_dataset("psi_zmin", (1,), data=psi_zmin, dtype="f8")
            g.create_dataset("psi_zmax", (1,), data=psi_zmax, dtype="f8")
            g.create_dataset("psi_nz",   (1,), data=psi_nz,   dtype="i4")
            g.create_dataset("axisr",    (1,), data=axisr,    dtype="f8")
            g.create_dataset("axisz",    (1,), data=axisz,    dtype="f8")
            g.create_dataset("psi0",     (1,), data=psi0,     dtype="f8")
            g.create_dataset("psi1",     (1,), data=psi1,     dtype="f8")

            g.create_dataset("psi",  (psi_nz,psi_nr),    data=psi,  dtype="f8")
            g.create_dataset("br",   (b_nz,b_nphi,b_nr), data=br,   dtype="f8")
            g.create_dataset("bphi", (b_nz,b_nphi,b_nr), data=bphi, dtype="f8")
            g.create_dataset("bz",   (b_nz,b_nphi,b_nr), data=bz,   dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        coefficients = np.array([ 2.218e-02, -1.288e-01, -4.177e-02, -6.227e-02,
                                  6.200e-03, -1.205e-03, -3.701e-05,  0,
                                  0,          0,          0,          0,
                                  -0.155])
        gs = {"rmin":1, "rmax":6, "nr":50, "zmin":-4, "zmax":4, "nz":100,
              "phimin":0, "phimax":360, "nphi":100,
              "r0":6.2, "z0":0, "bphi0":5.3, "psimult":200,
              "coefficients":coefficients, "nripple":18, "a0":2, "alpha0":2,
              "delta0":0.05}
        return B_3DS.write_hdf5(fn=fn, desc="DUMMY", **B_3DS.convert_B_GS(**gs))

    @staticmethod
    def convert_B_GS(rmin, rmax, nr, zmin, zmax, nz, phimin, phimax, nphi,
                     **kwargs):
        """Convert :class:`B_GS` input to `B_3DS` input.

        The resulting field is far from being divergence-free as the
        perturbation BR and Bz components are omitted.

        Parameters
        ----------
        rmin : float
            Magnetic field data R grid min edge [m].
        rmax : float
            Magnetic field data R grid max edge [m].
        nr : int
            Number of R grid points in magnetic field data.
        zmin : float
            Magnetic field data z grid min edge [m].
        zmax : float
            Magnetic field data z grid max edge [m].
        nz : int
            Number of z grid points in magnetic field data.
        phimin : float
            Beginning of the toroidal period [deg].
        phimax : float
            End of the toroidal period [deg].
        nphi : int
            Number of phi grid points in magnetic field data.
        **kwargs
            Arguments passed to :meth:`B_GS.write_hdf5` excluding ``fn`` and
            ``desc``.

        Returns
        -------
        out : dict
            :class:`B_GS` converted as an input for :meth:`write_hdf5`.
        """
        rgrid = np.linspace(rmin, rmax, nr)
        zgrid = np.linspace(zmin, zmax, nz)

        zg, rg = np.meshgrid(zgrid, rgrid);

        c = kwargs["coefficients"] # For shorter notation.
        psirz = kwargs["psimult"]*psifun.psi0(
            rg/kwargs["r0"],zg/kwargs["r0"],
            c[0],c[1],c[2],c[3],c[4],c[5],
            c[6],c[7],c[8],c[9],c[10],c[11],c[12])

        br = np.zeros((nr, nphi, nz))
        bz = np.zeros((nr, nphi, nz))
        bphi = np.zeros((nr, nphi, nz))

        # Ripple
        radius = np.sqrt( ( rg - kwargs["r0"] )**2 + ( zg - kwargs["z0"] )**2 )
        theta = np.arctan2( zg - kwargs["z0"], rg - kwargs["r0"] )
        delta = kwargs["delta0"] * np.exp(-0.5*theta**2) \
            * np.power( radius / kwargs["a0"], kwargs["alpha0"] )

        phigrid = np.linspace(phimin, phimax, nphi+1)[:-1]
        for i in range(nphi):
            cos = np.cos( kwargs["nripple"] * phigrid[i] * np.pi / 180 )
            bphi[:,i,:] = ( kwargs["r0"] / rg ) * kwargs["bphi0"] \
                * ( 1 + delta * cos )

        # search for magnetic axis if not given
        if not "psi0" in kwargs:
            x = psifun.find_axis(kwargs["r0"], kwargs["z0"], c[0], c[1], c[2],
                                 c[3], c[4], c[5],
                                 c[6], c[7], c[8], c[9], c[10], c[11], c[12])
            psi0 = kwargs["psimult"]*psifun.psi0(
                x[0], x[1], c[0], c[1], c[2], c[3], c[4],
                c[5], c[6], c[7], c[8], c[9], c[10], c[11], c[12]) # At axis.
            raxis = x[0]*kwargs["r0"]
            zaxis = x[1]*kwargs["r0"]

            psi1 = 0

        return {"b_rmin":rmin, "b_rmax":rmax, "b_nr":nr, "b_zmin":zmin,
                "b_zmax":zmax, "b_nz":nz, "b_phimin":phimin, "b_phimax":phimax,
                "b_nphi":nphi, "axisr":raxis, "axisz":zaxis, "psi":psirz,
                "psi0":psi0, "psi1":psi1, "br":br, "bphi":bphi, "bz":bz}

class B_3DST(DataGroup):
    """Time-dependent 3D tokamak field interpolated with splines.

    This input is equivalent to :class:`B_3DS` except that the data includes
    fourth dimension for time.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]

        out["psi"]  = np.transpose(out["psi"])
        out["br"]   = np.transpose(out["br"],   (3,2,1,0))
        out["bphi"] = np.transpose(out["bphi"], (3,2,1,0))
        out["bz"]   = np.transpose(out["bz"],   (3,2,1,0))
        return out

    @staticmethod
    def write_hdf5(fn, b_rmin, b_rmax, b_nr, b_zmin, b_zmax, b_nz,
                   b_phimin, b_phimax, b_nphi, b_tmin, b_tmax, b_nt,
                   axisr, axisz, psi, psi0, psi1, br, bphi, bz,
                   psi_rmin=None, psi_rmax=None, psi_nr=None,
                   psi_zmin=None, psi_zmax=None, psi_nz=None, desc=None):
        """Write input data to the HDF5 file.

        It is possible to use different (R,z) grids for psi and magnetic field
        components by giving the (R,z)-grid for psi separately. 3D data can be
        memory intensive which necessitates sparser grid for B components, but
        psi can still be evaluated on a dense grid.

        The toroidal angle phi is treated as a periodic coordinate, meaning
        ``A(phi=phimin) == A(phi=phimax)``. However, the phi grid, where input
        arrays are tabulated, is ``linspace(phimin, phimax, nphi+1)[:-1]``
        to avoid storing duplicate data.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        b_rmin : float
            Magnetic field data R grid min edge [m].
        b_rmax : float
            Magnetic field data R grid max edge [m].
        b_nr : int
            Number of R grid points in magnetic field data.
        b_zmin : float
            Magnetic field data z grid min edge [m].
        b_zmax : float
            Magnetic field data z grid max edge [m].
        b_nz : int
            Number of z grid points in magnetic field data.
        b_phimin : float
            Beginning of the toroidal period [deg].
        b_phimax : float
            End of the toroidal period [deg].
        b_nphi : int
            Number of phi grid points in magnetic field data.
        b_tmin : float
            Magnetic field data time grid min edge [s].
        b_tmax : float
            Magnetic field data time grid max edge [s].
        b_nt : int
            Number of t grid points in magnetic field data.
        axisr : float
            Magnetic axis R coordinate [m].
        axisz : float
            Magnetic axis z coordinate [m].
        psi0 : float
            On-axis poloidal flux value [Vs/m].
        psi1 : float
            Separatrix poloidal flux value [Vs/m].
        psi : array_like (nr, nz)
            Poloidal flux values on the Rz grid [Vs/m].
        br : array_like (nr,nphi,nz,nt)
            Magnetic field R component (excl. equilibrium comp.).
        bphi : array_like (nr,nphi,nz,nt)
            Magnetic field phi component on Rz grid [T].
        bz : array_like (nr,nphi,nz,nt)
            Magnetic field z component (excl. equilibrium comp.).
        psi_rmin : float, optional
            Psi data R grid min edge [m].
        psi_rmax : float, optional
            Psi data R grid max edge [m].
        psi_nr : int, optional
            Number of R grid points in psi data.
        psi_zmin : float, optional
            Psi data z grid min edge [m].
        psi_zmax : float, optional
            Psi data z grid max edge [m].
        psi_nz : int, optional
            Number of z grid points in psi data.
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        parent = "bfield"
        group  = "B_3DST"
        gname  = ""

        # Define psigrid to be same as Bgrid if not stated otherwise.
        if(psi_rmin is None or psi_rmax is None or psi_nr is None or
           psi_zmin is None or psi_zmax is None or psi_nz is None):
            psi_rmin = b_rmin
            psi_rmax = b_rmax
            psi_nr   = b_nr
            psi_zmin = b_zmin
            psi_zmax = b_zmax
            psi_nz   = b_nz

        if psi.shape  != (psi_nr,psi_nz):
            raise ValueError("Inconsistent shape for psi.")
        if br.shape   != (b_nr,b_nphi,b_nz,b_nt):
            raise ValueError("Inconsistent shape for br.")
        if bphi.shape != (b_nr,b_nphi,b_nz,b_nt):
            raise ValueError("Inconsistent shape for bphi.")
        if bz.shape   != (b_nr,b_nphi,b_nz,b_nt):
            raise ValueError("Inconsistent shape for bz.")

        psi  = np.transpose(psi)
        br   = np.transpose(br,   (3,2,1,0))
        bphi = np.transpose(bphi, (3,2,1,0))
        bz   = np.transpose(bz,   (3,2,1,0))

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("b_rmin",   (1,), data=b_rmin,   dtype="f8")
            g.create_dataset("b_rmax",   (1,), data=b_rmax,   dtype="f8")
            g.create_dataset("b_nr",     (1,), data=b_nr,     dtype="i4")
            g.create_dataset("b_phimin", (1,), data=b_phimin, dtype="f8")
            g.create_dataset("b_phimax", (1,), data=b_phimax, dtype="f8")
            g.create_dataset("b_nphi",   (1,), data=b_nphi,   dtype="i4")
            g.create_dataset("b_zmin",   (1,), data=b_zmin,   dtype="f8")
            g.create_dataset("b_zmax",   (1,), data=b_zmax,   dtype="f8")
            g.create_dataset("b_nz",     (1,), data=b_nz,     dtype="i4")
            g.create_dataset("b_tmin",   (1,), data=b_tmin,   dtype="f8")
            g.create_dataset("b_tmax",   (1,), data=b_tmax,   dtype="f8")
            g.create_dataset("b_nt",     (1,), data=b_nt,     dtype="i4")
            g.create_dataset("psi_rmin", (1,), data=psi_rmin, dtype="f8")
            g.create_dataset("psi_rmax", (1,), data=psi_rmax, dtype="f8")
            g.create_dataset("psi_nr",   (1,), data=psi_nr,   dtype="i4")
            g.create_dataset("psi_zmin", (1,), data=psi_zmin, dtype="f8")
            g.create_dataset("psi_zmax", (1,), data=psi_zmax, dtype="f8")
            g.create_dataset("psi_nz",   (1,), data=psi_nz,   dtype="i4")
            g.create_dataset("axisr",    (1,), data=axisr,    dtype="f8")
            g.create_dataset("axisz",    (1,), data=axisz,    dtype="f8")
            g.create_dataset("psi0",     (1,), data=psi0,     dtype="f8")
            g.create_dataset("psi1",     (1,), data=psi1,     dtype="f8")

            g.create_dataset("psi",  (psi_nz, psi_nr), data=psi,  dtype="f8")
            g.create_dataset("br",   (b_nt, b_nz, b_nphi, b_nr), data=br,
                             dtype="f8")
            g.create_dataset("bphi", (b_nt, b_nz, b_nphi, b_nr), data=bphi,
                             dtype="f8")
            g.create_dataset("bz",   (b_nt, b_nz, b_nphi, b_nr), data=bz,
                             dtype="f8")
        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        B_3DST.write_hdf5(
            fn=fn, b_rmin=4, b_rmax=8, b_nr=3, b_zmin=-2, b_zmax=2,
            b_nz=3, b_phimin=0, b_phimax=360, b_nphi=3, b_tmin=1,
            b_tmax=1, b_nt=3, axisr=4, axisz=0, psi0=0,
            psi1=1, br=np.zeros((3,3,3,3)), bphi=np.ones((3,3,3,3)),
            bz=np.zeros((3,3,3,3)), psi=0.5*np.ones((3,3)),
            psi_rmin=4, psi_rmax=8, psi_nr=3,
            psi_zmin=-2, psi_zmax=2, psi_nz=3, desc="DUMMY")

class B_STS(DataGroup):
    """Stellarator magnetic field.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]

        out["psi"]  = np.transpose(out["psi"],  (2,1,0))
        out["br"]   = np.transpose(out["br"],   (2,1,0))
        out["bphi"] = np.transpose(out["bphi"], (2,1,0))
        out["bz"]   = np.transpose(out["bz"],   (2,1,0))
        return out

    @staticmethod
    def write_hdf5(fn, b_rmin, b_rmax, b_nr, b_zmin, b_zmax, b_nz,
                   b_phimin, b_phimax, b_nphi, psi0, psi1,
                   br, bphi, bz, psi,
                   axis_phimin, axis_phimax, axis_nphi, axisr, axisz,
                   psi_rmin=None, psi_rmax=None, psi_nr=None,
                   psi_zmin=None, psi_zmax=None, psi_nz=None,
                   psi_phimin=None, psi_phimax=None, psi_nphi=None,
                   desc=None):
        """Write input data to the HDF5 file.

        It is possible to use different (R,z) grids for psi and magnetic field
        components by giving the (R,z)-grid for psi separately. 3D data can be
        memory intensive which necessitates sparser grid for B components, but
        psi can still be evaluated on a dense grid.

        The toroidal angle phi is treated as a periodic coordinate, meaning
        ``A(phi=phimin) == A(phi=phimax)``. However, the phi grid, where input
        arrays are tabulated, is ``linspace(phimin, phimax, nphi+1)[:-1]``
        to avoid storing duplicate data.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        b_rmin : float
            Magnetic field data R grid min edge [m].
        b_rmax : float
            Magnetic field data R grid max edge [m].
        b_nr : int
            Number of R grid points in magnetic field data.
        b_zmin : float
            Magnetic field data z grid min edge [m].
        b_zmax : float
            Magnetic field data z grid max edge [m].
        b_nz : int
            Number of z grid points in magnetic field data.
        b_phimin : float
            Beginning of the toroidal period [deg].
        b_phimax : float
            End of the toroidal period [deg].
        b_nphi : int
            Number of phi grid points in magnetic field data.
        axis_phimin : float
            Magnetic axis phi grid min value [deg].
        axis_phimax : float
            Magnetic axis phi grid max value [deg].
        axis_nphi : float
            Number of points in magnetic axis phi grid.
        axisr : float
            Magnetic axis R coordinates on axis phi grid [m].
        axisz : float
            Magnetic axis z coordinates on axis phi grid [m].
        psi0 : float
            On-axis poloidal flux value [Vs/m].
        psi1 : float
            Separatrix poloidal flux value [Vs/m].
        psi : array_like (psi_nr,psi_nphi,psi_nz)
            Poloidal flux values on the Rz grid [Vs/m].
        br : array_like (b_nr,b_nphi,b_nz)
            Magnetic field R component (excl. equilibrium comp.) on Rz grid [T].
        bphi : array_like (b_nr,b_nphi,b_nz)
            Magnetic field phi component on Rz grid [T].
        bz : array_like (b_nr,b_nphi,b_nz)
            Magnetic field z component (excl. equilibrium comp.) onRz grid [T].
        psi_rmin : float, optional
            Psi data R grid min edge [m].
        psi_rmax : float, optional
            Psi data R grid max edge [m].
        psi_nr : int, optional
            Number of R grid points in psi data.
        psi_zmin : float, optional
            Psi data z grid min edge [m].
        psi_zmax : float, optional
            Psi data z grid max edge [m].
        psi_nz : int, optional
            Number of z grid points in psi data.
        psi_phimin : float, optional
            Psi data beginning of the toroidal period.
        psi_phimax : float, optional
            Psi data end of the toroidal period.
        psi_nphi : int, optional
            Number of phi grid points in psi data.
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        parent = "bfield"
        group  = "B_STS"
        gname  = ""

        # Define psigrid to be same as Bgrid if not stated otherwise.
        if(psi_rmin is None or psi_rmax is None or psi_nr is None
           or psi_phimin is None or psi_phimax is None or psi_nphi is None
           or psi_zmin is None or psi_zmax is None or psi_nz is None):
            psi_rmin   = b_rmin
            psi_rmax   = b_rmax
            psi_nr     = b_nr
            psi_phimin = b_phimin
            psi_phimax = b_phimax
            psi_nphi   = b_nphi
            psi_zmin   = b_zmin
            psi_zmax   = b_zmax
            psi_nz     = b_nz

        if psi.shape  != (psi_nr,psi_nphi,psi_nz):
            raise ValueError("Inconsistent shape foor psi.")
        if br.shape   != (b_nr,b_nphi,b_nz):
            raise ValueError("Inconsistent shape foor br.")
        if bphi.shape != (b_nr,b_nphi,b_nz):
            raise ValueError("Inconsistent shape foor bphi.")
        if bz.shape   != (b_nr,b_nphi,b_nz):
            raise ValueError("Inconsistent shape foor bz.")

        psi  = np.transpose(psi,  (2,1,0))
        br   = np.transpose(br,   (2,1,0))
        bphi = np.transpose(bphi, (2,1,0))
        bz   = np.transpose(bz,   (2,1,0))

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("b_rmin",      (1,), data=b_rmin,      dtype="f8")
            g.create_dataset("b_rmax",      (1,), data=b_rmax,      dtype="f8")
            g.create_dataset("b_nr",        (1,), data=b_nr,        dtype="i4")
            g.create_dataset("b_phimin",    (1,), data=b_phimin,    dtype="f8")
            g.create_dataset("b_phimax",    (1,), data=b_phimax,    dtype="f8")
            g.create_dataset("b_nphi",      (1,), data=b_nphi,      dtype="i4")
            g.create_dataset("b_zmin",      (1,), data=b_zmin,      dtype="f8")
            g.create_dataset("b_zmax",      (1,), data=b_zmax,      dtype="f8")
            g.create_dataset("b_nz",        (1,), data=b_nz,        dtype="i4")
            g.create_dataset("psi_rmin",    (1,), data=psi_rmin,    dtype="f8")
            g.create_dataset("psi_rmax",    (1,), data=psi_rmax,    dtype="f8")
            g.create_dataset("psi_nr",      (1,), data=psi_nr,      dtype="i4")
            g.create_dataset("psi_phimin",  (1,), data=psi_phimin,  dtype="f8")
            g.create_dataset("psi_phimax",  (1,), data=psi_phimax,  dtype="f8")
            g.create_dataset("psi_nphi",    (1,), data=psi_nphi,    dtype="i4")
            g.create_dataset("psi_zmin",    (1,), data=psi_zmin,    dtype="f8")
            g.create_dataset("psi_zmax",    (1,), data=psi_zmax,    dtype="f8")
            g.create_dataset("psi_nz",      (1,), data=psi_nz,      dtype="i4")
            g.create_dataset("axis_phimin", (1,), data=axis_phimin, dtype="f8")
            g.create_dataset("axis_phimax", (1,), data=axis_phimax, dtype="f8")
            g.create_dataset("axis_nphi",   (1,), data=axis_nphi,   dtype="i4")
            g.create_dataset("psi0",        (1,), data=psi0,        dtype="f8")
            g.create_dataset("psi1",        (1,), data=psi1,        dtype="f8")

            g.create_dataset("axisr", (axis_nphi,), data=axisr, dtype="f8")
            g.create_dataset("axisz", (axis_nphi,), data=axisz, dtype="f8")

            g.create_dataset("br",         (b_nz,b_nphi,b_nr),       data=br,
                             dtype="f8")
            g.create_dataset("bphi",       (b_nz,b_nphi,b_nr),       data=bphi,
                             dtype="f8")
            g.create_dataset("bz",         (b_nz,b_nphi,b_nr),       data=bz,
                             dtype="f8")
            g.create_dataset("psi",        (psi_nz,psi_nphi,psi_nr), data=psi,
                             dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        B_STS.write_hdf5(fn=fn, b_rmin=4, b_rmax=8, b_nr=3, b_zmin=-2, b_zmax=2,
                         b_nz=3, b_phimin=0, b_phimax=360, b_nphi=3, psi0=0,
                         psi1=1, br=np.zeros((3,3,3)), bphi=np.ones((3,3,3)),
                         bz=np.zeros((3,3,3)), psi=0.5*np.ones((3,3,3)),
                         axis_phimin=0, axis_phimax=360, axis_nphi=3,
                         axisr=np.array([6,6,6]), axisz=np.array([0,0,0]),
                         psi_rmin=4, psi_rmax=8, psi_nr=3,
                         psi_zmin=-2, psi_zmax=2, psi_nz=3,
                         psi_phimin=0, psi_phimax=360, psi_nphi=3,
                         desc="DUMMY")

    @staticmethod
    def convert_B_GS(R0, z0, B_phi0, psi_mult, psi_coeff,
                     Rmin, Rmax, nR, zmin, zmax, nz, phimin, phimax, nphi,
                     psi0=None, raxis=None, zaxis=None, desc=None,
                     **kwargs):
        """Convert :class:`B_GS` input to `B_STS` input.

        The resulting field is far from being divergence-free as the
        perturbation BR and Bz components are omitted.

        Parameters
        ----------
        R0, z0 : real
            Major axis Rz-coordinates.
        B_phi0 : real
            Toroidal field at axis.
        psi_mult : real
            Scaling factor for psi.
        psi_coeff : real 13 x 1 numpy array
            Coefficients defining psi.
        Rmin, Rmax, zmin, zmax, phimin, phimax : real
            Edges of the uniform Rphiz-grid.
        nR, nz, nphi : int
            Number of Rphiz-grid points.
        psi0 : float, optional <br>
            Poloidal flux at magnetic axis.
        raxis : float, optional <br>
            Magnetic axis R coordinate [m].
        zaxis : float, optional <br>
            Magnetic axis z coordinate [m].
        **kwargs
            Arguments passed to :meth:`B_GS.write_hdf5` excluding ``fn`` and
            ``desc``.

        Returns
        -------
        out : dict
            :class:`B_GS` converted as an input for :meth:`write_hdf5`.
        """


        rgrid = np.linspace(Rmin, Rmax, nR)
        zgrid = np.linspace(zmin, zmax, nz)

        zg, Rg = np.meshgrid(zgrid, rgrid);

        c = psi_coeff # For shorter notation.
        psiRz = psi_mult*psifun.psi0(Rg/R0,zg/R0,c[0],c[1],c[2],c[3],c[4],c[5],c[6],
                                     c[7],c[8],c[9],c[10],c[11],c[12])
        psiRz_R = psi_mult*psifun.psiX(Rg/R0,zg/R0,c[0],c[1],c[2],c[3],c[4],c[5],c[6],
                                     c[7],c[8],c[9],c[10],c[11],c[12])
        psiRz_z = psi_mult*psifun.psiY(Rg/R0,zg/R0,c[0],c[1],c[2],c[3],c[4],c[5],c[6],
                                     c[7],c[8],c[9],c[10],c[11],c[12])

        Br = np.zeros((nR, nphi, nz))
        Bz = np.zeros((nR, nphi, nz))
        Bphi = np.zeros((nR, nphi, nz))
        psi  = np.zeros((nR, nphi, nz))
        axisr= np.zeros((nphi,))
        axisz= np.zeros((nphi,))


        phigrid = np.linspace(phimin,phimax,nphi+1)
        phigrid = phigrid[0:-1]

        # search for magnetic axis if not given
        if psi0 == None:
         x = psifun.find_axis(R0, z0, c[0], c[1], c[2], c[3], c[4], c[5], c[6],
             c[7], c[8], c[9], c[10], c[11], c[12])
         psi0 = psi_mult*psifun.psi0(x[0], x[1], c[0], c[1], c[2], c[3], c[4],
             c[5], c[6], c[7], c[8], c[9], c[10], c[11], c[12]) # At axis.
         raxis = x[0]*R0
         zaxis = x[1]*R0

        psi1=0 #??

        #Normalize psi (It tries to impersonate the VMEC psi)
        #print("Initial psi0: {}".format(psi0))
        #print("Initial psi1: {}".format(psi1))
        psiRz -= psi0 # move axis to 0
        psi1  -= psi0 # ..
        psiRz /= psi1 # scale LCFS to 1
        psi0 = 0.0
        psi1 = 1.0


        # The magnetic field one gets from psi as:
        #  B_R &= -\frac{1}{R}\frac{\partial\psi}{\partial z}\\
        #  B_z &=  \frac{1}{R}\frac{\partial\psi}{\partial R}
        #  f
        #  B_\phi = \frac{B_0R_0}{R}
        #  f

        for i in range(0,nphi):
            Bphi[:,i,:] = ((R0/Rg)*B_phi0 )
            psi[ :,i,:] = psiRz
            Br[  :,i,:] = ( -1.0 / Rg / R0 ) * psiRz_z
            Bz[  :,i,:] = (  1.0 / Rg / R0 ) * psiRz_R
            axisr[ i  ] = raxis
            axisz[ i  ] = zaxis


        return { "b_rmin"  : Rmin,   "b_rmax"  : Rmax,  "b_nr"  : nR,
                 "b_zmin"  : zmin,   "b_zmax"  : zmax,  "b_nz"  : nz,
                 "b_phimin": phimin, "b_phimax": phimax,"b_nphi": nphi,
                 "axisr" : axisr,  "axisz"   : axisz,
                 "psi" : psi, "psi0" : psi0, "psi1"  : psi1,
                 "br"      : Br,     "bphi"    : Bphi,  "bz"    : Bz,
                 "axis_phimin":phimin, "axis_phimax":phimax, "axis_nphi":nphi}

a5py/a5py/ascot5io/bfield.py
ascot4fusion/ascot5
"""Input representing mapping from real-space to Boozer coordinates.

This input is required in simulations where the input with MHD eigenfunctions
is used.
"""
import h5py
import numpy as np

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class Boozer(DataGroup):
    """Mapping between cylindrical and Boozer coordinates assuming that psi
    is axisymmetric.

    This input can be automatically created from tokamak field input.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key in ["npsi", "ntheta", "nthetag", "nr", "nz", "nrzs"]:
                    out[key] = int(out[key])

        # (Remove padding to theta_psithetageom)
        padding = 4
        nthetag = int(out["nthetag"] - padding*2)
        out["theta_psithetageom"] = out["theta_psithetageom"][padding:,:]
        out["theta_psithetageom"] = out["theta_psithetageom"][:nthetag,:]
        out["nthetag"] = nthetag

        out["psi_rz"]             = np.transpose(out["psi_rz"])
        out["theta_psithetageom"] = np.transpose(out["theta_psithetageom"])
        out["nu_psitheta"]        = np.transpose(out["nu_psitheta"])

        return out

    @staticmethod
    def write_hdf5(fn, psimin, psimax, npsi, ntheta, nthetag, rmin, rmax, nr,
               zmin, zmax, nz, r0, z0, psi0, psi1, psi_rz, theta_psithetageom,
               nu_psitheta, nrzs, rs, zs, desc=None):
        """Write input data to the HDF5 file.

        Note: the data in theta_psithetageom is assummed to span the whole
        interval (i.e. coinciding start and end points included) in the angular
        axis. This is needed so that we can add "padding" to this dataset.
        To be mores specic, the padding means that the dataset is artificially
        extended in the angular space. This is needed so that ASCOT5 can use
        the natural boundary condition when fitting the splines to the data:
        the data is effectively periodic but the periodic condition would
        assume x_end = x_ini when in fact x_end = x_ini + 2*pi e.g. for
        theta_psithetageom.

        Note that the nu_psitheta data is also periodic, but there it is valid
        to use the natural boundary condition which sets y'' = 0 (the nu data is
        has a local extrema at theta=0).

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        psimin : float
            Minimum psi grid value.
        psimax : float
            Maximum psi grid value.
        npsi : int
            Number of psi grid points.
        ntheta : int
            Number of boozer theta grid values.
        nthetag : int
            Number of geometric theta grid values.
        rmin : float
            Minimum R grid value.
        rmax : float
            Maximum R grid value.
        nr : int
            Number of R grid points.
        zmin : float
            Minimum z grid value.
        zmax : float
            Maximum z grid value.
        nz : int
            Number of z grid points.
        r0 : float
            Magnetic axis R coordinate.
        z0 : float
            Magnetic axis z coordinate.
        psi0 : float
            Coordinate psi on axis.
        psi1 : float
            Coordinate psi on separatrix.
        psi_rz : array_like (nr,nz)
            Coordinate psi(R,z).
        theta_psithetageom : array_like (npsi,nthetag)
            Coordinate theta(psi, thetag).
        nu_psitheta : array_like (npsi,ntheta)
            nu(psi, theta).
        nrsz : int
            Number of separatrix Rz points.
        rs : array_like (nrsz,1)
            Separatrix R coordinates, start and end points should coincide.
        zs : array_like (nrsz,1)
            Separatrix z coordinates, start and end points should coincide.
        desc : str, optional
            Input's description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if psi_rz.shape != (nr,nz):
            raise ValueError("psi has invalid shape.")
        if theta_psithetageom.shape != (npsi,nthetag):
            raise ValueError("theta has invalid shape.")
        if nu_psitheta.shape != (npsi,ntheta):
            raise ValueError("nu has invalid shape.")
        if rs.size != nrzs:
            raise ValueError("Separatrix r has invalid shape.")
        if zs.size != nrzs:
            raise ValueError("Separatrix z has invalid shape.")

        psi_rz             = np.transpose(psi_rz)
        theta_psithetageom = np.transpose(theta_psithetageom)
        nu_psitheta        = np.transpose(nu_psitheta)

        parent = "boozer"
        group  = "Boozer"
        gname  = ""

        # Add padding to theta_psithetageom (see boozer.c for explanation)
        padding = 4
        data = np.copy(theta_psithetageom)
        theta_psithetageom = np.concatenate(
            (data, data[-1,:] + data[1:padding+1,:]) )
        theta_psithetageom = np.concatenate(
            (data[int(nthetag-padding-1):-1,:] - data[-1,:],
             theta_psithetageom) )
        nthetag += padding*2

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            # grid specifications
            g.create_dataset("psimin",  (1,), data=psimin,  dtype="f8")
            g.create_dataset("psimax",  (1,), data=psimax,  dtype="f8")
            g.create_dataset("npsi",    (1,), data=npsi,    dtype="i4")
            g.create_dataset("ntheta",  (1,), data=ntheta,  dtype="i4")
            g.create_dataset("nthetag", (1,), data=nthetag, dtype="i4")
            g.create_dataset("rmin",    (1,), data=rmin,    dtype="f8")
            g.create_dataset("rmax",    (1,), data=rmax,    dtype="f8")
            g.create_dataset("nr",      (1,), data=nr,      dtype="i4")
            g.create_dataset("zmin",    (1,), data=zmin,    dtype="f8")
            g.create_dataset("zmax",    (1,), data=zmax,    dtype="f8")
            g.create_dataset("nz",      (1,), data=nz,      dtype="i4")
            g.create_dataset("r0",      (1,), data=r0,      dtype="f8")
            g.create_dataset("z0",      (1,), data=z0,      dtype="f8")
            g.create_dataset("nrzs",    (1,), data=nrzs,    dtype="i4")

            # the outermost poloidal psi-surface contour
            g.create_dataset("rs", (nrzs,), data=rs, dtype="f8")
            g.create_dataset("zs", (nrzs,), data=zs, dtype="f8")

            # psi data min and max values for normalization
            g.create_dataset("psi0", (1,), data=psi0, dtype="f8")
            g.create_dataset("psi1", (1,), data=psi1, dtype="f8")

            # tabulated coordinates maps
            g.create_dataset("psi_rz", (nz,nr), data=psi_rz, dtype="f8")
            g.create_dataset("theta_psithetageom", (nthetag,npsi),
                             data=theta_psithetageom, dtype="f8")
            g.create_dataset("nu_psitheta", (ntheta,npsi),
                             data=nu_psitheta, dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output is a very large rectangular wall.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        rs = np.cos(np.linspace(0, 2*np.math.pi, 10))
        zs = np.sin(np.linspace(0, 2*np.math.pi, 10))

        return Boozer.write_hdf5(
            fn=fn, psimin=0, psimax=1, npsi=6, ntheta=10, nthetag=10,
            rmin=0.1, rmax=10.0, nr=5, zmin=-10, zmax=10, nz=10, r0=5, z0=0,
            psi0=0, psi1=1, psi_rz=np.ones((5,10)),
            theta_psithetageom=np.ones((6,10)), nu_psitheta=np.ones((6,10)),
            nrzs=10, rs=rs, zs=zs, desc="DUMMY")

a5py/a5py/ascot5io/boozer.py
ascot4fusion/ascot5
"""Low level API for HDF5 data access and building the treeview.
"""

a5py/a5py/ascot5io/coreio/__init__.py
ascot4fusion/ascot5
"""
Make older ASCOT5 HDF5 files compatible with newer versions of ASCOT5.

Whenever a new version is released, add a new function here that converts
previous version to current version. Note that you might have to modify the HDF5
directly as this tool should be independent of ascot5io (which is subject to
changes).

Also update the CURRENT_VERSION number.

File compatibility.py
"""
import h5py
import numpy as np
import tempfile
import subprocess

## KEEP THIS UPDATED ##
CURRENT_VERSION = 4

def convert_oldtonew(fnin, origin, keeptemp=False):
    """
    Convert an old file to the current version.

    The old file remains as it is, and the converted version will be named
    fnin_{CURRENT_VERSION}.h5.

    If there are several versions between the current and the old version, the
    conversion is done incrementally e.g. 1->2, 2->3, 3->4 etc.

    ***********DEPRECATED*******************

    Args:
        fnin : str <br>
            Name of the input file.
        origin : int <br>
            Version number of the old file.
        keeptemp : bool, optional <br>
            Flag whether the intermediate files are kept.
    """

    version = origin
    fntemp = fnin
    while version < CURRENT_VERSION:
        fnout = fnin[:-3] + "_" + str(version+1) + ".h5"

        # Make the version conversion.
        funname = "convert_" + str(version) + "to" + str(version+1)
        globals()[funname](fntemp, fnout)

        if (not keeptemp) and (origin < version):
            # TODO remove temporary files
            pass

        version += 1
        fntemp  = fnout

def convert(fnin):
    """
    Update version 3 HDF5 to version 4.

    - Adds REVERSE_TIME option.
    - Renames ENDCOND_MAX_SIMTIME to ENDCOND_LIM_SIMTIME.
    """
    from a5py.ascot5io.ascot5file import get_qid
    from a5py.ascot5io.boozer import write_hdf5_dummy as boozer_write_hdf5_dummy
    from a5py.ascot5io.mhd    import write_hdf5_dummy as mhd_write_hdf5_dummy

    fnout = fnin[:-3] + "_" + str(CURRENT_VERSION) + ".h5"

    subprocess.call(["cp", fnin, fnout])

    def wrapper():
        import a5py.ascot5io.asigma_loc as asigmamod
        print("Add: Dummy atomic data (asigma)")
        asigma = asigmamod.write_hdf5_empty(fnout)
        with h5py.File(fnout, "a") as h5:
            if "options" in h5:
                print("Add: REVERSE_TIME = 0")
                print("ENDCOND_MAX_SIMTIME -> ENDCOND_LIM_SIMTIME")
                print("Add: ENABLE_ATOMIC = 0")
                print("Add: ENDCOND_IONIZED = 0")
                print("Add: ENDCOND_NEUTRALIZED = 0")
                for opt in h5["options"]:
                    opt = h5["options"][opt]
                    opt.create_dataset("ENDCOND_LIM_SIMTIME", (1,),
                                       data=opt["ENDCOND_MAX_SIMTIME"],
                                       dtype='f8')

                    opt.create_dataset("REVERSE_TIME", (1,), data=0,
                                       dtype='f8')
                    opt.create_dataset("ENABLE_ATOMIC", (1,), data=0,
                                       dtype='i8')
                    opt.create_dataset("ENDCOND_IONIZED", (1,), data=0,
                                       dtype='i8')
                    opt.create_dataset("ENDCOND_NEUTRALIZED", (1,), data=0,
                                       dtype='i8')
                    del opt["ENDCOND_MAX_SIMTIME"]

            if "results" in h5:
                for run in h5["results"]:
                    if run[:4] != "run": continue
                    run.attrs["qid_asigma"] = asigma.split["_"][-1]
                    print(asigma)

    wrapper()

    print("Conversion complete.")


def convert2to3(fnin):
    """
    Update version 2 HDF5 to version 3.

    - Adds dummy boozer and mhd groups.
    - Adds flags to wall input.
    - Adds ENDCOND_MAX_MILEAGE from options.
    - Adds ENABLE_MHD from options.
    - Replace velocity with momentum in distribution options.
    """
    from .coreio.treedata import get_qid
    from a5py.ascot5io.boozer import write_hdf5_dummy as boozer_write_hdf5_dummy
    from a5py.ascot5io.mhd    import write_hdf5_dummy as mhd_write_hdf5_dummy

    fnout = fnin[:-3] + "_" + str(CURRENT_VERSION) + ".h5"

    subprocess.call(["cp", fnin, fnout])

    print("Adding a dummy Boozer and MHD input groups.")
    boozer_qid = get_qid(boozer_write_hdf5_dummy(fnout))
    mhd_qid    = get_qid(mhd_write_hdf5_dummy(fnout))

    def wrapper():
        with h5py.File(fnout, "a") as h5:
            if "results" in h5:
                print("Adding dummy Boozer and MHD inputs for existing runs.")
                for run in h5["results"]:
                    h5["results"][run].attrs["qid_mhd"] = np.string_(mhd_qid)
                    h5["results"][run].attrs["qid_boozer"] = np.string_(boozer_qid)

            if "wall" in h5:
                print("Adding flags to wall inputs.")
                walls = list(h5["wall"].keys())
                for wall in walls:
                    if "wall_3D" in wall:
                        g = h5["wall"][wall]
                        if not "flag" in g:
                            nelements = int(g["nelements"][:])
                            flag = np.zeros(shape=(nelements,1), dtype=int)
                            g.create_dataset('flag',(nelements,1), data=flag,
                                         dtype='i4')
                        else:
                            print('      Already there...skipping.')

            if "options" in h5:
                print("Adding ENABLE_MHD=0 and ENDCOND_MAX_MILEAGE=100")
                print("and setting distribution momentum limits (DIST_MIN/MAX_P* ) to")
                print("(DIST_MIN/MAX_V*)*helium mass. Bin number is kept same and")
                print("the old velocity options are removed.")
                for opt in h5["options"]:
                    opt = h5["options"][opt]
                    opt.create_dataset("ENDCOND_MAX_MILEAGE", (1,), data=100,
                                       dtype='f8')

                    opt.create_dataset("ENABLE_MHD", (1,), data=0,
                                       dtype='f8')

                    for coord in ["PA", "PE", "R", "PHI", "Z"]:
                        helium_mass = 6.65e-27
                        nbin = opt["DIST_NBIN_V"+coord][:]
                        mini = opt["DIST_MIN_V"+coord][:] * helium_mass
                        maxi = opt["DIST_MAX_V"+coord][:] * helium_mass
                        opt.create_dataset("DIST_NBIN_P"+coord, (1,), data=nbin,
                                           dtype='f8')
                        opt.create_dataset("DIST_MIN_P"+coord, (1,), data=mini,
                                           dtype='f8')
                        opt.create_dataset("DIST_MAX_P"+coord, (1,), data=maxi,
                                           dtype='f8')

                        del opt["DIST_NBIN_V"+coord]
                        del opt["DIST_MIN_V"+coord]
                        del opt["DIST_MAX_V"+coord]

    wrapper()

    print("Conversion complete.")


def convert_1to2(fnin):
    """
    Update version 1 HDF5 to version 2.

    - Adds dummy NBI input to existing runs.
    """
    from .coreio.treedata import get_qid
    from a5py.ascot5io.nbi import write_hdf5_dummy

    fnout = fnin[:-3] + "_" + str(CURRENT_VERSION) + ".h5"

    subprocess.call(["cp", fnin, fnout])

    print("Adding a dummy NBI input group.")
    qid = get_qid(write_hdf5_dummy(fnout))

    print("Adding the dummy NBI input as a pseudo-input for existing runs.")
    def wrapper():
        with h5py.File(fnout, "a") as h5:
            if not "results" in h5:
                return

            for run in h5["results"]:
                h5["results"][run].attrs["qid_nbi"] = np.string_(qid)

    wrapper()

    print("Conversion complete.")

a5py/a5py/ascot5io/coreio/compatibility.py
ascot4fusion/ascot5
"""Module for creating and modifying the HDF5 file and for accessing meta data.

ASCOT5 HDF5 file is built as follows. At top level are parent groups, one (and
only one) for each type of input (e.g. bfield for magnetic field input) and one
parent group that contains all the results.

Parent groups contain groups that hold the actual datasets. Single parent can
have multiple groups e.g. bfield can store multiple magnetic field inputs.
The result group is a different in a sense that it contains run groups, one
for each simulation, which then contain the groups that store inistate,
endstate, and whatever diagnostics were used. However, here we refer to the
top level groups as parents, and the groups that are directly below them
as data containers.

A typical structure of a file can be like this:

> /
> /bfield
> /bfield/B_2DS_1234567890
> /bfield/B_2DS_2345678901
> /bfield/B_3DS_3456789012
> ...
> <other input groups>
> ...
> /results
> /results/run_4567890123
> /results/run_5678901234

The number in group name is unique identifier (QID) which is generated when a
group is created. It is generated from a 32 bit random integer and as such it
is almost quaranteed to be unique.

QID is used to determine which input group is active, i.e., to be used in
a simulation, and to which run group the simulation results are written.
QID of the active group is stored as an attribute (in a string format) in the
parent. For the results parent this is the QID of the most recent run (unless
manually changed). Run groups also store QIDs of all input fields used in that
particular run.

All groups also store date and time at which they where created and also a
description field that user can modify at will (see ascot5.py).

The groups should be created, removed, and their metadata accessed via this
module. The actual datasets are accessed via corresponding modules. This module
assumes the HDF5 file is open when these routines are called, and closed
afterwards.
"""

import numpy as np
import h5py
import unyt
import random
import datetime

from collections import OrderedDict
from a5py.exceptions import AscotNoDataException, AscotIOException

INPUTGROUPS = ["options", "bfield", "efield", "marker", "plasma", "neutral",
               "wall", "boozer", "mhd", "asigma", "nbi"]
"""Names of the input parent groups.
"""

OUTPUTGROUPS = ["inistate", "endstate", "dist5d", "distrho5d", "dist6d",
                "distrho6d", "orbit", "transcoef"]
"""Names of the output data containers in runs.
"""

VERSION = "5.5"
"""Current version of the code."""

def set_active(f, group):
    """Set given group as active.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        Group object or its name.

    Raises
    ------
    AscotNoDataException
        Raised if the group or its parent does not exist.
    """

    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group" + group)
        else:
            group = grp

    qid = get_qid(group)
    group.parent.attrs["active"] = np.string_(qid)

def get_active(f, parent):
    """Get active group.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    parent : str or `h5py.Group`
        Group object or its name.

    Returns
    -------
    group : `h5py.Group`
        The active group.

    Raises
    ------
    AscotNoDataException
        Raised if the parent or the active group does not exist.
    """

    if str(parent) == parent:
        if parent in f:
            parent = f[parent]
        else:
            raise AscotNoDataException(
            "Parent " + parent + " does not exist.")

    qid = parent.attrs["active"].decode('utf-8')
    group = get_group(f, qid)

    if group == None:
        raise AscotNoDataException("Active group does not exist.")
    else:
        return group

def set_activeqid(f, qid):
    """Set a group with the given QID as the active one.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    qid : str
        QID.
    """
    group = get_group(f, qid)
    set_active(f, group)

def get_activeqid(f, parent):
    """Get QID of the currently active group.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    parent : str or `h5py.Group`
        Parent whose active group is sought.

    Returns
    -------
    qid : str
        QID string.
    """
    group = get_active(f, parent)
    return group.parent.attrs["active"].decode('utf-8')


def set_desc(f, group, desc):
    """Set group description.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        Group object or its name.
    desc : str
        Description as a string.

    Raises
    ------
    AscotNoDataException
        Raised if group does not exist.
    """
    # Check the group exists and access it.
    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group" + group)
        else:
            group = grp

    group.attrs["description"] = np.string_(desc)

def get_desc(f, group):
    """Get group description.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        Group object or its name.

    Returns
    -------
    desc : str
        Description as a string.

    Raises
    ------
    AscotNoDataException
        Raised if group does not exist.
    """

    # Check the group exists and access it.
    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group" + group)
        else:
            group = grp

    return group.attrs["description"].decode('utf-8')

def _set_date(f, group, date):
    """Set group date.

    Note that this function should only be called when the group is created.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        Group object or its name.
    date : str
        Date as a string in a YYYY-MM-DD hh:mm:ss format.

    Raises
    ------
    AscotNoDataException
        Raised if group does not exist.
    """

    # Check the group exists and access it.
    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group" + group)
        else:
            group = grp

    group.attrs["date"] = np.string_(date)

def get_date(f, group):
    """Get date.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        Group object or its name.

    Returns
    -------
    date : str
        Date as a string in format YYYY-MM-DD hh:mm:ss.

    Raises
    ------
    AscotNoDataException
        Raised if group does not exist.
    """

    # Check the group exists and access it.
    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group" + group)
        else:
            group = grp

    return group.attrs["date"].decode('utf-8')

def _set_version(f, group):
    """Set group version.

    Note that this function should only be called when the group is created.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        Group object or its name.

    Raises
    ------
    AscotNoDataException
        Raised if group does not exist.
    """

    # Check the group exists and access it.
    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group" + group)
        else:
            group = grp

    group.attrs["version"] = np.string_(vers)

def get_version(f, group):
    """Get input version.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        Group object or its name.

    Returns
    -------
    version : str
        Version as a string.

    Raises
    ------
    AscotNoDataException
        Raised if group does not exist.
    """

    # Check the group exists and access it.
    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group" + group)
        else:
            group = grp

    return group.attrs["version"].decode('utf-8')

def get_qid(group):
    """Get QID from a given group.

    Parameters
    ----------
    group : str or `h5py.Group`
        Group object or its name.

    Returns
    -------
    qid : str
        QID string.

    Raises
    ------
    AscotNoDataException
        Raised if group does not have a valid QID.
    """
    if(str(group) != group):
        group = group.name

    if len(group) >= 10:
        qid = group[-10:]
        if qid.isdigit():
            # Seems like a valid QID
            return qid

    # Not a valid QID
    raise AscotNoDataException(group + " is not a valid QID.")

def get_type(group):
    """Get type from a given group or from its name.

    Parameters
    ----------
    group : str or `h5py.Group`
        Group object or its name.

    Returns
    -------
    type : str
         Group's type.

    Raises
    ------
    AscotNoDataException
        Raised if group does not have a valid type.
    """
    if(str(group) != group):
        group = group.name

    group = group.split("/")[-1]
    if len(group) > 12:
        type_ = group[:-11]
        if not type_.isdigit():
            # Seems like a valid QID
            return type_

    # Not a valid type
    raise AscotNoDataException(group + " does not contain a valid type.")

def get_group(f, qid):
    """Scan the file and return the group the QID corresponds to.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    qid : str
        QID string.

    Returns
    -------
    group : `h5py.Group`
        Group or None if the group was not present.
    """
    for parent in f.keys():
        for group in f[parent].keys():
            if qid in group:
                return f[parent][group]

    return None

def get_qids(f, parent):
    """Get QIDs of all the groups that a given parent contains.

    The QIDs are returned as list of no specific order.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    parent : str or `h5py.Group`
        Parent object or its name.

    Returns
    -------
    qids : list [str]
        A list of QID strings.

    Raises
    ------
    AscotNoDataException
        Raised if parent group does not exist.
    """
    # Check the parent exists and access it
    if(str(parent) == parent):
        if not parent in f:
            raise AscotNoDataException("Could not find parent " + parent)
        parent = f[parent]

    qids = []
    for group in parent:
        qids.append(get_qid(group))

    return qids

def get_inputqids(f, rungroup):
    """Get all QIDs that tell which input was used in the given run group.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    rungroup : str or `h5py.Group`
        Either the run group object or its name.

    Returns
    -------
    qids : `collections.OrderedDict` [str, str]
        Ordered dictionary with "parent name" - "qid" value pairs with the order
        being same as in `INPUTGROUPS`.

    Raises
    ------
    AscotNoDataException
        Raised if run group does not exist.
    """
    # Check the group exists and access it.
    if(str(rungroup) == rungroup):
        qid = get_qid(rungroup)
        grp = get_group(f, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group " + rungroup)
        else:
            rungroup = grp

    qids = OrderedDict();
    for inp in INPUTGROUPS:
        try:
            qids[inp] = rungroup.attrs["qid_" + inp].decode("utf-8")
        except KeyError as err:
            pass

    return qids

def add_group(f, parent, group, desc=None):
    """Create a new group. A parent is created if need be.

    If parent is created, then this new group is set as active.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    parent : str or `h5py.Group`
        Parent of the created group.
    group : str
        The group name.
    desc : str, optional
        Description for the group.

    Returns
    -------
    newgroup : `h5py.Group`
        The new group.
    """
    # Create a parent if one does not exists yet.
    if str(parent) == parent:
        parent = f.require_group(parent)

    # Generate metadata and include qid in group's name.
    qid, date, defdesc, vers = _generate_meta()
    group = group + "_" + qid
    if desc == None:
        desc = defdesc

    # Create group and set it as the active group if it is the only group.
    # Also set date and description.
    parent.create_group(group)
    set_desc(f, group, desc)
    _set_date(f, group, date)

    if len(parent.keys()) == 1:
        set_active(f, group)

    return get_group(f, qid)


def remove_group(f, group):
    """Remove a group.

    If this was an active group, a most recent group is set as an active
    instead. If no other groups exist the parent is also removed.

    Input groups are not removed if they have been used in an existing run.

    Note that to reclaim the disk space which the group occupied, one needs
    to call h5repack in a terminal.

    Parameters
    ----------
    f : `h5py.File`
        Open HDF5 file.
    group : str or `h5py.Group`
        The group to be removed.

    Raises
    ------
    AscotNoDataException
        Raised if group could not be found.
    AscotIOException
        Raised if the group is input used by a run.
    """
    # Check the group exists and access it.
    if(str(group) == group):
        if group in f.keys():
            # Group is a parent group
            group = f[group]
        else:
            # assume group is data group
            qid = get_qid(group)
            grp = get_group(f, qid)
            if grp is None:
                raise AscotNoDataException("Could not find group" + group)
            else:
                group = grp

    # Check if this is an input group and whether it has been used in run.
    parent = group.parent
    if "results" in f.keys() and parent.name != "results":
        if parent.name == "/":
            for run in f["results"].keys():
                for q in f["results"][run].attrs.keys():
                    if q[:4] == "qid_" and q[4:] == group.name[1:]:
                        raise AscotIOException(
                        "Input \"" + group.name[1:] + "\" is used by " + run)
        else:
            qid = get_qid(group)
            for run in f["results"].keys():
                for q in f["results"][run].attrs.keys():
                    t = f["results"][run].attrs[q].decode('utf-8')
                    if q[:4] == "qid_" and t == qid:
                        raise AscotIOException(
                            "Input " + qid + " is used by " + run)

    # Remove the group
    if parent.name!='/':
        was_active = get_active(f, parent) == group
        del f[group.name]
    else:
        was_active=False
        del f[group.name]

    # Set next group active (if removed group was) or remove the parent if no
    # other groups exist
    if was_active:
        if len(parent.keys()) == 0:
            del f[parent.name]
        else:
            date = "0"
            for grp in parent:
                grpdate = get_date(f, grp)
                if grpdate > date:
                    date  = grpdate
                    group = grp

            set_active(f, group)


def copy_group(fs, ft, group, newgroup=False):
    """Copy group from one file to another. A parent is also created if need be.

    The new group is set as active if the parent on the target file has no other
    groups. The copied group retains its original QID and date of creation.

    Parameters
    ----------
    fs : `h5py.File`
        File from which the group is copied.
    ft : `h5py.File`
        File to which the group is copied.
    group : str pr `h5py.Group`
        The group to be copied.
    newgroup : bool, optional
        Flag indicating if copied group should be given new QID and date.

    Returns
    -------
    groupt : `h5py.Group`
        The new group which is a copy of group at ft.

    Raises
    ------
    AscotNoDataException
        Raised if the copied group cannot be found.
    AscotIOException
        Raised if the group already exists on the target file.
    """
    # Check the group exists and access it.
    if(str(group) == group):
        qid = get_qid(group)
        grp = get_group(fs, qid)
        if grp is None:
            raise AscotNoDataException("Could not find group " + group)
        else:
            group = grp

    # Create target parent if none exists.
    parentname = group.parent.name

    newparent  = ft.require_group(parentname)
    if group.name in newparent:
        raise AscotIOException("Target already has the group " + group.name)

    # Copy
    if newgroup:
        qid, date, defdesc, vers = _generate_meta()
        newname = group.name[:-11] + "_" + qid
        fs.copy(group, newparent, name=newname)
        _set_date(ft, newname, date)
        newgroupobj = ft[parentname][newname]
    else:
        fs.copy(group, newparent)
        newgroupobj = ft[parentname][group.name]

    if(len(newparent)) == 1:
        set_active(ft, newgroupobj)

    return newgroupobj


def write_data(group, name, data, shape, dtype, unit=None):
    """Write a dataset.

    The shape of the written dataset is same as the input array.

    Parameters
    ----------
    group : `h5py.Group`
        HDF5 group where the dataset will be written.
    name : str
        Name of the new dataset.
    data : `np.array`
        Data to be written.
    dtype : str
        Data type.
    unit : str, optional
        Unit string if the data has units.
    """
    g = group.create_dataset(
        name  = name,
        shape = shape,
        data  = data,
        dtype = dtype
    )
    if unit is not None:
        g.attrs.create("unit", np.string_(unit))


def read_data(group, name):
    """Read a dataset and add units if present.

    Parameters
    ----------
    group : `h5py.Group`
        HDF5 group containing the dataset.
    name : str
        Name of the dataset.

    Returns
    -------
    data : `np.array` or `unyt.array`, `(N,)`
        Dataset with units read from the dataset attribute "unit".

        If dataset has no "unit" attribute, ordinary `np.array` is returned.
    """
    if "unit" in group[name].attrs.keys():
        unit_str = group[name].attrs["unit"]
        unit     = unyt.Unit(unit_str)

        return group[name][:].ravel() * unit
    else:
        return group[name][:].ravel()



def _generate_meta():
    """Generate QID, date and default description/tag.

    Calls random number generator to create 32 bit string which is then
    converted as a QID string using left-padding with zeroes if necessary.

    Returns
    -------
    meta : `tuple` [str, str, str, str]
        QID, date, description, and version number.
    """
    # Generate random unsigned 32 bit integer and convert it to string.
    qid = str(np.uint32(random.getrandbits(32)))

    # Left-padding with zeroes so that QID is always 10 characters long.
    while len(qid) < 10:
        qid = "0" + qid

    # Get date
    date = str(datetime.datetime.now())

    # Last digits are milliseconds which we don't need
    date = date[0:19]

    desc = "TAG"

    vers = VERSION

    return (qid, date, desc, vers)

a5py/a5py/ascot5io/coreio/fileapi.py
ascot4fusion/ascot5
"""
Tools to modify HDF5 files.

File: ascot5tools.py
"""
import numpy as np
import h5py
from a5py import Ascot
from . import fileapi

def call_fileapi(fn, method, *args):
    """
    Wrapper for calling fileapi methods.

    This wrapper handles opening and closing of HDF5 file.

    Args:
        fn : str <br>
            Full path to HDF5 file.
        method: str <br>
            Name of the method to be called.
        *args: <br>
            Arguments the called method requires.

    Returns:
        The value the called method returns.

    Raise:
        ValueError if the method does not exist.
    """
    if hasattr(fileapi, method):
        method_to_call = getattr(fileapi, method)
        with h5py.File(fn, "a") as f:
            return method_to_call(f, *args)
    else:
        raise ValueError(method + " is not a valid method.")


def removegroup(fn, group, force=False):
    """
    Remove a group or a parent.

    Tries to remove a group. If the group to be removed is an input group which
    has been used as an input in some of the run groups, the group won't be
    removed and exception is raised.

    Note that to reclaim the disk space which the group occupied, one needs
    to call h5repack in a terminal.

    Args:
        fn : str <br>
            Full path to HDF5 file.
        group: str <br>
            Name of the group to be removed.
        force: bool, optional <br>
            Remove group and don't raise an exception even if the group has been
            used as an input. Default is false.
    Raise:
        RuntimeError if the group has been used as an input.
    """
    with h5py.File(fn, "a") as f:

        try:
            # This will raise an exception if group is not a data group.
            qid = fileapi.get_qid(group)
            filegroup = fileapi.get_group(f, qid)

            # This is a data group. If it is a run group or removal is forced,
            # we can remove it directly.
            if force or fileapi.get_type(filegroup) == "run":
                fileapi.remove_group(f, group)
                return

            # For input data groups we have to check if any run group refers to
            # it.
            runqids = fileapi.get_qids(f, "results")

            for runqid in runqids:
                rungroup = fileapi.get_group(f, runqid)
                inqids   = fileapi.get_inputqids(f, rungroup)
                if qid in inqids:
                    raise RuntimeError("Run " + runqid
                                       + " has used group " + qid
                                       + " as an input. Removal aborted.")

            # No references, the group can be removed.
            fileapi.remove_group(f, group)

        except ValueError:
            # The group is a parent group. If it is a results group or removal
            # is forced, we can remove it directly.
            if force or group == "results":
                fileapi.remove_group(f, group)
                return

            try:
                runqids  = fileapi.get_qids(f, "results")
            except ValueError:
                # There is no results group, so we can safely remove the group
                pass
            else:
                # For input parent groups we have to check if any run group
                # refers to any of its data groups.
                dataqids = fileapi.get_qids(f, group)

                for dataqid in dataqids:
                    for runqid in runqids:
                        rungroup = fileapi.get_group(f, runqid)
                        inqids   = fileapi.get_inputqids(f, rungroup)
                        if dataqid in inqids:
                            raise RuntimeError("Run " + runqid
                                               + " has used group " + dataqid
                                               + " as an input. Removal aborted.")

            # No references, the group can be removed.
            fileapi.remove_group(f, group)


def copygroup(fns, fnt, group, newgroup=False):
    """
    Copy a group or a parent to a different HDF5 file.

    The copied field is set as the active field in target HDF5 file. The copied
    group retains it qid and description.

    Args:
        fns : str <br>
            Full path to source file.
        fnt : str <br>
            Full path to target file.
        group : str <br>
            Name of the group to be copied.
        newgroup : bool, optional <br>
            Flag indicating if copied group should be given new QID and creation
            date. Default is false.

    Returns:
        Name of the copied group or none if the group was parent.
    """

    # Get the target field and type from source
    with h5py.File(fns, "r") as fs, h5py.File(fnt, "a") as ft:

        try:
            # This will raise exception if group is not a data group
            fileapi.get_qid(group)

            # This is a data group
            grp = fileapi.copy_group(fs, ft, group, newgroup=newgroup)
            return grp.name
        except ValueError:
            # This is a parent group
            qids = fileapi.get_qids(fs, group)
            for qid in qids:
                grp = fileapi.get_group(fs, qid)
                fileapi.copy_group(fs, ft, grp, newgroup=newgroup)


def combineoutput(fnt, addfns=None, contfns=None):
    """
    Combine outputs of two HDF5 files.

    Depending on which argument is given, this either combines the output by
    assuming fnt and addfns are both part of a larger simulation ("add"), or
    markers in contfns were initialized from the endstate of fnt ("continue").

    The combined outputs will be stored on fnt on a run group that is active.
    The outputs are read from addfns (or contfns) from the group that is active.

    Args:
        fnt : str <br>
            Full path to HDF5 file where combined output will be stored.
        addfns : str, optional <br>
            Name of the HDF5 file where results are read if "add" mode is being
            used.
        contfns : str, optional <br>
            Name of the HDF5 file where results are read if "continue" mode is
            being used.
    """

    if addfns is not None:
        fns = addfns
    elif contfns is not None:
        fns = contfns
    else:
        return

    # Find the active groups
    source = Ascot(fns).data.active
    target = Ascot(fnt).data.active

    # Combine states
    if addfns is not None:
        if hasattr(target, "inistate") and hasattr(source, "inistate"):
            with target.inistate as tdata, source.inistate as sdata:
                for field in tdata:
                    tsize = tdata[field].size
                    ssize = sdata[field].size

                    tdata[field].resize((tsize+ssize, ))
                    tdata[field][tsize:] = sdata[field][:]

        if hasattr(target, "endstate") and hasattr(source, "endstate"):
            with target.endstate as tdata, source.endstate as sdata:
                for field in tdata:
                    tsize = tdata[field].size
                    ssize = sdata[field].size

                    tdata[field].resize((tsize+ssize, ))
                    tdata[field][tsize:] = sdata[field][:]
    else:
        # Sort target inistate by id.
        if hasattr(target, "inistate"):
            with target.inistate as tdata:
                idx = np.argsort(tdata["id"][:])
                for field in tdata:
                    tdata[field][:] = tdata[field][:][idx]

        # Replace target endstate with sorted by id source endstate.
        if hasattr(target, "endstate") and hasattr(source, "endstate"):
            with target.endstate as tdata, source.endstate as sdata:
                idx = np.argsort(sdata["id"][:])
                for field in tdata:
                    tdata[field][:] = sdata[field][:][idx]

    # Combine dists
    if hasattr(target, "dist5d") and \
       hasattr(source, "dist5d"):
        with target.dist5d as tdata, \
             source.dist5d as sdata:
            tdata["ordinate"][:] += sdata["ordinate"][:]

    if hasattr(target, "dist6d") and \
       hasattr(source, "dist6d"):
        with target.dist6d as tdata, \
             source.dist6d as sdata:
            tdata["ordinate"][:] += sdata["ordinate"][:]

    if hasattr(target, "distrho5d") and \
       hasattr(source, "distrho5d"):
        with target.distrho5d as tdata, \
             source.distrho5d as sdata:
            tdata["ordinate"][:] += sdata["ordinate"][:]

    if hasattr(target, "distrho6d") and \
       hasattr(source, "distrho6d"):
        with target.distrho6d as tdata, \
             source.distrho6d as sdata:
            tdata["ordinate"][:] += sdata["ordinate"][:]

    # Combine orbits
    if hasattr(target, "orbit") and hasattr(source, "orbit"):
        with target.orbit as tdata, source.orbit as sdata:
            for field in tdata:
                tsize = tdata[field].size
                ssize = sdata[field].size

                tdata[field].resize((tsize+ssize, ))
                tdata[field][tsize:] = sdata[field][:]

    # Combine transport coefficients
    if hasattr(target, "transcoef") and hasattr(source, "transcoef"):
        with target.transcoef as tdata, source.transcoef as sdata:
            for field in tdata:
                tsize = tdata[field].size
                ssize = sdata[field].size

                tdata[field].resize((tsize+ssize, ))
                tdata[field][tsize:] = sdata[field][:]

a5py/a5py/ascot5io/coreio/tools.py
ascot4fusion/ascot5
"""Abstract classes for objects containing data.
"""
import h5py

from a5py.exceptions import AscotIOException
from . import fileapi

class DataContainer():
    """Interface for providing access to underlying HDF5 data.

    This class provides methods to change and access the raw data. The
    raw data is accessed (via :obj:`h5py`) as

    .. code-block:: python

       with datacontainer as d:
           d["data"][:]

    where ``d`` is the :obj:`h5py.Group` corresponding to this container.

    Attributes
    ----------
    _root : :class:`RootNode`
        The rootnode this object belongs to.
    _path : str
        Path to this data within the HDF5 file.
    _opened : list [obj]
        List containing a single object, which is this data's HDF5 group if
        the data is being accessed and `None` otherwise.

        We are storing single item list, instead of storing the item directly,
        since this class is inherited by `Note` objects whose attributes cannot
        be changed directly. However, if the attribute is a list, the contents
        of that list can be changed ;).
    """

    def __init__(self, root, path, **kwargs):
        """Initialize data container.

        Parameters
        ----------
        root : :class:`RootNode`
            The root node this data container belongs to.
        path : str
            Path to this data within the HDF5 file.
        **kwargs
            Arguments passed to other constructors in case of multiple
            inheritance.
        """
        self._root   = root
        self._path   = path
        self._opened = [None]
        super().__init__(**kwargs)

    def __enter__(self):
        """Call _open when entering ``with`` clause.

        Returns
        -------
        data : :class:`h5py.Group`
            HDF5 group corresponding to this data.
        """
        return self._open()

    def __exit__(self, exception_type, exception_value, traceback):
        """Call _close when exiting ``with`` clause.
        """
        self._close()

    def _open(self, mode="r"):
        """Open and return the HDF5 group corresponding to this data.

        Returns
        -------
        data : :class:`h5py.Group`
            HDF5 group corresponding to this data.
        mode : {"r", "a"}
            Is the file opened for (r)eading or (a)ppending.

        Raises
        ------
        AscotIOException
            Raised if this instance has already opened the file.
        """
        if self._opened[0] is not None:
            raise AscotIOException(
                "File already opened by this instance")

        fn = self._root._ascot.file_getpath()
        self._opened[0] = h5py.File(fn, mode)[self._path]
        return self._opened[0]

    def _close(self):
        """Close the HDF5 group corresponding to this data.
        """
        self._opened[0].file.close()
        self._opened[0] = None

class DataGroup(DataContainer):
    """Data container that also has meta data (QID, date, description).
    """

    def set_desc(self, desc):
        """Set description for this group.

        Note that the first word in the description is this group's tag which
        can be used to refer to this group.

        Parameters
        ----------
        desc : str
            Short description for the user to document this group.
        """
        f = self._open("a")
        val = fileapi.set_desc(f.file, self._path, desc)
        self._close()
        self._root._build(self._root._ascot.file_getpath())
        return val

    def get_desc(self):
        """Get this group's description.

        Returns
        -------
        desc : str
            Documentation the user has used to describe this group.
        """
        f = self._open()
        val = fileapi.get_desc(f.file, self._path)
        self._close()
        return val

    def get_date(self):
        """Get the date when this group was created.

        Returns
        -------
        date : str
            The date in YYYY-MM-DD hh:mm:ss format.
        """
        f = self._open()
        val = fileapi.get_date(f.file, self._path)
        self._close()
        return val

    def get_qid(self):
        """Get QID of this group.

        Returns
        -------
        qid : str
            String with 10 characters containing numbers from 0-9 which is
            an unique identifier for this group.
        """
        val = fileapi.get_qid(self._path)
        return val

    def get_type(self):
        """Return type of this group.

        Returns
        -------
        gtype : str
            Group type.
        """
        path = self._path.split("/")[-1]
        return path[:-11]

    def get_name(self):
        """Return name of this group.

        Returns
        -------
        name : str
            The name of the group as "<group type>_<group qid>".
        """
        return self._path.split("/")[2]

    def activate(self):
        """Set this group as active.

        Active inputs are used when the simulation is run. Active groups are
        also used during post-processing.
        """
        self._root._activate_group(self.get_qid())

    def destroy(self, repack=True):
        """Remove this group from the HDF5 file.

        Parameters
        ----------
        repack : bool, optional
            If True, repack the HDF5 file.

            Removing data from the HDF5 file only removes references to it and
            repacking is required to free the disk space. Repacking makes a copy
            of the HDF5 file and destroys the original, and it also requires
            3rd party tool `h5repack` which is why it's use is optional here.
        """
        self._root._destroy_group(self.get_qid(), repack)

    def export(self,target_file,newgroup=False):
        """
        """

        import a5py.ascot5io.ascot5tools as tools
        group = tools.copygroup(self._root.file_getpath(), target_file,
                                self._path, newgroup=newgroup)
        return group

a5py/a5py/ascot5io/coreio/treedata.py
ascot4fusion/ascot5
"""Module for building the treeview that shows ASCOT5 HDF5 file contents.

Treeview consists of nodes acting as containers for other nodes and data. The
tree is spanned by the `RootNode`.
"""
import h5py
import subprocess

from . import fileapi
from .treedata import DataGroup

from a5py.exceptions import AscotIOException

class _FancyText():
    """Helper class to decorate output of `ls` commands so it is easier to read.
    """
    _green     = "\033[32m"
    _purple    = "\033[35m"
    _reset     = "\033[0m"
    _bold      = "\033[01m"
    _underline = "\033[04m"

    @staticmethod
    def title(string):
        """Convert string to title (underlined bold purple text).

        Parameters
        ----------
        string : str
            String to be converted.

        Returns
        -------
        title : str
            The converted string.
        """
        return _FancyText._bold + _FancyText._underline + _FancyText._purple \
            + string + _FancyText._reset

    @staticmethod
    def header(string):
        """Convert string to header (bold text).

        Parameters
        ----------
        string : str
            String to be converted.

        Returns
        -------
        header : str
            The converted string.
        """
        return _FancyText._bold + string + _FancyText._reset

    @staticmethod
    def active(string):
        """Convert string to active input (green text).

        Parameters
        ----------
        string : str
            String to be converted.

        Returns
        -------
        active : str
            The converted string.
        """
        return _FancyText._green + string + _FancyText._reset

class _Node():
    """Base class which all tree nodes inherit.

    This base class provides all nodes with two functionalities:

    1. Its attributes can be accessed in dictionary-like manner, e.g.
       `node.child` and `node["child"]` are equivalent.
    2. Freeze (and unfreeze) this instance preventing (allowing) adding
       or removing attributes. Once the tree is constructed, all nodes
       should be frozen.

    Attributes
    ----------
    _frozen : bool
        When True, attributes cannot be added or removed for this node.
    """

    def __init__(self, **kwargs):
        """Initialize a mutable (unfrozen) node.

        Parameters
        ----------
        **kwargs
            Arguments passed to other constructors in case of multiple
            inheritance.
        """
        self._frozen = False
        super().__init__(**kwargs)

    def _freeze(self):
        """Make this node immutable (frozen).
        """
        self._frozen = True

    def _unfreeze(self):
        """Make this node mutable.
        """
        self._frozen = False

    def __setitem__(self, key, value):
        """Add a new attribute this node in dictionary style.

        Parameters
        ----------
        key : str
            Name of the attribute.
        value
            Value of the attribute.

        Raises
        ------
        AscotIOException
            Raised if the node is frozen.
        """
        if self._frozen:
            raise AscotIOException(
                "Node is frozen and attributes cannot be set")

        setattr(self, key, value)

    def __setattr__(self, key, value):
        """Add a new attribute this node.

        Parameters
        ----------
        key: str
            Name of the attribute.
        value:
            Value of the attribute.

        Raises
        ------
        AscotIOException
            Raised if the node is frozen.
        """
        if key != "_frozen" and self._frozen:
            raise AscotIOException(
                "Node is frozen and attributes cannot be set")
        else:
            super().__setattr__(key, value)

    def __contains__(self, key):
        """Check whether this node contains the attribute.

        Parameters
        ----------
        key : str
            Name of the attribute.

        Returns
        -------
        contains : bool
            True if this node contains the attribute.
        """
        return hasattr(self, key)

    def __getitem__(self, key):
        """Retrieve attribute in dictionary-like manner.

        Parameters
        ----------
        key : str
            Name of the attribute.

        Returns
        -------
        value
            Value of the attribute.
        """
        return getattr(self, key)

class _ParentNode(_Node):
    """Node that has data groups as children.

    This node stores QIDs of its children and provides a way of accessing them
    through active-flag, QID, name, and tag.

    Attributes
    ----------
    _qids : list [str]
        QIDs of this node's children.
    _root : `RootNode`
        The `RootNode` this node belongs to.
    _path : str
        Path to this node within the HDF5 file.
    active : `DataGroup`
        The currently active group.
    """

    def __init__(self, root, path, **kwargs):
        """Initialize a parent node that initially has no children.

        Parameters
        ----------
        root : `RootNode`
            The `RootNode` this node belongs to.
        path : str
            Path to this node within the HDF5 file.
        **kwargs
            Arguments passed to other constructors in case of multiple
            inheritance.
        """
        super().__init__(**kwargs)
        self._qids  = []
        self._root  = root
        self._path  = path
        self.active = None

    def _add_child(self, key, datagroup):
        """Add a child to this node.

        This should be called only when the HDF5 file is open and the tree is
        being built. After all childs have been added, `_finalize` must be
        called to organize this node.

        Parameters
        ----------
        key : str
            Name of the child in format <type>_<QID>.
        datagroup : `DataGroup`
            The child data.
        """
        qid = fileapi.get_qid(key)
        self._qids.append(qid)

        # Add reference by name
        self[key] = datagroup

        # Add reference by QID
        self["q" + qid] = datagroup

    def _finalize(self, h5):
        """Organize contents and add references.

        This method should be called once after all childs have been added and
        the HDF5 file is still open.

        Parameters
        ----------
        h5 : `h5py.File`
            The HDF5 file from which the tree is constructed.
        """
        # Find active group and set it
        qid = fileapi.get_activeqid(h5, h5[self._path])
        self.active = self["q"+qid]

        # List all dates but remove the active
        dates = []
        for q in self._qids:
            grp = fileapi.get_group(h5, q)
            dates.append(fileapi.get_date(h5, grp))

        # Sort QIDs by date starting from most recent
        self._qids = [x for _, x in sorted(zip(dates, self._qids))][::-1]

        # Add references by tag
        def desc2tag(desc):
            """Convert desc to tag.
            """
            # Cut from first whitespace
            desc = desc.split(" ")[0]

            # Maximum length is 10 characters
            #if len(desc) > 10:
            #    desc = desc[:10]

            # Remove all special characters
            desc = "".join(c for c in desc if c.isalnum())

            # Make all caps
            desc = desc.upper()

            # Use default if empty first character is number
            if desc == "" or desc[0] in "1234567890":
                desc = "TAG"

            return desc

        replacedtags = [] # Tags that appear more than once are deleted
        for q in self._qids:
            tag = desc2tag(fileapi.get_desc(h5, q))
            if tag in self:
                # This tag is repeated

                if tag not in replacedtags:
                    # First time we notice a tag is repeated, rename the tag
                    # already present to "tag_0" and add the tag to list of tags
                    # we remove afterwards. (We cannot remove the tag here as
                    # then we wouldn't notice repeated entries).
                    replacedtags.append(tag)
                    self[tag + "_0"] = self[tag]

                # Find next available index and create "tag_i".
                tag0 = tag; i = 0
                while tag in self:
                    tag = tag0 + "_" + str(i)
                    i += 1

            self[tag] = self["q"+q]

        # Remove repeated tags as they have their renamed variants in place.
        for tag in replacedtags:
            delattr(self, tag)

    def destroy(self, repack=True):
        """Remove this group from the HDF5 file.

        Parameters
        ----------
        repack : bool, optional
            If True, repack the HDF5 file.

            Removing data from the HDF5 file only removes references to it and
            repacking is required to free the disk space. Repacking makes a copy
            of the HDF5 file and destroys the original, and it also requires
            3rd party tool `h5repack` which is why it's use is optional here.
        """
        self._root._destroy_group(self._path, repack=repack)

    def get_contents(self):
        """Return metadata for all childs sorted by date.

        Returns
        -------
        qids : array_like, str
            QIDs.
        dates : array_like, str
            Dates.
        descs : array_like, str
            Descriptions.
        types : array_like, str
            Data types.
        """
        nchild = len(self._qids)
        qids  = [None] * nchild
        dates = [None] * nchild
        descs = [None] * nchild
        types = [None] * nchild
        for i in range(nchild):
            qid = self._qids[i]
            qids[i]  = qid
            dates[i] = self["q"+qid].get_date()
            descs[i] = self["q"+qid].get_desc()
            types[i] = self["q"+qid].get_type()

        return qids, dates, descs, types

class InputNode(_ParentNode):
    """Node that represents an input parent group.

    Input parent groups are "bfield", "efield", etc. in the HDF5 file that can
    contain several input data groups.
    """

    def __init__(self, root, path, h5, **kwargs):
        """Create an input node and its children.

        Parameters
        ----------
        root : `RootNode`
            The `RootNode` this node belongs to.
        path : str
            Path to this node within the HDF5 file.
        h5 : `h5py.File`
            The HDF5 file from which the tree is constructed.
        **kwargs
            Arguments passed to other constructors in case of multiple
            inheritance.
        """
        super().__init__(root=root, path=path, **kwargs)
        parent = h5[path]
        for name, group in parent.items():
            inputtype = fileapi.get_type(name)
            inputobj  = root._create_datagroup(inputtype, group.name)
            self._add_child(name, inputobj)

        self._finalize(h5)
        self._freeze()

    def ls(self, show=True):
        """Get a string representation of the contents.

        Parameters
        ----------
        show : str, optional
            If True, the contents are also printed on screen.

        Returns
        -------
        contents : str
            Multiline string decorated with ANSI escape sequences that list
            all data groups within this node.
        """
        out = ""
        for i, q in enumerate(self._qids):
            date  = self["q"+q].get_date()
            desc  = self["q"+q].get_desc()
            gtype = self["q"+q].get_type()
            out += _FancyText.header(gtype.ljust(10) + " " + q)
            out += " " + date
            if q == self.active.get_qid():
                out += _FancyText.active(" [active]")

            out += "\n" + desc + "\n"

        if show: print(out)
        return out

class ResultNode(_Node, DataGroup):
    """Node that represents a results group.

    Result groups contain the data groups containing simulation results. Result
    node has its own meta data (QID, etc.) and its children can only be
    referenced by name e.g. `resultnode.inistate`. Result node also stores QIDs
    of the input groups used to obtain that result.
    """

    def __init__(self, root, path, h5, inputqids, **kwargs):
        """Create a result node and its children, and store references to
        the input data.

        Parameters
        ----------
        root : `RootNode`
            The `RootNode` this node belongs to.
        path : str
            Path to this node within the HDF5 file.
        h5 : `h5py.File`
            The HDF5 file from which the tree is constructed.
        inputqids : dict [str, str]
            Dictionary containing the name of the input parent group
            (e.g. "bfield") and the QID of the input used for this result.
        **kwargs
            Arguments passed to other constructors in case of multiple
            inheritance.
        """
        super().__init__(root=root, path=path, **kwargs)

        # Put references to the input data
        for name, qid in inputqids.items():
            self[name] = root[name]["q"+qid]

        # Store output data
        for name, group in h5[path].items():
            self["_" + name] = root._create_datagroup(name, group.name)

        self._freeze()

    def get_contents(self):
        """Return names of all childs.

        Returns
        -------
        names : list [str]
            List of names.
        """
        out = []
        for i in fileapi.OUTPUTGROUPS:
            if i in self: out.append(i)

        return out

    def ls(self, show=True):
        """Get a string representation of the contents.

        Parameters
        ----------
        show : str, optional
            If True, the contents are also printed on screen.

        Returns
        -------
        contents : str
            Multiline string decorated with ANSI escape sequences that list
            this node's meta data, all output data within this node, and inputs
            that were used.
        """
        qid   = self.get_qid()
        date  = self.get_date()
        desc  = self.get_desc()
        gtype = self.get_type()
        out = ""
        out += _FancyText.header(gtype.ljust(10) + " " + qid)
        out += " " + date
        out += "\n" + desc + "\n"

        out += _FancyText.title("Contents:\n")
        for cnt in self.get_contents():
            out += cnt + "\n"

        out += _FancyText.title("Input:\n")
        for inp in fileapi.INPUTGROUPS:
            if not inp in self:
                continue
            qid   = self[inp].get_qid()
            date  = self[inp].get_date()
            desc  = self[inp].get_desc()
            gtype = self[inp].get_type()
            out += _FancyText.active(inp.ljust(8))
            out += _FancyText.header(gtype.ljust(10) + " " + qid)
            out += " " + date
            out += "\n" + desc + "\n"

        if show: print(out)
        return out

class RootNode(_ParentNode):
    """The entry node for accessing data in the HDF5 file.

    The root node spawns the treeview creating all other nodes and data groups.

    Attributes
    ----------
    _ascot : `Ascot`
        The `Ascot` object whose HDF5 file is what this tree structure
        represents.
    """

    def __init__(self, ascot, **kwargs):
        """Initialize the tree structure based on the HDF5 file.

        This opens the HDF5 file obtained from `Ascot.file_getpath` and builds
        the tree structure. Nothing is done if the filename is `None`

        Parameters
        ----------
        ascot : `Ascot`
            The `Ascot` object whose HDF5 file is what this tree structure
            represents.
        **kwargs
            Arguments passed to other constructors in case of multiple
            inheritance.
        """
        super().__init__(root=self, path="/", **kwargs)
        self._ascot = ascot

        fn = self._ascot.file_getpath()
        if fn is not None:
            h5py.File(fn, "r") # Try opening the file
            self._build(fn)

    def _build(self, fn):
        """(Re-)build node structure from file.

        Use this liberally to update the treeview every time the contents have
        changed.

        Parameters
        ----------
        fn : str
            The HDF5 file from which the tree is constructed.
        """
        ascot = self._ascot
        for v in list(vars(self)):
            delattr(self, v)

        super().__init__(root=self, path="/results")
        self._ascot = ascot
        if fn is None:
            return

        with h5py.File(fn, "r") as h5:

            # Initialize input groups.
            for inp in h5.keys():
                if inp in fileapi.INPUTGROUPS:
                    self[inp] = self._create_inputgroup(inp, h5)

            if "results" in h5:
                results = h5["results"]
                for run in results.keys():

                    # Fetch those input groups that correspond to this run.
                    inputqids = fileapi.get_inputqids(h5, results[run])

                    # Make a result group
                    runnode = self._create_resultgroup(
                        results[run].name, h5, inputqids)
                    self._add_child(run, runnode)

                self._finalize(h5)

        self._freeze()

    def _destroy_group(self, group, repack=True):
        """Remove a group from the HDF5 file and rebuild the tree.

        Parameters
        ----------
        group : str
            Name of the group to be removed.
        repack : bool, optional
            If True, repack the HDF5 file.
        """
        fn = self._ascot.file_getpath()
        with h5py.File(fn, "a") as f:
            fileapi.remove_group(f, group)

        if repack:
            fntemp = fn + "_repack"
            subprocess.call(["h5repack", fn, fntemp], stdout=subprocess.DEVNULL)
            subprocess.call(["mv", fntemp, fn])

        self._build(fn)

    def _activate_group(self, group):
        """Set group as active and rebuild the tree.

        Parameters
        ----------
        group : str
            Name or QID of the group to be activated.
        """
        fn = self._ascot.file_getpath()
        with h5py.File(fn, "a") as f:
            fileapi.set_active(f, group)

        self._build(fn)

    def _get_group(self, name):
        """Fetch group based on its QID or name.

        Parameters
        ----------
        group : str
            Name or QID of the group to be fetched.

        Returns
        -------
        group : DataGroup
            The group that was fetched.
        """
        qid = "q" + fileapi.get_qid(name)
        if qid in self:
            return self[qid]
        for parent in fileapi.INPUTGROUPS:
            if parent in self and qid in self[parent]:
                return self[parent][qid]

    def _create_inputgroup(self, path, h5):
        """Create an input group to be added to the treeview.

        Parameters
        ----------
        path : str
            Path to the input group within the HDF5 file.
        h5 : `h5py.File`
            The HDF5 file from which the tree is constructed.

        Returns
        -------
        group : `InputNode`
            The input group that was created
        """
        pass

    def _create_resultgroup(self, path, h5, inputgroups):
        """Create a result group to be added to the treeview.

        Parameters
        ----------
        path : str
            Path to the result node within the HDF5 file.
        h5 : `h5py.File`
            The HDF5 file from which the tree is constructed.
        inputqids : dict [str, str]
            Dictionary containing the name of the input parent group
            (e.g. "bfield") and the QID of the input used for this result.

        Returns
        -------
        group : `ResultNode`
            The result group that was created.
        """
        pass

    def _create_datagroup(self, grouptype, group):
        """Create data group based on the given type.

        Parameters
        ----------
        grouptype : str
            Type of the group as it appears in `HDF5TOOBJ`.
        path : str
            Path to the data in the HDF5 file that corresponds to the group.

        Returns
        -------
        `DataContainer`
            The data group that was created.
        """
        pass

    @staticmethod
    def _create_file(fn):
        """Create an empty HDF5 file.
        """
        f = h5py.File(fn, "w")
        f.close()

    def get_runs(self, inputqid):
        """Fetch QIDs of runs using given input.

        Parameters
        ----------
        inputqid : `qid`
            QID of the input.

        Returns
        -------
        qids : list [str]
            List of run QIDs.
        """
        # Find the parent group first
        parent  = None
        qid = "q" + inputqid
        for p in fileapi.INPUTGROUPS:
            if p in self and qid in self[p]:
                parent = p
        if parent is None:
            return []

        # Find the runs
        runqids = []
        for qid in self._qids:
            if parent in self["q"+qid] and \
               self["q"+qid][parent].get_qid() == inputqid:
                runqids.append(qid)

        return runqids

    def destroy(self, repack=True):
        """Remove all results from the HDF5 file.

        Parameters
        ----------
        repack : bool, optional
            If True, repack the HDF5 file.

            Removing data from the HDF5 file only removes references to it and
            repacking is required to free the disk space. Repacking makes a copy
            of the HDF5 file and destroys the original, and it also requires
            3rd party tool `h5repack` which is why it's use is optional here.
        """
        self._destroy_group("results", repack=repack)

    def ls(self, show=True):
        """Get a string representation of the contents.

        Parameters
        ----------
        show : str, optional
            If True, the contents are also printed on screen.

        Returns
        -------
        contents : str
            Multiline string decorated with ANSI escape sequences that list
            all results and their meta data and currently active inputs.
        """
        out = ""
        out += _FancyText.title("Inputs:")
        out += _FancyText.active(" [only active shown]\n")
        for inp in fileapi.INPUTGROUPS:
            if not inp in self:
                continue

            ngrp  = len(self[inp]._qids)
            qid   = self[inp].active.get_qid()
            date  = self[inp].active.get_date()
            desc  = self[inp].active.get_desc()
            gtype = self[inp].active.get_type()
            out += _FancyText.active(inp.ljust(8))
            out += _FancyText.header(gtype.ljust(10) + " " + qid)
            out += " " + date
            out += "\n" + desc + "\n"
            out += "+ " + str(ngrp-1) + " other(s)\n"

        out += _FancyText.title("Results:\n")
        for i, q in enumerate(self._qids):
            date  = self["q"+q].get_date()
            desc  = self["q"+q].get_desc()
            gtype = self["q"+q].get_type()
            out += _FancyText.header(gtype.ljust(10) + " " + q)
            out += " " + date
            if q == self.active.get_qid():
                out += _FancyText.active(" [active]")

            out += "\n" + desc + "\n"

        if show: print(out)
        return out

a5py/a5py/ascot5io/coreio/treeview.py
ascot4fusion/ascot5
"""Distribution output.
"""
import numpy as np
import unyt

import itertools
from scipy.interpolate import griddata, RectBivariateSpline
from a5py import physlib

from .coreio.treedata import DataContainer

class DistData():
    """Distribution data object.

    Attributes
    ----------
    abscissae : list [str]
        Name of each abscissa in same order as they appear in the distribution.
    _distribution : array_like
        N-dimensional array, where N is the number of abscissae, containing
        the distribution data.
    **abscissa_edges : array_like
        Edges of each abscissa stored as attributes with name
        "_<name of the abscissa>" e.g. "_r".
    """

    def __init__(self, histogram, **abscissa_edges):
        """Initialize distribution data from given abscissae and histogram.

        Parameters
        ----------
        histogram : array_like
            N-dimensional array, where N is the number of abscissae, containing
            the distribution data as a histogram.
        **abscissa_edges : array_like
            Name and edges of each abscissa.
        """
        self.abscissae = list(abscissa_edges.keys())
        for k in abscissa_edges:
            setattr(self, "_" + k, abscissa_edges[k])

        self._distribution = histogram / self.phasespacevolume()

    def _copy(self):
        """Make a (deep) copy of this object.

        Returns
        -------
        copy : :class:`DistData`
            Copy of this object.
        """
        abscissa_edges = {}
        for name in self.abscissae:
            abscissa_edges[name] = getattr(self, "_" + name).copy()
        histogram = self.histogram().copy()

        return DistData(histogram, **abscissa_edges)

    def _multiply(self, fac, *abscissae):
        """Multiply distribution with a value.

        This method is used to calculate moments from the distribution.
        """
        ranges = []; indices = []
        for i, a in enumerate(self.abscissae):
            if a not in abscissae:
                ranges.append(range(self.abscissa(a).size))
                indices.append(i)

        idx = [slice(None)] * len(self.abscissae)
        for itr in itertools.product(*ranges):
            for k, i in enumerate(indices):
                idx[i] = itr[k]
            self._distribution[tuple(idx)] *= fac.v

        self._distribution *= fac.units

    def abscissa(self, name):
        """Return abscissa values.

        Parameters
        ----------
        name : str
            Name of the abscissa.

        Returns
        -------
        abscissa : array_like
            Abscissa values, i.e., the values at the center of each bin.

        Raises
        ------
        ValueError
            If ``name`` does not correspond to any abscissa.
        """
        abscissa = "_" + name
        if not hasattr(self, abscissa):
            raise ValueError("Unknown abscissa: " + name)
        edges = getattr(self, abscissa)
        return ( edges[1:] + edges[:-1] ) / 2

    def abscissa_edges(self, name):
        """Return abscissa edges.

        Parameters
        ----------
        name : str
            Name of the abscissa.

        Returns
        -------
        abscissa : array_like
            Abscissa bin edges.

        Raises
        ------
        ValueError
            If ``name`` does not correspond to any abscissa.
        """
        abscissa = "_" + name
        if not hasattr(self, abscissa):
            raise ValueError("Unknown abscissa: " + name)
        return getattr(self, abscissa).copy()

    def distribution(self):
        """Return the distribution function.

        Returns
        -------
        dist : array_like
            N-dimensional array, where N is the number of abscissae, containing
            the distribution data.
        """
        return self._distribution

    def histogram(self):
        """Return the distribution as a histogram (particles per bin).

        Returns
        -------
        dist : array_like
            N-dimensional array, where N is the number of abscissae, containing
            the distribution data as a histogram.
        """
        return self._distribution * self.phasespacevolume()

    def slice(self, copy=False, **abscissae):
        """Take slices of distribution.

        Parameters
        ----------
        copy : bool, optional
            Retain original distribution and return a copy which is sliced.
        **abscissae : slice
            Name of the coordinate and corresponding slice to be taken.

        Returns
        -------
        dist : :class:`DistData`
            The sliced distribution if ``copy=True``.
        """
        dist = self if not copy else self._copy()

        for k, s in abscissae.items():
            dim = dist.abscissae.index(k)
            idx = np.histogram(dist.abscissa(k)[s], dist.abscissa_edges(k))[0]
            idx = np.nonzero(idx)[0]

            endpoints = dist.abscissa_edges(k)[[0,-1]]
            midpoints = dist.abscissa(k)[s]
            edges = np.zeros((midpoints.size+1,))
            edges[0]  = endpoints[0]
            edges[-1] = endpoints[-1]
            if edges.size > 2:
                edges[1:-1] = ( midpoints[1:] + midpoints[:-1] ) / 2
            setattr(dist, "_" + k, edges)

            dist._distribution = dist._distribution.take(indices=idx, axis=dim)

        if copy: return dist

    def integrate(self, copy=False, **abscissae):
        """Integrate distribution along the given dimension.

        Parameters
        ----------
        copy : bool, optional
            Retain original distribution and return a copy which is integrated.
        **abscissae : slice or array_like
            Name of the coordinate and corresponding slice which is integrated.

            If argument is an array, it must have the same size as the
            corresponding dimension. The integration is then performed as
            int f(x,...) * array(x) * dx.

        Returns
        -------
        dist : :class:`DistData`
            The integrated distribution if ``copy=True``.
        """
        for k in abscissae.keys():
            if not hasattr(self, "_" + k):
                raise ValueError("Unknown abscissa: " + k)

        dist = self if not copy else self._copy()
        dim = 0
        for k in tuple(dist.abscissae):
            if k not in abscissae.keys():
                dim +=1
                continue

            idx = abscissae[k]
            if isinstance(idx, slice):
                ds   = np.diff(dist.abscissa_edges(k))
                mask = np.zeros(ds.shape)
                mask[idx] = 1
                ds  *= mask
            else:
                ds = np.diff(dist.abscissa_edges(k)) * idx

            dist._distribution, s = np.average(
                dist.distribution(), axis=dim, weights=ds, returned=True)

            # np.average seems to strip units from ds most of the time but not
            # always...
            if hasattr(s, "units"):
                dist._distribution *= s
            else:
                dist._distribution *= s * ds.units

            dist.abscissae.remove(k)
            delattr(dist, "_" + k)

        if copy: return dist

    def interpolate(self, **coordinates):
        """Perform (N-dim) linear interpolation on the distribution.

        Parameters
        ----------
        **coordinates : slice
            Name of the coordinate and value at which the distribution is
            interpolated.

        Returns
        -------
        dist : array_like
            The distribution interpolated at given point.
        """
        for k in coordinates.keys():
            if not hasattr(self, "_" + k):
                raise ValueError("Unknown abscissa: " + k)
            val = coordinates[k]
            edges = self.abscissa_edges(k)
            dim = edges.units.dimensions
            try:
                # Try to get argument units
                valdim = val.units.dimensions
            except AttributeError:
                # Argument doesn't have units, assign and add warning
                val = val*unit
                valdim = dim
                warnings.warn(AscotUnitWarning)
            if valdim != dim:
                raise ValueError(
                    "\"%s\" has incorrect dimensions: expected %s but got %s" %
                    (name, dim, valdim))

            if val < edges[0] or val > edges[-1]:
                raise ValueError(
                    "Coordinate out of bounds [%.2f, %.2f]: %s = %.2f",
                    (edges[0], edges[-1], k, val) )

        dim = 0
        dist = self.distribution().copy()
        for k in tuple(self.abscissae):
            if k not in coordinates.keys():
                dim += 1
                continue

            val = coordinates[k]
            edges = self.abscissa_edges(k)
            i = np.argmax(val < edges) - 1
            t = ( val - edges[i] ) / ( edges[i+1] - edges[i] )

            dist0 = dist.take(indices=i,   axis=dim)
            dist1 = dist.take(indices=i+1, axis=dim)

            dist = t * dist0 + (1-t) * dist1
        return dist

    def phasespacevolume(self):
        """Calculate phase-space volume of each bin.

        Returns
        -------
        vol : array_like
            Phase space volumes with units and same shape as the distribution.
        """
        vol = 1
        for name in self.abscissae:
            edges = self.abscissa_edges(name)
            dV    = (edges[1:] - edges[:-1])
            vol   = np.multiply.outer(vol, dV)
        return vol

class DistMoment:
    """Class that stores moments calculated from a distribution.
    """

    def __init__(self, x1, x2, x3, r, phi, z, area, volume, rhodist):
        """Initialize moment storage.

        The real space abscissa edges define where the moments are defined.
        The coordinates of the bin edges are required to calculate moments.
        Area and volume are needed for normalization and to transform from
        phase-space to real space.

        Parameters
        ----------
        x1 : array_like
            First real-space coordinate abscissa edges (R or rho).
        x2 : array_like
            First real-space coordinate abscissa edges (R or rho).
        x3 : array_like
            First real-space coordinate abscissa edges (R or rho).
        r : array_like
            R coordinates corresponding to bin centers.
        phi : array_like
            phi coordinates corresponding to bin centers.
        z : array_like
            z coordinates corresponding to bin centers.
        area : array_like
            Poloidal plane area of each bin.
        volume : array_like
            Volume of each bin.
        rhodist : bool
            Flag indicating whether the moment is in (rho,theta,phi) or
            (R,phi,z) coordinates.
        """
        if rhodist:
            self.rho   = x1
            self.theta = x2
            self.phi   = x3
        else:
            self.r     = x1
            self.phi   = x2
            self.z     = x3

        self.rc        = r
        self.phic      = phi
        self.zc        = z
        self.area      = area
        self.volume    = volume
        self.rhodist   = rhodist

    def ordinate(self, ordinate, toravg=False, polavg=False):
        """Return stored moment.

        Parameters
        ----------
        ordinate : str
            Name of the moment.
        toravg : bool, optional
            Return toroidal average of the ordinate.
        polavg : bool, optional
            Return poloidal average of the ordinate.

            Only valid if ``rhodist=True``. Both ``toravg`` and ``polavg`` can
            be set simultaneously, in which a radial profile is returned.

        Returns
        -------
        data : array_like
            Ordinate data.
        """
        name = "_ordinate_" + ordinate
        if not hasattr(self, name):
            raise ValueError("Unknown ordinate: " + ordinate)
        ordinate = getattr(self, name)

        if toravg:
            if self.rhodist:
                ordinate = np.sum(ordinate * self.volume, axis=2) \
                    / np.sum(self.volume, axis=2)
            else:
                ordinate = np.sum(ordinate * self.volume, axis=1) \
                    / np.sum(self.volume, axis=1)
        if polavg:
            if self.rhodist:
                volume = np.sum(self.volume, axis=2) if toravg else self.volume
                ordinate = np.sum(ordinate * volume, axis=1) \
                    / np.sum(volume, axis=1)
            else:
                raise ValueError("Cannot take poloidal average of non-rhodist")
        return ordinate

    def add_ordinates(self, **ordinates):
        """Add moments.

        Parameters
        ----------
        **ordinates : array_like
            Names and data for each ordinate to be added.
        """
        for ordinate, val in ordinates.items():
            if hasattr(self, "_ordinate_" + ordinate):
                raise ValueError("Ordinate %s is already present" % ordinate)
            setattr(self, "_ordinate_" + ordinate, val)

    def list_ordinates(self):
        """List all moments
        """
        ordinates = []
        for k in self.__dict__.keys():
            if "_ordinate_" in k:
                ordinates.append(k[10:])
        return ordinates

class Dist(DataContainer):

    def get(self):
        """Return the distribution data.

        Returns
        -------
        dist : class:`DistData`
            Distribution data.
        """
        with self as f:
            histogram = np.sum(f["ordinate"][:], axis=0) * unyt.particles
            abscissa_edges = {}
            for i in range(int(f["abscissa_ndim"][:])):
                abscissa = f["abscissa_vec_0"+str(i+1)]
                name     = abscissa.attrs["name_0"+str(i)].decode("utf-8")
                unit     = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
                try:
                    abscissa_edges[name] = abscissa[:] * unyt.Unit(unit)
                except:
                    unit = unit.replace(" ", "*")
                    abscissa_edges[name] = abscissa[:] * unyt.Unit(unit)

        return DistData(histogram, **abscissa_edges)

    @staticmethod
    def density(dist, moment):
        """Calculate number density.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        integrate = {}
        if moment.rhodist:
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi"]:
                    integrate[k] = np.s_[:]
        else:
            for k in dist.abscissae:
                if k not in ["r", "phi", "z"]:
                    integrate[k] = np.s_[:]
        dist = dist.integrate(copy=True, **integrate)
        moment.add_ordinates(density=dist.histogram() / moment.volume)

    @staticmethod
    def chargedensity(dist, moment):
        """Calculate charge density.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        integrate = {}
        dist = dist.integrate(copy=True, charge=dist.abscissa("charge"))
        if moment.rhodist:
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi"]:
                    integrate[k] = np.s_[:]
        else:
            for k in dist.abscissae:
                if k not in ["r", "phi", "z"]:
                    integrate[k] = np.s_[:]
        dist.integrate(**integrate)
        moment.add_ordinates(chargedensity=dist.histogram() / moment.volume)

    @staticmethod
    def energydensity(mass, dist, moment):
        """Calculate energy density.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        integrate = {}
        if moment.rhodist:
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi", "pperp", "ppar"]:
                    integrate[k] = np.s_[:]
        else:
            for k in dist.abscissae:
                if k not in ["r", "z", "phi", "pperp", "ppar"]:
                    integrate[k] = np.s_[:]
        dist = dist.integrate(copy = True, **integrate)
        ppa, ppe = np.meshgrid(dist.abscissa("ppar"), dist.abscissa("pperp"))
        pnorm = np.sqrt(ppa.ravel()**2 + ppe.ravel()**2)
        ekin  = (physlib.gamma_momentum(mass, pnorm) - 1) * mass * unyt.c**2
        dist._multiply(ekin.reshape(ppa.shape), "ppar", "pperp")
        dist.integrate(ppar=np.s_[:], pperp=np.s_[:])
        moment.add_ordinates(
            energydensity=dist.histogram().to("J") / moment.volume)

    @staticmethod
    def pressure(mass, dist, moment):
        """Calculate pressure.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        integrate = {}
        if moment.rhodist:
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi", "pperp", "ppar"]:
                    integrate[k] = np.s_[:]
        else:
            for k in dist.abscissae:
                if k not in ["r", "z", "phi", "pperp", "ppar"]:
                    integrate[k] = np.s_[:]
        dist = dist.integrate(copy = True, **integrate)
        ppa, ppe = np.meshgrid(dist.abscissa("ppar"),
                               dist.abscissa("pperp"))
        pnorm = np.sqrt(ppa.ravel()**2 + ppe.ravel()**2)
        vnorm = physlib.velocity_momentum(mass, pnorm)
        dist._multiply(vnorm.reshape(ppa.shape)**2 * mass / 3, "ppar", "pperp")
        dist.integrate(ppar=np.s_[:], pperp=np.s_[:])
        moment.add_ordinates(pressure=dist.histogram().to("J") / moment.volume)

    @staticmethod
    def toroidalcurrent(ascot, mass, dist, moment):
        """Calculate toroidal current density.

        Parameters
        ----------
        ascot : :class:`Ascot`
            Ascot object for interpolating input data.
        mass : float
            Test particle mass.
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        dist = dist.integrate(copy=True, charge=dist.abscissa("charge"),
                              ppar=dist.abscissa("ppar"))
        bphi, bnorm = ascot.input_eval(
            moment.rc, moment.phic, moment.zc, 0*unyt.s, "bphi", "bnorm")
        bphi  = bphi.reshape(moment.volume.shape)
        bnorm = bnorm.reshape(moment.volume.shape)

        integrate = {}
        if moment.rhodist:
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi"]:
                    integrate[k] = np.s_[:]
        else:
            for k in dist.abscissae:
                if k not in ["r", "phi", "z"]:
                    integrate[k] = np.s_[:]
        dist.integrate(**integrate)
        dist._distribution *= bphi / (bnorm * mass)
        moment.add_ordinates(
            toroidalcurrent=(dist.histogram() / moment.volume).to("A/m**2"))

    @staticmethod
    def parallelcurrent(ascot, mass, dist, moment):
        """Calculate parallel current density.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        dist = dist.integrate(copy=True, charge=dist.abscissa("charge"),
                              ppar=dist.abscissa("ppar"))
        integrate = {}
        if moment.rhodist:
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi"]:
                    integrate[k] = np.s_[:]
        else:
            for k in dist.abscissae:
                if k not in ["r", "phi", "z"]:
                    integrate[k] = np.s_[:]
        dist.integrate(**integrate)
        dist._distribution *= 1/mass

        moment.add_ordinates(
            parallelcurrent=(dist.histogram() / moment.volume).to("A/m**2"))

    @staticmethod
    def powerdep(ascot, mass, dist, moment):
        """Calculate collisional power deposition to plasma.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        dist = dist._copy()
        ppa, ppe = np.meshgrid(dist.abscissa("ppar"),
                               dist.abscissa("pperp"))
        pnorm = np.sqrt(ppa.ravel()**2 + ppe.ravel()**2)
        vnorm = physlib.velocity_momentum(mass, pnorm)
        for i, qa in enumerate(dist.abscissa("charge")):
            k = ascot.input_eval_collcoefs(
                mass, qa, moment.rc.ravel(), moment.phic.ravel(),
                moment.zc.ravel(), np.zeros(moment.rc.ravel().shape)*unyt.s,
                vnorm, "k", grid=True)
            k = -np.sum(k, axis=0) # Minus because k is from plasma to particle
            k = k.ravel().reshape(dist._distribution[:,:,:,:,:,i,0].shape)
            dist._distribution[:,:,:,:,:,i,0] *= k.v

        dist._distribution *= k.units * mass
        dist.integrate(charge=np.s_[:], time=np.s_[:])
        dist._multiply(vnorm.reshape(ppa.shape), "ppar", "pperp")
        dist.integrate(ppar=np.s_[:], pperp=np.s_[:])
        moment.add_ordinates(
            powerdep=(dist.histogram() / moment.volume ).to("W/m**3"))

    @staticmethod
    def electronpowerdep(ascot, mass, dist, moment):
        """Calculate power deposition to electrons.

        Parameters
        ----------
        ascot : :class:`Ascot`
            Ascot object for interpolating input data.
        mass : float
            Test particle mass.
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        dist = dist._copy()
        ppa, ppe = np.meshgrid(dist.abscissa("ppar"),
                               dist.abscissa("pperp"))
        pnorm = np.sqrt(ppa.ravel()**2 + ppe.ravel()**2)
        vnorm = physlib.velocity_momentum(mass, pnorm)
        for i, qa in enumerate(dist.abscissa("charge")):
            k = ascot.input_eval_collcoefs(
                mass, qa, moment.rc.ravel(), moment.phic.ravel(),
                moment.zc.ravel(), np.zeros(moment.rc.ravel().shape)*unyt.s,
                vnorm, "k", grid=True)
            k = -k[0,:,:] # Minus because k is from plasma to particle
            k = k.ravel().reshape(dist._distribution[:,:,:,:,:,i,0].shape)
            dist._distribution[:,:,:,:,:,i,0] *= k.v

        dist._distribution *= k.units * mass
        dist.integrate(charge=np.s_[:], time=np.s_[:])
        dist._multiply(vnorm.reshape(ppa.shape), "ppar", "pperp")
        dist.integrate(ppar=np.s_[:], pperp=np.s_[:])
        moment.add_ordinates(
            electronpowerdep=(dist.histogram() / moment.volume ).to("W/m**3"))

    @staticmethod
    def ionpowerdep(ascot, mass, dist, moment):
        """Calculate power deposition to ions.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        dist = dist._copy()
        ppa, ppe = np.meshgrid(dist.abscissa("ppar"),
                               dist.abscissa("pperp"))
        pnorm = np.sqrt(ppa.ravel()**2 + ppe.ravel()**2)
        vnorm = physlib.velocity_momentum(mass, pnorm)
        for i, qa in enumerate(dist.abscissa("charge")):
            k = ascot.input_eval_collcoefs(
                mass, qa, moment.rc.ravel(), moment.phic.ravel(),
                moment.zc.ravel(), np.zeros(moment.rc.ravel().shape)*unyt.s,
                vnorm, "k", grid=True)
            k = -np.sum(k[1:], axis=0) # Minus because k is from plasma to prt
            k = k.ravel().reshape(dist._distribution[:,:,:,:,:,i,0].shape)
            dist._distribution[:,:,:,:,:,i,0] *= k.v

        dist._distribution *= k.units * mass
        dist.integrate(charge=np.s_[:], time=np.s_[:])
        dist._multiply(vnorm.reshape(ppa.shape), "ppar", "pperp")
        dist.integrate(ppar=np.s_[:], pperp=np.s_[:])
        moment.add_ordinates(
            ionpowerdep=(dist.histogram() / moment.volume ).to("W/m**3"))

    @staticmethod
    def jxbtorque(ascot, mass, dist, moment):
        """Calculate j_rad x B_pol toroidal torque.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        dist = dist._copy()
        gradbr, gradbphi, gradbz, curlbr, curlbphi, curlbz, br, bphi, bz = \
            ascot.input_eval(moment.rc.ravel(), moment.phic.ravel(), moment.zc.ravel(), 0*unyt.s,
                             "gradbr", "gradbphi", "gradbz", "curlbr",
                             "curlbphi", "curlbz", "br", "bphi", "bz")
        bvec = unyt.unyt_array([br, bphi, bz]).T
        bnorm = np.sqrt(np.sum(bvec**2))
        curlb = unyt.unyt_array([gradbr, gradbphi, gradbz]).T
        gradb = unyt.unyt_array([curlbr, curlbphi, curlbz]).T
        curlbhat = ( np.cross(gradb, bvec/bnorm)*unyt.T/unyt.m - curlb ) / bnorm
        for iq, q in enumerate(dist.abscissa("charge")):
            for ippa, ppa in enumerate(dist.abscissa("ppar")):
                Bstar = bvec + (ppa / q) * curlbhat
                for ippe, ppe in enumerate(dist.abscissa("pperp")):
                    gamma = physlib.gamma_momentum(mass, np.sqrt(ppa**2+ppe**2))
                    dr = (ppa / (gamma*mass)) * Bstar[:,0] \
                        / (np.sum(Bstar * bvec / bnorm, axis=1))
                    dz = (ppa / (gamma*mass)) * Bstar[:,2] \
                        / (np.sum(Bstar * bvec / bnorm, axis=1))
                    deltapsi = bz * moment.rc.ravel() * dr - br * moment.rc.ravel() * dz
                    deltapsi = deltapsi.reshape(moment.volume.shape)
                    dist._distribution[:,:,:,ippa, ippe, iq, 0] *= (deltapsi * q).v
        integrate = {}
        dist._distribution *= (deltapsi * q).units
        if moment.rhodist:
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi"]:
                    integrate[k] = np.s_[:]
        else:
            for k in dist.abscissae:
                if k not in ["r", "phi", "z"]:
                    integrate[k] = np.s_[:]
        dist.integrate(**integrate)
        moment.add_ordinates(
            jxbtorque=(dist.histogram() / moment.volume).to("N*m/m**3"))

    @staticmethod
    def collTorque(ascot, mass, dist, moment):
        """Calculate power deposition to ions.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        if not "ekin" in dist.abscissae or not "pitch" in dist.abscissae:
            raise ValueError("Distribution must be in energy-pitch basis.")
        if moment.rhodist:
            va = physlib.vnorm_gamma(
                physlib.gamma_energy( mass, dist.abscissa("ekin") ) )
            pitch = dist.abscissa("pitch")
            dist = dist.integrate(copy=True, pitch=np.s_[:])
            for qa in dist.abscissa("charge"):
                coefs = ascot.input_eval_collcoefs(
                    mass.to("kg"), qa.to("C"), moment.rc.ravel(), moment.phic.ravel(), moment.zc.ravel(),
                    np.zeros(moment.rc.ravel().shape)*unyt.s, va)
                K = np.sum(coefs["K"],axis=0).reshape((*moment.volume.shape,va.size))
                nu = np.sum(coefs["nu"],axis=0).reshape((*moment.volume.shape,va.size))

                dpitch = -nu
                dPpara = mass*K + mass*va*dpitch
                units = dPpara.units
                dist._distribution[:,:,:,:,0,0] *= dPpara.v

            bphi, bnorm = ascot.input_eval(
                moment.rc, moment.phic, moment.zc, 0*unyt.s, "bphi", "bnorm")
            bphi   = bphi.reshape(moment.volume.shape)
            bnorm  = bnorm.reshape(moment.volume.shape)

            #dist.integrate(ekin=va, charge=np.s_[:], time=np.s_[:])
            dist.integrate(ekin=np.s_[:], charge=np.s_[:], time=np.s_[:])
            dist._distribution *= (bphi/bnorm) *moment.rc *units
        else:
            va = physlib.vnorm_gamma(
                physlib.gamma_energy( mass, dist.abscissa("ekin") ) )
            pitch = dist.abscissa("pitch")
            dist = dist.integrate(copy=True, pitch=np.s_[:])
            for qa in dist.abscissa("charge"):
                coefs = ascot.input_eval_collcoefs(
                    mass.to("kg"), qa.to("C"), moment.rc.ravel(), moment.phic.ravel(), moment.zc.ravel(),
                    np.zeros(moment.rc.ravel().shape)*unyt.s, va)
                K = np.sum(coefs["K"],axis=0).reshape((*moment.volume.shape,va.size))
                nu = np.sum(coefs["nu"],axis=0).reshape((*moment.volume.shape,va.size))

                dpitch = -nu
                dPpara = mass*K + mass*va*dpitch
                units = dPpara.units
                dist._distribution[:,:,:,:,0,0] *= dPpara.v

            bphi, bnorm = ascot.input_eval(
                moment.rc, moment.phic, moment.zc, 0*unyt.s, "bphi", "bnorm")
            bphi   = bphi.reshape(moment.volume.shape)
            bnorm  = bnorm.reshape(moment.volume.shape)
            dist.integrate(ekin=np.s_[:], charge=np.s_[:], time=np.s_[:])
            #dist.integrate(ekin=va, charge=np.s_[:], time=np.s_[:])

            r = np.transpose(moment.rc, (1,0,2))
            dist._distribution *= -(bphi/bnorm) *r *units

        moment.add_ordinates(collTorque=dist.histogram().to("J") / moment.volume)

    @staticmethod
    def canMomentTorque(dist, moment):
        """Calculate power deposition to ions.

        Parameters
        ----------
        dist : :class:`DistData`
            Distribution from where the moments are calculated.
        moment : class:`DistMoment`
            Moment data used in calculation and where the result is stored.
        """
        if not "ekin" in dist.abscissae or not "pitch" in dist.abscissae:
            raise ValueError("Distribution must be in energy-pitch basis.")
        if moment.rhodist:
            # Placeholder
            dist = dist.integrate(copy=True, charge=dist.abscissa("charge"))
            integrate = {}
            for k in dist.abscissae:
                if k not in ["rho", "theta", "phi"]:
                    integrate[k] = np.s_[:]
            dist.integrate(**integrate)
        else:
            # Placeholder
            dist = dist.integrate(copy=True, charge=dist.abscissa("charge"))
            integrate = {}
            for k in dist.abscissae:
                if k not in ["r", "phi", "z"]:
                    integrate[k] = np.s_[:]
            dist.integrate(**integrate)
        moment.add_ordinates(canMomentTorque=-dist.histogram() / moment.volume)

    @staticmethod
    def ppappe2ekinpitch(dist, mass, ekin_edges=10, pitch_edges=10):
        """Convert ppa and ppe distribution abscissae to energy and pitch.

        The conversion is done by first converting the input distribution to
        "markers" - one for each ppa-ppe cell and they have momentum equal
        to the midpoint of the cell. Then these "markers" are rebinned in
        the energy-pitch distribution. This transformation preserves the
        particle number unlike the alternative which would be to interpolate
        the particle density, but which otherwise could be more accurate.

        Energy is in electronvolts and pitch is ppa/(ppa^2 + ppe^2)^0.5. The
        transformation is relativistic.

        Parameters
        ----------
        dist : :class:`DistData`
            A ppar-pperp distribution.
        masskg : float
            Mass of the species (required for energy conversion).

            Note that distribution is assumed to consist of markers with equal
            mass for this conversion to work.
        ekin_edges : array_like or int, optional
            Energy grid edges or the number of bins in the new distribution.

            If edges are not given explicitly, the maximum energy is calculated
            from ppar or pperp abscissa (whichever gives the highest) and the
            minimum is set to zero.
        xi_edges : array_like or int, optional
            Pitch grid edges or the number of bins in the new distribution.

            If edges are not given explicitly, the pitch is binned in the
            [-1, 1] interval.

        Returns
        -------
        exidist : :class:`DistData`
            Energy-pitch distribution which is otherwise equivalent to the
            one given as input.
        """
        if isinstance(ekin_edges, int):
            p2max = np.maximum(dist.abscissa_edges("ppar")[-1]**2,
                               dist.abscissa_edges("ppar")[0]**2,
                               dist.abscissa_edges("pperp")[-1]**2)
            ekinmax = physlib.energy_momentum(mass, np.sqrt(p2max)).to("eV")
            ekin_edges = np.linspace(0, ekinmax, ekin_edges)
        if isinstance(pitch_edges, int):
            pitch_edges = np.linspace(-1, 1, pitch_edges)

        try:
            ekin_edges.units
        except AttributeError:
            ekin_edges *= unyt.eV
        try:
            pitch_edges.units
        except AttributeError:
            pitch_edges *= unyt.dimensionless

        # Create a new empty distribution where ppar and pperp are replaced
        # by ekin and pitch
        dim = []
        abscissa_edges = {}
        for k in dist.abscissae:
            if k == "ppar":
                dim.append(ekin_edges.size-1)
                abscissa_edges["ekin"] = ekin_edges
            elif k == "pperp":
                dim.append(pitch_edges.size-1)
                abscissa_edges["pitch"] = pitch_edges
            else:
                dim.append(dist.abscissa_edges(k).size-1)
                abscissa_edges[k] = dist.abscissa_edges(k)
        exdist = DistData(np.zeros(dim)*unyt.particles, **abscissa_edges)

        # Transform E-xi grid to points in (ppa,ppa) space that are used in
        # interpolation.
        xi, ekin = np.meshgrid(
            exdist.abscissa("pitch"), exdist.abscissa("ekin") )
        pnorm = physlib.pnorm_gamma(
            mass, physlib.gamma_energy(mass, ekin.ravel()) )
        ppa = ( xi.ravel() * pnorm ).to("amu*m/s").v
        ppe = ( np.sqrt(1 - xi.ravel()**2 ) * pnorm).to("amu*m/s").v

        ## NOTE: Jacobian and some other quantities would be needed only for
        ## the interpolation method, which was replaced by this histogram
        ## approach. They are kept here in case there is need to use the
        ## interpolation method in some cases.

        # Coordinate transform Jacobian: dppa dppe = |jac| dE dxi
        # Jacobian for transform (ppa, ppe) -> (p, xi) is p / sqrt(1-xi^2)
        # because jac = dppa / dp  = xi, dppe / dp  = sqrt(1-xi^2)
        #               dppa / dxi = p,  dppe / dxi = -xi p / sqrt(1-xi^2),
        # and the Jacobian for (p, xi) -> (E, xi) is m gamma / p.
        #
        # Therefore the combined Jacobian is:
        # ( m gamma / p ) / sqrt(1-xi*xi).
        jac = (mass + ekin / unyt.c**2) / np.sqrt(1 - xi**2)

        # Quantities needed in iteration on each loop
        ippa     = dist.abscissae.index("ppar")
        units    = dist.distribution().units
        exishape = (exdist.abscissa("ekin").size, exdist.abscissa("pitch").size)
        ppar     = dist.abscissa("ppar").to("amu*m/s").v
        pperp    = dist.abscissa("pperp").to("amu*m/s").v

        ppa0,ppe0 = np.meshgrid(dist.abscissa("ppar"), dist.abscissa("pperp"))
        ppa0 = ppa0.ravel()
        ppe0 = ppe0.ravel()
        pnorm0 = np.sqrt(ppa0**2 + ppe0**2)
        pitch0 = ppa0 / pnorm0
        ekin0 = ((physlib.gamma_momentum(mass, pnorm0) - 1) \
                 * mass * unyt.c**2).to("eV")

        # Use itertools to conveniently make N "for" loops into a single loop
        ranges = []
        for a in dist.abscissae:
            if a != "ppar" and a != "pperp":
                ranges.append(range(dist.abscissa(a).size))

        # We loop over all dimensions except ppar and pperp
        for itr in itertools.product(*ranges):

            # Consctruct indices to slice (ppa, ppa) at this iteration
            idx = [slice(None)] * (len(itr) + 2)
            idx[:ippa]   = itr[:ippa] # Coordinates before ppa
            idx[ippa+2:] = itr[ippa:] # Coordinates after ppa and ppe
            idx = tuple(idx)

            # Interpolate distribution at (ekin, xi) grid points
            #f = RectBivariateSpline(
            #    ppar, pperp, np.squeeze(dist.distribution()[idx]), kx=1, ky=1)
            #d = np.reshape(f.ev(ppa, ppe), exishape) * units

            # Apply jacobian and store the values to exidist
            #exdist._distribution[idx] = d * jac

            d = np.histogram2d(
                ekin0, pitch0, bins=(ekin_edges, pitch_edges),
                weights=np.squeeze(dist.histogram()[idx]).T.ravel())[0]
            exdist._distribution[idx] = d / exdist.phasespacevolume()[idx]

        return exdist

class Dist_5D(Dist):

    def read(self):
        """Read 5D distribution from a HDF5 file to a dictionary.
        """
        out = {}
        with self as f:
            # A Short helper function to calculate grid points from grid edges.
            def edges2grid(edges):
                return np.linspace(0.5*(edges[0]+edges[1]),
                                   0.5*(edges[-2]+edges[-1]), num=edges.size-1)

            abscissae = [0] * int(f["abscissa_ndim"][:])
            for i in range(0, len(abscissae)):
                abscissa     = f["abscissa_vec_0"+str(i+1)]
                name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
                abscissae[i] = name

                unit = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
                out[name + "_edges"] = abscissa[:] * unyt.Unit(unit)
                out[name]            = edges2grid(out[name + "_edges"])
                out["n" + name]      = out[name].size

            out["abscissae"] = abscissae
            out["histogram"] = f["ordinate"][0,:,:,:,:,:,:,:] \
                * unyt.dimensionless

        return out

class Dist_6D(Dist):
    pass

class Dist_rho5D(Dist):
    pass

class Dist_rho6D(Dist):
    pass

class Dist_COM(Dist):

    def read(self):
        """Read COM distribution from a HDF5 file to a dictionary.
        """
        out = {}
        with self as f:
            # A Short helper function to calculate grid points from grid edges.
            def edges2grid(edges):
                return np.linspace(0.5*(edges[0]+edges[1]),
                                   0.5*(edges[-2]+edges[-1]), num=edges.size-1)

            abscissae = [0] * int(f["abscissa_ndim"][:])
            for i in range(0, len(abscissae)):
                abscissa     = f["abscissa_vec_0"+str(i+1)]
                name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
                abscissae[i] = name

                unit = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
                out[name + "_edges"] = abscissa[:] * unyt.Unit(unit)
                out[name]            = edges2grid(out[name + "_edges"])
                out["n" + name]      = out[name].size

            out["abscissae"] = abscissae
            out["histogram"] = f["ordinate"][0,:,:,:] \
                * unyt.dimensionless

        return out

a5py/a5py/ascot5io/dist.py
ascot4fusion/ascot5
"""
5D distribution module.

File: dist_5D.py
"""
import numpy as np
import h5py
import scipy.constants as const

import a5py.dist as distmod
from a5py.physlib.alias import getalias as alias

from .coreio.treedata import DataContainer

def write_hdf5(fn, run, data):
    """
    Write dist5D data in HDF5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
    """

    gname = "results/" + run + "/dist5d"

    with h5py.File(fn, "a") as f:
        g = f.create_group(gname)

        for i in range(0, len(data["abscissae"])):
            name = data["abscissae"][i]
            g.create_dataset("abscissa_nbin_0"+str(i+1), (1,),
                             data=data["n" + name], dtype="i4")
            abscissa = g.create_dataset("abscissa_vec_0"+str(i+1),
                                        (data["n" + name]+1,),
                                        data=data[name + "_edges"], dtype="f8")

            abscissa.attrs["name_0"+str(i)] = np.string_(name)
            abscissa.attrs["unit_0"+str(i)] = np.string_(data[name + "_unit"])

        g.create_dataset("abscissa_ndim", (1,), data=7, dtype="i4")
        g.create_dataset("ordinate_ndim", (1,), data=1, dtype="i4")

        ordinate = g.create_dataset(
            "ordinate", data=np.expand_dims(data["histogram"], axis=0),
            dtype="f8")
        ordinate.attrs["name_00"] = np.string_("density")
        ordinate.attrs["unit_00"] = np.string_(data["ordinate_unit"])


def read_hdf5(fn, qid):
    """
    Read 5D distribution from a HDF5 file to a dictionary.

    Args:
        fn : str <br>
            HDF5 file filename.
        qid : str <br>
            QID of the run whose distribution is read.
    Returns:
        Distribution dictionary.
    """

    with h5py.File(fn,"r") as f:

        path = "/results/run_"+qid+"/dist5d/"
        dist = f[path]
        out = {}

        # A Short helper function to calculate grid points from grid edges.
        def edges2grid(edges):
            return np.linspace(0.5*(edges[0]+edges[1]),
                               0.5*(edges[-2]+edges[-1]), num=edges.size-1)

        abscissae = [0] * int(dist["abscissa_ndim"][:])
        for i in range(0, len(abscissae)):
            abscissa     = dist["abscissa_vec_0"+str(i+1)]
            name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
            abscissae[i] = name

            out[name + "_edges"] = abscissa[:]
            out[name + "_unit"]  = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
            out[name]            = edges2grid(out[name + "_edges"])
            out["n" + name]      = out[name].size

        out["abscissae"] = abscissae
        out["histogram"] = dist["ordinate"][0,:,:,:,:,:,:,:]
        out["histogram_unit"] = dist["ordinate"].attrs["unit_00"].decode("utf-8")

    return out


class Dist_5D(DataContainer):
    """
    Object representing 5D distribution data.
    """

    def __init__(self, root, hdf5):
        """
        Object representing orbit data.
        """
        super().__init__(root, hdf5)


    def read(self):
        """
        Read distribution data from HDF5 file to a dictionary.

        Returns:
            Distribution dictionary.
        """
        return read_hdf5(self._file, self.get_qid())


    def write(self, fn, run, data=None):
        """
        Write dist5D data to HDF5 file.
        """
        if data is None:
            data = self.read()

        write_hdf5(fn, run, data)


    def get_dist(self, dist=None, **kwargs):
        """
        Return distribution dictionary.

        The ordinate is density which is integrated over the dimensions which
        are given in kwargs.

        Args:
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file.
            kwargs : Name(s) (R, phi, z, ppa, ppe, time, charge) of those
               dimensions along which the distribution is either sliced or
               integrated over. Names are given as keyword value pairs where
               a tuple value (a, b) are indices for slicing and array [a, b]
               are indices for integration. A scalar zero means whole
               dimension is integrated.

        Returns:
            Distribution dictionary.
        """
        if not dist:
            dist = distmod.histogram2distribution(self.read())
        distmod.squeeze(dist, **kwargs)

        return dist

    def get_E_xi_dist(self, E_edges=None, xi_edges=None, dist=None, mass=None,
                      **kwargs):
        """
        Return distribution dictionary where (ppa, ppe) is converted to (E, xi).

        This is otherwise same function as Dist_5D.get_dist() except velocity
        components (ppa, ppe) are converted to Energy and pitch (E, xi). Energy
        is in electronvolts and pitch is ppa/(ppa^2 + ppe^2)^0.5.

        Args:
            E_edges : array  (optional) <br>
                Energy grid edges in the new distribution.
            xi_edges : (optional) <br>
                Pitch grid edges in the new distribution.
            dist :(optional) <br>
                Use this distribution instead of reading one from HDF5 file.
            kwargs : <br>
                Name(s) (R, phi, z, ppa, ppe, time, charge) of those dimensions
                along which the distribution is either sliced or integrated
                over. Names are given as keyword value pairs where a tuple value
                (a, b) are indices for slicing and array [a, b]are indices for
                integration. A scalar zero means whole dimension is integrated.

        Returns:
            Distribution dictionary.
        """

        if not dist:
            dist = distmod.histogram2distribution(self.read())

        if mass is None:
            # Mass from inistate.
            masskg = interpret.mass_kg(self._runnode.inistate["mass"][0])
        else:
            masskg = mass

        Exidist = distmod.convert_ppappe_to_Exi(dist, masskg, E_edges, xi_edges)
        distmod.squeeze(Exidist, **kwargs)

        return Exidist

    def plot_dist(self, *args, logscale=False, equal=False, axes=None,
                  dist=None):
        """
        Plot distribution.

        Either makes a 1D or a 2D plot depending on number of input arguments.

        Args:
            args : str, str <br>
                Name of the x-coordinate, and optionally y-coordinate if a 2D
                plot is desired.
            equal : bool, optional <br>
                Make axes equal.
            axes : Axes, optional <br>
                Axes where plotting happens. If None, a new figure will be
                created.
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file. Dimensions that are not x or y are integrated over.
        """
        abscissae = {"r" : 0, "phi" : 0, "z" : 0, "ppar" : 0,
                     "pperp" : 0, "time" : 0, "charge" : 0}

        x = alias(args[0])
        del abscissae[x]
        y = None
        if len(args) > 1:
            y = alias(args[1])
            del abscissae[y]

        if not dist:
            dist = self.get_dist()

        for k in list(abscissae.keys()):
            if k not in dist["abscissae"]:
                del abscissae[k]

        distmod.squeeze(dist, **abscissae)

        if not y:
            distmod.plot_dist_1D(dist, logscale=logscale, axes=axes)
        else:
            distmod.plot_dist_2D(dist, x, y, logscale=logscale, equal=equal,
                                 axes=axes)


    def eval_1d_dist(self, quantity, rhomin, rhomax, nrho):
        from a5py.ascotpy import Ascotpy
        a5 = Ascotpy(self._file)
        a5.init(bfield=self._runnode.bfield.get_qid(),
                plasma=self._runnode.plasma.get_qid())

        ma = self._runnode.inistate["mass"][0]
        #print('mass {} kg'.format(ma)) 
        # ma *= const.physical_constants["atomic mass constant"][0]
        qa = self._runnode.inistate["charge"][0]
        # qa *=  * const.e
        #print('charge {} C'.format(qa)) 
        
        dist = distmod.eval1d(a5, self.get_dist(), quantity,
                              rhomin, rhomax, nrho, ma=ma, qa=qa)

        a5.free(bfield=True, plasma=True)

        return dist


    def plot_1d_dist(self, quantity, rhomin, rhomax, nrho, axes=None):
        dist = self.eval_1d_dist(quantity, rhomin, rhomax, nrho)
        distmod.plot_dist_1D(dist, axes=axes)


    def plot_E_xi_dist(self, *args, E_edges=None, xi_edges=None,
                       logscale=False, equal=False, axes=None, dist=None):
        """
        Convert (ppa, ppe) to (E, xi) and plot the distribution.

        Either makes a 1D or a 2D plot depending on number of input arguments.

        Args:
            args : str, str <br>
                Name of the x-coordinate, and optionally y-coordinate if a 2D
                plot is desired.
            axes : Axes, optional <br>
                Axes where plotting happens. If None, a new figure will be
                created.
            equal : bool, optional <br>
                Make axes equal.
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file. Dimensions that are not x or y are integrated over.
        """
        abscissae = {"r" : 0, "phi" : 0, "z" : 0, "energy" : 0,
                     "pitch" : 0, "time" : 0, "charge" : 0}

        x = alias(args[0])
        del abscissae[x]
        y = None
        if len(args) > 1:
            y = alias(args[1])
            del abscissae[y]

        if not dist:
            dist = self.get_E_xi_dist(E_edges=E_edges, xi_edges=xi_edges)

        for k in list(abscissae.keys()):
            if k not in dist["abscissae"]:
                del abscissae[k]

        distmod.squeeze(dist, **abscissae)

        if not y:
            distmod.plot_dist_1D(dist, logscale=logscale, axes=axes)
        else:
            distmod.plot_dist_2D(dist, x, y, logscale=logscale, equal=equal,
                                 axes=axes)

a5py/a5py/ascot5io/dist_5D.py
ascot4fusion/ascot5
"""
Distribution HDF5 IO module.

File: dist_6D.py
"""
import numpy as np
import h5py

import a5py.dist as distmod
from a5py.physlib.alias import getalias as alias

from .coreio.treedata import DataContainer

def write_hdf5(fn, run, data):
    """
    Write dist6D data in HDF5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
    """

    gname = "results/" + run + "/dist6d"

    with h5py.File(fn, "a") as f:
        g = f.create_group(gname)

        for i in range(0, len(data["abscissae"])):
            name = data["abscissae"][i]
            g.create_dataset("abscissa_nbin_0"+str(i+1), (1,),
                             data=data["n" + name], dtype="i4")
            abscissa = g.create_dataset("abscissa_vec_0"+str(i+1),
                                        (data["n" + name]+1,),
                                        data=data[name + "_edges"], dtype="f8")

            abscissa.attrs["name_0"+str(i+1)] = np.string_(name)
            abscissa.attrs["unit_0"+str(i+1)] = np.string_(data[name + "_unit"])

        g.create_dataset("abscissa_ndim", (1,), data=7, dtype="i4")
        g.create_dataset("ordinate_ndim", (1,), data=1, dtype="i4")

        ordinate = g.create_dataset(
            "ordinate", data=np.expand_dims(data["histogram"], axis=0),
            dtype="f8")
        ordinate.attrs["name_00"] = np.string_("density")
        ordinate.attrs["unit_00"] = np.string_(data["ordinate_unit"])


def read_hdf5(fn, qid):
    """
    Read 6D distribution.

    Args:
        fn : str <br>
            Full path to HDF5 file.
        qid : str <br>
            QID of the run from where the distribution is read.
    Returns:
        Dictionary storing the distributions that were read.
    """

    with h5py.File(fn,"r") as f:

        path = "/results/run_"+qid+"/dist6d/"
        dist = f[path]
        out = {}

        # A Short helper function to calculate grid points from grid edges.
        def edges2grid(edges):
            return np.linspace(0.5*(edges[0]+edges[1]),
                               0.5*(edges[-2]+edges[-1]), num=edges.size-1)

        abscissae = [0] * int(dist["abscissa_ndim"][:])
        for i in range(0, len(abscissae)):
            abscissa     = dist["abscissa_vec_0"+str(i+1)]
            name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
            abscissae[i] = name

            out[name + "_edges"] = abscissa[:]
            out[name + "_unit"]  = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
            out[name]            = edges2grid(out[name + "_edges"])
            out["n" + name]      = out[name].size

        out["abscissae"] = abscissae
        out["histogram"] = dist["ordinate"][0,:,:,:,:,:,:,:,:]
        out["histogram_unit"] = dist["ordinate"].attrs["unit_00"].decode("utf-8")

    return out


class Dist_6D(DataContainer):

    def __init__(self, root, hdf5):
        """
        Object representing orbit data.
        """
        super().__init__(root, hdf5)


    def read(self):
        """
        Read distribution data from HDF5 file to a dictionary.

        Returns:
            Distribution dictionary.
        """
        return read_hdf5(self._file, self.get_qid())


    def write(self, fn, run, data=None):
        """
        Write dist6D data to HDF5 file.
        """
        if data is None:
            data = self.read()

        write_hdf5(fn, run, data)


    def get_dist(self, dist=None, **kwargs):
        """
        Return distribution dictionary.

        The ordinate is density which is integrated over the dimensions which
        are given in kwargs.

        Args:
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file.
            kwargs : Name(s) (R, phi, z, vpa, vpe, time, charge) of those
               dimensions along which the distribution is either sliced or
               integrated over. Names are given as keyword value pairs where
               a tuple value (a, b) are indices for slicing and array [a, b]
               are indices for integration. A scalar zero means whole
               dimension is integrated.

        Returns:
            Distribution dictionary.
        """
        if not dist:
            dist = distmod.histogram2distribution(self.read())
        distmod.squeeze(dist, **kwargs)

        return dist


    def plot_dist(self, *args, logscale=False, equal=False, axes=None,
                  dist=None):
        """
        Plot distribution.

        Either makes a 1D or a 2D plot depending on number of input arguments.

        Args:
            args : str, str <br>
                Name of the x-coordinate, and optionally y-coordinate if a 2D
                plot is desired.
            equal : bool, optional <br>
                Make axes equal.
            axes : Axes, optional <br>
                Axes where plotting happens. If None, a new figure will be
                created.
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file. Dimensions that are not x or y are integrated over.
        """
        abscissae = {"r" : 0, "phi" : 0, "z" : 0, "pr" : 0,
                     "pphi" : 0, "pz" : 0, "time" : 0, "charge" : 0}

        x = alias(args[0])
        del abscissae[x]
        y = None
        if len(args) > 1:
            y = alias(args[1])
            del abscissae[y]

        if not dist:
            dist = self.get_dist()

        for k in list(abscissae.keys()):
            if k not in dist["abscissae"]:
                del abscissae[k]

        distmod.squeeze(dist, **abscissae)

        if not y:
            distmod.plot_dist_1D(dist, logscale=logscale, axes=axes)
        else:
            distmod.plot_dist_2D(dist, x, y, logscale=logscale, equal=equal,
                                 axes=axes)

a5py/a5py/ascot5io/dist_6D.py
ascot4fusion/ascot5
"""
COM distribution module.
"""
import numpy as np
import h5py

import a5py.dist as distmod
from a5py.physlib.alias import getalias as alias

from .coreio.treedata import DataContainer

def write_hdf5(fn, run, data):
    """
    Write dist5D data in HDF5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
    """

    gname = "results/" + run + "/dist5d"

    with h5py.File(fn, "a") as f:
        g = f.create_group(gname)

        for i in range(0, len(data["abscissae"])):
            name = data["abscissae"][i]
            g.create_dataset("abscissa_nbin_0"+str(i+1), (1,),
                             data=data["n" + name], dtype="i4")
            abscissa = g.create_dataset("abscissa_vec_0"+str(i+1),
                                        (data["n" + name]+1,),
                                        data=data[name + "_edges"], dtype="f8")

            abscissa.attrs["name_0"+str(i)] = np.string_(name)
            abscissa.attrs["unit_0"+str(i)] = np.string_(data[name + "_unit"])

        g.create_dataset("abscissa_ndim", (1,), data=7, dtype="i4")
        g.create_dataset("ordinate_ndim", (1,), data=1, dtype="i4")

        ordinate = g.create_dataset(
            "ordinate", data=np.expand_dims(data["histogram"], axis=0),
            dtype="f8")
        ordinate.attrs["name_00"] = np.string_("density")
        ordinate.attrs["unit_00"] = np.string_(data["ordinate_unit"])


def read_hdf5(fn, qid):
    """
    Read 5D distribution from a HDF5 file to a dictionary.

    Args:
        fn : str <br>
            HDF5 file filename.
        qid : str <br>
            QID of the run whose distribution is read.
    Returns:
        Distribution dictionary.
    """

    with h5py.File(fn,"r") as f:

        path = "/results/run_"+qid+"/distcom/"
        dist = f[path]
        out = {}

        # A Short helper function to calculate grid points from grid edges.
        def edges2grid(edges):
            return np.linspace(0.5*(edges[0]+edges[1]),
                               0.5*(edges[-2]+edges[-1]), num=edges.size-1)

        abscissae = [0] * int(dist["abscissa_ndim"][:])
        for i in range(0, len(abscissae)):
            abscissa     = dist["abscissa_vec_0"+str(i+1)]
            name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
            abscissae[i] = name

            out[name + "_edges"] = abscissa[:]
            out[name + "_unit"]  = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
            out[name]            = edges2grid(out[name + "_edges"])
            out["n" + name]      = out[name].size

        out["abscissae"] = abscissae
        out["histogram"] = dist["ordinate"][0,:,:,:,:,:,:,:]
        out["histogram_unit"] = dist["ordinate"].attrs["unit_00"].decode("utf-8")

    return out


class Dist_COM(DataContainer):
    """
    Object representing 5D distribution data.
    """

    def read(self):
        """
        Read distribution data from HDF5 file to a dictionary.

        Returns:
            Distribution dictionary.
        """
        with self as f:

            #path = "/results/run_"+qid+"/distcom/"
            #dist = f[path]
            dist = f
            out = {}

            # A Short helper function to calculate grid points from grid edges.
            def edges2grid(edges):
                return np.linspace(0.5*(edges[0]+edges[1]),
                                   0.5*(edges[-2]+edges[-1]), num=edges.size-1)

            abscissae = [0] * int(dist["abscissa_ndim"][:])
            for i in range(0, len(abscissae)):
                abscissa     = dist["abscissa_vec_0"+str(i+1)]
                name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
                abscissae[i] = name

                out[name + "_edges"] = abscissa[:]
                out[name + "_unit"]  = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
                out[name]            = edges2grid(out[name + "_edges"])
                out["n" + name]      = out[name].size

                out["abscissae"] = abscissae
                out["histogram"] = dist["ordinate"][0,:,:,:]
                out["histogram_unit"] = dist["ordinate"].attrs["unit_00"].decode("utf-8")

        return out


    def write(self, fn, run, data=None):
        """
        Write dist5D data to HDF5 file.
        """
        if data is None:
            data = self.read()

        write_hdf5(fn, run, data)

        
    def plot_muptor(self, E, axes=None):
        """
        Plot constant of motion distribution on (mu, Ptor) plane for a given E.
        """
        import matplotlib.pyplot as plt
        data = self.read()
        hist = data["histogram"]
        mu_edges, Ptor_edges = data["mu_edges"], data["ptor_edges"]
        E_vector, E_edges = data["ekin"], data["ekin_edges"]
        if E is None:
            plot = np.sum(hist, axis=1)
        elif type(E) == int:
            i = E
            plot = hist[:,i,:]
        else:
            i = np.argmax(E < E_edges)-1
            plot1 = hist[:,i,:]
            plot2 = hist[:,i+1,:]

            k = E/(E_vector[i+1]-E_vector[i])
            plot = plot1*k + (1-k)*plot2

        h = axes.pcolormesh(Ptor_edges, mu_edges, plot, shading="flat")
        axes.set_ylabel("mu (J/T)")
        axes.set_xlabel("Ptor (Js)")
        plt.colorbar(h, ax=axes)
        
    def plot_muEkin(self, Ptor, axes=None):
        """
        Plot constant of motion distribution on (mu, Ptor) plane for a given E.
        """
        import matplotlib.pyplot as plt
        data = self.read()
        hist = data["histogram"]
        mu_edges, E_edges = data["mu_edges"], data["ekin_edges"]
        Ptor_vector, Ptor_edges = data["ptor"], data["ptor_edges"]
        if Ptor is None:
            plot = np.sum(hist, axis=2)
        elif type(Ptor) == int:
            i = Ptor
            plot = hist[:,:,i]
        else:
            i = np.argmax(Ptor < Ptor_edges)-1
            plot1 = hist[:,:,i]
            plot2 = hist[:,:,i+1]

            k = Ptor/(Ptor_vector[i+1]-Ptor_vector[i])
            plot = plot1*k + (1-k)*plot2
        
        h = axes.pcolormesh(E_edges, mu_edges, plot, shading="flat")
        axes.set_ylabel("mu (J/T)")
        axes.set_xlabel("E (J)")
        plt.colorbar(h, ax=axes)

        
    def plot_EkinPtor(self, mu, axes=None):
        """
        Plot constant of motion distribution on (mu, Ptor) plane for a given E.
        """
        import matplotlib.pyplot as plt
        data = self.read()
        hist = data["histogram"]
        E_edges, Ptor_edges = data["ekin_edges"], data["ptor_edges"]
        mu_vector, mu_edges = data["mu"], data["mu_edges"]
        if mu is None:
            plot = np.sum(hist, axis=0)
        elif type(mu) == int:
            i = mu
            plot = hist[i,:,:]
        else:
            i = np.argmax(mu < mu_edges)-1
            plot1 = hist[i,:,:]
            plot2 = hist[i+1,:,:]

            k = mu/(mu_vector[i+1]-mu_vector[i])
            plot = plot1*k + (1-k)*plot2
        
        h = axes.pcolormesh(Ptor_edges, E_edges, plot, shading="flat")
        axes.set_ylabel("E (J)")
        axes.set_xlabel("Ptor (Js)")
        plt.colorbar(h, ax=axes)

a5py/a5py/ascot5io/dist_com.py
ascot4fusion/ascot5
"""
Distribution HDF5 IO module.

File: dist_rho5D.py
"""
import numpy as np
import h5py

import a5py.dist as distmod
from a5py.physlib.alias import getalias as alias

from .coreio.treedata import DataContainer

def write_hdf5(fn, run, data):
    """
    Write distrho5D data in HDF5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
    """

    gname = "results/" + run + "/distrho5d"

    with h5py.File(fn, "a") as f:
        g = f.create_group(gname)

        for i in range(0, len(data["abscissae"])):
            name = data["abscissae"][i]
            g.create_dataset("abscissa_nbin_0"+str(i+1), (1,),
                             data=data["n" + name], dtype="i4")
            abscissa = g.create_dataset("abscissa_vec_0"+str(i+1),
                                        (data["n" + name]+1,),
                                        data=data[name + "_edges"], dtype="f8")

            abscissa.attrs["name_0"+str(i+1)] = np.string_(name)
            abscissa.attrs["unit_0"+str(i+1)] = np.string_(data[name + "_unit"])

        g.create_dataset("abscissa_ndim", (1,), data=7, dtype="i4")
        g.create_dataset("ordinate_ndim", (1,), data=1, dtype="i4")

        ordinate = g.create_dataset(
            "ordinate", data=np.expand_dims(data["histogram"], axis=0),
            dtype="f8")
        ordinate.attrs["name_00"] = np.string_("density")
        ordinate.attrs["unit_00"] = np.string_(data["ordinate_unit"])


def read_hdf5(fn, qid):
    """
    Read rho 5D distribution.

    Args:
        fn : str <br>
            Full path to HDF5 file.
        qid : str <br>
            QID of the run where distribution is read.

    Returns:
        Dictionary storing the distributions that were read.
    """

    with h5py.File(fn,"r") as f:

        path = "/results/run_"+qid+"/distrho5d/"
        dist = f[path]
        out = {}

        # A Short helper function to calculate grid points from grid edges.
        def edges2grid(edges):
            return np.linspace(0.5*(edges[0]+edges[1]),
                               0.5*(edges[-2]+edges[-1]), num=edges.size-1)

        abscissae = [0] * int(dist["abscissa_ndim"][:])
        for i in range(0, len(abscissae)):
            abscissa     = dist["abscissa_vec_0"+str(i+1)]
            name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
            abscissae[i] = name

            out[name + "_edges"] = abscissa[:]
            out[name + "_unit"]  = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
            out[name]            = edges2grid(out[name + "_edges"])
            out["n" + name]      = out[name].size

        out["abscissae"] = abscissae
        out["histogram"] = dist["ordinate"][0,:,:,:,:,:,:,:]
        out["histogram_unit"] = dist["ordinate"].attrs["unit_00"].decode("utf-8")

    return out


class Dist_rho5D(DataContainer):
    """
    """

    def __init__(self, root, hdf5):
        """
        Object representing orbit data.
        """
        super().__init__(root, hdf5)


    def read(self):
        """
        Read distribution data from HDF5 file to a dictionary.

        Returns:
            Distribution dictionary.
        """
        return read_hdf5(self._file, self.get_qid())


    def write(self, fn, run, data=None):
        """
        Write distrho5D data to HDF5 file.
        """
        if data is None:
            data = self.read()

        write_hdf5(fn, run, data)


    def get_dist(self, dist=None, **kwargs):
        """
        Return distribution dictionary.

        The ordinate is density which is integrated over the dimensions which
        are given in kwargs.

        Args:
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file.
            kwargs : Name(s) (R, phi, z, vpa, vpe, time, charge) of those
               dimensions along which the distribution is either sliced or
               integrated over. Names are given as keyword value pairs where
               a tuple value (a, b) are indices for slicing and array [a, b]
               are indices for integration. A scalar zero means whole
               dimension is integrated.

        Returns:
            Distribution dictionary.
        """
        if not dist:
            dist = distmod.histogram2distribution(self.read())
        distmod.squeeze(dist, **kwargs)

        return dist


    def get_E_xi_dist(self, E_edges=None, xi_edges=None, dist=None,
                      **kwargs):
        """
        Return distribution dictionary where (ppa, ppe) is converted to (E, xi).

        This is otherwise same function as Dist_5D.get_dist() except velocity
        components (vpa, vpe) are converted to Energy and pitch (E, xi). Energy
        is in electronvolts and pitch is ppa/(ppa^2 + ppe^2)^0.5.

        Args:
            E_edges : array  (optional) <br>
                Energy grid edges in the new distribution.
            xi_edges : (optional) <br>
                Pitch grid edges in the new distribution.
            dist :(optional) <br>
                Use this distribution instead of reading one from HDF5 file.
            kwargs : <br>
                Name(s) (R, phi, z, ppa, ppe, time, charge) of those dimensions
                along which the distribution is either sliced or integrated
                over. Names are given as keyword value pairs where a tuple value
                (a, b) are indices for slicing and array [a, b]are indices for
                integration. A scalar zero means whole dimension is integrated.

        Returns:
            Distribution dictionary.
        """

        if not dist:
            dist = distmod.histogram2distribution(self.read())

        # Mass from inistate.
        masskg = interpret.mass_kg(self._runnode.inistate["mass"][0])

        Exidist = distmod.convert_ppappe_to_Exi(dist, masskg, E_edges, xi_edges)
        distmod.squeeze(Exidist, **kwargs)

        return Exidist


    def plot_dist(self, *args, logscale=False, equal=False, axes=None,
                  dist=None,label=None):
        """
        Plot distribution.

        Either makes a 1D or a 2D plot depending on number of input arguments.

        Args:
            args : str, str <br>
                Name of the x-coordinate, and optionally y-coordinate if a 2D
                plot is desired.
            equal : bool, optional <br>
                Make axes equal.
            axes : Axes, optional <br>
                Axes where plotting happens. If None, a new figure will be
                created.
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file. Dimensions that are not x or y are integrated over.
        """
        abscissae = {"rho" : 0, "theta" : 0, "phi" : 0, "ppar" : 0,
                     "pperp" : 0, "time" : 0, "charge" : 0}

        x = alias(args[0])
        del abscissae[x]
        y = None
        if len(args) > 1:
            y = alias(args[1])
            del abscissae[y]

        if not dist:
            dist = self.get_dist()

        for k in list(abscissae.keys()):
            if k not in dist["abscissae"]:
                del abscissae[k]

        distmod.squeeze(dist, **abscissae)

        if not y:
            distmod.plot_dist_1D(dist, logscale=logscale, axes=axes,label=label)
        else:
            distmod.plot_dist_2D(dist, x, y, logscale=logscale, equal=equal,
                                 axes=axes)


    def plot_E_xi_dist(self, *args, E_edges=None, xi_edges=None,
                       logscale=False, equal=False, axes=None, dist=None):
        """
        Convert (vpa, vpe) to (E, xi) and plot the distribution.

        Either makes a 1D or a 2D plot depending on number of input arguments.

        Args:
            args : str, str <br>
                Name of the x-coordinate, and optionally y-coordinate if a 2D
                plot is desired.
            axes : Axes, optional <br>
                Axes where plotting happens. If None, a new figure will be
                created.
            equal : bool, optional <br>
                Make axes equal.
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file. Dimensions that are not x or y are integrated over.
        """
        abscissae = {"rho" : 0, "theta" : 0, "phi" : 0, "energy" : 0,
                     "pitch" : 0, "time" : 0, "charge" : 0}

        x = alias(args[0])
        del abscissae[x]
        y = None
        if len(args) > 1:
            y = alias(args[1])
            del abscissae[y]

        if not dist:
            dist = self.get_E_xi_dist(E_edges=E_edges, xi_edges=xi_edges)

        for k in abscissae.keys():
            if k not in dist["abscissae"]:
                del abscissae[k]

        distmod.squeeze(dist, **abscissae)

        if not y:
            distmod.plot_dist_1D(dist, logscale=logscale, axes=axes)
        else:
            distmod.plot_dist_2D(dist, x, y, logscale=logscale, equal=equal,
                                 axes=axes)



    def eval_1d_dist(self, quantity,rmin, rmax, zmin, zmax):
        '''
            The R/z min/max are needed for numerical calculation of the rho slot volumes.
            They should cover the whole plasma.
        '''
        from a5py.ascotpy import Ascotpy
        a5 = Ascotpy(self._file)
        a5.init(bfield=self._runnode.bfield.get_qid(),
                plasma=self._runnode.plasma.get_qid())

        ma = self._runnode.inistate["mass"][0]
        print('mass {} kg'.format(ma)) 
        # ma *= const.physical_constants["atomic mass constant"][0]
        qa = self._runnode.inistate["charge"][0]
        # qa *=  * const.e
        print('charge {} C'.format(qa)) 

        

        dist = distmod.eval1d_fromRhodist(a5, self.get_dist(), quantity, 
                                          rmin, rmax, zmin, zmax, ma=ma, qa=qa)


        a5.free(bfield=True, plasma=True)
        
        return dist
        
a5py/a5py/ascot5io/dist_rho5D.py
ascot4fusion/ascot5
"""
Distribution HDF5 IO module.

File: dist_rho6D.py
"""
import numpy as np
import h5py

import a5py.dist as distmod

from .coreio.treedata import DataContainer

def write_hdf5(fn, run, data):
    """
    Write distrho6D data in HDF5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
    """

    gname = "results/" + run + "/distrho6d"

    with h5py.File(fn, "a") as f:
        g = f.create_group(gname)

        for i in range(0, len(data["abscissae"])):
            name = data["abscissae"][i]
            g.create_dataset("abscissa_nbin_0"+str(i+1), (1,),
                             data=data["n" + name], dtype="i4")
            abscissa = g.create_dataset("abscissa_vec_0"+str(i+1),
                                        (data["n" + name]+1,),
                                        data=data[name + "_edges"], dtype="f8")

            abscissa.attrs["name_0"+str(i+1)] = np.string_(name)
            abscissa.attrs["unit_0"+str(i+1)] = np.string_(data[name + "_unit"])

        g.create_dataset("abscissa_ndim", (1,), data=7, dtype="i4")
        g.create_dataset("ordinate_ndim", (1,), data=1, dtype="i4")

        ordinate = g.create_dataset(
            "ordinate", data=np.expand_dims(data["histogram"], axis=0),
            dtype="f8")
        ordinate.attrs["name_00"] = np.string_("density")
        ordinate.attrs["unit_00"] = np.string_(data["ordinate_unit"])


def read_hdf5(fn, qid):
    """
    Read rho 6D distribution.

    Args:
        fn : str
            Full path to HDF5 file.
        qid : str
            QID of the run where distribution is read.
    Returns:
        Dictionary storing the distributions that were read.
    """

    with h5py.File(fn,"r") as f:

        path = "/results/run_"+qid+"/distrho6d/"
        dist = f[path]
        out = {}

        # A Short helper function to calculate grid points from grid edges.
        def edges2grid(edges):
            return np.linspace(0.5*(edges[0]+edges[1]),
                               0.5*(edges[-2]+edges[-1]), num=edges.size-1)

        abscissae = [0] * int(dist["abscissa_ndim"][:])
        for i in range(0, len(abscissae)):
            abscissa     = dist["abscissa_vec_0"+str(i+1)]
            name         = abscissa.attrs["name_0"+str(i)].decode("utf-8")
            abscissae[i] = name

            out[name + "_edges"] = abscissa[:]
            out[name + "_unit"]  = abscissa.attrs["unit_0"+str(i)].decode("utf-8")
            out[name]            = edges2grid(out[name + "_edges"])
            out["n" + name]      = out[name].size

        out["abscissae"] = abscissae
        out["histogram"] = dist["ordinate"][0,:,:,:,:,:,:,:,:]
        out["histogram_unit"] = dist["ordinate"].attrs["unit_00"].decode("utf-8")

    return out


class Dist_rho6D(DataContainer):
    """
    """

    def __init__(self, root, hdf5):
        """
        Object representing orbit data.
        """
        super().__init__(root, hdf5)


    def read(self):
        """
        Read distribution data from HDF5 file to a dictionary.

        Returns:
            Distribution dictionary.
        """
        return read_hdf5(self._file, self.get_qid())


    def write(self, fn, run, data=None):
        """
        Write distrho6D data to HDF5 file.
        """
        if data is None:
            data = self.read()

        write_hdf5(fn, run, data)


    def get_dist(self, dist=None, **kwargs):
        """
        Return distribution dictionary.

        The ordinate is density which is integrated over the dimensions which
        are given in kwargs.

        Args:
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file.
            kwargs : Name(s) (R, phi, z, vpa, vpe, time, charge) of those
               dimensions along which the distribution is either sliced or
               integrated over. Names are given as keyword value pairs where
               a tuple value (a, b) are indices for slicing and array [a, b]
               are indices for integration. A scalar zero means whole
               dimension is integrated.

        Returns:
            Distribution dictionary.
        """
        if not dist:
            dist = distmod.histogram2distribution(self.read())
        distmod.squeeze(dist, **kwargs)

        return dist


    def plot_dist(self, *args, logscale=False, equal=False, axes=None, dist=None):
        """
        Plot distribution.

        Either makes a 1D or a 2D plot depending on number of input arguments.

        Args:
            args : str, str <br>
                Name of the x-coordinate, and optionally y-coordinate if a 2D
                plot is desired.
            equal : bool, optional <br>
                Make axes equal.
            axes : Axes, optional <br>
                Axes where plotting happens. If None, a new figure will be
                created.
            dist : dict_like, optional <br>
               Give input distribution explicitly instead of reading one from
               HDF5 file. Dimensions that are not x or y are integrated over.
        """
        abscissae = {"rho" : 0, "theta" : 0, "phi" : 0, "pr" : 0,
                     "pphi" : 0, "pz" : 0, "time" : 0, "charge" : 0}

        x = args[0]
        del abscissae[x]
        y = None
        if len(args) > 1:
            y = args[1]
            del abscissae[y]

        if not dist:
            dist = self.get_dist()

        for k in list(abscissae.keys()):
            if k not in dist["abscissae"]:
                del abscissae[k]

        distmod.squeeze(dist, **abscissae)

        if not y:
            distmod.plot_dist_1D(dist, logscale=logscale, axes=axes)
        else:
            distmod.plot_dist_2D(dist, x, y, logscale=logscale, equal=equal,
                                 axes=axes)

a5py/a5py/ascot5io/dist_rho6D.py
ascot4fusion/ascot5
"""Electric field input.

Electric field is present in all simulations and it is automatically accounted
for when orbit-following is enabled. To "disable" electric field, set the values
to zero everywhere. The easiest and most computationally effective way to
accomplish this is to use E_TC input.
"""
import h5py
import numpy as np

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class E_TC(DataGroup):
    """Uniform electric field in Cartesian basis.

    This input fixes the electric field vector, in Cartesian basis, so that the
    field has same value and direction everywhere. This input is meant for
    testing purposes or for disabling electric field in simulations. To disable
    electric field, use :meth:`write_hdf5_dummy`.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        return out

    @staticmethod
    def write_hdf5(fn, exyz, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        exyz : array_like (3,1)
            Electric field value in cartesian coordinates [V/m].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if exyz.shape != (3,) and exyz.shape != (3,1):
            raise ValueError("Exyz has wrong shape.")

        parent = "efield"
        group  = "E_TC"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("exyz", (3,1), data=exyz, dtype="f8")

        return gname


    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return E_TC.write_hdf5(fn=fn, exyz=np.array([0,0,0]), desc="DUMMY")

class E_3D(DataGroup):
    """3D electric field that is linearly interpolated.

    This input tabulates the electric-field components on an uniform cylindrical
    grid, and then uses trilinear interpolation to calculate the values during
    the simulation.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        out["er"]   = np.transpose(out["er"],   (2,1,0))
        out["ephi"] = np.transpose(out["ephi"], (2,1,0))
        out["ez"]   = np.transpose(out["ez"],   (2,1,0))
        return out

    @staticmethod
    def write_hdf5(fn, rmin, rmax, nr, zmin, zmax, nz, phimin, phimax, nphi,
                   er, ephi, ez, desc=None):
        """Write input data to the HDF5 file.

        The toroidal angle phi is treated as a periodic coordinate, meaning
        ``A(phi=phimin) == A(phi=phimax)``. However, the phi grid, where input
        arrays are tabulated, is ``linspace(phimin, phimax, nphi+1)[:-1]``
        to avoid storing duplicate data.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        rmin : float
            Minimum value in R grid [m].
        rmax : float
            Maximum value in R grid [m].
        nr : int
            Number of R grid points.
        zmin : float
            Minimum value in z grid [m].
        zmax : float
            Maximum value in z grid [m].
        nz : int
            Number of z grid points.
        phimin : float
            Beginning of the toroidal period [deg].
        phimax : float
            End of the toroidal period [deg].
        nphi : int
            Number of phi grid points.
        er : array_like (nr,nphi,nz)
            Electric field R component [V/m].
        ephi : array_like (nr,nphi,nz)
            Electric field phi component [V/m].
        ez : array_like (nr,nphi,nz)
            Electric field z component [V/m].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if er.shape   != (nr,nphi,nz):
            raise ValueError("ER has an inconsinstent shape.")
        if ephi.shape != (nr,nphi,nz):
            raise ValueError("Ephi has an inconsinstent shape.")
        if ez.shape   != (nr,nphi,nz):
            raise ValueError("Ez has an inconsinstent shape.")

        parent = "efield"
        group  = "E_3D"
        gname  = ""

        er   = np.transpose(er,(2,1,0))
        ephi = np.transpose(ephi,(2,1,0))
        ez   = np.transpose(ez,(2,1,0))

        # Create a group for this input.
        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("rmin",          (1,),  data=rmin,   dtype="f8")
            g.create_dataset("rmax",          (1,),  data=rmax,   dtype="f8")
            g.create_dataset("nr",            (1,),  data=nr,     dtype="i4")
            g.create_dataset("phimin",        (1,),  data=phimin, dtype="f8")
            g.create_dataset("phimax",        (1,),  data=phimax, dtype="f8")
            g.create_dataset("nphi",          (1,),  data=nphi,   dtype="i4")
            g.create_dataset("zmin",          (1,),  data=zmin,   dtype="f8")
            g.create_dataset("zmax",          (1,),  data=zmax,   dtype="f8")
            g.create_dataset("nz",            (1,),  data=nz,     dtype="i4")
            g.create_dataset("er",  (nz, nphi, nr),  data=er,     dtype="f8")
            g.create_dataset("ephi",(nz, nphi, nr),  data=ephi,   dtype="f8")
            g.create_dataset("ez",  (nz, nphi, nr),  data=ez,     dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return E_3D.write_hdf5(fn=fn, rmin=1, rmax=10, nr=3, zmin=-10, zmax=10,
                               nz=3, phimin=0, phimax=360, nphi=3,
                               er=np.zeros((3,3,3)), ephi=np.zeros((3,3,3)),
                               ez=np.zeros((3,3,3)), desc="DUMMY")

class E_3DS(DataGroup):
    """3D electric field interpolated with cubic splines.

    This input tabulates the electric-field components on an uniform cylindrical
    grid, and then uses spline interpolation to calculate the values during
    the simulation. Slower and more memory-intensive than :class:`E_3D`, but
    potentially more accurate.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        out["er"]   = np.transpose(out["er"],   (2,1,0))
        out["ephi"] = np.transpose(out["ephi"], (2,1,0))
        out["ez"]   = np.transpose(out["ez"],   (2,1,0))
        return out

    @staticmethod
    def write_hdf5(fn, rmin, rmax, nr, zmin, zmax, nz, phimin, phimax, nphi,
                   er, ephi, ez, desc=None):
        """Write input data to the HDF5 file.

        The toroidal angle phi is treated as a periodic coordinate, meaning
        ``A(phi=phimin) == A(phi=phimax)``. However, the phi grid, where input
        arrays are tabulated, is ``linspace(phimin, phimax, nphi+1)[:-1]``
        to avoid storing duplicate data.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        rmin : float
            Minimum value in R grid [m].
        rmax : float
            Maximum value in R grid [m].
        nr : int
            Number of R grid points.
        zmin : float
            Minimum value in z grid [m].
        zmax : float
            Maximum value in z grid [m].
        nz : int
            Number of z grid points.
        phimin : float
            Beginning of the toroidal period [deg].
        phimax : float
            End of the toroidal period [deg].
        nphi : int
            Number of phi grid points.
        er : array_like (nr,nphi,nz)
            Electric field R component [V/m].
        ephi : array_like (nr,nphi,nz)
            Electric field phi component [V/m].
        ez : array_like (nr,nphi,nz)
            Electric field z component [V/m].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if er.shape   != (nr,nphi,nz):
            raise ValueError("ER has an inconsinstent shape.")
        if ephi.shape != (nr,nphi,nz):
            raise ValueError("Ephi has an inconsinstent shape.")
        if ez.shape   != (nr,nphi,nz):
            raise ValueError("Ez has an inconsinstent shape.")

        parent = "efield"
        group  = "E_3DS"
        gname  = ""

        er   = np.transpose(er,(2,1,0))
        ephi = np.transpose(ephi,(2,1,0))
        ez   = np.transpose(ez,(2,1,0))

        # Create a group for this input.
        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("rmin",          (1,),  data=rmin,   dtype="f8")
            g.create_dataset("rmax",          (1,),  data=rmax,   dtype="f8")
            g.create_dataset("nr",            (1,),  data=nr,     dtype="i4")
            g.create_dataset("phimin",        (1,),  data=phimin, dtype="f8")
            g.create_dataset("phimax",        (1,),  data=phimax, dtype="f8")
            g.create_dataset("nphi",          (1,),  data=nphi,   dtype="i4")
            g.create_dataset("zmin",          (1,),  data=zmin,   dtype="f8")
            g.create_dataset("zmax",          (1,),  data=zmax,   dtype="f8")
            g.create_dataset("nz",            (1,),  data=nz,     dtype="i4")
            g.create_dataset("er",  (nz, nphi, nr),  data=er,     dtype="f8")
            g.create_dataset("ephi",(nz, nphi, nr),  data=ephi,   dtype="f8")
            g.create_dataset("ez",  (nz, nphi, nr),  data=ez,     dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return E_3DS.write_hdf5(fn=fn, rmin=1, rmax=10, nr=3, zmin=-10, zmax=10,
                                nz=3, phimin=0, phimax=360, nphi=3,
                                er=np.zeros((3,3,3)), ephi=np.zeros((3,3,3)),
                                ez=np.zeros((3,3,3)), desc="DUMMY")

class E_3DST(DataGroup):
    """Time-dependent 3D electric field interpolated with cubic splines.

    This input tabulates the electric-field components on an uniform
    (R, phi, z, t) grid, and then uses spline interpolation to calculate
    the values during the simulation.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        out["er"]   = np.transpose(out["er"],   (3,2,1,0))
        out["ephi"] = np.transpose(out["ephi"], (3,2,1,0))
        out["ez"]   = np.transpose(out["ez"],   (3,2,1,0))
        return out

    @staticmethod
    def write_hdf5(fn, rmin, rmax, nr, zmin, zmax, nz, phimin, phimax, nphi,
                   tmin, tmax, nt, er, ephi, ez, desc=None):
        """Write input data to the HDF5 file.

        The toroidal angle phi is treated as a periodic coordinate, meaning
        ``A(phi=phimin) == A(phi=phimax)``. However, the phi grid, where input
        arrays are tabulated, is ``linspace(phimin, phimax, nphi+1)[:-1]``
        to avoid storing duplicate data.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        rmin : float
            Minimum value in R grid [m].
        rmax : float
            Maximum value in R grid [m].
        nr : int
            Number of R grid points.
        zmin : float
            Minimum value in z grid [m].
        zmax : float
            Maximum value in z grid [m].
        nz : int
            Number of z grid points.
        phimin : float
            Beginning of the toroidal period [deg].
        phimax : float
            End of the toroidal period [deg].
        nphi : int
            Number of phi grid points.
        tmin : float
            Minimum value in t grid [s].
        tmax : float
            Maximum value in t grid [s].
        nt : int
            Number of t grid points.
        er : array_like (nr,nphi,nz,nt)
            Electric field R component [V/m].
        ephi : array_like (nr,nphi,nz,nt)
            Electric field phi component [V/m].
        ez : array_like (nr,nphi,nz,nt)
            Electric field z component [V/m].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if er.shape   != (nr,nphi,nz,nt):
            raise ValueError("ER has an inconsinstent shape.")
        if ephi.shape != (nr,nphi,nz,nt):
            raise ValueError("Ephi has an inconsinstent shape.")
        if ez.shape   != (nr,nphi,nz,nt):
            raise ValueError("Ez has an inconsinstent shape.")

        parent = "efield"
        group  = "E_3DST"
        gname  = ""

        er   = np.transpose(er,(3,2,1,0))
        ephi = np.transpose(ephi,(3,2,1,0))
        ez   = np.transpose(ez,(3,2,1,0))

        # Create a group for this input.
        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("rmin",              (1,), data=rmin,   dtype="f8")
            g.create_dataset("rmax",              (1,), data=rmax,   dtype="f8")
            g.create_dataset("nr",                (1,), data=nr,     dtype="i4")
            g.create_dataset("phimin",            (1,), data=phimin, dtype="f8")
            g.create_dataset("phimax",            (1,), data=phimax, dtype="f8")
            g.create_dataset("nphi",              (1,), data=nphi,   dtype="i4")
            g.create_dataset("zmin",              (1,), data=zmin,   dtype="f8")
            g.create_dataset("zmax",              (1,), data=zmax,   dtype="f8")
            g.create_dataset("nz",                (1,), data=nz,     dtype="i4")
            g.create_dataset("tmin",              (1,), data=tmin,   dtype="f8")
            g.create_dataset("tmax",              (1,), data=tmax,   dtype="f8")
            g.create_dataset("nt",                (1,), data=nt,     dtype="i4")
            g.create_dataset("er",  (nt, nz, nphi, nr), data=er,     dtype="f8")
            g.create_dataset("ephi",(nt, nz, nphi, nr), data=ephi,   dtype="f8")
            g.create_dataset("ez",  (nt, nz, nphi, nr), data=ez,     dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return E_3DST.write_hdf5(
            fn=fn, rmin=1, rmax=10, nr=3, zmin=-10, zmax=10,
            nz=3, phimin=0, phimax=360, nphi=3, tmin=0, tmax=1,
            nt=3, er=np.zeros((3,3,3,3)),
            ephi=np.zeros((3,3,3,3)), ez=np.zeros((3,3,3,3)),
            desc="DUMMY")

class E_1DS(DataGroup):
    """One-dimensional electric field interpolated with cubic splines.

    This input tabulates the gradient of the electric field potential with
    respect to minor radius on 1D (radial) grid which is then interpolated with
    splines during the simulation.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key == "nrho":
                    out[key] = int(out[key])

        return out

    @staticmethod
    def write_hdf5(fn, nrho, rhomin, rhomax, dvdrho, reff, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        nrho : int
            Number of rho slots in data.
        rhomin : float
            Minimum rho value.
        rhomax : float
            Maximum rho value.
        dvdrho : array_like (nrho,1)
            Derivative of electric potential with respect to minor radius [V/m].

            If ``reff = 1 m``, this is essentially equal to ``dv/dr``.
        reff : float
            Effective minor radius of the plasma used to convert ``dv/drho`` to
            SI units as ``drho/dr=1/reff`` [m].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if dvdrho.shape != (nrho,) and dvdrho.shape != (nrho,1):
            raise ValueError("Input dv/drho has a wrong shape.")

        parent = "efield"
        group  = "E_1DS"
        gname  = ""
        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc)
            gname = g.name.split("/")[-1]

            g.create_dataset('nrho',   (1,1),     data=nrho,   dtype='i8')
            g.create_dataset('rhomin', (1,1),     data=rhomin, dtype='f8')
            g.create_dataset('rhomax', (1,1),     data=rhomax, dtype='f8')
            g.create_dataset('dvdrho', (nrho,1),  data=dvdrho, dtype='f8')
            g.create_dataset('reff',   (1,1),     data=reff,   dtype='f8')

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return E_1DS.write_hdf5(fn=fn, nrho=3, rhomin=0, rhomax=1,
                                dvdrho=np.zeros((3,)), reff=1, desc="DUMMY")

a5py/a5py/ascot5io/efield.py
ascot4fusion/ascot5
import numpy as np

class a5imas:


    def __init__(self):
        self.ids_name = "ids" #This should be overwritten by dervied classes

    def open(self, user, tokamak, version, shot, run, occurrence=0):

        # Put this inside the function, not to disturb usage where imas is not available
        import imas

        # ./idsdump.py g2jvarje test 3 92436 0272 distributions
        # user = "g2jvarje"
        # tokamak = "test"
        # version = "3"
        # shot = 92436
        # run = 272
        # ids = "wall"

        self.ids_coordinates = {'user'       : user,
                                'tokamak'    : tokamak,
                                'version'    : version,
                                'shot'       : shot,
                                'run'        : run,
                                'occurrence' : occurrence,
                                'ids_name'   : self.ids_name }

        self.ids = imas.ids(shot, run)
        self.ids.open_env(user, tokamak, version)

        #ids = ids.split('/')
        #if len(ids) == 1:
        #    occurrence = 0
        #else:
        #    occurrence = int(ids[1])
        #ids = ids[0]

        ids = self.ids_name

        idsdata = self.ids.__dict__[self.ids_name]
        if 'get' not in dir(idsdata):
            idsdata = self.ids.__dict__[self.ids_name + 'Array']
            idsdata.get(occurrence)
            #if idsdata.array:
            #    for slice in idsdata.array:
            #        print(slice)
        else:
            idsdata.get(occurrence)
            #print(idsdata)

        return self.ids

    def create(self, user, tokamak, version, shot, run, occurrence=0):

        # Put this inside the function, not to disturb usage where imas is not available
        import imas


        backend = imas.imasdef.MDSPLUS_BACKEND

        self.occurrence = 0
        self.data_entry = imas.DBEntry(backend_id=backend,
                                  db_name=tokamak, shot=shot, run=run,
                                  user_name=user, data_version=version )
        

        #data_entry.open()
        self.data_entry.create()
                                  
        if self.ids_name == 'wall':
            #ids = data_entry.get(ids_name,occurrence)
            self.ids = imas.wall()
        elif self.ids_name == 'equilibrium':
            self.ids = imas.equilibrium()
        else:
            raise NotImplementedError("The ids {} has not been implemented in ascot5 imas.py create().".format(self.ids_name))
            # Probably easy enough to implement another ids "newids":
            # "ids = imas.newids()"

    def write_data_entry(self):
                                       
        self.data_entry.put(self.ids,self.occurrence)

        self.data_entry.close()

    
    def close(self):
        self.ids.close()


    def description_string(self):
        c = self.ids_coordinates
        s = "DB:{} USER:{} IDS:{} SHOT:{} RUN:{} OCCUR:{}"
        return s.format(c['tokamak'], c['user'], c['ids_name'], c['shot'], c['run'], c['occurrence'])


class B_STS(a5imas):
    ''' Read stellarator 3D magnetic field with the conventions laid out in:
        git@github.com:sjjamsa/imas-ggd-b3d.git

        Returns a dict modelled after write_hdf5() in  B_STS.py 
    '''

    def __init__(self):
        super().__init__()
        self.ids_name = "equilibrium"


    def read(self, user, tokamak, version, shot, run, occurrence=0 ):


        # Put this inside the function, not to disturb usage where imas is not available
        import imas

        itm = self.open( user, tokamak, version, shot, run, occurrence )



        time_slice = 0
        ggd_index  = 0
        grids_ggd_index = 0

        psi0 = 0.0 # We assume so
        psi1 = 1.0 # We assume so


        ggd  = self.ids.equilibrium.time_slice[time_slice].ggd[ggd_index]
        grid = self.ids.equilibrium.grids_ggd[grids_ggd_index].grid[0] # Grid for cylindrical R phi z grids. Assumed (by convention) to be in index 0.


        # Check we have R,phi,z grid:
        if grid.identifier.index != 10:
            raise ValueError("Expecting index 10 in grid identifier, instead of {}.".format(grid.identifier.index))
        if len(grid.space) != 3 :
            raise ValueError("Should be 3-dimensional grid.")
        if grid.space[0].coordinates_type[0] != 4 :
            raise ValueError("Coordinate type mismatch [R]")
        if grid.space[1].coordinates_type[0] != 6 :
            raise ValueError("Coordinate type mismatch [phi]")
        if grid.space[2].coordinates_type[0] != 5 :
            raise ValueError("Coordinate type mismatch [z]")


        nR   = len(grid.space[0].objects_per_dimension[0].object)
        nphi = len(grid.space[1].objects_per_dimension[0].object)
        nz   = len(grid.space[2].objects_per_dimension[0].object)
        if nR   < 1 :
            raise ValueError("R is zero length")
        if nphi < 1 :
            raise ValueError("phi is zero length")
        if nz   < 1 :
            raise ValueError("z is zero length")


        R   = np.zeros(shape=(  nR, ) )
        phi = np.zeros(shape=(nphi, ) )
        z   = np.zeros(shape=(  nz, ) )

        for i in range(nR):
            R[i]   = grid.space[0].objects_per_dimension[0].object[i].geometry[0]
        for i in range(nphi):
            phi[i] = grid.space[1].objects_per_dimension[0].object[i].geometry[0]
        for i in range(nz):
            z[i]   = grid.space[2].objects_per_dimension[0].object[i].geometry[0]

        shape = (nR, nphi, nz)
        order = 'F'

        ldata = nR * nphi * nz
        if  len( ggd.b_field_r[   0 ].values) != ldata :
             raise ValueError("B_R size does not match R,phi,z size") 
        if  len( ggd.b_field_tor[   0 ].values) != ldata :
             raise ValueError("B_tor size does not match R,phi,z size") 
        if  len( ggd.b_field_z[   0 ].values) != ldata :
             raise ValueError("B_z size does not match R,phi,z size") 

        B_R   = np.reshape( ggd.b_field_r[   0 ].values, newshape=shape, order=order )
        B_tor = np.reshape( ggd.b_field_tor[ 0 ].values, newshape=shape, order=order )
        B_z   = np.reshape( ggd.b_field_z[   0 ].values, newshape=shape, order=order )

        if len( ggd.psi ) > 0 :
            if  len( ggd.psi[   0 ].values) != ldata :
                raise ValueError("psi data size does not match R,phi,z size") 
            psi_arr   = np.reshape( ggd.psi[   0 ].values, newshape=shape, order=order )
        else:
            psi_arr   = None

        if len( ggd.phi) > 0 :
            if  len( ggd.phi[   0 ].values) != ldata :
                raise ValueError("phi data size does not match R,phi,z size") 
            phi_arr   = np.reshape( ggd.phi[   0 ].values, newshape=shape, order=order )
        else:
            phi_arr   = None

        if len( ggd.theta) > 0 :
            if  len( ggd.theta[   0 ].values) != ldata :
                raise ValueError("theta data size does not match R,phi,z size") 
            theta_arr = np.reshape( ggd.theta[ 0 ].values, newshape=shape, order=order )
        else:
            theta_arr = None

        # Read the magnetix axis from grid[1] (index is "1" by convention)
        if len (   self.ids.equilibrium.grids_ggd[grids_ggd_index].grid    ) > 1 :
            grid = self.ids.equilibrium.grids_ggd[grids_ggd_index].grid[1]
            axis_R = None
            axis_z = None
            nphi = len(grid.space[0].objects_per_dimension[0].object)
            if nphi < 1 :
                raise ValueError("axis data is zero length")

            axis_R   = -1.0 * np.ones( shape=(nphi,) )
            axis_z   = -1.0 * np.ones( shape=(nphi,) )
            axis_phi = -1.0 * np.ones( shape=(nphi,) )

            if (
                    grid.space[0].coordinates_type[0] != 4 or  # R
                    grid.space[0].coordinates_type[1] != 6 or  # phi
                    grid.space[0].coordinates_type[2] != 5 ):  # z
                raise ValueError("expected coordinates_type=[4 6 5]")



            for i in range(nphi):
                axis_R[i]   = grid.space[0].objects_per_dimension[0].object[i].geometry[0]
                axis_phi[i] = grid.space[0].objects_per_dimension[0].object[i].geometry[1]
                axis_z[i]   = grid.space[0].objects_per_dimension[0].object[i].geometry[2]
        else:
            axis_R = None
            axis_z = None


        # Create the dictionary



        B = {}

        Rmin = R[ 0]
        Rmax = R[-1]
        nR   = len(R)
        Pmin = phi[ 0]
        Pmax = phi[-1]
        nP   = len(phi)
        zmin = z[ 0]
        zmax = z[-1]
        nz   = len(z)

        # For the 3D-arrays, the dimension order is (r,phi,z)
        # Check that initial data is the shape we expect it to be
        if (
                B_R.shape[0] != nR or
                B_R.shape[1] != nP or
                B_R.shape[2] != nz ):
            raise ValueError("unexpected array dimensions")





        B["b_rmin"]     =      np.array([Rmin   ])
        B["b_rmax"]     =      np.array([Rmax   ])
        B["b_nr"]       =      np.array([nR     ])
        B["b_phimin"]   =      np.array([Pmin   ])
        B["b_phimax"]   =      np.array([Pmax   ])
        B["b_nphi"]     =      np.array([nP     ])
        B["b_zmin"]     =      np.array([zmin   ])
        B["b_zmax"]     =      np.array([zmax   ])
        B["b_nz"]       =      np.array([nz     ])
        B["psi_rmin"]   =      np.array([Rmin   ])
        B["psi_rmax"]   =      np.array([Rmax   ])
        B["psi_nr"]     =      np.array([nR     ])
        B["psi_phimin"] =      np.array([Pmin   ])
        B["psi_phimax"] =      np.array([Pmax   ])
        B["psi_nphi"]   =      np.array([nP     ])
        B["psi_zmin"]   =      np.array([zmin   ])
        B["psi_zmax"]   =      np.array([zmax   ])
        B["psi_nz"]     =      np.array([nz     ])
        B["axis_phimin"]=      np.array([Pmin   ])
        B["axis_phimax"]=      np.array([Pmax   ])
        B["axis_nphi"]  =      np.array([nP     ])

        B["axisr"]      =      axis_R 
        B["axisz"]      =      axis_z 

        # For the 3D-arrays, the required dimension order is (r,phi,z)
        B["br"]         =      B_R
        B["bphi"]       =      B_tor
        B["bz"]         =      B_z
        B["psi"]        =      psi_arr

        #print('old ', psi_arr.shape)
        #print('new ', B['psi'].shape)
        #print('len ', (nR,nP,nz), 'R,phi,z')

        B["psi0"]       =      np.array([psi0])  # This is a bold assumption atm
        B["psi1"]       =      np.array([psi1])  # This is a bold assumption atm.


        B["desc"]       = self.description_string()

        return B

    def write(self, B, user, tokamak, version, shot, run, metadata={}):

        # Check that all arrays have same size
        if B['b_nr']   != B['psi_nr']:
            raise ValueError('b_nr != psi_nr')
        if B['b_nz']   != B['psi_nz']:
            raise ValueError('b_nz != psi_nz')
        if B['b_nphi'] != B['psi_nphi']:
            raise ValueError('b_nphi != psi_nphi')
        if B['b_nphi'] != B['axis_nphi']:
            raise ValueError('b_nphi != axis_nphi')

        if (
                B['b_rmin'] != B['psi_rmin'] or
                B['b_rmax'] != B['psi_rmax'] or
                B['b_zmin'] != B['psi_zmin'] or
                B['b_zmax'] != B['psi_zmax'] or
                B['b_phimin'] != B['psi_phimin'] or
                B['b_phimax'] != B['psi_phimax'] or
                B['b_phimax'] != B['psi_phimax'] or
                B['b_phimin'] != B['axis_phimin'] or
                B['b_phimax'] != B['axis_phimax']   ):
            raise ValueError('b/psi/axis r/z/phi min/max do not match Expecting identical grids.')
                
        
        index_grids_ggd = 0
        index_time_slice = 0
        nGrids = 2
        index_ggd = 0

        self.create( user, tokamak, version, shot, run)

        if 'ids_comment' in metadata :
            self.ids.ids_properties.comment = metadata['ids_comment']
        else:
            self.ids.ids_properties.comment = "equilibrium IDS for testing"

        # mandatory
        self.ids.ids_properties.homogeneous_time = 1
        self.ids.time = np.array([0.0])

        
        # Fill data

        if len( self.ids.grids_ggd ) <= index_grids_ggd :
            self.ids.grids_ggd.resize(index_grids_ggd+1)
        if len( self.ids.grids_ggd[index_grids_ggd].grid ) < nGrids :
            self.ids.grids_ggd[index_grids_ggd].grid.resize(nGrids)



        # The structured grid
        grid = self.ids.grids_ggd[index_grids_ggd].grid[0]

        if "description_ggd_identifier_name" in metadata :
            grid.identifier.name = metadata[ "description_ggd_identifier_name" ]
        else :
            grid.identifier.name = "structured_spaces"
        if "description_ggd_identifier_desc" in metadata :
            grid.identifier.description = metadata[ "description_ggd_identifier_desc" ]
        else :
            grid.identifier.description = "Cylindrical RphiZ grid"
        grid.identifier.index = 10


        nSpaces = 3 # nodes edges "faces"
        if len(grid.space) < nSpaces:
            grid.space.resize(nSpaces)

        grid.space[0].coordinates_type = np.array( [4] ) # R
        grid.space[1].coordinates_type = np.array( [6] ) # phi
        grid.space[2].coordinates_type = np.array( [5] ) # phi

        objdims=1 # GGD would create additional dimensions, but we don't really need them, do we?
        for iSpace in range(3):
            if len(grid.space[iSpace].objects_per_dimension) < objdims :
                grid.space[   iSpace].objects_per_dimension.resize(objdims)

        nR   = B['b_nr'][0]
        if len(grid.space[0].objects_per_dimension[0].object) != nR :
            grid.space[   0].objects_per_dimension[0].object.resize(nR)
        nPhi = B['b_nphi'][0]
        if len(grid.space[0].objects_per_dimension[0].object) != nPhi :
            grid.space[   1].objects_per_dimension[0].object.resize(nPhi)
        nz   = B['b_nz'][0]
        if len(grid.space[2].objects_per_dimension[0].object) != nz :
            grid.space[   2].objects_per_dimension[0].object.resize(nz)

        R   = np.linspace( B['b_rmin'  ][0], B['b_rmax'  ][0], B['b_nr'  ][0] )
        for i in range(nR):
            grid.space[0].objects_per_dimension[0].object[i].geometry = np.array( [R[i]  ] )
            grid.space[0].objects_per_dimension[0].object[i].nodes = np.array([i+1])
        phi = np.linspace( B['b_phimin'][0], B['b_phimax'][0], B['b_nphi'][0] )
        for i in range(nPhi):
            grid.space[1].objects_per_dimension[0].object[i].geometry = np.array( [phi[i]] )
            grid.space[1].objects_per_dimension[0].object[i].nodes = np.array([i+1])
        z   = np.linspace( B['b_zmin'  ][0], B['b_zmax'  ][0], B['b_nz'  ][0] )
        for i in range(nz):
            grid.space[2].objects_per_dimension[0].object[i].geometry = np.array( [z[i]  ] )
            grid.space[2].objects_per_dimension[0].object[i].nodes = np.array([i+1])
            
        # Todo... edges, sub-spaces

            
        # The magnetic axis grid
        grid = self.ids.grids_ggd[index_grids_ggd].grid[1]

        grid.identifier.index = 1
        grid.identifier.description = 'magnetic axis'

        
        nSpaces = 1 
        if len(grid.space) < nSpaces:
            grid.space.resize(nSpaces)
       
        grid.space[0].identifier.index = 1
        grid.space[0].coordinates_type = np.array([4, 6, 5])

        # nodes & edges
        objdims = 2

        # nodes
        if len(grid.space[0].objects_per_dimension) < objdims :
            grid.space[0].objects_per_dimension.resize(2)
            
        if len(grid.space[0].objects_per_dimension[0].object) != nPhi:
            grid.space[0].objects_per_dimension[0].object.resize(nPhi)

        for i in range(nPhi):
            grid.space[0].objects_per_dimension[0].object[i].geometry = np.array( [  B["axisr"][i], phi[i], B["axisz"][i] ] )
            grid.space[0].objects_per_dimension[0].object[i].nodes = np.array([i+1])

        # edges

        if len(grid.space[0].objects_per_dimension[1].object) != 1:
            grid.space[0].objects_per_dimension[1].object.resize(1)

        grid.space[0].objects_per_dimension[1].object[0].nodes = np.arange( 1, nPhi+1, dtype=int )
            

        if len(grid.grid_subset) < 1:
            grid.grid_subset.resize(1)

        grid.grid_subset[0].identifier.index = 100
        grid.grid_subset[0].dimension = 2

        if len(grid.grid_subset[0].element) < 1:
            grid.grid_subset[0].element.resize(1)
        if len(grid.grid_subset[0].element[0].object) < 1:
            grid.grid_subset[0].element[0].object.resize(1)
    
        grid.grid_subset[0].element[0].object[0].space     = 1
        grid.grid_subset[0].element[0].object[0].dimension = 2
        grid.grid_subset[0].element[0].object[0].index     = 1


        # The fields
        if len( self.ids.time_slice) <= index_time_slice :        
            self.ids.time_slice.resize(index_time_slice+1)
        if len(self.ids.time_slice[index_time_slice].ggd) <= index_ggd :
            self.ids.time_slice[index_time_slice].ggd.resize(index_ggd+1)
        ggd = self.ids.time_slice[index_time_slice].ggd[index_ggd]

        if len(ggd.psi) < 1:
            ggd.psi.resize(1)
        ggd.psi[0].values = B['psi'].flatten(order='F')

        if len(ggd.b_field_r) < 1:
            ggd.b_field_r.resize(1)
        ggd.b_field_r[0].values = B['br'].flatten(order='F')

        if len(ggd.b_field_z) < 1:
            ggd.b_field_z.resize(1)
        ggd.b_field_z[0].values = B['bz'].flatten(order='F')

        if len(ggd.b_field_tor) < 1:
            ggd.b_field_tor.resize(1)
        ggd.b_field_tor[0].values = B['bphi'].flatten(order='F')

        
        self.write_data_entry()
                                 

        
    
class wall_2d(a5imas):

    def __init__(self):
        super().__init__()
        self.ids_name = "wall"


    def read(self, user, tokamak, version, shot, run, occurrence=0 ):
        """
        Read an IMAS 2D-wall into a dictionary, that is a drop-in replacement for a dict read from hdf5.

        The data is read from
              wall.description_2d[0].limiter.unit[0].outline.r
        and
              wall.description_2d[0].limiter.unit[0].outline.z

        """

        # Put this inside the function, not to disturb usage where imas is not available
        import imas


        timeIndex = 0
        unit      = 0

        itm = self.open( user, tokamak, version, shot, run, occurrence )

        #itm.wall.get()
        r = self.ids.wall.description_2d[timeIndex].limiter.unit[unit].outline.r
        z = self.ids.wall.description_2d[timeIndex].limiter.unit[unit].outline.z

        self.close()

        w = {
            "r"         : r,
            "z"         : z,
            "nelements" : len(r),
        }

        return w



class wall_3d(a5imas):

    def __init__(self):
        super().__init__()
        self.ids_name = "wall"


    def read(self, user, tokamak, version, shot, run, occurrence=0 ):
        """
        Read an IMAS 3D-wall into a dictionary, that is a drop-in replacement for a dict read from hdf5.

        The data is read from
              wall.description_ggd[0].grid_ggd[0].grid_ggd.space[0]
        Flag is set to the index of description_ggd == 0

        """

        itm = self.open( user, tokamak, version, shot, run, occurrence )

        description_ggd_index = 0
        grid_ggd_index        = 0
        space_index           = 0
        
        
        desc_ggd = self.ids.wall.description_ggd[description_ggd_index]
        grid_ggd = desc_ggd.grid_ggd[grid_ggd_index]
        space    = grid_ggd.space[space_index]
        
        if not np.all( space.coordinates_type == np.array([1, 2, 3]) ):
            raise ValueError("Space coordinates not [1,2,3] in IMAS (3d)wall ids u:{} db:{} v:{} s:{} r:{} o:{} desc:{} grid:{} sp:{}".format(
                user,tokamak,version,shot,run,occurrence,description_ggd_index,grid_ggd_index,space_index ))


        node_coordinates,edge_indexes,tria_indexes = self.get_nodes_edges_polygons(space)

        self.close()


        # Convert 3D wall data to ASCOT5 format
                             
        xyz = self.get_xyz(node_coordinates,tria_indexes)
        

        nelements = xyz.shape[0]
        flag      = np.ones( (nelements,1) ) * description_ggd_index
        w = {
            "x1x2x3"         : xyz[:,:,0],
            "y1y2y3"         : xyz[:,:,1],
            "z1z2z3"         : xyz[:,:,2],
            "nelements"      : [[nelements]],
            "flag"           : flag,
            "flagIdList"     : flag,
            "flagIdStrings"  : ["IMAS 3D-wall ids u:{} db:{} v:{} s:{} r:{} o:{} desc:{} grid:{} sp:{}".format(
                user,tokamak,version,shot,run,occurrence,description_ggd_index,grid_ggd_index,space_index )],

        }

        return w


    def get_nodes_edges_polygons(self,space):
        '''
        Copy the node coordinates, edge indexes and polygons from the space to separae numpy arrays:

        returns node_coordinates[3,n_nodes], edge_indexes[2,n_edges],tria_indexes[3,n_tris]

        '''


        n_nodes=len(space.objects_per_dimension[0].object)
        n_edges=len(space.objects_per_dimension[1].object)
        n_tris =len(space.objects_per_dimension[2].object)
        node_coordinates = np.zeros( (3,n_nodes),dtype=float )
        edge_indexes     = np.zeros( (2,n_edges),dtype=int   )
        tria_indexes     = np.zeros( (3,n_tris), dtype=int   )

        for i in range(n_nodes):
            node_coordinates[:,i] = space.objects_per_dimension[0].object[i].geometry

        for i in range(n_edges):
            edge_indexes[:,i] = space.objects_per_dimension[1].object[i].nodes

        for i in range(n_tris):
            tria_indexes[:,i] = space.objects_per_dimension[2].object[i].nodes

        return node_coordinates,edge_indexes,tria_indexes

    def get_xyz(self,node_coordinates,tria_indexes):
        '''
        Parse out triangle corner coordinates. Assume one based indexing in the input (first = 1).

        Returns  xyz[ntri,3,3], where the last dimension [ntri,ncorn,:]  is x,y,z. 


        '''
        ntri =  tria_indexes.shape[1]

        xyz = np.zeros( (ntri, 3, 3 ),dtype=float )

        for itri in range( ntri ):
            for i in range(3):
                xyz[itri,i,:] = node_coordinates[:,tria_indexes[i,itri]-1]

        return xyz

    def fill_wall_3d_ids(self,xyz,metadata={}):
        
        # Let's first do the conversion
        nodes,edges,faces = self.xyz_to_nodes_edges_faces(xyz)

        ids = self.ids
        
        if 'ids_comment' in metadata :
            ids.ids_properties.comment = metadata['ids_comment']
        else:
            ids.ids_properties.comment = "wall IDS saved by ASCOT5 imas.py"

        # mandatory
        ids.ids_properties.homogeneous_time = 1
        ids.time = np.array([0.0])


        index_description_ggd = 0
        index_grid_ggd        = 0
        index_space           = 0

        if len( ids.description_ggd ) <= index_description_ggd :
            ids.description_ggd.resize( index_description_ggd + 1 )
        if len( ids.description_ggd[index_description_ggd].grid_ggd ) <= index_grid_ggd :
            ids.description_ggd[index_description_ggd].grid_ggd.resize( index_grid_ggd + 1 )
        if len( ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].space ) <= index_space :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].space.resize(index_space + 1)

        if "description_ggd_identifier_name" in metadata :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.name = metadata[ "description_ggd_identifier_name" ]
        else :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.name = "ascot5wall"
        if "description_ggd_identifier_desc" in metadata :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.description = metadata[ "description_ggd_identifier_desc" ]
        else :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.description = "Wall stored by ASCOT5 imas.py"
        ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.index = 1


        space = ids.\
            description_ggd[index_description_ggd].\
            grid_ggd[index_grid_ggd].\
            space[index_space]

        # x,y,z space
        space.coordinates_type = np.array([1,2,3])

        space.objects_per_dimension.resize(3)

        space.objects_per_dimension[0].object.resize(len(nodes))
        for i,node in enumerate(nodes):
            space.objects_per_dimension[0].object[i].geometry = node
            space.objects_per_dimension[0].object[i].nodes = np.ones( (1,),dtype=int )*(i+1)

        space.objects_per_dimension[1].object.resize(len(edges))
        for i,edge in enumerate(edges):
            space.objects_per_dimension[1].object[i].nodes = edge

        space.objects_per_dimension[2].object.resize(len(faces))
        for i,face in enumerate(faces):
            space.objects_per_dimension[2].object[i].nodes = face

    def xyz_to_nodes_edges_faces(self,xyz):
        '''
        Input must be xyz[ntri,3,3], where the last dimension [ntri,ncorn,:]  is x,y,z. 
        '''

        #remove duplicate nodes?
        remove_duplicate_nodes = True

        remove_duplicate_edges = True
        if remove_duplicate_edges:
            remove_duplicate_nodes = True

        # number of triangles
        n = xyz.shape[0]

        # Every node is independent. There may be duplicates.
        nodes = np.copy(xyz).reshape( (3*n,3) )

        # We have one index per corner = 3*n
        faces = np.arange(0,3*n,dtype=int).reshape(n,3)

        # a row of edges
        a = np.array([ [0,1],[1,2],[2,0] ]).reshape( (1,3,2) )
        # a column of triangle indexes
        b = np.arange(0,3*n,3,dtype=int).reshape((n,1,1))
        edges = a+b
        edges = edges.reshape(3*n,2)

        if remove_duplicate_nodes :
            nodes,reverse =  np.unique(nodes, return_inverse=True,return_index=False,
                                            axis=0)

            edges = reverse[edges]
            faces = reverse[faces]

        if remove_duplicate_edges :
            edges = np.unique(edges,axis=0)

        # remember to edges++ and faces++ to get one-based indexing 
        return nodes+1,edges+1,faces+1

    def write(self, w, user, tokamak, version, shot, run, metadata={}):

        self.create( user, tokamak, version, shot, run)

        n = w['nelements'][0][0]
        xyz=np.zeros( (n, 3, 3) )

        xyz[:,:,0] = w['x1x2x3'][:,:]
        xyz[:,:,1] = w['y1y2y3'][:,:]
        xyz[:,:,2] = w['z1z2z3'][:,:]
                                       
        # Save the data
        self.fill_wall_3d_ids(self.ids,xyz,metadata=metadata)

                                       
        self.write_data_entry()
                                 

    def fill_wall_3d_ids(self,ids,xyz,metadata={}):

        # Let's first do the conversion
        nodes,edges,faces = self.xyz_to_nodes_edges_faces(xyz)

        if 'ids_comment' in metadata :
            ids.ids_properties.comment = metadata['ids_comment']
        else:
            ids.ids_properties.comment = "wall IDS for testing"

        # mandatory
        ids.ids_properties.homogeneous_time = 1
        ids.time = np.array([0.0])


        index_description_ggd = 0
        index_grid_ggd        = 0
        index_space           = 0

        if len( ids.description_ggd ) <= index_description_ggd :
            ids.description_ggd.resize( index_description_ggd + 1 )
        if len( ids.description_ggd[index_description_ggd].grid_ggd ) <= index_grid_ggd :
            ids.description_ggd[index_description_ggd].grid_ggd.resize( index_grid_ggd + 1 )
        if len( ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].space ) <= index_space :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].space.resize(index_space + 1)

        if "description_ggd_identifier_name" in metadata :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.name = metadata[ "description_ggd_identifier_name" ]
        else :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.name = "MyMesh"
        if "description_ggd_identifier_desc" in metadata :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.description = metadata[ "description_ggd_identifier_desc" ]
        else :
            ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.description = "SMITER mesh MyMesh"
        ids.description_ggd[index_description_ggd].grid_ggd[index_grid_ggd].identifier.index = 1


        space = ids.\
            description_ggd[index_description_ggd].\
            grid_ggd[index_grid_ggd].\
            space[index_space]

        # x,y,z space
        space.coordinates_type = np.array([1,2,3])

        space.objects_per_dimension.resize(3)

        space.objects_per_dimension[0].object.resize(len(nodes))
        for i,node in enumerate(nodes):
            space.objects_per_dimension[0].object[i].geometry = node
            space.objects_per_dimension[0].object[i].nodes = np.ones( (1,),dtype=int )*(i+1)

        space.objects_per_dimension[1].object.resize(len(edges))
        for i,edge in enumerate(edges):
            space.objects_per_dimension[1].object[i].nodes = edge

        space.objects_per_dimension[2].object.resize(len(faces))
        for i,face in enumerate(faces):
            space.objects_per_dimension[2].object[i].nodes = face

    def xyz_to_nodes_edges_faces(self,xyz):
        '''
        Input must be xyz[ntri,3,3], where the last dimension [ntri,ncorn,:]  is x,y,z. 
        '''

        #remove duplicate nodes?
        remove_duplicate_nodes = True

        remove_duplicate_edges = True
        if remove_duplicate_edges:
            remove_duplicate_nodes = True

        # number of trianglesgg59
        n = xyz.shape[0]

        # Every node is independent. There may be duplicates.
        nodes = np.copy(xyz).reshape( (3*n,3) )

        # We have one index per corner = 3*n
        faces = np.arange(0,3*n,dtype=int).reshape(n,3)

        # a row of edges
        a = np.array([ [0,1],[1,2],[2,0] ]).reshape( (1,3,2) )
        # a column of triangle indexes
        b = np.arange(0,3*n,3,dtype=int).reshape((n,1,1))
        edges = a+b
        edges = edges.reshape(3*n,2)

        if remove_duplicate_nodes :
            nodes,reverse =  np.unique(nodes, return_inverse=True,return_index=False,
                                            axis=0)

            edges = reverse[edges]
            faces = reverse[faces]

        if remove_duplicate_edges :
            edges = np.unique(edges,axis=0)

        # remember to edges++ and faces++ to get one-based indexing 
        return nodes+1,edges+1,faces+1



a5py/a5py/ascot5io/imas.py
ascot4fusion/ascot5
"""Markers whose orbits are solved within the ASCOT5 code.
"""
import copy
import numpy as np
import h5py
import unyt

from .coreio.treedata import DataGroup
from .coreio.fileapi import read_data, add_group, write_data
from a5py.marker.plot import plot_histogram

from a5py.routines.plotting import openfigureifnoaxes
from a5py.physlib.species import species as getspecies

class Marker(DataGroup):
    """A class acting as a superclass for all marker types.
    """

    def prune(self, pick, normalize=True, mrk=None):
        """Take a subset of markers.

        Parameters
        ----------
        pick : int or array_like
            Either number of markers to be picked (in which case they are chosen
            randomly) or the marker IDs.
        normalize : bool, optional
            Reweight the output markers so that the total weight is
            the same as in input.
        mrk : dict, optional
            Marker data dictionary (same format as in corresponding
            ``write_hdf5`` function).

            If None, the data is read from the file.

        Returns
        -------
        out : dict
            Marker data dictionary (same format as in corresponding
            ``write_hdf5`` function) with picked markers.
        """

        if mrk is None:
            mrk = self.read()
        else:
            mrk = copy.deepcopy(mrk)

        weighttotal = np.sum(mrk["weight"])

        if isinstance(pick, int):
            rng  = np.random.default_rng()
            pick = rng.choice(mrk["ids"].ravel(), size=(pick,), replace=False)

        idx = np.where(np.in1d(mrk["ids"], pick))
        for k in mrk:
            if k == "n": continue
            mrk[k] = mrk[k][idx]

        if normalize:
            mrk["weight"][:] *= weighttotal / np.sum(mrk["weight"])

        mrk["n"] = pick.size
        return mrk

    @openfigureifnoaxes(projection=None)
    def plot_hist_rhophi(self, ascotpy, rbins=10, pbins=10, weighted=False,
                         axes=None):
        """
        Plot marker rho-phi histogram
        """
        weights = None
        if weighted:
            weights = self.read()["weight"]

        rho = self.eval_rho(ascotpy)
        phi = self.eval_phi(ascotpy)
        plot_histogram(rho, xbins=rbins, y=phi, ybins=pbins,
                       weights=weights,
                       logscale=False, xlabel="Normalized poloidal flux",
                       ylabel="Toroidal angle [deg]",
                       axes=axes)

    @openfigureifnoaxes(projection=None)
    def plot_hist_energypitch(self, ascotpy, ebins=10, pbins=10, weighted=False,
                              axes=None):
        """
        Plot marker energy-pitch histogram

        Energy is on logscale.
        """
        weights = None
        if weighted:
            weights = self.read()["weight"]

        pitch  = self.eval_pitch(ascotpy)
        energy = self.eval_energy(ascotpy)
        if energy is None:
            # Field lines don't have energy
            plot_histogram(x=pitch, xbins=pbins, weights=weights,
                           logscale=False, xlabel=r"Pitch [$p_\parallel/p$]",
                           axes=axes)
        else:
            energy = np.log10(energy.to("eV"))
            plot_histogram(energy, xbins=ebins, y=pitch, ybins=pbins,
                           weights=weights,
                           logscale=False, xlabel="Energy [eV]",
                           ylabel=r"Pitch [$p_\parallel/p$]",
                           axes=axes)

    @staticmethod
    def generate(mrktype, n, species=None):
        """Generate dummy marker input of given type and species.

        Parameters
        ----------
        mrktype : {"prt", "gc", "fl"}
            Type of marker input to be created.
        n : int
            Number of markers.
        species : str, optional
            Marker species to initialize anum, znum, mass, and charge.

        Returns
        -------
        mrk : dict
            Markers parameters that can be supplied to :meth:`Prt.write_hdf5`,
            :meth:`GC.write_hdf5` or :meth:`FL.write_hdf5` depending on
            ``mrktype`` value.
        """
        mrk = {
            "n"      : n,
            "ids"    : ( 1 + np.arange(n) ),
            "r"      : np.zeros((n,)),
            "z"      : np.zeros((n,)),
            "phi"    : np.zeros((n,)),
            "weight" : np.ones((n,)),
            "time"   : np.zeros((n,)),
        }
        if species is None:
            species = {"mass":1*unyt.atomic_mass_unit, "charge":1*unyt.e,
                       "anum":1, "znum":1}
        else:
            species = getspecies(species)

        if mrktype == "prt":
            mrk["vr"]     = np.zeros((n,))
            mrk["vz"]     = np.zeros((n,))
            mrk["vphi"]   = np.zeros((n,))
            mrk["mass"]   = ( species["mass"] * np.ones((n,)) ).to_value("amu")
            mrk["charge"] = species["charge"] * np.ones((n,), dtype=np.int16)
            mrk["anum"]   = species["anum"]   * np.ones((n,), dtype=np.int16)
            mrk["znum"]   = species["znum"]   * np.ones((n,), dtype=np.int16)
        if mrktype == "gc":
            mrk["pitch"]  = np.zeros((n,))
            mrk["energy"] = np.zeros((n,))
            mrk["zeta"]   = np.zeros((n,))
            mrk["mass"]   = ( species["mass"] * np.ones((n,)) ).to_value("amu")
            mrk["charge"] = species["charge"] * np.ones((n,), dtype=np.int16)
            mrk["anum"]   = species["anum"]   * np.ones((n,), dtype=np.int16)
            mrk["znum"]   = species["znum"]   * np.ones((n,), dtype=np.int16)
        if mrktype == "fl":
            mrk["pitch"]  = np.zeros((n,))

        return mrk

class FL(Marker):
    """Magnetic-field-line tracer input.

    These are markers that trace magnetic field lines exactly. They don't
    represent physical particles except that for calculating quantities like
    ``mileage``, they are assumed to be photons travelling at the speed of
    light.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key == "n":
                    out[key] = int(out[key])

        out["ids"] = out["id"]
        del out["id"]
        return out

    def eval_pitch(self, ascotpy):
        """
        Evaluate pitch.
        """
        with self as h5:
            pitch = read_data(h5, "pitch")

        return pitch

    @staticmethod
    def write_hdf5(fn, n, ids, r, phi, z, pitch, weight, time, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        n : int
            Number of markers.
        ids : array_like (n,)
            Unique identifier for each marker (must be a positive integer).
        r : array_like (n,)
            Magnetic field line R coordinate [m].
        phi : array_like (n,)
            Magnetic field line phi coordinate [deg].
        z : array_like (n,)
            Magnetic field line z coordinate [m].
        pitch : array_like (n,)
            Sign which defines the direction field line is traced, + is parallel
        weight : array_like (n,)
            Magnetic field line weight [markers/s].
        time : array_like (n,)
            Magnetic field line initial time [s].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if ids.size    != n: raise ValueError("Inconsistent size for ids.")
        if r.size      != n: raise ValueError("Inconsistent size for r.")
        if phi.size    != n: raise ValueError("Inconsistent size for phi.")
        if z.size      != n: raise ValueError("Inconsistent size for z.")
        if pitch.size  != n: raise ValueError("Inconsistent size for pitch.")
        if weight.size != n: raise ValueError("Inconsistent size for weight.")
        if time.size   != n: raise ValueError("Inconsistent size for time.")

        parent = "marker"
        group  = "fl"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            write_data(g, "n",      n,      (1,), "i8")
            write_data(g, "r",      r,      (n,), "f8", "m")
            write_data(g, "phi",    phi,    (n,), "f8", "deg")
            write_data(g, "z",      z,      (n,), "f8", "m")
            write_data(g, "pitch",  pitch,  (n,), "f8", "1")
            write_data(g, "weight", weight, (n,), "f8", "particles/s")
            write_data(g, "time",   time,   (n,), "f8", "s")
            write_data(g, "id",     ids,    (n,), "i8", "1")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        mrk = Marker.generate(n=1, mrktype="fl", species=None)
        return FL.write_hdf5(fn=fn, desc="DUMMY", **mrk)

class GC(Marker):
    """Particle input in guiding-center coordinates.

    The guiding center phase-space is (R, Phi, Z, mu, ppar, zeta), where
    R, Phi, Z are the guiding center position in cylindrical coordinates,
    mu is the magnetic moment (first adiabatic invariant), and ppar is
    the parallel momentum along the magnetic field line). The gyro-angle
    zeta is an ignorable quantity in guiding-center integration, but it is
    used in guiding-center-to-particle transformation.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key == "n":
                    out[key] = int(out[key])

        out["ids"] = out["id"]
        del out["id"]
        return out

    def eval_energy(self, ascotpy):
        with self as h5:
            energy = read_data(h5, "energy")
        return energy


    def eval_pitch(self, ascotpy):
        with self as h5:
            pitch = read_data(h5, "pitch")
        return pitch

    @staticmethod
    def write_hdf5(fn, n, ids, mass, charge, r, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        n : int
            Number of markers.
        ids : array_like (n,)
            Unique identifier for each marker (must be a positive integer).
        charge : array_like (n,)
            Charge [e].
        mass : array_like (n,)
            Mass [amu].
        r : array_like (n,)
            Guiding center R coordinate [m].
        phi : array_like (n,)
            Guiding center phi coordinate [deg].
        z : array_like (n,)
            Guiding center z coordinate [m].
        energy : array_like (n,)
            Guiding center energy [eV].
        pitch : array_like (n,)
            Guiding center pitch (v_para/v_tot).
        zeta : array_like (n,)
            Guiding center gyroangle [rad].
        anum : array_like (n,)
            Marker species atomic mass number.
        znum : array_like (n,)
            Marker species charge number.
        weight : array_like (n,)
            Guiding center weight [markers/s].
        time : array_like (n,)
            Guiding center initial time [s].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if ids.size    != n: raise ValueError("Inconsistent size for ids.")
        if mass.size   != n: raise ValueError("Inconsistent size for mass.")
        if charge.size != n: raise ValueError("Inconsistent size for charge.")
        if r.size      != n: raise ValueError("Inconsistent size for r.")
        if phi.size    != n: raise ValueError("Inconsistent size for phi.")
        if z.size      != n: raise ValueError("Inconsistent size for z.")
        if energy.size != n: raise ValueError("Inconsistent size for energy.")
        if pitch.size  != n: raise ValueError("Inconsistent size for pitch.")
        if zeta.size   != n: raise ValueError("Inconsistent size for zeta.")
        if anum.size   != n: raise ValueError("Inconsistent size for anum.")
        if znum.size   != n: raise ValueError("Inconsistent size for znum.")
        if weight.size != n: raise ValueError("Inconsistent size for weight.")
        if time.size   != n: raise ValueError("Inconsistent size for time.")

        parent = "marker"
        group  = "gc"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("n",      (1,), data=n,      dtype='i8').attrs['unit'] = '1';
            g.create_dataset("r",      (n,), data=r,      dtype='f8').attrs['unit'] = 'm';
            g.create_dataset("phi",    (n,), data=phi,    dtype='f8').attrs['unit'] = 'deg';
            g.create_dataset("z",      (n,), data=z,      dtype='f8').attrs['unit'] = 'm';
            g.create_dataset("energy", (n,), data=energy, dtype='f8').attrs['unit'] = 'ev';
            g.create_dataset("pitch",  (n,), data=pitch,  dtype='f8').attrs['unit'] = '1';
            g.create_dataset("zeta",   (n,), data=zeta,   dtype='f8').attrs['unit'] = 'rad';
            g.create_dataset("mass",   (n,), data=mass,   dtype='f8').attrs['unit'] = 'amu';
            g.create_dataset("charge", (n,), data=charge, dtype='i4').attrs['unit'] = 'e';
            g.create_dataset("anum",   (n,), data=anum,   dtype='i4').attrs['unit'] = '1';
            g.create_dataset("znum",   (n,), data=znum,   dtype='i4').attrs['unit'] = '1';
            g.create_dataset("weight", (n,), data=weight, dtype='f8').attrs['unit'] = 'markers/s';
            g.create_dataset("time",   (n,), data=time,   dtype='f8').attrs['unit'] = 's';
            g.create_dataset("id",     (n,), data=ids,    dtype='i8').attrs['unit'] = '1';

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        mrk = Marker.generate(n=1, mrktype="gc", species="alpha")
        return GC.write_hdf5(fn=fn, desc="DUMMY", **mrk)

class Prt(Marker):
    """Marker input representing physical particles.

    Particle phase-space is (r, phi, z, pr, pphi, pz), where r, phi, and z are
    particle position in cylindrical coordinates, and pr, phi, and pz are its
    momentum components.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key == "n":
                    out[key] = int(out[key])

        out["ids"] = out["id"]
        del out["id"]
        return out

    def eval_energy(self, ascotpy):
        with self as h5:
            vr   = read_data(h5, "vr")
            vz   = read_data(h5, "vz")
            vphi = read_data(h5, "vphi")
            mass = read_data(h5, "mass")

        v = np.sqrt(vr**2 + vz**2 + vphi**2)
        return energy_velocity(mass, v)

    def eval_pitch(self, ascotpy):
        with self as h5:
            r    = read_data(h5, "r")
            z    = read_data(h5, "z")
            phi  = read_data(h5, "phi")
            vr   = read_data(h5, "vr")
            vz   = read_data(h5, "vz")
            vphi = read_data(h5, "vphi")
            mass = read_data(h5, "mass")

        br    = ascotpy.evaluate(r, phi, z, 0, "br")
        bz    = ascotpy.evaluate(r, phi, z, 0, "bz")
        bphi  = ascotpy.evaluate(r, phi, z, 0, "bphi")
        b     = np.sqrt(br**2 + bz**2 + bphi**2)
        v     = np.sqrt(vr**2 + vz**2 + vphi**2)
        pitch = ( vr * br + vz * bz + vphi * bphi ) / ( b * v )
        return pitch


    @staticmethod
    def write_hdf5(fn, n, ids, mass, charge, r, phi, z, vr, vphi, vz,
                   anum, znum, weight, time, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        n : int
            Number of markers.
        ids : array_like (n,)
            Unique identifier for each marker (must be a positive integer).
        mass : array_like (n,)
            Mass [amu].
        charge : array_like (n,)
            Charge [e].
        r : array_like (n,)
            Particle R coordinate [m].
        phi : array_like (n,)
            Particle phi coordinate [deg].
        z : array_like (n,)
            Particle z coordinate [m].
        vr : array_like (n,)
            Particle velocity R-component [m/s].
        vphi : array_like (n,)
            Particle velocity phi-component [m/s].
        vz : array_like (n,)
            Particle velocity z-component [m/s].
        anum : array_like (n,)
            Marker species atomic mass number.
        znum : array_like (n,)
            Marker species charge number.
        weight : array_like (n,)
            Particle weight [markers/s].
        time : array_like (n,)
            Particle initial time [s].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if ids.size    != n: raise ValueError("Inconsistent size for ids.")
        if mass.size   != n: raise ValueError("Inconsistent size for mass.")
        if charge.size != n: raise ValueError("Inconsistent size for charge.")
        if r.size      != n: raise ValueError("Inconsistent size for r.")
        if phi.size    != n: raise ValueError("Inconsistent size for phi.")
        if z.size      != n: raise ValueError("Inconsistent size for z.")
        if vr.size     != n: raise ValueError("Inconsistent size for vR.")
        if vphi.size   != n: raise ValueError("Inconsistent size for vphi.")
        if vz.size     != n: raise ValueError("Inconsistent size for vz.")
        if anum.size   != n: raise ValueError("Inconsistent size for anum.")
        if znum.size   != n: raise ValueError("Inconsistent size for znum.")
        if weight.size != n: raise ValueError("Inconsistent size for weight.")
        if time.size   != n: raise ValueError("Inconsistent size for time.")

        parent = "marker"
        group  = "prt"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("n",      (1,), data=n,      dtype='i8').attrs['unit'] = '1';
            g.create_dataset("r",      (n,), data=r,      dtype='f8').attrs['unit'] = 'm';
            g.create_dataset("phi",    (n,), data=phi,    dtype='f8').attrs['unit'] = 'deg';
            g.create_dataset("z",      (n,), data=z,      dtype='f8').attrs['unit'] = 'm';
            g.create_dataset("vr",     (n,), data=vr,     dtype='f8').attrs['unit'] = 'm/s';
            g.create_dataset("vphi",   (n,), data=vphi,   dtype='f8').attrs['unit'] = 'm/s';
            g.create_dataset("vz",     (n,), data=vz,     dtype='f8').attrs['unit'] = 'm/s';
            g.create_dataset("mass",   (n,), data=mass,   dtype='f8').attrs['unit'] = 'amu';
            g.create_dataset("charge", (n,), data=charge, dtype='i4').attrs['unit'] = 'e';
            g.create_dataset("anum",   (n,), data=anum,   dtype='i4').attrs['unit'] = '1';
            g.create_dataset("znum",   (n,), data=znum,   dtype='i4').attrs['unit'] = '1';
            g.create_dataset("weight", (n,), data=weight, dtype='f8').attrs['unit'] = 'markers/s';
            g.create_dataset("time",   (n,), data=time,   dtype='f8').attrs['unit'] = 's';
            g.create_dataset("id",     (n,), data=ids,    dtype='i8').attrs['unit'] = '1';

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This dummy input sets electric field to zero everywhere.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        mrk = Marker.generate(n=1, mrktype="prt", species="alpha")
        return Prt.write_hdf5(fn=fn, desc="DUMMY", **mrk)

a5py/a5py/ascot5io/marker.py
ascot4fusion/ascot5
"""Input containing MHD eigenmodes.

This input is used in simulations studying particle response to MHD.

Using this data in a simulation requires that boozer data, containing
mapping between real-space and Boozer coordinates, is present.

This module reads and writes both MHD_STAT and MHD_NONSTAT data. The only
difference between these two is the time-dependency of the eigenfunctions,
that is, does the mode amplitude (alpha or Phi) depend only on psi or psi
and time. MHD_STAT assumes only psi dependency, making the interpolation
*much* faster. So don't use MHD_NONSTAT unless you really have to.
"""
import h5py
import numpy as np

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class MHD_STAT(DataGroup):
    """Stationary MHD eigenfunctions.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        out["alpha"] = np.transpose(out["alpha"], (1,0) )
        out["phi"]   = np.transpose(out["phi"],   (1,0) )
        return out

    def plot_amplitude(self, amplitude="alpha", mode=None, ax=None):
        """
        Plot radial profile of all (or given mode) amplitudes.

        For NONSTAT the profiles are shown at earliest time in dataset.
        TODO: Animate the NONSTAT to show profiles as a function of time.

        Args:
            amplitude : str <br>
                Which amplitude is plotted: "alpha" (magnetic) or "phi" (electric)
            mode : tuple(int,int) <br>
                Mode (n,m) to be plotted or None to plot all modes.
            ax : axes <br>
                Axes where plotting is done or None to create and show new fig.
        """
        import matplotlib.pyplot as plt

        data = self.read()
        isstat = self.isstat()

        axes = ax
        if ax is None:
            fig  = plt.figure()
            axes = fig.add_subplot(1,1,1)

        axes.set_xlabel(r"$\rho$")
        axes.set_xlim(0,1)
        if amplitude == "alpha":
            axes.set_ylabel("Tm")
        else:
            axes.set_ylabel("V")

        amplitude = data[amplitude]
        rho = np.linspace(data["rhomin"], data["rhomax"], int(data["nrho"]))

        for n in range(int(data["nmode"])):
            if mode is None or \
               ( mode[0] == data["nmodes"][n] and mode[1] == data["mmodes"][n] ):
                if isstat:
                    axes.plot(rho, amplitude[:,n] * data["amplitude"][n])
                else:
                    #t = data["tmin"]
                    #tgrid = np.linspace(data["tmin"], data["tmax"], data["ntime"])
                    axes.plot(rho, amplitude[n,0,:] * data["amplitude"][n])

        if ax is None:
            plt.show()

    @staticmethod
    def write_hdf5(fn, nmode, nmodes, mmodes, amplitude, omega, phase, alpha,
                   phi, nrho, rhomin, rhomax, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        nmode : int
            Number of modes.
        nmodes : array_like (nmode,)
            Mode n (toroidal) numbers.
        mmodes : array_like (nmode,)
            Mode m (poloidal) numbers.
        amplitude : array_like (nmode,)
            Mode amplitudies.
        omega : array_like (nmode,)
            Mode frequencies [rad/s].
        omega : array_like (nmode,)
            Mode phases [rad].
        alpha : array_like (nrho, nmode)
            Magnetic perturbation eigenfunctions.
        phi : array_like (nrho, nmode)
            Electric perturbation eigenfunctions.
        nrho : int
            Number of rho grid points.
        rhomin : float
            Minimum value in rho grid.
        rhomax : float
            Maximum value in rho grid.
        desc : str, optional
            Input's description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if nmodes.size != nmode:
            raise ValueError("Shape of nmodes is inconsistent with nmode")
        if mmodes.size != nmode:
            raise ValueError("Shape of mmodes is inconsistent with nmode")
        if alpha.shape != (nrho,nmode):
            raise ValueError("alpha has inconsistent shape")
        if phi.shape   != (nrho,nmode):
            raise ValueError("phi has inconsistent shape")
        alpha = np.transpose(alpha, (1,0) )
        phi   = np.transpose(phi,   (1,0) )

        parent = "mhd"
        group  = "MHD_STAT"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("nmode",  (1,), data=nmode,  dtype="i4")
            g.create_dataset("nrho",   (1,), data=nrho,   dtype="i4")
            g.create_dataset("rhomin", (1,), data=rhomin, dtype="f8")
            g.create_dataset("rhomax", (1,), data=rhomax, dtype="f8")

            g.create_dataset("nmodes", (nmode,), data=nmodes, dtype="i4")
            g.create_dataset("mmodes", (nmode,), data=mmodes, dtype="i4")

            g.create_dataset("amplitude", (nmode,), data=amplitude, dtype="f8")
            g.create_dataset("omega",     (nmode,), data=omega,     dtype="f8")
            g.create_dataset("phase",     (nmode,), data=phase,     dtype="f8")

            g.create_dataset("alpha", (nmode,nrho), data=alpha, dtype="f8")
            g.create_dataset("phi",   (nmode,nrho), data=phi,   dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This method writes two modes with constant eigenfunctions.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return MHD_STAT.write_hdf5(
            fn=fn, nmode=2, nmodes=np.array([1, 2]), mmodes=np.array([3, 4]),
            amplitude=np.array([0.1, 2]), omega=np.array([1, 1.5]),
            phase=np.array([0, 3.141/4]), alpha=np.ones((6,2)),
            phi=np.ones((6,2)), nrho=6, rhomin=0, rhomax=1, desc="DUMMY")

class MHD_NONSTAT(DataGroup):
    """Time-dependent MHD eigenfunctions.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        out["alpha"] = np.transpose(out["alpha"], (2,1,0) )
        out["phi"]   = np.transpose(out["phi"],   (2,1,0) )
        return out

    def plot_amplitude(self, amplitude="alpha", mode=None, ax=None):
        """
        Plot radial profile of all (or given mode) amplitudes.

        For NONSTAT the profiles are shown at earliest time in dataset.
        TODO: Animate the NONSTAT to show profiles as a function of time.

        Args:
            amplitude : str <br>
                Which amplitude is plotted: "alpha" (magnetic) or "phi" (electric)
            mode : tuple(int,int) <br>
                Mode (n,m) to be plotted or None to plot all modes.
            ax : axes <br>
                Axes where plotting is done or None to create and show new fig.
        """
        import matplotlib.pyplot as plt

        data = self.read()
        isstat = self.isstat()

        axes = ax
        if ax is None:
            fig  = plt.figure()
            axes = fig.add_subplot(1,1,1)

        axes.set_xlabel(r"$\rho$")
        axes.set_xlim(0,1)
        if amplitude == "alpha":
            axes.set_ylabel("Tm")
        else:
            axes.set_ylabel("V")

        amplitude = data[amplitude]
        rho = np.linspace(data["rhomin"], data["rhomax"], int(data["nrho"]))

        for n in range(int(data["nmode"])):
            if mode is None or \
               ( mode[0] == data["nmodes"][n] and mode[1] == data["mmodes"][n] ):
                if isstat:
                    axes.plot(rho, amplitude[:,n] * data["amplitude"][n])
                else:
                    #t = data["tmin"]
                    #tgrid = np.linspace(data["tmin"], data["tmax"], data["ntime"])
                    axes.plot(rho, amplitude[n,0,:] * data["amplitude"][n])

        if ax is None:
            plt.show()

    @staticmethod
    def write_hdf5(fn, nmode, nmodes, mmodes, amplitude, omega, phase, alpha,
                   phi, nrho, rhomin, rhomax, ntime, tmin, tmax,
                   desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        nmode : int
            Number of modes.
        nmodes : array_like (nmode,)
            Mode n (toroidal) numbers.
        mmodes : array_like (nmode,)
            Mode m (poloidal) numbers.
        amplitude : array_like (nmode,)
            Mode amplitudies.
        omega : array_like (nmode,)
            Mode frequencies [rad/s].
        omega : array_like (nmode,)
            Mode phases [rad].
        alpha : array_like (nrho, ntime, nmode)
            Magnetic perturbation eigenfunctions.
        phi : array_like (nrho, ntime, nmode)
            Electric perturbation eigenfunctions.
        nrho : int
            Number of rho grid points.
        rhomin : float
            Minimum value in rho grid.
        rhomax : float
            Maximum value in rho grid.
        ntime : int, optional
            Number of time grid points.
        tmin : float
            Minimum value in time grid.
        tmax : float, optional
            Maximum value in time grid.
        desc : str, optional
            Input's description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if nmodes.size != nmode:
            raise ValueError("Shape of nmodes is inconsistent with nmode")
        if mmodes.size != nmode:
            raise ValueError("Shape of mmodes is inconsistent with nmode")
        if alpha.shape != (nrho,ntime,nmode):
            raise ValueError("alpha has inconsistent shape")
        if phi.shape   != (nrho,ntime,nmode):
            raise ValueError("phi has inconsistent shape")
        alpha = np.transpose(alpha, (2,1,0) )
        phi   = np.transpose(phi,   (2,1,0) )

        parent = "mhd"
        group  = "MHD_NONSTAT"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("nmode",  (1,), data=nmode,  dtype="i4")
            g.create_dataset("nrho",   (1,), data=nrho,   dtype="i4")
            g.create_dataset("rhomin", (1,), data=rhomin, dtype="f8")
            g.create_dataset("rhomax", (1,), data=rhomax, dtype="f8")

            g.create_dataset("nmodes", (nmode,), data=nmodes, dtype="i4")
            g.create_dataset("mmodes", (nmode,), data=mmodes, dtype="i4")

            g.create_dataset("amplitude", (nmode,), data=amplitude, dtype="f8")
            g.create_dataset("omega",     (nmode,), data=omega,     dtype="f8")
            g.create_dataset("phase",     (nmode,), data=phase,     dtype="f8")

            g.create_dataset("ntime", (1,), data=ntime, dtype="i4")
            g.create_dataset("tmin",  (1,), data=tmin,  dtype="f8")
            g.create_dataset("tmax",  (1,), data=tmax,  dtype="f8")
            g.create_dataset("alpha", (nmode,ntime,nrho), data=alpha,dtype="f8")
            g.create_dataset("phi",   (nmode,ntime,nrho), data=phi,  dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        This method writes two modes with constant eigenfunctions.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return MHD_NONSTAT.write_hdf5(
            fn=fn, nmode=2, nmodes=np.array([1, 2]), mmodes=np.array([3, 4]),
            amplitude=np.array([0.1, 2]), omega=np.array([1, 1.5]),
            phase=np.array([0, 3.141/4]), alpha=np.ones((6,3,2)),
            phi=np.ones((6,3,2)), nrho=6, rhomin=0, rhomax=1, ntime=3,
            tmin=0, tmax=1, desc="DUMMY")

a5py/a5py/ascot5io/mhd.py
ascot4fusion/ascot5
"""
Marker IO.

File: mrk_prt_shined.py
"""
import h5py
import numpy as np

from .coreio.fileapi import add_group
from a5py.ascot5io.mrk import mrk
import a5py.ascot5io.mrk
import a5py.ascot5io.mrk_shined


def write_hdf5(fn, n, ids, mass, charge,
               r, phi, z, vr, vphi, vz,
               anum, znum, weight, time, desc=None):
    """
    Write particle marker input in hdf5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
        n : int <br>
            Number of markers.
        ids : array_like (n,1) <br>
            Unique identifier for each marker (must be a positive integer).
        mass : array_like (n,1) <br>
            Mass [amu].
        charge : array_like (n,1) <br>
            Charge [e].
        r : array_like (n,1) <br>
            Particle R coordinate [m].
        phi : array_like (n,1) <br>
            Particle phi coordinate [deg].
        z : array_like (n,1) <br>
            Particle z coordinate [m].
        vr : array_like (n,1) <br>
            Particle velocity R-component [m/s].
        vphi : array_like (n,1) <br>
            Particle velocity phi-component [m/s].
        vz : array_like (n,1) <br>
            Particle velocity z-component [m/s].
        anum : array_like (n,1) <br>
            Marker species atomic mass number.
        znum : array_like (n,1) <br>
            Marker species charge number.
        weight : array_like (n,1) <br>
            Particle weight [markers/s].
        time : array_like (n,1) <br>
            Particle initial time [s].
        desc : str, optional <br>
            Input description.

    Returns:
        Name of the new input that was written.
    """
    assert ids.size    == n
    assert mass.size   == n
    assert charge.size == n
    assert r.size      == n
    assert phi.size    == n
    assert z.size      == n
    assert vr.size     == n
    assert vphi.size   == n
    assert vz.size     == n
    assert anum.size   == n
    assert znum.size   == n
    assert weight.size == n
    assert time.size   == n

    parent = "marker_shined"
    group  = "prt_shined"
    gname  = ""

    with h5py.File(fn, "a") as f:
        g = add_group(f, parent, group, desc=desc)
        gname = g.name.split("/")[-1]

        g.create_dataset("n",      (1,1), data=n,      dtype='i8').attrs['unit'] = '1';
        g.create_dataset("r",      (n,1), data=r,      dtype='f8').attrs['unit'] = 'm';
        g.create_dataset("phi",    (n,1), data=phi,    dtype='f8').attrs['unit'] = 'deg';
        g.create_dataset("z",      (n,1), data=z,      dtype='f8').attrs['unit'] = 'm';
        g.create_dataset("vr",     (n,1), data=vr,     dtype='f8').attrs['unit'] = 'm/s';
        g.create_dataset("vphi",   (n,1), data=vphi,   dtype='f8').attrs['unit'] = 'm/s';
        g.create_dataset("vz",     (n,1), data=vz,     dtype='f8').attrs['unit'] = 'm/s';
        g.create_dataset("mass",   (n,1), data=mass,   dtype='f8').attrs['unit'] = 'amu';
        g.create_dataset("charge", (n,1), data=charge, dtype='i4').attrs['unit'] = 'e';
        g.create_dataset("anum",   (n,1), data=anum,   dtype='i4').attrs['unit'] = '1';
        g.create_dataset("znum",   (n,1), data=znum,   dtype='i4').attrs['unit'] = '1';
        g.create_dataset("weight", (n,1), data=weight, dtype='f8').attrs['unit'] = 'markers/s';
        g.create_dataset("time",   (n,1), data=time,   dtype='f8').attrs['unit'] = 's';
        g.create_dataset("id",     (n,1), data=ids,    dtype='i8').attrs['unit'] = '1';

    return gname

def read_hdf5(fn, qid):
    """
    Read particle marker input from HDF5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
        qid : str <br>
            QID of the data to be read.

    Returns:
        Dictionary containing input data.
    """
    prefix='prt_shined'
    return a5py.ascot5io.mrk_shined.read_hdf5(fn, qid, prefix)


class mrk_prt_shined(mrk):
    """
    Object representing particle marker data.
    """

    def read(self):
        return read_hdf5(self._root._ascot.file_getpath(), self.get_qid())


    def write(self, fn, data=None, desc=None):
        if data is None:
            data = self.read()

        return write_hdf5(fn, **data, desc=desc)

a5py/a5py/ascot5io/mrk_prt_shined.py
ascot4fusion/ascot5
'''
Created on Apr 14, 2020

@author: sjjamsa
'''

import copy
import numpy as np
import h5py

from .coreio.treedata import DataGroup

def read_hdf5(fn, qid, prefix):
    """
    Read particle marker input from HDF5 file.

    Args:
        fn : str <br>
            Full path to the HDF5 file.
        qid : str <br>
            QID of the data to be read.

    Returns:
        Dictionary containing input data.
    """

    path = "marker_shined/"+prefix+"_" + qid

    out = {}
    with h5py.File(fn,"r") as f:
        for key in f[path]:
            # Make all read-in datasets 1-d arrays,
            # regardless of their original dimensionality
            d=f[path][key][:]
            out[key] = np.reshape(d,newshape=(d.size,))

    out["ids"] = out["id"]
    del out["id"]
    return out



class mrk(DataGroup):
    '''
    A class acting as a superclass for all marker types.
    '''

    def prune(self,keepNmarkers, data=None, probabilities=None):
        '''
        Keep a subset of markers.
        
        Args:
            keepNmarkers : int <br>
                How many markers to keep.
                
            data=None         : dict <br>
                The data from .read() -method.
                If None, will be automatically read.
                
            probabilities=None
                What is the probability to include each marker
                If None, use the weights, (sum normalized to 1.0).
               
        Returns:
            A dictionary as if from .read()
            
            
            
        
        '''
        
        
        if data is None:
            data = self.read()
        else:
            data = copy.deepcopy(data)
        
        n = data['n'][0]
        totalWeight = np.sum(data["weight"])
        
        if probabilities is None:
            probabilities = data["weight"] / totalWeight
        
        
        
        fortune = np.random.choice(np.arange(n),size=(keepNmarkers,),replace=False,p=probabilities)
        
        for k in data.keys():
            if k=='n':
                data[k][0]=keepNmarkers
                continue
            data[k] = data[k][fortune]

        newWeight = np.sum(data["weight"])
        
        data['weight'] *= totalWeight / newWeight

        return data

a5py/a5py/ascot5io/mrk_shined.py
ascot4fusion/ascot5
"""Input data representing NBI injectors.

NBI injectors are used by BBNBI5 to generate NBI-ion source and calculate
shinethrough.
"""
import numpy as np
import h5py

import a5py.nbi.plot as plot

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class NBI(DataGroup):
    """Object representing nbi data.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = []
        with h5py.File(fn,"r") as f:
            ninj = f[path]["ninj"][0]
            for i in range(0,ninj):
                out.append({})
                for key in f[path+"/inj"+str(i+1)]:
                    out[i][key] = f[path+"/inj"+str(i+1)][key][:]

        return out

    def plot_grid_3D(self, ax=None):
        beams = self.read()

        for beam in beams:
            ax = plot.plot_scatter_3D(beam["beamletx"], beam["beamlety"], beam["beamletz"],
                                      equal = True, axes=ax,
                                      color="red",linewidth=0.75)
        return ax

    def plot_beamlet_3D(self, ax=None):
        beams = self.read()

        for beam in beams:
            ax = plot.plot_arrow_3D(beam["beamletx"], beam["beamlety"], beam["beamletz"],
                                    beam["beamletdx"], beam["beamletdy"], beam["beamletdz"],
                                    axes=ax, arrow_length_ratio=0,
                                    color="green", linewidth=0.1, length=10)

        return ax

    @staticmethod
    def write_hdf5(fn, nbi, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        nbi : array
            Array of dictionaries describing each injector with fields:

            id : int
                Numerical identifier for the injector
            nbeamlet : int
                Number of beamlets in this injector
            beamletx : float
                X coordinates of beamlets [m]
            beamlety : float
                Y coordinates of beamlets [m]
            beamletz : float
                Z coordinates of beamlets [m]
            beamletdx : float
                X components of the unit direction vector of beamlets
            beamletdy : float
                Y components of the unit direction vector of beamlets
            beamletdz : float
                Z components of the unit direction vector of beamlets
            div_h : float
                Horizontal divergence [radians]
            div_v : float
                Vertical divergence [radians]
            div_halo_frac : float
                Fraction of particles with halo divergence
            div_halo_h : float
                Horizontal divergence in halo [radians]
            div_halo_v : float
                Vertical divergence in halo [radians]
            anum : int
                Mass number of injected species
            znum : int
                Nuclear charge number of injected species
            mass : float
                Mass of the injected species [kg]
            energy : float
                Full injection energy [J]
            efrac : array_like (3)
                Particle fractions for full, 1/2 and 1/3 energies
            power : float
                Injected power [W]
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        parent = "nbi"
        group  = "nbi"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("ninj",  (1,), data=len(nbi),    dtype="i4")

            for i in range(len(nbi)):
                ginj = g.create_group("inj"+str(i+1))

                nbeamlet = nbi[i]["nbeamlet"]

                ginj.create_dataset("id",            (1,),
                                    data=nbi[i]["id"],            dtype="i4")
                ginj.create_dataset("nbeamlet",      (1,),
                                    data=nbi[i]["nbeamlet"],      dtype="i4")
                ginj.create_dataset("beamletx",      (nbeamlet,),
                                    data=nbi[i]["beamletx"],      dtype="f8")
                ginj.create_dataset("beamlety",      (nbeamlet,),
                                    data=nbi[i]["beamlety"],      dtype="f8")
                ginj.create_dataset("beamletz",      (nbeamlet,),
                                    data=nbi[i]["beamletz"],      dtype="f8")
                ginj.create_dataset("beamletdx",     (nbeamlet,),
                                    data=nbi[i]["beamletdx"],     dtype="f8")
                ginj.create_dataset("beamletdy",     (nbeamlet,),
                                    data=nbi[i]["beamletdy"],     dtype="f8")
                ginj.create_dataset("beamletdz",     (nbeamlet,),
                                    data=nbi[i]["beamletdz"],     dtype="f8")
                ginj.create_dataset("div_h",         (1,),
                                    data=nbi[i]["div_h"],         dtype="f8")
                ginj.create_dataset("div_v",         (1,),
                                    data=nbi[i]["div_v"],         dtype="f8")
                ginj.create_dataset("div_halo_frac", (1,),
                                    data=nbi[i]["div_halo_frac"], dtype="f8")

                ginj.create_dataset("div_halo_h", (1,),
                                    data=nbi[i]["div_halo_h"], dtype="f8")
                ginj.create_dataset("div_halo_v", (1,),
                                    data=nbi[i]["div_halo_v"], dtype="f8")
                ginj.create_dataset("anum",       (1,), data=nbi[i]["anum"],
                                    dtype="i4")
                ginj.create_dataset("znum",       (1,), data=nbi[i]["znum"],
                                    dtype="i4")
                ginj.create_dataset("mass",       (1,), data=nbi[i]["mass"],
                                    dtype="f8")
                ginj.create_dataset("energy",     (1,), data=nbi[i]["energy"],
                                    dtype="f8")
                ginj.create_dataset("efrac",      (3,), data=nbi[i]["efrac"],
                                    dtype="f8")
                ginj.create_dataset("power",      (1,), data=nbi[i]["power"],
                                    dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output is a very large rectangular wall.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        nbi = {
            "id" : 1,
            "nbeamlet" : 1,
            "beamletx" : np.array([1.0]),
            "beamlety" : np.array([0.0]),
            "beamletz" : np.array([0.0]),
            "beamletdx" : np.array([1.0]),
            "beamletdy" : np.array([0.0]),
            "beamletdz" : np.array([0.0]),
            "div_h" : 0.0,
            "div_v" : 0.0,
            "div_halo_frac" : 0.0,
            "div_halo_h" : 0.0,
            "div_halo_v" : 0.0,
            "anum" : 1.0,
            "znum" : 1.0,
            "mass" : 1.0,
            "energy" : 1.0,
            "efrac" : [1,0,0],
            "power" : 1
        }
        return NBI.write_hdf5(fn, [nbi], desc="DUMMY")

a5py/a5py/ascot5io/nbi.py
ascot4fusion/ascot5
"""Input representing plasma neutral bakcground.

Neutral data is used in simulations with atomic reactions enabled.
"""
import numpy as np
import h5py

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class N0_1D(DataGroup):
    """Constant-on-flux-surfaces neutral profile.
    """

    def read_hdf5(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        out["density"]     = np.transpose(out["density"],     (1,0))
        out["temperature"] = np.transpose(out["temperature"], (1,0))

        return out

    @staticmethod
    def write_hdf5(fn, rhomin, rhomax, nrho,
                   nspecies, anum, znum, density, temperature, maxwellian=1,
                   desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        rhomin : float
            Minimum value in rho grid [1].
        rhomax : float
            Maximum value in rho grid [1].
        nrho : int
            Number of rho grid points.
        nspecies : int
            Number of neutral species.
        anum : array_like (nspecies,1)
            Neutral species' atomic mass number.
        znum array_like (nspecies,1)
            Neutral species' charge number.
        density array_like (nrho,nspecies)
            Neutral species-wise density [m^-3].
        temperature array_like (nrho,nspecies)
            Neutral species-wise temperature [eV].
        maxwellian array_like (nspecies,1)
            Whether species distribution is Maxwellian (1) of monoenergetic (0)
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if density.shape != (nrho,nspecies):
            raise ValueError("Density has invalid shape.")
        if temperature.shape != (nrho,nspecies):
            raise ValueError("Temperature has invalid shape.")
        if anum.size != nspecies or znum.size != nspecies:
            raise ValueError("Anum or Znum has invalid shape.")
        if maxwellian != 1 and maxwellian.size != nspecies:
            raise ValueError("Failed to interpret maxwellian parameter.")

        parent = "neutral"
        group  = "N0_1D"
        gname  = ""

        # Transpose n0 from (rho, spec) to (spec, rho)
        density     = np.transpose(density)
        temperature = np.transpose(temperature)

        if maxwellian == 1:
            maxwellian = np.ones( (int(nspecies),1) )

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("rhomin",   (1,), data=rhomin,   dtype="f8")
            g.create_dataset("rhomax",   (1,), data=rhomax,   dtype="f8")
            g.create_dataset("nrho",     (1,), data=nrho,     dtype="i4")
            g.create_dataset("nspecies", (1,), data=nspecies, dtype="i4")

            g.create_dataset("anum",        (nspecies,),  data=anum,
                             dtype="i4")
            g.create_dataset("znum",        (nspecies,),  data=znum,
                             dtype="i4")
            g.create_dataset("maxwellian",  (nspecies,),  data=maxwellian,
                             dtype="i4")

            g.create_dataset("density",     (nspecies,nrho), data=density,
                             dtype="f8")
            g.create_dataset("temperature", (nspecies,nrho), data=temperature,
                             dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return N0_1D.write_hdf5(
            fn=fn, rhomin=0, rhomax=2, nrho=3, nspecies=1,
            anum=np.array([1]), znum=np.array([1]),
            density=np.ones((3, 1)), temperature=np.ones((3, 1)), desc="DUMMY")

class N0_3D(DataGroup):
    """Non-axisymmetric neutral data.

    This input represents neutral density and temperature that can have 3D
    profile. The data is interpolated on an uniform cylindrical grid using
    linear interpolation. The neutral temperature can either be Maxwellian
    or mono-energetic.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]

        out["density"]     = np.transpose(out["density"],     (3,1,2,0))
        out["temperature"] = np.transpose(out["temperature"], (3,1,2,0))

        return out

    @staticmethod
    def write_hdf5(fn, rmin, rmax, nr, zmin, zmax, nz, phimin, phimax, nphi,
               nspecies, anum, znum, density, temperature, maxwellian=1,
               desc=None):
        """Write input data to the HDF5 file.

        The toroidal angle phi is treated as a periodic coordinate, meaning
        ``A(phi=phimin) == A(phi=phimax)``. However, the phi grid, where input
        arrays are tabulated, is ``linspace(phimin, phimax, nphi+1)[:-1]``
        to avoid storing duplicate data.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        rmin : float
            Minimum value in R grid [m].
        rmax : float
            Maximum value in R grid [m].
        nr : int
            Number of R grid points.
        zmin : float
            Minimum value in z grid [m].
        zmax : float
            Maximum value in z grid [m].
        nz : int
            Number of z grid points.
        phimin : float
            Beginning of the toroidal period [deg].
        phimax : float
            End of the toroidal period [deg].
        nphi : int
            Number of phi grid points.
        nspecies : int
            Number of neutral species.
        anum : array_like (nspecies,1)
            Neutral species' atomic mass number.
        znum : array_like (nspecies,1)
            Neutral species' charge number.
        density : array_like (nr,nphi,nz,nspecies)
            Neutral species-wise density [m^-3].
        temperature : array_like (nr,nphi,nz,nspecies)
            Neutral species-wise temperature [eV].
        maxwellian : array_like (nspecies,1)
            Whether species distribution is Maxwellian (1) of monoenergetic (0)
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if density.shape != (nr,nphi,nz,nspecies):
            raise ValueError("Density has invalid shape.")
        if temperature.shape != (nr,nphi,nz,nspecies):
            raise ValueError("Temperature has invalid shape.")
        if anum.size != nspecies or znum.size != nspecies:
            raise ValueError("Anum or Znum has invalid shape.")
        if maxwellian != 1 and maxwellian.size != nspecies:
            raise ValueError("Maxwellian has invalid shape.")

        parent = "neutral"
        group  = "N0_3D"
        gname  = ""

        # Transpose n0 and t0 from (r, phi, z, spec) to (spec, phi, z, r)
        density     = np.transpose(density,(3,1,2,0))
        temperature = np.transpose(temperature,(3,1,2,0))

        if maxwellian == 1:
            maxwellian = np.ones( (int(nspecies),1) )

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("rmin",     (1,), data=rmin,     dtype="f8")
            g.create_dataset("rmax",     (1,), data=rmax,     dtype="f8")
            g.create_dataset("nr",       (1,), data=nr,       dtype="i4")
            g.create_dataset("phimin",   (1,), data=phimin,   dtype="f8")
            g.create_dataset("phimax",   (1,), data=phimax,   dtype="f8")
            g.create_dataset("nphi",     (1,), data=nphi,     dtype="i4")
            g.create_dataset("zmin",     (1,), data=zmin,     dtype="f8")
            g.create_dataset("zmax",     (1,), data=zmax,     dtype="f8")
            g.create_dataset("nz",       (1,), data=nz,       dtype="i4")
            g.create_dataset("nspecies", (1,), data=nspecies, dtype="i4")
            g.create_dataset("anum",        (nspecies,),
                             data=anum, dtype="i4")
            g.create_dataset("znum",        (nspecies,),
                             data=znum, dtype="i4")
            g.create_dataset("maxwellian",  (nspecies,),
                             data=maxwellian, dtype="i4")
            g.create_dataset("density",     (nspecies,nphi,nz,nr),
                             data=density, dtype="f8")
            g.create_dataset("temperature", (nspecies,nphi,nz,nr),
                             data=temperature, dtype="f8")

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return N0_3D.write_hdf5(
            fn=fn, rmin=0, rmax=100, nr=3, zmin=-100, zmax=100, nz=3, phimin=0,
            phimax=360, nphi=3, nspecies=1, anum=np.array([1]),
            znum=np.array([1]), density=np.ones((3,3,3,1)),
            temperature=np.ones((3,3,3,1)), desc="DUMMY")

a5py/a5py/ascot5io/neutral.py
ascot4fusion/ascot5
"""Simulation options input.
"""
import h5py
import numpy as np
import ast
import warnings

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class Opt(DataGroup):
    """Simulation options.

    Option parameters and default values are defined with properties and
    getters. These should not be accessed directly or altered in any way. Use
    :meth:`read` to obtain options from the HDF5 file.
    """

    def __init__(self, *args, **kwargs):
        """Initialize options attributes with default values.

        This class uses a hack that allows us to do two things simultaneously:

        1. Access and print the option parameters' meta data (name, description,
           and default value) when running the code.
        2. Show that meta data in documentation generated with Sphinx/autodoc.

        The hack is to implement all options attributes as properties even
        though they are not ever used other than to initialize default
        parameters. When adding new options, make sure both the attribute and
        the property is private (to discourage user calling them), and add
        "OPT" prefix for the attribute which allows us to parse options
        attributes from the rest.

        Note that the order in which attributes are declared is preserved.
        """
        super().__init__(*args, **kwargs)

        self._OPT_SIM_MODE                   = 2
        self._OPT_ENABLE_ADAPTIVE            = 1
        self._OPT_RECORD_MODE                = 0
        self._OPT_FIXEDSTEP_USE_USERDEFINED  = 0
        self._OPT_FIXEDSTEP_USERDEFINED      = 1.0e-8
        self._OPT_FIXEDSTEP_GYRODEFINED      = 20
        self._OPT_ADAPTIVE_TOL_ORBIT         = 1.0e-8
        self._OPT_ADAPTIVE_TOL_CCOL          = 1.0e-1
        self._OPT_ADAPTIVE_MAX_DRHO          = 0.1
        self._OPT_ADAPTIVE_MAX_DPHI          = 2.0
        self._OPT_ENDCOND_SIMTIMELIM         = 0
        self._OPT_ENDCOND_CPUTIMELIM         = 0
        self._OPT_ENDCOND_RHOLIM             = 0
        self._OPT_ENDCOND_ENERGYLIM          = 0
        self._OPT_ENDCOND_WALLHIT            = 0
        self._OPT_ENDCOND_MAXORBS            = 0
        self._OPT_ENDCOND_NEUTRALIZED        = 0
        self._OPT_ENDCOND_IONIZED            = 0
        self._OPT_ENDCOND_LIM_SIMTIME        = 1.0
        self._OPT_ENDCOND_MAX_MILEAGE        = 1.0
        self._OPT_ENDCOND_MAX_CPUTIME        = 3600.0
        self._OPT_ENDCOND_MAX_RHO            = 2.0
        self._OPT_ENDCOND_MIN_RHO            = 0.0
        self._OPT_ENDCOND_MIN_ENERGY         = 1.0e3
        self._OPT_ENDCOND_MIN_THERMAL        = 2.0
        self._OPT_ENDCOND_MAX_TOROIDALORBS   = 100
        self._OPT_ENDCOND_MAX_POLOIDALORBS   = 100
        self._OPT_ENABLE_ORBIT_FOLLOWING     = 0
        self._OPT_ENABLE_COULOMB_COLLISIONS  = 0
        self._OPT_ENABLE_MHD                 = 0
        self._OPT_ENABLE_ATOMIC              = 0
        self._OPT_DISABLE_FIRSTORDER_GCTRANS = 0
        self._OPT_DISABLE_ENERGY_CCOLL       = 0
        self._OPT_DISABLE_PITCH_CCOLL        = 0
        self._OPT_DISABLE_GCDIFF_CCOLL       = 0
        self._OPT_REVERSE_TIME               = 0
        self._OPT_ENABLE_DIST_5D             = 0
        self._OPT_ENABLE_DIST_6D             = 0
        self._OPT_ENABLE_DIST_RHO5D          = 0
        self._OPT_ENABLE_DIST_RHO6D          = 0
        self._OPT_DIST_MIN_R                 = 0.1
        self._OPT_DIST_MAX_R                 = 10.0
        self._OPT_DIST_NBIN_R                = 1
        self._OPT_DIST_MIN_PHI               = 0.0
        self._OPT_DIST_MAX_PHI               = 360.0
        self._OPT_DIST_NBIN_PHI              = 1
        self._OPT_DIST_MIN_Z                 = -6.0
        self._OPT_DIST_MAX_Z                 = 6.0
        self._OPT_DIST_NBIN_Z                = 1
        self._OPT_DIST_MIN_RHO               = 0.0
        self._OPT_DIST_MAX_RHO               = 1.0
        self._OPT_DIST_NBIN_RHO              = 1
        self._OPT_DIST_MIN_THETA             = 0.0
        self._OPT_DIST_MAX_THETA             = 360.0
        self._OPT_DIST_NBIN_THETA            = 1
        self._OPT_DIST_MIN_PPA               = -10.0e-20
        self._OPT_DIST_MAX_PPA               = 10.0e-20
        self._OPT_DIST_NBIN_PPA              = 1
        self._OPT_DIST_MIN_PPE               = 0.0
        self._OPT_DIST_MAX_PPE               = 10.0e-20
        self._OPT_DIST_NBIN_PPE              = 1
        self._OPT_DIST_MIN_PR                = 0.0
        self._OPT_DIST_MAX_PR                = 10.0e-20
        self._OPT_DIST_NBIN_PR               = 1
        self._OPT_DIST_MIN_PPHI              = 0.0
        self._OPT_DIST_MAX_PPHI              = 10.0e-20
        self._OPT_DIST_NBIN_PPHI             = 1
        self._OPT_DIST_MIN_PZ                = 0.0
        self._OPT_DIST_MAX_PZ                = 10.0e-20
        self._OPT_DIST_NBIN_PZ               = 1
        self._OPT_DIST_MIN_TIME              = 0.0
        self._OPT_DIST_MAX_TIME              = 10.0
        self._OPT_DIST_NBIN_TIME             = 1
        self._OPT_DIST_MIN_CHARGE            = -100
        self._OPT_DIST_MAX_CHARGE            = 100
        self._OPT_DIST_NBIN_CHARGE           = 1
        self._OPT_ENABLE_DIST_COM            = 0
        self._OPT_DIST_MIN_MU                = 0
        self._OPT_DIST_MAX_MU                = 2.5e-13
        self._OPT_DIST_NBIN_MU               = 100
        self._OPT_DIST_MIN_EKIN              = 0
        self._OPT_DIST_MAX_EKIN              = 1.0e-12
        self._OPT_DIST_NBIN_EKIN             = 50
        self._OPT_DIST_MIN_PTOR              = -1.0e-18
        self._OPT_DIST_MAX_PTOR              = 1.0e-18
        self._OPT_DIST_NBIN_PTOR             = 200
        self._OPT_ENABLE_ORBITWRITE          = 0
        self._OPT_ORBITWRITE_MODE            = 1
        self._OPT_ORBITWRITE_NPOINT          = 100
        self._OPT_ORBITWRITE_POLOIDALANGLES  = [0.0]
        self._OPT_ORBITWRITE_TOROIDALANGLES  = [0.0]
        self._OPT_ORBITWRITE_RADIALDISTANCES = [1.0]
        self._OPT_ORBITWRITE_INTERVAL        = 0.0
        self._OPT_ENABLE_TRANSCOEF           = 0
        self._OPT_TRANSCOEF_INTERVAL         = 0
        self._OPT_TRANSCOEF_NAVG             = 5
        self._OPT_TRANSCOEF_RECORDRHO        = 0

    @property
    def _SIM_MODE(self):
        """Simulation mode (1, 2, 3, 4)

        - 1 Gyro-orbit
        - 2 Guiding center
        - 3 Hybrid
        - 4 Magnetic field lines
        """
        return self._OPT_SIM_MODE

    @property
    def _ENABLE_ADAPTIVE(self):
        """Use adaptive time-step (0, 1)

        This option is used only if SIM_MODE = 2 or 3. Gyro-orbit
        simulations are always done with fixed time-step and magnetic field line
        simulations with adaptive time-step.

        - 0 Use fixed time-step
        - 1 Use adaptive time-step
        """
        return self._OPT_ENABLE_ADAPTIVE

    @property
    def _RECORD_MODE(self):
        """Record GOs as GCs in diagnostics (0, 1)

        - 0 Record GOs as GOs
        - 1 Record GOs as GCs
        """
        return self._OPT_RECORD_MODE

    @property
    def _FIXEDSTEP_USE_USERDEFINED(self):
        """Define fixed time-step value explicitly (0,1)

        Note: The adaptive scheme uses fixed time-step value as an initial step.

        - 0 Calculate time-step from FIXEDSTEP_NSTEPS_PER_GYROTIME
        - 1 Use opt.opt.FIXEDSTEP_USERDEFINED as a time-step
        """
        return self._OPT_FIXEDSTEP_USE_USERDEFINED

    @property
    def _FIXEDSTEP_USERDEFINED(self):
        """User-defined time-step [s]
        """
        return self._OPT_FIXEDSTEP_USERDEFINED

    @property
    def _FIXEDSTEP_GYRODEFINED(self):
        """Time-step is 2pi / ( gyrofrequency * N ) where N is this parameter
        """
        return self._OPT_FIXEDSTEP_GYRODEFINED

    @property
    def _ADAPTIVE_TOL_ORBIT(self):
        """Relative error tolerance for orbit following in adaptive scheme
        """
        return self._OPT_ADAPTIVE_TOL_ORBIT

    @property
    def _ADAPTIVE_TOL_CCOL(self):
        """Relative error tolerance for Coulomb collisions in adaptive scheme
        """
        return self._OPT_ADAPTIVE_TOL_CCOL

    @property
    def _ADAPTIVE_MAX_DRHO(self):
        """Maximum allowed change in rho during one time-step in adaptive scheme
        """
        return self._OPT_ADAPTIVE_MAX_DRHO

    @property
    def _ADAPTIVE_MAX_DPHI(self):
        """Maximum allowed change in phi during one time-step in adaptive scheme
        """
        return self._OPT_ADAPTIVE_MAX_DPHI

    @property
    def _ENDCOND_SIMTIMELIM(self):
        """Terminate when marker time passes ENDCOND_LIM_SIMTIME or when marker
        time has advanced ENDCOND_MAX_MILEAGE in a simulation

        Marker is terminated if t > ENDCOND_MAX_MILEAGE or
        t0 + t > ENDCOND_LIM_SIMTIME where t0 is marker's initial time and t
        the time it has been simulated. Note that if time is reversed, the
        simulation is instead terminated when t0 + t < ENDCOND_LIM_SIMTIME.
        See also ENDCOND_CPUTIMELIM.
        """
        return self._OPT_ENDCOND_SIMTIMELIM

    @property
    def _ENDCOND_CPUTIMELIM(self):
        """Terminate marker when the computer has spent ENDCOND_MAX_CPUTIME
        amount of real time to simulate it

        This limit should be rarely used  as its intended use is in debugging
        to stop markers stuck in a loop.
        """
        return self._OPT_ENDCOND_CPUTIMELIM

    @property
    def _ENDCOND_RHOLIM(self):
        """Terminate if marker goes outside given rho boundaries

        The boundaries are defined by ENDCOND_MAX_RHO and ENDCOND_MAX_RHO.
        """
        return self._OPT_ENDCOND_RHOLIM

    @property
    def _ENDCOND_ENERGYLIM(self):
        """Terminate when marker energy is below a user-specified value

        The user specified values are ENDCOND_MIN_ENERGY and
        ENDCOND_MIN_ENERGY_TIMES_THERMAL. Marker is terminated when either
        of these limits is reached.
        """
        return self._OPT_ENDCOND_ENERGYLIM

    @property
    def _ENDCOND_WALLHIT(self):
        """Terminate when marker intersects a wall element
        """
        return self._OPT_ENDCOND_WALLHIT

    @property
    def _ENDCOND_MAXORBS(self):
        """Terminate when marker has completed user-specified number of orbits

        Limit ENDCOND_MAX_TOROIDALORBS is used for a number of toroidal
        and ENDCOND_MAX_POLOIDALORBS for poloidal orbits.

        - 0 The end condition is not active.
        - 1 Marker is terminated when either of these limits is reached.
        - 2 Marker is terminated when both limits are reached.
        """
        return self._OPT_ENDCOND_MAXORBS

    @property
    def _ENDCOND_NEUTRALIZED(self):
        """Terminate when the marker becomes neutral
        """
        return self._OPT_ENDCOND_NEUTRALIZED

    @property
    def _ENDCOND_IONIZED(self):
        """Terminate when the marker becomes ionized.
        """
        return self._OPT_ENDCOND_IONIZED

    @property
    def _ENDCOND_LIM_SIMTIME(self):
        """Time when the simulation stops [s]
        """
        return self._OPT_ENDCOND_LIM_SIMTIME

    @property
    def _ENDCOND_MAX_MILEAGE(self):
        """The maximum amount of time this marker is simulated [s]
        """
        return self._OPT_ENDCOND_MAX_MILEAGE

    @property
    def _ENDCOND_MAX_CPUTIME(self):
        """Maximum cpu time per marker [s]
        """
        return self._OPT_ENDCOND_MAX_CPUTIME

    @property
    def _ENDCOND_MAX_RHO(self):
        """Maximum rho value
        """
        return self._OPT_ENDCOND_MAX_RHO

    @property
    def _ENDCOND_MIN_RHO(self):
        """Minimum rho value
        """
        return self._OPT_ENDCOND_MIN_RHO

    @property
    def _ENDCOND_MIN_ENERGY(self):
        """Minimum energy [eV]
        """
        return self._OPT_ENDCOND_MIN_ENERGY

    @property
    def _ENDCOND_MIN_THERMAL(self):
        """Minimum energy limit is local ion thermal energy times this value
        """
        return self._OPT_ENDCOND_MIN_THERMAL

    @property
    def _ENDCOND_MAX_TOROIDALORBS(self):
        """Maximum number of toroidal orbits
        """
        return self._OPT_ENDCOND_MAX_TOROIDALORBS

    @property
    def _ENDCOND_MAX_POLOIDALORBS(self):
        """Maximum number of poloidal orbits
        """
        return self._OPT_ENDCOND_MAX_POLOIDALORBS

    @property
    def _ENABLE_ORBIT_FOLLOWING(self):
        """Trace markers in an electromagnetic field
        """
        return self._OPT_ENABLE_ORBIT_FOLLOWING

    @property
    def _ENABLE_COULOMB_COLLISIONS(self):
        """Markers experience Coulomb collisions with background plasma
        """
        return self._OPT_ENABLE_COULOMB_COLLISIONS

    @property
    def _ENABLE_MHD(self):
        """Include MHD perturbations to orbit-following
        """
        return self._OPT_ENABLE_MHD

    @property
    def _ENABLE_ATOMIC(self):
        """Markers can undergo atomic reactions with background plasma
        or neutrals

        - 0 Atomic reactions are turned off.
        - 1 Atomic reactions are on but marker is terminated when outside
            the reaction data domain.
        - 2 Atomic reactions are on but they are ignored when marker is
            outside the reaction data domain.
        """
        return self._OPT_ENABLE_ATOMIC

    @property
    def _DISABLE_FIRSTORDER_GCTRANS(self):
        """Disable first order guiding center transformation in velocity space
        """
        return self._OPT_DISABLE_FIRSTORDER_GCTRANS

    @property
    def _DISABLE_ENERGY_CCOLL(self):
        """Disable guiding center energy collisions
        """
        return self._OPT_DISABLE_ENERGY_CCOLL

    @property
    def _DISABLE_PITCH_CCOLL(self):
        """Disable guiding center pitch collisions
        """
        return self._OPT_DISABLE_PITCH_CCOLL

    @property
    def _DISABLE_GCDIFF_CCOLL(self):
        """Disable guiding center spatial diffusion
        """
        return self._OPT_DISABLE_GCDIFF_CCOLL

    @property
    def _REVERSE_TIME(self):
         """Trace markers backwards in time.

         Collision operator isn't reversible so disable collisions if this
         option is used. Also when tracing markers, the simulation stops when
         marker time is below ENDCOND_LIM_SIMTIME.
         """
         return self._OPT_REVERSE_TIME


    @property
    def _ENABLE_DIST_5D(self):
        """Collect distribution histogram in [R, phi, z, ppa, ppe, t, q]

        The coordinates are:

        - R   major radius
        - phi toroidal angle
        - z   z-coordinate
        - ppa momentum component parallel to magnetic field
        - ppe momentum component perpendicular to magnetic field
        - t   time
        - q   charge
        """
        return self._OPT_ENABLE_DIST_5D

    @property
    def _ENABLE_DIST_6D(self):
        """Collect distribution histogram in [R, phi, z, pR, pphi, pz, t, q]

        The coordinates are:

        - R    major radius
        - phi  toroidal angle
        - z    z-coordinate
        - pR   momentum R-component
        - pphi momentum phi-component
        - pz   momentum z-component
        - t    time
        - q    charge
        """
        return self._OPT_ENABLE_DIST_6D

    @property
    def _ENABLE_DIST_RHO5D(self):
        """Collect distribution histogram in [rho, pol, phi, ppa, ppe, t, q]

        The coordinates are:

        - rho  flux surface
        - pol  poloidal angle
        - phi  toroidal angle
        - z    z-coordinate
        - ppa  momentum component parallel to magnetic field
        - ppe  momentum component perpendicular to magnetic field
        - t    time
        - q    charge
        """
        return self._OPT_ENABLE_DIST_RHO5D

    @property
    def _ENABLE_DIST_RHO6D(self):
        """Collect distribution histogram in [rho, pol, phi, pR, pphi, pz, t, q]

        The coordinates are:

        - rho  flux surface
        - pol  poloidal angle
        - phi  toroidal angle
        - z    z-coordinate
        - pR   momentum R-component
        - pphi momentum phi-component
        - pz   momentum z-component
        - t    time
        - q    charge
        """
        return self._OPT_ENABLE_DIST_RHO6D

    @property
    def _DIST_MIN_R(self):
        """Minimum bin edge for R coordinate [m]
        """
        return self._OPT_DIST_MIN_R

    @property
    def _DIST_MAX_R(self):
        """Maximum bin edge for R coordinate [m]
        """
        return self._OPT_DIST_MAX_R

    @property
    def _DIST_NBIN_R(self):
        """Number of bins the interval [MIN_R, MAX_R] is divided to
        """
        return self._OPT_DIST_NBIN_R

    @property
    def _DIST_MIN_PHI(self):
        """Minimum bin edge for phi coordinate [deg]
        """
        return self._OPT_DIST_MIN_PHI

    @property
    def _DIST_MAX_PHI(self):
        """Maximum bin edge for phi coordinate [deg]
        """
        return self._OPT_DIST_MAX_PHI

    @property
    def _DIST_NBIN_PHI(self):
        """Number of bins the interval [MIN_PHI, MAX_PHI] is divided to
        """
        return self._OPT_DIST_NBIN_PHI

    @property
    def _DIST_MIN_Z(self):
        """Minimum bin edge for z coordinate [m]
        """
        return self._OPT_DIST_MIN_Z

    @property
    def _DIST_MAX_Z(self):
        """Maximum bin edge for z coordinate [m]
        """
        return self._OPT_DIST_MAX_Z

    @property
    def _DIST_NBIN_Z(self):
        """Number of bins the interval [MIN_Z, MAX_Z] is divided to
        """
        return self._OPT_DIST_NBIN_Z

    @property
    def _DIST_MIN_RHO(self):
        """Minimum bin edge for rho coordinate [1]
        """
        return self._OPT_DIST_MIN_RHO

    @property
    def _DIST_MAX_RHO(self):
        """Maximum bin edge for rho coordinate [1]
        """
        return self._OPT_DIST_MAX_RHO

    @property
    def _DIST_NBIN_RHO(self):
        """Number of bins the interval [MIN_RHO, MAX_RHO] is divided to
        """
        return self._OPT_DIST_NBIN_RHO

    @property
    def _DIST_MIN_THETA(self):
        """Minimum bin edge for theta coordinate [deg]
        """
        return self._OPT_DIST_MIN_THETA

    @property
    def _DIST_MAX_THETA(self):
        """Maximum bin edge for theta coordinate [deg]
        """
        return self._OPT_DIST_MAX_THETA

    @property
    def _DIST_NBIN_THETA(self):
        """Number of bins the interval [MIN_THETA, MAX_THETA] is divided to
        """
        return self._OPT_DIST_NBIN_THETA

    @property
    def _DIST_MIN_PPA(self):
        """Minimum bin edge for ppa coordinate [kg m/s]
        """
        return self._OPT_DIST_MIN_PPA

    @property
    def _DIST_MAX_PPA(self):
        """Maximum bin edge for ppa coordinate [kg m/s]
        """
        return self._OPT_DIST_MAX_PPA

    @property
    def _DIST_NBIN_PPA(self):
        """Number of bins the interval [MIN_PPA, MAX_PPA] is divided to
        """
        return self._OPT_DIST_NBIN_PPA

    @property
    def _DIST_MIN_PPE(self):
        """Minimum bin edge for ppe coordinate [kg m/s]
        """
        return self._OPT_DIST_MIN_PPE

    @property
    def _DIST_MAX_PPE(self):
        """Maximum bin edge for ppe coordinate [kg m/s]
        """
        return self._OPT_DIST_MAX_PPE

    @property
    def _DIST_NBIN_PPE(self):
        """Number of bins the interval [MIN_PPE, MAX_PPE] is divided to
        """
        return self._OPT_DIST_NBIN_PPE

    @property
    def _DIST_MIN_PR(self):
        """Minimum bin edge for pR coordinate [kg m/s]
        """
        return self._OPT_DIST_MIN_PR

    @property
    def _DIST_MAX_PR(self):
        """Maximum bin edge for pR coordinate [kg m/s]
        """
        return self._OPT_DIST_MAX_PR

    @property
    def _DIST_NBIN_PR(self):
        """Number of bins the interval [MIN_PR, MAX_PR] is divided to
        """
        return self._OPT_DIST_NBIN_PR

    @property
    def _DIST_MIN_PPHI(self):
        """Minimum bin edge for pphi coordinate [kg m/s]
        """
        return self._OPT_DIST_MIN_PPHI

    @property
    def _DIST_MAX_PPHI(self):
        """Maximum bin edge for pphi coordinate [kg m/s]
        """
        return self._OPT_DIST_MAX_PPHI

    @property
    def _DIST_NBIN_PPHI(self):
        """Number of bins the interval [MIN_PPHI, MAX_PPHI] is divided to
        """
        return self._OPT_DIST_NBIN_PPHI

    @property
    def _DIST_MIN_PZ(self):
        """Minimum bin edge for pz coordinate [kg m/s]
        """
        return self._OPT_DIST_MIN_PZ

    @property
    def _DIST_MAX_PZ(self):
        """Maximum bin edge for pz coordinate [kg m/s]
        """
        return self._OPT_DIST_MAX_PZ

    @property
    def _DIST_NBIN_PZ(self):
        """Number of bins the interval [MIN_PZ, MAX_PZ] is divided to
        """
        return self._OPT_DIST_NBIN_PZ

    @property
    def _DIST_MIN_TIME(self):
        """Minimum bin edge for time coordinate [s]
        """
        return self._OPT_DIST_MIN_TIME

    @property
    def _DIST_MAX_TIME(self):
        """Maximum bin edge for time coordinate [s]
        """
        return self._OPT_DIST_MAX_TIME

    @property
    def _DIST_NBIN_TIME(self):
        """Number of bins the interval [MIN_TIME, MAX_TIME] is divided to
        """
        return self._OPT_DIST_NBIN_TIME

    @property
    def _DIST_MIN_CHARGE(self):
        """Minimum bin edge for charge coordinate [e]
        """
        return self._OPT_DIST_MIN_CHARGE

    @property
    def _DIST_MAX_CHARGE(self):
        """Maximum bin edge for charge coordinate [e]
        """
        return self._OPT_DIST_MAX_CHARGE

    @property
    def _DIST_NBIN_CHARGE(self):
        """Number of bins the interval [MIN_CHARGE, MAX_CHARGE] is divided to
        """
        return self._OPT_DIST_NBIN_CHARGE

    @property
    def _ENABLE_DIST_COM(self):
        """Collect constant-of-motion distribution histogram [mu, Ekin, Ptor]

        The coordinates are:

        - mu   magnetic moment (first adiabatic invariant)
        - Ekin kinetic energy
        - Ptor canonical toroidal angular momentum
        """
        return self._OPT_ENABLE_DIST_COM

    @property
    def _DIST_MIN_MU(self):
        """Minimum bin edge for the mu coordinate in COM-dist [J/T]
        """
        return self._OPT_DIST_MIN_MU

    @property
    def _DIST_MAX_MU(self):
        """Maximum bin edge for the mu coordinate in COM-dist [J/T]
        """
        return self._OPT_DIST_MAX_MU

    @property
    def _DIST_NBIN_MU(self):
        """Number of bins the interval [MIN_MU, MAX_MU] is divided to
        """
        return self._OPT_DIST_NBIN_MU

    @property
    def _DIST_MIN_EKIN(self):
        """Minimum bin edge for the energy coordinate [J]
        """
        return self._OPT_DIST_MIN_EKIN

    @property
    def _DIST_MAX_EKIN(self):
        """Maximum bin edge for the energy coordinate [J]
        """
        return self._OPT_DIST_MAX_EKIN

    @property
    def _DIST_NBIN_EKIN(self):
        """Number of bins the interval [MIN_EKIN, MAX_EKIN] is divided to
        """
        return self._OPT_DIST_NBIN_EKIN

    @property
    def _DIST_MIN_PTOR(self):
        """Minimum bin edge for the Ptor coordinate [kg m^2/s]
        """
        return self._OPT_DIST_MIN_PTOR

    @property
    def _DIST_MAX_PTOR(self):
        """Maximum bin edge for the Ptor coordinate [kg m^2/s]
        """
        return self._OPT_DIST_MAX_PTOR

    @property
    def _DIST_NBIN_PTOR(self):
        """Number of bins the interval [MIN_PTOR, MAX_PTOR] is divided to
        """
        return self._OPT_DIST_NBIN_PTOR

    @property
    def _ENABLE_ORBITWRITE(self):
        """Enable diagnostics that store marker orbit

        - 0 Marker orbit diagnostics are not collected
        - 1 Marker orbit diagnostics are collected
        """
        return self._OPT_ENABLE_ORBITWRITE

    @property
    def _ORBITWRITE_MODE(self):
        """What kind of marker orbit diagnostics are collected

        These are only used if ENABLE_ORBITWRITE is active.

        - 0 When marker crosses a plane (Poincare-plot)
        - 1 Between given time intervals
        """
        return self._OPT_ORBITWRITE_MODE

    @property
    def _ORBITWRITE_NPOINT(self):
        """Maximum number of points (per marker) to be written

        If this number is exceeded when marker is being simulated, the oldest
        points will be replaced as long as the simulation continues. Thus,
        this parameter is effectively the number of marker's last positions
        that are stored.
        """
        return self._OPT_ORBITWRITE_NPOINT

    @property
    def _ORBITWRITE_POLOIDALANGLES(self):
        """Poloidal angles of toroidal planes where toroidal plots are collected

        Used when ENABLE_ORBITWRITE = 1 and ORBITWRITE_MODE = 0.
        """
        return self._OPT_ORBITWRITE_POLOIDALANGLES

    @property
    def _ORBITWRITE_TOROIDALANGLES(self):
        """Toroidal angles of poloidal planes where poloidal plots are collected

        Used when ENABLE_ORBITWRITE = 1 and ORBITWRITE_MODE = 0.
        """
        return self._OPT_ORBITWRITE_TOROIDALANGLES

    @property
    def _ORBITWRITE_RADIALDISTANCES(self):
        """Minor radius coordinate where radial plots are collected

        Used when ENABLE_ORBITWRITE = 1 and ORBITWRITE_MODE = 0.
        """
        return self._OPT_ORBITWRITE_RADIALDISTANCES

    @property
    def _ORBITWRITE_INTERVAL(self):
        """Time interval for writing marker state [s]

        Used when ENABLE_ORBITWRITE = 1 and ORBITWRITE_MODE = 1.
        """
        return self._OPT_ORBITWRITE_INTERVAL

    @property
    def _ENABLE_TRANSCOEF(self):
        """Enable evaluation of transport coefficients.

        - 0 Transport coefficients are not collected
        - 1 Transport coefficients are collected
        """
        return self._OPT_ENABLE_TRANSCOEF

    @property
    def _TRANSCOEF_INTERVAL(self):
        """Time interval for recording data points

        The data points are recorded (at the outer mid-plane crossing) if this
        interval has passed from the previous recording.
        """
        return self._OPT_TRANSCOEF_INTERVAL

    @property
    def _TRANSCOEF_NAVG(self):
        """Number of subsequent data points that are averaged before calculating
        coefficients to reduce noise.
        """
        return self._OPT_TRANSCOEF_NAVG

    @property
    def _TRANSCOEF_RECORDRHO(self):
        """Record coefficients in terms of normalized poloidal flux instead of
        meters.
        """
        return self._OPT_TRANSCOEF_RECORDRHO

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.

        Raises
        ------
        AscotIOException
            If stored options were unviable.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        defopt = Opt.get_default()
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                if key not in defopt.keys():
                    warnings.warn("Unknown option " + key + " ignored.")
                    continue
                val = f[path][key][:]

                # Take type from the default parameter
                if isinstance(defopt[key], list):
                    try:
                        val[0]
                        out[key] = type(defopt[key])(val)
                    except Exception:
                        out[key] = val
                else:
                    out[key] = type(defopt[key])(val)

        for o in defopt.keys():
            if o not in out:
                raise ValueError("Missing parameter: " + o)

        return out

    def tostring(self, descriptions=True, aslist=False):
        """Convert options to string representation.

        Parameters
        ----------
        descriptions : bool, optional
            If True, section headers and descriptions are added above
            the parameters to make options more readable.
        aslist : bool, optional
            Instead of line breaks, separate each line as a list item.

        Returns
        -------
        opt : str or list [str]
            String where option parameters are written on each line as
            "<PARAMETER> = <VALUE>".
        """
        opt = self.read()

        def trim(val):
            """Make sure values are not too accurate or have too many leading
            zeroes.
            """
            if np.abs(val) >= 1e4 or (np.abs(val) <=1e-4 and val != 0):
                return "{0:e}".format(val)
            else:
                return "{0:g}".format(val)

        # Convert all parameters to strings first
        for o, val in opt.items():
            if isinstance(val, list):
                for i, v in enumerate(val):
                    val[i] = float(trim(v))
                opt[o] = str(val)
            else:
                opt[o] = trim(val)

        # Create an array of strings
        out = []
        for o in self.__dict__.keys():
            if len(o) < 5 or o[:5] != "_OPT_":
                continue

            # Extract parameter info
            attr = o[4:]
            desc = getattr(Opt, attr).__doc__
            name = getattr(Opt, attr).fget.__name__[1:]
            dval = opt[name]

            # Add decorations
            if descriptions:

                def makebanner(title):
                    """Makes a banner with given title.
                    """
                    nchar = len(title)
                    ljust = int((76 - nchar) / 2)
                    banner = []
                    banner += ["#"  + "*"*78 +  "#"]
                    banner += ["#*" + " "*ljust + title
                               + " "*(76 - nchar - ljust) + "*#"]
                    banner += ["#*" + " "*76 + "*#"]
                    banner += ["#"  + "*"*78 +  "#"]
                    banner += [""]
                    return banner

                if name == "SIM_MODE":
                    out.extend(makebanner("SIMULATION MODE AND TIME-STEP"))
                elif name == "ENDCOND_SIMTIMELIM":
                    out.extend(makebanner("END CONDITIONS"))
                elif name == "ENABLE_ORBIT_FOLLOWING":
                    out.extend(makebanner("PHYSICS"))
                elif name == "ENABLE_DIST_5D":
                    out.extend(makebanner("DISTRIBUTIONS"))
                elif name == "ENABLE_ORBITWRITE":
                    out.extend(makebanner("ORBIT WRITE"))
                elif name == "ENABLE_TRANSCOEF":
                    out.extend(makebanner("TRANSPORT COEFFICIENT"))

                # Clean docstrings a little bit by removing extra whispace and
                # empty lines.
                for d in desc.splitlines():
                    d = d.lstrip(" ")
                    if len(d) > 1: out.append("# " + d)

            # Write parameter and value after possible decorations
            out.append(name + " = " + str(dval))
            out.append("") # Empty line to separate parameters

        if not aslist:
            string = ""
            for o in out:
                string += o + "\n"

            out = string

        return out

    def new(self, desc=None, **kwargs):
        """Write new options with updated parameters.

        This method reads the current options, updates the given parameters,
        and writes the updated options as a new input.

        Parameters
        ----------
        desc : str, optional
            Input description.
        **kwargs
            <name> : <value> pairs for each options parameter that are updated.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If arguments contain unknown parameters.
        """
        options = self.read()
        for o, val in kwargs.items():
            if not o in options:
                raise ValueError("Unknown parameter: " + o)
            options[o] = val

        return self._root.create_input("opt", desc=desc, **options)

    @staticmethod
    def write_hdf5(fn, desc=None, **kwargs):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to HDF5 file.
        desc : str, optional
            Input description.
        **kwargs
            <name> : <value> pairs for each options parameter

            Hint: :meth:`get_default` generates a default options dictionary
            from which one can pick parameters to change.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If arguments contain unknown parameters or if there are parameters
            missing.
        """
        opt = Opt.get_default()
        for o in opt.keys():
            if not o in kwargs:
                raise ValueError("Missing parameter: " + o)

        for o in kwargs.keys():
            if not o in opt:
                raise ValueError("Unknown parameter: " + o)

        parent = "options"
        group  = "opt"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            # Convert all options to numpy float arrays and write
            for param, data in kwargs.items():
                data = np.asarray(data)

                data = data.astype("f8")
                d = g.create_dataset(param, (data.size,), data=data)

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The default options are written.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return Opt.write_hdf5(fn=fn, desc="DUMMY", **Opt.get_default())

    @staticmethod
    def get_default():
        """Get default option parameters.

        The default options have all "ENABLE" settings disabled. The default
        values are not suitable for any proper simulation.

        Returns
        -------
        options : dict
            Default options in a format that can be passed to
            :meth:`write_hdf5`.
        """
        out = {}
        defopt = Opt(None, None)
        for opt in defopt.__dict__.keys():
            if len(opt) < 5 or opt[:5] != "_OPT_":
                continue

            attr = opt[4:]
            name = getattr(Opt, attr).fget.__name__[1:]
            dval = getattr(defopt, attr)
            out[name] = dval

        return out

    @staticmethod
    def validate():
        """
        Validate options and return names of faulty parameters.
        """
        opt  = Opt.get_default()
        opt0 = {}
        for o in opt:
            if len(o) > 2:
                opt0[o[0]] = o[3]

        msg = []
        opt = self.read()
        for o in opt0.keys():
            if o in opt.keys() and not opt0[o].validate(opt[o]):
                msg += [o]

        return msg

    @staticmethod
    def convert_string(lines):
        """Converts a string to :class:`Opt` input.

        Parameters
        ----------
        lines : str or [str]
            String with line breaks or list of strings where each item is
            a single line.

            Options are interpreted from a string assuming that the relevant
            lines have format "<PARAMETER> = <VALUE>" and nothing else.
            The parameter name must be in all uppercase. Lines that don't start
            with a recognizable parameter are ignored.

        Returns
        -------
        out : dict
            String converted as an input for :meth:`write_hdf5`.
        """
        if not isinstance(lines, list):
            lines = lines.splitlines()

        opt = {}
        defopt = Opt.get_default()
        for line in lines:
            paramval = line.split("=")
            if len(paramval) != 2 or \
               len(paramval[0].split()) > 1 or len(paramval[1].split()) > 1:
                continue

            param = paramval[0].strip()
            val   = paramval[1].strip()
            if param != param.upper():
                continue
            if param not in defopt.keys():
                warnings.warn("Possibly unknown option " + param + " ignored.")
                continue
            if isinstance(defopt[param], list):
                try:
                    opt[param] = [float(val)]
                except ValueError:
                   opt[param] = ast.literal_eval(val)
            else:
                opt[param] = type(defopt[param])(val)

        return opt

a5py/a5py/ascot5io/options.py
ascot4fusion/ascot5
"""Orbit data IO module.
"""
import numpy as np
import h5py
import unyt
import a5py.physlib as physlib

from .coreio import fileapi
from .coreio.treedata import DataContainer


class Orbits(DataContainer):
    """Orbit diagnostics that collect marker phase-space coordinates and related
    quantities at certain points along the marker orbit.
    """

    GYROORBIT     = 1
    GUIDINGCENTER = 2
    FIELDLINE     = 3

    def read(self):
        """Read raw state data to a dictionary.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]

        return out

    def get(self, inistate, endstate, *qnt):
        """Return marker quantity.

        This function accesses the orbit data within the HDF5 file and uses that
        to evaluate the queried quantity. The evaluated quantity at a given
        position corresponds to that mode which was active at the simulation:
        GO simulations return particle phase-space, GC guiding-center
        phase-space and hybrid depends on whether marker was GO or GC at that
        moment when data point was written.

        Parameters
        ----------
        inistate : :class:`State`
            Inistate is needed to evaluate some orbit quantities.
        endstate : :class:`State`
            Endstate is needed to evaluate some orbit quantities.
        *qnt : str
            Names of the quantities.

        Returns
        -------
        *value : array_like
            The quantities as an array ordered by marker ID (major) and mileage
            (minor).
        """
        # Prepare helper variables and functions
        def _val(q, mask=None):
            """Read quantity from HDF5.
            """
            with self as h5:
                if q in h5:
                    return fileapi.read_data(h5, q)
            return None

        # Sort using the fact that inistate.get return values ordered by ID
        # and also np.unique returns indices that produce a sorted array.
        mode    = _val("simmode")
        _, idx  = np.unique(_val("ids").v, return_inverse=True)
        mass    = inistate.get("mass")[0][idx]
        time    = inistate.get("time")[0][idx]
        connlen = inistate.get("mileage")[0][idx] - _val("mileage")

        # Only field lines are constant in time
        if not Orbits.FIELDLINE in mode: time = time + _val("mileage")

        def _eval(q, mask=None):
            """Evaluate input quantities at marker position.
            """
            return self._root._ascot.input_eval(
                _val("r", mask=mask), _val("phi",  mask=mask),
                _val("z", mask=mask), time[mask], *[q])

        return Orbits._getactual(mass, time, connlen, mode, _val, _eval, *qnt)

    @staticmethod
    def _getactual(mass, time, totmil, mode, _val, _eval, *qnt):
        """Calculate orbit quantities using the helper functions and data.

        Parameters
        ----------
        mass : array_like, (n,)
            Marker mass repeated at each orbit point.
        time : array_like, (n,)
            Laboratory time at each orbit point.
        totmil : array_like, (n,)
            Total mileage (from endstate) repeated at each orbit point.
        mode : array_like, (n,)
            The simmode value at each orbit point.
        _val : callable
            Function that returns a stored orbit parameter at masked positions.

            ``_val(qnt : str, mask : array_like) -> value``
        _eval : callable
            Function that returns interpolated input quantity at masked
            positions.

            ``_eval(qnt : str, mask : array_like) -> value``
        *qnt : str
            Names of the quantities.

        Returns
        -------
        *value : array_like
            The quantities as an array ordered by marker ID (major) and mileage
            (minor).
        """
        items = [None]*len(qnt)
        def add(q, val):
            if q in qnt:
                items[qnt.index(q)] = val()

        # Some helper quantities as functions so they are evaluated only
        # when needed. Mask is used to separate GOs, GCs, and FLs since
        # they are all stored in the same data.
        bvec = lambda mask : unyt.unyt_array(
            [_val("br", mask), _val("bphi", mask), _val("bz", mask)])
        pvecprt = lambda mask : unyt.unyt_array(
            [_val("pr", mask), _val("pphi", mask), _val("pz", mask)])
        def pvecgc(mask):
            bnorm = np.sqrt(np.sum(bvec(mask)**2,axis=0))
            pnorm = physlib.momentum_muppar(
                mass[mask], _val("mu", mask),
                _val("ppar", mask), bnorm)
            bhat  = bvec(mask) / bnorm
            e1 = np.zeros(bhat.shape)
            e1[2,:] = 1
            e2 = np.cross(bhat.T, e1.T).T
            e1 = e2 / np.sqrt(np.sum(e2**2, axis=0))
            e2 = np.cross(bhat.T, e1.T).T
            perphat = -np.sin(_val("zeta", mask))*e1 \
                - np.cos(_val("zeta", mask))*e2
            return bhat * _val("ppar", mask) \
                + perphat * np.sqrt(pnorm**2 - _val("ppar", mask)**2)
        vvecprt = lambda mask : physlib.velocity_momentum(mass[mask],
                                                          pvecprt(mask))
        vvecgc = lambda mask : physlib.velocity_momentum(mass[mask],
                                                         pvecgc(mask))
        # Common quantities
        add("ids", lambda : _val("ids"))
        add("r", lambda : _val("r"))
        add("z", lambda : _val("z"))
        add("phi", lambda : _val("phi"))
        add("phimod", lambda : np.mod(_val("phi"), 2 * np.pi * unyt.rad))
        add("theta", lambda : _val("theta"))
        add("thetamod", lambda : np.mod(_val("theta"), 2 * np.pi * unyt.rad))
        add("x", lambda : physlib.pol2cart(_val("r"), _val("phi"))[0])
        add("y", lambda : physlib.pol2cart(_val("r"), _val("phi"))[1])
        add("br", lambda : _val("br"))
        add("bz", lambda : _val("bz"))
        add("bphi", lambda : _val("bphi"))
        add("bnorm", lambda : np.sqrt( _val("br")**2 + _val("bphi")**2
                                       + _val("bz")**2 ))
        add("rho", lambda : _eval("rho"))
        add("psi", lambda : _eval("psi"))
        add("mileage", lambda : _val("mileage"))
        add("mass", lambda : mass)
        add("charge", lambda : _val("charge"))
        add("weight", lambda : _val("weight"))
        add("time", lambda : time)
        add("pncrid", lambda : _val("pncrid"))
        add("pncrdir", lambda : _val("pncdir"))
        add("connlen", lambda : totmil)

        firstmask = 0; lastmask=0
        if Orbits.GYROORBIT in mode:
            # Record the index of first masked array so that we can later append
            # all other arrays created here with GC mask (in hybrid mode)
            firstmask = len(items)
            mask = mode == Orbits.GYROORBIT
            add("pr", lambda : _val("pr", mask))
            add("pz", lambda : _val("pz", mask))
            add("pphi", lambda : _val("pphi", mask))
            add("ppar", lambda : physlib.ppar_momentum(pvecprt(mask),
                                                       bvec(mask)))
            if "pperp" in qnt:
                pnorm2 = np.sum( pvecprt(mask)**2, axis=0 )
                pperp2 = physlib.ppar_momentum(pvecprt(mask), bvec(mask))**2
                add("pperp", lambda : np.sqrt(pnorm2 - pperp2))
            add("pnorm", lambda : np.sqrt( np.sum( pvecprt(mask)**2, axis=0 ) ))
            add("vr", lambda : vvecprt(mask)[0,:])
            add("vz", lambda : vvecprt(mask)[2,:])
            add("vphi", lambda : vvecprt(mask)[1,:])
            add("vpar", lambda : physlib.vpar_momentum(
                mass[mask], pvecprt(mask), bvec(mask)))
            if "vperp" in qnt:
                vpar = physlib.vpar_momentum(
                    mass[mask], pvecprt(mask), bvec(mask))
                add("vperp", lambda : np.sqrt(np.sum(pvecprt(mask)**2, axis=0)
                                              - vpar**2))
            if "vnorm" in qnt:
                pnorm = np.sqrt(np.sum(pvecprt(mask)**2, axis=0))
                add("vnorm", lambda : physlib.velocity_momentum(mass[mask],
                                                                pnorm))
            add("ekin", lambda : physlib.energy_momentum(mass[mask],
                                                         pvecprt(mask)))
            add("pitch", lambda : physlib.pitch_momentum(pvecprt(mask),
                                                         bvec(mask)))
            add("mu", lambda : physlib.mu_momentum(mass[mask], pvecprt(mask),
                                                   bvec(mask)))
            add("zeta", lambda : mass[mask]*np.nan) # TODO implement properly
            add("ptor", lambda : physlib.torcanangmom_momentum(
                _val("charge", mask), _val("r", mask), pvecprt(mask),
                _eval("psi", mask)))
            lastmask = len(items)

        if Orbits.GUIDINGCENTER in mode:
            mask = mode == Orbits.GUIDINGCENTER
            add("pr", lambda : pvecgc(mask)[0,:])
            add("pz", lambda : pvecgc(mask)[2,:])
            add("pphi", lambda : pvecgc(mask)[1,:])
            add("ppar", lambda : _val("ppar", mask))
            if "pperp" in qnt:
                pnorm = physlib.momentum_muppar(
                    mass[mask], _val("mu", mask),
                    _val("ppar", mask), bvec(mask))
                add("pperp", lambda : np.sqrt(pnorm**2 - _val("ppar")**2))
            add("pnorm", lambda : physlib.momentum_muppar(
                mass[mask], _val("mu", mask),
                _val("ppar", mask), bvec(mask)))
            add("vr", lambda : vvecgc(mask)[0,:])
            add("vz", lambda : vvecgc(mask)[2,:])
            add("vphi", lambda : vvecgc(mask)[1,:])
            add("vpar", lambda : physlib.vpar_muppar(
                mass[mask], _val("mu", mask), _val("ppar", mask), bvec(mask)))
            if "vperp" in qnt:
                pnorm = physlib.momentum_muppar(
                    mass[mask], _val("mu", mask),
                    _val("ppar", mask), bvec(mask))
                pperp = np.sqrt(pnorm**2 - _val("ppar", mask)**2)
                gamma = physlib.gamma_momentum(mass[mask], pnorm)
                add("vperp", lambda : pperp / (gamma * mass[mask]))
            if "vnorm" in qnt:
                pnorm = physlib.momentum_muppar(
                    mass[mask], _val("mu", mask),
                    _val("ppar", mask), bvec(mask))
                gamma = physlib.gamma_momentum(mass[mask], pnorm)
                add("vnorm", lambda : pnorm / (gamma * mass[mask]))
            add("ekin", lambda : physlib.energy_muppar(
                mass[mask], _val("mu", mask), _val("ppar", mask), bvec(mask)))
            add("pitch", lambda : physlib.pitch_muppar(
                mass[mask], _val("mu", mask), _val("ppar", mask), bvec(mask)))
            add("mu", lambda : _val("mu", mask))
            add("zeta", lambda : _val("zeta", mask))
            add("ptor", lambda : physlib.torcanangmom_ppar(
                _val("charge", mask), _val("r", mask), _val("ppar", mask),
                bvec(mask), _eval("psi", mask)))

        if Orbits.GYROORBIT in mode and Orbits.GUIDINGCENTER in mode:
            # Hybrid data, combine masked arrays
            nmasked = lastmask - firstmask
            for i in range(lastmask):
                if i < firstmask: continue
                items[i] = np.append(items[i], items[i+nmasked])
            for i in range(nmasked):
                del items[lastmask+i]

        for i in range(len(items)):
            if items[i] is None:
                raise ValueError("Unknown quantity in " + qnt[i])

        # Sort first by IDs and then by mileage
        ids  = _val("ids").v
        mile = _val("mileage").v
        idx  = np.lexsort((mile, ids))
        for i in range(len(items)):
            items[i] = items[i][idx]
            items[i].convert_to_base("ascot")
        return items

    @staticmethod
    def listqnts():
        """List all available quantities.
        """
        out = {
            "r":        "R coordinate",
            "z":        "z coordinate",
            "phi":      "Toroidal coordinate (cumulative)",
            "phimod":   "Toroidal coordinate",
            "theta":    "Poloidal coordinate (cumulative)",
            "thetamod": "Poloidal coordinate",
            "x":        "x coordinate",
            "y":        "y coordinate",
            "pr":       "Momentum R component",
            "pz":       "Momentum z component",
            "pphi":     "Momentum phi component",
            "ppar":     "Momentum component parallel to B",
            "pperp":    "Momentum component perpendicular to B",
            "pnorm":    "Momentum norm",
            "vr":       "Velocity R component",
            "vz":       "Velocity z component",
            "vphi":     "Velocity phi component",
            "vpar":     "Velocity component parallel to B",
            "vperp":    "Velocity component perpendicular to B",
            "vnorm":    "Velocity norm",
            "br":       "Magnetic field R component at the marker position",
            "bz":       "Magnetic field z component at the marker position",
            "bphi":     "Magnetic field phi component at the marker position",
            "bnorm":    "Magnetic field strength at the marker position",
            "ekin":     "Kinetic energy",
            "pitch":    "vpa / vnorm",
            "mu":       "Magnetic moment",
            "zeta":     "Gyroangle",
            "psi":      "Poloidal flux at the marker position",
            "rho":      "Square root of normalized psi",
            "ptor":     "Canonical toroidal angular momentum",
            "mass":     "Mass",
            "charge":   "Charge",
            "time":     "Current laboratory time",
            "mileage":  "Laboratory time elapsed in simulation",
            "weight":   "How many physical particles a marker represents",
            "ids":      "Marker ID",
            "connlen":  "Connection length for lost markers",
            "pncrid":   "Poincaré plane this point corresponds to",
            "pncrdir":  "Direction at which Poincaré plane was crossed",
        }
        return out

    @staticmethod
    def write_hdf5(fn, data, desc=None):
        """Write state data in HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        desc : str, optional
            Input description.
        """

        gname = "results/" + run + "/orbit"

        N = data["id"].size

        with h5py.File(fn, "a") as f:
            g = f.create_group(gname)

            g.create_dataset("id",   (N,1), data=data["id"],   dtype="i8")
            g.create_dataset("time", (N,1), data=data["time"], dtype="f8")

            if "pncrid" in data:
                g.create_dataset("pncrid", (N,1), data=data["pncrid"], dtype="i4")

            if "charge" in data:
                g.create_dataset("charge", (N,1), data=data["charge"], dtype="i4")

            if "weight" in data:
                g.create_dataset("weight", (N,1), data=data["weight"], dtype="f8")

            if "pr" in data:
                g.create_dataset("r",     (N,1), data=data["r"],     dtype="f8")
                g.create_dataset("phi",   (N,1), data=data["phi"],   dtype="f8")
                g.create_dataset("z",     (N,1), data=data["z"],     dtype="f8")

                g.create_dataset("pr",    (N,1), data=data["pr"],    dtype="f8")
                g.create_dataset("pphi",  (N,1), data=data["pphi"],  dtype="f8")
                g.create_dataset("pz",    (N,1), data=data["pz"],    dtype="f8")

                g.create_dataset("rho",   (N,1), data=data["rho"],   dtype="f8")
                g.create_dataset("theta", (N,1), data=data["theta"], dtype="f8")
                g.create_dataset("br",    (N,1), data=data["br"],    dtype="f8")
                g.create_dataset("bphi",  (N,1), data=data["bphi"],  dtype="f8")
                g.create_dataset("bz",    (N,1), data=data["bz"],    dtype="f8")
            else:
                g.create_dataset("r",     (N,1), data=data["r"],     dtype="f8")
                g.create_dataset("phi",   (N,1), data=data["phi"],   dtype="f8")
                g.create_dataset("z",     (N,1), data=data["z"],     dtype="f8")

                g.create_dataset("rho",   (N,1), data=data["rho"],   dtype="f8")
                g.create_dataset("theta", (N,1), data=data["theta"], dtype="f8")
                g.create_dataset("br",    (N,1), data=data["br"],    dtype="f8")
                g.create_dataset("bphi",  (N,1), data=data["bphi"],  dtype="f8")
                g.create_dataset("bz",    (N,1), data=data["bz"],    dtype="f8")

            if "ppar" in data:
                g.create_dataset("ppar",  (N,1), data=data["ppar"],  dtype="f8")
                g.create_dataset("mu",    (N,1), data=data["mu"],    dtype="f8")
                g.create_dataset("zeta",  (N,1), data=data["zeta"],  dtype="f8")

a5py/a5py/ascot5io/orbits.py
ascot4fusion/ascot5
"""Input data representing plasma background species.

Plasma input is required for simulations with collisions or collisional
diagnostics enabled. This includes BBNBI5 and ASCOT-BMC.
"""
import h5py
import numpy as np

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

class plasma_1D(DataGroup):
    """Plasma profiles that have only radial dependency.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key in ["nion", "nrho"]:
                    out[key] = int(out[key])

        out["idensity"] = np.transpose(out["idensity"])
        return out

    def plot(self, pls=None):
        import matplotlib.pyplot as plt

        if pls is None:
            pls = self.read()

        plotyy=False # This didn't come out nice, but keep it just in case.

        if plotyy:
            fig, ax1 = plt.subplots()

            color = 'tab:red'
            ax1.set_xlabel('$\\rho $')
            ax1.set_ylabel('Density (10$^{19}$/m$^3$)', color=color)
            ax1.plot(pls['rho'],pls['edensity']*1e-19,'-' ,color=color,label='n$_e$')
            ax1.plot(pls['rho'],pls['idensity']*1e-19,'--',label='n$_i$')
            ax1.tick_params(axis='y', labelcolor=color)

            ax2 = ax1.twinx()  # instantiate a second axes that shares the same x-axis
            color = 'tab:blue'
            ax2.set_ylabel('Temperature(eV)', color=color)  # we already handled the x-label with ax1
            ax2.plot(pls['rho'],pls['itemperature'],'-' ,color=color,label='T$_i$')
            ax2.plot(pls['rho'],pls['etemperature'],'--',label='T$_e$')
            ax2.tick_params(axis='y', labelcolor=color)

            ax1.legend()
            ax2.legend()

            fig.tight_layout()  # otherwise the right y-label is slightly clipped

            plt.show()

        else:

            plt.subplot(2,1,1)
            plt.plot(pls['rho'],pls['edensity']*1e-19,'-',label='n$_e$')
            plt.plot(pls['rho'],pls['idensity']*1e-19,'--',label='n$_i$')
            plt.legend()
            plt.ylabel('Density (10$^{19}$/m$^3$)')

            plt.subplot(2,1,2)
            plt.plot(pls['rho'],pls['etemperature'],'-',label='T$_e$')
            plt.plot(pls['rho'],pls['itemperature'],'--',label='T$_i$')
            plt.legend()
            plt.ylabel('Temperature (eV)')

            plt.xlabel('$\\rho $')

            plt.show()

    @staticmethod
    def write_hdf5(fn, nrho, nion, anum, znum, mass, charge, rho,
                   edensity, etemperature, idensity, itemperature, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Path to hdf5 file.
        nrho : int
            Number of rho grid points.
        nion : int
            Number of ion species.
        anum : array_like (nion,1)
            Ion species atomic mass number
        znum : array_like (nion,1)
            Ion species charge number.
        mass : array_like (nion,1)
            Ion species mass [amu].
        charge : array_like (nion,1)
            Ion species charge [e].
        rho : array_like (nrho,1)
            rho grid, doesn't have to be uniform.
        edensity : array_like (nrho,1)
            Electron density [m^-3].
        etemperature : array_like (nrho,1)
            Electron temperature [eV].
        idensity : array_like (nrho,nion)
            Ion density [m^-3].
        itemperature : array_like (nrho,1)
            Ion temperature [ev].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if etemperature.size != nrho:
            raise ValueError("Invalid size for electron temperature.")
        if itemperature.size != nrho:
            raise ValueError("Invalid size for ion temperature.")
        if edensity.size != nrho:
            raise ValueError("Invalid size for electron density.")
        if idensity.shape != (nrho,nion):
            raise ValueError("Invalid size for ion density.")

        idensity = np.transpose(idensity)

        parent = "plasma"
        group  = "plasma_1D"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset('nion',   (1,1),    data=nion,   dtype='i4')
            g.create_dataset('nrho',   (1,1),    data=nrho,   dtype='i4')
            g.create_dataset('znum',   (nion,1), data=znum,   dtype='i4')
            g.create_dataset('anum',   (nion,1), data=anum,   dtype='i4')
            g.create_dataset('charge', (nion,1), data=charge, dtype='i4')
            g.create_dataset('mass',   (nion,1), data=mass,   dtype='f8')
            g.create_dataset('rho',    (nrho,1), data=rho,    dtype='f8')

            g.create_dataset('etemperature', (nrho,1),    data=etemperature,
                             dtype='f8')
            g.create_dataset('edensity',     (nrho,1),    data=edensity,
                             dtype='f8')
            g.create_dataset('itemperature', (nrho,1),    data=itemperature,
                             dtype='f8')
            g.create_dataset('idensity',     (nion,nrho), data=idensity,
                             dtype='f8')

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output is an uniform hydrogen plasma.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return plasma_1D.write_hdf5(
            fn=fn, nrho=3, nion=1, znum=np.array([1]), anum=np.array([1]),
            mass=np.array([1]), charge=np.array([1]),
            rho=np.array([0, 0.5, 100]), edensity=1e20*np.ones((3,1)),
            etemperature=1e3*np.ones((3,1)), idensity=1e20*np.ones((3,1)),
            itemperature=1e20*np.ones((3,1)), desc="DUMMY")

class plasma_1DS(DataGroup):
    """Same input as :class:`plasma_1D` but interpolated with splines.

    Spline interpolation could yield negative value even though the inputs
    were all positive. To avoid this, ASCOT5 internally takes a logarithm
    of the input data before constructing the splines. The actual values
    are then computed after the spline-evaluation. The logarithmic also
    helps to interpolate the values at the edge more accurately.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key in ["nion", "nrho"]:
                    out[key] = int(out[key])

        out["idensity"] = np.transpose(out["idensity"])
        return out

    @staticmethod
    def write_hdf5(fn, nrho, nion, anum, znum, mass, charge, rhomin, rhomax,
                   edensity, etemperature, idensity, itemperature, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Path to hdf5 file.
        nrho : int
            Number of rho grid points.
        nion : int
            Number of ion species.
        anum : array_like (nion,1)
            Ion species atomic mass number
        znum : array_like (nion,1)
            Ion species charge number.
        mass : array_like (nion,1)
            Ion species mass [amu].
        charge : array_like (nion,1)
            Ion species charge [e].
        rhomin : float
            Minimum rho grid value.
        rhomax : float
            Maximum rho grid value.
        edensity : array_like (nrho,1)
            Electron density [m^-3].
        etemperature : array_like (nrho,1)
            Electron temperature [eV].
        idensity : array_like (nrho,nion)
            Ion density [m^-3].
        itemperature : array_like (nrho,1)
            Ion temperature [ev].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if etemperature.size != nrho:
            raise ValueError("Invalid size for electron temperature.")
        if itemperature.size != nrho:
            raise ValueError("Invalid size for ion temperature.")
        if edensity.size != nrho:
            raise ValueError("Invalid size for electron density.")
        if idensity.shape != (nrho,nion):
            raise ValueError("Invalid size for ion density.")

        idensity = np.transpose(idensity)

        parent = "plasma"
        group  = "plasma_1DS"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset('nion',   (1,1),    data=nion,   dtype='i4')
            g.create_dataset('nrho',   (1,1),    data=nrho,   dtype='i4')
            g.create_dataset('znum',   (nion,1), data=znum,   dtype='i4')
            g.create_dataset('anum',   (nion,1), data=anum,   dtype='i4')
            g.create_dataset('charge', (nion,1), data=charge, dtype='i4')
            g.create_dataset('mass',   (nion,1), data=mass,   dtype='f8')
            g.create_dataset('rhomin', (1,1),    data=rhomin, dtype='f8')
            g.create_dataset('rhomax', (1,1),    data=rhomax, dtype='f8')

            g.create_dataset('etemperature', (nrho,1),    data=etemperature,
                             dtype='f8')
            g.create_dataset('edensity',     (nrho,1),    data=edensity,
                             dtype='f8')
            g.create_dataset('itemperature', (nrho,1),    data=itemperature,
                             dtype='f8')
            g.create_dataset('idensity',     (nion,nrho), data=idensity,
                             dtype='f8')

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output is an uniform hydrogen plasma.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return plasma_1DS.write_hdf5(
            fn=fn, nrho=3, nion=1, znum=np.array([1]), anum=np.array([1]),
            mass=np.array([1]), charge=np.array([1]),
            rhomin=0, rhomax=100, edensity=1e20*np.ones((3,1)),
            etemperature=1e3*np.ones((3,1)), idensity=1e20*np.ones((3,1)),
            itemperature=1e20*np.ones((3,1)), desc="DUMMY")

class plasma_1Dt(DataGroup):
    """Time-dependent 1D plasma profiles.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key in ["nion", "nrho", "ntime"]:
                    out[key] = int(out[key])

        return out

    @staticmethod
    def write_hdf5(fn, nrho, ntime, nion, anum, znum, mass, charge, rho, time,
                   edensity, etemperature, idensity, itemperature, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Path to hdf5 file.
        nrho : int
            Number of rho grid points.
        ntime : int
            Number of time grid points.
        nion : int
            Number of ion species.
        anum : array_like (nion,1)
            Ion species atomic mass number
        znum : array_like (nion,1)
            Ion species charge number.
        mass : array_like (nion,1)
            Ion species mass [amu].
        charge : array_like (nion,1)
            Ion species charge [e].
        rho : array_like (nrho,1)
            rho grid, doesn't have to be uniform.
        time : array_like (nrho,1)
            time grid, doesn't have to be uniform.
        edensity : array_like (ntime,nrho)
            Electron density [m^-3].
        etemperature : array_like (ntime,nrho)
            Electron temperature [eV].
        idensity : array_like (ntime,nion,nrho)
            Ion density [m^-3].
        itemperature : array_like (ntime,nrho)
            Ion temperature [ev].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If inputs were not consistent.
        """
        if etemperature.shape != (ntime,nrho):
            raise ValueError("Invalid shape for electron temperature.")
        if itemperature.shape != (ntime,nrho):
            raise ValueError("Invalid shape for ion temperature.")
        if edensity.shape != (ntime,nrho):
            raise ValueError("Invalid shape for electron density.")
        if idensity.shape != (ntime,nion,nrho):
            raise ValueError("Invalid shape for ion density.")

        parent = "plasma"
        group  = "plasma_1Dt"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset('nion',   (1,1),     data=nion,   dtype='i4')
            g.create_dataset('nrho',   (1,1),     data=nrho,   dtype='i4')
            g.create_dataset('ntime',  (1,1),     data=ntime,  dtype='i4')
            g.create_dataset('znum',   (nion,1),  data=znum,   dtype='i4')
            g.create_dataset('anum',   (nion,1),  data=anum,   dtype='i4')
            g.create_dataset('charge', (nion,1),  data=charge, dtype='i4')
            g.create_dataset('mass',   (nion,1),  data=mass,   dtype='f8')
            g.create_dataset('rho',    (nrho,1),  data=rho,    dtype='f8')
            g.create_dataset('time',   (ntime,1), data=time,   dtype='f8')

            g.create_dataset('etemperature', (ntime,nrho),
                             data=etemperature, dtype='f8')
            g.create_dataset('edensity',     (ntime,nrho),
                             data=edensity, dtype='f8')
            g.create_dataset('itemperature', (ntime,nrho),
                             data=itemperature, dtype='f8')
            g.create_dataset('idensity',     (ntime,nion,nrho),
                             data=idensity, dtype='f8')

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output is an uniform hydrogen plasma.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return plasma_1Dt.write_hdf5(
            fn=fn, nrho=3, ntime=4, nion=1, znum=np.array([1]),
            anum=np.array([1]), mass=np.array([1]), charge=np.array([1]),
            rho=np.array([0, 0.5, 100]), time=np.array([0, 0.2, 0.4, 0.6]),
            edensity=1e20*np.ones((4,3)), etemperature=1e3*np.ones((4,3)),
            idensity=1e20*np.ones((4,1,3)), itemperature=1e20*np.ones((4,3)),
            desc="DUMMY")

a5py/a5py/ascot5io/plasma.py
ascot4fusion/ascot5
"""Marker initial and final phase-space positions and related quantities.
"""
import numpy as np
import h5py
import unyt
import a5py.physlib as physlib

from .coreio import fileapi
from .coreio.treedata import DataContainer

class State(DataContainer):
    """Marker initial and final phase-space positions and related quantities.
    """

    # These end conditions are used internally so we define them as class
    # variables. All end conditions are defined as properties so that they
    # can have docstrings.
    _NONE    = 0x1
    _ABORTED = 0x2
    _TLIM    = 0x4
    _EMIN    = 0x8
    _THERM   = 0x10
    _WALL    = 0x20
    _RHOMIN  = 0x40
    _RHOMAX  = 0x80
    _POLMAX  = 0x100
    _TORMAX  = 0x200
    _CPUMAX  = 0x400
    _NEUTR   = 0x800
    _IONIZ   = 0x1000

    @property
    def ABORTED(self):
        """Marker simulation terminated in an error.
        """
        return State._ABORTED

    @property
    def NONE(self):
        """No active end condition meaning the marker hasn't been simulated yet.
        """
        return State._NONE

    @property
    def TLIM(self):
        """Simulation time limit or maximum mileage reached.
        """
        return State._TLIM

    @property
    def EMIN(self):
        """Minimum energy reached.
        """
        return State._EMIN

    @property
    def THERM(self):
        """Local thermal energy reached.
        """
        return State._THERM

    @property
    def WALL():
        """Marker intersected a wall element.
        """
        return State._WALL

    @property
    def RHOMIN(self):
        """Minimum radial coordinate (rho) reached.
        """
        return State._RHOMIN

    @property
    def RHOMAX(self):
        """Maximum radial coordinate (rho) reached.
        """
        return _RHOMAX

    @property
    def POLMAX(self):
        """Maximum poloidal turns reached.
        """
        return State._POLMAX

    @property
    def TORMAX(self):
        """Maximum toroidal turns reached.
        """
        return State._TORMAX

    @property
    def CPUMAX(self):
        """Simulation for this marker exceeded the set CPU time.
        """
        return State._CPUMAX

    @property
    def NEUTR(self):
        """Ion marker neutralized.
        """
        return State._NEUTR

    @property
    def IONIZ(self):
        """Neutral marker ionized.
        """
        return State._IONIZ

    def write_hdf5(self):
        """Write state data in HDF5 file.

        Parameters
        ----------
        fn : str
        Full path to the HDF5 file.
        """

        gname = "results/" + run + "/" + name

        fields = [("mass",      "amu"),
                  ("time",      "s"),
                  ("cputime",   "s"),
                  ("weight",    "markers/s"),
                  ("rprt",      "m"),
                  ("phiprt",    "deg"),
                  ("zprt",      "m"),
                  ("prprt",     "kg*m/s"),
                  ("pphiprt",   "kg*m/s"),
                  ("pzprt",     "kg*m/s"),
                  ("r",         "m"),
                  ("phi",       "deg"),
                  ("z",         "m"),
                  ("mu",        "eV/T"),
                  ("ppar",      "kg*m/s"),
                  ("zeta",      "rad"),
                  ("rho",       "1"),
                  ("theta",     "deg"),
                  ("ids",       "1"),
                  ("walltile",  "1"),
                  ("endcond",   "1"),
                  ("anum",      "1"),
                  ("znum",      "1"),
                  ("charge",    "e"),
                  ("errormsg",  "1"),
                  ("errorline", "1"),
                  ("errormod",  "1"),
                  ("br",        "T"),
                  ("bphi",      "T"),
                  ("bz",        "T")]

        N = data["ids"].size

        with h5py.File(fn, "a") as f:
            g = f.create_group(gname)

            for field in fields:
                ds = g.create_dataset(field[0], (N,1), data=data[field[0]],
                                      dtype="f8")
                ds.attrs["unit"] = field[1]

    def read(self):
        """Read raw state data to a dictionary.
        """
        out = {}
        with self as f:
            for key in f:
                out[key] = f[key][:]

        return out

    def get(self, *qnt, mode="gc"):
        """Return marker quantity.

        This function accesses the state data within the HDF5 file and uses that
        to evaluate the queried quantity.

        Parameters
        ----------
        qnt : str
            Name of the quantities.
        mode : {"prt", "gc"}, optional
            Is the quantity evaluated in particle or guiding center phase-space.

        Returns
        -------
        value : array_like
            The quantities as an array ordered by marker ID.
        """
        def _val(q):
            """Read quantity from HDF5.
            """
            with self as h5:
                if q in h5:
                    return fileapi.read_data(h5, q)
            return None

        def _eval(r, phi, z, t, *q):
            """Evaluate quantity with ascotpy.
            """
            return self._root._ascot.input_eval(r, phi, z, t, *q)

        return State._getactual(mode, _val, _eval, *qnt)

    @staticmethod
    def _getactual(mode, _val, _eval, *qnt):
        """Calculate marker quantities using the helper functions and data.

        Parameters
        ----------
        mode : {"prt", "gc"}
            Phase-space where quantity is evaluated.
        _val : callable
            Function that returns a stored marker parameter.

            ``_val(qnt : str, mask : array_like) -> value``
        _eval : callable
            Function that returns interpolated input quantity at masked
            positions at given position.

            ``_eval(r, phi, z, qnt : str) -> value``
        *qnt : str
            Names of the quantities.

        Returns
        -------
        *value : array_like
            The quantities as an array ordered by marker ID.
        """
        items = [None]*len(qnt)
        def add(q, val):
            if q in qnt:
                items[qnt.index(q)] = val()

        evalprt = lambda *q : _eval(_val("rprt"), _val("phiprt"), _val("zprt"),
                                    _val("time"), *q)
        evalgc  = lambda *q : _eval(_val("r"), _val("phi"), _val("z"),
                                    _val("time"), *q)
        bvecprt = lambda : unyt.unyt_array(evalprt("br", "bphi", "bz"))
        bvecgc = lambda : unyt.unyt_array([_val("br"), _val("bphi"),
                                _val("bz")])
        pvecprt = lambda : unyt.unyt_array([_val("prprt"), _val("pphiprt"),
                                _val("pzprt")])
        def pvecgc():
            """Evaluate guiding center momentum vector.
            """
            bnorm = np.sqrt(np.sum(bvecgc()**2,axis=0))
            pnorm = physlib.momentum_muppar(_val("mass"), _val("mu"),
                                            _val("ppar"), bnorm)
            bhat  = bvecgc() / bnorm
            e1 = np.zeros(bhat.shape)
            e1[2,:] = 1
            e2 = np.cross(bhat.T, e1.T).T
            e1 = e2 / np.sqrt(np.sum(e2**2, axis=0))
            e2 = np.cross(bhat.T, e1.T).T
            perphat = -np.sin(_val("zeta"))*e1-np.cos(_val("zeta"))*e2
            return bhat * _val("ppar") \
                + perphat * np.sqrt(pnorm**2 - _val("ppar")**2)
        vvecprt = lambda : physlib.velocity_momentum(_val("mass"), pvecprt())
        vvecgc = lambda : physlib.velocity_momentum(_val("mass"), pvecgc())

        # Dimensionless quantities common for particle and guiding-center
        add("ids", lambda : _val("ids").v)
        add("anum", lambda : _val("anum").v)
        add("znum", lambda : _val("znum").v)
        add("walltile", lambda : _val("walltile").v)
        add("errormsg", lambda : _val("errormsg").v)
        add("errorline", lambda : _val("errorline").v)
        add("errormod", lambda : _val("errormod").v)
        if "endcond" in qnt:
            item = _val("endcond").v
            err = _val("errormsg").v
            item = item << 2
            item[err > 0] = item[err > 0] & State._ABORTED
            item[item==0] = State._NONE
            add("endcond", lambda : item)

        # Dimensional quantities common for particle and guiding-center
        add("mass", lambda : _val("mass"))
        add("charge", lambda : _val("charge"))
        add("time", lambda : _val("time"))
        add("cputime", lambda : _val("cputime"))
        add("mileage", lambda : _val("mileage"))
        add("weight", lambda : _val("weight"))

        # Quantities that have to be evaluated separately (keep this in same
        # order for both and in the list method if new quantities are added)
        if mode == "prt":
            add("r", lambda : _val("rprt"))
            add("z", lambda : _val("zprt"))
            add("phi", lambda : _val("phiprt"))
            add("phimod", lambda : np.mod(_val("phiprt"), 2 * np.pi * unyt.rad))
            if "theta" in qnt:
                axis = evalprt("axisr", "axisz")
                thetaprt = physlib.cart2pol(
                    _val("rprt")-axis[0], _val("zprt")-axis[1])[1]
                thetacum = np.floor(_val("phi") / (2*np.pi * unyt.rad))
                theta = 2*np.pi * thetacum * unyt.rad + thetaprt
                add("theta", lambda : np.mod(theta + np.pi*unyt.rad,
                                             2*np.pi*unyt.rad))
            if "thetamod" in qnt:
                axis = evalprt("axisr", "axisz")
                theta = physlib.cart2pol(
                    _val("rprt")-axis[0], _val("zprt")-axis[1])[1]
                add("thetamod", lambda : np.mod(theta + np.pi*unyt.rad,
                                                2*np.pi*unyt.rad))
            add("x", lambda : physlib.pol2cart(_val("rprt"), _val("phiprt"))[0])
            add("y", lambda : physlib.pol2cart(_val("rprt"), _val("phiprt"))[1])
            add("pr", lambda : _val("prprt"))
            add("pz", lambda : _val("pzprt"))
            add("pphi", lambda : _val("pphiprt"))
            add("ppar", lambda : physlib.ppar_momentum(pvecprt(), bvecprt()))
            if "pperp" in qnt:
                pnorm2 = np.sum( pvecprt()**2, axis=0 )
                ppar2  = physlib.ppar_momentum(pvecprt(), bvecprt())**2
                add("pperp", lambda : np.sqrt(pnorm2 - ppar2))
            if "pnorm" in qnt:
                add("pnorm", lambda : np.sqrt( np.sum( pvecprt()**2, axis=0 ) ))
            add("vr", lambda : vvecprt()[0,:])
            add("vz", lambda : vvecprt()[2,:])
            add("vphi", lambda : vvecprt()[1,:])
            add("vpar", lambda : physlib.vpar_momentum(_val("mass"), pvecprt(),
                                              bvecprt()))
            if "vperp" in qnt:
                vperp = physlib.vpar_momentum(
                    _val("mass"), pvecprt(), bvecprt())
                add("vperp", lambda : np.sum(pvecprt()**2, axis=0))
            if "vnorm" in qnt:
                pnorm = np.sqrt(np.sum(pvecprt()**2, axis=0))
                add("vnorm", lambda : physlib.velocity_momentum(_val("mass"),
                                                                pnorm))
            add("br", lambda : evalprt("br"))
            add("bz", lambda : evalprt("bz"))
            add("bphi", lambda : evalprt("bphi"))
            add("bnorm", lambda : evalprt("bnorm"))
            add("ekin", lambda : physlib.energy_momentum(_val("mass"),
                                                         pvecprt()))
            add("pitch", lambda : physlib.pitch_momentum(pvecprt(), bvecprt()))
            add("mu", lambda : physlib.mu_momentum(_val("mass"), pvecprt(),
                                          bvecprt()))
            add("zeta", lambda : _val("zeta"))
            add("psi", lambda : evalprt("psi"))
            add("rho", lambda : evalprt("rho"))
            add("ptor", lambda : physlib.torcanangmom_momentum(
                _val("charge"), _val("rprt"), pvecprt(), evalprt("psi")))

        if mode == "gc":
            add("r", lambda : _val("r"))
            add("z", lambda : _val("z"))
            add("phi", lambda : _val("phi"))
            add("phimod", lambda : np.mod(_val("phi"), 2*np.pi * unyt.rad))
            add("theta", lambda : _val("theta"))
            add("thetamod", lambda : np.mod(_val("theta"), 2*np.pi * unyt.rad))
            add("x", lambda : physlib.pol2cart(_val("r"), _val("phi"))[0])
            add("y", lambda : physlib.pol2cart(_val("r"), _val("phi"))[1])
            add("pr", lambda : pvecgc()[0,:])
            add("pz", lambda : pvecgc()[2,:])
            add("pphi", lambda : pvecgc()[1,:])
            add("ppar", lambda : _val("ppar"))
            if "pperp" in qnt:
                pnorm = physlib.momentum_muppar(_val("mass"), _val("mu"),
                                                _val("ppar"), bvecgc())
                add("pperp", lambda : np.sqrt(pnorm**2 - _val("ppar")**2))
            add("pnorm",lambda :  physlib.momentum_muppar(
                _val("mass"), _val("mu"), _val("ppar"), bvecgc()))
            add("vr", lambda : vvecgc()[0,:])
            add("vz", lambda : vvecgc()[2,:])
            add("vphi", lambda : vvecgc()[1,:])
            add("vpar", lambda : physlib.vpar_muppar(_val("mass"), _val("mu"),
                                                     _val("ppar"), bvecgc()))
            if "vperp" in qnt:
                pnorm = physlib.momentum_muppar(_val("mass"), _val("mu"),
                                                _val("ppar"), bvecgc())
                pperp = np.sqrt(pnorm**2 - _val("ppar")**2)
                gamma = physlib.gamma_momentum(_val("mass"), pnorm)
                add("vperp", lambda : pperp / (gamma * _val("mass")))
            if "vnorm" in qnt:
                pnorm = physlib.momentum_muppar(_val("mass"), _val("mu"),
                                                _val("ppar"), bvecgc())
                gamma = physlib.gamma_momentum(_val("mass"), pnorm)
                add("vnorm", lambda : pnorm / (gamma * _val("mass")))
            add("br", lambda : _val("br"))
            add("bz", lambda : _val("bz"))
            add("bphi", lambda : _val("bphi"))
            add("bnorm", lambda : np.sqrt(_val("br")**2 + _val("bphi")**2 +
                                 _val("bz")**2))
            add("ekin", lambda : physlib.energy_muppar(_val("mass"), _val("mu"),
                                              _val("ppar"), bvecgc()))
            add("pitch", lambda : physlib.pitch_muppar(_val("mass"), _val("mu"),
                                              _val("ppar"), bvecgc()))
            add("mu", lambda : _val("mu"))
            add("zeta", lambda : _val("zeta"))
            add("psi", lambda : evalgc("psi"))
            add("rho", lambda : _val("rho"))
            add("ptor", lambda : physlib.torcanangmom_ppar(
                _val("charge"), _val("r"), _val("ppar"),
                bvecgc(), evalgc("psi")))

        for i in range(len(items)):
            if items[i] is None:
                raise ValueError("Unknown quantity in " + qnt[i])

        # Order by ID
        idx = _val("ids").argsort()
        for i in range(len(items)):
            items[i] = items[i][idx]
            try:
                items[i].convert_to_base("ascot")
            except AttributeError:
                # Non-dimensional
                pass
        return items

    @staticmethod
    def endcond_check(bitarr, string):
        """Check if the binary end condition matches the human-readable.

        Parameters
        ----------
        bitarr : :obj:`np.uint32`
            Value of the end condition as it is stored in the HDF5 file.
        string : str
            Human-readable end condition (case-insensitive).

            Markers may have multiple end conditions active simultaneously.
            If just the name of the end condition e.g. "MAXPOL" is passed,
            then all markers with the ``MAXPOL`` end condition are returned.

            If the end cond is preceded by "NOT", e.g. "NOT MAXPOL", then
            markers that don't have that end condition are returned.

            Passing multiple end conditions returns markers that have all listed
            end conditions active, e.g. "MAXPOL MAXTOR" returns markers that
            have both ``MAXPOL`` and ``MAXTOR`` active simultaneously.

        Returns
        -------
        match : bool
            True if the two representations of end conditions match.
        """
        endconds = string.upper().split()
        ec_yes = np.array(0, dtype=np.uint32)
        ec_non = np.array(0, dtype=np.uint32)

        i = 0
        while i < len(endconds):
            ec = endconds[i]

            NOT = False
            if ec == "NOT":
                NOT = True
                i += 1
                ec = endconds[i]

            try:
                ec = getattr(State, "_" + ec)
            except AttributeError:
                raise ValueError("Unknown end condition: " + ec)

            i += 1
            if NOT:
                ec_non = ec_non | ec
            else:
                ec_yes = ec_yes | ec

        return (bitarr & ec_yes) == ec_yes and (bitarr & ec_non) == 0

    @staticmethod
    def endcond_tostring(bitarr):
        """Convert end condition bitarray to human readable.

        Parameters
        ----------
        bitarr : :obj:`np.uint32`
            Value of the end condition as it is stored in the HDF5 file.

        Returns
        -------
        string : str
            End condition in a human readable format.
        """
        endcond = ["NONE", "ABORTED", "TLIM", "EMIN", "THERM", "WALL", "RHOMIN",
                   "RHOMAX", "POLMAX", "TORMAX", "CPUMAX"]
        string = ""
        for ec in endcond:
            if bitarr & getattr(State, "_" + ec):
                string += ec + " and "
        return string[:-5] # remove last "and"

    @staticmethod
    def listqnts():
        """List available quantities.
        """
        out = {
            "ids":      "Marker ID",
            "anum":     "Mass number",
            "znum":     "Charge number",
            "walltile": "Element ID if the marker hit wall",
            "endcond":  "Bitarray showing active end conditions",
            "errormsg": "Error message if marker was aborted",
            "errorline":"Line where (if) marker was aborted",
            "errormod": "Name of the module where (if) marker was aborted",
            "mass":     "Mass",
            "charge":   "Charge",
            "time":     "Current laboratory time",
            "cputime":  "CPU time elapsed in simulation",
            "mileage":  "Laboratory time elapsed in simulation",
            "weight":   "How many physical particles a marker represents",
            "r":        "R coordinate",
            "z":        "z coordinate",
            "phi":      "Toroidal coordinate (cumulative)",
            "phimod":   "Toroidal coordinate",
            "theta":    "Poloidal coordinate (cumulative)",
            "thetamod": "Poloidal coordinate",
            "x":        "x coordinate",
            "y":        "y coordinate",
            "pr":       "Momentum R component",
            "pz":       "Momentum z component",
            "pphi":     "Momentum phi component",
            "ppar":     "Momentum component parallel to B",
            "pperp":    "Momentum component perpendicular to B",
            "pnorm":    "Momentum norm",
            "vr":       "Velocity R component",
            "vz":       "Velocity z component",
            "vphi":     "Velocity phi component",
            "vpar":     "Velocity component parallel to B",
            "vperp":    "Velocity component perpendicular to B",
            "vnorm":    "Velocity norm",
            "br":       "Magnetic field R component at the marker position",
            "bz":       "Magnetic field z component at the marker position",
            "bphi":     "Magnetic field phi component at the marker position",
            "bnorm":    "Magnetic field strength at the marker position",
            "ekin":     "Kinetic energy",
            "pitch":    "vpa / vnorm",
            "mu":       "Magnetic moment",
            "zeta":     "Gyroangle",
            "psi":      "Poloidal flux at the marker position",
            "rho":      "Square root of normalized psi",
            "ptor":     "Canonical toroidal angular momentum",
        }
        return out

a5py/a5py/ascot5io/state.py
ascot4fusion/ascot5
"""
Transport coefficient data IO module.

File: transcoef.py
"""
import numpy as np
import h5py

from .coreio.treedata import DataContainer
from a5py.marker.losstime import eval_coefficients

def read_hdf5(fn, qid):
    """
    Read transport coefficient data from HDF5 file to dictionary.

    Args:
        fn : str <br>
            Full path to HDF5 file.
        qid : str <br>
            QID of the run whose transport coefficient data is read.

    Returns:
        Dictionary storing the coefficients and ids (unsorted).
    """

    with h5py.File(fn,"r") as f:
        data = f["/results/run_"+qid+"/transcoef"]

        out = {}

        # Read data from file.
        for field in data:
            out[field] = data[field][:]

    return out


class Transcoef(DataContainer):
    """
    Object representing transport coefficient data.
    """

    def __init__(self, root, hdf5):
        """
        Initialize transcoef object from given HDF5 file to given RunNode.
        """
        super().__init__(root, hdf5)


    def read(self):
        """
        Read coefficient data to dictionary.
        """
        return read_hdf5(self._file, self.get_qid())


    def __getitem__(self, key):
        with self as h5:
            h5keys = list(h5.keys())
            item = None
            if key in h5keys:
                item = h5[key][:]

            # Sort by id
            idx  = h5["ids"][:].argsort()
            return item[idx]


    def get_matrix(self, dim=(1,), lossfrac=1.1, method="radindep",
                   userho=0):
        """
        Sort the data into a histogram and return the mean values in each bin.

        This function assumes markers were sorted during the creation.
        """
        ids = self["ids"]
        k   = self["k"]
        d   = self["d"]

        ids0 = self._runnode.inistate["id"]
        rho  = self._runnode.inistate["r"]
        if userho:
            rho = self._runnode.inistate["rho"]
        energy = self._runnode.inistate["energy"]
        pitch  = self._runnode.inistate["pitch"]
        ec     = self._runnode.endstate["endcond"]

        lost   = np.logical_or.reduce([ec==32, ec==128, ec==1])

        nsize = 1
        for i in np.arange( len(dim) ):
            nsize *= dim[i]

        ntot = self._runnode.inistate["ids"].size
        if ntot % nsize > 0:
            print("The number of markers must be a multiple of dimensions.")
            return

        nslot = int(ntot / nsize)
        K = np.zeros(dim).ravel()
        D = np.zeros(dim).ravel()
        A1 = np.zeros(dim).ravel()
        A2 = np.zeros(dim).ravel()
        A3 = np.zeros(dim).ravel()

        i = 0
        n = 0
        while n < ntot:
            idx  = np.logical_and(  ids > n,  ids <= n + nslot )
            idx0 = np.logical_and( ids0 > n, ids0 <= n + nslot )

            # Evaluate transport from losses instead
            lf = np.sum(lost[ids0[idx0]-1]) / nslot
            if lf >= lossfrac:
                coefs = eval_coefficients(
                    [self._runnode], method=method, endrho=1,
                    ids=[ids0[idx0]-1])
                K[i] = coefs[1]
                D[i] = coefs[2]

                A1[i] = np.mean(rho[ids0[idx0]-1])
                A2[i] = np.mean(energy[ids0[idx0]-1])
                A3[i] = np.mean(pitch[ids0[idx0]-1])

            else:
                K[i] = np.mean(k[idx])
                D[i] = np.mean(d[idx])

                A1[i] = np.mean(rho[ids[idx]-1])
                A2[i] = np.mean(energy[ids[idx]-1])
                A3[i] = np.mean(pitch[ids[idx]-1])

            n = n + nslot
            i = i + 1

        K = np.reshape(K, dim)
        D = np.reshape(D, dim)
        A1 = np.reshape(A1, dim)
        A2 = np.reshape(A2, dim)
        A3 = np.reshape(A3, dim)
        return (A1, A2, A3, K, D)

a5py/a5py/ascot5io/transcoef.py
ascot4fusion/ascot5
"""Input representing first wall and other plasma facing components.

Wall input is used in simulations to check intersections between a marker
trajectory and physical components like first wall of fast-ion-loss detector.
It is also used in BBNBI5 to estimate shinethrough and in ASCOT-BMC.

Wall models can also be used to mimic LCFS to catch markers that escape plasma.
This could prove more detailed information than just using the MAXRHO end
condition.
"""
import h5py
import numpy as np
import unyt

from .coreio.fileapi import add_group
from .coreio.treedata import DataGroup

import a5py.wall.plot as plot
import a5py.physlib as physlib

class wall_2D(DataGroup):
    """Contour in Rz-plane that represents an axisymmetric wall.

    This wall is simple and fast in simulations, but it does not allow
    evaluation of wall loads. For those, use 3D wall instead.

    The wall doesn't have to form a closed loop but it prevents markers
    from escaping the computational domain.
    """

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key == "nelements":
                    out[key] = int(out[key])

        return out

    @staticmethod
    def write_hdf5(fn, nelements, r, z, desc=None):
        """Write input data to the HDF5 file.

        First vertex shouldn't correspond to the last vertex, i.e., don't give
        (R,z) coordinates that make a closed loop.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        nelements : int
            Number of wall segments.
        r : array_like (nelements,1)
            R coordinates of wall segment vertices [m].
        z : array_like (nelements,1)
            z coordinates of wall segment vertices [m].
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If the input is inconsistent.
        """
        if r.size != nelements:
            raise ValueError("Size of r does not match the number of elements.")
        if z.size != nelements:
            raise ValueError("Size of z does not match the number of elements.")

        parent = "wall"
        group  = "wall_2D"
        gname  = ""

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset("r",         (nelements,1), data=r, dtype='f8')
            g.create_dataset("z",         (nelements,1), data=z, dtype='f8')
            g.create_dataset("nelements", (1,1),         data=nelements,
                             dtype='i4')

        return gname

    def getwallcontour(self, phi=0*unyt.deg):
        """Return a cross section of the wall with a given poloidal plane.

        Parameters
        ----------
        phi : float
            Toroidal angle of the plane.

        Returns
        -------
        lines : array_like (n,2,2)
            Line segments [[[r1,z1], [r2,z2]], ...] that form the cross section.
        """
        w = self.read()
        r = np.append(w["r"], w["r"][0])
        z = np.append(w["z"], w["z"][0])
        lines = np.zeros((r.size-1, 2, 2))
        for i in range(r.size-1):
            lines[i, 0, 0] = r[i]
            lines[i, 0, 1] = z[i]
            lines[i, 1, 0] = r[i+1]
            lines[i, 1, 1] = z[i+1]
        return lines

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output is a very large rectangular wall.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        r = np.array([0.01, 100, 100, 0.01])
        z = np.array([-100, -100, 100, 100])
        return wall_2D.write_hdf5(fn=fn, nelements=4, r=r, z=z, desc="DUMMY")


class wall_3D(DataGroup):
    """Triangle mesh that represents a 3D wall.

    This wall model can be obtained easily from a STEP file or from a CAD model
    with a little effort. If 3D model of the wall is not available, consider
    using the mock-up 3D wall generated by "rotating" the 2D wall contour
    toroidally in steps.

    3D wall model allows estimating the wall loads, but keep in mind that
    smaller triangles means more markers are required in simulation to gain
    sufficient statistics. On the other hand, too coarse model means that
    important details on wall load distribution could be missed.

    This wall doesn't have to form a closed "waterproof" structure, and one
    can use this to model FILD.
    """

    number_of_elements = None

    def read(self):
        """Read data from HDF5 file.

        Returns
        -------
        data : dict
            Data read from HDF5 stored in the same format as is passed to
            :meth:`write_hdf5`.
        """
        fn   = self._root._ascot.file_getpath()
        path = self._path

        out = {}
        with h5py.File(fn,"r") as f:
            for key in f[path]:
                out[key] = f[path][key][:]
                if key == "nelements":
                    out[key] = int(out[key])

            nTriangles = out['x1x2x3'].shape[0]

            if path+'/flag' in f:
                flagAttrs = ['flagIdStrings','flagIdList']
                for s in flagAttrs:
                    if s in f[ path+'/flag' ].attrs:
                        out[s] = f[ path+'/flag' ].attrs.get(s)

        if not 'flag' in out:
            out['flag'] = np.zeros(shape=(nTriangles,),dtype=int)

        if 'flagIdStrings' in out:
            # We need to decode the bytearrays into strings.
            s=[]
            for S in out['flagIdStrings']:
                if not isinstance(S,str):
                    s.append(S.decode('utf-8'))
                else:
                    s.append(S)
                out['flagIdStrings']=s
        else:
            # Generate some flag names.
            out['flagIdList']=np.unique(out['flag'])
            out['flagIdStrings']=[]
            for fl in out['flagIdList']:
                out['flagIdStrings'].append('Flag {}'.format(fl))

        return out

    def getNumberOfElements(self):
        if self.number_of_elements is None:
            self.read()
        return self.number_of_elements


    def noderepresentation(self):
        """
        Return an array of vertices and indices that define this wall.

        Returns:
            vertices : float, array_like <br>
                Triangle vertices in a format [[x1, y1, z1], ..., [xn, yn, zn]]
                where n is 3 x number of triangles.
            faces : int, array_like <br>
                Array of dimension (number of triangles, 4) where each four
                length array defines a wall elemenet as [number of vertices, i1,
                i2, i3] where number of vertices is always 3 and i1, i2, and i3
                define the indices of the vertices in the vertices array.
        """
        with self as h5:
            ntriangle = int(h5["nelements"][:])
            faces     = np.zeros((ntriangle,4), dtype="i8")
            vertices  = np.zeros((ntriangle*3,3), dtype="f8")

            faces[:, 0] = 3
            faces[:, 1] = np.arange(ntriangle) * 3 + 0
            faces[:, 2] = np.arange(ntriangle) * 3 + 1
            faces[:, 3] = np.arange(ntriangle) * 3 + 2

            vertices[0::3,0] = h5["x1x2x3"][:,0]
            vertices[0::3,1] = h5["y1y2y3"][:,0]
            vertices[0::3,2] = h5["z1z2z3"][:,0]
            vertices[1::3,0] = h5["x1x2x3"][:,1]
            vertices[1::3,1] = h5["y1y2y3"][:,1]
            vertices[1::3,2] = h5["z1z2z3"][:,1]
            vertices[2::3,0] = h5["x1x2x3"][:,2]
            vertices[2::3,1] = h5["y1y2y3"][:,2]
            vertices[2::3,2] = h5["z1z2z3"][:,2]

        return (vertices, faces)

    def remove_small_triangles(self,maximumAreaToRemove=0.0,data=None):
        "The modification happens in-place. No deep copy is made!"

        if data is None:
            w = self.read()
        else:
            w = data

        A = self.area(data=w)

        keep =  ( A > maximumAreaToRemove )

        fields = ["x1x2x3","y1y2y3","z1z2z3"]
        for f in fields:
            w[f] = w[f][keep,:]
        w['flag'] = w['flag'][keep]

        nOld = w['n'][:]
        nNew = len(w['flag'])
        w['n'][:]         = nNew
        w['nelements'][:] = nNew

        print('Removing {}/{} triangles'.format(nOld-nNew,nOld))

        return w

    def area(self, data=None):
        if data is None:
            w = self.read()
        else:
            w = data

        ab_x = w["x1x2x3"][:,1] - w["x1x2x3"][:,0]
        ab_y = w["y1y2y3"][:,1] - w["y1y2y3"][:,0]
        ab_z = w["z1z2z3"][:,1] - w["z1z2z3"][:,0]

        ac_x = w["x1x2x3"][:,2] - w["x1x2x3"][:,0]
        ac_y = w["y1y2y3"][:,2] - w["y1y2y3"][:,0]
        ac_z = w["z1z2z3"][:,2] - w["z1z2z3"][:,0]

        A = 0.5 * np.sqrt(   (ab_y * ac_z - ab_z * ac_y)**2
                           + (ab_z * ac_x - ab_x * ac_z)**2
                           + (ab_x * ac_y - ab_y * ac_x)**2 ) * unyt.m**2

        return A

    def getAspointsAndVertices(self, removeDuplcatePoints=True):
        a5wall = self.read()
        nTriangles = self.number_of_elements

        # Create an array
        points = np.zeros(shape=(3*nTriangles,3))

        points[ ::3,0]=a5wall['x1x2x3'][:,0]
        points[1::3,0]=a5wall['x1x2x3'][:,1]
        points[2::3,0]=a5wall['x1x2x3'][:,2]
        points[ ::3,1]=a5wall['y1y2y3'][:,0]
        points[1::3,1]=a5wall['y1y2y3'][:,1]
        points[2::3,1]=a5wall['y1y2y3'][:,2]
        points[ ::3,2]=a5wall['z1z2z3'][:,0]
        points[1::3,2]=a5wall['z1z2z3'][:,1]
        points[2::3,2]=a5wall['z1z2z3'][:,2]


        vertices = np.reshape( np.arange(0,3*nTriangles,1,dtype=int), (nTriangles,3) )

        if not removeDuplcatePoints:
            return points,vertices

        points,vertices = self._removeDuplicatePoints(points,vertices)
        return points,vertices

    def _removeDuplicatePoints(self,points,vertices):
        """
        Remove duplicate points from points and vertices representation.
        """

        upoints,inverse = np.unique(points, return_index=False,
                                    return_inverse=True, axis=0)


        uvertices = inverse[vertices]

        return upoints,uvertices

    def getwallcontour(self, phi=0*unyt.deg):
        """Return a cross section of the wall with a given poloidal plane.

        Parameters
        ----------
        phi : float
            Toroidal angle of the plane.

        Returns
        -------
        lines : array_like (n,2,2)
            Line segments [[[r1,z1], [r2,z2]], ...] that form the cross section.
        """
        import pyvista as pv
        phi = phi.to("rad").v
        s1 = self.tomesh()
        rmin = np.amin(np.sqrt( s1.points[:,0]**2 + s1.points[:,1]**2 ))
        rmax = np.amax(np.sqrt( s1.points[:,0]**2 + s1.points[:,1]**2 ))
        zmin = np.amin(s1.points[:,2])
        zmax = np.amax(s1.points[:,2])
        r0 = 0.5*(rmin+rmax)
        z0 = 0.5*(zmin+zmax)
        dr = (rmax - rmin)
        dz = (zmax - zmin)

        x, y, z = physlib.pol2cart(r0, phi, z0)
        s2 = pv.Plane(center=(x,y,z), direction=(-np.sin(phi),np.cos(phi),0),
                      i_size=dr*1.1, j_size=dz*1.1,
                      i_resolution=1, j_resolution=1).triangulate()

        cut,_,_ = s2.intersection(s1, split_first=False, split_second=False)
        r, phi, z = physlib.cart2pol(cut.points[:,0], cut.points[:,1],
                                     cut.points[:,2])
        i0 = cut.lines[1::3]
        i1 = cut.lines[2::3]
        lines = np.array( [
            ( (r[i0[i]], z[i0[i]]), (r[i1[i]], z[i1[i]]) ) \
            for i in range(i0.size) ] )
        return lines

    def tomesh(self):
        import pyvista as pv
        return pv.PolyData( *self.noderepresentation() )

    def toVtk(self):
        import a5py.wall.a5vtkwall

        a5VTKwall = a5py.wall.a5vtkwall.a5VtkWall()

        a5VTKwall.fromA5wall(self)

        W=self.read()

        # Add index of running triangle number
        a5VTKwall.addIndex(setAsActive=True)

        # Add coloring based on flags. (Get's activated only if there are more
        # than one flag.
        a5VTKwall.addFlag(W['flag'],W['flagIdList'],W['flagIdStrings'])

        return a5VTKwall


    def append(self,newWall,data=None):
        """
        Returns the wall with newWall["x1x2x3], newWall["y1y2y3"], newWall["z1z2z3"], newWall["flag"] to the existing triangles.
        If data=None (or not given) reads the old triangles from file.
        """
        if data is None:
            w = self.read()
        else:
            w = data

        fields = ["x1x2x3","y1y2y3","z1z2z3","flag"]
        for f in fields:
            #print('Field "{}" has the shapes (data,newWall):'.format(f))
            #print(w[f].shape, newWall[f].shape)
            if(len(w[f].shape)==1):
                if( len(newWall[f]) != len(newWall[f].ravel()) ):
                    raise Exception("Wrong size field {}".format(f))
                w[f] = np.concatenate( (w[f], newWall[f].ravel() ), axis=0 )
            else:
                w[f] = np.concatenate( (w[f], newWall[f]         ), axis=0 )

        w['n'] = w['n'] + newWall['n']

        return w


    def move_component(self, movement, direction, component):
        """
        @Params:
        filename    Filename where the wall will be saved.
        movement    How much the component is moved in metres
        direction    Which direction the component is moved [x, y, z]
        component     Bool array of the points of a component
        wall         Wall object where the wall is copied from, if left empty active 
        wall from filename will be used

        @return the new wall data
        """
        rwall = self.read()

        #Movement in metres
        t = movement/np.sqrt(direction[0]**2 + direction[1]**2 + direction[2]**2) 

        rwall['x1x2x3'][component,:] += t*direction[0]
        rwall['y1y2y3'][component,:] += t*direction[1]
        rwall['z1z2z3'][component,:] += t*direction[2]

        new_rwall = {'x1x2x3':rwall['x1x2x3'],'y1y2y3':rwall['y1y2y3'],
                     'z1z2z3':rwall['z1z2z3'],
                     'desc':'ICRHmv{}m'.format(movement),
            'flag':rwall['flag'], 'flagIdList':rwall['flagIdList'],
            'flagIdStrings':rwall['flagIdStrings'], 'nelements':rwall['n']}

        return new_rwall

    def set_flags( self, names, indexes):
        '''
        @Params:
        names       list of strings containing names of various flags
        indexes     which flag indexes (zero indexed, last index value not included) correspond to each flag, list of tuples

        example with 1000 triangles:
                 names   = ['foo','bar','baz']
                 indexes = [(0,10), (10,100), (100,1000)]

        @return the new wall data
        '''
        W = self.read()

        flagIdList    = []
        flagIdStrings = []

        for i,name in enumerate(names):
            W['flag'][indexes[i][0]:(indexes[i][1])] = i
            flagIdList.append(i)
            flagIdStrings.append(name)

        W['flagIdList']    = np.array(  flagIdList )
        W['flagIdStrings'] = flagIdStrings
        return W


    @staticmethod
    def write_hdf5(fn, nelements, x1x2x3, y1y2y3, z1z2z3, flag=None,
                   flagIdList=None, flagIdStrings=None, desc=None):
        """Write input data to the HDF5 file.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.
        nelements : int
            Number of wall triangles
        x1x2x3 : array_like (nelements,3)
            Each triangle's vertices' x coordinates [m].
        y1y2y3 : array_like (nelements,3)
            Each triangle's vertices' y coordinates [m].
        z1z2z3 : array_like (nelements,3)
            Each triangle's vertices' z coordinates [m].
        flag : array_like (nelements,1), optional
            Integer array depicting the wall component of each triangle.
        flagIdList : array_like (nUniqueFlags), optional
            List of keys (int) of the flagIdStrings.
        flagIdStrings : array_like (nUniqueFlags), optional
            List of values (str) of the flagIdStrings.
        desc : str, optional
            Input description.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.

        Raises
        ------
        ValueError
            If the triangle vertices or flags have incorrect shape.
        """
        if type(nelements) == int:
            pass
        elif type(nelements) == list:
            while type(nelements) == list:
                nelements = nelements[0]
        elif type(nelements) == np.ndarray:
            nelements = int(nelements)
        else:
            raise ValueError("Unsupported format for nelements.")

        if x1x2x3.shape != (nelements,3):
            raise ValueError(
                "Shape of x1x2x3 was " + str(x1x2x3.shape) + " but expected ("
                + str(nelements) + ",3)")
        if y1y2y3.shape != (nelements,3):
            raise ValueError(
                "Shape of y1y2y3 was " + str(y1y2y3.shape) + " but expected ("
                + str(nelements) + ",3)")
        if z1z2z3.shape != (nelements,3):
            raise ValueError(
                "Shape of z1z2z3 was " + str(z1z2z3.shape) + " but expected ("
                + str(nelements) + ",3)")

        if flag is None:
            flag = np.zeros(shape=(nelements,),dtype=int)
        elif flag.shape != (nelements,1):
            raise ValueError(
                "Shape of flag was " + str(flag.shape) + " but expected ("
                + str(nelements) + ",1)")

        parent = "wall"
        group  = "wall_3D"
        gname  = ""

        # Convert strings to the favorite format
        if flagIdStrings is not None:
            strlen = 0
            for s in flagIdStrings:
                if len(s) >  strlen:
                    strlen = len(s)
            fids = np.empty( shape=(len(flagIdStrings),),
                             dtype='|S{}'.format(strlen) )
            for istr,s in enumerate(flagIdStrings):
                fids[istr]=s

        with h5py.File(fn, "a") as f:
            g = add_group(f, parent, group, desc=desc)
            gname = g.name.split("/")[-1]

            g.create_dataset('x1x2x3',    (nelements,3), data=x1x2x3,
                             dtype='f8')
            g.create_dataset('y1y2y3',    (nelements,3), data=y1y2y3,
                             dtype='f8')
            g.create_dataset('z1z2z3',    (nelements,3), data=z1z2z3,
                             dtype='f8')
            g.create_dataset('nelements', (1,1),         data=nelements,
                             dtype='i4')

            fl = g.create_dataset('flag', (nelements,1), data=flag, dtype='i4')
            if flagIdList is not None and flagIdStrings is not None:
                flagIdList = np.array(flagIdList)
                fl.attrs.create(name='flagIdList', data=flagIdList, dtype='i4')
                fl.attrs.create(name='flagIdStrings', data=fids,
                                dtype=h5py.string_dtype(encoding='utf-8',
                                                        length=None))

        return gname

    @staticmethod
    def write_hdf5_dummy(fn):
        """Write dummy data that has correct format and is valid, but can be
        non-sensical.

        This method is intended for testing purposes or to provide data whose
        presence is needed but which is not actually used in simulation.

        The dummy output consists of a single triangle on a poloidal plane.

        Parameters
        ----------
        fn : str
            Full path to the HDF5 file.

        Returns
        -------
        name : str
            Name, i.e. "<type>_<qid>", of the new input that was written.
        """
        return wall_3D.write_hdf5(
            fn=fn, nelements=1, x1x2x3=np.array([[4],[8],[8]]).T,
            y1y2y3=np.array([[0],[0],[0]]).T,
            z1z2z3=np.array([[-4],[-4],[4]]).T, desc="DUMMY")

    @staticmethod
    def convert_wall_2D(nphi, **kwargs):
        """Convert :class:`wall_2D` input to :class:`wall_3D` input.

        This function repeats the 2D wall contour around the torus at intervals
        Delta phi = 2 pi / (nphi - 1). Then each (R, z) point is connected to a
        equivalent point in the adjacent contours by a line segment forming 3D
        wall made of rectangles. Finally each rectangle is divided into two
        triangles.

        Parameters
        ----------
        nphi : int
            Number of toroidal segments to be created.
        **kwargs
            Arguments passed to :meth:`wall_2D.write_hdf5` excluding ``fn`` and
            ``desc``.

        Returns
        -------
        out : dict
            :class:`wall_2D` converted as an input for :meth:`write_hdf5`.
        """
        r = kwargs["r"]
        z = kwargs["z"]
        n = kwargs["nelements"]
        x1x2x3 = np.ones((2*n*nphi, 3))
        y1y2y3 = np.ones((2*n*nphi, 3))
        r1r2r3 = np.ones((2*n*nphi, 3))
        p1p2p3 = np.ones((2*n*nphi, 3))
        z1z2z3 = np.ones((2*n*nphi, 3))

        def pol2cart(rho, phi):
            """Poloidal coordinates to Cartesian.
            """
            x = rho * np.cos(phi)
            y = rho * np.sin(phi)
            return(x, y)

        pv = np.linspace(0, 2*np.pi, nphi+1)
        for i in range(1,nphi+1):
            for j in range(n):
                r1r2r3[(i-1)*2*n + 2*(j-1),:] = [ r[j],    r[j],  r[j-1] ]
                p1p2p3[(i-1)*2*n + 2*(j-1),:] = [ pv[i-1], pv[i], pv[i] ]
                z1z2z3[(i-1)*2*n + 2*(j-1),:] = [ z[j],    z[j],  z[j-1] ]

                r1r2r3[(i-1)*2*n + 2*(j-1) + 1,:] = [ r[j],    r[j-1],  r[j-1] ]
                p1p2p3[(i-1)*2*n + 2*(j-1) + 1,:] = [ pv[i-1], pv[i-1], pv[i] ]
                z1z2z3[(i-1)*2*n + 2*(j-1) + 1,:] = [ z[j],    z[j-1],  z[j-1] ]

        x1x2x3,y1y2y3 = pol2cart(r1r2r3, p1p2p3)

        return {"nelements" : 2*n*nphi, "x1x2x3" : x1x2x3,
                "y1y2y3" : y1y2y3, "z1z2z3" : z1z2z3}

a5py/a5py/ascot5io/wall.py
ascot4fusion/ascot5
"""Python interface to libascot.so.
"""
import ctypes
import unyt
import numpy as np
import scipy
import warnings
import wurlitzer # For muting libascot.so

import a5py.physlib as physlib
import a5py.routines.plotting as a5plt
from a5py.routines.plotting import openfigureifnoaxes
from a5py.exceptions import AscotInitException

from .libascot    import LibAscot, _LIBASCOT
from .libsimulate import LibSimulate
from .libproviders import LibProviders

if _LIBASCOT:
    from . import ascot2py

class Ascotpy(LibAscot, LibSimulate, LibProviders):
    """Class with methods to initialize and access the data via Python.

    This class will be inherited by Ascot class. Here we hide all the dirty
    implementation of the libascot.so interface in private methods and
    attributes so that the Ascot front is clean for the users.

    MPI parallelism has not been tested.

    Attributes
    ----------
    _sim
        Simulation offload data struct.
    _offload_ready
        Flag indicating if inputs are packed.
    _nmrk
        Number of markers currently in the marker array.
    _inistate
        Marker input array for interactive simulations.
    _endstate
        Marker output array for interactive simulations.
    _diag_occupied
        Flag indicating if diagnostics array contains data.
    _diag_offload_array
        Diagnostics data offload array.
    _offload_array
        Offload array containing the input data when this instance is packed.
    _int_offload_array
        Offload array for integers containing the input data when this instance
        is packed.
    _bfield_offload_array
        Offload data for the magnetic field input.
    _efield_offload_array
        Offload data for the electric field input.
    _plasma_offload_array
        Offload data for the plasma input.
    _neutral_offload_array
        Offload data for the neutral input.
    _wall_offload_array
        Offload data for the wall (float) input.
    _wall_int_offload_array
        Offload data for the wall (int) input.
    _boozer_offload_array
        Offload data for the Boozer input.
    _mhd_offload_array
        Offload data for the MHD input.
    _asigma_offload_array
        Offload data for the atomic data input.
    _mpi_root
        Rank of the root MPI process.
    _mpi_rank
        Rank of this MPI process.
    _mpi_size
        Number of MPI processes.
    _mute
        Mute output from libascot.so: "yes" - stdout and stderr both muted,
        "no" - output is not muted, "err" - stderr is printed.
    """

    DUMMY_QID = "".encode('UTF-8')
    """Flag to use in _sim.qid_* to indicate that the input is not initialized.
    """

    def __init__(self):
        """Initialize pointers to data that is being passed between Python and C
        via libascot.
        """
        if not _LIBASCOT:
            return

        # Initialize attributes
        self._offload_ready    = False
        self._nmrk             = ctypes.c_int32()
        self._diag_occupied    = False
        self._offload_data      = ascot2py.struct_c__SA_offload_package()
        self._offload_array     = ctypes.POINTER(ctypes.c_double)()
        self._int_offload_array = ctypes.POINTER(ctypes.c_int   )()
        self._inistate = ctypes.POINTER(ascot2py.struct_c__SA_particle_state)()
        self._endstate = ctypes.POINTER(ascot2py.struct_c__SA_particle_state)()

        self._sim = ascot2py.struct_c__SA_sim_offload_data()
        self._bfield_offload_array  = ctypes.POINTER(ctypes.c_double)()
        self._efield_offload_array  = ctypes.POINTER(ctypes.c_double)()
        self._plasma_offload_array  = ctypes.POINTER(ctypes.c_double)()
        self._neutral_offload_array = ctypes.POINTER(ctypes.c_double)()
        self._wall_offload_array    = ctypes.POINTER(ctypes.c_double)()
        self._boozer_offload_array  = ctypes.POINTER(ctypes.c_double)()
        self._mhd_offload_array     = ctypes.POINTER(ctypes.c_double)()
        self._asigma_offload_array  = ctypes.POINTER(ctypes.c_double)()
        self._diag_offload_array    = ctypes.POINTER(ctypes.c_double)()

        self._wall_int_offload_array = ctypes.POINTER(ctypes.c_int)()

        # MPI init needs to be called (only) once so that we can run simulations
        argc       = ctypes.c_int32()
        argc.value = 0
        argv       = ctypes.POINTER(ctypes.c_char)()
        self._mpi_rank = ctypes.c_int32()
        self._mpi_root = ctypes.c_int32()
        self._mpi_size = ctypes.c_int32()
        ascot2py.mpi_interface_init(
            argc, ctypes.byref(argv), ctypes.byref(self._sim),
            ctypes.byref(self._mpi_rank), ctypes.byref(self._mpi_size),
            ctypes.byref(self._mpi_root))
        self._mute = "no"

    def _init(self, data, bfield=None, efield=None, plasma=None,
              wall=None, neutral=None, boozer=None, mhd=None, asigma=None,
              switch=False):
        """Read, offload, and initialize input data so it can be accessed
        by libascot.

        Assumes data is present in the hdf5 file and the provided QIDs are
        valid. If the parameter is given as a dict in same format as it is
        written, it is used directly instead of reading the from the file.

        Parameters
        ----------
        data : Ascot5IO
            Ascot5 data on disk which is used in initialization.
        bfield : str or dict
            QID of the magnetic field to be initialized or the data as a
            dictionary.
        efield : str or dict
            QID of the electric field to be initialized or the data as a
            dictionary.
        plasma : str or dict
            QID of the plasma data to be initialized or the data as a
            dictionary.
        wall : str or dict
            QID of the wall data to be initialized or the data as a dictionary.
        neutral : str or dict
            QID of the neutral data to be initialized or the data as a
            dictionary.
        boozer : str or dict
            QID of the boozer to be initialized or the data as a dictionary.
        mhd : str or dict
            QID of the MHD data to be initialized or the data as a dictionary.
        asigma : str or dict
            QID of the atomicdata to be initialized or the data as a dictionary.
        switch : bool
            If ``True``, free input that has been
        """
        if self._offload_ready:
            raise AscotInitException("This instance has been packed")

        # Iterate through all inputs and mark those that are initialized
        inputs2read = ctypes.c_int32()
        args = locals() # Contains function arguments and values in a dictionary

        # List here dependencies to be directly injected (provided)
        to_be_provided = []
        for inp in ["bfield", "efield", "plasma", "wall", "neutral", "boozer",
                    "mhd", "asigma"]:
            if args[inp] is None:
                # This input is not initialized
                continue

            # Inject/provide this input
            if isinstance(args[inp],dict):
                to_be_provided.append(inp)
                continue

            # Convert QID strings to bytes
            args[inp] = args[inp].encode("UTF-8")

            currentqid = getattr(self._sim, "qid_" + inp)
            if args[inp] == currentqid:
                # This input is already initialized
                continue

            if currentqid != Ascotpy.DUMMY_QID:
                # Other input of same type is initialized. Raise error or free
                # the old input.
                if not switch:
                    raise AscotInitException(
                        "Cannot initialize " + inp
                        + ": other input is already initialized.\n"
                        "Either use \"switch=True\" or free(" + inp + "=True)")

                self._free(**{inp : True})

            # Mark this input to be read
            inputs2read.value |= getattr(ascot2py, "hdf5_input_" + inp)

        # Separate loop to mark QIDs of the inputs that are read. (We couldn't
        # do this in the previous loop as that might get interrupted by
        # an exception, in which case sim.qid_* would point to data which is not
        # initialized.
        for inp in ["bfield", "efield", "plasma", "wall", "neutral", "boozer",
                    "mhd", "asigma"]:
            if inputs2read.value & getattr(ascot2py, "hdf5_input_" + inp):
                setattr(self._sim, "qid_" + inp, args[inp])

        def hdf5init():
            """Simple wrapper.
            """
            ascot2py.hdf5_interface_read_input(
                ctypes.byref(self._sim),
                inputs2read,
                ctypes.byref(self._bfield_offload_array),
                ctypes.byref(self._efield_offload_array),
                ctypes.byref(self._plasma_offload_array),
                ctypes.byref(self._neutral_offload_array),
                ctypes.byref(self._wall_offload_array),
                ctypes.byref(self._wall_int_offload_array),
                ctypes.byref(self._boozer_offload_array),
                ctypes.byref(self._mhd_offload_array),
                ctypes.byref(self._asigma_offload_array),
                None, # Marker array (ignore)
                None  # Number of markers that were read (ignore)
            )

        if self._mute == "no":
            hdf5init()
        else:
            with wurlitzer.pipes() as (out, err):
                hdf5init()
            err = err.read()
            if self._mute == "err" and len(err) > 1: print(err)

        for inp in to_be_provided:
            if inp == "wall":
                if 'x1x2x3' in args[inp]:
                    self.provide_wall_3d(  args[inp]['x1x2x3'],  args[inp]['y1y2y3'],  args[inp]['z1z2z3']  )
                elif 'r' in args[inp]:
                    self.provide_wall_2d(  args[inp]['r'],  args[inp]['z']  )
                else:
                    raise AscotInitException("Unsupported dict for input '{}' passed for injection.".format(inp))                    
            elif inp == "bfield":
                if 'b_rmin' in args[inp]:
                    self.provide_BSTS( b_rmin=args[inp]['b_rmin'], b_rmax=args[inp]['b_rmax'], b_nr=args[inp]['b_nr'],
                                       b_zmin=args[inp]['b_zmin'], b_zmax=args[inp]['b_zmax'], b_nz=args[inp]['b_nz'],
                                       b_phimin=args[inp]['b_phimin'], b_phimax=args[inp]['b_phimax'], b_nphi=args[inp]['b_nphi'],
                                       psi0=args[inp]['psi0'], psi1=args[inp]['psi1'],
                                       br=args[inp]['br'], bphi=args[inp]['bphi'], bz=args[inp]['bz'], psi=args[inp]['psi'],
                                       axis_phimin=args[inp]['axis_phimin'], axis_phimax=args[inp]['axis_phimax'],
                                       axis_nphi=args[inp]['axis_nphi'],
                                       axisr=args[inp]['axisr'], axisz=args[inp]['axisz'],
                                       psi_rmin=args[inp]['psi_rmin'],     psi_rmax=args[inp]['psi_rmax'], psi_nr=args[inp]['psi_nr'],
                                       psi_zmin=args[inp]['psi_zmin'],     psi_zmax=args[inp]['psi_zmax'], psi_nz=args[inp]['psi_nz'],
                                       psi_phimin=args[inp]['psi_phimin'], psi_phimax=args[inp]['psi_phimax'], psi_nphi=args[inp]['psi_nphi'] )
                else:
                    raise AscotInitException("Unsupported dict for input '{}' passed for injection.".format(inp))
            else:
                raise AscotInitException("Unsupported input to inject: '{}'".format(inp))

    def _free(self, bfield=False, efield=False, plasma=False, wall=False,
              neutral=False, boozer=False, mhd=False, asigma=False):
        """Free input data initialized in C-side.
        """
        if self._offload_ready:
            raise AscotInitException("This instance has been packed")

        args = locals() # Contains function arguments and values in a dictionary

        # Iterate through all inputs and free the data if the corresponding
        # argument is True
        for inp in ["bfield", "efield", "plasma", "wall", "neutral", "boozer",
                    "mhd", "asigma"]:
            if args[inp] and \
               getattr(self._sim, "qid_" + inp) != Ascotpy.DUMMY_QID:

                # Deallocate the data allocated in C side
                array = getattr(self, "_" + inp + "_offload_array")
                ascot2py.libascot_deallocate(array)
                if inp == "wall":
                    array = getattr(self, "_" + inp + "_int_offload_array")
                    ascot2py.libascot_deallocate(array)

                # Set the offload array length to zero
                if inp == "bfield":
                    self._sim.B_offload_data.offload_array_length = 0
                elif inp == "efield":
                    self._sim.E_offload_data.offload_array_length = 0
                else:
                    data = getattr(self._sim, inp + "_offload_data")
                    data.offload_array_length = 0
                    if inp == "wall":
                        data.int_offload_array_length = 0

                # Set QID to dummy value
                setattr(self._sim, "qid_" + inp, Ascotpy.DUMMY_QID)

    def _pack(self):
        """Pack offload arrays as one making this instance ready for simulation.

        Note that inputs cannot be changed or freed before calling _unpack. Make
        sure all required data is initialized before packing.
        """
        if self._offload_ready:
            raise AscotInitException("This instance is already packed")

        # This call internally frees individual offload arrays and initializes
        # the common ones.
        ascot2py.pack_offload_array(
            ctypes.byref(self._sim), ctypes.byref(self._offload_data),
            self._bfield_offload_array, self._efield_offload_array,
            self._plasma_offload_array, self._neutral_offload_array,
            self._wall_offload_array,   self._wall_int_offload_array,
            self._boozer_offload_array, self._mhd_offload_array,
            self._asigma_offload_array,
            ctypes.byref(self._offload_array),
            ctypes.byref(self._int_offload_array))
        self._offload_ready = True

        # Set pointers to correct locations (based on the order arrays are
        # packed) so that we can continue using evaluation routines.
        def advance(prevptr, increment):
            ptr = ctypes.addressof(prevptr) + 2 * increment
            ptr = ctypes.cast(ctypes.c_void_p(ptr),
                              ctypes.POINTER(ctypes.c_double))
            return ptr

        self._bfield_offload_array = self._offload_array
        self._efield_offload_array = advance(
            self._bfield_offload_array.contents,
            self._sim.B_offload_data.offload_array_length)
        self._plasma_offload_array = advance(
            self._efield_offload_array.contents,
            self._sim.E_offload_data.offload_array_length)
        self._neutral_offload_array = advance(
            self._plasma_offload_array.contents,
            self._sim.plasma_offload_data.offload_array_length)
        self._wall_offload_array = advance(
            self._neutral_offload_array.contents,
            self._sim.neutral_offload_data.offload_array_length)
        self._boozer_offload_array = advance(
            self._wall_offload_array.contents,
            self._sim.wall_offload_data.offload_array_length)
        self._mhd_offload_array = advance(
            self._boozer_offload_array.contents,
            self._sim.boozer_offload_data.offload_array_length)
        self._asigma_offload_array = advance(
            self._mhd_offload_array.contents,
            self._sim.mhd_offload_data.offload_array_length)

        self._wall_int_offload_array = self._int_offload_array


    def _unpack(self, bfield=True, efield=True, plasma=True, wall=True,
                neutral=True, boozer=True, mhd=True, asigma=True):
        """Unpack simulation arrays, i.e. free offload array and re-read data.

        After unpacking the inputs can be changed or freed again but simulations
        cannot be performed.
        """
        if not self._offload_ready:
            raise AscotInitException("This instance hasn't been packed")

        ascot2py.libascot_deallocate(self._offload_array)
        ascot2py.libascot_deallocate(self._int_offload_array)
        self._offload_data.offload_array_length     = 0
        self._offload_data.int_offload_array_length = 0
        self._offload_data.unpack_pos               = 0
        self._offload_data.int_unpack_pos           = 0

        inputs2read = {}
        def readornot(name, read, qid):
            if read:
                inputs2read[name] = qid.decode("utf-8")
            return Ascotpy.DUMMY_QID

        self._sim.qid_bfield  = readornot("bfield",  bfield,
                                          self._sim.qid_bfield)
        self._sim.qid_efield  = readornot("efield",  efield,
                                          self._sim.qid_efield)
        self._sim.qid_plasma  = readornot("plasma",  plasma,
                                          self._sim.qid_plasma)
        self._sim.qid_neutral = readornot("neutral", neutral,
                                          self._sim.qid_neutral)
        self._sim.qid_wall    = readornot("wall",    wall,
                                          self._sim.qid_wall)
        self._sim.qid_boozer  = readornot("boozer",  boozer,
                                          self._sim.qid_boozer)
        self._sim.qid_mhd     = readornot("mhd",     mhd,
                                          self._sim.qid_mhd)
        self._sim.qid_asigma  = readornot("asigma",  asigma,
                                          self._sim.qid_asigma)

        self._offload_ready = False
        self.input_init(**inputs2read)

    def _requireinit(self, *inputs):
        """Raise error if given input parent is not initialized.

        Parameters
        ----------
        *inputs : [str]
            Name(s) of the input parent group(s) e.g. "bfield".

        Raises
        ------
        AscotInitException
            If input parent group has no initialized data.
        """
        for inp in inputs:
            qid = getattr(self._sim, "qid_" + inp)
            if qid == Ascotpy.DUMMY_QID:
                raise AscotInitException(inp + " is not initialized")

    def input_initialized(self):
        """Get inputs that are currently initialized.

        Returns
        -------
        outputs : dict [str, str]
            Name of the input parent groups that are initialized and QID of the
            initialized input.
        """
        out = {}
        for inp in ["bfield", "efield", "plasma", "wall", "neutral", "boozer",
                    "mhd", "asigma"]:
            qid = getattr(self._sim, "qid_" + inp)
            if qid != Ascotpy.DUMMY_QID:
                out[inp] = qid.decode("utf-8")
        return out

    def input_eval_list(self, show=True):
        """Return quantities that can be evaluated from inputs.

        The returned quantities are those that can be evaluated from the inputs
        with `input_eval` method. Units are not listed as that function returns
        quantities with units included.

        The number ion species in plasma input is not fixed so "ni" refers to
        i'th plasma species as listed by `input_getplasmaspecies` (index starts
        from 1).

        Parameters
        ----------
        show : bool, optional
            Print output.

        Returns
        -------
        quantities : dict [str, str]
            Name of each quantity and a short description.
        """

        out = {
            "rho":
            "Square root of normalized poloidal flux",
            "psi":
            "Poloidal flux",
            "br":
            "Magnetic field R component (not including MHD)",
            "bphi":
            "Magnetic field phi component (not including MHD)",
            "bz":
            "Magnetic field z component (not including MHD)",
            "bnorm":
            "Magnetic field norm (not including MHD)",
            "brdr":
            "d/dr of magnetic field R component (not including MHD)",
            "brdphi":
            "d/dphi of magnetic field R component (not including MHD)",
            "brdz":
            "d/dz of magnetic field R component (not including MHD)",
            "bphidr":
            "d/dr of magnetic field phi component (not including MHD)",
            "bphidphi":
            "d/dphi of magnetic field phi component (not including MHD)",
            "bphidz":
            "d/dz of magnetic field phi component (not including MHD)",
            "bzdr":
            "d/dr of magnetic field z component (not including MHD)",
            "bzdphi":
            "d/dphi of magnetic field z component (not including MHD)",
            "bzdz":
            "d/dz of magnetic field z component (not including MHD)",
            "divb":
            "Magnetic field divergence (not including MHD)",
            "gradbr":
            "Magnetic field gradient R component (not including MHD)",
            "gradbphi":
            "Magnetic field gradient phi component (not including MHD)",
            "gradbz":
            "Magnetic field gradient z component (not including MHD)",
            "curlbr":
            "Magnetic field curl R component (not including MHD)",
            "curlbphi":
            "Magnetic field curl phi component (not including MHD)",
            "curlbz":
            "Magnetic field curl z component (not including MHD)",
            "jr":
            "Current density R component",
            "jphi":
            "Current density phi component",
            "jz":
            "Current density z component",
            "jnorm":
            "Current density",
            "er":
            "Electric field R component (not including MHD)",
            "ephi":
            "Electric field phi component (not including MHD)",
            "ez":
            "Electric field z component (not including MHD)",
            "ne":
            "Electron density",
            "te":
            "Electron temperature",
            "n0":
            "Neutral density",
            "alphaeig":
            "Eigenfunction of the magnetic field MHD perturbation",
            "phieig":
            "Eigenfunction of the electric field MHD perturbation",
            "mhd_br":
            "R component of B-field perturbation due to MHD",
            "mhd_bphi":
            "phi component of B-field perturbation due to MHD",
            "mhd_bz":
            "z component of B-field perturbation due to MHD",
            "mhd_er":
            "R component of E-field perturbation due to MHD",
            "mhd_ephi":
            "phi component of E-field perturbation due to MHD",
            "mhd_ez":
            "z component of E-field perturbation due to MHD",
            "mhd_phi":
            "Electric potential due to MHD",
            "db/b (mhd)":
            """|B_MHD| / |B| where B is evaluated from the magnetic
            field input and both contain all components""",
            "psi (bzr)":
            "Boozer poloidal flux",
            "rho (bzr)":
            "Boozer radial coordinate",
            "theta":
            "Boozer poloidal coordinate",
            "zeta":
            "Boozer toroidal coordinate",
            "dpsidr (bzr)":
            "d/dR of Boozer psi",
            "dpsidphi (bzr)":
            "d/dphi of Boozer psi",
            "dpsidz (bzr)":
            "d/dz of Boozer psi",
            "dthetadr":
            "d/dr of Boozer theta",
            "dthetadphi":
            "d/dphi of Boozer theta",
            "dthetadz":
            "d/dz of Boozer theta",
            "dzetadr":
            "d/dr of Boozer zeta",
            "dzetadphi":
            "d/dphi of Boozer zeta",
            "dzetadz":
            "d/dz of Boozer zeta",
            "qprof":
            "Local value of the q-profile (safety factor)",
            "bjac":
            "Determinant of the Jacobian in Boozer coordinate transformation",
            "bjacxb2":
            "bjac multiplied by B^2 which is constant along flux surfaces",
        }

        nion, mass, charge, anum, znum = self.input_getplasmaspecies()
        for i in range(nion):
            out["ni" + str(i+1)] = "Ion species (anum, znum) = (%d, %d) density" \
                % (anum[i], znum[i])
            out["ti" + str(i+1)] = "Ion species (anum, znum) = (%d, %d) temperature" \
                % (anum[i], znum[i])

        if show:
            for name, desc in out.items():
                print(name.ljust(15) + " : " + desc)

        return out

    @physlib.parseunits(r="m", phi="deg", z="m", t="s")
    def input_eval(self, r, phi, z, t, *qnt, grid=False):
        """Evaluate input quantities at given coordinates.

        This method uses ASCOT5 C-routines for evaluating inputs exactly as
        they are evaluated during a simulation. See :meth:`input_eval_list` for
        a complete list of available input and derived quantities.

        Parameters
        ----------
        r : array_like
            R coordinates where data is evaluated.
        phi : array_like
            phi coordinates where data is evaluated.
        z : array_like
            z coordinates where data is evaluated.
        t : array_like
            Time coordinates where data is evaluated.
        *qnt : str
            Names of evaluated quantities.
        grid : boolean, optional
            Treat input coordinates as abscissae and return the evaluated
            quantities on a grid instead.

        Returns
        -------
        val : array_like (n,) or (nr,nphi,nz,nt)
            Quantity evaluated in each queried point.

            NaN is returned at those points where the evaluation failed.

            4D matrix is returned when grid=``True``.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        r   = unyt.unyt_array(r).astype(dtype=np.double).ravel()
        phi = unyt.unyt_array(phi).astype(dtype=np.double).ravel()
        z   = unyt.unyt_array(z).astype(dtype=np.double).ravel()
        t   = unyt.unyt_array(t).astype(dtype=np.double).ravel()

        phi = phi.to("rad")

        if grid:
            arrsize = (r.size, phi.size, z.size, t.size)
            r, phi, z, t = np.meshgrid(r, phi, z, t, indexing="ij")
            r   = r.ravel()
            phi = phi.ravel()
            z   = z.ravel()
            t   = t.ravel()
        else:
            # Not a grid so check that dimensions are correct (and make
            # single-valued vectors correct size)
            arrsize = np.amax(np.array([r.size, phi.size, z.size, t.size]))
            errmsg = "Input arrays have inconsistent sizes ({}, {}, {}, {})"
            assert (r.size == 1 or r.size == arrsize) and  \
                (phi.size == 1 or phi.size == arrsize) and \
                (z.size == 1 or z.size == arrsize) and     \
                (t.size == 1 or t.size == arrsize),        \
                errmsg.format(r.size, phi.size, z.size, t.size)

            if r.size == 1:
                r = r[0]*np.ones((arrsize,))
            if phi.size == 1:
                phi = phi[0]*np.ones((arrsize,))
            if z.size == 1:
                z = z[0]*np.ones((arrsize,))
            if t.size == 1:
                t = t[0]*np.ones((arrsize,))

        out = {}
        if any(q in qnt for q in ["rho", "psi", "rhodpsi", "psidr", "psidphi",
                                  "psidz"]):
            out.update(**self._eval_bfield(r, phi, z, t, evalrho=True))
        if any(q in qnt for q in ["br", "bphi", "bz", "brdr", "brdphi", "brdz",
                                  "bphidr", "bphidphi", "bphidz", "bzdr",
                                  "bzdphi", "bzdz", "divb", "bnorm",
                                  "jnorm", "jr", "jphi", "jz", "gradbr",
                                  "gradbphi", "gradbz", "curlbr", "curlbphi",
                                  "curlbz"]):
            out.update(self._eval_bfield(r, phi, z, t, evalb=True))
            out["divb"] = out["br"]/r + out["brdr"] + out["bphidphi"]/r \
                + out["bzdz"]
            out["bnorm"] = np.sqrt(out["br"]**2 + out["bphi"]**2 + out["bz"]**2)
            out["jr"]    = (out["bzdphi"]/r - out["bphidz"]) / unyt.mu_0
            out["jphi"]  = (out["brdz"] - out["bzdr"]) / unyt.mu_0
            out["jz"]    = (out["bphi"]/r + out["bphidr"] - out["brdphi"]/r) \
                            / unyt.mu_0
            out["jnorm"] = np.sqrt(out["jr"]**2 + out["jphi"]**2 + out["jz"]**2)
            out["gradbr"]   = (out["br"]*out["brdr"] + out["bphi"]*out["bphidr"]
                               + out["bz"]*out["bzdr"]) / out["bnorm"]
            out["gradbphi"] = (out["br"]*out["brdphi"]
                               + out["bphi"]*out["bphidphi"]
                               + out["bz"]*out["bzdphi"]) / (out["bnorm"]*r)
            out["gradbz"]   = (out["br"]*out["brdz"] + out["bphi"]*out["bphidz"]
                               + out["bz"]*out["bzdz"]) / out["bnorm"]
            out["curlbr"]   = out["bzdphi"] / r - out["bphidz"]
            out["curlbphi"] = out["brdz"] - out["bzdr"]
            out["curlbz"]   = (out["bphi"] - out["brdphi"]) / r + out["bphidr"]
        if any(q in qnt for q in ["axisr", "axisz"]):
            out.update(self._eval_bfield(r, phi, z, t, evalaxis=True))
        if any(q in qnt for q in ["er", "ephi", "ez"]):
            out.update(self._eval_efield(r, phi, z, t))
        if any(q in qnt for q in ["n0"]):
            out.update(self._eval_neutral(r, phi, z, t))
        if any(q in qnt for q in ["psi (bzr)", "theta", "zeta", "dpsidr (bzr)",
                                  "dpsidphi (bzr)", "dpsidz (bzr)", "dthetadr",
                                  "dthetadphi", "dthetadz", "dzetadr",
                                  "dzetadphi", "dzetadz", "rho (bzr)"]):
            out.update(self._eval_boozer(r, phi, z, t))
        if any(q in qnt for q in ["qprof", "bjac", "bjacxb2"]):
            out.update(self._eval_boozer(r, phi, z, t, evalfun=True))
        if any(q in qnt for q in ["alphaeig", "phieig"]):
            out.update(self._eval_mhd(r, phi, z, t, evalpot=True))
        if any(q in qnt for q in ["mhd_br", "mhd_bphi", "mhd_bz", "mhd_er",
                                  "mhd_ephi", "mhd_ez", "mhd_phi"]):
            out.update(self._eval_mhd(r, phi, z, t))
        if any(q in qnt for q in ["db/b (mhd)"]):
            bpert = self._eval_mhd(r, phi, z, t)
            b = self._eval_bfield(r, phi, z, t, evalb=True)
            bpert = np.sqrt(  bpert["mhd_br"]**2
                            + bpert["mhd_bphi"]**2
                            + bpert["mhd_bz"]**2)
            b = np.sqrt(b["br"]**2 + b["bphi"]**2 + b["bz"]**2)
            out["db/b (mhd)"] = bpert/b

        ni = ["ni" + str(i+1) for i in range(99)]
        ti = ["ti" + str(i+1) for i in range(99)]
        if any(q in qnt for q in ["ne", "te"] + ni + ti):
            out.update(self._eval_plasma(r, phi, z, t))

        for q in list(out.keys()):
            if q not in qnt:
                del out[q]
            elif grid:
                out[q] = np.reshape(out[q], arrsize)
        if len(qnt) == 1:
            return out[qnt[0]]
        else:
            return [out[q] for q in qnt]

    def input_plotrz(self, r, z, qnt, phi=0*unyt.deg, t=0*unyt.s,
                     clim=[None, None], cmap=None, axes=None, cax=None):
        """Plot input quantity on a (R, z) plane at given coordinates.

        To plot quantity on a logarithmic scale (base 10), add "log" in
        the name of the quantity e.g. "log ne".

        Parameters
        ----------
        r : array_like (nr,)
             R abscissa where data is evaluated and plotted.
        z : array_like (nz,)
            z abscissa where data is evaluated and plotted.
        qnt : str
            Name of the plotted quantity.
        phi : float, optional
            Toroidal coordinate where data is evaluated.
        t : float, optional
            Time coordinate where data is evaluated.
        clim : list [float, float], optional
            Minimum and maximum color range values.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.
        """
        log = False
        if "log" in qnt:
            qnt = qnt.replace("log", "").strip()
            log = True

        out = np.squeeze(self.input_eval(r, phi, z, t, qnt, grid=True)[:,0,:,0])
        diverging = np.nanmin(out)*np.nanmax(out) < 0 # If data has both -+ vals
        try:
            r = r.v
        except AttributeError:
            pass
        try:
            z = z.v
        except AttributeError:
            pass
        a5plt.mesh2d(r, z, out.v, diverging=diverging, log=log,
                     axesequal=True, xlabel="R [m]", ylabel="z [m]",
                     clabel=qnt + " [" + str(out.units) + "]", clim=clim,
                     cmap=cmap, axes=axes, cax=cax)

    def input_plotseparatrix(self, phi=0*unyt.deg, t=0*unyt.s,
                             rlim=[0.1*unyt.m, 20*unyt.m],
                             zlim=[-10*unyt.m, 10*unyt.m], axes=None):
        """Plot separatrix on (R,z) plane.
        """
        r = np.linspace(rlim[0], rlim[1], 100)
        z = np.linspace(zlim[0], zlim[1], 100)
        v = np.squeeze(self.input_eval(r, phi, z, t, "rho", grid=True)[:,0,:,0])
        a5plt.contour2d(r, z, v, [1], xlabel="R [m]", ylabel="z [m]",
                        colors="black", linestyles="solid", linewidths=1,
                        axesequal=True, axes=axes)

    def get_plasmaquantities(self):
        """Return species present in plasma input.
        """
        spec = self.get_plasmaspecies()
        for i in range(1,spec["nspecies"]):
            quantities.append("ni" + str(i))
            quantities.append("ti" + str(i))

        return quantities

    def input_rhovolume(self, nrho=10, ntheta=360, nphi=10, method="mc",
                        tol=1e-1, t=0*unyt.s, return_area=False,
                        return_coords=False):
        """Evaluate flux surface volumes.

        Parameters
        ----------
        nrho : int, optional
            Number of radial grid edges between [0, 1].
        ntheta : int, optional
            Number of poloidal grid edges.
        nphi : int, optional
            Number of toroidal grid edges.
        method : {"prism", "mc"}, optional
            Method used to evaluate the volumes.

            "prism" divides the volume into structured 3D grid consisting of
            triangular prims where the triangular plane is poloidal and one
            of the vertices is at magnetic axis while the other two are along
            a given flux surface. The volume is the estimated as a sum of these
            prisms.

            "mc" divides the volume into toroidal sectors of same size, and on
            each sector fins the separatrix bounding box. Bounding box is
            assumed to be fixed along the sector and markers are randomly
            sampled within it. The volumes are then estimated via the Monte
            Carlo method.
        tol : float, optional
            If relative difference in volume in subsequent iterations falls
            below this value, the algorithm finishes.
        t : float, optional
            Time slice when the volumes are computed.
        return_area : bool, optional
            Return also the area on (R, z) plane.
        return_coords : bool, optional
            Return also the (R, phi, z) coordinates for grid center points.

        Returns
        -------
        volume : array_like (nrho-1, ntheta-1, nphi-1)
            Volume of each bin.
        area : array_like (nrho-1, ntheta-1, nphi-1)
            Area of the poloidal cross section of each bin if
            ``return_area=True``.
        r : array_like (nrho-1, ntheta-1, nphi-1)
            R coordinates of bin centers if ``return_coords=True``.
        p : array_like (nrho-1, ntheta-1, nphi-1)
            phi coordinates of bin centers if ``return_coords=True``.
        z : array_like (nrho-1, ntheta-1, nphi-1)
            z coordinates of bin centers if ``return_coords=True``.
        """
        if nrho < 2 or ntheta < 2 or nphi < 2:
            raise ValueError(
                "Minimum number of edges is 2 but nrho=%d, ntheta=%d, nphi%d"
                + " were given", (nrho, ntheta, nphi))

        # Grid edges and area and volume of the grid cells
        rho    = np.linspace(0, 1,   nrho) * unyt.dimensionless
        phi    = np.linspace(0, 2*np.pi, nphi) * unyt.rad
        theta  = np.linspace(0, 2*np.pi, ntheta) * unyt.rad
        area   = np.zeros((nrho-1, ntheta-1, nphi-1)) * unyt.m**2
        volume = np.zeros((nrho-1, ntheta-1, nphi-1)) * unyt.m**3

        # Compute grid center points in (R,phi,z)
        rhoc, thc, p = np.meshgrid(
            ( rho[1:]   + rho[:-1]   ) / 2,
            ( theta[1:] + theta[:-1] ) / 2,
            ( phi[1:]   + phi[:-1]   ) / 2, indexing="ij")
        r, z = self.input_rhotheta2rz(
            rhoc.ravel(), thc.ravel(), p.ravel(), t)
        del thc, rhoc

        def integrate_prism(rho, theta, phimin, phimax):
            """Integrate volume by dividing it into prisms.
            """
            # Integration variable theta
            th = np.linspace(0, 2*np.pi, 1000)

            # Get axis location
            out = self._eval_bfield(
                1*unyt.m, (phimin + phimax) / 2, 1*unyt.m, t, evalaxis=True)
            r0 = out["axisr"]
            z0 = out["axisz"]

            # Get contour points for each rho value
            rhoc, thc, phic = np.meshgrid(
                rho[1:], th, ( phimin + phimax ) / 2, indexing="ij" )
            rc, zc = self.input_rhotheta2rz(
                rhoc.ravel(), thc.ravel()*unyt.rad, phic.ravel(), t)
            rc = rc.reshape(rhoc.shape)[:,:,0]
            zc = zc.reshape(rhoc.shape)[:,:,0]

            # Calculate the triangle areas and triangle centroid
            area = 0.5 * np.abs(  r0        * ( zc[:,:-1] - zc[:,1:]  )
                                + rc[:,:-1] * ( zc[:,1:]  - z0        )
                                + rc[:,1:]  * ( z0        - zc[:,:-1] ) )
            h       = (phimax - phimin).v
            centerr = (r0 + rc[:,:-1] + rc[:,1:]) / 3
            volume  = h * centerr * area

            # Now divide these to theta sectors that were requested
            V = np.zeros((rho.size-1, theta.size-1)) * unyt.m**3
            A = np.zeros((rho.size-1, theta.size-1)) * unyt.m**2
            th = th[:-1] # Drop last index as it is same as first
            for it in range(theta.size-1):
                idx = np.logical_and(th >= theta[it], th < theta[it+1])
                V[:,it] = np.sum(volume[:, idx], axis=1)
                A[:,it] = np.sum(area[:, idx], axis=1)

            V[1:,:] -= V[:-1,:]
            A[1:,:] -= A[:-1,:]
            return V, A

        def integrate_mc(rho, theta, phimin, phimax):
            """Integrate volume by dividing it into prisms.
            """
            # Determine the bounding box
            th = np.linspace(0, 2*np.pi, 1000)
            rhoc, thc, phic = np.meshgrid(1.0*unyt.dimensionless, th,
                                          np.append(phimin, phimax),
                                          indexing="ij" )
            rc, zc = self.input_rhotheta2rz(
                rhoc.ravel(), thc.ravel()*unyt.rad, phic.ravel()*unyt.rad, t)
            bbox = np.ones((4,)) * unyt.m
            bbox[0] = np.nanmin(rc)
            bbox[1] = np.nanmax(rc)
            bbox[2] = np.nanmin(zc)
            bbox[3] = np.nanmax(zc)

            # Get axis location
            out = self._eval_bfield(
                1*unyt.m, (phimin + phimax) / 2, 1*unyt.m, t, evalaxis=True)
            r0 = out["axisr"]
            z0 = out["axisz"]

            # Number of markers drawn per iteration, total number of markers
            # drawn and area and volume of the test region
            ndraw  = 10000
            ntot   = 0
            atotal = (bbox[1] - bbox[0]) * (bbox[3] - bbox[2])
            vtotal = 0.5 * (bbox[1]**2 - bbox[0]**2) \
                * ( bbox[3] - bbox[2] ) * ( phimax - phimin ).v
            deltarho = 1.0 / (rho.size-1)
            deltath  = 2*np.pi / (theta.size-1)

            v0     = np.zeros((rho.size-1, theta.size-1)) * unyt.m**3
            area   = np.zeros((rho.size-1, theta.size-1)) * unyt.m**2
            volume = np.zeros((rho.size-1, theta.size-1)) * unyt.m**3
            points_in_bins = np.zeros((rho.size-1, theta.size-1))
            while np.sum(v0) == 0 or \
                  np.max(np.abs( volume[v0>0] - v0[v0>0] ) / v0[v0>0])  > tol:
                v0[:,:] = volume
                pi = np.random.uniform(phimin, phimax,   ndraw) * unyt.rad
                ri = np.random.uniform(bbox[0], bbox[1], ndraw) * unyt.m
                zi = np.random.uniform(bbox[2], bbox[3], ndraw) * unyt.m
                th = np.arctan2(zi-z0,ri-r0) + np.pi

                rhoi    = self.input_eval(ri, pi, zi, t, "rho").v
                ind     = rhoi <= 1 # Reject markers outside separatrix
                i_rho   = np.floor(rhoi[ind] / deltarho).astype(int)
                i_theta = np.floor(th[ind] / deltath).astype(int)

                np.add.at(points_in_bins, (i_rho, i_theta), 1)
                ntot += ndraw

                area[:,:]   = atotal * points_in_bins / ntot
                volume[:,:] = vtotal * points_in_bins / ntot

            return volume, area

        if method == "prism":
            integrate = integrate_prism
        elif method == "mc":
            integrate = integrate_mc

        for ip in range(nphi-1):
            volume[:,:,ip], area[:,:,ip] = \
                integrate(rho, theta, phi[ip], phi[ip+1])
            dphi = (phi[ip+1] - phi[ip]) / 2
            v, _   = integrate(rho, theta, phi[ip], phi[ip] + dphi)
            v1, _  = integrate(rho, theta, phi[ip] + dphi, phi[ip] + dphi*2)
            v += v1
            while(np.abs(np.sum(v) - np.sum(volume[:,:,ip])) > tol):
                volume[:,:,ip] = v
                dphi /= 2
                v[:] = 0 * unyt.m**3
                for i in range(int((phi[ip+1] - phi[ip])/dphi)):
                    v1, _ = integrate(rho, theta, phi[ip] + dphi*i,
                                      phi[ip] + dphi*(i+1))
                    v += v1
            volume[:,:,ip] = v

        if np.sum(volume == 0):
            warnings.warn(
                "Volume calculation resulted in one or more elements with zero"
              + " volume. This might happen when psi0 differs from the actual"
              + " value.")

        if return_area and return_coords:
            return volume, area, r, p, z
        elif return_area:
            return volume, area
        elif return_coords:
            return volume, r, p, z
        else:
            return volume

    def input_eval_safetyfactor(self, rho, nth=10000):
        """Evaluate safety factor and associated quantities.

        This function most likely works for tokamaks only.

        Parameters
        ----------
        rho : array_like, (n,)
            Radial positions where the safety factor is evaluated.

            Note that evaluating it near the axis or exactly at the separatrix
            may yield mess.

        Returns
        -------
        qprof : array_like, (n,)
            Safety factor.
        Iprof : array_like, (n,)
            Toroidal current term which, when multiplied with mu0/2, gives the
            enclosed toroidal current.
        gprof : array_like, (n,)
            This is just R * Bphi which is constant since Bphi ~ 1/R.
        """
        qprof = np.zeros(rho.shape)
        Iprof = np.zeros(rho.shape)
        gprof = np.zeros(rho.shape)
        thgrid = np.linspace(0, 2*np.pi, nth)
        for i in range(rho.size):
            # Interpolate the contour points on the fixed theta grid
            r, z = self.input_rhotheta2rz(
                rho[i]*np.ones(thgrid.shape), thgrid*unyt.rad,
                np.zeros(thgrid.shape)*unyt.rad, 0*unyt.s)

            # Magnetic field along the contour (psi can be used to check that
            # the contour was set properly). Drop the last element in r and z
            # as it is the same as first.
            br, bphi, bz, psi = self.input_eval(
                r[:-1], 0*unyt.rad, z[:-1], 0*unyt.s,
                "br", "bphi", "bz", "psi")

            bpol  = np.sqrt(br**2 + bz**2)
            bnorm = np.sqrt(br**2 + bphi**2 + bz**2)
            ds    = (np.diff(r) * br + np.diff(z) * bz) / bpol # darc dot e^_pol
            r = r[:-1]; z = z[:-1]

            # The toroidal current term (multiplying this with mu0/2 pi gets
            # enclosed toroidal current)
            Iprof[i] = np.sum( ds * bpol ) / ( 2*np.pi )

            # g = R*Bphi, since Bphi ~ 1/R this is a constant
            gprof[i] = r[0] * bphi[0]

            # The (global) safety factor q(psi)
            qprof[i] = np.sum( ds * gprof[i] / ( r**2 * bpol ) ) / ( 2*np.pi )

        return qprof, Iprof, gprof

    def input_eval_orbitboundaries(self, mugrid, ptorgrid, ekin):
        """Plot boundaries of different orbit topologies in (mu, ekin)
        phase-space when energy is fixed.

        Parameters
        ----------
        mugrid : array_like
            Grid for the magnetic moment abscissa.
        ptorgrid : array_like
            Grid for the the canonical toroidal angular momentum abscissa.
        ekin : float
            Particle energy.
        """
        pass

    def input_eval_ripple(
            self, rgrid, zgrid, rlarmor, rho0=1.0*unyt.dimensionless,
            theta0=0*unyt.rad, t=0*unyt.s, nphi=362, plot=True, axes1=None,
            axes2=None):
        """Evaluate various ripple quantities.

        The variation of toroidal field strength as a function of toroidal angle
        is called the toroidal field (TF) ripple. The TF ripple can lead to
        ripple-induced losses provided that i) the particle is trapped in
        a ripple well in which case it is promptly lost or ii) the ripple causes
        particle motion to become stochastic which leads to radial diffusion.

        This function evaluates the quantities relevant for ripple transport.
        Note that ``deltacrit`` and ``ripplewell`` are approximate quantities
        based on analytical estimates.

        Parameters
        ----------
        rgrid : array_like
            R points where ``delta``, ``deltacrit``, and ``ripplewell`` are
            evaluated.
        zgrid : array_like
            z points where ``delta``, ``deltacrit``, and ``ripplewell`` are
            evaluated.
        rlarmor : float or tuple (str, int, float)
            Larmor radius of the test particle or tuple with
            (species, charge, energy).
        rho0 : float, optional
            Radial coordinate where ``amplitude`` is evaluated.
        theta0 : float, optional
            Poloidal coordinate where ``amplitude`` is evaluated.
        t : float, optional
            Time instance for time-dependent fields.
        nphi : int, optional
            Number of toroidal grid points in ripple evaluation.

            There is probably never a need to change this number.
        plot : bool, optional
            If True, the evaluated quantities are also plotted.
        axes1 : , optional
            Axes for plotting ``amplitude`` if ``plot``=True.
        axes2 : , optional
            Axes for plotting ``delta``, ``deltacrit``, and ``ripplewell`` if
            ``plot``=True.

        Returns
        -------
        phigrid : array_like, (nphi)
            Toroidal grid [-2, 362] where ``amplitude`` is given.

            This grid is not limited to [0,360] as usual so that one can also
            use it to spot any discontinuities in field if the input was poorly
            constructed.
        amplitude : array_like, (nphi,)
            Variation in toroidal field strength, i.e.
            Bphi(phi) - Bphi_axisymmetric.
        delta : array_like, (nr,nz)
            Ripple magnitude defined as  (Bmax-Bmin) / (Bmax+Bmin), where Bmax
            and Bmin are toroidal extrema of Bphi.
        deltacrit : array_like, (nr,nz)
            Critical ripple magnitude.

            Particles are subject to stochastic ripple diffusion when
            ``deltacrit`` < 1.
        ripplewell : array_like, (nr,nz)
            Ripple well parameter.

            Particles are bound to be ripple-trapped when ``ripplewell`` < 1.
            However, this does not necessarily mean that the particle is lost as
            it may escape the well when ripple-trapping causes it to drift
            outwards. As particles that are ripple-trapped drift vertically,
            a good indication that a trapped particle will be lost if the ripple
            well region covers everything vertically below/above that point.
        """
        phi = np.linspace(-2, 362, nphi) * unyt.deg
        bphi, rho, drhodpsi = self.input_eval(
            rgrid, phi, zgrid, 0*unyt.s, "bphi", "rho", "rhodpsi", grid=True)
        bphi = np.squeeze(np.abs(bphi))
        bmax = np.nanmax(bphi, axis=1)
        bmin = np.nanmin(bphi, axis=1)
        delta = (bmax - bmin) / (bmax + bmin)
        rho = np.squeeze(rho[:,0,:,0])
        drhodpsi = np.squeeze(drhodpsi[:,0,:,0])

        if not isinstance(rlarmor, float):
            prt   = physlib.species.species(rlarmor[0], charge=rlarmor[1])
            ekin  = rlarmor[2] * unyt.eV
            bnorm = np.mean(bphi)
            rlarmor = physlib.gyrolength(
                prt["mass"], prt["charge"], ekin, 0.0, bnorm)

        r0, z0 = self.input_rhotheta2rz(rho0*np.ones((nphi,)),
                                        theta0*np.ones((nphi,)), phi, 0*unyt.s)
        amplitude = self.input_eval(r0, phi, z0, 0*unyt.s, "bphi")
        amplitude -= np.mean(amplitude)

        # Use fft to find ripple periodicity
        ff      = np.fft.fftfreq(len(phi), (phi[1]-phi[0]))
        Fyy     = abs(np.fft.fft(amplitude))
        Nripple = np.round(abs(ff[np.argmax(Fyy[1:])+1]) * (phi[-1] - phi[0]))

        rhogrid = np.linspace(0, 1, 100)
        q, _, _ = self.input_eval_safetyfactor(rhogrid)
        q = scipy.interpolate.CubicSpline(rhogrid, q, extrapolate=False)

        raxis, zaxis = self.input_rhotheta2rz(0.0*unyt.dimensionless, 0.0, 0.0, 0.0)

        rg, zg = np.meshgrid(rgrid, zgrid, indexing="ij")
        rminor = np.sqrt( (rg - raxis)**2 + (zg - zaxis)**2 )
        sinth  = (zg - zaxis) / rminor
        eps    = rminor / raxis
        qfac   = np.abs(q(rho))
        dqdpsi = np.abs(q(rho,1) * drhodpsi)

        deltacrit = np.sqrt( eps / (np.pi * Nripple * qfac) )**3 \
            / ( rlarmor * dqdpsi * delta )
        ripplewell = rg * np.abs(sinth) / ( Nripple * qfac * delta )

        if plot:
            axes1.plot(phi, amplitude)
            axes2.contour(rgrid, zgrid, deltacrit.T.v, np.array([1.0]), colors="C0")
            axes2.contour(rgrid, zgrid, ripplewell.T.v, np.array([1.0]), colors="C1")
            cs = axes2.contour(rgrid, zgrid, delta.T, colors="black")
            axes2.clabel(cs, cs.levels)

            axes2.set_aspect("equal", adjustable="box")

        return phi, amplitude, delta, deltacrit, ripplewell

    def input_eval_mhdalphafromphi(self, nrho=100,
                                   reverse=False, plot=True):
        """Evaluate magnetic potential from electric potential assuming that
        Epar = 0.

        This function can be used to verify that the MHD modes were porperly
        imported. Note the Epar=0 condition though.

        alphanm' = (n*q - m) * phinm / (g*q + I)

        Parameters
        ----------
        rgrid : array_like, (nr,)
        zgrid : array_like, (nz,)
        phi : float, optional
        """
        n, nmode, mmode, _, omega, _ = self.input_mhd_modes()
        rhogrid = np.linspace(0.05, 0.95, nrho)
        q, I, g = self.input_eval_safetyfactor(rhogrid)

        r, z = self.input_rhotheta2rz(rhogrid, rhogrid, rhogrid*0, 0*unyt.s)

        import matplotlib.pyplot as plt
        alphanm  = np.zeros((n, nrho)) + 1
        phinm    = np.zeros((n, nrho))
        alphanm_ = np.zeros((n, nrho)) + 2
        for i in range(n):
            out = self._eval_mhd(r, r, z, r*0*0, evalpot=True)
            alphanm[i,:]  = out["alphaeig"]
            phinm[i,:]    = out["phieig"]
            alphanm_[i,:] = (2*1.256e-6*np.pi) * (( nmode[i] * q - mmode[i] ) * phinm[i,:] / (g*q + I)) / omega[i]
            print(nmode[i], mmode[i], np.amax(alphanm[i,:]), np.amax(alphanm_[i,:]))

            plt.plot(alphanm[i,:])
            plt.plot(alphanm_[i,:], linestyle="--")
        plt.show()

a5py/a5py/ascotpy/__init__.py
ascot4fusion/ascot5
# -*- coding: utf-8 -*-
#
# TARGET arch is: ['-I/usr/include/hdf5/serial']
# WORD_SIZE is: 8
# POINTER_SIZE is: 8
# LONGDOUBLE_SIZE is: 16
#
import ctypes


class AsDictMixin:
    @classmethod
    def as_dict(cls, self):
        result = {}
        if not isinstance(self, AsDictMixin):
            # not a structure, assume it's already a python object
            return self
        if not hasattr(cls, "_fields_"):
            return result
        # sys.version_info >= (3, 5)
        # for (field, *_) in cls._fields_:  # noqa
        for field_tuple in cls._fields_:  # noqa
            field = field_tuple[0]
            if field.startswith('PADDING_'):
                continue
            value = getattr(self, field)
            type_ = type(value)
            if hasattr(value, "_length_") and hasattr(value, "_type_"):
                # array
                if not hasattr(type_, "as_dict"):
                    value = [v for v in value]
                else:
                    type_ = type_._type_
                    value = [type_.as_dict(v) for v in value]
            elif hasattr(value, "contents") and hasattr(value, "_type_"):
                # pointer
                try:
                    if not hasattr(type_, "as_dict"):
                        value = value.contents
                    else:
                        type_ = type_._type_
                        value = type_.as_dict(value.contents)
                except ValueError:
                    # nullptr
                    value = None
            elif isinstance(value, AsDictMixin):
                # other structure
                value = type_.as_dict(value)
            result[field] = value
        return result


class Structure(ctypes.Structure, AsDictMixin):

    def __init__(self, *args, **kwds):
        # We don't want to use positional arguments fill PADDING_* fields

        args = dict(zip(self.__class__._field_names_(), args))
        args.update(kwds)
        super(Structure, self).__init__(**args)

    @classmethod
    def _field_names_(cls):
        if hasattr(cls, '_fields_'):
            return (f[0] for f in cls._fields_ if not f[0].startswith('PADDING'))
        else:
            return ()

    @classmethod
    def get_type(cls, field):
        for f in cls._fields_:
            if f[0] == field:
                return f[1]
        return None

    @classmethod
    def bind(cls, bound_fields):
        fields = {}
        for name, type_ in cls._fields_:
            if hasattr(type_, "restype"):
                if name in bound_fields:
                    if bound_fields[name] is None:
                        fields[name] = type_()
                    else:
                        # use a closure to capture the callback from the loop scope
                        fields[name] = (
                            type_((lambda callback: lambda *args: callback(*args))(
                                bound_fields[name]))
                        )
                    del bound_fields[name]
                else:
                    # default callback implementation (does nothing)
                    try:
                        default_ = type_(0).restype().value
                    except TypeError:
                        default_ = None
                    fields[name] = type_((
                        lambda default_: lambda *args: default_)(default_))
            else:
                # not a callback function, use default initialization
                if name in bound_fields:
                    fields[name] = bound_fields[name]
                    del bound_fields[name]
                else:
                    fields[name] = type_()
        if len(bound_fields) != 0:
            raise ValueError(
                "Cannot bind the following unknown callback(s) {}.{}".format(
                    cls.__name__, bound_fields.keys()
            ))
        return cls(**fields)


class Union(ctypes.Union, AsDictMixin):
    pass



_libraries = {}
_libraries['libascot.so'] = ctypes.CDLL('libascot.so')
c_int128 = ctypes.c_ubyte*16
c_uint128 = c_int128
void = None
if ctypes.sizeof(ctypes.c_longdouble) == 16:
    c_long_double_t = ctypes.c_longdouble
else:
    c_long_double_t = ctypes.c_ubyte*16

def string_cast(char_pointer, encoding='utf-8', errors='strict'):
    value = ctypes.cast(char_pointer, ctypes.c_char_p).value
    if value is not None and encoding is not None:
        value = value.decode(encoding, errors=errors)
    return value


def char_pointer_cast(string, encoding='utf-8'):
    if encoding is not None:
        try:
            string = string.encode(encoding)
        except AttributeError:
            # In Python3, bytes has no encode attribute
            pass
    string = ctypes.c_char_p(string)
    return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))



class FunctionFactoryStub:
    def __getattr__(self, _):
      return ctypes.CFUNCTYPE(lambda y:y)

# libraries['FIXME_STUB'] explanation
# As you did not list (-l libraryname.so) a library that exports this function
# This is a non-working stub instead. 
# You can either re-run clan2py with -l /path/to/library.so
# Or manually fix this by comment the ctypes.CDLL loading
_libraries['FIXME_STUB'] = FunctionFactoryStub() #  ctypes.CDLL('FIXME_STUB')


integer = ctypes.c_int64
real = ctypes.c_double
class struct_c__SA_B_STS_offload_data(Structure):
    pass

struct_c__SA_B_STS_offload_data._pack_ = 1 # source:False
struct_c__SA_B_STS_offload_data._fields_ = [
    ('psigrid_n_r', ctypes.c_int32),
    ('psigrid_n_z', ctypes.c_int32),
    ('psigrid_n_phi', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('psigrid_r_min', ctypes.c_double),
    ('psigrid_r_max', ctypes.c_double),
    ('psigrid_z_min', ctypes.c_double),
    ('psigrid_z_max', ctypes.c_double),
    ('psigrid_phi_min', ctypes.c_double),
    ('psigrid_phi_max', ctypes.c_double),
    ('Bgrid_n_r', ctypes.c_int32),
    ('Bgrid_n_z', ctypes.c_int32),
    ('Bgrid_n_phi', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('Bgrid_r_min', ctypes.c_double),
    ('Bgrid_r_max', ctypes.c_double),
    ('Bgrid_z_min', ctypes.c_double),
    ('Bgrid_z_max', ctypes.c_double),
    ('Bgrid_phi_min', ctypes.c_double),
    ('Bgrid_phi_max', ctypes.c_double),
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('offload_array_length', ctypes.c_int32),
    ('n_axis', ctypes.c_int32),
    ('axis_min', ctypes.c_double),
    ('axis_max', ctypes.c_double),
    ('axis_grid', ctypes.c_double),
]

B_STS_offload_data = struct_c__SA_B_STS_offload_data
class struct_c__SA_B_STS_data(Structure):
    pass

class struct_c__SA_interp3D_data(Structure):
    pass

struct_c__SA_interp3D_data._pack_ = 1 # source:False
struct_c__SA_interp3D_data._fields_ = [
    ('n_x', ctypes.c_int32),
    ('n_y', ctypes.c_int32),
    ('n_z', ctypes.c_int32),
    ('bc_x', ctypes.c_int32),
    ('bc_y', ctypes.c_int32),
    ('bc_z', ctypes.c_int32),
    ('x_min', ctypes.c_double),
    ('x_max', ctypes.c_double),
    ('x_grid', ctypes.c_double),
    ('y_min', ctypes.c_double),
    ('y_max', ctypes.c_double),
    ('y_grid', ctypes.c_double),
    ('z_min', ctypes.c_double),
    ('z_max', ctypes.c_double),
    ('z_grid', ctypes.c_double),
    ('c', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_linint1D_data(Structure):
    pass

struct_c__SA_linint1D_data._pack_ = 1 # source:False
struct_c__SA_linint1D_data._fields_ = [
    ('n_x', ctypes.c_int32),
    ('bc_x', ctypes.c_int32),
    ('x_min', ctypes.c_double),
    ('x_max', ctypes.c_double),
    ('x_grid', ctypes.c_double),
    ('c', ctypes.POINTER(ctypes.c_double)),
]

struct_c__SA_B_STS_data._pack_ = 1 # source:False
struct_c__SA_B_STS_data._fields_ = [
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('axis_r', struct_c__SA_linint1D_data),
    ('axis_z', struct_c__SA_linint1D_data),
    ('psi', struct_c__SA_interp3D_data),
    ('B_r', struct_c__SA_interp3D_data),
    ('B_phi', struct_c__SA_interp3D_data),
    ('B_z', struct_c__SA_interp3D_data),
]

B_STS_data = struct_c__SA_B_STS_data
B_STS_init_offload = _libraries['libascot.so'].B_STS_init_offload
B_STS_init_offload.restype = ctypes.c_int32
B_STS_init_offload.argtypes = [ctypes.POINTER(struct_c__SA_B_STS_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]
B_STS_free_offload = _libraries['libascot.so'].B_STS_free_offload
B_STS_free_offload.restype = None
B_STS_free_offload.argtypes = [ctypes.POINTER(struct_c__SA_B_STS_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]
B_STS_init = _libraries['libascot.so'].B_STS_init
B_STS_init.restype = None
B_STS_init.argtypes = [ctypes.POINTER(struct_c__SA_B_STS_data), ctypes.POINTER(struct_c__SA_B_STS_offload_data), ctypes.POINTER(ctypes.c_double)]
a5err = ctypes.c_uint64
B_STS_eval_psi = _libraries['libascot.so'].B_STS_eval_psi
B_STS_eval_psi.restype = a5err
B_STS_eval_psi.argtypes = [ctypes.POINTER(ctypes.c_double), real, real, real, ctypes.POINTER(struct_c__SA_B_STS_data)]
B_STS_eval_psi_dpsi = _libraries['libascot.so'].B_STS_eval_psi_dpsi
B_STS_eval_psi_dpsi.restype = a5err
B_STS_eval_psi_dpsi.argtypes = [ctypes.c_double * 4, real, real, real, ctypes.POINTER(struct_c__SA_B_STS_data)]
B_STS_eval_rho_drho = _libraries['libascot.so'].B_STS_eval_rho_drho
B_STS_eval_rho_drho.restype = a5err
B_STS_eval_rho_drho.argtypes = [ctypes.c_double * 4, real, real, real, ctypes.POINTER(struct_c__SA_B_STS_data)]
B_STS_eval_B = _libraries['libascot.so'].B_STS_eval_B
B_STS_eval_B.restype = a5err
B_STS_eval_B.argtypes = [ctypes.c_double * 3, real, real, real, ctypes.POINTER(struct_c__SA_B_STS_data)]
B_STS_eval_B_dB = _libraries['libascot.so'].B_STS_eval_B_dB
B_STS_eval_B_dB.restype = a5err
B_STS_eval_B_dB.argtypes = [ctypes.c_double * 12, real, real, real, ctypes.POINTER(struct_c__SA_B_STS_data)]
B_STS_get_axis_rz = _libraries['libascot.so'].B_STS_get_axis_rz
B_STS_get_axis_rz.restype = a5err
B_STS_get_axis_rz.argtypes = [ctypes.c_double * 2, ctypes.POINTER(struct_c__SA_B_STS_data), real]

# values for enumeration 'B_field_type'
B_field_type__enumvalues = {
    0: 'B_field_type_GS',
    1: 'B_field_type_2DS',
    2: 'B_field_type_3DS',
    3: 'B_field_type_STS',
    4: 'B_field_type_TC',
}
B_field_type_GS = 0
B_field_type_2DS = 1
B_field_type_3DS = 2
B_field_type_STS = 3
B_field_type_TC = 4
B_field_type = ctypes.c_uint32 # enum
class struct_c__SA_B_field_offload_data(Structure):
    pass

class struct_c__SA_B_GS_offload_data(Structure):
    pass

struct_c__SA_B_GS_offload_data._pack_ = 1 # source:False
struct_c__SA_B_GS_offload_data._fields_ = [
    ('R0', ctypes.c_double),
    ('z0', ctypes.c_double),
    ('raxis', ctypes.c_double),
    ('zaxis', ctypes.c_double),
    ('B_phi0', ctypes.c_double),
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('psi_mult', ctypes.c_double),
    ('psi_coeff', ctypes.c_double * 13),
    ('Nripple', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('a0', ctypes.c_double),
    ('alpha0', ctypes.c_double),
    ('delta0', ctypes.c_double),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

class struct_c__SA_B_TC_offload_data(Structure):
    pass

struct_c__SA_B_TC_offload_data._pack_ = 1 # source:False
struct_c__SA_B_TC_offload_data._fields_ = [
    ('axisr', ctypes.c_double),
    ('axisz', ctypes.c_double),
    ('psival', ctypes.c_double),
    ('rhoval', ctypes.c_double),
    ('B', ctypes.c_double * 3),
    ('dB', ctypes.c_double * 9),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
]

class struct_c__SA_B_2DS_offload_data(Structure):
    pass

struct_c__SA_B_2DS_offload_data._pack_ = 1 # source:False
struct_c__SA_B_2DS_offload_data._fields_ = [
    ('n_r', ctypes.c_int32),
    ('n_z', ctypes.c_int32),
    ('r_min', ctypes.c_double),
    ('r_max', ctypes.c_double),
    ('z_min', ctypes.c_double),
    ('z_max', ctypes.c_double),
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('axis_r', ctypes.c_double),
    ('axis_z', ctypes.c_double),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
]

class struct_c__SA_B_3DS_offload_data(Structure):
    pass

struct_c__SA_B_3DS_offload_data._pack_ = 1 # source:False
struct_c__SA_B_3DS_offload_data._fields_ = [
    ('psigrid_n_r', ctypes.c_int32),
    ('psigrid_n_z', ctypes.c_int32),
    ('psigrid_r_min', ctypes.c_double),
    ('psigrid_r_max', ctypes.c_double),
    ('psigrid_z_min', ctypes.c_double),
    ('psigrid_z_max', ctypes.c_double),
    ('Bgrid_n_r', ctypes.c_int32),
    ('Bgrid_n_z', ctypes.c_int32),
    ('Bgrid_r_min', ctypes.c_double),
    ('Bgrid_r_max', ctypes.c_double),
    ('Bgrid_z_min', ctypes.c_double),
    ('Bgrid_z_max', ctypes.c_double),
    ('Bgrid_n_phi', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('Bgrid_phi_min', ctypes.c_double),
    ('Bgrid_phi_max', ctypes.c_double),
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('axis_r', ctypes.c_double),
    ('axis_z', ctypes.c_double),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

struct_c__SA_B_field_offload_data._pack_ = 1 # source:False
struct_c__SA_B_field_offload_data._fields_ = [
    ('type', B_field_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('BGS', struct_c__SA_B_GS_offload_data),
    ('B2DS', struct_c__SA_B_2DS_offload_data),
    ('B3DS', struct_c__SA_B_3DS_offload_data),
    ('BSTS', B_STS_offload_data),
    ('BTC', struct_c__SA_B_TC_offload_data),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

B_field_offload_data = struct_c__SA_B_field_offload_data
class struct_c__SA_B_field_data(Structure):
    pass

class struct_c__SA_B_3DS_data(Structure):
    pass

class struct_c__SA_interp2D_data(Structure):
    pass

struct_c__SA_interp2D_data._pack_ = 1 # source:False
struct_c__SA_interp2D_data._fields_ = [
    ('n_x', ctypes.c_int32),
    ('n_y', ctypes.c_int32),
    ('bc_x', ctypes.c_int32),
    ('bc_y', ctypes.c_int32),
    ('x_min', ctypes.c_double),
    ('x_max', ctypes.c_double),
    ('x_grid', ctypes.c_double),
    ('y_min', ctypes.c_double),
    ('y_max', ctypes.c_double),
    ('y_grid', ctypes.c_double),
    ('c', ctypes.POINTER(ctypes.c_double)),
]

struct_c__SA_B_3DS_data._pack_ = 1 # source:False
struct_c__SA_B_3DS_data._fields_ = [
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('axis_r', ctypes.c_double),
    ('axis_z', ctypes.c_double),
    ('psi', struct_c__SA_interp2D_data),
    ('B_r', struct_c__SA_interp3D_data),
    ('B_phi', struct_c__SA_interp3D_data),
    ('B_z', struct_c__SA_interp3D_data),
]

class struct_c__SA_B_2DS_data(Structure):
    pass

struct_c__SA_B_2DS_data._pack_ = 1 # source:False
struct_c__SA_B_2DS_data._fields_ = [
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('axis_r', ctypes.c_double),
    ('axis_z', ctypes.c_double),
    ('psi', struct_c__SA_interp2D_data),
    ('B_r', struct_c__SA_interp2D_data),
    ('B_phi', struct_c__SA_interp2D_data),
    ('B_z', struct_c__SA_interp2D_data),
]

class struct_c__SA_B_TC_data(Structure):
    pass

struct_c__SA_B_TC_data._pack_ = 1 # source:False
struct_c__SA_B_TC_data._fields_ = [
    ('axisr', ctypes.c_double),
    ('axisz', ctypes.c_double),
    ('psival', ctypes.c_double),
    ('rhoval', ctypes.c_double),
    ('B', ctypes.POINTER(ctypes.c_double)),
    ('dB', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_B_GS_data(Structure):
    pass

struct_c__SA_B_GS_data._pack_ = 1 # source:False
struct_c__SA_B_GS_data._fields_ = [
    ('R0', ctypes.c_double),
    ('z0', ctypes.c_double),
    ('raxis', ctypes.c_double),
    ('zaxis', ctypes.c_double),
    ('B_phi0', ctypes.c_double),
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('psi_mult', ctypes.c_double),
    ('psi_coeff', ctypes.POINTER(ctypes.c_double)),
    ('Nripple', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('a0', ctypes.c_double),
    ('alpha0', ctypes.c_double),
    ('delta0', ctypes.c_double),
]

struct_c__SA_B_field_data._pack_ = 1 # source:False
struct_c__SA_B_field_data._fields_ = [
    ('type', B_field_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('BGS', struct_c__SA_B_GS_data),
    ('B2DS', struct_c__SA_B_2DS_data),
    ('B3DS', struct_c__SA_B_3DS_data),
    ('BSTS', B_STS_data),
    ('BTC', struct_c__SA_B_TC_data),
]

B_field_data = struct_c__SA_B_field_data
B_field_init_offload = _libraries['libascot.so'].B_field_init_offload
B_field_init_offload.restype = ctypes.c_int32
B_field_init_offload.argtypes = [ctypes.POINTER(struct_c__SA_B_field_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]
B_field_free_offload = _libraries['libascot.so'].B_field_free_offload
B_field_free_offload.restype = None
B_field_free_offload.argtypes = [ctypes.POINTER(struct_c__SA_B_field_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]
B_field_init = _libraries['libascot.so'].B_field_init
B_field_init.restype = ctypes.c_int32
B_field_init.argtypes = [ctypes.POINTER(struct_c__SA_B_field_data), ctypes.POINTER(struct_c__SA_B_field_offload_data), ctypes.POINTER(ctypes.c_double)]
B_field_eval_psi = _libraries['libascot.so'].B_field_eval_psi
B_field_eval_psi.restype = a5err
B_field_eval_psi.argtypes = [ctypes.POINTER(ctypes.c_double), real, real, real, real, ctypes.POINTER(struct_c__SA_B_field_data)]
B_field_eval_psi_dpsi = _libraries['libascot.so'].B_field_eval_psi_dpsi
B_field_eval_psi_dpsi.restype = a5err
B_field_eval_psi_dpsi.argtypes = [ctypes.c_double * 4, real, real, real, real, ctypes.POINTER(struct_c__SA_B_field_data)]
B_field_eval_rho = _libraries['libascot.so'].B_field_eval_rho
B_field_eval_rho.restype = a5err
B_field_eval_rho.argtypes = [ctypes.c_double * 2, real, ctypes.POINTER(struct_c__SA_B_field_data)]
B_field_eval_rho_drho = _libraries['libascot.so'].B_field_eval_rho_drho
B_field_eval_rho_drho.restype = a5err
B_field_eval_rho_drho.argtypes = [ctypes.c_double * 4, real, real, real, ctypes.POINTER(struct_c__SA_B_field_data)]
B_field_eval_B = _libraries['libascot.so'].B_field_eval_B
B_field_eval_B.restype = a5err
B_field_eval_B.argtypes = [ctypes.c_double * 3, real, real, real, real, ctypes.POINTER(struct_c__SA_B_field_data)]
B_field_eval_B_dB = _libraries['libascot.so'].B_field_eval_B_dB
B_field_eval_B_dB.restype = a5err
B_field_eval_B_dB.argtypes = [ctypes.c_double * 15, real, real, real, real, ctypes.POINTER(struct_c__SA_B_field_data)]
B_field_get_axis_rz = _libraries['libascot.so'].B_field_get_axis_rz
B_field_get_axis_rz.restype = a5err
B_field_get_axis_rz.argtypes = [ctypes.c_double * 2, ctypes.POINTER(struct_c__SA_B_field_data), real]
class struct_c__SA_particle_state(Structure):
    pass

struct_c__SA_particle_state._pack_ = 1 # source:False
struct_c__SA_particle_state._fields_ = [
    ('r', ctypes.c_double),
    ('phi', ctypes.c_double),
    ('z', ctypes.c_double),
    ('ppar', ctypes.c_double),
    ('mu', ctypes.c_double),
    ('zeta', ctypes.c_double),
    ('rprt', ctypes.c_double),
    ('phiprt', ctypes.c_double),
    ('zprt', ctypes.c_double),
    ('p_r', ctypes.c_double),
    ('p_phi', ctypes.c_double),
    ('p_z', ctypes.c_double),
    ('mass', ctypes.c_double),
    ('charge', ctypes.c_double),
    ('anum', ctypes.c_int32),
    ('znum', ctypes.c_int32),
    ('weight', ctypes.c_double),
    ('time', ctypes.c_double),
    ('mileage', ctypes.c_double),
    ('cputime', ctypes.c_double),
    ('rho', ctypes.c_double),
    ('theta', ctypes.c_double),
    ('id', ctypes.c_int64),
    ('endcond', ctypes.c_int64),
    ('walltile', ctypes.c_int64),
    ('B_r', ctypes.c_double),
    ('B_phi', ctypes.c_double),
    ('B_z', ctypes.c_double),
    ('B_r_dr', ctypes.c_double),
    ('B_phi_dr', ctypes.c_double),
    ('B_z_dr', ctypes.c_double),
    ('B_r_dphi', ctypes.c_double),
    ('B_phi_dphi', ctypes.c_double),
    ('B_z_dphi', ctypes.c_double),
    ('B_r_dz', ctypes.c_double),
    ('B_phi_dz', ctypes.c_double),
    ('B_z_dz', ctypes.c_double),
    ('err', ctypes.c_uint64),
]

particle_state = struct_c__SA_particle_state
class struct_c__SA_particle(Structure):
    pass

struct_c__SA_particle._pack_ = 1 # source:False
struct_c__SA_particle._fields_ = [
    ('r', ctypes.c_double),
    ('phi', ctypes.c_double),
    ('z', ctypes.c_double),
    ('p_r', ctypes.c_double),
    ('p_phi', ctypes.c_double),
    ('p_z', ctypes.c_double),
    ('mass', ctypes.c_double),
    ('charge', ctypes.c_double),
    ('anum', ctypes.c_int32),
    ('znum', ctypes.c_int32),
    ('weight', ctypes.c_double),
    ('time', ctypes.c_double),
    ('mileage', ctypes.c_double),
    ('id', ctypes.c_int64),
]

particle = struct_c__SA_particle
class struct_c__SA_particle_gc(Structure):
    pass

struct_c__SA_particle_gc._pack_ = 1 # source:False
struct_c__SA_particle_gc._fields_ = [
    ('r', ctypes.c_double),
    ('phi', ctypes.c_double),
    ('z', ctypes.c_double),
    ('energy', ctypes.c_double),
    ('pitch', ctypes.c_double),
    ('zeta', ctypes.c_double),
    ('mass', ctypes.c_double),
    ('charge', ctypes.c_double),
    ('anum', ctypes.c_int32),
    ('znum', ctypes.c_int32),
    ('weight', ctypes.c_double),
    ('time', ctypes.c_double),
    ('id', ctypes.c_int64),
]

particle_gc = struct_c__SA_particle_gc
class struct_c__SA_particle_ml(Structure):
    pass

struct_c__SA_particle_ml._pack_ = 1 # source:False
struct_c__SA_particle_ml._fields_ = [
    ('r', ctypes.c_double),
    ('phi', ctypes.c_double),
    ('z', ctypes.c_double),
    ('pitch', ctypes.c_double),
    ('weight', ctypes.c_double),
    ('time', ctypes.c_double),
    ('id', ctypes.c_int64),
]

particle_ml = struct_c__SA_particle_ml
class struct_c__SA_particle_queue(Structure):
    pass

struct_c__SA_particle_queue._pack_ = 1 # source:False
struct_c__SA_particle_queue._fields_ = [
    ('n', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('p', ctypes.POINTER(ctypes.POINTER(struct_c__SA_particle_state))),
    ('next', ctypes.c_int32),
    ('finished', ctypes.c_int32),
]

particle_queue = struct_c__SA_particle_queue

# values for enumeration 'input_particle_type'
input_particle_type__enumvalues = {
    0: 'input_particle_type_p',
    1: 'input_particle_type_gc',
    2: 'input_particle_type_ml',
    3: 'input_particle_type_s',
}
input_particle_type_p = 0
input_particle_type_gc = 1
input_particle_type_ml = 2
input_particle_type_s = 3
input_particle_type = ctypes.c_uint32 # enum
class struct_c__SA_input_particle(Structure):
    pass

class union_c__SA_input_particle_0(Union):
    _pack_ = 1 # source:False
    _fields_ = [
    ('p', particle),
    ('p_gc', particle_gc),
    ('p_ml', particle_ml),
    ('p_s', particle_state),
     ]

struct_c__SA_input_particle._pack_ = 1 # source:False
struct_c__SA_input_particle._fields_ = [
    ('type', input_particle_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('c__SA_input_particle_0', union_c__SA_input_particle_0),
]

input_particle = struct_c__SA_input_particle
class struct_c__SA_particle_simd_fo(Structure):
    pass

struct_c__SA_particle_simd_fo._pack_ = 1 # source:False
struct_c__SA_particle_simd_fo._fields_ = [
    ('r', ctypes.c_double * 16),
    ('phi', ctypes.c_double * 16),
    ('z', ctypes.c_double * 16),
    ('p_r', ctypes.c_double * 16),
    ('p_phi', ctypes.c_double * 16),
    ('p_z', ctypes.c_double * 16),
    ('mass', ctypes.c_double * 16),
    ('charge', ctypes.c_double * 16),
    ('time', ctypes.c_double * 16),
    ('znum', ctypes.c_int32 * 16),
    ('anum', ctypes.c_int32 * 16),
    ('B_r', ctypes.c_double * 16),
    ('B_phi', ctypes.c_double * 16),
    ('B_z', ctypes.c_double * 16),
    ('B_r_dr', ctypes.c_double * 16),
    ('B_phi_dr', ctypes.c_double * 16),
    ('B_z_dr', ctypes.c_double * 16),
    ('B_r_dphi', ctypes.c_double * 16),
    ('B_phi_dphi', ctypes.c_double * 16),
    ('B_z_dphi', ctypes.c_double * 16),
    ('B_r_dz', ctypes.c_double * 16),
    ('B_phi_dz', ctypes.c_double * 16),
    ('B_z_dz', ctypes.c_double * 16),
    ('bounces', ctypes.c_int32 * 16),
    ('weight', ctypes.c_double * 16),
    ('cputime', ctypes.c_double * 16),
    ('rho', ctypes.c_double * 16),
    ('theta', ctypes.c_double * 16),
    ('id', ctypes.c_int64 * 16),
    ('endcond', ctypes.c_int64 * 16),
    ('walltile', ctypes.c_int64 * 16),
    ('mileage', ctypes.c_double * 16),
    ('running', ctypes.c_int64 * 16),
    ('err', ctypes.c_uint64 * 16),
    ('index', ctypes.c_int64 * 16),
]

particle_simd_fo = struct_c__SA_particle_simd_fo
class struct_c__SA_particle_simd_gc(Structure):
    pass

struct_c__SA_particle_simd_gc._pack_ = 1 # source:False
struct_c__SA_particle_simd_gc._fields_ = [
    ('r', ctypes.c_double * 16),
    ('phi', ctypes.c_double * 16),
    ('z', ctypes.c_double * 16),
    ('ppar', ctypes.c_double * 16),
    ('mu', ctypes.c_double * 16),
    ('zeta', ctypes.c_double * 16),
    ('mass', ctypes.c_double * 16),
    ('charge', ctypes.c_double * 16),
    ('time', ctypes.c_double * 16),
    ('B_r', ctypes.c_double * 16),
    ('B_phi', ctypes.c_double * 16),
    ('B_z', ctypes.c_double * 16),
    ('B_r_dr', ctypes.c_double * 16),
    ('B_phi_dr', ctypes.c_double * 16),
    ('B_z_dr', ctypes.c_double * 16),
    ('B_r_dphi', ctypes.c_double * 16),
    ('B_phi_dphi', ctypes.c_double * 16),
    ('B_z_dphi', ctypes.c_double * 16),
    ('B_r_dz', ctypes.c_double * 16),
    ('B_phi_dz', ctypes.c_double * 16),
    ('B_z_dz', ctypes.c_double * 16),
    ('bounces', ctypes.c_int32 * 16),
    ('weight', ctypes.c_double * 16),
    ('cputime', ctypes.c_double * 16),
    ('rho', ctypes.c_double * 16),
    ('theta', ctypes.c_double * 16),
    ('id', ctypes.c_int64 * 16),
    ('endcond', ctypes.c_int64 * 16),
    ('walltile', ctypes.c_int64 * 16),
    ('mileage', ctypes.c_double * 16),
    ('running', ctypes.c_int64 * 16),
    ('err', ctypes.c_uint64 * 16),
    ('index', ctypes.c_int64 * 16),
]

particle_simd_gc = struct_c__SA_particle_simd_gc
class struct_c__SA_particle_simd_ml(Structure):
    pass

struct_c__SA_particle_simd_ml._pack_ = 1 # source:False
struct_c__SA_particle_simd_ml._fields_ = [
    ('r', ctypes.c_double * 16),
    ('phi', ctypes.c_double * 16),
    ('z', ctypes.c_double * 16),
    ('pitch', ctypes.c_double * 16),
    ('time', ctypes.c_double * 16),
    ('B_r', ctypes.c_double * 16),
    ('B_phi', ctypes.c_double * 16),
    ('B_z', ctypes.c_double * 16),
    ('B_r_dr', ctypes.c_double * 16),
    ('B_phi_dr', ctypes.c_double * 16),
    ('B_z_dr', ctypes.c_double * 16),
    ('B_r_dphi', ctypes.c_double * 16),
    ('B_phi_dphi', ctypes.c_double * 16),
    ('B_z_dphi', ctypes.c_double * 16),
    ('B_r_dz', ctypes.c_double * 16),
    ('B_phi_dz', ctypes.c_double * 16),
    ('B_z_dz', ctypes.c_double * 16),
    ('weight', ctypes.c_double * 16),
    ('cputime', ctypes.c_double * 16),
    ('rho', ctypes.c_double * 16),
    ('theta', ctypes.c_double * 16),
    ('id', ctypes.c_int64 * 16),
    ('endcond', ctypes.c_int64 * 16),
    ('walltile', ctypes.c_int64 * 16),
    ('mileage', ctypes.c_double * 16),
    ('running', ctypes.c_int64 * 16),
    ('err', ctypes.c_uint64 * 16),
    ('index', ctypes.c_int64 * 16),
]

particle_simd_ml = struct_c__SA_particle_simd_ml
particle_to_fo_dummy = _libraries['libascot.so'].particle_to_fo_dummy
particle_to_fo_dummy.restype = None
particle_to_fo_dummy.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.c_int32]
particle_to_gc_dummy = _libraries['libascot.so'].particle_to_gc_dummy
particle_to_gc_dummy.restype = None
particle_to_gc_dummy.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.c_int32]
particle_to_ml_dummy = _libraries['libascot.so'].particle_to_ml_dummy
particle_to_ml_dummy.restype = None
particle_to_ml_dummy.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.c_int32]
particle_cycle_fo = _libraries['libascot.so'].particle_cycle_fo
particle_cycle_fo.restype = ctypes.c_int32
particle_cycle_fo.argtypes = [ctypes.POINTER(struct_c__SA_particle_queue), ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.POINTER(struct_c__SA_B_field_data), ctypes.POINTER(ctypes.c_int32)]
particle_cycle_gc = _libraries['libascot.so'].particle_cycle_gc
particle_cycle_gc.restype = ctypes.c_int32
particle_cycle_gc.argtypes = [ctypes.POINTER(struct_c__SA_particle_queue), ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.POINTER(struct_c__SA_B_field_data), ctypes.POINTER(ctypes.c_int32)]
particle_cycle_ml = _libraries['libascot.so'].particle_cycle_ml
particle_cycle_ml.restype = ctypes.c_int32
particle_cycle_ml.argtypes = [ctypes.POINTER(struct_c__SA_particle_queue), ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.POINTER(struct_c__SA_B_field_data), ctypes.POINTER(ctypes.c_int32)]
particle_input_to_state = _libraries['libascot.so'].particle_input_to_state
particle_input_to_state.restype = None
particle_input_to_state.argtypes = [ctypes.POINTER(struct_c__SA_input_particle), ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(struct_c__SA_B_field_data)]
particle_state_to_fo = _libraries['libascot.so'].particle_state_to_fo
particle_state_to_fo.restype = a5err
particle_state_to_fo.argtypes = [ctypes.POINTER(struct_c__SA_particle_state), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.c_int32, ctypes.POINTER(struct_c__SA_B_field_data)]
particle_fo_to_state = _libraries['libascot.so'].particle_fo_to_state
particle_fo_to_state.restype = None
particle_fo_to_state.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(struct_c__SA_B_field_data)]
particle_state_to_gc = _libraries['libascot.so'].particle_state_to_gc
particle_state_to_gc.restype = a5err
particle_state_to_gc.argtypes = [ctypes.POINTER(struct_c__SA_particle_state), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.c_int32, ctypes.POINTER(struct_c__SA_B_field_data)]
particle_gc_to_state = _libraries['libascot.so'].particle_gc_to_state
particle_gc_to_state.restype = None
particle_gc_to_state.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(struct_c__SA_B_field_data)]
particle_state_to_ml = _libraries['libascot.so'].particle_state_to_ml
particle_state_to_ml.restype = a5err
particle_state_to_ml.argtypes = [ctypes.POINTER(struct_c__SA_particle_state), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.c_int32, ctypes.POINTER(struct_c__SA_B_field_data)]
particle_ml_to_state = _libraries['libascot.so'].particle_ml_to_state
particle_ml_to_state.restype = None
particle_ml_to_state.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(struct_c__SA_B_field_data)]
particle_fo_to_gc = _libraries['libascot.so'].particle_fo_to_gc
particle_fo_to_gc.restype = ctypes.c_int32
particle_fo_to_gc.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.POINTER(struct_c__SA_B_field_data)]
particle_copy_fo = _libraries['libascot.so'].particle_copy_fo
particle_copy_fo.restype = None
particle_copy_fo.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.c_int32]
particle_copy_gc = _libraries['libascot.so'].particle_copy_gc
particle_copy_gc.restype = None
particle_copy_gc.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.c_int32]
particle_copy_ml = _libraries['libascot.so'].particle_copy_ml
particle_copy_ml.restype = None
particle_copy_ml.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.c_int32]

# values for enumeration 'wall_type'
wall_type__enumvalues = {
    0: 'wall_type_2D',
    1: 'wall_type_3D',
}
wall_type_2D = 0
wall_type_3D = 1
wall_type = ctypes.c_uint32 # enum
class struct_c__SA_wall_offload_data(Structure):
    pass

class struct_c__SA_wall_3d_offload_data(Structure):
    pass

struct_c__SA_wall_3d_offload_data._pack_ = 1 # source:False
struct_c__SA_wall_3d_offload_data._fields_ = [
    ('n', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('xmin', ctypes.c_double),
    ('xmax', ctypes.c_double),
    ('xgrid', ctypes.c_double),
    ('ymin', ctypes.c_double),
    ('ymax', ctypes.c_double),
    ('ygrid', ctypes.c_double),
    ('zmin', ctypes.c_double),
    ('zmax', ctypes.c_double),
    ('zgrid', ctypes.c_double),
    ('depth', ctypes.c_int32),
    ('ngrid', ctypes.c_int32),
    ('offload_array_length', ctypes.c_int32),
    ('int_offload_array_length', ctypes.c_int32),
]

class struct_c__SA_wall_2d_offload_data(Structure):
    pass

struct_c__SA_wall_2d_offload_data._pack_ = 1 # source:False
struct_c__SA_wall_2d_offload_data._fields_ = [
    ('n', ctypes.c_int32),
    ('offload_array_length', ctypes.c_int32),
]

struct_c__SA_wall_offload_data._pack_ = 1 # source:False
struct_c__SA_wall_offload_data._fields_ = [
    ('type', wall_type),
    ('w2d', struct_c__SA_wall_2d_offload_data),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('w3d', struct_c__SA_wall_3d_offload_data),
    ('offload_array_length', ctypes.c_int32),
    ('int_offload_array_length', ctypes.c_int32),
]

wall_offload_data = struct_c__SA_wall_offload_data
class struct_c__SA_wall_data(Structure):
    pass

class struct_c__SA_wall_2d_data(Structure):
    pass

struct_c__SA_wall_2d_data._pack_ = 1 # source:False
struct_c__SA_wall_2d_data._fields_ = [
    ('n', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('wall_r', ctypes.POINTER(ctypes.c_double)),
    ('wall_z', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_wall_3d_data(Structure):
    pass

struct_c__SA_wall_3d_data._pack_ = 1 # source:False
struct_c__SA_wall_3d_data._fields_ = [
    ('n', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('xmin', ctypes.c_double),
    ('xmax', ctypes.c_double),
    ('xgrid', ctypes.c_double),
    ('ymin', ctypes.c_double),
    ('ymax', ctypes.c_double),
    ('ygrid', ctypes.c_double),
    ('zmin', ctypes.c_double),
    ('zmax', ctypes.c_double),
    ('zgrid', ctypes.c_double),
    ('depth', ctypes.c_int32),
    ('ngrid', ctypes.c_int32),
    ('wall_tris', ctypes.POINTER(ctypes.c_double)),
    ('tree_array_size', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('tree_array', ctypes.POINTER(ctypes.c_int32)),
]

struct_c__SA_wall_data._pack_ = 1 # source:False
struct_c__SA_wall_data._fields_ = [
    ('type', wall_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('w2d', struct_c__SA_wall_2d_data),
    ('w3d', struct_c__SA_wall_3d_data),
]

wall_data = struct_c__SA_wall_data
wall_init_offload = _libraries['libascot.so'].wall_init_offload
wall_init_offload.restype = ctypes.c_int32
wall_init_offload.argtypes = [ctypes.POINTER(struct_c__SA_wall_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32))]
wall_free_offload = _libraries['libascot.so'].wall_free_offload
wall_free_offload.restype = None
wall_free_offload.argtypes = [ctypes.POINTER(struct_c__SA_wall_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32))]
wall_init = _libraries['libascot.so'].wall_init
wall_init.restype = ctypes.c_int32
wall_init.argtypes = [ctypes.POINTER(struct_c__SA_wall_data), ctypes.POINTER(struct_c__SA_wall_offload_data), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_int32)]
wall_hit_wall = _libraries['libascot.so'].wall_hit_wall
wall_hit_wall.restype = ctypes.c_int32
wall_hit_wall.argtypes = [real, real, real, real, real, real, ctypes.POINTER(struct_c__SA_wall_data), ctypes.POINTER(ctypes.c_double)]
diag_orb_check_plane_crossing = _libraries['libascot.so'].diag_orb_check_plane_crossing
diag_orb_check_plane_crossing.restype = real
diag_orb_check_plane_crossing.argtypes = [real, real, real]
diag_orb_check_radial_crossing = _libraries['libascot.so'].diag_orb_check_radial_crossing
diag_orb_check_radial_crossing.restype = real
diag_orb_check_radial_crossing.argtypes = [real, real, real]
class struct_c__SA_diag_orb_offload_data(Structure):
    pass

struct_c__SA_diag_orb_offload_data._pack_ = 1 # source:False
struct_c__SA_diag_orb_offload_data._fields_ = [
    ('record_mode', ctypes.c_int32),
    ('mode', ctypes.c_int32),
    ('Npnt', ctypes.c_int32),
    ('Nmrk', ctypes.c_int32),
    ('Nfld', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('writeInterval', ctypes.c_double),
    ('ntoroidalplots', ctypes.c_int32),
    ('npoloidalplots', ctypes.c_int32),
    ('nradialplots', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('toroidalangles', ctypes.c_double * 30),
    ('poloidalangles', ctypes.c_double * 30),
    ('radialdistances', ctypes.c_double * 30),
]

diag_orb_offload_data = struct_c__SA_diag_orb_offload_data
class struct_c__SA_diag_orb_data(Structure):
    pass

struct_c__SA_diag_orb_data._pack_ = 1 # source:False
struct_c__SA_diag_orb_data._fields_ = [
    ('id', ctypes.POINTER(ctypes.c_double)),
    ('mileage', ctypes.POINTER(ctypes.c_double)),
    ('r', ctypes.POINTER(ctypes.c_double)),
    ('phi', ctypes.POINTER(ctypes.c_double)),
    ('z', ctypes.POINTER(ctypes.c_double)),
    ('p_r', ctypes.POINTER(ctypes.c_double)),
    ('p_phi', ctypes.POINTER(ctypes.c_double)),
    ('p_z', ctypes.POINTER(ctypes.c_double)),
    ('ppar', ctypes.POINTER(ctypes.c_double)),
    ('mu', ctypes.POINTER(ctypes.c_double)),
    ('zeta', ctypes.POINTER(ctypes.c_double)),
    ('weight', ctypes.POINTER(ctypes.c_double)),
    ('charge', ctypes.POINTER(ctypes.c_double)),
    ('rho', ctypes.POINTER(ctypes.c_double)),
    ('theta', ctypes.POINTER(ctypes.c_double)),
    ('B_r', ctypes.POINTER(ctypes.c_double)),
    ('B_phi', ctypes.POINTER(ctypes.c_double)),
    ('B_z', ctypes.POINTER(ctypes.c_double)),
    ('simmode', ctypes.POINTER(ctypes.c_double)),
    ('pncrid', ctypes.POINTER(ctypes.c_double)),
    ('pncrdi', ctypes.POINTER(ctypes.c_double)),
    ('mrk_pnt', ctypes.POINTER(ctypes.c_int64)),
    ('mrk_recorded', ctypes.POINTER(ctypes.c_double)),
    ('mode', ctypes.c_int32),
    ('Npnt', ctypes.c_int32),
    ('Nmrk', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('writeInterval', ctypes.c_double),
    ('ntoroidalplots', ctypes.c_int32),
    ('npoloidalplots', ctypes.c_int32),
    ('nradialplots', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('toroidalangles', ctypes.c_double * 30),
    ('poloidalangles', ctypes.c_double * 30),
    ('radialdistances', ctypes.c_double * 30),
]

diag_orb_data = struct_c__SA_diag_orb_data
diag_orb_init = _libraries['libascot.so'].diag_orb_init
diag_orb_init.restype = None
diag_orb_init.argtypes = [ctypes.POINTER(struct_c__SA_diag_orb_data), ctypes.POINTER(struct_c__SA_diag_orb_offload_data), ctypes.POINTER(ctypes.c_double)]
diag_orb_free = _libraries['libascot.so'].diag_orb_free
diag_orb_free.restype = None
diag_orb_free.argtypes = [ctypes.POINTER(struct_c__SA_diag_orb_data)]
diag_orb_update_fo = _libraries['libascot.so'].diag_orb_update_fo
diag_orb_update_fo.restype = None
diag_orb_update_fo.argtypes = [ctypes.POINTER(struct_c__SA_diag_orb_data), ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.POINTER(struct_c__SA_particle_simd_fo)]
diag_orb_update_gc = _libraries['libascot.so'].diag_orb_update_gc
diag_orb_update_gc.restype = None
diag_orb_update_gc.argtypes = [ctypes.POINTER(struct_c__SA_diag_orb_data), ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.POINTER(struct_c__SA_particle_simd_gc)]
diag_orb_update_ml = _libraries['libascot.so'].diag_orb_update_ml
diag_orb_update_ml.restype = None
diag_orb_update_ml.argtypes = [ctypes.POINTER(struct_c__SA_diag_orb_data), ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.POINTER(struct_c__SA_particle_simd_ml)]
class struct_c__SA_diag_offload_data(Structure):
    pass

class struct_c__SA_dist_5D_offload_data(Structure):
    pass

struct_c__SA_dist_5D_offload_data._pack_ = 1 # source:False
struct_c__SA_dist_5D_offload_data._fields_ = [
    ('n_r', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_r', ctypes.c_double),
    ('max_r', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_z', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_z', ctypes.c_double),
    ('max_z', ctypes.c_double),
    ('n_ppara', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_ppara', ctypes.c_double),
    ('max_ppara', ctypes.c_double),
    ('n_pperp', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pperp', ctypes.c_double),
    ('max_pperp', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
]

class struct_c__SA_dist_6D_offload_data(Structure):
    pass

struct_c__SA_dist_6D_offload_data._pack_ = 1 # source:False
struct_c__SA_dist_6D_offload_data._fields_ = [
    ('n_r', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_r', ctypes.c_double),
    ('max_r', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_z', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_z', ctypes.c_double),
    ('max_z', ctypes.c_double),
    ('n_pr', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_pr', ctypes.c_double),
    ('max_pr', ctypes.c_double),
    ('n_pphi', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pphi', ctypes.c_double),
    ('max_pphi', ctypes.c_double),
    ('n_pz', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_pz', ctypes.c_double),
    ('max_pz', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_7', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
]

class struct_c__SA_diag_transcoef_offload_data(Structure):
    pass

struct_c__SA_diag_transcoef_offload_data._pack_ = 1 # source:False
struct_c__SA_diag_transcoef_offload_data._fields_ = [
    ('Nmrk', ctypes.c_int64),
    ('Navg', ctypes.c_int32),
    ('recordrho', ctypes.c_int32),
    ('interval', ctypes.c_double),
]

class struct_c__SA_dist_COM_offload_data(Structure):
    pass

struct_c__SA_dist_COM_offload_data._pack_ = 1 # source:False
struct_c__SA_dist_COM_offload_data._fields_ = [
    ('n_mu', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_mu', ctypes.c_double),
    ('max_mu', ctypes.c_double),
    ('n_Ekin', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_Ekin', ctypes.c_double),
    ('max_Ekin', ctypes.c_double),
    ('n_Ptor', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_Ptor', ctypes.c_double),
    ('max_Ptor', ctypes.c_double),
]

class struct_c__SA_dist_rho5D_offload_data(Structure):
    pass

struct_c__SA_dist_rho5D_offload_data._pack_ = 1 # source:False
struct_c__SA_dist_rho5D_offload_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_rho', ctypes.c_double),
    ('max_rho', ctypes.c_double),
    ('n_theta', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_theta', ctypes.c_double),
    ('max_theta', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_ppara', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_ppara', ctypes.c_double),
    ('max_ppara', ctypes.c_double),
    ('n_pperp', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pperp', ctypes.c_double),
    ('max_pperp', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
]

class struct_c__SA_dist_rho6D_offload_data(Structure):
    pass

struct_c__SA_dist_rho6D_offload_data._pack_ = 1 # source:False
struct_c__SA_dist_rho6D_offload_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_rho', ctypes.c_double),
    ('max_rho', ctypes.c_double),
    ('n_theta', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_theta', ctypes.c_double),
    ('max_theta', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_pr', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_pr', ctypes.c_double),
    ('max_pr', ctypes.c_double),
    ('n_pphi', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pphi', ctypes.c_double),
    ('max_pphi', ctypes.c_double),
    ('n_pz', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_pz', ctypes.c_double),
    ('max_pz', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_7', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
]

struct_c__SA_diag_offload_data._pack_ = 1 # source:False
struct_c__SA_diag_offload_data._fields_ = [
    ('diagorb_collect', ctypes.c_int32),
    ('dist5D_collect', ctypes.c_int32),
    ('dist6D_collect', ctypes.c_int32),
    ('distrho5D_collect', ctypes.c_int32),
    ('distrho6D_collect', ctypes.c_int32),
    ('distCOM_collect', ctypes.c_int32),
    ('diagtrcof_collect', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('diagorb', diag_orb_offload_data),
    ('dist5D', struct_c__SA_dist_5D_offload_data),
    ('dist6D', struct_c__SA_dist_6D_offload_data),
    ('distrho5D', struct_c__SA_dist_rho5D_offload_data),
    ('distrho6D', struct_c__SA_dist_rho6D_offload_data),
    ('distCOM', struct_c__SA_dist_COM_offload_data),
    ('diagtrcof', struct_c__SA_diag_transcoef_offload_data),
    ('offload_dist5D_index', ctypes.c_int32),
    ('offload_dist6D_index', ctypes.c_int32),
    ('offload_distrho5D_index', ctypes.c_int32),
    ('offload_distrho6D_index', ctypes.c_int32),
    ('offload_distCOM_index', ctypes.c_int32),
    ('offload_diagorb_index', ctypes.c_int32),
    ('offload_diagtrcof_index', ctypes.c_int32),
    ('offload_dist_length', ctypes.c_int32),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

diag_offload_data = struct_c__SA_diag_offload_data
class struct_c__SA_diag_data(Structure):
    pass

class struct_c__SA_dist_5D_data(Structure):
    pass

struct_c__SA_dist_5D_data._pack_ = 1 # source:False
struct_c__SA_dist_5D_data._fields_ = [
    ('n_r', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_r', ctypes.c_double),
    ('max_r', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_z', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_z', ctypes.c_double),
    ('max_z', ctypes.c_double),
    ('n_ppara', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_ppara', ctypes.c_double),
    ('max_ppara', ctypes.c_double),
    ('n_pperp', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pperp', ctypes.c_double),
    ('max_pperp', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
    ('histogram', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_dist_rho5D_data(Structure):
    pass

struct_c__SA_dist_rho5D_data._pack_ = 1 # source:False
struct_c__SA_dist_rho5D_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_rho', ctypes.c_double),
    ('max_rho', ctypes.c_double),
    ('n_theta', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_theta', ctypes.c_double),
    ('max_theta', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_ppara', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_ppara', ctypes.c_double),
    ('max_ppara', ctypes.c_double),
    ('n_pperp', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pperp', ctypes.c_double),
    ('max_pperp', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
    ('histogram', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_dist_COM_data(Structure):
    pass

struct_c__SA_dist_COM_data._pack_ = 1 # source:False
struct_c__SA_dist_COM_data._fields_ = [
    ('n_mu', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_mu', ctypes.c_double),
    ('max_mu', ctypes.c_double),
    ('n_Ekin', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_Ekin', ctypes.c_double),
    ('max_Ekin', ctypes.c_double),
    ('n_Ptor', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_Ptor', ctypes.c_double),
    ('max_Ptor', ctypes.c_double),
    ('histogram', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_diag_transcoef_data(Structure):
    pass

class struct_diag_transcoef_link(Structure):
    pass

struct_c__SA_diag_transcoef_data._pack_ = 1 # source:False
struct_c__SA_diag_transcoef_data._fields_ = [
    ('Navg', ctypes.c_int32),
    ('recordrho', ctypes.c_int32),
    ('interval', ctypes.c_double),
    ('datapoints', ctypes.POINTER(ctypes.POINTER(struct_diag_transcoef_link))),
    ('id', ctypes.POINTER(ctypes.c_double)),
    ('Kcoef', ctypes.POINTER(ctypes.c_double)),
    ('Dcoef', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_dist_6D_data(Structure):
    pass

struct_c__SA_dist_6D_data._pack_ = 1 # source:False
struct_c__SA_dist_6D_data._fields_ = [
    ('n_r', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_r', ctypes.c_double),
    ('max_r', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_z', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_z', ctypes.c_double),
    ('max_z', ctypes.c_double),
    ('n_pr', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_pr', ctypes.c_double),
    ('max_pr', ctypes.c_double),
    ('n_pphi', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pphi', ctypes.c_double),
    ('max_pphi', ctypes.c_double),
    ('n_pz', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_pz', ctypes.c_double),
    ('max_pz', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_7', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
    ('histogram', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_dist_rho6D_data(Structure):
    pass

struct_c__SA_dist_rho6D_data._pack_ = 1 # source:False
struct_c__SA_dist_rho6D_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_rho', ctypes.c_double),
    ('max_rho', ctypes.c_double),
    ('n_theta', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_theta', ctypes.c_double),
    ('max_theta', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_pr', ctypes.c_int32),
    ('PADDING_3', ctypes.c_ubyte * 4),
    ('min_pr', ctypes.c_double),
    ('max_pr', ctypes.c_double),
    ('n_pphi', ctypes.c_int32),
    ('PADDING_4', ctypes.c_ubyte * 4),
    ('min_pphi', ctypes.c_double),
    ('max_pphi', ctypes.c_double),
    ('n_pz', ctypes.c_int32),
    ('PADDING_5', ctypes.c_ubyte * 4),
    ('min_pz', ctypes.c_double),
    ('max_pz', ctypes.c_double),
    ('n_time', ctypes.c_int32),
    ('PADDING_6', ctypes.c_ubyte * 4),
    ('min_time', ctypes.c_double),
    ('max_time', ctypes.c_double),
    ('n_q', ctypes.c_int32),
    ('PADDING_7', ctypes.c_ubyte * 4),
    ('min_q', ctypes.c_double),
    ('max_q', ctypes.c_double),
    ('histogram', ctypes.POINTER(ctypes.c_double)),
]

struct_c__SA_diag_data._pack_ = 1 # source:False
struct_c__SA_diag_data._fields_ = [
    ('diagorb_collect', ctypes.c_int32),
    ('dist5D_collect', ctypes.c_int32),
    ('dist6D_collect', ctypes.c_int32),
    ('distrho5D_collect', ctypes.c_int32),
    ('distrho6D_collect', ctypes.c_int32),
    ('distCOM_collect', ctypes.c_int32),
    ('diagtrcof_collect', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('diagorb', diag_orb_data),
    ('dist5D', struct_c__SA_dist_5D_data),
    ('dist6D', struct_c__SA_dist_6D_data),
    ('distrho5D', struct_c__SA_dist_rho5D_data),
    ('distrho6D', struct_c__SA_dist_rho6D_data),
    ('distCOM', struct_c__SA_dist_COM_data),
    ('diagtrcof', struct_c__SA_diag_transcoef_data),
]

struct_diag_transcoef_link._pack_ = 1 # source:False
struct_diag_transcoef_link._fields_ = [
    ('rho', ctypes.c_double),
    ('time', ctypes.c_double),
    ('pitchsign', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('prevlink', ctypes.POINTER(struct_diag_transcoef_link)),
]

diag_data = struct_c__SA_diag_data
diag_init_offload = _libraries['libascot.so'].diag_init_offload
diag_init_offload.restype = ctypes.c_int32
diag_init_offload.argtypes = [ctypes.POINTER(struct_c__SA_diag_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.c_int32]
diag_free_offload = _libraries['libascot.so'].diag_free_offload
diag_free_offload.restype = None
diag_free_offload.argtypes = [ctypes.POINTER(struct_c__SA_diag_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]
diag_sum = _libraries['libascot.so'].diag_sum
diag_sum.restype = None
diag_sum.argtypes = [ctypes.POINTER(struct_c__SA_diag_offload_data), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)]
diag_init = _libraries['libascot.so'].diag_init
diag_init.restype = None
diag_init.argtypes = [ctypes.POINTER(struct_c__SA_diag_data), ctypes.POINTER(struct_c__SA_diag_offload_data), ctypes.POINTER(ctypes.c_double)]
diag_free = _libraries['libascot.so'].diag_free
diag_free.restype = None
diag_free.argtypes = [ctypes.POINTER(struct_c__SA_diag_data)]
diag_update_fo = _libraries['libascot.so'].diag_update_fo
diag_update_fo.restype = None
diag_update_fo.argtypes = [ctypes.POINTER(struct_c__SA_diag_data), ctypes.POINTER(struct_c__SA_B_field_data), ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.POINTER(struct_c__SA_particle_simd_fo)]
diag_update_gc = _libraries['libascot.so'].diag_update_gc
diag_update_gc.restype = None
diag_update_gc.argtypes = [ctypes.POINTER(struct_c__SA_diag_data), ctypes.POINTER(struct_c__SA_B_field_data), ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.POINTER(struct_c__SA_particle_simd_gc)]
diag_update_ml = _libraries['libascot.so'].diag_update_ml
diag_update_ml.restype = None
diag_update_ml.argtypes = [ctypes.POINTER(struct_c__SA_diag_data), ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.POINTER(struct_c__SA_particle_simd_ml)]
class struct_c__SA_offload_package(Structure):
    pass

struct_c__SA_offload_package._pack_ = 1 # source:False
struct_c__SA_offload_package._fields_ = [
    ('offload_array_length', ctypes.c_uint64),
    ('unpack_pos', ctypes.c_uint64),
    ('int_offload_array_length', ctypes.c_uint64),
    ('int_unpack_pos', ctypes.c_uint64),
]

offload_package = struct_c__SA_offload_package
offload_init_offload = _libraries['libascot.so'].offload_init_offload
offload_init_offload.restype = None
offload_init_offload.argtypes = [ctypes.POINTER(struct_c__SA_offload_package), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32))]
offload_free_offload = _libraries['libascot.so'].offload_free_offload
offload_free_offload.restype = None
offload_free_offload.argtypes = [ctypes.POINTER(struct_c__SA_offload_package), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32))]
size_t = ctypes.c_uint64
offload_pack = _libraries['libascot.so'].offload_pack
offload_pack.restype = None
offload_pack.argtypes = [ctypes.POINTER(struct_c__SA_offload_package), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.c_double), size_t, ctypes.POINTER(ctypes.POINTER(ctypes.c_int32)), ctypes.POINTER(ctypes.c_int32), size_t]
offload_unpack = _libraries['libascot.so'].offload_unpack
offload_unpack.restype = None
offload_unpack.argtypes = [ctypes.POINTER(struct_c__SA_offload_package), ctypes.POINTER(ctypes.c_double), size_t, ctypes.POINTER(ctypes.c_int32), size_t, ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32))]

# values for enumeration 'c__Ea_simulate_mode_fo'
c__Ea_simulate_mode_fo__enumvalues = {
    1: 'simulate_mode_fo',
    2: 'simulate_mode_gc',
    3: 'simulate_mode_hybrid',
    4: 'simulate_mode_ml',
}
simulate_mode_fo = 1
simulate_mode_gc = 2
simulate_mode_hybrid = 3
simulate_mode_ml = 4
c__Ea_simulate_mode_fo = ctypes.c_uint32 # enum
class struct_c__SA_sim_offload_data(Structure):
    pass

class struct_c__SA_mhd_offload_data(Structure):
    pass


# values for enumeration 'mhd_type'
mhd_type__enumvalues = {
    0: 'mhd_type_stat',
    1: 'mhd_type_nonstat',
}
mhd_type_stat = 0
mhd_type_nonstat = 1
mhd_type = ctypes.c_uint32 # enum
class struct_c__SA_mhd_nonstat_offload_data(Structure):
    pass

struct_c__SA_mhd_nonstat_offload_data._pack_ = 1 # source:False
struct_c__SA_mhd_nonstat_offload_data._fields_ = [
    ('n_modes', ctypes.c_int32),
    ('nrho', ctypes.c_int32),
    ('rho_min', ctypes.c_double),
    ('rho_max', ctypes.c_double),
    ('ntime', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('t_min', ctypes.c_double),
    ('t_max', ctypes.c_double),
    ('nmode', ctypes.c_int32 * 512),
    ('mmode', ctypes.c_int32 * 512),
    ('amplitude_nm', ctypes.c_double * 512),
    ('omega_nm', ctypes.c_double * 512),
    ('phase_nm', ctypes.c_double * 512),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

class struct_c__SA_mhd_stat_offload_data(Structure):
    pass

struct_c__SA_mhd_stat_offload_data._pack_ = 1 # source:False
struct_c__SA_mhd_stat_offload_data._fields_ = [
    ('n_modes', ctypes.c_int32),
    ('nrho', ctypes.c_int32),
    ('rho_min', ctypes.c_double),
    ('rho_max', ctypes.c_double),
    ('nmode', ctypes.c_int32 * 512),
    ('mmode', ctypes.c_int32 * 512),
    ('amplitude_nm', ctypes.c_double * 512),
    ('omega_nm', ctypes.c_double * 512),
    ('phase_nm', ctypes.c_double * 512),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
]

struct_c__SA_mhd_offload_data._pack_ = 1 # source:False
struct_c__SA_mhd_offload_data._fields_ = [
    ('type', mhd_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('stat', struct_c__SA_mhd_stat_offload_data),
    ('nonstat', struct_c__SA_mhd_nonstat_offload_data),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

class struct_c__SA_E_field_offload_data(Structure):
    pass


# values for enumeration 'E_field_type'
E_field_type__enumvalues = {
    0: 'E_field_type_TC',
    1: 'E_field_type_1DS',
}
E_field_type_TC = 0
E_field_type_1DS = 1
E_field_type = ctypes.c_uint32 # enum
class struct_c__SA_E_TC_offload_data(Structure):
    pass

struct_c__SA_E_TC_offload_data._pack_ = 1 # source:False
struct_c__SA_E_TC_offload_data._fields_ = [
    ('Exyz', ctypes.c_double * 3),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
]

class struct_c__SA_E_1DS_offload_data(Structure):
    pass

struct_c__SA_E_1DS_offload_data._pack_ = 1 # source:False
struct_c__SA_E_1DS_offload_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('rho_min', ctypes.c_double),
    ('rho_max', ctypes.c_double),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

struct_c__SA_E_field_offload_data._pack_ = 1 # source:False
struct_c__SA_E_field_offload_data._fields_ = [
    ('type', E_field_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('ETC', struct_c__SA_E_TC_offload_data),
    ('E1DS', struct_c__SA_E_1DS_offload_data),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

class struct_c__SA_plasma_offload_data(Structure):
    pass


# values for enumeration 'plasma_type'
plasma_type__enumvalues = {
    0: 'plasma_type_1D',
    1: 'plasma_type_1Dt',
    2: 'plasma_type_1DS',
}
plasma_type_1D = 0
plasma_type_1Dt = 1
plasma_type_1DS = 2
plasma_type = ctypes.c_uint32 # enum
class struct_c__SA_plasma_1D_offload_data(Structure):
    pass

struct_c__SA_plasma_1D_offload_data._pack_ = 1 # source:False
struct_c__SA_plasma_1D_offload_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('n_species', ctypes.c_int32),
    ('mass', ctypes.c_double * 8),
    ('charge', ctypes.c_double * 8),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
]

class struct_c__SA_plasma_1Dt_offload_data(Structure):
    pass

struct_c__SA_plasma_1Dt_offload_data._pack_ = 1 # source:False
struct_c__SA_plasma_1Dt_offload_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('n_time', ctypes.c_int32),
    ('n_species', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('mass', ctypes.c_double * 8),
    ('charge', ctypes.c_double * 8),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

class struct_c__SA_plasma_1DS_offload_data(Structure):
    pass

struct_c__SA_plasma_1DS_offload_data._pack_ = 1 # source:False
struct_c__SA_plasma_1DS_offload_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('rho_min', ctypes.c_double),
    ('rho_max', ctypes.c_double),
    ('n_species', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('mass', ctypes.c_double * 8),
    ('charge', ctypes.c_double * 8),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
]

struct_c__SA_plasma_offload_data._pack_ = 1 # source:False
struct_c__SA_plasma_offload_data._fields_ = [
    ('type', plasma_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('plasma_1D', struct_c__SA_plasma_1D_offload_data),
    ('plasma_1Dt', struct_c__SA_plasma_1Dt_offload_data),
    ('plasma_1DS', struct_c__SA_plasma_1DS_offload_data),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

class struct_c__SA_neutral_offload_data(Structure):
    pass

class struct_c__SA_N0_3D_offload_data(Structure):
    pass

struct_c__SA_N0_3D_offload_data._pack_ = 1 # source:False
struct_c__SA_N0_3D_offload_data._fields_ = [
    ('n_r', ctypes.c_int32),
    ('n_z', ctypes.c_int32),
    ('n_phi', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('r_min', ctypes.c_double),
    ('r_max', ctypes.c_double),
    ('z_min', ctypes.c_double),
    ('z_max', ctypes.c_double),
    ('phi_min', ctypes.c_double),
    ('phi_max', ctypes.c_double),
    ('n_species', ctypes.c_int32),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('maxwellian', ctypes.c_int32 * 8),
    ('offload_array_length', ctypes.c_int32),
]

class struct_c__SA_N0_1D_offload_data(Structure):
    pass

struct_c__SA_N0_1D_offload_data._pack_ = 1 # source:False
struct_c__SA_N0_1D_offload_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('rho_min', ctypes.c_double),
    ('rho_max', ctypes.c_double),
    ('n_species', ctypes.c_int32),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('maxwellian', ctypes.c_int32 * 8),
    ('offload_array_length', ctypes.c_int32),
]


# values for enumeration 'neutral_type'
neutral_type__enumvalues = {
    0: 'neutral_type_1D',
    1: 'neutral_type_3D',
}
neutral_type_1D = 0
neutral_type_3D = 1
neutral_type = ctypes.c_uint32 # enum
struct_c__SA_neutral_offload_data._pack_ = 1 # source:False
struct_c__SA_neutral_offload_data._fields_ = [
    ('type', neutral_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('N01D', struct_c__SA_N0_1D_offload_data),
    ('N03D', struct_c__SA_N0_3D_offload_data),
    ('offload_array_length', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

class struct_c__SA_boozer_offload_data(Structure):
    pass

struct_c__SA_boozer_offload_data._pack_ = 1 # source:False
struct_c__SA_boozer_offload_data._fields_ = [
    ('nr', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('r_min', ctypes.c_double),
    ('r_max', ctypes.c_double),
    ('nz', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('z_min', ctypes.c_double),
    ('z_max', ctypes.c_double),
    ('npsi', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('psi_min', ctypes.c_double),
    ('psi_max', ctypes.c_double),
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('ntheta', ctypes.c_int32),
    ('nthetag', ctypes.c_int32),
    ('r0', ctypes.c_double),
    ('z0', ctypes.c_double),
    ('nrzs', ctypes.c_int32),
    ('offload_array_length', ctypes.c_int32),
]

class struct_c__SA_asigma_offload_data(Structure):
    pass


# values for enumeration 'asigma_type'
asigma_type__enumvalues = {
    0: 'asigma_type_loc',
}
asigma_type_loc = 0
asigma_type = ctypes.c_uint32 # enum
class struct_c__SA_asigma_loc_offload_data(Structure):
    pass

struct_c__SA_asigma_loc_offload_data._pack_ = 1 # source:False
struct_c__SA_asigma_loc_offload_data._fields_ = [
    ('data_format', ctypes.c_int32),
    ('N_reac', ctypes.c_int32),
    ('offload_array_length', ctypes.c_int32),
]

struct_c__SA_asigma_offload_data._pack_ = 1 # source:False
struct_c__SA_asigma_offload_data._fields_ = [
    ('type', asigma_type),
    ('asigma_loc', struct_c__SA_asigma_loc_offload_data),
    ('offload_array_length', ctypes.c_int32),
]

struct_c__SA_sim_offload_data._pack_ = 1 # source:False
struct_c__SA_sim_offload_data._fields_ = [
    ('B_offload_data', B_field_offload_data),
    ('E_offload_data', struct_c__SA_E_field_offload_data),
    ('plasma_offload_data', struct_c__SA_plasma_offload_data),
    ('neutral_offload_data', struct_c__SA_neutral_offload_data),
    ('wall_offload_data', wall_offload_data),
    ('boozer_offload_data', struct_c__SA_boozer_offload_data),
    ('mhd_offload_data', struct_c__SA_mhd_offload_data),
    ('asigma_offload_data', struct_c__SA_asigma_offload_data),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('diag_offload_data', diag_offload_data),
    ('sim_mode', ctypes.c_int32),
    ('enable_ada', ctypes.c_int32),
    ('record_mode', ctypes.c_int32),
    ('fix_usrdef_use', ctypes.c_int32),
    ('fix_usrdef_val', ctypes.c_double),
    ('fix_gyrodef_nstep', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('ada_tol_orbfol', ctypes.c_double),
    ('ada_tol_clmbcol', ctypes.c_double),
    ('ada_max_drho', ctypes.c_double),
    ('ada_max_dphi', ctypes.c_double),
    ('enable_orbfol', ctypes.c_int32),
    ('enable_clmbcol', ctypes.c_int32),
    ('enable_mhd', ctypes.c_int32),
    ('enable_atomic', ctypes.c_int32),
    ('disable_gctransform', ctypes.c_int32),
    ('disable_energyccoll', ctypes.c_int32),
    ('disable_pitchccoll', ctypes.c_int32),
    ('disable_gcdiffccoll', ctypes.c_int32),
    ('reverse_time', ctypes.c_int32),
    ('endcond_active', ctypes.c_int32),
    ('endcond_lim_simtime', ctypes.c_double),
    ('endcond_max_mileage', ctypes.c_double),
    ('endcond_max_cputime', ctypes.c_double),
    ('endcond_min_rho', ctypes.c_double),
    ('endcond_max_rho', ctypes.c_double),
    ('endcond_min_ekin', ctypes.c_double),
    ('endcond_min_thermal', ctypes.c_double),
    ('endcond_max_tororb', ctypes.c_double),
    ('endcond_max_polorb', ctypes.c_double),
    ('endcond_torandpol', ctypes.c_int32),
    ('hdf5_in', ctypes.c_char * 256),
    ('hdf5_out', ctypes.c_char * 256),
    ('qid', ctypes.c_char * 256),
    ('description', ctypes.c_char * 256),
    ('mpi_rank', ctypes.c_int32),
    ('mpi_size', ctypes.c_int32),
    ('qid_options', ctypes.c_char * 256),
    ('qid_bfield', ctypes.c_char * 256),
    ('qid_efield', ctypes.c_char * 256),
    ('qid_marker', ctypes.c_char * 256),
    ('qid_wall', ctypes.c_char * 256),
    ('qid_plasma', ctypes.c_char * 256),
    ('qid_neutral', ctypes.c_char * 256),
    ('qid_boozer', ctypes.c_char * 256),
    ('qid_mhd', ctypes.c_char * 256),
    ('qid_asigma', ctypes.c_char * 256),
    ('PADDING_2', ctypes.c_ubyte * 4),
]

sim_offload_data = struct_c__SA_sim_offload_data
class struct_c__SA_sim_data(Structure):
    pass

class struct_c__SA_mhd_data(Structure):
    pass

class struct_c__SA_mhd_stat_data(Structure):
    pass

class struct_c__SA_interp1D_data(Structure):
    pass

struct_c__SA_interp1D_data._pack_ = 1 # source:False
struct_c__SA_interp1D_data._fields_ = [
    ('n_x', ctypes.c_int32),
    ('bc_x', ctypes.c_int32),
    ('x_min', ctypes.c_double),
    ('x_max', ctypes.c_double),
    ('x_grid', ctypes.c_double),
    ('c', ctypes.POINTER(ctypes.c_double)),
]

struct_c__SA_mhd_stat_data._pack_ = 1 # source:False
struct_c__SA_mhd_stat_data._fields_ = [
    ('n_modes', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('rho_min', ctypes.c_double),
    ('rho_max', ctypes.c_double),
    ('nmode', ctypes.c_int32 * 512),
    ('mmode', ctypes.c_int32 * 512),
    ('amplitude_nm', ctypes.c_double * 512),
    ('omega_nm', ctypes.c_double * 512),
    ('phase_nm', ctypes.c_double * 512),
    ('alpha_nm', struct_c__SA_interp1D_data * 512),
    ('phi_nm', struct_c__SA_interp1D_data * 512),
]

class struct_c__SA_mhd_nonstat_data(Structure):
    pass

struct_c__SA_mhd_nonstat_data._pack_ = 1 # source:False
struct_c__SA_mhd_nonstat_data._fields_ = [
    ('n_modes', ctypes.c_int32),
    ('nmode', ctypes.c_int32 * 512),
    ('mmode', ctypes.c_int32 * 512),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('amplitude_nm', ctypes.c_double * 512),
    ('omega_nm', ctypes.c_double * 512),
    ('phase_nm', ctypes.c_double * 512),
    ('alpha_nm', struct_c__SA_interp2D_data * 512),
    ('phi_nm', struct_c__SA_interp2D_data * 512),
]

struct_c__SA_mhd_data._pack_ = 1 # source:False
struct_c__SA_mhd_data._fields_ = [
    ('type', mhd_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('stat', struct_c__SA_mhd_stat_data),
    ('nonstat', struct_c__SA_mhd_nonstat_data),
]

class struct_c__SA_asigma_data(Structure):
    pass

class struct_c__SA_asigma_loc_data(Structure):
    pass

struct_c__SA_asigma_loc_data._pack_ = 1 # source:False
struct_c__SA_asigma_loc_data._fields_ = [
    ('N_reac', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('z_1', ctypes.POINTER(ctypes.c_int32)),
    ('a_1', ctypes.POINTER(ctypes.c_int32)),
    ('z_2', ctypes.POINTER(ctypes.c_int32)),
    ('a_2', ctypes.POINTER(ctypes.c_int32)),
    ('reac_type', ctypes.POINTER(ctypes.c_int32)),
    ('reac_avail', ctypes.POINTER(ctypes.c_int32)),
    ('sigma', ctypes.POINTER(struct_c__SA_interp1D_data)),
    ('sigmav', ctypes.POINTER(struct_c__SA_interp2D_data)),
    ('BMSsigmav', ctypes.POINTER(struct_c__SA_interp3D_data)),
]

struct_c__SA_asigma_data._pack_ = 1 # source:False
struct_c__SA_asigma_data._fields_ = [
    ('type', asigma_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('asigma_loc', struct_c__SA_asigma_loc_data),
]

class struct_c__SA_E_field_data(Structure):
    pass

class struct_c__SA_E_TC_data(Structure):
    pass

struct_c__SA_E_TC_data._pack_ = 1 # source:False
struct_c__SA_E_TC_data._fields_ = [
    ('Exyz', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_E_1DS_data(Structure):
    _pack_ = 1 # source:False
    _fields_ = [
    ('dV', struct_c__SA_interp1D_data),
     ]

struct_c__SA_E_field_data._pack_ = 1 # source:False
struct_c__SA_E_field_data._fields_ = [
    ('type', E_field_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('ETC', struct_c__SA_E_TC_data),
    ('E1DS', struct_c__SA_E_1DS_data),
]

class struct_c__SA_boozer_data(Structure):
    pass

struct_c__SA_boozer_data._pack_ = 1 # source:False
struct_c__SA_boozer_data._fields_ = [
    ('psi0', ctypes.c_double),
    ('psi1', ctypes.c_double),
    ('r0', ctypes.c_double),
    ('z0', ctypes.c_double),
    ('rs', ctypes.POINTER(ctypes.c_double)),
    ('zs', ctypes.POINTER(ctypes.c_double)),
    ('nrzs', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('nu_psitheta', struct_c__SA_interp2D_data),
    ('theta_psithetageom', struct_c__SA_interp2D_data),
    ('psi_rz', struct_c__SA_interp2D_data),
]

class struct_c__SA_plasma_data(Structure):
    pass

class struct_c__SA_plasma_1DS_data(Structure):
    pass

struct_c__SA_plasma_1DS_data._pack_ = 1 # source:False
struct_c__SA_plasma_1DS_data._fields_ = [
    ('n_species', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('mass', ctypes.c_double * 8),
    ('charge', ctypes.c_double * 8),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('temp', struct_c__SA_interp1D_data * 2),
    ('dens', struct_c__SA_interp1D_data * 8),
]

class struct_c__SA_plasma_1D_data(Structure):
    pass

struct_c__SA_plasma_1D_data._pack_ = 1 # source:False
struct_c__SA_plasma_1D_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('n_species', ctypes.c_int32),
    ('mass', ctypes.c_double * 8),
    ('charge', ctypes.c_double * 8),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('rho', ctypes.POINTER(ctypes.c_double)),
    ('temp', ctypes.POINTER(ctypes.c_double)),
    ('dens', ctypes.POINTER(ctypes.c_double)),
]

class struct_c__SA_plasma_1Dt_data(Structure):
    pass

struct_c__SA_plasma_1Dt_data._pack_ = 1 # source:False
struct_c__SA_plasma_1Dt_data._fields_ = [
    ('n_rho', ctypes.c_int32),
    ('n_time', ctypes.c_int32),
    ('n_species', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('mass', ctypes.c_double * 8),
    ('charge', ctypes.c_double * 8),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('rho', ctypes.POINTER(ctypes.c_double)),
    ('time', ctypes.POINTER(ctypes.c_double)),
    ('temp', ctypes.POINTER(ctypes.c_double)),
    ('dens', ctypes.POINTER(ctypes.c_double)),
]

struct_c__SA_plasma_data._pack_ = 1 # source:False
struct_c__SA_plasma_data._fields_ = [
    ('type', plasma_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('plasma_1D', struct_c__SA_plasma_1D_data),
    ('plasma_1Dt', struct_c__SA_plasma_1Dt_data),
    ('plasma_1DS', struct_c__SA_plasma_1DS_data),
]

class struct_c__SA_neutral_data(Structure):
    pass

class struct_c__SA_N0_3D_data(Structure):
    pass

class struct_c__SA_linint3D_data(Structure):
    pass

struct_c__SA_linint3D_data._pack_ = 1 # source:False
struct_c__SA_linint3D_data._fields_ = [
    ('n_x', ctypes.c_int32),
    ('n_y', ctypes.c_int32),
    ('n_z', ctypes.c_int32),
    ('bc_x', ctypes.c_int32),
    ('bc_y', ctypes.c_int32),
    ('bc_z', ctypes.c_int32),
    ('x_min', ctypes.c_double),
    ('x_max', ctypes.c_double),
    ('x_grid', ctypes.c_double),
    ('y_min', ctypes.c_double),
    ('y_max', ctypes.c_double),
    ('y_grid', ctypes.c_double),
    ('z_min', ctypes.c_double),
    ('z_max', ctypes.c_double),
    ('z_grid', ctypes.c_double),
    ('c', ctypes.POINTER(ctypes.c_double)),
]

struct_c__SA_N0_3D_data._pack_ = 1 # source:False
struct_c__SA_N0_3D_data._fields_ = [
    ('n_species', ctypes.c_int32),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('maxwellian', ctypes.c_int32 * 8),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('n0', struct_c__SA_linint3D_data * 8),
    ('t0', struct_c__SA_linint3D_data * 8),
]

class struct_c__SA_N0_1D_data(Structure):
    pass

struct_c__SA_N0_1D_data._pack_ = 1 # source:False
struct_c__SA_N0_1D_data._fields_ = [
    ('n_species', ctypes.c_int32),
    ('anum', ctypes.c_int32 * 8),
    ('znum', ctypes.c_int32 * 8),
    ('maxwellian', ctypes.c_int32 * 8),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('n0', struct_c__SA_linint1D_data * 8),
    ('t0', struct_c__SA_linint1D_data * 8),
]

struct_c__SA_neutral_data._pack_ = 1 # source:False
struct_c__SA_neutral_data._fields_ = [
    ('type', neutral_type),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('N01D', struct_c__SA_N0_1D_data),
    ('N03D', struct_c__SA_N0_3D_data),
]

class struct_c__SA_mccc_data(Structure):
    pass

struct_c__SA_mccc_data._pack_ = 1 # source:False
struct_c__SA_mccc_data._fields_ = [
    ('usetabulated', ctypes.c_int32),
    ('include_energy', ctypes.c_int32),
    ('include_pitch', ctypes.c_int32),
    ('include_gcdiff', ctypes.c_int32),
]

struct_c__SA_sim_data._pack_ = 1 # source:False
struct_c__SA_sim_data._fields_ = [
    ('B_data', B_field_data),
    ('E_data', struct_c__SA_E_field_data),
    ('plasma_data', struct_c__SA_plasma_data),
    ('neutral_data', struct_c__SA_neutral_data),
    ('wall_data', wall_data),
    ('boozer_data', struct_c__SA_boozer_data),
    ('mhd_data', struct_c__SA_mhd_data),
    ('asigma_data', struct_c__SA_asigma_data),
    ('diag_data', diag_data),
    ('random_data', ctypes.POINTER(None)),
    ('mccc_data', struct_c__SA_mccc_data),
    ('sim_mode', ctypes.c_int32),
    ('enable_ada', ctypes.c_int32),
    ('record_mode', ctypes.c_int32),
    ('fix_usrdef_use', ctypes.c_int32),
    ('fix_usrdef_val', ctypes.c_double),
    ('fix_gyrodef_nstep', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('ada_tol_orbfol', ctypes.c_double),
    ('ada_tol_clmbcol', ctypes.c_double),
    ('ada_max_drho', ctypes.c_double),
    ('ada_max_dphi', ctypes.c_double),
    ('enable_orbfol', ctypes.c_int32),
    ('enable_clmbcol', ctypes.c_int32),
    ('enable_mhd', ctypes.c_int32),
    ('enable_atomic', ctypes.c_int32),
    ('disable_gctransform', ctypes.c_int32),
    ('disable_energyccoll', ctypes.c_int32),
    ('disable_pitchccoll', ctypes.c_int32),
    ('disable_gcdiffccoll', ctypes.c_int32),
    ('reverse_time', ctypes.c_int32),
    ('endcond_active', ctypes.c_int32),
    ('endcond_lim_simtime', ctypes.c_double),
    ('endcond_max_mileage', ctypes.c_double),
    ('endcond_max_cputime', ctypes.c_double),
    ('endcond_min_rho', ctypes.c_double),
    ('endcond_max_rho', ctypes.c_double),
    ('endcond_min_ekin', ctypes.c_double),
    ('endcond_min_thermal', ctypes.c_double),
    ('endcond_max_tororb', ctypes.c_double),
    ('endcond_max_polorb', ctypes.c_double),
    ('endcond_torandpol', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
]

sim_data = struct_c__SA_sim_data
simulate_init_offload = _libraries['libascot.so'].simulate_init_offload
simulate_init_offload.restype = None
simulate_init_offload.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data)]
simulate = _libraries['libascot.so'].simulate
simulate.restype = None
simulate.argtypes = [ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.POINTER(struct_c__SA_offload_package), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_double)]

# values for enumeration 'c__Ea_hdf5_input_options'
c__Ea_hdf5_input_options__enumvalues = {
    1: 'hdf5_input_options',
    2: 'hdf5_input_bfield',
    4: 'hdf5_input_efield',
    8: 'hdf5_input_plasma',
    16: 'hdf5_input_neutral',
    32: 'hdf5_input_wall',
    64: 'hdf5_input_marker',
    128: 'hdf5_input_boozer',
    256: 'hdf5_input_mhd',
    512: 'hdf5_input_asigma',
}
hdf5_input_options = 1
hdf5_input_bfield = 2
hdf5_input_efield = 4
hdf5_input_plasma = 8
hdf5_input_neutral = 16
hdf5_input_wall = 32
hdf5_input_marker = 64
hdf5_input_boozer = 128
hdf5_input_mhd = 256
hdf5_input_asigma = 512
c__Ea_hdf5_input_options = ctypes.c_uint32 # enum
hdf5_interface_read_input = _libraries['libascot.so'].hdf5_interface_read_input
hdf5_interface_read_input.restype = ctypes.c_int32
hdf5_interface_read_input.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.c_int32, ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(struct_c__SA_input_particle)), ctypes.POINTER(ctypes.c_int32)]
hdf5_interface_init_results = _libraries['libascot.so'].hdf5_interface_init_results
hdf5_interface_init_results.restype = ctypes.c_int32
hdf5_interface_init_results.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.POINTER(ctypes.c_char)]
hdf5_interface_write_state = _libraries['libascot.so'].hdf5_interface_write_state
hdf5_interface_write_state.restype = ctypes.c_int32
hdf5_interface_write_state.argtypes = [ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), integer, ctypes.POINTER(struct_c__SA_particle_state)]
hdf5_interface_write_diagnostics = _libraries['libascot.so'].hdf5_interface_write_diagnostics
hdf5_interface_write_diagnostics.restype = ctypes.c_int32
hdf5_interface_write_diagnostics.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_char)]
hid_t = ctypes.c_int64
hdf5_get_active_qid = _libraries['libascot.so'].hdf5_get_active_qid
hdf5_get_active_qid.restype = ctypes.c_int32
hdf5_get_active_qid.argtypes = [hid_t, ctypes.POINTER(ctypes.c_char), ctypes.c_char * 11]
hdf5_generate_qid = _libraries['libascot.so'].hdf5_generate_qid
hdf5_generate_qid.restype = None
hdf5_generate_qid.argtypes = [ctypes.POINTER(ctypes.c_char)]
mpi_interface_init = _libraries['libascot.so'].mpi_interface_init
mpi_interface_init.restype = None
mpi_interface_init.argtypes = [ctypes.c_int32, ctypes.POINTER(ctypes.POINTER(ctypes.c_char)), ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32)]
mpi_interface_finalize = _libraries['libascot.so'].mpi_interface_finalize
mpi_interface_finalize.restype = None
mpi_interface_finalize.argtypes = []
mpi_my_particles = _libraries['libascot.so'].mpi_my_particles
mpi_my_particles.restype = None
mpi_my_particles.argtypes = [ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_int32), ctypes.c_int32, ctypes.c_int32, ctypes.c_int32]
mpi_gather_particlestate = _libraries['libascot.so'].mpi_gather_particlestate
mpi_gather_particlestate.restype = None
mpi_gather_particlestate.argtypes = [ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(ctypes.POINTER(struct_c__SA_particle_state)), ctypes.POINTER(ctypes.c_int32), ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32]
mpi_gather_diag = _libraries['libascot.so'].mpi_gather_diag
mpi_gather_diag.restype = None
mpi_gather_diag.argtypes = [ctypes.POINTER(struct_c__SA_diag_offload_data), ctypes.POINTER(ctypes.c_double), ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32]
pack_offload_array = _libraries['libascot.so'].pack_offload_array
pack_offload_array.restype = ctypes.c_int32
pack_offload_array.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.POINTER(struct_c__SA_offload_package), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32))]
prepare_markers = _libraries['libascot.so'].prepare_markers
prepare_markers.restype = ctypes.c_int32
prepare_markers.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(ctypes.POINTER(struct_c__SA_input_particle)), ctypes.POINTER(ctypes.POINTER(struct_c__SA_particle_state)), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.c_double)]
write_rungroup = _libraries['libascot.so'].write_rungroup
write_rungroup.restype = ctypes.c_int32
write_rungroup.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(ctypes.c_char)]
offload_and_simulate = _libraries['libascot.so'].offload_and_simulate
offload_and_simulate.restype = ctypes.c_int32
offload_and_simulate.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_state), ctypes.POINTER(struct_c__SA_offload_package), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_int32), ctypes.POINTER(ctypes.POINTER(struct_c__SA_particle_state)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]
write_output = _libraries['libascot.so'].write_output
write_output.restype = ctypes.c_int32
write_output.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(struct_c__SA_particle_state), ctypes.c_int32, ctypes.POINTER(ctypes.c_double)]
free_simulation_output = _libraries['FIXME_STUB'].free_simulation_output
free_simulation_output.restype = ctypes.c_int32
free_simulation_output.argtypes = [ctypes.POINTER(struct_c__SA_sim_offload_data), ctypes.POINTER(ctypes.POINTER(struct_c__SA_particle_state)), ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]
print_marker_summary = _libraries['libascot.so'].print_marker_summary
print_marker_summary.restype = None
print_marker_summary.argtypes = [ctypes.POINTER(struct_c__SA_particle_state), ctypes.c_int32]
libascot_allocate_input_particles = _libraries['libascot.so'].libascot_allocate_input_particles
libascot_allocate_input_particles.restype = ctypes.POINTER(struct_c__SA_input_particle)
libascot_allocate_input_particles.argtypes = [ctypes.c_int32]
libascot_allocate_particle_states = _libraries['libascot.so'].libascot_allocate_particle_states
libascot_allocate_particle_states.restype = ctypes.POINTER(struct_c__SA_particle_state)
libascot_allocate_particle_states.argtypes = [ctypes.c_int32]
libascot_allocate_reals = _libraries['libascot.so'].libascot_allocate_reals
libascot_allocate_reals.restype = ctypes.POINTER(ctypes.c_double)
libascot_allocate_reals.argtypes = [size_t]
libascot_deallocate = _libraries['libascot.so'].libascot_deallocate
libascot_deallocate.restype = None
libascot_deallocate.argtypes = [ctypes.POINTER(None)]

# values for enumeration 'c__Ea_endcond_tlim'
c__Ea_endcond_tlim__enumvalues = {
    1: 'endcond_tlim',
    2: 'endcond_emin',
    4: 'endcond_therm',
    8: 'endcond_wall',
    16: 'endcond_rhomin',
    32: 'endcond_rhomax',
    64: 'endcond_polmax',
    128: 'endcond_tormax',
    256: 'endcond_cpumax',
    512: 'endcond_hybrid',
    1024: 'endcond_neutr',
    2048: 'endcond_ioniz',
}
endcond_tlim = 1
endcond_emin = 2
endcond_therm = 4
endcond_wall = 8
endcond_rhomin = 16
endcond_rhomax = 32
endcond_polmax = 64
endcond_tormax = 128
endcond_cpumax = 256
endcond_hybrid = 512
endcond_neutr = 1024
endcond_ioniz = 2048
c__Ea_endcond_tlim = ctypes.c_uint32 # enum
endcond_check_gc = _libraries['libascot.so'].endcond_check_gc
endcond_check_gc.restype = None
endcond_check_gc.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.POINTER(struct_c__SA_particle_simd_gc), ctypes.POINTER(struct_c__SA_sim_data)]
endcond_check_fo = _libraries['libascot.so'].endcond_check_fo
endcond_check_fo.restype = None
endcond_check_fo.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.POINTER(struct_c__SA_particle_simd_fo), ctypes.POINTER(struct_c__SA_sim_data)]
endcond_check_ml = _libraries['libascot.so'].endcond_check_ml
endcond_check_ml.restype = None
endcond_check_ml.argtypes = [ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.POINTER(struct_c__SA_particle_simd_ml), ctypes.POINTER(struct_c__SA_sim_data)]
endcond_parse = _libraries['libascot.so'].endcond_parse
endcond_parse.restype = None
endcond_parse.argtypes = [ctypes.c_int32, ctypes.POINTER(ctypes.c_int32)]
endcond_parse2str = _libraries['libascot.so'].endcond_parse2str
endcond_parse2str.restype = None
endcond_parse2str.argtypes = [ctypes.c_int32, ctypes.POINTER(ctypes.c_char)]
hdf5_wall_init_offload = _libraries['libascot.so'].hdf5_wall_init_offload
hdf5_wall_init_offload.restype = ctypes.c_int32
hdf5_wall_init_offload.argtypes = [hid_t, ctypes.POINTER(struct_c__SA_wall_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.POINTER(ctypes.c_int32)), ctypes.POINTER(ctypes.c_char)]
hdf5_wall_2d_to_offload = _libraries['libascot.so'].hdf5_wall_2d_to_offload
hdf5_wall_2d_to_offload.restype = ctypes.c_int32
hdf5_wall_2d_to_offload.argtypes = [ctypes.POINTER(struct_c__SA_wall_2d_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.c_int32, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)]
hdf5_wall_3d_to_offload = _libraries['libascot.so'].hdf5_wall_3d_to_offload
hdf5_wall_3d_to_offload.restype = ctypes.c_int32
hdf5_wall_3d_to_offload.argtypes = [ctypes.POINTER(struct_c__SA_wall_3d_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.c_int32, ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double)]
hdf5_bfield_init_offload = _libraries['libascot.so'].hdf5_bfield_init_offload
hdf5_bfield_init_offload.restype = ctypes.c_int32
hdf5_bfield_init_offload.argtypes = [hid_t, ctypes.POINTER(struct_c__SA_B_field_offload_data), ctypes.POINTER(ctypes.POINTER(ctypes.c_double)), ctypes.POINTER(ctypes.c_char)]
class struct_c__SA_afsi_thermal_data(Structure):
    pass

struct_c__SA_afsi_thermal_data._pack_ = 1 # source:False
struct_c__SA_afsi_thermal_data._fields_ = [
    ('n_r', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('min_r', ctypes.c_double),
    ('max_r', ctypes.c_double),
    ('n_phi', ctypes.c_int32),
    ('PADDING_1', ctypes.c_ubyte * 4),
    ('min_phi', ctypes.c_double),
    ('max_phi', ctypes.c_double),
    ('n_z', ctypes.c_int32),
    ('PADDING_2', ctypes.c_ubyte * 4),
    ('min_z', ctypes.c_double),
    ('max_z', ctypes.c_double),
    ('temperature', ctypes.POINTER(ctypes.c_double)),
    ('density', ctypes.POINTER(ctypes.c_double)),
]

afsi_thermal_data = struct_c__SA_afsi_thermal_data
class struct_c__SA_afsi_data(Structure):
    pass

struct_c__SA_afsi_data._pack_ = 1 # source:False
struct_c__SA_afsi_data._fields_ = [
    ('type', ctypes.c_int32),
    ('PADDING_0', ctypes.c_ubyte * 4),
    ('dist_5D', ctypes.POINTER(struct_c__SA_dist_5D_data)),
    ('dist_thermal', ctypes.POINTER(struct_c__SA_afsi_thermal_data)),
]

afsi_data = struct_c__SA_afsi_data
afsi_run = _libraries['libascot.so'].afsi_run
afsi_run.restype = None
afsi_run.argtypes = [ctypes.c_int32, ctypes.c_int32, ctypes.POINTER(struct_c__SA_afsi_data), ctypes.POINTER(struct_c__SA_afsi_data), ctypes.POINTER(struct_c__SA_dist_5D_data), ctypes.POINTER(struct_c__SA_dist_5D_data)]
afsi_test_dist = _libraries['libascot.so'].afsi_test_dist
afsi_test_dist.restype = None
afsi_test_dist.argtypes = [ctypes.POINTER(struct_c__SA_dist_5D_data)]
afsi_test_thermal = _libraries['libascot.so'].afsi_test_thermal
afsi_test_thermal.restype = None
afsi_test_thermal.argtypes = []
__all__ = \
    ['B_STS_data', 'B_STS_eval_B', 'B_STS_eval_B_dB',
    'B_STS_eval_psi', 'B_STS_eval_psi_dpsi', 'B_STS_eval_rho_drho',
    'B_STS_free_offload', 'B_STS_get_axis_rz', 'B_STS_init',
    'B_STS_init_offload', 'B_STS_offload_data', 'B_field_data',
    'B_field_eval_B', 'B_field_eval_B_dB', 'B_field_eval_psi',
    'B_field_eval_psi_dpsi', 'B_field_eval_rho',
    'B_field_eval_rho_drho', 'B_field_free_offload',
    'B_field_get_axis_rz', 'B_field_init', 'B_field_init_offload',
    'B_field_offload_data', 'B_field_type', 'B_field_type_2DS',
    'B_field_type_3DS', 'B_field_type_GS', 'B_field_type_STS',
    'B_field_type_TC', 'E_field_type', 'E_field_type_1DS',
    'E_field_type_TC', 'a5err', 'afsi_data', 'afsi_run',
    'afsi_test_dist', 'afsi_test_thermal', 'afsi_thermal_data',
    'asigma_type', 'asigma_type_loc', 'c__Ea_endcond_tlim',
    'c__Ea_hdf5_input_options', 'c__Ea_simulate_mode_fo', 'diag_data',
    'diag_free', 'diag_free_offload', 'diag_init',
    'diag_init_offload', 'diag_offload_data',
    'diag_orb_check_plane_crossing', 'diag_orb_check_radial_crossing',
    'diag_orb_data', 'diag_orb_free', 'diag_orb_init',
    'diag_orb_offload_data', 'diag_orb_update_fo',
    'diag_orb_update_gc', 'diag_orb_update_ml', 'diag_sum',
    'diag_update_fo', 'diag_update_gc', 'diag_update_ml',
    'endcond_check_fo', 'endcond_check_gc', 'endcond_check_ml',
    'endcond_cpumax', 'endcond_emin', 'endcond_hybrid',
    'endcond_ioniz', 'endcond_neutr', 'endcond_parse',
    'endcond_parse2str', 'endcond_polmax', 'endcond_rhomax',
    'endcond_rhomin', 'endcond_therm', 'endcond_tlim',
    'endcond_tormax', 'endcond_wall', 'free_simulation_output',
    'hdf5_bfield_init_offload', 'hdf5_generate_qid',
    'hdf5_get_active_qid', 'hdf5_input_asigma', 'hdf5_input_bfield',
    'hdf5_input_boozer', 'hdf5_input_efield', 'hdf5_input_marker',
    'hdf5_input_mhd', 'hdf5_input_neutral', 'hdf5_input_options',
    'hdf5_input_plasma', 'hdf5_input_wall',
    'hdf5_interface_init_results', 'hdf5_interface_read_input',
    'hdf5_interface_write_diagnostics', 'hdf5_interface_write_state',
    'hdf5_wall_2d_to_offload', 'hdf5_wall_3d_to_offload',
    'hdf5_wall_init_offload', 'hid_t', 'input_particle',
    'input_particle_type', 'input_particle_type_gc',
    'input_particle_type_ml', 'input_particle_type_p',
    'input_particle_type_s', 'integer',
    'libascot_allocate_input_particles',
    'libascot_allocate_particle_states', 'libascot_allocate_reals',
    'libascot_deallocate', 'mhd_type', 'mhd_type_nonstat',
    'mhd_type_stat', 'mpi_gather_diag', 'mpi_gather_particlestate',
    'mpi_interface_finalize', 'mpi_interface_init',
    'mpi_my_particles', 'neutral_type', 'neutral_type_1D',
    'neutral_type_3D', 'offload_and_simulate', 'offload_free_offload',
    'offload_init_offload', 'offload_pack', 'offload_package',
    'offload_unpack', 'pack_offload_array', 'particle',
    'particle_copy_fo', 'particle_copy_gc', 'particle_copy_ml',
    'particle_cycle_fo', 'particle_cycle_gc', 'particle_cycle_ml',
    'particle_fo_to_gc', 'particle_fo_to_state', 'particle_gc',
    'particle_gc_to_state', 'particle_input_to_state', 'particle_ml',
    'particle_ml_to_state', 'particle_queue', 'particle_simd_fo',
    'particle_simd_gc', 'particle_simd_ml', 'particle_state',
    'particle_state_to_fo', 'particle_state_to_gc',
    'particle_state_to_ml', 'particle_to_fo_dummy',
    'particle_to_gc_dummy', 'particle_to_ml_dummy', 'plasma_type',
    'plasma_type_1D', 'plasma_type_1DS', 'plasma_type_1Dt',
    'prepare_markers', 'print_marker_summary', 'real', 'sim_data',
    'sim_offload_data', 'simulate', 'simulate_init_offload',
    'simulate_mode_fo', 'simulate_mode_gc', 'simulate_mode_hybrid',
    'simulate_mode_ml', 'size_t', 'struct_c__SA_B_2DS_data',
    'struct_c__SA_B_2DS_offload_data', 'struct_c__SA_B_3DS_data',
    'struct_c__SA_B_3DS_offload_data', 'struct_c__SA_B_GS_data',
    'struct_c__SA_B_GS_offload_data', 'struct_c__SA_B_STS_data',
    'struct_c__SA_B_STS_offload_data', 'struct_c__SA_B_TC_data',
    'struct_c__SA_B_TC_offload_data', 'struct_c__SA_B_field_data',
    'struct_c__SA_B_field_offload_data', 'struct_c__SA_E_1DS_data',
    'struct_c__SA_E_1DS_offload_data', 'struct_c__SA_E_TC_data',
    'struct_c__SA_E_TC_offload_data', 'struct_c__SA_E_field_data',
    'struct_c__SA_E_field_offload_data', 'struct_c__SA_N0_1D_data',
    'struct_c__SA_N0_1D_offload_data', 'struct_c__SA_N0_3D_data',
    'struct_c__SA_N0_3D_offload_data', 'struct_c__SA_afsi_data',
    'struct_c__SA_afsi_thermal_data', 'struct_c__SA_asigma_data',
    'struct_c__SA_asigma_loc_data',
    'struct_c__SA_asigma_loc_offload_data',
    'struct_c__SA_asigma_offload_data', 'struct_c__SA_boozer_data',
    'struct_c__SA_boozer_offload_data', 'struct_c__SA_diag_data',
    'struct_c__SA_diag_offload_data', 'struct_c__SA_diag_orb_data',
    'struct_c__SA_diag_orb_offload_data',
    'struct_c__SA_diag_transcoef_data',
    'struct_c__SA_diag_transcoef_offload_data',
    'struct_c__SA_dist_5D_data', 'struct_c__SA_dist_5D_offload_data',
    'struct_c__SA_dist_6D_data', 'struct_c__SA_dist_6D_offload_data',
    'struct_c__SA_dist_COM_data',
    'struct_c__SA_dist_COM_offload_data',
    'struct_c__SA_dist_rho5D_data',
    'struct_c__SA_dist_rho5D_offload_data',
    'struct_c__SA_dist_rho6D_data',
    'struct_c__SA_dist_rho6D_offload_data',
    'struct_c__SA_input_particle', 'struct_c__SA_interp1D_data',
    'struct_c__SA_interp2D_data', 'struct_c__SA_interp3D_data',
    'struct_c__SA_linint1D_data', 'struct_c__SA_linint3D_data',
    'struct_c__SA_mccc_data', 'struct_c__SA_mhd_data',
    'struct_c__SA_mhd_nonstat_data',
    'struct_c__SA_mhd_nonstat_offload_data',
    'struct_c__SA_mhd_offload_data', 'struct_c__SA_mhd_stat_data',
    'struct_c__SA_mhd_stat_offload_data', 'struct_c__SA_neutral_data',
    'struct_c__SA_neutral_offload_data',
    'struct_c__SA_offload_package', 'struct_c__SA_particle',
    'struct_c__SA_particle_gc', 'struct_c__SA_particle_ml',
    'struct_c__SA_particle_queue', 'struct_c__SA_particle_simd_fo',
    'struct_c__SA_particle_simd_gc', 'struct_c__SA_particle_simd_ml',
    'struct_c__SA_particle_state', 'struct_c__SA_plasma_1DS_data',
    'struct_c__SA_plasma_1DS_offload_data',
    'struct_c__SA_plasma_1D_data',
    'struct_c__SA_plasma_1D_offload_data',
    'struct_c__SA_plasma_1Dt_data',
    'struct_c__SA_plasma_1Dt_offload_data',
    'struct_c__SA_plasma_data', 'struct_c__SA_plasma_offload_data',
    'struct_c__SA_sim_data', 'struct_c__SA_sim_offload_data',
    'struct_c__SA_wall_2d_data', 'struct_c__SA_wall_2d_offload_data',
    'struct_c__SA_wall_3d_data', 'struct_c__SA_wall_3d_offload_data',
    'struct_c__SA_wall_data', 'struct_c__SA_wall_offload_data',
    'struct_diag_transcoef_link', 'union_c__SA_input_particle_0',
    'wall_data', 'wall_free_offload', 'wall_hit_wall', 'wall_init',
    'wall_init_offload', 'wall_offload_data', 'wall_type',
    'wall_type_2D', 'wall_type_3D', 'write_output', 'write_rungroup']

a5py/a5py/ascotpy/ascot2py.py
ascot4fusion/ascot5
"""Python side of the interactive Python interface.

This module defines LibAscot class whose methods can be used in Python to call
Ascot5 functions (written in C) directly. The callable functions are defined in
library module libascot.c which must be compiled first with make libascot. This
module acts as a wrapper for those functions. More advanced functionality should
be implemented in other modules.
"""
import ctypes
import sys
import os
import warnings
import numpy as np
import unyt

from scipy import interpolate
from numpy.ctypeslib import ndpointer

from a5py.physlib import parseunits

try:
    from . import ascot2py

    def _ndpointerwithnull(*args, **kwargs):
        """This produces ndpointer which can also be a NULL pointer.

        This wrapper is required since otherwise it is not possible to
        supply None (indicating NULL) as an argument that requires
        a numpy array.
        """
        base = ndpointer(*args, **kwargs)
        def from_param(cls, obj):
            if obj is None:
                return obj
            return base.from_param(obj)
        return type(base.__name__, (base,),
                    {'from_param': classmethod(from_param)})

    PTR_REAL = _ndpointerwithnull(ctypes.c_double, flags="C_CONTIGUOUS")
    PTR_INT  = _ndpointerwithnull(ctypes.c_int,    flags="C_CONTIGUOUS")
    PTR_SIM  = ctypes.POINTER(ascot2py.struct_c__SA_sim_offload_data)
    PTR_ARR  = ctypes.POINTER(ctypes.c_double)
    _LIBASCOT = ascot2py._libraries['libascot.so']

except OSError as error:
    _LIBASCOT = None
    msg = """
    Warning: Failed to import libascot.so. Some functionalities of Ascot
    are not available. Verify that libascot.so has been compiled, it can be
    found in LD_LIBRARY_PATH, and dependencies (e.g. HDF5) are available.
    """
    warnings.warn(msg)

class LibAscot:
    """Python wrapper of libascot.so.
    """

    def _eval_bfield(self, r, phi, z, t, evalb=False, evalrho=False,
                     evalaxis=False):
        """Evaluate magnetic field quantities at given coordinates.

        Parameters
        ----------
        r : array_like, (n,)
            R coordinates where data is evaluated [m].
        phi : array_like (n,)
            phi coordinates where data is evaluated [rad].
        z : array_like (n,)
            z coordinates where data is evaluated [m].
        t : array_like (n,)
            Time coordinates where data is evaluated [s].
        evalb : bool, optional
            Evaluate B_field_eval_B_dB.
        evalrho : bool, optional
            Evaluate B_field_eval_rho.
        evalaxis : bool, optional
            Evaluate B_field_get_axis.

        Returns
        -------
        out : dict [str, array_like (n,)]
            Evaluated quantities in a dictionary.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield")
        Neval = r.size
        out = {}

        if evalb:
            Tperm = unyt.T / unyt.m
            out["br"]       = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.T
            out["bphi"]     = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.T
            out["bz"]       = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.T
            out["brdr"]     = (np.zeros(r.shape, dtype="f8") + np.nan) * Tperm
            out["brdphi"]   = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.T
            out["brdz"]     = (np.zeros(r.shape, dtype="f8") + np.nan) * Tperm
            out["bphidr"]   = (np.zeros(r.shape, dtype="f8") + np.nan) * Tperm
            out["bphidphi"] = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.T
            out["bphidz"]   = (np.zeros(r.shape, dtype="f8") + np.nan) * Tperm
            out["bzdr"]     = (np.zeros(r.shape, dtype="f8") + np.nan) * Tperm
            out["bzdphi"]   = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.T
            out["bzdz"]     = (np.zeros(r.shape, dtype="f8") + np.nan) * Tperm

            fun = _LIBASCOT.libascot_B_field_eval_B_dB
            fun.restype  = None
            fun.argtypes = [PTR_SIM, PTR_ARR,
                            ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL]

            fun(ctypes.byref(self._sim), self._bfield_offload_array,
                Neval, r, phi, z, t, out["br"], out["bphi"], out["bz"],
                out["brdr"], out["brdphi"], out["brdz"], out["bphidr"],
                out["bphidphi"], out["bphidz"], out["bzdr"], out["bzdphi"],
                out["bzdz"])

        if evalrho:
            out["psi"] = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.Wb
            out["rho"] = (np.zeros(r.shape, dtype="f8") + np.nan) \
                * unyt.dimensionless
            out["rhodpsi"] = (np.zeros(r.shape, dtype="f8") + np.nan) / unyt.Wb
            out["psidr"] = (np.zeros(r.shape, dtype="f8") + np.nan) \
                * unyt.Wb / unyt.m
            out["psidphi"] = (np.zeros(r.shape, dtype="f8") + np.nan) \
                * unyt.Wb
            out["psidz"] = (np.zeros(r.shape, dtype="f8") + np.nan) \
                * unyt.Wb / unyt.m

            fun = _LIBASCOT.libascot_B_field_eval_rho
            fun.restype  = None
            fun.argtypes = [PTR_SIM, PTR_ARR,
                            ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL]

            fun(ctypes.byref(self._sim), self._bfield_offload_array,
                Neval, r, phi, z, t, out["rho"], out["rhodpsi"], out["psi"],
                out["psidr"], out["psidphi"], out["psidz"])

        if evalaxis:
            out["axisr"] = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.m
            out["axisz"] = (np.zeros(r.shape, dtype="f8") + np.nan) * unyt.m

            fun = _LIBASCOT.libascot_B_field_get_axis
            fun.restype  = None
            fun.argtypes = [PTR_SIM, PTR_ARR,
                            ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL]

            fun(ctypes.byref(self._sim), self._bfield_offload_array,
                Neval, phi, out["axisr"], out["axisz"])

        return out

    def _eval_efield(self, r, phi, z, t):
        """Evaluate electric field quantities at given coordinates.

        Parameters
        ----------
        r : array_like, (n,)
            R coordinates where data is evaluated [m].
        phi : array_like (n,)
            phi coordinates where data is evaluated [rad].
        z : array_like (n,)
            z coordinates where data is evaluated [m].
        t : array_like (n,)
            Time coordinates where data is evaluated [s].

        Returns
        -------
        out : dict [str, array_like (n,)]
            Evaluated quantities in a dictionary.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield", "efield")
        Neval = r.size
        out = {}
        Vperm = unyt.V / unyt.m
        out["er"]   = (np.zeros(r.shape, dtype="f8") + np.nan) * Vperm
        out["ephi"] = (np.zeros(r.shape, dtype="f8") + np.nan) * Vperm
        out["ez"]   = (np.zeros(r.shape, dtype="f8") + np.nan) * Vperm

        fun = _LIBASCOT.libascot_E_field_eval_E
        fun.restype  = None
        fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        PTR_REAL, PTR_REAL, PTR_REAL]
        fun(ctypes.byref(self._sim), self._efield_offload_array,
            self._bfield_offload_array,
            Neval, r, phi, z, t, out["er"], out["ephi"], out["ez"])

        return out

    def _eval_plasma(self, r, phi, z, t):
        """Evaluate plasma quantities at given coordinates.

        Parameters
        ----------
        r : array_like, (n,)
            R coordinates where data is evaluated [m].
        phi : array_like (n,)
            phi coordinates where data is evaluated [rad].
        z : array_like (n,)
            z coordinates where data is evaluated [m].
        t : array_like (n,)
            Time coordinates where data is evaluated [s].

        Returns
        -------
        out : dict [str, array_like (n,)]
            Evaluated quantities in a dictionary.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield", "plasma")
        Neval = r.size

        # Allocate enough space for electrons and all ion species.
        m3 = unyt.m**3
        eV = unyt.eV
        nspecies  = self.input_getplasmaspecies()[0] + 1
        rawdens = (np.zeros((Neval*nspecies,), dtype="f8") + np.nan) / m3
        rawtemp = (np.zeros((Neval*nspecies,), dtype="f8") + np.nan) * eV

        fun = _LIBASCOT.libascot_plasma_eval_background
        fun.restype  = None
        fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        PTR_REAL, PTR_REAL]

        fun(ctypes.byref(self._sim), self._bfield_offload_array,
            self._plasma_offload_array,
            Neval, r, phi, z, t, rawdens, rawtemp)

        out = {}
        out["ne"] = rawdens[0:Neval]
        out["te"] = rawtemp[0:Neval]
        for i in range(1, nspecies):
            out["ni"+str(i)] = rawdens[(Neval)*i:(Neval)*(i+1)]
            out["ti"+str(i)] = rawtemp[(Neval)*i:(Neval)*(i+1)]

        return out

    def _eval_neutral(self, r, phi, z, t):
        """Evaluate neutral input quantities at given coordinates.

        Parameters
        ----------
        r : array_like, (n,)
            R coordinates where data is evaluated [m].
        phi : array_like (n,)
            phi coordinates where data is evaluated [rad].
        z : array_like (n,)
            z coordinates where data is evaluated [m].
        t : array_like (n,)
            Time coordinates where data is evaluated [s].

        Returns
        -------
        out : dict [str, array_like (n,)]
            Evaluated quantities in a dictionary.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield", "neutral")
        Neval = r.size
        out = {}
        m3 = unyt.m**3
        out["n0"] = (np.zeros(r.shape, dtype="f8") + np.nan) / m3

        fun = _LIBASCOT.libascot_neutral_eval_density
        fun.restype  = None
        fun.argtypes = [PTR_SIM, PTR_ARR,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        PTR_REAL]

        fun(ctypes.byref(self._sim), self._neutral_offload_array,
            Neval, r, phi, z, t, out["n0"])

        return out

    def _eval_boozer(self, r, phi, z, t, evalfun=False):
        """Evaluate boozer coordinate quantities at given coordinates.

        Parameters
        ----------
        r : array_like, (n,)
            R coordinates where data is evaluated [m].
        phi : array_like (n,)
            phi coordinates where data is evaluated [rad].
        z : array_like (n,)
            z coordinates where data is evaluated [m].
        t : array_like (n,)
            Time coordinates where data is evaluated [s].

        Returns
        -------
        out : dict [str, array_like (n,)]
            Evaluated quantities in a dictionary.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield", "boozer")
        Neval = r.size
        out = {}
        temp = np.zeros(r.shape, dtype="f8")
        nodim = unyt.dimensionless
        T = unyt.T; Wb = unyt.Wb; m = unyt.m; rad = unyt.rad
        if evalfun:
            out["qprof"]   = (np.copy(temp) + np.nan) * nodim
            out["bjac"]    = (np.copy(temp) + np.nan) / T**2
            out["bjacxb2"] = (np.copy(temp) + np.nan) * nodim

            fun = _LIBASCOT.libascot_boozer_eval_fun
            fun.restype  = ctypes.c_int
            fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR,
                            ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL]

            fun(ctypes.byref(self._sim), self._bfield_offload_array,
                self._boozer_offload_array,
                Neval, r, phi, z, t, out["qprof"], out["bjac"],
                out["bjacxb2"])
        else:
            out["psi (bzr)"]      = (np.copy(temp) + np.nan) * Wb
            out["theta"]          = (np.copy(temp) + np.nan) * rad
            out["zeta"]           = (np.copy(temp) + np.nan) * rad
            out["dpsidr (bzr)"]   = (np.copy(temp) + np.nan) * Wb/m
            out["dpsidphi (bzr)"] = (np.copy(temp) + np.nan) * Wb
            out["dpsidz (bzr)"]   = (np.copy(temp) + np.nan) * Wb/m
            out["dthetadr"]       = (np.copy(temp) + np.nan) / m
            out["dthetadphi"]     = (np.copy(temp) + np.nan) * nodim
            out["dthetadz"]       = (np.copy(temp) + np.nan) / m
            out["dzetadr"]        = (np.copy(temp) + np.nan) / m
            out["dzetadphi"]      = (np.copy(temp) + np.nan) * nodim
            out["dzetadz"]        = (np.copy(temp) + np.nan) / m
            out["rho (bzr)"]      = (np.copy(temp) + np.nan) * nodim

            fun = _LIBASCOT.libascot_boozer_eval_psithetazeta
            fun.restype  = ctypes.c_int
            fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR,
                            ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL]

            fun(ctypes.byref(self._sim), self._bfield_offload_array,
                self._boozer_offload_array,
                Neval, r, phi, z, t, out["psi (bzr)"], out["theta"],
                out["zeta"], out["dpsidr (bzr)"], out["dpsidphi (bzr)"],
                out["dpsidz (bzr)"], out["dthetadr"], out["dthetadphi"],
                out["dthetadz"], out["dzetadr"], out["dzetadphi"],
                out["dzetadz"], out["rho (bzr)"])

        return out

    def _eval_mhd(self, r, phi, z, t, mode=None, evalpot=False):
        """Evaluate MHD quantities at given coordinates.

        Parameters
        ----------
        r : array_like, (n,)
            R coordinates where data is evaluated [m].
        phi : array_like (n,)
            phi coordinates where data is evaluated [rad].
        z : array_like (n,)
            z coordinates where data is evaluated [m].
        t : array_like (n,)
            Time coordinates where data is evaluated [s].
        mode : int
            Pick a specific mode that is included by giving its index or
            None to include all modes.
        evalpot : bool, optional
            Evaluate eigenfunctions as well.

        Returns
        -------
        out : dict [str, array_like (n,)]
            Evaluated quantities in a dictionary.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield", "boozer", "mhd")
        if mode is None: mode = -1 # equal to MHD_INCLUDE_ALL
        Neval = r.size
        out = {}
        T = unyt.T; V = unyt.V; m = unyt.m; s = unyt.s
        temp = np.zeros(r.shape, dtype="f8")
        out["mhd_br"]   = (np.copy(temp) + np.nan) * T
        out["mhd_bphi"] = (np.copy(temp) + np.nan) * T
        out["mhd_bz"]   = (np.copy(temp) + np.nan) * T
        out["mhd_er"]   = (np.copy(temp) + np.nan) * V/m
        out["mhd_ephi"] = (np.copy(temp) + np.nan) * V/m
        out["mhd_ez"]   = (np.copy(temp) + np.nan) * V/m
        out["mhd_phi"]  = (np.copy(temp) + np.nan) * V

        fun = _LIBASCOT.libascot_mhd_eval_perturbation
        fun.restype  = None
        fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR, PTR_ARR,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        PTR_REAL, PTR_REAL, PTR_REAL]

        fun(ctypes.byref(self._sim), self._bfield_offload_array,
            self._boozer_offload_array, self._mhd_offload_array,
            Neval, r, phi, z, t, mode, out["mhd_br"], out["mhd_bphi"],
            out["mhd_bz"], out["mhd_er"], out["mhd_ephi"], out["mhd_ez"],
            out["mhd_phi"])

        if evalpot:
            out["alphaeig"] = (np.copy(temp) + np.nan) * T*m
            out["phieig"]   = (np.copy(temp) + np.nan) * V
            out["dadt"]     = (np.copy(temp) + np.nan) * T*m/s
            out["dadr"]     = (np.copy(temp) + np.nan) * T
            out["dadphi"]   = (np.copy(temp) + np.nan) * T*m
            out["dadz"]     = (np.copy(temp) + np.nan) * T
            out["dphidt"]   = (np.copy(temp) + np.nan) * V/s
            out["dphidr"]   = (np.copy(temp) + np.nan) * V/m
            out["dphidphi"] = (np.copy(temp) + np.nan) * V
            out["dphidz"]   = (np.copy(temp) + np.nan) * V/m

            fun = _LIBASCOT.libascot_mhd_eval
            fun.restype  = None
            fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR, PTR_ARR,
                            ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, ctypes.c_int, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                            PTR_REAL, PTR_REAL, PTR_REAL]

            fun(ctypes.byref(self._sim), self._bfield_offload_array,
                self._boozer_offload_array, self._mhd_offload_array,
                Neval, r, phi, z, t, mode, out["alphaeig"],
                out["dadr"], out["dadphi"], out["dadz"], out["dadt"],
                out["phieig"], out["dphidr"], out["dphidphi"], out["dphidz"],
                out["dphidt"])

        return out

    def input_mhd_modes(self):
        """Return mode numbers, amplitudes and frequencies.

        Returns
        -------
        nmodes : int
            Number of modes present.
        nmode : array_like (nmodes,)
            Toroidal mode number of each mode.
        mmode : array_like (nmodes,)
            Poloidal mode number of each mode.
        amplitude : array_like (nmodes,)
            Amplitude of each mode.
        omega : array_like (nmodes,)
            Frequency of each mode.
        phase : array_like (nmodes,)
            Phase of each mode.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("mhd")
        fun = _LIBASCOT.libascot_mhd_get_n_modes
        fun.restype  = ctypes.c_int
        fun.argtypes = [PTR_SIM, PTR_ARR]

        out = {}
        out["nmodes"] = fun(
            ctypes.byref(self._sim), self._mhd_offload_array)

        out["nmode"]     = np.zeros((out["nmodes"],), dtype="i4")
        out["mmode"]     = np.zeros((out["nmodes"],), dtype="i4")
        out["amplitude"] = np.zeros((out["nmodes"],), dtype="f8")
        out["omega"] = np.zeros((out["nmodes"],), dtype="f8")
        out["phase"]     = np.zeros((out["nmodes"],), dtype="f8")

        fun = _LIBASCOT.libascot_mhd_get_mode_specs
        fun.restype  = ctypes.c_int
        fun.argtypes = [PTR_SIM, PTR_ARR, PTR_INT, PTR_INT, PTR_REAL, PTR_REAL,
                        PTR_REAL]
        fun(ctypes.byref(self._sim), self._mhd_offload_array,
            out["nmode"], out["mmode"], out["amplitude"], out["omega"],
            out["phase"])

        return out["nmodes"], out["nmode"], out["mmode"], out["amplitude"],\
            out["omega"], out["phase"]

    @parseunits(ma="kg", qa="C", r="m", phi="rad", z="m", t="s", va="m/s")
    def input_eval_collcoefs(self, ma, qa, r, phi, z, t, va, *coefs, grid=True):
        """Evaluate Coulomb collision coefficients for a given test particle.

        Collision coefficients are evaluated by interpolating the plasma
        parameters on given coordinates. The coefficients are returned as
        a function of the test particle velocity.

        Parameters
        ----------
        ma : float
            Test particle mass.
        qa : float
            Test particle charge.
        r : array_like, (n,)
            R coordinates where data is evaluated.
        phi : array_like, (n,)
            phi coordinates where data is evaluated.
        z : array_like, (n,)
            z coordinates where data is evaluated.
        t : array_like, (n,)
            Time coordinates where data is evaluated.
        va : array_like, (nv,) or (n,)
            Test particle velocity.
        *coefs : str
            Names of the coefficients to be evaluated.

            "clog"   - Coulomb logarithm.
            "f"      - Drag in particle Fokker-Planck equation.
            "k"      - Drag in guiding-center Fokker-Planck equation.
            "nu"     - Pitch collision frequency.
            "dpara"  - Parallel diffusion.
            "dperp"  - Perpendicular diffusion.
            "ddpara" -  d(Dpara) / dv.
            "q"      - From K = Q + dDpara + 2*Dpara / va.
            "dq"     - d(Q) / dv.
            "mu0"    - One of the special functions needed in evaluation.
            "mu1"    - One of the special functions needed in evaluation.
            "dmu0"   - One of the special functions needed in evaluation.
        grid : bool, optional
            If True, all velocity components are evaluated at each (R,phi,z)
            position, i.e., the returned values are of shape (nion+1,n,nv).

            If False, the coefficients are evaluated at positions (R,phi,z,va),
            e.g. along an orbit, and the returned values have shape (nion+1,n).

        Returns
        -------
        *out : array_like, (nion+1, n, nv) or (nion+1, n)
            Evaluated collision coefficients in same order as declared in
            ``*coefs``.

            The first dimension is the background plasma species where the first
            index is for the electrons followed by ions.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield", "plasma")
        Neval = r.size
        Nv = va.size if grid else 1
        n_species = self.input_getplasmaspecies()[0]

        m = unyt.m; s = unyt.s
        out = {"clog":None, "f":None, "k":None, "nu":None, "dpara":None,
               "dperp":None, "ddpara":None, "q":None, "dq":None, "mu0":None,
               "mu1":None, "dmu0":None}
        temp = np.zeros((n_species, Neval, Nv), dtype="f8")
        if "clog"   in coefs: out["clog"]   = np.copy(temp)
        if "f"      in coefs: out["f"]      = np.copy(temp) * m/s**2
        if "k"      in coefs: out["k"]      = np.copy(temp) * m/s**2
        if "nu"     in coefs: out["nu"]     = np.copy(temp) / s
        if "dpara"  in coefs: out["dpara"]  = np.copy(temp) * m**2/s**3
        if "dperp"  in coefs: out["dperp"]  = np.copy(temp) * m**2/s**3
        if "ddpara" in coefs: out["ddpara"] = np.copy(temp) * m/s**2
        if "q"      in coefs: out["q"]      = np.copy(temp) * m/s**2
        if "dq"     in coefs: out["dq"]     = np.copy(temp) / s
        if "mu0"    in coefs: out["mu0"]    = np.copy(temp)
        if "mu1"    in coefs: out["mu1"]    = np.copy(temp)
        if "dmu0"   in coefs: out["dmu0"]   = np.copy(temp)

        fun = _LIBASCOT.libascot_eval_collcoefs
        fun.restype  = ctypes.c_int
        fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        ctypes.c_int, PTR_REAL,
                        ctypes.c_double, ctypes.c_double,
                        PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL]
        fun(ctypes.byref(self._sim), self._bfield_offload_array,
            self._plasma_offload_array, Neval, r, phi, z, t, Nv, va, ma, qa,
            out["f"], out["dpara"], out["dperp"], out["k"], out["nu"], out["q"],
            out["dq"], out["ddpara"], out["clog"], out["mu0"], out["mu1"],
            out["dmu0"])

        for d in list(out.keys()):
            if d not in coefs:
                del out[d]
            else:
                if not grid: out[d] = out[d].reshape((n_species, Neval))

        out = [out[k] for k in coefs]
        if len(out) == 1: out = out[0]
        return out

    @parseunits(ma="kg", r="m", phi="rad", z="m", t="s", va="m/s")
    def input_eval_atomicsigma(self, ma, anum, znum, r, phi, z, t, va, ion,
                               reaction):
        """Evaluate atomic reaction rate cross-sections for a given test
        particle.

        This function is a work in progress.

        Parameters
        ----------
        ma : float
            Test particle mass
        anum : int
            Test particle atomic mass number.
        znum : int
            Test particle charge number.
        r : array_like, (n,)
            R coordinates where data is evaluated [m].
        phi : array_like (n,)
            phi coordinates where data is evaluated [rad].
        z : array_like (n,)
            z coordinates where data is evaluated [m].
        t : array_like (n,)
            Time coordinates where data is evaluated [s].
        va : array_like (n,)
            Test particle velocities where data is evaluated in each grid point.
        ion : int
            Index number of the background ion species in plasma input.
        reaction : int
            Reaction.

        Returns
        -------
        sigmav : array_like, (n,nv)
            Reaction cross-section.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        """
        self._requireinit("bfield", "plasma", "neutral", "asigma")

        fun = _LIBASCOT.libascot_eval_sigmav
        fun.restype  = ctypes.c_int
        fun.argtypes = [PTR_SIM, PTR_ARR, PTR_ARR, PTR_ARR, PTR_ARR,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL, PTR_REAL,
                        ctypes.c_int, PTR_REAL, ctypes.c_int, ctypes.c_int,
                        ctypes.c_double, ctypes.c_int, ctypes.c_int, PTR_REAL]

        Neval = r.size
        Nv    = va.size

        out = {}
        out["sigmav"] = (np.zeros(r.shape, dtype="f8") + np.nan) / unyt.m**2
        fun(ctypes.byref(self._sim), self._bfield_offload_array,
            self._plasma_offload_array, self._neutral_offload_array,
            self._asigma_offload_array, Neval, r, phi, z, t, Nv, va,
            anum, znum, ma, ion, reaction, out["sigmav"])

        return out["sigmav"]


    def input_getplasmaspecies(self):
        """Get species present in plasma input (electrons first).

        Returns
        -------
        nspecies : int
            Number of species.
        mass : array_like
            Species' mass.
        charge : array_like
            Species' charge state.
        anum : array_like
            Species' atomic mass number.
        znum : array_like
            Species' charge number.

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("plasma")
        fun = _LIBASCOT.libascot_plasma_get_n_species
        fun.restype  = ctypes.c_int
        fun.argtypes = [PTR_SIM, PTR_ARR]

        out = {}
        out["nspecies"] = fun(
            ctypes.byref(self._sim), self._plasma_offload_array)

        out["mass"]   = np.zeros((out["nspecies"],), dtype="f8")
        out["charge"] = np.zeros((out["nspecies"],), dtype="f8")
        out["anum"]   = np.zeros((out["nspecies"],), dtype="i4")
        out["znum"]   = np.zeros((out["nspecies"],), dtype="i4")

        fun = _LIBASCOT.libascot_plasma_get_species_mass_and_charge
        fun.restype  = ctypes.c_int
        fun.argtypes = [PTR_SIM, PTR_ARR, PTR_REAL, PTR_REAL, PTR_INT,
                        PTR_INT]
        fun(ctypes.byref(self._sim), self._plasma_offload_array,
            out["mass"], out["charge"], out["anum"], out["znum"])

        return out["nspecies"], out["mass"], out["charge"], out["anum"],\
            out["znum"]

    @parseunits(rhovals="1", theta="rad", phi="rad", time="s")
    def input_rhotheta2rz(self, rho, theta, phi, time, maxiter=100, tol=1e-5):
        """Convert (rho, theta, phi) coordinates to (R,z) positions.

        Parameters
        ----------
        rhovals : array_like (n,)
            Normalized poloidal flux coordinates to be converted.
        theta : array_like (n,)
            Poloidal angle coordinates to be converted.
        phi : array_like (n,)
            Toroidal angle coordinates to be converted.
        time : float
            Time slice (same for all).

        Returns
        -------
        r : array_like (n,)
            R coordinates at (rhovals, theta, phi)
        z : array_like (n,)
            z coordinates at (rhovals, theta, phi)

        Raises
        ------
        AssertionError
            If required data has not been initialized.
        RuntimeError
            If evaluation in libascot.so failed.
        """
        self._requireinit("bfield")
        #rho = np.asarray(rho).ravel().astype(dtype="f8")
        Neval = rho.size
        r   = np.NaN * np.zeros((Neval,), dtype="f8") * unyt.m
        z   = np.NaN * np.zeros((Neval,), dtype="f8") * unyt.m

        fun = _LIBASCOT.libascot_B_field_rhotheta2rz
        fun.restype  = None
        fun.argtypes = [PTR_SIM, PTR_ARR,
                        ctypes.c_int, PTR_REAL, PTR_REAL, PTR_REAL,
                        ctypes.c_double, ctypes.c_int, ctypes.c_double,
                        PTR_REAL, PTR_REAL]
        fun(ctypes.byref(self._sim), self._bfield_offload_array,
            Neval, rho, theta, phi, time, maxiter, tol, r, z)

        return (r, z)

a5py/a5py/ascotpy/libascot.py
ascot4fusion/ascot5
"""Methods to inject input dependencies to the C-structures directly
from python.

This is in contrast to reading inputs from HDF5 files.
"""
import ctypes
from math import pi
from .libascot import _LIBASCOT
if _LIBASCOT:
    from a5py.ascotpy import ascot2py

class LibProviders():
    """Mixin class to provide dependency injectors/providers.
    """

    def provide_wall_2d(self,r,z):

        if( len(r) != len(z) ):
            raise ValueError("R and z must be of equal length.")

        nelements = len(r)

        # Create temporary variable for the wall
        R_c = (ctypes.c_double * nelements)()
        z_c = (ctypes.c_double * nelements)()

        for i in range(nelements):
            R_c[i] = r[i]
            z_c[i] = z[i]

        ascot2py.hdf5_wall_2d_to_offload(
            ctypes.byref(self._sim.wall_offload_data.w2d),
            ctypes.byref(self._wall_offload_array),
            nelements,
            R_c, z_c
            )

        self._sim.wall_offload_data.type = ascot2py.wall_type_2D

        ascot2py.wall_init_offload(
            ctypes.byref(self._sim.wall_offload_data),
            self._wall_offload_array,
            self._wall_int_offload_array
            )

    def provide_wall_3d(self,x1x2x3,y1y2y3,z1z2z3):

        nelements = int(x1x2x3.shape[0])

        # Create temporary variable for the wall
        x1x2x3_c = (ctypes.c_double * (3*nelements) )()
        y1y2y3_c = (ctypes.c_double * (3*nelements) )()
        z1z2z3_c = (ctypes.c_double * (3*nelements) )()

        # Let's hope the ordering is correct...
        x1x2x3_c[:] = x1x2x3.flatten()[:]
        y1y2y3_c[:] = y1y2y3.flatten()[:]
        z1z2z3_c[:] = z1z2z3.flatten()[:]


        ascot2py.hdf5_wall_3d_to_offload(
            ctypes.byref(self._sim.wall_offload_data.w3d),
            ctypes.byref(self._wall_offload_array),
            nelements,
            x1x2x3_c, y1y2y3_c, z1z2z3_c,
            )

        self._sim.wall_offload_data.type = ascot2py.wall_type_3D

        ascot2py.wall_init_offload(
            ctypes.byref(self._sim.wall_offload_data),
            self._wall_offload_array,
            self._wall_int_offload_array
            )

    def provide_BSTS(self,
                     b_rmin, b_rmax, b_nr, b_zmin, b_zmax, b_nz,
                     b_phimin, b_phimax, b_nphi, psi0, psi1,
                     br, bphi, bz, psi,
                     axis_phimin, axis_phimax, axis_nphi, axisr, axisz,
                     psi_rmin=None,   psi_rmax=None, psi_nr=None,
                     psi_zmin=None,   psi_zmax=None, psi_nz=None,
                     psi_phimin=None, psi_phimax=None, psi_nphi=None ):
        # bsts is the dictionary that comes from reading the hdf5
        bsts=locals()

        # 1. First fill in the meta-data struct
        #--------------------------------------
        # Mimic the C-function    hdf5_bfield_read_STS()

        #phimin/max deg2rad

        #B_STS_offload_data
        #sts = struct_c__SA_B_STS_offload_data()
        BSTS = self._sim.B_offload_data.BSTS

        BSTS.psigrid_n_r     = bsts['psi_nr'][0]
        BSTS.psigrid_n_z     = bsts['psi_nz'][0]
        BSTS.psigrid_n_phi   = bsts['psi_nphi'][0]
        BSTS.psigrid_r_min   = bsts['psi_rmin'][0]
        BSTS.psigrid_r_max   = bsts['psi_rmax'][0]
        BSTS.psigrid_z_min   = bsts['psi_zmin'][0]
        BSTS.psigrid_z_max   = bsts['psi_zmax'][0]
        BSTS.psigrid_phi_min = bsts['psi_phimin'][0] * np.pi * 2.0 / 360.0
        BSTS.psigrid_phi_max = bsts['psi_phimax'][0] * np.pi * 2.0 / 360.0

        BSTS.Bgrid_n_r       = bsts['b_nr'][0]
        BSTS.Bgrid_n_z       = bsts['b_nz'][0]
        BSTS.Bgrid_n_phi     = bsts['b_nphi'][0]
        BSTS.Bgrid_r_min     = bsts['b_rmin'][0]
        BSTS.Bgrid_r_max     = bsts['b_rmax'][0]
        BSTS.Bgrid_z_min     = bsts['b_zmin'][0]
        BSTS.Bgrid_z_max     = bsts['b_zmax'][0]
        BSTS.Bgrid_phi_min   = bsts['b_phimin'][0] * np.pi * 2.0 / 360.0
        BSTS.Bgrid_phi_max   = bsts['b_phimax'][0] * np.pi * 2.0 / 360.0

        BSTS.psi0            = bsts['psi0'][0]
        BSTS.psi1            = bsts['psi1'][0]


        BSTS.n_axis          = bsts['axis_nphi'][0]
        BSTS.axis_min        = bsts['axis_phimin'][0] * np.pi * 2.0 / 360.0
        BSTS.axis_max        = bsts['axis_phimax'][0] * np.pi * 2.0 / 360.0
        # BSTS.axis_grid       = bsts['']   # Not really used

        # 2. Get the right sized offload array
        #--------------------------------------
        # Does this deed to be allocated by C code or can we do it with a python array?

        '''
        /* Allocate offload_array storing psi and the three components of B */
        int psi_size = offload_data->psigrid_n_r*offload_data->psigrid_n_z
           * offload_data->psigrid_n_phi;
        int B_size = offload_data->Bgrid_n_r * offload_data->Bgrid_n_z
           * offload_data->Bgrid_n_phi;
        int axis_size = offload_data->n_axis;

        *offload_array = (real*) malloc((psi_size + 3 * B_size + 2 * axis_size)
                                    * sizeof(real));
        offload_data->offload_array_length = psi_size + 3 * B_size + 2 * axis_size;
        '''

        offload_size = 0

        npsi = bsts['psi_nr'][0] * bsts['psi_nz'][0] * bsts[ 'psi_nphi'][0]
        nB   = bsts[  'b_nr'][0] * bsts[  'b_nz'][0] * bsts[   'b_nphi'][0]
        naxis=                                         bsts['axis_nphi'][0]

        # psi_size
        offload_size +=  1 * npsi
        # B_size
        offload_size +=  3 * nB
        # axis_size
        offload_size +=  2 * naxis

        BSTS.offload_array_length = offload_size

        # Python side allocation
        #-----------------------
        # B_offload_array = (ctypes.c_double * (offload_size) )()
        #self._B_offload_array.contents = B_offload_array

        # C side allocation
        #------------------
        self._bfield_offload_array = \
            ascot2py.libascot_allocate_reals(offload_size)

        # Cast the pointer into an array
        B_offload_array = ctypes.cast(
            self._bfield_offload_array,
            ctypes.POINTER(ctypes.c_double*offload_size) )[0]

        # 3. copy the large arrays to the offload array
        #----------------------------------------------

        offset = 0
        order = 'F'

        B_offload_array[(offset):(offset+nB  )] = \
            bsts['br'].flatten(order=order)
        offset += nB

        B_offload_array[(offset):(offset+nB  )] = \
            bsts['bphi'].flatten(order=order)
        offset += nB

        B_offload_array[(offset):(offset+nB  )] = \
            bsts['bz'].flatten(order=order)
        offset += nB

        B_offload_array[(offset):(offset+npsi)] = \
            bsts['psi'].flatten(order=order)
        offset += npsi

        B_offload_array[(offset):(offset+naxis)] = bsts['axisr']
        offset += naxis

        B_offload_array[(offset):(offset+naxis)] = bsts['axisz']
        offset += naxis

        # 4. Set the correct data type
        #------------------------------
        self._sim.B_offload_data.type = ascot2py.B_field_type_STS

        # 5. Do the init offload
        #-----------------------

        #B_field_init_offload.argtypes = \
        #    [ctypes.POINTER(struct_c__SA_B_field_offload_data),
        #     ctypes.POINTER(ctypes.POINTER(ctypes.c_double))]

        ascot2py.B_field_init_offload(
            ctypes.byref(self._sim.B_offload_data),
            ctypes.byref(self._bfield_offload_array) )

a5py/a5py/ascotpy/libproviders.py
ascot4fusion/ascot5
"""Methods to init, run, and process interactive simulations.

Interactive simulations are simulations that are run via libascot.so.
They are equivalent to normal ascot5_main simulations, except that the
output data is not written to the HDF5 file but the C structs are accessed
directly from Python.
"""
import unyt
import ctypes
import numpy as np

from numpy.ctypeslib import ndpointer

from a5py.physlib import momentum_velocity
from a5py.routines.virtualrun import VirtualRun
from a5py.exceptions import *

from .libascot import _LIBASCOT
if _LIBASCOT:
    from . import ascot2py

class LibSimulate():
    """Mixin class that introduces methods for active simulations.
    """

    def simulation_initoptions(self, **opt):
        """Set simulation options for the interactive simulation.

        There is no need to free options before setting new ones.

        Parameters
        ----------
        **opt :
            Options in same format as accepted by :meth:`Opt.write_hdf5`.
        """
        self._virtualoptions = opt
        if not _LIBASCOT:
            raise AscotInitException(
                "Python interface disabled as libascot.so is not found")
        # Simulation mode options
        self._sim.sim_mode    = int(opt["SIM_MODE"]);
        self._sim.enable_ada  = int(opt["ENABLE_ADAPTIVE"])
        self._sim.record_mode = int(opt["RECORD_MODE"])

        # Time step
        self._sim.fix_usrdef_use    = int(opt["FIXEDSTEP_USE_USERDEFINED"])
        self._sim.fix_usrdef_val    = opt["FIXEDSTEP_USERDEFINED"]
        self._sim.fix_gyrodef_nstep = int(opt["FIXEDSTEP_GYRODEFINED"])
        self._sim.ada_tol_orbfol    = opt["ADAPTIVE_TOL_ORBIT"]
        self._sim.ada_tol_clmbcol   = opt["ADAPTIVE_TOL_CCOL"]
        self._sim.ada_max_drho      = opt["ADAPTIVE_MAX_DRHO"]
        self._sim.ada_max_dphi      = opt["ADAPTIVE_MAX_DPHI"]

        # Physics
        self._sim.enable_orbfol       = int(opt["ENABLE_ORBIT_FOLLOWING"])
        self._sim.enable_clmbcol      = int(opt["ENABLE_COULOMB_COLLISIONS"])
        self._sim.enable_mhd          = int(opt["ENABLE_MHD"])
        self._sim.enable_atomic       = int(opt["ENABLE_ATOMIC"])
        self._sim.disable_gctransform = int(opt["DISABLE_FIRSTORDER_GCTRANS"])
        self._sim.disable_energyccoll = int(opt["DISABLE_ENERGY_CCOLL"])
        self._sim.disable_pitchccoll  = int(opt["DISABLE_PITCH_CCOLL"])
        self._sim.disable_gcdiffccoll = int(opt["DISABLE_GCDIFF_CCOLL"])
        self._sim.reverse_time        = int(opt["REVERSE_TIME"])

        # Which end conditions are active
        self._sim.endcond_active = 0;
        if opt["ENDCOND_SIMTIMELIM"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_tlim
        if opt["ENDCOND_CPUTIMELIM"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_cpumax
        if opt["ENDCOND_RHOLIM"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_rhomin
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_rhomax
        if opt["ENDCOND_ENERGYLIM"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_emin
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_therm
        if opt["ENDCOND_WALLHIT"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_wall
        if opt["ENDCOND_MAXORBS"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_polmax
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_tormax
        self._sim.endcond_torandpol  = 0 + 1 * int(opt["ENDCOND_MAXORBS"] == 2)
        if opt["ENDCOND_NEUTRALIZED"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_neutrz
        if opt["ENDCOND_IONIZED"]:
            self._sim.endcond_active = \
                self._sim.endcond_active | ascot2py.endcond_ioniz

        # End condition parameters
        eV2J = unyt.e.base_value
        self._sim.endcond_lim_simtime = opt["ENDCOND_LIM_SIMTIME"]
        self._sim.endcond_max_mileage = opt["ENDCOND_MAX_MILEAGE"]
        self._sim.endcond_max_cputime = opt["ENDCOND_MAX_CPUTIME"]
        self._sim.endcond_max_rho     = opt["ENDCOND_MAX_RHO"]
        self._sim.endcond_min_rho     = opt["ENDCOND_MIN_RHO"]
        self._sim.endcond_min_thermal = opt["ENDCOND_MIN_THERMAL"]
        self._sim.endcond_min_ekin    = opt["ENDCOND_MIN_ENERGY"]*eV2J
        self._sim.endcond_max_polorb  = 2*np.pi*opt["ENDCOND_MAX_POLOIDALORBS"]
        self._sim.endcond_max_tororb  = 2*np.pi*opt["ENDCOND_MAX_TOROIDALORBS"]

        diag = self._sim.diag_offload_data
        diag.dist5D_collect    = int(opt["ENABLE_DIST_5D"]) * 0    # Not impl.
        diag.dist6D_collect    = int(opt["ENABLE_DIST_6D"]) * 0    # Not impl.
        diag.distrho5D_collect = int(opt["ENABLE_DIST_RHO5D"]) * 0 # Not impl.
        diag.distrho6D_collect = int(opt["ENABLE_DIST_RHO6D"]) * 0 # Not impl.
        diag.dist5D_collect    = int(opt["ENABLE_DIST_COM"]) * 0   # Not impl.
        diag.diagtrcof_collect = int(opt["ENABLE_TRANSCOEF"]) * 0  # Not impl.
        diag.diagorb_collect   = int(opt["ENABLE_ORBITWRITE"])

        diagorb = diag.diagorb
        diagorb.mode          = int(opt["ORBITWRITE_MODE"])
        diagorb.Npnt          = int(opt["ORBITWRITE_NPOINT"])
        diagorb.writeInterval = opt["ORBITWRITE_INTERVAL"]

        diagorb.record_mode = self._sim.sim_mode
        if self._sim.record_mode and \
           (self._sim.sim_mode == ascot2py.simulate_mode_fo or
            self._sim.sim_mode == ascot2py.simulate_mode_hybrid):
            diagorb.record_mode = ascot2py.simulate_mode_gc

        torangs = opt["ORBITWRITE_TOROIDALANGLES"]
        torangs = torangs if isinstance(torangs, list) else [torangs]
        polangs = opt["ORBITWRITE_POLOIDALANGLES"]
        polangs = polangs if isinstance(polangs, list) else [polangs]
        radials = opt["ORBITWRITE_RADIALDISTANCES"]
        radials = radials if isinstance(radials, list) else [radials]

        diagorb.ntoroidalplots = len(torangs)
        for i in range(diagorb.ntoroidalplots):
            diagorb.toroidalangles[i] = torangs[i] * np.pi / 180

        diagorb.npoloidalplots = len(polangs)
        for i in range(diagorb.npoloidalplots):
            diagorb.poloidalangles[i] = polangs[i] * np.pi / 180

        diagorb.nradialplots = len(radials)
        for i in range(diagorb.nradialplots):
            diagorb.radialdistances[i] = radials[i]

    def simulation_initinputs(self, bfield=True, efield=True, plasma=True,
                              neutral=True, wall=True, boozer=True, mhd=True,
                              asigma=True, switch=True):
        """Prepare input fields for the interactive simulation.

        Initializes simulation inputs. The inputs used in the simulation are
        those that are active.

        This method differs from :meth:`input_init` in that here the inputs are
        "packed" internally in a single (offload) array as this is what ASCOT5
        does. Inputs cannot be changed while the they are packed. The unpacking
        is done with :meth:`simulation_free`.

        This method must be called before running the simulation.
        """
        self.input_init(
            bfield=bfield, efield=efield, plasma=plasma, neutral=neutral,
            wall=wall, boozer=boozer, mhd=mhd, asigma=asigma, switch=switch)
        self._pack()


    def simulation_initmarkers(self, **mrk):
        """Create markers for the interactive simulations.

        Any existing markers are deallocated when new ones are created.

        Parameters
        ----------
        **mrk :
            Marker input (all marker types are supported) as accepted by
            :meth:`Prt.write_hdf5`, :meth:`GC.write_hdf5`, or
            :meth:`FL.write_hdf5`.
        """
        if not _LIBASCOT:
            raise AscotInitException(
                "Python interface disabled as libascot.so is not found")
        if self._nmrk.value > 0:
            ascot2py.libascot_deallocate(self._inistate)
            self._nmrk.value = 0
            self._virtualmarkers = None
        self._virtualmarkers = mrk
        nmrk = mrk["n"]
        pin = ascot2py.libascot_allocate_input_particles(nmrk)
        prttypes = ascot2py.input_particle_type__enumvalues

        # particle
        if "vr" in mrk:
            for i in range(nmrk):
                pin[i].type = ascot2py.input_particle_type_p
                p = pin[i].c__SA_input_particle_0.p

                vvec = np.array([mrk["vr"][i], mrk["vphi"][i], mrk["vz"][i]]) * unyt.m/unyt.s
                pvec = momentum_velocity(mrk["mass"][i], vvec)

                p.r       = mrk["r"][i]
                p.phi     = mrk["phi"][i] * np.pi / 180
                p.z       = mrk["z"][i]
                p.p_r     = pvec[0]
                p.p_phi   = pvec[1]
                p.p_z     = pvec[2]
                p.mass    = mrk["mass"][i]
                p.charge  = mrk["charge"][i]
                print(mrk)
                print(mrk["anum"])
                print(mrk["anum"][i])
                print(mrk["anum"][i][0])
                p.anum    = mrk["anum"][i][0]
                p.znum    = mrk["znum"][i]
                p.weight  = mrk["weight"][i]
                p.time    = mrk["time"][i]
                p.mileage = mrk["mileage"][i]
                p.id      = mrk["ids"][i]

        # particle gc
        elif "energy" in mrk:
            for i in range(nmrk):
                pin[i].type = ascot2py.input_particle_type_gc
                p = pin[i].c__SA_input_particle_0.p_gc

                p.r       = mrk["r"][i]
                p.phi     = mrk["phi"][i] * np.pi / 180
                p.z       = mrk["z"][i]
                p.energy  = mrk["energy"][i] * unyt.elementary_charge.value
                p.pitch   = mrk["pitch"][i]
                p.zeta    = mrk["zeta"][i]
                p.mass    = mrk["mass"][i] * unyt.atomic_mass_unit.value
                p.charge  = mrk["charge"][i] * unyt.elementary_charge.value
                p.anum    = mrk["anum"][i]
                p.znum    = mrk["znum"][i]
                p.weight  = mrk["weight"][i]
                p.time    = mrk["time"][i]
                p.id      = mrk["ids"][i]

        # particle fl
        else:
            for i in range(nmrk):
                pin[i].type = ascot2py.input_particle_type_ml
                p = pin[i].c__SA_input_particle_0.p_ml

                p.r       = mrk["r"][i]
                p.phi     = mrk["phi"][i] * np.pi / 180
                p.z       = mrk["z"][i]
                p.pitch   = mrk["pitch"][i]
                p.weight  = mrk["weight"][i]
                p.time    = mrk["time"][i]
                p.id      = mrk["ids"][i]

        ascot2py.prepare_markers(
            ctypes.byref(self._sim), self._mpi_size, self._mpi_rank, nmrk,
            ctypes.byref(pin), ctypes.byref(self._inistate),
            ctypes.byref(self._nmrk), self._bfield_offload_array)

    def simulation_run(self, printsummary=True):
        """Run the interactive simulation using inputs, options and markers that
        were set.

        Parameters
        ----------
        printsummary : bool, optional
            If True, summary of marker endstates is printed after the simulation
            completes.

        Returns
        -------
        run : :class:`VirtualRun`
            An object that acts almost exactly as :class:`RunGroup` except that
            the data is read from C arrays in the memory instead of HDF5 file.

            The run can be used until the output is freed with
            :meth:`simulation_free`. Previous data must be freed before
            rerunning the simulation.

        Raises
        ------
        AscotInitException
            If inputs are not packed, markers are not initialized or previous
            results have not been freed.
        """
        if not _LIBASCOT:
            raise AscotInitException(
                "Python interface disabled as libascot.so is not found")
        if not self._offload_ready:
            raise AscotInitException(
                "Initialize inputs before running the simulation")
        if self._nmrk.value == 0:
            raise AscotInitException(
                "Initialize markers before running the simulation")
        if self._diag_occupied:
            raise AscotInitException(
                "Free previous results before running the simulation")

        # Copy inistate to "endstate" and give endstate to ASCOT5 as an input.
        # Otherwise the inistate would get overwritten as the code updates
        # the marker positions in the input array.
        self._endstate = ascot2py.libascot_allocate_particle_states(self._nmrk)
        for j in range(self._nmrk.value):
            for i in range(len(ascot2py.particle_state._fields_)):
                name = self._inistate[j]._fields_[i][0]
                val  = getattr(self._inistate[j], name)
                setattr(self._endstate[j], name, val)

        ascot2py.offload_and_simulate(
            ctypes.byref(self._sim), self._mpi_size, self._mpi_rank,
            self._mpi_root, self._nmrk, self._nmrk, self._endstate,
            ctypes.byref(self._offload_data), self._offload_array,
            self._int_offload_array, ctypes.byref(self._endstate),
            ctypes.byref(self._diag_offload_array))

        self._diag_occupied = True
        if printsummary:
            ascot2py.print_marker_summary(self._endstate, self._nmrk)

        class VirtualInput():
            """Wrapper for marker and options inputs.
            """

            def __init__(self, inp):
                self.inp = inp

            def read(self):
                return self.inp

        orbits = None
        if self._sim.diag_offload_data.diagorb_collect:
            orbits = self._diag_offload_array
        return VirtualRun(self, self._nmrk.value,
                          self._inistate, self._endstate,
                          self._sim.diag_offload_data.diagorb.Npnt,
                          self._sim.diag_offload_data.diagorb.record_mode,
                          self._sim.diag_offload_data.diagorb.mode,
                          VirtualInput(self._virtualoptions),
                          VirtualInput(self._virtualmarkers),
                          orbits)

    def simulation_free(self, inputs=False, markers=False, diagnostics=False):
        """Free resources used by the interactive simulation.

        If called without arguments, everything will be freed.

        Parameters
        ----------
        inputs : bool, optional
            If True, inputs are unpacked and freed.
        markers : bool, optional
            If True, markers are freed.
        diagnostics : bool, optional
            If True, diagnostics data is freed.
        """
        if not _LIBASCOT:
            raise AscotInitException(
                "Python interface disabled as libascot.so is not found")
        if not inputs and not markers and not diagnostics:
            inputs      = True
            markers     = True
            diagnostics = True

        if inputs and self._offload_ready:
            self._unpack()
        if markers and self._nmrk.value > 0:
            self._nmrk.value = 0
            ascot2py.libascot_deallocate(self._inistate)
            self._virtualmarkers = None
        if diagnostics and self._diag_occupied:
            self._diag_occupied = False
            ascot2py.libascot_deallocate(self._diag_offload_array)
            ascot2py.libascot_deallocate(self._endstate)

a5py/a5py/ascotpy/libsimulate.py
ascot4fusion/ascot5
"""
Checks to help user to ensure the inputs are ok before running the simulation.
"""

import numpy as np
import matplotlib.pyplot as plt
from a5py.ascotpy import Ascotpy
from a5py.routines.plotting import openfigureifnoaxes

def check_inputs_present(ascotpy):
    """
    Check required inputs are present for this run.
    """

    # Determine what we are simulating.
    runtype = "ascot"

    msg = []
    if runtype == "ascot":
        if not hasattr(ascotpy.hdf5, "bfield"):
            msg += ["Error: bfield data is missing"]
        if not hasattr(ascotpy.hdf5, "efield"):
            msg += ["Error: efield data is missing"]
        if not hasattr(ascotpy.hdf5, "plasma"):
            msg += ["Error: plasma data is missing"]
        if not hasattr(ascotpy.hdf5, "marker"):
            msg += ["Error: marker data is missing"]
        if not hasattr(ascotpy.hdf5, "options"):
            msg += ["Error: options are missing"]
        if not hasattr(ascotpy.hdf5, "boozer"):
            msg += ["Error: boozer data is missing"]
        if not hasattr(ascotpy.hdf5, "wall"):
            msg += ["Error: wall data is missing"]
        if not hasattr(ascotpy.hdf5, "mhd"):
            msg += ["Error: mhd data is missing"]


    return msg

def check_options_consistent(ascotpy):
    """
    Check that options are consistent with inputs and each other.
    """
    msg = []
    opt = ascotpy.hdf5.options.active

    msg0 = opt.validatevalues()
    if len(msg0) > 0:
        msg += ["Error: following options had invalid parameters:"]
        msg += msg0

    opt = opt.read()
    if opt["ENABLE_MHD"] == 1 and \
       ascotpy.hdf5.bfield.active.get_type() == "B_STS":
        msg += ["Error: cannot enable MHD for stellarators"]

    if opt["SIM_MODE"] in [1,2,3] and \
       ascotpy.hdf5.marker.active.get_type() == "mrk_fl":
        msg += ["Error: SIM_MODE invalid for field-line markers"]

    # Make rough estimates on how much memory is consumed by diagnostics
    high_memory_consumption = 1e9 # In bits
    # Orbit = Npoint * Nfields (~10) * Nmrk * 8 bit
    orb_mem = opt["ORBITWRITE_NPOINT"] * 10 \
              * ascotpy.hdf5.marker.active.read()["n"] * 8
    # Distributions = NR * Nz * ... * 8 bit
    rzp = opt["DIST_NBIN_R"] * opt["DIST_NBIN_Z"] * opt["DIST_NBIN_PHI"]
    rtp = opt["DIST_NBIN_RHO"] * opt["DIST_NBIN_THETA"] * opt["DIST_NBIN_PHI"]
    p2d = opt["DIST_NBIN_PPA"] * opt["DIST_NBIN_PPE"]
    p3d = opt["DIST_NBIN_PR"] * opt["DIST_NBIN_PZ"] * opt["DIST_NBIN_PPHI"]

    if opt["ENABLE_ORBITWRITE"] == 1 and orb_mem > high_memory_consumption:
        msg += ["Warning: orbit diagnostic memory consumption high (~" +
                str(int(orb_mem / 1e9)) + "Gb)"]

    if opt["ENABLE_DIST_5D"] == 1 and rzp * p2d * 8 > high_memory_consumption:
        msg += ["Warning: 5D distribution memory consumption high (~" +
                str(int(rzp * p2d * 8 / 1e9)) + "Gb)"]

    if opt["ENABLE_DIST_6D"] == 1 and rzp * p3d * 8 > high_memory_consumption:
        msg += ["Warning: 6D distribution memory consumption high (~" +
                str(int(rzp * p3d * 8 / 1e9)) + "Gb)"]

    if opt["ENABLE_DIST_RHO5D"] == 1 and rtp * p2d * 8 > high_memory_consumption:
        msg += ["Warning: rho5D distribution memory consumption high (~" +
                str(int(rtp * p2d * 8 / 1e9)) + "Gb)"]

    if opt["ENABLE_DIST_RHO6D"] == 1 and rtp * p3d * 8 > high_memory_consumption:
        msg += ["Warning: rho6D distribution memory consumption high (~" +
                str(int(rtp * p3d * 8 / 1e9)) + "Gb)"]

    return msg


def check_bfield_psi0(ascotpy):
    """
    Checks whether psi0 given in input is actually extreme value.

    Because psi is interpolated with splines, there might be numerical error
    that causes psi (near the axis) to have more extreme value than psi0
    which is given in input. This leads to imaginary rho and termination
    of the simulation if marker ends up there.

    This check uses Monte Carlo method to 1. Draw phi 2. Evaluate axis (R,z)
    3. Draw random (R,z) coordinates within 10 cm of the axis. 4. Evaluate psi
    at that point and compare to psi0. Process repeats N times. Check passes
    if all evaluations are valid.
    """

    data = ascotpy.hdf5.bfield.active.read()
    if "psi0" not in data:
        return []

    psi0 = data["psi0"]
    psi1 = data["psi1"]
    psi0 = -7.1

    N     = 10000
    phi   = np.random.rand(N,) * 360
    theta = np.random.rand(N,) * 2 * np.pi

    axis = ascotpy.evaluate(1, phi, 0, 0, "axis")
    z0 = axis["axisz"]
    r0 = axis["axisr"]

    R = 0.1 # 10 cm
    r = R * np.cos(theta) + r0
    z = R * np.sin(theta) + z0
    psi = ascotpy.evaluate(r, phi, z, 0, "psi")

    if psi0 < psi1 and any(psi < psi0):
        return ["Error: psi0 = %.2e but we found near axis that psi = %.2e" \
                % (psi0, np.amin(psi)) ]
    elif psi0 > psi1 and any(psi > psi0):
        return ["Error: psi0 = %.2e but we found near axis that psi = %.2e" \
                % (psi0, np.amax(psi)) ]

    return []


def plot_top_view(ascotpy, axes=None):
    """
    Plot top view of the machine showing Ip, Bphi, and markers.

    Assumes bfield is initialized in ascotpy.
    """
    showfig = False
    if axes == None:
        fig  = plt.figure()
        axes = fig.add_subplot(1,1,1)
        showfig = True

    axis = ascotpy.evaluate(1, 0, 0, 0, "axis")
    z0 = axis["axisz"]
    r0 = axis["axisr"]

    rmin = r0-r0/2
    rmax = r0+r0/2
    dphi = 10 * np.pi/180 # So that b and j quivers dont overlap

    r   = np.linspace(rmin, rmax, 10)
    phi = np.linspace(0, 360, 18, endpoint=False) * np.pi/180
    phi, r = np.meshgrid(phi, r)
    r   = r.ravel()
    phi = phi.ravel()

    br   = np.squeeze(ascotpy.evaluate(r, phi, z0, 0, "br"))
    bphi = np.squeeze(ascotpy.evaluate(r, phi, z0, 0, "bphi"))
    jr   = np.squeeze(ascotpy.evaluate(r, phi + dphi, z0, 0, "jr"))
    jphi = np.squeeze(ascotpy.evaluate(r, phi + dphi, z0, 0, "jphi"))

    x  = np.cos(phi) * r
    y  = np.sin(phi) * r
    bx = np.cos(phi) * br - np.sin(phi) * bphi
    by = np.sin(phi) * br + np.cos(phi) * bphi
    bnorm = np.sqrt(bx**2 + by**2)

    xj = np.cos(phi+dphi) * r
    yj = np.sin(phi+dphi) * r
    jx = np.cos(phi+dphi) * jr - np.sin(phi+dphi) * jphi
    jy = np.sin(phi+dphi) * jr + np.cos(phi+dphi) * jphi
    jnorm = np.sqrt(jx**2 + jy**2)

    axes.quiver(x,y, bx/bnorm, by/bnorm,
                color="blue", scale=40)
    axes.quiver(xj,yj, jx/jnorm, jy/jnorm,
                color="red", scale=40)

    axes.set_aspect("equal", adjustable="box")

    marker = ascotpy.hdf5.marker.active.read()
    x  = np.cos(marker["phi"] * np.pi/180) * marker["r"]
    y  = np.sin(marker["phi"] * np.pi/180) * marker["r"]

    axes.scatter(x,y, s=1, c="black", zorder=-2)

    axes.legend(("Magnetic field", "Plasma current", "Markers"))

    if showfig:
        plt.show()


@openfigureifnoaxes(projection=None)
def plot_energypitch(ascotpy, axes=None):
    """
    Plot marker energy-pitch histogram
    """
    ascotpy.hdf5.marker.active.plot_hist_energypitch(ascotpy, axes=axes)


@openfigureifnoaxes(projection=None)
def plot_rhophi(ascotpy, axes=None):
    """
    Plot marker rho-phi histogram
    """
    ascotpy.hdf5.marker.active.plot_hist_rhophi(ascotpy, axes=axes)


if __name__ == '__main__':
    ascotpy = Ascotpy("ascot.h5")
    ascotpy.init(bfield=True)

    msg = []
    msg += check_inputs_present(ascotpy)
    msg += check_options_consistent(ascotpy)
    msg += check_bfield_psi0(ascotpy)

    for s in msg:
        print(s)

    if len(msg) == 0:
        print("Preflight checks ok!")

    fig  = plt.figure()
    plot_top_view(ascotpy,    axes=fig.subplot(2,2,[1,3]))
    plot_rhophi(ascotpy,      axes=fig.subplot(2,2,2))
    plot_energypitch(ascotpy, axes=fig.subplot(2,2,4))
    plt.show()

    ascotpy.free(bfield=True)

a5py/a5py/ascotpy/preflight.py
ascot4fusion/ascot5
import numpy as np
from scipy import interpolate
import ctypes
from numpy.ctypeslib import ndpointer
from copy import deepcopy
import a5py.ascot5io.B_3DS as B_3D

libascot = ctypes.CDLL("libascot.so")
real_p = ndpointer(ctypes.c_double, flags="C_CONTIGUOUS")

libascot.biosaw_calc_B.restype = None
libascot.biosaw_calc_B.argtypes = [ctypes.c_int, real_p, real_p, real_p,
                                   ctypes.c_int, real_p, real_p, real_p,
                                   real_p, real_p, real_p]

def calc_B(Rv, phiv, zv, coil_x, coil_y, coil_z):
    coil_n = len(coil_x)
    coil_x = np.array(coil_x, dtype="f8")
    coil_y = np.array(coil_y, dtype="f8")
    coil_z = np.array(coil_z, dtype="f8")

    [R,phi,z] = np.meshgrid(Rv,phiv,zv)

    x = R * np.cos(phi)
    y = R * np.sin(phi)

    b3d = dict()

    b3d["b_rmin"] = Rv[0]
    b3d["b_rmax"] = Rv[-1]
    b3d["b_nr"] = Rv.size
    b3d["b_phimin"] = phiv[0]
    b3d["b_phimax"] = phiv[-1]
    b3d["b_nphi"] = phiv.size
    b3d["b_zmin"] = zv[0]
    b3d["b_zmax"] = zv[-1]
    b3d["b_nz"] = zv.size

    b3d["br"] = np.zeros(R.shape)
    b3d["bphi"] = np.zeros(R.shape)
    b3d["bz"] = np.zeros(R.shape)
    
    n = len(R.ravel())

    Bx = np.zeros(n, dtype="f8")
    By = np.zeros(n, dtype="f8")
    Bz = np.zeros(n, dtype="f8")

    libascot.biosaw_calc_B(n, x.ravel(), y.ravel(), z.ravel(), coil_n, coil_x, coil_y, coil_z, Bx, By, Bz)

    Bx = Bx.reshape((phiv.size,Rv.size,zv.size))
    By = By.reshape((phiv.size,Rv.size,zv.size))
    Bz = Bz.reshape((phiv.size,Rv.size,zv.size))

    b3d["br"] = Bx*np.cos(phi) + By*np.sin(phi)
    b3d["bphi"] = -Bx*np.sin(phi) + By*np.cos(phi)
    b3d["bz"] = Bz

    return b3d

def revolve_coils(b3d_coil, n_coil, B0=None, R0=None, width=None):
    b3d = deepcopy(b3d_coil)

    b3d["br"] = 0*b3d["br"]
    b3d["bphi"] = 0*b3d["bphi"]
    b3d["bz"] = 0*b3d["bz"]

    n_phi = b3d["b_nphi"]
    n_phi_coil = int(n_phi/n_coil)

    if width is None:
        width = 1

    for i in range(n_phi):
        for j in range(n_coil):
            for k in range(width):
                index = (k + i + j*n_phi_coil) % n_phi
                b3d["br"][index,:,:] += b3d_coil["br"][i,:,:]
                b3d["bphi"][index,:,:] += b3d_coil["bphi"][i,:,:]
                b3d["bz"][index,:,:] += b3d_coil["bz"][i,:,:]

    if B0 is not None and R0 is not None:
        Rv = np.linspace(b3d["b_rmin"],b3d["b_rmax"],b3d["b_nr"])
        zv = np.linspace(b3d["b_zmin"],b3d["b_zmax"],b3d["b_nz"])
        f=interpolate.interp2d(Rv, zv, np.mean(b3d["bphi"],0).T, kind='cubic')

        b3d["br"] = b3d["br"] * B0/f(R0,0)
        b3d["bphi"] = b3d["bphi"] * B0/f(R0,0)
        b3d["bz"] = b3d["bz"] * B0/f(R0,0)

    return b3d

def combine(fn, b2d, b3d):
    B_3D.write_hdf5(fn,
                    b3d["b_rmin"], b3d["b_rmax"], b3d["b_nr"],
                    b3d["b_zmin"], b3d["b_zmax"], b3d["b_nz"],
                    b3d["b_phimin"], b3d["b_phimax"], b3d["b_nphi"],
                    b2d["axisr"], b2d["axisz"],
                    b2d["psi"], b2d["psi0"], b2d["psi1"],
                    np.transpose(b3d["br"],(1,0,2)),
                    np.transpose(b3d["bphi"],(1,0,2)),
                    np.transpose(b3d["bz"],(1,0,2)),
                    psi_rmin=b2d["rmin"], psi_rmax=b2d["rmax"], psi_nr=b2d["nr"],
                    psi_zmin=b2d["zmin"], psi_zmax=b2d["zmax"], psi_nz=b2d["nz"])

a5py/a5py/biosaw/biosaw5.py
ascot4fusion/ascot5
from .basic import *
from .conversion import *
from .radial import *
from .plot import *

a5py/a5py/dist/__init__.py
ascot4fusion/ascot5
"""
Basic operations for processing distributions.

File: basic.py
"""
import numpy as np

from a5py.physlib.alias import getalias as alias

def histogram2distribution(dist):
    """
    Convert histogram to distribution.

    A histogram tells how many markers there are in a cell whereas density,
    i.e. histogram divided by cell volume, is a physical quantity.

    The density will not necessarily be in SI units, as the volume element is
    defined by the abscissae as vol = (a1[1:] - a1[:-1]) x (a2[1:] - a2[:-1])
    x ... .

    Args:
        dist: dict_like <br>
            Distribution whose histogram ordinate is converted to a density.
    Returns:
        Reference to the dist argument.
    """
    if "distribution" in dist:
        return dist

    # Calculate volumes for each cell
    vol = 1
    for coord in dist["abscissae"]:
        edges = dist[coord + "_edges"]
        dv    = (edges[1:] - edges[:-1])
        vol   = np.multiply.outer(vol, dv)

    dist["distribution"] = dist["histogram"] / vol
    del dist["histogram"]
    return dist

def squeeze(dist, **kwargs):
    """
    Integrate or slice dimensions.

    Integration means a sum sum(dist[a:b,...] * vol) will be calculated and
    the dimensionality of the distribution is reduced by one. Slicing means
    that abscissa is cut as abscissa[a:b] and the ordinate is modified
    accordingly.

    This function operates recursively on each given dimension.

    Args:
        dist : dict_like <br>
            Distribution so be integrated or sliced.
        kwargs <br>
            Name(s) (R, phi, z, vpa, vpe, time, charge) of those dimensions
            along which the distribution is either sliced or integrated over.
            Names are given as keyword value pairs, {str : tuple} or
            {str : array_like}, where a tuple value (a, b) are indices for
            slicing and array [a, b] are indices for integration. A scalar zero
            means whole dimension is integrated.
    Returns:
        Reference to the dist argument.
    """

    # Are kwargs included?
    if len(kwargs.keys()) == 0:
        return dist

    coord  = alias(list(kwargs.keys())[0])
    slices = kwargs[coord]
    del kwargs[coord]

    # Call recursively if additional dimensions still exist in kwargs.
    if len(kwargs.keys()) > 0:
        squeeze(dist, **kwargs)

    dim = dist["abscissae"].index(coord)

    if type(slices) is tuple:
        dist[coord] = dist[coord][slices[0]:slices[1]]
        dist[coord + "_edges"] = dist[coord][slices[0]:slices[1]+1]
        dist["n" + coord]      = dist[coord].size

        # Slice given dimension
        dist["distribution"] = np.take(dist["distribution"],
                                       range(slices[0], slices[1]),
                                       axis=dim)

    if type(slices) is int:
        slices = [0, dist["n" + coord]]

    if type(slices) is list or not slices:
        edges   = dist[coord + "_edges"]
        weights = (edges[1:] - edges[0:-1])

        if slices is not None:
            mask = np.zeros(weights.shape)
            mask[slices[0]:slices[1]] = 1
            weights = weights*mask

        dist["distribution"], s = np.average(dist["distribution"],
                                             axis=dim, weights=weights,
                                             returned=True)
        dist["distribution"] *= s
        del dist[coord]
        del dist[coord + "_edges"]
        del dist["n" + coord]
        del dist["abscissae"][dim]

    return dist

a5py/a5py/dist/basic.py
ascot4fusion/ascot5
"""
Routines for converting distribution abscissae to other coordinates.

File: conversion.py
"""

import numpy as np
import copy
import scipy.constants as constants
import itertools

from scipy.interpolate import griddata, RectBivariateSpline


def convert_ppappe_to_Exi(dist, masskg, E_edges=None, xi_edges=None):
    """
    Converts vpa and vpe distribution abscissae to energy and pitch.

    This function operates by looping through all other coordinates except
    ppa and ppe, and at each loop calculates
    f_Exi(E, xi) = f_ppappe(ppa(E_i, xi_i), ppe(E_i, xi_i)) where E_i and
    xi_i are grid points of the new energy-pitch distribution. Interpolation
    is done bilinearly.

    Energy is in electronvolts and pitch is ppa/(ppa^2 + ppe^2)^0.5. The
    transformation is not relativistic.

    Args:
        dist : dict_like <br>
            A ppa-ppe distribution. May hold other dimensions as well.
        masskg : float <br>
            Mass of the species (required for energy conversion) in kg. Note
            that distribution is assumed to consist of markers with equal mass.
        E_edges : array_like or int, optional <br>
            Energy grid edges in the new distribution. If not given,
            linspace(0, Emax, 10) will be used where Emax is
            e*0.5*masskg*max(vpa^2, vpe^2). If an integer is given then
            linspace(0, Emax, E_edges) is used.
        xi_edges : array_like or int, optional <br>
            Pitch grid edges in the new distribution. If not given,
            linspace(-1, 1, 10) will be used. If an integer is given then
            linspace(-1, 1, xi_edges) is used.

    Returns:
       Energy-pitch distribution dictionary whose other dimensions are same as
       in input.
    """

    if E_edges is None:
        pmax2 = np.maximum( dist["ppar_edges"][-1]*dist["ppar_edges"][-1],
                           dist["pperp_edges"][-1]*dist["pperp_edges"][-1] )
        gamma = np.sqrt( 1 +  pmax2 / ( masskg.v * constants.c) **2)
        Emax = (1/constants.e) * (gamma - 1) * masskg.v * constants.c**2
        E_edges = np.linspace(0, Emax, 10)
    if xi_edges is None:
        xi_edges = np.linspace(-1, 1, 10)

    if not isinstance(E_edges, np.ndarray):
        pmax2 = np.maximum( dist["ppar_edges"][-1]*dist["ppar_edges"][-1],
                           dist["pperp_edges"][-1]*dist["pperp_edges"][-1] )
        gamma = np.sqrt( 1 +  pmax2 / ( masskg.v * constants.c) **2)
        Emax = (1/constants.e) * (gamma - 1) * masskg.v * constants.c**2
        E_edges = np.linspace(0, Emax, E_edges)

    if not isinstance(xi_edges, np.ndarray):
        xi_edges = np.linspace(-1, 1, xi_edges)

    ## Create E-xi distribution ##
    Exidist = copy.deepcopy(dist)

    # Remove vpa and vpe components
    del Exidist["distribution"]
    Exidist["abscissae"].remove("ppar")
    Exidist["abscissae"].remove("pperp")
    for k in list(Exidist):
        if "ppar" in k or "pperp" in k:
            del Exidist[k]

    # Add E and xi abscissae and initialize a new density
    Exidist["abscissae"].insert(3, "energy")
    Exidist["abscissae"].insert(4, "pitch")
    abscissae = Exidist["abscissae"]

    Exidist["energy"]       = (E_edges[0:-1] + E_edges[1:]) / 2
    Exidist["energy_edges"] = E_edges
    Exidist["nenergy"]      = Exidist["energy"].size

    Exidist["pitch"]       = (xi_edges[0:-1] + xi_edges[1:]) / 2
    Exidist["pitch_edges"] = xi_edges
    Exidist["npitch"]      = Exidist["pitch"].size

    dims = []
    for a in abscissae:
        dims.append(Exidist["n" + a])

    Exidist["distribution"]  = np.zeros(tuple(dims))

    # Transform E-xi grid to points in (vpa,vpa) space that are used in
    # interpolation.
    xig, Eg = np.meshgrid(Exidist["pitch"], Exidist["energy"])
    pg   = np.sqrt( ( Eg * constants.e / constants.c + masskg.v * constants.c )**2
                     - (masskg.v * constants.c)**2 )
    ppag = ( xig * pg ).ravel()
    ppeg = (np.sqrt(1 - xig*xig) * pg).ravel()

    # Coordinate transform Jacobian: dvpa dvpe = |jac| dE dxi
    # Jacobian for transform (ppa, ppe) -> (p, xi) is p / sqrt(1-xi^2)
    # because jac = dppa / dp  = xi, dppe / dp  = sqrt(1-xi^2)
    #               dppa / dxi = p,  dppe / dxi = -xi p / sqrt(1-xi^2),
    # and the Jacobian for (p, xi) -> (E, xi) is e E0 / c^2 p when
    # E is in electronvolts. Therefore the combined Jacobian is
    # (e E0 / c^2) / sqrt(1-xi*xi).
    E0  = np.sqrt( (pg*constants.c)**2 + (masskg.v*constants.c**2)**2 )
    jac = (constants.e * E0 / constants.c**2) / np.sqrt(1 - xig*xig)

    # Interpolate.
    ranges = []
    for a in dist["abscissae"]:
        if a != "ppar" and a != "pperp":
            ranges.append(range(dist["n" + a]))

    iE   = Exidist["abscissae"].index("energy")
    ippa = dist["abscissae"].index("ppar")
    for itr in itertools.product(*ranges):

        idx = []
        for i in range(0, ippa):
            idx.append(itr[i])

        idx.append(slice(None))
        idx.append(slice(None))

        for i in range(ippa, len(dist["abscissae"])-2):
            idx.append(itr[i])

        idx = tuple(idx)

        if dist["ppar"].size == 1 and dist["pperp"].size == 1:
            d = np.ones( (Exidist["nenergy"], Exidist["npitch"]) ) \
                * dist["distribution"][idx]
        else:
            f = RectBivariateSpline(
                dist["ppar"], dist["pperp"],
                np.squeeze(dist["distribution"][idx]),
                kx=1, ky=1)
            d = np.reshape(f.ev(ppag, ppeg),
                           (Exidist["nenergy"], Exidist["npitch"]))

        Exidist["distribution"][idx] = d * jac

    return Exidist

a5py/a5py/dist/conversion.py
ascot4fusion/ascot5
"""
Routines for plotting distributions.

File: dist/plot.py
"""
import numpy as np
import importlib.util as util

plt = util.find_spec("matplotlib")
if plt:
    import matplotlib.pyplot as plt

def plot_dist_1D(dist, logscale=False, axes=None,label=None):
    """
    Plot distribution as a 1D plot.

    This function assumes the given distribution is squeezed so that a single
    dimension remains.

    Args:
        dist : dict_like <br>
            Distribution dictionary.
        axes : Axes, optional <br>
            Axes to which the distribution is plotted. If None, a new figure is
            created and displayed.
    """

    newfig = axes is None
    if newfig:
        plt.figure()
        axes = plt.gca()

    ordinate = None
    if "distribution" in dist:
        ordinate = dist["distribution"]
    elif "histogram" in dist:
        ordinate = dist["histogram"]
    elif "density" in dist:
        ordinate = dist["density"]
    elif "epowerdeposition" in dist:
        ordinate = dist["epowerdeposition"]

    x = dist["abscissae"][0]

    if logscale:
        ordinate = np.log10(ordinate)

    axes.plot(dist[x], ordinate,label=label)
    axes.set_xlabel(x);
    axes.tick_params(axis='x', direction='out')
    axes.tick_params(axis='y', direction='out')

    if newfig:
        plt.show(block=False)


def plot_dist_2D(dist, *args, logscale=False, equal=False, axes=None, cmap=False):
    """
    Plot distribution as a 2D plot (pcolormesh).

    This function assumes the given distribution is squeezed so that only two
    dimensions remain.

    Args:
        dist : dict_like <br>
            Distribution dictionary.
        args : str, str <br>
            Name of the x and y coordinates e.g. "R", "z".
        axes : Axes, optional <br>
            Axes to which the distribution is plotted. If None, a new figure is
            created and displayed.
        cmap : optional <br>
            If True, the cmap gnuplot is used
    """

    newfig = axes is None
    if newfig:
        plt.figure()
        axes = plt.gca()

    ordinate = None
    if "distribution" in dist:
        ordinate = dist["distribution"]
    elif "histogram" in dist:
        ordinate = dist["histogram"]

    if len(args) == 0:
        x = dist["abscissae"][0]
        y = dist["abscissae"][1]
    else:
        x = args[0]
        y = args[1]

    if x == dist["abscissae"][0]:
        ordinate = np.transpose(ordinate)

    if logscale:
        ordinate = np.log10(ordinate)

    ordinate = np.ma.masked_invalid(ordinate)

    if cmap:
        mesh = axes.pcolormesh(dist[x], dist[y], ordinate, cmap=plt.cm.gnuplot  , vmin=np.nanmin(ordinate), vmax=np.nanmax(ordinate))
    else:
        mesh = axes.pcolormesh(dist[x], dist[y], ordinate, vmin=np.nanmin(ordinate), vmax=np.nanmax(ordinate))
	
    # https://stackoverflow.com/a/16125413/190597 (Joe Kington)
    # and https://stackoverflow.com/a/35905483
    axes.patch.set(hatch='x', edgecolor=[0.9, 0.9, 0.9])
    plt.colorbar(mesh, ax=axes)

    axes.set_xlabel(x);
    axes.set_ylabel(y);
    axes.tick_params(axis='x', direction='out')
    axes.tick_params(axis='y', direction='out')

    if equal:
        axes.axis("image")
    else:
        axes.axis("tight")

    if newfig:
        plt.show(block=False)

a5py/a5py/dist/plot.py
ascot4fusion/ascot5
"""
Evaluate radial (1D) distributions from 5D or 6D distributions.

File: radial.py
"""
import copy
import numpy as np

from . import basic as distmod
from . import conversion as distconv
import unyt

def eval1d(ascotpy, dist, quantity, rhomin, rhomax, nrho, ma=None, qa=None, vol=None, area=None):
    '''
    Args units:
            ma : float <br>
                Test particle mass [kg].
            qa : float <br>
                Test particle charge [C].
    '''
    
    #if ma is not None:
    #    print('Marker mass {} kg'.format(ma))
    #if qa is not None:
    #    print('Marker charge {} C'.format(qa))
        
    
    # working copy
    dist = copy.deepcopy(dist)

    if vol is None or area is None:
        # Find rho volume
        vol,area = evalrhovol(
            ascotpy, rhomin, rhomax, nrho,
            dist["r_edges"][0], dist["r_edges"][-1], dist["r_edges"].size,
            dist["phi_edges"][0], dist["phi_edges"][-1], dist["phi_edges"].size,
            dist["z_edges"][0], dist["z_edges"][-1], dist["z_edges"].size )

    dist = eval_quantity_5d(ascotpy, dist, quantity, ma, qa)
    
    # Map  (R,phi,z) to rho
    rho = ascotpy.evaluate(dist["r"], np.deg2rad(dist["phi"]), dist["z"], 0,
                           "rho", grid=True)
    
    rho = np.squeeze(rho)
    rzpvol = (dist["r_edges"][1] - dist["r_edges"][0]) \
             * (dist["phi_edges"][1] - dist["phi_edges"][0]) \
             * (dist["z_edges"][1] - dist["z_edges"][0])

    dist = dist["distribution"] * rzpvol

    #idx   = np.logical_and(rho >= rhomin, rho <= rhomax)

    newdist = {}
    newdist["rho_edges"] = np.linspace(rhomin, rhomax, nrho+1)
    newdist["rho"] = newdist["rho_edges"][:-1] \
                     + (newdist["rho_edges"][1] - newdist["rho_edges"][0]) / 2
    newdist["nrho"] = nrho
    newdist["abscissae"] = ["rho"]
    newdist["vol"] = vol.ravel()
    newdist["area"] = area.ravel()

    newdist[quantity]  = np.histogram( rho.ravel(), newdist["rho_edges"],
                                       weights=dist.ravel() )[0]
    newdist[quantity] /= newdist["vol"]
    
    return newdist

def eval1d_fromRhodist(ascotpy, dist, quantity, rmin, rmax, zmin, zmax, ma=None, qa=None, vol=None, area=None):
    '''
    Args units:
            ma : float <br>
                Test particle mass [kg].
            qa : float <br>
                Test particle charge [C].
    '''
    
    #if ma is not None:
    #    print('Marker mass {} kg'.format(ma))
    #if qa is not None:
    #    print('Marker charge {} C'.format(qa))
        


    
    # working copy
    dist = copy.deepcopy(dist)

    if vol is None or area is None:
        # Find rho volume
        vol,area = evalrhovol(
            ascotpy, 
            dist["rho_edges"][0], dist["rho_edges"][-1], dist["rho_edges"].size-1,
            rmin, rmax, None,
            dist["phi_edges"][0], dist["phi_edges"][-1], None,
            zmin, zmax, None )

    dist = eval_quantity_5d_rho(ascotpy, dist, quantity, ma, qa)
    distmod.squeeze(dist, theta=0, phi=0)

    dist["vol"] = vol.ravel()
    dist["area"] = area.ravel()

    # Move from density per rho-slot to density per m3
    dist[quantity] = dist['distribution'] * (dist["rho_edges"][1] - dist["rho_edges"][0]) /  dist["vol"]
    del( dist['distribution'] )

    return dist

def eval_quantity_5d_rho(ascotpy, dist, quantity, ma, qa):
    # Evaluate the requested quantity in 5D
    if quantity == "density":
        dist = distmod.squeeze(dist, ppar=0, pperp=0, time=0, charge=0)

    elif quantity == "energydensity":
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        dist = distmod.squeeze(dist, pitch=0, time=0, charge=0)

        distE = copy.deepcopy(dist)
        for iphi in range(dist["phi"].size):
            for itheta in range(dist["theta"].size):
                for irho in range(dist["rho"].size):
                    distE["distribution"][irho, itheta, iphi, :] = dist["distribution"][irho, itheta, iphi,:] * dist["energy"] * unyt.e

        dist = distmod.squeeze(distE, energy=0)

    elif quantity == "pressure":
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        dist = distmod.squeeze(dist, pitch=0, time=0, charge=0)

        distp = copy.deepcopy(dist)
        for iphi in range(dist["phi"].size):
            for itheta in range(dist["theta"].size):
                for irho in range(dist["rho"].size):
                    velocity2 = 2 * dist["energy"] * unyt.e / ma
                    distp["distribution"][irho, itheta, iphi, :] = dist["distribution"][irho, itheta, iphi,:] * ma * velocity2 / 3.0

        dist = distmod.squeeze(distp, energy=0)

    elif quantity == "toroidalcurrent":
        dist = distmod.squeeze(dist, pperp=0, time=0, charge=0)

        distj = copy.deepcopy(dist)

        for iphi in range(dist["phi"].size):
            for itheta in range(dist["theta"].size):
                                #Calculate R,phi,z for all rho-values for single phi,theta
                (r, z) = ascotpy.get_rhotheta_rz(dist["rho"][:], dist["theta"][itheta], dist["phi"][iphi], time=0.0)
                phi = np.ones_like(r) * dist["phi"][iphi]

                for irho in range(dist["rho"].size):
                    bphi = ascotpy.evaluate( r[irho], phi[irho], z[irho], 0.0, "bphi")
                    bnorm = ascotpy.evaluate(r[irho], phi[irho], z[irho], 0.0, "bnorm")
                    for ippar in range(dist["ppar"].size):
                        distj["distribution"][irho, itheta, iphi, ippar] = dist["distribution"][irho, itheta, iphi, ippar] * dist["ppar"][ippar]/ma * qa * bphi / bnorm
        dist = distmod.squeeze(distj, ppar=0)

    
    elif quantity == "ipowerdeposition" or quantity == "epowerdeposition" or quantity == "powerdeposition":
        
                
        # Convert to pitch-energy distribution.
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        # Integrate over pitch, time and charge dimensions
        dist = distmod.squeeze(dist, pitch=0, time=0, charge=0)

        # Convert the energy abscissa into velocity
        va = np.sqrt( unyt.e*2*dist["energy"]/ma )

        # Evaluate the ASCOT collision operator at each spatial location for all velocities in va
        distP = copy.deepcopy(dist)
        for iphi in range(dist["phi"].size):
            for itheta in range(dist["theta"].size):

                #Calculate R,phi,z for all rho-values for single phi,theta
                (r, z) = ascotpy.get_rhotheta_rz(dist["rho"][:], dist["theta"][itheta], dist["phi"][iphi], time=0.0)
                phi = np.ones_like(r) * dist["phi"][iphi]

                for irho in range(dist["rho"].size):
                    
                    coefs = ascotpy.eval_collcoefs(ma, qa, r[irho],
                                                   phi[irho], z[irho], 0, va)
                    if   quantity == "ipowerdeposition":
                        distP["distribution"][irho, itheta, iphi, :] = -dist["distribution"][irho, itheta, iphi, :]*sum(coefs["K"][0,1:,:],0) *ma*va
                    elif quantity == "epowerdeposition": 
                        distP["distribution"][irho, itheta, iphi, :] = -dist["distribution"][irho, itheta, iphi, :]*    coefs["K"][0, 0,:]    *ma*va
                    elif quantity == "powerdeposition":
                        distP["distribution"][irho, itheta, iphi, :] = -dist["distribution"][irho, itheta, iphi, :]*sum(coefs["K"][0, :,:],0) *ma*va

        # Integrate over the energy dimension
        dist = distmod.squeeze(distP, energy=0)

    elif quantity == "colltorque":

        # Convert to pitch-energy distribution.
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        # Integrate over pitch, time and charge dimensions
        dist = distmod.squeeze(dist, time=0, charge=0)

        # Convert the energy abscissa into velocity
        va = np.sqrt( unyt.e*2*dist["energy"]/ma.v )

        # Evaluate the ASCOT collision operator at each spatial location for all velocities in va
        distT = copy.deepcopy(dist)
        for iphi in range(dist["phi"].size):
            for itheta in range(dist["theta"].size):

                #Calculate R,phi,z for all rho-values for single phi,theta
                (r, z) = ascotpy.get_rhotheta_rz(dist["rho"][:], dist["theta"][itheta], dist["phi"][iphi], time=0.0)
                phi = np.ones_like(r) * dist["phi"][iphi]

                for irho in range(dist["rho"].size):

                    coefs = ascotpy.eval_collcoefs(ma, qa, r[irho],
                                                   phi[irho], z[irho], 0, va)
                    Bnorm = ascotpy.evaluate(r[irho],phi[irho], z[irho], 0, "bnorm")
                    Bphi = ascotpy.evaluate(r[irho], phi[irho], z[irho], 0, "bphi")

                    v = np.tile(va,(49,1)).T
                    pitch = np.tile(dist["pitch"],(99,1))
                    K = np.tile(sum(coefs["K"][0,:,:],0),(49,1)).T
                    nu = np.tile(sum(coefs["nu"][0,:,:],0),(49,1)).T

                    deflectFreq = nu
                    dpitch = -deflectFreq*pitch
                    dp1 = K * ma
                    dPpara = pitch * dp1 + ma*v*dpitch
                    colltorque = dPpara * r[irho] * Bphi/Bnorm

                    distT["distribution"][irho, itheta, iphi, :] = dist["distribution"][irho, itheta, iphi, :]*-colltorque


        # Integrate over the energy dimension
        dist = distmod.squeeze(distT, energy=0, pitch=0)

    else:
        raise ValueError('Unknown quantity "{}".'.format(quantity))

    return dist


def eval_quantity_5d(ascotpy, dist, quantity, ma, qa):
    # Evaluate the requested quantity in 5D
    if quantity == "density":
        dist = distmod.squeeze(dist, ppar=0, pperp=0, time=0, charge=0)

    elif quantity == "energydensity":
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        dist = distmod.squeeze(dist, pitch=0, time=0, charge=0)

        distE = copy.deepcopy(dist)
        for ir in range(dist["r"].size):
            for ip in range(dist["phi"].size):
                for iz in range(dist["z"].size):
                    distE["distribution"][ir, ip, iz, :] = dist["distribution"][ir,ip,iz,:] * dist["energy"] * unyt.e

        dist = distmod.squeeze(distE, energy=0)

    elif quantity == "pressure":
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        dist = distmod.squeeze(dist, pitch=0, time=0, charge=0)

        distp = copy.deepcopy(dist)
        for ir in range(dist["r"].size):
            for ip in range(dist["phi"].size):
                for iz in range(dist["z"].size):
                    velocity2 = 2 * dist["energy"] * unyt.e / ma
                    distp["distribution"][ir, ip, iz, :] = dist["distribution"][ir,ip,iz,:] * ma * velocity2 / 3.0

        dist = distmod.squeeze(distp, energy=0)

    elif quantity == "poloidalcurrent":
        dist = distmod.squeeze(dist, pperp=0, time=0, charge=0)

        # current distribution
        distj = copy.deepcopy(dist)

        # integration over toroid angles
        for ir in range(dist["r"].size):
            for ip in range(dist["phi"].size):
                for iz in range(dist["z"].size):

                    bphi = ascotpy.evaluate(dist["r"][ir], dist["phi"][ip], dist["z"][iz], 0.0, "bphi")
                    bnorm = ascotpy.evaluate(dist["r"][ir], dist["phi"][ip], dist["z"][iz], 0.0, "bnorm")
                    
                    for ippar in range(dist["ppar"].size):
                        distj["distribution"][ir, ip, iz, ippar] = dist["distribution"][ir,ip,iz,ippar] * dist["ppar"][ippar]/ma * qa * bphi / bnorm
                        
        dist = distmod.squeeze(distj, ppar=0)   
    
    elif quantity == "toroidalcurrent":
        dist = distmod.squeeze(dist, pperp=0, time=0, charge=0)

        distj = copy.deepcopy(dist)

        for ir in range(dist["r"].size):
            for ip in range(dist["phi"].size):
                for iz in range(dist["z"].size):
                    bphi = ascotpy.evaluate(dist["r"][ir], dist["phi"][ip], dist["z"][iz], 0.0, "bphi")
                    bnorm = ascotpy.evaluate(dist["r"][ir], dist["phi"][ip], dist["z"][iz], 0.0, "bnorm")
                    for ippar in range(dist["ppar"].size):
                        distj["distribution"][ir, ip, iz, ippar] = dist["distribution"][ir,ip,iz,ippar] * dist["ppar"][ippar]/ma * qa * bphi / bnorm
        dist = distmod.squeeze(distj, ppar=0)

    elif quantity == "ipowerdeposition" or quantity == "epowerdeposition" or quantity == "powerdeposition":
        
        # Convert to pitch-energy distribution.
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        # Integrate over pitch, time and charge dimensions
        dist = distmod.squeeze(dist, pitch=0, time=0, charge=0)

        # Convert the energy abscissa into velocity
        va = np.sqrt( unyt.e*2*dist["energy"]/ma )

        # Evaluate the ASCOT collision operator at each spatial location for all velocities in va
        distP = copy.deepcopy(dist)
        for ir in range(dist["r"].size):
            for ip in range(dist["phi"].size):
                for iz in range(dist["z"].size):
                    coefs = ascotpy.eval_collcoefs(ma, qa, dist["r"][ir],
                                                   dist["phi"][ip], dist["z"][iz], 0, va)
                    
                    if   quantity == "ipowerdeposition":
                        distP["distribution"][ir, ip, iz, :] = -dist["distribution"][ir, ip, iz, :]*sum(coefs["K"][0,1:,:],0) *ma*va
                    elif quantity == "epowerdeposition": 
                        distP["distribution"][ir, ip, iz, :] = -dist["distribution"][ir, ip, iz, :]*    coefs["K"][0, 0,:]    *ma*va
                    elif quantity == "powerdeposition":
                        distP["distribution"][ir, ip, iz, :] = -dist["distribution"][ir, ip, iz, :]*sum(coefs["K"][0, :,:],0) *ma*va
                    

        # Integrate over the energy dimension
        dist = distmod.squeeze(distP, energy=0)

    elif quantity == "colltorque":

        # Convert to pitch-energy distribution.
        dist = distconv.convert_ppappe_to_Exi(dist, ma, 100, 50)
        # Integrate over pitch, time and charge dimensions
        dist = distmod.squeeze(dist, time=0, charge=0)

        # Convert the energy abscissa into velocity
        va = np.sqrt( unyt.e*2*dist["energy"]/ma )

        # Evaluate the ASCOT collision operator at each spatial location for all velocities in va
        distT = copy.deepcopy(dist)
        for ir in range(dist["r"].size):
            for ip in range(dist["phi"].size):
                for iz in range(dist["z"].size):
                    coefs = ascotpy.eval_collcoefs(ma, qa, dist["r"][ir],
                                                   dist["phi"][ip], dist["z"][iz], 0, va)
                    Bnorm = ascotpy.evaluate(dist["r"][ir],dist["phi"][ip], dist["z"][iz], 0, "bnorm")
                    Bphi = ascotpy.evaluate(dist["r"][ir],dist["phi"][ip], dist["z"][iz], 0, "bphi")

                    v = np.tile(va,(49,1)).T
                    pitch = np.tile(dist["pitch"],(99,1))
                    K = np.tile(sum(coefs["K"][0,:,:],0),(49,1)).T
                    nu = np.tile(sum(coefs["nu"][0,:,:],0),(49,1)).T

                    deflectFreq = nu
                    dpitch = -deflectFreq*pitch
                    dp1 = K * ma
                    dPpara = pitch * dp1 + ma*v*dpitch
                    colltorque = dPpara * dist["r"][ir] * Bphi/Bnorm

                    distT["distribution"][ir, ip, iz, :] = dist["distribution"][ir, ip, iz, :]*colltorque

        # Integrate over the energy dimension
        dist = distmod.squeeze(distT, energy=0, pitch=0)


    else:
        raise ValueError('Unknown quantity "{}".'.format(quantity))

    return dist

def evalrhovol(ascotpy, rhomin, rhomax, nrho, rmin, rmax, nr,
               phimin, phimax, nphi, zmin, zmax, nz):

    # Volume of the test space (rectangular toroid)
    V_tot = 0.5 * (rmax**2 - rmin**2) * np.deg2rad(phimax - phimin) * (zmax - zmin)
    A_tot = (rmax - rmin) * (zmax - zmin)
    drho = (rhomax - rhomin) / nrho

    points_in_bins = np.zeros((nrho, 1))

    n     = 100000
    n_tot = 0
    nmin  = 2

    while np.any(points_in_bins < nmin):
        r   = np.random.uniform(rmin, rmax, n)
        phi = np.random.uniform(phimin, phimax, n)
        z   = np.random.uniform(zmin, zmax, n)

        rho   = ascotpy.evaluate(r, np.deg2rad(phi), z, 0, "rho")
        ind   = np.logical_and(rho < rhomax, rho >= rhomin)
        i_rho = np.floor((rho[ind]-rhomin)/drho).astype(int)

        np.add.at(points_in_bins, (i_rho, 0), 1)

        n_tot = n_tot + n

    return (points_in_bins/n_tot) * V_tot, (points_in_bins/n_tot) * A_tot

a5py/a5py/dist/radial.py
ascot4fusion/ascot5
# coding: utf-8

# For Simppa at least, following recipe worked to get the below working:
# 0. compile ascotlib.so and add the path to LD_LIBRARY_PATH
# 1. in ascot5.yml, set the python version: python=3.6, and create the environment
# 2. module load imasenv
# 3. conda activate ascot5
# the below works.

import ctypes
from   a5py import Ascot
import a5py.ascot5io.imas
import a5py.ascot5io.wall
from   a5py.ascotpy import ascot2py
from   a5py.ascot5io.marker import Marker
from a5py.templates.optiontemplates import OptionTemplates
import numpy as np

# Read from IMAS

print('Reading 2D wall')
# 2D wall
w2d=a5py.ascot5io.imas.wall_2d()
#wdict2=w2d.read("g2jvarje","test","3",92436,272)
wdict2=w2d.read("akaslos","test","3",92436,272)
a5py.ascot5io.wall_2D.write_hdf5('from_imas.h5',**wdict2)

print('Reading 3D wall')
# 3D wall
w3d=a5py.ascot5io.imas.wall_3d()
wdict3=w3d.read("akaslos","mywall","3",201,101)
#print(wdict3['x1x2x3'].shape, wdict3['y1y2y3'].shape, wdict3['z1z2z3'].shape, )
a5py.ascot5io.wall_3D.write_hdf5('from_imas.h5',**wdict3)

print('Reading Bfield (stellarator)')
# 3D Bfield in Stellarator format
bsts=a5py.ascot5io.imas.B_STS()
bdict=bsts.read("akaslos","ggdtest","3",32,3)
a5py.ascot5io.B_STS.write_hdf5('from_imas.h5',**bdict)


print('Initializing ascot5')


a5 = Ascot('helloworld.h5')
#M=a5py.ascotpy.ascot5_main.ascot5_main(input_filename=b'helloworld.h5',output_filename=b'helloworld_out.h5')



#print("Inializing.")
#M.init()

# don't read wall input
#what_to_read = ctypes.c_int32()
#what_to_read.value =    ascotpy2.hdf5_input_options | \
#                        ascotpy2.hdf5_input_efield  | ascotpy2.hdf5_input_plasma | \
#                        ascotpy2.hdf5_input_neutral | ascotpy2.hdf5_input_marker |\
#                        ascotpy2.hdf5_input_boozer  | ascotpy2.hdf5_input_mhd

#  ascotpy2.hdf5_input_wall |
# ascotpy2.hdf5_input_bfield | \

#M.read_input(what_to_read=what_to_read)
print('Initializing inputs')

#a5.simulation_initinputs(bfield=True, efield=True, plasma=True, neutral=True,
#                        wall=True, boozer=True, mhd=True, switch=True)

bfield=bdict # True
wall=wdict3  # wdict2, True
a5.simulation_initinputs(bfield=bdict, efield=True, plasma=True, neutral=True,
                         wall=wdict3, boozer=True, mhd=True, switch=True)


print('generating markers')
nmarkers = 10
mrk = Marker.generate("gc", n=nmarkers, species="alpha")
mrk["energy"][:] = 3.5e6
mrk["pitch"][:]  = 0.99 - 1.98 * np.random.rand(nmarkers,)
mrk["r"][:]      = np.linspace(3.5, 5.4, nmarkers)
a5.data.create_input("gc", **mrk)

#print('Reading markers')
#mrk = a5.data.marker.active.read()
#print(mrk.keys())
print('providing markers')
a5.simulation_initmarkers(**mrk)


print('generating options')
#a5.data.create_input("options tutorial")
opt = OptionTemplates().options_tutorial()

#print('Reading options')
#opt = a5.data.options.active.read()
print('providing options')
a5.simulation_initoptions(**opt[1])



# Reduce the simulation time a bit:
max_simtime=1.0e-4
print("Reducing max_simtime to {}s.".format(max_simtime))
#M.sim.endcond_max_simtime = max_simtime
a5._sim.endcond_lim_simtime=max_simtime


#print("Injecting 2D wall from imas")
#M.inject_wall_2d( wdict['r'], wdict['z']  )

#print("Injecting 3D wall from imas")
#a5.provide_wall_3d( wdict3['x1x2x3'], wdict3['y1y2y3'], wdict3['z1z2z3']  )

#print("Injecting stellarator B-field from imas")
#a5.provide_BSTS(bdict)


print('starting simulation')
vrun_output = a5.simulation_run()





#print("offloading")
#M.offload()

#print("simulating")
#M.run_simulation()

#print("gathering output")
#M.gather_output()

#print("printing marker summary")
#M.print_marker_summary()

#print("writing output h5")
#M.write_output_h5()

print("freeing memory")
#M.free_c()
a5.simulation_free(diagnostics=True)

#print("finalizing MPI")
#M.finalize()

print('Finished test')


#print('Stopping test')
#quit()

a5py/a5py/examples/imas_run.py
ascot4fusion/ascot5
# coding: utf-8

# For Simppa at least, following recipe worked to get the below working:
# 1. in ascot5.yml, set the python version: python=3.6, and create the environment
# 2. module load imasenv
# 3. conda activate ascot5
# the below works.


import a5py.ascot5io.imas
import a5py.ascot5io.wall_2D
import a5py.ascot5io.wall_3D
import a5py.ascot5io.B_STS
w2d=a5py.ascot5io.imas.wall_2d()
w3d=a5py.ascot5io.imas.wall_3d()
bst=a5py.ascot5io.imas.B_STS()
# Gateway
# wdict=w2d.read("g2jvarje","test","3",92436,272)

# ITER SDCC
print("Reading 3D wall")
wdict3=w3d.read("akaslos","mywall","3",11,11)

print("Reading 2D wall")
wdict=w2d.read("akaslos","test","3",92436,272)

print("Reading 3D field")
#b3ddict = bst.read("akaslos","ggdtest","3",16,2)
b3ddict = bst.read("akaslos","ggdtest","3",32,1)

print("Writing to hdf5")
a5py.ascot5io.wall_3D.write_hdf5('from_imas.h5',**wdict3)
a5py.ascot5io.wall_2D.write_hdf5('from_imas.h5',**wdict)
a5py.ascot5io.B_STS.write_hdf5('from_imas.h5',**b3ddict)


print("Writing to IDS")
#bst.write( b3ddict, "akaslos", "ggdtest", "3", 32, 1, metadata={'comment':'written by read_from_imas.py'})
bst.write( b3ddict, "akaslos", "ggdtest", "3", 32, 2, metadata={'comment':'written by read_from_imas.py'})


a5py/a5py/examples/read_from_imas.py
ascot4fusion/ascot5
# coding: utf-8

# get helloworld.h5 with
# python ../preprocessing/simpleruns.py 

from a5py.ascot5io.ascot5 import Ascot
import a5py.ascot5io.imas
import a5py.ascot5io.wall_2D
import a5py.ascot5io.wall_3D

filename='helloworld.h5'
a=Ascot(filename)
wdict3 = a.wall.active.read()
bsts   = a.bfield.active.read()

w3dimas=a5py.ascot5io.imas.wall_3d()
bstsimas=a5py.ascot5io.imas.B_STS()
w3dimas.write( wdict3,'akaslos','mywall', '3',201,101, {'comment':filename} )
bstsimas.write(  bsts,'akaslos','ggdtest','3',32,   3, {'comment':filename} )



a5py/a5py/examples/write_to_imas.py
ascot4fusion/ascot5
"""Contains definitions of this package's exceptions.
"""

class AscotIOException(Exception):
    """Raised when there is an internal error in `Ascot5IO`.

    This exception should be raised in cases where the issue is most certainly
    a bug in code or inconsistent HDF5 file.
    """
    pass

class AscotNoDataException(Exception):
    """Raised when required input or output is not present.

    This exception should be raised when user tries to query input or results
    that is not present. Other case is when using methods that require the
    presence of certain input or output data (e.g. when plotting orbits but
    orbit data is not present).
    """
    pass

class AscotInitException(Exception):
    """Raised when there is an issue with data initialized by `Ascot`.

    This exception should be raised whenever there is an issue with `libascot`
    or any related data.
    """
    pass

class AscotUnitWarning(UserWarning):
    """Warning raised when quantities are provided without units.
    """
    pass

a5py/a5py/exceptions.py
ascot4fusion/ascot5
"""Graphical user interface for ASCOT5.
"""

a5py/a5py/gui/__init__.py
ascot4fusion/ascot5
"""Contains definition of various GUI component classes.
"""
import tkinter
import tkinter.ttk as ttk

import numpy as np
import matplotlib.pyplot as plt

from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk

class NumEntry(tkinter.Frame):
    """
    A component to receive numerical input.
    """

    def __init__(self, master, isint=False, labeltext=None, defval=None,
                 entrywidth=6, labelwidth=8, anchor="w"):
        """
        Initialize a frame which accepts inputs.

        The frame consists of label and entrybox side-by-side. The entry only
        accepts valid numerical input. The value is obtained with getval()
        method.
        """
        super().__init__(master)
        self.isint = isint

        self.choice = tkinter.StringVar(self)

        if defval is not None:
            self.choice.set(defval)

        self.label = tkinter.Label(self, text=labeltext, anchor=anchor,
                                   width=labelwidth, font=("Calibri 10"))

        vcmd  = self.register(self._validate)
        self.entry = tkinter.Entry(
            self, validate = "all", validatecommand=(vcmd, "%P"),
            width=entrywidth, textvariable=self.choice, font=("Calibri 10"))

        self.label.grid(row=0, column=0, sticky="W")
        self.entry.grid(row=0, column=1, sticky="E")


    def getval(self):
        """
        Parse entry to a valid float or int.
        """
        val = self.choice.get()

        val = val.split("e")

        s1 = val[0]
        if s1 in ["", ".", "+", "-"]:
            s1 = 0
        else:
            if self.isint:
                s1 = int(s1)
            else:
                s1 = float(s1)

        if len(val) == 2:
            s2 = val[1]
            if s2 in ["", "+", "-"]:
                s2 = 0
            else:
                s2 = int(s2)

            s1 = s1 * np.power(10.0, s2)

        if self.isint:
            return int(s1)
        else:
            return float(s1)


    def setval(self, val):
        self.choice.set(val)


    def isempty(self):
        """
        Check if entry is just an empty string.
        """
        val = self.choice.get()
        return val == ""


    def disable(self):
        self.label.configure(fg="gray")
        self.entry.config(state="disabled")


    def enable(self):
        self.label.configure(fg="black")
        self.entry.config(state="normal")

    def setlabel(self, text):
        self.label.configure(text=text)


    def _validate(self, P):
        """
        Check if entry is a valid float or int.
        """

        # Check if there is (exactly one) exponent
        if P.count("e") > 1:
            return False
        P = P.split("e")

        # Check that the prefix is valid
        s = P[0]
        if not ( (s == "" or s in "+-") or (not self.isint and s in ".") ) :
            try:
                if self.isint:
                    int(s)
                else:
                    float(s)
            except ValueError:
                return False

        # Check that exponent is valid
        if len(P) == 2:
            s = P[1]
            if not ( (s == "" or s in "+") or (not self.isint and s in "-") ):
                try:
                    int(s)
                except ValueError:
                    return False

        return True


class DropdownMenu(tkinter.Frame):
    """
    Dropdown menu where user can choose from given options.
    """

    def __init__(self, master, defval=None, values=None, log=False, trace=None,
                 label=None, width=6, labelwidth=2, labelanchor="c"):
        super().__init__(master)
        self.var = tkinter.StringVar(self)

        if defval is not None:
            self.var.set(defval)

        if trace is not None:
            self.var.trace('w', trace)

        if label is not None:
            label = tkinter.Label(self, text=label, width=labelwidth)
            label.pack(side="left", anchor=labelanchor)

        self.menu = ttk.Combobox(self, width=width, textvariable=self.var,
                                 state="readonly")
        self.menu.pack(side="left")
        if values is not None:
            self.menu["values"] = values

        self.log = None
        if log:
            self.log = tkinter.IntVar(self)
            logtick = tkinter.Checkbutton(
                self, text=" log10", onvalue=1,
                offvalue=0, variable=self.log, height=1, width=5)
            logtick.pack(side="left")

            if trace is not None:
                self.log.trace('w', trace)


    def islog(self):
        if self.log is None:
            return False

        return self.log.get()


    def getval(self):
        return self.var.get()

    def setvals(self, values, defval, log=None):
        self.menu["values"] = values
        self.var.set(defval)
        if log is not None:
            self.log.set(log)


    def settrace(self, trace):
        self.var.trace('w', trace)
        if self.log is not None:
            self.log.trace('w', trace)


class Tickbox(tkinter.Checkbutton):
    """
    A tickbox and label.
    """

    def __init__(self, master, defval=None, trace=None, label=None, width=8,
                 justify=None):
        self.var = tkinter.IntVar()
        super().__init__(master, text=label, variable=self.var,
                         onvalue=1, offvalue=0, height=1, width=width,
                         justify=justify)
        if defval is not None:
            self.var.set(defval)

        if trace is not None:
            self.var.trace('w', trace)


    def getval(self):
        return self.var.get()


class NavToolbarNocoord(NavigationToolbar2Tk):
    """
    Navigation toolbar but without the coordinate display
    """
    def set_message(self, msg):
        pass


class PlotFrame(tkinter.Frame):
    """
    Frame containing matplotlib plot and NavToolbarNocoord
    """

    def __init__(self, master, tight_layout=True):
        super().__init__(master)
        fig = plt.figure(tight_layout=tight_layout)
        figcanvas = FigureCanvasTkAgg(fig, self)

        toolbar = NavToolbarNocoord(figcanvas, self, pack_toolbar=False)
        toolbar.config(background="white")
        for button in toolbar.winfo_children():
            button.config(background="white")
        toolbar.update()

        toolbar.pack(side=tkinter.TOP, fill=tkinter.X)
        figcanvas.get_tk_widget().pack(fill=tkinter.BOTH, expand=True)

        self.fig       = fig
        self.axes      = self.set_axes()
        self.figcanvas = figcanvas

    def draw(self):
        self.figcanvas.draw()

    def clear(self):
        if self.fig.axes:
            for ax in self.fig.axes:
                self.fig.delaxes(ax)

        self.axes = self.set_axes()

    def set_axes(self):
        return self.fig.add_subplot(1,1,1)


class ScrollableFrame(ttk.Frame):

    def __init__(self, container, *args, framewidth=100, **kwargs):
        super().__init__(container, *args, **kwargs)
        canvas = tkinter.Canvas(self, width=framewidth)
        scrollbar = ttk.Scrollbar(self, orient="vertical", command=canvas.yview)
        self.scrollable_frame = ttk.Frame(canvas)

        self.scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(
                scrollregion=canvas.bbox("all")
            )
        )

        canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")

        canvas.configure(yscrollcommand=scrollbar.set)

        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")


class Switch(tkinter.Canvas):
    def __init__(self, *args, command=None, fg='black', bg='gray', width=35,
                     height=18, **kwargs):
        super().__init__(*args, **kwargs)

        self.configure(width=width, height=height, borderwidth=0,
                       highlightthickness=0)

        self.back_ground = self.create_arc((0, 0, 0, 0), start=90,
                                           extent=180, fill=bg, outline='')
        self.back_ground1 = self.create_arc((0, 0, 0, 0), start=270,
                                            extent=180, fill=bg, outline='')
        self.rect = self.create_rectangle(0, 0, 0, 0, fill=bg, outline='')

        self.btn = self.create_oval(0, 0, 0, 0, fill=fg, outline='')

        self.bind('<Configure>', self._resize)
        self.bind('<Button>', self._animate, add='+')
        self.bind('<Button>', command, add='+')

        self.var = tkinter.IntVar(self)
        self.var.set(0)
        self.isdisabled = False
        self.fillcolor = bg


    def _resize(self, event):
        self.coords(self.back_ground, 5, 5, event.height-5, event.height-5)
        self.coords(self.back_ground1, 5, 5, event.height-5, event.height-5)

        factor = event.width-(self.coords(self.back_ground1)[2] -
                              self.coords(self.back_ground1)[0])-10
        self.move(self.back_ground1, factor, 0)

        self.coords(self.rect, self.bbox(self.back_ground)[2]-1, 6,
                    self.bbox(self.back_ground1)[0]+2, event.height-5)

        self.coords(self.btn, 4, 5, event.height-6, event.height-5)

        if self.var.get():
            self.moveto(self.btn, self.coords(self.back_ground1)[0]+1, 4)


    def _animate(self, event):
        if self.isdisabled:
            return

        x, y, w, h = self.coords(self.btn)
        x = int(x-1)
        y = int(y-1)

        if x == self.coords(self.back_ground1)[0]:
            self.toggle("off")

        else:
            self.toggle("on")


    def toggle(self, onoff="on"):
        if onoff == "off":
            self.moveto(self.btn, 4, 4)
            self.var.set(0)
        elif onoff == "on":
            self.moveto(self.btn, self.coords(self.back_ground1)[0]+1, 4)
            self.var.set(1)

    def get_state(self):
        return self.var.get()


    def disable(self, onoff="off"):
        self.toggle(onoff)
        self.isdisabled = True
        self.itemconfigure(self.rect, fill="red")
        self.itemconfigure(self.back_ground, fill="red")
        self.itemconfigure(self.back_ground1, fill="red")


    def enable(self, onoff="off"):
        self.toggle(onoff)
        self.isdisabled = False
        self.itemconfigure(self.rect, fill=self.fillcolor)
        self.itemconfigure(self.back_ground, fill=self.fillcolor)
        self.itemconfigure(self.back_ground1, fill=self.fillcolor)


class ToggleButton(ttk.Frame):

    def __init__(self, *args, command=None, fg='black', bg='gray', width=35,
                 height=18, label1text="Off", label2text="On", **kwargs):
        super().__init__(*args, **kwargs)
        self.switch = Switch(self, command=command, fg=fg, bg=bg, width=width,
                             height=height)

        self.labeloff = tkinter.Label(self, text=label1text)
        self.labelon  = tkinter.Label(self, text=label2text)

        self.labeloff.pack(side="left")
        self.switch.pack(side="left")
        self.labelon.pack(side="left")
        self.var = self.switch.var


    def disable(self, onoff="off"):
        self.switch.disable(onoff)
        if onoff == "off":
            self.labelon.configure(fg="gray")
        elif onoff == "on":
            self.labeloff.configure(fg="gray")


    def enable(self, onoff="off"):
        self.switch.enable(onoff)
        self.labelon.configure(fg="black")
        self.labelon.configure(fg="black")

class ContentTab(ttk.Frame):
    """Leaf node in the NestedNotebook that is a frame widget.
    """

    def __init__(self, frame, canvas, gui, *args, tabselected=None, tabdisabled=None, **kwargs):
        super().__init__(frame, *args, tabselected=tabselected, tabdisabled=tabdisabled, **kwargs)
        self.frame  = frame
        self.canvas = canvas
        self.gui    = gui

class NestedNotebook(ttk.Notebook):
    """Notebook that can has other notebooks or frames in it's tabs.
    """

    def __init__(self, frame, *args, tabselected=None, **kwargs):
        super().__init__(frame, *args, **kwargs)
        self._children = {}
        self._sleep = True

        if tabselected is None:
            self.tabselected = lambda : 0
        else:
            self.tabselected = tabselected

        def eventhandler(event):
            #print(event.widget.tab("current")["text"])
            self.tabchanged()
        self.bind("<<NotebookTabChanged>>", eventhandler)

    def add(self, text, tabselected=None, tab=None):
        if tab == None:
            tab = NestedNotebook(self, tabselected=tabselected)
        self._children[text] = tab
        super().add(tab, text=text)

    def traverse(self, text):
        for txt, frame in self._children.items():
            if txt == text: return frame
        for c in self._children.values():
            if not isinstance(c, NestedNotebook): continue
            frame = c.traverse(text)
            if frame is not None: return frame
        return None

    def currenttab(self):
        tab = self.nametowidget(self.select())
        if isinstance(tab, NestedNotebook):
            return tab.currenttab()
        for name, item in self._children.items():
            if item == tab:
                return tab, name

    def tabchanged(self):
        if self._sleep: return
        tab = self.nametowidget(self.select())
        if isinstance(tab, NestedNotebook):
            tab.tabselected()
            tab.tabchanged()
        else:
            tab.selecttab()

    def wakeup(self):
        self._sleep = False
        for c in self._children.values():
            if not isinstance(c, NestedNotebook): continue
            c.wakeup()

a5py/a5py/gui/components.py
ascot4fusion/ascot5
import tkinter as tk
from tkinter import ttk

from .interactiveorbit import OrbitFrame, OrbitCanvas

class ContentInteractive:

    def __init__(self, gui, settings, canvas):
        self.gui    = gui
        self.canvas = canvas

        runselection = ttk.Notebook(settings)

        runpoincare = ttk.Frame(runselection)
        runorbit    = OrbitFrame(runselection)

        #runselection.add(runpoincare, text="Poincaré")
        runselection.add(runorbit,    text="Orbit")

        canvaspoincare = ttk.Frame(canvas)
        canvasorbit    = OrbitCanvas(canvas)

        def on_tab_change(event):
            tab = event.widget.tab('current')['text']
            self.view(tab)

        runselection.bind('<<NotebookTabChanged>>', on_tab_change)
        runselection.pack(fill="both", expand=True)

        runorbit.init(canvasorbit)
        canvasorbit.init()

        self.runpoincare    = runpoincare
        self.runorbit       = runorbit
        self.canvaspoincare = canvaspoincare
        self.canvasorbit    = canvasorbit

        self.active_canvas = canvasorbit

    def display(self):
        pass
        #self.pack()

    def view(self, tab):
        if tab == "Orbit":
            self.gui.ascot.input_free()
            self.active_canvas.pack_forget()
            self.runorbit.display(self.gui)
            self.active_canvas = self.runorbit.canvas
            self.active_canvas.pack(fill="both", expand=True)

a5py/a5py/gui/contentinteractive.py
ascot4fusion/ascot5
"""Defines ContentManager class for showing stuff at Canvas and Settings frames.
"""
import tkinter as tk
from tkinter import ttk

from a5py import AscotInitException

from .filecontents  import Info
from .inputcontents import Field, Preflight
from .resultcontents import Summary, StateScatter, StateHistogram, Orbit, \
    Poincare, LossSummary, Dists, Moments, WallLoad, Wall3D
from .runcontents import Trace
from .components import NestedNotebook

class ContentManager(NestedNotebook):
    """Manages the contents in SettingsFrame and in CanvasFrame.

    Settings frame is a notebook widget where changing the tab also changes
    what is shown on Canvas*. This is done so that we always pack_forget the
    current frame on Canvas and pack the new one (Note that the frames are
    kept as that way we preserve the settings user has given previously, which
    makes smoother user-experience when different frames are toggled).

    *For group contents the contents on Canvas also depends what is shown in
    the group treeview.

    This is fairly simple class as the actual contents are defined in Content*
    classes that are imported an used here. If those also have subcontents and
    implement a notebook to choose what is display, then those classes mimic the
    structure in this class. In theory, everything could be defined here but
    that would make this class huge...
    """

    def __init__(self, gui, settingsframe, canvasframe):
        """Generate all content widgets (which also generate their contents).
        """
        super().__init__(settingsframe)
        self.pack(fill="both", expand=True)
        self.gui = gui
        self.canvas = canvasframe

        def tabselected(tab):
            if self._sleep: return
            if tab == "Inputs":
                try:
                    self.gui.ascot.simulation_free()
                except AscotInitException:
                    pass
                init = {"wall" : False}
                for inp in ["bfield", "efield", "plasma", "neutral", "boozer",
                            "mhd"]:
                    init[inp] = inp in self.gui.ascot.data
                #msg = self.gui.pleasehold("Ascotpy is being initialized..."
                self.gui.ascot.input_init(**init, switch=True)
            elif tab == "Results":
                try:
                    self.gui.ascot.simulation_free()
                except AscotInitException:
                    pass
                self.gui.ascot.input_init(
                    run=True, bfield=True, efield=True, plasma=True,
                    neutral=True, boozer=True, mhd=True, switch=True)
            elif tab == "Run":
                self.gui.ascot.input_free()
                self.gui.ascot.simulation_initinputs()

        args = (self, self.canvas, self.gui)

        self.add("File", tab=Info(*args))
        self.add("Inputs", tabselected=lambda : tabselected("Inputs"))
        self.traverse("Inputs").add("Plot (R,z)", tab=Field(*args))
        self.traverse("Inputs").add("Preflight", tab=Preflight(*args))
        self.add("Results", tabselected=lambda : tabselected("Results"))
        self.traverse("Results").add("Summary", tab=Summary(*args))
        self.traverse("Results").add("Ini/Endstate")
        self.traverse("Ini/Endstate").add("Scatter", tab=StateScatter(*args))
        self.traverse("Ini/Endstate").add("Histogram", tab=StateHistogram(*args))
        self.traverse("Results").add("Orbits")
        self.traverse("Orbits").add("Trajectory", tab=Orbit(*args))
        self.traverse("Orbits").add("Poincaré", tab=Poincare(*args))
        self.traverse("Results").add("Dists")
        self.traverse("Dists").add("Distribution", tab=Dists(*args))
        self.traverse("Dists").add("Moments", tab=Moments(*args))
        self.traverse("Results").add("Losses")
        self.traverse("Losses").add("Total", tab=LossSummary(*args))
        self.traverse("Losses").add("Wall loads", tab=WallLoad(*args))
        self.traverse("Losses").add("View 3D", tab=Wall3D(*args))
        self.add("Run", tabselected=lambda : tabselected("Run"))
        self.traverse("Run").add("Trace", tab=Trace(*args))

    def update_content(self):
        """Redraw contents on settings and canvas frames.
        """
        tab, name = self.currenttab()
        if name == "File":
            tab.selecttab()

    def restart(self):
        # Free any used resources
        try:
            self.gui.ascot.simulation_free()
        except:
            self.gui.ascot.input_free()

        # Has results?
        if self.gui.ascot.data.active is None:
            self.tab(2, state="disabled")
        else:
            self.tab(2, state="normal")

            run = self.gui.ascot.data.active
            tab0 = self.traverse("Results")

            # Has orbit data?
            if not hasattr(run, "_orbit"):
                tab0.tab(2, state="disabled")
            else:
                tab0.tab(2, state="normal")
                tab1 = self.traverse("Orbits")

                # Has Poincare data?
                if run.getorbit_poincareplanes() is None:
                    tab1.tab(1, state="disabled")
                else:
                    tab1.tab(1, state="normal")

            # Has distribution data?
            if not ( hasattr(run, "_dist5d") or hasattr(run, "_distrho5d") or
                     hasattr(run, "_dist6d") or hasattr(run, "_distrho6d") or
                     hasattr(run, "_distcom") ):
                tab0.tab(3, state="disabled")
            else:
                tab0.tab(3, state="normal")
                tab1 = self.traverse("Dists")

                # Has 5D distribution data to calculate moments?
                if not( hasattr(run, "_dist5d") or hasattr(run, "_distrho5d") ):
                    tab1.tab(1, state="disabled")
                else:
                    tab1.tab(1, state="normal")

            # Has 3D wall?
            tab1 = self.traverse("Losses")
            if run.wall.get_type() != "wall_3D":
                tab1.tab(1, state="disabled")
                tab1.tab(2, state="disabled")
            else:
                tab1.tab(1, state="normal")
                tab1.tab(2, state="normal")

a5py/a5py/gui/contentmanager.py
ascot4fusion/ascot5
import os
import tkinter as tk
import tkinter.ttk as ttk

import numpy as np

from a5py.ascot5io.options import Opt
from .components import ContentTab, PlotFrame, ScrollableFrame, DropdownMenu

class Info(ContentTab):

    def __init__(self, frame, canvas, gui, *args, **kwargs):
        super().__init__(frame, canvas, gui, *args, **kwargs)
        emptycanvas   = ttk.Frame(self.canvas)
        optionscanvas = Info.OptionsCanvas(self.canvas)

        self.logo = Info.Logo(self)
        self.logo.grid(column=0, row=0, sticky="WE")

        self.desc = Info.Desc(self)
        self.desc.grid(column=0, row=1, sticky="WE")

        self.opt = Info.Options(self, optionscanvas, gui)
        self.opt.grid(column=0, row=2, sticky="WE")

        self.canvas.slideadd(self.logo, emptycanvas)
        self.canvas.slideadd(self.opt, optionscanvas)

        self.logo.grid_forget()
        self.desc.grid_forget()
        self.opt.grid_forget()

    def selecttab(self):
        tree   = self.gui.groups.tree
        qid    = tree.item(tree.selection(), "text")
        parent = tree.item(tree.parent(tree.selection()), "text")

        self.logo.grid_forget()
        self.desc.grid_forget()
        self.opt.grid_forget()

        if parent == "" and qid == "":
            # Show logo if no group is selected (on startup/opening new file)
            self.logo.grid()
            self.canvas.slideshow(self.logo)
            return

        try:
            int(qid)
            if parent == "results":
                group = self.gui.ascot.data["q"+qid]
            else:
                group = self.gui.ascot.data[parent]["q"+qid]
        except ValueError:
            # For parent groups there is nothing to display.
            self.canvas.slideshow(self.logo)
            return

        # Show group description and related information
        self.desc.grid()
        self.desc.view(group)

        if parent == "options":
            self.canvas.slideshow(self.opt)
            self.opt.grid()
            self.opt.view(group)
            return

    class Logo(tk.Canvas):

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            logo = os.path.join(
                os.path.dirname(__file__), "../../docs/figs/logo.png")
            logo = tk.PhotoImage(file=logo)
            self.logo = logo # Otherwise garbage collector eats this
            self.create_image(150, 50, image=logo)

    class Desc(ttk.Frame):

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

            # At top we always show box with the description and buttons
            # to save (or revert) the text.
            f1 = ttk.Frame(self)
            ttk.Label(f1, text="Description:").pack(side="left")
            f1.grid(row=0, column=0, sticky="ew")

            f2 = ttk.Frame(self)
            button_descsave = ttk.Button(f2, text="Save")
            button_descsave.pack(side="right")

            button_descrevert = ttk.Button(f2, text="Revert")
            button_descrevert.pack(side="right")

            f2.grid(row=0, column=1, sticky="ew")

            self.button_descsave   = button_descsave
            self.button_descrevert = button_descrevert

            descbox = tk.Text(self, height=5, width=50)
            descbox.grid(row=1, column=0, columnspan=2, sticky="nsew",
                         padx=2, pady=2)
            self.descbox = descbox

        def revert(self, group):
            """Revert description.
            """
            text = group.get_desc()
            self.descbox.delete("1.0", "end")
            self.descbox.insert("end", text)

        def save(self, group):
            """Save description.
            """
            group.set_desc(self.descbox.get("1.0", "end-1c"))

        def view(self, group):
            self.revert(group)
            self.button_descsave.configure(
                command=lambda:self.save(group))
            self.button_descrevert.configure(
                command=lambda:self.revert(group))

    class Options(ttk.Frame):

        def __init__(self, frame, canvas, gui, *args, **kwargs):
            super().__init__(frame, *args, **kwargs)
            self.gui    = gui
            self.canvas = canvas

            # On frame we have save/revert buttons
            tk.Label(self, text="Options:").pack(side="left")

            self.savebutton = tk.Button(self, text="Save")
            self.savebutton.pack(side="left")

            self.revertbutton = tk.Button(self, text="Revert")
            self.revertbutton.pack(side="left")

        def view(self, group):
            self.savebutton.configure(
                command=lambda:self.canvas.write(group, self.gui))
            self.revertbutton.configure(
                command=lambda:self.canvas.view(group))
            self.canvas.view(group)

    class OptionsCanvas(ttk.Frame):

        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            textbox = tk.Text(self, borderwidth=3, relief="sunken")
            textbox.pack(side="left", fill="both", expand=True)

            scrollbar = ttk.Scrollbar(self, command=textbox.yview)
            textbox["yscrollcommand"] = scrollbar.set
            scrollbar.pack(side="left", fill="y")

            # Styles for highlights
            textbox.tag_configure("HEADER", foreground="#004999")
            textbox.tag_configure("DESC",   foreground="#26abff")
            textbox.tag_configure("VALUE",  foreground="black")

            self.text = textbox
            self.text.delete("1.0", tk.END)

        def highlight(self):
            """Highlight values and descriptions.
            """
            lines = self.text.get('1.0', "end").splitlines()
            for i,line in enumerate(lines):
                if len(line) < 2:
                    pass
                elif line[0:2] == "#*":
                    self.text.tag_add("HEADER",str(i+1)+".0",
                                      str(i+1)+".end")
                elif line[0] == "#":
                    self.text.tag_add("DESC",str(i+1)+".0",
                                      str(i+1)+".end")
                else:
                    self.text.tag_add("VALUE",str(i+1)+".0",
                                      str(i+1)+".end")

        def view(self, group):
            """Read and display options in textbox.
            """
            opttext = group.tostring()
            self.text.delete("1.0", tk.END)
            self.text.insert("end", opttext)
            self.highlight()

        def write(self, group, gui):
            """Turn options string to a dict and write it.
            """
            lines = self.text.get("1.0", "end-1c").splitlines()
            opt = Opt.convert_string(lines)

            desc = group.get_desc()
            new = gui.ascot.data.create_input("opt", **opt, desc=desc,
                                              activate=True)
            gui.groups.add_group("options", gui.ascot.data._get_group(new))

a5py/a5py/gui/filecontents.py
ascot4fusion/ascot5
"""Main window for the GUI.
"""
import os
import sys
import subprocess
import threading
import tkinter as tk
from tkinter import ttk

from tkinter.filedialog import askopenfilename, askdirectory
from tkinter import messagebox, simpledialog

from a5py import Ascot, AscotIOException
from a5py.ascot5io.coreio import fileapi

from .contentmanager import ContentManager
from .guiparams import GUIparams

class GUI(tk.Tk):
    """Window where all GUI elements are stored.


    The GUI looks like following:
     ____________________________________________
    |  FileFrame   |                             |
    |______________|                             |
    |  GroupFrame  ^                             |
    | (Scrollable) |        CanvasFrame          |
    | (Resizable)  |        (Resizable)          |
    |              |                             |
    |              |                             |
    |              |                             |
    |______________v                             |
    |              ^                             |
    | SettingsFrame|                             |
    | (Scrollable) |                             |
    | (Resizable)  |                             |
    |______________v_____________________________|

    FileFrame has a pop-up window for choosing a file to be opened.

    GroupFrame also has an pop-up window for operations related to
    HDF5 file contents (like activating or removing a group).

    FileFrame and GroupFrame are always displayed. Contents of
    SettingsFrame and CanvasFrame depends on what is clicked on the
    FileFrame/GroupFrame/SettingsFrame and those are managed by
    the ContentManager object.

    Access to the actual data is handled via an Ascot object. This
    GUI should just pass the arguments to the plotting functions
    implemented elsewhere instead of implementing those itself.
    """

    # Minimum GUI size in pixels
    MINWIDTH  = 800
    MINHEIGHT = 600

    # FileFrame size (fixed) in pixels. Width is also used in other frames in
    # the same column.
    FILEFRAMEWIDTH  = 350
    FILEFRAMEHEIGHT = 50

    # Other frame heights
    INPUTFRAMEHEIGHT    = 300
    SETTINGSFRAMEHEIGHT = 250

    # Screenwidth method won't regonize dual monitors but instead returns their
    # combined width. If the screenwidth is above this limit (in pixels), assume
    # dual monitors are in use.
    DUALMONITORWIDTH = 2000

    BORDERWIDTH = 2


    def __init__(self, filename=None):
        tk.Tk.__init__(self)

        ## Add decorations: title and icons ##
        self.title("ASCOT5 GUI")
        icon = os.path.join(os.path.dirname(__file__),
                            "../../docs/figs/icon.png")
        self.tk.call("wm", "iconphoto", self._w, tk.PhotoImage(file=icon))

        ## Set window size and minimum size ##
        sw = self.winfo_screenwidth()
        sh = self.winfo_screenheight()

        # If the screen width is very large it might be due to dual
        # monitors. Divide the screen width by 2 to make GUI look nice
        # with dual monitors.
        if sw > GUI.DUALMONITORWIDTH:
            sw /= 2

        # GUI dimensions and locations
        w = sw*3/4
        h = sh*3/4
        x = (sw/2) - (w/2)
        y = (sh/2) - (h/2)

        # Make sure dimensions are below minimum values
        if w < GUI.MINWIDTH:
            w = GUI.MINWIDTH
        if h < GUI.MINHEIGHT:
            h = GUI.MINHEIGHT

        # Set window size and set minimum size
        self.geometry('%dx%d+%d+%d' % (w, h, x, y))
        self.minsize(GUI.MINWIDTH, GUI.MINHEIGHT)

        # Make sure everything is closed when the window is closed
        def close():
            """Close the gui and terminate the program.
            """
            self.destroy()
            exit()
        self.protocol("WM_DELETE_WINDOW", close)

        # Initialize content frames
        files    = FileFrame(self,     width=GUI.FILEFRAMEWIDTH,
                                       height=GUI.FILEFRAMEHEIGHT,
                                       borderwidth=GUI.BORDERWIDTH)
        groups   = GroupFrame(self,    width=GUI.FILEFRAMEWIDTH,
                                       height=GUI.INPUTFRAMEHEIGHT,
                                       borderwidth=GUI.BORDERWIDTH)
        settings = SettingsFrame(self, width=GUI.FILEFRAMEWIDTH,
                                       height=GUI.SETTINGSFRAMEHEIGHT,
                                       borderwidth=GUI.BORDERWIDTH)
        canvas   = CanvasFrame(self,   width=GUI.MINWIDTH-GUI.FILEFRAMEWIDTH,
                                       height=GUI.MINHEIGHT-GUI.FILEFRAMEHEIGHT)

        files.grid(   row=0, column=0, sticky="NWES")
        groups.grid(  row=1, column=0, sticky="NWES")
        settings.grid(row=2, column=0, sticky="NWES")
        canvas.grid(  row=0, column=1, sticky="NWSE", rowspan=3)

        settings.pack_propagate(False)

        self.grid_rowconfigure(1, weight=1)
        self.grid_rowconfigure(2, weight=1)
        self.grid_columnconfigure(1, weight=1)

        self.files  = files
        self.groups = groups

        # Set up content manager
        self.params  = GUIparams()
        self.content = ContentManager(self, settings, canvas)

        # Read file and show its contents
        self.ascot = Ascot()
        path = None if filename is None else os.path.abspath(filename)
        self.files.open_new_file(path)

        def wakecontentmanager(event):
            self.content.wakeup()
            self.unbind('<Visibility>')

        self.bind('<Visibility>', wakecontentmanager)

    def pleasehold(self, message, target):
        """Pops up a message and freezes GUI until the target function is done.

        This method can be called e.g. when ascotpy is being initialized
        to let the user know that it is happening.
        """

        # Dialog at the center of the main window
        wd = 300
        wh = 100
        dialog = tk.Toplevel(width=wd, height=wh)
        dialog.title("Please hold")
        x = self.winfo_x()
        y = self.winfo_y()
        w = self.winfo_width()
        h = self.winfo_height()
        dialog.geometry("+%d+%d" % (x + int(w/2-wd/2), y + int(h/2-wh/2)))
        f = tk.Frame(dialog, width=wd, height=wh)
        f.pack_propagate(0)
        tk.Label(dialog, text=message,
                 font=("Calibri",12)).pack(anchor="c", fill="both", expand=True)
        f.pack()

        # Show dialog
        dialog.update_idletasks()

        # Make it unclosable
        def do_nothing():
            pass
        dialog.protocol("WM_DELETE_WINDOW", do_nothing)

        # Focus on the dialog making main GUI unresponsive
        dialog.grab_set()

        # Do the thing
        target()

        # Flush all mouse clicks etc
        dialog.update()

        # Destroy dialog making main GUI responsive again
        dialog.destroy()

class SettingsFrame(ttk.Frame):
    """Frame that displays parameters the user can vary when viewing data.

    Contents of this frame changes depending on what tab is chosen.
    """
    pass

class CanvasFrame(ttk.Frame):
    """Frame that displays plots.

    Contents of this frame changes depending on what tab is chosen.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.slides = {}
        self.current = ttk.Frame(self)
        self.current.pack()

    def slideadd(self, widget, frame):
        self.slides[widget.winfo_name()] = frame

    def slideshow(self, widget):
        if self.current != widget:
            self.current.pack_forget()
            self.current = self.slides[widget.winfo_name()]
            self.current.pack(fill="both", expand=True)

class FileFrame(ttk.Frame):
    """Frame for accessing the HDF5 file.
    """

    def __init__(self, gui, *args, **kwargs):
        tk.Frame.__init__(self, gui, *args, **kwargs)
        self.gui = gui

        frame1 = tk.Frame(self)
        frame2 = tk.Frame(self)
        frame3 = tk.Frame(self, width=kwargs["width"]-5, height=100)

        # Label "File: XXX"
        tk.Label(frame1, text="File: ").pack(side="left", anchor="nw")

        # Textbox for showing the current file
        self.filenamefield = tk.Text(frame1, width=30, height=1)
        self.filenamefield.pack(side="left", anchor="nw")

        # Browse button. We need to do this frame trick to set it correct size
        tempframe = tk.Frame(frame1, width=24, height=20)
        tempframe.pack_propagate(False)
        self.browsebutton = tk.Button(tempframe, text="...")
        self.browsebutton.pack(side="left")
        tempframe.pack(side="left", anchor="nw")

        # Label "Size: XXX Mb"
        self.filesizelabel = tk.Label(frame2, text="Size: ")
        self.filesizelabel.pack(side="left", anchor="nw")

        frame1.pack(side="top", fill="x",    anchor="nw")
        frame2.pack(side="top", fill="x",    anchor="nw")

        ## Set functionality ##

        # Browse button opens filename dialog when clicked
        self.browsebutton.configure(
            command=self.browse_file
        )

        # Unmutable filename field
        self.filenamefield.configure(state="disabled", wrap="none")

    def browse_file(self):
        """Open dialog for choosing HDF5 file and open it.

        If a valid file is chosen, deliver the filename to GUI and refresh.
        """
        fn = askopenfilename( title="Select ASCOT5 HDF5 file",
                              filetypes = [("HDF5 files","*.h5")] )
        if len(fn) == 0:
            # Do nothing
            return

        # Get filename and refresh
        fn = os.path.abspath(fn)
        self.open_new_file(fn)

    def open_new_file(self, filename):
        """Update the contents of the fileframe if file has changed.
        """
        if filename is not None:

            try:
                self.gui.ascot.file_load(filename)
            except AscotInitException:
                messagebox.showerror(
                    "Error",
                    "Could not open file. The file could be corrupted.\n" +
                    "Try Ascot(fn.h5) in Python or h5ls fn.h5 in terminal.")
                return

            except AscotpyInitException:
                messagebox.showwarning(
                    "Warning",
                    "Could not initialize Ascotpy.\n" +
                    "Check that libascot is compiled and included in " +
                    "LD_LIBRARY_PATH")
                return

        self.gui.groups.build_groups()
        self.gui.params.retrieve(filename)
        self.display_file(filename)
        self.gui.content.restart()

    def display_file(self, filename):
        """Display file name and size.
        """
        if filename is None:
            filename = ""
            size = 0.0
        else:
            size = os.path.getsize(filename) / 1e6 # B -> MB

        # Update filename
        self.filenamefield.configure(state="normal")
        self.filenamefield.delete("1.0", tk.END)
        self.filenamefield.insert("end", filename)
        self.filenamefield.configure(state="disabled")
        self.filenamefield.see(tk.END)

        # Update size
        self.filesizelabel.configure(
            text="Size: " + "{0:.3f}".format(size) + " MB")

class GroupFrame(ttk.Frame):
    """Frame for displaying groups in a treeview.

    Also handles any operations regarding contents of the file.
    """

    def __init__(self, gui, *args, **kwargs):
        """Initializes tree which doesn't contain any data.
        """
        super().__init__(*args, **kwargs)

        # Set white background, header font, and default font for tree entries.
        style = ttk.Style()
        style.configure("Treeview.Heading", font=("Calibri", 10),
                        background="white")
        style.map('Treeview.Heading', background=[('focus', 'white')])
        style.configure("Treeview", font=("Calibri", 9),
                        background="white", fieldbackground="white", indent=1)

        # Tree where only one item can be selected at a time
        tree = ttk.Treeview(self, style="mystyle.Treeview", selectmode="browse")

        # Three columns: Group, Type, and Date
        tree["columns"]=("#1","#2")
        tree.column("#0", width= 90, minwidth= 90, stretch=False)
        tree.column("#1", width= 85, minwidth= 85, stretch=False)
        tree.column("#2", width=130, minwidth=130, stretch=True)
        tree.heading("#0", text="Group")
        tree.heading("#1", text="Type")
        tree.heading("#2", text="Date")

        # Styles for different entries:
        # - Every odd entry has a white background and every even a darker one
        # - Parents have white background with bold and larger font.
        # - If input parent don't exist in the file the font is set to red
        # - Active group has a green bolded medium font.
        tree.tag_configure("even",    background="#EFEFEF")
        tree.tag_configure("odd",     background="white")
        tree.tag_configure("nodata",  foreground="red")
        tree.tag_configure("parent",  background="white",
                           font=("Calibri", 10, "bold"))
        tree.tag_configure("active",  foreground="green",
                           font=("Calibri", 9, "bold"))

        # Prevents user from resizing columns
        tree.bind('<Motion>', 'break')
        tree.pack(side=tk.TOP, expand=True, fill="both")

        # Set actions
        tree.bind("<<TreeviewSelect>>", self.select_group)

        # Initialize popupmenu
        groupmenu = self.GroupMenu(gui, self, tree)

        self.tree = tree
        self.gui  = gui

    def build_groups(self):
        """(Re)Build tree from scratch from the file in GUI.
        """

        # Clear tree.
        self.tree.delete(*self.tree.get_children())

        # Check if we have an open HDF5 file and proceed.
        if self.gui.ascot.file_getpath() is None:
            return

        # Construct the tree by creating all parents and their children
        for index, parent in enumerate(fileapi.INPUTGROUPS + ["results"]):

            if parent == "results":
                qids, types, descs, dates = \
                    self.gui.ascot.data.get_contents()
                if len(qids) > 0:
                    activeqid = self.gui.ascot.data.active.get_qid()
                    item = self.tree.insert("", index, text=parent,
                                            tags=("parent"))
            else:
                if not parent in self.gui.ascot.data:
                    # The parent does not exist
                    item = self.tree.insert("", index, text=parent,
                                            tags=("parent", "nodata"))
                    continue

                item = self.tree.insert("", index, text=parent, tags=("parent"))
                qids, dates, descs, types = \
                    self.gui.ascot.data[parent].get_contents()
                if len(qids) > 0:
                    activeqid = self.gui.ascot.data[parent].active.get_qid()

            # Generate children:
            # - Rows are tagged "even" or "odd"
            # - Active groups are tagged "active"
            alternatingrowcolor = 0
            for i in range(len(qids)):
                tags = []
                tags.append("odd" if alternatingrowcolor % 2 else "even")
                alternatingrowcolor += 1
                if qids[i] == activeqid:
                    tags.append("active")

                self.tree.insert(item, 0, text=qids[i], tags=tuple(tags),
                                 values=(types[i], dates[i]))

    def select_group(self, event):
        """User selected a group.

        Update tree view accordingly and notify GUI.
        """
        qid    = self.tree.item(self.tree.selection(), "text")
        parent = self.tree.item(self.tree.parent(self.tree.selection()), "text")

        # Test if a valid QID or a parent group.
        try:
            int(qid)
        except:
            parent = qid
            qid    = None

        # If selection is a run, highlight the inputs it has used. If selection
        # is a input, highlight the runs that have used it.
        if qid is not None:
            if parent == "results":
                inputqids = []
                for p in fileapi.INPUTGROUPS:
                    if p in self.gui.ascot.data["q"+qid]:
                        inputqids.append(
                            self.gui.ascot.data["q"+qid][p].get_qid())

                self._highlightinputs(inputqids)

            else:
                outputqids = self.gui.ascot.data.get_runs(qid)
                self._highlightoutputs(outputqids)

        # Show contents for this group
        self.gui.content.update_content()

    def activate_group(self):
        """Set currently selected item (group) as active in the tree.
        """
        item       = self.tree.selection()
        parent     = self.tree.parent(item)
        qid        = self.tree.item(item, "text")
        parentname = self.tree.item(parent, "text")

        if parentname == "results":
            self.gui.ascot.data["q"+qid].activate()
        else:
            self.gui.ascot.data[parentname]["q"+qid].activate()
        self._activate(item, parent)
        self.gui.content.update_content()

    def remove_group(self):
        """Remove currently selected item (group) from tree.
        """
        item       = self.tree.selection()
        parent     = self.tree.parent(item)
        parentname = self.tree.item(parent, "text")
        qid        = self.tree.item(item, "text")

        # Really remove?
        mbox = messagebox.askquestion (
            "Remove group",
            "Are you sure you want to remove " + qid
            + " permanently from the HDF5 file?",
            icon = "warning")

        if mbox != "yes":
            return

        # Remove the group in HDF5
        try:
            self.gui.ascot.data._destroy_group(qid)
        except AscotIOException:
            messagebox.showerror(
                  "Error",
                  "Could not remove the input group as it has been used in\n"
                + "a simulation. Remove the simulation group first.")
            return

        # Update treeview

        # 1. Whole group was removed
        if qid in fileapi.INPUTGROUPS or qid == "results":
            for c in self.tree.get_children(item):
                self.tree.delete(c)

            if qid == "results":
                self.tree.delete(item)
                self._highlightinputs([])
            else:
                self.tree.item(item, tags=("parent", "nodata"))
            return

        # 2. Ordinary group: remove it and set the next one as active
        #    or remove the parent as well
        self.tree.delete(item)
        if parentname == "results":
            #if not hasattr(self.gui.ascot.data, "active"):
            if self.gui.ascot.data.active is None:
                # Last group
                self.tree.delete(parent)
                self._highlightinputs([])
            else:
                activeqid = self.gui.ascot.data.active.get_qid()
                self._highlightinputs([])
                for group in self.tree.get_children(parent):
                    if activeqid == self.tree.item(group, "text"):
                        self._activate(group, parent)
                        break
        else:
            if not hasattr(self.gui.ascot.data, parentname):
                # Last group
                self.tree.item(parent, tags=("parent", "nodata"))
            else:
                activeqid = self.gui.ascot.data[parentname].active.get_qid()
                for group in self.tree.get_children(parent):
                    if activeqid == self.tree.item(group, "text"):
                        self._activate(group, parent)
                        break
        self.gui.content.update_content()

    def export_group(self):
        """
        Export given group to another HDF5 file.

        Open a prompt for querying destination file.

        This operation requires no GUI re-initialization.
        """
        tree  = self.tree
        item  = tree.selection()
        group = tree.item(tree.selection(), "text")
        fnout = askopenfilename( title="Select ASCOT5 HDF5 file",
                                 filetypes = [("HDF5 files","*.h5")] )
        if len(fnout) == 0:
            # Do nothing
            return

        # Check that the target file is valid.
        try:
            Ascot(fnout)
        except:
            messagebox.showerror(
                    "Error",
                    "Could not open the file.")
            return

        try:
            tools.copygroup(self.gui.ascot.file_getpath(), fnout, group)

        except Exception as err:
            messagebox.showerror(
                    "Error",
                    "Could not copy the group as something went wrong.")
            return

    def add_group(self, parent, group):
        """
        Add new group to the tree.

        This function is called when the user has created a new group, so we can
        assume that the group is newest, belongs to the top of the list, and
        is the active group.
        """

        # Create / unflag parent if this will be the first child
        if parent == "results":
            parent = None
            for c in self.tree.get_children():
                if self.tree.item(c, "text") == "results":
                    parent = c

            if parent is None:
                parent = self.tree.insert("", "end", text="results",
                                          tags=("parent"))

        else:
            for c in self.tree.get_children():
                if self.tree.item(c, "text") == parent:
                    parent = c
                    if "nodata" in self.tree.item(parent, "tags"):
                        tags = list(self.tree.item(parent, "tags"))
                        tags.remove("nodata")
                        self.tree.item(parent, tags=tuple(tags))

        # Insert new group and set it active
        self.tree.insert(parent, 0, text=group.get_qid(),
                         values=(group.get_type(), group.get_date()))
        item = self.tree.get_children(parent)[0]
        self._activate(item, parent)

        # Recolor alternating rows
        alternatingrowcolor = 0
        for c in self.tree.get_children(parent):
            tags = list(self.tree.item(c, "tags"))

            if "odd" in tags:
                tags.remove("odd")
            if "even" in tags:
                tags.remove("even")

            tags.append("odd" if alternatingrowcolor % 2 else "even")
            alternatingrowcolor += 1

            self.tree.item(c, tags=tuple(tags))

        # Update contents
        self.gui.content.update_content()

    def add_dummygroup(self):
        """
        Adds dummy group for selected parent.
        """
        parent = self.tree.item(self.tree.selection(), "text")
        self.gui.ascot.data.add_dummyinputs(parent=parent, missing=False)

        # New group is the newest group (qids in Ascot are ordered by date)
        if parent == "results":
            tmp = self.gui.ascot.data
        else:
            tmp = self.gui.ascot.data[parent]

        qid = tmp._qids[-1]
        group = tmp[qid]
        group.activate()
        self.add_group(parent, group)

    def _highlightinputs(self, inputqids):
        """Highlight inputs used in the selected run.
        """
        children = self.tree.get_children()
        for c in children:
            if self.tree.item(c, "text") == "results":
                continue

            groups = self.tree.get_children(c)
            for g in groups:
                qid = self.tree.item(g, "text")
                vals = self.tree.item(g, "values")
                date = vals[1]
                if date[0] == "*":
                    date = date[1:-1]

                if qid in inputqids:
                    date = "*" + date + "*"

                self.tree.item(g, values=(vals[0], date))

    def _highlightoutputs(self, outputqids):
        """Highlight runs that have used the selected input.
        """
        children = self.tree.get_children()
        for c in children:
            if self.tree.item(c, "text") != "results":
                continue

            groups = self.tree.get_children(c)
            for g in groups:
                qid = self.tree.item(g, "text")
                vals = self.tree.item(g, "values")
                date = vals[1]
                if date[0] == "*":
                    date = date[1:-1]

                if qid in outputqids:
                    date = "*" + date + "*"

                self.tree.item(g, values=(vals[0], date))

    def _activate(self, item, parent):
        """Move active tag from the previous active to the current one.
        """
        for c in self.tree.get_children(parent):
            tags = list(self.tree.item(c, "tags"))

            if "active" in tags:
                tags.remove("active")
                self.tree.item(c, tags=tuple(tags))
            if c == item[0] or c == item:
                tags.append("active")
                self.tree.item(c, tags=tuple(tags))

    class GroupMenu(tk.Menu):
        """Popupmenu displayed when groups treeview is right-clicked.
        """

        def __init__(self, gui, groups, tree, *args, **kwargs):
            super().__init__(gui, *args, tearoff=0, **kwargs)

            self.add_command(label="Activate",
                             command=groups.activate_group)
            self.add_command(label="Remove",
                             command=groups.remove_group)
            #self.add_command(label="Export",
            #                 command=groups.export_group)
            #self.add_command(label="Add dummy input",
            #                 command=groups.add_dummygroup)

            # Show menu when tree is right-clicked
            tree.bind("<Button-3>", self.showmenu)

            # Hide menu when focus is lost or user clicks anywhere else
            self.bind("<FocusOut>", self.hidemenu)
            gui.bind("<Button-1>", self.hidemenu)

            self.groups = groups

        def showmenu(self, event):
            """
            Show menu if a group was right-clicked and select the group.
            """
            item       = self.groups.tree.identify("item", event.x, event.y)
            itemname   = self.groups.tree.item(item, "text")
            parent     = self.groups.tree.parent(item)
            parentname = self.groups.tree.item(parent, "text")
            if len(itemname) > 0:
                self.post(event.x_root, event.y_root)
                self.focus_set()

                # Disable "activate" option if selected item is a parent
                self.entryconfigure(0, state="normal")
                if itemname in fileapi.INPUTGROUPS + ["results"]:
                    self.entryconfigure(0, state="disabled")

                # Disable "remove" option if parent has no children
                self.entryconfigure(1, state="normal")
                if itemname in fileapi.INPUTGROUPS and \
                   len(self.groups.tree.get_children(item)) == 0:
                    self.entryconfigure(1, state="disabled")

                # Disable "export" option for parents and runs
                self.entryconfigure(2, state="normal")
                if itemname in fileapi.INPUTGROUPS + ["results"] \
                   or parentname == "results":
                    self.entryconfigure(2, state="disabled")

                # This will launch a event that will update the GUI.
                self.groups.tree.selection_set(item)

        def hidemenu(self, event):
            """Hide menu.
            """
            self.unpost()

a5py/a5py/gui/gui.py
ascot4fusion/ascot5
"""Store and retrieve GUI related parameters to/from HDF5
"""
import numpy as np
import h5py

class GUIparams():

    def __init__(self):
        self.params = {}

    def add(self, **kwargs):
        """Add parameters.
        """
        for k in kwargs.keys():
            self.params[k] = kwargs[k]

    def store(self, hdf5, params):
        """Stores parameter(s) to HDF5.

        Only stores the ones whose name is on the given list "params".
        """

        with h5py.File(hdf5, "a") as h5:
            if "gui" not in h5.keys():
                h5.create_group("gui")

            group = h5["gui"]
            def store_dataset(name, dtype):
                """For storing datasets
                """
                if name not in params:
                    return

                if name not in group.keys():
                    group.create_dataset(name, data=self.params[name].get(),
                                         dtype=dtype)
                else:
                    group[name][()] = self.params[name].get()

            S = h5py.string_dtype(encoding='utf-8', length=None)
            store_dataset("input_rzplot_minr", "f8")
            store_dataset("input_rzplot_maxr", "f8")
            store_dataset("input_rzplot_numr", "i8")
            store_dataset("input_rzplot_minz", "f8")
            store_dataset("input_rzplot_maxz", "f8")
            store_dataset("input_rzplot_numz", "i8")
            store_dataset("input_rzplot_qnt",  S)

    def retrieve(self, hdf5):
        """Retrieves parameter(s) from HDF5.
        """
        S = h5py.string_dtype(encoding='utf-8', length=None)
        try:
            with h5py.File(hdf5, "r") as h5:
                if "gui" in h5.keys():
                    for k in h5["gui"]:
                        if k in self.params.keys():
                            val = h5["gui"][k]
                            if val.dtype == S:
                                self.params[k].set(val.asstr()[()])
                            else:
                                self.params[k].set(val[()])

        except:
            pass

a5py/a5py/gui/guiparams.py
ascot4fusion/ascot5
import tkinter as tk
from tkinter import ttk

import numpy as np
import unyt

from .components import ContentTab, PlotFrame, NumEntry

class Field(ContentTab):

    def __init__(self, frame, canvas, gui):
        super().__init__(frame, canvas, gui)

        plotrz = PlotFrame(canvas)
        self.canvas.slideadd(self, plotrz)

        f1 = ttk.Frame(self)
        f2 = ttk.Frame(self)

        f1.pack(anchor="nw", side="left")
        f2.pack(anchor="ne", side="left", fill="x", expand=True)

        xmin_entry = NumEntry(f1, labeltext="R [m]     = ", entrywidth=5,
                              labelwidth=12, anchor="w", defval=0.1)
        xmax_entry = NumEntry(f1, labeltext="–",    entrywidth=5,
                              labelwidth=2, anchor="c", defval=10.0)
        xnum_entry = NumEntry(f1, labeltext="x",    entrywidth=5,
                              labelwidth=2, anchor="c", defval=50,
                              isint=True)

        ymin_entry = NumEntry(f1, labeltext="z [m]     = ", entrywidth=5,
                              labelwidth=12, anchor="w", defval=-8.0)
        ymax_entry = NumEntry(f1, labeltext="–",    entrywidth=5,
                              labelwidth=2, anchor="c", defval=8.0)
        ynum_entry = NumEntry(f1, labeltext="x",    entrywidth=5,
                              labelwidth=2, anchor="c", defval=50,
                              isint=True)

        xmin_entry.grid(row=0, column=0)
        xmax_entry.grid(row=0, column=1)
        xnum_entry.grid(row=0, column=2)

        ymin_entry.grid(row=1, column=0)
        ymax_entry.grid(row=1, column=1)
        ynum_entry.grid(row=1, column=2)

        phi_entry = NumEntry(f1, labeltext="phi [deg] = ", entrywidth=5,
                             labelwidth=12, anchor="w", defval=0)
        phi_entry.grid(row=2, column=0)

        time_entry = NumEntry(f1, labeltext="Time [s]  = ", entrywidth=5,
                              labelwidth=12, anchor="w", defval=0)
        time_entry.grid(row=3, column=0)

        cmin_entry = NumEntry(f1, labeltext="Color     = ", entrywidth=5,
                              labelwidth=12, anchor="w", defval="")
        cmax_entry = NumEntry(f1, labeltext="-", entrywidth=5,
                              labelwidth=2, anchor="c", defval="")
        cmin_entry.grid(row=4, column=0)
        cmax_entry.grid(row=4, column=1)

        qchoice = tk.StringVar(f1)
        qinput = ttk.Combobox(f1, width=12, textvariable=qchoice,
                                      state="readonly")
        qinput.grid(row=5,column=0, sticky="W")

        plotbutton = tk.Button(f2, text="Plot", width=3)
        plotbutton.pack(anchor="e")
        savebutton = tk.Button(f2, text="Store", width=3)
        savebutton.pack(anchor="e")

        # Store parameters
        gui.params.add(
            input_rzplot_minr=xmin_entry.choice,
            input_rzplot_maxr=xmax_entry.choice,
            input_rzplot_numr=xnum_entry.choice,
            input_rzplot_minz=ymin_entry.choice,
            input_rzplot_maxz=ymax_entry.choice,
            input_rzplot_numz=ynum_entry.choice,
            input_rzplot_qnt=qchoice
        )

        def storesettings():
            gui.params.store(
                gui.ascot.data,
                ["input_rzplot_minr", "input_rzplot_maxr", "input_rzplot_numr",
                 "input_rzplot_minz", "input_rzplot_maxz", "input_rzplot_numz",
                 "input_rzplot_qnt"] )

        def plot():
            r = np.linspace( xmin_entry.getval() ,
                             xmax_entry.getval() ,
                             xnum_entry.getval() )

            z = np.linspace( ymin_entry.getval() ,
                             ymax_entry.getval() ,
                             ynum_entry.getval() )

            r *= unyt.m
            z *= unyt.m
            phi  = phi_entry.getval()  * unyt.deg
            time = time_entry.getval() * unyt.s

            clim = [None, None]
            if not cmin_entry.isempty():
                clim[0] = cmin_entry.getval()
            if not cmax_entry.isempty():
                clim[1] = cmax_entry.getval()
                qnt  = qchoice.get()

            qnt = self.qchoice.get()

            plotrz.clear()
            self.gui.ascot.input_plotrz(r, z, qnt, phi=phi, t=time, clim=clim,
                                        axes=plotrz.axes)
            self.gui.ascot.input_plotseparatrix(phi, time, axes=plotrz.axes)
            self.gui.ascot.input_plotwallcontour(axes=plotrz.axes)
            plotrz.draw()

        savebutton.configure(command=storesettings)
        plotbutton.configure(command=plot)

        self.xmin_entry = xmin_entry
        self.xmax_entry = xmax_entry
        self.xnum_entry = xnum_entry
        self.ymin_entry = ymin_entry
        self.ymax_entry = ymax_entry
        self.ynum_entry = ynum_entry
        self.phi_entry  = phi_entry
        self.time_entry = time_entry
        self.cmin_entry = cmin_entry
        self.cmax_entry = cmax_entry
        self.qchoice    = qchoice
        self.qinput     = qinput
        self.savebutton = savebutton
        self.plotbutton = plotbutton

    def selecttab(self):
        qnt, desc = zip(*self.gui.ascot.input_eval_list(show=False).items())
        self.qinput["values"] = qnt
        if self.qchoice.get() is None or self.qchoice.get() not in qnt:
            self.qchoice.set(qnt[0])

        self.canvas.slideshow(self)

class Preflight(ContentTab):

    def __init__(self, frame, canvas, gui):
        super().__init__(frame, canvas, gui)
        self.text = tk.Text(self, height=25, width=50)
        self.text.pack()
        self.plot = Preflight.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

    def selecttab(self):
        """Run and display results of the preflight tests.
        """
        msg = []
        msg += self.gui.ascot.preflight_inputspresent()
        msg += self.gui.ascot.preflight_optionsconsistent()
        #msg += check_bfield_psi0(gui.ascot)
        if len(msg) == 0:
            msg += ["Preflight checks completed with no issues."]

        msg0 = ""
        for m in msg:
            msg0 += m + "\n"

        self.text.configure(state='normal')
        self.text.delete("1.0", "end")
        self.text.insert("end", msg0)
        self.text.configure(state='disabled')

        self.plot.clear()
        self.gui.ascot.preflight_plottopview(axes=self.plot.axes[0])
        #plot_energypitch(ascot, axes=self.axes[1])
        #plot_rhophi(ascot, axes=self.axes[2])
        self.plot.draw()
        self.canvas.slideshow(self)

    class Canvas(PlotFrame):

        def set_axes(self):
            ax1 = self.fig.add_subplot(2,2,(1,3))
            ax2 = self.fig.add_subplot(2,2,2)
            ax3 = self.fig.add_subplot(2,2,4)
            return [ax1, ax2, ax3]

a5py/a5py/gui/inputcontents.py
ascot4fusion/ascot5
import tkinter as tk
from tkinter import ttk

from .components import PlotFrame, NumEntry

from a5py.ascot5io import options
#from a5py.ascot5io import mrk as marker

class OrbitFrame(ttk.Frame):

    def init(self, canvas):
        f1 = ttk.Frame(self)
        f2 = ttk.Frame(self)
        f1.pack(fill="x")
        f2.pack(fill="both")

        f1.columnconfigure(1, weight=1)

        self.milent = NumEntry(f1, labeltext="Mileage [s]", entrywidth=10,
                               labelwidth=12, anchor="w", defval=0.0001)
        self.pntent = NumEntry(f1, labeltext="N points", entrywidth=10,
                               labelwidth=12, anchor="w", defval=1000,
                               isint=True)

        self.milent.grid(column=0, row=0)
        self.pntent.grid(column=0, row=1)

        self.plotbutton = tk.Button(f1, text="Run and plot", width=8)
        self.plotbutton.grid(column=1, row=0, rowspan=2)

        self.r_entry = NumEntry(f2, labeltext="R [m]", entrywidth=5,
                                labelwidth=12, anchor="w", defval=6.7)
        self.p_entry = NumEntry(f2, labeltext="phi [deg]", entrywidth=5,
                                labelwidth=12, anchor="w", defval=0.0)
        self.z_entry = NumEntry(f2, labeltext="z [m]", entrywidth=5,
                                labelwidth=12, anchor="w", defval=0.0)
        self.e_entry = NumEntry(f2, labeltext="Energy [eV]", entrywidth=10,
                                labelwidth=14, anchor="w", defval=10.0e6)
        self.k_entry = NumEntry(f2, labeltext="Pitch [vpa/v]", entrywidth=10,
                                labelwidth=14, anchor="w", defval=0.99)
        self.r_entry.grid(column=0, row=0)
        self.p_entry.grid(column=0, row=1)
        self.z_entry.grid(column=0, row=2)
        self.e_entry.grid(column=1, row=0)
        self.k_entry.grid(column=1, row=1)

        self.canvas     = canvas


    def display(self, gui):

        gui.ascot.simulation_initinput()

        def plot():
            r      = self.r_entry.getval()
            phi    = self.p_entry.getval()
            z      = self.z_entry.getval()
            energy = self.e_entry.getval()
            pitch  = self.k_entry.getval()

            mrk = marker.generatemrk(nmrk=1, mrktype="gc", species="electron")
            mrk["r"][:]      = r
            mrk["phi"][:]    = phi
            mrk["z"][:]      = z
            mrk["energy"][:] = energy
            mrk["pitch"][:]  = pitch

            mil = 1e-5
            npoint = 1000
            interval = mil/npoint

            ids = mrk["ids"]

            opt = options.generateopt(clean=True)
            opt.update({"SIM_MODE"               : 2,
                        "ENABLE_ADAPTIVE"        : 1,
                        "ENABLE_ORBIT_FOLLOWING" : 1,
                        "ENDCOND_SIMTIMELIM"     : 1,
                        "ENDCOND_MAX_SIMTIME"    : 100,
                        "ENDCOND_MAX_MILEAGE"    : mil,
                        "ENABLE_ORBITWRITE"      : 1,
                        "ORBITWRITE_MODE"        : 1,
                        "ORBITWRITE_NPOINT"      : npoint,
                        "ORBITWRITE_INTERVAL"    : interval
            })

            gui.ascot.simulation_initmarkers("gc", mrk)
            gui.ascot.simulation_initoptions(opt)
            gui.ascot.simulation_run()
            rorb = gui.ascot.simulation_getorbit("r", ids)
            zorb = gui.ascot.simulation_getorbit("z", ids)
            gui.ascot.simulation_freeoutput()

            self.canvas.fig_rzview.clear()
            self.canvas.fig_rzview.axes.plot(rorb, zorb)
            self.canvas.fig_rzview.draw()

        self.plotbutton.configure(command=plot)


class OrbitCanvas(ttk.Frame):

    def init(self):
        fig_rzview = PlotFrame(self)
        fig_rzview.place(relheight=0.8, anchor="nw")

        self.fig_rzview = fig_rzview
        return self

a5py/a5py/gui/interactiveorbit.py
ascot4fusion/ascot5
import tkinter as tk
from tkinter import ttk
import numpy as np

from matplotlib.gridspec import GridSpec
from .components import ContentTab, PlotFrame, NumEntry, DropdownMenu, Tickbox,\
    ToggleButton
from a5py.routines.plotting import defaultcamera

class Summary(ContentTab):
    """Settings frame summarizing simulation output.

    The summary canvas is not interactive, so this frame does not contain any
    settings. Instead, it has a text box that shows sum of markers by
    end condition and possible error messages.
    """

    def __init__(self, frame, canvas, gui):
        """Initializes frame by creating all the widgets.
        """
        super().__init__(frame, canvas, gui)

        self.plot = Summary.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)
        self.markersummary = tk.Text(self, width=50)
        self.markersummary.pack(fill="both", expand=True)
        self.markersummary.config(state="disabled")

    def selecttab(self):
        """Display this frame and redraws canvas.
        """
        run = self.gui.ascot.data.active
        self.markersummary.config(state="normal")
        self.markersummary.delete("1.0", "end")

        text = "Marker end conditions:\n"
        end, err = run.getstate_markersummary()
        for m in end:
            text += m[1] + " : " + str(m[0]) + "\n"
        text += "\nErrors:\n"
        for m in err:
            text += m[0] + " at line " + str(m[1]) + " in file " + m[2] + "\n"

        self.markersummary.insert("end", text)
        self.markersummary.config(state="disabled")
        self.plot.clear()

        self.plot.axes_inirho.set_xlim(0,1.1)
        self.plot.axes_inirho.set_title("Initial radial position")
        run.plotstate_histogram("ini rho", xbins=np.linspace(0,1.1,55),
                                weight=True, axes=self.plot.axes_inirho)

        self.plot.axes_endrho.set_xlim([0,1.1])
        self.plot.axes_endrho.set_title("Final radial position")
        run.plotstate_histogram("end rho", xbins=np.linspace(0,1.1,55),
                                weight=True, axes=self.plot.axes_endrho)

        self.plot.axes_mileage.set_title("Final mileage")
        run.plotstate_histogram("log end mileage", xbins=55, weight=True,
                                axes=self.plot.axes_mileage)

        self.plot.axes_energy.set_title("Final energy")
        run.plotstate_histogram("log ini ekin", xbins=55, weight=True,
                                axes=self.plot.axes_energy)

        self.plot.axes_rz.set_title("Final R-z positions")
        run.plotstate_scatter("end R", "end z", c="C0", endcond=None,
                              axesequal=True, axes=self.plot.axes_rz)
        #self.gui.ascot.input_plotseparatrix(0, 0, axes=self.plot.axes_rz)
        self.gui.ascot.input_plotwallcontour(axes=self.plot.axes_rz)

        self.plot.axes_rhophi.set_xlim([0,1.2])
        self.plot.axes_rhophi.set_ylim([0,360])
        self.plot.axes_rhophi.set_xticks([0, 0.5, 1.0])
        self.plot.axes_rhophi.set_yticks([0, 180, 360])
        self.plot.axes_rhophi.set_title("Final rho-phi positions")
        run.plotstate_scatter("end rho", "end phimod", c="C0", endcond=None,
                              axesequal=False, axes=self.plot.axes_rhophi)

        self.plot.draw()
        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        """Canvas for showing summary plots.

        This canvas creates several axes where the plots can be placed.
        """

        def __init__(self, frame):
            """Name axes just so that they are easier to reference
            when plotting.
            """
            super().__init__(frame)
            self.axes_inirho  = None
            self.axes_endrho  = None
            self.axes_mileage = None
            self.axes_energy  = None
            self.axes_rz      = None
            self.axes_rhophi  = None

        def set_axes(self):
            """Override the PlotFrame method in order to create multiple
            axes.
            """
            gs = GridSpec(2,4)
            self.axes_inirho  = self.fig.add_subplot(gs[0,0])
            self.axes_endrho  = self.fig.add_subplot(gs[0,1])
            self.axes_mileage = self.fig.add_subplot(gs[1,0])
            self.axes_energy  = self.fig.add_subplot(gs[1,1])
            self.axes_rz      = self.fig.add_subplot(gs[:,2])
            self.axes_rhophi  = self.fig.add_subplot(gs[:,3])

            ax = [self.axes_inirho, self.axes_endrho, self.axes_mileage,
                  self.axes_energy, self.axes_rz, self.axes_rhophi]
            return ax

class StateScatter(ContentTab):
    """Settings frame for visualizing ini/end state.

    Visualizing is done either via scatter plot (2D/3D + color) or histogram
    (1D/2D).

    This frame contains Notebook widget containing two nested frames: one for
    the scatter plot settings and the other for the histogram settings. Both
    share the same StateCanvas frame.
    """

    def __init__(self, frame, canvas, gui):
        """Define nested frames and initialize the notebook.
        """
        super().__init__(frame, canvas, gui)
        self.plot = StateScatter.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        # For the overall layout we use three frames
        f1 = ttk.Frame(self) # For coordinate selection widgets
        f2 = ttk.Frame(self) # Plot and store buttons
        f3 = ttk.Frame(self) # Additional settings

        f1.grid(row=0, column=0, sticky="e")
        f2.grid(row=0, column=1, sticky="ne")
        f3.grid(row=1, column=0, sticky="ew")

        self.columnconfigure(1, weight=1)

        # Choose coordinate for [x,y,z,c] and tick if log scale is used
        self.xcrd = DropdownMenu(f1, log=True, label="x: ", width=12)
        self.ycrd = DropdownMenu(f1, log=True, label="y: ", width=12)
        self.zcrd = DropdownMenu(f1, log=True, label="z: ", width=12)
        self.ccrd = DropdownMenu(f1, log=True, label="c: ", width=12)

        self.xcrd.grid(row=0, column=0)
        self.ycrd.grid(row=1, column=0)
        self.zcrd.grid(row=2, column=0)
        self.ccrd.grid(row=3, column=0)

        # Is coordinate from ini or endstate
        self.xbtn = ToggleButton(f1, label1text="Ini", label2text="End")
        self.ybtn = ToggleButton(f1, label1text="Ini", label2text="End")
        self.zbtn = ToggleButton(f1, label1text="Ini", label2text="End")
        self.cbtn = ToggleButton(f1, label1text="Ini", label2text="End")

        self.xbtn.grid(row=0, column=1)
        self.ybtn.grid(row=1, column=1)
        self.zbtn.grid(row=2, column=1)
        self.cbtn.grid(row=3, column=1)

        # Plot and store buttons
        self.plotbutton = tk.Button(f2, text="Plot", width=3)
        self.savebutton = tk.Button(f2, text="Store", width=3)
        self.plotbutton.pack(anchor="nw")
        self.savebutton.pack(anchor="nw")

        # Only markers with this endstate are plotted
        self.endc = DropdownMenu(f3, label="Endcond: ", width=12,
                                 labelwidth=8)
        # Set x,y,z axes equal aspect ration
        self.axeq = Tickbox(f3, label=" Axis equal")

        self.endc.pack(side="left", anchor="w")
        self.axeq.pack(side="left", anchor="w")

        def plot():
            """Plot the given run using the active settings in the frame.
            """
            # Read and process settings
            equal   = self.axeq.getval() == 1
            endcond = self.endc.getval()

            if endcond == "all":
                endcond = None

            coords = [None] * 4
            for i, coord in enumerate(["x", "y", "z", "c"]):
                coords[i] = getattr(self, coord + "crd").getval()
                if coords[i] == "None":
                    coords[i] = None
                    continue

                log = "log" if getattr(self, coord + "crd").islog() else ""
                ini = "end" if getattr(self, coord + "btn").var.get() else "ini"
                coords[i] = log + " " + ini + " " + coords[i]

            run = self.gui.ascot.data.active
            if coords[2] is None:
                # 2D plot
                self.plot.clear()
                run.plotstate_scatter(coords[0], coords[1], c=coords[3],
                                      axesequal=equal, endcond=endcond,
                                      axes=self.plot.axes)
            else:
                # 3D plot (here we have to first create 3D axes for canvas)
                self.plot.make3d()
                self.plot.clear()
                run.plotstate_scatter(coords[0], coords[1], coords[2],
                                      c=coords[3], axesequal=equal,
                                      endcond=endcond, axes=self.plot.axes)
            self.plot.draw()

        self.plotbutton.configure(command=plot)
        #self.savebutton.configure(command=save)

    def selecttab(self):
        run = self.gui.ascot.data.active

        # Set values for the coordinate drop-down menus
        quantities = run.getstate_list()
        self.xcrd.setvals(list(quantities.keys()), "r", 0)
        self.ycrd.setvals(list(quantities.keys()), "z", 0)

        # For z and c one can choose "None" which is also the default
        qwithnone = quantities.copy()
        qwithnone["None"] = ("None","")
        self.zcrd.setvals(list(qwithnone.keys()), "None", 0)
        self.ccrd.setvals(list(qwithnone.keys()), "None", 0)

        end, err = run.getstate_markersummary()
        _, end = zip(*end)
        end = ["all"] + list(end)
        self.endc.setvals(end, "all")

        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        """Canvas for plotting the state plots.
        """

        def __init__(self, canvas):
            """Add attribute axes3d which flags whether next axes should be 3d.
            """
            self.axes3d = False
            super().__init__(canvas)

        def make3d(self):
            """Flag next axes to be created 3d.
            """
            self.axes3d = True

        def set_axes(self):
            """Override this method to allow generation of 3d axes.
            """
            if self.axes3d:
                self.axes3d = False
                return self.fig.add_subplot(1,1,1, projection="3d")

            return self.fig.add_subplot(1,1,1)

class StateHistogram(ContentTab):
    """Settings for the histogram frame.
    """

    def __init__(self, frame, canvas, gui):
        """Initialize widgets.
        """
        super().__init__(frame, canvas, gui)
        self.plot = StateHistogram.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        # Use four frames for layout
        f1 = ttk.Frame(self) # Coordinate selection
        f2 = ttk.Frame(self) # Plot and save buttons
        f3 = ttk.Frame(self) # Bin selection
        f4 = ttk.Frame(self) # Other settings

        f1.grid(row=0, column=0, sticky="e")
        f2.grid(row=0, column=1, sticky="ne")
        f3.grid(row=1, column=0, columnspan=2, sticky="w")
        f4.grid(row=2, column=0, sticky="w")

        self.columnconfigure(1, weight=1)

        # Coordinate selection
        self.xcrd = DropdownMenu(f1, log=True, label="x: ", width=18,
                                 trace=self.newxcoord)
        self.ycrd = DropdownMenu(f1, log=True, label="y: ", width=18,
                                 trace=self.newycoord)

        # Coordinate taken from ini or endstate
        self.xbtn = ToggleButton(f1, label1text="Ini", label2text="End")
        self.ybtn = ToggleButton(f1, label1text="Ini", label2text="End")

        self.xcrd.grid(row=0, column=0)
        self.ycrd.grid(row=1, column=0)

        self.xbtn.grid(row=0, column=1)
        self.ybtn.grid(row=1, column=1)

        # Plot and save buttons
        self.plotbutton = tk.Button(f2, text="Plot", width=3)
        self.savebutton = tk.Button(f2, text="Store", width=3)
        self.plotbutton.pack(anchor="nw")
        self.savebutton.pack(anchor="nw")

        # Weighted histogram (or no), axes equal aspect ratio (or no)
        self.wght = Tickbox(f4, width=1)
        self.axeq = Tickbox(f4, width=1)
        self.zlog = Tickbox(f4, width=1)

        self.wght.grid(row=0,column=0,sticky="w")
        tk.Label(f4, text="With weights").grid(row=0,column=1,sticky="w")
        self.axeq.grid(row=1,column=0,sticky="w")
        tk.Label(f4, text="Axis equal").grid(row=1,column=1,sticky="w")
        self.zlog.grid(row=2,column=0,sticky="w")
        tk.Label(f4, text="log10 scale").grid(row=2,column=1,sticky="w")

        # Set NumEntries for xmin, xmax, and nx inputs
        self.xmin_entry = NumEntry(
            f3, labeltext="rho [1] = ", entrywidth=5, labelwidth=22,
            anchor="e", defval=0.0)
        self.xmax_entry = NumEntry(f3, labeltext="–", entrywidth=5,
                                   labelwidth=2, anchor="c", defval=1.0)
        self.xnum_entry = NumEntry(f3, labeltext="x", entrywidth=5,
                                   labelwidth=2, anchor="c", defval=50,
                                   isint=True)

        # Set NumEntries for ymin, ymax, and ny inputs
        self.ymin_entry = NumEntry(
            f3, labeltext="* = ", entrywidth=5, labelwidth=22,
            anchor="e", defval=0.0)
        self.ymax_entry = NumEntry(f3, labeltext="–", entrywidth=5,
                                   labelwidth=2, anchor="c", defval=1.0)
        self.ynum_entry = NumEntry(f3, labeltext="x", entrywidth=5,
                                   labelwidth=2, anchor="c", defval=50,
                                   isint=True)


        self.xmin_entry.grid(row=0, column=0)
        self.xmax_entry.grid(row=0, column=1)
        self.xnum_entry.grid(row=0, column=2)
        self.ymin_entry.grid(row=1, column=0)
        self.ymax_entry.grid(row=1, column=1)
        self.ynum_entry.grid(row=1, column=2)

        #gui.params.add(
        #output_stateplot_minr=xmin_entry.choice,
        #output_stateplot_maxr=xmax_entry.choice,
        #output_stateplot_numr=xnum_entry.choice,
        #output_stateplot_minz=ymin_entry.choice,
        #output_stateplot_maxz=ymax_entry.choice,
        #output_stateplot_numz=ynum_entry.choice,
        #output_stateplot_qnt=xcrd.var
        #)

        def plot():
            """Plot the histogram for the given run using active settings.
            """
            xcoord = self.xcrd.getval()
            xcoord = "end " + xcoord if self.xbtn.var.get() else "ini " + xcoord
            if self.xcrd.islog(): xcoord = "log " + xcoord

            ycoord = self.ycrd.getval()
            if ycoord == "None":
                ycoord = None
            else:
                ycoord = "end " + ycoord if self.ybtn.var.get() \
                    else "ini " + ycoord
                if self.ycrd.islog(): ycoord = "log " + ycoord

            logscale  = self.zlog.getval() == 1
            axesequal = self.axeq.getval() == 1
            weight    = self.wght.getval() == 1

            def getbins(vmin, vmax, nv, islog):
                """Short function to check user input and generate bins.
                """
                if vmax < vmin or nv <= 0: return None
                if islog and (vmin <= 0 or vmax <= 0): return None

                if islog:
                    return np.logspace(np.log10(vmin), np.log10(vmax), nv)
                else:
                    return np.linspace(vmin, vmax, nv)

            xbins = getbins(self.xmin_entry.getval(),
                            self.xmax_entry.getval(),
                            self.xnum_entry.getval(),
                            False)
            ybins = getbins(self.ymin_entry.getval(),
                            self.ymax_entry.getval(),
                            self.ynum_entry.getval(),
                            False)

            self.plot.clear()
            run = self.gui.ascot.data.active
            run.plotstate_histogram(
                xcoord, y=ycoord, xbins=xbins, ybins=ybins,
                endcond=None, weight=weight, logscale=logscale,
                axesequal=axesequal, axes=self.plot.axes)
            self.plot.draw()

        self.plotbutton.configure(command=plot)
        #self.savebutton.configure(command=save)

    def newxcoord(self, *args):
        """x-coordinate changed, update the label on the x bin selection.
        """
        val  = self.xcrd.var.get()
        try:
            unit = str(self.gui.ascot.data.active.getstate(val).units)
        except AttributeError:
            unit = "1"
        self.xmin_entry.setlabel(val + " [" + unit + "]" + " = ")


    def newycoord(self, *args):
        """y-coordinate changed, update the label on the y bin selection.

        If None, we disable the bin selection completely for y.
        """
        val = self.ycrd.var.get()
        if val == "None":
            self.ymin_entry.disable()
            self.ymax_entry.disable()
            self.ynum_entry.disable()
            self.ymin_entry.setlabel("* = ")
        else:
            try:
                unit = str(self.gui.ascot.data.active.getstate(val).units)
            except AttributeError:
                unit = "1"
            self.ymin_entry.enable()
            self.ymax_entry.enable()
            self.ynum_entry.enable()
            self.ymin_entry.setlabel(val + " [" + unit + "]" + " = ")

    def selecttab(self):
        """Display state frame for the run that is currently active.
        """
        run = self.gui.ascot.data.active
        # Set values for the drop-down menus and allow y coordinate to
        # have None option (in which case the histogram is 1D)
        quantities = run.getstate_list()
        self.xcrd.setvals(list(quantities.keys()), "rho", 0)
        qwithnone = quantities.copy()
        qwithnone["None"] = ("None","")
        self.ycrd.setvals(list(qwithnone.keys()), "None", 0)
        try:
            if self.xbtn.switch.isdisabled:
                self.xbtn.enable()
                self.ybtn.enable()
                self.zbtn.enable()
                self.cbtn.enable()
                self.xbtn.enable()
                self.ybtn.enable()

        except AttributeError:
            self.xbtn.disable()
            self.ybtn.disable()
            self.zbtn.disable()
            self.cbtn.disable()
            self.xbtn.disable()
            self.ybtn.disable()

        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        """Canvas for plotting the state plots.
        """
        pass

class Orbit(ContentTab):

    def __init__(self, frame, canvas, gui):
        """Initializes frame by creating all the widgets.
        """
        super().__init__(frame, canvas, gui)
        self.plot = Orbit.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        f1 = ttk.Frame(self)
        f2 = ttk.Frame(self)
        f3 = ttk.Frame(self)

        f1.grid(row=0, column=0, sticky="w")
        f2.grid(row=0, column=1, sticky="ne")
        f3.grid(row=1, column=0, sticky="ew")

        self.columnconfigure(1, weight=1)

        self.xcrd = DropdownMenu(f1, log=True, label="x: ", width=12,
                                 labelwidth=2)
        self.ycrd = DropdownMenu(f1, log=True, label="y: ", width=12,
                                 labelwidth=2)
        self.zcrd = DropdownMenu(f1, log=True, label="z: ", width=12,
                                 labelwidth=2)
        self.ccrd = DropdownMenu(f1, log=True, label="c: ", width=12,
                                 labelwidth=2)

        self.xcrd.grid(row=0, column=0)
        self.ycrd.grid(row=1, column=0)
        self.zcrd.grid(row=2, column=0)
        self.ccrd.grid(row=3, column=0)

        self.plotbutton = tk.Button(f2, text="Plot", width=3)
        self.savebutton = tk.Button(f2, text="Store", width=3)
        self.plotbutton.pack(anchor="nw")
        self.savebutton.pack(anchor="nw")

        self.endc = DropdownMenu(f3, label="Endcond: ", width=12,
                                 labelwidth=8)
        self.axeq = Tickbox(f3, label=" Axes equal")

        self.endc.pack(side="left", anchor="w")
        self.axeq.pack(side="left", anchor="w")

        def plot():
            equal   = self.axeq.getval() == 1
            endcond = self.endc.getval()

            if endcond == "all":
                endcond = None

            coords = [None] * 4
            for i, coord in enumerate(["x", "y", "z", "c"]):
                coords[i] = getattr(self, coord + "crd").getval()
                if coords[i] == "None":
                    coords[i] = None
                    continue

                if getattr(self, coord + "crd").islog():
                    coords[i] = "log " + coords[i]

            run = self.gui.ascot.data.active
            if coords[2] is None:
                self.plot.clear()
                run.plotorbit_trajectory(coords[0], coords[1], c=coords[3],
                                         axesequal=equal, endcond=endcond,
                                         axes=self.plot.axes)
                #if ccoord == "None":
                #    run.orbit.plot(xcoord, ycoord, equal=equal, log=log,
                #                   endcond=endcond, axes=self.plot.axes)
                #else:
                #    run.orbit.scatter(xcoord, ycoord, c=ccoord, equal=equal,
                #                      endcond=endcond, axes=self.plot.axes)
            else:
                self.plot.make3d()
                self.plot.clear()
                run.plotorbit_trajectory(coords[0], coords[1], coords[2],
                                         c=coords[3], axesequal=equal,
                                         endcond=endcond, axes=self.plot.axes)
                #if ccoord == "None":
                #    run.orbit.plot(xcoord, ycoord, zcoord, equal=equal,
                #                   endcond=endcond, axes=self.plot.axes)
                #else:
                #    run.orbit.scatter(xcoord, ycoord, zcoord, ccoord,
                #                      equal=equal, log=log,
                #                      endcond=endcond, axes=self.plot.axes)

            self.plot.draw()

        self.plotbutton.configure(command=plot)
        #self.savebutton.configure(command=save)

    def selecttab(self):
        run = self.gui.ascot.data.active

        # Set values for the coordinate drop-down menus
        quantities = run.getstate_list()
        self.xcrd.setvals(list(quantities.keys()), "r", 0)
        self.ycrd.setvals(list(quantities.keys()), "z", 0)

        # For z and c one can choose "None" which is also the default
        qwithnone = quantities.copy()
        qwithnone["None"] = ("None","")
        self.zcrd.setvals(list(qwithnone.keys()), "None", 0)
        self.ccrd.setvals(list(qwithnone.keys()), "None", 0)

        end, err = run.getstate_markersummary()
        _, end = zip(*end)
        end = ["all"] + list(end)
        self.endc.setvals(end, "all")

        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        """Canvas for plotting the state plots.
        """

        def __init__(self, canvas):
            """Add attribute axes3d which flags whether next axes should be 3d.
            """
            self.axes3d = False
            super().__init__(canvas)

        def make3d(self):
            """Flag next axes to be created 3d.
            """
            self.axes3d = True

        def set_axes(self):
            """Override this method to allow generation of 3d axes.
            """
            if self.axes3d:
                self.axes3d = False
                return self.fig.add_subplot(1,1,1, projection="3d")

            return self.fig.add_subplot(1,1,1)

class Poincare(ContentTab):

    def __init__(self, frame, canvas, gui):
        super().__init__(frame, canvas, gui)
        self.plot = Poincare.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        f1 = ttk.Frame(self)
        f2 = ttk.Frame(self)

        f1.grid(row=0, column=0, sticky="w")
        f2.grid(row=0, column=1, sticky="ne")

        self.columnconfigure(1, weight=1)

        self.plane = DropdownMenu(f1, label="           Plane: ", width=25,
                                  labelwidth=10, labelanchor="w")
        #self.coord = DropdownMenu(f1, label="Coordinates: ", width=8,
        #                          labelwidth=10, labelanchor="w")
        self.plane.pack(fill="x", anchor="w")
        #self.coord.pack(fill="x", anchor="w")

        self.plotbutton = tk.Button(f2, text="Plot", width=3)
        self.savebutton = tk.Button(f2, text="Store", width=3)
        self.plotbutton.pack(anchor="nw")
        self.savebutton.pack(anchor="nw")

        def plot():
            run = self.gui.ascot.data.active
            self.plot.clear()
            plane = self.plane.getval().split(" ")
            plane = plane[0] + " " + plane[1]

            run.plotorbit_poincare(plane, connlen=True, axes=self.plot.axes)
            self.plot.draw()

        self.plotbutton.configure(command=plot)
        #self.savebutton.configure(command=save)


    def selecttab(self):
        run = self.gui.ascot.data.active
        pol, tor, rad = run.getorbit_poincareplanes()
        planes = []
        for i, p in enumerate(pol):
            planes += ["pol %d at %f deg" % (i+1, p[0])]
        for i, p in enumerate(tor):
            planes += ["tor %d at %f deg" % (i+1, p[0])]
        for i, p in enumerate(rad):
            planes += ["rad %d at %f rho" % (i+1, p[0])]

        self.plane.setvals(planes, planes[0])
        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        pass

class Dists(ContentTab):

    def __init__(self, frame, canvas, gui):
        """Initializes frame by creating all the widgets.
        """
        super().__init__(frame, canvas, gui)
        self.plot = Dists.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        f1 = ttk.Frame(self)
        f2 = ttk.Frame(self)

        f1.grid(row=0, column=0, sticky="w")
        f2.grid(row=0, column=1, sticky="ne")

        self.columnconfigure(1, weight=1)

        self.source = DropdownMenu(f1, label="Source: ", width=24,
                                   labelwidth=8, labelanchor="w",
                                   trace=self.setcoords)
        self.xcrd = DropdownMenu(f1, label="x: ", width=5, labelwidth=2,
                                 trace=self.setxcoord)
        self.ycrd = DropdownMenu(f1, label="y: ", width=5, labelwidth=2)
        self.source.grid(row=0, column=0)
        self.xcrd.grid(row=1, column=0, sticky="e")
        self.ycrd.grid(row=2, column=0, sticky="e")

        self.plotbutton = tk.Button(f2, text="Plot", width=3)
        self.savebutton = tk.Button(f2, text="Store", width=3)
        self.plotbutton.pack(anchor="nw")
        self.savebutton.pack(anchor="nw")

        def plot():
            whichdist = self.source.getval()
            run  = self.gui.ascot.data.active
            dist = run.getdist(whichdist)
            x = self.xcrd.getval()
            y = self.ycrd.getval()

            integrate = {}
            for a in dist.abscissae:
                if x != a and y != a:
                    integrate[a] = np.s_[:]
            dist.integrate(**integrate)

            self.plot.clear()
            run.plotdist(dist, axes=self.plot.axes)
            self.plot.draw()

        self.plotbutton.configure(command=plot)

    def setcoords(self, *args):
        dist = self.source.var.get()
        coords = []
        if "com" in dist:
            coords += ["mu", "ekin", "ptor"]
        else:
            if "rho" in dist:
                coords += ["rho", "theta", "phi"]
            else:
                coords += ["r", "z", "phi"]

            if "5d" in dist:
                coords += ["ppar", "pperp"]
            else:
                coords += ["pr", "pz", "pphi"]

            coords += ["time", "charge"]

        self.xcrd.setvals(coords, coords[0])
        self.ycrd.setvals(coords[1:]+["None"], "None")

    def setxcoord(self, *args):
        dist = self.source.var.get()
        x = self.xcrd.var.get()
        y = self.ycrd.var.get()

        coords = []
        if "com" in dist:
            coords += ["mu", "ekin", "ptor"]
        else:
            if "rho" in dist:
                coords += ["rho", "theta", "phi"]
            else:
                coords += ["r", "z", "phi"]

            if "5d" in dist:
                coords += ["ppar", "pperp"]
            else:
                coords += ["pr", "pz", "pphi"]

            coords += ["time", "charge"]

        coords.remove(x)
        if x == y:
            self.ycrd.setvals(coords, "None")
        else:
            self.ycrd.setvals(coords, y)

    def selecttab(self):
        run = self.gui.ascot.data.active
        dists = run.getdist(None)
        self.source.setvals(dists, dists[0])

        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        pass

class Moments(ContentTab):

    def __init__(self, frame, canvas, gui):
        """Initializes frame by creating all the widgets.
        """
        super().__init__(frame, canvas, gui)
        self.plot = Moments.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        f1 = ttk.Frame(self)
        f2 = ttk.Frame(self)
        f1.grid(row=0, column=0, sticky="w")
        f2.grid(row=0, column=1, sticky="ne")

        self.columnconfigure(1, weight=1)

        self.source = DropdownMenu(f1, label="  Source: ", width=24,
                                   labelwidth=8, labelanchor="w")
        self.qnt = DropdownMenu(f1, label="Quantity: ", width=24,
                                labelwidth=8, labelanchor="w")
        self.source.grid(row=0, column=0)
        self.qnt.grid(row=1, column=0)

        self.plotbutton = tk.Button(f2, text="Plot", width=3)
        self.savebutton = tk.Button(f2, text="Store", width=3)
        self.plotbutton.pack(anchor="nw")
        self.savebutton.pack(anchor="nw")

        self.qnt.setvals(["density"], "density")

        def plot():
            dist = self.source.var.get()
            moment  = self.qnt.var.get()
            run  = self.gui.ascot.data.active
            dist = run.getdist(dist)
            mom  = run.getdist_moments(dist, moment)

            self.plot.clear()
            run.plotdist_moments(mom, moment, axes=self.plot.axes)
            self.plot.draw()

        self.plotbutton.configure(command=plot)

    def selecttab(self):
        run = self.gui.ascot.data.active
        dists = run.getdist(None)
        vals = []
        if "5d" in dists: vals += ["5d"]
        if "rho5d" in dists: vals += ["rho5d"]
        self.source.setvals(vals, vals[0])

        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        pass

class LossSummary(ContentTab):

    def __init__(self, frame, canvas, gui):
        super().__init__(frame, canvas, gui)
        self.plot = LossSummary.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        self.losssummary = tk.Text(self, width=50)
        self.losssummary.pack(fill="both", expand=True)
        self.losssummary.config(state="disabled")

    def selecttab(self):
        run = self.gui.ascot.data.active
        self.losssummary.config(state="normal")
        self.losssummary.delete("1.0", "end")

        text = ""
        msg = run.getstate_losssummary()
        for m in msg:
            text += m + "\n"

        self.losssummary.insert("end", text)
        self.losssummary.config(state="disabled")

        self.plot.clear()
        #run.endstate.plot_lossmap(5.0, axes=self.canvas.axes[4])
        self.plot.draw()
        self.canvas.slideshow(self)

    class Canvas(PlotFrame):

        def set_axes(self):
            """Override this method to allow generation of multiple axes.
            """
            gs = GridSpec(2,3)
            ax = []
            ax += [self.fig.add_subplot(gs[:,0])]
            ax += [self.fig.add_subplot(gs[0,1])]
            ax += [self.fig.add_subplot(gs[1,1])]
            ax += [self.fig.add_subplot(gs[0,2])]
            ax += [self.fig.add_subplot(gs[1,2])]
            return ax

class WallLoad(ContentTab):

    def __init__(self, frame, canvas, gui):
        super().__init__(frame, canvas, gui)
        self.plot = WallLoad.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        self.loadsummary = tk.Text(self, width=50)
        self.loadsummary.pack(fill="both", expand=True)
        self.loadsummary.config(state="disabled")

    def selecttab(self):
        run = self.gui.ascot.data.active
        self.loadsummary.config(state="normal")
        self.loadsummary.delete("1.0", "end")

        wetted, epeak = run.getwall_figuresofmerit()
        text = "Wetted area: " \
            + np.format_float_scientific(wetted, precision=2) \
            + " " + str(wetted.units) + "\n"\
            + "Peak load: " \
            + np.format_float_scientific(epeak, precision=2) \
            + " " + str(epeak.units) + "\n"\

        self.loadsummary.insert("end", text)
        self.loadsummary.config(state="disabled")

        self.plot.clear()
        run.plotwall_loadvsarea(axes=self.plot.axes[0])
        run.plotwall_torpol(axes=self.plot.axes[1])
        self.plot.draw()
        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        """Canvas for plotting.
        """

        def set_axes(self):
            """Override this method to allow generation of multiple axes.
            """
            gs = GridSpec(1,2)
            ax = []
            ax += [self.fig.add_subplot(gs[:,0])]
            ax += [self.fig.add_subplot(gs[0,1])]
            return ax


class Wall3D(ContentTab):

    def __init__(self, frame, canvas, gui):
        super().__init__(frame, canvas, gui)
        self.plot = Wall3D.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        # Use three frames for layout
        f1 = ttk.Frame(self) # Camera controls
        f2 = ttk.Frame(self) # Plot and save buttons
        f3 = ttk.Frame(f2) # Other settings (located inside f2)

        f1.grid(row=0, column=0, sticky="w")
        f2.grid(row=0, column=1, sticky="nwes")

        self.columnconfigure(1, weight=1)

        # Camera selection
        self.csel = DropdownMenu(f1, label="", width=12, labelwidth=0)
        self.csel.grid(row=0, column=0, columnspan=2, sticky="w")
        self.csel.setvals(["Camera 1"],#, "Camera 2", "Camera 3",
                          #"Camera 4", "Camera 5", "Camera 6",
                          #"Camera 7", "Camera 8", "Camera 9"],
                          "Camera 1")

        # Interactive plot
        self.intrbutton = tk.Button(f1, text="Interactive", width=6)
        self.intrbutton.grid(row=0, column=2, sticky="e")

        # Camera position control
        l = tk.Label(f1, text="Camera location:")
        l.grid(row=1, column=0, columnspan=2, sticky="w")
        self.cposr = NumEntry(f1, labeltext="R", entrywidth=7,
                              labelwidth=2, anchor="e", defval=0.0)
        self.cposp = NumEntry(f1, labeltext=" phi", entrywidth=7,
                              labelwidth=4, anchor="e", defval=0.0)
        self.cposz = NumEntry(f1, labeltext=" z", entrywidth=7,
                              labelwidth=2, anchor="e", defval=0.0)

        self.cposr.grid(row=2, column=0)
        self.cposp.grid(row=2, column=1)
        self.cposz.grid(row=2, column=2)

        # Camera focus control
        l = tk.Label(f1, text="Focus point:")
        l.grid(row=3, column=0, columnspan=2, sticky="w")
        self.cfocr = NumEntry(f1, labeltext="R", entrywidth=7,
                              labelwidth=2, anchor="e", defval=0.0)
        self.cfocp = NumEntry(f1, labeltext="phi", entrywidth=7,
                              labelwidth=4, anchor="e", defval=0.0)
        self.cfocz = NumEntry(f1, labeltext=" z", entrywidth=7,
                              labelwidth=2, anchor="e", defval=0.0)

        self.cfocr.grid(row=4, column=0)
        self.cfocp.grid(row=4, column=1)
        self.cfocz.grid(row=4, column=2)

        # Camera angle control
        l = tk.Label(f1, text="Azimuth:")
        l.grid(row=5, column=0, sticky="w")
        l = tk.Label(f1, text="Elevation:")
        l.grid(row=5, column=1, sticky="w")
        l = tk.Label(f1, text="Rotation:")
        l.grid(row=5, column=2, sticky="w")
        self.canga = NumEntry(f1, labeltext=" ", entrywidth=7,
                              labelwidth=2, anchor="e", defval=0.0)
        self.cange = NumEntry(f1, labeltext="   ", entrywidth=7,
                              labelwidth=4, anchor="e", defval=0.0)
        self.cangr = NumEntry(f1, labeltext="  ", entrywidth=7,
                              labelwidth=2, anchor="e", defval=0.0)

        self.canga.grid(row=6, column=0)
        self.cange.grid(row=6, column=1)
        self.cangr.grid(row=6, column=2)

        # Plot, save, and interactive buttons
        self.plotbutton = tk.Button(f2, text="Plot", width=3)
        self.savebutton = tk.Button(f2, text="Store", width=3)
        self.plotbutton.pack(anchor="ne")
        self.savebutton.pack(anchor="ne")

        f3.pack(fill="both")
        l = tk.Label(f3, text="Markers")
        l.grid(row=0, column=1, sticky="w")
        l = tk.Label(f3, text="Wall loads")
        l.grid(row=1, column=1, sticky="w")
        l = tk.Label(f3, text="log10 scale")
        l.grid(row=2, column=1, sticky="w")
        self.showmarkers = Tickbox(f3, width=1)
        self.showmarkers.grid(row=0, column=0)
        self.showloads = Tickbox(f3, width=1)
        self.showloads.grid(row=1, column=0)
        self.logscale = Tickbox(f3, width=1)
        self.logscale.grid(row=2, column=0)

        self.wallmesh = None

    def selecttab(self):
        """Setup data.
        """
        run = self.gui.ascot.data.active
        self.wallmesh = run.getwall_3dmesh()
        self.setcamera(*defaultcamera(self.wallmesh))

        self.plotbutton.configure(command=lambda:self.still(run))
        self.intrbutton.configure(command=lambda:self.interactive(run))

        self.canvas.slideshow(self)

    def still(self, run):
        """Plot still picture on the canvas.
        """
        self.plot.clear()
        cpos, cfoc, cang = self.getcamera()
        points = None
        if self.showmarkers.getval() == 1:
            points = run.getstate_pointcloud(endcond="wall")
        data = "eload" if self.showloads.getval() == 1 else None
        log = self.logscale.getval() == 1
        run.plotwall_3dstill(self.wallmesh,
                             points=points, data=data, log=log,
                             cpos=cpos, cfoc=cfoc, cang=cang,
                             axes=self.plot.axes)
        self.plot.draw()

    def interactive(self, run):
        """Open new window for interactive 3D view.
        """
        def recordcamera(plotter):
            cpos = plotter.camera_position[0]
            cfoc = plotter.camera_position[1]
            cang = plotter.camera_position[2]
            cang = [plotter.camera.azimuth,
                    plotter.camera.elevation,
                    plotter.camera.roll]

            self.setcamera(cpos, cfoc, cang)

        cpos, cfoc, cang = self.getcamera()
        points = None
        if self.showmarkers.getval() == 1:
            points = run.getstate_pointcloud(endcond="wall")
        data = "eload" if self.showloads.getval() == 1 else None
        log = self.logscale.getval() == 1
        run.plotwall_3dinteractive(self.wallmesh,
                                   ("k", recordcamera),
                                   points=points, data=data, log=log,
                                   cpos=cpos, cfoc=cfoc, cang=cang)

    def setcamera(self, cpos, cfoc, cang):
        """Record and display camera coordinates (given in xyz) in rpz.
        """

        def roundval(val):
            """
            Round to three decimals.
            """
            return int(val * 100) / 100

        cr = np.sqrt( cpos[0]**2 + cpos[1]**2 )
        cp = np.arctan2( cpos[1], cpos[0] )
        cz = cpos[2]

        self.cposr.setval( roundval(cr) )
        self.cposp.setval( roundval(cp * 180 / np.pi) )
        self.cposz.setval( roundval(cz) )

        cr = np.sqrt( cfoc[0]**2 + cfoc[1]**2 )
        cp = np.arctan2( cfoc[1], cfoc[0] )
        cz = cfoc[2]

        self.cfocr.setval( roundval(cr) )
        self.cfocp.setval( roundval(cp * 180 / np.pi) )
        self.cfocz.setval( roundval(cz) )

        self.canga.setval( roundval(cang[0]) )
        self.cange.setval( roundval(cang[1]) )
        self.cangr.setval( roundval(cang[2]) )

    def getcamera(self):
        """Get camera coordinates (stored as rpz) as xyz.
        """
        cr = self.cposr.getval()
        cp = self.cposp.getval() * np.pi / 180
        cz = self.cposz.getval()

        cpos = [cr * np.cos(cp), cr * np.sin(cp), cz]

        cr = self.cfocr.getval()
        cp = self.cfocp.getval() * np.pi / 180
        cz = self.cfocz.getval()

        cfoc = [cr * np.cos(cp), cr * np.sin(cp), cz]

        cang = [self.canga.getval(),
                self.cange.getval(),
                self.cangr.getval()]

        return (cpos, cfoc, cang)

    class Canvas(PlotFrame):
        pass

a5py/a5py/gui/resultcontents.py
ascot4fusion/ascot5
import tkinter as tk
from tkinter import ttk

from .components import ContentTab, PlotFrame, NumEntry

from a5py.ascot5io.options import Opt
from a5py.ascot5io.marker import Marker

class Trace(ContentTab):

    def __init__(self, frame, canvas, gui):
        super().__init__(frame, canvas, gui)
        self.plot = Trace.Canvas(canvas)
        self.canvas.slideadd(self, self.plot)

        f1 = ttk.Frame(self)
        f2 = ttk.Frame(self)
        f1.pack(fill="x")
        f2.pack(fill="both")

        f1.columnconfigure(1, weight=1)

        self.milent = NumEntry(f1, labeltext="Mileage [s]", entrywidth=10,
                               labelwidth=12, anchor="w", defval=0.0001)
        self.pntent = NumEntry(f1, labeltext="N points", entrywidth=10,
                               labelwidth=12, anchor="w", defval=1000,
                               isint=True)

        self.milent.grid(column=0, row=0)
        self.pntent.grid(column=0, row=1)

        self.plotbutton = tk.Button(f1, text="Run and plot", width=8)
        self.plotbutton.grid(column=1, row=0, rowspan=2)

        self.r_entry = NumEntry(f2, labeltext="R [m]", entrywidth=5,
                                labelwidth=12, anchor="w", defval=6.7)
        self.p_entry = NumEntry(f2, labeltext="phi [deg]", entrywidth=5,
                                labelwidth=12, anchor="w", defval=0.0)
        self.z_entry = NumEntry(f2, labeltext="z [m]", entrywidth=5,
                                labelwidth=12, anchor="w", defval=0.0)
        self.e_entry = NumEntry(f2, labeltext="Energy [eV]", entrywidth=10,
                                labelwidth=14, anchor="w", defval=10.0e6)
        self.k_entry = NumEntry(f2, labeltext="Pitch [vpa/v]", entrywidth=10,
                                labelwidth=14, anchor="w", defval=0.99)
        self.r_entry.grid(column=0, row=0)
        self.p_entry.grid(column=0, row=1)
        self.z_entry.grid(column=0, row=2)
        self.e_entry.grid(column=1, row=0)
        self.k_entry.grid(column=1, row=1)

        def plot():
            r      = self.r_entry.getval()
            phi    = self.p_entry.getval()
            z      = self.z_entry.getval()
            energy = self.e_entry.getval()
            pitch  = self.k_entry.getval()

            mrk = Marker.generate("gc", n=1, species="electron")
            mrk["r"][:]      = r
            mrk["phi"][:]    = phi
            mrk["z"][:]      = z
            mrk["energy"][:] = energy
            mrk["pitch"][:]  = pitch

            mil = 1e-5
            npoint = 1000
            interval = mil/npoint

            ids = mrk["ids"]

            opt = Opt.get_default()
            opt.update({"SIM_MODE"               : 2,
                        "ENABLE_ADAPTIVE"        : 1,
                        "ENABLE_ORBIT_FOLLOWING" : 1,
                        "ENDCOND_SIMTIMELIM"     : 1,
                        "ENDCOND_LIM_SIMTIME"    : 100,
                        "ENDCOND_MAX_MILEAGE"    : mil,
                        "ENABLE_ORBITWRITE"      : 1,
                        "ORBITWRITE_MODE"        : 1,
                        "ORBITWRITE_NPOINT"      : npoint,
                        "ORBITWRITE_INTERVAL"    : interval
            })

            self.gui.ascot.simulation_initmarkers(**mrk)
            self.gui.ascot.simulation_initoptions(**opt)
            vrun = self.gui.ascot.simulation_run()

            self.plot.clear()
            vrun.plotorbit_trajectory("r", "z", axes=self.plot.axes)
            self.plot.draw()

            self.gui.ascot.simulation_free(diagnostics=True)

        self.plotbutton.configure(command=plot)

    def selecttab(self):
        self.canvas.slideshow(self)

    class Canvas(PlotFrame):
        pass

a5py/a5py/gui/runcontents.py
ascot4fusion/ascot5
"""
Generate marker (and some trivial particle) distributions.

File: marker/generator.py
"""
import numpy as np
import scipy.constants as const

from scipy.interpolate import RegularGridInterpolator

import a5py.marker.phasespace as phasespace
from a5py.ascotpy.ascotpy import Ascotpy

import importlib.util as util

plt = util.find_spec("matplotlib")
if plt:
    import matplotlib.pyplot as plt
    from matplotlib.gridspec import GridSpec

def gen_markerdist(a5, mass, charge, energy, rgrid, zgrid, kgrid,
                   rhoksidist=None, plot=False):
    """
    Generate marker (r,zksi) distribution from a given distribution.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        rgrid : array_like (n,1) <br>
            R coordinate grid.
        zgrid : array_like (n,1) <br>
            z coordinate grid.
        kgrid : array_like (n,1) <br>
            ksi coordinate grid.
        rhoksidist : tuple, optional <br>
            Tuple (rhogrid, ksigrid, weights) if markers are to be generated
            based weighted (rho_omp, ksi_omp) distribution. See phasespace.py.

    Returns:
        Tuple (rgrid,zgrid,ksigrid,markerdist) where first three are the grid
        axes (same as was given) and fourth element is the marker distribution
        normalized to one.
    """
    free = False
    if isinstance(a5, str):
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    # Marker distribution is generated with a Monte Carlo method by generating
    # chuncks of markers until the change in each cell is less than the
    # tolerance.
    Ndraw = 100000
    tol = 1e-1
    err = tol+1
    newdist = np.zeros((rgrid.size-1, zgrid.size-1, kgrid.size-1))
    probdensity = np.append( [0], np.cumsum(rhoksidist[2].ravel()) )
    while(err > tol):
        olddist = np.copy(newdist)

        if rhoksidist is not None:
            rhogrid = rhoksidist[0]
            ksigrid = rhoksidist[1]

            rhovals, ksivals = np.meshgrid(rhogrid[:-1], ksigrid[:-1],
                                           indexing='ij')

            prob = np.random.rand(Ndraw)
            idx  = np.searchsorted(probdensity, prob, side='left')-1

            weights = np.ones(idx.shape)
            rhovals = rhovals.ravel()[idx] + np.random.rand(Ndraw) \
                      * (rhogrid[1]-rhogrid[0])
            ksivals = ksivals.ravel()[idx] + np.random.rand(Ndraw) \
                      * (ksigrid[1]-ksigrid[0])

            newdist += phasespace.maprhoksi2rzk(a5, mass, charge, energy,
                                                rhovals, ksivals,
                                                rgrid, zgrid, kgrid,
                                                weights=weights)

        if np.sum(olddist.ravel()) > 0:
            ratio = np.absolute(
                ( newdist.ravel() / np.sum(newdist.ravel()) )
                / ( olddist.ravel() / np.sum(olddist.ravel())) )
            err = np.mean(np.absolute( ratio[np.isfinite(ratio)] - 1))
            print("Iteration error: " + str(err) + " Tolerance: " + str(tol))

    if free:
        a5.free(bfield=True)


    # Normalize
    newdist /= np.sum(newdist.ravel())

    if plot:
        plotdist((rgrid, zgrid, kgrid, newdist))

    return (rgrid, zgrid, kgrid, newdist)


def gen_DTalphadist(a5, rgrid, zgrid, ksigrid, plot=False):
    """
    Generate DT fusion alpha particle distribution based on NRL formula.

    This function assumes the HDF5 file has a plasma where first two ion species
    are Deuterium and Tritium.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        rgrid : array_like (n,1) <br>
            R coordinate grid.
        zgrid : array_like (n,1) <br>
            z coordinate grid.
        ksigrid : array_like (n,1) <br>
            ksi coordinate grid.

    Returns:
        Tuple (rgrid,zgrid,ksigrid,P_DT) where first three are the grid axes
        (same as was given) and fourth element is the number of alpha particles
        (E = 3.5 MeV) born in second in each cell.
    """
    free = False
    if isinstance(a5, str):
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True, plasma=True)

    # Prepare grid
    dr = (rgrid[1] - rgrid[0])
    dz = (zgrid[1] - zgrid[0])
    rvals = rgrid[:-1] + dr / 2
    zvals = zgrid[:-1] + dz / 2
    ksivals = ksigrid[:-1] + (ksigrid[1] - ksigrid[0]) / 2
    R, Z, Ksi = np.meshgrid(rvals, zvals, ksivals, indexing='ij')

    # Evaluate ion temperature and D and T density
    ti   = a5.evaluate(R=R, phi=0, z=Z, t=0, quantity="ti1").reshape(
        rgrid.size-1, zgrid.size-1, ksigrid.size-1) / (1e3*const.e)
    n1   = a5.evaluate(R=R, phi=0, z=Z, t=0, quantity="ni1").reshape(
        rgrid.size-1, zgrid.size-1, ksigrid.size-1)
    n2   = a5.evaluate(R=R, phi=0, z=Z, t=0, quantity="ni2").reshape(
        rgrid.size-1, zgrid.size-1, ksigrid.size-1)

    # Fusion alpha power density (NRL)
    sigmaDT = 3.68e-12 * np.power(ti, -2.0/3) \
              * np.exp(-19.94*np.power(ti, -1.0/3) )
    PDT     = 5.6e-13*n1*n2*sigmaDT / 1e6 # W / m^3

    # Convert power density to total alpha power in each cell
    PDT = PDT * 2*np.pi*R * dr * dz / (ksigrid.size-1) # W/bin

    # Convert power to particle birth rate and replace NaNs with zeros.
    dist = PDT / (3.5e6*const.e)
    dist[np.isnan(dist)] = 0

    if free:
        a5.free(bfield=True, plasma=True)

    if plot:
        plotdist((rgrid, zgrid, ksigrid, dist))

    return (rgrid, zgrid, ksigrid, dist)


def generate(mass, charge, energy, markerdist, particledist, Nmrk,
             minweight=1e-3, plot=False):
    """
    Generate weighted markers from marker and particle distributions.

    Right now only guiding centers are generated but this can be fixed if
    necessary.

    Args:
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        markerdist : tuple <br>
            Tuple (rgrid,zgrid,ksigrid,ordinate) for marker distribution.
        particledist : tuple <br>
            Tuple (rgrid,zgrid,ksigrid,ordinate) for particle distribution.
        Nmrk : int <br>
            Number of markers to be generated.
        minweight : float, optional <br>
            Markers with weight < minweight * average_weight are rejected.

    Returns:
        Guiding center marker data on a dictionary and final marker
        distribution.
    """
    rgrid   = markerdist[0]
    zgrid   = markerdist[1]
    ksigrid = markerdist[2]

    if Nmrk < np.sum(markerdist[3]>0):
        print("More markers are needed to represent this distribution\n")
        return

    # Interpolate particle distribution to marker distribution grid.
    linint = RegularGridInterpolator(
        (particledist[0][:-1] + (particledist[0][1] - particledist[0][0]) / 2,
         particledist[1][:-1] + (particledist[1][1] - particledist[1][0]) / 2,
         particledist[2][:-1] + (particledist[2][1] - particledist[2][0]) / 2),
        particledist[3], method='linear', fill_value=None)

    R, Z, Ksi = np.meshgrid(rgrid[:-1], zgrid[:-1], ksigrid[:-1],
                            indexing='ij')
    R   = R.ravel()
    Z   = Z.ravel()
    Ksi = Ksi.ravel()

    particledist = linint( np.array([
        R   + (rgrid[1]   - rgrid[0]  ) / 2,
        Z   + (zgrid[1]   - zgrid[0]  ) / 2,
        Ksi + (ksigrid[1] - ksigrid[0]) / 2]).transpose() )

    # Turn marker distribution into 1D array that holds cumulative
    # probability P(i).
    check = markerdist[3].copy()
    markerdist = np.append( [0], np.cumsum(markerdist[3].ravel()) )

    # Markers are generated in a sets of Nsample markers. Once the total number
    # of markers is equal or greater than Nmrk, markers are shuffled and the
    # requested number of markers are drawn.
    Nsample = 10000

    indices = np.array([])
    Ntotal = indices.size
    while Ntotal < Nmrk:
        # Draw Nsample uniformly distributed (between 0 and 1) random numbers
        # and choose corresponding cells by finding last index where P is
        # smaller than the drawn number.
        prob = np.random.rand(Nsample)
        idx  = np.searchsorted(markerdist, prob, side='left')-1

        # Get weights for each marker
        weights = particledist[idx]

        # Remove those markers whose weight would be too small.
        idx = idx[weights > minweight * np.nanmean(weights)]

        indices = np.append(indices, idx)
        Ntotal = indices.size

    # Generate markers within these cells. Each marker represents a phase
    # space element enclosed inside the cell, so actual marker coordinates
    # are drawn from uniform distribution where limits are the cell edges.
    indices = indices.astype(int)
    r   = R[indices]   + np.random.rand(Ntotal) * (rgrid[1]   - rgrid[0]  )
    z   = Z[indices]   + np.random.rand(Ntotal) * (zgrid[1]   - zgrid[0]  )
    ksi = Ksi[indices] + np.random.rand(Ntotal) * (ksigrid[1] - ksigrid[0])

    # Make sure there is a marker in each cell
    check0 = check.copy().ravel()
    check0[np.unique(indices)] = 0
    Ncheck = np.sum(check0 > 0)
    if Ncheck > 0:
        Ncheck = np.sum(check > 0)
        r[:Ncheck]   = R[check.ravel() > 0]   \
                       + np.random.rand(Ncheck) * (rgrid[1]   - rgrid[0]  )
        z[:Ncheck]   = Z[check.ravel() > 0]   \
                       + np.random.rand(Ncheck) * (zgrid[1]   - zgrid[0]  )
        ksi[:Ncheck] = Ksi[check.ravel() > 0] \
                       + np.random.rand(Ncheck) * (ksigrid[1]   - ksigrid[0]  )

        ir = np.searchsorted(rgrid, r, side='left')-1
        iz = np.searchsorted(zgrid, z, side='left')-1
        ik = np.searchsorted(ksigrid, ksi, side='left')-1

        indices = (ik + iz*(ksigrid.size-1) + ir*(ksigrid.size-1)*(zgrid.size-1)).astype(int)

    # Weight is divided equally among markers within a given cell
    weights = particledist[indices]
    _, idx, ncount = np.unique(indices, return_inverse=True, return_counts=True)
    weights = weights/ncount[idx]

    # Shuffle and draw Nmrk markers
    ids = np.arange(Ntotal)
    np.random.shuffle(ids)
    r   = r[ids[:Nmrk]]
    z   = z[ids[:Nmrk]]
    ksi = ksi[ids[:Nmrk]]
    weights = weights[ids[:Nmrk]]

    # Store marker data
    mass   = mass / const.physical_constants["atomic mass constant"][0]
    charge = np.around(charge/const.e)
    anum   = np.around(mass)
    znum   = charge
    energy = energy / const.e
    markers = {
        "n"      : Nmrk,
        "ids"    : np.arange(Nmrk) + 1,
        "mass"   : mass * np.ones((Nmrk,1)),
        "charge" : charge * np.ones((Nmrk,1)),
        "r"      : r,
        "phi"    : 360 * np.random.rand(Nmrk),
        "z"      : z,
        "energy" : energy * np.ones((Nmrk,1)),
        "pitch"  : ksi,
        "zeta"   : 2*np.pi * np.random.rand(Nmrk),
        "anum"   : anum * np.ones((Nmrk,1)),
        "znum"   : znum * np.ones((Nmrk,1)),
        "weight" : weights,
        "time"   : 0 * np.ones((Nmrk,1))
    }

    dist = np.histogramdd((r,z,ksi), bins=(rgrid,zgrid,ksigrid),
                          weights=weights)[0]
    dist = (rgrid,zgrid,ksigrid,dist)
    if plot:
        plotdist(dist)

    return (markers, dist)


def plotdist(dist, axes=None):
    """
    Plot distribution in Rz and in ksi.
    """

    if axes is None:
        fig = plt.figure()

        gs = GridSpec(2,1, height_ratios=[3, 1])
        axesrz  = fig.add_subplot(gs[0,0])
        axesksi = fig.add_subplot(gs[1,0])
    else:
        axesrz  = axes[0]
        axesksi = axes[1]

    if axesrz is not None:
        mesh = axesrz.pcolormesh(dist[0], dist[1],
                                 np.sum(dist[3], axis=2).transpose())
        axesrz.set_aspect('equal', adjustable='box')
        plt.colorbar(mesh, ax=axesrz)

    if axesksi is not None:
        axesksi.bar(dist[2][:-1], np.sum(dist[3], axis=(0,1)),
                    align='edge', width=(dist[2][1] - dist[2][0]))

    if axes is None:
        plt.show(block=False)

a5py/a5py/marker/generator.py
ascot4fusion/ascot5
"""
Initialize markers for transport coefficient evaluations and checks.

Markers are initialized at the outer mid-plane. The point is to divide the
phase space in to a grid, and initialize a number of markers at each grid point.
This way one obtains the transport coefficients as a function phase-space at
these grid nodes.
"""

import numpy as np

import a5py.ascot5io.mrk_gc as mrk_gc
import a5py.ascot5io.mrk_fl as mrk_fl

from a5py.ascotpy import Ascotpy


def init_rho(fn, n, rhogrid, rgrid=None, randomize_rho=False, bfield=True,
             desc=None):
    """
    Initialize field line markers in radius.

    Markers are initialized so that at every rho coordinate there are n markers.
    The coordinates are given in rhogrid array at any order. If randomize_rho is
    True, n markers are randomly and uniformly distributed between min and max
    rhogrid values.

    Use rgrid to give R values explicitly in which case rhogrid is ignored.
    """
    rhogrid = np.atleast_1d(rhogrid)

    def getvals(grid):
        if randomize_rho:
            vals = np.amin(grid) + ( np.amax(grid)- np.amin(grid) ) \
                      * np.random.rand(n)
        else:
            vals = np.zeros((grid.size*n,))
            for i in range(grid.size):
                vals[i*n:(i+1)*n] = grid[i]

        return vals

    if rgrid is None:
        rhovals = getvals(rhogrid)
    else:
        rhovals = getvals(rgrid)

    # Find OMP R,z values
    a5 = Ascotpy(fn)
    a5.init(bfield=bfield)
    if rgrid is None:
        rz_omp = a5.get_rhotheta_rz( rhovals, 0, 0, 0 )
    else:
        rz_omp = a5.get_rhotheta_rz( 0.5*np.ones(rhovals.shape), 0, 0, 0 )
        rz_omp[0] = rhovals
    a5.free(bfield=bfield)

    mrk = {}
    mrk["n"]      = rz_omp[0].size
    mrk["ids"]    = np.arange(1, mrk["n"]+1)
    mrk["r"]      = rz_omp[0]
    mrk["z"]      = rz_omp[1]
    mrk["phi"]    = 360*np.random.rand(mrk["n"])
    mrk["pitch"]  = 1 * np.ones((mrk["n"],))
    mrk["time"]   = 0 * np.ones((mrk["n"],))
    mrk["weight"] = 1 * np.ones((mrk["n"],))

    mix = np.random.permutation(mrk["n"])
    mrk["ids"] = mrk["ids"][mix]
    mrk["r"]   = mrk["r"][mix]

    mrk_fl.write_hdf5(fn, **mrk, desc=desc)


def init_rhoenergypitch(fn, n, mass, charge, anum, znum, rhogrid, energygrid,
                        pitchgrid, rgrid=None,
                        randomize_rho=False, randomize_energy=False,
                        randomize_pitch=False, time=0, desc=None):
    """
    Initialize guiding center markers in radius, energy, and pitch.

    See init_rho for how markers are distributed. This function distributes
    markers also in energy and pitch. At each grid point n markers are
    initialized.

    Use rgrid to give R values explicitly in which case rhogrid is ignored.
    """

    rhogrid    = np.atleast_1d(rhogrid)
    energygrid = np.atleast_1d(energygrid)
    pitchgrid  = np.atleast_1d(pitchgrid)

    ntotal = n
    if rgrid is None and not randomize_rho:
        ntotal *= np.unique(rhogrid).size
    if not randomize_energy:
        ntotal *= np.unique(energygrid).size
    if not randomize_pitch:
        ntotal *= np.unique(pitchgrid).size

    if rgrid is not None:
        rgrid = np.atleast_1d(rgrid)
        rhogrid = rgrid
        if not randomize_rho:
            ntotal *= np.unique(rgrid).size

    rhovals    = np.zeros((ntotal,))
    energyvals = np.zeros((ntotal,))
    pitchvals  = np.zeros((ntotal,))

    irho = 0
    while irho < rhogrid.size:
        ienergy = 0
        while ienergy < energygrid.size:
            ipitch = 0
            while ipitch < pitchgrid.size:
                idx =   ipitch * (energygrid.size * rhogrid.size) \
                      + ienergy * rhogrid.size + irho

                if randomize_rho and rgrid is None:
                    rhovals[idx*n:(idx+1)*n] = \
                    np.amin(rhogrid) + ( np.amax(rhogrid)- np.amin(rhogrid) ) \
                        * np.random.rand(n)
                    irho = rhogrid.size
                elif rgrid is None:
                    rhovals[idx*n:(idx+1)*n] = rhogrid[irho]

                if randomize_rho and rgrid is not None:
                    rhovals[idx*n:(idx+1)*n] = \
                    np.amin(rgrid) + ( np.amax(rgrid)- np.amin(rgrid) ) \
                        * np.random.rand(n)
                    irho = rhogrid.size
                elif rgrid is not None:
                    rhovals[idx*n:(idx+1)*n] = rgrid[irho]

                if randomize_energy:
                    energyvals[idx*n:(idx+1)*n] = \
                      np.amin(energygrid) \
                    + ( np.amax(energygrid)- np.amin(energygrid) ) \
                    * np.random.rand(n)
                    ienergy = energygrid.size
                else:
                    energyvals[idx*n:(idx+1)*n] = energygrid[ienergy]

                if randomize_pitch:
                    pitchvals[idx*n:(idx+1)*n] = \
                      np.amin(pitchgrid)
                    + ( np.amax(pitchgrid)- np.amin(pitchgrid) ) \
                    * np.random.rand(n)
                    ipitch = pitchgrid.size
                else:
                    pitchvals[idx*n:(idx+1)*n] = pitchgrid[ipitch]

                ipitch += 1
            ienergy += 1
        irho += 1

    # Find OMP R,z values
    a5 = Ascotpy(fn)
    a5.init(bfield=True)
    if rgrid is None:
        rz_omp = a5.get_rhotheta_rz( rhovals, 0, 0, 0 )
    else:
        rz_omp = a5.get_rhotheta_rz( 0.5*np.ones(rhovals.shape), 0, 0, 0 )
        rz_omp = (rhovals, rz_omp[1])
    a5.free(bfield=True)

    mrk = {}
    mrk["n"]      = rz_omp[0].size
    mrk["ids"]    = np.arange(1, mrk["n"]+1)
    mrk["r"]      = rz_omp[0]
    mrk["z"]      = rz_omp[1]
    mrk["phi"]    = 360*np.random.rand(mrk["n"])
    mrk["zeta"]   = 2*np.pi*np.random.rand(mrk["n"])
    mrk["pitch"]  = pitchvals
    mrk["energy"] = energyvals
    mrk["time"]   = 0 * np.ones((mrk["n"],))
    mrk["weight"] = 1 * np.ones((mrk["n"],))

    mix = np.random.permutation(mrk["n"])
    mrk["ids"]    = mrk["ids"][mix]
    mrk["r"]      = mrk["r"][mix]
    mrk["pitch"]  = mrk["pitch"][mix]
    mrk["energy"] = mrk["energy"][mix]

    mrk["mass"]   = mass.to("amu").d   * np.ones((mrk["n"],))
    mrk["charge"] = charge * np.ones((mrk["n"],))
    mrk["anum"]   = anum   * np.ones((mrk["n"],))
    mrk["znum"]   = znum   * np.ones((mrk["n"],))

    mrk_gc.write_hdf5(fn=fn, **mrk, desc=desc)


def init_rhopparapperp(fn, n, mass, charge, anum, znum, rhogrid, pparagrid,
                       pperpgrid, rgrid=None,
                       randomize_rho=False, randomize_ppara=False,
                       randomize_pperp=False, time=0, desc=None):
    """
    Initialize guiding center markers in radius, ppara, and pperp.

    See init_rho for how markers are distributed. This function distributes
    markers also in ppara and pperp. At each grid point n markers are
    initialized.

    Use rgrid to give R values explicitly in which case rhogrid is ignored.
    """

    rhogrid   = np.atleast_1d(rhogrid)
    pparagrid = np.atleast_1d(pparagrid)
    pperpgrid = np.atleast_1d(pperpgrid)

    ntotal = n
    if rgrid is None and not randomize_rho:
        ntotal *= np.unique(rhogrid).size
    if not randomize_ppara:
        ntotal *= np.unique(pparagrid).size
    if not randomize_pperp:
        ntotal *= np.unique(pperpgrid).size

    if rgrid is not None:
        rgrid = np.atleast_1d(rgrid)
        rhogrid = rgrid
        if not randomize_rho:
            ntotal *= np.unique(rgrid).size

    rhovals   = np.zeros((ntotal,))
    pparavals = np.zeros((ntotal,))
    pperpvals = np.zeros((ntotal,))

    irho = 0
    while irho < rhogrid.size:
        ippara = 0
        while ippara < pparagrid.size:
            ipperp = 0
            while ipperp < pperpgrid.size:
                idx =   ipperp * (pparagrid.size * rhogrid.size) \
                      + ippara * rhogrid.size + irho
                if randomize_rho and rgrid is None:
                    rhovals[idx*n:(idx+1)*n] = \
                    np.amin(rhogrid) + ( np.amax(rhogrid)- np.amin(rhogrid) ) \
                        * np.random.rand(n)
                    irho = rhogrid.size
                elif rgrid is None:
                    rhovals[idx*n:(idx+1)*n] = rhogrid[irho]

                if randomize_rho and rgrid is not None:
                    rhovals[idx*n:(idx+1)*n] = \
                    np.amin(rgrid) + ( np.amax(rgrid)- np.amin(rgrid) ) \
                        * np.random.rand(n)
                    irho = rhogrid.size
                elif rgrid is not None:
                    rhovals[idx*n:(idx+1)*n] = rgrid[irho]

                if randomize_ppara:
                    pparavals[idx*n:(idx+1)*n] = \
                      np.amin(pparagrid) \
                    + ( np.amax(pparagrid)- np.amin(pparagrid) ) \
                    * np.random.rand(n)
                    ippara = pparagrid.size
                else:
                    pparavals[idx*n:(idx+1)*n] = pparagrid[ippara]

                if randomize_pperp:
                    pperpvals[idx*n:(idx+1)*n] = \
                      np.amin(pperpgrid)
                    + ( np.amax(pperpgrid)- np.amin(pperpgrid) ) \
                    * np.random.rand(n)
                    ipperp = pperpgrid.size
                else:
                    pperpvals[idx*n:(idx+1)*n] = pperpgrid[ipperp]

                ipperp += 1
            ippara += 1
        irho += 1

    # Find OMP R,z values
    a5 = Ascotpy(fn)
    a5.init(bfield=True)
    if rgrid is None:
        rz_omp = a5.get_rhotheta_rz( rhovals, 0, 0, 0 )
    else:
        rz_omp = a5.get_rhotheta_rz( 0.5*np.ones(rhovals.shape), 0, 0, 0 )
        rz_omp = (rhovals, rz_omp[1])
    a5.free(bfield=True)

    mass_amu = mass.to("amu").d
    mass_kg  = mass.to("kg").d
    c = 299792458

    mrk = {}
    mrk["n"]      = rz_omp[0].size
    mrk["ids"]    = np.arange(1, mrk["n"]+1)
    mrk["r"]      = rz_omp[0]
    mrk["z"]      = rz_omp[1]
    mrk["phi"]    = 360*np.random.rand(mrk["n"])
    mrk["zeta"]   = 2*np.pi*np.random.rand(mrk["n"])
    mrk["pitch"]  = pparavals / np.sqrt(pparavals**2 + pperpvals**2)
    mrk["time"]   = 0 * np.ones((mrk["n"],))
    mrk["weight"] = 1 * np.ones((mrk["n"],))

    pnorm2 =  pparavals**2 + pperpvals**2
    mrk["energy"] = (np.sqrt(mass_kg**2 + pnorm2 * c**2) - mass_kg) / 1.602e-19
    mrk["mass"]   = mass_amu * np.ones((mrk["n"],))
    mrk["charge"] = charge * np.ones((mrk["n"],))
    mrk["anum"]   = anum   * np.ones((mrk["n"],))
    mrk["znum"]   = znum   * np.ones((mrk["n"],))

    mix = np.random.permutation(mrk["n"])
    mrk["ids"]    = mrk["ids"][mix]
    mrk["r"]      = mrk["r"][mix]
    mrk["pitch"]  = mrk["pitch"][mix]
    mrk["energy"] = mrk["energy"][mix]

    mrk_gc.write_hdf5(fn=fn, **mrk, desc=desc)

a5py/a5py/marker/inittransportpop.py
ascot4fusion/ascot5
"""
Functions to convert user input to physical quantitites in known units.

The functions in this module try to guess in what units input is or what
quantity is being asked for, and give result in correct units. For example,
all following examples give alpha particle mass in kilograms (6.64e-27 kg):

- mass_kg("alpha")
- mass_kg("he4")
- mass_kg(6.64e-27)
- mass_kg(4)

In the last example the function deduces that 4 cannot be an elementary particle
mass and assumes it is in atomic mass units which makes more sense.

File: interpret.py
"""

from scipy.constants import physical_constants as const

kilogramlim = 1e-10
Joulelim    = 1e-5
Coulomblim  = 0.5


def mass_kg(m):
    """
    Convert mass to kilograms.

    Args:
        m : str, float <br>
            Mass
    Returns:
        Mass in kilograms.
    """
    if isinstance(m, str):
        m = m.lower()
        if m == "alpha" or \
           m == "he4":
            return const["alpha particle mass"][0]

        elif m == "proton" or   \
             m == "hydrogen" or \
             m == "p" or        \
             m == "h" or        \
             m == "h1":
            return const["proton mass"][0]

        elif m == "electron" or \
             m == "e":
            return const["electron mass"][0]

        elif m == "deuterium" or \
             m == "d" or         \
             m == "h2":
            return const["deuteron mass"][0]

        elif m == "tritium" or \
             m == "t" or       \
             m == "h3":
            return const["triton mass"][0]

        elif "yo mama":
            return 1.989e30

    else:
        if m > kilogramlim:
            return m * const["atomic mass constant"][0]
        else:
            return m

def mass_amu(m):
    """
    Convert mass to atomic mass units.

    Args:
        m : str, float <br>
            Mass
    Returns:
        Mass in atomic mass units.
    """
    return mass_kg(m) / const["atomic mass constant"][0]

def energy_J(E):
    """
    Convert energy to Joules.

    Args:
        E : str, float <br>
            Energy
    Returns:
        Energy in Joules.
    """
    if E > Joulelim:
        return E * const["elementary charge"][0]
    else:
        return E

def energy_eV(E):
    """
    Convert energy to electronvolts.

    Args:
        E : str, float <br>
            Energy
    Returns:
        Energy in electronvolts.
    """
    return energy_J(E) / const["elementary charge"][0]

def charge_C(q):
    """
    Convert charge to Coulombs.

    Args:
        q : str, float <br>
            charge
    Returns:
        Charge in Coulombs.
    """
    if isinstance(q, str):
        q = q.lower()
        if "elementary" in q:
            return const["elementary charge"][0]

    else:
        if q > Coulomblim:
            return q * const["elementary charge"][0]
        else:
            return q

def charge_e(q):
    """
    Convert charge to multiplies of elementary charge.

    Args:
        q : str, float <br>
            Charge
    Returns:
        Charge in multiplies of elementary charge.
    """
    return charge_C(q) / const["elementary charge"][0]

a5py/a5py/marker/interpret.py
ascot4fusion/ascot5
"""
Evaluate transport coefficients from marker loss-time.

The first-passage time is the time at which a marker crosses a given boundary
for the first time. If the transport is uniform on a infinite interval and the
the process obeys the Fokker-Planck equation, i.e., advection-diffusion
equation, the first-passage time distribution is the so-called inverse Gaussian
distribution.

Launching a delta distribution of markers from a same flux surface, and using
ASCOT5 to calculate the times at which they are lost, allows one to estimate
the transport coefficients via curve fitting.

In most cases the transport is not uniform and the coefficients have a radial
dependency. However, if we assume that the transport is uniform *locally* and
the transport is biased outwards (drift is positive), the radially dependent
coefficients can be solved:

1. Initialize several delta populations at different radial positions.
2. Solve the loss time for each population with ASCOT5.
3. Solve the transport coefficients for each population using the aforementioned
   method. These are not yet the final coefficients.
4. Now consider two first passage time distributions for a radial population i.
   The first one is the first-passage time at the edge, i.e., the loss-time
   which is known (L_i). The second one is the first-passage time (P_i) at the
   radial position where the next (going outwards) delta population is located.
   Because the loss time of the population i+1 is also known, L_{i+1}, we can
   write the following relation: <L_i> = <P_i> + <L_{i+1}> where brackets
   indicate mean value.
5. Solving for the distribution P_i, we can use that to fit the first-passage
   time distribution which gives as the (final) transport coefficients at
   position i. Formally, this is an inverse problem where we try to find
   K_i and D_i (drift and diffusion coefficient) that relation
   L_i = int_{-inf}^{inf} P_i(t-tau; K_i, D_i)*L_{i+1}(tau) dtau
   is fulfilled. If the relation above is confusing, google "convolution of
   probability distributions".
6. For practical purposes, it is easier to use the coefficients calculated at
   step 3 to express L_{i+1} analytically rather than to use the raw data.
"""
import numpy as np
import matplotlib.pyplot as plt

from scipy.optimize import curve_fit
from scipy.special import erf
from scipy.integrate import cumtrapz

def islost(endcond, losstime):
    """
    Small helper function to find markers that are lost: edcond=wall or maxrho.
    """
    return np.logical_and(np.logical_or.reduce([endcond==32, endcond==128,
                                                endcond==1]), losstime >1e-8)

def firstpassagetime(t, K, D, inirho, endrho):
    """
    Calculate first passage time distribution for given parameters.

    Args:
        t : array_like <br>
            Time coordinate [s].
        K : array_like <br>
            Drift coefficient [rho/s].
        D : array_like <br>
            Diffusion coefficient [rho^2/s].
        inirho : array_like <br>
            Initial radial coordinate [rho].
        endrho : array_like <br>
            Coordinate for which the first-passage time is evaluated [rho].
    Returns:
        The distribution f(t).
    """
    # Helper variables
    drho = endrho - inirho
    c1   = drho / K
    c2   = 2*drho**2/D

    temp = np.zeros(t.shape)
    x    = t[t>0]
    temp[t>0] = np.sqrt( c2 / ( 2*np.pi*x**3 ) ) \
                * np.exp( -c2*( x - c1 )**2 / ( 2*c1**2*x ) )
    return temp


def firstpassagetimecum(t, K, D, inirho, endrho):
    """
    Calculate cumulative first passage time distribution for given parameters.

    Args:
        t : array_like <br>
            Time coordinate [s].
        K : array_like <br>
            Drift coefficient [rho/s].
        D : array_like <br>
            Diffusion coefficient [rho^2/s].
        inirho : array_like <br>
            Initial radial coordinate [rho].
        endrho : array_like <br>
            Coordinate for which the first-passage time is evaluated [rho].
    Returns:
        The distribution F(t).
    """
    # Helper variables
    drho = endrho - inirho
    c1   = drho/K
    c2   = drho**2/(2*D)
    t1   =  np.sqrt(c2 / t) * (t/c1 - 1)
    t2   = -np.sqrt(c2 / t) * (t/c1 + 1)
    return 0.5*( 1 + erf( t1 / np.sqrt(2) ) ) \
        + np.exp( 2*c2/c1 ) * 0.5 * ( 1 + erf( t2 / np.sqrt(2)) )


def convolute(t, K, D, inirho, endrho, L2):
    """
    Convolution of two first-passage time distributions.

    The first distribution is evaluated from the given parameters and the second
    one is given as an input. This is because this function is intended for
    inding parameters K and D.

    Args:
        t : array_like <br>
            Time coordinate [s].
        K : array_like <br>
            Drift coefficient of first first-passage time distribution [rho/s].
        D : array_like <br>
            Diffusion coefficient of second first-passage time distribution
            [rho^2/s].
        inirho : array_like <br>
            Initial radial coordinate [rho].
        endrho : array_like <br>
            Coordinate for which the first-passage time is evaluated [rho].
        L2 : array_like <br>
            The second first-passage time distribution on t grid.
    """
    L1 = np.zeros(t.shape)
    for ti in np.arange(t.size):
        P1 = firstpassagetime(t[ti]-t, K, D, inirho, endrho)
        L1[ti] = np.trapz(P1*L2, t)

    temp = cumtrapz(L1, t, initial=0)
    return temp / temp[-1]


def eval_meanvar_coeffs(run=None, inirho=None, losstime=None, endcond=None,
                        endrho=1.0):
    """
    Evaluate coefficients without any fitting just using E[t] and Var[t].

    Args:
        run : RunNode <br>
            Ascot run node container.
        inirho : float, array_like <br>
            Marker initial rho or R (at OMP) coordinate(s).
        losstime : array_like, optional <br>
            Marker final time (including the confined markers).
        endcond : array_like, optional <br>
            Marker end conditions.
        endrho : float, optional <br>
            Radial position for which the passage time is evaluated. Note, use
            units of [m] here if inirho is also in meters.

    Returns:
        Tuple (rho, drift, diff, lost) where rho is the initial rho (or R)
        coordinate, drift is the drift coefficent [rho/s], diff is the diffusion
        coefficient [rho^2/s], and lost is the fraction of lost markers.
    """
    # Init data.
    dataprovided =     inirho   is not None \
                   and losstime is not None \
                   and endcond  is not None
    if run is not None and not dataprovided:
        losstime = run.endstate["mileage"].v
        inirho   = np.mean(run.inistate["rho"].v)
        endcond  = run.endstate["endcond"]
        lost     = np.sum(islost(endcond)) / losstime.size
        losstime = losstime[islost(endcond)]
    elif dataprovided:
        inirho   = np.mean(inirho)
        lost     = np.sum(islost(endcond)) / losstime.size
        losstime = losstime[islost(endcond)]
    else:
        raise Exception(
            "Please provide either run node or the data explicitly (not both).")

    # Return NaN if no losses.
    if lost == 0:
        return (inirho, np.nan, np.nan, lost)

    # Evaluate and return coefficients
    c1   = np.average(losstime)
    var  = np.average((losstime-c1)**2)
    c2   = c1**3/var
    drho = endrho - inirho

    coeffs = np.array([0,0])
    coeffs[0] = drho / c1
    coeffs[1] = 2*drho**2 / c2

    return (inirho, coeffs[0], coeffs[1], lost)


def eval_radindep_coeffs(run=None, inirho=None, losstime=None, endcond=None,
                         endrho=1.0):
    """
    Evaluates coefficients assuming no radial dependency.

    Provide either the run node (coefficients are fitted using all markers whose
    end condition is maxtime) or the requested data (inirho, losstime, endcond)
    explicitly.

    Args:
        run : RunNode <br>
            Ascot run node container.
        inirho : float, array_like <br>
            Marker initial rho or R (at OMP) coordinate(s).
        losstime : array_like, optional <br>
            Marker final time (including the confined markers).
        endcond : array_like, optional <br>
            Marker end conditions.
        endrho : float, optional <br>
            Radial position for which the passage time is evaluated. Note, use
            units of [m] here if inirho is also in meters.

    Returns:
        Tuple (rho, drift, diff, lost) where rho is the initial rho (or R)
        coordinate, drift is the drift coefficent [rho/s], diff is the diffusion
        coefficient [rho^2/s], and lost is the fraction of lost markers.
    """

    # Init data.
    dataprovided =     inirho   is not None \
                   and losstime is not None \
                   and endcond  is not None
    if run is not None and not dataprovided:
        inirho   = np.mean(run.inistate["rho"])
        endcond  = run.endstate["endcond"]
        lost     = np.sum(islost(endcond, losstime)) / losstime.size
        losstime = losstime[islost(endcond, losstime)]
        losstime = np.sort(losstime)
    elif dataprovided:
        inirho   = np.mean(inirho)
        lost     = np.sum(islost(endcond, losstime)) / losstime.size
        losstime = losstime[islost(endcond, losstime)]
        losstime = np.sort(losstime)
    else:
        raise Exception(
            "Please provide either run node or the data explicitly (not both).")

    # Return NaN if no losses.
    if lost == 0:
        return (inirho, np.nan, np.nan, lost)

    # Init fitted function.
    def fitfun(x, a, b):
        """
        Wrapper for firstpassagetimecum().
        """
        return firstpassagetimecum(x, a, b, inirho, endrho)

    # Fit and return the coefficients.
    cumlost = np.linspace(0, 1, losstime.size)
    coeffs, _ = curve_fit(fitfun, losstime, cumlost,
                          p0=[1e0, 1e5], bounds=(0,np.inf))

    return (inirho, coeffs[0], coeffs[1], lost)


def eval_raddep_coeffs(tgrid, rho2, K2, D2, run=None, inirho=None,
                       losstime=None, endcond=None, endrho=1.0):
    """
    Evaluates coefficients assuming a radial dependency.

    Provide either the run node (coefficients are fitted using all markers whose
    end condition is maxtime) or the requested data (inirho, losstime, endcond)
    explicitly.

    Args:
        tgrid : array_like <br>
            Time-grid used to calculate the convolution. Make sure it spans the
            whole regime of possible loss-times and is dense enough. Logarithmic
            grid is likely to provide better results. For example,
            tgrid=10**np.linspace(-6, -2, 1000) has worked well for RE studies.
        rho2 : float <br>
            Radial coordinate (either [rho] or [R]) for the outer point.
        K2 : float <br>
            Drift coefficient for the outer point [rho or m]/[s].
        D2 : float <br>
            Diffusion coefficient for the outer point [rho^2 or m^2]/[s].
        run : RunNode <br>
            Ascot run node container.
        inirho : float, array_like <br>
            Marker initial rho or R (at OMP) coordinate(s).
        losstime : array_like, optional <br>
            Marker final time (including the confined markers).
        endcond : array_like, optional <br>
            Marker end conditions.
        endrho : float, optional <br>
            Radial position for which the passage time is evaluated. Note, use
            units of [m] here if inirho is also in meters.

    Returns:
        Tuple (rho, drift, diff, lost) where rho is the initial rho (or R)
        coordinate, drift is the drift coefficent [rho/s], diff is the diffusion
        coefficient [rho^2/s], and lost is the fraction of lost markers.
    """
    # Init data.
    dataprovided =     inirho   is not None \
                   and losstime is not None \
                   and endcond  is not None
    if run is not None and not dataprovided:
        inirho   = np.mean(run.inistate["rho"].v)
        endcond  = run.endstate["endcond"]
        lost     = np.sum(islost(endcond)) / losstime.size
        losstime = losstime[islost(endcond)]
        losstime = np.sort(run.endstate["mileage"].v)
    elif dataprovided:
        inirho   = np.mean(inirho)
        lost     = np.sum(islost(endcond)) / losstime.size
        losstime = losstime[islost(endcond)]
        losstime = np.sort(losstime)
    else:
        raise Exception(
            "Please provide either run node or the data explicitly (not both).")

    # Return NaN if no losses.
    if lost == 0:
        return (inirho, np.nan, np.nan, lost)

    # Evaluate loss times on a given grid
    losstimeg = np.interp(tgrid, losstime, np.linspace(0, 1, losstime.size),
                         left=0,right=1)
    L2 = firstpassagetime(tgrid, K2, D2, rho2, endrho)

    # Init fitted function.
    def fitfun(x, a, b):
        return convolute(x, a, b, inirho, rho2, L2)

    # Fit and return the coefficients.
    coeffs, _ = curve_fit(fitfun, tgrid, losstimeg,
                          p0=[10000, 10000], bounds=(0,np.inf))
    return (inirho, coeffs[0], coeffs[1], lost)


def eval_coefficients(runs, method="raddep", lastpoint="raddep", gap=1,
                      endrho=1.0, ids=None, tgrid=None, SIunits=False):
    """
    Evaluate coefficients for given radial positions.

    Runs should be sorted so that the outermost run is given last.

    This function assumes that the data for each separate radial location is
    stored in a different run. However, if all data is in a single run, you can
    dublicate that run and use the ids field to control which markers correspond
    to different radial positions.

    The raddep method cannot be used directly to evaluate the last point, which
    means there is a different parameter to control how the last point is
    evaluated if method="raddep".

    Args:
        runs : array_like <br>
            List of runs for different radial points.
        method : str, optional <br>
            Either raddep, radindep, or meanvar.
        lastpoint : str, optional <br>
            What method to use for the last point (raddep means the last point
            is given the same value as the point before that).
        endrho : float, optional <br>
            Coordinate beyond which markers are lost. Give in meters if
            SIunits=True.
        ids : array_like, optional <br>
            ID of the markers for which the coefficients are evaluated.
            Default is all.
        tgrid : array_like, optional <br>
            Time grid needed if the method is raddep. See eval_raddep_coeffs().
        SIunits : bool, optional <br>
            Flag whether coefficients are in units of rho (default) or R.

    Returns:
        Tuple (rho, drift, diff, lost) where rho is the initial rho (or R)
        coordinate, drift is the drift coefficent [rho/s], diff is the diffusion
        coefficient [rho^2/s], and lost is the fraction of lost markers at each
        radial position.
    """

    MEANVAR  = "meanvar"
    RADDEP   = "raddep"
    RADINDEP = "radindep"

    if method not in [MEANVAR, RADDEP, RADINDEP] \
       or lastpoint not in [MEANVAR, RADDEP, RADINDEP]:
        raise Exception(
            "Please provide a valid method and lastpoint" \
            + " (meanvar,raddep,radindep).")

    if method == RADDEP and tgrid is None:
        raise Exception("Please provide tgrid when method=raddep.")

    rho      = np.zeros((len(runs),))
    drift    = np.zeros((len(runs),))
    diff     = np.zeros((len(runs),))
    lostfrac = np.zeros((len(runs),))

    def extractdata(run, ids):
        """
        Small helper for initializing the run data.
        """
        endcond  = run.endstate["endcond"]
        losstime = run.endstate["mileage"].v
        inirho   = run.inistate["rho"].v

        if SIunits:
            inirho = run.inistate["r"]

        if ids is not None:
            inirho   = inirho[ids]
            losstime = losstime[ids]
            endcond  = endcond[ids]

        return (inirho, losstime, endcond)

    if method == MEANVAR:
        for i in range(len(runs)):
            (inirho, losstime, endcond) = extractdata(runs[i], ids[i])

            rho[i], drift[i], diff[i], lostfrac[i] = eval_meanvar_coeffs(
                inirho=inirho, losstime=losstime, endcond=endcond,
                endrho=endrho)

        return (rho, drift, diff, lostfrac)

    # Radindep coefficients are needed for the raddep ones so calculate them.
    for i in range(len(runs)):
        (inirho, losstime, endcond) = extractdata(runs[i], ids[i])
        rho[i], drift[i], diff[i], lostfrac[i] = eval_radindep_coeffs(
            inirho=inirho, losstime=losstime, endcond=endcond,
            endrho=endrho)

    if method == RADINDEP:
        return (rho, drift, diff, lostfrac)

    drift0 = np.copy(drift)
    diff0  = np.copy(diff)

    for i in np.arange(len(runs)-1-gap, -1, -1):
        (inirho, losstime, endcond) = extractdata(runs[i], ids[i])
        rho[i], drift[i], diff[i], lostfrac[i] = eval_raddep_coeffs(
            tgrid=tgrid, rho2=rho[i+gap], K2=drift0[i+gap], D2=diff0[i+gap],
            inirho=inirho, losstime=losstime, endcond=endcond,
            endrho=endrho)

    if lastpoint == RADDEP:
        drift[-1] = drift[-2]
        diff[-1]  = diff[-2]
    elif lastpoint == RADINDEP:
        drift[-1] = drift0[-1]
        diff[-1]  = diff0[-1]
    elif lastpoint == MEANVAR:
        (inirho, losstime, endcond) = extractdata(runs[-1], ids[i])
        rho[-1], drift[-1], diff[-1], lostfrac[-1] = eval_meanvar_coeffs(
            inirho=inirho, losstime=losstime, endcond=endcond,
            endrho=endrho)

    return (rho, drift, diff, lostfrac)


def getplotdata_radindep(run, K, D, endrho=1.0, ids=None):
    def extractdata(run, ids):
        """
        Small helper for initializing the run data.
        """
        endcond  = run.endstate["endcond"]
        losstime = run.endstate["mileage"].v
        inirho   = run.inistate["rho"].v

        if ids is not None:
            inirho   = inirho[ids]
            losstime = losstime[ids]
            endcond  = endcond[ids]

        return (inirho, losstime, endcond)

    (inirho, losstime, endcond) = extractdata(run, ids)
    lf = islost(endcond, losstime)
    t = np.sort(losstime[lf])
    f = firstpassagetimecum(t, K, D, inirho[lf], endrho)

    return (np.log10(t), f, np.linspace(0,1,t.size))

a5py/a5py/marker/losstime.py
ascot4fusion/ascot5
"""
Transformations in phase-space and to/from real-space.

Phase-space transformations include:
- P,mu -> rho_omp, ksi_omp
- P,mu -> rho_omp, mu
and inverse operations. Here P is canonical toroidal momentum, mu is magnetic
moment, rho_omp is outer mid plane rho and ksi_omp is pitch at outer mid plane.
To clarify, rho_omp and ksi_omp are the values a marker with given P and mu
would have when crossing OMP. Since trapped particles have two crossing points,
the one where ksi_omp has the same sign as the current pitch is chosen. The sign
of pitch is transmitted as a sign of mu, i.e. mu = sign(ksi) * magnetic moment.

The real-space transformation is P,mu -> R,z,ksi and its inverse.

All parameters in this module are in SI units. Ascotpy is required and a file
with a magnetic field. Functions accept either pre-initialized Ascotpy object
or name of the HDF5 file, in which case magnetic field module is initialized
and freed inside the function.

This module is for tokamaks only.

File: transport/ppmappings.py
"""
import numpy as np
import scipy.constants as const
import importlib.util as util

from unyt import c as speed_of_light

plt = util.find_spec("matplotlib")
if plt:
    import matplotlib.pyplot as plt


def evalpnorm(mass, energy):
    """
    Small helper function to evaluate pnorm.
    """
    gamma = 1 + energy / (mass * speed_of_light.v**2)
    return np.sqrt(gamma**2 - 1) * mass * speed_of_light.v


def evalPmu(a5, mass, charge, energy, r, z, ksi):
    """
    Evaluate P and mu at the given coordinates.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        r : array_like (n,1) <br>
            Major radius coordinates [m].
        z : array_like (n,1) <br>
            z coordinates [m].
        ksi : array_like (n,1) <br>
            Pitch coordinates (v_para/v_tot).

    Returns:
        Tuple (P,mu) where P and mu are (n,1) arrays ([kg*m/s], [J/T]).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    psi   = a5.evaluate(r, 0, z, 0, "psi")
    bphi  = a5.evaluate(r, 0, z, 0, "bphi")
    bnorm = a5.evaluate(r, 0, z, 0, "bnorm")
    pnorm = evalpnorm(mass, energy)

    mu = np.sign(ksi) * ( 1 - ksi**2 ) * pnorm**2 / ( 2 * bnorm * mass )
    P  = r * pnorm * ksi * bphi / bnorm + charge * psi

    if free:
        a5.free(bfield=True)

    return (P, mu)


def istrapped(a5, mass, charge, energy, P, mu, rmin=None):
    """
    Check if a marker with given coordinates is trapped.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        P : array_like (n,1) <br>
            Canonical toroidal momentum coordinates [kg*m/s].
        mu : array_like (n,1) <br>
            Magnetic moment coordinates [J/T].
        rlim : tuple or None <br>
            Major grid minimum value that markers may have. If None, ther value
            is picked at inner midplane separatrix.

    Returns:
        Boolean array with True elements for coordinates that are for trapped
        markers.
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    if rmin is None:
        rz    = a5.get_rhotheta_rz( np.array([0,1]), 180, 0, 0 )
        rmin = rz[0][-1]

    # Make R grid for evaluating bnorm and psi.
    nr = 100
    axis  = a5.evaluate(0, 0, 0, 0, "axis")
    rgrid = np.linspace(rmin, axis["axisr"], nr).ravel()

    bphi  = a5.evaluate(rgrid, 0, axis["axisz"], 0, "bphi").ravel()
    bnorm = a5.evaluate(rgrid, 0, axis["axisz"], 0, "bnorm").ravel()
    psi   = a5.evaluate(rgrid, 0, axis["axisz"], 0, "psi").ravel()
    pnorm = evalpnorm(mass, energy)

    # Find which (P,mu) coordinates correspond to trapped markers. A marker is
    # trapped if a) mu is imaginary for all r < raxis or b) P(r,mu) - P has a
    # root for rmin < r <raxis.
    trapped = np.zeros(mu.shape) == 1
    for i in range(mu.size):
        ppar2  = pnorm*pnorm - 2 * np.absolute(mu[i]) * bnorm * mass
        if all(ppar2 < 0):
            # mu is imaginary for all r < raxis. Therefore marker orbit has no
            # solution there which means the marker is trapped.
            trapped[i] = True
            continue

        # Evaluate P(r,mu).  Note that np.abs() is used just to suppress warnings
        val = charge * psi + np.sign(mu[i]) * rgrid * np.sqrt(np.abs(ppar2)) \
              * bphi / bnorm
        val[ppar2 < 0] = np.nan

        if not (any(val - P[i] > 0) and any(val - P[i] < 0)):
            # No root for P(r,mu) - P, marker does not pass IMP and hence is
            # trapped
            trapped[i] = True

    if free:
        a5.free(bfield=True)

    return trapped


def initgridPmu(a5, mass, charge, energy, nP, nmu, padding):
    """
    Initialize (P,mu) grid that covers the whole phase-space.

    This function is useful as both P and mu depend not only on the marker
    parameters, but also on the magnetic background. Here values for P and mu
    are evaluated at OMP, and the extreme values (with some padding) are used
    to define the grid.

    mugrid is given in interval [-mumax, mumax].

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        nP : int <br>
            Number of P grid points.
        nmu : int <br>
            Number of mu grid points.
        padding : array_like (n,1) <br>
            How many percent grids are stretched beyond the extreme values.

    Returns:
        Tuple (Pgrid [kg*m/s], mugrid [J/T]).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    # OMP rz coordinates
    rz = a5.get_rhotheta_rz( np.array([0,1]), 0, 0, 0 )

    # mumax can be found assuming pperp = pnorm and using OMP separatrix value
    # for bmin
    bmin = a5.evaluate(rz[0][-1], 0, rz[1][-1], 0, "bnorm")
    pnorm = evalpnorm(mass, energy)

    mumax = pnorm**2 / ( 2 * bnorm * mass )
    mumin = -mumax

    mugrid = np.linspace(mumin*(1+padding), mumax*(1+padding), nmu)

    # Pmax and min are located at vpar = +-vnorm and psi = psimin or psimax.
    psi   = a5.evaluate(rz[0], 0, rz[1], 0, "psi")

    if psi[0] < psi[1]:
        psimin = psi[0]
        rmin   = rz[0][0]
        psimax = psi[1]
        rmax   = rz[0][1]
    else:
        psimin = psi[1]
        rmin   = rz[0][1]
        psimax = psi[0]
        rmax   = rz[0][0]

    bpmin = a5.evaluate(rmin, 0, rz[1][-1], 0, "bphi") / \
            a5.evaluate(rmin, 0, rz[1][-1], 0, "bnorm")
    bpmax = a5.evaluate(rmax, 0, rz[1][-1], 0, "bphi") / \
            a5.evaluate(rmax, 0, rz[1][-1], 0, "bnorm")

    P  = np.array([ rmin * pnorm * bpmin + charge * psimin,
                   -rmin * pnorm * bpmin + charge * psimin,
                    rmax * pnorm * bpmax + charge * psimax,
                   -rmax * pnorm * bpmax + charge * psimax])

    Pmin = np.nanmin(P)
    Pmax = np.nanmax(P)

    if Pmin < 0:
        Pmin = Pmin*(1+padding)
    else:
        Pmin = Pmin*(1-padding)

    if Pmax < 0:
        Pmax = Pmax*(1-padding)
    else:
        Pmax = Pmax*(1+padding)

    Pgrid = np.linspace(Pmin, Pmax, nP)

    if free:
        a5.free(bfield=True)

    return (Pgrid, mugrid)


def maprhomu2Pmu(a5, mass, charge, energy, rhovals, muvals,
                 Pgrid, mugrid, weights=None):
    """
    Map (rho_omp,mu) coordinates to (P,mu) grid.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        rhovals : array_like (n,1) <br>
            Marker rho_omp values.
        muvals : array_like (n,1) <br>
            Marker mu values (including sign).
        Pgrid : array_like (n,1) <br>
            P coordinate grid.
        mugrid : array_like (n,1) <br>
            mu coordinate grid.
        weights : array_like (n,1), optional <br>
            Weights of markers. Uniform weights if None.

    Returns:
        Tuple (ordinate,Pvals,muvals) where ordinate is the resulting histogram
        in (Pgrid,mugrid) grid and Pvals, and muvals are the exact (P,mu) values
        corresponding to (rhovals,muvals).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    rz    = a5.get_rhotheta_rz( rhovals, 0, 0, 0 )
    psi   = a5.evaluate(rz[0], 0, rz[1], 0, "psi")
    bnorm = a5.evaluate(rz[0], 0, rz[1], 0, "bnorm")
    bphi  = a5.evaluate(rz[0], 0, rz[1], 0, "bphi")
    pnorm = evalpnorm(mass, energy)

    ppar   = np.sign(muvals) * np.sqrt(pnorm*pnorm
                                       - 2*bnorm*np.absolute(muvals)*mass)
    Pvals  = rz[0] * ppar *bphi / bnorm + charge * psi

    if weights is None:
        weights = np.ones(muvals.shape)

    ordinate = np.histogram2d(Pvals, muvals,
                              bins=[Pgrid.ravel(), mugrid.ravel()],
                              weights=weights)[0]

    if free:
        a5.free(bfield=True)

    return (ordinate, Pvals, muvals)


def mapPmu2rhomu(a5, mass, charge, energy, Pvals, muvals, rhogrid, mugrid,
                 weights=None):
    """
    Map (P,mu) coordinates to (rho_omp,mu) grid.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        Pvals : array_like (n,1) <br>
            Marker P values.
        muvals : array_like (n,1) <br>
            Marker mu values (including sign).
        rhogrid : array_like (n,1) <br>
            rho_omp coordinate grid.
        mugrid : array_like (n,1) <br>
            mu coordinate grid.
        weights : array_like (n,1), optional <br>
            Weights of markers. Uniform weights if None.

    Returns:
        Tuple (ordinate,rhovals,muvals) where ordinate is the resulting
        histogram in (rhogrid,mugrid) grid and rhovals and muvals are values
        each Pvals,muvals pair correspond to (center value of the bin they
        belong to).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    rz    = a5.get_rhotheta_rz( rhogrid, 0, 0, 0 )
    psi   = a5.evaluate(rz[0], 0, rz[1], 0, "psi")
    bnorm = a5.evaluate(rz[0], 0, rz[1], 0, "bnorm")
    bphi  = a5.evaluate(rz[0], 0, rz[1], 0, "bphi")
    pnorm = evalpnorm(mass, energy)

    mu    = mugrid.transpose()
    psi   = psi.reshape(-1,1) * np.ones(mu.shape)
    bnorm = bnorm.reshape(-1,1) * np.ones(mu.shape)
    bphi  = bphi.reshape(-1,1) * np.ones(mu.shape)
    r     = rz[0].reshape(-1,1) * np.ones(mu.shape)
    ppar2 = pnorm * pnorm - 2 * bnorm * np.absolute(mu) * mass
    ppar  = np.sign(mu)*np.sqrt(ppar2)

    P  = r * ppar * bphi / bnorm + charge * psi

    if weights is None:
        weights = np.ones(muvals.shape)

    rhovals   = np.ones(Pvals.shape)
    newmuvals = np.ones(Pvals.shape)
    ordinate  = np.zeros( (rhogrid.size-1, mugrid.size-1) )
    for i in range(rhogrid.size-1):
        for j in range(mugrid.size-1):
            a = np.array([ P[i,j], P[i+1,j] ])

            ids = np.logical_and.reduce( [
                Pvals >= np.nanmin(a), Pvals < np.nanmax(a),
                muvals >= mugrid[j], muvals < mugrid[j+1]
            ] )

            rhovals[ids]   = (rhogrid[i] + rhogrid[i+1]) / 2
            newmuvals[ids] = (mugrid[j] + mugrid[j+1]) / 2
            ordinate[i,j]   = np.nansum(weights[ids])

    if free:
        a5.free(bfield=True)

    return (ordinate, rhovals, newmuvals)


def test_rhomu2Pmu(fn):
    """
    Verify rho_omp-mu & P-mu transformation.

    This function populates (rho,mu) mesh uniformly, transforms it to (P,mu)
    mesh, and transforms the resulting values back to (rho,mu) grid. Therefore
    the initial and final histograms should resemble each other. Process is
    repeated starting from uniform (P,mu) population. Final and intermediate
    histograms are plotted.

    Args:
        fn : str <br>
            Filename for ascot HDF5 file which must have magnetic field present.
    """
    from a5py.ascot5io.ascot5 import Ascot
    from a5py.ascotpy.ascotpy import Ascotpy
    h5 = Ascot(fn)
    a5 = Ascotpy(fn)
    a5.init(bfield=True)

    # Perform check for alpha particles
    mass   = const.physical_constants["alpha particle mass"][0]
    energy = 3.5e6 * const.e
    charge = 2 * const.e

    rhogrid = np.linspace( 0, 1, 100)

    nP  = 100
    nmu = 100
    try:
        fig  = plt.figure()
        Pgrid, mugrid = initgridPmu(a5, mass, charge, energy,
                                    nP, nmu, padding=1e-3)

        # Convert mesh midpoints to value pairs, transform, and plot
        rhovals, muvals = np.meshgrid(
            rhogrid[:-1] + (rhogrid[1] - rhogrid[0]) / 2,
            mugrid[:-1] + (mugrid[1] - mugrid[0]) / 2,
            indexing='ij')

        weights = np.ones(rhovals.shape)
        weights, Pvals, muvals = maprhomu2Pmu(a5, mass, charge, energy,
                                              rhovals.ravel(), muvals.ravel(),
                                              Pgrid, mugrid,
                                              weights=weights.ravel())

        axes = fig.add_subplot(2,2,1)
        h = axes.pcolormesh(Pgrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        # Repeat
        weights = np.ones(Pvals.shape)
        weights = mapPmu2rhomu(a5, mass, charge, energy,
                               Pvals.ravel(), muvals.ravel(),
                               rhogrid, mugrid, weights=weights.ravel())[0]

        axes = fig.add_subplot(2,2,2)
        h = axes.pcolormesh(rhogrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        # Now do the whole stuff otherway around
        Pvals, muvals = np.meshgrid(
            Pgrid[:-1] + (Pgrid[1] - Pgrid[0]) / 2,
            mugrid[:-1] + (mugrid[1] - mugrid[0]) / 2,
            indexing='ij')

        weights = np.ones(Pvals.shape)
        weights, rhovals, muvals = mapPmu2rhomu(a5, mass, charge, energy,
                                                Pvals.ravel(), muvals.ravel(),
                                                rhogrid, mugrid,
                                                weights=weights.ravel())

        axes = fig.add_subplot(2,2,3)
        h = axes.pcolormesh(rhogrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)


        # And repeat
        weights = maprhomu2Pmu(a5, mass, charge, energy,
                               rhovals.ravel(), muvals.ravel(),
                               Pgrid, mugrid, weights=weights.ravel())[0]

        axes = fig.add_subplot(2,2,4)
        h = axes.pcolormesh(Pgrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        plt.show(block=False)

    except Exception as e:
        a5.free(bfield=True)
        raise e

    a5.free(bfield=True)


def maprhoksi2Pmu(a5, mass, charge, energy, rhovals, ksivals,
                  Pgrid, mugrid, weights=None):
    """
    Map (rho_omp,rho_ksi) coordinates to (P,mu) grid.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        rhovals : array_like (n,1) <br>
            Marker rho_omp values.
        ksivals : array_like (n,1) <br>
            Marker pitch values.
        Pgrid : array_like (n,1) <br>
            P coordinate grid.
        mugrid : array_like (n,1) <br>
            mu coordinate grid.
        weights : array_like (n,1), optional <br>
            Weights of markers. Uniform weights if None.

    Returns:
        Tuple (ordinate,Pvals,muvals) where ordinate is the resulting histogram
        in (Pgrid,mugrid) grid and Pvals, and muvals are the exact (P,mu) values
        corresponding to (rhovals,ksivals).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    ksivals = ksivals
    weights = weights
    rhovals = rhovals

    rz    = a5.get_rhotheta_rz( rhovals, 0, 0, 0 )
    psi   = a5.evaluate(rz[0], 0, rz[1], 0, "psi")
    bnorm = a5.evaluate(rz[0], 0, rz[1], 0, "bnorm")
    bphi  = a5.evaluate(rz[0], 0, rz[1], 0, "bphi")
    pnorm = evalpnorm(mass, energy)

    muvals = (1 - 2*(ksivals<0)) * (1 - ksivals*ksivals) \
             * pnorm * pnorm / ( 2 * bnorm * mass )
    Pvals  = rz[0] * pnorm * ksivals * bphi / bnorm + charge * psi

    if weights is None:
        weights = np.ones(rhovals.shape)

    ordinate = np.histogram2d(Pvals,
                              muvals,
                              bins=[Pgrid.ravel(), mugrid.ravel()],
                              weights=weights)[0]

    if free:
        a5.free(bfield=True)

    return (ordinate, Pvals, muvals)


def mapPmu2rhoksi(a5, mass, charge, energy, Pvals, muvals, rhogrid, ksigrid,
                  weights=None):
    """
    Map (P,mu) coordinates to (rho_omp,ksi_omp) grid.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        Pvals : array_like (n,1) <br>
            Marker P values.
        muvals : array_like (n,1) <br>
            Marker mu values (including sign).
        rhogrid : array_like (n,1) <br>
            rho_omp coordinate grid.
        ksigrid : array_like (n,1) <br>
            Pitch coordinate grid.
        weights : array_like (n,1), optional <br>
            Weights of markers. Uniform weights if None.

    Returns:
        Tuple (ordinate,rhovals,ksivals) where ordinate is the resulting
        histogram in (rhogrid,ksigrid) grid and rhovals and ksivals are values
        each Pvals,muvals pair correspond to (center value of the bin they
        belong to).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    rz    = a5.get_rhotheta_rz( rhogrid, 0, 0, 0 )
    psi   = a5.evaluate(rz[0], 0, rz[1], 0, "psi")
    bnorm = a5.evaluate(rz[0], 0, rz[1], 0, "bnorm")
    bphi  = a5.evaluate(rz[0], 0, rz[1], 0, "bphi")
    pnorm = evalpnorm(mass, energy)

    ksi   = ksigrid.transpose()
    psi   = psi.reshape(-1,1) * np.ones(ksi.shape)
    bnorm = bnorm.reshape(-1,1) * np.ones(ksi.shape)
    bphi  = bphi.reshape(-1,1) * np.ones(ksi.shape)
    r     = rz[0].reshape(-1,1) * np.ones(ksi.shape)
    ksi   = np.ones(rz[0].reshape(-1,1).shape).reshape(-1,1)*ksigrid.transpose()
    ppar  = ksi*pnorm

    P  = r * ppar * bphi / bnorm + charge * psi
    mu = (1 - 2*(ksi < 0)) * (1 - ksi*ksi) * pnorm * pnorm / ( 2 * bnorm * mass )
    mumax = pnorm * pnorm / ( 2 * bnorm * mass )

    if weights is None:
        weights = np.ones(muvals.shape)

    rhovals = np.ones(Pvals.shape) + np.nan
    ksivals = np.ones(Pvals.shape) + np.nan
    ordinate = np.zeros( (rhogrid.size-1, ksigrid.size-1) )
    for i in range(rhogrid.size-1):
        for j in range(ksigrid.size-1):
            a = np.array([ P[i,j], P[i+1,j], P[i,j+1], P[i+1,j+1] ])
            b = np.array([ mu[i,j], mu[i,j+1], mu[i,j+1], mu[i+1,j+1]])
            mumaxval = np.nanmax([mumax[i,j], mumax[i+1,j]])

            ids = np.logical_and.reduce( [
                Pvals >= np.nanmin(a), Pvals < np.nanmax(a)
            ] )

            if ksigrid.size % 2 == 0 and j == ksigrid.size/2 - 1:
                a1 = np.logical_and.reduce( [
                    muvals < np.nanmax([mu[i,j], mu[i+1,j]]),
                    muvals >= -mumaxval ] )
                a2 = np.logical_and.reduce( [
                    muvals >= np.nanmin([mu[i,j+1], mu[i+1,j+1]]),
                    muvals <=  mumaxval ] )
                muids = np.logical_or.reduce([a1, a2])
            elif ksigrid.size % 2 == 1 and j == (ksigrid.size-1)/2-1:
                muids = np.logical_and.reduce(
                    [ muvals <  np.nanmax([mu[i,j], mu[i+1,j]]),
                      muvals >= -mumaxval] )
            elif ksigrid.size % 2 == 1 and j == (ksigrid.size-1)/2:
                muids = np.logical_and.reduce(
                    [ muvals >= np.nanmin([mu[i,j+1], mu[i+1,j+1]]),
                      muvals <= mumaxval] )
            else:
                muids = np.logical_and.reduce( [ muvals >= np.nanmin(b),
                                                 muvals < np.nanmax(b) ] )

            ids = np.logical_and.reduce( [ids, muids] )

            rhovals[ids] = (rhogrid[i] + rhogrid[i+1]) / 2
            ksivals[ids] = (ksigrid[j] + ksigrid[j+1]) / 2
            ordinate[i,j] = np.nansum(weights[ids])


    if free:
        a5.free(bfield=True)

    return (ordinate, rhovals, ksivals)


def test_rhoksi2Pmu(fn):
    """
    Verify rho_omp-ksi_omp & P-mu transformation.

    This function populates (rho,ksi) mesh uniformly, transforms it to (P,mu)
    mesh, and transforms the resulting values back to (rho,ksi) grid. Therefore
    the initial and final histograms should resemble each other. Process is
    repeated starting from uniform (P,mu) population. Final and intermediate
    histograms are plotted.

    Args:
        fn : str <br>
            Filename for ascot HDF5 file which must have magnetic field present.
    """
    from a5py.ascot5io.ascot5 import Ascot
    from a5py.ascotpy.ascotpy import Ascotpy
    h5 = Ascot(fn)
    a5 = Ascotpy(fn)
    a5.init(bfield=True)

    mass   = const.physical_constants["alpha particle mass"][0]
    energy = 3.5e6 * const.e
    charge = 2 * const.e

    rhogrid = np.linspace( 0, 1, 200)
    ksigrid = np.linspace(-1, 1, 50)

    nP  = 200
    nmu = 200
    try:
        fig  = plt.figure()
        Pgrid, mugrid = initgridPmu(a5, mass, charge, energy,
                                    nP, nmu, padding=1e-3)

        # Convert mesh midpoints to value pairs, transform, and plot
        rhovals, ksivals = np.meshgrid(
            rhogrid[:-1] + (rhogrid[1] - rhogrid[0]) / 2,
            ksigrid[:-1] + (ksigrid[1] - ksigrid[0]) / 2,
            indexing='ij')
        weights = np.ones(rhovals.shape)
        weights, Pvals, muvals = maprhoksi2Pmu(a5, mass, charge, energy,
                                               rhovals.ravel(), ksivals.ravel(),
                                               Pgrid, mugrid,
                                               weights=weights.ravel())

        axes = fig.add_subplot(2,2,1)
        h = axes.pcolormesh(Pgrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        # Convert back to rho_omp, ksi_omp space and plot
        weights = np.ones(Pvals.shape)
        weights, rhovals, ksivals = mapPmu2rhoksi(a5, mass, charge, energy,
                                Pvals.ravel(), muvals.ravel(),
                                rhogrid, ksigrid, weights=weights.ravel())

        axes = fig.add_subplot(2,2,2)
        h = axes.pcolormesh(rhogrid, ksigrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        # Repeat the process for a uniform (P,mu) population
        Pvals, muvals = np.meshgrid(
            Pgrid[:-1] + (Pgrid[1] - Pgrid[0]) / 2,
            mugrid[:-1] + (mugrid[1] - mugrid[0]) / 2,
            indexing='ij')
        weights = np.ones(Pvals.shape)
        weights, rhovals, ksivals = mapPmu2rhoksi(a5, mass, charge, energy,
                                                  Pvals.ravel(), muvals.ravel(),
                                                  rhogrid, ksigrid,
                                                  weights=weights.ravel())

        axes = fig.add_subplot(2,2,3)
        h = axes.pcolormesh(rhogrid, ksigrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        # And transform back again and plot
        rhovals = rhovals[np.isnan(rhovals) == False]
        ksivals = ksivals[np.isnan(ksivals) == False]
        weights = np.ones(rhovals.shape)
        weights = maprhoksi2Pmu(a5, mass, charge, energy,
                                rhovals.ravel(), ksivals.ravel(),
                                Pgrid, mugrid, weights=weights.ravel())[0]

        axes = fig.add_subplot(2,2,4)
        h = axes.pcolormesh(Pgrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        plt.show(block=False)

    except Exception as e:
        a5.free(bfield=True)
        raise e

    a5.free(bfield=True)


def maprzk2Pmu(a5, mass, charge, energy, rvals, zvals, ksivals,
               Pgrid, mugrid, weights=None):
    """
    Map (rho_omp,mu) coordinates to (P,mu) grid.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        rvals : array_like (n,1) <br>
            Marker R values.
        zvals : array_like (n,1) <br>
            Marker z values.
        ksivals : array_like (n,1) <br>
            Marker pitch values.
        Pgrid : array_like (n,1) <br>
            P coordinate grid.
        mugrid : array_like (n,1) <br>
            mu coordinate grid.
        weights : array_like (n,1), optional <br>
            Weights of markers. Uniform weights if None.

    Returns:
        Tuple (ordinate,Pvals,muvals) where ordinate is the resulting histogram
        in (Pgrid,mugrid) grid and Pvals, and muvals are the exact (P,mu) values
        corresponding to (rvals,zvals,ksivals).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    bnorm = a5.evaluate(rvals, 0, zvals, 0, "bnorm")
    bphi  = a5.evaluate(rvals, 0, zvals, 0, "bphi")
    psi   = a5.evaluate(rvals, 0, zvals, 0, "psi")
    pnorm = evalpnorm(mass, energy)

    muvals = (1-2*(ksivals<0)) * (1-ksivals*ksivals) \
             * pnorm * pnorm / ( 2 * bnorm * mass )
    Pvals  = rvals * ksivals * pnorm * bphi / bnorm + charge * psi

    if weights is None:
        weights = np.ones(muvals.shape)

    ordinate = np.histogram2d(Pvals, muvals,
                              bins=[Pgrid.ravel(), mugrid.ravel()],
                              weights=weights)[0]

    if free:
        a5.free(bfield=True)

    return (ordinate, Pvals, muvals)


def mapPmu2rzk(a5, mass, charge, energy, Pvals, muvals,
               rgrid, zgrid, ksigrid, weights=None):
    """
    Map (P,mu) coordinates to (R,z,ksi) grid.

    This mapping is not one-to-one as single (P,mu) value corresponds to a whole
    trajectory in (R,z,ksi) real-space. Therefore, this function finds all
    (R,z,ksi) cells a (P,mu) correspond to, and divides the weight evenly among
    them.

    Args:
        a5 : Ascotpy, str <br>
            Ascotpy object or HDF5 filename.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        Pvals : array_like (n,1) <br>
            Marker P values.
        muvals : array_like (n,1) <br>
            Marker mu values (including sign).
        rgrid : array_like (n,1) <br>
            R coordinate grid.
        zgrid : array_like (n,1) <br>
            z coordinate grid.
        ksigrid : array_like (n,1) <br>
            ksi coordinate grid.
        weights : array_like (n,1), optional <br>
            Weights of markers. Uniform weights if None.

    Returns:
        Ordinate. (R,z,ksi) values are not returned as they are not unique.
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    pnorm = evalpnorm(mass, energy)
    if weights is None:
        weights = np.ones(muvals.shape)

    Pvals   = Pvals.ravel()
    muvals  = muvals.ravel()
    weights = weights.ravel()

    # Evaluate what P and mu are at each Rzk-grid node
    R, Z, Ksi = np.meshgrid(rgrid, zgrid, ksigrid, indexing='ij')
    psi   = a5.evaluate(R, 0, Z, 0, "psi").reshape(R.shape)
    bnorm = a5.evaluate(R, 0, Z, 0, "bnorm").reshape(R.shape)
    bphi  = a5.evaluate(R, 0, Z, 0, "bphi").reshape(R.shape)
    Mu = (1-2*(Ksi<0)) * (1-Ksi*Ksi) * pnorm * pnorm / ( 2 * bnorm * mass )
    P  = R * Ksi * pnorm * bphi / bnorm  + charge * psi
    #Pmid = charge * psi

    # Maximum possible mu value at each node
    Mumax = pnorm * pnorm / ( 2 * bnorm * mass )

    # The index search algorithm, which finds 8R,z,ksi) cells (P,mu) correspond
    # to, is run twice: first to find how many cells each (P,mu) correspond to
    # so that their weights can be dividec accordingly, and then to distribute
    # the weight among the cells.
    ncells  = np.zeros(weights.shape)
    for i in range(rgrid.size-1):
        for j in range(zgrid.size-1):
            for k in range(ksigrid.size-1):
                # Check which (Pvals, muvals) values are within this cell by
                # finding those for which Mu_a < mu < Mu_b and P_c < P < P_d
                # where a, b, c, and d are some grid nodes.

                # However, there is a caveat. When ksi passes zero, mu is
                # discontinuous. Therefore we separate those ksi grid points
                # (k and k+1) so we can use Mu_a < mu < Mumax and
                # -Mumax < mu < Mu_b where ksi(k) < 0 and ksi(k+1) > 0.
                vals0 = [Mu[i,j,k], Mu[i+1,j,k], Mu[i,j+1,k], Mu[i+1,j+1,k]]
                vals1 = [Mu[i,j,k+1], Mu[i+1,j,k+1], Mu[i,j+1,k+1],
                         Mu[i+1,j+1,k+1]]

                vals2 = [P[i,j,k], P[i+1,j,k], P[i,j+1,k], P[i+1,j+1,k]]
                vals3 = [P[i,j,k+1], P[i+1,j,k+1], P[i,j+1,k+1], P[i+1,j+1,k+1]]

                Mumaxval = np.nanmax([Mumax[i,j,k], Mumax[i+1,j,k],
                                      Mumax[i,j+1,k], Mumax[i+1,j+1,k]])

                if ksigrid.size % 2 == 0 and k == ksigrid.size/2-1:
                    a1 = np.logical_and.reduce( [ muvals <= np.nanmin(vals0),
                                                  muvals >  -Mumaxval ] )
                    a2 = np.logical_and.reduce( [ muvals <= Mumaxval,
                                                  muvals > np.nanmin(vals1) ] )
                    a = np.logical_or.reduce([a1, a2])

                elif ksigrid.size % 2 == 1 and k == (ksigrid.size-1)/2-1 :
                    a = np.logical_and.reduce( [ muvals < np.nanmax(vals0),
                                                 muvals >= -Mumaxval ] )

                elif ksigrid.size % 2 == 1 and k == (ksigrid.size-1)/2 :
                    a = np.logical_and.reduce( [ muvals >= np.nanmin(vals1),
                                                 muvals <= Mumaxval ] )
                else:
                    a = np.logical_and.reduce(
                        [ muvals >= np.nanmin(np.append(vals0,vals1)),
                          muvals < np.nanmax(np.append(vals0,vals1)) ] )

                b = np.logical_and.reduce(
                    [Pvals >= np.nanmin(np.append(vals2,vals3)),
                     Pvals < np.nanmax(np.append(vals2,vals3)) ])

                # Add this cell to all (P,mu) values that it holds
                ids = np.logical_and.reduce([a, b])
                if np.sum(ids.ravel()) > 0:
                    ncells[ids] += 1

    # Update the weights and run the algorithm again
    weights = weights/ncells
    ordinate = np.zeros( (rgrid.size-1,zgrid.size-1,ksigrid.size-1) )
    for i in range(rgrid.size-1):
        for j in range(zgrid.size-1):
            for k in range(ksigrid.size-1):
                vals0 = [Mu[i,j,k], Mu[i+1,j,k], Mu[i,j+1,k], Mu[i+1,j+1,k]]
                vals1 = [Mu[i,j,k+1], Mu[i+1,j,k+1], Mu[i,j+1,k+1],
                         Mu[i+1,j+1,k+1]]

                vals2 = [P[i,j,k], P[i+1,j,k], P[i,j+1,k], P[i+1,j+1,k]]
                vals3 = [P[i,j,k+1], P[i+1,j,k+1], P[i,j+1,k+1], P[i+1,j+1,k+1]]

                Mumaxval = np.nanmax([Mumax[i,j,k], Mumax[i+1,j,k],
                                      Mumax[i,j+1,k], Mumax[i+1,j+1,k]])

                #if ksigrid.size % 2 == 1 and k == (ksigrid.size-1)/2-1 :
                if ksigrid.size % 2 == 0 and k == ksigrid.size/2-1:
                    a1 = np.logical_and.reduce( [ muvals <= np.nanmin(vals0),
                                                  muvals >=-Mumaxval ] )
                    a2 = np.logical_and.reduce( [ muvals <= Mumaxval,
                                                  muvals >  np.nanmin(vals1) ] )
                    a = np.logical_or.reduce([a1, a2])

                #elif ksigrid.size % 2 == 0 and k == ksigrid.size/2-1:
                elif ksigrid.size % 2 == 1 and k == (ksigrid.size-1)/2-1 :
                    a = np.logical_and.reduce( [ muvals <= np.nanmax(vals0),
                                                 muvals >= -Mumaxval ] )
                #elif ksigrid.size % 2 == 0 and k == ksigrid.size/2:
                elif ksigrid.size % 2 == 1 and k == (ksigrid.size-1)/2 :
                    a = np.logical_and.reduce( [ muvals > np.nanmin(vals1),
                                                 muvals <= Mumaxval ] )
                else:
                    a = np.logical_and.reduce(
                        [ muvals >= np.nanmin(np.append(vals0,vals1)),
                          muvals < np.nanmax(np.append(vals0,vals1)) ] )

                b = np.logical_and.reduce(
                    [ Pvals >= np.nanmin(np.append(vals2,vals3)),
                      Pvals < np.nanmax(np.append(vals2,vals3)) ])

                # Weight of this cell is equal to sum of weights of those (P,mu)
                # values that this cell holds.
                ids = np.logical_and.reduce([a, b])
                if np.sum(ids.ravel()) > 0:
                    ordinate[i,j,k] = np.sum( weights[ids].ravel() )

    if free:
        a5.free(bfield=True)

    return ordinate


def test_rzk2Pmu(fn):
    """
    Verify R-z-ksi & P-mu transformation.

    This function populates (r,z,ksi) mesh uniformly, transforms it to (P,mu)
    mesh, and transforms the resulting values back to (r,z,ksi) grid. Therefore
    the initial and final histograms should resemble each other. Process is
    repeated starting from uniform (P,mu) population. Final and intermediate
    histograms are plotted.

    Args:
        fn : str <br>
            Filename for ascot HDF5 file which must have magnetic field present.
    """
    from a5py.ascot5io.ascot5 import Ascot
    from a5py.ascotpy.ascotpy import Ascotpy
    h5 = Ascot(fn)
    a5 = Ascotpy(fn)
    a5.init(bfield=True)

    mass   = const.physical_constants["alpha particle mass"][0]
    energy = 3.5e6 * const.e
    charge = 2 * const.e

    # Begin by finding min/max values for r and z based on separatrix.
    rho   = np.array([0,1])
    theta = np.array([1])
    rz = a5.get_rhotheta_rz(rho, theta, 0, 0 )
    rmin = rz[0][0]
    rmax = rz[0][0]
    zmin = rz[1][0]
    zmax = rz[1][0]
    for i in range(360):
        rz = a5.get_rhotheta_rz(rho, i*theta, 0, 0 )
        rmin = np.minimum(rmin, rz[0][1])
        rmax = np.maximum(rmax, rz[0][1])
        zmin = np.minimum(zmin, rz[1][1])
        zmax = np.maximum(zmax, rz[1][1])

    rgrid = np.linspace(rmin, rmax, 20)
    zgrid = np.linspace(zmin, zmax, 40)
    kgrid = np.linspace(-1, 1, 51)

    nP  = 100
    nmu = 100
    try:
        fig  = plt.figure()

        Pgrid, mugrid, = initgridPmu(a5, mass, charge, energy,
                                     nP, nmu, padding=1e-3)

        # Make uniform population on (R,z,ksi) grid.
        rvals, zvals, kvals = np.meshgrid(
            rgrid[:-1] + (rgrid[1] - rgrid[0]) / 2,
            zgrid[:-1] + (zgrid[1] - zgrid[0]) / 2,
            kgrid[:-1] + (kgrid[1] - kgrid[0]) / 2, indexing='ij')
        weights = np.ones(rvals.shape)

        # Transform it to (P,mu) space and plot
        weights, Pvals, muvals = maprzk2Pmu(a5, mass, charge, energy,
                                            rvals.ravel(), zvals.ravel(),
                                            kvals.ravel(), Pgrid, mugrid,
                                            weights=weights.ravel())

        axes = fig.add_subplot(2,3,1)
        h = axes.pcolormesh(Pgrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        # Now transform the population back to (R,z,ksi) grid and plot
        weights = np.ones(Pvals.shape)
        weights = mapPmu2rzk(a5, mass, charge, energy,
                             Pvals.ravel(), muvals.ravel(),
                             rgrid, zgrid, kgrid, weights=weights.ravel())

        # Rz
        axes = fig.add_subplot(2,3,2)
        h = axes.pcolormesh(rgrid, zgrid, np.sum(weights, axis=2).transpose())
        plt.colorbar(h, ax=axes)

        # ksi
        axes = fig.add_subplot(2,3,3)
        h = axes.plot(kgrid[:-1] + (kgrid[1] - kgrid[0])/2,
                      np.sum(np.sum(weights, axis=0), axis=0))

        # Repeat the process other way around
        Pvals, muvals = np.meshgrid(
            Pgrid[:-1] + (Pgrid[1] - Pgrid[0]) / 2,
            mugrid[:-1] + (mugrid[1] - mugrid[0]) / 2,
            indexing='ij')
        weights = np.ones(Pvals.shape)
        weights = mapPmu2rzk(a5, mass, charge, energy,
                             Pvals.ravel(), muvals.ravel(),
                             rgrid, zgrid, kgrid, weights=weights.ravel())

        # Rz
        axes = fig.add_subplot(2,3,4)
        h = axes.pcolormesh(rgrid, zgrid, np.sum(weights, axis=2).transpose())
        plt.colorbar(h, ax=axes)

        # ksi
        axes = fig.add_subplot(2,3,5)
        h = axes.plot(kgrid[:-1] + (kgrid[1] - kgrid[0])/2,
                      np.sum(np.sum(weights, axis=0), axis=0))

        # Transform back to (P,mu) space and plot
        rvals, zvals, kvals = np.meshgrid(
            rgrid[:-1] + (rgrid[1] - rgrid[0]) / 2,
            zgrid[:-1] + (zgrid[1] - zgrid[0]) / 2,
            kgrid[:-1] + (kgrid[1] - kgrid[0]) / 2, indexing='ij')
        weights = maprzk2Pmu(a5, mass, charge, energy,
                             rvals.ravel(), zvals.ravel(), kvals.ravel(),
                             Pgrid, mugrid, weights=weights.ravel())[0]

        axes = fig.add_subplot(2,3,6)
        h = axes.pcolormesh(Pgrid, mugrid, weights.transpose())
        plt.colorbar(h, ax=axes)

        plt.show(block=False)

    except Exception as e:
        a5.free(bfield=True)
        raise e

    a5.free(bfield=True)


def maprhoksi2rzk(a5, mass, charge, energy, rhovals, ksivals, rgrid, zgrid,
                  ksigrid, weights=None):
    """
    Convert (rho_omp,ksi_omp) to (r,z,ksi) via (P,mu).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    Pgrid  = np.linspace(-1,1,5)
    mugrid = np.linspace(-1,1,5)
    _, Pvals, muvals = maprhoksi2Pmu(a5, mass, charge, energy, rhovals, ksivals,
                                     Pgrid, mugrid, weights=None)

    ordinate = mapPmu2rzk(a5, mass, charge, energy, Pvals, muvals,
                          rgrid, zgrid, ksigrid, weights=weights)

    if free:
        a5.free(bfield=True)

    return ordinate


def maprzk2rhoksi(a5, mass, charge, energy, rvals, zvals, ksivals, rhogrid,
                  ksigrid, weights=None):
    """
    Convert (r,z,ksi) to (rho_omp,ksi_omp) via (P,mu).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    Pgrid  = np.linspace(-1,1,5)
    mugrid = np.linspace(-1,1,5)
    #_, Pvals, muvals = maprzk2Pmu(a5, mass, charge, energy, rvals, zvals,
    #                              ksivals, Pgrid, mugrid, weights=None)
    Pvals, muvals = evalPmu(a5, mass, charge, energy, rvals, zvals, ksivals)

    ordinate, rhovals, ksivals = mapPmu2rhoksi(a5, mass, charge, energy, Pvals,
                                               muvals, rhogrid, ksigrid,
                                               weights=weights)

    if free:
        a5.free(bfield=True)

    return (ordinate, rhovals, ksivals)


def maprhoksi2rhomu(a5, mass, charge, energy, rhovals, ksivals, rhogrid,
                    mugrid, weights=None):
    """
    Convert (rho_omp,ksi_omp) to (rho,mu) via (P,mu).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    Pgrid  = np.linspace(-1,1,5)
    _, Pvals, muvals = maprhoksi2Pmu(a5, mass, charge, energy, rhovals, ksivals,
                                     Pgrid, mugrid, weights=None)

    ordinate, rhovals, muvals = mapPmu2rhomu(a5, mass, charge, energy, Pvals,
                                             muvals, rhogrid, mugrid,
                                             weights=weights)

    if free:
        a5.free(bfield=True)

    return (ordinate, rhovals, muvals)


def maprhomu2rhoksi(a5, mass, charge, energy, rhovals, muvals, rhogrid,
                    ksigrid, weights=None):
    """
    Convert (rho_omp,ksi_omp) to (rho,mu) via (P,mu).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    Pgrid  = np.linspace(-1,1,5)
    mugrid = np.linspace(-1,1,5)

    _, Pvals, muvals = maprhomu2Pmu(a5, mass, charge, energy, rhovals, muvals,
                                    Pgrid, mugrid, weights=None)

    ordinate, rhovals, ksivals = mapPmu2rhoksi(a5, mass, charge, energy, Pvals,
                                               muvals, rhogrid, ksigrid,
                                               weights=weights)

    if free:
        a5.free(bfield=True)

    return (ordinate, rhovals, ksivals)


def maprhomu2rzk(a5, mass, charge, energy, rhovals, muvals, rgrid, zgrid,
                 ksigrid, weights=None):
    """
    Convert (rho_omp,mu) to (r,z,ksi) via (P,mu).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    Pgrid  = np.linspace(-1,1,5)
    mugrid = np.linspace(-1,1,5)
    _, Pvals, muvals = maprhomu2Pmu(a5, mass, charge, energy, rhovals, muvals,
                                    Pgrid, mugrid, weights=None)

    ordinate = mapPmu2rzk(a5, mass, charge, energy, Pvals, muvals,
                          rgrid, zgrid, ksigrid, weights=weights)

    if free:
        a5.free(bfield=True)

    return ordinate


def maprzk2rhomu(a5, mass, charge, energy, rvals, zvals, ksivals,
                 rhogrid, mugrid, weights=None):
    """
    Convert (r,z,ksi) to (rho_omp,mu) via (P,mu).
    """
    free = False
    if isinstance(a5, str):
        from a5py.ascotpy.ascotpy import Ascotpy
        free = True
        a5 = Ascotpy(a5)
        a5.init(bfield=True)

    Pgrid  = np.linspace(-1,1,5)
    _, Pvals, muvals = maprzk2Pmu(a5, mass, charge, energy, rvals, zvals,
                                  ksivals, Pgrid, mugrid, weights=None)

    ordinate, rhovals, _ = mapPmu2rhomu(a5, mass, charge, energy, Pvals,
                                        muvals, rhogrid, mugrid,
                                        weights=weights)

    if free:
        a5.free(bfield=True)

    return (ordinate, rhovals, muvals)

a5py/a5py/marker/phasespace.py
ascot4fusion/ascot5
"""
Routines for plotting markers.

File: marker/plot.py
"""
import numpy as np
from numpy.random import randint
import importlib.util as util

plt = util.find_spec("matplotlib")
if plt:
    import matplotlib.pyplot as plt

def plot_line(x, y=None, z=None, ids=None, equal=False,
              xlabel=None, ylabel=None, zlabel=None, axes=None, **kwargs):
    """
    Plot continuous line f(x).

    Args:
        x : array_like <br>
            x data.
        y : array_like, optional <br>
            y data.
        z : array_like, optional <br>
            z data, in which case plot will be 3D.
        ids : array_like, optional <br>
            Array with same shape as input data containing marker ID for each
            point. If this is provided, each marker is plotted individually.
        equal : bool, optional <br>
            Flag for making axes equal.
        xlabel : str, optional <br>
            Label on x axis.
        ylabel : str, optional <br>
            Label on y axis.
        zlabel : str, optional <br>
            Label on z axis.
        axes : Axes, optional <br>
            Axes where plot is plotted. If None, a new figure is created.

    Returns:
        Axes where plot is plotted.
    """
    newfig = axes is None
    if newfig:
        plt.figure()
        axes = plt.gca()

    if y is None:
        y = np.linspace(0, x.size, x.size)

    # If markers are separated by their ids, call this function recursively
    # for each marker.
    if ids is not None:
        uids = np.unique(ids)
        for i in uids:
            if z is None:
                plot_line(x[i==ids], y[i==ids], equal=equal, xlabel=xlabel,
                          ylabel=ylabel, axes=axes, **kwargs)
            else:
                plot_line(x[i==ids], y[i==ids], z[i==ids], equal=equal,
                          xlabel=xlabel, ylabel=ylabel, zlabel=zlabel,
                          axes=axes, **kwargs)
    else:
        if z is None:
            axes.plot(x, y, **kwargs)
        else:
            axes.plot(x, y, z, **kwargs)

    if equal:
        axes.axis("scaled")

    if xlabel is not None:
        axes.set_xlabel(xlabel)
    if ylabel is not None:
        axes.set_ylabel(ylabel)
    if z is not None and zlabel is not None:
        axes.set_zlabel(zlabel)

    if newfig:
        plt.show(block=False)

    return axes

def plot_histogram(x, xbins=None, y=None, ybins=None, weights=None,
                   logscale=False, xlabel=None, ylabel=None, axes=None):
    """
    Plot histogram.

    Args:
        x : array_like <br>
            Histogram quantity.
        bins : int or array_like, optional <br>
            Number of bins or array storing bin edges.
        weights : array_like, optional <br>
            Weights for x.
        logy : bool, optional <br>
            Make y-axis logarithmic.
        xlabel : str, optional <br>
            x axis label.
        axes : Axes, optional <br>
            Axes where plot is plotted. If None, a new figure is created.

    Returns:
        Axes where plot is plotted.
    """
    newfig = axes is None
    if newfig:
        plt.figure()
        axes = plt.gca()

    if xbins is None:
        xbins = np.linspace(np.amin(x), np.amax(x), 10)

    if y is None:
        axes.hist(x, xbins, density=False, stacked=True, log=logscale,
                  weights=weights, rwidth=1)

        if xlabel is not None:
            axes.set_xlabel(xlabel)
        if weights is not None:
            axes.set_ylabel("Particles per bin")
        else:
            axes.set_ylabel("Markers per bin")

        if not logscale:
            axes.ticklabel_format(style="sci", axis="y", scilimits=(0,0))

    else:
        if ybins is None:
            ybins = np.linspace(np.amin(y), np.amax(y), 10)
        axes.hist2d(x, y, bins=[xbins, ybins], weights=weights)

        if xlabel is not None:
            axes.set_xlabel(xlabel)

        if ylabel is not None:
            axes.set_ylabel(ylabel)

    if newfig:
        plt.show(block=False)

    return axes

def plot_scatter(x, y=None, z=None, c=None, prune=1, equal=False, ids=None,
                 xlabel=None, ylabel=None, zlabel=None, axes=None, cmap=None,
                 **kwargs):
    """
    Plot a scatter plot.

    Args:
        x : array_like <br>
            x data.
        y : array_like, optional <br>
            y data.
        z : array_like, optional <br>
            z data. If None, scatter will be plotted in 2D.
        c : array_like, optional <br>
            Color data.
        prune : int, optional <br>
            Plot only x[::prune] data points.
        kwargs : <br>
            Same arguments as in plot_line.

    Returns:
        Axes where plot is plotted.
    """
    newfig = axes is None
    if newfig:
        axes = plt.figure()
        axes = plt.gca()

    if isinstance(prune, int):
        prune = slice(None,None,prune)

    if cmap is not None:
        cmap = plt.cm.get_cmap(cmap, 5)
    else:
        cmap = plt.cm.get_cmap("viridis", 5)
    if z is None and c is None:
        if ids is not None:
            uids  = np.unique(ids)
            cpick = np.arange(5)
            np.random.shuffle(cpick)
            for i in range(uids.size):
                c = np.asarray( [cmap( cpick[np.mod(i, 5)] )]
                                * np.sum( ids==uids[i] ) )
                h = axes.scatter(x[ids==uids[i]][prune],
                                 y[ids==uids[i]][prune], c=c[prune],
                                 edgecolors="none", **kwargs)
        else:
            h = axes.scatter(x[prune], y[prune], **kwargs)

    elif z is not None and c is None:
        if ids is not None:
            uids = np.unique(ids)
            for i in uids:
                h = axes.scatter(x[ids==i][prune], y[ids==i][prune],
                                 zs=z[ids==i][prune], edgecolors="none",
                                 **kwargs)

        else:
            h = axes.scatter(x, y, zs=z, **kwargs)

    elif z is None and c is not None:
        if ids is not None:
            uids = np.unique(ids)
            for i in uids:
                h = axes.scatter(x[ids==i][prune], y[ids==i][prune],
                                 c=c[ids==i][prune], edgecolors="none",
                                 **kwargs)

        else:
            h = axes.scatter(x, y, c=c, **kwargs)
            plt.colorbar(h, ax=axes)
    else:
        if ids is not None:
            uids = np.unique(ids)
            for i in uids:
                h = axes.scatter(x[ids==i][prune], y[ids==i][prune],
                                 zs=z[ids==i][prune], c=c[ids==i][prune],
                                 edgecolors="none", **kwargs)

        else:
            h = axes.scatter(x, y, zs=z, c=c, **kwargs)
            plt.colorbar(h, ax=axes)

    if equal:
        axes.axis("scaled")

    axes.set_xlabel(xlabel);
    axes.set_ylabel(ylabel);
    axes.tick_params(axis='x', direction='out')
    axes.tick_params(axis='y', direction='out')

    if z is not None:
        axes.set_zlabel(zlabel);
        axes.tick_params(axis='z', direction='out')

    if newfig:
        plt.show(block=False)

    return axes

a5py/a5py/marker/plot.py
ascot4fusion/ascot5
"""
Study particle losses in phase space.

This module contains tools to plot how losses originate in phase space, to
evaluate transport coefficients, and to use the coefficients for estimating
losses.

pploss.py
"""
import numpy as np
from scipy.interpolate import interp1d, RectBivariateSpline, interp2d

import a5py.marker.phasespace as phasespace

import importlib.util as util

plt = util.find_spec("matplotlib")
if plt:
    import matplotlib.pyplot as plt
    from matplotlib.gridspec import GridSpec
    from matplotlib.colors import ListedColormap


def plotlossmap(a5, mass, charge, energy, r, z, pitch, rhogrid, ksigrid,
                time, lost, weights, rmin, binmin=-5, binmax=-1, nbin=4, muin=False,
                uselinscale=False, axes=None):
    """
    Plot lost markers in (rho_omp,ksi_omp) or (rho_omp,mu) space.

    This plot constructs a loss map for tokamaks. Loss map shows the fraction of
    lost particles and their mean loss time. Loss time is shown in color while
    the fraction of particles lost is illustrated with shade and black contours
    (10 % losses = thin contour, 90 % losses = thick contour).  The plot also
    shows the passing-trapped boundary in red.

    The color scale is discrete and logarithmic (unless uselinscale=True).

    Args:
        a5 : Ascotpy <br>
            Ascotpy object with magnetic field initialized.
        mass : float <br>
            Marker mass [kg].
        charge : float <br>
            Marker charge [C].
        energy : float <br>
            Marker energy [J].
        r : array_like (n,1) <br>
            Marker major radius coordinates [m].
        z : array_like (n,1) <br>
            Marker z coordinates [m].
        pitch : array_like (n,1) <br>
            Marker pitch coordinates (v_para/v_tot).
        rhogrid : array_like <br>
            OMP rho grid values.
        ksigrid : array_like <br>
            OMP pitch grid values.
        time : array_like (n,1) <br>
            Marker end state time [s].
        lost : array_like (n,1) <br>
            Boolean array indicating which markers are lost.
        weights : array_like (n,1) <br>
            Marker weights.
        rmin : float <br>
            R coordinate for the inner mid plane separatrix.
        binmin : float, optional <br>
            Minimum time exponent (or value if uselinscale=True) for the color scale.
        binmax : float, optional <br>
            Maximum time exponent (or value if uselinscale=True) for the color scale.
        nbin : integer, optional <br>
            Number of bins on the color scale.
        muin : bool, optional <br>
            mu grid if the lossmap is plotted in (rho_omp,mu) instead.
        uselinscale : bool <br>
            Use linear scale for time instead.
        axes : Axes, optional <br>
            Axes on which the lossmap is plotted. Default is to open a new figure.
    """

    showfig = False
    if axes is None:
        plt.figure()
        axes = plt.gca()
        showfig = True

    # Find the phase-space coordinates for each marker
    if muin:
        xgrid = rhogrid
        ygrid = ksigrid
        _, x, y = phasespace.maprzk2rhomu(a5, mass, charge, energy, r, z,
                                          pitch, xgrid, ygrid,
                                          weights=None)
    else:
        xgrid = rhogrid
        ygrid = ksigrid
        _, x, y = phasespace.maprzk2rhoksi(a5, mass, charge, energy, r, z,
                                           pitch, xgrid, ygrid,
                                           weights=None)

    # Which particles are trapped
    P,mu = phasespace.evalPmu(a5, mass, charge, energy, r, z, pitch)
    trapped = phasespace.istrapped(a5, mass, charge, energy,
                                   P, mu, rmin=rmin)

    # Evaluate the plotted quantities
    lostxy = np.histogram2d(x[lost], y[lost], bins=[xgrid,ygrid],
                            weights=weights[lost])[0]
    timexy = np.histogram2d(x[lost], y[lost], bins=[xgrid,ygrid],
                            weights=(time*weights)[lost])[0]
    timexy = np.divide(timexy, lostxy,
                       out=np.zeros(timexy.shape, dtype=float), where=lostxy!=0)

    if not uselinscale:
        timexy = np.log10(timexy)

    weightxy = np.histogram2d(x, y, bins=[xgrid,ygrid], weights=weights)[0]
    lossfrac = np.divide(lostxy, weightxy,
                         out=np.zeros(lostxy.shape, dtype=float), where=weightxy!=0)

    trapxy = np.histogram2d(x, y, bins=[xgrid,ygrid],
                            weights=trapped*weights)[0]
    trapxy = np.divide(trapxy, weightxy,
                       out=np.zeros(trapxy.shape, dtype=float), where=weightxy!=0)

    # Plot loss time with different colours
    cmap = plt.cm.get_cmap("viridis_r", nbin)
    alphacmap = cmap(np.arange(cmap.N))
    alphacmap[:,:] = 1
    alphacmap[:,-1] = np.linspace(0, 1, cmap.N)
    alphacmap = ListedColormap(alphacmap)

    mesh = axes.pcolormesh(xgrid, ygrid, timexy.transpose(),
                           cmap=cmap, vmin=binmin, vmax=binmax)
    axes.figure.canvas.draw()

    # Get the colors and adjust shade so it corresponds to fraction of prts lost
    colors = mesh.get_facecolor()
    def alpha_to_white(color):
        white = np.array([1,1,1])
        alpha = color[-1]
        color = color[:-1]
        return alpha*color + (1 - alpha)*white

    colors[:,3] = lossfrac.transpose().ravel()
    colors = np.array([alpha_to_white(color) for color in colors])
    mesh.set_facecolors(colors)

    # Passing-trapped boundary
    axes.contour(xgrid[:-1] + (xgrid[1]-xgrid[0])/2,
                 ygrid[:-1] + (ygrid[1]-ygrid[0])/2,
                 (trapxy.transpose() > 0)*1, [1], alpha=0.7,
                 colors='red', linestyles='dotted', linewidths=1)

    # 10 % and 90 % contours
    axes.contour(xgrid[:-1] + (xgrid[1]-xgrid[0])/2,
                 ygrid[:-1] + (ygrid[1]-ygrid[0])/2,
                 lossfrac.transpose(), [0.1, 0.9], colors=['black', 'black'],
                 alpha=0.7, linewidths=[1, 3])

    axes.figure.canvas.draw()

    if showfig:
        plt.show()


def evallossfrac(a5, mass, charge, energy, r, z ,pitch, rhogrid, ksigrid,
                 time, lost, weights, endenergy=None):
    xgrid = rhogrid
    ygrid = ksigrid
    _, x, y = phasespace.maprzk2rhoksi(a5, mass, charge, energy, r, z,
                                       pitch, xgrid, ygrid,
                                       weights=None)
    P,mu = phasespace.evalPmu(a5, mass, charge, energy, r, z, pitch)

    lostxy = np.histogram2d(x[lost], y[lost], bins=[xgrid,ygrid],
                            weights=weights[lost])[0]
    timexy = np.histogram2d(x[lost], y[lost], bins=[xgrid,ygrid],
                            weights=(time*weights)[lost])[0] / lostxy

    weightxy = np.histogram2d(x, y, bins=[xgrid,ygrid], weights=weights)[0]
    lossfrac = lostxy / weightxy

    if endenergy is None:
        return (lossfrac, timexy)
    else:
        lostxy = np.histogram2d(x[lost], y[lost], bins=[xgrid,ygrid],
                                weights=(weights*endenergy)[lost])[0]
        weightxy = np.histogram2d(x, y, bins=[xgrid,ygrid],
                                  weights=weights*energy)[0]
        return (lossfrac, timexy, lostxy / weightxy)


def evalpdens(a5, mass, charge, energy, r, z ,pitch, rhogrid, ksigrid, weights):
    """
    """
    xgrid = rhogrid
    ygrid = ksigrid
    _, x, y = phasespace.maprzk2rhoksi(a5, mass, charge, energy, r, z,
                                       pitch, xgrid, ygrid,
                                       weights=None)
    P,mu = phasespace.evalPmu(a5, mass, charge, energy, r, z, pitch)


    weightxy = np.histogram2d(x, y, bins=[xgrid,ygrid],
                              weights=weights*energy)[0]

    return weightxy


def runtransportmodel(coefficients, markers, tmax, Nt,
                      xmin=0, xmax=1, Nmrk=1000, axis=None, ksigrid=None,
                      a5=None, mass=None, charge=None, energy=None):
    xgrid = coefficients[0]
    ygrid = coefficients[1]
    lossfrac = np.zeros((xgrid.size-1,ygrid.size-1))
    losstime = np.zeros((xgrid.size-1,ygrid.size-1))
    endstate = np.zeros((xgrid.size-1,ygrid.size-1))

    for j in range(ygrid.size-1):
        dt   = tmax/Nt

        x    = markers[0] + np.random.rand(Nmrk,1) * (markers[-1] - markers[0])
        x0   = x.copy()
        time = np.zeros(x.shape)
        lost = np.zeros(x.shape) == 1

        K = interp1d(xgrid[:-1] + (xgrid[0]-xgrid[0])/2,
                     coefficients[2][:,j], kind='nearest',
                     fill_value='extrapolate')
        D = interp1d(xgrid[:-1] + (xgrid[0]-xgrid[0])/2,
                     coefficients[3][:,j], kind='nearest',
                     fill_value='extrapolate')

        for i in range(Nt):
            # Prepare
            w = 1 - 2*np.random.rand(Nmrk,1)
            k = K(x)
            d = D(x)

            # Step
            x += (lost == False) * (k*dt + w*np.sqrt(2*d*dt))
            time[lost == False] += dt

            # Enforce boundary conditions
            reflect = x < xmin
            lost    = x >= xmax

            x[reflect] = 2*xmin - x[reflect]
            x[lost] = xmax

        for i in range(xgrid.size-1):
            ids = np.logical_and.reduce([x0 >= xgrid[i],
                                         x0 <  xgrid[i+1]])
            lossfrac[i,j] = np.sum(lost[ids])/np.sum(ids)

            ids = np.logical_and.reduce([lost, ids])
            losstime[i,j] = np.mean(time[ids])

            ids = np.logical_and.reduce([x >= xgrid[i],
                                         x <  xgrid[i+1]])
            endstate[i,j] = np.sum(ids)

    if ksigrid is not None:
        xvals, yvals = np.meshgrid(xgrid[:-1]+(xgrid[1]-xgrid[0])/2,
                                   ksigrid[:-1]+(ksigrid[1]-ksigrid[0])/2,
                                   indexing='ij')
        _, x, y = phasespace.maprhoksi2rhomu(a5, mass, charge, energy,
                                             xvals.ravel(), yvals.ravel(),
                                             xgrid, ygrid, weights=None)

        xvals, yvals = np.meshgrid(xgrid[:-1]+(xgrid[1]-xgrid[0])/2,
                                   ygrid[:-1]+(ygrid[1]-ygrid[0])/2,
                                   indexing='ij')
        f = RectBivariateSpline(xgrid[:-1]+(xgrid[1]-xgrid[0])/2,
                                ygrid[:-1]+(ygrid[1]-ygrid[0])/2,
                                lossfrac)
        lossfrac = f(x,y,grid=False).reshape(xgrid.size-1, ksigrid.size-1)
        #f = interp2d(xvals.ravel(),
        #             yvals.ravel(),
        #             lossfrac.ravel())
        #lossfrac = (np.diag(f(x,y, grid=False)).reshape(ksigrid.size-1, xgrid.size-1)).T
        #print(lossfrac.shape)
        #lossfrac = f(xgrid, y)


        lossfrac[lossfrac>1] = 1
        lossfrac[lossfrac<0] = 0

        losstime[np.isnan(losstime)] = 0
        f = RectBivariateSpline(xgrid[:-1]+(xgrid[1]-xgrid[0])/2,
                                ygrid[:-1]+(ygrid[1]-ygrid[0])/2,
                                losstime)
        losstime = f(x,y,grid=False).reshape(xgrid.size-1, ksigrid.size-1)

        #f = interp2d(xvals,
        #             yvals,
        #             losstime)
        #losstime = f(xgrid[:-1]+(xgrid[1]-xgrid[0])/2, ygrid[:-1]+(ygrid[1]-ygrid[0])/2).T#f(xgrid, ygrid)#(x, y).reshape(xgrid.size-1, ksigrid.size-1)
        #losstime = (np.diag(f(x,y)).reshape(ksigrid.size-1, xgrid.size-1)).T
        #losstime.setflags(write=1)

        ygrid = ksigrid

        #n = np.histogram2d(x, y, bins=[xgrid,ygrid])[0]
        #lossfrac = np.histogram2d(x, y, bins=[xgrid,ygrid],
        #                          weights=lossfrac.ravel())[0] / n
        #losstime = np.histogram2d(x, y, bins=[xgrid,ygrid],
        #                          weights=losstime.ravel())[0] / n
        #lossfrac[np.isnan(lossfrac)] = 0
        #losstime[np.isnan(losstime)] = 1


    if axis is not None:
        cmap = plt.cm.get_cmap("viridis_r", 5)

        mesh = axis.pcolormesh(xgrid, ygrid, np.log10(losstime.transpose()),
                               cmap=cmap, vmin=-5, vmax=-1)

        axis.figure.canvas.draw()
        colors = mesh.get_facecolor()

        def alpha_to_white(color):
            white = np.array([1,1,1])
            alpha = color[-1]
            color = color[:-1]
            return alpha*color + (1 - alpha)*white

        colors[:,3] = lossfrac.transpose().ravel()
        colors = np.array([alpha_to_white(color) for color in colors])
        mesh.set_facecolor(colors)

        axis.contour(xgrid[:-1] + (xgrid[1]-xgrid[0])/2,
                 ygrid[:-1] + (ygrid[1]-ygrid[0])/2,
                 lossfrac.transpose(), [0.1, 0.9], colors=['black', 'black'],
                 alpha=0.7, linewidths=[1, 3])


    return (lossfrac,losstime)



def evalcoefs(a5, mass, charge, energy, r, z, pitch, weights, xgrid, ygrid,
              drift, diff, lost, endtime, xedge, mutype=False):
    if mutype:
        _, x, y = phasespace.maprzk2rhomu(a5, mass, charge, energy, r, z,
                                          pitch, xgrid, ygrid, weights=None)
    else:
        _, x, y = phasespace.maprzk2rhoksi(a5, mass, charge, energy, r, z,
                                           pitch, xgrid, ygrid, weights=None)

    if weights is None:
        weights = np.ones(drift.shape)

    xi = np.searchsorted(xgrid, x, side='left')-1
    yi = np.searchsorted(ygrid, y, side='left')-1

    K = np.zeros((xgrid.size-1,ygrid.size-1))
    D = np.zeros((xgrid.size-1,ygrid.size-1))
    for i in range(xgrid.size-1):
        for j in range(ygrid.size-1):
            idx = np.logical_and.reduce([xi == i, yi == j])

            if np.sum(idx) == 0:
                # No markers in this cell
                continue

            iii = np.logical_and.reduce([idx, lost == False])
            if np.sum(lost[idx]) < np.sum(idx):
                #K[i,j] = np.nanmean(drift[iii])
                #D[i,j] = np.nanmean(diff[iii])
                K[i,j] = np.average(drift[iii], weights=weights[iii])
                D[i,j] = np.average(diff[iii], weights=weights[iii])
            else:
                K[i,j] = 0
                D[i,j] = 0

            if np.sum(lost[idx]) > 0:
                jjj = np.logical_and.reduce([lost, idx])
                c1 = np.average(endtime[jjj], weights=weights[jjj])

                v = np.average((endtime[jjj] - c1)**2, weights=weights[jjj])
                #c2 = c1*c1*c1 / np.var(endtime[jjj])
                c2 = c1*c1*c1/v
                delta  = xedge - ( xgrid[i] + xgrid[i+1] ) / 2

                frac = np.sum(lost[idx]) / np.sum(idx)

                K[i,j] = (1-frac) * K[i,j] + frac * (delta / c1)
                D[i,j] = (1-frac) * D[i,j] + frac * (2 * delta * delta / c2)


    return (xgrid, ygrid, K, D)

a5py/a5py/marker/pploss.py
ascot4fusion/ascot5
"""
Solve 1D advection diffusion equation with Monte Carlo method.

The advection-diffusion equation has a reflecting boundary condition at rho0 and
absorbing BC at rho1. The coefficients are stationary but may have spatial
dependency.
"""

import numpy as np
import matplotlib.pyplot as plt

def solve_model(xmin, xmax, simtime, timestep, numberofmarkers,
                driftfun, difffun, distfun, conffun, xgrid=None, tgrid=None):
    """
    Solve 1D advection diffusion equation with boundary conditions.

    Reflective boundary condition is at xmin and absorbing boundary is at xmax,
    where xmin < xmax and x is the radial coordinate.

    Args:
        xmin          : float <br>
            Minimum rho and where the reflecting boundary is located.
        xmax          : float <br>
            Maximum rho and where the absorbing boundary is located.
        simtime         : float <br>
            Simulation time [s].
        timestep        : float <br>
            Simulation time-step [s].
        numberofmarkers : float <br>
            Number of markers. Increase to reduce Monte Carlo noise.
        driftfun        : function <br>
            Function that evaluates drift at the given radial position.
        difffun         : function <br>
            Function that evaluates diffusion at the given radial position.
        distfun         : function <br>
            Function that evaluates initial marker distribution at the given
            radial position.
        conffun         : function <br>
            Function that evaluates which markers are born stationary.
        xgrid           : array_like, optional <br>
            Spatial grid for the density histogram.
        tgrid           : array_like <br>
            Temporal grid for the density histogram.

    Returns:
        Dictionary storing the final x coordinates for the confined and
        un-confined markers alike, losstimes, and the density histogram.
    """
    Nmrk      = numberofmarkers
    xcoords   = np.array([])
    xconfined = np.array([])
    losstime  = simtime*np.ones((Nmrk,))

    if xgrid is None:
        xgrid = np.linspace(xmin, xmax, 10)
    if tgrid is None:
        tgrid = np.linspace(timestep, simtime, 1000)
    density = np.zeros( (xgrid.size-1, tgrid.size-1) )

    # Evaluate initial profile
    N = 0
    while N < Nmrk:
        # Generate uniform distribution of coordinates and a random number
        # X c [0,1] for each coordinate. Accept coordinates for which
        # X(x) < P(x), where P is the given distribution function.
        x = xmin + (xmax - xmin)*np.random.rand(Nmrk - N,)
        accept  = distfun(x) - np.random.rand(Nmrk - N,)
        x       = x[accept>0]

        # Each coordinate has the probability of C(x) of being confined where
        # C is the given fraction of confined markers.
        conf      = conffun(x) - np.random.rand(x.size,)
        xcoords   = np.append(xcoords, x[conf<=0])
        xconfined = np.append(xconfined, x[conf>0])

        N = xcoords.size + xconfined.size

    # Set confined markers to correct bins
    weights = np.ones(xconfined.shape) * (tgrid[1]-tgrid[0])
    for ti in range(tgrid.size-1):
        density[:,ti] += np.histogram(xconfined, bins=xgrid, weights=weights)[0]

    # The simulation loop
    ismoving = np.ones((Nmrk,), dtype="i8")
    nsteps   = int(simtime/timestep)
    for i in range(nsteps):

        # Init step
        Nmoving  = np.sum(ismoving)
        x0 = xcoords[ismoving]
        K  = driftfun(x0)
        D  = difffun(x0)
        w  = np.random.randn(Nmoving,)

        # Take the step and apply boundary conditions
        x1 = x0 + K*timestep + np.sqrt(2*D*timestep)*w

        reflect = x1 < xmin
        absorb  = x1 >= xmax

        x1[reflect] = 2*xmin - x1[reflect]
        losstime[np.argwhere(ismoving)[absorb]] = (i+1)*timestep

        # Update coordinates and the histogram
        xcoords[ismoving] = x1
        ismoving = xcoords <= xmax

        ti = np.floor( i*timestep / ( tgrid[1] - tgrid[0] ) ).astype("i8")
        weights = np.ones(x1.shape) * timestep
        #density[:,ti] += np.histogram(x1, bins=xgrid, weights=weights)[0]

    density /= simtime
    losstime = np.sort(losstime[losstime<simtime])

    out = {"xcoord" : xcoords, "losstime" : losstime, "xconfined" : xconfined,
           "xgrid" : xgrid, "tgrid" : tgrid, "density" : density}
    return out


def solve_model_lincoeff(xmin, xmax, simtime, timestep, numberofmarkers,
                         datagrid, driftdata, diffdata, distdata, confdata,
                         xgrid=None, tgrid=None):
    """
    Solve the 1D model using linear fit on the data.

    If xgrid does not span the whole interval [xmin, xmax], the distribution is
    set to zero outside the grid while the closest value (i.e. first or last) is
    used for the transport coefficients.
    """

    def evaldrift(x):
        """
        Evaluate advection with linear interpolation.
        """
        return np.interp(x, datagrid, driftdata,
                         left=driftdata[0], right=driftdata[-1])

    def evaldiff(x):
        """
        Evaluate diffusion with linear interpolation.
        """
        return np.interp(x, datagrid, diffdata,
                         left=diffdata[0], right=diffdata[-1])

    def evaldist(x):
        """
        Evaluate initial distribution with linear interpolation.
        """
        return np.interp(x, datagrid, distdata, left=0, right=0)

    def evalconf(x):
        """
        Evaluate confined fraction with linear interpolation.
        """
        return np.interp(x, datagrid, confdata, left=0, right=0)

    return solve_model(xmin, xmax, simtime, timestep, numberofmarkers,
                       evaldrift, evaldiff, evaldist, evalconf,
                       xgrid=xgrid, tgrid=tgrid)

a5py/a5py/marker/solve1d.py
ascot4fusion/ascot5
"""
Generate a generic neutral beam injector geometry for BBNBI5.

Jari Varje <jari.varje@tokamakenergy.co.uk>
"""
import numpy as np
import scipy.constants as const
import a5py.ascot5io.nbi as nbi

def generate(fn, Rgrid, phigrid, zgrid, tanrad, focal_length, dgrid, nbeamlet, anum, znum, mass, energy, efrac, power, div, tilt=None, tanz=None, div_halo_frac=None, div_halo=None, desc=None):
    """
    Generate generic injector. The injector grid can be either circular or
    rectangular, and will be randomly populated with beamlets with a common
    focus point.

    Args:
        fn : str
            Full path to the HDF5 file.
        Rgrid, phigrid, zgrid : float
            Position of the center point of the injector grid [m], [deg], [m]
        tanrad : float
            Tangency radius of the injector centerline
        focal_length : float
            Distance to the focus point from the center of the grid
        dgrid : float or array_like (2,)
            Diameter of circular grid or width and height of rectangular grid
        nbeamlet : int
            Number of beamlets in this injector
        anum : int
            Mass number of injected species
        znum : int
            Nuclear charge number of injected species
        mass : float
            Mass of the injected species [amu]
        energy : float
            Full injection energy [keV]
        efrac : array_like (3,)
            Particle fractions for full, 1/2 and 1/3 energies
        power : float
            Injected power [W]
        div : array_like (2,)
            Horizontal and vertical divergences [rad]
        tilt : float, optional
            Vertical tilt angle of the beam centerline [deg]
        tanz : float, optional
            Vertical shift of the tangency radius point [m]
        div_halo_frac : float, optional
            Fraction of particles with halo divergence
        div_halo : float, optional
            Horizontal and vertical divergences [rad]
        desc : str, optional
            Input description
    """

    phigrid = phigrid * np.pi/180.0
    phidir = np.pi + phigrid - np.arcsin(tanrad/Rgrid);

    grid = np.array([Rgrid * np.cos(phigrid), Rgrid * np.sin(phigrid), zgrid])

    unitd = np.array([np.cos(phidir), np.sin(phidir), 0])
    unitz = np.array([0, 0, 1])
    unith = np.cross(unitd, unitz)

    if tilt is not None:
        tilt = tilt * np.pi/180.0
    elif tanz is not None:
        tilt = tanz / Rgrid*np.cos(phidir)
    else:
        tilt = 0

    focus = grid + focal_length*unitd + focal_length*np.tan(tilt)*unitz

    n=dict()
    n["id"] = 1
    n["nbeamlet"] = nbeamlet
    n["beamletx"] = np.zeros(nbeamlet)
    n["beamlety"] = np.zeros(nbeamlet)
    n["beamletz"] = np.zeros(nbeamlet)
    n["beamletdx"] = np.zeros(nbeamlet)
    n["beamletdy"] = np.zeros(nbeamlet)
    n["beamletdz"] = np.zeros(nbeamlet)
    n["div_h"] = div[0]
    n["div_v"] = div[1]

    if div_halo_frac is not None:
        n["div_halo_frac"] = div_halo_frac
        n["div_halo_h"] = div_halo[0]
        n["div_halo_v"] = div_halo[1]
    else:
        n["div_halo_frac"] = 0.0
        n["div_halo_h"] = 1e-10
        n["div_halo_v"] = 1e-10

    n["anum"] = anum
    n["znum"] = znum
    n["mass"] = mass*const.u
    n["energy"] = energy*const.e
    n["efrac"] = np.array(efrac)
    n["power"] = power

    for i in range(nbeamlet):
        if np.array(dgrid).size == 1:    # circular grid
            d = np.inf
            while d > dgrid/2:
                dx = dgrid/2 * (2*np.random.rand(1)-1)
                dy = dgrid/2 * (2*np.random.rand(1)-1)
                d = np.sqrt(dx**2+dy**2)
        else:    # rectangular grid
            dx = dgrid[0]/2 * (2*np.random.rand(1)-1)
            dy = dgrid[1]/2 * (2*np.random.rand(1)-1)

        beamletxyz = grid + dx*unith + dy*unitz
        dirxyz = focus - beamletxyz
        dirxyz = dirxyz / np.linalg.norm(dirxyz)

        n["beamletx"][i] = beamletxyz[0]
        n["beamlety"][i] = beamletxyz[1]
        n["beamletz"][i] = beamletxyz[2]
        n["beamletdx"][i] = dirxyz[0]
        n["beamletdy"][i] = dirxyz[1]
        n["beamletdz"][i] = dirxyz[2]

    nbi.write_hdf5(fn, [n], desc)

a5py/a5py/nbi/generate_injector.py
ascot4fusion/ascot5
"""
Routines for plotting nbi inputs.

File: nbi/plot.py

"""
import numpy as np
from numpy.random import randint
import importlib.util as util

plt = util.find_spec("matplotlib")
if plt:
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D

def plot_arrow_3D(x, y, z, dx, dy, dz, equal=False,
                  xlabel=None, ylabel=None, zlabel=None, axes=None, **kwargs):
    """
    Plot arrows in 3D plot.
    
    Args:
        x : array_like <br>
            x data.
        y : array_like <br>
            y data.
        z : array_like <br>
            z data.
        dx : array_like <br>
            dx data.
        dy : array_like <br>
            dy data.
        dz : array_like <br>
            dz data.
        equal : bool, optional <br>
            Flag for making axes equal.
        xlabel : str, optional <br>
            Label on x axis.
        ylabel : str, optional <br>
            Label on y axis.
        zlabel : str, optional <br>
            Label on z axis.
        axes : Axes, optional <br>
            Axes where plot is plotted. Remember it must have 3D projections enabled.
            If None, a new figure is created.

    Returns:
        Axes where plot is plotted.
    """
    newfig = axes is None
    if newfig:
        fig = plt.figure()
        axes = fig.add_subplot(1,1,1,projection="3d")

    axes.quiver3D(x,y,z,
                  dx,dy,dz, **kwargs)

    if equal:
        axes.axis("scaled")

    if xlabel is not None:
        axes.set_xlabel(xlabel)
    if ylabel is not None:
        axes.set_ylabel(ylabel)
    if zlabel is not None:
        axes.set_zlabel(zlabel)

    if newfig:
        plt.show(block=False)

    return axes

## Continue here

def plot_scatter_3D(x, y, z, equal=False, xlabel=None, ylabel=None, zlabel=None,
                    axes=None, **kwargs):
    """
    Plot a scatter plot.

    Args:
        x : array_like <br>
            x data.
        y : array_like, optional <br>
            y data.
        z : array_like, optional <br>
            z data. If None, scatter will be plotted in 2D.
        equal : bool, optional <br>
            Flag for making axes equal.
        xlabel : str, optional <br>
            Label on x axis.
        ylabel : str, optional <br>
            Label on y axis.
        zlabel : str, optional <br>
            Label on z axis.
        axes : Axes, optional <br>
            Axes where plot is plotted. Remember it must have 3D projections enabled.
            If None, a new figure is created.

    Returns:
        Axes where plot is plotted.
    """
    newfig = axes is None
    if newfig:
        fig = plt.figure()
        axes = fig.add_subplot(1,1,1,projection="3d")

    axes.scatter(x,y,z, **kwargs)

    if equal:
        axes.axis("scaled")

    if xlabel is not None:
        axes.set_xlabel(xlabel)
    if ylabel is not None:
        axes.set_ylabel(ylabel)
    if zlabel is not None:
        axes.set_zlabel(zlabel)

    if newfig:
        plt.show(block=False)

    return axes

a5py/a5py/nbi/plot.py
ascot4fusion/ascot5
"""Tools for assessing physical quantities and units.
"""
import unyt
import numpy as np
import scipy.constants as constants
import warnings
from functools import wraps

from a5py.exceptions import AscotUnitWarning

e     = constants.elementary_charge * unyt.C
m_e   = constants.physical_constants["electron mass"][0] * unyt.kg
m_p   = constants.physical_constants["proton mass"][0] * unyt.kg
m_a   = constants.physical_constants["alpha particle mass"][0] * unyt.kg
c     = unyt.c
eps_0 = unyt.eps_0

def cart2pol(x, y, z=None):
    """Cartesian x coordinate from cylindrical coordinates.
    """
    return np.sqrt(x**2 + y**2), np.arctan2(y, x) * unyt.rad, z

def pol2cart(r, phi, z=None):
    """Cartesian x coordinate from cylindrical coordinates.
    """
    return r * np.cos(phi), r * np.sin(phi), z

def anglemod(angle):
    """Transfer arbitrary angle to between interval [0, 2pi].
    """
    return np.mod(angle + 2*np.pi, 2*np.pi)

def velocity_momentum(m, p):
    """Evaluate velocity from momentum.

    If p is a vector (scalar) then returned value is a vector (scalar).
    """
    a = m * gamma_momentum(m=m, p=p)
    if p.shape[0] == 3:
        a = np.tile( a, (3,1))

    return p / a

def momentum_velocity(m, v):
    """Evaluate momentum from velocity.

    If v is a vector (scalar) then returned value is a vector (scalar).
    """
    return gamma_velocity(v=v) * m * v

def vpar_muppar(m, mu, ppar, b):
    """Evaluate vpar from mu and ppar.
    """
    gamma = gamma_muppar(m, mu, ppar, b)
    return ppar / (gamma * m)

def gamma_momentum(m, p):
    """Evaluate gamma from momentum.
    """
    if len(p.shape) > 0 and p.shape[0] == 3:
        p = np.sum(p**2, axis=0)
    else:
        p = p**2

    return np.sqrt( 1 + p / ( m * c )**2 )

def gamma_velocity(v):
    """Evaluate gamma from velocity.
    """
    if v.shape[0] == 3:
        v = np.sqrt(np.sum(v**2, axis=0))
    else:
        v = v
    return np.sqrt( 1.0 / ( ( 1 - v / c ) * ( 1 + v / c ) ) )

def gamma_muppar(m, mu, ppar, b):
    """Evaluate gamma from magnetic moment and parallel momentum.
    """
    if b.shape[0] == 3:
        b = np.sqrt( np.sum(b**2, axis=0) )

    return np.sqrt( 1 + 2 * mu * b / ( m * c**2 ) + ( ppar / ( m * c ) )**2 )

def gamma_muvpar(m, mu, vpar, b):
    """
    Evaluate gamma from magnetic moment and parallel velocity.
    """
    if b.shape[0] == 3:
        b = np.sqrt( np.sum(b**2, axis=0) )

    a1 = 1 + 2 * mu * b / ( m * c**2 )
    a2 = ( 1 - vpar / c ) * ( 1 + vpar * c )
    return np.sqrt( a1 / a2 )

def gamma_energy(m, energy):
    """Evaluate gamma from kinetic energy.
    """
    return 1 + energy / (m * c**2)

def energy_momentum(m, p):
    """Evaluate kinetic energy from momentum.
    """
    gamma = gamma_momentum(m=m, p=p)
    return energy_gamma(m=m, gamma=gamma)

def energy_velocity(m, v):
    """Evaluate kinetic energy from velocity.
    """
    gamma = gamma_velocity(v=v)
    return energy_gamma(m=m, gamma=gamma)

def energy_muppar(m, mu, ppar, b):
    """Evaluate kinetic energy from magnetic moment and parallel momentum.
    """
    gamma = gamma_muppar(m=m, mu=mu, ppar=ppar, b=b)
    return energy_gamma(m=m, gamma=gamma)

def energy_muvpar(m, mu, vpar, b):
    """Evaluate kinetic energy from magnetic moment and parallel velocity.
    """
    gamma = gamma_muvpar(m=m, mu=mu, vpar=vpar, b=b)
    return energy_gamma(m=m, gamma=gamma)

def energy_gamma(m, gamma):
    """Evaluate kinetic energy from Lorentz factor.
    """
    return (gamma - 1.0) * m * c**2

def vnorm_gamma(gamma):
    """Evaluate velocity norm from Lorentz factor.
    """
    return np.sqrt(1 - 1.0 / gamma**2) * c

def pnorm_gamma(m, gamma):
    """Evaluate momentum norm from Lorentz factor.
    """
    return np.sqrt(gamma**2 - 1) * m * c

def momentum_muppar(m, mu, ppar, b, zeta=None):
    """Evaluate momentum from mu and ppar.

    If gyroangle zeta is given, return the momentum vector. Otherwise return
    momentum norm.
    """
    if b.shape[0] == 3:
        b = np.sqrt( np.sum(b**2, axis=0) )

    return np.sqrt( 2 * m * mu * b + ppar**2 )

def velocity_muppar(m, mu, ppar, b, zeta=None):
    """Evaluate velocity from mu and ppar.

    If gyroangle zeta is given, return the velocity vector. Otherwise return
    velocity norm.
    """
    if b.shape[0] == 3:
        b = np.sqrt( np.sum(b**2, axis=0) )

    p = np.sqrt( 2 * m * mu * b + ppar**2 )
    return velocity_momentum(m=m, p=p)

def mu_momentum(m, p, b):
    """Evaluate magnetic moment from momentum vector.
    """
    bnorm = np.sqrt( np.sum(b**2, axis=0) )
    pnorm = np.sqrt( np.sum(p**2, axis=0) )
    pitch = np.sum( p * b, axis=0 ) / ( pnorm * bnorm )
    return ( 1 - pitch**2 ) * pnorm**2 / ( 2 * bnorm * m )

def ppar_momentum(p, b):
    """Evaluate parallel momentum from momentum vector.
    """
    bnorm = np.sqrt( np.sum(b**2, axis=0) )
    return np.sum( p * b, axis=0 ) / bnorm

def vpar_momentum(m, p, b):
    """Evaluate parallel velocity from momentum vector.
    """
    gamma = gamma_momentum(m=m, p=p)
    bnorm = np.sqrt( np.sum(b**2, axis=0) )
    return np.sum( p * b, axis=0 ) / ( gamma * m * bnorm)

def mu_velocity(m, v, b):
    """TODO: Evaluate magnetic moment from velocity vector.
    """
    gamma = gamma_velocity(v=v)
    return 0

def pitch_momentum(p, b):
    """Evaluate pitch from momentum vector.
    """
    return np.sum( p * b, axis=0 ) / np.sqrt(
        np.sum(p**2, axis=0) * np.sum(b**2, axis=0) )

def pitch_muppar(m, mu, ppar, b):
    """Evaluate pitch from mu and ppar.
    """
    if b.shape[0] == 3:
        b = np.sqrt( np.sum(b**2, axis=0) )
    return ppar / np.sqrt( 2 * m * mu * b + ppar**2 )

def torcanangmom_momentum(q, r, p, psi):
    """Evaluate toroidal canonical angular momentum from momentum.
    """
    return p[1,:] * r + q * psi

def torcanangmom_ppar(q, r, ppar, b, psi):
    """Evaluate toroidal canonical angular momentum from mu and ppar.
    """
    bnorm = np.sqrt( np.sum(b**2, axis=0) )
    return ppar * r * b[1,:] / bnorm + q * psi

def gyrolength(m, q, energy, pitch, bnorm):
    """Evaluate gyrolength from energy and pitch.
    """
    gamma = gamma_energy(m, energy)
    vnorm = vnorm_gamma(gamma)
    return ( gamma * np.sqrt(1 - pitch**2) * m * vnorm
             / (bnorm * np.abs(q)) ).to("m")

def gyrofrequency(m, q, energy, bnorm):
    """Evaluate gyrofrequency from energy.
    """
    gamma = gamma_energy(m, energy)
    return unyt.rad * np.abs(q) * bnorm / (gamma * m)

def bouncefrequency(m, ekin, minorradius, majorradius, safetyfactor):
    """Estimate bounce frequency from energy and aspect ratio.
    """
    eps = minorradius / majorradius
    gamma = gamma_energy(m, ekin)
    vnorm = vnorm_gamma(gamma)
    return np.sqrt(0.5 * eps) * vnorm / ( safetyfactor * majorradius )

def collfreq_ei(mi, qi, ne, Te, clog):
    """Evaluate electron-ion collision frequency.
    """
    return ( ( np.sqrt(2 / np.pi) / 3 )
             * ( unyt.e * qi / ( 4 * np.pi * unyt.eps_0 ) )**2 * ne * clog
             * 4 * np.pi / np.sqrt( unyt.me * Te**3 ) ).to("1/s")

def collfreq_ie(mi, qi, ne, Te, clog):
    """Evaluate ion-electron collision frequency.
    """
    return ( (unyt.me / mi) * ( np.sqrt(2 / np.pi) / 3 )
             * ( unyt.e * qi / ( 4 * np.pi * unyt.eps_0 ) )**2 * ne * clog
             * 4 * np.pi / np.sqrt( unyt.me * Te**3 ) ).to("1/s")

def parseunits(strip=False, **units):
    """Prepare arguments that are expected to have physical units.

    This decorator:

    - Makes sure every argument has expected physical dimensions.
    - Assigns units if units were not provided but they were expected.
    - Strips units if asked (after checking/assignment).

    Examples:

    .. code-block:: python

       @parseunits(x="m", strip=True)
       def fun(x):
           pass

    Parameters
    ----------
    strip : bool, optional
        Strip units so that dimensionless quantities (but in expected units)
        are passed to the function.
    **units
        Argument in wrapped function and the expected unit as a string.
    """
    for k in units.keys():
        # Convert unit strings to unyts
        units[k] = unyt.unyt_quantity.from_string(units[k]).units

    def actualdecorator(fun):
        """Parse fun arguments that are expected to have units.
        """
        argnames = list(fun.__code__.co_varnames)
        def checkandstrip(val, unit, name, assignedunits):
            """Check units of a given argument and strip if needed.
            """
            dim = unit.dimensions
            try:
                # Try to get argument units
                valdim = val.units.dimensions
            except AttributeError:
                # Argument doesn't have units, assign and add warning
                val = val*unit
                valdim = dim
                assignedunits[name] = unit

            if valdim != dim:
                raise ValueError(
                    "\"%s\" has incorrect dimensions: expected %s but got %s" %
                    (name, dim, valdim))
            if strip:
                return val.to(unit).v
            return val.to(unit)

        @wraps(fun)
        def wrapper(*args, **kwargs):
            """Replace args and kwargs with parsed arguments when necessary.
            """
            parsedargs = [None]*len(args)
            assignedunits = {} # Dimensionless units with assigned units
            for i, name in enumerate(argnames):
                if i == len(args): break
                if name in units:
                    parsedargs[i] = checkandstrip(args[i], units[name],
                                                  name, assignedunits)
                else:
                    parsedargs[i] = args[i]
            for name, val in kwargs.items():
                if name in units:
                    kwargs[name] = checkandstrip(kwargs[name], units[name],
                                                 name, assignedunits)
            if len(assignedunits) > 0:
                msg1 = "Argument(s) "
                msg2 = " given without dimensions (assumed "
                for name, unit in assignedunits.items():
                    msg1 += name + ", "
                    msg2 += str(unit) + ", "
                msg = msg1[:-2] + msg2[:-2] + ")"
                warnings.warn(msg, AscotUnitWarning, stacklevel=2)

            return fun(*parsedargs, **kwargs)

        return wrapper


    return actualdecorator

a5py/a5py/physlib/__init__.py
ascot4fusion/ascot5
"""
Module containing widely used synonyms for different quantities.

Examples:
    Check if quantity is mass:
    isalias("m", "mass")
    True

    Get standard name of the quantity:
    getalias("q")
    "charge"

Note that the synonyms are case insensitive and ignore spaces, lowercases and
dashes.

File: alias.py
"""

## Standard name : list of synonyms dictionary. Flag indicates whether quantity
## is common for both gyro-orbit and guiding-center.
synonyms = {
    "ids"      : ["id"],
    "mass"     : ["m"],
    "charge"   : ["q"],
    "vnorm"    : ["v", "velocity", "velocitynorm", "|v|"],
    "bnorm"    : ["b", "magneticfield", "magneticfieldnorm", "|b|"],
    "energy"   : ["ekin", "e"],
    "pitch"    : ["xi"],
    "mu"       : ["magneticmoment"],
    "gamma"    : ["lorentzfactor", "relativisticfactor"],
    "phi"      : ["tor", "toroidalangle"],
    "phimod"   : ["tormod", "phi(mod)", "tor(mod)"],
    "theta"    : ["pol", "poloidalangle"],
    "thetamod" : ["polmod", "pol(mod)", "theta(mod)"],
    "zeta"     : ["gyroangle"],
    "vpar"     : ["parallelvelocity", "vparallel", "vpa", "vpara"],
    "ppar"     : ["parallelmomentum", "pparallel", "ppa", "ppara"],
    "vperp"    : ["perpendicularvelocity", "vperpendicular", "vpe"],
    "pparp"    : ["perpendicularmomentum", "pperpendicular", "ppe"],
    "pnorm"    : ["p", "momentum", "momentumnorm", "|p|"],
    "ptor"     : ["canonicaltoroidalmomentum"],
    "wall"     : ["lost"],
    "therm"    : ["thermalized"]
}


def isalias(name, standardname):
    """
    Check if name is an alias of the standardname.
    """

    name = name.lower()
    name = name.replace(" ", "")
    name = name.replace("-", "")
    name = name.replace("_", "")

    suffix = ""
    if len(name) > 3 and name[-3:] == "prt":
        suffix = "prt"
        name   = name[:-3]

    return name in synonyms[standardname]


def getalias(name):
    """
    Get standard name of the given quantity.
    """

    name = name.lower()
    name = name.replace(" ", "")
    name = name.replace("-", "")
    name = name.replace("_", "")

    suffix = ""
    if len(name) > 3 and name[-3:] == "prt":
        suffix = "prt"
        name   = name[:-3]

    for key in synonyms.keys():
        if name in synonyms[key]:
            return key + suffix

    return name + suffix

a5py/a5py/physlib/alias.py
ascot4fusion/ascot5
"""
Analytic tokamak equilibrium.

Solution is based on this article:  http://dx.doi.org/10.1063/1.3328818

Note that the x and y coordinates here does not correspond to R and z directly!

File analyticequilibrium.py
"""
import math as m
import numpy as np
import scipy.optimize
from numpy import log as log
from numpy import power as p


def psi(x,y,i):
    """
    Evaluate basis functions Eqs. (8) and (27).
    """
    return {
        0 : 1.0,
        1 : p(x,2),
        2 : p(x,2)*log(x) - p(y,2),
        3 : p(x,4) - 4*p(x,2)*p(y,2),
        4 : 3.0  * p(x,4.0) * log(x)
          - 9.0  * p(x,2.0) * p(y,2.0)
          - 12.0 * p(x,2.0) * log(x) * p(y,2.0)
          + 2.0  * p(y,4.0),
        5 : p(x,6) - 12*p(x,4)*p(y,2) + 8*p(x,2)*p(y,4),
        6 : 8*p(y,6) - 140*p(x,2)*p(y,4) - 120*p(x,2)*log(x)*p(y,4)
        + 180*p(x,4)*log(x)*p(y,2) + 75*p(x,4)*p(y,2) - 15*p(x,6)*log(x),
        7 : y,
        8 : y*p(x,2),
        9 : p(y,3)-3*y*p(x,2)*log(x),
        10: 3*y*p(x,4)-4*p(y,3)*p(x,2),
        11: 8*p(y,5)-45*y*p(x,4)-80*p(y,3)*p(x,2)*log(x)+60*y*p(x,4)*log(x),
        }[i]


def psix(x,y,i):
    """
    Evaluate x-derivative of the basis functions.
    """
    return {
        0 : 0.0,
        1 : 2*x,
        2 : 2*x*log(x)+x,
        3 : 4*x**3-8*x*y**2,
        4 : 12*x**3*log(x)+3*x**3-30*x*y**2-24*x*log(x)*y**2,
        5 : 6*x**5-48*x**3*y**2+16*x*y**4,
        6 : -400*x*y**4-240*x*log(x)*y**4+720*x**3*log(x)*y**2
        +480*x**3*y**2-90*x**5*log(x)-15*x**5,
        7 : 0,
        8 : 2*y*x,
        9 : -6*y*x*log(x)-3*y*x,
        10: 12*y*x**3-8*y**3*x,
        11: -120*y*x**3-160*y**3*x*log(x)-80*y**3*x+240*y*x**3*log(x),
        }[i]


def psixx(x,y,i):
    """
    Evaluate x^2-derivative of the basis functions.
    """
    return {
        0 : 0,
        1 : 2.0,
        2 : 2*log(x)+3,
        3 : 12*x**2-8*y**2,
        4 : 36*x**2*log(x)+21*x**2-54*y**2-24*log(x)*y**2,
        5 : 30*x**4-144*x**2*y**2+16*y**4,
        6 : -640*y**4-240*log(x)*y**4+2160*x**2*log(x)*y**2+2160*x**2*y**2
        -450*x**4*log(x)-165*x**4,
        7 : 0,
        8 : 2*y,
        9 : -6*y*log(x)-9*y,
        10: 36*y*x**2-8*y**3,
        11: -120*y*x**2-160*y**3*log(x)-240*y**3+720*y*x**2*log(x),
        }[i]


def psiy(x,y,i):
    """
    Evaluate y-derivative of the basis functions.
    """
    return {
        0 : 0,
        1 : 0,
        2 : -2*y,
        3 : -8*x**2*y,
        4 : -18*x**2*y-24*x**2*log(x)*y+8*y**3,
        5 : -24*x**4*y+32*x**2*y**3,
        6 : 48*y**5-560*x**2*y**3-480*x**2*log(x)*y**3+360*x**4*log(x)*y
        +150*x**4*y,
        7 : 1,
        8 : x**2,
        9 : 3*y**2-3*x**2*log(x),
        10: 3*x**4-12*y**2*x**2,
        11: 40*y**4-45*x**4-240*y**2*x**2*log(x)+60*x**4*log(x) ,
        }[i]


def psiyy(x,y,i):
    """
    Evaluate y^2-derivative of the basis functions.
    """
    return {
        0 : 0,
        1 : 0,
        2 : -2,
        3 : -8*x**2,
        4 : -18*x**2-24*x**2*log(x)+24*y**2,
        5 : -24*x**4+96*x**2*y**2,
        6 : 240*y**4-1680*x**2*y**2-1440*x**2*log(x)*y**2+360*x**4*log(x)
        +150*x**4,
        7 : 0,
        8 : 0,
        9 : 6*y,
        10: -24*y*x**2,
        11: 160*y**3-480*y*x**2*log(x),
        }[i]


def psipart(x,y,i):
    """
    This is the (first) part in Eq. (8) that the basis functions don't cover.
    """
    return {
        0 : 1.0/2*p(x,2)*log(x),
        1 : 1.0/8*p(x,4),
        }[i]


def psipartx(x,y,i):
    """
    Evaluate x-derivative of the first part in Eq. (8).
    """
    return {
        0 : x*log(x)+1.0/2*x,
        1 :  1.0/2*x**3
        }[i]


def psipartxx(x,y,i):
    """
    Evaluate x^2-derivative of the first part in Eq. (8).
    """
    return {
        0 : log(x)+3.0/2,
        1 : 3.0/2*x**2,
        }[i]


def psiparty(x,y,i):
    """
    Evaluate y-derivative of the first part in Eq. (8).
    """
    return 0


def psipartyy(x,y,i):
    """
    Evaluate y^2-derivative of the first part in Eq. (8).
    """
    return 0


def psiX(x,y,c0,c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,c11,A):
    return   c0*psix(x,y,0)     \
        +    c1*psix(x,y,1)     \
        +    c2*psix(x,y,2)     \
        +    c3*psix(x,y,3)     \
        +    c4*psix(x,y,4)     \
        +    c5*psix(x,y,5)     \
        +    c6*psix(x,y,6)     \
        +    c7*psix(x,y,7)     \
        +    c8*psix(x,y,8)     \
        +    c9*psix(x,y,9)     \
        +   c10*psix(x,y,10)    \
        +   c11*psix(x,y,11)    \
        +     A*psipartx(x,y,0) \
        + (1-A)*psipartx(x,y,1)


def psiXX(x,y,c0,c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,c11,zeta):
    return c1*2                                                      \
         + c2*(2*log(x) + 3)                                         \
         + c3*(12*x**2 - 8*y**2)                                     \
         + c4*(36*x**2*log(x) + 21*x**2 - 54*y**2 - 24*log(x)*y**2)  \
         + c5*(30*x**4 - 144*x**2*y**2 + 16*y**4)                    \
         + c6*( -640*y**4 - 240*log(x)*y**4 + 2160*x**2*log(x)*y**2  \
               + 2160*x**2*y**2 - 450*x**4*log(x) - 165*x**4)        \
         + cos(zeta)*log(x) + 3.0/2*cos(zeta) + 3.0/2*sin(zeta)*x**2


def psiY(x,y,c0,c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,c11,A):
    return     c0*psiy(x,y,0)    \
         +     c1*psiy(x,y,1)    \
         +     c2*psiy(x,y,2)    \
         +     c3*psiy(x,y,3)    \
         +     c4*psiy(x,y,4)    \
         +     c5*psiy(x,y,5)    \
         +     c6*psiy(x,y,6)    \
         +     c7*psiy(x,y,7)    \
         +     c8*psiy(x,y,8)    \
         +     c9*psiy(x,y,9)    \
         +    c10*psiy(x,y,10)   \
         +    c11*psiy(x,y,11)   \
         +     A*psiparty(x,y,0) \
         + (1-A)*psiparty(x,y,1)


def psiYY(x,y,c0,c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,c11,zeta):
    return -c2*2                                                     \
          - c3*8*x**2                                                \
          + c4*( -18*x**2 - 24*x**2*log(x) + 24*y**2)                \
          + c5*( -24*x**4 + 96*x**2*y**2)                            \
          + c6*(  240*y**4 - 1680*x**2*y**2 - 1440*x**2*log(x)*y**2
                + 360*x**4*log(x) + 150*x**4)


def psi0(x,y,c0,c1,c2,c3,c4,c5,c6,c7,c8,c9,c10,c11,A):
    """
    Evaluate the (total) psi function at the given position.
    """
    return   c0*psi(x,y,0)     \
        +    c1*psi(x,y,1)     \
        +    c2*psi(x,y,2)     \
        +    c3*psi(x,y,3)     \
        +    c4*psi(x,y,4)     \
        +    c5*psi(x,y,5)     \
        +    c6*psi(x,y,6)     \
        +    c7*psi(x,y,7)     \
        +    c8*psi(x,y,8)     \
        +    c9*psi(x,y,9)     \
        +   c10*psi(x,y,10)    \
        +   c11*psi(x,y,11)    \
        +     A*psipart(x,y,0) \
        + (1-A)*psipart(x,y,1) \


def find_axis(R0, z0, c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, A):
    """
    Find minimum of psi (corresponding to magnetic axis)
    """
    pfun = lambda x : psi0(x[0], x[1],
                           c0, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, A)
    return scipy.optimize.minimize(pfun, np.array([R0/R0, z0/R0]).ravel(),
                                   tol=1e-8).x


def analyticGS(A, R0, epsilon, kappa, delta, Xpointx, Xpointy, sym):
    """
    Solve the eqquilibrium coefficients based in simple input.

    This function accepts intuitive tokamak parameters to solve for
    the coefficients.

    Args:
        A : float <br>
            The free parameter in the paper.
        R0 : float <br>
            Major axis R coordinate (not the same as the magnetic axis) [m].
        epsilon : float <br>
            Inverse aspect ratio.
        kappa : float <br>
            Elongation.
        delta : float <br>
            Triangularity.
        Xpointx : float <br>
            Separatrix x-coordinate [m] (if non-symmetric Eq.).
        Xpointy : float <br>
            Separatrix y-coordinate [m] (if non-symmetric Eq.).
        sym : bool <br>
            Make up-down symmetric background.

    Returns:
        The c coefficients.
    """

    # alpha as defined in the article @onko alpha radiaaneissa?
    alpha = m.asin(delta)

    C = aGS_allInputs(
        A                    = A,
        outerEqPoint         = [1+epsilon, 0],
        innerEqPoint         = [1-epsilon, 0],
        upperHighpoint       = [1-epsilon*delta, kappa*epsilon],
        lowerXPoint          = [Xpointx, Xpointy],
        outerSlope           = 0,
        innerSlope           = 0,
        curvatureOuterEq     = -(1+alpha)**2/(epsilon*kappa**2),
        curvatureInnerEq     = (1-alpha)**2/(epsilon*kappa**2),
        curvatureHighpointEq = -kappa/(epsilon*(m.cos(alpha))**2),
        sym                  = sym)

    return C


def aGS_allInputs(A, outerEqPoint, innerEqPoint, upperHighpoint, lowerXPoint,
                  outerSlope, innerSlope, curvatureOuterEq, curvatureInnerEq,
                  curvatureHighpointEq, sym):
    """
    Solve the equilibrium coefficients from the given boundary conditions.

    Args:
        outerEqPoint : float <br>
            xy-coordinates of the outer equatorial point at separatrix.
        innerEqPoint : float <br>
            xy-coordinates of the inner equatorial point at separatrix.
        upperHighPoint : float <br>
            xy-coordinates of the highest separatrix point.
        lowerXpoint : float <br>
            Outer equatorial point slope.
        outerSlope : float <br>
            Outer equatorial point slope.
        innerSlope : float <br>
            Inner equatorial point slope.
        curvatureOuterEq : float <br>
            Curvature at the outboard equatorial plane.
        curvatureInnerEq : float <br>
            Curvature at the inboard equatorial plane.
        curvatureHighPointEq : float <br>
            Curvature at the top point.
        sym : float <br>
            Make up-down symmetric background.

    Returns:
        The c coefficients.
    """

    # Rename A as Acoeff
    Acoeff = A

    ###########################################################################
    #                                                                         #
    #   Construct the matrix A of the boundary conditions for the funtions    #
    #   which are solutions to the homogeneous equation                       #
    #                                                                         #
    ###########################################################################
    if sym:
        A = np.zeros((7,7))
        for ipsi in range(7):
            A[0][ipsi] = psi(outerEqPoint[0], outerEqPoint[1], ipsi)
            A[1][ipsi] = psi(innerEqPoint[0], innerEqPoint[1], ipsi)
            A[2][ipsi] = psi(upperHighpoint[0], upperHighpoint[1], ipsi)
            A[3][ipsi] = psix(upperHighpoint[0], upperHighpoint[1], ipsi)
            A[4][ipsi] = psiyy(outerEqPoint[0], outerEqPoint[1], ipsi) \
                         + curvatureOuterEq * psix(outerEqPoint[0],
                                                   outerEqPoint[1], ipsi)
            A[5][ipsi] = psiyy(innerEqPoint[0], innerEqPoint[1], ipsi) \
                         + curvatureInnerEq * psix(innerEqPoint[0],
                                                   innerEqPoint[1], ipsi)
            A[6][ipsi] = psixx(upperHighpoint[0], upperHighpoint[1], ipsi) \
                         + curvatureHighpointEq*psiy(upperHighpoint[0],
                                                     upperHighpoint[1], ipsi)
    else:
        A = np.zeros((12,12))
        for ipsi in range(12):
            A[0][ipsi] = psi(outerEqPoint[0], outerEqPoint[1], ipsi)
            A[1][ipsi] = psi(innerEqPoint[0], innerEqPoint[1], ipsi)
            A[2][ipsi] = psi(upperHighpoint[0], upperHighpoint[1], ipsi)
            A[3][ipsi] = psi(lowerXPoint[0], lowerXPoint[1], ipsi)
            A[4][ipsi] = psix(outerEqPoint[0],
                              outerEqPoint[1], ipsi) * outerSlope \
                       + psiy(outerEqPoint[0], outerEqPoint[1], ipsi)
            A[5][ipsi] = psix(innerEqPoint[0],
                              innerEqPoint[1], ipsi) \
                       * innerSlope+psiy(innerEqPoint[0], innerEqPoint[1], ipsi)
            A[6][ipsi] = psix(upperHighpoint[0], upperHighpoint[1], ipsi)
            A[7][ipsi] = psix(lowerXPoint[0], lowerXPoint[1], ipsi)
            A[8][ipsi] = psiy(lowerXPoint[0], lowerXPoint[1], ipsi)
            A[9][ipsi] = psiyy(outerEqPoint[0], outerEqPoint[1], ipsi) \
                         + curvatureOuterEq*psix(outerEqPoint[0],
                                                 outerEqPoint[1], ipsi)
            A[10][ipsi]= psiyy(innerEqPoint[0], innerEqPoint[1], ipsi) \
                         + curvatureInnerEq*psix(innerEqPoint[0],
                                                 innerEqPoint[1], ipsi)
            A[11][ipsi]= psixx(upperHighpoint[0], upperHighpoint[1], ipsi) \
                         + curvatureHighpointEq*psiy(upperHighpoint[0],
                                                     upperHighpoint[1], ipsi)

    ###########################################################################
    #                                                                         #
    #   Construct the matrix B of the boundary conditions for the particular  #
    #   solutions to the equation                                             #
    #                                                                         #
    ###########################################################################

    if sym:
        B = np.zeros((7,1))
        f = np.array([Acoeff, 1-Acoeff])
        for ipart in range(2):
            B[0] = B[0] + f[ipart]*psipart(outerEqPoint[0],
                                           outerEqPoint[1], ipart)
            B[1] = B[1] + f[ipart]*psipart(innerEqPoint[0],
                                           innerEqPoint[1], ipart)
            B[2] = B[2] + f[ipart]*psipart(upperHighpoint[0],
                                           upperHighpoint[1], ipart)
            B[3] = B[3] + f[ipart]*psipartx(upperHighpoint[0],
                                            upperHighpoint[1], ipart)
            B[4] = B[4] + f[ipart] * (
                psipartyy(outerEqPoint[0], outerEqPoint[1], ipart)
                + curvatureOuterEq*psipartx(outerEqPoint[0],
                                            outerEqPoint[1], ipart) )
            B[5] = B[5] + f[ipart]*(
                psipartyy(innerEqPoint[0], innerEqPoint[1], ipart)
                + curvatureInnerEq*psipartx(innerEqPoint[0],
                                            innerEqPoint[1], ipart) )
            B[6] = B[6] + f[ipart]*(
                psipartxx(upperHighpoint[0], upperHighpoint[1], ipart)
                + curvatureHighpointEq*psiparty(upperHighpoint[0],
                                                upperHighpoint[1], ipart) )
    else:
        B = np.zeros((12,1))
        f = np.array([Acoeff, 1-Acoeff])
        for ipart in range(2):
            B[0] = B[0] + f[ipart]*psipart(outerEqPoint[0],
                                           outerEqPoint[1], ipart)
            B[1] = B[1] + f[ipart]*psipart(innerEqPoint[0],
                                           innerEqPoint[1], ipart)
            B[2] = B[2] + f[ipart]*psipart(upperHighpoint[0],
                                           upperHighpoint[1], ipart)
            B[3] = B[3] + f[ipart]*psipart(lowerXPoint[0], lowerXPoint, ipart)
            B[4] = B[4] + f[ipart]*(
                outerSlope*psipartx(outerEqPoint[0], outerEqPoint[1], ipart)
                + psiparty(outerEqPoint[0], outerEqPoint[1], ipart) )
            B[5] = B[5] + f[ipart]*(
                outerSlope*psipartx(innerEqPoint[0], innerEqPoint[1], ipart)
                + psiparty(innerEqPoint[0], innerEqPoint[1], ipart) )
            B[6] = B[6] + f[ipart]*psipartx(upperHighpoint[0],
                                            upperHighpoint[1], ipart)
            B[7] = B[7] + f[ipart]*psipartx(lowerXPoint[0],
                                            lowerXPoint[1], ipart)
            B[8] = B[8] + f[ipart]*psiparty(lowerXPoint[0],
                                            lowerXPoint[1], ipart)
            B[9] = B[9] + f[ipart]*(
                curvatureOuterEq*psipartx(outerEqPoint[0],
                                          outerEqPoint[1], ipart)
                + psipartyy(outerEqPoint[0], outerEqPoint[1], ipart) )
            B[10] = B[10] + f[ipart]*(
                curvatureInnerEq*psipartx(innerEqPoint[0],
                                          innerEqPoint[1], ipart)
                + psipartyy(innerEqPoint[0], innerEqPoint[1], ipart) )
            B[11] = B[11] + f[ipart]*(
                curvatureHighpointEq*psiparty(upperHighpoint[0],
                                              upperHighpoint[1], ipart)
                + psipartxx(upperHighpoint[0], upperHighpoint[1], ipart) )

    B = np.negative(B)

    ###########################################################################
    #                                                                         #
    #   Solve the linear system for the coefficients C of the general         #
    #   solution to the equation                                              #
    #                                                                         #
    ###########################################################################


    C = np.linalg.lstsq(A, B, rcond=None)[0].flatten()

    if sym:
        C = np.append(C,[0,0,0,0,0])

    return C


def plot_equilibrium(C, A, R0, epsilon, kappa):
    """
    Plot the analytical equilibrium.

    Args:
        C : float, array_like <br>
            The c coefficients.
        A : float <br>
            The free parameter in the paper.
        R0 : float <br>
            Major axis R coordinate (not the same as the magnetic axis) [m].
        epsilon : float <br>
            Inverse aspect ratio.
        kappa : float <br>
            Elongation.
    """
    import matplotlib.pyplot  as plt

    R = np.linspace((1-epsilon-0.1), 1+epsilon+0.1, 200)
    Z = np.linspace(-1.2*kappa*epsilon-0.1, 1.2*kappa*epsilon+0.1, 200)
    X,Y = np.meshgrid(R,Z)
    X[X < 0.2] = 0

    Z = psi0(X, Y, C[0], C[1], C[2], C[3], C[4], C[5], C[6], C[7], C[8],
             C[9], C[10], C[11], A)

    plt.pcolormesh(X*R0, Y*R0, Z)
    plt.contour(X*R0, Y*R0, Z, [0])
    plt.plot([R0, R0],
             np.array([-1.5*kappa*epsilon, 1.5*kappa*epsilon])*R0, '--')
    plt.axis('equal')
    plt.show(block=False)

    return C


if __name__ == "__main__":

    # Write and plot?
    write = False
    plot  = True

    # Choose a machine
    method  = "circular"

    # Free parameter from the paper
    # A = 1 # Force-free equilibrium
    # A = 0 # Vacuum toroidal field ("low-beta" equilibrium)
    A       = 0

    R0      = 0 # Major axis
    epsilon = 0 # Inverse aspect ratio
    kappa   = 0 # Elongation
    delta   = 0 # Triangularity
    xsep    = 0 # x-location of the separatrix
    ysep    = 0 # y-location of the separatrix
    D       = 0 # Values of the contours for plotting
    sym     = 0 # Make symmetric equilibrium

    if method == "itertok":
        A       = -0.155
        R0      = 6.2
        epsilon = 2/6.2
        kappa   = 1.7
        delta   = 0.33
        xsep    = 1-1.1*delta*epsilon
        ysep    = -1.1*kappa*epsilon
        sym     = False
    elif method == "nstx":
        A       = -0.155
        R0      = 0.86
        epsilon = 0.67/0.86
        kappa   = 2
        delta   = 0.35
        xsep    = 1-1.1*delta*epsilon
        ysep    = -1.1*kappa*epsilon
        sym     = False
    elif method == "aug":
        A       = -0.155
        R0      = 1.65
        epsilon = 0.5/R0
        kappa   = 1.6
        delta   = 0.2
        ysep    = -0.9041 /R0
        xsep    = 1.4528 /R0
        sym     = False
    elif method == "circular":
        A       = -0.155
        R0      = 6.2
        epsilon = 1/R0
        kappa   = 1
        delta   = 0
        sym     = True

    C = analyticGS(A, R0, epsilon, kappa, delta, xsep, ysep, sym)
    C = np.append(C, A)

    if write:
        np.savetxt("analyticBKG.txt", C)

    if plot:
        plot_equilibrium(C[:-1], A, R0, epsilon, kappa)

a5py/a5py/physlib/analyticequilibrium.py
ascot4fusion/ascot5
import numpy as np


def evaluate(quantity, m=None, q=None, r=None, phi=None, z=None, pvec=None,
             mu=None, ppar=None, zeta=None, b=None):
    """
    Evaluate derived quantities from the given quantities.

    Note that it is enough to pass enough parameters that are needed to evaluate
    the queried quantity.
    """

    if quantity in ["xprt", "xgc"]:
        return coords.xcoord(r=r, phi=phi)

    if quantity in ["yprt", "ygc"]:
        return coords.ycoord(r=r, phi=phi)

    if quantity in ["energyprt"]:
        return energy_momentum(m=m, p=pvec)

    if quantity in ["energygc"]:
        return energy_muppar(m=m, mu=mu, ppar=ppar, b=b)

    if quantity in ["pitchprt"]:
        return pitch_momentum(p=pvec, b=b)

    if quantity in ["pitchgc"]:
        return pitch_muppar(m=m, mu=mu, ppar=ppar, b=b)

    if quantity in ["muprt"]:
        return mu_momentum(m=m, p=pvec, b=b)

    if quantity in ["pparprt"]:
        return mu_momentum(m=m, p=pvec, b=b)

    if quantity in ["vparprt"]:
        return mu_momentum(m=m, p=pvec, b=b)

    if quantity in ["pnormprt"]:
        return np.sqrt( np.sum(pvec**2, axis=1) )

    if quantity in ["vvecprt", "vrprt", "vphiprt", "vzprt", "vnormprt"]:
        vvec = velocity_momentum(pvec)
        if quantity == "vvecprt":
            return vvec
        if quantity == "vrprt":
            return vvec[0]
        if quantity == "vphiprt":
            return vvec[1]
        if quantity == "vzprt":
            return vvec[2]
        if quantity == "vnormprt":
            return np.sqrt( np.sum(vvec**2) )

    if quantity in ["vrgc", "vphigc", "vzgc", "vnormgc", "vvecgc"]:
        if quantity == "vvecgc":
            return vvec
        if quantity == "vrgc":
            return vvec[0]
        if quantity == "vphigc":
            return vvec[1]
        if quantity == "vzgc":
            return vvec[2]
        if quantity == "vnormgc":
            return np.sqrt( np.sum(vvec**2) )

a5py/a5py/physlib/evaluate.py
ascot4fusion/ascot5
"""List of commonly encountered marker species and their properties.
"""
import numpy as np
import unyt

def species(name, charge=None):
    """Assumes the particle is fully ionized. Charge state can be given
    explicitly with charge parameter.
    """

    data = None
    if name in ["e", "electron"]:
        data = {"anum":0, "znum":0, "charge":-1*unyt.e, "mass":0.0005486*unyt.amu}

    if name in ["n", "neutron"]:
        data = {"anum":1, "znum":0, "charge":0*unyt.e,  "mass":1.009*unyt.amu}

    if name in ["p", "proton", "H", "H1"]:
        data = {"anum":1, "znum":1, "charge":1*unyt.e,  "mass":1.007*unyt.amu}

    if name in ["deuterium", "D", "H2"]:
        data = {"anum":2, "znum":1, "charge":1*unyt.e,  "mass":2.014*unyt.amu}

    if name in ["tritium", "T", "H3"]:
        data = {"anum":3, "znum":1, "charge":1*unyt.e,  "mass":3.016*unyt.amu}

    if name in ["He4", "alpha"]:
        data = {"anum":4, "znum":2, "charge":2*unyt.e,  "mass":4.003*unyt.amu}

    if charge is not None:
        data["charge"] = charge*unyt.e

    return data

a5py/a5py/physlib/species.py
ascot4fusion/ascot5
from functools import wraps
import warnings

import unyt
import numpy as np

def parseunits(strip=False, **units):
    """Prepare arguments that are expected to have physical units.

    This decorator:

    - Makes sure every argument has expected dimensions.
    - Assigns units if units were not provided but they were expected.
    - Strips units if asked (after checking/assignment).

    Examples:

    :: code-block python

       @parseunits(x="m", strip=True)
       def fun(x):
           pass

    Parameters
    ----------
    strip : bool, optional
        Strip units so that dimensionless quantities (but in expected units)
        are passed to the function.
    **units
        Argument in wrapped function and the expected unit as a string.
    """
    for k in units.keys():
        # Convert unit strings to unyts
        units[k] = unyt.unyt_quantity.from_string(units[k]).units

    def actualdecorator(fun):
        """Parse fun arguments that are expected to have units.
        """
        argnames = list(fun.__code__.co_varnames)
        if "self" in argnames:
            argnames.remove("self")

        def checkandstrip(val, unit, name, assignedunits):
            """Check units of a given argument and strip if needed.
            """
            dim = unit.dimensions
            try:
                # Try to get argument units
                valdim = val.units.dimensions
            except AttributeError:
                # Argument doesn't have units, assign and add warning
                val = val*unit
                valdim = dim
                assignedunits[name] = unit

            if valdim != dim:
                raise ValueError(
                    "\"%s\" has incorrect dimensions: expected %s but got %s" %
                    (name, dim, valdim))
            if strip:
                return val.to(unit).v
            return val

        @wraps(fun)
        def wrapper(*args, **kwargs):
            """Replace args and kwargs with parsed arguments when necessary.
            """
            parsedargs = [None]*len(args)
            assignedunits = {} # Dimensionless units with assigned units
            for i, name in enumerate(argnames):
                if i == len(args): break
                if name in units:
                    parsedargs[i] = checkandstrip(args[i], units[name],
                                                  name, assignedunits)
                else:
                    parsedargs[i] = args[i]
            for name, val in kwargs.items():
                if name in units:
                    kwargs[name] = checkandstrip(kwargs[name], units[name],
                                                 name, assignedunits)
            if len(assignedunits) > 0:
                msg1 = "Argument(s) "
                msg2 = " given without dimensions (assumed "
                for name, unit in assignedunits.items():
                    msg1 += name + ", "
                    msg2 += str(unit) + ", "
                msg = msg1[:-2] + msg2[:-2] + ")"
                warnings.warn(msg, UserWarning)

            return fun(*parsedargs, **kwargs)

        return wrapper


    return actualdecorator

a5py/a5py/physlib/units.py
ascot4fusion/ascot5
"""Advanced pre- and postprocessing routines.
"""

a5py/a5py/routines/__init__.py
ascot4fusion/ascot5
"""Module for generating common plots with ASCOT5.
"""
import numpy as np
import warnings

try:
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    import mpl_toolkits
    from mpl_toolkits.axes_grid1 import make_axes_locatable
except ImportError:
    warnings.warn("Could not import matplotlib. Plotting disabled.")

try:
    import pyvista as pv
except ImportError:
    warnings.warn("Could not import pyvista. 3D wall plotting disabled.")

from functools import wraps

def setpaperstyle(height=5, halfpage=False):
    """TODO make this style suitable for publications
    """
    mpl.style.use({
        "figure.autolayout":False,
        "font.family":"serif",
        "font.serif":"ComputerModern",
        "text.usetex":True,
        "pdf.fonttype":42,
        "ps.fonttype":42,
        "axes.labelsize":14,
        "axes.titlesize":16,
        "axes.titlepad":6,
        "xtick.labelsize":12,
        "ytick.labelsize":12,
        "axes.labelpad":6,
        "legend.fontsize":12,
        "legend.numpoints":1,
        "legend.scatterpoints":1,
        "grid.linewidth":0.8,
        "lines.linewidth":1.4,
        "patch.linewidth":0.24,
        "lines.markersize":5.6,
        "lines.markeredgewidth":0,
        "xtick.major.width":0.8,
        "ytick.major.width":0.8,
        "xtick.minor.width":0.4,
        "ytick.minor.width":0.4,
        "xtick.major.pad":5.6,
        "ytick.major.pad":5.6,
        "axes.formatter.limits":[-1,1]
    })

def setguistyle():
    """Style used in GUI.
    """
    mpl.style.use({
        "figure.autolayout":False,
        "font.family":"serif",
        "font.serif":"ComputerModern",
        "text.usetex":True,
        "axes.labelsize":18,
        "axes.titlesize":18,
        "axes.titlepad":12,
        "xtick.labelsize":16,
        "ytick.labelsize":16,
        "axes.labelpad":6,
        "legend.fontsize":16,
        "legend.numpoints":1,
        "legend.scatterpoints":1,
        "grid.linewidth":0.8,
        "lines.linewidth":1.4,
        "patch.linewidth":0.24,
        "lines.markersize":5.6,
        "lines.markeredgewidth":0,
        "xtick.major.width":0.8,
        "ytick.major.width":0.8,
        "xtick.minor.width":0.4,
        "ytick.minor.width":0.4,
        "xtick.major.pad":5.6,
        "ytick.major.pad":5.6,
    })

def figuresinglecolumn(aspectratio=3/2):
    """Return figure that has a size suitable for printing in A4 single-column
    width (when the paper has a double-column format).

    Parameters
    ----------
    aspectratio : float
        Width / height ratio of the returned figure.

    Returns
    -------
    """
    return plt.figure(figsize=(3.504, 3.504/aspectratio))

def figuredoublecolumn(aspectratio=3/2):
    """Return figure that has a size suitable for printing in A4 double-column
    width (when the paper has a double-column format).
    """
    return plt.figure(figsize=(7.205, 7.205/aspectratio))

def openfigureifnoaxes(projection="rectilinear"):
    """Decorator for creating and displaying a new figure if axes are not
    provided.

    Parameters
    ----------
    projection : str, {None, '3d', 'aitoff', 'hammer', 'lambert', 'mollweide',
    'polar', 'rectilinear'}, optional
        The axes projection type, see Matplotlib documentation for details.
    """
    def actualdecorator(plotfun):
        """Decorator that takes the plotting routine.
        """
        @wraps(plotfun)
        def wrapper(*args, axes=None, **kwargs):
            """Create a new figure if axes is None and pass *args and **kwargs
            for the plotter.
            """
            if axes is None:
                fig  = plt.figure()
                axes = fig.add_subplot(111, projection=projection)
                plotfun(*args, axes=axes, **kwargs)
                plt.show()
            else:
                if projection != None and axes.name != projection:
                    raise ValueError(
                        "Invalid projection \"%s\" on axes: expected \"%s\"" %
                        (axes.name, projection))
                plotfun(*args, axes=axes, **kwargs)

        return wrapper

    return actualdecorator

@openfigureifnoaxes(projection=None)
def scatter2d(x, y, c=None, xlog="linear", ylog="linear", clog="linear",
              xlabel=None, ylabel=None, clabel=None, cint=9, cmap=None,
              axesequal=False, axes=None, cax=None):
    """Make a scatter plot in 2D+1 where color can be one dimension.

    For better performance this function uses matplotlib's plot function instead
    of the scatter function. The only practical difference is that the marker
    size can't vary and the color can't be continuous.

    Parameters
    ----------
    x : array_like
        Marker x-coordinates.
    y : array_like
        Marker y-coordinates.
    c : str or array_like, optional
        Color data or string indicating the color.
    xlog : {"linear", "log", "symlog"}, optional
        x-axis scaling.
    ylog : {"linear", "log", "symlog"}, optional
        y-axis scaling.
    clog : {"linear", "log", "symlog"}, optional
        color-axis scaling.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    clabel : str, optional
        Label for the color axis.
    cint : int, optional
        Number of colors used if c contains data. Since we are using plot
        instead of data, the color scale can't be continuous.
    cmap : str, optional
        Name of the colormap where nc colors are picked if c contains data.
    axesequal : bool, optional
        Flag to set aspect ratio of [x,y] axes equal.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    axes.set_xscale(xlog)
    axes.set_yscale(ylog)
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    if axesequal: axes.set_aspect("equal", adjustable="box")

    if c is None or isinstance(c, str):
        # Simple plot with a single color
        axes.plot(x, y, color=c, linestyle="None", marker="o")
        return

    # Sort inputs by color values and then find the indices that divide the
    # color range in given intervals
    idx  = np.argsort(c)
    x    = x[idx]
    y    = y[idx]
    c    = c[idx]
    if isinstance(cint, int):
        cint = np.linspace(c[0], c[-1], cint+1)
    idx  = np.searchsorted(c, cint) + 1
    nc = cint.size - 1

    # Choose default colormap depending on whether values change sign, and pick
    # nc colors
    if cmap is None:
        cmap = "viridis"
        if any(c < 0) and any(c > 0): cmap = "bwr"
    cmap = plt.get_cmap(cmap, nc)

    # Now plot markers corresponding to each interval with a different color
    i1 = 0
    for i in range(nc):
        i2 = idx[i+1]
        axes.plot(x[i1:i2], y[i1:i2], color=cmap(i/nc), linestyle="None",
                  marker="o")
        i1 = i2

    # Make colorbar with where color scale has the intervals
    norm = mpl.colors.BoundaryNorm(cint, nc)
    smap = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    cbar = plt.colorbar(smap, ax=axes, cax=cax)
    ticks = [];
    for b in cint.v:
        log = np.floor(np.log10(np.abs(b)))
        mul = b / 10**log
        ticks.append(r"$%.2f\times 10^{%.0f}$" % (mul, log))
    cbar.ax.set_yticklabels(ticks)
    cbar.set_label(clabel)

@openfigureifnoaxes(projection="3d")
def scatter3d(x, y, z, c=None, xlog="linear", ylog="linear", zlog="linear",
              clog="linear", xlabel=None, ylabel=None, zlabel=None, clabel=None,
              cint=None, cmap=None, axesequal=False, axes=None, cax=None):
    """Make a scatter plot in 3D+1 where color can be one dimension.

    For better performance this function uses matplotlib's plot function instead
    of the scatter function. The only practical difference is that the marker
    size can't vary and the color can't be continuous.

    Parameters
    ----------
    x : array_like
        Marker x-coordinates.
    y : array_like
        Marker y-coordinates.
    z : array_like
        Marker z-coordinates.
    c : str or array_like, optional
        Color data or string indicating the color.
    xlog : {"linear", "log", "symlog"}, optional
        x-axis scaling.
    ylog : {"linear", "log", "symlog"}, optional
        y-axis scaling.
    zlog : {"linear", "log", "symlog"}, optional
        z-axis scaling.
    clog : {"linear", "log", "symlog"}, optional
        color-axis scaling.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    zlabel : str, optional
        Label for the z-axis.
    clabel : str, optional
        Label for the color axis.
    cint : int, optional
        Number of colors used if c contains data. Since we are using plot
        instead of data, the color scale can't be continuous.
    cmap : str, optional
        Name of the colormap where nc colors are picked if c contains data.
    axesequal : bool, optional
        Flag to set aspect ratio of [x,y] axes equal.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    axes.set_xscale(xlog)
    axes.set_yscale(ylog)
    axes.set_zscale(zlog)
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    axes.set_zlabel(zlabel)
    if axesequal: axes.set_aspect("equal", adjustable="box")

    if c is None or isinstance(c, str):
        # Simple plot with a single color
        axes.plot(x, y, z, color=c, linestyle="None", marker="o")
        return

    # Sort inputs by color values and then find the indices that divide the
    # color range in even intervals
    idx  = np.argsort(c)
    x    = x[idx]
    y    = y[idx]
    z    = z[idx]
    c    = c[idx]
    if cint is None:
        cint = np.linspace(c[0], c[-1], 10)
    elif isinstance(cint, int):
        cint = np.linspace(c[0], c[-1], cint+1)
    idx  = np.searchsorted(c, cint) + 1
    nc = cint.size - 1

    # Choose default colormap depending on whether values change sign, and pick
    # nc colors
    if cmap is None:
        cmap = "viridis"
        if any(c < 0) and any(c > 0): cmap = "bwr"
    cmap = plt.get_cmap(cmap, nc)

    # Now plot markers corresponding to each interval with a different color
    i1 = 0
    for i in range(nc):
        i2 = idx[i+1]
        axes.plot(x[i1:i2], y[i1:i2], z[i1:i2], color=cmap(i/nc),
                  linestyle="None", marker="o")
        i1 = i2

    # Make colorbar with where color scale has the intervals
    norm = mpl.colors.BoundaryNorm(cint, nc)
    smap = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    cbar = plt.colorbar(smap, ax=axes, cax=cax)

    ticks = [];
    for b in cint.v:
        log = np.floor(np.log10(np.abs(b)))
        mul = b / 10**log
        ticks.append(r"$%.2f\times 10^{%.0f}$" % (mul, log))
    cbar.ax.set_yticklabels(ticks)
    cbar.set_label(clabel)

@openfigureifnoaxes(projection=None)
def hist1d(x, xbins=None, weights=None, xlog="linear", logscale=False,
           xlabel=None, legend=None, axes=None):
    """Plot (stacked) marker histogram in 1D.

    Parameters
    ----------
    x : array_like or list [array_like]
        Array or a list of arrays to be binned and plotted.

        If list is given, the resulting histogram is stacked with each
        array in the list corresponding to one layer in the stacked histogram.
    xbins : int or array_like, optional
        Number of bins or array storing bin edges for the x coordinate.
    weights : array_like or list [array_like], optional
        Values the datapoints are weighted with.

        Same format as for x.
    xlog : {"linear", "log"}, optional
        x-axis scaling.
    logscale : bool, optional
        Show histogram in logarithmic scale.
    xlabel : str, optional
        Label for the x-axis.
    legend : str, array_like
        List of strings to label the data with.

        The length of the list must be same as the number of data arrays
        provided.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    """
    axes.set_xlabel(xlabel)
    axes.set_xscale(xlog)
    ylabel = "Markers per bin" if weights is None else "Particles per bin"
    axes.set_ylabel(ylabel)
    if not logscale:
        axes.ticklabel_format(style="sci", axis="y", scilimits=(0,0))

    if isinstance(xbins, int) and xlog == "log":
        xmin = x[0][0]; xmax = x[0][0]
        for i in range(len(x)):
            xmin = np.amin([np.amin(x[i]), xmin])
            xmax = np.amax([np.amax(x[i]), xmax])
        xbins = np.logspace(np.log10(xmin), np.log10(xmax), xbins)

    # Plot and legend
    axes.hist(x, xbins, density=False, stacked=True, log=logscale,
              weights=weights, rwidth=2)
    axes.legend(legend, frameon=False)

@openfigureifnoaxes(projection=None)
def hist2d(x, y, xbins=None, ybins=None, weights=None, xlog="linear",
           ylog="linear", logscale=False, xlabel=None, ylabel=None,
           axesequal=False, axes=None, cax=None):
    """Plot marker histogram in 2D.

    Parameters
    ----------
    x : array_like, (n,)
        x-coordinates of the data to be binned and plotted.
    y : array_like, (n,)
        y-coordinates of the data to be binned and plotted.
    xbins : int or array_like, optional
        Number of bins or array storing bin edges for the x coordinate.
    ybins : int or array_like, optional
        Number of bins or array storing bin edges for the y coordinate.
    weights : array_like, (n,), optional
        Values the datapoints are weighted with.

        If weights are included, the colorbar label changes from "Markers"
        to "Particles".
    xlog : {"linear", "log"}, optional
        x-axis scaling.
    ylog : {"linear", "log"}, optional
        y-axis scaling.
    logscale : bool, optional
        Show histogram in logarithmic scale.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    axesequal : bool, optional
        Flag to set the aspect ratio equal.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    axes.set_xlabel(xlabel)
    axes.set_xscale(xlog)
    axes.set_ylabel(ylabel)
    axes.set_yscale(ylog)
    clabel = "Markers per bin" if weights is None else "Particles per bin"
    if axesequal: axes.set_aspect("equal", adjustable="box")

    if isinstance(xbins, int) and xlog == "log":
        xmin = np.amin(x)
        xmax = np.amax(x)
        xbins = np.logspace(np.log10(xmin), np.log10(xmax), xbins)

    if isinstance(ybins, int) and ylog == "log":
        ymin = np.amin(y)
        ymax = np.amax(y)
        ybins = np.logspace(np.log10(ymin), np.log10(ymax), ybins)

    h,_,_,m = axes.hist2d(x, y, bins=[xbins, ybins], weights=weights)

    norm = None
    if logscale: norm = mpl.colors.LogNorm(np.amin(h), np.amax(h))
    cbar = plt.colorbar(m, norm=norm, ax=axes, cax=cax)
    cbar.set_label(clabel)

@openfigureifnoaxes(projection=None)
def mesh1d(x, y, log=False, xlabel=None, ylabel=None, axes=False):
    """Plot 1D distribution.

    Parameters
    ----------
    x : array_like (nx,)
        Abscissa edges for the x-axis.
    z : array_like (nx-1,)
        Data to be plotted.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    """
    xc = np.zeros((y.size*2,))
    xc[1:-1:2] = x[1:-1]
    xc[2:-1:2] = x[1:-1]
    xc[0]      = x[0]
    xc[-1]     = x[-1]
    yc = np.zeros((y.size*2,))
    yc[::2]    = y
    yc[1::2]   = y

    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    axes.set_xlim(x[0], x[-1])
    axes.plot(xc, yc)

@openfigureifnoaxes(projection=None)
def mesh2d(x, y, z, log=False, diverging=False, xlabel=None, ylabel=None,
           clabel=None, clim=None, cmap=None, axesequal=False,
           axes=None, cax=None):
    """Make a mesh (surface) plot in 2D.

    Parameters
    ----------
    x : array_like (nx,) or (nx+1,)
        Abscissa or abscissa edges for the x-axis.
    y : array_like (ny,) or (ny+1,)
        Abscissa or abscissa edges for the y-axis.
    z : array_like (nx,ny)
        Data to be plotted.
    log : bool, optional
        Make color axis logarithmic.
    diverging : bool, optional
        Use diverging colormapping which is centered at zero.

        Works with logarithmic scale as well.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    clabel : str, optional
        Label for the color-axis.
    clim : [float, float], optional
        Color [min, max] limits.
    cmap : str, optional
        Colormap.
    axesequal : bool, optional
        Flag to set the aspect ratio equal.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    if clim is None: clim = [None, None]
    if clim[0] is None:
        clim[0] = np.nanmin(z)
    if clim[1] is None:
        clim[1] = np.nanmax(z)
    z = np.ma.masked_invalid(z)

    if log:
        if diverging:
            if cmap == None: cmap = "bwr"
            norm = mpl.colors.SymLogNorm(linthresh=10, linscale=1.0,
                                 vmin=clim[0], vmax=clim[1], base=10)
        else:
            if cmap == None: cmap = "viridis"
            norm = mpl.colors.LogNorm(vmin=clim[0], vmax=clim[1])
    else:
        if diverging:
            if cmap == None: cmap = "bwr"
            norm = mpl.colors.CenteredNorm(halfrange=np.amax(np.abs(clim)))
        else:
            if cmap == None: cmap = "viridis"
            norm = mpl.colors.Normalize(vmin=clim[0], vmax=clim[1])

    shading = "nearest" if z.shape == (x.size,y.size) else "flat"
    pcm = axes.pcolormesh(x, y, z.T, norm=norm, cmap=cmap, shading=shading)
    axes.patch.set(hatch='x', edgecolor=[0.9, 0.9, 0.9])

    axes.set_xlim(x[0], x[-1])
    axes.set_ylim(y[0], y[-1])

    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)

    if axesequal:
        axes.set_aspect("equal", adjustable="box")

    cbar = plt.colorbar(pcm, ax=axes, cax=cax)
    cbar.set_label(clabel)

@openfigureifnoaxes(projection=None)
def contour2d(x, y, z, contours, xlabel=None, ylabel=None, axesequal=False,
              colors=None, linestyles=None, linewidths=None, axes=None):
    """Plot contour on 2D plane.
    """
    axes.contour(x, y, z.T, contours, colors=colors, linestyles=linestyles,
                 linewidths=linewidths)
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    if axesequal:
        axes.set_aspect("equal", adjustable="box")

@openfigureifnoaxes(projection=None)
def arrows2d():
    """Plot vector field on 2D plane.
    """
    pass

@openfigureifnoaxes(projection=None)
def line2d(x, y, c=None, xlog="linear", ylog="linear", clog="linear",
           xlabel=None, ylabel=None, clabel=None, bbox=None,
           cmap=None, axesequal=False, axes=None, cax=None):
    """Plot line segments in 2D.

    x : array_like
        Marker x-coordinates.
    y : array_like
        Marker y-coordinates.
    c : str or array_like, optional
        Color data or string indicating the color.
    xlog : {"linear", "log", "symlog"}, optional
        x-axis scaling.
    ylog : {"linear", "log", "symlog"}, optional
        y-axis scaling.
    clog : {"linear", "log", "symlog"}, optional
        color-axis scaling.
    cmap : str, optional
        Name of the colormap.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    clabel : str, optional
        Label for the color axis.
    axesequal : bool, optional
        Flag to set aspect ratio of [x,y] axes equal.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    axes.set_xscale(xlog)
    axes.set_yscale(ylog)
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    if axesequal: axes.set_aspect("equal", adjustable="box")
    if c is None or isinstance(c, str):
        # Simple plot with a single color
        for i in range(len(x)):
            axes.plot(x[i], y[i], color=c)
        return

    if bbox is None:
        bbox = "[xmin, xmax, ymin, ymax, cmin, cmax]"
        raise ValueError("Specify bounding box: bbox="+bbox)
    axes.set_xlim(bbox[0], bbox[1])
    axes.set_ylim(bbox[2], bbox[3])
    if clog == "linear":
        if bbox[-2]*bbox[-1] < 0:
            if cmap == None: cmap = "bwr"
            norm = mpl.colors.CenteredNorm(halfrange=np.amax(np.abs(bbox[-2:])))
        else:
            if cmap == None: cmap = "viridis"
            norm = plt.Normalize(bbox[-2], bbox[-1])
    elif clog == "log":
        if cmap == None: cmap = "viridis"
        norm = mpl.colors.LogNorm(vmin=bbox[-2], vmax=bbox[-1])
    elif clog == "symlog":
        if cmap == None: cmap = "bwr"
        norm = mpl.colors.SymLogNorm(linthresh=10, linscale=1.0,
                                     vmin=bbox[-2], vmax=bbox[-1], base=10)

    for i in range(len(x)):
        points = np.array([x[i], y[i]]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)

        lc = mpl.collections.LineCollection(segments, cmap=cmap, norm=norm)
        lc.set_array(c[i][1:])
        line = axes.add_collection(lc)
    smap = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    plt.colorbar(smap, ax=axes, cax=cax)

@openfigureifnoaxes(projection="3d")
def line3d(x, y, z, c=None, xlog="linear", ylog="linear", zlog="linear",
           clog="linear", xlabel=None, ylabel=None, zlabel=None, clabel=None,
           bbox=None, cmap=None, axesequal=False, axes=None, cax=None):
    """Plot line segments in 3D.

    Parameters
    ----------
    x : array_like
        Marker x-coordinates.
    y : array_like
        Marker y-coordinates.
    z : array_like
        Marker z-coordinates.
    c : str or array_like, optional
        Color data or string indicating the color.
    xlog : {"linear", "log", "symlog"}, optional
        x-axis scaling.
    ylog : {"linear", "log", "symlog"}, optional
        y-axis scaling.
    zlog : {"linear", "log", "symlog"}, optional
        z-axis scaling.
    clog : {"linear", "log", "symlog"}, optional
        color-axis scaling.
    cmap : str, optional
        Name of the colormap.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    zlabel : str, optional
        Label for the z-axis.
    clabel : str, optional
        Label for the color axis.
    axesequal : bool, optional
        Flag to set aspect ratio of [x,y,z] axes equal.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    axes.set_xscale(xlog)
    axes.set_yscale(ylog)
    axes.set_zscale(zlog)
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)
    axes.set_zlabel(zlabel)
    if axesequal: axes.set_aspect("equal", adjustable="box")
    if c is None or isinstance(c, str):
        # Simple plot with a single color
        for i in range(len(x)):
            axes.plot(x[i], y[i], z[i], color=c)
        return

    if bbox is None:
        bbox = "[xmin, xmax, ymin, ymax, cmin, cmax]"
        raise ValueError("Specify bounding box: bbox="+bbox)
    axes.set_xlim(bbox[0], bbox[1])
    axes.set_ylim(bbox[2], bbox[3])
    axes.set_zlim(bbox[4], bbox[5])
    if clog == "linear":
        if bbox[-2]*bbox[-1] < 0:
            if cmap == None: cmap = "bwr"
            norm = mpl.colors.CenteredNorm(halfrange=np.amax(np.abs(bbox[-2:])))
        else:
            if cmap == None: cmap = "viridis"
            norm = plt.Normalize(bbox[-2], bbox[-1])
    elif clog == "log":
        if cmap == None: cmap = "viridis"
        norm = mpl.colors.LogNorm(vmin=bbox[-2], vmax=bbox[-1])
    elif clog == "symlog":
        if cmap == None: cmap = "bwr"
        norm = mpl.colors.SymLogNorm(linthresh=10, linscale=1.0,
                                     vmin=bbox[-2], vmax=bbox[-1], base=10)
    for i in range(len(x)):
        points = np.array([x[i], y[i], z[i]]).T.reshape(-1, 1, 3)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)

        lc = mpl_toolkits.mplot3d.art3d.Line3DCollection(
            segments, cmap=cmap, norm=norm)
        lc.set_array(c[i][1:])
        line = axes.add_collection(lc)
    smap = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    plt.colorbar(smap, ax=axes, cax=cax)

@openfigureifnoaxes(projection=None)
def poincare(x, y, ids, connlen=None, xlim=None, ylim=None, xlabel=None,
             ylabel=None, clabel=None, axesequal=False, axes=None, cax=None):
    """Poincaré plot where color separates markers or shows the connection
    length.

    For better performance this function uses matplotlib's plot function instead
    of the scatter function. The only practical difference is that the marker
    color can't be continuous.

    Parameters
    ----------
    x : array_like
        Orbit x-coordinates.
    y : array_like
        Orbit y-coordinates.
    ids : array_like
        Array of marker IDs showing to which marker the points in x and y
        arrays correspond to.
    connlen : array_like, optional
        Connection length at the position (x,y).

        Negative if the marker is confined. If given, the color scale shows
        connection length instead of marker ID. The confined markers are still
        shown with shades of red.
    xlim : tuple(float, float), optional
        Min and max values for the x-axis.
    ylim : tuple(float, float), optional
        Min and max values for the y-axis.
    xlabel : str, optional
        Label for the x-axis.
    ylabel : str, optional
        Label for the y-axis.
    clabel : str, optional
        Label for the color axis.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    nc = 5 # How many colors are used

    if connlen is not None:
        # If we don't have lost markers then ignore the connection length.
        if np.argwhere(connlen > 0).size == 0:
            connlen = None

    if connlen is None:
        # Shuffle colors and plot
        cmap = plt.cm.get_cmap("Reds", nc)
        uids  = np.unique(ids)
        cpick = np.arange(nc)
        np.random.shuffle(cpick)
        for i in range(nc):
            idx = np.in1d(ids, uids[i::nc])
            axes.plot(x[idx], y[idx], color=cmap(cpick[i]/nc),
                      linestyle="None", marker=".", markersize=1)

    else:
        # Sort by connection length (confined markers are indicated with a
        # negative connlen).
        idx = np.argsort(connlen)
        x   = x[idx]
        y   = y[idx]
        ids = ids[idx]
        connlen = connlen[idx]

        # Find where the line between confined and lost markers is.
        # Set connlen positive again for confined markers and rearrange again
        # by the connection length.
        idx    = np.argwhere(connlen > 0).ravel()[0]
        x      = np.append(x[idx:],      np.flip(x[:idx]))
        y      = np.append(y[idx:],      np.flip(y[:idx]))
        ids    = np.append(ids[idx:],    np.flip(ids[:idx]))
        connlen = np.append(connlen[idx:], np.flip(-connlen[:idx]))
        idx = (connlen.size - 1) - idx

        # The color has meaning only for lost markers so find the scale
        cmin = connlen[0]
        cmax = connlen[idx]

        logscale = False
        if cmin / cmax < 0.1:
            logscale = True
            cmin   = np.log10(cmin)
            cmax   = np.log10(cmax)
            connlen = np.log10(connlen)

        cmin = np.floor(cmin)
        cmax = np.ceil(cmax)
        nc_b = int(cmax - cmin)
        clim = np.linspace(cmin, cmax, nc_b+1)

        # Confined markers are plotted separately.
        connlen[idx+1:] = cmax + 1/nc
        clim = np.append(clim, cmax + np.linspace(0, 1/nc, nc))

        # Create colormap and colorbar
        cmapred  = plt.cm.get_cmap("Reds_r")
        cmapblue = plt.cm.get_cmap("Blues")
        colours  = [None] * (nc_b + nc)
        for i in range(nc_b):
            colours[i]  = cmapblue(i/nc_b)
        for i in range(nc):
            colours[i+nc_b] = cmapred(i/nc)

        cmap = mpl.colors.ListedColormap(colours)
        norm = mpl.colors.BoundaryNorm(boundaries=clim, ncolors=nc_b+nc )
        ticks = clim[:nc_b+1]
        if logscale:
            ticklabels = list(clim[:nc_b]) + [r"$\inf$"]
        else:
            ticklabels = list(clim[:nc_b]) + [r"$\inf$"]
        cb = plt.colorbar(mpl.cm.ScalarMappable(norm=norm, cmap=cmap), ax=axes,
                          spacing='proportional', extend='min', ticks=ticks,
                          boundaries=[clim[0]-0.5] + clim + [clim[-1]+0.5])
        cb.ax.set_yticklabels(ticklabels)
        if clabel is not None:
            clabel = "log10( "+clabel+" )" if logscale else clabel
            cb.ax.set_ylabel(clabel)

        # Find the indices that divide the color range in even intervals
        idx  = np.searchsorted(connlen, clim) + 1

        # Now plot markers on to each interval with a different color
        i1 = 0
        for i in range(nc_b):
            i2 = idx[i+1]
            axes.plot(x[i1:i2], y[i1:i2], color=colours[i],
                      linestyle="None", marker="o", markersize=1)
            i1 = i2

        # Plot confined
        uids  = np.unique(ids[i2:])
        cpick = np.arange(nc)
        np.random.shuffle(cpick)
        for i in range(nc):
            idx = np.in1d(ids, uids[i::nc])
            axes.plot(x[idx], y[idx], color=colours[nc_b+cpick[i]],
                      linestyle="None", marker=".", markersize=1)

    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)

    axes.set_xlim(xlim)
    axes.set_ylim(ylim)
    if axesequal:
        axes.set_aspect("equal", adjustable="box")

@openfigureifnoaxes(projection=None)
def still(wallmesh, points=None, data=None, log=False, cpos=None, cfoc=None,
          cang=None, axes=None, cax=None):
    """Take a still shot of the mesh and display it using matplotlib backend.

    The rendering is done using vtk but the vtk (interactive) window is not
    displayed. It is recommended to use the interactive plot to find desired
    camera position and produce the actual plot using this method. The plot
    is shown using imshow and acts as a regular matplotlib plot.

    Parameters
    ----------
    wallmesh : Polydata
        Mesh representing the wall.
    points : array_like, optional
        Array Npoint x 3 defining points (markers) to be shown. For
        each point [x, y, z] coordinates are given.
    cpos : array_like, optional
        Camera position coordinates [x, y, z].
    cfoc : array_like, optional
        Camera focal point coordinates [x, y, z].
    cang : array_like, optional
        Camera angle [azimuth, elevation, roll].
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    p = pv.Plotter(off_screen=True)
    if data is None:
        p.add_mesh(wallmesh, color=[0.9,0.9,0.9])
    else:
        cmap = mpl.colormaps["Reds"].copy()
        cmap.set_bad(color=[0.9,0.9,0.9])
        p.add_mesh(wallmesh, scalars=data, cmap=cmap, log_scale=log)
        p.remove_scalar_bar()

        maxval = np.nanmax(wallmesh.cell_data[data])
        minval = np.nanmin(wallmesh.cell_data[data])

    if points is not None:
        p.theme.color = 'black'
        p.add_points(points, render_points_as_spheres=True, point_size=10)

    # Set camera
    if cpos is not None:
        p.camera.position = cpos
    if cfoc is not None:
        p.camera.focal_point = cfoc
    if cang is not None:
        p.camera.azimuth   = cang[0]
        p.camera.elevation = cang[1]
        p.camera.roll      = cang[2]

    p.show()
    axes.imshow(p.image)
    axes.set_xticks([])
    axes.set_yticks([])

    if data is not None:
        if log:
            norm = mpl.colors.LogNorm(vmin=minval, vmax=maxval)
        else:
            norm = mpl.colors.Normalize(vmin=0, vmax=maxval)
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
        cbar = plt.colorbar(sm, ax=axes, cax=cax)
        #cbar.set_label(r"Energy load J/m$^2$")


def interactive(wallmesh, *args, points=None, data=None, log=False, cpos=None,
                cfoc=None, cang=None):
    """Open VTK window to display interactive view of the wall mesh.

    Parameters
    ----------
    wallmesh : Polydata
        Mesh representing the wall.
    *args : tuple (str, method), optional
        Key (str) method pairs. When key is pressed when the plot is
        displayed, the associated method is called. The method should take
        Plotter instance as an argument.
    points : array_like, optional
        Array Npoint x 3 defining points (markers) to be shown. For
        each point [x, y, z] coordinates are given.
    cpos : array_like, optional
        Camera position coordinates [x, y, z].
    cfoc : array_like, optional
        Camera focal point coordinates [x, y, z].
    cang : array_like, optional
        Camera angle [azimuth, elevation, roll].
    """
    p = pv.Plotter()

    if data is None:
        p.add_mesh(wallmesh, color=[0.9,0.9,0.9], log_scale=log)
    else:
        cmap = mpl.colormaps["Reds"].copy()
        cmap.set_bad(color=[0.9,0.9,0.9])
        p.add_mesh(wallmesh, scalars=data, cmap=cmap, log_scale=log)

    if points is not None:
        p.theme.color = 'black'
        p.add_points(points, render_points_as_spheres=True, point_size=10)

    # Set events
    for i in range(len(args)):
        def wrapper(*wargs):
            args[i][1](p)

        p.add_key_event(args[i][0], wrapper)

    # Set camera
    if cpos is not None:
        p.camera.position = cpos
    if cfoc is not None:
        p.camera.focal_point = cfoc
    if cang is not None:
        p.camera.azimuth   = cang[0]
        p.camera.elevation = cang[1]
        p.camera.roll      = cang[2]

    p.show()

@openfigureifnoaxes(projection=None)
def loadvsarea(wetted, loads, axes=None):
    """Plot histogram showing minimum load vs area.
    """
    idx = np.argsort(-loads)
    wetted = np.cumsum(wetted[idx])
    loads  = loads[idx]

    axes.set_xscale('log')
    axes.set_yscale('linear')
    axes.plot(loads, wetted)
    axes.set_xlabel(r"Load above [" + str(loads.units) + "]")
    axes.set_ylabel(r"Wetted area [" + str(wetted.units) + "]")

@openfigureifnoaxes(projection=None)
def triangularpatch(
        patches, color, log=False, xlim=None, ylim=None, clim=None, xlabel=None,
        ylabel=None, clabel=None, cmap=None, axes=None, cax=None):
    """Plot triangular patches.
    """
    if clim    is None: clim = [None, None]
    if clim[0] is None: clim[0] = np.nanmin(color)
    if clim[1] is None: clim[1] = np.nanmax(color)

    if log:
        norm = mpl.colors.LogNorm(clim[0], clim[1])
    else:
        norm = mpl.colors.Normalize(clim[0], clim[1])

    coll = mpl.collections.PolyCollection(patches, array=color, norm=norm,
                                          cmap=cmap)
    axes.add_collection(coll)

    axes.set_xlim(xlim)
    axes.set_ylim(ylim)
    axes.set_xlabel(xlabel)
    axes.set_ylabel(ylabel)

    plt.colorbar(coll, norm=norm, ax=axes, cax=cax, label=clabel)

@openfigureifnoaxes(projection="polar")
def momentumpolarplot(pnorm_edges, pitch_edges, dist, axes=None, cax=None):
    """Plot momentum space distribution in polar coordinates.

    Parameters
    ----------
    pnorm_edges : array_like
        Momentum abscissa edges.
    pitch_edges : array_like
        Pitch abscissa edges.
    dist : array_like
        Values of the distribution.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    cax : :obj:`~matplotlib.axes.Axes`, optional
        The color bar axes or otherwise taken from the main axes.
    """
    theta = np.arccos(pitch_edges)
    h = axes.pcolormesh(theta, pnorm_edges, dist)
    plt.colorbar(h, ax=axes, cax=cax)
    axes.set_thetamin(0)
    axes.set_thetamax(180)
    axes.set_xticks(np.array([0, 45, 90, 135, 180])*np.pi/180)
    axes.set_xticklabels([1.0, 0.5, 0.0, -0.5, -1.0])

@openfigureifnoaxes(projection=None)
def momentumpolargrid(pnorm_edges, pitch_edges, axes=None):
    """Plot momentum space polar coordinate grid in Cartesian basis.

    Parameters
    ----------
    pnorm_edges : array_like
        Momentum abscissa edges.
    pitch_edges : array_like
        Pitch abscissa edges.
    axes : :obj:`~matplotlib.axes.Axes`, optional
        The axes where figure is plotted or otherwise new figure is created.
    """
    ang = np.linspace(0, np.pi, 60)
    for v in pnorm_edges:
        axes.plot(v*np.cos(ang), v*np.sin(ang), color="black")

    p = pnorm_edges[-1]
    for v in pitch_edges:
        axes.plot([0, v*p], [0, np.sqrt(1.0 - v**2)*p], color="black")

def defaultcamera(wallmesh):
    """Get default camera (helper function for the 3D plots).

    Default camera is located at R = (Rmax + Rmin) / 2, phi = 0 deg,
    z = (zmax, zmin) / 2, where the min/max values are taken from the bounding
    box of the wall mesh. The focal point is at same (R, z) coordinates but
    phi = 10 deg. The camera angle parameters are all set to zero.

    Parameters
    ----------
    wallmesh : :obj:`~pyvista.Polydata`
        The wall mesh.

    Returns
    -------
    cpos : array_like
        Camera position coordinates [x, y, z].
    cfoc : array_like
        Camera focal point coordinates [x, y, z].
    cang : array_like
        Camera angle [azimuth, elevation, roll].
    """
    # Find min/max values
    Rmax = np.sqrt( np.amax(wallmesh.points[:,0]**2 + wallmesh.points[:,1]**2) )
    Rmin = np.sqrt( np.amin(wallmesh.points[:,0]**2 + wallmesh.points[:,1]**2) )
    zmax = np.amax( wallmesh.points[:,2] )
    zmin = np.amin( wallmesh.points[:,2] )

    # Camera position in (R,phi,z)
    cpos = np.array([( Rmax + Rmin ) / 2,  0, ( zmax + zmin ) / 2])
    cfoc = np.array([( Rmax + Rmin ) / 2, 10, ( zmax + zmin ) / 2])
    cang = np.array([0, 0, 0])

    # Coordinate transformation (R,phi,z) -> (x,y,z)
    cpos = np.array([cpos[0] * np.cos( np.pi * cpos[1] / 180 ),
                     cpos[0] * np.sin( np.pi * cpos[1] / 180 ),
                     cpos[2]])
    cfoc = np.array([cfoc[0] * np.cos( np.pi * cfoc[1] / 180 ),
                     cfoc[0] * np.sin( np.pi * cfoc[1] / 180 ),
                     cfoc[2]])

    return (cpos, cfoc, cang)

a5py/a5py/routines/plotting.py
ascot4fusion/ascot5
"""Wrapper class for RunNode that adds methods to plot and access results.

The motivation for this class is following. Consider that user wants to plot
a figure that shows i) particle orbits in (R,z), ii) their final (R,z) position,
and iii) the wall tiles. It is not clear whether this plot should be implemented
as a method for State or Orbit class as data from both are needed. Furthermore,
neither has access to the wall data. The only logical place for this method
therefore is in the RunNode that has access to all relevant data.
"""
import numpy as np
import warnings
try:
    import pyvista as pv
except ImportError:
    warnings.warn("Could not import pyvista. 3D wall plotting disabled.")
import unyt

from a5py.exceptions import AscotNoDataException

import a5py.routines.plotting as a5plt
import a5py.wall as wall
from a5py.ascot5io import Marker, State, Orbits, Dist
from a5py.ascot5io.dist import DistMoment
import a5py.physlib as physlib

class RunMixin():
    """Class with methods to access and plot orbit and state data.

    This class assumes it is inherited by ResultsNode.
    """

    def _require(self, *args):
        """Check if required data is present and raise exception if not.

        This is a helper function to quickly check that the data is available.

        Parameters
        ----------
        *args : `str`
            Name(s) of the required data.

        Raises
        ------
        AscotNoDataException
            Raised if the required data is not present.
        """
        for arg in args:
            if not hasattr(self, arg):
                raise AscotNoDataException(
                    "Data for \"" +  arg + "\" is required but not present.")

    def getstate_list(self):
        """List quantities that can be evaluated with :meth:`getstate`.
        """
        self._require("_inistate")
        return State.listqnts()

    def getorbit_list(self):
        """List quantities that can be evaluated with :meth:`getorbit`.
        """
        self._require("_orbit")
        qnts = Orbits.listqnts()
        if self.options.read()["ORBITWRITE_MODE"] != 0:
            del qnts["pncrid"]
            del qnts["pncrdir"]
        return qnts

    def getstate(self, *qnt, mode="gc", state="ini", ids=None, endcond=None):
        """Evaluate a marker quantity based on its ini/endstate.

        Inistate is marker's phase-space position right at the start of
        the simulation and endstate is the position at the end of
        the simulation.

        This function not only returns the marker phase space coordinates but
        also other quantities that can be inferred from it and information that
        is stored along with coordinates. For a complete list of available
        quantities, see.

        ASCOT5 stores both particle and guiding center phase-space position in
        all simulations. To differentiate these, quantities with suffix "prt",
        e.g. "xprt", return particle quantities and without suffix the guiding
        center quantity is returned.

        Parameters
        ----------
        *qnt : str
            Names of the quantities.
        state : {"ini", "end"}, optional
            Is the quantity evaluated at the ini- or endstate.
        ids : array_like, optional
            Filter markers by their IDs.
        endcond : str or list [str], optional
            Filter markers by their end conditions.

            See for a list of all possible end conditions or to list end
            conditions that are currently present in the data.

            Markers may have multiple end conditions active simultaneously. If
            just the name of the end condition e.g. "POLMAX" is passed, then all
            markers that have (at least) the ``POLMAX`` end condition are
            returned.

            If the end condition is preceded by "NOT", e.g. "NOT POLMAX", then
            markers that don't have that end condition are returned.

            Passing multiple end conditions in a single string returns markers
            that have all listed end conditions active, e.g. "MAXPOL MAXTOR"
            returns markers that have both ``POLMAX`` and ``TORMAX`` active
            simultaneously.

            Passing end condition strings as separate list items acts as
            a logical OR, e.g. ["POLMAX", "TORMAX"] returns markers that have
            either ``POLMAX`` or ``TORMAX`` active.

        Returns
        -------
        val : array_like
            The evaluated quantity sorted by marker ID.

        If multiple quantities are queried, they are returned as a list in
            the order they were listed in ``*qnt``.

        Raises
        ------
        ValueError
            Raised when the queried quantity could not be interpreted.
        AscotNoDataException
            Raised when data required for the operation is not present.
        AscotInitException
            If evaluating quantity required interpolating an input that
            was not initialized.
        """
        self._require("_inistate")
        if endcond is not None: self._require("_endstate")
        if state not in ["ini", "end"]:
            raise ValueError("Unrecognized state: " + state)
        if state == "end": self._require("_endstate")

        # Get or evaluate the quantity
        data = getattr(self, "_" + state + "state").get(*qnt, mode=mode)

        # Parse by ids and endcond
        idx = np.ones(data[0].shape, dtype=bool)
        if endcond is not None:
            if not isinstance(endcond, list): endcond = [endcond]

            # Go through each unique end cond and mark that end cond valid or
            # not. This can then be used to make udix as boolean mask array.
            uecs, uidx = np.unique(self._endstate.get("endcond"),
                                   return_inverse=True)
            mask = np.zeros(uecs.shape, dtype=bool)
            for i, uec in enumerate(uecs):
                for ec in endcond:
                    accept  = State.endcond_check(uec, ec)
                    mask[i] = mask[i] or accept

            idx = mask[uidx]

        if ids is not None:
            idx = np.logical_and(idx, np.in1d(self._inistate.get("ids"), ids))

        for i in range(len(data)):
            data[i] = data[i][idx]
        if "mu" in qnt:
            data[qnt.index("mu")].convert_to_units("eV/T")
        return data if len(data) > 1 else data[0]

    def getorbit(self, *qnt, ids=None, pncrid=None, endcond=None):
        """Return orbit data.

        Returns marker phase space coordinates and derived quantities along
        the recorded orbit, if the orbit recording was enabled.

        Parameters
        ----------
        *qnt : str
            Names of the quantities.
        ids : array_like, optional
            Filter markers by their IDs.
        pncrid : array_like, optional
            Filter data points by the Poincaré plane they correspond to.
        endcond : str or list [str], optional
            Filter markers by their end conditions.

            See :meth:`getstate` for details on how this argument is parsed and
            for a list of end conditions present in the data.

        Returns
        -------
        val : array_like
            The queried quantity sorted first by marker ID and then by mileage.

            If multiple quantities are queried, they are returned as a list in
            the order they were listed in ``*qnt``.

        Raises
        ------
        ValueError
            Raised when the queried quantity could not be interpreted.
        AscotNoDataException
            Raised when data required for the operation is not present.
        AscotInitException
            If evaluating quantity required interpolating an input that
            was not initialized.
        """
        self._require("_orbit", "_inistate", "_endstate")
        data = self._orbit.get(self._inistate, self._endstate, *qnt)
        idarr = self._orbit.get(self._inistate, self._endstate, "ids")[0]
        idx = np.ones(data[0].shape, dtype=bool)
        if endcond is not None:
            eids = self.getstate("ids", endcond=endcond)
            idx = np.logical_and(idx, np.in1d(idarr, eids))

        if pncrid is not None:
            pncridarr = self._orbit.get(self._inistate, self._endstate,
                                       "pncrid")[0]
            idx = np.logical_and(idx, np.in1d(pncridarr, pncrid))

        if ids is not None:
            idx = np.logical_and(idx, np.in1d(idarr, ids))

        for i in range(len(data)):
            data[i] = data[i][idx]
        if "mu" in qnt:
            data[qnt.index("mu")].convert_to_units("eV/T")
        return data if len(data) > 1 else data[0]

    def getstate_markersummary(self):
        """Return a summary of marker end conditions and errors present in
        the data.

        Returns
        -------
        econds : list [(int, str)]
            List of present end conditions.

            Each list member is a tuple, where first item is the number of
            markers with the end condition specified in the second item.

        emsg : list [(str, int, str)]
            List of present errors.

            Each list member is a tuple, where first item is the error
            message, second the line number where error was raised, and
            third is the name of the file.

        Raises
        ------
        AscotNoDataException
            Raised when data required for the operation is not present.
        """
        self._require("_endstate")
        econd, emsg, emod, eline = self._endstate.get(
            "endcond", "errormsg", "errormod", "errorline")
        errors = np.unique(np.array([emsg, eline, emod]), axis=1).T

        ec, counts = np.unique(econd, return_counts=True)
        econds = []
        for i, e in enumerate(ec):
            econd = State.endcond_tostring(e)
            if econd == "none": econd = "aborted"
            econds.append( (counts[i], econd) )

        # It would be better to import these via libascot, but then this
        # function would require libascot and we don't want that.
        modules = [
            "mccc_wiener.c", "mccc_push.c", "mccc.c", "step_fo_vpa.c",
            "step_gc_cashkarp.c", "step_gc_rk4", "N0_3D.c", "N0_ST.c",
            "B_2DS.c", "B_2DS.c", "B_STS.c", "B_GS.c", "plasma_1D.c",
            "plasma_1DS.c", "plasma.c", "E_field.c", "neutral.c",
            "E_1DS.c", "B_field.c", "particle.c", "boozer.c", "mhd.c"
        ]
        messages = [
            "Input evaluation failed", "Unknown input type",
            "Unphysical quantity when evaluating input",
            "Unphysical marker quantity", "Time step too small/zero/NaN",
            "Wiener array is full or inconsistent",
            "Unphysical result when integrating marker coordinates"
        ]

        emsg = []
        for e in errors:
            if np.sum(e) > 0:
                emsg.append( (messages[e[0]-1], e[1], modules[e[2]-1]) )

        return econds, emsg

    def getstate_losssummary(self):
        """Return a summary of lost markers.
        """
        self._require("_endstate")

        wmrks, emrks = self.getstate("weight", "ekin", state="ini")
        wloss, eloss = self.getstate("weight", "ekin", state="end",
                                     endcond="wall")

        markers_lost   = wloss.size
        markers_frac   = wloss.size / wmrks.size
        particles_lost = np.sum(wloss)
        particles_frac = np.sum(wloss) / np.sum(wmrks)
        power_lost     = np.sum(wloss * eloss.to("J"))
        power_frac     = np.sum(wloss * eloss.to("J")) / np.sum(wmrks * emrks)

        msg = []
        msg += ["Markers lost: " + str(markers_lost) + " ("
                + str(np.around(markers_frac*100, decimals=1)) + "% of total)"]
        msg += ["Particles lost: " + str(particles_lost) + " ("
                + str(np.around(particles_frac*100, decimals=1)) + "% of total)"]
        msg += ["Energy lost: " + str(power_lost) + " ("
                + str(np.around(power_frac*100, decimals=1)) + "% of total)"]
        return msg

    def getstate_pointcloud(self, endcond=None):
        """Return marker endstate (x,y,z) coordinates in single array.

        Parameters
        ----------
        endcond : str, optional
            Only return markers that have given end condition.
        """
        self._require("_endstate")
        return np.array([self._endstate.get("x", endcond=endcond),
                         self._endstate.get("y", endcond=endcond),
                         self._endstate.get("z", endcond=endcond)]).T

    def getstate_markers(self, mrktype, ids=None):
        """Convert endstate to marker input.

        Parameters
        ----------
        mrktype : {"prt", "gc", "fl"}
            Type of marker input to be created.
        ids : array_like, optional
            Select only these markers for the output.

        Returns
        -------
        mrk : dict
            Markers parameters that can be supplied to :meth:`Prt.write_hdf5`,
            :meth:`GC.write_hdf5` or :meth:`FL.write_hdf5` depending on
            ``mrktype`` value.
        """
        ids = self.getstate("ids", state="ini", ids=ids)
        mrk = Marker.generate(mrktype, ids.size)
        mrk["ids"] = ids
        if mrktype == "prt":
            qnt = ["r", "phi", "z", "weight", "time", "vr", "vphi", "vz",
                   "mass", "charge", "anum", "znum"]
            state = self.getstate(*qnt, mode="prt", state="end", ids=ids)
            for i in qnt:
                mrk[i] = qnt[i]
        elif mrktype == "gc":
            qnt = ["r", "phi", "z", "weight", "time", "ekin", "pitch", "zeta",
                   "mass", "charge", "anum", "znum"]
            state = self.getstate(*qnt, mode="gc", state="end", ids=ids)
            for i, q in enumerate(qnt):
                if q == "ekin":
                    mrk["energy"] = state[i]
                else:
                    mrk[q] = state[i]
        elif mrktype == "fl":
            qnt = ["r", "phi", "z", "weight", "time", "pitch"]
            state = self.getstate(*qnt, mode="gc", state="end", ids=ids)
            for i in qnt:
                mrk[i] = qnt[i]
        return mrk

    def getorbit_poincareplanes(self):
        """Return a list of Poincaré planes that were present in the simulation
        and the corresponding pncrid.

        Returns
        -------
        pol : list [(float, int)]
            List of tuples with a toroidal angle and pncrid for each poloidal
            plane or empty list if there are none.
        tor : list [float]
            List of tuples with a poloidal angle and pncrid for each toroidal
            plane or empty list if there are none.
        rad : list [float]
            List of tuples with radius and pncrid for radial surfaces or empty
            list if there are none.
        """
        opt  = self.options.read()
        hasdata = opt["ENABLE_ORBITWRITE"] == 1 and opt["ORBITWRITE_MODE"] == 0
        if not hasdata: return None

        polval = opt["ORBITWRITE_POLOIDALANGLES"]
        torval = opt["ORBITWRITE_TOROIDALANGLES"]
        radval = opt["ORBITWRITE_RADIALDISTANCES"]

        pncrid = 0
        def opt2list(val):
            """Convert the option string to a list with tuples (val, pncrid)
            """
            nonlocal pncrid
            tuples = []
            if val[0] < 0: val = []
            for i in range(len(val)):
                tuples.append((val[i], pncrid))
                pncrid += 1
            return tuples

        return opt2list(polval), opt2list(torval), opt2list(radval)

    def getorbit_average(self, qnt, ids):
        """Calculate average of a quantity during a single poloidal transit.

        Parameters
        ----------
        qnt : str
            Name of the averaged quantity.
        """
        qnt, mileage, r, z, p, pitch, pol = \
            self.getorbit(qnt, "mileage", "r", "z", "phi", "pitch", "theta",
                          ids=ids)
        if any(pitch < 0) and any(pitch > 0):
            if pitch[0] < 0 or pitch[1] < 0:
                i1 = np.argmax(pitch > 0)
                i2 = np.argmax(pitch[i1:] < 0) + i1
                i2 = np.argmax(pitch[i2:] > 0) + i2
            else:
                i1 = np.argmax(pitch < 0)
                i2 = np.argmax(pitch[i1:] > 0) + i1
                i2 = np.argmax(pitch[i2:] < 0) + i2
        else:
            i1 = 1
            if pol[-1] > pol[0]:
                i2 = np.argmax(pol > (pol[0] + 360*unyt.deg))
            else:
                i2 = np.argmax(pol < (pol[0] - 360*unyt.deg))
        qnt = qnt[i1:i2]
        mileage = mileage[i1-1:i2]
        r = r[i1-1:i2]
        z = z[i1-1:i2]
        p = p[i1-1:i2]
        x, y, _ = physlib.pol2cart(r, p)
        ds = np.sqrt(np.diff(x)**2 + np.diff(y)**2 + np.diff(z)**2)
        qnt = np.sum(qnt*ds) / np.sum(ds)
        return mileage[1:], r, z, qnt

    def getwall_figuresofmerit(self):
        """Get peak power load and other 0D quantities related to wall loads.

        Returns
        -------
        warea : float
            Total wetted area.
        pload : float
            Peak power load.
        """
        wetted, area, edepo, pdepo, iangle = self.getwall_loads()
        wetted_total = np.sum(area)
        energy_peak  = np.amax(edepo/area)
        return wetted_total, energy_peak

    def getwall_loads(self, weights=True):
        """Get wall loads and associated quantities.

        This method does not return loads on all wall elements (as usually most
        of them receive no loads) but only those that are affected and their
        IDs.

        Parameters
        ----------
        weights : bool, optional
            Include marker weights to get physical results (otherwise particle
            deposition would be just the number of markers that hit the tile).

            Dropping weights is useful to check how many markers hit a tile
            which tells us how good the statistics are.

        Returns
        -------
        ids : array_like
            Indices of loaded wall tiles.
        area : array_like
            Area of each tile.
        edepo : array_like
            Energy/power deposition per tile.
        pdepo : array_like
            Particle/flux deposition per tile.
        iangle : array_like
            Angle of incidence (TODO).
        """
        self._require("_endstate")
        ids, energy, weight = self.getstate("walltile", "ekin", "weight",
                                            state="end", endcond="wall")
        energy.convert_to_units("J")
        eunit = (energy.units * weight.units)
        try:
            eunit = (1 * eunit).to("W")
        except unyt.exceptions.UnitConversionError:
            pass
        wetted = np.unique(ids)
        area   = self.wall.area()[wetted-1]
        edepo  = np.zeros(wetted.shape) * eunit
        pdepo  = np.zeros(wetted.shape) * weight.units
        iangle = np.zeros(wetted.shape)

        # Sort markers by wall tile ID to make processing faster
        idx    = np.argsort(ids)
        ids    = ids[idx]
        energy = energy[idx]
        weight = weight[idx]

        idx = np.append(np.argwhere(ids[1:] - ids[:-1]).ravel(), ids.size-1)

        i0 = 0
        for i in range(wetted.size):
            i1 = idx[i] + 1
            pdepo[i]  = np.sum(weight[i0:i1])
            edepo[i]  = np.sum(energy[i0:i1]*weight[i0:i1])
            iangle[i] = 0 #TBD
            i0 = i1

        return wetted, area, edepo, pdepo, iangle

    def getwall_3dmesh(self):
        """Return 3D mesh representation of 3D wall and associated loads.

        Returns
        -------
        wallmesh : Polydata
            Mesh representing the wall.

            The mesh cell data has fields:

            - "pload" particle load in units of prt/m^2 or prt/m^2s,
            - "eload" power/energy load in units of W/m^2 or J/m^2
        """
        wallmesh = pv.PolyData( *self.wall.noderepresentation() )
        ids, _, eload, pload, iangle = self.getwall_loads()
        ids = ids - 1 # Convert IDs to indices
        ntriangle = wallmesh.n_faces
        wallmesh.cell_data["pload"]       = np.zeros((ntriangle,)) + np.nan
        wallmesh.cell_data["pload"][ids]  = pload
        wallmesh.cell_data["eload"]       = np.zeros((ntriangle,)) + np.nan
        wallmesh.cell_data["eload"][ids]  = pload

        #- "iangle" angle of incidence in units of deg
        #wallmesh.cell_data["iangle"]      = np.zeros((ntriangle,)) + np.nan
        #wallmesh.cell_data["iangle"][ids] = iangle
        return wallmesh

    def getdist(self, dist, exi=False, ekin_edges=None, pitch_edges=None,
                plotexi=False):
        """Return distribution function.

        Parameters
        ----------
        dist : {"5d", "6d", "rho5d", "rho6d", "com"}
            Name of the distribution.

            If ``dist`` is ``None``, returns a list of available distributions.
        exi : bool, optional
            Convert the momentum space to energy-pitch.

            The distribution is normalized to conserve the particle number.

            Not applicable for ``dist`` = {"6d", "rho6d", "com"}.
        ekin_edges : int or array_like, optional
            Number of bins or bin edges in the energy abscissa if ``exi=True``.
        pitch_edges : int or array_like, optional
            Number of bins or bin edges in the pitch abscissa if ``exi=True``.
        plotexi : bool, optional
            Visualize the transformation from ppar-perp to energy-pitch if
            if ``exi=True``.

            Use this option to adjust energy and pitch abscissae to your liking.

        Returns
        -------
        data : :class:`DistData` or list [str]
            The distribution data object or a list of available distributions if
            ``dist`` is ``None``.
        """
        if dist is None:
            dists = ["5d", "6d", "rho5d", "rho6d", "com"]
            for d in dists:
                try:
                    self._require("_dist" + d)
                except AscotNoDataException:
                    dists.remove(d)
            return dists

        mass = np.mean(self.getstate("mass"))
        if dist == "5d":
            self._require("_dist5d")
            distout = self._dist5d.get()
            if exi:
                if ekin_edges is None:
                    ekin_edges  = int(distout.abscissa("ppar").size / 2)
                if pitch_edges is None:
                    pitch_edges = distout.abscissa("pperp").size
                distout = Dist.ppappe2ekinpitch(
                    distout, mass, ekin_edges=ekin_edges,
                    pitch_edges=pitch_edges)
        elif dist == "6d":
            self._require("_dist6d")
            if exi:
                raise ValueError("Energy-pitch transformation not valid for 6d")
            distout = self._dist6d.get()
        elif dist == "rho5d":
            self._require("_distrho5d")
            distout = self._distrho5d.get()
            if exi:
                if ekin_edges is None:
                    ekin_edges  = int(distout.abscissa("ppar").size / 2)
                if pitch_edges is None:
                    pitch_edges = distout.abscissa("pperp").size
                distout = Dist.ppappe2ekinpitch(
                    distout, mass, ekin_edges=ekin_edges,
                    pitch_edges=pitch_edges)
        elif dist == "rho6d":
            self._require("_distrho6d")
            if exi:
                raise ValueError("Energy-pitch transformation not valid for 6d")
            distout = self._distrho6d.get()
        elif dist == "com":
            self._require("_distcom")
            if exi:
                raise ValueError(
                    "Energy-pitch transformation not valid for COM")
            distout = self._distcom.get()
        else:
            raise ValueError("Unknown distribution")

        if exi and plotexi:
            if dist == "5d":    dist0 = self._dist5d.get()
            if dist == "rho5d": dist0 = self._distrho5d.get()
            integrate = {}
            for k in dist0.abscissae:
                if k not in ["ppar", "pperp"]:
                    integrate[k] = np.s_[:]
            dist0.integrate(**integrate)
            integrate = {}
            for k in distout.abscissae:
                if k not in ["ekin", "pitch"]:
                    integrate[k] = np.s_[:]
            dist1 = distout.integrate(copy=True, **integrate)

            g = physlib.gamma_energy(mass, dist1.abscissa_edges("ekin"))
            pnorm_edges = physlib.pnorm_gamma(mass, g).to("kg*m/s")
            pitch_edges = dist1.abscissa_edges("pitch")

            import matplotlib.pyplot as plt
            fig = plt.figure()
            ax1 = fig.add_subplot(3,1,1)
            ax2 = fig.add_subplot(3,1,2, projection='polar')
            ax3 = fig.add_subplot(3,1,3)

            self.plotdist(dist0, axes=ax1)
            a5plt.momentumpolargrid(pnorm_edges, pitch_edges, axes=ax1)
            a5plt.momentumpolarplot(pnorm_edges, pitch_edges,
                                    dist1.distribution(), axes=ax2)
            self.plotdist(dist1, axes=ax3)
            plt.show()

        return distout

    def getdist_moments(self, dist, *moments, volmethod="prism"):
        """Calculate moments of distribution.

        Parameters
        ----------
        dist : str or :class:`DistData`
            Distribution from which moments are calculated.
        *moments : str
            Moments to be calculated.
        volmethod : {"mc", "prism"}, optional
            Method used to calculate the volume.

            It is a good idea to verify that same results are obtained with both
            methods.

        Returns
        -------
        out : :class:`DistMoment`
            Distribution object containing moments as ordinates.
        """
        # Initialize the moment object and evaluate the volume.
        if all([x in dist.abscissae for x in ["rho", "theta", "phi"]]):
            rhodist = True
            nrho   = dist.abscissa_edges("rho").size
            ntheta = dist.abscissa_edges("theta").size
            nphi   = dist.abscissa_edges("phi").size
            volume, area, r, phi, z = self._root._ascot.input_rhovolume(
                method=volmethod, tol=1e-2, nrho=nrho, ntheta=ntheta, nphi=nphi,
                return_area=True, return_coords=True)
            volume[volume == 0] = 1e-8 # To avoid division by zero
            out = DistMoment(
                dist.abscissa_edges("rho"), dist.abscissa_edges("theta"),
                dist.abscissa_edges("phi"), r, phi, z, area, volume, rhodist)
        elif all([x in dist.abscissae for x in ["r", "phi", "z"]]):
            rhodist = False
            r = dist.abscissa_edges("r")
            p = dist.abscissa_edges("phi").to("rad").v
            z = dist.abscissa_edges("z")
            v1, v2 = np.meshgrid( r[1:] - r[:-1], z[1:] - z[:-1])
            area = v1 * v2
            v1, v2, v3 = np.meshgrid(
                p[1:] - p[:-1], r[1:]**2 - r[:-1]**2, z[1:] - z[:-1])
            volume = 0.5 * v1 * v2 * v3
            r, phi, z = np.meshgrid(
                dist.abscissa("r"), dist.abscissa("phi"), dist.abscissa("z"))
            out = DistMoment(
                dist.abscissa_edges("r"), dist.abscissa_edges("phi"),
                dist.abscissa_edges("z"), r, phi, z, area, volume, rhodist)
        else:
            raise ValueError(
                "Distribution has neither (rho, theta, phi) nor (R, phi, z)")

        ordinates = {}
        mass = np.mean(self.getstate("mass"))
        if "density" in moments:
            Dist.density(dist, out)
        if "chargedensity" in moments:
            Dist.chargedensity(dist, out)
        if "energydensity" in moments:
            Dist.energydensity(mass, dist, out)
        if "toroidalcurrent" in moments:
            Dist.toroidalcurrent(self._root._ascot, mass, dist, out)
        if "parallelcurrent" in moments:
            Dist.parallelcurrent(self._root._ascot, mass, dist, out)
        if "pressure" in moments:
            Dist.pressure(mass, dist, out)
        if "powerdep" in moments:
            Dist.powerdep(self._root._ascot, mass, dist, out)
        if "electronpowerdep" in moments:
            Dist.electronpowerdep(self._root._ascot, mass, dist, out)
        if "ionpowerdep" in moments:
            Dist.ionpowerdep(self._root._ascot, mass, dist, out)
        if "jxbtorque" in moments:
            Dist.jxbtorque(self._root._ascot, mass, dist, out)
        if "colltorque" in moments:
            Dist.collTorque(self._root._ascot, mass, dist, out)
        if "canmomtorque" in moments:
            Dist.canMomentTorque(dist, out)
        return out

    def getdist_list(self, show=True):
        """List all available distributions and moments.

        Parameters
        ----------
        show : bool, optional
            Print the output.

        Returns
        -------
        dists : list [str]
            List of available distributions.
        moms : list [(str, str)]
            List of distribution moments that can be evaluated and their
            meaning.
        """
        dists = []
        if hasattr(self, "_dist5d"):    dists.append("5d")
        if hasattr(self, "_distrho5d"): dists.append("rho5d")
        if hasattr(self, "_dist6d"):    dists.append("6d")
        if hasattr(self, "_distrho6d"): dists.append("rho6d")
        if hasattr(self, "_distcom"):   dists.append("com")

        moms = []
        if "5d" not in dists and "rho5d" not in dists:
            if show:
                print("Available distributions:")
                for d in dists:
                    print(d)
                print("\nMoments available only for 5d/rho5d dists.")
            return dists, moms

        moms = [
            ("density", "Particle number density"),
            ("chargedensity", "Charge density"),
            ("energydensity", "Energy density"),
            ("toroidalcurrent", "Toroidal current"),
            ("parallelcurrent", "Parallel current"),
            ("pressure", "Pressure"),
            ("powerdep", "Total deposited power"),
            ("ionpowerdep", "Power deposited to ions"),
            ("electronpowerdep", "Power deposited to electrons"),
            ("jxbtorque", "j_rad x B_pol torque"),
            ("colltorque", "Torque from collisions"),
            ("canmomtorque", "Torque from change in can. tor. ang. momentum"),
        ]
        if show:
            print("Available distributions:")
            for d in dists:
                print(d)
            print("Available Moments:")
            for d in moms:
                print(d[0] + " : " + d[1])
        return dists, moms

    def plotdist(self, dist, axes=None, cax=None):
        """Plot distribution in 1D or 2D.

        This method assumes that the input distribution has been integrated,
        sliced, and interpolated so that only one or two dimensions have
        a size above one.

        Parameters
        ----------
        dist : :class:`DistData`
            The distribution data object.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.
        """
        x = None; y = None;
        for key in dist.abscissae:
            val = dist.abscissa_edges(key)
            if val.size > 2:
                if x is None:
                    x = val
                    xlabel = key + " [" + str(x.units) + "]"
                elif y is None:
                    y = val
                    ylabel = key + " [" + str(y.units) + "]"
                else:
                    raise ValueError(
                        "The distribution has more than two dimensions with "
                        + "size greater than one")
        if x is None: raise ValueError("The distribution is zero dimensional")

        ordinate = np.squeeze(dist.distribution())
        if y is None:
            ylabel = "f" + " [" + str(ordinate.units) + "]"
            a5plt.mesh1d(x, ordinate, xlabel=xlabel, ylabel=ylabel, axes=axes)
        else:
            # Swap pitch and energy
            if "ekin" in xlabel and "pitch" in ylabel:
                temp = x; x = y; y = temp
                temp = xlabel; xlabel = ylabel; ylabel = temp
                ordinate = ordinate.T

            axesequal = x.units == y.units
            clabel = "f" + " [" + str(ordinate.units) + "]"
            a5plt.mesh2d(x, y, ordinate, axesequal=axesequal, xlabel=xlabel,
                         ylabel=ylabel, clabel=clabel, axes=axes, cax=cax)

    def plotdist_moments(self, moment, ordinate, axes=None, cax=None):
        """Plot radial or (R,z) profile of a distribution moment.

        The plotted profile is the average of (theta, phi) or phi depending
        on whether the input is calculated from a rho distribution or not.

        Parameters
        ----------
        moment: class:`DistMoment`
            Moments calculated from the distribution.
        ordinate : str
            Name of the moment to be plotted.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.
        """
        if moment.rhodist:
            ylabel = ordinate
            ordinate = moment.ordinate(ordinate, toravg=True, polavg=True)
            ylabel += " [" + str(ordinate.units) + "]"
            a5plt.mesh1d(moment.rho, ordinate,
                         xlabel="Normalized poloidal flux",
                         ylabel=ylabel, axes=axes)
        else:
            clabel = ordinate
            ordinate = moment.ordinate(ordinate, toravg=True)
            clabel += " [" + str(ordinate.units) + "]"
            a5plt.mesh2d(moment.r, moment.z, ordinate, axesequal=True,
                         xlabel="R [m]", ylabel="z [m]", clabel=clabel,
                         axes=axes, cax=cax)

    def plotstate_scatter(self, x, y, z=None, c=None, xmode="gc", ymode="gc",
                          zmode="gc", cmode="gc", endcond=None, ids=None,
                          cint=9, cmap=None, axesequal=False, axes=None,
                          cax=None):
        """Make a scatter plot of marker state coordinates.

        The plot is either 2D+1D or 3D+1D, where the extra coordinate is color,
        depending on the number of queried coordinates. The color scale is
        discrete, not continuous.

        The quantity ("qnt") must have one of the following formats:

        - "ini qnt" plots inistate of quantity qnt.
        - "end qnt" plots endstate of quantity qnt.
        - "diff qnt" plots the difference between ini and endstate.
        - "reldiff qnt" plots the relative difference (x1/x0 - 1).
        - "log ini/end/diff/reldiff qnt" plots the logarithmic value.

        Parameters
        ----------
        x : str
            Name of the quantity on x-axis.
        y : str
            Name of the quantity on y-axis.
        z : str, optional
            Name of the quantity on z-axis.

            If not None, the plot will be in 3D.
        c : str, optional
            Name of the quantity shown with color scale or name of a color
            to plot all markers with same color.
        xmode : {"prt", "gc"}, optional
            Evaluate x quantity in particle or guiding-center phase space.
        ymode : {"prt", "gc"}, optional
            Evaluate y quantity in particle or guiding-center phase space.
        zmode : {"prt", "gc"}, optional
            Evaluate z quantity in particle or guiding-center phase space.
        cmode : {"prt", "gc"}, optional
            Evaluate color quantity in particle or guiding-center phase space.
        endcond : str, array_like, optional
            Endcond of those markers which are plotted.
        ids : array_like
            IDs of the markers to be plotted.
        cint : int or array_like optional
            Number of colors to be used or bin edges.

            All markers that have ``cint[i] < c <= cint[i+1]`` are plotted with
            same color. If ``cint`` array is not given explicitly,
            cint = np.linspace(min(c), max(c), cint).
        cmap : str, optional
            Name of the colormap where colors are picked.
        axesequal : bool, optional
            Flag whether to set aspect ratio for x and y (and z) axes equal.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.

        Raises
        ------
        AscotNoDataException
            Raised when data required for the opreation is not present.
        ValueError
            If argument could not be parsed.
        """
        def parsearg(arg):
            """Parse arguments of from "log ini qnt" to (qnt, value, islog).
            """
            arg = arg.lower()
            log = "linear"
            if "log" in arg:
                arg = arg.replace("log", "")
                log = "log"

            if "ini" in arg:
                arg = arg.replace("ini", "").strip()
                val = self.getstate(arg, state="ini", endcond=endcond, ids=ids)
                arg = "Initial " + arg
            elif "end" in arg:
                arg = arg.replace("end", "").strip()
                val = self.getstate(arg, state="end", endcond=endcond, ids=ids)
                arg = "Final " + arg
            elif "reldiff" in arg:
                arg = arg.replace("reldiff", "").strip()
                val1 = self.getstate(arg, state="ini", endcond=endcond, ids=ids)
                val2 = self.getstate(arg, state="end", endcond=endcond, ids=ids)
                val = (val2 -val1) / val1
                arg = r"$\Delta x/x_0$ " + arg
            elif "diff" in arg:
                arg = arg.replace("diff", "").strip()
                val1 = self.getstate(arg, state="ini", endcond=endcond, ids=ids)
                val2 = self.getstate(arg, state="end", endcond=endcond, ids=ids)
                val = val2 - val1
                arg = r"$\Delta$ " + arg
            else:
                raise ValueError(
                    "Unclear if a quantity is evaluated from ini or endstate.\n"
                    + "Use \"ini/end/diff/reldiff %s\"" % (arg)
                )
            if any(val < 0) and log == "log":
                log = "symlog"

            # Make sure val is an unyt array
            try:
                val.units
            except AttributeError:
                val *= unyt.dimensionless
            return (arg, val, log)

        x, xc, xlog = parsearg(x)
        y, yc, ylog = parsearg(y)
        x = x + " [" + str(xc.units) + "]"
        y = y + " [" + str(yc.units) + "]"

        cc = None; clog = False;
        if c is not None:
            if len(c.split()) == 1:
                # c is a color string, not quantity
                cc = c
                c = None
            else:
                c, cc, clog = parsearg(c)
                c = c + " [" + str(cc.units) + "]"

        if z is None:
            a5plt.scatter2d(xc, yc, c=cc, xlog=xlog, ylog=ylog, clog=clog,
                            xlabel=x, ylabel=y, clabel=c, cint=cint, cmap=cmap,
                            axesequal=axesequal, axes=axes, cax=cax)
        else:
            z, zc, zlog = parsearg(z)
            z = z + " [" + str(zc.units) + "]"
            a5plt.scatter3d(xc, yc, zc, c=cc, xlog=xlog, ylog=ylog,
                            zlog=zlog, clog=clog, xlabel=x, ylabel=y, zlabel=z,
                            clabel=c, cint=cint, cmap=cmap,
                            axesequal=axesequal, axes=axes, cax=cax)

    def plotstate_histogram(self, x, y=None, xbins=10, ybins=10, xmode="gc",
                            ymode="gc", endcond=None, ids=None, weight=False,
                            logscale=False, cmap=None, axesequal=False,
                            axes=None, cax=None):
        """Make a histogram plot of marker state coordinates.

        The histogram is either 1D or 2D depending on if the y coordinate is
        provided. In the 1D histogram the markers with different endstate are
        separated by color if the endcond argument is None.

        The quantity ("qnt") must have one of the following formats:

        - "ini qnt" plots inistate of quantity qnt.
        - "end qnt" plots endstate of quantity qnt.
        - "diff qnt" plots the difference between ini and endstate.
        - "reldiff qnt" plots the relative difference (x1/x0 - 1).
        - "log ini/end/diff/reldiff qnt" plots the logarithmic value.

        Parameters
        ----------
        x : str
            Name of the quantity on x-axis.
        y : str, optional
            Name of the quantity on y-axis.

            If not None, the histogram will be in 2D.
        xbins : int or array_like, optional
            Bin edges for the x coordinate or the number of bins.
        ybins : int or array_like, optional
            Bin edges for the y coordinate or the number of bins.
        xmode : {"prt", "gc"}, optional
            Evaluate x quantity in particle or guiding-center phase space.
        ymode : {"prt", "gc"}, optional
            Evaluate y quantity in particle or guiding-center phase space.
        endcond : str, array_like, optional
            Endcond of those markers which are plotted.

            If None and the plot is in 1D, the histogram will be stacked with
            different colors indicating different end conditions.
        ids : array_like
            IDs of the markers to be plotted.
        weight : bool
            Whether to weight markers when they are binned to the histogram.

            The weighted histogram represents physical particles whereas
            the unweighted histogram corresponds to markers.
        cmap : str, optional
            Name of the colormap used in the 2D histogram.
        axesequal : bool, optional
            Flag whether to set aspect ratio for x and y axes equal in 2D.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.
        """
        def parsearg(arg, mode, endcond):
            """Parse arguments of from "log ini qnt" to (qnt, value, islog).
            """
            arg = arg.lower()
            log = "linear"
            if "log" in arg:
                arg = arg.replace("log", "")
                log = "log"

            if "ini" in arg:
                arg = arg.replace("ini", "").strip()
                val = self.getstate(arg, state="ini", endcond=endcond, ids=ids,
                                    mode=mode)
                if log == "log": arg = "|" + arg + "|"
                arg = "Initial " + arg
            elif "end" in arg:
                arg = arg.replace("end", "").strip()
                val = self.getstate(arg, state="end", endcond=endcond, ids=ids,
                                    mode=mode)
                if log == "log": arg = "|" + arg + "|"
                arg = "Final " + arg
            elif "reldiff" in arg:
                arg = arg.replace("reldiff", "").strip()
                val1 = self.getstate(arg, state="ini", endcond=endcond, ids=ids,
                                     mode=mode)
                val2 = self.getstate(arg, state="end", endcond=endcond, ids=ids,
                                     mode=mode)
                val = (val2 -val1) / val1
                arg = r"$\Delta x/x_0$ " + arg
                if log == "log": arg = "|" + arg + "|"
            elif "diff" in arg:
                arg = arg.replace("diff", "").strip()
                val1 = self.getstate(arg, state="ini", endcond=endcond, ids=ids,
                                     mode=mode)
                val2 = self.getstate(arg, state="end", endcond=endcond, ids=ids,
                                     mode=mode)
                val = val2 - val1
                arg = r"$\Delta$ " + arg
                if log == "log": arg = "|" + arg + "|"
            else:
                raise ValueError(
                    "Unclear if a quantity is evaluated from ini or endstate.\n"
                    + "Use \"ini/end/diff/reldiff %s\"" % (arg)
                )
            if log == "log": val = np.abs(val)
            # Make sure val is an unyt array
            try:
                val.units
            except AttributeError:
                val *= unyt.dimensionless
            return (arg, val, log)

        if y is None:
            # 1D plot
            xcs = []
            weights = []

            #ecs, _ = self.getstate_markersummary()
            #for ec in ecs:
            #    if endcond is not None and ec[1] not in [endcond]:
            #        continue

            #    ids, w = self.getstate("ids", "weight", endcond=ec[1], ids=ids)
            #    xcs.append(xc[ids-1])
            #    weights.append(w)

            #xc = [xc]

            x0       = x
            xcs      = []
            weights  = []
            endconds = []
            ecs, _ = self.getstate_markersummary()
            for ec in ecs:
                if endcond is None or ec[1] in endcond:
                    w = self.getstate("weight", endcond=ec[1], ids=ids)
                    x, xc, xlog = parsearg(x0, xmode, ec[1])
                    xcs.append(xc)
                    weights.append(w)
                    endconds.append(ec)

            if len(xcs) == 0: return

            # Sort data so that when the stacked histogram is plotted, the stack
            # with most markers is at the bottom.
            idx = np.argsort([len(i) for i in xcs])[::-1]
            xcs = [xcs[i].v for i in idx]
            ecs = [endconds[i][1] + " : %.2e" % endconds[i][0] for i in idx]
            weights = [weights[i] for i in idx]
            if not weight: weights = None

            a5plt.hist1d(x=xcs, xbins=xbins, weights=weights, xlog=xlog,
                         logscale=logscale, xlabel=x, axes=axes, legend=ecs)

        else:
            # 2D plot
            x, xc, xlog = parsearg(x, xmode, endcond)
            y, yc, ylog = parsearg(y, ymode, endcond)
            weights = self.getstate("weight", state="ini", endcond=endcond,
                                    ids=ids)
            if not weight: weights = None

            a5plt.hist2d(xc, yc, xbins=xbins, ybins=ybins, weights=weights,
                         xlog=xlog, ylog=ylog, logscale=logscale, xlabel=x,
                         ylabel=y, axesequal=axesequal, axes=axes, cax=cax)

    def plotorbit_trajectory(self, x, y, z=None, c=None, endcond=None, ids=None,
                             cmap=None, axesequal=False, axes=None, cax=None):
        """Plot orbit trajectories in arbitrary coordinates.

        The plot is either 2D+1D or 3D+1D, where the extra coordinate is color,
        depending on the number of queried coordinates. The color scale is
        discrete, not continuous.

        The quantity ("qnt") must have one of the following formats:

        - "diff qnt" plots the difference between inistate and current value.
        - "reldiff qnt" plots the relative difference (x1/x0 - 1).
        - "log <diff/reldiff> qnt" plots the logarithmic value.

        Parameters
        ----------
        x : str
            Name of the quantity on x-axis.
        y : str
            Name of the quantity on y-axis.
        z : str, optional
            Name of the quantity on z-axis.

            If not None, the plot will be in 3D.
        c : str, optional
            The color used to plot the markers or name of the quantity shown
            with color.

            If None, markers are plotted with different colors.
        endcond : str, array_like, optional
            Endcond of those  markers which are plotted.
        ids : array_like
            IDs of the markers to be plotted.
        cmap : str, optional
            Colormap.
        axesequal : bool, optional
            Flag whether to set aspect ratio for x and y (and z) axes equal.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.
        """
        def parsearg(arg):
            """Parse arguments of from string to (qnt, label, islog).
            """
            arg = arg.lower()
            log = "linear"
            label = arg
            if "log" in arg:
                arg = arg.replace("log", "")
                log = "log"
            if "reldiff" in arg:
                arg = arg.replace("reldiff", "")
                label = r"$\Delta x/x_0$ " + arg
            elif "diff" in arg:
                arg = arg.replace("diff", "")
                label = r"$\Delta$ " + arg

            arg = arg.strip()
            return (arg, label, log)

        cc = c; clabel = None; clog = "linear"# Default values passed to plotter
        x, xlabel, xlog = parsearg(x)
        y, ylabel, ylog = parsearg(y)
        if c is not None: c, clabel, clog = parsearg(c)
        if z is not None: z, zlabel, zlog = parsearg(z)

        if c is not None and not c in self.getorbit_list():
            # c is a color string, not quantity
            c = None

        # Orbit evaluation can be expensive so we try to get all coordinates
        # with a single call
        if z is None and c is None:
            idarr, xc, yc = self.getorbit(
                "ids", x, y, endcond=endcond, ids=ids)
        elif z is not None and c is None:
            idarr, xc, yc, zc = self.getorbit(
                "ids", x, y, z, endcond=endcond, ids=ids)
        elif z is None and c is not None:
            idarr, xc, yc, cc = self.getorbit(
                "ids", x, y, c, endcond=endcond, ids=ids)
        elif z is not None and c is not None:
            idarr, xc, yc, zc, cc = self.getorbit(
                "ids", x, y, z, c, endcond=endcond, ids=ids)

        # Find indices to map values from inistate to orbit array
        _, idarr = np.unique(idarr, return_inverse=True)
        idx = np.where(idarr[:-1] != idarr[1:])[0] + 1
        def parsevals(val, log, label, qnt):
            """Compute values and split them by orbit
            """
            if "x/x_0" in label:
                val0 = self.getstate(qnt, state="ini", endcond=endcond, ids=ids)
                val = ( val / val0[idarr] ) - 1
            elif "Delta" in label:
                val0 = self.getstate(qnt, state="ini", endcond=endcond, ids=ids)
                val = val - val0[idarr]

            if any(val < 0) and log == "log":
                log = "symlog"

            # Make sure val is an unyt array
            try:
                val.units
            except AttributeError:
                val *= unyt.dimensionless
            vmin = np.amin(val)
            vmax = np.amax(val)
            val = np.split(val, idx)
            return (val, log, vmin, vmax)

        xc, xlog, xmin, xmax = parsevals(xc, xlog, xlabel, x)
        yc, ylog, ymin, ymax = parsevals(yc, ylog, ylabel, y)
        xlabel = xlabel + " [" + str(xc[0].units) + "]"
        ylabel = ylabel + " [" + str(yc[0].units) + "]"
        bbox = [xmin, xmax, ymin, ymax, None, None]
        if z is not None:
            zc, zlog, zmin, zmax = parsevals(zc, zlog, zlabel, z)
            zlabel + " [" + str(zc[0].units) + "]"
            bbox = [xmin, xmax, ymin, ymax, zmin, zmax, None, None]
        if c is not None:
            cc, clog, bbox[-2], bbox[-1] = parsevals(cc, clog, clabel, c)
            clabel + " [" + str(cc[0].units) + "]"

        if z is None:
            a5plt.line2d(xc, yc, c=cc, xlog=xlog, ylog=ylog, clog=clog,
                         xlabel=xlabel, ylabel=ylabel, clabel=clabel, bbox=bbox,
                         cmap=cmap, axesequal=axesequal, axes=axes, cax=cax)
        else:
            a5plt.line3d(xc, yc, zc, c=cc, xlog=xlog, ylog=ylog, zlog=zlog,
                         clog=clog, xlabel=xlabel, ylabel=ylabel, zlabel=zlabel,
                         clabel=clabel, bbox=bbox, cmap=cmap,
                         axesequal=axesequal, axes=axes, cax=cax)

    def plotorbit_poincare(self, plane, connlen=True, axes=None, cax=None):
        """Create a Poincaré plot where the color separates different markers
        or shows the connection length.

        Parameters
        ----------
        plane : str
            The Poincaré plane to be plotted.

            The argument is expected to have format "pol/tor/rad i" where the
            first part specifies the type of the plane (poloidal, toroidal,
            radial) and the second argument is the plane index. For example,
            ``plane="pol 2"`` plots the second poloidal plane. The order of
            the planes is same as given by :meth:`getorbit_poincareplanes`.
        connlen : bool, optional
            Show connection length and separated lost and confined markers with
            color.

            If true, trajectories of lost markers are colored (in blue
            shades) where the color shows the connection length at that
            position. Confined (or all if conlen=False) markers are shown
            with shades of red where color separates subsequent
            trajectories.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.

        Raises
        ------
        ValueError
            Raised when the plane is unknown.
        """
        # Set quantities corresponding to the planetype
        pol, tor, rad = self.getorbit_poincareplanes()
        if "pol" in plane:
            plane = int(plane.split()[1]) - 1
            if plane < len(pol) and plane >= 0:
                plane = pol[plane][1]
                x = "r"
                y = "z"
                xlabel = "R [m]"
                ylabel = "z [m]"
                xlim = None # Determined spearately
                ylim = None # Determined separately
                axesequal = True
            else:
                raise ValueError("Unknown plane.")
        elif "tor" in plane:
            plane = int(plane.split()[1]) - 1
            if plane < len(tor) and plane >= 0:
                plane = tor[plane][1]
                x = "rho"
                y = "phimod"
                xlabel = "Normalized poloidal flux"
                ylabel = "Toroidal angle [deg]"
                xlim = [0, 1.1]
                ylim = [0, 360]
                axesequal = False
            else:
                raise ValueError("Unknown plane.")
        elif "rad" in plane:
            plane = int(plane.split()[1]) - 1
            if plane < len(rad) and plane >= 0:
                plane = rad[plane][1]
                x = "thetamod"
                y = "phimod"
                xlabel = "Poloidal angle [deg]"
                ylabel = "Toroidal angle [deg]"
                xlim = [0, 360]
                ylim = [0, 360]
                axesequal = True
            else:
                raise ValueError("Unknown plane.")
        else:
            raise ValueError("Unknown plane.")

        plotconnlen = connlen
        ids, x, y, connlen = self.getorbit("ids", x, y, "connlen", pncrid=plane)

        if xlim == None:
            # Determine (R,z) limits for poloidal plane by making sure the data
            # fits nicely and then rounding to nearest decimal.
            xlim = [np.floor(np.amin(x)*9) / 10, np.ceil(np.amax(x)*11) / 10]
            ylim = [np.floor(np.amin(y)*11) / 10, np.ceil(np.amax(y)*11) / 10]

        if plotconnlen:
            # Now set confined markers as having negative connection length
            connlen *= -1
            lost1 = self.getstate("ids", state="end", endcond="rhomax wall")

            idx = ~np.in1d(ids, lost1)
            connlen[idx] *= -1
            clabel = "Connection length [" + str(connlen.units) + "]"
        else:
            connlen = None
            clabel = None

        a5plt.poincare(x, y, ids, connlen=connlen, xlim=xlim, ylim=ylim,
                       xlabel=xlabel, ylabel=ylabel, clabel=clabel,
                       axesequal=axesequal, axes=axes, cax=cax)

    def plotwall_loadvsarea(self, axes=None):
        """Plot histogram showing area affected by at least a given load.

        Parameters
        ----------
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        """
        _, area, eload, _, _ = self.getwall_loads()
        a5plt.loadvsarea(area, eload/area, axes=axes)

    @a5plt.openfigureifnoaxes(projection=None)
    def plotwall_torpol(self, log=True, clim=None, axes=None, cax=None):
        """Plot the toroidal-poloidal projection of the wall loads.

        Note that the resulting doesn't present the areas accurately and even
        the locations are approximate. The purpose of the plot is to roughly
        illustrate how the losses are distributed.

        Parameters
        ----------
        log : bool, optional
            Make the color scale logarithmic.
        clim : [float, float], optional
            Minimum and maximum values in the color scale.
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.
        """
        wetted, area, edepo, pdepo, iangle = self.getwall_loads()
        d = self.wall.read()
        nelement = wetted.size
        color    = edepo/area
        x1x2x3   = d["x1x2x3"][wetted]
        y1y2y3   = d["y1y2y3"][wetted]
        z1z2z3   = d["z1z2z3"][wetted]

        # Toroidal angle for each vertex
        tor = np.rad2deg( np.arctan2( y1y2y3, x1x2x3 ))

        # Poloidal angle for each vertex
        out = self._root._ascot._eval_bfield(
            1*unyt.m, tor*unyt.deg, 1*unyt.m, 0*unyt.s, evalaxis=True)
        axisr  = out["axisr"].v
        axisz  = out["axisz"].v
        rmajor = np.sqrt( x1x2x3**2 + y1y2y3**2 )
        pol    = np.rad2deg( np.arctan2(z1z2z3 - axisz, rmajor - axisr) )

        # If triangle would otherwise be split on different ends of [0, 360]
        # range in tor or pol, move it to top/right.
        dx = (np.amax(tor, axis=1) - np.amin(tor, axis=1)) > 180.0
        tor[dx, :] = np.mod(tor[dx, :], 360)
        dy = (np.amax(pol, axis=1) - np.amin(pol, axis=1)) > 180.0
        pol[dy, :] = np.mod(pol[dy, :], 360)

        # Move the horizontal axis to 0...360
        tor[np.amin(tor, axis=1) < 0, :] += 360.0

        # Sort the patches so that triangles closest to the axis are drawn
        # last (so that they are on top of others)
        rminor = np.sqrt( (rmajor - axisr)**2 + (z1z2z3 - axisz)**2 )
        idx    = np.flipud( np.argsort(np.amin(rminor, axis=1), axis=0) )
        tor    = tor[idx, :]
        pol    = pol[idx, :]
        color  = color[idx]

        # Convert the data into format [vert0, vert1, ...] where verti contains
        # vertices of triangle i in (3,2) array.
        patches = np.zeros((nelement,3,2))
        patches[:,:,0] = tor[:,:]
        patches[:,:,1] = pol[:,:]
        a5plt.triangularpatch(
            patches, color, xlim=[-2, 362], ylim=[-182, 182], clim=clim,
            log=log, axes=axes, cax=cax, clabel=r"Wall load [W/m$^2$]",
            cmap="Reds",
            xlabel="Toroidal angle [deg]", ylabel="Poloidal angle [deg]")
        axes.set_xticks([0, 90, 180, 270, 360])
        axes.set_yticks([-180, -90, 0, 90, 180])

    def plotwall_3dstill(self, wallmesh=None, points=None, data=None, log=False,
                         cpos=None, cfoc=None, cang=None, axes=None, cax=None):
        """Take a still shot of the mesh and display it using matplotlib
        backend.

        The rendering is done using vtk but the vtk (interactive) window is not
        displayed. It is recommended to use the interactive plot to find desired
        camera position and produce the actual plot using this method. The plot
        is shown using imshow and acts as a regular matplotlib plot.

        Parameters
        ----------
        wallmesh : Polydata
            Mesh representing the wall. If not given then construct the mesh
            from the wall data.
        points : {array_like, bool}, optional
            Array Npoint x 3 defining points (markers) to be shown. For each
            point [x, y, z] coordinates are given. If boolean True is given,
            then markers are read from the endstate.
        cpos : array_like, optional
            Camera position coordinates [x, y, z].
        cfoc : array_like, optional
            Camera focal point coordinates [x, y, z].
        cang : array_like, optional
            Camera angle [azimuth, elevation, roll].
        axes : :obj:`~matplotlib.axes.Axes`, optional
            The axes where figure is plotted or otherwise new figure is created.
        cax : :obj:`~matplotlib.axes.Axes`, optional
            The color bar axes or otherwise taken from the main axes.
        """
        if wallmesh is None:
            wallmesh = self.getwall_3dmesh()
        if isinstance(points, bool) and points == True:
            points = self.getstate_pointcloud(endcond="wall")

        (cpos0, cfoc0, cang0) = a5plt.defaultcamera(wallmesh)
        if cpos is None: cpos = cpos0
        if cfoc is None: cfoc = cfoc0
        if cang is None: cang = cang0

        a5plt.still(wallmesh, points=points, data=data, log=log, cpos=cpos,
                    cfoc=cfoc, cang=cang, axes=axes, cax=cax)


    def plotwall_3dinteractive(self, wallmesh=None, *args, points=None,
                               data=None, log=False, cpos=None, cfoc=None,
                               cang=None):
        """Open vtk window to display interactive view of the wall mesh.

        Parameters
        ----------
        wallmesh, optional : Polydata
            Mesh representing the wall. If not given then construct the
            mesh from the wall data.
        *args : tuple (str, method), optional
            Key (str) method pairs. When key is pressed when the plot is
            displayed, the associated method is called. The method should
            take Plotter instance as an argument.
        points : array_like, optional
            Array Npoint x 3 defining points (markers) to be shown. For
            each point [x, y, z] coordinates are given. If boolean True is
            given, then markers are read from the endstate.
        cpos : array_like, optional
            Camera position coordinates [x, y, z].
        cfoc : array_like, optional
            Camera focal point coordinates [x, y, z].
        cang : array_like, optional
            Camera angle [azimuth, elevation, roll].
        """
        if wallmesh is None:
            wallmesh = self.getwall_3dmesh()
        if isinstance(points, bool) and points == True:
            points = self.getstate_pointcloud(endcond="wall")

        (cpos0, cfoc0, cang0) = a5plt.defaultcamera(wallmesh)
        if cpos is None: cpos = cpos0
        if cfoc is None: cfoc = cfoc0
        if cang is None: cang = cang0

        a5plt.interactive(wallmesh, *args, points=points, data=data, log=log,
                          cpos=cpos, cfoc=cfoc, cang=cang)

a5py/a5py/routines/runmixin.py
ascot4fusion/ascot5
"""Defines virtual run which reads data from the C arrays instead of HDF5.
"""
import numpy as np
import unyt

from a5py.ascot5io import State, Orbits
from .runmixin import RunMixin

from a5py.ascotpy.libascot import _LIBASCOT
if _LIBASCOT:
    import a5py.ascotpy.ascot2py as ascot2py

class VirtualRun(RunMixin):
    """Virtual :class:`RunGroup` whose data exists solely in the memory.
    """

    def __init__(self, ascot, nmrk, inistate, endstate, npoint, mode, orbmode,
                 options, markers, orbits=None):
        """Initialize fields that allow this instance to replicate
        :class:`RunGroup` behavior.

        Sets state and orbit attributes which are the used by the methods in
        :class:`RunMixin`.

        Parameters
        ----------
        ascot : :class:`Ascot`
            Ascot instance for input interpolation.
        nmrk : int
            Number of markers in the simulation.
        npoint : int
            Number of data points reserved per marker.
        inistate :
            Pointer to marker inistate array.
        endstate :
            Pointer to marker endstate array.
        mode : int
            What simulation mode was active.
        orbmode : int
            What orbit collection mode was active.
        orbits : array_like, optional
            The diagnostics array containing the orbit data if present.
        """
        self.options = options
        self.markers = markers
        for inp in ["bfield", "efield", "plasma", "neutral", "wall", "boozer",
                    "mhd", "asigma"]:
            grp = ascot.data[inp].active
            setattr(self, inp, grp)
        self._inistate = VirtualState(ascot, nmrk, inistate)
        self._endstate = VirtualState(ascot, nmrk, endstate)
        if orbits is not None:
            self._orbit = VirtualOrbits(
                ascot, nmrk, npoint, mode, orbmode, orbits)

class VirtualState():
    """Like :class:`State` but the data is in C array.
    """

    def __init__(self, ascot, nmrk, state):
        """Initialize fields that allow this instance to replicate
        :class:`RunGroup` behavior.

        Parameters
        ----------
        ascot : :class:`Ascot`
            Ascot instance for input interpolation.
        nmrk : int
            Number of markers in the simulation.
        state :
            Pointer to marker state array.
        """
        self.ascot = ascot
        self.state = state
        self.nmrk  = nmrk

    def get(self, *qnt, mode="gc"):

        def _val(q):
            """Read quantity from HDF5.
            """
            arr = lambda q : np.array([getattr(self.state[i], q)
                                       for i in range(self.nmrk)])

            m = unyt.m; rad = unyt.rad; kg = unyt.kg; s = unyt.s; J = unyt.J
            T = unyt.T; C = unyt.C; nodim = unyt.dimensionless
            if   q == "ids":       return arr("id") * nodim
            elif q == "rprt":      return arr(q) * m
            elif q == "phiprt":    return arr(q) * rad
            elif q == "zprt":      return arr(q) * m
            elif q == "prprt":     return arr("p_r") * kg*m/s
            elif q == "pphiprt":   return arr("p_phi") * kg*m/s
            elif q == "pzprt":     return arr("p_z") * kg*m/s
            elif q == "r":         return arr(q) * m
            elif q == "phi":       return arr(q) * rad
            elif q == "z":         return arr(q) * m
            elif q == "ppar":      return arr(q) * kg*m/s
            elif q == "mu":        return arr(q) * J/T
            elif q == "zeta":      return arr(q) * rad
            elif q == "weight":    return arr(q) * nodim
            elif q == "time":      return arr(q) * s
            elif q == "mileage":   return arr(q) * s
            elif q == "cputime":   return arr(q) * s
            elif q == "rho":       return arr(q) * nodim
            elif q == "theta":     return arr(q) * rad
            elif q == "mass":      return arr(q) * kg
            elif q == "br":        return arr("B_r") * T
            elif q == "bphi":      return arr("B_phi") * T
            elif q == "bz":        return arr("B_z") * T
            elif q == "ids":       return arr(q) * nodim
            elif q == "endcond":   return arr(q) * nodim
            elif q == "walltile":  return arr(q) * nodim
            elif q == "charge":    return arr(q) * C
            elif q == "anum":      return arr(q) * nodim
            elif q == "znum":      return arr(q) * nodim
            elif q == "errormsg":  return arr("err") * nodim
            elif q == "errorline": return arr("err") * nodim
            elif q == "errormod":  return arr("err") * nodim
            return arr

        def _eval(r, phi, z, t, *q):
            """Evaluate input quantities at marker position.
            """
            return self.ascot.input_eval(r, phi, z, t, *q)

        return State._getactual(mode, _val, _eval, *qnt)

class VirtualOrbits():
    """Like :class:`Orbits` but the data is in C array.
    """

    def __init__(self, ascot, nmrk, npoint, mode, orbmode, orbits):
        """Initialize fields that allow this instance to replicate
        :class:`Orbits` behavior.

        Parameters
        ----------
        ascot : :class:`Ascot`
            Ascot instance for input interpolation.
        nmrk : int
            Number of markers in the simulation.
        npoint : int
            Number of data points reserved per marker.
        mode : int
            What simulation mode was active.
        orbmode : int
            What orbit collection mode was active.
        orbits : array_like
            The diagnostics array containing the orbit data.
        """
        self.ascot   = ascot
        self.nmrk    = nmrk
        self.npnt    = npoint
        self.mode    = mode
        self.orbmode = orbmode
        self.arr     = orbits

    def get(self, inistate, endstate, *qnt):
        """Return marker quantity.

        This function accesses the orbit data within the C array and uses that
        to evaluate the queried quantity. The evaluated quantity at a given
        position corresponds to that mode which was active at the simulation:
        GO simulations return particle phase-space, GC guiding-center
        phase-space and hybrid depends on whether marker was GO or GC at that
        moment when data point was written.

        Parameters
        ----------
        inistate : State
            Inistate is needed to evaluate some orbit quantities.
        endstate : State
            Endstate is needed to evaluate some orbit quantities.
        *qnt : str
            Names of the quantities.

        Returns
        -------
        *value : array_like
            The quantities as an array ordered by marker ID (major) and mileage
            (minor).
        """
        # Prepare helper variables and functions
        arrlen = self.nmrk * self.npnt
        def _val(q, mask=None):
            """Read quantity from the orbit array.
            """
            # C array contains dummy data where ids == 0
            idx = np.array(self.arr[arrlen*0:arrlen*1]) > 0
            m = unyt.m; rad = unyt.rad; kg = unyt.kg; s = unyt.s; J = unyt.J
            T = unyt.T; C = unyt.C; nodim = unyt.dimensionless
            if q == "ids":
                val = np.array(self.arr[arrlen*0:arrlen*1]) * nodim
            elif q == "mileage":
                val = np.array(self.arr[arrlen*1:arrlen*2]) * s
            elif q == "r":
                val = np.array(self.arr[arrlen*2:arrlen*3]) * m
            elif q == "phi":
                val = np.array(self.arr[arrlen*3:arrlen*4]) * rad
            elif q == "z":
                val = np.array(self.arr[arrlen*4:arrlen*5]) * m

            if self.mode == ascot2py.simulate_mode_fo:
                if q == "pr":
                    val = np.array(self.arr[arrlen*5:arrlen*6]) * kg*m/s
                elif q == "pphi":
                    val = np.array(self.arr[arrlen*6:arrlen*7]) * kg*m/s
                elif q == "pz":
                    val = np.array(self.arr[arrlen*7:arrlen*8]) * kg*m/s
                elif q == "weight":
                    val = np.array(self.arr[arrlen*8:arrlen*9]) * nodim
                elif q == "charge":
                    val = np.array(self.arr[arrlen*9:arrlen*10]) * C
                elif q == "rho":
                    val = np.array(self.arr[arrlen*10:arrlen*11]) * nodim
                elif q == "theta":
                    val = np.array(self.arr[arrlen*11:arrlen*12]) * rad
                elif q == "br":
                    val = np.array(self.arr[arrlen*12:arrlen*13]) * T
                elif q == "bphi":
                    val = np.array(self.arr[arrlen*13:arrlen*14]) * T
                elif q == "bz":
                    val = np.array(self.arr[arrlen*14:arrlen*15]) * T
                elif q == "simmode":
                    val = np.array(self.arr[arrlen*15:arrlen*16]) * nodim
                if q == "pncrid" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*16:arrlen*17]) * nodim
                if q == "pncrdir" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*17:arrlen*18]) * nodim
            elif self.mode == ascot2py.simulate_mode_gc:
                if q == "ppar":
                    val = np.array(self.arr[arrlen*5:arrlen*6]) * kg*m/s
                elif q == "mu":
                    val = np.array(self.arr[arrlen*6:arrlen*7]) * J/T
                elif q == "zeta":
                    val = np.array(self.arr[arrlen*7:arrlen*8]) * rad
                elif q == "weight":
                    val = np.array(self.arr[arrlen*8:arrlen*9]) * nodim
                elif q == "charge":
                    val = np.array(self.arr[arrlen*9:arrlen*10]) * C
                elif q == "rho":
                    val = np.array(self.arr[arrlen*10:arrlen*11]) * nodim
                elif q == "theta":
                    val = np.array(self.arr[arrlen*11:arrlen*12]) * rad
                elif q == "br":
                    val = np.array(self.arr[arrlen*12:arrlen*13]) * T
                elif q == "bphi":
                    val = np.array(self.arr[arrlen*13:arrlen*14]) * T
                elif q == "bz":
                    val = np.array(self.arr[arrlen*14:arrlen*15]) * T
                elif q == "simmode":
                    val = np.array(self.arr[arrlen*15:arrlen*16]) * nodim
                if q == "pncrid" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*16:arrlen*17]) * nodim
                if q == "pncrdir" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*17:arrlen*18]) * nodim
            elif self.mode == ascot2py.simulate_mode_ml:
                if q == "rho":
                    val = np.array(self.arr[arrlen*5:arrlen*6]) * nodim
                elif q == "theta":
                    val = np.array(self.arr[arrlen*6:arrlen*7]) * rad
                elif q == "br":
                    val = np.array(self.arr[arrlen*7:arrlen*8]) * T
                elif q == "bphi":
                    val = np.array(self.arr[arrlen*8:arrlen*9]) * T
                elif q == "bz":
                    val = np.array(self.arr[arrlen*9:arrlen*10]) * T
                elif q == "simmode":
                    val = np.array(self.arr[arrlen*10:arrlen*11]) * nodim
                if q == "pncrid" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*11:arrlen*12]) * nodim
                if q == "pncrdir" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*12:arrlen*13]) * nodim
            elif self.mode == ascot2py.simulate_mode_hybrid:
                if q == "pr":
                    val = np.array(self.arr[arrlen*5:arrlen*6]) * kg*m/s
                elif q == "pphi":
                    val = np.array(self.arr[arrlen*6:arrlen*7]) * kg*m/s
                elif q == "pz":
                    val = np.array(self.arr[arrlen*7:arrlen*8]) * kg*m/s
                elif q == "ppar":
                    val = np.array(self.arr[arrlen*8:arrlen*9]) * kg*m/s
                elif q == "mu":
                    val = np.array(self.arr[arrlen*9:arrlen*10]) * J/T
                elif q == "zeta":
                    val = np.array(self.arr[arrlen*10:arrlen*11]) * rad
                elif q == "weight":
                    val = np.array(self.arr[arrlen*11:arrlen*12]) * nodim
                elif q == "charge":
                    val = np.array(self.arr[arrlen*12:arrlen*13]) * C
                elif q == "rho":
                    val = np.array(self.arr[arrlen*13:arrlen*14]) * nodim
                elif q == "theta":
                    val = np.array(self.arr[arrlen*14:arrlen*15]) * rad
                elif q == "br":
                    val = np.array(self.arr[arrlen*15:arrlen*16]) * T
                elif q == "bphi":
                    val = np.array(self.arr[arrlen*16:arrlen*17]) * T
                elif q == "bz":
                    val = np.array(self.arr[arrlen*17:arrlen*18]) * T
                elif q == "simmode":
                    val = np.array(self.arr[arrlen*18:arrlen*19]) * nodim
                if q == "pncrid" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*19:arrlen*20]) * nodim
                if q == "pncrdir" and self.orbmode == 0:
                    val = np.array(self.arr[arrlen*20:arrlen*21]) * nodim
            return val[idx][mask].ravel()

        # Sort using the fact that inistate.get return values ordered by ID
        # and also np.unique returns indices that produce a sorted array.
        mode    = _val("simmode")
        _, idx  = np.unique(_val("ids").v, return_inverse=True)
        mass    = inistate.get("mass")[0][idx]
        time    = inistate.get("time")[0][idx]
        connlen = inistate.get("mileage")[0][idx] - _val("mileage")

        # Only field lines are constant in time
        if not Orbits.FIELDLINE in mode: time = time + _val("mileage")

        def _eval(q, mask=None):
            """Evaluate input quantities at marker position.
            """
            return self.ascot.input_eval(
                _val("r", mask=mask), _val("phi",  mask=mask),
                _val("z", mask=mask), time[mask], *[q])

        return Orbits._getactual(mass, time, connlen, mode, _val, _eval, *qnt)

a5py/a5py/routines/virtualrun.py
ascot4fusion/ascot5
"""Package for generating inputs from templates and imported data.
"""
from .analyticalinputs import AnalyticalInputs
from .optiontemplates  import OptionTemplates
from .poincare  import PoincareTemplates
from .importdata import ImportAdas

class InputFactory(AnalyticalInputs, OptionTemplates, PoincareTemplates,
                   ImportAdas):
    """Class for creating input data from templates or imported data.

    The templates are constructed by calling :meth:`construct` and specifying
    the template. The templates are constructed in other (public) methods
    of this class. The name of the template is the name of the method, and
    details of the template can be found from the method description.
    """

    def __init__(self, ascot):
        """Store class:`Ascot` object needed for input creation.
        """
        self._ascot = ascot

    def construct(self, template, **kwargs):
        """Create input from a template or import data.

        This method is just a wrapper for calling other methods that generate
        the inputs. The name of the template should be the name of
        the corresponding method (whitespace instead of underscores and capital
        letters are allowed).

        Parameters
        ----------
        template : str
            Name of the template.
        **kwargs
            Parameters passed to the template.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.

        Raises
        ------
        ValueError
            When given template was not found.
        """
        template = template.lower().replace(" ", "_")
        try:
            template = getattr(self, template)
        except AttributeError:
            raise ValueError("Unknown template: " + template)
        return template(**kwargs)

a5py/a5py/templates/__init__.py
ascot4fusion/ascot5
"""Analytical inputs.
"""
import numpy as np
import unyt

import a5py.ascot5io.marker as marker
import a5py.ascot5io.options as options

from a5py.ascot5io.bfield import B_GS, B_2DS, B_3DS
from a5py.ascot5io.wall import wall_2D, wall_3D

class AnalyticalInputs():
    """Inputs that can be constructed analytically or otherwise from scratch.
    """

    def bfield_analytical_iter_circular(self, axisymmetric=True, splines=False):
        """Create ITER like but circular magnetic field.

        Parameters
        ----------
        axisymmetric : bool, optional
            If True, the field contains toroidal ripple.

            Note that the 3D field is not divergence free.
        splines : bool, optional
            If True, the analytical magnetic field format class:`B_GS` is
            converted to :class:`B_2DS` or :class:`B_3DS` depending on whether
            the field is axisymmetric or not.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.
        """

        out = {}
        out["coefficients"] = np.array([
             2.21808016e-02,  -1.28841781e-01,  -4.17718173e-02,
            -6.22680280e-02,   6.20083978e-03,  -1.20524711e-03,
            -3.70147050e-05,   0.00000000e+00,   0.00000000e+00,
             0.00000000e+00,   0.00000000e+00,   0.00000000e+00,
            -0.155])
        out.update({"r0" : 6.2, "z0" : 0, "bphi0" : 5.3, "psimult" : 200})

        if axisymmetric:
            if splines:
                out.update({"rmin" : 4, "rmax" : 8.5, "nr" : 120, "zmin" : -4,
                            "zmax" : 4, "nz" : 200})
                return ("B_2DS", B_2DS.convert_B_GS(**out))
            return ("B_GS", out)

        out.update({"nripple" : 18, "a0" : 2, "alpha0" : 0.2, "delta0" : 0.5})
        if splines:
            out.update({"rmin" : 4, "rmax" : 8.5, "nr" : 120, "zmin" : -4,
                        "zmax" : 4, "nz" : 200, "phimin" : 0, "phimax" : 360,
                        "nphi" : 360})
            return ("B_3DS", B_3DS.convert_B_GS(**out))

        return ("B_GS", out)

    def plasma_flat(self, density=10e20, temperature=10e3, anum=1, znum=1,
                    charge=1, mass=1):
        """Create uniform single-species plasma that is flat inside the
        separatrix but inexistent outside.

        Parameters
        ----------
        density : float, optional
            Plasma density.
        temperature : float, optional
            Plasma temperature.
        anum : int, optional
            Ion atomic number.
        znum : int, optional
            Ion charge number.
        charge : int, optional
            Ion charge.
        mass : int, optional
            Ion mass.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.
        """
        nrho   = 100
        rho    = np.transpose( np.linspace(0, 10, nrho) )
        prof   = np.ones((nrho, 1))
        edens  = density     * prof
        etemp  = temperature * prof
        idens  = density     * prof
        itemp  = temperature * prof
        edens[rho>1] = 1
        idens[rho>1] = 1

        out = {"nrho" : nrho, "nion" : 1, "anum" : np.array([anum]),
               "znum" : np.array([znum]), "mass" : np.array([mass]),
               "charge" : np.array([charge]), "rho" : rho,
               "edensity" : edens, "etemperature" : etemp, "idensity" : idens,
               "itemperature" : itemp}
        return ("plasma_1D", out)

    def neutral_flat(self, density=10e20, temperature=10e3, anum=1, znum=1):
        """Create uniform single-species constant neutral data.
        """
        density     = np.ones((100,1)) * density
        temperature = np.ones((100,1)) * temperature
        out = {"rhomin" : 0, "rhomax" : 10, "nrho" : 100, "nspecies" : 1,
               "anum" : np.array([anum]), "znum" : np.array([znum]),
               "density" : density, "temperature" : temperature,
               "maxwellian" : 1}
        return ("N0_1D", out)

    def wall_rectangular(self, nphi=1):
        """Create wall with a rectangular cross section.

        Parameters
        ----------
        nphi : int, optional
            Number of sectors in 3D wall.

            If larger than one, the wall input will be :class:`wall_3D` instead
            of :class:`wall_2D`.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.
        """
        out = {"nelements" : 20,
               "r" :  np.concatenate( (np.linspace(4.1, 8.4, 6)[1:],
                                       np.linspace(8.4, 8.4, 6)[1:],
                                       np.linspace(8.4, 4.1, 6)[1:],
                                       np.linspace(4.1, 4.1, 6)[1:]) ),
               "z" :  np.concatenate( (np.linspace(-3.9, -3.9, 6)[1:],
                                       np.linspace(-3.9, 3.9, 6)[1:],
                                       np.linspace(3.9, 3.9, 6)[1:],
                                       np.linspace(3.9, -3.9, 6)[1:]) )
        }
        if nphi > 1:
            out.update({"nphi" : nphi})
            return ("wall_3D", wall_3D.convert_wall_2D(**out))
        return ("wall_2D", out)

a5py/a5py/templates/analyticalinputs.py
ascot4fusion/ascot5
"""
Generate boozer input from a 2D bfield data.

This script takes the psi contours of a 2D equilibrium, and uses
those to construct the corresponding boozer grid. The output
is directly applicable as a boozer input for the ASCOT5.

The script evaluates the q profile in the process.

You need to have B_2DS active before running the script. If you only have
3D data, use write_hdf5_B2DS in ascot5io.B_3DS module to generate 2D
input.

File: generateboozer.py
"""
import sys
import numpy as np

from skimage import measure
from scipy.interpolate import griddata, interp1d
from scipy.integrate import solve_ivp, cumtrapz

import a5py.ascot5io.boozer as boozer
from a5py.ascot5io.ascot5 import Ascot
from a5py.ascotpy.ascotpy import Ascotpy


def generate(a5, rgrid, zgrid, npsi, nthgeo, nthbzr, zaxis, raxis, psi0, psi1,
             evalq=True, tstep=1e-2, nint=8000):

    # Resolutions for grids used internally
    rres  = 400
    zres  = 800

    # Set up the grids
    R = np.linspace(rgrid[0], rgrid[-1], rres) # For contours
    Z = np.linspace(zgrid[0], zgrid[-1], zres) # For contours

    # Evaluating contours right at the axis or separatrix might cause issues
    psimin = np.amin([psi0, psi1])
    psimax = np.amax([psi0, psi1])
    dpsi   = psimax - psimin
    psipad = 0.01
    psimin += psipad*dpsi
    psimax -= psipad*dpsi

    # For contours and thbzr and nu datas
    psigrid   = np.linspace(psimin, psimax, npsi)
    # For thtbzr data
    thgeogrid = np.linspace(0, 2*np.pi, nthgeo)
    # For nu data
    thbzrgrid = np.linspace(0, 2*np.pi, nthbzr+1)[:-1]
    # Angles at which contour points are evaluated
    thgrid = np.linspace(0, 2*np.pi, nint)

    ## Set up the data tables (psi can be evaluated directly) ##
    psitable = np.squeeze(a5.evaluate(R=rgrid, phi=0, z=zgrid, t=0,
                                      quantity="psi", grid=True))
    thtable  = np.zeros((psigrid.size,thgeogrid.size))
    nutable  = np.zeros((psigrid.size,thbzrgrid.size))

    # Also store q values to an array which can be returned
    qprof = np.zeros(psigrid.shape)
    Iprof = np.zeros(psigrid.shape)
    gprof = np.zeros(psigrid.shape)

    if psimin < psi0:
        psigrid = np.flip(psigrid)

    # Functions needed to trace psi = const. surfaces.
    def event(t, y):
        """Detect the midplane crossing.
        """
        return zaxis - y[1]

    # Terminate when midplane crossing is detected (see docmentation on
    # solve_ivp at scipy) and at right diection. delta is to ensure that
    # we ignore the first crossing by starting the simulation slightly off
    # plane.
    event.terminal  = True
    event.direction = 1 - 2 * (psi1 > psi0)
    delta = -1e-8 * event.direction

    def tracepsi(t, y):
        """Calculate new position when tracing Bpol.
        """
        br = a5.evaluate(R=y[0], phi=0, z=y[1], t=0, quantity="br")
        bz = a5.evaluate(R=y[0], phi=0, z=y[1], t=0, quantity="bz")
        return np.array([br / np.sqrt(br**2 + bz**2),
                         bz / np.sqrt(br**2 + bz**2)]).ravel()

    for i in range(psigrid.size):
        # Find (R,z) location at the outer midplane where psi(R,z) = psi(i)
        rzomp = a5.get_rhotheta_rz(
            np.sqrt((psigrid[i]-psi0) / (psi1 - psi0)), 0, 0, 0)

        # Solve the contour by tracing Bpol. The integration upper limit
        # (2 pi R) is set high as the integration is actually terminated
        # when OMP is crossed (this is taken care by the event. z0 + delta
        # ensures that we don't accidentally terminate it when we start
        # the integration.
        sol = solve_ivp(tracepsi, [0, 2*np.pi*raxis],
                        np.array([rzomp[0], rzomp[1]+delta[0]]).ravel(),
                        max_step=tstep, events=event)
        r = sol.y[0,:]
        z = sol.y[1,:]

        # Interpolate the contour points on the fixed (geometrical) theta
        # grid (at OMP we set thetageom=thetabzr=0)
        theta = np.mod(np.arctan2(z - zaxis, r - raxis), 2*np.pi)
        r = np.interp(thgrid, theta, r, period=2*np.pi)
        z = np.interp(thgrid, theta, z, period=2*np.pi)

        # Magnetic field along the contour (psi can be used to check that
        # the contour was set properly). Drop the last element in r and z
        # as it is the same as first.
        br   = a5.evaluate(R=r[:-1], phi=0, z=z[:-1], t=0, quantity="br")
        bphi = a5.evaluate(R=r[:-1], phi=0, z=z[:-1], t=0, quantity="bphi")
        bz   = a5.evaluate(R=r[:-1], phi=0, z=z[:-1], t=0, quantity="bz")

        bpol  = np.sqrt(br**2 + bz**2)
        bnorm = np.sqrt(br**2 + bphi**2 + bz**2)
        ds    = (np.diff(r) * br + np.diff(z) * bz) / bpol # darc dot e^_pol
        r = r[:-1]; z = z[:-1]

        # The toroidal current term (multiplying this with mu0/2 pi gets
        # enclosed toroidal current)
        Iprof[i] = np.sum( ds * bpol ) / ( 2*np.pi )

        # g = R*Bphi, since Bphi ~ 1/R this is a constant
        gprof[i] = r[0] * bphi[0]

        # The (global) safety factor q(psi)
        qprof[i] = np.sum( ds * gprof[i] / ( r**2 * bpol ) ) / ( 2*np.pi )

        # Boozer coordinate Jacobian is (I - qg) / B^2. Setting it fixes the
        # Boozer poloidal angle which we can now solve.
        jac = (Iprof + qprof*gprof)[i] / bnorm**2
        btheta = np.append(0, np.cumsum( ds / ( jac * bpol ) ))

        # The above Jacobian is for a periodical theta, so theta[-1] should
        # equal to 2 pi already, but normalize it to remove numerical error
        # (note that the new Jacobian would be J / a)
        a = 2*np.pi / btheta[-1]
        btheta *= a
        thtable[i, :] = interp1d(thgrid, btheta, "linear")(thgeogrid)

        # For Boozer toroidal coordinate, we need to integrate the local
        # safety factor along the contour
        nu = gprof[i] * np.append(0, np.cumsum( ds / ( r**2 * bpol ) ) )

        # Interpolate nu used in zeta = phi + nu(psi, theta)
        nutable[i,:] = -interp1d(btheta, nu, 'linear')(thbzrgrid) \
            + qprof[i] * thbzrgrid

    # Flip the data grids to set indices right
    if psimin < psi0:
        thtable = np.flip(thtable,axis=0)
        nutable = np.flip(nutable,axis=0)


    ## Construct the boozer input ##

    # Note that the last contour can be used to define separatrix location
    cr = np.append(r, r[0])
    cz = np.append(z, z[0])
    booz = {
        "psimin":psimin,
        "psimax":psimax,
        "npsi":int(psigrid.size),
        "ntheta":int(thbzrgrid.size),
        "nthetag":int(thgeogrid.size),
        "rmin":rgrid[0],
        "rmax":rgrid[-1],
        "nr":int(rgrid.size),
        "zmin":zgrid[0],
        "zmax":zgrid[-1],
        "nz":int(zgrid.size),
        "r0":raxis,
        "z0":zaxis,
        "psi0":psi0,
        "psi1":psi1,
        "psi_rz":psitable,
        "theta_psithetageom":thtable,
        "nu_psitheta":nutable,
        "nrzs":int(cr.size),
        "rs":cr,
        "zs":cz}

    if evalq:
        rhogrid = np.sqrt( (psigrid - psi0) / (psi1 - psi0) )
        return booz, rhogrid, qprof, gprof, Iprof
    else:
        return booz

if __name__ == "__main__":

    if len(sys.argv) < 2:
        print("Please give filename of the HDF5 file as an input. Aborting...")
        exit()

    fn = sys.argv[1]

    h5    = Ascot(fn)
    a5    = Ascotpy(fn)
    data  = h5.bfield.active.read()
    rgrid = np.linspace(data["rmin"], data["rmax"], int(data["nr"]))
    zgrid = np.linspace(data["zmin"], data["zmax"], int(data["nz"]))

    a5.init(bfield=True)
    booz, rhoprof, qprof, gprof, Iprof = generate(a5,
                                                  rgrid  = rgrid,
                                                  zgrid  = zgrid,
                                                  npsi   = 100,
                                                  nthgeo = 180,
                                                  nthbzr = 180,
                                                  raxis  = data["axisr"],
                                                  zaxis  = data["axisz"],
                                                  psi0   = data["psi0"],
                                                  psi1   = data["psi1"])
    a5.free(bfield=True)

    print("rho and q(rho):\n")
    print(repr( rhoprof.ravel() ))
    print(repr( qprof.ravel() ))

    boozer.write_hdf5(fn, desc="BOOZER", **booz)
    print("Boozerdata written.")

a5py/a5py/templates/generateboozer.py
ascot4fusion/ascot5
import os
import numpy as np

class ImportAdas():

    def import_adas(self,
                    fn_CX_DD0   = "input.sigmavCX_DD0",
                    fn_BMS_H0H  = "input.BMSsigmav_H0H",
                    fn_BMS_H0He = "input.BMSsigmav_H0He",
                    fn_BMS_H0C  = "input.BMSsigmav_H0C"):
        """Import data from ADAS files.

        Note: this method should be redone completely as right now it is just
        the script asigma_loc_read2hdf5.py. The problem with that script
        (and this method) is that the input files are expected to have in-house
        format. Instead, what we would like to have is that this script would
        take the ADAS files directly.

        Parameters
        ----------
        files : str or list [str]
            Folder where input files are located or a list of input files.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.
        """
        # Count the number of reactions for which data files can be found,
        # and account for possible multiples due to isotope combinations
        N_reac = 0
        CX_DD0_found = 0
        # Name HH0 instead of DD0 since we may, in the future, use this variable
        # for the total number of hydrogen-hydrogen CX reactions when different
        # isotope combinations are included.
        N_CX_HH0 = 0
        # NOTE: This list of BMS data filenames must be up to date with the
        # function parameter list.
        fn_BMSs = [fn_BMS_H0H,fn_BMS_H0He,fn_BMS_H0C]
        N_fn_BMS = len(fn_BMSs)
        BMS_founds = [0]*N_fn_BMS
        N_BMS = 0
        if(os.path.isfile(fn_CX_DD0)):
            CX_DD0_found = 1
            N_CX_HH0 = 1
        for i_BMS in range(0,N_fn_BMS):
            if(os.path.isfile(fn_BMSs[i_BMS])):
                BMS_founds[i_BMS] = 1
                N_BMS += 1
        N_reac = N_CX_HH0 + N_BMS

        # Initialize arrays for reaction identifiers and abscissae
        z_1       = np.zeros(N_reac, dtype=int)
        a_1       = np.zeros(N_reac, dtype=int)
        z_2       = np.zeros(N_reac, dtype=int)
        a_2       = np.zeros(N_reac, dtype=int)
        reac_type = np.zeros(N_reac, dtype=int)
        N_E       = np.zeros(N_reac, dtype=int)
        E_min     = np.zeros(N_reac, dtype=float)
        E_max     = np.zeros(N_reac, dtype=float)
        N_n       = np.zeros(N_reac, dtype=int)
        n_min     = np.zeros(N_reac, dtype=float)
        n_max     = np.zeros(N_reac, dtype=float)
        N_T       = np.zeros(N_reac, dtype=int)
        T_min     = np.zeros(N_reac, dtype=float)
        T_max     = np.zeros(N_reac, dtype=float)
        N         = np.zeros(N_reac, dtype=int)

        # Read atomic reaction data from local files
        #print("Reading atomic reaction data from local files")
        i_reac = 0
        # Read CX reaction data of the form sigmav(E,T)
        #print("  Trying to read D-D0 CX rate coefficient data")
        if(CX_DD0_found):
            f = open(fn_CX_DD0, "r")
            if f.mode == 'r':
                print("    Reading file " + fn_CX_DD0)
                lines = f.readlines()
                metadata = lines[0].split(' ')
                data_format = int(metadata[0]) # Not used
                z_1[i_reac]       = int(metadata[1])
                a_1[i_reac]       = int(metadata[2])
                z_2[i_reac]       = int(metadata[3])
                a_2[i_reac]       = int(metadata[4])
                reac_type[i_reac] = int(metadata[5])
                N_E[i_reac]       = int(metadata[6])
                N_n[i_reac]       = 1
                N_T[i_reac]       = int(metadata[7])
                N[i_reac]         = N_E[i_reac]*N_n[i_reac]*N_T[i_reac]
                abscissa1 = lines[1].split(' ')
                abscissa2 = lines[2].split(' ')
                ordinate  = lines[3].split(' ')
                # Currently, the ASCOT5 spline interpolation module only
                # supports uniform grid spacing. Hence, only min, max and N
                # is needed to reconstruct the abscissa on the ASCOT5 side.
                E_min[i_reac] = float(abscissa1[0])
                E_max[i_reac] = float(abscissa1[N_E[i_reac]-1])
                T_min[i_reac] = float(abscissa2[0])
                T_max[i_reac] = float(abscissa2[N_T[i_reac]-1])
                # Read the rate coefficient data
                sigmavCX_DD0 = np.zeros(N[i_reac])
                for i in range(0, N[i_reac]):
                    sigmavCX_DD0[i] = float(ordinate[i])
                # Increment the reaction index
                i_reac += 1
        else:
            print("    File not found: " + fn_CX_DD0)
        # Read BMS reaction data of the form BMSsigmav(E, n, T)
        #print("  Trying to read BMS data")
        # Initialize a list that can be filled with lists of unknown length below
        BMS_sigmavs = []
        i_BMS = 0
        for i_fn_BMS in range(0,N_fn_BMS):
            if(BMS_founds[i_fn_BMS]):
                f = open(fn_BMSs[i_fn_BMS], "r")
                if f.mode == 'r':
                    #print("    Reading file " + fn_BMSs[i_fn_BMS])
                    lines = f.readlines()
                    metadata = lines[0].split(' ')
                    data_format = int(metadata[0]) # Overwritten, because not used
                    z_1[i_reac]       = int(metadata[1])
                    a_1[i_reac]       = int(metadata[2])
                    z_2[i_reac]       = int(metadata[3])
                    a_2[i_reac]       = int(metadata[4])
                    reac_type[i_reac] = int(metadata[5])
                    N_E[i_reac]       = int(metadata[6])
                    N_n[i_reac]       = int(metadata[7])
                    N_T[i_reac]       = int(metadata[8])
                    N[i_reac]         = N_E[i_reac]*N_n[i_reac]*N_T[i_reac]
                    abscissa1 = lines[1].split(' ')
                    abscissa2 = lines[2].split(' ')
                    abscissa3 = lines[3].split(' ')
                    ordinate  = lines[4].split(' ')
                    # Currently, the ASCOT5 spline interpolation module only
                    # supports uniform grid spacing. Hence, only min, max and N is
                    # needed to reconstruct the abscissa on the ASCOT5 side.
                    # The density unit is converted (cm^-3 --> m^-3).
                    E_min[i_reac] =     float(abscissa1[0])
                    E_max[i_reac] =     float(abscissa1[N_E[i_reac]-1])
                    n_min[i_reac] = 1e6*float(abscissa2[0])
                    n_max[i_reac] = 1e6*float(abscissa2[N_n[i_reac]-1])
                    T_min[i_reac] =     float(abscissa3[0])
                    T_max[i_reac] =     float(abscissa3[N_T[i_reac]-1])
                    # Read the BMS coefficients and convert
                    # the unit (cm^3s^-1 --> m^3s^-1)
                    BMS_sigmavs.append(ordinate[0:N[i_reac]])
                    for i in range(0, N[i_reac]):
                        BMS_sigmavs[i_BMS][i] = 1e-6*float(BMS_sigmavs[i_BMS][i])
                    # Increment the BMS reaction index
                    i_BMS += 1
                    # Increment the reaction index
                    i_reac += 1
            else:
                print("    File not found: " + fn_BMSs[i_fn_BMS])

        # Move reaction data to one long array
        N_tot = 0
        for i_reac in range(0, N_reac):
            N_tot += N[i_reac]
        sigma = np.zeros((1,N_tot))
        i_reac  = 0
        i_sigma = 0
        if(CX_DD0_found):
            for i in range(0, N[i_reac]):
                sigma[0,i_sigma+i] = sigmavCX_DD0[i]
            i_sigma += N[i_reac]
            i_reac  += 1
        # We write no 'else' for this 'if' here because we already did above
        i_BMS = 0
        for i_fn_BMS in range(0,N_fn_BMS):
            if(BMS_founds[i_fn_BMS]):
                for i in range(0, N[i_reac]):
                    sigma[0,i_sigma+i] = BMS_sigmavs[i_BMS][i]
                i_sigma += N[i_reac]
                i_BMS   += 1
                i_reac  += 1

        out = {
            "nreac" : N_reac, "z1" : z_1, "a1" : a_1, "z2" : z_2, "a2" : a_2,
            "reactype" : reac_type, "nenergy" : N_E, "energymin" : E_min,
            "energymax" : E_max, "ndensity" : N_n, "densitymin" : n_min,
            "densitymax" : n_max, "ntemperature" : N_T,
            "temperaturemin" : T_min, "temperaturemax" : T_max, "sigma" : sigma
        }
        return ("asigma_loc", out)

a5py/a5py/templates/importdata.py
ascot4fusion/ascot5
"""Options templates for various simulations
"""
from a5py.ascot5io.options import Opt

class OptionTemplates():
    """All options templates.
    """

    def options_tutorial(self):
        """Create fast slowing-down simulation options with orbits and
        distributions enabled.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.
        """
        out = Opt.get_default()
        out.update(
            {"SIM_MODE":2, "ENABLE_ADAPTIVE":1,
             "ENDCOND_SIMTIMELIM":1,      "ENDCOND_CPUTIMELIM":1,
             "ENDCOND_ENERGYLIM":1,       "ENDCOND_WALLHIT":1,
             "ENDCOND_LIM_SIMTIME":0.5,   "ENDCOND_MAX_MILEAGE":0.5,
             "ENDCOND_MAX_CPUTIME":1.0e1, "ENDCOND_MIN_ENERGY":2.0e3,
             "ENDCOND_MIN_THERMAL":2.0,
             "ENABLE_ORBIT_FOLLOWING":1,  "ENABLE_COULOMB_COLLISIONS":1,
             "ENABLE_DIST_5D":1,          "ENABLE_DIST_6D":1,
             "ENABLE_DIST_RHO5D":1,       "ENABLE_DIST_RHO6D":1,
             "ENABLE_DIST_COM":1,
             "DIST_MIN_R":3.5, "DIST_MAX_R":8.5, "DIST_NBIN_R":12,
             "DIST_MIN_PHI":0, "DIST_MAX_PHI":360, "DIST_NBIN_PHI":20,
             "DIST_MIN_Z":-2.45, "DIST_MAX_Z":2.45, "DIST_NBIN_Z":24,
             "DIST_MIN_RHO":0, "DIST_MAX_RHO":1, "DIST_NBIN_RHO":11,
             "DIST_MIN_THETA":0, "DIST_MAX_THETA":360, "DIST_NBIN_THETA":13,
             "DIST_MIN_PPA":-10e-20,"DIST_MAX_PPA":10e-20,"DIST_NBIN_PPA":36,
             "DIST_MIN_PPE":0, "DIST_MAX_PPE":10e-20, "DIST_NBIN_PPE":18,
             "DIST_MIN_PR":-10e-20, "DIST_MAX_PR":10e-20,
             "DIST_NBIN_PR":14,
             "DIST_MIN_PPHI":-10e-20, "DIST_MAX_PPHI":10e-20,
             "DIST_NBIN_PPHI":15,
             "DIST_MIN_PZ":-10e-20, "DIST_MAX_PZ":10e-20, "DIST_NBIN_PZ":16,
             "DIST_MIN_TIME":0, "DIST_MAX_TIME":3e-2, "DIST_NBIN_TIME":2,
             "DIST_MIN_CHARGE":-1.5, "DIST_MAX_CHARGE":2.5,
             "DIST_NBIN_CHARGE":1,
             "ENABLE_ORBITWRITE":1, "ORBITWRITE_MODE":1,
             "ORBITWRITE_NPOINT":100, "ORBITWRITE_INTERVAL":0,}
        )
        return ("opt", out)

a5py/a5py/templates/optiontemplates.py
ascot4fusion/ascot5
"""Poincare run templates.
"""
import numpy as np
import unyt

from scipy.interpolate import interp1d
from scipy.integrate import solve_ivp, cumtrapz

from a5py.ascot5io.marker  import Marker
from a5py.ascot5io.options import Opt
from a5py.physlib import parseunits
from a5py.exceptions import *

class PoincareTemplates():

    def marker_poincare(self, nmrk=100, rhomin=0, rhomax=1, tor=None,
                        pol=0*unyt.deg, time=0*unyt.s, pitch=1.0, species=None,
                        energy=None):
        """Generate markers at fixed intervals in radius.

        By default, the generated markers are field lines. Physical particles
        can be generated by passing marker species and mass.

        Parameters
        ----------
        nmrk : int, optional
            Number of markers.
        rhomin : float, optional
            Minimum marker initial rho coordinate.
        rhomax : float, optional
            Maximum marker initial rho coordinate.
        tor : float, optional
            Initial marker toroidal coordinate.

            If None, markers are distributed uniformly in toroidal direction.
        pol : float, optional
            Initial marker poloidal coordinate.

            If None, markers are distributed uniformly in poloidal direction.
        time : float
            Time slice at when field line markers are traced.

            If physical particles are traced, this is the time-instant at when
            the simulation begins.
        pitch : float
            Field line marker direction (1) co- (2) counter-parallel to
            the magnetic field lines.

            If physical particles are traced, this parameter is the guiding
            center pitch (ppar/p) instead.

        species : str, optional
            Marker species to be simulated instead of field lines.

            **IMPORTANT** The simulation mode also needs to be set to GO or GC.
        energy : float
            Marker energy if physical particles are traced.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.

        Raises
        ------
        ValueError
            If species is provided but not energy.
        """
        if species is not None and energy is None:
            raise ValueError("Please provide both species and energy")

        if tor is None:
            tor = 360 * np.random.rand(nmrk,) * unyt.deg
        else:
            tor = tor * np.ones((nmrk,))

        if pol is None:
            pol = 360 * np.random.rand(nmrk,) * unyt.deg
        else:
            pol = pol * np.ones((nmrk,))

        rhovals = np.linspace(rhomin, rhomax, nmrk)
        r, z = self._ascot.input_rhotheta2rz(rhovals, pol, tor, time)

        if species is None:
            mrk = Marker.generate("fl", nmrk)
            mrk["r"][:]     = r
            mrk["z"][:]     = z
            mrk["phi"][:]   = tor
            mrk["pitch"][:] = pitch
            mrk["time"][:]  = time
            return ("fl", mrk)
        else:
            mrk = Marker.generate("gc", nmrk, species=species)
            mrk["r"][:]      = r
            mrk["z"][:]      = z
            mrk["phi"][:]    = tor
            mrk["pitch"][:]  = pitch
            mrk["time"][:]   = time
            mrk["energy"][:] = energy
            return ("gc", mrk)

    @parseunits(tor="deg", pol="deg")
    def options_poincare(self, simmode=4, tor=0*unyt.deg, pol=0*unyt.deg,
                         ntorb=500, nporb=500, both=False, maxrho=False,
                         mhd=False, cpulim=None):
        """Generate options to generate Poincaré plot data.

        Collisionless orbits are traced for a fixed number of toroidal and
        poloidal transits. The points, where the given poloidal and toroidal
        planes were crossed, are recorded.

        Parameters
        ----------
        simmode : {1, 2, 4}, optional
            Simulation mode (1: gyro-orbit, 2: guiding-center, 4: field line).
        tor : array_like, optional
            Toroidal angle(s) where data is gathered for (R,z) plot(s).
        pol : array_like, optional
            Poloidal angle(s) where data is gathered for (rho,phi) plot(s).
        ntorb : int, optional
            Number of full toroidal orbits before simulation is terminated.
        nporb : int, optional
            Number of full poloidal orbits before simulation is terminated.
        both : bool, optional
            Both ``ntorb`` and ``nporb`` conditions must be met before
            simulation is terminated.
        maxrho : bool, optional
            Terminate marker when it escapes plasma.

            Otherwise the marker is terminated at the wall.
        mhd : bool, optional
            Set ``True`` to enable MHD.
        cpulim : float, optional
            Maximum CPU time spent on simulating a single marker.

            This parameter can be used to ensure the simulation is terminated.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.
        """
        out = Opt.get_default()
        out.update({
            "SIM_MODE":simmode,
            "ENABLE_ADAPTIVE":1,
            "ENABLE_ORBIT_FOLLOWING":1,
            "ENABLE_ORBITWRITE":1,
            "ORBITWRITE_MODE":0,
            "ORBITWRITE_TOROIDALANGLES":tor,
            "ORBITWRITE_POLOIDALANGLES":pol,
            "ORBITWRITE_NPOINT":nporb*pol.size + ntorb*tor.size,
            "ENDCOND_MAX_POLOIDALORBS":nporb,
            "ENDCOND_MAX_TOROIDALORBS":ntorb,
            "ENDCOND_MAXORBS":1,
            "ENDCOND_WALLHIT":1,
            "ENDCOND_MAX_RHO":1.0,
        })

        if both:
            out["ENDCOND_MAXORBS"] = 2
        if maxrho:
            out["ENDCOND_RHOLIM"] = 1
        if cpulim is not None:
            out["ENDCOND_CPUTIMELIM"]  = 1
            out["ENDCOND_MAX_CPUTIME"] = cpulim
        if mhd:
            out["ENABLE_MHD":mhd]

        return ("opt", out)

    def boozer_tokamak(self, npsi=100, nthgeo=200, nthbzr=200, tol=1e-5,
                       nint=10000, rhomin=0.3, rhomax=0.9):
        """Build mapping from real-space to Boozer coordinates assuming
        axisymmetric tokamak field.

        Parameters
        ----------
        npsi : int, optional
            Number of psi grid points.
        nthgeo : int, optional
            Number of geometrical theta grid points.
        nthbzr : int, optional
            Number of boozer theta grid points.
        nint : int, optional
            Number of points in line integral evaluations.
        rhomin : float, optional
            Minimum rho in the radial grid which cannot be zero since
            the Boozer coordinates cannot be computed there.
        rhomax : float, optional
            Maximum rho in the radial grid which cannot be one since
            the separatrix does not form a closed surface when there is
            an X-point.

        Returns
        -------
        gtype : str
            Type of the generated input data.
        data : dict
            Input data that can be passed to ``write_hdf5`` method of
            a corresponding type.
        """
        inp = self._ascot.input_initialized()
        if "bfield" not in inp:
            raise AscotInitException("bfield not initialized")
        grp   = self._ascot.data._get_group(inp["bfield"])
        gtype = grp.get_type()
        if gtype == "B_2DS":
            d = grp.read()
            rmin   = d["rmin"]
            rmax   = d["rmax"]
            nrcntr = int(d["nr"])
            zmin   = d["zmin"]
            zmax   = d["zmax"]
            nzcntr = int(d["nz"])
            psi0   = d["psi0"]
            psi1   = d["psi1"]
            raxis  = d["axisr"]
            zaxis  = d["axisz"]
            del d
        else:
            raise AscotNoDataException("bfield is not B_2DS")

        # The boozer coordinate system is constructed by integrating along
        # psi = constant surfaces. Use this grid to evaluate psi values
        # on a (R,z) plane.
        rgrid = np.linspace(rmin, rmax, nrcntr)
        zgrid = np.linspace(zmin, zmax, nzcntr)

        # ...and this poloidal grid to evaluate values along the contour
        thgrid = np.linspace(0, 2*np.pi, nint)

        # Boozer coordinate psi-grid. Add a little bit of padding to psi0 and
        # psi1 values as otherwise making the contour at those points could
        # yield funny results.
        pmin  = float(psi0 + rhomin**2 * ( psi1 - psi0 ))
        pmax  = float(psi0 + rhomax**2 * ( psi1 - psi0 ))
        psimin  = np.amin([pmin, pmax])
        psimax  = np.amax([pmin, pmax])
        psigrid = np.linspace(psimin, psimax, npsi)
        if psi1 < psi0: psigrid = np.flip(psigrid) # Ensure grid start at psi0

        # Boozer coordinate theta grid
        thgeogrid = np.linspace(0, 2*np.pi, nthgeo)

        # Boozer coordinate nu grid
        thbzrgrid = np.linspace(0, 2*np.pi, nthbzr+1)[:-1]

        # Set up the data tables (psi can be evaluated directly)
        thtable  = np.zeros( (psigrid.size, thgeogrid.size) )
        nutable  = np.zeros( (psigrid.size, thbzrgrid.size) )
        psitable = self._ascot.input_eval(
            rgrid*unyt.m, 0*unyt.deg, zgrid*unyt.m, 0*unyt.s, "psi", grid=True)
        psitable = np.squeeze(psitable)

        # Helper quantities evaluated when the coordinate transform is made
        qprof = np.zeros(psigrid.shape)
        Iprof = np.zeros(psigrid.shape)
        gprof = np.zeros(psigrid.shape)

        # Calculate Boozer angular coordinates for each psi
        for i in range(psigrid.size):

            # Interpolate the contour points on the fixed (geometrical) theta
            # grid (at OMP we set thetageom=thetabzr=0)
            rhogrid = np.sqrt((psigrid[i]-psi0) / (psi1 - psi0)) \
                * np.ones(thgrid.shape)*unyt.dimensionless
            r, z = self._ascot.input_rhotheta2rz(
                rhogrid, thgrid*unyt.rad, np.zeros(thgrid.shape)*unyt.rad,
                0*unyt.s, tol=tol)

            # Magnetic field along the contour (psi can be used to check that
            # the contour was set properly). Drop the last element in r and z
            # as it is the same as first.
            br, bphi, bz, psi = self._ascot.input_eval(
                r[:-1], 0*unyt.rad, z[:-1], 0*unyt.s,
                "br", "bphi", "bz", "psi")

            bpol  = np.sqrt(br**2 + bz**2)
            bnorm = np.sqrt(br**2 + bphi**2 + bz**2)
            ds    = (np.diff(r) * br + np.diff(z) * bz) / bpol # darc dot e^_pol
            r = r[:-1]; z = z[:-1]

            # The toroidal current term (multiplying this with mu0/2 pi gets
            # enclosed toroidal current)
            Iprof[i] = np.sum( ds * bpol ) / ( 2*np.pi )

            # g = R*Bphi, since Bphi ~ 1/R this is a constant
            gprof[i] = r[0] * bphi[0]

            # The (global) safety factor q(psi)
            qprof[i] = np.sum( ds * gprof[i] / ( r**2 * bpol ) ) / ( 2*np.pi )

            # Boozer coordinate Jacobian is (I - qg) / B^2. Setting it fixes the
            # Boozer poloidal angle which we can now solve.
            jac = (Iprof + qprof*gprof)[i] / bnorm**2
            btheta = np.append(0, np.cumsum( ds / ( jac * bpol ) ))

            # The above Jacobian is for a periodical theta, so theta[-1] should
            # equal to 2 pi already, but normalize it to remove numerical error
            # (note that the new Jacobian would be J / a)
            a = 2*np.pi / btheta[-1]
            btheta *= a
            thtable[i, :] = interp1d(thgrid, btheta, "linear")(thgeogrid)

            # For Boozer toroidal coordinate, we need to integrate the local
            # safety factor along the contour
            nu = gprof[i] * np.append(0, np.cumsum( ds / ( r**2 * bpol ) ) )

            # Interpolate nu used in zeta = phi + nu(psi, theta)
            nutable[i,:] = -interp1d(btheta, nu, 'linear')(thbzrgrid) \
                + qprof[i] * thbzrgrid

        # Flip the data grids to set indices right
        if psi1 < psi0:
            thtable = np.flip(thtable,axis=0)
            nutable = np.flip(nutable,axis=0)

        # The last contour can be used to define separatrix location
        cr = np.append(r, r[0])
        cz = np.append(z, z[0])

        #Create input
        return ("Boozer", {
            "psimin":psimin, "psimax":psimax, "npsi":int(psigrid.size),
            "ntheta":int(thbzrgrid.size), "nthetag":int(thgeogrid.size),
            "rmin":rgrid[0], "rmax":rgrid[-1], "nr":int(rgrid.size),
            "zmin":zgrid[0], "zmax":zgrid[-1], "nz":int(zgrid.size),
            "r0":raxis, "z0":zaxis, "psi0":psi0, "psi1":psi1,
            "psi_rz":psitable, "theta_psithetageom":thtable,
            "nu_psitheta":nutable, "nrzs":int(cr.size), "rs":cr, "zs":cz} )

a5py/a5py/templates/poincare.py
ascot4fusion/ascot5
"""
Helper routines and attributes for test modules.

File: testascot/helpers.py
"""
import subprocess
import numpy as np

from time import sleep
from time import perf_counter as timer

import a5py.ascot5io.ascot5      as ascot5
import a5py.ascot5io.ascot5tools as tools
import a5py.ascot5io.N0_3D       as N0_3D

## Name of the test HDF5 file
testfn = "test_ascot.h5"

## Name of the ascot5 binary
testbin = "ascot5_main"

def clean_opt(odict):
    """
    Turn all tunable features off in given options dictionary.

    Args:
        odict: dict Options dictionary
    """
    odict["RECORD_MODE"] = np.array([0],dtype='i4')
    for o in odict:
        if o.startswith("ENABLE"):
            odict[o] = np.array([0],dtype='i4')
        if o.startswith("ENDCOND"):
            odict[o] = np.array([0],dtype='i4')
        if o.startswith("DISABLE"):
            odict[o] = np.array([0],dtype='i4')

    odict["ENDCOND_MAX_MILEAGE"] = 1


def set_correct_input(parent, test):
    """
    Set that input field active that contains test name in its description.

    Args:
        parent:  str Input parent
        test:   str Name of the tests
    """
    a5 = ascot5.Ascot(testfn)
    typ = a5[parent][test].get_type()
    qid = a5[parent][test].get_qid()
    group = typ + "-" + qid
    tools.call_ascot5file(testfn, "set_active", group)


def set_and_run(test):
    """
    Set correct inputs active and carry out the test simulation.

    Args:
        testfn: str Name of the HDF5 file with tests
        test:   str Name of the test
    """
    set_correct_input("bfield",  test)
    set_correct_input("efield",  test)
    set_correct_input("marker",  test)
    set_correct_input("plasma",  test)
    set_correct_input("neutral", test)
    set_correct_input("wall",    test)
    set_correct_input("options", test)
    set_correct_input("boozer",  test)
    set_correct_input("mhd",     test)
    
    sleep(1.01) # Sleep for one second so that each run gets unique QID

    frm   = lambda x: "%.3f s" % x
    start = timer()
    subprocess.call(["./"+testbin, "--in="+testfn[:-3], "--d="+test],
                    stdout=subprocess.DEVNULL)
    print("Completed test " + test + " in " + frm(timer() - start)) 

a5py/a5py/testascot/helpers.py
ascot4fusion/ascot5
"""Module for verifying that ASCOT5 can model the physics it claims.

Implemented tests:

  - elementary: verify that ASCOT5 reproduces correct gyromotion and drifts
    in a simple electromagnetic field.
  - orbitfollowing: verify the conservation properties of orbit-integrators.
  - gctransform: verify the guiding center transformation implementation.
  - ccoll: verify the Coulomb collision operator.
  - classical: verify that ASCOT5 reproduces classical transport correctly.
  - neoclassical: verify that ASCOT5 reproduces the neoclassical transport
    correctly
  - boozer: verify Boozer coordinate transformation.
  - mhd: verify inclusion of MHD modes.
  - atomic: verify implementation of ionization and neutralization reactions.
"""
import copy
import unyt
import subprocess
import warnings
import numpy as np
import matplotlib.pyplot as plt

from matplotlib.gridspec import GridSpec

from a5py import Ascot, physlib
from a5py.routines import plotting as a5plt
from a5py.ascot5io.options import Opt
from a5py.ascot5io.marker import Marker
from a5py.ascot5io.bfield import B_2DS

class PhysTest():

    tag_elementary_gyro    = "TESTELEMENTARYGYRO"
    tag_elementary_exbgo   = "TESTELEMENTARYEXBGO"
    tag_elementary_exbgc   = "TESTELEMENTARYEXBGC"
    tag_elementary_gradbgo = "TESTELEMENTARYGRADBGO"
    tag_elementary_gradbgc = "TESTELEMENTARYGRADBGC"
    tag_orbfol_go          = "TESTORBFOLGO"
    tag_orbfol_gcf         = "TESTORBFOLGCF"
    tag_orbfol_gca         = "TESTORBFOLGCA"
    tag_gctransform_go     = "TESTGCTRANSFORMGO"
    tag_gctransform_gc     = "TESTGCTRANSFORMGC"
    tag_gctransform_go2gc  = "TESTGCTRANSFORMGO2GC"
    tag_gctransform_zeroth = "TESTGCTRANSFORMZEROTH"
    tag_gctransform_first  = "TESTGCTRANSFORMFIRST"
    tag_ccoll_thermalgo    = "TESTCCOLLTHERMALGO"
    tag_ccoll_thermalgcf   = "TESTCCOLLTHERMALGCF"
    tag_ccoll_thermalgca   = "TESTCCOLLTHERMALGCA"
    tag_ccoll_slowinggo    = "TESTCCOLLSLOWINGGO"
    tag_ccoll_slowinggcf   = "TESTCCOLLSLOWINGGCF"
    tag_ccoll_slowinggca   = "TESTCCOLLSLOWINGGCA"
    tag_classical_go       = "TESTCLASSGO"
    tag_classical_gcf      = "TESTCLASSGCF"
    tag_classical_gca      = "TESTCLASSGCA"
    tag_neoclassical_go    = "TESTNEOCLASSGO"
    tag_neoclassical_gcf   = "TESTNEOCLASSGCF"
    tag_neoclassical_gca   = "TESTNEOCLASSGCA"
    tag_boozer             = "TESTBOOZER"
    tag_mhd_go             = "TESTMHDGO"
    tag_mhd_gcf            = "TESTMHDGCF"
    tag_mhd_gca            = "TESTMHDGCA"
    tag_atomic_cx          = "TESTATOMICCX"
    tag_atomic_ionz        = "TESTATOMICIONZ"

    def __init__(self, fn="testascot.h5"):
        try:
            self.ascot = Ascot(fn)
        except FileNotFoundError:
            self.ascot = Ascot(fn, create=True)
            init = self.ascot.data.create_input
            init("opt")
            init("gc")
            init("B_TC")
            init("E_TC")
            init("wall_2D")
            init("plasma_1D")
            init("N0_1D")
            init("Boozer")
            init("MHD_STAT")
            init("asigma_loc")

    def execute(self, init=True, run=True, check=True, tests=None):
        """Execute test(s).

        Parameters
        ----------
        """
        if tests is None:
            tests = ["elementary", "orbitfollowing", "gctransform", "ccoll",
                     "classical", "neoclassical", "boozer", "mhd", "atomic"]
        elif isinstance(tests, str):
            tests = [tests]

        for test in tests:
            if init:
                getattr(self, "init_" + test)()
                print("Test %s initialized" % test)
            if run:
                getattr(self, "run_" + test)()
                print("Test %s simulation complete" % test)
            if check:
                a5plt.setpaperstyle()
                getattr(self, "check_" + test)()
                print("Test %s check finished" % test)

    def init_elementary(self):
        """Initialize data for the elementary test.

        This test covers elementary results such as the gyromotion and drifts by
        verifying that ASCOT5 reproduces:

        1. correct gyroradius and gyrofrequency in an uniform magnetic field.
        2. correct E X B drift in uniform electromagnetic field.
        3. correct gradB drift in a magnetic field with constant gradient.

        Test 1 is done with GO mode and 2 and 3 with both GO and GC modes,
        latter using the fixed step scheme. Tests are done in Cartesian
        electromagnetic field (B_TC and E_TC) and without collisions.
        The test particles are an energetic electron and a positron so
        the tests also verify that ASCOT5 is valid in the relativistic regime.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_elementary_gyro):
            warnings.warn("Inputs already present: Test elementary")
            return
        init = self.ascot.data.create_input

        # Options
        opt = Opt.get_default()
        opt.update({
            "SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
            "FIXEDSTEP_USERDEFINED" : 1e-11, "ENDCOND_SIMTIMELIM" : 1,
            "ENDCOND_LIM_SIMTIME" : 2e-9, "ENABLE_ORBIT_FOLLOWING" : 1,
            "ENABLE_ORBITWRITE" : 1, "ORBITWRITE_MODE" : 1,
            "ORBITWRITE_INTERVAL" : 1e-11, "ORBITWRITE_NPOINT" : 202
        })
        init("opt", **opt, desc=PhysTest.tag_elementary_gyro)
        opt.update({
            "FIXEDSTEP_USERDEFINED" : 1e-10, "ENDCOND_LIM_SIMTIME" : 1e-7,
            "ORBITWRITE_NPOINT" : 10002
        })
        init("opt", **opt, desc=PhysTest.tag_elementary_exbgo)
        init("opt", **opt, desc=PhysTest.tag_elementary_gradbgo)
        opt.update({
            "SIM_MODE" : 2, "FIXEDSTEP_USERDEFINED" : 1e-9,
            "ORBITWRITE_INTERVAL" : 1e-9, "ORBITWRITE_NPOINT" : 102
        })
        init("opt", **opt, desc=PhysTest.tag_elementary_exbgc)
        init("opt", **opt, desc=PhysTest.tag_elementary_gradbgc)

        # Magnetic field
        d = {"bxyz" : np.array([5, 0, 0]),
             "jacobian" : np.array([0,0,0,0,0,0,0,0,0]),
             "rhoval" : 1.5}
        init("B_TC", **d, desc=PhysTest.tag_elementary_gyro)
        init("B_TC", **d, desc=PhysTest.tag_elementary_exbgo)
        init("B_TC", **d, desc=PhysTest.tag_elementary_exbgc)

        d.update({"jacobian" : np.array([0,0,0.1,0,0,0,0,0,0])})
        init("B_TC", **d, desc=PhysTest.tag_elementary_gradbgo)
        init("B_TC", **d, desc=PhysTest.tag_elementary_gradbgc)

        # Electric field
        d = {"exyz" : np.array([0, 0, 0])}
        init("E_TC", **d, desc=PhysTest.tag_elementary_gyro)
        init("E_TC", **d, desc=PhysTest.tag_elementary_gradbgo)
        init("E_TC", **d, desc=PhysTest.tag_elementary_gradbgc)

        d.update({"exyz" : np.array([0, 1e6, 0])})
        init("E_TC", **d, desc=PhysTest.tag_elementary_exbgo)
        init("E_TC", **d, desc=PhysTest.tag_elementary_exbgc)

        # Marker input is an electron and positron
        mrk = Marker.generate("gc", n=2, species="electron")
        mrk["charge"]    = np.array([1, -1])
        mrk["r"][:]      = 5
        mrk["phi"][:]    = 90
        mrk["z"][:]      = 0
        mrk["zeta"][:]   = 0
        mrk["energy"][:] = 100e6
        mrk["pitch"][:]  = 0.5

        for tag in [PhysTest.tag_elementary_gyro,
                    PhysTest.tag_elementary_exbgo,
                    PhysTest.tag_elementary_exbgc,
                    PhysTest.tag_elementary_gradbgo,
                    PhysTest.tag_elementary_gradbgc]:
            init("gc", **mrk, desc=tag)

    def run_elementary(self):
        """Run elementary test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_elementary_gyro):
            warnings.warn("Results already present: Test elementary")
            return
        for tag in [PhysTest.tag_elementary_gyro,
                    PhysTest.tag_elementary_exbgo,
                    PhysTest.tag_elementary_exbgc,
                    PhysTest.tag_elementary_gradbgo,
                    PhysTest.tag_elementary_gradbgc]:
            self._activateinputs(tag)
            self._runascot(tag)

    def check_elementary(self):
        """Verify and plot elementary test results.

        Returns
        -------
        passed : bool
            True if the test passed.
        """
        run_gyro    = self.ascot.data[PhysTest.tag_elementary_gyro]
        run_exbgo   = self.ascot.data[PhysTest.tag_elementary_exbgo]
        run_exbgc   = self.ascot.data[PhysTest.tag_elementary_exbgc]
        run_gradbgo = self.ascot.data[PhysTest.tag_elementary_gradbgo]
        run_gradbgc = self.ascot.data[PhysTest.tag_elementary_gradbgc]

        # Initialize plots
        fig = a5plt.figuredoublecolumn()
        gs = GridSpec(2, 3, figure=fig)
        h1a = fig.add_subplot(gs[0,0])
        h2a = fig.add_subplot(gs[0,1])
        h2b = fig.add_subplot(gs[1,1])
        h3a = fig.add_subplot(gs[0,2])
        h3b = fig.add_subplot(gs[1,2])

        h1a.set_title("Gyro-motion")
        h2a.set_title("ExB drift")
        h3a.set_title("Grad-B drift")

        h1a.set_ylabel("z [m]")
        h1a.set_xlabel("y [m]")

        h1a.set_aspect("equal", adjustable="box")
        h2a.set_aspect("equal", adjustable="box")
        h2b.set_aspect("equal", adjustable="box")
        h3a.set_aspect("equal", adjustable="box")
        h3b.set_aspect("equal", adjustable="box")

        def plotarrow(ax, ini, end, txt):
            """Plot line with an arrow
            """
            ax.annotate("", xytext=ini, xy=end,
                        arrowprops={"arrowstyle":"->", "color":"black"})
            ax.plot([ini[0], end[0]], [ini[1], end[1]], color="black")
            ax.annotate(txt, ini + .005)

        ## Gyro test ##
        m, q, pitch, ekin, x0, y0 = \
            run_gyro.getstate("mass", "charge", "pitch", "ekin", "y", "z")
        x, y, time = run_gyro.getorbit("y", "z", "time", ids=1)

        # Analytical values
        bnorm = np.sqrt(np.sum(run_gyro.bfield.read()["bxyz"]**2)) * unyt.T
        larmorrad_ana = physlib.gyrolength(m, q, ekin, pitch, bnorm).to("m")
        gyrofreq_ana  = physlib.gyrofrequency(m, q, ekin, bnorm).to("rad/s")

        # Numerical values
        larmorrad_go = np.mean( np.sqrt( (x - x0[0])**2 + (y - y0[0])**2 ) )
        gyrofreq_go  = np.sum( np.sqrt( np.diff(x)**2 + np.diff(y)**2 ) ) \
            * unyt.rad / (larmorrad_go * time[-1])

        # Plot
        orbx = x - x0[0]
        orby = y - y0[0]

        anax   = larmorrad_ana[0] * np.sin(np.linspace(0, 2*np.pi, 360))
        anay   = larmorrad_ana[0] * np.cos(np.linspace(0, 2*np.pi, 360))
        fangle = np.arctan2(orby[0], orbx[0]) * unyt.rad \
            + gyrofreq_ana[0] * time[-1]
        fx = larmorrad_ana[0] * np.cos(-fangle) # Minus sign because this we are
        fy = larmorrad_ana[0] * np.sin(-fangle) # in LHS coordinate system.

        h1a.plot(orbx, orby, linewidth=3)
        h1a.plot(anax, anay, linestyle="--", color="black", alpha=0.7)
        h1a.scatter(orbx[-1], orby[-1], marker="o", color="C0", zorder=3)
        h1a.scatter(fx, fy, marker="x", color="black", zorder=4)

        ## E x B drift
        yi_go            = run_exbgo.getstate("z", state="ini")
        yf_go, deltat_go = run_exbgo.getstate("z", "mileage", state="end")
        x_go1, y_go1     = run_exbgo.getorbit("y", "z", ids=1)
        x_go2, y_go2     = run_exbgo.getorbit("y", "z", ids=2)

        yi_gc            = run_exbgc.getstate("z", state="ini")
        yf_gc, deltat_gc = run_exbgc.getstate("z", "mileage", state="end")
        x_gc1, y_gc1     = run_exbgc.getorbit("y", "z", ids=1)
        x_gc2, y_gc2     = run_exbgc.getorbit("y", "z", ids=2)

        # Analytical values
        bvec  = run_exbgo.bfield.read()["bxyz"].ravel()
        evec  = run_exbgo.efield.read()["exyz"].ravel()
        v_ExB = np.cross(evec, bvec) / np.inner(bvec, bvec) * unyt.m / unyt.s
        time  = run_exbgo.options.read()["ENDCOND_LIM_SIMTIME"] * unyt.s

        # Numerical values
        vgo_ExB = (yf_go - yi_go) / deltat_go
        vgc_ExB = (yf_gc - yi_gc) / deltat_gc

        # Plot
        h2a.plot(x_go1, y_go1)
        h2b.plot(x_go2, y_go2)
        h2a.plot(x_gc1, y_gc1)
        h2b.plot(x_gc2, y_gc2)

        ini = np.array([x_gc1[0].v + .01, y_gc1[0].v])
        end = np.array([x_gc1[0].v + .01, (y_gc1[0] + v_ExB[2]*deltat_gc[1]).v])
        plotarrow(h2a, ini, end, r"$e^+$")
        ini = np.array([x_gc2[0].v + .01, y_gc2[0].v])
        end = np.array([x_gc2[0].v + .01, (y_gc2[0] + v_ExB[2]*deltat_gc[0]).v])
        plotarrow(h2b, ini, end, r"$e^-$")

        ## gradB drift
        pitch, ekin, q, m = run_gradbgo.getstate(
            "pitch", "ekin", "charge", "mass", ids=1)
        xi_go            = run_gradbgo.getstate("y", state="ini")
        xf_go, deltat_go = run_gradbgo.getstate("y", "mileage", state="end")
        x_go1, y_go1     = run_gradbgo.getorbit("y", "z", ids=1)
        x_go2, y_go2     = run_gradbgo.getorbit("y", "z", ids=2)

        xi_gc            = run_gradbgc.getstate("y", state="ini")
        xf_gc, deltat_gc = run_gradbgc.getstate("y", "mileage", state="end")
        x_gc1, y_gc1     = run_gradbgc.getorbit("y", "z", ids=1)
        x_gc2, y_gc2     = run_gradbgc.getorbit("y", "z", ids=2)

        # Analytical values
        gamma   = physlib.gamma_energy(m, ekin)
        vnorm   = physlib.vnorm_gamma(gamma)
        bvec    = run_gradbgo.bfield.read()["bxyz"].ravel()
        gradb   = run_gradbgo.bfield.read()["jacobian"][0,:]
        bnorm   = np.sqrt(np.sum(bvec**2)) * unyt.T
        v_gradb = gamma * m * ( (1.0 - pitch**2) * vnorm**2 / ( q * bnorm ) ) \
            * 0.5 * np.cross(bvec, gradb) / bnorm**2 * unyt.T**2 / unyt.m
        v_gradb.convert_to_units("m/s")
        time    = run_gradbgo.options.read()["ENDCOND_LIM_SIMTIME"] * unyt.s

        # Numerical values
        vgo_gradb = (xf_go - xi_go) / deltat_go
        vgc_gradb = (xf_gc - xi_gc) / deltat_gc

        # Plot
        h3a.plot(x_go1, y_go1)
        h3b.plot(x_go2, y_go2)
        h3a.plot(x_gc1, y_gc1)
        h3b.plot(x_gc2, y_gc2)

        ini = np.array([x_gc1[0].v, y_gc1[0].v + .01])
        end = np.array([x_gc1[0] + v_gradb[1]*deltat_gc[0], y_gc1[0].v + .01])
        plotarrow(h3a, ini, end, r"$e^+$")
        ini = np.array([x_gc2[0].v, y_gc2[0] + .01])
        end = np.array([x_gc2[0] - v_gradb[1]*deltat_gc[0], y_gc2[0].v + .01])
        plotarrow(h3b, ini, end, r"$e^-$")

        print("Test elementary:")
        passed = True

        print("  Gyroradius %e, Gyrofrequency %e (gyro-orbit)" \
              % (larmorrad_go, gyrofreq_go))
        err = np.abs(larmorrad_go - larmorrad_ana[0]) / larmorrad_ana[0]
        if err > 5e-8: print("Error: %e (FAILED)" % err); passed = False

        print("  Gyroradius %e, Gyrofrequency %e (expected)" \
              % (larmorrad_ana[0], gyrofreq_ana[0]))
        err = np.abs(gyrofreq_go - gyrofreq_ana[0]) / gyrofreq_ana[0]
        if err > 5e-4: print("Error: %e (FAILED)" % err); passed = False

        print("  ExB drift for positrons %e electrons %e (gyro-orbit)" \
              % (vgo_ExB[0], vgo_ExB[1]) )
        print("  ExB drift for positrons %e electrons %e (guiding center)" \
              % (vgc_ExB[0], vgc_ExB[1]) )
        print("  ExB drift for positrons %e electrons %e (expected)" \
              % (v_ExB[2], v_ExB[2]) )
        err = np.amax([np.abs( ( vgo_ExB[0] - v_ExB[2] ) / v_ExB[2] ),
                       np.abs( ( vgo_ExB[1] - v_ExB[2] ) / v_ExB[2] ),
                       np.abs( ( vgc_ExB[0] - v_ExB[2] ) / v_ExB[2] ),
                       np.abs( ( vgc_ExB[1] - v_ExB[2] ) / v_ExB[2] )])
        if err > 5e-5: print("Error: %e (FAILED)" % err); passed = False

        print("  Grad-B drift for positrons %e electrons %e (gyro-orbit)" \
              % (vgo_gradb[0], vgo_gradb[1]) )
        print("  Grad-B drift for positrons %e electrons %e (guiding center)" \
              % (vgc_gradb[0], vgc_gradb[1]) )
        print("  Grad-B drift for positrons %e electrons %e (expected)" \
              % (v_gradb[1], -v_gradb[1]) )
        err = np.amax([np.abs( ( vgo_gradb[0] - v_gradb[1] ) / v_gradb[1] ),
                       np.abs( ( vgo_gradb[1] + v_gradb[1] ) / v_gradb[1] ),
                       np.abs( ( vgc_gradb[0] - v_gradb[1] ) / v_gradb[1] ),
                       np.abs( ( vgc_gradb[1] + v_gradb[1] ) / v_gradb[1] )])
        if err > 1e-3: print("Error: %e (FAILED)" % err); passed = False

        return passed

    def init_orbitfollowing(self):
        """Initialize data for the  test.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_orbfol_go):
            warnings.warn("Inputs already present: Test orbit-following")
            return
        init = self.ascot.data.create_input

        # Options
        opt = Opt.get_default()
        opt.update({
            "SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
            "FIXEDSTEP_USERDEFINED" : 1e-12, "ENDCOND_SIMTIMELIM" : 1,
            "ENDCOND_LIM_SIMTIME" : 5e-6, "ENABLE_ORBIT_FOLLOWING" : 1,
            "ENABLE_ORBITWRITE" : 1, "ORBITWRITE_MODE" : 1,
            "ORBITWRITE_INTERVAL" : 1e-10, "ORBITWRITE_NPOINT" : 50002
        })
        init("opt", **opt, desc=PhysTest.tag_orbfol_go)
        opt.update({
            "SIM_MODE" : 2, "FIXEDSTEP_USERDEFINED" : 1e-12,
            "ENABLE_ADAPTIVE" : 0,
            "ORBITWRITE_INTERVAL" : 1e-8, "ORBITWRITE_NPOINT" : 502
        })
        init("opt", **opt, desc=PhysTest.tag_orbfol_gcf)
        opt.update({
            "ENABLE_ADAPTIVE" : 1, "ADAPTIVE_MAX_DRHO" : 0.1,
            "ADAPTIVE_MAX_DPHI" : 10, "ADAPTIVE_TOL_ORBIT" : 1e-10,
            "FIXEDSTEP_USERDEFINED" : 1e-8
        })
        init("opt", **opt, desc=PhysTest.tag_orbfol_gca)

        # Magnetic field is just some tokamak
        init("bfield_analytical_iter_circular", desc=PhysTest.tag_orbfol_go)
        init("bfield_analytical_iter_circular", desc=PhysTest.tag_orbfol_gcf)
        init("bfield_analytical_iter_circular", desc=PhysTest.tag_orbfol_gca)

        # Marker input is a trapped positron and a passing electron
        mrk = Marker.generate("gc", n=2, species="electron")
        mrk["charge"]    = np.array([1, -1])
        mrk["r"][:]      = 7.6
        mrk["phi"][:]    = 90
        mrk["z"][:]      = 0
        mrk["zeta"][:]   = 2
        mrk["energy"][:] = 10e6
        mrk["pitch"]     = np.array([0.4, 0.9])
        for tag in [PhysTest.tag_orbfol_go, PhysTest.tag_orbfol_gcf,
                    PhysTest.tag_orbfol_gca]:
            init("gc", **mrk, desc=tag)

    def run_orbitfollowing(self):
        """Run orbit-following test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_orbfol_go):
            warnings.warn("Results already present: Test orbit-following")
            return
        for tag in [PhysTest.tag_orbfol_go, PhysTest.tag_orbfol_gcf,
                    PhysTest.tag_orbfol_gca]:
            self._activateinputs(tag)
            self._runascot(tag)

    def check_orbitfollowing(self):
        """Check test.
        """
        run_go  = self.ascot.data[PhysTest.tag_orbfol_go]
        run_gcf = self.ascot.data[PhysTest.tag_orbfol_gcf]
        run_gca = self.ascot.data[PhysTest.tag_orbfol_gca]

        # Initialize plots
        fig = a5plt.figuredoublecolumn(3/2)
        gs = GridSpec(3, 4, figure=fig)
        h1a = fig.add_subplot(gs[0,0])
        h2a = fig.add_subplot(gs[1,0])
        h3a = fig.add_subplot(gs[2,0])
        h4a = fig.add_subplot(gs[:,1])
        h1b = fig.add_subplot(gs[0,2])
        h2b = fig.add_subplot(gs[1,2])
        h3b = fig.add_subplot(gs[2,2])
        h4b = fig.add_subplot(gs[:,3])

        h4a.set_aspect("equal", adjustable="box")
        h4b.set_aspect("equal", adjustable="box")

        h1a.set_xlim(0, 5)
        h2a.set_xlim(0, 5)
        h3a.set_xlim(0, 5)
        h1b.set_xlim(0, 5)
        h2b.set_xlim(0, 5)
        h3b.set_xlim(0, 5)

        h4a.set_xlim(5, 8)
        h4a.set_ylim(-2, 2)
        h4b.set_xlim(5, 8)
        h4b.set_ylim(-2, 2)

        h1a.set_xticklabels([])
        h2a.set_xticklabels([])
        h1b.set_xticklabels([])
        h2b.set_xticklabels([])

        h4a.set_yticklabels([])
        h4b.set_xlabel("R [m]")
        h4b.set_ylabel("z [m]")
        h4b.yaxis.set_label_position("right")
        h4b.yaxis.tick_right()

        h1a.set_ylabel(r"$(E-E_0)/E_0$")
        h2a.set_ylabel(r"$(\mu-\mu_0)/\mu_0$")
        h3a.set_ylabel(r"$(P-P_0)/P_0$")

        h3a.set_xlabel("Time [µs]")
        h3b.set_xlabel("Time [µs]")

        def plotreldiff(ax1, ax2, ax3, t, q1, q2, q3, **kwargs):
            """Plot relative diffenrence of given quantities on given axes.
            """
            ax1.plot(t, ( q1 - q1[0] ) / q1[0], **kwargs)
            ax2.plot(t, ( q2 - q2[0] ) / q2[0], **kwargs)
            ax3.plot(t, ( q3 - q3[0] ) / q3[0], **kwargs)

        def fails(t, q, eps, qnt, otype, mode):
            """Check if the change in time of a given quantity is below given
            tolerance
            """
            err = np.polyfit(t, (q - q[0]) / q[0], 1)[0]
            msg = "Rate of change in %6s (%s/%11s): %e Tolerance: %e" \
                % (qnt, otype, mode, err, eps)
            if err > eps:
                msg += " (FAILED)"
                print(msg)
                return False
            print(msg)
            return True

        # Numerical values
        self.ascot.input_init(run=run_go.get_qid(), bfield=True)
        tgo1, ego1, mugo1, pgo1, rgo1, zgo1 = run_go.getorbit(
            "mileage", "ekin", "mu", "ptor", "r", "z", ids=1)
        tgo2, ego2, mugo2, pgo2, rgo2, zgo2 = run_go.getorbit(
            "mileage", "ekin", "mu", "ptor", "r", "z", ids=2)
        tgcf1, egcf1, mugcf1, pgcf1, rgcf1, zgcf1 = run_gcf.getorbit(
            "mileage", "ekin", "mu", "ptor", "r", "z", ids=1)
        tgcf2, egcf2, mugcf2, pgcf2, rgcf2, zgcf2 = run_gcf.getorbit(
            "mileage", "ekin", "mu", "ptor", "r", "z", ids=2)
        tgca1, egca1, mugca1, pgca1, rgca1, zgca1 = run_gca.getorbit(
            "mileage", "ekin", "mu", "ptor", "r", "z", ids=1)
        tgca2, egca2, mugca2, pgca2, rgca2, zgca2 = run_gca.getorbit(
            "mileage", "ekin", "mu", "ptor", "r", "z", ids=2)
        self.ascot.input_free()

        # Plot
        plotreldiff(h1a, h2a, h3a, tgo1.to("µs"),  ego1,  mugo1,  pgo1,
                    color="C0")
        plotreldiff(h1b, h2b, h3b, tgo2.to("µs"),  ego2,  mugo2,  pgo2,
                    color="C0")
        plotreldiff(h1a, h2a, h3a, tgcf1.to("µs"), egcf1, mugcf1, pgcf1,
                    color="C1")
        plotreldiff(h1b, h2b, h3b, tgcf2.to("µs"), egcf2, mugcf2, pgcf2,
                    color="C1")
        plotreldiff(h1a, h2a, h3a, tgca1.to("µs"), egca1, mugca1, pgca1,
                    color="C2")
        plotreldiff(h1b, h2b, h3b, tgca2.to("µs"), egca2, mugca2, pgca2,
                    color="C2")

        h4a.plot(rgo1,  zgo1,  color="C0")
        h4a.plot(rgcf1, zgcf1, color="C1")
        h4a.plot(rgca1, zgca1, color="C2")
        h4b.plot(rgo2,  zgo2,  color="C0")
        h4b.plot(rgcf2, zgcf2, color="C1")
        h4b.plot(rgca2, zgca2, color="C2")

        print("Test orbit-following:")
        passed = True

        otype1 = "passing"; otype2 = "trapped"
        simmode1 = "gyro-orbit"; simmode2 = "fixed GC"; simmode3 = "adaptive GC"
        if fails(tgo1,  ego1,   5e-6, "energy", otype1, simmode1):passed = False
        if fails(tgo2,  ego2,   5e-6, "energy", otype2, simmode1):passed = False
        if fails(tgcf1, egcf1,  5e-3, "energy", otype1, simmode2):passed = False
        if fails(tgcf2, egcf2,  5e-5, "energy", otype2, simmode2):passed = False
        if fails(tgca1, egca1,  5e-3, "energy", otype1, simmode3):passed = False
        if fails(tgca2, egca2,  5e-5, "energy", otype2, simmode3):passed = False
        if fails(tgo1,  mugo1,  2e2,  "mu",     otype1, simmode1):passed = False
        if fails(tgo2,  mugo2,  1e-1, "mu",     otype2, simmode1):passed = False
        if fails(tgcf1, mugcf1, 1e-9, "mu",     otype1, simmode2):passed = False
        if fails(tgcf2, mugcf2, 1e-9, "mu",     otype2, simmode2):passed = False
        if fails(tgca1, mugca1, 1e-9, "mu",     otype1, simmode3):passed = False
        if fails(tgca2, mugca2, 1e-9, "mu",     otype2, simmode3):passed = False
        if fails(tgo1,  pgo1,   5e-3, "Ptor",   otype1, simmode1):passed = False
        if fails(tgo2,  pgo2,   5e-5, "Ptor",   otype2, simmode1):passed = False
        if fails(tgcf1, pgcf1,  5e-2, "Ptor",   otype1, simmode2):passed = False
        if fails(tgcf2, pgcf2,  5e-2, "Ptor",   otype2, simmode2):passed = False
        if fails(tgca1, pgca1,  5e-2, "Ptor",   otype1, simmode3):passed = False
        if fails(tgca2, pgca2,  5e-2, "Ptor",   otype2, simmode3):passed = False

        return passed

    def init_gctransform(self):
        """Initialize data for the guiding-center transformation test.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_gctransform_go):
            warnings.warn("Inputs already present: Test GC transform")
            return
        init = self.ascot.data.create_input

        # Options
        opt = Opt.get_default()
        opt.update({
            "SIM_MODE" : 2, "FIXEDSTEP_USE_USERDEFINED" : 1,
            "FIXEDSTEP_USERDEFINED" : 1e-10, "ENDCOND_SIMTIMELIM" : 1,
            "ENDCOND_LIM_SIMTIME" : 3e-5, "ENABLE_ORBIT_FOLLOWING" : 1,
            "ENABLE_ORBITWRITE" : 1, "ORBITWRITE_MODE" : 1,
            "ORBITWRITE_INTERVAL" : 4e-10, "ORBITWRITE_NPOINT" : 75002
        })
        init("opt", **opt, desc=PhysTest.tag_gctransform_gc)
        init("opt", **opt, desc=PhysTest.tag_gctransform_first)
        opt.update({"DISABLE_FIRSTORDER_GCTRANS" : 1})
        init("opt", **opt, desc=PhysTest.tag_gctransform_zeroth)

        opt.update({"DISABLE_FIRSTORDER_GCTRANS" : 0, "SIM_MODE" : 1 })
        init("opt", **opt, desc=PhysTest.tag_gctransform_go)
        opt.update({"RECORD_MODE" : 1 })
        init("opt", **opt, desc=PhysTest.tag_gctransform_go2gc)

        # Magnetic field is just some tokamak
        init("bfield_analytical_iter_circular",
             desc=PhysTest.tag_gctransform_go)
        init("bfield_analytical_iter_circular",
             desc=PhysTest.tag_gctransform_gc)
        init("bfield_analytical_iter_circular",
             desc=PhysTest.tag_gctransform_go2gc)
        init("bfield_analytical_iter_circular",
             desc=PhysTest.tag_gctransform_zeroth)
        init("bfield_analytical_iter_circular",
             desc=PhysTest.tag_gctransform_first)

        # Use single alpha particle in tests
        mrk = Marker.generate("gc", n=1, species="alpha")
        mrk["r"][:]      = 7.6
        mrk["phi"][:]    = 90
        mrk["z"][:]      = 0
        mrk["zeta"][:]   = 2
        mrk["energy"][:] = 3.5e6
        mrk["pitch"][:]  = 0.4
        init("gc", **mrk, desc=PhysTest.tag_gctransform_go)
        init("gc", **mrk, desc=PhysTest.tag_gctransform_gc)
        init("gc", **mrk, desc=PhysTest.tag_gctransform_go2gc)

    def run_gctransform(self):
        """Run GC transform test.
        """
        init = self.ascot.data.create_input
        if hasattr(self.ascot.data, PhysTest.tag_gctransform_go):
            warnings.warn("Results already present: Test GC transform")
            return
        for tag in [PhysTest.tag_gctransform_go, PhysTest.tag_gctransform_gc,
                    PhysTest.tag_gctransform_go2gc]:
            self._activateinputs(tag)
            self._runascot(tag)

        # Create new marker input from results
        nrep = 10
        mrk = Marker.generate("prt", n=nrep, species="alpha")
        run = self.ascot.data[PhysTest.tag_gctransform_go]

        # Pick initial coordinates along the orbit trajectory
        dt   = 20
        idx = np.s_[:nrep*dt:dt]
        time, r, phi, z, vr, vphi, vz = \
            run.getorbit("time", "r", "phi", "z", "vr", "vphi", "vz")
        mrk.update({
            "time" : time[idx], "r" : r[idx], "phi" : phi[idx], "z" : z[idx],
            "vr" : vr[idx], "vphi" : vphi[idx], "vz" : vz[idx]
        })
        init("prt", **mrk, desc=PhysTest.tag_gctransform_zeroth)
        init("prt", **mrk, desc=PhysTest.tag_gctransform_first)

        for tag in [PhysTest.tag_gctransform_zeroth,
                    PhysTest.tag_gctransform_first]:
            self._activateinputs(tag)
            self._runascot(tag)

    def check_gctransform(self):
        """Check test.
        """
        run_go     = self.ascot.data[PhysTest.tag_gctransform_go]
        run_gc     = self.ascot.data[PhysTest.tag_gctransform_gc]
        run_go2gc  = self.ascot.data[PhysTest.tag_gctransform_go2gc]
        run_zeroth = self.ascot.data[PhysTest.tag_gctransform_zeroth]
        run_first  = self.ascot.data[PhysTest.tag_gctransform_first]

        # Initialize plots
        fig = a5plt.figuredoublecolumn(3/2)
        gs  = GridSpec(3, 3, figure=fig)
        h1a = fig.add_subplot(gs[0,0])
        h1b = fig.add_subplot(gs[1,0])
        h1c = fig.add_subplot(gs[2,0])
        h2  = fig.add_subplot(gs[:,1])
        h3  = fig.add_subplot(gs[:,2])

        h1a.set_xlim(0, 30)
        h1b.set_xlim(0, 30)
        h1c.set_xlim(0, 30)
        h1a.set_xticklabels([])
        h1b.set_xticklabels([])

        h2.set_xlim(6.2, 6.7)
        h2.set_ylim(1.0, 1.6)
        h2.set_aspect("equal", adjustable="box")
        h3.set_xlim(6.2, 6.7)
        h3.set_ylim(1.0, 1.6)
        h3.set_aspect("equal", adjustable="box")
        h3.set_yticklabels([])

        h1c.set_xlabel("Time [µs]")
        h1a.set_ylabel(r"$(E-E_0)/E_0$")
        h1b.set_ylabel(r"$(\mu-\mu_0)/\mu_0$")
        h1c.set_ylabel(r"$(P-P_0)/P_0$")
        h2.set_xlabel("R [m]")
        h2.set_ylabel("z [m]")

        # Get data and plot
        self.ascot.input_init(run=run_go.get_qid(), bfield=True)
        rgo, zgo, tgo, ego, mugo, pgo = run_go.getorbit(
            "r", "z", "mileage", "ekin", "mu", "ptor")
        rgc, zgc, tgc, egc, mugc, pgc = run_gc.getorbit(
            "r", "z", "mileage", "ekin", "mu", "ptor")
        rgo2gc, zgo2gc, tgo2gc, ego2gc, mugo2gc, pgo2gc = run_go2gc.getorbit(
            "r", "z", "mileage", "ekin", "mu", "ptor")
        self.ascot.input_free()

        h1a.plot(tgo.to("µs"), mugo)
        h1a.plot(tgo2gc.to("µs"), mugo2gc)
        h1a.plot(tgc.to("µs"), mugc)

        h1b.plot(tgo.to("µs"), ego)
        h1b.plot(tgo2gc.to("µs"), ego2gc)
        h1b.plot(tgc.to("µs"), egc)

        h1c.plot(tgo.to("µs"), pgo)
        h1c.plot(tgo2gc.to("µs"), pgo2gc)
        h1c.plot(tgc.to("µs"), pgc)

        h2.plot(rgo, zgo)
        h2.plot(rgo2gc, zgo2gc)
        h2.plot(rgc, zgc)

        nrep = run_zeroth.getstate("ids").size
        for i in range(nrep):
            r, z = run_zeroth.getorbit("r", "z", ids=i)
            h3.plot(r, z, color="C1")
        for i in range(nrep):
            r, z = run_first.getorbit("r", "z", ids=i)
            h3.plot(r.v+0.01, z, color="C2")
        h3.plot(rgo2gc, zgo2gc, color="black")

        # Verify results by interpolating the orbits at fixed intervals and then
        # calculating sum-of-squares of the difference between go and gc and
        # go2gc and gc. The latter should be smaller if the guiding center
        # transformation works.
        t = np.linspace(0, 1e-6, 1000)
        mugo    = np.interp(t, tgo,    mugo)
        mugc    = np.interp(t, tgc,    mugc)
        mugo2gc = np.interp(t, tgo2gc, mugo2gc)
        ego    = np.interp(t, tgo,    ego)
        egc    = np.interp(t, tgc,    egc)
        ego2gc = np.interp(t, tgo2gc, ego2gc)
        pgo    = np.interp(t, tgo,    pgo)
        pgc    = np.interp(t, tgc,    pgc)
        pgo2gc = np.interp(t, tgo2gc, pgo2gc)

        err1 = np.sum( (mugo    - mugc)**2 )
        err2 = np.sum( (mugo2gc - mugc)**2 )
        print(err2/err1)
        err1 = np.sum( (ego    - egc)**2 )
        err2 = np.sum( (ego2gc - egc)**2 )
        print(err2/err1)
        err1 = np.sum( (pgo    - pgc)**2 )
        err2 = np.sum( (pgo2gc - pgc)**2 )
        print(err2/err1)

    def init_ccoll(self):
        """Initialize data for the Coulomb collision test.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_ccoll_thermalgo):
            warnings.warn("Inputs already present: Test Coulomb collision")
            return
        init = self.ascot.data.create_input

        # Options
        opt = Opt.get_default()
        opt.update({
            "ENDCOND_SIMTIMELIM" : 1, "ENDCOND_LIM_SIMTIME" : 2e-2,
            "ENABLE_ORBIT_FOLLOWING" : 1, "ENABLE_COULOMB_COLLISIONS" : 1,
            "ENABLE_DIST_5D" : 1,
            "DIST_MIN_R" : 4, "DIST_MAX_R" : 10, "DIST_NBIN_R" : 1,
            "DIST_MIN_PHI" : 0, "DIST_MAX_PHI" : 360, "DIST_NBIN_PHI" : 1,
            "DIST_MIN_Z" : -5, "DIST_MAX_Z" : 5, "DIST_NBIN_Z" : 1,
            "DIST_MIN_TIME" : 0, "DIST_MAX_TIME" : 2e-2, "DIST_NBIN_TIME" : 2,
            "DIST_MIN_PPA" : -2.5e-21, "DIST_MAX_PPA" : 2.5e-21,
            "DIST_NBIN_PPA" : 140, "DIST_MIN_PPE" : 0, "DIST_MAX_PPE" : 2.5e-21,
            "DIST_NBIN_PPE" : 80
        })

        opt.update({"SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
                    "FIXEDSTEP_USERDEFINED" : 1e-8})
        init("opt", **opt, desc=PhysTest.tag_ccoll_thermalgo)
        opt.update({"SIM_MODE" : 2, "FIXEDSTEP_USE_USERDEFINED" : 1,
                    "FIXEDSTEP_USERDEFINED" : 2e-8})
        init("opt", **opt, desc=PhysTest.tag_ccoll_thermalgcf)
        opt.update({"ENABLE_ADAPTIVE" : 1, "ADAPTIVE_TOL_ORBIT" : 1e-6,
                    "ADAPTIVE_TOL_CCOL" : 1e-2, "ADAPTIVE_MAX_DRHO" : 0.1,
                    "ADAPTIVE_MAX_DPHI" : 10, "FIXEDSTEP_USERDEFINED" : 1e-8})
        init("opt", **opt, desc=PhysTest.tag_ccoll_thermalgca)

        opt = Opt.get_default()
        opt.update({
            "ENDCOND_ENERGYLIM" : 1, "ENDCOND_MIN_ENERGY" : 50e3,
            "ENDCOND_MIN_THERMAL" : 0, "ENABLE_ORBIT_FOLLOWING" : 1,
            "ENABLE_COULOMB_COLLISIONS" : 1, "ENABLE_DIST_5D" : 1,
            "DIST_MIN_R" : 4, "DIST_MAX_R" : 10, "DIST_NBIN_R" : 1,
            "DIST_MIN_PHI" : 0, "DIST_MAX_PHI" : 360, "DIST_NBIN_PHI" : 1,
            "DIST_MIN_Z" : -5, "DIST_MAX_Z" : 5, "DIST_NBIN_Z" : 1,
            "DIST_MIN_TIME" : 0, "DIST_MAX_TIME" : 2e-2, "DIST_NBIN_TIME" : 2,
            "DIST_MIN_PPA" : -1.3e-19, "DIST_MAX_PPA" : 1.3e-19,
            "DIST_NBIN_PPA" : 200, "DIST_MIN_PPE" : 0, "DIST_MAX_PPE" : 1.3e-19,
            "DIST_NBIN_PPE" : 100
        })

        opt.update({"SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
                    "FIXEDSTEP_USERDEFINED" : 2e-9})
        init("opt", **opt, desc=PhysTest.tag_ccoll_slowinggo)
        opt.update({"SIM_MODE" : 2, "FIXEDSTEP_USE_USERDEFINED" : 1,
                    "FIXEDSTEP_USERDEFINED" : 3e-8})
        init("opt", **opt, desc=PhysTest.tag_ccoll_slowinggcf)
        opt.update({"ENABLE_ADAPTIVE" : 1, "ADAPTIVE_TOL_ORBIT" : 1e-6,
                    "ADAPTIVE_TOL_CCOL" : 1e-2, "ADAPTIVE_MAX_DRHO" : 0.1,
                    "ADAPTIVE_MAX_DPHI" : 10, "FIXEDSTEP_USERDEFINED" : 1e-8})
        init("opt", **opt, desc=PhysTest.tag_ccoll_slowinggca)

        # Magnetic field is just some tokamak and plasma is uniform
        for tag in [PhysTest.tag_ccoll_thermalgo, PhysTest.tag_ccoll_thermalgcf,
                    PhysTest.tag_ccoll_thermalgca, PhysTest.tag_ccoll_slowinggo,
                    PhysTest.tag_ccoll_slowinggcf, PhysTest.tag_ccoll_slowinggca
                    ]:
            init("bfield_analytical_iter_circular", desc=tag)
            init("plasma_flat", density=1e20, temperature=1e3, desc=tag)

        mrk = Marker.generate("gc", n=20, species="proton")
        pol = 2 * np.pi * np.random.rand(mrk["n"],)
        mrk["r"][:]      = 6.2 + 0.8 * np.cos(pol)
        mrk["phi"][:]    = 90
        mrk["z"][:]      = 0.8 * np.sin(pol)
        mrk["zeta"][:]   = 2 * np.pi * np.random.rand(mrk["n"],)
        mrk["energy"][:] = 1.e3
        mrk["pitch"][:]  = 0.5
        init("gc", **mrk, desc=PhysTest.tag_ccoll_thermalgo)
        init("gc", **mrk, desc=PhysTest.tag_ccoll_thermalgcf)
        init("gc", **mrk, desc=PhysTest.tag_ccoll_thermalgca)

        mrk = Marker.generate("gc", n=200, species="alpha")
        pol = 2 * np.pi * np.random.rand(mrk["n"],)
        mrk["r"][:]      = 6.2 + 0.8 * np.cos(pol)
        mrk["phi"][:]    = 90
        mrk["z"][:]      = 0.8 * np.sin(pol)
        mrk["zeta"][:]   = 2 * np.pi * np.random.rand(mrk["n"],)
        mrk["energy"][:] = 3.5e6
        mrk["pitch"][:]  = 1.0 - 2.0 * np.random.rand(mrk["n"],)
        init("gc", **mrk, desc=PhysTest.tag_ccoll_slowinggo)
        init("gc", **mrk, desc=PhysTest.tag_ccoll_slowinggcf)
        init("gc", **mrk, desc=PhysTest.tag_ccoll_slowinggca)

    def run_ccoll(self):
        """Run Coulmb collision test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_ccoll_thermalgo):
            warnings.warn("Results already present: Test Coulomb collision")
            return
        for tag in [PhysTest.tag_ccoll_thermalgo, PhysTest.tag_ccoll_thermalgcf,
                    PhysTest.tag_ccoll_thermalgca, PhysTest.tag_ccoll_slowinggo,
                    PhysTest.tag_ccoll_slowinggcf, PhysTest.tag_ccoll_slowinggca
                    ]:
            self._activateinputs(tag)
            self._runascot(tag)

    def check_ccoll(self):
        """Check Coulomb collision test.
        """
        run_tgo  = self.ascot.data[PhysTest.tag_ccoll_thermalgo]
        run_tgcf = self.ascot.data[PhysTest.tag_ccoll_thermalgcf]
        run_tgca = self.ascot.data[PhysTest.tag_ccoll_thermalgca]
        run_sgo  = self.ascot.data[PhysTest.tag_ccoll_slowinggo]
        run_sgcf = self.ascot.data[PhysTest.tag_ccoll_slowinggcf]
        run_sgca = self.ascot.data[PhysTest.tag_ccoll_slowinggca]

        fig = a5plt.figuredoublecolumn()
        gs = GridSpec(2, 2, figure=fig)
        h1 = fig.add_subplot(gs[0,0])
        h2 = fig.add_subplot(gs[0,1])
        h3 = fig.add_subplot(gs[1,0])
        h4 = fig.add_subplot(gs[1,1])

        # Analytical results
        alphaZ = 2
        clog   = 16
        vth    = np.sqrt(2*Te*e / m_e)
        vcrit  = vth * np.power( (3.0*np.sqrt(np.pi)/4.0) * (m_e / m_p) , 1/3.0)
        Ecrit  = 0.5 * m_a * vcrit * vcrit / e
        ts     = 3 * np.sqrt( np.power(2*np.pi * Te * e, 3) / m_e ) * eps_0 * eps_0 \
            * m_a /( alphaZ * alphaZ * np.power(e, 4) * ne * clog)

        heaviside = np.logical_and(SLOWING["Egrid"] <= Esd,
                                   SLOWING["Egrid"] >= 50*Te)

        THERMAL["analytical"]  = 2 * np.sqrt(THERMAL["Egrid"]/np.pi) \
            * np.power(Te,-3.0/2) \
            * np.exp(-THERMAL["Egrid"]/Te)
        THERMAL["analytical"] *= (simtime_th/2)
        SLOWING["analytical"] = heaviside * ts \
            / ( ( 1 + np.power(Ecrit/SLOWING["Egrid"], 3.0/2) )\
                * 2 * SLOWING["Egrid"] )

        # ts is slowing down rate which gives the slowing down time as
        # t_sd = ts*log(v_0 / v_th) = 0.5*ts*log(E_0/E_th)
        slowingdowntime = 0.5*ts*np.log(Esd/(50*Te))

        for run in [run_tgo, run_tgcf, run_tgca]:
            dist = run.getdist("5d", exi=True)
            edist = dist.integrate(
                r=np.s_[:], phi=np.s_[:], z=np.s_[:], time=np.s_[:],
                charge=np.s_[:], pitch=np.s_[:], copy=True)
            xdist = dist.integrate(
                r=np.s_[:], phi=np.s_[:], z=np.s_[:], time=np.s_[:],
                charge=np.s_[:], ekin=np.s_[:], copy=True)
            run.plotdist(edist, axes=h1)
            run.plotdist(xdist, axes=h2)

        for run in [run_sgo, run_sgcf, run_sgca]:
            dist = run.getdist("5d", exi=True)
            edist = dist.integrate(
                r=np.s_[:], phi=np.s_[:], z=np.s_[:], time=np.s_[:],
                charge=np.s_[:], pitch=np.s_[:], copy=True)
            xdist = dist.integrate(
                r=np.s_[:], phi=np.s_[:], z=np.s_[:], time=np.s_[:],
                charge=np.s_[:], ekin=np.s_[:], copy=True)
            run.plotdist(edist, axes=h3)
            run.plotdist(xdist, axes=h4)

        plt.show()

    def init_classical(self):
        """Initialize data for the classical transport test.
        """
        if hasattr(self.ascot.data.options,
                   PhysTest.tag_classical_go + "0"):
            warnings.warn("Inputs already present: Test classical transport")
            return
        init = self.ascot.data.create_input

        # Options
        optgo = Opt.get_default()
        optgo.update({
            "SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
            "FIXEDSTEP_USERDEFINED" : 1e-10, "ENDCOND_SIMTIMELIM" : 1,
            "ENDCOND_LIM_SIMTIME" : 5e-6, "ENABLE_ORBIT_FOLLOWING" : 1,
            "ENABLE_COULOMB_COLLISIONS" : 1
        })
        optgcf = copy.deepcopy(optgo)
        optgcf.update({
            "SIM_MODE" : 2, "FIXEDSTEP_USERDEFINED" : 1e-9
        })
        optgca = copy.deepcopy(optgcf)
        optgca.update({
            "ENABLE_ADAPTIVE" : 1, "ADAPTIVE_MAX_DRHO" : 0.1,
            "ADAPTIVE_TOL_ORBIT" : 1e-8, "ADAPTIVE_TOL_CCOL" : 1e-1,
            "ADAPTIVE_MAX_DPHI" : 10, "FIXEDSTEP_USERDEFINED" : 1e-8
        })

        # Marker input consists of protons
        mrk = Marker.generate("gc", n=200, species="proton")
        mrk["r"][:]      = 5.0
        mrk["phi"][:]    = 0.0
        mrk["z"][:]      = 0.0
        mrk["energy"][:] = 1e3
        mrk["pitch"][:]  = 1.0 - 2.0 * np.random.rand(mrk["n"],)
        mrk["zeta"][:]   = 2 * np.pi * np.random.rand(mrk["n"],)

        # Plasma consisting of electrons only to avoid proton-proton collisions
        pls = init("plasma_flat", density=1e22, temperature=1e3, dryrun=True)
        pls["idensity"][:] = 1

        for i in range(6):
            init("opt", **optgo,  desc=PhysTest.tag_classical_go  + str(i))
            init("opt", **optgcf, desc=PhysTest.tag_classical_gcf + str(i))
            init("opt", **optgca, desc=PhysTest.tag_classical_gca + str(i))
            for tag in [PhysTest.tag_classical_go, PhysTest.tag_classical_gcf,
                        PhysTest.tag_classical_gca]:
                # Transport is scanned as a function of magnetic field strength
                d = {"bxyz" : np.array([1, 0, 0]), "rhoval" : 0.5,
                     "jacobian" : np.array([0,0,0,0,0,0,0,0,0])}
                d["bxyz"][0] = 1.0 + i * (10.0 - 1.0) / 5
                init("B_TC", **d, desc=tag + str(i))
                init("gc", **mrk, desc=tag + str(i))
                init("plasma_1D", **pls, desc=tag + str(i))

    def run_classical(self):
        """Run classical transport test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_classical_go):
            warnings.warn("Results already present: Test classical transport")
            return
        for tag in [PhysTest.tag_classical_go, PhysTest.tag_classical_gcf,
                    PhysTest.tag_classical_gca]:
            i = 0
            while tag + str(i) in self.ascot.data.bfield.ls(show=False):
                self._activateinputs(tag+str(i))
                self._runascot(tag+str(i))
                i += 1

    def check_classical(self):
        """Check classical transport test.
        """
        nscan = 0
        while PhysTest.tag_classical_go + str(nscan) in \
              self.ascot.data.bfield.ls(show=False):
            nscan += 1

        # Initialize plots
        fig = a5plt.figuredoublecolumn(3/2)
        ax = fig.add_subplot(1,1,1)

        # Numerical values
        ndim = 2 # Diffusion happens on 2D plane
        bnorm = np.zeros((nscan,)) * unyt.T
        Dgo   = np.zeros((nscan,)) * unyt.m**2 / unyt.s
        Dgcf  = np.zeros((nscan,)) * unyt.m**2 / unyt.s
        Dgca  = np.zeros((nscan,)) * unyt.m**2 / unyt.s
        for i in range(nscan):
            run_go  = self.ascot.data[PhysTest.tag_classical_go  + str(i)]
            run_gcf = self.ascot.data[PhysTest.tag_classical_gcf + str(i)]
            run_gca = self.ascot.data[PhysTest.tag_classical_gca + str(i)]

            yi, zi, ti = run_go.getstate("y", "z", "mileage", state="ini")
            yf, zf, tf = run_go.getstate("y", "z", "mileage", state="end")
            Dgo[i] = np.mean( ( (yi - yf)**2 + (zi - zf)**2 ) / (tf - ti) ) \
                / (2*ndim)
            yi, zi, ti = run_gcf.getstate("y", "z", "mileage", state="ini")
            yf, zf, tf = run_gcf.getstate("y", "z", "mileage", state="end")
            Dgcf[i] = np.mean( ( (yi - yf)**2 + (zi - zf)**2 ) / (tf - ti) ) \
                / (2*ndim)
            yi, zi, ti = run_gca.getstate("y", "z", "mileage", state="ini")
            yf, zf, tf = run_gca.getstate("y", "z", "mileage", state="end")
            Dgca[i] = np.mean( ( (yi - yf)**2 + (zi - zf)**2 ) / (tf - ti) ) \
                / (2*ndim)

            bnorm[i] = np.sqrt(np.sum(run_go.bfield.read()["bxyz"]**2))

        # Analytical
        clog = 13.4
        ekin = run_go.getstate("ekin")[0]
        self.ascot.input_init(run=run_go.get_qid(), bfield=True, plasma=True)
        ne, Te = self.ascot.input_eval(6.2, 0, 0, 0, "ne", "te")
        self.ascot.input_free()
        #collfreq = ( (unyt.me / unyt.mp) * ( np.sqrt(2 / np.pi) / 3 )
        #             * ( unyt.e**2 / ( 4 * np.pi * unyt.eps_0 ) )**2 * ne * clog
        #             * 4 * np.pi / np.sqrt( unyt.me * (Te)**3 ) ).to("1/s")
        collfreq = physlib.collfreq_ie(unyt.mp, unyt.e, ne, Te, clog)
        rhog = physlib.gyrolength(unyt.mp, 1*unyt.e, ekin, 0.0, bnorm).to("m")
        Dana = collfreq * rhog**2 / 2

        # Plotting
        ax.scatter(1/bnorm**2, Dgo)
        ax.scatter(1/bnorm**2, Dgcf)
        ax.scatter(1/bnorm**2, Dgca)
        ax.plot(1/bnorm**2, Dana, color="black")

    def init_neoclassical(self):
        """Initialize data for the neoclassical transport test.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_neoclassical_go):
            warnings.warn("Inputs already present: Test neoclass. transport")
            return
        init = self.ascot.data.create_input

        # Ion densities to be scanned
        ni = np.power( 10, np.linspace(17.5, 22.0, 20) )

        # Options (some parameters are changed between the scans)
        optgo = Opt.get_default()
        optgo.update({
            "SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
            "FIXEDSTEP_USERDEFINED" : 3e-10, "ENDCOND_SIMTIMELIM" : 1,
            "ENABLE_ORBIT_FOLLOWING" : 1, "ENABLE_COULOMB_COLLISIONS" : 1
        })
        optgcf = copy.deepcopy(optgo)
        optgcf.update({
            "SIM_MODE" : 2
        })
        optgca = copy.deepcopy(optgcf)
        optgca.update({
            "ENABLE_ADAPTIVE" : 1, "ADAPTIVE_MAX_DRHO" : 0.1,
            "ADAPTIVE_TOL_ORBIT" : 1e-8, "ADAPTIVE_TOL_CCOL" : 1e-1,
            "ADAPTIVE_MAX_DPHI" : 10, "FIXEDSTEP_USERDEFINED" : 1e-10
        })

        # Marker input consists of electrons
        mrk = Marker.generate("gc", n=100, species="electron")
        mrk["r"][:]      = 7.2
        mrk["phi"][:]    = 0.0
        mrk["z"][:]      = 0.0
        mrk["energy"][:] = 1e3
        mrk["pitch"][:]  = 1.0 - 2.0 * np.random.rand(mrk["n"],)
        mrk["zeta"][:]   = 2 * np.pi * np.random.rand(mrk["n"],)

        # Plasma consisting of protons only to avoid e-e collisions
        pls = init("plasma_flat", temperature=1e3, dryrun=True)
        pls["edensity"][:] = 1

        for i in range(ni.size):

            # Adjust simulation time and time step as the density changes
            # (otherwise simulations with low density would take very long)
            simtime = np.maximum( 5e-4, 4e-2 / ( ni[i-1] / ni[0] ) )
            optgo.update({
                "ENDCOND_LIM_SIMTIME" : simtime,
                "FIXEDSTEP_USERDEFINED" :
                np.minimum( 2e-9, 3e-10 / ( ni[i-1] / ni[-1] ) )
            })
            optgcf.update({
                "ENDCOND_LIM_SIMTIME" : simtime,
                "FIXEDSTEP_USERDEFINED" :
                np.minimum( 2e-8, 5e-10 / ( ni[i-1] / ni[-1] ) )
            })
            optgca.update({
                "ENDCOND_LIM_SIMTIME" : simtime
            })
            init("opt", **optgo,  desc=PhysTest.tag_neoclassical_go  + str(i))
            init("opt", **optgcf, desc=PhysTest.tag_neoclassical_gcf + str(i))
            init("opt", **optgca, desc=PhysTest.tag_neoclassical_gca + str(i))

            pls["idensity"][:] = ni[i]
            for tag in [PhysTest.tag_neoclassical_go,
                        PhysTest.tag_neoclassical_gcf,
                        PhysTest.tag_neoclassical_gca]:
                init("gc", **mrk, desc=tag + str(i))
                init("plasma_1D", **pls, desc=tag + str(i))
                init("bfield_analytical_iter_circular", desc=tag + str(i))

    def run_neoclassical(self):
        """Run neoclassical transport test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_neoclassical_go):
            warnings.warn("Results already present: Test neoclass. transport")
            return
        for tag in [PhysTest.tag_neoclassical_go, PhysTest.tag_neoclassical_gcf,
                    PhysTest.tag_neoclassical_gca]:
            i = 0
            while tag + str(i) in self.ascot.data.bfield.ls(show=False):
                print(i)
                self._activateinputs(tag+str(i))
                self._runascot(tag+str(i))
                i += 1

    def check_neoclassical(self):
        """Check neoclassical transport test.
        """
        nscan = 0
        items = self.ascot.data.bfield.ls(show=False)
        while PhysTest.tag_neoclassical_go + str(nscan) in items:
            nscan += 1
            print(nscan)

        # Evaluate Ti and R_omp(rho) as these are needed
        run_go = self.ascot.data[PhysTest.tag_neoclassical_go  + "0"]
        self.ascot.input_init(run=run_go.get_qid(), bfield=True, plasma=True)
        Ti         = self.ascot.input_eval(6.2, 0, 0, 0, "ti1")
        rhoomp     = np.linspace(0, 1, 100)
        romp, zomp = self.ascot.input_rhotheta2rz(rhoomp, 0, 0, 0)
        self.ascot.input_free()

        # Initialize plots
        fig = a5plt.figuredoublecolumn(3/2)
        ax = fig.add_subplot(1,1,1)
        ax.set_xscale("log")
        ax.set_yscale("log")

        # Numerical values
        ni    = np.zeros((nscan,)) / unyt.m**3
        Dgo   = np.zeros((nscan,)) * unyt.m**2 / unyt.s
        Dgcf  = np.zeros((nscan,)) * unyt.m**2 / unyt.s
        Dgca  = np.zeros((nscan,)) * unyt.m**2 / unyt.s
        for i in range(nscan):
            run_go  = self.ascot.data[PhysTest.tag_neoclassical_go  + str(i)]
            run_gcf = self.ascot.data[PhysTest.tag_neoclassical_gcf + str(i)]
            run_gca = self.ascot.data[PhysTest.tag_neoclassical_gca + str(i)]

            ri, ti = run_go.getstate("rho", "mileage", state="ini")
            rf, tf = run_go.getstate("rho", "mileage", state="end")
            ri = np.interp(ri, rhoomp, romp) * unyt.m
            rf = np.interp(rf, rhoomp, romp) * unyt.m
            Dgo[i] = 0.5 * np.mean( (rf - ri)**2 / (tf - ti) )

            ri, ti = run_gcf.getstate("rho", "mileage", state="ini")
            rf, tf = run_gcf.getstate("rho", "mileage", state="end")
            ri = np.interp(ri, rhoomp, romp) * unyt.m
            rf = np.interp(rf, rhoomp, romp) * unyt.m
            Dgcf[i] = 0.5 * np.mean( (rf - ri)**2 / (tf - ti) )

            ri, ti = run_gca.getstate("rho", "mileage", state="ini")
            rf, tf = run_gca.getstate("rho", "mileage", state="end")
            ri = np.interp(ri, rhoomp, romp) * unyt.m
            rf = np.interp(rf, rhoomp, romp) * unyt.m
            Dgca[i] = 0.5 * np.mean( (rf - ri)**2 / (tf - ti) )

            ni[i] = run_go.plasma.read()["idensity"][0, 0]

        r0    = run_go.getstate("r")[0]
        axisr = run_go.bfield.read()["raxis"][0] * unyt.m
        eps   = (r0 - axisr) / axisr
        qfac  = 1.7 # Safety factor at r0 was verified numerically
        bnorm = run_go.bfield.read()["bphi0"][0] * unyt.T

        ekin   = run_go.getstate("ekin")[0]
        omegat = physlib.bouncefrequency(unyt.me, ekin, r0, axisr, qfac)
        rhog   = physlib.gyrolength(unyt.me, 1*unyt.e, ekin, 0.0, bnorm).to("m")

        clog     = 15
        density  = np.power( 10, np.linspace(np.log10(ni[0]) - 1,
                                             np.log10(ni[-1]) + 1, 50) ) / unyt.m**3
        collfreq = physlib.collfreq_ie(unyt.mp, unyt.e, density, Ti, clog) \
            * ( unyt.mp / unyt.me )
        veff = collfreq / omegat
        # Add intermediate values needed for plotting a continuous curve
        veff = np.append(veff, [1, np.power(eps, 3.0/2.0)])
        veff.sort()

        Dps = qfac**2 * veff * omegat * rhog**2 / 2
        Dp  = 0.5 * qfac**2 * omegat * rhog**2 * np.ones(veff.shape)
        Db  = np.power(eps, -3.0/2.0) * Dps

        # x coordinate for plotting the numerical coefficients
        collfreq = physlib.collfreq_ie(unyt.mp, unyt.e, ni, Ti, clog) \
            * ( unyt.mp / unyt.me )
        veff_x = collfreq / omegat

        ax.plot(veff, Dps, color="black")
        ax.plot(veff, Dp, color="black")
        ax.plot(veff, Db, color="black")
        ax.scatter(veff_x, Dgo, marker="*")
        ax.scatter(veff_x, Dgcf, marker="o")
        ax.scatter(veff_x, Dgca, marker="^")
        print(Dgo)
        print(Dgcf)
        print(Dgca)
        plt.show()

    def init_boozer(self):
        """Initialize data for the Boozer transformation test.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_boozer):
            warnings.warn("Inputs already present: Test Boozer transformation")
            return
        init = self.ascot.data.create_input

        # Options
        opt = Opt.get_default()
        opt.update({
            "SIM_MODE" : 4, "ENDCOND_SIMTIMELIM" : 1,
            "ENDCOND_MAX_MILEAGE" : 1e2/3e8, "ENABLE_ORBIT_FOLLOWING" : 1,
            "ENABLE_ORBITWRITE" : 1, "ORBITWRITE_MODE" : 1,
            "ORBITWRITE_INTERVAL" : 1e-1/3e8, "ORBITWRITE_NPOINT" : 10002
        })

        # Use field line markers
        mrk = Marker.generate("fl", n=1)
        mrk["r"][:]      = 8.0
        mrk["phi"][:]    = 0
        mrk["z"][:]      = 0
        mrk["pitch"][:]  = 1.0

        # Magnetic field is just some tokamak with Boozer data
        b = init("bfield_analytical_iter_circular", dryrun=True)
        b.update({"rmin" : 4, "rmax" : 8.5, "nr" : 120, "zmin" : -4,
                  "zmax" : 4, "nz" : 200})

        bphi = [5.3, -5.3, 5.3, -5.3]
        psimult = [200, 200, -200, -200]
        for i in range(4):
            b.update({"bphi0" : bphi[i], "psimult" : psimult[i]})
            out = B_2DS.convert_B_GS(**b)
            qid = init("B_2DS", **out, desc=PhysTest.tag_boozer+str(i))
            qid = self.ascot.data.bfield[qid].get_qid()
            self.ascot.input_init(bfield=qid)
            qid = init("boozer_tokamak", desc=PhysTest.tag_boozer+str(i))
            self.ascot.input_free(bfield=True)

            init("fl", **mrk, desc=PhysTest.tag_boozer+str(i))
            init("opt", **opt, desc=PhysTest.tag_boozer+str(i))

        # This perturbation is used only in the post-processing
        mhd = {"nmode" : 1, "nmodes" : np.array([2]), "mmodes" : np.array([3]),
               "amplitude" : np.array([1.0]), "omega" : np.array([1.0]),
               "phase" : np.array([0.0]), "nrho" : 100, "rhomin" : 0.0,
               "rhomax" : 1.0}
        rhogrid = np.linspace(mhd["rhomin"], mhd["rhomax"], mhd["nrho"])
        alpha   = np.exp( -(rhogrid-0.85)**2/0.1 )
        phi     = alpha*0
        mhd["phi"]   = np.tile(phi, (mhd["nmode"],1)).T
        mhd["alpha"] = np.tile(alpha, (mhd["nmode"],1)).T
        init("MHD_STAT", **mhd, desc=PhysTest.tag_boozer+"0")

    def run_boozer(self):
        """Run Boozer transformation test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_boozer):
            warnings.warn("Results already present: Test Boozer transformation")
            return
        for i in range(4):
            self._activateinputs(PhysTest.tag_boozer + str(i))
            self._runascot(PhysTest.tag_boozer + str(i))

    def check_boozer(self):
        """Check Boozer transformation test.
        """
        run1 = self.ascot.data[PhysTest.tag_boozer+"0"]
        run2 = self.ascot.data[PhysTest.tag_boozer+"1"]
        run3 = self.ascot.data[PhysTest.tag_boozer+"2"]
        run4 = self.ascot.data[PhysTest.tag_boozer+"3"]

        # Initialize plots
        fig = a5plt.figuredoublecolumn(3/2)
        ax1 = fig.add_subplot(3,2,1)
        ax2 = fig.add_subplot(3,2,3)
        ax3 = fig.add_subplot(3,2,5)
        ax4 = fig.add_subplot(3,2,2)
        ax5 = fig.add_subplot(3,2,4)

        ax3.set_xlabel("Boozer theta [rad]")
        ax1.set_ylabel("Error in Bpol [T]")
        ax2.set_ylabel("Error in Bphi [T]")
        ax3.set_ylabel("Jacobian")
        ax4.set_ylabel("q")
        ax5.set_ylabel("q")
        ax5.set_xlabel("Orbit parameter s")

        colors = ["C0", "C1", "C2", "C3"]
        for i, run in enumerate([run1, run2, run3, run4]):
            self.ascot.input_init(run=run.get_qid(), bfield=True, boozer=True,
                                  mhd=True)
            r, phi, z, t, pol = run.getorbit("r", "phi", "z", "time", "theta")
            theta, zeta, alpha, jacb2 = self.ascot.input_eval(
                r, phi, z, t, "theta", "zeta", "alphaeig", "bjacxb2")

            # Results are plotted as a function of poloidal angle so find the
            # "discontinuity" points at 0/2pi.
            idx = np.nonzero(np.abs(np.diff(theta)) > np.pi)[0]
            theta = theta.v
            zeta  = zeta.v

            # Numerical safety factor has some noise so evaluate the average
            dz = np.diff(zeta)
            dz[dz > np.pi]  = dz[dz > np.pi]  - 2*np.pi
            dz[dz < -np.pi] = dz[dz < -np.pi] + 2*np.pi
            dt = np.diff(theta)
            dt[dt > np.pi]  = dt[dt > np.pi]  - 2*np.pi
            dt[dt < -np.pi] = dt[dt < -np.pi] + 2*np.pi
            qfac = dz/dt

            # Evaluate gradients and field components so we can compare those
            a, b, c = self.ascot.input_eval(
                r, phi, z, t, "dpsidr (bzr)", "dpsidphi (bzr)", "dpsidz (bzr)")
            gradpsi = np.array([a, b/r, c]).T
            a, b, c = self.ascot.input_eval(
                r, phi, z, t, "dthetadr", "dthetadphi", "dthetadz")
            gradtheta = np.array([a, b/r, c]).T
            a, b, c = self.ascot.input_eval(
                r, phi, z, t, "dzetadr", "dzetadphi", "dzetadz")
            gradzeta = np.array([a, b/r, c]).T

            br, bphi, bz = self.ascot.input_eval(
                r, phi, z, t, "br", "bphi", "bz")
            self.ascot.input_free()

            # Magnetic field vector from Boozer coordinates
            veca = np.cross(gradpsi, gradtheta)
            vecb = np.cross(gradzeta, gradpsi)
            bvec = -veca*np.mean(qfac) - vecb

            idx = np.nonzero(np.abs(np.diff(theta)) > np.pi)[0]
            dbpol = np.sqrt((bvec[:,0] - br.v)**2 + (bvec[:,2] - bz.v)**2)
            dbphi = bvec[:,1] - bphi.v
            bnorm = br**2 + bz**2 + bphi**2

            j0 = 1
            for j in idx[1:]:
                ax1.plot(theta[j0:j], dbpol[j0:j], color=colors[i])
                ax2.plot(theta[j0:j], dbphi[j0:j], color=colors[i])
                ax3.plot(theta[j0:j], jacb2[j0:j], color=colors[i])
                j0 = j
            if qfac[0] > 0:
                ax4.plot(qfac, color=colors[i])
            else:
                ax5.plot(qfac, color=colors[i])

    def init_mhd(self):
        """Initialize data for the MHD test.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_mhd_go):
            warnings.warn("Inputs already present: Test MHD")
            return
        init = self.ascot.data.create_input

        # Options
        opt = Opt.get_default()
        opt.update({
            "SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
            "FIXEDSTEP_USERDEFINED" : 1e-11, "ENDCOND_SIMTIMELIM" : 1,
            "ENDCOND_LIM_SIMTIME" : 1e-3, "ENABLE_ORBIT_FOLLOWING" : 1,
            "ENABLE_MHD" : 1, "ENABLE_ORBITWRITE" : 1, "ORBITWRITE_MODE" : 1,
            "ORBITWRITE_INTERVAL" : 1e-7, "ORBITWRITE_NPOINT" : 10000
        })
        init("opt", **opt, desc=PhysTest.tag_mhd_go)
        opt.update({
            "SIM_MODE" : 2, "FIXEDSTEP_USERDEFINED" : 1e-10
        })
        init("opt", **opt, desc=PhysTest.tag_mhd_gcf)
        opt.update({
            "ENABLE_ADAPTIVE" : 1, "FIXEDSTEP_USERDEFINED" : 1e-11,
            "ADAPTIVE_TOL_ORBIT" : 1e-10, "ADAPTIVE_MAX_DRHO" : 0.1,
            "ADAPTIVE_MAX_DPHI" : 10
        })
        init("opt", **opt, desc=PhysTest.tag_mhd_gca)

        # Use field line markers
        mrk = Marker.generate("gc", n=2, species="electron")
        mrk["r"][:]      = np.array([7.0, 8.0])
        mrk["phi"][:]    = 0
        mrk["z"][:]      = 0
        mrk["pitch"][:]  = np.array([0.4, 0.9])
        mrk["energy"][:] = 1e6

        # ITER-like field with boozer data and field line markers
        for tag in [PhysTest.tag_mhd_go, PhysTest.tag_mhd_gcf,
                    PhysTest.tag_mhd_gca]:
            qid = init("bfield_analytical_iter_circular", splines=True,
                       desc=tag)
            init("gc", **mrk, desc=tag)

        qid = self.ascot.data.bfield[qid].get_qid()
        self.ascot.input_init(bfield=qid)
        bzr = init("boozer_tokamak", dryrun=True)
        self.ascot.input_free(bfield=True)
        for tag in [PhysTest.tag_mhd_go, PhysTest.tag_mhd_gcf,
                    PhysTest.tag_mhd_gca]:
            init("Boozer", **bzr, desc=tag)

    def run_mhd(self):
        """Run MHD test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_mhd_go):
            warnings.warn("Results already present: Test MHD")
            return
        for tag in [PhysTest.tag_mhd_go, PhysTest.tag_mhd_gcf,
                    PhysTest.tag_mhd_gca]:
            self._activateinputs(tag)
            self._runascot(tag)

    def check_mhd(self):
        """Check MHD test.
        """
        run_go  = self.ascot.data[PhysTest.tag_mhd_go]
        run_gcf = self.ascot.data[PhysTest.tag_mhd_gcf]
        run_gca = self.ascot.data[PhysTest.tag_mhd_gca]

    def init_atomic(self):
        """Initialize data for the atomic reaction test.
        """
        if hasattr(self.ascot.data.options, PhysTest.tag_atomic_cx):
            warnings.warn("Inputs already present: Test atomic reaction")
            return
        init = self.ascot.data.create_input

        # Options
        opt = Opt.get_default()
        opt.update({
            "SIM_MODE" : 1, "FIXEDSTEP_USE_USERDEFINED" : 1,
            "ENDCOND_SIMTIMELIM" : 1, "ENDCOND_LIM_SIMTIME" : 1e-3, # 1e-3
            "ENABLE_ORBIT_FOLLOWING" : 1, "ENABLE_ATOMIC" : 1
        })
        opt.update({
            "FIXEDSTEP_USERDEFINED" : 1e-9, "ENDCOND_NEUTRALIZED" : 1
        })
        init("opt", **opt, desc=PhysTest.tag_atomic_ionz)
        opt.update({
            "FIXEDSTEP_USERDEFINED" : 1e-9, "ENDCOND_NEUTRALIZED" : 0,#1e-13
            "ENDCOND_IONIZED" : 1
        })
        init("opt", **opt, desc=PhysTest.tag_atomic_cx)

        mrk = Marker.generate("gc", n=100, species="deuterium")
        mrk["r"][:]      = 7.0
        mrk["phi"][:]    = 0.0
        mrk["z"][:]      = 0.0
        mrk["zeta"][:]   = 2 * np.pi * np.random.rand(mrk["n"],)
        mrk["energy"][:] = 1e5
        mrk["pitch"][:]  = 1.0 - 2.0 * np.random.rand(mrk["n"],)
        init("gc", **mrk, desc=PhysTest.tag_atomic_cx)

        mrk = Marker.generate("gc", n=100, species="deuterium")
        mrk["charge"][:] = 0
        mrk["r"][:]      = 7.0
        mrk["phi"][:]    = 0.0
        mrk["z"][:]      = 0.0
        mrk["zeta"][:]   = 1.0 - 2.0 * np.random.rand(mrk["n"],)
        mrk["energy"][:] = 1e5
        mrk["pitch"][:]  = 2 * np.pi * np.random.rand(mrk["n"],)
        init("gc", **mrk, desc=PhysTest.tag_atomic_ionz)

        for tag in [PhysTest.tag_atomic_cx, PhysTest.tag_atomic_ionz]:
            # Some tokamak magnetic field
            init("bfield_analytical_iter_circular", desc=tag)

            # Plasma and neutral data
            init("plasma_flat", anum=2, znum=1, mass=2.0135532,
                 density=1e20, temperature=1e3, desc=tag)
            init("neutral_flat", anum=2, znum=1, density=1e16, temperature=1e3,
                 desc=tag)

            # ADAS data
            init("import_adas", desc=tag)

    def run_atomic(self):
        """Run atomic reaction test.
        """
        if hasattr(self.ascot.data, PhysTest.tag_atomic_cx):
            warnings.warn("Results already present: Test atomic reaction")
            return
        for tag in [PhysTest.tag_atomic_cx, PhysTest.tag_atomic_ionz]:
            self._activateinputs(tag)
            self._runascot(tag)

    def check_atomic(self):
        """Check atomic reaction test.
        """
        run_cx   = self.ascot.data[PhysTest.tag_atomic_cx]
        run_ionz = self.ascot.data[PhysTest.tag_atomic_ionz]

    def _activateinputs(self, tag):
        data = self.ascot.data
        tag0 = tag if hasattr(data.bfield, tag) else "DUMMY"
        data.bfield[tag0].activate()
        tag0 = tag if hasattr(data.efield, tag) else "DUMMY"
        data.efield[tag0].activate()
        tag0 = tag if hasattr(data.wall, tag) else "DUMMY"
        data.wall[tag0].activate()
        tag0 = tag if hasattr(data.plasma, tag) else "DUMMY"
        data.plasma[tag0].activate()
        tag0 = tag if hasattr(data.neutral, tag) else "DUMMY"
        data.neutral[tag0].activate()
        tag0 = tag if hasattr(data.boozer, tag) else "DUMMY"
        data.boozer[tag0].activate()
        tag0 = tag if hasattr(data.mhd, tag) else "DUMMY"
        data.mhd[tag0].activate()
        tag0 = tag if hasattr(data.asigma, tag) else "DUMMY"
        data.asigma[tag0].activate()
        tag0 = tag if hasattr(data.options, tag) else "DUMMY"
        data.options[tag0].activate()
        tag0 = tag if hasattr(data.marker, tag) else "DUMMY"
        data.marker[tag0].activate()

    def _runascot(self, test):
        subprocess.call(["./ascot5_main", "--in=testascot.h5", "--d="+test],
                        stdout=subprocess.DEVNULL)
        self.ascot = Ascot(self.ascot.file_getpath())

if __name__ == '__main__':
    test = PhysTest()
    test.execute(init=True, run=True, check=True, tests=["boozer"])
    plt.show()

a5py/a5py/testascot/physicstests.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test mean free paths for atomic reactions.

This file initializes, runs, and plots test cases for checking that the
ASCOT5 atomic reactions module reproduces analytical mean free paths for
charge-exchange (CX) neutralization and ionization in a plasma.

The tests are done using deuterons and deuterium atoms as markers. We
include orbit following to measure the distance of the path travelled
before a reaction occurs.

Similar tests were reported in
P Ollus et al 2022 Plasma Phys. Control. Fusion 64 035014.

NOTE: This script was originally written with test_coulombcollisions.py as
      a template, i.e., by editing a copy of test_coulombcollisions.py.

File: test_atomicreactions.py
"""


import sys
import copy

import numpy                   as np
import unyt
import scipy

import a5py.ascot5io.options   as options
import a5py.ascot5io.B_GS      as B_GS
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_1D     as N0_1D
import a5py.ascot5io.mrk_gc    as mrk_gc
import a5py.ascot5io.mrk_prt   as mrk_prt
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhd

import a5py.testascot.helpers as helpers
import a5py.physlib.gamma as get_gamma

from a5py.preprocessing.analyticequilibrium import psi0 as psifun
from a5py.preprocessing.asigma_loc_read2hdf5 import read2hdf5 as \
    asigma_read2hdf5

from a5py.ascot5io.ascot5 import Ascot
from a5py.physlib import e, m_e, m_p, c, eps_0


simtime = 1e-1

Te  = 1e3
T0  = 1e3
ne  = 1e20
n0  = 1e16
Eion = 1e5
Entl = 1e5

R0 = 6.2
z0 = 0
Bphi0 = 5.3
psi_mult = 200

# ITER-like but circular equilibrium
psi_coeff = np.array([ 8.629e-02,  3.279e-01,  5.268e-01, -2.366e-01,
                       3.825e-01, -3.573e-01, -1.484e-02,  1.506e-01,
                       7.428e-01, -4.447e-01, -1.084e-01,  1.281e-02, -0.155])

m_Dn = scipy.constants.physical_constants["deuteron mass"][0] * unyt.kg
# An attempt at the exact mass of a deuterium atom
m_Dm = 3.34449439e-27 * unyt.kg


def init():
    """
    Initialize tests

    This function initializes two tests: one that simulates the
    CX neutralization of fast deuterons (TEST_CX), and one that simulates the
    ionization of fast deuterium atoms (TEST_IONZ).

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*              Generate options for TEST_CX and TEST_IONZ                 #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    # Shared simulation settings
    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_MAX_SIMTIME"]       = simtime
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 0
    odict["ENABLE_ATOMIC"]             = 1
    odict["ENABLE_DIST_5D"]            = 1

    # Shared distribution configurations
    odict["DIST_MIN_R"]     = 4
    odict["DIST_MAX_R"]     = 10
    odict["DIST_NBIN_R"]    = 1
    odict["DIST_MIN_PHI"]   = 0
    odict["DIST_MAX_PHI"]   = 360
    odict["DIST_NBIN_PHI"]  = 1
    odict["DIST_MIN_Z"]     = -5
    odict["DIST_MAX_Z"]     = 5
    odict["DIST_NBIN_Z"]    = 1
    odict["DIST_NBIN_PPA"]  = 140
    odict["DIST_NBIN_PPE"]  = 80
    odict["DIST_MIN_TIME"]  = 0
    odict["DIST_MAX_TIME"]  = simtime
    odict["DIST_NBIN_TIME"] = 2

    # Distinct settings for TEST_CX
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-9
    odict["ENDCOND_NEUTRALIZED"]       = 1
    odict["DIST_MIN_PPA"]   = -3.2e6 * m_Dn
    odict["DIST_MAX_PPA"]   =  3.2e6 * m_Dn
    odict["DIST_MIN_PPE"]   = 0      * m_Dn
    odict["DIST_MAX_PPE"]   = 3.2e6  * m_Dn
    options.write_hdf5(helpers.testfn, odict, desc="TEST_CX")

    # Distinct settings for TEST_IONZ
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-13
    odict["ENDCOND_NEUTRALIZED"]       = 0
    odict["ENDCOND_IONIZED"]           = 1
    odict["DIST_MIN_PPA"]   = -3.2e6 * m_Dm
    odict["DIST_MAX_PPA"]   =  3.2e6 * m_Dm
    odict["DIST_MIN_PPE"]   = 0      * m_Dm
    odict["DIST_MAX_PPE"]   = 3.2e6  * m_Dm
    options.write_hdf5(helpers.testfn, odict, desc="TEST_IONZ")

    #**************************************************************************#
    #*       Marker input consisting of deuterons and deuterium atoms          #
    #*                                                                         #
    #**************************************************************************#
    Nmrk   = 1000
    ids    = np.linspace(1,Nmrk,Nmrk)
    weight = (1/Nmrk)* np.ones(ids.shape)
    anum   = 2       * np.ones(ids.shape)
    znum   = 1       * np.ones(ids.shape)
    time   = 0       * np.ones(ids.shape)
    phi    = 90      * np.ones(ids.shape)

    mass   = m_Dn.to("amu") * np.ones(ids.shape)
    pol    = 2 * np.pi      * np.random.rand(1, Nmrk)
    R      = 6.2 + 0.8      * np.cos(pol)
    z      = 0.8            * np.sin(pol)
    charge = 1              * np.ones(ids.shape)
    vtot = get_gamma.vnorm_gamma(get_gamma.gamma_energy(m_Dn,Eion*unyt.eV))
    vr     = -vtot/np.sqrt(2) * np.ones(ids.shape)
    vphi   =  vtot/np.sqrt(2) * np.ones(ids.shape)
    vz     = 0                * np.ones(ids.shape)
    mrk_prt.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, vr, vphi, vz,
                       anum, znum, weight, time, desc="TEST_CX")

    # Place neutral markers closely around the axis with velocity vectors
    # pointing towards the inner-midplane edge to give neutrals enough
    # space to ionize before exiting plasma.
    mass   = m_Dm.to("amu")   * np.ones(ids.shape)
    pol    = 2 * np.pi        * np.random.rand(1, Nmrk)
    R      = 6.2 + 0.1        * np.cos(pol)
    z      = 0.1              * np.sin(pol)
    charge = 0                * np.ones(ids.shape)
    vtot = get_gamma.vnorm_gamma(get_gamma.gamma_energy(m_Dm,Entl*unyt.eV))
    vr     = -vtot/np.sqrt(2) * np.ones(ids.shape)
    vphi   =  vtot/np.sqrt(2) * np.ones(ids.shape)
    vz     = 0                * np.ones(ids.shape)
    mrk_prt.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, vr, vphi, vz,
                       anum, znum, weight, time, desc="TEST_IONZ")

    #**************************************************************************#
    #*                 Analytical ITER-like magnetic field                     #
    #*                                                                         #
    #**************************************************************************#
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="TEST_CX")
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="TEST_IONZ")

    #**************************************************************************#
    #*                               Plasma                                    #
    #*                                                                         #
    #**************************************************************************#
    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([2])
    mass   = np.array([m_Dn.to("amu").v[()]])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = ne * np.ones(rho.shape)
    etemp  = Te * np.ones(rho.shape)
    idens  = ne * np.ones((rho.size, Nion))
    itemp  = Te * np.ones(rho.shape)
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="TEST_CX")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="TEST_IONZ")

    #**************************************************************************#
    #*                               Neutrals                                  #
    #*                                                                         #
    #**************************************************************************#
    Nntl   = 1
    znum0  = np.array([1])
    anum0  = np.array([2])
    dens0  = n0 * np.ones((rho.size, Nntl))
    temp0  = T0 * np.ones(rho.shape)
    N0_1D.write_hdf5(helpers.testfn, rho[0], rho[-1], Nrho, Nntl, anum0, znum0,
                     dens0, temp0, 1, desc="TEST_CX")

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    Exyz   = np.array([0, 0, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="TEST_CX")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="TEST_IONZ")

    nwall = 4
    Rwall = np.array([0.1, 100, 100, 0.1])
    zwall = np.array([-100, -100, 100, 100])
    asigma_dir = "../../../../ascot5/asigma/asigma_loc/"
    for tname in ["TEST_CX", "TEST_IONZ"]:
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall, desc=tname)
        boozer.write_hdf5_dummy(helpers.testfn, desc=tname)
        mhd.write_hdf5_dummy(helpers.testfn, desc=tname)
        asigma_read2hdf5(
            fnh5 = helpers.testfn,
            fn_CX_DD0   = asigma_dir + "input.sigmavCX_DD0",
            fn_BMS_H0H  = asigma_dir + "input.BMSsigmav_H0H",
            fn_BMS_H0He = asigma_dir + "input.BMSsigmav_H0He",
            fn_BMS_H0C  = asigma_dir + "input.BMSsigmav_H0C",
            desc = tname)
    N0_1D.write_hdf5_dummy(helpers.testfn, desc="TEST_IONZ")


def run():
    """
    Run tests.
    """
    for test in ["TEST_CX", "TEST_IONZ"]:
        helpers.set_and_run(test)


def check():
    """
    Plot the results of these tests.

    This function makes four plots.
    - One that shows marker equilibrium energy distributions (for all modes) and
      compares them with a Maxwellian distribution.
    - One that shows the pitch distributions in equilibrium.
    - One that shows marker slowing-down energy distributions and compares them
      to the analytical result.
    - One that shows pitch distribution during the slowing-down.
    """

    a5 = Ascot(helpers.testfn)

    #**************************************************************************#
    #* Read, interpolate and evaluate reaction data for analytical calculation #
    #*                                                                         #
    #**************************************************************************#
    asigma_dir = "../../../../ascot5/asigma/asigma_loc/"
    fn_CX_DD0   = asigma_dir + "input.sigmavCX_DD0"
    fn_BMS_H0H  = asigma_dir + "input.BMSsigmav_H0H"

    # Read CX reaction data of the form sigmav(E, T)
    f = open(fn_CX_DD0, "r")
    if f.mode == 'r':
        lines = f.readlines()
        metadata = lines[0].split(' ')
        N_E = int(metadata[6])
        N_T = int(metadata[7])
        abscissa1 = lines[1].split(' ')
        # Convert the 1st abscissa into floats
        E = np.zeros(N_E)
        for i in range(0, N_E):
            E[i] = float(abscissa1[i])
        abscissa2 = lines[2].split(' ')
        # Convert the 2nd abscissa into floats
        T = np.zeros(N_T)
        for i in range(0, N_T):
            T[i] = float(abscissa2[i])
        ordinate = lines[3].split(' ')
        # Move the rate coefficient data into a 2D array
        sigmavCX_DD0 = np.zeros((N_E,N_T))
        for i_E in range(0, N_E):
            for i_T in range(0, N_T):
                sigmavCX_DD0[i_E,i_T] = float(ordinate[i_T*N_E + i_E])
        # Initialize cubic bivariate spline
        sigmavCX_DD0_interp = scipy.interpolate.RectBivariateSpline(
            E, T, sigmavCX_DD0)

    # Read BMS reaction data of the form BMSsigmav(E, n, T)
    # NOTE: Some of the variables from above are reused
    f = open(fn_BMS_H0H, "r")
    if f.mode == 'r':
        lines = f.readlines()
        metadata = lines[0].split(' ')
        N_E = int(metadata[6])
        N_n = int(metadata[7])
        N_T = int(metadata[8])
        abscissa1 = lines[1].split(' ')
        # Convert the 1st abscissa into floats
        E = np.zeros(N_E)
        for i in range(0, N_E):
            E[i] = float(abscissa1[i])
        abscissa2 = lines[2].split(' ')
        # Convert the 2nd abscissa into floats
        n = np.zeros(N_n)
        for i in range(0, N_n):
            # Convert density unit from cm^-3 to m^-3
            n[i] = 1e6*float(abscissa2[i])
        abscissa3 = lines[3].split(' ')
        # Convert the 3rd abscissa into floats
        T = np.zeros(N_T)
        for i in range(0, N_T):
            T[i] = float(abscissa3[i])
        ordinate = lines[4].split(' ')
        # Move the BMS coefficient data into a 3D array
        BMSsigmav_H0H = np.zeros((N_E, N_n, N_T))
        for i_E in range(0, N_E):
            for i_n in range(0,N_n):
                for i_T in range(0,N_T):
                    # Convert BMS coefficient unit from cm^3s^-1 to m^3s^-1
                    BMSsigmav_H0H[i_E, i_n, i_T] = 1e-6*float(ordinate[
                        i_T*N_n*N_E + i_n*N_E + i_E])
        # Python does not initialize cubic trivariate spline function
        # as it does bivariate spline functions, instead it directly
        # evaluates the target values. Hence, the 3D interpolation is
        # done in a single step below.

    #**************************************************************************#
    #*                      Calculate analytical results                       #
    #*                                                                         #
    #**************************************************************************#
    # The analytical mean free path is calculated as
    #   d_MFP = u/nu = u/(<sigma*v>*n),
    # where u is the fast-particle speed , nu the reaction frequency,
    # <sigma*v> the reaction rate coefficients, and n  the reaction
    # counterpart density.
    u_ion = get_gamma.vnorm_gamma(get_gamma.gamma_energy(m_Dn,Eion*unyt.eV))
    u_ntl = get_gamma.vnorm_gamma(get_gamma.gamma_energy(m_Dm,Entl*unyt.eV))
    sigmav = sigmavCX_DD0_interp(Eion, T0)
    d_MFP_anl_CX   = u_ion/(sigmav*n0)
    # Also calculate analytical mean free time
    t_MFT_anl_CX   = u_ion/u_ion/(sigmav*n0)
    # BMS data energy abscissa in eV/amu. Mass-normalize fast-neutral energy.
    Entl_per_amu = Entl/m_Dm.to("amu")
    BMSsigmav = scipy.interpolate.interpn(
        (E, n, T), BMSsigmav_H0H, np.array([Entl_per_amu, ne, Te]), 'cubic')
    d_MFP_anl_IONZ = u_ntl/(BMSsigmav*ne)
    # Also calculate analytical mean free time
    t_MFT_anl_IONZ = u_ntl/u_ntl/(BMSsigmav*ne)

    #**************************************************************************#
    #*                      Calculate simulated results                        #
    #*                                                                         #
    #**************************************************************************#
    ini_CX   = a5.TEST_CX.inistate.read()
    end_CX   = a5.TEST_CX.endstate.read()
    ini_IONZ = a5.TEST_IONZ.inistate.read()
    end_IONZ = a5.TEST_IONZ.endstate.read()

    # Extract termination times, i.e. time taken before a reaction
    t_CX   = end_CX["time"]
    t_IONZ = end_IONZ["time"]

    # Get number of markers for each test
    Nmrk_CX   = np.size(t_CX)
    Nmrk_IONZ = np.size(t_IONZ)

    # Extract speeds and check that they did not change
    # TEST_CX
    v_CX     = np.zeros(Nmrk_CX)
    v_CX_end = np.zeros(Nmrk_CX)
    for i in range(0,Nmrk_CX):
        # Inistate
        pvec = ( unyt.me*unyt.c/(unyt.me.v[()]*unyt.c.v[()])*
                 np.array([[ini_CX["prprt"][i]],
                           [ini_CX["pphiprt"][i]],
                           [ini_CX["pzprt"][i]]]) )
        vvec = get_gamma.velocity_momentum(ini_CX["mass"][i]*unyt.amu,pvec)
        v_CX[i] = unyt.array.unorm(vvec)
        # Endstate
        pvec = ( unyt.me*unyt.c/(unyt.me.v[()]*unyt.c.v[()])*
                 np.array([[end_CX["prprt"][i]],
                           [end_CX["pphiprt"][i]],
                           [end_CX["pzprt"][i]]]) )
        vvec = get_gamma.velocity_momentum(end_CX["mass"][i]*unyt.amu,pvec)
        v_CX_end[i] = unyt.array.unorm(vvec)
    v_CX     = v_CX*unyt.c/(unyt.c.v[()])
    v_CX_end = v_CX_end*unyt.c/(unyt.c.v[()])
    # TEST_IONZ
    v_IONZ     = np.zeros(Nmrk_IONZ)
    v_IONZ_end = np.zeros(Nmrk_IONZ)
    for i in range(0,Nmrk_IONZ):
        # Inistate
        pvec = ( unyt.me*unyt.c/(unyt.me.v[()]*unyt.c.v[()])*
                 np.array([[ini_IONZ["prprt"][i]],
                           [ini_IONZ["pphiprt"][i]],
                           [ini_IONZ["pzprt"][i]]]) )
        vvec = get_gamma.velocity_momentum(ini_IONZ["mass"][i]*unyt.amu,pvec)
        v_IONZ[i] = unyt.array.unorm(vvec)
        # Endstate
        pvec = ( unyt.me*unyt.c/(unyt.me.v[()]*unyt.c.v[()])*
                 np.array([[end_IONZ["prprt"][i]],
                           [end_IONZ["pphiprt"][i]],
                           [end_IONZ["pzprt"][i]]]) )
        vvec = get_gamma.velocity_momentum(end_IONZ["mass"][i]*unyt.amu,pvec)
        v_IONZ_end[i] = unyt.array.unorm(vvec)
    v_IONZ     = v_IONZ*unyt.c/(unyt.c.v[()])
    v_IONZ_end = v_IONZ_end*unyt.c/(unyt.c.v[()])
    # Check conservation
    if((np.abs(v_CX_end/v_CX - 1) > 1e-6).any()):
        raise(Exception("ERROR, TEST_CX: Marker speed(s) changed. Exiting."))
    if((np.abs(v_IONZ_end/v_IONZ - 1) > 1e-6).any()):
        raise(Exception("ERROR, TEST_IONZ: Marker speed(s) changed. Exiting."))

    # Calculate distances travelled before reaction, i.e. Free Paths, and
    # their mean, i.e. Mean Free Path, for each test
    d_FP_sim_CX    = v_CX*t_CX
    d_FP_sim_IONZ  = v_IONZ*t_IONZ
    d_MFP_sim_CX   = np.mean(d_FP_sim_CX)
    d_MFP_sim_IONZ = np.mean(d_FP_sim_IONZ)
    # Also get Free Times and calculate Mean Free Time (MFT) for each test
    t_FT_sim_CX    = v_CX/v_CX*t_CX
    t_FT_sim_IONZ  = v_IONZ/v_IONZ*t_IONZ
    t_MFT_sim_CX   = np.mean(t_FT_sim_CX)
    t_MFT_sim_IONZ = np.mean(t_FT_sim_IONZ)

    # Estimate standard error of the mean
    err_d_MFP_sim_CX   = np.sqrt(np.sum((d_FP_sim_CX-d_MFP_sim_CX)**2))/Nmrk_CX
    err_d_MFP_sim_IONZ = np.sqrt(
        np.sum((d_FP_sim_IONZ-d_MFP_sim_IONZ)**2))/Nmrk_IONZ
    # Also for the times
    err_t_MFT_sim_CX   = np.sqrt(np.sum((t_FT_sim_CX-t_MFT_sim_CX)**2))/Nmrk_CX
    err_t_MFT_sim_IONZ = np.sqrt(
        np.sum((t_FT_sim_IONZ-t_MFT_sim_IONZ)**2))/Nmrk_IONZ

    #**************************************************************************#
    #*                                Print                                    #
    #*                                                                         #
    #**************************************************************************#
    print(f"Mean free paths (simulated | analytical):\n"
          "  CX neutralization: {:.3f} \u00B1 {:.3f} m | {:.3f} m\n"
          "  ionization:        {:.3f} \u00B1 {:.3f} mm | {:.3f} mm"
          .format(d_MFP_sim_CX.v[()], err_d_MFP_sim_CX.v[()],
                  d_MFP_anl_CX.v[()][0][0],
                  1e3*d_MFP_sim_IONZ.v[()], 1e3*err_d_MFP_sim_IONZ.v[()],
                  1e3*d_MFP_anl_IONZ.v[()][0]))
    print(f"Mean free times (simulated | analytical):\n"
          "  CX neutralization: {:.3f} \u00B1 {:.3f} mus | {:.3f} mus\n"
          "  ionization:        {:.3f} \u00B1 {:.3f} ns | {:.3f} ns\n"
          .format(1e6*t_MFT_sim_CX.v[()], 1e6*err_t_MFT_sim_CX.v[()],
                  1e6*t_MFT_anl_CX.v[()][0][0],
                  1e9*t_MFT_sim_IONZ.v[()], 1e9*err_t_MFT_sim_IONZ.v[()],
                  1e9*t_MFT_anl_IONZ.v[()][0]))


if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if( sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_atomicreactions.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test the evaluation and implementation of the boozer coordinates.

This file initializes, runs, and plots test case for checking that the tools in
a5py correctly construct boozer coordinates, and that these coordinates are
correctly implemented in ASCOT5.

Boozer coordinates have the following properties:

1. Field lines appear as straight lines whose slope equals to the safety factor.

2. The Jacobian has the following dependency h(psi) / B^2.

These properties are tested by implementing an analytical equilibrium,
constructing the boozer coordinates, and evaluating the field line trajectories
with ASCOT5. Furthermore, it is shown that the helical perturbations appear
correctly.

To init, run and check this test, call this script without any arguments. To
do only one of the above, call this script with an argument "init", "run", or
"check".

File: test_orbitfollowing.py
"""
import sys

import numpy                   as np
import scipy.constants         as constants
import matplotlib.pyplot       as plt

import a5py.ascot5io.ascot5    as ascot5
import a5py.ascot5io.orbits    as orbits
import a5py.ascot5io.options   as options
import a5py.ascot5io.B_GS      as B_GS
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_3D     as N0_3D
import a5py.ascot5io.mrk_fl    as mrk
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhdmod
import a5py.ascot5io.asigma_loc as asigma_loc

import a5py.testascot.helpers as helpers

from a5py.ascotpy.ascotpy import Ascotpy
from a5py.preprocessing.analyticequilibrium import psi0 as psifun
from a5py.preprocessing.generateboozer import generate as genbooz

from matplotlib.gridspec import GridSpec

psi_mult  = 200
R0        = 6.2
z0        = 0
Bphi0     = 5.3

# ITER-like but circular equilibrium
psi_coeff = np.array([ 2.218e-02, -1.288e-01, -4.177e-02, -6.227e-02,
                       6.200e-03, -1.205e-03, -3.701e-05,  0,
                       0,          0,          0,          0,         -0.155])

# MHd modes
nmode = 2
mmode = 3

# boozer Rz grid
rgrid = np.linspace(4, 8.5, 200)
zgrid = np.linspace(-2.5, 2.5, 200)

def init():
    """
    Initialize tests

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*             Generate options for field line tracing                     #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 4
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_MAX_MILEAGE"]       = 1e3/3e8
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-1/3e8
    odict["ORBITWRITE_NPOINT"]         = 10002

    options.write_hdf5(helpers.testfn, odict, desc="BOOZER")


    #**************************************************************************#
    #*                            Field line markers                           #
    #*                                                                         #
    #**************************************************************************#
    Nmrk   = 2
    ids    = np.array([1, 2])
    R      = np.array([7.0, 8.0])
    weight = 1       * np.array([1, 1])
    pitch  = 1       * np.array([1, 1])
    time   = 0       * np.array([1, 1])
    phi    = 0       * np.array([1, 1])
    z      = 0       * np.array([1, 1])
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, R, phi, z, pitch, weight, time,
                   desc="BOOZER")

    #**************************************************************************#
    #*                     Construct ITER-like magnetic field                  #
    #*                                                                         #
    #**************************************************************************#
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="BOOZER")

    #**************************************************************************#
    #*                     Construct the boozer input                          #
    #*                                                                         #
    #**************************************************************************#
    h5 = ascot5.Ascot(helpers.testfn)
    a5 = Ascotpy(helpers.testfn)
    a5.init(bfield=h5.bfield["BOOZER"].get_qid())
    booz, rhoprof, qprof,_,_ = genbooz(a5,
                                   rgrid=rgrid,
                                   zgrid=zgrid,
                                   npsi=100,
                                   nthgeo=200,
                                   nthbzr=200,
                                   raxis=h5.bfield["BOOZER"].read()["raxis"],
                                   zaxis=h5.bfield["BOOZER"].read()["zaxis"],
                                   psi0=h5.bfield["BOOZER"].read()["psi0"],
                                   psi1=h5.bfield["BOOZER"].read()["psi1"])
    a5.free(bfield=True)
    boozer.write_hdf5(helpers.testfn, desc="BOOZER", **booz)

    # Get qprofile value easily by uncommenting this line
    #print(np.interp(np.array([0.21, 0.85]), rhoprof.ravel(), qprof.ravel()))

    #**************************************************************************#
    #* Helical perturbation which is not used in the simulation but is used    #
    #* during the postprocessing                                               #
    #*                                                                         #
    #**************************************************************************#
    mhd = {}
    mhd["nmode"]     = 1
    mhd["nmodes"]    = np.array([nmode])
    mhd["mmodes"]    = np.array([mmode])
    mhd["amplitude"] = np.array([1])
    mhd["omega"]     = np.array([1])
    mhd["phase"]     = np.array([0])

    mhd["nrho"]   = 100
    mhd["rhomin"] = 0.0
    mhd["rhomax"] = 1

    rhogrid = np.linspace(mhd["rhomin"], mhd["rhomax"], mhd["nrho"])
    alpha   = np.exp( -(rhogrid-0.85)**2/0.1 )
    phi     = alpha*0
    mhd["phi"]   = np.tile(phi, (mhd["nmode"],1)).T
    mhd["alpha"] = np.tile(alpha, (mhd["nmode"],1)).T
    mhdmod.write_hdf5(helpers.testfn, desc="BOOZER", **mhd)

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    Exyz   = np.array([0, 0, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="BOOZER")

    nwall = 4
    Rwall = np.array([0.1, 100, 100, 0.1])
    zwall = np.array([-100, -100, 100, 100])
    W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall, desc="BOOZER")
    N0_3D.write_hdf5_dummy(helpers.testfn, desc="BOOZER")
    asigma_loc.write_hdf5_empty(helpers.testfn, desc="BOOZER")

    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([1])
    mass   = np.array([1])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = 1e20 * np.ones(rho.shape)
    etemp  = 1e3  * np.ones(rho.shape)
    idens  = 1e20 * np.ones((rho.size, Nion))
    itemp  = 1e3  * np.ones(rho.shape)
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="BOOZER")


def run():
    """
    Run tests.
    """
    for test in ["BOOZER"]:
        helpers.set_and_run(test)

def check():
    """
    Plot the results of these tests.

    This function makes five plots.
    - One for both markers showing that indeed they trace straight field lines
      in Boozer coordinates. On top of these a line with safety factor as slope
      is plotted.
    - One for both markers showing that the MHD perturbation along the field has
      a sinusoidal shape with correct amplitude.
    - Final plot shows the product Jacobian * B^2 which should be a flux surface
      function. To verify this, rho contours are also plotted.
    """
    h5 = ascot5.Ascot(helpers.testfn)
    orb = h5["BOOZER"].orbit

    a5 = Ascotpy(helpers.testfn)
    a5.init(bfield=h5.bfield["BOOZER"].get_qid(),
            boozer=h5.boozer["BOOZER"].get_qid(),
            mhd=h5.mhd["BOOZER"].get_qid())
    theta = a5.evaluate(orb["r"], phi=orb["phi"].to("rad"), z=orb["z"],
                        t=orb["time"], quantity="theta")
    zeta  = a5.evaluate(orb["r"], phi=orb["phi"].to("rad"), z=orb["z"],
                        t=orb["time"], quantity="zeta")
    alpha = a5.evaluate(orb["r"], phi=orb["phi"].to("rad"), z=orb["z"],
                        t=orb["time"]*0, quantity="alpha")
    jacB2 = a5.evaluate(rgrid, phi=0, z=zgrid, t=0, grid=True,
                        quantity="jacobianb2")
    #B =  a5.evaluate(rgrid, phi=0, z=zgrid, t=0, grid=True, quantity="bnorm")
    rhovals = a5.evaluate(rgrid, phi=0, z=zgrid, t=0, grid=True, quantity="rho")
    a5.free(bfield=True, boozer=True, mhd=True)

    ids = orb["id"]

    fig = plt.figure(figsize=(11.9/2.54, 8/2.54), constrained_layout=True)
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    gs = GridSpec(ncols=3, nrows=2, figure=fig)

    # Plot field lines for the first point #
    th = np.linspace(0, 2*np.pi*3, 5000)
    s1 = fig.add_subplot(gs[0,0])
    s1.plot(theta[ids==1][:160], zeta[ids==1][:160],
            linestyle="none", marker=".")
    line = 1.67*th
    idx = np.argwhere(np.logical_or(
        np.diff(np.mod(th, 2*np.pi)) < 0,
        np.diff(np.mod(line, 2*np.pi)) < 0
    ) )
    i0 = 0
    for i in range(idx.size):
        i1 = int(idx[i])
        s1.plot(np.mod(th, 2*np.pi)[i0:i1],
                np.mod(line, 2*np.pi)[i0:i1], color="black")
        i0 = i1+1

    s1.set_xlim(0, 2*np.pi)
    s1.set_ylim(0, 2*np.pi)
    s1.set_xticks([0, 2*np.pi])
    s1.set_xticklabels(["", ""])
    s1.set_yticks([0, 2*np.pi])
    s1.set_yticklabels(["0", r"$2\pi$"])
    s1.set_ylabel("Boozer phi [rad]")
    s1.set_title("Field lines")

    # Plot field lines for the second point #
    th = np.linspace(0, 2*np.pi*5, 5000)
    s2 = fig.add_subplot(gs[1,0])
    s2.plot(theta[ids==2][:350], zeta[ids==2][:350],
            linestyle="none", marker=".")
    line = 2.198*th
    idx = np.argwhere(np.logical_or(
        np.diff(np.mod(th, 2*np.pi)) < 0,
        np.diff(np.mod(line, 2*np.pi)) < 0
    ) )
    i0 = 0
    for i in range(idx.size):
        i1 = int(idx[i])
        s2.plot(np.mod(th, 2*np.pi)[i0:i1],
                np.mod(line, 2*np.pi)[i0:i1], color="black")
        i0 = i1+1

    s2.set_xlim(0, 2*np.pi)
    s2.set_ylim(0, 2*np.pi)
    s2.set_xticks([0, 2*np.pi])
    s2.set_xticklabels(["0", r"$2\pi$"])
    s2.set_yticks([0, 2*np.pi])
    s2.set_yticklabels(["0", r"$2\pi$"])
    s2.set_xlabel("Boozer theta [rad]")
    s2.set_ylabel("Boozer phi [rad]")

    # Plot perturbation for the first point #
    s3 = fig.add_subplot(gs[0,1])
    s3.plot(np.mod(nmode*zeta[ids==1] - mmode*theta[ids==1], 2*np.pi),
            alpha[ids==1], linestyle="none", marker=".")

    rho        = orb["rho"][ids==1][0]
    amplitude  = np.exp( -(rho-0.85)**2/0.1 )
    th         = np.linspace(0, 2*np.pi, 100)
    mode       = amplitude * np.cos(th)
    s3.plot(th, mode, color="black")

    s3.set_xlim(0, 2*np.pi)
    s3.set_xticks([0, 2*np.pi])
    s3.set_xticklabels(["", ""])
    s3.set_title("Perturbation")

    # Plot perturbation for the second point #
    s4 = fig.add_subplot(gs[1,1])
    s4.plot(np.mod(nmode*zeta[ids==2] - mmode*theta[ids==2], 2*np.pi),
            alpha[ids==2], linestyle="none", marker=".")

    rho        = orb["rho"][ids==2][0]
    amplitude  = np.exp( -(rho-0.85)**2/0.1 )
    th         = np.linspace(0, 2*np.pi, 100)
    mode       = amplitude * np.cos(th)
    s4.plot(th, mode, color="black")

    s4.set_xlim(0, 2*np.pi)
    s4.set_xticks([0, 2*np.pi])
    s4.set_xticklabels(["0", r"$2\pi$"])
    s4.set_xlabel("Phase [rad]")
    s4.set_ylabel("Amplitude [A.U.]")

    # Plot Jacobian #
    jacB2   = np.squeeze(jacB2)
    rhovals = np.squeeze(rhovals)
    s5 = fig.add_subplot(gs[0:2,2])
    s5.contour(rgrid, zgrid, jacB2.T, colors="C0" )
    s5.contour(rgrid, zgrid, rhovals.T, colors="black" )
    s5.set_xlabel(r"$R$ [m]")
    s5.set_ylabel(r"$z$ [m]")
    s5.set_title("Jacobian x B")

    plt.savefig("test_boozercoordinates.png", dpi=300)
    plt.show()


if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if( sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_boozercoordinates.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test collisional transport in uniform magnetic field.

This file initializes, runs, and plots test case for checking that ASCOT5
reproduces collsional transport in uniform magnetic field, i.e., the
so-called classical transport, correctly. This transport is diffusive
with the diffusion coefficient being inversely proportional to magnetic
field strength squared.

Tests are done in a Cartesian (B_TC) magnetic field using protons as markers.

To init, run and check this test, call this script without any arguments. To
do only one of the above, call this script with an argument "init", "run", or
"check".

File: test_classicaltransport.py
"""

import sys

import numpy                   as np
import matplotlib.pyplot       as plt
import unyt

import a5py.ascot5io.options   as options
import a5py.ascot5io.B_TC      as B_TC
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_3D     as N0_3D
import a5py.ascot5io.mrk_gc    as mrk
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhd
import a5py.ascot5io.asigma_loc as asigma_loc

from a5py.ascot5io.ascot5 import Ascot
from a5py.physlib import e, m_e, m_p, c, eps_0

import a5py.testascot.helpers as helpers


Te   = 1e3
ne   = 1e22
Bmin = 1
Bmax = 10
nB   = 6

# Number of markers. Adjust to change the time it takes to simulate tests.
Nmrk = 200

def init():
    """
    Initialize tests

    This function initializes parameter scan for three test cases:
    - CLASS_GO tests gyro-orbit scheme
    - CLASS_GCF tests guiding center fixed-scheme
    - CLASS_GCA tests guiding center adaptive scheme

    Each test case is initialized nB times; each with an decreasing value for
    the magnetic field strength. The test cases are named with running index on
    their prefix e.g. CLASS_GO1, CLASS_GO2, and so on.

    All cases are run without energy collisions.

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*                     Generate options for CLASS-GO                       #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 5e-6
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1

    for i in range(1, nB+1):
        options.write_hdf5(helpers.testfn, odict, desc="CLASS_GO" + str(i))

    #**************************************************************************#
    #*                     Generate options for CLASS-GCF                      #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-9
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 5e-6
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1

    for i in range(1, nB+1):
        options.write_hdf5(helpers.testfn, odict, desc="CLASS_GCF" + str(i))

    #**************************************************************************#
    #*                     Generate options for CLASS-GCA                      #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["ENABLE_ADAPTIVE"]           = 1
    odict["ADAPTIVE_TOL_ORBIT"]        = 1e-8
    odict["ADAPTIVE_TOL_COL"]          = 1e-1
    odict["ADAPTIVE_MAX_DRHO"]         = 0.1
    odict["ADAPTIVE_MAX_DPHI"]         = 10
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-8
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 5e-6
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1

    for i in range(1, nB+1):
        options.write_hdf5(helpers.testfn, odict, desc="CLASS_GCA" + str(i))

    #**************************************************************************#
    #*                    Marker input consisting of protons                   #
    #*                                                                         #
    #**************************************************************************#
    ids    = np.linspace(1,Nmrk,Nmrk)
    weight = np.ones(ids.shape)
    mass   = m_p.to("amu") * np.ones(ids.shape)
    charge = 1             * np.ones(ids.shape)
    anum   = 1             * np.ones(ids.shape)
    znum   = 1             * np.ones(ids.shape)
    time   = 0             * np.ones(ids.shape)
    R      = 5             * np.ones(ids.shape)
    phi    = 0             * np.ones(ids.shape)
    z      = 0             * np.ones(ids.shape)
    energy = 1e3           * np.ones(ids.shape)
    zeta   = 2 * np.pi * np.random.rand(1,Nmrk)
    pitch  = 1 - 2 * np.random.rand(1,Nmrk)
    for i in range(1, nB+1):
        mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, energy, pitch, zeta,
                       anum, znum, weight, time, desc="CLASS_GO" + str(i))
        mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, energy, pitch, zeta,
                       anum, znum, weight, time, desc="CLASS_GCF" + str(i))
        mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, energy, pitch, zeta,
                       anum, znum, weight, time, desc="CLASS_GCA" + str(i))

    #**************************************************************************#
    #*     Uniform magnetic field with values scanned from Bmin to Bmax        #
    #*                                                                         #
    #**************************************************************************#
    gradB  = np.array([0,0,0,0,0,0,0,0,0])
    rhoval = 1.5
    Bvals = np.linspace( 1 / ( Bmax * Bmax ), 1 / ( Bmin * Bmin ), nB )
    Bvals = 1 / np.sqrt(Bvals)
    for i in range(1, nB+1):
        Bxyz   = np.array([Bvals[i-1], 0, 0])
        B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval,
                        desc="CLASS_GO" + str(i))
        B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval,
                        desc="CLASS_GCF" + str(i))
        B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval,
                        desc="CLASS_GCA" + str(i))

    #**************************************************************************#
    #* Plasma consisting of electrons only (to avoid proton-proton collisions) #
    #*                                                                         #
    #**************************************************************************#
    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([1])
    mass   = np.array([1])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = ne  * np.ones(rho.shape)
    etemp  = Te  * np.ones(rho.shape)
    idens  = 1   * np.ones((rho.size, Nion))
    itemp  = 1e3 * np.ones(rho.shape)
    for i in range(1, nB+1):
        P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge,
                        rho, edens, etemp, idens, itemp,
                        desc="CLASS_GO" + str(i))
        P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge,
                        rho, edens, etemp, idens, itemp,
                        desc="CLASS_GCF" + str(i))
        P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge,
                        rho, edens, etemp, idens, itemp,
                        desc="CLASS_GCA" + str(i))

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    for i in range(1, nB+1):
        Exyz   = np.array([0, 0, 0])
        E_TC.write_hdf5(helpers.testfn, Exyz, desc="CLASS_GO"  + str(i))
        E_TC.write_hdf5(helpers.testfn, Exyz, desc="CLASS_GCF" + str(i))
        E_TC.write_hdf5(helpers.testfn, Exyz, desc="CLASS_GCA" + str(i))

        nwall = 4
        Rwall = np.array([0.1, 100, 100, 0.1])
        zwall = np.array([-100, -100, 100, 100])
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall,
                        desc="CLASS_GO" + str(i))
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall,
                        desc="CLASS_GCF" + str(i))
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall,
                        desc="CLASS_GCA" + str(i))

        N0_3D.write_hdf5_dummy(helpers.testfn, desc="CLASS_GO"  + str(i))
        N0_3D.write_hdf5_dummy(helpers.testfn, desc="CLASS_GCF" + str(i))
        N0_3D.write_hdf5_dummy(helpers.testfn, desc="CLASS_GCA" + str(i))

        boozer.write_hdf5_dummy(helpers.testfn, desc="CLASS_GO"  + str(i))
        boozer.write_hdf5_dummy(helpers.testfn, desc="CLASS_GCF" + str(i))
        boozer.write_hdf5_dummy(helpers.testfn, desc="CLASS_GCA" + str(i))

        mhd.write_hdf5_dummy(helpers.testfn,    desc="CLASS_GO"  + str(i))
        mhd.write_hdf5_dummy(helpers.testfn,    desc="CLASS_GCF" + str(i))
        mhd.write_hdf5_dummy(helpers.testfn,    desc="CLASS_GCA" + str(i))

        asigma_loc.write_hdf5_empty(helpers.testfn, desc="CLASS_GO"  + str(i))
        asigma_loc.write_hdf5_empty(helpers.testfn, desc="CLASS_GCF" + str(i))
        asigma_loc.write_hdf5_empty(helpers.testfn, desc="CLASS_GCA" + str(i))


def run():
    """
    Run tests.
    """
    for test in ["CLASS_GO", "CLASS_GCF", "CLASS_GCA"]:
        for i in range(1, nB+1):
            helpers.set_and_run(test + str(i))


def check():
    """
    Plot the results of these tests.

    This function makes one plot that shows how diffusion coefficient evaluated
    numerically with different modes scales with magnetic field strength, and
    compares that to the classical diffusion coefficient.
    """
    a5 = Ascot(helpers.testfn)

    # Diffusion occurs on a 2D plane
    ndim = 2

    # Evaluate diffusion coefficient according to equation
    # D = var(Delta x) / (2 * ndim * Delta t)
    # Because E[Delta x] = 0, var[Delta x] = E[(Delta x)^2]
    DGO  = np.zeros(nB)
    DGCF = np.zeros(nB)
    DGCA = np.zeros(nB)
    for i in range(1, nB+1):
        inistate = a5["CLASS_GO" + str(i)].inistate
        yi = inistate["y"]
        zi = inistate["z"]
        ti = inistate["time"]

        endstate = a5["CLASS_GO" + str(i)].endstate
        ye = endstate["y"]
        ze = endstate["z"]
        te = endstate["time"]

        dr2 = ( yi - ye ) * ( yi - ye ) + ( zi - ze ) * ( zi - ze )
        DGO[i-1] = np.mean( dr2 / ( te - ti ) ) / (2*ndim)

        inistate = a5["CLASS_GCF" + str(i)].inistate
        yi = inistate["y"]
        zi = inistate["z"]
        ti = inistate["time"]

        endstate = a5["CLASS_GCF" + str(i)].endstate
        ye = endstate["y"]
        ze = endstate["z"]
        te = endstate["time"]

        dr2 = ( yi - ye ) * ( yi - ye ) + ( zi - ze ) * ( zi - ze )
        DGCF[i-1] = np.mean( dr2 / ( te - ti ) ) / (2*ndim)

        inistate = a5["CLASS_GCA" + str(i)].inistate
        yi = inistate["y"]
        zi = inistate["z"]
        ti = inistate["time"]

        endstate = a5["CLASS_GCA" + str(i)].endstate
        ye = endstate["y"]
        ze = endstate["z"]
        te = endstate["time"]

        dr2 = ( yi - ye ) * ( yi - ye ) + ( zi - ze ) * ( zi - ze )
        DGCA[i-1] = np.mean( dr2 / ( te - ti ) ) / (2*ndim)

    #**************************************************************************#
    #*                  Evaluate the analytical estimate                       #
    #*                                                                         #
    #**************************************************************************#

    clog     = 13.4
    collfreq = (m_e/m_p) * (np.sqrt(2/np.pi)/3) \
               * np.power(e*e/(4*np.pi*eps_0), 2) \
               * (4*np.pi / np.sqrt(m_e*np.power(Te*e,3) ) ) * ne * clog

    xvals    = np.linspace( 1 / ( Bmax * Bmax ), 1 / ( Bmin * Bmin ), nB )
    Bvals    = 1 / np.sqrt(xvals)
    gamma  = 1 + 1e3*unyt.eV / ( m_p * c * c )
    v      = np.sqrt(1.0 - 1.0 / ( gamma * gamma ) ) * c
    rhog   = gamma * m_p * v / (Bvals * e)
    Dclass = collfreq * rhog * rhog / 2

    #**************************************************************************#
    #*                            Plot                                         #
    #*                                                                         #
    #**************************************************************************#

    f = plt.figure(figsize=(11.9/2.54, 5/2.54))
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    h = plt.gca()
    h.set_position([0.15, 0.25, 0.82, 0.7], which='both')

    scaling = 100
    h.plot(xvals, Dclass * scaling, 'black', linewidth=1, linestyle='-',
           label='Analytical')
    h.plot(xvals, DGO * scaling, linestyle='none', marker='*', markersize=11,
           alpha=0.5, label='GO')
    h.plot(xvals, DGCF * scaling, linestyle='none', marker='.', markersize=10,
           alpha=0.5, label='GCF')
    h.plot(xvals, DGCA * scaling, linestyle='none', marker='^', markersize=6,
           alpha=0.5, label='GCA')

    h.tick_params(axis='y', direction='out')
    h.tick_params(axis='x', direction='out')
    h.xaxis.set(ticks=np.linspace(0,1,11),
                ticklabels=[0, '', '', '', '', 5, '', '', '', '', 10])
    h.yaxis.set(ticks=np.linspace(0,8,9),
                ticklabels=[0, '', 2, '',  4, '', 6, '', 8])
    h.set(xlabel="$1/B^2$ [T$^{-2}$]", ylabel="$D$ [m$^2$/s]")

    h.legend(loc='upper left', frameon=False, fontsize=10, numpoints=1)

    plt.savefig("test_classicaltransport.png", dpi=300)
    plt.show()


if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if(   sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_classicaltransport.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test collisional slowing-down and equilibrium.

This file initializes, runs, and plots test case for checking that ASCOT5
collision operator produces correct distribution for fast-particle slowing-down
and thermal equilibrium.

The tests are done using protons as markers. Even though these tests only
require collision operator, we also include orbit following to be sure that
this does not bias the results.

File: test_coulombcollisions.py
"""

import sys
import copy

import numpy                    as np
import unyt
import scipy
import matplotlib.pyplot        as plt
import matplotlib.lines         as mlines

import a5py.ascot5io.options    as options
import a5py.ascot5io.B_GS       as B_GS
import a5py.ascot5io.E_TC       as E_TC
import a5py.ascot5io.plasma_1D  as P_1D
import a5py.ascot5io.wall_2D    as W_2D
import a5py.ascot5io.N0_3D      as N0_3D
import a5py.ascot5io.mrk_gc     as mrk
import a5py.ascot5io.boozer     as boozer
import a5py.ascot5io.mhd        as mhd
import a5py.ascot5io.asigma_loc as asigma_loc

import a5py.testascot.helpers as helpers

from a5py.preprocessing.analyticequilibrium import psi0 as psifun

from a5py.ascot5io.ascot5 import Ascot
from a5py.physlib import e, m_e, m_p, m_a, c, eps_0

simtime_th = 2e-2

Te  = 1e3
ne  = 1e20
Eth = 1e3
Esd = 3.5e6

R0 = 6.2
z0 = 0
Bphi0 = 5.3
psi_mult = 200

# ITER-like but circular equilibrium
psi_coeff = np.array([ 8.629e-02,  3.279e-01,  5.268e-01, -2.366e-01,
                       3.825e-01, -3.573e-01, -1.484e-02,  1.506e-01,
                       7.428e-01, -4.447e-01, -1.084e-01,  1.281e-02, -0.155])

def init():
    """
    Initialize tests

    This function initializes three tests that simulate a thermal equilibrium
    in different modes: THERMAL_GO, THERMAL_GCF, and THERMAL_GCA. Three
    additional simulations are initialized that simulate alpha particle
    slowing-down in different modes: SLOWING_GO, SLOWING_GCF, and SLOWING_GCA.

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*       Generate options for THERMAL_GO, THERMAL_GCF, and THERMAL_GCA     #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = simtime_th
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1
    odict["ENABLE_DIST_5D"]            = 1

    odict["DIST_MIN_R"]     = 4
    odict["DIST_MAX_R"]     = 10
    odict["DIST_NBIN_R"]    = 1
    odict["DIST_MIN_PHI"]   = 0
    odict["DIST_MAX_PHI"]   = 360
    odict["DIST_NBIN_PHI"]  = 1
    odict["DIST_MIN_Z"]     = -5
    odict["DIST_MAX_Z"]     = 5
    odict["DIST_NBIN_Z"]    = 1
    odict["DIST_MIN_PPA"]   = -1.5e6 * m_p
    odict["DIST_MAX_PPA"]   =  1.5e6 * m_p
    odict["DIST_NBIN_PPA"]  = 140
    odict["DIST_MIN_PPE"]   = 0      * m_p
    odict["DIST_MAX_PPE"]   = 1.5e6  * m_p
    odict["DIST_NBIN_PPE"]  = 80
    odict["DIST_MIN_TIME"]  = 0
    odict["DIST_MAX_TIME"]  = simtime_th
    odict["DIST_NBIN_TIME"] = 2

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-8
    options.write_hdf5(helpers.testfn, odict, desc="THERMAL_GO")

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 2e-8
    options.write_hdf5(helpers.testfn, odict, desc="THERMAL_GCF")

    odict["SIM_MODE"]                  = 2
    odict["ENABLE_ADAPTIVE"]           = 1
    odict["ADAPTIVE_TOL_ORBIT"]        = 1e-6
    odict["ADAPTIVE_TOL_COL"]          = 1e-2
    odict["ADAPTIVE_MAX_DRHO"]         = 0.1
    odict["ADAPTIVE_MAX_DPHI"]         = 10
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-8
    options.write_hdf5(helpers.testfn, odict, desc="THERMAL_GCA")


    #**************************************************************************#
    #*  Generate options for SLOWING_GO, SLOWING_GCF, and SLOWING_GCA          #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["ENDCOND_ENERGYLIM"]         = 1
    odict["ENDCOND_MIN_ENERGY"]        = 50e3
    odict["ENDCOND_MIN_THERMAL"]       = 0
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1
    odict["ENABLE_DIST_5D"]            = 1

    odict["DIST_MIN_R"]    = 4
    odict["DIST_MAX_R"]    = 10
    odict["DIST_NBIN_R"]   = 1
    odict["DIST_MIN_PHI"]  = 0
    odict["DIST_MAX_PHI"]  = 360
    odict["DIST_NBIN_PHI"] = 1
    odict["DIST_MIN_Z"]    = -5
    odict["DIST_MAX_Z"]    = 5
    odict["DIST_NBIN_Z"]   = 1
    odict["DIST_MIN_PPA"]  = -2e7 * m_a
    odict["DIST_MAX_PPA"]  =  2e7 * m_a
    odict["DIST_NBIN_PPA"] = 200
    odict["DIST_MIN_PPE"]  = 0    * m_a
    odict["DIST_MAX_PPE"]  = 2e7  * m_a
    odict["DIST_NBIN_PPE"] = 100

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 2e-9
    options.write_hdf5(helpers.testfn, odict, desc="SLOWING_GO")

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 3e-8
    options.write_hdf5(helpers.testfn, odict, desc="SLOWING_GCF")

    odict["SIM_MODE"]                  = 2
    odict["ENABLE_ADAPTIVE"]           = 1
    odict["ADAPTIVE_TOL_ORBIT"]        = 1e-6
    odict["ADAPTIVE_TOL_COL"]          = 1e-2
    odict["ADAPTIVE_MAX_DRHO"]         = 0.1
    odict["ADAPTIVE_MAX_DPHI"]         = 10
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-8
    options.write_hdf5(helpers.testfn, odict, desc="SLOWING_GCA")

    #**************************************************************************#
    #*            Marker input consisting of protons and alfas                 #
    #*                                                                         #
    #**************************************************************************#
    Nmrk   = 20
    ids    = np.linspace(1,Nmrk,Nmrk)
    weight = (1/Nmrk)* np.ones(ids.shape)
    mass   = m_p.to("amu") * np.ones(ids.shape)
    charge = 1       * np.ones(ids.shape)
    anum   = 1       * np.ones(ids.shape)
    znum   = 1       * np.ones(ids.shape)
    time   = 0       * np.ones(ids.shape)
    phi    = 90      * np.ones(ids.shape)
    zeta   = 2 * np.pi * np.random.rand(1,Nmrk)
    energy = Eth     * np.ones(ids.shape)
    pitch  = 0.5     * np.ones(ids.shape)

    pol    = 2*np.pi*np.random.rand(1, Nmrk)
    R      = 6.2 + 0.8 * np.cos(pol)
    z      = 0.8 * np.sin(pol)
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                   charge, R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="THERMAL_GO")

    pol    = 2*np.pi*np.random.rand(1, Nmrk)
    R      = 6.2 + 0.8 * np.cos(pol)
    z      = 0.8 * np.sin(pol)
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                   charge, R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="THERMAL_GCF")

    pol    = 2*np.pi*np.random.rand(1, Nmrk)
    R      = 6.2 + 0.8 * np.cos(pol)
    z      = 0.8 * np.sin(pol)
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                   charge, R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="THERMAL_GCA")

    Nmrk   = 200
    ids    = np.linspace(1,Nmrk,Nmrk)
    weight = (1/Nmrk)* np.ones(ids.shape)
    mass   = m_a.to("amu") * np.ones(ids.shape)
    charge = 2       * np.ones(ids.shape)
    anum   = 4       * np.ones(ids.shape)
    znum   = 2       * np.ones(ids.shape)
    time   = 0       * np.ones(ids.shape)
    phi    = 90      * np.ones(ids.shape)
    zeta   = 2 * np.pi * np.random.rand(1,Nmrk)
    energy = Esd * np.ones(ids.shape)

    pol    = 2*np.pi*np.random.rand(1, Nmrk)
    R      = 6.2 + 0.8 * np.cos(pol)
    z      = 0.8 * np.sin(pol)
    pitch  = 1 - 2 * np.random.rand(1,Nmrk)
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                   charge, R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="SLOWING_GO")

    pol    = 2*np.pi*np.random.rand(1, Nmrk)
    R      = 6.2 + 0.8 * np.cos(pol)
    z      = 0.8 * np.sin(pol)
    pitch  = 1 - 2 * np.random.rand(1,Nmrk)
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                   charge, R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="SLOWING_GCF")

    pol    = 2*np.pi*np.random.rand(1, Nmrk)
    R      = 6.2 + 0.8 * np.cos(pol)
    z      = 0.8 * np.sin(pol)
    pitch  = 1 - 2 * np.random.rand(1,Nmrk)
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                   charge, R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="SLOWING_GCA")

    #**************************************************************************#
    #*                 Analytical ITER-like magnetic field                     #
    #*                                                                         #
    #**************************************************************************#
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="THERMAL_GO")
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="THERMAL_GCF")
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="THERMAL_GCA")
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="SLOWING_GO")
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="SLOWING_GCF")
    B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                    desc="SLOWING_GCA")

    #**************************************************************************#
    #*                               Plasma                                    #
    #*                                                                         #
    #**************************************************************************#
    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([1])
    mass   = np.array([1])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = ne  * np.ones(rho.shape)
    etemp  = Te  * np.ones(rho.shape)
    idens  = ne  * np.ones((rho.size, Nion))
    itemp  = 1e3 * np.ones(rho.shape)
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="THERMAL_GO")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="THERMAL_GCF")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="THERMAL_GCA")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="SLOWING_GO")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="SLOWING_GCF")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="SLOWING_GCA")

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    Exyz   = np.array([0, 0, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="THERMAL_GO")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="THERMAL_GCF")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="THERMAL_GCA")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="SLOWING_GO")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="SLOWING_GCF")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="SLOWING_GCA")

    nwall = 4
    Rwall = np.array([0.1, 100, 100, 0.1])
    zwall = np.array([-100, -100, 100, 100])
    for tname in ["THERMAL_GO", "THERMAL_GCF", "THERMAL_GCA", "SLOWING_GO",
                  "SLOWING_GCF", "SLOWING_GCA"]:
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall, desc=tname)
        N0_3D.write_hdf5_dummy(helpers.testfn, desc=tname)
        boozer.write_hdf5_dummy(helpers.testfn, desc=tname)
        mhd.write_hdf5_dummy(helpers.testfn, desc=tname)
        asigma_loc.write_hdf5_empty(helpers.testfn, desc=tname)


def run():
    """
    Run tests.
    """
    for test in ["THERMAL_GO", "THERMAL_GCF", "THERMAL_GCA", \
                 "SLOWING_GO", "SLOWING_GCF", "SLOWING_GCA"]:
        helpers.set_and_run(test)

def check():
    """
    Plot the results of these tests.

    This function makes four plots.
    - One that shows marker equilibrium energy distributions (for all modes) and
      compares them with a Maxwellian distribution.
    - One that shows the pitch distributions in equilibrium.
    - One that shows marker slowing-down energy distributions and compares them
      to the analytical result.
    - One that shows pitch distribution during the slowing-down.
    """

    a5 = Ascot(helpers.testfn)

    ts_GO  = np.mean(a5.SLOWING_GO.endstate["time"])
    ts_GCF = np.mean(a5.SLOWING_GCF.endstate["time"])
    ts_GCA = np.mean(a5.SLOWING_GCA.endstate["time"])

    #**************************************************************************#
    #*      Initialize histogram grids and evaluate analytical results         #
    #*                                                                         #
    #**************************************************************************#
    THERMAL = {}
    THERMAL["Egrid"] = np.linspace( 0,   4e4, 200)
    THERMAL["xigrid"]  = np.linspace(-1,     1, 50)
    SLOWING = {}
    SLOWING["Egrid"] = np.linspace( 1, 4.0e6, 250)
    SLOWING["xigrid"]  = np.linspace(-1,     1, 20)

    alphaZ = 2
    clog   = 16
    vth    = np.sqrt(2*Te*e / m_e)
    vcrit  = vth * np.power( (3.0*np.sqrt(np.pi)/4.0) * (m_e / m_p) , 1/3.0)
    Ecrit  = 0.5 * m_a * vcrit * vcrit / e
    ts     = 3 * np.sqrt( np.power(2*np.pi * Te * e, 3) / m_e ) * eps_0 * eps_0 \
             * m_a /( alphaZ * alphaZ * np.power(e, 4) * ne * clog)

    heaviside = np.logical_and(SLOWING["Egrid"] <= Esd,
                               SLOWING["Egrid"] >= 50*Te)

    THERMAL["analytical"]  = 2 * np.sqrt(THERMAL["Egrid"]/np.pi) \
                             * np.power(Te,-3.0/2) \
                             * np.exp(-THERMAL["Egrid"]/Te)
    THERMAL["analytical"] *= (simtime_th/2)
    SLOWING["analytical"] = heaviside * ts \
                            / ( ( 1 + np.power(Ecrit/SLOWING["Egrid"], 3.0/2) )\
                                * 2 * SLOWING["Egrid"] )

    # ts is slowing down rate which gives the slowing down time as
    # t_sd = ts*log(v_0 / v_th) = 0.5*ts*log(E_0/E_th)
    slowingdowntime = 0.5*ts*np.log(Esd/(50*Te))

    #**************************************************************************#
    #*            Evaluate thermal distributions in energy and pitch           #
    #*                                                                         #
    #**************************************************************************#

    for mode in ["GO", "GCF", "GCA"]:
        distobj = a5["THERMAL_" + mode]["dist5d"]
        Exidist = distobj.get_E_xi_dist(E_edges=THERMAL["Egrid"],
                                        xi_edges=THERMAL["xigrid"],
                                        r=0, phi=0, z=0, charge=0, time=[1, 2])

        Edist  = Exidist
        xidist = copy.deepcopy(Exidist)
        distobj.get_dist(dist=Edist, pitch=0)
        distobj.get_dist(dist=xidist, energy=0)

        THERMAL[mode + "_Edist"]  = Edist["distribution"]
        THERMAL[mode + "_xidist"] = xidist["distribution"]
        THERMAL["E"]  = Edist["energy"]
        THERMAL["xi"] = xidist["pitch"]

    #**************************************************************************#
    #*        Evaluate slowing-down distributions in energy and pitch          #
    #*                                                                         #
    #**************************************************************************#

    for mode in ["GO", "GCF", "GCA"]:
        distobj = a5["SLOWING_" + mode]["dist5d"]

        Exidist = distobj.get_E_xi_dist(E_edges=SLOWING["Egrid"],
                                        xi_edges=SLOWING["xigrid"],
                                        r=0, phi=0, z=0, charge=0, time=0)

        Edist  = Exidist
        xidist = copy.deepcopy(Exidist)
        distobj.get_dist(dist=Edist,   pitch=0)
        distobj.get_dist(dist=xidist, energy=0)

        SLOWING[mode + "_Edist"]  = Edist["distribution"]
        SLOWING[mode + "_xidist"] = xidist["distribution"]
        SLOWING["E"]  = Edist["energy"]
        SLOWING["xi"] = xidist["pitch"]

    #**************************************************************************#
    #*                                 Plot                                    #
    #*                                                                         #
    #**************************************************************************#

    cy = plt.rcParams['axes.prop_cycle'].by_key()['color']

    f = plt.figure(figsize=(11.9/2.54, 8/2.54))
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    c = ['b', 'g', 'r']
    h1 = f.add_subplot(1,4,1)
    h1.set_position([0.12, 0.67, 0.45, 0.3], which='both')

    h2 = f.add_subplot(1,4,2)
    h2.set_position([0.72, 0.67, 0.25, 0.3], which='both')

    h3 = f.add_subplot(1,4,3)
    h3.set_position([0.12, 0.17, 0.45, 0.3], which='both')

    h4 = f.add_subplot(1,4,4)
    h4.set_position([0.72, 0.17, 0.25, 0.3], which='both')

    h1.plot(THERMAL["Egrid"], THERMAL["analytical"], 'black')
    h1.plot(THERMAL["E"], THERMAL["GO_Edist"],
            marker='*', markersize=8, markevery=3, linestyle='none', alpha=0.5)
    h1.plot(THERMAL["E"], THERMAL["GCF_Edist"],
            marker='.', markersize=10, markevery=3, linestyle='none', alpha=0.5)
    h1.plot(THERMAL["E"], THERMAL["GCA_Edist"],
            marker='^', markersize=5, markevery=3, linestyle='none', alpha=0.5)

    h2.plot(np.array([-1, 1]), np.array([0.5, 0.5])*simtime_th/2, 'black')
    h2.plot(THERMAL["xi"], THERMAL["GO_xidist"], alpha=0.5)
    h2.plot(THERMAL["xi"], THERMAL["GCF_xidist"], alpha=0.5)
    h2.plot(THERMAL["xi"], THERMAL["GCA_xidist"], alpha=0.5)

    dense = 16
    h3.plot(SLOWING["Egrid"], SLOWING["analytical"], 'black')
    h3.plot(SLOWING["E"][:dense], SLOWING["GO_Edist"][:dense], marker='*',
            markersize=8, markevery=2, linestyle='none', alpha=0.5, color=cy[0])
    h3.plot(SLOWING["E"][dense:], SLOWING["GO_Edist"][dense:], marker='*',
            markersize=8, markevery=20, linestyle='none', alpha=0.5,color=cy[0])
    h3.plot(SLOWING["E"][:dense], SLOWING["GCF_Edist"][:dense], marker='.',
            markersize=10, markevery=2, linestyle='none', alpha=0.5,color=cy[1])
    h3.plot(SLOWING["E"][dense:], SLOWING["GCF_Edist"][dense:], marker='.',
            markersize=10, markevery=20, linestyle='none', alpha=0.5,color=cy[1])
    h3.plot(SLOWING["E"][:dense], SLOWING["GCA_Edist"][:dense], marker='^',
            markersize=5, markevery=2, linestyle='none', alpha=0.5, color=cy[2])
    h3.plot(SLOWING["E"][dense:], SLOWING["GCA_Edist"][dense:], marker='^',
            markersize=5, markevery=20, linestyle='none', alpha=0.5, color=cy[2])

    h4.plot(np.array([-1, 1]), np.array([0.5, 0.5])*slowingdowntime, 'black')
    h4.plot(SLOWING["xi"], SLOWING["GO_xidist"], alpha=0.5)
    h4.plot(SLOWING["xi"], SLOWING["GCF_xidist"], alpha=0.5)
    h4.plot(SLOWING["xi"], SLOWING["GCA_xidist"], alpha=0.5)

    #**************************************************************************#
    #*                         Decorate and finish                             #
    #*                                                                         #
    #**************************************************************************#
    h1.set_xlim(0, 1e4)
    h1.xaxis.set(ticks=[0, 2.5e3, 5e3, 7.5e3, 1e4],
                 ticklabels=['0', '2.5', '5', '7.5', '10'])
    h1.yaxis.set(ticks=np.array([0, 2, 4, 6])*1e-6,
                 ticklabels=['0', '2', '4', '6'])
    h1.tick_params(axis='y', direction='out')
    h1.tick_params(axis='x', direction='out')
    h1.spines['right'].set_visible(False)
    h1.spines['top'].set_visible(False)
    h1.yaxis.set_ticks_position('left')
    h1.xaxis.set_ticks_position('bottom')
    h1.set(ylabel=r"$f(E) \times 10^{-6}$ [1/eV]")
    h1.set(xlabel=r"Energy $\times 10^{3}$ [eV]")

    h2.set_xlim(-1, 1)
    h2.set_ylim(0, 1*simtime_th/2)
    h2.xaxis.set(ticks=[-1, -0.5, 0, 0.5, 1])
    h2.yaxis.set(ticks=np.array([0, 0.5, 1])*simtime_th/2,
                 ticklabels=['0', '5', '10'])
    h2.tick_params(axis='y', direction='out')
    h2.tick_params(axis='x', direction='out')
    h2.spines['right'].set_visible(False)
    h2.spines['top'].set_visible(False)
    h2.yaxis.set_ticks_position('left')
    h2.xaxis.set_ticks_position('bottom')
    h2.set(ylabel=r"$f(\xi) \times 10^{-4}$")
    h2.set(xlabel=r"Pitch")

    h3.set_xlim(0, 4e6)
    h3.xaxis.set(ticks=[0, 1e6, 2e6, 3e6, 4e6],
                 ticklabels=['0', '1', '2', '3', '4'])
    h3.yaxis.set(ticks=np.array([0, 2, 4, 6, 8])*1e-8,
                 ticklabels=['0', '2', '4', '6', '8'])
    h3.tick_params(axis='y', direction='out')
    h3.tick_params(axis='x', direction='out')
    h3.spines['right'].set_visible(False)
    h3.spines['top'].set_visible(False)
    h3.yaxis.set_ticks_position('left')
    h3.xaxis.set_ticks_position('bottom')
    h3.set(ylabel=r"$f(E) \times 10^{-8}$ [1/eV]")
    h3.set(xlabel=r"Energy $\times 10^{6}$ [eV]")

    h4.set_xlim(-1, 1)
    h4.set_ylim(0, 1 * ts)
    h4.xaxis.set(ticks=[-1, -0.5, 0, 0.5, 1])
    h4.yaxis.set(ticks=[0, 0.015, 0.03],
                 ticklabels=['0', '15', '30'])
    h4.tick_params(axis='y', direction='out')
    h4.tick_params(axis='x', direction='out')
    h4.spines['right'].set_visible(False)
    h4.spines['top'].set_visible(False)
    h4.yaxis.set_ticks_position('left')
    h4.xaxis.set_ticks_position('bottom')
    h4.set(ylabel=r"$f(\xi)\times 10^{-3}$")
    h4.set(xlabel=r"Pitch")

    l1 = mlines.Line2D([], [], color='black', linestyle='-',
                       label='Analytical',axes=h1)
    l2 = mlines.Line2D([], [], color=cy[0], marker='*', linestyle='-',
                       markersize=6, label='GO',axes=h1)
    l3 = mlines.Line2D([], [], color=cy[1], marker='.', linestyle='-',
                       markersize=8, label='GCF',axes=h1)
    l4 = mlines.Line2D([], [], color=cy[2], marker='^', linestyle='-',
                       markersize=4, label='GCA',axes=h1)

    h1.legend(handles=[l1, l2, l3, l4], loc='upper right', frameon=False,
              fontsize=8)

    frm  = lambda x: "%3.0f ms" % (x*1e3)
    text = r"$\tau_s$:      " + frm(slowingdowntime) + "\n" \
           + "GO:    " + frm(ts_GO) + "\n" \
           + "GCF:  " + frm(ts_GCF) + "\n" \
           + "GCA: " + frm(ts_GCA)
    h3.text(2e6, 3e-8, text, fontsize=9)

    plt.savefig("test_coulombcollisions.png", dpi=300)
    plt.show()

if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if( sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_coulombcollisions.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test elementary results such as gyromotion and drifts

This file initializes, runs, and plots test case for checking that ASCOT5
reproduces:

1. correct gyroradius and gyrofrequency in uniform magnetic field
2. correct E X B drift in uniform electromagnetic field
3. correct gradB drift in a magnetic field with constant gradient

Test 1 is done with GO mode and 2 and 3 with both GO and GC modes, latter
using the fixed step scheme. Tests are done in Cartesian electromagnetic
field (B_TC and E_TC) and without collisions. Test particle is a energetic
electron (and positron) so tests also verify that ASCOT5 is valid in
relativistic regime.

To init, run and check this test, call this script without any arguments. To
do only one of the above, call this script with an argument "init", "run", or
"check".

File: test_elementary.py
"""

import sys

import numpy                   as np
import matplotlib.pyplot       as plt
import unyt

import a5py.ascot5io.orbits    as orbits
import a5py.ascot5io.options   as options
import a5py.ascot5io.B_TC      as B_TC
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_3D     as N0_3D
import a5py.ascot5io.mrk_gc    as mrk
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhd
import a5py.ascot5io.asigma_loc as asigma_loc

import a5py.testascot.helpers as helpers

from a5py.ascot5io.ascot5 import Ascot
from a5py.physlib import e, m_e, c

def init():
    """
    Initialize tests

    This function initializes five test cases:
    - GYROMOTION tests that gyroradius and gyrofrequency are correct
    - EXB-GO tests that ExB-drift is correct for GO
    - EXB-GC tests that ExB-drift is correct for GC
    - GRADB-GO tests that gradB-drift is correct for GO
    - GRADB-GC tests that gradB-drift is correct for GC

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*                     Generate options for GYROMOTION                     #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-11
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 2e-9
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-11
    odict["ORBITWRITE_NPOINT"]         = 202

    options.write_hdf5(helpers.testfn, odict, desc="GYROMOTION")

    #**************************************************************************#
    #*              Generate options for EXB-GO and GRAD-GO                    #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 1e-7
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-11
    odict["ORBITWRITE_NPOINT"]         = 10002

    options.write_hdf5(helpers.testfn, odict, desc="EXB_GO")
    options.write_hdf5(helpers.testfn, odict, desc="GRADB_GO")

    #**************************************************************************#
    #*       Generate options for EXB-GC and GRADB-GC                          #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-9
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 1e-7
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-9
    odict["ORBITWRITE_NPOINT"]         = 102

    options.write_hdf5(helpers.testfn, odict, desc="EXB_GC")
    options.write_hdf5(helpers.testfn, odict, desc="GRADB_GC")

    #**************************************************************************#
    #*                      Markers are same for all tests                     #
    #*                                                                         #
    #**************************************************************************#
    Nmrk   = 2
    ids    = np.array([1, 2])
    weight = np.array([1, 1])
    mass   = m_e.to("amu") * np.array([1, 1])
    charge = 1       * np.array([1,-1])
    anum   = 1       * np.array([1, 0])
    znum   = 1       * np.array([1, 0])
    time   = 0       * np.array([1, 2])
    R      = 5       * np.array([1, 1])
    phi    = 90      * np.array([1, 1])
    z      = 0       * np.array([1, 1])
    zeta   = 0       * np.array([1, 1])
    energy = 100e6   * np.array([1, 1])
    pitch  = 0.5     * np.array([1, 1])
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="GYROMOTION")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="EXB_GO")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="EXB_GC")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="GRADB_GO")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="GRADB_GC")

    #**************************************************************************#
    #*             Magnetic and electric fields for GYROMOTION                 #
    #*                                                                         #
    #**************************************************************************#
    Bxyz   = np.array([5, 0, 0])
    gradB  = np.array([0,0,0,0,0,0,0,0,0])
    rhoval = 1.5
    B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval, desc="GYROMOTION")

    Exyz   = np.array([0, 0, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GYROMOTION")

    #**************************************************************************#
    #*            Magnetic and electric fields for EXB-GO and EXB-GC           #
    #*                                                                         #
    #**************************************************************************#
    Bxyz   = np.array([5, 0, 0])
    gradB  = np.array([0,0,0,0,0,0,0,0,0])
    rhoval = 1.5
    B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval, desc="EXB_GO")
    B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval, desc="EXB_GC")

    Exyz   = np.array([0, 1e6, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="EXB_GO")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="EXB_GC")

    #**************************************************************************#
    #*           Magnetic and electric fields for GRADB-GO and GRADB-GC        #
    #*                                                                         #
    #**************************************************************************#
    Bxyz   = np.array([5, 0, 0])
    gradB  = np.array([0,0,0.1,0,0,0,0,0,0])
    rhoval = 1.5
    B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval, desc="GRADB_GO")
    B_TC.write_hdf5(helpers.testfn, Bxyz, gradB, rhoval, desc="GRADB_GC")

    Exyz   = np.array([0, 0, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GRADB_GO")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GRADB_GC")

    #**************************************************************************#
    #*              Other inputs are trivial and same for all tests            #
    #*                                                                         #
    #**************************************************************************#
    nwall = 4
    Rwall = np.array([0.1, 100, 100, 0.1])
    zwall = np.array([-100, -100, 100, 100])
    for tname in ["GYROMOTION", "EXB_GO", "EXB_GC", "GRADB_GO", "GRADB_GC"]:
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall, desc=tname)
        N0_3D.write_hdf5_dummy(helpers.testfn, desc=tname)
        boozer.write_hdf5_dummy(helpers.testfn, desc=tname)
        mhd.write_hdf5_dummy(helpers.testfn, desc=tname)
        asigma_loc.write_hdf5_empty(helpers.testfn, desc=tname)

    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([1])
    mass   = np.array([1])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = 1e20 * np.ones(rho.shape)
    etemp  = 1e3  * np.ones(rho.shape)
    idens  = 1e20 * np.ones((rho.size, Nion))
    itemp  = 1e3  * np.ones(rho.shape)
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GYROMOTION")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="EXB_GO")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="EXB_GC")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GRADB_GO")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GRADB_GC")

def run():
    """
    Run tests.
    """
    for test in ["GYROMOTION", "EXB_GO", "EXB_GC", "GRADB_GO", "GRADB_GC"]:
        helpers.set_and_run(test)

def check():
    """
    Plot the results of these tests.

    This function makes three plots.
    - One that shows gyromotion viewed from above (to chect test GYROMOTION)
    - Two that shows gyromotion viewed from side (to show drifts in EXB-GO,
      EXB_GC, GRADB-GO, and GRADB-GC)
    - Also the analytical an numerical values are shown
    """
    a5 = Ascot(helpers.testfn)

    GYROMOTION = {}
    EXB        = {}
    GRADB      = {}
    GYROMOTION["GO"] = a5["GYROMOTION"].orbit
    EXB["GO"]        = a5["EXB_GO"].orbit
    EXB["GC"]        = a5["EXB_GC"].orbit
    GRADB["GO"]      = a5["GRADB_GO"].orbit
    GRADB["GC"]      = a5["GRADB_GC"].orbit

    # Electron energy in Joules
    E = 100e6 * unyt.eV
    # Lorentz factor
    gamma = 1 + E / ( m_e * c * c )
    # Pitch, magnetic field strength and gradient, electric field, and velocity
    xi    = 0.5
    B     = 5
    gradB = 0.1
    E     = 1e6
    v     = np.sqrt(1.0 - 1.0 / ( gamma * gamma ) ) * c

    f = plt.figure(figsize=(11.9/2.54, 8/2.54))
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    h1 = f.add_subplot(1,3,1)
    h1.set_position([0.12, 0.27, 0.26, 1], which='both')

    h2 = f.add_subplot(1,3,2)
    h2.set_position([0.56, 0.30, 0.42, 1], which='both')

    h3 = f.add_subplot(1,3,3)
    h3.set_position([0.56, -0.13, 0.42, 1], which='both')

    #**************************************************************************#
    #*                           Check GYROMOTION                              #
    #*                                                                         #
    #**************************************************************************#

    # Analytical values
    rhog   = gamma * np.sqrt(1 - 0.5 * 0.5) * m_e * v / (B * e)
    omegag = (unyt.kg / unyt.C)*e * B / ( gamma * m_e)
    m = unyt.m

    # Numerical values
    ang  = GYROMOTION["GO"]["phi"] * np.pi / 180
    igo  = GYROMOTION["GO"]["id"]
    x    = GYROMOTION["GO"]["y"].v
    y    = GYROMOTION["GO"]["z"].v
    time = GYROMOTION["GO"]["time"][igo==1]

    rho   = np.max( np.sqrt((x-5)**2 + y**2) )
    zero_crossings = np.where(np.diff(np.sign(x[igo==1])))[0].size
    omega = zero_crossings * np.pi * 2 / time[-1]

    # Plot
    h1.plot(x[igo==1] - 5, y[igo==1], linewidth=3)
    h1.plot(rhog * np.sin(np.linspace(0,2*np.pi,360)),
            rhog * np.cos(np.linspace(0,2*np.pi,360)), linestyle="--",
            color="black", alpha=0.7)

    #**************************************************************************#
    #*                           Check EXB                                     #
    #*                                                                         #
    #**************************************************************************#

    # Analytical values
    v_ExB = E * B / (B*B)

    # Numerical values
    ang = EXB["GO"]["phi"] * np.pi / 180
    igo = EXB["GO"]["id"]
    xgo = EXB["GO"]["y"] / m
    ygo = EXB["GO"]["z"] / m

    igo0  = a5["EXB_GO"]["inistate"]["id"]
    time  = a5["EXB_GO"]["endstate"]["time"]
    ang   = a5["EXB_GO"]["inistate"]["phi"] * np.pi / 180
    xgo0  = a5["EXB_GO"]["inistate"]["y"] / m
    ygo0  = a5["EXB_GO"]["inistate"]["z"] / m
    ang   = a5["EXB_GO"]["endstate"]["phi"] * np.pi / 180
    xgo1  = a5["EXB_GO"]["endstate"]["y"] / m
    ygo1  = a5["EXB_GO"]["endstate"]["z"] / m

    ang = EXB["GC"]["phi"] * np.pi / 180
    igc = EXB["GC"]["id"]
    xgc = EXB["GC"]["y"] / m
    ygc = EXB["GC"]["z"] / m

    vgo1_ExB = ((ygo1[igo0==1] - ygo0[igo0==1]) / (time[igo0==1]))[0]

    y          = ygc[igc==1]
    t          = EXB["GC"]["time"][igc==1]
    vgc1_ExB = (y[-1] - y[0]) / (t[-1] - t[0])

    # Plot
    h2.plot(xgo[igo==1] - 0.07 - 5, ygo[igo==1])
    h2.plot(xgo[igo==2][1:-1] + 0.07 - 5, ygo[igo==2][1:-1])
    h2.plot(xgc[igc==1] - 0.07 - 5, ygc[igc==1], color="red")
    h2.plot(xgc[igc==2] + 0.07 - 5, ygc[igc==2], color="red")

    #**************************************************************************#
    #*                           Check GRADB                                   #
    #*                                                                         #
    #**************************************************************************#

    # Analytical values
    v_gradB = ( gamma * m_e * (1.0 - xi*xi) * v * v / ( 2 * e * B) ) \
              * gradB * B / (B*B)

    # Numerical values
    ang = GRADB["GO"]["phi"] * np.pi / 180
    igo = GRADB["GO"]["id"]
    xgo = GRADB["GO"]["y"] / m
    ygo = GRADB["GO"]["z"] / m

    igo0  = a5["GRADB_GO"]["inistate"]["id"]
    time  = a5["GRADB_GO"]["endstate"]["time"]
    ang   = a5["GRADB_GO"]["inistate"]["phi"]
    xgo0  = a5["GRADB_GO"]["inistate"]["y"] / m
    zgo0  = a5["GRADB_GO"]["inistate"]["z"] / m
    ang   = a5["GRADB_GO"]["endstate"]["phi"]
    xgo1  = a5["GRADB_GO"]["endstate"]["y"] / m
    ygo1  = a5["GRADB_GO"]["endstate"]["z"] / m

    ang = GRADB["GC"]["phi"] * np.pi / 180
    igc = GRADB["GC"]["id"]
    xgc = GRADB["GC"]["y"] / m
    ygc = GRADB["GC"]["z"] / m

    vgo1_gradB = ((xgo1[igo0==1] - xgo0[igo0==1]) / (time[igo0==1]))[0]

    x          = xgc[igc==1]
    t          = GRADB["GC"]["time"][igc==1]
    vgc1_gradB = (x[-1] - x[0]) / (t[-1] - t[0])

    # Plot
    h3.plot(ygo[igo==1][1:-1] - 0.07, xgo[igo==1][1:-1] - 5)
    h3.plot(ygo[igo==2] + 0.07, xgo[igo==2] - 5)
    h3.plot(ygc[igc==1] - 0.07, xgc[igc==1] - 5, color="red")
    h3.plot(ygc[igc==2] + 0.07, xgc[igc==2] - 5, color="red")

    # Print analytical values
    text1  = r"$\rho_{g}$ = %2.3f cm" % (rhog*100)
    text1 += r" (%2.3f cm)" % (rho*100)
    text1 += "\n"+ r"$\omega_{g}$ = " + latex_float(omegag.v)
    text1 += r" $\frac{\mathrm{rad}}{\mathrm{s}}$ "
    text1 += r"(" + latex_float(omegag.v)
    text1 += r" $\frac{\mathrm{rad}}{\mathrm{s}}$)"
    text1 += "\n" + r"$v_{E \times B}$ = " + latex_float(v_ExB) + r" m/s"
    text1 += "\n" + r"(" + latex_float(-vgo1_ExB.v) + r" m/s, "
    text1 += latex_float(-vgc1_ExB.v) + r" m/s)" + "\n"
    text1 += r"$v_{\nabla B}$ = " + latex_float(v_gradB.v) + r" m/s"
    text1 += "\n" + r"(" + latex_float(-vgo1_gradB.v) + r" m/s, "
    text1 += latex_float(-vgc1_gradB.v) + r" m/s)" + "\n"
    h1.text(-0.1, -0.26, text1, fontsize=9)

    #**************************************************************************#
    #*                           Finalize                                      #
    #*                                                                         #
    #**************************************************************************#

    # Decorate all plots
    h1.axis('scaled')
    h1.xaxis.set(ticks=[-0.06, 0, 0.06], ticklabels=[-6, 0, 6])
    h1.yaxis.set(ticks=[-0.06, 0, 0.06], ticklabels=[-6, 0, 6])
    h1.tick_params(axis='y', direction='out')
    h1.tick_params(axis='x', direction='out')
    h1.set(xlabel="$x$ [cm]", ylabel="$y$ [cm]")

    h2.axis('scaled')
    h2.set(xlim=[-0.15, 0.15], ylim=[-0.08, 0.08])
    h2.xaxis.set(ticks=[-0.12,-0.07,-0.02, 0.02, 0.07, 0.12], ticklabels=[])
    h2.yaxis.set(ticks=[-0.08, 0, 0.08], ticklabels=[-8, 0, 8])
    h2.tick_params(axis='y', direction='out')
    h2.tick_params(axis='x', direction='out')
    h2.set(ylabel="$y$ [cm]")

    h3.axis('scaled')
    h3.set(xlim=[-0.15, 0.15], ylim=[-0.08, 0.08])
    h3.xaxis.set(ticks=[-0.12,-0.07,-0.02, 0.02, 0.07, 0.12],
                 ticklabels=[-5, 0, 5, -5, 0, 5])
    h3.yaxis.set(ticks=[-0.08, 0, 0.08], ticklabels=[-8, 0, 8])
    h3.tick_params(axis='y', direction='out')
    h3.tick_params(axis='x', direction='out')
    h3.set(xlabel="$x$ [cm]", ylabel="$y$ [cm]")

    h2.plot([0,0], [-8, 8],color="black", linewidth=1)
    h3.plot([0,0], [-8, 8],color="black", linewidth=1)

    plt.savefig("test_elementary.png", dpi=300)
    plt.show()

def latex_float(f):
    float_str = "{0:.4g}".format(f)
    if "e" in float_str:
        base, exponent = float_str.split("e")
        return r"${0} \times 10^{{{1}}}$".format(base, int(exponent))
    else:
        return float_str

if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if(   sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_elementary.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test guiding center transformation.

This file initializes, runs, and plots test case for checking that ASCOT5
performs firstorder guiding center correctly. Since there is no analytical
way to verify this, except for comparing to the very same formulas that ASCOT5
implements, instead we:

1. Initialize a guiding center marker and trace it in GC mode ITER-like
   axisymmetric analytic (B_GS) field without collisions.
2. Trace the very same guiding center with FO mode, so that we get to perform
   guiding center to particle transformation.
3. Repeat step 2. but this time perform particle to guiding center at each time-
   step, and record the guiding center coordinates while continuing the FO
   simulation. This way we can compare the "actual" guiding center trajectory
   to the one computed with GC mode.

Ideally, the actual guiding center trajectory should be close to the GC
trajectory and the former should have almost eliminated the gyroangle
dependency. Do note that the guiding center transformation is not an exact but
done to the first order. Also note that because magnetic moment is not an exact
invariant, the actual guiding center and GC trajectories never match completely.
Test particle is a energetic electron so this test verifies that the
transformation is valid also in the relativstic regime.

This test also shows the difference in first and zeroth order velocity space*
transformation. This is done by taking the calculated orbit from step 2. and
using that to launch additional GC simulations. These simulations are launched
within the first gyroperiod, so ideally the resulting GC trajectories should
be identical.

To init, run and check this test, call this script without any arguments. To
do only one of the above, call this script with an argument "init", "run", or
"check".

*Zeroth order spatial transformation is trivial as then guiding center and
particle coordinates would be identical. Therefore that is not examined here.

File: test_gctransform.py
"""
import sys

import numpy                   as np
import unyt
import matplotlib.pyplot       as plt
import matplotlib.lines        as mlines

import a5py.ascot5io.orbits    as orbits
import a5py.ascot5io.options   as options
import a5py.ascot5io.B_GS      as B_GS
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_3D     as N0_3D
import a5py.ascot5io.mrk_gc    as mrk
import a5py.ascot5io.mrk_prt   as prt
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhd
import a5py.ascot5io.asigma_loc as asigma_loc

import a5py.testascot.helpers as helpers

from a5py.preprocessing.analyticequilibrium import psi0 as psifun

from a5py.ascot5io.ascot5 import Ascot
from a5py.physlib import e, m_a, c

psi_mult  = 200
R0        = 6.2
z0        = 0
Bphi0     = 5.3

# ITER-like equilibrium
psi_coeff = np.array([ 8.629e-02,  3.279e-01,  5.268e-01, -2.366e-01,
                       3.825e-01, -3.573e-01, -1.484e-02,  1.506e-01,
                       7.428e-01, -4.447e-01, -1.084e-01,  1.281e-02, -0.155])

nrep = 10

def init():
    """
    Initialize tests

    This function initializes four test cases:
    - GCTRANSFORM_GC initializes a guiding center marker and simulates it in GC
      mode.
    - GCTRANSFORM_GO initializes a guiding center marker and simulates it in FO
      mode.
    - GCTRANSFORM_GO2GC is same as GO but guiding center transformation is done
      at each
      time-step to get the actual guiding center trajectory.
    - GCTRANSFORM_FIRST initializes particle markers from the orbit calculated
      in GCTRANSFORM_GO. Markers are initialized at different time but all are
      launched within the first gyroperiod. Markers are simulated in GC mode
      with the initial guiding center transformation done to first order.
    - GCTRANSFORM_ZEROTH is same as GCTRANSFORM_FIRST but the guiding center
      transformation is only done to zeroth order in velocity space.

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*                     Generate options for GC runs                        #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 3e-5
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 4e-10
    odict["ORBITWRITE_NPOINT"]         = 75002

    options.write_hdf5(helpers.testfn, odict, desc="GCTRANSFORM_GC")
    options.write_hdf5(helpers.testfn, odict, desc="GCTRANSFORM_FIRST")

    odict["DISABLE_FIRSTORDER_GCTRANS"] = 1
    options.write_hdf5(helpers.testfn, odict, desc="GCTRANSFORM_ZEROTH")

    #**************************************************************************#
    #*                     Generate options for GO                             #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 3e-5
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 4e-10
    odict["ORBITWRITE_NPOINT"]         = 75002

    options.write_hdf5(helpers.testfn, odict, desc="GCTRANSFORM_GO")

    #**************************************************************************#
    #*                     Generate options for GO2GC                          #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["RECORD_MODE"]               = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 3e-5
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 4e-10
    odict["ORBITWRITE_NPOINT"]         = 75002

    options.write_hdf5(helpers.testfn, odict, desc="GCTRANSFORM_GO2GC")

    #**************************************************************************#
    #*                 Marker input consisting of an alpha particle            #
    #*                                                                         #
    #**************************************************************************#
    Nmrk   = 1
    ids    = np.array([1])
    weight = 1       * np.ones(ids.shape)
    pitch  = 0.4     * np.ones(ids.shape)
    mass   = m_a.to("amu") * np.ones(ids.shape)
    charge = 2       * np.ones(ids.shape)
    anum   = 4       * np.ones(ids.shape)
    znum   = 2       * np.ones(ids.shape)
    time   = 0       * np.ones(ids.shape)
    R      = 7.6     * np.ones(ids.shape)
    phi    = 90      * np.ones(ids.shape)
    z      = 0       * np.ones(ids.shape)
    zeta   = 2       * np.ones(ids.shape)
    energy = 3.5e6   * np.ones(ids.shape)
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="GCTRANSFORM_GC")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="GCTRANSFORM_GO")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="GCTRANSFORM_GO2GC")

    #**************************************************************************#
    #*                     Construct ITER-like magnetic field                  #
    #*                                                                         #
    #**************************************************************************#
    Rmin = 4; Rmax = 8.5; nR = 120; zmin = -4; zmax = 4; nz = 200;
    B_GS.write_hdf5_B_2DS(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                          Rmin, Rmax, nR, zmin, zmax, nz,
                          desc="GCTRANSFORM_GC")
    B_GS.write_hdf5_B_2DS(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                          Rmin, Rmax, nR, zmin, zmax, nz,
                          desc="GCTRANSFORM_GO")
    B_GS.write_hdf5_B_2DS(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                          Rmin, Rmax, nR, zmin, zmax, nz,
                          desc="GCTRANSFORM_GO2GC")
    B_GS.write_hdf5_B_2DS(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                          Rmin, Rmax, nR, zmin, zmax, nz,
                          desc="GCTRANSFORM_ZEROTH")
    B_GS.write_hdf5_B_2DS(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                          Rmin, Rmax, nR, zmin, zmax, nz,
                          desc="GCTRANSFORM_FIRST")

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    Exyz   = np.array([0, 0, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GCTRANSFORM_GC")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GCTRANSFORM_GO")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GCTRANSFORM_GO2GC")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GCTRANSFORM_ZEROTH")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="GCTRANSFORM_FIRST")

    nwall = 4
    Rwall = np.array([0.1, 100, 100, 0.1])
    zwall = np.array([-100, -100, 100, 100])
    for tname in ["GCTRANSFORM_GC", "GCTRANSFORM_GO", "GCTRANSFORM_GO2GC",
                  "GCTRANSFORM_ZEROTH", "GCTRANSFORM_FIRST"]:
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall, desc=tname)
        N0_3D.write_hdf5_dummy(helpers.testfn, desc=tname)
        boozer.write_hdf5_dummy(helpers.testfn, desc=tname)
        mhd.write_hdf5_dummy(helpers.testfn, desc=tname)
        asigma_loc.write_hdf5_empty(helpers.testfn, desc=tname)

    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([1])
    mass   = np.array([1])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = 1e20 * np.ones(rho.shape)
    etemp  = 1e3  * np.ones(rho.shape)
    idens  = 1e20 * np.ones((rho.size, Nion))
    itemp  = 1e3  * np.ones(rho.shape)
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GCTRANSFORM_GC")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GCTRANSFORM_GO")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GCTRANSFORM_GO2GC")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GCTRANSFORM_ZEROTH")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="GCTRANSFORM_FIRST")

def run():
    """
    Run tests.

    This first runs the GCTRANSFORM_GC, GCTRANSFORM_GO, and GCTRANSFORM_GO2GC
    cases, then uses results of GCTRANSFORM_GO to generate markers for
    GCTRANSFORM_ZEROTH and GCTRANSFORM_FIRST which are then run.
    """
    for test in ["GCTRANSFORM_GC", "GCTRANSFORM_GO", "GCTRANSFORM_GO2GC"]:
        helpers.set_and_run(test)

    a5 = Ascot(helpers.testfn)

    dt = 20
    Nmrk   = nrep
    ids    = np.linspace(1, Nmrk, Nmrk)
    weight = 1       * np.ones(ids.shape)
    mass   = m_a.to("amu") * np.ones(ids.shape)
    charge = 2       * np.ones(ids.shape)
    znum   = 4       * np.ones(ids.shape)
    anum   = 2       * np.ones(ids.shape)
    time   = a5["GCTRANSFORM_GO"]["orbit"]["time"][0:Nmrk*dt:dt]
    R      = a5["GCTRANSFORM_GO"]["orbit"]["r"][0:Nmrk*dt:dt]
    phi    = a5["GCTRANSFORM_GO"]["orbit"]["phi"][0:Nmrk*dt:dt]
    z      = a5["GCTRANSFORM_GO"]["orbit"]["z"][0:Nmrk*dt:dt]
    vR     = a5["GCTRANSFORM_GO"]["orbit"]["vr"][0:Nmrk*dt:dt]
    vphi   = a5["GCTRANSFORM_GO"]["orbit"]["vphi"][0:Nmrk*dt:dt]
    vz     = a5["GCTRANSFORM_GO"]["orbit"]["vz"][0:Nmrk*dt:dt]
    prt.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, vR, vphi, vz,
                   anum, znum, weight, time, desc="GCTRANSFORM_ZEROTH")
    prt.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, vR, vphi, vz,
                   anum, znum, weight, time, desc="GCTRANSFORM_FIRST")

    for test in ["GCTRANSFORM_ZEROTH", "GCTRANSFORM_FIRST"]:
        helpers.set_and_run(test)

def check():
    """
    Plot the results of these tests.

    This function makes four plots.
    - One that shows difference in magnetic moment with respect to time between
      cases GCTRANSFORM_GC and GCTRANSFORM_GO, and GCTRANSFORM_GC and
      GCTRANSFORM_GO2GC.
    - The second one is the same but for the parallel velocity.
    - Third plot shows the orbits of these three cases on a Rz plane with focus
      on the banana turning point.
    - Fourth one is the same as previous plot but showing cases GCTRANSFORM_GC,
      GCTRANSFORM_ZEROTH and GCTRANSFORM_FIRST. The latter two have orbits from
      multiple markers but these are plotted with same color.
    """
    a5 = Ascot(helpers.testfn)

    f = plt.figure(figsize=(11.9/2.54, 8/2.54))
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    h1 = f.add_subplot(1,4,1)
    h1.set_position([0.12, 0.58, 0.26, 0.38], which='both')

    h2 = f.add_subplot(1,4,2)
    h2.set_position([0.12, 0.15, 0.26, 0.38], which='both')

    h3 = f.add_subplot(1,4,3)
    h3.set_position([0.5, 0.15, 0.24, 0.8], which='both')

    h4 = f.add_subplot(1,4,4)
    h4.set_position([0.74, 0.15, 0.24, 0.8], which='both')

    # Make plots, scale the plotted quantities.
    # For some reason GCTRANSFORM_GO2GC GCTRANSFORM_GC are not equal in length
    # and we need to have [:-1]?
    ccyc = plt.rcParams['axes.prop_cycle'].by_key()['color'] # default colors

    h1.plot(a5["GCTRANSFORM_GO"]["orbit"]["time"]*1e6,
            ( a5["GCTRANSFORM_GO"]["orbit"]["mu"]
              - a5["GCTRANSFORM_GC"]["orbit"]["mu"] )/e / 1e4 )
    h1.plot(a5["GCTRANSFORM_GO2GC"]["orbit"]["time"]*1e6,
            ( a5["GCTRANSFORM_GO2GC"]["orbit"]["mu"]
              - a5["GCTRANSFORM_GC"]["orbit"]["mu"][:-1] )/e / 1e4 )

    dppar = ( a5["GCTRANSFORM_GO"]["orbit"]["ppar"]
             - a5["GCTRANSFORM_GC"]["orbit"]["ppar"] )
    dppar.convert_to_mks()
    h2.plot(a5["GCTRANSFORM_GO"]["orbit"]["time"]*1e6, dppar / 1e-21 )

    dppar = ( a5["GCTRANSFORM_GO2GC"]["orbit"]["ppar"]
              - a5["GCTRANSFORM_GC"]["orbit"]["ppar"][:-1] )
    dppar.convert_to_mks()
    h2.plot(a5["GCTRANSFORM_GO2GC"]["orbit"]["time"]*1e6, dppar / 1e-21 )

    h3.plot(a5["GCTRANSFORM_GO"]["orbit"]["r"],
            a5["GCTRANSFORM_GO"]["orbit"]["z"])
    h3.plot(a5["GCTRANSFORM_GO2GC"]["orbit"]["r"],
            a5["GCTRANSFORM_GO2GC"]["orbit"]["z"])
    h3.plot(a5["GCTRANSFORM_GC"]["orbit"]["r"],
            a5["GCTRANSFORM_GC"]["orbit"]["z"])

    for i in range(0, nrep):
        id0 = a5["GCTRANSFORM_ZEROTH"]["orbit"]["id"]
        id1 = a5["GCTRANSFORM_FIRST"]["orbit"]["id"]
        h4.plot(a5["GCTRANSFORM_ZEROTH"]["orbit"]["r"][id0==i+1],
                a5["GCTRANSFORM_ZEROTH"]["orbit"]["z"][id0==i+1], ccyc[4])
        h4.plot(a5["GCTRANSFORM_FIRST"]["orbit"]["r"][id1==i+1],
                a5["GCTRANSFORM_FIRST"]["orbit"]["z"][id1==i+1], ccyc[3])

    h4.plot(a5["GCTRANSFORM_GO2GC"]["orbit"]["r"],
            a5["GCTRANSFORM_GO2GC"]["orbit"]["z"], ccyc[1])

    #**************************************************************************#
    #*                 Finalize and print and show the figure                  #
    #*                                                                         #
    #**************************************************************************#

    h1.set_xlim(0, 30)
    h1.set_ylim(-2, 2)
    h1.spines['right'].set_visible(False)
    h1.spines['top'].set_visible(False)
    h1.spines['bottom'].set_visible(False)
    h1.yaxis.set_ticks_position('left')
    h1.xaxis.set_ticks_position('bottom')
    h1.tick_params(axis='y', direction='out')
    h1.tick_params(axis='x', direction='inout')
    h1.xaxis.set(ticks=[0, 10, 20, 30], ticklabels=[])
    h1.yaxis.set(ticks=[-2, 0, 2])
    h1.set(ylabel=r"$\Delta \mu$ [$10^{4}$ eV/T]")

    h2.set_xlim(0, 30)
    h2.set_ylim(-1, 2)
    h2.spines['right'].set_visible(False)
    h2.spines['top'].set_visible(False)
    h2.yaxis.set_ticks_position('left')
    h2.xaxis.set_ticks_position('bottom')
    h2.tick_params(axis='y', direction='out')
    h2.tick_params(axis='x', direction='out')
    h2.xaxis.set(ticks=[0, 10, 20, 30])
    h2.yaxis.set(ticks=[-1, 0, 2],
                 ticklabels=[-1, 0, 2])
    h2.set(ylabel=r"$\Delta p_\parallel$ [$10^{-21}$ kg$\cdot$m/s]",
           xlabel=r"Time [$10^{-6}$ s]")

    h3.axis('scaled')
    h3.set_xlim(6.18, 6.68)
    h3.set_ylim(1.6, 2.6)
    h3.tick_params(axis='y', direction='out')
    h3.tick_params(axis='x', direction='out')
    h3.xaxis.set(ticks=[6.2, 6.4, 6.6])
    h3.yaxis.set(ticks=[1.7, 1.9, 2.1, 2.3, 2.5],
                 ticklabels=[1.7, '', 2.1, '', 2.5])
    h3.set(ylabel=r"$z$ [m]")
    h3.set_xlabel(r"$R$ [m]", position=[1,0,1,1]);

    h4.axis('scaled')
    h4.set_xlim(6.18, 6.68)
    h4.set_ylim(1.6, 2.6)
    h4.tick_params(axis='y', direction='inout')
    h4.tick_params(axis='x', direction='out')
    h4.xaxis.set(ticks=[6.2, 6.4, 6.6])
    h4.yaxis.set(ticks=[1.7, 1.9, 2.1, 2.3, 2.5], ticklabels=[])

    l1 = mlines.Line2D([], [], color=ccyc[2], linestyle='-', label='GC',
                       axes=h3)
    l2 = mlines.Line2D([], [], color=ccyc[0], linestyle='-', label='GO',
                       axes=h3)
    l3 = mlines.Line2D([], [], color=ccyc[1], linestyle='-', label='GO2GC',
                       axes=h3)
    h3.legend(handles=[l1, l2, l3], loc='lower left', frameon=False,
              fontsize=9)

    l1 = mlines.Line2D([], [], color=ccyc[1], linestyle='-', label='GO2GC',
                       axes=h4)
    l2 = mlines.Line2D([], [], color=ccyc[4], linestyle='-',
                       label=r'0th order GC', axes=h4)
    l3 = mlines.Line2D([], [], color=ccyc[3], linestyle='-',
                       label=r'1st order GC', axes=h4)
    h4.legend(handles=[l1, l2, l3], loc='lower left', frameon=False,
              fontsize=9)

    plt.savefig("test_gctransform.png", dpi=300)
    plt.show()

if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if( sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_gctransform.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test the implementation of the MHD module.

This file initializes, runs, and plots test case for checking that the helical
perturbations introduced by the MHD module conserve the quantity

H = K - omega*P/n,

where K is particle hamiltonian, P canonical angular toroidal momentum, omega
mode frequency and n mode toroidal number.

This is validated for all particle simulation modes.

To init, run and check this test, call this script without any arguments. To
do only one of the above, call this script with an argument "init", "run", or
"check".

File: test_orbitfollowing.py
"""
import sys

import numpy                   as np
import matplotlib.pyplot       as plt

import a5py.ascot5io.ascot5    as ascot5
import a5py.ascot5io.orbits    as orbits
import a5py.ascot5io.options   as options
import a5py.ascot5io.B_GS      as B_GS
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_3D     as N0_3D
import a5py.ascot5io.mrk_gc    as mrk
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhdmod
import a5py.ascot5io.asigma_loc as asigma_loc

import a5py.testascot.helpers as helpers

from a5py.ascot5io.ascot5 import write_dummy
from a5py.physlib import m_e, m_p

from a5py.ascotpy.ascotpy import Ascotpy
from a5py.preprocessing.generateboozer import generate as genbooz

from matplotlib.gridspec import GridSpec

psi_mult  = 200
R0        = 6.2
z0        = 0
Bphi0     = 5.3

# ITER-like but circular equilibrium
psi_coeff = np.array([ 2.218e-02, -1.288e-01, -4.177e-02, -6.227e-02,
                       6.200e-03, -1.205e-03, -3.701e-05,  0,
                       0,          0,          0,          0,         -0.155])

# MHD mode specifications
nmode = 2
mmode = 3
modefreq = 1e6
amplitude = 1e-6 # 1e-2

def init():
    """
    Initialize tests

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*             Generate options for GO                                     #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-11
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 1e-3
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_MHD"]                = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-7
    odict["ORBITWRITE_NPOINT"]         = 1e4

    options.write_hdf5(helpers.testfn, odict, desc="MHD_GO")

    #**************************************************************************#
    #*             Generate options for GCF                                    #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 1e-3
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_MHD"]                = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-7
    odict["ORBITWRITE_NPOINT"]         = 1e4

    options.write_hdf5(helpers.testfn, odict, desc="MHD_GCF")

    #**************************************************************************#
    #*             Generate options for GCA                                    #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["ENABLE_ADAPTIVE"]           = 1
    odict["ADAPTIVE_TOL_ORBIT"]        = 1e-10
    odict["ADAPTIVE_MAX_DRHO"]         = 0.1
    odict["ADAPTIVE_MAX_DPHI"]         = 10
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-11
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 1e-3
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_MHD"]                = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-7
    odict["ORBITWRITE_NPOINT"]         = 1e4

    options.write_hdf5(helpers.testfn, odict, desc="MHD_GCA")


    #**************************************************************************#
    #*                            Field line markers                           #
    #*                                                                         #
    #**************************************************************************#
    Nmrk   = 2
    ids    = np.array([1, 2])
    R      = np.array([7.0, 8.0])
    pitch  = np.array([0.4, 0.9])
    zeta   = 2       * np.array([1, 1])
    energy = 1e6     * np.array([1, 1])
    mass   = m_e.to("amu") * np.array([1, 1])
    charge = 1       * np.array([1, 1])
    anum   = 1       * np.array([1, 1])
    znum   = 1       * np.array([1, 1])
    weight = 1       * np.array([1, 1])
    time   = 0       * np.array([1, 1])
    phi    = 0       * np.array([1, 1])
    z      = 0       * np.array([1, 1])
    for d in ["MHD_GO", "MHD_GCF", "MHD_GCA"]:
        mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                       R, phi, z, energy, pitch, zeta,
                       anum, znum, weight, time, desc=d)

    #**************************************************************************#
    #*                     Construct ITER-like magnetic field                  #
    #*                                                                         #
    #**************************************************************************#
    for d in ["MHD_GO", "MHD_GCF", "MHD_GCA"]:
        B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                        desc=d)

    #**************************************************************************#
    #*                     Construct the boozer input                          #
    #*                                                                         #
    #**************************************************************************#
    h5 = ascot5.Ascot(helpers.testfn)
    a5 = Ascotpy(helpers.testfn)
    a5.init(bfield=h5.bfield["MHD_GCF"].get_qid())
    booz, rhoprof, qprof = genbooz(a5,
                                   rgrid=np.linspace(4, 8.5, 200),
                                   zgrid=np.linspace(-2.5, 2.5, 200),
                                   npsi=200,
                                   nthgeo=360,
                                   nthbzr=360,
                                   raxis=h5.bfield["MHD_GCF"].read()["raxis"],
                                   zaxis=h5.bfield["MHD_GCF"].read()["zaxis"],
                                   psi0=h5.bfield["MHD_GCF"].read()["psi0"],
                                   psi1=h5.bfield["MHD_GCF"].read()["psi1"])
    a5.free(bfield=True)
    for d in ["MHD_GO", "MHD_GCF", "MHD_GCA"]:
        boozer.write_hdf5(helpers.testfn, desc=d, **booz)

    #**************************************************************************#
    #* Helical perturbation which is not used in the simulation but is used    #
    #* during the postprocessing                                               #
    #*                                                                         #
    #**************************************************************************#
    mhd = {}
    mhd["nmode"]     = 1
    mhd["nmodes"]    = np.array([nmode])
    mhd["mmodes"]    = np.array([mmode])
    mhd["amplitude"] = np.array([amplitude])
    mhd["omega"]     = np.array([modefreq])
    mhd["phase"]     = np.array([np.pi/4])

    mhd["nrho"]   = 100
    mhd["rhomin"] = 0.0
    mhd["rhomax"] = 1

    psigrid = np.linspace(mhd["rhomin"], mhd["rhomax"], mhd["nrho"])
    alpha   = np.exp( -(psigrid-0.85)**2/0.1 )
    phi     = np.exp( -(psigrid-0.85)**2/0.1 )*1e9
    mhd["phi"]   = np.tile(phi, (mhd["nmode"],1)).T
    mhd["alpha"] = np.tile(alpha, (mhd["nmode"],1)).T
    for d in ["MHD_GO", "MHD_GCF", "MHD_GCA"]:
        mhdmod.write_hdf5(helpers.testfn, desc=d, **mhd)

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    h5 = ascot5.Ascot(helpers.testfn)
    for d in ["MHD_GO", "MHD_GCF", "MHD_GCA"]:
        write_dummy(helpers.testfn, "efield", desc=d)
        write_dummy(helpers.testfn, "plasma", desc=d)
        write_dummy(helpers.testfn, "wall", desc=d)
        write_dummy(helpers.testfn, "neutral", desc=d)
        write_dummy(helpers.testfn, "nbi", desc=d)
        asigma_loc.write_hdf5_empty(helpers.testfn, desc=d)


def run():
    """
    Run tests.
    """
    for test in ["MHD_GO", "MHD_GCF", "MHD_GCA"]:
        helpers.set_and_run(test)

def check():
    """
    Plot the results of these tests.

    This function makes four plots.
    - One that shows conservation of energy for all cases
    - One that shows conservation of magnetic moment for all cases
    - One that shows conservation of toroidal canonical momentum for all cases
    - And one that shows trajectories on a Rz plane for all cases
    """
    h5 = ascot5.Ascot(helpers.testfn)

    a5 = Ascotpy(helpers.testfn)
    a5.init(bfield=h5.bfield["MHD_GCF"].get_qid(),
            boozer=h5.boozer["MHD_GCF"].get_qid(),
            mhd=h5.mhd["MHD_GCF"].get_qid())
    a5.free(bfield=True, boozer=True, mhd=True)

    f = plt.figure(figsize=(11.9/2.54, 8/2.54))
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    gs = GridSpec(ncols=1, nrows=1, left=0.15, bottom=0.15, figure=f)
    s1 = f.add_subplot(gs[0,0])

    i = 0
    for run in ["MHD_GO", "MHD_GCF", "MHD_GCA"]:
        orb = h5[run].orbit

        H = (orb["energy"] + orb["mhd_epot"] * orb["charge"]).to("J").v

        ctor = orb["ptor"] + orb["r"] * orb["charge"] * orb["mhdalpha"] * orb["bphi"]
        ctor.convert_to_mks()
        ctor = ctor.v

        ids  = orb["id"] == 2
        time = orb["time"][ids]
        H    = H[ids]
        P    = (modefreq)*ctor[ids]/nmode
        K    = H - P

        ls = [":", "--", "-"]

        h1 = s1.plot(time, (P-P[0])/H[0], color="C0", linestyle=ls[i],
                       alpha=0.3, label=r"$\omega_nP_\phi/n$")[0]
        h2 = s1.plot(time, (H-H[0])/H[0], color="C1", linestyle=ls[i],
                       alpha=0.3, label=r"$E$")[0]
        h3 = s1.plot(time, (K-K[0])/H[0], color="C2", linestyle=ls[i],
                       alpha=0.3, linewidth=3, label=r"$E-\omega_nP_\phi/n$")[0]
        i += 1

    a5.free(bfield=True, boozer=True, mhd=True)

    s1.ticklabel_format(scilimits=(0,0))

    s1.set_xlim(0, 1e-3)
    s1.set_xticks([0, 0.5e-3, 1e-3])

    s1.set_xlabel(r"Time [s]")
    s1.set_ylabel(r"Relative difference $(y-y_0)/E_0$")

    plt.legend(handles=[h1,h2,h3], loc='upper center',
               bbox_to_anchor=(0.6, 1.15),
               ncol=3, fancybox=True, shadow=True)

    plt.savefig("test_mhd.png", dpi=300)
    plt.show()


if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if( sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_mhd.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test collisional transport in tokamak magnetic field.

This file initializes, runs, and plots test case for checking that ASCOT5
reproduces collisional transport in tokamak magnetic field, i.e., the
so-called neoclassical transport, correctly. This transport is diffusive
with the diffusion coefficient being determined by the ratio between
orbit bounce frequency and collision frequency. If collision frequency is a
lot higher than the bounce frequency, the transport is so-called
Pfirsch-Schlüter diffusion. If opposite is true, then the transport is called
banana transport. Between these two regimes is the plateau regime where
transport does not depend on collision frequency.

The neoclassical transport is tested on an analytical tokamak field that has
circular cross-section. Electrons are used as markers.

To init, run and check this test, call this script without any arguments. To
do only one of the above, call this script with an argument "init", "run", or
"check".

File: test_neoclassicaltransport.py
"""

import sys

import numpy                   as np
import matplotlib.pyplot       as plt

import a5py.ascot5io.options   as options
import a5py.ascot5io.B_GS      as B_GS
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_3D     as N0_3D
import a5py.ascot5io.mrk_gc    as mrk
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhd
import a5py.ascot5io.asigma_loc as asigma_loc

import a5py.testascot.helpers as helpers

from a5py.ascot5io.ascot5 import Ascot
from a5py.physlib import e, m_e, m_p, c, eps_0

from a5py.preprocessing.analyticequilibrium import psi0 as psifun

Nmrk  = 100
nscan = 20

Ti   = 1e3
ni   = np.power( 10, np.linspace(17.5, 22.0, nscan) )
Ekin = 1e3

R0 = 6.2
r0 = 1.0
z0 = 0
B0 = 5.3
psi_mult = 200

# ITER-like but circular equilibrium
psi_coeff = np.array([ 2.218e-02, -1.288e-01, -4.177e-02, -6.227e-02,
                       6.200e-03, -1.205e-03, -3.701e-05,  0,
                       0,          0,          0,          0,         -0.155])

def init():
    """
    Initialize tests

    This function initializes parameter scan for three test cases:
    - NEOCLASS-GO tests gyro-orbit scheme
    - NEOCLASS-GCF tests guiding center fixed-scheme
    - NEOCLASS-GCA tests guiding center adaptive scheme

    Each test case is initialized nscan times; each with a increasing value for
    the plasma density. The test cases are named with running index on their
    prefix e.g. NEOCLASS-GO1, NEOCLASS-GO2, and so on.

    All cases are run without energy collisions.

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*                  Generate options for NEOCLASS-GO                       #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 3e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1

    for i in range(1, nscan+1):
        odict["ENDCOND_LIM_SIMTIME"]   = np.maximum( 5e-4,
                                                     4e-2 / (ni[i-1]/ni[0]) )
        odict["FIXEDSTEP_USERDEFINED"] = np.minimum( 2e-9,
                                                     3e-10 / (ni[i-1]/ni[-1]) )
        options.write_hdf5(helpers.testfn, odict,
                           desc="NEOCLASS_GO" + str(i))

    #**************************************************************************#
    #*                  Generate options for NEOCLASS-GCF                      #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1

    for i in range(1, nscan+1):
        odict["ENDCOND_LIM_SIMTIME"]   = np.maximum( 5e-4,
                                                     4e-2 / (ni[i-1]/ni[0]) )
        odict["FIXEDSTEP_USERDEFINED"] = np.minimum( 2e-8,
                                                     5e-10 / (ni[i-1]/ni[-1]) )
        options.write_hdf5(helpers.testfn, odict,
                           desc="NEOCLASS_GCF" + str(i))

    #**************************************************************************#
    #*                  Generate options for NEOCLASS-GCA                      #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["ENABLE_ADAPTIVE"]           = 1
    odict["ADAPTIVE_TOL_ORBIT"]        = 5e-9
    odict["ADAPTIVE_TOL_COL"]          = 9e-1
    odict["ADAPTIVE_MAX_DRHO"]         = 0.1
    odict["ADAPTIVE_MAX_DPHI"]         = 10
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_COULOMB_COLLISIONS"] = 1

    for i in range(1, nscan+1):
        odict["ENDCOND_LIM_SIMTIME"] = np.maximum( 5e-4,
                                                   4e-2 / (ni[i-1]/ni[0]) )
        options.write_hdf5(helpers.testfn, odict,
                           desc="NEOCLASS_GCA" + str(i))

    #**************************************************************************#
    #*                    Marker input consisting of electrons                 #
    #*                                                                         #
    #**************************************************************************#
    ids    = np.linspace(1,Nmrk,Nmrk)
    weight = np.ones(ids.shape)
    mass   = m_e.to("amu") * np.ones(ids.shape)
    charge = 1       * np.ones(ids.shape)
    anum   = 0       * np.ones(ids.shape)
    znum   = 0       * np.ones(ids.shape)
    time   = 0       * np.ones(ids.shape)
    R      = (R0+r0) * np.ones(ids.shape)
    phi    = 90      * np.ones(ids.shape)
    z      = 0       * np.ones(ids.shape)
    energy = Ekin    * np.ones(ids.shape)
    zeta   = 2 * np.pi * np.random.rand(1,Nmrk)
    pitch  = 1 - 2 * np.random.rand(1,Nmrk)
    for i in range(1, nscan+1):
        mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, energy, pitch, zeta,
                       anum, znum, weight, time,
                       desc="NEOCLASS_GO" + str(i))
        mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, energy, pitch, zeta,
                       anum, znum, weight, time,
                       desc="NEOCLASS_GCF" + str(i))
        mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass,
                       charge, R, phi, z, energy, pitch, zeta,
                       anum, znum, weight, time,
                       desc="NEOCLASS_GCA" + str(i))

    #**************************************************************************#
    #*             Analytical ITER-like but circular magnetic field            #
    #*                                                                         #
    #**************************************************************************#
    for i in range(1, nscan+1):
        B_GS.write_hdf5(helpers.testfn, R0, z0, B0, psi_mult, psi_coeff,
                        desc="NEOCLASS_GO" + str(i))
        B_GS.write_hdf5(helpers.testfn, R0, z0, B0, psi_mult, psi_coeff,
                        desc="NEOCLASS_GCF" + str(i))
        B_GS.write_hdf5(helpers.testfn, R0, z0, B0, psi_mult, psi_coeff,
                        desc="NEOCLASS_GCA" + str(i))

    #**************************************************************************#
    #*     Plasma consisting of protons only (to avoid e-e collisions)         #
    #*                                                                         #
    #**************************************************************************#
    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([1])
    mass   = np.array([1])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = 1    * np.ones(rho.shape)
    etemp  = Ti * np.ones(rho.shape)
    itemp  = Ti  * np.ones(rho.shape)
    for i in range(1, nscan+1):
        idens = ni[i-1] * np.ones((rho.size, Nion))
        P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge,
                        rho, edens, etemp, idens, itemp,
                        desc="NEOCLASS_GO" + str(i))
        P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge,
                        rho, edens, etemp, idens, itemp,
                        desc="NEOCLASS_GCF" + str(i))
        P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge,
                        rho, edens, etemp, idens, itemp,
                        desc="NEOCLASS_GCA" + str(i))

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    for i in range(1, nscan+1):
        Exyz   = np.array([0, 0, 0])
        E_TC.write_hdf5(helpers.testfn, Exyz, desc="NEOCLASS_GO" + str(i))
        E_TC.write_hdf5(helpers.testfn, Exyz, desc="NEOCLASS_GCF" + str(i))
        E_TC.write_hdf5(helpers.testfn, Exyz, desc="NEOCLASS_GCA" + str(i))

        nwall = 4
        Rwall = np.array([0.1, 100, 100, 0.1])
        zwall = np.array([-100, -100, 100, 100])
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall,
                        desc="NEOCLASS_GO" + str(i))
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall,
                        desc="NEOCLASS_GCF" + str(i))
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall,
                        desc="NEOCLASS_GCA" + str(i))

        N0_3D.write_hdf5_dummy(helpers.testfn, desc="NEOCLASS_GO" + str(i))
        N0_3D.write_hdf5_dummy(helpers.testfn, desc="NEOCLASS_GCF" + str(i))
        N0_3D.write_hdf5_dummy(helpers.testfn, desc="NEOCLASS_GCA" + str(i))

        boozer.write_hdf5_dummy(helpers.testfn, desc="NEOCLASS_GO"  + str(i))
        boozer.write_hdf5_dummy(helpers.testfn, desc="NEOCLASS_GCF" + str(i))
        boozer.write_hdf5_dummy(helpers.testfn, desc="NEOCLASS_GCA" + str(i))

        mhd.write_hdf5_dummy(helpers.testfn,    desc="NEOCLASS_GO"  + str(i))
        mhd.write_hdf5_dummy(helpers.testfn,    desc="NEOCLASS_GCF" + str(i))
        mhd.write_hdf5_dummy(helpers.testfn,    desc="NEOCLASS_GCA" + str(i))

        asigma_loc.write_hdf5_empty(helpers.testfn, desc="NEOCLASS_GO" + str(i))
        asigma_loc.write_hdf5_empty(helpers.testfn, desc="NEOCLASS_GCF" + str(i))
        asigma_loc.write_hdf5_empty(helpers.testfn, desc="NEOCLASS_GCA" + str(i))


def run():
    """
    Run tests.
    """
    for test in ["NEOCLASS_GO", "NEOCLASS_GCF", "NEOCLASS_GCA"]:
        for i in range(1, nscan+1):
            helpers.set_and_run(test + str(i))


def check():
    """
    Plot the results of these tests.

    This function makes a one plot that shows numerically evaluated diffusion
    coefficients for all modes as a function of effective collision frequency
    (collision frequency divided by orbit transit frequency). For comparison,
    the analytical results for different regimes are also shown and also
    the boundaries of the regimes themselves.
    """

    a5 = Ascot(helpers.testfn)
    axisr = a5["NEOCLASS_GO1"].bfield.read()["raxis"][:][0]

    # Map rho values to R outer mid-plane values
    R_omp    = np.linspace(axisr,8.4,1000)
    z_omp    = 0*np.ones(R_omp.shape)
    psi_omp  = psifun(R_omp/axisr, z_omp/axisr, psi_coeff[0], psi_coeff[1],
                      psi_coeff[2], psi_coeff[3], psi_coeff[4], psi_coeff[5],
                      psi_coeff[6], psi_coeff[7], psi_coeff[8], psi_coeff[9],
                      psi_coeff[10], psi_coeff[11], psi_coeff[12]) * psi_mult
    psi_axis = psifun(axisr/axisr, 0, psi_coeff[0], psi_coeff[1],
                      psi_coeff[2], psi_coeff[3], psi_coeff[4], psi_coeff[5],
                      psi_coeff[6], psi_coeff[7], psi_coeff[8], psi_coeff[9],
                      psi_coeff[10], psi_coeff[11], psi_coeff[12]) * psi_mult
    rho_omp  = np.sqrt(np.absolute( (psi_omp - psi_axis) / psi_axis ))

    #**************************************************************************#
    #*               Evaluate the numerical diffusion coefficients             #
    #*                                                                         #
    #**************************************************************************#
    DGO  = np.zeros(nscan)
    DGCF = np.zeros(nscan)
    DGCA = np.zeros(nscan)
    for i in range(1, nscan+1):
        inistate = a5["NEOCLASS_GO" + str(i)].inistate
        endstate = a5["NEOCLASS_GO" + str(i)].endstate
        ri = np.interp(inistate["rho"], rho_omp, R_omp)
        rf = np.interp(endstate["rho"], rho_omp, R_omp)
        t  = endstate["time"]

        DGO[i-1] = np.mean( np.power(ri - rf, 2) / (2*t) )

        inistate = a5["NEOCLASS_GCF" + str(i)].inistate
        endstate = a5["NEOCLASS_GCF" + str(i)].endstate
        ri = np.interp(inistate["rho"], rho_omp, R_omp)
        rf = np.interp(endstate["rho"], rho_omp, R_omp)
        t  = endstate["time"]

        DGCF[i-1] = np.mean( np.power(ri - rf, 2) / (2*t) )

        inistate = a5["NEOCLASS_GCA" + str(i)].inistate
        endstate = a5["NEOCLASS_GCA" + str(i)].endstate
        ri = np.interp(inistate["rho"], rho_omp, R_omp)
        rf = np.interp(endstate["rho"], rho_omp, R_omp)
        t  = endstate["time"]

        DGCA[i-1] = np.mean( np.power(ri - rf, 2) / (2*t) )

    #**************************************************************************#
    #*                  Evaluate the analytical estimate                       #
    #*                                                                         #
    #**************************************************************************#
    eps = (r0 - (axisr - R0)) / axisr
    q   = 1.7 # This was verified numerically
    B   = 5.3

    gamma  = 1 + (Ekin * e / ( m_e * c * c )).v
    v      = np.sqrt(1.0 - 1.0 / ( gamma * gamma ) ) * c
    omegat = (v / (q * axisr)) * np.sqrt(0.5*eps)
    rhog   = gamma * m_e * v / (B * e)

    clog     = 15
    density  = np.power( 10, np.linspace(np.log10(ni[0]) - 1,
                                         np.log10(ni[-1]) + 1, 50) )
    collfreq = (np.sqrt(2/np.pi) / 3) \
               * np.power(e*e / ( 4*np.pi*eps_0 ), 2) \
               * ( 4*np.pi / np.sqrt( m_e*np.power(Ti*e, 3) ) ) * density * clog
    veff = collfreq/omegat
    # Add values needed for plotting a continuous curve
    veff = np.append(veff, [1, np.power(eps, 3.0/2.0)])
    veff.sort()

    Dps = q * q * veff * omegat * np.power(rhog, 2) / 2
    Dp  = q * q * rhog * rhog * omegat \
          * np.ones(veff.shape) / 2
    Db  = np.power(eps, -3.0/2.0) * Dps

    # x coordinate for plotting the numerical coefficients
    collfreq = (np.sqrt(2/np.pi) / 3) \
               * np.power(e*e / ( 4*np.pi*eps_0 ), 2) \
               * ( 4*np.pi / np.sqrt( m_e*np.power(Ti*e, 3) ) ) * ni * clog
    veff_x   = collfreq/omegat

    #**************************************************************************#
    #*                                  Plot                                   #
    #*                                                                         #
    #**************************************************************************#

    f = plt.figure(figsize=(11.9/2.54, 5/2.54))
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    h = plt.gca()
    h.set_position([0.15, 0.25, 0.82, 0.7], which='both')

    h.plot(np.log10(np.power([eps, eps], 3.0/2.0)), np.array([-6, 0]), 'grey')
    h.plot(np.log10(np.array([1, 1])), np.array([-6, 0]), 'grey')

    h.plot(np.log10(veff_x), np.log10(DGO), linestyle='none', marker='*',
           markersize=12, alpha=0.5, label='GO')
    h.plot(np.log10(veff_x), np.log10(DGCF), linestyle='none', marker='.',
           markersize=10, alpha=0.5, label='GCF')
    h.plot(np.log10(veff_x), np.log10(DGCA), linestyle='none', marker='^',
           markersize=6, alpha=0.5, label='GCA')

    ind = veff >= 1
    h.plot(np.log10(veff[ind]), np.log10(Dps[ind]), 'black')
    ind = np.logical_and(veff >= np.power(eps, 3.0/2.0), veff <= 1)
    h.plot(np.log10(veff[ind]), np.log10(Dp[ind]), 'black')
    ind = veff <= np.power(eps, 3.0/2.0)
    h.plot(np.log10(veff[ind]), np.log10(Db[ind]), 'black')

    h.set_xlim(-3, 2)
    h.set_ylim(-6, -2)
    h.tick_params(axis='y', direction='out')
    h.tick_params(axis='x', direction='out')
    h.xaxis.set(ticks=np.linspace(-3, 2, 6))
    h.yaxis.set(ticks=np.linspace(-6, -1, 6),
                ticklabels=['$10^{-6}$', '$10^{-5}$', '$10^{-4}$', '$10^{-3}$',
                            '$10^{-2}$', '$10^{-1}$'])
    h.set(xlabel=r"$\log_{10}\nu^*$", ylabel=r"$D$ [m$^2$/s]")

    h.text(-1.9, -5.7, r"$\nu^*=\epsilon^{3/2}$", fontsize=10,
           bbox={'facecolor':'white', 'edgecolor':'none', 'pad':0})
    h.text(-0.25, -5.7, r"$\nu^*=1$", fontsize=10,
           bbox={'facecolor':'white', 'edgecolor':'none', 'pad':0})
    h.text(-2.5, -4, r"$D_{B}$", fontsize=10)
    h.text(-0.8, -3, r"$D_{P}$", fontsize=10)
    h.text(0.4, -2.1, r"$D_{PS}$", fontsize=10)

    h.legend(loc='upper left', frameon=False, fontsize=10, numpoints=1)

    plt.savefig("test_neoclassicaltransport.png", dpi=300)
    plt.show()

if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if(   sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_neoclassicaltransport.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test conservation properties of the integration of the Hamiltonian motion.

This file initializes, runs, and plots test case for checking that ASCOT5
orbit integrators have the following properties:

1. gyro-orbit integrator (VPA) conserves energy, and magnetic moment and
   canonical toroidal momentum are bounded
2. guiding center fixed step scheme (RK4) conserves magnetic moment
   (trivially) and drift in energy and canonical toroidal momentum decrese
   with time-step
3. guiding center adaptive step scheme (Cash-Karp) conserves magnetic moment
   (trivially) and drift in energy and canonical toroidal momentum decrese
   with error-tolerance

Tests are done in ITER-like but analytic (B_GS) magnetic field without
collisions. Test particle is a energetic electron (and positron) so tests also
verify that ASCOT5 is valid in relativistic regime.

To init, run and check this test, call this script without any arguments. To
do only one of the above, call this script with an argument "init", "run", or
"check".

File: test_orbitfollowing.py
"""
import sys

import numpy                   as np
import matplotlib.pyplot       as plt
import unyt

import a5py.ascot5io.orbits    as orbits
import a5py.ascot5io.options   as options
import a5py.ascot5io.B_GS      as B_GS
import a5py.ascot5io.E_TC      as E_TC
import a5py.ascot5io.plasma_1D as P_1D
import a5py.ascot5io.wall_2D   as W_2D
import a5py.ascot5io.N0_3D     as N0_3D
import a5py.ascot5io.mrk_gc    as mrk
import a5py.ascot5io.boozer    as boozer
import a5py.ascot5io.mhd       as mhd
import a5py.ascot5io.asigma_loc as asigma_loc

import a5py.testascot.helpers as helpers

from a5py.preprocessing.analyticequilibrium import psi0 as psifun

from a5py.ascot5io.ascot5 import Ascot
from a5py.physlib import e, m_e, c

psi_mult  = 200
R0        = 6.2
z0        = 0
Bphi0     = 5.3

# ITER-like but circular equilibrium
psi_coeff = np.array([ 2.218e-02, -1.288e-01, -4.177e-02, -6.227e-02,
                       6.200e-03, -1.205e-03, -3.701e-05,  0,
                       0,          0,          0,          0,         -0.155])

# Change this if you wish to use spline-interpolated magnetic field instead
use_spline = True

def init():
    """
    Initialize tests

    This function initializes three test cases:
    - ORBFOL-GO tests VPA algorithm used in integrating gyro-orbit motion
    - ORBFOL-GCF tests RK4 used in integrating guiding center motion with fixed
      time-step
    - ORBFOL-GCA tests Cash-Karp used in integrating guiding center motion with
      adaptive time-step

    Input fields contain the test case name (to which the input corresponds to)
    as a description.
    """

    #**************************************************************************#
    #*                     Generate options for ORBFOL-GO                      #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 1
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-11
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 5e-6
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-10
    odict["ORBITWRITE_NPOINT"]         = 50002

    options.write_hdf5(helpers.testfn, odict, desc="ORBFOL_GO")

    #**************************************************************************#
    #*                     Generate options for ORBFOL-GCF                     #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-10
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 5e-6
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-8
    odict["ORBITWRITE_NPOINT"]         = 502

    options.write_hdf5(helpers.testfn, odict, desc="ORBFOL_GCF")

    #**************************************************************************#
    #*                     Generate options for ORBFOL-GCA                     #
    #*                                                                         #
    #**************************************************************************#
    odict = options.generateopt()
    helpers.clean_opt(odict)

    odict["SIM_MODE"]                  = 2
    odict["ENABLE_ADAPTIVE"]           = 1
    odict["ADAPTIVE_TOL_ORBIT"]        = 1e-10
    odict["ADAPTIVE_MAX_DRHO"]         = 0.1
    odict["ADAPTIVE_MAX_DPHI"]         = 10
    odict["FIXEDSTEP_USE_USERDEFINED"] = 1
    odict["FIXEDSTEP_USERDEFINED"]     = 1e-8
    odict["ENDCOND_SIMTIMELIM"]        = 1
    odict["ENDCOND_LIM_SIMTIME"]       = 5e-6
    odict["ENABLE_ORBIT_FOLLOWING"]    = 1
    odict["ENABLE_ORBITWRITE"]         = 1
    odict["ORBITWRITE_MODE"]           = 1
    odict["ORBITWRITE_INTERVAL"]       = 1e-8
    odict["ORBITWRITE_NPOINT"]         = 502

    options.write_hdf5(helpers.testfn, odict, desc="ORBFOL_GCA")

    #**************************************************************************#
    #*           Marker input consisting of an electron and positron           #
    #*                                                                         #
    #**************************************************************************#
    Nmrk   = 2
    ids    = np.array([1, 2])
    weight = np.array([1, 1])
    pitch  = np.array([0.4, 0.9])
    mass   = m_e.to("amu") * np.array([1, 1])
    charge = 1       * np.array([1,-1])
    anum   = 1       * np.array([1, 0])
    znum   = 1       * np.array([1, 0])
    time   = 0       * np.array([1, 1])
    R      = 7.6     * np.array([1, 1])
    phi    = 90      * np.array([1, 1])
    z      = 0       * np.array([1, 1])
    zeta   = 2       * np.array([1, 1])
    energy = 10e6    * np.array([1, 1])
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="ORBFOL_GO")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="ORBFOL_GCF")
    mrk.write_hdf5(helpers.testfn, Nmrk, ids, mass, charge,
                   R, phi, z, energy, pitch, zeta,
                   anum, znum, weight, time, desc="ORBFOL_GCA")

    #**************************************************************************#
    #*                     Construct ITER-like magnetic field                  #
    #*                                                                         #
    #**************************************************************************#
    if use_spline:
        B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                        desc="ORBFOL_GO")
        B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                        desc="ORBFOL_GCF")
        B_GS.write_hdf5(helpers.testfn, R0, z0, Bphi0, psi_mult, psi_coeff,
                        desc="ORBFOL_GCA")
    else:
        Rmin = 4; Rmax = 8.5; nR = 120; zmin = -4; zmax = 4; nz = 200;
        B_GS.write_hdf5_B_2D(helpers.testfn, R0, z0, Bphi0, psi_mult,
                             psi_coeff, Rmin, Rmax, nR, zmin, zmax, nz,
                             desc="ORBFOL_GO")
        B_GS.write_hdf5_B_2D(helpers.testfn, R0, z0, Bphi0, psi_mult,
                             psi_coeff, Rmin, Rmax, nR, zmin, zmax, nz,
                             desc="ORBFOL_GCF")
        B_GS.write_hdf5_B_2D(helpers.testfn, R0, z0, Bphi0, psi_mult,
                             psi_coeff, Rmin, Rmax, nR, zmin, zmax, nz,
                             desc="ORBFOL_GCA")

    #**************************************************************************#
    #*                     Rest of the inputs are trivial                      #
    #*                                                                         #
    #**************************************************************************#
    Exyz   = np.array([0, 0, 0])
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="ORBFOL_GO")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="ORBFOL_GCF")
    E_TC.write_hdf5(helpers.testfn, Exyz, desc="ORBFOL_GCA")

    nwall = 4
    Rwall = np.array([0.1, 100, 100, 0.1])
    zwall = np.array([-100, -100, 100, 100])
    for tname in ["ORBFOL_GO", "ORBFOL_GCF", "ORBFOL_GCA"]:
        W_2D.write_hdf5(helpers.testfn, nwall, Rwall, zwall, desc=tname)
        N0_3D.write_hdf5_dummy(helpers.testfn, desc=tname)
        boozer.write_hdf5_dummy(helpers.testfn, desc=tname)
        mhd.write_hdf5_dummy(helpers.testfn, desc=tname)
        asigma_loc.write_hdf5_empty(helpers.testfn, desc=tname)

    Nrho   = 3
    Nion   = 1
    znum   = np.array([1])
    anum   = np.array([1])
    mass   = np.array([1])
    charge = np.array([1])
    rho    = np.array([0, 0.5, 100])
    edens  = 1e20 * np.ones(rho.shape)
    etemp  = 1e3  * np.ones(rho.shape)
    idens  = 1e20 * np.ones((rho.size, Nion))
    itemp  = 1e3  * np.ones(rho.shape)
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="ORBFOL_GO")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="ORBFOL_GCF")
    P_1D.write_hdf5(helpers.testfn, Nrho, Nion, anum, znum, mass, charge, rho,
                    edens, etemp, idens, itemp, desc="ORBFOL_GCA")


def run():
    """
    Run tests.
    """
    for test in ["ORBFOL_GO", "ORBFOL_GCF", "ORBFOL_GCA"]:
        helpers.set_and_run(test)

def check():
    """
    Plot the results of these tests.

    This function makes four plots.
    - One that shows conservation of energy for all cases
    - One that shows conservation of magnetic moment for all cases
    - One that shows conservation of toroidal canonical momentum for all cases
    - And one that shows trajectories on a Rz plane for all cases
    """
    a5 = Ascot(helpers.testfn)

    raxis = R0

    f = plt.figure(figsize=(11.9/2.54, 8/2.54))
    plt.rc('xtick', labelsize=10)
    plt.rc('ytick', labelsize=10)
    plt.rc('axes', labelsize=10)
    plt.rcParams['mathtext.fontset'] = 'stix'
    plt.rcParams['font.family'] = 'STIXGeneral'

    h1 = f.add_subplot(1,4,1)
    h1.set_position([0.12, 0.72, 0.4, 0.25], which='both')

    h2 = f.add_subplot(1,4,2)
    h2.set_position([0.12, 0.44, 0.4, 0.25], which='both')

    h3 = f.add_subplot(1,4,3)
    h3.set_position([0.12, 0.155, 0.4, 0.25], which='both')

    h4 = f.add_subplot(1,4,4)
    h4.set_position([0.6, 0.3, 0.45, 0.45], which='both')

    colors = ["C0", "C9", "C2", "C8", "C3", "C1"]

    #**************************************************************************#
    #*     Evaluate and plot conservation quantities for ORBFOL_GO             #
    #*                                                                         #
    #**************************************************************************#
    ORBFOL = {}
    ORBFOL["GO"] = {}
    orb = a5["ORBFOL_GO"]["orbit"]

    ORBFOL["GO"]["time"] = orb["time"]
    ORBFOL["GO"]["id"]   = orb["id"]
    ORBFOL["GO"]["r"]    = orb["r"]
    ORBFOL["GO"]["z"]    = orb["z"]
    ORBFOL["GO"]["ekin"] = orb["ekin"]
    ORBFOL["GO"]["mu"]   = orb["mu"]
    ORBFOL["GO"]["ctor"] = orb["ctor"]

    id1 = ORBFOL["GO"]["id"] == 1
    id2 = ORBFOL["GO"]["id"] == 2
    plot_relerr(h1, ORBFOL["GO"]["time"][id1], ORBFOL["GO"]["ekin"][id1],
                colors[0])
    plot_relerr(h1, ORBFOL["GO"]["time"][id2], ORBFOL["GO"]["ekin"][id2],
                colors[1])
    plot_relerr(h2, ORBFOL["GO"]["time"][id2], ORBFOL["GO"]["mu"][id2],
                colors[0])
    plot_relerr(h2, ORBFOL["GO"]["time"][id1], ORBFOL["GO"]["mu"][id1],
                colors[1])
    plot_relerr(h3, ORBFOL["GO"]["time"][id1], ORBFOL["GO"]["ctor"][id1],
                colors[0])
    plot_relerr(h3, ORBFOL["GO"]["time"][id2], ORBFOL["GO"]["ctor"][id2],
                colors[1])
    h4.plot(        ORBFOL["GO"]["r"][id1],    ORBFOL["GO"]["z"][id1],
                    colors[0], alpha=0.7)
    h4.plot(        ORBFOL["GO"]["r"][id2],    ORBFOL["GO"]["z"][id2],
                    colors[1], alpha=0.7)

    #**************************************************************************#
    #*     Evaluate and plot conservation quantities for ORBFOL_GCF            #
    #*                                                                         #
    #**************************************************************************#
    ORBFOL["GCF"] = {}
    orb = a5["ORBFOL_GCF"]["orbit"]

    ORBFOL["GCF"]["time"] = orb["time"]
    ORBFOL["GCF"]["id"]   = orb["id"]
    ORBFOL["GCF"]["r"]    = orb["r"]
    ORBFOL["GCF"]["z"]    = orb["z"]
    ORBFOL["GCF"]["ekin"] = orb["ekin"]
    ORBFOL["GCF"]["mu"]   = orb["mu"]
    ORBFOL["GCF"]["ctor"] = orb["ctor"]

    id1 = ORBFOL["GCF"]["id"] == 1
    id2 = ORBFOL["GCF"]["id"] == 2
    plot_relerr(h1, ORBFOL["GCF"]["time"][id1], ORBFOL["GCF"]["ekin"][id1],
                colors[2])
    plot_relerr(h1, ORBFOL["GCF"]["time"][id2], ORBFOL["GCF"]["ekin"][id2],
                colors[3])
    plot_relerr(h2, ORBFOL["GCF"]["time"][id1], ORBFOL["GCF"]["mu"][id1],
                colors[2])
    plot_relerr(h2, ORBFOL["GCF"]["time"][id2], ORBFOL["GCF"]["mu"][id2],
                colors[3])
    plot_relerr(h3, ORBFOL["GCF"]["time"][id1], ORBFOL["GCF"]["ctor"][id1],
                colors[2])
    plot_relerr(h3, ORBFOL["GCF"]["time"][id2], ORBFOL["GCF"]["ctor"][id2],
                colors[3])
    h4.plot(        ORBFOL["GCF"]["r"][id1],    ORBFOL["GCF"]["z"][id1],
                colors[2], alpha=0.7, linewidth=1)
    h4.plot(        ORBFOL["GCF"]["r"][id2],    ORBFOL["GCF"]["z"][id2],
                colors[3], alpha=0.7, linewidth=1)

    #**************************************************************************#
    #*     Evaluate and plot conservation quantities for ORBFOL_GCA            #
    #*                                                                         #
    #**************************************************************************#
    ORBFOL["GCA"] = {}
    orb = a5["ORBFOL_GCA"]["orbit"]

    ORBFOL["GCA"]["time"] = orb["time"]
    ORBFOL["GCA"]["id"]   = orb["id"]
    ORBFOL["GCA"]["r"]    = orb["r"]
    ORBFOL["GCA"]["z"]    = orb["z"]
    ORBFOL["GCA"]["ekin"] = orb["ekin"]
    ORBFOL["GCA"]["mu"]   = orb["mu"]
    ORBFOL["GCA"]["ctor"] = orb["ctor"]

    id1 = ORBFOL["GCA"]["id"] == 1
    id2 = ORBFOL["GCA"]["id"] == 2
    plot_relerr(h1, ORBFOL["GCA"]["time"][id1], ORBFOL["GCA"]["ekin"][id1],
                colors[4])
    plot_relerr(h1, ORBFOL["GCA"]["time"][id2], ORBFOL["GCA"]["ekin"][id2],
                colors[5])
    plot_relerr(h2, ORBFOL["GCA"]["time"][id1], ORBFOL["GCA"]["mu"][id1],
                colors[4])
    plot_relerr(h2, ORBFOL["GCA"]["time"][id2], ORBFOL["GCA"]["mu"][id2],
                colors[5])
    plot_relerr(h3, ORBFOL["GCA"]["time"][id1], ORBFOL["GCA"]["ctor"][id1],
                colors[4])
    plot_relerr(h3, ORBFOL["GCA"]["time"][id2], ORBFOL["GCA"]["ctor"][id2],
                colors[5])
    h4.plot(        ORBFOL["GCA"]["r"][id1],    ORBFOL["GCA"]["z"][id1],
                    colors[4], alpha=0.7, linewidth=1)
    h4.plot(        ORBFOL["GCA"]["r"][id2],    ORBFOL["GCA"]["z"][id2],
                colors[5], alpha=0.7, linewidth=1)

    #**************************************************************************#
    #*                 Finalize and print and show the figure                  #
    #*                                                                         #
    #**************************************************************************#

    h1.set_xlim(0, 5e-6)
    h1.xaxis.set(ticks=[0, 1e-6, 2e-6, 3e-6, 4e-6, 5e-6], ticklabels=[])
    h1.yaxis.set(ticks=np.array([-6,-4, -2, 0, 2])*1e-11,
                 ticklabels=[-6, '', '', 0, 2])
    h1.tick_params(axis='y', direction='out')
    h1.tick_params(axis='x', direction='out')
    h1.spines['right'].set_visible(False)
    h1.spines['top'].set_visible(False)
    h1.yaxis.set_ticks_position('left')
    h1.xaxis.set_ticks_position('bottom')
    h1.set(ylabel=r"$\Delta E/E_0\;[10^{-11}]$")

    h2.set_xlim(0, 5e-6)
    h2.xaxis.set(ticks=[0, 1e-6, 2e-6, 3e-6, 4e-6, 5e-6], ticklabels=[])
    h2.yaxis.set(ticks=np.array([-4, 0, 4])*1e-3, ticklabels=[-4, 0, 4])
    h2.tick_params(axis='y', direction='out')
    h2.tick_params(axis='x', direction='out')
    h2.spines['right'].set_visible(False)
    h2.spines['top'].set_visible(False)
    h2.yaxis.set_ticks_position('left')
    h2.xaxis.set_ticks_position('bottom')
    h2.set(ylabel=r"$\Delta \mu/\mu_0\;[10^{-3}]$")

    h3.set_xlim(0, 5e-6)
    h3.xaxis.set(ticks=[0, 1e-6, 2e-6, 3e-6, 4e-6, 5e-6],
                 ticklabels=[0, 1, 2, 3, 4, 5])
    h3.yaxis.set(ticks=np.array([-6, 0, 6])*1e-7, ticklabels=[-6, 0, 6])
    h3.tick_params(axis='y', direction='out')
    h3.tick_params(axis='x', direction='out')
    h3.spines['right'].set_visible(False)
    h3.spines['top'].set_visible(False)
    h3.yaxis.set_ticks_position('left')
    h3.xaxis.set_ticks_position('bottom')
    h3.set(ylabel=r"$\Delta P/P_0\;[10^{-7}]$", xlabel=r"Time [$10^{-6}$ s]")

    h4.axis('scaled')
    h4.xaxis.set(ticks=[5, 6.5, 8])
    h4.yaxis.set(ticks=[-1.5, 0, 1.5])
    h4.tick_params(axis='y', direction='out')
    h4.tick_params(axis='x', direction='out')
    h4.set(xlabel="$R$ [m]", ylabel="$z$ [m]")

    legend = [r"GO-p", r"GCF-p", r"GCA-p", r"GO-t", r"GCF-t", r"GCA-t"]
    h4.text(5.0, 2.5, legend[0], fontsize=9, color=colors[1])
    h4.text(6.0, 2.5, legend[1], fontsize=9, color=colors[3])
    h4.text(7.0, 2.5, legend[2], fontsize=9, color=colors[5])
    h4.text(5.0, 2.0, legend[3], fontsize=9, color=colors[0])
    h4.text(6.0, 2.0, legend[4], fontsize=9, color=colors[2])
    h4.text(7.0, 2.0, legend[5], fontsize=9, color=colors[4])

    plt.savefig("test_orbitfollowing.png", dpi=300)
    plt.show()

def plot_relerr(axis, x, y, color):
    axis.plot(x, y/y[0] - 1, color, alpha=0.7)


if __name__ == '__main__':
    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if( sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/a5py/testascot/test_orbitfollowing.py
ascot4fusion/ascot5
"""Test Ascot5IO.
"""
import numpy as np
import h5py
import time
import unittest
import datetime
import subprocess
import importlib

from a5py import Ascot
from a5py.exceptions import AscotIOException
from a5py.ascot5io.coreio import fileapi
from a5py.ascot5io.state import State

class TestAscot5IO(unittest.TestCase):
    """Class to test `ascot5io` module.
    """

    def setUp(self):
        """Initialize name for the file used in testing.
        """
        self.testfilename = "unittest.h5"

    def test_io_operations(self):
        """Test Ascot5IO operations that affect the contents of the HDF5 file.
        """

        # Start with an empty object.
        a5 = Ascot(self.testfilename, create=True)
        ETC = {"exyz" : np.array([0,0,0])}

        # Add one group
        grp = a5.data.create_input("E_TC", ETC)

        # Verify it was added to Ascot5IO
        self.assertTrue("efield" in a5.data,
                        "Parent group not found in Ascot5IO.")
        self.assertTrue(grp in a5.data["efield"],
                        "Data group not found in Ascot5IO.")

        # Verify it was added to HDF5.
        with h5py.File(self.testfilename, "r") as h5:
            self.assertTrue("efield" in h5,
                            "Parent group not found in HDF5.")
            self.assertTrue(grp in h5["efield"],
                            "Data group not found in HDF5.")

        # Verify it became the active group.
        self.assertTrue("active" in a5.data["efield"],
                        "No active field in Ascot5IO.")
        self.assertEqual(a5.data["efield"]["active"], a5.data["efield"][grp],
                         "Active group not properly set.")

        # Remove group
        a5.data["efield"][grp].destroy(repack=False)

        # Since it was the only group, the parent should have been removed.
        self.assertFalse("efield" in a5.data,
                        "Parent group not removed from Ascot5IO.")

        # Verify the parent was also removed from the HDF5.
        with h5py.File(self.testfilename, "r") as h5:
            self.assertFalse("efield" in h5,
                            "Parent group not removed from HDF5.")

        # Add three new groups (sleep between so that dates are not same)
        grp1 = a5.data.create_input("E_TC", ETC)
        time.sleep(1.0)
        grp2 = a5.data.create_input("E_TC", ETC)
        time.sleep(1.0)
        grp3 = a5.data.create_input("E_TC", ETC)

        # Groups should be ordered by date
        dates = a5.data["efield"].get_contents()[1]
        self.assertEqual(dates, sorted(dates)[::-1],
                         "Groups are not sorted by date.")

        # First group should be set as active
        self.assertEqual(a5.data["efield"]["active"], a5.data["efield"][grp1],
                         "Active group not properly set.")

        # Switch active
        a5.data["efield"][grp2].activate()
        self.assertEqual(a5.data["efield"]["active"], a5.data["efield"][grp2],
                         "Active group not properly set.")

        # Remove active group
        a5.data["efield"][grp2].destroy(repack=False)

        # Parent group and active field should still be present
        with h5py.File(self.testfilename, "r") as h5:
            self.assertTrue("efield" in h5,
                            "Parent group not found in HDF5.")
        self.assertTrue("efield" in a5.data,
                        "Parent group not found in Ascot5IO.")
        self.assertTrue("active" in a5.data["efield"],
                        "No active field in Ascot5IO.")

        # Group itself should have been removed
        self.assertFalse(grp2 in a5.data["efield"],
                        "Data group not removed from Ascot5IO.")
        with h5py.File(self.testfilename, "r") as h5:
            self.assertFalse(grp2 in h5["efield"],
                            "Data group not removed from HDF5.")

        # The most recent group should be active
        self.assertEqual(a5.data["efield"]["active"], a5.data["efield"][grp3],
                         "Active group not properly set.")

        # Remove parent
        a5.data["efield"].destroy(repack=False)

        # Parent group should have been removed now
        self.assertFalse("efield" in a5.data,
                        "Parent group not removed from Ascot5IO.")
        with h5py.File(self.testfilename, "r") as h5:
            self.assertFalse("efield" in h5,
                            "Parent group not removed from HDF5.")

    def test_metadata(self):
        """Test object metadata.
        """
        # Create group and store its type, QID, and date.
        a5 = Ascot(self.testfilename, create=True)
        ETC = {"exyz" : np.array([0,0,0])}
        date1 = datetime.datetime.strptime(str(datetime.datetime.now())[0:19],
                                           "%Y-%m-%d %H:%M:%S")
        grp = a5.data.create_input("E_TC", ETC)
        qid = grp.split("_")[-1]
        gtype = grp[:-11]
        date2 = datetime.datetime.strptime(str(datetime.datetime.now())[0:19],
                                           "%Y-%m-%d %H:%M:%S")

        # Set description
        desc = "Test"
        a5.data.efield.active.set_desc(desc)

        # Check that the values in HDF5 correspond to ones that were set.
        self.assertEqual(a5.data.efield.active.get_qid(), qid,
                         "QID invalid.")
        self.assertEqual(a5.data.efield.active.get_type(), gtype,
                         "Type invalid.")
        date = datetime.datetime.strptime(a5.data.efield.active.get_date(),
                                          "%Y-%m-%d %H:%M:%S")
        self.assertTrue(date >= date1 and date <=date2,
                         "Date invalid.")
        self.assertEqual(a5.data.efield.active.get_desc(), desc,
                         "Desc invalid.")

    def test_reference(self):
        """Test object referencing in treeview.
        """
        # Create object with multiple groups.
        a5 = Ascot(self.testfilename, create=True)
        ETC = {"exyz" : np.array([0,0,0])}
        grp1 = a5.data.create_input("E_TC", **ETC)
        time.sleep(1.0)
        grp2 = a5.data.create_input("E_TC", **ETC)
        time.sleep(1.0)
        grp3 = a5.data.create_input("E_TC", **ETC)

        # Test with parent group that dict and attribute references are equal
        self.assertEqual(a5.data.efield, a5.data["efield"],
                         "Referencing parent group with dict and attr failed.")

        # Reference by name
        qid   = a5.data.efield[grp1].get_qid()
        gtype = a5.data.efield[grp1].get_type()
        self.assertEqual(grp1, gtype + "_" + qid,
                         "Reference by name failed.")

        # Reference by QID
        self.assertEqual(a5.data.efield[grp1], a5.data.efield["q"+qid],
                         "Reference by QID failed.")

        # Reference by tag: same tags should have running index from newest
        self.assertEqual(a5.data.efield[grp3], a5.data.efield["TAG_0"],
                         "Reference by tag failed.")
        self.assertEqual(a5.data.efield[grp2], a5.data.efield["TAG_1"],
                         "Reference by tag failed.")
        self.assertEqual(a5.data.efield[grp1], a5.data.efield["TAG_2"],
                         "Reference by tag failed.")

        # Reference by tag: same tag running index should update dynamically
        a5.data.efield[grp2].set_desc("test test")
        self.assertEqual(a5.data.efield[grp3], a5.data.efield["TAG_0"],
                         "Reference by tag failed.")
        self.assertEqual(a5.data.efield[grp1], a5.data.efield["TAG_1"],
                         "Reference by tag failed.")

        # Reference by tag: tag is the first word and on uppercase
        self.assertEqual(a5.data.efield[grp2], a5.data.efield["TEST"],
                         "Reference by tag failed.")

        # Reference by tag: special chars omitted
        a5.data.efield[grp2].set_desc("N1C3T^GBA_T")
        self.assertFalse("TEST" in a5.data.efield,
                         "Reference by tag failed.")
        self.assertEqual(a5.data.efield[grp2], a5.data.efield["N1C3TGBAT"],
                         "Reference by tag failed.")

        # Reference by tag: no tag reverts to default and if first char is num.
        a5.data.efield[grp2].set_desc("")
        self.assertEqual(a5.data.efield[grp2], a5.data.efield["TAG_1"],
                         "Reference by tag failed.")
        a5.data.efield[grp2].set_desc("1TAGGY")
        self.assertEqual(a5.data.efield[grp2], a5.data.efield["TAG_1"],
                         "Reference by tag failed.")

    def test_resultnode(self):
        """Test result nodes.
        """
        # Create file and two inputs
        a5 = Ascot(self.testfilename, create=True)
        ETC = {"exyz" : np.array([0,0,0])}
        grp1 = a5.data.create_input("E_TC", ETC)
        time.sleep(1.0)
        grp2 = a5.data.create_input("E_TC", ETC)

        # Create two runs; one for each input
        qid1 = a5.data.efield[grp1].get_qid()
        qid2 = a5.data.efield[grp2].get_qid()
        with h5py.File(self.testfilename, "a") as h5:
            group = fileapi.add_group(h5, "results", "run", desc="RUN1")
            group.attrs["qid_efield"] = np.string_(qid1)
            time.sleep(1.0)
            group = fileapi.add_group(h5, "results", "run", desc="RUN2")
            group.attrs["qid_efield"] = np.string_(qid2)

        a5 = Ascot(self.testfilename)
        run1 = a5.data["RUN1"].get_qid()
        run2 = a5.data["RUN2"].get_qid()

        # Test print functions (that they don't raise errors)
        a5.data.efield.ls()
        a5.data.active.ls()
        a5.data.ls()

        # Run-to-input reference
        self.assertEqual(a5.data["RUN1"].efield.get_qid(), qid1,
                         "Run does not reference input properly.")
        self.assertEqual(a5.data["RUN2"].efield.get_qid(), qid2,
                         "Run does not reference input properly.")

        # Input-to-run reference
        self.assertEqual(a5.data.get_runs(qid1)[0], run1,
                         "Failed to find the run where input was used.")
        self.assertEqual(a5.data.get_runs(qid2)[0], run2,
                         "Failed to find the run where input was used.")

        # One cannot remove input that is used by a run
        with self.assertRaises(AscotIOException,
                               msg="Input was incorrectly removed."):
            a5.data["efield"][grp1].destroy(repack=False)

        # Removing one run should preserve results group
        a5.data["q"+run1].destroy(repack=False)
        self.assertFalse(("q"+run1) in a5.data,
                        "Run group not removed in Ascot5IO.")
        with h5py.File(self.testfilename, "r") as h5:
            self.assertTrue("results" in h5,
                            "Results not found in HDF5.")
            self.assertFalse(run1 in h5["results"],
                            "Run not removed in HDF5.")

        # Removing both runs removes the group as well
        a5.data["q"+run2].destroy(repack=False)
        self.assertFalse("results" in a5.data,
                        "Results not removed in Ascot5IO.")
        with h5py.File(self.testfilename, "r") as h5:
            self.assertFalse("results" in h5,
                            "Results not removed in HDF5.")

        # Create a new run and test the remove all results method
        with h5py.File(self.testfilename, "a") as h5:
            group = fileapi.add_group(h5, "results", "run", desc="RUN1")
            group.attrs["qid_efield"] = np.string_(qid1)

        a5.data.destroy(repack=False)
        self.assertFalse("results" in a5.data,
                        "Results not removed in Ascot5IO.")
        with h5py.File(self.testfilename, "r") as h5:
            self.assertFalse("results" in h5,
                            "Results not removed in HDF5.")

    def test_endcond(self):
        """Test end conditions are parsed properly
        """
        self.assertTrue(State.endcond_check(0x2, "aborted"),
                        "Endcond did not match its binary repr.")
        self.assertFalse(State.endcond_check(0x2, "none"),
                         "Endcond matched wrong binary repr.")
        self.assertFalse(State.endcond_check(0x2, "aborted none"),
                         "Endcond AND operation failed.")
        self.assertFalse(State.endcond_check(0x2, "not aborted"),
                         "Endcond NOT operation failed.")
        self.assertTrue(State.endcond_check(0x2, "aborted not none"),
                        "Endcond AND NOT operation failed.")
        self.assertTrue(State.endcond_check(0x2 | 0x4, "aborted"),
                        "Endcond failed when multiple end conditions active.")
        with self.assertRaises(
                ValueError,
                msg="Failed to raise exception when endcond unknown"):
            State.endcond_check(0x2 | 0x4, "vanished")

    def test_inputs(self):
        inputs = {
            "bfield"  : ["B_TC", "B_GS", "B_2DS", "B_3DS", "B_3DST", "B_STS",],
            "efield"  : ["E_TC", "E_1DS", "E_3D", "E_3DS", "E_3DST",],
            "marker"  : ["Prt", "GC", "FL",],
            "wall"    : ["wall_2D", "wall_3D",],
            "plasma"  : ["plasma_1D", "plasma_1DS", "plasma_1Dt"],
            "neutral" : ["N0_3D",],
            "boozer"  : ["Boozer",],
            "mhd"     : ["MHD_STAT", "MHD_NONSTAT",],
            "asigma"  : ["Asigma_loc"],
            "options" : ["Opt",],
            #"nbi"     : ["NBI",],
            }
        a5 = Ascot(self.testfilename, create=True)
        for parent, inp in inputs.items():
            mod = importlib.import_module("a5py.ascot5io." + parent)

            for grp in inp:
                fun = getattr(getattr(mod, grp), "write_hdf5_dummy")
                fun(self.testfilename)
                a5 = Ascot(self.testfilename)
                data = a5.data[parent].DUMMY.read()

                fun = getattr(getattr(mod, grp), "write_hdf5")
                fun(fn=self.testfilename, desc="DUMMY2", **data)
                a5 = Ascot(self.testfilename)
                data = a5.data[parent].DUMMY2.read()
                a5.data[parent].DUMMY.destroy()
                a5.data[parent].DUMMY2.destroy()

    def tearDown(self):
        """Remove the file used in testing.
        """
        subprocess.run(["rm", "-f", self.testfilename])

if __name__ == '__main__':
    unittest.main()

a5py/a5py/testascot/unit_ascot5io.py
ascot4fusion/ascot5
"""Unit tests for most post-processing routines.
"""
import numpy as np
import unittest
import subprocess
import unyt

from a5py import Ascot, AscotInitException, AscotIOException

class TestAscot(unittest.TestCase):
    """Class for testing :class:`Ascot` object.
    """

    @classmethod
    def setUpClass(cls):
        """Create and run a test case.

        Assumes ascot5_main is located in the same folder.
        """
        super(TestAscot, cls).setUpClass()
        a5 = Ascot("unittest.h5", create=True)
        a5.data.create_input("options tutorial")
        a5.data.create_input("bfield analytical iter circular", splines=False)
        a5.data.create_input("wall rectangular")
        a5.data.create_input("plasma flat")

        from a5py.ascot5io.marker import Marker
        mrk = Marker.generate("gc", n=100, species="alpha")
        mrk["energy"][:] = 3.5e6
        #mrk["pitch"][:]  = 0.99 - 1.98 * np.random.rand(100,)
        mrk["pitch"][:]  = -0.99 - 0 * np.random.rand(100,)
        mrk["r"][:]      = np.linspace(6.2, 8.2, 100)
        a5.data.create_input("gc", **mrk)
        a5.data.create_input("E_TC", exyz=np.array([0,0,0]))
        a5.data.create_input("E_TC", exyz=np.array([0,0,0]), desc="UNUSED")

        a5.data.create_input("N0_3D")
        a5.data.create_input("Boozer")
        a5.data.create_input("MHD_STAT")
        a5.data.create_input("asigma_loc")

        name = a5.data.options.active.new(
            ENDCOND_MAX_MILEAGE=0.5e-4, DIST_MIN_CHARGE=1.5,
            DIST_MAX_CHARGE=2.5, DIST_NBIN_PPE=10, DIST_NBIN_PPA=20,
            DIST_NBIN_PHI=3, DIST_NBIN_R=10, DIST_NBIN_Z=10, DIST_NBIN_RHO=10,
            desc="Fast")
        a5.data.options[name].activate()

        subprocess.run(["./ascot5_main", "--in=unittest.h5"])

    @classmethod
    def tearDownClass(cls):
        super(TestAscot, cls).tearDownClass()
        subprocess.run(["rm", "-f", "unittest.h5"])

    def _test_initandsim(self):
        """Test :class:`Ascotpy` initialization and simulation routines.
        """
        a5 = Ascot("unittest.h5")
        a5.data.active.efield.activate()
        allinps = ["bfield", "efield", "plasma", "wall", "neutral", "boozer",
                    "mhd", "asigma"]

        # Test basic init and free
        a5.input_init()
        inps = a5.input_initialized()
        self.assertCountEqual(inps.keys(), allinps,
                              "Not all inputs were initialized.")
        a5.input_free()
        inps = a5.input_initialized()
        self.assertTrue(len(inps) == 0,
                        "Not all inputs were free'd.")

        # Test explicit init and free
        a5.input_init(efield=True)
        inps = a5.input_initialized()
        self.assertCountEqual(inps.keys(), ["efield"],
                              "Explicit initialization failed.")
        a5.input_free(bfield=True)
        inps = a5.input_initialized()
        self.assertCountEqual(inps.keys(), ["efield"],
                              "Explicit free failed.")
        a5.input_free(efield=True)
        inps = a5.input_initialized()
        self.assertTrue(len(inps) == 0,
                        "Explicit free failed.")

        # Test init and free with multiple inputs of same type
        a5.input_init(efield=True)
        inps = a5.input_initialized()
        self.assertEqual(inps["efield"], a5.data.efield.active.get_qid(),
                         "Initialized input was not the active one.")

        try:
            a5.input_init(efield=True)
        except AscotInitException:
            self.fail("Re-initializing same input failed.")

        with self.assertRaises(
                AscotInitException,
                msg="Initializing occupied input did not raise exception."):
            a5.input_init(efield=a5.data.efield["UNUSED"].get_qid())

        with self.assertRaises(
                AscotIOException,
                msg="Initializing nonexistent input did not raise exception."):
            a5.input_init(efield="0123456789")

        try:
            a5.input_init(efield=a5.data.efield["UNUSED"].get_qid(),
                          switch=True)
        except AscotInitException:
            self.fail("Switch did not work.")

        a5.input_free(efield=True)
        inps = a5.input_initialized()
        self.assertTrue(len(inps) == 0,
                        "Not all inputs were free'd.")

        # Double free should be ok
        a5.input_free(bfield=True)
        a5.input_free()

        # Test init from run
        a5.data.efield["UNUSED"].activate()
        a5.input_init(run=True, efield=True)
        inps = a5.input_initialized()
        self.assertEqual(inps["efield"], a5.data.active.efield.get_qid(),
                         "Initialized input was not the correct one.")
        a5.input_free(efield=True)

        # Give run QID, init all inputs
        a5.input_init(run=a5.data.active.get_qid())
        inps = a5.input_initialized()
        self.assertEqual(inps["efield"], a5.data.active.efield.get_qid(),
                         "Initialized input was not the correct one.")
        a5.input_free()
        a5.data.active.efield.activate()

        # Test packing
        a5.simulation_initinputs()
        with self.assertRaises(
                AscotInitException,
                msg="Trying to init packed data should raise an exception."):
            a5.input_init(efield=a5.data.efield["UNUSED"].get_qid())

        # Test marker and options initialization
        opt = a5.data.options.active.read()
        opt.update({"ENDCOND_MAX_MILEAGE" : 1e-10})
        a5.simulation_initoptions(**opt)

        mrk = a5.data.marker.active.read()
        a5.simulation_initmarkers(**mrk)

        # Run simulation and test output initialization
        a5.simulation_run(printsummary=False)

        # Verify that simulations can't be run when data is occupied
        with self.assertRaises(
                AscotInitException,
                msg="Trying to run with occupied output should cause an error"):
            a5.simulation_run(printsummary=False)

        # This should run
        a5.simulation_free(diagnostics=True)
        a5.simulation_run(printsummary=False)
        a5.simulation_free(diagnostics=True)

        # Can't run when input is missing
        a5.simulation_free(markers=True, diagnostics=True)
        with self.assertRaises(
                AscotInitException,
                msg="Trying to run without markers should cause an error"):
            a5.simulation_run(printsummary=False)

        # Init markers but free inputs
        a5.simulation_initmarkers(**mrk)
        a5.simulation_free(inputs=True)
        with self.assertRaises(
                AscotInitException,
                msg="Trying to run without input should cause an error"):
            a5.simulation_run(printsummary=False)

        # Free everyting
        a5.simulation_free()
        with self.assertRaises(
                AscotInitException,
                msg="Trying to run without input should cause an error"):
            a5.simulation_run(printsummary=False)

    def _test_eval(self):
        """Test :class:`Ascotpy` evaluation routines.
        """
        a5 = Ascot("unittest.h5")

        # Input evaluations
        a5.input_init()
        inputqnts = a5.input_eval_list(show=False)
        r = 6.2 * unyt.m; phi = 0 * unyt.deg; z = 0 * unyt.m; t = 0 * unyt.s
        br, bphi, bz = a5.input_eval(r, phi, z, t, "br", "bphi", "bz")
        for q in inputqnts:
            out = a5.input_eval(r, phi, z, t, q)
        a5.input_free()

        # State evaluations
        a5.input_init(bfield=True)
        e, v, m = a5.data.active.getstate(
            "ekin", "vpar", "mu", endcond="tlim", ids=1)
        self.assertTrue(
            e.units == unyt.eV and v.units == unyt.m/unyt.s and \
            m.units == unyt.eV/unyt.T and e.size==1 and v.size==1 and \
            m.size==1,
            "State evaluation failed")
        outputqnts = a5.data.active.getstate_list()
        for q in outputqnts:
            out = a5.data.active.getstate(q)

        # Orbit evaluations
        e, v, m = a5.data.active.getorbit(
            "ekin", "vpar", "mu", endcond="tlim", ids=1)
        self.assertTrue(
            e.units == unyt.eV and v.units == unyt.m/unyt.s and \
            m.units == unyt.eV/unyt.T,
            "Orbit evaluation failed")
        outputqnts = a5.data.active.getorbit_list()
        for q in outputqnts:
            out = a5.data.active.getorbit(q)
        a5.input_free()

        # Packing should not prevent input evaluation
        a5.simulation_initinputs()
        out = a5.input_eval(r, phi, z, t, "bnorm")

        # Prepare simulation data
        opt = a5.data.options.active.read()
        opt.update({"ENDCOND_MAX_MILEAGE" : 1e-7})
        a5.simulation_initoptions(**opt)
        mrk = a5.data.marker.active.read()
        a5.simulation_initmarkers(**mrk)
        vrun = a5.simulation_run(printsummary=False)

        # Virtual state evaluations
        e, v, m = vrun.getstate(
            "ekin", "vpar", "mu", endcond="tlim", ids=1)
        self.assertTrue(
            e.units == unyt.eV and v.units == unyt.m/unyt.s and \
            m.units == unyt.eV/unyt.T and e.size==1 and v.size==1 and \
            m.size==1,
            "Virtual state evaluation failed")
        outputqnts = a5.data.active.getstate_list()
        for q in outputqnts:
            out = vrun.getstate(q)

        # Virtual orbit evaluations
        out = vrun.getorbit("r", "phi", "z")
        e, v, m = vrun.getorbit(
            "ekin", "vpar", "mu", endcond="tlim", ids=1)
        self.assertTrue(
            e.units == unyt.eV and v.units == unyt.m/unyt.s and \
            m.units == unyt.eV/unyt.T,
            "Virtual orbit evaluation failed")
        outputqnts = vrun.getorbit_list()
        for q in outputqnts:
            out = a5.data.active.getorbit(q)
        a5.simulation_free()

    def test_dist(self):
        """Test distribution postprocessing.
        """
        a5 = Ascot("unittest.h5")

        # Test volume calculation
        R,Z = np.meshgrid(np.linspace(3,7,101), np.linspace(-2,2,101))
        psi = np.sqrt((R-5)**2 + (Z-0)**2)
        bdata = {"rmin":3, "rmax":7, "nr":101, "zmin":-2, "zmax":2, "nz":101,
                 "axisr":5, "axisz":0, "psi":psi, "psi0":0, "psi1":1,
                 "br":psi*0, "bphi":psi*0, "bz":psi*0}
        qid = a5.data.create_input("B_2DS", **bdata).split("_")[-1]

        try:
            a5.input_init(bfield=qid)
            vol1, area1, r1, p1, z1 = a5.input_rhovolume(
                method="prism", nrho=5, ntheta=4, nphi=3, return_area=True,
                return_coords=True)
            vol2, area2, r2, p2, z2 = a5.input_rhovolume(
                method="mc", nrho=5, ntheta=4, nphi=3, return_area=True,
                return_coords=True)
            a5.input_free()
        except:
            a5.input_free()
            raise

        print(r1.shape, r2.shape)
        print(p1.shape, p2.shape)
        print(z1.shape, z2.shape)
        delta = 1
        vol0 = 2*np.pi*5*np.pi*1.0**2
        print(
            """
            Volume calculation:
            %3.3f %3.3f %3.3f
            """ % (np.sum(vol1).v, np.sum(vol2).v, vol0)
        )
        self.assertAlmostEqual(
            np.sum(vol1).v, vol0, None, "Volume calculation with prism failed",
            delta)
        self.assertAlmostEqual(
            np.sum(vol2).v, vol0, None, "Volume calculation with MC failed",
            delta)
        area0 = np.pi*1.0**2
        print(
            """
            Area calculation:
            %3.3f %3.3f %3.3f
            """ % (np.sum(area1[:,:,0]).v, np.sum(area2[:,:,0]).v, area0)
        )
        self.assertAlmostEqual(
            np.sum(area1[:,:,0]).v, area0, None,
            "Area calculation with prism failed", delta)
        self.assertAlmostEqual(
            np.sum(area2[:,:,0]).v, area0, None,
            "Area calculation with MC failed", delta)

        # Test distribution initializing
        a5.data.active.getdist("5d")
        a5.data.active.getdist("rho5d")
        a5.data.active.getdist("6d")
        a5.data.active.getdist("rho6d")
        a5.data.active.getdist("com")

        # Test slicing, interpolating, and integrating
        dist = a5.data.active.getdist("5d")
        dist.integrate(charge=np.s_[:], time=np.s_[:])
        self.assertFalse("charge" in dist.abscissae, "Integration failed")
        self.assertFalse("time"   in dist.abscissae, "Integration failed")

        dist.slice(ppar=np.s_[0], pperp=np.s_[0])
        self.assertTrue(dist.abscissa("ppar").size == 1, "Slicing failed")
        self.assertTrue(dist.abscissa("pperp").size == 1, "Slicing failed")

        out = dist.interpolate(phi=2*unyt.deg, r=6.2*unyt.m)
        self.assertEqual(out.size, dist.abscissa("z").size,
                         "Integration/slicing/interpolation failed")

        # Get 5D distribution in both momentum spaces and verify the number of
        # particles is same
        dist    = a5.data.active.getdist("5d")
        exidist = a5.data.active.getdist("5d", exi=True, ekin_edges=10,
                                         pitch_edges=20, plotexi=False)
        time, weight = a5.data.active.getstate("mileage", "weight", state="end")
        self.assertAlmostEqual(np.sum(dist.histogram().v),
                               np.sum(time*weight).v, None,
                               "Distribution not valid",
                               0.001)
        self.assertAlmostEqual(np.sum(exidist.histogram().v),
                               np.sum(time*weight).v, None,
                               "Converting ppar-pperp to ekin-pitch failed",
                               0.001)

        a5.input_init(bfield=True)
        rhodist    = a5.data.active.getdist("rho5d")
        rhoexidist = a5.data.active.getdist("rho5d", exi=True)
        mom1 = a5.data.active.getdist_moments(dist, "density")
        mom2 = a5.data.active.getdist_moments(exidist, "density")
        mom3 = a5.data.active.getdist_moments(rhodist, "density")
        mom4 = a5.data.active.getdist_moments(rhoexidist, "density")
        a5.input_free()
        print(np.sum(mom1.ordinate("density") * mom1.volume),
              np.sum(mom2.ordinate("density") * mom2.volume),
              np.sum(mom3.ordinate("density") * mom3.volume),
              np.sum(mom4.ordinate("density") * mom4.volume))

        self.assertAlmostEqual(np.sum(mom1.ordinate("density") * mom1.volume).v,
                               np.sum(time*weight).v, None,
                               "Converting ppar-pperp to ekin-pitch failed",
                               0.001)
        self.assertAlmostEqual(np.sum(mom2.ordinate("density") * mom2.volume).v,
                               np.sum(time*weight).v, None,
                               "Converting ppar-pperp to ekin-pitch failed",
                               0.001)
        self.assertAlmostEqual(np.sum(mom3.ordinate("density") * mom3.volume).v,
                               np.sum(time*weight).v, None,
                               "Converting ppar-pperp to ekin-pitch failed",
                               0.001)
        self.assertAlmostEqual(np.sum(mom4.ordinate("density") * mom4.volume).v,
                               np.sum(time*weight).v, None,
                               "Converting ppar-pperp to ekin-pitch failed",
                               0.001)
if __name__ == '__main__':
    unittest.main()

a5py/a5py/testascot/unit_ascotpy.py
ascot4fusion/ascot5
"""Unit tests for distribution moments.
"""
import numpy as np
import unittest
import subprocess
import unyt

from a5py import Ascot, AscotInitException, AscotIOException

class TestAscot(unittest.TestCase):
    """Class for testing distribution moments.
    """

    @classmethod
    def setUpClass(cls):
        """Create and run a test case.

        Assumes ascot5_main is located in the same folder.
        """
        super(TestAscot, cls).setUpClass()
        a5 = Ascot("unittest.h5", create=True)
        a5.data.create_input("options tutorial")
        a5.data.create_input("bfield analytical iter circular")
        a5.data.create_input("wall rectangular")
        a5.data.create_input("plasma flat")

        from a5py.ascot5io.marker import Marker
        mrk = Marker.generate("gc", n=1, species="alpha")
        mrk["energy"][:] = 3.5e6
        mrk["pitch"][:]  = 0.7
        mrk["r"][:]      = 7.8
        a5.data.create_input("gc", **mrk)
        a5.data.create_input("E_TC", exyz=np.array([0,0,0]))

        a5.data.create_input("N0_3D")
        a5.data.create_input("Boozer")
        a5.data.create_input("MHD_STAT")
        a5.data.create_input("asigma_loc")

        name = a5.data.options.active.new(
            ENDCOND_MAX_MILEAGE=1.5e-3, DIST_NBIN_TIME=1, DIST_MIN_CHARGE=1.5,
            DIST_MAX_CHARGE=2.5, DIST_NBIN_CHARGE=1, DIST_MIN_R=4, DIST_MAX_R=8,
            DIST_MIN_Z=-2, DIST_MAX_Z=2,
            ENABLE_DIST_6D=0, ENABLE_DIST_RHO6D=1, ENABLE_DIST_COM=0,
            DIST_NBIN_PHI=1, DIST_NBIN_R=50, DIST_NBIN_Z=100,
            DIST_NBIN_PPE=50, DIST_NBIN_PPA=50,
            DIST_NBIN_THETA=2,  DIST_NBIN_RHO=10, ENABLE_COULOMB_COLLISIONS=1,
            ORBITWRITE_NPOINT=60000)
        a5.data.options[name].activate()

        subprocess.run(["./ascot5_main", "--in=unittest.h5"])

    @classmethod
    def tearDownClass(cls):
        super(TestAscot, cls).tearDownClass()
        subprocess.run(["rm", "-f", "unittest.h5"])

    def test_moments(self):
        a5 = Ascot("unittest.h5")
        #ordinates = ["density", "chargedensity", "energydensity", "pressure",
        #             "toroidalcurrent", "parallelcurrent", "powerdep",
        #             "electronpowerdep", "ionpowerdep"]
        ordinates = ["jxbtorque"]

        a5.input_init(bfield=True, plasma=True)
        dist    = a5.data.active.getdist("5d")
        mom     = a5.data.active.getdist_moments(dist, *ordinates)
        rhodist = a5.data.active.getdist("rho5d")
        rhomom  = a5.data.active.getdist_moments(rhodist, *ordinates)

        re = dist.abscissa_edges("r")
        ze = dist.abscissa_edges("z")
        v1, v2 = np.meshgrid(re[1:]**2 - re[:-1]**2, ze[1:] - ze[:-1])
        vol = np.pi * v1 * v2

        rho, r, z, phi, psi, weight, time,  \
        charge, energy, vphi, vnorm, vpar,  \
        mass, p, pitch, bnorm, bphi, ppar = \
            a5.data.active.getorbit("rho", "r", "z", "phi", "psi", "weight",
                                    "mileage", "charge", "ekin", "vphi",
                                    "vnorm", "vpar", "mass", "pnorm", "pitch",
                                    "bnorm", "bphi", "ppar")
        ei, psii, Pphii = a5.data.active.getstate("ekin", "psi", "pphi",
                                                  state="ini")
        ef, tf = a5.data.active.getstate("ekin", "mileage", state="end")
        dt = np.diff(time, prepend=0)

        k, nu = a5.input_eval_collcoefs(
            mass[0], charge[0], r, phi, z, time, vnorm, "k", "nu",
            grid=False)
        dEtot_d = p * dt * np.sum(k, axis=0)
        dEele_d = p * dt * k[0,:]
        dEion_d = p * dt * np.sum(k[1:,:], axis=0)
        dpsi    = np.zeros(psi.shape) * psi.units
        dpsi[0] = psi[0] - psii
        dpsi[1:] = psi[1:] - psi[:-1]
        nu      = np.sum(nu, axis=0)
        k       = np.sum(k, axis=0)
        dppar   = mass * k * pitch - p * pitch * nu
        Pphi    = ppar * r * (bphi/bnorm) + charge * psi
        dPphi   = np.diff(Pphi, prepend=Pphi[0])

        a5.input_free()

        mrkdist = {}
        mrkdist["density"]          = weight * dt
        mrkdist["chargedensity"]    = weight * charge * dt
        mrkdist["energydensity"]    = weight * ( energy * dt ).to("J")
        mrkdist["pressure"]         = weight * (mass * vnorm**2 * dt).to("J")/3
        mrkdist["toroidalcurrent"]  = weight * ( charge * vphi * dt ).to("A*m")
        mrkdist["parallelcurrent"]  = weight * ( charge * vpar * dt ).to("A*m")
        mrkdist["powerdep"]         = weight * dEtot_d.to("W")
        mrkdist["electronpowerdep"] = weight * dEele_d.to("W")
        mrkdist["ionpowerdep"]      = weight * dEion_d.to("W")
        mrkdist["jxbtorque"]        = weight * (-charge * dpsi/unyt.s).to("N*m")
        mrkdist["colltorque"]       = weight * (r*dppar*(bphi/bnorm)*dt).to("J")
        mrkdist["canmomtorque"]     = weight * -charge * dPphi

        print(weight[0]*tf)
        print(((ef-ei)*weight[0]/unyt.s).to("W"))
        for o in ordinates:
            a1 = np.sum(rhomom.ordinate(o) * rhomom.volume)
            a2 = np.sum(mom.ordinate(o) * mom.volume)
            a3 = np.sum(mrkdist[o])
            print(o, a1, a2, a3)
        """
        for w in mrkdist:
            prt = np.histogram2d(
                r, z, [dist.abscissa_edges("r"), dist.abscissa_edges("z")],
                weights=w)[0]
            rhoprt = np.histogram(
                rho, rhodist.abscissa_edges("rho"), weights=w)[0]

        import matplotlib.pyplot as plt
        fig = plt.figure()
        ax1 = fig.add_subplot(1,3,1)
        plt.title("Test")
        pm = ax1.pcolormesh(
            dist.abscissa_edges("r").v, dist.abscissa_edges("z").v,
            prt.T / vol, shading="flat")
        plt.colorbar(pm)
        ax1.set_aspect("equal", adjustable="box")
        ax2 = fig.add_subplot(1,3,2)
        plt.title("Ascot <3")
        a5.data.active.plotdist_moments(mom, ordinate, axes=ax2)

        ax3 = fig.add_subplot(1,3,3)
        pm = ax3.plot(rhodist.abscissa("rho").v,
                      rhoprt / np.sum(np.sum(rhomom.volume, axis=1), axis=1))
        a5.data.active.plotdist_moments(rhomom, ordinate, axes=ax3)
        plt.show()
        """

if __name__ == '__main__':
    unittest.main()

a5py/a5py/testascot/unit_moments.py
ascot4fusion/ascot5
"""Test plotting routines.

The plotting routines obtain the data using routines that are tested elsewhere,
so here we can focus on making sure that the plots look correct. These tests
only fail if there was an exception but otherwise the actual "testing" is to
look at the figures.
"""
import numpy as np
import unittest
import subprocess
import matplotlib.pyplot as plt
import unyt

from a5py import Ascot
import a5py.routines.plotting as a5plt
from a5py.ascot5io.marker import Marker

class TestPlotting(unittest.TestCase):

    def tearDown(self):
        try:
            a5 = Ascot("unittest.h5")
            a5.simulation_free()
        except:
            pass
        subprocess.run(["rm", "-f", "unittest.h5"])

    def initandrunsim(self, simtime):
        a5 = Ascot("unittest.h5", create=True)
        a5.data.create_input("options tutorial")
        a5.data.create_input("bfield analytical iter circular")
        a5.data.create_input("wall rectangular")
        a5.data.create_input("plasma flat")
        a5.data.create_input("E_TC", exyz=np.array([0,0,0]))
        a5.data.create_input("N0_3D")
        a5.data.create_input("Boozer")
        a5.data.create_input("MHD_STAT")
        a5.data.create_input("asigma_loc")

        a5.simulation_initinputs()

        mrk = Marker.generate("gc", n=100, species="alpha")
        mrk["weight"][:] = 10**(10 +10*np.random.rand(100,))
        mrk["energy"][:] = 3.5e6
        mrk["pitch"][:]  = 0.99 - 1.98 * np.random.rand(100,)
        mrk["r"][:]      = np.linspace(6.2, 8.5, 100)
        a5.simulation_initmarkers(**mrk)

        opt = a5.data.options.active.read()
        opt.update(ENDCOND_MAX_MILEAGE=simtime)
        a5.simulation_initoptions(**opt)
        return a5.simulation_run()

    def test_input_plotrz(self):
        a5 = Ascot("unittest.h5", create=True)
        a5.data.create_input("bfield analytical iter circular", splines=True,
                             axisymmetric=True)
        a5.data.create_input("plasma flat")
        a5.input_init(bfield=True, plasma=True)

        fig = plt.figure(figsize=(30,40))

        r = np.linspace( 3.0, 9.0,  50) * unyt.m
        z = np.linspace(-2.0, 2.0, 100) * unyt.m

        ax = fig.add_subplot(2,4,1)
        a5.input_plotrz(r, z, "bnorm", clim=[None, None], cmap=None, axes=ax)

        ax = fig.add_subplot(2,4,5)
        a5.input_plotrz(r, z, "bnorm", clim=[0, 10], cmap="Greys", axes=ax)

        ax = fig.add_subplot(2,4,2)
        a5.input_plotrz(r, z, "br", clim=[None, None], cmap=None, axes=ax)

        ax = fig.add_subplot(2,4,6)
        a5.input_plotrz(r, z, "br", clim=[-1, 1], cmap="PuOr", axes=ax)

        ax = fig.add_subplot(2,4,3)
        a5.input_plotrz(r, z, "log ne", clim=[None, None], cmap=None, axes=ax)

        ax = fig.add_subplot(2,4,7)
        a5.input_plotrz(r, z, "log ne", clim=[1e10, 1e22],
                        cmap="Greys", axes=ax)

        ax = fig.add_subplot(2,4,4)
        a5.input_plotrz(r, z, "log divergence", clim=[None, None],
                        cmap=None, axes=ax)

        ax = fig.add_subplot(2,4,8)
        a5.input_plotrz(r, z, "log divergence",clim=[-1e-14, 1e-14],
                        cmap="PuOr",axes=ax)

        a5.input_free()
        plt.show(block=False)

    def test_plotstate_scatter(self):
        vr = self.initandrunsim(1e-3)
        fig = plt.figure(figsize=(30,40))

        ax = fig.add_subplot(2,3,1)
        vr.plotstate_scatter("end rho", "log end ekin",
                             axesequal=False, axes=ax)

        ax = fig.add_subplot(2,3,2)
        vr.plotstate_scatter("ini rho", "log diff ekin",
                             axesequal=False, axes=ax)

        ax = fig.add_subplot(2,3,3)
        vr.plotstate_scatter("log end x", "log end y", c="log diff phi",
                             axesequal=True, axes=ax)

        ax = fig.add_subplot(2,3,4, projection="3d")
        vr.plotstate_scatter("end x", "end y", "end z",
                             axesequal=True, axes=ax)

        ax = fig.add_subplot(2,3,5, projection="3d")
        vr.plotstate_scatter("end x", "end y", "end z", c="end rho",
                             axesequal=True, axes=ax)

        ax = fig.add_subplot(2,3,6, projection="3d")
        vr.plotstate_scatter("log end x", "log end y", "log end z",
                             c="log diff phi", axesequal=True, axes=ax)

        plt.show(block=False)

    def test_plotstate_histogram(self):
        vr = self.initandrunsim(1e-4)
        fig = plt.figure(figsize=(30,40))

        ax = fig.add_subplot(2,3,1)
        vr.plotstate_histogram("end phimod", axes=ax)

        ax = fig.add_subplot(2,3,2)
        vr.plotstate_histogram("log end phi", xbins=10, endcond="WALL", axes=ax)

        ax = fig.add_subplot(2,3,3)
        vr.plotstate_histogram("end ekin", xbins=np.linspace(0,4e6,100),
                               logscale=True, weight=True, axes=ax)

        ax = fig.add_subplot(2,3,4)
        vr.plotstate_histogram("end phimod", "end thetamod", axes=ax)

        ax = fig.add_subplot(2,3,5)
        vr.plotstate_histogram("log end phi", "log end theta", endcond="WALL",
                               axes=ax)

        ax = fig.add_subplot(2,3,6)
        vr.plotstate_histogram("end phimod", "end thetamod", logscale=True,
                               weight=True, axes=ax)

        plt.show(block=False)

    def test_plotorbit_trajectory(self):
        vr = self.initandrunsim(1e-4)
        fig = plt.figure(figsize=(30,40))

        ax = fig.add_subplot(2,3,1)
        vr.plotorbit_trajectory("r", "z",
                                axesequal=True, axes=ax)

        ax = fig.add_subplot(2,3,2)
        vr.plotorbit_trajectory("mileage", "reldiff ekin", c="rho",
                                axesequal=False, axes=ax)

        ax = fig.add_subplot(2,3,3)
        vr.plotorbit_trajectory("log mileage", "log reldiff ekin",
                                c="log rho",
                                axesequal=False, axes=ax)

        ax = fig.add_subplot(2,3,4, projection="3d")
        vr.plotorbit_trajectory("x", "y", "z",
                                axesequal=True, axes=ax)

        ax = fig.add_subplot(2,3,5, projection="3d")
        vr.plotorbit_trajectory("x", "y", "z", c="diff ekin",
                                axesequal=True, axes=ax)

        ax = fig.add_subplot(2,3,6, projection="3d")
        vr.plotorbit_trajectory("log x", "log y", "log z", c="log diff ekin",
                                axesequal=True, axes=ax)

        plt.show(block=False)

if __name__ == '__main__':
    unittest.main()

a5py/a5py/testascot/unit_plotting.py
ascot4fusion/ascot5
from .loads import loads, figuresofmerit, loadvsarea

a5py/a5py/wall/__init__.py
ascot4fusion/ascot5
'''
This module is the interface to VTK for walls.

Created on Mar 6, 2020

@author: sjjamsa
'''

import vtk
import numpy as np



class a5VtkWall(object):
    '''
    This class represents an ASCOT5 wall in VTK polydata datatype.
    The data can include zero or more scalar values per polygon.
    '''


    def __init__(self, W=None):
        '''
        @param W: (optional) The wall_3D class:  a5py.ascot5io.ascot5.Ascot(filename).active.wall
        '''
        
        if W is not None:
            self.fromA5wall(W)
            return
    
        
    def fromA5wall(self, W):
        '''
        @param W: The wall_3D class:  a5py.ascot5io.ascot5.Ascot(filename).active.wall
        '''
        
        points,vertices  = W.getAspointsAndVertices()
        
        self.fromPointsAndVertices(points,vertices)
    
    def fromPointsAndVertices(self,points,vertices):
        '''
        @param points: nx3 array of triangle corners; a list of (x,y,z)-tuples
        @param vertices: nx3 array of 0-based indexes in the points list 
        '''
        
        
        print('Creating VTK points from triangle vertices.')

        
        nTriangles = vertices.shape[0]
        nPoints    = points.shape[0]
        print('Points: {}, Triangles {}'.format(nPoints,nTriangles))
        
        Points = vtk.vtkPoints()
        Points.SetDataTypeToFloat()
        #Points.SetDataTypeToDouble()
        Points.SetNumberOfPoints(nPoints) 
        for i,P in enumerate(points):
            Points.SetPoint(i,P[0], P[1], P[2])



            
        
        print('Creating VTK triangles referencing the points.')
        
        
        Triangles = vtk.vtkCellArray()
        
        # Prepare a connectivity array
        Cells= vtk.vtkIdTypeArray()
        Cells.SetNumberOfTuples(4*nTriangles)
        
        # Fill the array with nVertices, vert1, vert2, ..., vertn, nVertices, vert1, vert2, ...        
        for iTri in range(nTriangles):
            ind=iTri*4
            Cells.SetValue(ind,3) # triangle (3 vertices)
            for iP in range(3):
                ind = iTri*4+iP+1
                Cells.SetValue(ind,vertices[iTri,iP])


        Triangles.SetCells(   nTriangles,   Cells ) 

        
        print('Combining the points and triangles to a single VTK PolyData datastructure.')
        
            
        Polydata = vtk.vtkPolyData()
        Polydata.SetPoints(Points)
        Polydata.SetPolys(Triangles)
        
        Polydata.Modified()
        
        self.nTriangles = nTriangles
        self.vtkPolyData = Polydata

    def addFlag(self,flag,flagIdList,flagIdStrings,setAsActive=None):
        
        fieldname = 'Triangle Flag'
        
        # Save these for future need.
        self.flagIdList    = flagIdList
        self.flagIdStrings = flagIdStrings
        
        print('\nFlag descriptions:\n-------------')
        for i,s in enumerate(flagIdStrings):
            print('{:3} {:12}'.format(i,s))
        print('--')

        if setAsActive is None:
            setAsActive =  len(np.unique(flag)) > 1
            
        self.addScalarField(fieldname, flag, setAsActive=setAsActive)

    
    def addIndex(self,setAsActive=True):
        
        fieldname = 'Triangle index'
        data=np.arange(0,self.nTriangles,1,dtype=float)
        self.addScalarField(fieldname, data, setAsActive=setAsActive)
    
    def addScalarField(self,fieldname,data,setAsActive=True):
        '''
        @param fieldname: A string describing the data
        @param data: length-n array of values for coloring the triangles 
        '''

        print('Creating a VTK array representing "'+fieldname+'" and attaching it to the VTK PolyData.')
        
     
        
        field=vtk.vtkDoubleArray()
        field.SetNumberOfTuples(self.nTriangles)
        field.SetNumberOfComponents(1)
        field.SetName(fieldname)
        for i,l in enumerate(data):
            field.SetTuple(i,[l])
            
        #self.vtkPolyData.GetCellData().SetScalars(field) # This one replaces.
        self.vtkPolyData.GetCellData().AddArray(field)
        if setAsActive:
            self.vtkPolyData.GetCellData().SetActiveScalars(fieldname);
        self.vtkPolyData.Modified()
        
        
        
    def writeVtp(self, filename):
        print('Writing "' +filename+'".')
        
        writer = vtk.vtkXMLPolyDataWriter()
        writer.SetFileName(filename)
        writer.SetInputData(self.vtkPolyData)
        writer.Write()
        


    def readVtp(self, filename):
        print('Reading "'+filename+'".')
        reader = vtk.vtkXMLPolyDataReader()
        reader.SetFileName(filename)
        reader.Update()
        self.vtkPolyData = reader.GetOutput()


    def plot(self,manual_range=None,logarithmicColorScale=True,camControl=None,
             invertColormap=True,replaceBottom=True):
        '''
        @param manual_range: (opt) length two vector/tuple defining the coloraxis range as 10**manual_range[0] -- 10**manual_range[0] 
        @param camControl: (opt) an instance of a5VtkWall.camControl. 
               It will be associated with the new camera in the plot. The requested camera settings should be set before calling this routine.
        '''

        output = self.vtkPolyData

        # GetScalars() returns the active scalar field.
        title = output.GetCellData().GetScalars().GetName()
        print('Plotting "'+title+'".')

        scalar_range = output.GetScalarRange()

        print('Original color value range {} - {}'.format(scalar_range[0],scalar_range[1]))
        if manual_range is not None:
            #manual_range=(2,7)
            scalar_range = (10.0**manual_range[0], 10.0**manual_range[1])
            #print(scalar_range)

        mapper = vtk.vtkPolyDataMapper()
        mapper.SetInputData(output)

        mapper.SetScalarRange(scalar_range)


        '''
        Get the vtkPolyData -object. This could be useful for e.g.
        as input for
            A5vtkw = a5VtkWall(W)
            Mapper = vtk.vtkPolyDataMapper()
            Mapper.SetInputData(A5vtkw.getVtkPolyData())
        '''

        lut = vtk.vtkLookupTable()
        if logarithmicColorScale:
            lut.SetScaleToLog10()
        lut.Build()
        # Invert the colormap
        if invertColormap:
            nColors=lut.GetNumberOfTableValues()
            colors=list( range( nColors ) )
            for iColor in range( nColors ):
                colors[iColor]=lut.GetTableValue(iColor)
            for iColor in range( nColors ):
                lut.SetTableValue(nColors-iColor-1, colors[iColor] )
        # replace bottom with white
        if replaceBottom:
            lut.SetTableValue(0,(1.0,1.0,1.0,1.0))

        mapper.SetLookupTable(lut)

        actor = vtk.vtkActor()
        actor.SetMapper(mapper)
        ren = vtk.vtkRenderer()

        renWin = vtk.vtkRenderWindow()
        renWin.AddRenderer(ren)

        ren.GradientBackgroundOn();
        ren.SetBackground(1,1,1);
        ren.SetBackground2(0,0,0);

        iren = vtk.vtkRenderWindowInteractor()
        iren.SetRenderWindow(renWin)

        ren.AddActor(actor)

        # create the scalar_bar
        scalar_bar = vtk.vtkScalarBarActor()
        scalar_bar.SetOrientationToHorizontal()
        scalar_bar.SetLookupTable(lut)
        scalar_bar.SetTitle(title)
        if manual_range is not None:
            nTicks = int(round(manual_range[1]-manual_range[0]+1))
            if nTicks >=3:
                scalar_bar.SetNumberOfLabels(nTicks)

        # Enable user interface interactor
        print("Initializing vtkRenderWindowInteractor.")
        iren.Initialize()
        print("Rendering the vtkRenderWindow (1)")
        renWin.Render()


        # create the scalar_bar_widget
        scalar_bar_widget = vtk.vtkScalarBarWidget()
        scalar_bar_widget.SetInteractor(iren)
        scalar_bar_widget.SetScalarBarActor(scalar_bar)
        scalar_bar_widget.On()



        # Different interactor styles
        # Type "j", "t", "c" or "a" to switch
        # Or "m"??
        style = vtk.vtkInteractorStyleSwitch()
        style.SetCurrentStyleToTrackballCamera()
        iren.SetInteractorStyle(style)

        print("Rendering the vtkRenderWindow (1)")
        renWin.Render()


        ren.ResetCamera()
        if camControl is None:
            print('Setting default camera position.')
            ren.GetActiveCamera().SetPosition(6.0, 0.0, 0.0)
            ren.GetActiveCamera().SetViewUp(0.0, 0.0, 1.0)
        else:
            # Set the camera using the parameters in the camControl object given as a parameter.
            camControl.vtkCamera = ren.GetActiveCamera()
            camControl.applyValuesToVtkCamera()
            
        print("Rendering the vtkRenderWindow (2)")
        renWin.Render()


        print(self.plotHelpText)

        print("Rendering the vtkRenderWindow (3)")
        renWin.Render()

        #print("Before Start()")
        iren.Start()
        #print("After Start()")





    plotHelpText="""
    For VTK 5.0:

    vtkInteractorStyle implements the "joystick" style of interaction. That is, holding down the mouse keys generates a stream of events that cause continuous actions (e.g., rotate, translate, pan, zoom). (The class vtkInteractorStyleTrackball implements a grab and move style.) The event bindings for this class include the following:

    Keypress j / Keypress t: toggle between joystick (position sensitive) and trackball (motion sensitive) styles. In joystick style, motion occurs continuously as long as a mouse button is pressed. In trackball style, motion occurs when the mouse button is pressed and the mouse pointer moves.
    Keypress c / Keypress a: toggle between camera and actor modes. In camera mode, mouse events affect the camera position and focal point. In actor mode, mouse events affect the actor that is under the mouse pointer.
    Button 1: rotate the camera around its focal point (if camera mode) or rotate the actor around its origin (if actor mode). The rotation is in the direction defined from the center of the renderer's viewport towards the mouse position. In joystick mode, the magnitude of the rotation is determined by the distance the mouse is from the center of the render window.
    Button 2: pan the camera (if camera mode) or translate the actor (if actor mode). In joystick mode, the direction of pan or translation is from the center of the viewport towards the mouse position. In trackball mode, the direction of motion is the direction the mouse moves. (Note: with 2-button mice, pan is defined as <shift>-Button 1.)
    Button 3: zoom the camera (if camera mode) or scale the actor (if actor mode). Zoom in/increase scale if the mouse position is in the top half of the viewport; zoom out/decrease scale if the mouse position is in the bottom half. In joystick mode, the amount of zoom is controlled by the distance of the mouse pointer from the horizontal centerline of the window.
    Keypress 3: toggle the render window into and out of stereo mode. By default, red-blue stereo pairs are created. Some systems support Crystal Eyes LCD stereo glasses; you have to invoke SetStereoTypeToCrystalEyes() on the rendering window.
    Keypress e: exit the application.
    Keypress f: fly to the picked point
    Keypress p: perform a pick operation. The render window interactor has an internal instance of vtkCellPicker that it uses to pick.
    Keypress r: reset the camera view along the current view direction. Centers the actors and moves the camera so that all actors are visible.
    Keypress s: modify the representation of all actors so that they are surfaces.
    Keypress u: invoke the user-defined function. Typically, this keypress will bring up an interactor that you can type commands in.
    Keypress w: modify the representation of all actors so that they are wireframe.

    """
    # Above is from https://vtk.org/doc/release/5.0/html/a01650.html

    # For vtk 5.8, check the subpages via "see also" in https://vtk.org/doc/release/5.8/html/a01087.html



class camControl():
    
    
    def __init__(self,vtkCamera=None):
        self.focalPoint    = None  # (x,y,z)
        self.position      = None  # (x,y,z)
        self.viewUp        = None  # (x,y,z)
        self.__viewAngle     = None  # degrees (vertical)
        self.__clipRange     = None  # (min,max)
        
        self.vtkCamera = vtkCamera
        if self.vtkCamera is not None:
            self.loadValuesFromVtkCamera()

    def loadValuesFromVtkCamera(self):
        '''
        Load the settings from the VtkCamera
        '''
        if self.vtkCamera is None:
            raise ValueError('No camera to load from.')
            
        self.focalPoint = self.vtkCamera.GetFocalPoint()
        self.position   = self.vtkCamera.GetPosition()
        self.viewUp     = self.vtkCamera.GetViewUp()
        self.viewAngle  = self.vtkCamera.GetViewAngle()
    
    def applyValuesToVtkCamera(self):
        if self.vtkCamera is None:
            raise ValueError('No camera to apply to.')
        
        #self.printValues()
        self.vtkCamera.SetViewAngle( self.viewAngle )
        self.vtkCamera.SetPosition(  self.position  )
        self.vtkCamera.SetViewUp(    self.viewUp    )
        self.vtkCamera.SetFocalPoint(self.focalPoint)

        self.vtkCamera.OrthogonalizeViewUp()
        #minClip = self.vtkCamera.GetDistance() * 1.0e-3
        #maxClip = 1.0e2
        #self.vtkCamera.SetClippingRange(minClip,maxClip)
        self.vtkCamera.SetClippingRange(self.clipRange)
        #self.printValues()

    def printValues(self):
        print('Cam: '+str(self.vtkCamera))
        print('Cam: position   ({:6.3f},{:6.3f},{:6.3f})\n'.format(
            self.position[0],self.position[1],self.position[2]))
        print('Cam: focalpoint ({:6.3f},{:6.3f},{:6.3f})\n'.format(
            self.focalPoint[0],self.focalPoint[1],self.focalPoint[2]))
        print('Cam: viewUp     ({:6.3f},{:6.3f},{:6.3f})\n'.format(
            self.viewUp[0],self.viewUp[1],self.viewUp[2]))
        print('Cam: viewAngle   {:6.3f}\n'.format(
            self.viewAngle))
    
    @property
    def clipRange(self):
        return self.__clipRange

    def clipRange(self,cr):
        if len(cr) != 2:
            raise ValueError('ClipRange should be length two tuple.')
        if cr[0] >= cr[1]:
            raise ValueError('ClipRange should be increasing')
        self.__clipRange = cr

     
#     @property
#     def focalPoint(self):
#         return self.__focalPoint
#     
#     @property
#     def position(self):
#         return self.__position
#     
#     @property
#     def viewUp(self):
#         return self.__viewUp
#     
    @property
    def viewAngle(self):
        return self.__viewAngle
    
    @viewAngle.setter
    def viewAngle(self,angle):
        if angle <= 0.0:
            raise ValueError('View angle must be positive "{}".'.angle)        
        self.__viewAngle = angle
        



a5py/a5py/wall/a5vtkwall.py
ascot4fusion/ascot5
import numpy as np

def loads(ids, energy, weight, area):
    """
    Get wall loads per triangle.

    Returns:
            ids
            edepo
            eload
            pdepo
            pload
            mdepo
            iangle
    """
    # Indices of wetted triangles
    wetted = np.unique(ids)
    area   = area[wetted-1]
    edepo  = np.zeros(wetted.shape)
    eload  = np.zeros(wetted.shape)
    pdepo  = np.zeros(wetted.shape)
    eload  = np.zeros(wetted.shape)
    mdepo  = np.zeros(wetted.shape)
    iangle = np.zeros(wetted.shape)

    # Sort markers by wall tile ID to make processing faster
    idx    = np.argsort(ids)
    ids    = ids[idx]
    energy = energy[idx]
    weight = weight[idx]

    idx = np.append(np.argwhere(ids[1:] - ids[:-1]).ravel(), ids.size-1)

    i0 = 0
    for i in range(wetted.size):
        i1 = idx[i] + 1
        mdepo[i]  = i1 - i0
        pdepo[i]  = np.sum(weight[i0:i1])
        edepo[i]  = np.sum(energy[i0:i1]*weight[i0:i1])
        iangle[i] = 0 #TBD
        i0 = i1
    pload = pdepo / area
    eload = edepo / area

    return (wetted, edepo, eload, pdepo, pload, mdepo, iangle)


def figuresofmerit(ids, energy, weight, area, plim=None):
    """
    Return figures of merit (0D quantities) for wall loads.

    Returns:
    peak power/energy load
    peak particle load
    number of markers
    total wetted area
    wetted area above limi
    """
    wetted, _, eload, _, pload, mdepo, _ = loads(ids, energy, weight, area)

    wetted_total = np.sum(area[wetted-1])
    energy_peak  = np.amax(eload)

    return (wetted_total, energy_peak)


def loadvsarea(ids, energy, weight, area):
    """
    Return a curve showing affected area as a function of (minimum) load.
    """
    wetted, _, eload, _, _, _, _ = loads(ids, energy, weight, area)
    idx = np.argsort(-eload)
    y = (eload[idx])
    x = np.cumsum(w[idx])

a5py/a5py/wall/loads.py
ascot4fusion/ascot5
"""
Routines for plotting walls.

File: wall/plot.py
"""
import numpy as np

import importlib.util as util

plt = util.find_spec("matplotlib")
if plt:
        import matplotlib.pyplot as plt

def plot_segments(x, y, axes=None):
    """
    """
    newfig = axes is None
    if newfig:
        plt.figure()
        axes = plt.gca()

    x = np.append(x, x[0])
    y = np.append(y, y[0])
    axes.plot(x, y, color="black", linewidth=2)

    if newfig:
        axes.axis("scaled")
        plt.show(block=False)

def plot_projection(x1x2x3, y1y2y3, z1z2z3, axes=None):
    """
    """
    newfig = axes is None
    if newfig:
        plt.figure()
        axes = plt.gca()

    r1r2r3 = np.sqrt(x1x2x3 * x1x2x3 + y1y2y3 * y1y2y3)

    r1r2r3 = r1r2r3.ravel()
    z1z2z3 = z1z2z3.ravel()
    axes.plot(r1r2r3, z1z2z3, marker=".", markeredgecolor="black",
              linestyle="None")

    if newfig:
        axes.axis("scaled")
        plt.show(block=False)

def plot_intersection(x1x2x3, y1y2y3, z1z2z3, phi, axes=None):
    """
    """
    newfig = axes is None
    if newfig:
        plt.figure()
        axes = plt.gca()

    # Polar coordinates, p1p2p3 in range [-pi, pi]
    r1r2r3, p1p2p3 = mathlib.cart2pol(x1x2x3, y1y2y3)
    # Wrap phi to range [-pi, pi) and set p1p2p3 == phi to zero
    phi = np.mod(phi + np.pi, 2*np.pi) - np.pi
    p1p2p3 = p1p2p3 - phi
    # Mod p1p2p3 again to range [-pi, pi)
    p1p2p3 = np.mod(p1p2p3 + np.pi, 2*np.pi) - np.pi
    # Filter out triangles crossing the -pi/pi boundary
    ind_valid = np.amax(p1p2p3, 1) - np.amin(p1p2p3, 1) < np.pi
    # Intersection happens if any(p1p2p3) < phi = 0 and any(p1p2p3) >= phi = 0
    p1p2p3_sign = np.copysign(1, p1p2p3)
    p1p2p3_polarity = np.sum(p1p2p3_sign, 1)
    ind_inter = np.logical_and(np.abs(p1p2p3_polarity) < 3, ind_valid)

    # Find intersection line for all crossing triangles
    p1p2p3_polarity = np.tile(p1p2p3_polarity, (3,1)).T
    ind_inter = np.tile(ind_inter, (3,1)).T
    ind_diff = np.logical_and(p1p2p3_sign != p1p2p3_polarity, ind_inter)
    ind_same = np.logical_and(p1p2p3_sign == p1p2p3_polarity, ind_inter)

    phi_diff = p1p2p3[ind_diff]
    phi_same = np.reshape(p1p2p3[ind_same], (-1,2))
    phi_diffs = (- phi_same / (phi_diff[:,None] - phi_same))

    r_diff = r1r2r3[ind_diff]
    r_same = np.reshape(r1r2r3[ind_same], (-1,2))
    r = r_same + (r_diff[:,None] - r_same ) * phi_diffs

    z_diff = z1z2z3[ind_diff]
    z_same = np.reshape(z1z2z3[ind_same], (-1,2))
    z = z_same + (z_diff[:,None] - z_same ) * phi_diffs

    # Add triangles for which one side lies entirely on phi = 0
    ind_2zero = np.sum(p1p2p3 == 0, axis=1) == 2
    ind_2zero = np.tile(ind_2zero, (3,1)).T
    ind_zero = p1p2p3 == 0
    ind_valid = np.tile(ind_valid, (3,1)).T
    ind_onplane = np.logical_and.reduce((ind_2zero, ind_zero, ind_valid))

    r2 = np.reshape(r1r2r3[ind_onplane], (-1,2))
    z2 = np.reshape(z1z2z3[ind_onplane], (-1,2))

    # Add triangles for which all 3 sides are entirely on phi = 0
    ind_3zero = np.all(p1p2p3 == 0, axis=1)
    r3 = np.concatenate((r1r2r3[ind_3zero][:,(0,1)],
                         r1r2r3[ind_3zero][:,(1,2)],
                         r1r2r3[ind_3zero][:,(2,0)]))
    z3 = np.concatenate((z1z2z3[ind_3zero][:,(0,1)],
                         z1z2z3[ind_3zero][:,(1,2)],
                         z1z2z3[ind_3zero][:,(2,0)]))

    axes.plot(r.T, z.T, 'k')
    axes.plot(r2.T, z2.T, 'k')
    axes.plot(r3.T, z3.T, 'k')

    if newfig:
        axes.axis("scaled")
        plt.show(block=False)

a5py/a5py/wall/plot.py
ascot4fusion/ascot5
"""
Created on Mar 9, 2020

@author: sjjamsa
"""

import numpy as np
import a5py.marker.endcond
from scipy.constants import physical_constants as const

def wallLoad3DEndstate(run):
    """
    Calculates the wall load from 3D wall for the specified run.
    @param hdf5: a5py.ascot5io.ascot5.Ascot(filename)

    @return: A float array of power load per triangle (W/m^2).
    """

    cW = run.wall  # The wall class
    es = run.endstate  # The Endstates

    A = cW.area()  # Triangle area
    nWallTris = cW.getNumberOfElements()

    E = es.get("endcond")
    wall_hit_endcond = a5py.marker.endcond.getbin("wall")
    # Pick the wall hit particles
    WH = (E & wall_hit_endcond) > 0

    T = es.get("walltile")[WH] - 1
    weight = es.get("weight")[WH]
    ene = es.get("energy")[WH]
    # The energies are now in eV. Need to convert to Joules.
    ene = ene * const["elementary charge"][0]

    # Calculate the power represented by each marker
    P = np.multiply(weight, ene)

    # Calculate the wall load
    wallLoad = np.bincount(
        T, weights=P, minlength=nWallTris
    )  # This is a histogram, where triangle indexes are the bins
    nonZero = wallLoad > 0.0
    wallLoad[nonZero] = np.divide(wallLoad[nonZero], A[nonZero])

    # This isn't accessed before function return. Should probably be removed
    # but leaving in case it is a useful variable that might be wanted in the
    # future
    # Nmarkers = np.bincount(T, minlength=nWallTris)

    return wallLoad

a5py/a5py/wallloads/calculate.py
ascot4fusion/ascot5
'''
Created on Mar 10, 2020

@author: sjjamsa
'''

def as3DpolyData(hdf5):
    '''
    Calculate the wall 3D loads to a VTK polydata datatype.
    @param hdf5: a5py.ascot5io.ascot5.Ascot(filename)
    '''

    import a5py.wallloads.calculate
    import a5py.wall.a5vtkwall
    
    wallLoad = a5py.wallloads.calculate.wallLoad3DEndstate(hdf5.active)
    wall = hdf5.active.wall
    
    avtk = a5py.wall.a5vtkwall.a5VtkWall(wall) 
    avtk.addIndex() # This is just the triangle index, included for just in case...
    
    fieldname = 'Wall load (W/m^2)'
    avtk.addScalarField(fieldname=fieldname,data=wallLoad)

    return avtk

def save3DAsVTP(hdf5,filename):
    '''
    Calculate and save the wall 3D loads to a VTK file.
    @param hdf5: a5py.ascot5io.ascot5.Ascot(filename)
    @param filename: The output filename, usually with the suffix .vtp
    '''
    avtk = as3DpolyData(hdf5)

    avtk.writeVtp(filename)


a5py/a5py/wallloads/toVTK.py
ascot4fusion/ascot5
#!/usr/bin/env python3

"""
Test ASCOT5 physics.

This module contains the routines to carry out all physics tests. These routines
call the init, run, and check routines of each test module.

All tests are done using a single HDF5 file test_ascot.h5. It is user's
responsibility to remove that file before repeating the tests. Running
tests requires that ascot5_main binary is in the same folder as from which
this script is called.

To init, run and check all tests, call this script as
> python test_ascot.py

To only init (i.e. write input data to the HDF5 file), call
> python test_ascot.py init

To only run (i.e. carry out the ascot simulations), call
> python test_ascot.py run

To only check (i.e. plot the results), call
> python test_ascot.py check

File: test_ascot.py
"""

import sys
import subprocess
import os.path
from time import perf_counter as timer

import numpy as np

import a5py.ascot5io.ascot5      as ascot5
import a5py.ascot5io.ascot5tools as tools

import a5py.testascot.helpers as helpers
from a5py.testascot import test_elementary
from a5py.testascot import test_orbitfollowing
from a5py.testascot import test_gctransform
from a5py.testascot import test_coulombcollisions
from a5py.testascot import test_classicaltransport
from a5py.testascot import test_neoclassicaltransport
from a5py.testascot import test_boozercoordinates

# You can specify here which tests to run

## Test test_elementary.py
dotest_elementary            = True
## Test test_orbitfollowing.py
dotest_orbitfollowing        = True
## Test test_gctransform.py
dotest_gctransform           = True
## Test test_coulombcollisions.py
dotest_coulombcollisions     = False
## Test test_classicaltransport.py
dotest_classicaltransport    = True
## Test test_neoclassicaltransport.py
dotest_neoclassicaltransport = False
## Test test_neoclassicaltransport.py
dotest_boozercoordinate      = True

def init():
    """
    Initialize all tests to test_ascot.h5
    """

    # Make sure no existing data interferes with new tests.
    subprocess.call(["rm", "-f", helpers.testfn])

    if dotest_elementary:
        print("Initializing test_elementary.")
        test_elementary.init()

    if dotest_orbitfollowing:
        print("Initializing test_orbitfollowing.")
        test_orbitfollowing.init()

    if dotest_gctransform:
        print("Initializing test_gctransform.")
        test_gctransform.init()

    if dotest_coulombcollisions:
        print("Initializing test_coulombcollisions.")
        test_coulombcollisions.init()

    if dotest_classicaltransport:
        print("Initializing test_classicaltransport.")
        test_classicaltransport.init()

    if dotest_neoclassicaltransport:
        print("Initializing test_neoclassicaltransport.")
        test_neoclassicaltransport.init()

    if dotest_boozercoordinate:
        print("Initializing test_boozercoordinates.")
        test_boozercoordinates.init()


def run():
    """
    Run each test

    This function goes through each test by first finding which inputs are
    required for that particular test (from their description) and using the
    qids to carry out the simulation. A new simulation is started when the
    previous one finishes.
    """
    if not os.path.isfile(helpers.testfn):
        print(helpers.testfn + " is missing.")
        print("Aborting runs.")
        sys.exit()

    if not os.path.isfile(helpers.testbin):
        print(helpers.testbin + " is missing.")
        print("Aborting runs.")
        sys.exit()

    frm   = lambda x: "%.3f s" % x
    start = timer()
    inter = start

    if dotest_elementary:
        print("Running test_elementary.")
        test_elementary.run()
        print("Elapsed time is " + frm(timer() - inter))
        print("")
        inter = timer()

    if dotest_orbitfollowing:
        print("Running test_orbitfollowing.")
        test_orbitfollowing.run()
        print("Elapsed time is " + frm(timer() - inter))
        print("")
        inter = timer()

    if dotest_gctransform:
        print("Running test_gctransform.")
        test_gctransform.run()
        print("Elapsed time is " + frm(timer() - inter))
        print("")
        inter = timer()

    if dotest_coulombcollisions:
        print("Running test_coulombcollisions.")
        test_coulombcollisions.run()
        print("Elapsed time is " + frm(timer() - inter))
        print("")
        inter = timer()

    if dotest_classicaltransport:
        print("Running test_classicaltransport.")
        test_classicaltransport.run()
        print("Elapsed time is " + frm(timer() - inter))
        print("")
        inter = timer()

    if dotest_neoclassicaltransport:
        print("Running test_neoclassicaltransport.")
        test_neoclassicaltransport.run()
        print("Elapsed time is " + frm(timer() - inter))
        print("")
        inter = timer()

    if dotest_boozercoordinate:
        print("Running test_boozercoordinates.")
        test_boozercoordinates.run()
        print("Elapsed time is " + frm(timer() - inter))
        print("")
        inter = timer()

    print("Total elapsed time is " + frm(timer() - start))


def check():
    """
    Plot the results of these tests.
    """
    if not os.path.isfile(helpers.testfn):
        print(testfn + " is missing.")
        print("Aborting checks.")
        sys.exit()

    if dotest_elementary:
        print("Checking results of test_elementary.")
        test_elementary.check()

    if dotest_orbitfollowing:
        print("Checking results of test_orbitfollowing.")
        test_orbitfollowing.check()

    if dotest_gctransform:
        print("Checking results of test_gctransform.")
        test_gctransform.check()

    if dotest_coulombcollisions:
        print("Checking results of test_coulombcollisions.")
        test_coulombcollisions.check()

    if dotest_classicaltransport:
        print("Checking results of test_classicaltransport.")
        test_classicaltransport.check()

    if dotest_neoclassicaltransport:
        print("Checking results of test_neoclassicaltransport.")
        test_neoclassicaltransport.check()

    if dotest_boozercoordinate:
        print("Checking results of test_boozercoordinates.")
        test_boozercoordinates.check()

if __name__ == '__main__':

    if( len(sys.argv) == 1 ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        print("")
        print("Running tests.")
        run()
        print("Runs complete.")
        print("")
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    if(len(sys.argv) > 2):
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

    if(   sys.argv[1] == "init" ):
        print("Initializing tests.")
        init()
        print("Initialization complete.")
        sys.exit()

    elif( sys.argv[1] == "run" ):
        print("Running tests.")
        run()
        print("Runs complete.")
        sys.exit()

    elif( sys.argv[1] == "check" ):
        print("Checking test results.")
        check()
        print("Testing complete.")
        sys.exit()

    else:
        print("Too many arguments.")
        print("Only \"init\", \"run\" or \"check\" is accepted.")
        print("Aborting.")
        sys.exit()

a5py/bin/test_ascot.py
ascot4fusion/ascot5
# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Set path to where the Python source can be found---------------------------
import os
import sys
sys.path.insert(0, os.path.abspath('../../a5py'))

import a5py

# -- Project information -------------------------------------------------------
project   = 'ASCOT5'
copyright = '2023, Ascot Group'
author    = 'Ascot Group'
release   = a5py.ascot5io.coreio.fileapi.VERSION

# -- General configuration -----------------------------------------------------
extensions = ['sphinx.ext.autodoc',      # For generating doc from Python source
              'numpydoc',                # Source docs are done in numpy style
              'nbsphinx',                # Embed Jupyter notebooks
              'breathe',                 # Sphinx can access Doxygen output
              'sphinxcontrib.bibtex',    # Can use bibtex
              'sphinx.ext.autosummary',  # Creating summary tables
              'sphinx.ext.intersphinx',] # Link to external libraries

exclude_patterns = []
numpydoc_xref_param_type = True # Automatically link str, array_like, etc.
numpydoc_show_class_members = False # Removes table summarizing class methods

# -- Where Doxygen generated xml files are located -----------------------------
breathe_projects = {'ascot5': '_static/xml'}
breathe_default_project = 'ascot5'

intersphinx_mapping = {
    'python'    : ('https://docs.python.org/3', None),
    'numpy'     : ('https://docs.scipy.org/doc/numpy', None),
    'scipy'     : ('https://docs.scipy.org/doc/scipy/reference', None),
    'matplotlib': ('https://matplotlib.org/stable', None)}

# -- Options for HTML output ---------------------------------------------------
html_theme = 'sphinx_rtd_theme'
html_theme_options = {
    'logo_only': True,
    'navigation_depth':6,
    }
html_logo = 'figs/logo.png'

html_static_path = ['_static']
html_css_files = [
    'custom.css',
]

# Bibtex
bibtex_bibfiles = ['ascotwork.bib']
nbsphinx_execute = 'never'

a5py/docs/conf.py
ascot4fusion/ascot5
from os import path
from setuptools import setup, find_packages

here = path.abspath(path.dirname(__file__))

with open(path.join(here, 'requirements.txt')) as requirements_file:
    # Parse requirements.txt, ignoring any commented-out lines.
    requirements = [line for line in requirements_file.read().splitlines()
                    if not line.startswith('#')]

setup(name='a5py',
      version='0.0',
      description='ASCOT5 python library',
      url='http://github.com/ascot/python/a5py',
      license='LGPL',
      packages=find_packages(),
      zip_safe=False,
      install_requires=requirements,
      scripts=[
        'bin/a5removegroup',
        'bin/a5copygroup',
        'bin/a5editoptions',
        'bin/a5combine',
        'bin/a5ascot4input',
        'bin/a5setactive',
        'bin/a5gui',
        'bin/a5ls',
        'bin/a5makecompatible',
        'bin/a5doxygen',
        'bin/test_ascot.py'
      ],
      package_data={'a5py/ascotpy': ['lib.so']},
      include_package_data=True)

a5py/setup.py
ascot4fusion/ascot5