You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

58 lines
1.7 KiB

import models
6 years ago
import logging
6 years ago
import config
6 years ago
import pandas as pd
from detectors import Detector
6 years ago
6 years ago
logger = logging.getLogger('PATTERN_DETECTOR')
6 years ago
def resolve_model_by_pattern(pattern: str) -> models.Model:
if pattern == 'PEAK':
return models.PeaksModel()
if pattern == 'DROP':
return models.StepModel()
if pattern == 'JUMP':
return models.JumpModel()
if pattern == 'CUSTOM':
6 years ago
return models.CustomModel()
raise ValueError('Unknown pattern "%s"' % pattern)
6 years ago
class PatternDetector(Detector):
6 years ago
def __init__(self, pattern_type):
self.pattern_type = pattern_type
self.model = resolve_model_by_pattern(self.pattern_type)
window_size = 100
6 years ago
async def train(self, dataframe: pd.DataFrame, segments: list, cache: dict):
# TODO: pass only part of dataframe that has segments
self.model.fit(dataframe, segments)
# TODO: save model after fit
return cache
6 years ago
async def predict(self, dataframe: pd.DataFrame, cache: dict):
predicted_indexes = await self.model.predict(dataframe)
6 years ago
segments = []
# for time_value in predicted_times:
# ts1 = int(time_value[0].timestamp() * 1000)
# ts2 = int(time_value[1].timestamp() * 1000)
# segments.append({
# 'start': min(ts1, ts2),
# 'finish': max(ts1, ts2)
# })
6 years ago
last_dataframe_time = dataframe.iloc[-1]['timestamp']
6 years ago
last_prediction_time = int(last_dataframe_time.timestamp() * 1000)
return {
'cache': cache,
'segments': segments,
'last_prediction_time': last_prediction_time
}