appl121105_iotorus.simulator

'''Simulator class

.. todo::

    Io phase setting method.

.. todo::

    Future, but plasma model not in the JEqS frame.
'''
import logging
logging.basicConfig()

import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate

from . import neutralmodel
from . import plasmamodel
from . import xsect
from . import trajectory
from . import jupiter
from . import moon_io
from . import f2j

import pyana.pep.iotorus



class SimulatorCore:
    logger = logging.getLogger('SimulatorCore')
    #logger.setLevel(logging.DEBUG)
    def __init__(self):
        self.observer = None
        self.trajectory = None

        self.tlist = None

        self.enamass = 16.

        tpm = pyana.pep.iotorus.TwoPeakModel()
        self.nmodel = neutralmodel.NeutralModelTwoPeak2(moon_io.Io())
#        self.pmodel = plasmamodel.SimpleCorotationalFlow2(m_amu=self.enamass)
        self.pmodel = plasmamodel.CorotationalFlowMJconv(m_amu=self.enamass)
        self.xsectmodel = xsect.CrossSectionSimplest()

    def set_iophase(self, iophase):
        ''' Set the io phase.
        '''
        self.nmodel = neutralmodel.NeutralModelTwoPeak2(moon_io.Io(initial_phase_degree=iophase))

    def set_jupiterphase(self, jupiter_phase):
        self.pmodel = plasmamodel.CorotationalFlowMJconv(m_amu=self.enamass, jupiter_phase=jupiter.JupiterPhase(jupiter_phase))

    def set_observer(self, observer):
        ''' Set the observer.
        ''' 
        self.observer = observer
        self.trajectory = trajectory.Trajectory(self.observer)

    def set_tlist(self, tlist):
        ''' Set the calculation time list
        '''
        self.tlist = tlist

    def run(self, mpl=False, debug=False, mpl_prefix=None):
        ''' Run the calculation.
        '''
        if self.observer == None:
            raise ValueError('!!! Observer not defined.\n'
                '     Use set_observer method\n'
                '     to set the observer.\n')

        if self.tlist == None:
            raise ValueError('!!! Time list not defined.\n'
                '     Use set_tlist method\n'
                '     to set the Time list.\n')

        nt = len(self.tlist)
        # First, trajectory is calculated.
        dist_along, traj_coord = self.trajectory.get_length(self.tlist)  # (3, N), km
        velvecs = traj_coord[3:, :]
        traj_coord = traj_coord[:3, :]
        # Then, calculate the neutral density there.
        # Then, also calculate the plasma distribution function there.
        neutral_alongtraj = self.nmodel.get_density(traj_coord[0, :],
                            traj_coord[1, :], traj_coord[2, :], tarr=self.tlist)
        plasma_distfunc = self.pmodel.get_distribution_function(
                            traj_coord[0, :], traj_coord[1, :],
                            traj_coord[2, :],
                            velvecs[0, :], velvecs[1, :], velvecs[2, :], tarr=self.tlist)
        crosssection = np.zeros([nt])
        for i_t in range(nt):
            # For generic purpose, argument should be |vp-vn|
            crosssection[i_t] = self.xsectmodel.get_crosssection()

        ### ENA distribution function, df/dl, in MKSA.  s3/m6 * /m3 * m2
        denaflux = plasma_distfunc * (neutral_alongtraj * 1e6) * (crosssection / 1e4)
        ### Integral over L.  f.
        dist_along_meter = dist_along * 1e3  # in m

        ### Integral
        f_ENA = scipy.integrate.trapz(denaflux, x=dist_along_meter)

        if debug:  # Too power hungry.  Only for debugging.
            self.logger.debug('xyz [km] = %g %g %g' % str(traj_coord))
            self.logger.debug('vxyz [km/s] = %g %g %g' % str(velvecs))
            self.logger.debug('l [m] = %g %g %g' % str(dist_along_meter))
            self.logger.debug('nn [/cm3] = %g %g %g' % str(neutral_alongtraj))
            self.logger.debug('fp [s6/m3] = %g %g %g' % str(plasma_distfunc))
            self.logger.debug('sig [cm2] = %g %g %g' % str(crosssection))

        # Conversion to differential flux

        vel_t0 = self.observer.sensor.get_velocity_vector()  # Velocity vector in spacecraft frame.in km/s
        vel = np.sqrt((vel_t0 ** 2).sum())
        ene = f2j.ene(vel, self.enamass)

        j_ENA = f2j.f2j(f_ENA, ene, self.enamass)

        if mpl:
            rj = 71492.
            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.plot(self.tlist, traj_coord[0, :] / rj)
            ax.plot(self.tlist, traj_coord[1, :] / rj)
            ax.plot(self.tlist, traj_coord[2, :] / rj)
            ax.set_title('Timeline of ENA trajectory (inertia)')
            if mpl_prefix != None:
                fig.savefig('%s-00.png' % mpl_prefix)

            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.plot(traj_coord[0, :] / rj, traj_coord[1, :] / rj)
            ax.plot([0], [0], 'ro')
            ax.set_aspect('equal')
            ax.set_title('XY of ENA trajectory (inertia)')
            if mpl_prefix != None:
                fig.savefig('%s-01.png' % mpl_prefix)

            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.plot(dist_along / rj, neutral_alongtraj, label='Nn /cm3')
            ax.plot(dist_along / rj, plasma_distfunc, label='fp s3/m6')
            ax.plot(dist_along / rj, denaflux, label='fena s3/m6 /m')
            ax.set_yscale('log')
            ax.legend(loc='best')
            ax.set_title('Fluxes')
            if mpl_prefix != None:
                fig.savefig('%s-02.png' % mpl_prefix)

        return f_ENA, j_ENA

class Visualizer:
    def __init__(self, core):
        raise NotImplementedError('')