Skip to content

Commit

Permalink
modify entityset function
Browse files Browse the repository at this point in the history
  • Loading branch information
Sara Pido authored and SaraPido committed Apr 2, 2024
1 parent a190905 commit 7e2e174
Showing 1 changed file with 15 additions and 4 deletions.
19 changes: 15 additions & 4 deletions zephyr_ml/entityset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down

0 comments on commit 7e2e174

Please sign in to comment.