# -*- coding: utf-8 -*-
"""Untitled0.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1yRWSvCoFVhPuWK6B71vLJzYhcG-sYSmH
"""

class Net(nn.Module): # our model
  def __init__(self, ncomp, NN1, NN2, bidi=True):
    super(Net, self).__init__()

    # play with some of the options in the RNN!
    self.rnn1 = nn.RNN(NN1, ncomp, num_layers = 1, dropout = 0, # MO
                      bidirectional = bidi, nonlinearity = 'tanh')
    self.rnn2 = nn.RNN(NN2,ncomp,num_layers = 1, dropout = 0, bidirectional = bidi, nonlinearity = 'tanh') #TH

    self.fclatent = nn.Linear(ncomp,ncomp)
    self.fc = nn.Linear(ncomp,NN1)

  def forward(self, x0,x1):

    y2 = self.rnn2(x1)[0] # ncomp TH
    y =  self.rnn1(x0)[0] # rnn1 weight is frozen MOs
    y += self.fclatent(y2)

    if self.rnn.bidirectional:
      # if the rnn is bidirectional, it concatenates the activations from the forward and backward pass
      # we want to add them instead, so as to enforce the latents to match between the forward and backward pass
      q = (y[:, :, :ncomp] + y[:, :, ncomp:])/2
    else:
      q = y

    # the softplus function is just like a relu but it's smoothed out so we can't predict 0
    # if we predict 0 and there was a spike, that's an instant Inf in the Poisson log-likelihood which leads to failure
    z = F.softplus(self.fc(q), 10)

    return z, q

mynetwork = Net(10,40,30)
print(mynetwork)
# -*- coding: utf-8 -*-
"""Untitled0.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1yRWSvCoFVhPuWK6B71vLJzYhcG-sYSmH
"""

class Net(nn.Module): # our model
  def __init__(self, ncomp, NN1, NN2, bidi=True):
    super(Net, self).__init__()

    # play with some of the options in the RNN!
    self.rnn1 = nn.RNN(NN1, ncomp, num_layers = 1, dropout = 0, # MO
                      bidirectional = bidi, nonlinearity = 'tanh')
    self.rnn2 = nn.RNN(NN2,ncomp,num_layers = 1, dropout = 0, bidirectional = bidi, nonlinearity = 'tanh') #TH

    self.fclatent = nn.Linear(ncomp,ncomp)
    self.fc = nn.Linear(ncomp,NN1)

  def forward(self, x0,x1):

    y2 = self.rnn2(x1)[0] # ncomp TH
    y =  self.rnn1(x0)[0] # rnn1 weight is frozen MOs
    y += self.fclatent(y2)

    if self.rnn.bidirectional:
      # if the rnn is bidirectional, it concatenates the activations from the forward and backward pass
      # we want to add them instead, so as to enforce the latents to match between the forward and backward pass
      q = (y[:, :, :ncomp] + y[:, :, ncomp:])/2
    else:
      q = y

    # the softplus function is just like a relu but it's smoothed out so we can't predict 0
    # if we predict 0 and there was a spike, that's an instant Inf in the Poisson log-likelihood which leads to failure
    z = F.softplus(self.fc(q), 10)

    return z, q

mynetwork = Net(10,40,30)
print(mynetwork)
# -*- coding: utf-8 -*- Untitled0.ipynb  Automatically generated by Colaboratory.  Original file is located at     https://colab.research.google.com/drive/1yRWSvCoFVhPuWK6B71vLJzYhcG-sYSmH # our model # play with some of the options in the RNN! # MO #TH # ncomp TH # rnn1 weight is frozen MOs # if the rnn is bidirectional, it concatenates the activations from the forward and backward pass # we want to add them instead, so as to enforce the latents to match between the forward and backward pass # the softplus function is just like a relu but it's smoothed out so we can't predict 0 # if we predict 0 and there was a spike, that's an instant Inf in the Poisson log-likelihood which leads to failure