Files
drift/training/walk_forward/inference_parallel.py
T
Mark Aron Szulyovszky 8dd2d88740 chore(Linter): reformatted code with black (#211)
* chore(Linter): reformatted code with black

* Create black.yaml
2022-02-17 19:22:17 +01:00

115 lines
3.6 KiB
Python

import pandas as pd
from models.base import Model
from training.types import (
ModelOverTime,
TransformationsOverTime,
PredictionsSeries,
ProbabilitiesDataFrame,
)
from transformations.base import Transformation
from utils.helpers import get_first_valid_return_index
from tqdm import tqdm
from typing import Optional
from data_loader.types import XDataFrame
import ray
def walk_forward_inference(
model_name: str,
model_over_time: ModelOverTime,
transformations_over_time: TransformationsOverTime,
X: XDataFrame,
expanding_window: bool,
window_size: int,
retrain_every: int,
from_index: Optional[pd.Timestamp],
) -> tuple[PredictionsSeries, ProbabilitiesDataFrame]:
predictions = pd.Series(index=X.index, dtype="object").rename(model_name)
probabilities = pd.DataFrame(index=X.index)
inference_from = (
max(
get_first_valid_return_index(model_over_time),
get_first_valid_return_index(X.iloc[:, 0]),
)
if from_index is None
else X.index.to_list().index(from_index)
)
inference_till = X.shape[0]
first_model = model_over_time[inference_from]
if first_model.only_column is not None:
X = X[[column for column in X.columns if first_model.only_column in column]]
if first_model.data_transformation == "original":
transformations_over_time = []
batch_size = int((inference_till - inference_from) / 10)
batched_results = ray.get(
[
__inference_from_window.remote(
index,
index + batch_size,
inference_from,
retrain_every,
X,
model_over_time,
transformations_over_time,
expanding_window,
window_size,
)
for index in range(inference_from, inference_till)
]
)
for batch in batched_results:
for index, prediction, probs in batch:
predictions[X.index[index]] = prediction
probabilities.loc[X.index[index]] = probs
return predictions, probabilities
@ray.remote
def __inference_from_window(
index_start: int,
index_end: int,
inference_from: int,
retrain_every: int,
X: XDataFrame,
model_over_time: ModelOverTime,
transformations_over_time: TransformationsOverTime,
expanding_window: bool,
window_size: int,
) -> list[tuple[int, float, pd.Series]]:
results = []
for index in range(index_start, index_end):
last_model_index = index - ((index - inference_from) % retrain_every)
train_window_start = (
X.index[inference_from]
if expanding_window
else X.index[index - window_size - 1]
)
current_model = model_over_time[X.index[last_model_index]]
current_transformations = [
transformation_over_time[X.index[last_model_index]]
for transformation_over_time in transformations_over_time
]
if current_model.predict_window_size == "window_size":
next_timestep = X.loc[train_window_start : X.index[index]]
else:
# we need to get a Dataframe out of it, since the transformation step always expects a 2D array, but it's equivalent to X.iloc[index]
next_timestep = X.loc[X.index[index] : X.index[index]]
for transformation in current_transformations:
next_timestep = transformation.transform(next_timestep)
next_timestep = next_timestep.to_numpy()
prediction, probs = current_model.predict(next_timestep)
results.append((index, prediction, probs))
return results