HABIT — Soil Water Retention Predictor

HABIT (Hierarchical Attention-Based Inference with Transfer learning) predicts soil water retention curves from basic soil properties, using a 20-member deep ensemble with per-prediction uncertainty.

Paper: Ghezzehei TA (2026). Water Resources Research, 62, e2025WR042833. doi:10.1029/2025WR042833

Web app: soil-habit.streamlit.app — same ensemble, no install required.

Features

• ✅ Easy to use: simple Python API, CSV interface, and CLI
• ✅ Ensemble predictions: mean, standard deviation, and 95% interval
• ✅ Adaptive inputs: each row uses whichever properties you supply
• ✅ Lightweight: ONNX Runtime inference, ~50 MB of weights, no TensorFlow

Installation

pip install habit-ptf

Or from source:

git clone https://github.com/Teamrat/habit.git
cd habit/habit-ptf
pip install -e .

Model weights

Weights are not bundled. On first use the 20 ONNX ensemble members (~50 MB total) are downloaded from huggingface.co/Teamrat/habit and cached in ~/.cache/habit-ptf/onnx. Subsequent runs load from the cache and need no network access.

To work fully offline, download the members yourself and point the loader at the directory:

predictor = load_ensemble(ensemble_dir='/path/to/onnx_weights')

Quick Start

Python API

from habit_ptf import load_ensemble

# Downloads the ensemble on first use, then loads from cache
predictor = load_ensemble()

# Predict from CSV
predictions = predictor.predict_from_csv(
    input_csv='my_soils.csv',
    output_csv='predictions.csv'
)

print(predictions.head())

By default the stage is inferred per row from the columns you supply — a single CSV may mix rows with and without bulk density, organic carbon, or Ksat. Pass stage=1 to cap every row at texture + bulk density regardless.

Command Line Interface

habit-predict --input my_soils.csv --output predictions.csv

Input Format

Your CSV file should contain the following columns:

Units are fixed and are NOT auto-detected or converted. Supply each column in exactly the units below; values in other units are either rejected or silently wrong.

Required Columns

• soil_id: Unique identifier for each soil sample
• sand: Sand, percent by mass
• silt: Silt, percent by mass
• clay: Clay, percent by mass

sand + silt + clay must sum to ~100 (tolerance ±10). Fractions summing to 1 raise a ValueError rather than being reinterpreted.

Optional Columns

• bd: Bulk density, g/cm³
• oc: Organic carbon, percent by mass — 1.2 means 1.2%, 0.8 means 0.8%
• ksat: Saturated hydraulic conductivity, cm/day

Example input CSV:

soil_id,sand,silt,clay,bd,oc,ksat
soil_001,40.0,35.0,25.0,1.35,2.0,15.2
soil_002,55.0,30.0,15.0,1.45,1.5,
soil_003,25.0,45.0,30.0,,,

Output Format

Long-form: one row per soil per water potential.

soil_id,stage,water_potential_kPa,water_content_mean,water_content_std,water_content_q025,water_content_q975
soil_001,Stage 3,33.0,0.3668,0.0315,0.3141,0.4001
soil_001,Stage 3,1500.0,0.1624,0.0316,0.1131,0.1910

stage records which properties were actually used for that row. water_content_std is the spread among the 20 independently trained members. It is not a calibrated uncertainty interval, but larger spread was empirically associated with larger prediction error on held-out data (Ghezzehei, 2026).

Columns:

• soil_id: Soil identifier (matches input)
• water_potential_kPa: Water potential in kPa
• water_content_mean: Mean volumetric water content (cm³/cm³) across ensemble
• water_content_std: Standard deviation (uncertainty estimate)

Default Water Potentials

By default, predictions are made at these water potentials (kPa):

0.01, 0.1, 1.0, 3.0, 10.0, 33.0, 100.0, 300.0, 1000.0, 15000.0

Custom Water Potentials

You can specify custom water potentials:

import numpy as np

# Custom water potentials (in kPa)
custom_wp = np.array([10, 33, 100, 1500, 15000])

predictions = predictor.predict(
    soil_data=my_dataframe,
    water_potentials=custom_wp
)

Training Stages

HABIT is trained hierarchically with different levels of input data:

Stage	Properties Available	Use Case
0	Texture only	Minimal data available
1	Texture + BD	Common field measurements
2	Texture + BD + OC	Enhanced predictions
3	Texture + BD + OC + Ksat	Maximum accuracy

A single ONNX model serves every stage — the stage is selected by a mask, not by loading different weights. By default it is inferred per row from the columns you supply, so one CSV may mix stages freely.

# Default: each row uses whatever properties it has
predictor = load_ensemble()

# Cap every row at texture + BD, ignoring any OC/Ksat columns present
predictor = load_ensemble(stage=1)

Advanced Usage

Programmatic Prediction

import pandas as pd
from habit_ptf import load_ensemble

# Create soil data
soils = pd.DataFrame({
    'soil_id': ['A', 'B', 'C'],
    'sand': [40.0, 50.0, 30.0],     # percent
    'silt': [35.0, 30.0, 40.0],     # percent
    'clay': [25.0, 20.0, 30.0],     # percent
    'bd':   [1.35, 1.45, 1.30],     # g/cm3
    'oc':   [2.0, 1.5, 2.5],        # percent
    'ksat': [15.2, 20.5, 10.8]      # cm/day
})

# Load predictor
predictor = load_ensemble('path/to/ensemble', stage=3)

# Predict
results = predictor.predict(soils)

# Filter specific water potential
field_capacity = results[results['water_potential_kPa'] == 33.0]
print(field_capacity)

Handling Missing Properties

The model gracefully handles missing properties:

# Some soils have all properties, some don't
soils = pd.DataFrame({
    'soil_id': ['complete', 'no_ksat', 'texture_only'],
    'sand': [40.0, 50.0, 30.0],       # percent
    'silt': [35.0, 30.0, 40.0],       # percent
    'clay': [25.0, 20.0, 30.0],       # percent
    'bd':   [1.35, 1.45, np.nan],     # g/cm3, missing BD
    'oc':   [2.0, 1.5, np.nan],       # percent, missing OC
    'ksat': [15.2, np.nan, np.nan]    # cm/day, missing Ksat
})

# Model automatically adapts to available data
predictions = predictor.predict(soils)

Uncertainty Quantification

The ensemble provides uncertainty estimates:

# Get predictions for a soil
soil_predictions = results[results['soil_id'] == 'soil_001']

# High uncertainty indicates:
# - Unusual soil property combinations
# - Extrapolation beyond training data
# - Model disagreement

# Filter high-uncertainty predictions
uncertain = results[results['water_content_std'] > 0.05]
print(f"Found {len(uncertain)} high-uncertainty predictions")

Model Details

HABIT uses a hierarchical attention-based architecture that:

• Captures interactions between soil properties (texture × BD, texture × OC)
• Uses multi-head attention to learn diverse patterns
• Enforces physical constraints (monotonic water retention curves)
• Provides ensemble uncertainty quantification

For more details, see Ghezzehei (2026) and the main HABIT repository.

Requirements

• Python >= 3.8
• onnxruntime >= 1.16
• NumPy >= 1.21.0
• Pandas >= 1.3.0
• huggingface_hub >= 0.20

No TensorFlow, no scikit-learn, no bundled weights.

License

MIT License - see LICENSE file for details

Citation

If you use HABIT in your research, please cite:

@article{Ghezzehei2026HABIT,
  title   = {Hierarchical Attention-Based Inference with Transfer Learning
             for Soil Water Retention Prediction},
  author  = {Ghezzehei, Teamrat A.},
  journal = {Water Resources Research},
  volume  = {62},
  pages   = {e2025WR042833},
  year    = {2026},
  doi     = {10.1029/2025WR042833}
}

Support

For questions or issues:

• Open an issue on GitHub
• Soil Physics Lab, UC Merced: soilphysics.ucmerced.edu
• Contact: taghezzehei@ucmerced.edu

Contributing

Contributions welcome! Please see CONTRIBUTING.md for guidelines.

Changelog

Version 1.0.0 (2026-08-19)

• Fixed organic carbon preprocessing. Earlier versions scaled OC without the training-time log transform, so every Stage 2 and Stage 3 prediction was wrong. The correct transform is log(1 + 10*OC%) / log(11) with OC in percent. Verified against the archived training tensors.
• Fixed Ksat preprocessing. Removed a heuristic that treated any Ksat column with max < 10 as already log-scaled.
• Fixed input units. Units are now fixed and validated rather than guessed from the data: texture and OC in percent, BD in g/cm³, Ksat in cm/day. Previous versions inspected column maxima to decide whether values were fractions or percentages, so a single outlying row could reinterpret an entire column. Out-of-spec texture now raises a clear error.
• ONNX Runtime inference. Replaced the bundled 2.9 GB of Keras weights with ~50 MB of ONNX members downloaded from HuggingFace and cached locally. Drops the TensorFlow and scikit-learn dependencies.
• Stage is now inferred per row by default; rows in one batch may differ.
• Added water_content_q025 / water_content_q975 and a stage column.
• Chunked inference keeps memory bounded on large batches.
• Added verify_preprocessing.py, which reproduces the archived training tensors for all four stages.

Version 0.1.0 (2025-10-31)

• Initial release
• Ensemble prediction support
• CSV input/output interface
• Uncertainty quantification