Application: Cavity Flow
One of the most common validation cases in CFD is the lid-driven cavity flow. We take a square cavity filled with a fluid and set the velocity of the lid to some constant value. The flow within the cavity is driven by the lid, a spiral flow pattern develops and two distinctive pressure zones are visible in the upper corners against the lid.
import numpy
The Poisson equation is an elliptic PDE which almost always means using an iterative solver. We're going to use the Jacobi method. There are better ways, but that's beside the point.
Here's the pressure Poisson equation:
That looks a little nasty, but we only care about the top line when we iterate, since the bottom three lines depend only on values that don't change when we're correcting the pressure field. Because it doesn't change, we break it out into a separate function.
def velocity_term(b, rho, dt, u, v, dx):
b[1:-1, 1:-1] = (
rho * dx / 16 *
(2 / dt * (u[1:-1, 2:] -
u[1:-1, :-2] +
v[2:, 1:-1] -
v[:-2, 1:-1]) -
2 / dx * (u[2:, 1:-1] - u[:-2, 1:-1]) *
(v[1:-1, 2:] - v[1:-1, :-2]) -
(u[1:-1, 2:] - u[1:-1, :-2])**2 / dx -
(v[2:, 1:-1] - v[:-2, 1:-1])**2 / dx)
)
return b
Now, to calculate the pressure field, we pass in the original pressure field, the value b (which is the result of the velocity_term function above) and a target value for difference between two iterates. We repeatedly update the pressure field until the difference of the L2 norm between two successive iterations is less than that target value.
def pressure_poisson(p, b, l2_target):
iter_diff = l2_target + 1
n = 0
while iter_diff > l2_target and n <= 500:
pn = p.copy()
p[1:-1,1:-1] = (.25 * (pn[1:-1, 2:] +
pn[1:-1, :-2] +
pn[2:, 1:-1] +
pn[:-2, 1:-1]) -
b[1:-1, 1:-1])
p[:, 0] = p[:, 1] #dp/dx = 0 at x = 0
p[:, -1] = p[:, -2] #dp/dx = 0 at x = 2
p[0, :] = p[1, :] #dp/dy = 0 at y = 0
p[-1, :] = 0 #p = 0 at y = 2
if n % 10 == 0:
iter_diff = numpy.sqrt(numpy.sum((p - pn)**2)/numpy.sum(pn**2))
n += 1
return p
In the interests of brevity, we're only going to worry about the pressure Poisson solver. The rest of the 2D Navier-Stokes solution is encapsulated in the function cavity_flow, which we've prepared ahead of time and saved in a helper file. If you want to dig deeper into how the cavity_flow function works, check out "The 12 Steps to Navier-Stokes" at CFD Python.
For now, though, we just need to import the function:
from snippets.ns_helper import cavity_flow
We'll also load up pickled initial conditions, so we can reliably compare final solutions.
import pickle
def run_cavity():
nx = 41
ny = 41
with open('IC.pickle', 'rb') as f:
u, v, p, b = pickle.load(f)
dx = 2 / (nx - 1)
dt = .005
nt = 1000
u, v, p = cavity_flow(u, v, p, nt, dt, dx,
velocity_term,
pressure_poisson,
rtol=1e-4)
return u, v, p
So what does this all do? Let's check it out.
u, v, p = run_cavity()
%matplotlib inline
from snippets.ns_helper import quiver_plot
quiver_plot(u, v, p)
Save NumPy answers for comparison
This will create a binary file, saved to disk as numpy_ans.pickle that contains the values of u, v, and p.
with open('numpy_ans.pickle', 'wb') as f:
pickle.dump((u, v, p), f)
Let's profile the cavity_flow function and see if there's a specific place that's really hurting our performance.
%timeit run_cavity()
1 loop, best of 3: 462 ms per loop
Remember that we need to load the line_profiler extension before we can use the lprun magic.
%load_ext line_profiler
%lprun -f cavity_flow run_cavity()
Where is the bottleneck?
Clearly the PPE is the problem here, so let's use numba to rewrite it.
Exercise: Speed up the PPE
from numba import jit
%load snippets/ppe_numba.py
Since we have redefined the pressure_poisson function, we can run run_cavity again and we'll be using the new and improved jitted PPE.
We don't want to overwrite our results, so we can choose different variable names to store the final velocity and pressure fields.
u_numba, v_numba, p_numba = run_cavity()
We use numpy.allclose to check that each value of the Python and Numba fields match to within a specified tolerance (the default tolerance is 108).
assert numpy.allclose(p, p_numba)
assert numpy.allclose(u, u_numba)
assert numpy.allclose(v, v_numba)
If there were no errors raised by the previous cell, then the two answers match. Hooray! If the answers match, we should see the same plot as above. If they don't match, we can also check the plot to see if we're just a little bit off or completely wrong (although that can be hard to judge sometimes).
quiver_plot(u_numba, v_numba, p_numba)
And now we can check to see how much the performance has been improved:
%timeit run_cavity()
1 loop, best of 3: 460 ms per loop
Not as spectacular as some of the other examples, but remember, we started off using NumPy arrays, not regular nested Python loops, so this is still a very respectable speedup.
Now if we re-profile the code, we can see how the runtime percentages have shifted around the improved PPE. It's almost certainly going to remain the most expensive part of the solve, but it should be a smaller percentage overall.
%lprun -f cavity_flow run_cavity()
One more bit of optimization?
As we improve the performance of the Pressure Poisson function, the profiling results will change and indicate other possible hotspots that we want to improve. It can be hard to know when to stop improving, to recognize that the gains are not worth what you are investing in rewriting code.
In the profile above (after rewriting the PPE) you'll notice that the velocity_term now occupies a larger percentage of the total run time. Here's a rewritten version of velocity_term that makes use of Numba so you can compare how much more performance we might be able to eke out of this simple cavity flow solver.
@jit(nopython=True)
def velocity_term(b, rho, dt, u, v, dx):
J, I = b.shape
for i in range(1, I):
for j in range(1, J):
b[j, i] = (
rho * dx / 16 *
(2 / dt * (u[j, i + 1] -
u[j, i - 1] +
v[j + 1, i] -
v[j - 1, i]) -
2 / dx * (u[j + 1, i] - u[j - 1, i]) *
(v[j, i + 1] - v[j, i - 1]) -
(u[j, i + 1] - u[j, i - 1])**2 / dx -
(v[j + 1, i] - v[j - 1, i])**2 / dx)
)
return b
%timeit run_cavity()
1 loop, best of 3: 415 ms per loop