NIST parser new HDF output

carsus/io/nist/__init__.py

@@ -1,2 +1,8 @@
-from .weightscomp import download_weightscomp, NISTWeightsCompPyparser, NISTWeightsCompIngester
-from .ionization import download_ionization_energies, NISTIonizationEnergiesParser, NISTIonizationEnergiesIngester
+from carsus.io.nist.weightscomp import (download_weightscomp, 
+                                   NISTWeightsCompPyparser, 
+                                   NISTWeightsCompIngester,
+                                   NISTWeightsComp)
+from carsus.io.nist.ionization import (download_ionization_energies, 
+                                       NISTIonizationEnergiesParser, 
+                                       NISTIonizationEnergiesIngester,
+                                       NISTIonizationEnergies)

carsus/io/nist/ionization.py

@@ -4,6 +4,7 @@
 """
 
 import requests
+import numpy as np
 import pandas as pd
 
 from bs4 import BeautifulSoup
@@ -290,3 +291,59 @@ def ingest(self, ionization_energies=True, ground_levels=True):
         if ground_levels:
             self.ingest_ground_levels()
             self.session.flush()
+
+
+class NISTIonizationEnergies(BaseParser):
+    """
+    Attributes
+    ----------
+    base : pandas.Series
+
+    Methods
+    -------
+    to_hdf(fname)
+        Dump the `base` attribute into an HDF5 file
+
+    """       
+    def __init__(self, spectra):
+        input_data = download_ionization_energies(spectra)
+        self.parser = NISTIonizationEnergiesParser(input_data)
+        self._prepare_data()
+
+    def _prepare_data(self):
+        ionization_data = pd.DataFrame()
+        ionization_data['atomic_number'] = self.parser.base['atomic_number']
+        ionization_data['ion_number'] = self.parser.base['ion_charge'] + 1
+        ionization_data['ionization_energy'] = self.parser.base['ionization_energy_str'].str.strip('[]()').astype(np.float64)
+        ionization_data.set_index(['atomic_number', 'ion_number'], inplace=True)
+
+        # `base` attribute is a Series object
+        self.base = ionization_data['ionization_energy']
+
+    def get_ground_levels(self):
+        """Returns a DataFrame with the ground levels for the selected spectra
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with ground levels
+        """
+        levels = self.parser.prepare_ground_levels()
+        levels['g'] = 2*levels['J'] + 1
+        levels['g'] = levels['g'].astype(np.int)
+        levels['energy'] = 0.
+        levels = levels[['g','energy']]
+        levels = levels.reset_index()
+
+        return levels
+
+    def to_hdf(self, fname):
+        """Dump the `base` attribute into an HDF5 file
+
+        Parameters
+        ----------
+        fname : path
+           Path to the HDF5 output file
+        """
+        with pd.HDFStore(fname, 'a') as f:
+            f.append('/ionization_data', self.base)

carsus/io/nist/weightscomp.py

@@ -8,9 +8,12 @@
 
 from bs4 import BeautifulSoup
 from astropy import units as u
+from carsus.base import basic_atomic_data_fname
 from carsus.model import AtomWeight
-from carsus.io.base import BasePyparser, BaseIngester
+from carsus.util import parse_selected_atoms
+from carsus.io.base import BasePyparser, BaseIngester, BaseParser
 from carsus.io.util import to_nom_val_and_std_dev
+from carsus.util.helpers import ATOMIC_SYMBOLS_DATA
 from carsus.io.nist.weightscomp_grammar import isotope, COLUMNS, ATOM_NUM_COL, MASS_NUM_COL,\
     AM_VAL_COL, AM_SD_COL, INTERVAL, STABLE_MASS_NUM, ATOM_WEIGHT_COLS, AW_STABLE_MASS_NUM_COL,\
     AW_TYPE_COL, AW_VAL_COL, AW_SD_COL, AW_LWR_BND_COL, AW_UPR_BND_COL
@@ -183,3 +186,56 @@ def ingest(self, atomic_weights=True):
         if atomic_weights:
             self.ingest_atomic_weights()
             self.session.flush()
+
+
+class NISTWeightsComp(BaseParser):
+    """
+    Attributes
+    ----------
+    base : pandas.DataFrame
+
+    columns : list of str
+
+    Methods
+    -------
+    to_hdf(fname)
+        Dump the `base` attribute into an HDF5 file
+
+    """    
+    def __init__(self, atoms):
+        input_data = download_weightscomp()
+        self.parser = NISTWeightsCompPyparser(input_data=input_data)
+        self._prepare_data(atoms)
+
+    def _prepare_data(self, atoms):
+        atomic_numbers = parse_selected_atoms(atoms)
+        atom_data_list = []
+
+        for atomic_number in atomic_numbers:
+            basic_atomic_data = pd.read_csv(basic_atomic_data_fname)
+            basic_atomic_data = basic_atomic_data.loc[atomic_number-1]
+
+            atom_masses = self.parser.prepare_atomic_dataframe()
+            atom_masses = atom_masses.drop(columns='atomic_weight_std_dev')
+            atom_masses = atom_masses.rename(columns={'atomic_weight_nominal_value': 'mass'})
+
+            data = atom_masses.loc[[(atomic_number)]]
+            data['symbol'] = basic_atomic_data['symbol']
+            data['name'] = basic_atomic_data['name']
+
+            atom_data_list.append(data)
+
+        atom_data = pd.concat(atom_data_list)
+        self.base = atom_data[['symbol', 'name', 'mass']]
+        self.columns = atom_data.columns
+
+    def to_hdf(self, fname):
+        """Dump the `base` attribute into an HDF5 file
+
+        Parameters
+        ----------
+        fname : path
+           Path to the HDF5 output file
+        """
+        with pd.HDFStore(fname, 'a') as f:
+            f.append('/atom_data', self.base, min_itemsize={'symbol': 2, 'name': 15})