| from fastapi import FastAPI, File, UploadFile, HTTPException |
| from fastapi.responses import JSONResponse |
| from typing import List, Optional |
| import uvicorn |
| import tempfile |
| import os |
| import wntr |
| import matplotlib.pyplot as plt |
| import pandas as pd |
|
|
| app = FastAPI(title="EPANET Simulation API") |
|
|
| @app.post("/simulate") |
| async def simulate_epanet( |
| inp_file: UploadFile = File(...), |
| simulation_duration: Optional[int] = 24, |
| output_results: List[str] = ["pressure", "flow"], |
| nodes_to_report: Optional[List[str]] = None, |
| links_to_report: Optional[List[str]] = None, |
| generate_plot: Optional[bool] = True |
| ): |
| try: |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".inp") as temp_file: |
| temp_file.write(await inp_file.read()) |
| inp_path = temp_file.name |
|
|
| wn = wntr.network.WaterNetworkModel(inp_path) |
| sim = wntr.sim.WNTRSimulator(wn) |
| results = sim.run_sim() |
|
|
| response_data = {} |
|
|
| if "pressure" in output_results: |
| pressures = results.node['pressure'] |
| if nodes_to_report: |
| pressures = pressures[nodes_to_report] |
| response_data["pressure"] = pressures.describe().to_dict() |
|
|
| if "flow" in output_results: |
| flows = results.link['flowrate'] |
| if links_to_report: |
| flows = flows[links_to_report] |
| response_data["flow"] = flows.describe().to_dict() |
|
|
| if generate_plot: |
| plot_dir = tempfile.mkdtemp() |
| plot_paths = [] |
| if "pressure" in output_results and nodes_to_report: |
| for node in nodes_to_report: |
| plt.figure() |
| results.node['pressure'][node].plot(title=f"Pressure at {node}") |
| img_path = os.path.join(plot_dir, f"pressure_{node}.png") |
| plt.savefig(img_path) |
| plot_paths.append(img_path) |
| plt.close() |
| if "flow" in output_results and links_to_report: |
| for link in links_to_report: |
| plt.figure() |
| results.link['flowrate'][link].plot(title=f"Flow in {link}") |
| img_path = os.path.join(plot_dir, f"flow_{link}.png") |
| plt.savefig(img_path) |
| plot_paths.append(img_path) |
| plt.close() |
| response_data["plots"] = plot_paths |
|
|
| return JSONResponse(content=response_data) |
|
|
| except Exception as e: |
| raise HTTPException(status_code=500, detail=str(e)) |
| finally: |
| if os.path.exists(inp_path): |
| os.remove(inp_path) |
