<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
# TODO rovib in H2O has wrong format for header
import urllib
import requests
import numpy as np
from astropy.table import Table

__all__ = ['LAMDAQuery']


class LAMDAQuery(object):
    """
    LAMDA query base class
    """
    def __init__(self):
        self.url = "http://home.strw.leidenuniv.nl/~moldata/datafiles/{}.dat"
        self.mols = {
            # Atoms
            'C': ['catom'],
            'C+': ['c+', 'c+@uv'],
            'O': ['oatom'],
            # Molecules
            'CO': ['co', '13co', 'c17o', 'c18o', 'co@neufold'],
            'CS': ['cs@xpol', '13cs@xpol', 'c34s@xpol'],
            'HCl': ['hcl', 'hcl@hfs'],
            'OCS': ['ocs@xpol'],
            'SO': ['so'],
            'SO2': ['so2@xpol'],
            'SiO': ['sio', '29sio'],
            'SiS': ['sis@xpol'],
            'SiC2': ['o-sic2'],
            'HCO+': ['hco+@xpol', 'h13co+@xpol', 'hc17o+@xpol', 'hc18o+@xpol',
                'dco+@xpol'],
            'N2H+': ['n2h+@xpol', 'n2h+_hfs'],
            'HCS+': ['hcs+@xpol'],
            'HC3N': ['hc3n'],
            'HCN': ['hcn', 'hcn@xpol', 'hcn@hfs', 'h13cn@xpol', 'hc15n@xpol'],
            'HNC': ['hnc'],
            'C3H2': ['p-c3h2', 'o-c3h2'],
            'H2O': ['<EMAIL>', '<EMAIL>', '<EMAIL>', '<EMAIL>'],
            'H2CO': ['p-h2co', 'o-h2co'],
            'OH': ['oh', 'oh@hfs'],
            'CH3OH': ['e-ch3oh', 'a-ch3oh'],
            'NH3': ['p-nh3', 'o-nh3'],
            'HDO': ['hdo'],
            'H3O+': ['p-h3o+', 'o-h3o+'],
            'HNCO': ['hnco'],
            'NO': ['no'],
            'CN': ['cn'],
            'CH3CN': ['ch3cn'],
            'O2': ['o2'],
            'HF': ['hf']
             }
        self.query_types = {
             'erg_levels': '!NUMBER OF ENERGY LEVELS',
             'rad_trans': '!NUMBER OF RADIATIVE TRANSITIONS',
             'coll_rates': '!COLLISIONS BETWEEN'
             }

    def print_mols(self):
        """
        Print molecule names available for query.
        """
        mols = self.mols
        for mol_family in mols.keys():
            print '-- {} :'.format(mol_family)
            print mols[mol_family], '\n'

    def lamda_query(self, mol, query_type, coll_partner_index=0):
        """
        Query the LAMDA database.

        Parameters
        ----------
        mol : string
            Molecule or atom designation. For a list of valid designations see
            the :meth:``print_mols`` method.
        query_type : string
            Query type to execute. Valid options:
                'erg_levels' -> Energy Levels
                'rad_trans'  -> Radiative transitions
                'coll_rates' -> Collisional rates
        coll_partner_index : number, default 0
            If collisional rates are queried, set the index of the collisional
            partner to the order found in the LAMDA datafile. If no index is
            specified, the first collsional partner is taken.

        Returns
        -------
        table : astropy.table.Table

        Examples
        --------
        >>> t = lamda_query(mol='co', query_type='erg_levels')
        >>> t.pprint()
        LEVEL ENERGIES(cm^-1) WEIGHT  J
        ----- --------------- ------ ---
            2     3.845033413    3.0   1
            3    11.534919938    5.0   2
          ...             ...    ... ...
        """
        if query_type not in self.query_types.keys():
            raise ValueError
        # Send HTTP request to open URL
        datafile = np.array([s.strip() for s in
            urllib.urlopen(self.url.format(mol)).readlines()])
        # Parse datafile string list and return a table
        table = self._parse_datafile(datafile, query_type=query_type,
            coll_partner_index=coll_partner_index)
        return table

    def _parse_datafile(self, datafile, query_type, coll_partner_index=0):
        """
        Parse datafile according to query type and return a Table instance of
        the data.

        Parameters
        ----------
        datafile : np.ndarray
            Raw data array pulled from LAMDA
        query_type : string
            Valid query_type in ['coll_rates', 'rad_trans', 'erg_levels']
        coll_partner_index : number, default 0
            Collision partner index, order of partner in file

        Returns<commit_msg>Convert to requests and formatting edits<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
# TODO rovib in H2O has wrong format for header
import requests
import numpy as np
from astropy.table import Table

__all__ = ['LAMDAQuery']


class LAMDAQuery(object):
    """
    LAMDA query base class
    """
    def __init__(self):
        self.url = "http://home.strw.leidenuniv.nl/~moldata/datafiles/{}.dat"
        self.mols = {
            # Atoms
            'C': ['catom'],
            'C+': ['c+', 'c+@uv'],
            'O': ['oatom'],
            # Molecules
            'CO': ['co', '13co', 'c17o', 'c18o', 'co@neufold'],
            'CS': ['cs@xpol', '13cs@xpol', 'c34s@xpol'],
            'HCl': ['hcl', 'hcl@hfs'],
            'OCS': ['ocs@xpol'],
            'SO': ['so'],
            'SO2': ['so2@xpol'],
            'SiO': ['sio', '29sio'],
            'SiS': ['sis<EMAIL>'],
            'SiC2': ['o-sic2'],
            'HCO+': ['hco+@xpol', 'h13co+@xpol', 'hc17o+@xpol', 'hc18o+@xpol',
                'dco+@xpol'],
            'N2H+': ['n2h+@xpol', 'n2h+_hfs'],
            'HCS+': ['hcs+@xpol'],
            'HC3N': ['hc3n'],
            'HCN': ['hcn', 'hcn@xpol', 'hcn@hfs', 'h13cn@xpol', 'hc15n@xpol'],
            'HNC': ['hnc'],
            'C3H2': ['p-c3h2', 'o-c3h2'],
            'H2O': ['ph<EMAIL>', 'oh2o@<EMAIL>', 'ph2o@<EMAIL>', 'oh2o@<EMAIL>'],
            'H2CO': ['p-h2co', 'o-h2co'],
            'OH': ['oh', 'oh@hfs'],
            'CH3OH': ['e-ch3oh', 'a-ch3oh'],
            'NH3': ['p-nh3', 'o-nh3'],
            'HDO': ['hdo'],
            'H3O+': ['p-h3o+', 'o-h3o+'],
            'HNCO': ['hnco'],
            'NO': ['no'],
            'CN': ['cn'],
            'CH3CN': ['ch3cn'],
            'O2': ['o2'],
            'HF': ['hf']
             }
        self.query_types = {
             'erg_levels': '!NUMBER OF ENERGY LEVELS',
             'rad_trans': '!NUMBER OF RADIATIVE TRANSITIONS',
             'coll_rates': '!COLLISIONS BETWEEN'
             }

    def print_mols(self):
        """
        Print molecule names available for query.
        """
        mols = self.mols
        for mol_family in mols.keys():
            print '-- {} :'.format(mol_family)
            print mols[mol_family], '\n'

    def lamda_query(self, mol, query_type, coll_partner_index=0):
        """
        Query the LAMDA database.

        Parameters
        ----------
        mol : string
            Molecule or atom designation. For a list of valid designations see
            the :meth:``print_mols`` method.
        query_type : string
            Query type to execute. Valid options:
                'erg_levels' -> Energy Levels
                'rad_trans'  -> Radiative transitions
                'coll_rates' -> Collisional rates
        coll_partner_index : number, default 0
            If collisional rates are queried, set the index of the collisional
            partner to the order found in the LAMDA datafile. If no index is
            specified, the first collsional partner is taken.

        Returns
        -------
        table : astropy.table.Table

        Examples
        --------
        >>> Q = LAMDAQuery()
        >>> t = Q.lamda_query(mol='co', query_type='erg_levels')
        >>> t.pprint()
        LEVEL ENERGIES(cm^-1) WEIGHT  J
        ----- --------------- ------ ---
            2     3.845033413    3.0   1
            3    11.534919938    5.0   2
          ...             ...    ... ...
        """
        if query_type not in self.query_types.keys():
            raise ValueError
        # Send HTTP request to open URL
        datafile = np.array([s.strip() for s in
            requests.get(self.url.format(mol)).iter_lines()])
        # Parse datafile string list and return a table
        table = self._parse_datafile(datafile, query_type=query_type,
            coll_partner_index=coll_partner_index)
        return table

    def _parse_datafile(self, datafile, query_type, coll_partner_index=0):
        """
        Parse datafile according to query type and return a Table instance of
        the data.

        Parameters
        ----------
        datafile : np.ndarray
            Raw data array pulled from LAMDA
        query_type : string
            Valid query_type in ['coll_rates', 'rad_trans', 'erg_levels']
        coll_partner_index : number, default 0
            Collision partner index, order of partner in file

        Returns
{'fraction_non_alphanumeric': 0.16150542495479203, 'fraction_numerical': 0.020682640144665462, 'mean_word_length': 1.6359845099791481, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '627735', 'n_tokens_mistral': 2991, 'n_tokens_neox': 2844, 'n_words': 704}