{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":56537,"databundleVersionId":8877088,"sourceType":"competition"},{"sourceId":8409068,"sourceType":"datasetVersion","datasetId":5004471},{"sourceId":8901446,"sourceType":"datasetVersion","datasetId":5351317}],"dockerImageVersionId":30698,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## <div  style=\"color:blue;  font-weight:bold; font-size:100%; text-align:center;padding:12.0px; background:#ffffff\"> Thank you for your attention! Please upvote this kernel if you like it. It motivates me to produce more quality content) </div>","metadata":{}},{"cell_type":"markdown","source":"<center>\n<img src=\"https://i.postimg.cc/pXG9GZRR/4-Untitled.png\" width=1100>\n</center>","metadata":{}},{"cell_type":"markdown","source":"\n\n\n# <div  style=\"color:white; letter-spacing: 2px;  font-weight:bold; font-size:120%; text-align:center;padding:12.0px; background:black; border:#ff0066 solid; \">    1. OVERVIEW</div>\n","metadata":{}},{"cell_type":"markdown","source":"# Task overview\nClimate models are essential to understanding Earth’s climate system. Because of the complexity of Earth’s climate, these models rely on parameterizations to approximate the effects of physical processes that occur at scales smaller than the size of their grid cells. These approximations are imperfect, however, and their imperfections are a leading source of uncertainty in expected warming, changing precipitation patterns, and the frequency and severity of extreme events. The Multi-scale Modeling Framework (MMF) approach, by contrast, more explicitly represents these subgrid processes, but at a cost too high to be used for operational climate prediction.\n\nThis competition accompanies an upcoming 2024 ICML Machine Learning for Earth System Modeling (ML4ESM) Workshop and is based on the ClimSim paper and dataset which won the Outstanding Datasets and Benchmarks Paper award at NeurIPS 2023. Winning submissions will be highlighted at the upcoming ML4ESM ICML workshop, and participants in this Kaggle competition are also encouraged to submit workshop papers.","metadata":{}},{"cell_type":"markdown","source":"# Goal\n\n\nThe goal is to develop a machine learning models that accurately emulate subgrid-scale atmospheric physics in an operational climate model — an important step in improving climate projections and reducing uncertainty surrounding future climate trends. \n\nML models have to  emulate subgrid atmospheric processes–such as storms, clouds, turbulence, rainfall, and radiation–within E3SM-MMF, a multi-scale climate model backed by the U.S. Department of Energy. Because ML emulators are significantly cheaper to inference than MMF, progress on this front can help scientists realize a future in which high-resolution and physically credible long-term climate projections are broadly accessible, bringing greater clarity to the hazards associated with climate change and empowering policymakers with the knowledge necessary to mitigate them.\n\n","metadata":{}},{"cell_type":"markdown","source":"# Data overview","metadata":{}},{"cell_type":"markdown","source":"We have the following information about the data:\n- There are **input** and **output** data\n- Variables are **arrays** and **scalars**\n- Variables start with: **state**, **ptend**, **pbuf**, **cam_in** and **cam_out**\n\n\n\n","metadata":{}},{"cell_type":"markdown","source":"## In details: all variables in table format","metadata":{}},{"cell_type":"markdown","source":"\n\n| Input variables (State)         | Name       | Output variables (Tendencies/Target)         | Name          |\n|-------------------------------------------|-----------|-----------------------------------------|-----------------|\n|                                ||  <span style=\"color:blue;\"> Arrays  </span> ||   |\n| Air Temperature                 | state_t      | Heating Tendency                                   | ptend_t         |\n| Specific Humidity                 | state_q0001 | Moistening Tendency                                | ptend_q0001    |\n| Cloud Liquid Mixing Ratio             | state_q0002 | Cloud Liquid Mixing Ratio Change                   | ptend_q0002    |\n| Cloud Ice Mixing Ratio                        | state_q0003 | Cloud Ice Mixing Ratio Change                      | ptend_q0003    |\n| Zonal Wind Speed                              | state_u      | Zonal Wind Acceleration                            | ptend_u         |\n| Meridional Wind Speed                         | state_v      | Meridional Wind Acceleration                       | ptend_v         |\n| Ozone Volume Mixing Ratio                     | pbuf_ozone   |                                                  |                 |\n| Methane Volume Mixing Ratio                   | pbuf_CH4     |                                                  |                 |\n| Nitrous Oxide Volume Mixing Ratio             | pbuf_N2O     |                                                  |                 |\n|||                                              <span style=\"color:blue;\"> Scalars    </span>    |||\n| Surface Pressure                              | state_PS     | Net Shortwave Flux at Surface                      | cam_out_NetTSw  |\n| Solar Insolation                              | pbuf_SolIn  | Downward Longwave Flux at Surface                  | cam_out_FlWDS  |\n| Surface Latent Heat Flux                      | pbuf_LHFlx  | Snow Rate (Liquid Water Equivalent)                | cam_out_PRECSC |\n| Surface Sensible Heat Flux                    | pbuf_SHFlx  | Rain Rate                                          | cam_out_PRECC  |\n| Zonal Surface Stress                          | pbuf_TauX    | Downward Visible Direct Solar Flux to Surface      | cam_out_SolS   |\n| Meridional Surface Stress                     | pbuf_TauY    | Downward Near-Infrared Direct Solar Flux to Surface| cam_out_SolL   |\n| Cosine of Solar Zenith Angle                  | pbuf_CosZRS  | Downward Diffuse Solar Flux to Surface             | cam_out_SolSD  |\n| Albedo for Diffuse Longwave Radiation         | cam_in_ALDif | Downward Diffuse Near-Infrared Solar Flux to Surface| cam_out_SolLD  |\n| Albedo for Direct Longwave Radiation          | cam_in_ALDir |                                                    |                 |\n| Albedo for Diffuse Shortwave Radiation        | cam_in_ASDif |                                                    |                 |\n| Albedo for Direct Shortwave Radiation         | cam_in_ASDir |                                                    |                 |\n| Upward Longwave Flux                          | cam_in_LWUp  |                                                    |                 |\n| Sea-Ice Areal Fraction                        | cam_in_IceFrac |                                                  |                 |\n| Land Areal Fraction                           | cam_in_LandFrac |                                                 |                 |\n| Ocean Areal Fraction                          | cam_in_OcnFrac |                                                  |                 |\n| Snow Depth over Land                          | cam_in_SnowHLand |                                                 |                 |\n","metadata":{}},{"cell_type":"markdown","source":"## Description of **input** and **output** variables:\n\n### **Input Arrays variables**:\n**State**:\n- **State_t (Air Temperature):** Represents air temperature at different vertical levels within the atmosphere (60 levels). Measured in Kelvin (K).\n- **State_q0001 (Specific Humidity):** Indicates specific humidity at different vertical levels (60 levels). Measured in kilograms of water vapor per kilogram of air (kg/kg).\n- **State_q0002 (Cloud Liquid Mixing Ratio):** Represents the mixing ratio of cloud liquid water content to total air mass at different vertical levels (60 levels). Measured in kg/kg.\n- **State_q0003 (Cloud Ice Mixing Ratio):** Similar to cloud liquid mixing ratio but for ice clouds. Measured in kg/kg.\n- **State_u (Zonal Wind Speed):** Represents the zonal (east-west) component of wind speed at different vertical levels (60 levels). Measured in meters per second (m/s).\n- **State_v (Meridional Wind Speed):** Represents the meridional (north-south) component of wind speed at different vertical levels (60 levels). Measured in meters per second (m/s).\n- **State_ps (Surface Pressure):** Indicates surface pressure at the Earth's surface. Measured in Pascals (Pa).\n\n**pbuf:**\n- **pbuf_ozone**: ozone volume mixing ratio\n- **pbuf_CH4**: methane volume mixing ratio\n- **pbuf_N2O**: nitros oxide volume mixing ratio\n\n\n\n### **Input Scalars variables**:\n**State**:\n- **state_PS**: surface pressure\n\n**pbuf:**\n- **pbuf_SolIn** : solar insolation\n- **pbuf_LHFlx** : surface latent heat flux\n- **pbuf_SHFlx** : surface sensible heat flux\n- **pbuf_TauX** : zonal surface stress\n- **pbuf_TauY** : meridional surface stress\n- **pbuf_CosZRS**: cosine of solar zenith angle\n\n**cam_in:**\n- **cam_in_ALDif**: albedo for diffuse longwave radiation\n- **cam_in_ALDir**: albedo for direct longwave radiation\n- **cam_in_ASDif**: albedo for diffuse shortwave radiation\n- **cam_in_ASDir**: albedo for direct shortwave radiation\n- **cam_in_LwUp**: upward longwave flux\n- **cam_in_IceFrac:** sea-ice areal fraction\n- **cam_in_LandFrac**: land areal fraction\n- **cam_in_OcnFrac**: ocean areal fraction\n- **cam_in_SnowHLand**: snow depth over land\n\n---\n\n### **Output (Target) Arrays variables**:\n**Ptend:**\n- **Ptend_t (Heating Tendency):** Represents the rate of change of air temperature with respect to time at different vertical levels (60 levels). Measured in Kelvin per second (K/s).\n- **Ptend_q0001 (Moistening Tendency):** Indicates the rate of change of specific humidity with respect to time at different vertical levels (60 levels). Measured in kg/kg per second (kg/kg/s).\n- **Ptend_q0002 and Ptend_q0003:** Represent the rate of change of cloud liquid and ice mixing ratios with respect to time, respectively, at different vertical levels (60 levels). Measured in kg/kg per second (kg/kg/s).\n- **ptend_q0003**: cloud ice mixing ratio change over time\n- **Ptend_u and Ptend_v:** zonal wind acceleration and meridional wind acceleration. Represent the acceleration of zonal and meridional wind speeds with respect to time, respectively, at different vertical levels (60 levels). Measured in meters per second squared (m/s^2).\n\n\n### **Output (Target) Scalars variables**:\n**cam_out:**\n- **cam_out_NetSw**: Net shortwave flux at surface\n- **cam_out_FLwDS**: Downward longwave flux at surface\n- **cam_out_PRECSC**: Snow rate (liquid water equivalent)\n- **cam_out_PRECC**: rain rate\n- **cam_out_SolS**: downward visible direct solar flux to surface\n- **cam_out_SolL**: downward near-infrared direct solar flux to surface\n- **cam_out_SolSD**: downward diffuse solar flux to surface\n- **cam_out_SolLD**: downward diffuse near-infrared solar flux to surface\n\n---","metadata":{}},{"cell_type":"markdown","source":"## Relationships between Input and Output variables\n\n### Air temperature and heating tendency:\nHigher air temperatures typically lead to increased heating tendencies, as warmer air has a greater capacity to absorb solar radiation and transfer heat to the surrounding environment. Conversely, cooler temperatures may result in decreased heating tendencies.\n\n### Specific humidity and moistening tendency:\nHigher specific humidity levels indicate greater moisture content in the atmosphere, which can lead to increased moistening tendencies. Conversely, lower specific humidity levels may result in decreased moistening tendencies.\n\n### Cloud liquid/ice mixing ratios and cloud formation/change tendencies:\nHigher mixing ratios of cloud liquid and ice indicate the presence of clouds in the atmosphere. Changes in these mixing ratios over time reflect processes such as cloud formation, growth, or dissipation. Positive tendencies indicate increasing cloud content, while negative tendencies suggest cloud dissipation.\n\n### Wind speeds and wind acceleration tendencies:\nChanges in wind speeds over time can result from various atmospheric dynamics, including pressure gradients, temperature gradients, and surface friction. Acceleration tendencies represent changes in wind speeds with respect to time, reflecting the influence of atmospheric forces on wind motion.\n\n### Surface pressure and atmospheric dynamics:\nSurface pressure variations influence atmospheric circulation patterns and weather systems. Changes in surface pressure may affect wind patterns, temperature distributions, and precipitation regimes.\n\n### Surface fluxes (solar insolation, latent heat flux, sensible heat flux) and surface energy budget:\nSurface fluxes play a crucial role in the exchange of energy between the Earth's surface and the atmosphere. Changes in solar insolation, latent heat flux (from evaporation), and sensible heat flux (from temperature gradients) influence the surface energy budget, which in turn impacts atmospheric dynamics and climate.\n\n### Atmospheric gases mixing ratios (ozone, methane, nitrous oxide) and radiative forcing:\nChanges in the concentrations of greenhouse gases such as ozone, methane, and nitrous oxide can affect the Earth's radiative balance and contribute to climate change. Alterations in these mixing ratios may lead to changes in radiative forcing and atmospheric temperatures over time.\n","metadata":{}},{"cell_type":"markdown","source":"\n# <div  style=\"color:white; letter-spacing: 2px;  font-weight:bold; font-size:120%; text-align:center;padding:12.0px; background:black; border:#ff0066 solid; \">   2. DATA LOADING</div>","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport polars as pl # faster alternative to pandas\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nimport random, sys, gc, warnings, math, os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n\nfrom sklearn.model_selection import train_test_split\n\nimport jax\nimport tensorflow as tf\nimport keras\nfrom keras.regularizers import l2\nfrom keras.layers import Conv1D, GroupNormalization, Activation\n\nfrom sklearn import metrics\nfrom sklearn.metrics import mean_squared_error\n\nfrom tqdm.notebook import tqdm\n\nprint(tf.__version__)\nprint(jax.__version__)\n\nnp.random.seed(1)\nrandom.seed(1)\npd.set_option('display.max_columns', None) # we want to display all columns in this notebook","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:24.38179Z","iopub.execute_input":"2024-07-15T13:28:24.382202Z","iopub.status.idle":"2024-07-15T13:28:24.391231Z","shell.execute_reply.started":"2024-07-15T13:28:24.38217Z","shell.execute_reply":"2024-07-15T13:28:24.39031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# file overview\n!ls -l '../input/leap-atmospheric-physics-ai-climsim/'","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:24.392815Z","iopub.execute_input":"2024-07-15T13:28:24.39311Z","iopub.status.idle":"2024-07-15T13:28:25.620802Z","shell.execute_reply.started":"2024-07-15T13:28:24.393079Z","shell.execute_reply":"2024-07-15T13:28:25.619627Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Datasets are **huge**, let's for the first step and simplicity start with a small subset. We select only **\"N_columns_X\"** in X, **\"N_column_y\"** in y and **\"N_rows\"** in rows","metadata":{}},{"cell_type":"code","source":"%%time\n#polars to pandas dataframe\n\nN_rows=10000\nfolder = 'leap-atmospheric-physics-ai-climsim/'\n\n\ndf_train = pl.read_csv('../input/'+folder+'/train.csv', n_rows=N_rows).to_pandas()\ndf_test  = pl.read_csv('../input/'+folder+'/test.csv', n_rows=N_rows).to_pandas()\ndf_sub   = pl.read_csv('../input/'+folder+'/sample_submission.csv', n_rows=N_rows).to_pandas()\n\n","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:25.622502Z","iopub.execute_input":"2024-07-15T13:28:25.622826Z","iopub.status.idle":"2024-07-15T13:28:26.457718Z","shell.execute_reply.started":"2024-07-15T13:28:25.622796Z","shell.execute_reply":"2024-07-15T13:28:26.456398Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# <div  style=\"color:white; letter-spacing: 2px;  font-weight:bold; font-size:120%; text-align:center;padding:12.0px; background:black; border:#ff0066 solid; \">   3. EDA</div>","metadata":{}},{"cell_type":"markdown","source":"# Dataset - first glance","metadata":{}},{"cell_type":"code","source":"df_train.head(5).round(1)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:26.460616Z","iopub.execute_input":"2024-07-15T13:28:26.461315Z","iopub.status.idle":"2024-07-15T13:28:27.349224Z","shell.execute_reply.started":"2024-07-15T13:28:26.461269Z","shell.execute_reply":"2024-07-15T13:28:27.348278Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train.tail(5).round(1)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:27.35053Z","iopub.execute_input":"2024-07-15T13:28:27.350878Z","iopub.status.idle":"2024-07-15T13:28:28.228163Z","shell.execute_reply.started":"2024-07-15T13:28:27.350845Z","shell.execute_reply":"2024-07-15T13:28:28.227173Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train.shape","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:28.229338Z","iopub.execute_input":"2024-07-15T13:28:28.229636Z","iopub.status.idle":"2024-07-15T13:28:28.236009Z","shell.execute_reply.started":"2024-07-15T13:28:28.22961Z","shell.execute_reply":"2024-07-15T13:28:28.234945Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# preview - test\ndf_test.head(5).round(2)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:28.23728Z","iopub.execute_input":"2024-07-15T13:28:28.237604Z","iopub.status.idle":"2024-07-15T13:28:28.765947Z","shell.execute_reply.started":"2024-07-15T13:28:28.237579Z","shell.execute_reply":"2024-07-15T13:28:28.764965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_train.shape","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:28.767637Z","iopub.execute_input":"2024-07-15T13:28:28.768042Z","iopub.status.idle":"2024-07-15T13:28:28.774325Z","shell.execute_reply.started":"2024-07-15T13:28:28.768009Z","shell.execute_reply":"2024-07-15T13:28:28.773349Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Extracting Targets and Features","metadata":{}},{"cell_type":"code","source":"# targets (extract from submission file)\ntargets = [x for x in df_sub.columns.tolist() if x not in ['sample_id']]\n\n# numerical features\nfeatures_num = [x for x in df_train.columns.tolist() if x not in ['sample_id']+targets]\n\n# categorical features\nfeatures_cat = []\n\n# all features combined\nfeatures = features_num + features_cat","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:28.775711Z","iopub.execute_input":"2024-07-15T13:28:28.776103Z","iopub.status.idle":"2024-07-15T13:28:28.79345Z","shell.execute_reply.started":"2024-07-15T13:28:28.776068Z","shell.execute_reply":"2024-07-15T13:28:28.792715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# ouput of dimensions\nprint(\"Number of numerical   features:\", len(features_num))\nprint(\"Number of categorical features:\", len(features_cat))\nprint(\"Number of targets:\", len(targets))\nprint()\nprint(\"Size of subset:\", N_rows)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:28.797224Z","iopub.execute_input":"2024-07-15T13:28:28.797605Z","iopub.status.idle":"2024-07-15T13:28:28.805592Z","shell.execute_reply.started":"2024-07-15T13:28:28.797581Z","shell.execute_reply":"2024-07-15T13:28:28.804592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Targets analysis","metadata":{}},{"cell_type":"code","source":"# basic statistics - targets\ndf_train[targets].describe().round(4)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:28.806637Z","iopub.execute_input":"2024-07-15T13:28:28.806886Z","iopub.status.idle":"2024-07-15T13:28:29.90121Z","shell.execute_reply.started":"2024-07-15T13:28:28.806864Z","shell.execute_reply":"2024-07-15T13:28:29.900197Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Target distribution for array variables ","metadata":{}},{"cell_type":"code","source":"# plot target distributions \n\nname_target = ['Heating tendency', 'Moistening tendency', 'Cloud liquid mixing ratio change', \n               'Cloud ice mixing ratio change', 'Zonal wind acceleration', 'Meridional wind acceleration']\n\n\nTitle_size=18\nbins=30\nc=3 # n of columns\nr=3 # n of rows\nN=c*r\ndefault_color_3='white'\nred_color='#ff0066'\ncolor2='blue'\nfor j in range(len(name_target)):\n    fig, axs = plt.subplots(r, c, figsize=(5*c, 5*r))\n    plt.suptitle(f\"{name_target[j]}\", fontsize=Title_size+4, fontweight='bold', y=0.93)\n    i = 0\n    for t in targets[0+j*60: 0+j*60+9]:\n        current_ax = axs.flat[i]\n        current_ax.hist(df_train[t],  color=default_color_3,  edgecolor = color2, bins = bins, density=True)\n        current_ax.set_title(str(t), fontsize=Title_size)\n        current_ax.grid()\n        i = i + 1\n    fig.patch.set_linewidth(5)\n    fig.patch.set_edgecolor('black')  # substitute 'k' for black\n    #fig.patch.set_facecolor('#dddddd') #gray\n    fig.patch.set_facecolor('#eeeeee') #gray\n    plt.show()\n    print()\n    print()\n    ","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:29.902449Z","iopub.execute_input":"2024-07-15T13:28:29.902738Z","iopub.status.idle":"2024-07-15T13:28:41.510631Z","shell.execute_reply.started":"2024-07-15T13:28:29.902713Z","shell.execute_reply":"2024-07-15T13:28:41.509678Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Target distribution for scalar variables","metadata":{}},{"cell_type":"code","source":"# plot target distributions in compact matrix form\n\ncol = ['NETSW \\n Net shortwave flux\\n at surface', 'FLWDS \\n Downward longwave flux \\nat surface', \n       'PRECSC\\n Snow rate \\n(liquid water equivalent)', 'PRECC\\n Rain rate', \n       'SOLS\\n Downward visible direct \\nsolar flux to surface', \n       'SOLL\\n Downward near-infrared direct\\n solar flux to surface', 'SOLSD\\n Downward diffuse \\nsolar flux to surface', \n       'SOLLD\\n Downward diffuse near-infrared\\n solar flux to surface']\n                                \n    \nbins=50\nc=3 # n of columns\nr=3 # n of rows\n#default_color_3 = 'white'\nfig, axs = plt.subplots(r, c, figsize=(7*c, 7*r))\nplt.suptitle(\"Cam_out\", fontsize=Title_size+4, fontweight='bold', y=0.95)\ni = 0\nfor t in targets[-8:]:\n    current_ax = axs.flat[i]\n    current_ax.grid()\n    current_ax.hist(df_train[t],  color=default_color_3,  edgecolor = 'blue', bins = bins, density=True )\n    current_ax.set_title(f'{col[i]}', fontsize=Title_size-4, y=1.00)\n    \n    i = i + 1\n    print()\nfig.patch.set_linewidth(5)\nfig.patch.set_edgecolor('black') \nfig.patch.set_facecolor('#eeeeee') #gray\nplt.show()\n","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:41.51206Z","iopub.execute_input":"2024-07-15T13:28:41.512451Z","iopub.status.idle":"2024-07-15T13:28:43.935846Z","shell.execute_reply.started":"2024-07-15T13:28:41.512418Z","shell.execute_reply":"2024-07-15T13:28:43.934835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Features analysis","metadata":{}},{"cell_type":"code","source":"# basic stats - train\ndf_train[features_num].describe().round(1)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:43.937108Z","iopub.execute_input":"2024-07-15T13:28:43.937396Z","iopub.status.idle":"2024-07-15T13:28:45.614108Z","shell.execute_reply.started":"2024-07-15T13:28:43.937371Z","shell.execute_reply":"2024-07-15T13:28:45.613123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# basic stats - test\ndf_test[features_num].describe().round(2)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:45.61526Z","iopub.execute_input":"2024-07-15T13:28:45.615532Z","iopub.status.idle":"2024-07-15T13:28:47.315856Z","shell.execute_reply.started":"2024-07-15T13:28:45.615509Z","shell.execute_reply":"2024-07-15T13:28:47.314892Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Features distribution for array variables","metadata":{}},{"cell_type":"markdown","source":"### Show first five sets in every block","metadata":{}},{"cell_type":"code","source":"# col = ['Air temperature  \\n  (state_t)', 'Specific humidity \\n (state_q0001)', 'Cloud liquid mixing ratio  \\n (state_q0002)', \n# 'Cloud ice mixing ratio \\n (state_q0003)', 'Zonal wind speed \\n (state_u)', 'Meridional wind speed \\n (state_v)',\n#        'Ozone volume mixing ratio \\n (pbuf_ozone)', 'Methane volume mixing ratio \\n (pbuf_CH4)', \n#        'Nitros oxide volume mixing ratio \\n (pbuf_N2O)']\n\ncol = ['Air temperature  \\n  (state_t)', 'Specific humidity \\n (state_q0001)', 'Cloud liquid mixing ratio  \\n (state_q0002)', \n'Cloud ice mixing ratio \\n (state_q0003)', 'Zonal wind speed \\n (state_u)', 'Meridional wind speed \\n (state_v)'       ]\n                          \n    \nr=5\nbins=30\ndefault_color_1 = 'white'\nfor j in range(len(col)):\n    i=1\n    fig = plt.figure(figsize=(2*6, r*5))\n    plt.suptitle(f\"{col[j]}\", fontsize=Title_size+4, fontweight='bold', y=0.95) \n    \n    for f in features_num[0+j*60: 0+j*60+r]:\n        ax1 = plt.subplot(r,2,i)\n        df_train[f].plot(kind='hist', bins=bins, color=default_color_1,  edgecolor = 'blue', density = True)\n        plt.title(f + ' - Train')\n        plt.grid()\n        i=i+1\n        \n        ax2 = plt.subplot(r,2,i, sharex=ax1)\n        df_test[f].plot(kind='hist', bins=bins, color=default_color_1,  edgecolor = red_color, density = True)\n        plt.title(f + ' - Test')\n        plt.grid()\n        i=i+1\n        \n    fig.patch.set_linewidth(5)\n    fig.patch.set_edgecolor('black') \n    fig.patch.set_facecolor('#eeeeee') #gray\n    plt.show()\n    print()\n    print()","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:28:47.317247Z","iopub.execute_input":"2024-07-15T13:28:47.317636Z","iopub.status.idle":"2024-07-15T13:29:04.998624Z","shell.execute_reply.started":"2024-07-15T13:28:47.317602Z","shell.execute_reply":"2024-07-15T13:29:04.99763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# <div  style=\"color:white; letter-spacing: 2px;  font-weight:bold; font-size:120%; text-align:center;padding:12.0px; background:black; border:#ff0066 solid; \">   4. MODELLING</div>","metadata":{}},{"cell_type":"markdown","source":"The overall architecture is copied from the code in this Kaggle notebook.\nhttps://www.kaggle.com/code/abiolatti/keras-baseline-seq2seq?scriptVersionId=180403717\n\n\nThe model architecture is based on the paper discussed in this Kaggle discussion and is detailed in this arXiv paper.\n\nhttps://www.kaggle.com/competitions/leap-atmospheric-physics-ai-climsim/discussion/516605\nhttps://arxiv.org/abs/2407.00124","metadata":{}},{"cell_type":"code","source":"isTrain =False #True #\nisTrain","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.000015Z","iopub.execute_input":"2024-07-15T13:29:05.000353Z","iopub.status.idle":"2024-07-15T13:29:05.006662Z","shell.execute_reply.started":"2024-07-15T13:29:05.000324Z","shell.execute_reply":"2024-07-15T13:29:05.005719Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def is_interactive():\n    return 'runtime' in get_ipython().config.IPKernelApp.connection_file\n\nprint('Interactive?', is_interactive())","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.007805Z","iopub.execute_input":"2024-07-15T13:29:05.008166Z","iopub.status.idle":"2024-07-15T13:29:05.020697Z","shell.execute_reply.started":"2024-07-15T13:29:05.00814Z","shell.execute_reply":"2024-07-15T13:29:05.019705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SEED = 2024\nkeras.utils.set_random_seed(SEED)\ntf.random.set_seed(SEED)\ntf.config.experimental.enable_op_determinism()","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.021748Z","iopub.execute_input":"2024-07-15T13:29:05.022022Z","iopub.status.idle":"2024-07-15T13:29:05.038171Z","shell.execute_reply.started":"2024-07-15T13:29:05.021999Z","shell.execute_reply":"2024-07-15T13:29:05.037257Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"DATA = \"/kaggle/input/leap-atmospheric-physics-ai-climsim\"\nDATA_TFREC = \"/kaggle/input/leap-train-tfrecords\"","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.039158Z","iopub.execute_input":"2024-07-15T13:29:05.040256Z","iopub.status.idle":"2024-07-15T13:29:05.049733Z","shell.execute_reply.started":"2024-07-15T13:29:05.040229Z","shell.execute_reply":"2024-07-15T13:29:05.048951Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample = pl.read_csv(os.path.join(DATA, \"sample_submission.csv\"), n_rows=1)\nTARGETS = sample.select(pl.exclude('sample_id')).columns\nprint(len(TARGETS))","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.050719Z","iopub.execute_input":"2024-07-15T13:29:05.051175Z","iopub.status.idle":"2024-07-15T13:29:05.193649Z","shell.execute_reply.started":"2024-07-15T13:29:05.051152Z","shell.execute_reply":"2024-07-15T13:29:05.192714Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def _parse_function(example_proto):\n    feature_description = {\n        'x': tf.io.FixedLenFeature([556], tf.float32),\n        'targets': tf.io.FixedLenFeature([368], tf.float32)\n    }\n    e = tf.io.parse_single_example(example_proto, feature_description)\n    return e['x'], e['targets']","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.19526Z","iopub.execute_input":"2024-07-15T13:29:05.19562Z","iopub.status.idle":"2024-07-15T13:29:05.201444Z","shell.execute_reply.started":"2024-07-15T13:29:05.195586Z","shell.execute_reply":"2024-07-15T13:29:05.20041Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if is_interactive():\n    train_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(1)]\n    valid_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(100, 101)]\nelse: \n    train_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(100)]\n    valid_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(100, 101)]","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.202658Z","iopub.execute_input":"2024-07-15T13:29:05.202943Z","iopub.status.idle":"2024-07-15T13:29:05.213395Z","shell.execute_reply.started":"2024-07-15T13:29:05.202913Z","shell.execute_reply":"2024-07-15T13:29:05.212589Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 512\n\ntrain_options = tf.data.Options()\ntrain_options.deterministic = True\n\nds_train = (\n    tf.data.Dataset.from_tensor_slices(train_files)\n    .with_options(train_options)\n    .shuffle(100)\n    .interleave(\n        lambda file: tf.data.TFRecordDataset(file).map(_parse_function, num_parallel_calls=tf.data.AUTOTUNE),\n        num_parallel_calls=tf.data.AUTOTUNE,\n        cycle_length=10,\n        block_length=100,\n        #block_length=1000,\n        deterministic=True\n    )\n    .shuffle(4 * BATCH_SIZE)\n    .batch(BATCH_SIZE)\n    .prefetch(tf.data.AUTOTUNE)\n)\n\nds_valid = (\n    tf.data.TFRecordDataset(valid_files)\n    .map(_parse_function)\n    .batch(BATCH_SIZE)\n    .prefetch(tf.data.AUTOTUNE)\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.214544Z","iopub.execute_input":"2024-07-15T13:29:05.214814Z","iopub.status.idle":"2024-07-15T13:29:05.291348Z","shell.execute_reply.started":"2024-07-15T13:29:05.214791Z","shell.execute_reply":"2024-07-15T13:29:05.290546Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"norm_x = keras.layers.Normalization()\nnorm_x.adapt(ds_train.map(lambda x, y: x).take(20 if is_interactive() else 10000))\n\nplt.scatter(\n    norm_x.mean.squeeze(),\n    norm_x.variance.squeeze() ** 0.5,\n    marker=\".\",\n    alpha=0.5\n)\nplt.xscale('log')\nplt.yscale('log')","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:05.292424Z","iopub.execute_input":"2024-07-15T13:29:05.292693Z","iopub.status.idle":"2024-07-15T13:29:06.159344Z","shell.execute_reply.started":"2024-07-15T13:29:05.292669Z","shell.execute_reply":"2024-07-15T13:29:06.158428Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"norm_y = keras.layers.Normalization()\nnorm_y.adapt(ds_train.map(lambda x, y: y).take(20 if is_interactive() else 10000))\n\nmean_y = norm_y.mean\nstdd_y = keras.ops.maximum(1e-10, norm_y.variance ** 0.5)\n\nplt.scatter(\n    mean_y.squeeze(),\n    stdd_y.squeeze(),\n    marker=\".\",\n    alpha=0.5\n)\nplt.xscale('log')\nplt.yscale('log')\n#plt.grid()","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:06.160747Z","iopub.execute_input":"2024-07-15T13:29:06.161202Z","iopub.status.idle":"2024-07-15T13:29:07.020041Z","shell.execute_reply.started":"2024-07-15T13:29:06.161165Z","shell.execute_reply":"2024-07-15T13:29:07.019071Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"min_y = np.min(np.stack([np.min(yb, 0) for _, yb in ds_train.take(20 if is_interactive() else 10000)], 0), 0, keepdims=True)\nmax_y = np.max(np.stack([np.max(yb, 0) for _, yb in ds_train.take(20 if is_interactive() else 10000)], 0), 0, keepdims=True)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:07.021111Z","iopub.execute_input":"2024-07-15T13:29:07.021377Z","iopub.status.idle":"2024-07-15T13:29:07.810256Z","shell.execute_reply.started":"2024-07-15T13:29:07.021353Z","shell.execute_reply":"2024-07-15T13:29:07.809252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if is_interactive(): \n    epochs = 1 # \nelse:  \n    epochs = 30 #11 # 25  # 15  # 12 \n    \nlearning_rate = 1e-3\nearly_patience = 5\n# epochs_warmup = 1\n# epochs_ending = 2\n# steps_per_epoch = int(np.ceil(len(train_files) * 100_000 / BATCH_SIZE))\n\n# lr_scheduler = keras.optimizers.schedules.CosineDecay(\n#     1e-4, \n#     (epochs - epochs_warmup - epochs_ending) * steps_per_epoch, \n#     warmup_target=learning_rate,\n#     warmup_steps=steps_per_epoch * epochs_warmup,\n#     alpha=0.1\n# )\n\n# plt.plot([lr_scheduler(it) for it in range(0, epochs * steps_per_epoch, steps_per_epoch)]);","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:07.815834Z","iopub.execute_input":"2024-07-15T13:29:07.816318Z","iopub.status.idle":"2024-07-15T13:29:07.821334Z","shell.execute_reply.started":"2024-07-15T13:29:07.816289Z","shell.execute_reply":"2024-07-15T13:29:07.820363Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow.keras.layers import LayerNormalization, MultiHeadAttention, Dense, Dropout\n\n@keras.saving.register_keras_serializable()\nclass TransformerEncoderLayer(tf.keras.layers.Layer):\n    def __init__(self, head_size, num_heads, ff_dim, dropout=0.1, **kwargs):\n        super(TransformerEncoderLayer, self).__init__(**kwargs)\n        self.att = MultiHeadAttention(key_dim=head_size, num_heads=num_heads, dropout=dropout)\n        self.ffn = tf.keras.Sequential([\n            Dense(ff_dim, activation='gelu'),\n            #Dense(ff_dim, activation='relu'),\n            Dense(25)\n        ])\n        self.layernorm1 = LayerNormalization(epsilon=1e-6)\n        self.layernorm2 = LayerNormalization(epsilon=1e-6)\n        self.dropout1 = Dropout(dropout)\n        self.dropout2 = Dropout(dropout)\n\n    def call(self, inputs, training=False):\n        attn_output = self.att(inputs, inputs)\n        attn_output = self.dropout1(attn_output, training=training)\n        out1 = self.layernorm1(inputs + attn_output)\n\n        ffn_output = self.ffn(out1)\n        ffn_output = self.dropout2(ffn_output, training=training)\n        return self.layernorm2(out1 + ffn_output) ","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:07.822659Z","iopub.execute_input":"2024-07-15T13:29:07.823021Z","iopub.status.idle":"2024-07-15T13:29:07.833332Z","shell.execute_reply.started":"2024-07-15T13:29:07.82299Z","shell.execute_reply":"2024-07-15T13:29:07.832589Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Transformer block\n@keras.saving.register_keras_serializable()\ndef transformer_block(inputs, head_size, num_heads, ff_dim, dropout=0.2):\n    attention_output = tf.keras.layers.MultiHeadAttention(\n        key_dim=head_size, num_heads=num_heads, dropout=dropout)(inputs, inputs)\n    attention_output = tf.keras.layers.Dropout(dropout)(attention_output)\n    attention_output = tf.keras.layers.LayerNormalization(epsilon=1e-6)(attention_output + inputs)\n\n    ff_output = tf.keras.layers.Dense(ff_dim, activation='gelu')(attention_output)\n    ff_output = tf.keras.layers.Dropout(dropout)(ff_output)\n    ff_output = tf.keras.layers.Dense(inputs.shape[-1])(ff_output)\n    ff_output = tf.keras.layers.LayerNormalization(epsilon=1e-6)(ff_output + attention_output)\n\n    return ff_output\n\n# ResBlock function\n@keras.saving.register_keras_serializable()\ndef res_block(x, filters, output_filters=None, groups=8):\n    if output_filters is None:\n        output_filters = filters\n    norm1 = GroupNormalization(groups=groups, axis=-1)(x)\n    silu1 = tf.keras.layers.Activation('swish')(norm1)\n    conv1 = tf.keras.layers.Conv1D(filters, kernel_size=3, padding='same')(silu1)\n\n    norm2 = GroupNormalization(groups=groups, axis=-1)(conv1)\n    silu2 = tf.keras.layers.Activation('swish')(norm2)\n    conv2 = tf.keras.layers.Conv1D(output_filters, kernel_size=3, padding='same')(silu2)\n\n    if x.shape[-1] != conv2.shape[-1]:\n        x = tf.keras.layers.Conv1D(output_filters, kernel_size=1, padding='same')(x)\n    output = tf.keras.layers.Add()([conv2, x])\n    return output\n\n# Downsample block\n@keras.saving.register_keras_serializable()\ndef repeat_block(x, filters, repeat):\n    for _ in range(repeat):\n        x = res_block(x, filters) \n    return x\n\n# Upsample block\n@keras.saving.register_keras_serializable()\ndef upsample_block(x, filters, repeat, concat_layer):\n    x = tf.keras.layers.Conv1DTranspose(filters, kernel_size=2, strides=2, padding='same')(x)\n    x = tf.keras.layers.Concatenate()([x, concat_layer])\n    for _ in range(repeat):\n        x = res_block(x, filters)\n    return x\n\n# Define a custom Lambda function to print and remove specific data slices\n# TPU does not support printing in this manner, for compatibility, direct value retrieval is used instead.\n# This function was previously used but later removed. Hence, there might be inconsistencies.\n@keras.saving.register_keras_serializable()\ndef slice_and_print(x):\n    # Print the first 2 time steps that are being removed\n    tf.print(\"Removed start:\", x[:, :2, :], summarize=-1)\n    # Print the last 2 time steps that are being removed\n    tf.print(\"Removed end:\", x[:, -2:, :], summarize=-1)\n    # Return the data after removing the padding from the start and end\n    return x[:, 2:-2, :] ","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:07.834461Z","iopub.execute_input":"2024-07-15T13:29:07.834734Z","iopub.status.idle":"2024-07-15T13:29:07.851002Z","shell.execute_reply.started":"2024-07-15T13:29:07.834712Z","shell.execute_reply":"2024-07-15T13:29:07.850257Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# Path to the model file\nmodel_filepath = '/kaggle/input/leap-u-net-models/model_epoch_09.keras'\n\n# Check if the model file exists\nif os.path.exists(model_filepath):\n    print(f\"Loading model from {model_filepath}\")\n    keras.utils.clear_session()\n    custom_objects = {\n        'R2Score': tf.keras.metrics.R2Score(class_aggregation=\"variance_weighted_average\"),\n    }\n    model = keras.models.load_model(model_filepath, custom_objects)\nelse:\n    print(\"Saved model not found, training a new model\")\n    # Clear the current Keras session\n    keras.utils.clear_session() \n\n    def x_to_seq(x):\n        x_seq0 = keras.ops.transpose(keras.ops.reshape(x[:, 0:60 * 6], (-1, 6, 60)), (0, 2, 1))\n        x_seq1 = keras.ops.transpose(keras.ops.reshape(x[:, 60 * 6 + 16:60 * 9 + 16], (-1, 3, 60)), (0, 2, 1))\n        x_flat = keras.ops.reshape(x[:, 60 * 6:60 * 6 + 16], (-1, 1, 16))\n        x_flat = keras.ops.repeat(x_flat, 60, axis=1)\n        return keras.ops.concatenate([x_seq0, x_seq1, x_flat], axis=-1) \n  \n   \n    # Build 1D U-Net model\n    def create_unet(input_shape):\n        inputs = keras.layers.Input(shape=input_shape)\n\n        # Encoder\n        encoder_1 = repeat_block(inputs, 128, 2)  # 64 x 128 x 2\n        encoder_1_down = keras.layers.MaxPooling1D(pool_size=2, strides=2)(encoder_1)  # Downsample # 32 x 128\n\n        encoder_2 = repeat_block(encoder_1_down, 256, 2)  # 32 x 256 x 2\n        encoder_2_down = keras.layers.MaxPooling1D(pool_size=2, strides=2)(encoder_2)  # Downsample # 16 x 256\n\n        encoder_3 = repeat_block(encoder_2_down, 256, 2)  # 16 x 256 x 2\n        encoder_3_down = keras.layers.MaxPooling1D(pool_size=2, strides=2)(encoder_3)  # Downsample # 8 x 256\n\n        encoder_4 = repeat_block(encoder_3_down, 256, 2)  # 8 x 256 x 2\n\n        # Bottleneck (Transformer)\n        bottleneck = transformer_block(encoder_4, head_size=4, num_heads=64, ff_dim=512)\n        \n        decoder_1 = keras.layers.Concatenate()([bottleneck, encoder_4])\n        decoder_1_block = repeat_block(decoder_1, 256, 3)  # 8 x 256 x 3\n        decoder_1_upsample = keras.layers.Conv1DTranspose(256, kernel_size=2, strides=2, padding='same')(decoder_1_block)  # Upsample # 16 x 256\n\n        decoder_2 = keras.layers.Concatenate()([decoder_1_upsample, encoder_3])\n        decoder_2_block = repeat_block(decoder_2, 256, 3)  # 16 x 256 x 3\n        decoder_2_upsample = keras.layers.Conv1DTranspose(256, kernel_size=2, strides=2, padding='same')(decoder_2_block)  # Upsample # 32 x 256\n\n        decoder_3 = keras.layers.Concatenate()([decoder_2_upsample, encoder_2])\n        decoder_3_block = repeat_block(decoder_3, 256, 3)  # 32 x 256 x 3\n        decoder_3_upsample = keras.layers.Conv1DTranspose(256, kernel_size=2, strides=2, padding='same')(decoder_3_block)  # Upsample # 64 x 256\n\n        decoder_4 = keras.layers.Concatenate()([decoder_3_upsample, encoder_1])\n        decoder_4_block = repeat_block(decoder_4, 256, 3)  # 64 x 256 x 3\n\n        model = keras.models.Model(inputs, decoder_4_block)\n        return model\n    \n    X_input = x = keras.layers.Input(ds_train.element_spec[0].shape[1:]) \n    x = keras.layers.Normalization(mean=norm_x.mean, variance=norm_x.variance)(x)\n    x = x_to_seq(x) \n  \n    # Zero-padding at the beginning and end of the sequence to extend the length from 60 to 64\n    x = keras.layers.ZeroPadding1D(padding=(2, 2))(x)\n    # \n    e = keras.layers.Conv1D(48, 1, padding='same')(x)   \n    e = create_unet(e.shape[1:])(e)      \n      \n    # Use a Lambda layer to remove the first and last 2 time steps \n    e = e[:, 2:-2, :]\n\n    p_all = keras.layers.Conv1D(14, 1, padding='same')(e)\n    print(p_all.shape)\n    \n    p_seq = p_all[:, :, :6]\n    p_seq = keras.ops.transpose(p_seq, (0, 2, 1))\n    p_seq = keras.layers.Flatten()(p_seq)\n    assert p_seq.shape[-1] == 360\n\n    p_flat = p_all[:, :, 6:6 + 8]\n    p_flat = keras.ops.mean(p_flat, axis=1)\n    assert p_flat.shape[-1] == 8\n\n    P = keras.ops.concatenate([p_seq, p_flat], axis=1)\n\n    # Build & compile the model\n    model = keras.Model(X_input, P)\n    model.compile(\n        loss='mse', \n        optimizer=keras.optimizers.Adam(0.0002),\n        metrics=[keras.metrics.MeanSquaredError(), \n                     keras.metrics.R2Score(class_aggregation=\"variance_weighted_average\"), \n        ]  # Updated R2Score\n    )\n    model.build(tuple(ds_train.element_spec[0].shape))\n    model.summary()","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:07.852065Z","iopub.execute_input":"2024-07-15T13:29:07.852333Z","iopub.status.idle":"2024-07-15T13:29:15.078505Z","shell.execute_reply.started":"2024-07-15T13:29:07.85231Z","shell.execute_reply":"2024-07-15T13:29:15.07753Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n\n# Normalize target values for training and validation datasets\nds_train_target_normalized = ds_train.map(lambda x, y: (x, (y - mean_y) / stdd_y))\nds_valid_target_normalized = ds_valid.map(lambda x, y: (x, (y - mean_y) / stdd_y))\n\nif isTrain:\n    from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping \n\n    # Model checkpoint callback\n    model_checkpoint = ModelCheckpoint(\n        filepath='model_epoch_{epoch:02d}.keras',  # Save model with epoch number\n        monitor='val_loss',  # Monitor validation loss\n        save_best_only=False,  # Save all models, not just the best one\n        save_weights_only=False,  # Save the entire model structure and weights\n        mode='min',  # 'min' indicates saving when the monitored value decreases\n        verbose=2  # Provide detailed logging\n    )\n\n    # Early stopping callback\n    early_stopping = EarlyStopping(\n        monitor='val_loss',  # Monitor validation loss\n        patience=early_patience,  # Number of epochs to wait before stopping\n        restore_best_weights=True  # Restore model weights from the epoch with the best value of the monitored quantity\n    )\n\n    # Train the model\n    history = model.fit(\n        ds_train_target_normalized,  # Training data\n        validation_data=ds_valid_target_normalized,  # Validation data\n        epochs=epochs,  # Number of epochs to train\n        verbose=1 if is_interactive() else 2,  # Verbose output\n        callbacks=[early_stopping, model_checkpoint]  # List of callbacks to apply during training\n    )\n","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:15.079975Z","iopub.execute_input":"2024-07-15T13:29:15.080253Z","iopub.status.idle":"2024-07-15T13:29:15.11438Z","shell.execute_reply.started":"2024-07-15T13:29:15.080228Z","shell.execute_reply":"2024-07-15T13:29:15.113576Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"if isTrain:\n    plt.plot(history.history['loss'], color='tab:blue')\n    plt.plot(history.history['val_loss'], color='tab:red')\n    plt.yscale('log');","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:15.115346Z","iopub.execute_input":"2024-07-15T13:29:15.115605Z","iopub.status.idle":"2024-07-15T13:29:15.12032Z","shell.execute_reply.started":"2024-07-15T13:29:15.115581Z","shell.execute_reply":"2024-07-15T13:29:15.119519Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\ny_valid = np.concatenate([yb for _, yb in ds_valid])\np_valid = model.predict(ds_valid, batch_size=BATCH_SIZE) * stdd_y + mean_y","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:15.121527Z","iopub.execute_input":"2024-07-15T13:29:15.121792Z","iopub.status.idle":"2024-07-15T13:29:29.957275Z","shell.execute_reply.started":"2024-07-15T13:29:15.121764Z","shell.execute_reply":"2024-07-15T13:29:29.955912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nscores_valid = np.array([metrics.r2_score(y_valid[:, i], p_valid[:, i]) for i in range(len(TARGETS))])\nplt.plot(scores_valid.clip(-1, 1))","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:29.958674Z","iopub.execute_input":"2024-07-15T13:29:29.959001Z","iopub.status.idle":"2024-07-15T13:29:31.238544Z","shell.execute_reply.started":"2024-07-15T13:29:29.958973Z","shell.execute_reply":"2024-07-15T13:29:31.237628Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nmask = scores_valid <= 1e-3\nf\"Number of under-performing targets: {sum(mask)}\"","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:31.23971Z","iopub.execute_input":"2024-07-15T13:29:31.240018Z","iopub.status.idle":"2024-07-15T13:29:31.247355Z","shell.execute_reply.started":"2024-07-15T13:29:31.239993Z","shell.execute_reply":"2024-07-15T13:29:31.246486Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\nf\"Clipped score: {scores_valid.clip(0, 1).mean()}\"","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:31.248492Z","iopub.execute_input":"2024-07-15T13:29:31.248753Z","iopub.status.idle":"2024-07-15T13:29:31.260535Z","shell.execute_reply.started":"2024-07-15T13:29:31.248731Z","shell.execute_reply":"2024-07-15T13:29:31.259704Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\ndel y_valid, p_valid\ngc.collect();","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:31.261604Z","iopub.execute_input":"2024-07-15T13:29:31.26189Z","iopub.status.idle":"2024-07-15T13:29:31.716629Z","shell.execute_reply.started":"2024-07-15T13:29:31.261866Z","shell.execute_reply":"2024-07-15T13:29:31.71572Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"\n# <div  style=\"color:white; letter-spacing: 2px;  font-weight:bold; font-size:120%; text-align:center;padding:12.0px; background:black; border:#ff0066 solid; \">   5. SUBMISSION</div>","metadata":{}},{"cell_type":"code","source":"sample = pl.read_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv\")","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:31.719697Z","iopub.execute_input":"2024-07-15T13:29:31.720008Z","iopub.status.idle":"2024-07-15T13:29:34.522781Z","shell.execute_reply.started":"2024-07-15T13:29:31.719982Z","shell.execute_reply":"2024-07-15T13:29:34.52195Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_test = (\n    pl.scan_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/test.csv\")\n    .select(pl.exclude(\"sample_id\"))\n    .cast(pl.Float32)\n    .collect()\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:34.523944Z","iopub.execute_input":"2024-07-15T13:29:34.52424Z","iopub.status.idle":"2024-07-15T13:29:45.831438Z","shell.execute_reply.started":"2024-07-15T13:29:34.524216Z","shell.execute_reply":"2024-07-15T13:29:45.830359Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gc.collect();","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:45.832738Z","iopub.execute_input":"2024-07-15T13:29:45.833048Z","iopub.status.idle":"2024-07-15T13:29:46.186096Z","shell.execute_reply.started":"2024-07-15T13:29:45.833022Z","shell.execute_reply":"2024-07-15T13:29:46.184923Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\np_test = model.predict(df_test.to_numpy(), batch_size=4 * BATCH_SIZE) * stdd_y + mean_y\np_test = np.array(p_test)\np_test[:, mask] = mean_y[:, mask]","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:29:46.188032Z","iopub.execute_input":"2024-07-15T13:29:46.188486Z","iopub.status.idle":"2024-07-15T13:30:43.145689Z","shell.execute_reply.started":"2024-07-15T13:29:46.188446Z","shell.execute_reply":"2024-07-15T13:30:43.144654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# correction of ptend_q0002 targets (from 12 to 29)\ndf_p_test = pd.DataFrame(p_test, columns=TARGETS)\n\nfor idx in range(12, 30):\n    df_p_test[f\"ptend_q0002_{idx}\"] = -df_test[f\"state_q0002_{idx}\"].to_numpy() / 1200\n    \np_test = df_p_test.values","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:30:43.14707Z","iopub.execute_input":"2024-07-15T13:30:43.1474Z","iopub.status.idle":"2024-07-15T13:30:44.956375Z","shell.execute_reply.started":"2024-07-15T13:30:43.147372Z","shell.execute_reply":"2024-07-15T13:30:44.955557Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"submission = sample.to_pandas()\nsubmission[TARGETS] = submission[TARGETS] * p_test\npl.from_pandas(submission[[\"sample_id\"] + TARGETS]).write_csv(\"submission.csv\")","metadata":{"execution":{"iopub.status.busy":"2024-07-15T13:30:44.957502Z","iopub.execute_input":"2024-07-15T13:30:44.957856Z","iopub.status.idle":"2024-07-15T13:31:16.228698Z","shell.execute_reply.started":"2024-07-15T13:30:44.957823Z","shell.execute_reply":"2024-07-15T13:31:16.227812Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## <div  style=\"color:blue;  font-weight:bold; font-size:100%; text-align:center;padding:12.0px; background:#ffffff\"> Thank you for your attention! Please upvote this kernel if you like it. It motivates me to produce more quality content) </div>","metadata":{}}]}