Spaces:
Running
Running
File size: 14,609 Bytes
7eb1167 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 | # visualization_node.py
import os
import subprocess
import sys
from typing import List, Optional
from pydantic import BaseModel, Field
from utils import save_file
import glob
# Helper to get the .foam file name
def get_foam_file(case_dir):
case_dir_name = os.path.basename(os.path.normpath(case_dir))
return f"{case_dir_name}.foam"
VISUALIZATION_SYSTEM_PROMPT = (
"You are an expert in OpenFOAM post-processing and PyVista Python scripting. "
"Your task is to generate a PyVista Python script that visualizes the specified data from the OpenFOAM case. "
"The script should load the OpenFOAM case data by reading the .foam file (e.g., 'runs.foam') in the case directory using PyVista, display the geometry, and color the surface by the specified field (e.g., 'U' for velocity). "
"Ensure the script shows the geometry, sets up the colorbar, and saves the visualization as a PNG file. "
"Use coolwarm colormap by default."
"The script must save the visualization as a PNG file, and the output image must contain the geometry and the colorbar, not just the colorbar. "
"IMPORTANT: Return ONLY the Python code without any markdown formatting, code block markers, or explanatory text. "
"The script should start with the necessary imports, read the .foam file using PyVista, and end with the screenshot saving."
)
ERROR_FIX_SYSTEM_PROMPT = (
"You are an expert in PyVista Python scripting and OpenFOAM visualization. "
"Your task is to fix the provided PyVista Python script that encountered an error. "
"Ensure the script loads the OpenFOAM case data by reading the .foam file (e.g., 'runs.foam') in the case directory using PyVista, displays the geometry, and colors the surface by the specified field. "
"Make sure the script shows the geometry, sets up the colorbar, and saves the visualization as a PNG file. "
"Use coolwarm colormap by default."
"The script must save the visualization as a PNG file, and the output image must contain the geometry and the colorbar, not just the colorbar. "
"IMPORTANT: Return ONLY the Python code without any markdown formatting, code block markers, or explanatory text. "
"The script should start with the necessary imports, read the .foam file using PyVista, and end with the screenshot saving."
)
class PlotConfigPydantic(BaseModel):
"""Configuration for plotting parameters"""
plot_type: str = Field(description="Type of plot (e.g., 'contour', 'vector', 'streamline', 'time_series')")
field_name: str = Field(description="Field to plot (e.g., 'U', 'p', 'T', 'rho')")
time_step: Optional[str] = Field(default=None, description="Time step to plot (if None, use latest)")
output_format: str = Field(default="png", description="Output format for plots")
output_path: str = Field(description="Path to save the plot")
class VisualizationPlanPydantic(BaseModel):
"""Plan for visualization tasks"""
plots: List[PlotConfigPydantic] = Field(description="List of plots to generate")
class VisualizationAnalysisPydantic(BaseModel):
"""Analysis of user requirements for visualization needs"""
primary_field: str = Field(description="Primary field to visualize (e.g., 'U', 'p', 'T', 'rho')")
plot_type: str = Field(description="Type of plot requested (e.g., 'contour', 'vector', 'streamline', 'time_series')")
time_step: Optional[str] = Field(default=None, description="Specific time step to plot (if mentioned)")
plane_info: Optional[str] = Field(default=None, description="Plane information if 2D slice is requested (e.g., 'Z plane', 'X=0.5')")
additional_fields: List[str] = Field(default=[], description="Additional fields that might be useful to visualize")
visualization_priority: str = Field(description="Priority of visualization (e.g., 'high', 'medium', 'low')")
def visualization_node(state):
"""
Visualization node: Creates PyVista visualizations from the successfully generated OpenFOAM case.
This node uses the successfully generated code and user_requirement to create PyVista visualizations.
Updates state with:
- plot_configs: List of plot configurations
- plot_outputs: List of generated plot file paths
- visualization_summary: Summary of generated visualizations
- pyvista_visualization: PyVista visualization results
"""
config = state["config"]
user_requirement = state["user_requirement"]
case_dir = state["case_dir"]
print(f"============================== Visualization (PyVista) ==============================")
# Ensure case_dir is absolute
case_dir = os.path.abspath(case_dir)
if not os.path.exists(case_dir):
print(f"Case directory does not exist: {case_dir}")
return {
**state,
"plot_configs": [],
"plot_outputs": [],
"visualization_summary": {"error": f"Case directory does not exist: {case_dir}"},
"pyvista_visualization": {"success": False, "error": f"Case directory does not exist: {case_dir}"}
}
# Touch the .foam file before generating the visualization script
foam_file = get_foam_file(case_dir)
foam_file_path = os.path.join(case_dir, foam_file)
with open(foam_file_path, 'a'):
os.utime(foam_file_path, None)
# Initialize loop counter
current_loop = 0
error_logs = []
max_loop = state['config'].max_loop
while current_loop < max_loop:
current_loop += 1
print(f"Attempt {current_loop} of {max_loop}")
# Create visualization script
viz_prompt = (
f"<case_directory>{case_dir}</case_directory>\n"
f"<foam_file>{foam_file}</foam_file>\n"
f"<visualization_requirements>{state['user_requirement']}</visualization_requirements>\n"
f"<previous_errors>{error_logs}</previous_errors>\n"
f"Please create a PyVista Python script that visualizes the specified data by reading the .foam file ('{foam_file}')."
"Save the visualization as PNG file named visualization.png if not specified otherwise in the user requirement."
)
viz_script = state["llm_service"].invoke(viz_prompt, VISUALIZATION_SYSTEM_PROMPT)
# Save the visualization script
script_path = os.path.join(case_dir, "visualization.py")
save_file(script_path, viz_script)
# Execute the script using Python
try:
result = subprocess.run(
[sys.executable, script_path],
cwd=case_dir,
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
print(f"Finished command: Return Code {result.returncode}")
error_logs = []
# Check if any PNG output image was created
png_files = glob.glob(os.path.join(case_dir, "*.png"))
if png_files:
output_image = png_files[0] # Use the first PNG file found
print(f"PyVista visualization created successfully: {output_image}")
# Create plot configs and outputs in the expected format
plot_configs = [
{
"plot_type": "pyvista_2d",
"field_name": "U", # Default field, could be enhanced to detect from script
"time_step": "latest",
"output_format": "png",
"output_path": output_image
}
]
plot_outputs = [output_image]
visualization_summary = {
"total_plots_generated": len(plot_outputs),
"plot_types": ["pyvista_2d"],
"fields_visualized": ["U"],
"output_directory": case_dir,
"pyvista_success": True
}
pyvista_result = {
"success": True,
"output_image": output_image,
"script": viz_script
}
print(f"Generated {len(plot_outputs)} plots")
print(f"PyVista visualization saved to: {output_image}")
print("============================== Visualization Complete ==============================")
return {
**state,
"plot_configs": plot_configs,
"plot_outputs": plot_outputs,
"visualization_summary": visualization_summary,
"pyvista_visualization": pyvista_result
}
else:
error_logs.append("Visualization script executed but no PNG output image was created")
except subprocess.CalledProcessError as e:
error_message = f"Error executing visualization script: {str(e)}"
if e.stdout:
error_message += f"\nSTDOUT:\n{e.stdout.decode() if isinstance(e.stdout, bytes) else e.stdout}"
if e.stderr:
error_message += f"\nSTDERR:\n{e.stderr.decode() if isinstance(e.stderr, bytes) else e.stderr}"
error_logs.append(error_message)
# If we have errors and haven't reached max loops, try to fix them
if error_logs and current_loop < max_loop:
error_fix_prompt = (
f"<error_logs>{error_logs}</error_logs>\n"
f"<foam_file>{foam_file}</foam_file>\n"
f"<original_script>{viz_script}</original_script>\n"
f"<attempt_number>{current_loop}</attempt_number>\n"
f"Please fix the PyVista Python script based on the error messages. The script should read the .foam file ('{foam_file}') in the case directory."
)
fixed_script = state["llm_service"].invoke(error_fix_prompt, ERROR_FIX_SYSTEM_PROMPT)
# Save the fixed script
save_file(script_path, fixed_script)
# Try executing the fixed script
try:
result = subprocess.run(
[sys.executable, script_path],
cwd=case_dir,
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
print(f"Finished command: Return Code {result.returncode}")
error_logs = []
# Check if any PNG output image was created
png_files = glob.glob(os.path.join(case_dir, "*.png"))
if png_files:
output_image = png_files[0] # Use the first PNG file found
print(f"PyVista visualization created successfully: {output_image}")
# Create plot configs and outputs in the expected format
plot_configs = [
{
"plot_type": "pyvista_3d",
"field_name": "U", # Default field, could be enhanced to detect from script
"time_step": "latest",
"output_format": "png",
"output_path": output_image
}
]
plot_outputs = [output_image]
visualization_summary = {
"total_plots_generated": len(plot_outputs),
"plot_types": ["pyvista_3d"],
"fields_visualized": ["U"],
"output_directory": case_dir,
"pyvista_success": True
}
pyvista_result = {
"success": True,
"output_image": output_image,
"script": fixed_script
}
print(f"Generated {len(plot_outputs)} plots")
print(f"PyVista visualization saved to: {output_image}")
print("============================== Visualization Complete ==============================")
return {
**state,
"plot_configs": plot_configs,
"plot_outputs": plot_outputs,
"visualization_summary": visualization_summary,
"pyvista_visualization": pyvista_result
}
else:
error_logs.append("Fixed visualization script executed but no PNG output image was created")
except subprocess.CalledProcessError as e:
error_message = f"Error executing fixed visualization script: {str(e)}"
if e.stdout:
error_message += f"\nSTDOUT:\n{e.stdout.decode() if isinstance(e.stdout, bytes) else e.stdout}"
if e.stderr:
error_message += f"\nSTDERR:\n{e.stderr.decode() if isinstance(e.stderr, bytes) else e.stderr}"
error_logs.append(error_message)
# If we've exhausted all attempts
if current_loop >= max_loop:
print(f"Failed to create visualization after {max_loop} attempts")
error_message = f"Maximum number of attempts ({max_loop}) reached without success"
error_logs.append(error_message)
# Return failure state in the expected format
plot_configs = []
plot_outputs = []
visualization_summary = {
"total_plots_generated": 0,
"plot_types": [],
"fields_visualized": [],
"output_directory": case_dir,
"pyvista_success": False,
"error": error_message if 'error_message' in locals() else "Unknown error"
}
pyvista_result = {
"success": False,
"error": error_message if 'error_message' in locals() else "Unknown error",
"error_logs": error_logs
}
print("============================== Visualization Failed ==============================")
return {
**state,
"plot_configs": plot_configs,
"plot_outputs": plot_outputs,
"visualization_summary": visualization_summary,
"pyvista_visualization": pyvista_result
}
|