text: import objax
  from jax import vmap
  import jax.numpy as np
  from jax.scipy.linalg import cho_factor, cho_solve, block_diag, expm
  from jax.ops import index_add, index
  from .utils import scaled_squared_euclid_dist, softplus, softplus_inv, rotation_matrix
  from warnings import warn
  
  
  class Kernel(objax.Module):
      """
      """
  
      def __call__(self, X, X2):
          return self.K(X, X2)
  
      def K(self, X, X2):
          raise NotImplementedError('kernel function not implemented')
  
      def measurement_model(self):
          raise NotImplementedError
  
      def inducing_precision(self):
          return None, None
  
      def kernel_to_state_space(self, R=None):
          raise NotImplementedError
  
      def spatial_conditional(self, R=None, predict=False):
          """
          """
          return None, None
  
      def get_meanfield_block_index(self):
          raise Exception('Either the mean-field method is not applicable to this kernel, '
                          'or this kernel\'s get_meanfield_block_index() method has not been implemented')
  
      def feedback_matrix(self):
          raise NotImplementedError
  
      def state_transition(self, dt):
          F = self.feedback_matrix()
          A = expm(F * dt)
          return A
  
  
  class StationaryKernel(Kernel):
      """
      """
  
      def __init__(self,
                   variance=1.0,
                   lengthscale=1.0,
                   fix_variance=False,
                   fix_lengthscale=False):
          # check whether the parameters are to be optimised
          if fix_lengthscale:
              self.transformed_lengthscale = objax.StateVar(softplus_inv(np.array(lengthscale)))
          else:
              self.transformed_lengthscale = objax.TrainVar(softplus_inv(np.array(lengthscale)))
          if fix_variance:
              self.transformed_variance = objax.StateVar(softplus_inv(np.array(variance)))
          else:
              self.transformed_variance = objax.TrainVar(softplus_inv(np.array(variance)))
  
      @property
      def variance(self):
          return softplus(self.transformed_variance.value)
  
      @property
      def lengthscale(self):
          return softplus(self.transformed_lengthscale.value)
  
      def K(self, X, X2):
          r2 = scaled_squared_euclid_dist(X, X2, self.lengthscale)
          return self.K_r2(r2)
  
      def K_r2(self, r2):
          # Clipping around the (single) float precision which is ~1e-45.
          r = np.sqrt(np.maximum(r2, 1e-36))
          return self.K_r(r)
  
      @staticmethod
      def K_r(r):
          raise NotImplementedError('kernel not implemented')
  
      def kernel_to_state_space(self, R=None):
          raise NotImplementedError
  
      def measurement_model(self):
          raise NotImplementedError
  
      def stationary_covariance(self):
          raise NotImplementedError
  
      def feedback_matrix(self):
          raise NotImplementedError
  
  
  class Matern12(StationaryKernel):
      """
      The Matern 1/2 kernel. Functions drawn from a GP with this kernel are not
      differentiable anywhere. The kernel equation is
  
      k(r) = σ² exp{-r}
  
      where:
      r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ.
      σ² is the variance parameter
      """
  
      @property
      def state_dim(self):
          return 1
  
      def K_r(self, r):
          return self.variance * np.exp(-r)
  
      def kernel_to_state_space(self, R=None):
          F = np.array([[-1.0 / self.lengthscale]])
          L = np.array([[1.0]])
          Qc = np.array([[2.0 * self.variance / self.lengthscale]])
          H = np.array([[1.0]])
          Pinf = np.array([[self.variance]])
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          Pinf = np.array([[self.variance]])
          return Pinf
  
      def measurement_model(self):
          H = np.array([[1.0]])
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the discrete-time state transition matrix A = expm(FΔt) for the exponential prior.
          :param dt: step size(s), Δtₙ = tₙ - tₙ₋₁ [scalar]
          :return: state transition matrix A [1, 1]
          """
          A = np.broadcast_to(np.exp(-dt / self.lengthscale), [1, 1])
          return A
  
      def feedback_matrix(self):
          F = np.array([[-1.0 / self.lengthscale]])
          return F
  
  
  Exponential = Matern12
  
  
  class Matern32(StationaryKernel):
      """
      The Matern 3/2 kernel. Functions drawn from a GP with this kernel are once
      differentiable. The kernel equation is
  
      k(r) = σ² (1 + √3r) exp{-√3 r}
  
      where:
      r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ,
      σ² is the variance parameter.
      """
  
      @property
      def state_dim(self):
          return 2
  
      def K_r(self, r):
          sqrt3 = np.sqrt(3.0)
          return self.variance * (1.0 + sqrt3 * r) * np.exp(-sqrt3 * r)
  
      def kernel_to_state_space(self, R=None):
          lam = 3.0 ** 0.5 / self.lengthscale
          F = np.array([[0.0,       1.0],
                        [-lam ** 2, -2 * lam]])
          L = np.array([[0],
                        [1]])
          Qc = np.array([[12.0 * 3.0 ** 0.5 / self.lengthscale ** 3.0 * self.variance]])
          H = np.array([[1.0, 0.0]])
          Pinf = np.array([[self.variance, 0.0],
                           [0.0, 3.0 * self.variance / self.lengthscale ** 2.0]])
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          Pinf = np.array([[self.variance, 0.0],
                           [0.0, 3.0 * self.variance / self.lengthscale ** 2.0]])
          return Pinf
  
      def measurement_model(self):
          H = np.array([[1.0, 0.0]])
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the discrete-time state transition matrix A = expm(FΔt) for the Matern-3/2 prior.
          :param dt: step size(s), Δtₙ = tₙ - tₙ₋₁ [scalar]
          :return: state transition matrix A [2, 2]
          """
          lam = np.sqrt(3.0) / self.lengthscale
          A = np.exp(-dt * lam) * (dt * np.array([[lam, 1.0], [-lam**2.0, -lam]]) + np.eye(2))
          return A
  
      def feedback_matrix(self):
          lam = 3.0 ** 0.5 / self.lengthscale
          F = np.array([[0.0, 1.0],
                        [-lam ** 2, -2 * lam]])
          return F
  
  
  class Matern52(StationaryKernel):
      """
      The Matern 5/2 kernel. Functions drawn from a GP with this kernel are twice
      differentiable. The kernel equation is
  
      k(r) = σ² (1 + √5r + 5/3r²) exp{-√5 r}
  
      where:
      r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ,
      σ² is the variance parameter.
      """
  
      @property
      def state_dim(self):
          return 3
  
      def K_r(self, r):
          sqrt5 = np.sqrt(5.0)
          return self.variance * (1.0 + sqrt5 * r + 5.0 / 3.0 * np.square(r)) * np.exp(-sqrt5 * r)
  
      def kernel_to_state_space(self, R=None):
          lam = 5.0**0.5 / self.lengthscale
          F = np.array([[0.0, 1.0, 0.0],
                        [0.0, 0.0, 1.0],
                        [-lam**3.0, -3.0*lam**2.0, -3.0*lam]])
          L = np.array([[0.0],
                        [0.0],
                        [1.0]])
          Qc = np.array([[self.variance * 400.0 * 5.0 ** 0.5 / 3.0 / self.lengthscale ** 5.0]])
          H = np.array([[1.0, 0.0, 0.0]])
          kappa = 5.0 / 3.0 * self.variance / self.lengthscale**2.0
          Pinf = np.array([[self.variance,    0.0,   -kappa],
                           [0.0,    kappa, 0.0],
                           [-kappa, 0.0,   25.0*self.variance / self.lengthscale**4.0]])
          return F, L, Qc, H, Pinf
  
      def measurement_model(self):
          H = np.array([[1.0, 0.0, 0.0]])
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the discrete-time state transition matrix A = expm(FΔt) for the Matern-5/2 prior.
          :param dt: step size(s), Δtₙ = tₙ - tₙ₋₁ [scalar]
          :return: state transition matrix A [3, 3]
          """
          lam = np.sqrt(5.0) / self.lengthscale
          dtlam = dt * lam
          A = np.exp(-dtlam) \
              * (dt * np.array([[lam * (0.5 * dtlam + 1.0),      dtlam + 1.0,            0.5 * dt],
                                [-0.5 * dtlam * lam ** 2,        lam * (1.0 - dtlam),    1.0 - 0.5 * dtlam],
                                [lam ** 3 * (0.5 * dtlam - 1.0), lam ** 2 * (dtlam - 3), lam * (0.5 * dtlam - 2.0)]])
                 + np.eye(3))
          return A
  
      def stationary_covariance(self):
          kappa = 5.0 / 3.0 * self.variance / self.lengthscale**2.0
          Pinf = np.array([[self.variance,    0.0,   -kappa],
                           [0.0,    kappa, 0.0],
                           [-kappa, 0.0,   25.0*self.variance / self.lengthscale**4.0]])
          return Pinf
  
      def feedback_matrix(self):
          lam = 5.0**0.5 / self.lengthscale
          F = np.array([[0.0, 1.0, 0.0],
                        [0.0, 0.0, 1.0],
                        [-lam**3.0, -3.0*lam**2.0, -3.0*lam]])
          return F
  
  
  class Matern72(StationaryKernel):
      """
      The Matern 7/2 kernel. Functions drawn from a GP with this kernel are three times differentiable.
  
      where:
      r  is the Euclidean distance between the input points, scaled by the lengthscales parameter ℓ,
      σ² is the variance parameter.
      """
  
      @property
      def state_dim(self):
          return 4
  
      def K_r(self, r):
          sqrt7 = np.sqrt(7.0)
          return self.variance * (1. + sqrt7 * r + 14. / 5. * np.square(r) + 7. * sqrt7 / 15. * r**3) * np.exp(-sqrt7 * r)
  
      def kernel_to_state_space(self, R=None):
          lam = 7.0**0.5 / self.lengthscale
          F = np.array([[0.0,       1.0,           0.0,           0.0],
                        [0.0,       0.0,           1.0,           0.0],
                        [0.0,       0.0,           0.0,           1.0],
                        [-lam**4.0, -4.0*lam**3.0, -6.0*lam**2.0, -4.0*lam]])
          L = np.array([[0.0],
                        [0.0],
                        [0.0],
                        [1.0]])
          Qc = np.array([[self.variance * 10976.0 * 7.0 ** 0.5 / 5.0 / self.lengthscale ** 7.0]])
          H = np.array([[1, 0, 0, 0]])
          kappa = 7.0 / 5.0 * self.variance / self.lengthscale**2.0
          kappa2 = 9.8 * self.variance / self.lengthscale**4.0
          Pinf = np.array([[self.variance,   0.0,    -kappa, 0.0],
                           [0.0,    kappa,   0.0,    -kappa2],
                           [-kappa, 0.0,     kappa2, 0.0],
                           [0.0,    -kappa2, 0.0,    343.0*self.variance / self.lengthscale**6.0]])
          return F, L, Qc, H, Pinf
  
      def measurement_model(self):
          H = np.array([[1.0, 0.0, 0.0, 0.0]])
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the discrete-time state transition matrix A = expm(FΔt) for the Matern-7/2 prior.
          :param dt: step size(s), Δtₙ = tₙ - tₙ₋₁ [scalar]
          :return: state transition matrix A [4, 4]
          """
          lam = np.sqrt(7.0) / self.lengthscale
          lam2 = lam * lam
          lam3 = lam2 * lam
          dtlam = dt * lam
          dtlam2 = dtlam ** 2
          A = np.exp(-dtlam) \
              * (dt * np.array([[lam * (1.0 + 0.5 * dtlam + dtlam2 / 6.0),      1.0 + dtlam + 0.5 * dtlam2,
                                0.5 * dt * (1.0 + dtlam),                       dt ** 2 / 6],
                                [-dtlam2 * lam ** 2.0 / 6.0,                    lam * (1.0 + 0.5 * dtlam - 0.5 * dtlam2),
                                1.0 + dtlam - 0.5 * dtlam2,                     dt * (0.5 - dtlam / 6.0)],
                                [lam3 * dtlam * (dtlam / 6.0 - 0.5),            dtlam * lam2 * (0.5 * dtlam - 2.0),
                                lam * (1.0 - 2.5 * dtlam + 0.5 * dtlam2),       1.0 - dtlam + dtlam2 / 6.0],
                                [lam2 ** 2 * (dtlam - 1.0 - dtlam2 / 6.0),      lam3 * (3.5 * dtlam - 4.0 - 0.5 * dtlam2),
                                lam2 * (4.0 * dtlam - 6.0 - 0.5 * dtlam2),      lam * (1.5 * dtlam - 3.0 - dtlam2 / 6.0)]])
                 + np.eye(4))
          return A
  
      def stationary_covariance(self):
          kappa = 7.0 / 5.0 * self.variance / self.lengthscale ** 2.0
          kappa2 = 9.8 * self.variance / self.lengthscale ** 4.0
          Pinf = np.array([[self.variance, 0.0, -kappa, 0.0],
                           [0.0, kappa, 0.0, -kappa2],
                           [-kappa, 0.0, kappa2, 0.0],
                           [0.0, -kappa2, 0.0, 343.0 * self.variance / self.lengthscale ** 6.0]])
          return Pinf
  
      def feedback_matrix(self):
          lam = 7.0 ** 0.5 / self.lengthscale
          F = np.array([[0.0, 1.0, 0.0, 0.0],
                        [0.0, 0.0, 1.0, 0.0],
                        [0.0, 0.0, 0.0, 1.0],
                        [-lam ** 4.0, -4.0 * lam ** 3.0, -6.0 * lam ** 2.0, -4.0 * lam]])
          return F
  
  
  class SpatioTemporalKernel(Kernel):
      """
      The Spatio-Temporal GP class
      :param temporal_kernel: the temporal prior, must be a member of the Prior class
      :param spatial_kernel: the kernel used for the spatial dimensions
      :param z: the initial spatial locations
      :param conditional: specifies which method to use for computing the covariance of the spatial conditional;
                          must be one of ['DTC', 'FIC', 'Full']
      :param sparse: boolean specifying whether the model is sparse in space
      :param opt_z: boolean specifying whether to optimise the spatial input locations z
      """
      def __init__(self,
                   temporal_kernel,
                   spatial_kernel,
                   z=None,
                   conditional=None,
                   sparse=True,
                   opt_z=False,
                   spatial_dims=None):
          self.temporal_kernel = temporal_kernel
          self.spatial_kernel = spatial_kernel
          if conditional is None:
              if sparse:
                  conditional = 'Full'
              else:
                  conditional = 'DTC'
          if opt_z and (not sparse):  # z should not be optimised if the model is not sparse
              warn("spatial inducing inputs z will not be optimised because sparse=False")
              opt_z = False
          self.sparse = sparse
          if z is None:  # initialise z
              # TODO: smart initialisation
              if spatial_dims == 1:
                  z = np.linspace(-3., 3., num=15)
              elif spatial_dims == 2:
                  z1 = np.linspace(-3., 3., num=5)
                  zA, zB = np.meshgrid(z1, z1)  # Adding additional dimension to inducing points grid
                  z = np.hstack((zA.reshape(-1, 1), zB.reshape(-1, 1)))  # Flattening grid for use in kernel functions
              else:
                  raise NotImplementedError('please provide an initialisation for inducing inputs z')
          if z.ndim < 2:
              z = z[:, np.newaxis]
          if spatial_dims is None:
              spatial_dims = z.ndim - 1
          assert spatial_dims == z.ndim - 1
          self.M = z.shape[0]
          if opt_z:
              self.z = objax.TrainVar(np.array(z))  # .reshape(-1, 1)
          else:
              self.z = objax.StateVar(np.array(z))
          if conditional in ['DTC', 'dtc']:
              self.conditional_covariance = self.deterministic_training_conditional
          elif conditional in ['FIC', 'FITC', 'fic', 'fitc']:
              self.conditional_covariance = self.fully_independent_conditional
          elif conditional in ['Full', 'full']:
              self.conditional_covariance = self.full_conditional
          else:
              raise NotImplementedError('conditional method not recognised')
          if (not sparse) and (conditional != 'DTC'):
              warn("You chose a non-deterministic conditional, but \'DTC\' will be used because the model is not sparse")
  
      @property
      def variance(self):
          return self.temporal_kernel.variance
  
      @property
      def temporal_lengthscale(self):
          return self.temporal_kernel.lengthscale
  
      @property
      def spatial_lengthscale(self):
          return self.spatial_kernel.lengthscale
  
      @property
      def state_dim(self):
          return self.temporal_kernel.state_dim
  
      def K(self, X, X2):
          T = X[:, :1]
          T2 = X2[:, :1]
          R = X[:, 1:]
          R2 = X2[:, 1:]
          return self.temporal_kernel(T, T2) * self.spatial_kernel(R, R2)
  
      @staticmethod
      def deterministic_training_conditional(X, R, Krz, K):
          cov = np.array([[0.0]])
          return cov
  
      def fully_independent_conditional(self, X, R, Krz, K):
          Krr = self.spatial_kernel(R, R)
          X = X.reshape(-1, 1)
          cov = self.temporal_kernel.K(X, X) * (np.diag(np.diag(Krr - K @ Krz.T)))
          return cov
  
      def full_conditional(self, X, R, Krz, K):
          Krr = self.spatial_kernel(R, R)
          X = X.reshape(-1, 1)
          cov = self.temporal_kernel.K(X, X) * (Krr - K @ Krz.T)
          return cov
  
      def spatial_conditional(self, X=None, R=None, predict=False):
          """
          Compute the spatial conditional, i.e. the measurement model projecting the latent function u(t) to f(X,R)
              f(X,R) | u(t) ~ N(f(X,R) | B u(t), C)
          """
          Qzz, Lzz = self.inducing_precision()  # pre-calculate inducing precision and its Cholesky factor
          if self.sparse or predict:
              # TODO: save compute if R is constant:
              # gridded_data = np.all(np.abs(np.diff(R, axis=0)) < 1e-10)
              # if gridded_data:
              #     R = R[:1]
              R = R.reshape((R.shape[0],) + (-1,) + self.z.value.shape[1:])
              Krz = vmap(self.spatial_kernel, [0, None])(R, self.z.value)
              K = Krz @ Qzz  # Krz / Kzz
              B = K @ Lzz
              C = vmap(self.conditional_covariance)(X, R, Krz, K)  # conditional covariance
          else:
              B = Lzz
              # conditional covariance (deterministic mapping is exact in non-sparse case)
              C = np.zeros([B.shape[0], B.shape[0]])
          return B, C
  
      def inducing_precision(self):
          """
          Compute the covariance and precision of the inducing spatial points to be used during filtering
          """
          Kzz = self.spatial_kernel(self.z.value, self.z.value)
          Lzz, low = cho_factor(Kzz, lower=True)  # K_zz^(1/2)
          Qzz = cho_solve((Lzz, low), np.eye(self.M))  # K_zz^(-1)
          return Qzz, Lzz
  
      def stationary_covariance(self):
          """
          Compute the covariance of the stationary state distribution. Since the latent components are independent
          under the prior, this is a block-diagonal matrix
          """
          Pinf_time = self.temporal_kernel.stationary_covariance()
          Pinf = np.kron(np.eye(self.M), Pinf_time)
          return Pinf
  
      def stationary_covariance_meanfield(self):
          """
          Stationary covariance as a tensor of blocks, as required when using a mean-field assumption
          """
          Pinf_time = self.temporal_kernel.stationary_covariance()
          Pinf = np.tile(Pinf_time, [self.M, 1, 1])
          return Pinf
  
      def measurement_model(self):
          """
          Compute the spatial conditional, i.e. the measurement model projecting the state x(t) to function space
              f(t, R) = H x(t)
          """
          H_time = self.temporal_kernel.measurement_model()
          H = np.kron(np.eye(self.M), H_time)
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the discrete-time state transition matrix A = expm(FΔt) for the spatio-temporal prior.
          :param dt: step size(s), Δtₙ = tₙ - tₙ₋₁ [scalar]
          :return: state transition matrix A
          """
          A_time = self.temporal_kernel.state_transition(dt)
          A = np.kron(np.eye(self.M), A_time)
          return A
  
      def state_transition_meanfield(self, dt):
          """
          State transition matrix in the form required for mean-field inference.
          :param dt: step size(s), Δtₙ = tₙ - tₙ₋₁ [scalar]
          :return: state transition matrix A
          """
          A_time = self.temporal_kernel.state_transition(dt)
          A = np.tile(A_time, [self.M, 1, 1])
          return A
  
      def kernel_to_state_space(self, R=None):
          F_t, L_t, Qc_t, H_t, Pinf_t = self.temporal_kernel.kernel_to_state_space()
          Kzz = self.spatial_kernel(self.z.value, self.z.value)
          F = np.kron(np.eye(self.M), F_t)
          Qc = None
          L = None
          H = self.measurement_model()
          Pinf = np.kron(Kzz, Pinf_t)
          return F, L, Qc, H, Pinf
  
      def get_meanfield_block_index(self):
          Pinf = self.stationary_covariance_meanfield()
          num_latents = Pinf.shape[0]
          sub_state_dim = Pinf.shape[1]
          state = np.ones([sub_state_dim, sub_state_dim])
          for i in range(1, num_latents):
              state = block_diag(state, np.ones([sub_state_dim, sub_state_dim]))
          block_index = np.where(np.array(state, dtype=bool))
          return block_index
  
      def feedback_matrix(self):
          F_t = self.temporal_kernel.feedback_matrix()
          F = np.kron(np.eye(self.M), F_t)
          return F
  
  
  class SpatioTemporalMatern12(SpatioTemporalKernel):
      """
      Spatio-Temporal Matern-1/2 kernel in SDE form.
      Hyperparameters:
          variance, σ²
          temporal lengthscale, lt
          spatial lengthscale, ls
      """
      def __init__(self,
                   variance=1.0,
                   lengthscale_time=1.0,
                   lengthscale_space=1.0,
                   z=None,
                   sparse=True,
                   opt_z=False,
                   conditional=None):
          super().__init__(temporal_kernel=Matern12(variance=variance, lengthscale=lengthscale_time),
                           spatial_kernel=Matern12(variance=1., lengthscale=lengthscale_space, fix_variance=True),
                           z=z,
                           conditional=conditional,
                           sparse=sparse,
                           opt_z=opt_z)
          self.name = 'Spatio-Temporal Matern-1/2'
  
  
  class SpatioTemporalMatern32(SpatioTemporalKernel):
      """
      Spatio-Temporal Matern-3/2 kernel in SDE form.
      Hyperparameters:
          variance, σ²
          temporal lengthscale, lt
          spatial lengthscale, ls
      """
      def __init__(self,
                   variance=1.0,
                   lengthscale_time=1.0,
                   lengthscale_space=1.0,
                   z=None,
                   sparse=True,
                   opt_z=False,
                   conditional=None):
          super().__init__(temporal_kernel=Matern32(variance=variance, lengthscale=lengthscale_time),
                           spatial_kernel=Matern32(variance=1., lengthscale=lengthscale_space, fix_variance=True),
                           z=z,
                           conditional=conditional,
                           sparse=sparse,
                           opt_z=opt_z)
          self.name = 'Spatio-Temporal Matern-3/2'
  
  
  class SpatioTemporalMatern52(SpatioTemporalKernel):
      """
      Spatio-Temporal Matern-5/2 kernel in SDE form.
      Hyperparameters:
          variance, σ²
          temporal lengthscale, lt
          spatial lengthscale, ls
      """
      def __init__(self,
                   variance=1.0,
                   lengthscale_time=1.0,
                   lengthscale_space=1.0,
                   z=None,
                   sparse=True,
                   opt_z=False,
                   conditional=None):
          super().__init__(temporal_kernel=Matern52(variance=variance, lengthscale=lengthscale_time),
                           spatial_kernel=Matern52(variance=1., lengthscale=lengthscale_space, fix_variance=True),
                           z=z,
                           conditional=conditional,
                           sparse=sparse,
                           opt_z=opt_z)
          self.name = 'Spatio-Temporal Matern-5/2'
  
  
  class SpatialMatern12(SpatioTemporalKernel):
      """
      Spatial Matern-1/2 kernel in SDE form. Similar to the spatio-temporal kernel but the
      lengthscale is shared across dimensions.
      Hyperparameters:
          variance, σ²
          lengthscale, l
      """
      def __init__(self,
                   variance=1.0,
                   lengthscale=1.0,
                   z=None,
                   sparse=True,
                   opt_z=False,
                   conditional=None):
          super().__init__(temporal_kernel=Matern12(variance=variance, lengthscale=lengthscale),
                           spatial_kernel=Matern12(variance=1., lengthscale=lengthscale, fix_variance=True),
                           z=z,
                           conditional=conditional,
                           sparse=sparse,
                           opt_z=opt_z)
          # --- couple the lengthscales ---
          self.spatial_kernel.transformed_lengthscale = self.temporal_kernel.transformed_lengthscale
          # -------------------------------
          self.name = 'Spatial Matern-1/2'
  
  
  class SpatialMatern32(SpatioTemporalKernel):
      """
      Spatial Matern-3/2 kernel in SDE form. Similar to the spatio-temporal kernel but the
      lengthscale is shared across dimensions.
      Hyperparameters:
          variance, σ²
          lengthscale, l
      """
      def __init__(self,
                   variance=1.0,
                   lengthscale=1.0,
                   z=None,
                   sparse=True,
                   opt_z=False,
                   conditional=None):
          super().__init__(temporal_kernel=Matern32(variance=variance, lengthscale=lengthscale),
                           spatial_kernel=Matern32(variance=1., lengthscale=lengthscale, fix_variance=True),
                           z=z,
                           conditional=conditional,
                           sparse=sparse,
                           opt_z=opt_z)
          # --- couple the lengthscales ---
          self.spatial_kernel.transformed_lengthscale = self.temporal_kernel.transformed_lengthscale
          # -------------------------------
          self.name = 'Spatial Matern-3/2'
  
  
  class SpatialMatern52(SpatioTemporalKernel):
      """
      Spatial Matern-5/2 kernel in SDE form. Similar to the spatio-temporal kernel but the
      lengthscale is shared across dimensions.
      Hyperparameters:
          variance, σ²
          lengthscale, l
      """
      def __init__(self,
                   variance=1.0,
                   lengthscale=1.0,
                   z=None,
                   sparse=True,
                   opt_z=False,
                   conditional=None):
          super().__init__(temporal_kernel=Matern52(variance=variance, lengthscale=lengthscale),
                           spatial_kernel=Matern52(variance=1., lengthscale=lengthscale, fix_variance=True),
                           z=z,
                           conditional=conditional,
                           sparse=sparse,
                           opt_z=opt_z)
          # --- couple the lengthscales ---
          self.spatial_kernel.transformed_lengthscale = self.temporal_kernel.transformed_lengthscale
          # -------------------------------
          self.name = 'Spatial Matern-5/2'
  
  
  class Cosine(Kernel):
      """
      Cosine kernel in SDE form.
      Hyperparameters:
          radial frequency, ω
      The associated continuous-time state space model matrices are:
      F      = ( 0   -ω
                 ω    0 )
      L      = N/A
      Qc     = N/A
      H      = ( 1  0 )
      Pinf   = ( 1  0
                 0  1 )
      and the discrete-time transition matrix is (for step size Δt),
      A      = ( cos(ωΔt)   -sin(ωΔt)
                 sin(ωΔt)    cos(ωΔt) )
      """
      def __init__(self, frequency=1.0):
          self.transformed_frequency = objax.TrainVar(np.array(softplus_inv(frequency)))
          super().__init__()
          self.name = 'Cosine'
  
      @property
      def frequency(self):
          return softplus(self.transformed_frequency.value)
  
      def kernel_to_state_space(self, R=None):
          F = np.array([[0.0,   -self.frequency],
                        [self.frequency, 0.0]])
          H = np.array([[1.0, 0.0]])
          L = []
          Qc = []
          Pinf = np.eye(2)
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          Pinf = np.eye(2)
          return Pinf
  
      def measurement_model(self):
          H = np.array([[1.0, 0.0]])
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the closed form discrete-time state
          transition matrix A = expm(FΔt) for the Cosine prior
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [M+1, 1]
          :return: state transition matrix A [M+1, D, D]
          """
          state_transitions = rotation_matrix(dt, self.frequency)  # [2, 2]
          return state_transitions
  
      def feedback_matrix(self):
          F = np.array([[0.0, -self.frequency],
                        [self.frequency, 0.0]])
          return F
  
  
  class QuasiPeriodicMatern12(Kernel):
      """
      TODO: implement a general 'Product' class to reduce code duplication
      Quasi-periodic kernel in SDE form (product of Periodic and Matern-1/2).
      Hyperparameters:
          variance, σ²
          lengthscale of Periodic, l_p
          period, p
          lengthscale of Matern, l_m
      The associated continuous-time state space model matrices are constructed via
      a sum of cosines times a Matern-1/2.
      """
      def __init__(self, variance=1.0, lengthscale_periodic=1.0, period=1.0, lengthscale_matern=1.0, order=6):
          self.transformed_lengthscale_periodic = objax.TrainVar(np.array(softplus_inv(lengthscale_periodic)))
          self.transformed_variance = objax.TrainVar(np.array(softplus_inv(variance)))
          self.transformed_period = objax.TrainVar(np.array(softplus_inv(period)))
          self.transformed_lengthscale_matern = objax.TrainVar(np.array(softplus_inv(lengthscale_matern)))
          super().__init__()
          self.name = 'Quasi-periodic Matern-1/2'
          self.order = order
          self.igrid = np.meshgrid(np.arange(self.order + 1), np.arange(self.order + 1))[1]
          factorial_mesh_K = np.array([[1., 1., 1., 1., 1., 1., 1.],
                                       [1., 1., 1., 1., 1., 1., 1.],
                                       [2., 2., 2., 2., 2., 2., 2.],
                                       [6., 6., 6., 6., 6., 6., 6.],
                                       [24., 24., 24., 24., 24., 24., 24.],
                                       [120., 120., 120., 120., 120., 120., 120.],
                                       [720., 720., 720., 720., 720., 720., 720.]])
          b = np.array([[1., 0., 0., 0., 0., 0., 0.],
                        [0., 2., 0., 0., 0., 0., 0.],
                        [2., 0., 2., 0., 0., 0., 0.],
                        [0., 6., 0., 2., 0., 0., 0.],
                        [6., 0., 8., 0., 2., 0., 0.],
                        [0., 20., 0., 10., 0., 2., 0.],
                        [20., 0., 30., 0., 12., 0., 2.]])
          self.b_fmK_2igrid = b * (1. / factorial_mesh_K) * (2. ** -self.igrid)
  
      @property
      def variance(self):
          return softplus(self.transformed_variance.value)
  
      @property
      def lengthscale_periodic(self):
          return softplus(self.transformed_lengthscale_periodic.value)
  
      @property
      def lengthscale_matern(self):
          return softplus(self.transformed_lengthscale_matern.value)
  
      @property
      def period(self):
          return softplus(self.transformed_period.value)
  
      def K(self, X, X2):
          raise NotImplementedError
  
      def kernel_to_state_space(self, R=None):
          var_p = 1.
          ell_p = self.lengthscale_periodic
          a = self.b_fmK_2igrid * ell_p ** (-2. * self.igrid) * np.exp(-1. / ell_p ** 2.) * var_p
          q2 = np.sum(a, axis=0)
          # The angular frequency
          omega = 2 * np.pi / self.period
          # The model
          F_p = np.kron(np.diag(np.arange(self.order + 1)), np.array([[0., -omega], [omega, 0.]]))
          L_p = np.eye(2 * (self.order + 1))
          # Qc_p = np.zeros(2 * (self.N + 1))
          Pinf_p = np.kron(np.diag(q2), np.eye(2))
          H_p = np.kron(np.ones([1, self.order + 1]), np.array([1., 0.]))
          F_m = np.array([[-1.0 / self.lengthscale_matern]])
          L_m = np.array([[1.0]])
          Qc_m = np.array([[2.0 * self.variance / self.lengthscale_matern]])
          H_m = np.array([[1.0]])
          Pinf_m = np.array([[self.variance]])
          F = np.kron(F_m, np.eye(2 * (self.order + 1))) + np.kron(np.eye(1), F_p)
          L = np.kron(L_m, L_p)
          Qc = np.kron(Qc_m, Pinf_p)
          H = np.kron(H_m, H_p)
          # Pinf = np.kron(Pinf_m, Pinf_p)
          Pinf = block_diag(
              np.kron(Pinf_m, q2[0] * np.eye(2)),
              np.kron(Pinf_m, q2[1] * np.eye(2)),
              np.kron(Pinf_m, q2[2] * np.eye(2)),
              np.kron(Pinf_m, q2[3] * np.eye(2)),
              np.kron(Pinf_m, q2[4] * np.eye(2)),
              np.kron(Pinf_m, q2[5] * np.eye(2)),
              np.kron(Pinf_m, q2[6] * np.eye(2)),
          )
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          var_p = 1.
          ell_p = self.lengthscale_periodic
          a = self.b_fmK_2igrid * ell_p ** (-2. * self.igrid) * np.exp(-1. / ell_p ** 2.) * var_p
          q2 = np.sum(a, axis=0)
          Pinf_m = np.array([[self.variance]])
          Pinf = block_diag(
              np.kron(Pinf_m, q2[0] * np.eye(2)),
              np.kron(Pinf_m, q2[1] * np.eye(2)),
              np.kron(Pinf_m, q2[2] * np.eye(2)),
              np.kron(Pinf_m, q2[3] * np.eye(2)),
              np.kron(Pinf_m, q2[4] * np.eye(2)),
              np.kron(Pinf_m, q2[5] * np.eye(2)),
              np.kron(Pinf_m, q2[6] * np.eye(2)),
          )
          return Pinf
  
      def measurement_model(self):
          H_p = np.kron(np.ones([1, self.order + 1]), np.array([1., 0.]))
          H_m = np.array([[1.0]])
          H = np.kron(H_m, H_p)
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the closed form discrete-time state
          transition matrix A = expm(FΔt) for the Quasi-Periodic Matern-3/2 prior
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [M+1, 1]
          :return: state transition matrix A [M+1, D, D]
          """
          # The angular frequency
          omega = 2 * np.pi / self.period
          harmonics = np.arange(self.order + 1) * omega
          R0 = rotation_matrix(dt, harmonics[0])
          R1 = rotation_matrix(dt, harmonics[1])
          R2 = rotation_matrix(dt, harmonics[2])
          R3 = rotation_matrix(dt, harmonics[3])
          R4 = rotation_matrix(dt, harmonics[4])
          R5 = rotation_matrix(dt, harmonics[5])
          R6 = rotation_matrix(dt, harmonics[6])
          A = np.exp(-dt / self.lengthscale_matern) * block_diag(R0, R1, R2, R3, R4, R5, R6)
          return A
  
      def feedback_matrix(self):
          # The angular frequency
          omega = 2 * np.pi / self.period
          # The model
          F_p = np.kron(np.diag(np.arange(self.order + 1)), np.array([[0., -omega], [omega, 0.]]))
          F_m = np.array([[-1.0 / self.lengthscale_matern]])
          F = np.kron(F_m, np.eye(2 * (self.order + 1))) + np.kron(np.eye(1), F_p)
          return F
  
  
  class QuasiPeriodicMatern32(Kernel):
      """
      Quasi-periodic kernel in SDE form (product of Periodic and Matern-3/2).
      Hyperparameters:
          variance, σ²
          lengthscale of Periodic, l_p
          period, p
          lengthscale of Matern, l_m
      The associated continuous-time state space model matrices are constructed via
      a sum of cosines times a Matern-3/2.
      """
      def __init__(self, variance=1.0, lengthscale_periodic=1.0, period=1.0, lengthscale_matern=1.0, order=6):
          self.transformed_lengthscale_periodic = objax.TrainVar(np.array(softplus_inv(lengthscale_periodic)))
          self.transformed_variance = objax.TrainVar(np.array(softplus_inv(variance)))
          self.transformed_period = objax.TrainVar(np.array(softplus_inv(period)))
          self.transformed_lengthscale_matern = objax.TrainVar(np.array(softplus_inv(lengthscale_matern)))
          super().__init__()
          self.name = 'Quasi-periodic Matern-3/2'
          self.order = order
          self.igrid = np.meshgrid(np.arange(self.order + 1), np.arange(self.order + 1))[1]
          factorial_mesh_K = np.array([[1., 1., 1., 1., 1., 1., 1.],
                                       [1., 1., 1., 1., 1., 1., 1.],
                                       [2., 2., 2., 2., 2., 2., 2.],
                                       [6., 6., 6., 6., 6., 6., 6.],
                                       [24., 24., 24., 24., 24., 24., 24.],
                                       [120., 120., 120., 120., 120., 120., 120.],
                                       [720., 720., 720., 720., 720., 720., 720.]])
          b = np.array([[1., 0., 0., 0., 0., 0., 0.],
                        [0., 2., 0., 0., 0., 0., 0.],
                        [2., 0., 2., 0., 0., 0., 0.],
                        [0., 6., 0., 2., 0., 0., 0.],
                        [6., 0., 8., 0., 2., 0., 0.],
                        [0., 20., 0., 10., 0., 2., 0.],
                        [20., 0., 30., 0., 12., 0., 2.]])
          self.b_fmK_2igrid = b * (1. / factorial_mesh_K) * (2. ** -self.igrid)
  
      @property
      def variance(self):
          return softplus(self.transformed_variance.value)
  
      @property
      def lengthscale_periodic(self):
          return softplus(self.transformed_lengthscale_periodic.value)
  
      @property
      def lengthscale_matern(self):
          return softplus(self.transformed_lengthscale_matern.value)
  
      @property
      def period(self):
          return softplus(self.transformed_period.value)
  
      def K(self, X, X2):
          raise NotImplementedError
  
      def kernel_to_state_space(self, R=None):
          var_p = 1.
          ell_p = self.lengthscale_periodic
          a = self.b_fmK_2igrid * ell_p ** (-2. * self.igrid) * np.exp(-1. / ell_p ** 2.) * var_p
          q2 = np.sum(a, axis=0)
          # The angular frequency
          omega = 2 * np.pi / self.period
          # The model
          F_p = np.kron(np.diag(np.arange(self.order + 1)), np.array([[0., -omega], [omega, 0.]]))
          L_p = np.eye(2 * (self.order + 1))
          # Qc_p = np.zeros(2 * (self.N + 1))
          Pinf_p = np.kron(np.diag(q2), np.eye(2))
          H_p = np.kron(np.ones([1, self.order + 1]), np.array([1., 0.]))
          lam = 3.0 ** 0.5 / self.lengthscale_matern
          F_m = np.array([[0.0, 1.0],
                          [-lam ** 2, -2 * lam]])
          L_m = np.array([[0],
                          [1]])
          Qc_m = np.array([[12.0 * 3.0 ** 0.5 / self.lengthscale_matern ** 3.0 * self.variance]])
          H_m = np.array([[1.0, 0.0]])
          Pinf_m = np.array([[self.variance, 0.0],
                             [0.0, 3.0 * self.variance / self.lengthscale_matern ** 2.0]])
          # F = np.kron(F_p, np.eye(2)) + np.kron(np.eye(14), F_m)
          F = np.kron(F_m, np.eye(2 * (self.order + 1))) + np.kron(np.eye(2), F_p)
          L = np.kron(L_m, L_p)
          Qc = np.kron(Qc_m, Pinf_p)
          H = np.kron(H_m, H_p)
          # Pinf = np.kron(Pinf_m, Pinf_p)
          Pinf = block_diag(
              np.kron(Pinf_m, q2[0] * np.eye(2)),
              np.kron(Pinf_m, q2[1] * np.eye(2)),
              np.kron(Pinf_m, q2[2] * np.eye(2)),
              np.kron(Pinf_m, q2[3] * np.eye(2)),
              np.kron(Pinf_m, q2[4] * np.eye(2)),
              np.kron(Pinf_m, q2[5] * np.eye(2)),
              np.kron(Pinf_m, q2[6] * np.eye(2)),
          )
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          var_p = 1.
          ell_p = self.lengthscale_periodic
          a = self.b_fmK_2igrid * ell_p ** (-2. * self.igrid) * np.exp(-1. / ell_p ** 2.) * var_p
          q2 = np.sum(a, axis=0)
          Pinf_m = np.array([[self.variance, 0.0],
                             [0.0, 3.0 * self.variance / self.lengthscale_matern ** 2.0]])
          Pinf = block_diag(
              np.kron(Pinf_m, q2[0] * np.eye(2)),
              np.kron(Pinf_m, q2[1] * np.eye(2)),
              np.kron(Pinf_m, q2[2] * np.eye(2)),
              np.kron(Pinf_m, q2[3] * np.eye(2)),
              np.kron(Pinf_m, q2[4] * np.eye(2)),
              np.kron(Pinf_m, q2[5] * np.eye(2)),
              np.kron(Pinf_m, q2[6] * np.eye(2)),
          )
          return Pinf
  
      def measurement_model(self):
          H_p = np.kron(np.ones([1, self.order + 1]), np.array([1., 0.]))
          H_m = np.array([[1.0, 0.0]])
          H = np.kron(H_m, H_p)
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the closed form discrete-time state
          transition matrix A = expm(FΔt) for the Quasi-Periodic Matern-3/2 prior
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [M+1, 1]
          :return: state transition matrix A [M+1, D, D]
          """
          lam = np.sqrt(3.0) / self.lengthscale_matern
          # The angular frequency
          omega = 2 * np.pi / self.period
          harmonics = np.arange(self.order + 1) * omega
          R0 = self.subband_mat32(dt, lam, harmonics[0])
          R1 = self.subband_mat32(dt, lam, harmonics[1])
          R2 = self.subband_mat32(dt, lam, harmonics[2])
          R3 = self.subband_mat32(dt, lam, harmonics[3])
          R4 = self.subband_mat32(dt, lam, harmonics[4])
          R5 = self.subband_mat32(dt, lam, harmonics[5])
          R6 = self.subband_mat32(dt, lam, harmonics[6])
          A = np.exp(-dt * lam) * block_diag(R0, R1, R2, R3, R4, R5, R6)
          return A
  
      @staticmethod
      def subband_mat32(dt, lam, omega):
          R = rotation_matrix(dt, omega)
          Ri = np.block([
              [(1. + dt * lam) * R, dt * R],
              [-dt * lam ** 2 * R,  (1. - dt * lam) * R]
          ])
          return Ri
  
      def feedback_matrix(self):
          # The angular frequency
          omega = 2 * np.pi / self.period
          # The model
          F_p = np.kron(np.diag(np.arange(self.order + 1)), np.array([[0., -omega], [omega, 0.]]))
          lam = 3.0 ** 0.5 / self.lengthscale_matern
          F_m = np.array([[0.0, 1.0],
                          [-lam ** 2, -2 * lam]])
          F = np.kron(F_m, np.eye(2 * (self.order + 1))) + np.kron(np.eye(2), F_p)
          return F
  
  
  class SubbandMatern12(Kernel):
      """
      Subband Matern-1/2 (i.e. Exponential) kernel in SDE form (product of Cosine and Matern-1/2).
      Hyperparameters:
          variance, σ²
          lengthscale, l
          radial frequency, ω
      The associated continuous-time state space model matrices are constructed via
      kronecker sums and products of the exponential and cosine components:
      F      = F_exp ⊕ F_cos  =  ( -1/l  -ω
                                   ω     -1/l )
      L      = L_exp ⊗ I      =  ( 1      0
                                   0      1 )
      Qc     = I ⊗ Qc_exp     =  ( 2σ²/l  0
                                   0      2σ²/l )
      H      = H_exp ⊗ H_cos  =  ( 1      0 )
      Pinf   = Pinf_exp ⊗ I   =  ( σ²     0
                                   0      σ² )
      and the discrete-time transition matrix is (for step size Δt),
      A      = exp(-Δt/l) ( cos(ωΔt)   -sin(ωΔt)
                            sin(ωΔt)    cos(ωΔt) )
      """
      def __init__(self, variance=1.0, lengthscale=1.0, radial_frequency=1.0, fix_variance=False):
          self.transformed_lengthscale = objax.TrainVar(np.array(softplus_inv(lengthscale)))
          if fix_variance:
              self.transformed_variance = objax.StateVar(np.array(softplus_inv(variance)))
          else:
              self.transformed_variance = objax.TrainVar(np.array(softplus_inv(variance)))
          self.transformed_radial_frequency = objax.TrainVar(np.array(softplus_inv(radial_frequency)))
          super().__init__()
          self.name = 'Subband Matern-1/2'
  
      @property
      def variance(self):
          return softplus(self.transformed_variance.value)
  
      @property
      def lengthscale(self):
          return softplus(self.transformed_lengthscale.value)
  
      @property
      def radial_frequency(self):
          return softplus(self.transformed_radial_frequency.value)
  
      def K(self, X, X2):
          raise NotImplementedError
  
      def kernel_to_state_space(self, R=None):
          F_mat = np.array([[-1.0 / self.lengthscale]])
          L_mat = np.array([[1.0]])
          Qc_mat = np.array([[2.0 * self.variance / self.lengthscale]])
          H_mat = np.array([[1.0]])
          Pinf_mat = np.array([[self.variance]])
          F_cos = np.array([[0.0, -self.radial_frequency],
                            [self.radial_frequency, 0.0]])
          H_cos = np.array([[1.0, 0.0]])
          # F = (-1/l -ω
          #      ω    -1/l)
          F = np.kron(F_mat, np.eye(2)) + F_cos
          L = np.kron(L_mat, np.eye(2))
          Qc = np.kron(np.eye(2), Qc_mat)
          H = np.kron(H_mat, H_cos)
          Pinf = np.kron(Pinf_mat, np.eye(2))
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          Pinf_mat = np.array([[self.variance]])
          Pinf = np.kron(Pinf_mat, np.eye(2))
          return Pinf
  
      def measurement_model(self):
          H_mat = np.array([[1.0]])
          H_cos = np.array([[1.0, 0.0]])
          H = np.kron(H_mat, H_cos)
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the closed form discrete-time state
          transition matrix A = expm(FΔt) for the Subband Matern-1/2 prior:
          A = exp(-Δt/l) ( cos(ωΔt)   -sin(ωΔt)
                           sin(ωΔt)    cos(ωΔt) )
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [1]
          :return: state transition matrix A [2, 2]
          """
          R = rotation_matrix(dt, self.radial_frequency)
          A = np.exp(-dt / self.lengthscale) * R  # [2, 2]
          return A
  
      def feedback_matrix(self):
          F_mat = np.array([[-1.0 / self.lengthscale]])
          F_cos = np.array([[0.0, -self.radial_frequency],
                            [self.radial_frequency, 0.0]])
          # F = (-1/l -ω
          #      ω    -1/l)
          F = np.kron(F_mat, np.eye(2)) + F_cos
          return F
  
  
  SubbandExponential = SubbandMatern12
  
  
  class SubbandMatern32(Kernel):
      """
      Subband Matern-3/2 kernel in SDE form (product of Cosine and Matern-3/2).
      Hyperparameters:
          variance, σ²
          lengthscale, l
          radial frequency, ω
      The associated continuous-time state space model matrices are constructed via
      kronecker sums and products of the Matern3/2 and cosine components:
      letting λ = √3 / l
      F      = F_mat3/2 ⊕ F_cos  =  ( 0     -ω     1     0
                                      ω      0     0     1
                                     -λ²     0    -2λ   -ω
                                      0     -λ²    ω    -2λ )
      L      = L_mat3/2 ⊗ I      =  ( 0      0
                                      0      0
                                      1      0
                                      0      1 )
      Qc     = I ⊗ Qc_mat3/2     =  ( 4λ³σ²  0
                                      0      4λ³σ² )
      H      = H_mat3/2 ⊗ H_cos  =  ( 1      0     0      0 )
      Pinf   = Pinf_mat3/2 ⊗ I   =  ( σ²     0     0      0
                                      0      σ²    0      0
                                      0      0     3σ²/l² 0
                                      0      0     0      3σ²/l²)
      and the discrete-time transition matrix is (for step size Δt),
      R = ( cos(ωΔt)   -sin(ωΔt)
            sin(ωΔt)    cos(ωΔt) )
      A = exp(-Δt/l) ( (1+Δtλ)R   ΔtR
                       -Δtλ²R    (1-Δtλ)R )
      """
      def __init__(self, variance=1.0, lengthscale=1.0, radial_frequency=1.0, fix_variance=False):
          self.transformed_lengthscale = objax.TrainVar(np.array(softplus_inv(lengthscale)))
          if fix_variance:
              self.transformed_variance = objax.StateVar(np.array(softplus_inv(variance)))
          else:
              self.transformed_variance = objax.TrainVar(np.array(softplus_inv(variance)))
          self.transformed_radial_frequency = objax.TrainVar(np.array(softplus_inv(radial_frequency)))
          super().__init__()
          self.name = 'Subband Matern-3/2'
  
      @property
      def variance(self):
          return softplus(self.transformed_variance.value)
  
      @property
      def lengthscale(self):
          return softplus(self.transformed_lengthscale.value)
  
      @property
      def radial_frequency(self):
          return softplus(self.transformed_radial_frequency.value)
  
      def K(self, X, X2):
          raise NotImplementedError
  
      def kernel_to_state_space(self, R=None):
          lam = 3.0 ** 0.5 / self.lengthscale
          F_mat = np.array([[0.0, 1.0],
                            [-lam ** 2, -2 * lam]])
          L_mat = np.array([[0],
                            [1]])
          Qc_mat = np.array([[12.0 * 3.0 ** 0.5 / self.lengthscale ** 3.0 * self.variance]])
          H_mat = np.array([[1.0, 0.0]])
          Pinf_mat = np.array([[self.variance, 0.0],
                               [0.0, 3.0 * self.variance / self.lengthscale ** 2.0]])
          F_cos = np.array([[0.0, -self.radial_frequency],
                            [self.radial_frequency, 0.0]])
          H_cos = np.array([[1.0, 0.0]])
          # F = (0   -ω   1   0
          #      ω    0   0   1
          #      -λ²  0  -2λ -ω
          #      0   -λ²  ω  -2λ)
          F = np.kron(F_mat, np.eye(2)) + np.kron(np.eye(2), F_cos)
          L = np.kron(L_mat, np.eye(2))
          Qc = np.kron(np.eye(2), Qc_mat)
          H = np.kron(H_mat, H_cos)
          Pinf = np.kron(Pinf_mat, np.eye(2))
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          Pinf_mat = np.array([[self.variance, 0.0],
                               [0.0, 3.0 * self.variance / self.lengthscale ** 2.0]])
          Pinf = np.kron(Pinf_mat, np.eye(2))
          return Pinf
  
      def measurement_model(self):
          H_mat = np.array([[1.0, 0.0]])
          H_cos = np.array([[1.0, 0.0]])
          H = np.kron(H_mat, H_cos)
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the closed form discrete-time state
          transition matrix A = expm(FΔt) for the Subband Matern-3/2 prior
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [1]
          :return: state transition matrix A [4, 4]
          """
          lam = np.sqrt(3.0) / self.lengthscale
          R = rotation_matrix(dt, self.radial_frequency)
          A = np.exp(-dt * lam) * np.block([
              [(1. + dt * lam) * R, dt * R],
              [-dt * lam ** 2 * R, (1. - dt * lam) * R]
          ])
          return A
  
      def feedback_matrix(self):
          lam = 3.0 ** 0.5 / self.lengthscale
          F_mat = np.array([[0.0, 1.0],
                            [-lam ** 2, -2 * lam]])
          F_cos = np.array([[0.0, -self.radial_frequency],
                            [self.radial_frequency, 0.0]])
          # F = (0   -ω   1   0
          #      ω    0   0   1
          #      -λ²  0  -2λ -ω
          #      0   -λ²  ω  -2λ)
          F = np.kron(F_mat, np.eye(2)) + np.kron(np.eye(2), F_cos)
          return F
  
  
  class SubbandMatern52(Kernel):
      """
      Subband Matern-5/2 kernel in SDE form (product of Cosine and Matern-5/2).
      Hyperparameters:
          variance, σ²
          lengthscale, l
          radial frequency, ω
      The associated continuous-time state space model matrices are constructed via
      kronecker sums and products of the Matern5/2 and cosine components:
      letting λ = √5/l
      F      = F_mat5/2 ⊕ F_cos  =  ( 0    -ω     1     0     0     0
                                      ω     0     0     1     0     0
                                      0     0     0    -ω     1     0
                                      0     0     ω     0     0     1
                                     -λ³    0    -3λ²   0    -3λ   -ω
                                      0    -λ³    0     3λ²   w    -3λ )
      L      = L_mat5/2 ⊗ I      =  ( 0     0
                                      0     0
                                      0     0
                                      0     0
                                      1     0
                                      0     1 )
      Qc     = I ⊗ Qc_mat5/2     =  ( 16λ⁵σ²/3  0
                                      0         16λ⁵σ²/3 )
      H      = H_mat5/2 ⊗ H_cos  =  ( 1     0     0      0     0    0 )
      letting κ = λ²σ²/3
      Pinf   = Pinf_mat5/2 ⊗ I   =  ( σ²    0     0      0    -κ     0
                                      0     σ²    0      0     0    -κ
                                      0     0     κ      0     0     0
                                      0     0     0      κ     0     0
                                     -κ     0     0      0     λ⁴σ²  0
                                      0    -κ     0      0     0     λ⁴σ² )
      and the discrete-time transition matrix is (for step size Δt),
      R = ( cos(ωΔt)   -sin(ωΔt)
            sin(ωΔt)    cos(ωΔt) )
      A = exp(-Δt/l) ( 1/2(2+Δtλ(2+Δtλ))R   Δt(1+Δtλ)R         1/2Δt²R
                      -1/2Δt²λ³R           (1+Δtλ(1-Δtλ))R    -1/2Δt(-2+Δtλ)R
                       1/2Δtλ³(-2+Δtλ)R     Δt²(-3+Δtλ)R       1/2(2+Δtλ(-4+Δtλ))R )
      """
      def __init__(self, variance=1.0, lengthscale=1.0, radial_frequency=1.0, fix_variance=False):
          self.transformed_lengthscale = objax.TrainVar(np.array(softplus_inv(lengthscale)))
          if fix_variance:
              self.transformed_variance = objax.StateVar(np.array(softplus_inv(variance)))
          else:
              self.transformed_variance = objax.TrainVar(np.array(softplus_inv(variance)))
          self.transformed_radial_frequency = objax.TrainVar(np.array(softplus_inv(radial_frequency)))
          super().__init__()
          self.name = 'Subband Matern-5/2'
  
      @property
      def variance(self):
          return softplus(self.transformed_variance.value)
  
      @property
      def lengthscale(self):
          return softplus(self.transformed_lengthscale.value)
  
      @property
      def radial_frequency(self):
          return softplus(self.transformed_radial_frequency.value)
  
      def K(self, X, X2):
          raise NotImplementedError
  
      def kernel_to_state_space(self, R=None):
          lam = 5.0 ** 0.5 / self.lengthscale
          F_mat = np.array([[0.0, 1.0, 0.0],
                            [0.0, 0.0, 1.0],
                            [-lam ** 3.0, -3.0 * lam ** 2.0, -3.0 * lam]])
          L_mat = np.array([[0.0],
                            [0.0],
                            [1.0]])
          Qc_mat = np.array([[self.variance * 400.0 * 5.0 ** 0.5 / 3.0 / self.lengthscale ** 5.0]])
          H_mat = np.array([[1.0, 0.0, 0.0]])
          kappa = 5.0 / 3.0 * self.variance / self.lengthscale ** 2.0
          Pinf_mat = np.array([[self.variance, 0.0, -kappa],
                               [0.0, kappa, 0.0],
                               [-kappa, 0.0, 25.0 * self.variance / self.lengthscale ** 4.0]])
          F_cos = np.array([[0.0, -self.radial_frequency],
                            [self.radial_frequency, 0.0]])
          H_cos = np.array([[1.0, 0.0]])
          # F = (0   -ω   1    0    0   0
          #      ω    0   0    1    0   0
          #      0    0   0   -ω    1   0
          #      0    0   ω    0    0   1
          #      -λ³  0  -3λ²  0   -3λ -ω
          #      0   -λ³  0   -3λ²  ω  -3λ )
          F = np.kron(F_mat, np.eye(2)) + np.kron(np.eye(3), F_cos)
          L = np.kron(L_mat, np.eye(2))
          Qc = np.kron(np.eye(2), Qc_mat)
          H = np.kron(H_mat, H_cos)
          Pinf = np.kron(Pinf_mat, np.eye(2))
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          kappa = 5.0 / 3.0 * self.variance / self.lengthscale ** 2.0
          Pinf_mat = np.array([[self.variance, 0.0, -kappa],
                               [0.0, kappa, 0.0],
                               [-kappa, 0.0, 25.0 * self.variance / self.lengthscale ** 4.0]])
          Pinf = np.kron(Pinf_mat, np.eye(2))
          return Pinf
  
      def measurement_model(self):
          H_mat = np.array([[1.0, 0.0, 0.0]])
          H_cos = np.array([[1.0, 0.0]])
          H = np.kron(H_mat, H_cos)
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the closed form discrete-time state
          transition matrix A = expm(FΔt) for the Subband Matern-5/2 prior
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [1]
          :return: state transition matrix A [6, 6]
          """
          lam = 5.0 ** 0.5 / self.lengthscale
          R = rotation_matrix(dt, self.radial_frequency)
          A = np.exp(-dt * lam) * np.block([
              [0.5*(2. + dt*lam*(2. + dt*lam)) * R, dt * (1. + dt * lam) * R,        0.5 * dt**2 * R],
              [-0.5*dt ** 2 * lam**3 * R,           (1. + dt*lam*(1. - dt*lam)) * R, -0.5 * dt * (-2. + dt * lam) * R],
              [0.5*dt*lam**3 * (-2. + dt*lam) * R,  dt * lam**2*(-3. + dt*lam) * R,  0.5*(2. + dt*lam*(-4. + dt*lam)) * R]
          ])
          return A
  
      def feedback_matrix(self):
          lam = 5.0 ** 0.5 / self.lengthscale
          F_mat = np.array([[0.0, 1.0, 0.0],
                            [0.0, 0.0, 1.0],
                            [-lam ** 3.0, -3.0 * lam ** 2.0, -3.0 * lam]])
          F_cos = np.array([[0.0, -self.radial_frequency],
                            [self.radial_frequency, 0.0]])
          # F = (0   -ω   1    0    0   0
          #      ω    0   0    1    0   0
          #      0    0   0   -ω    1   0
          #      0    0   ω    0    0   1
          #      -λ³  0  -3λ²  0   -3λ -ω
          #      0   -λ³  0   -3λ²  ω  -3λ )
          F = np.kron(F_mat, np.eye(2)) + np.kron(np.eye(3), F_cos)
          return F
  
  
  class Periodic(Kernel):
      """
      Periodic kernel in SDE form.
      Hyperparameters:
          variance, σ²
          lengthscale, l
          period, p
      The associated continuous-time state space model matrices are constructed via
      a sum of cosines.
      TODO: allow for orders other than 6
      """
      def __init__(self, variance=1.0, lengthscale=1.0, period=1.0, order=6, fix_variance=False):
          self.transformed_lengthscale = objax.TrainVar(np.array(softplus_inv(lengthscale)))
          if fix_variance:
              self.transformed_variance = objax.StateVar(np.array(softplus_inv(variance)))
          else:
              self.transformed_variance = objax.TrainVar(np.array(softplus_inv(variance)))
          self.transformed_period = objax.TrainVar(np.array(softplus_inv(period)))
          super().__init__()
          self.name = 'Periodic'
          self.order = order
          self.M = np.meshgrid(np.arange(self.order + 1), np.arange(self.order + 1))[1]
          factorial_mesh_M = np.array([[1., 1., 1., 1., 1., 1., 1.],
                                       [1., 1., 1., 1., 1., 1., 1.],
                                       [2., 2., 2., 2., 2., 2., 2.],
                                       [6., 6., 6., 6., 6., 6., 6.],
                                       [24., 24., 24., 24., 24., 24., 24.],
                                       [120., 120., 120., 120., 120., 120., 120.],
                                       [720., 720., 720., 720., 720., 720., 720.]])
          b = np.array([[1., 0., 0., 0., 0., 0., 0.],
                        [0., 2., 0., 0., 0., 0., 0.],
                        [2., 0., 2., 0., 0., 0., 0.],
                        [0., 6., 0., 2., 0., 0., 0.],
                        [6., 0., 8., 0., 2., 0., 0.],
                        [0., 20., 0., 10., 0., 2., 0.],
                        [20., 0., 30., 0., 12., 0., 2.]])
          self.b_fmK_2M = b * (1. / factorial_mesh_M) * (2. ** -self.M)
  
      @property
      def variance(self):
          return softplus(self.transformed_variance.value)
  
      @property
      def lengthscale(self):
          return softplus(self.transformed_lengthscale.value)
  
      @property
      def period(self):
          return softplus(self.transformed_period.value)
  
      def kernel_to_state_space(self, R=None):
          a = self.b_fmK_2M * self.lengthscale ** (-2. * self.M) * np.exp(-1. / self.lengthscale ** 2.) * self.variance
          q2 = np.sum(a, axis=0)
          # The angular frequency
          omega = 2 * np.pi / self.period
          # The model
          F = np.kron(np.diag(np.arange(self.order + 1)), np.array([[0., -omega], [omega, 0.]]))
          L = np.eye(2 * (self.order + 1))
          Qc = np.zeros(2 * (self.order + 1))
          Pinf = np.kron(np.diag(q2), np.eye(2))
          H = np.kron(np.ones([1, self.order + 1]), np.array([1., 0.]))
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          a = self.b_fmK_2M * self.lengthscale ** (-2. * self.M) * np.exp(-1. / self.lengthscale ** 2.) * self.variance
          q2 = np.sum(a, axis=0)
          Pinf = np.kron(np.diag(q2), np.eye(2))
          return Pinf
  
      def measurement_model(self):
          H = np.kron(np.ones([1, self.order + 1]), np.array([1., 0.]))
          return H
  
      def state_transition(self, dt):
          """
          Calculation of the closed form discrete-time state
          transition matrix A = expm(FΔt) for the Periodic prior
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [1]
          :return: state transition matrix A [2(N+1), 2(N+1)]
          """
          omega = 2 * np.pi / self.period  # The angular frequency
          harmonics = np.arange(self.order + 1) * omega
          R0 = rotation_matrix(dt, harmonics[0])
          R1 = rotation_matrix(dt, harmonics[1])
          R2 = rotation_matrix(dt, harmonics[2])
          R3 = rotation_matrix(dt, harmonics[3])
          R4 = rotation_matrix(dt, harmonics[4])
          R5 = rotation_matrix(dt, harmonics[5])
          R6 = rotation_matrix(dt, harmonics[6])
          A = np.block([
              [R0, np.zeros([2, 12])],
              [np.zeros([2, 2]),  R1, np.zeros([2, 10])],
              [np.zeros([2, 4]),  R2, np.zeros([2, 8])],
              [np.zeros([2, 6]),  R3, np.zeros([2, 6])],
              [np.zeros([2, 8]),  R4, np.zeros([2, 4])],
              [np.zeros([2, 10]), R5, np.zeros([2, 2])],
              [np.zeros([2, 12]), R6]
          ])
          return A
  
      def feedback_matrix(self):
          # The angular frequency
          omega = 2 * np.pi / self.period
          # The model
          F = np.kron(np.diag(np.arange(self.order + 1)), np.array([[0., -omega], [omega, 0.]]))
          return F
  
  
  class Independent(Kernel):
      """
      A stack of independent GP priors. 'kernels' is a list of GP kernels, and this class stacks
      the state space models such that each component is fed to the likelihood.
      This class differs from Sum only in the measurement model.
      """
      def __init__(self, kernels):
          self.num_kernels = len(kernels)
          for i in range(self.num_kernels):
              selfdotkerneli = "self.kernel" + str(i)
              exec(selfdotkerneli + " = kernels[i]")
          self.name = 'Independent'
  
      def K(self, X, X2):
          zeros = np.zeros(self.num_kernels)
          K0 = self.kernel0.K(X, X2)
          index_vector = index_add(zeros, index[0], 1.)
          Kstack = np.kron(K0, np.diag(index_vector))
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              index_vector = index_add(zeros, index[i], 1.)
              Kstack += np.kron(kerneli.K(X, X2), np.diag(index_vector))
          return Kstack
  
      def kernel_to_state_space(self, R=None):
          F, L, Qc, H, Pinf = self.kernel0.kernel_to_state_space(R)
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              F_, L_, Qc_, H_, Pinf_ = kerneli.kernel_to_state_space(R)
              F = block_diag(F, F_)
              L = block_diag(L, L_)
              Qc = block_diag(Qc, Qc_)
              H = block_diag(H, H_)
              Pinf = block_diag(Pinf, Pinf_)
          return F, L, Qc, H, Pinf
  
      def measurement_model(self):
          H = self.kernel0.measurement_model()
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              H_ = kerneli.measurement_model()
              H = block_diag(H, H_)
          return H
  
      def stationary_covariance(self):
          Pinf = self.kernel0.stationary_covariance()
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              Pinf_ = kerneli.stationary_covariance()
              Pinf = block_diag(Pinf, Pinf_)
          return Pinf
  
      def stationary_covariance_meanfield(self):
          """
          Stationary covariance as a tensor of blocks, as required when using a mean-field assumption
          """
          raise NotImplementedError('blocks are not all the same shape')
          # Pinf = self.kernel0.stationary_covariance()[None]
          # for i in range(1, self.num_kernels):
          #     kerneli = eval("self.kernel" + str(i))
          #     Pinf_ = kerneli.stationary_covariance()[None]
          #     Pinf = np.concatenate([Pinf, Pinf_])
          # return Pinf
  
      def inducing_precision(self):
          Qzz0, Lzz0 = self.kernel0.inducing_precision()
          Qzz, Lzz = [Qzz0], [Lzz0]
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              Qzz_, Lzz_ = kerneli.inducing_precision()
              Qzz, Lzz = Qzz + [Qzz_], Lzz + [Lzz_]
          return Qzz, Lzz
  
      def state_transition(self, dt):
          """
          Calculation of the discrete-time state transition matrix A = expm(FΔt) for a set of stacked of GPs
          :param dt: step size(s), Δt = tₙ - tₙ₋₁ [1]
          :return: state transition matrix A [D, D]
          """
          A = self.kernel0.state_transition(dt)
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              A_ = kerneli.state_transition(dt)
              A = block_diag(A, A_)
          return A
  
      def state_transition_meanfield(self, dt):
          """
          State transition matrix in the form required for mean-field inference.
          :param dt: step size(s), Δtₙ = tₙ - tₙ₋₁ [scalar]
          :return: state transition matrix A
          """
          raise NotImplementedError('blocks are not all the same shape')
          # A = self.kernel0.state_transition(dt)[None]
          # for i in range(1, self.num_kernels):
          #     kerneli = eval("self.kernel" + str(i))
          #     A_ = kerneli.state_transition(dt)[None]
          #     A = np.concatenate([A, A_])
          # return A
  
      def get_meanfield_block_index(self):
          raise NotImplementedError('blocks are not all the same shape')
          # Pinf = self.stationary_covariance_meanfield()
          # num_latents = Pinf.shape[0]
          # sub_state_dim = Pinf.shape[1]
          # state = np.ones([sub_state_dim, sub_state_dim])
          # for i in range(1, num_latents):
          #     state = block_diag(state, np.ones([sub_state_dim, sub_state_dim]))
          # block_index = np.where(np.array(state, dtype=bool))
          # return block_index
  
      def feedback_matrix(self):
          F = self.kernel0.feedback_matrix()
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              F_ = kerneli.feedback_matrix()
              F = block_diag(F, F_)
          return F
  
  
  Separate = Independent
  
  Stack = Independent
  
  
  class Sum(Independent):
      """
      A sum of GP priors. 'components' is a list of GP kernels, and this class stacks
      the state space models to produce their sum.
      This class differs from Independent only in the measurement model.
      """
      def __init__(self, kernels):
          super().__init__(kernels=kernels)
          self.name = 'Sum'
  
      def K(self, X, X2):
          Ksum = self.kernel0.K(X, X2)
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              Ksum = Ksum + kerneli.K(X, X2)
          return Ksum
  
      def kernel_to_state_space(self, R=None):
          F, L, Qc, H, Pinf = self.kernel0.kernel_to_state_space(R)
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              F_, L_, Qc_, H_, Pinf_ = kerneli.kernel_to_state_space(R)
              F = block_diag(F, F_)
              L = block_diag(L, L_)
              Qc = block_diag(Qc, Qc_)
              H = np.block([
                  H, H_
              ])
              Pinf = block_diag(Pinf, Pinf_)
          return F, L, Qc, H, Pinf
  
      def measurement_model(self):
          H = self.kernel0.measurement_model()
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              H_ = kerneli.measurement_model()
              H = np.block([
                  H, H_
              ])
          return H
  
  
  class Separable(Independent):
      """
      A product of separable GP priors. 'components' is a list of GP kernels, and this class stacks
      the state space models to produce their product.
      This class differs from Independent only in the measurement model.
      TODO: this assumes that each kernel acts on a different dimension. Generalise.
      TODO: implement state space form of product kernels
      """
      def __init__(self, kernels):
          super().__init__(kernels=kernels)
          self.name = 'Product'
  
      def K(self, X, X2):
          Kprod = self.kernel0.K(X[:, :1], X2[:, :1])
          for i in range(1, self.num_kernels):
              kerneli = eval("self.kernel" + str(i))
              Kprod = Kprod * kerneli.K(X[:, i:i+1], X2[:, i:i+1])
          return Kprod
  
      # def measurement_model(self):
      #     H = self.kernel0.measurement_model()
      #     for i in range(1, self.num_kernels):
      #         kerneli = eval("self.kernel" + str(i))
      #         H_ = kerneli.measurement_model()
      #         H = np.block([
      #             H, H_
      #         ])
      #     return H
  
  
  class SpectroTemporal(Independent):
  
      def __init__(self,
                   subband_lengthscales,
                   subband_frequencies,
                   modulator_variances,
                   modulator_lengthscales,
                   subband_kernel=SubbandMatern12,
                   modulator_kernel=Matern32):
          assert len(subband_lengthscales) == len(subband_frequencies)
          assert len(modulator_lengthscales) == len(modulator_variances)
          num_subbands = len(subband_frequencies)
          num_modulators = len(modulator_lengthscales)
          radial_freq = 2 * np.pi * subband_frequencies  # radial freq = 2pi * f
          kernels = [subband_kernel(variance=1, lengthscale=subband_lengthscales[0], radial_frequency=radial_freq[0],
                                    fix_variance=True)]
          for i in range(1, num_subbands):
              kernels.append(
                  subband_kernel(variance=1, lengthscale=subband_lengthscales[i], radial_frequency=radial_freq[i],
                                 fix_variance=True)
              )
          for j in range(num_modulators):
              kernels.append(
                  modulator_kernel(variance=modulator_variances[j], lengthscale=modulator_lengthscales[j])
              )
          super().__init__(kernels=kernels)
  
  
  class LatentExponentiallyGenerated(Kernel):
  
      def __init__(self,
                   N,
                   R):
          assert N.shape == R.shape
          self.N_ = objax.StateVar(np.array(N))
          self.R_ = objax.StateVar(np.array(R))
  
      @property
      def N(self):
          return self.N_.value
  
      @property
      def R(self):
          return self.R_.value
  
      def kernel_to_state_space(self, R=None):
          F = -0.5 * (self.N @ self.N.T + self.R - self.R.T)
          L = self.N
          Qc = np.array([[1.]])
          H = np.concatenate(
              [np.array([[1.0]]), np.zeros([1, self.N.shape[0] - 1])],
              axis=1
          )
          Pinf = np.eye(self.N.shape[0])
          return F, L, Qc, H, Pinf
  
      def stationary_covariance(self):
          Pinf = np.eye(self.N.shape[0])
          return Pinf
  
      def measurement_model(self):
          H = np.concatenate(
              [np.array([[1.0]]), np.zeros([1, self.N.shape[0] - 1])],
              axis=1
          )
          return H
  
      def feedback_matrix(self):
          F = -0.5 * (self.N @ self.N.T + self.R - self.R.T)
          return F
  
  
  LEG = LatentExponentiallyGenerated
  