files: import numpy as np
  import numba
  import umap.distances as dist
  from umap.utils import tau_rand_int
  
  
  @numba.njit()
  def clip(val):
      """Standard clamping of a value into a fixed range (in this case -4.0 to
      4.0)
  
      Parameters
      ----------
      val: float
          The value to be clamped.
  
      Returns
      -------
      The clamped value, now fixed to be in the range -4.0 to 4.0.
      """
      if val > 4.0:
          return 4.0
      elif val < -4.0:
          return -4.0
      else:
          return val
  
  
  @numba.njit(
      "f4(f4[::1],f4[::1])",
      fastmath=True,
      cache=True,
      locals={
          "result": numba.types.float32,
          "diff": numba.types.float32,
          "dim": numba.types.intp,
      },
  )
  def rdist(x, y):
      """Reduced Euclidean distance.
  
      Parameters
      ----------
      x: array of shape (embedding_dim,)
      y: array of shape (embedding_dim,)
  
      Returns
      -------
      The squared euclidean distance between x and y
      """
      result = 0.0
      dim = x.shape[0]
      for i in range(dim):
          diff = x[i] - y[i]
          result += diff * diff
  
      return result
  
  
  def _optimize_layout_euclidean_single_epoch(
      head_embedding,
      tail_embedding,
      head,
      tail,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma,
      dim,
      move_other,
      alpha,
      epochs_per_negative_sample,
      epoch_of_next_negative_sample,
      epoch_of_next_sample,
      n,
      densmap_flag,
      dens_phi_sum,
      dens_re_sum,
      dens_re_cov,
      dens_re_std,
      dens_re_mean,
      dens_lambda,
      dens_R,
      dens_mu,
      dens_mu_tot,
  ):
      for i in numba.prange(epochs_per_sample.shape[0]):
          if epoch_of_next_sample[i] <= n:
              j = head[i]
              k = tail[i]
  
              current = head_embedding[j]
              other = tail_embedding[k]
  
              dist_squared = rdist(current, other)
  
              if densmap_flag:
                  phi = 1.0 / (1.0 + a * pow(dist_squared, b))
                  dphi_term = (
                      a * b * pow(dist_squared, b - 1) / (1.0 + a * pow(dist_squared, b))
                  )
  
                  q_jk = phi / dens_phi_sum[k]
                  q_kj = phi / dens_phi_sum[j]
  
                  drk = q_jk * (
                      (1.0 - b * (1 - phi)) / np.exp(dens_re_sum[k]) + dphi_term
                  )
                  drj = q_kj * (
                      (1.0 - b * (1 - phi)) / np.exp(dens_re_sum[j]) + dphi_term
                  )
  
                  re_std_sq = dens_re_std * dens_re_std
                  weight_k = (
                      dens_R[k]
                      - dens_re_cov * (dens_re_sum[k] - dens_re_mean) / re_std_sq
                  )
                  weight_j = (
                      dens_R[j]
                      - dens_re_cov * (dens_re_sum[j] - dens_re_mean) / re_std_sq
                  )
  
                  grad_cor_coeff = (
                      dens_lambda
                      * dens_mu_tot
                      * (weight_k * drk + weight_j * drj)
                      / (dens_mu[i] * dens_re_std)
                      / n_vertices
                  )
  
              if dist_squared > 0.0:
                  grad_coeff = -2.0 * a * b * pow(dist_squared, b - 1.0)
                  grad_coeff /= a * pow(dist_squared, b) + 1.0
              else:
                  grad_coeff = 0.0
  
              for d in range(dim):
                  grad_d = clip(grad_coeff * (current[d] - other[d]))
  
                  if densmap_flag:
                      grad_d += clip(2 * grad_cor_coeff * (current[d] - other[d]))
  
                  current[d] += grad_d * alpha
                  if move_other:
                      other[d] += -grad_d * alpha
  
              epoch_of_next_sample[i] += epochs_per_sample[i]
  
              n_neg_samples = int(
                  (n - epoch_of_next_negative_sample[i]) / epochs_per_negative_sample[i]
              )
  
              for p in range(n_neg_samples):
                  k = tau_rand_int(rng_state) % n_vertices
  
                  other = tail_embedding[k]
  
                  dist_squared = rdist(current, other)
  
                  if dist_squared > 0.0:
                      grad_coeff = 2.0 * gamma * b
                      grad_coeff /= (0.001 + dist_squared) * (
                          a * pow(dist_squared, b) + 1
                      )
                  elif j == k:
                      continue
                  else:
                      grad_coeff = 0.0
  
                  for d in range(dim):
                      if grad_coeff > 0.0:
                          grad_d = clip(grad_coeff * (current[d] - other[d]))
                      else:
                          grad_d = 4.0
                      current[d] += grad_d * alpha
  
              epoch_of_next_negative_sample[i] += (
                  n_neg_samples * epochs_per_negative_sample[i]
              )
  
  
  def _optimize_layout_euclidean_densmap_epoch_init(
      head_embedding,
      tail_embedding,
      head,
      tail,
      a,
      b,
      re_sum,
      phi_sum,
  ):
      re_sum.fill(0)
      phi_sum.fill(0)
  
      for i in numba.prange(head.size):
          j = head[i]
          k = tail[i]
  
          current = head_embedding[j]
          other = tail_embedding[k]
          dist_squared = rdist(current, other)
  
          phi = 1.0 / (1.0 + a * pow(dist_squared, b))
  
          re_sum[j] += phi * dist_squared
          re_sum[k] += phi * dist_squared
          phi_sum[j] += phi
          phi_sum[k] += phi
  
      epsilon = 1e-8
      for i in range(re_sum.size):
          re_sum[i] = np.log(epsilon + (re_sum[i] / phi_sum[i]))
  
  
  def optimize_layout_euclidean(
      head_embedding,
      tail_embedding,
      head,
      tail,
      n_epochs,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma=1.0,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      parallel=False,
      verbose=False,
      densmap=False,
      densmap_kwds={},
      move_other=False,
  ):
      """Improve an embedding using stochastic gradient descent to minimize the
      fuzzy set cross entropy between the 1-skeletons of the high dimensional
      and low dimensional fuzzy simplicial sets. In practice this is done by
      sampling edges based on their membership strength (with the (1-p) terms
      coming from negative sampling similar to word2vec).
      Parameters
      ----------
      head_embedding: array of shape (n_samples, n_components)
          The initial embedding to be improved by SGD.
      tail_embedding: array of shape (source_samples, n_components)
          The reference embedding of embedded points. If not embedding new
          previously unseen points with respect to an existing embedding this
          is simply the head_embedding (again); otherwise it provides the
          existing embedding to embed with respect to.
      head: array of shape (n_1_simplices)
          The indices of the heads of 1-simplices with non-zero membership.
      tail: array of shape (n_1_simplices)
          The indices of the tails of 1-simplices with non-zero membership.
      n_epochs: int, or list of int
          The number of training epochs to use in optimization, or a list of
          epochs at which to save the embedding. In case of a list, the optimization
          will use the maximum number of epochs in the list, and will return a list
          of embedding in the order of increasing epoch, regardless of the order in
          the epoch list.
      n_vertices: int
          The number of vertices (0-simplices) in the dataset.
      epochs_per_samples: array of shape (n_1_simplices)
          A float value of the number of epochs per 1-simplex. 1-simplices with
          weaker membership strength will have more epochs between being sampled.
      a: float
          Parameter of differentiable approximation of right adjoint functor
      b: float
          Parameter of differentiable approximation of right adjoint functor
      rng_state: array of int64, shape (3,)
          The internal state of the rng
      gamma: float (optional, default 1.0)
          Weight to apply to negative samples.
      initial_alpha: float (optional, default 1.0)
          Initial learning rate for the SGD.
      negative_sample_rate: int (optional, default 5)
          Number of negative samples to use per positive sample.
      parallel: bool (optional, default False)
          Whether to run the computation using numba parallel.
          Running in parallel is non-deterministic, and is not used
          if a random seed has been set, to ensure reproducibility.
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
      densmap: bool (optional, default False)
          Whether to use the density-augmented densMAP objective
      densmap_kwds: dict (optional, default {})
          Auxiliary data for densMAP
      move_other: bool (optional, default False)
          Whether to adjust tail_embedding alongside head_embedding
      Returns
      -------
      embedding: array of shape (n_samples, n_components)
          The optimized embedding.
      """
  
      dim = head_embedding.shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = epochs_per_sample / negative_sample_rate
      epoch_of_next_negative_sample = epochs_per_negative_sample.copy()
      epoch_of_next_sample = epochs_per_sample.copy()
  
      optimize_fn = numba.njit(
          _optimize_layout_euclidean_single_epoch, fastmath=True, parallel=parallel
      )
  
      if densmap:
          dens_init_fn = numba.njit(
              _optimize_layout_euclidean_densmap_epoch_init,
              fastmath=True,
              parallel=parallel,
          )
  
          dens_mu_tot = np.sum(densmap_kwds["mu_sum"]) / 2
          dens_lambda = densmap_kwds["lambda"]
          dens_R = densmap_kwds["R"]
          dens_mu = densmap_kwds["mu"]
          dens_phi_sum = np.zeros(n_vertices, dtype=np.float32)
          dens_re_sum = np.zeros(n_vertices, dtype=np.float32)
          dens_var_shift = densmap_kwds["var_shift"]
      else:
          dens_mu_tot = 0
          dens_lambda = 0
          dens_R = np.zeros(1, dtype=np.float32)
          dens_mu = np.zeros(1, dtype=np.float32)
          dens_phi_sum = np.zeros(1, dtype=np.float32)
          dens_re_sum = np.zeros(1, dtype=np.float32)
  
      epochs_list = None
      embedding_list = []
      if isinstance(n_epochs, list):
          epochs_list = n_epochs
          n_epochs = max(epochs_list)
  
      for n in range(n_epochs):
  
          densmap_flag = (
              densmap
              and (densmap_kwds["lambda"] > 0)
              and (((n + 1) / float(n_epochs)) > (1 - densmap_kwds["frac"]))
          )
  
          if densmap_flag:
              dens_init_fn(
                  head_embedding,
                  tail_embedding,
                  head,
                  tail,
                  a,
                  b,
                  dens_re_sum,
                  dens_phi_sum,
              )
  
              dens_re_std = np.sqrt(np.var(dens_re_sum) + dens_var_shift)
              dens_re_mean = np.mean(dens_re_sum)
              dens_re_cov = np.dot(dens_re_sum, dens_R) / (n_vertices - 1)
          else:
              dens_re_std = 0
              dens_re_mean = 0
              dens_re_cov = 0
  
          optimize_fn(
              head_embedding,
              tail_embedding,
              head,
              tail,
              n_vertices,
              epochs_per_sample,
              a,
              b,
              rng_state,
              gamma,
              dim,
              move_other,
              alpha,
              epochs_per_negative_sample,
              epoch_of_next_negative_sample,
              epoch_of_next_sample,
              n,
              densmap_flag,
              dens_phi_sum,
              dens_re_sum,
              dens_re_cov,
              dens_re_std,
              dens_re_mean,
              dens_lambda,
              dens_R,
              dens_mu,
              dens_mu_tot,
          )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
          if epochs_list is not None and n in epochs_list:
              embedding_list.append(head_embedding.copy())
  
      # Add the last embedding to the list as well
      if epochs_list is not None:
          embedding_list.append(head_embedding.copy())
  
      return head_embedding if epochs_list is None else embedding_list
  
  
  @numba.njit(fastmath=True)
  def optimize_layout_generic(
      head_embedding,
      tail_embedding,
      head,
      tail,
      n_epochs,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma=1.0,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      output_metric=dist.euclidean,
      output_metric_kwds=(),
      verbose=False,
      move_other=False,
  ):
      """Improve an embedding using stochastic gradient descent to minimize the
      fuzzy set cross entropy between the 1-skeletons of the high dimensional
      and low dimensional fuzzy simplicial sets. In practice this is done by
      sampling edges based on their membership strength (with the (1-p) terms
      coming from negative sampling similar to word2vec).
  
      Parameters
      ----------
      head_embedding: array of shape (n_samples, n_components)
          The initial embedding to be improved by SGD.
  
      tail_embedding: array of shape (source_samples, n_components)
          The reference embedding of embedded points. If not embedding new
          previously unseen points with respect to an existing embedding this
          is simply the head_embedding (again); otherwise it provides the
          existing embedding to embed with respect to.
  
      head: array of shape (n_1_simplices)
          The indices of the heads of 1-simplices with non-zero membership.
  
      tail: array of shape (n_1_simplices)
          The indices of the tails of 1-simplices with non-zero membership.
  
      weight: array of shape (n_1_simplices)
          The membership weights of the 1-simplices.
  
      n_epochs: int
          The number of training epochs to use in optimization.
  
      n_vertices: int
          The number of vertices (0-simplices) in the dataset.
  
      epochs_per_sample: array of shape (n_1_simplices)
          A float value of the number of epochs per 1-simplex. 1-simplices with
          weaker membership strength will have more epochs between being sampled.
  
      a: float
          Parameter of differentiable approximation of right adjoint functor
  
      b: float
          Parameter of differentiable approximation of right adjoint functor
  
      rng_state: array of int64, shape (3,)
          The internal state of the rng
  
      gamma: float (optional, default 1.0)
          Weight to apply to negative samples.
  
      initial_alpha: float (optional, default 1.0)
          Initial learning rate for the SGD.
  
      negative_sample_rate: int (optional, default 5)
          Number of negative samples to use per positive sample.
  
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
  
      move_other: bool (optional, default False)
          Whether to adjust tail_embedding alongside head_embedding
  
      Returns
      -------
      embedding: array of shape (n_samples, n_components)
          The optimized embedding.
      """
  
      dim = head_embedding.shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = epochs_per_sample / negative_sample_rate
      epoch_of_next_negative_sample = epochs_per_negative_sample.copy()
      epoch_of_next_sample = epochs_per_sample.copy()
  
      for n in range(n_epochs):
          for i in range(epochs_per_sample.shape[0]):
              if epoch_of_next_sample[i] <= n:
                  j = head[i]
                  k = tail[i]
  
                  current = head_embedding[j]
                  other = tail_embedding[k]
  
                  dist_output, grad_dist_output = output_metric(
                      current, other, *output_metric_kwds
                  )
                  _, rev_grad_dist_output = output_metric(
                      other, current, *output_metric_kwds
                  )
  
                  if dist_output > 0.0:
                      w_l = pow((1 + a * pow(dist_output, 2 * b)), -1)
                  else:
                      w_l = 1.0
                  grad_coeff = 2 * b * (w_l - 1) / (dist_output + 1e-6)
  
                  for d in range(dim):
                      grad_d = clip(grad_coeff * grad_dist_output[d])
  
                      current[d] += grad_d * alpha
                      if move_other:
                          grad_d = clip(grad_coeff * rev_grad_dist_output[d])
                          other[d] += grad_d * alpha
  
                  epoch_of_next_sample[i] += epochs_per_sample[i]
  
                  n_neg_samples = int(
                      (n - epoch_of_next_negative_sample[i])
                      / epochs_per_negative_sample[i]
                  )
  
                  for p in range(n_neg_samples):
                      k = tau_rand_int(rng_state) % n_vertices
  
                      other = tail_embedding[k]
  
                      dist_output, grad_dist_output = output_metric(
                          current, other, *output_metric_kwds
                      )
  
                      if dist_output > 0.0:
                          w_l = pow((1 + a * pow(dist_output, 2 * b)), -1)
                      elif j == k:
                          continue
                      else:
                          w_l = 1.0
  
                      grad_coeff = gamma * 2 * b * w_l / (dist_output + 1e-6)
  
                      for d in range(dim):
                          grad_d = clip(grad_coeff * grad_dist_output[d])
                          current[d] += grad_d * alpha
  
                  epoch_of_next_negative_sample[i] += (
                      n_neg_samples * epochs_per_negative_sample[i]
                  )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
      return head_embedding
  
  
  @numba.njit(fastmath=True)
  def optimize_layout_inverse(
      head_embedding,
      tail_embedding,
      head,
      tail,
      weight,
      sigmas,
      rhos,
      n_epochs,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma=1.0,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      output_metric=dist.euclidean,
      output_metric_kwds=(),
      verbose=False,
      move_other=False,
  ):
      """Improve an embedding using stochastic gradient descent to minimize the
      fuzzy set cross entropy between the 1-skeletons of the high dimensional
      and low dimensional fuzzy simplicial sets. In practice this is done by
      sampling edges based on their membership strength (with the (1-p) terms
      coming from negative sampling similar to word2vec).
  
      Parameters
      ----------
      head_embedding: array of shape (n_samples, n_components)
          The initial embedding to be improved by SGD.
  
      tail_embedding: array of shape (source_samples, n_components)
          The reference embedding of embedded points. If not embedding new
          previously unseen points with respect to an existing embedding this
          is simply the head_embedding (again); otherwise it provides the
          existing embedding to embed with respect to.
  
      head: array of shape (n_1_simplices)
          The indices of the heads of 1-simplices with non-zero membership.
  
      tail: array of shape (n_1_simplices)
          The indices of the tails of 1-simplices with non-zero membership.
  
      weight: array of shape (n_1_simplices)
          The membership weights of the 1-simplices.
  
      n_epochs: int
          The number of training epochs to use in optimization.
  
      n_vertices: int
          The number of vertices (0-simplices) in the dataset.
  
      epochs_per_sample: array of shape (n_1_simplices)
          A float value of the number of epochs per 1-simplex. 1-simplices with
          weaker membership strength will have more epochs between being sampled.
  
      a: float
          Parameter of differentiable approximation of right adjoint functor
  
      b: float
          Parameter of differentiable approximation of right adjoint functor
  
      rng_state: array of int64, shape (3,)
          The internal state of the rng
  
      gamma: float (optional, default 1.0)
          Weight to apply to negative samples.
  
      initial_alpha: float (optional, default 1.0)
          Initial learning rate for the SGD.
  
      negative_sample_rate: int (optional, default 5)
          Number of negative samples to use per positive sample.
  
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
  
      move_other: bool (optional, default False)
          Whether to adjust tail_embedding alongside head_embedding
  
      Returns
      -------
      embedding: array of shape (n_samples, n_components)
          The optimized embedding.
      """
  
      dim = head_embedding.shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = epochs_per_sample / negative_sample_rate
      epoch_of_next_negative_sample = epochs_per_negative_sample.copy()
      epoch_of_next_sample = epochs_per_sample.copy()
  
      for n in range(n_epochs):
          for i in range(epochs_per_sample.shape[0]):
              if epoch_of_next_sample[i] <= n:
                  j = head[i]
                  k = tail[i]
  
                  current = head_embedding[j]
                  other = tail_embedding[k]
  
                  dist_output, grad_dist_output = output_metric(
                      current, other, *output_metric_kwds
                  )
  
                  w_l = weight[i]
                  grad_coeff = -(1 / (w_l * sigmas[k] + 1e-6))
  
                  for d in range(dim):
                      grad_d = clip(grad_coeff * grad_dist_output[d])
  
                      current[d] += grad_d * alpha
                      if move_other:
                          other[d] += -grad_d * alpha
  
                  epoch_of_next_sample[i] += epochs_per_sample[i]
  
                  n_neg_samples = int(
                      (n - epoch_of_next_negative_sample[i])
                      / epochs_per_negative_sample[i]
                  )
  
                  for p in range(n_neg_samples):
                      k = tau_rand_int(rng_state) % n_vertices
  
                      other = tail_embedding[k]
  
                      dist_output, grad_dist_output = output_metric(
                          current, other, *output_metric_kwds
                      )
  
                      # w_l = 0.0 # for negative samples, the edge does not exist
                      w_h = np.exp(-max(dist_output - rhos[k], 1e-6) / (sigmas[k] + 1e-6))
                      grad_coeff = -gamma * ((0 - w_h) / ((1 - w_h) * sigmas[k] + 1e-6))
  
                      for d in range(dim):
                          grad_d = clip(grad_coeff * grad_dist_output[d])
                          current[d] += grad_d * alpha
  
                  epoch_of_next_negative_sample[i] += (
                      n_neg_samples * epochs_per_negative_sample[i]
                  )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
      return head_embedding
  
  
  def _optimize_layout_aligned_euclidean_single_epoch(
      head_embeddings,
      tail_embeddings,
      heads,
      tails,
      epochs_per_sample,
      a,
      b,
      regularisation_weights,
      relations,
      rng_state,
      gamma,
      lambda_,
      dim,
      move_other,
      alpha,
      epochs_per_negative_sample,
      epoch_of_next_negative_sample,
      epoch_of_next_sample,
      n,
  ):
      n_embeddings = len(heads)
      window_size = (relations.shape[1] - 1) // 2
  
      max_n_edges = 0
      for e_p_s in epochs_per_sample:
          if e_p_s.shape[0] >= max_n_edges:
              max_n_edges = e_p_s.shape[0]
  
      embedding_order = np.arange(n_embeddings).astype(np.int32)
      np.random.seed(abs(rng_state[0]))
      np.random.shuffle(embedding_order)
  
      for i in range(max_n_edges):
          for m in embedding_order:
              if i < epoch_of_next_sample[m].shape[0] and epoch_of_next_sample[m][i] <= n:
                  j = heads[m][i]
                  k = tails[m][i]
  
                  current = head_embeddings[m][j]
                  other = tail_embeddings[m][k]
  
                  dist_squared = rdist(current, other)
  
                  if dist_squared > 0.0:
                      grad_coeff = -2.0 * a * b * pow(dist_squared, b - 1.0)
                      grad_coeff /= a * pow(dist_squared, b) + 1.0
                  else:
                      grad_coeff = 0.0
  
                  for d in range(dim):
                      grad_d = clip(grad_coeff * (current[d] - other[d]))
  
                      for offset in range(-window_size, window_size):
                          neighbor_m = m + offset
                          if (
                              neighbor_m >= 0
                              and neighbor_m < n_embeddings
                              and offset != 0
                          ):
                              identified_index = relations[m, offset + window_size, j]
                              if identified_index >= 0:
                                  grad_d -= clip(
                                      (lambda_ * np.exp(-(np.abs(offset) - 1)))
                                      * regularisation_weights[m, offset + window_size, j]
                                      * (
                                          current[d]
                                          - head_embeddings[neighbor_m][
                                              identified_index, d
                                          ]
                                      )
                                  )
  
                      current[d] += clip(grad_d) * alpha
                      if move_other:
                          other_grad_d = clip(grad_coeff * (other[d] - current[d]))
  
                          for offset in range(-window_size, window_size):
                              neighbor_m = m + offset
                              if (
                                  neighbor_m >= 0
                                  and neighbor_m < n_embeddings
                                  and offset != 0
                              ):
                                  identified_index = relations[m, offset + window_size, k]
                                  if identified_index >= 0:
                                      grad_d -= clip(
                                          (lambda_ * np.exp(-(np.abs(offset) - 1)))
                                          * regularisation_weights[
                                              m, offset + window_size, k
                                          ]
                                          * (
                                              other[d]
                                              - head_embeddings[neighbor_m][
                                                  identified_index, d
                                              ]
                                          )
                                      )
  
                          other[d] += clip(other_grad_d) * alpha
  
                  epoch_of_next_sample[m][i] += epochs_per_sample[m][i]
  
                  if epochs_per_negative_sample[m][i] > 0:
                      n_neg_samples = int(
                          (n - epoch_of_next_negative_sample[m][i])
                          / epochs_per_negative_sample[m][i]
                      )
                  else:
                      n_neg_samples = 0
  
                  for p in range(n_neg_samples):
                      k = tau_rand_int(rng_state) % tail_embeddings[m].shape[0]
  
                      other = tail_embeddings[m][k]
  
                      dist_squared = rdist(current, other)
  
                      if dist_squared > 0.0:
                          grad_coeff = 2.0 * gamma * b
                          grad_coeff /= (0.001 + dist_squared) * (
                              a * pow(dist_squared, b) + 1
                          )
                      elif j == k:
                          continue
                      else:
                          grad_coeff = 0.0
  
                      for d in range(dim):
                          if grad_coeff > 0.0:
                              grad_d = clip(grad_coeff * (current[d] - other[d]))
                          else:
                              grad_d = 4.0
  
                          for offset in range(-window_size, window_size):
                              neighbor_m = m + offset
                              if (
                                  neighbor_m >= 0
                                  and neighbor_m < n_embeddings
                                  and offset != 0
                              ):
                                  identified_index = relations[m, offset + window_size, j]
                                  if identified_index >= 0:
                                      grad_d -= clip(
                                          (lambda_ * np.exp(-(np.abs(offset) - 1)))
                                          * regularisation_weights[
                                              m, offset + window_size, j
                                          ]
                                          * (
                                              current[d]
                                              - head_embeddings[neighbor_m][
                                                  identified_index, d
                                              ]
                                          )
                                      )
  
                          current[d] += clip(grad_d) * alpha
  
                  epoch_of_next_negative_sample[m][i] += (
                      n_neg_samples * epochs_per_negative_sample[m][i]
                  )
  
  
  def optimize_layout_aligned_euclidean(
      head_embeddings,
      tail_embeddings,
      heads,
      tails,
      n_epochs,
      epochs_per_sample,
      regularisation_weights,
      relations,
      rng_state,
      a=1.576943460405378,
      b=0.8950608781227859,
      gamma=1.0,
      lambda_=5e-3,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      parallel=True,
      verbose=False,
      move_other=False,
  ):
      dim = head_embeddings[0].shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = numba.typed.List.empty_list(numba.types.float32[::1])
      epoch_of_next_negative_sample = numba.typed.List.empty_list(
          numba.types.float32[::1]
      )
      epoch_of_next_sample = numba.typed.List.empty_list(numba.types.float32[::1])
  
      for m in range(len(heads)):
          epochs_per_negative_sample.append(
              epochs_per_sample[m].astype(np.float32) / negative_sample_rate
          )
          epoch_of_next_negative_sample.append(
              epochs_per_negative_sample[m].astype(np.float32)
          )
          epoch_of_next_sample.append(epochs_per_sample[m].astype(np.float32))
  
      optimize_fn = numba.njit(
          _optimize_layout_aligned_euclidean_single_epoch,
          fastmath=True,
          parallel=parallel,
      )
  
      for n in range(n_epochs):
          optimize_fn(
              head_embeddings,
              tail_embeddings,
              heads,
              tails,
              epochs_per_sample,
              a,
              b,
              regularisation_weights,
              relations,
              rng_state,
              gamma,
              lambda_,
              dim,
              move_other,
              alpha,
              epochs_per_negative_sample,
              epoch_of_next_negative_sample,
              epoch_of_next_sample,
              n,
          )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
      return head_embeddings
  
  
  [... the rest is omitted ...]
files: umap/layouts.py
files: import numpy as np
  import numba
  import umap.distances as dist
  from umap.utils import tau_rand_int
  
  
  @numba.njit()
  def clip(val):
      """Standard clamping of a value into a fixed range (in this case -4.0 to
      4.0)
  
      Parameters
      ----------
      val: float
          The value to be clamped.
  
      Returns
      -------
      The clamped value, now fixed to be in the range -4.0 to 4.0.
      """
      if val > 4.0:
          return 4.0
      elif val < -4.0:
          return -4.0
      else:
          return val
  
  
  @numba.njit(
      "f4(f4[::1],f4[::1])",
      fastmath=True,
      cache=True,
      locals={
          "result": numba.types.float32,
          "diff": numba.types.float32,
          "dim": numba.types.intp,
      },
  )
  def rdist(x, y):
      """Reduced Euclidean distance.
  
      Parameters
      ----------
      x: array of shape (embedding_dim,)
      y: array of shape (embedding_dim,)
  
      Returns
      -------
      The squared euclidean distance between x and y
      """
      result = 0.0
      dim = x.shape[0]
      for i in range(dim):
          diff = x[i] - y[i]
          result += diff * diff
  
      return result
  
  
  def _optimize_layout_euclidean_single_epoch(
      head_embedding,
      tail_embedding,
      head,
      tail,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma,
      dim,
      move_other,
      alpha,
      epochs_per_negative_sample,
      epoch_of_next_negative_sample,
      epoch_of_next_sample,
      n,
      densmap_flag,
      dens_phi_sum,
      dens_re_sum,
      dens_re_cov,
      dens_re_std,
      dens_re_mean,
      dens_lambda,
      dens_R,
      dens_mu,
      dens_mu_tot,
  ):
      for i in numba.prange(epochs_per_sample.shape[0]):
          if epoch_of_next_sample[i] <= n:
              j = head[i]
              k = tail[i]
  
              current = head_embedding[j]
              other = tail_embedding[k]
  
              dist_squared = rdist(current, other)
  
              if densmap_flag:
                  phi = 1.0 / (1.0 + a * pow(dist_squared, b))
                  dphi_term = (
                      a * b * pow(dist_squared, b - 1) / (1.0 + a * pow(dist_squared, b))
                  )
  
                  q_jk = phi / dens_phi_sum[k]
                  q_kj = phi / dens_phi_sum[j]
  
                  drk = q_jk * (
                      (1.0 - b * (1 - phi)) / np.exp(dens_re_sum[k]) + dphi_term
                  )
                  drj = q_kj * (
                      (1.0 - b * (1 - phi)) / np.exp(dens_re_sum[j]) + dphi_term
                  )
  
                  re_std_sq = dens_re_std * dens_re_std
                  weight_k = (
                      dens_R[k]
                      - dens_re_cov * (dens_re_sum[k] - dens_re_mean) / re_std_sq
                  )
                  weight_j = (
                      dens_R[j]
                      - dens_re_cov * (dens_re_sum[j] - dens_re_mean) / re_std_sq
                  )
  
                  grad_cor_coeff = (
                      dens_lambda
                      * dens_mu_tot
                      * (weight_k * drk + weight_j * drj)
                      / (dens_mu[i] * dens_re_std)
                      / n_vertices
                  )
  
              if dist_squared > 0.0:
                  grad_coeff = -2.0 * a * b * pow(dist_squared, b - 1.0)
                  grad_coeff /= a * pow(dist_squared, b) + 1.0
              else:
                  grad_coeff = 0.0
  
              for d in range(dim):
                  grad_d = clip(grad_coeff * (current[d] - other[d]))
  
                  if densmap_flag:
                      grad_d += clip(2 * grad_cor_coeff * (current[d] - other[d]))
  
                  current[d] += grad_d * alpha
                  if move_other:
                      other[d] += -grad_d * alpha
  
              epoch_of_next_sample[i] += epochs_per_sample[i]
  
              n_neg_samples = int(
                  (n - epoch_of_next_negative_sample[i]) / epochs_per_negative_sample[i]
              )
  
              for p in range(n_neg_samples):
                  k = tau_rand_int(rng_state) % n_vertices
  
                  other = tail_embedding[k]
  
                  dist_squared = rdist(current, other)
  
                  if dist_squared > 0.0:
                      grad_coeff = 2.0 * gamma * b
                      grad_coeff /= (0.001 + dist_squared) * (
                          a * pow(dist_squared, b) + 1
                      )
                  elif j == k:
                      continue
                  else:
                      grad_coeff = 0.0
  
                  for d in range(dim):
                      if grad_coeff > 0.0:
                          grad_d = clip(grad_coeff * (current[d] - other[d]))
                      else:
                          grad_d = 4.0
                      current[d] += grad_d * alpha
  
              epoch_of_next_negative_sample[i] += (
                  n_neg_samples * epochs_per_negative_sample[i]
              )
  
  
  def _optimize_layout_euclidean_densmap_epoch_init(
      head_embedding,
      tail_embedding,
      head,
      tail,
      a,
      b,
      re_sum,
      phi_sum,
  ):
      re_sum.fill(0)
      phi_sum.fill(0)
  
      for i in numba.prange(head.size):
          j = head[i]
          k = tail[i]
  
          current = head_embedding[j]
          other = tail_embedding[k]
          dist_squared = rdist(current, other)
  
          phi = 1.0 / (1.0 + a * pow(dist_squared, b))
  
          re_sum[j] += phi * dist_squared
          re_sum[k] += phi * dist_squared
          phi_sum[j] += phi
          phi_sum[k] += phi
  
      epsilon = 1e-8
      for i in range(re_sum.size):
          re_sum[i] = np.log(epsilon + (re_sum[i] / phi_sum[i]))
  
  
  def optimize_layout_euclidean(
      head_embedding,
      tail_embedding,
      head,
      tail,
      n_epochs,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma=1.0,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      parallel=False,
      verbose=False,
      densmap=False,
      densmap_kwds={},
      move_other=False,
  ):
      """Improve an embedding using stochastic gradient descent to minimize the
      fuzzy set cross entropy between the 1-skeletons of the high dimensional
      and low dimensional fuzzy simplicial sets. In practice this is done by
      sampling edges based on their membership strength (with the (1-p) terms
      coming from negative sampling similar to word2vec).
      Parameters
      ----------
      head_embedding: array of shape (n_samples, n_components)
          The initial embedding to be improved by SGD.
      tail_embedding: array of shape (source_samples, n_components)
          The reference embedding of embedded points. If not embedding new
          previously unseen points with respect to an existing embedding this
          is simply the head_embedding (again); otherwise it provides the
          existing embedding to embed with respect to.
      head: array of shape (n_1_simplices)
          The indices of the heads of 1-simplices with non-zero membership.
      tail: array of shape (n_1_simplices)
          The indices of the tails of 1-simplices with non-zero membership.
      n_epochs: int
          The number of training epochs to use in optimization.
      n_vertices: int
          The number of vertices (0-simplices) in the dataset.
      epochs_per_samples: array of shape (n_1_simplices)
          A float value of the number of epochs per 1-simplex. 1-simplices with
          weaker membership strength will have more epochs between being sampled.
      a: float
          Parameter of differentiable approximation of right adjoint functor
      b: float
          Parameter of differentiable approximation of right adjoint functor
      rng_state: array of int64, shape (3,)
          The internal state of the rng
      gamma: float (optional, default 1.0)
          Weight to apply to negative samples.
      initial_alpha: float (optional, default 1.0)
          Initial learning rate for the SGD.
      negative_sample_rate: int (optional, default 5)
          Number of negative samples to use per positive sample.
      parallel: bool (optional, default False)
          Whether to run the computation using numba parallel.
          Running in parallel is non-deterministic, and is not used
          if a random seed has been set, to ensure reproducibility.
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
      densmap: bool (optional, default False)
          Whether to use the density-augmented densMAP objective
      densmap_kwds: dict (optional, default {})
          Auxiliary data for densMAP
      move_other: bool (optional, default False)
          Whether to adjust tail_embedding alongside head_embedding
      Returns
      -------
      embedding: array of shape (n_samples, n_components)
          The optimized embedding.
      """
  
      dim = head_embedding.shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = epochs_per_sample / negative_sample_rate
      epoch_of_next_negative_sample = epochs_per_negative_sample.copy()
      epoch_of_next_sample = epochs_per_sample.copy()
  
      optimize_fn = numba.njit(
          _optimize_layout_euclidean_single_epoch, fastmath=True, parallel=parallel
      )
  
      if densmap:
          dens_init_fn = numba.njit(
              _optimize_layout_euclidean_densmap_epoch_init,
              fastmath=True,
              parallel=parallel,
          )
  
          dens_mu_tot = np.sum(densmap_kwds["mu_sum"]) / 2
          dens_lambda = densmap_kwds["lambda"]
          dens_R = densmap_kwds["R"]
          dens_mu = densmap_kwds["mu"]
          dens_phi_sum = np.zeros(n_vertices, dtype=np.float32)
          dens_re_sum = np.zeros(n_vertices, dtype=np.float32)
          dens_var_shift = densmap_kwds["var_shift"]
      else:
          dens_mu_tot = 0
          dens_lambda = 0
          dens_R = np.zeros(1, dtype=np.float32)
          dens_mu = np.zeros(1, dtype=np.float32)
          dens_phi_sum = np.zeros(1, dtype=np.float32)
          dens_re_sum = np.zeros(1, dtype=np.float32)
  
      for n in range(n_epochs):
  
          densmap_flag = (
              densmap
              and (densmap_kwds["lambda"] > 0)
              and (((n + 1) / float(n_epochs)) > (1 - densmap_kwds["frac"]))
          )
  
          if densmap_flag:
              dens_init_fn(
                  head_embedding,
                  tail_embedding,
                  head,
                  tail,
                  a,
                  b,
                  dens_re_sum,
                  dens_phi_sum,
              )
  
              dens_re_std = np.sqrt(np.var(dens_re_sum) + dens_var_shift)
              dens_re_mean = np.mean(dens_re_sum)
              dens_re_cov = np.dot(dens_re_sum, dens_R) / (n_vertices - 1)
          else:
              dens_re_std = 0
              dens_re_mean = 0
              dens_re_cov = 0
  
          optimize_fn(
              head_embedding,
              tail_embedding,
              head,
              tail,
              n_vertices,
              epochs_per_sample,
              a,
              b,
              rng_state,
              gamma,
              dim,
              move_other,
              alpha,
              epochs_per_negative_sample,
              epoch_of_next_negative_sample,
              epoch_of_next_sample,
              n,
              densmap_flag,
              dens_phi_sum,
              dens_re_sum,
              dens_re_cov,
              dens_re_std,
              dens_re_mean,
              dens_lambda,
              dens_R,
              dens_mu,
              dens_mu_tot,
          )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
      return head_embedding
  
  
  @numba.njit(fastmath=True)
  def optimize_layout_generic(
      head_embedding,
      tail_embedding,
      head,
      tail,
      n_epochs,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma=1.0,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      output_metric=dist.euclidean,
      output_metric_kwds=(),
      verbose=False,
      move_other=False,
  ):
      """Improve an embedding using stochastic gradient descent to minimize the
      fuzzy set cross entropy between the 1-skeletons of the high dimensional
      and low dimensional fuzzy simplicial sets. In practice this is done by
      sampling edges based on their membership strength (with the (1-p) terms
      coming from negative sampling similar to word2vec).
  
      Parameters
      ----------
      head_embedding: array of shape (n_samples, n_components)
          The initial embedding to be improved by SGD.
  
      tail_embedding: array of shape (source_samples, n_components)
          The reference embedding of embedded points. If not embedding new
          previously unseen points with respect to an existing embedding this
          is simply the head_embedding (again); otherwise it provides the
          existing embedding to embed with respect to.
  
      head: array of shape (n_1_simplices)
          The indices of the heads of 1-simplices with non-zero membership.
  
      tail: array of shape (n_1_simplices)
          The indices of the tails of 1-simplices with non-zero membership.
  
      weight: array of shape (n_1_simplices)
          The membership weights of the 1-simplices.
  
      n_epochs: int
          The number of training epochs to use in optimization.
  
      n_vertices: int
          The number of vertices (0-simplices) in the dataset.
  
      epochs_per_sample: array of shape (n_1_simplices)
          A float value of the number of epochs per 1-simplex. 1-simplices with
          weaker membership strength will have more epochs between being sampled.
  
      a: float
          Parameter of differentiable approximation of right adjoint functor
  
      b: float
          Parameter of differentiable approximation of right adjoint functor
  
      rng_state: array of int64, shape (3,)
          The internal state of the rng
  
      gamma: float (optional, default 1.0)
          Weight to apply to negative samples.
  
      initial_alpha: float (optional, default 1.0)
          Initial learning rate for the SGD.
  
      negative_sample_rate: int (optional, default 5)
          Number of negative samples to use per positive sample.
  
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
  
      move_other: bool (optional, default False)
          Whether to adjust tail_embedding alongside head_embedding
  
      Returns
      -------
      embedding: array of shape (n_samples, n_components)
          The optimized embedding.
      """
  
      dim = head_embedding.shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = epochs_per_sample / negative_sample_rate
      epoch_of_next_negative_sample = epochs_per_negative_sample.copy()
      epoch_of_next_sample = epochs_per_sample.copy()
  
      for n in range(n_epochs):
          for i in range(epochs_per_sample.shape[0]):
              if epoch_of_next_sample[i] <= n:
                  j = head[i]
                  k = tail[i]
  
                  current = head_embedding[j]
                  other = tail_embedding[k]
  
                  dist_output, grad_dist_output = output_metric(
                      current, other, *output_metric_kwds
                  )
                  _, rev_grad_dist_output = output_metric(
                      other, current, *output_metric_kwds
                  )
  
                  if dist_output > 0.0:
                      w_l = pow((1 + a * pow(dist_output, 2 * b)), -1)
                  else:
                      w_l = 1.0
                  grad_coeff = 2 * b * (w_l - 1) / (dist_output + 1e-6)
  
                  for d in range(dim):
                      grad_d = clip(grad_coeff * grad_dist_output[d])
  
                      current[d] += grad_d * alpha
                      if move_other:
                          grad_d = clip(grad_coeff * rev_grad_dist_output[d])
                          other[d] += grad_d * alpha
  
                  epoch_of_next_sample[i] += epochs_per_sample[i]
  
                  n_neg_samples = int(
                      (n - epoch_of_next_negative_sample[i])
                      / epochs_per_negative_sample[i]
                  )
  
                  for p in range(n_neg_samples):
                      k = tau_rand_int(rng_state) % n_vertices
  
                      other = tail_embedding[k]
  
                      dist_output, grad_dist_output = output_metric(
                          current, other, *output_metric_kwds
                      )
  
                      if dist_output > 0.0:
                          w_l = pow((1 + a * pow(dist_output, 2 * b)), -1)
                      elif j == k:
                          continue
                      else:
                          w_l = 1.0
  
                      grad_coeff = gamma * 2 * b * w_l / (dist_output + 1e-6)
  
                      for d in range(dim):
                          grad_d = clip(grad_coeff * grad_dist_output[d])
                          current[d] += grad_d * alpha
  
                  epoch_of_next_negative_sample[i] += (
                      n_neg_samples * epochs_per_negative_sample[i]
                  )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
      return head_embedding
  
  
  @numba.njit(fastmath=True)
  def optimize_layout_inverse(
      head_embedding,
      tail_embedding,
      head,
      tail,
      weight,
      sigmas,
      rhos,
      n_epochs,
      n_vertices,
      epochs_per_sample,
      a,
      b,
      rng_state,
      gamma=1.0,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      output_metric=dist.euclidean,
      output_metric_kwds=(),
      verbose=False,
      move_other=False,
  ):
      """Improve an embedding using stochastic gradient descent to minimize the
      fuzzy set cross entropy between the 1-skeletons of the high dimensional
      and low dimensional fuzzy simplicial sets. In practice this is done by
      sampling edges based on their membership strength (with the (1-p) terms
      coming from negative sampling similar to word2vec).
  
      Parameters
      ----------
      head_embedding: array of shape (n_samples, n_components)
          The initial embedding to be improved by SGD.
  
      tail_embedding: array of shape (source_samples, n_components)
          The reference embedding of embedded points. If not embedding new
          previously unseen points with respect to an existing embedding this
          is simply the head_embedding (again); otherwise it provides the
          existing embedding to embed with respect to.
  
      head: array of shape (n_1_simplices)
          The indices of the heads of 1-simplices with non-zero membership.
  
      tail: array of shape (n_1_simplices)
          The indices of the tails of 1-simplices with non-zero membership.
  
      weight: array of shape (n_1_simplices)
          The membership weights of the 1-simplices.
  
      n_epochs: int
          The number of training epochs to use in optimization.
  
      n_vertices: int
          The number of vertices (0-simplices) in the dataset.
  
      epochs_per_sample: array of shape (n_1_simplices)
          A float value of the number of epochs per 1-simplex. 1-simplices with
          weaker membership strength will have more epochs between being sampled.
  
      a: float
          Parameter of differentiable approximation of right adjoint functor
  
      b: float
          Parameter of differentiable approximation of right adjoint functor
  
      rng_state: array of int64, shape (3,)
          The internal state of the rng
  
      gamma: float (optional, default 1.0)
          Weight to apply to negative samples.
  
      initial_alpha: float (optional, default 1.0)
          Initial learning rate for the SGD.
  
      negative_sample_rate: int (optional, default 5)
          Number of negative samples to use per positive sample.
  
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
  
      move_other: bool (optional, default False)
          Whether to adjust tail_embedding alongside head_embedding
  
      Returns
      -------
      embedding: array of shape (n_samples, n_components)
          The optimized embedding.
      """
  
      dim = head_embedding.shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = epochs_per_sample / negative_sample_rate
      epoch_of_next_negative_sample = epochs_per_negative_sample.copy()
      epoch_of_next_sample = epochs_per_sample.copy()
  
      for n in range(n_epochs):
          for i in range(epochs_per_sample.shape[0]):
              if epoch_of_next_sample[i] <= n:
                  j = head[i]
                  k = tail[i]
  
                  current = head_embedding[j]
                  other = tail_embedding[k]
  
                  dist_output, grad_dist_output = output_metric(
                      current, other, *output_metric_kwds
                  )
  
                  w_l = weight[i]
                  grad_coeff = -(1 / (w_l * sigmas[k] + 1e-6))
  
                  for d in range(dim):
                      grad_d = clip(grad_coeff * grad_dist_output[d])
  
                      current[d] += grad_d * alpha
                      if move_other:
                          other[d] += -grad_d * alpha
  
                  epoch_of_next_sample[i] += epochs_per_sample[i]
  
                  n_neg_samples = int(
                      (n - epoch_of_next_negative_sample[i])
                      / epochs_per_negative_sample[i]
                  )
  
                  for p in range(n_neg_samples):
                      k = tau_rand_int(rng_state) % n_vertices
  
                      other = tail_embedding[k]
  
                      dist_output, grad_dist_output = output_metric(
                          current, other, *output_metric_kwds
                      )
  
                      # w_l = 0.0 # for negative samples, the edge does not exist
                      w_h = np.exp(-max(dist_output - rhos[k], 1e-6) / (sigmas[k] + 1e-6))
                      grad_coeff = -gamma * ((0 - w_h) / ((1 - w_h) * sigmas[k] + 1e-6))
  
                      for d in range(dim):
                          grad_d = clip(grad_coeff * grad_dist_output[d])
                          current[d] += grad_d * alpha
  
                  epoch_of_next_negative_sample[i] += (
                      n_neg_samples * epochs_per_negative_sample[i]
                  )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
      return head_embedding
  
  
  def _optimize_layout_aligned_euclidean_single_epoch(
      head_embeddings,
      tail_embeddings,
      heads,
      tails,
      epochs_per_sample,
      a,
      b,
      regularisation_weights,
      relations,
      rng_state,
      gamma,
      lambda_,
      dim,
      move_other,
      alpha,
      epochs_per_negative_sample,
      epoch_of_next_negative_sample,
      epoch_of_next_sample,
      n,
  ):
      n_embeddings = len(heads)
      window_size = (relations.shape[1] - 1) // 2
  
      max_n_edges = 0
      for e_p_s in epochs_per_sample:
          if e_p_s.shape[0] >= max_n_edges:
              max_n_edges = e_p_s.shape[0]
  
      embedding_order = np.arange(n_embeddings).astype(np.int32)
      np.random.seed(abs(rng_state[0]))
      np.random.shuffle(embedding_order)
  
      for i in range(max_n_edges):
          for m in embedding_order:
              if i < epoch_of_next_sample[m].shape[0] and epoch_of_next_sample[m][i] <= n:
                  j = heads[m][i]
                  k = tails[m][i]
  
                  current = head_embeddings[m][j]
                  other = tail_embeddings[m][k]
  
                  dist_squared = rdist(current, other)
  
                  if dist_squared > 0.0:
                      grad_coeff = -2.0 * a * b * pow(dist_squared, b - 1.0)
                      grad_coeff /= a * pow(dist_squared, b) + 1.0
                  else:
                      grad_coeff = 0.0
  
                  for d in range(dim):
                      grad_d = clip(grad_coeff * (current[d] - other[d]))
  
                      for offset in range(-window_size, window_size):
                          neighbor_m = m + offset
                          if (
                              neighbor_m >= 0
                              and neighbor_m < n_embeddings
                              and offset != 0
                          ):
                              identified_index = relations[m, offset + window_size, j]
                              if identified_index >= 0:
                                  grad_d -= clip(
                                      (lambda_ * np.exp(-(np.abs(offset) - 1)))
                                      * regularisation_weights[m, offset + window_size, j]
                                      * (
                                          current[d]
                                          - head_embeddings[neighbor_m][
                                              identified_index, d
                                          ]
                                      )
                                  )
  
                      current[d] += clip(grad_d) * alpha
                      if move_other:
                          other_grad_d = clip(grad_coeff * (other[d] - current[d]))
  
                          for offset in range(-window_size, window_size):
                              neighbor_m = m + offset
                              if (
                                  neighbor_m >= 0
                                  and neighbor_m < n_embeddings
                                  and offset != 0
                              ):
                                  identified_index = relations[m, offset + window_size, k]
                                  if identified_index >= 0:
                                      grad_d -= clip(
                                          (lambda_ * np.exp(-(np.abs(offset) - 1)))
                                          * regularisation_weights[
                                              m, offset + window_size, k
                                          ]
                                          * (
                                              other[d]
                                              - head_embeddings[neighbor_m][
                                                  identified_index, d
                                              ]
                                          )
                                      )
  
                          other[d] += clip(other_grad_d) * alpha
  
                  epoch_of_next_sample[m][i] += epochs_per_sample[m][i]
  
                  if epochs_per_negative_sample[m][i] > 0:
                      n_neg_samples = int(
                          (n - epoch_of_next_negative_sample[m][i])
                          / epochs_per_negative_sample[m][i]
                      )
                  else:
                      n_neg_samples = 0
  
                  for p in range(n_neg_samples):
                      k = tau_rand_int(rng_state) % tail_embeddings[m].shape[0]
  
                      other = tail_embeddings[m][k]
  
                      dist_squared = rdist(current, other)
  
                      if dist_squared > 0.0:
                          grad_coeff = 2.0 * gamma * b
                          grad_coeff /= (0.001 + dist_squared) * (
                              a * pow(dist_squared, b) + 1
                          )
                      elif j == k:
                          continue
                      else:
                          grad_coeff = 0.0
  
                      for d in range(dim):
                          if grad_coeff > 0.0:
                              grad_d = clip(grad_coeff * (current[d] - other[d]))
                          else:
                              grad_d = 4.0
  
                          for offset in range(-window_size, window_size):
                              neighbor_m = m + offset
                              if (
                                  neighbor_m >= 0
                                  and neighbor_m < n_embeddings
                                  and offset != 0
                              ):
                                  identified_index = relations[m, offset + window_size, j]
                                  if identified_index >= 0:
                                      grad_d -= clip(
                                          (lambda_ * np.exp(-(np.abs(offset) - 1)))
                                          * regularisation_weights[
                                              m, offset + window_size, j
                                          ]
                                          * (
                                              current[d]
                                              - head_embeddings[neighbor_m][
                                                  identified_index, d
                                              ]
                                          )
                                      )
  
                          current[d] += clip(grad_d) * alpha
  
                  epoch_of_next_negative_sample[m][i] += (
                      n_neg_samples * epochs_per_negative_sample[m][i]
                  )
  
  
  def optimize_layout_aligned_euclidean(
      head_embeddings,
      tail_embeddings,
      heads,
      tails,
      n_epochs,
      epochs_per_sample,
      regularisation_weights,
      relations,
      rng_state,
      a=1.576943460405378,
      b=0.8950608781227859,
      gamma=1.0,
      lambda_=5e-3,
      initial_alpha=1.0,
      negative_sample_rate=5.0,
      parallel=True,
      verbose=False,
      move_other=False,
  ):
      dim = head_embeddings[0].shape[1]
      alpha = initial_alpha
  
      epochs_per_negative_sample = numba.typed.List.empty_list(numba.types.float32[::1])
      epoch_of_next_negative_sample = numba.typed.List.empty_list(
          numba.types.float32[::1]
      )
      epoch_of_next_sample = numba.typed.List.empty_list(numba.types.float32[::1])
  
      for m in range(len(heads)):
          epochs_per_negative_sample.append(
              epochs_per_sample[m].astype(np.float32) / negative_sample_rate
          )
          epoch_of_next_negative_sample.append(
              epochs_per_negative_sample[m].astype(np.float32)
          )
          epoch_of_next_sample.append(epochs_per_sample[m].astype(np.float32))
  
      optimize_fn = numba.njit(
          _optimize_layout_aligned_euclidean_single_epoch,
          fastmath=True,
          parallel=parallel,
      )
  
      for n in range(n_epochs):
          optimize_fn(
              head_embeddings,
              tail_embeddings,
              heads,
              tails,
              epochs_per_sample,
              a,
              b,
              regularisation_weights,
              relations,
              rng_state,
              gamma,
              lambda_,
              dim,
              move_other,
              alpha,
              epochs_per_negative_sample,
              epoch_of_next_negative_sample,
              epoch_of_next_sample,
              n,
          )
  
          alpha = initial_alpha * (1.0 - (float(n) / float(n_epochs)))
  
          if verbose and n % int(n_epochs / 10) == 0:
              print("\tcompleted ", n, " / ", n_epochs, "epochs")
  
      return head_embeddings
  
  
  [... the rest is omitted ...]
files: umap/layouts.py
files: # Author: Leland McInnes <leland.mcinnes@gmail.com>
  #
  # License: BSD 3 clause
  from __future__ import print_function
  
  import locale
  from warnings import warn
  import time
  
  from scipy.optimize import curve_fit
  from sklearn.base import BaseEstimator
  from sklearn.utils import check_random_state, check_array
  from sklearn.utils.validation import check_is_fitted
  from sklearn.metrics import pairwise_distances
  from sklearn.preprocessing import normalize
  from sklearn.neighbors import KDTree
  
  try:
      import joblib
  except ImportError:
      # sklearn.externals.joblib is deprecated in 0.21, will be removed in 0.23
      from sklearn.externals import joblib
  
  import numpy as np
  import scipy.sparse
  from scipy.sparse import tril as sparse_tril, triu as sparse_triu
  import scipy.sparse.csgraph
  import numba
  
  import umap.distances as dist
  
  import umap.sparse as sparse
  
  from umap.utils import (
      submatrix,
      ts,
      csr_unique,
      fast_knn_indices,
  )
  from umap.spectral import spectral_layout
  from umap.layouts import (
      optimize_layout_euclidean,
      optimize_layout_generic,
      optimize_layout_inverse,
  )
  
  from pynndescent import NNDescent
  from pynndescent.distances import named_distances as pynn_named_distances
  from pynndescent.sparse import sparse_named_distances as pynn_sparse_named_distances
  
  locale.setlocale(locale.LC_NUMERIC, "C")
  
  INT32_MIN = np.iinfo(np.int32).min + 1
  INT32_MAX = np.iinfo(np.int32).max - 1
  
  SMOOTH_K_TOLERANCE = 1e-5
  MIN_K_DIST_SCALE = 1e-3
  NPY_INFINITY = np.inf
  
  DISCONNECTION_DISTANCES = {
      "correlation": 2,
      "cosine": 2,
      "hellinger": 1,
      "jaccard": 1,
      "dice": 1,
  }
  
  
  def flatten_iter(container):
      for i in container:
          if isinstance(i, (list, tuple)):
              for j in flatten_iter(i):
                  yield j
          else:
              yield i
  
  
  def flattened(container):
      return tuple(flatten_iter(container))
  
  
  def breadth_first_search(adjmat, start, min_vertices):
      explored = []
      queue = [start]
      levels = {}
      levels[start] = 0
      max_level = np.inf
      visited = [start]
  
      while queue:
          node = queue.pop(0)
          explored.append(node)
          if max_level == np.inf and len(explored) > min_vertices:
              max_level = max(levels.values())
  
          if levels[node] + 1 < max_level:
              neighbors = adjmat[node].indices
              for neighbour in neighbors:
                  if neighbour not in visited:
                      queue.append(neighbour)
                      visited.append(neighbour)
  
                      levels[neighbour] = levels[node] + 1
  
      return np.array(explored)
  
  
  def raise_disconnected_warning(
      edges_removed,
      vertices_disconnected,
      disconnection_distance,
      total_rows,
      threshold=0.1,
      verbose=False,
  ):
      """A simple wrapper function to avoid large amounts of code repetition."""
      if verbose & (vertices_disconnected == 0) & (edges_removed > 0):
          print(
              f"Disconnection_distance = {disconnection_distance} has removed {edges_removed} edges.  "
              f"This is not a problem as no vertices were disconnected."
          )
      elif (vertices_disconnected > 0) & (
          vertices_disconnected <= threshold * total_rows
      ):
          warn(
              f"A few of your vertices were disconnected from the manifold.  This shouldn't cause problems.\n"
              f"Disconnection_distance = {disconnection_distance} has removed {edges_removed} edges.\n"
              f"It has only fully disconnected {vertices_disconnected} vertices.\n"
              f"Use umap.utils.disconnected_vertices() to identify them.",
          )
      elif vertices_disconnected > threshold * total_rows:
          warn(
              f"A large number of your vertices were disconnected from the manifold.\n"
              f"Disconnection_distance = {disconnection_distance} has removed {edges_removed} edges.\n"
              f"It has fully disconnected {vertices_disconnected} vertices.\n"
              f"You might consider using find_disconnected_points() to find and remove these points from your data.\n"
              f"Use umap.utils.disconnected_vertices() to identify them.",
          )
  
  
  @numba.njit(
      locals={
          "psum": numba.types.float32,
          "lo": numba.types.float32,
          "mid": numba.types.float32,
          "hi": numba.types.float32,
      },
      fastmath=True,
  )  # benchmarking `parallel=True` shows it to *decrease* performance
  def smooth_knn_dist(distances, k, n_iter=64, local_connectivity=1.0, bandwidth=1.0):
      """Compute a continuous version of the distance to the kth nearest
      neighbor. That is, this is similar to knn-distance but allows continuous
      k values rather than requiring an integral k. In essence we are simply
      computing the distance such that the cardinality of fuzzy set we generate
      is k.
  
      Parameters
      ----------
      distances: array of shape (n_samples, n_neighbors)
          Distances to nearest neighbors for each samples. Each row should be a
          sorted list of distances to a given samples nearest neighbors.
  
      k: float
          The number of nearest neighbors to approximate for.
  
      n_iter: int (optional, default 64)
          We need to binary search for the correct distance value. This is the
          max number of iterations to use in such a search.
  
      local_connectivity: int (optional, default 1)
          The local connectivity required -- i.e. the number of nearest
          neighbors that should be assumed to be connected at a local level.
          The higher this value the more connected the manifold becomes
          locally. In practice this should be not more than the local intrinsic
          dimension of the manifold.
  
      bandwidth: float (optional, default 1)
          The target bandwidth of the kernel, larger values will produce
          larger return values.
  
      Returns
      -------
      knn_dist: array of shape (n_samples,)
          The distance to kth nearest neighbor, as suitably approximated.
  
      nn_dist: array of shape (n_samples,)
          The distance to the 1st nearest neighbor for each point.
      """
      target = np.log2(k) * bandwidth
      rho = np.zeros(distances.shape[0], dtype=np.float32)
      result = np.zeros(distances.shape[0], dtype=np.float32)
  
      mean_distances = np.mean(distances)
  
      for i in range(distances.shape[0]):
          lo = 0.0
          hi = NPY_INFINITY
          mid = 1.0
  
          # TODO: This is very inefficient, but will do for now. FIXME
          ith_distances = distances[i]
          non_zero_dists = ith_distances[ith_distances > 0.0]
          if non_zero_dists.shape[0] >= local_connectivity:
              index = int(np.floor(local_connectivity))
              interpolation = local_connectivity - index
              if index > 0:
                  rho[i] = non_zero_dists[index - 1]
                  if interpolation > SMOOTH_K_TOLERANCE:
                      rho[i] += interpolation * (
                          non_zero_dists[index] - non_zero_dists[index - 1]
                      )
              else:
                  rho[i] = interpolation * non_zero_dists[0]
          elif non_zero_dists.shape[0] > 0:
              rho[i] = np.max(non_zero_dists)
  
          for n in range(n_iter):
  
              psum = 0.0
              for j in range(1, distances.shape[1]):
                  d = distances[i, j] - rho[i]
                  if d > 0:
                      psum += np.exp(-(d / mid))
                  else:
                      psum += 1.0
  
              if np.fabs(psum - target) < SMOOTH_K_TOLERANCE:
                  break
  
              if psum > target:
                  hi = mid
                  mid = (lo + hi) / 2.0
              else:
                  lo = mid
                  if hi == NPY_INFINITY:
                      mid *= 2
                  else:
                      mid = (lo + hi) / 2.0
  
          result[i] = mid
  
          # TODO: This is very inefficient, but will do for now. FIXME
          if rho[i] > 0.0:
              mean_ith_distances = np.mean(ith_distances)
              if result[i] < MIN_K_DIST_SCALE * mean_ith_distances:
                  result[i] = MIN_K_DIST_SCALE * mean_ith_distances
          else:
              if result[i] < MIN_K_DIST_SCALE * mean_distances:
                  result[i] = MIN_K_DIST_SCALE * mean_distances
  
      return result, rho
  
  
  def nearest_neighbors(
      X,
      n_neighbors,
      metric,
      metric_kwds,
      angular,
      random_state,
      low_memory=True,
      use_pynndescent=True,
      n_jobs=-1,
      verbose=False,
  ):
      """Compute the ``n_neighbors`` nearest points for each data point in ``X``
      under ``metric``. This may be exact, but more likely is approximated via
      nearest neighbor descent.
  
      Parameters
      ----------
      X: array of shape (n_samples, n_features)
          The input data to compute the k-neighbor graph of.
  
      n_neighbors: int
          The number of nearest neighbors to compute for each sample in ``X``.
  
      metric: string or callable
          The metric to use for the computation.
  
      metric_kwds: dict
          Any arguments to pass to the metric computation function.
  
      angular: bool
          Whether to use angular rp trees in NN approximation.
  
      random_state: np.random state
          The random state to use for approximate NN computations.
  
      low_memory: bool (optional, default True)
          Whether to pursue lower memory NNdescent.
  
      verbose: bool (optional, default False)
          Whether to print status data during the computation.
  
      Returns
      -------
      knn_indices: array of shape (n_samples, n_neighbors)
          The indices on the ``n_neighbors`` closest points in the dataset.
  
      knn_dists: array of shape (n_samples, n_neighbors)
          The distances to the ``n_neighbors`` closest points in the dataset.
  
      rp_forest: list of trees
          The random projection forest used for searching (if used, None otherwise)
      """
      if verbose:
          print(ts(), "Finding Nearest Neighbors")
  
      if metric == "precomputed":
          # Note that this does not support sparse distance matrices yet ...
          # Compute indices of n nearest neighbors
          knn_indices = fast_knn_indices(X, n_neighbors)
          # knn_indices = np.argsort(X)[:, :n_neighbors]
          # Compute the nearest neighbor distances
          #   (equivalent to np.sort(X)[:,:n_neighbors])
          knn_dists = X[np.arange(X.shape[0])[:, None], knn_indices].copy()
          # Prune any nearest neighbours that are infinite distance apart.
          disconnected_index = knn_dists == np.inf
          knn_indices[disconnected_index] = -1
  
          knn_search_index = None
      else:
          # TODO: Hacked values for now
          n_trees = min(64, 5 + int(round((X.shape[0]) ** 0.5 / 20.0)))
          n_iters = max(5, int(round(np.log2(X.shape[0]))))
  
          knn_search_index = NNDescent(
              X,
              n_neighbors=n_neighbors,
              metric=metric,
              metric_kwds=metric_kwds,
              random_state=random_state,
              n_trees=n_trees,
              n_iters=n_iters,
              max_candidates=60,
              low_memory=low_memory,
              n_jobs=n_jobs,
              verbose=verbose,
          )
          knn_indices, knn_dists = knn_search_index.neighbor_graph
  
      if verbose:
          print(ts(), "Finished Nearest Neighbor Search")
      return knn_indices, knn_dists, knn_search_index
  
  
  @numba.njit(
      locals={
          "knn_dists": numba.types.float32[:, ::1],
          "sigmas": numba.types.float32[::1],
          "rhos": numba.types.float32[::1],
          "val": numba.types.float32,
      },
      parallel=True,
      fastmath=True,
  )
  def compute_membership_strengths(
      knn_indices,
      knn_dists,
      sigmas,
      rhos,
      return_dists=False,
      bipartite=False,
  ):
      """Construct the membership strength data for the 1-skeleton of each local
      fuzzy simplicial set -- this is formed as a sparse matrix where each row is
      a local fuzzy simplicial set, with a membership strength for the
      1-simplex to each other data point.
  
      Parameters
      ----------
      knn_indices: array of shape (n_samples, n_neighbors)
          The indices on the ``n_neighbors`` closest points in the dataset.
  
      knn_dists: array of shape (n_samples, n_neighbors)
          The distances to the ``n_neighbors`` closest points in the dataset.
  
      sigmas: array of shape(n_samples)
          The normalization factor derived from the metric tensor approximation.
  
      rhos: array of shape(n_samples)
          The local connectivity adjustment.
  
      return_dists: bool (optional, default False)
          Whether to return the pairwise distance associated with each edge
  
      bipartite: bool (optional, default False)
          Does the nearest neighbour set represent a bipartite graph?  That is are the
          nearest neighbour indices from the same point set as the row indices?
  
      Returns
      -------
      rows: array of shape (n_samples * n_neighbors)
          Row data for the resulting sparse matrix (coo format)
  
      cols: array of shape (n_samples * n_neighbors)
          Column data for the resulting sparse matrix (coo format)
  
      vals: array of shape (n_samples * n_neighbors)
          Entries for the resulting sparse matrix (coo format)
  
      dists: array of shape (n_samples * n_neighbors)
          Distance associated with each entry in the resulting sparse matrix
      """
      n_samples = knn_indices.shape[0]
      n_neighbors = knn_indices.shape[1]
  
      rows = np.zeros(knn_indices.size, dtype=np.int32)
      cols = np.zeros(knn_indices.size, dtype=np.int32)
      vals = np.zeros(knn_indices.size, dtype=np.float32)
      if return_dists:
          dists = np.zeros(knn_indices.size, dtype=np.float32)
      else:
          dists = None
  
      for i in range(n_samples):
          for j in range(n_neighbors):
              if knn_indices[i, j] == -1:
                  continue  # We didn't get the full knn for i
              # If applied to an adjacency matrix points shouldn't be similar to themselves.
              # If applied to an incidence matrix (or bipartite) then the row and column indices are different.
              if (bipartite == False) & (knn_indices[i, j] == i):
                  val = 0.0
              elif knn_dists[i, j] - rhos[i] <= 0.0 or sigmas[i] == 0.0:
                  val = 1.0
              else:
                  val = np.exp(-((knn_dists[i, j] - rhos[i]) / (sigmas[i])))
  
              rows[i * n_neighbors + j] = i
              cols[i * n_neighbors + j] = knn_indices[i, j]
              vals[i * n_neighbors + j] = val
              if return_dists:
                  dists[i * n_neighbors + j] = knn_dists[i, j]
  
      return rows, cols, vals, dists
  
  
  def fuzzy_simplicial_set(
      X,
      n_neighbors,
      random_state,
      metric,
      metric_kwds={},
      knn_indices=None,
      knn_dists=None,
      angular=False,
      set_op_mix_ratio=1.0,
      local_connectivity=1.0,
      apply_set_operations=True,
      verbose=False,
      return_dists=None,
  ):
      """Given a set of data X, a neighborhood size, and a measure of distance
      compute the fuzzy simplicial set (here represented as a fuzzy graph in
      the form of a sparse matrix) associated to the data. This is done by
      locally approximating geodesic distance at each point, creating a fuzzy
      simplicial set for each such point, and then combining all the local
      fuzzy simplicial sets into a global one via a fuzzy union.
  
      Parameters
      ----------
      X: array of shape (n_samples, n_features)
          The data to be modelled as a fuzzy simplicial set.
  
      n_neighbors: int
          The number of neighbors to use to approximate geodesic distance.
          Larger numbers induce more global estimates of the manifold that can
          miss finer detail, while smaller values will focus on fine manifold
          structure to the detriment of the larger picture.
  
      random_state: numpy RandomState or equivalent
          A state capable being used as a numpy random state.
  
      metric: string or function (optional, default 'euclidean')
          The metric to use to compute distances in high dimensional space.
          If a string is passed it must match a valid predefined metric. If
          a general metric is required a function that takes two 1d arrays and
          returns a float can be provided. For performance purposes it is
          required that this be a numba jit'd function. Valid string metrics
          include:
              * euclidean (or l2)
              * manhattan (or l1)
              * cityblock
              * braycurtis
              * canberra
              * chebyshev
              * correlation
              * cosine
              * dice
              * hamming
              * jaccard
              * kulsinski
              * ll_dirichlet
              * mahalanobis
              * matching
              * minkowski
              * rogerstanimoto
              * russellrao
              * seuclidean
              * sokalmichener
              * sokalsneath
              * sqeuclidean
              * yule
              * wminkowski
  
          Metrics that take arguments (such as minkowski, mahalanobis etc.)
          can have arguments passed via the metric_kwds dictionary. At this
          time care must be taken and dictionary elements must be ordered
          appropriately; this will hopefully be fixed in the future.
  
      metric_kwds: dict (optional, default {})
          Arguments to pass on to the metric, such as the ``p`` value for
          Minkowski distance.
  
      knn_indices: array of shape (n_samples, n_neighbors) (optional)
          If the k-nearest neighbors of each point has already been calculated
          you can pass them in here to save computation time. This should be
          an array with the indices of the k-nearest neighbors as a row for
          each data point.
  
      knn_dists: array of shape (n_samples, n_neighbors) (optional)
          If the k-nearest neighbors of each point has already been calculated
          you can pass them in here to save computation time. This should be
          an array with the distances of the k-nearest neighbors as a row for
          each data point.
  
      angular: bool (optional, default False)
          Whether to use angular/cosine distance for the random projection
          forest for seeding NN-descent to determine approximate nearest
          neighbors.
  
      set_op_mix_ratio: float (optional, default 1.0)
          Interpolate between (fuzzy) union and intersection as the set operation
          used to combine local fuzzy simplicial sets to obtain a global fuzzy
          simplicial sets. Both fuzzy set operations use the product t-norm.
          The value of this parameter should be between 0.0 and 1.0; a value of
          1.0 will use a pure fuzzy union, while 0.0 will use a pure fuzzy
          intersection.
  
      local_connectivity: int (optional, default 1)
          The local connectivity required -- i.e. the number of nearest
          neighbors that should be assumed to be connected at a local level.
          The higher this value the more connected the manifold becomes
          locally. In practice this should be not more than the local intrinsic
          dimension of the manifold.
  
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
  
      return_dists: bool or None (optional, default None)
          Whether to return the pairwise distance associated with each edge.
  
      Returns
      -------
      fuzzy_simplicial_set: coo_matrix
          A fuzzy simplicial set represented as a sparse matrix. The (i,
          j) entry of the matrix represents the membership strength of the
          1-simplex between the ith and jth sample points.
      """
      if knn_indices is None or knn_dists is None:
          knn_indices, knn_dists, _ = nearest_neighbors(
              X,
              n_neighbors,
              metric,
              metric_kwds,
              angular,
              random_state,
              verbose=verbose,
          )
  
      knn_dists = knn_dists.astype(np.float32)
  
      sigmas, rhos = smooth_knn_dist(
          knn_dists,
          float(n_neighbors),
          local_connectivity=float(local_connectivity),
      )
  
      rows, cols, vals, dists = compute_membership_strengths(
          knn_indices, knn_dists, sigmas, rhos, return_dists
      )
  
      result = scipy.sparse.coo_matrix(
          (vals, (rows, cols)), shape=(X.shape[0], X.shape[0])
      )
      result.eliminate_zeros()
  
      if apply_set_operations:
          transpose = result.transpose()
  
          prod_matrix = result.multiply(transpose)
  
          result = (
              set_op_mix_ratio * (result + transpose - prod_matrix)
              + (1.0 - set_op_mix_ratio) * prod_matrix
          )
  
      result.eliminate_zeros()
  
      if return_dists is None:
          return result, sigmas, rhos
      else:
          if return_dists:
              dmat = scipy.sparse.coo_matrix(
                  (dists, (rows, cols)), shape=(X.shape[0], X.shape[0])
              )
  
              dists = dmat.maximum(dmat.transpose()).todok()
          else:
              dists = None
  
          return result, sigmas, rhos, dists
  
  
  @numba.njit()
  def fast_intersection(rows, cols, values, target, unknown_dist=1.0, far_dist=5.0):
      """Under the assumption of categorical distance for the intersecting
      simplicial set perform a fast intersection.
  
      Parameters
      ----------
      rows: array
          An array of the row of each non-zero in the sparse matrix
          representation.
  
      cols: array
          An array of the column of each non-zero in the sparse matrix
          representation.
  
      values: array
          An array of the value of each non-zero in the sparse matrix
          representation.
  
      target: array of shape (n_samples)
          The categorical labels to use in the intersection.
  
      unknown_dist: float (optional, default 1.0)
          The distance an unknown label (-1) is assumed to be from any point.
  
      far_dist float (optional, default 5.0)
          The distance between unmatched labels.
  
      Returns
      -------
      None
      """
      for nz in range(rows.shape[0]):
          i = rows[nz]
          j = cols[nz]
          if (target[i] == -1) or (target[j] == -1):
              values[nz] *= np.exp(-unknown_dist)
          elif target[i] != target[j]:
              values[nz] *= np.exp(-far_dist)
  
      return
  
  
  @numba.jit()
  def fast_metric_intersection(
      rows, cols, values, discrete_space, metric, metric_args, scale
  ):
      """Under the assumption of categorical distance for the intersecting
      simplicial set perform a fast intersection.
  
      Parameters
      ----------
      rows: array
          An array of the row of each non-zero in the sparse matrix
          representation.
  
      cols: array
          An array of the column of each non-zero in the sparse matrix
          representation.
  
      values: array of shape
          An array of the values of each non-zero in the sparse matrix
          representation.
  
      discrete_space: array of shape (n_samples, n_features)
          The vectors of categorical labels to use in the intersection.
  
      metric: numba function
          The function used to calculate distance over the target array.
  
      scale: float
          A scaling to apply to the metric.
  
      Returns
      -------
      None
      """
      for nz in range(rows.shape[0]):
          i = rows[nz]
          j = cols[nz]
          dist = metric(discrete_space[i], discrete_space[j], *metric_args)
          values[nz] *= np.exp(-(scale * dist))
  
      return
  
  
  @numba.njit()
  def reprocess_row(probabilities, k=15, n_iters=32):
      target = np.log2(k)
  
      lo = 0.0
      hi = NPY_INFINITY
      mid = 1.0
  
      for n in range(n_iters):
  
          psum = 0.0
          for j in range(probabilities.shape[0]):
              psum += pow(probabilities[j], mid)
  
          if np.fabs(psum - target) < SMOOTH_K_TOLERANCE:
              break
  
          if psum < target:
              hi = mid
              mid = (lo + hi) / 2.0
          else:
              lo = mid
              if hi == NPY_INFINITY:
                  mid *= 2
              else:
                  mid = (lo + hi) / 2.0
  
      return np.power(probabilities, mid)
  
  
  @numba.njit()
  def reset_local_metrics(simplicial_set_indptr, simplicial_set_data):
      for i in range(simplicial_set_indptr.shape[0] - 1):
          simplicial_set_data[
              simplicial_set_indptr[i] : simplicial_set_indptr[i + 1]
          ] = reprocess_row(
              simplicial_set_data[simplicial_set_indptr[i] : simplicial_set_indptr[i + 1]]
          )
      return
  
  
  def reset_local_connectivity(simplicial_set, reset_local_metric=False):
      """Reset the local connectivity requirement -- each data sample should
      have complete confidence in at least one 1-simplex in the simplicial set.
      We can enforce this by locally rescaling confidences, and then remerging the
      different local simplicial sets together.
  
      Parameters
      ----------
      simplicial_set: sparse matrix
          The simplicial set for which to recalculate with respect to local
          connectivity.
  
      Returns
      -------
      simplicial_set: sparse_matrix
          The recalculated simplicial set, now with the local connectivity
          assumption restored.
      """
      simplicial_set = normalize(simplicial_set, norm="max")
      if reset_local_metric:
          simplicial_set = simplicial_set.tocsr()
          reset_local_metrics(simplicial_set.indptr, simplicial_set.data)
          simplicial_set = simplicial_set.tocoo()
      transpose = simplicial_set.transpose()
      prod_matrix = simplicial_set.multiply(transpose)
      simplicial_set = simplicial_set + transpose - prod_matrix
      simplicial_set.eliminate_zeros()
  
      return simplicial_set
  
  
  def discrete_metric_simplicial_set_intersection(
      simplicial_set,
      discrete_space,
      unknown_dist=1.0,
      far_dist=5.0,
      metric=None,
      metric_kws={},
      metric_scale=1.0,
  ):
      """Combine a fuzzy simplicial set with another fuzzy simplicial set
      generated from discrete metric data using discrete distances. The target
      data is assumed to be categorical label data (a vector of labels),
      and this will update the fuzzy simplicial set to respect that label data.
  
      TODO: optional category cardinality based weighting of distance
  
      Parameters
      ----------
      simplicial_set: sparse matrix
          The input fuzzy simplicial set.
  
      discrete_space: array of shape (n_samples)
          The categorical labels to use in the intersection.
  
      unknown_dist: float (optional, default 1.0)
          The distance an unknown label (-1) is assumed to be from any point.
  
      far_dist: float (optional, default 5.0)
          The distance between unmatched labels.
  
      metric: str (optional, default None)
          If not None, then use this metric to determine the
          distance between values.
  
      metric_scale: float (optional, default 1.0)
          If using a custom metric scale the distance values by
          this value -- this controls the weighting of the
          intersection. Larger values weight more toward target.
  
      Returns
      -------
      simplicial_set: sparse matrix
          The resulting intersected fuzzy simplicial set.
      """
      simplicial_set = simplicial_set.tocoo()
  
      if metric is not None:
          # We presume target is now a 2d array, with each row being a
          # vector of target info
          if metric in dist.named_distances:
              metric_func = dist.named_distances[metric]
          else:
              raise ValueError("Discrete intersection metric is not recognized")
  
          fast_metric_intersection(
              simplicial_set.row,
              simplicial_set.col,
              simplicial_set.data,
              discrete_space,
              metric_func,
              tuple(metric_kws.values()),
              metric_scale,
          )
      else:
          fast_intersection(
              simplicial_set.row,
              simplicial_set.col,
              simplicial_set.data,
              discrete_space,
              unknown_dist,
              far_dist,
          )
  
      simplicial_set.eliminate_zeros()
  
      return reset_local_connectivity(simplicial_set)
  
  
  def general_simplicial_set_intersection(
      simplicial_set1, simplicial_set2, weight=0.5, right_complement=False
  ):
  
      if right_complement:
          result = simplicial_set1.tocoo()
      else:
          result = (simplicial_set1 + simplicial_set2).tocoo()
      left = simplicial_set1.tocsr()
      right = simplicial_set2.tocsr()
  
      sparse.general_sset_intersection(
          left.indptr,
          left.indices,
          left.data,
          right.indptr,
          right.indices,
          right.data,
          result.row,
          result.col,
          result.data,
          mix_weight=weight,
          right_complement=right_complement,
      )
  
      return result
  
  
  def general_simplicial_set_union(simplicial_set1, simplicial_set2):
      result = (simplicial_set1 + simplicial_set2).tocoo()
      left = simplicial_set1.tocsr()
      right = simplicial_set2.tocsr()
  
      sparse.general_sset_union(
          left.indptr,
          left.indices,
          left.data,
          right.indptr,
          right.indices,
          right.data,
          result.row,
          result.col,
          result.data,
      )
  
      return result
  
  
  def make_epochs_per_sample(weights, n_epochs):
      """Given a set of weights and number of epochs generate the number of
      epochs per sample for each weight.
  
      Parameters
      ----------
      weights: array of shape (n_1_simplices)
          The weights ofhow much we wish to sample each 1-simplex.
  
      n_epochs: int
          The total number of epochs we want to train for.
  
      Returns
      -------
      An array of number of epochs per sample, one for each 1-simplex.
      """
      result = -1.0 * np.ones(weights.shape[0], dtype=np.float64)
      n_samples = n_epochs * (weights / weights.max())
      result[n_samples > 0] = float(n_epochs) / n_samples[n_samples > 0]
      return result
  
  
  def simplicial_set_embedding(
      data,
      graph,
      n_components,
      initial_alpha,
      a,
      b,
      gamma,
      negative_sample_rate,
      n_epochs,
      init,
      random_state,
      metric,
      metric_kwds,
      densmap,
      densmap_kwds,
      output_dens,
      output_metric=dist.named_distances_with_gradients["euclidean"],
      output_metric_kwds={},
      euclidean_output=True,
      parallel=False,
      verbose=False,
  ):
      """Perform a fuzzy simplicial set embedding, using a specified
      initialisation method and then minimizing the fuzzy set cross entropy
      between the 1-skeletons of the high and low dimensional fuzzy simplicial
      sets.
  
      Parameters
      ----------
      data: array of shape (n_samples, n_features)
          The source data to be embedded by UMAP.
  
      graph: sparse matrix
          The 1-skeleton of the high dimensional fuzzy simplicial set as
          represented by a graph for which we require a sparse matrix for the
          (weighted) adjacency matrix.
  
      n_components: int
          The dimensionality of the euclidean space into which to embed the data.
  
      initial_alpha: float
          Initial learning rate for the SGD.
  
      a: float
          Parameter of differentiable approximation of right adjoint functor
  
      b: float
          Parameter of differentiable approximation of right adjoint functor
  
      gamma: float
          Weight to apply to negative samples.
  
      negative_sample_rate: int (optional, default 5)
          The number of negative samples to select per positive sample
          in the optimization process. Increasing this value will result
          in greater repulsive force being applied, greater optimization
          cost, but slightly more accuracy.
  
      n_epochs: int (optional, default 0), or list of int
          The number of training epochs to be used in optimizing the
          low dimensional embedding. Larger values result in more accurate
          embeddings. If 0 is specified a value will be selected based on
          the size of the input dataset (200 for large datasets, 500 for small).
          If a list of int is specified, then the intermediate embeddings at the
          different epochs specified in that list are returned in
          ``aux_data["embedding_list"]``.
  
      init: string
          How to initialize the low dimensional embedding. Options are:
              * 'spectral': use a spectral embedding of the fuzzy 1-skeleton
              * 'random': assign initial embedding positions at random.
              * A numpy array of initial embedding positions.
  
      random_state: numpy RandomState or equivalent
          A state capable being used as a numpy random state.
  
      metric: string or callable
          The metric used to measure distance in high dimensional space; used if
          multiple connected components need to be layed out.
  
      metric_kwds: dict
          Key word arguments to be passed to the metric function; used if
          multiple connected components need to be layed out.
  
      densmap: bool
          Whether to use the density-augmented objective function to optimize
          the embedding according to the densMAP algorithm.
  
     
  
  [... the rest is omitted ...]
files: umap/umap_.py
files: # Author: Leland McInnes <leland.mcinnes@gmail.com>
  #
  # License: BSD 3 clause
  from __future__ import print_function
  
  import locale
  from warnings import warn
  import time
  
  from scipy.optimize import curve_fit
  from sklearn.base import BaseEstimator
  from sklearn.utils import check_random_state, check_array
  from sklearn.utils.validation import check_is_fitted
  from sklearn.metrics import pairwise_distances
  from sklearn.preprocessing import normalize
  from sklearn.neighbors import KDTree
  
  try:
      import joblib
  except ImportError:
      # sklearn.externals.joblib is deprecated in 0.21, will be removed in 0.23
      from sklearn.externals import joblib
  
  import numpy as np
  import scipy.sparse
  from scipy.sparse import tril as sparse_tril, triu as sparse_triu
  import scipy.sparse.csgraph
  import numba
  
  import umap.distances as dist
  
  import umap.sparse as sparse
  
  from umap.utils import (
      submatrix,
      ts,
      csr_unique,
      fast_knn_indices,
  )
  from umap.spectral import spectral_layout
  from umap.layouts import (
      optimize_layout_euclidean,
      optimize_layout_generic,
      optimize_layout_inverse,
  )
  
  from pynndescent import NNDescent
  from pynndescent.distances import named_distances as pynn_named_distances
  from pynndescent.sparse import sparse_named_distances as pynn_sparse_named_distances
  
  locale.setlocale(locale.LC_NUMERIC, "C")
  
  INT32_MIN = np.iinfo(np.int32).min + 1
  INT32_MAX = np.iinfo(np.int32).max - 1
  
  SMOOTH_K_TOLERANCE = 1e-5
  MIN_K_DIST_SCALE = 1e-3
  NPY_INFINITY = np.inf
  
  DISCONNECTION_DISTANCES = {
      "correlation": 2,
      "cosine": 2,
      "hellinger": 1,
      "jaccard": 1,
      "dice": 1,
  }
  
  
  def flatten_iter(container):
      for i in container:
          if isinstance(i, (list, tuple)):
              for j in flatten_iter(i):
                  yield j
          else:
              yield i
  
  
  def flattened(container):
      return tuple(flatten_iter(container))
  
  
  def breadth_first_search(adjmat, start, min_vertices):
      explored = []
      queue = [start]
      levels = {}
      levels[start] = 0
      max_level = np.inf
      visited = [start]
  
      while queue:
          node = queue.pop(0)
          explored.append(node)
          if max_level == np.inf and len(explored) > min_vertices:
              max_level = max(levels.values())
  
          if levels[node] + 1 < max_level:
              neighbors = adjmat[node].indices
              for neighbour in neighbors:
                  if neighbour not in visited:
                      queue.append(neighbour)
                      visited.append(neighbour)
  
                      levels[neighbour] = levels[node] + 1
  
      return np.array(explored)
  
  
  def raise_disconnected_warning(
      edges_removed,
      vertices_disconnected,
      disconnection_distance,
      total_rows,
      threshold=0.1,
      verbose=False,
  ):
      """A simple wrapper function to avoid large amounts of code repetition."""
      if verbose & (vertices_disconnected == 0) & (edges_removed > 0):
          print(
              f"Disconnection_distance = {disconnection_distance} has removed {edges_removed} edges.  "
              f"This is not a problem as no vertices were disconnected."
          )
      elif (vertices_disconnected > 0) & (
          vertices_disconnected <= threshold * total_rows
      ):
          warn(
              f"A few of your vertices were disconnected from the manifold.  This shouldn't cause problems.\n"
              f"Disconnection_distance = {disconnection_distance} has removed {edges_removed} edges.\n"
              f"It has only fully disconnected {vertices_disconnected} vertices.\n"
              f"Use umap.utils.disconnected_vertices() to identify them.",
          )
      elif vertices_disconnected > threshold * total_rows:
          warn(
              f"A large number of your vertices were disconnected from the manifold.\n"
              f"Disconnection_distance = {disconnection_distance} has removed {edges_removed} edges.\n"
              f"It has fully disconnected {vertices_disconnected} vertices.\n"
              f"You might consider using find_disconnected_points() to find and remove these points from your data.\n"
              f"Use umap.utils.disconnected_vertices() to identify them.",
          )
  
  
  @numba.njit(
      locals={
          "psum": numba.types.float32,
          "lo": numba.types.float32,
          "mid": numba.types.float32,
          "hi": numba.types.float32,
      },
      fastmath=True,
  )  # benchmarking `parallel=True` shows it to *decrease* performance
  def smooth_knn_dist(distances, k, n_iter=64, local_connectivity=1.0, bandwidth=1.0):
      """Compute a continuous version of the distance to the kth nearest
      neighbor. That is, this is similar to knn-distance but allows continuous
      k values rather than requiring an integral k. In essence we are simply
      computing the distance such that the cardinality of fuzzy set we generate
      is k.
  
      Parameters
      ----------
      distances: array of shape (n_samples, n_neighbors)
          Distances to nearest neighbors for each samples. Each row should be a
          sorted list of distances to a given samples nearest neighbors.
  
      k: float
          The number of nearest neighbors to approximate for.
  
      n_iter: int (optional, default 64)
          We need to binary search for the correct distance value. This is the
          max number of iterations to use in such a search.
  
      local_connectivity: int (optional, default 1)
          The local connectivity required -- i.e. the number of nearest
          neighbors that should be assumed to be connected at a local level.
          The higher this value the more connected the manifold becomes
          locally. In practice this should be not more than the local intrinsic
          dimension of the manifold.
  
      bandwidth: float (optional, default 1)
          The target bandwidth of the kernel, larger values will produce
          larger return values.
  
      Returns
      -------
      knn_dist: array of shape (n_samples,)
          The distance to kth nearest neighbor, as suitably approximated.
  
      nn_dist: array of shape (n_samples,)
          The distance to the 1st nearest neighbor for each point.
      """
      target = np.log2(k) * bandwidth
      rho = np.zeros(distances.shape[0], dtype=np.float32)
      result = np.zeros(distances.shape[0], dtype=np.float32)
  
      mean_distances = np.mean(distances)
  
      for i in range(distances.shape[0]):
          lo = 0.0
          hi = NPY_INFINITY
          mid = 1.0
  
          # TODO: This is very inefficient, but will do for now. FIXME
          ith_distances = distances[i]
          non_zero_dists = ith_distances[ith_distances > 0.0]
          if non_zero_dists.shape[0] >= local_connectivity:
              index = int(np.floor(local_connectivity))
              interpolation = local_connectivity - index
              if index > 0:
                  rho[i] = non_zero_dists[index - 1]
                  if interpolation > SMOOTH_K_TOLERANCE:
                      rho[i] += interpolation * (
                          non_zero_dists[index] - non_zero_dists[index - 1]
                      )
              else:
                  rho[i] = interpolation * non_zero_dists[0]
          elif non_zero_dists.shape[0] > 0:
              rho[i] = np.max(non_zero_dists)
  
          for n in range(n_iter):
  
              psum = 0.0
              for j in range(1, distances.shape[1]):
                  d = distances[i, j] - rho[i]
                  if d > 0:
                      psum += np.exp(-(d / mid))
                  else:
                      psum += 1.0
  
              if np.fabs(psum - target) < SMOOTH_K_TOLERANCE:
                  break
  
              if psum > target:
                  hi = mid
                  mid = (lo + hi) / 2.0
              else:
                  lo = mid
                  if hi == NPY_INFINITY:
                      mid *= 2
                  else:
                      mid = (lo + hi) / 2.0
  
          result[i] = mid
  
          # TODO: This is very inefficient, but will do for now. FIXME
          if rho[i] > 0.0:
              mean_ith_distances = np.mean(ith_distances)
              if result[i] < MIN_K_DIST_SCALE * mean_ith_distances:
                  result[i] = MIN_K_DIST_SCALE * mean_ith_distances
          else:
              if result[i] < MIN_K_DIST_SCALE * mean_distances:
                  result[i] = MIN_K_DIST_SCALE * mean_distances
  
      return result, rho
  
  
  def nearest_neighbors(
      X,
      n_neighbors,
      metric,
      metric_kwds,
      angular,
      random_state,
      low_memory=True,
      use_pynndescent=True,
      n_jobs=-1,
      verbose=False,
  ):
      """Compute the ``n_neighbors`` nearest points for each data point in ``X``
      under ``metric``. This may be exact, but more likely is approximated via
      nearest neighbor descent.
  
      Parameters
      ----------
      X: array of shape (n_samples, n_features)
          The input data to compute the k-neighbor graph of.
  
      n_neighbors: int
          The number of nearest neighbors to compute for each sample in ``X``.
  
      metric: string or callable
          The metric to use for the computation.
  
      metric_kwds: dict
          Any arguments to pass to the metric computation function.
  
      angular: bool
          Whether to use angular rp trees in NN approximation.
  
      random_state: np.random state
          The random state to use for approximate NN computations.
  
      low_memory: bool (optional, default True)
          Whether to pursue lower memory NNdescent.
  
      verbose: bool (optional, default False)
          Whether to print status data during the computation.
  
      Returns
      -------
      knn_indices: array of shape (n_samples, n_neighbors)
          The indices on the ``n_neighbors`` closest points in the dataset.
  
      knn_dists: array of shape (n_samples, n_neighbors)
          The distances to the ``n_neighbors`` closest points in the dataset.
  
      rp_forest: list of trees
          The random projection forest used for searching (if used, None otherwise)
      """
      if verbose:
          print(ts(), "Finding Nearest Neighbors")
  
      if metric == "precomputed":
          # Note that this does not support sparse distance matrices yet ...
          # Compute indices of n nearest neighbors
          knn_indices = fast_knn_indices(X, n_neighbors)
          # knn_indices = np.argsort(X)[:, :n_neighbors]
          # Compute the nearest neighbor distances
          #   (equivalent to np.sort(X)[:,:n_neighbors])
          knn_dists = X[np.arange(X.shape[0])[:, None], knn_indices].copy()
          # Prune any nearest neighbours that are infinite distance apart.
          disconnected_index = knn_dists == np.inf
          knn_indices[disconnected_index] = -1
  
          knn_search_index = None
      else:
          # TODO: Hacked values for now
          n_trees = min(64, 5 + int(round((X.shape[0]) ** 0.5 / 20.0)))
          n_iters = max(5, int(round(np.log2(X.shape[0]))))
  
          knn_search_index = NNDescent(
              X,
              n_neighbors=n_neighbors,
              metric=metric,
              metric_kwds=metric_kwds,
              random_state=random_state,
              n_trees=n_trees,
              n_iters=n_iters,
              max_candidates=60,
              low_memory=low_memory,
              n_jobs=n_jobs,
              verbose=verbose,
          )
          knn_indices, knn_dists = knn_search_index.neighbor_graph
  
      if verbose:
          print(ts(), "Finished Nearest Neighbor Search")
      return knn_indices, knn_dists, knn_search_index
  
  
  @numba.njit(
      locals={
          "knn_dists": numba.types.float32[:, ::1],
          "sigmas": numba.types.float32[::1],
          "rhos": numba.types.float32[::1],
          "val": numba.types.float32,
      },
      parallel=True,
      fastmath=True,
  )
  def compute_membership_strengths(
      knn_indices,
      knn_dists,
      sigmas,
      rhos,
      return_dists=False,
      bipartite=False,
  ):
      """Construct the membership strength data for the 1-skeleton of each local
      fuzzy simplicial set -- this is formed as a sparse matrix where each row is
      a local fuzzy simplicial set, with a membership strength for the
      1-simplex to each other data point.
  
      Parameters
      ----------
      knn_indices: array of shape (n_samples, n_neighbors)
          The indices on the ``n_neighbors`` closest points in the dataset.
  
      knn_dists: array of shape (n_samples, n_neighbors)
          The distances to the ``n_neighbors`` closest points in the dataset.
  
      sigmas: array of shape(n_samples)
          The normalization factor derived from the metric tensor approximation.
  
      rhos: array of shape(n_samples)
          The local connectivity adjustment.
  
      return_dists: bool (optional, default False)
          Whether to return the pairwise distance associated with each edge
  
      bipartite: bool (optional, default False)
          Does the nearest neighbour set represent a bipartite graph?  That is are the
          nearest neighbour indices from the same point set as the row indices?
  
      Returns
      -------
      rows: array of shape (n_samples * n_neighbors)
          Row data for the resulting sparse matrix (coo format)
  
      cols: array of shape (n_samples * n_neighbors)
          Column data for the resulting sparse matrix (coo format)
  
      vals: array of shape (n_samples * n_neighbors)
          Entries for the resulting sparse matrix (coo format)
  
      dists: array of shape (n_samples * n_neighbors)
          Distance associated with each entry in the resulting sparse matrix
      """
      n_samples = knn_indices.shape[0]
      n_neighbors = knn_indices.shape[1]
  
      rows = np.zeros(knn_indices.size, dtype=np.int32)
      cols = np.zeros(knn_indices.size, dtype=np.int32)
      vals = np.zeros(knn_indices.size, dtype=np.float32)
      if return_dists:
          dists = np.zeros(knn_indices.size, dtype=np.float32)
      else:
          dists = None
  
      for i in range(n_samples):
          for j in range(n_neighbors):
              if knn_indices[i, j] == -1:
                  continue  # We didn't get the full knn for i
              # If applied to an adjacency matrix points shouldn't be similar to themselves.
              # If applied to an incidence matrix (or bipartite) then the row and column indices are different.
              if (bipartite == False) & (knn_indices[i, j] == i):
                  val = 0.0
              elif knn_dists[i, j] - rhos[i] <= 0.0 or sigmas[i] == 0.0:
                  val = 1.0
              else:
                  val = np.exp(-((knn_dists[i, j] - rhos[i]) / (sigmas[i])))
  
              rows[i * n_neighbors + j] = i
              cols[i * n_neighbors + j] = knn_indices[i, j]
              vals[i * n_neighbors + j] = val
              if return_dists:
                  dists[i * n_neighbors + j] = knn_dists[i, j]
  
      return rows, cols, vals, dists
  
  
  def fuzzy_simplicial_set(
      X,
      n_neighbors,
      random_state,
      metric,
      metric_kwds={},
      knn_indices=None,
      knn_dists=None,
      angular=False,
      set_op_mix_ratio=1.0,
      local_connectivity=1.0,
      apply_set_operations=True,
      verbose=False,
      return_dists=None,
  ):
      """Given a set of data X, a neighborhood size, and a measure of distance
      compute the fuzzy simplicial set (here represented as a fuzzy graph in
      the form of a sparse matrix) associated to the data. This is done by
      locally approximating geodesic distance at each point, creating a fuzzy
      simplicial set for each such point, and then combining all the local
      fuzzy simplicial sets into a global one via a fuzzy union.
  
      Parameters
      ----------
      X: array of shape (n_samples, n_features)
          The data to be modelled as a fuzzy simplicial set.
  
      n_neighbors: int
          The number of neighbors to use to approximate geodesic distance.
          Larger numbers induce more global estimates of the manifold that can
          miss finer detail, while smaller values will focus on fine manifold
          structure to the detriment of the larger picture.
  
      random_state: numpy RandomState or equivalent
          A state capable being used as a numpy random state.
  
      metric: string or function (optional, default 'euclidean')
          The metric to use to compute distances in high dimensional space.
          If a string is passed it must match a valid predefined metric. If
          a general metric is required a function that takes two 1d arrays and
          returns a float can be provided. For performance purposes it is
          required that this be a numba jit'd function. Valid string metrics
          include:
              * euclidean (or l2)
              * manhattan (or l1)
              * cityblock
              * braycurtis
              * canberra
              * chebyshev
              * correlation
              * cosine
              * dice
              * hamming
              * jaccard
              * kulsinski
              * ll_dirichlet
              * mahalanobis
              * matching
              * minkowski
              * rogerstanimoto
              * russellrao
              * seuclidean
              * sokalmichener
              * sokalsneath
              * sqeuclidean
              * yule
              * wminkowski
  
          Metrics that take arguments (such as minkowski, mahalanobis etc.)
          can have arguments passed via the metric_kwds dictionary. At this
          time care must be taken and dictionary elements must be ordered
          appropriately; this will hopefully be fixed in the future.
  
      metric_kwds: dict (optional, default {})
          Arguments to pass on to the metric, such as the ``p`` value for
          Minkowski distance.
  
      knn_indices: array of shape (n_samples, n_neighbors) (optional)
          If the k-nearest neighbors of each point has already been calculated
          you can pass them in here to save computation time. This should be
          an array with the indices of the k-nearest neighbors as a row for
          each data point.
  
      knn_dists: array of shape (n_samples, n_neighbors) (optional)
          If the k-nearest neighbors of each point has already been calculated
          you can pass them in here to save computation time. This should be
          an array with the distances of the k-nearest neighbors as a row for
          each data point.
  
      angular: bool (optional, default False)
          Whether to use angular/cosine distance for the random projection
          forest for seeding NN-descent to determine approximate nearest
          neighbors.
  
      set_op_mix_ratio: float (optional, default 1.0)
          Interpolate between (fuzzy) union and intersection as the set operation
          used to combine local fuzzy simplicial sets to obtain a global fuzzy
          simplicial sets. Both fuzzy set operations use the product t-norm.
          The value of this parameter should be between 0.0 and 1.0; a value of
          1.0 will use a pure fuzzy union, while 0.0 will use a pure fuzzy
          intersection.
  
      local_connectivity: int (optional, default 1)
          The local connectivity required -- i.e. the number of nearest
          neighbors that should be assumed to be connected at a local level.
          The higher this value the more connected the manifold becomes
          locally. In practice this should be not more than the local intrinsic
          dimension of the manifold.
  
      verbose: bool (optional, default False)
          Whether to report information on the current progress of the algorithm.
  
      return_dists: bool or None (optional, default None)
          Whether to return the pairwise distance associated with each edge.
  
      Returns
      -------
      fuzzy_simplicial_set: coo_matrix
          A fuzzy simplicial set represented as a sparse matrix. The (i,
          j) entry of the matrix represents the membership strength of the
          1-simplex between the ith and jth sample points.
      """
      if knn_indices is None or knn_dists is None:
          knn_indices, knn_dists, _ = nearest_neighbors(
              X,
              n_neighbors,
              metric,
              metric_kwds,
              angular,
              random_state,
              verbose=verbose,
          )
  
      knn_dists = knn_dists.astype(np.float32)
  
      sigmas, rhos = smooth_knn_dist(
          knn_dists,
          float(n_neighbors),
          local_connectivity=float(local_connectivity),
      )
  
      rows, cols, vals, dists = compute_membership_strengths(
          knn_indices, knn_dists, sigmas, rhos, return_dists
      )
  
      result = scipy.sparse.coo_matrix(
          (vals, (rows, cols)), shape=(X.shape[0], X.shape[0])
      )
      result.eliminate_zeros()
  
      if apply_set_operations:
          transpose = result.transpose()
  
          prod_matrix = result.multiply(transpose)
  
          result = (
              set_op_mix_ratio * (result + transpose - prod_matrix)
              + (1.0 - set_op_mix_ratio) * prod_matrix
          )
  
      result.eliminate_zeros()
  
      if return_dists is None:
          return result, sigmas, rhos
      else:
          if return_dists:
              dmat = scipy.sparse.coo_matrix(
                  (dists, (rows, cols)), shape=(X.shape[0], X.shape[0])
              )
  
              dists = dmat.maximum(dmat.transpose()).todok()
          else:
              dists = None
  
          return result, sigmas, rhos, dists
  
  
  @numba.njit()
  def fast_intersection(rows, cols, values, target, unknown_dist=1.0, far_dist=5.0):
      """Under the assumption of categorical distance for the intersecting
      simplicial set perform a fast intersection.
  
      Parameters
      ----------
      rows: array
          An array of the row of each non-zero in the sparse matrix
          representation.
  
      cols: array
          An array of the column of each non-zero in the sparse matrix
          representation.
  
      values: array
          An array of the value of each non-zero in the sparse matrix
          representation.
  
      target: array of shape (n_samples)
          The categorical labels to use in the intersection.
  
      unknown_dist: float (optional, default 1.0)
          The distance an unknown label (-1) is assumed to be from any point.
  
      far_dist float (optional, default 5.0)
          The distance between unmatched labels.
  
      Returns
      -------
      None
      """
      for nz in range(rows.shape[0]):
          i = rows[nz]
          j = cols[nz]
          if (target[i] == -1) or (target[j] == -1):
              values[nz] *= np.exp(-unknown_dist)
          elif target[i] != target[j]:
              values[nz] *= np.exp(-far_dist)
  
      return
  
  
  @numba.jit()
  def fast_metric_intersection(
      rows, cols, values, discrete_space, metric, metric_args, scale
  ):
      """Under the assumption of categorical distance for the intersecting
      simplicial set perform a fast intersection.
  
      Parameters
      ----------
      rows: array
          An array of the row of each non-zero in the sparse matrix
          representation.
  
      cols: array
          An array of the column of each non-zero in the sparse matrix
          representation.
  
      values: array of shape
          An array of the values of each non-zero in the sparse matrix
          representation.
  
      discrete_space: array of shape (n_samples, n_features)
          The vectors of categorical labels to use in the intersection.
  
      metric: numba function
          The function used to calculate distance over the target array.
  
      scale: float
          A scaling to apply to the metric.
  
      Returns
      -------
      None
      """
      for nz in range(rows.shape[0]):
          i = rows[nz]
          j = cols[nz]
          dist = metric(discrete_space[i], discrete_space[j], *metric_args)
          values[nz] *= np.exp(-(scale * dist))
  
      return
  
  
  @numba.njit()
  def reprocess_row(probabilities, k=15, n_iters=32):
      target = np.log2(k)
  
      lo = 0.0
      hi = NPY_INFINITY
      mid = 1.0
  
      for n in range(n_iters):
  
          psum = 0.0
          for j in range(probabilities.shape[0]):
              psum += pow(probabilities[j], mid)
  
          if np.fabs(psum - target) < SMOOTH_K_TOLERANCE:
              break
  
          if psum < target:
              hi = mid
              mid = (lo + hi) / 2.0
          else:
              lo = mid
              if hi == NPY_INFINITY:
                  mid *= 2
              else:
                  mid = (lo + hi) / 2.0
  
      return np.power(probabilities, mid)
  
  
  @numba.njit()
  def reset_local_metrics(simplicial_set_indptr, simplicial_set_data):
      for i in range(simplicial_set_indptr.shape[0] - 1):
          simplicial_set_data[
              simplicial_set_indptr[i] : simplicial_set_indptr[i + 1]
          ] = reprocess_row(
              simplicial_set_data[simplicial_set_indptr[i] : simplicial_set_indptr[i + 1]]
          )
      return
  
  
  def reset_local_connectivity(simplicial_set, reset_local_metric=False):
      """Reset the local connectivity requirement -- each data sample should
      have complete confidence in at least one 1-simplex in the simplicial set.
      We can enforce this by locally rescaling confidences, and then remerging the
      different local simplicial sets together.
  
      Parameters
      ----------
      simplicial_set: sparse matrix
          The simplicial set for which to recalculate with respect to local
          connectivity.
  
      Returns
      -------
      simplicial_set: sparse_matrix
          The recalculated simplicial set, now with the local connectivity
          assumption restored.
      """
      simplicial_set = normalize(simplicial_set, norm="max")
      if reset_local_metric:
          simplicial_set = simplicial_set.tocsr()
          reset_local_metrics(simplicial_set.indptr, simplicial_set.data)
          simplicial_set = simplicial_set.tocoo()
      transpose = simplicial_set.transpose()
      prod_matrix = simplicial_set.multiply(transpose)
      simplicial_set = simplicial_set + transpose - prod_matrix
      simplicial_set.eliminate_zeros()
  
      return simplicial_set
  
  
  def discrete_metric_simplicial_set_intersection(
      simplicial_set,
      discrete_space,
      unknown_dist=1.0,
      far_dist=5.0,
      metric=None,
      metric_kws={},
      metric_scale=1.0,
  ):
      """Combine a fuzzy simplicial set with another fuzzy simplicial set
      generated from discrete metric data using discrete distances. The target
      data is assumed to be categorical label data (a vector of labels),
      and this will update the fuzzy simplicial set to respect that label data.
  
      TODO: optional category cardinality based weighting of distance
  
      Parameters
      ----------
      simplicial_set: sparse matrix
          The input fuzzy simplicial set.
  
      discrete_space: array of shape (n_samples)
          The categorical labels to use in the intersection.
  
      unknown_dist: float (optional, default 1.0)
          The distance an unknown label (-1) is assumed to be from any point.
  
      far_dist: float (optional, default 5.0)
          The distance between unmatched labels.
  
      metric: str (optional, default None)
          If not None, then use this metric to determine the
          distance between values.
  
      metric_scale: float (optional, default 1.0)
          If using a custom metric scale the distance values by
          this value -- this controls the weighting of the
          intersection. Larger values weight more toward target.
  
      Returns
      -------
      simplicial_set: sparse matrix
          The resulting intersected fuzzy simplicial set.
      """
      simplicial_set = simplicial_set.tocoo()
  
      if metric is not None:
          # We presume target is now a 2d array, with each row being a
          # vector of target info
          if metric in dist.named_distances:
              metric_func = dist.named_distances[metric]
          else:
              raise ValueError("Discrete intersection metric is not recognized")
  
          fast_metric_intersection(
              simplicial_set.row,
              simplicial_set.col,
              simplicial_set.data,
              discrete_space,
              metric_func,
              tuple(metric_kws.values()),
              metric_scale,
          )
      else:
          fast_intersection(
              simplicial_set.row,
              simplicial_set.col,
              simplicial_set.data,
              discrete_space,
              unknown_dist,
              far_dist,
          )
  
      simplicial_set.eliminate_zeros()
  
      return reset_local_connectivity(simplicial_set)
  
  
  def general_simplicial_set_intersection(
      simplicial_set1, simplicial_set2, weight=0.5, right_complement=False
  ):
  
      if right_complement:
          result = simplicial_set1.tocoo()
      else:
          result = (simplicial_set1 + simplicial_set2).tocoo()
      left = simplicial_set1.tocsr()
      right = simplicial_set2.tocsr()
  
      sparse.general_sset_intersection(
          left.indptr,
          left.indices,
          left.data,
          right.indptr,
          right.indices,
          right.data,
          result.row,
          result.col,
          result.data,
          mix_weight=weight,
          right_complement=right_complement,
      )
  
      return result
  
  
  def general_simplicial_set_union(simplicial_set1, simplicial_set2):
      result = (simplicial_set1 + simplicial_set2).tocoo()
      left = simplicial_set1.tocsr()
      right = simplicial_set2.tocsr()
  
      sparse.general_sset_union(
          left.indptr,
          left.indices,
          left.data,
          right.indptr,
          right.indices,
          right.data,
          result.row,
          result.col,
          result.data,
      )
  
      return result
  
  
  def make_epochs_per_sample(weights, n_epochs):
      """Given a set of weights and number of epochs generate the number of
      epochs per sample for each weight.
  
      Parameters
      ----------
      weights: array of shape (n_1_simplices)
          The weights ofhow much we wish to sample each 1-simplex.
  
      n_epochs: int
          The total number of epochs we want to train for.
  
      Returns
      -------
      An array of number of epochs per sample, one for each 1-simplex.
      """
      result = -1.0 * np.ones(weights.shape[0], dtype=np.float64)
      n_samples = n_epochs * (weights / weights.max())
      result[n_samples > 0] = float(n_epochs) / n_samples[n_samples > 0]
      return result
  
  
  def simplicial_set_embedding(
      data,
      graph,
      n_components,
      initial_alpha,
      a,
      b,
      gamma,
      negative_sample_rate,
      n_epochs,
      init,
      random_state,
      metric,
      metric_kwds,
      densmap,
      densmap_kwds,
      output_dens,
      output_metric=dist.named_distances_with_gradients["euclidean"],
      output_metric_kwds={},
      euclidean_output=True,
      parallel=False,
      verbose=False,
  ):
      """Perform a fuzzy simplicial set embedding, using a specified
      initialisation method and then minimizing the fuzzy set cross entropy
      between the 1-skeletons of the high and low dimensional fuzzy simplicial
      sets.
  
      Parameters
      ----------
      data: array of shape (n_samples, n_features)
          The source data to be embedded by UMAP.
  
      graph: sparse matrix
          The 1-skeleton of the high dimensional fuzzy simplicial set as
          represented by a graph for which we require a sparse matrix for the
          (weighted) adjacency matrix.
  
      n_components: int
          The dimensionality of the euclidean space into which to embed the data.
  
      initial_alpha: float
          Initial learning rate for the SGD.
  
      a: float
          Parameter of differentiable approximation of right adjoint functor
  
      b: float
          Parameter of differentiable approximation of right adjoint functor
  
      gamma: float
          Weight to apply to negative samples.
  
      negative_sample_rate: int (optional, default 5)
          The number of negative samples to select per positive sample
          in the optimization process. Increasing this value will result
          in greater repulsive force being applied, greater optimization
          cost, but slightly more accuracy.
  
      n_epochs: int (optional, default 0)
          The number of training epochs to be used in optimizing the
          low dimensional embedding. Larger values result in more accurate
          embeddings. If 0 is specified a value will be selected based on
          the size of the input dataset (200 for large datasets, 500 for small).
  
      init: string
          How to initialize the low dimensional embedding. Options are:
              * 'spectral': use a spectral embedding of the fuzzy 1-skeleton
              * 'random': assign initial embedding positions at random.
              * A numpy array of initial embedding positions.
  
      random_state: numpy RandomState or equivalent
          A state capable being used as a numpy random state.
  
      metric: string or callable
          The metric used to measure distance in high dimensional space; used if
          multiple connected components need to be layed out.
  
      metric_kwds: dict
          Key word arguments to be passed to the metric function; used if
          multiple connected components need to be layed out.
  
      densmap: bool
          Whether to use the density-augmented objective function to optimize
          the embedding according to the densMAP algorithm.
  
      densmap_kwds: dict
          Key word arguments to be used by the densMAP optimization.
  
      output_dens: bool
          Whether to output local radii in the original data and the embedding.
  
     
  
  [... the rest is omitted ...]
files: umap/umap_.py