Upload 90 files

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/6_steps_finetune/model2/10_steps_finetune/training_checkpoints/.ipynb_checkpoints/readme-checkpoint.txt

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

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/.ipynb_checkpoints/readme-checkpoint.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/6_steps_finetune/training_checkpoints_atmos/best_ckpt.tar

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

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

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

exp/NeuralOM/20250309-195251/config.yaml

@@ -0,0 +1,78 @@
+### 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'
+  test_data_path_atmos:  './data/test_atmos'
+
+  # 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'
+
+  global_means_path_atmos: './data/mean_atmos.npy' 
+  global_stds_path_atmos:  './data/std_atmos.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]
+
+  in_channels_atmos: [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]
+
+  out_channels_atmos: [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]  
+
+
+  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"]
+  

inference_forecasting.py

@@ -0,0 +1,366 @@
+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.nn.functional as F
+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)
+
+        params.means_atmos = np.load(params.global_means_path_atmos)
+        params.stds_atmos = np.load(params.global_stds_path_atmos)
+
+    if params.nettype == 'NeuralOM':
+        from networks.MIGNN1 import MIGraph as model
+        from networks.MIGNN2 import MIGraph_stage2 as model2
+        from networks.OneForecast import OneForecast as model_atmos
+    else:
+        raise Exception("not implemented")
+
+    checkpoint_file  = params['best_checkpoint_path']
+    checkpoint_file2  = params['best_checkpoint_path2']
+    checkpoint_file_atmos = params['best_checkpoint_path_atmos']
+    logging.info('Loading trained model checkpoint from {}'.format(checkpoint_file))
+    logging.info('Loading trained model2 checkpoint from {}'.format(checkpoint_file2))
+    logging.info('Loading trained model_atmos checkpoint from {}'.format(checkpoint_file_atmos))
+    
+    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')
+
+    model_atmos = model_atmos(params).to(device) 
+    model_atmos = load_model(model_atmos, params, checkpoint_file_atmos)
+    model_atmos = model_atmos.to(device)
+
+    print('model_atmos is ok')
+    
+    files_paths = glob.glob(params.test_data_path + "/*.h5")
+    files_paths.sort()
+
+    files_paths_atmos = glob.glob(params.test_data_path_atmos + "/*.h5")
+    files_paths_atmos.sort()
+
+    # which year
+    yr = 0
+    logging.info('Loading inference data')
+    logging.info('Inference data from {}'.format(files_paths[yr]))
+    logging.info('Inference data_atmos from {}'.format(files_paths_atmos[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
+
+    valid_data_full_atmos = h5py.File(files_paths_atmos[yr], 'r')['fields'][2:1460:4, :, :, :]
+
+    return valid_data_full, valid_data_full_atmos, model, model2, model_atmos
+
+    
+def autoregressive_inference(params, init_condition, valid_data_full, valid_data_full_atmos, model, model2, model_atmos): 
+    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)
+    in_channels_atmos    = np.array(params.in_channels_atmos)
+    out_channels_atmos   = np.array(params.out_channels_atmos)
+    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]
+    valid_data_atmos = valid_data_full_atmos[icd:(icd+prediction_length*dt+n_history*dt):dt][:, params.in_channels_atmos][:,:,0:120]
+    logging.info(f'valid_data_full: {valid_data_full.shape}')
+    logging.info(f'valid_data: {valid_data.shape}')
+    logging.info(f'valid_data_full_atmos: {valid_data_full_atmos.shape}')
+    logging.info(f'valid_data_atmos: {valid_data_atmos.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_atmos = (valid_data_atmos - params.means_atmos[:,params.in_channels_atmos])/params.stds_atmos[:,params.in_channels_atmos]
+        valid_data_atmos = np.nan_to_num(valid_data_atmos, nan=0)
+        
+    valid_data = torch.as_tensor(valid_data)
+    valid_data_atmos = torch.as_tensor(valid_data_atmos)
+
+    # 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]
+                first_atmos = valid_data_atmos[0:n_history+1]
+                ic(valid_data.shape, first.shape)
+                ic(valid_data_atmos.shape, first_atmos.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_atmos = first_atmos.to(device, dtype=torch.float)
+                first_ocean = first[:, params.ocean_channels, :, :]
+                ic(first_ocean.shape)
+                future_force0 = first_atmos[:, [65, 66, 67, 68], :120, :240]
+                # future_force0 = torch.unsqueeze(future_force0, dim=0).to(device, dtype=torch.float)
+                future_force0 = F.interpolate(future_force0, size=(360, 720), mode='bilinear', align_corners=False)
+
+                model_input_atmos = first_atmos
+                ic(model_input_atmos.shape)
+                for k in range(4):
+                    if k ==0:   
+                        model_atmos_future_pred = model_atmos(model_input_atmos)
+                    else:
+                        model_atmos_future_pred = model_atmos(model_atmos_future_pred)
+               
+                future_force = model_atmos_future_pred[:, [65, 66, 67, 68], :120, :240]
+                # future_force = torch.unsqueeze(future_force, dim=0).to(device, dtype=torch.float)
+                future_force = F.interpolate(future_force, size=(360, 720), mode='bilinear', align_corners=False)
+                
+                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 = model_atmos_future_pred[:, [65, 66, 67, 68], :120, :240]
+                # future_force0 = torch.unsqueeze(future_force0, dim=0).to(device, dtype=torch.float)
+                future_force0 = F.interpolate(future_force0, size=(360, 720), mode='bilinear', align_corners=False)
+
+                for k in range(4):
+                    model_atmos_future_pred = model_atmos(model_atmos_future_pred)
+                        
+                future_force = model_atmos_future_pred[:, [65, 66, 67, 68], :120, :240]
+                # future_force = torch.unsqueeze(future_force, dim=0).to(device, dtype=torch.float)
+                future_force = F.interpolate(future_force, size=(360, 720), mode='bilinear', align_corners=False)
+                
+                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')
+
+    params['best_checkpoint_path_atmos'] = os.path.join(expDir, 'training_checkpoints_atmos/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, 50, 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, valid_data_full_atmos, model, model2, model_atmos = 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, valid_data_full_atmos, model, model2, model_atmos)
+
+        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_forecasting.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_forecasting.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=2 python inference_forecasting.py --exp_dir=${exp_dir} --config=${config} --run_num=${run_num} --finetune_dir=$finetune_dir --prediction_length=${prediction_length} --ics_type=${ics_type}
+
+
+

inference_simulation.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_simulation.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_simulation.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_simulation.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/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("---------------------------------------------------")

my_utils/data_loader.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/logging_utils.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/norm.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
+    
+    
+
+
+

networks/.ipynb_checkpoints/CirT1-checkpoint.py

@@ -0,0 +1,301 @@
+from functools import lru_cache
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import trunc_normal_, Block
+from torch.jit import Final
+
+import torch.nn.functional as F
+from typing import Optional
+from timm.layers import DropPath, use_fused_attn, Mlp
+
+class PatchEmbed(nn.Module):
+    def __init__(
+            self,
+            img_size=[121, 240],
+            in_chans=63,
+            embed_dim=768,
+            norm_layer=None,
+            flatten=True,
+            bias=True,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.num_patches = img_size[0]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=[1, img_size[1]], stride=1, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // self.num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+        self.dim = dim
+        self.attn_bias = nn.Parameter(torch.zeros(121, 121, 2), requires_grad=True)
+
+        # self.qkv = CLinear(dim, dim * 3, bias=qkv_bias)
+        self.qkv = nn.Linear((dim // 2 + 1) * 2, dim * 3, bias=qkv_bias)
+        self.q = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.k = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.v = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+
+
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        # self.proj = CLinear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.proj = nn.Linear(dim, dim)
+        # self.proj_drop = nn.Dropout(proj_drop)
+
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+
+        x = torch.fft.rfft(x, norm="forward")
+        x = torch.view_as_real(x)
+        x = torch.cat((x[:, :, :, 0], -x[:, :, :, 1]), dim=-1)
+
+        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        # q, k, v = qkv.unbind(0)
+        q = self.q(x).reshape(B, self.num_heads, N, self.head_dim)
+        k = self.k(x).reshape(B, self.num_heads, N, self.head_dim)
+        v = self.v(x).reshape(B, self.num_heads, N, self.head_dim)
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1) 
+
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = attn @ v
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        real, img = torch.split(x, x.shape[-1] // 2, dim=-1)
+        x = torch.stack([real,-img], dim=-1)
+        x = torch.view_as_complex(x)
+        x = torch.fft.irfft(x, self.dim, norm="forward")
+
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+
+class Block(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int,
+            mlp_ratio: float = 4.,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            proj_drop: float = 0.,
+            attn_drop: float = 0.,
+            init_values: Optional[float] = None,
+            drop_path: float = 0.,
+            act_layer: nn.Module = nn.GELU,
+            norm_layer: nn.Module = nn.LayerNorm,
+            mlp_layer: nn.Module = Mlp,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = mlp_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=proj_drop,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+class CirT(nn.Module):
+    def __init__(
+        self,
+        params,
+        img_size=[360, 720],
+        input_size=101,
+        output_size=93,
+        patch_size=124, #124
+        embed_dim=256,
+        depth=8,
+        decoder_depth=2,
+        num_heads=16,
+        mlp_ratio=4.0,
+        drop_path=0.1,
+        drop_rate=0.1
+    ):
+        super().__init__()
+
+        # TODO: remove time_history parameter
+        self.img_size = img_size
+        self.patch_size = img_size[1]
+        self.input_size = input_size
+        self.output_size = output_size
+        self.token_embeds = PatchEmbed(img_size, input_size, embed_dim)
+        # self.token_embeds = nn.Linear(img_size[0] * 2, embed_dim)
+        self.num_patches = self.token_embeds.num_patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim), requires_grad=True)
+        # self.pos_embed = PosEmbed(embed_dim=embed_dim)
+        
+
+        # --------------------------------------------------------------------------
+
+        # ViT backbone
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = [x.item() for x in torch.linspace(0, drop_path, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    embed_dim,
+                    num_heads,
+                    mlp_ratio,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                    norm_layer=nn.LayerNorm,
+                    # drop=drop_rate,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+
+        # --------------------------------------------------------------------------
+
+        # prediction head
+        self.head = nn.ModuleList()
+        for _ in range(decoder_depth):
+            self.head.append(nn.Linear(embed_dim, embed_dim))
+            self.head.append(nn.GELU())
+        self.head.append(nn.Linear(embed_dim, output_size * self.img_size[1]))
+        self.head = nn.Sequential(*self.head)
+
+        # --------------------------------------------------------------------------
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # token embedding layer
+        w = self.token_embeds.proj.weight.data
+        trunc_normal_(w.view([w.shape[0], -1]), std=0.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+    def unpatchify(self, x: torch.Tensor, h=None, w=None):
+        """
+        x: (B, L, V * patch_size)
+        return imgs: (B, V, H, W)
+        """
+        p = self.patch_size
+        c_out = self.output_size
+        h = self.img_size[0] // 1
+        w = self.img_size[1] // p
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, c_out))
+        x = torch.einsum("nhwpc->nchpw", x)
+        imgs = x.reshape(shape=(x.shape[0], c_out, h, w * p))
+        return imgs
+
+    def forward_encoder(self, x: torch.Tensor):
+        # x: `[B, V, H, W]` shape.
+
+        # tokenize each variable separately
+        # x = torch.fft.rfft(x, norm="forward")
+        # x = torch.view_as_real(x)
+        # x = torch.cat((x[:, :, :, :, 0], -x[:, :, :, :, 1]), dim=-1)
+
+        x = self.token_embeds(x)
+
+        # pos_embed = self.pos_embed()
+        # add pos embedding
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x
+
+    def forward(self, x):
+        B, V, H, W = x.shape
+        # print(x.shape)
+        out_transformers = self.forward_encoder(x)  # B, L, D
+        preds = self.head(out_transformers)  # B, L, V*p*p
+        preds = self.unpatchify(preds)
+
+        # real, img = torch.split(preds, preds.shape[-1] // 2, dim=-1)
+        # preds = torch.cat([real, -img], dim=-1)
+        # preds = torch.fft.irfft(preds, W, norm="forward")
+        return preds

networks/.ipynb_checkpoints/CirT2-checkpoint.py

@@ -0,0 +1,301 @@
+from functools import lru_cache
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import trunc_normal_, Block
+from torch.jit import Final
+
+import torch.nn.functional as F
+from typing import Optional
+from timm.layers import DropPath, use_fused_attn, Mlp
+
+class PatchEmbed(nn.Module):
+    def __init__(
+            self,
+            img_size=[121, 240],
+            in_chans=63,
+            embed_dim=768,
+            norm_layer=None,
+            flatten=True,
+            bias=True,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.num_patches = img_size[0]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=[1, img_size[1]], stride=1, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // self.num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+        self.dim = dim
+        self.attn_bias = nn.Parameter(torch.zeros(121, 121, 2), requires_grad=True)
+
+        # self.qkv = CLinear(dim, dim * 3, bias=qkv_bias)
+        self.qkv = nn.Linear((dim // 2 + 1) * 2, dim * 3, bias=qkv_bias)
+        self.q = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.k = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.v = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+
+
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        # self.proj = CLinear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.proj = nn.Linear(dim, dim)
+        # self.proj_drop = nn.Dropout(proj_drop)
+
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+
+        x = torch.fft.rfft(x, norm="forward")
+        x = torch.view_as_real(x)
+        x = torch.cat((x[:, :, :, 0], -x[:, :, :, 1]), dim=-1)
+
+        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        # q, k, v = qkv.unbind(0)
+        q = self.q(x).reshape(B, self.num_heads, N, self.head_dim)
+        k = self.k(x).reshape(B, self.num_heads, N, self.head_dim)
+        v = self.v(x).reshape(B, self.num_heads, N, self.head_dim)
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1) 
+
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = attn @ v
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        real, img = torch.split(x, x.shape[-1] // 2, dim=-1)
+        x = torch.stack([real,-img], dim=-1)
+        x = torch.view_as_complex(x)
+        x = torch.fft.irfft(x, self.dim, norm="forward")
+
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+
+class Block(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int,
+            mlp_ratio: float = 4.,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            proj_drop: float = 0.,
+            attn_drop: float = 0.,
+            init_values: Optional[float] = None,
+            drop_path: float = 0.,
+            act_layer: nn.Module = nn.GELU,
+            norm_layer: nn.Module = nn.LayerNorm,
+            mlp_layer: nn.Module = Mlp,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = mlp_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=proj_drop,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+class CirT_stage2(nn.Module):
+    def __init__(
+        self,
+        params,
+        img_size=[360, 720],
+        input_size=93,
+        output_size=93,
+        patch_size=124, #124
+        embed_dim=256,
+        depth=8,
+        decoder_depth=2,
+        num_heads=16,
+        mlp_ratio=4.0,
+        drop_path=0.1,
+        drop_rate=0.1
+    ):
+        super().__init__()
+
+        # TODO: remove time_history parameter
+        self.img_size = img_size
+        self.patch_size = img_size[1]
+        self.input_size = input_size
+        self.output_size = output_size
+        self.token_embeds = PatchEmbed(img_size, input_size, embed_dim)
+        # self.token_embeds = nn.Linear(img_size[0] * 2, embed_dim)
+        self.num_patches = self.token_embeds.num_patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim), requires_grad=True)
+        # self.pos_embed = PosEmbed(embed_dim=embed_dim)
+        
+
+        # --------------------------------------------------------------------------
+
+        # ViT backbone
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = [x.item() for x in torch.linspace(0, drop_path, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    embed_dim,
+                    num_heads,
+                    mlp_ratio,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                    norm_layer=nn.LayerNorm,
+                    # drop=drop_rate,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+
+        # --------------------------------------------------------------------------
+
+        # prediction head
+        self.head = nn.ModuleList()
+        for _ in range(decoder_depth):
+            self.head.append(nn.Linear(embed_dim, embed_dim))
+            self.head.append(nn.GELU())
+        self.head.append(nn.Linear(embed_dim, output_size * self.img_size[1]))
+        self.head = nn.Sequential(*self.head)
+
+        # --------------------------------------------------------------------------
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # token embedding layer
+        w = self.token_embeds.proj.weight.data
+        trunc_normal_(w.view([w.shape[0], -1]), std=0.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+    def unpatchify(self, x: torch.Tensor, h=None, w=None):
+        """
+        x: (B, L, V * patch_size)
+        return imgs: (B, V, H, W)
+        """
+        p = self.patch_size
+        c_out = self.output_size
+        h = self.img_size[0] // 1
+        w = self.img_size[1] // p
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, c_out))
+        x = torch.einsum("nhwpc->nchpw", x)
+        imgs = x.reshape(shape=(x.shape[0], c_out, h, w * p))
+        return imgs
+
+    def forward_encoder(self, x: torch.Tensor):
+        # x: `[B, V, H, W]` shape.
+
+        # tokenize each variable separately
+        # x = torch.fft.rfft(x, norm="forward")
+        # x = torch.view_as_real(x)
+        # x = torch.cat((x[:, :, :, :, 0], -x[:, :, :, :, 1]), dim=-1)
+
+        x = self.token_embeds(x)
+
+        # pos_embed = self.pos_embed()
+        # add pos embedding
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x
+
+    def forward(self, x):
+        B, V, H, W = x.shape
+        # print(x.shape)
+        out_transformers = self.forward_encoder(x)  # B, L, D
+        preds = self.head(out_transformers)  # B, L, V*p*p
+        preds = self.unpatchify(preds)
+
+        # real, img = torch.split(preds, preds.shape[-1] // 2, dim=-1)
+        # preds = torch.cat([real, -img], dim=-1)
+        # preds = torch.fft.irfft(preds, W, norm="forward")
+        return preds

networks/CirT1.py

@@ -0,0 +1,301 @@
+from functools import lru_cache
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import trunc_normal_, Block
+from torch.jit import Final
+
+import torch.nn.functional as F
+from typing import Optional
+from timm.layers import DropPath, use_fused_attn, Mlp
+
+class PatchEmbed(nn.Module):
+    def __init__(
+            self,
+            img_size=[121, 240],
+            in_chans=63,
+            embed_dim=768,
+            norm_layer=None,
+            flatten=True,
+            bias=True,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.num_patches = img_size[0]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=[1, img_size[1]], stride=1, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // self.num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+        self.dim = dim
+        self.attn_bias = nn.Parameter(torch.zeros(121, 121, 2), requires_grad=True)
+
+        # self.qkv = CLinear(dim, dim * 3, bias=qkv_bias)
+        self.qkv = nn.Linear((dim // 2 + 1) * 2, dim * 3, bias=qkv_bias)
+        self.q = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.k = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.v = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+
+
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        # self.proj = CLinear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.proj = nn.Linear(dim, dim)
+        # self.proj_drop = nn.Dropout(proj_drop)
+
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+
+        x = torch.fft.rfft(x, norm="forward")
+        x = torch.view_as_real(x)
+        x = torch.cat((x[:, :, :, 0], -x[:, :, :, 1]), dim=-1)
+
+        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        # q, k, v = qkv.unbind(0)
+        q = self.q(x).reshape(B, self.num_heads, N, self.head_dim)
+        k = self.k(x).reshape(B, self.num_heads, N, self.head_dim)
+        v = self.v(x).reshape(B, self.num_heads, N, self.head_dim)
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1) 
+
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = attn @ v
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        real, img = torch.split(x, x.shape[-1] // 2, dim=-1)
+        x = torch.stack([real,-img], dim=-1)
+        x = torch.view_as_complex(x)
+        x = torch.fft.irfft(x, self.dim, norm="forward")
+
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+
+class Block(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int,
+            mlp_ratio: float = 4.,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            proj_drop: float = 0.,
+            attn_drop: float = 0.,
+            init_values: Optional[float] = None,
+            drop_path: float = 0.,
+            act_layer: nn.Module = nn.GELU,
+            norm_layer: nn.Module = nn.LayerNorm,
+            mlp_layer: nn.Module = Mlp,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = mlp_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=proj_drop,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+class CirT(nn.Module):
+    def __init__(
+        self,
+        params,
+        img_size=[360, 720],
+        input_size=101,
+        output_size=93,
+        patch_size=124, #124
+        embed_dim=256,
+        depth=8,
+        decoder_depth=2,
+        num_heads=16,
+        mlp_ratio=4.0,
+        drop_path=0.1,
+        drop_rate=0.1
+    ):
+        super().__init__()
+
+        # TODO: remove time_history parameter
+        self.img_size = img_size
+        self.patch_size = img_size[1]
+        self.input_size = input_size
+        self.output_size = output_size
+        self.token_embeds = PatchEmbed(img_size, input_size, embed_dim)
+        # self.token_embeds = nn.Linear(img_size[0] * 2, embed_dim)
+        self.num_patches = self.token_embeds.num_patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim), requires_grad=True)
+        # self.pos_embed = PosEmbed(embed_dim=embed_dim)
+        
+
+        # --------------------------------------------------------------------------
+
+        # ViT backbone
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = [x.item() for x in torch.linspace(0, drop_path, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    embed_dim,
+                    num_heads,
+                    mlp_ratio,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                    norm_layer=nn.LayerNorm,
+                    # drop=drop_rate,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+
+        # --------------------------------------------------------------------------
+
+        # prediction head
+        self.head = nn.ModuleList()
+        for _ in range(decoder_depth):
+            self.head.append(nn.Linear(embed_dim, embed_dim))
+            self.head.append(nn.GELU())
+        self.head.append(nn.Linear(embed_dim, output_size * self.img_size[1]))
+        self.head = nn.Sequential(*self.head)
+
+        # --------------------------------------------------------------------------
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # token embedding layer
+        w = self.token_embeds.proj.weight.data
+        trunc_normal_(w.view([w.shape[0], -1]), std=0.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+    def unpatchify(self, x: torch.Tensor, h=None, w=None):
+        """
+        x: (B, L, V * patch_size)
+        return imgs: (B, V, H, W)
+        """
+        p = self.patch_size
+        c_out = self.output_size
+        h = self.img_size[0] // 1
+        w = self.img_size[1] // p
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, c_out))
+        x = torch.einsum("nhwpc->nchpw", x)
+        imgs = x.reshape(shape=(x.shape[0], c_out, h, w * p))
+        return imgs
+
+    def forward_encoder(self, x: torch.Tensor):
+        # x: `[B, V, H, W]` shape.
+
+        # tokenize each variable separately
+        # x = torch.fft.rfft(x, norm="forward")
+        # x = torch.view_as_real(x)
+        # x = torch.cat((x[:, :, :, :, 0], -x[:, :, :, :, 1]), dim=-1)
+
+        x = self.token_embeds(x)
+
+        # pos_embed = self.pos_embed()
+        # add pos embedding
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x
+
+    def forward(self, x):
+        B, V, H, W = x.shape
+        # print(x.shape)
+        out_transformers = self.forward_encoder(x)  # B, L, D
+        preds = self.head(out_transformers)  # B, L, V*p*p
+        preds = self.unpatchify(preds)
+
+        # real, img = torch.split(preds, preds.shape[-1] // 2, dim=-1)
+        # preds = torch.cat([real, -img], dim=-1)
+        # preds = torch.fft.irfft(preds, W, norm="forward")
+        return preds

networks/CirT2.py

@@ -0,0 +1,301 @@
+from functools import lru_cache
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import trunc_normal_, Block
+from torch.jit import Final
+
+import torch.nn.functional as F
+from typing import Optional
+from timm.layers import DropPath, use_fused_attn, Mlp
+
+class PatchEmbed(nn.Module):
+    def __init__(
+            self,
+            img_size=[121, 240],
+            in_chans=63,
+            embed_dim=768,
+            norm_layer=None,
+            flatten=True,
+            bias=True,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.num_patches = img_size[0]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=[1, img_size[1]], stride=1, bias=bias)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // self.num_heads
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+        self.dim = dim
+        self.attn_bias = nn.Parameter(torch.zeros(121, 121, 2), requires_grad=True)
+
+        # self.qkv = CLinear(dim, dim * 3, bias=qkv_bias)
+        self.qkv = nn.Linear((dim // 2 + 1) * 2, dim * 3, bias=qkv_bias)
+        self.q = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.k = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+        self.v = nn.Linear((dim // 2 + 1) * 2, dim, bias=qkv_bias)
+
+
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        # self.proj = CLinear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.proj = nn.Linear(dim, dim)
+        # self.proj_drop = nn.Dropout(proj_drop)
+
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+
+        x = torch.fft.rfft(x, norm="forward")
+        x = torch.view_as_real(x)
+        x = torch.cat((x[:, :, :, 0], -x[:, :, :, 1]), dim=-1)
+
+        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        # q, k, v = qkv.unbind(0)
+        q = self.q(x).reshape(B, self.num_heads, N, self.head_dim)
+        k = self.k(x).reshape(B, self.num_heads, N, self.head_dim)
+        v = self.v(x).reshape(B, self.num_heads, N, self.head_dim)
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1) 
+
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = attn @ v
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        real, img = torch.split(x, x.shape[-1] // 2, dim=-1)
+        x = torch.stack([real,-img], dim=-1)
+        x = torch.view_as_complex(x)
+        x = torch.fft.irfft(x, self.dim, norm="forward")
+
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+
+class Block(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int,
+            mlp_ratio: float = 4.,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            proj_drop: float = 0.,
+            attn_drop: float = 0.,
+            init_values: Optional[float] = None,
+            drop_path: float = 0.,
+            act_layer: nn.Module = nn.GELU,
+            norm_layer: nn.Module = nn.LayerNorm,
+            mlp_layer: nn.Module = Mlp,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = mlp_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=proj_drop,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x))))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x
+
+
+class CirT_stage2(nn.Module):
+    def __init__(
+        self,
+        params,
+        img_size=[360, 720],
+        input_size=93,
+        output_size=93,
+        patch_size=124, #124
+        embed_dim=256,
+        depth=8,
+        decoder_depth=2,
+        num_heads=16,
+        mlp_ratio=4.0,
+        drop_path=0.1,
+        drop_rate=0.1
+    ):
+        super().__init__()
+
+        # TODO: remove time_history parameter
+        self.img_size = img_size
+        self.patch_size = img_size[1]
+        self.input_size = input_size
+        self.output_size = output_size
+        self.token_embeds = PatchEmbed(img_size, input_size, embed_dim)
+        # self.token_embeds = nn.Linear(img_size[0] * 2, embed_dim)
+        self.num_patches = self.token_embeds.num_patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim), requires_grad=True)
+        # self.pos_embed = PosEmbed(embed_dim=embed_dim)
+        
+
+        # --------------------------------------------------------------------------
+
+        # ViT backbone
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = [x.item() for x in torch.linspace(0, drop_path, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    embed_dim,
+                    num_heads,
+                    mlp_ratio,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                    norm_layer=nn.LayerNorm,
+                    # drop=drop_rate,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+
+        # --------------------------------------------------------------------------
+
+        # prediction head
+        self.head = nn.ModuleList()
+        for _ in range(decoder_depth):
+            self.head.append(nn.Linear(embed_dim, embed_dim))
+            self.head.append(nn.GELU())
+        self.head.append(nn.Linear(embed_dim, output_size * self.img_size[1]))
+        self.head = nn.Sequential(*self.head)
+
+        # --------------------------------------------------------------------------
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # token embedding layer
+        w = self.token_embeds.proj.weight.data
+        trunc_normal_(w.view([w.shape[0], -1]), std=0.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+
+    def unpatchify(self, x: torch.Tensor, h=None, w=None):
+        """
+        x: (B, L, V * patch_size)
+        return imgs: (B, V, H, W)
+        """
+        p = self.patch_size
+        c_out = self.output_size
+        h = self.img_size[0] // 1
+        w = self.img_size[1] // p
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, c_out))
+        x = torch.einsum("nhwpc->nchpw", x)
+        imgs = x.reshape(shape=(x.shape[0], c_out, h, w * p))
+        return imgs
+
+    def forward_encoder(self, x: torch.Tensor):
+        # x: `[B, V, H, W]` shape.
+
+        # tokenize each variable separately
+        # x = torch.fft.rfft(x, norm="forward")
+        # x = torch.view_as_real(x)
+        # x = torch.cat((x[:, :, :, :, 0], -x[:, :, :, :, 1]), dim=-1)
+
+        x = self.token_embeds(x)
+
+        # pos_embed = self.pos_embed()
+        # add pos embedding
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # apply Transformer blocks
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x
+
+    def forward(self, x):
+        B, V, H, W = x.shape
+        # print(x.shape)
+        out_transformers = self.forward_encoder(x)  # B, L, D
+        preds = self.head(out_transformers)  # B, L, V*p*p
+        preds = self.unpatchify(preds)
+
+        # real, img = torch.split(preds, preds.shape[-1] // 2, dim=-1)
+        # preds = torch.cat([real, -img], dim=-1)
+        # preds = torch.fft.irfft(preds, W, norm="forward")
+        return preds