{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"tpu1vmV38","dataSources":[{"sourceId":56537,"databundleVersionId":8877088,"sourceType":"competition"}],"dockerImageVersionId":30734,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Explanations and Insights\nAs the [paper](https://arxiv.org/abs/2306.08754) outlines:\n\nCNN performed well on multi-dimensional features (the features with 60 dimensions),\n\nand MLP performed well on single-dimensional features (the features with 1 dimension).\n\nWhat I've done is build a **U-Net** that predicts 60-dimensional features and build an **MLP** that predicts 1-dimensional features.\n\nI have not tested the model on LB yet.\n\n# My Insight\n1. The model cannot predict ptend_q0002_25, ptend_q0002_26, ptend_q0002_27.\n2. I have tried to weight the loss, but it tends to focus more on poorly predicted features and skips the rest, which makes the overall R2 too low.\n```python\ndef weighted_loss(y_true, y_pred):\n    losses = K.mean(K.square(y_true - y_pred), axis=0)\n    w = losses / K.sum(losses)\n    return K.sum(w * losses)\n```\n3. Standard scaler is the right choice for the MLP model, but with U-Net, it is preferred to use Min-Max scaler.\n\n# Final Note\nPlease, if you have any comments or suggestions on my work, I will be happy to hear from you. Thanks, and do not forget to upvote.","metadata":{}},{"cell_type":"code","source":"!pip install polars \n!pip install -U tensorflow\n!pip install -U keras","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint\nfrom tensorflow.keras.callbacks import LearningRateScheduler, ReduceLROnPlateau\nfrom tensorflow.keras import backend as K\nimport tensorflow as tf\n\nimport os\nimport polars as pl\nimport pandas as pd\nimport numpy as np\nfrom sklearn.model_selection import KFold\nfrom sklearn.metrics import r2_score\n# from kan import *\nimport torch\nfrom matplotlib import pyplot as plt","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"\ntpu = tf.distribute.cluster_resolver.TPUClusterResolver()\nprint('Running on TPU ', tpu.master())\ntf.config.experimental_connect_to_cluster(tpu)\ntf.tpu.experimental.initialize_tpu_system(tpu)\nstrategy = tf.distribute.TPUStrategy(tpu)\n\nAUTO = tf.data.experimental.AUTOTUNE\nprint(\"REPLICAS: \", strategy.num_replicas_in_sync)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"EPOCH = 50\nBATCH_SIZE = 2048*strategy.num_replicas_in_sync\nEARLY_PATIENCE = 20\nMAX_LR = 0.0001\nPCT = 0.1\nMAX_LEARNING_RATE = 0.001\nMODEL_TYPE = 'UNET'\nVERSION = 11\nCLIP_NORM = 0.1\nDIV_FACTOR = 100","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Read data","metadata":{}},{"cell_type":"code","source":"%%time\ntrain_df = pl.read_csv('/kaggle/input/leap-atmospheric-physics-ai-climsim/train.csv')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"FEAT_COLS = train_df.columns[1:557]\nTARGET_COLS = train_df.columns[557:]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"X_train = train_df.select(FEAT_COLS).to_numpy()\ny = train_df.select(TARGET_COLS).to_numpy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Standard scalar","metadata":{}},{"cell_type":"code","source":"# # norm X\ny_min = y.min(axis=0)\ny_max = y.max(axis=0)\ny_mean = y.mean(axis=0)\nmin_std = 1e-24\n\nX_std = X_train.std(axis=0)\nX_std = np.maximum(X_std, min_std)\n\nX_mean = X_train.mean(axis=0)\n\ny_std = np.maximum(y.std(axis=0), min_std)\n\ny = y / y_std\nX_train = X_train / X_std","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model builder","metadata":{}},{"cell_type":"code","source":"from tensorflow.keras.layers import LayerNormalization, Layer, Dense, ReLU,ELU, Dropout\nfrom tensorflow import math, matmul, reshape, shape, transpose, cast, float32\nfrom tensorflow.keras.layers import Dense, Layer,Embedding\nfrom keras.backend import softmax\n\n# Implementing the Scaled-Dot Product Attention\nclass DotProductAttention(Layer):\n    def __init__(self, **kwargs):\n        super(DotProductAttention, self).__init__(**kwargs)\n\n    def call(self, queries, keys, values, d_k, mask=None):\n        # Scoring the queries against the keys after transposing the latter, and scaling\n        scores = matmul(queries, keys, transpose_b=True) / math.sqrt(cast(d_k, float32))\n\n        # Apply mask to the attention scores\n        if mask is not None:\n            scores += -1e9 * mask\n\n        # Computing the weights by a softmax operation\n        weights = softmax(scores)\n\n        # Computing the attention by a weighted sum of the value vectors\n        return matmul(weights, values)\n\n# Implementing the Multi-Head Attention\nclass MultiHeadAttention(Layer):\n    def __init__(self, h, d_k, d_v, d_model, **kwargs):\n        super(MultiHeadAttention, self).__init__(**kwargs)\n        self.attention = DotProductAttention()  # Scaled dot product attention\n        self.heads = h  # Number of attention heads to use\n        self.d_k = d_k  # Dimensionality of the linearly projected queries and keys\n        self.d_v = d_v  # Dimensionality of the linearly projected values\n        self.d_model = d_model  # Dimensionality of the model\n        self.W_q = Dense(d_k)  # Learned projection matrix for the queries\n        self.W_k = Dense(d_k)  # Learned projection matrix for the keys\n        self.W_v = Dense(d_v)  # Learned projection matrix for the values\n        self.W_o = Dense(d_model)  # Learned projection matrix for the multi-head output\n\n    def reshape_tensor(self, x, heads, flag):\n        if flag:\n            # Tensor shape after reshaping and transposing: (batch_size, heads, seq_length, -1)\n            x = reshape(x, shape=(shape(x)[0], shape(x)[1], heads, -1))\n            x = transpose(x, perm=(0, 2, 1, 3))\n        else:\n            # Reverting the reshaping and transposing operations: (batch_size, seq_length, d_k)\n            x = transpose(x, perm=(0, 2, 1, 3))\n            x = reshape(x, shape=(shape(x)[0], shape(x)[1], self.d_k))\n        return x\n\n    def call(self, queries, keys, values, mask=None):\n        # Rearrange the queries to be able to compute all heads in parallel\n        q_reshaped = self.reshape_tensor(self.W_q(queries), self.heads, True)\n        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)\n\n        # Rearrange the keys to be able to compute all heads in parallel\n        k_reshaped = self.reshape_tensor(self.W_k(keys), self.heads, True)\n        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)\n\n        # Rearrange the values to be able to compute all heads in parallel\n        v_reshaped = self.reshape_tensor(self.W_v(values), self.heads, True)\n        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)\n\n        # Compute the multi-head attention output using the reshaped queries, keys and values\n        o_reshaped = self.attention(q_reshaped, k_reshaped, v_reshaped, self.d_k, mask)\n        # Resulting tensor shape: (batch_size, heads, input_seq_length, -1)\n\n        # Rearrange back the output into concatenated form\n        output = self.reshape_tensor(o_reshaped, self.heads, False)\n        # Resulting tensor shape: (batch_size, input_seq_length, d_v)\n\n        # Apply one final linear projection to the output to generate the multi-head attention\n        # Resulting tensor shape: (batch_size, input_seq_length, d_model)\n        return self.W_o(output)\n\nclass PositionEmbeddingLayer(Layer):\n    def __init__(self, sequence_length, output_dim, **kwargs):\n        super(PositionEmbeddingLayer, self).__init__(**kwargs)\n        self.position_embedding_layer = Embedding(\n            input_dim=sequence_length, output_dim=output_dim\n        )\n\n    def call(self, inputs):\n        position_indices = tf.range(tf.shape(inputs)[-2])\n        embedded_indices = self.position_embedding_layer(position_indices)\n        return inputs + embedded_indices\n\n# Implementing the Add & Norm Layer\nclass AddNormalization(Layer):\n    def __init__(self, **kwargs):\n        super(AddNormalization, self).__init__(**kwargs)\n        self.layer_norm = LayerNormalization()  # Layer normalization layer\n\n    def call(self, x, sublayer_x):\n        # The sublayer input and output need to be of the same shape to be summed\n        add = x + sublayer_x\n\n        # Apply layer normalization to the sum\n        return self.layer_norm(add)\n\n# Implementing the Feed-Forward Layer\nclass FeedForward(Layer):\n    def __init__(self, d_ff, d_model, **kwargs):\n        super(FeedForward, self).__init__(**kwargs)\n        self.fully_connected1 = Dense(d_ff)  # First fully connected layer\n        self.fully_connected2 = Dense(d_model)  # Second fully connected layer\n        self.activation = ReLU()  # ReLU activation layer\n\n    def call(self, x):\n        # The input is passed into the two fully-connected layers, with a ReLU in between\n        x_fc1 = self.fully_connected1(x)\n\n        return self.fully_connected2(self.activation(x_fc1))\n\n# Implementing the Encoder Layer\nclass EncoderLayer(Layer):\n    def __init__(self, h, d_k, d_v, d_model, d_ff, rate, **kwargs):\n        super(EncoderLayer, self).__init__(**kwargs)\n        self.multihead_attention = MultiHeadAttention(h, d_k, d_v, d_model)\n        self.dropout1 = Dropout(rate)\n        self.add_norm1 = AddNormalization()\n        self.feed_forward = FeedForward(d_ff, d_model)\n        self.dropout2 = Dropout(rate)\n        self.add_norm2 = AddNormalization()\n\n    def call(self, x, padding_mask, training):\n        # Multi-head attention layer\n        multihead_output = self.multihead_attention(x, x, x, padding_mask)\n        # Expected output shape = (batch_size, sequence_length, d_model)\n\n        # Add in a dropout layer\n        multihead_output = self.dropout1(multihead_output, training=training)\n\n        # Followed by an Add & Norm layer\n        addnorm_output = self.add_norm1(x, multihead_output)\n        # Expected output shape = (batch_size, sequence_length, d_model)\n\n        # Followed by a fully connected layer\n        feedforward_output = self.feed_forward(addnorm_output)\n        # Expected output shape = (batch_size, sequence_length, d_model)\n\n        # Add in another dropout layer\n        feedforward_output = self.dropout2(feedforward_output, training=training)\n\n        # Followed by another Add & Norm layer\n        return self.add_norm2(addnorm_output, feedforward_output)\n\n# Implementing the Encoder\nclass Encoder(Layer):\n    def __init__(self, sequence_length, h, d_k, d_v, d_model, d_ff, n, rate, **kwargs):\n        super(Encoder, self).__init__(**kwargs)\n        self.pos_encoding = PositionEmbeddingLayer(sequence_length, d_model)\n        self.dropout = Dropout(rate)\n        self.encoder_layer = [EncoderLayer(h, d_k, d_v, d_model, d_ff, rate) for _ in range(n)]\n\n    def call(self, input_sentence, padding_mask, training):\n        # Generate the positional encoding\n        pos_encoding_output = self.pos_encoding(input_sentence)\n        # Expected output shape = (batch_size, sequence_length, d_model)\n\n        # Add in a dropout layer\n        x = self.dropout(pos_encoding_output, training=training)\n\n        # Pass on the positional encoded values to each encoder layer\n        for i, layer in enumerate(self.encoder_layer):\n            x = layer(x, padding_mask, training)\n\n        return x","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow.keras import layers, Model\n\ndef conv_block(inputs, num_filters):\n    x = layers.Conv1D(num_filters, kernel_size=3, padding=\"same\", activation=\"relu\")(inputs)\n    x = layers.Conv1D(num_filters, kernel_size=3, padding=\"same\", activation=\"relu\")(x)\n    return x\n\ndef encoder_block(inputs, num_filters):\n    x = conv_block(inputs, num_filters)\n    p = layers.MaxPooling1D(pool_size=2)(x)\n    return x, p\n\ndef attention_block(inputs, num_filters):\n    return MultiHeadAttention(8, 64, 64, num_filters)(inputs,inputs,inputs)\n\ndef decoder_block(inputs, skip_features, num_filters):\n    x = layers.Conv1DTranspose(num_filters, kernel_size=2, strides=2, padding=\"same\")(inputs)\n    x = layers.Concatenate()([x, skip_features])\n    x = conv_block(x, num_filters)\n    return x\n\ndef unet_model(input_shape):\n    inputs = layers.Input(shape=input_shape)\n    x = layers.ZeroPadding1D(padding=2)(inputs)\n    # Encoder\n    s00, p00 = encoder_block(x, 128)\n    s0, p0 = encoder_block(p00, 128)\n    s1, p1 = encoder_block(p0, 256)\n\n    # Bottleneck\n    p1 = attention_block(p1, 256)\n    p1 = attention_block(p1, 256)\n    p1 = attention_block(p1, 256)\n    \n    b1 = conv_block(p1, 256)\n\n    # Decoder\n\n    d4 = decoder_block(b1, s1, 256)\n    d5 = decoder_block(d4, s0, 128)\n    d6 = decoder_block(d5, s00, 128)\n\n    # Output\n    outputs = layers.Conv1D(14, kernel_size=1, padding=\"same\")(d6)\n    outputs = layers.Cropping1D(cropping=2)(outputs)\n    model = Model(inputs, outputs, name=\"U-Net\")\n    return model","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"unet_model((60,25)).summary()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MLP(tf.keras.Model):\n    def __init__(self):\n        super(MLP, self).__init__()\n        # self.mean = mx\n        # self.var = vx\n        # self.hidden_size = len(FEAT_COLS) + len(TARGET_COLS)\n        self.hidden_size = HIDDEN_SIZE\n        # 1. Define the layers.\n        self.inp_shape = len(FEAT_COLS)\n        # 2. Suggest values of the hyperparameters using a trial object.\n        self.n_layers = N_LAYERS\n        self.activation_fn = ACTIVATION_FUNCTION\n        inputs = tf.keras.Input(shape=(self.inp_shape,),name='input')\n        x = inputs\n        x1 = tf.keras.layers.Dense( self.hidden_size)(x)\n        x = x1\n        for i in range(self.n_layers):\n            x = tf.keras.layers.Dense(self.hidden_size)(x)\n            if self.activation_fn=='relu':\n                x = tf.keras.layers.ReLU()(x)\n            elif self.activation_fn=='elu':\n                x = tf.keras.layers.ELU()(x)\n            elif self.activation_fn=='leakyRelu':\n                x = tf.keras.layers.LeakyReLU(alpha=.15)(x)\n            x = tf.keras.layers.Dropout(0.1)(x)\n            x = tf.keras.layers.BatchNormalization()(x)\n            if (i+1) % 4 == 0:\n                x = tf.keras.layers.add([x, x1])\n                x1 = x\n        x = tf.keras.layers.Dense(self.hidden_size)(x)\n        outputs = tf.keras.layers.Dense(len(TARGET_COLS))(x)\n        self.model  = tf.keras.Model(inputs,outputs)\n        self.output_layer = tf.keras.layers.Dense(len(TARGET_COLS),activation='linear')\n\n        # self.norm_layer_y = tf.keras.layers.Normalization(axis=-1,invert=True)\n        # self.norm_layer_y.adapt(y)\n\n    # def call(self, x, training=False):\n    def call(self, x):\n        x = self.model(x)\n        return x\n\n\nclass MLP_UNET(tf.keras.Model):\n    def __init__(self):\n        super(MLP_UNET, self).__init__()\n\n        self.unet = unet_model((60,25))\n        self.hidden_size = HIDDEN_SIZE\n        # 1. Define the layers.\n        self.inp_shape = len(FEAT_COLS)\n        # 2. Suggest values of the hyperparameters using a trial object.\n        self.n_layers = N_LAYERS\n        self.activation_fn = ACTIVATION_FUNCTION\n        inputs = tf.keras.Input(shape=(self.inp_shape,),name='input')\n        x = inputs\n        x = tf.keras.layers.Dense(3 * self.hidden_size)(x)\n        x1 = tf.keras.layers.Dense(2 * self.hidden_size)(x)\n        for i in range(self.n_layers):\n            x = tf.keras.layers.Dense((2*self.hidden_size))(x)\n            if self.activation_fn=='relu':\n                x = tf.keras.layers.ReLU()(x)\n            elif self.activation_fn=='elu':\n                x = tf.keras.layers.ELU()(x)\n            elif self.activation_fn=='leakyRelu':\n                x = tf.keras.layers.LeakyReLU(alpha=.15)(x)\n            # x = tf.keras.layers.Dropout(0.1)(x)\n            x = tf.keras.layers.BatchNormalization()(x)\n            if (i+1) % 2 == 0:\n                x = tf.keras.layers.add([x, x1])\n                x1 = x\n\n        x = tf.keras.layers.Dense(3*self.hidden_size,activation= 'relu')(x)\n        outputs = tf.keras.layers.Dense(len(TARGET_COLS) - (60 * 6),activation= 'linear')(x)\n        self.ann  = tf.keras.Model(inputs,outputs)\n\n    def call(self, x):\n        x1 = x[:,0:360]\n        x2 = x[:,376:]\n        x3 = x[:,360:376]\n        x3 = tf.keras.layers.Reshape((1,16))(x3)\n        x3 =  tf.keras.layers.Concatenate(axis=1)([x3 for i in range(60)])\n        x1 = tf.keras.layers.Reshape((60,6))(x1)\n        x2 = tf.keras.layers.Reshape((60,3))(x2)\n        cnn_input = tf.keras.layers.Concatenate(axis=-1)([x1,x3,x2])\n\n        cnn_out = self.unet(cnn_input)\n        ann_out = self.ann(x)\n        cnn_out = tf.keras.layers.Flatten()(cnn_out)\n        x = tf.keras.layers.Concatenate()([cnn_out,ann_out])\n        return x\n\n\nclass CNN(tf.keras.Model):\n    def __init__(self):\n        super(CNN, self).__init__()\n        self.cnn = build_cnn((60,64))\n        self.conv = tf.keras.layers.Conv1D(64, kernel_size=1, padding=\"same\", activation=\"relu\")\n        self.batch_norm = tf.keras.layers.BatchNormalization()\n        self.conv_last = tf.keras.layers.Conv1D(14, kernel_size=1, padding=\"same\")\n        self.gavg = tf.keras.layers.GlobalAveragePooling1D(keepdims=True)\n        self.gmax = tf.keras.layers.GlobalMaxPooling1D(keepdims=True)\n        self.concat = tf.keras.layers.Concatenate()\n        self.concat_1 = tf.keras.layers.Concatenate(axis=-1)\n        self.flatten = tf.keras.layers.Flatten()\n        self.add = tf.keras.layers.Add()\n        self.avg = tf.keras.layers.GlobalAveragePooling1D()\n\n    def call(self, x):\n        x1 = x[:,0:360]\n        x2 = x[:,376:]\n        x3 = x[:,360:376]\n        x3 = tf.keras.layers.Reshape((1,16))(x3)\n        x3 = tf.tile(x3, [1, 60, 1])\n        # x3 =  tf.keras.layers.Concatenate(axis=1)([x3 for i in range(60)])\n        x1 = tf.keras.layers.Reshape((60,6))(x1)\n        x2 = tf.keras.layers.Reshape((60,3))(x2)\n        cnn_input = self.concat_1([x1,x3,x2])\n        # cnn_input = tf.transpose(x, perm=[0, 2, 1])\n        cnn_input = self.conv(cnn_input)\n        cnn_out1 = self.cnn(cnn_input)\n\n        gavg1 = self.gavg(cnn_out1)\n        # gmax1 = self.gmax(cnn_out1)\n        cnn_input2 = self.add([cnn_input,cnn_out1,gavg1])\n        cnn_input2 = self.batch_norm(cnn_input2)\n        cnn_out2 = self.cnn(cnn_input2)\n        cnn_out2 = self.add([cnn_out2,cnn_input2])\n        x = self.conv_last(cnn_out2)\n\n        seq_out = self.flatten(x[:,:,0:6])\n        lin_out = self.avg(x[:,:,6:])\n        lin_out = self.flatten(lin_out)\n        x = self.concat([seq_out,lin_out])\n        return x\nclass Transformers(tf.keras.Model):\n    def __init__(self):\n        super(Transformers, self).__init__()\n        self.concat = tf.keras.layers.Concatenate(axis=-1)\n        self.flatten = tf.keras.layers.Flatten()\n        self.avg = tf.keras.layers.GlobalAveragePooling1D(keepdims=True)\n        self.dense = tf.keras.layers.Dense(len(TARGET_COLS))\n        # if NORMALIZER == 'std':\n        #     self.norm_layer.adapt(X_train)\n\n        input_seq_length = 60  # Maximum length of the input sequence\n        h = 8  # Number of self-attention heads\n        d_k = 64  # Dimensionality of the linearly projected queries and keys\n        d_v = 64  # Dimensionality of the linearly projected values\n        d_ff = 512  # Dimensionality of the inner fully connected layer\n        d_model = 25  # Dimensionality of the model sub-layers' outputs\n        n = 6  # Number of layers in the encoder stack  # Batch size from the training process\n        dropout_rate = 0.1  # Frequency of dropping the input units in the dropout layers\n\n        self.transformer_encoder = Encoder(input_seq_length, h, d_k, d_v, d_model, d_ff, n, dropout_rate)\n    def call(self, x):\n        x1 = x[:,0:360]\n        x2 = x[:,376:]\n        x3 = x[:,360:376]\n        x3 = tf.keras.layers.Reshape((1,16))(x3)\n        x3 = tf.tile(x3, [1,60,1])\n        x1 = tf.keras.layers.Reshape((60,6))(x1)\n        x2 = tf.keras.layers.Reshape((60,3))(x2)\n        seq_input = self.concat([x1,x3,x2])\n        seq_out = self.transformer_encoder(seq_input,None,True)\n        # seq_out = self.avg(seq_out)\n        seq_out = self.flatten(seq_out)\n        seq_out = self.dense(seq_out)\n        return seq_out\n\nclass UNET(tf.keras.Model):\n    def __init__(self):\n        super(UNET, self).__init__()\n        self.concat = tf.keras.layers.Concatenate(axis=-1)\n        self.flatten = tf.keras.layers.Flatten()\n        self.avg = tf.keras.layers.GlobalAveragePooling1D()\n        self.dense = tf.keras.layers.Dense(len(TARGET_COLS))\n        self.unet = unet_model((60,25))\n\n    def call(self, x):\n        x1 = x[:,0:360]\n        x2 = x[:,376:]\n        x3 = x[:,360:376]\n        x3 = tf.keras.layers.Reshape((1,16))(x3)\n        x3 = tf.tile(x3, [1,60,1])\n        x1 = tf.keras.layers.Reshape((60,6))(x1)\n        x2 = tf.keras.layers.Reshape((60,3))(x2)\n        seq_input = self.concat([x1,x3,x2])\n        out = self.unet(seq_input)\n        seq_out = self.flatten(out[:,:,0:6])\n        lin_out = self.avg(out[:,:,6:])\n        out = self.concat([seq_out,lin_out])\n        return out","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def build_model():\n    if MODEL_TYPE == 'CNN':\n        model = CNN()\n    elif MODEL_TYPE == 'MLP':\n        model = MLP()\n    elif MODEL_TYPE == 'UNET':\n        model = UNET()\n    elif MODEL_TYPE == 'MLP_UNET':\n        model = MLP_UNET()\n    elif MODEL_TYPE == 'Transformers':\n        model = Transformers()\n    else:\n        raise ValueError(\"Invalid model type\")\n    return model\n\ndef predict(model, data):\n    y_pred = model.predict(data)\n    y_pred = y_pred * y_std\n    return y_pred","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del train_df","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\ngc.collect()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Cosine scheduler","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport logging\n\nlogging.getLogger('tensorflow').setLevel(logging.ERROR)\n\nfrom tensorflow.keras.callbacks import Callback\n\nclass CosineAnnealer:\n\n    def __init__(self, start, end, steps):\n        self.start = start\n        self.end = end\n        self.steps = steps\n        self.n = 0\n\n    def step(self):\n        self.n += 1\n        cos = np.cos(np.pi * (self.n / self.steps)) + 1\n        return self.end + (self.start - self.end) / 2. * cos\n\n\nclass OneCycleScheduler(Callback):\n    \"\"\"\n    \"\"\"\n\n    def __init__(self, lr_max, steps, mom_min=0.85, mom_max=0.95, phase_1_pct=PCT, div_factor=DIV_FACTOR):\n        super(OneCycleScheduler, self).__init__()\n        lr_min = lr_max / div_factor\n        final_lr = lr_max / (div_factor * 10)\n        phase_1_steps = steps * phase_1_pct\n        phase_2_steps = steps - phase_1_steps\n\n        self.phase_1_steps = phase_1_steps\n        self.phase_2_steps = phase_2_steps\n        self.phase = 0\n        self.step = 0\n\n        self.phases = [[CosineAnnealer(lr_min, lr_max, phase_1_steps), CosineAnnealer(mom_max, mom_min, phase_1_steps)],\n                 [CosineAnnealer(lr_max, final_lr, phase_2_steps), CosineAnnealer(mom_min, mom_max, phase_2_steps)]]\n\n        self.lrs = []\n        self.moms = []\n\n    def on_train_begin(self, logs=None):\n        self.phase = 0\n        self.step = 0\n\n        self.set_lr(self.lr_schedule().start)\n        self.set_momentum(self.mom_schedule().start)\n\n    def on_train_batch_begin(self, batch, logs=None):\n        self.lrs.append(self.get_lr())\n        self.moms.append(self.get_momentum())\n\n    def on_train_batch_end(self, batch, logs=None):\n        self.step += 1\n        if self.step >= self.phase_1_steps:\n            self.phase = 1\n\n        self.set_lr(self.lr_schedule().step())\n        self.set_momentum(self.mom_schedule().step())\n\n    def get_lr(self):\n        try:\n            return tf.keras.backend.get_value(self.model.optimizer.lr)\n        except AttributeError:\n            return None\n\n    def get_momentum(self):\n        try:\n            return tf.keras.backend.get_value(self.model.optimizer.momentum)\n        except AttributeError:\n            return None\n\n    def set_lr(self, lr):\n        try:\n            tf.keras.backend.set_value(self.model.optimizer.lr, lr)\n        except AttributeError:\n            pass # ignore\n\n    def set_momentum(self, mom):\n        try:\n            tf.keras.backend.set_value(self.model.optimizer.momentum, mom)\n        except AttributeError:\n            pass # ignore\n\n    def lr_schedule(self):\n        return self.phases[self.phase][0]\n\n    def mom_schedule(self):\n        return self.phases[self.phase][1]\n\n    def plot(self):\n        ax = plt.subplot(1, 2, 1)\n        ax.plot(self.lrs)\n        ax.set_title('Learning Rate')\n        ax = plt.subplot(1, 2, 2)\n        ax.plot(self.moms)\n        ax.set_title('Momentum')\ndef scheduler(epoch, lr):\n    if epoch < 1 or epoch > 1:\n        return lr\n    else:\n        return lr / (25 * 10)\n\ndef lr_warmup_cosine_decay(global_step,\n                           warmup_steps,\n                           hold = 0,\n                           total_steps=0,\n                           start_lr=0.0,\n                           target_lr=1e-3):\n    # Cosine decay\n    learning_rate = 0.5 * target_lr * (1 + np.cos(np.pi * (global_step - warmup_steps - hold) / float(total_steps - warmup_steps - hold)))\n\n    # Target LR * progress of warmup (=1 at the final warmup step)\n    warmup_lr = target_lr * (global_step / warmup_steps)\n\n    # Choose between `warmup_lr`, `target_lr` and `learning_rate` based on whether `global_step < warmup_steps` and we're still holding.\n    # i.e. warm up if we're still warming up and use cosine decayed lr otherwise\n    if hold > 0:\n        learning_rate = np.where(global_step > warmup_steps + hold,\n                                 learning_rate, target_lr)\n\n    learning_rate = np.where(global_step < warmup_steps, warmup_lr, learning_rate)\n    return learning_rate","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# R2 score function\nCalculate R2 for every target seperatly, then take the average.","metadata":{}},{"cell_type":"code","source":"def R2(y_true, y_pred):\n    SS_res =  K.sum(K.square( y_true - y_pred ),axis=0)\n    SS_tot = K.sum(K.square( y_true - K.mean(y_true,axis=0)),axis=0)\n    R2 = ( 1 - SS_res/(SS_tot + K.epsilon()) )\n    return K.mean( R2)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"kf = KFold(n_splits=20, shuffle=True, random_state=42)\n\nfor fold, (train_idx, val_idx) in enumerate(kf.split(X_train, y)):\n\n    checkpoint_filepath = f\"folds{fold}.weights.h5\"\n\n    K.clear_session()\n    with strategy.scope():\n\n        model = build_model()\n        optimizer = tf.keras.optimizers.Adam(learning_rate=MAX_LR,clipnorm=0.1)\n\n        model.compile(optimizer=optimizer, loss='mse',metrics=[R2])\n\n    monitor = \"val_loss\"\n    sv = ModelCheckpoint(\n            checkpoint_filepath, monitor=monitor, verbose=1, save_best_only=True,\n            save_weights_only=True, mode='min', save_freq='epoch'\n    )\n    steps = (len(train_idx) // BATCH_SIZE) * EPOCH\n    lr_schedule = OneCycleScheduler(lr_max=MAX_LR,steps=steps)\n\n    early_stop = tf.keras.callbacks.EarlyStopping(\n        monitor=monitor,\n        min_delta=0.0,\n        patience=EARLY_PATIENCE,\n        mode=\"min\"\n    )\n\n    history = model.fit(X_train[train_idx], y[train_idx], verbose=1,\n                        validation_data=( X_train[val_idx], y[val_idx]),\n                        epochs=EPOCH, batch_size=BATCH_SIZE, callbacks=[lr_schedule,sv, early_stop]) \n    break","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Evaluation","metadata":{}},{"cell_type":"code","source":"val_pred = predict(model,X_train[val_idx])    \nval_true = y[val_idx].copy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"val_true = val_true * y_std","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"R2 : \",r2_score(y_true =  val_true ,y_pred = val_pred ,multioutput='uniform_average'))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"c = 0\nfor i,j in zip(TARGET_COLS,r2_score(y_true =  val_true,y_pred = val_pred,multioutput='raw_values').tolist()):\n    if j < 0.5:\n        print(i,' = ',j,' - ',c)\n    c = c +1","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"test_df = pl.read_csv('/kaggle/input/leap-atmospheric-physics-ai-climsim/test.csv')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_feats = test_df.select(FEAT_COLS).to_numpy()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_feats = test_feats  / X_std","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"test_pred = predict(model,test_feats) ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import gc\ngc.collect()\nsub = pd.read_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv\")\nsub.iloc[:,1:] = test_pred\nuse_cols = []\nfor i in range(27):\n    use_cols.append(f\"ptend_q0002_{i}\")\n\nw = pd.read_csv('/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv')\ntest_df = test_df.to_pandas()\nfor col in use_cols:\n    sub[col] = (-test_df[col.replace(\"ptend\", \"state\")]*w[col])/1200.\n\nsub2 = pd.read_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv\")\nsub.iloc[:,1:] = sub.iloc[:,1:] * sub2.iloc[:,1:]\n    \ntest_polars = pl.from_pandas(sub[[\"sample_id\"]+TARGET_COLS])\ntest_polars.write_csv(\"submission.csv\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}