201 downscaling

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

@@ -53,6 +53,10 @@ data:
   input_size_lat: 60 # 6x coarsening 
   input_size_lon: 96 # 6x coarsening 
   apply_smoothen: True
+  data_path_surface: /home/romeokienzler/Downloads/merra-2
+  data_path_vertical: /home/romeokienzler/Downloads/merra-2
+  climatology_path_surface: /home/romeokienzler/Downloads/climatology
+  climatology_path_vertical: /home/romeokienzler/Downloads/climatology
 
 model:
     # Platform independent config
@@ -82,6 +86,10 @@ model:
     encoder_decoder_conv_channels: 128
     freeze_backbone: True
     freeze_decoder: False
+    input_scalers_surface_path: /home/romeokienzler/Downloads/climatology/musigma_surface.nc
+    input_scalers_vertical_path: /home/romeokienzler/Downloads/climatology/musigma_vertical.nc
+    output_scalers_surface_path: /home/romeokienzler/Downloads/climatology/anomaly_variance_surface.nc
+    output_scalers_vertical_path: /home/romeokienzler/Downloads/climatology/anomaly_variance_vertical.nc
 
 
 

examples/confs/granite-wxc-merra2-downscale-config_ccc.yaml

@@ -0,0 +1,120 @@
+data:
+  type: merra2
+
+  # 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
+  data_path_surface: /dccstor/wfm/shared/datasets/training/merra-2_v1/merra-2
+  data_path_vertical: /dccstor/wfm/shared/datasets/training/merra-2_v1/merra-2
+  climatology_path_surface: /dccstor/wfm/shared/datasets/training/merra-2_v1/climatology
+  climatology_path_vertical: /dccstor/wfm/shared/datasets/training/merra-2_v1/climatology
+
+model:
+    # Platform independent 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
+
+    # Accepted values: temporal, climate, none
+    residual: climate
+
+    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
+    input_scalers_surface_path: /dccstor/wfm/shared/datasets/training/merra-2_v1/climatology/musigma_surface.nc
+    input_scalers_vertical_path: /dccstor/wfm/shared/datasets/training/merra-2_v1/climatology/musigma_vertical.nc
+    output_scalers_surface_path: /dccstor/wfm/shared/datasets/training/merra-2_v1/climatology/anomaly_variance_surface.nc
+    output_scalers_vertical_path: /dccstor/wfm/shared/datasets/training/merra-2_v1/climatology/anomaly_variance_vertical.nc
+
+
+
+job_id: inference-test
+batch_size: 1
+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

examples/notebooks/InspectMerra2Data.ipynb

@@ -0,0 +1,174 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "id": "initial_id",
+   "metadata": {
+    "collapsed": true
+   },
+   "source": [
+    "import xarray as xr\n",
+    "from granitewxc.utils.config import get_config\n",
+    "from granitewxc.utils.data import _get_transforms\n",
+    "\n",
+    "config = get_config(\"../../integrationtests/test_prithvi_wxc_model_factory_config.yaml\")"
+   ],
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "ds = xr.open_dataset(f'{config.data.data_path_surface}/MERRA2_sfc_19800101.nc')",
+   "id": "491de2f5a1fcacc7",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "ds",
+   "id": "f3aaad1bfe70c524",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "",
+   "id": "45c6800d47520f1a",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "from terratorch.datamodules.merra2_downscale import Merra2DownscaleNonGeoDataModule\n",
+    "dm = Merra2DownscaleNonGeoDataModule(\n",
+    "    time_range=('2001-01-01', '2001-01-02'),\n",
+    "    data_path_surface = config.data.data_path_surface,\n",
+    "    data_path_vertical = config.data.data_path_vertical,\n",
+    "    climatology_path_surface = config.data.climatology_path_surface,\n",
+    "    climatology_path_vertical = config.data.climatology_path_vertical,\n",
+    "    input_surface_vars = config.data.input_surface_vars,\n",
+    "    input_static_surface_vars = config.data.input_static_surface_vars,\n",
+    "    input_vertical_vars = config.data.input_vertical_vars,\n",
+    "    input_levels = config.data.input_levels,\n",
+    "    n_input_timestamps = config.data.n_input_timestamps,\n",
+    "    output_vars=config.data.output_vars,\n",
+    "    transforms=_get_transforms(config),\n",
+    ")\n",
+    "dm.setup('predict')"
+   ],
+   "id": "25c2c082084f5070",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "item = dm.predict_dataloader().dataset[0]",
+   "id": "7b49b7ab084fc7db",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "type(item)",
+   "id": "c5a0d2342430727c",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "item.keys()",
+   "id": "31dc81de5131fc82",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "type(item['x'])",
+   "id": "d41df3dde88dc010",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "item['x'].shape",
+   "id": "c1143d39a3c39cc5",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "item['y'].shape\n",
+   "id": "b6a25e5b6046d9c7",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": [
+    "for key in item.keys():\n",
+    "    print(key)\n",
+    "    print(item[key].shape)"
+   ],
+   "id": "b491aa2e883634ab",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "ds = xr.open_dataset(f'{config.data.climatology_path_surface}/climate_surface_doy241_hour09.nc')\n",
+   "id": "3b6210fb62905090",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "ds",
+   "id": "e81e22919f5103d0",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "metadata": {},
+   "cell_type": "code",
+   "source": "",
+   "id": "a8a23c68d9b94c4b",
+   "outputs": [],
+   "execution_count": null
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}