CoST_pipeline.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "fa20273c28f14a6fb408d382d4ac98e6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_81488352900249259e4c0f3f116fa08c",
              "IPY_MODEL_ff494d4ed0464ef48bd780aa333f8902",
              "IPY_MODEL_ca08c18f3af64124bc1cf2bc45121c1e"
            ],
            "layout": "IPY_MODEL_f2bf006d8412490f8cea31c6a589ece2"
          }
        },
        "81488352900249259e4c0f3f116fa08c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7b8bf6d3ed804d75b2ea50494048bd37",
            "placeholder": "​",
            "style": "IPY_MODEL_a324c27f1e2d4c5b8b547899a50aa13f",
            "value": "data encoding: 100%"
          }
        },
        "ff494d4ed0464ef48bd780aa333f8902": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_eca61e01d24144b48c2eff401765b1c8",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_4b1136aa31cd45c38d3c55b2fed75a46",
            "value": 1
          }
        },
        "ca08c18f3af64124bc1cf2bc45121c1e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_479c27d780b84af7af586d7ad754086c",
            "placeholder": "​",
            "style": "IPY_MODEL_4453bcf9179b4eb18d21750e3717ca7a",
            "value": " 1/1 [23:18&lt;00:00, 1398.28s/it]"
          }
        },
        "f2bf006d8412490f8cea31c6a589ece2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7b8bf6d3ed804d75b2ea50494048bd37": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a324c27f1e2d4c5b8b547899a50aa13f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "eca61e01d24144b48c2eff401765b1c8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4b1136aa31cd45c38d3c55b2fed75a46": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "479c27d780b84af7af586d7ad754086c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4453bcf9179b4eb18d21750e3717ca7a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "43eb6a5cb7f44ee89a72abf0cd5c9723": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_e66a5509d30b43ff80d0b0af3c6f9c8d",
              "IPY_MODEL_00f66ad6e8b34fe19c1a1bf65159451c",
              "IPY_MODEL_c9f3af47ffb84c259ffc633e554c2e0f"
            ],
            "layout": "IPY_MODEL_14789f347e174395ac6f314a38afb721"
          }
        },
        "e66a5509d30b43ff80d0b0af3c6f9c8d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_c5da477f433f40349751688dfeddfb8e",
            "placeholder": "​",
            "style": "IPY_MODEL_0be77f4742cc419991f65c85e173c02d",
            "value": "sequence encoding: 100%"
          }
        },
        "00f66ad6e8b34fe19c1a1bf65159451c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_555016b5553446a595d0d582f750119b",
            "max": 69680,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_77126262e4804dd293c348a73451c92d",
            "value": 69680
          }
        },
        "c9f3af47ffb84c259ffc633e554c2e0f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4b5b599c1ef649c29c48c41dc5ca86e8",
            "placeholder": "​",
            "style": "IPY_MODEL_cbdc4be27e694ef49e4f02de28414625",
            "value": " 69680/69680 [23:17&lt;00:00, 48.36it/s]"
          }
        },
        "14789f347e174395ac6f314a38afb721": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "c5da477f433f40349751688dfeddfb8e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "0be77f4742cc419991f65c85e173c02d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "555016b5553446a595d0d582f750119b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "77126262e4804dd293c348a73451c92d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "4b5b599c1ef649c29c48c41dc5ca86e8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "cbdc4be27e694ef49e4f02de28414625": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "17e0c718123245e0894cb62237f054ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_0cb3cb25168c423390185db0b39538ba",
              "IPY_MODEL_5989a8ecbcc34145965922753d3d6648",
              "IPY_MODEL_ec203b01b7f34051913482e8851b24e2"
            ],
            "layout": "IPY_MODEL_cda8bdf2d8da435aba7fcde8e368b4d9"
          }
        },
        "0cb3cb25168c423390185db0b39538ba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_5fa3dd57be684025806c95921a551037",
            "placeholder": "​",
            "style": "IPY_MODEL_7261b4f95a91428f901d284efc03a2ce",
            "value": "forecasting evaluation [24, 48, 96, 288, 672]: 100%"
          }
        },
        "5989a8ecbcc34145965922753d3d6648": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_720b8aba5a1745428a67f994a5fd6051",
            "max": 5,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a87cec2cb96d495ab0fc7ab68e12bc78",
            "value": 5
          }
        },
        "ec203b01b7f34051913482e8851b24e2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7a849ebf6c754ca89fa554d756e69470",
            "placeholder": "​",
            "style": "IPY_MODEL_14443025f49d43ceba6cf2d0d5b71bd2",
            "value": " 5/5 [01:36&lt;00:00, 27.11s/it]"
          }
        },
        "cda8bdf2d8da435aba7fcde8e368b4d9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "5fa3dd57be684025806c95921a551037": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "7261b4f95a91428f901d284efc03a2ce": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "720b8aba5a1745428a67f994a5fd6051": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a87cec2cb96d495ab0fc7ab68e12bc78": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7a849ebf6c754ca89fa554d756e69470": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "14443025f49d43ceba6cf2d0d5b71bd2": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/whoami-Lory271/thesis-project/blob/main/CoST_pipeline.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Symbol legend\n",
        "\n",
        "* B: batch size\n",
        "* L: lookback window (aka input_size)\n"
      ],
      "metadata": {
        "id": "7s9odzFFQWyj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Installations and imports\n"
      ],
      "metadata": {
        "id": "w7opc0NsjlNx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install pytorch-lightning==2.0.1.post0 --quiet\n",
        "!pip install einops==0.6.1 --quiet\n",
        "!pip install ipdb --quiet\n",
        "!pip install wandb --quiet\n",
        "# !pip install objsize --quiet"
      ],
      "metadata": {
        "id": "ehQC2AKyci-4",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1976a271-3679-48ef-89bb-9ddd8b04b9d9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m718.6/718.6 kB\u001b[0m \u001b[31m9.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m805.2/805.2 kB\u001b[0m \u001b[31m46.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.2/42.2 kB\u001b[0m \u001b[31m1.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m18.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.1/2.1 MB\u001b[0m \u001b[31m29.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.0/190.0 kB\u001b[0m \u001b[31m23.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m224.8/224.8 kB\u001b[0m \u001b[31m26.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 kB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "import logging\n",
        "# https://github.com/gotcha/ipdb\n",
        "import ipdb\n",
        "import copy\n",
        "from google.colab import drive\n",
        "from tqdm.notebook import tqdm\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.utils.data import TensorDataset\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.profiler import profile, record_function, ProfilerActivity\n",
        "import numpy as np\n",
        "import random\n",
        "from collections import defaultdict\n",
        "import functools\n",
        "# https://theaisummer.com/einsum-attention/\n",
        "import einops\n",
        "import math\n",
        "from sklearn.model_selection import train_test_split\n",
        "from sklearn.linear_model import Ridge\n",
        "# https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.StandardScaler.html\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "import pytorch_lightning as pl\n",
        "from pytorch_lightning.loggers import TensorBoardLogger\n",
        "from pytorch_lightning.loggers import WandbLogger\n",
        "from pytorch_lightning.callbacks import ModelCheckpoint\n",
        "from pytorch_lightning.callbacks.progress import TQDMProgressBar\n",
        "from pytorch_lightning.callbacks import EarlyStopping\n",
        "from pytorch_lightning.profilers import PyTorchProfiler\n",
        "from pytorch_lightning.tuner import Tuner\n",
        "\n",
        "import wandb\n",
        "import sys\n",
        "import pickle\n",
        "import time\n",
        "import os\n",
        "import torch.fft as fft\n",
        "# import objsize"
      ],
      "metadata": {
        "id": "WuaX4Ts_jqmd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "drive.mount('/content/drive')\n",
        "wandb.login()"
      ],
      "metadata": {
        "id": "4F86VEC4VtL8",
        "outputId": "90443e4b-dfbf-4856-a575-3a7a847fe7e2",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mdesantis-1849114\u001b[0m (\u001b[33mdesantis\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Helper Functions"
      ],
      "metadata": {
        "id": "J2UVU6VqgizJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# setup logger function\n",
        "def setup_log(self, level):\n",
        "    log = logging.getLogger(self.__class__.__name__)\n",
        "    log.setLevel(level)\n",
        "    return log"
      ],
      "metadata": {
        "id": "gkBTe3nko46m"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# write log on a file\n",
        "def write_log(a):\n",
        "    with open(\"log.txt\", 'w') as file:\n",
        "        for row in a:\n",
        "            file.write(str(row))\n",
        "        log.debug(\"object logged\")"
      ],
      "metadata": {
        "id": "i8LE_3TogiZn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def get_patch_num(patch_len, num_var, stride):\n",
        "    return (max(patch_len, num_var)-patch_len) // stride + 2"
      ],
      "metadata": {
        "id": "Ke58BUkLo-dd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# pad array with nan\n",
        "def pad_nan(arr, left=0, right=0, dim=0):\n",
        "    # padding the right side\n",
        "    if left > 0:\n",
        "        padshape = list(arr.shape)\n",
        "        padshape[dim] = left\n",
        "        arr = torch.cat((torch.full(padshape, np.nan).to(arr.device), arr), dim=dim)\n",
        "\n",
        "    # padding the left side\n",
        "    if right > 0:\n",
        "        padshape = list(arr.shape)\n",
        "        padshape[dim] = right\n",
        "        arr = torch.cat((arr, torch.full(padshape, np.nan)).to(arr.device), dim=dim)\n",
        "    return arr"
      ],
      "metadata": {
        "id": "T_LjdLhhiyGm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# split with nan\n",
        "def split_with_nan(x, sections, axis=0):\n",
        "    assert x.dtype in [np.float16, np.float32, np.float64]\n",
        "    arrs = np.array_split(x, sections, axis=axis)\n",
        "    target_length = arrs[0].shape[axis]\n",
        "    for i in range(len(arrs)):\n",
        "        arrs[i] = pad_nan_to_target(arrs[i], target_length, axis=axis)\n",
        "    return arrs"
      ],
      "metadata": {
        "id": "8Z2FotiSArtJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# pad with nan\n",
        "def pad_nan_to_target(array, target_length, axis=0):\n",
        "    assert array.dtype in [np.float16, np.float32, np.float64]\n",
        "    pad_size = target_length - array.shape[axis]\n",
        "    if pad_size <= 0:\n",
        "        return array\n",
        "    npad = [(0, 0)] * array.ndim\n",
        "    npad[axis] = (0, pad_size)\n",
        "\n",
        "    return np.pad(array, pad_width=npad, mode='constant', constant_values=np.nan)"
      ],
      "metadata": {
        "id": "1SAjsR_-A3RB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def centerize_vary_length_series(x):\n",
        "    prefix_zeros = np.argmax(~np.isnan(x).all(axis=-1), axis=1)\n",
        "    suffix_zeros = np.argmax(~np.isnan(x[:, ::-1]).all(axis=-1), axis=1)\n",
        "    offset = (prefix_zeros + suffix_zeros) // 2 - prefix_zeros\n",
        "    rows, column_indices = np.ogrid[:x.shape[0], :x.shape[1]]\n",
        "    offset[offset < 0] += x.shape[1]\n",
        "    column_indices = column_indices - offset[:, np.newaxis]\n",
        "\n",
        "    return x[rows, column_indices]"
      ],
      "metadata": {
        "id": "eK2M1qV3JBdl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# random seed functions\n",
        "def seed_everything(seed):\n",
        "    pl.seed_everything(seed, workers=True)\n",
        "    random.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.cuda.manual_seed_all(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = False"
      ],
      "metadata": {
        "id": "mzBQAeoIi8l2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# pickle helper\n",
        "def pkl_save(name, var):\n",
        "    os.makedirs(os.path.dirname(name), exist_ok=True)\n",
        "    with open(name, 'wb') as f:\n",
        "        pickle.dump(var, f)\n",
        "\n",
        "def pkl_load(name):\n",
        "    with open(name, 'rb') as f:\n",
        "        return pickle.load(f)"
      ],
      "metadata": {
        "id": "NxBH76S1gdKm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Constants"
      ],
      "metadata": {
        "id": "94xSbfaKkOmq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# logger\n",
        "LOG_LEVEL = logging.DEBUG\n",
        "\n",
        "# datasets name\n",
        "ELECTRICITY = \"electricity\"\n",
        "M5 = \"M5\"\n",
        "ETTh1 = \"ETTh1\"\n",
        "ETTh2 = \"ETTh2\"\n",
        "ETTm1 = \"ETTm1\"\n",
        "ETTm2 = \"ETTm2\"\n",
        "WEATHER = \"WTH\"\n",
        "\n",
        "# backbones\n",
        "\n",
        "Pyraformer = \"Pyraformer\"\n",
        "TCN = \"TCN\"\n",
        "NLinear = \"NLinear\"\n",
        "Linear = \"Linear\"\n",
        "\n",
        "\n",
        "# models\n",
        "CoST = \"CoST\"\n",
        "### TODO ###\n",
        "# choose the backbone (encoder) to use and the dataset to set as benchmark\n",
        "BACKBONE = \"\"\n",
        "MODEL = CoST + '-' + BACKBONE\n",
        "DATASET = \"\"\n",
        "###########\n",
        "\n",
        "# device\n",
        "\n",
        "DEVICE = torch.device('cpu') if not torch.cuda.is_available() else torch.device('cuda')\n",
        "\n",
        "#hyperparameters\n",
        "\n",
        "LR = 1e-3 #1e-3 #0.04365158322401657\n",
        "\n",
        "# Train\n",
        "BATCH_SIZE = 128\n",
        "EPOCHS = None\n",
        "ITERS = None\n",
        "L = 201\n",
        "UNIVARIATE = False\n",
        "\n",
        "# Eval\n",
        "EVALUATE = True\n",
        "MAX_TRAIN_LENGTH = 201\n",
        "PADDING = MAX_TRAIN_LENGTH-1\n",
        "ENCODE_BATCH_SIZE = 256\n",
        "\n",
        "\n",
        "# training\n",
        "TRAIN = True\n",
        "DETERMINISTIC = True\n",
        "LOAD_MODEL = False\n",
        "RESUME_TRAINING = False\n",
        "MEMORY_PROFILING = False\n",
        "LOAD_ENCODE = False\n",
        "\n",
        "# training techniques\n",
        "\n",
        "GRADIENT_CLIPPING = None # None is default, 0.5 is by norm\n",
        "FIND_LR = False\n",
        "\n",
        "# wandb\n",
        "\n",
        "SETTINGS_STRING = \"univariate\" if UNIVARIATE else \"multivariate\"\n",
        "RUN_ID = \"n4tagkey\"\n",
        "RESUME_RUN = False\n",
        "RUN_ID = wandb.util.generate_id() if not RESUME_RUN else RUN_ID\n",
        "print(f\"current RUN_ID is {RUN_ID}\")\n",
        "MODEL_ID = MODEL+'_'+DATASET+'_'+SETTINGS_STRING\n",
        "\n",
        "\n",
        "#paths\n",
        "MODEL_SETTINGS_FOLDER = \"/univariate\" if UNIVARIATE else \"/multivariate\"\n",
        "### TODO ###\n",
        "ROOT_FOLDER = \"\" # insert the root folder\n",
        "###########\n",
        "MODEL_FOLDER = ROOT_FOLDER + \"/models\"\n",
        "CHECKPOINT_FOLDER = ROOT_FOLDER + \"/checkpoints\" + \"/\" + MODEL + \"/\" + DATASET + MODEL_SETTINGS_FOLDER\n",
        "# LOGS_FOLDER = ROOT_FOLDER + \"/logs\"\n",
        "ENCODING_FOLDER = ROOT_FOLDER + \"/encoding/\" + MODEL + \"/\" + DATASET + MODEL_SETTINGS_FOLDER\n",
        "FORECASTING_RESULT = ROOT_FOLDER + \"/forecasting_result/\" + MODEL + \"/\" + DATASET + MODEL_SETTINGS_FOLDER\n",
        "datasets_path = {\n",
        "    ELECTRICITY: ROOT_FOLDER + \"/datasets/electricity\",\n",
        "    ETTh1: ROOT_FOLDER + \"/datasets/ETT\",\n",
        "    ETTh2: ROOT_FOLDER + \"/datasets/ETT\",\n",
        "    ETTm1: ROOT_FOLDER + \"/datasets/ETT\",\n",
        "    ETTm2: ROOT_FOLDER + \"/datasets/ETT\"\n",
        "}\n",
        "\n",
        "datasets_name = {\n",
        "    ELECTRICITY: \"/LD2011_2014.txt\"\n",
        "}\n",
        "datasets_processed_name = {\n",
        "    ELECTRICITY: \"/electricity.csv\",\n",
        "    ETTh1: \"/ETTh1.csv\",\n",
        "    ETTh2: \"/ETTh2.csv\",\n",
        "    ETTm1: \"/ETTm1.csv\",\n",
        "    ETTm2: \"/ETTm2.csv\"\n",
        "}\n",
        "\n",
        "datasets_pred_lens = {\n",
        "    ELECTRICITY: [24, 48, 168, 336, 720], # [24, 48, 168, 336, 720]\n",
        "    ETTh1: [24, 48, 168, 336, 720],\n",
        "    ETTh2: [24, 48, 168, 336, 720],\n",
        "    WEATHER: [24, 48, 168, 336, 720],\n",
        "    M5: [28],\n",
        "    ETTm1: [24, 48, 96, 288, 672],\n",
        "    ETTm2: [24, 48, 96, 288, 672]\n",
        "}\n",
        "\n",
        "datasets_n_iters = {\n",
        "    ELECTRICITY: 600,\n",
        "    ETTh1: 200 if UNIVARIATE else 600,\n",
        "    ETTh2: 200 if UNIVARIATE else 600,\n",
        "    ETTm1: 200 if UNIVARIATE else 600,\n",
        "    ETTm2: 200 if UNIVARIATE else 600\n",
        "}\n",
        "\n",
        "input_dim = {\n",
        "    ELECTRICITY: 8,\n",
        "    ETTh1: 14,\n",
        "    ETTh2: 14,\n",
        "    ETTm1: 14,\n",
        "    ETTm2: 14\n",
        "}\n",
        "\n",
        "ITERS = datasets_n_iters[DATASET] if ITERS is None and EPOCHS is None else ITERS\n",
        "\n",
        "config = dict(\n",
        "    iters = ITERS,\n",
        "    epochs = 0 if EPOCHS is None else EPOCHS,\n",
        "    batch_size=BATCH_SIZE,\n",
        "    learning_rate = LR,\n",
        "    backbone = BACKBONE,\n",
        "    dataset=DATASET,\n",
        "    architecture=MODEL,\n",
        "    model_id = MODEL_ID,\n",
        "    run_id = RUN_ID)"
      ],
      "metadata": {
        "id": "ySxfaOHQkQ_B",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "27c3afd6-cfc5-4a1a-d524-90c3a9ff4e74"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "current RUN_ID is nt30vazb\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# WanDB"
      ],
      "metadata": {
        "id": "WXYPOAZ8O2zL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# start a new experiment\n",
        "run = wandb.init(project=MODEL_ID, config = config, id = config[\"run_id\"], resume = 'allow')\n",
        "ENCODING_FOLDER += \"/\" + run.name\n",
        "FORECASTING_RESULT += \"/\" + run.name\n",
        "CHECKPOINT_FOLDER += \"/\" + run.name"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 104
        },
        "id": "sC_id8hdO5i_",
        "outputId": "b0e4654f-7cd8-4172-d497-09975270960c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Tracking run with wandb version 0.15.12"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Run data is saved locally in <code>/content/wandb/run-20231007_141828-nt30vazb</code>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              "Syncing run <strong><a href='https://wandb.ai/desantis/CoST-TCN_ETTm2_multivariate/runs/nt30vazb' target=\"_blank\">denim-shadow-1</a></strong> to <a href='https://wandb.ai/desantis/CoST-TCN_ETTm2_multivariate' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              " View project at <a href='https://wandb.ai/desantis/CoST-TCN_ETTm2_multivariate' target=\"_blank\">https://wandb.ai/desantis/CoST-TCN_ETTm2_multivariate</a>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ],
            "text/html": [
              " View run at <a href='https://wandb.ai/desantis/CoST-TCN_ETTm2_multivariate/runs/nt30vazb' target=\"_blank\">https://wandb.ai/desantis/CoST-TCN_ETTm2_multivariate/runs/nt30vazb</a>"
            ]
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Logger"
      ],
      "metadata": {
        "id": "8j5-8dn5ppGo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# create logger\n",
        "log = logging.getLogger('APP')\n",
        "log.setLevel(LOG_LEVEL)\n",
        "logging.basicConfig(level=LOG_LEVEL)\n",
        "\n",
        "# # create console handler and set level to debug\n",
        "# ch = logging.StreamHandler()\n",
        "# ch.setLevel(logging.INFO)\n",
        "\n",
        "# # create formatter\n",
        "# formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')\n",
        "\n",
        "# # add formatter to ch\n",
        "# ch.setFormatter(formatter)\n",
        "\n",
        "# # add ch to logger\n",
        "# logger.addHandler(ch)"
      ],
      "metadata": {
        "id": "iRLWiTu4mlx9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Dataset and Datamodule"
      ],
      "metadata": {
        "id": "ma655OWbiZ0h"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class CustomDataset(Dataset):\n",
        "    def __init__(self, data, sigma, eval_mode = False,p = 0.5, multiplier = 10):\n",
        "        super().__init__()\n",
        "        self.data = data\n",
        "        self.sigma = sigma\n",
        "        self.eval_mode = eval_mode\n",
        "        self.p = p\n",
        "        self.multiplier = multiplier\n",
        "        self.N, self.T, self.D = data.shape # num_ts, time, dim\n",
        "\n",
        "    def __len__(self):\n",
        "        # return self.data.shape[0] // self.look_window\n",
        "        return self.data.shape[0] * self.multiplier\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        ts = self.data[idx % self.N]\n",
        "\n",
        "        return self.transform(ts), self.transform(ts)\n",
        "\n",
        "    def get_len(self):\n",
        "        return self.__len__()\n",
        "\n",
        "    # def get_channels(self):\n",
        "    #     return self.data.iloc[0, 1:].astype(str).str.replace(',', '.').astype('float32').shape[0]\n",
        "\n",
        "    def transform(self, x):\n",
        "        return self.jitter(self.shift(self.scale(x)))\n",
        "\n",
        "    def jitter(self, x):\n",
        "        if random.random() > self.p:\n",
        "            return x\n",
        "        return x + (torch.empty(x.shape).normal_(mean = 0, std = 0.5) * self.sigma)\n",
        "\n",
        "    def scale(self, x):\n",
        "        if random.random() > self.p:\n",
        "            return x\n",
        "        return x * (torch.empty(x.size(-1)).normal_(mean = 0, std = 0.5) * self.sigma + 1)\n",
        "\n",
        "    def shift(self, x):\n",
        "        if random.random() > self.p:\n",
        "            return x\n",
        "        return x + (torch.empty(x.size(-1)).normal_(mean = 0, std = 0.5) * self.sigma)\n",
        "\n",
        "class CustomDataModule(pl.LightningDataModule):\n",
        "    def __init__(self, name, path, batch_size, max_train_length,\n",
        "                 encode_batch_size = 256, train_size = 0.6, test_size = 0.2, univariate = False):\n",
        "        super().__init__()\n",
        "        self.path = path # path to csv file\n",
        "        self.batch_size = batch_size\n",
        "        self.encode_batch_size = encode_batch_size\n",
        "        self.max_train_length = max_train_length\n",
        "        self.padding = max_train_length-1\n",
        "        self.train_size = train_size\n",
        "        self.test_size = test_size\n",
        "        self.univariate = univariate\n",
        "        self.name = name\n",
        "        # self.data = np.load(path)\n",
        "\n",
        "    # def prepare_data(self):\n",
        "    #     # download\n",
        "\n",
        "    def _get_time_features(self,dt):\n",
        "        return np.stack([\n",
        "                dt.minute.to_numpy(),\n",
        "                dt.hour.to_numpy(),\n",
        "                dt.dayofweek.to_numpy(),\n",
        "                dt.day.to_numpy(),\n",
        "                dt.dayofyear.to_numpy(),\n",
        "                dt.month.to_numpy(),\n",
        "                dt.weekofyear.to_numpy(),\n",
        "            ], axis=1).astype(float)\n",
        "\n",
        "    def _load_forecast_csv(self, path, name):\n",
        "        data = pd.read_csv(path, index_col='date', parse_dates=True)\n",
        "        dt_embed = self._get_time_features(data.index)\n",
        "        n_covariate_cols = dt_embed.shape[-1]\n",
        "        if self.univariate:\n",
        "            if name in ('ETTh1', 'ETTh2', 'ETTm1', 'ETTm2'):\n",
        "                data = data[['OT']]\n",
        "            elif name == 'electricity':\n",
        "                data = data[['MT_001']]\n",
        "            elif name == 'WTH':\n",
        "                data = data[['WetBulbCelsius']]\n",
        "            else:\n",
        "                data = data.iloc[:, -1:]\n",
        "        data = data.to_numpy()\n",
        "        # compute slices\n",
        "        if name == 'ETTh1' or name == 'ETTh2':\n",
        "            self.train_slice = slice(None, 12 * 30 * 24)\n",
        "            self.valid_slice = slice(12 * 30 * 24, 16 * 30 * 24)\n",
        "            self.test_slice = slice(16 * 30 * 24, 20 * 30 * 24)\n",
        "        elif name == 'ETTm1' or name == 'ETTm2':\n",
        "            self.train_slice = slice(None, 12 * 30 * 24 * 4)\n",
        "            self.valid_slice = slice(12 * 30 * 24 * 4, 16 * 30 * 24 * 4)\n",
        "            self.test_slice = slice(16 * 30 * 24 * 4, 20 * 30 * 24 * 4)\n",
        "        elif name.startswith('M5'):\n",
        "            self.train_slice = slice(None, int(0.8 * (1913 + 28)))\n",
        "            self.valid_slice = slice(int(0.8 * (1913 + 28)), 1913 + 28)\n",
        "            self.test_slice = slice(1913 + 28 - 1, 1913 + 2 * 28)\n",
        "        else:\n",
        "            self.train_slice = slice(None, int(self.train_size * len(data)))\n",
        "            self.valid_slice = slice(int(self.train_size * len(data)), - int(self.test_size * len(data)))\n",
        "            self.test_slice = slice(- int(self.test_size * len(data)), None)\n",
        "\n",
        "        return data, dt_embed, n_covariate_cols\n",
        "\n",
        "    def _scale_and_transform(self, data, dt_embed, n_covariate_cols, name):\n",
        "        # scale data\n",
        "        scaler = StandardScaler().fit(data[self.train_slice])\n",
        "        data = scaler.transform(data)\n",
        "        # NUM_ROWS x NUM_FEATURES -> NUM_FEATURES x NUM_ROWS x 1\n",
        "        if name in ('electricity') or name.startswith('M5'):\n",
        "            data = np.expand_dims(data.T, -1)  # Each variable is an instance rather than a feature\n",
        "        else:\n",
        "            data = np.expand_dims(data, 0)\n",
        "        if n_covariate_cols > 0:\n",
        "            dt_scaler = StandardScaler().fit(dt_embed[self.train_slice])\n",
        "            dt_embed = np.expand_dims(dt_scaler.transform(dt_embed), 0)\n",
        "            data = np.concatenate([np.repeat(dt_embed, data.shape[0], axis=0), data], axis=-1)\n",
        "\n",
        "        return data, scaler, n_covariate_cols\n",
        "\n",
        "    def _fit_setup(self, train_data):\n",
        "        if self.max_train_length is not None:\n",
        "            sections = train_data.shape[1] // self.max_train_length\n",
        "        if sections >= 2:\n",
        "            train_data = np.concatenate(split_with_nan(train_data, sections, axis=1), axis=0)\n",
        "\n",
        "        temporal_missing = np.isnan(train_data).all(axis=-1).any(axis=0)\n",
        "        if temporal_missing[0] or temporal_missing[-1]:\n",
        "            train_data = centerize_vary_length_series(train_data)\n",
        "\n",
        "        train_data = train_data[~np.isnan(train_data).all(axis=2).all(axis=1)]\n",
        "        # pkl_save(\"/content/drive/MyDrive/Tesi/code/test/my_train_data.pkl\", train_data)\n",
        "\n",
        "        multiplier = 1 if train_data.shape[0] >= self.batch_size else math.ceil(self.batch_size / train_data.shape[0])\n",
        "\n",
        "        return train_data, multiplier\n",
        "\n",
        "\n",
        "    def setup(self, stage: str):\n",
        "        # load csv\n",
        "        data, dt_embed, n_covariate_cols = self._load_forecast_csv(self.path, self.name)\n",
        "\n",
        "        # scale and transform\n",
        "        self.data, self.scaler, self.n_covariate_cols = self._scale_and_transform(data, dt_embed, n_covariate_cols, self.name)\n",
        "\n",
        "        # Assign train/val datasets for use in dataloaders\n",
        "        if stage == \"fit\":\n",
        "            train_data = self.data[:, self.train_slice]\n",
        "            train_data, multiplier = self._fit_setup(train_data)\n",
        "            # fit setup\n",
        "            self.train = CustomDataset(torch.from_numpy(train_data).to(torch.float), sigma = 0.5, multiplier = multiplier)\n",
        "            # self.validate = ElectricityDataset(data[self.valid_slice], self.look_window)\n",
        "        if stage == 'encoding':\n",
        "            self.encode = TensorDataset(torch.from_numpy(self.data).to(torch.float))\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "        # Assign test dataset for use in dataloader(s)\n",
        "        # if stage == \"test\":\n",
        "        #     self.test = ElectricityDataset(data[self.test_slice], self.look_window)\n",
        "\n",
        "        # if stage == \"predict\":\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        return DataLoader(self.train, batch_size=min(self.batch_size, len(self.train)), drop_last = True, shuffle = True)\n",
        "\n",
        "    def encode_dataloader(self):\n",
        "        return DataLoader(self.encode, batch_size=self.encode_batch_size)\n",
        "\n",
        "    # def test_dataloader(self):\n",
        "    #     return DataLoader(self.test, batch_size=self.batch_size, drop_last = True)\n",
        "\n",
        "    # def predict_dataloader(self):"
      ],
      "metadata": {
        "id": "37YF2vGhSYjl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Preprocessing"
      ],
      "metadata": {
        "id": "S84zQ9UjigKj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Electricity"
      ],
      "metadata": {
        "id": "o-wdZhMWZybk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# df = pd.read_csv(datasets_path[ELECTRICITY] + datasets_name[ELECTRICITY], sep = ';')\n",
        "# df.rename(columns={df.columns[0]: 'Date'},inplace=True)\n",
        "# values = df.values\n",
        "# values = values[:, 1:].astype(str)\n",
        "# for i, value in enumerate(values):\n",
        "#   values[i] = np.char.replace(value, \",\", \".\")\n",
        "# values = values.astype(np.float32)\n",
        "# np.save(datasets_path[ELECTRICITY] + \"/electricity\", values)"
      ],
      "metadata": {
        "id": "YsyolJAOjZVS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# # # call this function to initialize the csv file\n",
        "# def electricity_preprocess(path):\n",
        "#     data_ecl = pd.read_csv(path + '/LD2011_2014.txt', parse_dates=True, sep=';', decimal=',', index_col=0)\n",
        "#     data_ecl = data_ecl.resample('1h', closed='right').sum()\n",
        "#     data_ecl = data_ecl.loc[:, data_ecl.cumsum(axis=0).iloc[8920] != 0]  # filter out instances with missing values\n",
        "#     data_ecl.index = data_ecl.index.rename('date')\n",
        "#     data_ecl = data_ecl['2012':]\n",
        "#     data_ecl.to_csv(path + '/electricity.csv')\n",
        "#     log.info(\"electriciy.csv created!\")"
      ],
      "metadata": {
        "id": "HqM-Bf5vWJDg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# electricity_preprocess(datasets_path[ELECTRICITY])"
      ],
      "metadata": {
        "id": "YS1CZBhuWei-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Models"
      ],
      "metadata": {
        "id": "Oq9gmXnKGmVP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Vanilla Transformer Encoder"
      ],
      "metadata": {
        "id": "lnbN5rUYocvE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# VanillaTransformer encoder\n",
        "\n",
        "\"\"\"\n",
        "https://arxiv.org/pdf/1706.03762.pdf\n",
        "\"\"\"\n",
        "\n",
        "class VanillaTransformerEncoder(nn.Module):\n",
        "    def __init__(self, d_model, n_heads = 16, dropout = 0.2):\n",
        "        super(VanillaTransformerEncoder, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.mha = MultiHeadAttention(d_model, n_heads)\n",
        "        self.norm1 = nn.LayerNorm(d_model) # maybe batch normalization\n",
        "        self.dropout1 = nn.Dropout(dropout)\n",
        "\n",
        "        self.pffn = PositionWiseFeedForwardNetwork(d_model)\n",
        "        self.norm2 = nn.LayerNorm(d_model)\n",
        "        self.dropout2 = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # new variable because of residual connection\n",
        "        z = self.mha(x,x,x)\n",
        "        z = self.dropout1(z)\n",
        "        z = self.norm1(z + x)\n",
        "\n",
        "        # set the new value for the residual connection\n",
        "        x = z\n",
        "        z = self.pffn(z)\n",
        "        z = self.dropout2(z)\n",
        "        return self.norm2(z + x)\n",
        "\n",
        "\"\"\"\n",
        "ref: https://d2l.ai/chapter_attention-mechanisms-and-transformers/multihead-attention.html\n",
        "\"\"\"\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, d_model, d_k, n_head = 6):\n",
        "        super(MultiHeadAttention, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.n_heads = n_head\n",
        "        self.d_k = d_k\n",
        "\n",
        "        self.W_q = nn.Linear(d_model, n_head * d_k, bias = False)\n",
        "        self.W_k = nn.Linear(d_model, n_head * d_k, bias = False)\n",
        "        self.W_v = nn.Linear(d_model, n_head * d_k, bias = False)\n",
        "        self.W_o = nn.Linear(n_head * d_k, d_model, bias = False)\n",
        "\n",
        "    # reshape to compute in parallel the several heads\n",
        "    def reshape_vector(self, x, inverse = False):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, d_model) || (batch, n_head, input_size, d_k)\n",
        "        output: (batch, n_head, input_size, d_k) || (batch, input_size, d_model)\n",
        "        \"\"\"\n",
        "        out = None\n",
        "\n",
        "        if not inverse:\n",
        "            out = einops.rearrange(x, 'b l (dim h) -> b h l dim', h=self.n_heads)\n",
        "        else:\n",
        "            out = einops.rearrange(x, 'b h l dim -> b l (dim h)')\n",
        "\n",
        "        return out\n",
        "\n",
        "    \"\"\"\n",
        "    ref: https://machinelearningmastery.com/the-transformer-attention-mechanism/\n",
        "    \"\"\"\n",
        "\n",
        "    def scaled_attention(self, q, k, v, dk, mask = None):\n",
        "        \"\"\"\n",
        "        q, k, v: (batch, n_head, input_size, d_k)\n",
        "        output: (batch, n_head, input_size, d_k), (batch, n_head, input_size, input_size)\n",
        "        \"\"\"\n",
        "        sqrt_d_k = math.sqrt(dk)\n",
        "\n",
        "        # using einsum to perform batch matrix multiplication\n",
        "        score = einops.einsum(q, k, 'b h l d_k, b h l_1 d_k -> b h l l_1') / sqrt_d_k\n",
        "\n",
        "        if mask is not None:\n",
        "            score = score.masked_fill(mask, -1e9)\n",
        "\n",
        "        weights = F.softmax(score, dim = -1)\n",
        "        attn = weights\n",
        "\n",
        "        res = einops.einsum(weights, v, 'b h l l_1, b h l_1 d_k -> b h l d_k')\n",
        "\n",
        "        return res, attn\n",
        "\n",
        "    def forward(self, q, k, v, mask = None):\n",
        "        \"\"\"\n",
        "        q, k, v: (batch, input_size, d_model)\n",
        "        output: (batch, input_size, d_model), (batch, n_head, input_size, input_size)\n",
        "        \"\"\"\n",
        "        residual = q\n",
        "\n",
        "        q = self.reshape_vector(self.W_q(q))\n",
        "        k = self.reshape_vector(self.W_k(k))\n",
        "        v = self.reshape_vector(self.W_v(v))\n",
        "\n",
        "        if mask is not None:\n",
        "            if len(mask.size()) == 3:\n",
        "                mask = mask.unsqueeze(1)  # For head axis broadcasting.\n",
        "\n",
        "        # parallel computation\n",
        "        out, attn = self.scaled_attention(q, k, v, self.d_k, mask)\n",
        "        out_concat = self.reshape_vector(out, inverse = True)\n",
        "\n",
        "        return self.W_o(out_concat) + residual, attn\n",
        "\n",
        "class PositionWiseFeedForwardNetwork(nn.Module):\n",
        "    def __init__(self, d_model, d_inner = 512):\n",
        "        super(PositionWiseFeedForwardNetwork, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.W_1 = nn.Linear(d_model, d_inner)\n",
        "        self.act = nn.GELU()\n",
        "        self.W_2 = nn.Linear(d_inner, d_model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, d_model)\n",
        "        output: (batch, input_size, d_model)\n",
        "        \"\"\"\n",
        "        res = x\n",
        "        x = self.W_1(x)\n",
        "        x = self.act(x)\n",
        "        return res + self.W_2(x)"
      ],
      "metadata": {
        "id": "f7JI-qgOpWbs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## PyraformerEncoder"
      ],
      "metadata": {
        "id": "j4VKdJgKsshh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Helper Functions"
      ],
      "metadata": {
        "id": "d0TB55l1yKvG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_mask(input_size, window_size, inner_size, device):\n",
        "    \"\"\"Get the attention mask of PAM-Naive\"\"\"\n",
        "    # Get the size of all layers\n",
        "    all_size = []\n",
        "    all_size.append(input_size)\n",
        "    # we split the nodes in according with the number of children\n",
        "    for i in range(len(window_size)):\n",
        "        layer_size = math.floor(all_size[i] / window_size[i])\n",
        "        all_size.append(layer_size)\n",
        "\n",
        "    # length of the flattened graph\n",
        "    seq_length = sum(all_size)\n",
        "    # mask matrix\n",
        "    mask = torch.zeros(seq_length, seq_length, device=device)\n",
        "\n",
        "    # get intra-scale mask\n",
        "    inner_window = inner_size // 2\n",
        "    for layer_idx in range(len(all_size)):\n",
        "        start = sum(all_size[:layer_idx])\n",
        "        for i in range(start, start + all_size[layer_idx]):\n",
        "            left_side = max(i - inner_window, start)\n",
        "            right_side = min(i + inner_window + 1, start + all_size[layer_idx])\n",
        "            mask[i, left_side:right_side] = 1\n",
        "\n",
        "    # get inter-scale mask\n",
        "    for layer_idx in range(1, len(all_size)):\n",
        "        start = sum(all_size[:layer_idx])\n",
        "        for i in range(start, start + all_size[layer_idx]):\n",
        "            left_side = (start - all_size[layer_idx - 1]) + (i - start) * window_size[layer_idx - 1]\n",
        "            if i == ( start + all_size[layer_idx] - 1):\n",
        "                right_side = start\n",
        "            else:\n",
        "                right_side = (start - all_size[layer_idx - 1]) + (i - start + 1) * window_size[layer_idx - 1]\n",
        "            mask[i, left_side:right_side] = 1\n",
        "            mask[left_side:right_side, i] = 1\n",
        "\n",
        "    mask = (1 - mask).bool()\n",
        "\n",
        "    return mask, all_size\n",
        "\n",
        "def get_graph_dim(input_size, window_size):\n",
        "    \"\"\" get the dimension of the graph computed by CSCM\"\"\"\n",
        "    res = input_size\n",
        "    for w in window_size:\n",
        "        input_size = math.floor(input_size / w)\n",
        "        res += input_size\n",
        "\n",
        "    return res\n",
        "\n"
      ],
      "metadata": {
        "id": "rZXf-0odyPbJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Model"
      ],
      "metadata": {
        "id": "d6HSFST8yIkK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class Bottleneck_Construct(nn.Module):\n",
        "    \"\"\"Bottleneck convolution CSCM\"\"\"\n",
        "    def __init__(self, d_model, window_size, d_inner):\n",
        "        super(Bottleneck_Construct, self).__init__()\n",
        "        if not isinstance(window_size, list):\n",
        "            self.conv_layers = nn.ModuleList([\n",
        "                ConvLayer(d_inner, window_size),\n",
        "                ConvLayer(d_inner, window_size),\n",
        "                ConvLayer(d_inner, window_size)\n",
        "                ])\n",
        "        else:\n",
        "            self.conv_layers = []\n",
        "            for i in range(len(window_size)):\n",
        "                self.conv_layers.append(ConvLayer(d_inner, window_size[i]))\n",
        "            self.conv_layers = nn.ModuleList(self.conv_layers)\n",
        "        self.up = nn.Linear(d_inner, d_model)\n",
        "        self.down = nn.Linear(d_model, d_inner)\n",
        "        self.norm = nn.LayerNorm(d_model)\n",
        "\n",
        "    def forward(self, enc_input):\n",
        "\n",
        "        temp_input = self.down(enc_input).permute(0, 2, 1)\n",
        "        all_inputs = []\n",
        "        for i in range(len(self.conv_layers)):\n",
        "            temp_input = self.conv_layers[i](temp_input)\n",
        "            all_inputs.append(temp_input)\n",
        "\n",
        "        all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)\n",
        "        all_inputs = self.up(all_inputs)\n",
        "        all_inputs = torch.cat([enc_input, all_inputs], dim=1)\n",
        "\n",
        "        all_inputs = self.norm(all_inputs)\n",
        "\n",
        "        return all_inputs\n",
        "\"\"\" For Electricity Dataset\"\"\"\n",
        "class PositionalEmbedding(nn.Module):\n",
        "    def __init__(self, d_model, max_len=5000):\n",
        "        super(PositionalEmbedding, self).__init__()\n",
        "        pe = torch.zeros(max_len, d_model)\n",
        "        # create a positional array\n",
        "        position = torch.arange(0, max_len).unsqueeze(1)\n",
        "        # div term for half of positions\n",
        "        div_term = torch.pow(10000.0, torch.arange(0, d_model, 2) / d_model)\n",
        "        # even positions\n",
        "        pe[:, 0::2] = torch.sin(position * div_term)\n",
        "        # odd positions\n",
        "        pe[:, 1::2] = torch.cos(position * div_term)\n",
        "\n",
        "        pe = pe.unsqueeze(0)\n",
        "        self.register_buffer('pe', pe)\n",
        "\n",
        "        # if normalize:\n",
        "        #     pe = pe - pe.mean()\n",
        "        #     pe = pe / (pe.std() * 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        output: (1, input_size, d_model)\n",
        "        \"\"\"\n",
        "        return self.pe[:, :x.size(1)]\n",
        "\n",
        "class TokenEmbedding(nn.Module):\n",
        "    def __init__(self, c_in, d_model):\n",
        "        super(TokenEmbedding, self).__init__()\n",
        "        padding = 1 if torch.__version__>='1.5.0' else 2\n",
        "        self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,\n",
        "                                    kernel_size=3, padding=padding, padding_mode='circular')\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.Conv1d):\n",
        "                nn.init.kaiming_normal_(m.weight,mode='fan_in',nonlinearity='leaky_relu')\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, enc_in)\n",
        "        output: (batch, input_size, d_model)\n",
        "        \"\"\"\n",
        "        x = self.tokenConv(einops.rearrange(x, 'b l e -> b e l')).transpose(1,2)\n",
        "        return x\n",
        "\n",
        "class CustomEmbedding(nn.Module):\n",
        "    def __init__(self, c_in, d_model, temporal_size, seq_num, dropout=0.1):\n",
        "        super(CustomEmbedding, self).__init__()\n",
        "\n",
        "        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)\n",
        "        self.position_embedding = PositionalEmbedding(d_model=d_model)\n",
        "        self.temporal_embedding = nn.Linear(temporal_size, d_model)\n",
        "        self.seqid_embedding = nn.Embedding(seq_num, d_model)\n",
        "\n",
        "        self.dropout = nn.Dropout(p=dropout)\n",
        "\n",
        "    def forward(self, x, x_mark):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, enc_in)\n",
        "        x_mark: (batch, input_size)\n",
        "        output: (batch, input_size, d_model)\n",
        "        \"\"\"\n",
        "        x = self.value_embedding(x) + self.position_embedding(x) + self.temporal_embedding(x_mark[:, :, :-1]) \\\n",
        "            + self.seqid_embedding(x_mark[:, :, -1].long())\n",
        "\n",
        "\n",
        "\n",
        "        return self.dropout(x)\n",
        "\n",
        "class EncoderLayer(nn.Module):\n",
        "    \"\"\" Compose with two layers \"\"\"\n",
        "\n",
        "    def __init__(self, d_model, d_inner, d_k, n_head, dropout=0.1, normalize_before=True, q_k_mask=None, k_q_mask=None):\n",
        "        super(EncoderLayer, self).__init__()\n",
        "        self.slf_attn = MultiHeadAttention(d_model, d_k, n_head)\n",
        "\n",
        "        self.pos_ffn = PositionWiseFeedForwardNetwork(\n",
        "            d_model, d_inner)\n",
        "\n",
        "    def forward(self, enc_input, slf_attn_mask=None):\n",
        "        \"\"\"\n",
        "        enc_input: (batch, input_size, d_model)\n",
        "        output: (batch, input_size, d_model), (batch, n_head, input_size, input_size)\n",
        "        \"\"\"\n",
        "        enc_output, enc_slf_attn = self.slf_attn(enc_input, enc_input, enc_input, mask=slf_attn_mask)\n",
        "\n",
        "        enc_output = self.pos_ffn(enc_output)\n",
        "\n",
        "        return enc_output, enc_slf_attn\n",
        "\n",
        "class ConvLayer(nn.Module):\n",
        "    def __init__(self, c_in, window_size):\n",
        "        super(ConvLayer, self).__init__()\n",
        "        self.downConv = nn.Conv1d(in_channels=c_in,\n",
        "                                  out_channels=c_in,\n",
        "                                  kernel_size=window_size,\n",
        "                                  stride=window_size)\n",
        "        self.norm = nn.BatchNorm1d(c_in)\n",
        "        self.activation = nn.ELU()\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.downConv(x)\n",
        "        x = self.norm(x)\n",
        "        x = self.activation(x)\n",
        "        return x\n",
        "\n",
        "class Bottleneck_Construct(nn.Module):\n",
        "    \"\"\"Bottleneck convolution CSCM\"\"\"\n",
        "    def __init__(self, d_model, window_size, d_inner):\n",
        "        super(Bottleneck_Construct, self).__init__()\n",
        "        if not isinstance(window_size, list):\n",
        "            self.conv_layers = nn.ModuleList([\n",
        "                ConvLayer(d_inner, window_size),\n",
        "                ConvLayer(d_inner, window_size),\n",
        "                ConvLayer(d_inner, window_size)\n",
        "                ])\n",
        "        else:\n",
        "            self.conv_layers = []\n",
        "            for i in range(len(window_size)):\n",
        "                self.conv_layers.append(ConvLayer(d_inner, window_size[i]))\n",
        "            self.conv_layers = nn.ModuleList(self.conv_layers)\n",
        "        self.up = nn.Linear(d_inner, d_model)\n",
        "        self.down = nn.Linear(d_model, d_inner)\n",
        "        self.norm = nn.LayerNorm(d_model)\n",
        "\n",
        "    def forward(self, enc_input):\n",
        "        \"\"\"\n",
        "        enc_input: (batch, input_size, d_model)\n",
        "        output: (batch, graph_size, d_model)\n",
        "        \"\"\"\n",
        "\n",
        "        temp_input = self.down(enc_input).permute(0, 2, 1)\n",
        "        all_inputs = []\n",
        "        for i in range(len(self.conv_layers)):\n",
        "            temp_input = self.conv_layers[i](temp_input)\n",
        "            all_inputs.append(temp_input)\n",
        "\n",
        "        all_inputs = torch.cat(all_inputs, dim=2).transpose(1, 2)\n",
        "        all_inputs = self.up(all_inputs)\n",
        "        # concat the computed new nodes with the input nodes\n",
        "        all_inputs = torch.cat([enc_input, all_inputs], dim=1)\n",
        "\n",
        "        all_inputs = self.norm(all_inputs)\n",
        "\n",
        "        return all_inputs\n",
        "\n",
        "class PyraformerEncoder(nn.Module):\n",
        "    \"\"\" A encoder model with self attention mechanism. \"\"\"\n",
        "\n",
        "    def __init__(self, d_model = 320, d_k = 160, window_size = [4,4,4], inner_size = 3,\n",
        "                 input_size = 201, d_inner_hid = 512, n_head = 6, n_layer = 10,\n",
        "                 # Dataloader parameters\n",
        "                 enc_in = 1, covariate_size = 7, seq_num = 321,\n",
        "                 CSCM = \"Bottleneck_Construct\", d_bottleneck = 128, device = 'cpu'):\n",
        "        super(PyraformerEncoder, self).__init__()\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.d_model = d_model # size of the latent vector\n",
        "        self.d_k = d_k # size of the inner dimension of q, k, v\n",
        "        self.window_size = window_size # The number of children of a parent node\n",
        "        self.inner_size = inner_size # The number of ajacent nodes\n",
        "        self.input_size = input_size # length of the sequence\n",
        "        self.d_inner_hid = d_inner_hid # inner size of the PostitionalFeedForward\n",
        "        self.n_head = n_head\n",
        "        self.n_layer = n_layer\n",
        "        self.enc_in = enc_in\n",
        "        self.covariate_size = covariate_size # number of temporal covariate\n",
        "        self.seq_num = seq_num # size of the time series\n",
        "        self.CSCM = CSCM # called coarser-scale construction module\n",
        "        self.g_size = get_graph_dim(input_size, window_size)\n",
        "        self.d_bottleneck = d_bottleneck #\n",
        "        self.mask, self.all_size = get_mask(self.input_size, self.window_size, self.inner_size, device)\n",
        "        self.layers = nn.ModuleList([\n",
        "                EncoderLayer(self.d_model, self.d_inner_hid, self.d_k, self.n_head) for i in range(self. n_layer)\n",
        "                ])\n",
        "        self.enc_embedding = nn.Linear(enc_in + covariate_size, d_model)#CustomEmbedding(self.enc_in, self.d_model, self.covariate_size, self.seq_num)\n",
        "\n",
        "        self.conv_layers = eval(self.CSCM)(self.d_model, self.window_size, self.d_bottleneck)\n",
        "\n",
        "        self.fc = nn.Linear(self.g_size, self.input_size)\n",
        "\n",
        "        self.test = nn.Linear(enc_in + covariate_size, d_model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, enc_in + covariate_size)\n",
        "        output: (batch, input_size, d_model)\n",
        "        \"\"\"\n",
        "\n",
        "        seq_enc = self.enc_embedding(x)\n",
        "\n",
        "        # Repeat the mask for all the batch\n",
        "        mask = self.mask.repeat(len(seq_enc), 1, 1)\n",
        "\n",
        "        seq_enc = self.conv_layers(seq_enc)\n",
        "\n",
        "        for i in range(len(self.layers)):\n",
        "            seq_enc, _ = self.layers[i](seq_enc, mask)\n",
        "\n",
        "        seq_enc = einops.rearrange(seq_enc, 'b g d -> b d g')\n",
        "\n",
        "        seq_enc = self.fc(seq_enc)\n",
        "\n",
        "        seq_enc = einops.rearrange(seq_enc, 'b d l -> b l d')\n",
        "\n",
        "        return seq_enc"
      ],
      "metadata": {
        "id": "LiOdyvP8syOH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## PatchTST"
      ],
      "metadata": {
        "id": "QD52xt7xGsd9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Utility functions\n",
        "\n",
        "def create_patches(xb, patch_len, stride):\n",
        "    \"\"\"\n",
        "    xb -> [B x L x M] // [B x L x M x T]\n",
        "    output -> [B x N x M x P] // [B x N x M x T x P], N\n",
        "    \"\"\"\n",
        "    _, num_var, _, _ = xb.shape\n",
        "    # compute number of patches\n",
        "    patch_num = (max(patch_len, num_var)-patch_len) // stride + 2\n",
        "\n",
        "    # we repeat the last variable of the sequence to have equal patches\n",
        "    tail = torch.repeat_interleave(xb[:,-1:,...], stride, dim = 1)\n",
        "    xb = torch.concatenate((xb, tail), axis = 1)\n",
        "\n",
        "    # create patches\n",
        "    xb = xb.unfold(dimension=1, size=patch_len, step=stride)\n",
        "\n",
        "    assert patch_num == xb.shape[1], f\"wrong number of computed patches, expected {patch_num} but computed {xb.shape[1]}\"\n",
        "\n",
        "    return xb, patch_num\n",
        "\n",
        "\"\"\"\n",
        "ref: https://kazemnejad.com/blog/transformer_architecture_positional_encoding/\n",
        "\"\"\"\n",
        "\n",
        "def positional_encoding(batch_size, max_len, d_model):\n",
        "    \"\"\"\n",
        "    output\n",
        "    \"\"\"\n",
        "    pe = torch.zeros(batch_size, max_len, d_model)\n",
        "    # create a positional array\n",
        "    position = torch.arange(0, max_len).unsqueeze(1)\n",
        "    # div term for half of positions\n",
        "    div_term = torch.pow(10000.0, torch.arange(0, d_model, 2) / d_model)\n",
        "    # even positions\n",
        "    pe[:, :, 0::2] = torch.sin(position * div_term)\n",
        "    # odd positions\n",
        "    pe[:, :, 1::2] = torch.cos(position * div_term)\n",
        "\n",
        "    # if normalize:\n",
        "    #     pe = pe - pe.mean()\n",
        "    #     pe = pe / (pe.std() * 10)\n",
        "\n",
        "    return nn.parameter.Parameter(pe, requires_grad= False)"
      ],
      "metadata": {
        "id": "hbRO5NzrNy8b"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#PatchTST\n",
        "\n",
        "class PatchTSTEncoder(nn.Module):\n",
        "    def __init__(self, num_channels, num_var, patch_len, stride, batch_size, time_dimension = 8, d_model = 128, n_layers = 3, n_heads = 16, dropout = 0.2):\n",
        "        super(PatchTSTEncoder, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.num_channels = num_channels\n",
        "        self.patch_num = (max(patch_len, num_var)-patch_len) // stride + 2\n",
        "        self.patch_len = patch_len\n",
        "        self.stride = stride\n",
        "        self.batch_size = batch_size\n",
        "        self.d_model = d_model\n",
        "        self.n_layers = n_layers\n",
        "        self.n_heads = n_heads\n",
        "        self.dropout = dropout\n",
        "\n",
        "        # instance normalization\n",
        "        \"\"\"\n",
        "        ref: https://wandb.ai/wandb_fc/Normalization-Series/reports/Instance-Normalization-in-PyTorch-With-Examples---VmlldzoxNDIyNTQx\n",
        "        \"\"\"\n",
        "        self.inst_norm = nn.InstanceNorm2d(num_channels)\n",
        "\n",
        "        # patch creation\n",
        "        self.create_patch = create_patches\n",
        "\n",
        "        # embedding\n",
        "        self.W_p = nn.Linear(patch_len * time_dimension, d_model, bias = False)\n",
        "\n",
        "        # positional encoding\n",
        "        self.W_pos = positional_encoding(batch_size * num_channels, self.patch_num, d_model)\n",
        "\n",
        "        # dropout\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "        # encoder\n",
        "        self.encoders = nn.ModuleList([VanillaTransformerEncoder(d_model) for _ in range(n_layers)])\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x -> [B x L x M] // [(B x M) x MAX_TRAIN_LENGTH x TIME_DIM]\n",
        "        output -> [(B M) x N x D]\n",
        "        \"\"\"\n",
        "        b_m, _, _ = x\n",
        "        assert b_m / self.num_channel == self.batch_size, f\"invalid fisrt dimension {b_m / self.num_channel} != {self.batch_size}\"\n",
        "        # [(B M) x MAX_TRAIN_LENGTH x TIME_DIM] -> [B x M x MAX_TRAIN_LENGTH x TIME_DIM]\n",
        "        x = einops.rearrange(x, '(b m) l t -> b m l t', m=self.num_channels)\n",
        "        # we need to reshape dimensione before apply instance normalization\n",
        "        x = einops.rearrange(self.inst_norm(x), 'b m l t -> b l m t')\n",
        "\n",
        "        # create patches\n",
        "        x, patch_num = self.create_patch(x, self.patch_len, self.stride)\n",
        "\n",
        "        # x: [B x N x M x T x P]\n",
        "\n",
        "        assert self.patch_num == patch_num, f\"wrong number for patch_num {self.patch_num} != {patch_num}\"\n",
        "\n",
        "        # reshape the tensor from [B x N x M x T x P] -> [(B M) x N x (P T)]\n",
        "        x = einops.rearrange(x, 'b n m t p -> (b m) n (p t)')\n",
        "        # now it can be provided to our transformer implementation\n",
        "\n",
        "        # project into transformer latent space\n",
        "        x = self.W_p(x) + self.W_pos\n",
        "\n",
        "        for layer in self.encoders:\n",
        "            x = layer(x)\n",
        "\n",
        "        return x\n",
        "\n"
      ],
      "metadata": {
        "id": "bhiJYUZcGvr_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Cost"
      ],
      "metadata": {
        "id": "DTlGPLvHrJ5g"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.nn.modules.loss import TripletMarginWithDistanceLoss\n",
        "#Cost\n",
        "\n",
        "class Cost(nn.Module):\n",
        "    def __init__(self, input_size, d_model = 320, d_s = 160, d_t = 160, n_layers = 3, n_heads = 6, dropout = 0.2):\n",
        "        super(Cost, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        # Dropout for seasonal representation output\n",
        "        self.seasonal_drop = nn.Dropout(0.1)\n",
        "\n",
        "        # Trend Feature Disentangler\n",
        "        self.tfd = TrendFeatureDisentangler(d_model, d_t, input_size)\n",
        "\n",
        "        # Seasonal Feature Disentangler\n",
        "        self.sfd = SeasonalFeatureDisentangler(d_model, d_s, input_size)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, d_model)\n",
        "        outputs: (batch, input_size, d_t), (batch, input_size, d_s)\n",
        "        \"\"\"\n",
        "        out_tfd = self.tfd(x)\n",
        "\n",
        "        out_svd = self.sfd(x)\n",
        "\n",
        "        out_svd = self.seasonal_drop(out_svd)\n",
        "\n",
        "        return out_tfd, out_svd\n",
        "\n",
        "\n",
        "\n",
        "# Causal Convolution (dilated)\n",
        "\n",
        "class CausalConv1d(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, kernel_size, dilation=1):\n",
        "        super(CausalConv1d, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.kernel_size = kernel_size\n",
        "        pad = (kernel_size - 1) * dilation\n",
        "        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=pad, dilation=dilation)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        input: (batch, input_size, in_channels)\n",
        "        output: (batch, conv_out, out_channels)\n",
        "        \"\"\"\n",
        "        # we need to reshape before applying the convolution\n",
        "        x = einops.rearrange(x, 'b l i_c -> b i_c l')\n",
        "        x = self.conv(x)\n",
        "\n",
        "        # we need to remove the trailing padding zeros (except for the fist layer) from the values\n",
        "        if self.kernel_size > 1:\n",
        "            x = x[...,0:-(self.kernel_size-1)]\n",
        "\n",
        "        # rearrange to the original shape\n",
        "        x = einops.rearrange(x, 'b o_c l -> b l o_c')\n",
        "\n",
        "        return x\n",
        "\n",
        "# TFD\n",
        "\n",
        "class TrendFeatureDisentangler(nn.Module):\n",
        "    def __init__(self, d_model, d_t, input_size):\n",
        "        super(TrendFeatureDisentangler, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "        self.d_model = d_model\n",
        "        self.d_t = d_t\n",
        "        self.input_size = input_size\n",
        "\n",
        "        # https://discuss.pytorch.org/t/causal-convolution/3456/3\n",
        "        # https://arxiv.org/pdf/1609.03499v2.pdf\n",
        "\n",
        "        # floor(log(N/2)) autoregressive expert\n",
        "        self.conv_num = math.floor(math.log2(input_size / 2)) + 1\n",
        "        self.kernel = [2**i for i in range(self.conv_num)]\n",
        "        self.convolutions = nn.ModuleList([CausalConv1d(d_model, d_t, k) for k in self.kernel])\n",
        "\n",
        "    def avg_pooling(self, input):\n",
        "        \"\"\"\n",
        "        input: (list, batch, input_size, d_t)\n",
        "        \"\"\"\n",
        "        return einops.reduce(input, 'list b l d_t -> b l d_t', 'mean')\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, d_model)\n",
        "        output: (batch, input_size, d_t)\n",
        "        \"\"\"\n",
        "        batch_size, input_size, d_model = x.shape\n",
        "\n",
        "        assert input_size == self.input_size and d_model == self.d_model, \"wrong input dimensions\"\n",
        "\n",
        "        # create the result tensor\n",
        "        trend = torch.zeros(self.conv_num, batch_size, input_size, self.d_t, device = x.device)\n",
        "\n",
        "        for i, conv in enumerate(self.convolutions):\n",
        "            out = conv(x)\n",
        "            trend[i,...] = out\n",
        "\n",
        "        # apply the average pooling operation\n",
        "        trend = self.avg_pooling(trend)\n",
        "\n",
        "        return trend\n",
        "\n",
        "# SVD\n",
        "\n",
        "class SeasonalFeatureDisentangler(nn.Module):\n",
        "    def __init__(self, d_model, d_s, input_size):\n",
        "        super(SeasonalFeatureDisentangler, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.input_size = input_size\n",
        "\n",
        "        # number of frequencies after dft\n",
        "        self.f = input_size // 2 + 1\n",
        "\n",
        "        # discrete fast fourier transform, rfft output contains only the positive frequencies below the Nyquist frequency\n",
        "        self.dft = torch.fft.rfft\n",
        "\n",
        "        # Learnable Fourier Layer\n",
        "        self.fl = FourierLayer(self.f, d_model, d_s, input_size)\n",
        "\n",
        "        # inverse of discrete fast fourier transform\n",
        "        self.idft = torch.fft.irfft\n",
        "\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, d_model)\n",
        "        output: (batch, input_size, d_s)\n",
        "        \"\"\"\n",
        "        # we apply dft along the temporal dimension\n",
        "        x = self.dft(x, dim = 1)\n",
        "\n",
        "        assert self.f == x.shape[1], \"wrong dimension of dft\"\n",
        "\n",
        "        # apply fourier layer\n",
        "        x = self.fl(x)\n",
        "\n",
        "        # compute the inverse of dft to come back to time domain\n",
        "        x = self.idft(x, n = self.input_size, dim = 1) # pass also the legth in order to avoid odd-length problems\n",
        "\n",
        "        return x\n",
        "\n",
        "class FourierLayer(nn.Module):\n",
        "    def __init__(self, f, d_model, d_s, input_size):\n",
        "        super(FourierLayer, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.f = f\n",
        "        self.d_model = d_model\n",
        "\n",
        "        self.A = nn.Parameter(torch.empty((f, d_model, d_s), dtype=torch.cfloat))\n",
        "        self.B = nn.Parameter(torch.empty((f, d_s), dtype=torch.cfloat))\n",
        "        self._reset_parameters()\n",
        "\n",
        "    def _reset_parameters(self):\n",
        "        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))\n",
        "        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.A)\n",
        "        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0\n",
        "        nn.init.uniform_(self.B, -bound, bound)\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, f, d_model)\n",
        "        out: (batch, f, d_s)\n",
        "        \"\"\"\n",
        "        batch_size, f, _ = x.shape\n",
        "\n",
        "        assert f == self.f, \"wrong dimensions of x\"\n",
        "\n",
        "        out = einops.einsum(x, self.A, 'b f d, f d d_s -> b f d_s') + self.B\n",
        "\n",
        "        return out\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "-U9fYJUKrN4S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## CoSpy Encoder"
      ],
      "metadata": {
        "id": "omgfLK9mFicv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# %debug\n",
        "class CoSpyEncoder(nn.Module):\n",
        "    def __init__(self, input_size,\n",
        "                 d_model = 320, d_s = 160, d_t = 160,\n",
        "                 dropout = 0.2, device = 'cpu', enc = \"Pyraformer\",\n",
        "                 input_dims = 8, hidden_dims = 64, output_dims = 320, depth = 3):\n",
        "        super(CoSpyEncoder, self).__init__()\n",
        "        # self.save_hyperparameters()\n",
        "\n",
        "        self.input_size = input_size\n",
        "        self.device = device\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        # Pyraformer or TCN layer (backbone encoder)\n",
        "        self.en = PyraformerEncoder(d_model, device = device) if enc == \"Pyraformer\" \\\n",
        "            else TCN(input_dims, hidden_dims, output_dims, depth)\n",
        "\n",
        "\n",
        "        # CoST layer (disentangler)\n",
        "        self.cost = Cost(self.input_size)\n",
        "\n",
        "    def encode(self, data_shape, loader, batch_size = 256, sliding_length=1, padding=200):\n",
        "\n",
        "        encoding_window = None\n",
        "        slicing = None\n",
        "\n",
        "        n_samples, ts_l, _ = data_shape\n",
        "\n",
        "        org_training = self.training\n",
        "        self.eval()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            output = []\n",
        "            for batch in tqdm(loader, desc=\"data encoding\"):\n",
        "                x = batch[0]\n",
        "                reprs = []\n",
        "                # x = x.to(device)\n",
        "                if n_samples < batch_size:\n",
        "                    calc_buffer = []\n",
        "                    calc_buffer_l = 0\n",
        "                # self.log.debug(type(x))\n",
        "                # self.log.debug(f\"shape of batch: {x.shape}\")\n",
        "\n",
        "                # if self.batch_size != batch_size:\n",
        "                #     self.log.debug(\"different batch size return same batch\")\n",
        "                #     return batch\n",
        "\n",
        "                for i in tqdm(range(0, ts_l, sliding_length),\n",
        "                              desc = \"sequence encoding\"):\n",
        "                    l = i - padding # sliding_padding=200\n",
        "                    r = i + sliding_length\n",
        "                    # self.log.debug(x.device)\n",
        "                    x_sliding = pad_nan(\n",
        "                        x[:, max(l, 0) : min(r, ts_l)],\n",
        "                        left=-l if l<0 else 0,\n",
        "                        right=r-ts_l if r>ts_l else 0,\n",
        "                        dim=1\n",
        "                    )\n",
        "                    if n_samples < batch_size:\n",
        "                        if calc_buffer_l + n_samples > batch_size:\n",
        "                            out = self._eval_with_pooling(\n",
        "                                torch.cat(calc_buffer, dim=0)\n",
        "                            )\n",
        "                            reprs += torch.split(out, n_samples)\n",
        "                            calc_buffer = []\n",
        "                            calc_buffer_l = 0\n",
        "                        calc_buffer.append(x_sliding)\n",
        "                        calc_buffer_l += n_samples\n",
        "                    else:\n",
        "                        reprs.append(self._eval_with_pooling(x_sliding))\n",
        "\n",
        "                if n_samples < batch_size:\n",
        "                    if calc_buffer_l > 0:\n",
        "                        out = self._eval_with_pooling(\n",
        "                            torch.cat(calc_buffer, dim=0)\n",
        "                        )\n",
        "                        reprs += torch.split(out, n_samples)\n",
        "                        calc_buffer = []\n",
        "                        calc_buffer_l = 0\n",
        "\n",
        "                out = torch.cat(reprs, dim=1)\n",
        "                output.append(out)\n",
        "\n",
        "            output = torch.cat(output, dim=0)\n",
        "\n",
        "        self.train(org_training)\n",
        "        return output.numpy()\n",
        "\n",
        "    def _eval_with_pooling(self, x):\n",
        "        out_t, out_s = self(x.to(self.device, non_blocking=True))\n",
        "        out = torch.cat([out_t[:, -1], out_s[:, -1]], dim=-1)\n",
        "        return einops.rearrange(out.cpu(), 'b d -> b () d')\n",
        "\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, enc_in + covariate_size)\n",
        "        outputs: (batch, input_size, d_t), (batch, input_size, d_s)\n",
        "        \"\"\"\n",
        "\n",
        "        nan_mask = ~x.isnan().any(axis=-1)\n",
        "        x[~nan_mask] = 0\n",
        "\n",
        "        x = self.en(x)\n",
        "\n",
        "        trend, season = self.cost(x)\n",
        "\n",
        "        return trend, season"
      ],
      "metadata": {
        "id": "uSr7enmSFptR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## OFFCIAL IMPLEMENTATIONS 🔽"
      ],
      "metadata": {
        "id": "wriuLMemIEtl"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## NLinear"
      ],
      "metadata": {
        "id": "9VKcjh66BBAJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "\n",
        "class NLinear(nn.Module):\n",
        "    \"\"\"\n",
        "    Normalization-Linear\n",
        "    \"\"\"\n",
        "    def __init__(self, covariate_size, d_model, layers, individual = True):\n",
        "        super(NLinear, self).__init__()\n",
        "        self.seq_len = covariate_size\n",
        "        self.d_model = d_model\n",
        "\n",
        "        # Use this line if you want to visualize the weights\n",
        "        # self.Linear.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))\n",
        "        self.channels = layers\n",
        "        self.individual = individual\n",
        "        if self.individual:\n",
        "            self.Linear = nn.ModuleList()\n",
        "            for i in range(self.channels):\n",
        "                self.Linear.append(nn.Linear(self.seq_len,self.d_model))\n",
        "        else:\n",
        "            self.Linear = nn.Linear(self.seq_len, self.d_model)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # x: [Batch, Input length, Channel]\n",
        "        seq_last = x[:,-1:,:].detach()\n",
        "        x = x - seq_last\n",
        "        if self.individual:\n",
        "            output = torch.zeros([x.size(0),self.d_model,x.size(2)],dtype=x.dtype).to(x.device)\n",
        "            for i in range(self.channels):\n",
        "                output[:,:,i] = self.Linear[i](x[:,:,i])\n",
        "            x = output\n",
        "        else:\n",
        "            x = self.Linear(x.permute(0,2,1)).permute(0,2,1)\n",
        "        x = x + seq_last\n",
        "        return x # [Batch, Output length, Channel]"
      ],
      "metadata": {
        "id": "wb7zI3HuBDo0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "x = torch.rand((16, 8, 201))\n",
        "model = NLinear(8, 320, 4)\n",
        "x = model(x)\n",
        "x.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_gxE8HLUF_xn",
        "outputId": "22b03ff0-b4cb-425f-ea02-45c9f23a2525"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([16, 320, 201])"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## TCN"
      ],
      "metadata": {
        "id": "2xgjSvFBQUiS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TCN(nn.Module):\n",
        "    def __init__(self, input_dims, hidden_dims, output_dims, depth):\n",
        "        super(TCN, self).__init__()\n",
        "\n",
        "        self.log = setup_log(self, LOG_LEVEL)\n",
        "\n",
        "        self.input_fc = nn.Linear(input_dims, hidden_dims)\n",
        "\n",
        "        self.feature_extractor = DilatedConvEncoder(\n",
        "                hidden_dims,\n",
        "                [hidden_dims] * depth + [output_dims],\n",
        "                kernel_size=3\n",
        "            )\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, enc_in + covariate_size)\n",
        "        outputs: (batch, input_size, d_model)\n",
        "        \"\"\"\n",
        "        x = self.input_fc(x)\n",
        "        x = x.transpose(1, 2)\n",
        "        x = self.feature_extractor(x)\n",
        "        x = x.transpose(1, 2)\n",
        "        return x\n",
        "\n",
        "class DilatedConvEncoder(nn.Module):\n",
        "    def __init__(self, in_channels, channels, kernel_size, extract_layers=None):\n",
        "        super().__init__()\n",
        "\n",
        "        if extract_layers is not None:\n",
        "            assert len(channels) - 1 in extract_layers\n",
        "\n",
        "        self.extract_layers = extract_layers\n",
        "        self.net = nn.Sequential(*[\n",
        "            ConvBlock(\n",
        "                channels[i-1] if i > 0 else in_channels,\n",
        "                channels[i],\n",
        "                kernel_size=kernel_size,\n",
        "                dilation=2**i,\n",
        "                final=(i == len(channels)-1)\n",
        "            )\n",
        "            for i in range(len(channels))\n",
        "        ])\n",
        "\n",
        "    def forward(self, x):\n",
        "        if self.extract_layers is not None:\n",
        "            outputs = []\n",
        "            for idx, mod in enumerate(self.net):\n",
        "                x = mod(x)\n",
        "                if idx in self.extract_layers:\n",
        "                    outputs.append(x)\n",
        "            return outputs\n",
        "        return self.net(x)\n",
        "\n",
        "class ConvBlock(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, kernel_size, dilation, final=False):\n",
        "        super().__init__()\n",
        "        self.conv1 = SamePadConv(in_channels, out_channels, kernel_size, dilation=dilation)\n",
        "        self.conv2 = SamePadConv(out_channels, out_channels, kernel_size, dilation=dilation)\n",
        "        self.projector = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels or final else None\n",
        "\n",
        "    def forward(self, x):\n",
        "        residual = x if self.projector is None else self.projector(x)\n",
        "        x = F.gelu(x)\n",
        "        x = self.conv1(x)\n",
        "        x = F.gelu(x)\n",
        "        x = self.conv2(x)\n",
        "        return x + residual\n",
        "\n",
        "class SamePadConv(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, kernel_size, dilation=1, groups=1):\n",
        "        super().__init__()\n",
        "        self.receptive_field = (kernel_size - 1) * dilation + 1\n",
        "        padding = self.receptive_field // 2\n",
        "        self.conv = nn.Conv1d(\n",
        "            in_channels, out_channels, kernel_size,\n",
        "            padding=padding,\n",
        "            dilation=dilation,\n",
        "            groups=groups\n",
        "        )\n",
        "        self.remove = 1 if self.receptive_field % 2 == 0 else 0\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.conv(x)\n",
        "        if self.remove > 0:\n",
        "            out = out[:, :, : -self.remove]\n",
        "        return out"
      ],
      "metadata": {
        "id": "YtqxPp42QXsJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## CoSTEncoder"
      ],
      "metadata": {
        "id": "HuGIr5N1HpkN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_continuous_mask(B, T, n=5, l=0.1):\n",
        "    res = torch.full((B, T), True, dtype=torch.bool)\n",
        "    if isinstance(n, float):\n",
        "        n = int(n * T)\n",
        "    n = max(min(n, T // 2), 1)\n",
        "\n",
        "    if isinstance(l, float):\n",
        "        l = int(l * T)\n",
        "    l = max(l, 1)\n",
        "\n",
        "    for i in range(B):\n",
        "        for _ in range(n):\n",
        "            t = np.random.randint(T-l+1)\n",
        "            res[i, t:t+l] = False\n",
        "    return res\n",
        "\n",
        "\n",
        "def generate_binomial_mask(B, T, p=0.5):\n",
        "    return torch.from_numpy(np.random.binomial(1, p, size=(B, T))).to(torch.bool)\n",
        "\n",
        "\n",
        "class BandedFourierLayer(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, band, num_bands, length=201):\n",
        "        super().__init__()\n",
        "\n",
        "        self.length = length\n",
        "        self.total_freqs = (self.length // 2) + 1\n",
        "\n",
        "        self.in_channels = in_channels\n",
        "        self.out_channels = out_channels\n",
        "\n",
        "        self.band = band  # zero indexed\n",
        "        self.num_bands = num_bands\n",
        "\n",
        "        self.num_freqs = self.total_freqs // self.num_bands + (self.total_freqs % self.num_bands if self.band == self.num_bands - 1 else 0)\n",
        "\n",
        "        self.start = self.band * (self.total_freqs // self.num_bands)\n",
        "        self.end = self.start + self.num_freqs\n",
        "\n",
        "\n",
        "        # case: from other frequencies\n",
        "        self.weight = nn.Parameter(torch.empty((self.num_freqs, in_channels, out_channels), dtype=torch.cfloat))\n",
        "        self.bias = nn.Parameter(torch.empty((self.num_freqs, out_channels), dtype=torch.cfloat))\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def forward(self, input):\n",
        "        # input - b t d\n",
        "        b, t, _ = input.shape\n",
        "        input_fft = fft.rfft(input, dim=1)\n",
        "        output_fft = torch.zeros(b, t // 2 + 1, self.out_channels, device=input.device, dtype=torch.cfloat)\n",
        "        output_fft[:, self.start:self.end] = self._forward(input_fft)\n",
        "        return fft.irfft(output_fft, n=input.size(1), dim=1)\n",
        "\n",
        "    def _forward(self, input):\n",
        "        output = torch.einsum('bti,tio->bto', input[:, self.start:self.end], self.weight)\n",
        "        return output + self.bias\n",
        "\n",
        "    def reset_parameters(self) -> None:\n",
        "        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))\n",
        "        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight)\n",
        "        bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0\n",
        "        nn.init.uniform_(self.bias, -bound, bound)\n",
        "\n",
        "class OfficialCoSTEncoder(nn.Module):\n",
        "    def __init__(self, input_dims, output_dims,\n",
        "                 kernels,\n",
        "                 length,\n",
        "                 d_model = 320,\n",
        "                 hidden_dims=64, depth=4,\n",
        "                 enc_in = 1, covariate_size = 7,\n",
        "                 mask_mode='binomial',\n",
        "                 device = 'cuda',\n",
        "                 backbone = 'NLinear'):\n",
        "        super().__init__()\n",
        "\n",
        "        component_dims = output_dims // 2\n",
        "\n",
        "        self.input_dims = input_dims\n",
        "        self.output_dims = output_dims\n",
        "        self.component_dims = component_dims\n",
        "        self.hidden_dims = hidden_dims\n",
        "        self.mask_mode = mask_mode\n",
        "        self.input_fc = nn.Linear(input_dims, hidden_dims)\n",
        "        self.device = device\n",
        "        self.enc_in = enc_in\n",
        "        self.covariate_size = covariate_size\n",
        "\n",
        "        self.repr_dropout = nn.Dropout(p=0.1)\n",
        "\n",
        "        self.kernels = kernels\n",
        "\n",
        "        self.backbone = backbone\n",
        "\n",
        "        # Select backbone encoder\n",
        "        if backbone == \"NLinear\":\n",
        "            self.enc = NLinear(self.input_dims, d_model, depth, individual=True)\n",
        "        elif backbone == \"Pyraformer\":\n",
        "            self.enc = PyraformerEncoder(d_model, enc_in=self.enc_in,\n",
        "                                         covariate_size=self.covariate_size,device = device)\n",
        "        elif backbone == \"Linear\":\n",
        "            self.enc = nn.Sequential(nn.Linear(input_dims,d_model))\n",
        "            for i in range(depth-1):\n",
        "                self.enc.append(nn.Linear(d_model,d_model))\n",
        "        elif backbone == \"TCN\":\n",
        "            self.enc = DilatedConvEncoder(\n",
        "                hidden_dims,\n",
        "                [hidden_dims] * depth + [output_dims],\n",
        "                kernel_size=3\n",
        "            )\n",
        "        else:\n",
        "            raise Exception(\"Backbone not implemented\")\n",
        "\n",
        "        self.tfd = nn.ModuleList(\n",
        "            [nn.Conv1d(output_dims, component_dims, k, padding=k-1) for k in kernels]\n",
        "        )\n",
        "\n",
        "        self.sfd = nn.ModuleList(\n",
        "            [BandedFourierLayer(output_dims, component_dims, b, 1, length=length) for b in range(1)]\n",
        "        )\n",
        "\n",
        "    def encode(self, data_shape, loader, batch_size = 256, sliding_length=1, padding=200):\n",
        "\n",
        "        encoding_window = None\n",
        "        slicing = None\n",
        "\n",
        "        n_samples, ts_l, _ = data_shape\n",
        "\n",
        "        org_training = self.training\n",
        "        self.eval()\n",
        "\n",
        "        with torch.no_grad():\n",
        "            output = []\n",
        "            for batch in tqdm(loader, desc=\"data encoding\"):\n",
        "                x = batch[0]\n",
        "                reprs = []\n",
        "                # x = x.to(device)\n",
        "                if n_samples < batch_size:\n",
        "                    calc_buffer = []\n",
        "                    calc_buffer_l = 0\n",
        "                # self.log.debug(type(x))\n",
        "                # self.log.debug(f\"shape of batch: {x.shape}\")\n",
        "\n",
        "                # if self.batch_size != batch_size:\n",
        "                #     self.log.debug(\"different batch size return same batch\")\n",
        "                #     return batch\n",
        "\n",
        "                for i in tqdm(range(0, ts_l, sliding_length),\n",
        "                              desc = \"sequence encoding\"):\n",
        "                    l = i - padding # sliding_padding=200\n",
        "                    r = i + sliding_length\n",
        "                    # self.log.debug(x.device)\n",
        "                    x_sliding = pad_nan(\n",
        "                        x[:, max(l, 0) : min(r, ts_l)],\n",
        "                        left=-l if l<0 else 0,\n",
        "                        right=r-ts_l if r>ts_l else 0,\n",
        "                        dim=1\n",
        "                    )\n",
        "                    if n_samples < batch_size:\n",
        "                        if calc_buffer_l + n_samples > batch_size:\n",
        "                            out = self._eval_with_pooling(\n",
        "                                torch.cat(calc_buffer, dim=0)\n",
        "                            )\n",
        "                            reprs += torch.split(out, n_samples)\n",
        "                            calc_buffer = []\n",
        "                            calc_buffer_l = 0\n",
        "                        calc_buffer.append(x_sliding)\n",
        "                        calc_buffer_l += n_samples\n",
        "                    else:\n",
        "                        reprs.append(self._eval_with_pooling(x_sliding))\n",
        "\n",
        "                if n_samples < batch_size:\n",
        "                    if calc_buffer_l > 0:\n",
        "                        out = self._eval_with_pooling(\n",
        "                            torch.cat(calc_buffer, dim=0)\n",
        "                        )\n",
        "                        reprs += torch.split(out, n_samples)\n",
        "                        calc_buffer = []\n",
        "                        calc_buffer_l = 0\n",
        "\n",
        "                out = torch.cat(reprs, dim=1)\n",
        "                output.append(out)\n",
        "\n",
        "            output = torch.cat(output, dim=0)\n",
        "\n",
        "        self.train(org_training)\n",
        "        return output.numpy()\n",
        "\n",
        "    def _eval_with_pooling(self, x):\n",
        "        out_t, out_s = self(x.to(self.device, non_blocking=True))\n",
        "        out = torch.cat([out_t[:, -1], out_s[:, -1]], dim=-1)\n",
        "        return einops.rearrange(out.cpu(), 'b d -> b () d')\n",
        "\n",
        "    def forward(self, x, mask='all_true'):  # x: B x T x input_dims\n",
        "        \"\"\"\n",
        "        x: (batch, input_size, enc_in + covariate_size)\n",
        "        outputs: (batch, input_size, d_t), (batch, input_size, d_s)\n",
        "        \"\"\"\n",
        "        nan_mask = ~x.isnan().any(axis=-1)\n",
        "        x[~nan_mask] = 0\n",
        "        # x = self.input_fc(x)\n",
        "\n",
        "\n",
        "        if self.backbone == \"NLinear\":\n",
        "            x = x.transpose(1, 2)\n",
        "            x = self.enc(x)\n",
        "        elif self.backbone == \"Pyraformer\":\n",
        "            x = self.enc(x)\n",
        "            x = x.transpose(1, 2)\n",
        "        elif self.backbone == \"Linear\":\n",
        "            x = self.enc(x)\n",
        "            x = x.transpose(1, 2)\n",
        "        elif self.backbone == \"TCN\":\n",
        "            x = self.input_fc(x)\n",
        "            x = x.transpose(1, 2)\n",
        "            x = self.enc(x)\n",
        "        else:\n",
        "            raise Exception(\"Backbone not implemented\")\n",
        "\n",
        "        trend = []\n",
        "        for idx, mod in enumerate(self.tfd):\n",
        "            out = mod(x)  # b d t\n",
        "            if self.kernels[idx] != 1:\n",
        "                out = out[..., :-(self.kernels[idx] - 1)]\n",
        "            trend.append(out.transpose(1, 2))  # b t d\n",
        "        trend = einops.reduce(\n",
        "            einops.rearrange(trend, 'list b t d -> list b t d'),\n",
        "            'list b t d -> b t d', 'mean'\n",
        "        )\n",
        "\n",
        "        x = x.transpose(1, 2)  # B x T x Co\n",
        "\n",
        "        season = []\n",
        "        for mod in self.sfd:\n",
        "            out = mod(x)  # b t d\n",
        "            season.append(out)\n",
        "        season = season[0]\n",
        "\n",
        "        return trend, self.repr_dropout(season)\n"
      ],
      "metadata": {
        "id": "fm99hMyEHoYN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## CoST Model"
      ],
      "metadata": {
        "id": "FfsnoB2YQJOq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class OfficialCoSTModel(pl.LightningModule):\n",
        "    def __init__(self,\n",
        "                 encoder_q: nn.Module, encoder_k: nn.Module,\n",
        "                 kernels,\n",
        "                 dim = 128,\n",
        "                 alpha = 0.05,\n",
        "                 K = 65536,\n",
        "                 m = 0.999,\n",
        "                 T = 0.07,\n",
        "                 lr = 1e-3, om = 0.9, wd = 1e-4,\n",
        "                 epochs = 10, n_iters = 600, enc = \"Pyraformer\",\n",
        "                 max_train_length = 201):\n",
        "        super().__init__()\n",
        "\n",
        "        self.save_hyperparameters()\n",
        "\n",
        "        self.cum_loss = 0\n",
        "        self.n_epoch_iters = 0\n",
        "        self.epochs = epochs\n",
        "        self.n_iters = n_iters\n",
        "\n",
        "        self.K = K\n",
        "        self.m = m\n",
        "        self.T = T\n",
        "\n",
        "        self.kernels = kernels\n",
        "\n",
        "        self.alpha = alpha\n",
        "\n",
        "        self.lr = lr\n",
        "        self.om = om\n",
        "        self.wd = wd\n",
        "        self.max_train_length = max_train_length\n",
        "\n",
        "        self.encoder_q = encoder_q\n",
        "        self.encoder_k = encoder_k\n",
        "\n",
        "        # create the encoders\n",
        "        self.head_q = nn.Sequential(\n",
        "            nn.Linear(dim, dim),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(dim, dim)\n",
        "        )\n",
        "        self.head_k = nn.Sequential(\n",
        "            nn.Linear(dim, dim),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(dim, dim)\n",
        "        )\n",
        "\n",
        "        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):\n",
        "            param_k.data.copy_(param_q.data)  # initialize\n",
        "            param_k.requires_grad = False  # not update by gradient\n",
        "        for param_q, param_k in zip(self.head_q.parameters(), self.head_k.parameters()):\n",
        "            param_k.data.copy_(param_q.data)  # initialize\n",
        "            param_k.requires_grad = False  # not update by gradient\n",
        "\n",
        "        self.register_buffer('queue', F.normalize(torch.randn(dim, K), dim=0))\n",
        "        self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))\n",
        "\n",
        "\n",
        "    def compute_loss(self, q, k, k_negs):\n",
        "        # compute logits\n",
        "        # positive logits: Nx1\n",
        "        l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)\n",
        "        # negative logits: NxK\n",
        "        l_neg = torch.einsum('nc,ck->nk', [q, k_negs])\n",
        "\n",
        "        # logits: Nx(1+K)\n",
        "        logits = torch.cat([l_pos, l_neg], dim=1)\n",
        "\n",
        "        # apply temperature\n",
        "        logits /= self.T\n",
        "\n",
        "        # labels: positive key indicators - first dim of each batch\n",
        "        labels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()\n",
        "        loss = F.cross_entropy(logits, labels)\n",
        "\n",
        "        return loss\n",
        "\n",
        "    def convert_coeff(self, x, eps=1e-6):\n",
        "        amp = torch.sqrt((x.real + eps).pow(2) + (x.imag + eps).pow(2))\n",
        "        phase = torch.atan2(x.imag, x.real + eps)\n",
        "        return amp, phase\n",
        "\n",
        "    def instance_contrastive_loss(self, z1, z2):\n",
        "        B, T = z1.size(0), z1.size(1)\n",
        "        z = torch.cat([z1, z2], dim=0)  # 2B x T x C\n",
        "        z = z.transpose(0, 1)  # T x 2B x C\n",
        "        sim = torch.matmul(z, z.transpose(1, 2))  # T x 2B x 2B\n",
        "        logits = torch.tril(sim, diagonal=-1)[:, :, :-1]  # T x 2B x (2B-1)\n",
        "        logits += torch.triu(sim, diagonal=1)[:, :, 1:]\n",
        "        logits = -F.log_softmax(logits, dim=-1)\n",
        "\n",
        "        i = torch.arange(B, device=z1.device)\n",
        "        loss = (logits[:, i, B + i - 1].mean() + logits[:, B + i, i].mean()) / 2\n",
        "        return loss\n",
        "\n",
        "    def forward(self, x_q, x_k):\n",
        "        # compute query features\n",
        "        rand_idx = np.random.randint(0, x_q.shape[1])\n",
        "\n",
        "        q_t, q_s = self.encoder_q(x_q)\n",
        "        if q_t is not None:\n",
        "            q_t = F.normalize(self.head_q(q_t[:, rand_idx]), dim=-1)\n",
        "\n",
        "        # compute key features\n",
        "        with torch.no_grad():  # no gradient for keys\n",
        "            self._momentum_update_key_encoder()  # update key encoder\n",
        "            k_t, k_s = self.encoder_k(x_k)\n",
        "            if k_t is not None:\n",
        "                k_t = F.normalize(self.head_k(k_t[:, rand_idx]), dim=-1)\n",
        "\n",
        "        loss = 0\n",
        "\n",
        "        loss += self.compute_loss(q_t, k_t, self.queue.clone().detach())\n",
        "        self._dequeue_and_enqueue(k_t)\n",
        "\n",
        "        q_s = F.normalize(q_s, dim=-1)\n",
        "        _, k_s = self.encoder_q(x_k)\n",
        "        k_s = F.normalize(k_s, dim=-1)\n",
        "\n",
        "        q_s_freq = fft.rfft(q_s, dim=1)\n",
        "        k_s_freq = fft.rfft(k_s, dim=1)\n",
        "        q_s_amp, q_s_phase = self.convert_coeff(q_s_freq)\n",
        "        k_s_amp, k_s_phase = self.convert_coeff(k_s_freq)\n",
        "\n",
        "        seasonal_loss = self.instance_contrastive_loss(q_s_amp, k_s_amp) + \\\n",
        "                        self.instance_contrastive_loss(q_s_phase,k_s_phase)\n",
        "        loss += (self.alpha * (seasonal_loss/2))\n",
        "\n",
        "        return loss\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def _momentum_update_key_encoder(self):\n",
        "        \"\"\"\n",
        "        Momentum update for key encoder\n",
        "        \"\"\"\n",
        "        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):\n",
        "            param_k.data = param_k.data * self.m + param_q.data * (1 - self.m)\n",
        "        for param_q, param_k in zip(self.head_q.parameters(), self.head_k.parameters()):\n",
        "            param_k.data = param_k.data * self.m + param_q.data * (1 - self.m)\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def _dequeue_and_enqueue(self, keys):\n",
        "        batch_size = keys.shape[0]\n",
        "\n",
        "        ptr = int(self.queue_ptr)\n",
        "        assert self.K % batch_size == 0\n",
        "\n",
        "        # replace keys at ptr (dequeue and enqueue)\n",
        "        self.queue[:, ptr:ptr + batch_size] = keys.T\n",
        "\n",
        "        ptr = (ptr + batch_size) % self.K\n",
        "        self.queue_ptr[0] = ptr\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        # set parameters of SGD\n",
        "        optimizer = torch.optim.SGD([p for p in self.parameters() if p.requires_grad],\n",
        "                                    lr = self.lr, momentum = self.om, weight_decay = self.wd)\n",
        "        # optimizer = torch.optim.Adam([p for p in self.parameters() if p.requires_grad],\n",
        "        #                             lr = self.lr)\n",
        "        # cosine annelling is a wrapper for SGD\n",
        "        # cosine_anneling = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer = optimizer, T_max = 100)\n",
        "        return optimizer\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        return self.train_dataloader\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        x_q, x_k = batch\n",
        "        if self.max_train_length is not None and x_q.size(1) > self.max_train_length:\n",
        "            window_offset = np.random.randint(x_q.size(1) - self.max_train_length + 1)\n",
        "            x_q = x_q[:, window_offset : window_offset + self.max_train_length]\n",
        "            x_k = x_k[:, window_offset : window_offset + self.max_train_length]\n",
        "\n",
        "        loss = self.forward(x_q, x_k)\n",
        "\n",
        "        self.cum_loss += loss.item()\n",
        "        self.n_epoch_iters += 1\n",
        "\n",
        "        if self.n_iters is not None:\n",
        "            optimizer = self.optimizers().optimizer\n",
        "            self._adjust_learning_rate(self.lr, optimizer, self.global_step, self.n_iters)\n",
        "\n",
        "        # logs metrics for each training_step,\n",
        "        # and the average across the epoch, to the progress bar and logger\n",
        "        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True)\n",
        "        wandb.log({\"train\": {\"loss\":loss}})\n",
        "        return loss\n",
        "\n",
        "    def on_train_epoch_end(self):\n",
        "        # self.log(self.log(\"cum_loss\", self.cum_loss, on_step=True, on_epoch=True, prog_bar=True, logger = True))\n",
        "\n",
        "        wandb.log({\"train\": {\"epoch\": self.current_epoch ,\"cum_loss_epoch\": self.cum_loss / self.n_epoch_iters}})\n",
        "        # adjust learning rate\n",
        "        optimizer = self.optimizers().optimizer\n",
        "        if self.epochs is not None:\n",
        "            self._adjust_learning_rate(self.lr, optimizer, self.current_epoch, self.epochs)\n",
        "        self.n_epoch_iters = 0\n",
        "        self.cum_loss = 0\n",
        "\n",
        "    def _adjust_learning_rate(self, lr, optimizer, epoch, epochs):\n",
        "        \"\"\"Decay the learning rate based on schedule\"\"\"\n",
        "        lr *= 0.5 * (1. + math.cos(math.pi * epoch / epochs))\n",
        "        for param_group in optimizer.param_groups:\n",
        "            param_group['lr'] = lr"
      ],
      "metadata": {
        "id": "TmbzDjrkQLDH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## My Model 🔽"
      ],
      "metadata": {
        "id": "YP8y2ZUZP0ew"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## CoSpy Model"
      ],
      "metadata": {
        "id": "niFGAd70x41j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Constrastive model similar to MoCo\n",
        "\"\"\"\n",
        "https://arxiv.org/pdf/1911.05722.pdf\n",
        "https://github.com/facebookresearch/moco/blob/main/moco/builder.py\n",
        "\"\"\"\n",
        "\n",
        "class CoSpyModel(pl.LightningModule):\n",
        "    def __init__(self, max_train_length, comp_dimension = 160, alpha = 5e-4, K = 65536, m = 0.999, T = 0.07,\n",
        "                 lr = 1e-3, om = 0.9, wd = 1e-4, epochs = 10, n_iters = 600, device = 'cpu', enc = \"Pyraformer\"):\n",
        "        super(CoSpyModel, self).__init__()\n",
        "\n",
        "        self.cum_loss = 0\n",
        "        self.n_epoch_iters = 0\n",
        "        self.epochs = epochs\n",
        "        self.n_iters = n_iters\n",
        "\n",
        "        self.input_size = max_train_length\n",
        "        self.max_train_length = max_train_length\n",
        "\n",
        "        self.K = K\n",
        "        self.m = m\n",
        "        self.T = T\n",
        "\n",
        "        self.alpha = alpha\n",
        "\n",
        "        self.lr = lr\n",
        "        self.om = om\n",
        "        self.wd = wd\n",
        "\n",
        "        self.encoder_q = CoSpyEncoder(self.input_size, device = device, enc = enc)\n",
        "        self.encoder_k = copy.deepcopy(self.encoder_q)\n",
        "\n",
        "        # projections head for queries and keyes\n",
        "        self.head_q = nn.Sequential(\n",
        "            nn.Linear(comp_dimension, comp_dimension),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(comp_dimension, comp_dimension)\n",
        "        )\n",
        "        self.head_k = nn.Sequential(\n",
        "            nn.Linear(comp_dimension, comp_dimension),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(comp_dimension, comp_dimension)\n",
        "        )\n",
        "\n",
        "        # initialize the parameters of the keyes encoder and projection head\n",
        "        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):\n",
        "            param_k.data.copy_(param_q.data)\n",
        "            param_k.requires_grad = False # the keyes encoder will be updated by the momentum update\n",
        "\n",
        "        for param_q, param_k in zip(self.head_q.parameters(), self.head_k.parameters()):\n",
        "            param_k.data.copy_(param_q.data)\n",
        "            param_k.requires_grad = False # the head_k will be updated by the momentum update\n",
        "\n",
        "        # register a dictionary buffer as a queue (decouped from the minibatch size)\n",
        "        # https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer\n",
        "        self.register_buffer('queue', F.normalize(torch.randn(comp_dimension, K), dim=0))\n",
        "        self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))\n",
        "\n",
        "        self.save_hyperparameters(ignore=['encoder_q', 'encoder_k'])\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        # set parameters of SGD\n",
        "        optimizer = torch.optim.SGD([p for p in self.parameters() if p.requires_grad],\n",
        "                                    lr = self.lr, momentum = self.om, weight_decay = self.wd)\n",
        "        # cosine annelling is a wrapper for SGD\n",
        "        # cosine_anneling = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer = optimizer, T_max = 100)\n",
        "        return optimizer\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        return self.train_dataloader\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def _momentum_update_key_encoder(self):\n",
        "        \"\"\"\n",
        "        Momentum update for key encoder\n",
        "        \"\"\"\n",
        "        for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):\n",
        "            param_k.data = param_k.data * self.m + param_q.data * (1 - self.m)\n",
        "\n",
        "        for param_q, param_k in zip(self.head_q.parameters(), self.head_k.parameters()):\n",
        "            param_k.data = param_k.data * self.m + param_q.data * (1 - self.m)\n",
        "\n",
        "    def compute_loss(self, q, k, k_negs):\n",
        "        # compute logits\n",
        "        # positive logits: Bx1 (one timestamp as postive)\n",
        "        l_pos = einops.einsum(q, k, 'b c,b c->b').unsqueeze(-1)\n",
        "        # negative logits: BxK\n",
        "        l_neg = einops.einsum(q, k_negs, 'b c,c k->b k')\n",
        "\n",
        "        # logits: Bx(1+K)\n",
        "        logits = torch.cat([l_pos, l_neg], dim=1)\n",
        "\n",
        "        # apply temperature\n",
        "        logits /= self.T\n",
        "\n",
        "        # labels: positive key indicators - first dim of each batch (it will be considered the positive sample)\n",
        "        # so we can consider this as a classification problem and use the CE\n",
        "        labels = torch.zeros(logits.shape[0], dtype=torch.long, device = logits.device)\n",
        "        loss = F.cross_entropy(logits, labels)\n",
        "\n",
        "        return loss\n",
        "\n",
        "    def get_polar(self, x, eps=1e-6):\n",
        "        amp = torch.sqrt((x.real + eps).pow(2) + (x.imag + eps).pow(2))\n",
        "        phase = torch.atan2(x.imag, x.real + eps)\n",
        "\n",
        "        return amp, phase\n",
        "\n",
        "    def instance_contrastive_loss(self, z1, z2):\n",
        "        B = z1.shape[0]\n",
        "        z = torch.cat([z1, z2], dim=0)  # 2B x F x d_s\n",
        "        z = einops.rearrange(z, 'b f d_s -> f b d_s')  # F x 2B x d_s\n",
        "        sim = einops.einsum(z, z, 'f b_1 d_s, f b_2 d_s -> f b_1 b_2')  # F x 2B x 2B\n",
        "        logits = torch.tril(sim, diagonal=-1)[:, :, :-1]  # F x 2B x (2B-1)\n",
        "        logits += torch.triu(sim, diagonal=1)[:, :, 1:]\n",
        "        logits = -F.log_softmax(logits, dim=-1)\n",
        "        # log.debug(logits)\n",
        "\n",
        "        i = torch.arange(B)\n",
        "        loss = (logits[:, i, B + i - 1].mean() + logits[:, B + i, i].mean()) / 2\n",
        "        return loss\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def _dequeue_and_enqueue(self, keys):\n",
        "        batch_size = keys.shape[0]\n",
        "\n",
        "        ptr = int(self.queue_ptr)\n",
        "        assert self.K % batch_size == 0, \"K must be a multiple of batch_size\"\n",
        "\n",
        "        # replace keys at ptr (dequeue and enqueue)\n",
        "        self.queue[:, ptr:ptr + batch_size] = keys.T\n",
        "\n",
        "        ptr = (ptr + batch_size) % self.K\n",
        "        self.queue_ptr[0] = ptr\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        x_q, x_k = batch\n",
        "        if self.max_train_length is not None and x_q.size(1) > self.max_train_length:\n",
        "            window_offset = np.random.randint(x_q.size(1) - self.max_train_length + 1)\n",
        "            x_q = x_q[:, window_offset : window_offset + self.max_train_length]\n",
        "            x_k = x_k[:, window_offset : window_offset + self.max_train_length]\n",
        "\n",
        "        loss = self.forward(x_q, x_k)\n",
        "\n",
        "        self.cum_loss += loss.item()\n",
        "        self.n_epoch_iters += 1\n",
        "\n",
        "        if self.n_iters is not None:\n",
        "            optimizer = self.optimizers().optimizer\n",
        "            self._adjust_learning_rate(self.lr, optimizer, self.global_step, self.n_iters)\n",
        "\n",
        "        # logs metrics for each training_step,\n",
        "        # and the average across the epoch, to the progress bar and logger\n",
        "        self.log(\"train_loss\", loss, on_step=True, on_epoch=True, prog_bar=True)\n",
        "        wandb.log({\"train\": {\"loss\":loss}})\n",
        "        return loss\n",
        "\n",
        "    def on_train_epoch_end(self):\n",
        "        # self.log(self.log(\"cum_loss\", self.cum_loss, on_step=True, on_epoch=True, prog_bar=True, logger = True))\n",
        "        wandb.log({\"train\": {\"epoch\": self.current_epoch ,\"cum_loss_epoch\": self.cum_loss / self.n_epoch_iters}})\n",
        "        # adjust learning rate\n",
        "        optimizer = self.optimizers().optimizer\n",
        "        if self.epochs is not None:\n",
        "            self._adjust_learning_rate(self.lr, optimizer, self.current_epoch, self.epochs)\n",
        "        self.n_epoch_iters = 0\n",
        "        self.cum_loss = 0\n",
        "\n",
        "\n",
        "\n",
        "    def _adjust_learning_rate(self, lr, optimizer, epoch, epochs):\n",
        "        \"\"\"Decay the learning rate based on schedule\"\"\"\n",
        "        lr *= 0.5 * (1. + math.cos(math.pi * epoch / epochs))\n",
        "        for param_group in optimizer.param_groups:\n",
        "            param_group['lr'] = lr\n",
        "\n",
        "\n",
        "    # def validation_step(self, batch, batch_idx):\n",
        "    #     x_q, x_k = batch\n",
        "    #     loss = self.forward(x_q, x_k)\n",
        "\n",
        "    #     # logs metrics for each training_step,\n",
        "    #     # and the average across the epoch, to the progress bar and logger\n",
        "    #     self.log(\"val_loss\", loss)\n",
        "    #     return loss\n",
        "\n",
        "    # def test_step(self, batch, batch_idx):\n",
        "    #     x_q, x_k = batch\n",
        "    #     loss = self.forward(x_q, x_k)\n",
        "\n",
        "    #     # logs metrics for each training_step,\n",
        "    #     # and the average across the epoch, to the progress bar and logger\n",
        "    #     self.log(\"test_loss\", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)\n",
        "    #     return loss\n",
        "\n",
        "    def forward(self, x_q, x_k):\n",
        "        \"\"\"\n",
        "        x_q, x_k: (batch, input_size, enc_in + covariate_size)\n",
        "        \"\"\"\n",
        "        # select a random timestamp\n",
        "        rand_idx = np.random.randint(0, self.input_size)\n",
        "\n",
        "        # trend and seasonal queries\n",
        "        q_t, q_s = self.encoder_q(x_q)\n",
        "\n",
        "        if q_t is not None:\n",
        "            q_t = F.normalize(self.head_q(q_t[:, rand_idx]), dim=-1)\n",
        "\n",
        "        # compute key features\n",
        "        with torch.no_grad():  # no gradient update for keys (momentum update will be used)\n",
        "            self._momentum_update_key_encoder()  # update key encoder using momentum\n",
        "            k_t, k_s = self.encoder_k(x_k)\n",
        "            if k_t is not None:\n",
        "                k_t = F.normalize(self.head_k(k_t[:, rand_idx]), dim=-1)\n",
        "\n",
        "        loss = 0\n",
        "\n",
        "        loss += self.compute_loss(q_t, k_t, self.queue.clone().detach())\n",
        "        self._dequeue_and_enqueue(k_t)\n",
        "\n",
        "        q_s = F.normalize(q_s, dim=-1)\n",
        "        _, k_s = self.encoder_q(x_k)\n",
        "        k_s = F.normalize(k_s, dim=-1)\n",
        "\n",
        "        # the frequency and phase lost must be computed in the frequency domain\n",
        "        q_s_freq = torch.fft.rfft(q_s, dim=1)\n",
        "        k_s_freq = torch.fft.rfft(k_s, dim=1)\n",
        "        q_s_amp, q_s_phase = self.get_polar(q_s_freq)\n",
        "        k_s_amp, k_s_phase = self.get_polar(k_s_freq)\n",
        "\n",
        "        seasonal_loss = self.instance_contrastive_loss(q_s_amp, k_s_amp) + \\\n",
        "                        self.instance_contrastive_loss(q_s_phase,k_s_phase)\n",
        "        loss += (self.alpha * (seasonal_loss/2))\n",
        "\n",
        "        return loss"
      ],
      "metadata": {
        "id": "xSBWEMomx4KG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Train"
      ],
      "metadata": {
        "id": "DYNOokNm3LCY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# train function\n",
        "def train(batch_size, datamodule, model, model_name, max_epochs = None, max_steps = -1, checkpoint_every_n_epochs = 50,\n",
        "          check_val_every_n_epoch = 5, gradient_clip_val=0, resume_training = True, load_model = False,\n",
        "          enable_checkpoint = True, monitor_metric = \"val_loss\", checkpoint_dir = None,\n",
        "          log_flag = False, logs_dir = None,\n",
        "          early_stopping = True, deterministic = False, profiler = None, find_lr = False):\n",
        "\n",
        "    # check monitor metric\n",
        "    assert monitor_metric in [\"train_loss\", \"vall_loss\"], \"metric to monitor is invalid\"\n",
        "\n",
        "    # initialize callbacks array\n",
        "    callbacks = [TQDMProgressBar(refresh_rate=20)]\n",
        "\n",
        "    # add checkpoints to callbacks\n",
        "    checkpoint_callback = None\n",
        "    if enable_checkpoint and checkpoint_dir is not None:\n",
        "        checkpoint_callback = ModelCheckpoint(dirpath=checkpoint_dir,  monitor = monitor_metric, filename=model_name + '{epoch:02d}-{' + monitor_metric + ':.2f}',\n",
        "                                         save_last =True, every_n_epochs = checkpoint_every_n_epochs, save_on_train_epoch_end = True)\n",
        "        callbacks.append(checkpoint_callback)\n",
        "\n",
        "    # add early stopping to the callbacks\n",
        "    if early_stopping:\n",
        "        callbacks.append(EarlyStopping(monitor=\"val_loss\", min_delta = 0.1, patience = 3, mode=\"min\", check_on_train_epoch_end = False))\n",
        "\n",
        "    # define the logger object\n",
        "    logger = None\n",
        "    if log_flag:\n",
        "        # logger = TensorBoardLogger(logs_dir, name=model_name)\n",
        "        wandb_logger = WandbLogger(name = model_name, log_model = 'all')\n",
        "\n",
        "    # create the Trainer\n",
        "    trainer = pl.Trainer(enable_checkpointing=enable_checkpoint, devices=1, accelerator=\"auto\",\n",
        "                         max_epochs=max_epochs, max_steps=max_steps, logger=logger, callbacks=callbacks,  ## remove max_step\n",
        "                         check_val_every_n_epoch = check_val_every_n_epoch, gradient_clip_val=gradient_clip_val,## remove\n",
        "                         deterministic = deterministic, profiler = profiler)\n",
        "\n",
        "    if find_lr:\n",
        "        model.train_dataloader\n",
        "        tuner = Tuner(trainer)\n",
        "        # Run learning rate finder\n",
        "        lr_finder = tuner.lr_find(model, datamodule = datamodule, early_stop_threshold=None)\n",
        "\n",
        "        # Plot with\n",
        "        fig = lr_finder.plot(suggest=True)\n",
        "        fig.show()\n",
        "        # Pick point based on plot, or get suggestion\n",
        "        new_lr = lr_finder.suggestion()\n",
        "        log.info(f\"the suggested lr is {new_lr}\")\n",
        "        model.lr = new_lr\n",
        "\n",
        "    ckpt_path = None\n",
        "    if resume_training:\n",
        "        ckpt_path = checkpoint_dir + \"/last.ckpt\"\n",
        "    trainer.fit(ckpt_path = ckpt_path, model=model, datamodule=datamodule)\n",
        "    if checkpoint_callback is not None:\n",
        "        log.info(checkpoint_callback.best_model_path)\n",
        "\n",
        "    return trainer"
      ],
      "metadata": {
        "id": "q-936er_-snn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#  Encoding"
      ],
      "metadata": {
        "id": "6lzsSVC6kmJy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def prepare_encoding(datamodule):\n",
        "    datamodule.setup(\"encoding\")\n",
        "    loader = datamodule.encode_dataloader()\n",
        "    return loader, datamodule.data.shape\n",
        "\n",
        "def encode(model, data_shape, loader, batch_size, device, save_path):\n",
        "    model.to(device)\n",
        "    model = model.encoder_q\n",
        "    res = model.encode(data_shape, loader, batch_size = batch_size, padding = PADDING)\n",
        "    file_name = f\"encoding_{time.time()}.pkl\"\n",
        "    pkl_save(f'{save_path}/{file_name}', res)\n",
        "    pkl_save(f'{save_path}/last.pkl', res)\n",
        "    log.info(f\"encoding {file_name} saved\")\n",
        "\n",
        "    return res"
      ],
      "metadata": {
        "id": "BIhjMuNq2nOY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Forecasting Evaluation"
      ],
      "metadata": {
        "id": "sLMxIOWrOckk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Helper functions\n",
        "def generate_pred_samples(features, data, pred_len, drop=0):\n",
        "    n = data.shape[1]\n",
        "    features = features[:, :-pred_len]\n",
        "    labels = np.stack([ data[:, i:1+n+i-pred_len] for i in range(pred_len)], axis=2)[:, 1:]\n",
        "    features = features[:, drop:]\n",
        "    labels = labels[:, drop:]\n",
        "    return features.reshape(-1, features.shape[-1]), labels.reshape(-1, labels.shape[2]*labels.shape[3])\n",
        "\n",
        "def fit_ridge(train_features, train_y, valid_features, valid_y, MAX_SAMPLES=100000):\n",
        "    # If the training set is too large, subsample MAX_SAMPLES examples\n",
        "    if train_features.shape[0] > MAX_SAMPLES:\n",
        "        split = train_test_split(\n",
        "            train_features, train_y,\n",
        "            train_size=MAX_SAMPLES, random_state=0\n",
        "        )\n",
        "        train_features = split[0]\n",
        "        train_y = split[2]\n",
        "    if valid_features.shape[0] > MAX_SAMPLES:\n",
        "        split = train_test_split(\n",
        "            valid_features, valid_y,\n",
        "            train_size=MAX_SAMPLES, random_state=0\n",
        "        )\n",
        "        valid_features = split[0]\n",
        "        valid_y = split[2]\n",
        "    alphas = [1, 2, 5, 10, 20, 50, 100, 200, 500, 1000] if UNIVARIATE \\\n",
        "        else [0.1, 0.2, 0.5, 1, 2, 5, 10, 20, 50, 100, 200, 500, 1000]\n",
        "    valid_results = []\n",
        "    for alpha in alphas:\n",
        "        # ipdb.set_trace(context=6)\n",
        "        # log.debug(f\"alpha: {alpha}\")\n",
        "        lr = Ridge(alpha=alpha).fit(train_features, train_y)\n",
        "        valid_pred = lr.predict(valid_features)\n",
        "        score = np.sqrt(((valid_pred - valid_y) ** 2).mean()) + np.abs(valid_pred - valid_y).mean()\n",
        "        valid_results.append(score)\n",
        "    best_alpha = alphas[np.argmin(valid_results)]\n",
        "    log.info(f\"best alpha: {best_alpha}\")\n",
        "\n",
        "    lr = Ridge(alpha=best_alpha)\n",
        "    lr.fit(train_features, train_y)\n",
        "    return lr\n",
        "\n",
        "def cal_metrics(pred, target):\n",
        "    return {\n",
        "        'MSE': ((pred - target) ** 2).mean(),\n",
        "        'MAE': np.abs(pred - target).mean()\n",
        "    }"
      ],
      "metadata": {
        "id": "aw_fXa2wRZxr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def eval_forecasting(repr, data, train_slice, valid_slice, test_slice, n_covariate_cols, scaler, padding, pred_lens):\n",
        "    train_repr = repr[:, train_slice]\n",
        "    valid_repr = repr[:, valid_slice]\n",
        "    test_repr = repr[:, test_slice]\n",
        "\n",
        "    train_data = data[:, train_slice, n_covariate_cols:]\n",
        "    valid_data = data[:, valid_slice, n_covariate_cols:]\n",
        "    test_data = data[:, test_slice, n_covariate_cols:]\n",
        "\n",
        "    ours_result = {}\n",
        "    out_log = {}\n",
        "    y_labels_mse = []\n",
        "    y_labels_mae = []\n",
        "    for pred_len in tqdm(pred_lens, desc=f\"forecasting evaluation {pred_lens}\"):\n",
        "        train_features, train_labels = generate_pred_samples(train_repr, train_data, pred_len, drop=padding)\n",
        "        valid_features, valid_labels = generate_pred_samples(valid_repr, valid_data, pred_len)\n",
        "        test_features, test_labels = generate_pred_samples(test_repr, test_data, pred_len)\n",
        "\n",
        "        lr = fit_ridge(train_features, train_labels, valid_features, valid_labels)\n",
        "\n",
        "        test_pred = lr.predict(test_features)\n",
        "\n",
        "        ori_shape = test_data.shape[0], -1, pred_len, test_data.shape[2]\n",
        "        test_pred = test_pred.reshape(ori_shape)\n",
        "        test_labels = test_labels.reshape(ori_shape)\n",
        "\n",
        "        test_shape = test_pred.shape\n",
        "        test_shape_swap = (test_shape[3], test_shape[1], test_shape[2], test_shape[0])\n",
        "        if test_data.shape[0] > 1:\n",
        "            test_pred_inv = scaler.inverse_transform(test_pred.swapaxes(0, 3)\n",
        "                .reshape(-1, test_shape[0])).reshape(test_shape_swap).swapaxes(0, 3)\n",
        "            test_labels_inv = scaler.inverse_transform(test_labels.swapaxes(0, 3)\n",
        "                .reshape(-1, test_shape[0])).reshape(test_shape_swap).swapaxes(0, 3)\n",
        "        else:\n",
        "            test_pred_inv = scaler.inverse_transform(test_pred.reshape(-1, test_shape[3])).reshape(test_shape)\n",
        "            test_labels_inv = scaler.inverse_transform(test_labels.reshape(-1, test_shape[3])).reshape(test_shape)\n",
        "\n",
        "        # out_log[pred_len] = {\n",
        "        #     # 'norm': test_pred,\n",
        "        #     # 'raw': test_pred_inv\n",
        "        #     # 'norm_gt': test_labels,\n",
        "        #     # 'raw_gt': test_labels_inv\n",
        "        # }\n",
        "        ours_result[pred_len] = {\n",
        "            'norm': cal_metrics(test_pred, test_labels),\n",
        "            'raw': cal_metrics(test_pred_inv, test_labels_inv)\n",
        "        }\n",
        "        log.info(ours_result[pred_len])\n",
        "        y_labels_mse.append(ours_result[pred_len]['norm']['MSE'])\n",
        "        y_labels_mae.append(ours_result[pred_len]['norm']['MAE'])\n",
        "\n",
        "\n",
        "    eval_res = {\n",
        "        'ours': ours_result\n",
        "    }\n",
        "    return out_log, eval_res, y_labels_mse, y_labels_mae"
      ],
      "metadata": {
        "id": "0aJRiW2lOhoh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Main"
      ],
      "metadata": {
        "id": "Oqxoxjh9C0rC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "seeds = [1,2,3]\n",
        "results = []\n",
        "\n",
        "for seed in seeds:\n",
        "    # initialize dataset\n",
        "    # datamodule = ElectricityDataModule(datasets_path[ELECTRICITY] + datasets_processed_name[ELECTRICITY], BATCH_SIZE, L)\n",
        "    datamodule = CustomDataModule(DATASET, datasets_path[DATASET] + datasets_processed_name[DATASET],\n",
        "                        BATCH_SIZE, L, encode_batch_size=ENCODE_BATCH_SIZE, univariate=UNIVARIATE)\n",
        "\n",
        "    # set seed if deterministic\n",
        "    if DETERMINISTIC:\n",
        "        seed_everything(seed)\n",
        "    # initialize model, or load an extisting one\n",
        "    if LOAD_MODEL:\n",
        "        log.info(\"loading model...\")\n",
        "        model = OfficialCoSTModel.load_from_checkpoint(CHECKPOINT_FOLDER + \"/last.ckpt\")\n",
        "        net = model.encoder_q\n",
        "    else:\n",
        "        net = OfficialCoSTEncoder(\n",
        "                input_dims=input_dim[DATASET], output_dims=320,\n",
        "                kernels=[1, 2, 4, 8, 16, 32, 64, 128],\n",
        "                enc_in=1, covariate_size=input_dim[DATASET]-1,\n",
        "                length=201,\n",
        "                hidden_dims=64, depth=4,\n",
        "                backbone = BACKBONE\n",
        "            )\n",
        "        model = OfficialCoSTModel(\n",
        "            net,\n",
        "            copy.deepcopy(net),\n",
        "            kernels=[1, 2, 4, 8, 16, 32, 64, 128],\n",
        "            dim=net.component_dims,\n",
        "            alpha=0.0005,\n",
        "            K=256,\n",
        "            max_train_length = 201,\n",
        "            epochs = EPOCHS,\n",
        "            n_iters = ITERS,\n",
        "            lr = LR\n",
        "        )#CoSpyModel(L, epochs = EPOCHS, device = DEVICE, n_iters = ITERS, enc = ENCODER, lr = LR)\n",
        "\n",
        "    if MEMORY_PROFILING:\n",
        "        profiler = PyTorchProfiler()\n",
        "    if TRAIN:\n",
        "        trainer = train(BATCH_SIZE, datamodule, model, MODEL, max_epochs = EPOCHS, max_steps = ITERS, check_val_every_n_epoch = None, load_model = LOAD_MODEL,\n",
        "                        gradient_clip_val=GRADIENT_CLIPPING,\n",
        "            resume_training = RESUME_TRAINING, monitor_metric = \"train_loss\", checkpoint_dir = CHECKPOINT_FOLDER,\n",
        "            early_stopping = False, deterministic = DETERMINISTIC, find_lr = FIND_LR)\n",
        "    if EVALUATE:\n",
        "        encoding_loader, data_shape = prepare_encoding(datamodule)\n",
        "        if LOAD_ENCODE:\n",
        "            log.info(\"load encoding...\")\n",
        "            repr = pkl_load(ENCODING_FOLDER + \"/last.pkl\")\n",
        "        else:\n",
        "            repr = encode(model, data_shape, encoding_loader, ENCODE_BATCH_SIZE, DEVICE, ENCODING_FOLDER)\n",
        "    if EVALUATE:\n",
        "        train_slice = datamodule.train_slice\n",
        "        valid_slice = datamodule.valid_slice\n",
        "        test_slice = datamodule.test_slice\n",
        "        data = datamodule.data\n",
        "        n_covariate_cols = datamodule.n_covariate_cols\n",
        "        scaler = datamodule.scaler\n",
        "        padding = PADDING\n",
        "        pred_lens = datasets_pred_lens[DATASET]\n",
        "\n",
        "        out, eval_res, y_labels_mse, y_labels_mae = eval_forecasting(repr, data, train_slice, valid_slice, test_slice,\n",
        "                                        n_covariate_cols, scaler, padding, pred_lens)\n",
        "\n",
        "        results.append([y_labels_mse, y_labels_mae])\n",
        "\n",
        "        # data = [[x, y] for (x, y) in zip(pred_lens, y_labels_mse)]\n",
        "        # table = wandb.Table(data=data, columns = [\"pred_lens\", \"mse\"])\n",
        "        # wandb.log(\n",
        "        # {f\"{SETTINGS_STRING} forecasting MSE plot\" : wandb.plot.scatter(table, \"pred_lens\", \"mse\",\n",
        "        #     title=f\"{SETTINGS_STRING} forecasting MSE norm plot\")})\n",
        "\n",
        "        # data = [[x, y] for (x, y) in zip(pred_lens, y_labels_mae)]\n",
        "        # table = wandb.Table(data=data, columns = [\"pred_lens\", \"mae\"])\n",
        "        # wandb.log(\n",
        "        # {f\"{SETTINGS_STRING} forecasting MAE plot\" : wandb.plot.scatter(table, \"pred_lens\", \"mae\",\n",
        "        #     title=f\"{SETTINGS_STRING} forecasting MAE norm plot\")})\n",
        "        # wandb.log({\"eval_forecasting\": eval_res})\n",
        "        pkl_save(FORECASTING_RESULT + \"/out.pkl\", out)\n",
        "        pkl_save(FORECASTING_RESULT + \"/eval_res.pkl\", eval_res)"
      ],
      "metadata": {
        "id": "AvftO5iOK9ZR",
        "outputId": "3ee3e547-a2b0-40d7-eda5-298f072a3275",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 445,
          "referenced_widgets": [
            "fa20273c28f14a6fb408d382d4ac98e6",
            "81488352900249259e4c0f3f116fa08c",
            "ff494d4ed0464ef48bd780aa333f8902",
            "ca08c18f3af64124bc1cf2bc45121c1e",
            "f2bf006d8412490f8cea31c6a589ece2",
            "7b8bf6d3ed804d75b2ea50494048bd37",
            "a324c27f1e2d4c5b8b547899a50aa13f",
            "eca61e01d24144b48c2eff401765b1c8",
            "4b1136aa31cd45c38d3c55b2fed75a46",
            "479c27d780b84af7af586d7ad754086c",
            "4453bcf9179b4eb18d21750e3717ca7a",
            "43eb6a5cb7f44ee89a72abf0cd5c9723",
            "e66a5509d30b43ff80d0b0af3c6f9c8d",
            "00f66ad6e8b34fe19c1a1bf65159451c",
            "c9f3af47ffb84c259ffc633e554c2e0f",
            "14789f347e174395ac6f314a38afb721",
            "c5da477f433f40349751688dfeddfb8e",
            "0be77f4742cc419991f65c85e173c02d",
            "555016b5553446a595d0d582f750119b",
            "77126262e4804dd293c348a73451c92d",
            "4b5b599c1ef649c29c48c41dc5ca86e8",
            "cbdc4be27e694ef49e4f02de28414625",
            "17e0c718123245e0894cb62237f054ef",
            "0cb3cb25168c423390185db0b39538ba",
            "5989a8ecbcc34145965922753d3d6648",
            "ec203b01b7f34051913482e8851b24e2",
            "cda8bdf2d8da435aba7fcde8e368b4d9",
            "5fa3dd57be684025806c95921a551037",
            "7261b4f95a91428f901d284efc03a2ce",
            "720b8aba5a1745428a67f994a5fd6051",
            "a87cec2cb96d495ab0fc7ab68e12bc78",
            "7a849ebf6c754ca89fa554d756e69470",
            "14443025f49d43ceba6cf2d0d5b71bd2"
          ]
        }
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:lightning_fabric.utilities.seed:Global seed set to 1\n",
            "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/utilities/parsing.py:197: UserWarning: Attribute 'encoder_q' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['encoder_q'])`.\n",
            "  rank_zero_warn(\n",
            "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/utilities/parsing.py:197: UserWarning: Attribute 'encoder_k' is an instance of `nn.Module` and is already saved during checkpointing. It is recommended to ignore them using `self.save_hyperparameters(ignore=['encoder_k'])`.\n",
            "  rank_zero_warn(\n",
            "<ipython-input-16-19f5bcc49657>:70: FutureWarning: weekofyear and week have been deprecated, please use DatetimeIndex.isocalendar().week instead, which returns a Series. To exactly reproduce the behavior of week and weekofyear and return an Index, you may call pd.Int64Index(idx.isocalendar().week)\n",
            "  dt.weekofyear.to_numpy(),\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "data encoding:   0%|          | 0/1 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "fa20273c28f14a6fb408d382d4ac98e6"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "sequence encoding:   0%|          | 0/69680 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "43eb6a5cb7f44ee89a72abf0cd5c9723"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:APP:encoding encoding_1696689714.4892964.pkl saved\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "forecasting evaluation [24, 48, 96, 288, 672]:   0%|          | 0/5 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "17e0c718123245e0894cb62237f054ef"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:APP:best alpha: 1\n",
            "INFO:APP:{'norm': {'MSE': 0.12564270224803148, 'MAE': 0.2548746581672258}, 'raw': {'MSE': 8.215635876279993, 'MAE': 1.971135493806163}}\n",
            "INFO:APP:best alpha: 2\n",
            "INFO:APP:{'norm': {'MSE': 0.1866364517380724, 'MAE': 0.3201052548710548}, 'raw': {'MSE': 12.02393211370452, 'MAE': 2.470399551050955}}\n",
            "INFO:APP:best alpha: 5\n",
            "INFO:APP:{'norm': {'MSE': 0.29406416662523127, 'MAE': 0.4055741548901386}, 'raw': {'MSE': 17.730418496845736, 'MAE': 3.071767129267302}}\n",
            "INFO:APP:best alpha: 10\n",
            "INFO:APP:{'norm': {'MSE': 0.7549016602700783, 'MAE': 0.6663402889854041}, 'raw': {'MSE': 47.958239695058914, 'MAE': 5.105067619041073}}\n",
            "INFO:APP:best alpha: 200\n",
            "INFO:APP:{'norm': {'MSE': 1.7324578238936703, 'MAE': 1.0441874147484846}, 'raw': {'MSE': 64.80876565788205, 'MAE': 6.781790260528435}}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "mse = []\n",
        "mae = []\n",
        "for i, pred_len in enumerate(pred_lens):\n",
        "    mse.append(0)\n",
        "    mae.append(0)\n",
        "    for res in results:\n",
        "        mse[i] += res[0][i] / len(results)\n",
        "        mae[i] += res[1][i] / len(results)\n",
        "\n",
        "log.info(mse)\n",
        "log.info(mae)\n",
        "\n",
        "data = [[x, y] for (x, y) in zip(pred_lens, mse)]\n",
        "table = wandb.Table(data=data, columns = [\"pred_lens\", \"mse\"])\n",
        "wandb.log(\n",
        "{f\"{SETTINGS_STRING} forecasting MSE plot\" : wandb.plot.scatter(table, \"pred_lens\", \"mse\",\n",
        "    title=f\"{SETTINGS_STRING} forecasting MSE norm plot\")})\n",
        "\n",
        "data = [[x, y] for (x, y) in zip(pred_lens, mae)]\n",
        "table = wandb.Table(data=data, columns = [\"pred_lens\", \"mae\"])\n",
        "wandb.log(\n",
        "{f\"{SETTINGS_STRING} forecasting MAE plot\" : wandb.plot.scatter(table, \"pred_lens\", \"mae\",\n",
        "    title=f\"{SETTINGS_STRING} forecasting MAE norm plot\")})\n"
      ],
      "metadata": {
        "id": "NXZsEiDVmyrB",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f513a64a-a080-4be0-a01f-9a15736dd0cb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "INFO:APP:[0.12564270224803148, 0.1866364517380724, 0.29406416662523127, 0.7549016602700783, 1.7324578238936703]\n",
            "INFO:APP:[0.2548746581672258, 0.3201052548710548, 0.4055741548901386, 0.6663402889854041, 1.0441874147484846]\n"
          ]
        }
      ]
    }
  ]
}