v2.0.0 - #5
Open
aaTman wants to merge 116 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>
* 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 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> * Expand front target classes to 9: merge forming/dissipating variants, add trough/tropical-trough/instability Forming (CF-F..OF-F) and dissipating (CF-D..OF-D) front codes now collapse into their parent front class instead of being dropped as background. Trough (TROF), tropical trough (TT), and instability axis (INST) become new independent classes 6, 7, and 8, bringing the total to 9 (was 6: background + CF/WF/SF/OF/DL). n_classes, class_weights, loss_class_weights, and front_types are updated across all training/eval/pipeline configs. The CONUS->full-domain warm-start config keeps skip_mismatch=True so the 6-class CONUS checkpoint's output layer is dropped and the 3 new classes train from scratch alongside the rest of the fine-tuned decoder. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * swap to completed full domain pretrain in schooner eval * Add trough, tropical trough, and instability axis to training-viz callback callbacks.py had its own local 5-class FRONT_TYPE_CLASS_INDEX that was missed when the rest of the codebase (targets.py, evaluate.py, plot.py, configs) was expanded to 9 classes. Training-time visualization images and lite performance-diagram stats were silently excluding TROF/TT/INST. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * update pretrained run name * update all run names * typo in run name * Fix GPU OOM in TestVisualizationCallback by streaming predict() to CPU per batch model.predict() batches its forward passes internally but still accumulates every batch's output into one GPU-resident tensor before returning. On full-CONUS-domain training runs (large spatial extent x 200-timestep visualization subsample x 9 front classes), that accumulated buffer on top of training's already-resident GPU memory reliably OOMs (observed crash: OOM in ConcatV2 inside TestVisualizationCallback._predict at shape[200,320,960,9]). Loop over predict_on_batch and move each batch to CPU immediately instead, so only one batch's output is ever GPU-resident at a time. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Add self-contained aies1702 harness for model_1702 comparison Runs the legacy AIES FrontFinder checkpoint (Keras 2.10 h5) through the 2.0 eval/plot machinery for a head-to-head against current models on identical 2019 test data: - loader: patches the stored model_config (TFOpLambda squeeze -> LevelSqueeze, standard inbound-node format), rebuilds under Keras 3, loads weights by name with full weight-consumption verification (8,745,968 params, 4 softmax heads) - legacy_formulas: verbatim numpy port of utils/variables.py@c59ac2b with docstring golden-value tests (derived.py diverges: Lv 2.501e6 vs 2.257e6) - normalization: embedded min-max table, parity-tested against model_1702_properties.txt - adapter: raw (batch, lat, lon, level, var) -> legacy normalization -> lon-major transpose -> sup1 head -> 6-to-9-class zero padding, matching evaluate.compute_stats' contract; ClassPaddingAdapter for 6-class baselines - store: side icechunk store generation (surface+1000/950/900/850, legacy units) from Arraylake ERA5 or ARCO fallback, resume-safe batched writes - run_eval: config-driven driver for both model kinds over full/land/ocean and the five office regions (with wrap-crossing longitude fix), writing evaluate-compatible stats_*.nc - figures/compare: batch 4-panel rendering and cross-model max-CSI summary with out-of-domain footnotes Zero core-module edits; verified end-to-end on the 2019-03-13 bomb cyclone (coherent CF/WF/DL structures from a real ARCO ERA5 timestep). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Add config-driven case-study figures mimicking AIES Fig. 14 fronts.aies1702.case_study renders multi-panel model_1702 prediction maps in the paper's presentation: filled front-probability contours at 10% intervals (CF/WF/SF/OF cmaps from fronts.plot), state-border backgrounds via plot_background, yellow a)/b)/c) panel labels, and per-panel UTC captions. Inputs are derived on the fly from Arraylake or ARCO ERA5 with the legacy formulas and cached to NetCDF so repeat renders skip the fetch entirely. Ships configs/aies1702/case_study_xmas2023.yaml reproducing the 2023 Christmas storm panels (00Z/12Z 26-27 Dec 2023); verified end-to-end locally against the real checkpoint — panel structures match the paper's Fig. 14, including the d-panel occlusion remnants over the UP of Michigan and the Midwest called out in the text. Raw (uncalibrated) probabilities from ERA5 analysis vs the paper's calibrated GFS f000, noted in the figure title. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Switch side-store source to ARCO and add single-job comparison SLURM script The Arraylake earthmover-public ERA5 repo lacks the 950 and 900 hPa levels model_1702 requires, so both generate configs now point at the anonymous Google ARCO ERA5 zarr (no auth needed). scripts/aies1702_comparison.sbatch runs the whole 2019 comparison in one job: side-store generation (resume-safe), all five evaluations, figures, and the comparison summary. Resubmit-safe — completed eval stages are skipped via their final region's stats_derived marker file. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * Rename harness namespace from aies1702 to model_1702 The artifact under comparison is model_1702, so the package, tests, configs, SLURM script, model_kind string, keras serialization package, render-test env var (MODEL_1702_RENDER_TESTS), and all schooner paths now use model_1702. No functional changes; 83 tests pass unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * add cpu and mem resources to slurm script * update vars * simplify evaluate to avoid recomputing and running redundant inference * numpy optimizations * figures slurm script * Remove trough/forming/dissipating classes, restore 6-class scheme harness/model_1702-comparison was branched from feat/include-forming-dissipating-and-troughs, so every shared file it touched carried that branch's 9-class scheme (TROF/TT/INST plus forming/dissipating collapse) even though the actual model_1702 comparison harness itself never needed more than the original 6 classes. Reverts config.py, model.py, plot.py, targets.py, and the schooner_*.yaml training configs to their feat/2.0.0 state; keeps the callbacks.py predict_on_batch OOM fix and the evaluate.py predict_batches/accumulate_stats split plus numpy optimizations, which are unrelated to class count; and fixes the new model_1702/ package (adapter.py, eval configs, tests) to assume 6 classes throughout. 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
1 participant
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.