simulation_utils.py

import numpy as np
import json
import os
import scipy
import scipy.io as spio
import scipy.linalg as spla
import scipy.sparse as sps
import scipy.sparse.linalg as spsla
import time

def solve_transient_nse(Re, u1, u2, palpha=1e-3, tE=8, Nts=2**12):
    """
    Solve the transient incompressible Navier-Stokes equations for a given Reynolds number
    and boundary velocities (u1, u2), using a IMEX Euler time-stepping.

    Parameters
    ----------
    Re : float
        Reynolds number
    u1, u2 : float
        Boundary velocities
    palpha : float, optional
        Penalty for Robin boundary condition
    tE : float, optional
        End time of simulation, default 8.0
    Nts : int, optional
        Number of time steps, default 2**12

    """

    # Setup paths for data storage
    data_str = 'data'
    setup_str = os.path.join(data_str, f"Re_{Re}_u1_{u1}_u2_{u2}")
    if not os.path.exists(setup_str):
        os.makedirs(setup_str)

    # Load matrices
    mats = spio.loadmat(data_str + '/mats')
    for k, v in mats.items():
        if sps.issparse(v):
            mats[k] = v.tocsc()

    M = mats['M']
    J = mats['J']
    hmat = mats['H']
    fv = mats['fv'] + 1./Re*mats['fv_diff'] + mats['fv_conv']
    fp = mats['fp'] + mats['fp_div']
    NV, NP = fv.shape[0], fp.shape[0]

    A = 1./Re*mats['A'] + mats['L1'] + mats['L2'] + 1./palpha*mats['Arob']
    B = 1./palpha*mats['Brob']

    # Steady-state NSE solution
    if not os.path.isfile(setup_str + '/ss_nse_sol'):
        ss_nse_v, _ = solve_steadystate_nse(mats, Re, palpha=palpha)
        scipy.io.savemat(setup_str + '/ss_nse_sol', {'ss_nse_v': ss_nse_v})
    else:
        ss_nse_sol = scipy.io.loadmat(setup_str + '/ss_nse_sol')
        ss_nse_v = ss_nse_sol['ss_nse_v']

    # Files for results and visualization
    if not os.path.exists(setup_str + '/timesteps'):
        os.makedirs(setup_str + '/timesteps')

    vfile = lambda t, u1, u2: f"{setup_str}/timesteps/v_t{t}.vtu"
    pfile = lambda t, u1, u2: f"{setup_str}/timesteps/p_t{t}.vtu"
    vfilerel = lambda t, u1, u2: f"timesteps/v_t{t}.vtu"
    pfilerel = lambda t, u1, u2: f"timesteps/p_t{t}.vtu"
    strtojson = data_str + '/visualization.jsn'

    # Precompute time-stepping variables
    t0 = 0.
    DT = (tE - t0) / Nts
    trange = np.linspace(t0, tE, Nts + 1)
    vfilelist = []
    pfilelist = []
    Bu = B @ np.array([[u1], [u2]])
    old_v = ss_nse_v

    sysmat = sps.vstack([
        sps.hstack([M + DT*A, -J.T]),
        sps.hstack([J, sps.csc_matrix((NP, NP))])
    ]).tocsc()
    sysmati = spsla.factorized(sysmat)

    # Time-stepping loop
    tic = time.time()
    for k, t in enumerate(trange):
        crhsv = M*old_v + DT*(fv - eva_quadterm(hmat, old_v) + Bu)
        crhs = np.vstack([crhsv, fp])
        vp_new = np.atleast_2d(sysmati(crhs.flatten())).T
        old_v = vp_new[:NV]
        p = vp_new[NV:]

        if np.mod(k, round(Nts / 128)) == 0:
            toc = time.time()
            v_mag = np.linalg.norm(old_v)
            if v_mag > 1000 or np.isnan(v_mag):
                print('Velocity magnitude diverging.')
                break
            print('u1 {0:f} u2 {1:f} timestep {2:4d}/{3}, t={4:f}, |v|={5:e}, elapsed {6:f}'
                  .format(u1, u2, k, Nts, t, v_mag, toc-tic))
            writevp_paraview(velvec=old_v, pvec=p,
                             vfile=vfile(t, u1, u2),
                             pfile=pfile(t, u1, u2),
                             strtojson=strtojson)
            vfilelist.append(vfilerel(t, u1, u2))
            pfilelist.append(pfilerel(t, u1, u2))
            tic = time.time()

    # Create .pvd files from time-stepping results
    collect_vtu_files(vfilelist, f"{setup_str}/v_result.pvd")
    collect_vtu_files(pfilelist, f"{setup_str}/p_result.pvd")


def solve_steadystate_nse(mats, Re, control='bc', tol=1e-10, npicardstps=5, maxit=30,
                          palpha=1, uvec=np.array([[0], [0]]), v0=None):
    """Compute steady-state NSE solution."""
    J = mats['J']
    hmat = mats['H']
    fp = mats['fp'] + mats['fp_div']

    if control == 'bc':
        A = 1./Re*mats['A'] + mats['L1'] + mats['L2'] + 1./palpha*mats['Arob']
        Brob = mats['Brob']
        fv = mats['fv'] + 1./Re*mats['fv_diff'] + mats['fv_conv'] \
                        + 1. / palpha*np.dot(Brob, uvec)
    else:
        A = 1. / Re * mats['A'] + mats['L1'] + mats['L2']
        fv = mats['fv'] + 1./Re*mats['fv_diff'] + mats['fv_conv']

    updnorm = 1
    NV, NP = fv.shape[0], fp.shape[0]
    if v0 is None:
        curv = np.zeros((NV, 1))
    else:
        curv = v0
    stpcount = 0
    while updnorm > tol:
        picard = stpcount < npicardstps
        H1k, H2k = linearized_convection(hmat, curv, retparts=True)
        if picard:
            currhs = np.vstack([fv, fp])
            HL = H1k
        else:
            currhs = np.vstack([fv+eva_quadterm(hmat, curv), fp])
            HL = H1k + H2k
        cursysmat = sps.vstack([
                        sps.hstack([A+HL, -J.T]),
                        sps.hstack([J, sps.csc_matrix((NP, NP))])
                    ]).tocsc()
        nextvp = spsla.spsolve(cursysmat, currhs).reshape((NV+NP, 1))

        print('Iteration step {0} ({1})'.format(stpcount, 'Picard' if picard else 'Newton'))
        nextv = nextvp[:NV].reshape((NV, 1))
        nextp = nextvp[NV:].reshape((NP, 1))
        curnseres = A*nextv + eva_quadterm(hmat, nextv) - J.T*nextp - fv
        print('Norm of nse residual:   {0:e}'.format(np.linalg.norm(curnseres)))
        updnorm = np.linalg.norm(nextv - curv) / np.linalg.norm(nextv)
        print('Norm of current update: {0:e}'.format(updnorm))
        print('\n')
        curv = nextv
        stpcount += 1
        if stpcount > maxit:
            raise RuntimeError('Could not compute steady-state NSE solution.')

    return nextv, nextp


def linearized_convection(H, linv, retparts=False):
    """Compute linearized convection term."""
    nv = linv.size
    if retparts:
        H1 = H * (sps.kron(sps.eye(nv), linv))
        H2 = H * (sps.kron(linv, sps.eye(nv)))
        return H1, H2
    else:
        H1 = H * (sps.kron(sps.eye(nv), linv))
        H2 = H * (sps.kron(linv, sps.eye(nv)))
        return H1 + H2

def collect_vtu_files(filelist, pvdfilestr):
    with open(pvdfilestr, 'w') as colfile:
        colfile.write(u'<?xml version="1.0"?>\n<VTKFile type="Collection" version="0.1"> <Collection>\n')
        for tsp, vtufile in enumerate(filelist):
            dtst = u'<DataSet timestep="{0}" part="0" file="{1}"/>'.format(tsp, vtufile)
            colfile.write(dtst)

        colfile.write(u'</Collection> </VTKFile>')

def eva_quadterm(H, v):
    ''' function to evaluate `H*kron(v, v)` without forming `kron(v, v)`

    Parameters:
    ---
    H : (nv, nv*nv) sparse array
        the tensor (as a matrix) that evaluates the convection term

    '''
    nv = v.size
    hvv = np.zeros((nv,1))

    for k in range(nv):
        hvv += H[:, k*nv:(k+1)*nv] @ (v[k] * v)
    
    return hvv

def writevp_paraview(velvec=None, pvec=None, strtojson=None, visudict=None, vfile='vel.vtu', pfile='p.vtu', include_bc=True):
    if visudict is None:
        jsfile = open(strtojson)
        visudict = json.load(jsfile)
        vaux = np.zeros((visudict['vdim'], 1))
        if include_bc:
            for bcdict in visudict['bclist']:
                intbcidx = [int(bci) for bci in bcdict.keys()]
                vaux[intbcidx, 0] = list(bcdict.values())
        vaux[visudict['invinds']] = velvec

    vxvtxdofs = visudict['vxvtxdofs']
    vyvtxdofs = visudict['vyvtxdofs']

    with open(vfile, 'w') as velfile:
        velfile.write(visudict['vtuheader_v'])
        for xvtx, yvtx in zip(vxvtxdofs, vyvtxdofs):
            velfile.write(u'{0} {1} {2} '.format(vaux[xvtx][0], vaux[yvtx][0], 0.))
        velfile.write(visudict['vtufooter_v'])

    if pvec is not None:
        pvtxdofs = visudict['pvtxdofs']
        with open(pfile, 'w') as prefile:
            prefile.write(visudict['vtuheader_p'])
            for pval in pvec[pvtxdofs, 0]:
                prefile.write(u'{0} '.format(pval))
            prefile.write(visudict['vtufooter_p'])