test-stokes.py

from nutils import mesh, function as fn, log, _, plot
import numpy as np

CASE = 'channel'               # 'backstep', or 'channel'
DEGREE = 3
VISCOSITY = 1/20

if CASE == 'backstep':
    # Three-patch domain
    domain, geom = mesh.multipatch(
        patches=[[[0,1],[3,4]], [[3,4],[6,7]], [[2,3],[5,6]]],
        nelems={
            (0,1): 10, (3,4): 10, (6,7): 10,
            (2,5): 100, (3,6): 100, (4,7): 100,
            (0,3): 10, (1,4): 10,
            (2,3): 10, (5,6): 10,
        },
        patchverts=[
            [-1, 0], [-1, 1],
            [0, -1], [0, 0], [0, 1],
            [10, -1], [10, 0], [10, 1]
        ],
    )
    flow_boundaries = [
        'patch0-bottom', 'patch0-top', 'patch0-left',
        'patch1-top', 'patch2-bottom', 'patch2-left',
    ]
    press_boundaries = ['patch1-right', 'patch2-right']

elif CASE == 'channel':
    domain, geom = mesh.rectilinear(
        (np.linspace(0, 10, 101), np.linspace(0, 1, 11))
    )
    flow_boundaries = ['left', 'bottom', 'top']
    press_boundaries = ['right']


# Bases
bases = [
    domain.basis('spline', degree=(DEGREE, DEGREE-1)),
    domain.basis('spline', degree=(DEGREE-1, DEGREE)),
    domain.basis('spline', degree=DEGREE-1),
]
vxbasis, vybasis, pbasis = fn.chain(bases)
vbasis = vxbasis[:,_] * (1,0) + vybasis[:,_] * (0,1)
vgrad = vbasis.grad(geom)

# Matrix
stokes_integrand = (
    VISCOSITY * fn.outer(vgrad).sum((-1, -2))
    - fn.outer(vbasis.div(geom), pbasis)
    - fn.outer(pbasis, vbasis.div(geom))
)
stokes_matrix = domain.integrate(stokes_integrand, geometry=geom, ischeme='gauss9')

# Lift function
x, y = geom
profile = fn.max(0, y*(1-y) * 4)[_] * (1, 0)
lift = domain.project(profile, onto=vbasis, geometry=geom, ischeme='gauss9')
lift[np.where(np.isnan(lift))] = 0.0

# Homogeneous Dirichlet boundary constraints
boundary = domain.boundary[','.join(flow_boundaries)]
constrain = boundary.project((0, 0), onto=vbasis, geometry=geom, ischeme='gauss9')
boundary = domain.boundary[','.join(press_boundaries)]
constrain = boundary.project(0, onto=pbasis, geometry=geom, ischeme='gauss9', constrain=constrain)

# Solution
lhs = stokes_matrix.solve(-stokes_matrix.matvec(lift), constrain=constrain) + lift

# Plotting
vsol, psol = vbasis.dot(lhs), pbasis.dot(lhs)
points, v, vnorm, p = domain.elem_eval(
    [geom, vsol, fn.norm2(vsol), psol],
    ischeme='bezier9', separate=True,
)
with plot.PyPlot('vsol', index=0, figsize=(15,5)) as plt:
    plt.mesh(points, vnorm)
    plt.colorbar()
    plt.streamplot(points, v, spacing=0.2, color='black')
with plot.PyPlot('psol', index=0, figsize=(15,5)) as plt:
    plt.mesh(points, p)
    plt.colorbar()
    plt.clim(-3, 3)