108 lines
No EOL
2.9 KiB
Python
108 lines
No EOL
2.9 KiB
Python
import numpy as np
|
|
from scipy.signal import convolve
|
|
|
|
|
|
class StaggeredArray2D:
|
|
def __init__(self, x_n, y_n):
|
|
"""
|
|
creates a staggered array
|
|
:param x_n: x size of the array
|
|
:param y_n: y size of the array
|
|
"""
|
|
self.x_n = x_n
|
|
self.y_n = y_n
|
|
self.p = np.zeros((x_n, y_n), dtype=np.float)
|
|
|
|
self.u_x = np.zeros((x_n + 1, y_n), dtype=np.float)
|
|
self.u_y = np.zeros((x_n, y_n + 1), dtype=np.float)
|
|
|
|
self.has_fluid = np.zeros((x_n, y_n), dtype=np.bool)
|
|
self.density = np.zeros((x_n, y_n), dtype=np.float)
|
|
|
|
def get_velocity(self, x, y):
|
|
"""
|
|
get mid point value for the coordinates
|
|
:param x: x coordinate
|
|
:param y: y coordinate
|
|
:return:
|
|
"""
|
|
assert 0 <= x < self.x_n, 'x value out of bounds!'
|
|
assert 0 <= y < self.y_n, 'y value out of bounds!'
|
|
|
|
lower_x = self.u_x[x, y]
|
|
upper_x = self.u_x[x + 1, y]
|
|
|
|
lower_y = self.u_y[x, y]
|
|
upper_y = self.u_y[x, y + 1]
|
|
|
|
return (lower_x + upper_x) / 2.0, (lower_y + upper_y) / 2.0
|
|
|
|
def get_velocity_arrays(self):
|
|
c_x = np.array([[0.5], [0.5]])
|
|
u_x = convolve(self.u_x, c_x, mode='valid')
|
|
|
|
c_y = np.array([[0.5, 0.5]])
|
|
u_y = convolve(self.u_y, c_y, mode='valid')
|
|
|
|
return u_x, u_y
|
|
|
|
|
|
class StaggeredArray3D:
|
|
def __init__(self, x_n, y_n, z_n):
|
|
"""
|
|
creates a staggered array
|
|
:param x_n: x size of the array
|
|
:param y_n: y size of the array
|
|
:param z_n: z size of the array
|
|
"""
|
|
|
|
self.x_n = x_n
|
|
self.y_n = y_n
|
|
self.z_n = z_n
|
|
|
|
self.p = np.zeros((x_n, y_n, z_n), dtype=np.float)
|
|
|
|
self.u_x = np.zeros((x_n + 1, y_n, z_n), dtype=np.float)
|
|
self.u_y = np.zeros((x_n, y_n + 1, z_n), dtype=np.float)
|
|
self.u_z = np.zeros((x_n, y_n, z_n + 1), dtype=np.float)
|
|
|
|
self.has_fluid = np.zeros((x_n, y_n, z_n), dtype=np.bool)
|
|
|
|
def get_velocity(self, x, y, z):
|
|
"""
|
|
get mid point value for the coordinates
|
|
:param x: x coordinate
|
|
:param y: y coordinate
|
|
:param z: z coordinate
|
|
:return:
|
|
"""
|
|
assert 0 <= x < self.x_n, 'x value out of bounds!'
|
|
assert 0 <= y < self.y_n, 'y value out of bounds!'
|
|
assert 0 <= z < self.z_n, 'y value out of bounds!'
|
|
|
|
lower_x = self.u_x[x, y, z]
|
|
upper_x = self.u_x[x + 1, y, z]
|
|
|
|
lower_y = self.u_y[x, y, z]
|
|
upper_y = self.u_y[x, y + 1, z]
|
|
|
|
lower_z = self.u_z[x, y, z]
|
|
upper_z = self.u_z[x, y, z + 1]
|
|
|
|
return (lower_x + upper_x) / 2.0, (lower_y + upper_y) / 2.0, (lower_z + upper_z) / 2.0
|
|
|
|
|
|
if __name__ == '__main__':
|
|
sa = StaggeredArray2D(10, 10)
|
|
|
|
for x in range(11):
|
|
for y in range(10):
|
|
sa.u_x[x, y] = y
|
|
|
|
for x in range(10):
|
|
for y in range(11):
|
|
sa.u_y[x, y] = x
|
|
|
|
ux, uy = sa.get_velocity_arrays()
|
|
|
|
ux2, uy2 = sa.get_velocity(0, 0) |