Combining Time Series Data and Summary Statistics for RNN/CNN Embedding Networks

[ali-akhavan89]

My model generates data in the shape of (batch_size, n_time_steps, n_time_series). I’ve also defined a function to calculate hand-crafted summary statistics, resulting in data of shape (batch_size, n_manual_features). How should I effectively combine these datasets for use with RNN/CNN networks or other embedding nets? Would you please provide a minimum code showing how to do this? Thank you very much!

[janfb]

Hi @ali-akhavan89 thanks for starting this discussion!

The embedding net is trained end-to-end with the density estimator. Internally, the output of the embedding net, i.e., the embedded times series of image features, are fed into the density estimator network at some point. I assume that this is the point where you want to add your hand-crafted features as input to the density estimator as well. However, it is difficult to access this point from the outside.

Thus, my suggestion would be that you provide a custom embedding net that can take both, the raw time series and your hand crafted summary stats as input. The net then just passes on the summary stats to the output layer, and embeds the time series using an appropriate embedding net and outputs the concatenation of embedded time series features and summary stats.

Note that you probably need to find a way to combine the shapes (batch_size, n_time_steps, n_time_series) with (batch_size, *summary_stats_shape), e.g., by expanding and repeating, in order to be able to feed them as one tensor into the embedding net and into .append_simulations

From the perspective of sbi, you are then using one big embedding net that takes as input (batch_size, *combined_shape_of_time_series_and_stats) and returns some latent embedding shape, which is passed on to the density estimator.

I hope this helps!

Best,
Jan

[ali-akhavan89]

Thank you, @janfb . I appreciate your detailed response. I was under the impression that I could avoid passing the hand-crafted summaries to the embedding net. Your explanation makes sense. I’ll try to isolate summaries inside the embedding network.

I would also appreciate it if you could provide a minimal code example that demonstrates this approach.

Thank you so much,
Ali

[janfb]

Here is a minimal example showing what I sketched above:

import torch
import torch.nn as nn
import torch.nn.functional as F

class CustomEmbeddingNet(nn.Module):
    def __init__(self, n_time_series: int, n_manual_features: int, hidden_dim: int, output_dim: int):
        super(CustomEmbeddingNet, self).__init__()
        # RNN for embedding time series data
        self.rnn = nn.LSTM(input_size=n_time_series, hidden_size=hidden_dim, batch_first=True)
        # Fully connected layer for combining embedded time series and manual features
        self.fc = nn.Linear(hidden_dim + n_manual_features, output_dim)

    def forward(self, time_series_data: torch.Tensor, manual_features: torch.Tensor) -> torch.Tensor:
        # time_series_data shape: (batch_size, n_time_steps, n_time_series)
        # manual_features shape: (batch_size, n_manual_features)
        
        # Process time series data with RNN
        _, (hn, _) = self.rnn(time_series_data)
        embedded_time_series = hn[-1]  # Take the last hidden state
        # embedded_time_series shape: (batch_size, hidden_dim)
        
        # Concatenate embedded time series features with manual features
        combined_features = torch.cat((embedded_time_series, manual_features), dim=1)
        # combined_features shape: (batch_size, hidden_dim + n_manual_features)
        
        # Pass combined features through fully connected layer
        output = F.relu(self.fc(combined_features))
        return output

# Example usage
batch_size = 32
n_time_steps = 100
n_time_series = 10
n_manual_features = 5
hidden_dim = 64
output_dim = 128

# Instantiate the custom embedding net
embedding_net = CustomEmbeddingNet(n_time_series, n_manual_features, hidden_dim, output_dim)

# Generate some example data
time_series_data = torch.randn(batch_size, n_time_steps, n_time_series)
manual_features = torch.randn(batch_size, n_manual_features)

# Pass the data through the embedding net
output = embedding_net(time_series_data, manual_features)
print(output.shape)  # Should be (batch_size, output_dim)

[janfb]

This example is a bit independent from sbi. To use this setup in sbi, you would need to find a way to concatenate time_series_data and manual_features into one Tensor, to be able to pass it as one x into append_simulations.

forward would then accept one argument and extract the time series and the manual features from the Tensor.