apply black

jvshields · Aug 19, 2024 · 0a4831f · 0a4831f
1 parent 63eb762
commit 0a4831f
Showing 1 changed file with 25 additions and 0 deletions.
diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py
@@ -221,6 +221,17 @@ def from_hdf(cls, fname=None):
             if "linelist" in store:
                 dataframes["linelist"] = store["linelist"]
 
+            if "molecules" in store:
+                dataframes["molecular_equilibrium_constants"] = store[
+                    "molecules/equilibrium_constants"
+                ]
+                dataframes["molecular_partition_functions"] = store[
+                    "molecules/partition_functions"
+                ]
+                dataframes["molecular_dissociation_energies"] = store[
+                    "molecules/dissociation_energies"
+                ]
+
             atom_data = cls(**dataframes)
 
             try:
@@ -284,6 +295,9 @@ def __init__(
         two_photon_data=None,
         linelist=None,
         decay_radiation_data=None,
+        molecular_equilibrium_constants=None,
+        molecular_partition_functions=None,
+        molecular_dissociation_energies=None,
     ):
         self.prepared = False
 
@@ -346,6 +360,17 @@ def __init__(
         if linelist is not None:
             self.linelist = linelist
 
+        if molecular_equilibrium_constants is not None:
+            self.molecular_equilibrium_constants = (
+                molecular_equilibrium_constants
+            )
+        if molecular_partition_functions is not None:
+            self.molecular_partition_functions = molecular_partition_functions
+        if molecular_dissociation_energies is not None:
+            self.molecular_dissociation_energies = (
+                molecular_dissociation_energies
+            )
+
         if decay_radiation_data is not None:
             self.decay_radiation_data = decay_radiation_data
         self._check_related()