Wykrywanie linii rdzenia wirowego

def Lambda2(macroscopic_velocities, negative_cutoff = 0):
  """Use vortex detection algorithm "Lambda2" to get which nodes in your velocity field are inside a vortex.
  """
  # First define the gradients of the velocity field: Δu
  _gradients = np.gradient(macroscopic_velocities, 
                           axis = (0, 1, 2))
  gradients = np.einsum('dD...-> ...dD', 
                        np.array(_gradients))

  # Get the strain rate tensor: S = (Δu + Δuᵀ)/2
  transposed_gradients = np.einsum('...dD -> ...Dd', 
                                   gradients)
  strain_rate = (gradients + transposed_gradients)/2

  # Get the vorticity tensor: Ω = (Δu - Δuᵀ)/2
  vorticity = (gradients - transposed_gradients)/2 

  # Get the eigenvalues (λ) of S² + Ω². Use np.eigh to get ordered eigenvalues. 
  S2_O2 = strain_rate**2 + vorticity**2
  eigen_values, errors = np.linalg.eigh(S2_O2)

  # Select the locations where λ₂ < the negative cutoff. 
  # Negative cutoff must be 0 or less.
  in_vortex = (eigen_values < negative_cutoff)[..., 1]
  return in_vortex
Bart D