Update visualization/visualization.py

visualization/visualization.py

@@ -1,45 +1,15 @@
-# visualization/visualization.py
+import pyvista as pv
 import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D
 
-def visualize_results(solver_name, length, width, thickness, stress, deformation):
-    """
-    Visualize the simulation results in both 2D and 3D.
-
-    Parameters:
-        solver_name (str): Name of the solver.
-        length (float): Length of the plate/beam.
-        width (float): Width of the plate/beam.
-        thickness (float): Thickness of the plate/beam.
-        stress (float): Stress value from the simulation.
-        deformation (float): Deformation value from the simulation.
-
-    Returns:
-        tuple: Paths to the 2D and 3D visualization images.
-    """
-    # 2D visualization (stress distribution over the length)
+def visualize_results(simulator, length, width, thickness, stress, deformation):
     fig, ax = plt.subplots()
-    ax.plot([0, length], [0, stress], label="Stress Distribution")
-    ax.set_title("2D Stress Distribution")
-    ax.set_xlabel("Length (mm)")
-    ax.set_ylabel("Stress (MPa)")
-    ax.legend()
-
-    graph_path = "/tmp/2d_stress_distribution.png"
-    plt.savefig(graph_path)
-    plt.close()
-
-    # 3D visualization (stress and deformation)
-    fig = plt.figure()
-    ax = fig.add_subplot(111, projection='3d')
-    ax.scatter(length, width, stress, label="Stress Distribution")
-    ax.set_title("3D Stress Visualization")
-    ax.set_xlabel("Length (mm)")
-    ax.set_ylabel("Width (mm)")
-    ax.set_zlabel("Stress (MPa)")
-
-    three_d_path = "/tmp/3d_stress_visualization.png"
-    plt.savefig(three_d_path)
-    plt.close()
-
-    return graph_path, three_d_path
+    ax.bar(["Stress", "Deformation"], [stress, deformation], color=["red", "blue"])
+    ax.set_title(f"Results ({simulator})")
+    plt.savefig("results_2d.png")
+    plt.close(fig)
+
+    mesh = pv.Box(bounds=(0, length, 0, width, 0, thickness))
+    plotter = pv.Plotter(off_screen=True)
+    plotter.add_mesh(mesh, color="white", show_edges=True)
+    plotter.screenshot("results_3d.png")
+    return "results_2d.png", "results_3d.png"