Update app.py

app.py

@@ -7,41 +7,87 @@ from hardy_cross_plot import plot_hardy_cross_network
 import os
 import numpy as np
 
+
 def hardy_cross_tab(file, loops_json, max_iterations, tolerance):
+    """
+    Runs the Hardy Cross solver tab:
+    - Validates inputs
+    - Computes flows & iterations
+    - Saves a PNG plot (if possible) and a CSV ready for head-loss tab
+    - Returns Markdown tables + image path (or None) + CSV path (or None)
+    """
     try:
+        # Basic presence checks
+        if file is None:
+            return "❌ **Error:** No CSV uploaded.", "", None, None
+
+        try:
+            loops = json.loads(loops_json)
+        except json.JSONDecodeError as e:
+            return f"❌ **Invalid Loops JSON:** {e}", "", None, None
+
+        # Read CSV
         df = pd.read_csv(file.name)
+
+        # Validate loop pipe IDs exist in CSV
+        if "pipe_id" not in df.columns:
+            return "❌ **Error:** CSV must include a 'pipe_id' column.", "", None, None
+
+        csv_ids = set(df["pipe_id"].astype(str))
+        missing = {}
+        for loop_id, ids in loops.items():
+            ids = [str(x) for x in ids]
+            miss = [p for p in ids if p not in csv_ids]
+            if miss:
+                missing[loop_id] = miss
+        if missing:
+            msg = "❌ **Loops reference unknown pipes**:\n" + "\n".join(
+                f"- {lid}: {', '.join(miss)}" for lid, miss in missing.items()
+            )
+            return msg, "", None, None
+
+        # Prepare data for solver
         pipes_data = df.to_dict(orient='records')
-        loops = json.loads(loops_json)
+
+        # Run solver
         results, final_flows = hardy_cross_solver(
             pipes_data,
             loops,
             max_iterations=int(max_iterations),
             tolerance=float(tolerance)
         )
+
+        # Plot and verify
         plot_filename = "hardy_cross_network_result.png"
         plot_path = plot_hardy_cross_network(final_flows, save_path=plot_filename)
+        if not (isinstance(plot_path, str) and os.path.isfile(plot_path) and plot_path.lower().endswith(".png")):
+            plot_path = None  # critical: don't return "", use None for gr.Image
 
-        # Create a downloadable CSV for head loss calculation
+        # Build CSV for head-loss tab
         auto_hl_csv = "final_flows_for_head_loss.csv"
-        # Merge original parameters with computed flows
-        merge_df = df[["pipe_id", "length", "diameter", "friction_factor"]].copy()
-        merge_df = merge_df.drop_duplicates(subset="pipe_id")  # remove potential duplicates
+        needed_cols = ["pipe_id", "length", "diameter", "friction_factor"]
+        missing_cols = [c for c in needed_cols if c not in df.columns]
+        if missing_cols:
+            # If geometry columns are missing, we still return flows, but skip head-loss CSV
+            flows_md = final_flows.to_markdown(index=False)
+            iters_md = pd.DataFrame(results).to_markdown(index=False)
+            warn = (
+                f"⚠️ **Note:** Could not create Head Loss CSV. Missing columns in input CSV: "
+                f"{', '.join(missing_cols)}. Include these to enable head-loss download."
+            )
+            flows_md = warn + "\n\n" + flows_md
+            return flows_md, iters_md, plot_path, None
+
+        merge_df = df[needed_cols].copy().drop_duplicates(subset="pipe_id")
         merge_df = pd.merge(
             merge_df,
             final_flows.rename(columns={"Q": "flow_m3s"}),
             on="pipe_id",
             how="left"
-        )
-        # Rename columns for head loss calc
-        merge_df = merge_df.rename(columns={
-            "length": "length_m",
-            "diameter": "diameter_m"
-        })
+        ).rename(columns={"length": "length_m", "diameter": "diameter_m"})
         merge_df.to_csv(auto_hl_csv, index=False)
-
-        # Safety check
-        if not (os.path.isfile(plot_path) and plot_path.endswith(".png")):
-            plot_path = ""
+        if not os.path.isfile(auto_hl_csv):
+            auto_hl_csv = None  # critical: gr.File expects a real file path or None
 
         return (
             final_flows.to_markdown(index=False),
@@ -49,44 +95,77 @@ def hardy_cross_tab(file, loops_json, max_iterations, tolerance):
             plot_path,
             auto_hl_csv
         )
+
     except Exception as e:
         tb = traceback.format_exc()
         error_message = f"❌ **Error:** {str(e)}\n\n```\n{tb}\n```"
-        return error_message, "", "", None
+        # critical: return None for image/file slots on error
+        return error_message, "", None, None
+
 
 def head_loss_tab(file):
+    """
+    Head loss calculator tab:
+    - Expects CSV with: pipe_id, length_m, diameter_m, friction_factor, flow_m3s
+    - Computes velocity and head loss (Darcy–Weisbach) and returns Markdown + downloadable CSV
+    """
     try:
+        if file is None:
+            return "❌ **Error:** No CSV uploaded.", None
+
         df = pd.read_csv(file.name)
+
+        required = ["pipe_id", "length_m", "diameter_m", "friction_factor", "flow_m3s"]
+        missing_cols = [c for c in required if c not in df.columns]
+        if missing_cols:
+            return (
+                f"❌ **Error:** Head loss CSV is missing required columns: {', '.join(missing_cols)}",
+                None
+            )
+
         g = 9.81
         df["area_m2"] = np.pi * (df["diameter_m"] / 2) ** 2
+        # avoid division by zero
+        df["area_m2"] = df["area_m2"].replace({0.0: np.nan})
         df["velocity_ms"] = np.abs(df["flow_m3s"]) / df["area_m2"]
         df["head_loss_m"] = (
-            df["friction_factor"]
-            * df["length_m"]
-            / df["diameter_m"]
-            * df["velocity_ms"] ** 2
-            / (2 * g)
+            df["friction_factor"] * df["length_m"] / df["diameter_m"] *
+            df["velocity_ms"] ** 2 / (2 * g)
         )
+
         output_csv = "head_loss_results.csv"
         df.to_csv(output_csv, index=False)
+        if not os.path.isfile(output_csv):
+            output_csv = None  # critical: gr.File expects path or None
+
         return (
             df[["pipe_id", "velocity_ms", "head_loss_m"]].to_markdown(index=False),
             output_csv
         )
+
     except Exception as e:
         tb = traceback.format_exc()
         error_message = f"❌ **Error:** {str(e)}\n\n```\n{tb}\n```"
         return error_message, None
 
+
 with gr.Blocks(title="Hardy Cross Method and Head Loss Calculator") as demo:
     gr.Markdown("# Hardy Cross Method Solver and Head Loss Calculator")
+
     with gr.Tab("Hardy Cross Solver"):
-        gr.Markdown("Upload your pipe network CSV and loops JSON to compute flows and visualize using the Hardy Cross Method. A CSV ready for head loss calculation will also be generated for you.")
+        gr.Markdown(
+            "Upload your pipe network CSV and loops JSON to compute flows and visualize using the Hardy Cross Method. "
+            "A CSV ready for head loss calculation will also be generated for you (if geometry columns are present)."
+        )
         file_input = gr.File(label="Upload Pipe Network CSV")
-        loops_json_input = gr.Textbox(label="Loops JSON", placeholder='{"L1": ["AB", "BC", "CD"], "L2": ["DE", "EF", "BC"]}')
+        loops_json_input = gr.Textbox(
+            label="Loops JSON",
+            placeholder='{"L1": ["AB", "BC", "CD"], "L2": ["DE", "EF", "BC"]}'
+        )
         max_iter_input = gr.Number(label="Max Iterations", value=10)
         tol_input = gr.Number(label="Tolerance", value=0.001)
         run_button = gr.Button("Run Hardy Cross")
+
         flows_output = gr.Markdown(label="Final Pipe Flows")
         iterations_output = gr.Markdown(label="Iteration Details")
         plot_output = gr.Image(label="Flow Direction Plot")
@@ -99,7 +178,10 @@ with gr.Blocks(title="Hardy Cross Method and Head Loss Calculator") as demo:
         )
 
     with gr.Tab("Head Loss Calculator"):
-        gr.Markdown("Upload your final flows CSV with columns: `pipe_id`, `length_m`, `diameter_m`, `friction_factor`, `flow_m3s` to compute head losses.")
+        gr.Markdown(
+            "Upload your final flows CSV with columns: "
+            "`pipe_id`, `length_m`, `diameter_m`, `friction_factor`, `flow_m3s` to compute head losses."
+        )
         file_input_hl = gr.File(label="Upload Final Flows CSV")
         run_button_hl = gr.Button("Compute Head Losses")
         head_loss_output = gr.Markdown(label="Head Loss Results")
@@ -111,5 +193,8 @@ with gr.Blocks(title="Hardy Cross Method and Head Loss Calculator") as demo:
             outputs=[head_loss_output, download_output]
         )
 
+
 if __name__ == "__main__":
-    demo.launch()
+    # If you see intermittent SSR caching issues on Spaces, you can set ssr_mode=False:
+    # demo.launch(ssr_mode=False)
+    demo.launch()