""" HABIT — Interactive Soil Water Retention Predictor HuggingFace Space (Gradio) Downloads ensemble weights from huggingface.co/Teamrat/habit on startup, then predicts water retention curves from user-supplied soil properties. """ import os os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" os.environ["KERAS_BACKEND"] = "tensorflow" import time import shutil import numpy as np import pandas as pd import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import gradio as gr from huggingface_hub import hf_hub_download import tensorflow as tf # Import model from co-located files (exact copies from HABIT-training) from habit_model import HABIT # ═══════════════════════════════════════════════════════════════════════════ # Configuration # ═══════════════════════════════════════════════════════════════════════════ HF_REPO_ID = "Teamrat/habit" NUM_MEMBERS = 20 MODEL_CONFIG = { "embedding_dim": 192, "num_heads": 4, "num_monotonic_basis": 40, "dropout_rate": 0.15, } SCALER_PARAMS = { "texture": {"center": [0.2712, 0.413, 0.172], "scale": [0.456, 0.4543, 0.183]}, "bd": {"center": [1.4], "scale": [0.31]}, "oc": {"center": [1.28], "scale": [1.9902]}, "ksat": {"center": [2.1206], "scale": [1.5133]}, } DEFAULT_WP_KPA = np.array( [1, 3, 6, 10, 33, 100, 300, 500, 1000, 5000, 10000, 15000], dtype=np.float64 ) # ═══════════════════════════════════════════════════════════════════════════ # Load ensemble on startup # ═══════════════════════════════════════════════════════════════════════════ def download_and_load_ensemble(): """Download weights from HF Hub and load all ensemble members.""" cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "habit-ptf", "weights") os.makedirs(cache_dir, exist_ok=True) models = [] dummy = [ np.zeros((1, 3), dtype=np.float32), np.zeros((1, 1), dtype=np.float32), np.zeros((1, 1), dtype=np.float32), np.zeros((1, 1), dtype=np.float32), np.ones((1, 4), dtype=np.float32), np.zeros((1, 12), dtype=np.float32), ] for i in range(1, NUM_MEMBERS + 1): name = f"member_{i:02d}.h5" print(f"Loading {name}...", flush=True) local_path = os.path.join(cache_dir, name) if not os.path.exists(local_path): downloaded = hf_hub_download( repo_id=HF_REPO_ID, filename=f"weights/{name}", ) # Copy — not symlink — Keras 3 h5 loader can't follow symlinks shutil.copy2(downloaded, local_path) model = HABIT(**MODEL_CONFIG) model(dummy, training=False) model.load_weights(local_path) models.append(model) print(f"Loaded {len(models)}-member ensemble.", flush=True) return models print("Starting HABIT ensemble loading...") t0 = time.time() ENSEMBLE = download_and_load_ensemble() print(f"Ensemble ready in {time.time() - t0:.1f}s") # ═══════════════════════════════════════════════════════════════════════════ # Prediction logic # ═══════════════════════════════════════════════════════════════════════════ def robust_scale(values, center, scale): return ((values - np.array(center)) / np.array(scale)).astype(np.float32) def predict_retention(sand, silt, clay, bd, oc, ksat, wp_min, wp_max, n_points): """Run ensemble prediction and return plot + table + CSV path.""" # Validate texture if sand is None or silt is None or clay is None: return None, None, "Sand, silt, and clay are required." sand_f, silt_f, clay_f = float(sand), float(silt), float(clay) if sand_f + silt_f + clay_f < 1: return None, None, "Texture fractions must sum to ~100% (or ~1.0)." # Normalise texture if sand_f + silt_f + clay_f > 5: # percentages sand_f, silt_f, clay_f = sand_f / 100, silt_f / 100, clay_f / 100 total = sand_f + silt_f + clay_f sand_f, silt_f, clay_f = sand_f / total, silt_f / total, clay_f / total texture_sc = robust_scale( np.array([[sand_f, silt_f, clay_f]]), SCALER_PARAMS["texture"]["center"], SCALER_PARAMS["texture"]["scale"], ) # Build mask and optional properties mask = np.zeros((1, 4), dtype=np.float32) mask[0, 0] = 1.0 # texture always if bd is not None and bd > 0: bd_sc = robust_scale( np.array([[float(bd)]]), SCALER_PARAMS["bd"]["center"], SCALER_PARAMS["bd"]["scale"], ) mask[0, 1] = 1.0 else: bd_sc = np.zeros((1, 1), dtype=np.float32) if oc is not None and oc > 0: oc_val = float(oc) if oc_val > 1.0: oc_val /= 100 oc_log = np.log1p(oc_val) oc_sc = robust_scale( np.array([[oc_log]]), SCALER_PARAMS["oc"]["center"], SCALER_PARAMS["oc"]["scale"], ) mask[0, 2] = 1.0 else: oc_sc = np.zeros((1, 1), dtype=np.float32) if ksat is not None and ksat > 0: ksat_log = np.log10(max(float(ksat), 1e-6)) ksat_sc = robust_scale( np.array([[ksat_log]]), SCALER_PARAMS["ksat"]["center"], SCALER_PARAMS["ksat"]["scale"], ) mask[0, 3] = 1.0 else: ksat_sc = np.zeros((1, 1), dtype=np.float32) # Stage label stage_names = { (1, 0, 0, 0): "Stage 0 — texture only", (1, 1, 0, 0): "Stage 1 — texture + BD", (1, 1, 1, 0): "Stage 2 — texture + BD + OC", (1, 1, 1, 1): "Stage 3 — all properties", } mask_key = tuple(int(m) for m in mask[0]) stage_label = stage_names.get(mask_key, f"Custom mask: {mask_key}") # Water potentials n_pts = int(n_points) if n_points else 50 wp_kpa = np.logspace( np.log10(max(float(wp_min), 0.1)), np.log10(float(wp_max)), n_pts ) wp_log = np.log10(wp_kpa).astype(np.float32).reshape(1, -1) # Run ensemble inputs = [texture_sc, bd_sc, oc_sc, ksat_sc, mask, wp_log] all_preds = [] for model in ENSEMBLE: pred = model(inputs, training=False).numpy() all_preds.append(pred[0]) all_preds = np.array(all_preds) # (members, n_wp) mean = np.mean(all_preds, axis=0) std = np.std(all_preds, axis=0) lower = np.percentile(all_preds, 2.5, axis=0) upper = np.percentile(all_preds, 97.5, axis=0) # ── Plot ────────────────────────────────────────────────────── fig, ax = plt.subplots(figsize=(8, 5)) ax.fill_between( wp_kpa, lower, upper, alpha=0.25, color="#2196F3", label="95% interval" ) ax.plot(wp_kpa, mean, color="#1565C0", linewidth=2, label="Ensemble mean") for m in range(len(all_preds)): ax.plot(wp_kpa, all_preds[m], color="#90CAF9", linewidth=0.4, alpha=0.6) ax.set_xscale("log") ax.set_xlabel("Water potential |ψ| (kPa)", fontsize=12) ax.set_ylabel("Volumetric water content θ (cm³/cm³)", fontsize=12) ax.set_title(f"HABIT Prediction — {stage_label}", fontsize=13) ax.legend(loc="upper right") ax.set_ylim(bottom=0) ax.grid(True, alpha=0.3) fig.tight_layout() # ── Table at standard tensions ──────────────────────────────── standard_kpa = [1, 3, 6, 10, 33, 100, 300, 500, 1000, 5000, 10000, 15000] standard_kpa = [p for p in standard_kpa if float(wp_min) <= p <= float(wp_max)] table_rows = [] for target_kpa in standard_kpa: idx = np.argmin(np.abs(wp_kpa - target_kpa)) table_rows.append( { "ψ (kPa)": int(target_kpa), "θ mean": f"{mean[idx]:.4f}", "θ std": f"{std[idx]:.4f}", "θ lower 95%": f"{lower[idx]:.4f}", "θ upper 95%": f"{upper[idx]:.4f}", } ) table_df = pd.DataFrame(table_rows) # ── CSV download ────────────────────────────────────────────── full_df = pd.DataFrame( { "water_potential_kPa": wp_kpa, "water_content_mean": mean, "water_content_std": std, "water_content_lower95": lower, "water_content_upper95": upper, } ) for m in range(len(all_preds)): full_df[f"member_{m + 1:02d}"] = all_preds[m] csv_path = "/tmp/habit_prediction.csv" full_df.to_csv(csv_path, index=False) return fig, table_df, csv_path # ═══════════════════════════════════════════════════════════════════════════ # Batch prediction from CSV # ═══════════════════════════════════════════════════════════════════════════ def predict_from_csv(file): """Run predictions for all soils in an uploaded CSV.""" if file is None: return None, None, "Please upload a CSV file." df = pd.read_csv(file.name if hasattr(file, "name") else file) cols_lower = {c.lower(): c for c in df.columns} if not all(k in cols_lower for k in ["sand", "silt", "clay"]): return None, None, f"CSV must have sand, silt, clay columns. Found: {list(df.columns)}" results_all = [] for idx, row in df.iterrows(): sand = row[cols_lower["sand"]] silt = row[cols_lower["silt"]] clay = row[cols_lower["clay"]] bd = row.get(cols_lower.get("bd")) if "bd" in cols_lower else None oc = row.get(cols_lower.get("oc")) if "oc" in cols_lower else None ksat = row.get(cols_lower.get("ksat")) if "ksat" in cols_lower else None soil_id = row.get(cols_lower.get("soil_id", ""), idx + 1) bd = None if bd is not None and (pd.isna(bd) or bd <= 0) else bd oc = None if oc is not None and (pd.isna(oc) or oc <= 0) else oc ksat = None if ksat is not None and (pd.isna(ksat) or ksat <= 0) else ksat fig, table, csv_path = predict_retention( sand, silt, clay, bd, oc, ksat, 1, 15000, 50 ) plt.close(fig) pred = pd.read_csv(csv_path) pred.insert(0, "soil_id", soil_id) results_all.append(pred) combined = pd.concat(results_all, ignore_index=True) out_path = "/tmp/habit_batch_predictions.csv" combined.to_csv(out_path, index=False) summary = ( combined.groupby("soil_id") .agg( n_points=("water_content_mean", "count"), theta_sat=("water_content_mean", "max"), theta_15000=("water_content_mean", "min"), ) .reset_index() ) return summary, out_path, f"Predicted {len(df)} soils successfully." # ═══════════════════════════════════════════════════════════════════════════ # Gradio interface # ═══════════════════════════════════════════════════════════════════════════ EXAMPLE_SOILS = { "Clay (heavy)": {"sand": 10, "silt": 30, "clay": 60, "bd": 1.2, "oc": 2.0, "ksat": 5}, "Sandy loam": {"sand": 65, "silt": 25, "clay": 10, "bd": 1.5, "oc": 0.5, "ksat": 200}, "Silt loam": {"sand": 15, "silt": 65, "clay": 20, "bd": 1.3, "oc": 1.5, "ksat": 25}, "Loam (average)": {"sand": 40, "silt": 40, "clay": 20, "bd": 1.35, "oc": 1.2, "ksat": 50}, "Sand (texture only)": {"sand": 90, "silt": 5, "clay": 5, "bd": None, "oc": None, "ksat": None}, } def load_preset(name): if name and name in EXAMPLE_SOILS: s = EXAMPLE_SOILS[name] return ( s["sand"], s["silt"], s["clay"], s["bd"] if s["bd"] else 0, s["oc"] if s["oc"] else 0, s["ksat"] if s["ksat"] else 0, ) return gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update() with gr.Blocks( title="HABIT — Soil Water Retention Predictor", theme=gr.themes.Soft(), ) as demo: gr.Markdown( """ # HABIT — Soil Water Retention Predictor **Hierarchical Attention-Based Inference with Transfer Learning** Predict soil water retention curves from basic soil properties using a 20-member deep learning ensemble. Provide whatever properties you have — the model adapts automatically. *Ghezzehei TA (2025). Water Resources Research.*  |  [Model weights](https://huggingface.co/Teamrat/habit)  |  [pip install habit-ptf](https://pypi.org/project/habit-ptf/) """ ) with gr.Tabs(): # ── Tab 1: Single soil ──────────────────────────────────── with gr.TabItem("Single Soil"): with gr.Row(): with gr.Column(scale=1): gr.Markdown("### Soil Properties") preset = gr.Dropdown( choices=list(EXAMPLE_SOILS.keys()), label="Load example soil", interactive=True, ) gr.Markdown("**Texture** (required — % or fraction)") with gr.Row(): sand_in = gr.Number(label="Sand", value=40, precision=1) silt_in = gr.Number(label="Silt", value=40, precision=1) clay_in = gr.Number(label="Clay", value=20, precision=1) gr.Markdown("**Optional properties** (leave at 0 to omit)") bd_in = gr.Number( label="Bulk density (g/cm³)", value=1.35, precision=2 ) oc_in = gr.Number( label="Organic carbon (%)", value=1.2, precision=2 ) ksat_in = gr.Number( label="Ksat (cm/day)", value=50, precision=1 ) gr.Markdown("**Water potential range**") with gr.Row(): wp_min_in = gr.Number(label="Min (kPa)", value=1, precision=0) wp_max_in = gr.Number( label="Max (kPa)", value=15000, precision=0 ) n_pts_in = gr.Number(label="Points", value=50, precision=0) predict_btn = gr.Button("Predict", variant="primary", size="lg") with gr.Column(scale=2): plot_out = gr.Plot(label="Water Retention Curve") table_out = gr.Dataframe(label="Predictions at Standard Tensions") csv_out = gr.File(label="Download Full Results (CSV)") preset.change( fn=load_preset, inputs=[preset], outputs=[sand_in, silt_in, clay_in, bd_in, oc_in, ksat_in], ) predict_btn.click( fn=predict_retention, inputs=[ sand_in, silt_in, clay_in, bd_in, oc_in, ksat_in, wp_min_in, wp_max_in, n_pts_in, ], outputs=[plot_out, table_out, csv_out], ) # ── Tab 2: Batch from CSV ───────────────────────────────── with gr.TabItem("Batch (CSV Upload)"): gr.Markdown( """ ### Batch Prediction Upload a CSV with columns: `sand`, `silt`, `clay` (required), plus optional `bd`, `oc`, `ksat`, `soil_id`. Values can be percentages (0–100) or fractions (0–1). Missing optional properties should be blank or 0. """ ) csv_upload = gr.File(label="Upload CSV", file_types=[".csv"]) batch_btn = gr.Button("Predict All", variant="primary") batch_status = gr.Textbox(label="Status") batch_summary = gr.Dataframe(label="Summary") batch_download = gr.File(label="Download Results") batch_btn.click( fn=predict_from_csv, inputs=[csv_upload], outputs=[batch_summary, batch_download, batch_status], ) # ── Tab 3: About ────────────────────────────────────────── with gr.TabItem("About"): gr.Markdown( """ ### About HABIT HABIT is a deep learning model for predicting soil water retention curves from basic soil properties. It uses a transformer-based architecture with: - **Property-specific encoders** for each soil property - **Cross-attention layers** that learn interactions between properties - **Monotonic output layer** ensuring physically correct behavior (water content decreases with increasing tension) - **Hierarchical training** so one model handles any combination of inputs #### Performance (test set, 95% CI from cluster bootstrap) | Inputs | R² | RMSE (cm³/cm³) | |---|---|---| | Texture only | 0.78 [0.74, 0.82] | 0.067 | | + Bulk density | 0.85 [0.75, 0.91] | 0.056 | | + Organic carbon | 0.86 [0.78, 0.92] | 0.052 | | + Ksat | 0.92 [0.90, 0.94] | 0.043 | #### Python package ```bash pip install habit-ptf ``` ```python from habit_ptf import load_ensemble predictor = load_ensemble() result = predictor.predict(soil_dataframe) ``` #### Citation Ghezzehei TA (2025). Interpretable Soil Water Retention Prediction Using Hierarchical Attention Networks with Uncertainty Quantification. *Water Resources Research*. #### License MIT (code and weights). Training data: CC BY 4.0. """ ) if __name__ == "__main__": demo.launch()