Upload 99 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+data/SEVIR/CATALOG.csv filter=lfs diff=lfs merge=lfs -text

.ipynb_checkpoints/demo-checkpoint.ipynb

@@ -0,0 +1,37 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6e716a85-c1aa-4b14-90e5-181840d54cc6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
+    "import numpy as np\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.18"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 sqfoo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,155 @@
----
-license: mit
----
+# STLDM: Spatio-Temporal Latent Diffusion Model for Precipitation Nowcasting
+
+This is the official code implementation of the paper [**STLDM: Spatio-Temporal Latent Diffusion Model for Precipitation Nowcasting**](https://openreview.net/forum?id=f4oJwXn3qg) submitted to TMLR.
+
+## Setup Environment
+
+Create a new conda environment:
+
+```bash
+conda create -n stldm python=3.9
+conda activate stldm
+```
+
+Install related packages:
+
+```bash
+pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 -f https://download.pytorch.org/whl/torch_stable.html
+pip install -r requirements.txt
+```
+
+## Data Preparation
+
+There are three radar reflectivity datasets being evaluated with **STLDM** and other baselines: SEVIR, HKO-7, and MeteoNet.
+
+### SEVIR
+
+For the SEVIR dataset, please refer to [https://github.com/amazon-science/earth-forecasting-transformer](https://github.com/amazon-science/earth-forecasting-transformer) for downloading the SEVIR dataset. Please make sure the downloaded files are stored in the following way:
+
+```
+data/
+├─ SEVIR/
+│   ├─ data/
+│   │   └─ vil/
+|   |      ├─ 2017/
+|   |      ├─ 2018/
+|   |      └─ 2019/
+│   └─ CATALOG.csv
+├─ ...
+```
+
+### HKO-7
+
+For the HKO-7 dataset, please refer to [https://github.com/sxjscience/HKO-7](https://github.com/sxjscience/HKO-7) for downloading the HKO-7 dataset. Please make sure the downloaded files are stored in the following way:
+
+```
+data/
+├─ HKO-7/
+│   ├─ hko_data/
+│   │   └─ mask_dat.npz
+│   ├─ radarPNG/
+│   │   ├─ 2009/
+│   │   └─ ...
+│   ├─ radarPNG_mask/
+│   │   ├─ 2009/
+│   │   └─ ...
+│   └─ samplers/
+├─ ...
+```
+
+The samplers files have been saved inside ```data/HKO-7/samplers/``` for you.
+
+### MeteoNet
+
+For the MeteoNet dataset, please refer to [https://github.com/DeminYu98/DiffCast](https://github.com/DeminYu98/DiffCast) to download the pre-processed h5 file, or you can follow the provided instruction to pre-process the raw dataset found on [the official MeteoNet website](https://meteofrance.github.io/meteonet/english/data/rain-radar/). Please make sure the file is named as ```meteo.h5``` and stored in the following way:
+
+```
+data/
+├─ meteonet/
+│   ├─ meteo.h5
+├─ ...
+```
+
+### FYI 💡
+
+If you really want to change the suggested way to save those data files above, remember to update the corresponding file directories as well in the following ways:
+
+- The SEVIR dataset: ```SEVIR_ROOT_DIR``` in ```data/dutils.py```
+- The HKO-7 dataset: ```__C.ROOT_DIR```, ```possible_hko_png_paths``` and ```possible_hko_mask_paths``` in ```nowcasting/config.py```
+- The MeteoNet dataset: ```METEO_FILE_DIR``` in ```data/dutils.py```
+
+## Training
+
+You can train the **STLDM** with the script ```train.py``` with the following command:
+
+``` bash
+python train.py -d HKO7_5_20 --seq_len 5 --out_len 20 -m STLDM_HKO --type "3D" 
+```
+
+In particular, there are a few arguments to set:
+
+- ```-d``` / ```--dataset``` : The dataset config found in ```data/config.py```. Please set the corresponding ```--seq_len``` (input sequence length) and ```--out_len``` (output sequence length) as well.
+- ```-m``` / ```--model``` : The STLDM config found in ```stldm/__init__.py```
+- ```--type```: is to specify whether is ```"3D"``` (**Spatiotemporal**) or ```"2S"``` (**Spatial**) Visual Enhancement.
+
+## Evaluation and Sampling
+
+### Evaluation 
+
+For the evaluation of **STLDM**, we generate ten ensemble predictions of **STLDM** and evaluate them.
+
+First, let's run the ensemble generation script, ```ens_gen.py``` to generate the ensemble prediction and save it as an npy file, with the following command:
+
+```bash
+python ens_gen.py -d HKO7_5_20 -m STLDM_HKO --type "3D" -f "model_checkpoint" --c_str 1.0 --e_id 0
+```
+
+Other than the arguments above, there are stll a few parameters to set:
+
+- ```-f```: the relative/absolute path to **STLDM** checkpoint
+- ```--c_str```: Classifier-Free Guidance strength, it is disabled when set to 0.0
+- ```--e_id```: Represent the $e\_id$ th ensemble prediction, starting from 0
+
+Then, we run the evaluation script, ```ens_eval.py``` to evaluate the generated ensemble predictions (labeled starting from 0) with the following command:
+
+```bash
+python ens_eval.py -d HKO7_5_20 --out_len 20 --e_file "filepath_{}.npy" --ens_no 10
+```
+
+Again, other than the arguments specified above, there are still a few parameters to set:
+
+- ```--e_file```: The format of the ensemble predictions, replace the $e\_id$ by $\{ \}$. Make sure the predictions are labelled, starting from 0.
+- ```--ens_no```: Total number of ensemble predictions, i.e., 10
+
+### Sampling
+
+Other than the evaluation process, we also provide a demo file, ```demo.ipynb``` to show you how to set up and call the **STLDM$** to generate samples for your side implementation. In this demo, we include three different configurations:
+
+- **SpatioTemporal** Visual Enhancement with image size of *128*
+- **Spatial** Visual Enhancement with image size of *128*
+- **SpatioTemporal** Visual Enhancement with image size of *256*
+
+You can download their corresponding modek checkpoints from [this link](https://hkustconnect-my.sharepoint.com/:f:/g/personal/sqfoo_connect_ust_hk/IgATefXlByydRaKlqYnC3hIyAUNk5ftNZBXJz0yKa7d89yE?e=BLk0V3) ([Alternative link](https://drive.google.com/drive/folders/1bCQBt5JPQ-JzHSy8ruYhj6p32Q5uEECM?usp=sharing)).
+
+## Credits and Acknowledgment
+
+We would like to thank these developers and credit their code.
+
+- [FACL](https://github.com/argenycw/FACL)
+- [OpenSTL](https://github.com/chengtan9907/OpenSTL/blob/OpenSTL-Lightning/README.md)
+- [DiffCast](https://github.com/DeminYu98/DiffCast)
+- [denoising_diffusion_pytorch](https://github.com/lucidrains/denoising-diffusion-pytorch/tree/main/denoising_diffusion_pytorch)
+
+## Citation
+
+If you find this work helpful, please cite the following:
+
+```bib
+@article{foo2025stldm,
+  author  = {Foo, Shi Quan and Wong, Chi-Ho and Gao, Zhihan and Yeung, Dit-Yan and Wong, Ka-Hing and Wong, Wai-Kin},
+  title   = {STLDM: Spatio-Temporal Latent Diffusion Model for Precipitation Nowcasting},
+  journal = {Transactions on Machine Learning Research},
+  year    = {2025},
+  url     = {https://openreview.net/forum?id=f4oJwXn3qg},
+}
+```

ckpts/hko-7/stldm/n2n_timesteps-50_sampling_timesteps-20_objective-pred_v_final.pt

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

data/HKO-7/samplers/hko7_cloudy_days_t20_test.txt

@@ -0,0 +1,199 @@
+20150110, 43.75619097948626
+20150111, 149.95275616718968
+20150112, 374.587865563982
+20150113, 216.4265648608205
+20150215, 23.97378254300325
+20150221, 47.148593604108314
+20150222, 109.70054189911669
+20150223, 61.28945820984426
+20150312, 98.79946209611806
+20150319, 31.994431187529052
+20150320, 60.65151818921431
+20150323, 28.129440031220827
+20150325, 22.872230575894932
+20150331, 21.820746346176193
+20150403, 34.57228669078336
+20150404, 24.82365106345886
+20150405, 31.665026804393307
+20150407, 33.26660849235528
+20150408, 94.84768094781496
+20150409, 22.203001174428096
+20150410, 82.69853265922829
+20150411, 353.58096869915147
+20150412, 31.1800565144119
+20150417, 44.918818645397494
+20150418, 24.080770753428627
+20150419, 64.81862887174572
+20150420, 134.63497737244305
+20150421, 21.996475658124133
+20150423, 32.6925866602743
+20150428, 23.7152796014239
+20150429, 54.57041110965831
+20150430, 44.70511117648767
+20150501, 46.1183677286727
+20150502, 44.90144663819153
+20150503, 49.670905610762425
+20150504, 67.67321630927475
+20150505, 67.95217779820261
+20150506, 92.13506781003022
+20150507, 66.49604560574231
+20150508, 41.54096078132263
+20150509, 141.6098750437675
+20150510, 118.69222291230821
+20150511, 123.33292957808122
+20150512, 22.11332483437936
+20150515, 31.402990251627152
+20150516, 190.73493669368895
+20150517, 123.26633275706115
+20150518, 55.93378915330079
+20150519, 75.29142622617388
+20150520, 286.1066578771502
+20150521, 327.0232328350187
+20150522, 197.9260324049861
+20150523, 492.82381178085785
+20150524, 266.1958120859484
+20150525, 172.35522310486695
+20150526, 170.2591961224431
+20150527, 55.43129521879359
+20150530, 87.61113998576245
+20150531, 101.77555679044627
+20150601, 90.33209807938168
+20150602, 38.3629325688609
+20150603, 26.60613086936309
+20150604, 26.28305530131335
+20150605, 76.74834212330765
+20150606, 61.74615059129474
+20150607, 32.53982813226406
+20150609, 41.16127709495584
+20150610, 41.9677021588796
+20150611, 51.02534211601307
+20150612, 47.988766671025104
+20150613, 38.858003544862854
+20150614, 26.969955794817928
+20150615, 22.181291719049337
+20150621, 53.58603615925537
+20150622, 98.10795233321713
+20150623, 178.18817356026273
+20150624, 129.40208224372384
+20150625, 38.546225592747554
+20150702, 28.222513147954437
+20150703, 21.610705049976755
+20150704, 26.45230561366806
+20150705, 55.110333856345875
+20150706, 25.856551631069088
+20150709, 235.03884457428808
+20150710, 102.79465580503033
+20150714, 24.60319509821014
+20150715, 46.689866363211955
+20150716, 63.028152240237105
+20150717, 108.734602365179
+20150718, 73.07054603963272
+20150719, 70.11114648984594
+20150720, 287.5367683272315
+20150721, 322.20042222367505
+20150722, 186.14272948762843
+20150723, 118.66465379474661
+20150724, 127.66547082025804
+20150725, 48.57482203045095
+20150726, 44.40593104429272
+20150727, 32.61881900860065
+20150728, 32.448406627731295
+20150729, 32.201990898128784
+20150730, 30.077929415097636
+20150809, 59.961966454556034
+20150810, 126.86668727386788
+20150811, 76.57628991602742
+20150812, 96.5194674223856
+20150813, 252.00741877480158
+20150814, 179.07121160942586
+20150815, 152.54204565793245
+20150816, 46.21115141369047
+20150817, 36.35543464781746
+20150820, 48.975007445664815
+20150825, 29.65178480246265
+20150826, 102.09460952028128
+20150827, 120.86394627498836
+20150828, 104.83123928550674
+20150829, 276.51425990091815
+20150830, 289.6720114087301
+20150831, 225.63438807531384
+20150901, 231.3144142622618
+20150902, 319.48056572524405
+20150903, 135.36383892098507
+20150904, 21.66859854428173
+20150905, 24.239027087691777
+20150906, 23.229757345646593
+20150907, 159.19597334815205
+20150908, 91.84391777544349
+20150909, 48.96262528448881
+20150910, 43.52720809361923
+20150911, 41.92730688487913
+20150912, 43.19016518479776
+20150913, 34.703902617968396
+20150914, 29.62716214841933
+20150915, 51.113737251569034
+20150916, 100.1013191538819
+20150917, 51.71855559932305
+20150918, 59.26835556136683
+20150919, 40.77983042044398
+20150920, 78.79494620961182
+20150921, 161.67126445548584
+20150922, 49.63578440986749
+20150923, 36.10682937237395
+20150924, 50.429000532504666
+20150925, 59.50495138159431
+20150926, 59.904413770716616
+20150927, 58.42647587604603
+20150928, 78.84617601551604
+20150929, 41.22325771443514
+20150930, 44.00618280741514
+20151001, 50.48162457691409
+20151002, 25.243249226774044
+20151003, 184.0588648811599
+20151004, 518.6718274203859
+20151005, 258.02547016649237
+20151006, 107.7635876118666
+20151007, 84.55949267782428
+20151008, 53.157816495234776
+20151009, 26.70240894496746
+20151010, 83.35799482798699
+20151011, 110.94956088737794
+20151012, 60.71514775104601
+20151013, 88.16609571129707
+20151014, 87.18272267512077
+20151015, 52.31939377760342
+20151016, 49.53965161552765
+20151017, 35.443420211529514
+20151018, 31.41580968880754
+20151019, 35.210431376394695
+20151023, 23.18438644089958
+20151024, 28.495018850244072
+20151025, 23.631041521385402
+20151027, 32.73702674628079
+20151028, 20.834829817343607
+20151030, 23.063971500058358
+20151102, 21.893503748256617
+20151109, 40.70959491951417
+20151110, 23.001540157641813
+20151113, 121.34067551499768
+20151114, 30.08454733990005
+20151116, 22.963587075455177
+20151118, 65.64878399581589
+20151119, 86.3667786001462
+20151120, 48.5797208727241
+20151121, 40.62472124157369
+20151122, 31.929468670095304
+20151123, 46.070741987738245
+20151124, 43.33614906584146
+20151125, 38.5326454265458
+20151203, 44.5021005854835
+20151205, 156.30654074413056
+20151206, 31.48730096466759
+20151208, 33.42894220711297
+20151209, 473.5921466396443
+20151220, 49.35841541724779
+20151221, 30.52428394496746
+20151226, 55.58330554829149
+20151227, 28.596221597512784
+20151229, 42.04237763685495
+20151230, 49.52976704149231

data/HKO-7/samplers/hko7_cloudy_days_t20_test.txt.pkl

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

data/HKO-7/samplers/hko7_cloudy_days_t20_train.txt

@@ -0,0 +1,1224 @@
+20090305, 136.62595858667146
+20090306, 228.80309496601032
+20090307, 61.28797688575081
+20090313, 50.83967177330312
+20090323, 27.882427867852165
+20090324, 98.90460069444444
+20090325, 186.85773640893774
+20090327, 83.00352533110642
+20090328, 85.88496538673874
+20090329, 95.76494835251044
+20090404, 55.43760260489309
+20090405, 51.40799718880753
+20090406, 51.97289596408646
+20090407, 41.94900518654117
+20090408, 22.849034969200364
+20090409, 96.84607150151723
+20090410, 34.88548858089261
+20090411, 26.66028119188749
+20090412, 73.90951701243608
+20090413, 83.42903845595072
+20090414, 57.75939424976757
+20090415, 163.3683099793701
+20090416, 180.4777461791027
+20090417, 38.150535906264544
+20090418, 136.80413615760114
+20090419, 88.94150068282197
+20090420, 37.34591705166202
+20090421, 35.38847592689446
+20090422, 46.74223053957462
+20090424, 21.330452405857738
+20090425, 251.52358005724082
+20090426, 30.056307386099487
+20090429, 41.453953426932365
+20090430, 27.32378381421432
+20090505, 21.583478977801022
+20090509, 22.712410470420732
+20090513, 32.6148927469136
+20090516, 85.55672561453974
+20090517, 109.60553775860065
+20090518, 116.95029219694324
+20090519, 89.23390664952346
+20090520, 106.5634751889944
+20090521, 82.6532027254765
+20090522, 181.06651902458165
+20090523, 501.6702046286611
+20090524, 400.99202769060906
+20090525, 173.28997614670865
+20090526, 101.19955307851002
+20090527, 53.12012853759879
+20090528, 60.35818260402138
+20090529, 25.776541979021392
+20090603, 160.16725414778008
+20090604, 97.28280759675732
+20090607, 71.03364441684101
+20090608, 99.24742143915621
+20090609, 116.8017359295095
+20090610, 82.46225544905461
+20090611, 152.62945577638308
+20090612, 90.96381378506264
+20090613, 88.4151001714319
+20090614, 118.5259904550209
+20090615, 130.086386600709
+20090616, 87.38530679771037
+20090617, 25.00878032302665
+20090620, 30.45724553986518
+20090621, 40.03228839783821
+20090622, 106.4602314330544
+20090623, 83.79152683344954
+20090624, 70.1403104660623
+20090625, 85.02154085309161
+20090626, 75.30984571129706
+20090627, 335.9534569676896
+20090628, 186.45501920803784
+20090629, 42.00070443253138
+20090703, 63.63226896646908
+20090704, 237.3798625639238
+20090705, 193.21849521298233
+20090706, 48.60273110617154
+20090710, 32.801129743433286
+20090711, 73.77359440376569
+20090712, 27.58495957548814
+20090714, 44.6260489307299
+20090715, 54.747222222222206
+20090716, 22.874410779936973
+20090718, 235.99509966294747
+20090719, 341.8592338955717
+20090720, 65.83892666201766
+20090723, 22.81097400191771
+20090724, 34.50459739803094
+20090725, 35.20270513714551
+20090726, 51.135496026557405
+20090727, 37.10169459541063
+20090728, 42.00904757235007
+20090729, 44.38810836529521
+20090730, 63.88379169572293
+20090731, 33.26692928094232
+20090803, 57.20850566250586
+20090804, 195.11130405287113
+20090805, 320.39658806950257
+20090806, 177.06629202260575
+20090807, 23.55537268276384
+20090808, 30.765311555671776
+20090810, 35.96150210866714
+20090811, 253.41603796199436
+20090812, 341.82729526237796
+20090813, 191.62245008680554
+20090814, 53.51327997344771
+20090823, 35.0751107769642
+20090824, 51.215284315434666
+20090825, 45.89880397198978
+20090828, 24.214795262377965
+20090829, 26.32466514861783
+20090830, 54.47782354137608
+20090831, 32.23210352742911
+20090901, 62.516496505985586
+20090902, 26.58246055519574
+20090905, 20.926650939969775
+20090907, 22.240835448762837
+20090908, 34.00065395950047
+20090909, 57.947631552185015
+20090910, 91.215858820922
+20090911, 82.86404510226207
+20090912, 70.03017207278481
+20090913, 56.37304160855416
+20090914, 272.6651468369081
+20090915, 285.3876847353525
+20090916, 49.86654958858544
+20090918, 151.93731350368924
+20090919, 69.07707661407484
+20090920, 36.43623096815435
+20090921, 43.216834880541164
+20090922, 46.76212733477011
+20090923, 49.67186882554626
+20090924, 52.48864079419726
+20090925, 61.88516623808692
+20090926, 37.474039327638295
+20090927, 69.27637835599721
+20090928, 463.9436457606926
+20090929, 332.59933334059735
+20090930, 66.80611588150651
+20091001, 54.379151775337064
+20091002, 28.19407234280568
+20091003, 44.98317751917714
+20091004, 37.289169463331014
+20091007, 37.08707105706647
+20091008, 61.00016235181311
+20091009, 44.652421475476515
+20091010, 23.698430417538354
+20091011, 112.88366802504649
+20091012, 35.83612582083915
+20091013, 36.11428387232682
+20091016, 60.58141761825895
+20091017, 51.14802181398186
+20091018, 71.10693136913063
+20091019, 101.7541401528359
+20091020, 201.2627237331474
+20091021, 64.73917236983361
+20091022, 63.29760327994324
+20091023, 39.95842286291259
+20091024, 35.976269213447246
+20091025, 31.75665043148535
+20091026, 55.861363900511385
+20091027, 46.790897765574144
+20091028, 30.22591944880288
+20091029, 21.87970993084734
+20091030, 23.724973122966055
+20091031, 40.22858427039748
+20091101, 33.49137192730126
+20091102, 24.3667207839377
+20091109, 22.49475020151289
+20091110, 28.50735213999301
+20091111, 66.11944244682707
+20091112, 154.01058955137148
+20091113, 59.36137043816828
+20091115, 66.11488043351932
+20091116, 102.36392502760343
+20091123, 24.5346406831125
+20091124, 34.07892132293119
+20091125, 46.277792245370364
+20091126, 39.02615244363087
+20091128, 35.47976431746862
+20091206, 24.131424337517426
+20091207, 98.9516879866922
+20091208, 66.62441923814505
+20091209, 34.3988533676197
+20091215, 56.45255150220826
+20091216, 24.40855307124593
+20091227, 81.33299646240121
+20091228, 23.943436737273355
+20091229, 88.94060556863089
+20091230, 96.75108851987449
+20100102, 134.41120754155045
+20100103, 26.636634668468155
+20100105, 37.79852993520455
+20100106, 40.25822146675964
+20100107, 57.98259385169688
+20100110, 49.90426454991865
+20100111, 95.87596648361225
+20100121, 52.24734371367966
+20100122, 48.799663673872615
+20100206, 32.67162693224081
+20100207, 209.8687973980125
+20100208, 60.44115473181079
+20100212, 27.867008077638303
+20100213, 27.533020433809863
+20100214, 66.34328727917247
+20100216, 32.234679545850774
+20100217, 25.200238987680148
+20100218, 51.988591970304505
+20100219, 88.0366902385518
+20100222, 78.82401481142492
+20100306, 26.108932618549513
+20100307, 35.84544125552069
+20100308, 20.593918344665273
+20100309, 50.63448759298
+20100324, 55.8357261157601
+20100325, 39.645025751104136
+20100401, 20.778464050151097
+20100402, 102.23865208042773
+20100405, 24.484351391794515
+20100406, 59.879344636215706
+20100407, 75.98970573678116
+20100408, 141.23548440405628
+20100409, 25.15233466992097
+20100410, 34.2181490803696
+20100411, 78.74429970798465
+20100412, 73.08766689185262
+20100413, 53.82095809361923
+20100415, 155.0404605071195
+20100417, 127.06235825421942
+20100418, 102.04714594955833
+20100420, 44.37650511390051
+20100421, 40.503602975360295
+20100422, 189.74148597309394
+20100426, 85.08088334640868
+20100427, 25.063565274872154
+20100428, 59.77439016350212
+20100429, 324.0794886633429
+20100430, 34.49387021427267
+20100503, 29.15010260489308
+20100506, 124.59078898622735
+20100507, 103.47258360936776
+20100508, 24.618598544281735
+20100509, 21.73658647486772
+20100510, 157.08738735945764
+20100512, 22.094342621894043
+20100514, 84.33014786000699
+20100515, 88.7563665882148
+20100516, 20.627662460774058
+20100519, 204.0931331834312
+20100520, 148.64851395426544
+20100521, 37.800374825662495
+20100522, 117.86880774930266
+20100523, 181.03243712953275
+20100526, 41.56335881635037
+20100527, 27.768272204715217
+20100528, 75.1496062748016
+20100529, 206.38358703074147
+20100530, 175.6673069938401
+20100531, 118.4293410041841
+20100601, 237.94262170937938
+20100602, 334.39367899378186
+20100608, 55.2173480357973
+20100609, 207.89246825604368
+20100610, 319.07378980706653
+20100611, 168.00780432792885
+20100612, 72.57029488464669
+20100613, 126.20548273332172
+20100614, 118.64496182734774
+20100615, 210.03491617271035
+20100616, 181.45767230357976
+20100617, 73.54427446536496
+20100618, 25.24199954236402
+20100621, 120.53997976958391
+20100622, 82.61656079294514
+20100623, 104.60782877604167
+20100624, 120.09928467137378
+20100625, 218.53039047971882
+20100626, 407.84667342224554
+20100627, 376.6243358830195
+20100628, 250.77930849750115
+20100629, 35.55333202580195
+20100713, 20.72497693659926
+20100714, 24.57621055613668
+20100715, 64.12604257467457
+20100716, 94.36938324471177
+20100717, 118.22245939388658
+20100718, 25.131903584088786
+20100720, 24.465838200255696
+20100721, 213.62310389789633
+20100722, 206.2271132873631
+20100723, 127.32937133164812
+20100724, 41.27161966817759
+20100725, 41.94855426981636
+20100726, 114.66956012465128
+20100727, 358.45817969839607
+20100728, 234.21029375291786
+20100729, 115.44730793816831
+20100730, 55.26047169194562
+20100801, 23.52279190638075
+20100802, 51.18088879445607
+20100803, 22.69003860849605
+20100804, 25.242931712962964
+20100805, 131.22104725999537
+20100806, 47.94726707781264
+20100807, 92.85841741486517
+20100808, 55.812995772315205
+20100810, 21.316553347280333
+20100811, 79.84416296199443
+20100812, 43.735309703335666
+20100813, 20.47297060233613
+20100814, 39.483855801080885
+20100815, 72.57462335832172
+20100816, 119.99333086355185
+20100817, 79.66040751804195
+20100818, 52.19738548204323
+20100819, 91.94117597919576
+20100820, 67.82326788412367
+20100822, 59.08391282397721
+20100823, 107.90557934536264
+20100824, 271.26208825231475
+20100825, 158.71470170454546
+20100826, 52.67590156909284
+20100827, 45.05917614626917
+20100828, 21.632718430381214
+20100829, 111.65726660564856
+20100830, 124.24277933955138
+20100831, 43.97353320403301
+20100901, 40.16028300790331
+20100902, 142.00870325575315
+20100903, 498.5411797928289
+20100904, 266.90315673144255
+20100905, 51.12990215306834
+20100906, 32.760605169107386
+20100907, 31.888638823512327
+20100908, 55.89353366427952
+20100909, 152.4942677767997
+20100910, 142.45224347461686
+20100911, 421.70778871019286
+20100912, 234.28636390051145
+20100913, 52.64439450400976
+20100914, 84.66343270571826
+20100915, 76.15386468114241
+20100916, 46.37933737215248
+20100917, 61.118001983089236
+20100918, 68.72911563662251
+20100919, 97.44578212168761
+20100920, 332.5900812848675
+20100921, 579.4747775380636
+20100922, 378.79682342515105
+20100923, 124.51707036698046
+20100924, 85.51370801371456
+20100925, 78.88520437441889
+20100926, 62.14001699790795
+20100927, 59.77074071652721
+20100928, 56.54303221450616
+20100929, 68.64558376736112
+20100930, 38.8137911872385
+20101001, 26.032749665853093
+20101002, 24.983342413412363
+20101003, 74.97003192555788
+20101004, 147.84034751278477
+20101005, 320.81967598645986
+20101006, 282.0738953175848
+20101007, 191.48288677504644
+20101008, 72.47399828568108
+20101009, 144.1631623009647
+20101010, 155.19021930206875
+20101011, 75.0252271835774
+20101012, 58.9400700255695
+20101013, 58.17412940202231
+20101014, 46.33378225244073
+20101015, 75.42381559449093
+20101016, 32.24196231403998
+20101017, 39.06391013627381
+20101018, 41.11644093590189
+20101020, 27.779980742296928
+20101021, 89.352832258252
+20101022, 169.78000384995352
+20101023, 69.3806056775918
+20101027, 40.52576958038521
+20101028, 29.311661182298934
+20101104, 167.72977993520453
+20101105, 279.8310407949792
+20101106, 158.45435916434218
+20101107, 40.77464769293352
+20101108, 30.161046824151555
+20101109, 44.06837121212121
+20101110, 65.25919721205254
+20101111, 42.84965041695723
+20101112, 25.77128070049694
+20101113, 86.23978596090822
+20101114, 61.09089703916783
+20101116, 27.438181732333796
+20101215, 83.16066727684799
+20101216, 76.18976547971873
+20101225, 52.48511448163644
+20110103, 35.99297964609484
+20110104, 91.52568191393539
+20110111, 71.14335537976523
+20110112, 99.04646240120874
+20110120, 22.689244283182234
+20110213, 221.73691360849605
+20110214, 40.80066320897257
+20110215, 149.97527676080898
+20110216, 22.64409540620641
+20110219, 33.252693393836445
+20110309, 91.90697621641252
+20110317, 55.72166615818224
+20110318, 142.97228550274278
+20110319, 94.70543569851233
+20110327, 96.05199209515366
+20110329, 91.72684224731557
+20110330, 259.44745757787075
+20110331, 54.02866871513249
+20110415, 22.540139760576473
+20110416, 24.26677653562297
+20110417, 125.34206364772196
+20110418, 44.91978966904928
+20110419, 22.963660288665253
+20110420, 23.317585633497426
+20110422, 57.15267880491631
+20110423, 61.766474895397494
+20110425, 21.202812645281266
+20110426, 24.233773898768014
+20110427, 41.24985853236867
+20110428, 137.31416801051836
+20110429, 295.45362149872153
+20110430, 95.75766830834493
+20110501, 152.57700361023944
+20110502, 117.76413259821015
+20110503, 177.39995587081592
+20110504, 75.2182224474082
+20110505, 73.69405418264763
+20110506, 127.82177475592748
+20110507, 104.55865658414693
+20110508, 84.0641158908647
+20110509, 20.580801902090975
+20110511, 21.399591941248257
+20110512, 79.2632547361692
+20110513, 191.61483376191308
+20110514, 146.87530164022547
+20110515, 91.66247693659926
+20110516, 203.88988951359835
+20110517, 48.00782158007903
+20110518, 76.23233597745235
+20110520, 25.131326999070197
+20110521, 161.98638423988842
+20110522, 306.4063389847746
+20110523, 131.7929106374942
+20110524, 55.43122512058344
+20110531, 32.16935289740896
+20110605, 23.848160194386328
+20110607, 33.29129329381683
+20110608, 38.98019979805906
+20110609, 21.67068333042771
+20110611, 128.9752891097164
+20110612, 167.42113388394932
+20110613, 41.01110094142259
+20110614, 27.571970340829843
+20110615, 47.03335839435147
+20110616, 177.92250715510227
+20110617, 81.67146349081827
+20110618, 34.76633651499303
+20110621, 63.823974495874005
+20110622, 182.5606137043817
+20110623, 218.6694170226058
+20110624, 34.418169964551375
+20110625, 39.781262893712224
+20110626, 123.53867786785216
+20110627, 125.84756871803812
+20110628, 210.20931743229892
+20110629, 323.0898435683983
+20110630, 188.5001558141562
+20110701, 37.033898477452354
+20110709, 60.17592090161552
+20110710, 43.41258789516504
+20110711, 161.83613054248025
+20110712, 183.02187009675737
+20110713, 92.7071393029405
+20110714, 136.99119668177593
+20110715, 383.15455819967457
+20110716, 381.922978411204
+20110717, 198.67304197175736
+20110718, 164.5868838258368
+20110719, 144.3219961282543
+20110720, 203.32798353236865
+20110721, 49.25801589377032
+20110722, 25.731379481926997
+20110728, 107.40300260053466
+20110729, 142.66062387407015
+20110730, 100.56940757932357
+20110731, 26.589137501452807
+20110805, 29.346624396135265
+20110806, 20.18941389905857
+20110807, 43.96820283443747
+20110808, 139.8291853716295
+20110809, 208.90599285216177
+20110810, 175.34678056717806
+20110811, 56.27505284605997
+20110812, 20.006154114365422
+20110817, 25.62793449994189
+20110818, 21.587661080602047
+20110819, 30.14290101261041
+20110820, 35.011025395165035
+20110821, 20.12696020746165
+20110824, 40.05301749186426
+20110825, 22.119369624593205
+20110826, 30.547850018886557
+20110830, 34.09250966120409
+20110831, 52.65844156787542
+20110901, 130.54113203161316
+20110902, 128.1669594084147
+20110903, 116.22911418380986
+20110904, 37.470217522373325
+20110905, 47.890565979486276
+20110906, 30.946702841701526
+20110907, 34.402428194735016
+20110908, 31.603840510227805
+20110909, 78.73658291201767
+20110910, 100.50139034170155
+20110911, 156.5110798756392
+20110912, 66.90991326708507
+20110913, 46.516962314039965
+20110914, 41.780873234710555
+20110915, 64.64849579410738
+20110916, 77.78759879125988
+20110917, 58.228567563052074
+20110918, 112.29172586878197
+20110919, 60.13979453597164
+20110920, 27.00396381770107
+20110922, 29.113824601929327
+20110923, 84.00498296577173
+20110924, 141.03301316974665
+20110925, 230.6655854108554
+20110926, 78.97429810989075
+20110927, 35.68688818427476
+20110928, 136.85599775540447
+20110929, 259.23085592456994
+20110930, 75.60321144235238
+20111001, 61.53155200342862
+20111002, 214.85392967660394
+20111003, 176.42546925848444
+20111004, 61.928290802243126
+20111005, 42.20670182473268
+20111006, 90.13077529346815
+20111007, 68.09544016591119
+20111008, 29.100231360413762
+20111009, 51.94596844490933
+20111010, 137.91833086355183
+20111011, 127.01731498430965
+20111012, 170.94246189998842
+20111013, 286.50326265399815
+20111014, 114.70196547390748
+20111015, 47.84846256101814
+20111016, 39.638181913935384
+20111017, 28.59887061976987
+20111018, 32.294380157484895
+20111019, 31.440034613261272
+20111020, 39.14295803550673
+20111024, 28.146944553405397
+20111104, 35.26615709263133
+20111107, 71.93888816248257
+20111108, 331.386634305265
+20111109, 379.73886037744074
+20111112, 39.01522039167829
+20111113, 34.4868455079033
+20111114, 35.933232726057646
+20111115, 26.571033639876806
+20111116, 26.743759080079037
+20111117, 198.26452394961643
+20111118, 133.26355946362156
+20111119, 86.71600527370987
+20111120, 35.035476050383544
+20111121, 57.03569506188985
+20111122, 28.57397849110879
+20111123, 35.00330605677591
+20111125, 23.649023346699206
+20111126, 67.68473711355185
+20111127, 30.021795095304512
+20111205, 48.25002124738492
+20120105, 104.29902189388655
+20120106, 21.676583565783353
+20120111, 46.484685720304505
+20120112, 155.01304698396095
+20120113, 148.11024614278244
+20120115, 265.70077870757785
+20120116, 157.90245852219903
+20120122, 34.969506188981875
+20120123, 64.36847218559541
+20120124, 35.67474159780824
+20120125, 30.020862571187816
+20120207, 66.90006538120565
+20120208, 23.766170894351465
+20120222, 21.957772184449095
+20120224, 56.05528164390757
+20120225, 29.70407423146211
+20120227, 43.75812358350767
+20120228, 146.18389430061598
+20120304, 26.300959764353788
+20120305, 34.86704457955603
+20120306, 59.563980416085535
+20120307, 38.043944676894476
+20120308, 54.86796003312412
+20120311, 73.48692305177823
+20120312, 42.601446274988376
+20120318, 20.28321583711065
+20120327, 20.15923425877499
+20120403, 33.63068885111576
+20120405, 199.76914952347744
+20120406, 181.56111873837747
+20120407, 22.05863079672245
+20120408, 243.07268748547187
+20120409, 141.99248950342863
+20120410, 24.95706648070665
+20120411, 34.93628381421433
+20120412, 46.378410114481625
+20120413, 73.36980219414893
+20120415, 27.01078804189911
+20120416, 86.76931859454905
+20120417, 253.06885242329153
+20120418, 158.07843917073458
+20120419, 216.44427755259184
+20120420, 356.705582614191
+20120421, 69.50133889328411
+20120422, 25.895080049976745
+20120423, 62.498505418991165
+20120424, 35.07947938458856
+20120425, 79.38964326185493
+20120426, 46.88026808025337
+20120427, 246.83180188720362
+20120428, 177.62698944531613
+20120429, 98.48469970362622
+20120430, 35.45415776818921
+20120501, 38.946295146152956
+20120502, 62.265384922710346
+20120503, 98.89222781555094
+20120504, 59.85654473936542
+20120505, 115.06919256305207
+20120506, 69.89122809884935
+20120510, 52.98434067730125
+20120511, 82.38079926487681
+20120512, 33.287890080195254
+20120513, 100.77997188807528
+20120514, 39.74556837662715
+20120515, 93.45692682909112
+20120516, 179.34509784693168
+20120517, 150.1514281148303
+20120518, 283.1472258542538
+20120519, 162.8450588592945
+20120520, 84.38526884298003
+20120521, 55.78281831125058
+20120523, 30.494975629067877
+20120525, 42.9129234948861
+20120526, 87.15840252353556
+20120527, 111.9839186352278
+20120528, 155.79253399581592
+20120529, 81.49536661727102
+20120530, 73.16466087720828
+20120531, 66.30515494246862
+20120601, 65.95614049424687
+20120607, 20.85536124186425
+20120608, 30.526771886622516
+20120609, 125.87894311512086
+20120610, 60.65825760547421
+20120611, 112.87685033850535
+20120612, 176.38053267375636
+20120613, 225.37360402864948
+20120614, 95.9233399799512
+20120615, 112.87394634762902
+20120616, 199.91418707868434
+20120617, 123.39643516097166
+20120618, 82.30777291085542
+20120619, 52.6865451388889
+20120620, 73.52798008193864
+20120621, 227.8351153896444
+20120622, 200.99316033966764
+20120623, 130.71126093241514
+20120624, 145.42582483437937
+20120625, 45.958252520629955
+20120626, 45.63854874912831
+20120627, 22.53596165301022
+20120629, 76.21593517549977
+20120630, 67.28973624186425
+20120701, 30.23978636390051
+20120705, 97.22105180003486
+20120706, 67.02140174628079
+20120707, 36.87964391562063
+20120708, 22.167915359135282
+20120713, 31.47329403475128
+20120714, 24.112315311192468
+20120715, 23.03197931194793
+20120716, 29.393501205834497
+20120717, 31.363913223500706
+20120718, 42.11399966585309
+20120719, 30.423471459495584
+20120720, 32.19373801429568
+20120721, 38.418141271501625
+20120722, 111.534527000523
+20120723, 305.4445363348443
+20120725, 425.3922848384472
+20120726, 270.14713341904934
+20120727, 143.0784675005811
+20120728, 39.630659794862844
+20120729, 23.80491940521851
+20120803, 65.87163310669456
+20120804, 174.61665994740815
+20120805, 127.87644718299626
+20120806, 52.96096837226872
+20120808, 24.279802998605298
+20120809, 21.723397729253836
+20120810, 56.92519594810553
+20120811, 148.85347222222217
+20120812, 204.3983773898768
+20120813, 160.65132732595308
+20120814, 28.701588832229195
+20120816, 98.85505270077871
+20120817, 141.65484167974196
+20120818, 73.6276789865179
+20120820, 22.215848369944215
+20120821, 48.31004220420735
+20120822, 114.96252433461177
+20120823, 21.592981280509065
+20120825, 58.30701000261506
+20120826, 92.99211268014876
+20120829, 26.10597381736705
+20120830, 30.602471597512785
+20120831, 48.01793242968387
+20120901, 41.164819778591344
+20120902, 36.74191092079264
+20120903, 52.148103534693156
+20120904, 77.8733388903417
+20120905, 33.99965459379359
+20120906, 60.53588192991633
+20120907, 74.14039636360997
+20120909, 33.5636689693747
+20120910, 45.419486394409574
+20120911, 41.957370844955825
+20120912, 34.730700836820084
+20120913, 61.21184986634123
+20120914, 33.420998045966996
+20120915, 26.844679618491394
+20120916, 54.41844618055556
+20120917, 71.693800666841
+20120919, 24.288151404869826
+20120920, 30.290509319793124
+20120921, 44.52357624360763
+20120922, 46.52463007758021
+20120923, 60.30840070751975
+20120924, 102.84257503777312
+20120925, 116.06859727307068
+20120926, 77.39527763249652
+20120927, 46.295262196362145
+20120930, 39.05014368603801
+20121001, 43.8218346844491
+20121002, 45.0548095362622
+20121003, 55.270770135983255
+20121004, 64.44542382176896
+20121005, 43.024875602917255
+20121006, 21.134233532368672
+20121008, 32.69247788092748
+20121010, 30.112397758310088
+20121011, 29.853273731694564
+20121012, 64.561951744828
+20121013, 100.94356912482566
+20121014, 71.15270459234077
+20121015, 29.164262988145044
+20121016, 50.05106436686424
+20121017, 32.32784406235472
+20121018, 35.83185237244305
+20121019, 23.56877088418177
+20121020, 41.49348014005114
+20121021, 40.56814653794745
+20121022, 34.88908138656439
+20121023, 24.15241021617852
+20121025, 25.634765988203167
+20121026, 42.98020288528591
+20121027, 51.68258350040679
+20121028, 34.085144990702005
+20121029, 70.69760158792423
+20121030, 414.6291510489307
+20121031, 246.88588501714315
+20121101, 88.46325292015341
+20121102, 60.19221546664342
+20121103, 58.46229333740121
+20121104, 28.681218946129714
+20121105, 41.65760765341701
+20121106, 22.620909787598798
+20121110, 52.8720482463819
+20121111, 47.08687238493724
+20121113, 44.9120341919456
+20121114, 39.10897965335891
+20121116, 90.44041960861227
+20121117, 238.69832726784048
+20121118, 22.83449976028591
+20121119, 23.235088875813577
+20121120, 23.16098144758252
+20121121, 34.50703597164109
+20121123, 104.08733910099954
+20121124, 36.144737912598785
+20121125, 58.245344643479775
+20121126, 138.49525910913528
+20121127, 116.1277260213273
+20121128, 153.1060002978266
+20121129, 62.34876438284517
+20121130, 106.80491123314735
+20121201, 29.832474357856817
+20121202, 51.42532379561831
+20121203, 43.87085330950721
+20121204, 37.574779354079496
+20121205, 59.6601429930846
+20121207, 24.548445490469543
+20121208, 38.5434260227801
+20121218, 33.644407034518835
+20121226, 64.43589990847282
+20121227, 77.76107515399815
+20121229, 55.823466919456074
+20121230, 22.691671025104604
+20130126, 68.32953658908647
+20130203, 40.21736147431425
+20130204, 69.91238250377732
+20130205, 57.8372849837285
+20130313, 37.18495957548814
+20130319, 46.00838563458857
+20130320, 61.9663423262436
+20130323, 33.65348184710601
+20130324, 64.3982873154928
+20130326, 168.771519424105
+20130327, 41.56376340219664
+20130328, 265.6667685451534
+20130329, 79.51368476871221
+20130330, 355.0431574340425
+20130331, 67.06829345362623
+20130402, 149.4377527894003
+20130403, 110.74030156758481
+20130404, 66.24243357014181
+20130405, 156.6795004503719
+20130406, 91.90465898855184
+20130408, 23.98437699761739
+20130409, 136.47248481810786
+20130410, 99.25893407136216
+20130411, 145.10497624651325
+20130412, 38.984963752324504
+20130413, 25.025353759879135
+20130414, 44.8889865905393
+20130415, 43.28978545589259
+20130416, 27.412897707461646
+20130417, 118.70350817933519
+20130418, 98.94480619479309
+20130419, 78.98868458711065
+20130420, 81.90855979050443
+20130421, 22.120508448105525
+20130425, 121.91539109716413
+20130426, 252.95637327847808
+20130428, 25.334676458623893
+20130429, 62.346051824763585
+20130430, 142.26964329817525
+20130501, 33.8277409126569
+20130502, 119.36171929480476
+20130503, 68.33329737622034
+20130504, 54.23608768450721
+20130506, 21.33145303056718
+20130507, 22.99018280015109
+20130508, 178.57049848552012
+20130509, 137.87810556863087
+20130510, 156.389030901325
+20130511, 46.95578782397722
+20130512, 48.34755909315434
+20130514, 31.646399022257093
+20130515, 141.49384643043936
+20130516, 231.89553605154583
+20130517, 108.04281911029754
+20130518, 32.8831909576941
+20130519, 113.01127255491636
+20130520, 148.83095816625988
+20130521, 243.68406282688287
+20130522, 447.63819553405403
+20130523, 96.4477764702464
+20130524, 42.805950902196656
+20130525, 163.88438226406328
+20130526, 80.59039618200836
+20130527, 89.96674293933053
+20130528, 52.900444197466285
+20130529, 49.58047056601582
+20130530, 26.962475302185027
+20130603, 54.61715099808228
+20130604, 153.63807785622967
+20130605, 186.80010859774524
+20130606, 61.38567726493492
+20130608, 35.266644511273824
+20130609, 58.68488003399581
+20130610, 139.65129427446539
+20130611, 309.9337880636913
+20130612, 120.38993564039978
+20130613, 116.76242281932822
+20130614, 233.6716854079498
+20130615, 320.22559765080206
+20130616, 139.76226337314037
+20130617, 30.093121840132504
+20130621, 48.08229983873781
+20130622, 82.57217791143655
+20130623, 146.94995877644118
+20130624, 282.88758952957926
+20130625, 95.0396178376337
+20130701, 115.02560437006045
+20130702, 61.928661814272424
+20130706, 37.02248790533473
+20130707, 97.14188077493026
+20130708, 33.69977971728266
+20130709, 55.034494857043235
+20130710, 35.16722363871455
+20130714, 127.9919087415737
+20130715, 180.73405374680374
+20130716, 199.5831769743724
+20130717, 191.53194281003022
+20130718, 65.19927994973267
+20130719, 133.9287567918991
+20130720, 240.20309303812172
+20130721, 76.10018487040911
+20130722, 50.93545008135751
+20130723, 55.84513471205254
+20130724, 179.46736546954904
+20130725, 261.0974208943515
+20130726, 271.54983855619474
+20130727, 136.99305464754767
+20130728, 92.09601166608553
+20130731, 24.211531881973492
+20130801, 64.45616210483497
+20130802, 185.41478363987684
+20130803, 120.62649948425155
+20130807, 37.23730023826127
+20130810, 33.29708638423989
+20130811, 21.927062630753134
+20130813, 322.93491290388187
+20130814, 307.2372183359484
+20130815, 220.49118614888434
+20130816, 297.62268022140853
+20130817, 267.12960251046024
+20130818, 111.35105564998838
+20130819, 45.49491134210832
+20130820, 57.66521966527196
+20130821, 71.53396149320083
+20130822, 101.2599978934217
+20130823, 140.78798832665038
+20130824, 89.85586246222688
+20130825, 40.11625461268016
+20130829, 51.15046163121804
+20130830, 147.82179708672084
+20130831, 246.754305955079
+20130901, 113.76087139702464
+20130902, 79.31063386215712
+20130903, 168.32009875493958
+20130904, 272.31384984454905
+20130905, 91.41031605939098
+20130906, 69.47553917509299
+20130907, 39.205069952928866
+20130908, 21.31527142172246
+20130909, 31.588420720013954
+20130910, 44.12261992968386
+20130911, 52.30735867764992
+20130912, 79.7153144612971
+20130913, 103.39279368607623
+20130914, 72.49962335832173
+20130915, 48.97323882787076
+20130916, 76.44151357653416
+20130917, 53.45604969345652
+20130918, 54.43395295792656
+20130919, 29.099646603324036
+20130922, 367.8701579570549
+20130923, 249.235258854893
+20130924, 32.61869243229893
+20130925, 35.36432091904928
+20130928, 73.08063255462574
+20130929, 243.00749633164807
+20130930, 306.4345184652487
+20131001, 40.19865033705253
+20131002, 34.479447059507216
+20131004, 21.415825669746635
+20131009, 25.32022790136535
+20131010, 32.18095126539982
+20131011, 31.99367899378197
+20131012, 30.35447956619015
+20131013, 30.706253450430033
+20131014, 142.5235437369247
+20131015, 43.649596299686195
+20131016, 25.725661756159933
+20131017, 21.21159925615992
+20131018, 28.128024392724317
+20131022, 32.32290086732914
+20131023, 35.4515408894119
+20131029, 38.405682676662
+20131030, 31.816779986053007
+20131031, 31.90093615614831
+20131101, 27.61662816713157
+20131102, 115.37988308490236
+20131103, 149.52505230125522
+20131104, 78.23612200720596
+20131105, 41.96016533007903
+20131106, 26.114996295327753
+20131107, 42.38168293816829
+20131110, 136.3468579294514
+20131112, 94.41665213854021
+20131115, 39.82262229050443
+20131116, 78.61473497065319
+20131117, 48.18344810553231
+20131118, 29.973857907659227
+20131120, 20.604020114191066
+20131121, 51.136010503835415
+20131122, 66.37020844229428
+20131124, 76.96372798988841
+20131127, 50.039894416841015
+20131128, 75.84964823773826
+20131213, 79.16618905450954
+20131214, 211.67638525685723
+20131215, 420.3117846713737
+20131216, 342.96863105096475
+20131217, 478.720224241632
+20140208, 23.630335454439784
+20140209, 75.14158494595539
+20140210, 27.744407034518833
+20140213, 94.30788241660855
+20140219, 172.0230234483961
+20140222, 69.79637305323105
+20140309, 65.87043417305905
+20140310, 64.25986641387725
+20140311, 24.917366922361698
+20140312, 29.19463512610414
+20140313, 25.458290293758722
+20140314, 31.69204403475128
+20140318, 23.423604755055795
+20140319, 39.21903130084844
+20140327, 23.41754380230125
+20140329, 91.91434470885638
+20140330, 292.9196107914923
+20140331, 253.80143901092512
+20140401, 106.97886139440962
+20140402, 181.600617990179
+20140403, 252.41761190289398
+20140404, 20.795628849953516
+20140405, 21.720432357043236
+20140406, 300.43859778155513
+20140408, 78.17030487273361
+20140409, 32.890573061947926
+20140410, 22.470932850999535
+20140413, 25.260595907426776
+20140418, 20.41067762813808
+20140426, 113.97868723849369
+20140427, 44.45263358612272
+20140429, 52.85484658298466
+20140430, 126.49033480067413
+20140501, 128.88813433432125
+20140502, 35.50418646123896
+20140503, 68.03873616195956
+20140504, 124.95661828800556
+20140505, 174.92232119508373
+20140506, 88.36784852975359
+20140507, 253.31558577405852
+20140508, 281.3919763217787
+20140509, 480.06120300151093
+20140510, 230.49119141533015
+20140511, 205.4820327028126
+20140512, 156.76944699754904
+20140513, 81.30908298465832
+20140514, 21.89773052504649
+20140515, 39.947487178928405
+20140516, 88.09127694967457
+20140517, 183.43590390370755
+20140518, 105.40262759327061
+20140519, 112.52293954846583
+20140520, 143.1907039966876
+20140521, 123.768075713331
+20140522, 146.64557309826824
+20140523, 194.98976075807767
+20140524, 40.74365030073222
+20140526, 36.89792792596466
+20140527, 25.41117921170386
+20140528, 32.128269373256636
+20140529, 42.571904237854476
+20140530, 30.894538332461654
+20140603, 119.3601629692585
+20140604, 28.660918250232445
+20140605, 106.6490854544398
+20140606, 169.37256871803808
+20140607, 154.40959328509996
+20140608, 53.36430057967224
+20140609, 67.63162573367039
+20140610, 89.99786563662249
+20140611, 37.91046715190608
+20140612, 36.04271668700604
+20140613, 22.212055620932123
+20140614, 60.21817922623198
+20140615, 160.63765908298467
+20140616, 54.423535564853566
+20140617, 58.11623173088098
+20140618, 38.36407049046953
+20140619, 74.36847196798001
+20140620, 77.7738283066016
+20140621, 119.29349666579498
+20140622, 187.1327440362041
+20140623, 111.1517991268596
+20140624, 65.32679131799162
+20140625, 71.67983804771039
+20140626, 23.243867677824277
+20140629, 50.502375711878194
+20140630, 44.06397496803812
+20140701, 70.78839747501164
+20140702, 21.4579365338796
+20140705, 20.564100636331936
+20140706, 59.28443547332637
+20140707, 55.36132866980474
+20140708, 94.93222883251977
+20140709, 24.205018559681545
+20140710, 65.32519667451186
+20140711, 57.94388838040446
+20140712, 61.87451693979545
+20140713, 36.246024196594604
+20140714, 24.11177268566946
+20140717, 86.05393040301021
+20140718, 170.54987887174568
+20140719, 70.18380277341934
+20140720, 21.389537206531852
+20140721, 24.904192817294284
+20140723, 33.28208209844259
+20140724, 98.2812189461297
+20140725, 66.55174991283123
+20140726, 89.153669804742
+20140727, 55.76670643741284
+20140801, 77.50097955892608
+20140802, 47.23127941945607
+20140803, 24.28849445025569
+20140804, 27.641382823396093
+20140805, 59.530193986808456
+20140806, 107.75092253602975
+20140807, 94.35307978120642
+20140808, 57.08875014528127
+20140809, 76.8537945650279
+20140810, 34.28849717427941
+20140811, 51.66818957752208
+20140812, 197.83517677097862
+20140813, 357.7546878632031
+20140814, 122.80910241602739
+20140818, 40.27715887959089
+20140819, 327.54424032426783
+20140820, 286.2331065129592
+20140821, 63.01878777312877
+20140822, 32.8901639135867
+20140825, 21.239476006799162
+20140826, 28.87913172375236
+20140827, 124.34361852045556
+20140828, 78.05018559681542
+20140829, 23.721888438516967
+20140830, 34.31858492416317
+20140831, 62.38237102655741
+20140903, 25.001808933344957
+20140904, 34.299766278765695
+20140905, 64.1351253777313
+20140906, 40.08666572233845
+20140907, 86.35310429741982
+20140908, 49.52347545473036
+20140909, 29.29712361256393
+20140910, 37.02630753138075
+20140911, 40.57987055439331
+20140912, 80.21033512755693
+20140913, 45.57955347803347
+20140914, 56.27964918206646
+20140915, 131.16330304218968
+20140916, 303.1881555817061
+20140917, 76.85736303608788
+20140918, 51.148375755462574
+20140919, 67.16068652661554
+20140920, 24.000558606462118
+20140923, 21.921376866864247
+20140924, 25.080225948686653
+20140925, 35.09464892782425
+20140926, 39.96193594549049
+20140927, 43.513512792015334
+20140928, 48.600577129823336
+20140929, 65.18053576098326
+20140930, 78.84402938258636
+20141001, 72.51337659082985
+20141002, 50.72223438952814
+20141003, 76.47042702497669
+20141004, 126.25225493843371
+20141005, 64.90787206531846
+20141011, 20.986212263191543
+20141012, 20.71796657078103
+20141016, 23.756614474372387
+20141017, 35.27394997966062
+20141018, 23.492493317061832
+20141019, 20.595106107887954
+20141021, 40.95922790271968
+20141022, 23.438181187529064
+20141023, 43.47170483931891
+20141024, 46.727642645599914
+20141025, 51.37816004910506
+20141026, 94.35662713001871
+20141027, 68.03675961035566
+20141028, 21.424196957810324
+20141030, 24.947406410480855
+20141031, 34.73852623053231
+20141101, 42.339237927126916
+20141104, 21.830416267143192
+20141106, 43.89685048378662
+20141107, 69.33576352568573
+20141108, 87.40360715219663
+20141110, 29.08102771966528
+20141111, 37.43220394928804
+20141112, 29.107394453161323
+20141113, 20.995226239249185
+20141114, 23.49454341730591
+20141115, 26.82312696129707
+20141116, 29.49105357973036
+20141201, 33.785750450371914
+20141202, 43.8162371980042
+20141203, 21.638188269990703
+20141204, 52.66889273884239
+20141205, 20.34541219926778
+20141207, 22.521633833100882
+20141218, 25.065987839958158
+20141219, 117.93202650852008
+20141225, 87.132361220072
+20141226, 20.54271959263134
+20141227, 48.39155815972224
+20141228, 34.44478494740817

data/HKO-7/samplers/hko7_cloudy_days_t20_train.txt.pkl

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

data/SEVIR/CATALOG.csv

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

data/config.py

@@ -0,0 +1,62 @@
+from data import dutils 
+
+SEVIR_13_12 = {
+    'meta': {
+        'dataset': 'SEVIR',    
+        'seq_len': 13,
+        'out_len': 12,
+        'metrics': ['mae', 'mse', 'ssim', 'psnr', 'lpips64', 
+                    'csi-16', 'csi-74', 'csi-133', 'csi-160', 'csi-181', 'csi-219',
+                    'csi_4-16', 'csi_4-74', 'csi_4-133', 'csi_4-160', 'csi_4-181', 'csi_4-219',
+                    'csi_16-16', 'csi_16-74', 'csi_16-133', 'csi_16-160', 'csi_16-181', 'csi_16-219',
+                    'hss-16', 'hss-74', 'hss-133', 'hss-160', 'hss-181', 'hss-219' ],
+    },
+    'param': {        
+        'seq_len': 25,
+        'data_types': ['vil'], 
+        'sample_mode': 'sequent',
+        'layout': 'NTCHW', 
+        'raw_seq_len': 25, 
+        'start_date': dutils.SEVIR_TRAIN_TEST_SPLIT_DATE,
+        'end_date': None,
+    },
+    'savedir': 'sevir'
+}
+
+HKO7_5_20 = {
+    'meta': {
+        'dataset': 'HKO-7',
+        'seq_len': 5,
+        'out_len': 20,
+        'metrics': ['mae', 'mse', 'ssim', 'psnr', 'lpips64', 
+                    'csi-84', 'csi-117', 'csi-140', 'csi-158', 'csi-185',
+                    'csi_4-84', 'csi_4-117', 'csi_4-140', 'csi_4-158', 'csi_4-185',
+                    'csi_16-84', 'csi_16-117', 'csi_16-140', 'csi_16-158', 'csi_16-185',
+                    'hss-84', 'hss-117', 'hss-140', 'hss-158', 'hss-185'],
+    },
+    'param': {            
+        'pd_path': 'data/HKO-7/samplers/hko7_cloudy_days_t20_test.txt.pkl',
+        'sample_mode': 'sequent',
+        'seq_len': 25,
+        'stride': 13,
+    },
+    'savedir': 'hko-7'
+}
+
+METEONET_5_20 = {
+    'meta': {
+        'dataset': 'meteonet',
+        'seq_len': 5,
+        'out_len': 20,
+        'metrics': ['mae', 'mse', 'ssim', 'psnr', 'lpips64', 
+                    'csi-44', 'csi-64', 'csi-87', 'csi-117',
+                    'csi_4-44', 'csi_4-64', 'csi_4-87', 'csi_4-117', 
+                    'csi_16-44', 'csi_16-64', 'csi_16-87', 'csi_16-117',
+                    'hss-44', 'hss-64', 'hss-87', 'hss-117']
+    },
+    'param': {        
+        'img_size': 128,
+        'in_len': 5, 
+    },
+    'savedir': 'meteonet'
+}

data/dutils.py

@@ -0,0 +1,1212 @@
+'''
+dutils.py
+A utility library for customized data loading functions
+'''
+import os
+import gzip
+import numpy as np
+import pandas as pd
+
+import os
+import cv2
+from typing import List, Union, Dict, Sequence
+import numpy as np
+import numpy.random as nprand
+import datetime
+import pandas as pd
+import h5py
+import torch
+import torch.nn.functional as F
+from torch.nn.functional import avg_pool2d
+import random
+from torchvision import transforms as T
+from torchvision import datasets
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+
+SEVIR_ROOT_DIR = "data/SEVIR"
+METEO_FILE_DIR = "data/meteonet"
+
+def resize(seq, size):
+    # seq shape : (B, T, 1, H, W)
+    seq = F.interpolate(seq.squeeze(dim=2), size=size, mode='bilinear', align_corners=False) # (B, T, H, W)
+    seq = seq.clamp(0,1)
+    return seq.unsqueeze(2) # (B, T, 1, H, W)
+
+# =====================================================================================
+# HKO-7 data
+# =====================================================================================
+def pixel_to_dBZ_nonlinear(img):
+    '''
+    [0, 255] OR [0, 1] pixel => [0, 80] dBZ
+    '''
+    if img.mean() > 1.0:
+        img = img / 255.0
+    ashift = 31.0
+    afact = 4.0
+    atan_dBZ_min = -1.482
+    atan_dBZ_max = 1.412
+    tan_pix = np.tan(img * (atan_dBZ_max - atan_dBZ_min) + atan_dBZ_min) 
+    return tan_pix * afact + ashift
+
+def dbZ_to_pixel_nonlinear(dbZ):
+    '''
+    [0, 80] dBZ => [0, 255] OR [0, 1] pixel
+    '''
+    ashift = 31.0
+    afact = 4.0
+    atan_dBZ_min = -1.482
+    atan_dBZ_max = 1.412
+    dbZ_adjusted = (dbZ - ashift) / afact
+    return (np.arctan(dbZ_adjusted) - atan_dBZ_min) / (atan_dBZ_max - atan_dBZ_min)
+
+def dbZ_to_pixel(dbZ):
+    '''
+    [0, 80] dbZ => [0, 1] pixel
+    '''
+    return np.floor((dbZ + 10) * 255 / 70 + 0.5) / 255.0
+
+def pixel_to_dBZ(pixel):
+    '''
+    [0, 255] (or [0, 1]) pixel => [0, 80] dBZ
+    '''
+    if pixel.mean() > 1.0:
+        pixel = pixel / 255.0
+    return (70 * pixel) - 10
+
+def nonlinear_to_linear(im):
+    return dbZ_to_pixel(pixel_to_dBZ_nonlinear(im))
+
+def nonlinear_to_linear_batched(seq, datetime):
+    seq_linear = np.zeros_like(seq)
+    for i, (seq_b, dt_b) in enumerate(zip(seq, datetime)):
+        if dt_b[0].year >= 2016:
+            seq_linear[i] = nonlinear_to_linear(seq_b)
+        else:
+            seq_linear[i] = seq_b
+    seq_linear = np.clip(seq_linear, 0.0, 1.0)
+    return seq_linear
+
+def linear_to_nonlinear(im):
+    return dbZ_to_pixel_nonlinear(pixel_to_dBZ(im))
+
+def linear_to_nonlinear_batched(seq, datetime):
+    seq_nonlinear = np.zeros_like(seq)
+    for i, (seq_b, dt_b) in enumerate(zip(seq, datetime)):
+        if dt_b[0].year < 2016:
+            seq_nonlinear[i] = linear_to_nonlinear(seq_b)
+        else:
+            seq_nonlinear[i] = seq_b
+    seq_nonlinear = np.clip(seq_nonlinear, 0.0, 1.0)
+    return seq_nonlinear
+
+
+# =====================================================================================
+# SEVIR data
+# Code is adapted from https://github.com/MIT-AI-Accelerator/neurips-2020-sevir. Their license is MIT License.
+# (From Earthformer's implementation)
+# =====================================================================================
+
+
+# SEVIR Dataset constants
+SEVIR_DATA_TYPES = ['vis', 'ir069', 'ir107', 'vil', 'lght']
+SEVIR_RAW_DTYPES = {'vis': np.int16,
+                    'ir069': np.int16,
+                    'ir107': np.int16,
+                    'vil': np.uint8,
+                    'lght': np.int16}
+LIGHTING_FRAME_TIMES = np.arange(- 120.0, 125.0, 5) * 60
+SEVIR_DATA_SHAPE = {'lght': (48, 48), }
+PREPROCESS_SCALE_SEVIR = {'vis': 1,  # Not utilized in original paper
+                          'ir069': 1 / 1174.68,
+                          'ir107': 1 / 2562.43,
+                          'vil': 1 / 47.54,
+                          'lght': 1 / 0.60517}
+PREPROCESS_OFFSET_SEVIR = {'vis': 0,  # Not utilized in original paper
+                           'ir069': 3683.58,
+                           'ir107': 1552.80,
+                           'vil': - 33.44,
+                           'lght': - 0.02990}
+PREPROCESS_SCALE_01 = {'vis': 1,
+                       'ir069': 1,
+                       'ir107': 1,
+                       'vil': 1 / 255,  # currently the only one implemented
+                       'lght': 1}
+PREPROCESS_OFFSET_01 = {'vis': 0,
+                        'ir069': 0,
+                        'ir107': 0,
+                        'vil': 0,  # currently the only one implemented
+                        'lght': 0}
+
+# sevir
+SEVIR_CATALOG = os.path.join(SEVIR_ROOT_DIR, "CATALOG.csv")
+SEVIR_DATA_DIR = os.path.join(SEVIR_ROOT_DIR, "data")
+SEVIR_RAW_SEQ_LEN = 49
+
+SEVIR_TRAIN_VAL_SPLIT_DATE = datetime.datetime(2019, 1, 1)
+SEVIR_TRAIN_TEST_SPLIT_DATE = datetime.datetime(2019, 6, 1)
+
+def change_layout_np(data,
+                     in_layout='NHWT', out_layout='NHWT',
+                     ret_contiguous=False):
+    # first convert to 'NHWT'
+    if in_layout == 'NHWT':
+        pass
+    elif in_layout == 'NTHW':
+        data = np.transpose(data,
+                            axes=(0, 2, 3, 1))
+    elif in_layout == 'NWHT':
+        data = np.transpose(data,
+                            axes=(0, 2, 1, 3))
+    elif in_layout == 'NTCHW':
+        data = data[:, :, 0, :, :]
+        data = np.transpose(data,
+                            axes=(0, 2, 3, 1))
+    elif in_layout == 'NTHWC':
+        data = data[:, :, :, :, 0]
+        data = np.transpose(data,
+                            axes=(0, 2, 3, 1))
+    elif in_layout == 'NTWHC':
+        data = data[:, :, :, :, 0]
+        data = np.transpose(data,
+                            axes=(0, 3, 2, 1))
+    elif in_layout == 'TNHW':
+        data = np.transpose(data,
+                            axes=(1, 2, 3, 0))
+    elif in_layout == 'TNCHW':
+        data = data[:, :, 0, :, :]
+        data = np.transpose(data,
+                            axes=(1, 2, 3, 0))
+    else:
+        raise NotImplementedError
+
+    if out_layout == 'NHWT':
+        pass
+    elif out_layout == 'NTHW':
+        data = np.transpose(data,
+                            axes=(0, 3, 1, 2))
+    elif out_layout == 'NWHT':
+        data = np.transpose(data,
+                            axes=(0, 2, 1, 3))
+    elif out_layout == 'NTCHW':
+        data = np.transpose(data,
+                            axes=(0, 3, 1, 2))
+        data = np.expand_dims(data, axis=2)
+    elif out_layout == 'NTHWC':
+        data = np.transpose(data,
+                            axes=(0, 3, 1, 2))
+        data = np.expand_dims(data, axis=-1)
+    elif out_layout == 'NTWHC':
+        data = np.transpose(data,
+                            axes=(0, 3, 2, 1))
+        data = np.expand_dims(data, axis=-1)
+    elif out_layout == 'TNHW':
+        data = np.transpose(data,
+                            axes=(3, 0, 1, 2))
+    elif out_layout == 'TNCHW':
+        data = np.transpose(data,
+                            axes=(3, 0, 1, 2))
+        data = np.expand_dims(data, axis=2)
+    else:
+        raise NotImplementedError
+    if ret_contiguous:
+        data = data.ascontiguousarray()
+    return data
+
+def change_layout_torch(data,
+                        in_layout='NHWT', out_layout='NHWT',
+                        ret_contiguous=False):
+    # first convert to 'NHWT'
+    if in_layout == 'NHWT':
+        pass
+    elif in_layout == 'NTHW':
+        data = data.permute(0, 2, 3, 1)
+    elif in_layout == 'NTCHW':
+        data = data[:, :, 0, :, :]
+        data = data.permute(0, 2, 3, 1)
+    elif in_layout == 'NTHWC':
+        data = data[:, :, :, :, 0]
+        data = data.permute(0, 2, 3, 1)
+    elif in_layout == 'TNHW':
+        data = data.permute(1, 2, 3, 0)
+    elif in_layout == 'TNCHW':
+        data = data[:, :, 0, :, :]
+        data = data.permute(1, 2, 3, 0)
+    else:
+        raise NotImplementedError
+
+    if out_layout == 'NHWT':
+        pass
+    elif out_layout == 'NTHW':
+        data = data.permute(0, 3, 1, 2)
+    elif out_layout == 'NTCHW':
+        data = data.permute(0, 3, 1, 2)
+        data = torch.unsqueeze(data, dim=2)
+    elif out_layout == 'NTHWC':
+        data = data.permute(0, 3, 1, 2)
+        data = torch.unsqueeze(data, dim=-1)
+    elif out_layout == 'TNHW':
+        data = data.permute(3, 0, 1, 2)
+    elif out_layout == 'TNCHW':
+        data = data.permute(3, 0, 1, 2)
+        data = torch.unsqueeze(data, dim=2)
+    else:
+        raise NotImplementedError
+    if ret_contiguous:
+        data = data.contiguous()
+    return data
+
+class SEVIRDataLoader:
+    r"""
+    DataLoader that loads SEVIR sequences, and spilts each event
+    into segments according to specified sequence length.
+
+    Event Frames:
+        [-----------------------raw_seq_len----------------------]
+        [-----seq_len-----]
+        <--stride-->[-----seq_len-----]
+                    <--stride-->[-----seq_len-----]
+                                        ...
+    """
+    def __init__(self,
+                 data_types: Sequence[str] = None,
+                 seq_len: int = 49,
+                 raw_seq_len: int = 49,
+                 sample_mode: str = 'sequent',
+                 stride: int = 12,
+                 batch_size: int = 1,
+                 layout: str = 'NHWT',
+                 num_shard: int = 1,
+                 rank: int = 0,
+                 split_mode: str = "uneven",
+                 sevir_catalog: Union[str, pd.DataFrame] = None,
+                 sevir_data_dir: str = None,
+                 start_date: datetime.datetime = None,
+                 end_date: datetime.datetime = None,
+                 datetime_filter=None,
+                 catalog_filter='default',
+                 shuffle: bool = False,
+                 shuffle_seed: int = 1,
+                 output_type=np.float32,
+                 preprocess: bool = True,
+                 rescale_method: str = '01',
+                 downsample_dict: Dict[str, Sequence[int]] = None,
+                 verbose: bool = False):
+        r"""
+        Parameters
+        ----------
+        data_types
+            A subset of SEVIR_DATA_TYPES.
+        seq_len
+            The length of the data sequences. Should be smaller than the max length raw_seq_len.
+        raw_seq_len
+            The length of the raw data sequences.
+        sample_mode
+            'random' or 'sequent'
+        stride
+            Useful when sample_mode == 'sequent'
+            stride must not be smaller than out_len to prevent data leakage in testing.
+        batch_size
+            Number of sequences in one batch.
+        layout
+            str: consists of batch_size 'N', seq_len 'T', channel 'C', height 'H', width 'W'
+            The layout of sampled data. Raw data layout is 'NHWT'.
+            valid layout: 'NHWT', 'NTHW', 'NTCHW', 'TNHW', 'TNCHW'.
+        num_shard
+            Split the whole dataset into num_shard parts for distributed training.
+        rank
+            Rank of the current process within num_shard.
+        split_mode: str
+            if 'ceil', all `num_shard` dataloaders have the same length = ceil(total_len / num_shard).
+            Different dataloaders may have some duplicated data batches, if the total size of datasets is not divided by num_shard.
+            if 'floor', all `num_shard` dataloaders have the same length = floor(total_len / num_shard).
+            The last several data batches may be wasted, if the total size of datasets is not divided by num_shard.
+            if 'uneven', the last datasets has larger length when the total length is not divided by num_shard.
+            The uneven split leads to synchronization error in dist.all_reduce() or dist.barrier().
+            See related issue: https://github.com/pytorch/pytorch/issues/33148
+            Notice: this also affects the behavior of `self.use_up`.
+        sevir_catalog
+            Name of SEVIR catalog CSV file.
+        sevir_data_dir
+            Directory path to SEVIR data.
+        start_date
+            Start time of SEVIR samples to generate.
+        end_date
+            End time of SEVIR samples to generate.
+        datetime_filter
+            function
+            Mask function applied to time_utc column of catalog (return true to keep the row).
+            Pass function of the form   lambda t : COND(t)
+            Example:  lambda t: np.logical_and(t.dt.hour>=13,t.dt.hour<=21)  # Generate only day-time events
+        catalog_filter
+            function or None or 'default'
+            Mask function applied to entire catalog dataframe (return true to keep row).
+            Pass function of the form lambda catalog:  COND(catalog)
+            Example:  lambda c:  [s[0]=='S' for s in c.id]   # Generate only the 'S' events
+        shuffle
+            bool, If True, data samples are shuffled before each epoch.
+        shuffle_seed
+            int, Seed to use for shuffling.
+        output_type
+            np.dtype, dtype of generated tensors
+        preprocess
+            bool, If True, self.preprocess_data_dict(data_dict) is called before each sample generated
+        downsample_dict:
+            dict, downsample_dict.keys() == data_types. downsample_dict[key] is a Sequence of (t_factor, h_factor, w_factor),
+            representing the downsampling factors of all dimensions.
+        verbose
+            bool, verbose when opening raw data files
+        """
+        super(SEVIRDataLoader, self).__init__()
+        if sevir_catalog is None:
+            sevir_catalog = SEVIR_CATALOG
+        if sevir_data_dir is None:
+            sevir_data_dir = SEVIR_DATA_DIR
+        if data_types is None:
+            data_types = SEVIR_DATA_TYPES
+        else:
+            assert set(data_types).issubset(SEVIR_DATA_TYPES)
+
+        # configs which should not be modified
+        self._dtypes = SEVIR_RAW_DTYPES
+        self.lght_frame_times = LIGHTING_FRAME_TIMES
+        self.data_shape = SEVIR_DATA_SHAPE
+
+        self.raw_seq_len = raw_seq_len
+        assert seq_len <= self.raw_seq_len, f'seq_len must not be larger than raw_seq_len = {raw_seq_len}, got {seq_len}.'
+        self.seq_len = seq_len
+        assert sample_mode in ['random', 'sequent'], f'Invalid sample_mode = {sample_mode}, must be \'random\' or \'sequent\'.'
+        self.sample_mode = sample_mode
+        self.stride = stride
+        self.batch_size = batch_size
+        valid_layout = ('NHWT', 'NTHW', 'NTCHW', 'NTHWC', 'TNHW', 'TNCHW')
+        if layout not in valid_layout:
+            raise ValueError(f'Invalid layout = {layout}! Must be one of {valid_layout}.')
+        self.layout = layout
+        self.num_shard = num_shard
+        self.rank = rank
+        valid_split_mode = ('ceil', 'floor', 'uneven')
+        if split_mode not in valid_split_mode:
+            raise ValueError(f'Invalid split_mode: {split_mode}! Must be one of {valid_split_mode}.')
+        self.split_mode = split_mode
+        self._samples = None
+        self._hdf_files = {}
+        self.data_types = data_types
+        if isinstance(sevir_catalog, str):
+            self.catalog = pd.read_csv(sevir_catalog, parse_dates=['time_utc'], low_memory=False)
+        else:
+            self.catalog = sevir_catalog
+        self.sevir_data_dir = sevir_data_dir
+        self.datetime_filter = datetime_filter
+        self.catalog_filter = catalog_filter
+        self.start_date = start_date
+        self.end_date = end_date
+        self.shuffle = shuffle
+        self.shuffle_seed = int(shuffle_seed)
+        self.output_type = output_type
+        self.preprocess = preprocess
+        self.downsample_dict = downsample_dict
+        self.rescale_method = rescale_method
+        self.verbose = verbose
+
+        if self.start_date is not None:
+            self.catalog = self.catalog[self.catalog.time_utc > self.start_date]
+        if self.end_date is not None:
+            self.catalog = self.catalog[self.catalog.time_utc <= self.end_date]
+        if self.datetime_filter:
+            self.catalog = self.catalog[self.datetime_filter(self.catalog.time_utc)]
+
+        if self.catalog_filter is not None:
+            if self.catalog_filter == 'default':
+                self.catalog_filter = lambda c: c.pct_missing == 0
+            self.catalog = self.catalog[self.catalog_filter(self.catalog)]
+
+        self._compute_samples()
+        self._open_files(verbose=self.verbose)
+        self.reset()
+
+    def _compute_samples(self):
+        """
+        Computes the list of samples in catalog to be used. This sets self._samples
+        """
+        # locate all events containing colocated data_types
+        imgt = self.data_types
+        imgts = set(imgt)
+        filtcat = self.catalog[ np.logical_or.reduce([self.catalog.img_type==i for i in imgt]) ]
+        # remove rows missing one or more requested img_types
+        filtcat = filtcat.groupby('id').filter(lambda x: imgts.issubset(set(x['img_type'])))
+        # If there are repeated IDs, remove them (this is a bug in SEVIR)
+        # TODO: is it necessary to keep one of them instead of deleting them all
+        filtcat = filtcat.groupby('id').filter(lambda x: x.shape[0]==len(imgt))
+        self._samples = filtcat.groupby('id').apply(lambda df: self._df_to_series(df,imgt) )
+        if self.shuffle:
+            self.shuffle_samples()
+
+    def shuffle_samples(self):
+        self._samples = self._samples.sample(frac=1, random_state=self.shuffle_seed)
+
+    def _df_to_series(self, df, imgt):
+        d = {}
+        df = df.set_index('img_type')
+        for i in imgt:
+            s = df.loc[i]
+            idx = s.file_index if i != 'lght' else s.id
+            d.update({f'{i}_filename': [s.file_name],
+                      f'{i}_index': [idx]})
+
+        return pd.DataFrame(d)
+
+    def _open_files(self, verbose=True):
+        """
+        Opens HDF files
+        """
+        imgt = self.data_types
+        hdf_filenames = []
+        for t in imgt:
+            hdf_filenames += list(np.unique( self._samples[f'{t}_filename'].values ))
+        self._hdf_files = {}
+        for f in hdf_filenames:
+            if verbose:
+                print('Opening HDF5 file for reading', f)
+            self._hdf_files[f] = h5py.File(self.sevir_data_dir + '/' + f, 'r')
+
+    def close(self):
+        """
+        Closes all open file handles
+        """
+        for f in self._hdf_files:
+            self._hdf_files[f].close()
+        self._hdf_files = {}
+
+    @property
+    def num_seq_per_event(self):
+        return 1 + (self.raw_seq_len - self.seq_len) // self.stride
+
+    @property
+    def total_num_seq(self):
+        """
+        The total number of sequences within each shard.
+        Notice that it is not the product of `self.num_seq_per_event` and `self.total_num_event`.
+        """
+        return int(self.num_seq_per_event * self.num_event)
+
+    @property
+    def total_num_event(self):
+        """
+        The total number of events in the whole dataset, before split into different shards.
+        """
+        return int(self._samples.shape[0])
+
+    @property
+    def start_event_idx(self):
+        """
+        The event idx used in certain rank should satisfy event_idx >= start_event_idx
+        """
+        return self.total_num_event // self.num_shard * self.rank
+
+    @property
+    def end_event_idx(self):
+        """
+        The event idx used in certain rank should satisfy event_idx < end_event_idx
+
+        """
+        if self.split_mode == 'ceil':
+            _last_start_event_idx = self.total_num_event // self.num_shard * (self.num_shard - 1)
+            _num_event = self.total_num_event - _last_start_event_idx
+            return self.start_event_idx + _num_event
+        elif self.split_mode == 'floor':
+            return self.total_num_event // self.num_shard * (self.rank + 1)
+        else:  # self.split_mode == 'uneven':
+            if self.rank == self.num_shard - 1:  # the last process
+                return self.total_num_event
+            else:
+                return self.total_num_event // self.num_shard * (self.rank + 1)
+
+    @property
+    def num_event(self):
+        """
+        The number of events split into each rank
+        """
+        return self.end_event_idx - self.start_event_idx
+
+    def _read_data(self, row, data):
+        """
+        Iteratively read data into data dict. Finally data[imgt] gets shape (batch_size, height, width, raw_seq_len).
+
+        Parameters
+        ----------
+        row
+            A series with fields IMGTYPE_filename, IMGTYPE_index, IMGTYPE_time_index.
+        data
+            Dict, data[imgt] is a data tensor with shape = (tmp_batch_size, height, width, raw_seq_len).
+
+        Returns
+        -------
+        data
+            Updated data. Updated shape = (tmp_batch_size + 1, height, width, raw_seq_len).
+        """
+        imgtyps = np.unique([x.split('_')[0] for x in list(row.keys())])
+        for t in imgtyps:
+            fname = row[f'{t}_filename']
+            idx = row[f'{t}_index']
+            t_slice = slice(0, None)
+            # Need to bin lght counts into grid
+            if t == 'lght':
+                lght_data = self._hdf_files[fname][idx][:]
+                data_i = self._lght_to_grid(lght_data, t_slice)
+            else:
+                data_i = self._hdf_files[fname][t][idx:idx + 1, :, :, t_slice]
+            data[t] = np.concatenate((data[t], data_i), axis=0) if (t in data) else data_i
+
+        return data
+
+    def _lght_to_grid(self, data, t_slice=slice(0, None)):
+        """
+        Converts Nx5 lightning data matrix into a 2D grid of pixel counts
+        """
+        # out_size = (48,48,len(self.lght_frame_times)-1) if isinstance(t_slice,(slice,)) else (48,48)
+        out_size = (*self.data_shape['lght'], len(self.lght_frame_times)) if t_slice.stop is None else (*self.data_shape['lght'], 1)
+        if data.shape[0] == 0:
+            return np.zeros((1,) + out_size, dtype=np.float32)
+
+        # filter out points outside the grid
+        x, y = data[:, 3], data[:, 4]
+        m = np.logical_and.reduce([x >= 0, x < out_size[0], y >= 0, y < out_size[1]])
+        data = data[m, :]
+        if data.shape[0] == 0:
+            return np.zeros((1,) + out_size, dtype=np.float32)
+
+        # Filter/separate times
+        t = data[:, 0]
+        if t_slice.stop is not None:  # select only one time bin
+            if t_slice.stop > 0:
+                if t_slice.stop < len(self.lght_frame_times):
+                    tm = np.logical_and(t >= self.lght_frame_times[t_slice.stop - 1],
+                                        t < self.lght_frame_times[t_slice.stop])
+                else:
+                    tm = t >= self.lght_frame_times[-1]
+            else:  # special case:  frame 0 uses lght from frame 1
+                tm = np.logical_and(t >= self.lght_frame_times[0], t < self.lght_frame_times[1])
+            # tm=np.logical_and( (t>=FRAME_TIMES[t_slice],t<FRAME_TIMES[t_slice+1]) )
+
+            data = data[tm, :]
+            z = np.zeros(data.shape[0], dtype=np.int64)
+        else:  # compute z coordinate based on bin location times
+            z = np.digitize(t, self.lght_frame_times) - 1
+            z[z == -1] = 0  # special case:  frame 0 uses lght from frame 1
+
+        x = data[:, 3].astype(np.int64)
+        y = data[:, 4].astype(np.int64)
+
+        k = np.ravel_multi_index(np.array([y, x, z]), out_size)
+        n = np.bincount(k, minlength=np.prod(out_size))
+        return np.reshape(n, out_size).astype(np.int16)[np.newaxis, :]
+
+    def _old_save_downsampled_dataset(self, save_dir, downsample_dict, verbose=True):
+        """
+        This method does not save .h5 dataset correctly. There are some batches missed due to unknown error.
+        E.g., the first converted .h5 file `SEVIR_VIL_RANDOMEVENTS_2017_0501_0831.h5` only has batch_dim = 1414,
+        while it should be 1440 in the original .h5 file.
+        """
+        import os
+        from skimage.measure import block_reduce
+        assert not os.path.exists(save_dir), f"save_dir {save_dir} already exists!"
+        os.makedirs(save_dir)
+        sample_counter = 0
+        for index, row in self._samples.iterrows():
+            if verbose:
+                print(f"Downsampling {sample_counter}-th data item.", end='\r')
+            for data_type in self.data_types:
+                fname = row[f'{data_type}_filename']
+                idx = row[f'{data_type}_index']
+                t_slice = slice(0, None)
+                if data_type == 'lght':
+                    lght_data = self._hdf_files[fname][idx][:]
+                    data_i = self._lght_to_grid(lght_data, t_slice)
+                else:
+                    data_i = self._hdf_files[fname][data_type][idx:idx + 1, :, :, t_slice]
+                # Downsample t
+                t_slice = [slice(None, None), ] * 4
+                t_slice[-1] = slice(None, None, downsample_dict[data_type][0])  # layout = 'NHWT'
+                data_i = data_i[tuple(t_slice)]
+                # Downsample h, w
+                data_i = block_reduce(data_i,
+                                      block_size=(1, *downsample_dict[data_type][1:], 1),
+                                      func=np.max)
+                # Save as new .h5 file
+                new_file_path = os.path.join(save_dir, fname)
+                if not os.path.exists(new_file_path):
+                    if not os.path.exists(os.path.dirname(new_file_path)):
+                        os.makedirs(os.path.dirname(new_file_path))
+                    # Create dataset
+                    with h5py.File(new_file_path, 'w') as hf:
+                        hf.create_dataset(
+                            data_type, data=data_i,
+                            maxshape=(None, *data_i.shape[1:]))
+                else:
+                    # Append
+                    with h5py.File(new_file_path, 'a') as hf:
+                        hf[data_type].resize((hf[data_type].shape[0] + data_i.shape[0]), axis=0)
+                        hf[data_type][-data_i.shape[0]:] = data_i
+
+            sample_counter += 1
+
+    def save_downsampled_dataset(self, save_dir, downsample_dict, verbose=True):
+        """
+        Parameters
+        ----------
+        save_dir
+        downsample_dict:    Dict[Sequence[int]]
+            Notice that this is different from `self.downsample_dict`, which is used during runtime.
+        """
+        import os
+        from skimage.measure import block_reduce
+        from ...utils.utils import path_splitall
+        assert not os.path.exists(save_dir), f"save_dir {save_dir} already exists!"
+        os.makedirs(save_dir)
+        for fname, hdf_file in self._hdf_files.items():
+            if verbose:
+                print(f"Downsampling data in {fname}.")
+            data_type = path_splitall(fname)[0]
+            if data_type == 'lght':
+                # TODO: how to get idx?
+                raise NotImplementedError
+                # lght_data = self._hdf_files[fname][idx][:]
+                # t_slice = slice(0, None)
+                # data_i = self._lght_to_grid(lght_data, t_slice)
+            else:
+                data_i = self._hdf_files[fname][data_type]
+            # Downsample t
+            t_slice = [slice(None, None), ] * 4
+            t_slice[-1] = slice(None, None, downsample_dict[data_type][0])  # layout = 'NHWT'
+            data_i = data_i[tuple(t_slice)]
+            # Downsample h, w
+            data_i = block_reduce(data_i,
+                                  block_size=(1, *downsample_dict[data_type][1:], 1),
+                                  func=np.max)
+            # Save as new .h5 file
+            new_file_path = os.path.join(save_dir, fname)
+            if not os.path.exists(os.path.dirname(new_file_path)):
+                os.makedirs(os.path.dirname(new_file_path))
+            # Create dataset
+            with h5py.File(new_file_path, 'w') as hf:
+                hf.create_dataset(
+                    data_type, data=data_i,
+                    maxshape=(None, *data_i.shape[1:]))
+
+    @property
+    def sample_count(self):
+        """
+        Record how many times self.__next__() is called.
+        """
+        return self._sample_count
+
+    def inc_sample_count(self):
+        self._sample_count += 1
+
+    @property
+    def curr_event_idx(self):
+        return self._curr_event_idx
+
+    @property
+    def curr_seq_idx(self):
+        """
+        Used only when self.sample_mode == 'sequent'
+        """
+        return self._curr_seq_idx
+
+    def set_curr_event_idx(self, val):
+        self._curr_event_idx = val
+
+    def set_curr_seq_idx(self, val):
+        """
+        Used only when self.sample_mode == 'sequent'
+        """
+        self._curr_seq_idx = val
+
+    def reset(self, shuffle: bool = None):
+        self.set_curr_event_idx(val=self.start_event_idx)
+        self.set_curr_seq_idx(0)
+        self._sample_count = 0
+        if shuffle is None:
+            shuffle = self.shuffle
+        if shuffle:
+            self.shuffle_samples()
+
+    def __len__(self):
+        """
+        Used only when self.sample_mode == 'sequent'
+        """
+        return self.total_num_seq // self.batch_size
+
+    @property
+    def use_up(self):
+        """
+        Check if dataset is used up in 'sequent' mode.
+        """
+        if self.sample_mode == 'random':
+            return False
+        else:   # self.sample_mode == 'sequent'
+            # compute the remaining number of sequences in current event
+            curr_event_remain_seq = self.num_seq_per_event - self.curr_seq_idx
+            all_remain_seq = curr_event_remain_seq + (
+                        self.end_event_idx - self.curr_event_idx - 1) * self.num_seq_per_event
+            if self.split_mode == "floor":
+                # This approach does not cover all available data, but avoid dealing with masks
+                return all_remain_seq < self.batch_size
+            else:
+                return all_remain_seq <= 0
+
+    def _load_event_batch(self, event_idx, event_batch_size):
+        """
+        Loads a selected batch of events (not batch of sequences) into memory.
+
+        Parameters
+        ----------
+        idx
+        event_batch_size
+            event_batch[i] = all_type_i_available_events[idx:idx + event_batch_size]
+        Returns
+        -------
+        event_batch
+            list of event batches.
+            event_batch[i] is the event batch of the i-th data type.
+            Each event_batch[i] is a np.ndarray with shape = (event_batch_size, height, width, raw_seq_len)
+        """
+        event_idx_slice_end = event_idx + event_batch_size
+        pad_size = 0
+        if event_idx_slice_end > self.end_event_idx:
+            pad_size = event_idx_slice_end - self.end_event_idx
+            event_idx_slice_end = self.end_event_idx
+        pd_batch = self._samples.iloc[event_idx:event_idx_slice_end]
+        data = {}
+        for index, row in pd_batch.iterrows():
+            data = self._read_data(row, data)
+        if pad_size > 0:
+            event_batch = []
+            for t in self.data_types:
+                pad_shape = [pad_size, ] + list(data[t].shape[1:])
+                data_pad = np.concatenate((data[t].astype(self.output_type),
+                                           np.zeros(pad_shape, dtype=self.output_type)),
+                                          axis=0)
+                event_batch.append(data_pad)
+        else:
+            event_batch = [data[t].astype(self.output_type) for t in self.data_types]
+        return event_batch
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self.sample_mode == 'random':
+            self.inc_sample_count()
+            ret_dict = self._random_sample()
+        else:
+            if self.use_up:
+                raise StopIteration
+            else:
+                self.inc_sample_count()
+                ret_dict = self._sequent_sample()
+        ret_dict = self.data_dict_to_tensor(data_dict=ret_dict,
+                                            data_types=self.data_types)
+        if self.preprocess:
+            ret_dict = self.preprocess_data_dict(data_dict=ret_dict,
+                                                 data_types=self.data_types,
+                                                 layout=self.layout,
+                                                 rescale=self.rescale_method)
+        if self.downsample_dict is not None:
+            ret_dict = self.downsample_data_dict(data_dict=ret_dict,
+                                                 data_types=self.data_types,
+                                                 factors_dict=self.downsample_dict,
+                                                 layout=self.layout)
+        return ret_dict
+
+    def __getitem__(self, index):
+        data_dict = self._idx_sample(index=index)
+        return data_dict
+
+    @staticmethod
+    def preprocess_data_dict(data_dict, data_types=None, layout='NHWT', rescale='01'):
+        """
+        Parameters
+        ----------
+        data_dict:  Dict[str, Union[np.ndarray, torch.Tensor]]
+        data_types: Sequence[str]
+            The data types that we want to rescale. This mainly excludes "mask" from preprocessing.
+        layout: str
+            consists of batch_size 'N', seq_len 'T', channel 'C', height 'H', width 'W'
+        rescale:    str
+            'sevir': use the offsets and scale factors in original implementation.
+            '01': scale all values to range 0 to 1, currently only supports 'vil'
+        Returns
+        -------
+        data_dict:  Dict[str, Union[np.ndarray, torch.Tensor]]
+            preprocessed data
+        """
+        if rescale == 'sevir':
+            scale_dict = PREPROCESS_SCALE_SEVIR
+            offset_dict = PREPROCESS_OFFSET_SEVIR
+        elif rescale == '01':
+            scale_dict = PREPROCESS_SCALE_01
+            offset_dict = PREPROCESS_OFFSET_01
+        else:
+            raise ValueError(f'Invalid rescale option: {rescale}.')
+        if data_types is None:
+            data_types = data_dict.keys()
+        for key, data in data_dict.items():
+            if key in data_types:
+                if isinstance(data, np.ndarray):
+                    data = scale_dict[key] * (
+                            data.astype(np.float32) +
+                            offset_dict[key])
+                    data = change_layout_np(data=data,
+                                            in_layout='NHWT',
+                                            out_layout=layout)
+                elif isinstance(data, torch.Tensor):
+                    data = scale_dict[key] * (
+                            data.float() +
+                            offset_dict[key])
+                    data = change_layout_torch(data=data,
+                                               in_layout='NHWT',
+                                               out_layout=layout)
+                data_dict[key] = data
+        return data_dict
+
+    @staticmethod
+    def process_data_dict_back(data_dict, data_types=None, rescale='01'):
+        """
+        Parameters
+        ----------
+        data_dict
+            each data_dict[key] is a torch.Tensor.
+        rescale
+            str:
+                'sevir': data are scaled using the offsets and scale factors in original implementation.
+                '01': data are all scaled to range 0 to 1, currently only supports 'vil'
+        Returns
+        -------
+        data_dict
+            each data_dict[key] is the data processed back in torch.Tensor.
+        """
+        if rescale == 'sevir':
+            scale_dict = PREPROCESS_SCALE_SEVIR
+            offset_dict = PREPROCESS_OFFSET_SEVIR
+        elif rescale == '01':
+            scale_dict = PREPROCESS_SCALE_01
+            offset_dict = PREPROCESS_OFFSET_01
+        else:
+            raise ValueError(f'Invalid rescale option: {rescale}.')
+        if data_types is None:
+            data_types = data_dict.keys()
+        for key in data_types:
+            data = data_dict[key]
+            data = data.float() / scale_dict[key] - offset_dict[key]
+            data_dict[key] = data
+        return data_dict
+
+    @staticmethod
+    def data_dict_to_tensor(data_dict, data_types=None):
+        """
+        Convert each element in data_dict to torch.Tensor (copy without grad).
+        """
+        ret_dict = {}
+        if data_types is None:
+            data_types = data_dict.keys()
+        for key, data in data_dict.items():
+            if key in data_types:
+                if isinstance(data, torch.Tensor):
+                    ret_dict[key] = data.detach().clone()
+                elif isinstance(data, np.ndarray):
+                    ret_dict[key] = torch.from_numpy(data)
+                else:
+                    raise ValueError(f"Invalid data type: {type(data)}. Should be torch.Tensor or np.ndarray")
+            else:   # key == "mask"
+                ret_dict[key] = data
+        return ret_dict
+
+    @staticmethod
+    def downsample_data_dict(data_dict, data_types=None, factors_dict=None, layout='NHWT'):
+        """
+        Parameters
+        ----------
+        data_dict:  Dict[str, Union[np.array, torch.Tensor]]
+        factors_dict:   Optional[Dict[str, Sequence[int]]]
+            each element `factors` is a Sequence of int, representing (t_factor, h_factor, w_factor)
+
+        Returns
+        -------
+        downsampled_data_dict:  Dict[str, torch.Tensor]
+            Modify on a deep copy of data_dict instead of directly modifying the original data_dict
+        """
+        if factors_dict is None:
+            factors_dict = {}
+        if data_types is None:
+            data_types = data_dict.keys()
+        downsampled_data_dict = SEVIRDataLoader.data_dict_to_tensor(
+            data_dict=data_dict,
+            data_types=data_types)    # make a copy
+        for key, data in data_dict.items():
+            factors = factors_dict.get(key, None)
+            if factors is not None:
+                downsampled_data_dict[key] = change_layout_torch(
+                    data=downsampled_data_dict[key],
+                    in_layout=layout,
+                    out_layout='NTHW')
+                # downsample t dimension
+                t_slice = [slice(None, None), ] * 4
+                t_slice[1] = slice(None, None, factors[0])
+                downsampled_data_dict[key] = downsampled_data_dict[key][tuple(t_slice)]
+                # downsample spatial dimensions
+                downsampled_data_dict[key] = avg_pool2d(
+                    input=downsampled_data_dict[key],
+                    kernel_size=(factors[1], factors[2]))
+
+                downsampled_data_dict[key] = change_layout_torch(
+                    data=downsampled_data_dict[key],
+                    in_layout='NTHW',
+                    out_layout=layout)
+
+        return downsampled_data_dict
+
+    def _random_sample(self):
+        """
+        Returns
+        -------
+        ret_dict
+            dict. ret_dict.keys() == self.data_types.
+            If self.preprocess == False:
+                ret_dict[imgt].shape == (batch_size, height, width, seq_len)
+        """
+        num_sampled = 0
+        event_idx_list = nprand.randint(low=self.start_event_idx,
+                                        high=self.end_event_idx,
+                                        size=self.batch_size)
+        seq_idx_list = nprand.randint(low=0,
+                                      high=self.num_seq_per_event,
+                                      size=self.batch_size)
+        seq_slice_list = [slice(seq_idx * self.stride,
+                                seq_idx * self.stride + self.seq_len)
+                          for seq_idx in seq_idx_list]
+        ret_dict = {}
+        while num_sampled < self.batch_size:
+            event = self._load_event_batch(event_idx=event_idx_list[num_sampled],
+                                           event_batch_size=1)
+            for imgt_idx, imgt in enumerate(self.data_types):
+                sampled_seq = event[imgt_idx][[0, ], :, :, seq_slice_list[num_sampled]]  # keep the dim of batch_size for concatenation
+                if imgt in ret_dict:
+                    ret_dict[imgt] = np.concatenate((ret_dict[imgt], sampled_seq),
+                                                    axis=0)
+                else:
+                    ret_dict.update({imgt: sampled_seq})
+        return ret_dict
+
+    def _sequent_sample(self):
+        """
+        Returns
+        -------
+        ret_dict:   Dict
+            `ret_dict.keys()` contains `self.data_types`.
+            `ret_dict["mask"]` is a list of bool, indicating if the data entry is real or padded.
+            If self.preprocess == False:
+                ret_dict[imgt].shape == (batch_size, height, width, seq_len)
+        """
+        assert not self.use_up, 'Data loader used up! Reset it to reuse.'
+        event_idx = self.curr_event_idx
+        seq_idx = self.curr_seq_idx
+        num_sampled = 0
+        sampled_idx_list = []   # list of (event_idx, seq_idx) records
+        while num_sampled < self.batch_size:
+            sampled_idx_list.append({'event_idx': event_idx,
+                                     'seq_idx': seq_idx})
+            seq_idx += 1
+            if seq_idx >= self.num_seq_per_event:
+                event_idx += 1
+                seq_idx = 0
+            num_sampled += 1
+
+        start_event_idx = sampled_idx_list[0]['event_idx']
+        event_batch_size = sampled_idx_list[-1]['event_idx'] - start_event_idx + 1
+
+        event_batch = self._load_event_batch(event_idx=start_event_idx,
+                                             event_batch_size=event_batch_size)
+        ret_dict = {"mask": []}
+        all_no_pad_flag = True
+        for sampled_idx in sampled_idx_list:
+            batch_slice = [sampled_idx['event_idx'] - start_event_idx, ]  # use [] to keepdim
+            seq_slice = slice(sampled_idx['seq_idx'] * self.stride,
+                              sampled_idx['seq_idx'] * self.stride + self.seq_len)
+            for imgt_idx, imgt in enumerate(self.data_types):
+                sampled_seq = event_batch[imgt_idx][batch_slice, :, :, seq_slice]
+                if imgt in ret_dict:
+                    ret_dict[imgt] = np.concatenate((ret_dict[imgt], sampled_seq),
+                                                    axis=0)
+                else:
+                    ret_dict.update({imgt: sampled_seq})
+            # add mask
+            no_pad_flag = sampled_idx['event_idx'] < self.end_event_idx
+            if not no_pad_flag:
+                all_no_pad_flag = False
+            ret_dict["mask"].append(no_pad_flag)
+        if all_no_pad_flag:
+            # if there is no padded data items at all, set `ret_dict["mask"] = None` for convenience.
+            ret_dict["mask"] = None
+        # update current idx
+        self.set_curr_event_idx(event_idx)
+        self.set_curr_seq_idx(seq_idx)
+        return ret_dict
+
+    def _idx_sample(self, index):
+        """
+        Parameters
+        ----------
+        index
+            The index of the batch to sample.
+        Returns
+        -------
+        ret_dict
+            dict. ret_dict.keys() == self.data_types.
+            If self.preprocess == False:
+                ret_dict[imgt].shape == (batch_size, height, width, seq_len)
+        """
+        event_idx = (index * self.batch_size) // self.num_seq_per_event
+        seq_idx = (index * self.batch_size) % self.num_seq_per_event
+        num_sampled = 0
+        sampled_idx_list = []  # list of (event_idx, seq_idx) records
+        while num_sampled < self.batch_size:
+            sampled_idx_list.append({'event_idx': event_idx,
+                                     'seq_idx': seq_idx})
+            seq_idx += 1
+            if seq_idx >= self.num_seq_per_event:
+                event_idx += 1
+                seq_idx = 0
+            num_sampled += 1
+
+        start_event_idx = sampled_idx_list[0]['event_idx']
+        event_batch_size = sampled_idx_list[-1]['event_idx'] - start_event_idx + 1
+
+        event_batch = self._load_event_batch(event_idx=start_event_idx,
+                                             event_batch_size=event_batch_size)
+        ret_dict = {}
+        for sampled_idx in sampled_idx_list:
+            batch_slice = [sampled_idx['event_idx'] - start_event_idx, ]  # use [] to keepdim
+            seq_slice = slice(sampled_idx['seq_idx'] * self.stride,
+                              sampled_idx['seq_idx'] * self.stride + self.seq_len)
+            for imgt_idx, imgt in enumerate(self.data_types):
+                sampled_seq = event_batch[imgt_idx][batch_slice, :, :, seq_slice]
+                if imgt in ret_dict:
+                    ret_dict[imgt] = np.concatenate((ret_dict[imgt], sampled_seq),
+                                                    axis=0)
+                else:
+                    ret_dict.update({imgt: sampled_seq})
+
+        ret_dict = self.data_dict_to_tensor(data_dict=ret_dict,
+                                            data_types=self.data_types)
+        if self.preprocess:
+            ret_dict = self.preprocess_data_dict(data_dict=ret_dict,
+                                                 data_types=self.data_types,
+                                                 layout=self.layout,
+                                                 rescale=self.rescale_method)
+
+        if self.downsample_dict is not None:
+            ret_dict = self.downsample_data_dict(data_dict=ret_dict,
+                                                 data_types=self.data_types,
+                                                 factors_dict=self.downsample_dict,
+                                                 layout=self.layout)
+        return ret_dict
+
+
+class SEVIRDataIterator():
+    '''
+    A wrapper s.t. it implements the function sample().
+    Every arguments in this class will be redirected to the inner SEVIRDataLoader object.
+    If you expect a pythonic iterator, use SEVIRDataLoader instead.
+    '''
+    def __init__(self, **kwargs):
+        self.loader = SEVIRDataLoader(**kwargs)
+        self.sample_mode = kwargs['sample_mode'] if 'sample_mode' in kwargs else 'random'
+    
+    def reset(self):
+        self.loader.reset()
+    
+    def sample(self, batch_size=None):
+        '''
+        The input param batch_size here is not used
+        '''
+        out = next(self.loader, None)
+        if out is None and self.sample_mode == 'random':
+            self.loader.reset()
+            out = next(self.loader, None)
+        return out
+    
+    def __len__(self):
+        """
+        Used only when self.sample_mode == 'sequent'
+        """
+        return len(self.loader)
+
+# =====================================================================================
+# MeteoNet data
+# Reshape it to 256x256, with in_len=4, out_len=10
+# https://meteofrance.github.io/meteonet/
+# dwonload from https://meteonet.umr-cnrm.fr/dataset/data/NW/radar/reflectivity_old_product/
+# =====================================================================================
+
+class Meteo(Dataset):
+    def __init__(self, data_path, img_size, type='train', trans=None, in_len=-1):
+        super().__init__()
+        
+        self.pixel_scale = 70.0
+        
+        self.data_path = data_path
+        self.img_size = img_size
+        self.in_len = in_len
+
+        assert type in ['train', 'test', 'val']
+        self.type = type if type!='val' else 'test'
+        with h5py.File(data_path,'r') as f:
+            self.all_len = int(f[f'{self.type}_len'][()]) #  10000-3000 for train, 2000 for test, 1000 for val
+        if trans is not None:
+            self.transform = trans
+        else:
+            self.transform = T.Compose([
+                        T.Resize((img_size, img_size)),
+                        # transforms.ToTensor(),
+                        # trans.Lambda(lambda x: x/255.0),
+                        # transforms.Normalize(mean=[0.5], std=[0.5]),
+                        # trans.RandomCrop(data_config["img_size"]),
+
+                    ])
+                    
+    def __len__(self):
+        return self.all_len
+
+    def sample(self):
+        index = np.random.randint(0, self.all_len)
+        return self.__getitem__(index)
+    
+    
+    def __getitem__(self, index):
+
+        with h5py.File(self.data_path,'r') as f:
+            imgs = f[self.type][str(index)][()]   # numpy array: (25, 565, 784), dtype=uint8, range(0,70)
+
+            frames = torch.from_numpy(imgs).float().squeeze() 
+            frames = frames / self.pixel_scale
+            frames = self.transform(frames).unsqueeze(1)
+            
+        # return frames.unsqueeze(1) # (25,1,128,128
+        return frames[:self.in_len], frames[self.in_len:]
+ 
+
+def load_meteonet(batch_size, val_batch_size, in_len, train=False, num_workers=0, img_size=128):
+    meteo_filepath = os.path.join(METEO_FILE_DIR, "meteo.h5")
+    if train:
+        train_set = Meteo(meteo_filepath, img_size, 'train', in_len=in_len)
+        valid_set = Meteo(meteo_filepath, img_size, 'val', in_len=in_len)
+        dataloader_train = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=num_workers)
+        dataloader_valid = torch.utils.data.DataLoader(valid_set, batch_size=val_batch_size, shuffle=False, drop_last=True, num_workers=num_workers)
+        return dataloader_train, dataloader_valid
+    else:
+        test_set = Meteo(meteo_filepath, img_size, 'test', in_len=in_len)
+        dataloader_test = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+        return None, dataloader_test

data/loader.py

@@ -0,0 +1,49 @@
+import torch.nn.functional as F
+
+from data import dutils
+from nowcasting.hko_iterator import HKOIterator
+
+def GET_TrainLoader(meta, param, batch_size, in_len, out_len):
+    if meta['dataset'] == 'SEVIR':
+        total_seq_len = in_len + out_len
+        train_config = {
+            'data_types': ['vil'],
+            'layout': 'NTCHW',
+            'seq_len': total_seq_len,
+            'raw_seq_len': total_seq_len,
+            'end_date': dutils.SEVIR_TRAIN_TEST_SPLIT_DATE,
+            'start_date': None
+        }
+        test_config = {
+            'data_types': ['vil'],
+            'layout': 'NTCHW',
+            'seq_len': total_seq_len,
+            'raw_seq_len': total_seq_len,
+            'end_date': None,
+            'start_date': dutils.SEVIR_TRAIN_TEST_SPLIT_DATE
+        }
+        train_loader = dutils.SEVIRDataIterator(**train_config, batch_size=batch_size)
+        test_loader = dutils.SEVIRDataIterator(**test_config, batch_size=8 if batch_size > 8 else batch_size)
+        return train_loader, test_loader
+    elif meta['dataset'].startswith('HKO'):
+        total_seq_len = in_len + out_len
+        pkl_path = param['pd_path']
+        train_loader = HKOIterator(pd_path=pkl_path.replace('test', 'train'), sample_mode="random", seq_len=total_seq_len, stride=1)
+        test_loader = HKOIterator(pd_path=pkl_path, sample_mode="sequent", seq_len=total_seq_len, stride=in_len)
+        return train_loader, test_loader
+    elif meta['dataset'] == 'meteonet':
+        train_loader, test_loader = dutils.load_meteonet(batch_size=batch_size, val_batch_size=8 if batch_size > 8 else batch_size, train=True, **param)
+        return train_loader, test_loader
+    else:
+        raise Exception(f'Undefined dataset config name: {dataset_config["dataset"]}')
+
+def GET_TestLoader(meta, param, batch_size):
+    if meta['dataset'] == 'SEVIR':
+        return dutils.SEVIRDataIterator(**param, batch_size=batch_size)
+    elif meta['dataset'].startswith('HKO'):
+        return HKOIterator(**param)
+    elif meta['dataset'] == 'meteonet':
+        _, test_iter = dutils.load_meteonet(batch_size=batch_size, val_batch_size=8, train=False, **param)
+        return iter(test_iter)
+    else:
+        raise Exception(f'Undefined dataset config name: {dataset_config["dataset"]}')

data/sample_data.npy

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

ens_eval.py

@@ -0,0 +1,164 @@
+import os, sys, logging, argparse
+import torch
+from torch import nn
+import torch.nn.functional as F
+import numpy as np
+
+import utilspp as utpp
+from utilspp import mae, mse, ssim, psnr, lpips64, csi, hss
+from data.config import SEVIR_13_12, HKO7_5_20, METEONET_5_20
+from data.loader import GET_TestLoader
+from data.dutils import resize
+
+class MetricListEvaluator():    
+    '''
+    To evaluate a list of metrics. Supported metrics:
+    - CSI, HSS (Eg. `csi-84, hss-84`)
+    - CSI-pooled (Eg. `csi_4-84`)
+    - MAE
+    - MSE
+    - SSIM
+    - PSNR  
+    '''
+    def __init__(self, metric_list):        
+        self.metric_holder = {}
+        self.batch_count = 0
+        for metric_name in metric_list:
+            threshold = ''
+            radius = ''
+            if '-' in metric_name:
+                metric, threshold = metric_name.split('-')
+                if '_' in metric:
+                    metric, radius = metric.split('_')
+                    radius = int(radius)
+            # initialize metrics
+            threshold = float(threshold) / 255 if threshold.isdigit() else threshold
+            self.metric_holder[metric_name] = self.init_metric(metric_name, threshold=threshold, radius=radius)
+
+    def init_metric(self, metric_name, **kwarg):
+        '''
+        return a tuple of three items in order:
+        - the function to call during eval
+        - the value(s) to keep track of
+        - a dict of any additional item to pass into the function
+        '''
+        if metric_name.split('-')[0] in ['csi', 'hss']:
+            # use tfpn instead
+            return [utpp.tfpn, np.array([0, 0, 0, 0], dtype=np.float32), {'threshold': kwarg['threshold']}] # tp, 
+        elif '_' in metric_name.split('-')[0]: # Indicate Pooling
+            return [utpp.tfpn_pool, np.array([0, 0, 0, 0], dtype=np.float32), {'threshold': kwarg['threshold'], 'radius': kwarg['radius']}]
+        else:
+            # directly convert the string name into function call
+            return [eval(metric_name), 0, {}]
+
+    def eval(self, y_pred, y):
+        self.batch_count += 1
+        for _, metric in self.metric_holder.items():
+            temp = metric[0](y_pred, y, **metric[-1])      
+            if temp is list:
+                temp = np.array(temp)
+            elif type(temp) == torch.Tensor:
+                temp = temp.detach().cpu().numpy()
+            metric[1] += temp
+            
+    def get_results(self):
+        output_holder = {}
+        for key, metric in self.metric_holder.items():
+            val = metric[1]
+            # special handle of tfpn => compute the final score now
+            if metric[0] is utpp.tfpn:
+                metric_name, threshold = key.split('-')
+                val = eval(metric_name)(*list(metric[1]))
+            elif metric[0] is utpp.tfpn_pool:
+                metric_name, info = key.split('_')
+                val = eval(metric_name)(*list(metric[1]))
+            else:
+                val /= self.batch_count
+            output_holder[key] = val
+        return output_holder
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    # Dataset related
+    parser.add_argument('-d', '--dataset', type=str, default='', help='the dataset definition to be set')
+    parser.add_argument('--out_len',type=int, required=True, help='The actual prediction length')
+    # ensemble npy filename with {}
+    parser.add_argument('--e_file', default='', type=str, help='Ensemble npy filename with included \{ \}')
+    parser.add_argument('--ens_no', default=1, type=int, help='Total ensemble number')
+    # hyperparams
+    parser.add_argument('-s', '--step', type=int, default=-1, help='The number of steps to run. -1: the entire dataloader')
+    parser.add_argument('-b', '--batch_size', type=int, default=16, help='The batch size')
+    # config override
+    parser.add_argument('--metrics', type=str, default=None, help='A list of metrics to be evaluated, separated by character /')    
+    # logging related
+    parser.add_argument('--print_every', type=int, default=100, help='The number of steps to log the training loss')
+    args = parser.parse_args()
+
+    # Prepare logger
+    path_list = args.e_file.split("/")
+    logfile_name = os.path.join(*path_list[:-1], 'ensemble_eval.log')
+    logging.basicConfig(level=logging.NOTSET, handlers=[logging.FileHandler(logfile_name), logging.StreamHandler()], format='%(message)s')
+    logging.info(f'Steps: {args.step}')
+
+    dataset_config = globals()[args.dataset]
+    dataset_param, dataset_meta = dataset_config['param'], dataset_config['meta']
+    loader = GET_TestLoader(dataset_meta, dataset_param, args.batch_size)
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    # prepare metrics
+    metric_list = dataset_meta['metrics']
+    if args.metrics is not None: 
+        metric_list = args.metrics.lower().split('/')
+        logging.info(f'Overwriting metrics list with: {metric_list}')
+    evaluator = MetricListEvaluator(metric_list)
+
+    for e in range(args.ens_no):
+        prediction = np.load(args.e_file.format(str(e)))
+        prediction = torch.tensor(prediction, device=device)
+        
+        step = 1
+        if dataset_meta['dataset'] in ['SEVIR', 'HKO-7']:
+            loader.reset() # Reset it, otherwise alignment error
+        else:
+            pass
+            
+        while args.step < 0 or step <= args.step:
+            if dataset_meta['dataset'] == 'SEVIR':
+                data = loader.sample(batch_size=args.batch_size)
+                if data is None:
+                    break
+                y = data['vil'][:, -args.out_len:] # Expected to be same as prediction
+            elif dataset_meta['dataset'] == 'HKO-7':
+                setattr(args, 'seq_len', dataset_meta['seq_len'])
+                try:
+                    data = loader.sample(batch_size=args.batch_size)
+                except Exception as e:
+                    logging.error(e)
+                    break
+                x_seq, x_mask, dt_clip, _ = data
+                x, y = utpp.hko7_preprocess(x_seq, x_mask, dt_clip, args)
+            elif dataset_meta['dataset'].startswith('meteo'):
+                try:
+                    x, y = next(loader)
+                except Exception as e:
+                    logging.error(e)
+                    break
+
+            with torch.no_grad():
+                y = y.to(device)
+                y_pred = prediction[(step-1)*args.batch_size:step*args.batch_size]
+
+                if y.shape[-1] != y_pred.shape[-1]:
+                    y = resize(y, y_pred.shape[-1])
+
+                y, y_pred = y.clamp(0,1), y_pred.clamp(0,1) # B T C H W
+                evaluator.eval(y_pred, y)
+            # log/print every
+            if step == 1 or step % args.print_every == 0:
+                logging.info(f'E_ID:{e} -> {step} Steps evaluated')
+            step += 1
+    # log the final scores
+    final_results = evaluator.get_results()
+    for k, v in final_results.items():
+        logging.info(f'{k}: {v}')

ens_gen.py

@@ -0,0 +1,124 @@
+"""
+Sample Command
+"""
+import os, sys, logging, argparse
+import torch
+from torch import nn
+import numpy as np
+import torch.nn.functional as F
+
+from stldm import *
+import utilspp as utpp
+from data.config import SEVIR_13_12, HKO7_5_20, METEONET_5_20
+from data.loader import GET_TestLoader
+from data.dutils import resize
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    # Dataset related
+    parser.add_argument('-d', '--dataset', type=str, default='', help='the dataset definition to be set')
+    # model related
+    parser.add_argument('-f', dest='checkpt', type=str, default='', help='model checkpoint to be loaded from (Empty = not loading)')
+    parser.add_argument('-m', '--model', type=str, default='', help='the model definition to be created')
+    parser.add_argument('--type', type=str, default='3D', help='Determine which kind of model to use, 2D or 3D')
+    parser.add_argument('--c_str', type=float, default=0.0, help='CFG strength')
+    parser.add_argument('--e_id', type=int, default=0, help='Ensemble ID')
+    # hyperparams
+    parser.add_argument('-s', '--step', type=int, default=-1, help='The number of steps to run. -1: the entire dataloader')
+    parser.add_argument('-b', '--batch_size', type=int, default=16, help='The batch size')
+    # logging related
+    parser.add_argument('--print_every', type=int, default=100, help='The number of steps to log the training loss')
+    parser.add_argument('-o', '--output', default=None, help='The path to save the log files')
+    args = parser.parse_args()
+
+    # prepare logger
+    if args.output is None:
+        path_list = args.checkpt.split("/")
+        logfile_name = os.path.join(*path_list[:-1], 'logs', f'{path_list[-1][:-3]}.log')
+    else:
+        logfile_name = f'{args.output}.log'
+    logging.basicConfig(level=logging.NOTSET, handlers=[logging.FileHandler(logfile_name), logging.StreamHandler()], format='%(message)s')
+    logging.info(f'Model checkpoint: {args.checkpt}')
+    logging.info(f'Steps: {args.step}')
+
+    sampler_dir = os.path.join(*logfile_name.split("/")[:-2], f'CFG={args.c_str}_samples')
+    os.makedirs(sampler_dir, exist_ok=True)
+
+    # Prepare Dataloader
+    dataset_config = globals()[args.dataset]
+    dataset_param, dataset_meta = dataset_config['param'], dataset_config['meta']
+    loader = GET_TestLoader(dataset_meta, dataset_param, args.batch_size)
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    # Prepare Model
+    assert args.type in ['2D', '3D'], 'Please specify either 2D or 3D'
+    model_config = globals()[args.model]
+    model = n2n_setup[args.type](model_config, print_info=True, cfg_str=args.c_str if args.c_str != 0.0 else None).to(device)
+    logging.info(f'CFG Scheduler: Const-{args.c_str}')
+    
+    data = torch.load(args.checkpt, map_location=device)
+    if 'model' in data.keys():
+        model.load_state_dict(data['model'])
+    else:
+        model.load_state_dict(data)
+    
+    
+    in_len, out_len = model_config['vp_param']['shape_in'][0], model_config['vp_param']['shape_out'][0]
+    img_size = model_config['vp_param']['shape_in'][-1]
+    
+    step = 0
+    out = []
+    while args.step < 0 or step <=args.step:
+        model.eval()
+
+        if dataset_meta['dataset'] == 'HKO-7':  
+            setattr(args, 'seq_len', in_len)
+            try:
+                data = loader.sample(batch_size=args.batch_size)
+            except Exception as e:
+                logging.error(e)
+                break
+            x_seq, x_mask, dt_clip, _ = data                  
+            x, y = utpp.hko7_preprocess(x_seq, x_mask, dt_clip, args)
+        elif dataset_meta['dataset'] == 'SEVIR':
+            data = loader.sample(batch_size=args.batch_size)
+            if data is None:
+                break
+            x, y = data['vil'][:, :in_len], data['vil'][:, in_len:]
+        elif dataset_meta['dataset'].startswith('meteo'):
+            try:
+                x, y = next(loader)
+            except Exception as e:
+                logging.error(e)
+                break
+                
+        x, y = x.to(device), y.to(device)
+
+        with torch.no_grad():
+            if x.shape[-1] != img_size:
+                x = resize(x, img_size)
+                y = resize(y, img_size) # TO compare with DiffCast paper
+            if model_config['pre'] is not None:
+                x = model_config['pre'](x)
+
+            y_pred = model(x)
+                
+            if model_config['post'] is not None:
+                x = model_config['post'](x)
+                y_pred = model_config['post'](y_pred)
+            y_pred = y_pred.clamp(0,1)
+            
+            out.append(y_pred.detach().cpu())
+
+        step += 1
+        # log/print every
+        if step == 1 or step % args.print_every == 0:
+            logging.info(f'{step} Steps Generated, {len(out)} in out array')  
+        
+    logging.info(f'{step} Steps Generated, {len(out)} in out array')  
+    out = torch.cat(out, dim=0)
+    out = out.numpy()
+    save_path = os.path.join(sampler_dir, f'BTCHW_total-no:{len(out)}_e={args.e_id}.npy')
+    np.save(save_path, out)
+    print('Output saved in', save_path)

nowcasting/config.py

@@ -0,0 +1,302 @@
+import numpy as np
+import os
+import yaml
+import logging
+from collections import OrderedDict
+from .helpers.ordered_easydict import OrderedEasyDict as edict
+
+__C = edict()
+cfg = __C  # type: edict()
+
+# Random seed
+__C.SEED = None
+
+# Dataset name
+# Used by symbols factories who need to adjust for different
+# inputs based on dataset used. Should be set by the script.
+__C.DATASET = None
+
+# Project directory, since config.py is supposed to be in $ROOT_DIR/nowcasting
+#__C.ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+# ====== Changed the root dir to a fixed path =======
+__C.ROOT_DIR = 'data/HKO-7'
+
+__C.MNIST_PATH = os.path.join(__C.ROOT_DIR, 'mnist_data')
+if not os.path.exists(__C.MNIST_PATH):
+    os.makedirs(__C.MNIST_PATH)
+__C.HKO_DATA_BASE_PATH = os.path.join(__C.ROOT_DIR, 'hko_data')
+
+# # Append your path to the possible paths
+possible_hko_png_paths = [os.path.join('E:\\datasets\\HKO-data\\radarPNG\\radarPNG'),
+                          os.path.join(__C.HKO_DATA_BASE_PATH, 'radarPNG'),
+                          'data/HKO-7/radarPNG']
+possible_hko_mask_paths = [os.path.join('E:\\datasets\\HKO-data\\radarPNG\\radarPNG_mask'),
+                           os.path.join(__C.HKO_DATA_BASE_PATH, 'radarPNG_mask'),
+                           'data/HKO-7/radarPNG_mask']
+
+# Search for the radarPNG
+find_hko_png_path = False
+for ele in possible_hko_png_paths:
+    if os.path.exists(ele):
+        find_hko_png_path = True
+        __C.HKO_PNG_PATH = ele
+        break
+if not find_hko_png_path:
+    raise RuntimeError("radarPNG is not found! You can download the radarPNG using"
+                       " `bash download_radar_png.bash`")
+# Search for the radarPNG_mask
+find_hko_mask_path = False
+for ele in possible_hko_mask_paths:
+    if os.path.exists(ele):
+        find_hko_mask_path = True
+        __C.HKO_MASK_PATH = ele
+        break
+if not find_hko_mask_path:
+    raise RuntimeError("radarPNG_mask is not found! You can download the radarPNG_mask using"
+                       " `bash download_radar_png.bash`")
+if not os.path.exists(__C.HKO_DATA_BASE_PATH):
+    os.makedirs(__C.HKO_DATA_BASE_PATH)
+__C.HKO_PD_BASE_PATH = os.path.join(__C.HKO_DATA_BASE_PATH, 'pd')
+if not os.path.exists(__C.HKO_PD_BASE_PATH):
+    os.makedirs(__C.HKO_PD_BASE_PATH)
+__C.HKO_VALID_DATETIME_PATH = os.path.join(__C.HKO_DATA_BASE_PATH, 'valid_datetime.pkl')
+__C.HKO_SORTED_DAYS_PATH = os.path.join(__C.HKO_DATA_BASE_PATH, 'sorted_day.pkl')
+__C.HKO_RAINY_TRAIN_DAYS_PATH = os.path.join(__C.HKO_DATA_BASE_PATH, 'hko7_rainy_train_days.txt')
+__C.HKO_RAINY_VALID_DAYS_PATH = os.path.join(__C.HKO_DATA_BASE_PATH, 'hko7_rainy_valid_days.txt')
+__C.HKO_RAINY_TEST_DAYS_PATH = os.path.join(__C.HKO_DATA_BASE_PATH, 'hko7_rainy_test_days.txt')
+
+__C.HKO_PD = edict()
+__C.HKO_PD.ALL = os.path.join(__C.HKO_PD_BASE_PATH, 'hko7_all.pkl')
+__C.HKO_PD.ALL_09_14 = os.path.join(__C.HKO_PD_BASE_PATH, 'hko7_all_09_14.pkl')
+__C.HKO_PD.ALL_15 = os.path.join(__C.HKO_PD_BASE_PATH, 'hko7_all_15.pkl')
+__C.HKO_PD.RAINY_TRAIN = os.path.join(__C.HKO_PD_BASE_PATH, 'hko7_rainy_train.pkl')
+__C.HKO_PD.RAINY_VALID = os.path.join(__C.HKO_PD_BASE_PATH, 'hko7_rainy_valid.pkl')
+__C.HKO_PD.RAINY_TEST = os.path.join(__C.HKO_PD_BASE_PATH, 'hko7_rainy_test.pkl')
+
+__C.HKO = edict()
+__C.HKO.ITERATOR = edict()
+__C.HKO.ITERATOR.WIDTH = 480
+__C.HKO.ITERATOR.HEIGHT = 480
+__C.HKO.ITERATOR.FILTER_RAINFALL = True           # Whether to discard part of the rainfall, has a denoising effect
+__C.HKO.ITERATOR.FILTER_RAINFALL_THRESHOLD = 0.28 # All the pixel values that are smaller than round(threshold * 255) will be discarded
+
+
+# The Benchmark parameters
+__C.HKO.BENCHMARK = edict()
+__C.HKO.BENCHMARK.STAT_PATH = os.path.join(__C.HKO_DATA_BASE_PATH, 'benchmark_stat')
+if not os.path.exists(__C.HKO.BENCHMARK.STAT_PATH):
+    os.makedirs(__C.HKO.BENCHMARK.STAT_PATH)
+__C.HKO.BENCHMARK.VISUALIZE_SEQ_NUM = 10  # Number of sequences that will be plotted and saved to the benchmark directory
+__C.HKO.BENCHMARK.IN_LEN = 5   # The maximum input length to ensure that all models are tested on the same set of input data
+__C.HKO.BENCHMARK.OUT_LEN = 20  # The maximum output length to ensure that all models are tested on the same set of input data
+__C.HKO.BENCHMARK.STRIDE = 5   # The stride
+
+
+__C.HKO.EVALUATION = edict()
+__C.HKO.EVALUATION.ZR = edict()
+__C.HKO.EVALUATION.ZR.a = 58.53  # The a factor in the Z-R relationship
+__C.HKO.EVALUATION.ZR.b = 1.56  # The b factor in the Z-R relationship
+__C.HKO.EVALUATION.THRESHOLDS = (0.5, 2, 5, 10, 30)
+__C.HKO.EVALUATION.BALANCING_WEIGHTS = (1, 1, 2, 5, 10, 30)  # The corresponding balancing weights
+__C.HKO.EVALUATION.CENTRAL_REGION = (120, 120, 360, 360)
+
+__C.MOVINGMNIST = edict()
+__C.MOVINGMNIST.DISTRACTOR_NUM = 0
+__C.MOVINGMNIST.VELOCITY_LOWER = 0.0
+__C.MOVINGMNIST.VELOCITY_UPPER = 3.6
+__C.MOVINGMNIST.SCALE_VARIATION_LOWER = 1/1.1
+__C.MOVINGMNIST.SCALE_VARIATION_UPPER = 1.1
+__C.MOVINGMNIST.ROTATION_LOWER = -30
+__C.MOVINGMNIST.ROTATION_UPPER = 30
+__C.MOVINGMNIST.ILLUMINATION_LOWER = 0.6
+__C.MOVINGMNIST.ILLUMINATION_UPPER = 1.0
+__C.MOVINGMNIST.DIGIT_NUM = 3
+__C.MOVINGMNIST.IN_LEN = 10
+__C.MOVINGMNIST.OUT_LEN = 10
+__C.MOVINGMNIST.TESTING_LEN = 20
+__C.MOVINGMNIST.IMG_SIZE = 64
+__C.MOVINGMNIST.TEST_FILE = os.path.join(__C.MNIST_PATH, "movingmnist_10000_nodistr.npz")
+
+__C.MODEL = edict()
+__C.MODEL.RESUME = False  # If True, load LOAD_ITER parameters from LOAD_DIR
+__C.MODEL.TESTING = False # If True, run in Testing mode
+__C.MODEL.LOAD_DIR = "" # The directory to load the pre-trained parameters
+                        # Could be like `D:\\HKUST\\3-2\\NIPS2017\\hko_0502\\bal_loss_direct`
+__C.MODEL.LOAD_ITER = 79999           # Only applicable when LOAD_DIR is non-empty
+__C.MODEL.SAVE_DIR = ""
+__C.MODEL.CNN_ACT_TYPE = "leaky"
+__C.MODEL.RNN_ACT_TYPE = "leaky"
+__C.MODEL.FRAME_STACK = 1          # Stack multiple frames as the input
+__C.MODEL.FRAME_SKIP = 1           # The frame skip size
+__C.MODEL.IN_LEN = 5               # Size of the input
+__C.MODEL.OUT_LEN = 20             # Size of the output
+__C.MODEL.OUT_TYPE = "direct"      # Can be "direct", or "DFN"
+__C.MODEL.NORMAL_LOSS_GLOBAL_SCALE = 0.00005
+__C.MODEL.USE_BALANCED_LOSS = True
+__C.MODEL.TEMPORAL_WEIGHT_TYPE = "same"  # Can be "same", "linear" or "exponential"
+__C.MODEL.TEMPORAL_WEIGHT_UPPER = 5      # Only applicable when temporal_weights_type is "linear" or "exponential"
+                                         # If linear
+                                         #   the weights will be increased following (1 + i * (upper - 1) / (T - 1))
+                                         # If exponential
+                                         #   the weights will be increased following exp^{i * \ln(upper) / (T-1)}
+__C.MODEL.L1_LAMBDA = 1.0
+__C.MODEL.L2_LAMBDA = 1.0
+__C.MODEL.GDL_LAMBDA = 0.0
+__C.MODEL.USE_SEASONALITY = False          # Whether to use seasonality
+
+__C.MODEL.TRAJRNN = edict()
+__C.MODEL.TRAJRNN.INIT_GRID = True
+__C.MODEL.TRAJRNN.FLOW_LR_MULT = 1.0
+__C.MODEL.TRAJRNN.SAVE_MID_RESULTS = False
+
+__C.MODEL.ENCODER_FORECASTER = edict()
+__C.MODEL.ENCODER_FORECASTER.HAS_MASK = True
+__C.MODEL.ENCODER_FORECASTER.FEATMAP_SIZE = [96, 32, 16]
+__C.MODEL.ENCODER_FORECASTER.FIRST_CONV = (8, 7, 5, 1)  # Num filter, kernel, stride, pad
+__C.MODEL.ENCODER_FORECASTER.LAST_DECONV = (8, 7, 5, 1)  # Num filter, kernel, stride, pad
+__C.MODEL.ENCODER_FORECASTER.DOWNSAMPLE = [(5, 3, 1),
+                                           (3, 2, 1)]  # (kernel, stride, pad) for conv2d
+__C.MODEL.ENCODER_FORECASTER.UPSAMPLE = [(5, 3, 1),
+                                         (4, 2, 1)]  # (kernel, stride, pad) for deconv2d
+
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS = edict()    # Define the RNN blocks for the encoder RNN
+                                                     # In our network, the forecaster RNN will always have the reverse structure of encoder RNN
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.RES_CONNECTION = True
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.LAYER_TYPE = ["ConvGRU", "ConvGRU", "ConvGRU"]
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.STACK_NUM = [2, 3, 3]
+# These features are used for both ConvGRU
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.NUM_FILTER = [32, 64, 64]
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.H2H_KERNEL = [(5, 5), (5, 5), (3, 3)]
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.H2H_DILATE = [(1, 1), (1, 1), (1, 1)]
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.I2H_KERNEL = [(3, 3), (3, 3), (3, 3)]
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.I2H_PAD = [(1, 1), (1, 1), (1, 1)]
+# These features are only used in TrajGRU
+__C.MODEL.ENCODER_FORECASTER.RNN_BLOCKS.L = [5, 5, 5]
+
+__C.MODEL.DECONVBASELINE = edict()
+__C.MODEL.DECONVBASELINE.BASE_NUM_FILTER = 16
+__C.MODEL.DECONVBASELINE.USE_3D = True
+__C.MODEL.DECONVBASELINE.ENCODER = "separate"
+__C.MODEL.DECONVBASELINE.BN = True
+__C.MODEL.DECONVBASELINE.BN_GLOBAL_STATS = False
+__C.MODEL.DECONVBASELINE.COMPAT = edict() # Compatibility flags to recover behavior of previous versions
+__C.MODEL.DECONVBASELINE.COMPAT.CONV_INSTEADOF_FC_IN_ENCODER = False # Until 6th May 2017
+__C.MODEL.DECONVBASELINE.FC_BETWEEN_ENCDEC = 0
+
+__C.MODEL.TRAIN = edict()
+__C.MODEL.TRAIN.BATCH_SIZE = 3
+__C.MODEL.TRAIN.TBPTT = False
+__C.MODEL.TRAIN.OPTIMIZER = "adam"
+__C.MODEL.TRAIN.LR = 1E-4
+__C.MODEL.TRAIN.GAMMA1 = 0.9   # Used in RMSProp
+__C.MODEL.TRAIN.BETA1 = 0.5    # When using ADAM, momentum is called beta1
+__C.MODEL.TRAIN.EPS = 1E-8
+__C.MODEL.TRAIN.MIN_LR = 1E-6
+__C.MODEL.TRAIN.GRAD_CLIP = 50.0
+__C.MODEL.TRAIN.WD = 0
+__C.MODEL.TRAIN.MAX_ITER = 180000
+__C.MODEL.VALID_ITER = 5000
+__C.MODEL.SAVE_ITER = 15000
+__C.MODEL.TRAIN.LR_DECAY_ITER = 10000
+__C.MODEL.TRAIN.LR_DECAY_FACTOR = 0.7
+
+__C.MODEL.TEST = edict()
+__C.MODEL.TEST.FINETUNE = True
+__C.MODEL.TEST.MAX_ITER = 1      # Number of samples to generate in testing mode
+__C.MODEL.TEST.MODE = "online"    # Can be `online` or `fixed`
+__C.MODEL.TEST.DISABLE_TBPTT = True
+__C.MODEL.TEST.ONLINE = edict()
+__C.MODEL.TEST.ONLINE.OPTIMIZER = "adagrad"
+__C.MODEL.TEST.ONLINE.LR = 1E-4
+__C.MODEL.TEST.ONLINE.FINETUNE_MIN_MSE = 0.0
+__C.MODEL.TEST.ONLINE.GAMMA1 = 0.9    # Used in RMSProp
+__C.MODEL.TEST.ONLINE.BETA1 = 0.5     # Used in ADAM!
+__C.MODEL.TEST.ONLINE.EPS = 1E-6
+__C.MODEL.TEST.ONLINE.GRAD_CLIP = 50.0
+__C.MODEL.TEST.ONLINE.WD = 0
+
+
+def _merge_two_config(user_cfg, default_cfg):
+    """ Merge user's config into default config dictionary, clobbering the
+        options in b whenever they are also specified in a.
+        Need to ensure the type of two val under same key are the same
+        Do recursive merge when encounter hierarchical dictionary
+    """
+    if type(user_cfg) is not edict:
+        return
+    for key, val in user_cfg.items():
+        # Since user_cfg is a sub-file of default_cfg
+        if not key in default_cfg:
+            raise KeyError('{} is not a valid config key'.format(key))
+
+        if (type(default_cfg[key]) is not type(val) and
+                default_cfg[key] is not None):
+            if isinstance(default_cfg[key], np.ndarray):
+                val = np.array(val, dtype=default_cfg[key].dtype)
+            else:
+                raise ValueError(
+                     'Type mismatch ({} vs. {}) '
+                     'for config key: {}'.format(type(default_cfg[key]),
+                                                 type(val), key))
+        # Recursive merge config
+        if type(val) is edict:
+            try:
+                _merge_two_config(user_cfg[key], default_cfg[key])
+            except:
+                print('Error under config key: {}'.format(key))
+                raise
+        else:
+            default_cfg[key] = val
+
+
+def cfg_from_file(file_name, target=__C):
+    """ Load a config file and merge it into the default options.
+    """
+    import yaml
+    with open(file_name, 'r') as f:
+        print('Loading YAML config file from %s' %f)
+        yaml_cfg = edict(yaml.load(f))
+
+    _merge_two_config(yaml_cfg, target)
+
+
+def ordered_dump(data, stream=None, Dumper=yaml.SafeDumper, **kwds):
+    class OrderedDumper(Dumper):
+        pass
+
+    def _dict_representer(dumper, data):
+        return dumper.represent_mapping(
+            yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
+            data.items(), flow_style=False)
+
+    def _ndarray_representer(dumper, data):
+        return dumper.represent_list(data.tolist())
+
+    OrderedDumper.add_representer(OrderedDict, _dict_representer)
+    OrderedDumper.add_representer(edict, _dict_representer)
+    OrderedDumper.add_representer(np.ndarray, _ndarray_representer)
+    return yaml.dump(data, stream, OrderedDumper, **kwds)
+
+
+def save_cfg(dir_path, source=__C):
+    cfg_count = 0
+    file_path = os.path.join(dir_path, 'cfg%d.yml' %cfg_count)
+    while os.path.exists(file_path):
+        cfg_count += 1
+        file_path = os.path.join(dir_path, 'cfg%d.yml' % cfg_count)
+    with open(file_path, 'w') as f:
+        logging.info("Save YAML config file to %s" %file_path)
+        ordered_dump(source, f, yaml.SafeDumper, default_flow_style=None)
+
+
+def load_latest_cfg(dir_path, target=__C):
+    import re
+    cfg_count = None
+    source_cfg_path = None
+    for fname in os.listdir(dir_path):
+        ret = re.search('cfg(\d+)\.yml', fname)
+        if ret != None:
+            if cfg_count is None or (int(re.group(1)) > cfg_count):
+                cfg_count = int(re.group(1))
+                source_cfg_path = os.path.join(dir_path, ret.group(0))
+    cfg_from_file(file_name=source_cfg_path, target=target)

nowcasting/encoder_forecaster.py

@@ -0,0 +1,556 @@
+import mxnet as mx
+import mxnet.ndarray as nd
+import nowcasting.config as cfg
+from nowcasting.ops import reset_regs
+from nowcasting.operators.common import grid_generator
+from nowcasting.operators import *
+from nowcasting.ops import *
+from nowcasting.prediction_base_factory import PredictionBaseFactory
+from nowcasting.operators.transformations import DFN
+from nowcasting.my_module import MyModule
+
+
+def get_encoder_forecaster_rnn_blocks(batch_size):
+    encoder_rnn_blocks = []
+    forecaster_rnn_blocks = []
+    gan_rnn_blocks = []
+    CONFIG = cfg.MODEL.ENCODER_FORECASTER.RNN_BLOCKS
+    for vec, block_prefix in [(encoder_rnn_blocks, "ebrnn"),
+                              (forecaster_rnn_blocks, "fbrnn"),
+                              (gan_rnn_blocks, "dbrnn")]:
+        for i in range(len(CONFIG.NUM_FILTER)):
+            name = "%s%d" % (block_prefix, i + 1)
+            if CONFIG.LAYER_TYPE[i] == "ConvGRU":
+                rnn_block = BaseStackRNN(base_rnn_class=ConvGRU,
+                                         stack_num=CONFIG.STACK_NUM[i],
+                                         name=name,
+                                         residual_connection=CONFIG.RES_CONNECTION,
+                                         num_filter=CONFIG.NUM_FILTER[i],
+                                         b_h_w=(batch_size,
+                                                cfg.MODEL.ENCODER_FORECASTER.FEATMAP_SIZE[i],
+                                                cfg.MODEL.ENCODER_FORECASTER.FEATMAP_SIZE[i]),
+                                         h2h_kernel=CONFIG.H2H_KERNEL[i],
+                                         h2h_dilate=CONFIG.H2H_DILATE[i],
+                                         i2h_kernel=CONFIG.I2H_KERNEL[i],
+                                         i2h_pad=CONFIG.I2H_PAD[i],
+                                         act_type=cfg.MODEL.RNN_ACT_TYPE)
+            elif CONFIG.LAYER_TYPE[i] == "TrajGRU":
+                rnn_block = BaseStackRNN(base_rnn_class=TrajGRU,
+                                         stack_num=CONFIG.STACK_NUM[i],
+                                         name=name,
+                                         L=CONFIG.L[i],
+                                         residual_connection=CONFIG.RES_CONNECTION,
+                                         num_filter=CONFIG.NUM_FILTER[i],
+                                         b_h_w=(batch_size,
+                                                cfg.MODEL.ENCODER_FORECASTER.FEATMAP_SIZE[i],
+                                                cfg.MODEL.ENCODER_FORECASTER.FEATMAP_SIZE[i]),
+                                         h2h_kernel=CONFIG.H2H_KERNEL[i],
+                                         h2h_dilate=CONFIG.H2H_DILATE[i],
+                                         i2h_kernel=CONFIG.I2H_KERNEL[i],
+                                         i2h_pad=CONFIG.I2H_PAD[i],
+                                         act_type=cfg.MODEL.RNN_ACT_TYPE)
+            else:
+                raise NotImplementedError
+            vec.append(rnn_block)
+    return encoder_rnn_blocks, forecaster_rnn_blocks, gan_rnn_blocks
+
+class EncoderForecasterBaseFactory(PredictionBaseFactory):
+    def __init__(self,
+                 batch_size,
+                 in_seq_len,
+                 out_seq_len,
+                 height,
+                 width,
+                 ctx_num=1,
+                 name="encoder_forecaster"):
+        super(EncoderForecasterBaseFactory, self).__init__(batch_size=batch_size,
+                                                           in_seq_len=in_seq_len,
+                                                           out_seq_len=out_seq_len,
+                                                           height=height,
+                                                           width=width,
+                                                           name=name)
+        self._ctx_num = ctx_num
+
+    def _init_rnn(self):
+        self._encoder_rnn_blocks, self._forecaster_rnn_blocks, self._gan_rnn_blocks =\
+            get_encoder_forecaster_rnn_blocks(batch_size=self._batch_size)
+        return self._encoder_rnn_blocks + self._forecaster_rnn_blocks + self._gan_rnn_blocks
+
+    @property
+    def init_encoder_state_info(self):
+        init_state_info = []
+        for block in self._encoder_rnn_blocks:
+            for state in block.init_state_vars():
+                init_state_info.append({'name': state.name,
+                                        'shape': state.attr('__shape__'),
+                                        '__layout__': state.list_attr()['__layout__']})
+        return init_state_info
+
+    @property
+    def init_forecaster_state_info(self):
+        init_state_info = []
+        for block in self._forecaster_rnn_blocks:
+            for state in block.init_state_vars():
+                init_state_info.append({'name': state.name,
+                                        'shape': state.attr('__shape__'),
+                                        '__layout__': state.list_attr()['__layout__']})
+        return init_state_info
+
+    @property
+    def init_gan_state_info(self):
+        init_gan_state_info = []
+        for block in self._gan_rnn_blocks:
+            for state in block.init_state_vars():
+                init_gan_state_info.append({'name': state.name,
+                                            'shape': state.attr('__shape__'),
+                                            '__layout__': state.list_attr()['__layout__']})
+        return init_gan_state_info
+
+    def stack_rnn_encode(self, data):
+        CONFIG = cfg.MODEL.ENCODER_FORECASTER
+        pre_encoded_data = self._pre_encode_frame(frame_data=data, seqlen=self._in_seq_len)
+        reshape_data = mx.sym.Reshape(pre_encoded_data, shape=(-1, 0, 0, 0), reverse=True)
+
+        # Encoder Part
+        conv1 = conv2d_act(data=reshape_data,
+                           num_filter=CONFIG.FIRST_CONV[0],
+                           kernel=(CONFIG.FIRST_CONV[1], CONFIG.FIRST_CONV[1]),
+                           stride=(CONFIG.FIRST_CONV[2], CONFIG.FIRST_CONV[2]),
+                           pad=(CONFIG.FIRST_CONV[3], CONFIG.FIRST_CONV[3]),
+                           act_type=cfg.MODEL.CNN_ACT_TYPE,
+                           name="econv1")
+        rnn_block_num = len(CONFIG.RNN_BLOCKS.NUM_FILTER)
+        encoder_rnn_block_states = []
+        for i in range(rnn_block_num):
+            if i == 0:
+                inputs = conv1
+            else:
+                inputs = downsample
+            rnn_out, states = self._encoder_rnn_blocks[i].unroll(
+                length=self._in_seq_len,
+                inputs=inputs,
+                begin_states=None,
+                ret_mid=False)
+            encoder_rnn_block_states.append(states)
+            if i < rnn_block_num - 1:
+                downsample = downsample_module(data=rnn_out[-1],
+                                               num_filter=CONFIG.RNN_BLOCKS.NUM_FILTER[i + 1],
+                                               kernel=(CONFIG.DOWNSAMPLE[i][0],
+                                                       CONFIG.DOWNSAMPLE[i][0]),
+                                               stride=(CONFIG.DOWNSAMPLE[i][1],
+                                                       CONFIG.DOWNSAMPLE[i][1]),
+                                               pad=(CONFIG.DOWNSAMPLE[i][2],
+                                                    CONFIG.DOWNSAMPLE[i][2]),
+                                               b_h_w=(self._batch_size,
+                                                      CONFIG.FEATMAP_SIZE[i + 1],
+                                                      CONFIG.FEATMAP_SIZE[i + 1]),
+                                               name="edown%d" %(i + 1))
+        return encoder_rnn_block_states
+
+    def stack_rnn_forecast(self, block_state_list, last_frame):
+        CONFIG = cfg.MODEL.ENCODER_FORECASTER
+        block_state_list = [self._forecaster_rnn_blocks[i].to_split(block_state_list[i])
+                            for i in range(len(self._forecaster_rnn_blocks))]
+        rnn_block_num = len(CONFIG.RNN_BLOCKS.NUM_FILTER)
+        rnn_block_outputs = []
+        # RNN Forecaster Part
+        curr_inputs = None
+        for i in range(rnn_block_num - 1, -1, -1):
+            rnn_out, rnn_state = self._forecaster_rnn_blocks[i].unroll(
+                length=self._out_seq_len, inputs=curr_inputs,
+                begin_states=block_state_list[i][::-1],  # Reverse the order of states for the forecaster
+                ret_mid=False)
+            rnn_block_outputs.append(rnn_out)
+            if i > 0:
+                upsample = upsample_module(data=rnn_out[-1],
+                                           num_filter=CONFIG.RNN_BLOCKS.NUM_FILTER[i],
+                                           kernel=(CONFIG.UPSAMPLE[i - 1][0],
+                                                   CONFIG.UPSAMPLE[i - 1][0]),
+                                           stride=(CONFIG.UPSAMPLE[i - 1][1],
+                                                   CONFIG.UPSAMPLE[i - 1][1]),
+                                           pad=(CONFIG.UPSAMPLE[i - 1][2],
+                                                CONFIG.UPSAMPLE[i - 1][2]),
+                                           b_h_w=(self._batch_size, CONFIG.FEATMAP_SIZE[i - 1]),
+                                           name="fup%d" %i)
+                curr_inputs = upsample
+        # Output
+        if cfg.MODEL.OUT_TYPE == "DFN":
+            concat_fbrnn1_out = mx.sym.concat(*rnn_out[-1], dim=0)
+            dynamic_filter = deconv2d(data=concat_fbrnn1_out,
+                                      num_filter=121,
+                                      kernel=(CONFIG.LAST_DECONV[1], CONFIG.LAST_DECONV[1]),
+                                      stride=(CONFIG.LAST_DECONV[2], CONFIG.LAST_DECONV[2]),
+                                      pad=(CONFIG.LAST_DECONV[3], CONFIG.LAST_DECONV[3]))
+            flow = dynamic_filter
+            dynamic_filter = mx.sym.SliceChannel(dynamic_filter, axis=0, num_outputs=self._out_seq_len)
+            prev_frame = last_frame
+            preds = []
+            for i in range(self._out_seq_len):
+                pred_ele = DFN(data=prev_frame, local_kernels=dynamic_filter[i], K=11, batch_size=self._batch_size)
+                preds.append(pred_ele)
+                prev_frame = pred_ele
+            pred = mx.sym.concat(*preds, dim=0)
+        elif cfg.MODEL.OUT_TYPE == "direct":
+            flow = None
+            deconv1 = deconv2d_act(data=mx.sym.concat(*rnn_out[-1], dim=0),
+                                   num_filter=CONFIG.LAST_DECONV[0],
+                                   kernel=(CONFIG.LAST_DECONV[1], CONFIG.LAST_DECONV[1]),
+                                   stride=(CONFIG.LAST_DECONV[2], CONFIG.LAST_DECONV[2]),
+                                   pad=(CONFIG.LAST_DECONV[3], CONFIG.LAST_DECONV[3]),
+                                   act_type=cfg.MODEL.CNN_ACT_TYPE,
+                                   name="fdeconv1")
+            conv_final = conv2d_act(data=deconv1,
+                                    num_filter=CONFIG.LAST_DECONV[0],
+                                    kernel=(3, 3), stride=(1, 1), pad=(1, 1),
+                                    act_type=cfg.MODEL.CNN_ACT_TYPE, name="conv_final")
+            pred = conv2d(data=conv_final,
+                          num_filter=1, kernel=(1, 1), name="out")
+        else:
+            raise NotImplementedError
+        pred = mx.sym.Reshape(pred,
+                              shape=(self._out_seq_len, self._batch_size,
+                                     1, self._height, self._width),
+                              __layout__="TNCHW")
+        return pred, flow
+
+    def encoder_sym(self):
+        self.reset_all()
+        data = mx.sym.Variable('data')  # Shape: (in_seq_len, batch_size, C, H, W)
+        block_state_list = self.stack_rnn_encode(data=data)
+        states = []
+        for i, rnn_block in enumerate(self._encoder_rnn_blocks):
+            states.extend(rnn_block.flatten_add_layout(block_state_list[i]))
+        return mx.sym.Group(states)
+
+    def encoder_data_desc(self):
+        ret = list()
+        ret.append(mx.io.DataDesc(name='data',
+                                  shape=(self._in_seq_len,
+                                         self._batch_size * self._ctx_num,
+                                         1,
+                                         self._height,
+                                         self._width),
+                                  layout="TNCHW"))
+        for info in self.init_encoder_state_info:
+            state_shape = safe_eval(info['shape'])
+            assert info['__layout__'].find('N') == 0,\
+                "Layout=%s is not supported!" %info["__layout__"]
+            state_shape = (state_shape[0] * self._ctx_num, ) + state_shape[1:]
+            ret.append(mx.io.DataDesc(name=info['name'],
+                                      shape=state_shape,
+                                      layout=info['__layout__']))
+        return ret
+
+    def forecaster_sym(self):
+        self.reset_all()
+        block_state_list = []
+        for block in self._forecaster_rnn_blocks:
+            block_state_list.append(block.init_state_vars())
+
+        if cfg.MODEL.OUT_TYPE == "direct":
+            pred, _ = self.stack_rnn_forecast(block_state_list=block_state_list,
+                                              last_frame=None)
+            return mx.sym.Group([pred])
+        else:
+            last_frame = mx.sym.Variable('last_frame')  # Shape: (batch_size, C, H, W)
+            pred, flow = self.stack_rnn_forecast(block_state_list=block_state_list,
+                                                 last_frame=last_frame)
+            return mx.sym.Group([pred, mx.sym.BlockGrad(flow)])
+
+    def forecaster_data_desc(self):
+        ret = list()
+        for info in self.init_forecaster_state_info:
+            state_shape = safe_eval(info['shape'])
+            assert info['__layout__'].find('N') == 0, \
+                "Layout=%s is not supported!" % info["__layout__"]
+            state_shape = (state_shape[0] * self._ctx_num,) + state_shape[1:]
+            ret.append(mx.io.DataDesc(name=info['name'],
+                                      shape=state_shape,
+                                      layout=info['__layout__']))
+        if cfg.MODEL.OUT_TYPE != "direct":
+            ret.append(mx.io.DataDesc(name="last_frame",
+                                      shape=(self._ctx_num * self._batch_size,
+                                             1, self._height, self._width),
+                                      layout="NCHW"))
+        return ret
+
+    def loss_sym(self):
+        raise NotImplementedError
+
+    def loss_data_desc(self):
+        ret = list()
+        ret.append(mx.io.DataDesc(name='pred',
+                                  shape=(self._out_seq_len,
+                                         self._ctx_num * self._batch_size,
+                                         1,
+                                         self._height,
+                                         self._width),
+                                  layout="TNCHW"))
+        return ret
+
+    def loss_label_desc(self):
+        ret = list()
+        ret.append(mx.io.DataDesc(name='target',
+                                  shape=(self._out_seq_len,
+                                         self._ctx_num * self._batch_size,
+                                         1,
+                                         self._height,
+                                         self._width),
+                                  layout="TNCHW"))
+        if cfg.MODEL.ENCODER_FORECASTER.HAS_MASK:
+            ret.append(mx.io.DataDesc(name='mask',
+                                      shape=(self._out_seq_len,
+                                             self._ctx_num * self._batch_size,
+                                             1,
+                                             self._height,
+                                             self._width),
+                                      layout="TNCHW"))
+        return ret
+
+
+
+def init_optimizer_using_cfg(net, for_finetune):
+    if not for_finetune:
+        lr_scheduler = mx.lr_scheduler.FactorScheduler(step=cfg.MODEL.TRAIN.LR_DECAY_ITER,
+                                                       factor=cfg.MODEL.TRAIN.LR_DECAY_FACTOR,
+                                                       stop_factor_lr=cfg.MODEL.TRAIN.MIN_LR)
+        if cfg.MODEL.TRAIN.OPTIMIZER.lower() == "adam":
+            net.init_optimizer(optimizer="adam",
+                               optimizer_params={'learning_rate': cfg.MODEL.TRAIN.LR,
+                                                 'beta1': cfg.MODEL.TRAIN.BETA1,
+                                                 'rescale_grad': 1.0,
+                                                 'epsilon': cfg.MODEL.TRAIN.EPS,
+                                                 'lr_scheduler': lr_scheduler,
+                                                 'wd': cfg.MODEL.TRAIN.WD})
+        elif cfg.MODEL.TRAIN.OPTIMIZER.lower() == "rmsprop":
+            net.init_optimizer(optimizer="rmsprop",
+                               optimizer_params={'learning_rate': cfg.MODEL.TRAIN.LR,
+                                                 'gamma1': cfg.MODEL.TRAIN.GAMMA1,
+                                                 'rescale_grad': 1.0,
+                                                 'epsilon': cfg.MODEL.TRAIN.EPS,
+                                                 'lr_scheduler': lr_scheduler,
+                                                 'wd': cfg.MODEL.TRAIN.WD})
+        elif cfg.MODEL.TRAIN.OPTIMIZER.lower() == "sgd":
+            net.init_optimizer(optimizer="sgd",
+                               optimizer_params={'learning_rate': cfg.MODEL.TRAIN.LR,
+                                                 'momentum': 0.0,
+                                                 'rescale_grad': 1.0,
+                                                 'lr_scheduler': lr_scheduler,
+                                                 'wd': cfg.MODEL.TRAIN.WD})
+        elif cfg.MODEL.TRAIN.OPTIMIZER.lower() == "adagrad":
+            net.init_optimizer(optimizer="adagrad",
+                               optimizer_params={'learning_rate': cfg.MODEL.TRAIN.LR,
+                                                 'eps': cfg.MODEL.TRAIN.EPS,
+                                                 'rescale_grad': 1.0,
+                                                 'wd': cfg.MODEL.TRAIN.WD})
+        else:
+            raise NotImplementedError
+    else:
+        if cfg.MODEL.TEST.ONLINE.OPTIMIZER.lower() == "adam":
+            net.init_optimizer(optimizer="adam",
+                               optimizer_params={'learning_rate': cfg.MODEL.TEST.ONLINE.LR,
+                                                 'beta1': cfg.MODEL.TEST.ONLINE.BETA1,
+                                                 'rescale_grad': 1.0,
+                                                 'epsilon': cfg.MODEL.TEST.ONLINE.EPS,
+                                                 'wd': cfg.MODEL.TEST.ONLINE.WD})
+        elif cfg.MODEL.TEST.ONLINE.OPTIMIZER.lower() == "rmsprop":
+            net.init_optimizer(optimizer="rmsprop",
+                               optimizer_params={'learning_rate': cfg.MODEL.TEST.ONLINE.LR,
+                                                 'gamma1': cfg.MODEL.TEST.ONLINE.GAMMA1,
+                                                 'rescale_grad': 1.0,
+                                                 'epsilon': cfg.MODEL.TEST.ONLINE.EPS,
+                                                 'wd': cfg.MODEL.TEST.ONLINE.WD})
+        elif cfg.MODEL.TEST.ONLINE.OPTIMIZER.lower() == "sgd":
+            net.init_optimizer(optimizer="sgd",
+                               optimizer_params={'learning_rate': cfg.MODEL.TEST.ONLINE.LR,
+                                                 'momentum': 0.0,
+                                                 'rescale_grad': 1.0,
+                                                 'wd': cfg.MODEL.TEST.ONLINE.WD})
+        elif cfg.MODEL.TEST.ONLINE.OPTIMIZER.lower() == "adagrad":
+            net.init_optimizer(optimizer="adagrad",
+                               optimizer_params={'learning_rate': cfg.MODEL.TEST.ONLINE.LR,
+                                                 'eps': cfg.MODEL.TRAIN.EPS,
+                                                 'rescale_grad': 1.0,
+                                                 'wd': cfg.MODEL.TEST.ONLINE.WD})
+    return net
+
+
+def encoder_forecaster_build_networks(factory, context,
+                                      shared_encoder_net=None,
+                                      shared_forecaster_net=None,
+                                      shared_loss_net=None,
+                                      for_finetune=False):
+    """
+    
+    Parameters
+    ----------
+    factory : EncoderForecasterBaseFactory
+    context : list
+    shared_encoder_net : MyModule or None
+    shared_forecaster_net : MyModule or None
+    shared_loss_net : MyModule or None
+    for_finetune : bool
+
+    Returns
+    -------
+
+    """
+    encoder_net = MyModule(factory.encoder_sym(),
+                           data_names=[ele.name for ele in factory.encoder_data_desc()],
+                           label_names=[],
+                           context=context,
+                           name="encoder_net")
+    encoder_net.bind(data_shapes=factory.encoder_data_desc(),
+                     label_shapes=None,
+                     inputs_need_grad=True,
+                     shared_module=shared_encoder_net)
+    if shared_encoder_net is None:
+        encoder_net.init_params(mx.init.MSRAPrelu(slope=0.2))
+        init_optimizer_using_cfg(encoder_net, for_finetune=for_finetune)
+    forecaster_net = MyModule(factory.forecaster_sym(),
+                                   data_names=[ele.name for ele in
+                                               factory.forecaster_data_desc()],
+                                   label_names=[],
+                                   context=context,
+                                   name="forecaster_net")
+    forecaster_net.bind(data_shapes=factory.forecaster_data_desc(),
+                        label_shapes=None,
+                        inputs_need_grad=True,
+                        shared_module=shared_forecaster_net)
+    if shared_forecaster_net is None:
+        forecaster_net.init_params(mx.init.MSRAPrelu(slope=0.2))
+        init_optimizer_using_cfg(forecaster_net, for_finetune=for_finetune)
+
+    loss_net = MyModule(factory.loss_sym(),
+                        data_names=[ele.name for ele in
+                                    factory.loss_data_desc()],
+                        label_names=[ele.name for ele in
+                                     factory.loss_label_desc()],
+                        context=context,
+                        name="loss_net")
+    loss_net.bind(data_shapes=factory.loss_data_desc(),
+                  label_shapes=factory.loss_label_desc(),
+                  inputs_need_grad=True,
+                  shared_module=shared_loss_net)
+    if shared_loss_net is None:
+        loss_net.init_params()
+    return encoder_net, forecaster_net, loss_net
+
+
+class EncoderForecasterStates(object):
+    def __init__(self, factory, ctx):
+        self._factory = factory
+        self._ctx = ctx
+        self._encoder_state_info = factory.init_encoder_state_info
+        self._forecaster_state_info = factory.init_forecaster_state_info
+        self._states_nd = []
+        for info in self._encoder_state_info:
+            state_shape = safe_eval(info['shape'])
+            state_shape = (state_shape[0] * factory._ctx_num, ) + state_shape[1:]
+            self._states_nd.append(mx.nd.zeros(shape=state_shape, ctx=ctx))
+
+    def reset_all(self):
+        for ele, info in zip(self._states_nd, self._encoder_state_info):
+            ele[:] = 0
+
+    def reset_batch(self, batch_id):
+        for ele, info in zip(self._states_nd, self._encoder_state_info):
+            ele[batch_id][:] = 0
+
+    def update(self, states_nd):
+        for target, src in zip(self._states_nd, states_nd):
+            target[:] = src
+
+    def get_encoder_states(self):
+        return self._states_nd
+
+    def get_forecaster_state(self):
+        return self._states_nd
+
+
+def train_step(batch_size, encoder_net, forecaster_net,
+               loss_net, init_states,
+               data_nd, gt_nd, mask_nd, iter_id=None):
+    """Finetune the encoder, forecaster and GAN for one step
+
+    Parameters
+    ----------
+    batch_size : int
+    encoder_net : MyModule
+    forecaster_net : MyModule
+    loss_net : MyModule
+    init_states : EncoderForecasterStates
+    data_nd : mx.nd.ndarray
+    gt_nd : mx.nd.ndarray
+    mask_nd : mx.nd.ndarray
+    iter_id : int
+
+    Returns
+    -------
+    init_states: EncoderForecasterStates
+    loss_dict: dict
+    """
+    # Forward Encoder
+    encoder_net.forward(is_train=True,
+                        data_batch=mx.io.DataBatch(data=[data_nd] + init_states.get_encoder_states()))
+    encoder_states_nd = encoder_net.get_outputs()
+    init_states.update(encoder_states_nd)
+    # Forward Forecaster
+    if cfg.MODEL.OUT_TYPE == "direct":
+        forecaster_net.forward(is_train=True,
+                               data_batch=mx.io.DataBatch(data=init_states.get_forecaster_state()))
+    else:
+        last_frame_nd = data_nd[data_nd.shape[0] - 1]
+        forecaster_net.forward(is_train=True,
+                               data_batch=mx.io.DataBatch(data=init_states.get_forecaster_state() +
+                                                                    [last_frame_nd]))
+    forecaster_outputs = forecaster_net.get_outputs()
+    pred_nd = forecaster_outputs[0]
+
+    # Calculate the gradient of the loss functions
+    if cfg.MODEL.ENCODER_FORECASTER.HAS_MASK:
+        loss_net.forward_backward(data_batch=mx.io.DataBatch(data=[pred_nd],
+                                                             label=[gt_nd, mask_nd]))
+    else:
+        loss_net.forward_backward(data_batch=mx.io.DataBatch(data=[pred_nd],
+                                                             label=[gt_nd]))
+    pred_grad = loss_net.get_input_grads()[0]
+    loss_dict = loss_net.get_output_dict()
+    for k in loss_dict:
+        loss_dict[k] = nd.mean(loss_dict[k]).asscalar()
+    # Backward Forecaster
+    forecaster_net.backward(out_grads=[pred_grad])
+    if cfg.MODEL.OUT_TYPE == "direct":
+        encoder_states_grad_nd = forecaster_net.get_input_grads()
+    else:
+        encoder_states_grad_nd = forecaster_net.get_input_grads()[:-1]
+    # Backward Encoder
+    encoder_net.backward(encoder_states_grad_nd)
+    # Update forecaster and encoder
+    forecaster_grad_norm = forecaster_net.clip_by_global_norm(max_norm=cfg.MODEL.TRAIN.GRAD_CLIP)
+    encoder_grad_norm = encoder_net.clip_by_global_norm(max_norm=cfg.MODEL.TRAIN.GRAD_CLIP)
+    forecaster_net.update()
+    encoder_net.update()
+    loss_str = ", ".join(["%s=%g" %(k, v) for k, v in loss_dict.items()])
+    if iter_id is not None:
+        logging.info("Iter:%d, %s, e_gnorm=%g, f_gnorm=%g"
+                     % (iter_id, loss_str, encoder_grad_norm, forecaster_grad_norm))
+    return init_states, loss_dict
+
+
+def load_encoder_forecaster_params(load_dir, load_iter, encoder_net, forecaster_net):
+    logging.info("Loading parameters from {}, Iter = {}"
+                 .format(os.path.realpath(load_dir), load_iter))
+    encoder_arg_params, encoder_aux_params = load_params(prefix=os.path.join(load_dir,
+                                                                             "encoder_net"),
+                                                         epoch=load_iter)
+    encoder_net.init_params(arg_params=encoder_arg_params, aux_params=encoder_aux_params,
+                            allow_missing=False, force_init=True)
+    forecaster_arg_params, forecaster_aux_params = load_params(prefix=os.path.join(load_dir,
+                                                                             "forecaster_net"),
+                                                               epoch=load_iter)
+    forecaster_net.init_params(arg_params=forecaster_arg_params,
+                               aux_params=forecaster_aux_params,
+                               allow_missing=False,
+                               force_init=True)
+    logging.info("Loading Complete!")