Source code for cripticpy.parseinput

"""
A little utility to parse criptic input files
"""

import warnings
import numpy as np
try:
    import astropy.units as u
    import astropy.constants as const
    from astropy.constants import m_p, m_e, c
except:
    u = None



def parseinput(fname = "criptic.in", units = "cgs"):
    """
    Parse a criptic input file

    Parameters
       fname : string
          name of input file to parse
       units : "cgs" | "mks" | None
          unit system to use when assigning units to quantities; if
          None, then quantities are not assigned units

    Returns
       params : dict
          A dict containing the input parameters, with the parameter
          names as keys
    """

    # Open the file
    fp = open(fname, "r")

    # Set unit choices
    if units is None:
        lunit = 1
        tunit = 1
        munit = 1
        Bunit = 1
        Tunit = 1
        m_p = 1
        c_ = 1
    else:
        if u is None:
            warnings.warn("unit'ed output requested, but could not "
                          "import astropy.units; results will be "
                          "without units")
            lunit = 1
            tunit = 1
            munit = 1
            Bunit = 1
            Tunit = 1
            m_p = 1
            c_ = 1
        elif units == 'cgs':
            lunit = u.cm
            tunit = u.s
            munit = u.g
            Bunit = u.G
            Tunit = u.K
            c_ = c
        else:
            lunit = u.m
            tunit = u.s
            munit = u.kg
            Bunit = u.T
            Tunit = u.K
            c_ = c

    # Make output holder
    out = { }

    # Loop through it
    for line in fp:
        l = line.split('#')[0].strip()    # Strip comments and whitespace
        if len(l) == 0: continue          # Skip blank lines
        s = l.split()                     # Split by whitespace

        # Store keyword and value
        if len(s) < 2:
            raise ValueError("input file format not understood "
                             "at line " + line)
        k = s[0]
        out[k] = s[1:]
        
        # Try type conversions where possible
        for i in range(len(out[k])):
            try:
                out[k][i] = int(out[k][i])
            except:
                try:
                    out[k][i] = float(out[k][i])
                except:
                    pass

        # Reduce 1-element lists to non-lists
        if len(out[k]) == 1:
            out[k] = out[k][0]

    # Close file and return
    fp.close()

    # Apply units to known keywords; silently skip anything that
    # doesn't work, since users may define their own keywords that
    # don't match our expectations
    for k in out.keys():
        if k == 'dt_init' or \
           k == 'max_time' or \
           k == 'min_dt' or \
           k == 'output.chk_dt':
            try: out[k] = out[k] * tunit
            except: pass
        elif k == 'integrator.packet_rate':
            try: out[k] = out[k] / tunit
            except: pass
        elif k == 'geometry.prob_lo' or \
             k == 'geometry.prob_hi' or \
             k == 'tree.source_dither' or \
             k == 'gas.dx' or \
             k == 'gas.r0' or \
             k == 'gas.rflat':
            try: out[k] = out[k] * lunit
            except: pass
        elif k == 'cr.kPar0' or \
             k == 'cr.kPerp0':
            try: out[k] = out[k] * lunit**2 / tunit
            except: pass
        elif k == 'cr.vStr0':
            try:
                if 'cr.vAStream' in out:
                    if out['cr.vAStream'] == 0:
                        out[k] = out[k] * lunit / tunit
            except: pass
        elif k == 'cr.kPP0':
            try: out[k] = out[k] * (u.GeV / c_)**2 / tunit
            except: pass
        elif k == 'cr.kMu0':
            try: out[k] = out[k] / tunit
            except: pass
        elif k == 'gas.density' or \
             k == 'gas.ionDensity' or \
             k == 'gas.rho0':
            try: out[k] = out[k] * munit / lunit**3
            except: pass
        elif k == 'gas.velocity':
            try:
                for i in range(len(out[k])):
                    out[k][i] = out[k][i] * lunit / tunit
            except: pass
        elif k == 'gas.B0':
            try: out[k] = out[k] * Bunit
            except: pass
        elif k == 'gas.magField':
            try:
                out[k] = np.array(out[k]) * Bunit
            except: pass
        elif k == 'gas.TBB':
            try:
                out[k] = np.array(out[k]) * Tunit
            except: pass
        elif k == 'gas.WBB':
            try:
                out[k] = np.array(out[k])
            except: pass

    # Return
    return out