genformer 0.1.0__py3-none-any.whl

genformer/__init__.py

@@ -0,0 +1,9 @@
+from genformer.models import Enformer, GEnformer
+from genformer.noise import GaussianNoise, UniformNoise
+
+__all__ = [
+    "Enformer",
+    "GEnformer",
+    "GaussianNoise",
+    "UniformNoise",
+]

genformer/data.py

@@ -0,0 +1,193 @@
+import numpy as np
+import pandas as pd
+from typing import Tuple, Optional, List
+from darts import TimeSeries, concatenate
+from gluonts.dataset.repository.datasets import get_dataset
+from gluonts.dataset.multivariate_grouper import MultivariateGrouper
+from darts.dataprocessing.transformers import Scaler
+from pandas.tseries.frequencies import to_offset
+from gluonts.time_feature import norm_freq_str
+
+def gluonts_item_to_darts_mv(item, freq: str) -> TimeSeries:
+    start = item["start"].to_timestamp() if hasattr(item["start"], "to_timestamp") else pd.Timestamp(item["start"])
+    target = np.asarray(item["target"])
+    if target.ndim != 2:
+        raise ValueError(f"Expected multivariate target with ndim=2, got shape {target.shape}")
+
+    values = target.T 
+    times = pd.date_range(start=start, periods=values.shape[0], freq=freq)
+    cols = [f"dim_{i}" for i in range(values.shape[1])]
+
+    return TimeSeries.from_times_and_values(times, values, columns=cols)
+
+def lag_covs_from_scaled_target(ts_sc: TimeSeries, lags=(1,24,168)) -> TimeSeries:
+    shifted = []
+    for L in lags:
+        s = ts_sc.shift(L).with_columns_renamed(
+            ts_sc.components, [f"{c}_lag{L}" for c in ts_sc.components]
+        )
+        shifted.append(s)
+
+    common = shifted[0]
+    for s in shifted[1:]:
+        common = common.slice_intersect(s)
+    shifted = [s.slice_intersect(common) for s in shifted]
+    return concatenate(shifted, axis=1)
+
+def fourier_from_index_min(idx) -> TimeSeries:
+    minute = idx.minute.to_numpy()
+    hour   = idx.hour.to_numpy()
+    dow    = idx.dayofweek.to_numpy()
+
+    X = np.vstack([
+        np.sin(2 * np.pi * minute / 60.0),
+        np.cos(2 * np.pi * minute / 60.0),
+        np.sin(2 * np.pi * hour / 24.0),
+        np.cos(2 * np.pi * hour / 24.0),
+        np.sin(2 * np.pi * dow / 7.0),
+        np.cos(2 * np.pi * dow / 7.0),
+    ]).T
+
+    return TimeSeries.from_times_and_values(
+        idx,
+        X,
+        columns=["min_sin", "min_cos", "h_sin", "h_cos", "dow_sin", "dow_cos"]
+    )
+
+def fourier_from_index_day(idx) -> TimeSeries:
+    dow  = idx.dayofweek.to_numpy()
+    X = np.vstack([
+        np.sin(2*np.pi*dow/7.0),
+        np.cos(2*np.pi*dow/7.0),
+    ]).T
+    return TimeSeries.from_times_and_values(idx, X, columns=["dow_sin","dow_cos"])
+
+def fourier_from_index(idx) -> TimeSeries:
+    hour = idx.hour.to_numpy()
+    dow  = idx.dayofweek.to_numpy()
+    X = np.vstack([
+        np.sin(2*np.pi*hour/24.0),
+        np.cos(2*np.pi*hour/24.0),
+        np.sin(2*np.pi*dow/7.0),
+        np.cos(2*np.pi*dow/7.0),
+    ]).T
+    return TimeSeries.from_times_and_values(idx, X, columns=["h_sin","h_cos","dow_sin","dow_cos"])
+
+def dim_indicator_norm(idx, D: int) -> TimeSeries:
+    v = (np.arange(D, dtype=np.float32) / (D-1)).astype(np.float32)  # 0..1
+    X = np.tile(v, (len(idx), 1))
+    cols = [f"dim_id_{i}" for i in range(D)]
+    return TimeSeries.from_times_and_values(idx, X, columns=cols)
+
+def build_past_covs_552(ts_sc: TimeSeries, lags=(1,24,168), fourier_func=fourier_from_index_min) -> TimeSeries:
+    lag_covs = lag_covs_from_scaled_target(ts_sc, lags)
+    idx = lag_covs.time_index
+    time_covs = fourier_func(idx)
+    dim_covs  = dim_indicator_norm(idx, ts_sc.width)
+    return concatenate([lag_covs, dim_covs, time_covs], axis=1)
+
+def ts_upto(ts, end_time):
+    if hasattr(ts, "slice_end"):
+        return ts.slice_end(end_time)
+    if hasattr(ts, "drop_after"):
+        return ts.drop_after(end_time)
+    if hasattr(ts, "split_after"):
+        return ts.split_after(end_time)[0]
+    return ts.slice(ts.start_time(), end_time)
+
+def gluon_to_wide_df(dataset):
+    series_list = []
+    
+    for i, entry in enumerate(dataset):
+        idx = pd.date_range(
+            start=entry["start"].to_timestamp(), 
+            periods=len(entry["target"]), 
+            freq=entry["start"].freqstr
+        )
+        series = pd.Series(entry["target"], index=idx, name=f"node_{i}")
+        series_list.append(series)
+    
+    return pd.concat(series_list, axis=1)
+
+def get_7_test_windows(dataset, num_nodes=137):
+    all_series = []
+    
+    for entry in dataset:
+        idx = pd.date_range(
+            start=entry["start"].to_timestamp(), 
+            periods=len(entry["target"]), 
+            freq=entry["start"].freqstr
+        )
+        all_series.append(pd.Series(entry["target"], index=idx))
+    
+    num_windows = len(all_series) // num_nodes
+    windows = []
+    
+    for w in range(num_windows):
+        start_idx = w * num_nodes
+        end_idx = (w + 1) * num_nodes
+        window_df = pd.concat(all_series[start_idx:end_idx], axis=1)
+        window_df.columns = [f"node_{i}" for i in range(num_nodes)]
+        windows.append(window_df)
+        
+    return windows, all_series
+
+def load_and_prepare_data(dataset_name: str, 
+                          lags: Tuple, 
+                          fourier_func_dict: dict):
+    
+    ds = get_dataset(dataset_name, regenerate=False)
+    freq = ds.metadata.freq
+    offset = to_offset(freq)
+    granularity = norm_freq_str(offset.name) 
+    
+    train_list = list(ds.train)
+    test_list  = list(ds.test)
+    num_test_dates = len(test_list) // len(train_list)  
+    
+    target_dim = int(ds.metadata.feat_static_cat[0].cardinality) 
+    train_grouper = MultivariateGrouper(max_target_dim=target_dim)
+    test_grouper  = MultivariateGrouper(num_test_dates=num_test_dates, max_target_dim=target_dim)
+    
+    train_mv_items = list(train_grouper(train_list))   
+    
+    if dataset_name == 'kdd_cup_2018_without_missing':  
+        for i in range(len(test_list)):
+            if len(test_list[i]['target']) == 10898:
+                test_list[i]['target'] = np.concatenate(
+                    (test_list[i]['target'], np.zeros(8)), axis=0)
+        dataset_test = test_grouper(test_list)
+    else:
+        dataset_test = test_grouper(test_list)
+
+    test_mv_items  = list(dataset_test)     
+    
+    train_ts = gluonts_item_to_darts_mv(train_mv_items[0], freq)
+    test_ts_list = [gluonts_item_to_darts_mv(it, freq) for it in test_mv_items]  
+
+    y_scaler = Scaler()                     
+    train_y_sc = y_scaler.fit_transform(train_ts)
+
+    test_windows, all_series = get_7_test_windows(test_list, num_nodes=target_dim)
+    
+    fourier_func = fourier_func_dict.get(dataset_name, fourier_from_index)
+    train_pc = build_past_covs_552(train_y_sc, lags=lags, fourier_func=fourier_func)
+    
+    train_y_sc = train_y_sc.slice_intersect(train_pc)
+    train_pc   = train_pc.slice_intersect(train_y_sc)
+
+    return {
+        'scaler': y_scaler, 'train_ts': train_ts, 'train_y_sc': train_y_sc, 
+        'test_ts_list': test_ts_list, 'test_windows': test_windows, 
+        'all_series': all_series, 'train_pc': train_pc, 'freq': granularity,
+        'fourier_func': fourier_func
+    }
+
+FOURIER_DICT = {
+    'solar_nips': fourier_from_index, 
+    'wiki2000_nips': fourier_from_index_day, 
+    'electricity_nips': fourier_from_index, 
+    'taxi_30min': fourier_from_index_min,
+    'kdd_cup_2018_without_missing': fourier_from_index, 
+    'traffic_nips': fourier_from_index
+}

genformer/metrics.py

@@ -0,0 +1,144 @@
+import numpy as np
+import pandas as pd
+import torch
+from gluonts.model.forecast import SampleForecast
+from gluonts.evaluation import MultivariateEvaluator
+from darts import TimeSeries
+
+from genformer.data import ts_upto, build_past_covs_552
+
+def to_gluonts_multivariate_inputs(
+    preds: np.ndarray,
+    targets: np.ndarray,
+    start_dates,
+    freq: str = "H",
+    item_ids=None,
+    past_targets: np.ndarray | None = None,
+    dtype=np.float32,
+):
+    preds = np.asarray(preds, dtype=dtype)
+    targets = np.asarray(targets, dtype=dtype)
+
+    B, N, T, D = preds.shape
+
+    if isinstance(start_dates, (str, pd.Timestamp, pd.Period)) or not hasattr(start_dates, "__len__"):
+        start_dates = [start_dates] * B
+
+    if item_ids is None:
+        item_ids = [f"item_{i}" for i in range(B)]
+
+    targets_list = []
+    forecasts_list = []
+
+    columns = list(range(D))  
+
+    for i in range(B):
+        start_period = pd.Period(start_dates[i], freq=freq)
+        forecast = SampleForecast(
+            samples=preds[i],          
+            start_date=start_period,
+            item_id=item_ids[i],
+        )
+        forecasts_list.append(forecast)
+
+        if past_targets is None:
+            target_index = pd.period_range(start=start_period, periods=T, freq=freq)
+            target_values = targets[i]  
+        else:
+            H = past_targets.shape[1]
+            target_index = pd.period_range(start=start_period - H, periods=H + T, freq=freq)
+            target_values = np.concatenate([past_targets[i], targets[i]], axis=0)  
+
+        target_df = pd.DataFrame(target_values, index=target_index, columns=columns)
+        targets_list.append(target_df)
+
+    return targets_list, forecasts_list
+
+def crps(preds, targets, quantiles=(np.arange(20) / 20.0)[1:]):
+    x = np.quantile(preds, quantiles, axis=1, method="nearest")  
+    quantiles = np.expand_dims(quantiles, axis=list(range(1, len(preds.shape))))  
+    loss = 2 * np.sum(np.abs((x - targets) * ((targets <= x) - quantiles)), axis=2)  
+    return loss.mean() / np.abs(targets).sum(axis=1).mean()
+
+def crps_sum(preds, targets, quantiles=(np.arange(20) / 20.0)[1:], frequency='D'):
+    preds_sum = preds.sum(axis=-1)      
+    targets_sum = targets.sum(axis=-1)  
+    return crps(preds_sum, targets_sum, quantiles=quantiles)
+
+def get_metric_and_prediction(model, test_windows, y_scaler, pred_len=24, lags=(1, 24, 168), num_samples=100, seed=42, std=None, fourier_func=None, is_clip=False, frequency='D'):
+    all_forecasts = []
+    all_targets = []
+    to_save_forecast = []
+    to_save_targets = []
+
+    for i, window_df in enumerate(test_windows):
+        full_ts = TimeSeries.from_dataframe(window_df).astype(np.float32)
+        full_sc = y_scaler.transform(full_ts).astype(np.float32)
+        full_pc = build_past_covs_552(full_sc, lags=lags, fourier_func=fourier_func).astype(np.float32)
+
+        full_sc = full_sc.slice_intersect(full_pc)
+        full_ts = full_ts.slice_intersect(full_sc)
+        full_pc = full_pc.slice_intersect(full_sc)
+
+        past_true_sc = full_sc[:-pred_len]
+        gt_future    = full_ts[-pred_len:]
+        
+        forecast_start = gt_future.start_time()
+        pc_past = ts_upto(full_pc, forecast_start)
+
+        model.model.encoder[0].reset_seed(seed)
+        if std is not None:
+            model.model.encoder[0].reset_std(std)
+        
+        fc_sc = model.predict(
+            n=pred_len,
+            series=past_true_sc,
+            past_covariates=pc_past,
+            num_samples=num_samples,
+            verbose=False,
+            random_state=1456445
+        )
+
+        fc = y_scaler.inverse_transform(fc_sc)
+        if is_clip:
+            fc = fc.with_values(np.clip(fc.all_values(), a_min=0, a_max=None))
+
+        to_save_forecast.append(fc)
+        to_save_targets.append(gt_future)
+        
+        all_forecasts.append(fc.all_values(copy=False))
+        all_targets.append(gt_future.all_values(copy=False))
+
+    stacked_forecasts = np.stack(all_forecasts, axis=0)      
+    preds_reshaped = np.transpose(stacked_forecasts, (0, 3, 1, 2))  
+
+    stacked_targets = np.stack(all_targets, axis=0)          
+    targets_reshaped = np.squeeze(stacked_targets, axis=-1)  
+
+    target_list , forecast_list = to_gluonts_multivariate_inputs(preds_reshaped, targets_reshaped, pd.Timestamp(gt_future.start_time()), freq=frequency)
+
+    evaluator = MultivariateEvaluator(quantiles=(np.arange(20) / 20.0)[1:], target_agg_funcs={'sum': np.sum})
+    agg_metric, item_metrics = evaluator(target_list, forecast_list)
+
+    print(f"======= Evaluation metrics for models =======")
+    print("CRPS:", agg_metric["mean_wQuantileLoss"])
+    print("ND:", agg_metric["ND"])
+    print("NRMSE:", agg_metric["NRMSE"])
+    print("")
+    print("CRPS-Sum:", agg_metric["m_sum_mean_wQuantileLoss"])
+    print("ND-Sum:", agg_metric["m_sum_ND"])
+    print("NRMSE-Sum:", agg_metric["m_sum_NRMSE"])
+    crps_ours = crps_sum(preds_reshaped, targets_reshaped)
+    print("CRPS Ours:", crps_ours)
+    
+    return agg_metric["m_sum_mean_wQuantileLoss"], agg_metric["m_sum_ND"], agg_metric["m_sum_NRMSE"], crps_ours, to_save_forecast, to_save_targets
+
+def energy_score_loss(samples: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    M = samples.size(0)
+    dist_to_target = torch.linalg.norm(samples - target.unsqueeze(0), dim=-1).mean(0)
+    s = samples.reshape(M, -1, samples.size(-1))
+    diff = s.unsqueeze(1) - s.unsqueeze(0)
+    pairwise_dist = torch.linalg.norm(diff, dim=-1).mean(dim=(0, 1))
+    dist_samples = pairwise_dist.view(target.size(0), target.size(1))
+    loss = dist_to_target - 0.5 * dist_samples
+    return loss.mean()