├── .gitignore ├── LICENSE ├── README.md ├── conda_env.yml ├── config ├── largescale │ ├── dcrnn_cer.yaml │ ├── dcrnn_pv.yaml │ ├── gatedgn_cer.yaml │ ├── gatedgn_pv.yaml │ ├── gwnet_cer.yaml │ ├── gwnet_pv.yaml │ ├── sgp_cer.yaml │ └── sgp_pv.yaml ├── largescale_100nn │ ├── dcrnn_cer.yaml │ ├── dcrnn_pv.yaml │ ├── gatedgn_cer.yaml │ ├── gatedgn_pv.yaml │ ├── gwnet_cer.yaml │ ├── gwnet_pv.yaml │ ├── sgp_cer.yaml │ └── sgp_pv.yaml └── traffic │ ├── dcrnn.yaml │ ├── gatedgn_bay.yaml │ ├── gatedgn_la.yaml │ ├── gesn_bay.yaml │ ├── gesn_la.yaml │ ├── gwnet.yaml │ ├── rnn.yaml │ ├── sgp_bay.yaml │ ├── sgp_bay_abl1.yaml │ ├── sgp_bay_abl2.yaml │ ├── sgp_bay_abl3.yaml │ ├── sgp_la.yaml │ ├── sgp_la_abl1.yaml │ ├── sgp_la_abl2.yaml │ └── sgp_la_abl3.yaml ├── experiments ├── __init__.py ├── run_closed_form.py ├── run_largescale_baselines.py ├── run_largescale_sgp.py ├── run_traffic_baselines.py └── run_traffic_sgp.py ├── lib ├── __init__.py ├── dataloader │ ├── __init__.py │ ├── iid_dataloader.py │ ├── sgp_dataloader.py │ └── subgraph_dataloader.py ├── datamodule │ ├── __init__.py │ ├── sgp_datamodule.py │ └── subgraph_datamodule.py ├── datasets │ ├── __init__.py │ ├── cer_en.py │ ├── iid_dataset.py │ └── pv.py ├── nn │ ├── __init__.py │ ├── encoders │ │ ├── __init__.py │ │ ├── dyn_gesn_encoder.py │ │ ├── sgp_encoder.py │ │ ├── sgp_spatial_encoder.py │ │ └── sgp_temporal_encoder.py │ ├── models │ │ ├── __init__.py │ │ ├── esn_model.py │ │ ├── gated_gn_model.py │ │ ├── gwnet_model.py │ │ ├── sgp_model.py │ │ └── sgp_online.py │ └── reservoir │ │ ├── __init__.py │ │ ├── graph_reservoir.py │ │ └── reservoir.py ├── predictors │ ├── __init__.py │ ├── profiling_predictor.py │ └── subgraph_predictor.py ├── sgp_preprocessing.py └── utils.py ├── scalable.png ├── sgp_paper.pdf ├── sgp_poster.pdf └── tsl ├── __init__.py ├── data ├── __init__.py ├── batch.py ├── batch_map.py ├── data.py ├── datamodule │ ├── __init__.py │ ├── spatiotemporal_datamodule.py │ └── splitters.py ├── imputation_stds.py ├── loader │ ├── __init__.py │ └── dataloader.py ├── mixin.py ├── preprocessing │ ├── __init__.py │ └── scalers.py ├── spatiotemporal_dataset.py └── utils.py ├── datasets ├── __init__.py ├── metr_la.py ├── mts_benchmarks.py ├── pems_bay.py └── prototypes │ ├── __init__.py │ ├── casting.py │ ├── dataset.py │ ├── mixin.py │ ├── pd_dataset.py │ └── tabular_dataset.py ├── global_scope ├── __init__.py ├── config.py ├── lazy_loader.py └── logger.py ├── imputers ├── __init__.py └── imputer.py ├── nn ├── __init__.py ├── base │ ├── __init__.py │ ├── attention │ │ ├── __init__.py │ │ ├── attention.py │ │ └── linear_attention.py │ ├── dense.py │ ├── embedding.py │ ├── graph_conv.py │ └── temporal_conv.py ├── blocks │ ├── __init__.py │ ├── decoders │ │ ├── __init__.py │ │ ├── att_pool.py │ │ ├── gcn_decoder.py │ │ ├── linear_readout.py │ │ ├── mlp_decoder.py │ │ └── multi_step_mlp_decoder.py │ └── encoders │ │ ├── __init__.py │ │ ├── conditional.py │ │ ├── dcrnn.py │ │ ├── dense_dcrnn.py │ │ ├── gcgru.py │ │ ├── gclstm.py │ │ ├── gcrnn.py │ │ ├── input_encoder.py │ │ ├── mlp.py │ │ ├── nri_dcrnn.py │ │ ├── rnn.py │ │ ├── stcn.py │ │ ├── tcn.py │ │ └── transformer.py ├── functional.py ├── layers │ ├── __init__.py │ ├── graph_convs │ │ ├── __init__.py │ │ ├── dense_spatial_conv.py │ │ ├── diff_conv.py │ │ ├── gat_conv.py │ │ ├── gated_gn.py │ │ ├── graph_attention.py │ │ ├── grin_cell.py │ │ └── spatio_temporal_att.py │ ├── link_predictor.py │ ├── norm │ │ ├── __init__.py │ │ ├── batch_norm.py │ │ ├── instance_norm.py │ │ ├── layer_norm.py │ │ └── norm.py │ ├── positional_encoding.py │ ├── spatial_attention.py │ └── temporal_attention.py ├── metrics │ ├── __init__.py │ ├── metric_base.py │ ├── metric_wrappers.py │ ├── metrics.py │ ├── multi_loss.py │ └── pinball_loss.py ├── models │ ├── __init__.py │ ├── imputation │ │ ├── __init__.py │ │ ├── grin_model.py │ │ └── rnni_models.py │ ├── rnn_model.py │ ├── stgn │ │ ├── __init__.py │ │ ├── dcrnn_model.py │ │ ├── gated_gn_model.py │ │ ├── graph_wavenet_model.py │ │ ├── nri_model.py │ │ ├── rnn2gcn_model.py │ │ └── stcn_model.py │ ├── tcn_model.py │ └── transformer_model.py ├── ops │ ├── __init__.py │ ├── grad_norm.py │ └── ops.py └── utils │ ├── __init__.py │ ├── casting.py │ └── utils.py ├── ops ├── __init__.py ├── connectivity.py ├── framearray.py ├── imputation.py ├── pattern.py ├── similarities.py └── test.py ├── predictors ├── __init__.py └── base_predictor.py ├── typing.py └── utils ├── __init__.py ├── download.py ├── experiment.py ├── io.py ├── neptune_utils.py ├── numpy_metrics.py ├── parser_utils.py ├── preprocessing.py └── python_utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | *.DS_STORE -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Graph Machine Learning Group 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /conda_env.yml: -------------------------------------------------------------------------------- 1 | name: sgp 2 | channels: 3 | - pytorch 4 | - pyg 5 | - defaults 6 | - conda-forge 7 | dependencies: 8 | - pip 9 | - pyg=2.0 10 | - python=3.8 11 | - pytorch=1.9 12 | - torchvision 13 | - torchaudio 14 | - wheel 15 | - pip: 16 | - einops 17 | - numpy 18 | - pandas 19 | - pytorch-lightning=1.5 20 | - scikit-learn 21 | - scipy 22 | - tables 23 | - test-tube 24 | - torchmetrics=0.7 25 | - tqdm 26 | -------------------------------------------------------------------------------- /config/largescale/dcrnn_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 500 8 | 9 | #### Training params ########################################################## 10 | batch_size: 1 11 | batch_inference: 12 12 | max_edges: 2200000 13 | cut_edges_uniformly: True 14 | 15 | epochs: 128 16 | batches_epoch: 32 17 | 18 | lr: 0.001 19 | use_lr_schedule: False 20 | 21 | #### Model params ############################################################# 22 | hidden_size: 64 23 | ff_size: 128 24 | kernel_size: 2 25 | n_layers: 1 26 | dropout: 0 27 | -------------------------------------------------------------------------------- /config/largescale/dcrnn_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | 8 | #### Training params ########################################################## 9 | batch_size: 1 10 | batch_inference: 12 11 | max_edges: 2000000 12 | cut_edges_uniformly: True 13 | 14 | epochs: 155 15 | batches_epoch: 32 16 | 17 | lr: 0.001 18 | use_lr_schedule: False 19 | 20 | #### Model params ############################################################# 21 | hidden_size: 64 22 | ff_size: 128 23 | kernel_size: 2 24 | n_layers: 1 25 | dropout: 0 26 | -------------------------------------------------------------------------------- /config/largescale/gatedgn_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 500 8 | 9 | #### Training params ########################################################## 10 | batch_size: 1 11 | batch_inference: 1 12 | max_edges: 2500000 13 | cut_edges_uniformly: True 14 | 15 | epochs: 987 16 | batches_epoch: 32 17 | 18 | lr: 0.001 19 | use_lr_schedule: False 20 | 21 | #### Model params ############################################################# 22 | hidden_size: 64 23 | enc_layers: 2 24 | gnn_layers: 2 25 | full_graph: False 26 | positional_encoding: True 27 | activation: 'silu' 28 | -------------------------------------------------------------------------------- /config/largescale/gatedgn_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | 8 | #### Training params ########################################################## 9 | batch_size: 1 10 | batch_inference: 1 11 | max_edges: 2500000 12 | cut_edges_uniformly: True 13 | 14 | epochs: 993 15 | batches_epoch: 32 16 | 17 | lr: 0.001 18 | use_lr_schedule: False 19 | 20 | #### Model params ############################################################# 21 | hidden_size: 64 22 | enc_layers: 2 23 | gnn_layers: 2 24 | full_graph: False 25 | positional_encoding: True 26 | activation: 'silu' 27 | -------------------------------------------------------------------------------- /config/largescale/gwnet_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 500 8 | 9 | #### Training params ########################################################## 10 | batch_size: 1 11 | batch_inference: 6 12 | 13 | epochs: 142 14 | batches_epoch: 32 15 | 16 | lr: 0.001 17 | use_lr_schedule: False 18 | 19 | #### Model params ############################################################# 20 | hidden_size: 32 21 | ff_size: 256 22 | n_layers: 8 23 | dropout: 0.3 24 | temporal_kernel_size: 2 25 | spatial_kernel_size: 2 26 | dilation: 2 27 | dilation_mod: 2 28 | norm: batch 29 | learned_adjacency: True 30 | -------------------------------------------------------------------------------- /config/largescale/gwnet_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | 8 | #### Training params ########################################################## 9 | batch_size: 2 10 | batch_inference: 10 11 | 12 | epochs: 87 13 | batches_epoch: 32 14 | 15 | lr: 0.001 16 | use_lr_schedule: False 17 | 18 | #### Model params ############################################################# 19 | hidden_size: 32 20 | ff_size: 256 21 | n_layers: 8 22 | dropout: 0.3 23 | temporal_kernel_size: 2 24 | spatial_kernel_size: 2 25 | dilation: 2 26 | dilation_mod: 2 27 | norm: batch 28 | learned_adjacency: True 29 | -------------------------------------------------------------------------------- /config/largescale/sgp_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 1 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 500 8 | 9 | #### Reservoir params ######################################################### 10 | reservoir_size: 16 11 | reservoir_layers: 6 12 | leaking_rate: 1. 13 | spectral_radius: 0.99 14 | density: 0.7 15 | alpha_decay: True 16 | preprocess_exogenous: True 17 | keep_raw: True 18 | 19 | #### SGP params ############################################################## 20 | bidirectional: False 21 | receptive_field: 2 22 | undirected: False 23 | add_self_loops: False 24 | global_attr: True 25 | 26 | #### Training params ########################################################## 27 | iid_sampling: True 28 | sgp_preprocessing: False 29 | workers: 16 30 | batch_size: 4096 31 | batch_inference: 16 32 | 33 | epochs: 12422 # 1:00:00 - 2:49 of 13033 34 | batches_epoch: 32 35 | scale_target: False 36 | 37 | lr: 0.001 38 | use_lr_schedule: False 39 | 40 | #### Model params ############################################################# 41 | emb_size: 32 42 | hidden_size: 960 43 | mlp_size: 256 44 | n_layers: 2 45 | dropout: 0. 46 | positional_encoding: True 47 | resnet: True 48 | fully_connected: False -------------------------------------------------------------------------------- /config/largescale/sgp_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 1 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | 8 | #### Reservoir params ######################################################### 9 | reservoir_size: 16 10 | reservoir_layers: 8 # 4 11 | leaking_rate: 1. # 0.8 12 | spectral_radius: 0.99 13 | density: 0.7 14 | alpha_decay: True 15 | preprocess_exogenous: True 16 | keep_raw: True 17 | 18 | #### SGP params ############################################################## 19 | bidirectional: False 20 | receptive_field: 2 21 | undirected: False 22 | add_self_loops: False 23 | global_attr: True 24 | 25 | #### Training params ########################################################## 26 | iid_sampling: True 27 | sgp_preprocessing: False 28 | workers: 16 29 | batch_size: 4096 30 | batch_inference: 16 31 | 32 | epochs: 12496 # 1:00:00 - 3:50 of 13348 33 | batches_epoch: 32 34 | scale_target: False 35 | 36 | lr: 0.001 37 | use_lr_schedule: False 38 | 39 | #### Model params ############################################################# 40 | emb_size: 32 41 | hidden_size: 960 42 | mlp_size: 256 43 | n_layers: 2 44 | dropout: 0. 45 | positional_encoding: True 46 | resnet: True 47 | fully_connected: False -------------------------------------------------------------------------------- /config/largescale_100nn/dcrnn_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 100 8 | 9 | #### Training params ########################################################## 10 | batch_size: 2 11 | batch_inference: 12 12 | 13 | epochs: 162 14 | batches_epoch: 32 15 | 16 | lr: 0.001 17 | use_lr_schedule: False 18 | 19 | #### Model params ############################################################# 20 | hidden_size: 64 21 | ff_size: 128 22 | kernel_size: 2 23 | n_layers: 1 24 | dropout: 0 25 | -------------------------------------------------------------------------------- /config/largescale_100nn/dcrnn_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | adj_knn: 100 8 | 9 | #### Training params ########################################################## 10 | batch_size: 2 11 | batch_inference: 12 12 | 13 | epochs: 230 14 | batches_epoch: 32 15 | scale_target: False 16 | 17 | lr: 0.001 18 | use_lr_schedule: False 19 | 20 | #### Model params ############################################################# 21 | hidden_size: 64 22 | ff_size: 128 23 | kernel_size: 2 24 | n_layers: 1 25 | dropout: 0 26 | -------------------------------------------------------------------------------- /config/largescale_100nn/gatedgn_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 100 8 | 9 | #### Training params ########################################################## 10 | batch_size: 4 11 | batch_inference: 8 12 | 13 | epochs: 924 14 | batches_epoch: 32 15 | scale_target: False 16 | 17 | lr: 0.001 18 | use_lr_schedule: False 19 | 20 | #### Model params ############################################################# 21 | hidden_size: 64 22 | enc_layers: 2 23 | gnn_layers: 2 24 | full_graph: False 25 | positional_encoding: True 26 | activation: 'silu' 27 | -------------------------------------------------------------------------------- /config/largescale_100nn/gatedgn_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | adj_knn: 100 8 | 9 | #### Training params ########################################################## 10 | batch_size: 5 11 | batch_inference: 8 12 | 13 | epochs: 947 14 | batches_epoch: 32 15 | scale_target: False 16 | 17 | lr: 0.001 18 | use_lr_schedule: False 19 | 20 | #### Model params ############################################################# 21 | hidden_size: 64 22 | enc_layers: 2 23 | gnn_layers: 2 24 | full_graph: False 25 | positional_encoding: True 26 | activation: 'silu' 27 | -------------------------------------------------------------------------------- /config/largescale_100nn/gwnet_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 100 8 | 9 | #### Training params ########################################################## 10 | batch_size: 1 11 | batch_inference: 6 12 | 13 | epochs: 271 14 | batches_epoch: 32 15 | scale_target: False 16 | 17 | lr: 0.001 18 | use_lr_schedule: False 19 | 20 | #### Model params ############################################################# 21 | hidden_size: 32 22 | ff_size: 256 23 | n_layers: 8 24 | dropout: 0.3 25 | temporal_kernel_size: 2 26 | spatial_kernel_size: 2 27 | dilation: 2 28 | dilation_mod: 2 29 | norm: batch 30 | learned_adjacency: True 31 | -------------------------------------------------------------------------------- /config/largescale_100nn/gwnet_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 36 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | adj_knn: 100 8 | 9 | #### Training params ########################################################## 10 | batch_size: 2 11 | batch_inference: 10 12 | 13 | epochs: 227 14 | batches_epoch: 32 15 | scale_target: False 16 | 17 | lr: 0.001 18 | use_lr_schedule: False 19 | 20 | #### Model params ############################################################# 21 | hidden_size: 32 22 | ff_size: 256 23 | n_layers: 8 24 | dropout: 0.3 25 | temporal_kernel_size: 2 26 | spatial_kernel_size: 2 27 | dilation: 2 28 | dilation_mod: 2 29 | norm: batch 30 | learned_adjacency: True 31 | -------------------------------------------------------------------------------- /config/largescale_100nn/sgp_cer.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 1 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: cer 7 | adj_knn: 100 8 | 9 | #### Reservoir params ######################################################### 10 | reservoir_size: 16 11 | reservoir_layers: 6 12 | leaking_rate: 1. 13 | spectral_radius: 0.99 14 | density: 0.7 15 | alpha_decay: True 16 | preprocess_exogenous: True 17 | keep_raw: True 18 | 19 | #### SGP params ############################################################## 20 | bidirectional: False 21 | receptive_field: 2 22 | undirected: False 23 | add_self_loops: False 24 | global_attr: True 25 | 26 | #### Training params ########################################################## 27 | iid_sampling: True 28 | sgp_preprocessing: False 29 | workers: 16 30 | batch_size: 4096 31 | batch_inference: 16 32 | 33 | epochs: 12980 # 1:00:00 - 1:00 of 13199 34 | batches_epoch: 32 35 | scale_target: False 36 | 37 | lr: 0.001 38 | use_lr_schedule: False 39 | 40 | #### Model params ############################################################# 41 | emb_size: 32 42 | hidden_size: 960 43 | mlp_size: 256 44 | n_layers: 2 45 | dropout: 0. 46 | positional_encoding: True 47 | resnet: True 48 | fully_connected: False -------------------------------------------------------------------------------- /config/largescale_100nn/sgp_pv.yaml: -------------------------------------------------------------------------------- 1 | #### Dataset params ########################################################### 2 | window: 1 3 | horizon: 22 4 | horizon_lag: 7 5 | 6 | dataset_name: pv 7 | adj_knn: 100 8 | 9 | #### Reservoir params ######################################################### 10 | reservoir_size: 16 11 | reservoir_layers: 8 # 4 12 | leaking_rate: 1. # 0.8 13 | spectral_radius: 0.99 14 | density: 0.7 15 | alpha_decay: True 16 | preprocess_exogenous: True 17 | keep_raw: True 18 | 19 | #### SGP params ############################################################## 20 | bidirectional: False 21 | receptive_field: 2 22 | undirected: False 23 | add_self_loops: False 24 | global_attr: True 25 | 26 | #### Training params ########################################################## 27 | iid_sampling: True 28 | sgp_preprocessing: False 29 | workers: 16 30 | batch_size: 4096 31 | batch_inference: 16 32 | 33 | epochs: 12897 # 1:00:00 - 1:00 of 13115 34 | batches_epoch: 32 35 | scale_target: False 36 | 37 | lr: 0.001 38 | use_lr_schedule: False 39 | 40 | #### Model params ############################################################# 41 | emb_size: 32 42 | hidden_size: 960 43 | mlp_size: 256 44 | n_layers: 2 45 | dropout: 0. 46 | positional_encoding: True 47 | resnet: True 48 | fully_connected: False -------------------------------------------------------------------------------- /config/traffic/dcrnn.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 12 3 | horizon: 12 4 | 5 | #### Splitting params ######################################################### 6 | val_len: 0.1 7 | test_len: 0.2 8 | 9 | #### Training params ########################################################## 10 | epochs: 300 11 | patience: 50 12 | batch_size: 64 13 | 14 | lr: 0.01 15 | use_lr_schedule: True 16 | lr_gamma: 0.1 17 | lr_milestones: [ 20, 30, 40 ] 18 | 19 | #### Model params ############################################################# 20 | hidden_size: 64 21 | ff_size: 128 22 | kernel_size: 2 23 | n_layers: 1 24 | dropout: 0 25 | -------------------------------------------------------------------------------- /config/traffic/gatedgn_bay.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 12 3 | horizon: 12 4 | 5 | #### Splitting params ######################################################### 6 | val_len: 0.1 7 | test_len: 0.2 8 | 9 | #### Training params ########################################################## 10 | batch_size: 16 11 | epochs: 300 12 | patience: 50 13 | 14 | lr: 0.0002 15 | use_lr_schedule: True 16 | lr_milestones: [ 50 ] 17 | lr_gamma: 0.1 18 | 19 | #### Model params ############################################################# 20 | hidden_size: 64 21 | enc_layers: 2 22 | gnn_layers: 2 23 | full_graph: False # True for FC 24 | activation: 'silu' 25 | -------------------------------------------------------------------------------- /config/traffic/gatedgn_la.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 12 3 | horizon: 12 4 | 5 | #### Splitting params ######################################################### 6 | val_len: 0.1 7 | test_len: 0.2 8 | 9 | #### Training params ########################################################## 10 | batch_size: 16 11 | epochs: 300 12 | patience: 50 13 | 14 | lr: 0.005 15 | use_lr_schedule: True 16 | lr_milestones: [ 20, 30, 40 ] 17 | lr_gamma: 0.1 18 | 19 | #### Model params ############################################################# 20 | hidden_size: 64 21 | enc_layers: 2 22 | gnn_layers: 2 23 | full_graph: False # True for FC 24 | activation: 'silu' 25 | -------------------------------------------------------------------------------- /config/traffic/gesn_bay.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Training params ########################################################## 6 | workers: 4 7 | l2_reg: 0.001 8 | 9 | #### SGP params ############################################################## 10 | reservoir_layers: 4 11 | reservoir_size: 128 12 | leaking_rate: 0.9 13 | spectral_radius: 0.9 14 | density: 1. 15 | preprocess_exogenous: true 16 | alpha_decay: true 17 | -------------------------------------------------------------------------------- /config/traffic/gesn_la.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Training params ########################################################## 6 | workers: 4 7 | l2_reg: 0.001 8 | 9 | #### SGP params ############################################################## 10 | reservoir_layers: 3 11 | reservoir_size: 320 12 | leaking_rate: 0.9 13 | spectral_radius: 0.9 14 | density: 1. 15 | preprocess_exogenous: true 16 | alpha_decay: true 17 | -------------------------------------------------------------------------------- /config/traffic/gwnet.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 12 3 | horizon: 12 4 | 5 | #### Splitting params ######################################################### 6 | val_len: 0.1 7 | test_len: 0.2 8 | 9 | #### Training params ########################################################## 10 | batch_size: 64 11 | l2_reg: 0.0001 12 | epochs: 300 13 | patience: 50 14 | 15 | lr: 0.001 16 | use_lr_schedule: False 17 | 18 | #### Model params ############################################################# 19 | hidden_size: 32 20 | ff_size: 256 21 | n_layers: 8 22 | dropout: 0.3 23 | temporal_kernel_size: 2 24 | spatial_kernel_size: 2 25 | dilation: 2 26 | dilation_mod: 2 27 | norm: batch 28 | learned_adjacency: True 29 | -------------------------------------------------------------------------------- /config/traffic/rnn.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 12 3 | horizon: 12 4 | 5 | #### Splitting params ######################################################### 6 | val_len: 0.1 7 | test_len: 0.2 8 | 9 | #### Training params ########################################################## 10 | batch_size: 64 11 | epochs: 300 12 | patience: 50 13 | 14 | lr: 0.001 15 | use_lr_schedule: True 16 | lr_gamma: 0.1 17 | lr_milestones: [ 20 ] 18 | 19 | #### Model params ############################################################# 20 | hidden_size: 128 21 | ff_size: 256 22 | rec_layers: 1 23 | ff_layers: 1 24 | rec_dropout: 0 25 | ff_dropout: 0.1 26 | cell_type: 'lstm' 27 | -------------------------------------------------------------------------------- /config/traffic/sgp_bay.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | dataset_name: bay 5 | 6 | #### Reservoir params ######################################################### 7 | reservoir_size: 128 8 | reservoir_layers: 1 # 4 9 | leaking_rate: 0.8 10 | spectral_radius: 0.9 11 | density: 0.7 12 | alpha_decay: True 13 | preprocess_exogenous: True 14 | keep_raw: True 15 | 16 | #### SGP params ############################################################## 17 | bidirectional: True 18 | receptive_field: 4 19 | undirected: False 20 | add_self_loops: False 21 | global_attr: True 22 | 23 | #### Performance params ####################################################### 24 | iid_sampling: False 25 | sgp_preprocessing: False 26 | batch_size: 64 27 | batch_inference: 64 28 | 29 | #### Training params ########################################################## 30 | l2_reg: 0. 31 | use_lr_schedule: True 32 | lr_milestones: [ 100, 150 ] 33 | lr: 0.0035 34 | epochs: 200 35 | patience: 100 36 | batches_epoch: 300 37 | workers: 4 38 | 39 | #### Model params ############################################################# 40 | hidden_size: 960 41 | mlp_size: 256 42 | n_layers: 2 43 | dropout: 0.3 44 | positional_encoding: True 45 | resnet: true 46 | fully_connected: False -------------------------------------------------------------------------------- /config/traffic/sgp_bay_abl1.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Reservoir params ######################################################### 6 | reservoir_size: 128 7 | reservoir_layers: 1 # 4 8 | leaking_rate: 0.8 9 | spectral_radius: 0.9 10 | density: 0.7 11 | alpha_decay: True 12 | preprocess_exogenous: True 13 | keep_raw: True 14 | 15 | #### SGP params ############################################################## 16 | bidirectional: True 17 | receptive_field: 0 18 | undirected: False 19 | add_self_loops: False 20 | global_attr: False 21 | 22 | #### Performance params ####################################################### 23 | iid_sampling: False 24 | sgp_preprocessing: False 25 | batch_size: 64 26 | batch_inference: 64 27 | 28 | #### Training params ########################################################## 29 | l2_reg: 0. 30 | use_lr_schedule: True 31 | lr_milestones: [ 100, 150 ] 32 | lr: 0.0035 33 | epochs: 200 34 | patience: 100 35 | batches_epoch: 300 36 | workers: 4 37 | 38 | #### Model params ############################################################# 39 | hidden_size: 960 40 | mlp_size: 256 41 | n_layers: 2 42 | dropout: 0.3 43 | positional_encoding: True 44 | resnet: true 45 | fully_connected: False -------------------------------------------------------------------------------- /config/traffic/sgp_bay_abl2.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Reservoir params ######################################################### 6 | reservoir_size: 128 7 | reservoir_layers: 1 # 4 8 | leaking_rate: 0.8 # 0.8 9 | spectral_radius: 0.9 10 | density: 0.7 11 | alpha_decay: True 12 | preprocess_exogenous: True 13 | keep_raw: True 14 | 15 | #### SGP params ############################################################## 16 | bidirectional: False 17 | receptive_field: 1 18 | undirected: False 19 | add_self_loops: False 20 | global_attr: False 21 | 22 | #### Performance params ####################################################### 23 | iid_sampling: False 24 | sgp_preprocessing: False 25 | batch_size: 64 26 | batch_inference: 64 27 | 28 | #### Training params ########################################################## 29 | l2_reg: 0. 30 | use_lr_schedule: True 31 | lr_milestones: [ 100, 150 ] 32 | lr: 0.0035 33 | epochs: 200 34 | patience: 100 35 | batches_epoch: 300 36 | workers: 4 37 | 38 | #### Model params ############################################################# 39 | hidden_size: 960 40 | mlp_size: 256 41 | n_layers: 2 42 | dropout: 0.3 43 | positional_encoding: True 44 | resnet: true 45 | fully_connected: False -------------------------------------------------------------------------------- /config/traffic/sgp_bay_abl3.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Reservoir params ######################################################### 6 | reservoir_size: 128 7 | reservoir_layers: 1 # 4 8 | leaking_rate: 0.8 # 0.8 9 | spectral_radius: 0.9 10 | density: 0.7 11 | alpha_decay: True 12 | preprocess_exogenous: True 13 | keep_raw: True 14 | 15 | 16 | #### SGP params ############################################################## 17 | bidirectional: True 18 | receptive_field: 4 19 | undirected: False 20 | add_self_loops: False 21 | global_attr: True 22 | 23 | #### Performance params ####################################################### 24 | iid_sampling: False 25 | sgp_preprocessing: False 26 | batch_size: 64 27 | batch_inference: 64 28 | 29 | #### Training params ########################################################## 30 | l2_reg: 0. 31 | use_lr_schedule: True 32 | lr_milestones: [ 100, 150 ] 33 | lr: 0.0035 34 | epochs: 200 35 | patience: 100 36 | batches_epoch: 300 37 | workers: 4 38 | 39 | #### Model params ############################################################# 40 | hidden_size: 960 41 | mlp_size: 256 42 | n_layers: 2 43 | dropout: 0.3 44 | positional_encoding: True 45 | resnet: true 46 | fully_connected: True -------------------------------------------------------------------------------- /config/traffic/sgp_la.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | dataset_name: la 5 | 6 | #### Reservoir params ######################################################### 7 | reservoir_size: 64 8 | reservoir_layers: 2 # 4 9 | leaking_rate: 0.9 # 0.8 10 | spectral_radius: 0.9 11 | density: 0.7 12 | alpha_decay: True 13 | preprocess_exogenous: True 14 | keep_raw: True 15 | 16 | #### SGP params ############################################################## 17 | bidirectional: True 18 | receptive_field: 4 19 | undirected: False 20 | add_self_loops: False 21 | global_attr: True 22 | 23 | #### Performance params ####################################################### 24 | iid_sampling: False 25 | sgp_preprocessing: False 26 | batch_size: 64 27 | batch_inference: 64 28 | 29 | #### Training params ########################################################## 30 | l2_reg: 0. 31 | use_lr_schedule: True 32 | lr_milestones: [ 40, 80, 120 ] # [100, 150] 33 | lr: 0.004 # 0.0035 34 | epochs: 200 35 | patience: 100 36 | batches_epoch: 300 37 | workers: 4 38 | 39 | #### Model params ############################################################# 40 | hidden_size: 960 41 | mlp_size: 256 42 | n_layers: 2 43 | dropout: 0.3 44 | positional_encoding: True 45 | resnet: true 46 | fully_connected: False -------------------------------------------------------------------------------- /config/traffic/sgp_la_abl1.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Reservoir params ######################################################### 6 | reservoir_size: 64 7 | reservoir_layers: 2 # 4 8 | leaking_rate: 0.9 # 0.8 9 | spectral_radius: 0.9 10 | density: 0.7 11 | alpha_decay: True 12 | preprocess_exogenous: True 13 | keep_raw: True 14 | 15 | #### SGP params ############################################################## 16 | bidirectional: True 17 | receptive_field: 0 18 | undirected: False 19 | add_self_loops: False 20 | global_attr: False 21 | 22 | #### Performance params ####################################################### 23 | iid_sampling: False 24 | sgp_preprocessing: False 25 | batch_size: 64 26 | batch_inference: 64 27 | 28 | #### Training params ########################################################## 29 | l2_reg: 0. 30 | use_lr_schedule: True 31 | lr_milestones: [ 40, 80, 120 ] # [100, 150] 32 | lr: 0.004 # 0.0035 33 | epochs: 200 34 | patience: 100 35 | batches_epoch: 300 36 | workers: 4 37 | 38 | #### Model params ############################################################# 39 | hidden_size: 960 40 | mlp_size: 256 41 | n_layers: 2 42 | dropout: 0.3 43 | positional_encoding: True 44 | resnet: true 45 | fully_connected: False -------------------------------------------------------------------------------- /config/traffic/sgp_la_abl2.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Reservoir params ######################################################### 6 | reservoir_size: 64 7 | reservoir_layers: 2 8 | leaking_rate: 0.9 9 | spectral_radius: 0.9 10 | density: 0.7 11 | alpha_decay: True 12 | preprocess_exogenous: True 13 | keep_raw: True 14 | 15 | #### SGP params ############################################################## 16 | bidirectional: True 17 | receptive_field: 4 18 | undirected: False 19 | add_self_loops: False 20 | global_attr: True 21 | 22 | #### Performance params ####################################################### 23 | iid_sampling: False 24 | sgp_preprocessing: False 25 | batch_size: 64 26 | batch_inference: 64 27 | 28 | #### Training params ########################################################## 29 | l2_reg: 0. 30 | use_lr_schedule: True 31 | lr_milestones: [ 40, 80, 120 ] 32 | lr: 0.004 33 | epochs: 200 34 | patience: 100 35 | batches_epoch: 300 36 | workers: 4 37 | 38 | #### Model params ############################################################# 39 | hidden_size: 960 40 | mlp_size: 256 41 | n_layers: 2 42 | dropout: 0.3 43 | positional_encoding: True 44 | resnet: true 45 | fully_connected: True -------------------------------------------------------------------------------- /config/traffic/sgp_la_abl3.yaml: -------------------------------------------------------------------------------- 1 | #### Windowing params ######################################################### 2 | window: 1 3 | horizon: 12 4 | 5 | #### Reservoir params ######################################################### 6 | reservoir_size: 64 7 | reservoir_layers: 2 # 4 8 | leaking_rate: 0.9 # 0.8 9 | spectral_radius: 0.9 10 | density: 0.7 11 | alpha_decay: True 12 | preprocess_exogenous: True 13 | keep_raw: True 14 | 15 | #### SGP params ############################################################## 16 | bidirectional: False 17 | receptive_field: 1 18 | undirected: False 19 | add_self_loops: False 20 | global_attr: False 21 | 22 | #### Performance params ####################################################### 23 | iid_sampling: False 24 | sgp_preprocessing: False 25 | batch_size: 64 26 | batch_inference: 64 27 | 28 | #### Training params ########################################################## 29 | l2_reg: 0. 30 | use_lr_schedule: True 31 | lr_milestones: [ 40, 80, 120 ] # [100, 150] 32 | lr: 0.004 # 0.0035 33 | epochs: 200 34 | patience: 100 35 | batches_epoch: 300 36 | workers: 4 37 | 38 | #### Model params ############################################################# 39 | hidden_size: 960 40 | mlp_size: 256 41 | n_layers: 2 42 | dropout: 0.3 43 | positional_encoding: True 44 | resnet: true 45 | fully_connected: False -------------------------------------------------------------------------------- /experiments/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Graph-Machine-Learning-Group/sgp/ae1abd55e15cf21d5ac0c5289a508751ccb3d589/experiments/__init__.py -------------------------------------------------------------------------------- /lib/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | base_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) 4 | config_file = os.path.join(base_dir, 'config.yaml') 5 | config = { 6 | 'config_dir': os.path.join(base_dir, 'config'), 7 | 'data_dir': os.path.join(base_dir, 'datasets'), 8 | 'logs_dir': os.path.join(base_dir, 'log') 9 | } 10 | -------------------------------------------------------------------------------- /lib/dataloader/__init__.py: -------------------------------------------------------------------------------- 1 | from .iid_dataloader import IIDLoader 2 | from .sgp_dataloader import SGPLoader 3 | from .subgraph_dataloader import SubsetLoader, SubgraphLoader 4 | -------------------------------------------------------------------------------- /lib/dataloader/iid_dataloader.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional 2 | 3 | import torch 4 | from torch.utils import data 5 | from torch.utils.data import Sampler 6 | from tsl.data import Data 7 | 8 | from lib.datasets import IIDDataset 9 | 10 | 11 | class IIDSampler(Sampler): 12 | 13 | def __init__(self, num_batches): 14 | super(IIDSampler, self).__init__(torch.arange(num_batches)) 15 | self.num_batches = num_batches 16 | 17 | def __iter__(self): 18 | for i in range(self.num_batches): 19 | yield i 20 | 21 | def __len__(self): 22 | return self.num_batches 23 | 24 | 25 | class IIDLoader(data.DataLoader): 26 | 27 | def __init__(self, dataset: IIDDataset, 28 | batch_size: Optional[int] = 1024, 29 | num_batches: int = 1000, 30 | num_workers: int = 0, 31 | **kwargs): 32 | if 'collate_fn' in kwargs: 33 | del kwargs['collate_fn'] 34 | self._batch_size = batch_size 35 | dataset.make_random_iid(batch_size) 36 | super().__init__(dataset, 37 | batch_size=1, 38 | sampler=IIDSampler(num_batches), 39 | num_workers=num_workers, 40 | collate_fn=self.collate, 41 | **kwargs) 42 | 43 | def collate(self, data_list: List[Data]): 44 | batch = data_list[0] 45 | batch.__dict__['batch_size'] = self._batch_size 46 | return batch 47 | -------------------------------------------------------------------------------- /lib/dataloader/sgp_dataloader.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional 2 | 3 | import torch 4 | from torch.utils import data 5 | from tsl.data import static_graph_collate, Batch, Data, SpatioTemporalDataset 6 | 7 | from lib.sgp_preprocessing import sgp_spatial_support 8 | 9 | 10 | class SGPLoader(data.DataLoader): 11 | 12 | def __init__(self, dataset: SpatioTemporalDataset, keys: List = None, 13 | k: int = 2, 14 | undirected: bool = False, 15 | add_self_loops: bool = False, 16 | remove_self_loops: bool = False, 17 | bidirectional: bool = False, 18 | global_attr: bool = False, 19 | batch_size: Optional[int] = 1, 20 | shuffle: bool = False, 21 | num_workers: int = 0, 22 | **kwargs): 23 | if 'collate_fn' in kwargs: 24 | del kwargs['collate_fn'] 25 | self.keys = keys 26 | self.k = k 27 | self.undirected = undirected 28 | self.add_self_loops = add_self_loops 29 | self.remove_self_loops = remove_self_loops 30 | self.bidirectional = bidirectional 31 | self.global_attr = global_attr 32 | super().__init__(dataset, 33 | shuffle=shuffle, 34 | batch_size=batch_size, 35 | num_workers=num_workers, 36 | collate_fn=self.collate, 37 | **kwargs) 38 | 39 | def collate(self, data_list: List[Data]): 40 | elem = data_list[0] 41 | edge_index, edge_weight = elem.edge_index, elem.edge_weight 42 | num_nodes = elem.num_nodes 43 | support = sgp_spatial_support(edge_index, edge_weight, 44 | num_nodes=num_nodes, 45 | k=self.k, 46 | undirected=self.undirected, 47 | add_self_loops=self.add_self_loops, 48 | remove_self_loops=self.remove_self_loops, 49 | bidirectional=self.bidirectional, 50 | global_attr=self.global_attr) 51 | keys = self.keys 52 | if keys is None: 53 | keys = [k for k, pattern in elem.pattern.items() 54 | if 'n' in pattern and k in elem.input] 55 | # subsample every item in batch 56 | for sample in data_list: 57 | for key in keys: 58 | x = sample[key] 59 | sample[key] = torch.cat([x] + [adj @ x for adj in support], 60 | dim=-1) 61 | # collate tensors in batch 62 | batch = static_graph_collate(data_list, Batch) 63 | 64 | # subset sampler can only be used over set of nodes (without edges) 65 | if 'edge_index' in batch: 66 | del batch['edge_index'] 67 | if 'edge_weight' in batch: 68 | del batch['edge_weight'] 69 | batch.__dict__['batch_size'] = len(data_list) 70 | 71 | return batch 72 | -------------------------------------------------------------------------------- /lib/datamodule/__init__.py: -------------------------------------------------------------------------------- 1 | from .sgp_datamodule import SGPDataModule 2 | from .subgraph_datamodule import SubgraphDataModule 3 | -------------------------------------------------------------------------------- /lib/datamodule/subgraph_datamodule.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, Mapping 2 | 3 | from tsl.data import SpatioTemporalDataModule, SpatioTemporalDataset, Splitter 4 | from tsl.data.loader import StaticGraphLoader 5 | 6 | from ..dataloader import SubgraphLoader, SubsetLoader 7 | 8 | 9 | class SubgraphDataModule(SpatioTemporalDataModule): 10 | def __init__(self, dataset: SpatioTemporalDataset, 11 | max_nodes_training: int = None, 12 | receptive_field: int = 1, 13 | max_edges: Optional[int] = None, 14 | cut_edges_uniformly: bool = False, 15 | val_stride: int = None, 16 | scalers: Optional[Mapping] = None, 17 | splitter: Optional[Splitter] = None, 18 | batch_size: int = 32, 19 | batch_inference: int = 32, 20 | workers: int = 0, 21 | pin_memory: bool = False): 22 | super(SubgraphDataModule, self).__init__(dataset, 23 | scalers=scalers, 24 | splitter=splitter, 25 | mask_scaling=True, 26 | batch_size=batch_size, 27 | workers=workers, 28 | pin_memory=pin_memory) 29 | if max_nodes_training is not None: 30 | if receptive_field > 0: 31 | self._trainloader_class = SubgraphLoader 32 | self.train_loader_kwargs = dict( 33 | num_nodes=max_nodes_training, 34 | k=receptive_field, 35 | max_edges=max_edges, 36 | cut_edges_uniformly=cut_edges_uniformly) 37 | else: 38 | self._trainloader_class = SubsetLoader 39 | self.train_loader_kwargs = dict(max_nodes=max_nodes_training) 40 | else: 41 | self._trainloader_class = StaticGraphLoader 42 | self.train_loader_kwargs = dict() 43 | self.batch_inference = batch_inference 44 | self.val_stride = val_stride 45 | 46 | def setup(self, stage=None): 47 | super(SubgraphDataModule, self).setup(stage) 48 | if self.val_stride is not None: 49 | self.valset.indices = self.valset.indices[::self.val_stride] 50 | 51 | def train_dataloader(self, shuffle=True, batch_size=None): 52 | if self.trainset is None: 53 | return None 54 | return self._trainloader_class(self.trainset, 55 | batch_size=batch_size or self.batch_size, 56 | shuffle=shuffle, 57 | num_workers=self.workers, 58 | pin_memory=self.pin_memory, 59 | drop_last=True, 60 | **self.train_loader_kwargs) 61 | 62 | def val_dataloader(self, shuffle=False, batch_size=None): 63 | if self.valset is None: 64 | return None 65 | return StaticGraphLoader(self.valset, 66 | batch_size=batch_size or self.batch_inference, 67 | shuffle=shuffle, 68 | num_workers=self.workers) 69 | 70 | def test_dataloader(self, shuffle=False, batch_size=None): 71 | if self.testset is None: 72 | return None 73 | return StaticGraphLoader(self.testset, 74 | batch_size=batch_size or self.batch_inference, 75 | shuffle=shuffle, 76 | num_workers=self.workers) 77 | 78 | @staticmethod 79 | def add_argparse_args(parser, **kwargs): 80 | parser.add_argument('--receptive-field', type=int, default=1) 81 | parser.add_argument('--max-nodes-training', type=int, default=None) 82 | parser.add_argument('--max-edges', type=int, default=None) 83 | parser.add_argument('--cut-edges-uniformly', type=bool, default=False) 84 | parser.add_argument('--batch-size', type=int, default=32) 85 | parser.add_argument('--batch-inference', type=int, default=32) 86 | parser.add_argument('--mask-scaling', type=bool, default=True) 87 | parser.add_argument('--workers', type=int, default=0) 88 | return parser 89 | -------------------------------------------------------------------------------- /lib/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from .cer_en import CEREn 2 | from .iid_dataset import IIDDataset 3 | from .pv import PvUS 4 | -------------------------------------------------------------------------------- /lib/datasets/pv.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Union, List 3 | 4 | import pandas as pd 5 | from tsl.datasets.prototypes import PandasDataset 6 | from tsl.utils.python_utils import ensure_list 7 | 8 | from .. import config 9 | 10 | 11 | class PvUS(PandasDataset): 12 | r"""Simulated solar power production from more than 5,000 photovoltaic 13 | plants in the US. 14 | 15 | Data are provided by `National Renewable Energy Laboratory (NREL) 16 | `_'s `Solar Power Data for Integration Studies 17 | `_. Original raw data 18 | consist of 1 year (2006) of 5-minute solar power (in MW) for approximately 19 | 5,000 synthetic PV plants in the United States. 20 | 21 | Preprocessed data are resampled in 10-minutes intervals taking the average. 22 | The entire dataset contains 5016 plants, divided in two macro zones (east 23 | and west). The "east" zone contains 4084 plants, the "west" zone has 1082 24 | plants. Some states appear in both zones, with plants at same geographical 25 | position. When loading the entire datasets, duplicated plants in "east" zone 26 | are dropped. 27 | """ 28 | 29 | available_zones = ['east', 'west'] 30 | similarity_options = {'distance', 'correntropy'} 31 | 32 | def __init__(self, zones: Union[str, List] = None, mask_zeros: bool = False, 33 | root: str = None, freq: str = None): 34 | # set root path 35 | if zones is None: 36 | zones = self.available_zones 37 | else: 38 | zones = ensure_list(zones) 39 | if not set(zones).issubset(self.available_zones): 40 | invalid_zones = set(zones).difference(self.available_zones) 41 | raise ValueError(f"Invalid zones {invalid_zones}. " 42 | f"Allowed zones are {self.available_zones}.") 43 | self.zones = zones 44 | self.mask_zeros = mask_zeros 45 | self.root = config['data_dir'] + '/pv_us/' if root is None else root 46 | # set name 47 | name = "PvUS" if len(zones) == 2 else f"PvUS-{zones[0]}" 48 | # load dataset 49 | actual, mask, metadata = self.load(mask_zeros) 50 | super().__init__(target=actual, mask=mask, freq=freq, 51 | similarity_score="distance", 52 | spatial_aggregation="sum", 53 | temporal_aggregation="mean", 54 | name=name) 55 | self.add_covariate('metadata', metadata, pattern='n f') 56 | 57 | @property 58 | def raw_file_names(self): 59 | return [f'{zone}.h5' for zone in self.zones] 60 | 61 | @property 62 | def required_file_names(self): 63 | return self.raw_file_names 64 | 65 | def load_raw(self): 66 | actual, metadata = [], [] 67 | for zone in self.zones: 68 | # load zone data 69 | zone_path = os.path.join(self.root_dir, f'{zone}.h5') 70 | actual.append(pd.read_hdf(zone_path, key='actual')) 71 | metadata.append(pd.read_hdf(zone_path, key='metadata')) 72 | # concat zone and sort by plant id 73 | actual = pd.concat(actual, axis=1).sort_index(axis=1, level=0) 74 | metadata = pd.concat(metadata, axis=0).sort_index() 75 | # drop duplicated farms when loading whole dataset 76 | if len(self.zones) == 2: 77 | duplicated_farms = metadata.index[[s_id.endswith('-east') 78 | for s_id in metadata.state_id]] 79 | metadata = metadata.drop(duplicated_farms, axis=0) 80 | actual = actual.drop(duplicated_farms, axis=1, level=0) 81 | return actual, metadata 82 | 83 | def load(self, mask_zeros): 84 | actual, metadata = self.load_raw() 85 | mask = (actual > 0) if mask_zeros else None 86 | return actual, mask, metadata 87 | 88 | def compute_similarity(self, method: str, theta: float = 150, **kwargs): 89 | if method == "distance": 90 | from tsl.ops.similarities import (gaussian_kernel, 91 | geographical_distance) 92 | # compute distances from latitude and longitude degrees 93 | loc_coord = self.metadata.loc[:, ['lat', 'lon']] 94 | dist = geographical_distance(loc_coord, to_rad=True).values 95 | return gaussian_kernel(dist, theta=theta) 96 | -------------------------------------------------------------------------------- /lib/nn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Graph-Machine-Learning-Group/sgp/ae1abd55e15cf21d5ac0c5289a508751ccb3d589/lib/nn/__init__.py -------------------------------------------------------------------------------- /lib/nn/encoders/__init__.py: -------------------------------------------------------------------------------- 1 | from .dyn_gesn_encoder import GESNEncoder 2 | from .sgp_encoder import SGPEncoder 3 | from .sgp_spatial_encoder import SGPSpatialEncoder 4 | from .sgp_temporal_encoder import SGPTemporalEncoder 5 | -------------------------------------------------------------------------------- /lib/nn/encoders/dyn_gesn_encoder.py: -------------------------------------------------------------------------------- 1 | from einops import rearrange 2 | from torch import nn 3 | from torch_geometric.utils import add_self_loops 4 | from torch_sparse import SparseTensor 5 | from tsl.ops.connectivity import normalize 6 | from tsl.utils.parser_utils import ArgParser, str_to_bool 7 | 8 | from lib.nn.reservoir.graph_reservoir import GraphESN 9 | 10 | 11 | class GESNEncoder(nn.Module): 12 | def __init__(self, 13 | input_size, 14 | reservoir_size, 15 | reservoir_layers, 16 | leaking_rate, 17 | spectral_radius, 18 | density, 19 | input_scaling, 20 | alpha_decay, 21 | reservoir_activation='tanh' 22 | ): 23 | super(GESNEncoder, self).__init__() 24 | self.reservoir = GraphESN(input_size=input_size, 25 | hidden_size=reservoir_size, 26 | input_scaling=input_scaling, 27 | num_layers=reservoir_layers, 28 | leaking_rate=leaking_rate, 29 | spectral_radius=spectral_radius, 30 | density=density, 31 | activation=reservoir_activation, 32 | alpha_decay=alpha_decay) 33 | 34 | def forward(self, x, edge_index, edge_weight): 35 | # x : [t n f] 36 | x = rearrange(x, 't n f -> 1 t n f') 37 | edge_index, edge_weight = add_self_loops(edge_index, edge_weight) 38 | if not isinstance(edge_index, SparseTensor): 39 | _, edge_weight = normalize(edge_index, edge_weight, dim=1) 40 | col, row = edge_index 41 | edge_index = SparseTensor(row=row, col=col, value=edge_weight, 42 | sparse_sizes=(x.size(-2), x.size(-2))) 43 | x, _ = self.reservoir(x, edge_index) 44 | return x[0] 45 | 46 | @staticmethod 47 | def add_model_specific_args(parser: ArgParser): 48 | parser.opt_list('--reservoir-size', type=int, default=32, tunable=True, 49 | options=[16, 32, 64, 128, 256]) 50 | parser.opt_list('--reservoir-layers', type=int, default=1, tunable=True, 51 | options=[1, 2, 3]) 52 | parser.opt_list('--spectral-radius', type=float, default=0.9, 53 | tunable=True, options=[0.7, 0.8, 0.9]) 54 | parser.opt_list('--leaking-rate', type=float, default=0.9, tunable=True, 55 | options=[0.7, 0.8, 0.9]) 56 | parser.opt_list('--density', type=float, default=0.7, tunable=True, 57 | options=[0.7, 0.8, 0.9]) 58 | parser.opt_list('--input-scaling', type=float, default=1., tunable=True, 59 | options=[1., 1.5, 2.]) 60 | parser.add_argument('--reservoir-activation', type=str, default='tanh') 61 | parser.opt_list('--alpha-decay', type=str_to_bool, nargs='?', 62 | const=True, default=False) 63 | return parser 64 | -------------------------------------------------------------------------------- /lib/nn/encoders/sgp_encoder.py: -------------------------------------------------------------------------------- 1 | from einops import rearrange 2 | from torch import nn 3 | from tsl.utils.parser_utils import ArgParser, str_to_bool 4 | 5 | from lib.nn.encoders.sgp_spatial_encoder import SGPSpatialEncoder 6 | from lib.nn.reservoir import Reservoir 7 | 8 | 9 | class SGPEncoder(nn.Module): 10 | def __init__(self, 11 | input_size, 12 | reservoir_size, 13 | reservoir_layers, 14 | leaking_rate, 15 | spectral_radius, 16 | density, 17 | input_scaling, 18 | receptive_field, 19 | bidirectional, 20 | alpha_decay, 21 | global_attr, 22 | add_self_loops=False, 23 | undirected=False, 24 | reservoir_activation='tanh' 25 | ): 26 | super(SGPEncoder, self).__init__() 27 | self.reservoir = Reservoir(input_size=input_size, 28 | hidden_size=reservoir_size, 29 | input_scaling=input_scaling, 30 | num_layers=reservoir_layers, 31 | leaking_rate=leaking_rate, 32 | spectral_radius=spectral_radius, 33 | density=density, 34 | activation=reservoir_activation, 35 | alpha_decay=alpha_decay) 36 | 37 | self.sgp_encoder = SGPSpatialEncoder( 38 | receptive_field=receptive_field, 39 | bidirectional=bidirectional, 40 | undirected=undirected, 41 | add_self_loops=add_self_loops, 42 | global_attr=global_attr 43 | ) 44 | 45 | def forward(self, x, edge_index, edge_weight): 46 | # x : [t n f] 47 | x = rearrange(x, 't n f -> 1 t n f') 48 | x = self.reservoir(x) 49 | x = x[0] 50 | x = self.sgp_encoder(x, edge_index, edge_weight) 51 | return x 52 | 53 | @staticmethod 54 | def add_model_specific_args(parser: ArgParser): 55 | parser.opt_list('--reservoir-size', type=int, default=32, tunable=True, 56 | options=[16, 32, 64, 128, 256]) 57 | parser.opt_list('--reservoir-layers', type=int, default=1, tunable=True, 58 | options=[1, 2, 3]) 59 | parser.opt_list('--receptive-field', type=int, default=1, tunable=True, 60 | options=[1, 2, 3]) 61 | parser.opt_list('--spectral-radius', type=float, default=0.9, 62 | tunable=True, options=[0.7, 0.8, 0.9]) 63 | parser.opt_list('--leaking-rate', type=float, default=0.9, tunable=True, 64 | options=[0.7, 0.8, 0.9]) 65 | parser.opt_list('--density', type=float, default=0.7, tunable=True, 66 | options=[0.7, 0.8, 0.9]) 67 | parser.opt_list('--input-scaling', type=float, default=1., tunable=True, 68 | options=[1., 1.5, 2.]) 69 | parser.opt_list('--bidirectional', type=str_to_bool, nargs='?', 70 | const=True, default=False) 71 | parser.opt_list('--undirected', type=str_to_bool, nargs='?', const=True, 72 | default=False) 73 | parser.opt_list('--add-self-loops', type=str_to_bool, nargs='?', 74 | const=True, default=False) 75 | parser.opt_list('--alpha-decay', type=str_to_bool, nargs='?', 76 | const=True, default=False) 77 | parser.opt_list('--global-attr', type=str_to_bool, nargs='?', 78 | const=True, default=False) 79 | parser.add_argument('--reservoir-activation', type=str, default='tanh') 80 | return parser 81 | -------------------------------------------------------------------------------- /lib/nn/encoders/sgp_spatial_encoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | from tsl.utils.parser_utils import ArgParser, str_to_bool 4 | 5 | from lib.sgp_preprocessing import sgp_spatial_embedding 6 | 7 | 8 | class SGPSpatialEncoder(nn.Module): 9 | def __init__(self, 10 | receptive_field, 11 | bidirectional, 12 | undirected, 13 | global_attr, 14 | add_self_loops=False): 15 | super(SGPSpatialEncoder, self).__init__() 16 | self.receptive_field = receptive_field 17 | self.bidirectional = bidirectional 18 | self.undirected = undirected 19 | self.add_self_loops = add_self_loops 20 | self.global_attr = global_attr 21 | 22 | def forward(self, x, edge_index, edge_weight): 23 | num_nodes = x.size(-2) 24 | out = sgp_spatial_embedding(x, 25 | num_nodes=num_nodes, 26 | edge_index=edge_index, 27 | edge_weight=edge_weight, 28 | k=self.receptive_field, 29 | bidirectional=self.bidirectional, 30 | undirected=self.undirected, 31 | add_self_loops=self.add_self_loops) 32 | if self.global_attr: 33 | g = torch.ones_like(x) * x.mean(-2, keepdim=True) 34 | out.append(g) 35 | return torch.cat(out, -1) 36 | 37 | @staticmethod 38 | def add_model_specific_args(parser: ArgParser): 39 | parser.opt_list('--receptive-field', type=int, default=1, tunable=True, 40 | options=[1, 2, 3]) 41 | parser.opt_list('--bidirectional', type=str_to_bool, nargs='?', 42 | const=True, default=False) 43 | parser.opt_list('--undirected', type=str_to_bool, nargs='?', const=True, 44 | default=False) 45 | parser.opt_list('--add-self-loops', type=str_to_bool, nargs='?', 46 | const=True, default=False) 47 | parser.opt_list('--global-attr', type=str_to_bool, nargs='?', 48 | const=True, default=False) 49 | return parser 50 | -------------------------------------------------------------------------------- /lib/nn/encoders/sgp_temporal_encoder.py: -------------------------------------------------------------------------------- 1 | from einops import rearrange 2 | from torch import nn, Tensor 3 | from tsl.utils.parser_utils import ArgParser, str_to_bool 4 | 5 | from lib.nn.reservoir import Reservoir 6 | 7 | 8 | class SGPTemporalEncoder(nn.Module): 9 | def __init__(self, input_size, 10 | reservoir_size=32, 11 | reservoir_layers=1, 12 | leaking_rate=0.9, 13 | spectral_radius=0.9, 14 | density=0.7, 15 | input_scaling=1., 16 | alpha_decay=False, 17 | reservoir_activation='tanh'): 18 | super(SGPTemporalEncoder, self).__init__() 19 | self.reservoir = Reservoir(input_size=input_size, 20 | hidden_size=reservoir_size, 21 | input_scaling=input_scaling, 22 | num_layers=reservoir_layers, 23 | leaking_rate=leaking_rate, 24 | spectral_radius=spectral_radius, 25 | density=density, 26 | activation=reservoir_activation, 27 | alpha_decay=alpha_decay) 28 | 29 | def forward(self, x: Tensor, *args, **kwargs): 30 | # x : [t n f] 31 | x = rearrange(x, 't n f -> 1 t n f') 32 | x = self.reservoir(x) 33 | x = x[0] 34 | return x 35 | 36 | @staticmethod 37 | def add_model_specific_args(parser: ArgParser): 38 | parser.opt_list('--reservoir-size', type=int, default=32, tunable=True, 39 | options=[16, 32, 64, 128, 256]) 40 | parser.opt_list('--reservoir-layers', type=int, default=1, tunable=True, 41 | options=[1, 2, 3]) 42 | parser.opt_list('--spectral-radius', type=float, default=0.9, 43 | tunable=True, options=[0.7, 0.8, 0.9]) 44 | parser.opt_list('--leaking-rate', type=float, default=0.9, tunable=True, 45 | options=[0.7, 0.8, 0.9]) 46 | parser.opt_list('--density', type=float, default=0.7, tunable=True, 47 | options=[0.7, 0.8, 0.9]) 48 | parser.opt_list('--input-scaling', type=float, default=1., tunable=True, 49 | options=[1., 1.5, 2.]) 50 | parser.opt_list('--alpha-decay', type=str_to_bool, nargs='?', 51 | const=True, default=False) 52 | parser.add_argument('--reservoir-activation', type=str, default='tanh') 53 | # for sgp spatial preprocessing 54 | parser.opt_list('--receptive-field', type=int, default=1, tunable=True, 55 | options=[1, 2, 3]) 56 | parser.opt_list('--bidirectional', type=str_to_bool, nargs='?', 57 | const=True, default=False) 58 | parser.opt_list('--undirected', type=str_to_bool, nargs='?', const=True, 59 | default=False) 60 | parser.opt_list('--add-self-loops', type=str_to_bool, nargs='?', 61 | const=True, default=False) 62 | parser.opt_list('--global-attr', type=str_to_bool, nargs='?', 63 | const=True, default=False) 64 | return parser 65 | -------------------------------------------------------------------------------- /lib/nn/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .esn_model import ESNModel 2 | from .gated_gn_model import GatedGraphNetworkMLPModel, \ 3 | GatedGraphNetworkConvModel 4 | from .gwnet_model import GraphWaveNetModel 5 | from .sgp_model import SGPModel, OnlineSGPModel 6 | from .sgp_online import SGPOnlineModel 7 | -------------------------------------------------------------------------------- /lib/nn/models/esn_model.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from tsl.nn.blocks.decoders import LinearReadout 3 | from tsl.nn.utils.utils import maybe_cat_exog 4 | from tsl.utils.parser_utils import ArgParser 5 | 6 | from lib.nn.reservoir import Reservoir 7 | 8 | 9 | class ESNModel(nn.Module): 10 | def __init__(self, 11 | input_size, 12 | hidden_size, 13 | output_size, 14 | exog_size, 15 | rec_layers, 16 | horizon, 17 | activation='tanh', 18 | spectral_radius=0.9, 19 | leaking_rate=0.9, 20 | density=0.7): 21 | super(ESNModel, self).__init__() 22 | 23 | self.reservoir = Reservoir(input_size=input_size + exog_size, 24 | hidden_size=hidden_size, 25 | num_layers=rec_layers, 26 | leaking_rate=leaking_rate, 27 | spectral_radius=spectral_radius, 28 | density=density, 29 | activation=activation) 30 | 31 | self.readout = LinearReadout( 32 | input_size=hidden_size * rec_layers, 33 | output_size=output_size, 34 | horizon=horizon, 35 | ) 36 | 37 | def forward(self, x, u=None, **kwargs): 38 | """""" 39 | # x: [batches steps nodes features] 40 | # u: [batches steps (nodes) features] 41 | x = maybe_cat_exog(x, u) 42 | 43 | x = self.reservoir(x, return_last_state=True) 44 | 45 | return self.readout(x) 46 | 47 | @staticmethod 48 | def add_model_specific_args(parser: ArgParser): 49 | parser.opt_list('--hidden-size', type=int, default=32, tunable=True, 50 | options=[16, 32, 64, 128, 256]) 51 | parser.opt_list('--rec-layers', type=int, default=1, tunable=True, 52 | options=[1, 2, 3]) 53 | parser.opt_list('--spectral-radius', type=float, default=0.9, 54 | tunable=True, options=[0.7, 0.8, 0.9]) 55 | parser.opt_list('--leaking-rate', type=float, default=0.9, tunable=True, 56 | options=[0.7, 0.8, 0.9]) 57 | parser.opt_list('--density', type=float, default=0.7, tunable=True, 58 | options=[0.7, 0.8, 0.9]) 59 | return parser 60 | -------------------------------------------------------------------------------- /lib/nn/models/gwnet_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from einops import repeat 3 | from torch.nn import functional as F 4 | from tsl.nn.models.stgn import GraphWaveNetModel as GWN 5 | 6 | 7 | class GraphWaveNetModel(GWN): 8 | 9 | def get_learned_adj(self, node_index=None): 10 | logits = F.relu(self.source_embeddings(token_index=node_index) @ 11 | self.target_embeddings(token_index=node_index).T) 12 | adj = torch.softmax(logits, dim=1) 13 | return adj 14 | 15 | def forward(self, x, edge_index, edge_weight=None, u=None, node_index=None, 16 | **kwargs): 17 | """""" 18 | # x: [batches, steps, nodes, channels] 19 | 20 | if u is not None: 21 | if u.dim() == 3: 22 | u = repeat(u, 'b s c -> b s n c', n=x.size(-2)) 23 | x = torch.cat([x, u], -1) 24 | 25 | if self.receptive_field > x.size(1): 26 | # pad temporal dimension 27 | x = F.pad(x, (0, 0, 0, 0, self.receptive_field - x.size(1), 0)) 28 | 29 | if len(self.dense_sconvs): 30 | adj_z = self.get_learned_adj(node_index) 31 | 32 | x = self.input_encoder(x) 33 | 34 | out = torch.zeros(1, x.size(1), 1, 1, device=x.device) 35 | for i, (tconv, sconv, skip_conn, norm) in enumerate( 36 | zip(self.tconvs, self.sconvs, self.skip_connections, 37 | self.norms)): 38 | res = x 39 | # temporal conv 40 | x = tconv(x) 41 | # residual connection -> out 42 | out = skip_conn(x) + out[:, -x.size(1):] 43 | # spatial conv 44 | xs = sconv(x, edge_index, edge_weight) 45 | if len(self.dense_sconvs): 46 | x = xs + self.dense_sconvs[i](x, adj_z) 47 | else: 48 | x = xs 49 | x = self.dropout(x) 50 | # residual connection -> next layer 51 | x = x + res[:, -x.size(1):] 52 | x = norm(x) 53 | 54 | return self.readout(out) 55 | -------------------------------------------------------------------------------- /lib/nn/models/sgp_online.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from einops import rearrange 3 | from torch import nn 4 | from tsl.nn.base import StaticGraphEmbedding 5 | from tsl.nn.blocks.decoders import MLPDecoder 6 | from tsl.nn.blocks.encoders import MLP 7 | from tsl.utils.parser_utils import ArgParser 8 | 9 | from lib.sgp_preprocessing import sgp_spatial_embedding 10 | 11 | 12 | class SGPOnlineModel(nn.Module): 13 | def __init__(self, 14 | input_size, 15 | n_nodes, 16 | hidden_size, 17 | output_size, 18 | n_layers, 19 | window, 20 | horizon, 21 | k=2, 22 | bidirectional=True, 23 | dropout=0., 24 | activation='relu'): 25 | super(SGPOnlineModel, self).__init__() 26 | 27 | n_features = input_size * window 28 | input_size = n_features * (1 + k * (int(bidirectional) + 1)) 29 | self.k = k 30 | self.bidirectional = bidirectional 31 | 32 | self.mlp = MLP( 33 | input_size=input_size, 34 | n_layers=n_layers, 35 | hidden_size=hidden_size, 36 | activation=activation, 37 | dropout=dropout 38 | ) 39 | 40 | self.node_emb = StaticGraphEmbedding( 41 | n_tokens=n_nodes, 42 | emb_size=hidden_size 43 | ) 44 | 45 | self.readout = MLPDecoder( 46 | input_size=hidden_size, 47 | hidden_size=hidden_size, 48 | output_size=output_size, 49 | horizon=horizon, 50 | ) 51 | 52 | def forward(self, x, edge_index, edge_weight, **kwargs): 53 | """""" 54 | # x: [batches steps nodes features] 55 | x = rearrange(x, 'b t n f -> b n (t f)') 56 | 57 | x = sgp_spatial_embedding(x, num_nodes=x.size(1), 58 | edge_index=edge_index, 59 | edge_weight=edge_weight, 60 | k=self.k, bidirectional=self.bidirectional) 61 | x = torch.cat(x, -1) 62 | 63 | x = self.mlp(x) + self.node_emb() 64 | 65 | return self.readout(x) 66 | 67 | @staticmethod 68 | def add_model_specific_args(parser: ArgParser): 69 | parser.opt_list('--hidden-size', type=int, default=32, tunable=True, 70 | options=[16, 32, 64, 128, 256]) 71 | parser.opt_list('--n-layers', type=int, default=1, tunable=True, 72 | options=[1, 2, 3]) 73 | parser.opt_list('--dropout', type=float, default=0., tunable=True, 74 | options=[0., 0.2, 0.3]) 75 | return parser 76 | -------------------------------------------------------------------------------- /lib/nn/reservoir/__init__.py: -------------------------------------------------------------------------------- 1 | from .graph_reservoir import GraphESN, GESNLayer 2 | from .reservoir import Reservoir, ReservoirLayer 3 | -------------------------------------------------------------------------------- /lib/predictors/__init__.py: -------------------------------------------------------------------------------- 1 | from .profiling_predictor import ProfilingPredictor 2 | from .subgraph_predictor import SubgraphPredictor 3 | -------------------------------------------------------------------------------- /lib/predictors/profiling_predictor.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from tsl.predictors import Predictor 4 | 5 | 6 | class ProfilingPredictor(Predictor): 7 | step_start = None 8 | 9 | def on_after_batch_transfer(self, batch, dataloader_idx: int): 10 | self.step_start = time.time() 11 | return batch 12 | 13 | def on_after_backward(self) -> None: 14 | if self.step_start is not None: 15 | elapsed = time.time() - self.step_start 16 | self.loggers[0].log_metric('backward_time', elapsed) 17 | self.log('backward_time', elapsed, on_step=False, 18 | on_epoch=True, batch_size=1) 19 | 20 | def on_train_batch_end(self, outputs, batch, batch_idx, unused=0) -> None: 21 | if self.step_start is not None: 22 | elapsed = time.time() - self.step_start 23 | self.loggers[0].log_metric('train_step_time', elapsed) 24 | self.log('train_step_time', elapsed, on_step=False, 25 | on_epoch=True, batch_size=1) 26 | self.step_start = None 27 | 28 | def on_validation_batch_end(self, outputs, batch, 29 | batch_idx: int, dataloader_idx: int) -> None: 30 | if self.step_start is not None: 31 | elapsed = time.time() - self.step_start 32 | self.log('val_step_time', elapsed, on_step=True, 33 | on_epoch=True, batch_size=1) 34 | self.step_start = None 35 | 36 | def on_test_batch_end(self, outputs, batch, 37 | batch_idx: int, dataloader_idx: int) -> None: 38 | if self.step_start is not None: 39 | elapsed = time.time() - self.step_start 40 | self.log('test_step_time', elapsed, on_step=True, 41 | on_epoch=True, batch_size=1) 42 | self.step_start = None 43 | -------------------------------------------------------------------------------- /lib/predictors/subgraph_predictor.py: -------------------------------------------------------------------------------- 1 | from tsl.predictors import Predictor 2 | 3 | 4 | class SubgraphPredictor(Predictor): 5 | 6 | def training_step(self, batch, batch_idx): 7 | y = y_loss = batch.y 8 | mask = batch.get('mask') 9 | 10 | # Compute predictions and compute loss 11 | y_hat_loss = self.predict_batch(batch, preprocess=False, 12 | postprocess=not self.scale_target) 13 | 14 | if 'target_nodes' in batch: 15 | y_hat_loss = y_hat_loss[..., batch.target_nodes, :] 16 | 17 | y_hat = y_hat_loss.detach() 18 | 19 | # Scale target and output, eventually 20 | if self.scale_target: 21 | y_loss = batch.transform['y'].transform(y) 22 | y_hat = batch.transform['y'].inverse_transform(y_hat) 23 | 24 | # Compute loss 25 | loss = self.loss_fn(y_hat_loss, y_loss, mask) 26 | 27 | # Logging 28 | self.train_metrics.update(y_hat, y, mask) 29 | self.log_metrics(self.train_metrics, batch_size=batch.batch_size) 30 | self.log_loss('train', loss, batch_size=batch.batch_size) 31 | return loss 32 | 33 | def validation_step(self, batch, batch_idx): 34 | 35 | y = y_loss = batch.y 36 | mask = batch.get('mask') 37 | 38 | # Compute predictions 39 | y_hat_loss = self.predict_batch(batch, preprocess=False, 40 | postprocess=not self.scale_target) 41 | 42 | if 'target_nodes' in batch: 43 | y_hat_loss = y_hat_loss[..., batch.target_nodes, :] 44 | 45 | y_hat = y_hat_loss.detach() 46 | 47 | # Scale target and output, eventually 48 | if self.scale_target: 49 | y_loss = batch.transform['y'].transform(y) 50 | y_hat = batch.transform['y'].inverse_transform(y_hat) 51 | 52 | # Compute loss 53 | val_loss = self.loss_fn(y_hat_loss, y_loss, mask) 54 | 55 | # Logging 56 | self.val_metrics.update(y_hat, y, mask) 57 | self.log_metrics(self.val_metrics, batch_size=batch.batch_size) 58 | self.log_loss('val', val_loss, batch_size=batch.batch_size) 59 | return val_loss 60 | 61 | def test_step(self, batch, batch_idx): 62 | 63 | # Compute outputs and rescale 64 | y_hat = self.predict_batch(batch, preprocess=False, postprocess=True) 65 | 66 | if 'target_nodes' in batch: 67 | y_hat = y_hat[..., batch.target_nodes, :] 68 | 69 | y, mask = batch.y, batch.get('mask') 70 | test_loss = self.loss_fn(y_hat, y, mask) 71 | 72 | # Logging 73 | self.test_metrics.update(y_hat.detach(), y, mask) 74 | self.log_metrics(self.test_metrics, batch_size=batch.batch_size) 75 | self.log_loss('test', test_loss, batch_size=batch.batch_size) 76 | return test_loss 77 | -------------------------------------------------------------------------------- /lib/utils.py: -------------------------------------------------------------------------------- 1 | from time import time 2 | 3 | import torch 4 | from torch import Tensor 5 | from torch.nn import functional as F 6 | from tsl import logger 7 | from tsl.utils.python_utils import ensure_list 8 | 9 | 10 | def encode_dataset( 11 | dataset, 12 | encoder_class, 13 | encoder_kwargs, 14 | encode_exogenous=True, 15 | keep_raw=False, 16 | save_path=None 17 | ): 18 | # if preprocess_exogenous is True, preprocess all exogenous 19 | if isinstance(encode_exogenous, bool): 20 | preprocess_exogenous = dataset.exogenous.keys() \ 21 | if encode_exogenous else [] 22 | preprocess_exogenous = ensure_list(preprocess_exogenous) 23 | 24 | x, _ = dataset.get_tensors(['data'] + preprocess_exogenous, 25 | preprocess=True, cat_dim=-1) 26 | 27 | encoder = encoder_class(**encoder_kwargs) 28 | 29 | start = time() 30 | encoded_x = encoder(x, edge_index=dataset.edge_index, 31 | edge_weight=dataset.edge_weight) 32 | elapsed = int(time() - start) 33 | 34 | if save_path is not None: 35 | torch.save(encoded_x, save_path) 36 | 37 | logger.info( 38 | f"Dataset encoded in {elapsed // 60}:{elapsed % 60:02d} minutes.") 39 | 40 | dataset.add_exogenous('encoded_x', encoded_x, add_to_input_map=False) 41 | 42 | input_map = {'x': ['encoded_x']} 43 | u = ([] if encode_exogenous else ['u']) + (['data'] if keep_raw else []) 44 | if len(u): 45 | input_map['u'] = u 46 | dataset.set_input_map(input_map) 47 | return dataset 48 | 49 | 50 | def self_normalizing_activation(x: Tensor, r: float = 1.0): 51 | return r * F.normalize(x, p=2, dim=-1) 52 | -------------------------------------------------------------------------------- /scalable.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Graph-Machine-Learning-Group/sgp/ae1abd55e15cf21d5ac0c5289a508751ccb3d589/scalable.png -------------------------------------------------------------------------------- /sgp_paper.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Graph-Machine-Learning-Group/sgp/ae1abd55e15cf21d5ac0c5289a508751ccb3d589/sgp_paper.pdf -------------------------------------------------------------------------------- /sgp_poster.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Graph-Machine-Learning-Group/sgp/ae1abd55e15cf21d5ac0c5289a508751ccb3d589/sgp_poster.pdf -------------------------------------------------------------------------------- /tsl/__init__.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import tsl.global_scope 4 | from tsl.global_scope import * 5 | 6 | data = LazyLoader('data', globals(), 'tsl.data') 7 | datasets = LazyLoader('datasets', globals(), 'tsl.datasets') 8 | nn = LazyLoader('nn', globals(), 'tsl.nn') 9 | predictors = LazyLoader('predictors', globals(), 'tsl.predictors') 10 | imputers = LazyLoader('imputers', globals(), 'tsl.imputers') 11 | 12 | __version__ = '0.1.0' 13 | 14 | epsilon = 5e-8 15 | config = Config() 16 | 17 | config_file = os.path.join(config.curr_dir, 'tsl_config.yaml') 18 | if os.path.exists(config_file): 19 | config.load_config_file(config_file) 20 | 21 | __all__ = [ 22 | '__version__', 23 | 'config', 24 | 'epsilon', 25 | 'logger', 26 | 'tsl', 27 | 'data', 28 | 'datasets', 29 | 'nn', 30 | 'predictors', 31 | 'imputers' 32 | ] 33 | -------------------------------------------------------------------------------- /tsl/data/__init__.py: -------------------------------------------------------------------------------- 1 | from .batch import Batch, static_graph_collate 2 | from .batch_map import BatchMap, BatchMapItem 3 | from .data import Data, DataView 4 | from .datamodule import * 5 | from .spatiotemporal_dataset import SpatioTemporalDataset 6 | from .imputation_stds import ImputationDataset 7 | 8 | data_classes = ['Data', 'DataView', 'Batch'] 9 | dataset_classes = ['SpatioTemporalDataset', 'ImputationDataset'] 10 | 11 | __all__ = [ 12 | *data_classes, 13 | *dataset_classes 14 | ] 15 | -------------------------------------------------------------------------------- /tsl/data/batch.py: -------------------------------------------------------------------------------- 1 | from typing import Callable 2 | from typing import (Optional, Any, Union, List, Mapping) 3 | 4 | import torch 5 | from torch import Tensor 6 | from torch.utils.data.dataloader import default_collate 7 | 8 | from .data import Data 9 | from .preprocessing import ScalerModule 10 | 11 | 12 | def _collate_scaler_modules(batch: List[Mapping[str, Any]]): 13 | transform = batch[0] 14 | for k, v in transform.items(): 15 | # scaler params are supposed to be the same for all elements in 16 | # minibatch, just add a fake, 1-sized, batch dimension 17 | scaler = ScalerModule() 18 | if v.bias is not None: 19 | scaler.bias = transform[k].bias[None] 20 | if v.scale is not None: 21 | scaler.scale = transform[k].scale[None] 22 | if v.trend is not None: 23 | trend = torch.stack([b[k].trend for b in batch], 0) 24 | scaler.trend = trend 25 | transform[k] = scaler 26 | return transform 27 | 28 | 29 | def static_graph_collate(batch: List[Data], cls: Optional[type] = None) -> Data: 30 | # collate subroutine 31 | def _collate(items: List[Union[Tensor, Mapping[str, Any]]], key: str, 32 | pattern: str): 33 | if key == 'transform': 34 | return _collate_scaler_modules(items), None 35 | # if key.startswith('edge_'): 36 | # return items[0] 37 | if pattern is not None: 38 | if 's' in pattern: 39 | return default_collate(items), 'b ' + pattern 40 | return items[0], pattern 41 | return default_collate(items), None 42 | 43 | # collate all sample-wise elements 44 | elem = batch[0] 45 | if cls is None: 46 | cls = elem.__class__ 47 | out = cls() 48 | out = out.stores_as(elem) 49 | for k in elem.keys: 50 | pattern = elem.pattern.get(k) 51 | out[k], pattern = _collate([b[k] for b in batch], k, pattern) 52 | if pattern is not None: 53 | out.pattern[k] = pattern 54 | 55 | out.__dict__['batch_size'] = len(batch) 56 | return out 57 | 58 | 59 | class Batch(Data): 60 | _collate_fn: Callable = static_graph_collate 61 | 62 | @classmethod 63 | def from_data_list(cls, data_list: List[Data]): 64 | r"""Constructs a :class:`~tsl.data.Batch` object from a Python list of 65 | :class:`~tsl.data.Data`, representing temporal signals on static 66 | graphs.""" 67 | 68 | batch = cls._collate_fn(data_list, cls) 69 | 70 | batch.__dict__['batch_size'] = len(data_list) 71 | 72 | return batch 73 | -------------------------------------------------------------------------------- /tsl/data/batch_map.py: -------------------------------------------------------------------------------- 1 | from typing import Iterator 2 | from typing import (Optional, Union, List, Tuple, Mapping) 3 | 4 | from tsl.utils.python_utils import ensure_list 5 | from .utils import SynchMode 6 | 7 | 8 | class BatchMapItem: 9 | def __init__(self, keys: Union[List, str], 10 | synch_mode: SynchMode = SynchMode.WINDOW, 11 | preprocess: bool = True, 12 | cat_dim: Optional[int] = -1, 13 | n_channels: Optional[int] = None): 14 | super(BatchMapItem, self).__init__() 15 | self.keys = ensure_list(keys) 16 | assert isinstance(synch_mode, SynchMode) 17 | self.synch_mode = synch_mode 18 | self.preprocess = preprocess 19 | if len(self.keys) > 1: 20 | assert cat_dim is not None, \ 21 | '"cat_dim" cannot be None with multiple keys.' 22 | self.cat_dim = cat_dim 23 | self.n_channels = n_channels 24 | 25 | def __repr__(self): 26 | return "([{}], {})".format(', '.join(self.keys), self.synch_mode.name) 27 | 28 | def kwargs(self): 29 | return self.__dict__ 30 | 31 | 32 | class BatchMap(Mapping): 33 | 34 | def __init__(self, **kwargs): 35 | super().__init__() 36 | for k, v in kwargs.items(): 37 | self[k] = v 38 | 39 | def __setitem__(self, key: str, value: Union[BatchMapItem, Tuple, Mapping]): 40 | # cast item 41 | if isinstance(value, BatchMapItem): 42 | pass 43 | elif isinstance(value, Tuple): 44 | value = BatchMapItem(*value) 45 | elif isinstance(value, (List, str)): 46 | value = BatchMapItem(value) 47 | elif isinstance(value, Mapping): 48 | value = BatchMapItem(**value) 49 | else: 50 | raise TypeError('Invalid type for InputMap item "{}"' 51 | .format(type(value))) 52 | self.__dict__[key] = value 53 | 54 | def __getitem__(self, k): 55 | return self.__dict__[k] 56 | 57 | def __len__(self) -> int: 58 | return len(self.__dict__) 59 | 60 | def __iter__(self) -> Iterator: 61 | return iter(self.__dict__) 62 | 63 | def __repr__(self): 64 | s = ['({}={}, {})'.format(key, value.keys, value.synch_mode.name) 65 | for key, value in self.items()] 66 | return "{}[{}]".format(self.__class__.__name__, ', '.join(s)) 67 | 68 | def update(self, **kwargs): 69 | for k, v in kwargs.items(): 70 | self[k] = v 71 | 72 | def by_synch_mode(self, synch_mode: SynchMode): 73 | return {k: v for k, v in self.items() if 74 | v.synch_mode is synch_mode} 75 | -------------------------------------------------------------------------------- /tsl/data/datamodule/__init__.py: -------------------------------------------------------------------------------- 1 | from .spatiotemporal_datamodule import SpatioTemporalDataModule 2 | from .splitters import * 3 | from . import splitters 4 | 5 | datamodule_classes = ['SpatioTemporalDataModule'] 6 | splitter_classes = splitters.__all__ 7 | 8 | __all__ = datamodule_classes + splitter_classes 9 | 10 | -------------------------------------------------------------------------------- /tsl/data/imputation_stds.py: -------------------------------------------------------------------------------- 1 | from typing import Union, Optional, Mapping, Tuple 2 | 3 | import numpy as np 4 | 5 | from tsl.data import SpatioTemporalDataset, BatchMap, BatchMapItem 6 | from tsl.data.preprocessing import Scaler 7 | from tsl.typing import (TensArray, TemporalIndex) 8 | 9 | 10 | class ImputationDataset(SpatioTemporalDataset): 11 | 12 | def __init__(self, data: TensArray, 13 | index: Optional[TemporalIndex] = None, 14 | training_mask: Optional[TensArray] = None, 15 | eval_mask: Optional[TensArray] = None, 16 | connectivity: Optional[ 17 | Union[TensArray, Tuple[TensArray]]] = None, 18 | exogenous: Optional[Mapping[str, TensArray]] = None, 19 | attributes: Optional[Mapping[str, TensArray]] = None, 20 | input_map: Optional[Union[Mapping, BatchMap]] = None, 21 | trend: Optional[TensArray] = None, 22 | scalers: Optional[Mapping[str, Scaler]] = None, 23 | window: int = 24, 24 | stride: int = 1, 25 | window_lag: int = 1, 26 | horizon_lag: int = 1, 27 | precision: Union[int, str] = 32, 28 | name: Optional[str] = None): 29 | if training_mask is None: 30 | training_mask = np.isnan(data) 31 | if exogenous is None: 32 | exogenous = dict() 33 | if eval_mask is not None: 34 | exogenous['eval_mask'] = eval_mask 35 | if input_map is not None: 36 | input_map['eval_mask'] = BatchMapItem('eval_mask', preprocess=False) 37 | super(ImputationDataset, self).__init__(data, 38 | index=index, 39 | mask=training_mask, 40 | connectivity=connectivity, 41 | exogenous=exogenous, 42 | attributes=attributes, 43 | input_map=input_map, 44 | trend=trend, 45 | scalers=scalers, 46 | window=window, 47 | horizon=window, 48 | delay=-window, 49 | stride=stride, 50 | window_lag=window_lag, 51 | horizon_lag=horizon_lag, 52 | precision=precision, 53 | name=name) 54 | 55 | @staticmethod 56 | def add_argparse_args(parser, **kwargs): 57 | parser.add_argument('--window', type=int, default=24) 58 | parser.add_argument('--stride', type=int, default=1) 59 | parser.add_argument('--window-lag', type=int, default=1) 60 | parser.add_argument('--horizon-lag', type=int, default=1) 61 | return parser 62 | -------------------------------------------------------------------------------- /tsl/data/loader/__init__.py: -------------------------------------------------------------------------------- 1 | from .dataloader import StaticGraphLoader 2 | -------------------------------------------------------------------------------- /tsl/data/loader/dataloader.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | from torch.utils import data 4 | 5 | from ..spatiotemporal_dataset import SpatioTemporalDataset 6 | from ..batch import Batch 7 | 8 | 9 | class StaticGraphLoader(data.DataLoader): 10 | 11 | def __init__(self, dataset: SpatioTemporalDataset, 12 | batch_size: Optional[int] = 1, 13 | shuffle: bool = False, 14 | num_workers: int = 0, 15 | **kwargs): 16 | if 'collate_fn' in kwargs: 17 | del kwargs['collate_fn'] 18 | super().__init__(dataset, 19 | shuffle=shuffle, 20 | batch_size=batch_size, 21 | num_workers=num_workers, 22 | collate_fn=Batch.from_data_list, 23 | **kwargs) 24 | -------------------------------------------------------------------------------- /tsl/data/preprocessing/__init__.py: -------------------------------------------------------------------------------- 1 | from . import scalers 2 | from .scalers import * 3 | 4 | scaler_classes = scalers.__all__ 5 | 6 | __all__ = scaler_classes 7 | -------------------------------------------------------------------------------- /tsl/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | # Interfaces 2 | from .prototypes import Dataset, TabularDataset, PandasDataset 3 | from .prototypes import classes as prototype_classes 4 | # Datasets 5 | from .metr_la import MetrLA 6 | from .pems_bay import PemsBay 7 | from .mts_benchmarks import ( 8 | ElectricityBenchmark, 9 | TrafficBenchmark, 10 | SolarBenchmark, 11 | ExchangeBenchmark 12 | ) 13 | 14 | dataset_classes = [ 15 | 'MetrLA', 16 | 'PemsBay', 17 | 'ElectricityBenchmark', 18 | 'TrafficBenchmark', 19 | 'SolarBenchmark', 20 | 'ExchangeBenchmark' 21 | ] 22 | 23 | __all__ = prototype_classes + dataset_classes 24 | -------------------------------------------------------------------------------- /tsl/datasets/prototypes/__init__.py: -------------------------------------------------------------------------------- 1 | from .dataset import Dataset 2 | from .tabular_dataset import TabularDataset 3 | from .pd_dataset import PandasDataset 4 | 5 | __all__ = [ 6 | 'Dataset', 7 | 'TabularDataset', 8 | 'PandasDataset' 9 | ] 10 | 11 | classes = __all__ 12 | -------------------------------------------------------------------------------- /tsl/global_scope/__init__.py: -------------------------------------------------------------------------------- 1 | from .config import Config 2 | from .logger import logger 3 | from .lazy_loader import LazyLoader 4 | -------------------------------------------------------------------------------- /tsl/global_scope/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Mapping, Optional 3 | 4 | 5 | class Config(dict): 6 | """Manage the package configuration from a single object. 7 | 8 | With a :obj:`Config` object you can edit settings within the tsl scope, like 9 | directory in which you store configuration files for experiments 10 | (:obj:`config_dir`), logs (:obj:`log_dir`), and data (:obj:`data_dir`). 11 | """ 12 | 13 | def __init__(self, **kwargs): 14 | super(Config, self).__init__() 15 | # configure paths for config files and logs 16 | self.config_dir = kwargs.pop('config_dir', 'config') 17 | self.log_dir = kwargs.pop('log_dir', 'log') 18 | # set 'data_dir' as directory for data loading and downloading 19 | # defaults to '{tsl_path}/.storage' 20 | default_storage = os.path.join(self.root_dir, '.storage') 21 | self.data_dir = kwargs.pop('data_dir', default_storage) 22 | self.update(**kwargs) 23 | 24 | def __setitem__(self, key: str, value): 25 | # when adding a directory, transform it to an absolute path (if it is 26 | # not already) considering the path relative to the current directory 27 | if key.endswith('_dir') and value is not None: 28 | if not os.path.isabs(value): 29 | value = os.path.join(self.curr_dir, value) 30 | super(Config, self).__setitem__(key, value) 31 | 32 | def __setattr__(self, key, value): 33 | self[key] = value 34 | 35 | def __getattr__(self, item): 36 | return self[item] 37 | 38 | def __delattr__(self, item): 39 | del self[item] 40 | 41 | def __repr__(self): 42 | type_name = type(self).__name__ 43 | arg_strings = [] 44 | for name, value in sorted(self.items()): 45 | arg_strings.append('%s=%r' % (name, value)) 46 | return '%s(%s)' % (type_name, ', '.join(arg_strings)) 47 | 48 | @property 49 | def root_dir(self): 50 | """Path to tsl installation.""" 51 | return os.path.dirname(os.path.dirname(os.path.realpath(__file__))) 52 | 53 | @property 54 | def curr_dir(self): 55 | """System current directory.""" 56 | return os.getcwd() 57 | 58 | def update(self, mapping: Optional[Mapping] = None, **kwargs) -> None: 59 | mapping = dict(mapping or {}, **kwargs) 60 | for k, v in mapping.items(): 61 | self[k] = v 62 | 63 | def disable_logging(self): 64 | from .logger import logger 65 | logger.disabled = True 66 | 67 | def load_config_file(self, filename: str): 68 | """Load a configuration from a json or yaml file.""" 69 | with open(filename, 'r') as fp: 70 | if filename.endswith('.json'): 71 | import json 72 | data = json.load(fp) 73 | elif filename.endswith('.yaml') or filename.endswith('.yml'): 74 | import yaml 75 | data = yaml.load(fp, Loader=yaml.FullLoader) 76 | else: 77 | raise RuntimeError('Config file format not supported.') 78 | self.update(data) 79 | return self 80 | 81 | @classmethod 82 | def from_config_file(cls, filename: str): 83 | """Create new configuration from a json or yaml file.""" 84 | config = cls() 85 | config.load_config_file(filename) 86 | return config 87 | -------------------------------------------------------------------------------- /tsl/global_scope/lazy_loader.py: -------------------------------------------------------------------------------- 1 | from importlib import import_module 2 | from types import ModuleType 3 | 4 | 5 | # https://github.com/tensorflow/tensorflow/blob/master/tensorflow/ 6 | # python/util/lazy_loader.py 7 | class LazyLoader(ModuleType): 8 | def __init__(self, local_name, parent_module_globals, name): 9 | self._local_name = local_name 10 | self._parent_module_globals = parent_module_globals 11 | super(LazyLoader, self).__init__(name) 12 | 13 | def _load(self): 14 | module = import_module(self.__name__) 15 | self._parent_module_globals[self._local_name] = module 16 | self.__dict__.update(module.__dict__) 17 | return module 18 | 19 | def __getattr__(self, item): 20 | module = self._load() 21 | return getattr(module, item) 22 | 23 | def __dir__(self): 24 | module = self._load() 25 | return dir(module) 26 | -------------------------------------------------------------------------------- /tsl/global_scope/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import logging.config 3 | 4 | DEFAULT_LOGGING = { 5 | 'version': 1, 6 | 'disable_existing_loggers': False, 7 | 'formatters': { 8 | 'standard': { 9 | 'format': '%(asctime)s [%(levelname)s]: %(message)s' 10 | }, 11 | }, 12 | 'handlers': { 13 | 'default': { 14 | 'level': 'INFO', 15 | 'formatter': 'standard', 16 | 'class': 'logging.StreamHandler', 17 | 'stream': 'ext://sys.stdout', 18 | }, 19 | }, 20 | 'loggers': { 21 | 'log': { 22 | 'handlers': ['default'], 23 | 'level': 'INFO', 24 | 'propagate': True 25 | } 26 | } 27 | } 28 | 29 | logging.config.dictConfig(DEFAULT_LOGGING) 30 | logger = logging.getLogger('log') 31 | -------------------------------------------------------------------------------- /tsl/imputers/__init__.py: -------------------------------------------------------------------------------- 1 | from .imputer import Imputer 2 | 3 | imputer_classes = ['Imputer'] 4 | 5 | __all__ = imputer_classes 6 | -------------------------------------------------------------------------------- /tsl/nn/__init__.py: -------------------------------------------------------------------------------- 1 | from . import layers 2 | from . import models 3 | -------------------------------------------------------------------------------- /tsl/nn/base/__init__.py: -------------------------------------------------------------------------------- 1 | from . import attention 2 | from .embedding import StaticGraphEmbedding 3 | from .graph_conv import GraphConv 4 | from .temporal_conv import TemporalConv2d, GatedTemporalConv2d 5 | from .attention import * 6 | 7 | __all__ = [ 8 | 'attention', 9 | 'GraphConv', 10 | 'TemporalConv2d', 11 | 'GatedTemporalConv2d', 12 | *attention.classes 13 | ] 14 | 15 | classes = __all__[1:] 16 | -------------------------------------------------------------------------------- /tsl/nn/base/attention/__init__.py: -------------------------------------------------------------------------------- 1 | from .attention import AttentionEncoder, MultiHeadAttention 2 | from .linear_attention import CausalLinearAttention 3 | 4 | __all__ = [ 5 | 'AttentionEncoder', 6 | 'MultiHeadAttention', 7 | 'CausalLinearAttention' 8 | ] 9 | 10 | classes = __all__ 11 | -------------------------------------------------------------------------------- /tsl/nn/base/dense.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from tsl.nn.utils import utils 4 | 5 | 6 | class Dense(nn.Module): 7 | r""" 8 | A simple fully-connected layer. 9 | 10 | Args: 11 | input_size (int): Size of the input. 12 | output_size (int): Size of the output. 13 | activation (str, optional): Activation function. 14 | dropout (float, optional): Dropout rate. 15 | bias (bool, optional): Whether to use a bias. 16 | """ 17 | def __init__(self, input_size, output_size, activation='linear', dropout=0., bias=True): 18 | super(Dense, self).__init__() 19 | self.layer = nn.Sequential( 20 | nn.Linear(input_size, output_size, bias=bias), 21 | utils.get_layer_activation(activation)(), 22 | nn.Dropout(dropout) if dropout > 0. else nn.Identity() 23 | ) 24 | 25 | def forward(self, x): 26 | return self.layer(x) 27 | -------------------------------------------------------------------------------- /tsl/nn/base/graph_conv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from torch import nn 4 | from torch.nn import Parameter 5 | from torch_geometric.nn.conv import MessagePassing 6 | from torch_geometric.typing import Adj, OptTensor 7 | 8 | from tsl.ops.connectivity import normalize 9 | 10 | 11 | class GraphConv(MessagePassing): 12 | r"""A simple graph convolutional operator where the message function is a simple linear projection and aggregation 13 | a simple average. In other terms: 14 | 15 | .. math:: 16 | \mathbf{X}^{\prime} = \mathbf{\hat{D}}^{-1} \mathbf{A} \mathbf{X} \boldsymbol{\Theta} 17 | 18 | Args: 19 | input_size (int): Size of the input features. 20 | output_size (int): Size of each output features. 21 | add_self_loops (bool, optional): If set to :obj:`True`, will add 22 | self-loops to the input graph. (default: :obj:`False`) 23 | bias (bool, optional): If set to :obj:`False`, the layer will not learn 24 | an additive bias. (default: :obj:`True`) 25 | **kwargs (optional): Additional arguments of 26 | :class:`torch_geometric.nn.conv.MessagePassing`. 27 | """ 28 | 29 | def __init__(self, input_size: int, output_size: int, bias: bool = True, root_weight: bool = True, **kwargs): 30 | super(GraphConv, self).__init__(aggr="add", node_dim=-2) 31 | super().__init__(**kwargs) 32 | 33 | self.in_channels = input_size 34 | self.out_channels = output_size 35 | 36 | self.lin = nn.Linear(input_size, output_size, bias=False) 37 | 38 | if root_weight: 39 | self.root_lin = nn.Linear(input_size, output_size, bias=False) 40 | else: 41 | self.register_parameter('root_lin', None) 42 | 43 | if bias: 44 | self.bias = Parameter(torch.Tensor(output_size)) 45 | else: 46 | self.register_parameter('bias', None) 47 | 48 | self.reset_parameters() 49 | 50 | def reset_parameters(self): 51 | self.lin.reset_parameters() 52 | if self.root_lin is not None: 53 | self.root_lin.reset_parameters() 54 | if self.bias is not None: 55 | torch.nn.init.zeros_(self.bias) 56 | 57 | def forward(self, x: Tensor, edge_index: Adj, 58 | edge_weight: OptTensor = None) -> Tensor: 59 | """""" 60 | n = x.size(-2) 61 | out = self.lin(x) 62 | 63 | _, edge_weight = normalize(edge_index, edge_weight, dim=1, num_nodes=n) 64 | out = self.propagate(edge_index, x=out, edge_weight=edge_weight) 65 | 66 | if self.root_lin is not None: 67 | out += self.root_lin(x) 68 | 69 | if self.bias is not None: 70 | out += self.bias 71 | 72 | return out 73 | 74 | def message(self, x_j: Tensor, edge_weight) -> Tensor: 75 | return edge_weight.view(-1, 1) * x_j -------------------------------------------------------------------------------- /tsl/nn/base/temporal_conv.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | import torch.nn as nn 4 | from tsl.nn.functional import gated_tanh 5 | 6 | from einops import rearrange 7 | 8 | 9 | class TemporalConv2d(nn.Module): 10 | r""" 11 | Learns a standard temporal convolutional filter. 12 | 13 | Args: 14 | input_channels (int): Input size. 15 | output_channels (int): Output size. 16 | kernel_size (int): Size of the convolution kernel. 17 | dilation (int, optional): Spacing between kernel elements. 18 | stride (int, optional): Stride of the convolution. 19 | bias (bool, optional): Whether to add a learnable bias to the output of the convolution. 20 | padding (int or tuple, optional): Padding of the input. Used only of `causal_pad` is `False`. 21 | causal_pad (bool, optional): Whether to pad the input as to preserve causality. 22 | weight_norm (bool, optional): Wheter to apply weight normalization to the parameters of the filter. 23 | """ 24 | def __init__(self, 25 | input_channels, 26 | output_channels, 27 | kernel_size, 28 | dilation=1, 29 | stride=1, 30 | bias=True, 31 | padding=0, 32 | causal_pad=True, 33 | weight_norm=False, 34 | channel_last=False): 35 | super().__init__() 36 | if causal_pad: 37 | self.padding = ((kernel_size - 1) * dilation, 0, 0, 0) 38 | else: 39 | self.padding = padding 40 | self.pad_layer = nn.ZeroPad2d(self.padding) 41 | # we use conv2d here to accommodate multiple input sequences 42 | self.conv = nn.Conv2d(input_channels, output_channels, (1, kernel_size), 43 | stride=stride, padding=(0, 0), dilation=(1, dilation), bias=bias) 44 | if weight_norm: 45 | self.conv = nn.utils.weight_norm(self.conv) 46 | self.channel_last = channel_last 47 | 48 | def forward(self, x): 49 | """""" 50 | if self.channel_last: 51 | x = rearrange(x, 'b s n c -> b c n s') 52 | # batch, channels, nodes, steps 53 | x = self.pad_layer(x) 54 | x = self.conv(x) 55 | if self.channel_last: 56 | x = rearrange(x, 'b c n s -> b s n c') 57 | return x 58 | 59 | 60 | class GatedTemporalConv2d(TemporalConv2d): 61 | def __init__(self, 62 | input_channels, 63 | output_channels, 64 | kernel_size, 65 | dilation=1, 66 | stride=1, 67 | bias=True, 68 | padding=0, 69 | causal_pad=True, 70 | weight_norm=False, 71 | channel_last=False): 72 | super(GatedTemporalConv2d, self).__init__(input_channels=input_channels, 73 | output_channels=2 * output_channels, 74 | kernel_size=kernel_size, 75 | dilation=dilation, 76 | stride=stride, 77 | bias=bias, 78 | padding=padding, 79 | causal_pad=causal_pad, 80 | weight_norm=weight_norm, 81 | channel_last=channel_last) 82 | 83 | def forward(self, x): 84 | """""" 85 | # batch, channels, nodes, steps 86 | x = super(GatedTemporalConv2d, self).forward(x) 87 | dim = -1 if self.channel_last else 1 88 | return gated_tanh(x, dim=dim) 89 | -------------------------------------------------------------------------------- /tsl/nn/blocks/__init__.py: -------------------------------------------------------------------------------- 1 | from . import decoders, encoders 2 | 3 | encoder_classes = encoders.classes 4 | decoder_classes = decoders.classes 5 | 6 | __all__ = [ 7 | 'encoders', 8 | 'decoders', 9 | *encoder_classes, 10 | *decoder_classes 11 | ] 12 | 13 | classes = __all__ 14 | -------------------------------------------------------------------------------- /tsl/nn/blocks/decoders/__init__.py: -------------------------------------------------------------------------------- 1 | from .att_pool import AttPool 2 | from .gcn_decoder import GCNDecoder 3 | from .linear_readout import LinearReadout 4 | from .mlp_decoder import MLPDecoder 5 | from .multi_step_mlp_decoder import MultiHorizonMLPDecoder 6 | 7 | __all__ = [ 8 | 'AttPool', 9 | 'GCNDecoder', 10 | 'LinearReadout', 11 | 'MLPDecoder', 12 | 'MultiHorizonMLPDecoder' 13 | ] 14 | 15 | classes = __all__ 16 | -------------------------------------------------------------------------------- /tsl/nn/blocks/decoders/att_pool.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch.nn import functional as F 3 | 4 | 5 | class AttPool(nn.Module): 6 | r""" 7 | Pool representations along a dimension with learned softmax scores. 8 | 9 | Args: 10 | input_size (int): Input size. 11 | dim (int): Dimension on which to apply the attention pooling. 12 | """ 13 | def __init__(self, input_size, dim): 14 | super(AttPool, self).__init__() 15 | self.lin = nn.Linear(input_size, 1) 16 | self.dim = dim 17 | 18 | def forward(self, x): 19 | scores = F.softmax(self.lin(x), dim=self.dim) 20 | x = (scores * x).sum(dim=self.dim, keepdim=True) 21 | return x -------------------------------------------------------------------------------- /tsl/nn/blocks/decoders/gcn_decoder.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from torch.nn import functional as F 3 | 4 | from tsl.nn.base.graph_conv import GraphConv 5 | from tsl.nn.blocks.decoders.mlp_decoder import MLPDecoder 6 | from tsl.nn.utils import utils 7 | 8 | 9 | class GCNDecoder(nn.Module): 10 | r""" 11 | GCN decoder for multi-step forecasting. 12 | Applies multiple graph convolutional layers followed by a feed-forward layer amd a linear readout. 13 | 14 | If the input representation has a temporal dimension, this model will simply take as input the representation 15 | corresponding to the last step. 16 | 17 | Args: 18 | input_size (int): Input size. 19 | hidden_size (int): Hidden size. 20 | output_size (int): Output size. 21 | horizon (int): Output steps. 22 | n_layers (int, optional): Number of layers in the decoder. (default: 1) 23 | activation (str, optional): Activation function to use. 24 | dropout (float, optional): Dropout probability applied in the hidden layers. 25 | """ 26 | def __init__(self, 27 | input_size, 28 | hidden_size, 29 | output_size, 30 | horizon=1, 31 | n_layers=1, 32 | activation='relu', 33 | dropout=0.): 34 | super(GCNDecoder, self).__init__() 35 | graph_convs = [] 36 | for l in range(n_layers): 37 | graph_convs.append( 38 | GraphConv(input_size=input_size if l == 0 else hidden_size, 39 | output_size=hidden_size) 40 | ) 41 | self.convs = nn.ModuleList(graph_convs) 42 | self.activation = utils.get_functional_activation(activation) 43 | self.dropout = nn.Dropout(dropout) 44 | self.readout = MLPDecoder(input_size=hidden_size, 45 | hidden_size=hidden_size, 46 | output_size=output_size, 47 | activation=activation, 48 | horizon=horizon) 49 | 50 | def forward(self, h, edge_index, edge_weight=None): 51 | """""" 52 | # h: [batches (steps) nodes features] 53 | if h.dim() == 4: 54 | # take last step representation 55 | h = h[:, -1] 56 | for conv in self.convs: 57 | h = self.dropout(self.activation(conv(h, edge_index, edge_weight))) 58 | return self.readout(h) 59 | -------------------------------------------------------------------------------- /tsl/nn/blocks/decoders/linear_readout.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from einops.layers.torch import Rearrange 3 | 4 | 5 | class LinearReadout(nn.Module): 6 | r""" 7 | Simple linear readout for multi-step forecasting. 8 | 9 | If the input representation has a temporal dimension, this model will simply take the representation corresponding 10 | to the last step. 11 | 12 | Args: 13 | input_size (int): Input size. 14 | output_size (int): Output size. 15 | horizon(int): Number of steps predict. 16 | """ 17 | def __init__(self, 18 | input_size, 19 | output_size, 20 | horizon=1): 21 | super(LinearReadout, self).__init__() 22 | 23 | self.readout = nn.Sequential( 24 | nn.Linear(input_size, output_size * horizon), 25 | Rearrange('b n (h c) -> b h n c', c=output_size, h=horizon) 26 | ) 27 | 28 | def forward(self, h): 29 | # h: [batches (steps) nodes features] 30 | if h.dim() == 4: 31 | # take last step representation 32 | h = h[:, -1] 33 | return self.readout(h) 34 | -------------------------------------------------------------------------------- /tsl/nn/blocks/decoders/mlp_decoder.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from tsl.nn.blocks.encoders.mlp import MLP 4 | from einops.layers.torch import Rearrange 5 | 6 | from einops import rearrange 7 | 8 | 9 | class MLPDecoder(nn.Module): 10 | r""" 11 | Simple MLP decoder for multi-step forecasting. 12 | 13 | If the input representation has a temporal dimension, this model will take the flatten representations corresponding 14 | to the last `receptive_field` time steps. 15 | 16 | Args: 17 | input_size (int): Input size. 18 | hidden_size (int): Hidden size. 19 | output_size (int): Output size. 20 | horizon (int): Output steps. 21 | n_layers (int, optional): Number of layers in the decoder. (default: 1) 22 | receptive_field (int, optional): Number of steps to consider for decoding. (default: 1) 23 | activation (str, optional): Activation function to use. 24 | dropout (float, optional): Dropout probability applied in the hidden layers. 25 | """ 26 | def __init__(self, 27 | input_size, 28 | hidden_size, 29 | output_size, 30 | horizon=1, 31 | n_layers=1, 32 | receptive_field=1, 33 | activation='relu', 34 | dropout=0.): 35 | super(MLPDecoder, self).__init__() 36 | 37 | self.receptive_field = receptive_field 38 | self.readout = nn.Sequential( 39 | MLP(input_size=receptive_field * input_size, 40 | hidden_size=hidden_size, 41 | output_size=output_size * horizon, 42 | n_layers=n_layers, 43 | dropout=dropout, 44 | activation=activation), 45 | Rearrange('b n (h c) -> b h n c', c=output_size, h=horizon) 46 | ) 47 | 48 | def forward(self, h): 49 | # h: [batches (steps) nodes features] 50 | if h.dim() == 4: 51 | # take last step representation 52 | h = rearrange(h[:, -self.receptive_field:], 'b s n c -> b n (s c)') 53 | else: 54 | assert self.receptive_field == 1 55 | return self.readout(h) 56 | -------------------------------------------------------------------------------- /tsl/nn/blocks/decoders/multi_step_mlp_decoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from tsl.nn.blocks.encoders.mlp import MLP 5 | from einops import rearrange, repeat 6 | 7 | 8 | class MultiHorizonMLPDecoder(nn.Module): 9 | r""" 10 | Decoder for multistep forecasting based on 11 | 12 | Wen et al., "A Multi-Horizon Quantile Recurrent Forecaster", 2018. 13 | 14 | It requires exogenous variables synched with the forecasting horizon. 15 | 16 | Args: 17 | input_size (int): Size of the input. 18 | exog_size (int): Size of the horizon exogenous variables. 19 | hidden_size (int): Number of hidden units. 20 | context_size (int): Number of units used to condition the forecasting of each step. 21 | output_size (int): Output channels. 22 | n_layers (int): Number of hidden layers. 23 | horizon (int): Forecasting horizon. 24 | activation (str, optional): Activation function. 25 | dropout (float, optional): Dropout probability. 26 | """ 27 | def __init__(self, 28 | input_size, 29 | exog_size, 30 | hidden_size, 31 | context_size, 32 | output_size, 33 | n_layers, 34 | horizon, 35 | activation='relu', 36 | dropout=0.): 37 | super(MultiHorizonMLPDecoder, self).__init__() 38 | global_d_out = horizon * context_size + context_size 39 | self.d_context = context_size 40 | self.horizon = horizon 41 | self.global_mlp = MLP(input_size=input_size, hidden_size=hidden_size, output_size=global_d_out, 42 | n_layers=n_layers, activation=activation, dropout=dropout) 43 | self.local_mlp = MLP(input_size=exog_size + 2 * context_size, hidden_size=hidden_size, output_size=output_size, 44 | n_layers=n_layers, activation=activation, dropout=dropout) 45 | 46 | def forward(self, x: torch.Tensor, u: torch.Tensor): 47 | """""" 48 | # x: [batch, (steps), nodes, channels] 49 | # u: [batch, horizon, (nodes), channels] 50 | # out: [batch, steps, nodes, channels] 51 | if x.dim() == 4: 52 | x = x[:, -1] 53 | n = x.size(1) 54 | 55 | if u.dim() == 3: 56 | u = repeat(u, 'b h c -> b h n c', n=n) 57 | u = rearrange(u, 'b h n c -> b n h c') 58 | 59 | out = self.global_mlp(x) 60 | global_context, contexts = torch.split(out, [self.d_context, self.horizon * self.d_context], -1) 61 | global_context = repeat(global_context, 'b n c -> b n h c', h=self.horizon) 62 | contexts = rearrange(contexts, 'b n (h c) -> b n h c', c=self.d_context, h=self.horizon) 63 | x_dec = torch.cat([contexts, global_context, u], -1) 64 | x_dec = self.local_mlp(x_dec) 65 | 66 | return rearrange(x_dec, 'b n h c -> b h n c') 67 | 68 | def reset_parameters(self): 69 | self.global_mlp.reset_parameters() 70 | self.local_mlp.reset_parameters() 71 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/__init__.py: -------------------------------------------------------------------------------- 1 | from .conditional import ConditionalBlock, ConditionalTCNBlock 2 | from .dcrnn import DCRNNCell, DCRNN 3 | from .dense_dcrnn import DenseDCRNNCell, DenseDCRNN 4 | from .gcgru import GraphConvGRUCell, GraphConvGRU 5 | from .gclstm import GraphConvLSTMCell, GraphConvLSTM 6 | from .mlp import MLP, ResidualMLP 7 | from .rnn import RNN 8 | from .stcn import SpatioTemporalConvNet 9 | from .tcn import TemporalConvNet 10 | from .transformer import (TransformerLayer, SpatioTemporalTransformerLayer, 11 | Transformer) 12 | 13 | general_classes = [ 14 | 'ConditionalBlock', 15 | 'ConditionalTCNBlock', 16 | 'MLP', 17 | 'ResidualMLP', 18 | 'RNN', 19 | ] 20 | 21 | cell_classes = [ 22 | 'DCRNNCell', 23 | 'DenseDCRNNCell', 24 | 'GraphConvGRUCell', 25 | 'GraphConvLSTMCell' 26 | ] 27 | 28 | grnn_classes = [ 29 | 'DCRNN', 30 | 'DenseDCRNN', 31 | 'GraphConvGRU', 32 | 'GraphConvLSTM' 33 | ] 34 | 35 | conv_classes = [ 36 | 'TemporalConvNet', 37 | 'SpatioTemporalConvNet' 38 | ] 39 | 40 | transformer_classes = [ 41 | 'TransformerLayer', 42 | 'SpatioTemporalTransformerLayer', 43 | 'Transformer' 44 | ] 45 | 46 | classes = [ 47 | *general_classes, 48 | *cell_classes, 49 | *grnn_classes, 50 | *conv_classes, 51 | *transformer_classes 52 | ] 53 | 54 | __all__ = classes 55 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/dcrnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from tsl.nn.layers.graph_convs.diff_conv import DiffConv 4 | from tsl.nn.blocks.encoders.gcrnn import _GraphGRUCell, _GraphRNN 5 | 6 | 7 | class DCRNNCell(_GraphGRUCell): 8 | """ 9 | Diffusion Convolutional Recurrent Cell. 10 | 11 | Args: 12 | input_size: Size of the input. 13 | output_size: Number of units in the hidden state. 14 | k: Size of the diffusion kernel. 15 | root_weight: Whether to learn a separate transformation for the central node. 16 | """ 17 | 18 | def __init__(self, input_size, output_size, k=2, root_weight=True): 19 | super(DCRNNCell, self).__init__() 20 | # instantiate gates 21 | self.forget_gate = DiffConv(input_size + output_size, output_size, k=k, 22 | root_weight=root_weight) 23 | self.update_gate = DiffConv(input_size + output_size, output_size, k=k, 24 | root_weight=root_weight) 25 | self.candidate_gate = DiffConv(input_size + output_size, output_size, 26 | k=k, root_weight=root_weight) 27 | 28 | 29 | class DCRNN(_GraphRNN): 30 | r"""Diffusion Convolutional Recurrent Network, from the paper 31 | `"Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting" `_. 32 | 33 | Args: 34 | input_size: Size of the input. 35 | hidden_size: Number of units in the hidden state. 36 | n_layers: Number of layers. 37 | k: Size of the diffusion kernel. 38 | root_weight: Whether to learn a separate transformation for the central node. 39 | """ 40 | _n_states = 1 41 | 42 | def __init__(self, 43 | input_size, 44 | hidden_size, 45 | n_layers=1, 46 | k=2, 47 | root_weight=True): 48 | super(DCRNN, self).__init__() 49 | self.input_size = input_size 50 | self.hidden_size = hidden_size 51 | self.n_layers = n_layers 52 | self.k = k 53 | self.rnn_cells = torch.nn.ModuleList() 54 | for i in range(self.n_layers): 55 | self.rnn_cells.append(DCRNNCell( 56 | input_size=self.input_size if i == 0 else self.hidden_size, 57 | output_size=self.hidden_size, k=self.k, 58 | root_weight=root_weight)) 59 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/dense_dcrnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from tsl.nn.layers.graph_convs.dense_spatial_conv import SpatialConvOrderK 4 | from tsl.nn.blocks.encoders.gcrnn import _GraphGRUCell, _GraphRNN 5 | 6 | 7 | class DenseDCRNNCell(_GraphGRUCell): 8 | r""" 9 | Diffusion Convolutional Recurrent Cell. 10 | 11 | Args: 12 | input_size: Size of the input. 13 | output_size: Number of units in the hidden state. 14 | k: Size of the diffusion kernel. 15 | root_weight (bool): Whether to learn a separate transformation for the 16 | central node. 17 | """ 18 | 19 | def __init__(self, input_size, output_size, k=2, root_weight=False): 20 | super(DenseDCRNNCell, self).__init__() 21 | # instantiate gates 22 | self.forget_gate = SpatialConvOrderK( 23 | input_size=input_size + output_size, 24 | output_size=output_size, 25 | support_len=2, 26 | order=k, 27 | include_self=root_weight, 28 | channel_last=True) 29 | self.update_gate = SpatialConvOrderK( 30 | input_size=input_size + output_size, 31 | output_size=output_size, 32 | support_len=2, 33 | order=k, 34 | include_self=root_weight, 35 | channel_last=True) 36 | self.candidate_gate = SpatialConvOrderK( 37 | input_size=input_size + output_size, 38 | output_size=output_size, 39 | support_len=2, 40 | order=k, 41 | include_self=root_weight, 42 | channel_last=True) 43 | 44 | 45 | class DenseDCRNN(_GraphRNN): 46 | r""" 47 | Diffusion Convolutional Recurrent Network. 48 | 49 | From Li et al., ”Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting”, ICLR 2018 50 | 51 | Args: 52 | input_size: Size of the input. 53 | hidden_size: Number of units in the hidden state. 54 | n_layers: Number of layers. 55 | k: Size of the diffusion kernel. 56 | root_weight: Whether to learn a separate transformation for the central node. 57 | """ 58 | _n_states = 1 59 | 60 | def __init__(self, 61 | input_size, 62 | hidden_size, 63 | n_layers=1, 64 | k=2, 65 | root_weight=False): 66 | super(DenseDCRNN, self).__init__() 67 | self.input_size = input_size 68 | self.hidden_size = hidden_size 69 | self.n_layers = n_layers 70 | self.k = k 71 | self.rnn_cells = torch.nn.ModuleList() 72 | for i in range(self.n_layers): 73 | self.rnn_cells.append(DenseDCRNNCell( 74 | input_size=self.input_size if i == 0 else self.hidden_size, 75 | output_size=self.hidden_size, k=self.k, 76 | root_weight=root_weight)) 77 | 78 | def forward(self, x, adj, h=None): 79 | support = SpatialConvOrderK.compute_support(adj) 80 | return super(DenseDCRNN, self).forward(x, support, h=h) 81 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/gcgru.py: -------------------------------------------------------------------------------- 1 | from tsl.nn.base import GraphConv 2 | from tsl.nn.blocks.encoders.gcrnn import _GraphGRUCell, _GraphRNN 3 | 4 | from torch import nn 5 | 6 | 7 | class GraphConvGRUCell(_GraphGRUCell): 8 | r""" 9 | Gate Recurrent Unit with `GraphConv` gates. 10 | Loosely based on Seo et al., ”Structured Sequence Modeling with Graph Convolutional Recurrent Networks”, ICONIP 2017 11 | 12 | Args: 13 | input_size: Size of the input. 14 | out_size: Number of units in the hidden state. 15 | root_weight: Whether to learn a separate transformation for the central node. 16 | """ 17 | def __init__(self, in_size, out_size, root_weight=True): 18 | super(GraphConvGRUCell, self).__init__() 19 | # instantiate gates 20 | self.forget_gate = GraphConv(in_size + out_size, out_size, root_weight=root_weight) 21 | self.update_gate = GraphConv(in_size + out_size, out_size, root_weight=root_weight) 22 | self.candidate_gate = GraphConv(in_size + out_size, out_size, root_weight=root_weight) 23 | 24 | 25 | class GraphConvGRU(_GraphRNN): 26 | r""" 27 | GraphConv GRU network. 28 | 29 | Loosely based on Seo et al., ”Structured Sequence Modeling with Graph Convolutional Recurrent Networks”, ICONIP 2017 30 | 31 | Args: 32 | input_size (int): Size of the input. 33 | hidden_size (int): Number of units in the hidden state. 34 | n_layers (int, optional): Number of hidden layers. 35 | root_weight (bool, optional): Whether to learn a separate transformation for the central node. 36 | """ 37 | _n_states = 1 38 | 39 | def __init__(self, 40 | input_size, 41 | hidden_size, 42 | n_layers=1, 43 | root_weight=True): 44 | super(GraphConvGRU, self).__init__() 45 | self.input_size = input_size 46 | self.hidden_size = hidden_size 47 | self.n_layers = n_layers 48 | self.rnn_cells = nn.ModuleList() 49 | for i in range(self.n_layers): 50 | self.rnn_cells.append(GraphConvGRUCell(in_size=self.input_size if i == 0 else self.hidden_size, 51 | out_size=self.hidden_size, 52 | root_weight=root_weight)) 53 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/gclstm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from tsl.nn.base import GraphConv 4 | from tsl.nn.blocks.encoders.gcrnn import _GraphLSTMCell, _GraphRNN 5 | 6 | from torch import nn 7 | 8 | class GraphConvLSTMCell(_GraphLSTMCell): 9 | r""" 10 | LSTM with `GraphConv` gates. 11 | Loosely based on Seo et al., ”Structured Sequence Modeling with Graph Convolutional Recurrent Networks”, ICONIP 2017 12 | 13 | Args: 14 | input_size: Size of the input. 15 | out_size: Number of units in the hidden state. 16 | root_weight: Whether to learn a separate transformation for the central node. 17 | """ 18 | def __init__(self, in_size, out_size, root_weight=True): 19 | super(GraphConvLSTMCell, self).__init__() 20 | # instantiate gates 21 | self.input_gate = GraphConv(in_size + out_size, out_size, root_weight=root_weight) 22 | self.forget_gate = GraphConv(in_size + out_size, out_size, root_weight=root_weight) 23 | self.cell_gate = GraphConv(in_size + out_size, out_size, root_weight=root_weight) 24 | self.output_gate = GraphConv(in_size + out_size, out_size, root_weight=root_weight) 25 | 26 | 27 | class GraphConvLSTM(_GraphRNN): 28 | r""" 29 | GraphConv LSTM network. 30 | 31 | Loosely based on Seo et al., ”Structured Sequence Modeling with Graph Convolutional Recurrent Networks”, ICONIP 2017 32 | 33 | Args: 34 | input_size (int): Size of the input. 35 | hidden_size (int): Number of units in the hidden state. 36 | n_layers (int, optional): Number of hidden layers. 37 | root_weight (bool, optional): Whether to learn a separate transformation for the central node. 38 | """ 39 | _n_states = 2 40 | 41 | def __init__(self, 42 | input_size, 43 | hidden_size, 44 | n_layers=1, 45 | root_weight=True): 46 | super(GraphConvLSTM, self).__init__() 47 | self.input_size = input_size 48 | self.hidden_size = hidden_size 49 | self.n_layers = n_layers 50 | self.rnn_cells = nn.ModuleList() 51 | for i in range(self.n_layers): 52 | self.rnn_cells.append(GraphConvLSTMCell(in_size=self.input_size if i == 0 else self.hidden_size, 53 | out_size=self.hidden_size, 54 | root_weight=root_weight)) -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/gcrnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from einops import rearrange 4 | 5 | 6 | class _GraphGRUCell(torch.nn.Module): 7 | r""" 8 | Base class for implementing `GraphGRU` cells. 9 | """ 10 | def forward(self, x, h, *args, **kwargs): 11 | """""" 12 | # x: [batch, nodes, channels] 13 | # h: [batch, nodes, channels] 14 | x_gates = torch.cat([x, h], dim=-1) 15 | r = torch.sigmoid(self.forget_gate(x_gates, *args, **kwargs)) 16 | u = torch.sigmoid(self.update_gate(x_gates, *args, **kwargs)) 17 | x_c = torch.cat([x, r * h], dim=-1) 18 | c = torch.tanh(self.candidate_gate(x_c, *args, **kwargs)) 19 | return u * h + (1. - u) * c 20 | 21 | 22 | class _GraphLSTMCell(torch.nn.Module): 23 | r""" 24 | Base class for implementing `GraphLSTM` cells. 25 | """ 26 | def forward(self, x, hs, *args, **kwargs): 27 | """""" 28 | # x: [batch, nodes, channels] 29 | # hs: (h, c) 30 | # h: [batch, nodes, channels] 31 | # c: [batch, nodes, channels] 32 | h, c = hs 33 | x_gates = torch.cat([x, h], dim=-1) 34 | i = torch.sigmoid(self.input_gate(x_gates, *args, **kwargs)) 35 | f = torch.sigmoid(self.forget_gate(x_gates, *args, **kwargs)) 36 | g = torch.tanh(self.cell_gate(x_gates, *args, **kwargs)) 37 | o = torch.sigmoid(self.output_gate(x_gates, *args, **kwargs)) 38 | c_new = f * c + i * g 39 | h_new = o * torch.tan(c) 40 | return (h_new, c_new) 41 | 42 | 43 | class _GraphRNN(torch.nn.Module): 44 | r""" 45 | Base class for GraphRNNs 46 | """ 47 | _n_states = None 48 | hidden_size: int 49 | _cat_states_layers = False 50 | 51 | def _init_states(self, x): 52 | assert 'hidden_size' in self.__dict__, \ 53 | f"Class {self.__class__.__name__} must have the attribute " \ 54 | f"`hidden_size`." 55 | return torch.zeros(size=(self.n_layers, x.shape[0], x.shape[-2], self.hidden_size), device=x.device) 56 | 57 | def single_pass(self, x, h, *args, **kwargs): 58 | # x: [batch, nodes, channels] 59 | # h: [layers, batch, nodes, channels] 60 | h_new = [] 61 | out = x 62 | for i, cell in enumerate(self.rnn_cells): 63 | out = cell(out, h[i], *args, **kwargs) 64 | h_new.append(out) 65 | return torch.stack(h_new) 66 | 67 | def forward(self, x, *args, h=None, **kwargs): 68 | # x: [batch, steps, nodes, channels] 69 | steps = x.size(1) 70 | if h is None: 71 | *h, = self._init_states(x) 72 | if not len(h): 73 | h = h[0] 74 | # temporal conv 75 | out = [] 76 | for step in range(steps): 77 | h = self.single_pass(x[:, step], h, *args, **kwargs) 78 | if not isinstance(h, torch.Tensor): 79 | h_out, _ = h 80 | else: 81 | h_out = h 82 | # append hidden state of the last layer 83 | if self._cat_states_layers: 84 | h_out = rearrange(h_out, 'l b n f -> b n (l f)') 85 | else: 86 | h_out = h_out[-1] 87 | 88 | out.append(h_out) 89 | out = torch.stack(out) 90 | # out: [steps, batch, nodes, channels] 91 | out = rearrange(out, 's b n c -> b s n c') 92 | # h: [l b n c] 93 | return out, h 94 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/input_encoder.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from .mlp import MLP 4 | from .rnn import RNN 5 | from .conditional import ConditionalBlock 6 | from .tcn import TemporalConvNet 7 | 8 | 9 | class InputEncoder(nn.Module): 10 | def __init__(self, 11 | enc_type, 12 | input_size, 13 | exog_size, 14 | output_size, 15 | dropout=0., 16 | activation=None, **kwargs): 17 | super(InputEncoder, self).__init__() 18 | if enc_type == 'mlp': 19 | self.input_encoder = MLP( 20 | input_size=input_size, 21 | exog_size=exog_size, 22 | hidden_size=output_size, 23 | activation=activation, 24 | dropout=dropout, 25 | **kwargs 26 | ) 27 | elif enc_type == 'conditional': 28 | self.input_encoder = ConditionalBlock( 29 | input_size=input_size, 30 | exog_size=exog_size, 31 | output_size=output_size, 32 | dropout=dropout, 33 | activation=activation, 34 | **kwargs 35 | ) 36 | elif enc_type == 'rnn': 37 | assert activation is None 38 | self.input_encoder = RNN( 39 | input_size=input_size, 40 | exog_size=exog_size, 41 | output_size=output_size, 42 | dropout=dropout, 43 | **kwargs 44 | ) 45 | elif enc_type == 'tcn': 46 | self.input_encoder = TemporalConvNet( 47 | input_channels=input_size, 48 | exog_channels=exog_size, 49 | output_channels=output_size, 50 | activation=activation, 51 | dropout=dropout, 52 | **kwargs 53 | ) 54 | else: 55 | raise NotImplementedError(f"Encoder type {enc_type} not implemented.") 56 | 57 | def forward(self, x, u=None): 58 | return self.input_encoder(x, u) 59 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/nri_dcrnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from tsl.nn.base.embedding import StaticGraphEmbedding 5 | from tsl.nn.layers.link_predictor import LinkPredictor 6 | 7 | from tsl.nn.blocks.encoders.dense_dcrnn import DenseDCRNN 8 | 9 | import tsl 10 | 11 | 12 | class DifferentiableBinarySampler(nn.Module): 13 | """ 14 | This module exploits the GumbelMax trick to sample from a Bernoulli distribution in differentiable fashion. 15 | 16 | Adapted from https://github.com/yaringal/ConcreteDropout 17 | """ 18 | def __init__(self): 19 | super(DifferentiableBinarySampler, self).__init__() 20 | 21 | def forward(self, scores, tau): 22 | unif_noise = torch.rand_like(scores) 23 | eps = tsl.epsilon 24 | 25 | logit = torch.log(scores + eps) - torch.log(1 - scores + eps) + \ 26 | torch.log(unif_noise + eps) - torch.log(1 - unif_noise + eps) 27 | 28 | soft_out = torch.sigmoid(logit / tau) 29 | return soft_out 30 | 31 | 32 | class NeuRelInfDCRNN(DenseDCRNN): 33 | r""" 34 | Diffusion Convolutional Recurrent Network with graph learned through neural relational inference. 35 | 36 | Loosely inspired by: 37 | - Kipf et al. "Neural relational inference for interacting systems". ICLR 2018. 38 | - Shang et al. "Discrete graph structure learning for forecasting multiple time series". ICLR 2021. 39 | 40 | Args: 41 | input_size: Size of the input. 42 | hidden_size: Number of units in the hidden state. 43 | n_layers: Number of layers. 44 | k: Size of the diffusion kernel. 45 | root_weight: Whether to learn a separate transformation for the central node. 46 | """ 47 | def __init__(self, 48 | input_size, 49 | hidden_size, 50 | emb_size, 51 | n_nodes, 52 | n_layers=1, 53 | k=2, 54 | root_weight=False): 55 | super(NeuRelInfDCRNN, self).__init__(input_size=input_size, 56 | hidden_size=hidden_size, 57 | n_layers=n_layers, 58 | k=k, 59 | root_weight=root_weight) 60 | 61 | self.node_emb = StaticGraphEmbedding(n_tokens=n_nodes, 62 | emb_size=emb_size) 63 | self.link_predictor = LinkPredictor(emb_size=emb_size, 64 | ff_size=hidden_size, 65 | hidden_size=hidden_size 66 | ) 67 | self.sampler = DifferentiableBinarySampler() 68 | 69 | def forward(self, x, h=None, tau=0.25): 70 | emb = self.node_emb() 71 | adj_p = torch.sigmoid(self.link_predictor(emb)) 72 | adj = self.sampler(adj_p, tau) 73 | return super(NeuRelInfDCRNN, self).forward(x, adj, h) 74 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/rnn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | from torch import nn 4 | from einops import rearrange 5 | 6 | from ...utils.utils import maybe_cat_exog 7 | 8 | class RNN(nn.Module): 9 | r""" 10 | Simple RNN encoder with optional linear readout. 11 | 12 | Args: 13 | input_size (int): Input size. 14 | hidden_size (int): Units in the hidden layers. 15 | exog_size (int, optional): Size of the optional exogenous variables. 16 | output_size (int, optional): Size of the optional readout. 17 | n_layers (int, optional): Number of hidden layers. (default: 1) 18 | cell (str, optional): Type of cell that should be use (options: [`gru`, `lstm`]). (default: `gru`) 19 | dropout (float, optional): Dropout probability. 20 | """ 21 | def __init__(self, 22 | input_size, 23 | hidden_size, 24 | exog_size=None, 25 | output_size=None, 26 | n_layers=1, 27 | dropout=0., 28 | cell='gru'): 29 | super(RNN, self).__init__() 30 | 31 | if cell == 'gru': 32 | cell = nn.GRU 33 | elif cell == 'lstm': 34 | cell = nn.LSTM 35 | else: 36 | raise NotImplementedError(f'"{cell}" cell not implemented.') 37 | 38 | if exog_size is not None: 39 | input_size += exog_size 40 | 41 | self.rnn = cell(input_size=input_size, 42 | hidden_size=hidden_size, 43 | num_layers=n_layers, 44 | dropout=dropout) 45 | 46 | if output_size is not None: 47 | self.readout = nn.Linear(hidden_size, output_size) 48 | else: 49 | self.register_parameter('readout', None) 50 | 51 | def forward(self, x, u=None, return_last_state=False): 52 | """ 53 | 54 | Args: 55 | x (torch.Tensor): Input tensor. 56 | return_last_state: Whether to return only the state corresponding to the last time step. 57 | """ 58 | # x: [batches, steps, nodes, features] 59 | x = maybe_cat_exog(x, u) 60 | b, *_ = x.size() 61 | x = rearrange(x, 'b s n f -> s (b n) f') 62 | x, *_ = self.rnn(x) 63 | # [steps batches * nodes, features] -> [steps batches, nodes, features] 64 | x = rearrange(x, 's (b n) f -> b s n f', b=b) 65 | if return_last_state: 66 | x = x[:, -1] 67 | if self.readout is not None: 68 | return self.readout(x) 69 | return x 70 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/stcn.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | from tsl.nn.layers.graph_convs.diff_conv import DiffConv 4 | from tsl.nn.blocks.encoders.tcn import TemporalConvNet 5 | from tsl.nn.layers.norm.norm import Norm 6 | 7 | from tsl.nn.utils import utils 8 | 9 | 10 | class SpatioTemporalConvNet(nn.Module): 11 | r""" 12 | SpatioTemporalConvolutional encoder with optional linear readout. 13 | Applies several temporal convolutions followed by diffusion convolution over a graph. 14 | 15 | Args: 16 | input_size (int): Input size. 17 | output_size (int): Channels in the output representation. 18 | temporal_kernel_size (int): Size of the temporal convolutional kernel. 19 | spatial_kernel_size (int): Size of the spatial diffusion kernel. 20 | temporal_convs (int, optional): Number of temporal convolutions. (default: 2) 21 | spatial_convs (int, optional): Number of spatial convolutions. (default: 1) 22 | dilation (int): Dilation coefficient of the temporal convolutional kernel. 23 | norm (str, optional): Type of normalization applied to the hidden units. 24 | dropout (float, optional): Dropout probability. 25 | gated (bool, optional): Whether to used the GatedTanH activation function after temporal convolutions. 26 | (default: `False`) 27 | pad (bool, optional): Whether to pad the input sequence to preserve the sequence length. 28 | activation (str, optional): Activation function. (default: `relu`) 29 | """ 30 | def __init__(self, 31 | input_size, 32 | output_size, 33 | temporal_kernel_size, 34 | spatial_kernel_size, 35 | temporal_convs=2, 36 | spatial_convs=1, 37 | dilation=1, 38 | norm='none', 39 | dropout=0., 40 | gated=False, 41 | pad=True, 42 | activation='relu'): 43 | super(SpatioTemporalConvNet, self).__init__() 44 | self.pad = pad 45 | 46 | self.tcn = nn.Sequential( 47 | Norm(norm_type=norm, in_channels=input_size), 48 | TemporalConvNet( 49 | input_channels=input_size, 50 | hidden_channels=output_size, 51 | kernel_size=temporal_kernel_size, 52 | dilation=dilation, 53 | exponential_dilation=True, 54 | n_layers=temporal_convs, 55 | activation=activation, 56 | causal_padding=pad, 57 | dropout=dropout, 58 | gated=gated 59 | )) 60 | 61 | self.skip_conn = nn.Linear(input_size, output_size) 62 | 63 | self.spatial_convs = nn.ModuleList(DiffConv(in_channels=output_size, 64 | out_channels=output_size, 65 | k=spatial_kernel_size) for _ in range(spatial_convs)) 66 | self.spatial_norms = nn.ModuleList(Norm(norm_type=norm, in_channels=output_size) 67 | for _ in range(spatial_convs)) 68 | self.dropout = nn.Dropout(dropout) 69 | self.activation = utils.get_functional_activation(activation) 70 | 71 | def forward(self, x, edge_index, edge_weight=None): 72 | """""" 73 | # temporal conv 74 | x = self.skip_conn(x) + self.tcn(x) 75 | # spatial conv 76 | for filter, norm in zip(self.spatial_convs, self.spatial_norms): 77 | x = x + self.dropout(self.activation(filter(norm(x), edge_index, edge_weight))) 78 | return x 79 | -------------------------------------------------------------------------------- /tsl/nn/blocks/encoders/tcn.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from tsl.nn.base import TemporalConv2d, GatedTemporalConv2d 3 | from tsl.nn.utils import utils 4 | from tsl.nn.utils.utils import maybe_cat_exog 5 | 6 | from einops import rearrange 7 | 8 | 9 | class TemporalConvNet(nn.Module): 10 | r""" 11 | Simple TCN encoder with optional linear readout. 12 | 13 | Args: 14 | input_channels (int): Input size. 15 | hidden_channels (int): Channels in the hidden layers. 16 | kernel_size (int): Size of the convolutional kernel. 17 | dilation (int): Dilation coefficient of the convolutional kernel. 18 | stride (int, optional): Stride of the convolutional kernel. 19 | output_channels (int, optional): Channels of the optional exogenous variables. 20 | output_channels (int, optional): Channels in the output layer. 21 | n_layers (int, optional): Number of hidden layers. (default: 1) 22 | gated (bool, optional): Whether to used the GatedTanH activation function. (default: `False`) 23 | dropout (float, optional): Dropout probability. 24 | activation (str, optional): Activation function. (default: `relu`) 25 | exponential_dilation (bool, optional): Whether to increase exponentially the dilation factor at each layer. 26 | weight_norm (bool, optional): Whether to apply weight normalization to the temporal convolutional filters. 27 | causal_padding (bool, optional): Whether to pad the input sequence to preserve causality. 28 | bias (bool, optional): Whether to add a learnable bias to the output. 29 | channel_last (bool, optional): If `True` input must have layout (b s n c), (b c n s) otherwise. 30 | """ 31 | def __init__(self, 32 | input_channels, 33 | hidden_channels, 34 | kernel_size, 35 | dilation, 36 | stride=1, 37 | exog_channels=None, 38 | output_channels=None, 39 | n_layers=1, 40 | gated=False, 41 | dropout=0., 42 | activation='relu', 43 | exponential_dilation=False, 44 | weight_norm=False, 45 | causal_padding=True, 46 | bias=True, 47 | channel_last=True): 48 | super(TemporalConvNet, self).__init__() 49 | self.channel_last = channel_last 50 | base_conv = TemporalConv2d if not gated else GatedTemporalConv2d 51 | 52 | if exog_channels is not None: 53 | input_channels += exog_channels 54 | 55 | layers = [] 56 | d = dilation 57 | for i in range(n_layers): 58 | if exponential_dilation: 59 | d = dilation ** i 60 | layers.append(base_conv(input_channels=input_channels if i == 0 else hidden_channels, 61 | output_channels=hidden_channels, 62 | kernel_size=kernel_size, 63 | dilation=d, 64 | stride=stride, 65 | causal_pad=causal_padding, 66 | weight_norm=weight_norm, 67 | bias=bias 68 | )) 69 | 70 | self.convs = nn.ModuleList(layers) 71 | self.f = utils.get_functional_activation(activation) if not gated else nn.Identity() 72 | self.dropout = nn.Dropout(dropout) if dropout > 0. else nn.Identity() 73 | 74 | if output_channels is not None: 75 | self.readout = TemporalConv2d(input_channels=hidden_channels, 76 | output_channels=output_channels, 77 | kernel_size=1) 78 | else: 79 | self.register_parameter('readout', None) 80 | 81 | def forward(self, x, u=None): 82 | """""" 83 | if self.channel_last: 84 | x = maybe_cat_exog(x, u, -1) 85 | x = rearrange(x, 'b s n c -> b c n s') 86 | else: 87 | x = maybe_cat_exog(x, u, 1) 88 | 89 | for conv in self.convs: 90 | x = self.dropout(self.f(conv(x))) 91 | if self.readout is not None: 92 | x = self.readout(x) 93 | if self.channel_last: 94 | x = rearrange(x, 'b c n s -> b s n c') 95 | return x 96 | -------------------------------------------------------------------------------- /tsl/nn/layers/__init__.py: -------------------------------------------------------------------------------- 1 | from .link_predictor import LinkPredictor 2 | from .positional_encoding import PositionalEncoding 3 | from . import norm, graph_convs 4 | 5 | __all__ = [ 6 | 'graph_convs', 7 | 'norm', 8 | 'LinkPredictor', 9 | 'PositionalEncoding' 10 | ] 11 | 12 | classes = __all__[2:] 13 | -------------------------------------------------------------------------------- /tsl/nn/layers/graph_convs/__init__.py: -------------------------------------------------------------------------------- 1 | from .dense_spatial_conv import SpatialConv, SpatialConvOrderK 2 | from .diff_conv import DiffConv 3 | from .graph_attention import AttentionScores, MultiHeadGraphAttention, GATLayer 4 | from .gat_conv import GATConv 5 | from .grin_cell import GRIL 6 | from .spatio_temporal_att import SpatioTemporalAtt 7 | from .gated_gn import GatedGraphNetwork 8 | 9 | __all__ = [ 10 | 'SpatialConv', 11 | 'SpatialConvOrderK', 12 | 'DiffConv', 13 | 'MultiHeadGraphAttention', 14 | 'GATConv', 15 | 'GATLayer', 16 | 'GRIL', 17 | 'SpatioTemporalAtt', 18 | 'GatedGraphNetwork' 19 | ] 20 | 21 | classes = __all__ 22 | -------------------------------------------------------------------------------- /tsl/nn/layers/graph_convs/diff_conv.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | import torch 4 | from torch import nn, Tensor 5 | from torch_geometric.nn import MessagePassing 6 | from torch_geometric.typing import Adj, OptTensor 7 | from torch_sparse import SparseTensor, matmul 8 | 9 | from tsl.ops.connectivity import transpose, normalize 10 | 11 | 12 | class DiffConv(MessagePassing): 13 | r"""An implementation of the Diffusion Convolution Layer from `"Diffusion 14 | Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting" 15 | `_. 16 | 17 | Args: 18 | in_channels (int): Number of input features. 19 | out_channels (int): Number of output features. 20 | k (int): Filter size :math:`K`. 21 | bias (bool, optional): If set to :obj:`False`, the layer 22 | will not learn an additive bias (default :obj:`True`). 23 | add_backward (bool): If :obj:`True`, additional :math:`K` filters are 24 | learnt for the transposed connectivity. 25 | (default :obj:`True`) 26 | 27 | """ 28 | 29 | def __init__(self, in_channels, out_channels, k, 30 | root_weight: bool = True, 31 | add_backward: bool = True, 32 | bias: bool=True): 33 | super(DiffConv, self).__init__(aggr="add", node_dim=-2) 34 | self.in_channels = in_channels 35 | self.out_channels = out_channels 36 | self.k = k 37 | 38 | self.root_weight = root_weight 39 | self.add_backward = add_backward 40 | 41 | n_filters = 2 * k if not root_weight else 2 * k + 1 42 | 43 | self.filters = nn.Linear(in_channels * n_filters, out_channels, 44 | bias=bias) 45 | 46 | self._support = None 47 | self.reset_parameters() 48 | 49 | @staticmethod 50 | def compute_support_index(edge_index: Adj, edge_weight: OptTensor = None, 51 | num_nodes: int = None, 52 | add_backward: bool = True) -> List: 53 | norm_edge_index, \ 54 | norm_edge_weight = normalize(edge_index, edge_weight, 55 | dim=1, num_nodes=num_nodes) 56 | # Add backward matrices 57 | if add_backward: 58 | return [(norm_edge_index, norm_edge_weight)] + \ 59 | DiffConv.compute_support_index(transpose(edge_index), 60 | edge_weight=edge_weight, 61 | num_nodes=num_nodes, 62 | add_backward=False) 63 | # Normalize 64 | return [(norm_edge_index, norm_edge_weight)] 65 | 66 | def reset_parameters(self): 67 | self.filters.reset_parameters() 68 | self._support = None 69 | 70 | def message(self, x_j: Tensor, weight: Tensor) -> Tensor: 71 | # x_j: [batch, edges, channels] 72 | return weight.view(-1, 1) * x_j 73 | 74 | def message_and_aggregate(self, adj_t: SparseTensor, x: Tensor) -> Tensor: 75 | # adj_t: SparseTensor [nodes, nodes] 76 | # x: [(batch,) nodes, channels] 77 | return matmul(adj_t, x, reduce=self.aggr) 78 | 79 | def forward(self, x: Tensor, edge_index: Adj, 80 | edge_weight: OptTensor = None, cache_support: bool = False) \ 81 | -> Tensor: 82 | """""" 83 | # x: [batch, (steps), nodes, nodes] 84 | n = x.size(-2) 85 | if self._support is None: 86 | support = self.compute_support_index(edge_index, edge_weight, 87 | add_backward=self.add_backward, 88 | num_nodes=n) 89 | if cache_support: 90 | self._support = support 91 | else: 92 | support = self._support 93 | 94 | out = [] 95 | if self.root_weight: 96 | out += [x] 97 | 98 | for sup_index, sup_weights in support: 99 | x_sup = x 100 | for _ in range(self.k): 101 | x_sup = self.propagate(sup_index, x=x_sup, weight=sup_weights) 102 | out += [x_sup] 103 | 104 | out = torch.cat(out, -1) 105 | return self.filters(out) 106 | -------------------------------------------------------------------------------- /tsl/nn/layers/graph_convs/gated_gn.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from torch_geometric.nn import MessagePassing 5 | 6 | from tsl.nn.utils import get_layer_activation 7 | 8 | 9 | class GatedGraphNetwork(MessagePassing): 10 | r""" 11 | 12 | Gate Graph Neural Network model inspired by 13 | Satorras et al., "Multivariate Time Series Forecasting with Latent Graph Inference", arxiv 2022. 14 | 15 | Args: 16 | input_size (int): Input channels. 17 | output_size (int): Output channels. 18 | activation (str, optional): Activation function. 19 | """ 20 | 21 | def __init__(self, 22 | input_size: int, 23 | output_size: int, 24 | activation='silu'): 25 | super(GatedGraphNetwork, self).__init__(aggr="add", node_dim=-2) 26 | 27 | self.in_channels = input_size 28 | self.out_channels = output_size 29 | 30 | self.msg_mlp = nn.Sequential( 31 | nn.Linear(2 * input_size, output_size // 2), 32 | get_layer_activation(activation)(), 33 | nn.Linear(output_size // 2, output_size), 34 | get_layer_activation(activation)(), 35 | ) 36 | 37 | self.gate_mlp = nn.Sequential( 38 | nn.Linear(output_size, 1), 39 | nn.Sigmoid() 40 | ) 41 | 42 | self.update_mlp = nn.Sequential( 43 | nn.Linear(input_size + output_size, output_size), 44 | get_layer_activation(activation)(), 45 | nn.Linear(output_size, output_size) 46 | ) 47 | 48 | if input_size != output_size: 49 | self.skip_conn = nn.Linear(input_size, output_size) 50 | else: 51 | self.skip_conn = nn.Identity() 52 | 53 | def forward(self, x, edge_index): 54 | """""" 55 | 56 | out = self.propagate(edge_index, x=x) 57 | 58 | out = self.update_mlp(torch.cat([out, x], -1)) + self.skip_conn(x) 59 | 60 | return out 61 | 62 | def message(self, x_i, x_j): 63 | mij = self.msg_mlp(torch.cat([x_i, x_j], -1)) 64 | return self.gate_mlp(mij) * mij 65 | -------------------------------------------------------------------------------- /tsl/nn/layers/graph_convs/spatio_temporal_att.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from torch.nn import MultiheadAttention 3 | 4 | from einops import rearrange, reduce 5 | 6 | 7 | class SpatioTemporalAtt(nn.Module): 8 | def __init__(self, 9 | d_in, 10 | d_model, 11 | d_ff, 12 | n_heads, 13 | dropout, 14 | pool_size=1, 15 | pooling_op='mean'): 16 | super(SpatioTemporalAtt, self).__init__() 17 | self.d_in = d_in 18 | self.d_model = d_model 19 | self.d_ff = d_ff 20 | self.n_heads = n_heads 21 | self.pool_size = pool_size 22 | self.pooling_op = pooling_op 23 | 24 | if self.d_in != self.d_model: 25 | self.input_encoder = nn.Linear(self.d_in, self.d_model) 26 | else: 27 | self.input_encoder = nn.Identity() 28 | 29 | self.temporal_attn = MultiheadAttention(self.d_model, self.n_heads, dropout=dropout) 30 | self.spatial_attn = MultiheadAttention(self.d_model, self.n_heads, dropout=dropout) 31 | # Implementation of Feedforward model 32 | self.linear1 = nn.Linear(self.d_model, self.d_ff) 33 | self.linear2 = nn.Linear(self.d_ff, self.d_model) 34 | 35 | self.norm1 = nn.LayerNorm(self.d_model) 36 | self.norm2 = nn.LayerNorm(self.d_model) 37 | self.norm3 = nn.LayerNorm(self.d_model) 38 | self.dropout = nn.Dropout(dropout) 39 | self.dropout1 = nn.Dropout(dropout) 40 | self.dropout2 = nn.Dropout(dropout) 41 | self.dropout3 = nn.Dropout(dropout) 42 | 43 | def forward(self, x, **kwargs): 44 | # x: [batch, steps, nodes, features] 45 | # u: [batch, steps, nodes, features] 46 | b, s, n, f = x.size() 47 | x = rearrange(x, 'b s n f -> s (b n) f') 48 | 49 | x = self.input_encoder(x) 50 | if (self.pool_size > 1) and (s >= self.pool_size): 51 | q = reduce(x, '(s1 s2) m f -> s1 m f', self.pooling_op, s2=self.pool_size) 52 | else: 53 | q = x 54 | # temporal module 55 | x2 = self.temporal_attn(q, x, x)[0] 56 | x = x + self.dropout1(x2) 57 | x = self.norm1(x) 58 | x = rearrange(x, 's (b n) f -> n (b s) f', b=b, n=n) 59 | 60 | # spatial module 61 | x2 = self.spatial_attn(x, x, x)[0] 62 | x = x + self.dropout2(x2) 63 | x = self.norm2(x) 64 | 65 | # feed-forward network 66 | x2 = self.linear2(self.dropout(self.activation(self.linear1(x)))) 67 | x = x + self.dropout3(x2) 68 | x = self.norm3(x) 69 | return x 70 | 71 | -------------------------------------------------------------------------------- /tsl/nn/layers/link_predictor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import nn 3 | 4 | from tsl.nn.utils.utils import get_layer_activation 5 | 6 | 7 | class LinkPredictor(nn.Module): 8 | r""" 9 | Output a pairwise score for each couple of input elements. 10 | Can be used as a building block for a graph learning model. 11 | 12 | .. math:: 13 | \mathbf{S} = \left(\text{MLP}_s(\mathbf{E})\right) \left(\text{MLP}_t(\mathbf{E})\right)^T 14 | 15 | Args: 16 | emb_size: Size of the input embeddings. 17 | ff_size: Size of the hidden layer used to learn the scores. 18 | dropout: Dropout probability. 19 | activation: Activation function used in the hidden layer. 20 | """ 21 | def __init__(self, 22 | emb_size, 23 | ff_size, 24 | hidden_size, 25 | dropout=0., 26 | activation='relu'): 27 | super(LinkPredictor, self).__init__() 28 | self.source_mlp = nn.Sequential( 29 | nn.Linear(emb_size, ff_size), 30 | get_layer_activation(activation)(), 31 | nn.Dropout(dropout), 32 | nn.Linear(ff_size, hidden_size) 33 | ) 34 | 35 | self.target_mlp = nn.Sequential( 36 | nn.Linear(emb_size, ff_size), 37 | get_layer_activation(activation)(), 38 | nn.Dropout(dropout), 39 | nn.Linear(ff_size, hidden_size) 40 | ) 41 | 42 | def forward(self, x): 43 | """""" 44 | # x: [*, nodes, channels] 45 | z_s = self.source_mlp(x) 46 | z_t = self.target_mlp(x) 47 | # scores = z_s @ z_t.T 48 | logits = torch.einsum('... ik, ... jk -> ... ij', z_s, z_t) 49 | return logits 50 | -------------------------------------------------------------------------------- /tsl/nn/layers/norm/__init__.py: -------------------------------------------------------------------------------- 1 | from .layer_norm import LayerNorm 2 | from .instance_norm import InstanceNorm 3 | from .batch_norm import BatchNorm 4 | from .norm import Norm 5 | 6 | __all__ = [ 7 | 'Norm', 8 | 'LayerNorm', 9 | 'InstanceNorm', 10 | 'BatchNorm' 11 | ] 12 | 13 | classes = __all__ 14 | -------------------------------------------------------------------------------- /tsl/nn/layers/norm/batch_norm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from einops import rearrange 3 | from torch import Tensor 4 | 5 | 6 | class BatchNorm(torch.nn.Module): 7 | r"""Applies graph-wise batch normalization. 8 | 9 | Args: 10 | in_channels (int): Size of each input sample. 11 | eps (float, optional): A value added to the denominator for numerical 12 | stability. (default: :obj:`1e-5`) 13 | momentum (float, bool): Running stats momentum. 14 | affine (bool, optional): If set to :obj:`True`, this module has 15 | learnable affine parameters :math:`\gamma` and :math:`\beta`. 16 | (default: :obj:`True`) 17 | track_running_stats (bool, optional): Whether to track stats to perform 18 | batch norm. 19 | (default: :obj:`True`) 20 | """ 21 | 22 | def __init__(self, in_channels, eps: float = 1e-5, momentum: float = 0.1, 23 | affine: bool = True, 24 | track_running_stats: bool = True): 25 | super().__init__() 26 | self.module = torch.nn.BatchNorm1d(in_channels, eps, momentum, affine, 27 | track_running_stats) 28 | 29 | def reset_parameters(self): 30 | self.module.reset_parameters() 31 | 32 | def forward(self, x: Tensor) -> Tensor: 33 | """""" 34 | b, *_ = x.size() 35 | x = rearrange(x, 'b ... n c -> (b n) c ...') 36 | x = self.module(x) 37 | return rearrange(x, '(b n) c ... -> b ... n c', b=b) 38 | 39 | def __repr__(self): 40 | return f'{self.__class__.__name__}({self.module.num_features})' 41 | -------------------------------------------------------------------------------- /tsl/nn/layers/norm/instance_norm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from torch.nn import Parameter, Linear 4 | from torch_geometric.nn import inits 5 | 6 | 7 | class InstanceNorm(torch.nn.Module): 8 | r"""Applies graph-wise instance normalization. 9 | 10 | Args: 11 | in_channels (int): Size of each input sample. 12 | eps (float, optional): A value added to the denominator for numerical 13 | stability. (default: :obj:`1e-5`) 14 | affine (bool, optional): If set to :obj:`True`, this module has 15 | learnable affine parameters :math:`\gamma` and :math:`\beta`. 16 | (default: :obj:`True`) 17 | """ 18 | def __init__(self, in_channels, eps=1e-5, affine=True): 19 | super().__init__() 20 | 21 | self.in_channels = in_channels 22 | self.eps = eps 23 | 24 | if affine: 25 | self.weight = Parameter(torch.Tensor(in_channels)) 26 | self.bias = Parameter(torch.Tensor(in_channels)) 27 | else: 28 | self.register_parameter('weight', None) 29 | self.register_parameter('bias', None) 30 | 31 | self.reset_parameters() 32 | 33 | def reset_parameters(self): 34 | inits.ones(self.weight) 35 | inits.zeros(self.bias) 36 | 37 | def forward(self, x: Tensor) -> Tensor: 38 | # x : [*, nodes, features] 39 | mean = torch.mean(x, dim=-2, keepdim=True) 40 | std = torch.std(x, dim=-2, unbiased=False, keepdim=True) 41 | 42 | out = (x - mean) / (std + self.eps) 43 | 44 | if self.weight is not None and self.bias is not None: 45 | out = out * self.weight + self.bias 46 | 47 | return out 48 | 49 | def __repr__(self): 50 | return f'{self.__class__.__name__}({self.in_channels})' 51 | -------------------------------------------------------------------------------- /tsl/nn/layers/norm/layer_norm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from torch.nn import Parameter 4 | from torch_geometric.nn import inits 5 | 6 | 7 | class LayerNorm(torch.nn.Module): 8 | r"""Applies layer normalization. 9 | 10 | Args: 11 | in_channels (int): Size of each input sample. 12 | eps (float, optional): A value added to the denominator for numerical 13 | stability. (default: :obj:`1e-5`) 14 | affine (bool, optional): If set to :obj:`True`, this module has 15 | learnable affine parameters :math:`\gamma` and :math:`\beta`. 16 | (default: :obj:`True`) 17 | """ 18 | 19 | def __init__(self, in_channels, eps=1e-5, affine=True): 20 | super().__init__() 21 | 22 | self.in_channels = in_channels 23 | self.eps = eps 24 | 25 | if affine: 26 | self.weight = Parameter(torch.Tensor(in_channels)) 27 | self.bias = Parameter(torch.Tensor(in_channels)) 28 | else: 29 | self.register_parameter('weight', None) 30 | self.register_parameter('bias', None) 31 | 32 | self.reset_parameters() 33 | 34 | def reset_parameters(self): 35 | inits.ones(self.weight) 36 | inits.zeros(self.bias) 37 | 38 | def forward(self, x: Tensor) -> Tensor: 39 | """""" 40 | mean = torch.mean(x, dim=-1, keepdim=True) 41 | std = torch.std(x, dim=-1, unbiased=False, keepdim=True) 42 | 43 | out = (x - mean) / (std + self.eps) 44 | 45 | if self.weight is not None and self.bias is not None: 46 | out = out * self.weight + self.bias 47 | 48 | return out 49 | 50 | def __repr__(self): 51 | return f'{self.__class__.__name__}({self.in_channels})' 52 | -------------------------------------------------------------------------------- /tsl/nn/layers/norm/norm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch import Tensor 3 | from .batch_norm import BatchNorm 4 | from .layer_norm import LayerNorm 5 | from .instance_norm import InstanceNorm 6 | 7 | from torch import nn 8 | 9 | 10 | class Norm(torch.nn.Module): 11 | r"""Applies a normalization of the specified type. 12 | 13 | Args: 14 | in_channels (int): Size of each input sample. 15 | """ 16 | def __init__(self, norm_type, in_channels, **kwargs): 17 | super().__init__() 18 | self.norm_type = norm_type 19 | self.in_channels = in_channels 20 | 21 | if norm_type == 'instance': 22 | norm_layer = InstanceNorm 23 | elif norm_type == 'batch': 24 | norm_layer = BatchNorm 25 | elif norm_type == 'layer': 26 | norm_layer = LayerNorm 27 | elif norm_type == 'none': 28 | norm_layer = nn.Identity 29 | else: 30 | raise NotImplementedError(f'"{norm_type}" is not a valid normalization option.') 31 | 32 | self.norm = norm_layer(in_channels, **kwargs) 33 | 34 | def forward(self, x: Tensor) -> Tensor: 35 | """""" 36 | return self.norm(x) 37 | 38 | def __repr__(self): 39 | return f'{self.__class__.__name__}({self.norm_type}, {self.in_channels})' 40 | -------------------------------------------------------------------------------- /tsl/nn/layers/positional_encoding.py: -------------------------------------------------------------------------------- 1 | import math 2 | 3 | import torch 4 | from torch import nn 5 | 6 | 7 | class PositionalEncoding(nn.Module): 8 | """ 9 | Implementation of the positional encoding from Vaswani et al. 2017 10 | """ 11 | def __init__(self, d_model, dropout=0., max_len=5000, affinity=False, batch_first=True): 12 | super(PositionalEncoding, self).__init__() 13 | self.dropout = nn.Dropout(p=dropout) 14 | if affinity: 15 | self.affinity = nn.Linear(d_model, d_model) 16 | else: 17 | self.affinity = None 18 | pe = torch.zeros(max_len, d_model) 19 | position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) 20 | div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) 21 | pe[:, 0::2] = torch.sin(position * div_term) 22 | pe[:, 1::2] = torch.cos(position * div_term) 23 | pe = pe.unsqueeze(0).transpose(0, 1) 24 | self.register_buffer('pe', pe) 25 | self.batch_first = batch_first 26 | 27 | def forward(self, x): 28 | if self.affinity is not None: 29 | x = self.affinity(x) 30 | pe = self.pe[:x.size(1), :] if self.batch_first else self.pe[:x.size(0), :] 31 | x = x + pe 32 | return self.dropout(x) -------------------------------------------------------------------------------- /tsl/nn/layers/spatial_attention.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | from torch import nn, Tensor 4 | from torch_geometric.nn.dense import Linear 5 | 6 | from tsl.nn.base.attention.attention import MultiHeadAttention 7 | 8 | 9 | class SpatialSelfAttention(nn.Module): 10 | """ 11 | Spatial Self Attention layer. 12 | 13 | Args: 14 | embed_dim (int): Size of the hidden dimension associeted with each node at each time step. 15 | num_heads (int): Number of parallel attention heads. 16 | dropout (float): Dropout probability. 17 | bias (bool, optional): Whther to add a learnable bias. 18 | device (optional): Device on which store the model. 19 | dtype (optional): Data Type of the parameters. 20 | Examples:: 21 | >>> import torch 22 | >>> m = SpatialSelfAttention(32, 4, -1) 23 | >>> input = torch.randn(128, 24, 10, 20) 24 | >>> output, _ = m(input) 25 | >>> print(output.size()) 26 | torch.Size([128, 24, 10, 32]) 27 | """ 28 | 29 | def __init__(self, embed_dim, num_heads, 30 | in_channels=None, 31 | dropout=0., 32 | bias=True, 33 | device=None, 34 | dtype=None) -> None: 35 | super(SpatialSelfAttention, self).__init__() 36 | 37 | self.embed_dim = embed_dim 38 | 39 | if in_channels is not None: 40 | self.input_encoder = Linear(in_channels, self.embed_dim) 41 | else: 42 | self.input_encoder = nn.Identity() 43 | 44 | self.attention = MultiHeadAttention(embed_dim, num_heads, 45 | axis='nodes', 46 | dropout=dropout, 47 | bias=bias, 48 | device=device, 49 | dtype=dtype) 50 | 51 | def forward(self, x, attn_mask: Optional[Tensor] = None, 52 | key_padding_mask: Optional[Tensor] = None, 53 | need_weights: bool = True): 54 | """""" 55 | # x: [batch, steps, nodes, in_channels] 56 | x = self.input_encoder(x) # -> [batch, steps, nodes, embed_dim] 57 | return self.attention(x, 58 | attn_mask=attn_mask, 59 | key_padding_mask=key_padding_mask, 60 | need_weights=need_weights) 61 | -------------------------------------------------------------------------------- /tsl/nn/layers/temporal_attention.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | from torch import nn, Tensor 4 | from torch_geometric.nn.dense import Linear 5 | 6 | from tsl.nn.base.attention.attention import MultiHeadAttention 7 | 8 | 9 | class TemporalSelfAttention(nn.Module): 10 | """ 11 | Temporal Self Attention layer. 12 | 13 | Args: 14 | embed_dim (int): Size of the hidden dimension associeted with each node at each time step. 15 | num_heads (int): Number of parallel attention heads. 16 | dropout (float): Dropout probability. 17 | bias (bool, optional): Whther to add a learnable bias. 18 | device (optional): Device on which store the model. 19 | dtype (optional): Data Type of the parameters. 20 | Examples:: 21 | >>> import torch 22 | >>> m = SpatialSelfAttention(32, 4, -1) 23 | >>> input = torch.randn(128, 24, 10, 20) 24 | >>> output, _ = m(input) 25 | >>> print(output.size()) 26 | torch.Size([128, 24, 10, 32]) 27 | """ 28 | def __init__(self, embed_dim, num_heads, 29 | in_channels=None, 30 | dropout=0., 31 | bias=True, 32 | device=None, 33 | dtype=None) -> None: 34 | super(TemporalSelfAttention, self).__init__() 35 | 36 | self.embed_dim = embed_dim 37 | 38 | if in_channels is not None: 39 | self.input_encoder = Linear(in_channels, self.embed_dim) 40 | else: 41 | self.input_encoder = nn.Identity() 42 | 43 | self.attention = MultiHeadAttention(embed_dim, num_heads, 44 | axis='steps', 45 | dropout=dropout, 46 | bias=bias, 47 | device=device, 48 | dtype=dtype) 49 | 50 | def forward(self, x, attn_mask: Optional[Tensor] = None, 51 | key_padding_mask: Optional[Tensor] = None, 52 | need_weights: bool = True): 53 | """""" 54 | # x: [batch, steps, nodes, in_channels] 55 | x = self.input_encoder(x) # -> [batch, steps, nodes, embed_dim] 56 | return self.attention(x, 57 | attn_mask=attn_mask, 58 | key_padding_mask=key_padding_mask, 59 | need_weights=need_weights) 60 | -------------------------------------------------------------------------------- /tsl/nn/metrics/__init__.py: -------------------------------------------------------------------------------- 1 | from .metric_base import MaskedMetric 2 | from .multi_loss import MaskedMultiLoss 3 | from .metrics import MaskedMAE, MaskedMRE, MaskedMSE, MaskedMAPE 4 | from .pinball_loss import MaskedPinballLoss 5 | -------------------------------------------------------------------------------- /tsl/nn/metrics/metric_wrappers.py: -------------------------------------------------------------------------------- 1 | from tsl.nn.metrics.metric_base import MaskedMetric 2 | from tsl.utils.python_utils import ensure_list 3 | 4 | 5 | class MaskedMetricWrapper(MaskedMetric): 6 | def __init__(self, 7 | metric: MaskedMetric, 8 | input_preprocessing=None, 9 | target_preprocessing=None, 10 | mask_preprocessing=None): 11 | super(MaskedMetricWrapper, self).__init__(None) 12 | self.metric = metric 13 | 14 | if input_preprocessing is None: 15 | input_preprocessing = lambda x: x 16 | 17 | if target_preprocessing is None: 18 | target_preprocessing = lambda x: x 19 | 20 | if mask_preprocessing is None: 21 | mask_preprocessing = lambda x: x 22 | 23 | self.input_preprocessing = input_preprocessing 24 | self.target_preprocessing = target_preprocessing 25 | self.mask_preprocessing = mask_preprocessing 26 | 27 | def update(self, y_hat, y, mask=None): 28 | y_hat = self.input_preprocessing(y_hat) 29 | y = self.target_preprocessing(y) 30 | mask = self.mask_preprocessing(mask) 31 | return self.metric.update(y_hat, y, mask) 32 | 33 | def compute(self): 34 | return self.metric.compute() 35 | 36 | def reset(self): 37 | self.metric.reset() 38 | super(MaskedMetricWrapper, self).reset() 39 | 40 | 41 | class SplitMetricWrapper(MaskedMetricWrapper): 42 | def __init__(self, metric, input_idx=None, target_idx=None, mask_idx=None): 43 | if input_idx is not None: 44 | input_preprocessing = lambda x: x[input_idx] 45 | else: 46 | input_preprocessing = None 47 | 48 | if target_idx is not None: 49 | target_preprocessing = lambda x: x[target_idx] 50 | else: 51 | target_preprocessing = None 52 | 53 | if mask_idx is not None: 54 | map_preprocessing = lambda x: x[mask_idx] 55 | else: 56 | map_preprocessing = None 57 | super(SplitMetricWrapper, self).__init__(metric, 58 | input_preprocessing=input_preprocessing, 59 | target_preprocessing=target_preprocessing, 60 | mask_preprocessing=map_preprocessing) 61 | 62 | 63 | class ChannelSplitMetricWrapper(MaskedMetricWrapper): 64 | def __init__(self, metric, input_channels=None, target_channels=None, map_channels=None): 65 | if input_channels is not None: 66 | input_preprocessing = lambda x: x[..., ensure_list(input_channels)] 67 | else: 68 | input_preprocessing = None 69 | 70 | if target_channels is not None: 71 | target_preprocessing = lambda x: x[..., ensure_list(target_channels)] 72 | else: 73 | target_preprocessing = None 74 | 75 | if map_channels is not None: 76 | map_preprocessing = lambda x: x[..., ensure_list(map_channels)] 77 | else: 78 | map_preprocessing = None 79 | super(ChannelSplitMetricWrapper, self).__init__(metric, 80 | input_preprocessing=input_preprocessing, 81 | target_preprocessing=target_preprocessing, 82 | mask_preprocessing=map_preprocessing) 83 | -------------------------------------------------------------------------------- /tsl/nn/metrics/multi_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from tsl.nn.metrics.metric_base import MaskedMetric 3 | import torch.nn as nn 4 | 5 | class MaskedMultiLoss(MaskedMetric): 6 | r""" 7 | Adapted from: https://github.com/jdb78/pytorch-forecasting/blob/master/pytorch_forecasting/metrics.py 8 | Metric that can be used to combine multiple metrics. 9 | 10 | Args: 11 | metrics: List of metrics. 12 | weights (optional): List of weights for the corresponding metrics. 13 | """ 14 | def __init__(self, metrics, weights=None): 15 | super().__init__(None, compute_on_step=True) 16 | assert len(metrics) > 0, "at least one metric has to be specified" 17 | if weights is None: 18 | weights = [1.0 for _ in metrics] 19 | assert len(weights) == len(metrics), "Number of weights has to match number of metrics" 20 | 21 | self.metrics = nn.ModuleList(metrics) 22 | self.weights = weights 23 | 24 | def __repr__(self): 25 | name = ( 26 | f"{self.__class__.__name__}(" 27 | + ", ".join([f"{w:.3g} * {repr(m)}" if w != 1.0 else repr(m) for w, m in zip(self.weights, self.metrics)]) 28 | + ")" 29 | ) 30 | return name 31 | 32 | def __iter__(self): 33 | """ 34 | Iterate over metrics. 35 | """ 36 | return iter(self.metrics) 37 | 38 | def __len__(self) -> int: 39 | """ 40 | Number of metrics. 41 | Returns: 42 | int: number of metrics 43 | """ 44 | return len(self.metrics) 45 | 46 | def update(self, y_hat: torch.Tensor, y: torch.Tensor, mask=None): 47 | """ 48 | Update composite metric 49 | Args: 50 | y_hat: network output 51 | y: actual values 52 | Returns: 53 | torch.Tensor: metric value on which backpropagation can be applied 54 | """ 55 | assert len(self) == y_hat.size(0) 56 | for idx, metric in enumerate(self.metrics): 57 | metric.update(y_hat[idx], y, mask) 58 | 59 | def compute(self) -> torch.Tensor: 60 | """ 61 | Get metric 62 | Returns: 63 | torch.Tensor: metric 64 | """ 65 | results = [] 66 | for weight, metric in zip(self.weights, self.metrics): 67 | results.append(metric.compute() * weight) 68 | 69 | if len(results) == 1: 70 | results = results[0] 71 | else: 72 | results = torch.stack(results, dim=0).sum(0) 73 | return results 74 | 75 | def reset(self) -> None: 76 | for m in self.metrics: 77 | m.reset() 78 | super(MaskedMultiLoss, self).reset() 79 | 80 | def __getitem__(self, idx: int): 81 | """ 82 | Return metric. 83 | Args: 84 | idx (int): metric index 85 | """ 86 | return self.metrics[idx] 87 | -------------------------------------------------------------------------------- /tsl/nn/metrics/pinball_loss.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from .metric_base import MaskedMetric 3 | from tsl.utils.python_utils import ensure_list 4 | 5 | 6 | def _pinball_loss(y_hat, y, q): 7 | err = y - y_hat 8 | return torch.maximum((q - 1) * err, q * err) 9 | 10 | def _multi_quantile_pinball_loss(y_hat, y, q): 11 | q = ensure_list(q) 12 | assert y_hat.size(0) == len(q) 13 | loss = torch.zeros_like(y_hat) 14 | for i, qi in enumerate(q): 15 | loss += _pinball_loss(y_hat[i], y, qi) 16 | return loss 17 | 18 | 19 | class MaskedPinballLoss(MaskedMetric): 20 | """ 21 | Quantile loss. 22 | 23 | Args: 24 | q (float): Target quantile. 25 | mask_nans (bool, optional): Whether to automatically mask nan values. 26 | mask_inf (bool, optional): Whether to automatically mask infinite values. 27 | compute_on_step (bool, optional): Whether to compute the metric right-away or if accumulate the results. 28 | This should be `True` when using the metric to compute a loss function, `False` if the metric 29 | is used for logging the aggregate error across different minibatches. 30 | at (int, optional): Whether to compute the metric only w.r.t. a certain time step. 31 | """ 32 | def __init__(self, 33 | q, 34 | mask_nans=False, 35 | mask_inf=False, 36 | compute_on_step=True, 37 | dist_sync_on_step=False, 38 | process_group=None, 39 | dist_sync_fn=None, 40 | at=None): 41 | super(MaskedPinballLoss, self).__init__(metric_fn=_pinball_loss, 42 | mask_nans=mask_nans, 43 | mask_inf=mask_inf, 44 | compute_on_step=compute_on_step, 45 | dist_sync_on_step=dist_sync_on_step, 46 | process_group=process_group, 47 | dist_sync_fn=dist_sync_fn, 48 | metric_kwargs={'q': q}, 49 | at=at) 50 | -------------------------------------------------------------------------------- /tsl/nn/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .rnn_model import RNNModel, FCRNNModel 2 | from .tcn_model import TCNModel 3 | from .transformer_model import TransformerModel 4 | 5 | from . import imputation 6 | from . import stgn 7 | 8 | __all__ = [ 9 | 'imputation', 10 | 'stgn', 11 | 'FCRNNModel', 12 | 'RNNModel', 13 | 'TCNModel', 14 | 'TransformerModel' 15 | ] 16 | 17 | classes = __all__[2:] 18 | -------------------------------------------------------------------------------- /tsl/nn/models/imputation/__init__.py: -------------------------------------------------------------------------------- 1 | from .rnni_models import RNNImputerModel, BiRNNImputerModel 2 | from .grin_model import GRINModel 3 | 4 | __all__ = [ 5 | 'RNNImputerModel', 6 | 'BiRNNImputerModel', 7 | 'GRINModel' 8 | ] 9 | 10 | classes = __all__ 11 | -------------------------------------------------------------------------------- /tsl/nn/models/stgn/__init__.py: -------------------------------------------------------------------------------- 1 | from .dcrnn_model import DCRNNModel 2 | from .graph_wavenet_model import GraphWaveNetModel 3 | from .gated_gn_model import GatedGraphNetworkModel 4 | from .rnn2gcn_model import RNNEncGCNDecModel 5 | from .stcn_model import STCNModel 6 | 7 | __all__ = [ 8 | 'DCRNNModel', 9 | 'GraphWaveNetModel', 10 | 'GatedGraphNetworkModel', 11 | 'RNNEncGCNDecModel', 12 | 'STCNModel' 13 | ] 14 | 15 | classes = __all__ 16 | -------------------------------------------------------------------------------- /tsl/nn/models/stgn/dcrnn_model.py: -------------------------------------------------------------------------------- 1 | from tsl.nn.blocks.encoders.dcrnn import DCRNN 2 | from tsl.utils.parser_utils import ArgParser 3 | 4 | from einops import rearrange 5 | from torch import nn 6 | 7 | from tsl.nn.blocks.encoders import ConditionalBlock 8 | from tsl.nn.blocks.decoders.mlp_decoder import MLPDecoder 9 | 10 | 11 | class DCRNNModel(nn.Module): 12 | r""" 13 | Diffusion ConvolutionalRecurrent Neural Network with a nonlinear readout. 14 | 15 | From Li et al., "Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting", ICLR 2018. 16 | 17 | Args: 18 | input_size (int): Size of the input. 19 | hidden_size (int): Number of units in the DCRNN hidden layer. 20 | ff_size (int): Number of units in the nonlinear readout. 21 | output_size (int): Number of output channels. 22 | n_layers (int): Number DCRNN cells. 23 | exog_size (int): Number of channels in the exogenous variable. 24 | horizon (int): Number of steps to forecast. 25 | activation (str, optional): Activation function in the readout. 26 | dropout (float, optional): Dropout probability. 27 | """ 28 | 29 | def __init__(self, 30 | input_size, 31 | hidden_size, 32 | ff_size, 33 | output_size, 34 | n_layers, 35 | exog_size, 36 | horizon, 37 | activation='relu', 38 | dropout=0., 39 | kernel_size=2): 40 | super(DCRNNModel, self).__init__() 41 | if exog_size: 42 | self.input_encoder = ConditionalBlock(input_size=input_size, 43 | exog_size=exog_size, 44 | output_size=hidden_size, 45 | activation=activation) 46 | else: 47 | self.input_encoder = nn.Linear(input_size, hidden_size) 48 | 49 | self.dcrnn = DCRNN(input_size=hidden_size, 50 | hidden_size=hidden_size, 51 | n_layers=n_layers, 52 | k=kernel_size) 53 | 54 | self.readout = MLPDecoder(input_size=hidden_size, 55 | hidden_size=ff_size, 56 | output_size=output_size, 57 | horizon=horizon, 58 | activation=activation, 59 | dropout=dropout) 60 | 61 | def forward(self, x, edge_index, edge_weight=None, u=None, **kwargs): 62 | if u is not None: 63 | if u.dim() == 3: 64 | u = rearrange(u, 'b s c -> b s 1 c') 65 | x = self.input_encoder(x, u) 66 | else: 67 | x = self.input_encoder(x) 68 | 69 | h, _ = self.dcrnn(x, edge_index, edge_weight) 70 | return self.readout(h) 71 | 72 | @staticmethod 73 | def add_model_specific_args(parser: ArgParser): 74 | parser.opt_list('--hidden-size', type=int, default=32, tunable=True, options=[16, 32, 64, 128]) 75 | parser.opt_list('--ff-size', type=int, default=256, tunable=True, options=[64, 128, 256, 512]) 76 | parser.opt_list('--n-layers', type=int, default=1, tunable=True, options=[1, 2]) 77 | parser.opt_list('--dropout', type=float, default=0., tunable=True, options=[0., 0.1, 0.25, 0.5]) 78 | parser.opt_list('--kernel-size', type=int, default=2, tunable=True, options=[1, 2]) 79 | return parser 80 | -------------------------------------------------------------------------------- /tsl/nn/models/stgn/nri_model.py: -------------------------------------------------------------------------------- 1 | from tsl.nn.blocks.encoders.nri_dcrnn import NeuRelInfDCRNN 2 | from tsl.utils.parser_utils import ArgParser 3 | 4 | from einops import rearrange 5 | from torch import nn 6 | 7 | from tsl.nn.blocks.encoders import ConditionalBlock 8 | from tsl.nn.blocks.decoders.mlp_decoder import MLPDecoder 9 | 10 | 11 | class NRIModel(nn.Module): 12 | """ 13 | Simple model performing graph learning with a binary sampler and a DCRNN backbone. 14 | 15 | Args: 16 | input_size (int): Size of the input. 17 | hidden_size (int): Number of units in the recurrent cell. 18 | ff_size (int): Number of units in the link predictor. 19 | emb_size (int): Number of features for the node embeddings. 20 | output_size (int): Number of output channels. 21 | n_layers (int): Number of DCRNN layers. 22 | exog_size (int): Size of the exogenous variables. 23 | horizon (int): Number of forecasting steps. 24 | n_nodes (int): Number of nodes in the input graph. 25 | sampler_tau (float, optional): Temperature of the binary sampler. 26 | activation (str, optional): Activation function. 27 | dropout (float, optional): Dropout probability. 28 | kernel_size (int, optional): Order of the spatial diffusion process. 29 | """ 30 | def __init__(self, 31 | input_size, 32 | hidden_size, 33 | ff_size, 34 | emb_size, 35 | output_size, 36 | n_layers, 37 | exog_size, 38 | horizon, 39 | n_nodes, 40 | sampler_tau=0.25, 41 | activation='relu', 42 | dropout=0., 43 | kernel_size=2): 44 | super(NRIModel, self).__init__() 45 | self.tau = sampler_tau 46 | if exog_size: 47 | self.input_encoder = ConditionalBlock(input_size=input_size, 48 | exog_size=exog_size, 49 | output_size=hidden_size, 50 | activation=activation) 51 | else: 52 | self.input_encoder = nn.Linear(input_size, hidden_size) 53 | 54 | self.nri_dcrnn = NeuRelInfDCRNN(input_size=hidden_size, 55 | hidden_size=hidden_size, 56 | n_layers=n_layers, 57 | n_nodes=n_nodes, 58 | emb_size=emb_size, 59 | k=kernel_size) 60 | 61 | self.readout = MLPDecoder(input_size=hidden_size, 62 | hidden_size=ff_size, 63 | output_size=output_size, 64 | horizon=horizon, 65 | activation=activation, 66 | dropout=dropout) 67 | 68 | def forward(self, x, u=None, **kwargs): 69 | """""" 70 | if u is not None: 71 | if u.dim() == 3: 72 | u = rearrange(u, 'b s c -> b s 1 c') 73 | x = self.input_encoder(x, u) 74 | else: 75 | x = self.input_encoder(x) 76 | 77 | h, _ = self.nri_dcrnn(x, tau=self.tau) 78 | return self.readout(h) 79 | 80 | @staticmethod 81 | def add_model_specific_args(parser: ArgParser): 82 | parser.opt_list('--hidden-size', type=int, default=32, tunable=True, options=[16, 32, 64, 128]) 83 | parser.opt_list('--ff-size', type=int, default=256, tunable=True, options=[64, 128, 256, 512]) 84 | parser.opt_list('--emb-size', type=int, default=10, tunable=True, options=[8, 10, 16, 32, 64]) 85 | parser.opt_list('--sampler-tau', type=float, default=0.25, tunable=True, options=[0.1, 0.25, 0.5, 0.75, 1]) 86 | parser.opt_list('--n-layers', type=int, default=1, tunable=True, options=[1, 2]) 87 | parser.opt_list('--dropout', type=float, default=0., tunable=True, options=[0., 0.1, 0.25, 0.5]) 88 | parser.opt_list('--kernel-size', type=int, default=2, tunable=True, options=[1, 2]) 89 | return parser 90 | -------------------------------------------------------------------------------- /tsl/nn/models/stgn/rnn2gcn_model.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from tsl.utils.parser_utils import ArgParser 3 | 4 | from tsl.nn.blocks.encoders import ConditionalBlock 5 | from tsl.nn.blocks.decoders.gcn_decoder import GCNDecoder 6 | from tsl.nn.blocks.encoders.rnn import RNN 7 | 8 | from einops import rearrange 9 | 10 | 11 | class RNNEncGCNDecModel(nn.Module): 12 | """ 13 | Simple time-then-space model. 14 | 15 | Input time series are encoded in vectors using an RNN and then decoded using a stack of GCN layers. 16 | 17 | Args: 18 | input_size (int): Input size. 19 | hidden_size (int): Units in the hidden layers. 20 | output_size (int): Size of the optional readout. 21 | exog_size (int): Size of the exogenous variables. 22 | rnn_layers (int): Number of recurrent layers in the encoder. 23 | gcn_layers (int): Number of graph convolutional layers in the decoder. 24 | rnn_dropout (float, optional): Dropout probability in the RNN encoder. 25 | gcn_dropout (float, optional): Dropout probability int the GCN decoder. 26 | horizon (int): Forecasting horizon. 27 | cell_type (str, optional): Type of cell that should be use (options: [`gru`, `lstm`]). (default: `gru`) 28 | activation (str, optional): Activation function. 29 | """ 30 | def __init__(self, 31 | input_size, 32 | hidden_size, 33 | output_size, 34 | exog_size, 35 | rnn_layers, 36 | gcn_layers, 37 | rnn_dropout, 38 | gcn_dropout, 39 | horizon, 40 | cell_type='gru', 41 | activation='relu'): 42 | super(RNNEncGCNDecModel, self).__init__() 43 | 44 | if exog_size > 0: 45 | self.input_encoder = ConditionalBlock(input_size=input_size, 46 | exog_size=exog_size, 47 | output_size=hidden_size, 48 | activation=activation) 49 | else: 50 | self.input_encoder = nn.Sequential( 51 | nn.Linear(input_size, hidden_size), 52 | ) 53 | 54 | self.encoder = RNN(input_size=hidden_size, 55 | hidden_size=hidden_size, 56 | n_layers=rnn_layers, 57 | dropout=rnn_dropout, 58 | cell=cell_type) 59 | 60 | self.decoder = GCNDecoder( 61 | input_size=hidden_size, 62 | hidden_size=hidden_size, 63 | output_size=output_size, 64 | horizon=horizon, 65 | n_layers=gcn_layers, 66 | activation=activation, 67 | dropout=gcn_dropout 68 | ) 69 | 70 | def forward(self, x, edge_index, edge_weight, u=None, **kwargs): 71 | """""" 72 | # x: [batches steps nodes features] 73 | # u: [batches steps (nodes) features] 74 | if u is not None: 75 | if u.dim() == 3: 76 | u = rearrange(u, 'b s f -> b s 1 f') 77 | x = self.input_encoder(x, u) 78 | else: 79 | x = self.input_encoder(x) 80 | 81 | x = self.encoder(x, return_last_state=True) 82 | 83 | return self.decoder(x, edge_index, edge_weight) 84 | 85 | @staticmethod 86 | def add_model_specific_args(parser: ArgParser): 87 | parser.opt_list('--hidden-size', type=int, default=32, tunable=True, options=[16, 32, 64, 128, 256]) 88 | parser.opt_list('--rnn-layers', type=int, default=1, tunable=True, options=[1, 2, 3]) 89 | parser.opt_list('--gcn-layers', type=int, default=1, tunable=True, options=[1, 2, 3]) 90 | parser.opt_list('--rnn-dropout', type=float, default=0., tunable=True, options=[0., 0.1, 0.2]) 91 | parser.opt_list('--gcn-dropout', type=float, default=0., tunable=True, options=[0., 0.1, 0.25, 0.5]) 92 | parser.opt_list('--cell-type', type=str, default='gru', tunable=True, options=['gru', 'lstm']) 93 | return parser 94 | -------------------------------------------------------------------------------- /tsl/nn/models/transformer_model.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | from einops import rearrange 3 | 4 | from tsl.utils.parser_utils import ArgParser 5 | 6 | from tsl.nn.blocks.encoders import ConditionalBlock 7 | from tsl.nn.blocks.encoders.mlp import MLP 8 | from tsl.nn.blocks.encoders.transformer import Transformer 9 | from tsl.nn.ops.ops import Select 10 | from tsl.nn.layers.positional_encoding import PositionalEncoding 11 | 12 | from einops.layers.torch import Rearrange 13 | 14 | 15 | class TransformerModel(nn.Module): 16 | r""" 17 | Simple Transformer for multi-step time series forecasting. 18 | 19 | Args: 20 | input_size (int): Input size. 21 | hidden_size (int): Dimension of the learned representations. 22 | output_size (int): Dimension of the output. 23 | ff_size (int): Units in the MLP after self attention. 24 | exog_size (int): Dimension of the exogenous variables. 25 | horizon (int): Number of forecasting steps. 26 | n_heads (int, optional): Number of parallel attention heads. 27 | n_layers (int, optional): Number of layers. 28 | dropout (float, optional): Dropout probability. 29 | axis (str, optional): Dimension on which to apply attention to update the representations. 30 | activation (str, optional): Activation function. 31 | """ 32 | 33 | def __init__(self, 34 | input_size, 35 | hidden_size, 36 | output_size, 37 | ff_size, 38 | exog_size, 39 | horizon, 40 | n_heads, 41 | n_layers, 42 | dropout, 43 | axis, 44 | activation='elu'): 45 | super(TransformerModel, self).__init__() 46 | 47 | if exog_size > 0: 48 | self.input_encoder = ConditionalBlock(input_size=input_size, 49 | exog_size=exog_size, 50 | output_size=hidden_size, 51 | activation=activation) 52 | else: 53 | self.input_encoder = nn.Linear(input_size, hidden_size) 54 | 55 | self.pe = PositionalEncoding(hidden_size, max_len=100) 56 | 57 | self.transformer_encoder = nn.Sequential( 58 | Transformer(input_size=hidden_size, 59 | hidden_size=hidden_size, 60 | ff_size=ff_size, 61 | n_heads=n_heads, 62 | n_layers=n_layers, 63 | activation=activation, 64 | dropout=dropout, 65 | axis=axis), 66 | Select(1, -1) 67 | ) 68 | 69 | self.readout = nn.Sequential( 70 | MLP(input_size=hidden_size, 71 | hidden_size=ff_size, 72 | output_size=output_size * horizon, 73 | dropout=dropout), 74 | Rearrange('b n (h c) -> b h n c', c=output_size, h=horizon) 75 | ) 76 | 77 | def forward(self, x, u=None, **kwargs): 78 | # x: [batches steps nodes features] 79 | # u: [batches steps (nodes) features] 80 | b, *_ = x.size() 81 | if u is not None: 82 | if u.dim() == 3: 83 | u = rearrange(u, 'b s f -> b s 1 f') 84 | x = self.input_encoder(x, u) 85 | else: 86 | x = self.input_encoder(x) 87 | x = self.pe(x) 88 | x = self.transformer_encoder(x) 89 | 90 | return self.readout(x) 91 | 92 | @staticmethod 93 | def add_model_specific_args(parser: ArgParser): 94 | parser.opt_list('--hidden-size', type=int, default=32, tunable=True, options=[16, 32, 64, 128, 256]) 95 | parser.opt_list('--ff-size', type=int, default=32, tunable=True, options=[32, 64, 128, 256, 512, 1024]) 96 | parser.opt_list('--n-layers', type=int, default=1, tunable=True, options=[1, 2, 3]) 97 | parser.opt_list('--n-heads', type=int, default=1, tunable=True, options=[1, 2, 3]) 98 | parser.opt_list('--dropout', type=float, default=0., tunable=True, options=[0., 0.1, 0.25, 0.5]) 99 | parser.opt_list('--axis', type=str, default='steps', tunable=True, options=['steps', 'both']) 100 | return parser 101 | -------------------------------------------------------------------------------- /tsl/nn/ops/__init__.py: -------------------------------------------------------------------------------- 1 | from .ops import expand_then_cat 2 | -------------------------------------------------------------------------------- /tsl/nn/ops/grad_norm.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class GradNorm(torch.autograd.Function): 5 | """ 6 | Scales the gradient in brackprop. 7 | In the forward pass is an identity operation. 8 | """ 9 | @staticmethod 10 | def forward(ctx, x, norm): 11 | ctx.save_for_backward(x) 12 | ctx.norm = norm # save normalization coefficient 13 | return x # identity 14 | 15 | @staticmethod 16 | def backward(ctx, grad_output): 17 | norm = ctx.norm 18 | return grad_output / norm, None # return the normalized gradient 19 | -------------------------------------------------------------------------------- /tsl/nn/ops/ops.py: -------------------------------------------------------------------------------- 1 | from typing import Union, Tuple, List 2 | 3 | import torch 4 | from torch import nn, Tensor 5 | import numpy as np 6 | 7 | from ..functional import expand_then_cat 8 | 9 | class Lambda(nn.Module): 10 | 11 | def __init__(self, action): 12 | super(Lambda, self).__init__() 13 | self.action = action 14 | 15 | def forward(self, input: Tensor) -> Tensor: 16 | return self.action(input) 17 | 18 | 19 | class Concatenate(nn.Module): 20 | 21 | def __init__(self, dim: int = 0): 22 | super(Concatenate, self).__init__() 23 | self.dim = dim 24 | 25 | def forward(self, tensors: Union[Tuple[Tensor, ...], List[Tensor]]) \ 26 | -> Tensor: 27 | return expand_then_cat(tensors, self.dim) 28 | 29 | 30 | class Select(nn.Module): 31 | """ 32 | Select one element along a dimension. 33 | """ 34 | def __init__(self, dim, index): 35 | super(Select, self).__init__() 36 | self.dim = dim 37 | self.index = index 38 | 39 | def forward(self, tensor: Tensor) \ 40 | -> Tensor: 41 | return tensor.select(self.dim, self.index) 42 | -------------------------------------------------------------------------------- /tsl/nn/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .utils import get_functional_activation, get_layer_activation -------------------------------------------------------------------------------- /tsl/nn/utils/casting.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import torch 3 | from typing import Any, Union, List, Mapping 4 | 5 | 6 | def numpy(tensors: Union[List[torch.Tensor], Mapping[Any, torch.Tensor], torch.Tensor]) \ 7 | -> Union[List[np.ndarray], Mapping[Any, np.ndarray], np.ndarray]: 8 | """ 9 | Cast tensors to numpy arrays. 10 | 11 | Args:' 12 | tensors: A tensor or a list or dictinary containing tensors. 13 | 14 | Returns: 15 | Tensors casted to numpy arrays. 16 | """ 17 | if isinstance(tensors, list): 18 | return [t.detach().cpu().numpy() for t in tensors] 19 | if isinstance(tensors, dict): 20 | for k, v in tensors.items(): 21 | tensors[k] = v.detach().cpu().numpy() 22 | return tensors 23 | if isinstance(tensors, torch.Tensor): 24 | return tensors.detach().cpu().numpy() 25 | raise ValueError 26 | -------------------------------------------------------------------------------- /tsl/nn/utils/utils.py: -------------------------------------------------------------------------------- 1 | from typing import Optional 2 | 3 | import torch 4 | from einops import rearrange 5 | from torch import nn 6 | from torch.nn import functional as F 7 | 8 | from ..ops import expand_then_cat 9 | 10 | _torch_activations_dict = { 11 | 'elu': 'ELU', 12 | 'leaky_relu': 'LeakyReLU', 13 | 'prelu': 'PReLU', 14 | 'relu': 'ReLU', 15 | 'rrelu': 'RReLU', 16 | 'selu': 'SELU', 17 | 'celu': 'CELU', 18 | 'gelu': 'GELU', 19 | 'glu': 'GLU', 20 | 'mish': 'Mish', 21 | 'sigmoid': 'Sigmoid', 22 | 'softplus': 'Softplus', 23 | 'tanh': 'Tanh', 24 | 'silu': 'SiLU', 25 | 'swish': 'SiLU', 26 | 'linear': 'Identity' 27 | } 28 | 29 | 30 | def _identity(x): 31 | return x 32 | 33 | 34 | def get_functional_activation(activation: Optional[str] = None): 35 | if activation is None: 36 | return _identity 37 | activation = activation.lower() 38 | if activation == 'linear': 39 | return _identity 40 | if activation in ['tanh', 'sigmoid']: 41 | return getattr(torch, activation) 42 | if activation in _torch_activations_dict: 43 | return getattr(F, activation) 44 | raise ValueError(f"Activation '{activation}' not valid.") 45 | 46 | 47 | def get_layer_activation(activation: Optional[str] = None): 48 | if activation is None: 49 | return nn.Identity 50 | activation = activation.lower() 51 | if activation in _torch_activations_dict: 52 | return getattr(nn, _torch_activations_dict[activation]) 53 | raise ValueError(f"Activation '{activation}' not valid.") 54 | 55 | 56 | def maybe_cat_exog(x, u, dim=-1): 57 | r""" 58 | Concatenate `x` and `u` if `u` is not `None`. 59 | 60 | We assume `x` to be a 4-dimensional tensor, if `u` has only 3 dimensions we 61 | assume it to be a global exog variable. 62 | 63 | Args: 64 | x: Input 4-d tensor. 65 | u: Optional exogenous variable. 66 | dim (int): Concatenation dimension. 67 | 68 | Returns: 69 | Concatenated `x` and `u`. 70 | """ 71 | if u is not None: 72 | if u.dim() == 3: 73 | u = rearrange(u, 'b s f -> b s 1 f') 74 | x = expand_then_cat([x, u], dim) 75 | return x 76 | -------------------------------------------------------------------------------- /tsl/ops/__init__.py: -------------------------------------------------------------------------------- 1 | from . import * 2 | -------------------------------------------------------------------------------- /tsl/ops/pattern.py: -------------------------------------------------------------------------------- 1 | import re 2 | from collections import Counter 3 | from typing import Iterable, Optional, Union, List 4 | 5 | import torch 6 | from torch import Tensor 7 | 8 | PATTERN_MATCH = re.compile('^(t?){2}(n?){2}[f]*$') 9 | 10 | 11 | def check_pattern(pattern: str, split: bool = False) -> Union[str, list]: 12 | pattern_squeezed = pattern.replace(' ', '').replace('c', 'f') 13 | # check 'c'/'f' follows 'n', 'n' follows 't' 14 | # allow for duplicate 'n' or 't' dims (e.g., 'n n', 't t n f') 15 | # allow for limitless 'c'/'f' dims (e.g., 't n f f') 16 | if not PATTERN_MATCH.match(pattern_squeezed): 17 | raise RuntimeError(f'Pattern "{pattern}" not allowed.') 18 | if split: 19 | return list(pattern_squeezed) 20 | return ' '.join(pattern_squeezed) 21 | 22 | 23 | def outer_pattern(patterns: Iterable[str]): 24 | dims = dict(t=0, n=0, f=0) 25 | for pattern in patterns: 26 | dim_count = Counter(check_pattern(pattern)) 27 | for dim, count in dim_count.items(): 28 | dims[dim] = max(dims[dim], count) 29 | dims = [d for dim, count in dims.items() for d in [dim] * count] 30 | return ' '.join(dims) 31 | 32 | 33 | def broadcast(x, pattern: str, 34 | t: Optional[int] = None, n: Optional[int] = None, 35 | time_index: Union[List, Tensor] = None, 36 | node_index: Union[List, Tensor] = None): 37 | # check patterns 38 | left, rght = pattern.split('->') 39 | left_dims = check_pattern(left, split=True) 40 | rght_dims = check_pattern(rght, split=True) 41 | if not set(left_dims).issubset(rght_dims): 42 | raise RuntimeError(f"Shape {left_dims} cannot be " 43 | f"broadcasted to {rght.strip()}.") 44 | 45 | dim_map = dict(t=t, n=n) 46 | if time_index is not None: 47 | time_index = torch.as_tensor(time_index, dtype=torch.long) 48 | dim_map['t'] = time_index.size(0) 49 | if node_index is not None: 50 | node_index = torch.as_tensor(node_index, dtype=torch.long) 51 | dim_map['n'] = node_index.size(0) 52 | if 't' in rght_dims and 't' not in left_dims and t is None: 53 | raise RuntimeError("Cannot infer dimension for t") 54 | if 'n' in rght_dims and 'n' not in left_dims and n is None: 55 | raise RuntimeError("Cannot infer dimension for n") 56 | 57 | for pos, rght_dim in enumerate(rght_dims): 58 | left_dim = left_dims[pos] 59 | if left_dim != rght_dim: 60 | x = x.unsqueeze(pos) 61 | shape = [dim_map[rght_dim] if i == pos else -1 62 | for i in range(x.ndim)] 63 | x = x.expand(shape) 64 | left_dims.insert(pos, rght_dim) 65 | elif rght_dim == 't' and time_index is not None: 66 | x = x.index_select(pos, time_index) 67 | elif rght_dim == 'n' and node_index is not None: 68 | x = x.index_select(pos, node_index) 69 | return x 70 | -------------------------------------------------------------------------------- /tsl/predictors/__init__.py: -------------------------------------------------------------------------------- 1 | from .base_predictor import Predictor 2 | 3 | predictor_classes = ['Predictor'] 4 | 5 | __all__ = predictor_classes 6 | -------------------------------------------------------------------------------- /tsl/typing.py: -------------------------------------------------------------------------------- 1 | import threading 2 | from typing import Union, Tuple, List, Optional 3 | 4 | from numpy import ndarray 5 | from pandas import DatetimeIndex, PeriodIndex, TimedeltaIndex, DataFrame 6 | from scipy.sparse import coo_matrix, csr_matrix, csc_matrix 7 | from torch import Tensor 8 | from torch_sparse import SparseTensor 9 | 10 | # Tensor = "Tensor" 11 | # SparseTensor = "SparseTensor" 12 | # 13 | # 14 | # def lazy_load_types(): 15 | # from torch import Tensor 16 | # from torch_sparse import SparseTensor 17 | # global Tensor, SparseTensor 18 | # Tensor = Tensor 19 | # SparseTensor = SparseTensor 20 | # 21 | # 22 | # download_thread = threading.Thread(target=lazy_load_types) 23 | # download_thread.start() 24 | 25 | TensArray = Union[Tensor, ndarray] 26 | OptTensArray = Optional[TensArray] 27 | 28 | ScipySparseMatrix = Union[coo_matrix, csr_matrix, csc_matrix] 29 | SparseTensArray = Union[Tensor, SparseTensor, ndarray, ScipySparseMatrix] 30 | OptSparseTensArray = Optional[SparseTensArray] 31 | 32 | FrameArray = Union[DataFrame, ndarray] 33 | OptFrameArray = Optional[FrameArray] 34 | 35 | DataArray = Union[DataFrame, ndarray, Tensor] 36 | OptDataArray = Optional[DataArray] 37 | 38 | TemporalIndex = Union[DatetimeIndex, PeriodIndex, TimedeltaIndex] 39 | 40 | Index = Union[List, Tuple, TensArray] 41 | IndexSlice = Union[slice, Index] 42 | -------------------------------------------------------------------------------- /tsl/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .experiment import TslExperiment 2 | from .download import download_url 3 | from .io import extract_zip 4 | from .parser_utils import ArgParser -------------------------------------------------------------------------------- /tsl/utils/download.py: -------------------------------------------------------------------------------- 1 | import os 2 | import urllib.request 3 | from typing import Optional 4 | 5 | from tqdm import tqdm 6 | 7 | from tsl import logger 8 | 9 | 10 | class DownloadProgressBar(tqdm): 11 | # From https://stackoverflow.com/a/53877507 12 | def update_to(self, b=1, bsize=1, tsize=None): 13 | if tsize is not None: 14 | self.total = tsize 15 | self.update(b * bsize - self.n) 16 | 17 | 18 | def download_url(url: str, folder: str, filename: Optional[str] = None, 19 | log: bool = True): 20 | r"""Downloads the content of an URL to a specific folder. 21 | 22 | Args: 23 | url (string): The url. 24 | folder (string): The folder. 25 | filename (string, optional): The filename. If :obj:`None`, inferred from 26 | url. 27 | log (bool, optional): If :obj:`False`, will not log anything. 28 | (default: :obj:`True`) 29 | """ 30 | if filename is None: 31 | filename = url.rpartition('/')[2].split('?')[0] 32 | path = os.path.join(folder, filename) 33 | 34 | if os.path.exists(path): 35 | if log: 36 | logger.warning(f'Using existing file {filename}') 37 | return path 38 | 39 | if log: 40 | logger.info(f'Downloading {url}') 41 | 42 | os.makedirs(folder, exist_ok=True) 43 | 44 | # From https://stackoverflow.com/a/53877507 45 | with DownloadProgressBar(unit='B', unit_scale=True, 46 | miniters=1, desc=url.split('/')[-1]) as t: 47 | urllib.request.urlretrieve(url, filename=path, reporthook=t.update_to) 48 | return path 49 | -------------------------------------------------------------------------------- /tsl/utils/experiment.py: -------------------------------------------------------------------------------- 1 | import os 2 | import warnings 3 | 4 | import numpy as np 5 | from test_tube import HyperOptArgumentParser as ArgParser 6 | 7 | import tsl 8 | from tsl.utils import parser_utils 9 | from tsl.utils.python_utils import ensure_list 10 | 11 | class TslExperiment: 12 | r""" 13 | Simple class to handle the routines used to run experiments. 14 | 15 | Args: 16 | run_fn: Python function that actually runs the experiment when called. 17 | The run function must accept single argument being the experiment hyperparameters. 18 | parser: Parser used to read the hyperparameters for the experiment. 19 | debug: Whether to run the experiment in debug mode. 20 | config_path: Path to configuration files, if not specified the default will be used. 21 | """ 22 | def __init__(self, run_fn, parser: ArgParser, debug=False, config_path=None): 23 | self.run_fn = run_fn 24 | self.parser = parser 25 | self.debug = debug 26 | self.config_root = config_path if config_path is not None else tsl.config.config_dir 27 | 28 | def _check_config(self, hparams): 29 | config_file = hparams.__dict__.get('config', None) 30 | if config_file is not None: 31 | # read config file 32 | import yaml 33 | 34 | config_file = os.path.join(self.config_root, config_file) 35 | with open(config_file, 'r') as fp: 36 | experiment_config = yaml.load(fp, Loader=yaml.FullLoader) 37 | 38 | # update hparams 39 | hparams = parser_utils.update_from_config(hparams, experiment_config) 40 | if hasattr(self.parser, 'parsed_args'): 41 | self.parser.parsed_args.update(experiment_config) 42 | return hparams 43 | 44 | def make_run_dir(self): 45 | """Create directory to store run logs and artifacts.""" 46 | raise NotImplementedError 47 | 48 | def run(self): 49 | hparams = self.parser.parse_args() 50 | hparams = self._check_config(hparams) 51 | 52 | return self.run_fn(hparams) 53 | 54 | def run_many_times_sequential(self, n): 55 | hparams = self.parser.parse_args() 56 | hparams = self._check_config(hparams) 57 | warnings.warn('Running multiple times. Make sure that randomness is handled properly') 58 | for i in range(n): 59 | print(f"**************Trial n.{i}**************") 60 | np.random.seed() 61 | self.run_fn(hparams) 62 | 63 | def run_search_sequential(self, n): 64 | hparams = self.parser.parse_args() 65 | hparams = self._check_config(hparams) 66 | for i, h in enumerate(hparams.trials(n)): 67 | print(f'**************Trial n.{i}**************') 68 | try: 69 | np.random.seed() 70 | self.run_fn(h) 71 | except RuntimeError as err: 72 | print(f'Trial n. {i} failed due to a Runtime error: {err}') 73 | 74 | def run_search_parallel(self, n, workers, gpus=None): 75 | hparams = self.parser.parse_args() 76 | hparams = self._check_config(hparams) 77 | if gpus is None: 78 | hparams.optimize_parallel_cpu(self.run_fn, nb_trials=n, 79 | nb_workers=workers) 80 | else: 81 | gpus = ensure_list(gpus) 82 | hparams.optimize_parallel_gpu(self.run_fn, max_nb_trials=n, 83 | gpu_ids=gpus) 84 | -------------------------------------------------------------------------------- /tsl/utils/io.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import zipfile 4 | from typing import Any 5 | 6 | from tsl import logger 7 | 8 | 9 | def extract_zip(path: str, folder: str, log: bool = True): 10 | r"""Extracts a zip archive to a specific folder. 11 | 12 | Args: 13 | path (string): The path to the zip archive. 14 | folder (string): The folder. 15 | log (bool, optional): If :obj:`False`, will not log anything. 16 | (default: :obj:`True`) 17 | """ 18 | if log: 19 | logger.info(f"Extracting {path}") 20 | with zipfile.ZipFile(path, 'r') as f: 21 | f.extractall(folder) 22 | 23 | 24 | def save_pickle(obj: Any, filename: str) -> str: 25 | """Save obj to path as pickle. 26 | 27 | Args: 28 | obj: Object to be saved. 29 | filename (string): Where to save the file. 30 | 31 | Returns: 32 | path (string): The absolute path to the saved pickle 33 | """ 34 | abspath = os.path.abspath(filename) 35 | directory = os.path.dirname(abspath) 36 | os.makedirs(directory, exist_ok=True) 37 | with open(abspath, 'wb') as fp: 38 | pickle.dump(obj, fp) 39 | return abspath 40 | 41 | 42 | def load_pickle(filename: str) -> Any: 43 | """Load object from pickle filename. 44 | 45 | Args: 46 | filename (string): The absolute path to the saved pickle. 47 | 48 | Returns: 49 | data (any): The loaded object. 50 | """ 51 | with open(filename, 'rb') as fp: 52 | data = pickle.load(fp) 53 | return data 54 | 55 | 56 | def save_figure(fig, filename: str, as_html=False, as_pickle=False): 57 | if filename.endswith('html'): 58 | as_html = True 59 | filename = filename[:-5] 60 | elif filename.endswith('pkl'): 61 | as_pickle = True 62 | filename = filename[:-4] 63 | if not (as_html or as_pickle): 64 | as_html = False # save as html if nothing is specified 65 | if as_html: 66 | import mpld3 67 | with open(filename + '.html', 'w') as fp: 68 | mpld3.save_html(fig, fp) 69 | if as_pickle: 70 | import pickle 71 | with open(filename + '.pkl', 'wb') as fp: 72 | pickle.dump(fig, fp) -------------------------------------------------------------------------------- /tsl/utils/numpy_metrics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import tsl 4 | 5 | 6 | def mae(y_hat, y): 7 | return np.abs(y_hat - y).mean() 8 | 9 | 10 | def nmae(y_hat, y): 11 | delta = np.max(y) - np.min(y) + tsl.epsilon 12 | return mae(y_hat, y) * 100 / delta 13 | 14 | 15 | def mape(y_hat, y): 16 | return 100 * np.abs((y_hat - y) / (y + tsl.epsilon)).mean() 17 | 18 | 19 | def mse(y_hat, y): 20 | return np.square(y_hat - y).mean() 21 | 22 | 23 | def rmse(y_hat, y): 24 | return np.sqrt(mse(y_hat, y)) 25 | 26 | 27 | def nrmse(y_hat, y): 28 | delta = np.max(y) - np.min(y) + tsl.epsilon 29 | return rmse(y_hat, y) * 100 / delta 30 | 31 | 32 | def nrmse_2(y_hat, y): 33 | nrmse_ = np.sqrt(np.square(y_hat - y).sum() / np.square(y).sum()) 34 | return nrmse_ * 100 35 | 36 | 37 | def r2(y_hat, y): 38 | return 1. - np.square(y_hat - y).sum() / (np.square(y.mean(0) - y).sum()) 39 | 40 | 41 | def masked_mae(y_hat, y, mask=None): 42 | if mask is None: 43 | mask = slice(None) 44 | else: 45 | mask = np.asarray(mask, dtype=bool) 46 | err = y_hat[mask] - y[mask] 47 | return np.abs(err).mean() 48 | 49 | 50 | def masked_mape(y_hat, y, mask=None): 51 | if mask is None: 52 | mask = slice(None) 53 | else: 54 | mask = np.asarray(mask, dtype=bool) 55 | err = (y_hat[mask] - y[mask]) / (y[mask] + tsl.epsilon) 56 | return np.abs(err).mean() 57 | 58 | 59 | def masked_mse(y_hat, y, mask=None): 60 | if mask is None: 61 | mask = slice(None) 62 | else: 63 | mask = np.asarray(mask, dtype=bool) 64 | err = y_hat[mask] - y[mask] 65 | return np.square(err).mean() 66 | 67 | 68 | def masked_rmse(y_hat, y, mask=None): 69 | if mask is None: 70 | mask = slice(None) 71 | else: 72 | mask = np.asarray(mask, dtype=bool) 73 | err = np.square(y_hat[mask] - y[mask]) 74 | return np.sqrt(err.mean()) 75 | 76 | 77 | def masked_mre(y_hat, y, mask=None): 78 | if mask is None: 79 | mask = slice(None) 80 | else: 81 | mask = np.asarray(mask, dtype=bool) 82 | err = np.abs(y_hat[mask] - y[mask]) 83 | return err.sum() / (y[mask].sum() + tsl.epsilon) 84 | -------------------------------------------------------------------------------- /tsl/utils/preprocessing.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import pandas as pd 3 | 4 | import tsl 5 | 6 | 7 | def aggregate(dataframe, idx=None, type='sum'): 8 | # todo: make it work with multindex dataframes 9 | aggregation_fn = getattr(np, type) 10 | if idx is None: 11 | aggr = aggregation_fn(dataframe.values, axis=1) 12 | return pd.DataFrame(aggr, index=dataframe.index, columns=['seq']) 13 | 14 | else: 15 | ids = np.unique(idx) 16 | aggregates = [] 17 | x = dataframe.values.T 18 | for i in ids: 19 | aggregates.append(aggregation_fn(x[idx == i], axis=0)) 20 | 21 | cols = ['seq' + '_' + str(i) for i in ids] 22 | return pd.DataFrame(dict(zip(cols, aggregates)), index=dataframe.index) 23 | 24 | 25 | def get_trend(df, period='week', train_len=None, valid_mask=None): 26 | """ 27 | Perform detrending on a time series by subtrating from each value of the input dataframe 28 | the average value computed over the training dataset for each hour/weekday 29 | :param df: dataframe 30 | :param period: period of the trend ('day', 'week', 'month') 31 | :param train_len: train length, 32 | :return: 33 | - the detrended datasets 34 | - the trend values that has to be added back after computing the prediction 35 | """ 36 | df = df.copy() 37 | if train_len is not None: 38 | df[train_len:] = np.nan 39 | if valid_mask is not None: 40 | df[~valid_mask] = np.nan 41 | idx = [df.index.hour, df.index.minute] 42 | if period == 'week': 43 | idx = [df.index.weekday, ] + idx 44 | elif period == 'month': 45 | idx = [df.index.month, df.index.weekday] + idx 46 | elif period != 'day': 47 | raise NotImplementedError("Period must be in ('day', 'week', 'month')") 48 | 49 | means = df.groupby(idx).transform(np.nanmean) 50 | return means 51 | 52 | 53 | def normalize_by_group(df, by): 54 | """ 55 | Normalizes a dataframe using mean and std of a specified group. 56 | 57 | :param df: the data 58 | :param by: used to determine the groups for the groupby 59 | :return: the normalized df 60 | """ 61 | groups = df.groupby(by) 62 | # computes group-wise mean/std, 63 | # then auto broadcasts to size of group chunk 64 | mean = groups.transform(np.nanmean) 65 | std = groups.transform(np.nanstd) + tsl.epsilon # add epsilon to avoid division by zero 66 | return (df[mean.columns] - mean) / std 67 | -------------------------------------------------------------------------------- /tsl/utils/python_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Any, Sequence, List 3 | 4 | 5 | def ensure_list(value: Any) -> List: 6 | # if isinstance(value, Sequence) and not isinstance(value, str): 7 | if hasattr(value, '__iter__') and not isinstance(value, str): 8 | return list(value) 9 | else: 10 | return [value] 11 | 12 | 13 | def files_exist(files: Sequence[str]) -> bool: 14 | files = ensure_list(files) 15 | return len(files) != 0 and all([os.path.exists(f) for f in files]) 16 | 17 | 18 | def hash_dict(obj: dict): 19 | from hashlib import md5 20 | obj = {k: obj[k] for k in sorted(obj)} 21 | return md5(str(obj).encode()).hexdigest() 22 | 23 | 24 | def set_property(obj, name, prop_function): 25 | """Add property :obj:`prop_function` to :obj:`obj`. 26 | 27 | :obj:`prop_function` must be a function taking only one argument, i.e., 28 | :obj:`obj`. 29 | 30 | Args: 31 | obj (object): object on which the property has to be added. 32 | name (str): the name of the property. 33 | prop_function (function): function taking only :obj:`obj` as argument. 34 | """ 35 | 36 | class_name = obj.__class__.__name__ 37 | new_class = type(class_name, (obj.__class__,), 38 | {name: property(prop_function)}) 39 | obj.__class__ = new_class 40 | --------------------------------------------------------------------------------