text: import logging
  import math
  import multiprocessing as mp
  import numpy as np
  import time
  
  from typing import List, Callable
  from flare.kernels.utils import from_mask_to_args, from_grad_to_mask
  
  _global_training_data = {}
  _global_training_labels = {}
  _global_training_structures = {}
  _global_energy_labels = {}
  
  
  def queue_wrapper(result_queue, wid, func, args):
      """
      wrapper function for multiprocessing queue
      """
      result_queue.put((wid, func(*args)))
  
  
  def partition_matrix(n_sample, size, n_cpus):
      """
      partition the training data for matrix calculation
      the number of blocks are close to n_cpus
      since mp.Process does not allow to change the thread number
      """
  
      # divide the block by n_cpu partitions, with size n_sample0
      # if the divided chunk is smaller than the requested chunk n_sample
      # use the requested chunk size
      n_unique_elements = size * (size + 1) / 2
      n_sample0 = int(math.ceil(np.sqrt(n_unique_elements / n_cpus)))
      if n_sample0 > n_sample:
          n_sample = n_sample0
  
      block_id = []
      nbatch = 0
      nbatch1 = 0
      nbatch2 = 0
      e1 = 0
      while e1 < size:
          s1 = int(n_sample * nbatch1)
          e1 = int(np.min([s1 + n_sample, size]))
          nbatch2 = nbatch1
          nbatch1 += 1
          e2 = 0
          while e2 < size:
              s2 = int(n_sample * nbatch2)
              e2 = int(np.min([s2 + n_sample, size]))
              block_id += [(s1, e1, s2, e2)]
              nbatch2 += 1
              nbatch += 1
  
      return block_id, nbatch
  
  
  def partition_matrix_custom(n_sample: int, start1, end1, start2, end2, n_cpus):
      """
      Partition a specified portion of a matrix.
      """
  
      n_rows = end1 - start1
      n_cols = end2 - start2
      n_unique_elements = n_rows * n_cols
      n_sample0 = int(math.ceil(np.sqrt(n_unique_elements / n_cpus)))
      if n_sample0 > n_sample:
          n_sample = n_sample0
  
      block_id = []
      n_block = 0  # total number of blocks
      nbatch1 = 0
      nbatch2 = 0
      e1 = 0
      while e1 < end1:
          s1 = start1 + n_sample * nbatch1
          e1 = np.min([s1 + n_sample, end1])
          nbatch1 += 1
  
          e2 = 0
          nbatch2 = 0
          while e2 < end2:
              s2 = start2 + n_sample * nbatch2
              e2 = np.min([s2 + n_sample, end2])
              block_id += [(s1, e1, s2, e2)]
              nbatch2 += 1
              n_block += 1  # update block count
  
      return block_id, n_block
  
  
  def partition_vector(n_sample, size, n_cpus):
      """
      partition the training data for vector calculation
      the number of blocks are the same as n_cpus
      since mp.Process does not allow to change the thread number
      """
      n_sample0 = int(math.ceil(size / n_cpus))
      if n_sample0 > n_sample:
          n_sample = n_sample0
  
      block_id = []
      nbatch = 0
      e = 0
      while e < size:
          s = n_sample * nbatch
          e = np.min([s + n_sample, size])
          block_id += [(s, e)]
          nbatch += 1
      return block_id, nbatch
  
  
  def partition_force_energy_block(n_sample: int, size1: int, size2: int, n_cpus: int):
      """Special partition method for the force/energy block. Because the number
      of environments in a structure can vary, we only split up the environment
      list, which has length size1.
  
      Note that two sizes need to be specified: the size of the envionment
      list and the size of the structure list.
  
      Args:
          n_sample (int): Number of environments per processor.
          size1 (int): Size of the environment list.
          size2 (int): Size of the structure list.
          n_cpus (int): Number of cpus.
      """
  
      n_sample0 = int(math.ceil(size1 / n_cpus))
      if n_sample0 > n_sample:
          n_sample = n_sample0
  
      block_id = []
      nbatch = 0
      e = 0
  
      while e < size1:
          s = n_sample * nbatch
          e = np.min([s + n_sample, size1])
          block_id += [(s, e, 0, size2)]
          nbatch += 1
  
      return block_id, nbatch
  
  
  def partition_update(n_sample, size, old_size, n_cpus):
  
      n_unique_elements = (size - old_size) * size
      n_sample0 = int(math.ceil(np.sqrt(n_unique_elements / n_cpus)))
      if n_sample0 > n_sample:
          n_sample = n_sample0
  
      ns_new = int(math.ceil((size - old_size) / n_sample))
      old_ns = int(math.ceil(old_size / n_sample))
  
      nbatch = 0
      block_id = []
      for ibatch1 in range(old_ns):
          s1 = int(n_sample * ibatch1)
          e1 = int(np.min([s1 + n_sample, old_size]))
  
          for ibatch2 in range(ns_new):
              s2 = int(n_sample * ibatch2) + old_size
              e2 = int(np.min([s2 + n_sample, size]))
              block_id += [(s1, e1, s2, e2)]
              nbatch += 1
  
      for ibatch1 in range(ns_new):
          s1 = int(n_sample * ibatch1) + old_size
          e1 = int(np.min([s1 + n_sample, size]))
  
          for ibatch2 in range(ns_new):
              s2 = int(n_sample * ibatch2) + old_size
              e2 = int(np.min([s2 + n_sample, size]))
              block_id += [(s1, e1, s2, e2)]
              nbatch += 1
  
      return block_id, nbatch
  
  
  def obtain_noise_len(hyps, hyps_mask):
      """
      obtain the noise parameter from hyps and mask
      """
  
      # assume sigma_n is the final hyperparameter
      sigma_n = hyps[-1]
      # correct it if map is defined
      non_noise_hyps = len(hyps) - 1
      train_noise = True
      if hyps_mask is not None:
          train_noise = hyps_mask.get("train_noise", True)
          if train_noise is False:
              sigma_n = hyps_mask["original_hyps"][-1]
              non_noise_hyps = len(hyps)
  
      return sigma_n, non_noise_hyps, train_noise
  
  
  # --------------------------------------------------------------------------
  #                   Parallel matrix/vector construction
  # --------------------------------------------------------------------------
  
  
  def parallel_matrix_construction(
      pack_function,
      hyps,
      name,
      kernel,
      cutoffs,
      hyps_mask,
      block_id,
      nbatch,
      size1,
      size2,
      m1,
      m2,
      symm=True,
  ):
      result_queue = mp.Queue()
      children = []
      for wid in range(nbatch):
          s1, e1, s2, e2 = block_id[wid]
          children.append(
              mp.Process(
                  target=queue_wrapper,
                  args=(
                      result_queue,
                      wid,
                      pack_function,
                      (hyps, name, s1, e1, s2, e2, s1 == s2, kernel, cutoffs, hyps_mask),
                  ),
              )
          )
  
      # Run child processes.
      for c in children:
          c.start()
  
      # Wait for all results to arrive.
      matrix = np.zeros((size1, size2))
      for _ in range(nbatch):
          wid, result_chunk = result_queue.get(block=True)
          s1, e1, s2, e2 = block_id[wid]
          matrix[s1 * m1 : e1 * m1, s2 * m2 : e2 * m2] = result_chunk
          # Note that the force/energy block is not symmetric; in this case
          # symm is False.
          if (s1 != s2) and (symm is True):
              matrix[s2 * m2 : e2 * m2, s1 * m1 : e1 * m1] = result_chunk.T
  
      # Join child processes (clean up zombies).
      for c in children:
          c.join()
  
      # matrix manipulation
      del result_queue
      del children
  
      return matrix
  
  
  def parallel_vector_construction(
      pack_function,
      name,
      x,
      kernel,
      hyps,
      cutoffs,
      hyps_mask,
      block_id,
      nbatch,
      size,
      mult,
      d_1=None,
  ):
      result_queue = mp.Queue()
      children = []
      for wid in range(nbatch):
          s, e = block_id[wid]
          children.append(
              mp.Process(
                  target=queue_wrapper,
                  args=(
                      result_queue,
                      wid,
                      pack_function,
                      (name, s, e, x, kernel, hyps, cutoffs, hyps_mask, d_1),
                  ),
              )
          )
  
      # Run child processes.
      for c in children:
          c.start()
  
      # Wait for all results to arrive.
      vector = np.zeros(size * mult)
      for _ in range(nbatch):
          wid, result_chunk = result_queue.get(block=True)
          s, e = block_id[wid]
          vector[s * mult : e * mult] = result_chunk
  
      # Join child processes (clean up zombies).
      for c in children:
          c.join()
  
      return vector
  
  
  def multiple_array_construction(
      pack_function,
      name,
      x,
      kernel,
      hyps,
      cutoffs,
      hyps_mask,
      block_id,
      nbatch,
      size,
      mult,
      array_sizes,
  ):
      result_queue = mp.Queue()
      children = []
      for wid in range(nbatch):
          s, e = block_id[wid]
          children.append(
              mp.Process(
                  target=queue_wrapper,
                  args=(
                      result_queue,
                      wid,
                      pack_function,
                      (name, s, e, x, kernel, hyps, cutoffs, hyps_mask),
                  ),
              )
          )
  
      # Run child processes.
      for c in children:
          c.start()
  
      # Initialize arrays.
      n_arrays = len(array_sizes)
      arrays = []
      for n in range(n_arrays):
          arrays.append(np.zeros(array_sizes[n]))
  
      # Wait for all results to arrive.
      for _ in range(nbatch):
          wid, result_chunk = result_queue.get(block=True)
          s, e = block_id[wid]
          for n in range(n_arrays):
              arrays[n][:, s * mult : e * mult] = result_chunk[n]
  
      # Join child processes (clean up zombies).
      for c in children:
          c.join()
  
      return arrays
  
  
  # --------------------------------------------------------------------------
  #                            Ky construction
  # --------------------------------------------------------------------------
  
  
  def get_distance_mat_pack(
      hyps: np.ndarray,
      name: str,
      s1: int,
      e1: int,
      s2: int,
      e2: int,
      same: bool,
      kernel,
      cutoffs,
      hyps_mask,
  ):
      """Compute covariance matrix element between set1 and set2
      :param hyps: list of hyper-parameters
      :param name: name of the gp instance.
      :param same: whether the row and column are the same
      :param kernel: function object of the kernel
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
      :return: covariance matrix
      """
  
      # initialize matrices
      training_data = _global_training_data[name]
      size1 = e1 - s1
      size2 = e2 - s2
      dist_mat = np.zeros([size1, size2])
  
      # calculate elements
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      for m_index in range(size1):
          x_1 = training_data[m_index + s1]
          if same:
              lowbound = m_index
          else:
              lowbound = 0
          for n_index in range(lowbound, size2):
              x_2 = training_data[n_index + s2]
              kern_curr = kernel(x_1, x_2, *args)
              # store kernel value
              dist_mat[m_index, n_index] = kern_curr
              if same:
                  dist_mat[n_index, m_index] = kern_curr
  
      return dist_mat
  
  
  def kernel_distance_mat(
      hyps: np.ndarray,
      name: str,
      kernel,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      """parallel version of get_ky_mat
      :param hyps: list of hyper-parameters
      :param name: name of the gp instance.
      :param kernel: function object of the kernel
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
      :return: covariance matrix
      """
  
      training_data = _global_training_data[name]
      size = len(training_data)
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          k_mat = get_distance_mat_pack(
              hyps, name, 0, size, 0, size, True, kernel, cutoffs, hyps_mask
          )
      else:
          # initialize matrices
          block_id, nbatch = partition_matrix(n_sample, size, n_cpus)
          mult = 1
  
          k_mat = parallel_matrix_construction(
              get_distance_mat_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id,
              nbatch,
              size,
              size,
              mult,
              mult,
          )
  
      return k_mat
  
  
  def get_force_block_pack(
      hyps: np.ndarray,
      name: str,
      s1: int,
      e1: int,
      s2: int,
      e2: int,
      same: bool,
      kernel,
      cutoffs,
      hyps_mask,
  ):
      """Compute covariance matrix element between set1 and set2
      :param hyps: list of hyper-parameters
      :param name: name of the gp instance.
      :param same: whether the row and column are the same
      :param kernel: function object of the kernel
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
  
      :return: covariance matrix
      """
  
      # initialize matrices
      training_data = _global_training_data[name]
      size1 = (e1 - s1) * 3
      size2 = (e2 - s2) * 3
      force_block = np.zeros([size1, size2])
  
      ds = [1, 2, 3]
  
      # calculate elements
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      for m_index in range(size1):
          x_1 = training_data[int(math.floor(m_index / 3)) + s1]
          d_1 = ds[m_index % 3]
          if same:
              lowbound = m_index
          else:
              lowbound = 0
          for n_index in range(lowbound, size2):
              x_2 = training_data[int(math.floor(n_index / 3)) + s2]
              d_2 = ds[n_index % 3]
              kern_curr = kernel(x_1, x_2, d_1, d_2, *args)
              # store kernel value
              force_block[m_index, n_index] = kern_curr
              if same:
                  force_block[n_index, m_index] = kern_curr
  
      return force_block
  
  
  def get_energy_block_pack(
      hyps: np.ndarray,
      name: str,
      s1: int,
      e1: int,
      s2: int,
      e2: int,
      same: bool,
      kernel,
      cutoffs,
      hyps_mask,
  ):
  
      # initialize matrices
      training_structures = _global_training_structures[name]
      size1 = e1 - s1
      size2 = e2 - s2
      energy_block = np.zeros([size1, size2])
  
      # calculate elements
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      for m_index in range(size1):
          struc_1 = training_structures[m_index + s1]
          if same:
              lowbound = m_index
          else:
              lowbound = 0
  
          for n_index in range(lowbound, size2):
              struc_2 = training_structures[n_index + s2]
  
              # Loop over environments in both structures to compute the
              # energy/energy kernel.
              kern_curr = 0
              for environment_1 in struc_1:
                  for environment_2 in struc_2:
                      kern_curr += kernel(environment_1, environment_2, *args)
  
              # Store kernel value.
              energy_block[m_index, n_index] = kern_curr
              if same:
                  energy_block[n_index, m_index] = kern_curr
  
      return energy_block
  
  
  def get_force_energy_block_pack(
      hyps: np.ndarray,
      name: str,
      s1: int,
      e1: int,
      s2: int,
      e2: int,
      same: bool,
      kernel,
      cutoffs,
      hyps_mask,
  ):
      # initialize matrices
      training_data = _global_training_data[name]
      training_structures = _global_training_structures[name]
      size1 = (e1 - s1) * 3
      size2 = e2 - s2
      force_energy_block = np.zeros([size1, size2])
  
      ds = [1, 2, 3]
  
      # calculate elements
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      for m_index in range(size1):
          environment_1 = training_data[int(math.floor(m_index / 3)) + s1]
          d_1 = ds[m_index % 3]
  
          for n_index in range(size2):
              structure = training_structures[n_index + s2]
  
              # Loop over environments in the training structure.
              kern_curr = 0
              for environment_2 in structure:
                  kern_curr += kernel(environment_1, environment_2, d_1, *args)
  
              # store kernel value
              force_energy_block[m_index, n_index] = kern_curr
  
      return force_energy_block
  
  
  def get_force_block(
      hyps: np.ndarray,
      name: str,
      kernel,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      """parallel version of get_ky_mat
      :param hyps: list of hyper-parameters
      :param name: name of the gp instance.
      :param kernel: function object of the kernel
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
  
      :return: covariance matrix
      """
  
      training_data = _global_training_data[name]
      size = len(training_data)
      size3 = 3 * size
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          k_mat = get_force_block_pack(
              hyps, name, 0, size, 0, size, True, kernel, cutoffs, hyps_mask
          )
      else:
          # initialize matrices
          block_id, nbatch = partition_matrix(n_sample, size, n_cpus)
          mult = 3
  
          k_mat = parallel_matrix_construction(
              get_force_block_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id,
              nbatch,
              size3,
              size3,
              mult,
              mult,
          )
  
      sigma_n, _, __ = obtain_noise_len(hyps, hyps_mask)
      force_block = k_mat
      force_block += sigma_n ** 2 * np.eye(size3)
  
      return force_block
  
  
  def get_energy_block(
      hyps: np.ndarray,
      name: str,
      kernel,
      energy_noise,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      training_structures = _global_training_structures[name]
      size = len(training_structures)
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          k_mat = get_energy_block_pack(
              hyps, name, 0, size, 0, size, True, kernel, cutoffs, hyps_mask
          )
      else:
          block_id, nbatch = partition_matrix(n_sample, size, n_cpus)
          mult = 1
  
          k_mat = parallel_matrix_construction(
              get_energy_block_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id,
              nbatch,
              size,
              size,
              mult,
              mult,
          )
  
      energy_block = k_mat
      energy_block += (energy_noise ** 2) * np.eye(size)
  
      return energy_block
  
  
  def get_force_energy_block(
      hyps: np.ndarray,
      name: str,
      kernel,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      training_data = _global_training_data[name]
      training_structures = _global_training_structures[name]
      size1 = len(training_data) * 3
      size2 = len(training_structures)
      size3 = len(training_data)
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          force_energy_block = get_force_energy_block_pack(
              hyps, name, 0, size3, 0, size2, True, kernel, cutoffs, hyps_mask
          )
      else:
          # initialize matrices
          block_id, nbatch = partition_force_energy_block(n_sample, size3, size2, n_cpus)
          m1 = 3  # 3 force components per environment
          m2 = 1  # 1 energy per structure
  
          force_energy_block = parallel_matrix_construction(
              get_force_energy_block_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id,
              nbatch,
              size1,
              size2,
              m1,
              m2,
              symm=False,
          )
  
      return force_energy_block
  
  
  def get_Ky_mat(
      hyps: np.ndarray,
      name: str,
      force_kernel: Callable,
      energy_kernel: Callable,
      force_energy_kernel: Callable,
      energy_noise,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      training_data = _global_training_data[name]
      training_structures = _global_training_structures[name]
      size1 = len(training_data) * 3
      size2 = len(training_structures)
  
      # Initialize Ky.
      ky_mat = np.zeros((size1 + size2, size1 + size2))
  
      # Assemble the full covariance matrix block-by-block.
      force_block = get_force_block(
          hyps, name, force_kernel, cutoffs, hyps_mask, n_cpus, n_sample
      )
  
      energy_block = get_energy_block(
          hyps, name, energy_kernel, energy_noise, cutoffs, hyps_mask, n_cpus, n_sample
      )
  
      force_energy_block = get_force_energy_block(
          hyps, name, force_energy_kernel, cutoffs, hyps_mask, n_cpus, n_sample
      )
  
      ky_mat[0:size1, 0:size1] = force_block
      ky_mat[size1:, size1:] = energy_block
      ky_mat[0:size1, size1:] = force_energy_block
      ky_mat[size1:, 0:size1] = force_energy_block.transpose()
  
      return ky_mat
  
  
  # --------------------------------------------------------------------------
  #                              Ky updates
  # --------------------------------------------------------------------------
  
  
  def update_force_block(
      ky_mat_old: np.ndarray,
      n_envs_prev: int,
      hyps: np.ndarray,
      name: str,
      kernel,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      old_size = n_envs_prev
      old_size3 = n_envs_prev * 3
      mult = 3
      size = len(_global_training_data[name])
      size3 = 3 * size
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
  
      # serial version
      if n_cpus == 1:
          force_block = np.zeros((size3, size3))
          new_mat = get_force_block_pack(
              hyps, name, 0, size, old_size, size, False, kernel, cutoffs, hyps_mask
          )
          force_block[:, old_size3:] = new_mat
          force_block[old_size3:, :] = new_mat.transpose()
  
      # parallel version
      else:
          block_id, nbatch = partition_update(n_sample, size, old_size, n_cpus)
  
          force_block = parallel_matrix_construction(
              get_force_block_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id,
              nbatch,
              size3,
              size3,
              mult,
              mult,
          )
  
      # insert previous covariance matrix
      force_block[:old_size3, :old_size3] = ky_mat_old[:old_size3, :old_size3]
  
      # add the noise parameter
      sigma_n, _, _ = obtain_noise_len(hyps, hyps_mask)
      force_block[old_size3:, old_size3:] += sigma_n ** 2 * np.eye(size3 - old_size3)
  
      return force_block
  
  
  def update_energy_block(
      ky_mat_old: np.ndarray,
      n_envs_prev: int,
      hyps: np.ndarray,
      name: str,
      kernel,
      energy_noise: float,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      old_size = ky_mat_old.shape[0] - 3 * n_envs_prev
      mult = 1
      size = len(_global_training_structures[name])
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
  
      # serial version
      if n_cpus == 1:
          energy_block = np.zeros((size, size))
          new_mat = get_energy_block_pack(
              hyps, name, 0, size, old_size, size, False, kernel, cutoffs, hyps_mask
          )
          energy_block[:, old_size:] = new_mat
          energy_block[old_size:, :] = new_mat.transpose()
  
      # parallel version
      else:
          block_id, nbatch = partition_update(n_sample, size, old_size, n_cpus)
  
          energy_block = parallel_matrix_construction(
              get_energy_block_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id,
              nbatch,
              size,
              size,
              mult,
              mult,
          )
  
      # insert previous covariance matrix (if it has nonzero size)
      if old_size > 0:
          energy_block[:old_size, :old_size] = ky_mat_old[-old_size:, -old_size:]
  
      # add the noise parameter
      energy_block[old_size:, old_size:] += (energy_noise ** 2) * np.eye(size - old_size)
  
      return energy_block
  
  
  def update_force_energy_block(
      ky_mat_old: np.ndarray,
      n_envs_prev: int,
      hyps: np.ndarray,
      name: str,
      kernel,
      energy_noise: float,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      n_strucs_prev = ky_mat_old.shape[0] - 3 * n_envs_prev
      n_envs = len(_global_training_data[name])
      n_strucs = len(_global_training_structures[name])
      force_energy_block = np.zeros((n_envs * 3, n_strucs))
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
  
      # serial version
      if n_cpus == 1:
          mat_1 = get_force_energy_block_pack(
              hyps,
              name,
              0,
              n_envs,
              n_strucs_prev,
              n_strucs,
              False,
              kernel,
              cutoffs,
              hyps_mask,
          )
  
          mat_2 = get_force_energy_block_pack(
              hyps,
              name,
              n_envs_prev,
              n_envs,
              0,
              n_strucs_prev,
              False,
              kernel,
              cutoffs,
              hyps_mask,
          )
  
      # parallel version
      else:
          force_energy_block = np.zeros((n_envs * 3, n_strucs))
  
          block_id_1, nbatch_1 = partition_matrix_custom(
              n_sample, 0, n_envs, n_strucs_prev, n_strucs, n_cpus
          )
          block_id_2, nbatch_2 = partition_matrix_custom(
              n_sample, n_envs_prev, n_envs, 0, n_strucs_prev, n_cpus
          )
  
          size1 = n_envs * 3
          size2 = n_strucs
          m1 = 3
          m2 = 1
  
          mat_1 = parallel_matrix_construction(
              get_force_energy_block_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id_1,
              nbatch_1,
              size1,
              size2,
              m1,
              m2,
              False,
          )
          mat_2 = parallel_matrix_construction(
              get_force_energy_block_pack,
              hyps,
              name,
              kernel,
              cutoffs,
              hyps_mask,
              block_id_2,
              nbatch_2,
              size1,
              size2,
              m1,
              m2,
              False,
          )
  
          # reduce the size of the matrices
          mat_1 = mat_1[: n_envs * 3, n_strucs_prev:n_strucs]
          mat_2 = mat_2[n_envs_prev * 3 : n_envs * 3, :n_strucs_prev]
  
      force_energy_block[: n_envs * 3, n_strucs_prev:n_strucs] = mat_1
      force_energy_block[n_envs_prev * 3 : n_envs * 3, :n_strucs_prev] = mat_2
  
      # insert previous covariance matrix
      force_energy_block[: n_envs_prev * 3, :n_strucs_prev] = ky_mat_old[
          : n_envs_prev * 3, n_envs_prev * 3 : n_envs_prev * 3 + n_strucs_prev
      ]
  
      return force_energy_block
  
  
  def get_ky_mat_update(
      ky_mat_old: np.ndarray,
      n_envs_prev: int,
      hyps: np.ndarray,
      name: str,
      force_kernel: Callable,
      energy_kernel: Callable,
      force_energy_kernel: Callable,
      energy_noise: float,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      n_envs = len(_global_training_data[name])
      n_strucs = len(_global_training_structures[name])
      size1 = n_envs * 3
      size2 = n_strucs
      ky_mat = np.zeros((size1 + size2, size1 + size2))
  
      force_block = update_force_block(
          ky_mat_old,
          n_envs_prev,
          hyps,
          name,
          force_kernel,
          cutoffs,
          hyps_mask,
          n_cpus,
          n_sample,
      )
  
      energy_block = update_energy_block(
          ky_mat_old,
          n_envs_prev,
          hyps,
          name,
          energy_kernel,
          energy_noise,
          cutoffs,
          hyps_mask,
          n_cpus,
          n_sample,
      )
  
      force_energy_block = update_force_energy_block(
          ky_mat_old,
          n_envs_prev,
          hyps,
          name,
          force_energy_kernel,
          energy_noise,
          cutoffs,
          hyps_mask,
          n_cpus,
          n_sample,
      )
  
      ky_mat[0:size1, 0:size1] = force_block
      ky_mat[size1:, size1:] = energy_block
      ky_mat[0:size1, size1:] = force_energy_block
      ky_mat[size1:, 0:size1] = force_energy_block.transpose()
  
      return ky_mat
  
  
  # --------------------------------------------------------------------------
  #                            Kernel vectors
  # --------------------------------------------------------------------------
  
  
  def energy_energy_vector_unit(
      name, s, e, x, kernel, hyps, cutoffs=None, hyps_mask=None, d_1=None
  ):
      """
      Gets part of the energy/energy vector.
      """
  
      training_structures = _global_training_structures[name]
  
      size = e - s
      energy_energy_unit = np.zeros(
          size,
      )
  
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      for m_index in range(size):
          structure = training_structures[m_index + s]
          kern_curr = 0
          for environment in structure:
              kern_curr += kernel(x, environment, *args)
  
          energy_energy_unit[m_index] = kern_curr
  
      return energy_energy_unit
  
  
  def energy_force_vector_unit(
      name, s, e, x, kernel, hyps, cutoffs=None, hyps_mask=None, d_1=None
  ):
      """
      Gets part of the energy/force vector.
      """
  
      training_data = _global_training_data[name]
  
      ds = [1, 2, 3]
      size = (e - s) * 3
      k_v = np.zeros(
          size,
      )
  
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      for m_index in range(size):
          x_2 = training_data[int(math.floor(m_index / 3)) + s]
          d_2 = ds[m_index % 3]
          k_v[m_index] = kernel(x_2, x, d_2, *args)
  
      return k_v
  
  
  def force_energy_vector_unit(name, s, e, x, kernel, hyps, cutoffs, hyps_mask, d_1):
      """
      Gets part of the force/energy vector.
      """
  
      training_structures = _global_training_structures[name]
  
      size = e - s
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
      force_energy_unit = np.zeros(
          size,
      )
  
      for m_index in range(size):
          training_structure = training_structures[m_index + s]
          kern_curr = 0
          for environment in training_structure:
              kern_curr += kernel(x, environment, d_1, *args)
  
          force_energy_unit[m_index] = kern_curr
  
      return force_energy_unit
  
  
  def force_force_vector_unit(name, s, e, x, kernel, hyps, cutoffs, hyps_mask, d_1):
      """
      Gets part of the force/force vector.
      """
  
      training_data = _global_training_data[name]
  
      size = e - s
      ds = [1, 2, 3]
  
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      k_v = np.zeros(size * 3)
  
      for m_index in range(size):
          x_2 = training_data[m_index + s]
          for d_2 in ds:
              k_v[m_index * 3 + d_2 - 1] = kernel(x, x_2, d_1, d_2, *args)
  
      return k_v
  
  
  def efs_force_vector_unit(name, s, e, x, efs_force_kernel, hyps, cutoffs, hyps_mask):
  
      training_data = _global_training_data[name]
  
      size = e - s
      k_ef = np.zeros((1, size * 3))
      k_ff = np.zeros((3, size * 3))
      k_sf = np.zeros((6, size * 3))
  
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      for m_index in range(size):
          x_2 = training_data[m_index + s]
          ef, ff, sf = efs_force_kernel(x, x_2, *args)
  
          ind1 = m_index * 3
          ind2 = (m_index + 1) * 3
  
          k_ef[:, ind1:ind2] = ef
          k_ff[:, ind1:ind2] = ff
          k_sf[:, ind1:ind2] = sf
  
      return k_ef, k_ff, k_sf
  
  
  def efs_energy_vector_unit(name, s, e, x, efs_energy_kernel, hyps, cutoffs, hyps_mask):
  
      training_structures = _global_training_structures[name]
  
      size = e - s
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      k_ee = np.zeros((1, size))
      k_fe = np.zeros((3, size))
      k_se = np.zeros((6, size))
  
      for m_index in range(size):
          training_structure = training_structures[m_index + s]
  
          ee_curr = 0
          fe_curr = np.zeros(3)
          se_curr = np.zeros(6)
  
          for environment in training_structure:
              ee, fe, se = efs_energy_kernel(x, environment, *args)
              ee_curr += ee
              fe_curr += fe
              se_curr += se
  
          k_ee[:, m_index] = ee_curr
          k_fe[:, m_index] = fe_curr
          k_se[:, m_index] = se_curr
  
      return k_ee, k_fe, k_se
  
  
  def energy_energy_vector(
      name, kernel, x, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
  ):
      """
      Get a vector of covariances between the local energy of a test environment
      and the total energy labels in the training set.
      """
  
      size = len(_global_training_structures[name])
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          return energy_energy_vector_unit(
              name, 0, size, x, kernel, hyps, cutoffs, hyps_mask
          )
  
      block_id, nbatch = partition_vector(n_sample, size, n_cpus)
      pack_function = energy_energy_vector_unit
      mult = 1
  
      force_energy_vector = parallel_vector_construction(
          pack_function,
          name,
          x,
          kernel,
          hyps,
          cutoffs,
          hyps_mask,
          block_id,
          nbatch,
          size,
          mult,
      )
  
      return force_energy_vector
  
  
  def energy_force_vector(
      name, kernel, x, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
  ):
      """
      Get the vector of covariances between the local energy of a test
      environment and the force labels in the training set.
      """
  
      training_data = _global_training_data[name]
      size = len(training_data)
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          return energy_force_vector_unit(
              name, 0, size, x, kernel, hyps, cutoffs, hyps_mask
          )
  
      block_id, nbatch = partition_vector(n_sample, size, n_cpus)
      pack_function = energy_force_vector_unit
      mult = 3
  
      k12_v = parallel_vector_construction(
          pack_function,
          name,
          x,
          kernel,
          hyps,
          cutoffs,
          hyps_mask,
          block_id,
          nbatch,
          size,
          mult,
      )
  
      return k12_v
  
  
  def force_energy_vector(
      name, kernel, x, d_1, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
  ):
      """
      Get a vector of covariances between a force component of a test environment
      and the total energy labels in the training set.
      """
  
      size = len(_global_training_structures[name])
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          return force_energy_vector_unit(
              name, 0, size, x, kernel, hyps, cutoffs, hyps_mask, d_1
          )
  
      block_id, nbatch = partition_vector(n_sample, size, n_cpus)
      pack_function = force_energy_vector_unit
      mult = 1
  
      force_energy_vector = parallel_vector_construction(
          pack_function,
          name,
          x,
          kernel,
          hyps,
          cutoffs,
          hyps_mask,
          block_id,
          nbatch,
          size,
          mult,
          d_1,
      )
  
      return force_energy_vector
  
  
  def force_force_vector(
      name, kernel, x, d_1, hyps, cutoffs=None, hyps_mask=None, n_cpus=1, n_sample=100
  ):
      """
      Get a vector of covariances between a force component of a test environment
      and the force labels in the training set.
      """
  
      size = len(_global_training_data[name])
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          return force_force_vector_unit(
              name, 0, size, x, kernel, hyps, cutoffs, hyps_mask, d_1
          )
  
      block_id, nbatch = partition_vector(n_sample, size, n_cpus)
      pack_function = force_force_vector_unit
      mult = 3
  
      k12_v = parallel_vector_construction(
          pack_function,
          name,
          x,
          kernel,
          hyps,
          cutoffs,
          hyps_mask,
          block_id,
          nbatch,
          size,
          mult,
          d_1,
      )
  
      return k12_v
  
  
  def efs_force_vector(
      name,
      efs_force_kernel,
      x,
      hyps,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      """
      Returns covariances between the local eneregy, force components, and
      partial stresses of a test environment and the force labels in the
      training set.
      """
  
      size = len(_global_training_data[name])
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
  
      # Perform serial calculation if n_cpus = 1.
      if n_cpus == 1:
          return efs_force_vector_unit(
              name, 0, size, x, efs_force_kernel, hyps, cutoffs, hyps_mask
          )
  
      # Otherwise, perform parallel calculation.
      block_id, nbatch = partition_vector(n_sample, size, n_cpus)
      pack_function = efs_force_vector_unit
      mult = 3
      n_comps = size * mult
      array_sizes = [(1, n_comps), (3, n_comps), (6, n_comps)]
  
      efs_arrays = multiple_array_construction(
          pack_function,
          name,
          x,
          efs_force_kernel,
          hyps,
          cutoffs,
          hyps_mask,
          block_id,
          nbatch,
          size,
          mult,
          array_sizes,
      )
  
      return efs_arrays
  
  
  def efs_energy_vector(
      name,
      efs_energy_kernel,
      x,
      hyps,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      """
      Returns covariances between the local eneregy, force components, and
      partial stresses of a test environment and the total energy labels in the
      training set.
      """
  
      size = len(_global_training_structures[name])
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
  
      # Perform serial calculation if n_cpus = 1.
      if n_cpus == 1:
          return efs_energy_vector_unit(
              name, 0, size, x, efs_energy_kernel, hyps, cutoffs, hyps_mask
          )
  
      # Otherwise, perform parallel calculation.
      block_id, nbatch = partition_vector(n_sample, size, n_cpus)
      pack_function = efs_energy_vector_unit
      mult = 1
      array_sizes = [(1, size), (3, size), (6, size)]
  
      efs_arrays = multiple_array_construction(
          pack_function,
          name,
          x,
          efs_energy_kernel,
          hyps,
          cutoffs,
          hyps_mask,
          block_id,
          nbatch,
          size,
          mult,
          array_sizes,
      )
  
      return efs_arrays
  
  
  def get_kernel_vector(
      name,
      force_force_kernel: Callable,
      force_energy_kernel: Callable,
      x,
      d_1,
      hyps,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      size1 = len(_global_training_data[name])
      size2 = len(_global_training_structures[name])
      kernel_vector = np.zeros(size1 * 3 + size2)
  
      force_vector = force_force_vector(
          name, force_force_kernel, x, d_1, hyps, cutoffs, hyps_mask, n_cpus, n_sample
      )
      energy_vector = force_energy_vector(
          name, force_energy_kernel, x, d_1, hyps, cutoffs, hyps_mask, n_cpus, n_sample
      )
  
      kernel_vector[0 : size1 * 3] = force_vector
      kernel_vector[size1 * 3 :] = energy_vector
  
      return kernel_vector
  
  
  def en_kern_vec(
      name,
      energy_force_kernel,
      energy_energy_kernel,
      x,
      hyps,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      size1 = len(_global_training_data[name])
      size2 = len(_global_training_structures[name])
      kernel_vector = np.zeros(size1 * 3 + size2)
  
      force_vector = energy_force_vector(
          name, energy_force_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
      )
      energy_vector = energy_energy_vector(
          name, energy_energy_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
      )
  
      kernel_vector[0 : size1 * 3] = force_vector
      kernel_vector[size1 * 3 :] = energy_vector
  
      return kernel_vector
  
  
  def efs_kern_vec(
      name,
      efs_force_kernel,
      efs_energy_kernel,
      x,
      hyps,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
  
      size1 = len(_global_training_data[name])
      size2 = len(_global_training_structures[name])
      tot_size = size1 * 3 + size2
  
      # Initialize arrays.
      energy_vector = np.zeros(tot_size)
      force_array = np.zeros((3, tot_size))
      stress_array = np.zeros((6, tot_size))
  
      # Compute force and energy arrays.
      force_arrays = efs_force_vector(
          name, efs_force_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
      )
  
      energy_arrays = efs_energy_vector(
          name, efs_energy_kernel, x, hyps, cutoffs, hyps_mask, n_cpus, n_sample
      )
  
      # Populate arrays.
      energy_vector[0 : size1 * 3] = force_arrays[0]
      energy_vector[size1 * 3 :] = energy_arrays[0]
      force_array[:, 0 : size1 * 3] = force_arrays[1]
      force_array[:, size1 * 3 :] = energy_arrays[1]
      stress_array[:, 0 : size1 * 3] = force_arrays[2]
      stress_array[:, size1 * 3 :] = energy_arrays[2]
  
      return energy_vector, force_array, stress_array
  
  
  # --------------------------------------------------------------------------
  #                      Ky hyperparameter gradient
  # --------------------------------------------------------------------------
  
  
  def get_ky_and_hyp_pack(
      name,
      s1,
      e1,
      s2,
      e2,
      same: bool,
      hyps: np.ndarray,
      kernel_grad,
      cutoffs=None,
      hyps_mask=None,
  ):
      """
      computes a block of ky matrix and its derivative to hyper-parameter
      If the cpu set up is None, it uses as much as posible cpus
  
      :param hyps: list of hyper-parameters
      :param name: name of the gp instance.
      :param kernel_grad: function object of the kernel gradient
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
  
      :return: hyp_mat, ky_mat
      """
  
      # assume sigma_n is the final hyperparameter
      _, non_noise_hyps, _ = obtain_noise_len(hyps, hyps_mask)
  
      # initialize matrices
      size1 = (e1 - s1) * 3
      size2 = (e2 - s2) * 3
      k_mat = np.zeros([size1, size2])
      hyp_mat = np.zeros([non_noise_hyps, size1, size2])
  
      args = from_mask_to_args(hyps, cutoffs, hyps_mask)
  
      ds = [1, 2, 3]
  
      training_data = _global_training_data[name]
      # calculate elements
      for m_index in range(size1):
          x_1 = training_data[int(math.floor(m_index / 3)) + s1]
          d_1 = ds[m_index % 3]
  
          if same:
              lowbound = m_index
          else:
              lowbound = 0
          for n_index in range(lowbound, size2):
              x_2 = training_data[int(math.floor(n_index / 3)) + s2]
              d_2 = ds[n_index % 3]
  
              # calculate kernel and gradient
              cov = kernel_grad(x_1, x_2, d_1, d_2, *args)
  
              # store kernel value
              k_mat[m_index, n_index] = cov[0]
              grad = from_grad_to_mask(cov[1], hyps_mask)
  
              hyp_mat[:, m_index, n_index] = grad
              if same:
                  k_mat[n_index, m_index] = cov[0]
                  hyp_mat[:, n_index, m_index] = grad
  
      return hyp_mat, k_mat
  
  
  def get_ky_and_hyp(
      hyps: np.ndarray,
      name,
      kernel_grad,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      """
      parallel version of get_ky_and_hyp
  
      :param hyps: list of hyper-parameters
      :param name: name of the gp instance.
      :param kernel_grad: function object of the kernel gradient
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
      :param n_cpus: number of cpus to use.
      :param n_sample: the size of block for matrix to compute
  
      :return: hyp_mat, ky_mat
      """
  
      training_data = _global_training_data[name]
      size = len(training_data)
      size3 = size * 3
  
      sigma_n, non_noise_hyps, train_noise = obtain_noise_len(hyps, hyps_mask)
  
      if n_cpus is None:
          n_cpus = mp.cpu_count()
      if n_cpus == 1:
          hyp_mat0, k_mat = get_ky_and_hyp_pack(
              name, 0, size, 0, size, True, hyps, kernel_grad, cutoffs, hyps_mask
          )
      else:
  
          block_id, nbatch = partition_matrix(n_sample, size, n_cpus)
  
          result_queue = mp.Queue()
          children = []
          for wid in range(nbatch):
              s1, e1, s2, e2 = block_id[wid]
              children.append(
                  mp.Process(
                      target=queue_wrapper,
                      args=(
                          result_queue,
                          wid,
                          get_ky_and_hyp_pack,
                          (
                              name,
                              s1,
                              e1,
                              s2,
                              e2,
                              s1 == s2,
                              hyps,
                              kernel_grad,
                              cutoffs,
                              hyps_mask,
                          ),
                      ),
                  )
              )
  
          # Run child processes.
          for c in children:
              c.start()
  
          k_mat = np.zeros([size3, size3])
          hyp_mat0 = np.zeros([non_noise_hyps, size3, size3])
          for _ in range(nbatch):
              wid, result_chunk = result_queue.get(block=True)
              s1, e1, s2, e2 = block_id[wid]
              h_mat_block, k_mat_block = result_chunk
              k_mat[s1 * 3 : e1 * 3, s2 * 3 : e2 * 3] = k_mat_block
              hyp_mat0[:, s1 * 3 : e1 * 3, s2 * 3 : e2 * 3] = h_mat_block
              if s1 != s2:
                  k_mat[s2 * 3 : e2 * 3, s1 * 3 : e1 * 3] = k_mat_block.T
                  for idx in range(hyp_mat0.shape[0]):
                      hyp_mat0[idx, s2 * 3 : e2 * 3, s1 * 3 : e1 * 3] = h_mat_block[idx].T
  
          # Join child processes (clean up zombies).
          for c in children:
              c.join()
  
      # add gradient of noise variance
      if train_noise:
          hyp_mat = np.zeros(
              [hyp_mat0.shape[0] + 1, hyp_mat0.shape[1], hyp_mat0.shape[2]]
          )
          hyp_mat[:-1, :, :] = hyp_mat0
          hyp_mat[-1, :, :] = np.eye(size3) * 2 * sigma_n
      else:
          hyp_mat = hyp_mat0
  
      # matrix manipulation
      ky_mat = k_mat + sigma_n ** 2 * np.eye(size3)
  
      return hyp_mat, ky_mat
  
  
  # --------------------------------------------------------------------------
  #                        Computing the likelihood
  # --------------------------------------------------------------------------
  
  
  def get_like_from_mats(ky_mat, l_mat, alpha, name):
      """compute the likelihood from the covariance matrix
  
      :param ky_mat: the covariance matrix
  
      :return: float, likelihood
      """
      # catch linear algebra errors
      force_labels = _global_training_labels[name]
      energy_labels = _global_energy_labels[name]
      labels = np.concatenate((force_labels, energy_labels))
  
      # calculate likelihood
      like = (
          -0.5 * np.matmul(labels, alpha)
          - np.sum(np.log(np.diagonal(l_mat)))
          - math.log(2 * np.pi) * ky_mat.shape[1] / 2
      )
  
      return like
  
  
  def get_neg_like(
      hyps: np.ndarray,
      name: str,
      force_kernel,
      logger_name=None,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      """compute the log likelihood and its gradients
      :param hyps: list of hyper-parameters
      :type hyps: np.ndarray
      :param name: name of the gp instance.
      :param kernel_grad: function object of the kernel gradient
      :param output: Output object for dumping every hyper-parameter
                     sets computed
      :type output: logger_name
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
      :param n_cpus: number of cpus to use.
      :param n_sample: the size of block for matrix to compute
      :return: float, np.array
      """
  
      time0 = time.time()
  
      ky_mat = get_force_block(
          hyps,
          name,
          force_kernel,
          cutoffs=cutoffs,
          hyps_mask=hyps_mask,
          n_cpus=n_cpus,
          n_sample=n_sample,
      )
  
      logger = logging.getLogger(logger_name)
      logger.debug(f"{name} get_ky_and_hyp {time.time()-time0}")
  
      time0 = time.time()
  
      # catch linear algebra errors
      try:
          ky_mat_inv = np.linalg.inv(ky_mat)
          l_mat = np.linalg.cholesky(ky_mat)
      except np.linalg.LinAlgError:
          number_of_hyps = len(hyps)
          return -1e8, np.zeros(number_of_hyps)
  
      labels = _global_training_labels[name]
  
      alpha = np.matmul(ky_mat_inv, labels)
  
      like = get_like_from_mats(ky_mat,l_mat,alpha,name)
  
      logger.debug(f"get_like_from_mats {time.time()-time0}")
  
      logger.debug(f"{name} Hyperparameters: {list(hyps)}")
      logger.debug(f"{name} Likelihood: {like}")
  
      return -like
  
  def get_neg_like_grad(
      hyps: np.ndarray,
      name: str,
      kernel_grad,
      logger_name: str = None,
      cutoffs=None,
      hyps_mask=None,
      n_cpus=1,
      n_sample=100,
  ):
      """compute the log likelihood and its gradients
  
      :param hyps: list of hyper-parameters
      :type hyps: np.ndarray
      :param name: name of the gp instance.
      :param kernel_grad: function object of the kernel gradient
      :param logger_name: name of logger object for dumping every hyper-parameter
                     sets computed
      :type logger_name: str
      :param cutoffs: The cutoff values used for the atomic environments
      :type cutoffs: list of 2 float numbers
      :param hyps_mask: dictionary used for multi-group hyperparmeters
      :param n_cpus: number of cpus to use.
      :param n_sample: the size of block for matrix to compute
  
      :return: float, np.array
      """
  
      time0 = time.time()
  
      hyp_mat, ky_mat = get_ky_and_hyp(
          hyps,
          name,
          kernel_grad,
          cutoffs=cutoffs,
          hyps_mask=hyps_mask,
          n_cpus=n_cpus,
          n_sample=n_sample,
      )
  
      logger = logging.getLogger(logger_name)
      logger.debug(f"get_ky_and_hyp {time.time()-time0}")
  
      time0 = time.time()
  
      like, like_grad = get_like_grad_from_mats(ky_mat, hyp_mat, name)
  
      logger.debug(f"get_like_grad_from_mats {time.time()-time0}")
  
      logger.debug("")
      logger.info(f"Hyperparameters: {list(hyps)}")
      logger.info(f"Likelihood: {like}")
      logger.info(f"Likelihood Gradient: {list(like_grad)}")
  
      return -like, -like_grad
  
  
  def get_like_grad_from_mats(ky_mat, hyp_mat, name):
      """compute the gradient of likelihood to hyper-parameters
      from covariance matrix and its gradient
  
      :param ky_mat: covariance matrix
      :type ky_mat: np.array
      :param hyp_mat: dky/d(hyper parameter) matrix
      :type hyp_mat: np.array
      :param name: name of the gp instance.
  
      :return: float, list. the likelihood and its gradients
      """
  
      number_of_hyps = hyp_mat.shape[0]
  
      # catch linear algebra errors
      try:
          ky_mat_inv = np.linalg.inv(ky_mat)
          l_mat = np.linalg.cholesky(ky_mat)
      except np.linalg.LinAlgError:
          return -1e8, np.zeros(number_of_hyps)
  
      labels = _global_training_labels[name]
  
      alpha = np.matmul(ky_mat_inv, labels)
      alpha_mat = np.matmul(alpha.reshape(-1, 1), alpha.reshape(1, -1))
      like_mat = alpha_mat - ky_mat_inv
  
      # calculate likelihood
      like = (
          -0.5 * np.matmul(labels, alpha)
          - np.sum(np.log(np.diagonal(l_mat)))
          - math.log(2 * np.pi) * ky_mat.shape[1] / 2
      )
  
      # calculate likelihood gradient
      like_grad = np.zeros(number_of_hyps)
      for n in range(number_of_hyps):
          like_grad[n] = 0.5 * np.trace(np.matmul(like_mat, hyp_mat[n, :, :]))
  
      return like, like_grad
  