Upload 107 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+img/fig_csi.jpg filter=lfs diff=lfs merge=lfs -text
+img/fig_NeuralOM.jpg filter=lfs diff=lfs merge=lfs -text
+img/fig_rmse_acc.jpg filter=lfs diff=lfs merge=lfs -text
+img/fig_visual.jpg filter=lfs diff=lfs merge=lfs -text
+img/tab_acc_rmse.jpg filter=lfs diff=lfs merge=lfs -text
+networks/__pycache__/GraphCast.cpython-310.pyc filter=lfs diff=lfs merge=lfs -text

data/climate_mean_s_t_ssh.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8cdfd60af19771408e952a2214593593452e82ccb45d877ee1310bfb5f13d85
+size 36809870528

data/land_mask.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b21a4f8d43afff2391d95f4b7afa3242da33b6680ce93607193a3a106c50927
+size 96424448

data/mean_s_t_ssh.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2ed37c0e9f106ad2d403af1d6ceff5d41630a90207764200b4d6dd7aee9c9beb
+size 904

data/std_s_t_ssh.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fecbc3deb38ae6ed8f98ac46126959361c59ebfb33ea46e3de0774c97291325f
+size 904

data/test/.ipynb_checkpoints/test_data-checkpoint.txt

@@ -0,0 +1 @@
+Please put the test data in this folder

data/test/test_data.txt

@@ -0,0 +1 @@
+Please put the test data in this folder, the intact project is available at the Hugging Face.

data/train/.ipynb_checkpoints/train_data-checkpoint.txt

@@ -0,0 +1 @@
+Please put the train data in this folder

data/train/train_data.txt

@@ -0,0 +1 @@
+Please put the train data in this folder, the intact project is available at the Hugging Face.

data/valid/.ipynb_checkpoints/valid_data-checkpoint.txt

@@ -0,0 +1 @@
+Please put the valid data in this folder

data/valid/valid_data.txt

@@ -0,0 +1 @@
+Please put the valid data in this folder, the intact project is available at the Hugging Face.

environment.yml

@@ -0,0 +1,248 @@
+name: neuralom
+channels:
+  - pytorch
+  - dglteam/label/th24_cu118
+  - nvidia
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - blas=1.0=mkl
+  - brotli-python=1.0.9=py310h6a678d5_8
+  - bzip2=1.0.8=h5eee18b_6
+  - ca-certificates=2024.9.24=h06a4308_0
+  - certifi=2024.8.30=py310h06a4308_0
+  - charset-normalizer=3.3.2=pyhd3eb1b0_0
+  - cuda-cudart=11.8.89=0
+  - cuda-cupti=11.8.87=0
+  - cuda-libraries=11.8.0=0
+  - cuda-nvrtc=11.8.89=0
+  - cuda-nvtx=11.8.86=0
+  - cuda-runtime=11.8.0=0
+  - dgl=2.4.0.th24.cu118=py310_0
+  - ffmpeg=4.3=hf484d3e_0
+  - filelock=3.13.1=py310h06a4308_0
+  - freetype=2.12.1=h4a9f257_0
+  - gmp=6.2.1=h295c915_3
+  - gmpy2=2.1.2=py310heeb90bb_0
+  - gnutls=3.6.15=he1e5248_0
+  - idna=3.7=py310h06a4308_0
+  - intel-openmp=2023.1.0=hdb19cb5_46306
+  - jinja2=3.1.4=py310h06a4308_0
+  - jpeg=9e=h5eee18b_3
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libcublas=11.11.3.6=0
+  - libcufft=10.9.0.58=0
+  - libcufile=1.9.1.3=0
+  - libcurand=10.3.5.147=0
+  - libcusolver=11.4.1.48=0
+  - libcusparse=11.7.5.86=0
+  - libdeflate=1.17=h5eee18b_1
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgfortran-ng=11.2.0=h00389a5_1
+  - libgfortran5=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.16=h5eee18b_3
+  - libidn2=2.3.4=h5eee18b_0
+  - libjpeg-turbo=2.0.0=h9bf148f_0
+  - libnpp=11.8.0.86=0
+  - libnvjpeg=11.9.0.86=0
+  - libpng=1.6.39=h5eee18b_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libtasn1=4.19.0=h5eee18b_0
+  - libtiff=4.5.1=h6a678d5_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libwebp-base=1.3.2=h5eee18b_0
+  - llvm-openmp=14.0.6=h9e868ea_0
+  - lz4-c=1.9.4=h6a678d5_1
+  - markupsafe=2.1.3=py310h5eee18b_0
+  - mkl=2023.1.0=h213fc3f_46344
+  - mkl-service=2.4.0=py310h5eee18b_1
+  - mkl_fft=1.3.10=py310h5eee18b_0
+  - mkl_random=1.2.7=py310h1128e8f_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.2=hb69a4c5_1
+  - mpmath=1.3.0=py310h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nettle=3.7.3=hbbd107a_1
+  - networkx=3.3=py310h06a4308_0
+  - numpy=1.26.4=py310h5f9d8c6_0
+  - numpy-base=1.26.4=py310hb5e798b_0
+  - openh264=2.1.1=h4ff587b_0
+  - openjpeg=2.5.2=he7f1fd0_0
+  - openssl=3.0.15=h5eee18b_0
+  - pillow=10.4.0=py310h5eee18b_0
+  - pip=24.2=py310h06a4308_0
+  - psutil=5.9.0=py310h5eee18b_0
+  - pybind11-abi=4=hd3eb1b0_1
+  - pysocks=1.7.1=py310h06a4308_0
+  - python=3.10.15=he870216_1
+  - pytorch=2.4.0=py3.10_cuda11.8_cudnn9.1.0_0
+  - pytorch-cuda=11.8=h7e8668a_5
+  - pytorch-mutex=1.0=cuda
+  - pyyaml=6.0.1=py310h5eee18b_0
+  - readline=8.2=h5eee18b_0
+  - requests=2.32.3=py310h06a4308_0
+  - scipy=1.13.1=py310h5f9d8c6_0
+  - setuptools=75.1.0=py310h06a4308_0
+  - sqlite=3.45.3=h5eee18b_0
+  - sympy=1.13.2=py310h06a4308_0
+  - tbb=2021.8.0=hdb19cb5_0
+  - tk=8.6.14=h39e8969_0
+  - torchaudio=2.4.0=py310_cu118
+  - torchtriton=3.0.0=py310
+  - torchvision=0.19.0=py310_cu118
+  - tqdm=4.66.5=py310h2f386ee_0
+  - typing_extensions=4.11.0=py310h06a4308_0
+  - urllib3=2.2.3=py310h06a4308_0
+  - wheel=0.44.0=py310h06a4308_0
+  - xz=5.4.6=h5eee18b_1
+  - yaml=0.2.5=h7b6447c_0
+  - zlib=1.2.13=h5eee18b_1
+  - zstd=1.5.5=hc292b87_2
+  - pip:
+      - aiobotocore==2.15.1
+      - aiohappyeyeballs==2.4.3
+      - aiohttp==3.10.8
+      - aioitertools==0.12.0
+      - aiosignal==1.3.1
+      - anyio==4.6.0
+      - argon2-cffi==23.1.0
+      - argon2-cffi-bindings==21.2.0
+      - arrow==1.3.0
+      - asttokens==2.4.1
+      - async-lru==2.0.4
+      - async-timeout==4.0.3
+      - attrs==24.2.0
+      - babel==2.16.0
+      - beautifulsoup4==4.12.3
+      - bleach==6.1.0
+      - blessed==1.20.0
+      - botocore==1.35.23
+      - cartopy==0.24.1
+      - cffi==1.17.1
+      - cftime==1.6.4.post1
+      - cmocean==4.0.3
+      - colorama==0.4.6
+      - comm==0.2.2
+      - contourpy==1.3.0
+      - cycler==0.12.1
+      - debugpy==1.8.6
+      - decorator==5.1.1
+      - defusedxml==0.7.1
+      - einops==0.8.0
+      - exceptiongroup==1.2.2
+      - executing==2.1.0
+      - fastjsonschema==2.20.0
+      - fonttools==4.54.1
+      - fqdn==1.5.1
+      - frozenlist==1.4.1
+      - fsspec==2024.9.0
+      - gpustat==1.1.1
+      - h11==0.14.0
+      - h5netcdf==1.4.0
+      - h5py==3.12.1
+      - httpcore==1.0.6
+      - httpx==0.27.2
+      - huggingface-hub==0.25.1
+      - icecream==2.1.3
+      - importlib-metadata==8.5.0
+      - ipykernel==6.29.5
+      - ipython==8.28.0
+      - isoduration==20.11.0
+      - jedi==0.19.1
+      - jmespath==1.0.1
+      - joblib==1.4.2
+      - json5==0.9.25
+      - jsonpointer==3.0.0
+      - jsonschema==4.23.0
+      - jsonschema-specifications==2024.10.1
+      - jupyter-client==8.6.3
+      - jupyter-core==5.7.2
+      - jupyter-events==0.10.0
+      - jupyter-lsp==2.2.5
+      - jupyter-server==2.14.2
+      - jupyter-server-terminals==0.5.3
+      - jupyterlab==4.2.5
+      - jupyterlab-pygments==0.3.0
+      - jupyterlab-server==2.27.3
+      - kiwisolver==1.4.7
+      - matplotlib==3.9.2
+      - matplotlib-inline==0.1.7
+      - mistune==3.0.2
+      - multidict==6.1.0
+      - nbclient==0.10.0
+      - nbconvert==7.16.4
+      - nbformat==5.10.4
+      - nest-asyncio==1.6.0
+      - netcdf4==1.7.2
+      - notebook==7.2.2
+      - notebook-shim==0.2.4
+      - nvfuser-cu118-torch24==0.2.9.dev20240808
+      - nvidia-cuda-cupti-cu11==11.8.87
+      - nvidia-cuda-nvrtc-cu11==11.8.89
+      - nvidia-cuda-runtime-cu11==11.8.89
+      - nvidia-ml-py==12.560.30
+      - nvidia-nvtx-cu11==11.8.86
+      - overrides==7.7.0
+      - packaging==24.1
+      - pandas==2.2.3
+      - pandocfilters==1.5.1
+      - parso==0.8.4
+      - pexpect==4.9.0
+      - platformdirs==4.3.6
+      - prometheus-client==0.21.0
+      - prompt-toolkit==3.0.48
+      - ptyprocess==0.7.0
+      - pure-eval==0.2.3
+      - pycparser==2.22
+      - pygments==2.18.0
+      - pyparsing==3.2.0
+      - pyproj==3.7.0
+      - pyshp==2.3.1
+      - python-dateutil==2.9.0.post0
+      - python-json-logger==2.0.7
+      - pytz==2024.2
+      - pyzmq==26.2.0
+      - referencing==0.35.1
+      - rfc3339-validator==0.1.4
+      - rfc3986-validator==0.1.1
+      - rpds-py==0.20.0
+      - ruamel-yaml==0.18.6
+      - ruamel-yaml-clib==0.2.8
+      - s3fs==2024.9.0
+      - safetensors==0.4.5
+      - scikit-learn==1.5.2
+      - send2trash==1.8.3
+      - shapely==2.0.6
+      - six==1.16.0
+      - sniffio==1.3.1
+      - soupsieve==2.6
+      - stack-data==0.6.3
+      - terminado==0.18.1
+      - thop==0.1.1-2209072238
+      - threadpoolctl==3.5.0
+      - timm==1.0.9
+      - tinycss2==1.3.0
+      - tomli==2.0.2
+      - torchsummary==1.5.1
+      - tornado==6.4.1
+      - traitlets==5.14.3
+      - treelib==1.7.0
+      - types-python-dateutil==2.9.0.20241003
+      - tzdata==2024.2
+      - uri-template==1.3.0
+      - wcwidth==0.2.13
+      - webcolors==24.8.0
+      - webencodings==0.5.1
+      - websocket-client==1.8.0
+      - wrapt==1.16.0
+      - xarray==2024.9.0
+      - yarl==1.13.1
+      - zipp==3.20.2
+prefix: /miniconda3/envs/neuralom

exp/NeuralOM/20250309-195251/.ipynb_checkpoints/config-checkpoint.yaml

@@ -0,0 +1,70 @@
+### base config ###
+# -*- coding: utf-8 -*-
+full_field: &FULL_FIELD
+  num_data_workers: 4
+  dt: 1                     
+  n_history: 0              
+  prediction_length: 41     
+  ics_type: "default"       
+
+  exp_dir: './exp'
+
+  # data
+  train_data_path: './data/train'
+  valid_data_path: './data/valid'  
+  test_data_path:  './data/test'
+
+  # land mask
+  land_mask: !!bool True
+  land_mask_path: './data/land_mask.h5'
+
+  # normalization
+  normalize: !!bool True
+  normalization: 'zscore'
+  global_means_path: './data/mean_s_t_ssh.npy' 
+  global_stds_path:  './data/std_s_t_ssh.npy'
+
+  # orography
+  orography: !!bool False
+
+  # noise
+  add_noise: !!bool False
+  noise_std: 0
+
+  # crop
+  crop_size_x: None
+  crop_size_y: None
+
+  log_to_screen: !!bool True
+  log_to_wandb: !!bool False
+  save_checkpoint: !!bool True
+  plot_animations: !!bool False
+
+
+#############################################
+NeuralOM: &NeuralOM
+  <<: *FULL_FIELD
+  nettype: 'NeuralOM'
+  log_to_wandb: !!bool False
+
+  # Train params
+  lr: 1e-3
+  batch_size: 32 
+  scheduler: 'CosineAnnealingLR'
+
+  loss_channel_wise: True
+  loss_scale: False 
+  use_loss_scaler_from_metnet3: True
+
+
+  atmos_channels: [93, 94, 95, 96]
+
+  ocean_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92]
+
+  in_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96]
+
+  out_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92]
+
+
+  out_variables: ["S0", "S2", "S5", "S7", "S11", "S15", "S21", "S29", "S40", "S55", "S77", "S92", "S109", "S130", "S155", "S186", "S222", "S266", "S318", "S380", "S453", "S541", "S643", "U0", "U2", "U5", "U7", "U11", "U15", "U21", "U29", "U40", "U55", "U77", "U92", "U109", "U130", "U155", "U186", "U222", "U266", "U318", "U380", "U453", "U541", "U643", "V0", "V2", "V5", "V7", "V11", "V15", "V21", "V29", "V40", "V55", "V77", "V92", "V109", "V130", "V155", "V186", "V222", "V266", "V318", "V380", "V453", "V541", "V643", "T0", "T2", "T5", "T7", "T11", "T15", "T21", "T29", "T40", "T55", "T77", "T92", "T109", "T130", "T155", "T186", "T222", "T266", "T318", "T380", "T453", "T541", "T643", "SSH"]
+  

exp/NeuralOM/20250309-195251/6_steps_finetune/model2/10_steps_finetune/training_checkpoints/best_ckpt.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9fe78a12419997b9deaf0dd2ec912c1c936e96cc418bc507acdd2baecf908a2
+size 661771939

exp/NeuralOM/20250309-195251/6_steps_finetune/model2/10_steps_finetune/training_checkpoints/readme.txt

@@ -0,0 +1 @@
+The intact project is available at the Hugging Face.

exp/NeuralOM/20250309-195251/6_steps_finetune/training_checkpoints/best_ckpt.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8fb1c1827478608c96b490662f9b59351a52b1ee423883158c0a9e7c09b7d919
+size 661813411

exp/NeuralOM/20250309-195251/6_steps_finetune/training_checkpoints/readme.txt

@@ -0,0 +1 @@
+The intact project is available at the Hugging Face.

exp/NeuralOM/20250309-195251/config.yaml

@@ -0,0 +1,70 @@
+### base config ###
+# -*- coding: utf-8 -*-
+full_field: &FULL_FIELD
+  num_data_workers: 4
+  dt: 1                     
+  n_history: 0              
+  prediction_length: 41     
+  ics_type: "default"       
+
+  exp_dir: './exp'
+
+  # data
+  train_data_path: './data/train'
+  valid_data_path: './data/valid'  
+  test_data_path:  './data/test'
+
+  # land mask
+  land_mask: !!bool True
+  land_mask_path: './data/land_mask.h5'
+
+  # normalization
+  normalize: !!bool True
+  normalization: 'zscore'
+  global_means_path: './data/mean_s_t_ssh.npy' 
+  global_stds_path:  './data/std_s_t_ssh.npy'
+
+  # orography
+  orography: !!bool False
+
+  # noise
+  add_noise: !!bool False
+  noise_std: 0
+
+  # crop
+  crop_size_x: None
+  crop_size_y: None
+
+  log_to_screen: !!bool True
+  log_to_wandb: !!bool False
+  save_checkpoint: !!bool True
+  plot_animations: !!bool False
+
+
+#############################################
+NeuralOM: &NeuralOM
+  <<: *FULL_FIELD
+  nettype: 'NeuralOM'
+  log_to_wandb: !!bool False
+
+  # Train params
+  lr: 1e-3
+  batch_size: 32 
+  scheduler: 'CosineAnnealingLR'
+
+  loss_channel_wise: True
+  loss_scale: False 
+  use_loss_scaler_from_metnet3: True
+
+
+  atmos_channels: [93, 94, 95, 96]
+
+  ocean_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92]
+
+  in_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96]
+
+  out_channels: [0, 1 ,2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20 ,21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92]
+
+
+  out_variables: ["S0", "S2", "S5", "S7", "S11", "S15", "S21", "S29", "S40", "S55", "S77", "S92", "S109", "S130", "S155", "S186", "S222", "S266", "S318", "S380", "S453", "S541", "S643", "U0", "U2", "U5", "U7", "U11", "U15", "U21", "U29", "U40", "U55", "U77", "U92", "U109", "U130", "U155", "U186", "U222", "U266", "U318", "U380", "U453", "U541", "U643", "V0", "V2", "V5", "V7", "V11", "V15", "V21", "V29", "V40", "V55", "V77", "V92", "V109", "V130", "V155", "V186", "V222", "V266", "V318", "V380", "V453", "V541", "V643", "T0", "T2", "T5", "T7", "T11", "T15", "T21", "T29", "T40", "T55", "T77", "T92", "T109", "T130", "T155", "T186", "T222", "T266", "T318", "T380", "T453", "T541", "T643", "SSH"]
+  

exp/baselines/readme.txt

@@ -0,0 +1 @@
+The intact project is available at the Hugging Face.

inference.py

@@ -0,0 +1,312 @@
+import os
+import sys
+import time
+import glob
+import h5py
+import logging
+import argparse
+import numpy as np
+from icecream import ic
+from datetime import datetime
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import torch.cuda.amp as amp
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+
+sys.path.append(os.path.dirname(os.path.realpath(__file__)) + '/../')
+from my_utils.YParams import YParams
+from my_utils.data_loader import get_data_loader
+from my_utils import logging_utils
+logging_utils.config_logger()
+
+
+def load_model(model, params, checkpoint_file):
+    model.zero_grad()
+    checkpoint_fname = checkpoint_file
+    checkpoint = torch.load(checkpoint_fname)
+    try:
+        new_state_dict = OrderedDict()
+        for key, val in checkpoint['model_state'].items():
+            name = key[7:]
+            if name != 'ged':
+                new_state_dict[name] = val  
+        model.load_state_dict(new_state_dict)
+    except:
+        model.load_state_dict(checkpoint['model_state'])
+    model.eval()
+    return model
+
+def setup(params):
+    device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
+
+    # get data loader
+    valid_data_loader, valid_dataset = get_data_loader(params, params.test_data_path, dist.is_initialized(), train=False)
+
+    img_shape_x = valid_dataset.img_shape_x
+    img_shape_y = valid_dataset.img_shape_y
+    params.img_shape_x = img_shape_x
+    params.img_shape_y = img_shape_y
+
+    in_channels = np.array(params.in_channels)
+    out_channels = np.array(params.out_channels)
+    n_in_channels = len(in_channels)
+    n_out_channels = len(out_channels)
+
+    params['N_in_channels'] = n_in_channels
+    params['N_out_channels'] = n_out_channels
+
+    if params.normalization == 'zscore': 
+        params.means = np.load(params.global_means_path)
+        params.stds = np.load(params.global_stds_path)
+
+    if params.nettype == 'NeuralOM':
+        from networks.MIGNN1 import MIGraph as model
+        from networks.MIGNN2 import MIGraph_stage2 as model2
+    else:
+        raise Exception("not implemented")
+
+    checkpoint_file  = params['best_checkpoint_path']
+    checkpoint_file2  = params['best_checkpoint_path2']
+    logging.info('Loading trained model checkpoint from {}'.format(checkpoint_file))
+    logging.info('Loading trained model2 checkpoint from {}'.format(checkpoint_file2))
+    
+    model = model(params).to(device) 
+    model = load_model(model, params, checkpoint_file)
+    model = model.to(device)
+
+    print('model is ok')
+
+    model2 = model2(params).to(device) 
+    model2 = load_model(model2, params, checkpoint_file2)
+    model2 = model2.to(device)
+
+    print('model2 is ok')
+    
+    files_paths = glob.glob(params.test_data_path + "/*.h5")
+    files_paths.sort()
+
+    # which year
+    yr = 0
+    logging.info('Loading inference data')
+    logging.info('Inference data from {}'.format(files_paths[yr]))
+    climate_mean = np.load('./data/climate_mean_s_t_ssh.npy')
+    valid_data_full = h5py.File(files_paths[yr], 'r')['fields'][:365, :, :, :]
+    valid_data_full = valid_data_full - climate_mean
+
+    return valid_data_full, model, model2
+
+    
+def autoregressive_inference(params, init_condition, valid_data_full, model, model2): 
+    device = torch.cuda.current_device() if torch.cuda.is_available() else 'cpu'
+        
+    icd = int(init_condition) 
+    
+    exp_dir = params['experiment_dir'] 
+    dt                = int(params.dt)
+    prediction_length = int(params.prediction_length/dt)
+    n_history      = params.n_history
+    img_shape_x    = params.img_shape_x
+    img_shape_y    = params.img_shape_y
+    in_channels    = np.array(params.in_channels)
+    out_channels   = np.array(params.out_channels)
+    atmos_channels = np.array(params.atmos_channels)
+    n_in_channels  = len(in_channels)
+    n_out_channels = len(out_channels)
+    
+    seq_real        = torch.zeros((prediction_length, n_out_channels, img_shape_x, img_shape_y))
+    seq_pred        = torch.zeros((prediction_length, n_out_channels, img_shape_x, img_shape_y))
+
+
+    valid_data = valid_data_full[icd:(icd+prediction_length*dt+n_history*dt):dt][:, params.in_channels][:,:,0:360]
+    logging.info(f'valid_data_full: {valid_data_full.shape}')
+    logging.info(f'valid_data: {valid_data.shape}')
+    
+    # normalize
+    if params.normalization == 'zscore': 
+        valid_data = (valid_data - params.means[:,params.in_channels])/params.stds[:,params.in_channels]
+        valid_data = np.nan_to_num(valid_data, nan=0)
+        
+    valid_data = torch.as_tensor(valid_data)
+
+    # autoregressive inference
+    logging.info('Begin autoregressive inference')
+    
+    
+    with torch.no_grad():
+        for i in range(valid_data.shape[0]): 
+            if i==0: # start of sequence, t0 --> t0'
+                first = valid_data[0:n_history+1]
+                ic(valid_data.shape, first.shape)
+                future = valid_data[n_history+1]
+                ic(future.shape)
+
+                for h in range(n_history+1):
+                    
+                    seq_real[h] = first[h*n_in_channels : (h+1)*n_in_channels, :93]
+                    
+                    seq_pred[h] = seq_real[h]
+
+                first = first.to(device, dtype=torch.float)
+                first_ocean = first[:, params.ocean_channels, :, :]
+                ic(first_ocean.shape)
+                future_force0 = first[:, params.atmos_channels, :, :]
+               
+                future_force = future[params.atmos_channels, :360, :720]
+                future_force = torch.unsqueeze(future_force, dim=0).to(device, dtype=torch.float)
+                model_input = torch.cat((first_ocean, future_force0, future_force.cuda()), axis=1)
+                ic(model_input.shape)
+                model1_future_pred = model(model_input)
+                with h5py.File(params.land_mask_path, 'r') as _f: 
+                    mask_data = torch.as_tensor(_f['fields'][:,out_channels, :360, :720], dtype=bool).to(device, dtype=torch.bool)
+                model1_future_pred = torch.masked_fill(input=model1_future_pred, mask=~mask_data, value=0)
+                future_pred = model2(model1_future_pred) + model1_future_pred
+
+            else:
+                if i < prediction_length-1:
+                    future0 = valid_data[n_history+i]
+                    future = valid_data[n_history+i+1]
+
+                inf_one_step_start = time.time()
+                future_force0 = future0[params.atmos_channels, :360, :720]
+                future_force = future[params.atmos_channels, :360, :720]
+                future_force0 = torch.unsqueeze(future_force0, dim=0).to(device, dtype=torch.float)
+                future_force = torch.unsqueeze(future_force, dim=0).to(device, dtype=torch.float)
+                model1_future_pred = model(torch.cat((future_pred.cuda(), future_force0, future_force), axis=1)) #autoregressive step
+                with h5py.File(params.land_mask_path, 'r') as _f: 
+                    mask_data = torch.as_tensor(_f['fields'][:,out_channels, :360, :720], dtype=bool).to(device, dtype=torch.bool)
+                model1_future_pred = torch.masked_fill(input=model1_future_pred, mask=~mask_data, value=0)
+                future_pred = model2(model1_future_pred) + model1_future_pred
+                inf_one_step_time = time.time() - inf_one_step_start
+
+                logging.info(f'inference one step time: {inf_one_step_time}')
+
+            if i < prediction_length - 1: # not on the last step
+                with h5py.File(params.land_mask_path, 'r') as _f: 
+                    mask_data = torch.as_tensor(_f['fields'][:,out_channels, :360, :720], dtype=bool)
+                seq_pred[n_history+i+1] = torch.masked_fill(input=future_pred.cpu(), mask=~mask_data, value=0)
+                seq_real[n_history+i+1] = future[:93]
+                history_stack = seq_pred[i+1:i+2+n_history]
+
+            future_pred = history_stack
+
+            pred = torch.unsqueeze(seq_pred[i], 0)
+            tar  = torch.unsqueeze(seq_real[i], 0)
+
+            with h5py.File(params.land_mask_path, 'r') as _f: 
+                mask_data = torch.as_tensor(_f['fields'][:,out_channels, :360, :720], dtype=bool)
+                ic(mask_data.shape, pred.shape, tar.shape)
+            pred = torch.masked_fill(input=pred, mask=~mask_data, value=0)
+            tar  = torch.masked_fill(input=tar,  mask=~mask_data, value=0)
+
+            print(torch.mean((pred-tar)**2))
+
+    
+    seq_real = seq_real * params.stds[:,params.out_channels] + params.means[:,params.out_channels]
+    seq_real = seq_real.numpy()
+    seq_pred = seq_pred * params.stds[:,params.out_channels] + params.means[:,params.out_channels]
+    seq_pred = seq_pred.numpy()
+   
+
+    return (np.expand_dims(seq_real[n_history:], 0), 
+            np.expand_dims(seq_pred[n_history:], 0), 
+           )     
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--exp_dir", default='../exp_15_levels', type=str)
+    parser.add_argument("--config", default='full_field', type=str)
+    parser.add_argument("--run_num", default='00', type=str)
+    parser.add_argument("--prediction_length", default=61, type=int)
+    parser.add_argument("--finetune_dir", default='', type=str)
+    parser.add_argument("--ics_type", default='default', type=str)
+    args = parser.parse_args()
+
+    config_path = os.path.join(args.exp_dir, args.config, args.run_num, 'config.yaml')
+    params = YParams(config_path, args.config)
+
+    params['resuming']           = False
+    params['interp']             = 0 
+    params['world_size']         = 1
+    params['local_rank']         = 0
+    params['global_batch_size']  = params.batch_size
+    params['prediction_length']  = args.prediction_length
+    params['multi_steps_finetune'] = 1
+
+    torch.cuda.set_device(0)
+    torch.backends.cudnn.benchmark = True
+
+    # Set up directory
+    if args.finetune_dir == '':
+        expDir = os.path.join(params.exp_dir, args.config, str(args.run_num))
+    else:
+        expDir = os.path.join(params.exp_dir, args.config, str(args.run_num), args.finetune_dir)
+    logging.info(f'expDir: {expDir}')
+    params['experiment_dir']       = expDir 
+    params['best_checkpoint_path'] = os.path.join(expDir, 'training_checkpoints/best_ckpt.tar')
+    params['best_checkpoint_path2'] = os.path.join(expDir, 'model2/10_steps_finetune/training_checkpoints/best_ckpt.tar') 
+
+    # set up logging
+    logging_utils.log_to_file(logger_name=None, log_filename=os.path.join(expDir, 'inference.log'))
+    logging_utils.log_versions()
+    params.log()
+
+    if params["ics_type"] == 'default':
+        ics = np.arange(0, 240, 1)
+        n_ics = len(ics)
+        print('init_condition:', ics)
+
+    logging.info("Inference for {} initial conditions".format(n_ics))
+
+    try:
+      autoregressive_inference_filetag = params["inference_file_tag"]
+    except:
+      autoregressive_inference_filetag = ""
+    if params.interp > 0:
+        autoregressive_inference_filetag = "_coarse"
+
+    valid_data_full, model, model2 = setup(params)
+
+
+    seq_pred = []
+    seq_real = []
+
+    # run autoregressive inference for multiple initial conditions
+    for i, ic_ in enumerate(ics):
+        logging.info("Initial condition {} of {}".format(i+1, n_ics))
+        seq_real, seq_pred = autoregressive_inference(params, ic_, valid_data_full, model, model2)
+
+        prediction_length = seq_real[0].shape[0]
+        n_out_channels = seq_real[0].shape[1]
+        img_shape_x = seq_real[0].shape[2]
+        img_shape_y = seq_real[0].shape[3]
+
+        # save predictions and loss
+        save_path = os.path.join(params['experiment_dir'], 'results.h5')
+        logging.info("Saving to {}".format(save_path))
+        print(f'saving to {save_path}')
+        if i==0:
+            f = h5py.File(save_path, 'w')
+            f.create_dataset(
+                    "ground_truth",
+                    data=seq_real,
+                    maxshape=[None, prediction_length, n_out_channels, img_shape_x, img_shape_y], 
+                    dtype=np.float32)
+            f.create_dataset(
+                    "predicted",       
+                    data=seq_pred, 
+                    maxshape=[None, prediction_length, n_out_channels, img_shape_x, img_shape_y], 
+                    dtype=np.float32)
+            f.close()
+        else:
+            f = h5py.File(save_path, 'a')
+
+            f["ground_truth"].resize((f["ground_truth"].shape[0] + 1), axis = 0)
+            f["ground_truth"][-1:] = seq_real 
+
+            f["predicted"].resize((f["predicted"].shape[0] + 1), axis = 0)
+            f["predicted"][-1:] = seq_pred 
+            f.close()
+

inference.sh

@@ -0,0 +1,13 @@
+prediction_length=61 # 31
+
+exp_dir='./exp'
+config='NeuralOM' 
+run_num='20250309-195251'
+finetune_dir='6_steps_finetune'
+
+ics_type='default'
+
+CUDA_VISIBLE_DEVICES=0 python inference.py --exp_dir=${exp_dir} --config=${config} --run_num=${run_num} --finetune_dir=$finetune_dir --prediction_length=${prediction_length} --ics_type=${ics_type}
+
+
+

my_utils/.ipynb_checkpoints/YParams-checkpoint.py

@@ -0,0 +1,55 @@
+# coding: utf-8
+
+import importlib
+import sys
+import os
+importlib.reload(sys)
+
+from ruamel.yaml import YAML
+import logging
+
+class YParams():
+  """ Yaml file parser """
+  def __init__(self, yaml_filename, config_name, print_params=False):
+    self._yaml_filename = yaml_filename
+    self._config_name = config_name
+    self.params = {}
+
+    if print_params:
+      print(os.system('hostname'))
+      print("------------------ Configuration ------------------ ", yaml_filename)
+
+    with open(yaml_filename, 'rb') as _file:
+      yaml = YAML().load(_file)
+      for key, val in yaml[config_name].items():
+        if print_params: print(key, val)
+        if val =='None': val = None
+
+        self.params[key] = val
+        self.__setattr__(key, val)
+
+    if print_params:
+      print("---------------------------------------------------")
+
+  def __getitem__(self, key):
+    return self.params[key]
+
+  def __setitem__(self, key, val):
+    self.params[key] = val
+    self.__setattr__(key, val)
+
+  def __contains__(self, key):
+    return (key in self.params)
+
+  def update_params(self, config):
+    for key, val in config.items():
+      self.params[key] = val
+      self.__setattr__(key, val)
+
+  def log(self):
+    logging.info("------------------ Configuration ------------------")
+    logging.info("Configuration file: "+str(self._yaml_filename))
+    logging.info("Configuration name: "+str(self._config_name))
+    for key, val in self.params.items():
+        logging.info(str(key) + ' ' + str(val))
+    logging.info("---------------------------------------------------")

my_utils/.ipynb_checkpoints/data_loader-checkpoint.py

@@ -0,0 +1,205 @@
+import logging
+import glob
+import torch
+import random
+import numpy as np
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from torch import Tensor
+import h5py
+import math
+from my_utils.norm import reshape_fields
+import os
+
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+climate_mean_path = os.path.join(parent_dir, 'data/climate_mean_s_t_ssh.npy')
+
+def get_data_loader(params, files_pattern, distributed, train):
+    dataset = GetDataset(params, files_pattern, train)
+    sampler = DistributedSampler(dataset, shuffle=train) if distributed else None
+
+
+    dataloader = DataLoader(dataset,
+                            batch_size  = int(params.batch_size),
+                            num_workers = params.num_data_workers,
+                            shuffle     = False,
+                            sampler     = sampler if train else None,
+                            drop_last   = True,
+                            pin_memory  = True) 
+
+    if train:
+        return dataloader, dataset, sampler
+    else:
+        return dataloader, dataset
+
+
+class GetDataset(Dataset):
+    def __init__(self, params, location, train):
+        self.params = params
+        self.location = location
+        self.train = train
+        self.orography = params.orography
+        self.normalize = params.normalize
+        self.dt = params.dt 
+        self.n_history = params.n_history 
+        self.in_channels = np.array(params.in_channels)
+        self.out_channels = np.array(params.out_channels)
+        self.ocean_channels = np.array(params.ocean_channels)
+        self.atmos_channels = np.array(params.atmos_channels)
+        self.n_in_channels = len(self.in_channels)
+        self.n_out_channels = len(self.out_channels)
+
+        self._get_files_stats()
+        self.add_noise = params.add_noise if train else False
+        self.climate_mean = np.load(climate_mean_path, mmap_mode='r')
+
+
+    def _get_files_stats(self):
+        self.files_paths = glob.glob(self.location + "/*.h5")
+        self.files_paths.sort()
+        self.n_years = len(self.files_paths)
+
+        with h5py.File(self.files_paths[0], 'r') as _f: 
+            logging.info("Getting file stats from {}".format(self.files_paths[0]))
+
+            self.n_samples_per_year = _f['fields'].shape[0] - self.params.multi_steps_finetune 
+
+            self.img_shape_x = _f['fields'].shape[2] - 1 
+            self.img_shape_y = _f['fields'].shape[3]
+
+        self.n_samples_total = self.n_years * self.n_samples_per_year
+        self.files = [None for _ in range(self.n_years)]
+
+        logging.info("Number of samples per year: {}".format(self.n_samples_per_year))
+        logging.info("Found data at path {}. Number of examples: {}. Image Shape: {} x {} x {}".format(self.location,
+                                                                                                       self.n_samples_total,
+                                                                                                       self.img_shape_x,
+                                                                                                       self.img_shape_y,
+                                                                                                       self.n_in_channels))
+        logging.info("Delta t: {} days".format(1 * self.dt))
+        logging.info("Including {} days of past history in training at a frequency of {} days".format(
+            1 * self.dt * self.n_history, 1 * self.dt))
+
+    def _open_file(self, year_idx):
+        _file = h5py.File(self.files_paths[year_idx], 'r')
+        self.files[year_idx] = _file['fields'] 
+
+        if self.orography and self.params.normalization == 'zscore': 
+            _orog_file = h5py.File(self.params.orography_norm_zscore_path, 'r')
+        if self.orography and self.params.normalization == 'maxmin': 
+            _orog_file = h5py.File(self.params.orography_norm_maxmin_path, 'r')
+
+    def __len__(self):
+        return self.n_samples_total
+
+    def __getitem__(self, global_idx):
+        year_idx  = int(global_idx / self.n_samples_per_year)  # which year
+        local_idx = int(global_idx % self.n_samples_per_year)  # which sample in a year
+
+        if self.files[year_idx] is None:
+            self._open_file(year_idx)
+
+        if local_idx < self.dt * self.n_history:
+            local_idx += self.dt * self.n_history
+
+        step = 0 if local_idx >= self.n_samples_per_year - self.dt else self.dt
+        
+        orog = None
+    
+
+        if self.params.multi_steps_finetune == 1:
+            if local_idx == 365:
+                local_idx = 364
+            
+            climate_mean_ocean = self.climate_mean[(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.ocean_channels, :360, :720]
+            ocean = reshape_fields( 
+                    self.files[year_idx][(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.ocean_channels, :360, :720] - climate_mean_ocean, 
+                    'ocean', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog, 
+                    self.add_noise 
+                )
+
+            force_future0 = reshape_fields( 
+                    self.files[year_idx][local_idx, self.atmos_channels, :360, :720], 
+                    'force', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog, 
+                    self.add_noise 
+                )
+            
+            force_future1 = reshape_fields( 
+                    self.files[year_idx][local_idx+step, self.atmos_channels, :360, :720], 
+                    'force', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog, 
+                    self.add_noise 
+                )
+
+            climate_mean_tar = self.climate_mean[local_idx+step, self.out_channels, :360, :720]
+            tar = reshape_fields(
+                    self.files[year_idx][local_idx+step, self.out_channels, :360, :720] - climate_mean_tar, 
+                    'tar', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog 
+                )
+        else:
+            climate_mean_ocean = self.climate_mean[(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.ocean_channels, :360, :720]
+            ocean = reshape_fields( 
+                    self.files[year_idx][(local_idx-self.dt*self.n_history):(local_idx+1):self.dt, self.ocean_channels, :360, :720] - climate_mean_ocean, 
+                    'ocean', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog, 
+                    self.add_noise 
+                )
+            
+            force_future0 = reshape_fields( 
+                    self.files[year_idx][local_idx, self.atmos_channels, :360, :720], 
+                    'force', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog, 
+                    self.add_noise 
+                )
+            
+            force_future1 = reshape_fields( 
+                    self.files[year_idx][local_idx+step, self.atmos_channels, :360, :720], 
+                    'force', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog, 
+                    self.add_noise 
+                )
+
+            climate_mean_tar = self.climate_mean[local_idx+step:local_idx+step+self.params.multi_steps_finetune, self.in_channels, :360, :720]
+            tar_data = self.files[year_idx][local_idx+step:local_idx+step+self.params.multi_steps_finetune, self.in_channels, :360, :720]
+            tar = reshape_fields( 
+                    tar_data - climate_mean_tar, 
+                    'inp', 
+                    self.params, 
+                    self.train, 
+                    self.normalize, 
+                    orog 
+                )
+            
+        ocean = np.nan_to_num(ocean, nan=0)
+        force_future0 = np.nan_to_num(force_future0, nan=0)
+        force_future1 = np.nan_to_num(force_future1, nan=0)
+        tar = np.nan_to_num(tar, nan=0)
+
+
+        return np.concatenate((ocean, force_future0, force_future1), axis=0), tar 

my_utils/.ipynb_checkpoints/logging_utils-checkpoint.py

@@ -0,0 +1,26 @@
+import os
+import logging
+
+_format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+
+def config_logger(log_level=logging.INFO):
+  logging.basicConfig(format=_format, level=log_level)
+
+def log_to_file(logger_name=None, log_level=logging.INFO, log_filename='tensorflow.log'):
+
+  if not os.path.exists(os.path.dirname(log_filename)):
+    os.makedirs(os.path.dirname(log_filename))
+
+  if logger_name is not None:
+    log = logging.getLogger(logger_name)
+  else:
+    log = logging.getLogger()
+
+  fh = logging.FileHandler(log_filename)
+  fh.setLevel(log_level)
+  fh.setFormatter(logging.Formatter(_format))
+  log.addHandler(fh)
+
+def log_versions():
+  import torch
+  import subprocess

my_utils/.ipynb_checkpoints/norm-checkpoint.py

@@ -0,0 +1,114 @@
+import logging
+import glob
+from types import new_class
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import random
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+from torch.utils.data.distributed import DistributedSampler
+from torch import Tensor
+import h5py
+import math
+import torchvision.transforms.functional as TF
+# import matplotlib
+# import matplotlib.pyplot as plt
+
+class PeriodicPad2d(nn.Module):
+    """ 
+        pad longitudinal (left-right) circular 
+        and pad latitude (top-bottom) with zeros
+    """
+    def __init__(self, pad_width):
+       super(PeriodicPad2d, self).__init__()
+       self.pad_width = pad_width
+
+    def forward(self, x):
+        # pad left and right circular
+        out = F.pad(x, (self.pad_width, self.pad_width, 0, 0), mode="circular") 
+        # pad top and bottom zeros
+        out = F.pad(out, (0, 0, self.pad_width, self.pad_width), mode="constant", value=0) 
+        return out
+
+def reshape_fields(img, inp_or_tar, params, train, normalize=True, orog=None, add_noise=False):
+    # Takes in np array of size (n_history+1, c, h, w) 
+    # returns torch tensor of size ((n_channels*(n_history+1), crop_size_x, crop_size_y)
+
+    if len(np.shape(img)) == 3:
+        img = np.expand_dims(img, 0)
+    
+    if np.shape(img)[2] == 721:
+        img = img[:,:, 0:720, :] # remove last pixel
+
+    n_history = np.shape(img)[0] - 1
+    img_shape_x = np.shape(img)[-2]
+    img_shape_y = np.shape(img)[-1]
+    n_channels = np.shape(img)[1] # this will either be N_in_channels or N_out_channels
+    
+    if inp_or_tar == 'inp':
+        channels = params.in_channels
+    elif inp_or_tar == 'ocean':
+        channels = params.ocean_channels
+    elif inp_or_tar == 'force':
+        channels = params.atmos_channels
+    else:
+        channels = params.out_channels
+
+    if normalize and params.normalization == 'minmax':
+        maxs = np.load(params.global_maxs_path)[:, channels]
+        mins = np.load(params.global_mins_path)[:, channels]
+        img = (img - mins) / (maxs - mins)
+
+    if normalize and params.normalization == 'zscore':
+        means = np.load(params.global_means_path)[:, channels]
+        stds = np.load(params.global_stds_path)[:, channels]
+        img -=means
+        img /=stds
+
+    if normalize and params.normalization == 'zscore_lat':
+        means = np.load(params.global_lat_means_path)[:, channels,:720]
+        stds = np.load(params.global_lat_stds_path)[:, channels,:720]
+        img -=means
+        img /=stds
+
+    if params.orography and inp_or_tar == 'inp':
+        # print('img:', img.shape, 'orog:', orog.shape)
+        orog = np.expand_dims(orog, axis = (0,1))
+        orog = np.repeat(orog, repeats=img.shape[0], axis=0)
+        # print('img:', img.shape, 'orog:', orog.shape)
+        img = np.concatenate((img, orog), axis = 1)
+        n_channels += 1
+
+    img = np.squeeze(img)
+    # if inp_or_tar == 'inp':
+    #     img = np.reshape(img, (n_channels*(n_history+1))) # ??
+    # elif inp_or_tar == 'tar':
+    #     img = np.reshape(img, (n_channels, crop_size_x, crop_size_y)) #??
+
+    if add_noise:
+        img = img + np.random.normal(0, scale=params.noise_std, size=img.shape)
+
+    return torch.as_tensor(img)
+
+def vis_precip(fields):
+    pred, tar = fields
+    fig, ax = plt.subplots(1, 2, figsize=(24,12))
+    ax[0].imshow(pred, cmap="coolwarm")
+    ax[0].set_title("tp pred")
+    ax[1].imshow(tar, cmap="coolwarm")
+    ax[1].set_title("tp tar")
+    fig.tight_layout()
+    return fig
+
+def read_max_min_value(min_max_val_file_path):
+    with h5py.File(min_max_val_file_path, 'r') as f:
+        max_values = f['max_values']
+        min_values = f['min_values']
+    return max_values, min_values
+    
+    
+
+
+

my_utils/YParams.py

@@ -0,0 +1,55 @@
+# coding: utf-8
+
+import importlib
+import sys
+import os
+importlib.reload(sys)
+
+from ruamel.yaml import YAML
+import logging
+
+class YParams():
+  """ Yaml file parser """
+  def __init__(self, yaml_filename, config_name, print_params=False):
+    self._yaml_filename = yaml_filename
+    self._config_name = config_name
+    self.params = {}
+
+    if print_params:
+      print(os.system('hostname'))
+      print("------------------ Configuration ------------------ ", yaml_filename)
+
+    with open(yaml_filename, 'rb') as _file:
+      yaml = YAML().load(_file)
+      for key, val in yaml[config_name].items():
+        if print_params: print(key, val)
+        if val =='None': val = None
+
+        self.params[key] = val
+        self.__setattr__(key, val)
+
+    if print_params:
+      print("---------------------------------------------------")
+
+  def __getitem__(self, key):
+    return self.params[key]
+
+  def __setitem__(self, key, val):
+    self.params[key] = val
+    self.__setattr__(key, val)
+
+  def __contains__(self, key):
+    return (key in self.params)
+
+  def update_params(self, config):
+    for key, val in config.items():
+      self.params[key] = val
+      self.__setattr__(key, val)
+
+  def log(self):
+    logging.info("------------------ Configuration ------------------")
+    logging.info("Configuration file: "+str(self._yaml_filename))
+    logging.info("Configuration name: "+str(self._config_name))
+    for key, val in self.params.items():
+        logging.info(str(key) + ' ' + str(val))
+    logging.info("---------------------------------------------------")