{"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},{"sourceId":8962332,"sourceType":"datasetVersion","datasetId":5386327},{"sourceId":8973457,"sourceType":"datasetVersion","datasetId":5402556}],"dockerImageVersionId":30699,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nos.environ[\"KERAS_BACKEND\"] = \"jax\"\n\nimport gc\nimport numpy as np\nimport pandas as pd\nimport polars as pl\nimport matplotlib.pyplot as plt\n \nfrom keras.regularizers import l2\n\nimport tensorflow as tf\nimport jax\nimport keras\nfrom keras.layers import Conv1D, GroupNormalization, Activation\n\nfrom sklearn import metrics\n\nfrom tqdm.notebook import tqdm\n\nprint(tf.__version__)\nprint(jax.__version__)","metadata":{"execution":{"iopub.status.busy":"2024-07-08T02:07:06.663166Z","iopub.execute_input":"2024-07-08T02:07:06.663999Z","iopub.status.idle":"2024-07-08T02:07:20.234405Z","shell.execute_reply.started":"2024-07-08T02:07:06.66397Z","shell.execute_reply":"2024-07-08T02:07:20.23348Z"},"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-08T02:07:20.236123Z","iopub.execute_input":"2024-07-08T02:07:20.236626Z","iopub.status.idle":"2024-07-08T02:07:20.241642Z","shell.execute_reply.started":"2024-07-08T02:07:20.236599Z","shell.execute_reply":"2024-07-08T02:07:20.240795Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"isTrain = False\nSEED = 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-08T02:07:20.242627Z","iopub.execute_input":"2024-07-08T02:07:20.242872Z","iopub.status.idle":"2024-07-08T02:07:20.270716Z","shell.execute_reply.started":"2024-07-08T02:07:20.242849Z","shell.execute_reply":"2024-07-08T02:07:20.269938Z"},"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-08T02:07:20.272722Z","iopub.execute_input":"2024-07-08T02:07:20.273032Z","iopub.status.idle":"2024-07-08T02:07:20.279786Z","shell.execute_reply.started":"2024-07-08T02:07:20.27301Z","shell.execute_reply":"2024-07-08T02:07:20.278937Z"},"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-08T02:07:20.280756Z","iopub.execute_input":"2024-07-08T02:07:20.280996Z","iopub.status.idle":"2024-07-08T02:07:20.386127Z","shell.execute_reply.started":"2024-07-08T02:07:20.280974Z","shell.execute_reply":"2024-07-08T02:07:20.385276Z"},"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-08T02:07:20.387155Z","iopub.execute_input":"2024-07-08T02:07:20.387464Z","iopub.status.idle":"2024-07-08T02:07:20.392277Z","shell.execute_reply.started":"2024-07-08T02:07:20.38744Z","shell.execute_reply":"2024-07-08T02:07:20.391382Z"},"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-08T02:07:20.393447Z","iopub.execute_input":"2024-07-08T02:07:20.393764Z","iopub.status.idle":"2024-07-08T02:07:20.402809Z","shell.execute_reply.started":"2024-07-08T02:07:20.393736Z","shell.execute_reply":"2024-07-08T02:07:20.401851Z"},"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=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-08T02:07:20.403911Z","iopub.execute_input":"2024-07-08T02:07:20.40423Z","iopub.status.idle":"2024-07-08T02:07:22.510548Z","shell.execute_reply.started":"2024-07-08T02:07:20.404201Z","shell.execute_reply":"2024-07-08T02:07:22.509574Z"},"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-08T02:07:22.511722Z","iopub.execute_input":"2024-07-08T02:07:22.513587Z","iopub.status.idle":"2024-07-08T02:07:26.188274Z","shell.execute_reply.started":"2024-07-08T02:07:22.51356Z","shell.execute_reply":"2024-07-08T02:07:26.187215Z"},"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","metadata":{"execution":{"iopub.status.busy":"2024-07-08T02:07:26.191801Z","iopub.execute_input":"2024-07-08T02:07:26.192142Z","iopub.status.idle":"2024-07-08T02:07:27.903658Z","shell.execute_reply.started":"2024-07-08T02:07:26.192116Z","shell.execute_reply":"2024-07-08T02:07:27.90278Z"},"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-08T02:07:27.904641Z","iopub.execute_input":"2024-07-08T02:07:27.904893Z","iopub.status.idle":"2024-07-08T02:07:28.780201Z","shell.execute_reply.started":"2024-07-08T02:07:27.90487Z","shell.execute_reply":"2024-07-08T02:07:28.779381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model definition & Training","metadata":{}},{"cell_type":"code","source":"if is_interactive(): \n    epochs = 1 # \nelse:  \n    epochs = 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-08T02:07:28.781596Z","iopub.execute_input":"2024-07-08T02:07:28.782057Z","iopub.status.idle":"2024-07-08T02:07:28.787572Z","shell.execute_reply.started":"2024-07-08T02:07:28.782025Z","shell.execute_reply":"2024-07-08T02:07:28.786612Z"},"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(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-08T02:07:28.788724Z","iopub.execute_input":"2024-07-08T02:07:28.78904Z","iopub.status.idle":"2024-07-08T02:07:28.803392Z","shell.execute_reply.started":"2024-07-08T02:07:28.789015Z","shell.execute_reply":"2024-07-08T02:07:28.802629Z"},"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-08T02:07:28.804574Z","iopub.execute_input":"2024-07-08T02:07:28.805214Z","iopub.status.idle":"2024-07-08T02:07:28.823566Z","shell.execute_reply.started":"2024-07-08T02:07:28.805178Z","shell.execute_reply":"2024-07-08T02:07:28.822621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# Path to the model file\nmodel_filepath = '/kaggle/input/als-b-140-unet-1024-e10/model_epoch_08.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, 256, 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, 512, 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, 512, 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, 512, 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, 512, 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, 512, 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, 512, 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, 512, 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":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-08T02:07:28.825126Z","iopub.execute_input":"2024-07-08T02:07:28.825659Z","iopub.status.idle":"2024-07-08T02:07:40.822851Z","shell.execute_reply.started":"2024-07-08T02:07:28.82561Z","shell.execute_reply":"2024-07-08T02:07:40.821908Z"},"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":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2024-07-08T02:07:40.824055Z","iopub.execute_input":"2024-07-08T02:07:40.824374Z","iopub.status.idle":"2024-07-08T02:07:40.857941Z","shell.execute_reply.started":"2024-07-08T02:07:40.824348Z","shell.execute_reply":"2024-07-08T02:07:40.857053Z"},"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-08T02:07:40.858918Z","iopub.execute_input":"2024-07-08T02:07:40.85915Z","iopub.status.idle":"2024-07-08T02:07:40.896752Z","shell.execute_reply.started":"2024-07-08T02:07:40.85913Z","shell.execute_reply":"2024-07-08T02:07:40.896Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"y_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-08T02:07:40.897775Z","iopub.execute_input":"2024-07-08T02:07:40.898032Z","iopub.status.idle":"2024-07-08T02:08:10.978117Z","shell.execute_reply.started":"2024-07-08T02:07:40.898009Z","shell.execute_reply":"2024-07-08T02:08:10.977334Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"scores_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-08T02:08:10.979217Z","iopub.execute_input":"2024-07-08T02:08:10.979546Z","iopub.status.idle":"2024-07-08T02:08:12.333695Z","shell.execute_reply.started":"2024-07-08T02:08:10.979521Z","shell.execute_reply":"2024-07-08T02:08:12.332768Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mask = scores_valid <= 1e-3\nf\"Number of under-performing targets: {sum(mask)}\"","metadata":{"execution":{"iopub.status.busy":"2024-07-08T02:08:12.335046Z","iopub.execute_input":"2024-07-08T02:08:12.33552Z","iopub.status.idle":"2024-07-08T02:08:12.342059Z","shell.execute_reply.started":"2024-07-08T02:08:12.335485Z","shell.execute_reply":"2024-07-08T02:08:12.3411Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"f\"Clipped score: {scores_valid.clip(0, 1).mean()}\"","metadata":{"execution":{"iopub.status.busy":"2024-07-08T02:08:12.343105Z","iopub.execute_input":"2024-07-08T02:08:12.343421Z","iopub.status.idle":"2024-07-08T02:08:12.3562Z","shell.execute_reply.started":"2024-07-08T02:08:12.343396Z","shell.execute_reply":"2024-07-08T02:08:12.355297Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del y_valid, p_valid\ngc.collect();","metadata":{"execution":{"iopub.status.busy":"2024-07-08T02:08:12.357412Z","iopub.execute_input":"2024-07-08T02:08:12.357679Z","iopub.status.idle":"2024-07-08T02:08:12.591654Z","shell.execute_reply.started":"2024-07-08T02:08:12.357656Z","shell.execute_reply":"2024-07-08T02:08:12.590602Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Submission","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-08T02:08:12.592824Z","iopub.execute_input":"2024-07-08T02:08:12.59314Z","iopub.status.idle":"2024-07-08T02:08:16.172098Z","shell.execute_reply.started":"2024-07-08T02:08:12.593116Z","shell.execute_reply":"2024-07-08T02:08:16.171298Z"},"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    .collect()\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-08T02:08:16.173168Z","iopub.execute_input":"2024-07-08T02:08:16.173454Z","iopub.status.idle":"2024-07-08T02:08:40.653239Z","shell.execute_reply.started":"2024-07-08T02:08:16.173431Z","shell.execute_reply":"2024-07-08T02:08:40.652186Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"gc.collect();","metadata":{"execution":{"iopub.status.busy":"2024-07-08T02:08:40.654582Z","iopub.execute_input":"2024-07-08T02:08:40.654947Z","iopub.status.idle":"2024-07-08T02:08:40.880131Z","shell.execute_reply.started":"2024-07-08T02:08:40.654914Z","shell.execute_reply":"2024-07-08T02:08:40.879064Z"},"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-08T02:08:40.881485Z","iopub.execute_input":"2024-07-08T02:08:40.881776Z","iopub.status.idle":"2024-07-08T02:09:47.312073Z","shell.execute_reply.started":"2024-07-08T02:08:40.881751Z","shell.execute_reply":"2024-07-08T02:09:47.311106Z"},"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, 29):\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-08T02:09:47.313524Z","iopub.execute_input":"2024-07-08T02:09:47.313902Z","iopub.status.idle":"2024-07-08T02:09:48.949616Z","shell.execute_reply.started":"2024-07-08T02:09:47.31387Z","shell.execute_reply":"2024-07-08T02:09:48.948389Z"},"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-08T02:09:48.954992Z","iopub.execute_input":"2024-07-08T02:09:48.955338Z","iopub.status.idle":"2024-07-08T02:10:00.890008Z","shell.execute_reply.started":"2024-07-08T02:09:48.955298Z","shell.execute_reply":"2024-07-08T02:10:00.88891Z"},"trusted":true},"execution_count":null,"outputs":[]}]}