polar = de.PolarCoordinates("phi", "r")
axial = de.Coordinate("z")
dist = de.Distributor((axial, polar), dtype=np.float64)
disk = de.DiskBasis(polar, (12, 20), dtype=np.float64, radius=1.0)
axis = de.RealFourier(axial, 10, bounds=(0, 2))
phi, r = disk.local_grids()
z = axis.local_grids()[0]
u = dist.Field(bases=(axis, disk), name="u")
v = dist.Field(bases=(axis, disk), name="v")
w = dist.Field(bases=(axis, disk), name="w")
tau_u = dist.Field(bases=(axis, disk.edge), name="tau_u")
tau_v = dist.Field(bases=(axis, disk.edge), name="tau_u")
tau_w = dist.Field(bases=(axis, disk.edge), name="tau_w")
# Substitutions
dz = lambda a: de.Differentiate(a, axial)
lap = lambda a: de.Laplacian(a, polar)
lift = lambda a: de.Lift(a, disk, -1)
# Problem
problem = de.IVP([u, v, w, tau_u, tau_v, tau_w], namespace=locals())
problem.add_equation("dt(u) - lap(u) - dz(dz(u)) + lift(tau_u) = 0")
problem.add_equation("dt(v) - lap(v) - dz(dz(v)) + lift(tau_v) = 0")
problem.add_equation("dt(w) - lap(w) - dz(dz(w)) + lift(tau_w) = -1")
problem.add_equation("u(r=1) = 0")
problem.add_equation("v(r=1) = 0")
problem.add_equation("w(r=1) = 0")
solver = problem.build_solver(de.SBDF2)
solver.stop_sim_time = 5
flow = de.GlobalFlowProperty(solver, cadence=100)
flow.add_property(u * u, name="u2")
while solver.proceed:
solver.step(1e-2)
if (solver.iteration - 1) % 100 == 0:
log_prop = np.sqrt(flow.max("u2"))
print(
f"Iteration={solver.iteration}, Time={solver.sim_time:e}, "
f"flow max={log_prop}"
)
solver.log_stats()