Add molecular data exposure to tardis atom data

tardis/io/atom_data/base.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 from astropy.units import Quantity
+from dataclasses import dataclass
 
 from tardis import constants as const
 from tardis.io.atom_data.collision_data import (
@@ -103,6 +104,11 @@
         columns: atomic_number, element, Rad energy, Rad intensity decay mode.
         Curated from nndc
 
+    molecule_data : MolecularData
+    A class containing the *molecular data* with:
+        equilibrium_constants, partition_functions, dissociation_energies
+
+
     Attributes
     ----------
     prepared : bool
@@ -117,6 +123,7 @@
     photoionization_data : pandas.DataFrame
     two_photon_data : pandas.DataFrame
     decay_radiation_data : pandas.DataFrame
+    molecule_data : MolecularData
 
     Methods
     -------
@@ -222,7 +229,16 @@
             if "linelist" in store:
                 dataframes["linelist"] = store["linelist"]
 
-            atom_data = cls(**dataframes)
+            if "molecules" in store:
+                molecule_data = MoleculeData(
+                    store["molecules/equilibrium_constants"],
+                    store["molecules/partition_functions"],
+                    store["molecules/dissociation_energies"],
+                )
+            else:
+                molecule_data = None
+
+            atom_data = cls(**dataframes, molecule_data=molecule_data)
 
             atom_data.uuid1 = cls.get_attributes_from_store(store, "uuid1")
             atom_data.md5 = cls.get_attributes_from_store(store, "md5")
@@ -261,6 +277,7 @@
         two_photon_data=None,
         linelist=None,
         decay_radiation_data=None,
+        molecule_data=None,
     ):
         self.prepared = False
 
@@ -332,6 +349,9 @@
         if linelist is not None:
             self.linelist = linelist
 
+        if molecule_data is not None:
+            self.molecule_data = molecule_data
+
         if decay_radiation_data is not None:
             self.decay_radiation_data = decay_radiation_data
         self._check_related()
@@ -672,3 +692,29 @@
             attribute = None
 
         return attribute
+
+
+@dataclass
+class MoleculeData:
+    """
+    Class to hold molecular data. Held by the AtomData object.
+
+    equilibrium_constants : pandas.DataFrame
+    A DataFrame containing the *molecular equilibrium constants* with:
+        index: molecule
+        columns: temperatures
+
+    partition_functions : pandas.DataFrame
+    A DataFrame containing the *molecular partition functions* with:
+        index: molecule
+        columns: temperatures
+
+    dissociation_energies : pandas.DataFrame
+    A DataFrame containing the *molecular dissociation energies* with:
+        index: molecule
+
+    """
+
+    equilibrium_constants: pd.DataFrame
+    partition_functions: pd.DataFrame
+    dissociation_energies: pd.DataFrame