INWORK: lowtran interface (needs test) and concat/merge

PlotEmissionProfiles.py

@@ -1,9 +1,12 @@
 #!/usr/bin/env python
 """
 using excitation rates and other factors, creates volume emission rate profiles.
+
+python PlotEmissionProfiles.py ../transcarread/tests/data
 """
 from pathlib import Path
 import logging
+import numpy as np
 import h5py
 from matplotlib.pyplot import show
 import seaborn as sns
@@ -15,47 +18,38 @@
 from gridaurora.filterload import getSystemT
 from gridaurora.plots import writeplots
 # github.com/scivision/transcarread
-from transcarread import readTranscarInput,readexcrates, SimpleSim
+import transcarread as tr
 
 if __name__ == '__main__':
-    import cProfile,pstats
-
-    #
     from argparse import ArgumentParser
     p = ArgumentParser(description = 'using excitation rates and other factors, creates volume emission rate profiles.')
     p.add_argument('path',help='root path where simulation inputs/outputs are')
     p.add_argument('beamenergy',help='beam energy [eV] to plot',type=float)
     p.add_argument('-r','--reacreq',help='reactions to include e.g. metastable atomic',nargs='+',default=['metastable','atomic','n21ng','n2meinel','n22pg','n21pg'])
     p.add_argument('-m','--makeplot',help='specify plots to make e.g. vjinc vjinc1d',nargs='+',default=['eigtime','eigtime1d','spectra'])
     p.add_argument('--datcarfn',help='path to dir.input/DATCAR',default='dir.input/DATCAR')
-    p.add_argument('--profile',help='profile performance',action='store_true')
     p.add_argument('-o','--outfile',help='HDF5 filename to write')
     p.add_argument('-t','--tind',help='time index to use',type=int,default=0)
     p=p.parse_args()
 
     tReqInd = p.tind
     #NOTE: make sure tReqInd after precipitation starts or you're looking at airglow instead of aurora!
     path = Path(p.path).expanduser()
-    simpath = path/('beam'+str(p.beamenergy))
+    simpath = path / f'beam{p.beamenergy:.0f}'
 
     excrpath = simpath/'dir.output/emissions.dat'
-    excrates = readexcrates(excrpath)[0]
-    sim = SimpleSim(filt=None,inpath=simpath,reacreq=p.reacreq)
-#%% testing code timing only
-    if p.profile:
-        proffn = 'calcemissions.pstats'
-        cProfile.run('calcemissions(excrates,tReqInd,sim)',proffn)
-        pstats.Stats(proffn).sort_stats('time','cumulative').print_stats(50)
-        exit()
-#%% normal usage, continue processing data
-    tctime = readTranscarInput(simpath/p.datcarfn)
+    excrates = tr.readexcrates(excrpath)
+    sim = tr.SimpleSim(filt=None,inpath=simpath,reacreq=p.reacreq)
+#%%
+    tctime = tr.readTranscarInput(simpath/p.datcarfn)
 
-    t=excrates.minor_axis.to_datetime().to_pydatetime()
+    excrates = excrates.isel(time=tReqInd)
+    t = excrates.time
 
-    if t[tReqInd]<tctime['tstartPrecip']:
+    if t < np.datetime64(tctime['tstartPrecip']):
         logging.warning('you picked a time before precipitation started, so youre looking at AIRGLOW instead of AURORA!')
 
-    tver,ver,br=calcemissions(excrates, tReqInd, sim)
+    tver,ver,br = calcemissions(excrates['excitation'], sim)
     optT = getSystemT(tver.index,sim.bg3fn, sim.windowfn, sim.qefn, sim.obsalt_km, sim.zenang)
 #%% write as hdf5
     if p.outfile:

gridaurora/__init__.py

@@ -29,7 +29,7 @@ def readmonthlyApF107(yearmon:Union[str,datetime], fn:Union[str,Path]=None, forc
         yearmon = parse(yearmon)
     #not elif
     if isinstance(yearmon, datetime):
-        yearmon = int(str(yearmon.year) + f'{yearmon.month:02d}')
+        yearmon = int(f'{yearmon.year:d}{yearmon.month:02d}')
 
     assert isinstance(yearmon,int)
 #%% load data

gridaurora/calcemissions.py

@@ -1,8 +1,8 @@
-#!/usr/bin/env python3
+#!/usr/bin/env python
 import logging
 from numpy import trapz, isfinite, concatenate, nansum
 import h5py
-from xarray import DataArray
+import xarray
 from matplotlib.pyplot import figure,close,subplots
 #from matplotlib.colors import LogNorm
 from matplotlib.ticker import MultipleLocator
@@ -20,11 +20,10 @@
 spectraAminmax = (1e-1,8e5) #for plotting
 spectrallines=(391.44, 427.81, 557.7, 630.0, 777.4, 844.6) #297.2, 636.4,762.0, #for plotting
 
-def calcemissions(spec,tReqInd,sim):
+def calcemissions(rates:xarray, sim):
     if not sim.reacreq:
         return
 
-    assert isinstance(spec,DataArray), 'I expect an xarray DataArray as input'
 
     ver = None; lamb = None; br=None
     """
@@ -34,67 +33,62 @@ def calcemissions(spec,tReqInd,sim):
     """
 #%% METASTABLE
     if 'metastable' in sim.reacreq:
-        ver, lamb,br = getMetastable(spec,ver,lamb,br, sim.reactionfn)
+        ver, lamb,br = getMetastable(rates,ver,lamb,br, sim.reactionfn)
 #%% PROMPT ATOMIC OXYGEN EMISSIONS
     if 'atomic' in sim.reacreq:
-        ver, lamb,br = getAtomic(spec, ver, lamb, br, sim.reactionfn)
+        ver, lamb,br = getAtomic(rates, ver, lamb, br, sim.reactionfn)
 #%% N2 1N EMISSIONS
     if 'n21ng' in sim.reacreq:
-        ver, lamb,br = getN21NG(spec,ver,lamb, br, sim.reactionfn)
+        ver, lamb,br = getN21NG(rates,ver,lamb, br, sim.reactionfn)
 #%% N2+ Meinel band
     if 'n2meinel' in sim.reacreq:
-        ver, lamb,br = getN2meinel(spec,ver,lamb, br, sim.reactionfn)
+        ver, lamb,br = getN2meinel(rates,ver,lamb, br, sim.reactionfn)
 #%% N2 2P (after Vallance Jones, 1974)
     if 'n22pg' in sim.reacreq:
-        ver, lamb,br = getN22PG(spec,ver,lamb, br, sim.reactionfn)
+        ver, lamb,br = getN22PG(rates,ver,lamb, br, sim.reactionfn)
 #%% N2 1P
     if 'n21pg' in sim.reacreq:
-        ver, lamb,br = getN21PG(spec,ver,lamb, br, sim.reactionfn)
+        ver, lamb,br = getN21PG(rates,ver,lamb, br, sim.reactionfn)
 #%% Remove NaN wavelength entries
     if ver is None:
         raise ValueError('you have not selected any reactions to generate VER')
 #%% sort by wavelength, eliminate NaN
     lamb, ver, br = sortelimlambda(lamb,ver,br)
-    try:
-        tver = ver[tReqInd,...]
-        br = br[tReqInd,:]
-    except IndexError as e:
-        logging.error('error in time index, falling back to last time value')
-        print(str(e))
-        tver = ver[-1,...]
-        br = br[-1,:]
 #%% assemble output
-    dfver = DataArray(data=tver, coords=[('alt_km',spec.alt_km),('wavelength_nm',lamb)])
+    dfver = DataArray(data=ver, coords=[('alt_km',rates.alt_km),
+                                        ('wavelength_nm',lamb)])
 
     return dfver, ver,br
 
-def getMetastable(spec,ver,lamb,br,reactfn):
-    with h5py.File(str(reactfn),'r',libver='latest') as f:
-        A = f['/metastable/A'].value
+def getMetastable(rates,ver,lamb,br,reactfn):
+    with h5py.File(reactfn,'r') as f:
+        A = f['/metastable/A'][:]
         lambnew = f['/metastable/lambda'].value.ravel(order='F') #some are not 1-D!
 
     """
     concatenate along the reaction dimension, axis=-1
     """
-    vnew = concatenate((A[None,None,:2] * spec.loc[...,'no1s'].values[...,None],
-                        A[None,None,2:4]* spec.loc[...,'no1d'].values[...,None],
-                        A[None,None,4:] * spec.loc[...,'noii2p'].values[...,None]), axis=-1)
+    vnew = concatenate((A[:2] * rates.loc[...,'no1s'].values[:,None],
+                        A[2:4]* rates.loc[...,'no1d'].values[:,None],
+                        A[4:] * rates.loc[...,'noii2p'].values[:,None]), axis=-1)
+
+    assert vnew.shape == (rates.shape[0], A.size)
 
-    return catvl(spec.alt_km,ver,vnew,lamb,lambnew,br)
+    return catvl(rates.alt_km, ver, vnew, lamb, lambnew, br)
 
-def getAtomic(spec,ver,lamb, br,reactfn):
+def getAtomic(rates,ver,lamb, br,reactfn):
     """ prompt atomic emissions (nm)
     844.6 777.4
     """
-    with h5py.File(str(reactfn),'r',libver='latest') as f:
+    with h5py.File(reactfn,'r') as f:
         lambnew = f['/atomic/lambda'].value.ravel(order='F') #some are not 1-D!
 
-    vnew = concatenate((spec.loc[...,'po3p3p'].values[...,None],
-                        spec.loc[...,'po3p5p'].values[...,None]),axis=-1)
+    vnew = concatenate((rates.loc[...,'po3p3p'].values[...,None],
+                        rates.loc[...,'po3p5p'].values[...,None]),axis=-1)
 
-    return catvl(spec.alt_km,ver,vnew,lamb,lambnew,br)
+    return catvl(rates.alt_km,ver,vnew,lamb,lambnew,br)
 
-def getN21NG(spec,ver,lamb,br, reactfn):
+def getN21NG(rates,ver,lamb,br, reactfn):
     """
     excitation Franck-Condon factors (derived from Vallance Jones, 1974)
     """
@@ -103,28 +97,28 @@ def getN21NG(spec,ver,lamb,br, reactfn):
         lambdaA = f['/N2+1NG/lambda'].value.ravel(order='F')
         franckcondon = f['/N2+1NG/fc'].value
 
-    return doBandTrapz(A,lambdaA,franckcondon, spec.loc[...,'p1ng'],lamb,ver,spec.alt_km,br)
+    return doBandTrapz(A,lambdaA,franckcondon, rates.loc[...,'p1ng'],lamb,ver,rates.alt_km,br)
 
-def getN2meinel(spec,ver,lamb,br,reactfn):
+def getN2meinel(rates,ver,lamb,br,reactfn):
     with h5py.File(str(reactfn),'r',libver='latest') as f:
         A = f['/N2+Meinel/A'].value
         lambdaA = f['/N2+Meinel/lambda'].value.ravel(order='F')
         franckcondon = f['/N2+Meinel/fc'].value
     #normalize
     franckcondon = franckcondon/franckcondon.sum() #special to this case
 
-    return doBandTrapz(A,lambdaA,franckcondon, spec.loc[...,'pmein'],lamb,ver,spec.alt_km,br)
+    return doBandTrapz(A,lambdaA,franckcondon, rates.loc[...,'pmein'],lamb,ver,rates.alt_km,br)
 
-def getN22PG(spec,ver,lamb,br,reactfn):
+def getN22PG(rates,ver,lamb,br,reactfn):
     """ from Benesch et al, 1966a """
     with h5py.File(str(reactfn),'r',libver='latest') as f:
         A = f['/N2_2PG/A'].value
         lambdaA = f['/N2_2PG/lambda'].value.ravel(order='F')
         franckcondon = f['/N2_2PG/fc'].value
 
-    return doBandTrapz(A,lambdaA,franckcondon,spec.loc[...,'p2pg'],lamb,ver,spec.alt_km,br)
+    return doBandTrapz(A,lambdaA,franckcondon,rates.loc[...,'p2pg'],lamb,ver,rates.alt_km,br)
 
-def getN21PG(spec,ver,lamb,br,reactfn):
+def getN21PG(rates,ver,lamb,br,reactfn):
 
     with h5py.File(str(reactfn),'r',libver='latest') as fid:
         A = fid['/N2_1PG/A'].value
@@ -141,13 +135,13 @@ def getN21PG(spec,ver,lamb,br,reactfn):
 
     consfac = franckcondon/franckcondon.sum() #normalize
     losscoef = (consfac / tau1PG).sum()
-    N01pg=spec.loc[...,'p1pg'] / losscoef
+    N01pg=rates.loc[...,'p1pg'] / losscoef
 
     scalevec = (A * consfac[:,None]).ravel(order='F') #for clarity (verified with matlab)
 
     vnew = scalevec[None,None,:] * N01pg.values[...,None]
 
-    return catvl(spec.alt_km,ver,vnew,lamb,lambnew,br)
+    return catvl(rates.alt_km,ver,vnew,lamb,lambnew,br)
 
 def doBandTrapz(Aein,lambnew,fc,kin,lamb,ver,z,br):
     """
@@ -166,19 +160,23 @@ def doBandTrapz(Aein,lambnew,fc,kin,lamb,ver,z,br):
 
     return catvl(z,ver,vnew,lamb,lambnew,br)
 
-def catvl(z,ver,vnew,lamb,lambnew,br):
+def catvl(z, ver, vnew, lamb, lambnew, br):
     """
-    trapz integrates over axis=1, the altitude dimension
-    concatenate over reaction dimension, axis=-1
+    trapz integrates over altitude axis, axis = -2
+    concatenate over reaction dimension, axis = -1
+
+    br: column integrated brightness
+    lamb: wavelength [nm]
+    ver: volume emission rate  [photons / cm^-3 s^-3 ...]
     """
     if ver is not None:
-        br = concatenate((br,trapz(vnew,z,axis=1)), axis=-1) #must come first!
+        br = concatenate((br, trapz(vnew,z,axis=-2)), axis=-1) #must come first!
         ver=concatenate((ver,vnew), axis=-1)
         lamb=concatenate((lamb,lambnew))
     else:
         ver = vnew.copy(order='F')
         lamb = lambnew.copy()
-        br = trapz(ver,z,axis=1)
+        br = trapz(ver, z, axis=-2)
 
     return ver, lamb, br
 

gridaurora/filterload.py

@@ -7,10 +7,10 @@
 from xarray import DataArray
 # consider atmosphere
 try:
-    from lowtran import golowtran
+    import lowtran
 except ImportError as e:
     logging.error(f'failure to load LOWTRAN, proceeding without atmospheric absorption model.  {e}')
-    golowtran=None
+    lowtran = None
 '''
 gets optical System Transmittance from filter, sensor window, and QE spec.
 Michael Hirsch 2014
@@ -30,12 +30,12 @@ def getSystemT(newLambda, bg3fn:Path, windfn,qefn,obsalt_km,zenang_deg, verbose=
 
     newLambda = asarray(newLambda)
 #%% atmospheric absorption
-    if golowtran is not None:
-        c1={'model':5,'itype':3,'iemsct':0,'h1':obsalt_km,'angle':zenang_deg,
+    if lowtran is not None:
+        c1={'model':5,'h1':obsalt_km,'angle':zenang_deg,
             'wlnmlim':(newLambda[0],newLambda[-1])}
         if verbose:
             print('loading LOWTRAN7 atmosphere model...')
-        atmT = golowtran(c1).loc[...,'transmission']
+        atmT = lowtran.transmittance(c1)['transmission'].squeeze()
         try:
             atmTcleaned = atmT.values.squeeze()
             atmTcleaned[atmTcleaned==0] = spacing(1) # to avoid log10(0)

setup.py

@@ -9,16 +9,17 @@
       packages=find_packages(),
       author='Michael Hirsch, Ph.D.',
       description='Gridding for auroral and ionospheric modeling',
+      long_description=open('README.rst').read(),
       url='https://github.com/scivision/gridaurora',
       version='1.1.8',
       package_data={'gridaurora':['data/RecentIndices.txt'],
                     'gridaurora.precompute': ['*.h5']},
       classifiers=[
-          'Intended Audience :: Science/Research',
           'Development Status :: 5 - Production/Stable',
+          'Intended Audience :: Science/Research',
           'License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)',
-          'Topic :: Scientific/Engineering :: Atmospheric Science',
           'Programming Language :: Python :: 3',
+          'Topic :: Scientific/Engineering :: Atmospheric Science',
           ],
       install_requires=install_requires,
       python_requires='>=3.6',