Merge pull request #14 from sonachitchyan/carsus_format_support

Carsus format support
Rodot- · May 16, 2022 · 916b291 · 916b291
2 parents 2483516 + 2121719
commit 916b291
Showing 1 changed file with 34 additions and 21 deletions.
diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py
@@ -2,6 +2,8 @@
 
 
 import logging
+from multiprocessing.sharedctypes import Value
+from matplotlib import collections
 import numpy as np
 import pandas as pd
 
@@ -165,44 +167,55 @@ def from_hdf(cls, fname=None):
         fname = resolve_atom_data_fname(fname)
 
         with pd.HDFStore(fname, "r") as store:
+            if "metadata" in store:
+                carsus_version = store["metadata"].loc[('format', 'version')].value
+                if carsus_version == "1.0":
+                    if "collisions_data" in store:
+                        try:
+
+                            dataframes["collision_data_temperatures"] = store[
+                                "collisions_metadata"
+                            ].temperatures
+                            if "cmfgen" in store["collisions_metadata"].dataset:
+                                dataframes["yg_data"] = store["collisions_data"]
+                                dataframes["collision_data"] = "dummy value"
+                            elif "chianti" in store["collisions_metadata"].dataset:
+                                dataframes["collision_data"] = store["collisions_data"]
+                            else:
+                                raise KeyError(
+                                    "Atomic Data Collisions Not a Valid Chanti or CMFGEN Carsus Data File"
+                                )
+                        except KeyError as e:
+                            logger.warn(
+                                "Atomic Data is not a Valid Carsus Atomic Data File"
+                            )
+                            raise
+                    dataframes['levels'] = store['levels_data']
+                    dataframes['lines'] = store['lines_data']
+                else:
+                    raise ValueError(f"Current carsus version, {carsus_version}, is not supported.")
+
+
             for name in cls.hdf_names:
                 try:
                     dataframes[name] = store[name]
                 except KeyError:
                     logger.debug(f"Dataframe does not contain {name} column")
                     nonavailable.append(name)
             # Checks for various collisional data from Carsus files
-            if "collisions" in store:
-                try:
-                    dataframes["collision_data_temperatures"] = store[
-                        "collisions_metadata"
-                    ].temperatures
-                    if store["collisions_metadata"].dataset == "cmfgen":
-                        dataframes["yg_data"] = store["collisions"]
-                        dataframes["collision_data"] = "dummy value"
-                    elif store["collisions_metadata"].dataset == "chianti":
-                        dataframes["collision_data"] = store["collisions"]
-                    else:
-                        raise KeyError(
-                            "Atomic Data Collisions Not a Valid Chanti or CMFGEN Carsus Data File"
-                        )
-                except KeyError as e:
-                    logger.warn(
-                        "Atomic Data is not a Valid Carsus Atomic Data File"
-                    )
-                    raise
+
             atom_data = cls(**dataframes)
 
             try:
-                atom_data.uuid1 = store.root._v_attrs["uuid1"].decode("ascii")
+                atom_data.uuid1 = store.root._v_attrs["uuid1"]
             except KeyError:
                 logger.debug(
                     "UUID not available for Atom Data. Setting value to None"
                 )
                 atom_data.uuid1 = None
 
             try:
-                atom_data.md5 = store.root._v_attrs["md5"].decode("ascii")
+                atom_data.md5 = store.root._v_attrs["md5"]
             except KeyError:
                 logger.debug(
                     "MD5 not available for Atom Data. Setting value to None"