text: import torch
  import torch.nn as nn
  import torch.utils.data as data
  
  import pytorch_lightning as pl
  from pytorch_lightning.loggers import TensorBoardLogger
  
  
  from src.maths.dennet import DENNet
  from src.mechanics.hamiltonian import HNNMassSpring, HNNMassSpringSeparable
  from src.time_integrator import TimeIntegrator
  from src.utils.activation import Siren
  
  import matplotlib.pyplot as plt
  
  
  class Learner(pl.LightningModule):
      def __init__(self, model: nn.Module):
          super().__init__()
          self.model = model
          self.c = 0
  
      def forward(self, x):
          return self.model.de_function(0, x)
  
      def backward(self, use_amp, loss, optimizer, optimizer_idx):
          loss.backward(retain_graph=True)
          return
  
      @staticmethod
      def loss(y, y_hat):
          return ((y - y_hat) ** 2).sum()
  
      @staticmethod
      def calculate_f(x):
          q_dot = x[:, 1].unsqueeze(1)
          p_dot = -x[:, 0].unsqueeze(1)
          f = torch.cat([q_dot, p_dot], 1).to(x)
          return f
  
      def training_step(self, batch, batch_idx):
          x = batch[0]
          y_hat = self.model.de_function(0, x)
          # will need to calculate a y here from the equation residual
          y = self.calculate_f(x)
          loss = self.loss(y_hat, y)
          logs = {'train_loss': loss}
          return {'loss': loss, 'log': logs}
  
      def configure_optimizers(self):
          return torch.optim.Adam(self.model.parameters(), lr=0.01)
  
      @staticmethod
      def train_dataloader():
          return trainloader
  
  
  class SirenNet(nn.Module):
      def __init__(self, dim_in, dim_hidden, dim_out, num_layers, w0=30., w0_initial=30.):
          super().__init__()
          layers = list()
          for ind in range(num_layers):
              is_first = ind == 0
              layer_w0 = w0_initial if is_first else w0
              layer_dim_in = dim_in if is_first else dim_hidden
  
              layers.append(Siren(
                  dim_in=layer_dim_in,
                  dim_out=dim_hidden,
                  w0=layer_w0,
                  is_first=is_first
              ))
  
          self.net = nn.Sequential(*layers)
          self.last_layer = nn.Linear(dim_hidden, dim_out)
  
      def forward(self, x):
          x = self.net(x)
          return self.last_layer(x)
  
  
  def basic_hnn():
      """
      Simple Hamiltonian network.
  
      :return:
      """
      h = HNNMassSpring(nn.Sequential(
          nn.Linear(2, 50),
          nn.Tanh(),
          nn.Linear(50, 1))).to(device)
  
      model = DENNet(h, case='springmass').to(device)
      learn = Learner(model)
      logger = TensorBoardLogger('HNN_logs')
      trainer = pl.Trainer(gpus=1, min_epochs=50, max_epochs=300, logger=logger)
      trainer.fit(learn)
  
      return h, model
  
  
  def separable_hnn(input_h_s=None, input_model=None):
      """
      Separable Hamiltonian network.
  
      :return:
      """
      if input_h_s:
          h_s = input_h_s
          model = input_model
      else:
          h_s = HNNMassSpringSeparable(nn.Sequential(
              nn.Linear(2, 100),
              nn.Tanh(),
              nn.Linear(100, 1))).to(device)
          model = DENNet(h_s, case='springmass').to(device)
  
      learn_sep = Learner(model)
      logger = TensorBoardLogger('separable_logs')
      trainer_sep = pl.Trainer(min_epochs=50, max_epochs=100, logger=logger)
      trainer_sep.fit(learn_sep)
  
      return h_s, model
  
  
  def separable_hnn_siren(input_h_s=None, input_model=None):
      """
      Separable Hamiltonian network using Siren network.
  
      :return:
      """
      if input_h_s:
          h_s = input_h_s
          model = input_model
      else:
          network = SirenNet(
              dim_in=2,
              dim_hidden=100,
              dim_out=1,
              num_layers=5,
              w0_initial=30.
          )
  
          h_s = HNNMassSpringSeparable(network).to(device)
          model = DENNet(h_s).to(device)
  
      learn_sep = Learner(model)
      logger = TensorBoardLogger('separable_logs_siren')
      trainer_sep = pl.Trainer(min_epochs=50, max_epochs=100, logger=logger, gpus=1)
      trainer_sep.fit(learn_sep)
  
      return h_s, model
  
  
  if __name__ == '__main__':
      device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  
      # Training conditions
      num_train_data = 100
      num_tSteps_training = 5
  
      X_sv = torch.cat([  # Training initial conditions
          (3 * torch.rand(num_train_data) - 1.5).unsqueeze(1),
          (3 * torch.rand(num_train_data) - 1.5).unsqueeze(1)
      ], 1).to(device)
  
      dt_train = 0.05  # Training time step
  
      # Testing conditions
      q_init = torch.linspace(0.2, 1.5, 3)
      p_init = torch.zeros(q_init.shape)
      x_init = torch.cat([q_init.unsqueeze(1), p_init.unsqueeze(1)], 1).to(device)
  
      t_span_test = torch.linspace(0, 20, 400).to(device)  # Testing time span
  
      # Wrap in for loop and change inputs to the stepped forward p's and q's
      for tStep in range(num_tSteps_training):
  
          train = data.TensorDataset(X_sv)
          trainloader = data.DataLoader(train, batch_size=len(X_sv), shuffle=False)
  
          # hamiltonian, basic_model = basic_hnn()
          if tStep == 0:
              separable, separable_model = separable_hnn_siren()
          else:
              separable, separable_model = separable_hnn_siren(input_h_s=separable, input_model=separable_model)
  
          time_integrator_sv = TimeIntegrator(separable).to(device)
  
          X_sv = time_integrator_sv.sv_step(X_sv, dt_train).detach()
  
      # set up time integrator that uses our separable HNN
      time_integrator_sv = TimeIntegrator(separable).to(device)
      # calculate trajectory
      traj_HNN_sv = time_integrator_sv.integrate(x_init, t_span_test, method='SV').detach().cpu()
  
      n_grid = 50
      x = torch.linspace(-2, 2, n_grid)
      Q, P = torch.meshgrid(x, x)
      H, U, V = torch.zeros(Q.shape), torch.zeros(Q.shape), torch.zeros(Q.shape)
      for i in range(n_grid):
          for j in range(n_grid):
              x = torch.cat([Q[i, j].reshape(1, 1), P[i, j].reshape(1, 1)], 1).to(device)
              H[i, j] = separable_model.de_function.model.H(x).detach().cpu()
              O = separable_model.de_function(0, x).detach().cpu()
              U[i, j], V[i, j] = O[0, 0], O[0, 1]
  
      fig = plt.figure(figsize=(8, 8))
      ax = fig.add_subplot(111)
      ax.contourf(Q, P, H, 100, cmap='seismic')
      # ax.streamplot(Q.T.numpy(), P.T.numpy(), U.T.numpy(), V.T.numpy(), color='black')
      # ax.plot(traj[:, 0, 0], traj[:, 0, 1], color='k')
      for count in range(len(traj_HNN_sv[:, 0, 0]) - 1):
          # ax.plot(traj_HNN_Euler[count, 0, :], traj_HNN_Euler[count, 1, :], color='y')
          ax.plot(traj_HNN_sv[count, 0, :], traj_HNN_sv[count, 1, :], color='g')
      # plot last index with a label for legend
      count = count + 1
      # ax.plot(traj_HNN_Euler[count, 0, :], traj_HNN_Euler[count, 1, :], color='y', label='Euler')
      ax.plot(traj_HNN_sv[count, 0, :], traj_HNN_sv[count, 1, :], color='g', label='Stormer-Verlet')
  
      ax.legend()
      ax.set_xlim([Q.min(), Q.max()])
      ax.set_ylim([P.min(), P.max()])
      ax.set_xlabel(r"$q$")
      ax.set_ylabel(r"$p$")
      plt.show()
  