Skip to content

v2.0.0#5

Open
aaTman wants to merge 111 commits into
ai2es:masterfrom
aaTman:feat/2.0.0
Open

v2.0.0#5
aaTman wants to merge 111 commits into
ai2es:masterfrom
aaTman:feat/2.0.0

Conversation

@aaTman

@aaTman aaTman commented Jun 30, 2026

Copy link
Copy Markdown

This is a comprehensive refactor that transforms the fronts project into a well-structured, tested Python package. It introduces proper project organization, CI/CD pipelines, a complete test suite, and modernized dependencies via pixi.

Key Structural Changes:

  1. Project Organization
  • New package structure: All source code moved to src/fronts/ with organized modules:
    • src/fronts/data/ — data generation, sources, targets, datasets, and derived variables
    • src/fronts/layers/ — custom Keras layers, losses, metrics, activations, and modules
    • src/fronts/plot/ — visualization and plotting utilities
    • src/fronts/ — core modules: train.py, model.py, evaluate.py, calibrate.py, pipeline.py, utils.py
  • Removed legacy files: All old root-level scripts deleted (train_model.py, models.py, file_manager.py, custom loss/activation files, etc.)
  • Test suite added: Comprehensive tests under tests/ mirroring the source structure with fixtures and config files
  1. Dependency Management
  • Moved from requirements.txtpyproject.toml: Modern Python packaging with pixi as the lock/environment manager
  • Generated pixi.lock: 10,617 lines of locked dependencies for reproducible environments
  • Pre-commit hooks: Added .pre-commit-config.yaml for code quality checks
  1. CI/CD & Automation
  • GitHub Actions workflows:
    • .github/workflows/ci.yml — continuous integration pipeline
    • .github/workflows/tests.yml — automated test suite execution
    • .github/workflows/update-locks.yml — dependency updates and lock file maintenance
  1. Configuration & Scripts
  • Centralized YAML configs:
    • configs/schooner_train.yaml — training configuration
    • configs/schooner_eval.yaml — evaluation configuration
    • configs/schooner_pipeline.yaml — full pipeline orchestration
    • configs/generate_icechunk.yaml — data generation for icechunk storage
  • Utility scripts in scripts/:
    • benchmark_dataloaders.py — performance testing
    • compare_era5_sources.py — data source validation
    • rechunk_era5_store.py — storage optimization
    • scan_nan_channels.py — data quality checks

Core Module Highlights

Data Pipeline (src/fronts/data/)

  • Config-driven data generation with DataGenerationConfig
  • Modular data sources (ERA5, fronts) and targets
  • Derived variable computation (equivalent potential temperature, etc.)
  • PyDataset integration for efficient TensorFlow training (more info here)

Model Architecture (src/fronts/layers/)

  • Custom activations (including Gelu, Swish)
  • Advanced loss functions (physics-informed losses)
  • Evaluation metrics
  • Keras layer builders and custom modules (attention, etc.)

Training & Evaluation (src/fronts/train.py / src/fronts/evaluate.py)

  • Comprehensive training loop with early stopping, model checkpointing
  • Memory logging and garbage collection management
  • Evaluation pipeline with performance statistics
  • Integration with Weights & Biases for experiment tracking

Visualization (src/fronts/plot/)

  • Plotting utilities for model predictions and performance diagnostics
  • Customizable colormaps and geographic projections

Test Coverage

  • 1000+ lines of tests covering:
    • Data generation and validation (test_generate.py)
    • Input/output handling (test_inputs.py, test_targets.py)
    • Derived variables and transformations (test_derived.py)
    • Loss functions and metrics (test_losses.py, test_metrics.py)
    • Training workflows (test_train.py)
    • Calibration (test_calibrate.py)

Documentation & Quality

  • Updated README.md with setup and usage instructions
  • pyproject.toml with full project metadata and build configuration
  • .gitignore enhanced for Python, data, and logs
  • Pre-commit integration for automated code quality

Notable Commits

  • Initial reorganization (May 27–28): Project structure and CI setup
  • Data pipeline refactor (May–June): Modularized data generation with improved variable handling
  • Model training improvements (June 1–9): Checkpointing, memory logging, vertical velocity channel
  • PyDataset adoption (June 19): Modern TensorFlow data pipeline integration
  • Evaluation refactor (June 26–29): Comprehensive evaluation and visualization pipeline
  • Testing setup (Ongoing): Test suite expansion and validation

Statistics

  • Files: 100+ new/modified files
  • Commits: 50+ commits driving incremental improvements
  • Code: ~4,000 lines of new source code, ~3,000+ lines of tests
  • Deletions: ~6,500 lines of deprecated scripts and modules

Purpose

v2.0.0 establishes a production-ready foundation for the fronts deep learning project with proper software engineering practices: modularity, testability, reproducibility via lock files, and continuous integration.

aaTman added 30 commits May 27, 2026 17:35
* add gc callback

* fix wandb logging and add memory logging

* add model checkpointing

* add model checkpoint path to config
* init commit; convert old stats generation for new approach

* remove default coordinates from evalconfig

* add eval config

* add imports

* update imports and eval domain

* turn compile off for loading model

* place tf imports at top of files where possible

* use periodicboundaryindex

* reorder the selection and subsetting for front targets

* handle dask weirdness

* fixing zerodivision weirdness

* latitude sel flip

* more fixes, typing, add drop duplicate func

* trying to resolve the lat/lon slicing

* remove dask refs

* add iso8601 config date option

* improve throughput, load data into memory

* logging to figure out speed up opportunities

* logging

* move bool mask outside of loop

* attempt 1 at speedup

* ruff

* update cmap function

* change central lon calc

* modify colormaps

* add logs and compute stats prealloc

* fix contourf longitude coords and have predict_config use coordinates instead of domain. Add tests

* update dependences

* comment

* remove constants module and put relevant vars into plot and compute_stats
* change run name to vertical_velocity

* change model checkpoint dir

* dedupe era5

* regenerate icechunk

* change path of icechunk data, create variable addition check

* swap to to_icechunk and defer compute until that method

* update icechunk dir

* update num channels to account for new w variable (30 -> 36)

* correct num channels in default yaml

* Merge branch 'feat/2.0.0' into train/vertical-velocity
@aaTman
aaTman marked this pull request as ready for review July 9, 2026 19:29
aaTman and others added 7 commits July 13, 2026 11:05
* Add latitude-aware neighborhood Brier loss; remove label dilation

Replace fractions_skill_score as the training objective with a proper
scoring rule: the MSE between observed and forecast neighborhood
fractions (FSS's numerator without the hedgeable skill-score
normalization), summed over 25/100/250 km tolerance scales. Zonal
windows widen by 1/cos(lat) so tolerance stays isotropic in km, and
domain edges use valid-cell normalization (no reflect or wrap; the
longitude strip ends are not adjacent). No sigmoid discretization,
which would break propriety.

The 25 km scale reproduces the old 1-pixel label dilation, so
front_dilation is removed from the dataset, evaluation, and configs.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* run name update

* bump learning rate based on slow training

* update run names

* use min-max instead of standardization for compute_norm_stats

* brier loss but 3d

* set gradient clipping max, fix callbacks to end run with lr decay

* add brier conus

* more precise lat dependent pool

* test bucketing for lat dependency

* remove lat dependent pooling

* drop batch size and multiple neighborhoods

* add spatial subset internals

* merge conus and brier branches

* swap loss to fractions for test

* bugfix conv3dbackpropfilterv2

* swap to brier loss

* Revert "bugfix conv3dbackpropfilterv2"

This reverts commit bdcb14e.

* fix remainder issue

* swap to fractions for comparison test

* swap back to brier

* remove comment

* typo

* add 0.5 threshold hss metric

* use keras meanmetricwrapper for hss threshold

* try 32 for batch size

* generalize consolidation callback

* swap to fractions loss (again, for updated metrics and batch size)

* brier

* Add CONUS-to-full-domain pretrained warm-start and layer freezing (#41)

* Add pretrained-weight warm-start and layer-freezing for CONUS-to-full-domain transfer

Adds ModelConfig.pretrained_weights_path/freeze_layer_prefixes and train.py helpers
to load a prior checkpoint's weights, reset input_normalization to the current run's
domain stats, and optionally freeze layers by name prefix before compiling. Includes
a schooner_train_full_pretrained_3d.yaml config wiring this up for a full-domain
fine-tune from a CONUS-trained checkpoint.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Zero-weight the dryline class in the full-domain pretrained config

The full domain has no dryline labels because label coverage doesn't extend
there, not because drylines don't occur, so leaving that class weighted
would actively train the model toward "dryline never happens" and erode
what the CONUS checkpoint already learned for it.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* change eval run name

* Fix schooner_eval.yaml to match 3D CONUS model's volume-input config

The eval config was never updated for the brier_conus_3d_comparison
3D U-Net: it was missing volume_inputs (defaulting to the flattened
2D path), listed the wrong variable set/order, and pointed at the
full-domain coordinates instead of the CONUS box the model trained on.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant