diff --git "a/run_AIFS_v0_2_1.ipynb" "b/run_AIFS_v0_2_1.ipynb" --- "a/run_AIFS_v0_2_1.ipynb" +++ "b/run_AIFS_v0_2_1.ipynb" @@ -153,7 +153,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 7, "metadata": { "colab": { "base_uri": "https://localhost:8080/" @@ -176,7 +176,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Initial date is 2024-12-20 06:00:00\n" + "Initial date is 2025-02-18 00:00:00\n" ] } ], @@ -195,7 +195,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 66, "metadata": { "executionInfo": { "elapsed": 215, @@ -217,6 +217,7 @@ " for date in [DATE - datetime.timedelta(hours=6), DATE]:\n", " data = ekd.from_source(\"ecmwf-open-data\", date=date, param=param, levelist=levelist)\n", " for f in data:\n", + " assert f.to_numpy().shape == (721,1440)\n", " # Open data is between -180 and 180, we need to shift it to 0-360\n", " values = np.roll(f.to_numpy(), -f.shape[1] // 2, axis=1)\n", " # Interpolate the data to from 0.25 to N320\n", @@ -243,7 +244,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 68, "metadata": { "executionInfo": { "elapsed": 186, @@ -273,7 +274,16 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 69, + "metadata": {}, + "outputs": [], + "source": [ + "sfc_data=get_open_data(param=PARAM_SFC)" + ] + }, + { + "cell_type": "code", + "execution_count": 70, "metadata": { "colab": { "base_uri": "https://localhost:8080/", @@ -325,11 +335,12 @@ "user_tz": -60 }, "id": "b8sjaN5SLEtf", - "outputId": "228d3a4d-ba64-4500-9b7e-64cfe1430401" + "outputId": "228d3a4d-ba64-4500-9b7e-64cfe1430401", + "scrolled": true }, "outputs": [], "source": [ - "fields.update(get_open_data(param=PARAM_SFC))" + "fields.update(sfc_data)" ] }, { @@ -343,7 +354,16 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 76, + "metadata": {}, + "outputs": [], + "source": [ + "pl_data=get_open_data(param=PARAM_PL, levelist=LEVELS)" + ] + }, + { + "cell_type": "code", + "execution_count": 77, "metadata": { "colab": { "base_uri": "https://localhost:8080/", @@ -388,7 +408,7 @@ }, "outputs": [], "source": [ - "fields.update(get_open_data(param=PARAM_PL, levelist=LEVELS))" + "fields.update(pl_data)" ] }, { @@ -402,7 +422,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 78, "metadata": { "executionInfo": { "elapsed": 206, @@ -424,6 +444,26 @@ " fields[f\"z_{level}\"] = gh * 9.80665" ] }, + { + "cell_type": "code", + "execution_count": 74, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "dict_keys(['10u', '10v', '2d', '2t', 'msl', 'skt', 'sp', 'tcw', 'lsm', 'z', 'slor', 'sdor', 't_1000', 't_925', 't_850', 't_700', 't_600', 't_500', 't_400', 't_300', 't_250', 't_200', 't_150', 't_100', 't_50', 'u_1000', 'u_925', 'u_850', 'u_700', 'u_600', 'u_500', 'u_400', 'u_300', 'u_250', 'u_200', 'u_150', 'u_100', 'u_50', 'v_1000', 'v_925', 'v_850', 'v_700', 'v_600', 'v_500', 'v_400', 'v_300', 'v_250', 'v_200', 'v_150', 'v_100', 'v_50', 'w_1000', 'w_925', 'w_850', 'w_700', 'w_600', 'w_500', 'w_400', 'w_300', 'w_250', 'w_200', 'w_150', 'w_100', 'w_50', 'q_1000', 'q_925', 'q_850', 'q_700', 'q_600', 'q_500', 'q_400', 'q_300', 'q_250', 'q_200', 'q_150', 'q_100', 'q_50', 'z_1000', 'z_925', 'z_850', 'z_700', 'z_600', 'z_500', 'z_400', 'z_300', 'z_250', 'z_200', 'z_150', 'z_100', 'z_50'])" + ] + }, + "execution_count": 74, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "fields.keys()" + ] + }, { "cell_type": "markdown", "metadata": { @@ -435,7 +475,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 79, "metadata": { "executionInfo": { "elapsed": 187, @@ -454,6 +494,26 @@ "input_state = dict(date=DATE, fields=fields)" ] }, + { + "cell_type": "code", + "execution_count": 82, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "dict_keys(['10u', '10v', '2d', '2t', 'msl', 'skt', 'sp', 'tcw', 'lsm', 'z', 'slor', 'sdor', 't_1000', 't_925', 't_850', 't_700', 't_600', 't_500', 't_400', 't_300', 't_250', 't_200', 't_150', 't_100', 't_50', 'u_1000', 'u_925', 'u_850', 'u_700', 'u_600', 'u_500', 'u_400', 'u_300', 'u_250', 'u_200', 'u_150', 'u_100', 'u_50', 'v_1000', 'v_925', 'v_850', 'v_700', 'v_600', 'v_500', 'v_400', 'v_300', 'v_250', 'v_200', 'v_150', 'v_100', 'v_50', 'w_1000', 'w_925', 'w_850', 'w_700', 'w_600', 'w_500', 'w_400', 'w_300', 'w_250', 'w_200', 'w_150', 'w_100', 'w_50', 'q_1000', 'q_925', 'q_850', 'q_700', 'q_600', 'q_500', 'q_400', 'q_300', 'q_250', 'q_200', 'q_150', 'q_100', 'q_50', 'z_1000', 'z_925', 'z_850', 'z_700', 'z_600', 'z_500', 'z_400', 'z_300', 'z_250', 'z_200', 'z_150', 'z_100', 'z_50'])" + ] + }, + "execution_count": 82, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "input_state['fields'].keys()" + ] + }, { "cell_type": "markdown", "metadata": { @@ -474,7 +534,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 83, "metadata": { "colab": { "base_uri": "https://localhost:8080/", @@ -508,12 +568,32 @@ }, "outputs": [], "source": [ - "checkpoint = {\"huggingface\":\"ecmwf/aifs-single-0.2.1\"}" + "checkpoint = {\"huggingface\":\"ecmwf/aifs-single\"}" ] }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 84, + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "{'huggingface': 'ecmwf/aifs-single'}" + ] + }, + "execution_count": 84, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "checkpoint" + ] + }, + { + "cell_type": "code", + "execution_count": 85, "metadata": { "executionInfo": { "elapsed": 193, @@ -527,24 +607,94 @@ }, "id": "FAw6hansLEth" }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Fetching 10 files: 100%|██████████| 10/10 [00:00<00:00, 20311.40it/s]\n" + ] + } + ], + "source": [ + "runner = SimpleRunner(checkpoint, device=\"cuda\")" + ] + }, + { + "cell_type": "code", + "execution_count": 40, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Thu Feb 13 09:15:28 2025 \n", + "+-----------------------------------------------------------------------------------------+\n", + "| NVIDIA-SMI 550.54.14 Driver Version: 550.54.14 CUDA Version: 12.4 |\n", + "|-----------------------------------------+------------------------+----------------------+\n", + "| GPU Name Persistence-M | Bus-Id Disp.A | Volatile Uncorr. ECC |\n", + "| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |\n", + "| | | MIG M. |\n", + "|=========================================+========================+======================|\n", + "| 0 NVIDIA A100-SXM4-40GB On | 00000000:03:00.0 Off | 0 |\n", + "| N/A 38C P0 61W / 400W | 40068MiB / 40960MiB | 0% Default |\n", + "| | | Disabled |\n", + "+-----------------------------------------+------------------------+----------------------+\n", + " \n", + "+-----------------------------------------------------------------------------------------+\n", + "| Processes: |\n", + "| GPU GI CI PID Type Process name GPU Memory |\n", + "| ID ID Usage |\n", + "|=========================================================================================|\n", + "| 0 N/A N/A 2097036 C ..._perm_a/ecm2559/hf_venv/bin/python3 40060MiB |\n", + "+-----------------------------------------------------------------------------------------+\n" + ] + } + ], + "source": [ + "!nvidia-smi" + ] + }, + { + "cell_type": "code", + "execution_count": 101, + "metadata": {}, "outputs": [ { "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "62c500d276a94f14ba3f2bef80688f4a", - "version_major": 2, - "version_minor": 0 - }, "text/plain": [ - "Fetching 10 files: 0%| | 0/10 [00:00" ] @@ -677,6 +847,13 @@ "plt.title(\"Temperature at 2m {}\".format(state[\"date\"]))\n", "plt.show()" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -692,9 +869,9 @@ "toc_visible": true }, "kernelspec": { - "display_name": "venv-aifs0.2.1", + "display_name": "hf_venv", "language": "python", - "name": "python3" + "name": "hf_venv" }, "language_info": { "codemirror_mode": { @@ -706,7 +883,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.6" + "version": "3.11.10" }, "widgets": { "application/vnd.jupyter.widget-state+json": {