From 0e536e146ffa0f2e699c9b7f2e51c8865f4d2000 Mon Sep 17 00:00:00 2001
From: Sarah Alnegheimish <sarahalnegheimish@gmail.com>
Date: Thu, 4 Jan 2024 22:56:23 +0700
Subject: [PATCH] add example vibrations

---
 notebooks/data/vibrations.csv |  5 +++++
 zephyr_ml/entityset.py        | 10 ++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)
 create mode 100644 notebooks/data/vibrations.csv

diff --git a/notebooks/data/vibrations.csv b/notebooks/data/vibrations.csv
new file mode 100644
index 0000000..6f3dd1d
--- /dev/null
+++ b/notebooks/data/vibrations.csv
@@ -0,0 +1,5 @@
+COD_ELEMENT,signal_id,timestamp,xvalues,yvalues
+0,S1,2022-01-02 13:21:01,"[2022-01-02 13:21:01, 2022-01-02 13:41:01, 2022-01-02 14:01:01]","[0.5, 0.4, 0.2]"
+0,S2,2022-03-08 11:00:00,"[2022-03-08 11:00:00, 2022-03-08 11:10:00, 2022-03-08 11:20:00]","[10, 10, 9]"
+1,S1,2022-01-02 13:21:01,"[2022-01-02 13:21:01, 2022-01-02 13:41:01, 2022-01-02 14:01:01]","[0.9, 0.95, 0.8]"
+1,S2,2022-03-08 11:00:00,"[2022-03-08 11:00:00, 2022-03-08 11:10:00, 2022-03-08 11:20:00]","[15, 11, 9]"
\ No newline at end of file
diff --git a/zephyr_ml/entityset.py b/zephyr_ml/entityset.py
index 533262f..d3e02dc 100644
--- a/zephyr_ml/entityset.py
+++ b/zephyr_ml/entityset.py
@@ -1,3 +1,5 @@
+from itertools import chain
+
 import featuretools as ft
 
 from zephyr_ml.metadata import get_mapped_kwargs
@@ -85,7 +87,7 @@ def create_vibrations_entityset(dfs, new_kwargs_mapping=None):
         **get_mapped_kwargs('pidata', new_kwargs_mapping),
         **get_mapped_kwargs('vibrations', new_kwargs_mapping),
     }
-    _validate_data(dfs, 'vibrations', entity_kwargs)
+    _validate_data(dfs, ['pidata', 'vibrations'], entity_kwargs)
 
     es = _create_entityset(dfs, 'vibrations', entity_kwargs)
     es.id = 'Vibrations data'
@@ -96,7 +98,11 @@ def create_vibrations_entityset(dfs, new_kwargs_mapping=None):
 def _validate_data(dfs, es_type, es_kwargs):
     '''Validate data by checking for required columns in each entity
     '''
-    entities = set(['alarms', 'stoppages', 'work_orders', 'notifications', 'turbines', es_type])
+    if not isinstance(es_type, list):
+        es_type = [es_type]
+
+    entities = set(chain(['alarms', 'stoppages', 'work_orders', 'notifications', 'turbines', *es_type]))
+
     if set(dfs.keys()) != entities:
         missing = entities.difference(set(dfs.keys()))
         extra = set(dfs.keys()).difference(entities)