{"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":"none","dataSources":[{"sourceId":56537,"databundleVersionId":8877088,"sourceType":"competition"}],"dockerImageVersionId":30746,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"\n# Description\n\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\nYour task is to develop ML models that 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.","metadata":{}},{"cell_type":"markdown","source":"# Evaluation\nWe will be using a custom R-squared metric for evaluation, but on a weighted solution. Prior to submitting your prediction, please multiply your prediction data element-wise by the data found in `sample_submission.csv`, which serves dual purpose as both a \"sample submission\" and a \"weighting file\".\n\nAs a reminder, R-squared is defined as:\n\n\\[ R^2 = 1 - \\frac{SS_{res}}{SS_{tot}} \\]\nwhere \\( SS_{res} \\) is the sum of squares of the residual errors and \\( SS_{tot} \\) is the total sum of squares.\n\nThe highest R-squared possible is 1. However, negative R-squared values are possible if the model performs worse than predicting the average value across all variables.","metadata":{}},{"cell_type":"markdown","source":"# My Thoughts\nOur task is to train a machine learning model to replace the MMF grid model in predicting atmospheric physics information.\nObservations of the given data are as follows:\n1. The fundamental task is multi-to-multi prediction, where multi-dimensional indicators and single-dimensional indicators are mixed.\n2. Multi-dimensional data is divided into 60 dimensions based on altitude.\n3. Some dependent variables are the rates of change of their corresponding independent variables at time t. That is, \\( \\Delta y = \\frac{\\Delta x}{t} \\).\n\nRelated research project:\nhttps://github.com/leap-stc/ClimSim/blob/main/demo_notebooks/water_conservation.ipynb \n\nTherefore, regarding the type of model, the following considerations are made:\n1. Reshape the dimensions to three dimensions to capture sequence information.\n2. Perform first-order differencing on the features to capture gradient information.\n3. Choose an appropriate time series model.","metadata":{}},{"cell_type":"markdown","source":"# Summary: Key Technical Implementation of Atmospheric Physics Surrogate Models\n\n#### 1. **The Necessity and Implementation of Reshaping Dimensions**\n**Core Significance**:\n- **Heterogeneous Feature Integration**: The original data includes 60-dimensional height-layer features (temperature, humidity, wind speed, etc.) and global single-dimensional features (heat flux, radiation). Reshaping creates a unified 3D tensor structure.\n- **Preserving the Physical Meaning of the Vertical Dimension**: The 60 height layers are treated as time steps, fully retaining the physical continuity of the atmospheric vertical profile.\n- **Adapting to the Input Requirements of Time Series Models**: The data is transformed from `(batch, feature)` to `(batch, timesteps, channels)` to meet the input requirements of LSTM/Transformer models.\n\n**Code Implementation Highlights**:\n```python\n# Direct mapping of multi-dimensional features\nreshaped_data[:, :, i] = x[:, indices]\n\n# Broadcasting single-dimensional features\nreshaped_data[:, :, i] = np.tile(x[:, index], (60, 1)).T\n```\n- Creating a `(batch, 60, 33)` tensor (33 = 25 60-dimensional features + 8 single-dimensional features)\n- Using `np.tile` to broadcast global features across the vertical dimension\n\n#### 2. **The Physical Significance of First-Order Differencing**\n**Atmospheric Dynamics Foundation**:\n- Vertical gradients drive atmospheric motion (temperature gradient → thermodynamic circulation, humidity gradient → cloud formation).\n- Five types of differencing calculations in the code:\n  ```python\n  # Temperature gradient\n  (pl.col(temp_cols[i]) - pl.col(temp_cols[i - 1])).alias(f'diff_state_t_{i}')\n  # Wind shear\n  (pl.col(wind_speed_diff_cols[i]) - pl.col(...)).alias(f'diff_wind_speed_{i}')\n  ```\n**Core Value**:\n1. Revealing atmospheric instability energy (Convective Available Potential Energy, CAPE)\n2. Quantifying turbulence exchange intensity (basis for calculating the Richardson number)\n3. A key physical process that the surrogate model must capture\n\n#### 3. **Scientific Basis for the Choice of Time Series Models**\n**Architecture Evolution Analysis**:\n```python\n# Initial Transformer solution (potentially over-parameterized)\nself.enc_embedding = DataEmbedding_inverted(configs.in_channel, configs.d_model)\nself.encoder = Encoder([EncoderLayer(...) for _ in range(6)])\n\n# Final CNN-LSTM solution (optimal choice)\nself.conv1 = nn.Conv1d(in_channel, cnn_out_channels, kernel_size=3)\nself.lstm = nn.LSTM(cnn_out_channels, d_model, bidirectional=True)\n```\n**Reasons for Choosing CNN-LSTM**:\n1. **Vertical Dimension Characteristics**: Atmospheric variables have strong local correlations in the vertical direction (advantage of CNN)\n   - 3x1 convolutional kernel captures physical interactions between adjacent height layers\n2. **Bidirectional Dependency Modeling**: Atmospheric processes are influenced by both upper and lower layers (bidirectional LSTM)\n   ```python\n   self.lstm = nn.LSTM(..., bidirectional=True)  # Considering both low-level to high-level and high-level to low-level influences\n   ```\n3. **Computational Efficiency**: Compared to the O(L²) complexity of Transformer, CNN-LSTM is more suitable for sequences of medium length (60 layers)\n\n#### 4. **Physical Insights in Feature Engineering**\n**Synthesizing Key Atmospheric Parameters**:\n```python\n# Scalar wind speed (foundation of dynamics)\n(pl.col(f\"state_u_{i}\")**2 + pl.col(f\"state_v_{i}\")**2).sqrt()\n\n# Surface energy balance (core of thermodynamics)\ntotal_thermal_flux = df['pbuf_LHFLX'] + df['pbuf_SHFLX']  # Sensible heat + latent heat\neffective_solar_radiation = df['pbuf_SOLIN'] * df['pbuf_COSZRS'] * (1 - df['cam_in_ASDIR'])\n```\n**Physical Significance**:\n- Scalar wind speed: Converting vector wind to scalar (basis for turbulent kinetic energy calculations)\n- Radiation calculation: Considering solar zenith angle (cosz) and surface albedo feedback\n\n#### 5. **Special Considerations in Normalization**\n**Adaptation to Atmospheric Data Characteristics**:\n```python\n# Robust normalization for sparse features (e.g., high cloud ice content)\nnp.maximum(X.std(axis=0), self.min_std)  # min_std=1e-8 to avoid division by zero\n\n# L2 normalization for target variables\nnp.maximum(np.sqrt((Y * Y).mean(axis=0)), self.min_std)  # Handling non-Gaussian distributions\n```\n\n### Key Problem-Solving Pathways\n1. **Preserving Vertical Dimension Physics**  \n   Maintaining the 60-layer structure during tensor reshaping → Retaining vertical processes such as gravity waves and convection\n\n2. **Modeling Cross-Scale Processes**  \n   CNN captures local turbulence (small scale) + LSTM captures stratification changes (medium scale)\n\n3. **Embedding Conservation Constraints**  \n   Explicitly introducing energy/water vapor conservation quantities in feature engineering (e.g., LHFLX + SHFLX)\n\n4. **Handling Sparse Data**  \n   Enhancing signals with differencing operations + Preventing normalization distortion with min_std\n\n### Model Improvement Suggestions\n1. **Physics-Constrained Loss Function**  \n   Adding mass/energy conservation penalty terms (refer to ClimSim's water conservation example)\n   ```python\n   loss += λ * (predicted_flux - observed_flux).abs().mean()\n   ```\n\n2. **Hierarchical Feature Extraction**  \n   ```python\n   # Different convolutional kernels for different height ranges\n   low_level = nn.Conv1d(0:30, kernel_size=5)  # Larger kernel for boundary layer\n   upper_level = nn.Conv1d(30:60, kernel_size=3)\n   ```\n\n3. **Multi-Task Decoding Design**  \n   Independent output heads for different types of variables (scalars, vectors, conserved quantities)\n\nThis implementation successfully integrates atmospheric physics prior knowledge into deep learning architectures, providing an efficient computational path for surrogate complex MMF models.","metadata":{}},{"cell_type":"markdown","source":"# import libs","metadata":{}},{"cell_type":"code","source":"import random, sys, gc, warnings, math\nimport numpy as np\nimport pandas as pd\nimport polars as pl\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport matplotlib.pyplot as plt\nimport warnings\nfrom torchmetrics.regression import R2Score\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torch.optim.lr_scheduler import CosineAnnealingLR\nfrom torch import nn, save\nfrom torch.optim import Adam, SGD\nimport torch.nn.functional as F\n\nimport logging\n\n\nwarnings.filterwarnings('ignore', category=FutureWarning)\n\nSCHEDULER_FACTOR = 10**(-1)  # The learning rate will be multiplied by 0.1 each time it is reduced.\n\nSCHEDULER_PATIENCE = 3\nPATIENCE = 6  # PATIENCE > SCHEDULER_PATIENCE\n\nMIN_STD = 1e-8\nBATCH_SIZE = 1024\nEPOCH = 10\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\nLOCAL = False\n\ndevice","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# load data","metadata":{}},{"cell_type":"code","source":"if LOCAL:\n    n_rows = 100000\nelse:\n    n_rows = 1000000\n    \ndf_train1 = pl.read_parquet(\"autodl-tmp/train_sampled_2m_part_0.parquet\", n_rows=n_rows)\ndf_train2 = pl.read_parquet(\"autodl-tmp/train_sampled_2m_part_1.parquet\", n_rows=n_rows)\ndf_train3 = pl.read_parquet(\"autodl-tmp/train_sampled_2m_part_2.parquet\", n_rows=n_rows)\n#df_train4 = pl.read_parquet(\"autodl-tmp/train_sampled_2m_part_3.parquet\", n_rows=n_rows)\n\ndf_train = pl.concat([df_train1, df_train2, df_train3])\n\ndel df_train1\ndel df_train2\ndel df_train3\n#del df_train4\ngc.collect()\n\n'''\ndf_train = pl.read_parquet(\"autodl-tmp/df_train.parquet\", n_rows=2000000)\ndf_train.head()\n'''","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"FEAT_COLS = df_train.columns[1:557]\nTARGET_COLS = df_train.columns[557:]\n\nx_train = df_train.select(FEAT_COLS)\ny_train = df_train.select(TARGET_COLS)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# feature engineering","metadata":{}},{"cell_type":"code","source":"def process_data(df):\n    # Temperature gradient, perform differencing by column\n    temp_cols = [f'state_t_{i}' for i in range(60)]\n    temp_diff_exprs = [pl.lit(0).alias('diff_state_t_0')] + [\n        (pl.col(temp_cols[i]) - pl.col(temp_cols[i - 1])).fill_null(0).alias(f'diff_state_t_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(temp_diff_exprs)\n\n    # Synthesize wind speed\n    wind_speed_exprs = [\n        (pl.col(f\"state_u_{i}\")**2 + pl.col(f\"state_v_{i}\")**2).sqrt().alias(f\"wind_speed_{i}\")\n        for i in range(60)\n    ]\n    df = df.with_columns(wind_speed_exprs)\n\n    # Humidity gradient, perform differencing by column\n    humidity_cols = [f'state_q0001_{i}' for i in range(60)]\n    humidity_diff_exprs = [pl.lit(0).alias('diff_state_q0001_0')] + [\n        (pl.col(humidity_cols[i]) - pl.col(humidity_cols[i - 1])).fill_null(0).alias(f'diff_state_q0001_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(humidity_diff_exprs)\n\n    # Liquid gradient, perform differencing by column\n    liquid_cols = [f'state_q0002_{i}' for i in range(60)]\n    liquid_diff_exprs = [pl.lit(0).alias('diff_state_q0002_0')] + [\n        (pl.col(liquid_cols[i]) - pl.col(liquid_cols[i - 1])).fill_null(0).alias(f'diff_state_q0002_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(liquid_diff_exprs)\n\n    # Ice gradient, perform differencing by column\n    ice_cols = [f'state_q0003_{i}' for i in range(60)]\n    ice_diff_exprs = [pl.lit(0).alias('diff_state_q0003_0')] + [\n        (pl.col(ice_cols[i]) - pl.col(ice_cols[i - 1])).fill_null(0).alias(f'diff_state_q0003_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(ice_diff_exprs)\n\n    # Wind shear/gradient\n    wind_speed_diff_cols = [f'wind_speed_{i}' for i in range(60)]\n    wind_speed_diff_exprs = [pl.lit(0).alias('diff_wind_speed_0')] + [\n        (pl.col(wind_speed_diff_cols[i]) - pl.col(wind_speed_diff_cols[i - 1])).fill_null(0).alias(f'diff_wind_speed_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(wind_speed_diff_exprs)\n\n    # Total thermal flux = Sensible heat + Latent heat\n    total_thermal_flux = (df['pbuf_LHFLX'] + df['pbuf_SHFLX']).alias(\"total_thermal_flux\")\n    df = df.with_columns(total_thermal_flux)\n\n    # Effective solar radiation = Solar insolation × cos(Solar zenith angle) × (1 - Surface albedo)\n    effective_solar_radiation = (\n        df['pbuf_SOLIN'] * df['pbuf_COSZRS'] * (1 - df['cam_in_ASDIR'])\n    ).alias(\"effective_solar_radiation\")\n    df = df.with_columns(effective_solar_radiation)\n\n    return df\n\n\nx_train = process_data(x_train)\nFEAT_COLS = list(x_train.columns)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"x_train = x_train.to_numpy()\ny_train = y_train.to_numpy()\ndel df_train\ngc.collect()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class StandardScaler:\n    def __init__(self, min_std=MIN_STD, test=False):\n        self.min_std = min_std\n        self.test = test\n        self.mean_x = None\n        self.std_x = None\n        self.mean_y = None\n        self.std_y = None\n\n    def fit(self, X, Y=None):\n        # If processing test data, ignore Y\n        if self.test:\n            Y = None\n\n        # Calculate the mean and standard deviation of X\n        self.mean_x = X.mean(axis=0)\n        self.std_x = np.maximum(X.std(axis=0), self.min_std)\n\n        # If not processing test data, calculate the mean and standard deviation of Y\n        if not self.test and Y is not None:\n            print('将处理y')\n            self.mean_y = Y.mean(axis=0)\n            self.std_y = np.maximum(np.sqrt((Y * Y).mean(axis=0)), self.min_std)\n\n    def transform(self, X, Y=None):\n        # If processing test data, ignore Y\n        if self.test:\n            Y = None\n\n        # Standardize X using the saved mean and standard deviation\n        X_norm = (X - self.mean_x.reshape(1, -1)) / self.std_x.reshape(1, -1)\n\n        # If not processing test data, standardize Y\n        if not self.test and Y is not None:\n            Y_norm = (Y - self.mean_y.reshape(1, -1)) / self.std_y.reshape(1, -1)\n            return X_norm, Y_norm\n\n        return X_norm\n\n    def fit_transform(self, X, Y=None):\n         # Call the fit method based on whether it is in test mode\n        self.fit(X, Y)\n        return self.transform(X, Y)\n    \n    def get_info(self):\n        # Return the values of mean_y and std_y\n        return self.mean_y, self.std_y, self.mean_x, self.std_x\n\nscaler = StandardScaler(min_std=1e-8, test=False)\nx_train, y_train = scaler.fit_transform(x_train, y_train)\nmy, sy, mx, sx = scaler.get_info()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Determine the indices for multi-dimensional and single-dimensional features\nmulti_feat_indices = []\nsingle_feat_indices = []\n\nfor i, col in enumerate(FEAT_COLS):\n    if '_' in col and col.split('_')[-1].isdigit():\n         # Part of a 60-dimensional feature\n        base_name = col.rsplit('_', 1)[0]\n        if col.endswith('_0'):   # Only add index range when encountering a new group\n            start_index = i\n            end_index = i + 60\n            multi_feat_indices.append((base_name, list(range(start_index, end_index))))\n    else:\n        # Single-dimensional feature\n        single_feat_indices.append((col, i))\n\n# Check indices\n# Display example of the first two groups\n#print(\"Multi-dimensional feature indices (sample):\", multi_feat_indices[:2])  \n\n# Display example of the first two single-dimensional features\n#print(\"Single-dimensional feature indices (sample):\", single_feat_indices[:2])  \n\n\ndef reshape_data(x, multi_feat_indices, single_feat_indices):\n    num_samples = x.shape[0]\n    # Create a new array with the second dimension being the time steps (60) and the third dimension being the number of features (33)\n    reshaped_data = np.zeros((num_samples, 60, 33), dtype=np.float32)\n    \n    i = 0\n    # Iterate through multi-dimensional feature groups\n    for _, indices in multi_feat_indices:\n        reshaped_data[:, :, i] = x[:, indices]\n        i += 1\n\n    # Iterate through single-dimensional features\n    for _, index in single_feat_indices:\n        reshaped_data[:, :, i] = np.tile(x[:, index], (60, 1)).T\n        i += 1\n\n    return reshaped_data\n\n\n# Reshape\nx_train = reshape_data(x_train, multi_feat_indices, single_feat_indices)\n\nmulti_target_indices = []\nsingle_target_indices = []\n\nfor i, col in enumerate(TARGET_COLS):\n    if '_' in col and col.split('_')[-1].isdigit():\n        # Part of a 60-dimensional feature\n        base_name = col.rsplit('_', 1)[0]\n        if col.endswith('_0'):  # Only add index range when encountering a new group\n            start_index = i\n            end_index = i + 60\n            multi_target_indices.append((base_name, list(range(start_index, end_index))))\n    else:\n        # Single-dimensional feature\n        single_target_indices.append((col, i))\n\ndef reshape_targets(y, multi_target_indices, single_target_indices):\n    num_samples = y.shape[0]\n    reshaped_y = np.zeros((num_samples, 60, 14), dtype=np.float32)  # New shape\n\n    i = 0\n    \n    # Process multi-dimensional features\n    for _, indices in multi_target_indices:\n        reshaped_y[:, :, i] = y[:, indices]\n        i += 1\n\n    # Process single-dimensional features\n    for _, index in single_target_indices:\n        reshaped_y[:, :, i] = np.tile(y[:, index], (60, 1)).T\n        i += 1\n\n    return reshaped_y\n\n# Apply the target reshaping function\ny_train = reshape_targets(y_train, multi_target_indices, single_target_indices)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"class TimeseriesDataset(Dataset):\n    def __init__(self, data, labels):\n        self.data = data\n        self.labels = labels\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        x = self.data[idx]  \n        y = self.labels[idx]\n        return torch.tensor(x, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)\n\n# Split the data into training and validation sets\nsplit_index = int(0.8 * len(x_train))\nx_train_set, x_valid_set = x_train[:split_index], x_train[split_index:]\ny_train_set, y_valid_set = y_train[:split_index], y_train[split_index:]\n\ntrain_dataset = TimeseriesDataset(x_train_set, y_train_set)\nvalid_dataset = TimeseriesDataset(x_valid_set, y_valid_set)\n\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\nvalid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\nfor data, target in train_loader:\n    # Should output (batch_size, 60, 25 + 6 + 2) and (batch_size, 60, 14)\n    print(data.shape, target.shape)  \n    break","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# class models","metadata":{}},{"cell_type":"markdown","source":"## model: iTransformer","metadata":{}},{"cell_type":"code","source":"class DataEmbedding_inverted(nn.Module):\n    def __init__(self, c_in, d_model, dropout=0.1):\n        super(DataEmbedding_inverted, self).__init__()\n        self.value_embedding = nn.Linear(c_in, d_model)  # Linear transformation to map input features to d_model dimensions\n        self.dropout = nn.Dropout(p=dropout)  # Dropout layer to prevent overfitting\n        \n    def forward(self, x, x_mark):\n        # x: [batch_size, 60, 33] - Input data with shape [batch_size, sequence_length, feature_dim]\n        # value_embedding expects input of shape [batch_size, 60, 33]\n        x = self.value_embedding(x)  # Transform input to [batch_size, 60, d_model]\n        x = self.dropout(x)  \n        return x  # Maintain the shape [batch_size, 60, d_model]\n\nclass FullAttention(nn.Module):\n    def __init__(self, mask_flag=True, scale=1.0, attention_dropout=0.1, output_attention=False):\n        super(FullAttention, self).__init__()\n        self.scale = scale  # Scaling factor for attention scores\n        self.mask_flag = mask_flag  # Flag to indicate if masking is required\n        self.output_attention = output_attention  # Flag to determine if attention weights should be returned\n        self.dropout = nn.Dropout(attention_dropout)\n\n    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n        B, L, H, E = queries.shape  # Batch size, sequence length, number of heads, embedding dimension\n        _, S, _, D = values.shape  # Sequence length of values, embedding dimension of values\n        \n        # Compute attention scores using einsum for efficient computation\n        scores = torch.einsum(\"blhe,bshe->bhls\", queries, keys) * self.scale\n\n        if self.mask_flag:\n            if attn_mask is None:\n                # Create a causal mask if not provided\n                attn_mask = TriangularCausalMask(B, L, device=queries.device)\n\n            # Apply masking to prevent information leakage\n            scores.masked_fill_(attn_mask.mask, -np.inf)\n\n        A = self.dropout(torch.softmax(scores, dim=-1))  # Compute attention weights and apply dropout\n        V = torch.einsum(\"bhls,bshd->blhd\", A, values)  # Compute weighted sum of values\n\n        if self.output_attention:\n            return V.contiguous(), A # Return both the output and attention weights\n        else:\n            return V.contiguous(), None\n\nclass AttentionLayer(nn.Module):\n    def __init__(self, attention, d_model, n_heads, d_keys=None, d_values=None):\n        super(AttentionLayer, self).__init__()\n\n        d_keys = d_keys or (d_model // n_heads)\n        d_values = d_values or (d_model // n_heads)\n\n        self.inner_attention = attention\n        self.query_projection = nn.Linear(d_model, d_keys * n_heads)\n        self.key_projection = nn.Linear(d_model, d_keys * n_heads)\n        self.value_projection = nn.Linear(d_model, d_values * n_heads)\n        self.out_projection = nn.Linear(d_values * n_heads, d_model)\n        self.n_heads = n_heads\n\n    def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):\n        B, L, _ = queries.shape\n        _, S, _ = keys.shape\n        H = self.n_heads\n\n        # Project queries, keys, and values to higher dimensions\n        queries = self.query_projection(queries).view(B, L, H, -1)\n        keys = self.key_projection(keys).view(B, S, H, -1)\n        values = self.value_projection(values).view(B, S, H, -1)\n\n        # Compute attention using the inner attention mechanism\n        out, attn = self.inner_attention(\n            queries,\n            keys,\n            values,\n            attn_mask,\n            tau=tau,\n            delta=delta\n        )\n        out = out.view(B, L, -1)\n\n        return self.out_projection(out), attn\n\nclass EncoderLayer(nn.Module):\n    def __init__(self, attention, d_model, d_ff=None, dropout=0.1, activation=\"relu\"):\n        super(EncoderLayer, self).__init__()\n        d_ff = d_ff or 4 * d_model\n        self.attention = attention\n        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)\n        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)\n        self.norm1 = nn.LayerNorm(d_model)  # Layer normalization for attention output\n        self.norm2 = nn.LayerNorm(d_model)  # Layer normalization for feed-forward output\n        self.dropout = nn.Dropout(dropout)\n        self.activation = F.relu if activation == \"relu\" else F.gelu\n\n    def forward(self, x, attn_mask=None, tau=None, delta=None):\n        # Apply attention mechanism\n        new_x, attn = self.attention(\n            x, x, x,\n            attn_mask=attn_mask,\n            tau=tau, delta=delta\n        )\n        x = x + self.dropout(new_x)  # Add dropout and residual connection\n\n        y = x = self.norm1(x)\n        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))  # First convolutional layer with activation\n        y = self.dropout(self.conv2(y).transpose(-1, 1))  # Second convolutional layer\n\n        return self.norm2(x + y), attn\n        \nclass Encoder(nn.Module):\n    def __init__(self, attn_layers, conv_layers=None, norm_layer=None):\n        super(Encoder, self).__init__()\n        # List of attention layers\n        self.attn_layers = nn.ModuleList(attn_layers)\n        # List of convolutional layers\n        self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None\n        self.norm = norm_layer\n\n    def forward(self, x, attn_mask=None, tau=None, delta=None):\n        # x [B, L, D] - Input data with shape [batch_size, sequence_length, feature_dim]\n        attns = []\n        if self.conv_layers is not None:\n            for i, (attn_layer, conv_layer) in enumerate(zip(self.attn_layers, self.conv_layers)):\n                delta = delta if i == 0 else None # Use delta only for the first layer\n                x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)\n                x = conv_layer(x)\n                attns.append(attn)\n            x, attn = self.attn_layers[-1](x, tau=tau, delta=None)\n            attns.append(attn)\n        else:\n            for attn_layer in self.attn_layers:\n                x, attn = attn_layer(x, attn_mask=attn_mask, tau=tau, delta=delta)\n                attns.append(attn)\n\n        if self.norm is not None:\n            x = self.norm(x)\n\n        return x, attns\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nclass Model(nn.Module):\n    \"\"\"\n    Paper link: https://arxiv.org/abs/2310.06625\n    \"\"\"\n    def __init__(self, configs):\n        super(Model, self).__init__()\n        self.seq_len = configs.seq_len\n        self.output_attention = configs.output_attention\n        \n        # Embedding\n        self.enc_embedding = DataEmbedding_inverted(configs.in_channel, configs.d_model, configs.dropout)\n        \n        # Encoder\n        self.encoder = Encoder(\n            [\n                EncoderLayer(\n                    AttentionLayer(\n                        FullAttention(False, attention_dropout=configs.dropout, output_attention=configs.output_attention), \n                        configs.d_model, configs.n_heads\n                    ),\n                    configs.d_model,\n                    configs.d_ff,\n                    dropout=configs.dropout,\n                    activation=configs.activation\n                ) for _ in range(configs.n_layers)\n            ],\n            norm_layer=torch.nn.LayerNorm(configs.d_model)\n        )\n        \n        # Projection\n        self.projection = nn.Linear(configs.d_model, configs.out_channel, bias=True)\n\n    def forecast(self, x_enc):\n        # Embedding\n        enc_out = self.enc_embedding(x_enc, None)\n        enc_out, attns = self.encoder(enc_out, attn_mask=None)\n\n        dec_out = self.projection(enc_out)\n        return dec_out\n\n    def forward(self, x_enc):\n        dec_out = self.forecast(x_enc)\n        return dec_out  # [B, L, C]\n\nclass Args():\n    def __init__(self):\n        self.in_channel = 33  # Number of input channels\n        self.d_model = 512  # Hidden layer dimension of the Transformer\n        self.n_heads = 8  # Number of heads in multi-head attention mechanism\n        self.d_ff = 2048  # Dimension of the feed-forward network\n        self.n_layers = 6  # Number of encoder layers\n        self.dropout = 0.2  # Dropout probability\n        self.output_attention = False  # Whether to output attention weights\n        self.activation = \"relu\"  # Activation function\n        self.out_channel = 14  # Output feature dimension\n        self.seq_len = 60  # Length of the input sequence\n        \nconfigs = Args()\nmodel = Model(configs).to(device)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## model: CNN+LSTM","metadata":{}},{"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self, configs):\n        super(Model, self).__init__()\n        # Add a 1D convolutional layer\n        self.conv1 = nn.Conv1d(\n            in_channels=configs.in_channel, \n            out_channels=configs.cnn_out_channels, \n            kernel_size=configs.kernel_size, \n            padding=configs.kernel_size // 2  # Keep the sequence length unchanged\n        )\n        self.gelu = nn.GELU()\n        self.lstm = nn.LSTM(\n            input_size=configs.cnn_out_channels,  # LSTM input = CNN output\n            hidden_size=configs.d_model, \n            num_layers=configs.n_layers, \n            batch_first=True, \n            dropout=configs.dropout, \n            bidirectional=configs.bidirectional\n        )\n        self.projection = nn.Linear(\n            in_features=configs.d_model * 2 if configs.bidirectional else configs.d_model, \n            out_features=configs.out_channel\n        )\n\n    def forward(self, x):\n        x = x.transpose(1, 2)  # Transpose dimensions to (batch_size, in_channel, seq_len)\n        x = self.conv1(x)\n        x = self.gelu(x)\n        x = x.transpose(1, 2)  # Transpose dimensions back to (batch_size, seq_len, cnn_out_channels)\n        x, _ = self.lstm(x)\n        x = self.gelu(x)  \n        return self.projection(x)\n\nclass Args():\n    def __init__(self):\n        self.in_channel = 33 \n        self.cnn_out_channels = 128  # Number of output channels for the convolutional layer\n        self.kernel_size = 3  # Size of the convolutional kernel\n        self.d_model = 512  # Number of hidden units in the LSTM\n        self.n_layers = 4  # Number of LSTM layers\n        self.dropout = 0.2  \n        self.bidirectional = True  # Whether the LSTM is bidirectional\n        self.out_channel = 6 + 8  # Output feature dimension\n\nargs = Args()\nmodel = Model(args).to(device)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# train","metadata":{}},{"cell_type":"code","source":"torch.cuda.empty_cache()\n\n# Setup basic configuration for logging\nlogging.basicConfig(filename='training_log.txt', level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n\ndef train_model(model, train_loader, criterion, optimizer, device):\n    model.train()\n    total_loss = 0\n    for batch_idx, (data, target) in enumerate(train_loader):\n        data, target = data.to(device), target.to(device)\n        optimizer.zero_grad()\n        output = model(data)\n        loss = criterion(output, target)\n        loss.backward()\n        optimizer.step()\n        total_loss += loss.item()\n        if (batch_idx + 1) % 100 == 0:  # Print training loss every 100 batches\n            logging.info(f'Batch {batch_idx + 1}, Train Loss: {total_loss / (batch_idx + 1):.4f}')\n            #print(f'Batch {batch_idx + 1}, Train Loss: {total_loss / (batch_idx + 1):.4f}')\n    return total_loss / len(train_loader)\n\ndef validate_model(model, valid_loader, criterion, device):\n    model.eval()\n    total_loss = 0\n    with torch.no_grad():\n        for data, target in valid_loader:\n            data, target = data.to(device), target.to(device)\n            output = model(data)\n            loss = criterion(output, target)\n            total_loss += loss.item()\n    return total_loss / len(valid_loader)\n\n# Separate weight decay\n#optimizer = AdamW(model.parameters(), lr=0.001, weight_decay=0.01)  \noptimizer = Adam(model.parameters(), lr=0.0001)\n\ncriterion = nn.MSELoss()\n\n# Define the number of epochs as the T_max for the scheduler\nscheduler = CosineAnnealingLR(optimizer, T_max=EPOCH)\n\npatience = PATIENCE\npatience_counter = 0\nbest_loss = float('inf')\nbest_model_state = None\n\n\ndef save_model_if_improved(validation_loss, best_loss, model, epoch):\n    if validation_loss < best_loss:\n        logging.info(f'Epoch {epoch + 1}: New best model saved with validation loss {validation_loss:.4f}')\n        torch.save(model.state_dict(), 'best_model.pth')\n        return validation_loss, 0  # reset patience counter\n    return best_loss, 1  # increment patience counter\n\nfor epoch in range(EPOCH):\n    train_loss = train_model(model, train_loader, criterion, optimizer, device)\n    valid_loss = validate_model(model, valid_loader, criterion, device)\n    \n    logging.info(f'Epoch {epoch+1}: Train Loss: {train_loss:.4f}, Valid Loss: {valid_loss:.4f}')\n    \n    scheduler.step()\n    current_lr = optimizer.param_groups[0]['lr']  # Get the current learning rate\n    logging.info(f'Epoch {epoch+1}: Learning Rate adjusted to {current_lr:.6f}')\n    \n    best_loss, improvement = save_model_if_improved(valid_loss, best_loss, model, epoch)\n    patience_counter = 0 if improvement == 0 else patience_counter + 1\n    \n    if patience_counter >= patience:\n        logging.info(\"Stopping early due to no improvement in validation loss.\")\n        break\n\n# Restore the best model state\nif best_model_state:\n    model.load_state_dict(best_model_state)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Save model\ntorch.save(model.state_dict(), 'best_model.pth')\n\n# Save scalers and feature indices\nnp.save('mx.npy', mx)\nnp.save('sx.npy', sx)\nnp.save('my.npy', my)\nnp.save('sy.npy', sy)\n\nimport json\nwith open('cols_config.json', 'w') as f:\n    config = {\n        'FEAT_COLS': FEAT_COLS,\n        'multi_feat_indices': multi_feat_indices,\n        'single_feat_indices': single_feat_indices,\n        'TARGET_COLS': TARGET_COLS\n    }\n    json.dump(config, f)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# test","metadata":{}},{"cell_type":"code","source":"import random, sys, gc, warnings, math\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nimport warnings\n\nfrom torch.optim import Adam\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nfrom torch import nn, save\n\nwarnings.filterwarnings('ignore', category=FutureWarning)\n\nMIN_STD = 1e-8\nBATCH_SIZE = 1024\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n\nLOCAL = False\n\n# Load model\nmodel.load_state_dict(torch.load('/kaggle/input/leap-infer/best_model.pth'))\nmodel.eval()\nmodel.to(device)\n\n# Load scalers and feature indices\nmx = np.load('/kaggle/input/leap-infer/mx.npy')\nsx = np.load('/kaggle/input/leap-infer/sx.npy')\nmy = np.load('/kaggle/input/leap-infer/my.npy')\nsy = np.load('/kaggle/input/leap-infer/sy.npy')\n\nimport json\nwith open('/kaggle/input/leap-infer/cols_config.json', 'r') as f:\n    config = json.load(f)\nFEAT_COLS = config['FEAT_COLS']\nmulti_feat_indices = config['multi_feat_indices']\nsingle_feat_indices = config['single_feat_indices']\nTARGET_COLS = config['TARGET_COLS']\n\ndef reshape_data(x, multi_feat_indices, single_feat_indices):\n    num_samples = x.shape[0]\n    # Create a new array with the second dimension being the time steps (60) and the third dimension being the number of features (33)\n    reshaped_data = np.zeros((num_samples, 60, 33), dtype=np.float32)\n    \n    i = 0\n\n    # Iterate through multi-dimensional feature groups\n    for _, indices in multi_feat_indices:\n        reshaped_data[:, :, i] = x[:, indices]\n        i += 1\n\n    # Iterate through single-dimensional features\n    for _, index in single_feat_indices:\n        reshaped_data[:, :, i] = np.tile(x[:, index], (60, 1)).T\n        i += 1\n\n    return reshaped_data\n    \ndef process_data(df):\n    # Temperature gradient, perform differencing by column\n    temp_cols = [f'state_t_{i}' for i in range(60)]\n    temp_diff_exprs = [pl.lit(0).alias('diff_state_t_0')] + [\n        (pl.col(temp_cols[i]) - pl.col(temp_cols[i - 1])).fill_null(0).alias(f'diff_state_t_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(temp_diff_exprs)\n\n    # Synthesize wind speed\n    wind_speed_exprs = [\n        (pl.col(f\"state_u_{i}\")**2 + pl.col(f\"state_v_{i}\")**2).sqrt().alias(f\"wind_speed_{i}\")\n        for i in range(60)\n    ]\n    df = df.with_columns(wind_speed_exprs)\n\n    # Humidity gradient, perform differencing by column\n    humidity_cols = [f'state_q0001_{i}' for i in range(60)]\n    humidity_diff_exprs = [pl.lit(0).alias('diff_state_q0001_0')] + [\n        (pl.col(humidity_cols[i]) - pl.col(humidity_cols[i - 1])).fill_null(0).alias(f'diff_state_q0001_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(humidity_diff_exprs)\n\n    # Liquid gradient, perform differencing by column\n    liquid_cols = [f'state_q0002_{i}' for i in range(60)]\n    liquid_diff_exprs = [pl.lit(0).alias('diff_state_q0002_0')] + [\n        (pl.col(liquid_cols[i]) - pl.col(liquid_cols[i - 1])).fill_null(0).alias(f'diff_state_q0002_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(liquid_diff_exprs)\n    \n    # Ice gradient, perform differencing by column\n    ice_cols = [f'state_q0003_{i}' for i in range(60)]\n    ice_diff_exprs = [pl.lit(0).alias('diff_state_q0003_0')] + [\n        (pl.col(ice_cols[i]) - pl.col(ice_cols[i - 1])).fill_null(0).alias(f'diff_state_q0003_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(ice_diff_exprs)\n\n    # Wind shear/gradient\n    wind_speed_diff_cols = [f'wind_speed_{i}' for i in range(60)]\n    wind_speed_diff_exprs = [pl.lit(0).alias('diff_wind_speed_0')] + [\n        (pl.col(wind_speed_diff_cols[i]) - pl.col(wind_speed_diff_cols[i - 1])).fill_null(0).alias(f'diff_wind_speed_{i}')\n        for i in range(1, 60)\n    ]\n    df = df.with_columns(wind_speed_diff_exprs)\n\n    # Total thermal flux = Sensible heat + Latent heat\n    total_thermal_flux = (df['pbuf_LHFLX'] + df['pbuf_SHFLX']).alias(\"total_thermal_flux\")\n    df = df.with_columns(total_thermal_flux)\n\n    # Effective solar radiation = Solar insolation × cos(Solar zenith angle) × (1 - Surface albedo)\n    effective_solar_radiation = (\n        df['pbuf_SOLIN'] * df['pbuf_COSZRS'] * (1 - df['cam_in_ASDIR'])\n    ).alias(\"effective_solar_radiation\")\n    df = df.with_columns(effective_solar_radiation)\n\n    return df","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not LOCAL:\n    df_test = pl.read_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/test.csv\")\n    \n    df_test = process_data(df_test)\n    \n    for col in FEAT_COLS:\n        df_test = df_test.with_columns(pl.col(col).cast(pl.Float32))\n\n    x_test = df_test.select(FEAT_COLS).to_numpy()\n\n    # Standardize the data\n    x_test = (x_test - mx.reshape(1, -1)) / sx.reshape(1, -1)\n\n    # Reshape the test data using the reshaping function defined during training\n    x_test = reshape_data(x_test, multi_feat_indices, single_feat_indices)\n\n    # Batched prediction\n    batch_size = 1024  \n    predt = np.zeros((x_test.shape[0], 368), dtype=np.float32)\n    for i in range(0, x_test.shape[0], batch_size):\n        i2 = min(i + batch_size, x_test.shape[0])\n        inputs = torch.from_numpy(x_test[i:i2]).float().to(device)\n        with torch.no_grad():\n            outputs = model(inputs).cpu().numpy()\n            #print(\"Shape of outputs:\", outputs.shape)\n            # Process the outputs to match the submission format\n            for j in range(6):  # Process 60 outputs for each multi-dimensional feature\n                predt[i:i2, j*60:(j+1)*60] = outputs[:, :, j]\n            for k in range(6, 14):  # Directly assign values for single-dimensional features\n                predt[i:i2, 360+k-6] = outputs[:, 0, k]\n    \n    \n    # Reverse standardization and process specific columns\n    predt = predt * sy.reshape(1, -1) + my.reshape(1, -1)\n    \n    for i in range(sy.shape[0]):\n        if sy[i] < MIN_STD * 1.1:\n            predt[:, i] = 0  # Set columns with little change to zero\n\n    # Process specific columns\n    ss = pd.read_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv\")\n    ss.iloc[:, 1:] = predt\n    \n    del predt\n    gc.collect()\n    \n    ss2 = pd.read_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv\")\n    df_test = df_test.to_pandas()\n    use_cols = [f\"ptend_q0002_{i}\" for i in range(27)]\n    for col in use_cols:\n        ss[col] = -df_test[col.replace(\"ptend\", \"state\")] * ss2[col] / 1200.\n\n    # Prepare the submission file\n    test_polars = pl.from_pandas(ss[[\"sample_id\"] + TARGET_COLS])\n    test_polars.write_csv(\"submission.csv\")\n    print(\"Submission file created.\")\n\n\n\ntest_polars.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}