{"metadata":{"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":56537,"databundleVersionId":8877088,"sourceType":"competition"},{"sourceId":8409068,"sourceType":"datasetVersion","datasetId":5004471},{"sourceId":77207,"sourceType":"modelInstanceVersion","modelInstanceId":64908,"modelId":86710}],"dockerImageVersionId":30732,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true},"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.13"},"papermill":{"default_parameters":{},"duration":28455.972829,"end_time":"2024-07-06T17:19:44.894618","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2024-07-06T09:25:28.921789","version":"2.5.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# LEAP Atmospheric Physics AI ClimSim\n\nThis notebook was created base on the \"leap-train-tfrecords\" dataset\n\nthe model was changed according to the next references \n* https://arxiv.org/abs/2306.08754\n* https://leap-stc.github.io/ml4esm-workshop/","metadata":{}},{"cell_type":"markdown","source":"Credits:\n\nNotebook inspo\n*         https://www.kaggle.com/code/enzosebiane/keras-baseline-seq2seq\n\nNotebook & tf DataSet\n*         https://www.kaggle.com/code/abiolatti/keras-baseline-seq2seq\n*         https://www.kaggle.com/datasets/abiolatti/leap-train-tfrecords","metadata":{}},{"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\nimport tensorflow as tf\nimport jax\nimport keras\n\nfrom sklearn import metrics\n\nfrom tqdm.notebook import tqdm\n\nprint(tf.__version__)\nprint(jax.__version__)","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","papermill":{"duration":12.986537,"end_time":"2024-07-06T09:25:44.628314","exception":false,"start_time":"2024-07-06T09:25:31.641777","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:30:50.703678Z","iopub.execute_input":"2024-07-13T08:30:50.704226Z","iopub.status.idle":"2024-07-13T08:31:02.588935Z","shell.execute_reply.started":"2024-07-13T08:30:50.704191Z","shell.execute_reply":"2024-07-13T08:31:02.588027Z"},"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":{"papermill":{"duration":0.017648,"end_time":"2024-07-06T09:25:44.65541","exception":false,"start_time":"2024-07-06T09:25:44.637762","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:02.590732Z","iopub.execute_input":"2024-07-13T08:31:02.591223Z","iopub.status.idle":"2024-07-13T08:31:02.596322Z","shell.execute_reply.started":"2024-07-13T08:31:02.591197Z","shell.execute_reply":"2024-07-13T08:31:02.59533Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"SEED = 42\nkeras.utils.set_random_seed(SEED)\ntf.random.set_seed(SEED)\ntf.config.experimental.enable_op_determinism()","metadata":{"papermill":{"duration":0.016453,"end_time":"2024-07-06T09:25:44.681045","exception":false,"start_time":"2024-07-06T09:25:44.664592","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:02.597823Z","iopub.execute_input":"2024-07-13T08:31:02.598507Z","iopub.status.idle":"2024-07-13T08:31:02.615629Z","shell.execute_reply.started":"2024-07-13T08:31:02.598472Z","shell.execute_reply":"2024-07-13T08:31:02.614823Z"},"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":{"papermill":{"duration":0.015675,"end_time":"2024-07-06T09:25:44.706051","exception":false,"start_time":"2024-07-06T09:25:44.690376","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:02.617504Z","iopub.execute_input":"2024-07-13T08:31:02.617916Z","iopub.status.idle":"2024-07-13T08:31:02.623242Z","shell.execute_reply.started":"2024-07-13T08:31:02.617891Z","shell.execute_reply":"2024-07-13T08:31:02.622412Z"},"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\n#print(len(TARGETS))","metadata":{"papermill":{"duration":0.111416,"end_time":"2024-07-06T09:25:44.826693","exception":false,"start_time":"2024-07-06T09:25:44.715277","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:02.723235Z","iopub.execute_input":"2024-07-13T08:31:02.723492Z","iopub.status.idle":"2024-07-13T08:31:02.829156Z","shell.execute_reply.started":"2024-07-13T08:31:02.72347Z","shell.execute_reply":"2024-07-13T08:31:02.828449Z"},"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":{"papermill":{"duration":0.017203,"end_time":"2024-07-06T09:25:44.852825","exception":false,"start_time":"2024-07-06T09:25:44.835622","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:03.937436Z","iopub.execute_input":"2024-07-13T08:31:03.938214Z","iopub.status.idle":"2024-07-13T08:31:03.943235Z","shell.execute_reply.started":"2024-07-13T08:31:03.938184Z","shell.execute_reply":"2024-07-13T08:31:03.942256Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(100)]\nvalid_files = [os.path.join(DATA_TFREC, \"train_%.3d.tfrec\" % i) for i in range(100, 101)]","metadata":{"papermill":{"duration":0.017534,"end_time":"2024-07-06T09:25:44.880582","exception":false,"start_time":"2024-07-06T09:25:44.863048","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:06.212186Z","iopub.execute_input":"2024-07-13T08:31:06.212965Z","iopub.status.idle":"2024-07-13T08:31:06.219379Z","shell.execute_reply.started":"2024-07-13T08:31:06.212924Z","shell.execute_reply":"2024-07-13T08:31:06.218394Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 2048\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":{"papermill":{"duration":2.384317,"end_time":"2024-07-06T09:25:47.27428","exception":false,"start_time":"2024-07-06T09:25:44.889963","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:08.286512Z","iopub.execute_input":"2024-07-13T08:31:08.287102Z","iopub.status.idle":"2024-07-13T08:31:10.650506Z","shell.execute_reply.started":"2024-07-13T08:31:08.287074Z","shell.execute_reply":"2024-07-13T08:31:10.649677Z"},"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 200))","metadata":{"papermill":{"duration":10.186178,"end_time":"2024-07-06T09:25:57.470774","exception":false,"start_time":"2024-07-06T09:25:47.284596","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:10.651868Z","iopub.execute_input":"2024-07-13T08:31:10.652142Z","iopub.status.idle":"2024-07-13T08:31:14.097985Z","shell.execute_reply.started":"2024-07-13T08:31:10.652119Z","shell.execute_reply":"2024-07-13T08:31:14.096976Z"},"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 200))\n\nmean_y = norm_y.mean\nstdd_y = keras.ops.maximum(1e-10, norm_y.variance ** 0.5)","metadata":{"papermill":{"duration":9.508181,"end_time":"2024-07-06T09:26:07.00668","exception":false,"start_time":"2024-07-06T09:25:57.498499","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:16.816962Z","iopub.execute_input":"2024-07-13T08:31:16.81757Z","iopub.status.idle":"2024-07-13T08:31:19.746612Z","shell.execute_reply.started":"2024-07-13T08:31:16.817522Z","shell.execute_reply":"2024-07-13T08:31:19.745602Z"},"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 1000)], 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 1000)], 0), 0, keepdims=True)","metadata":{"papermill":{"duration":73.684169,"end_time":"2024-07-06T09:27:20.718828","exception":false,"start_time":"2024-07-06T09:26:07.034659","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:19.748358Z","iopub.execute_input":"2024-07-13T08:31:19.748721Z","iopub.status.idle":"2024-07-13T08:31:23.763877Z","shell.execute_reply.started":"2024-07-13T08:31:19.748691Z","shell.execute_reply":"2024-07-13T08:31:23.763011Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Model definition & Training","metadata":{"papermill":{"duration":0.009229,"end_time":"2024-07-06T09:27:20.737728","exception":false,"start_time":"2024-07-06T09:27:20.728499","status":"completed"},"tags":[]}},{"cell_type":"code","source":"@keras.saving.register_keras_serializable(package=\"MyMetrics\", name=\"ClippedR2Score\")\nclass ClippedR2Score(keras.metrics.Metric):\n    def __init__(self, name='r2_score', **kwargs):\n        super().__init__(name=name, **kwargs)\n        self.base_metric = keras.metrics.R2Score(class_aggregation=None)\n        \n    def update_state(self, y_true, y_pred, sample_weight=None):\n        self.base_metric.update_state(y_true, y_pred, sample_weight=None)\n        \n    def result(self):\n        return keras.ops.mean(keras.ops.clip(self.base_metric.result(), 0.0, 1.0))\n        \n    def reset_states(self):\n        self.base_metric.reset_states()","metadata":{"papermill":{"duration":0.020741,"end_time":"2024-07-06T09:27:20.767519","exception":false,"start_time":"2024-07-06T09:27:20.746778","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:23.765448Z","iopub.execute_input":"2024-07-13T08:31:23.765741Z","iopub.status.idle":"2024-07-13T08:31:23.772812Z","shell.execute_reply.started":"2024-07-13T08:31:23.765717Z","shell.execute_reply":"2024-07-13T08:31:23.771902Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = 20 # 25  # 15  # 12\nlearning_rate = 1e-3\n\nepochs_warmup = 1\nepochs_ending = 2\nsteps_per_epoch = int(np.ceil(len(train_files) * 100_000 / BATCH_SIZE))\n\nlr_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":{"papermill":{"duration":1.992459,"end_time":"2024-07-06T09:27:22.769505","exception":false,"start_time":"2024-07-06T09:27:20.777046","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:24.333314Z","iopub.execute_input":"2024-07-13T08:31:24.33368Z","iopub.status.idle":"2024-07-13T08:31:24.339159Z","shell.execute_reply.started":"2024-07-13T08:31:24.33365Z","shell.execute_reply":"2024-07-13T08:31:24.3383Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"keras.utils.clear_session()","metadata":{"papermill":{"duration":0.209404,"end_time":"2024-07-06T09:27:22.988422","exception":false,"start_time":"2024-07-06T09:27:22.779018","status":"completed"},"tags":[],"execution":{"iopub.status.busy":"2024-07-13T08:31:28.245284Z","iopub.execute_input":"2024-07-13T08:31:28.245624Z","iopub.status.idle":"2024-07-13T08:31:28.432605Z","shell.execute_reply.started":"2024-07-13T08:31:28.245597Z","shell.execute_reply":"2024-07-13T08:31:28.431687Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = keras.models.load_model('/kaggle/input/leap-segmented-features/keras/leap_pret291520/1/model.keras')","metadata":{"execution":{"iopub.status.busy":"2024-07-13T08:31:31.027206Z","iopub.execute_input":"2024-07-13T08:31:31.02818Z","iopub.status.idle":"2024-07-13T08:31:34.292373Z","shell.execute_reply.started":"2024-07-13T08:31:31.028134Z","shell.execute_reply":"2024-07-13T08:31:34.291437Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"ds_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\nhistory = model.fit(\n    ds_train_target_normalized,\n    validation_data=ds_valid_target_normalized,\n    epochs=epochs,\n    verbose=1 if is_interactive() else 2,\n    callbacks=[\n        keras.callbacks.ModelCheckpoint(filepath='model.keras')\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-13T08:31:35.761136Z","iopub.execute_input":"2024-07-13T08:31:35.762056Z","iopub.status.idle":"2024-07-13T09:39:25.644681Z","shell.execute_reply.started":"2024-07-13T08:31:35.762016Z","shell.execute_reply":"2024-07-13T09:39:25.641782Z"}}},{"cell_type":"markdown","source":"plt.plot(history.history['loss'], color='tab:blue')\nplt.plot(history.history['val_loss'], color='tab:red')\nplt.yscale('log');","metadata":{}},{"cell_type":"markdown","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":{}},{"cell_type":"markdown","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":{}},{"cell_type":"markdown","source":"mask = scores_valid <= 1e-3\nf\"Number of under-performing targets: {sum(mask)}\"","metadata":{}},{"cell_type":"markdown","source":"f\"Clipped score: {scores_valid.clip(0, 1).mean()}\"","metadata":{}},{"cell_type":"markdown","source":"del y_valid, p_valid\ngc.collect();","metadata":{}},{"cell_type":"markdown","source":"# Submission","metadata":{"papermill":{"duration":0.02798,"end_time":"2024-07-06T17:17:40.189581","exception":false,"start_time":"2024-07-06T17:17:40.161601","status":"completed"},"tags":[]}},{"cell_type":"code","source":"sample = pl.read_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/sample_submission.csv\")","metadata":{"papermill":{"duration":5.02882,"end_time":"2024-07-06T17:17:45.246884","exception":false,"start_time":"2024-07-06T17:17:40.218064","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_test = (\n    pl.scan_csv(\"/kaggle/input/leap-atmospheric-physics-ai-climsim/test.csv\")\n    .select(pl.exclude(\"sample_id\"))\n    .cast(pl.Float32)\n    .collect()\n)","metadata":{"papermill":{"duration":27.72388,"end_time":"2024-07-06T17:18:13.001","exception":false,"start_time":"2024-07-06T17:17:45.27712","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#p_test = model.predict(df_test.to_numpy(), batch_size=4 * BATCH_SIZE) * stdd_y + mean_y\np_test = model.predict(df_test.to_numpy()) * stdd_y + mean_y\np_test = np.array(p_test)\n#p_test[:, mask] = mean_y[:, mask]","metadata":{"papermill":{"duration":38.984873,"end_time":"2024-07-06T17:18:52.015591","exception":false,"start_time":"2024-07-06T17:18:13.030718","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# correction of ptend_q0002 targets (from 12 to 29)\ndf_p_test = pd.DataFrame(p_test, columns=TARGETS)\n\nfor idx in range(12, 30):\n    df_p_test[f\"ptend_q0002_{idx}\"] = -df_test[f\"state_q0002_{idx}\"].to_numpy() / 1200\n    \np_test = df_p_test.values","metadata":{"papermill":{"duration":6.492216,"end_time":"2024-07-06T17:18:58.542448","exception":false,"start_time":"2024-07-06T17:18:52.050232","status":"completed"},"tags":[],"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":{"papermill":{"duration":42.73755,"end_time":"2024-07-06T17:19:41.316878","exception":false,"start_time":"2024-07-06T17:18:58.579328","status":"completed"},"tags":[],"trusted":true},"execution_count":null,"outputs":[]}]}