Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions cfd_viz/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,15 @@
)
from .common import VTKData, ensure_dirs, find_vtk_files, read_vtk_file
from .convert import from_cfd_python, from_simulation_result, to_cfd_python
from .defaults import (
UNSET,
PlotDefaults,
get_defaults,
load_config_file,
plot_context,
reset_defaults,
set_defaults,
)
from .info import get_recommended_settings, get_system_info, print_system_info
from .quick import quick_plot, quick_plot_data, quick_plot_result
from .stats import (
Expand All @@ -67,6 +76,14 @@

__all__ = [
"__version__",
# Configuration
"UNSET",
"PlotDefaults",
"get_defaults",
"set_defaults",
"reset_defaults",
"plot_context",
"load_config_file",
# I/O
"VTKData",
"read_vtk_file",
Expand Down
56 changes: 40 additions & 16 deletions cfd_viz/animation/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,17 @@
from matplotlib import animation
from matplotlib.figure import Figure

from cfd_viz.defaults import UNSET, get_defaults, resolve

from .frames import AnimationFrames, FrameData
from .renderers import create_velocity_colormap


def export_frame_to_image(
frame: FrameData,
output_path: Union[str, Path],
figsize: Tuple[int, int] = (16, 12),
dpi: int = 150,
figsize: Tuple[int, int] = UNSET,
dpi: int = UNSET,
include_vectors: bool = True,
include_streamlines: bool = True,
include_pressure: bool = True,
Expand All @@ -41,6 +43,9 @@ def export_frame_to_image(
include_streamlines: Whether to include streamlines subplot.
include_pressure: Whether to include pressure subplot.
"""
figsize = resolve(figsize, "figsize")
dpi = resolve(dpi, "dpi")
defaults = get_defaults()
fig, axes = plt.subplots(2, 2, figsize=figsize)
axes = axes.flatten()

Expand All @@ -55,7 +60,9 @@ def export_frame_to_image(

# Panel 1: Velocity magnitude contours
if velocity_mag is not None:
im1 = axes[0].contourf(X, Y, velocity_mag, levels=20, cmap="viridis")
im1 = axes[0].contourf(
X, Y, velocity_mag, levels=defaults.levels, cmap=defaults.cmap
)
axes[0].set_title("Velocity Magnitude")
axes[0].axis("equal")
plt.colorbar(im1, ax=axes[0])
Expand All @@ -77,7 +84,7 @@ def export_frame_to_image(
u_sub,
v_sub,
np.sqrt(u_sub**2 + v_sub**2),
cmap="plasma",
cmap=defaults.sequential_cmap,
scale_units="xy",
angles="xy",
)
Expand All @@ -91,27 +98,33 @@ def export_frame_to_image(
try:
if velocity_mag is not None:
axes[2].streamplot(
X, Y, u, v, color=velocity_mag, cmap="viridis", density=2
X, Y, u, v, color=velocity_mag, cmap=defaults.cmap, density=2
)
else:
axes[2].streamplot(X, Y, u, v, density=2)
except ValueError:
# Streamplot can fail with certain grid configurations
if velocity_mag is not None:
axes[2].contourf(X, Y, velocity_mag, levels=20, cmap="viridis")
axes[2].contourf(
X, Y, velocity_mag, levels=defaults.levels, cmap=defaults.cmap
)
axes[2].set_title("Flow Streamlines")
axes[2].axis("equal")
else:
axes[2].axis("off")

# Panel 4: Pressure or combined view
if include_pressure and p is not None:
im4 = axes[3].contourf(X, Y, p, levels=20, cmap="RdBu_r")
im4 = axes[3].contourf(
X, Y, p, levels=defaults.levels, cmap=defaults.diverging_cmap
)
axes[3].set_title("Pressure Field")
axes[3].axis("equal")
plt.colorbar(im4, ax=axes[3])
elif velocity_mag is not None and u is not None and v is not None:
im4 = axes[3].contourf(X, Y, velocity_mag, levels=20, cmap="viridis", alpha=0.7)
im4 = axes[3].contourf(
X, Y, velocity_mag, levels=defaults.levels, cmap=defaults.cmap, alpha=0.7
)
axes[3].streamplot(X, Y, u, v, color="white", density=1, linewidth=0.8)
axes[3].set_title("Combined: Magnitude + Streamlines")
plt.colorbar(im4, ax=axes[3])
Expand All @@ -129,8 +142,8 @@ def export_animation_frames(
animation_frames: AnimationFrames,
output_dir: Union[str, Path],
prefix: str = "frame",
figsize: Tuple[int, int] = (16, 12),
dpi: int = 150,
figsize: Tuple[int, int] = UNSET,
dpi: int = UNSET,
) -> List[str]:
"""Export all frames from AnimationFrames to image files.

Expand All @@ -144,6 +157,8 @@ def export_animation_frames(
Returns:
List of exported file paths.
"""
figsize = resolve(figsize, "figsize")
dpi = resolve(dpi, "dpi")
output_dir = Path(output_dir)
output_dir.mkdir(parents=True, exist_ok=True)

Expand All @@ -164,7 +179,7 @@ def save_animation(
output_path: Union[str, Path],
writer: str = "pillow",
fps: int = 5,
dpi: int = 100,
dpi: int = UNSET,
) -> None:
"""Save a matplotlib animation to file.

Expand All @@ -175,6 +190,7 @@ def save_animation(
fps: Frames per second.
dpi: Resolution.
"""
dpi = resolve(dpi, "dpi")
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)

Expand All @@ -183,7 +199,7 @@ def save_animation(

def create_comprehensive_frame_figure(
frame: FrameData,
figsize: Tuple[int, int] = (18, 10),
figsize: Tuple[int, int] = UNSET,
) -> Figure:
"""Create a comprehensive 2x3 figure for a single frame.

Expand All @@ -194,6 +210,8 @@ def create_comprehensive_frame_figure(
Returns:
Matplotlib Figure object.
"""
figsize = resolve(figsize, "figsize")
defaults = get_defaults()
fig, axes = plt.subplots(2, 3, figsize=figsize)
axes = axes.flatten()

Expand All @@ -207,15 +225,19 @@ def create_comprehensive_frame_figure(
velocity_cmap = create_velocity_colormap()

# Velocity magnitude
im0 = axes[0].contourf(X, Y, velocity_mag, levels=20, cmap=velocity_cmap)
im0 = axes[0].contourf(
X, Y, velocity_mag, levels=defaults.levels, cmap=velocity_cmap
)
axes[0].set_title("Velocity Magnitude", fontweight="bold")
axes[0].set_xlabel("X")
axes[0].set_ylabel("Y")
axes[0].set_aspect("equal")
plt.colorbar(im0, ax=axes[0])

# Pressure
im1 = axes[1].contourf(X, Y, p, levels=20, cmap="RdBu_r")
im1 = axes[1].contourf(
X, Y, p, levels=defaults.levels, cmap=defaults.diverging_cmap
)
contours = axes[1].contour(X, Y, p, levels=8, colors="black", linewidths=0.5)
axes[1].clabel(contours, inline=True, fontsize=8)
axes[1].set_title("Pressure Field", fontweight="bold")
Expand All @@ -225,7 +247,7 @@ def create_comprehensive_frame_figure(
plt.colorbar(im1, ax=axes[1])

# Vorticity
im2 = axes[2].contourf(X, Y, vorticity, levels=20, cmap="seismic")
im2 = axes[2].contourf(X, Y, vorticity, levels=defaults.levels, cmap="seismic")
axes[2].set_title("Vorticity", fontweight="bold")
axes[2].set_xlabel("X")
axes[2].set_ylabel("Y")
Expand Down Expand Up @@ -262,7 +284,9 @@ def create_comprehensive_frame_figure(
axes[4].set_aspect("equal")

# Combined analysis
axes[5].contourf(X, Y, velocity_mag, levels=20, cmap=velocity_cmap, alpha=0.8)
axes[5].contourf(
X, Y, velocity_mag, levels=defaults.levels, cmap=velocity_cmap, alpha=0.8
)
axes[5].contour(X, Y, p, levels=6, colors="white", linewidths=1.5)
if np.any(vorticity != 0):
vort_threshold = np.percentile(np.abs(vorticity), 85)
Expand Down
Loading