More updates to metadata column names

kboone · Apr 23, 2019 · f51bf4d · f51bf4d
1 parent c392892
commit f51bf4d
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 10 deletions.
diff --git a/avocado/dataset.py b/avocado/dataset.py
@@ -3,7 +3,7 @@
 from sklearn.model_selection import StratifiedKFold
 
 from .astronomical_object import AstronomicalObject
-from .logging import logger
+from .utils import logger
 from .settings import settings
 
 class Dataset():

diff --git a/avocado/plasticc.py b/avocado/plasticc.py
@@ -3,9 +3,10 @@
 import pandas as pd
 import os
 
-from . import Dataset, settings
+from .dataset import Dataset
+from .utils import settings, AvocadoException
 
-def update_plasticc_names(metadata, observations, datsaet_kind):
+def update_plasticc_names(metadata, observations, dataset_kind):
     """Rename columns in PLAsTiCC tables to follow the avocado naming scheme.
 
     Parameters
@@ -16,7 +17,7 @@ def update_plasticc_names(metadata, observations, datsaet_kind):
     observations : pandas.DataFrame
         Original observations DataFrame
 
-    dataset_kind : str {'trainng'}
+    dataset_kind : str {'training'}
 
     Returns
     -------
@@ -46,17 +47,16 @@ def update_plasticc_names(metadata, observations, datsaet_kind):
         'target': 'class',
         'mjd': 'time',
         'flux_err': 'flux_error',
+        'hostgal_specz': 'host_spectroscopic_redshift',
         'hostgal_photoz': 'host_photometric_redshift',
         'hostgal_photoz_err': 'host_photometric_redshift_error',
     }
 
-    if kind == 'training':
-        metadata_name_map['hostgal_specz'] = 'redshift'
-    else:
-        raise AvocadoException("Unknown dataset kind %s!" % kind)
-
     metadata.rename(metadata_name_map, axis=1, inplace=True)
 
+    if dataset_kind == 'training':
+        metadata['redshift'] = metadata['host_spectroscopic_redshift']
+
     return metadata, observations
 
 def load_training_set():
@@ -75,7 +75,8 @@ def load_training_set():
     metadata_path = os.path.join(data_directory, 'training_set_metadata.csv')
     metadata = pd.read_csv(metadata_path)
 
-    metadata, observations = update_plasticc_names(metadata, observations)
+    metadata, observations = update_plasticc_names(metadata, observations,
+                                                   'training')
 
     # Create a Dataset object
     dataset = Dataset('plasticc_trainng', metadata, observations)