! Copyright (c) 2011-2013 Manuel Hasert <m.hasert@grs-sim.de>
! Copyright (c) 2011-2016 Jiaxing Qi <jiaxing.qi@uni-siegen.de>
! Copyright (c) 2011 Jan Hueckelheim <j.hueckelheim@grs-sim.de>
! Copyright (c) 2011-2014 Simon Zimny <s.zimny@grs-sim.de>
! Copyright (c) 2012-2015, 2018-2019 Kannan Masilamani <kannan.masilamani@uni-siegen.de>
! Copyright (c) 2012-2014 Kartik Jain <kartik.jain@uni-siegen.de>
! Copyright (c) 2012, 2014 Harald Klimach <harald.klimach@uni-siegen.de>
! Copyright (c) 2016 Tobias Schneider <tobias1.schneider@student.uni-siegen.de>
! Copyright (c) 2018 Raphael Haupt <raphael.haupt@uni-siegen.de>
! Copyright (c) 2018 Jana Gericke <jana.gericke@uni-siegen.de>
! Copyright (c) 2019-2020 Peter Vitt <peter.vitt2@uni-siegen.de>
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions are met:
!
! 1. Redistributions of source code must retain the above copyright notice,
! this list of conditions and the following disclaimer.
!
! 2. Redistributions in binary form must reproduce the above copyright notice,
! this list of conditions and the following disclaimer in the documentation
! and/or other materials provided with the distribution.
!
! THIS SOFTWARE IS PROVIDED BY THE UNIVERSITY OF SIEGEN “AS IS” AND ANY EXPRESS
! OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
! OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
! IN NO EVENT SHALL UNIVERSITY OF SIEGEN OR CONTRIBUTORS BE LIABLE FOR ANY
! DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
! (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
! LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
! ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
! (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
! SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
! ****************************************************************************** !
!> author: Jiaxing Qi
!! This module provides functions for calculating macroscopic quantities
!!
?? include 'header/lbm_macros.inc'
module mus_derivedQuantities_module2

  ! include treelm modules
  use env_module,               only: rk
  use tem_stencil_module,       only: tem_stencilHeader_type
  use tem_param_module,         only: div1_3, div3_4, cs2inv, div1_9, div1_54, &
    &                                 cs2, cs4inv, div1_2, sqrt3, rho0
  use tem_logging_module,       only: logUnit
  use tem_float_module,         only: operator(.fne.)

  ! include musubi modules
  use mus_scheme_layout_module, only: mus_scheme_layout_type
  use mus_moments_module,       only: get_moment
  use mus_gradData_module,      only: mus_gradData_type

  implicit none
  private

  public :: getDensity, getVelocity, getVelocity_incomp
  public :: getEquilibriumIncomp, getEquilibrium
  public :: getWSS2D, getWSS
  public :: getShearStressTensor
  public :: getShearRateTensor_diffusive
  public :: getShearRateTensor_acoustic
  public :: getEqDistr
  public :: set_pdfDiffusive
  public :: set_pdfAcoustic
  public :: getNEq_diffusive
  public :: getNEq_acoustic
  public :: convPrePost
  public :: setThirdOrder_diffusive
  public :: secondMom
  public :: getShearRate
  public :: getStrainFacDffs
  public :: getNonEqFac
  public :: getNonEqFac_intp
  public :: getEqByDensVelPerDir
  public :: geteqbydensvel
  public :: getGradU

  interface getShearStressTensor
    module procedure getShearStressTensor_forElemFromState
    module procedure getShearStressTensor_forPdfSubset
    module procedure getShearStressTensorIncomp_forPdfSubset
  end interface

  interface getShearRateTensor_diffusive
    module procedure getShearRateTensor_diffusive_forPdfSubset
  end interface

  interface getShearRateTensor_acoustic
    ! module procedure getShearRateTensor_acoustic_forPdfSubset
    module procedure getShearRateTensor_acoustic_lbm
    module procedure getShearRateTensor_acoustic_incomp
  end interface

  interface getEquilibrium
    module procedure getEquilibrium_forElemfromState
    module procedure getEqByDensVel
    module procedure getEqByDensVelPerDir
    module procedure getEquilibrium_forPdfSubset
  end interface

  interface getDensity
    module procedure getDensity_forElemfromState
    module procedure getDensity_forPdfSubset
  end interface

  interface getVelocity
    module procedure getVelocity_forElemFromState_noForce
    ! module procedure getVelocity_forElemFromState_Force
    module procedure getVelocity_forPdfSubset
  end interface

  interface getVelocity_incomp
    module procedure getVelocity_forPdfSubset_incomp
  end interface getVelocity_incomp

contains

! ****************************************************************************** !
  !> Calculate the density of a given subset of pdfs
  !!        vector (sum up all values)
  !!
  pure function getDensity_forPdfSubset( subset, stencil, varPos ) result( res )
    ! ---------------------------------------------------------------------------
    type(tem_stencilHeader_type), intent(in) :: stencil
    real(kind=rk), intent(in) :: subset(:)
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    real(kind=rk)             :: res !< return value
    ! ---------------------------------------------------------------------------
    ! local variables
    integer :: iDir
    ! ---------------------------------------------------------------------------

    res = 0._rk
    do iDir = 1,stencil%QQ
      res = res + subset( varPos( iDir ) )
    enddo

  end function getDensity_forPdfSubset
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the density of a given element number with the given state
  !!        vector (sum up all values)
  !!
  pure function getDensity_forElemFromState( state, elem, stencil, varPos,     &
    &                                        nScalars ) result( res )
    ! ---------------------------------------------------------------------------
    type(tem_stencilHeader_type), intent(in) :: stencil
    real(kind=rk), intent(in) :: state(:)
    integer, intent(in)       :: elem
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk)             :: res !< return value
    ! ---------------------------------------------------------------------------
    ! local variables
    integer :: iDir
    integer :: nElems
    ! ---------------------------------------------------------------------------
    nElems = size( state ) / nScalars

    res = 0._rk
    do iDir = 1, stencil%QQ
      res = res + state( ?IDX?( varPos(iDir), elem, nScalars, nElems ))
    enddo

  end function getDensity_forElemFromState
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the velocity in all 3 directions
  !!        from a subset given, ordered according to the stencil
  !!
  pure function getVelocity_forPdfSubset( subset, stencil, varPos )            &
    &                             result( vel )
    ! ---------------------------------------------------------------------------
    type(tem_stencilHeader_type), intent(in) :: stencil !< stencil information
    real(kind=rk), intent(in) :: subset(:) !< complete state of one level
    integer,       intent(in) :: varPos(:) !< varPos of current field variable
    real(kind=rk)             :: vel(3)    !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: dens
    integer       :: iDir
    ! ---------------------------------------------------------------------------

    vel = 0._rk
    dens = 0._rk
    do iDir = 1,stencil%QQ
      vel(:) = vel(:) + subset(varPos(iDir)) * stencil%cxDirRK(:,iDir)
      dens   = dens   + subset(varPos(iDir))
    enddo
    vel = vel / dens

  end function getVelocity_forPdfSubset
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the velocity in all 3 directions
  !!        from a subset given, ordered according to the stencil
  !!
  pure function getVelocity_forPdfSubset_incomp( subset, stencil, varPos )     &
    &                             result( vel )
    ! ---------------------------------------------------------------------------
    type(tem_stencilHeader_type), intent(in) :: stencil !< stencil information
    real(kind=rk), intent(in) :: subset(:) !< complete state of one level
    integer,       intent(in) :: varPos(:) !< varPos of current field variable
    real(kind=rk)             :: vel(3)    !< return value
    ! ---------------------------------------------------------------------------
    integer       :: iDir
    ! ---------------------------------------------------------------------------

    vel = 0._rk
    do iDir = 1,stencil%QQ
      vel(:) = vel(:) + subset(varPos(iDir)) * stencil%cxDirRK(:,iDir)
    enddo
    vel = vel / rho0

  end function getVelocity_forPdfSubset_incomp
! ****************************************************************************** !

! ****************************************************************************** !
  !> Calculate the velocity in all 3 directions
  !! from the element indicated (elem),
  !! reading the pdf (state information) from the state array.
  !! state array includes all the pdfs of all elements.
  !! The access to the state array has to be done via the generic
  !! access macro IDX, as we want to access post-collision values.
  !!
  pure function getVelocity_forElemFromState_noForce( state, elem, stencil,    &
    &                                           varPos, nScalars ) result( vel )
    ! ---------------------------------------------------------------------------
    type(tem_stencilHeader_type), intent(in) :: stencil !< stencil information
    real(kind=rk), intent(in) :: state(:)  !< complete state of one level
    !> element index, for which to calc velocity
    integer, intent(in) :: elem
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk)       :: vel(3)    !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: dens
    integer       :: iDir
    integer :: nElems
    ! ---------------------------------------------------------------------------
    nElems = size( state ) / nScalars

    vel = 0._rk
    dens = 0._rk
    do iDir = 1,stencil%QQ
      vel(:) = vel(:) + state( ?IDX?( varPos(iDir), elem, nScalars,nElems ))&
        &               * stencil%cxDirRK(:,iDir)
      dens   = dens   + state( ?IDX?( varPos(iDir), elem, nScalars,nElems ))
    enddo

    vel = vel/dens

  end function getVelocity_forElemFromState_noForce
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the velocity in all 3 directions
  !!
  pure function getVelocity_forElemFromState_Force( state, elem, force,    &
    &                       dtLB, stencil, varPos, nScalars ) result( vel )
    ! ---------------------------------------------------------------------------
    real(kind=rk), intent(in) :: state(:) !< state array of a level
    type(tem_stencilHeader_type), intent(in) :: stencil
    real(kind=rk), intent(in) :: force(3)
    real(kind=rk), intent(in) :: dtLB
    integer, intent(in) :: elem
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk)       :: vel(3) !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: dens
    integer       :: iDir
    ! ---------------------------------------------------------------------------
    integer :: nElems
    ! ---------------------------------------------------------------------------
    nElems = size( state ) / nScalars

    vel = 0._rk
    dens = 0._rk
    do iDir = 1,stencil%QQ
      vel(:) = vel(:) + state( ?IDX?( varPos(iDir), elem, nScalars, nElems))  &
        &               * stencil%cxDirRK( :, iDir )           &
        ! Forcing Term
        & + dtLB*0.5_rk*force(:)

      dens = dens + state( ?IDX?( varPos(iDir), elem, nScalars, nElems))
    enddo
    vel = vel / dens

  end function getVelocity_forElemFromState_Force
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the momentum in all 3 directions  for a subset of pdfs
  !!        the subset is provided and must be of the size of DOFs per element
  !!
  pure function getMomentum_forPdfSubset_force( subset, force, dtLB, &
    &                                     stencil, varPos ) result( mom )
    ! ---------------------------------------------------------------------------
    real(kind=rk), intent(in) :: force(3)
    type(tem_stencilHeader_type), intent(in) :: stencil
    real(kind=rk), intent(in) :: dtLB
    real(kind=rk), intent(in) :: subset(:)
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    real(kind=rk) :: mom(3)!< return value
    ! ---------------------------------------------------------------------------
    integer :: iDir
    ! ---------------------------------------------------------------------------

    mom = 0._rk
    do iDir = 1,stencil%QQ
      mom(:) = mom(:) + subset(varPos(iDir)) * real(stencil%cxDir(:,iDir), rk) &
      ! Forcing Term
        &             + dtLB*div1_2*force(:)
    enddo

  end function getMomentum_forPdfSubset_force
! ****************************************************************************** !


! ****************************************************************************** !
  !> Get the equilibrium distribution in the specified direction iDir
  !!
  pure function getEqDistr( iDir, rho, vel, layout ) result( fEq )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: rho
    real(kind=rk), intent(in) :: vel(3)
    real(kind=rk) :: fEq!< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: usq, ucx, cxReal(3)
    integer, intent(in) :: iDir
    ! ---------------------------------------------------------------------------

    cxReal = real(layout%fStencil%cxDir(:,iDir), rk)
    usq = vel(1)*vel(1) + vel(2)*vel(2) + vel(3)*vel(3)
    ucx = cxReal(1)*vel(1) + cxReal(2)*vel(2) + cxReal(3)*vel(3)

    fEq = layout%weight( iDir )*rho*( 1.d0 +    &
      &     ucx*cs2inv                          &
      &   + ucx*ucx*cs2inv*cs2inv*div1_2        &
      &   - usq*div1_2*cs2inv )

  end function getEqDistr
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the equilibrium distribution function in all directions
  !!
  !! The equilibrium distribution function is:\n
  !! \[ f^{eq}_i = w_i \rho ( 1 + \frac{\vec c_i \cdot \vec u}{c^2_s}
  !!                      + \frac{ {(\vec c_i \cdot \vec u)}^2}{2c^4_s}
  !!                      - \frac{\vec u \cdot \vec u}{2c^2_s}) \]
  !!
  !! where \(w_i\) is the weight in each direction,\n
  !! \(\rho\) is the macroscopic value of density,\n
  !! \(c_s\) is the speed of sound,\n
  !! \(\vec c_i\) is the lattice unit velocity in each direction,\n
  !! \(\vec u\) is the macroscopic value of velocity.
  !!
  pure function getEqByDensVel( dens, vel, layout ) result( equil )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: dens
    real(kind=rk), intent(in) :: vel(3)
    real(kind=rk)             :: equil(layout%fStencil%QQ) !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: ucx, usq
    integer :: iDir
    ! ---------------------------------------------------------------------------

    ! square of velocity
    usq = vel(1)*vel(1) + vel(2)*vel(2) + vel(3)*vel(3)

    do iDir = 1, layout%fStencil%QQ

      ! velocity times lattice unit velocity
      ucx =   layout%fStencil%cxDirRK(1, iDir) * vel(1) &
        &   + layout%fStencil%cxDirRK(2, iDir) * vel(2) &
        &   + layout%fStencil%cxDirRK(3, iDir) * vel(3)

      ! calculate equilibrium density
      equil(iDir) =   layout%weight(iDir) * dens * ( 1.d0 + ucx*cs2inv &
        &           + ucx*ucx*cs2inv*cs2inv*div1_2                      &
        &           - usq*cs2inv*div1_2 )

    enddo

  end function getEqByDensVel
! ****************************************************************************** !


! ****************************************************************************** !
  pure function getEqByDensVelPerDir( dens, vel, cx, weight ) result( equil )
    ! --------------------------------------------------------------------------
    real(kind=rk), intent(in) :: dens
    real(kind=rk), intent(in) :: vel(3)
    real(kind=rk), intent(in) :: cx(3)
    real(kind=rk), intent(in) :: weight
    real(kind=rk)             :: equil !< return value
    ! --------------------------------------------------------------------------
    real(kind=rk) :: ucx, usq
    ! --------------------------------------------------------------------------
    ! square of velocity
    usq = vel(1)*vel(1) + vel(2)*vel(2) + vel(3)*vel(3)

    ! velocity times lattice unit velocity
    ucx =   cx(1) * vel(1) + cx(2) * vel(2) + cx(3) * vel(3)

    ! calculate equilibrium density
    equil = weight * dens * ( 1.d0 + ucx*cs2inv &
      &   + ucx*ucx*cs2inv*cs2inv*div1_2        &
      &   - usq*cs2inv*div1_2 )

  end function getEqByDensVelPerDir
! ****************************************************************************** !

! ****************************************************************************** !
  !> Incompressible formulation of the equilibrium for specific direction.
  !!
  !! @warning This function is deprecated. Use the macro EqDir_incomp instead.
  pure function getEqByDensVelPerDir_incomp( dens, vel, cx, weight ) &
    & result( equil )
    ! --------------------------------------------------------------------------
    real(kind=rk), intent(in) :: dens
    real(kind=rk), intent(in) :: vel(3)
    real(kind=rk), intent(in) :: cx(3)
    real(kind=rk), intent(in) :: weight
    real(kind=rk)             :: equil !< return value
    ! --------------------------------------------------------------------------
    real(kind=rk) :: ucx, usq
    ! --------------------------------------------------------------------------
    ! square of velocity
    usq = vel(1)*vel(1) + vel(2)*vel(2) + vel(3)*vel(3)

    ! velocity times lattice unit velocity
    ucx =   cx(1) * vel(1) + cx(2) * vel(2) + cx(3) * vel(3)

    ! calculate equilibrium density
    equil = weight * ( dens + rho0*(ucx*cs2inv &
      &   + ucx*ucx*cs2inv*cs2inv*div1_2       &
      &   - usq*cs2inv*div1_2 ) )

  end function getEqByDensVelPerDir_incomp
! ****************************************************************************** !

! ****************************************************************************** !
  !> Calculate the equilibrium distribution function in all directions
  !!
  !! The equilibrim distribution function is:\n
  !! \[ f^{eq}_i = w_i \rho ( 1 + \frac{\vec c_i \cdot \vec u}{c^2_s}
  !!                      + \frac{ {(\vec c_i \cdot \vec u)}^2}{2c^4_s}
  !!                      - \frac{\vec u \cdot \vec u}{2c^2_s}) \]\n
  !! where \(w_i\) is the weight in each direction,\n
  !! \(\rho\) is the macroscopic value of density,\n
  !! \(c_s\) is the speed of sound,\n
  !! \(\vec c_i\) is the lattice unit velocity in each direction,\n
  !! \(\vec u\) is the macroscopic value of velocity.
  !!
  pure function getEquilibrium_forPdfSubset( subset, layout, varPos )          &
    &                                                            result( equil )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: subset(:)  !< pdf array
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    real(kind=rk) :: equil(layout%fStencil%QQ) !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: rho, vel(3)
    real(kind=rk) :: ucx, usq
    integer :: iDir
    ! ---------------------------------------------------------------------------

    rho = getDensity_forPdfSubset(  subset, layout%fStencil, varPos )
    vel = getVelocity_forPdfSubset( subset, layout%fStencil, varPos )
    ! square of velocity
    usq = vel(1)*vel(1) + vel(2)*vel(2) + vel(3)*vel(3)

    do iDir = 1, layout%fStencil%QQ

      ! velocity times lattice unit velocity
      ucx =   layout%fStencil%cxDirRK(1, iDir)*vel(1) &
        &   + layout%fStencil%cxDirRK(2, iDir)*vel(2) &
        &   + layout%fStencil%cxDirRK(3, iDir)*vel(3)

      ! calculate equilibrium density
      equil( iDir ) = layout%weight( iDir ) * rho * ( 1.d0 + ucx*cs2inv        &
        &           + ucx*ucx*cs2inv*cs2inv*0.5d0                              &
        &           - usq*cs2inv*0.5d0 )

    enddo

  end function getEquilibrium_forPdfSubset
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the equilibrium distribution function in all directions
  !!
  !! The equilibrim distribution function is:\n
  !! \[ f^{eq}_i = w_i \rho ( 1 + \frac{\vec c_i \cdot \vec u}{c^2_s}
  !!                      + \frac{ {(\vec c_i \cdot \vec u)}^2}{2c^4_s}
  !!                      - \frac{\vec u \cdot \vec u}{2c^2_s}) \]\n
  !! where \(w_i\) is the weight in each direction,\n
  !! \(\rho\) is the macroscopic value of density,\n
  !! \(c_s\) is the speed of sound,\n
  !! \(\vec c_i\) is the lattice unit velocity in each direction,\n
  !! \(\vec u\) is the macroscopic value of velocity.
  !!
  pure function getEquilibrium_forElemfromState( state, elem, layout,          &
    &                                  varPos, nScalars, neigh ) result( equil )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: state(:)   !< pdf array
    integer, intent(in)       :: elem       !< treeID of the target element
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    integer, intent(in) :: nScalars !< number of scalars in global system
    integer, intent(in) :: neigh(:)   !< connectivity vector
    real(kind=rk) :: equil(layout%fStencil%QQ) !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: rho, vel(3)
    real(kind=rk) :: ucx, usq
    integer :: iDir
    ! ---------------------------------------------------------------------------
    rho = getDensity( state, elem, layout%fStencil, varPos, nScalars )
    vel(:) = getVelocity( state, elem, layout%fStencil, varPos, nScalars )

    ! square of velocity
    usq = vel(1)*vel(1) + vel(2)*vel(2) + vel(3)*vel(3)

    do iDir = 1, layout%fStencil%QQ

      ! velocity times lattice unit velocity
      ucx =   layout%fStencil%cxDirRK( 1, iDir )*vel(1)          &
        &   + layout%fStencil%cxDirRK( 2, iDir )*vel(2)          &
        &   + layout%fStencil%cxDirRK( 3, iDir )*vel(3)

      ! calculate equilibrium density
      equil( iDir ) = layout%weight( iDir ) * rho * ( 1.d0 + ucx*cs2inv        &
        &           + ucx*ucx*cs2inv*cs2inv*0.5d0                              &
        &           - usq*cs2inv*0.5d0 )

    enddo

  end function getEquilibrium_forElemfromState
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Manuel Hasert
  !! Calculate the viscous shear stress (exclude pressure)
  !! This function returns a one-dimensional array with 6 entries:
  !! tau(1:6) = [ Sxx, Syy, Szz, Sxy, Syz, Sxz ]
  !!
  !! The formula is:
  !! \[ \tau_{\alpha \beta}=
  !! -(1-\frac{\omega}{2}) \sum_{i} f^{neq}_{i} c_{i\alpha} c_{i\beta} \]\n
  !! where \( \tau_{\alpha \beta}\) is the stress
  !! in the \(\beta\)-direction on a face normal to the \(\alpha\)-axis,\n
  !! \( f^{neq}_i = f_i - f^{eq}_i\) is the non-equilibirium density.\n
  !! For more information, please refer to:\n
  !! Krueger T, Varnik F, Raabe D. Shear stress in lattice Boltzmann
  !! simulations. Physical Review E. 2009;79(4):1-14.
  !!
  pure function getShearStressTensorIncomp_forPdfSubset( subset, omega, layout,&
    &                                rho0 ) result( tau )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: subset(:)  !< pdf array
    real(kind=rk), intent(in) :: omega, rho0
    real(kind=rk)  :: tau(6) !< Shear stress includes xx, xy, xz, yy, yz, zz
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: fEq(layout%fStencil%QQ), density,vel(3)
    real(kind=rk) :: fNeq(layout%fStencil%QQ)
    integer       :: iDir
    ! ---------------------------------------------------------------------------

    density = sum(subset)
    do iDir = 1, layout%fStencil%nDims
      vel(iDir)  = sum( subset(:) * dble(layout%fStencil%cxDir(iDir,:)) )
    end do

    fEq(:) = getEquilibriumIncomp( density, vel, layout, rho0 )
    fNeq(:) = ( subset(:) - fEq(:) ) * convPrePost( omega )
    tau(:) =   secondMom( layout%fStencil%cxcx, fNeq, layout%fStencil%QQ ) &
      &      * ( omega * div1_2 - 1._rk )

  end function getShearstressTensorIncomp_forPdfSubset
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Manuel Hasert
  !! Calculate the viscous shear stress (exclude pressure)
  !! This function returns a one-dimensional array with 6 entries:
  !! tau(1:6) = [ Sxx, Syy, Szz, Sxy, Syz, Sxz ]
  !!
  !! The formula is:
  !! \[ \tau_{\alpha \beta}=
  !! -(1-\frac{\omega}{2}) \sum_{i} f^{neq}_{i} c_{i\alpha} c_{i\beta} \]\n
  !! where \( \tau_{\alpha \beta}\) is the stress
  !! in the \(\beta\)-direction on a face normal to the \(\alpha\)-axis,\n
  !! \( f^{neq}_i = f_i - f^{eq}_i\) is the non-equilibirium density.\n
  !! For more information, please refer to:\n
  !! Krueger T, Varnik F, Raabe D. Shear stress in lattice Boltzmann
  !! simulations. Physical Review E. 2009;79(4):1-14.
  !!
  pure function getShearStressTensor_forPdfSubset( subset, omega, layout,      &
    &                                                     varPos ) result( tau )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: subset(:)  !< pdf array
    real(kind=rk), intent(in) :: omega
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    real(kind=rk)  :: tau(6) !< Shear stress includes xx, xy, xz, yy, yz, zz
    ! ---------------------------------------------------------------------------
    real(kind=rk) ::  fEq(layout%fStencil%QQ)
    real(kind=rk) :: fNeq(layout%fStencil%QQ)
    ! ---------------------------------------------------------------------------

    fEq  = getEquilibrium( subset, layout, varPos )
    fNeq = ( subset - fEq ) * convPrePost( omega )
    tau  =   secondMom( layout%fStencil%cxcx, fNeq, layout%fStencil%QQ ) &
      &    * ( omega * div1_2 - 1._rk )

  end function getShearstressTensor_forPdfSubset
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the viscous shear stress (exclude pressure)
  !! This function returns a one-dimensional array with 6 entries:
  !! tau(1:6) = [ Sxx, Syy, Szz, Sxy, Syz, Sxz ]
  !!
  !! The formula is:
  !! \[ \tau_{\alpha \beta}=
  !! -(1-\frac{\omega}{2}) \sum_{i} f^{neq}_{i} c_{i\alpha} c_{i\beta} \]\n
  !! where \( \tau_{\alpha \beta}\) is the stress
  !! in the \(\beta\)-direction on a face normal to the \(\alpha\)-axis,\n
  !! \( f^{neq}_i = f_i - f^{eq}_i\) is the non-equilibirium density.\n
  !! For more information, please refer to:\n
  !! Krueger T, Varnik F, Raabe D. Shear stress in lattice Boltzmann
  !! simulations. Physical Review E. 2009;79(4):1-14.
  !!
  pure function getShearStressTensor_forElemFromState( state, neigh, elem,     &
    &                    omega, layout, iField, varPos, nScalars ) result( tau )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: state(:)  !< pdf array
    real(kind=rk), intent(in) :: omega
    integer, intent(in) :: neigh(:)   !< connectivity vector
    integer, intent(in) :: elem   !< treeID of target element
    integer, intent(in) :: iField !< current field
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk)  :: tau(6) !< Shear stress includes xx, xy, xz, yy, yz, zz
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: fEq( layout%fStencil%QQ)
    real(kind=rk) :: fNeq(layout%fStencil%QQ)
    integer       :: iDir
    ! ---------------------------------------------------------------------------
    integer :: nElems, QQ
    ! ---------------------------------------------------------------------------
    QQ = layout%fStencil%QQ
    nElems = size( neigh ) / QQ

    ! Get the equilibrium density which is required by shear stress
    fEq = getEquilibrium( state, elem, layout, varPos, nScalars, neigh )

    do iDir = 1, QQ
      fNeq(iDir) = state( ?SAVE?( iDir, iField, elem, QQ, nScalars, nElems, neigh ) ) &
        &          - fEq( iDir )
    enddo

    ! compute shear stress
    tau = secondMom( layout%fStencil%cxcx, fNeq, QQ )
    tau = ( omega * div1_2 - 1._rk ) * tau * convPrePost( omega )

  end function getShearstressTensor_forElemFromState
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the shear rate tensor (strain rate) by acoustic scaling.
  !! @todo: This function actually uses diffusive scaling to calculate shear
  !!       rate.
  !!
  pure function getShearRateTensor_acoustic_forPdfSubset( subset, omega,       &
    &                            layout, incompressible ) result( S )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: subset(:)  !< pdf array
    real(kind=rk), intent(in) :: omega
    real(kind=rk) :: S(3,3)
    logical, intent(in), optional :: incompressible
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: vel(3) ! velocity components
    real(kind=rk) :: diagVal(3,3), rho, rho0
    real(kind=rk) :: fTmp(layout%fStencil%QQ), rhoD
    integer       :: iDir, iVal, jVal, order(3)
    logical       :: incomp
    ! ---------------------------------------------------------------------------
    incomp = .false.
    if ( present( incompressible )) incomp = incompressible

    Rho  = 0._rk
    Vel  = 0._rk
    do iDir = 1, layout%fStencil%QQ
      fTmp( iDir ) = subset( iDir )
      vel(:) = vel(:) + fTmp( iDir ) * dble(layout%fStencil%cxDir(:,iDir))
    end do
    rho = sum( fTmp(:) )

    if ( .not. incomp ) then
      vel  = vel / rho
      rhoD = rho
    else
      rhoD = rho0
    end if

    ! Get the second velocity moments of the source element's pdf
    diagVal(:,:) = 0._rk
    diagVal(1,1) = cs2*Rho
    diagVal(2,2) = cs2*Rho
    diagVal(3,3) = cs2*Rho
    do jVal = 1, layout%fStencil%nDims
      do iVal = 1, layout%fStencil%nDims
        order = 0
        order(iVal) = order(iVal) +1
        order(jVal) = order(jVal) +1
        S(iVal,jVal) = get_moment( layout%fStencil%QQ, layout%fStencil%cxDir, order, fTmp)

        ! Now compute the S(1) from the second velocity moment
        S(iVal,jVal) = rhoD * Vel(iVal) * Vel(jVal) + diagVal(iVal,jVal) - S(iVal,jVal)
      end do
    end do

    S(:,:) = S(:,:) * omega * cs2inv * convPrePost(omega) * div1_2

  end function getShearRateTensor_acoustic_forPdfSubset
! ****************************************************************************** !


! ****************************************************************************** !
  !> This routine calculates shear rate tensor (i.e. strain rate tensor) by
  !! acoustic scaling (i.e. CE analysis)
  pure function getShearRateTensor_acoustic_incomp( subset, omega, layout,&
    &                                                       rho0)    result( S )
    ! ---------------------------------------------------------------------------
    real(kind=rk), intent(in) :: subset(:)  !< pdf array
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: omega
    real(kind=rk), intent(in) :: rho0
    real(kind=rk) :: S(3,3) !< output
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: rho, vel(3), strain(6)
    real(kind=rk) :: fTmp(layout%fStencil%QQ)
    real(kind=rk) :: fEq(layout%fStencil%QQ)
    real(kind=rk) :: fNeq(layout%fStencil%QQ)
    integer       :: iDir
    ! ---------------------------------------------------------------------------

    Rho  = 0._rk
    Vel  = 0._rk
    S  = 0._rk

    do iDir = 1, layout%fStencil%QQ
      fTmp( iDir ) = subset( iDir )
      vel(:) = vel(:) + fTmp( iDir ) * layout%fStencil%cxDirRK(:,iDir)
    end do
    rho = sum( fTmp(:) )

    fEq(:) = getEquilibriumIncomp( rho, vel, layout, rho0 )
    fNeq   = ( fTmp - fEq ) * convPrePost( omega )
    strain(:) = secondMom( layout%fStencil%cxcx, fNeq, layout%fStencil%QQ )
    strain(:) =  -1.5_rk * omega / rho0 * strain

    ! re-arrange vector into tensor
    if ( layout%fStencil%nDims == 3 ) then
      S(1,1) = strain(1)
      S(2,2) = strain(2)
      S(3,3) = strain(3)
      S(1,2) = strain(4)
      S(2,1) = strain(4)
      S(2,3) = strain(5)
      S(3,2) = strain(5)
      S(1,3) = strain(6)
      S(3,1) = strain(6)
    else
      S(1,1) = strain(1)
      S(2,2) = strain(2)
      S(1,2) = strain(3)
      S(2,1) = strain(3)
    end if

  end function getShearRateTensor_acoustic_incomp
! ****************************************************************************** !

! ****************************************************************************** !
  !> This routine calculates shear rate tensor (i.e. strain rate tensor) by
  !! acoustic scaling (i.e. CE analysis)
  pure function getShearRateTensor_acoustic_lbm( subset, omega, layout ) &
    &                                                                result( S )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: subset(:)  !< pdf array
    real(kind=rk), intent(in) :: omega
    real(kind=rk) :: S(3,3)
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: rho, vel(3), strain(6)
    real(kind=rk) :: fTmp(layout%fStencil%QQ)
    real(kind=rk) :: fEq(layout%fStencil%QQ)
    real(kind=rk) :: fNeq(layout%fStencil%QQ)
    integer       :: iDir
    ! ---------------------------------------------------------------------------

    rho = 0._rk
    Vel = 0._rk
    S   = 0._rk

    do iDir = 1, layout%fStencil%QQ
      fTmp( iDir ) = subset( iDir )
      vel(:) = vel(:) + fTmp( iDir ) * layout%fStencil%cxDirRK(:,iDir)
    end do
    rho = sum( fTmp(:) )
    vel = vel / rho

    fEq(:) = getEquilibrium( rho, vel, layout )
    fNeq   = ( fTmp - fEq ) * convPrePost( omega )
    strain(:) = secondMom( layout%fStencil%cxcx, fNeq, layout%fStencil%QQ )
    strain(:) =  -1.5_rk * omega / rho * strain

    ! re-arrange vector into tensor
    if ( layout%fStencil%nDims == 3 ) then
      S(1,1) = strain(1)
      S(2,2) = strain(2)
      S(3,3) = strain(3)
      S(1,2) = strain(4)
      S(2,1) = strain(4)
      S(2,3) = strain(5)
      S(3,2) = strain(5)
      S(1,3) = strain(6)
      S(3,1) = strain(6)
    else
      S(1,1) = strain(1)
      S(2,2) = strain(2)
      S(1,2) = strain(3)
      S(2,1) = strain(3)
    end if

  end function getShearRateTensor_acoustic_lbm
! ****************************************************************************** !

! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the strain rate tensor through 2nd moment.
  !! This function returns a two-dimensional 3 x 3 symmetirc array:
  !!```
  !! S = | Sxx, Sxy, Sxz,|
  !!     | Syx, Syy, Syz |
  !!     | Szx, Szy, Szz |
  !!```
  !! @todo Move the equation latex to a seperate place, where it can be refered
  !!       by serevel places.
  !!
  !! The formula to compute this can be found in paper of Junk 2005
  !! after equation A.15
  !! \[
  !!    S^{(1)} = \frac{\omega}{\kappa c_s^2}
  !!                  ( u^{(1)} \otimes u^{(1)} + c_s^2 \rho^{(2)}I - p^{(2)} )
  !! \]
  !! where \( p^{(2)} \) is the second velocity moment,
  !! \( I \) is the identity matrix and \( \kappa=1 \).\n
  !! Note that \( S^{(1)} \) is not shear rate \( S \) yet.
  !! \[
  !!    S = \frac{1}{2} S^{(1)}
  !! \]
  !!
  pure function getShearRateTensor_diffusive_forPdfSubset( f, omega, layout ) &
    &                                              result( S )
    ! ---------------------------------------------------------------------------
    !> stencil layout
    type( mus_scheme_layout_type ), intent(in) :: layout
    !> pdf array ( post-collision value )
    real(kind=rk), intent(in) :: f(layout%fStencil%QQ)
    !> relaxation parameter
    real(kind=rk), intent(in) :: omega
    !> output array: strain rate tensor
    real(kind=rk) :: S(3,3)
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: vel(3), diagVal(3,3), rho
    integer       :: iVal, jVal, order(3)
    ! ---------------------------------------------------------------------------

    Rho    = sum( f(:) )
    Vel(1) = sum( f * layout%fStencil%cxDirRK(1,:) )
    Vel(2) = sum( f * layout%fStencil%cxDirRK(2,:) )
    Vel(3) = sum( f * layout%fStencil%cxDirRK(3,:) )

    !! \todo: this section can become a function that be called by several
    !!        places.
    diagVal(:,:) = 0._rk
    diagVal(1,1) = cs2 * Rho
    diagVal(2,2) = cs2 * Rho
    diagVal(3,3) = cs2 * Rho

    ! Get the second velocity moments of the source element's pdf
    do jVal = 1, layout%fStencil%nDims
      do iVal = 1, layout%fStencil%nDims

        order = 0
        order( iVal ) = order( iVal ) + 1
        order( jVal ) = order( jVal ) + 1

        ! Calculate second velocity moment
        S(iVal,jVal) = - get_moment( layout%fStencil%QQ, layout%fStencil%cxDir, order, f) &
          &            + vel( iVal ) * vel( jVal )      &
          &            + diagVal( iVal, jVal )

      end do
    end do

    ! Multiply factor and convert to pre-collision value
    s = s * omega * cs2inv * convPrePost( omega ) * div1_2

  end function getShearRateTensor_diffusive_forPdfSubset
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the viscous shear stress (exclude pressure)
  !!
  !! The shear stress formula is:\n
  !! \[
  !! \tau_{\alpha \beta}=
  !!      -(1-\frac{\omega}{2}) \sum_{i} f^{neq}_{i} c_{i\alpha} c_{i\beta}
  !! \]\n
  !! where \( \tau_{\alpha \beta} \) is the stress
  !! in the \(\beta\)-direction on a face normal to the \(\alpha\)-axis,\n
  !! \( f^{neq}_i = f_i - f^{eq}_i \) is the non-equilibirium density.\n
  !! For more information, please refer to:\n
  !! Krueger T, Varnik F, Raabe D. Shear stress in lattice Boltzmann
  !! simulations. Physical Review E. 2009;79(4):1-14.
  !!
  !! Notice:\n
  !! Normally shear stress should be calculated by PRE-collision pdf.\n
  !! But this routine utilize POST-collision pdf (SAVE)!
  !!
  pure function getShearStressPost( state, neigh, elem, omega, layout, iField, &
    &                          nScalars ) result( tau )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: state(:)  !< pdf array
    real(kind=rk), intent(in) :: omega
    integer, intent(in) :: neigh(:)   !< connectivity vector
    integer, intent(in) :: elem   !< treeID of target element
    integer, intent(in) :: iField !< current field
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk)  :: tau(6) !< Shear stress includes xx, yy, zz, xy, yz, xz
    ! ---------------------------------------------------------------------------
    ! local variables
    real(kind=rk) :: equil(layout%fStencil%QQ), nonEq(layout%fStencil%QQ)
    real(kind=rk) :: rho, vel(3)
    real(kind=rk) :: fTmp(layout%fStencil%QQ) !< pdf array of current element
    integer       :: iDir
    ! ---------------------------------------------------------------------------
    integer :: nElems, QQ
    ! ---------------------------------------------------------------------------
    QQ = layout%fStencil%QQ
    nElems = size( neigh ) / QQ

    rho = 0._rk
    vel = 0._rk
    do iDir = 1, QQ
      ! use POST-collision pdf
      fTmp(iDir) = state( ?SAVE?( iDir, iField, elem, QQ, nScalars, nElems, neigh ) )
      ! compute density and velocity
      rho = rho + fTmp(iDir)
      vel(:) = vel(:) + fTmp(iDir) * layout%fStencil%cxDirRK(:,iDir)
    end do

    ! compute Eq and nonEq
    equil(:) = getEquilibrium( rho, vel, layout )
    nonEq = equil - fTmp

    ! compute shear stress
    tau(:) = secondMom( layout%fStencil%cxcx(:,:), nonEq(:), QQ )
    tau(:) = ( omega * div1_2 - 1._rk ) * tau(:)

  end function getShearStressPost
! ****************************************************************************** !

! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the Shear Rate
  !!
  !! The Shear Rate is defined as
  !! \[
  !!  \dot{\gamma} = 2\sqrt{ D_{II} }
  !! \]
  !! where \( D_{II} \) is the second invariant of the strain rate tensor and
  !! defined as
  !! \[
  !!    D_{II} = \sum^{l}_{\alpha,\beta=l} S_{\alpha\beta} S_{\alpha\beta}
  !! \]
  !! where \( S_{\alpha\beta} \) is the strain rate tensor.
  !!
  pure function getShearRate( strain ) result( shear )
    ! ---------------------------------------------------------------------------
    !> strain rate tensor: xx, yy, zz, xy, yz, zx
    real(kind=rk), intent(in) :: strain(:)
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: shear !< shear rate
    ! ---------------------------------------------------------------------------

    shear = sqrt( sum( strain(:) * strain(:) ) ) * 2._rk

  end function getShearRate
! ****************************************************************************** !

! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the wall shear stress (WSS)
  !!
  !! @todo: add formule and reference!
  function getWSS( state, neigh, elem, omega, layout, iField, varPos,          &
    &                                                   nScalars ) result( wss )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: state(:)
    integer, intent(in) :: neigh(:)   !< connectivity vector
    real(kind=rk), intent(in) :: omega
    integer, intent(in)       :: elem
    integer, intent(in) :: iField !< current field
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk) :: wss !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: tau(6) ! shear stress components
    real(kind=rk) :: a0, a1, a2
    real(kind=rk) :: q_stress, q_stress_cube, q_term, r_stress, d_stress
    real(kind=rk) :: theta_stress, cos_term, shear_1, shear_2, shear_3
    ! ---------------------------------------------------------------------------
    wss = 0._rk
    tau(:) = getShearstressTensor(state, neigh, elem, omega, layout, iField,   &
      &                           varPos, nScalars)
!    write(logUnit(1),*)'WARNING: This is still the old order of the shear stress!!'
!    write(logUnit(1),*)'PLease change in getWSS !'
    a2 = -1.d0 * ( tau(1) + tau(4) + tau(6) )
    a1 = tau(1) * tau(4) + tau(4) * tau(6) + tau(1) * tau(6)                   &
      &- tau(2) ** 2 - tau(3) ** 2 - tau(5) ** 2
    a0 = -1.d0 * ( tau(1)*tau(4)*tau(6) + 2.d0*tau(2)*tau(3)*tau(5)            &
      &- tau(1)*tau(5)**2 - tau(4)*tau(3)**2 - tau(6)*tau(2)**2 )

    q_stress = div1_9 * ( 3._rk*a1 - a2*a2 )
    if ( q_stress < 0._rk ) then
      q_term = 2._rk * sqrt( -q_stress )
    else
      q_term = 0._rk
    end if
    r_stress = div1_54 * ( 9._rk*a2*a1 - 27._rk*a0 - 2._rk*a2*a2*a2 )
    q_stress_cube = q_stress * q_stress * q_stress
    d_stress = q_stress_cube + r_stress * r_stress

    if ( d_stress < 0._rk ) then
      ! the solutions are real and unequal
      theta_stress = div1_3 * acos( r_stress / sqrt( -q_stress_cube ))

      cos_term = div3_4 * q_term * cos( theta_stress )

      ! three values of maximum shear stress
      shear_1 = abs( sqrt3 * div1_2 * q_term * sin( theta_stress ))

      shear_2 = abs( cos_term - shear_1*div1_2 )

      shear_3 = abs( cos_term + shear_1*div1_2 )

      !! choose the max shear as Wall Shear Stress(wss)
      wss = max(shear_1, shear_2, shear_3)
    end if

  end function getWSS
! ****************************************************************************** !

! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the wall shear stress (WSS)
  !!
  !! @todo: add formule and reference!
  function getWSS_vonMises( state, neigh, elem, omega, layout, iField,         &
    &                       nScalars ) result( wss )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: state(:)
    integer, intent(in) :: neigh(:)   !< connectivity vector
    real(kind=rk), intent(in) :: omega
    integer, intent(in)       :: elem
    integer, intent(in) :: iField !< current field
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk) :: wss !< return value
    ! ---------------------------------------------------------------------------
    ! local variables
    real(kind=rk) :: tau(6) ! shear stress components
    ! ---------------------------------------------------------------------------
    wss = 0._rk
    tau(:) = getShearstressPost( state, neigh, elem, omega, layout, iField,    &
      &                                                              nScalars )
    ! Check wikipedia article on von Mises stress criterion
    wss = (tau(1)-tau(2))**2 + (tau(2)-tau(3))**2 + (tau(1)-tau(3))**2         &
      & + 6._rk*( tau(4)*tau(4) + tau(5)*tau(5) + tau(6)*tau(6))
    wss = sqrt( wss * div1_2 )

  end function getWSS_vonMises
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the wall shear stress (WSS) 2D
  !!
  !! @todo: add formule and reference!
  function getWSS2D( state, neigh, elem, omega, layout, iField, varPos,        &
    &                                                   nScalars ) result( wss )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: state(:)
    integer, intent(in) :: neigh(:)   !< connectivity vector
    real(kind=rk), intent(in) :: omega
    integer, intent(in)       :: elem
    integer, intent(in) :: iField !< current field
    integer, intent(in) :: varPos(:) !< varPos of current field variable
    integer, intent(in) :: nScalars !< number of scalars in global system
    real(kind=rk) :: wss !< return value
    ! ---------------------------------------------------------------------------
    ! local variables
    real(kind=rk) :: tau(6) !< shear stress components
    real(kind=rk) :: a0, a1
    ! ---------------------------------------------------------------------------

    tau(:) = getShearstressTensor(state, neigh, elem, omega, layout, iField, &
      &                           varPos, nScalars)

    a1 = - tau(1) - tau(2)
    a0 =   tau(1) * tau(2) - tau(3) * tau(3)

    wss = sqrt( a1 * a1 - 4._rk * a0 )

  end function getWSS2D
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Setting the non-equilibrium part based on the acoustic scaling
  !!
  !! The non-equilibirium part of pdf is computed from strain rate tensor
  !! by \cite Latt:2011vr
  !! \[
  !!  f_i^{neq} = - \frac{t_i \rho_0}{c_{s}^2 \omega)}
  !!                            Q_{i\alpha\beta}:S_{\alpha\beta}
  !! \]
  !! where \( \boldsymbol{A} : \boldsymbol{B} \) is Frobenius inner product.
  !! \[
  !!   \boldsymbol{A} : \boldsymbol{B} = \sum_{i,j} A_{ij}B_{ij}
  !! \]
  !! \[
  !!  Q_{i\alpha\beta} = c_{i\alpha}c_{i\beta} - c_s^2 \delta_{\alpha\beta}
  !! \]
  !! and \( S \) is the strain rate tensor
  !! \[
  !!  S_{\alpha\beta} = -\frac{1}{2}
  !!          ( \partial_{\alpha}u_{\beta} + \partial_{\alpha}u_{\beta} )
  !! \]
  !!
  function getNEq_acoustic( layout, omega, Sxx ) result( nEq )
    ! ---------------------------------------------------------------------------
    type( mus_scheme_layout_type ), intent(in) :: layout
    real(kind=rk), intent(in) :: omega
    !> strain rate tensor
    real(kind=rk), intent(in) :: Sxx(3,3)
    real(kind=rk) :: nEq(layout%fStencil%QQ)
    ! ---------------------------------------------------------------------------
    integer :: iVal, jVal, iDir
    real(kind=rk) :: tau(3,3)
    real(kind=rk) :: nu, coeff
    ! ---------------------------------------------------------------------------

    nu = cs2 * ( 1._rk / omega - div1_2 )
    ! convert strain rate to stress
    tau(:,:) = 2._rk * nu * Sxx(:,:)
    coeff = cs4inv / ( 2._rk - omega )

    ! Recover the non-equilibrium part from stress (tau)
    nEq(:) = 0._rk
    do iDir = 1, layout%fStencil%QQ
      do jVal = 1, layout%fStencil%nDims
        do iVal = 1, layout%fStencil%nDims
          Neq( iDir ) = Neq( iDir ) + &
          &   tau(iVal,jVal) * layout%fStencil%cxDirRK(iVal,iDir) &
          &                   *layout%fStencil%cxDirRK(jVal,iDir)
        end do ! iVal
        neq( iDir ) = neq( iDir ) - cs2 * tau( jVal, jVal )
      end do ! jVal
    end do ! iDir
    neq(:) = -layout%weight(:) * neq(:) * coeff

    ! Convert from post to pre-collision
    ! KM: convert factor is zero for omega = 1.0 and dividing with 0.0
    ! leads to NaN so divide conv factor only if omega is not equal to 1
    if (omega .fne. 1.0_rk) Neq = Neq / convPrePost( omega )

  end function getNEq_acoustic
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Calculate the non-equilibrium part of pdf from strain rate tensor
  !! based on the diffusive scaling
  !!
  !! According to \cite Junk:2005cr \n
  !! The non-equilibrium part of pdf \( f_i^{neq} \) is set by \n
  !! \[
  !!   f_i^{neq} = -\frac{t_i}{2\kappa c_{s}^4} \nu' S^{(1)}:\Lambda
  !!             = -\frac{t_i}{2\kappa c_{s}^4} \frac{\kappa c_s^2}{\omega} S^{(1)}:\Lambda
  !!             = -\frac{t_i}{2 c_{s}^2} S^{(1)}:\Lambda
  !!             = -\frac{t_i}{c_{s}^2} \bm S:\Lambda
  !! \]
  !! where \( \nu' = \frac{\kappa c_s^2}{\omega} \) is the viscosity,
  !! \( \bm S = \frac{1}{2}S^{(1)} \) is the strain rate tensor and
  !! \[
  !! \Lambda_{i\alpha\beta} =
  !!    c_{i\alpha} c_{i\beta} - \frac{1}{D}\sum_{\gamma}(c_{i\gamma}c_{i\gamma})
  !!                                                      \delta_{\alpha\beta}
  !! \]
  !! and \( D\) is the number of dimension. \n
  !! Notice here that strain rate tensor above has to be a traceless tensor, i.e.
  !! \( Tr(S) = 0 \).
  !! In current implementation,
  !! the above equation is slightly modified so that
  !! the strain rate tensor is not required to be traceless anymore.
  !! In this way, \( f_i^{neq} \) calculated by this routine can recover the
  !! input strain rate tensor no matter it is traceless or not.\n
  !! Specificly the \( \Lambda \) in above equation is modified slightly, i.e.
  !! \[
  !! \Lambda_{i\alpha\beta} =
  !!                       c_{i\alpha} c_{i\beta} - c_s^2 \delta_{\alpha\beta}
  !! \]
  !! This routine has a unit test program utest/mus_fNeq_diffusive_test
  !!
  function getNEq_diffusive( layout, omega, Sxx ) result( nEq )
    ! ---------------------------------------------------------------------------
    type( mus_scheme_layout_type ), intent(in) :: layout
    real(kind=rk), intent(in) :: omega
    !> Strain rate tensor. It is a symmetric 3x3 matrix
    real(kind=rk), intent(in) :: Sxx(3,3)
    real(kind=rk) :: nEq(layout%fStencil%QQ)
    ! ---------------------------------------------------------------------------
    integer :: iVal, jVal, iDir
    real(kind=rk) :: strain(3,3)
    ! ---------------------------------------------------------------------------

    strain = Sxx
    ! First calculate the part of f_neq = strain : Lambda
    nEq(:) = 0._rk
    do iDir = 1, layout%fStencil%QQ
      ! cx2Sum = 0
      ! do iVal = 1, layout%fStencil%nDims
      !   cx2Sum = cx2Sum + layout%fStencil%cxDir(iVal,iDir)**2
      ! end do
      ! cx2Sum = cx2Sum / nDims
      do jVal = 1, layout%fStencil%nDims
        do iVal = 1, layout%fStencil%nDims
          Neq( iDir ) = Neq( iDir ) + strain( iVal, jVal )     &
            &           *layout%fStencil%cxDirRK( iVal,iDir) &
            &           *layout%fStencil%cxDirRK( jVal,iDir)
        end do

        ! By this equation, strain rate tensor is required to be traceless
        ! Neq( iDir ) = Neq( iDir ) - strain( jVal, jVal ) * real(cx2sum, rk)

        ! By the following, strain rate tensor is NOT required to be traceless
        Neq( iDir ) = Neq( iDir ) - strain( jVal, jVal ) * cs2

      end do
    end do
    ! Then multiply the pre-factor
    ! f_neq = -t_i/cs^2/omega * (strain:Lambda)
    Neq(:) = -layout%weight(:) / omega / cs2 * Neq(:)

    ! Convert from post to pre-collision
    ! KM: convert factor is zero for omega = 1.0 and dividing with 0.0
    ! leads to NaN so divide conv factor only if omega is not equal to 1
    if (omega .fne. 1.0_rk) Neq = Neq / convPrePost( omega )

  end function getNEq_diffusive
! ****************************************************************************** !


! ****************************************************************************** !
  !> Set third order for diffusive scaling (for push algorithm?!)
  !! Have a look at the initial condition file for TGV in
  !! Hasert/2013/multilevel_diffusive/ic
  !! Note that this is only for a 2d taylor green test case in the region
  !! x = y = 0:2pi
  !! with reference velocity u0 = 1
  !!
  function setThirdOrder_diffusive( layout, omega, x, dx, dt ) result( f3 )
    ! ---------------------------------------------------------------------------
    type( mus_scheme_layout_type ), intent(in) :: layout
    real(kind=rk), intent(in) :: omega
    real(kind=rk), intent(in) :: x(3), dx, dt
    real(kind=rk) :: f3(layout%fStencil%QQ)
    ! ---------------------------------------------------------------------------
    integer :: iDir
    real(kind=rk) :: cx(3), nu, u0
    real(kind=rk) :: A, BI, BII, CI, CII, D
    ! ---------------------------------------------------------------------------
    ! Reference velocity
    u0 = 1._rk
    ! A term
    nu = ( 1._rk/omega - div1_2 ) * div1_3
    do iDir = 1, layout%fStencil%QQ
      cx = real(layout%fStencil%cxDir(:,iDir),kind=rk)
      ! k = 1
      A = u0**2*nu*cs2inv*cs2inv*0.5_rk*(cx(1)*sin(2._rk*x(1))                 &
                                        +cx(2)*sin(2._rk*x(2)))*(dt**3/dx**3)
      BI  = 6._rk*nu/omega*u0*(-cx(1)*sin(x(1))*cos(x(2))+cx(2)*cos(x(1))      &
        & * sin(x(2)))*(dx*dt)
      BII = 3._rk/omega*u0**2*(cx(1)*cos(x(1))*sin(x(1))+cx(2)*cos(x(2))       &
        & * sin(x(2))) &
           *(dt**2/dx)
      CI  = -9._rk/omega*u0**2*(cx(1)**2 - cx(2)**2)                           &
        & * (cx(1)*sin(x(1))*cos(x(1))*(cos(x(2)))**2                          &
        & - cx(2)*sin(x(2))*cos(x(2))*(cos(x(1)))**2 )*(dt**2/dx)
      CII = 3._rk/omega*u0**2*( cx(1)*sin(x(1))*cos(x(1))*(cos(x(2))**2        &
        & - sin(x(2))**2) + cx(2)*sin(x(2))*cos(x(2))*(cos(x(1))**2            &
        & - sin(x(1))**2))*(dt**2/dx)
      D   = cs2inv*(1._rk/(omega**2)-0.5_rk/omega)*u0*( sin(x(1))*cos(x(2))    &
        & * cx(1) + cos(x(1))*sin(x(2))*cx(2))*(cx(2)**2 - cx(1)**2)*(dt*dx)

      f3(iDir ) = layout%weight( iDir )*(A + BI + BII + CI + CII + D)
      ! A, BI: error smaller, BII, CI: error larger-, CII: error larger+,
      ! D: error ! little larger-
      !f3(iDir ) = layout%weight( iDir )*(A)
    end do
    ! ... and convert for pre/post collision correction
    f3 = f3/convPrePost( omega )

  end function setThirdOrder_diffusive
! ****************************************************************************** !


! ****************************************************************************** !
  !> author: Jiaxing Qi
  !! Conversion factor betwen the pre- and post-collision quantity for the shear
  !! stress.
  !!
  !! Shear stress calculation requires the non-equilibirium value of pdf before
  !! collision. However that value not may be accessable directly when PULL
  !! scheme is uitilized, as only pdf after collision is available.
  !! So this conversion factor is introduced to help
  !! recover fNeq before collision from fNeq after collision as long as
  !! relaxation parameter (omage) does not equal to 1.0. When omage equals to 1,
  !! this conversion factor is set to be 0.
  !!
  !! How to use this pre-factor?
  !!```
  !!  shearstress = convPrePost(omega) * omega * cs2inv * shearLB_postColl
  !!```
  !!
  pure function convPrePost( omega ) result( conv )
    ! ---------------------------------------------------------------------------
    real(kind=rk), intent(in) :: omega !< relaxation parameter
    real(kind=rk) :: conv !< conversion factor
    ! ---------------------------------------------------------------------------

??  if( PUSH ) then
      conv = 1._rk
??  else
      conv = 1._rk / ( 1._rk - omega )
??  endif
  end function convPrePost
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the conversion factor for nonEq in difference levels
  pure function getNonEqFac( omegaS, omegaT ) result( fac )
    ! ---------------------------------------------------------------------------
    real(kind=rk), intent(in) :: omegaS  !< omega value on source level
    real(kind=rk), intent(in) :: omegaT  !< omage value on target level
    real(kind=rk) :: fac
    ! ---------------------------------------------------------------------------
!KM: \todo check nonEqfac with paper "Interpretation and derivation of lattice
! bolzmann method with string spliting
!    fac = (2.0_rk-omegaS) / (2.0_rk-omegaT)
??  if( PUSH ) then
      fac = ( 1._rk - 0.5_rk * omegaS ) / ( 1._rk - 0.5_rk * omegaT )
??  else
      fac =   ( 1._rk - 0.5_rk * omegaS ) * ( 1._rk - omegaT ) &
        &   / ( 1._rk - omegaS ) / ( 1._rk - 0.5_rk * omegaT )
??  endif

  end function getNonEqFac
! ****************************************************************************** !

  !> Calculate the conversion factor to convert nonEq moments
  !! between fine and coarser.
  pure function getNonEqFac_intp( omegaS, omegaT ) result( fac )
    ! ---------------------------------------------------------------------------
    real(kind=rk), intent(in) :: omegaS  !< omega value on source level
    real(kind=rk), intent(in) :: omegaT  !< omage value on target level
    real(kind=rk) :: fac
    ! ---------------------------------------------------------------------------
??  if( PUSH ) then ! for pre-collision PDF
      fac = omegaS / omegaT
??  else ! for post-collision PDF
      fac =   omegaS * ( 1._rk - omegaT ) &
        &   / ( ( 1._rk - omegaS ) * omegaT )
??  endif

  end function getNonEqFac_intp
! ****************************************************************************** !

! ****************************************************************************** !
  !> Calculate the equilibrium distribution function in all directions
  !! This is the incompressible formulation with reference density rho0
  !!
  !! The equilibrium distribution function is:\n
  !! \[ f^{eq}_i = w_i  ( \rho + \frac{\vec c_i \cdot \vec u}{c^2_s}
  !!                      + \frac{ {(\vec c_i \cdot \vec u)}^2}{2c^4_s}
  !!                      - \frac{\vec u \cdot \vec u}{2c^2_s}) \]\n
  !! where \(w_i\) is the weight in each direction,\n
  !! \(\rho = \sum_i f_i\) is the macroscopic density,\n
  !! \(c_s\) is the speed of sound,\n
  !! \(\vec c_i\) is the lattice unit velocity in each direction,\n
  !! \(\vec u = \sum_i c_i f_i\) is the macroscopic value of velocity.
  !!
  pure function getEquilibriumIncomp( dens, vel, layout, rho0 )  result( equil )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: dens
    real(kind=rk), intent(in) :: rho0
    real(kind=rk), intent(in) :: vel(3)
    real(kind=rk)             :: equil(layout%fStencil%QQ) !< return value
    ! ---------------------------------------------------------------------------
    ! local variables
    real(kind=rk) :: ucx, usq
    integer :: iDir
    ! ---------------------------------------------------------------------------

    ! square of velocity
    usq = vel(1)*vel(1) + vel(2)*vel(2) + vel(3)*vel(3)

    do iDir = 1, layout%fStencil%QQ

      ! velocity times lattice unit velocity
      ucx =   layout%fStencil%cxDirRK(1,iDir) * vel(1) &
        &   + layout%fStencil%cxDirRK(2,iDir) * vel(2) &
        &   + layout%fStencil%cxDirRK(3,iDir) * vel(3)

      ! calculate equilibrium density
      equil( iDir ) =   layout%weight( iDir )         &
        &             * ( dens + rho0*( ucx*cs2inv    &
        &             + ucx*ucx*cs2inv*cs2inv*div1_2  &
        &             - usq*cs2inv*div1_2 ))

    end do

  end function getEquilibriumIncomp
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the distribution function in all directions
  !! by using the fEq + fNeq
  !!
  !! The input vector must look like
  !! 2D:
  !!```
  !! mom = [ rho, ux, uy, Sxx, Syy, Sxy ]
  !! mom = [ rho, ux, uy, uz, Sxx, Syy, Szz, Sxy, Syz, Sxz ]
  !!```
  !!
  function set_pdfDiffusive( layout, omega, rho0, mom )  result( pdf )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: mom(layout%fStencil%QQ) !< input moments
    real(kind=rk), intent(in) :: omega
    real(kind=rk), intent(in) :: rho0
    real(kind=rk)             :: pdf(layout%fStencil%QQ) !< return value
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: fEq(layout%fStencil%QQ)  ! eq distribution
    real(kind=rk) :: fnEq(layout%fStencil%QQ) ! neq distribution
    ! ---------------------------------------------------------------------------
    if ( layout%fStencil%nDims == 3 ) then
      ! Compute the equilibrium
      fEq(:) = getEquilibriumIncomp( dens = mom(1),               &
        &                            vel  = mom(2:4),             &
        &                            rho0 = rho0, layout = layout )
      ! Compute the non-equilibrium
      fNeq = getNeq_diffusive( layout = layout,     &
        &  omega = omega, Sxx = [                   &
        &               mom(5), mom(8), mom(10),    &
        &               mom(8), mom(6), mom(9),     &
        &               mom(10), mom(9), mom(7) ] )
    else
      ! Compute the equilibrium
      fEq(:) = getEquilibriumIncomp( dens = mom(1),               &
        &                            vel  = [ mom(2), mom(3), 0._rk ],    &
        &                            rho0 = rho0, layout = layout )
      ! Compute the non-equilibrium
      fNeq = getNeq_diffusive(  layout = layout,                &
        &                       omega = omega,                  &
        &                       Sxx = [ mom(4), mom(6), 0._rk,  &
        &                               mom(6), mom(5), 0._rk,  &
        &                               0._rk,  0._rk,  0._rk ] )
    end if
    pdf(:) = fEq(:) + fNeq(:)

  end function set_pdfDiffusive
! ****************************************************************************** !


! ****************************************************************************** !
  !> Calculate the distribution function in all directions
  !! by using the fEq + fNeq
  !!
  !! The input vector must look like
  !! 2D:
  !!```
  !! mom = [ rho, ux, uy, Sxx, Syy, Sxy ]
  !! mom = [ rho, ux, uy, uz, Sxx, Syy, Szz, Sxy, Syz, Sxz ]
  !!```
  !!
  function set_pdfAcoustic( layout, omega, rho0, mom, incompressible )         &
    &                                                             result( pdf )
    ! ---------------------------------------------------------------------------
    type(mus_scheme_layout_type), intent(in) :: layout !scheme layout
    real(kind=rk), intent(in) :: mom(layout%fStencil%QQ) !< input moments
    real(kind=rk), intent(in) :: omega
    real(kind=rk), intent(in) :: rho0
    real(kind=rk)             :: pdf(layout%fStencil%QQ) !< return value
    logical, intent(in), optional :: incompressible
    ! ---------------------------------------------------------------------------
    real(kind=rk) :: fEq(layout%fStencil%QQ)  ! eq distribution
    real(kind=rk) :: fnEq(layout%fStencil%QQ) ! neq distribution
    real(kind=rk) :: moms(layout%fStencil%QQ) ! neq distribution
    logical :: incomp = .false.
    ! ---------------------------------------------------------------------------
    ! Compute the equilibrium
    if ( present(incompressible) ) incomp = incompressible

    moms = mom
    if( layout%fStencil%nDims == 2 ) then
      ! Set for 2d problems the z-velocity to zero
      moms(4) = 0._rk
    end if
    if( incomp ) then
      fEq(:) = getEquilibriumIncomp( dens = mom(1),            &
        &                            vel = moms(2:4),             &
        &                            rho0 = rho0, layout = layout )
    else
      fEq(:) = getEquilibrium( dens = mom(1), vel = moms(2:4), layout = layout )
    end if
    if( layout%fStencil%nDims == 3 ) then
      ! Compute the non-equilibrium
      fNeq = getNeq_acoustic( layout = layout,                   &
        &                     omega = omega,                     &
        &                     Sxx = [                            &
        &                             mom(5), mom(8), mom(10),   &
        &                             mom(8), mom(6), mom(9),    &
        &                             mom(10), mom(9), mom(7) ]  )
    else
      fNeq = getNeq_acoustic( layout = layout,                   &
        &                     omega = omega,                     &
        &                     Sxx = [ mom(4), mom(6), 0._rk,     &
        &                             mom(6), mom(5), 0._rk,     &
        &                             0._rk,  0._rk,  0._rk ]    )
    end if

    pdf(:) = fEq(:) + fNeq(:)

  end function set_pdfAcoustic
! ****************************************************************************** !

! ****************************************************************************** !
  !> Calculate second moments of some quantity \( f \)
  !! \[
  !!    M_{\alpha\beta} = \sum_{i=1}^{Q} c_{i\alpha} c_{i\beta} f_i
  !! \]
  !! where Q is number of discrete velocity.\n
  !! The output is 1 dimentional array which has 6 componenents.\n
  !! Specifically,
  !! \[ m_1 = \sum_{i=1}^{Q} c_{i1} c_{i1} f_i \]
  !! \[ m_2 = \sum_{i=1}^{Q} c_{i2} c_{i2} f_i \]
  !! \[ m_3 = \sum_{i=1}^{Q} c_{i3} c_{i3} f_i \]
  !! \[ m_4 = \sum_{i=1}^{Q} c_{i1} c_{i2} f_i \]
  !! \[ m_5 = \sum_{i=1}^{Q} c_{i2} c_{i3} f_i \]
  !! \[ m_6 = \sum_{i=1}^{Q} c_{i3} c_{i1} f_i \]
  !! This function is used by shear stress and strain rate.
  !!
  pure function secondMom( cxcx, f, QQ ) result ( m )
    ! ---------------------------------------------------------------------------
    integer,       intent(in) :: QQ !< number of discrete directions (=QQ)
    real(kind=rk), intent(in) :: cxcx(6,QQ)   !< discrete velocity of stencil
    !> quantity to which second moment is applied
    real(kind=rk), intent(in) :: f(QQ)
    real(kind=rk)             :: m(6) !< output array
    ! ---------------------------------------------------------------------------

    m(1) = sum( cxcx(1,:) * f(:) )
    m(2) = sum( cxcx(2,:) * f(:) )
    m(3) = sum( cxcx(3,:) * f(:) )
    m(4) = sum( cxcx(4,:) * f(:) )
    m(5) = sum( cxcx(5,:) * f(:) )
    m(6) = sum( cxcx(6,:) * f(:) )

  end function secondMom
! ****************************************************************************** !

! ****************************************************************************** !
  !> Calculate the conversion factor for nonEq in difference levels
  pure function getStrainFacDffs( omegaS, omegaT ) result( fac )
    ! ---------------------------------------------------------------------------
    real(kind=rk), intent(in) :: omegaS  !< omega value on source level
    real(kind=rk), intent(in) :: omegaT  !< omage value on target level
    real(kind=rk) :: fac
    ! ---------------------------------------------------------------------------

??  if( PUSH ) then
      fac = omegaS / omegaT
??  else
      fac =   omegaS * ( 1._rk - omegaT ) &
        &   / ( 1._rk - omegaS ) / omegaT
??  endif

  end function getStrainFacDffs
! ****************************************************************************** !

  ! ************************************************************************** !
  !> This function computes gradient of velocity from gradient and veleocity
  !! data.
  !! Gradient is computed using central difference.
  !! if an element has an boundary then neighbor refers to current element
  !! then forward difference is used
  pure function getGradU( iElem, auxField, gradData, velPos, nAuxScalars, &
    &                     nDims ) result(gradU)
    ! --------------------------------------------------------------------------
    !> element index to compute velocity gradient
    integer, intent(in) :: iElem
    !> auxField
    real(kind=rk), intent(in) :: auxField(:)
    !> gradient data
    type(mus_gradData_type), intent(in) :: gradData
    !> Position of velocity field in auxField
    integer, intent(in) :: velPos(3)
    !> Number of scalars in auxField array
    integer, intent(in) :: nAuxScalars
    !> Dimensions
    integer, intent(in) :: nDims
    !> output: gradient of velocity
    real(kind=rk) :: gradU(nDims, nDims)
    ! --------------------------------------------------------------------------
    integer :: irow, icol
    integer :: leftngh, rightngh
    real(kind=rk) :: leftvel(3), rightvel(3)
    ! --------------------------------------------------------------------------

    gradu = 0.0_rk
    do icol = 1, ndims
      leftngh = (graddata%neighpos(ielem, icol, 1) - 1) * nauxscalars
      rightngh = (graddata%neighpos(ielem, icol, 2) - 1) * nauxscalars
      ! left and right velocity for each direction
      leftvel(1) = auxfield( leftngh+velpos(1) )
      leftvel(2) = auxfield( leftngh+velpos(2) )
      leftvel(3) = auxfield( leftngh+velpos(3) )

      rightvel(1) = auxfield( rightngh+velpos(1) )
      rightvel(2) = auxfield( rightngh+velpos(2) )
      rightvel(3) = auxfield( rightngh+velpos(3) )

      do irow = 1, ndims
        gradu(irow, icol) = (rightvel(irow)-leftvel(irow)) &
          &                 * graddata%fdcoeff(ielem, icol)
      end do !irow
    end do !icol

  end function getGradU
  ! ************************************************************************** !

end module mus_derivedQuantities_module2
! ****************************************************************************** !
