modify entityset function

zephyr_ml/entityset.py

@@ -80,14 +80,25 @@ def create_vibrations_entityset(dfs, new_kwargs_mapping=None):
 
     Args:
         data_paths (dict): Dictionary mapping entity names ('alarms', 'notifications',
-            'stoppages', 'work_orders', 'pidata', 'vibrations', 'turbines') to the pandas
-            dataframe for that entity.
+            'stoppages', 'work_orders', 'vibrations', 'turbines') to the pandas
+            dataframe for that entity. Optionally 'pidata' and 'scada' can be included.
     '''
+    entities = ['vibrations']
+
+    pidata_kwargs, scada_kwargs = {}, {}
+    if 'pidata' in dfs:
+        pidata_kwargs = get_mapped_kwargs('pidata', new_kwargs_mapping)
+        entities.append('pidata')
+    if 'scada' in dfs:
+        pidata_kwargs = get_mapped_kwargs('scada', new_kwargs_mapping)
+        entities.append('scada')
+
     entity_kwargs = {
-        **get_mapped_kwargs('pidata', new_kwargs_mapping),
+        **pidata_kwargs,
+        **scada_kwargs,
         **get_mapped_kwargs('vibrations', new_kwargs_mapping),
     }
-    _validate_data(dfs, ['pidata', 'vibrations'], entity_kwargs)
+    _validate_data(dfs, entities , entity_kwargs)
 
     es = _create_entityset(dfs, 'vibrations', entity_kwargs)
     es.id = 'Vibrations data'