Refactor datetime and timezone utils
This commit is contained in:
Urtzi Alfaro
270
services/training/app/utils/ml_datetime.py
Normal file
270
services/training/app/utils/ml_datetime.py
Normal file
@@ -0,0 +1,270 @@
|
||||
"""
|
||||
ML-Specific DateTime Utilities
|
||||
|
||||
DateTime utilities for machine learning operations, specifically for:
|
||||
- Prophet forecasting model (requires timezone-naive datetimes)
|
||||
- Pandas DataFrame datetime operations
|
||||
- Time series data processing
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Union
|
||||
import pandas as pd
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def ensure_timezone_aware(dt: datetime, default_tz=timezone.utc) -> datetime:
|
||||
"""
|
||||
Ensure a datetime is timezone-aware.
|
||||
|
||||
Args:
|
||||
dt: Datetime to check
|
||||
default_tz: Timezone to apply if datetime is naive (default: UTC)
|
||||
|
||||
Returns:
|
||||
Timezone-aware datetime
|
||||
"""
|
||||
if dt is None:
|
||||
return None
|
||||
|
||||
if dt.tzinfo is None:
|
||||
return dt.replace(tzinfo=default_tz)
|
||||
return dt
|
||||
|
||||
|
||||
def ensure_timezone_naive(dt: datetime) -> datetime:
|
||||
"""
|
||||
Remove timezone information from a datetime.
|
||||
|
||||
Args:
|
||||
dt: Datetime to process
|
||||
|
||||
Returns:
|
||||
Timezone-naive datetime
|
||||
"""
|
||||
if dt is None:
|
||||
return None
|
||||
|
||||
if dt.tzinfo is not None:
|
||||
return dt.replace(tzinfo=None)
|
||||
return dt
|
||||
|
||||
|
||||
def normalize_datetime_to_utc(dt: Union[datetime, pd.Timestamp]) -> datetime:
|
||||
"""
|
||||
Normalize any datetime to UTC timezone-aware datetime.
|
||||
|
||||
Args:
|
||||
dt: Datetime or pandas Timestamp to normalize
|
||||
|
||||
Returns:
|
||||
UTC timezone-aware datetime
|
||||
"""
|
||||
if dt is None:
|
||||
return None
|
||||
|
||||
if isinstance(dt, pd.Timestamp):
|
||||
dt = dt.to_pydatetime()
|
||||
|
||||
if dt.tzinfo is None:
|
||||
return dt.replace(tzinfo=timezone.utc)
|
||||
|
||||
return dt.astimezone(timezone.utc)
|
||||
|
||||
|
||||
def normalize_dataframe_datetime_column(
|
||||
df: pd.DataFrame,
|
||||
column: str,
|
||||
target_format: str = 'naive'
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Normalize a datetime column in a dataframe to consistent format.
|
||||
|
||||
Args:
|
||||
df: DataFrame to process
|
||||
column: Name of datetime column
|
||||
target_format: 'naive' or 'aware' (UTC)
|
||||
|
||||
Returns:
|
||||
DataFrame with normalized datetime column
|
||||
"""
|
||||
if column not in df.columns:
|
||||
logger.warning(f"Column {column} not found in dataframe")
|
||||
return df
|
||||
|
||||
df[column] = pd.to_datetime(df[column])
|
||||
|
||||
if target_format == 'naive':
|
||||
if df[column].dt.tz is not None:
|
||||
df[column] = df[column].dt.tz_localize(None)
|
||||
elif target_format == 'aware':
|
||||
if df[column].dt.tz is None:
|
||||
df[column] = df[column].dt.tz_localize(timezone.utc)
|
||||
else:
|
||||
df[column] = df[column].dt.tz_convert(timezone.utc)
|
||||
else:
|
||||
raise ValueError(f"Invalid target_format: {target_format}. Must be 'naive' or 'aware'")
|
||||
|
||||
return df
|
||||
|
||||
|
||||
def prepare_prophet_datetime(df: pd.DataFrame, datetime_col: str = 'ds') -> pd.DataFrame:
|
||||
"""
|
||||
Prepare datetime column for Prophet (requires timezone-naive datetimes).
|
||||
|
||||
Args:
|
||||
df: DataFrame with datetime column
|
||||
datetime_col: Name of datetime column (default: 'ds')
|
||||
|
||||
Returns:
|
||||
DataFrame with Prophet-compatible datetime column
|
||||
"""
|
||||
df = df.copy()
|
||||
df = normalize_dataframe_datetime_column(df, datetime_col, target_format='naive')
|
||||
return df
|
||||
|
||||
|
||||
def safe_datetime_comparison(dt1: datetime, dt2: datetime) -> int:
|
||||
"""
|
||||
Safely compare two datetimes, handling timezone mismatches.
|
||||
|
||||
Args:
|
||||
dt1: First datetime
|
||||
dt2: Second datetime
|
||||
|
||||
Returns:
|
||||
-1 if dt1 < dt2, 0 if equal, 1 if dt1 > dt2
|
||||
"""
|
||||
dt1_utc = normalize_datetime_to_utc(dt1)
|
||||
dt2_utc = normalize_datetime_to_utc(dt2)
|
||||
|
||||
if dt1_utc < dt2_utc:
|
||||
return -1
|
||||
elif dt1_utc > dt2_utc:
|
||||
return 1
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
def get_current_utc() -> datetime:
|
||||
"""
|
||||
Get current datetime in UTC with timezone awareness.
|
||||
|
||||
Returns:
|
||||
Current UTC datetime
|
||||
"""
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def convert_timestamp_to_datetime(timestamp: Union[int, float, str]) -> datetime:
|
||||
"""
|
||||
Convert various timestamp formats to datetime.
|
||||
|
||||
Args:
|
||||
timestamp: Unix timestamp (seconds or milliseconds) or ISO string
|
||||
|
||||
Returns:
|
||||
UTC timezone-aware datetime
|
||||
"""
|
||||
if isinstance(timestamp, str):
|
||||
dt = pd.to_datetime(timestamp)
|
||||
return normalize_datetime_to_utc(dt)
|
||||
|
||||
if timestamp > 1e10:
|
||||
timestamp = timestamp / 1000
|
||||
|
||||
dt = datetime.fromtimestamp(timestamp, tz=timezone.utc)
|
||||
return dt
|
||||
|
||||
|
||||
def align_dataframe_dates(
|
||||
dfs: list[pd.DataFrame],
|
||||
date_column: str = 'ds',
|
||||
method: str = 'inner'
|
||||
) -> list[pd.DataFrame]:
|
||||
"""
|
||||
Align multiple dataframes to have the same date range.
|
||||
|
||||
Args:
|
||||
dfs: List of DataFrames to align
|
||||
date_column: Name of the date column
|
||||
method: 'inner' (intersection) or 'outer' (union)
|
||||
|
||||
Returns:
|
||||
List of aligned DataFrames
|
||||
"""
|
||||
if not dfs:
|
||||
return []
|
||||
|
||||
if len(dfs) == 1:
|
||||
return dfs
|
||||
|
||||
all_dates = None
|
||||
|
||||
for df in dfs:
|
||||
if date_column not in df.columns:
|
||||
continue
|
||||
|
||||
dates = set(pd.to_datetime(df[date_column]).dt.date)
|
||||
|
||||
if all_dates is None:
|
||||
all_dates = dates
|
||||
else:
|
||||
if method == 'inner':
|
||||
all_dates = all_dates.intersection(dates)
|
||||
elif method == 'outer':
|
||||
all_dates = all_dates.union(dates)
|
||||
|
||||
aligned_dfs = []
|
||||
for df in dfs:
|
||||
if date_column not in df.columns:
|
||||
aligned_dfs.append(df)
|
||||
continue
|
||||
|
||||
df = df.copy()
|
||||
df[date_column] = pd.to_datetime(df[date_column])
|
||||
df['_date_only'] = df[date_column].dt.date
|
||||
df = df[df['_date_only'].isin(all_dates)]
|
||||
df = df.drop('_date_only', axis=1)
|
||||
aligned_dfs.append(df)
|
||||
|
||||
return aligned_dfs
|
||||
|
||||
|
||||
def fill_missing_dates(
|
||||
df: pd.DataFrame,
|
||||
date_column: str = 'ds',
|
||||
freq: str = 'D',
|
||||
fill_value: float = 0.0
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Fill missing dates in a DataFrame with a specified frequency.
|
||||
|
||||
Args:
|
||||
df: DataFrame with date column
|
||||
date_column: Name of the date column
|
||||
freq: Pandas frequency string ('D' for daily, 'H' for hourly, etc.)
|
||||
fill_value: Value to fill for missing dates
|
||||
|
||||
Returns:
|
||||
DataFrame with filled dates
|
||||
"""
|
||||
df = df.copy()
|
||||
df[date_column] = pd.to_datetime(df[date_column])
|
||||
|
||||
df = df.set_index(date_column)
|
||||
|
||||
full_range = pd.date_range(
|
||||
start=df.index.min(),
|
||||
end=df.index.max(),
|
||||
freq=freq
|
||||
)
|
||||
|
||||
df = df.reindex(full_range, fill_value=fill_value)
|
||||
|
||||
df = df.reset_index()
|
||||
df = df.rename(columns={'index': date_column})
|
||||
|
||||
return df
|
||||
Reference in New Issue
Block a user