tris_models.py

from periodic2d import *
import subprocess
import glob
import scipy.special as spec
from PIL import Image
import glob

plt.rcParams.update({'lines.linewidth': 1, 'lines.markersize': 6,'lines.markeredgewidth': 0.2})
plt.rcParams.update({'font.size': 12, 'legend.fontsize': 8})
plt.rc('font', family='serif')
plt.rc('text',usetex=True)

# sine integral function
def Si(x):
  return spec.sici(x)[0]

def Ci(x):
  return spec.sici(x)[1]

# input parameters
class Params:
  """Data structure for storing and passing input parameters for TRIS closures."""

  def __init__(self,kappa=0.41,B=5.2,Re=300.,Pi_ref=0.15,Cuv=0.6,Cs0=0.04,Cs1=0.04,Cp=6.0,Cn=0.0,n=0.0,Ctau=1.0,CPi=1.0,linearize=True,flow_angle=0.):
    self.kappa = kappa
    self.B   = B
    self.Re  = Re
    self.Pi_ref  = Pi_ref
    self.Cuv = Cuv
    self.Cs0 = Cs0
    self.Cs1 = Cs1
    self.Cp  = Cp
    self.Cn  = Cn
    self.n   = n
    self.Ctau = Ctau
    self.CPi = CPi
    self.linearize = linearize
    self.flow_angle = flow_angle

class Animation:
    """Data structure for creating an animation of a given quantity."""

    def __init__(self, vname, vdir='animation', fmt='png', dpi=300):
        self.name = vname
        self.dir = vdir
        self.fmt = fmt
        self.dpi = dpi
        self.filename = self.dir+'vid_'+self.name
        self.N = 0

    def update(self, i, field):
        Lx = field.msh.Lx
        Lz = field.msh.Lz
        colormap = plt.get_cmap('RdBu_r')
        plt.figure(1, figsize=(6,1.5))
        im = plt.imshow((field.val.T - np.mean(field.val.T))/np.std(field.val.T), interpolation='bicubic', origin='lower',\
                        extent=(0,Lx,0,Lz), cmap=colormap, vmin=-3, vmax=3)
        cb = plt.colorbar(im)
        cb.set_ticks(np.linspace(-3, 3, num=3))  # Adjust number of ticks
        plt.xticks(np.linspace(0, Lx, 5), [r"$"+pi_formatter(val, None)+r"$" for val in np.linspace(0, Lx, 5)],fontsize=12)
        plt.yticks(np.linspace(0, Lz, 4), [r"$"+pi_formatter(val, None)+r"$" for val in np.linspace(0, Lz, 4)],fontsize=12)
        if self.name == 'u0':
          plt.title(r'$\left< \widetilde{u} \right>_0^s$')
        if self.name == 'u1':
          plt.title(r'$\left< \widetilde{u} \right>_1^s$')
        if self.name == 'w0':
          plt.title(r'$\left< \widetilde{w} \right>_0^s$')
        if self.name == 'w1':
          plt.title(r'$\left< \widetilde{w} \right>_1^s$')
        if self.name == 'v0':
          plt.title(r'$\left< \widetilde{v} \right>_0^s$')
        if self.name == 'p0':
          plt.title(r'$\left< \widetilde{p} \right>_0^s$')
        if self.name == 'p0':
          plt.title(r'$\left< \widetilde{p} \right>_1^s$')
        plt.savefig(self.filename+'_'+str(i)+'.'+self.fmt, format=self.fmt, bbox_inches='tight', dpi=self.dpi)
        plt.close(1)
        self.N += 1

    def compile(self,animation_variable):
        image_files = sorted(glob.glob('animation/*.png'))
        images = [Image.open(image) for image in image_files]

        # Save as GIF
        gif_path = 'animation/'+animation_variable+'.gif'
        images[0].save(
            gif_path,
            save_all=True,
            append_images=images[1:],
            duration=100,
            loop=0)

        cmd = ['rm']
        files = glob.glob(self.filename+'_*.'+self.fmt)
        subprocess.run(cmd+files)

def calc_v0(u1, w1):
  return div(u1, w1).smult(0.5)

def calc_rhs(u0, w0, u1, w1, params, verbose=False):
  ## 0. calculate v0 field
  v0 = calc_v0(u1, w1)

  ## obtain u,v,w fields in physical space
  for field in [u0, w0, u1, w1, v0]:
    field.ifft()

  ## 1. calculate assumed profile parameters from moments
  (Re, thRe, Pi, thPi) = profile_params(u0.val, w0.val, u1.val, w1.val, params)

  ## 2. closure models for nonlinear flux terms
  (r0xx,r0zz,r0xz,r1xx,r1zz,r1xz) = structural_model(Re, thRe, Pi, thPi, params)
  # place it back on the mesh
  uu0 = Var(u0.msh, val=r0xx)
  uw0 = Var(u0.msh, val=r0xz)
  ww0 = Var(w0.msh, val=r0zz)
  uu1 = Var(u1.msh, val=r1xx)
  uw1 = Var(u1.msh, val=r1xz)
  ww1 = Var(w1.msh, val=r1zz)
  # 
  for field in [uu0, uw0, ww0, uu1, uw1, ww1]:
    field.fft()
  # negative divergence on RHS
  rhs_u0 = div(uu0, uw0).smult(-1.)
  rhs_w0 = div(uw0, ww0).smult(-1.)
  rhs_p0 = div(rhs_u0, rhs_w0)
  rhs_u1 = div(uu1, uw1).smult(-1.)
  rhs_w1 = div(uw1, ww1).smult(-1.)
  rhs_p1 = div(rhs_u1, rhs_w1)

  ## 3. calculate resolved wall-normal momentum fluxes
  uv0 = u0.mult(v0)
  wv0 = w0.mult(v0)
  for field in [uv0, wv0]:
    field.fft()
  rhs_u1 = rhs_u1.add(uv0.smult(2.))
  rhs_w1 = rhs_w1.add(wv0.smult(2.))
  rhs_p1 = rhs_p1.add(div(uv0,wv0).smult(4.))

  ## 4. calculate wall shear stress model
  (taux, tauz, ztaux, ztauz) = stress_model(Re, thRe, Pi, thPi, params)
  # put it back on the mesh
  tx = Var(u0.msh, val=taux)
  tz = Var(w0.msh, val=tauz)
  ztx = Var(u0.msh, val=ztaux)
  ztz = Var(w0.msh, val=ztauz)

  # return to wavenumber space
  for field in [tx, tz, ztx, ztz]:
    field.fft()
  # subtract from RHS
  rhs_u0 = rhs_u0.add(tx.smult(-1.))
  rhs_w0 = rhs_w0.add(tz.smult(-1.))
  rhs_p0 = rhs_p0.add(div(tx, tz).smult(-1.))
  # integral of the viscous and unresolved turbulent stress 
  rhs_u1 = rhs_u1.add(ztx.smult(-1.))
  rhs_w1 = rhs_w1.add(ztz.smult(-1.))
  rhs_p1 = rhs_p1.add(div(ztx, ztz).smult(-2.))

  ## 4b. calculate viscous stress integral
  (utopx, utopz) = top_velocity_model(Re, thRe, Pi, thPi, params)
  visc_x = Var(u0.msh, val=-2.*utopx/params.Re)
  visc_z = Var(u0.msh, val=-2.*utopz/params.Re)
  for field in [visc_x, visc_z]:
    field.fft()
  rhs_u1 = rhs_u1.add(visc_x)
  rhs_w1 = rhs_w1.add(visc_z)

  ## 5. add body force driving the flow
  rhs_u0 = rhs_u0.sadd(np.cos(params.flow_angle))
  rhs_u1 = rhs_u1.sadd(np.cos(params.flow_angle))
  rhs_w0 = rhs_w0.sadd(np.sin(params.flow_angle))
  rhs_w1 = rhs_w1.sadd(np.sin(params.flow_angle))

  ## 6. wall-parallel eddy viscosity
  (S0xx, S0zz, S0xz, S02) = strain_rate(u0, w0)
  (S1xx, S1zz, S1xz, S12) = strain_rate(u1, w1)
  if params.linearize:
    Lx = u0.msh.Lx
    Lz = u0.msh.Lz
    nx = u0.msh.nx
    nz = u0.msh.nz
    dx = Lx/nx
    dz = Lz/nz
    Smag = S02.sqrt()
    nuT0 = Smag.smult(params.Cs0).smult(dx).smult(dz)
    Smag = S12.sqrt()
    nuT1 = Smag.smult(params.Cs1).smult(dx).smult(dz)
  else:
    nuT0 = eddy_viscosity_model(S02, S12, params)
    nuT1 = nuT0
  tau0xx = S0xx.mult(nuT0)
  tau0zz = S0zz.mult(nuT0)
  tau0xz = S0xz.mult(nuT0)
  tau1xx = S1xx.mult(nuT1)
  tau1zz = S1zz.mult(nuT1)
  tau1xz = S1xz.mult(nuT1)
  f0x = div(tau0xx, tau0xz)
  f0z = div(tau0xz, tau0zz)
  f1x = div(tau1xx, tau1xz)
  f1z = div(tau1xz, tau1zz)
  rhs_u0 = rhs_u0.add(f0x)
  rhs_w0 = rhs_w0.add(f0z)
  rhs_u1 = rhs_u1.add(f1x)
  rhs_w1 = rhs_w1.add(f1z)
  rhs_p0 = rhs_p0.add(div(f0x, f0z))
  rhs_p1 = rhs_p1.add(div(f1x, f1z))

  nu_eff = (1./params.Re) # molecular viscosity
  rhs_u0 = rhs_u0.add(u0.laplacian().smult(nu_eff))  
  rhs_w0 = rhs_w0.add(w0.laplacian().smult(nu_eff))
  rhs_u1 = rhs_u1.add(u1.laplacian().smult(nu_eff))
  rhs_w1 = rhs_w1.add(w1.laplacian().smult(nu_eff))

  # project RHS to make divergence-free
  p0 = calc_p0(rhs_p0)
  rhs_u0 = rhs_u0.add(p0.diff_x().smult(-1.))
  rhs_w0 = rhs_w0.add(p0.diff_z().smult(-1.))
  p1 = new_calc_p1(rhs_p1, p0, params)
  rhs_u1 = rhs_u1.add(p1.diff_x().smult(-1.))
  rhs_w1 = rhs_w1.add(p1.diff_z().smult(-1.))

  return (rhs_u0, rhs_w0, rhs_u1, rhs_w1, p0, p1)

def profile_params(U0, W0, U1, W1, params, verbose=False):
  # unpack model parameters
  kappa = params.kappa
  B = params.B
  #
  # first vector for obtaining Re_* and theta_*
  q1 = (1 + np.pi**2/4)*U0 - np.pi**2/4*U1
  q2 = (1 + np.pi**2/4)*W0 - np.pi**2/4*W1
  qmag = np.sqrt(q1**2 + q2**2)
  Re = np.exp(kappa*qmag - kappa*B + np.pi**2/8 + 1)  
  thRe = np.sign(q2)*np.arccos(q1/qmag)
  #
  # second vector for obtaining Pi and theta_Pi
  q1 = np.pi**2/4*kappa*(U1 - U0) - np.pi**2/8*np.cos(thRe)
  q2 = np.pi**2/4*kappa*(W1 - W0) - np.pi**2/8*np.sin(thRe)
  Pi = np.sqrt(q1**2 + q2**2)
  thPi = np.sign(q2)*np.arccos(q1/Pi)
  #
  return (Re, thRe, Pi, thPi)

# solve pressure Poisson for <p>_0
def calc_p0(rhs):
  msh = rhs.msh 
  kx = msh.KX
  kz = msh.KZ
  k2 = kx*kx + kz*kz + 1e-99
  p0_hat = -rhs.hat/k2
  p0_hat[k2==0.] = 0. # take the mean pressure out
  return Var(msh, hat=p0_hat)

# solve pressure Helmholtz for <p>_1
def old_calc_p1(rhs, p0, params):
  # unpack model parameters
  Cp = params.Cp
  # setup Helmholtz equation
  rhs = rhs.add(p0.smult(-2.*Cp))
  msh = rhs.msh 
  kx = msh.KX
  kz = msh.KZ
  k2 = kx*kx + kz*kz# + 1e-99
  p1_hat = -rhs.hat/(k2+2*Cp)
  return Var(msh, hat=p1_hat)

def calc_p1(rhs):
  msh = rhs.msh
  kx = msh.KX
  kz = msh.KZ
  k2 = kx*kx + kz*kz + 1e-99
  p1_hat = -rhs.hat/k2
  p1_hat[k2==0.] = 0.
  return Var(msh, hat=p1_hat)

# solve pressure Helmholtz for <p>_1 + large scale relax to zero divergence
def new_calc_p1(rhs, p0, params):
  # unpack model parameters
  Cp  = params.Cp
  # setup Helmholtz equation
  rhs = rhs.add(p0.smult(-2.*Cp))
  msh = rhs.msh 
  kx = msh.KX
  kz = msh.KZ
  k2 = kx*kx + kz*kz# + 1e-99
  return Var(msh, hat=-rhs.hat/(k2+2*Cp))

# return nonlinear for zeroth and first moments
def structural_model(Re, thRe, Pi, thPi, params, verbose=False):
  # unpack model parameters
  kappa = params.kappa
  B = params.B
  Re_ref = params.Re
  Pi_ref = params.Pi_ref
  lin = params.linearize
  #
  ff = np.log(Re)/kappa + B
  ## zeroth moment closure
  ff_ref = np.log(Re_ref)/kappa + B
  r0RR = (2*ff_ref - 2/kappa)*ff + (2/kappa**2 - ff_ref**2) 
  r0RP = Pi/kappa*ff_ref + Pi_ref/kappa*(ff-ff_ref) - Pi/(kappa**2)*(1-Si(np.pi)/np.pi)
  r0PP = 3*Pi_ref / (kappa**2) * (Pi - 0.5*Pi_ref)
  # 
  # i=1, j=1
  r0xx  = r0RR*np.cos(thRe)**2 
  r0xx += r0RP*(2*np.cos(thRe)*np.cos(thPi)) 
  r0xx += r0PP*np.cos(thPi)**2
  # i=2, j=2
  r0zz  = r0RR*np.sin(thRe)**2 
  r0zz += r0RP*(2*np.sin(thRe)*np.sin(thPi)) 
  r0zz += r0PP*np.sin(thPi)**2
  # i=1, j=2 OR i=2, j=1 (symmetric)
  r0xz  = r0RR*np.cos(thRe)*np.sin(thRe) 
  r0xz += r0RP*(np.cos(thRe)*np.sin(thPi) + np.sin(thRe)*np.cos(thPi))
  r0xz += r0PP*np.cos(thPi)*np.sin(thPi)
  #
  # first moment closure
  r1RR = (2*ff_ref - 1/kappa)*ff + (0.5/(kappa**2) - ff_ref**2)
  r1RP = (1 + 4/np.pi**2)*((Pi/kappa)*ff_ref + Pi_ref/kappa*(ff - ff_ref))
  r1RP -= 2*Pi/kappa**2*(0.25-2/np.pi**2-(Ci(np.pi)-np.euler_gamma-np.log(np.pi))/np.pi**2)
  r1PP = (3 + 16/np.pi)*Pi_ref/(kappa**2) * (Pi - 0.5*Pi_ref)
  # 
  # i=1, j=1
  r1xx  = r1RR*np.cos(thRe)**2 
  r1xx += r1RP*(2*np.cos(thRe)*np.cos(thPi)) 
  r1xx += r1PP*np.cos(thPi)**2
  # i=2, j=2
  r1zz  = r1RR*np.sin(thRe)**2 
  r1zz += r1RP*(2*np.sin(thRe)*np.sin(thPi)) 
  r1zz += r1PP*np.sin(thPi)**2
  # i=1, j=2 OR i=2, j=1 (symmetric)
  r1xz  = r1RR*np.cos(thRe)*np.sin(thRe) 
  r1xz += r1RP*(np.cos(thRe)*np.sin(thPi) + np.sin(thRe)*np.cos(thPi))
  r1xz += r1PP*np.cos(thPi)*np.sin(thPi)
  #
  return (r0xx, r0zz, r0xz, r1xx, r1zz, r1xz)

def stress_model(Re, thRe, Pi, thPi, params, verbose=False):
  # unpack model parameters
  kappa = params.kappa
  B = params.B
  lin = params.linearize
  Re_ref = params.Re
  Pi_ref = params.Pi_ref
  Cuv = params.Cuv
  Ctau = params.Ctau
  CPi  = params.CPi
  Cn = params.Cn
  n = params.n
  # calculate wall shear stress in local reference frame
  tmag = (Re/Re_ref)**(2*Ctau)
  tX = tmag * np.cos(thRe)
  tZ = tmag * np.sin(thRe)

  # viscous + unresolved turbulent stress integral for moment of momentum
  ttR = 1. + (1-CPi)*Pi_ref
  ttP = 1.*CPi*Pi
  ttref = 1. + Pi_ref
  ttX = Cuv * (ttR*np.cos(thRe) + ttP*np.cos(thPi))/ttref
  ttZ = Cuv * (ttR*np.sin(thRe) + ttP*np.sin(thPi))/ttref

  return (tX, tZ, ttX, ttZ)

def top_velocity_model(Re, thRe, Pi, thPi, params):
  kappa = params.kappa
  B = params.B
  Re_ref = params.Re
  Pi_ref = params.Pi_ref
  utR = np.log(Re)/kappa+B 
  utP = 2*Pi/kappa
  utx = utR*np.cos(thRe) + utP*np.cos(thPi)
  utz = utR*np.sin(thRe) + utP*np.sin(thPi)
  return (utx, utz)
 
def ref_zeroth_moment(Re, thRe, Pi, thPi, params, verbose=False):
  # unpack model parameters
  kappa = params.kappa
  B = params.B
  Re_ref = params.Re
  Pi_ref = params.Pi_ref
  R0_ref = (np.log(Re_ref)-1.)/kappa + B
  P0_ref = Pi_ref / kappa 
  u0_ref = R0_ref * np.cos(thRe) + P0_ref * np.cos(thPi)
  w0_ref = R0_ref * np.sin(thRe) + P0_ref * np.sin(thPi)
  return (u0_ref, w0_ref)

def strain_rate(u, w):
  dudx = u.diff_x()
  dudz = u.diff_z()
  dwdx = w.diff_x()
  dwdz = w.diff_z()
  Sxx = dudx
  Sxz = dudz.add(dwdx)
  Sxz = Sxz.smult(0.5)
  Szz = dwdz
  Smag2 = Sxx.mult(Sxx) 
  Smag2 = Smag2.add(Szz.mult(Szz))
  Smag2 = Smag2.add(Sxz.mult(Sxz.smult(2.)))
  return (Sxx, Szz, Sxz, Smag2)

def eddy_viscosity_model(S02, S12, params):
  # unpack model parameters
  Cs = params.Cs0
  Smag = S02.add(S12)
  Smag = Smag.sqrt()
  return Smag.smult(Cs)

# Pi formatter function
def pi_formatter(x, pos):
  N = int(np.round(x / np.pi))
  if N == 0:
    return "0"
  elif N == 1:
    return r"\pi"
  elif N == -1:
    return r"-\pi"
  elif N % 2 == 0:
    return r"{0}\pi".format(N)
  else:
    return r"{0}\pi".format(N)