Coverage for src / prepare_times_nz / utilities / timeslices.py: 100%

1""" 

2Shared helpers for constructing TIMES timeslices. 

3""" 

4 

5from __future__ import annotations 

6 

7from functools import cache 

8 

9import numpy as np 

10import pandas as pd 

11from prepare_times_nz.utilities.filepaths import DATA_RAW 

12 

13TIME_OF_DAY_FILE = DATA_RAW / "user_config/settings/time_of_day_types.csv" 

14 

15 

16@cache 

17def get_hour_to_time_map() -> dict[int, str]: 

18 """ 

19 Load the project hour-to-time-of-day mapping once per process. 

20 """ 

21 time_of_day_types = pd.read_csv(TIME_OF_DAY_FILE) 

22 return dict(zip(time_of_day_types["Hour"], time_of_day_types["Time_Of_Day"])) 

23 

24 

25def convert_hour_to_timeofday(df: pd.DataFrame, hour_col: str = "Hour") -> pd.DataFrame: 

26 """ 

27 Map integer hours to the project time-of-day categories. 

28 """ 

29 df["Time_Of_Day"] = df[hour_col].map(get_hour_to_time_map()) 

30 df["Time_Of_Day"] = df["Time_Of_Day"].fillna("N") 

31 return df 

32 

33 

34def convert_date_to_daytype( 

35 df: pd.DataFrame, date_col: str = "Trading_Date" 

36) -> pd.DataFrame: 

37 """ 

38 Map dates to project weekday and weekend labels. 

39 """ 

40 df[date_col] = pd.to_datetime(df[date_col]) 

41 weekday = df[date_col].dt.weekday 

42 df["Day_Type"] = np.select( 

43 [weekday.isin([5, 6]), weekday.isin([0, 1, 2, 3, 4])], 

44 ["WE-", "WK-"], 

45 default="ERROR", 

46 ) 

47 return df 

48 

49 

50def convert_date_to_season( 

51 df: pd.DataFrame, date_col: str = "Trading_Date" 

52) -> pd.DataFrame: 

53 """ 

54 Map dates to the project season labels. 

55 """ 

56 df[date_col] = pd.to_datetime(df[date_col]) 

57 month = df[date_col].dt.month 

58 df["Season"] = np.select( 

59 [ 

60 month.isin([12, 1, 2]), 

61 month.isin([3, 4, 5]), 

62 month.isin([6, 7, 8]), 

63 month.isin([9, 10, 11]), 

64 ], 

65 ["SUM-", "FAL-", "WIN-", "SPR-"], 

66 default="ERROR", 

67 ) 

68 return df 

69 

70 

71def create_timeslices( 

72 df: pd.DataFrame, date_col: str = "Trading_Date", hour_col: str = "Hour" 

73) -> pd.DataFrame: 

74 """ 

75 Add a TIMES-style `TimeSlice` column using the shared project mapping. 

76 """ 

77 df = convert_hour_to_timeofday(df, hour_col=hour_col) 

78 df = convert_date_to_daytype(df, date_col) 

79 df = convert_date_to_season(df, date_col) 

80 df["TimeSlice"] = df["Season"] + df["Day_Type"] + df["Time_Of_Day"] 

81 df = df.drop(columns=["Season", "Day_Type", "Time_Of_Day"]) 

82 return df