Metadata-Version: 2.4
Name: wet-net
Version: 0.4.0b1
Summary: WetNet helps Aigües de Barcelona (Barcelona Water Company) predict anomalous water consumption, using machine learning.
Author: Antonio Lobo Santos, Zachary Parent, Bruno Sánchez Gómez
Author-email: Antonio Lobo Santos <alobosantos10@gmail.com>, Zachary Parent <zachparent@duck.com>, Bruno Sánchez Gómez <brunosangom@gmail.com>
License-Expression: MIT
License-File: LICENSE
Classifier: Programming Language :: Python :: 3
Classifier: Operating System :: OS Independent
Requires-Dist: huggingface-hub>=1.2.1
Requires-Dist: matplotlib>=3.10.7
Requires-Dist: numpy>=2.3.3
Requires-Dist: pandas>=2.3.2
Requires-Dist: plotly>=6.4.0
Requires-Dist: pyarrow>=22.0.0
Requires-Dist: requests>=2.32.5
Requires-Dist: safetensors>=0.7.0
Requires-Dist: torch>=2.9.1
Requires-Dist: torchvision>=0.24.1
Requires-Dist: typer>=0.19.2
Requires-Dist: umap-learn>=0.5.7
Requires-Dist: rich>=13.9.2
Requires-Dist: pyyaml>=6.0.2
Requires-Python: >=3.11
Project-URL: Homepage, https://github.com/ZachParent/wet-net
Description-Content-Type: text/markdown

# WetNet

WetNet helps Aigües de Barcelona (Barcelona Water Company) predict anomalous water consumption, using machine learning.

![PyPI - Version](https://img.shields.io/pypi/v/wet-net)
![PyPI - License](https://img.shields.io/pypi/l/wet-net)
![PyPI - Python Version](https://img.shields.io/pypi/pyversions/wet-net)

![WetNet logo](https://github.com/ZachParent/wet-net/raw/main/assets/wet-net-logo.jpeg)

## Setup

### 1. Install uv

Our project uses uv to manage dependencies in a reproducible way. See [Installing uv](https://docs.astral.sh/uv/getting-started/installation/) documentation for installation instructions.

> [!TIP]
> You can skip the rest of the setup if you just want to run the scripts and see the project in action! Run `uvx wet-net --help` to see the available commands in the CLI. This will install the latest release of the project from PyPI in an isolated virtual environment.

### 2. Clone the repository

```bash
git clone https://github.com/ZachParent/wet-net.git
cd wet-net
```

### 3. Install the project dependencies

```bash
uv sync --locked
```

> Torch install hint (if not pulled via lockfile):
> - CPU-only: `pip install torch torchvision torchaudio`
> - CUDA 12.8 (recommended for GPUs): `pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128`

### HuggingFace token (only needed if you plan to push trained models)
- Copy `.env.example` to `.env` and set `HUGGINGFACE_TOKEN=<your_token>`. If you already placed your token in `.env`, nothing else to do.
- One-off export: `export HUGGINGFACE_TOKEN=hf_xxx`.
- CLI alternative: `huggingface-cli login` (same token).

End-to-end train (full data) + push example:
```bash
# assuming data is preprocessed and CUDA available
uv run wet-net train --seq-len 96 --optimize-for recall
uv run python - <<'PY'
import os, torch
from huggingface_hub import HfApi, HfFolder
from wet_net.paths import RESULTS_DIR
run_dir = RESULTS_DIR / "wetnet" / "seq96_recall"
repo = os.environ.get("HF_REPO", "your-username/wetnet-seq96-recall")
api = HfApi()
api.create_repo(repo, exist_ok=True)
api.upload_file(path_or_fileobj=run_dir / "wetnet.pt", path_in_repo="wetnet.pt", repo_id=repo)
api.upload_file(path_or_fileobj=run_dir / "vib.pt", path_in_repo="vib.pt", repo_id=repo)
api.upload_file(path_or_fileobj=run_dir / "config.json", path_in_repo="config.json", repo_id=repo)
PY
```
Make sure `HUGGINGFACE_TOKEN` is set in your environment before running the upload step.

## Usage

### Run the scripts

The CLI entry points are unchanged, now targeting the WetNet tri-task pipeline (no grid search; cached best configs per sequence length):

```bash
# Generate data (real: pass your private URL; mock: bundled synthetic)
# Accepts a parquet URL or the official challenge zip URL
uv run wet-net pre-process --data-url https://<your-parquet-or-zip-url>   # real
uv run wet-net pre-process --mock                                         # synthetic

# Train (picks the best recall/false_alarm config for that seq_len)
uv run wet-net train --seq-len 192 --optimize-for recall
uv run wet-net train --seq-len 192 --optimize-for false_alarm --mock

# Optional: push trained artifacts to Hugging Face
uv run wet-net train --seq-len 96 --optimize-for recall --push-to-hub --hub-model-name WetNet/wet-net

# Optional: override hyperparameters by editing src/wet_net/config/best_configs.yaml
# (any fields you change there are picked up automatically at runtime)

# Evaluate saved checkpoints (no retraining)
uv run wet-net evaluate --seq-len 192 --optimize-for recall
```

The code for the scripts can be found in the [src/wet_net/scripts](src/wet_net/scripts) directory.

### Tri-Task WetNet

- Model: Cyclic-UniTS backbone + 4 heads (reconstruction, 24h forecast, short/long anomaly logits), class `wet_net.models.wetnet.WetNet`.
- Cached best configs: one per sequence length (96, 192, 360, 720, 1440) for two objectives (recall, false_alarm). Each stores schedule variant, PCGrad flag, transformer width/depth/heads/dropout, and the fusion threshold to report.
- VIB fusion: lightweight dual-VIB probe is trained after the main model to provide uncertainty mass for Dempster–Shafer fusion.
- Data: preprocessing lives in `wet_net.data.preprocess`. For real runs supply the parquet URL via `--data-url` or `WETNET_DATA_URL`. A mock parquet is shipped so the full pipeline can run without private data.

### Notebooks

We use Jupyter notebooks to show the process of preparing the data, training the model, and evaluating the model, with descriptions and code. You can run the notebooks by opening them up in your favorite IDE. Be sure to choose the `.venv` kernel which is created and managed by uv. The notebooks can be found in the [notebooks](notebooks) directory.

### 01_pre_process.ipynb

The [01_pre_process.ipynb](notebooks/01_pre_process.ipynb) notebook shows the process of preparing the data. It includes:

- Loading the data
- Pre-processing the data
- Saving the pre-processed data

### 02_train.ipynb

The [02_train.ipynb](notebooks/02_train.ipynb) notebook shows the process of training the model. It includes:

- Loading the pre-processed data
- Training the model
- Saving the trained model

### 03_evaluate.ipynb

The [03_evaluate.ipynb](notebooks/03_evaluate.ipynb) notebook shows the process of evaluating the model. It includes:

- Loading the trained model
- Evaluating the model
- Saving the evaluation results

## Contributing

We welcome contributions to the project. Please feel free to submit an issue or pull request.

For more information on contributing, including how to set up pre-commit hooks and how to cut a new release, see [CONTRIBUTING.md](CONTRIBUTING.md).
