Creating a coupled LSTM

Investor101: <p>Hi all,</p>
<p>I am attempting to reproduce the Neural Assimilation model architecture described in this paper: <a href="https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7304751/" class="inline-onebox" rel="noopener nofollow ugc">Neural Assimilation - PMC</a>. They are essentially trying to design an architecture that would be much faster than using a Kalman Filter by bypassing the correction step. Despite the authors providing some code in Keras, I have had some issues trying to implement this with PyTorch as I simply cannot see how they have built the model. (link to the code is here: <a href="https://drive.google.com/drive/folders/1TBg9ZfPyXiegoXbn2sZURPjPml8TWQIO" class="inline-onebox" rel="noopener nofollow ugc">FINAL CODE - Google Drive</a>)</p>
<p>From my understanding, we need two LSTMs with one using the model forecasted data and the other using the observational data. The tricky part is that each LSTM requires as input the hidden state and cell state from the other LSTM which I am having some difficulty trying to implement. Anyone have any ideas on how I could do this? So far I have only used simple LSTMs for many-to-many prediction.</p>
<p>This is a diagram of the architecture:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/9/099d359783e1b4514ab7712a8b1d935e447559fd.png" data-download-href="https://discuss.pytorch.org/uploads/default/099d359783e1b4514ab7712a8b1d935e447559fd" title="Screenshot at May 10 12-11-27"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/9/099d359783e1b4514ab7712a8b1d935e447559fd_2_690x433.png" alt="Screenshot at May 10 12-11-27" data-base62-sha1="1n37eAvairP2sx4l4jdVJPGQgCF" width="690" height="433" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/9/099d359783e1b4514ab7712a8b1d935e447559fd_2_690x433.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/9/099d359783e1b4514ab7712a8b1d935e447559fd_2_1035x649.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/0/9/099d359783e1b4514ab7712a8b1d935e447559fd.png 2x" data-dominant-color="F508F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot at May 10 12-11-27</span><span class="informations">1176×738 95.9 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>and what I am working on thus far as its implementation:</p>
<pre><code class="lang-auto">class Neural_Assimilation(nn.Module):
  def __init__(self, n_inputs, n_hidden, num_layers, n_outputs):
    super(Neural_Assimilation, self).__init__()
    self.D = n_inputs
    self.M = n_hidden
    self.K = n_outputs
    self.L = num_layers

    self.rnn_f = nn.LSTM(
        input_size=self.D,
        hidden_size=self.M,
        num_layers=self.L,
        batch_first=True)
    
    self.rnn_o = nn.LSTM(
        input_size=self.D,
        hidden_size=self.M,
        num_layers=self.L,
        batch_first=True)
    self.fc = nn.Linear(self.M, self.K) # For BiDir LSTM need to multiply self.M x 2
  
  def forward(self, X_f, X_o):
    # initial hidden states
    h0_f = torch.zeros(self.L, X_f.size(0), self.M).to(device) # For BiDir LSTM need to multiply self.L x 2
    c0_f = torch.zeros(self.L, X_f.size(0), self.M).to(device) # For BiDir LSTM need to multiply self.L x 2

    # get RNN unit output
    out_f, (h_t_f, c_t_f) = self.rnn_f(X_f, (h0_f, c0_f))
    out_o, (h_t_o, c_t_o) = self.rnn_o(X_o, (h_t_f, c_t_f))
    out_f, (h_t_o, c_t_o) = self.rnn_f(X_f, (h_t_o, c_t_o))
    
    # we want h(t) and c(t) at the final time step so we can feed this into the other RNN
    out_f = self.fc(out_f[:, -1, :])

    hidden_state = h_t[:, -1, :]
    hidden_out = self.fc(hidden_state)

    cell_state = c_t[:, -1, :]
    cell_out = self.fc(cell_state)

    return out, hidden_state, cell_out 
</code></pre>
<p>Would really appreciate your expertise!</p>