refactor library

.github/workflows/ci.yml

@@ -0,0 +1,38 @@
+name: CI
+
+on:
+  push:
+    branches: [main, refactor]
+  pull_request:
+    branches: [main, refactor]
+
+permissions:
+  contents: read
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v8.1.0
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - run: uv sync --locked --all-extras --group dev
+      - run: uv run ruff check . --exclude "*.ipynb"
+      - run: uv run ruff format --check . --exclude "*.ipynb"
+      - run: uv run ty check claymodel tests --exclude "*.ipynb"
+
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v8.1.0
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - run: uv sync --locked --all-extras --group dev
+      - run: uv run pytest tests/ -v --cov=claymodel --cov-report=term-missing --cov-fail-under=90

.pre-commit-config.yaml

@@ -1,26 +1,39 @@
-# See https://pre-commit.com for more information
-# See https://pre-commit.com/hooks.html for more hooks
 repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v6.0.0
-  hooks:
-    - id: check-added-large-files
-      args: [ '--maxkb=512', '--enforce-all' ]
-      exclude: '^docs/tutorials/.*\.ipynb$'
-    - id: check-yaml
-    - id: end-of-file-fixer
-    - id: trailing-whitespace
-- repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.13.3
-  hooks:
-    - id: ruff  # Run the linter
-      args: [ --fix ]
-      types_or: [ python, pyi ]
-    - id: ruff  # Run the linter for Jupyter notebooks with the PLR0913 rule ignored
-      args: [ --fix, --ignore=PLR0913 ]
-      types: [ jupyter ]
-    - id: ruff-format  # Run the formatter
-      types_or: [ python, pyi, jupyter ]
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: check-added-large-files
+        args: ["--maxkb=512", "--enforce-all"]
+        exclude: "^(docs/tutorials/.*\\.ipynb|uv\\.lock)$"
+      - id: check-yaml
+      - id: end-of-file-fixer
+        exclude: "\\.ipynb$"
+      - id: trailing-whitespace
+        exclude: "\\.ipynb$"
+      - id: check-merge-conflict
+      - id: debug-statements
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.11
+    hooks:
+      - id: ruff
+        args: ["--fix", "--show-fixes"]
+        types_or: [python, pyi]
+      - id: ruff-format
+        types_or: [python, pyi]
+
+  - repo: local
+    hooks:
+      - id: ty-check
+        name: ty check
+        language: system
+        entry: uv run ty check claymodel tests --exclude "*.ipynb"
+        pass_filenames: false
+      - id: uv-lock
+        name: Lock dependencies with uv
+        language: system
+        entry: uv lock
+        pass_filenames: false
 
 # https://pre-commit.ci/#configuration
 ci:

CONTRIBUTING.md

@@ -0,0 +1,135 @@
+# Contributing to Clay Foundation Model
+
+Thank you for your interest in contributing to Clay! This guide covers how to set up a development environment, run tests, and submit changes.
+
+## Development Setup
+
+### Prerequisites
+
+- Python 3.11 or later
+- Git
+- (Optional) CUDA-capable GPU for model training/inference
+
+### Installation
+
+```bash
+# Clone the repository
+git clone https://github.com/Clay-foundation/model.git
+cd model
+
+# Install in development mode with all extras
+uv pip install -e ".[dev]"
+```
+
+### Verify Installation
+
+```bash
+# Run the test suite
+uv run pytest tests/ -v
+
+# Check linting
+uv run ruff check claymodel/ tests/
+
+# Check formatting
+uv run ruff format --check claymodel/ tests/
+```
+
+## Project Structure
+
+```
+claymodel/
+    __init__.py          # Package exports
+    api.py               # High-level API: embed(), load_model(), normalize()
+    cli.py               # `clay` commands: info, benchmark
+    model.py             # Core model: Encoder, Decoder, ClayMAE, factory functions
+    module.py            # Lightning module: ClayMAEModule
+    utils.py             # Utilities: position embeddings, weight loading
+    metadata.py          # PlatformMetadata Pydantic model, YAML loading
+    configs/
+        metadata.yaml    # Bundled sensor metadata (wavelengths, normalization stats)
+    inference/
+        deterministic.py # DeterministicInference context manager
+        elle.py          # ELLE quality scoring probe
+        masking.py       # PatchAnalyzer for chip quality filtering
+training/                # Training data loading, callbacks
+finetune/                # Downstream task examples
+tests/                   # Test suite
+docs/                    # Documentation (Jupyter Book)
+configs/                 # Training configs
+```
+
+## Running Tests
+
+```bash
+# Run all tests
+pytest tests/ -v
+
+# Run a specific test file
+pytest tests/test_model.py -v
+
+# Run with coverage
+pytest tests/ --cov=claymodel --cov-report=term-missing
+```
+
+Tests use a tiny encoder (dim=192, random weights) so they run in seconds without a GPU or checkpoint.
+
+## Code Style
+
+We use [ruff](https://docs.astral.sh/ruff/) for linting and formatting.
+
+```bash
+# Check for lint errors
+uv run ruff check claymodel/ tests/
+
+# Auto-fix fixable errors
+uv run ruff check claymodel/ tests/ --fix
+
+# Check formatting
+uv run ruff format --check claymodel/ tests/
+
+# Auto-format
+uv run ruff format claymodel/ tests/
+```
+
+Configuration is in `pyproject.toml`. Key rules:
+- Max line length: 100
+- Max function arguments: 6 (with exceptions for model constructors)
+- Import sorting enforced (isort-compatible)
+
+## Making Changes
+
+### Before You Start
+
+1. Check [existing issues](https://github.com/Clay-foundation/model/issues) for related work
+2. For significant changes, open an issue first to discuss the approach
+
+### Workflow
+
+1. Create a branch from `main`
+2. Make your changes
+3. Run `ruff check` and `ruff format`
+4. Run `pytest tests/`
+5. Submit a pull request
+
+### Key Principles
+
+- **No changes to model computation**: Clay v1.5 must produce identical embeddings. Any refactoring must be verified with before/after numerical comparison.
+- **Test new functionality**: Add tests for new features. Tests should be fast (<30s total).
+- **Follow existing patterns**: Look at how similar features are implemented before adding new ones.
+- **Sensor metadata in metadata.yaml**: When adding new sensor support, add entries to `claymodel/configs/metadata.yaml`.
+
+## Adding New Sensors
+
+To add support for a new satellite sensor:
+
+1. Compute normalization statistics (mean/std per band) from a representative sample
+2. Find the central wavelength of each band in micrometers
+3. Add an entry to `claymodel/configs/metadata.yaml`
+4. Test with `clay info --sensor your-sensor` and `normalize(pixels, "your-sensor")`
+5. Submit a PR with the metadata and a brief description of the instrument
+
+## Questions?
+
+- Open an [issue](https://github.com/Clay-foundation/model/issues)
+- Start a [discussion](https://github.com/Clay-foundation/model/discussions)
+- Email: hello@madewithclay.org

README.md

@@ -21,54 +21,42 @@ Launch into a [JupyterLab](https://jupyterlab.readthedocs.io) environment on
 
 ## Installation
 
-### Pip Installation (Recommended)
+### uv Installation (Recommended)
 
-The easiest way to install Clay Foundation Model is via pip:
+The easiest way to install Clay Foundation Model is via `uv`:
 
-    pip install git+https://github.com/Clay-foundation/model.git
+    uv pip install git+https://github.com/Clay-foundation/model.git
 
 This will install the `claymodel` package and all its dependencies. You can then import and use it in your Python code:
 
 ```python
-from claymodel.datamodule import ClayDataModule
-from claymodel.module import ClayMAEModule
+from claymodel import load_model, embed
 ```
 
 ### Development Installation
 
-For development or advanced usage, you can set up the full development environment:
-
-To help out with development, start by cloning this [repo-url](/../../)
+For development or advanced usage, clone the repository and install with dev extras:
 
     git clone <repo-url>
     cd model
-
-Then we recommend [using mamba](https://mamba.readthedocs.io/en/latest/installation/mamba-installation.html)
-to install the dependencies. A virtual environment will also be created with Python and
-[JupyterLab](https://github.com/jupyterlab/jupyterlab) installed.
-
-    mamba env create --file environment.yml
-
-> [!NOTE]
-> The command above has been tested on Linux devices with CUDA GPUs.
-
-Activate the virtual environment first.
-
-    mamba activate claymodel
+    uv pip install -e ".[dev]"
 
 Finally, double-check that the libraries have been installed.
 
-    mamba list
+    uv run pytest tests/test_imports.py -q
 
 
 ## Usage
 
 ### Running jupyter lab
 
-    mamba activate claymodel
-    python -m ipykernel install --user --name claymodel  # to install virtual env properly
-    jupyter kernelspec list --json                       # see if kernel is installed
-    jupyter lab &
+    uv run jupyter lab
+
+### Using the clay CLI
+
+    clay info
+    clay info --sensor sentinel-2-l2a
+    clay benchmark
 
 
 ### Running the model
@@ -77,22 +65,22 @@ The neural network model can be run via
 [LightningCLI v2](https://pytorch-lightning.medium.com/introducing-lightningcli-v2supercharge-your-training-c070d43c7dd6).
 
 > [!NOTE]
-> If you installed via pip, you'll need to clone the repository to access the trainer script and config files.
+> If you installed from the package, you'll need to clone the repository to access the trainer script and config files.
 
 To check out the different options available, and look at the hyperparameter
 configurations, run:
 
-    python trainer.py --help
+    uv run python trainer.py --help
 
 To quickly test the model on one batch in the validation set:
 
-    python trainer.py fit --model ClayMAEModule --data ClayDataModule --config configs/config.yaml --trainer.fast_dev_run=True
+    uv run python trainer.py fit --model ClayMAEModule --data ClayDataModule --config configs/config.yaml --trainer.fast_dev_run=True
 
 To train the model:
 
-    python trainer.py fit --model ClayMAEModule --data ClayDataModule --config configs/config.yaml
+    uv run python trainer.py fit --model ClayMAEModule --data ClayDataModule --config configs/config.yaml
 
-More options can be found using `python trainer.py fit --help`, or at the
+More options can be found using `uv run python trainer.py fit --help`, or at the
 [LightningCLI docs](https://lightning.ai/docs/pytorch/2.1.0/cli/lightning_cli.html).
 
 ## Contributing