modules/visualization.py

import matplotlib.pyplot as plt
import io
from PIL import Image

def plot_simulation_results(data):
    try:
        stress = data.get("stress", [])
        strain = data.get("strain", [])
        
        plt.figure(figsize=(8, 6))
        plt.plot(stress, strain, marker="o", label="Stress-Strain Curve")
        plt.title("Simulation Results")
        plt.xlabel("Stress (MPa)")
        plt.ylabel("Strain")
        plt.legend()
        plt.grid(True)
        
        buf = io.BytesIO()
        plt.savefig(buf, format="png")
        buf.seek(0)
        plt.close()
        
        return Image.open(buf)
    except Exception as e:
        return f"Error in visualization: {str(e)}"