code: def _stump(
      T_A,
      T_B,
      m,
      μ_Q,
      M_T,
      σ_Q_inverse,
      Σ_T_inverse,
      μ_Q_m_1,
      M_T_m_1,
      T_A_subseq_isfinite,
      T_B_subseq_isfinite,
      T_A_subseq_isconstant,
      T_B_subseq_isconstant,
      diags,
      ignore_trivial,
      k,
  ):
      """
      A Numba JIT-compiled version of STOMPopt with Pearson correlations for parallel
      computation of the (top-k) matrix profile, the (top-k) matrix profile indices,
      the top-1 left matrix profile and its matrix profile index, and the top-1 right
      matrix profile and its matrix profile index.
  
      Parameters
      ----------
      T_A : numpy.ndarray
          The time series or sequence for which to compute the matrix profile
  
      T_B : numpy.ndarray
          The time series or sequence that will be used to annotate `T_A`. For every
          subsequence in `T_A`, its nearest neighbor in `T_B` will be recorded.
  
      m : int
          Window size
  
      μ_Q : numpy.ndarray
          Mean of the query sequence, `Q`, relative to the current sliding window
  
      M_T : numpy.ndarray
          Sliding mean of time series, `T`
  
      σ_Q_inverse : numpy.ndarray
          Inverse standard deviation of the query sequence, `Q`, relative to the current
          sliding window
  
      Σ_T_inverse : numpy.ndarray
          Inverse sliding standard deviation of time series, `T`
  
      μ_Q_m_1 : numpy.ndarray
          Mean of the query sequence, `Q`, relative to the current sliding window and
          using a window size of `m-1`
  
      M_T_m_1 : numpy.ndarray
          Sliding mean of time series, `T`, using a window size of `m-1`
  
      T_A_subseq_isfinite : numpy.ndarray
          A boolean array that indicates whether a subsequence in `T_A` contains a
          `np.nan`/`np.inf` value (False)
  
      T_B_subseq_isfinite : numpy.ndarray
          A boolean array that indicates whether a subsequence in `T_B` contains a
          `np.nan`/`np.inf` value (False)
  
      T_A_subseq_isconstant : numpy.ndarray
          A boolean array that indicates whether a subsequence in `T_A` is constant (True)
  
      T_B_subseq_isconstant : numpy.ndarray
          A boolean array that indicates whether a subsequence in `T_B` is constant (True)
  
      diags : numpy.ndarray
          The diagonal indices
  
      ignore_trivial : bool
          Set to `True` if this is a self-join. Otherwise, for AB-join, set this to
          `False`. Default is `True`.
  
      k : int
          The number of top `k` smallest distances used to construct the matrix profile.
          Note that this will increase the total computational time and memory usage
          when k > 1.
  
      Returns
      -------
      out1 : numpy.ndarray
          The (top-k) matrix profile
  
      out2 : numpy.ndarray
          The (top-1) left matrix profile
  
      out3 : numpy.ndarray
          The (top-1) right matrix profile
  
      out4 : numpy.ndarray
          The (top-k) matrix profile indices
  
      out5 : numpy.ndarray
          The (top-1) left matrix profile indices
  
      out6 : numpy.ndarray
          The (top-1) right matrix profile indices
  
      Notes
      -----
      `DOI: 10.1007/s10115-017-1138-x \
      <https://www.cs.ucr.edu/~eamonn/ten_quadrillion.pdf>`__
  
      See Section 4.5
  
      The above reference outlines a general approach for traversing the distance
      matrix in a diagonal fashion rather than in a row-wise fashion.
  
      `DOI: 10.1145/3357223.3362721 \
      <https://www.cs.ucr.edu/~eamonn/public/GPU_Matrix_profile_VLDB_30DraftOnly.pdf>`__
  
      See Section 3.1 and Section 3.3
  
      The above reference outlines the use of the Pearson correlation via Welford's
      centered sum-of-products along each diagonal of the distance matrix in place of the
      sliding window dot product found in the original STOMP method.
  
      `DOI: 10.1109/ICDM.2016.0085 \
      <https://www.cs.ucr.edu/~eamonn/STOMP_GPU_final_submission_camera_ready.pdf>`__
  
      See Table II
  
      Timeseries, T_A, will be annotated with the distance location
      (or index) of all its subsequences in another times series, T_B.
  
      Return: For every subsequence, Q, in T_A, you will get a distance
      and index for the closest subsequence in T_B. Thus, the array
      returned will have length T_A.shape[0]-m+1. Additionally, the
      left and right matrix profiles are also returned.
  
      Note: Unlike in the Table II where T_A.shape is expected to be equal
      to T_B.shape, this implementation is generalized so that the shapes of
      T_A and T_B can be different. In the case where T_A.shape == T_B.shape,
      then our algorithm reduces down to the same algorithm found in Table II.
  
      Additionally, unlike STAMP where the exclusion zone is m/2, the default
      exclusion zone for STOMP is m/4 (See Definition 3 and Figure 3).
  
      For self-joins, set `ignore_trivial = True` in order to avoid the
      trivial match.
  
      Note that left and right matrix profiles are only available for self-joins.
      """
      n_A = T_A.shape[0]
      n_B = T_B.shape[0]
      l = n_A - m + 1
      n_threads = numba.config.NUMBA_NUM_THREADS
  
      ρ = np.full((n_threads, l, k), np.NINF, dtype=np.float64)
      I = np.full((n_threads, l, k), -1, dtype=np.int64)
  
      ρL = np.full((n_threads, l), np.NINF, dtype=np.float64)
      IL = np.full((n_threads, l), -1, dtype=np.int64)
  
      ρR = np.full((n_threads, l), np.NINF, dtype=np.float64)
      IR = np.full((n_threads, l), -1, dtype=np.int64)
  
      ndist_counts = core._count_diagonal_ndist(diags, m, n_A, n_B)
      diags_ranges = core._get_array_ranges(
          ndist_counts, n_threads, False
      )
  
      cov_a = T_B[m - 1 :] - M_T_m_1[:-1]
      cov_b = T_A[m - 1 :] - μ_Q_m_1[:-1]
      # The next lines are equivalent and left for reference
      # cov_c = np.roll(T_A, 1)
      # cov_ = cov_c[:M_T_m_1.shape[0]] - M_T_m_1[:]
      cov_c = np.empty(M_T_m_1.shape[0], dtype=np.float64)
      cov_c[1:] = T_B[: M_T_m_1.shape[0] - 1]
      cov_c[0] = T_B[-1]
      cov_c[:] = cov_c - M_T_m_1
      # The next lines are equivalent and left for reference
      # cov_d = np.roll(T_B, 1)
      # cov_d = cov_d[:μ_Q_m_1.shape[0]] - μ_Q_m_1[:]
      cov_d = np.empty(μ_Q_m_1.shape[0], dtype=np.float64)
      cov_d[1:] = T_A[: μ_Q_m_1.shape[0] - 1]
      cov_d[0] = T_A[-1]
      cov_d[:] = cov_d - μ_Q_m_1
  
      for thread_idx in prange(n_threads):
          # Compute and update pearson correlations and matrix profile indices
          # within a single thread and avoiding race conditions
          _compute_diagonal(
              T_A,
              T_B,
              m,
              μ_Q,
              M_T,
              σ_Q_inverse,
              Σ_T_inverse,
              cov_a,
              cov_b,
              cov_c,
              cov_d,
              T_A_subseq_isfinite,
              T_B_subseq_isfinite,
              T_A_subseq_isconstant,
              T_B_subseq_isconstant,
              diags,
              diags_ranges[thread_idx, 0],
              diags_ranges[thread_idx, 1],
              thread_idx,
              ρ,
              ρL,
              ρR,
              I,
              IL,
              IR,
              ignore_trivial,
          )
  
      # Reduction of results from all threads
      for thread_idx in range(1, n_threads):
          # update top-k arrays
          core._merge_topk_ρI(ρ[0], ρ[thread_idx], I[0], I[thread_idx])
  
          # update left matrix profile and matrix profile indices
          mask = ρL[0] < ρL[thread_idx]
          ρL[0][mask] = ρL[thread_idx][mask]
          IL[0][mask] = IL[thread_idx][mask]
  
          # update right matrix profile and matrix profile indices
          mask = ρR[0] < ρR[thread_idx]
          ρR[0][mask] = ρR[thread_idx][mask]
          IR[0][mask] = IR[thread_idx][mask]
  
      # Reverse top-k rho (and its associated I) to be in descending order and
      # then convert from Pearson correlations to Euclidean distances (ascending order)
      p_norm = np.abs(2 * m * (1 - ρ[0, :, ::-1]))
      I = I[0, :, ::-1]
  
      p_norm_L = np.abs(2 * m * (1 - ρL[0, :]))
      p_norm_R = np.abs(2 * m * (1 - ρR[0, :]))
  
      for i in prange(p_norm.shape[0]):
          for j in range(p_norm.shape[1]):
              if p_norm[i, j] < config.STUMPY_P_NORM_THRESHOLD:
                  p_norm[i, j] = 0.0
  
          if p_norm_L[i] < config.STUMPY_P_NORM_THRESHOLD:
              p_norm_L[i] = 0.0
  
          if p_norm_R[i] < config.STUMPY_P_NORM_THRESHOLD:
              p_norm_R[i] = 0.0
  
      return (
          np.sqrt(p_norm),
          np.sqrt(p_norm_L),
          np.sqrt(p_norm_R),
          I,
          IL[0],
          IR[0],
      )
  
file_path: stumpy/stumpy/stump.py
imports: import numba
  import numpy
  from numba import njit, prange
  from .None import config, core
  from .aamp import aamp
label_desc: Use image node_id 2 for calling a global function with example usage: _stump(T_A, T_B, m, μ_Q, M_T, σ_Q_inverse, Σ_T_inverse, μ_Q_m_1, M_T_m_1, T_A_subseq_isfinite, T_B_subseq_isfinite, T_A_subseq_isconstant, T_B_subseq_isconstant, diags, ignore_trivial, k) and returns: np,  np,  np,  I,  IL,  IR