dataloading

seq2seq_time/data/dataset.py

@@ -0,0 +1,100 @@
+import pandas as pd
+import torch.utils.data
+import numpy as np
+
+def assert_normalized(df):
+    stats = df.describe().T
+    np.testing.assert_allclose(stats['mean'].values, 0, atol=0.1), 'means should be normalized to ~0'
+    np.testing.assert_allclose(stats['std'].values, 1, atol=0.1), 'standard deviations should be normalized to ~0'
+
+def assert_no_objects(df):
+    for name, dtype in df.dtypes.iteritems():
+        assert dtype.name!='object', f'all objects should be pd.categories. {name} is not'
+
+
+class Seq2SeqDataSet(torch.utils.data.Dataset):
+    """
+    Takes in dataframe and returns sequences through time.
+    
+    Returns x_past, y_past, x_future, etc.
+    """
+    
+    def __init__(self, df: pd.DataFrame, window_past=40, window_future=10, columns_target=['energy(kWh/hh)'], columns_blank=[],):
+        """
+        Args:
+        - df: DataFrame with time index, already scaled
+        - columns_blank: The columns we will blank, in the future
+        """
+        super().__init__()
+        # TODO auto categorical columns
+        # TODO specify blank future columns
+        assert isinstance(df.index, pd.DatetimeIndex), 'should have a datetime index'
+        assert df.index.freq is not None, 'should have freq'
+        # assert_normalized(df)
+        assert_no_objects(df)
+
+        # Use numpy instead of pandas, for speed
+        self.x = df.drop(columns=columns_target).copy().values
+        self.y = df[columns_target].copy().values
+        self.t = df.index.copy()
+        self.columns = list(df.columns)
+        self.icol_blank = [df.drop(columns=columns_target).columns.tolist().index(n) for n in columns_blank]
+
+        self.window_past = window_past
+        self.window_future = window_future
+        self.columns_target = columns_target
+
+    def get_components(self, i):
+        """Get past and future rows."""
+        x = self.x[i : i + (self.window_past + self.window_future)].copy()
+        y = self.y[i:i + (self.window_past + self.window_future)].copy()
+        t = self.t[i:i + (self.window_past + self.window_future)].copy()
+        t = t.astype(int) * 1e-9 / 60 / 60 / 24  # days
+        t = t.values
+        now = t[self.window_past]
+        
+        # Add a features: relative hours since present time, is future
+        tstp = (t - now)[:, None]
+        is_past = tstp < 0
+        x = np.concatenate([x, tstp, is_past], -1)
+        
+        # Split into future and past
+        x_past = x[:self.window_past]
+        y_past = y[:self.window_past]
+        x_future = x[self.window_past:]
+        y_future = y[self.window_past:]
+
+        # Stop it cheating by using future weather measurements
+        x_future[:, self.icol_blank] = 0
+        return x_past, y_past, x_future, y_future
+
+
+    def __getitem__(self, i):
+        """This is how python implements square brackets"""
+        if i<0:
+            # Handle negative integers
+            i = len(self)+i
+        data = self.get_components(i)
+        # From dataframe to torch
+        return [d.astype(np.float32) for d in data]
+    
+    
+    def get_rows(self, i):
+        """
+        Output pandas dataframes for display purposes.
+        """
+        x_cols = list(self.columns)[1:] + ['tsp_days', 'is_past']
+        x_past, y_past, x_future, y_future = self.get_components(i)
+        t_past = self.t[i:i+self.window_past]
+        t_future = self.t[i+self.window_past:i+self.window_past + self.window_future]
+        x_past = pd.DataFrame(x_past, columns=x_cols, index=t_past)
+        x_future = pd.DataFrame(x_future, columns=x_cols, index=t_future)
+        y_past = pd.DataFrame(y_past, columns=self.columns_target, index=t_past)
+        y_future = pd.DataFrame(y_future, columns=self.columns_target, index=t_future)
+        return x_past, y_past, x_future, y_future
+        
+    def __len__(self):
+        return len(self.x) - (self.window_past + self.window_future)
+    
+    def __repr__(self):
+        return f'<{type(self).__name__}(shape={self.x.shape}, times={self.t[0]} to {self.t[1]} at {self.t.freq.freqstr})>'

seq2seq_time/predict.py

@@ -0,0 +1,72 @@
+import xarray as xr 
+import torch
+from tqdm.auto import tqdm
+import pandas as pd
+
+from .util import to_numpy
+
+def predict(model, ds_test, batch_size, device='cpu', scaler=None):
+    """
+    Gather all predictions into xarray.
+    
+    When we generate prediction in a sequence to sequence model we start at a time then predict
+    N steps into the future. So we have 2 dimensions: source time, target time.
+
+    But we also care about how far we were predicting into the future, so we have 3 dimensions: source time, target time, time ahead.
+
+    It's hard to use pandas for data with virtual dimensions so we will use xarray. Xarray has an interface similar to pandas but also allows coordinates which are virtual dimensions.
+    """
+    load_test = torch.utils.data.dataloader.DataLoader(ds_test, batch_size=batch_size)
+    freq = ds_test.t.freq
+    xrs = []
+    for i, batch in enumerate(tqdm(load_test, desc='predict')):
+        model.eval()
+        with torch.no_grad():
+            x_past, y_past, x_future, y_future = [d.to(device) for d in batch]
+            y_dist = model(x_past, y_past, x_future, y_future)
+            nll = -y_dist.log_prob(y_future)
+
+            # Convert to numpy
+            mean = to_numpy(y_dist.loc.squeeze(-1))
+            std = to_numpy(y_dist.scale.squeeze(-1))
+            nll = to_numpy(nll.squeeze(-1))
+            y_future = to_numpy(y_future.squeeze(-1))
+            y_past = to_numpy(y_past.squeeze(-1))    
+
+        # Make an xarray.Dataset for the data
+        bs = y_future.shape[0]
+        t_source = ds_test.t[i:i+bs].values
+        t_ahead = pd.timedelta_range(0, periods=ds_test.window_future, freq=freq).values
+        t_behind = pd.timedelta_range(end=-pd.Timedelta(freq), periods=ds_test.window_past, freq=freq)
+        xr_out = xr.Dataset(
+            {
+                # Format> name: ([dimensions,...], array),
+                "y_past": (["t_source", "t_behind",], y_past),
+                "nll": (["t_source", "t_ahead",], nll),
+                "y_pred": (["t_source", "t_ahead",], mean),
+                "y_pred_std": (["t_source", "t_ahead",], std),
+                "y_true": (["t_source", "t_ahead",], y_future),
+            },
+            coords={"t_source": t_source, "t_ahead": t_ahead, "t_behind": t_behind},
+        )
+        xrs.append(xr_out)
+
+    # Join all batches
+    ds_preds = xr.concat(xrs, dim="t_source")
+    
+    # undo scaling on y
+    if scaler:
+        ds_preds['y_pred_std'].values = ds_preds.y_pred_std * scaler.scale_
+        ds_preds['y_past'].values =  scaler.inverse_transform(ds_preds.y_past)
+        ds_preds['y_pred'].values =  scaler.inverse_transform(ds_preds.y_pred)
+        ds_preds['y_true'].values =  scaler.inverse_transform(ds_preds.y_true)
+
+    # Add some derived coordinates, they will be the ones not in bold
+    # The target time, is a function of the source time, and how far we predict ahead
+    ds_preds = ds_preds.assign_coords(t_target=ds_preds.t_source+ds_preds.t_ahead)
+
+    ds_preds = ds_preds.assign_coords(t_past=ds_preds.t_source+ds_preds.t_behind)
+
+    # Some plots don't like timedeltas, so lets make a coordinate for time ahead in hours
+    ds_preds = ds_preds.assign_coords(t_ahead_hours=(ds_preds.t_ahead*1.0e-9/60/60).astype(float))
+    return ds_preds

seq2seq_time/util.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+import torch
+
+project_dir = Path(__file__).parent.parent
+
+def to_numpy(x):
+    """Helper function to avoid repeating code"""
+    if isinstance(x, torch.Tensor):
+        x = x.cpu().detach().numpy()
+    return x