import xarray as xr
import pandas as pd
import numpy as np
import xskillscore as xs
from sklearn.metrics import mean_squared_error
stores = np.arange(100)
skus = np.arange(100)
dates = pd.date_range("1/1/2020", "1/10/2020", freq="D")
rows = []
for _, date in enumerate(dates):
for _, store in enumerate(stores):
for _, sku in enumerate(skus):
rows.append(
dict(
{
"DATE": date,
"STORE": store,
"SKU": sku,
"QUANTITY_SOLD": np.random.randint(9) + 1,
}
)
)
df = pd.DataFrame(rows)
df.rename(columns={"QUANTITY_SOLD": "y"}, inplace=True)
df.set_index(['DATE', 'STORE', 'SKU'], inplace=True)
noise = np.random.uniform(-1, 1, size=len(df['y']))
df['yhat'] = (df['y'] + (df['y'] * noise)).astype(int)
df.groupby(['STORE', 'SKU']).apply(lambda x: mean_squared_error(x.y, x.yhat))
ds = df.to_xarray()
ds.xs.mse('y', 'yhat', 'DATE')