{"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}],"dockerImageVersionId":30733,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\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\nfrom tqdm.notebook import tqdm\n\nprint(tf.__version__)\nprint(jax.__version__)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2024-07-02T02:50:05.226436Z","iopub.execute_input":"2024-07-02T02:50:05.226806Z","iopub.status.idle":"2024-07-02T02:50:18.189595Z","shell.execute_reply.started":"2024-07-02T02:50:05.226775Z","shell.execute_reply":"2024-07-02T02:50:18.188615Z"},"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-02T02:50:22.71444Z","iopub.execute_input":"2024-07-02T02:50:22.714796Z","iopub.status.idle":"2024-07-02T02:50:22.720168Z","shell.execute_reply.started":"2024-07-02T02:50:22.714767Z","shell.execute_reply":"2024-07-02T02:50:22.719287Z"},"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":{"execution":{"iopub.status.busy":"2024-07-02T02:50:44.052085Z","iopub.execute_input":"2024-07-02T02:50:44.052451Z","iopub.status.idle":"2024-07-02T02:50:44.057832Z","shell.execute_reply.started":"2024-07-02T02:50:44.052424Z","shell.execute_reply":"2024-07-02T02:50:44.056605Z"},"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-02T02:51:14.251301Z","iopub.execute_input":"2024-07-02T02:51:14.251681Z","iopub.status.idle":"2024-07-02T02:51:14.255985Z","shell.execute_reply.started":"2024-07-02T02:51:14.25165Z","shell.execute_reply":"2024-07-02T02:51:14.255035Z"},"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-02T02:51:40.617317Z","iopub.execute_input":"2024-07-02T02:51:40.617716Z","iopub.status.idle":"2024-07-02T02:51:40.733524Z","shell.execute_reply.started":"2024-07-02T02:51:40.617688Z","shell.execute_reply":"2024-07-02T02:51:40.732596Z"},"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    return tf.io.parse_single_example(example_proto, feature_description)['x'], tf.io.parse_single_example(example_proto, feature_description)['targets']","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:52:06.975476Z","iopub.execute_input":"2024-07-02T02:52:06.976176Z","iopub.status.idle":"2024-07-02T02:52:06.981704Z","shell.execute_reply.started":"2024-07-02T02:52:06.976144Z","shell.execute_reply":"2024-07-02T02:52:06.980746Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_files = [os.path.join(DATA_TFREC, f\"train_{i:03d}.tfrec\") for i in range(100)]\nvalid_files = [os.path.join(DATA_TFREC, \"train_100.tfrec\")]","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:52:25.555447Z","iopub.execute_input":"2024-07-02T02:52:25.555812Z","iopub.status.idle":"2024-07-02T02:52:25.560924Z","shell.execute_reply.started":"2024-07-02T02:52:25.555785Z","shell.execute_reply":"2024-07-02T02:52:25.560008Z"},"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(lambda file: tf.data.TFRecordDataset(file).map(_parse_function, num_parallel_calls=tf.data.AUTOTUNE),\n                num_parallel_calls=tf.data.AUTOTUNE, cycle_length=10, block_length=1000, deterministic=True)\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)\n","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:53:56.086611Z","iopub.execute_input":"2024-07-02T02:53:56.087142Z","iopub.status.idle":"2024-07-02T02:53:58.524854Z","shell.execute_reply.started":"2024-07-02T02:53:56.087091Z","shell.execute_reply":"2024-07-02T02:53:58.524082Z"},"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 1000))","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:57:26.944424Z","iopub.execute_input":"2024-07-02T02:57:26.945343Z","iopub.status.idle":"2024-07-02T02:57:28.995681Z","shell.execute_reply.started":"2024-07-02T02:57:26.94531Z","shell.execute_reply":"2024-07-02T02:57:28.994676Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Convert tensors to numpy arrays and then squeeze\nmean_x = norm_x.mean.numpy().squeeze()\nvariance_x = (norm_x.variance.numpy() ** 0.5).squeeze()\n\nplt.scatter(mean_x, variance_x, marker=\".\", alpha=0.5)\nplt.xscale('log')\nplt.yscale('log')","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:57:42.742492Z","iopub.execute_input":"2024-07-02T02:57:42.742856Z","iopub.status.idle":"2024-07-02T02:57:43.266813Z","shell.execute_reply.started":"2024-07-02T02:57:42.742829Z","shell.execute_reply":"2024-07-02T02:57:43.265817Z"},"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 1000))","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:58:13.458778Z","iopub.execute_input":"2024-07-02T02:58:13.459154Z","iopub.status.idle":"2024-07-02T02:58:15.488588Z","shell.execute_reply.started":"2024-07-02T02:58:13.459125Z","shell.execute_reply":"2024-07-02T02:58:15.487762Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mean_y = norm_y.mean.numpy()\nstdd_y = np.maximum(1e-10, norm_y.variance.numpy() ** 0.5)\n\nplt.scatter(mean_y.squeeze(), stdd_y.squeeze(), marker=\".\", alpha=0.5)\nplt.xscale('log')\nplt.yscale('log')","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:58:29.894045Z","iopub.execute_input":"2024-07-02T02:58:29.894422Z","iopub.status.idle":"2024-07-02T02:58:30.259066Z","shell.execute_reply.started":"2024-07-02T02:58:29.894393Z","shell.execute_reply":"2024-07-02T02:58:30.258149Z"},"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":{"execution":{"iopub.status.busy":"2024-07-02T02:58:45.691072Z","iopub.execute_input":"2024-07-02T02:58:45.691805Z","iopub.status.idle":"2024-07-02T02:58:49.808756Z","shell.execute_reply.started":"2024-07-02T02:58:45.691774Z","shell.execute_reply":"2024-07-02T02:58:49.80791Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"epochs = 10\nlearning_rate = 1e-3\n\nepochs_warmup = 2\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\nplt.plot([lr_scheduler(it) for it in range(0, epochs * steps_per_epoch, steps_per_epoch)])","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:59:39.624832Z","iopub.execute_input":"2024-07-02T02:59:39.625252Z","iopub.status.idle":"2024-07-02T02:59:40.362618Z","shell.execute_reply.started":"2024-07-02T02:59:39.625163Z","shell.execute_reply":"2024-07-02T02:59:40.361726Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = keras.Sequential([\n    keras.layers.Normalization(mean=norm_x.mean, variance=norm_x.variance),\n    keras.layers.Dense(1024, activation='relu'),\n    keras.layers.Dense(512, activation='relu'),\n    keras.layers.Dense(len(TARGETS))\n])\nmodel.compile(loss='mse', optimizer=keras.optimizers.Adam(lr_scheduler))\nmodel.build(tuple(ds_train.element_spec[0].shape))\nmodel.summary()","metadata":{"execution":{"iopub.status.busy":"2024-07-02T02:59:55.474164Z","iopub.execute_input":"2024-07-02T02:59:55.474999Z","iopub.status.idle":"2024-07-02T02:59:55.547122Z","shell.execute_reply.started":"2024-07-02T02:59:55.474955Z","shell.execute_reply":"2024-07-02T02:59:55.546108Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","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))","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:00:19.489693Z","iopub.execute_input":"2024-07-02T03:00:19.490071Z","iopub.status.idle":"2024-07-02T03:00:19.548737Z","shell.execute_reply.started":"2024-07-02T03:00:19.490043Z","shell.execute_reply":"2024-07-02T03:00:19.547725Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"history = 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=[keras.callbacks.ModelCheckpoint(filepath='model.keras')]\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:00:48.234989Z","iopub.execute_input":"2024-07-02T03:00:48.235352Z","iopub.status.idle":"2024-07-02T03:35:56.993004Z","shell.execute_reply.started":"2024-07-02T03:00:48.235324Z","shell.execute_reply":"2024-07-02T03:35:56.9901Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(history.history['loss'], color='tab:blue')\nplt.plot(history.history['val_loss'], color='tab:red')\nplt.yscale('log')","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:36:43.329024Z","iopub.execute_input":"2024-07-02T03:36:43.329903Z","iopub.status.idle":"2024-07-02T03:36:43.683869Z","shell.execute_reply.started":"2024-07-02T03:36:43.32986Z","shell.execute_reply":"2024-07-02T03:36:43.683023Z"},"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-02T03:36:47.305233Z","iopub.execute_input":"2024-07-02T03:36:47.305847Z","iopub.status.idle":"2024-07-02T03:36:53.521345Z","shell.execute_reply.started":"2024-07-02T03:36:47.305816Z","shell.execute_reply":"2024-07-02T03:36:53.519326Z"},"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-02T03:36:57.393709Z","iopub.execute_input":"2024-07-02T03:36:57.394554Z","iopub.status.idle":"2024-07-02T03:36:58.7509Z","shell.execute_reply.started":"2024-07-02T03:36:57.394521Z","shell.execute_reply":"2024-07-02T03:36:58.750029Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"mask = scores_valid <= 1e-3\nprint(f\"Number of under-performing targets: {sum(mask)}\")","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:37:02.10749Z","iopub.execute_input":"2024-07-02T03:37:02.107823Z","iopub.status.idle":"2024-07-02T03:37:02.113404Z","shell.execute_reply.started":"2024-07-02T03:37:02.107799Z","shell.execute_reply":"2024-07-02T03:37:02.112466Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Clipped score: {scores_valid.clip(0, 1).mean()}\")","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:37:03.595267Z","iopub.execute_input":"2024-07-02T03:37:03.595607Z","iopub.status.idle":"2024-07-02T03:37:03.600659Z","shell.execute_reply.started":"2024-07-02T03:37:03.595582Z","shell.execute_reply":"2024-07-02T03:37:03.599702Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del y_valid, p_valid\ngc.collect()","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:37:04.736705Z","iopub.execute_input":"2024-07-02T03:37:04.737419Z","iopub.status.idle":"2024-07-02T03:37:05.069083Z","shell.execute_reply.started":"2024-07-02T03:37:04.737388Z","shell.execute_reply":"2024-07-02T03:37:05.068082Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample = pl.read_csv(os.path.join(DATA, \"sample_submission.csv\"))","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:37:08.074436Z","iopub.execute_input":"2024-07-02T03:37:08.075291Z","iopub.status.idle":"2024-07-02T03:37:12.172505Z","shell.execute_reply.started":"2024-07-02T03:37:08.075261Z","shell.execute_reply":"2024-07-02T03:37:12.171688Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"df_test = (\n    pl.scan_csv(os.path.join(DATA, \"test.csv\"))\n    .select(pl.exclude(\"sample_id\"))\n    .cast(pl.Float32)\n    .collect()\n)","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:37:15.976504Z","iopub.execute_input":"2024-07-02T03:37:15.977431Z","iopub.status.idle":"2024-07-02T03:37:46.939657Z","shell.execute_reply.started":"2024-07-02T03:37:15.977396Z","shell.execute_reply":"2024-07-02T03:37:46.938622Z"},"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 = np.array(p_test)\np_test[:, mask] = mean_y[:, mask]\n\ndf_p_test = pd.DataFrame(p_test, columns=TARGETS)","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:37:50.537393Z","iopub.execute_input":"2024-07-02T03:37:50.53824Z","iopub.status.idle":"2024-07-02T03:38:24.275887Z","shell.execute_reply.started":"2024-07-02T03:37:50.538208Z","shell.execute_reply":"2024-07-02T03:38:24.275091Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for idx in range(12, 30):\n    df_p_test[f\"ptend_q0002_{idx}\"] = -df_test[f\"state_q0002_{idx}\"].to_numpy() / 1200\n\np_test = df_p_test.values","metadata":{"execution":{"iopub.status.busy":"2024-07-02T03:38:26.576066Z","iopub.execute_input":"2024-07-02T03:38:26.576432Z","iopub.status.idle":"2024-07-02T03:38:32.720155Z","shell.execute_reply.started":"2024-07-02T03:38:26.576404Z","shell.execute_reply":"2024-07-02T03:38:32.718982Z"},"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-02T03:38:34.995123Z","iopub.execute_input":"2024-07-02T03:38:34.995749Z","iopub.status.idle":"2024-07-02T03:38:57.137356Z","shell.execute_reply.started":"2024-07-02T03:38:34.995719Z","shell.execute_reply":"2024-07-02T03:38:57.136185Z"},"trusted":true},"execution_count":null,"outputs":[]}]}