VoxelEngine/FluidSim/StaggeredArray.py
2021-12-20 17:21:46 +01:00

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)