v2.0.0#5
Open
aaTman wants to merge 111 commits into
Open
Conversation
…tch chunks attribute
* 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
marked this pull request as ready for review
July 9, 2026 19:29
* 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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
src/fronts/data/— data generation, sources, targets, datasets, and derived variablessrc/fronts/layers/— custom Keras layers, losses, metrics, activations, and modulessrc/fronts/plot/— visualization and plotting utilitiessrc/fronts/— core modules:train.py,model.py,evaluate.py,calibrate.py,pipeline.py,utils.pytrain_model.py,models.py,file_manager.py, custom loss/activation files, etc.)requirements.txt→pyproject.toml: Modern Python packaging withpixias the lock/environment managerpixi.lock: 10,617 lines of locked dependencies for reproducible environments.pre-commit-config.yamlfor code quality checks.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 maintenanceconfigs/schooner_train.yaml— training configurationconfigs/schooner_eval.yaml— evaluation configurationconfigs/schooner_pipeline.yaml— full pipeline orchestrationconfigs/generate_icechunk.yaml— data generation for icechunk storagescripts/:benchmark_dataloaders.py— performance testingcompare_era5_sources.py— data source validationrechunk_era5_store.py— storage optimizationscan_nan_channels.py— data quality checksCore Module Highlights
Data Pipeline (
src/fronts/data/)DataGenerationConfigPyDatasetintegration for efficient TensorFlow training (more info here)Model Architecture (
src/fronts/layers/)Training & Evaluation (
src/fronts/train.py/src/fronts/evaluate.py)Visualization (
src/fronts/plot/)Test Coverage
test_generate.py)test_inputs.py,test_targets.py)test_derived.py)test_losses.py,test_metrics.py)test_train.py)test_calibrate.py)Documentation & Quality
README.mdwith setup and usage instructionspyproject.tomlwith full project metadata and build configuration.gitignoreenhanced for Python, data, and logsNotable Commits
PyDatasetadoption (June 19): Modern TensorFlow data pipeline integrationStatistics
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.