Metadata-Version: 2.4
Name: habit-ptf
Version: 1.0.0
Summary: HABIT: predict soil water retention curves from basic soil properties
Home-page: https://github.com/Teamrat/habit
Author: Teamrat A. Ghezzehei
Author-email: taghezzehei@ucmerced.edu
License: MIT
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: onnxruntime>=1.16
Requires-Dist: numpy>=1.21.0
Requires-Dist: pandas>=1.3.0
Requires-Dist: huggingface_hub>=0.20
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: license
Dynamic: license-file
Dynamic: provides-extra
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# 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](https://doi.org/10.1029/2025WR042833)

**Web app:** [soil-habit.streamlit.app](https://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

```bash
pip install habit-ptf
```

Or from source:

```bash
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](https://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:

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

## Quick Start

### Python API

```python
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

```bash
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:**
```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.

```csv
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](https://doi.org/10.1029/2025WR042833)).


**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:

```python
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.

```python
# 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

```python
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:

```python
# 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:

```python
# 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)](https://doi.org/10.1029/2025WR042833) and the [main HABIT repository](https://github.com/Teamrat/habit).

## 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:

```bibtex
@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](https://github.com/Teamrat/habit/issues)
- Soil Physics Lab, UC Merced: [soilphysics.ucmerced.edu](https://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
