stokes1.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. """This demo solves the Stokes equations, using quadratic elements for
  2. the velocity and first degree elements for the pressure (Taylor-Hood elements)"""
  3. # Credits:
  4. # https://github.com/pf4d/fenics_scripts/blob/master/cbc_block/stokes.py
  5. from dolfin import *
  6. import numpy as np
  7. from vedo.dolfin import plot
  8. from vedo import Latex, dataurl, download
  9. # Load mesh and subdomains
  10. fpath = download(dataurl + "dolfin_fine.xml")
  11. mesh = Mesh(fpath)
  12. fpath = download(dataurl + "dolfin_fine_subdomains.xml.gz")
  13. sub_domains = MeshFunction("size_t", mesh, fpath)
  14. # Define function spaces
  15. P2 = VectorElement("Lagrange", mesh.ufl_cell(), 2)
  16. P1 = FiniteElement("Lagrange", mesh.ufl_cell(), 1)
  17. TH = P2 * P1
  18. W = FunctionSpace(mesh, TH)
  19. # No-slip boundary condition for velocity
  20. noslip = Constant((0, 0))
  21. bc0 = DirichletBC(W.sub(0), noslip, sub_domains, 0)
  22. # Inflow boundary condition for velocity
  23. inflow = Expression(("-sin(x[1]*pi)", "0.0"), degree=2)
  24. bc1 = DirichletBC(W.sub(0), inflow, sub_domains, 1)
  25. bcs = [bc0, bc1]
  26. # Define variational problem
  27. (u, p) = TrialFunctions(W)
  28. (v, q) = TestFunctions(W)
  29. f = Constant((0, 0))
  30. a = (inner(grad(u), grad(v)) - div(v) * p + q * div(u)) * dx
  31. L = inner(f, v) * dx
  32. w = Function(W)
  33. solve(a == L, w, bcs)
  34. # Split the mixed solution using a shallow copy
  35. (u, p) = w.split()
  36. ##################################################################### vedo
  37. f = r"-\nabla \cdot(\nabla u+p I)=f ~\mathrm{in}~\Omega"
  38. formula = Latex(f, pos=(0.55, 0.45, -0.05), s=0.1)
  39. plot(
  40. u,
  41. N=2,
  42. mode="mesh and arrows",
  43. scale=0.03,
  44. wireframe=True,
  45. scalarbar=False,
  46. style=1,
  47. ).close()
  48. plot(p, text="pressure", cmap="rainbow").close()