apps120811_flux_recalc.precipitation_notrace

''' Calculate precipitation flux from JIAMHD data

Calculate precipitation following the scheme below.

Iterate over the surface

1. Specify the point at surface.
2. Trace back to the point at 1.2 Rg above (zenith)
3. Obtain the plasma parameters (n, vxyz, vth, Bxyz) at the point.
4. Calculate the velocity vector along magnetic field.
5. Calculate the precipitation flux along the magnetic field.

.. image:: ../../../src/scripts/apps120811_flux_recalc/precipitation_notrace1.png

'''


import os
import sys
import logging
logging.basicConfig()
logger = logging.getLogger('project_at_500km.py')
import datetime
import math

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

from pyana.util import streamline
from pyana.util import maxwell
from pyana.pep import mhddata


def main(mf=None, pp=None, dlon=45, dlat=10):

    if mf == None:
        mf = mhddata.MagneticField1201()
    mf.set_interpolate_strategy('project')  # Using scaling projection.

    if pp == None:
        pp = mhddata.PlasmaParameter1205()

    lats = np.arange(-90, 91, dlat)
    lons = np.arange(0, 361, dlon)

    projlonlist = np.zeros([len(lons), len(lats)]) + np.nan
    projlatlist = np.zeros([len(lons), len(lats)]) + np.nan
    projnlist = np.zeros([len(lons), len(lats)]) + np.nan
    projvlist = np.zeros([3, len(lons), len(lats)]) + np.nan
    projvthlist = np.zeros([len(lons), len(lats)]) + np.nan
    projblist = np.zeros([3, len(lons), len(lats)]) + np.nan

    projfluxlist = np.zeros([len(lons), len(lats)]) + np.nan
    projvblist = np.zeros([len(lons), len(lats)]) + np.nan

    # Suppress warnings.
    logging.getLogger('MagneticField1201').setLevel(logging.WARN)

    for ilon in range(len(lons)):
        lon = lons[ilon]
        lonr = np.deg2rad(lon)
        for ilat in range(len(lats)):
            lat = lats[ilat]
            latr = np.deg2rad(lat)

            pos1 = [np.cos(lonr) * np.cos(latr), np.sin(lonr) * np.cos(latr), np.sin(latr)]
            pos1 = np.array(pos1) * 1.2
        
            # Trace back to r=1.2.
            if np.isnan(pos1).any():
                continue
            r1 = np.sqrt((pos1 ** 2).sum())

            # Longitude and latitude
            projlon = np.rad2deg(np.arctan2(pos1[1], pos1[0]))
            projlat = 90 - np.rad2deg(np.arccos(pos1[2] / r1))
            projlonlist[ilon, ilat] = projlon
            projlatlist[ilon, ilat] = projlat

            # Parameters
            print(pos1)
            n, vx, vy, vz, temp, pth = pp.interpolate3d(pos1[0], pos1[1], pos1[2])
            projnlist[ilon, ilat] = n * 1e6
            projvlist[:, ilon, ilat] = vx * 1e3, vy * 1e3, vz * 1e3   # In m/s
            vth = mhddata.t2vth(temp, mass=16.)   # In m/s
            projvthlist[ilon, ilat] = vth
            # Bfield
            bx, by, bz = mf.interpolate3d(pos1[0], pos1[1], pos1[2])
            b = np.array([bx, by, bz])
            # Change oxrientation toward surface!!
            if (b * pos1).sum() > 0:   # if b vector pointing upward
                b = -b
            projblist[:, ilon, ilat] = b

            # Velocity component along b.
            bn = b / np.sqrt((b * b).sum())
            vb = (np.array([vx, vy, vz]) * bn).sum() * 1e3   # in m/s
            projvblist[ilon, ilat] = vb
            # Precipitating flux
            precip = maxwell.flux_precip(n * 1e6, vb, vth) * 1e-4   # In /cm2 s
            projfluxlist[ilon, ilat] = precip

            print(lon, lat, projlon, projlat, n, vx, vy, vz, temp, projvthlist[ilon, ilat], bx, by, bz, precip)

    return mf, pp, lons, lats, projfluxlist, projlonlist, projlatlist, projnlist, projvlist, projvthlist, projblist, projvblist,

def main_all():
    retval = main()

    lvalue = mhddata.LValue1201()
    ocb = lvalue.opencloseboundaries()

    plt.figure()
    plt.pcolor(retval[2], retval[3], np.ma.log10(retval[4].T), vmin=7, vmax=8)
    plt.plot(lvalue.lonlist(), ocb[0], 'w:', lw=2)
    plt.plot(lvalue.lonlist(), ocb[1], 'w:', lw=2)
    plt.title('Downgoing Flux (log) /cm2 s')
    plt.xlabel('Longitude')
    plt.ylabel('Latitude')
    plt.xlim(0, 360)
    plt.ylim(-90, 90)
    plt.colorbar()
    plt.savefig('precipitation_notrace1.png')

    plt.figure()
    plt.pcolor(retval[2], retval[3], retval[7].T / 1e6, vmin=0)
    plt.plot(lvalue.lonlist(), ocb[0], 'w:', lw=2)
    plt.plot(lvalue.lonlist(), ocb[1], 'w:', lw=2)
    plt.title('Density /cm3')
    plt.colorbar()
    plt.savefig('precipitation_notrace2.png')

    plt.figure()
    v = np.sqrt((retval[8] ** 2).sum(0))
    plt.pcolor(retval[2], retval[3], v.T / 1e3, vmin=0)
    plt.plot(lvalue.lonlist(), ocb[0], 'w:', lw=2)
    plt.plot(lvalue.lonlist(), ocb[1], 'w:', lw=2)
    plt.title('Velocity km/s')
    plt.colorbar()
    plt.savefig('precipitation_notrace3.png')

    plt.figure()
    plt.pcolor(retval[2], retval[3], retval[9].T / 1e3, vmin=0)
    plt.plot(lvalue.lonlist(), ocb[0], 'w:', lw=2)
    plt.plot(lvalue.lonlist(), ocb[1], 'w:', lw=2)
    plt.title('Vth km/s')
    plt.colorbar()
    plt.savefig('precipitation_notrace4.png')

    plt.figure()
    plt.pcolor(retval[2], retval[3], retval[11].T / 1e3, vmin=0)
    plt.plot(lvalue.lonlist(), ocb[0], 'w:', lw=2)
    plt.plot(lvalue.lonlist(), ocb[1], 'w:', lw=2)
    plt.title('V// km/s')
    plt.colorbar()
    plt.savefig('precipitation_notrace5.png')

    return retval

if __name__ == "__main__":
    
    retval = main_all()