fixes init

.github/dependabot.yaml

@@ -10,6 +10,7 @@ updates:
     directory: "/"
     schedule:
       interval: "daily"
+    target-branch: "maintenance"
     groups:
       # torchvision pins torch, must update in unison
       torch:
@@ -21,4 +22,4 @@ updates:
       - dependency-name: "setuptools"
         update-types: ["version-update:semver-patch"]
       # segmentation-models-pytorch pins timm, must update in unison
-      - dependency-name: "timm"
+      - dependency-name: "timm"

.github/workflows/pylint.yml

@@ -11,7 +11,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10","3.11"]
+        python-version: ["3.10","3.11","3.12"]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}

.github/workflows/test.yaml

@@ -12,32 +12,10 @@ on:
 jobs:
   build:
     runs-on: ubuntu-latest
+    timeout-minutes: 20
     strategy:
       matrix:
-        python-version: ["3.10", "3.11"]
-
-    steps:
-      - name: Clone repo
-        uses: actions/checkout@v3
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
-        with:
-          python-version: ${{ matrix.python-version }}
-          cache: 'pip'
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -r requirements/required.txt -r requirements/test.txt -r requirements/optional.txt
-      - name: List pip dependencies
-        run: pip list
-      - name: Test with pytest
-        run: |
-          pytest -s tests
-  build_weather:
-    runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: ["3.11"]
+        python-version: ["3.10", "3.11", "3.12"]
 
     steps:
       - name: Clone repo
@@ -58,3 +36,30 @@ jobs:
       - name: Test with pytest
         run: |
           pytest -s tests
+
+#  build_weather:
+#    runs-on: ubuntu-latest
+#    strategy:
+#      matrix:
+#        python-version: ["3.11"]
+
+#    steps:
+#      - name: Clone repo
+#        uses: actions/checkout@v3
+#      - name: Set up Python ${{ matrix.python-version }}
+#        uses: actions/setup-python@v4
+#        with:
+#          python-version: ${{ matrix.python-version }}
+#          cache: 'pip'
+#      - name: Install dependencies
+#        run: |
+#          python -m pip install --upgrade pip
+#          pip install -r requirements/required.txt -r requirements/test.txt -r requirements/optional.txt
+#          pip install git+https://github.com/NASA-IMPACT/Prithvi-WxC.git
+#          pip install git+https://github.com/IBM/granite-wxc.git
+#      - name: List pip dependencies
+#        run: pip list
+#      - name: Test with pytest
+#        run: |
+#          pytest -s tests
+

.gitignore

@@ -9,4 +9,12 @@ dist/*
 **/*pth
 .venv/*
 venv/*
+examples/notebooks/config.yaml
+examples/notebooks/wxc_input_u_v_t_p_output_theta_uw_vw_era5_training_data_hourly_2015_constant_mu_sigma_scaling05.nc
 tests/all_ecos_random/*
+**/climatology/*
+**/lightning_logs/*
+**/merra-2/*
+**/*.bin
+*.stdout
+*.log

examples/confs/granite-wxc-merra2-downscale-large-config.yaml

@@ -0,0 +1,245 @@
+# lightning.pytorch==2.1.1
+seed_everything: 0
+trainer:
+  accelerator: auto
+  strategy: auto
+  devices: auto
+  num_nodes: 1
+  precision: 16-mixed
+  logger: True # will use tensorboardlogger
+  callbacks:
+    - class_path: RichProgressBar
+    - class_path: LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+  max_epochs: 200
+  check_val_every_n_epoch: 1
+  log_every_n_steps: 50
+  enable_checkpointing: true
+  default_root_dir: examples/
+  #num_epochs: 400
+  #dl_num_workers: 19
+  #dl_prefetch_size: 1
+  #learning_rate: 0.0001
+  #limit_steps_train: 250
+  #limit_steps_valid: 25
+  #min_lr: 0.00001
+  #max_lr: 0.0002
+  #warm_up_steps: 0
+  #mask_unit_size:
+  #  - 15
+  #  - 16
+  #mask_ratio_inputs: 0.0
+  #mask_ratio_targets: 0.0
+  #max_batch_size: 16
+
+  #path_experiment: experiment
+  #loss: rmse
+  #freeze_backbone: True
+  #freeze_decoder: False
+  #backbone_prefix: encoder.
+  #finetune_w_static: True
+  #strict_matching: true
+data:
+  class_path: terratorch.datamodules.Merra2DownscaleNonGeoDataModule
+  init_args: 
+    transforms_fn: granitewxc.utils.data._get_transforms
+    # Input variables definition
+    input_surface_vars:
+      - EFLUX
+      - GWETROOT
+      - HFLUX
+      - LAI
+      - LWGAB # surface absorbed longwave radiation
+      - LWGEM # longwave flux emitted from surface
+      - LWTUP # upwelling longwave flux at toa
+      - PS # surface pressure
+      - QV2M # 2-meter specific humidity
+      - SLP # sea level pressure
+      - SWGNT # surface net downward shortwave flux
+      - SWTNT # toa net downward shortwave flux
+      - T2M # near surface temperature
+      - TQI # total precipitable ice water
+      - TQL # total precipitable liquid water
+      - TQV # total precipitable water vapor
+      - TS # surface skin temperature
+      - U10M # 10m eastward wind
+      - V10M # 10m northward wind
+      - Z0M # surface roughness
+    input_static_surface_vars: [FRACI, FRLAND, FROCEAN, PHIS]
+    input_vertical_vars:
+      - CLOUD # cloud feraction for radiation
+      - H # geopotential/ mid layer heights
+      - OMEGA # vertical pressure velocity
+      - PL # mid level pressure
+      - QI # mass fraction of clous ice water
+      - QL # mass fraction of cloud liquid water
+      - QV # specific humidity
+      - T # tempertaure
+      - U # eastward wind
+      - V # northward wind
+    # (model level/ml ~ pressure level/hPa)
+    # 52ml ~ 562.5hPa, 56ml ~ 700hPa, 63 ml ~ 850hPa
+    input_levels: [34.0, 39.0, 41.0, 43.0, 44.0, 45.0, 48.0, 53.0, 56.0, 63.0, 68.0, 72.0]
+    ## remove: n_input_timestamps: 1
+    # Output variables definition
+    output_vars:
+      - T2M # near surface temperature
+
+    n_input_timestamps: 2
+
+    # Data transformations
+    # Initial crop before any other processing
+    crop_lat: [0, 1]
+    # crop_lon: [0, 0]
+    # coarsening of target -- applied after crop
+    input_size_lat: 60 # 6x coarsening 
+    input_size_lon: 96 # 6x coarsening 
+    apply_smoothen: True
+  dict_kwargs:
+    data_path_surface: examples/scripts/merra-2/ 
+    data_path_vertical: examples/scripts/merra-2/ 
+    climatology_path_surface: examples/scripts/climatology/ 
+    climatology_path_vertical: examples/scripts/climatology/ 
+    time_range:
+      - '2020-01-01'
+      - '2020-01-02'
+model:
+  class_path: WxCDownscalingTask 
+  init_args:
+    model_args:
+      checkpoint_path: examples/pytorch_model.bin
+      num_static_channels: 7
+      embed_dim: 2560
+      token_size:
+      - 1
+      - 1
+      n_blocks_encoder: 12
+      mlp_multiplier: 4
+      n_heads: 16
+      dropout_rate: 0.0
+      drop_path: 0.05
+      mask_unit_size: 
+        - 15
+        - 16
+      residual_connection: True
+    model_factory: WxCModelFactory 
+    extra_kwargs:
+        encoder_decoder_kernel_size_per_stage: [[3], [3]] # Optional, default = 3 for conv_tanspose [[3], [2]]
+        output_vars:
+          - T2M # near surface temperature
+        type: merra2
+        input_scalers_surface_path: examples/scripts/climatology/musigma_surface.nc
+        input_scalers_vertical_path: examples/scripts/climatology/musigma_vertical.nc
+        output_scalers_surface_path: examples/scripts/climatology/anomaly_variance_surface.nc
+        output_scalers_vertical_path: examples/scripts/climatology/anomaly_variance_vertical.nc
+        input_levels: [34.0, 39.0, 41.0, 43.0, 44.0, 45.0, 48.0, 53.0, 56.0, 63.0, 68.0, 72.0]
+        downscaling_patch_size: [2, 2]
+        n_input_timestamps: 2
+        downscaling_embed_dim: 256
+        encoder_decoder_conv_channels: 128
+        encoder_decoder_scale_per_stage: [[2], [3]] # First list determines before/after backbone
+        encoder_decoder_upsampling_mode: pixel_shuffle # ['nearest', 'bilinear', 'pixel_shuffle', 'conv_transpose']
+        encoder_shift: False
+        drop_path: 0.05
+        encoder_decoder_type: 'conv' # ['conv', 'transformer']
+        input_size_lat: 60 # 6x coarsening 
+        input_size_lon: 96 # 6x coarsening 
+        freeze_backbone: True
+        freeze_decoder: False
+        data_path_surface: examples/scripts/merra-2/ 
+        data_path_vertical: examples/scripts/merra-2/ 
+        climatology_path_surface: examples/scripts/climatology/ 
+        climatology_path_vertical: examples/scripts/climatology/ 
+        residual: climate
+    model_config:
+      num_epochs: 200   
+      limit_steps_train: 250 
+      limit_steps_valid: 50
+      batch_size: 1 #16 
+      learning_rate: 0.0001 
+      min_lr: 0.00001 
+      dl_num_workers: 19 
+      dl_prefetch_size: 1 
+      path_experiment: experiment/
+      warm_up_steps: 0 
+      mask_ratio_inputs: 0.0
+      mask_ratio_targets: 0.0    # Accepted values: temporal, climate, none
+      job_id: inference-test
+      model_config:    
+        num_static_channels: 7
+        embed_dim: 2560
+        token_size:
+        - 1
+        - 1
+        n_blocks_encoder: 12
+        mlp_multiplier: 4
+        n_heads: 16
+        dropout_rate: 0.0
+        #drop_path: 0.05
+        residual: True
+      data_config: 
+        surface_vars:
+            - EFLUX
+            - GWETROOT
+            - HFLUX
+            - LAI
+            - LWGAB # surface absorbed longwave radiation
+            - LWGEM # longwave flux emitted from surface
+            - LWTUP # upwelling longwave flux at toa
+            - PS # surface pressure
+            - QV2M # 2-meter specific humidity
+            - SLP # sea level pressure
+            - SWGNT # surface net downward shortwave flux
+            - SWTNT # toa net downward shortwave flux
+            - T2M # near surface temperature
+            - TQI # total precipitable ice water
+            - TQL # total precipitable liquid water
+            - TQV # total precipitable water vapor
+            - TS # surface skin temperature
+            - U10M # 10m eastward wind
+            - V10M # 10m northward wind
+            - Z0M # surface roughness
+        static_surface_vars: [FRACI, FRLAND, FROCEAN, PHIS]
+        vertical_vars:
+            - CLOUD # cloud feraction for radiation
+            - H # geopotential/ mid layer heights
+            - OMEGA # vertical pressure velocity
+            - PL # mid level pressure
+            - QI # mass fraction of clous ice water
+            - QL # mass fraction of cloud liquid water
+            - QV # specific humidity
+            - T # tempertaure
+            - U # eastward wind
+            - V # northward wind
+         # (model level/ml ~ pressure level/hPa)
+         # 52ml ~ 562.5hPa, 56ml ~ 700hPa, 63 ml ~ 850hPa
+         #    levels: [34.0, 39.0, 41.0, 43.0, 44.0, 45.0, 48.0, 53.0, 56.0, 63.0, 68.0, 72.0]
+         ## remove: n_input_aimestamps: 1
+         # Output variables definition
+             #residual_connection: True
+      #encoder_shift: False
+
+      #downscaling_patch_size: [2, 2]
+      #downscaling_embed_dim: 256
+      #encoder_decoder_type: 'conv' # ['conv', 'transformer']
+      #encoder_decoder_upsampling_mode: pixel_shuffle # ['nearest', 'bilinear', 'pixel_shuffle', 'conv_transpose']
+      #encoder_decoder_kernel_size_per_stage: [[3], [3]] # Optional, default = 3 for conv_tanspose [[3], [2]]
+      #encoder_decoder_scale_per_stage: [[2], [3]] # First list determines before/after backbone
+      #encoder_decoder_conv_channels: 128
+      #freeze_backbone: True
+      #freeze_decoder: False
+      #ignore_index: -1
+      #loss: rmse 
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 6.e-5
+    weight_decay: 0.05
+lr_scheduler:
+  class_path: ReduceLROnPlateau
+  init_args:
+    monitor: val/loss
+
+
+