Fix assignment to unbound variable

[nickto]

When specifying is_balanced=True in aggregate function, we get the following error:

UnboundLocalError: local variable 'dates' referenced before assignment

This happens because dates is defined only in not balanced case.

    if not is_balanced:
        dates         = Y_bottom_df['ds'].unique()
        [...]
    else:
        balanced_df = Y_bottom_df.copy()
        # `dates` is not defined in this if-else branch 

Suggested fix: define it in the balanced case too:

    if not is_balanced:
        dates         = Y_bottom_df['ds'].unique()
        [...]
    else:
        dates       = Y_bottom_df['ds'].unique()
        balanced_df = Y_bottom_df.copy()

Test

Code which fails before the proposed fix:

import random
from datetime import date, timedelta

import pandas as pd

from hierarchicalforecast.utils import aggregate

# Create a sample data frame
random.seed(2023)
T = 5
n_bottom = 4

df = pd.DataFrame(
    {
        "level_1": ["a"] * T * 2 + ["b"] * T * 2,
        "level_2": [chr(ord("c") + i) for i in range(n_bottom) for _ in range(T)],
        "y": [random.randint(2, 4) for i in range(n_bottom) for _ in range(T)],
        "ds": [date(year=2023, month=1, day=1) + timedelta(days=Δ) for Δ in range(T)] * n_bottom,
    }
)
print(df)
#   level_1 level_2  y          ds
# 0       a       c  3  2023-01-01
# 1       a       c  4  2023-01-02
# 2       a       d  3  2023-01-01
# 3       a       d  3  2023-01-02
# 4       b       e  3  2023-01-01
# 5       b       e  4  2023-01-02
# 6       b       f  4  2023-01-01
# 7       b       f  3  2023-01-02

# Aggregate
hierarchy_levels = [
    ["level_1"],
    ["level_1", "level_2"],
]

df, S, tags = aggregate(df=df, spec=hierarchy_levels, is_balanced=True)

results in

UnboundLocalError: local variable 'dates' referenced before assignment

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[AzulGarza]

hey @nickto! Thank you for letting us know about the issue and solving it.🎉

LGTM! :)