code: 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 n_embeddings > neighbor_m >= 0 != offset:
                              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 (
                                  n_embeddings
                                  > neighbor_m
                                  >= 0
                                  != offset
                              ):
                                  identified_index = relations[
                                      m, offset + window_size, k
                                  ]
                                  if identified_index >= 0:
                                      other_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 = 0.0
  
                          for offset in range(
                              -window_size, window_size
                          ):
                              neighbor_m = m + offset
                              if (
                                  n_embeddings
                                  > neighbor_m
                                  >= 0
                                  != offset
                              ):
                                  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]
                  )
  
file_path: umap/umap/layouts.py
imports: import numpy
  import numba
  import umap.distances
  from umap.utils import tau_rand_int
  from tqdm.auto import tqdm
label_desc: Use image node_id 10 for calling a global function with example usage: _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) without return types