Merge main

Author: ryan-williams

.gitignore

@@ -38,3 +38,7 @@ output.log
 .data_cache/
 configs/schemas/
 /playpen
+*.env
+build/
+.pytype/
+.claude/

AGENTS.md

@@ -0,0 +1,125 @@
+# AGENTS.md
+
+This file provides guidance to automated agents when working with code in this repository.
+
+## Overview
+
+Ocean Emulator is a PyTorch-based machine learning project for training and evaluating models that emulate ocean physics. Currently, it implements a ConvNeXt U-Net neural network architecture.
+
+## Development Commands
+
+### Environment Setup
+
+```bash
+# Install dependencies with uv
+uv sync --dev
+source .venv/bin/activate
+```
+
+### Running Tests
+
+```bash
+# Run standard tests (excluding manual and CUDA tests)
+uv run pytest -m "not manual and not cuda"
+
+# Run tests with multiple CPU cores
+uv run pytest -m "not manual and not cuda" -n auto
+
+# Run CUDA tests (requires GPU)
+uv run pytest -m cuda
+
+# Run manual tests
+uv run pytest -m manual -k test_whatever
+
+# Run specific test file
+uv run pytest tests/test_train.py
+
+# Run benchmarks
+uv run pytest --benchmark-only --benchmark-autosave
+```
+
+### Code Quality Checks
+
+```bash
+# Run all pre-commit checks (linting, formatting, type checking)
+uvx pre-commit run --all-files
+```
+
+### Training, Evaluation, and Visualization
+
+Please ask if you need to train or evaluate a model.
+
+For vizualization or other long-running tasks:
+
+* Run with `PYTHONUNBUFFERED=1 uv run ... > /tmp/logfile.txt 2>&1`
+* Wait for a while with `timeout $time tail --pid=$pid -f /tmp/logfile.txt`
+* You may need to run this repeatedly.
+
+## High-Level Architecture
+
+### Core Components
+
+1. **Model Architecture** (`src/ocean_emulators/models/`)
+   * `convnext_unet.py`: Main neural network implementing ocean predictions
+   * `modules/`: Reusable network blocks (ConvNext, etc.)
+   * Models predict ocean variables from ocean model data
+
+2. **Data Pipeline** (`src/ocean_emulators/datasets.py`)
+   * Handles OM4 and CM4 ocean model data via Zarr format
+   * Supports time-based train/validation splits
+   * Variables include temperature, salinity, u/v velocities, sea surface height, and heat
+   * Data normalization and preprocessing
+
+3. **Training Loop** (`src/ocean_emulators/train.py`)
+   * Distributed training support via PyTorch DDP
+   * Checkpointing with model state and optimizer
+   * Weights & Biases integration for experiment tracking
+   * Learning rate scheduling and warmup
+   * Initializes data loaders and training datasets.
+
+4. **Evaluation System** (`src/ocean_emulators/eval.py`, `aggregator/`)
+   * Comprehensive metrics including RMSE, bias, correlations
+   * Ocean heat content (OHC) analysis
+   * ENSO metrics and basin-specific analysis
+   * Visualization tools for maps, time series, PDFs
+   * Aggregator pattern for metric collection
+
+5. **Configuration System** (`src/ocean_emulators/config.py`)
+   * YAML-based configuration with JSON schema validation
+   * Hierarchical configs with `!include` directives
+   * Pydantic models for type safety
+   * Command-line overrides supported
+
+### Key Design Patterns
+
+1. **Multiton Pattern**: Used for managing global state in tests via `MultitonScope`
+2. **Factory Pattern**: For creating network blocks dynamically
+3. **Configuration-Driven**: All major components configured via YAML
+4. **Aggregator Pattern**: For collecting distributed metrics during evaluation
+
+### Project Structure
+
+```text
+src/ocean_emulators/
+├── train.py              # Training entry point
+├── eval.py               # Evaluation entry point
+├── config.py             # Configuration classes
+├── config_schema.py      # JSON schema generation
+├── datasets.py           # Data loading
+├── models/               # Neural network architectures
+├── aggregator/           # Metric aggregation
+└── utils/                # Utilities for distributed training, logging
+
+configs/                  # YAML configuration files
+tests/                    # Comprehensive test suite
+scripts/                  # Data download and preprocessing
+```
+
+### Important Considerations
+
+1. **Data Format**: Uses Zarr format for efficient array storage
+2. **Distributed Training**: Supports multi-GPU via PyTorch DDP
+3. **Performance Mindset**: We include profiling tools (memray, py-spy, scalene) and aim to keep the code performant core train and eval loops.
+4. **Testing Philosophy**: Tests marked as `manual` or `cuda` for selective execution
+5. **Pre-commit Hooks**: Enforces code quality (ruff, mypy, detect-secrets)
+6. **Cloud Training**: Supports SkyPilot for remote job execution

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

CONTRIBUTING.md

@@ -162,20 +162,28 @@ git push --force-with-lease
 ```bash
 DATA_PATH=path/to/save/data
 uv run scripts/clone_data.py $DATA_PATH --compact_variables
-uv run -m ocean_emulators.train configs/train_om4.yaml --data_root $DATA_PATH
+uv run -m ocean_emulators.train configs/train_om4.yaml --experiment.data_root $DATA_PATH
 ```
 
 You can run `uv run -m ocean_emulators.train --help` to see all the options available.
 
 To learn more about other datasets used during training, please see the _Data Engineering_ section below.
 
+To run a remote training job with Skypilot, use the following command:
+```shell
+# export WANDB_API_KEY=<my-key>  # Get your key at https://wandb.ai/authorize
+# sky launch -c fomo-cluster train.sky.yaml  --env WANDB_API_KEY --env-file <my-vars>.env
+```
+
+Please read the docstring in the `train.sky.yaml` for more information.
+
 ### Evaluating the model
 
 ```bash
 DATA_PATH=path/to/save/data
 uv run scripts/clone_data.py $DATA_PATH --compact_variables
 # (then put a checkpoint of the model at path/to/checkpoint)
-uv run -m ocean_emulators.eval configs/eval_om4.yaml --ckpt_path path/to/checkpoint --data_root $DATA_PATH
+uv run -m ocean_emulators.eval configs/eval_om4.yaml --ckpt_path path/to/checkpoint --experiment.data_root $DATA_PATH
 ```
 
 You can run `uv run -m ocean_emulators.eval --help` to see all the options available.

configs/data/om4_sky_local.yaml

@@ -0,0 +1,17 @@
+# yaml-language-server: $schema=../schemas/DataConfig.json
+
+data_location:
+  type: local
+  path: OM4.zarr
+
+data_means_location:
+  type: local
+  path: OM4_means.zarr
+
+data_stds_location:
+  type: local
+  path: OM4_stds.zarr
+
+static_data_vars:
+  - sea_surface_fraction
+  - hfgeou

configs/slurm_perlmutter_train_cm4.yaml

@@ -7,7 +7,7 @@ save_freq: 5
 epochs: 120
 batch_size: 3
 learning_rate: 0.0001
-scheduler: true
+scheduler: {type: cosine}
 loss: mse
 finetune: false
 resume_ckpt_path: null

configs/slurm_perlmutter_train_om4.yaml

@@ -7,7 +7,7 @@ save_freq: 5
 epochs: 70
 batch_size: 4
 learning_rate: 0.0002
-scheduler: true
+scheduler: {type: cosine}
 loss: mse
 finetune: false
 resume_ckpt_path: null

configs/train_cm4.yaml

@@ -7,7 +7,7 @@ save_freq: 5
 epochs: 120
 batch_size: 3 # 4
 learning_rate: 0.0001
-scheduler: false
+scheduler: null
 loss: mse
 finetune: false
 resume_ckpt_path: null

configs/train_default.test.yaml

@@ -10,7 +10,7 @@ epochs: 2
 batch_size: 1
 preemptible: false
 learning_rate: 0.0001
-scheduler: false
+scheduler: null
 loss: mse
 finetune: false
 resume_ckpt_path: null

configs/train_default_2step.test.yaml

@@ -10,7 +10,7 @@ epochs: 1
 batch_size: 1
 preemptible: false
 learning_rate: 0.0001
-scheduler: false
+scheduler: null
 loss: mse
 finetune: false
 resume_ckpt_path: null