working

seq2seq_time/models/lstm_seq2seq.py

@@ -36,4 +36,4 @@ class LSTMSeq2Seq(nn.Module):
         log_sigma = self.std(outputs)
         sigma = self._min_std + (1 - self._min_std) * F.softplus(log_sigma)
         y_dist = torch.distributions.Normal(mean, sigma)
-        return y_dist
+        return y_dist, {}

seq2seq_time/models/neural_process.py

@@ -7,6 +7,7 @@ import math
 
 
 class LSTMBlock(nn.Module):
+    """Wrapper to return only lstm output."""
     def __init__(
         self,
         in_channels,
@@ -437,14 +438,21 @@ class RANP(nn.Module):
         if self._use_rnn:
             # see https://arxiv.org/abs/1910.09323 where x is substituted with h = RNN(x)
             # x need to be provided as [B, T, H]
-            future_x, _ = self._lstm(future_x)
-            past_x, _ = self._lstm(past_x)
+            S = past_x.shape[1]
+            x = torch.cat([past_x, future_x], 1)
+            x, _ = self._lstm(x)
+            past_x = x[:, :S]
+            future_x = x[:, S:]
+            # future_x, _ = self._lstm(future_x)
+            # past_x, _ = self._lstm(past_x)
 
         dist_prior, log_var_prior = self._latent_encoder(past_x, past_y)
 
-        if future_y is not None:
+        if (future_y is not None):
             dist_post, log_var_post = self._latent_encoder(future_x, future_y)
-            z = dist_post.loc
+
+        if self.training:
+            z = dist_prior.rsample()
         else:
             z = dist_prior.loc
 
@@ -471,3 +479,4 @@ class RANP(nn.Module):
             ].mean()
             loss = (kl_loss - log_p).mean()
         return dist, {'loss':loss}
+

seq2seq_time/predict.py

@@ -35,9 +35,10 @@ def predict(model, ds_test, batch_size, device='cpu', scaler=None):
 
         # Make an xarray.Dataset for the data
         bs = y_future.shape[0]
-        t_source = ds_test.df.index[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)
+        wp = ds_test.window_past
+        t_source = ds_test.df.index[wp + i*bs -1:wp+ i*bs+bs -1].values
+        t_ahead = pd.timedelta_range(1, periods=ds_test.window_future, freq=freq).values
+        t_behind = pd.timedelta_range(end=0, periods=ds_test.window_past, freq=freq)
         xr_out = xr.Dataset(
             {
                 # Format> name: ([dimensions,...], array),
@@ -77,5 +78,5 @@ def predict_multi(model, datasets, batch_size, device='cpu', scaler=None):
                         d,
                         batch_size,
                         device=device,
-                        scaler=output_scaler) for d in tqdm(datasets)]
+                        scaler=scaler) for d in tqdm(datasets, desc='predict_multi')]
     return xr.concat(ds_preds, dim='block')