def compute_pressure(N, n1, n2, uk, vk):
pk = np.zeros_like(uk)
fku = compute_fk.compute_fk(N, n1, n2, uk, vk, dim=0)
fkv = compute_fk.compute_fk(N, n1, n2, uk, vk, dim=1)
for i in range(N):
for j in range(N):
if n1[i] == n2[j] == 0:
pk[i, j] = 1
else:
pk[i,j] = 1j*(n1[i]*fku[i,j] + n2[j]*fkv[i,j]\
)/(n1[i]**2 + n2[j]**2)
return np.real(np.fft.ifftn(np.fft.ifftshift(pk)) * N**2)
def fv(uk, vk):
fku = compute_fk.compute_fk(N, n1, n2, uk, vk, dim=0)
fkv = compute_fk.compute_fk(N, n1, n2, uk, vk, dim=1)
a = np.zeros_like(vk)
d = np.zeros_like(vk)
for i in range(N):
for j in range(N):
if n1[i] == n2[j] == 0:
a[i, j] = 0
d[i, j] = 0
else:
a[i,j] = (n1[i]*fku[i,j] + n2[j]*fkv[i,j]\
)*n2[j]/(n1[i]**2 + n2[j]**2)
d[i, j] = nu * (n1[i]**2 + n2[j]**2) * vk[i, j]
return fkv - a - d
def fu(uk, vk):
fku = compute_fk.compute_fk(N, n1, n2, uk, vk, dim=0)
fkv = compute_fk.compute_fk(N, n1, n2, uk, vk, dim=1)
a = np.zeros_like(uk)
d = np.zeros_like(uk)
for i in range(N):
for j in range(N):
# singular k**2
if n1[i] == n2[j] == 0:
a[i, j] = 0
d[i, j] = 0
else:
a[i,j] = (n1[i]*fku[i,j] + n2[j]*fkv[i,j])*\
n1[i]/(n1[i]**2 + n2[j]**2)
d[i, j] = nu * (n1[i]**2 + n2[j]**2) * uk[i, j]
return fku - a - d
