Add data loading pipeline for training and prediction

.gitattributes

@@ -0,0 +1,2 @@
+# SCM syntax highlighting & preventing 3-way merges
+pixi.lock merge=binary linguist-language=YAML linguist-generated=true

.gitignore

@@ -57,7 +57,7 @@ coverage.xml
 
 # Django stuff:
 *.log
-local_settings.py
+local_data_utils.py
 db.sqlite3
 db.sqlite3-journal
 
@@ -130,3 +130,12 @@ dmypy.json
 
 # For PyCharm
 .idea
+# pixi environments
+.pixi/*
+!.pixi/config.toml
+
+# wandb
+wandb/
+
+# osx files
+.DS_Store

configs/1702.yaml

@@ -0,0 +1,124 @@
+epochs: 5000
+training_steps_per_epoch: 30
+validation_steps_per_epoch:
+validation_frequency: 1
+verbose: 1
+repeat: true
+seed: 42
+
+model:
+  name: "unet_3plus"
+  batch_normalization: true
+  num_filters: [16, 32, 64, 128]
+  kernel_size: [5, 5, 5]
+  pool_size: [2, 2, 1]
+  upsample_size: [2, 2, 1]
+  depth: 4
+  modules_per_node: 2
+  padding: "same"
+  bias: true
+  loss:
+    name: "fractions_skill_score"
+    config:
+      mask_size: [3, 3]
+  metric:
+    name: "critical_success_index"
+    config:
+      class_weights: [0, 1, 1, 1, 1, 1]
+
+  optimizer:
+    name: "Adam"
+    config:
+      beta_1: 0.9
+      beta_2: 0.999
+
+  convolution_activity_regularizer:
+    regularizer:
+
+  bias_vector:
+    constraint:
+    initializer: 
+      name: "zeros"
+      config:
+    regularizer:
+
+  kernel_matrix:
+    constraint:
+    initializer: 
+      name: "glorot_uniform"
+      config:
+    regularizer:
+
+  activation:
+    name: "gelu"
+    config:
+
+
+
+
+
+
+data:
+  train_years: [2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019]
+  val_years: [2020]
+  test_years: [2021]
+  shuffle: true
+  normalization_method: "standard"
+
+  era5:
+    domain_extent: [-140.0, -60.0, 20.0, 60.0]
+    # Use canonical (pressure-level) names for variables that have both surface
+    # and pressure-level representations.  Surface-only variables (e.g.
+    # "mean_sea_level_pressure") are listed here directly.
+    # The mapping from pressure-level name to surface name is defined in
+    # fronts.data.config.SURFACE_VARIABLE_MAP and applied automatically when
+    # "surface" appears in the levels list.
+    variables:
+      - "temperature"            # surface counterpart: 2m_temperature
+      - "specific_humidity"      # surface counterpart: surface_specific_humidity
+      - "u_component_of_wind"    # surface counterpart: 10m_u_component_of_wind
+      - "v_component_of_wind"    # surface counterpart: 10m_v_component_of_wind
+      - "mean_sea_level_pressure"  # surface-only
+    # levels may contain "surface" and/or integer hPa values in any order.
+    # "surface" triggers loading of the surface counterpart for each variable.
+    levels: ["surface", 1000, 950, 900, 850]
+    store: "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+    chunks:
+      time: 48
+    consolidated: true
+
+  fronts:
+    directory: "/path/to/fronts/netcdf"
+    front_types: "MERGED-ALL"
+
+  batch:
+    input_sizes:
+      time: 1
+      latitude: 128
+      longitude: 128
+    target_sizes:
+      time: 1
+      latitude: 128
+      longitude: 128
+    prefetch_number: 3
+    preload_batch: false
+
+wandb:
+  project_name: 'fronts'
+  model_run_name: '1702_retrain'
+  # -- Optional parameters -- 
+  # log_frequency: 
+  # upload_checkpoints: 
+  # api_key: 
+  # wandb_filepath:
+
+callbacks:
+  monitor: "val_loss"
+  verbose: 1
+  save_best_only: true
+  save_weights_only: false
+  save_freq: "epoch"
+  # model_checkpoint_path:
+  # csv_logger_path:
+  # patience:
+

configs/1702_tf.yaml

@@ -0,0 +1,93 @@
+epochs: 5000
+training_steps_per_epoch: 30
+validation_steps_per_epoch:
+validation_frequency: 1
+verbose: 1
+repeat: true
+seed: 42
+
+model:
+  name: "unet_3plus"
+  batch_normalization: true
+  num_filters: [16, 32, 64, 128]
+  kernel_size: [5, 5, 5]
+  pool_size: [2, 2, 1]
+  upsample_size: [2, 2, 1]
+  depth: 4
+  modules_per_node: 2
+  padding: "same"
+  bias: true
+  loss:
+    name: "fractions_skill_score"
+    config:
+      mask_size: [3, 3]
+  metric:
+    name: "critical_success_index"
+    config:
+      class_weights: [0, 1, 1, 1, 1, 1]
+
+  optimizer:
+    name: "Adam"
+    config:
+      beta_1: 0.9
+      beta_2: 0.999
+
+  convolution_activity_regularizer:
+    regularizer:
+
+  bias_vector:
+    constraint:
+    initializer:
+      name: "zeros"
+      config:
+    regularizer:
+
+  kernel_matrix:
+    constraint:
+    initializer:
+      name: "glorot_uniform"
+      config:
+    regularizer:
+
+  activation:
+    name: "gelu"
+    config:
+
+
+data:
+  train_years: [2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2020, 2021, 2022, 2024]
+  val_years: [2023]
+  test_years: []
+  shuffle: true
+
+  # Load pre-built TF dataset snapshots directly from disk.
+  # Subdirectory naming convention: {year}-{month}_tf  (e.g. 2010-1_tf)
+  # train_years/val_years/test_years above are injected at build time —
+  # leave them as [] here.
+  tf_dataset:
+    directory: "/ourdisk/hpc/ai2es/fronts/raw_front_data/tf_datasets/5class-7lvl-conus"
+    train_years: []
+    val_years: []
+    test_years: []
+    shuffle: true
+    shuffle_buffer: 1000
+    prefetch: 3
+
+wandb:
+  project_name: 'fronts'
+  model_run_name: '1702_retrain'
+  wandb_filepath: '/ourdisk/hpc/ai2es/tman/models/1702_retrain.keras'
+  # -- Optional parameters --
+  # log_frequency:
+  # upload_checkpoints:
+  # api_key:
+
+callbacks:
+  monitor: "val_loss"
+  verbose: 1
+  save_best_only: true
+  save_weights_only: false
+  save_freq: "epoch"
+  model_checkpoint_path: "/ourdisk/hpc/ai2es/tman/models/1702_tf.keras"
+  csv_logger_path: "/ourdisk/hpc/ai2es/tman/logs/1702_tf.csv"
+  # patience: 55

configs/predict_1702.yaml

@@ -0,0 +1,53 @@
+# Prediction configuration for FrontFinder model 1702.
+#
+# This config is loaded as a PredictConfig via:
+#   from fronts.train import open_config_yaml_as_dataclass
+#   from fronts.data.config import PredictConfig
+#   cfg = open_config_yaml_as_dataclass("configs/predict_1702.yaml", PredictConfig)
+#   inputs_ds = cfg.build()  # -> normalized xr.Dataset ready for model inference
+#
+# TIME SELECTION — uncomment exactly one block below.
+# Use ISO 8601 format with a "T" separator (e.g. "2024-06-01T12:00:00") so that
+# PyYAML treats the value as a string rather than a date, allowing dacite to cast
+# it to datetime.datetime correctly.
+
+# Option A: most recent timestep available in the zarr store (no date needed)
+time_selection:
+  most_recent: true
+
+# Option B: explicit individual analysis times (date + hour each)
+# time_selection:
+#   timestamps:
+#     - "2024-06-01T12:00:00"
+#     - "2024-06-02T00:00:00"
+
+# Option C: all timesteps in an inclusive date range
+# time_selection:
+#   date_range:
+#     - "2024-06-01T00:00:00"
+#     - "2024-06-07T18:00:00"
+
+normalization_method: "standard"
+
+era5:
+  domain_extent: [-140.0, -60.0, 20.0, 60.0]   # [lon_min, lon_max, lat_min, lat_max]
+  # Use canonical (pressure-level) names for variables that have both surface and
+  # pressure-level representations.  Surface-only variables (e.g.
+  # "mean_sea_level_pressure") are listed directly.
+  # The mapping from pressure-level name to surface name is defined in
+  # fronts.data.config.SURFACE_VARIABLE_MAP and applied automatically when
+  # "surface" appears in the levels list.
+  variables:
+    - "temperature"              # surface counterpart: 2m_temperature
+    - "specific_humidity"        # surface counterpart: surface_specific_humidity
+    - "u_component_of_wind"      # surface counterpart: 10m_u_component_of_wind
+    - "v_component_of_wind"      # surface counterpart: 10m_v_component_of_wind
+    - "mean_sea_level_pressure"  # surface-only
+  # levels may contain "surface" and/or integer hPa values.
+  # "surface" triggers loading of the surface counterpart for each variable.
+  levels: ["surface", 1000, 950, 900, 850]
+  store: "gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3"
+  chunks:
+    time: 48
+  consolidated: true
+  # Note: era5.years is unused by PredictConfig — time is controlled by time_selection above.