def initialize():
"""Wrapper that initializes all necessary data structures.
Returns
U (3D array) : the state-variables-3D-matrix (populating a x,y grid)
- shape: (3, Nx + 2 * Ng, Ny + 2 * Ng)
- U[0] : state varables [h, hu, hv]
- U[1] : y dimention (rows)
- U[2] : x dimention (columns)
h_hist (array) : holds the step-wise height solutions for the
post-processing animation
t_hist (array) : holds the step-wise times for the post-
processing animation
U_ds (memmap) : holds the state-variables 3D matrix data for all
the timesteps
(conf.MAX_ITERS, 3, Nx + 2 * Ng, Ny + 2 * Ng)
"""
logger.log('Initialization...')
U = _init_U()
h_hist = _init_h_hist(U)
t_hist = t_hist = np.zeros(len(h_hist), dtype=conf.DTYPE)
if conf.SAVE_DS_FOR_ML:
U_ds = _init_U_ds(U)
else:
U_ds = None
return U, h_hist, t_hist, U_ds
def _plot_basin(sub):
"""Plots the basin that contains the fluid.
Args:
sub (axes.SubplotBase) : Axes3D subplot object
"""
if conf.SHOW_BASIN is True:
# Make the basin a bit wider, because water appears to be out of the
# basin because of the perspective mode.
X_bas, Y_bas = np.meshgrid(conf.CX[conf.Ng - 1:conf.Nx + conf.Ng + 1],
conf.CY[conf.Ng - 1:conf.Ny + conf.Ng + 1])
# BASIN
BASIN = np.zeros((conf.Ny + 2, conf.Nx + 2))
# left-right walls
BASIN[:, 0] = 2.5
BASIN[:, conf.Nx + 1] = 2.5
# top-bottom walls
BASIN[0, :] = 2.5
BASIN[conf.Ny + 1, :] = 2.5
sub.plot_surface(X_bas,
Y_bas,
BASIN,
rstride=2,
cstride=2,
linewidth=0,
color=(0.4, 0.4, 0.5, 0.1))
elif conf.SHOW_BASIN is False:
pass
else:
logger.log("Configure SHOW_BASIN. Options: True, False")
def drop_iters_list():
"""list with the simulation iters at which a drop is going to fall."""
drop_iters = [0]
iters_cumsum = 0
i = 0
if conf.ITERS_BETWEEN_DROPS_MODE == "custom":
while iters_cumsum <= conf.MAX_ITERS:
iters_cumsum += conf.CUSTOM_ITERS_BETWEEN_DROPS[i % 10]
drop_iters.append(iters_cumsum)
i += 1
elif conf.ITERS_BETWEEN_DROPS_MODE == "random":
while iters_cumsum <= conf.MAX_ITERS:
iters_cumsum += random.randint(60, 120)
drop_iters.append(iters_cumsum)
else:
logger.log("Configure ITERS_BETWEEN_DROPS_MODE | options:"
" 'fixed', 'custom', 'random'")
# # Remove drop iterations that fall after MAX_ITERS.
# last_drop_idx = np.searchsorted(drop_iters, conf.MAX_ITERS)
# drop_iters = drop_iters[:last_drop_idx]
# # Overwrite max number of drops.
# conf.MAX_N_DROPS = len(drop_iters)
if conf.SAVE_DS_FOR_ML:
# It is needed to retrieve the new drop frames, because these frames
# cannot be used as labels (the previous frame cannot know when and
# where a new drop will fall).
# ds_shape
dss = ds_shape()
file_name = f"drop_iters_list_{dss[0]}x{dss[1]}x{dss[2]}x{dss[3]}.npy"
np.save(os.path.join(os.getcwd(), file_name), drop_iters)
return drop_iters
def _save_ani(ani, fps, dpi):
"""Saves the animation in .mp4 and .gif formats.
Args:
ani (obj) : animation.FuncAnimation() object
fps (int) : frames per second
dpi (int) : dots per inch
"""
if conf.SAVE_ANIMATION is True:
# file name
date_n_time = str(datetime.now())[:19]
# Replace ':' with '-' for compatibility with windows file formating.
date_n_time = date_n_time.replace(':', '-').replace(' ', '_')
file_name = conf.MODE + '_animation_' + date_n_time
# Configure the writer
plt.rcParams['animation.ffmpeg_path'] = conf.PATH_TO_FFMPEG
FFwriter = animation.FFMpegWriter(fps=fps,
bitrate=-1,
extra_args=[
'-r',
str(fps), '-pix_fmt', 'yuv420p',
'-vcodec', 'libx264',
'-qscale:v', '1'
])
# Save
try:
ani.save(file_name + '.' + conf.VID_FORMAT,
writer=FFwriter,
dpi=dpi)
# log only if a log file is already initialzed
if isinstance(logger.find_open_log(), str):
logger.log('Animation saved as: ' + file_name + '.' +
conf.VID_FORMAT + ' | fps: ' + str(fps))
# convert to a lighter gif
cmd = 'ffmpeg -i ' + file_name + '.' + conf.VID_FORMAT + ' -vf ' \
'"fps=' + str(fps) + ',scale=240:-1:flags=lanczos,split' \
'[s0][s1];[s0]palettegen[p];[s1][p]paletteuse" -hide_banner' \
' -loglevel panic -loop 0 ' + file_name + '.gif'
os.system(cmd)
if isinstance(logger.find_open_log(), str):
logger.log('Animation saved as: ' + file_name + '.gif' +
' | fps: ' + str(fps))
except FileNotFoundError:
logger.log('Configure PATH_TO_FFMPEG')
elif conf.SAVE_ANIMATION is False:
pass
else:
logger.log("Configure SAVE_ANIMATION | Options: True, False")
def _save_ani(ani):
"""Saves the animation in .mp4 and .gif formats.
Args:
ani (obj) : animation.FuncAnimation() object
"""
dpi = conf.DPI
fps = conf.FPS
# file name
date_n_time = str(datetime.now())[:19]
# Replace ':' with '-' for compatibility with windows file formating.
date_n_time = date_n_time.replace(':', '-').replace(' ', '_')
file_name = conf.MODE + '_animation_' + date_n_time
# Configure the writer
plt.rcParams['animation.ffmpeg_path'] = conf.PATH_TO_FFMPEG
FFwriter = animation.FFMpegWriter(
fps=fps, bitrate=-1,
extra_args=['-r', str(fps), '-pix_fmt', 'yuv420p', '-vcodec',
'libx264', '-qscale:v', '1']
)
# Save
try:
ani.save(os.path.join(conf.SAVE_DIR, f"{file_name}.{conf.VID_FORMAT}"),
writer=FFwriter,
dpi=dpi)
# log only if a log file is already initialzed
if isinstance(logger.find_open_log(), str):
logger.log(f'Animation saved as: {file_name}.'
f'{conf.VID_FORMAT} | fps: {fps}')
# convert to a lighter gif
cmd = (f'ffmpeg -i {file_name}.{conf.VID_FORMAT} -vf '
f'"fps={fps},scale=240:-1:flags=lanczos,split'
'[s0][s1];[s0]palettegen[p];[s1][p]paletteuse" -hide_banner'
f' -loglevel panic -loop 0 {file_name}.gif')
os.system(cmd)
if isinstance(logger.find_open_log(), str):
logger.log(f'Animation saved as: {file_name}.gif | fps: {fps}')
except FileNotFoundError:
logger.log('Configure PATH_TO_FFMPEG')
def plot_from_dat(time, it):
"""Creates and saves a frame as .png, reading data from a .dat file.
Args:
time (float) : current time
it (int) : current itereration
"""
# Create ./session directory for saving the results.
utils.child_dir("session")
# Extract data from dat.
X, Y, Z, Nx, Ny = _data_from_dat(it)
# plot {
#
fig = plt.figure(figsize=(9.6, 6.4), dpi=112) # 1080x720
sub = fig.add_subplot(111, projection="3d")
plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0)
if conf.PLOTTING_STYLE == 'water':
if conf.ROTATION:
# Azimuthal rotate every 8 frames and vetical every 20 frames
azimuth = 45 + it / 8
elevation = 55 - it / 20
sub.view_init(elevation, azimuth)
else:
sub.view_init(45, 55)
sub.plot_surface(X,
Y,
Z,
rstride=1,
cstride=1,
linewidth=0,
color=(0.251, 0.643, 0.875, 0.95),
shade=True,
antialiased=False)
elif conf.PLOTTING_STYLE == 'contour':
sub.view_init(45, 55)
sub.contour3D(X, Y, Z, 140, cmap='plasma', vmin=0.6, vmax=1.4)
elif conf.PLOTTING_STYLE == 'wireframe':
if conf.ROTATION:
# Azimuthal rotate every 3 frames and vetical every 4 frames
azimuth = 45 + it / 3
elevation = 55 - it / 4
sub.view_init(elevation, azimuth)
else:
sub.view_init(45, 55)
sub.plot_wireframe(
X,
Y,
Z,
rstride=2,
cstride=2,
linewidth=1,
)
else:
styles = ['water', 'contour', 'wireframe']
logger.log(f"Configure PLOTTING_STYLE | options: {styles}")
fig.gca().set_zlim([-0.5, 4])
plt.title(f"time: {time: >{6}.3f}\titer: {it: >{4}d}", y=0.8, fontsize=18)
sub.title.set_position([0.51, 0.80])
plt.rcParams.update({'font.size': 20})
plt.axis('off')
# Render the basin that contains the fluid.
_plot_basin(sub)
# save
# fig.tight_layout()
fig_file = os.path.join("session", f"iter_{it:0>{4}}.png")
plt.savefig(fig_file)
plt.close()
def animate(h_hist, t_hist=None):
"""Generates and saves an animation of the simulation.
Args:
h_hist (array) : array of iter-wise heights solutions
t_hist (array) : holds the iter-wise times
Returns:
ani (animation.FuncAnimation) : It is returned in case of ipython
"""
# resolution = figsize * dpi
# --------------------------
# example:
# figsize = (9.6, 5.4), dpi=200
# resolution: 1920x1080 (1920/200=9.6)
fps = conf.FPS
dpi = conf.DPI
figsize = conf.FIGSIZE
# total frames
frames = len(h_hist)
# X, Y, Z
X, Y = np.meshgrid(conf.CX[conf.Ng:-conf.Ng], conf.CY[conf.Ng:-conf.Ng])
Z = h_hist
# Plot configuration
fig = plt.figure(figsize=figsize, dpi=dpi)
sub = fig.add_subplot(111, projection="3d")
if conf.ROTATION:
sub.view_init(45, 55)
else:
sub.view_init(30, 20)
plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0)
fig.gca().set_zlim([-0.5, 4])
plt.axis('off')
if t_hist is None:
ani_title = f"iter: {0:>{5}d}"
else:
ani_title = f"time: {t_hist[0]:>{6}.3f} iter: {0:>{5}d}"
plt.title(ani_title, y=0.8, fontsize=18)
sub.title.set_position([0.51, 0.80])
plt.rcParams.update({'font.size': 20})
# Program name and version text
fig.text(0.85, 0.06, s=f"MattFlow v{__version__}", fontsize=16, c='navy')
# Plot initialization
if conf.PLOTTING_STYLE == 'water':
plot = [
sub.plot_surface(X,
Y,
Z[0],
rstride=1,
cstride=1,
linewidth=0,
color=(0.251, 0.643, 0.875, 0.9),
shade=True,
antialiased=False)
]
elif conf.PLOTTING_STYLE == 'contour':
plot = [
sub.contour3D(X, Y, Z[0], 150, cmap='ocean', vmin=0.5, vmax=1.5)
]
elif conf.PLOTTING_STYLE == 'wireframe':
plot = [
sub.plot_wireframe(X, Y, Z[0], rstride=2, cstride=2, linewidth=1)
]
else:
logger.log("Configure PLOTTING_STYLE | options: 'water', 'contour',",
"'wireframe'")
# Render the basin that contains the fluid.
_plot_basin(sub)
# Generate the animation.
ani = animation.FuncAnimation(fig,
_update_plot,
frames,
fargs=(X, Y, Z, plot, fig, sub, t_hist,
ani_title),
interval=1000 / fps,
repeat=True)
# Save the animation.
_save_ani(ani, fps, dpi)
# Play the animation.
if conf.SHOW_ANIMATION is True:
logger.log('Playing animation...')
try:
# In case of jupyter notebook, don't run plt.show(), to prevent
# displaying a static figure. Instead, return the Funcanimation
# object.
get_ipython()
return ani
except NameError:
plt.show()
elif conf.SHOW_ANIMATION is False:
pass
else:
logger.log("Configure SHOW_ANIMATION | Options: True, False")
def _solve(U, delta_t, it, drops_count, drop_its_iterator, next_drop_it):
"""Evaluates the state variables (h, hu, hv) at a new time-step.
It can be used in a for/while loop, iterating through each time-step.
Args:
U (3D array) : the state variables, populating a x,y grid
delta_t (float) : time discretization step
it (int) : current iteration
drops_count (int) : number of drops been generated
drop_its_iterator (iterator)
: iterator of the drop_its list (the list with
the iters at which a new drop falls)
next_drop_it (int) : the next iteration at which a new drop will fall
Returns:
U, drops_count, drop_its_iterator, next_drop_it
"""
# Retrieve the mesh
Ng = conf.Ng
cx = conf.CX
cy = conf.CY
cellArea = conf.dx * conf.dy
# Simulation mode
# ---------------
# 'single drop': handled at the initialization
if conf.MODE == 'single drop':
pass
# 'drops': specified number of drops are generated at specified frequency
elif conf.MODE == 'drops':
if conf.ITERS_BETWEEN_DROPS_MODE == "fixed":
drop_condition = (it % conf.FIXED_ITERS_BETWEEN_DROPS == 0
and drops_count < conf.MAX_N_DROPS)
else: # conf.ITERS_TO_NEXT_DROP_MODE in ["custom", "random"]
drop_condition = (it == next_drop_it
and drops_count < conf.MAX_N_DROPS)
if drop_condition:
U[0, :, :] = initializer.drop(U[0, :, :], drops_count + 1)
drops_count += 1
if (conf.ITERS_BETWEEN_DROPS_MODE in ["custom", "random"]
and drops_count < conf.MAX_N_DROPS):
next_drop_it = next(drop_its_iterator)
# 'rain': random number of drops are generated at random frequency
elif conf.MODE == 'rain':
if it % random.randrange(1, 15) == 0:
simultaneous_drops = range(random.randrange(1, 2))
for _ in simultaneous_drops:
U[0, :, :] = initializer.drop(U[0, :, :])
else:
modes = ['single drop', 'drops', 'rain']
logger.log(f"Configure MODE | options: {modes}")
# Numerical scheme
# flux.flux() returns the total flux entering and leaving each cell.
if conf.SOLVER_TYPE == 'Lax-Friedrichs Riemann':
U[:, Ng:-Ng, Ng:-Ng] += delta_t / cellArea * flux.flux(U)
elif conf.SOLVER_TYPE == '2-stage Runge-Kutta':
# 1st stage
U_pred = U
U_pred[:, Ng:-Ng, Ng:-Ng] += delta_t / cellArea * flux.flux(U)
# 2nd stage
U[:, Ng: -Ng, Ng: -Ng] = \
0.5 * (U[:, Ng: -Ng, Ng: -Ng]
+ U_pred[:, Ng: -Ng, Ng: -Ng]
+ delta_t / cellArea * flux.flux(U_pred)
)
else:
solver_types = ['Lax-Friedrichs Riemann', '2-stage Runge-Kutta']
logger.log(f"Configure SOLVER_TYPE | Options: {solver_types}")
return U, drops_count, drop_its_iterator, next_drop_it
'''
def simulate():
time = 0
U, h_hist, t_hist, U_ds = initializer.initialize()
drops_count = 1
# idx of the frame saved in h_hist
saving_frame_idx = 0
# Counts up to conf.FRAMES_PER_PERIOD (1st frame saved at initialization).
consecutive_frames_counter = 1
if conf.ITERS_BETWEEN_DROPS_MODE in ["custom", "random"]:
# List with the simulation iterations at which a drop is going to fall
drop_its = utils.drop_iters_list()
# Drop the 0th drop
drop_its_iterator = iter(drop_its[1:])
# The iteration at which the next drop will fall
try:
next_drop_it = next(drop_its_iterator)
except StopIteration:
drop_its_iterator = None
next_drop_it = None
else:
drop_its_iterator = None
next_drop_it = None
for it in range(1, conf.MAX_ITERS):
# Time discretization step (CFL condition)
delta_t = dt(U)
# Update current time
time += delta_t
if time > conf.STOPPING_TIME:
break
# Apply boundary conditions (reflective)
U = bcmanager.update_ghost_cells(U)
# Numerical iterative scheme
U, drops_count, drop_its_iterator, next_drop_it = _solve(
U=U,
delta_t=delta_t,
it=it,
drops_count=drops_count,
drop_its_iterator=drop_its_iterator,
next_drop_it=next_drop_it)
if conf.WRITE_DAT:
dat_writer.write_dat(
U[0, conf.Ng:conf.Ny + conf.Ng, conf.Ng:conf.Nx + conf.Ng],
time, it)
mattflow_post.plot_from_dat(time, it)
elif not conf.WRITE_DAT:
# Append current frame to the list, to be animated at
# post-processing.
if it % conf.FRAME_SAVE_FREQ == 0:
# Zero the counter, when a perfect division occurs.
consecutive_frames_counter = 0
if consecutive_frames_counter < conf.FRAMES_PER_PERIOD:
saving_frame_idx += 1
h_hist[saving_frame_idx] = \
U[0, conf.Ng: -conf.Ng, conf.Ng: -conf.Ng]
# time * 10 is insertd, because space is scaled about x10.
t_hist[saving_frame_idx] = time * 10
consecutive_frames_counter += 1
if conf.SAVE_DS_FOR_ML:
U_ds[it] = U[:, conf.Ng:-conf.Ng, conf.Ng:-conf.Ng]
else:
logger.log("Configure WRITE_DAT | Options: True, False")
logger.log_timestep(it, time)
# Clean-up the memmap
if conf.DUMP_MEMMAP and conf.WORKERS > 1:
utils.delete_memmap()
return h_hist, t_hist, U_ds